mirror of
https://github.com/ROCm/composable_kernel.git
synced 2026-05-18 03:49:41 +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
|
||||
|
||||
@@ -6,9 +6,7 @@
|
||||
"""
|
||||
Example 02: Batch GEMM
|
||||
|
||||
Runs multiple GEMM operations with different sizes.
|
||||
|
||||
Complexity: ★★☆☆☆
|
||||
Runs multiple GEMM operations with different sizes using JIT compilation.
|
||||
|
||||
Usage:
|
||||
python3 02_batch_gemm.py
|
||||
@@ -25,9 +23,8 @@ import numpy as np
|
||||
|
||||
from ctypes_utils import (
|
||||
KernelConfig,
|
||||
setup_gemm_dispatcher,
|
||||
cleanup_gemm,
|
||||
reset_for_example,
|
||||
Registry,
|
||||
detect_gpu_arch,
|
||||
)
|
||||
|
||||
|
||||
@@ -55,20 +52,20 @@ Examples:
|
||||
help="Maximum problem size (default: 4096)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--arch", default="gfx942", help="Target architecture (default: gfx942)"
|
||||
"--arch",
|
||||
default=detect_gpu_arch(),
|
||||
help="Target architecture (auto-detected from rocminfo)",
|
||||
)
|
||||
args = parser.parse_args()
|
||||
|
||||
reset_for_example()
|
||||
|
||||
print("=" * 60)
|
||||
print("Example 02: Batch GEMM")
|
||||
print("=" * 60)
|
||||
|
||||
# =========================================================================
|
||||
# Step 1: Setup dispatcher
|
||||
# Step 1: JIT build dispatcher
|
||||
# =========================================================================
|
||||
print("\nStep 1: Setup Dispatcher")
|
||||
print("\nStep 1: JIT Build Dispatcher")
|
||||
|
||||
config = KernelConfig(
|
||||
dtype_a=args.dtype,
|
||||
@@ -80,19 +77,22 @@ Examples:
|
||||
gfx_arch=args.arch,
|
||||
)
|
||||
|
||||
setup = setup_gemm_dispatcher(config, registry_name="batch_gemm", verbose=True)
|
||||
if not setup.success:
|
||||
print(f" ERROR: {setup.error}")
|
||||
reg = Registry(name="batch_gemm")
|
||||
reg.register_kernel(config)
|
||||
|
||||
setups = reg.build(verbose=True)
|
||||
if not setups or not setups[0].success:
|
||||
error = setups[0].error if setups else "No kernels built"
|
||||
print(f" ERROR: {error}")
|
||||
return 1
|
||||
|
||||
dispatcher = setup.dispatcher
|
||||
dispatcher = setups[0].dispatcher
|
||||
|
||||
# =========================================================================
|
||||
# Step 2: Run batch of different sizes
|
||||
# =========================================================================
|
||||
print("\nStep 2: Run Batch")
|
||||
|
||||
# Generate sizes up to max_size
|
||||
all_sizes = [
|
||||
(256, 256, 256),
|
||||
(512, 512, 512),
|
||||
@@ -135,9 +135,6 @@ Examples:
|
||||
avg_tflops = (total_ops / 1e12) / (total_time / 1000)
|
||||
print(f"\n Total: {total_time:.2f} ms, Average: {avg_tflops:.2f} TFLOPS")
|
||||
|
||||
# Cleanup
|
||||
cleanup_gemm()
|
||||
|
||||
print("\n" + "=" * 60)
|
||||
print("Batch GEMM complete!")
|
||||
print("=" * 60)
|
||||
|
||||
@@ -6,9 +6,8 @@
|
||||
"""
|
||||
Example 03: Benchmark
|
||||
|
||||
Performance benchmarking with compute-optimized kernel configuration.
|
||||
|
||||
Complexity: ★★★☆☆
|
||||
Performance benchmarking with compute-optimized kernel configuration
|
||||
using JIT compilation.
|
||||
|
||||
Usage:
|
||||
python3 03_benchmark.py
|
||||
@@ -26,9 +25,8 @@ import numpy as np
|
||||
|
||||
from ctypes_utils import (
|
||||
KernelConfig,
|
||||
setup_gemm_dispatcher,
|
||||
cleanup_gemm,
|
||||
reset_for_example,
|
||||
Registry,
|
||||
detect_gpu_arch,
|
||||
)
|
||||
|
||||
|
||||
@@ -63,20 +61,20 @@ Examples:
|
||||
"--iterations", type=int, default=10, help="Benchmark iterations (default: 10)"
|
||||
)
|
||||
parser.add_argument(
|
||||
"--arch", default="gfx942", help="Target architecture (default: gfx942)"
|
||||
"--arch",
|
||||
default=detect_gpu_arch(),
|
||||
help="Target architecture (auto-detected from rocminfo)",
|
||||
)
|
||||
args = parser.parse_args()
|
||||
|
||||
reset_for_example()
|
||||
|
||||
print("=" * 60)
|
||||
print("Example 03: Benchmark")
|
||||
print("=" * 60)
|
||||
|
||||
# =========================================================================
|
||||
# Step 1: Setup dispatcher with compute-optimized config
|
||||
# Step 1: JIT build dispatcher with compute-optimized config
|
||||
# =========================================================================
|
||||
print("\nStep 1: Setup Dispatcher")
|
||||
print("\nStep 1: JIT Build Dispatcher")
|
||||
|
||||
config = KernelConfig(
|
||||
dtype_a=args.dtype,
|
||||
@@ -90,12 +88,16 @@ Examples:
|
||||
gfx_arch=args.arch,
|
||||
)
|
||||
|
||||
setup = setup_gemm_dispatcher(config, registry_name="benchmark", verbose=True)
|
||||
if not setup.success:
|
||||
print(f" ERROR: {setup.error}")
|
||||
reg = Registry(name="benchmark")
|
||||
reg.register_kernel(config)
|
||||
|
||||
setups = reg.build(verbose=True)
|
||||
if not setups or not setups[0].success:
|
||||
error = setups[0].error if setups else "No kernels built"
|
||||
print(f" ERROR: {error}")
|
||||
return 1
|
||||
|
||||
dispatcher = setup.dispatcher
|
||||
dispatcher = setups[0].dispatcher
|
||||
|
||||
# =========================================================================
|
||||
# Step 2: Benchmark
|
||||
@@ -130,11 +132,9 @@ Examples:
|
||||
A = np.random.randn(M, K).astype(np_dtype) * 0.1
|
||||
B = np.random.randn(K, N).astype(np_dtype) * 0.1
|
||||
|
||||
# Warmup
|
||||
for _ in range(args.warmup):
|
||||
dispatcher.run(A, B, M, N, K)
|
||||
|
||||
# Benchmark
|
||||
times = []
|
||||
for _ in range(args.iterations):
|
||||
result = dispatcher.run(A, B, M, N, K)
|
||||
@@ -150,9 +150,6 @@ Examples:
|
||||
f" {M:>4}x{N:>4}x{K:<4} | {min_time:>10.4f} | {avg_time:>10.4f} | {tflops:>10.2f}"
|
||||
)
|
||||
|
||||
# Cleanup
|
||||
cleanup_gemm()
|
||||
|
||||
# Summary
|
||||
print("\n" + "=" * 60)
|
||||
print("Summary")
|
||||
|
||||
@@ -6,9 +6,7 @@
|
||||
"""
|
||||
Example 04: Validation
|
||||
|
||||
Validates GPU GEMM against NumPy reference.
|
||||
|
||||
Complexity: ★★★☆☆
|
||||
Validates GPU GEMM against NumPy reference using JIT compilation.
|
||||
|
||||
Usage:
|
||||
python3 04_validation.py
|
||||
@@ -26,9 +24,8 @@ import numpy as np
|
||||
from ctypes_utils import (
|
||||
KernelConfig,
|
||||
Validator,
|
||||
setup_gemm_dispatcher,
|
||||
cleanup_gemm,
|
||||
reset_for_example,
|
||||
Registry,
|
||||
detect_gpu_arch,
|
||||
)
|
||||
|
||||
|
||||
@@ -56,20 +53,20 @@ Examples:
|
||||
"--atol", type=float, default=1e-2, help="Absolute tolerance (default: 1e-2)"
|
||||
)
|
||||
parser.add_argument(
|
||||
"--arch", default="gfx942", help="Target architecture (default: gfx942)"
|
||||
"--arch",
|
||||
default=detect_gpu_arch(),
|
||||
help="Target architecture (auto-detected from rocminfo)",
|
||||
)
|
||||
args = parser.parse_args()
|
||||
|
||||
reset_for_example()
|
||||
|
||||
print("=" * 60)
|
||||
print("Example 04: Validation")
|
||||
print("=" * 60)
|
||||
|
||||
# =========================================================================
|
||||
# Step 1: Setup dispatcher
|
||||
# Step 1: JIT build dispatcher
|
||||
# =========================================================================
|
||||
print("\nStep 1: Setup Dispatcher")
|
||||
print("\nStep 1: JIT Build Dispatcher")
|
||||
|
||||
config = KernelConfig(
|
||||
dtype_a=args.dtype,
|
||||
@@ -81,12 +78,16 @@ Examples:
|
||||
gfx_arch=args.arch,
|
||||
)
|
||||
|
||||
setup = setup_gemm_dispatcher(config, registry_name="validation", verbose=True)
|
||||
if not setup.success:
|
||||
print(f" ERROR: {setup.error}")
|
||||
reg = Registry(name="validation")
|
||||
reg.register_kernel(config)
|
||||
|
||||
setups = reg.build(verbose=True)
|
||||
if not setups or not setups[0].success:
|
||||
error = setups[0].error if setups else "No kernels built"
|
||||
print(f" ERROR: {error}")
|
||||
return 1
|
||||
|
||||
dispatcher = setup.dispatcher
|
||||
dispatcher = setups[0].dispatcher
|
||||
|
||||
# =========================================================================
|
||||
# Step 2: Run validation tests
|
||||
@@ -139,9 +140,6 @@ Examples:
|
||||
print(f" {name:<15} | {M}x{N}x{K:<5} | {max_err:>10.2e} | FAILED")
|
||||
failed += 1
|
||||
|
||||
# Cleanup
|
||||
cleanup_gemm()
|
||||
|
||||
# Summary
|
||||
print("\n" + "=" * 60)
|
||||
total = passed + failed
|
||||
|
||||
@@ -8,7 +8,6 @@ Example 05: NumPy Integration
|
||||
|
||||
Shows how to create a GPU-accelerated matmul wrapper.
|
||||
|
||||
Complexity: ★★☆☆☆
|
||||
|
||||
Usage:
|
||||
python3 05_numpy_integration.py
|
||||
@@ -29,6 +28,7 @@ from ctypes_utils import (
|
||||
setup_gemm_dispatcher,
|
||||
cleanup_gemm,
|
||||
reset_for_example,
|
||||
detect_gpu_arch,
|
||||
)
|
||||
|
||||
|
||||
@@ -70,7 +70,9 @@ Examples:
|
||||
help="Data type (default: fp16)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--arch", default="gfx942", help="Target architecture (default: gfx942)"
|
||||
"--arch",
|
||||
default=detect_gpu_arch(),
|
||||
help="Target architecture (auto-detected from rocminfo)",
|
||||
)
|
||||
args = parser.parse_args()
|
||||
|
||||
|
||||
@@ -8,7 +8,6 @@ Example 06: JSON Export
|
||||
|
||||
Exports registry configuration to JSON.
|
||||
|
||||
Complexity: ★★☆☆☆
|
||||
|
||||
Usage:
|
||||
python3 06_json_export.py
|
||||
@@ -28,6 +27,7 @@ from ctypes_utils import (
|
||||
setup_gemm_dispatcher,
|
||||
cleanup_gemm,
|
||||
reset_for_example,
|
||||
detect_gpu_arch,
|
||||
)
|
||||
|
||||
|
||||
@@ -54,7 +54,9 @@ Examples:
|
||||
help="Data type (default: fp16)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--arch", default="gfx942", help="Target architecture (default: gfx942)"
|
||||
"--arch",
|
||||
default=detect_gpu_arch(),
|
||||
help="Target architecture (auto-detected from rocminfo)",
|
||||
)
|
||||
args = parser.parse_args()
|
||||
|
||||
|
||||
@@ -18,7 +18,6 @@ This tests:
|
||||
- Multiple data types (fp16, bf16)
|
||||
- Different schedulers (intrawave, interwave)
|
||||
|
||||
Complexity: ★★★★☆
|
||||
|
||||
Usage:
|
||||
python3 07_stress_test.py
|
||||
@@ -43,6 +42,7 @@ from ctypes_utils import (
|
||||
cleanup_gemm,
|
||||
reset_for_example,
|
||||
Validator,
|
||||
detect_gpu_arch,
|
||||
)
|
||||
|
||||
|
||||
@@ -413,8 +413,8 @@ Examples:
|
||||
)
|
||||
parser.add_argument(
|
||||
"--arch",
|
||||
default="gfx942",
|
||||
help="Target architecture (default: gfx942)",
|
||||
default=detect_gpu_arch(),
|
||||
help="Target architecture (auto-detected from rocminfo, override with --arch gfxNNN)",
|
||||
)
|
||||
args = parser.parse_args()
|
||||
|
||||
|
||||
@@ -19,7 +19,6 @@ Heuristic strategies:
|
||||
- Memory-bound: Optimize memory access for bandwidth-limited cases
|
||||
- Latency-focused: Minimize kernel launch overhead for small problems
|
||||
|
||||
Complexity: ★★★★☆
|
||||
|
||||
Usage:
|
||||
python3 08_heuristics.py
|
||||
@@ -43,6 +42,7 @@ from ctypes_utils import (
|
||||
setup_gemm_dispatcher,
|
||||
cleanup_gemm,
|
||||
reset_for_example,
|
||||
detect_gpu_arch,
|
||||
)
|
||||
|
||||
|
||||
@@ -561,8 +561,8 @@ Examples:
|
||||
)
|
||||
parser.add_argument(
|
||||
"--arch",
|
||||
default="gfx942",
|
||||
help="Target architecture (default: gfx942)",
|
||||
default=detect_gpu_arch(),
|
||||
help="Target architecture (auto-detected from rocminfo, override with --arch gfxNNN)",
|
||||
)
|
||||
args = parser.parse_args()
|
||||
|
||||
|
||||
@@ -8,7 +8,6 @@ Example 09: Multiple Registries
|
||||
|
||||
Demonstrates multiple registries for different optimization targets.
|
||||
|
||||
Complexity: ★★★★★
|
||||
|
||||
Usage:
|
||||
python3 09_multi_registry.py
|
||||
@@ -30,6 +29,7 @@ from ctypes_utils import (
|
||||
setup_gemm_dispatcher,
|
||||
cleanup_gemm,
|
||||
reset_for_example,
|
||||
detect_gpu_arch,
|
||||
)
|
||||
|
||||
|
||||
@@ -50,7 +50,9 @@ Examples:
|
||||
help="Data type (default: fp16)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--arch", default="gfx942", help="Target architecture (default: gfx942)"
|
||||
"--arch",
|
||||
default=detect_gpu_arch(),
|
||||
help="Target architecture (auto-detected from rocminfo)",
|
||||
)
|
||||
args = parser.parse_args()
|
||||
|
||||
|
||||
@@ -33,6 +33,7 @@ from ctypes_utils import (
|
||||
setup_gemm_dispatcher,
|
||||
cleanup_gemm,
|
||||
reset_for_example,
|
||||
detect_gpu_arch,
|
||||
)
|
||||
|
||||
|
||||
@@ -69,7 +70,11 @@ def parse_args():
|
||||
# Kernel configuration
|
||||
parser.add_argument("--dtype", default="fp16", help="Data type")
|
||||
parser.add_argument("--pipeline", default="compv4", help="Pipeline type")
|
||||
parser.add_argument("--arch", default="gfx942", help="GPU architecture")
|
||||
parser.add_argument(
|
||||
"--arch",
|
||||
default=detect_gpu_arch(),
|
||||
help="GPU architecture (auto-detected from rocminfo)",
|
||||
)
|
||||
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
@@ -15,7 +15,6 @@ Key Features:
|
||||
- Use arch_filter validation on loaded configs
|
||||
- Export to C++ DECL_KERNEL_SET format
|
||||
|
||||
Complexity: ★★★☆☆
|
||||
|
||||
Usage:
|
||||
python3 11_json_import.py
|
||||
@@ -45,6 +44,7 @@ from ctypes_utils import ( # noqa: E402
|
||||
cleanup_gemm,
|
||||
reset_for_example,
|
||||
validate_kernel_config,
|
||||
detect_gpu_arch,
|
||||
)
|
||||
|
||||
# Sample JSON configuration (embedded for demonstration)
|
||||
@@ -141,8 +141,8 @@ Examples:
|
||||
)
|
||||
parser.add_argument(
|
||||
"--arch",
|
||||
default="gfx942",
|
||||
help="Target GPU architecture (default: gfx942)",
|
||||
default=detect_gpu_arch(),
|
||||
help="Target GPU architecture (auto-detected from rocminfo, override with --arch gfxNNN)",
|
||||
)
|
||||
args = parser.parse_args()
|
||||
|
||||
@@ -236,13 +236,13 @@ Examples:
|
||||
else:
|
||||
invalid_count += 1
|
||||
if invalid_count <= 3: # Show first 3 invalid
|
||||
print(f"\n ✗ Invalid: {config.kernel_name()}")
|
||||
print(f"\n FAIL Invalid: {config.kernel_name()}")
|
||||
for error in result.errors:
|
||||
print(f" Error: {error}")
|
||||
|
||||
print("\n Validation Summary:")
|
||||
print(f" ✓ Valid: {valid_count}")
|
||||
print(f" ✗ Invalid: {invalid_count}")
|
||||
print(f" OK Valid: {valid_count}")
|
||||
print(f" FAIL Invalid: {invalid_count}")
|
||||
print(f" Total: {len(configs)}")
|
||||
|
||||
# =========================================================================
|
||||
@@ -275,12 +275,12 @@ Examples:
|
||||
disp_config, registry_name="json_import", verbose=False
|
||||
)
|
||||
if setup.success:
|
||||
print(" ✓ Dispatcher setup successful")
|
||||
print(" OK Dispatcher setup successful")
|
||||
print(
|
||||
f" Kernel header: {setup.kernel_header.name if setup.kernel_header else 'N/A'}"
|
||||
)
|
||||
else:
|
||||
print(f" ⚠ Dispatcher setup: {setup.error}")
|
||||
print(f" WARNING Dispatcher setup: {setup.error}")
|
||||
print(" (This is expected if kernels aren't generated)")
|
||||
|
||||
# =========================================================================
|
||||
|
||||
@@ -295,5 +295,5 @@ Compilation time scales roughly linearly with kernel count.
|
||||
## Related Documentation
|
||||
|
||||
- [C++ GEMM Examples](../cpp/README.md)
|
||||
- [Python Conv Examples](../../conv/python/README.md)
|
||||
- [Python Utilities](../../../python/README.md)
|
||||
- [Main Dispatcher README](../../../README.md)
|
||||
|
||||
Reference in New Issue
Block a user