mirror of
https://github.com/ROCm/composable_kernel.git
synced 2026-07-03 21:58:13 +00:00
60b276647b
[CK Tile] Rule-based configuration generation in CK Dispatcher codegen (#8157) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Motivation The CK Tile Dispatcher code generation for CK Tile Profiler relies on flat JSON files to list the generated configurations. This approach has the following problems - The JSON files are verbose - The JSON files get easily out of sync with the CK Builder .config files from which they were generated from. - The JSON file based configuration make it hard to list explicitly the rules that govern the instance generation. ## Technical Details Replaced the JSON files with a rule based configuration. To preserve the existing functionality, the `profiler` and the `tests` instance sets are generated directly from the CK Builder config files. The JSON config files are removed from source control, and the "on-the-fly" generation guarantees that the Dispatcher codegen uses up to date configurations. This is PR introduces six different rule sets for the CK Tile Dispatcher code generation 1. `profiler`: matches with the old JSON set of profiler configurations. 2. `tests`: matches with the old JSON set of tests configurations. 3. `full`: full configuration set created from a rule-based config selection 4. `full-tests`: a subset of `full` for generating configurations for convolution integration tests. 5. `tiny`: a subset of `full-tests` to produce the minimal set of configurations to test the Dispatcher codegen. 6. `default`: the default rules, which corresponds to the existing heuristic rules for configuration selection. This ensures that ML based kernel selection doesn't get broken. The main use of the `full` rule set is to define a reasonable solution space for the possible implicit GEMM configurations. We start from the configurations that allowed by the device architecture. The `full` rule set defines the relevant tile sizes for each convolution direction. From the tile size we have a curated mapping to the number of waves over the different GEMM axes, i.e., we describe how many waves each GEMM dimensions corresponds to. The GEMM-K wave tile dimension can be computed from the other parameters and does not need to be listed explicitly. An orthogonal axis to the tiling strategy is the vectorization strategy. This mainly defined by the data type and hardware as in general, we want to use the maximum possible load widths. The maximum sizes for each convolution direction variant are defined by the implicit GEMM matrix dimensions. For cases where have a low number of channels per convolution group, we need smaller vector load sizes. These are captured by the `VecStrategy` enumeration in the codegen rules. The problem with the rule based configuration selection is that we "over generate" configurations. The old JSON configurations compose approximately 25% of all configuration that the `full` rule set creates. The additional configurations are valid, but they many not provide any performance benefits. Hence, we keep the `profiler` and `tests` rule set for now to avoid building an excessive amount configurations by default. The `full` rule set can be taken into use by specifying CMake configuration flag `-D DISPATCHER_RULE_SET=full`. By default, the `tests` rule set is used, i.e., we don't change the existing bahaviour. ## Test Plan Added a new stage in the CI/CD pipeline that ensures the Dispatcher codegen rules are up to date. Otherwise the functionality is covered by the existing CI/CD tests. There are no functional changes to the convolution kernels. Only how the different instances are generated. ## Test Result If the CK Tile conv instances build without errors, the Dispatcher codegen is generating valid code. If all tests in CI/CD pipeline are passing, the Dispatcher codegen generates valid instances. ## Submission Checklist - [x] Look over the contributing guidelines at https://github.com/ROCm/ROCm/blob/develop/CONTRIBUTING.md#pull-requests.
ML Heuristic Validation Tools
This directory contains validation scripts for testing ML-based kernel selection heuristics.
Directory Structure
validation/
├── README.md # This file
├── validate_ml_heuristic.py # GEMM universal validation
└── grouped_conv/ # Grouped convolution specific
├── validate_training_shapes.py # Training data sanity check
└── validate_backward_models.py # Backward pass prediction quality
Scripts Overview
1. validate_ml_heuristic.py - GEMM Universal Validation
Purpose: Validate ML heuristic for GEMM universal operations (not grouped conv).
Usage:
python validate_ml_heuristic.py --dtype fp16 --layout rcr
python validate_ml_heuristic.py --dtype bf16 --model_dir models/gemm_universal_bf16_gfx950
What it does:
- Loads benchmark data (oracle-best results for each GEMM shape)
- Uses ML model to predict best kernel for each shape
- Compares ML selection with oracle-best to compute efficiency
- Outputs mean/median/P10/P90 efficiency statistics
When to use: Testing GEMM universal ML models on new training data or architectures.
Grouped Convolution Validation
2. grouped_conv/validate_training_shapes.py - Training Data Sanity Check
Purpose: Quick sanity check on shapes WITH multiple kernels in training data.
Usage:
cd dispatcher/heuristics/validation/grouped_conv
python validate_training_shapes.py
What it does:
- Selects 5 random training shapes with ≥5 kernels each
- For each shape:
- Gets oracle-best from training data
- Uses ML to predict best kernel
- Builds BOTH kernels (oracle + ML)
- Runs both on hardware
- Compares actual TFLOPS
Output:
- Per-shape efficiency (ML vs Oracle on hardware)
- Prediction accuracy (ML predicted TFLOPS vs actual)
- Mean efficiency across test shapes
Runtime: ~5-10 minutes (builds 10 kernels, runs on hardware)
When to use:
- Quick sanity check after model training
- Verify model isn't overfitting to training data
- Debug prediction accuracy issues
3. grouped_conv/validate_backward_models.py - Backward Pass Prediction Quality
Purpose: Quick prediction quality check for bwd_data and bwd_weight ML models.
Usage:
cd dispatcher/heuristics/validation/grouped_conv
python validate_backward_models.py
What it does:
- Loads bwd_data and bwd_weight ML models
- Tests on 5-7 hardcoded representative problems
- For each problem:
- Predicts TFLOPS for all backward kernels (compv3, mem pipelines)
- Shows top-3 predicted kernels
- Reports prediction statistics
Output:
- Top-3 predicted kernels for each problem
- Average predicted TFLOPS
- Pipeline preference (compv3 vs mem)
- Prediction confidence (gap between best and 3rd)
Runtime: <1 minute (NO hardware - prediction only)
When to use:
- Quick check after training backward models
- Verify model predictions are reasonable
- Debug backward pass heuristic issues
Note: This does NOT run on hardware - it only checks prediction quality.
Comparison Matrix
| Script | Operation | Hardware? | Shapes Tested | Runtime | Use Case |
|---|---|---|---|---|---|
validate_ml_heuristic.py |
GEMM universal | ✗ | All training | <1 min | GEMM model validation |
validate_training_shapes.py |
Grouped conv fwd | ✓ | 5 training | 5-10 min | Quick sanity check |
validate_backward_models.py |
Grouped conv bwd | ✗ | 5-7 hardcoded | <1 min | Backward prediction quality |
Typical Workflow
-
After training forward model:
# Quick check python grouped_conv/validate_training_shapes.py -
After training backward models:
python grouped_conv/validate_backward_models.py
Target Metrics
Forward Pass (Tier-1 Model)
- Mean efficiency: >90% (currently 93.05%)
- P10 efficiency: >75% (currently 79.21%)
- Kernel match rate: >70%
Backward Pass
- Mean efficiency: >85%
- Prediction accuracy: >90%
Dependencies
All scripts require:
- Trained ML models in
../models/ - Training data in
../data/ - Python packages: pandas, numpy, LightGBM, matplotlib (for plotting)
Grouped conv hardware validation scripts additionally require:
- GPU hardware (gfx950 default)
- Compiled kernels or JIT compilation support
tile_engine/ops/grouped_conv/utilities