mirror of
https://github.com/ROCm/composable_kernel.git
synced 2026-05-01 12:11:19 +00:00
f1746955fd
[CK_TILE] Reduce Register Spills in Stream-K Reductions (#4984) ## Motivation In CK Tile Stream-K, kernels using one of two non-atomic reduction strategies (i.e., linear, tree) have high register spill count, with the tree reduction generally being worse. These changes act a first step to help decrease the register spill count. ## Technical Details ### Problem 1: Unvectorized access to partials In both the linear and tree reductions, workgroups write partials results to a global buffer; another workgroup will later read this data. When the initial logic to support reading and writing to the partials buffer was added (see https://github.com/ROCm/composable_kernel/pull/3107), the tile distribution encoding used to read from and write to partials matches the register layout for the accumulator of the mfma instruction used for the kernel. Since we do not currently use the transposed register layout for the accumulator, we end with an encoding that is not optimized for writing to HBM. For example: Consider the register layout of the `v_mfma_f32_16x16x32_fp8_fp8` instruction. ```bash ./matrix_calculator.py --architecture gfx942 --instruction v_mfma_f32_16x16x32_fp8_fp8 --register-layout --C-matrix ``` <img width="1113" height="537" alt="image" src="https://github.com/user-attachments/assets/afc8f556-08cc-4224-a6e5-b5edabc5fc02" /> The above shows that threads are responsible for consecutive elements down a column of the C tile. If we use this distribution to read and write to partials with C in row major, then threads are unable to perform vectorized reads and writes. Note: thread 0 is shown in red and thread 1 is shown in green. Since the C-shuffle Epilogue only supports C in row major, reading and writing to partials is highly unoptimized. ### Problem 2: Missed opportunity for SPGR use in tree reduction loop Since the reduction occurs between workgroups, all threads in the workgroup follow the same execution paths in the tree reduction logic, hence various variables should be using SGPRs, but they are not. ### Implemented Solutions 1. Add a new tile distribution encoding that is optimized for accessing partials in HBM. This encoding does not change the data assignment to threads, it merely changes the addresses to which they write/read in the partials buffer. For example, continuing with the `v_mfma_f32_16x16x32_fp8_fp8` instruction, the new encoding would result in threads writing in the following layout: <img width="517" height="342" alt="image" src="https://github.com/user-attachments/assets/93b5e0ea-bafc-47b8-89bb-c40ba75cb202" /> This layout ensures that each thread writes along a row, enabling `buffer_{store|load}_dwordx4` instructions (i.e., vectorized accesses). This helps reduce register usage due to requiring fewer offset calculations. 2. To force SGPR usage in the tree reduction loop, I make use of CK Tile's `amd_wave_read_first_lane` which is a wrapper around `__builtin_amdgcn_readfirstlane`. This helps reduce VGPR spills in the tree reduction. _These changes do not fully eliminate register spills. Future work will aim to further reduce spills. But these changes make good progress._ ## Test Plan Added tests for different warp tile sizes to validate that the new encoding works with different `WarpGemm` variants. ## Test Result All tests pass locally on all gfx9 architectures. Some results for decreases in register spills on gfx942: (BL = baseline) | Kernel | SGPR Spill (BL) | SGPR Spill (new) | SGPR Delta | SGPR % | VGPR Spill (BL) | VGPR Spill (new) | VGPR Delta | VGPR % | |--------|------------------:|------------------:|-----------:|-------:|-------------------:|------------------:|-----------:|-------:| | fp16 linear F/F/F/T 256x256x32 2x2x1 32x32x16 | 223 | 0 | -223 | -100.0% | 21 | 20 | -1 | -4.8% | | fp16 tree F/F/F/T 256x256x32 2x2x1 32x32x16 | 233 | 11 | -222 | -95.3% | 443 | 23 | -420 | -94.8% | | fp8 linear F/F/F/F 256x256x32 2x2x1 32x32x32 | 221 | 3 | -218 | -98.6% | 12 | 6 | -6 | -50.0% | | fp8 tree F/F/F/F 256x256x32 2x2x1 32x32x32 | 230 | 14 | -216 | -93.9% | 396 | 12 | -384 | -97.0% | ## Submission Checklist - [x] Look over the contributing guidelines at https://github.com/ROCm/ROCm/blob/develop/CONTRIBUTING.md#pull-requests.
Stream-K GEMM Tile Engine Unit Tests
How It Works
This unit test system integrates tile_engine's kernel generation into automated testing:
- Uses tile_engine scripts directly: Same Python scripts that generate tile_engine kernels
- JSON-based configuration: Define test parameters in JSON files (like tile_engine)
- Build-time generation: CMake calls tile_engine scripts to generate kernel headers
- Individual test executables: Each kernel configuration becomes a separate test
- Tile_engine verification: Uses exact same error thresholds and validation as tile_engine
Tile Engine Integration
JSON Config → tile_engine Python scripts → Generated Headers → Test Executables
--list_kernels: Get available kernel configurations from JSON--gen_individual: Generate all kernel headers in parallel during CMake configuration--gen_single: Generate individual kernel header for each configuration- Same verification: Uses tile_engine's adaptive error thresholds and reference calculations
- Same patterns: Follows tile_engine's tensor initialization, stride calculation, and kernel launching
Config-Specific Test Parameters
Each test configuration can specify optimized problem sizes in its JSON file:
test_params.problem_sizes: Array of{m, n, k, split_k}configurations- CMake extraction:
extract_test_params.pygenerates config-specific test parameter files - Build integration: Each test target uses parameters appropriate for its kernel configuration
- Optimized testing: Different configs test different problem sizes that showcase their strengths
The key idea: Unit tests that use tile_engine's exact kernel generation and verification methodology instead of creating separate test infrastructure.
Test Configurations
Test configs are generated during the Generation Phase. They are stored under the build directory at test/ck_tile/gemm_streamk_tile_engine/configs. The Compute Unit (CU) count of the device is required to generate the configs. If the Generation Phase occurs on a machine without a GPU or does not contain same GPU architecture on which you will run the tests, you can manually set the CU count using the CU_COUNT option:
# Assuming you are at the root of the repo
cd build
../script/cmake-ck-dev.sh .. gfx90a -G Ninja -DCU_COUNT=100
You can reference the public whitepaper for your specific GPU to get the appropriate CU count.
If no CU_COUNT option is given and no HIP device is found, then the default value of 100 CUs will be used to determine the problem sizes tested.
1. Smoke Tests
- Purpose: Basic functionality validation for fp16/bf16/fp8/bf8 data types
- Config: 256x256x32 (for bf16/fp16) or 128x128x32 (for bf8/fp8), warp 2x2x1, warp_tile 32x32x16
- Traits: compv3 pipeline only
- Coverage: All 4 layouts (rcr, rrr, ccr, crr)
Data Type Support
- ✅ fp16, bf16, fp8, bf8: Fully supported - all layouts (rcr, rrr, ccr, crr)
- ❌ fp64: Not supported (hardware MFMA limitation)
- ⏳ fp32, pk-int4-t: Not yet supported by gemm_instance_builder (will be added later)
Test Result Behavior
Tests automatically handle unsupported configurations through runtime validation:
- PASSED: Kernel executed correctly with results within error thresholds ✅
- SKIPPED: Kernel validation returned "Arguments not supported" (expected for certain problem sizes/configurations) ⚠️
- FAILED: Actual error or incorrect computation results ❌
When a kernel's IsSupportedArgument() check fails (e.g., due to vector alignment requirements, dimension constraints, or padding limitations), the test is automatically skipped rather than failed. This allows comprehensive testing across various problem sizes while gracefully handling configurations that don't meet specific kernel requirements.