[rocm-libraries] ROCm/rocm-libraries#6327 (commit 1e7a12e)

[CK][CK TILE] Dispatcher kernel selection heuristic for
 grouped conv (#6327)

## Motivation
The ML heuristic in dispatcher does not support grouped-conv operator
yet. In this PR, the support for fwd, bdw-data, and bwd-weight
grouped-conv kernels have been added. A tile_engine utility has also
been added to compile and run any selected kernel configuration through
dispatcher infrastructure.

## Technical Details

1. Tile engine utility is added to benchmark each shape with all the
possible kernel+tile_size combinations here -
[https://github.com/ROCm/rocm-libraries/blob/users/yraparti/ck/dispatcher-grouped-conv-heuristics/projects/composablekernel/tile_engine/ops/grouped_conv/grouped_conv_full_benchmark.py](url)
2. New LGBM regressor models for grouped conv are added to models
directory. We have 3 separate models for fwd, bwd-data, and bwd-weights
[https://github.com/ROCm/rocm-libraries/tree/users/yraparti/ck/dispatcher-grouped-conv-heuristics/projects/composablekernel/dispatcher/heuristics/models](url)
3. Implemented lazy GPU initialization (dispatcher/python)
- **Issue**: ProcessPoolExecutor fork() + GPU context caused memory
access faults
- **Solution**: Mirror FMHA pattern - defer GPU initialization until
first run()
  - **Changes**:
- setup_multiple_grouped_conv_dispatchers() returns List[Path], not
loaded libs
    - GpuGroupedConvRunner.__init__() no longer calls ctypes.CDLL
    - Added _ensure_initialized() method for lazy GPU loading
    - GPU context created only on first run() call
  - **Benefit**: Parallel compilation now works without GPU conflicts
4. Addressed few miscellaneous issues such as:
  - Fixed BF16->FP16 naming bug in the dispatcher wrapper
- Added new tile sizes, and comp_v5 pipeline to the arch spec to expand
the kernel selection
- Added automatic padding support for unsupported shapes in dispatcher
runner
- Created a single source of truth between tile_engine and dispatcher
about the architecture and tile_size details
- Build a validation scripts to compare oracle_best vs ml_heuristic
comparison

## Test Plan

1. Validated fwd, bwd-data, and bwd-weight kernels with both known and
unseen data sets with up to 300 problems.
2. Ensured that test cases are added in both dispatcher and tile_engine
to validate the heuristic.

## Test Result
Results on Unseen shapes validated on gfx950
#### Forward Pass Model
- **Training Data**: 48,845 measurements across 1,372 unique problem
shapes
- **Validation Set**: 300 unseen problems from model crawler
- **Validation Performance** (vs. oracle):
  - Mean Efficiency: **93.05%**
  - Median Efficiency: **96.8%**
  - P10 Efficiency: **79.9%**

#### Backward Data Gradient (bwd_data) Model
- **Training Data**: 18,773 measurements across 891 unique problem
shapes
- **Validation Set**: 300 unseen problems from model crawler
- **Validation Performance** (vs. oracle):
  - Mean Efficiency: **93.8%**
  - Median Efficiency: **96.5%**
  - P10 Efficiency: **82.9%**

#### Backward Weight Gradient (bwd_weight) Model
- **Training Data**: 34,900 measurements across 1,508 unique problem
shapes
- **Validation Set**: 300 unseen problems from model crawler
- **Validation Performance** (vs. oracle):
  - Mean Efficiency: **96.1%**
  - Median Efficiency: **99.2%**
  - P10 Efficiency: **89.4%**

## Submission Checklist

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

Author: yraparti

dispatcher/bindings/ctypes/conv_bwdw_ctypes_lib.cpp

@@ -129,7 +129,22 @@ float conv_bwdw_run(const void* input_ptr,
         return -1.0f;
     if(!input_ptr || !grad_output_ptr || !grad_weight_ptr)
         return -1.0f; // Null data pointer would cause kernel crash
-    return run_bwd_weight_impl(input_ptr, grad_output_ptr, grad_weight_ptr, prob, stream);
+
+    try
+    {
+        return run_bwd_weight_impl(input_ptr, grad_output_ptr, grad_weight_ptr, prob, stream);
+    }
+    catch(const std::exception&)
+    {
+        // Kernel rejected args (e.g. unsupported tile/channel combo)
+        // -3.0f matches conv_ctypes_lib.cpp:316 convention
+        // -2.0f is reserved for "no kernel / not compiled for this direction"
+        return -3.0f;
+    }
+    catch(...)
+    {
+        return -3.0f;
+    }
 #else
     return -1.0f;
 #endif

dispatcher/codegen/arch_specs.json

@@ -81,7 +81,9 @@
       "warp_configs": [
         [1, 4, 1],
         [2, 2, 1],
-        [4, 1, 1]
+        [4, 1, 1],
+        [8, 2, 1],
+        [4, 4, 1]
       ],
       "warp_tile_combos": {
         "fp32_fp32_fp32": [[16, 16, 4], [16, 16, 16]],
@@ -256,8 +258,8 @@
       "int8_int8_int32": [[16, 16, 32], [32, 32, 16]]
     },
     "gfx950": {
-      "fp16_fp16_fp32": [[32, 32, 8], [16, 16, 16], [32, 32, 16], [16, 16, 32], [64, 4, 16]],
-      "bf16_bf16_fp32": [[32, 32, 8], [16, 16, 16], [32, 32, 16], [16, 16, 32], [64, 4, 16]],
+      "fp16_fp16_fp32": [[32, 32, 8], [16, 16, 16], [32, 32, 16], [16, 16, 32], [64, 4, 16], [32, 32, 32], [16, 16, 64]],
+      "bf16_bf16_fp32": [[32, 32, 8], [16, 16, 16], [32, 32, 16], [16, 16, 32], [64, 4, 16], [32, 32, 32], [16, 16, 64]],
       "fp8_fp8_fp32": [[32, 32, 16], [32, 32, 32], [16, 16, 32], [16, 16, 64], [16, 16, 128], [32, 32, 64]],
       "bf8_bf8_fp32": [[32, 32, 16], [32, 32, 32], [16, 16, 64], [16, 16, 32], [16, 16, 128], [32, 32, 64]]
     }

dispatcher/codegen/arch_specs_generated.py

@@ -1,11 +1,10 @@
-# Copyright (c) Advanced Micro Devices, Inc., or its affiliates.
 # SPDX-License-Identifier: MIT
 
 """
 AUTO-GENERATED FILE - DO NOT EDIT DIRECTLY!
 
 Generated from: arch_specs.json
-Generated at: 2026-01-05T19:34:01.224422
+Generated at: 2026-04-10T20:07:11.665064
 
 To update this file:
 1. Edit arch_specs.json
@@ -50,7 +49,7 @@
     "gfx908": [[1, 4, 1], [2, 2, 1], [4, 1, 1]],
     "gfx90a": [[1, 4, 1], [2, 2, 1], [4, 1, 1]],
     "gfx942": [[1, 4, 1], [2, 2, 1], [4, 1, 1]],
-    "gfx950": [[1, 4, 1], [2, 2, 1], [4, 1, 1]],
+    "gfx950": [[1, 4, 1], [2, 2, 1], [4, 1, 1], [8, 2, 1], [4, 4, 1]],
     "gfx1100": [[2, 4, 1], [1, 8, 1], [8, 1, 1], [4, 2, 1]],
     "gfx1200": [[2, 4, 1], [1, 8, 1], [8, 1, 1], [4, 2, 1]],
     "gfx1201": [[2, 4, 1], [1, 8, 1], [8, 1, 1], [4, 2, 1]],
@@ -226,13 +225,17 @@
             [32, 32, 16],
             [16, 16, 32],
             [64, 4, 16],
+            [32, 32, 32],
+            [16, 16, 64],
         ],
         "bf16_bf16_fp32": [
             [32, 32, 8],
             [16, 16, 16],
             [32, 32, 16],
             [16, 16, 32],
             [64, 4, 16],
+            [32, 32, 32],
+            [16, 16, 64],
         ],
         "fp8_fp8_fp32": [
             [32, 32, 16],

dispatcher/codegen/generate_arch_specs.py

@@ -230,7 +230,7 @@ def generate_cpp_header(specs: Dict[str, Any], output_path: Path):
 
     for arch, data in archs.items():
         enum_name = arch.upper().replace("GFX", "GFX_")
-        arch_enums.append(f"    {enum_name},  // {data['description']}")
+        arch_enums.append(f"    {enum_name},")
         arch_to_string_cases.append(
             f'        case GpuArch::{enum_name}: return "{arch}";'
         )
@@ -288,12 +288,12 @@ def generate_cpp_header(specs: Dict[str, Any], output_path: Path):
                 f"    if (pipeline == Pipeline::{pipeline_enum_map[pipeline]}) return {limit};"
             )
 
-    content = f"""// SPDX-License-Identifier: MIT
-// Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved.
+    content = f"""// Copyright (c) Advanced Micro Devices, Inc., or its affiliates.
+// SPDX-License-Identifier: MIT
 
 /**
  * AUTO-GENERATED FILE - DO NOT EDIT DIRECTLY!
- * 
+ *
  * Generated from: arch_specs.json
  * Generated at: {timestamp}
  *