Files
composable_kernel/dispatcher/codegen
damien-lejeune 2c0b7cbb0a [rocm-libraries] ROCm/rocm-libraries#8424 (commit debb669)
Add missing constraint in the FMHA qr async pipeline to
 enforce bk0=bk1  (#8424)

## Motivation

The purpose of this change is to add a guardrail to what values bk0 and
bk1 can take. This is to avoid ill defined sizes, silently failing and
generating NaN (or other error) at runtime.

An example of such failure can be obtained using the tile engine:

```
cd rocm-libraries/projects/composablekernel/tile_engine/ops/fmha
python fmha_benchmark.py configs/batch_prefill.json \
  --problems "1,4,1,8000,8000,256" \
  --filter "c.data_type=='bf16' and c.hdim_q==256 and c.pipeline=='qr_async' and c.mode=='group' and c.tile_n0==32 and c.tile_k0==64"
 ```

## Technical Details

The qr_async pipeline stages data in the K dimensions into LDS using a bk1-descriptor, while the (Q*K^T) gemm0 consumes bk0

## Test Plan

See command above

## Test Result

Before the change: (invalid) generate instances, error at runtime
After this change: no instance generated

## Submission Checklist

- [X] Look over the contributing guidelines at https://github.com/ROCm/ROCm/blob/develop/CONTRIBUTING.md#pull-requests.

Co-authored-by: Damien Lejeune <damien.lejeune@amd.com>
2026-06-16 07:41:58 +00:00
..

CK Tile Unified Code Generators

Single source of truth for GEMM and Grouped Convolution kernel generation.

See also: Main Dispatcher README for installation and core concepts.

Shared Infrastructure

Both GEMM and Grouped Conv generators share common code via codegen_common.py:

  • TileConfig - Dataclass for tile dimensions
  • TraitConfigBase - Base for kernel trait configurations with arch-aware validation
  • CommonTypeMappings - Dtype-to-C++ type mappings
  • parallel_generate() - Parallel kernel generation with per-kernel progress logging
  • Arch-aware expansion helpers (valid_wave_configs, valid_warp_configs, etc.)

Quick Start

GEMM

cd dispatcher/codegen

# Generate standard FP16 kernels
python3 unified_gemm_codegen.py \
    --output-dir ../build/generated_kernels \
    --datatype fp16 \
    --layout rcr \
    --variants standard

# Generate all variants
python3 unified_gemm_codegen.py \
    --output-dir ../build/generated_kernels \
    --variants standard preshuffle multi_d

Grouped Convolution

cd dispatcher/codegen

# Generate forward FP16 grouped conv kernels
python3 unified_grouped_conv_codegen.py \
    --output-dir ../build/generated_kernels \
    --datatype fp16 \
    --variant forward \
    --ndim-spatial 2

# Generate backward data kernels
python3 unified_grouped_conv_codegen.py \
    --output-dir ../build/generated_kernels \
    --variant backward_data \
    --ndim-spatial 2

Using from Python

from ctypes_utils import CodegenRunner, KernelConfig

# Generate from specific config
config = KernelConfig(tile_m=256, tile_n=256, tile_k=64)
codegen = CodegenRunner()
result = codegen.generate_from_config(config)

# Generate variant
result = codegen.generate("preshuffle")

# Generate all
results = codegen.generate_all()

Command Line Options

Option Values Description
--output-dir path Output directory
--datatype fp16, bf16, fp32, int8 Data type
--layout rcr, rrr, crr, ccr Matrix layouts
--gpu-target gfx942, gfx90a, gfx950 Target GPU
--variants standard, preshuffle, multi_d Kernel variants
--preselected fp16_rcr_essential, etc. Predefined kernel set

Layout Notation

  • R = Row-major, C = Column-major
  • Order: A, B, C (e.g., rcr = A row, B col, C row)

Variants

Standard

Basic GEMM: C = A x B

PreShuffle

Optimized weight access with LDS pre-shuffling. Best for large matrices.

Multi-D

Element-wise fusion: C = op(A x B + D0 + D1 + ...)

Supported ops: PassThrough, MultiDAdd, Relu, Gelu, Sigmoid, Tanh

Output Structure

generated_kernels/
|---- gemm_fp16_rcr_compv4_..._128x128x32_....hpp          # GEMM kernels
|---- gemm_fp16_rcr_compv4_..._preshuffle.hpp
|---- gemm_fp16_rcr_compv4_..._multid_Relu_d1.hpp
|---- grouped_conv_fwd_fp16_nhwgc_..._128x128x32_....hpp   # Grouped conv kernels
+---- ...

Configuration Files

arch_specs.json

GPU architecture specifications (single source of truth):

{
  "architectures": {
    "gfx942": {
      "family": "cdna3",
      "warp_size": 64,
      "warp_configs": [[2, 2, 1], [4, 4, 1]],
      ...
    }
  }
}

preselected_kernels.py

Curated kernel sets for common use cases.

Adding New GPU Support

See ADDING_NEW_GPU.md for complete guide.

Quick steps:

  1. Edit arch_specs.json
  2. Run python generate_arch_specs.py
  3. Rebuild

Troubleshooting

Issue Solution
"Arguments not supported" Check tile config validity
Missing element-wise op Check elementwise_ops.hpp
Compilation errors Verify C++17, include paths

More info: See ../README.md for full documentation.