diff --git a/projects/composablekernel/dispatcher/bindings/ctypes/batched_gemm_ctypes_lib.cpp b/projects/composablekernel/dispatcher/bindings/ctypes/batched_gemm_ctypes_lib.cpp new file mode 100644 index 000000000000..95b8bff4cf93 --- /dev/null +++ b/projects/composablekernel/dispatcher/bindings/ctypes/batched_gemm_ctypes_lib.cpp @@ -0,0 +1,382 @@ +// Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +// SPDX-License-Identifier: MIT + +/** + * Batched GEMM Dispatcher ctypes Library + * + * Provides a C API for Python ctypes integration for the BATCHED GEMM bridge. + * + * Unlike gemm_ctypes_lib.cpp (single-problem GEMM), batched GEMM has a + * divergent ABI: it carries a batch dimension with per-batch strides. The + * registry / Dispatcher::run() path only knows the single-problem + * (A, B, C, M, N, K) signature, so this library BYPASSES the registry and + * launches the force-included kernel directly via + * ``SelectedKernel::launch(ck_tile::BatchedGemmHostArgs{...}, stream)`` -- + * the same launch entry the Tile Engine batched_gemm benchmark uses. This + * mirrors the registry-bypass pattern used by the stream-K bridge. + * + * Usage from Python: + * lib = ctypes.CDLL("libbatched_gemm_....so") + * lib.dispatcher_initialize() // the name the Python wrapper binds + * lib.dispatcher_run_batched(A, B, C, M, N, K, batch_count, k_batch, + * stride_A, stride_B, stride_C, + * batch_stride_A, batch_stride_B, batch_stride_C, + * warmup, repeat, flush_cache, rotating_count, + * &time_ms) + * + * The ABI threads k_batch (split-K, Old-TE supports it) and the full + * benchmarking knobs (warmup/repeat/flush_cache/rotating_count) so this bridge + * drives ``SelectedKernel::launch`` with the SAME ``stream_config`` the Tile + * Engine batched_gemm profiler uses -- a prerequisite for a fair A/B timing + * comparison (the earlier version launched with defaults cold=3/repeat=10). + */ + +#include +#include +#include +#include +#include +#include + +#include "ck_tile/core.hpp" +#include "ck_tile/host.hpp" +#include "ck_tile/ops/gemm.hpp" + +// Kernel header included via -include compiler flag under +// CK_TILE_SINGLE_KERNEL_INCLUDE. Defines: ADataType, BDataType, CDataType, +// AccDataType, ALayout, BLayout, CLayout, SelectedKernel, KERNEL_NAME. + +// GPU architecture - MUST be supplied at compile time via -DGFX_ARCH="". +// Do not default to a specific GPU architecture: the arch is resolved from the +// host (get_arch / rocminfo) and threaded through as a compile flag so the +// kernel is never silently built for the wrong target. +#ifndef GFX_ARCH +#error \ + "GFX_ARCH must be defined at compile time (pass -DGFX_ARCH=); do not default to a specific GPU architecture." +#endif + +namespace { + +// The batched bridge is single-kernel-per-.so: the force-included header fully +// determines the kernel. There is no registry to initialize, but the init entry +// is kept for ABI symmetry with the single-problem GEMM library so the Python +// runner can call initialize() uniformly. +bool g_initialized = false; + +// Log a failing HIP call (with hipGetErrorString) before the caller returns an +// error status: a bare -1 with no diagnostic makes a real GPU sweep failure +// impossible to triage. Returns -1 so call sites read `return fail_hip(...)`. +inline int fail_hip(const char* what, hipError_t err) +{ + std::cerr << "dispatcher_run_batched: " << what << " failed: " << hipGetErrorString(err) + << std::endl; + return -1; +} + +// Default (contiguous) stride for a row/col-major operand, matching +// ck_tile::get_default_stride used by the Tile Engine profiler: a value of 0 +// means "packed", so derive it from the problem shape. +inline std::int64_t +default_stride(std::int64_t rows, std::int64_t cols, std::int64_t provided, bool row_major) +{ + // Match ck_tile::get_default_stride / the TE profiler exactly: a provided + // value of <= 0 means "packed" and the leading stride is derived from the + // shape. (N2: the earlier `> 0` check disagreed with TE's `<= 0` sentinel; + // they are equivalent for legal strides but this is the exact form.) + if(provided <= 0) + { + return row_major ? cols : rows; + } + return provided; +} + +// Per-operand layout is emitted by the codegen as single-char strings +// GEMM_KEY_LAYOUT_A/B/C ("r"/"c"). The generated header only exports the data +// types (not the layout aliases) into the global namespace, so we read layout +// from the macros rather than std::is_same_v. +inline bool operand_is_row_major(const char* layout_char) { return layout_char[0] == 'r'; } + +#ifdef GEMM_KEY_LAYOUT_A +constexpr const char* kLayoutA = GEMM_KEY_LAYOUT_A; +constexpr const char* kLayoutB = GEMM_KEY_LAYOUT_B; +constexpr const char* kLayoutC = GEMM_KEY_LAYOUT_C; +#else +constexpr const char* kLayoutA = "r"; +constexpr const char* kLayoutB = "c"; +constexpr const char* kLayoutC = "r"; +#endif + +} // namespace + +extern "C" { + +/** + * Initialize the library. No registry is used for batched GEMM (the kernel is + * force-included), so this simply flips a flag. Returns 0 on success. + */ +int dispatcher_initialize() +{ + g_initialized = true; + return 0; +} + +// Back-compat alias. `dispatcher_initialize` is the canonical entry (the one the +// Python BatchedGemmDispatcherLib binds); `dispatcher_init` is kept only so the +// same symbol name works across every sibling bridge .so that historically +// exported it. New callers should use dispatcher_initialize. +int dispatcher_init() { return dispatcher_initialize(); } + +/** + * Report the compile-time kernel name of the force-included batched kernel. + * The batched bridge is always one kernel per .so. + */ +const char* dispatcher_get_kernel_name() { return KERNEL_NAME; } + +/** + * Multi-kernel ABI shim: the batched .so exposes exactly one kernel, so index 0 + * returns KERNEL_NAME and every other index fails. Mirrors the single-problem + * library so the shared Python wrapper can query names uniformly. + */ +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; +} + +int dispatcher_get_kernel_count() { return 1; } + +/** + * Run a batched GEMM on the GPU via the force-included kernel. + * + * Takes HOST pointers and manages GPU memory internally (hipMalloc/hipMemcpy/ + * hipFree), matching the single-problem GEMM ABI. The per-batch strides let the + * caller lay out A/B/C as [batch_count, rows, cols] tensors; a stride argument + * of 0 falls back to the packed/default stride. + * + * Returns: 0 on success; -1 on a HIP error or bad/guarded arguments (incl. an + * out-of-range 32-bit index); -2 if the kernel launch throws. + */ +int dispatcher_run_batched(const void* A, + const void* B, + void* C, + std::int64_t M, + std::int64_t N, + std::int64_t K, + std::int64_t batch_count, + std::int64_t k_batch, + std::int64_t stride_A, + std::int64_t stride_B, + std::int64_t stride_C, + std::int64_t batch_stride_A, + std::int64_t batch_stride_B, + std::int64_t batch_stride_C, + std::int64_t warmup, + std::int64_t repeat, + std::int64_t flush_cache, + std::int64_t rotating_count, + float* time_ms) +{ + if(!g_initialized || !A || !B || !C || M <= 0 || N <= 0 || K <= 0 || batch_count <= 0) + { + if(time_ms) + { + *time_ms = -1.0f; + } + return -1; + } + + // split-K count: <= 0 falls back to 1 (no split), matching the TE default. + const ck_tile::index_t kbatch = static_cast(k_batch > 0 ? k_batch : 1); + + // Resolve leading strides: <= 0 -> packed default (matches TE profiler / + // ck_tile::get_default_stride behaviour), respecting each operand's own + // row/col-major layout. For rcr this yields stride_A=K, stride_B=K, + // stride_C=N. + const std::int64_t sa = default_stride(M, K, stride_A, operand_is_row_major(kLayoutA)); + const std::int64_t sb = default_stride(K, N, stride_B, operand_is_row_major(kLayoutB)); + const std::int64_t sc = default_stride(M, N, stride_C, operand_is_row_major(kLayoutC)); + + // Batch strides: <= 0 -> packed default (contiguous per-batch slab). The + // packed slab spans the RESOLVED leading stride (sa/sb/sc), not the bare + // shape, so a non-packed leading stride (padded rows/cols) is honoured: a + // row-major operand's slab is rows*leading and a col-major operand's slab is + // cols*leading. Deriving the default from M*K/K*N/M*N would under-size the + // buffer and cause out-of-bounds access whenever stride_A/B/C is padded. A + // caller may still pass an even LARGER batch stride (extra inter-slab + // padding); the element-count / allocation below is sized from the resolved + // value so that layout is honoured too. + const std::int64_t bsa = + batch_stride_A > 0 ? batch_stride_A : (operand_is_row_major(kLayoutA) ? M : K) * sa; + const std::int64_t bsb = + batch_stride_B > 0 ? batch_stride_B : (operand_is_row_major(kLayoutB) ? K : N) * sb; + const std::int64_t bsc = + batch_stride_C > 0 ? batch_stride_C : (operand_is_row_major(kLayoutC) ? M : N) * sc; + + // Total element counts across all batches, sized from the RESOLVED batch + // stride so that a padded (non-packed) batch stride allocates enough for the + // last batch's slab plus its padding. + const std::int64_t a_elems = bsa * batch_count; + const std::int64_t b_elems = bsb * batch_count; + const std::int64_t c_elems = bsc * batch_count; + + const ADataType* A_host = static_cast(A); + const BDataType* B_host = static_cast(B); + CDataType* C_host = static_cast(C); + + ADataType* A_dev = nullptr; + BDataType* B_dev = nullptr; + CDataType* C_dev = nullptr; + + auto cleanup = [&]() { + if(A_dev) + (void)hipFree(A_dev); + if(B_dev) + (void)hipFree(B_dev); + if(C_dev) + (void)hipFree(C_dev); + }; + + // ck_tile::index_t is 32-bit: reject any dimension / stride / batch value + // that would not fit before the static_casts into BatchedGemmHostArgs below, + // so an out-of-range problem fails loudly instead of silently truncating. + { + auto fits_i32 = [](std::int64_t v) { return v <= static_cast(INT32_MAX); }; + if(!fits_i32(M) || !fits_i32(N) || !fits_i32(K) || !fits_i32(sa) || !fits_i32(sb) || + !fits_i32(sc) || !fits_i32(bsa) || !fits_i32(bsb) || !fits_i32(bsc) || + !fits_i32(batch_count) || !fits_i32(kbatch)) + { + std::cerr << "dispatcher_run_batched: a dimension/stride/batch value exceeds the " + "32-bit index range\n"; + if(time_ms) + { + *time_ms = -1.0f; + } + return -1; + } + } + + hipError_t err; + if((err = hipMalloc(&A_dev, a_elems * sizeof(ADataType))) != hipSuccess) + { + cleanup(); + return fail_hip("hipMalloc(A)", err); + } + if((err = hipMalloc(&B_dev, b_elems * sizeof(BDataType))) != hipSuccess) + { + cleanup(); + return fail_hip("hipMalloc(B)", err); + } + if((err = hipMalloc(&C_dev, c_elems * sizeof(CDataType))) != hipSuccess) + { + cleanup(); + return fail_hip("hipMalloc(C)", err); + } + + if((err = hipMemcpy(A_dev, A_host, a_elems * sizeof(ADataType), hipMemcpyHostToDevice)) != + hipSuccess) + { + cleanup(); + return fail_hip("hipMemcpy(A H2D)", err); + } + if((err = hipMemcpy(B_dev, B_host, b_elems * sizeof(BDataType), hipMemcpyHostToDevice)) != + hipSuccess) + { + cleanup(); + return fail_hip("hipMemcpy(B H2D)", err); + } + if((err = hipMemset(C_dev, 0, c_elems * sizeof(CDataType))) != hipSuccess) + { + cleanup(); + return fail_hip("hipMemset(C)", err); + } + + float exec_time = -1.0f; + try + { + // k_batch (split-K) is a real capability: Old-TE batched_gemm passes + // split_k_ through BatchedGemmHostArgs, so we thread it here too + // (default 1 == no split). + ck_tile::BatchedGemmHostArgs args{A_dev, + B_dev, + C_dev, + kbatch, + static_cast(M), + static_cast(N), + static_cast(K), + static_cast(sa), + static_cast(sb), + static_cast(sc), + static_cast(bsa), + static_cast(bsb), + static_cast(bsc), + static_cast(batch_count)}; + + // Build the stream_config to MATCH the Tile Engine batched_gemm profiler + // (batched_gemm_profiler.hpp: {nullptr, true, log, n_warmup, n_repeat, + // is_gpu_timer, flush_cache, rotating_count}). The struct field order is + // {stream_id, time_kernel, log_level, cold_niters, nrepeat, + // is_gpu_timer, flush_cache, rotating_count} (stream_config.hpp). The + // benchmark defaults are warmup=50 / repeat=100 / flush_cache=true / + // rotating_count=1000 / gpu_timer=true; callers thread the exact values + // through so both sides are driven identically. A caller passing <= 0 + // for warmup/repeat falls back to the TE benchmark defaults. + const int n_warmup = warmup > 0 ? static_cast(warmup) : 50; + const int n_repeat = repeat > 0 ? static_cast(repeat) : 100; + const bool do_flush = flush_cache != 0; + const int n_rotating = rotating_count > 0 ? static_cast(rotating_count) : 1; + const ck_tile::stream_config stream{nullptr, + /*time_kernel=*/true, + /*log_level=*/0, + /*cold_niters=*/n_warmup, + /*nrepeat=*/n_repeat, + /*is_gpu_timer=*/true, + /*flush_cache=*/do_flush, + /*rotating_count=*/n_rotating}; + exec_time = SelectedKernel::launch(args, stream); + } + catch(const std::exception& e) + { + std::cerr << "dispatcher_run_batched: kernel launch failed: " << e.what() << std::endl; + cleanup(); + if(time_ms) + { + *time_ms = -1.0f; + } + return -2; + } + catch(...) + { + std::cerr << "dispatcher_run_batched: kernel launch failed (unknown exception)" + << std::endl; + cleanup(); + if(time_ms) + { + *time_ms = -1.0f; + } + return -2; + } + + if((err = hipMemcpy(C_host, C_dev, c_elems * sizeof(CDataType), hipMemcpyDeviceToHost)) != + hipSuccess) + { + cleanup(); + return fail_hip("hipMemcpy(C D2H)", err); + } + + if(time_ms) + { + *time_ms = exec_time; + } + + cleanup(); + return 0; +} + +void dispatcher_cleanup() { g_initialized = false; } + +} // extern "C" diff --git a/projects/composablekernel/dispatcher/codegen/unified_gemm_codegen.py b/projects/composablekernel/dispatcher/codegen/unified_gemm_codegen.py index e93c212d491c..7dc3de7dbf5d 100755 --- a/projects/composablekernel/dispatcher/codegen/unified_gemm_codegen.py +++ b/projects/composablekernel/dispatcher/codegen/unified_gemm_codegen.py @@ -202,6 +202,7 @@ class GemmVariant(Enum): STANDARD = "standard" PRESHUFFLE = "preshuffle" MULTI_D = "multi_d" + BATCHED = "batched" MULTI_ABD = "multi_abd" GROUPED = "grouped" # Stream-K. COVERAGE LIMITATION: the dispatcher does NOT yet emit the full @@ -312,6 +313,10 @@ def key_name(self, datatype: str, layout: str) -> str: if self.preshuffle: parts.append("preshuffle") + # Batched variant + if self.variant == GemmVariant.BATCHED: + parts.append("batched") + # Multi-D variant: include elementwise op, num tensors, and D layout if self.variant == GemmVariant.MULTI_D: parts.append(f"ew_{self.elementwise_op}") @@ -391,6 +396,8 @@ def generate(config: KernelConfig, datatype: str, layout: str) -> str: name += "_preshuffle" elif config.variant == GemmVariant.MULTI_D: name += f"_multid_{config.elementwise_op}_d{config.num_d_tensors}" + elif config.variant == GemmVariant.BATCHED: + name += "_batched" elif config.variant == GemmVariant.MULTI_ABD: # Encode every behavior-affecting multi-abd parameter so distinct # tensor counts / elementwise ops get distinct kernel names. @@ -478,6 +485,11 @@ def _header(self, kernel_name: str, config: KernelConfig) -> str: includes += """ #include "ck_tile/ops/gemm/pipeline/wp_pipeline_agmem_bgmem_creg_v2.hpp" #include "ck_tile/ops/gemm/pipeline/wp_pipeline_agmem_bgmem_creg_base_policy.hpp" +""" + + if config.variant == GemmVariant.BATCHED: + includes += """ +#include "ck_tile/ops/gemm/kernel/batched_gemm_kernel.hpp" """ if config.variant == GemmVariant.STREAM_K: @@ -764,6 +776,7 @@ def _selected_kernel_struct(self, config: KernelConfig, kernel_name: str) -> str #define GEMM_KEY_PRESHUFFLE {int(config.preshuffle)} #define GEMM_KEY_TRANSPOSE_C 0 #define GEMM_KEY_GROUPED 0 +#define GEMM_KEY_BATCHED {int(config.variant == GemmVariant.BATCHED)} #define GEMM_KEY_SPLIT_K 1 {self._multi_d_single_include(config)}#endif // CK_TILE_SINGLE_KERNEL_INCLUDE """ @@ -792,6 +805,8 @@ def _launch_function(self, config: KernelConfig) -> str: return self._launch_function_multi_abd(config) if config.variant == GemmVariant.MULTI_D: return self._launch_function_multi_d(config) + if config.variant == GemmVariant.BATCHED: + return self._launch_function_batched(config) if config.variant == GemmVariant.GROUPED: return self._launch_function_grouped(config) if config.variant == GemmVariant.STREAM_K: @@ -848,6 +863,90 @@ def _launch_function_standard(self, config: KernelConfig) -> str: return ave_time; }}""" + def _launch_function_batched(self, config: KernelConfig) -> str: + """Generate launch function for batched GEMM. + + Mirrors the Tile Engine batched_gemm instance + (gemm_instance_builder.py, kernel_name_prefix == "batched_gemm"): + * takes ``ck_tile::BatchedGemmHostArgs`` (adds batch_count + per-batch + strides on top of the single-problem args); + * uses the batched-specific ``TileGemmUniversalTraits`` (no persistent/ + preshuffle/wave-group template tail); + * builds ``ck_tile::BatchedGemmKernel`` and sizes the grid with the + 4-arg ``GridSize(M, N, k_batch, batch_count)``. + Unlike the standard path there is no hot-loop TailHandler branch -- the + batched kernel launches directly, matching Old-TE. + """ + return f""" + static float launch(const ck_tile::BatchedGemmHostArgs& args, const stream_config& stream) {{ + float ave_time{{0}}; + + constexpr auto scheduler = {self.tm.SCHEDULER_TO_CK[config.trait.scheduler]}; + + using GemmUniversalTraits = ck_tile::TileGemmUniversalTraits< + kPadM, kPadN, kPadK, DoubleSmemBuffer, + ALayout, BLayout, CLayout, TransposeC>; + + using UniversalGemmProblem = UniversalGemmPipelineProblem< + ADataType, BDataType, AccDataType, TileShape, + GemmUniversalTraits, + scheduler>; + + using GemmPipeline = {self.tm.PIPELINE_TO_CK[config.trait.pipeline]}; + {self._epilogue_code(config)} + + using GemmKernel = ck_tile::BatchedGemmKernel; + + auto kargs = GemmKernel::MakeKernelArgs(args); + + if (!GemmKernel::IsSupportedArgument(kargs)) {{ + throw std::runtime_error("Arguments not supported for batched gemm kernel!"); + }} + + const dim3 grids = GemmKernel::GridSize(args.M, args.N, args.k_batch, args.batch_count); + const dim3 blocks = GemmKernel::BlockSize(); + + constexpr int kBlockPerCu = {config.k_block_per_cu}; + if(args.k_batch > 1) {{ + // Split-K path: the universal GEMM epilogue stores C with + // memory_operation_enum::atomic_add when k_batch > 1, so C must be + // zero at the start of EVERY kernel invocation. The benchmarking + // timing loop re-runs the kernel cold_niters + nrepeat times, so a + // single pre-launch memset is not enough -- C accumulates across + // iterations. Pass a hipMemsetAsync callable as the first callable + // to launch_kernel (the documented ck_tile pattern, see + // kernel_launch.hpp) so C is re-zeroed before each timed launch. + // Use hipMemsetAsync on the launch stream (sc.stream_id_): the + // synchronous hipMemset ignores stream_config and would serialise + // onto the default stream, skewing timing/correctness whenever a + // non-default stream is used. Only the batch slices actually + // written (M*N per batch) need clearing; use the batch stride to + // cover padded layouts. + const std::size_t c_bytes = + static_cast(args.batch_stride_E) * + static_cast(args.batch_count) * sizeof(CDataType); + ave_time = launch_kernel(stream, + [=](const stream_config& sc) {{ + // Check the return: a failed reset would let the split-K + // atomic-add epilogue accumulate into a non-zero C and + // silently produce wrong results (mirrors the stream-K path). + if(hipMemsetAsync(args.e_ptr, 0, c_bytes, sc.stream_id_) != hipSuccess) {{ + throw std::runtime_error( + "batched split-K: hipMemsetAsync failed to reset C between iterations"); + }} + }}, + make_kernel(GemmKernel{{}}, grids, blocks, 0, kargs)); + }} else {{ + // k_batch == 1 (the Old-TE parity default): byte-identical to the + // Tile Engine batched_gemm launch -- single kernel callable, no + // memset (the epilogue uses memory_operation_enum::set). + ave_time = launch_kernel(stream, + make_kernel(GemmKernel{{}}, grids, blocks, 0, kargs)); + }} + + return ave_time; + }}""" + def _launch_function_grouped(self, config: KernelConfig) -> str: """Generate launch function for grouped GEMM. @@ -1242,6 +1341,30 @@ def _launch_function_streamk(self, config: KernelConfig) -> str: def _epilogue_code(self, config: KernelConfig) -> str: """Generate epilogue code""" + if config.variant == GemmVariant.BATCHED: + # Respect the requested epilogue: the batched sweep space includes + # BOTH "cshuffle" and "default" (see the shipped default_config.json + # trait_config.epilogue). Emitting CShuffle unconditionally would + # silently build a cshuffle kernel under a "_epi_default" name and + # skip the cshuffle-store correctness gate (_cshuffle_repeat_ok) for + # configs that requested "default" -- a name/impl mismatch that also + # breaks the config->codegen->runtime byte-parity invariant. + if config.trait.epilogue == "cshuffle": + return """ + using EpilogueProblem = CShuffleEpilogueProblem< + ADataType, BDataType, tuple<>, AccDataType, CDataType, + tuple<>, CLayout, element_wise::PassThrough, + TilePartitioner::MPerBlock, TilePartitioner::NPerBlock, + WarpPerBlock_M, WarpPerBlock_N, WarpTileM, WarpTileN, WarpTileK, + UniversalGemmProblem::TransposeC>; + using GemmEpilogue = CShuffleEpilogue;""" + return """ + using EpilogueProblem = DefaultGemm2DEpilogueProblem< + ADataType, BDataType, tuple<>, AccDataType, CDataType, + tuple<>, CLayout, element_wise::PassThrough, + TilePartitioner::MPerBlock, TilePartitioner::NPerBlock, + kPadM, kPadN, WarpTileM, WarpTileN, WarpTileK, TransposeC>; + using GemmEpilogue = DefaultGemm2DEpilogue;""" if config.variant == GemmVariant.MULTI_ABD: # Multi-ABD epilogue: tuple A/B/D dtypes and D layouts, EDataType as # output, and the CDE element-wise function. Matches the TE builder's @@ -1595,28 +1718,36 @@ def _get_preselected_configs(self) -> List[KernelConfig]: @staticmethod def _cshuffle_repeat_ok(tile: TileConfig) -> bool: - """CShuffle-store correctness gate. - - The CShuffle epilogue stores the accumulator back through LDS in - power-of-two MRepeat/NRepeat chunks, so a tile whose per-wave repeat - count -- tile / (warp * warp_tile) -- is not a power of two is - mis-stored and yields numerically WRONG results at runtime. The kernel - still compiles (the epilogue's static_asserts only check divisibility, - which such tiles satisfy), so it must be filtered in codegen. Observed - on MI350 for tile_m=192 (MRepeat = 192 / (2*32) = 3): verified incorrect - on BOTH the bridge and Tile Engine at every shape, including shapes - divisible by 192. Power-of-two tiles (64/128/256) are unaffected. + """CShuffle-store correctness gate (narrowed; see issue #9684). + + The CShuffle epilogue mis-stores the accumulator only for ONE specific + combination: an ODD per-wave repeat (>1) -- tile / (warp * warp_tile) -- + paired with a 32-wide warp tile in that dimension. GPU-verified on + gfx942: the tile_m=192 / warp_m=2 / warp_tile_m=32 configs + (MRepeat = 192/(2*32) = 3) return garbage (built-in verify=fail, + max_rel ~1.3), while EVERY other non-power-of-two repeat is numerically + correct -- including MRepeat=3 with warp_tile_m=16 (192/(4*16)) and even + non-pow2 repeats such as 6 and 12. Confirmed by building all 116 + non-pow2-repeat cshuffle/192 configs and running the op's CPU validation: + exactly the 26 matching the signature above fail; the other 90 pass. The + earlier "per-wave repeat must be a power of two" rule was therefore too + broad and needlessly dropped those 90 valid configs. This is CShuffle-specific: the "default" (DefaultGemm2DEpilogue) path stores directly (not through the LDS repack) and is numerically correct - for non-pow2 repeats -- verified on gfx942 at tile_m=192/MRepeat=3 - (max_rel ~5e-4 across shapes divisible by 192, while the same tile under - CShuffle returns garbage, max_rel ~1.3). Only call this for kernels - whose resolved epilogue is "cshuffle". + for non-pow2 repeats. Only call this for kernels whose resolved epilogue + is "cshuffle". """ + + def _dim_bad(repeat: int, warp_tile: int) -> bool: + return repeat > 1 and repeat % 2 == 1 and warp_tile == 32 + m_repeat = tile.tile_m // (tile.warp_m * tile.warp_tile_m) n_repeat = tile.tile_n // (tile.warp_n * tile.warp_tile_n) - return _is_power_of_two(m_repeat) and _is_power_of_two(n_repeat) + return not ( + _dim_bad(m_repeat, tile.warp_tile_m) + or _dim_bad(n_repeat, tile.warp_tile_n) + ) def _get_configs_for_variant(self, variant: GemmVariant) -> List[KernelConfig]: """Get all configurations for a variant @@ -1645,11 +1776,12 @@ def _get_configs_for_variant(self, variant: GemmVariant) -> List[KernelConfig]: ): continue - if variant == GemmVariant.STANDARD: + if variant in (GemmVariant.STANDARD, GemmVariant.BATCHED): # CShuffle-store correctness gate: skip non-pow2 repeat tiles # only for the cshuffle epilogue (see _cshuffle_repeat_ok). The # "default" epilogue is correct with non-pow2 repeats, so it is - # NOT gated here. + # NOT gated here. Batched GEMM shares the standard tile/trait + # sweep -- it only differs in the launch/kernel type. if trait.epilogue == "cshuffle" and not self._cshuffle_repeat_ok(tile): continue configs.append(KernelConfig(tile=tile, trait=trait, variant=variant)) @@ -2122,6 +2254,7 @@ def main(): "preshuffle", "multi_d", "stream_k", + "batched", "multi_abd", "grouped", ], diff --git a/projects/composablekernel/dispatcher/python/batched_gemm_utils.py b/projects/composablekernel/dispatcher/python/batched_gemm_utils.py new file mode 100644 index 000000000000..49f8a126cd6e --- /dev/null +++ b/projects/composablekernel/dispatcher/python/batched_gemm_utils.py @@ -0,0 +1,768 @@ +# Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +# SPDX-License-Identifier: MIT +""" +Batched GEMM Tile Engine <-> Dispatcher bridge. + +Batched counterpart of ``gemm_utils.py``. Batched GEMM adds a leading batch +dimension with per-batch strides, giving it a DIVERGENT ABI from the +single-problem GEMM library: the ctypes entry point +``dispatcher_run_batched`` carries ``batch_count``, ``k_batch`` (split-K), the +three batch strides, and the benchmarking knobs (warmup/repeat/flush_cache/ +rotating_count) in addition to M/N/K, and the .so launches the force-included +kernel directly via ``SelectedKernel::launch(BatchedGemmHostArgs{...}, ...)`` +(registry bypass). The benchmarking knobs let the bridge build the SAME +``stream_config`` the Tile Engine batched_gemm profiler uses, so A/B timing is +fair; ``k_batch`` mirrors Old-TE's split-K support. + +Public surface (mirrors gemm_utils): + + BatchedGemmKernelConfig -- shared contract dataclass (variant="batched") + .name -- registry/runtime lookup key (byte-exact vs the + codegen kernel header stem) + .to_codegen_json() -- feeds unified_gemm_codegen.py + BatchedGemmProblem -- a single (batch, M, N, K) problem + setup_multiple_batched_gemm_dispatchers -- codegen + hipcc -> .so (NO GPU) + BatchedGemmDispatcherLib -- thin ctypes ABI wrapper + GpuBatchedGemmRunner -- GPU run + time (from a .so path) + expand_sweep -- TE JSON sweep config -> [BatchedGemmKernelConfig] + +The build half reuses gemm_utils' proven codegen/compile flag machinery +(``_tile_engine_codegen_flags`` etc.) so both bridges compile kernels with +byte-identical AMDGPU backend flags -- a prerequisite for fair A/B parity. +""" + +from __future__ import annotations + +import ctypes +import multiprocessing +from concurrent.futures import ProcessPoolExecutor, as_completed +from dataclasses import dataclass +from pathlib import Path +from typing import Any, Dict, List, Optional, Tuple + +import numpy as np + +import ctypes_utils as _cu +import gemm_utils as _gu + + +# ============================================================================ +# GPU architecture resolution (never default to gfx942) +# ============================================================================ + +_SUPPORTED_ARCHES: Tuple[str, ...] = ("gfx90a", "gfx942", "gfx950") + +# Byte size of each C output dtype as the compiled kernel writes it +# (sizeof(CDataType)). The host numpy buffer is memcpy'd verbatim to/from the +# device, so its element size must match these exactly -- see the F6 check in +# GpuBatchedGemmRunner.run(). ck_tile bf16_t is a 2-byte type (mirrored on the +# host by np.uint16); fp16 is 2 bytes; the int8 path accumulates into int32. +_C_SIZEOF: Dict[str, int] = {"fp16": 2, "bf16": 2, "int32": 4} + + +def _get_arch() -> str: + """Detect GPU arch via rocminfo and validate; raise on failure. Never defaults.""" + import subprocess + + arch = "" + try: + out = subprocess.check_output(["rocminfo"], text=True, stderr=subprocess.DEVNULL) + for line in out.splitlines(): + if "Name:" in line and "gfx" in line: + arch = line.split()[-1].strip() + break + except Exception: + arch = "" + if not arch: + raise RuntimeError( + "Could not detect GPU architecture from rocminfo; refusing to default " + "to gfx942. Pass gfx_arch explicitly." + ) + if arch not in _SUPPORTED_ARCHES: + raise ValueError( + f"Unsupported GPU architecture {arch!r}; supported: {list(_SUPPORTED_ARCHES)}" + ) + return arch + + +def _resolve_arch(arch: Optional[str]) -> str: + """Resolve an explicit arch (validated) or auto-detect via rocminfo.""" + if arch is None: + return _get_arch() + if arch not in _SUPPORTED_ARCHES: + raise ValueError( + f"Unsupported GPU architecture {arch!r}; supported: {list(_SUPPORTED_ARCHES)}" + ) + return arch + + +# ============================================================================ +# The shared contract: BatchedGemmKernelConfig +# ============================================================================ + + +@dataclass +class BatchedGemmKernelConfig(_gu.GemmKernelConfig): + """GEMM kernel config specialised to the batched variant. + + Inherits every field/property from GemmKernelConfig and only pins the + variant to ``"batched"`` so ``.name`` gains the ``_batched`` suffix that the + codegen (unified_gemm_codegen.py::KernelNaming.generate) also appends -- the + single thread tying config -> codegen -> runtime name together. + + ``gfx_arch`` is overridden to default to ``None`` (rather than the parent's + hardcoded ``"gfx942"``) so the batched bridge never silently builds for the + wrong GPU: the build path resolves ``None`` via ``_get_arch()`` (rocminfo) + and raises on an unsupported / undetectable arch. + """ + + variant: str = "batched" + gfx_arch: Optional[str] = None # type: ignore[assignment] + + +# Extend GemmKernelConfig.name handling for the batched variant. The parent +# already appends "_streamk"/"_preshuffle" for those variants; "batched" needs +# the same treatment, so we patch the property lookup by overriding name here. +def _batched_name(self: BatchedGemmKernelConfig) -> str: # pragma: no cover - thin + base = _gu.GemmKernelConfig.name.fget(self) # type: ignore[attr-defined] + if self.variant == "batched" and not base.endswith("_batched"): + return base + "_batched" + return base + + +BatchedGemmKernelConfig.name = property(_batched_name) # type: ignore[assignment] + + +# ============================================================================ +# Problem / result +# ============================================================================ + + +@dataclass +class BatchedGemmProblem: + """A batched GEMM problem: C[b,M,N] = A[b,M,K] @ B[b,K,N] for b in [0,batch).""" + + batch_count: int + M: int + N: int + K: int + + @property + def flops(self) -> float: + return 2.0 * self.batch_count * self.M * self.N * self.K + + def to_dict(self) -> Dict[str, int]: + return { + "batch_count": self.batch_count, + "M": self.M, + "N": self.N, + "K": self.K, + } + + @classmethod + def from_dict(cls, d: Dict[str, int]) -> "BatchedGemmProblem": + return cls( + batch_count=int(d["batch_count"]), + M=int(d["M"]), + N=int(d["N"]), + K=int(d["K"]), + ) + + +@dataclass +class BatchedGemmResult: + output: np.ndarray + time_ms: float + status: int + tflops: float + kernel_name: str + + @property + def success(self) -> bool: + return self.status == 0 + + +# ============================================================================ +# ctypes ABI wrapper (divergent from single-problem GEMM) +# ============================================================================ + + +class BatchedGemmDispatcherLib: + """Thin ctypes wrapper around a compiled batched GEMM dispatcher .so. + + The batched .so exposes exactly one kernel (the force-included header), so + ``kernel_names`` returns a single-element list. + """ + + def __init__(self, so_path: Path): + self._path = Path(so_path) + self._lib = ctypes.CDLL(str(self._path)) + self._setup_functions() + + def _setup_functions(self) -> None: + lib = self._lib + + lib.dispatcher_initialize.argtypes = [] + lib.dispatcher_initialize.restype = ctypes.c_int + + lib.dispatcher_get_kernel_count.argtypes = [] + lib.dispatcher_get_kernel_count.restype = ctypes.c_int + + lib.dispatcher_get_kernel_name.argtypes = [] + lib.dispatcher_get_kernel_name.restype = ctypes.c_char_p + + lib.dispatcher_get_kernel_name_at.argtypes = [ + ctypes.c_int, + ctypes.c_char_p, + ctypes.c_int, + ] + lib.dispatcher_get_kernel_name_at.restype = ctypes.c_int + + # Divergent ABI: batch_count + k_batch (split-K) + three batch strides + # on top of M/N/K, plus the benchmarking knobs (warmup/repeat/ + # flush_cache/rotating_count) so the .so drives stream_config identically + # to the Tile Engine batched_gemm profiler for fair A/B timing. + lib.dispatcher_run_batched.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.c_int64, # batch_count + ctypes.c_int64, # k_batch (split-K; 0/1 -> no split) + ctypes.c_int64, # stride_A + ctypes.c_int64, # stride_B + ctypes.c_int64, # stride_C + ctypes.c_int64, # batch_stride_A + ctypes.c_int64, # batch_stride_B + ctypes.c_int64, # batch_stride_C + ctypes.c_int64, # warmup (<=0 -> TE default 50) + ctypes.c_int64, # repeat (<=0 -> TE default 100) + ctypes.c_int64, # flush_cache (0/1) + ctypes.c_int64, # rotating_count (<=0 -> 1) + ctypes.POINTER(ctypes.c_float), # time_ms + ] + lib.dispatcher_run_batched.restype = ctypes.c_int + + lib.dispatcher_cleanup.argtypes = [] + lib.dispatcher_cleanup.restype = None + + @property + def path(self) -> Path: + return self._path + + def initialize(self) -> bool: + return self._lib.dispatcher_initialize() == 0 + + def get_kernel_count(self) -> int: + return int(self._lib.dispatcher_get_kernel_count()) + + @property + def kernel_names(self) -> List[str]: + buf = ctypes.create_string_buffer(256) + if self._lib.dispatcher_get_kernel_name_at(0, buf, 256) == 0: + return [buf.value.decode("utf-8")] + raw = self._lib.dispatcher_get_kernel_name() + return [raw.decode("utf-8")] if raw else [] + + def run( + self, + A: np.ndarray, + B: np.ndarray, + C: np.ndarray, + M: int, + N: int, + K: int, + batch_count: int, + k_batch: int = 1, + stride_A: int = 0, + stride_B: int = 0, + stride_C: int = 0, + batch_stride_A: int = 0, + batch_stride_B: int = 0, + batch_stride_C: int = 0, + warmup: int = 50, + repeat: int = 100, + flush_cache: bool = True, + rotating_count: int = 1000, + ) -> Tuple[int, float]: + time_ms = ctypes.c_float(0.0) + status = self._lib.dispatcher_run_batched( + A.ctypes.data_as(ctypes.c_void_p), + B.ctypes.data_as(ctypes.c_void_p), + C.ctypes.data_as(ctypes.c_void_p), + M, + N, + K, + batch_count, + k_batch, + stride_A, + stride_B, + stride_C, + batch_stride_A, + batch_stride_B, + batch_stride_C, + warmup, + repeat, + 1 if flush_cache else 0, + rotating_count, + ctypes.byref(time_ms), + ) + return status, time_ms.value + + def cleanup(self) -> None: + self._lib.dispatcher_cleanup() + + +# ============================================================================ +# GPU runner +# ============================================================================ + + +class GpuBatchedGemmRunner: + """High-level runner: construct from a .so path, call run(A, B, problem). + + A/B/C are batched tensors laid out per the compiled kernel's layout. The C + ABI takes HOST pointers and manages GPU memory internally, so this runner + hands numpy arrays (in the kernel's element dtype + memory order) straight + to the .so. fp16/rcr is the only TE-supported batched signature today, but + the dtype-encode helpers are reused from gemm_utils so extending to more + dtypes only requires codegen support. + """ + + def __init__(self, lib_path: Path): + self.lib = BatchedGemmDispatcherLib(lib_path) + if not self.lib.initialize(): + raise RuntimeError(f"Failed to initialize batched dispatcher .so: {lib_path}") + names = self.lib.kernel_names + self._kernel_name = names[0] if names else "unknown" + + @property + def kernel_name(self) -> str: + return self._kernel_name + + def run( + self, + A: np.ndarray, + B: np.ndarray, + problem: BatchedGemmProblem, + k_batch: int = 1, + batch_pad_A: int = 0, + batch_pad_B: int = 0, + batch_pad_C: int = 0, + warmup: int = 50, + repeat: int = 100, + flush_cache: bool = True, + rotating_count: int = 1000, + ) -> BatchedGemmResult: + """Run a batched GEMM. + + ``k_batch`` selects split-K (default 1 == no split, matching Old-TE). + + ``batch_pad_{A,B,C}`` add padding elements between consecutive per-batch + slabs so the batch stride is NON-packed + (batch_stride = packed_slab + pad). When any pad is > 0 the host buffer + is over-allocated to slab*batch + pad*(batch) and each logical slab is + scattered into its padded slot, exercising the batch_stride path in the + .so (B2). With all pads 0 this is the packed case (batch stride 0). + + ``warmup/repeat/flush_cache/rotating_count`` are threaded into the .so's + stream_config so timing matches the Tile Engine batched_gemm profiler + (B1). + """ + batch, M, N, K = problem.batch_count, problem.M, problem.N, problem.K + + dtype = _gu._dtype_from_kernel_name(self._kernel_name) + la, lb, lc = _gu._layout_from_kernel_name(self._kernel_name) + + # Per-batch layout transform: a 'c' (column-major) operand is stored as + # the transpose of its logical (rows, cols) so the contiguous per-batch + # slab matches column-major memory order (leading batch axis untouched). + A_lay = A if la == "r" else np.transpose(A, (0, 2, 1)) + B_lay = B if lb == "r" else np.transpose(B, (0, 2, 1)) + C_shape = (batch, M, N) if lc == "r" else (batch, N, M) + + if dtype == "bf16": + A_h = _gu._fp32_to_bf16_u16(A_lay) + B_h = _gu._fp32_to_bf16_u16(B_lay) + elif dtype == "fp8": + A_h = _gu._fp32_to_fp8_u8(A_lay) + B_h = _gu._fp32_to_fp8_u8(B_lay) + elif dtype == "bf8": + A_h = _gu._fp32_to_bf8_u8(A_lay) + B_h = _gu._fp32_to_bf8_u8(B_lay) + elif dtype == "int8": + A_h = np.ascontiguousarray(A_lay, dtype=np.int8) + B_h = np.ascontiguousarray(B_lay, dtype=np.int8) + else: # fp16 (default / only TE batched dtype) + A_h = np.ascontiguousarray(A_lay, dtype=np.float16) + B_h = np.ascontiguousarray(B_lay, dtype=np.float16) + + out_dtype = _gu._output_dtype(dtype) + _C_NP = {"fp16": np.float16, "bf16": np.uint16, "int32": np.int32} + if out_dtype not in _C_NP: + raise ValueError( + f"unsupported C dtype {out_dtype!r} (from input dtype {dtype!r}); " + "add it to _C_NP so the host buffer matches sizeof(CDataType)" + ) + # F6: the host C buffer is memcpy'd byte-for-byte to/from the device + # buffer the kernel writes as CDataType, so the numpy element size MUST + # equal the C++ sizeof(CDataType). Assert it here so a future _C_NP edit + # (e.g. mapping fp16 -> np.float32) fails loudly instead of silently + # copying the wrong byte count. + if np.dtype(_C_NP[out_dtype]).itemsize != _C_SIZEOF[out_dtype]: + raise ValueError( + f"host C dtype size mismatch for {out_dtype!r}: numpy " + f"{np.dtype(_C_NP[out_dtype]).itemsize} bytes != kernel " + f"sizeof(CDataType) {_C_SIZEOF[out_dtype]} bytes" + ) + + # Packed per-batch slab element counts (row-major flattened). + slab_A = M * K + slab_B = K * N + slab_C = M * N + + if batch_pad_A or batch_pad_B or batch_pad_C: + # NON-PACKED (B2): interleave a padding gap between per-batch slabs. + # Batch stride = packed slab + pad; host buffer is over-allocated to + # batch_stride * batch and each slab is copied into its slot. The .so + # sizes its device allocation from the same batch stride. + bsa = slab_A + int(batch_pad_A) + bsb = slab_B + int(batch_pad_B) + bsc = slab_C + int(batch_pad_C) + + A_flat = A_h.reshape(batch, slab_A) + B_flat = B_h.reshape(batch, slab_B) + A_buf = np.zeros(bsa * batch, dtype=A_h.dtype) + B_buf = np.zeros(bsb * batch, dtype=B_h.dtype) + for b in range(batch): + A_buf[b * bsa : b * bsa + slab_A] = A_flat[b] + B_buf[b * bsb : b * bsb + slab_B] = B_flat[b] + C_buf = np.zeros(bsc * batch, dtype=_C_NP[out_dtype]) + + status, time_ms = self.lib.run( + A_buf, + B_buf, + C_buf, + M, + N, + K, + batch, + k_batch, + 0, + 0, + 0, + bsa, + bsb, + bsc, + warmup, + repeat, + flush_cache, + rotating_count, + ) + # Gather the C slabs back out of their padded slots. + C_h = np.empty((batch, slab_C), dtype=_C_NP[out_dtype]) + for b in range(batch): + C_h[b] = C_buf[b * bsc : b * bsc + slab_C] + C_h = C_h.reshape(C_shape) + else: + # PACKED: batch strides default (0) -> kernel/.so derive them. + C_h = np.zeros(C_shape, dtype=_C_NP[out_dtype]) + status, time_ms = self.lib.run( + A_h, + B_h, + C_h, + M, + N, + K, + batch, + k_batch, + warmup=warmup, + repeat=repeat, + flush_cache=flush_cache, + rotating_count=rotating_count, + ) + + if out_dtype == "bf16": + C_dec = _gu._bf16_u16_to_fp32(C_h) + else: + C_dec = C_h + C_out = C_dec if lc == "r" else np.transpose(C_dec, (0, 2, 1)) + + tflops = (problem.flops / (time_ms * 1e-3)) / 1e12 if time_ms > 0 else 0.0 + return BatchedGemmResult( + output=C_out, + time_ms=time_ms, + status=status, + tflops=tflops, + kernel_name=self._kernel_name, + ) + + +# ============================================================================ +# Build API: codegen + hipcc -> .so paths (no GPU) +# ============================================================================ + + +def _build_batched_compile_jobs( + config: BatchedGemmKernelConfig, header: Path +) -> Tuple[Dict[str, Any], Path]: + """Compile+link commands for the batched ctypes lib. Identical to the + single-problem job (same flags, same includes) except it force-links the + batched ctypes source (divergent ABI) instead of gemm_ctypes_lib.cpp.""" + root = _cu.get_dispatcher_root() + ck_root = root.parent + build_dir = _cu.get_build_dir() + output_dir = _cu.get_generated_kernels_dir() + ctypes_source = root / "bindings" / "ctypes" / "batched_gemm_ctypes_lib.cpp" + + lib_path = build_dir / "examples" / f"lib{config.name}.so" + obj_file = lib_path.with_suffix(".o") + + # Never default to gfx942: resolve None via rocminfo (_get_arch) and validate + # an explicit arch. GFX_ARCH is threaded to the .cpp as a hard requirement + # (the source #errors if it is missing). + gfx_arch = _resolve_arch(config.gfx_arch) + + compile_cmd = [ + _gu._resolve_hipcc(), + "-c", + "-fPIC", + "-O3", + f"-I{root / 'include'}", + f"-I{ck_root / 'include'}", + f"-I{ck_root}", + f"-I{str(output_dir)}", + "-DCK_TILE_SINGLE_KERNEL_INCLUDE", + f"-include{header}", + "-D__HIP_PLATFORM_AMD__", + f"--offload-arch={gfx_arch}", + f'-DGFX_ARCH="{gfx_arch}"', + # Byte-identical AMDGPU backend flags to the single-problem bridge and + # Old-TE (see gemm_utils._tile_engine_codegen_flags) -- required for a + # fair A/B parity comparison. + *_gu._tile_engine_codegen_flags(), + "-Wno-undefined-func-template", + "-Wno-float-equal", + str(ctypes_source), + "-o", + str(obj_file), + ] + # The batched ctypes lib bypasses the registry (it launches the kernel + # directly), so unlike the single-problem lib it does NOT need to link the + # dispatcher static archive -- only the kernel + ck_tile headers. + link_cmd = [ + _gu._resolve_hipcc(), + "-shared", + "-fPIC", + f"--offload-arch={gfx_arch}", + "--hip-link", + str(obj_file), + "-o", + str(lib_path), + ] + job = {"compile_cmd": compile_cmd, "link_cmd": link_cmd, "lib_path": str(lib_path)} + return job, lib_path + + +def setup_multiple_batched_gemm_dispatchers( + configs: List[BatchedGemmKernelConfig], + verbose: bool = True, + max_workers: Optional[int] = None, +) -> List[Optional[Path]]: + """Codegen + compile each batched config into its own .so. Returns .so paths + aligned to input order (None for configs that failed). Pure CPU (no GPU).""" + import sys + + n = len(configs) + results: List[Optional[Path]] = [None] * n + if n == 0: + return results + + max_workers = max_workers or min(multiprocessing.cpu_count(), 8) + + first_index: Dict[str, int] = {} + unique: List[int] = [] + for i, c in enumerate(configs): + key = c.name + if key not in first_index: + first_index[key] = i + unique.append(i) + + codegen_script = _cu.get_codegen_path() + output_dir = _cu.get_generated_kernels_dir() + ctypes_source = ( + _cu.get_dispatcher_root() + / "bindings" + / "ctypes" + / "batched_gemm_ctypes_lib.cpp" + ) + if not ctypes_source.exists(): + raise FileNotFoundError( + f"Missing batched ctypes source required for compilation:\n {ctypes_source}" + ) + + # -- Step 1: parallel codegen (one header per unique config) -------------- + codegen_args = [] + for i in unique: + c = configs[i] + # Resolve arch here too so codegen and compile target the same detected + # (or explicitly-validated) GPU -- never a silent gfx942 default. + gpu_target = _resolve_arch(c.gfx_arch) + codegen_args.append( + { + "index": i, + "python": sys.executable, + "codegen_script": str(codegen_script), + "output_dir": str(output_dir), + "dtype": c.dtype_a, + "layout": c.layout, + "gpu_target": gpu_target, + "tile_config_json": c.to_codegen_json(), + "hpp_glob_pattern": f"{c.name}.hpp", + # Thread the batched variant so the emitted header is the batched + # kernel and its name carries the _batched suffix (matching the + # glob pattern above). + "variant": c.variant, + } + ) + + if verbose: + print( + f"[batched-gemm-bridge] codegen: {len(codegen_args)} headers " + f"(workers={max_workers})..." + ) + + headers: Dict[int, Path] = {} + with ProcessPoolExecutor(max_workers=max_workers) as ex: + futs = { + ex.submit(_cu._generate_single_kernel_subprocess, a): a["index"] + for a in codegen_args + } + for fut in as_completed(futs): + i = futs[fut] + ok, hdr, err = fut.result() + if ok and hdr: + headers[i] = Path(hdr) + if verbose: + print(f" OK codegen [{i}] {configs[i].name}") + elif verbose: + print(f" FAIL codegen [{i}] {configs[i].name}: {err}") + + # -- Step 2: parallel compile + link -------------------------------------- + compile_jobs = [] + job_index: List[int] = [] + for i in unique: + hdr = headers.get(i) + if hdr is None: + continue + job, _ = _build_batched_compile_jobs(configs[i], hdr) + compile_jobs.append(job) + job_index.append(i) + + if verbose and compile_jobs: + print( + f"[batched-gemm-bridge] compile: {len(compile_jobs)} .so " + f"(workers={max_workers})..." + ) + + with ProcessPoolExecutor(max_workers=max_workers) as ex: + futs = { + ex.submit(_cu._run_hipcc_subprocess, job): job_index[j] + for j, job in enumerate(compile_jobs) + } + for fut in as_completed(futs): + i = futs[fut] + ok, lp, err = fut.result() + if ok and lp: + results[i] = Path(lp) + if verbose: + print(f" OK compile [{i}] {Path(lp).name}") + elif verbose: + print(f" FAIL compile [{i}] {configs[i].name}: {err}") + + for i, c in enumerate(configs): + if results[i] is None: + results[i] = results[first_index[c.name]] + + if verbose: + ok_count = sum(1 for r in results if r is not None) + print(f"[batched-gemm-bridge] setup complete: {ok_count}/{n} configs -> .so") + + return results + + +# ============================================================================ +# TE sweep config expansion +# ============================================================================ + + +def expand_sweep( + config_path: str, + arch: Optional[str] = None, + dtype: str = "fp16", + layout: str = "rcr", +) -> List[BatchedGemmKernelConfig]: + """Expand a Tile Engine batched GEMM JSON sweep config into + [BatchedGemmKernelConfig]. Reuses gemm_utils.expand_sweep (same tile/trait + sweep + validity gates) and re-stamps each result as the batched variant. + + ``arch`` is resolved via ``_get_arch()`` (rocminfo) when omitted and + validated against the supported set otherwise -- never a silent gfx942 + default.""" + arch = _resolve_arch(arch) + # Match Old-TE's validated set EXACTLY: the batched_gemm instance builder + # (tile_engine/ops/gemm/batched_gemm/batched_gemm_instance_builder.py) + # declares --datatype choices=["fp16"] and --layout choices=["rcr"], so + # anything else would codegen/compile/launch a kernel Old-TE never validated + # (claimed parity != exercised parity). Reject it up front with a clear error + # rather than silently building an untested signature. + if dtype != "fp16": + raise ValueError( + f"batched_gemm bridge supports only dtype 'fp16' (Old-TE " + f"batched_gemm_instance_builder declares --datatype choices=['fp16']); " + f"got {dtype!r}" + ) + if layout != "rcr": + raise ValueError( + f"batched_gemm bridge supports only layout 'rcr' (Old-TE " + f"batched_gemm_instance_builder declares --layout choices=['rcr']); " + f"got {layout!r}" + ) + base_configs = _gu.expand_sweep(config_path, arch, dtype=dtype, layout=layout) + out: List[BatchedGemmKernelConfig] = [] + seen: set = set() + for b in base_configs: + c = BatchedGemmKernelConfig( + dtype_a=b.dtype_a, + dtype_b=b.dtype_b, + dtype_c=b.dtype_c, + dtype_acc=b.dtype_acc, + layout_a=b.layout_a, + layout_b=b.layout_b, + layout_c=b.layout_c, + tile_m=b.tile_m, + tile_n=b.tile_n, + tile_k=b.tile_k, + wave_m=b.wave_m, + wave_n=b.wave_n, + wave_k=b.wave_k, + warp_tile_m=b.warp_tile_m, + warp_tile_n=b.warp_tile_n, + warp_tile_k=b.warp_tile_k, + pipeline=b.pipeline, + scheduler=b.scheduler, + epilogue=b.epilogue, + pad_m=b.pad_m, + pad_n=b.pad_n, + pad_k=b.pad_k, + persistent=b.persistent, + gfx_arch=b.gfx_arch, + variant="batched", + ) + if c.name in seen: + continue + seen.add(c.name) + out.append(c) + return out diff --git a/projects/composablekernel/dispatcher/python/gemm_utils.py b/projects/composablekernel/dispatcher/python/gemm_utils.py index 50367bd26b7a..4ef7bc06b6a8 100644 --- a/projects/composablekernel/dispatcher/python/gemm_utils.py +++ b/projects/composablekernel/dispatcher/python/gemm_utils.py @@ -2095,6 +2095,23 @@ def _is_power_of_two(x: int) -> bool: return x > 0 and (x & (x - 1)) == 0 +def _cshuffle_store_ok( + m_repeat: int, n_repeat: int, warp_tile_m: int, warp_tile_n: int +) -> bool: + """Return False for the one CShuffle-store combination that is numerically + wrong (issue #9684): an ODD per-wave repeat (>1) paired with a 32-wide warp + tile in that dimension. GPU-verified on gfx942 -- e.g. tile_m=192 / wave_m=2 + / warp_tile_m=32 (MRepeat=3) returns garbage, while every other non-power-of- + two repeat (incl. MRepeat=3 with warp_tile_m=16, and even repeats like 6/12) + is correct. Only relevant for the CShuffle epilogue; the default epilogue is + exempt.""" + + def _dim_bad(repeat: int, warp_tile: int) -> bool: + return repeat > 1 and repeat % 2 == 1 and warp_tile == 32 + + return not (_dim_bad(m_repeat, warp_tile_m) or _dim_bad(n_repeat, warp_tile_n)) + + # --- Warp-configuration gate (parity with Old-TE) -------------------------- # Old-TE's gemm_validation_utils.validate_warp_configuration restricts the # warps-per-block triple (wave_m/n/k) to WARP_SUPPORTED_COMBINATIONS[arch]. diff --git a/projects/composablekernel/dispatcher/tests/CMakeLists.txt b/projects/composablekernel/dispatcher/tests/CMakeLists.txt index 6f562ee9609c..5bd4c35d8207 100644 --- a/projects/composablekernel/dispatcher/tests/CMakeLists.txt +++ b/projects/composablekernel/dispatcher/tests/CMakeLists.txt @@ -10,6 +10,18 @@ cmake_minimum_required(VERSION 3.16) # Find Python find_package(Python3 COMPONENTS Interpreter REQUIRED) +# Bridge on-device GPU correctness test (SKIPs cleanly without a GPU/built lib). +add_test( + NAME dispatcher_test_batched_gemm_gpu_correctness + COMMAND ${Python3_EXECUTABLE} -m unittest discover -s ${CMAKE_CURRENT_SOURCE_DIR} -p test_batched_gemm_gpu_correctness.py -v + WORKING_DIRECTORY ${CMAKE_CURRENT_SOURCE_DIR} +) +set_tests_properties(dispatcher_test_batched_gemm_gpu_correctness PROPERTIES + LABELS "dispatcher;python;bridge;gpu" + TIMEOUT 600 + ENVIRONMENT "PYTHONPATH=${CMAKE_CURRENT_SOURCE_DIR}/../python:${CMAKE_CURRENT_SOURCE_DIR}/../codegen:${CMAKE_CURRENT_SOURCE_DIR}/../scripts" +) + # ============================================================================= # Python Tests # ============================================================================= diff --git a/projects/composablekernel/dispatcher/tests/test_batched_bridge.py b/projects/composablekernel/dispatcher/tests/test_batched_bridge.py new file mode 100644 index 000000000000..92d4c54d38ac --- /dev/null +++ b/projects/composablekernel/dispatcher/tests/test_batched_bridge.py @@ -0,0 +1,357 @@ +#!/usr/bin/env python3 + +# Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +# SPDX-License-Identifier: MIT + +"""CPU-only unit tests for the batched GEMM bridge. + +Batched GEMM reuses the plain-GEMM config/codegen machinery but pins the +variant to ``"batched"`` so ``BatchedGemmKernelConfig.name`` gains a trailing +``_batched`` token. That token is the byte-parity invariant tying config -> +codegen -> the compiled kernel name the runtime reports, so distinct batched +kernels never collapse onto the plain-GEMM name. + +Everything under test is pure host-side logic (name generation, the codegen +JSON projection, the batched problem flop count, and the shipped +configs/*.json). No GPU, hipcc, or dispatcher build is required. + +Run: python3 -m pytest tests/test_batched_bridge.py -v +""" + +import ctypes +import json +import sys +import unittest +from pathlib import Path +from unittest import mock + +import numpy as np + +SCRIPT_DIR = Path(__file__).parent.resolve() +DISPATCHER_DIR = SCRIPT_DIR.parent +REPO_ROOT = DISPATCHER_DIR.parent +sys.path.insert(0, str(DISPATCHER_DIR / "python")) +sys.path.insert(0, str(DISPATCHER_DIR / "codegen")) + +import batched_gemm_utils # noqa: E402 +from batched_gemm_utils import ( # noqa: E402 + BatchedGemmDispatcherLib, + BatchedGemmKernelConfig, + BatchedGemmProblem, + _C_SIZEOF, + _get_arch, + _resolve_arch, + expand_sweep, +) +from gemm_utils import ( # noqa: E402 + GemmKernelConfig, + _output_dtype, + _dtype_from_kernel_name, +) + +_CONFIG_DIR = ( + REPO_ROOT + / "tile_engine" + / "ops" + / "gemm" + / "batched_gemm" + / "configs" +) + + +def _make_config(**overrides): + kw = dict( + dtype_a="fp16", + dtype_b="fp16", + dtype_c="fp16", + dtype_acc="fp32", + layout_a="row", + layout_b="col", + layout_c="row", + ) + kw.update(overrides) + return BatchedGemmKernelConfig(**kw) + + +class TestBatchedName(unittest.TestCase): + """The batched kernel-name contract (the byte-parity invariant).""" + + def test_variant_is_batched(self): + self.assertEqual(_make_config().variant, "batched") + + def test_name_gains_batched_suffix(self): + cfg = _make_config() + self.assertTrue(cfg.name.endswith("_batched"), cfg.name) + + def test_name_is_plain_gemm_plus_suffix(self): + # The batched name must be exactly the plain-GEMM name + "_batched", + # so the two bridges share codegen but never collide. + common = dict( + dtype_a="bf16", dtype_b="bf16", dtype_c="bf16", dtype_acc="fp32", + layout_a="row", layout_b="row", layout_c="row", + tile_m=128, tile_n=128, tile_k=32, + wave_m=2, wave_n=2, wave_k=1, + warp_tile_m=32, warp_tile_n=32, warp_tile_k=16, + pipeline="compv4", scheduler="intrawave", epilogue="cshuffle", + pad_m=True, pad_n=True, pad_k=True, persistent=False, + ) + plain = GemmKernelConfig(**common) + batched = BatchedGemmKernelConfig(**common) + self.assertEqual(batched.name, plain.name + "_batched") + + def test_suffix_not_doubled(self): + # Re-deriving the name must not append a second "_batched". + cfg = _make_config() + self.assertEqual(cfg.name.count("_batched"), 1) + + def test_dtype_recovers_from_name(self): + for dtype in ("fp16", "bf16"): + cfg = _make_config( + dtype_a=dtype, dtype_b=dtype, dtype_c=_output_dtype(dtype), + ) + self.assertEqual(_dtype_from_kernel_name(cfg.name), dtype) + + +class TestBatchedCodegenJson(unittest.TestCase): + """The codegen JSON projection is inherited from the plain-GEMM config.""" + + def test_codegen_json_core_blocks(self): + j = _make_config().to_codegen_json() + self.assertIn("tile_config", j) + self.assertIn("trait_config", j) + # dispatcher wave_* is projected onto codegen warp_* (warps per block). + self.assertEqual(j["tile_config"]["warp_m"], [2]) + + +class TestBatchedProblem(unittest.TestCase): + """The batched problem carries the batch dimension into the flop count.""" + + def test_flops_scale_with_batch(self): + one = BatchedGemmProblem(batch_count=1, M=128, N=128, K=64) + eight = BatchedGemmProblem(batch_count=8, M=128, N=128, K=64) + self.assertEqual(eight.flops, 8 * one.flops) + + def test_flops_formula(self): + p = BatchedGemmProblem(batch_count=4, M=32, N=16, K=8) + self.assertEqual(p.flops, 2.0 * 4 * 32 * 16 * 8) + + def test_dict_roundtrip(self): + p = BatchedGemmProblem(batch_count=3, M=64, N=48, K=16) + back = BatchedGemmProblem.from_dict(p.to_dict()) + self.assertEqual(back.batch_count, 3) + self.assertEqual((back.M, back.N, back.K), (64, 48, 16)) + + +class TestBatchedShippedConfigs(unittest.TestCase): + """The configs/*.json the bridge codegens from must be well-formed sweeps.""" + + def test_config_dir_exists(self): + self.assertTrue(_CONFIG_DIR.is_dir(), _CONFIG_DIR) + + def test_configs_are_valid_sweeps(self): + configs = sorted(_CONFIG_DIR.glob("*.json")) + self.assertGreater(len(configs), 0, "no batched configs shipped") + for path in configs: + with self.subTest(config=path.name): + with open(path) as f: + data = json.load(f) + self.assertIn("tile_config", data) + self.assertIn("trait_config", data) + tc = data["tile_config"] + for key in ( + "tile_m", "tile_n", "tile_k", + "warp_m", "warp_n", "warp_k", + "warp_tile_m", "warp_tile_n", "warp_tile_k", + ): + self.assertIn(key, tc, f"{path.name} missing {key}") + tr = data["trait_config"] + for key in ("pipeline", "scheduler", "epilogue"): + self.assertIn(key, tr, f"{path.name} missing {key}") + + +class TestBatchedArchResolution(unittest.TestCase): + """`_get_arch` / `_resolve_arch`: detect + validate, never a silent default. + + `_get_arch` shells out to `rocminfo`; these tests mock + `subprocess.check_output` so they run GPU-free, covering the three outcomes + the arch feedback demanded: valid detect, undetectable -> RuntimeError, + unsupported -> ValueError. + """ + + _ROCMINFO = "Agent 1\n Name: AMD EPYC\nAgent 2\n Name: gfx942\n" + + def test_get_arch_detects_supported(self): + with mock.patch("subprocess.check_output", return_value=self._ROCMINFO): + self.assertEqual(_get_arch(), "gfx942") + + def test_get_arch_undetectable_raises_runtimeerror(self): + # rocminfo missing / no gfx line -> refuse to guess. + with mock.patch("subprocess.check_output", side_effect=FileNotFoundError()): + with self.assertRaises(RuntimeError): + _get_arch() + + def test_get_arch_unsupported_raises_valueerror(self): + out = "Agent 1\n Name: gfx1030\n" + with mock.patch("subprocess.check_output", return_value=out): + with self.assertRaises(ValueError): + _get_arch() + + def test_resolve_arch_explicit_valid_passthrough(self): + # An explicit, supported arch must not touch rocminfo at all. + with mock.patch("subprocess.check_output", side_effect=AssertionError("called")): + self.assertEqual(_resolve_arch("gfx950"), "gfx950") + + def test_resolve_arch_explicit_invalid_raises(self): + with self.assertRaises(ValueError): + _resolve_arch("gfx942x") + + def test_resolve_arch_none_autodetects(self): + with mock.patch("subprocess.check_output", return_value=self._ROCMINFO): + self.assertEqual(_resolve_arch(None), "gfx942") + + +class TestBatchedAbiMarshalling(unittest.TestCase): + """GPU-free tests of the C-ABI surface in ``BatchedGemmDispatcherLib``. + + The ``.so`` is mocked (``ctypes.CDLL`` patched), so no generated kernel or + GPU is needed. Locks the declared ABI (19 args + time_ms) and the positional + order/values ``run()`` forwards -- the surface most likely to drift silently + against the C entry point ``dispatcher_run_batched``. + """ + + def _make_lib(self, status=0): + fake = mock.MagicMock() + fake.dispatcher_run_batched.return_value = status + with mock.patch("ctypes.CDLL", return_value=fake): + lib = BatchedGemmDispatcherLib(Path("/nonexistent/batched.so")) + return lib, fake + + def test_run_abi_argtypes(self): + _, fake = self._make_lib() + argt = fake.dispatcher_run_batched.argtypes + # 18 scalar/pointer args + the time_ms out-pointer. + self.assertEqual(len(argt), 19) + for i in range(3): # A/B/C host pointers + self.assertEqual(argt[i], ctypes.c_void_p) + for i in range(3, 18): # M..rotating_count are all int64 + self.assertEqual(argt[i], ctypes.c_int64) + self.assertEqual(argt[-1], ctypes.POINTER(ctypes.c_float)) + self.assertEqual(fake.dispatcher_run_batched.restype, ctypes.c_int) + + def test_run_marshals_positional_order(self): + lib, fake = self._make_lib(status=0) + A = np.zeros((2, 2, 2), np.float16) + B = np.zeros((2, 2, 2), np.float16) + C = np.zeros((2, 2, 2), np.float16) + status, _ = lib.run( + A, B, C, M=2, N=2, K=2, batch_count=2, k_batch=1, + stride_A=2, stride_B=2, stride_C=2, + batch_stride_A=4, batch_stride_B=4, batch_stride_C=4, + warmup=50, repeat=100, flush_cache=True, rotating_count=1000, + ) + self.assertEqual(status, 0) + args = fake.dispatcher_run_batched.call_args[0] + self.assertEqual(len(args), 19) + # M/N/K/batch_count/k_batch occupy positions 3..7. + self.assertEqual((args[3], args[4], args[5]), (2, 2, 2)) + self.assertEqual((args[6], args[7]), (2, 1)) + # flush_cache marshals to the int 1 (position 16), not the bool True. + self.assertEqual(args[16], 1) + + def test_run_forwards_nonzero_status(self): + # A thin shim must surface the C error code verbatim (e.g. -2 launch). + lib, fake = self._make_lib(status=-2) + A = np.zeros((1, 1, 1), np.float16) + status, _ = lib.run(A, A, A, M=1, N=1, K=1, batch_count=1) + self.assertEqual(status, -2) + + +class TestBatchedDtypeLayoutGate(unittest.TestCase): + """`expand_sweep` must reject anything outside Old-TE's fp16/rcr set. + + Old-TE ``batched_gemm_instance_builder`` declares ``--datatype + choices=['fp16']`` / ``--layout choices=['rcr']``; the bridge must match that + EXACTLY rather than silently building a signature Old-TE never validated. An + explicit (supported) arch is passed so the gate fires before rocminfo/config + I/O is touched. + """ + + def test_rejects_non_fp16_dtype(self): + with self.assertRaises(ValueError): + expand_sweep("/nonexistent/config.json", arch="gfx942", dtype="bf16") + + def test_rejects_non_rcr_layout(self): + with self.assertRaises(ValueError): + expand_sweep("/nonexistent/config.json", arch="gfx942", layout="rrr") + + +class TestBatchedSplitKContract(unittest.TestCase): + """F2: lock the generated split-K reduction contract. + + For ``k_batch > 1`` the ck_tile batched kernel switches its epilogue to + ``memory_operation_enum::atomic_add``, so C MUST be zeroed before the launch + (a stale C would be accumulated into and silently corrupt the result). For + ``k_batch == 1`` the epilogue is ``set`` and C must NOT be pre-zeroed (that + is the Old-TE byte-identical parity path). This test renders the batched + launch via the codegen and asserts both halves of the contract, so a future + codegen edit that drops the memset (or the return check) fails loudly here + instead of only on a GPU under split-K. + """ + + @classmethod + def setUpClass(cls): + from codegen_common import TileConfig + from unified_gemm_codegen import ( + CKTileKernelGenerator, + GemmVariant, + KernelConfig, + TraitConfig, + ) + + 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, + ) + cfg = KernelConfig(tile=tile, trait=trait, variant=GemmVariant.BATCHED) + cls.src = CKTileKernelGenerator("fp16", "rcr").generate(cfg) + + def test_split_k_branch_zeroes_c(self): + # The k_batch>1 branch exists and zeroes C via hipMemsetAsync. + self.assertIn("if(args.k_batch > 1)", self.src) + self.assertIn("hipMemsetAsync(args.e_ptr", self.src) + # The reason is the atomic-add reduction; the contract is documented. + self.assertIn("atomic_add", self.src) + + def test_split_k_memset_return_is_checked(self): + # The Copilot follow-up: a failed reset must throw, not run silently. + self.assertIn("!= hipSuccess", self.src) + self.assertIn("failed to reset C", self.src) + + def test_k_batch_one_uses_set_no_memset(self): + # The k_batch==1 parity path uses memory_operation_enum::set and must not + # pre-zero C (byte-identical to the Old-TE batched launch). + self.assertIn("memory_operation_enum::set", self.src) + # Exactly one memset in the batched launch: the split-K path only. + self.assertEqual(self.src.count("hipMemsetAsync(args.e_ptr"), 1) + + +class TestBatchedCDtypeSizeMap(unittest.TestCase): + """F6: the host numpy C dtype size must equal sizeof(CDataType).""" + + def test_c_sizeof_matches_numpy_itemsize(self): + _C_NP = {"fp16": np.float16, "bf16": np.uint16, "int32": np.int32} + for out_dtype, nbytes in _C_SIZEOF.items(): + self.assertIn(out_dtype, _C_NP) + self.assertEqual( + np.dtype(_C_NP[out_dtype]).itemsize, nbytes, + f"{out_dtype}: numpy itemsize != declared sizeof(CDataType)", + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/projects/composablekernel/dispatcher/tests/test_batched_gemm_gpu_correctness.py b/projects/composablekernel/dispatcher/tests/test_batched_gemm_gpu_correctness.py new file mode 100644 index 000000000000..240263950813 --- /dev/null +++ b/projects/composablekernel/dispatcher/tests/test_batched_gemm_gpu_correctness.py @@ -0,0 +1,179 @@ +#!/usr/bin/env python3 + +# Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +# SPDX-License-Identifier: MIT + +""" +GPU correctness test for the batched GEMM dispatcher bridge. + +Builds the batched GEMM dispatcher .so (fp16 / rcr — the only signature Old-TE's +batched_gemm_instance_builder validates), runs a small batched problem on-device +via GpuBatchedGemmRunner, and compares the GPU output to a per-batch fp32 numpy +reference within an fp16-appropriate tolerance. Skips cleanly (exit 0) when no +GPU / hipcc is available so it is safe in a CPU-only CI lane. + +The kernel computes, for each batch b: + C[b] = A[b] @ B[b] A[b] is M x K, B[b] is K x N, C[b] is M x N +with the rcr layout: A row-major, B column-major, C row-major. The runner +handles the operand layout transform, so the test hands it logical row-major A +and logical (K, N) B per batch. + +Run: + python3 test_batched_gemm_gpu_correctness.py + python3 test_batched_gemm_gpu_correctness.py -v + python3 test_batched_gemm_gpu_correctness.py --gfx gfx942 +""" + +import argparse +import logging +import sys +import tempfile +from pathlib import Path + +import numpy as np + +sys.path.insert(0, str(Path(__file__).resolve().parents[1] / "python")) + +from batched_gemm_utils import ( # noqa: E402 + BatchedGemmKernelConfig, + BatchedGemmProblem, + GpuBatchedGemmRunner, + setup_multiple_batched_gemm_dispatchers, + _resolve_arch, +) + +log = logging.getLogger(__name__) + +# fp16 accumulate-in-fp32 GEMM: worst-case relative error is well under 1e-2 for +# the small K used here. Use the standard fp16 bar. +TOLERANCE = 1e-2 + +PASS = "PASS" +FAIL = "FAIL" +SKIP = "SKIP" + + +def _has_gpu() -> bool: + """True iff a supported GPU is visible to rocminfo (so the .so can run).""" + try: + _resolve_arch(None) + return True + except Exception: + return False + + +def _max_rel_err(C_gpu: np.ndarray, C_ref: np.ndarray) -> float: + """Max absolute error normalized by the largest reference magnitude. + + A GEMM output has elements that partially cancel toward zero, so a naive + per-element relative error is dominated by those near-zero entries and does + NOT measure whether the kernel computed the right matrix. Normalizing the + worst absolute error by the global reference scale (max |ref|) is the honest + correctness bar: any structural error (wrong accumulation, transposed + operand, wrong batch stride) blows far past 1e-2, while correct fp16 math + lands at ~1e-3-1e-4 here. + """ + g = C_gpu.astype(np.float32) + r = C_ref.astype(np.float32) + ref_scale = max(float(np.abs(r).max()), 1e-6) + return float(np.max(np.abs(g - r)) / ref_scale) + + +def _reference_batched(A: np.ndarray, B: np.ndarray) -> np.ndarray: + """Per-batch fp32 reference: C[b] = A[b] @ B[b] (A: b,M,K B: b,K,N).""" + return np.matmul(A.astype(np.float32), B.astype(np.float32)) + + +def _make_fp16_config(gfx_arch: str) -> BatchedGemmKernelConfig: + """A single small, valid fp16/rcr batched kernel. + + 128x128x32 tile, 2x2x1 waves, 32x32x16 warp-tile, compv3/intrawave/cshuffle + — a divisibility-valid combination on both gfx942 and gfx950. Padding is on + so non-tile-multiple shapes still run. + """ + return BatchedGemmKernelConfig( + dtype_a="fp16", dtype_b="fp16", dtype_c="fp16", dtype_acc="fp32", + layout_a="row", layout_b="col", layout_c="row", + tile_m=128, tile_n=128, tile_k=32, + wave_m=2, wave_n=2, wave_k=1, + warp_tile_m=32, warp_tile_n=32, warp_tile_k=16, + pipeline="compv3", scheduler="intrawave", epilogue="cshuffle", + pad_m=True, pad_n=True, pad_k=True, persistent=False, + gfx_arch=gfx_arch, + ) + + +def test_batched_fp16(gfx_arch: str) -> tuple[str, str]: + # Small multi-batch problem; K=128 gives 4 tile-K iterations (128/32). + batch, M, N, K = 3, 128, 128, 128 + cfg = _make_fp16_config(gfx_arch) + + so_paths = setup_multiple_batched_gemm_dispatchers([cfg], verbose=False) + if not so_paths or so_paths[0] is None: + return FAIL, "batched/fp16: kernel build failed" + + runner = GpuBatchedGemmRunner(so_paths[0]) + + rng = np.random.default_rng(7) + A = rng.uniform(-1.0, 1.0, (batch, M, K)).astype(np.float32) + B = rng.uniform(-1.0, 1.0, (batch, K, N)).astype(np.float32) + + problem = BatchedGemmProblem(batch_count=batch, M=M, N=N, K=K) + result = runner.run(A, B, problem, warmup=5, repeat=10) + + if result.status != 0: + return FAIL, f"batched/fp16: run status={result.status} (nonzero)" + C_gpu = result.output + if C_gpu.shape != (batch, M, N): + return FAIL, f"batched/fp16: output shape {C_gpu.shape} != {(batch, M, N)}" + if np.all(C_gpu == 0): + return FAIL, "batched/fp16: GPU output is all-zero" + if not np.all(np.isfinite(C_gpu.astype(np.float32))): + return FAIL, "batched/fp16: GPU output contains NaN/Inf" + + # A[b]@B[b]: the runner already produced logical (batch, M, N). B is passed + # logical (K, N); the runner transposes it for the col-major kernel, so the + # reference multiplies the same logical operands. + C_ref = _reference_batched(A, B) + mre = _max_rel_err(C_gpu, C_ref) + if mre > TOLERANCE: + return FAIL, (f"batched/fp16: max_rel_err={mre:.4e} > tol={TOLERANCE:.1e} " + f"(batch={batch} M={M} N={N} K={K})") + if result.time_ms <= 0.0: + return FAIL, f"batched/fp16: time_ms={result.time_ms:.4f} not positive" + + return PASS, (f"batched/fp16: max_rel_err={mre:.4e}, " + f"time_ms={result.time_ms:.3f}, batch={batch} MNK={M}/{N}/{K}") + + +def main() -> int: + parser = argparse.ArgumentParser(description="Batched GEMM GPU correctness test") + parser.add_argument("--gfx", default=None, help="GPU arch (default: auto-detect)") + parser.add_argument("-v", "--verbose", action="store_true") + args = parser.parse_args() + + logging.basicConfig( + level=logging.DEBUG if args.verbose else logging.INFO, + format="%(levelname)s: %(message)s", + ) + + if not _has_gpu(): + print("SKIP: no supported GPU detected (rocminfo); batched GPU test skipped") + return 0 + + gfx = args.gfx or _resolve_arch(None) + log.info("Running batched GEMM GPU correctness on %s", gfx) + + try: + status, detail = test_batched_fp16(gfx) + except Exception as exc: # noqa: BLE001 + status, detail = FAIL, f"batched/fp16: exception: {exc}" + + print("\n=== Summary ===") + print(f" [{status:4s}] {detail}") + print(f"\n{1 if status == PASS else 0}/1 passed") + return 0 if status == PASS else 1 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/projects/composablekernel/dispatcher/tests/test_gemm_utils.py b/projects/composablekernel/dispatcher/tests/test_gemm_utils.py index 3dd2f50ace88..14afc24a2b19 100644 --- a/projects/composablekernel/dispatcher/tests/test_gemm_utils.py +++ b/projects/composablekernel/dispatcher/tests/test_gemm_utils.py @@ -40,6 +40,7 @@ _output_dtype, _dtype_from_kernel_name, _layout_from_kernel_name, + _cshuffle_store_ok, ) @@ -202,6 +203,37 @@ def test_name_roundtrips_through_parsers(self): self.assertEqual(_layout_from_kernel_name(name), cfg.layout) +class TestCShuffleStoreGate(unittest.TestCase): + """Narrowed CShuffle-store correctness gate (issue #9684). + + Only an ODD per-wave repeat (>1) with a 32-wide warp tile in that dimension + is numerically wrong; every other non-power-of-two repeat is correct. These + expectations were GPU-verified on gfx942 (26 broken / 90 correct across the + tile_m=192 cshuffle config space). + """ + + def test_broken_signature_rejected(self): + # tile_m=192 / wave_m=2 / warp_tile_m=32 -> MRepeat = 192/(2*32) = 3. + # The 26 verified-wrong configs all match this (odd repeat + 32 warp). + self.assertFalse(_cshuffle_store_ok(3, 2, 32, 32)) + self.assertFalse(_cshuffle_store_ok(3, 4, 32, 16)) # M side triggers + self.assertFalse(_cshuffle_store_ok(4, 3, 16, 32)) # N side triggers + + def test_odd_repeat_with_16_warp_tile_allowed(self): + # MRepeat=3 via wave_m=4 / warp_tile_m=16 is numerically correct. + self.assertTrue(_cshuffle_store_ok(3, 2, 16, 16)) + + def test_even_nonpow2_repeat_allowed(self): + # Repeats 6 and 12 are non-power-of-two but verified correct, incl. w/32. + self.assertTrue(_cshuffle_store_ok(6, 4, 32, 16)) + self.assertTrue(_cshuffle_store_ok(12, 2, 16, 32)) + + def test_power_of_two_repeats_allowed(self): + for rep in (1, 2, 4, 8): + self.assertTrue(_cshuffle_store_ok(rep, rep, 32, 32)) + self.assertTrue(_cshuffle_store_ok(rep, rep, 16, 16)) + + class TestModuleImportsAndRunnerShape(unittest.TestCase): """Guards against a merge truncating gemm_utils (regression: #9308 dropped the tail of GpuMultiDGemmRunner.run, leaving an unterminated diff --git a/projects/composablekernel/tile_engine/ops/gemm/batched_gemm_full_benchmark.py b/projects/composablekernel/tile_engine/ops/gemm/batched_gemm_full_benchmark.py new file mode 100644 index 000000000000..20dc0b0338bd --- /dev/null +++ b/projects/composablekernel/tile_engine/ops/gemm/batched_gemm_full_benchmark.py @@ -0,0 +1,471 @@ +#!/usr/bin/env python3 +# Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +# SPDX-License-Identifier: MIT +"""Full BATCHED GEMM benchmark sweep driven through the Dispatcher bridge. + +Batched counterpart of gemm_full_benchmark.py. + +Phases: + Phase 1: Compile all batched kernels (parallel, returns .so paths -- no GPU) + Phase 2: Load batched problems (batch_count, M, N, K) + Phase 3: Benchmark via subprocess isolation, fanned across all visible GPUs + +Tile Engine generates NO binaries here: it expands its sweep config into shared +``BatchedGemmKernelConfig`` objects and hands them to the dispatcher, which +codegens + compiles each into a .so. Each kernel runs in a disposable worker +subprocess so a GPU fault takes down only one worker. + +Examples: + python batched_gemm_full_benchmark.py + python batched_gemm_full_benchmark.py \ + batched_gemm/configs/default_config.json --devices 4 --csv out.csv +""" + +import argparse +import csv +import json +import os +import queue +import subprocess +import sys +import threading +import time +from pathlib import Path + +_THIS_DIR = Path(__file__).resolve().parent +_COMMON_DIR = _THIS_DIR.parent / "common" +_DISPATCHER_ROOT = _THIS_DIR.parents[2] / "dispatcher" +sys.path.insert(0, str(_DISPATCHER_ROOT / "python")) +sys.path.insert(0, str(_COMMON_DIR)) +sys.path.insert(0, str(_THIS_DIR)) + +from batched_gemm_utils import ( # noqa: E402 + setup_multiple_batched_gemm_dispatchers, + expand_sweep, +) +from smi_utils import detect_gpu_ids # noqa: E402 + +# The batched op keeps its sweep configs in tile_engine/ops/gemm/batched_gemm/configs. +CONFIG_DIR = _THIS_DIR / "batched_gemm" / "configs" +CI_CONFIG_NAME = "default_ci_config.json" +DEFAULT_CONFIG_NAME = "default_config.json" +EXAMPLE_PROBLEMS_NAME = "example_problems.json" + +# Fallback batched problem set if none is supplied. +DEFAULT_PROBLEMS = [ + {"batch_count": 8, "M": 1024, "N": 1024, "K": 1024}, + {"batch_count": 4, "M": 2048, "N": 2048, "K": 2048}, + {"batch_count": 16, "M": 512, "N": 512, "K": 512}, + {"batch_count": 2, "M": 3840, "N": 4096, "K": 2048}, +] + +# Batched GEMM TE capability set: fp16 / rcr only. +SUPPORTED_DTYPES = ("fp16",) +SUPPORTED_LAYOUTS = ("rcr",) + + +def detect_devices(): + """Return a list of visible GPU id strings (best-effort).""" + return detect_gpu_ids() + + +def resolve_devices(spec): + """Resolve --devices into a concrete list of device id strings.""" + detected = detect_devices() + if spec is None: + return detected + spec = str(spec).strip() + if "," in spec: + return [s.strip() for s in spec.split(",") if s.strip() != ""] + if spec.isdigit(): + n = int(spec) + if n <= 0: + return detected + if len(detected) >= n: + return detected[:n] + if os.environ.get("HIP_VISIBLE_DEVICES") or os.environ.get( + "CUDA_VISIBLE_DEVICES" + ): + return detected + return [str(i) for i in range(n)] + return [spec] + + +def resolve_configs(args): + """Resolve positional configs -> concrete list of config paths.""" + if args.configs: + return args.configs + return [str(CONFIG_DIR / CI_CONFIG_NAME)] + + +def load_problems(path): + if path: + with open(path) as f: + data = json.load(f) + return data["problems"] if isinstance(data, dict) else data + example = CONFIG_DIR / EXAMPLE_PROBLEMS_NAME + if example.exists(): + with open(example) as f: + data = json.load(f) + return data["problems"] if isinstance(data, dict) else data + return DEFAULT_PROBLEMS + + +def _run_batch_on_device(device_id, unit, args, worker_path, base_env): + """Run one (problem, kernel-batch) unit in a device-pinned subprocess.""" + prob_idx, prob_dict, batch = unit + bc, M, N, K = ( + prob_dict["batch_count"], + prob_dict["M"], + prob_dict["N"], + prob_dict["K"], + ) + + items = [ + {"so_path": str(lib), "problem": prob_dict, "kernel_name": cfg.name} + for _, cfg, lib in batch + ] + payload = json.dumps( + {"items": items, "verify": args.verify, "verify_tol": args.verify_tol} + ) + + env = base_env.copy() + env["HIP_VISIBLE_DEVICES"] = str(device_id) + + rows, lines, n_fail = [], [], 0 + proc = None + try: + proc = subprocess.Popen( + [sys.executable, str(worker_path)], + stdin=subprocess.PIPE, + stdout=subprocess.PIPE, + stderr=subprocess.DEVNULL, + env=env, + ) + stdout_bytes, _ = proc.communicate( + input=payload.encode("utf-8"), + timeout=args.kernel_timeout * len(batch), + ) + + reported = set() + for line in stdout_bytes.decode("utf-8").strip().split("\n"): + if not line: + continue + try: + result = json.loads(line) + except json.JSONDecodeError: + lines.append( + f" [gpu{device_id}] Warning: bad result line: {line[:50]}" + ) + n_fail += 1 + continue + bidx = result.get("idx", 0) + _, cfg, _ = batch[bidx] + reported.add(bidx) + if result.get("ok", False): + status = "OK" if result.get("non_zero", 0) > 0 else "ZERO" + mismatch = False + if args.verify and "verified" in result: + if result["verified"]: + status = "VERIFY" + else: + status = "MISMATCH" + mismatch = True + extra = ( + f" rel={result['max_rel']:.2e}" if "max_rel" in result else "" + ) + lines.append( + f" [gpu{device_id}] {cfg.name:<62} {result['ms']:>10.3f} " + f"{result['tflops']:>10.2f} {status:>8}{extra}" + ) + rows.append( + { + "kernel": cfg.name, + "problem_idx": prob_idx, + "batch_count": bc, + "M": M, + "N": N, + "K": K, + "device": device_id, + "latency_ms": result["ms"], + "tflops": result["tflops"], + "non_zero": result.get("non_zero", 0), + "max_rel": result.get("max_rel", ""), + "verified": result.get("verified", ""), + } + ) + if mismatch: + n_fail += 1 + else: + lines.append(f" [gpu{device_id}] {cfg.name:<62} FAILED") + lines.append(f" Error: {result.get('error', 'unknown')[:100]}") + n_fail += 1 + + missing = set(range(len(batch))) - reported + if missing or proc.returncode != 0: + if proc.returncode != 0: + lines.append( + f" [gpu{device_id}] worker exited code {proc.returncode}" + ) + for idx in sorted(missing): + _, cfg, _ = batch[idx] + lines.append(f" [gpu{device_id}] {cfg.name:<62} MISSING (crash)") + n_fail += len(missing) + + except subprocess.TimeoutExpired: + lines.append(f" [gpu{device_id}] batch timeout ({len(batch)} kernels)") + try: + proc.kill() + proc.communicate(timeout=5) + except Exception: + pass + n_fail += len(batch) + except Exception as e: + lines.append(f" [gpu{device_id}] batch error: {e}") + try: + if proc and proc.poll() is None: + proc.kill() + except Exception: + pass + n_fail += len(batch) + + return rows, lines, n_fail + + +def main(): + parser = argparse.ArgumentParser( + description="Batched GEMM Benchmark Sweep (via Dispatcher)" + ) + parser.add_argument( + "configs", + nargs="*", + help="TE sweep config JSON files (default: batched_gemm/configs/default_ci_config.json)", + ) + parser.add_argument("--arch", default="gfx942") + parser.add_argument( + "--dtype", + default="fp16", + choices=SUPPORTED_DTYPES, + help=f"Input dtype (supported: {', '.join(SUPPORTED_DTYPES)})", + ) + parser.add_argument( + "--layout", + default="rcr", + choices=SUPPORTED_LAYOUTS, + help=f"A/B/C layout (supported: {', '.join(SUPPORTED_LAYOUTS)})", + ) + parser.add_argument( + "--problems", default=None, help="JSON file of batched problems" + ) + parser.add_argument("--csv", type=str, default="batched_gemm_results.csv") + parser.add_argument( + "--workers", type=int, default=8, help="Parallel build workers" + ) + parser.add_argument( + "--devices", + default=None, + help="GPUs to use: int count or comma-list of ids; default auto-detect", + ) + 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=30, help="Per-kernel timeout (s)" + ) + parser.add_argument( + "--max-kernels", type=int, default=0, help="Limit to first N kernels (0=all)" + ) + parser.add_argument( + "--verify", + action="store_true", + help="Check each kernel's output against an fp32 numpy reference " + "(per-batch A @ B); a mismatch counts as a failure", + ) + parser.add_argument( + "--verify-tol", + type=float, + default=2e-2, + help="Relative tolerance for --verify (default 2e-2, suits fp16)", + ) + args = parser.parse_args() + + # --batch-size is the step of range(0, len(built_kernels), args.batch_size); + # a zero step raises ValueError and a negative one silently yields no batches, + # so reject non-positive values up front with a clear message. + if args.batch_size <= 0: + parser.error("--batch-size must be a positive integer (kernels per subprocess)") + + config_paths = resolve_configs(args) + devices = resolve_devices(args.devices) + + # ======================================================================== + # Phase 1: Compile kernels (parallel, no GPU) + # ======================================================================== + print(f"\n{'=' * 80}") + print("Phase 1: Compile batched kernels") + print(f"{'=' * 80}") + print(f" Configs: {', '.join(config_paths)}") + + all_configs = [] + for cfg_path in config_paths: + all_configs.extend( + expand_sweep(cfg_path, args.arch, dtype=args.dtype, layout=args.layout) + ) + + 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() + lib_paths = setup_multiple_batched_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 + ] + + 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 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 across all visible GPUs (subprocess isolation, batched) + # ======================================================================== + print(f"\n{'=' * 80}") + print("Phase 3: Benchmark (multi-GPU, subprocess isolation, batched)") + print(f"{'=' * 80}") + print(f" Devices: {len(devices)} -> {', '.join(devices)}") + 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", + "batch_count", + "M", + "N", + "K", + "device", + "latency_ms", + "tflops", + "non_zero", + "max_rel", + "verified", + ] + csv_file = open(csv_path, "w", newline="") + writer = csv.DictWriter(csv_file, fieldnames=csv_fields) + writer.writeheader() + + worker_path = _THIS_DIR / "run_one_batched_gemm_kernel.py" + base_env = os.environ.copy() + base_env["GEMM_PYPATH"] = os.pathsep.join( + [str(_DISPATCHER_ROOT / "python"), str(_THIS_DIR)] + ) + + work_q = queue.Queue() + for prob_idx, prob in enumerate(problems): + prob_dict = { + "batch_count": int(prob["batch_count"]), + "M": int(prob["M"]), + "N": int(prob["N"]), + "K": int(prob["K"]), + } + for start in range(0, len(built_kernels), args.batch_size): + end = min(start + args.batch_size, len(built_kernels)) + batch = [ + (start + j, cfg, lib) + for j, (cfg, lib) in enumerate(built_kernels[start:end]) + ] + work_q.put((prob_idx, prob_dict, batch)) + + io_lock = threading.Lock() + stats = {"measurements": 0, "failures": 0} + bench_t0 = time.perf_counter() + + def device_thread(device_id): + while True: + try: + unit = work_q.get_nowait() + except queue.Empty: + return + rows, lines, n_fail = _run_batch_on_device( + device_id, unit, args, worker_path, base_env + ) + with io_lock: + for ln in lines: + print(ln) + for row in rows: + writer.writerow(row) + csv_file.flush() + stats["measurements"] += len(rows) + stats["failures"] += n_fail + work_q.task_done() + + threads = [ + threading.Thread(target=device_thread, args=(d,), daemon=True) for d in devices + ] + for t in threads: + t.start() + for t in threads: + t.join() + + 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" Devices used: {len(devices)}") + print(f" Successful measurements: {stats['measurements']}") + print(f" Failed measurements: {stats['failures']}") + print(f" Output: {csv_path}") + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/projects/composablekernel/tile_engine/ops/gemm/run_one_batched_gemm_kernel.py b/projects/composablekernel/tile_engine/ops/gemm/run_one_batched_gemm_kernel.py new file mode 100644 index 000000000000..1b0bcd4b4663 --- /dev/null +++ b/projects/composablekernel/tile_engine/ops/gemm/run_one_batched_gemm_kernel.py @@ -0,0 +1,152 @@ +#!/usr/bin/env python3 +# Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +# SPDX-License-Identifier: MIT +"""Worker script for running BATCHED GEMM kernels in an isolated subprocess. + +Batched counterpart of run_one_gemm_kernel.py: +- Receives kernel config + batched 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 + +Input JSON format: + Single: {"so_path": "...", + "problem": {"batch_count":.., "M":.., "N":.., "K":..}, + "kernel_name": "..."} + Batch: {"items": [{...}, ...]} + +Optional top-level keys ``verify`` (bool) and ``verify_tol`` (float) enable an +fp32 numpy reference check (per-batch A @ B); when set, each OK result also +carries ``verified`` and ``max_rel``. +""" + +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 batched_gemm_utils import ( # noqa: E402 + BatchedGemmProblem, + GpuBatchedGemmRunner, +) +import numpy as np # noqa: E402 + + +def _run_one(idx, so_path, prob_dict, kernel_name, verify=False, verify_tol=2e-2): + """Run a single batched kernel and emit its result as one JSON line. + + When ``verify`` is set, the batched output is checked against an fp32 numpy + reference (per-batch ``A @ B``) using the global relative metric + ``max|out - ref| / max|ref|``. + """ + try: + problem = BatchedGemmProblem.from_dict(prob_dict) + + # Cache host tensors per (batch, M, N, K) so batch mode doesn't + # regenerate large inputs per kernel. + cache = getattr(_run_one, "_ab_cache", {}) + key = (problem.batch_count, problem.M, problem.N, problem.K) + if key not in cache: + rng = np.random.RandomState(42) + cache[key] = ( + (rng.randn(problem.batch_count, problem.M, problem.K) * 0.1).astype( + np.float32 + ), + (rng.randn(problem.batch_count, problem.K, problem.N) * 0.1).astype( + np.float32 + ), + ) + _run_one._ab_cache = cache + A, B = cache[key] + + # CRITICAL: load the library ONLY inside this subprocess. + runner = GpuBatchedGemmRunner(lib_path=so_path) + result = runner.run(A, B, problem) + + if result.success: + non_zero = ( + int(np.count_nonzero(result.output)) + if result.output is not None + else 0 + ) + out = { + "idx": idx, + "ok": True, + "ms": result.time_ms, + "tflops": result.tflops, + "non_zero": non_zero, + "kernel": kernel_name, + } + if verify: + ref = np.matmul(A.astype(np.float32), B.astype(np.float32)) + got = result.output.astype(np.float32) + denom = float(np.max(np.abs(ref))) or 1.0 + max_rel = float(np.max(np.abs(got - ref)) / denom) + out["max_rel"] = max_rel + out["verified"] = bool(max_rel <= verify_tol) + print(json.dumps(out), 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 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) + + verify = bool(d.get("verify", False)) + verify_tol = float(d.get("verify_tol", 2e-2)) + + if "items" in d: + for i, item in enumerate(d["items"]): + _run_one( + i, + item["so_path"], + item["problem"], + item.get("kernel_name", "unknown"), + verify=verify, + verify_tol=verify_tol, + ) + else: + _run_one( + 0, + d["so_path"], + d["problem"], + d.get("kernel_name", "unknown"), + verify=verify, + verify_tol=verify_tol, + ) + + +if __name__ == "__main__": + main()