diff --git a/dispatcher/bindings/README.md b/dispatcher/bindings/README.md index e460b38b5b..407ac7b697 100644 --- a/dispatcher/bindings/README.md +++ b/dispatcher/bindings/README.md @@ -8,6 +8,7 @@ This directory contains language bindings for the CK Tile Dispatcher. bindings/ |---- ctypes/ # Python ctypes bindings (C API) | |---- gemm_ctypes_lib.cpp # GEMM dispatcher C API +| |---- grouped_gemm_ctypes_lib.cpp # Grouped (multi-problem) GEMM bridge C API -- see GROUPED_GEMM_BRIDGE.md | |---- conv_ctypes_lib.cpp # Grouped conv dispatcher C API (fwd + bwd_data) | |---- conv_bwdw_ctypes_lib.cpp # Grouped conv backward weight C API (separate library) | |---- fmha_ctypes_lib.cpp # FMHA dispatcher C API (fwd + bwd) diff --git a/dispatcher/bindings/ctypes/GROUPED_GEMM_BRIDGE.md b/dispatcher/bindings/ctypes/GROUPED_GEMM_BRIDGE.md new file mode 100644 index 0000000000..f6a7780b83 --- /dev/null +++ b/dispatcher/bindings/ctypes/GROUPED_GEMM_BRIDGE.md @@ -0,0 +1,121 @@ + + +# Grouped GEMM: Tile Engine -> Dispatcher Bridge + +This document describes the **grouped_gemm** variant of the Tile Engine (TE) -> +Dispatcher bridge (PR #8130). It is the grouped counterpart of the regular-GEMM +bridge (#8123/#8479), the fp8/bf8/int8 bridge (#8887), and the Stream-K bridge +(#8136). + +## What the bridge is + +In the bridge model the **Dispatcher is the single source of truth** for +codegen, build, and runtime; **Tile Engine only generates configs and +benchmarks** them. TE no longer carries its own kernel-instance build path — +it shells out to the dispatcher codegen and runs the resulting kernel. + +For most variants the dispatcher runs a kernel through its registry/backend. +Grouped GEMM cannot use that path (see below), so the grouped bridge takes the +same approach as Stream-K: a dedicated ctypes library that **bypasses the +registry** and calls the generated `SelectedKernel::launch(...)` directly. + +## Why grouped needs special handling + +Grouped GEMM is **multi-problem**: a single launch runs a *list* of `(M, N, K)` +sub-problems, each with its own A/B/C device pointers. Two consequences: + +1. The single-problem run path (`g_dispatcher->run` / `GemmHostArgs`) cannot + express a list of problems. +2. The generated registry wrapper (`generated_tile_backend.hpp::run()`) + hard-codes the single-problem `SelectedKernel::launch(GemmHostArgs, ...)` + signature and will not compile against a grouped `SelectedKernel`. + +So the grouped kernel header exposes a different launch signature + +```cpp +static float launch(const std::vector>& descs, + const stream_config& stream); +``` + +and the grouped ctypes lib force-includes one generated kernel header +(`-include ..._grouped.hpp` with `CK_TILE_SINGLE_KERNEL_INCLUDE`), calls that +`launch` directly, and reports the kernel name from the compile-time +`KERNEL_NAME` macro. + +## Components + +| Layer | File | Role | +|---|---|---| +| Codegen | `dispatcher/codegen/unified_gemm_codegen.py` | `GemmVariant.GROUPED`; `_launch_function_grouped` (DeviceMem internal workspace, `MakeKargs`, persistent/non-persistent grid). Kept in lockstep with PR #8075. | +| Codegen | `dispatcher/codegen/arch_filter.py` | `GEMM_GROUPED` operator tile constraints. | +| C API | `dispatcher/bindings/ctypes/grouped_gemm_ctypes_lib.cpp` | Multi-problem ABI; per-group device alloc/copy; layout-derived strides; warmup/repeat timing. | +| Python | `dispatcher/python/gemm_utils.py` | `GroupedGemmProblem` / `GroupedGemmResult`, `GpuGroupedGemmRunner`, `run_grouped`, `build_grouped`, dtype/layout codecs. | +| Python | `dispatcher/python/ctypes_utils.py` | Threads the `grouped` variant into the codegen `--variants` flag. | +| TE driver | `tile_engine/ops/gemm/grouped_gemm_full_benchmark.py` | Generates configs, builds `.so`s in parallel, benchmarks in disposable workers. | +| TE worker | `tile_engine/ops/gemm/run_one_grouped_gemm_kernel.py` | Runs one grouped kernel; dtype/layout-aware operand generation. | + +## C ABI + +```c +int dispatcher_init(void); // lightweight no-op (no registry) +int dispatcher_run_grouped_gemm( + int group_count, + const int64_t* Ms, // [group_count] + const int64_t* Ns, // [group_count] + const int64_t* Ks, // [group_count] + const void** A_ptrs, // host A buffers, one per group + const void** B_ptrs, // host B buffers, one per group + void** C_ptrs, // host C out buffers, one per group + float* time_ms); // out: average kernel time +// returns 0 ok, -1 HIP/throw, -2 arguments unsupported by the kernel +``` + +The lib `hipMalloc`s A/B/C per group, copies A and B host->device, memsets C, +builds `std::vector>` with **strides derived from +the compile-time `ALayout`/`BLayout`/`CLayout`** of the `-include`d header +(`std::is_same_v<…, RowMajor>`), launches once, then copies each C back. The ABI +is `void*` + element-size, so it is dtype-agnostic; the Python runner owns the +numpy codecs. + +## Coverage + +The bridge runnable set is exactly the Old-TE grouped_gemm runnable set on +`develop` — no more, no less: + +| Layout \ Dtype | fp16 | bf16 | fp8 (E4M3) | bf8 (E5M2) | +|---|---|---|---|---| +| rcr | ✓ | ✓ | ✓ | ✓ | +| rrr | ✓ | ✓ | ✓ | ✓ | +| ccr | ✓ | ✓ | ✓ | ✓ | +| crr | ✓ | ✓ | ✓ | ✓ | + +- **Matrix C is always row-major** (grouped builder constraint), so the layout + string varies A/B only. +- **Excluded:** `int8` (rejected by the TE grouped builder), `fp32`/`fp64` + (no MFMA warp tiles). These are excluded on both sides. +- fp8/bf8 use the **FNUZ** encoding on gfx942 (matches the regular #8887 path); + the Python codecs require `ml_dtypes`. + +## Building and running + +Generate + build one grouped `.so` and run the A/B parity sweep vs Old-TE: + +```bash +# Codegen smoke (no GPU): one variant/dtype/layout +python3 dispatcher/codegen/unified_gemm_codegen.py \ + --output-dir /tmp/grp --datatype bf16 --layout ccr \ + --variants grouped --config dispatcher/codegen/default_config.json + +# Full TE-driven parity sweep (build + benchmark) +python3 tile_engine/ops/gemm/grouped_gemm_full_benchmark.py \ + --arch gfx942 --dtype fp16 --layout rcr --csv grouped_results.csv +``` + +Timing knobs `CK_TILE_BENCH_WARMUP` (default 50) and `CK_TILE_BENCH_REPEAT` +(default 100) are honored by **both** the grouped ctypes lib and the registry +backend, so bridge-vs-Old-TE A/B comparisons stay matched. For fair parity keep +`flush_cache=false`, `rotating_count=1`, run on a single GPU, and re-measure any +`|gap|>15%` outlier standalone. diff --git a/dispatcher/bindings/ctypes/grouped_gemm_ctypes_lib.cpp b/dispatcher/bindings/ctypes/grouped_gemm_ctypes_lib.cpp new file mode 100644 index 0000000000..af1dc854af --- /dev/null +++ b/dispatcher/bindings/ctypes/grouped_gemm_ctypes_lib.cpp @@ -0,0 +1,295 @@ +// Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +// SPDX-License-Identifier: MIT + +/** + * Grouped GEMM Dispatcher ctypes Library + * + * Provides C API for Python ctypes integration for the GROUPED GEMM variant. + * Kernel header included via -include at compile time. + * + * The grouped kernel has a genuinely different ABI from regular GEMM: it takes a + * LIST of (M,N,K) sub-problems plus arrays of A/B/C device pointers, and its + * generated launch() builds the per-group arg workspace internally: + * + * static float launch(const std::vector>& descs, + * const stream_config& stream); + * + * The single-problem dispatcher run path (g_dispatcher->run / GemmHostArgs) cannot + * express this, and the generated_tile_backend wrapper hard-codes the single-problem + * launch signature, so this lib calls SelectedKernel::launch(descs, stream) directly + * and reports the kernel name from the compile-time KERNEL_NAME macro instead of the + * registry. + * + * Usage from Python: + * lib = ctypes.CDLL("libdispatcher_grouped_gemm.so") + * lib.dispatcher_init() + * lib.dispatcher_run_grouped_gemm(...) + */ + +#include +#include +#include +#include +#include +#include +#include +#include +#include + +// Kernel header included via -include compiler flag (with CK_TILE_SINGLE_KERNEL_INCLUDE). +// Defines: ADataType, BDataType, CDataType, AccDataType, SelectedKernel, KERNEL_NAME +// and transitively brings in ck_tile::GroupedGemmHostArgs and ck_tile::stream_config. + +// GPU architecture - can be overridden via -DGFX_ARCH="gfx90a" at compile time +#ifndef GFX_ARCH +#define GFX_ARCH "gfx942" +#endif + +static bool g_initialized = false; + +// Read an integer benchmark knob from the environment, falling back to +// `fallback` when unset or unparseable. Mirrors generated_tile_backend.hpp so +// both bridge sides honor the same CK_TILE_BENCH_* env vars. +static int env_int(const char* name, int fallback) +{ + const char* v = std::getenv(name); + if(v == nullptr || *v == '\0') + return fallback; + char* end = nullptr; + const long out = std::strtol(v, &end, 10); + if(end == v) + return fallback; + return static_cast(out); +} + +extern "C" { + +/** + * Initialize the grouped GEMM library. + * + * The grouped path does not use the dispatcher/registry (it launches the + * force-included kernel directly), so this is a lightweight no-op kept for ABI + * parity with the regular GEMM lib. Returns 0 on success. + */ +int dispatcher_initialize() +{ + g_initialized = true; + return 0; +} + +/** + * Initialize dispatcher (alias) + */ +int dispatcher_init() { return dispatcher_initialize(); } + +/** + * Run grouped GEMM on GPU by launching the force-included kernel directly. + * + * For each group: hipMalloc A/B/C, copy A and B host->device, memset C, then build + * a std::vector> with strides derived from the + * compile-time ALayout/BLayout/CLayout of the -include'd kernel header (k_batch=1) + * and launch. After the launch the per-group C buffers are copied back to the + * caller's host buffers. + * + * Layout contract: A is MxK, B is KxN, C is MxN; leading dimensions follow each + * operand's row/col-major layout (CLayout is always RowMajor for grouped). + * + * Returns: 0 on success, -1 on HIP error / generic throw, -2 if the kernel reports + * the arguments are unsupported. + */ +int dispatcher_run_grouped_gemm(int group_count, + const int64_t* Ms, + const int64_t* Ns, + const int64_t* Ks, + const void** A_ptrs, + const void** B_ptrs, + void** C_ptrs, + float* time_ms) +{ + if(!g_initialized || group_count <= 0 || !Ms || !Ns || !Ks || !A_ptrs || !B_ptrs || !C_ptrs) + { + return -1; + } + + std::vector A_dev(group_count, nullptr); + std::vector B_dev(group_count, nullptr); + std::vector C_dev(group_count, nullptr); + + auto cleanup_gpu_mem = [&]() { + for(int g = 0; g < group_count; ++g) + { + if(A_dev[g]) + (void)hipFree(A_dev[g]); + if(B_dev[g]) + (void)hipFree(B_dev[g]); + if(C_dev[g]) + (void)hipFree(C_dev[g]); + } + }; + + std::vector> descs; + descs.reserve(group_count); + + for(int g = 0; g < group_count; ++g) + { + const int64_t M = Ms[g]; + const int64_t N = Ns[g]; + const int64_t K = Ks[g]; + + if(M <= 0 || N <= 0 || K <= 0 || !A_ptrs[g] || !B_ptrs[g] || !C_ptrs[g]) + { + cleanup_gpu_mem(); + return -1; + } + + if(hipMalloc(&A_dev[g], M * K * sizeof(ADataType)) != hipSuccess) + { + cleanup_gpu_mem(); + return -1; + } + if(hipMalloc(&B_dev[g], K * N * sizeof(BDataType)) != hipSuccess) + { + cleanup_gpu_mem(); + return -1; + } + if(hipMalloc(&C_dev[g], M * N * sizeof(CDataType)) != hipSuccess) + { + cleanup_gpu_mem(); + return -1; + } + + if(hipMemcpy(A_dev[g], A_ptrs[g], M * K * sizeof(ADataType), hipMemcpyHostToDevice) != + hipSuccess) + { + cleanup_gpu_mem(); + return -1; + } + if(hipMemcpy(B_dev[g], B_ptrs[g], K * N * sizeof(BDataType), hipMemcpyHostToDevice) != + hipSuccess) + { + cleanup_gpu_mem(); + return -1; + } + if(hipMemset(C_dev[g], 0, M * N * sizeof(CDataType)) != hipSuccess) + { + cleanup_gpu_mem(); + return -1; + } + + // Derive leading dimensions from the compile-time layouts the kernel was + // generated with (ALayout/BLayout/CLayout from the -include'd header), + // matching Old-TE gemm_validation_utils.get_abc_layouts: + // stride_A = ALayout row-major ? K : M + // stride_B = BLayout row-major ? N : K + // stride_E = CLayout row-major ? N : M (CLayout is always RowMajor for grouped) + using RowMajor = ck_tile::tensor_layout::gemm::RowMajor; + const auto stride_A = std::is_same_v ? static_cast(K) + : static_cast(M); + const auto stride_B = std::is_same_v ? static_cast(N) + : static_cast(K); + const auto stride_E = std::is_same_v ? static_cast(N) + : static_cast(M); + // k_batch=1 for numeric parity. + descs.emplace_back(static_cast(A_dev[g]), + static_cast(B_dev[g]), + std::array{}, + static_cast(C_dev[g]), + /*k_batch=*/1, + static_cast(M), + static_cast(N), + static_cast(K), + stride_A, + stride_B, + std::array{}, + stride_E); + } + + ck_tile::stream_config stream_cfg; + stream_cfg.stream_id_ = nullptr; + stream_cfg.time_kernel_ = true; + stream_cfg.log_level_ = 0; + stream_cfg.cold_niters_ = env_int("CK_TILE_BENCH_WARMUP", 50); + stream_cfg.nrepeat_ = env_int("CK_TILE_BENCH_REPEAT", 100); + stream_cfg.is_gpu_timer_ = true; + stream_cfg.flush_cache_ = false; + stream_cfg.rotating_count_ = 1; + + float exec_time = 0.0f; + try + { + exec_time = SelectedKernel::launch(descs, stream_cfg); + } + catch(const std::exception& e) + { + cleanup_gpu_mem(); + if(std::string(e.what()).find("not supported") != std::string::npos) + { + if(time_ms) + { + *time_ms = -1.0f; + } + return -2; // Arguments not supported by this kernel + } + return -1; + } + + // Copy each group's result back to host. + for(int g = 0; g < group_count; ++g) + { + const int64_t M = Ms[g]; + const int64_t N = Ns[g]; + if(hipMemcpy(C_ptrs[g], C_dev[g], M * N * sizeof(CDataType), hipMemcpyDeviceToHost) != + hipSuccess) + { + cleanup_gpu_mem(); + return -1; + } + } + + if(time_ms) + { + *time_ms = exec_time; + } + + cleanup_gpu_mem(); + return 0; +} + +/** + * Get kernel information (legacy single-kernel ABI). + * + * Returns the compile-time KERNEL_NAME of the force-included kernel header. + */ +const char* dispatcher_get_kernel_name() { return KERNEL_NAME; } + +/** + * Get the name of the kernel at a given registry index (multi-kernel ABI). + * + * Each grouped .so force-includes exactly one kernel header, so index 0 reports + * KERNEL_NAME and any other index is out of range. Mirrors the regular GEMM lib's + * name ABI so the Python bridge can use the same name-lookup path. + * Returns 0 on success, -1 on bad args or out-of-range index. + */ +int dispatcher_get_kernel_name_at(int index, char* buffer, int buffer_size) +{ + if(!buffer || buffer_size <= 0 || index != 0) + { + return -1; + } + + std::strncpy(buffer, KERNEL_NAME, static_cast(buffer_size) - 1); + buffer[buffer_size - 1] = '\0'; + return 0; +} + +/** + * Get the number of kernels in this .so (always 1 for the grouped single-include lib). + */ +int dispatcher_get_kernel_count() { return 1; } + +/** + * Cleanup library resources (no-op; kept for ABI parity). + */ +void dispatcher_cleanup() { g_initialized = false; } + +} // extern "C" diff --git a/dispatcher/codegen/README.md b/dispatcher/codegen/README.md index 40a9b7b8c1..965abb1213 100644 --- a/dispatcher/codegen/README.md +++ b/dispatcher/codegen/README.md @@ -77,7 +77,7 @@ results = codegen.generate_all() | `--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 | +| `--variants` | `standard`, `preshuffle`, `multi_d`, `grouped` | Kernel variants | | `--preselected` | `fp16_rcr_essential`, etc. | Predefined kernel set | ### Layout Notation @@ -98,6 +98,28 @@ Element-wise fusion: `C = op(A x B + D0 + D1 + ...)` Supported ops: `PassThrough`, `MultiDAdd`, `Relu`, `Gelu`, `Sigmoid`, `Tanh` +### Grouped +Batched GEMM over a list of independently-shaped groups in a single launch +(`ck_tile::GroupedGemmKernel`). Brings the dispatcher to parity with the Tile Engine +`grouped_gemm` op. The per-group argument vector is built with `MakeKargs`, copied to an +internally-allocated `DeviceMem` workspace, and the device pointer + group count are passed +to the kernel (the dispatcher workspace idiom — no external `kargs_ptr`). + +- Datatypes: `fp16`, `bf16`, `fp8`, `bf8` (matches the Tile Engine grouped runnable set; + `fp8`/`bf8` accumulate in `fp32` and emit an `fp16` C output). +- Layouts: `rcr`, `rrr`, `ccr`, `crr` (C is always row-major). + +```bash +python3 unified_gemm_codegen.py \ + --datatype fp16 \ + --layout rcr \ + --variants grouped \ + --gpu-target gfx942 \ + --output-dir generated_kernels +``` + +Build and run end-to-end with [`examples/gemm/cpp/02_grouped_gemm_driver.cpp`](../examples/gemm/cpp/README.md). + ## Output Structure ``` @@ -105,6 +127,7 @@ generated_kernels/ |---- gemm_fp16_rcr_compv4_..._128x128x32_....hpp # GEMM kernels |---- gemm_fp16_rcr_compv4_..._preshuffle.hpp |---- gemm_fp16_rcr_compv4_..._multid_Relu_d1.hpp +|---- gemm_fp16_rcr_compv3_..._128x128x64_..._grouped.hpp # Grouped GEMM kernels |---- grouped_conv_fwd_fp16_nhwgc_..._128x128x32_....hpp # Grouped conv kernels +---- ... ``` diff --git a/dispatcher/codegen/arch_filter.py b/dispatcher/codegen/arch_filter.py index 6a2dff0b39..a6299b0b3f 100644 --- a/dispatcher/codegen/arch_filter.py +++ b/dispatcher/codegen/arch_filter.py @@ -50,6 +50,7 @@ class OperatorType(Enum): GEMM = "gemm" GEMM_PRESHUFFLE = "gemm_preshuffle" GEMM_MULTI_D = "gemm_multi_d" + GEMM_GROUPED = "gemm_grouped" GEMM_STREAMK = "gemm_streamk" CONV_FWD = "conv_fwd" CONV_BWD_DATA = "conv_bwd_data" @@ -86,6 +87,7 @@ OPERATOR_TILE_CONSTRAINTS = { "tile_n_alignment": 16, "tile_k_alignment": 8, }, + OperatorType.GEMM_GROUPED: { # NOTE: these are copied from plain GEMM and only gate tile *shape* validity. # They do NOT express Stream-K's real feasibility requirement -- that a problem # has enough output tiles to partition K-work across the CUs. That gate is diff --git a/dispatcher/codegen/unified_gemm_codegen.py b/dispatcher/codegen/unified_gemm_codegen.py index 3b26e5a112..05e773442e 100755 --- a/dispatcher/codegen/unified_gemm_codegen.py +++ b/dispatcher/codegen/unified_gemm_codegen.py @@ -202,6 +202,7 @@ class GemmVariant(Enum): STANDARD = "standard" PRESHUFFLE = "preshuffle" MULTI_D = "multi_d" + GROUPED = "grouped" # Stream-K. COVERAGE LIMITATION: the dispatcher does NOT yet emit the full # Old-TE Stream-K tile surface. The kernels generated here are driven by the # tile list passed to this codegen, which is narrower than tile_engine's: @@ -367,6 +368,8 @@ class KernelNaming: name += "_preshuffle" elif config.variant == GemmVariant.MULTI_D: name += f"_multid_{config.elementwise_op}_d{config.num_d_tensors}" + elif config.variant == GemmVariant.GROUPED: + name += "_grouped" elif config.variant == GemmVariant.STREAM_K: name += "_streamk" # Atomic keeps the bare "_streamk" suffix for name parity with the @@ -420,6 +423,15 @@ class CKTileKernelGenerator: includes += """ #include "ck_tile/ops/elementwise/unary_element_wise_operation.hpp" #include "ck_tile/ops/gemm/kernel/gemm_multi_d_kernel.hpp" +""" + + if config.variant == GemmVariant.GROUPED: + includes += """ +#include +#include +#include "ck_tile/host/device_memory.hpp" +#include "ck_tile/host/hip_check_error.hpp" +#include "ck_tile/ops/gemm/kernel/grouped_gemm_kernel.hpp" """ if config.preshuffle: @@ -604,6 +616,9 @@ using CLayout = {ns_name}::CLayout; #define GEMM_KEY_TRANSPOSE_C 0 #define GEMM_KEY_GROUPED 0 #define GEMM_KEY_SPLIT_K 1 +using ALayout = {ns_name}::ALayout; +using BLayout = {ns_name}::BLayout; +using CLayout = {ns_name}::CLayout; #endif // CK_TILE_SINGLE_KERNEL_INCLUDE """ @@ -629,6 +644,8 @@ using CLayout = {ns_name}::CLayout; """Generate launch function""" if config.variant == GemmVariant.MULTI_D: return self._launch_function_multi_d(config) + if config.variant == GemmVariant.GROUPED: + return self._launch_function_grouped(config) if config.variant == GemmVariant.STREAM_K: return self._launch_function_streamk(config) if config.preshuffle: @@ -683,6 +700,69 @@ using CLayout = {ns_name}::CLayout; return ave_time; }}""" + def _launch_function_grouped(self, config: KernelConfig) -> str: + """Generate launch function for grouped GEMM. + + Follows the dispatcher's workspace idiom (see grouped_conv stream-K launch in + unified_grouped_conv_codegen.py): signature is (args, stream); the device + workspace is allocated internally via DeviceMem rather than passed in. The + grouped kernel's per-group arg vector is built with MakeKargs, copied to the + workspace, and the device pointer + group count are passed to the kernel. + """ + persistent = config.trait.persistent + grid_expr = ( + "GemmKernel::MaxOccupancyGridSize(stream)" + if persistent + else "dim3(kargs.empty() ? 0 : kargs.back().block_end, 1, 1)" + ) + return f""" + static float launch(const std::vector>& gemm_descs, + const stream_config& stream) {{ + if(gemm_descs.empty()) return 0.0f; + + float ave_time{{0}}; + + constexpr auto scheduler = {self.tm.SCHEDULER_TO_CK[config.trait.scheduler]}; + + using UniversalGemmProblem = UniversalGemmPipelineProblem< + ADataType, BDataType, AccDataType, TileShape, + TileGemmUniversalTraits, + scheduler>; + + using GemmPipeline = {self.tm.PIPELINE_TO_CK[config.trait.pipeline]}; + {self._epilogue_code(config)} + + using GemmKernel = ck_tile::GroupedGemmKernel; + + auto kargs = GemmKernel::MakeKargs(gemm_descs); + if(!GemmKernel::IsSupportedArgument(kargs)) {{ + throw std::runtime_error("Arguments not supported for grouped gemm kernel"); + }} + + // Workspace allocated internally (dispatcher idiom, mirrors grouped_conv stream-K). + const std::size_t ws_size = kargs.size() * sizeof(ck_tile::GemmTransKernelArg<>); + ck_tile::DeviceMem workspace_dev(ws_size); + HIP_CHECK_ERROR(hipMemcpyWithStream(workspace_dev.GetDeviceBuffer(), + kargs.data(), + ws_size, + hipMemcpyHostToDevice, + stream.stream_id_)); + + const dim3 grids = {grid_expr}; + const dim3 blocks = GemmKernel::BlockSize(); + + constexpr int kBlockPerCu = {config.k_block_per_cu}; + ave_time = launch_kernel(stream, + make_kernel(GemmKernel{{}}, grids, blocks, 0, + cast_pointer_to_constant_address_space(workspace_dev.GetDeviceBuffer()), + kargs.size())); + + return ave_time; + }}""" + def _launch_function_preshuffle(self, config: KernelConfig) -> str: """Generate launch function for preshuffle GEMM (weight preshuffle variant) @@ -985,7 +1065,7 @@ using CLayout = {ns_name}::CLayout; tuple<>, CLayout, element_wise::PassThrough, TilePartitioner::MPerBlock, TilePartitioner::NPerBlock, WarpPerBlock_M, WarpPerBlock_N, WarpTileM, WarpTileN, WarpTileK, - TransposeC, NumWaveGroups>; + TransposeC, NumWaveGroups, false, 1, 1, DoubleSmemBuffer>; using GemmEpilogue = CShuffleEpilogue;""" else: return """ @@ -1327,7 +1407,7 @@ class UnifiedGemmCodegen: """Get all configurations for a variant Args: - variant: GEMM variant (STANDARD, PRESHUFFLE, MULTI_D) + variant: GEMM variant (STANDARD, PRESHUFFLE, MULTI_D, GROUPED) Returns: List of valid kernel configurations for the variant @@ -1428,6 +1508,11 @@ class UnifiedGemmCodegen: ) ) + elif variant == GemmVariant.GROUPED: + # Grouped GEMM uses the same tile/trait configs as STANDARD — + # the only difference is the kernel type (GroupedGemmKernel vs GemmKernel) + configs.append(KernelConfig(tile=tile, trait=trait, variant=variant)) + return configs def _get_tile_configs(self) -> List[TileConfig]: @@ -1534,6 +1619,7 @@ class UnifiedGemmCodegen: GemmVariant.STANDARD: OperatorType.GEMM, GemmVariant.PRESHUFFLE: OperatorType.GEMM_PRESHUFFLE, GemmVariant.MULTI_D: OperatorType.GEMM_MULTI_D, + GemmVariant.GROUPED: OperatorType.GEMM_GROUPED, GemmVariant.STREAM_K: OperatorType.GEMM_STREAMK, } operator = variant_to_operator.get(variant, OperatorType.GEMM) @@ -1784,7 +1870,7 @@ def main(): parser.add_argument( "--variants", nargs="+", - choices=["standard", "preshuffle", "multi_d", "stream_k"], + choices=["standard", "preshuffle", "multi_d", "stream_k" ,"grouped"], default=["standard"], help="Variants to generate", ) diff --git a/dispatcher/examples/gemm/cpp/02_grouped_gemm_driver.cpp b/dispatcher/examples/gemm/cpp/02_grouped_gemm_driver.cpp new file mode 100644 index 0000000000..6f8b20216e --- /dev/null +++ b/dispatcher/examples/gemm/cpp/02_grouped_gemm_driver.cpp @@ -0,0 +1,192 @@ +// Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +// SPDX-License-Identifier: MIT + +/** + * Minimal standalone grouped-GEMM driver (dispatcher way). + * + * Grouped GEMM cannot ride the standard dispatcher.run(A,B,C,problem) path: + * that backend hardcodes a single GemmHostArgs. Instead, this driver includes a + * single generated grouped kernel header (CK_TILE_SINGLE_KERNEL_INCLUDE) and + * calls SelectedKernel::launch(descs, stream) directly with a vector of + * descriptors -- the same 2-arg signature the dispatcher generates (workspace is + * allocated INSIDE launch()). It builds per-group tensors, runs, and verifies + * each group against ck_tile::reference_gemm. + * + * Build (single-kernel include style): + * hipcc -std=c++17 --offload-arch=gfx942 \ + * -DCK_TILE_SINGLE_KERNEL_INCLUDE \ + * -I /include -I \ + * -include /_grouped.hpp \ + * 02_grouped_gemm_driver.cpp -o grouped_gemm_driver + */ + +#include + +#include +#include +#include +#include +#include +#include +#include + +#include "ck_tile/core.hpp" +#include "ck_tile/host.hpp" +#include "ck_tile/ops/gemm.hpp" + +// The generated grouped kernel header is injected on the command line with +// -include and -DCK_TILE_SINGLE_KERNEL_INCLUDE. It exports into the global +// namespace: SelectedKernel, ADataType, BDataType, CDataType, AccDataType, +// ALayout, BLayout, CLayout, and KERNEL_NAME. + +template +static constexpr inline auto is_row_major(Layout) +{ + return ck_tile::bool_constant< + std::is_same_v, ck_tile::tensor_layout::gemm::RowMajor>>{}; +} + +static std::vector parse_csv_ints(const std::string& s) +{ + std::vector out; + std::string cur; + for(char c : s) + { + if(c == ',') + { + if(!cur.empty()) + { + out.push_back(std::stoi(cur)); + cur.clear(); + } + } + else + cur.push_back(c); + } + if(!cur.empty()) + out.push_back(std::stoi(cur)); + return out; +} + +static std::string get_opt(int argc, char** argv, const std::string& key, const std::string& def) +{ + for(int i = 1; i < argc - 1; ++i) + if(key == argv[i]) + return argv[i + 1]; + return def; +} + +int main(int argc, char** argv) +{ + const int group_count = std::stoi(get_opt(argc, argv, "--groups", "8")); + const int kbatch = std::stoi(get_opt(argc, argv, "--kbatch", "1")); + const int warmup = std::stoi(get_opt(argc, argv, "--warmup", "10")); + const int repeat = std::stoi(get_opt(argc, argv, "--repeat", "50")); + const bool validate = get_opt(argc, argv, "--validate", "1") != "0"; + + std::vector Ms = parse_csv_ints(get_opt(argc, argv, "--Ms", "")); + std::vector Ns = parse_csv_ints(get_opt(argc, argv, "--Ns", "")); + std::vector Ks = parse_csv_ints(get_opt(argc, argv, "--Ks", "")); + + const int dm = std::stoi(get_opt(argc, argv, "--m", "256")); + const int dn = std::stoi(get_opt(argc, argv, "--n", "256")); + const int dk = std::stoi(get_opt(argc, argv, "--k", "512")); + + auto sz = static_cast(group_count); + if(Ms.size() != sz || Ns.size() != sz || Ks.size() != sz) + { + Ms.assign(group_count, dm); + Ns.assign(group_count, dn); + Ks.assign(group_count, dk); + } + + std::cout << "Kernel: " << KERNEL_NAME << "\n"; + std::cout << "groups=" << group_count << " kbatch=" << kbatch << "\n"; + + std::vector> a_host, b_host; + std::vector> c_host; + std::vector> a_dev, b_dev, c_dev; + std::vector sA(group_count), sB(group_count), sC(group_count); + + std::vector> descs; + descs.reserve(group_count); + + for(int i = 0; i < group_count; ++i) + { + const ck_tile::index_t M = Ms[i], N = Ns[i], K = Ks[i]; + sA[i] = ck_tile::get_default_stride(M, K, 0, is_row_major(ALayout{})); + sB[i] = ck_tile::get_default_stride(K, N, 0, is_row_major(BLayout{})); + sC[i] = ck_tile::get_default_stride(M, N, 0, is_row_major(CLayout{})); + + a_host.push_back(ck_tile::HostTensor( + ck_tile::host_tensor_descriptor(M, K, sA[i], is_row_major(ALayout{})))); + b_host.push_back(ck_tile::HostTensor( + ck_tile::host_tensor_descriptor(K, N, sB[i], is_row_major(BLayout{})))); + c_host.push_back(ck_tile::HostTensor( + ck_tile::host_tensor_descriptor(M, N, sC[i], is_row_major(CLayout{})))); + + ck_tile::FillUniformDistribution{-1.f, 1.f}(a_host[i]); + ck_tile::FillUniformDistribution{-1.f, 1.f}(b_host[i]); + c_host[i].SetZero(); + + a_dev.push_back(std::make_unique(a_host[i])); + b_dev.push_back(std::make_unique(b_host[i])); + c_dev.push_back(std::make_unique(c_host[i])); + c_dev[i]->SetZero(); + + descs.push_back(ck_tile::GroupedGemmHostArgs<>{a_dev[i]->GetDeviceBuffer(), + b_dev[i]->GetDeviceBuffer(), + {}, + c_dev[i]->GetDeviceBuffer(), + kbatch, + M, + N, + K, + sA[i], + sB[i], + {}, + sC[i]}); + } + + const ck_tile::stream_config s{nullptr, true, /*log=*/0, warmup, repeat}; + float ave_time = SelectedKernel::launch(descs, s); + + std::size_t flop = 0, bytes = 0; + for(int i = 0; i < group_count; ++i) + { + flop += std::size_t(2) * Ms[i] * Ns[i] * Ks[i]; + bytes += sizeof(ADataType) * Ms[i] * Ks[i] + sizeof(BDataType) * Ks[i] * Ns[i] + + sizeof(CDataType) * Ms[i] * Ns[i]; + } + const float tflops = static_cast(flop) / 1.E9 / ave_time; + const float gbps = static_cast(bytes) / 1.E6 / ave_time; + std::cout << "Perf: " << std::setw(10) << ave_time << " ms, " << tflops << " TFlops, " << gbps + << " GB/s\n"; + + for(int i = 0; i < group_count; ++i) + c_dev[i]->FromDevice(c_host[i].data()); + + bool pass = true; + if(validate) + { + for(int i = 0; i < group_count; ++i) + { + ck_tile::HostTensor ref( + ck_tile::host_tensor_descriptor(Ms[i], Ns[i], sC[i], is_row_major(CLayout{}))); + ref.SetZero(); + ck_tile::reference_gemm( + a_host[i], b_host[i], ref); + const float maxv = *std::max_element(ref.mData.begin(), ref.mData.end()); + const auto rtol = ck_tile::get_relative_threshold( + ck_tile::integer_divide_ceil(Ks[i], kbatch)); + const auto atol = ck_tile::get_absolute_threshold( + maxv / kbatch, ck_tile::integer_divide_ceil(Ks[i], kbatch)); + bool ok = + ck_tile::check_err(c_host[i], ref, "group[" + std::to_string(i) + "]", rtol, atol); + pass &= ok; + } + std::cout << "Verification: " << (pass ? "PASS" : "FAIL") << "\n"; + } + + return pass ? 0 : 1; +} diff --git a/dispatcher/examples/gemm/cpp/README.md b/dispatcher/examples/gemm/cpp/README.md index 79d60d1198..6b84443ea4 100644 --- a/dispatcher/examples/gemm/cpp/README.md +++ b/dispatcher/examples/gemm/cpp/README.md @@ -37,6 +37,7 @@ cd examples | [04_heuristics.cpp](04_heuristics.cpp) | Heuristic-based kernel selection | | [05_json_export.cpp](05_json_export.cpp) | Registry JSON export for external tools | | [06_multi_registry.cpp](06_multi_registry.cpp) | Multiple registries with named kernel sets | +| [02_grouped_gemm_driver.cpp](02_grouped_gemm_driver.cpp) | Standalone grouped (batched) GEMM driver: builds per-group descriptors, launches `GroupedGemmKernel`, verifies each group | ## Example Details @@ -113,6 +114,31 @@ Consolidated example combining performance benchmarking with correctness validat - GPU reference validation using `ck_tile::reference_gemm_gpu` - Configurable tolerances +### 02_grouped_gemm_driver.cpp - Grouped (Batched) GEMM +Standalone driver for the `grouped` codegen variant. One generated grouped kernel header is +injected on the command line (`-include _grouped.hpp -DCK_TILE_SINGLE_KERNEL_INCLUDE`); +the driver reads the kernel's `ADataType/BDataType/CDataType/ALayout/BLayout/CLayout` so it +works for any datatype/layout the kernel was generated for (no hardcoded `fp16`/`rcr`). + +```bash +# Build against one generated grouped kernel +hipcc -std=c++17 --offload-arch=gfx942 -O3 -DCK_TILE_SINGLE_KERNEL_INCLUDE \ + -I ../../../include -I ../../../../include \ + -I generated_kernels \ + -include gemm_fp16_rcr_compv3_..._grouped.hpp \ + 02_grouped_gemm_driver.cpp -o gemm_02_grouped + +# Run: 8 groups, verify each group against the CPU reference +./gemm_02_grouped --groups 8 --Ms 3840 --Ns 4096 --Ks 2048 \ + --warmup 50 --repeat 100 --validate 1 +``` + +**Features:** +- Builds a vector of per-group `GroupedGemmHostArgs` descriptors (per-group M/N/K) +- Layout-driven leading dimensions via `get_default_stride(is_row_major(Layout))` +- Per-group correctness check using `ck_tile::reference_gemm` +- Reports per-group PASS/FAIL plus aggregate TFLOPS + ### 04_heuristics.cpp - Heuristic Selection Demonstrates custom kernel selection based on problem characteristics: diff --git a/dispatcher/python/gemm_utils.py b/dispatcher/python/gemm_utils.py index 160c0f0694..26e68c061b 100644 --- a/dispatcher/python/gemm_utils.py +++ b/dispatcher/python/gemm_utils.py @@ -59,6 +59,88 @@ def _cap(flag: bool) -> str: return "True" if flag else "False" +# --------------------------------------------------------------------------- +# Dtype codecs: map a bridge dtype token -> numpy dtype for host operands. +# +# fp16 maps to plain numpy; bf16/fp8/bf8 need ml_dtypes. fp8/bf8 use the FNUZ +# encodings (E4M3FNUZ / E5M2FNUZ) that the gfx942 MFMA path expects -- matching +# the regular bridge's fp8/bf8 codec (PR #8887). ml_dtypes is imported lazily so +# the fp16-only path keeps working where ml_dtypes is unavailable. +# --------------------------------------------------------------------------- + +# Canonicalize common spellings to a single token. +_DTYPE_ALIASES = { + "fp16": "fp16", + "f16": "fp16", + "half": "fp16", + "float16": "fp16", + "bf16": "bf16", + "bfloat16": "bf16", + "fp8": "fp8", + "fp8_e4m3": "fp8", + "e4m3": "fp8", + "bf8": "bf8", + "fp8_e5m2": "bf8", + "e5m2": "bf8", +} + + +def numpy_dtype_for(dtype: str): + """Return the numpy dtype object used for host operands of ``dtype``. + + fp16 -> np.float16; bf16/fp8/bf8 require the ``ml_dtypes`` package (imported + lazily) and use FNUZ fp8 encodings for gfx942 parity. + """ + token = _DTYPE_ALIASES.get(str(dtype).lower()) + if token is None: + raise ValueError(f"Unsupported grouped GEMM dtype: {dtype!r}") + if token == "fp16": + return np.float16 + try: + import ml_dtypes # noqa: WPS433 (lazy: optional dep) + except ImportError as exc: # pragma: no cover - env-dependent + raise RuntimeError( + f"dtype {dtype!r} requires the 'ml_dtypes' package (pip install ml_dtypes)" + ) from exc + if token == "bf16": + return np.dtype(ml_dtypes.bfloat16) + if token == "fp8": + return np.dtype(ml_dtypes.float8_e4m3fnuz) + if token == "bf8": + return np.dtype(ml_dtypes.float8_e5m2fnuz) + raise ValueError(f"Unsupported grouped GEMM dtype: {dtype!r}") # pragma: no cover + + +def output_dtype_for(dtype: str) -> str: + """Return the bridge dtype token of a kernel's OUTPUT for input ``dtype``. + + Mirrors ``codegen_common.CommonTypeMappings.get_output_dtype`` (fp8/bf8 -> + fp16, else identity): the generated grouped kernel emits an fp16 ``CDataType`` + for fp8/bf8 inputs, so the host C buffer must be sized/typed by the OUTPUT + dtype, not the INPUT dtype. ``codegen_common`` lives on the dispatcher + ``codegen`` dir which ctypes_utils already puts on ``sys.path``; import it + lazily so the fp16-only path has no extra dependency. + """ + token = _DTYPE_ALIASES.get(str(dtype).lower()) + if token is None: + raise ValueError(f"Unsupported grouped GEMM dtype: {dtype!r}") + try: + from codegen_common import CommonTypeMappings # noqa: WPS433 (lazy) + except ImportError: # pragma: no cover - fall back to the documented mapping + return "fp16" if token in ("fp8", "bf8") else token + return CommonTypeMappings.get_output_dtype(token) + + +def output_numpy_dtype_for(dtype: str): + """Numpy dtype of a kernel's OUTPUT buffer for input ``dtype``. + + Composition of :func:`output_dtype_for` + :func:`numpy_dtype_for`. For + fp8/bf8 this resolves to ``np.float16`` (2 bytes) because the kernel's + ``CDataType`` is fp16; for fp16/bf16 it equals the input dtype. + """ + return numpy_dtype_for(output_dtype_for(dtype)) + + # ============================================================================ # The shared contract: GemmKernelConfig # ============================================================================ @@ -151,6 +233,8 @@ class GemmKernelConfig: name += "_preshuffle" elif self.variant == "streamk": name += "_streamk" + elif self.variant == "grouped": + name += "_grouped" return name # ------------------------------------------------------------------ # @@ -264,6 +348,39 @@ class GemmProblem: return cls(M=int(d["M"]), N=int(d["N"]), K=int(d["K"])) +@dataclass +class GroupedGemmProblem: + """A grouped GEMM problem: a list of independent (M, N, K) sub-problems + all run by a single grouped kernel launch. + + Each group g computes C_g[M_g x N_g] = A_g[M_g x K_g] @ B_g[K_g x N_g]. + """ + + groups: List[Tuple[int, int, int]] + + @classmethod + def uniform( + cls, group_count: int, M: int, N: int, K: int + ) -> "GroupedGemmProblem": + """All groups share the same (M, N, K) shape.""" + return cls(groups=[(int(M), int(N), int(K)) for _ in range(int(group_count))]) + + @property + def group_count(self) -> int: + return len(self.groups) + + @property + def flops(self) -> float: + return sum(2.0 * m * n * k for (m, n, k) in self.groups) + + def to_dict(self) -> Dict[str, Any]: + return {"groups": [[int(m), int(n), int(k)] for (m, n, k) in self.groups]} + + @classmethod + def from_dict(cls, d: Dict[str, Any]) -> "GroupedGemmProblem": + return cls(groups=[(int(m), int(n), int(k)) for (m, n, k) in d["groups"]]) + + @dataclass class GemmResult: output: np.ndarray @@ -277,6 +394,22 @@ class GemmResult: return self.status == 0 +@dataclass +class GroupedGemmResult: + """Result of a grouped GEMM launch: one output per group plus aggregate + timing/throughput across the whole batch.""" + + outputs: List[np.ndarray] + time_ms: float + status: int + tflops: float + kernel_name: str + + @property + def success(self) -> bool: + return self.status == 0 + + # ============================================================================ # ctypes ABI wrapper # ============================================================================ @@ -294,6 +427,8 @@ class GemmDispatcherLib: self._path = Path(so_path) self._lib = ctypes.CDLL(str(self._path)) self._has_indexed = hasattr(self._lib, "dispatcher_get_kernel_name_at") + self._has_grouped = hasattr(self._lib, "dispatcher_run_grouped_gemm") + self._has_single = hasattr(self._lib, "dispatcher_run_gemm") self._setup_functions() def _setup_functions(self) -> None: @@ -316,16 +451,32 @@ class GemmDispatcherLib: ] lib.dispatcher_get_kernel_name_at.restype = ctypes.c_int - lib.dispatcher_run_gemm.argtypes = [ - ctypes.c_void_p, # A (host) - ctypes.c_void_p, # B (host) - ctypes.c_void_p, # C (host) - ctypes.c_int64, # M - ctypes.c_int64, # N - ctypes.c_int64, # K - ctypes.POINTER(ctypes.c_float), # time_ms - ] - lib.dispatcher_run_gemm.restype = ctypes.c_int + # Single-problem ABI (regular GEMM .so). Absent on grouped libs. + if self._has_single: + lib.dispatcher_run_gemm.argtypes = [ + ctypes.c_void_p, # A (host) + ctypes.c_void_p, # B (host) + ctypes.c_void_p, # C (host) + ctypes.c_int64, # M + ctypes.c_int64, # N + ctypes.c_int64, # K + ctypes.POINTER(ctypes.c_float), # time_ms + ] + lib.dispatcher_run_gemm.restype = ctypes.c_int + + # Multi-problem ABI (grouped GEMM .so). Absent on regular libs. + if self._has_grouped: + lib.dispatcher_run_grouped_gemm.argtypes = [ + ctypes.c_int, # group_count + ctypes.POINTER(ctypes.c_int64), # Ms[] + ctypes.POINTER(ctypes.c_int64), # Ns[] + ctypes.POINTER(ctypes.c_int64), # Ks[] + ctypes.POINTER(ctypes.c_void_p), # A_ptrs[] + ctypes.POINTER(ctypes.c_void_p), # B_ptrs[] + ctypes.POINTER(ctypes.c_void_p), # C_ptrs[] + ctypes.POINTER(ctypes.c_float), # time_ms + ] + lib.dispatcher_run_grouped_gemm.restype = ctypes.c_int lib.dispatcher_cleanup.argtypes = [] lib.dispatcher_cleanup.restype = None @@ -371,6 +522,52 @@ class GemmDispatcherLib: ) return status, time_ms.value + def run_grouped( + self, + A_list: List[np.ndarray], + B_list: List[np.ndarray], + C_list: List[np.ndarray], + Ms: List[int], + Ns: List[int], + Ks: List[int], + ) -> Tuple[int, float]: + """Launch the grouped kernel over a batch of (M, N, K) sub-problems. + + Each A/B/C entry is a host numpy array already laid out (dtype + row/col + transpose) as the kernel expects for its compile-time layout; the caller + (GpuGroupedGemmRunner) does that per-dtype/per-layout packing. Pointers + are marshalled into ctypes pointer arrays. + """ + if not self._has_grouped: + raise RuntimeError( + f"{self._path} does not expose dispatcher_run_grouped_gemm" + ) + + g = len(A_list) + c_int64_arr = (ctypes.c_int64 * g) + c_void_arr = (ctypes.c_void_p * g) + + ms = c_int64_arr(*[int(m) for m in Ms]) + ns = c_int64_arr(*[int(n) for n in Ns]) + ks = c_int64_arr(*[int(k) for k in Ks]) + + a_ptrs = c_void_arr(*[A.ctypes.data_as(ctypes.c_void_p) for A in A_list]) + b_ptrs = c_void_arr(*[B.ctypes.data_as(ctypes.c_void_p) for B in B_list]) + c_ptrs = c_void_arr(*[C.ctypes.data_as(ctypes.c_void_p) for C in C_list]) + + time_ms = ctypes.c_float(0.0) + status = self._lib.dispatcher_run_grouped_gemm( + g, + ms, + ns, + ks, + a_ptrs, + b_ptrs, + c_ptrs, + ctypes.byref(time_ms), + ) + return status, time_ms.value + def cleanup(self) -> None: self._lib.dispatcher_cleanup() @@ -619,6 +816,94 @@ class GpuGemmRunner: ) +class GpuGroupedGemmRunner: + """High-level runner for the GROUPED variant: construct from a grouped .so + path, call run(A_list, B_list, problem). + + Like GpuGemmRunner, the ctypes ABI takes HOST pointers and manages GPU + memory internally (per group), so this runner only marshals the host operand + arrays. The runner is parameterized by ``(dtype, layout)`` (mirroring + ``GpuGemmRunner``/``GemmProblem``): the A/B operands are cast to the per-dtype + INPUT numpy codec (fp16/bf16/fp8-E4M3FNUZ/bf8-E5M2FNUZ) and transposed per the + A/B/C layout so the contiguous host buffer matches the layout the kernel was + generated with (the ctypes lib derives strides from the same layouts). + + The C/output buffer is sized/typed by the kernel's OUTPUT dtype, not the input + dtype: for fp8/bf8 inputs the generated kernel's ``CDataType`` is fp16, so the + host C buffer is fp16 (2 bytes) even though A/B are 1-byte fp8/bf8. Sizing C by + the input dtype would under-allocate by 2x and the ctypes copy-back would + overrun the host buffer (heap corruption). See :func:`output_numpy_dtype_for`. + """ + + def __init__(self, lib_path: Path, dtype: str = "fp16", layout: str = "rcr"): + self.lib = GemmDispatcherLib(lib_path) + if not self.lib.initialize(): + raise RuntimeError( + f"Failed to initialize grouped dispatcher .so: {lib_path}" + ) + names = self.lib.kernel_names + self._kernel_name = names[0] if names else "unknown" + self._dtype = dtype + # A/B (input) codec vs C (output) codec: they differ for fp8/bf8 + # (output is fp16), so keep them distinct to size the C buffer correctly. + self._np_dtype = numpy_dtype_for(dtype) + self._c_np_dtype = output_numpy_dtype_for(dtype) + if len(layout) != 3 or any(ch not in ("r", "c") for ch in layout): + raise ValueError(f"layout must be a 3-char r/c string, got {layout!r}") + self._layout = layout + + @property + def kernel_name(self) -> str: + return self._kernel_name + + def run( + self, + A_list: List[np.ndarray], + B_list: List[np.ndarray], + problem: GroupedGemmProblem, + ) -> GroupedGemmResult: + groups = problem.groups + if len(A_list) != len(groups) or len(B_list) != len(groups): + raise ValueError( + "A_list/B_list length must match the number of groups " + f"({len(A_list)}/{len(B_list)} vs {len(groups)})" + ) + + Ms = [g[0] for g in groups] + Ns = [g[1] for g in groups] + Ks = [g[2] for g in groups] + + la, lb, _lc = self._layout[0], self._layout[1], self._layout[2] + nd = self._np_dtype + c_nd = self._c_np_dtype # OUTPUT dtype (fp16 for fp8/bf8); see __init__. + + A_h: List[np.ndarray] = [] + B_h: List[np.ndarray] = [] + C_h: List[np.ndarray] = [] + for A, B, (M, N, _K) in zip(A_list, B_list, groups): + # A logically MxK, B logically KxN, C row-major MxN (CLayout is always + # RowMajor for grouped). Store each operand so its contiguous buffer + # matches its layout: row-major -> as-is, col-major -> transpose. + A_buf = A if la == "r" else A.T + B_buf = B if lb == "r" else B.T + A_h.append(np.ascontiguousarray(A_buf, dtype=nd)) + B_h.append(np.ascontiguousarray(B_buf, dtype=nd)) + # Size C by the kernel's CDataType (output dtype), NOT the input dtype: + # fp8/bf8 inputs produce fp16 output, so a 1-byte C would be overrun. + C_h.append(np.zeros((M, N), dtype=c_nd)) + + status, time_ms = self.lib.run_grouped(A_h, B_h, C_h, Ms, Ns, Ks) + + tflops = (problem.flops / (time_ms * 1e-3)) / 1e12 if time_ms > 0 else 0.0 + return GroupedGemmResult( + outputs=C_h, + time_ms=time_ms, + status=status, + tflops=tflops, + kernel_name=self._kernel_name, + ) + + # ============================================================================ # Build API: codegen + hipcc -> .so paths (no GPU) # ============================================================================ @@ -697,6 +982,18 @@ def _tile_engine_codegen_flags() -> Tuple[str, ...]: return tuple(flags) +def _ctypes_source_name(config: GemmKernelConfig) -> str: + """Pick the ctypes ABI source for a config's variant. + + The grouped kernel has a multi-problem launch signature that the + single-problem ``gemm_ctypes_lib.cpp`` cannot express, so grouped configs + compile against the dedicated ``grouped_gemm_ctypes_lib.cpp``. + """ + if config.variant == "grouped": + return "grouped_gemm_ctypes_lib.cpp" + return "gemm_ctypes_lib.cpp" + + def _build_compile_jobs( config: GemmKernelConfig, header: Path ) -> Tuple[Dict[str, Any], Path]: @@ -705,7 +1002,7 @@ def _build_compile_jobs( ck_root = root.parent build_dir = _cu.get_build_dir() output_dir = _cu.get_generated_kernels_dir() - ctypes_source = root / "bindings" / "ctypes" / "gemm_ctypes_lib.cpp" + ctypes_source = root / "bindings" / "ctypes" / _ctypes_source_name(config) static_lib = build_dir / "libck_tile_dispatcher.a" lib_path = build_dir / "examples" / f"lib{config.name}.so" @@ -786,13 +1083,13 @@ def setup_multiple_gemm_dispatchers( codegen_script = _cu.get_codegen_path() output_dir = _cu.get_generated_kernels_dir() static_lib = _cu.get_build_dir() / "libck_tile_dispatcher.a" - ctypes_source = ( - _cu.get_dispatcher_root() / "bindings" / "ctypes" / "gemm_ctypes_lib.cpp" - ) - if not static_lib.exists() or not ctypes_source.exists(): + ctypes_dir = _cu.get_dispatcher_root() / "bindings" / "ctypes" + needed_sources = {ctypes_dir / _ctypes_source_name(c) for c in configs} + missing = [str(p) for p in needed_sources if not p.exists()] + if not static_lib.exists() or missing: raise FileNotFoundError( "Missing static lib or ctypes source required for compilation:\n" - f" {static_lib}\n {ctypes_source}\n" + f" {static_lib}\n " + "\n ".join(missing) + "\n" "Build the dispatcher first (cmake + make)." ) @@ -915,6 +1212,7 @@ def expand_sweep( arch: str, dtype: str = "fp16", layout: str = "rcr", + variant: str = "standard", ) -> List[GemmKernelConfig]: """Expand a Tile Engine GEMM JSON sweep config into GemmKernelConfig list. @@ -924,8 +1222,8 @@ def expand_sweep( one GemmKernelConfig. Invalid combinations are dropped via the dispatcher's own validator, and duplicates (by .name) are collapsed. - The signature is controlled by the `dtype` and `layout` arguments (defaults - to fp16 / rcr). + The operand signature (``dtype``, ``layout``) is applied to every emitted + GemmKernelConfig, so the same sweep expands across any supported dtype/layout. """ with open(config_path) as f: cfg = json.load(f) @@ -1015,6 +1313,7 @@ def expand_sweep( pad_k=bool(pk), persistent=bool(persist), gfx_arch=arch, + variant=variant, ) if c.name in seen: continue diff --git a/dispatcher/tests/test_grouped_gemm_codegen.py b/dispatcher/tests/test_grouped_gemm_codegen.py new file mode 100644 index 0000000000..4dd2891f63 --- /dev/null +++ b/dispatcher/tests/test_grouped_gemm_codegen.py @@ -0,0 +1,130 @@ +#!/usr/bin/env python3 + +# Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +# SPDX-License-Identifier: MIT + +"""CPU-only unit tests for the grouped GEMM variant of the TE -> Dispatcher bridge. + +Grouped GEMM has a genuinely different ABI from regular GEMM (a list of +sub-problems, a launch(descs, stream) signature, and an internally-owned device +workspace), so it is generated by a dedicated path in the unified GEMM codegen +and launched by grouped_gemm_ctypes_lib.cpp. These tests lock in the host-side +contract of that path without touching a GPU: + + * the GROUPED variant generates the grouped launch/kernel with fully-qualified + ck_tile:: types (so the header does not silently depend on a using-directive); + * the grouped ctypes bridge includes the standard headers it uses directly + rather than relying on transitive includes from the force-included kernel; + * the FNUZ decode table shared by the fp8/bf8 path is cached and read-only. + +No GPU is touched -- codegen is pure string formatting and the helpers are pure. +Run: python3 -m pytest tests/test_grouped_gemm_codegen.py -v +""" + +import sys +import unittest +from pathlib import Path + +SCRIPT_DIR = Path(__file__).parent.resolve() +DISPATCHER_DIR = SCRIPT_DIR.parent +sys.path.insert(0, str(DISPATCHER_DIR / "codegen")) +sys.path.insert(0, str(DISPATCHER_DIR / "python")) + +from codegen_common import TileConfig # noqa: E402 + +from unified_gemm_codegen import ( # noqa: E402 + GemmVariant, + KernelConfig, + TraitConfig, + CKTileKernelGenerator, +) + +from gemm_utils import _fnuz_decode_table # noqa: E402 + +GROUPED_CTYPES_LIB = ( + DISPATCHER_DIR / "bindings" / "ctypes" / "grouped_gemm_ctypes_lib.cpp" +) + + +def _grouped_config() -> KernelConfig: + """A minimal valid grouped-GEMM kernel config (fp16/rcr, cshuffle).""" + tile = TileConfig( + tile_m=128, tile_n=128, tile_k=32, + warp_m=2, warp_n=2, warp_k=1, + warp_tile_m=32, warp_tile_n=32, warp_tile_k=16, + ) + trait = TraitConfig( + pipeline="compv3", epilogue="cshuffle", scheduler="intrawave", + pad_m=False, pad_n=False, pad_k=False, persistent=False, + ) + return KernelConfig(tile=tile, trait=trait, variant=GemmVariant.GROUPED) + + +class TestGroupedVariant(unittest.TestCase): + def test_enum_value(self): + self.assertEqual(GemmVariant.GROUPED.value, "grouped") + + +class TestGroupedCodegen(unittest.TestCase): + """The generated grouped header must be self-contained and correctly typed.""" + + def setUp(self): + gen = CKTileKernelGenerator("fp16", "rcr") + self.src = gen.generate(_grouped_config()) + + def test_generates_grouped_launch_and_kernel(self): + # The grouped path emits the (descs, stream) launch and the grouped kernel. + self.assertIn( + "launch(const std::vector>&", self.src + ) + self.assertIn("ck_tile::GroupedGemmKernel<", self.src) + + def test_device_memory_is_ck_tile_qualified(self): + # DeviceMem lives in the ck_tile namespace; the workspace allocation must + # be fully qualified so the header does not depend on a using-directive. + self.assertIn("ck_tile::DeviceMem workspace_dev", self.src) + # And no bare, unqualified `DeviceMem ` declaration should remain. + import re + + self.assertIsNone( + re.search(r"(?{} descriptors are built here; must be direct. + self.assertIn("std::array", text) + self.assertIn("#include ", text) + + +class TestFnuzTableCaching(unittest.TestCase): + """The 256-entry FNUZ decode table is pure per (exp_bits, mant_bits): it must + be cached and handed out read-only so callers cannot mutate the shared copy.""" + + def test_table_is_cached_same_object(self): + self.assertIs(_fnuz_decode_table(4, 3), _fnuz_decode_table(4, 3)) + self.assertIs(_fnuz_decode_table(5, 2), _fnuz_decode_table(5, 2)) + + def test_table_is_read_only(self): + table = _fnuz_decode_table(4, 3) + self.assertFalse(table.flags.writeable) + with self.assertRaises(ValueError): + table[0] = 1.0 + + +if __name__ == "__main__": + unittest.main() diff --git a/tile_engine/ops/gemm/grouped_gemm_full_benchmark.py b/tile_engine/ops/gemm/grouped_gemm_full_benchmark.py new file mode 100644 index 0000000000..eaed58c3f8 --- /dev/null +++ b/tile_engine/ops/gemm/grouped_gemm_full_benchmark.py @@ -0,0 +1,327 @@ +#!/usr/bin/env python3 +"""Full GROUPED GEMM benchmark sweep driven through the Dispatcher bridge. + +The grouped counterpart of gemm_full_benchmark.py. Same 3-phase architecture: + Phase 1: Compile all grouped kernels (parallel, returns .so paths only -- no GPU) + Phase 2: Load grouped problems (lists of (M, N, K) sub-problems) + Phase 3: Benchmark via subprocess isolation (serial GPU, batched) + +Tile Engine generates NO binaries here: it expands its grouped sweep config into +shared ``GemmKernelConfig(variant="grouped")`` objects and hands them to the +dispatcher, which codegens + compiles each into a grouped .so (built against +``grouped_gemm_ctypes_lib.cpp``). Each kernel runs in a disposable worker +subprocess so a GPU fault takes down only one worker. + +Usage: + python grouped_gemm_full_benchmark.py grouped_gemm/configs/default_config.json \ + --arch gfx942 --csv grouped_gemm_results.csv + +A grouped "problem" is a batch of sub-problems run by one launch. The default set +uses uniform shapes at two group counts; override with --problems pointing at a +JSON file of [{"groups": [[M,N,K], ...]}, ...]. +""" + +import argparse +import csv +import json +import os +import subprocess +import sys +import time +from pathlib import Path + +_THIS_DIR = Path(__file__).resolve().parent +_DISPATCHER_ROOT = _THIS_DIR.parents[2] / "dispatcher" +sys.path.insert(0, str(_DISPATCHER_ROOT / "python")) +sys.path.insert(0, str(_THIS_DIR)) + +from gemm_utils import setup_multiple_gemm_dispatchers, expand_sweep # noqa: E402 + +# Default grouped problem set: uniform groups at two batch sizes (Phase 3 parity). +DEFAULT_PROBLEMS = [ + {"groups": [[1024, 1024, 1024]] * 4}, + {"groups": [[2048, 2048, 2048]] * 4}, + {"groups": [[1024, 1024, 1024]] * 8}, + {"groups": [[512, 1024, 2048], [1024, 512, 1024], [2048, 2048, 512], [256, 768, 1024]]}, +] + + +def load_problems(path): + if not path: + return DEFAULT_PROBLEMS + with open(path) as f: + data = json.load(f) + # Accept either a bare list or {"problems": [...]}. + return data["problems"] if isinstance(data, dict) else data + + +def _problem_dims(prob): + """Aggregate (group_count, total_flops) for reporting.""" + groups = prob["groups"] + flops = sum(2.0 * m * n * k for (m, n, k) in groups) + return len(groups), flops + + +def main(): + parser = argparse.ArgumentParser( + description="Grouped GEMM Benchmark Sweep (via Dispatcher)" + ) + parser.add_argument("configs", nargs="+", help="TE grouped sweep config JSON files") + parser.add_argument("--arch", default="gfx942") + parser.add_argument("--dtype", default="fp16") + parser.add_argument("--layout", default="rcr") + parser.add_argument( + "--problems", default=None, help="JSON file of grouped problems" + ) + parser.add_argument("--csv", type=str, default="grouped_gemm_results.csv") + parser.add_argument("--workers", type=int, default=8, help="Parallel build workers") + parser.add_argument( + "--batch-size", + type=int, + default=20, + help="Kernels per subprocess (overhead vs fault isolation)", + ) + parser.add_argument( + "--kernel-timeout", type=int, default=60, help="Per-kernel timeout (s)" + ) + parser.add_argument( + "--max-kernels", type=int, default=0, help="Limit to first N kernels (0=all)" + ) + args = parser.parse_args() + + # ======================================================================== + # Phase 1: Compile kernels (parallel, no GPU) + # ======================================================================== + print(f"\n{'=' * 80}") + print("Phase 1: Compile grouped kernels") + print(f"{'=' * 80}") + + all_configs = [] + for cfg_path in args.configs: + all_configs.extend( + expand_sweep( + cfg_path, + args.arch, + dtype=args.dtype, + layout=args.layout, + variant="grouped", + ) + ) + + if args.max_kernels > 0: + all_configs = all_configs[: args.max_kernels] + + print(f" Expanded configs: {len(all_configs)}") + print(f" Build workers: {args.workers}") + + t0 = time.perf_counter() + # CRITICAL: returns Path objects only, does NOT load any .so. + lib_paths = setup_multiple_gemm_dispatchers( + all_configs, verbose=True, max_workers=args.workers + ) + build_time = time.perf_counter() - t0 + + built_kernels = [ + (cfg, lib) for cfg, lib in zip(all_configs, lib_paths) if lib is not None + ] + + # Dedupe by .so path (distinct configs can map to the same physical kernel). + seen_libs = set() + unique_kernels = [] + duplicate_count = 0 + for cfg, lib in built_kernels: + lib_key = str(lib.resolve()) + if lib_key not in seen_libs: + seen_libs.add(lib_key) + unique_kernels.append((cfg, lib)) + else: + duplicate_count += 1 + built_kernels = unique_kernels + + print( + f"\n Built {len(all_configs)} configs -> {len(built_kernels)} unique kernels " + f"({duplicate_count} duplicates filtered) in {build_time:.0f}s" + ) + + if not built_kernels: + print(" ERROR: No kernels built successfully") + return 1 + + # ======================================================================== + # Phase 2: Load problems + # ======================================================================== + print(f"\n{'=' * 80}") + print("Phase 2: Load grouped test problems") + print(f"{'=' * 80}") + + problems = load_problems(args.problems) + print(f" Problems: {len(problems)}") + print( + f" Total measurements: {len(built_kernels)} x {len(problems)} = " + f"{len(built_kernels) * len(problems)}" + ) + + # ======================================================================== + # Phase 3: Benchmark via subprocess (serial GPU, batched) + # ======================================================================== + print(f"\n{'=' * 80}") + print("Phase 3: Benchmark (subprocess isolation, batched)") + print(f"{'=' * 80}") + print(f" Batch size: {args.batch_size} kernels per subprocess") + print(f" Timeout: {args.kernel_timeout}s per kernel\n") + + csv_path = Path(args.csv) + csv_fields = [ + "kernel", + "problem_idx", + "group_count", + "total_flops", + "latency_ms", + "tflops", + "non_zero", + ] + csv_file = open(csv_path, "w", newline="") + writer = csv.DictWriter(csv_file, fieldnames=csv_fields) + writer.writeheader() + + worker_path = _THIS_DIR / "run_one_grouped_gemm_kernel.py" + worker_env = os.environ.copy() + worker_env["GEMM_PYPATH"] = os.pathsep.join( + [str(_DISPATCHER_ROOT / "python"), str(_THIS_DIR)] + ) + + total_measurements = 0 + total_failures = 0 + bench_t0 = time.perf_counter() + + for prob_idx, prob in enumerate(problems): + group_count, total_flops = _problem_dims(prob) + print( + f"\nProblem [{prob_idx + 1}/{len(problems)}]: groups={group_count} " + f"flops={total_flops / 1e9:.1f}G ({len(built_kernels)} kernels)" + ) + print(f" {'Kernel':<60} {'Time(ms)':>10} {'TFLOPS':>10} {'Status':>10}") + print(f" {'-' * 95}") + + prob_dict = {"groups": [list(g) for g in prob["groups"]]} + + for batch_start in range(0, len(built_kernels), args.batch_size): + batch_end = min(batch_start + args.batch_size, len(built_kernels)) + batch = built_kernels[batch_start:batch_end] + + items = [ + { + "so_path": str(lib_path), + "problem": prob_dict, + "kernel_name": cfg.name, + "dtype": args.dtype, + "layout": args.layout, + } + for cfg, lib_path in batch + ] + payload = json.dumps({"items": items}) + + try: + proc = subprocess.Popen( + [sys.executable, str(worker_path)], + stdin=subprocess.PIPE, + stdout=subprocess.PIPE, + stderr=subprocess.DEVNULL, + env=worker_env, + ) + timeout_total = args.kernel_timeout * len(batch) + stdout_bytes, _ = proc.communicate( + input=payload.encode("utf-8"), timeout=timeout_total + ) + + reported_indices = set() + for line in stdout_bytes.decode("utf-8").strip().split("\n"): + if not line: + continue + try: + result = json.loads(line) + batch_idx = result.get("idx") + if not isinstance(batch_idx, int) or not ( + 0 <= batch_idx < len(batch) + ): + print( + f" [warn] worker returned out-of-range idx " + f"{batch_idx!r}; skipping line" + ) + continue + cfg, lib_path = batch[batch_idx] + reported_indices.add(batch_idx) + + if result.get("ok", False): + status = "OK" if result.get("non_zero", 0) > 0 else "ZERO" + print( + f" {cfg.name:<60} {result['ms']:>10.3f} " + f"{result['tflops']:>10.2f} {status:>10}" + ) + writer.writerow( + { + "kernel": cfg.name, + "problem_idx": prob_idx, + "group_count": group_count, + "total_flops": total_flops, + "latency_ms": result["ms"], + "tflops": result["tflops"], + "non_zero": result.get("non_zero", 0), + } + ) + csv_file.flush() + total_measurements += 1 + else: + print(f" {cfg.name:<60} FAILED") + print(f" Error: {result.get('error', 'unknown')[:100]}") + total_failures += 1 + except json.JSONDecodeError: + print(f" Warning: could not parse result line: {line[:50]}") + total_failures += 1 + + missing_indices = set(range(len(batch))) - reported_indices + if missing_indices or proc.returncode != 0: + if proc.returncode != 0: + print(f" Worker exited with code {proc.returncode}") + for idx in sorted(missing_indices): + cfg, _ = batch[idx] + print(f" {cfg.name:<60} MISSING (worker crash)") + total_failures += len(missing_indices) + + except subprocess.TimeoutExpired: + print(f" Batch timeout ({len(batch)} kernels)") + try: + proc.kill() + proc.communicate(timeout=5) + except Exception: + pass + total_failures += len(batch) + except Exception as e: + print(f" Batch error: {e}") + try: + if proc and proc.poll() is None: + proc.kill() + except Exception: + pass + total_failures += len(batch) + + bench_time = time.perf_counter() - bench_t0 + csv_file.close() + + # ======================================================================== + # Summary + # ======================================================================== + print(f"\n{'=' * 80}") + print("BENCHMARK COMPLETE") + print(f"{'=' * 80}") + print(f" Build time: {build_time:.0f}s") + print(f" Benchmark time: {bench_time:.0f}s") + print(f" Total time: {build_time + bench_time:.0f}s") + print(f" Successful measurements: {total_measurements}") + print(f" Failed measurements: {total_failures}") + print(f" Output: {csv_path}") + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/tile_engine/ops/gemm/run_one_grouped_gemm_kernel.py b/tile_engine/ops/gemm/run_one_grouped_gemm_kernel.py new file mode 100644 index 0000000000..0a2ea7ed0f --- /dev/null +++ b/tile_engine/ops/gemm/run_one_grouped_gemm_kernel.py @@ -0,0 +1,130 @@ +#!/usr/bin/env python3 +"""Worker script for running GROUPED GEMM kernels in an isolated subprocess. + +Mirrors run_one_gemm_kernel.py, but for the grouped variant: +- Receives a grouped kernel config + grouped problem via stdin as JSON +- Loads the .so library ONLY inside this subprocess +- Outputs timing results as JSON to stdout (one line per kernel, flushed) +- A GPU fault kills only this process; the parent driver can continue + +A grouped problem is a LIST of (M, N, K) sub-problems run by one grouped launch. + +Input JSON format (dtype/layout default to fp16/rcr if absent): + Single: {"so_path": "...", "problem": {"groups": [[M,N,K], ...]}, + "kernel_name": "...", "dtype": "fp16", "layout": "rcr"} + Batch: {"items": [{"so_path": "...", "problem": {...}, "kernel_name": "...", + "dtype": "...", "layout": "..."}, ...]} + +Output JSON format (one line per kernel): + {"idx": 0, "ok": true, "ms": 0.123, "tflops": 456.7, "non_zero": 1, + "group_count": 8, "kernel": "..."} + {"idx": 1, "ok": false, "error": "...", "kernel": "..."} +""" + +import json +import os +import sys + +# Add dispatcher python paths from environment (os.pathsep-separated). +gemm_pypath = os.environ.get("GEMM_PYPATH", "") +if gemm_pypath: + for p in gemm_pypath.split(os.pathsep): + if p and p not in sys.path: + sys.path.insert(0, p) + +from gemm_utils import GroupedGemmProblem, GpuGroupedGemmRunner # noqa: E402 +import numpy as np # noqa: E402 + + +def _run_one(idx, so_path, prob_dict, kernel_name, dtype="fp16", layout="rcr"): + """Run a single grouped kernel and emit its result as one JSON line.""" + try: + problem = GroupedGemmProblem.from_dict(prob_dict) + + # Operands are generated as fp32-ish floats; the runner casts them to the + # per-dtype codec (fp16/bf16/fp8/bf8) and applies layout transposes. + np.random.seed(42) + A_list = [] + B_list = [] + for (M, N, K) in problem.groups: + A_list.append((np.random.randn(M, K) * 0.1).astype(np.float32)) + B_list.append((np.random.randn(K, N) * 0.1).astype(np.float32)) + + # CRITICAL: load the library ONLY inside this subprocess. + runner = GpuGroupedGemmRunner(lib_path=so_path, dtype=dtype, layout=layout) + result = runner.run(A_list, B_list, problem) + + if result.success: + non_zero = sum( + int(np.count_nonzero(o)) for o in result.outputs if o is not None + ) + print( + json.dumps( + { + "idx": idx, + "ok": True, + "ms": result.time_ms, + "tflops": result.tflops, + "non_zero": 1 if non_zero > 0 else 0, + "group_count": problem.group_count, + "kernel": kernel_name, + } + ), + flush=True, + ) + else: + print( + json.dumps( + { + "idx": idx, + "ok": False, + "error": f"kernel returned status {result.status}", + "kernel": kernel_name, + } + ), + flush=True, + ) + + except Exception as e: + print( + json.dumps( + {"idx": idx, "ok": False, "error": str(e), "kernel": kernel_name} + ), + flush=True, + ) + + +def main(): + """Read JSON from stdin, run grouped kernel(s), output results.""" + try: + d = json.loads(sys.stdin.buffer.read()) + except Exception as e: + print( + json.dumps({"idx": 0, "ok": False, "error": f"JSON parse error: {e}"}), + flush=True, + ) + sys.exit(1) + + if "items" in d: + for i, item in enumerate(d["items"]): + _run_one( + i, + item["so_path"], + item["problem"], + item.get("kernel_name", "unknown"), + item.get("dtype", "fp16"), + item.get("layout", "rcr"), + ) + else: + _run_one( + 0, + d["so_path"], + d["problem"], + d.get("kernel_name", "unknown"), + d.get("dtype", "fp16"), + d.get("layout", "rcr"), + ) + + +if __name__ == "__main__": + main()