mirror of
https://github.com/ROCm/composable_kernel.git
synced 2026-05-18 12:00:07 +00:00
[CK] [CK_Tile] Add GroupConv to Kernel Dispatcher (#5168)
## Motivation This PR adds CK Tile group convolution (forward, backward-data, backward-weight) support to the kernel dispatcher, matching and unifying with the existing dispatcher GEMM infrastructure in architecture and usability. The dispatcher provides a unified kernel dispatch system with both C++ and Python frontends, and until now only supported GEMM operations. This PR enables framework integrators to use the same declarative kernel workflow for convolutions as they do for GEMM: declare kernels, build a registry JIT, select kernels within the registry at runtime, and dispatch to GPU. Future PRs will include runtime kernel selection heuristics for autotuning of kernel parameters based on (problem, hardware arch). ## Technical Details Grouped convolution support has been added to the CK Tile Dispatcher with generated_conv_backend.hpp enabling dispatcher.run(in, wei, out, problem) for all 6 conv variants (fwd/bwdd/bwdw x 2D/3D), runtime heuristic kernel selection, and GroupedConvKernelKey with full ConvConfigBase fields. Python side adds parallel JIT via registry.build(max_workers) and heuristic registry.select(). Includes 7 C++ and 6 Python examples covering all directions with CPU reference validation, and shared infrastructure improvements (BaseRegistry CRTP, structured exceptions). As a sanity check, JIT compile times for a single kernel remains the same and for multiple kernels there is better parallelism: Kernels | 1 worker | 8 workers 1 | 7.7 s | 7.7 s 2 | 15.9 s | 8.2 s 4 | 33.4 s | 9.7 s 6 | 52.3 s | 10.2 s ## Test Plan 145 ephemeral unit tests have been added to test basic functionality. All 30 examples/integration tests run end-to-end on gfx950 (MI350): 7 C++ conv, 7 C++ GEMM, 6 Python conv, 10 Python GEMM. CPU reference validation for forward, backward-data, and backward-weight (2D) in both C++ and Python examples pass. ## Test Result 30 examples pass. Peak performance: 132 TFLOPS (Batch-32 forward 56x56), 53 TFLOPS (pointwise 1x1). CPU reference accuracy: max_abs_diff < 0.002 for all directions (fp16 vs fp32 reference). ## Submission Checklist - [x] Look over the contributing guidelines at https://github.com/ROCm/ROCm/blob/develop/CONTRIBUTING.md#pull-requests. --------- Co-authored-by: Yaswanth Raparti <113389104+yraparti@users.noreply.github.com>
This commit is contained in:
Vidyasagar Ananthan
committed by
GitHub
GitHub
parent
fb22cd0c69
commit
a2b844d335
@@ -7,41 +7,37 @@
|
||||
Example 01: Basic GEMM with Multiple Kernels
|
||||
|
||||
Demonstrates:
|
||||
1. Declaring multiple kernel configurations
|
||||
2. Printing all registered kernels
|
||||
3. Running each kernel and validating output
|
||||
1. Building a Registry with multiple kernel configurations
|
||||
2. Parallel JIT compilation via registry.build()
|
||||
3. Running each kernel and validating output against NumPy reference
|
||||
4. Comparing performance across kernels
|
||||
|
||||
Complexity: ★★☆☆☆
|
||||
|
||||
Usage:
|
||||
python3 01_basic_gemm.py
|
||||
python3 01_basic_gemm.py --help
|
||||
python3 01_basic_gemm.py --dtype bf16
|
||||
python3 01_basic_gemm.py --size 2048
|
||||
python3 01_basic_gemm.py --num-kernels 4
|
||||
python3 01_basic_gemm.py --workers 4
|
||||
"""
|
||||
|
||||
import sys
|
||||
import time
|
||||
import argparse
|
||||
from pathlib import Path
|
||||
from dataclasses import dataclass
|
||||
from typing import List
|
||||
|
||||
sys.path.insert(0, str(Path(__file__).parent.parent.parent.parent / "python"))
|
||||
import numpy as np
|
||||
|
||||
from ctypes_utils import (
|
||||
KernelConfig,
|
||||
setup_gemm_dispatcher,
|
||||
cleanup_gemm,
|
||||
reset_for_example,
|
||||
Registry,
|
||||
detect_gpu_arch,
|
||||
)
|
||||
|
||||
|
||||
@dataclass
|
||||
class KernelSpec:
|
||||
"""Specification for a kernel configuration"""
|
||||
|
||||
name: str
|
||||
tile_m: int
|
||||
tile_n: int
|
||||
@@ -50,80 +46,37 @@ class KernelSpec:
|
||||
scheduler: str = "intrawave"
|
||||
|
||||
|
||||
# Define multiple kernel configurations to test (50+ kernels)
|
||||
KERNEL_SPECS = [
|
||||
# Small tiles - compv3
|
||||
# Small tiles
|
||||
KernelSpec("small_64x64_k32", 64, 64, 32, "compv3"),
|
||||
KernelSpec("small_64x64_k64", 64, 64, 64, "compv3"),
|
||||
# Small tiles - compv4
|
||||
KernelSpec("small_64x64_v4_k32", 64, 64, 32, "compv4"),
|
||||
KernelSpec("small_64x64_v4_k64", 64, 64, 64, "compv4"),
|
||||
# Medium tiles - compv3
|
||||
# Medium tiles
|
||||
KernelSpec("med_128x128_k32", 128, 128, 32, "compv3"),
|
||||
KernelSpec("med_128x128_k64", 128, 128, 64, "compv3"),
|
||||
KernelSpec("med_128x128_k128", 128, 128, 128, "compv3"),
|
||||
# Medium tiles - compv4
|
||||
KernelSpec("med_128x128_v4_k32", 128, 128, 32, "compv4"),
|
||||
KernelSpec("med_128x128_v4_k64", 128, 128, 64, "compv4"),
|
||||
KernelSpec("med_128x128_v4_k128", 128, 128, 128, "compv4"),
|
||||
# Rectangular tiles - compv3
|
||||
# Rectangular tiles
|
||||
KernelSpec("rect_64x128_k32", 64, 128, 32, "compv3"),
|
||||
KernelSpec("rect_64x128_k64", 64, 128, 64, "compv3"),
|
||||
KernelSpec("rect_128x64_k32", 128, 64, 32, "compv3"),
|
||||
KernelSpec("rect_128x64_k64", 128, 64, 64, "compv3"),
|
||||
# Rectangular tiles - compv4
|
||||
KernelSpec("rect_64x128_v4_k32", 64, 128, 32, "compv4"),
|
||||
KernelSpec("rect_64x128_v4_k64", 64, 128, 64, "compv4"),
|
||||
KernelSpec("rect_128x64_v4_k32", 128, 64, 32, "compv4"),
|
||||
KernelSpec("rect_128x64_v4_k64", 128, 64, 64, "compv4"),
|
||||
# Large tiles - compv3
|
||||
# Large tiles
|
||||
KernelSpec("large_256x128_k32", 256, 128, 32, "compv3"),
|
||||
KernelSpec("large_256x128_k64", 256, 128, 64, "compv3"),
|
||||
KernelSpec("large_128x256_k32", 128, 256, 32, "compv3"),
|
||||
KernelSpec("large_128x256_k64", 128, 256, 64, "compv3"),
|
||||
KernelSpec("large_256x256_k32", 256, 256, 32, "compv3"),
|
||||
KernelSpec("large_256x256_k64", 256, 256, 64, "compv3"),
|
||||
# Large tiles - compv4
|
||||
KernelSpec("large_256x128_v4_k32", 256, 128, 32, "compv4"),
|
||||
KernelSpec("large_256x128_v4_k64", 256, 128, 64, "compv4"),
|
||||
KernelSpec("large_128x256_v4_k32", 128, 256, 32, "compv4"),
|
||||
KernelSpec("large_128x256_v4_k64", 128, 256, 64, "compv4"),
|
||||
KernelSpec("large_256x256_v4_k32", 256, 256, 32, "compv4"),
|
||||
KernelSpec("large_256x256_v4_k64", 256, 256, 64, "compv4"),
|
||||
# Interwave scheduler variants
|
||||
KernelSpec("int_64x64_k32", 64, 64, 32, "compv3", "interwave"),
|
||||
# Interwave scheduler
|
||||
KernelSpec("int_128x128_k32", 128, 128, 32, "compv3", "interwave"),
|
||||
KernelSpec("int_128x128_k64", 128, 128, 64, "compv3", "interwave"),
|
||||
KernelSpec("int_256x128_k32", 256, 128, 32, "compv3", "interwave"),
|
||||
# More tile_k variations - compv3
|
||||
KernelSpec("med_128x128_k16", 128, 128, 16, "compv3"),
|
||||
KernelSpec("rect_64x128_k16", 64, 128, 16, "compv3"),
|
||||
KernelSpec("rect_128x64_k16", 128, 64, 16, "compv3"),
|
||||
# More tile_k variations - compv4
|
||||
KernelSpec("med_128x128_v4_k16", 128, 128, 16, "compv4"),
|
||||
KernelSpec("rect_64x128_v4_k16", 64, 128, 16, "compv4"),
|
||||
KernelSpec("rect_128x64_v4_k16", 128, 64, 16, "compv4"),
|
||||
# Additional rectangular
|
||||
KernelSpec("rect_32x64_k32", 32, 64, 32, "compv3"),
|
||||
KernelSpec("rect_64x32_k32", 64, 32, 32, "compv3"),
|
||||
KernelSpec("rect_32x128_k32", 32, 128, 32, "compv3"),
|
||||
KernelSpec("rect_128x32_k32", 128, 32, 32, "compv3"),
|
||||
# Additional compv4 variants
|
||||
KernelSpec("rect_32x64_v4_k32", 32, 64, 32, "compv4"),
|
||||
KernelSpec("rect_64x32_v4_k32", 64, 32, 32, "compv4"),
|
||||
KernelSpec("rect_32x128_v4_k32", 32, 128, 32, "compv4"),
|
||||
KernelSpec("rect_128x32_v4_k32", 128, 32, 32, "compv4"),
|
||||
]
|
||||
|
||||
|
||||
def create_kernel_config(spec: KernelSpec, dtype: str, arch: str) -> KernelConfig:
|
||||
"""Create a KernelConfig from a spec"""
|
||||
# Adjust warp tiles based on tile size
|
||||
if spec.tile_m <= 64:
|
||||
warp_m, warp_n = 16, 16
|
||||
else:
|
||||
warp_m, warp_n = 32, 32
|
||||
|
||||
def spec_to_config(spec: KernelSpec, dtype: str, arch: str) -> KernelConfig:
|
||||
warp_m, warp_n = (16, 16) if spec.tile_m <= 64 else (32, 32)
|
||||
return KernelConfig(
|
||||
dtype_a=dtype,
|
||||
dtype_b=dtype,
|
||||
@@ -148,180 +101,118 @@ def create_kernel_config(spec: KernelSpec, dtype: str, arch: str) -> KernelConfi
|
||||
)
|
||||
|
||||
|
||||
def print_kernel_table(specs: List[KernelSpec], dtype: str):
|
||||
"""Print a formatted table of kernel configurations"""
|
||||
print("\n" + "=" * 70)
|
||||
print(f" DECLARED KERNEL CONFIGURATIONS ({len(specs)} kernels)")
|
||||
print("=" * 70)
|
||||
print(f"\n {'#':<3} {'Name':<18} {'Tile':<14} {'Pipeline':<10} {'Scheduler':<12}")
|
||||
print(" " + "-" * 68)
|
||||
|
||||
for i, spec in enumerate(specs, 1):
|
||||
tile = f"{spec.tile_m}x{spec.tile_n}x{spec.tile_k}"
|
||||
print(
|
||||
f" {i:<3} {spec.name:<18} {tile:<14} {spec.pipeline:<10} {spec.scheduler:<12}"
|
||||
)
|
||||
|
||||
print(" " + "-" * 68)
|
||||
print(f" Data type: {dtype}")
|
||||
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser(
|
||||
description="Basic GEMM Example with Multiple Kernels",
|
||||
formatter_class=argparse.RawDescriptionHelpFormatter,
|
||||
epilog="""
|
||||
Examples:
|
||||
python3 01_basic_gemm.py # Default FP16 with 4 kernels
|
||||
python3 01_basic_gemm.py --dtype bf16 # BF16 mode
|
||||
python3 01_basic_gemm.py --size 2048 # Larger problem size
|
||||
python3 01_basic_gemm.py --num-kernels 2 # Test only 2 kernels
|
||||
""",
|
||||
)
|
||||
parser = argparse.ArgumentParser(description="Basic GEMM with Multiple Kernels")
|
||||
parser.add_argument("--dtype", default="fp16", choices=["fp16", "bf16"])
|
||||
parser.add_argument("--arch", default=detect_gpu_arch())
|
||||
parser.add_argument("--size", type=int, default=512, help="Problem size MxNxK")
|
||||
parser.add_argument("--num-kernels", type=int, default=0, help="0 = all")
|
||||
parser.add_argument(
|
||||
"--dtype",
|
||||
default="fp16",
|
||||
choices=["fp16", "bf16", "fp32"],
|
||||
help="Data type (default: fp16)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--arch",
|
||||
default="gfx942",
|
||||
help="Target architecture (default: gfx942)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--size",
|
||||
type=int,
|
||||
default=512,
|
||||
help="Problem size MxNxK (default: 512)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--num-kernels",
|
||||
type=int,
|
||||
default=0,
|
||||
help="Number of kernels to test (0 = all)",
|
||||
"--workers", type=int, default=0, help="Max parallel JIT workers (0 = auto)"
|
||||
)
|
||||
args = parser.parse_args()
|
||||
|
||||
reset_for_example()
|
||||
|
||||
print("=" * 70)
|
||||
print("Example 01: Basic GEMM with Multiple Kernels")
|
||||
print("=" * 70)
|
||||
|
||||
# Select kernels to test
|
||||
specs = KERNEL_SPECS[: args.num_kernels] if args.num_kernels > 0 else KERNEL_SPECS
|
||||
|
||||
# =========================================================================
|
||||
# Step 1: Print all kernel configurations
|
||||
# =========================================================================
|
||||
print_kernel_table(specs, args.dtype)
|
||||
|
||||
# =========================================================================
|
||||
# Step 2: Setup and test each kernel
|
||||
# =========================================================================
|
||||
print("\n" + "=" * 70)
|
||||
print(" RUNNING KERNELS")
|
||||
print("=" * 70)
|
||||
|
||||
np_dtype = np.float16 if args.dtype in ["fp16", "bf16"] else np.float32
|
||||
M, N, K = args.size, args.size, args.size
|
||||
|
||||
results = []
|
||||
|
||||
print(f"\n Problem size: {M}x{N}x{K}\n")
|
||||
# Step 1: Build registry
|
||||
print(
|
||||
f" {'#':<3} {'Name':<18} {'Tile':<14} {'Time (ms)':>10} {'TFLOPS':>10} {'Max Err':>10} {'Status':<8}"
|
||||
f"\n {len(specs)} kernel configurations, dtype={args.dtype}, arch={args.arch}"
|
||||
)
|
||||
print(" " + "-" * 78)
|
||||
|
||||
for i, spec in enumerate(specs, 1):
|
||||
# Create unique test data per kernel
|
||||
np.random.seed(42 + i * 1000)
|
||||
A = (np.random.randn(M, K) * 0.1).astype(np_dtype)
|
||||
B = (np.random.randn(K, N) * 0.1).astype(np_dtype)
|
||||
|
||||
# Create config and setup dispatcher
|
||||
config = create_kernel_config(spec, args.dtype, args.arch)
|
||||
|
||||
setup = setup_gemm_dispatcher(
|
||||
config=config,
|
||||
registry_name=f"kernel_{spec.name}",
|
||||
verbose=False,
|
||||
auto_rebuild=True,
|
||||
print(f"\n {'#':<3} {'Name':<22} {'Tile':<14} {'Pipeline':<10} {'Scheduler':<12}")
|
||||
print(" " + "-" * 64)
|
||||
for i, s in enumerate(specs, 1):
|
||||
print(
|
||||
f" {i:<3} {s.name:<22} {s.tile_m}x{s.tile_n}x{s.tile_k:<6} {s.pipeline:<10} {s.scheduler:<12}"
|
||||
)
|
||||
|
||||
reg = Registry(name="basic_gemm")
|
||||
for s in specs:
|
||||
reg.register_kernel(spec_to_config(s, args.dtype, args.arch))
|
||||
|
||||
# Step 2: Parallel JIT build via registry.build()
|
||||
workers = args.workers if args.workers > 0 else None
|
||||
print(
|
||||
f"\n--- Parallel JIT Build ({len(specs)} kernels{f', workers={workers}' if workers else ''}) ---"
|
||||
)
|
||||
|
||||
t0 = time.perf_counter()
|
||||
setups = reg.build(verbose=False, max_workers=workers)
|
||||
jit_build_s = time.perf_counter() - t0
|
||||
|
||||
built = sum(1 for s in setups if s.success)
|
||||
print(f" Built: {built}/{len(specs)} kernels in {jit_build_s:.1f} s")
|
||||
|
||||
if built == 0:
|
||||
print(" ERROR: No kernels built")
|
||||
return 1
|
||||
|
||||
# Step 3: Run each kernel and validate
|
||||
print(f"\n--- Running Kernels (problem {args.size}x{args.size}x{args.size}) ---")
|
||||
np_dtype = np.float16 if args.dtype in ["fp16", "bf16"] else np.float32
|
||||
M = N = K = args.size
|
||||
|
||||
np.random.seed(42)
|
||||
A = (np.random.randn(M, K) * 0.1).astype(np_dtype)
|
||||
B = (np.random.randn(K, N) * 0.1).astype(np_dtype)
|
||||
C_ref = np.matmul(A.astype(np.float32), B.astype(np.float32)).astype(np_dtype)
|
||||
|
||||
print(
|
||||
f"\n {'#':<3} {'Name':<22} {'Tile':<14} {'Time(ms)':>10} {'TFLOPS':>10} {'MaxErr':>10} {'Status':<6}"
|
||||
)
|
||||
print(" " + "-" * 80)
|
||||
|
||||
results = []
|
||||
for i, (spec, setup) in enumerate(zip(specs, setups), 1):
|
||||
tile = f"{spec.tile_m}x{spec.tile_n}x{spec.tile_k}"
|
||||
|
||||
if not setup.success:
|
||||
print(
|
||||
f" {i:<3} {spec.name:<18} {tile:<14} {'N/A':>10} {'N/A':>10} {'N/A':>10} {'FAIL':<8}"
|
||||
f" {i:<3} {spec.name:<22} {tile:<14} {'---':>10} {'---':>10} {'---':>10} {'SKIP':<6}"
|
||||
)
|
||||
results.append((spec.name, False, 0, 0, 0))
|
||||
cleanup_gemm()
|
||||
results.append((spec.name, False, 0.0, 0.0, 0.0))
|
||||
continue
|
||||
|
||||
dispatcher = setup.dispatcher
|
||||
|
||||
# Check if size is supported
|
||||
if not dispatcher.is_supported(M, N, K):
|
||||
disp = setup.dispatcher
|
||||
if not disp.is_supported(M, N, K):
|
||||
print(
|
||||
f" {i:<3} {spec.name:<18} {tile:<14} {'N/A':>10} {'N/A':>10} {'N/A':>10} {'SKIP':<8}"
|
||||
f" {i:<3} {spec.name:<22} {tile:<14} {'---':>10} {'---':>10} {'---':>10} {'SKIP':<6}"
|
||||
)
|
||||
results.append((spec.name, False, 0, 0, 0))
|
||||
cleanup_gemm()
|
||||
results.append((spec.name, False, 0.0, 0.0, 0.0))
|
||||
continue
|
||||
|
||||
# Run GEMM
|
||||
result = dispatcher.run(A, B, M, N, K)
|
||||
|
||||
if not result.success:
|
||||
res = disp.run(A, B, M, N, K)
|
||||
if not res.success:
|
||||
print(
|
||||
f" {i:<3} {spec.name:<18} {tile:<14} {'N/A':>10} {'N/A':>10} {'N/A':>10} {'FAIL':<8}"
|
||||
f" {i:<3} {spec.name:<22} {tile:<14} {'---':>10} {'---':>10} {'---':>10} {'FAIL':<6}"
|
||||
)
|
||||
results.append((spec.name, False, 0, 0, 0))
|
||||
cleanup_gemm()
|
||||
results.append((spec.name, False, 0.0, 0.0, 0.0))
|
||||
continue
|
||||
|
||||
# Validate against NumPy reference
|
||||
C_ref = np.matmul(A.astype(np.float32), B.astype(np.float32)).astype(np_dtype)
|
||||
max_err = np.max(np.abs(result.output - C_ref))
|
||||
|
||||
# Check if within tolerance
|
||||
passed = max_err < 1e-2
|
||||
status = "PASS" if passed else "FAIL"
|
||||
|
||||
max_err = float(np.max(np.abs(res.output - C_ref)))
|
||||
ok = max_err < 1e-2
|
||||
tag = "PASS" if ok else "FAIL"
|
||||
print(
|
||||
f" {i:<3} {spec.name:<18} {tile:<14} {result.time_ms:>10.4f} {result.tflops:>10.2f} {max_err:>10.2e} {status:<8}"
|
||||
f" {i:<3} {spec.name:<22} {tile:<14} {res.time_ms:>10.4f} {res.tflops:>10.2f} {max_err:>10.2e} {tag:<6}"
|
||||
)
|
||||
results.append((spec.name, passed, result.time_ms, result.tflops, max_err))
|
||||
|
||||
cleanup_gemm()
|
||||
|
||||
# =========================================================================
|
||||
# Step 3: Summary
|
||||
# =========================================================================
|
||||
print("\n" + "=" * 70)
|
||||
print(" SUMMARY")
|
||||
print("=" * 70)
|
||||
results.append((spec.name, ok, res.time_ms, res.tflops, max_err))
|
||||
|
||||
# Step 4: Summary
|
||||
passed = sum(1 for r in results if r[1])
|
||||
failed = len(results) - passed
|
||||
valid = [r for r in results if r[1]]
|
||||
|
||||
print(f"\n Results: {passed}/{len(results)} kernels passed")
|
||||
print(f" Problem: {M}x{N}x{K}, dtype={args.dtype}")
|
||||
|
||||
if results:
|
||||
valid_results = [r for r in results if r[1]]
|
||||
if valid_results:
|
||||
best = max(valid_results, key=lambda x: x[3])
|
||||
print(f"\n Best kernel: {best[0]} ({best[3]:.2f} TFLOPS)")
|
||||
|
||||
if failed == 0:
|
||||
print("\n *** ALL KERNELS PASSED ***")
|
||||
else:
|
||||
print(f"\n *** {failed} KERNELS FAILED ***")
|
||||
|
||||
print("\n" + "=" * 70)
|
||||
print(f" Results: {passed}/{len(results)} passed")
|
||||
print(f" Problem: {M}x{N}x{K}, dtype={args.dtype}")
|
||||
print(f" JIT time: {jit_build_s:.1f} s (parallel)")
|
||||
if valid:
|
||||
best = max(valid, key=lambda x: x[3])
|
||||
print(f" Best: {best[0]} ({best[3]:.2f} TFLOPS)")
|
||||
print(f" Status: {'PASS' if failed == 0 else 'FAIL'}")
|
||||
print("=" * 70)
|
||||
|
||||
return 0 if failed == 0 else 1
|
||||
|
||||
Reference in New Issue
Block a user