From 80aafb95a2583a87e2097d0e6ff49282a1f98dc8 Mon Sep 17 00:00:00 2001 From: Muhammed Emin Ozturk <3836908+ozturkosu@users.noreply.github.com> Date: Sat, 11 Jul 2026 04:01:36 -0400 Subject: [PATCH 1/8] Add gemm_multi_abd to the TileEngine->Dispatcher bridge Bridge the multi-tensor gemm_multi_abd op (fp16, rcrr layout, arrays of A/B/D tensors with per-group element-wise ops) so the Dispatcher generates and launches it at parity with the legacy Tile Engine version. Mirrors the gemm_universal (#8997) bridge, using the grouped/divergent-ABI pattern (#9000) since multi_abd needs array-of-pointer arguments. - unified_gemm_codegen.py: new GemmVariant.MULTI_ABD (tuple As/Bs/Ds types and layouts, three element-wise fns, GemmKernelMultiABD launch, CShuffle epilogue); tensor counts + all three ops folded into key_name()/kernel name; multi_abd_config search block; global tensor-count re-exports for the registry-bypass ctypes lib. - gemm_multi_abd_ctypes_lib.cpp: new registry-bypass C ABI (dispatcher_run_multi_abd) that calls SelectedKernel::launch directly on the force-included kernel with GemmMultiABDHostArgs; owns GPU memory. - gemm_utils.py: multi_abd fields on GemmKernelConfig (byte-exact name), 4-char layout handling, multi_abd codegen json, variant-aware build (own ctypes source, no static lib), MultiABDDispatcherLib + GpuMultiABDRunner, expand_sweep variant support. - gemm_full_benchmark.py / run_one_gemm_kernel.py: gemm_multi_abd variant wired into the 3-phase driver and routed to GpuMultiABDRunner by name. - ctypes CMakeLists: dispatcher_gemm_multi_abd_lib target. Validated on gfx942 (MI300X): fp16/rcrr default_config kernel codegen (0 failures), ctypes .so compiles + links, all C-ABI symbols exported; standard/multi_d/multi_abd codegen regression clean; name parity between GemmKernelConfig.name and the generated header confirmed byte-for-byte. --- .../dispatcher/bindings/ctypes/CMakeLists.txt | 32 ++ .../ctypes/gemm_multi_abd_ctypes_lib.cpp | 338 ++++++++++++++++++ .../codegen/unified_gemm_codegen.py | 218 ++++++++++- .../dispatcher/python/gemm_utils.py | 316 +++++++++++++++- .../ops/gemm/gemm_full_benchmark.py | 26 +- .../ops/gemm/run_one_gemm_kernel.py | 53 ++- 6 files changed, 947 insertions(+), 36 deletions(-) create mode 100644 projects/composablekernel/dispatcher/bindings/ctypes/gemm_multi_abd_ctypes_lib.cpp diff --git a/projects/composablekernel/dispatcher/bindings/ctypes/CMakeLists.txt b/projects/composablekernel/dispatcher/bindings/ctypes/CMakeLists.txt index 18314017f225..dd50c4941af8 100644 --- a/projects/composablekernel/dispatcher/bindings/ctypes/CMakeLists.txt +++ b/projects/composablekernel/dispatcher/bindings/ctypes/CMakeLists.txt @@ -71,6 +71,38 @@ else() target_link_libraries(dispatcher_gemm_lib PRIVATE hip::device) endif() +# ============================================================================= +# GEMM Multi-ABD ctypes Library (divergent, array-pointer ABI; registry-bypass) +# ============================================================================= +# +# gemm_multi_abd takes ARRAYS of A/B/D tensors, so it has its own C ABI +# (dispatcher_run_multi_abd) and calls SelectedKernel::launch directly rather +# than going through the name-keyed dispatcher registry. It is compiled against a +# force-included generated multi_abd kernel header (name carries the +# ``_multiabd_`` suffix). The Python bridge normally builds these .so files at +# runtime via hipcc; this target exists for CMake-driven / CI builds. +file(GLOB GEMM_MULTI_ABD_KERNEL_HEADERS + "${CMAKE_BINARY_DIR}/generated_kernels/gemm_*_multiabd_*.hpp") +if(GEMM_MULTI_ABD_KERNEL_HEADERS) + list(GET GEMM_MULTI_ABD_KERNEL_HEADERS 0 GEMM_MULTI_ABD_KERNEL_HEADER) + message(STATUS "Found GEMM Multi-ABD kernel for ctypes lib: ${GEMM_MULTI_ABD_KERNEL_HEADER}") + + add_ctypes_library(dispatcher_gemm_multi_abd_lib + gemm_multi_abd_ctypes_lib.cpp + KERNEL_HEADER ${GEMM_MULTI_ABD_KERNEL_HEADER} + ) + target_compile_definitions(dispatcher_gemm_multi_abd_lib PRIVATE + CK_TILE_SINGLE_KERNEL_INCLUDE) +else() + message(STATUS "No GEMM Multi-ABD kernel found for ctypes lib - building without kernel") + add_library(dispatcher_gemm_multi_abd_lib SHARED gemm_multi_abd_ctypes_lib.cpp) + target_include_directories(dispatcher_gemm_multi_abd_lib PRIVATE + ${PROJECT_SOURCE_DIR}/include + ${PROJECT_SOURCE_DIR}/dispatcher/include + ) + target_link_libraries(dispatcher_gemm_multi_abd_lib PRIVATE hip::device) +endif() + # ============================================================================= # Convolution ctypes Library (supports forward + bwd_data) # ============================================================================= diff --git a/projects/composablekernel/dispatcher/bindings/ctypes/gemm_multi_abd_ctypes_lib.cpp b/projects/composablekernel/dispatcher/bindings/ctypes/gemm_multi_abd_ctypes_lib.cpp new file mode 100644 index 000000000000..70790597966d --- /dev/null +++ b/projects/composablekernel/dispatcher/bindings/ctypes/gemm_multi_abd_ctypes_lib.cpp @@ -0,0 +1,338 @@ +// Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +// SPDX-License-Identifier: MIT + +/** + * GEMM Multi-ABD Dispatcher ctypes Library + * + * Provides a C API for Python ctypes integration for the gemm_multi_abd op. + * + * WHY A SEPARATE .so (divergent ABI): + * gemm_multi_abd is a multi-tensor op: it takes ARRAYS of A, B and D device + * pointers (NumATensors / NumBTensors / NumDTensors), not the single A/B/C + * triple the regular GEMM ABI (dispatcher_run_gemm) exposes. Its launch takes + * a ck_tile::GemmMultiABDHostArgs, so this library bypasses + * the name-keyed dispatcher registry and calls SelectedKernel::launch(...) + * directly on the force-included kernel -- exactly the divergent-ABI bridge + * pattern used for grouped GEMM (#9000). The kernel header is force-included + * via the -include compiler flag and defines SelectedKernel, KERNEL_NAME, and + * the NumA/NumB/NumD tensor counts (inside the kernel's namespace, re-exported + * to global scope under CK_TILE_SINGLE_KERNEL_INCLUDE). + * + * Usage from Python: + * lib = ctypes.CDLL("libgemm_multi_abd_.so") + * lib.dispatcher_initialize() + * lib.dispatcher_run_multi_abd(as_ptrs, bs_ptrs, ds_ptrs, e_ptr, + * num_a, num_b, num_d, M, N, K, &time_ms) + */ + +#include +#include +#include +#include +#include +#include +#include +#include + +#include "ck_tile/core.hpp" +#include "ck_tile/host.hpp" +#include "ck_tile/ops/gemm.hpp" +#include "ck_tile/ops/gemm/kernel/gemm_multi_abd_kernel.hpp" + +// Kernel header force-included via -include compiler flag. +// Defines (under CK_TILE_SINGLE_KERNEL_INCLUDE): +// SelectedKernel -- the generated Kernel_ struct with a static +// launch(const GemmMultiABDHostArgs<...>&, stream_config&) +// KERNEL_NAME -- the byte-exact runtime kernel name +// NumATensors / NumBTensors / NumDTensors -- tensor counts baked into the type + +// GPU architecture - can be overridden via -DGFX_ARCH="gfx90a" at compile time +#ifndef GFX_ARCH +#define GFX_ARCH "gfx942" +#endif + +// The force-included header exports the tensor counts into the global namespace +// under CK_TILE_SINGLE_KERNEL_INCLUDE. Guard with a fallback so the file is +// still self-describing if the macros ever change. +#ifndef GEMM_MULTI_ABD_NUM_A +#define GEMM_MULTI_ABD_NUM_A NumATensors +#endif +#ifndef GEMM_MULTI_ABD_NUM_B +#define GEMM_MULTI_ABD_NUM_B NumBTensors +#endif +#ifndef GEMM_MULTI_ABD_NUM_D +#define GEMM_MULTI_ABD_NUM_D NumDTensors +#endif + +namespace { + +constexpr ck_tile::index_t kNumA = GEMM_MULTI_ABD_NUM_A; +constexpr ck_tile::index_t kNumB = GEMM_MULTI_ABD_NUM_B; +constexpr ck_tile::index_t kNumD = GEMM_MULTI_ABD_NUM_D; + +using HostArgs = ck_tile::GemmMultiABDHostArgs; + +bool g_initialized = false; + +} // namespace + +extern "C" { + +/** + * Initialize the library. Multi-ABD is registry-bypass (it calls + * SelectedKernel::launch directly), so there is no registry to populate here; + * this just flips the ready flag and is kept for ABI symmetry with the regular + * GEMM ctypes lib. + * + * Returns: 0 on success. + */ +int dispatcher_initialize() +{ + g_initialized = true; + return 0; +} + +int dispatcher_init() { return dispatcher_initialize(); } + +/** + * Number of tensors the force-included kernel expects (compile-time constants). + */ +int dispatcher_get_num_a_tensors() { return static_cast(kNumA); } +int dispatcher_get_num_b_tensors() { return static_cast(kNumB); } +int dispatcher_get_num_d_tensors() { return static_cast(kNumD); } + +/** + * Kernel name (single-kernel ABI): the byte-exact KERNEL_NAME baked into the + * force-included header. Mirrors dispatcher_get_kernel_name in the regular lib. + */ +const char* dispatcher_get_kernel_name() { return KERNEL_NAME; } + +/** + * Multi-kernel ABI shim: this .so holds exactly one force-included kernel, so + * index 0 returns KERNEL_NAME and any other index is out of range. Kept so the + * Python GemmDispatcherLib multi-kernel path works uniformly across variants. + */ +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 multi-ABD GEMM on the GPU via the force-included kernel. + * + * The Python runner hands HOST buffers (contiguous, already laid out for the + * kernel's compiled layout); this shim owns the GPU side -- it hipMallocs one + * device buffer per A/B/D tensor and for E, copies inputs up, launches, and + * copies the result back. This mirrors the regular gemm_ctypes_lib's + * host-pointer contract so GpuGemmRunner-style callers just pass numpy arrays. + * + * as_hosts / bs_hosts / ds_hosts : arrays of host pointers (length num_a/b/d). + * e_host : output host pointer (M*N * sizeof(EDataType)). + * elem_a/b/d/e : element size in bytes for each group's dtype + * (so this shim need not know the CK dtype). + * stride_* : per-operand leading stride; a value <= 0 lets + * the kernel derive the default from M,N,K and + * the compiled layout. + * time_ms : filled with the average kernel time (ms). + * + * Returns 0 on success, negative on error. num_a/num_b/num_d MUST equal the + * kernel's compiled tensor counts or the call is rejected (-3). + */ +int dispatcher_run_multi_abd(const void** as_hosts, + const void** bs_hosts, + const void** ds_hosts, + void* e_host, + const int64_t* stride_as, + const int64_t* stride_bs, + const int64_t* stride_ds, + int64_t stride_e, + int elem_a, + int elem_b, + int elem_d, + int elem_e, + int num_a, + int num_b, + int num_d, + int64_t M, + int64_t N, + int64_t K, + float* time_ms) +{ + if(!g_initialized || !as_hosts || !bs_hosts || !e_host) + { + return -1; + } + if(M <= 0 || N <= 0 || K <= 0 || elem_a <= 0 || elem_b <= 0 || elem_e <= 0) + { + return -1; + } + // The tensor counts are baked into the kernel type at compile time; a + // mismatch would silently read past the caller's arrays, so reject it. + if(num_a != static_cast(kNumA) || num_b != static_cast(kNumB) || + num_d != static_cast(kNumD)) + { + return -3; + } + if(kNumD > 0 && (!ds_hosts || elem_d <= 0)) + { + return -1; + } + + const size_t a_bytes = static_cast(M) * K * elem_a; + const size_t b_bytes = static_cast(K) * N * elem_b; + const size_t d_bytes = static_cast(M) * N * elem_d; + const size_t e_bytes = static_cast(M) * N * elem_e; + + std::vector a_dev(kNumA, nullptr), b_dev(kNumB, nullptr), d_dev(kNumD, nullptr); + void* e_dev = nullptr; + + auto cleanup = [&]() { + for(auto p : a_dev) + if(p) + (void)hipFree(p); + for(auto p : b_dev) + if(p) + (void)hipFree(p); + for(auto p : d_dev) + if(p) + (void)hipFree(p); + if(e_dev) + (void)hipFree(e_dev); + }; + + // Allocate + upload each operand tensor. + for(int i = 0; i < num_a; ++i) + { + if(hipMalloc(&a_dev[i], a_bytes) != hipSuccess) + { + cleanup(); + return -1; + } + if(hipMemcpy(a_dev[i], as_hosts[i], a_bytes, hipMemcpyHostToDevice) != hipSuccess) + { + cleanup(); + return -1; + } + } + for(int i = 0; i < num_b; ++i) + { + if(hipMalloc(&b_dev[i], b_bytes) != hipSuccess) + { + cleanup(); + return -1; + } + if(hipMemcpy(b_dev[i], bs_hosts[i], b_bytes, hipMemcpyHostToDevice) != hipSuccess) + { + cleanup(); + return -1; + } + } + for(int i = 0; i < num_d; ++i) + { + if(hipMalloc(&d_dev[i], d_bytes) != hipSuccess) + { + cleanup(); + return -1; + } + if(hipMemcpy(d_dev[i], ds_hosts[i], d_bytes, hipMemcpyHostToDevice) != hipSuccess) + { + cleanup(); + return -1; + } + } + if(hipMalloc(&e_dev, e_bytes) != hipSuccess) + { + cleanup(); + return -1; + } + if(hipMemset(e_dev, 0, e_bytes) != hipSuccess) + { + cleanup(); + return -1; + } + + // Pack the device pointers / strides into the std::array shapes the + // GemmMultiABDHostArgs constructor requires. + std::array as{}; + std::array bs{}; + std::array ds{}; + std::array str_as{}; + std::array str_bs{}; + std::array str_ds{}; + + for(ck_tile::index_t i = 0; i < kNumA; ++i) + { + as[i] = a_dev[i]; + str_as[i] = stride_as ? static_cast(stride_as[i]) : 0; + } + for(ck_tile::index_t i = 0; i < kNumB; ++i) + { + bs[i] = b_dev[i]; + str_bs[i] = stride_bs ? static_cast(stride_bs[i]) : 0; + } + for(ck_tile::index_t i = 0; i < kNumD; ++i) + { + ds[i] = d_dev[i]; + str_ds[i] = stride_ds ? static_cast(stride_ds[i]) : 0; + } + + // Multi-ABD supports only k_batch = 1. + HostArgs args{as, + bs, + ds, + e_dev, + /*k_batch=*/1, + static_cast(M), + static_cast(N), + static_cast(K), + str_as, + str_bs, + str_ds, + static_cast(stride_e)}; + + float exec_time = -1.0f; + try + { + // Registry bypass: launch the force-included kernel directly. The + // launch() returns a single float average time (unlike some Old-TE + // launchers that return a tuple), so no tuple normalization is needed. + ck_tile::stream_config stream{nullptr, /*time_kernel=*/true}; + exec_time = SelectedKernel::launch(args, stream); + } + catch(const std::exception& e) + { + std::cerr << "gemm_multi_abd launch failed: " << e.what() << std::endl; + cleanup(); + if(time_ms) + { + *time_ms = -1.0f; + } + return -2; + } + + // Copy result back to the host output buffer. + if(hipMemcpy(e_host, e_dev, e_bytes, hipMemcpyDeviceToHost) != hipSuccess) + { + cleanup(); + return -1; + } + + cleanup(); + if(time_ms) + { + *time_ms = exec_time; + } + 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 47cb3680f152..4214aef598e3 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" + MULTI_ABD = "multi_abd" # TileConfig imported from codegen_common @@ -228,6 +229,15 @@ class KernelConfig: num_d_tensors: int = 0 d_layout: str = "r" # Layout for D tensors (r=row, c=col) - same for all D tensors + # Multi-ABD variant: arrays of A/B tensors and per-group element-wise ops. + # These are behavior-affecting and MUST participate in key_name() so distinct + # tensor counts / elementwise ops never alias to the same kernel. + num_a_tensors: int = 1 + num_b_tensors: int = 1 + a_elementwise_op: str = "PassThrough" + b_elementwise_op: str = "PassThrough" + cde_elementwise_op: str = "PassThrough" + # Fixed parameters block_size: int = 256 k_block_per_cu: int = 1 @@ -289,6 +299,18 @@ def key_name(self, datatype: str, layout: str) -> str: parts.append(f"nd{self.num_d_tensors}") parts.append(f"dly_{self.d_layout}") + # Multi-ABD variant: include per-group tensor counts, all three + # element-wise ops, and the D layout. Every one of these changes the + # generated kernel's types, so they must be in the unique key. + if self.variant == GemmVariant.MULTI_ABD: + parts.append(f"na{self.num_a_tensors}") + parts.append(f"nb{self.num_b_tensors}") + parts.append(f"nd{self.num_d_tensors}") + parts.append(f"aew_{self.a_elementwise_op}") + parts.append(f"bew_{self.b_elementwise_op}") + parts.append(f"cdew_{self.cde_elementwise_op}") + parts.append(f"dly_{self.d_layout}") + # Occupancy parameters (only if non-default) if self.num_wave_groups != 1: parts.append(f"wg{self.num_wave_groups}") @@ -324,8 +346,10 @@ def generate(config: KernelConfig, datatype: str, layout: str) -> str: t = config.tile tr = config.trait - # For multi-d, use 4-char layout (abcd), otherwise use 3-char layout (abc) - if config.variant == GemmVariant.MULTI_D: + # For multi-d / multi-abd, use 4-char layout (abcd), otherwise 3-char (abc). + # For multi-abd the 4th char is the D layout (A,B,E,D convention); the + # incoming ``layout`` is the 3-char A,B,E slice. + if config.variant in (GemmVariant.MULTI_D, GemmVariant.MULTI_ABD): full_layout = layout + config.d_layout # e.g., "rcr" + "r" = "rcrr" else: full_layout = layout @@ -344,6 +368,15 @@ 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.MULTI_ABD: + # Encode every behavior-affecting multi-abd parameter so distinct + # tensor counts / elementwise ops get distinct kernel names. + name += ( + f"_multiabd_a{config.num_a_tensors}_b{config.num_b_tensors}" + f"_d{config.num_d_tensors}" + f"_{config.a_elementwise_op}_{config.b_elementwise_op}" + f"_{config.cde_elementwise_op}" + ) return name @@ -391,6 +424,14 @@ def _header(self, kernel_name: str, config: KernelConfig) -> str: includes += """ #include "ck_tile/ops/elementwise/unary_element_wise_operation.hpp" #include "ck_tile/ops/gemm/kernel/gemm_multi_d_kernel.hpp" +""" + + if config.variant == GemmVariant.MULTI_ABD: + includes += """ +#include +#include +#include "ck_tile/ops/elementwise/unary_element_wise_operation.hpp" +#include "ck_tile/ops/gemm/kernel/gemm_multi_abd_kernel.hpp" """ if config.preshuffle: @@ -446,6 +487,72 @@ def _multi_d_types(self, config: KernelConfig) -> str: using ElementWiseFn = element_wise::{config.elementwise_op}; static constexpr index_t NumDTensor = {config.num_d_tensors}; using GemmMultiDArgs = GemmMultiDHostArgs; +""" + + def _multi_abd_types(self, config: KernelConfig) -> str: + """Generate multi-ABD type definitions (inside namespace to avoid conflicts). + + Multi-ABD uses tuples of A, B and D tensors (each group homogeneous in + dtype/layout here, matching the Tile Engine op) plus three separate + element-wise functions (A, B, CDE). The layouts mirror the TE builder's + 4-char ``rcrr`` convention: A=layout[0], B=layout[1], E(=C)=layout[2], + D=config.d_layout. + """ + if config.variant != GemmVariant.MULTI_ABD: + return "" + + a_layout_ck = self.tm.LAYOUT_TO_CK[self.layout[0]] + b_layout_ck = self.tm.LAYOUT_TO_CK[self.layout[1]] + e_layout_ck = self.tm.LAYOUT_TO_CK[self.layout[2]] + d_layout_ck = self.tm.LAYOUT_TO_CK[config.d_layout] + + a_types = ", ".join(["ADataType"] * config.num_a_tensors) + b_types = ", ".join(["BDataType"] * config.num_b_tensors) + d_types = ", ".join(["CDataType"] * config.num_d_tensors) + a_layouts = ", ".join([a_layout_ck] * config.num_a_tensors) + b_layouts = ", ".join([b_layout_ck] * config.num_b_tensors) + d_layouts = ", ".join([d_layout_ck] * config.num_d_tensors) + + return f""" +// Multi-ABD types (defined in namespace to avoid conflicts). +// EDataType is the output type (aliased to CDataType by the standard path). +using EDataType = CDataType; +using AsDataType = tuple<{a_types}>; +using BsDataType = tuple<{b_types}>; +using DsDataType = tuple<{d_types}>; +using AsLayout = tuple<{a_layouts}>; +using BsLayout = tuple<{b_layouts}>; +using ELayout = {e_layout_ck}; +using DsLayout = tuple<{d_layouts}>; +static constexpr index_t NumATensors = {config.num_a_tensors}; +static constexpr index_t NumBTensors = {config.num_b_tensors}; +static constexpr index_t NumDTensors = {config.num_d_tensors}; +using AElementWiseFn = element_wise::{config.a_elementwise_op}; +using BElementWiseFn = element_wise::{config.b_elementwise_op}; +using CDEElementWiseFn = element_wise::{config.cde_elementwise_op}; +using GemmMultiABDArgs = + GemmMultiABDHostArgs; +""" + + def _multi_abd_global_exports(self, config: KernelConfig, ns_name: str) -> str: + """Re-export the multi-ABD tensor counts to the global namespace. + + The gemm_multi_abd ctypes lib is registry-bypass: it force-includes this + header and constructs GemmMultiABDHostArgs at global scope. Those counts live inside the kernel's + namespace, so under CK_TILE_SINGLE_KERNEL_INCLUDE we surface them (and a + few convenience macros) globally. Empty for non-multi-abd variants. + """ + if config.variant != GemmVariant.MULTI_ABD: + return "" + return f""" +// Multi-ABD tensor counts (re-exported for the registry-bypass ctypes lib). +constexpr index_t NumATensors = {ns_name}::NumATensors; +constexpr index_t NumBTensors = {ns_name}::NumBTensors; +constexpr index_t NumDTensors = {ns_name}::NumDTensors; +#define GEMM_MULTI_ABD_NUM_A {config.num_a_tensors} +#define GEMM_MULTI_ABD_NUM_B {config.num_b_tensors} +#define GEMM_MULTI_ABD_NUM_D {config.num_d_tensors} """ def _selected_kernel_struct(self, config: KernelConfig, kernel_name: str) -> str: @@ -460,7 +567,7 @@ def _selected_kernel_struct(self, config: KernelConfig, kernel_name: str) -> str # Create valid C++ namespace name (replace invalid chars) ns_name = "ns_" + kernel_name.replace("-", "_") - multi_d_types = self._multi_d_types(config) + multi_d_types = self._multi_d_types(config) + self._multi_abd_types(config) return f""" namespace {ns_name} {{ @@ -526,7 +633,7 @@ def _selected_kernel_struct(self, config: KernelConfig, kernel_name: str) -> str using BDataType = {self.tm.DTYPE_TO_CK_QUALIFIED[self.datatype]}; using CDataType = {self.tm.DTYPE_TO_CK_QUALIFIED[output_dtype]}; using AccDataType = {self.tm.DTYPE_TO_CK_QUALIFIED[acc_dtype]}; - +{self._multi_abd_global_exports(config, ns_name)} // KernelKey field descriptors for the force-included kernel. // The ctypes library builds the registry KernelKey from these so the // registered entry reflects this kernel's real traits (not a hard-coded @@ -586,6 +693,8 @@ def _tile_types(self, config: KernelConfig, ns_name: str) -> str: def _launch_function(self, config: KernelConfig) -> str: """Generate launch function""" + if config.variant == GemmVariant.MULTI_ABD: + return self._launch_function_multi_abd(config) if config.variant == GemmVariant.MULTI_D: return self._launch_function_multi_d(config) if config.preshuffle: @@ -768,8 +877,65 @@ def _launch_function_multi_d(self, config: KernelConfig) -> str: return launch(multi_d_args, stream); }}""" + def _launch_function_multi_abd(self, config: KernelConfig) -> str: + """Generate launch function for Multi-ABD GEMM. + + Mirrors the Tile Engine gemm_multi_abd instance builder's ``launch``: + tuple A/B/E layouts in the traits, tuple A/B dtypes plus the A and B + element-wise functions in the pipeline problem, and GemmKernelMultiABD + as the kernel. Multi-ABD supports only k_batch = 1, so it launches the + kernel directly (no hot-loop tail handler, matching Old-TE). + """ + return f""" + // Multi-ABD launch function - takes GemmMultiABDHostArgs with tuple A/B/D. + static float launch(const GemmMultiABDArgs& args, const stream_config& stream) {{ + float ave_time{{0}}; + + constexpr auto scheduler = {self.tm.SCHEDULER_TO_CK[config.trait.scheduler]}; + + // Traits use tuple layouts for multi-abd (AsLayout/BsLayout/ELayout). + using Traits = TileGemmUniversalTraits; + + using UniversalGemmProblem = UniversalGemmPipelineProblem< + AsDataType, BsDataType, AccDataType, TileShape, Traits, scheduler, + AElementWiseFn, BElementWiseFn>; + + using GemmPipeline = {self.tm.PIPELINE_TO_CK[config.trait.pipeline]}; + {self._epilogue_code(config)} + + using GemmKernel = ck_tile::GemmKernelMultiABD; + + auto kargs = GemmKernel::MakeKernelArgs(args); + + if (!GemmKernel::IsSupportedArgument(kargs)) {{ + throw std::runtime_error("Arguments not supported! Multi-ABD only supports k_batch = 1"); + }} + + const dim3 grids = GemmKernel::GridSize(args.M, args.N, args.k_batch); + const dim3 blocks = GemmKernel::BlockSize(); + + constexpr int kBlockPerCu = {config.k_block_per_cu}; + ave_time = launch_kernel(stream, + make_kernel(GemmKernel{{}}, grids, blocks, 0, kargs)); + + return ave_time; + }}""" + def _epilogue_code(self, config: KernelConfig) -> str: """Generate epilogue code""" + 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 + # CShuffleEpilogueProblem for gemm_multi_abd. + return """ + using EpilogueProblem = CShuffleEpilogueProblem< + AsDataType, BsDataType, DsDataType, AccDataType, EDataType, + DsLayout, ELayout, CDEElementWiseFn, + TilePartitioner::MPerBlock, TilePartitioner::NPerBlock, + WarpPerBlock_M, WarpPerBlock_N, WarpTileM, WarpTileN, WarpTileK, + TransposeC>; + using GemmEpilogue = CShuffleEpilogue;""" if config.variant == GemmVariant.MULTI_D: return """ using EpilogueProblem = CShuffleEpilogueProblem< @@ -985,6 +1151,16 @@ def _load_config(self, config_file: Optional[Path]) -> Dict: "elementwise_ops": ["MultiDAdd", "MultiDMultiply"], "num_d_tensors": [1, 2], }, + "multi_abd_config": { + # Match the Tile Engine gemm_multi_abd instance builder defaults: + # 2 A tensors, 2 B tensors, 2 D tensors, all PassThrough ops. + "num_a_tensors": 2, + "num_b_tensors": 2, + "num_d_tensors": 2, + "a_elementwise_op": "PassThrough", + "b_elementwise_op": "PassThrough", + "cde_elementwise_op": "PassThrough", + }, } def generate_all(self, parallel: bool = True) -> Dict: @@ -1170,6 +1346,34 @@ def _get_configs_for_variant(self, variant: GemmVariant) -> List[KernelConfig]: ) ) + elif variant == GemmVariant.MULTI_ABD: + # Multi-ABD always uses the CShuffle epilogue (it fuses the D + # tensors + CDE elementwise through LDS), so the CShuffle-store + # power-of-two repeat correctness gate always applies. + if not self._cshuffle_repeat_ok(tile): + continue + multi_abd = self.config.get("multi_abd_config", {}) + num_a = multi_abd.get("num_a_tensors", 2) + num_b = multi_abd.get("num_b_tensors", 2) + num_d = multi_abd.get("num_d_tensors", 2) + a_ew = multi_abd.get("a_elementwise_op", "PassThrough") + b_ew = multi_abd.get("b_elementwise_op", "PassThrough") + cde_ew = multi_abd.get("cde_elementwise_op", "PassThrough") + configs.append( + KernelConfig( + tile=tile, + trait=trait, + variant=variant, + num_a_tensors=num_a, + num_b_tensors=num_b, + num_d_tensors=num_d, + a_elementwise_op=a_ew, + b_elementwise_op=b_ew, + cde_elementwise_op=cde_ew, + d_layout=self.d_layout, + ) + ) + return configs def _get_tile_configs(self) -> List[TileConfig]: @@ -1276,6 +1480,10 @@ def _is_tile_arch_valid( GemmVariant.STANDARD: OperatorType.GEMM, GemmVariant.PRESHUFFLE: OperatorType.GEMM_PRESHUFFLE, GemmVariant.MULTI_D: OperatorType.GEMM_MULTI_D, + # Multi-ABD shares Multi-D's XDL/tile constraints (same + # UniversalGemm pipeline + CShuffle epilogue), so validate it + # against the GEMM_MULTI_D operator rules. + GemmVariant.MULTI_ABD: OperatorType.GEMM_MULTI_D, } operator = variant_to_operator.get(variant, OperatorType.GEMM) @@ -1525,7 +1733,7 @@ def main(): parser.add_argument( "--variants", nargs="+", - choices=["standard", "preshuffle", "multi_d"], + choices=["standard", "preshuffle", "multi_d", "multi_abd"], default=["standard"], help="Variants to generate", ) diff --git a/projects/composablekernel/dispatcher/python/gemm_utils.py b/projects/composablekernel/dispatcher/python/gemm_utils.py index 160c0f0694bc..ed5a5f614946 100644 --- a/projects/composablekernel/dispatcher/python/gemm_utils.py +++ b/projects/composablekernel/dispatcher/python/gemm_utils.py @@ -106,6 +106,19 @@ class GemmKernelConfig: gfx_arch: str = "gfx942" variant: str = "standard" + # --- Multi-ABD only ---------------------------------------------------- + # Arrays of A/B/D tensors and per-group element-wise ops. These are + # behavior-affecting and appear in .name (and thus in the codegen kernel + # name) so distinct tensor counts / ops never collapse to one kernel. + # layout_d is the 4th ('D') char of the multi_abd rcrr layout code. + num_a_tensors: int = 2 + num_b_tensors: int = 2 + num_d_tensors: int = 2 + a_elementwise_op: str = "PassThrough" + b_elementwise_op: str = "PassThrough" + cde_elementwise_op: str = "PassThrough" + layout_d: str = "row" + # ------------------------------------------------------------------ # # Derived string fragments # ------------------------------------------------------------------ # @@ -118,6 +131,11 @@ def layout(self) -> str: + _LAYOUT_CHAR[self.layout_c] ) + @property + def layout4(self) -> str: + """4-char multi_abd layout string (A,B,E,D), e.g. 'rcrr'.""" + return self.layout + _LAYOUT_CHAR[self.layout_d] + @property def tile_str(self) -> str: return f"{self.tile_m}x{self.tile_n}x{self.tile_k}" @@ -140,8 +158,11 @@ def name(self) -> str: ``dispatcher_get_kernel_name``). This is the single thread tying config -> codegen -> runtime together. """ + # Multi-ABD uses the 4-char layout (A,B,E,D); every other variant the + # 3-char (A,B,C). This mirrors KernelNaming.generate in the codegen. + layout_str = self.layout4 if self.variant == "multi_abd" else self.layout name = ( - f"gemm_{self.dtype_a}_{self.layout}" + f"gemm_{self.dtype_a}_{layout_str}" f"_{self.pipeline}_{self.epilogue}_{self.scheduler}" f"_{_cap(self.pad_m)}_{_cap(self.pad_n)}_{_cap(self.pad_k)}" f"_{_cap(self.persistent)}" @@ -151,6 +172,15 @@ def name(self) -> str: name += "_preshuffle" elif self.variant == "streamk": name += "_streamk" + elif self.variant == "multi_abd": + # Byte-for-byte match to codegen KernelNaming.generate's multiabd + # suffix: tensor counts then the three element-wise ops. + name += ( + f"_multiabd_a{self.num_a_tensors}_b{self.num_b_tensors}" + f"_d{self.num_d_tensors}" + f"_{self.a_elementwise_op}_{self.b_elementwise_op}" + f"_{self.cde_elementwise_op}" + ) return name # ------------------------------------------------------------------ # @@ -163,7 +193,7 @@ def to_codegen_json(self) -> Dict[str, Any]: triple ``warp_*`` and the MFMA triple ``warp_tile_*``. We translate from dispatcher semantics here so the mapping cannot drift. """ - return { + cfg = { "tile_config": { "tile_m": [self.tile_m], "tile_n": [self.tile_n], @@ -187,6 +217,19 @@ def to_codegen_json(self) -> Dict[str, Any]: "persistent": [self.persistent], }, } + # Multi-ABD codegen reads its tensor counts / element-wise ops from a + # dedicated ``multi_abd_config`` block. These are scalars (one kernel per + # config), matching the codegen's _get_configs_for_variant reader. + if self.variant == "multi_abd": + cfg["multi_abd_config"] = { + "num_a_tensors": self.num_a_tensors, + "num_b_tensors": self.num_b_tensors, + "num_d_tensors": self.num_d_tensors, + "a_elementwise_op": self.a_elementwise_op, + "b_elementwise_op": self.b_elementwise_op, + "cde_elementwise_op": self.cde_elementwise_op, + } + return cfg def to_dict(self) -> Dict[str, Any]: return { @@ -619,6 +662,211 @@ def run( ) +# ============================================================================ +# Multi-ABD ctypes ABI wrapper + runner (divergent, array-pointer ABI) +# ============================================================================ + +# Element size (bytes) per CK dtype -- mirrors the codegen's ELEMENT_SIZE_MAP and +# lets the ctypes shim size its device buffers without knowing the CK type. +_ELEM_BYTES = {"fp16": 2, "bf16": 2, "fp32": 4, "fp8": 1, "bf8": 1, "int8": 1, "int32": 4} + + +class MultiABDDispatcherLib: + """Thin ctypes wrapper around a compiled gemm_multi_abd dispatcher .so. + + Multi-ABD is registry-bypass with a divergent ABI: ``dispatcher_run_multi_abd`` + takes ARRAYS of host pointers (one per A/B/D tensor) plus per-group element + sizes, and the .so owns all GPU memory (hipMalloc/Memcpy/Free) internally. + """ + + 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_name.argtypes = [] + lib.dispatcher_get_kernel_name.restype = ctypes.c_char_p + for fn in ( + "dispatcher_get_num_a_tensors", + "dispatcher_get_num_b_tensors", + "dispatcher_get_num_d_tensors", + ): + getattr(lib, fn).argtypes = [] + getattr(lib, fn).restype = ctypes.c_int + lib.dispatcher_run_multi_abd.argtypes = [ + ctypes.POINTER(ctypes.c_void_p), # as_hosts + ctypes.POINTER(ctypes.c_void_p), # bs_hosts + ctypes.POINTER(ctypes.c_void_p), # ds_hosts + ctypes.c_void_p, # e_host + ctypes.POINTER(ctypes.c_int64), # stride_as + ctypes.POINTER(ctypes.c_int64), # stride_bs + ctypes.POINTER(ctypes.c_int64), # stride_ds + ctypes.c_int64, # stride_e + ctypes.c_int, # elem_a + ctypes.c_int, # elem_b + ctypes.c_int, # elem_d + ctypes.c_int, # elem_e + ctypes.c_int, # num_a + ctypes.c_int, # num_b + ctypes.c_int, # num_d + ctypes.c_int64, # M + ctypes.c_int64, # N + ctypes.c_int64, # K + ctypes.POINTER(ctypes.c_float), # time_ms + ] + lib.dispatcher_run_multi_abd.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 + + @property + def kernel_name(self) -> str: + raw = self._lib.dispatcher_get_kernel_name() + return raw.decode("utf-8") if raw else "unknown" + + @property + def tensor_counts(self) -> Tuple[int, int, int]: + return ( + int(self._lib.dispatcher_get_num_a_tensors()), + int(self._lib.dispatcher_get_num_b_tensors()), + int(self._lib.dispatcher_get_num_d_tensors()), + ) + + def run( + self, + as_arrays: List[np.ndarray], + bs_arrays: List[np.ndarray], + ds_arrays: List[np.ndarray], + e_array: np.ndarray, + M: int, + N: int, + K: int, + elem_a: int, + elem_b: int, + elem_d: int, + elem_e: int, + ) -> Tuple[int, float]: + def _ptr_array(arrays): + arr = (ctypes.c_void_p * max(len(arrays), 1))() + for i, a in enumerate(arrays): + arr[i] = a.ctypes.data_as(ctypes.c_void_p) + return arr + + as_ptrs = _ptr_array(as_arrays) + bs_ptrs = _ptr_array(bs_arrays) + ds_ptrs = _ptr_array(ds_arrays) + # Strides <= 0 tell the kernel to use its default layout-derived stride. + null_i64 = ctypes.POINTER(ctypes.c_int64)() + time_ms = ctypes.c_float(0.0) + status = self._lib.dispatcher_run_multi_abd( + as_ptrs, + bs_ptrs, + ds_ptrs, + e_array.ctypes.data_as(ctypes.c_void_p), + null_i64, + null_i64, + null_i64, + 0, + elem_a, + elem_b, + elem_d, + elem_e, + len(as_arrays), + len(bs_arrays), + len(ds_arrays), + M, + N, + K, + ctypes.byref(time_ms), + ) + return status, time_ms.value + + def cleanup(self) -> None: + self._lib.dispatcher_cleanup() + + +class GpuMultiABDRunner: + """High-level multi-ABD runner: construct from a .so, call run(problem). + + All A/B/D operands share the group dtype/layout (matching the Tile Engine + gemm_multi_abd op). The runner builds the host operand buffers in the + kernel's element dtype + layout and hands raw pointer arrays to the .so, + which owns GPU memory. fp16 is the only supported multi-abd dtype. + """ + + def __init__(self, lib_path: Path): + self.lib = MultiABDDispatcherLib(lib_path) + if not self.lib.initialize(): + raise RuntimeError(f"Failed to initialize multi_abd .so: {lib_path}") + self._kernel_name = self.lib.kernel_name + self._num_a, self._num_b, self._num_d = self.lib.tensor_counts + + @property + def kernel_name(self) -> str: + return self._kernel_name + + def run(self, problem: GemmProblem, seed: int = 0) -> GemmResult: + M, N, K = problem.M, problem.N, problem.K + dtype = _dtype_from_kernel_name(self._kernel_name) + # Kernel name for multi_abd is gemm__<4char-layout>_...; take the + # A/B/D layout chars (E is the C position, index 2). + parts = self._kernel_name.split("_") + layout4 = parts[2] if len(parts) > 2 else "rcrr" + la, lb = layout4[0], layout4[1] + ld = layout4[3] if len(layout4) >= 4 else "r" + + rng = np.random.default_rng(seed) + + def _mk(rows, cols, layout_char, n_tensors, lo, hi): + outs = [] + for _ in range(n_tensors): + x = rng.uniform(lo, hi, size=(rows, cols)).astype(np.float32) + x_lay = x if layout_char == "r" else x.T + if dtype == "fp16": + outs.append(np.ascontiguousarray(x_lay, dtype=np.float16)) + else: + raise ValueError(f"multi_abd runner supports fp16 only, got {dtype}") + return outs + + as_arrays = _mk(M, K, la, self._num_a, -5.0, 5.0) + bs_arrays = _mk(K, N, lb, self._num_b, -5.0, 5.0) + ds_arrays = _mk(M, N, ld, self._num_d, -1.0, 1.0) + elem = _ELEM_BYTES.get(dtype, 2) + e_array = np.zeros((M, N), dtype=np.float16) + + status, time_ms = self.lib.run( + as_arrays, + bs_arrays, + ds_arrays, + e_array, + M, + N, + K, + elem, + elem, + elem, + elem, + ) + tflops = (problem.flops / (time_ms * 1e-3)) / 1e12 if time_ms > 0 else 0.0 + return GemmResult( + output=e_array, + time_ms=time_ms, + status=status, + tflops=tflops, + kernel_name=self._kernel_name, + ) + + # ============================================================================ # Build API: codegen + hipcc -> .so paths (no GPU) # ============================================================================ @@ -705,7 +953,13 @@ def _build_compile_jobs( ck_root = root.parent build_dir = _cu.get_build_dir() output_dir = _cu.get_generated_kernels_dir() - ctypes_source = root / "bindings" / "ctypes" / "gemm_ctypes_lib.cpp" + # Multi-ABD has a divergent (array-pointer) ABI and is registry-bypass, so it + # is compiled against its own ctypes lib (dispatcher_run_multi_abd) instead of + # the regular single-pointer gemm_ctypes_lib.cpp. + if config.variant == "multi_abd": + ctypes_source = root / "bindings" / "ctypes" / "gemm_multi_abd_ctypes_lib.cpp" + else: + ctypes_source = root / "bindings" / "ctypes" / "gemm_ctypes_lib.cpp" static_lib = build_dir / "libck_tile_dispatcher.a" lib_path = build_dir / "examples" / f"lib{config.name}.so" @@ -745,10 +999,14 @@ def _build_compile_jobs( f"--offload-arch={config.gfx_arch}", "--hip-link", str(obj_file), - str(static_lib), - "-o", - str(lib_path), ] + # The regular GEMM ABI goes through the dispatcher registry and must link the + # dispatcher static lib. Multi-ABD is registry-bypass (it calls + # SelectedKernel::launch directly), so it needs only the force-included + # kernel -- no static lib -- keeping its .so self-contained. + if config.variant != "multi_abd": + link_cmd.append(str(static_lib)) + link_cmd += ["-o", str(lib_path)] job = {"compile_cmd": compile_cmd, "link_cmd": link_cmd, "lib_path": str(lib_path)} return job, lib_path @@ -786,13 +1044,24 @@ def setup_multiple_gemm_dispatchers( codegen_script = _cu.get_codegen_path() output_dir = _cu.get_generated_kernels_dir() static_lib = _cu.get_build_dir() / "libck_tile_dispatcher.a" - ctypes_source = ( - _cu.get_dispatcher_root() / "bindings" / "ctypes" / "gemm_ctypes_lib.cpp" + ctypes_dir = _cu.get_dispatcher_root() / "bindings" / "ctypes" + # Multi-ABD is registry-bypass: it links only the force-included kernel, so it + # needs its own ctypes source but NOT the dispatcher static lib. Every other + # variant goes through the registry and requires both. + variants = {c.variant for c in configs} + all_multi_abd = variants == {"multi_abd"} + ctypes_source = ctypes_dir / ( + "gemm_multi_abd_ctypes_lib.cpp" if all_multi_abd else "gemm_ctypes_lib.cpp" ) - if not static_lib.exists() or not ctypes_source.exists(): + if not ctypes_source.exists(): raise FileNotFoundError( - "Missing static lib or ctypes source required for compilation:\n" - f" {static_lib}\n {ctypes_source}\n" + f"Missing ctypes source required for compilation:\n {ctypes_source}\n" + "Build the dispatcher first (cmake + make)." + ) + if not all_multi_abd and not static_lib.exists(): + raise FileNotFoundError( + "Missing dispatcher static lib required for compilation:\n" + f" {static_lib}\n" "Build the dispatcher first (cmake + make)." ) @@ -807,7 +1076,9 @@ def setup_multiple_gemm_dispatchers( "codegen_script": str(codegen_script), "output_dir": str(output_dir), "dtype": c.dtype_a, - "layout": c.layout, + # Multi-ABD codegen expects the 4-char (A,B,E,D) layout so it can + # split off the D layout; every other variant uses 3-char. + "layout": c.layout4 if c.variant == "multi_abd" else c.layout, "gpu_target": c.gfx_arch, "tile_config_json": c.to_codegen_json(), "hpp_glob_pattern": f"{c.name}.hpp", @@ -915,6 +1186,13 @@ def expand_sweep( arch: str, dtype: str = "fp16", layout: str = "rcr", + variant: str = "standard", + num_a_tensors: int = 2, + num_b_tensors: int = 2, + num_d_tensors: int = 2, + a_elementwise_op: str = "PassThrough", + b_elementwise_op: str = "PassThrough", + cde_elementwise_op: str = "PassThrough", ) -> List[GemmKernelConfig]: """Expand a Tile Engine GEMM JSON sweep config into GemmKernelConfig list. @@ -926,6 +1204,10 @@ def expand_sweep( The signature is controlled by the `dtype` and `layout` arguments (defaults to fp16 / rcr). + + For ``variant='multi_abd'`` the ``layout`` is the 4-char (A,B,E,D) code + (e.g. ``rcrr``); the tensor counts and per-group element-wise ops are carried + onto every produced config so they participate in the kernel name. """ with open(config_path) as f: cfg = json.load(f) @@ -952,6 +1234,8 @@ def expand_sweep( persistents = _expand_values(tr.get("persistent"), [False]) la, lb, lc = layout[0], layout[1], layout[2] + # Multi-ABD carries a 4th (D) layout char; default D to C's layout otherwise. + ld = layout[3] if (variant == "multi_abd" and len(layout) >= 4) else lc configs: List[GemmKernelConfig] = [] seen: set = set() @@ -1015,6 +1299,14 @@ def expand_sweep( pad_k=bool(pk), persistent=bool(persist), gfx_arch=arch, + variant=variant, + num_a_tensors=num_a_tensors, + num_b_tensors=num_b_tensors, + num_d_tensors=num_d_tensors, + a_elementwise_op=a_elementwise_op, + b_elementwise_op=b_elementwise_op, + cde_elementwise_op=cde_elementwise_op, + layout_d=_LAYOUT_WORD[ld], ) if c.name in seen: continue diff --git a/projects/composablekernel/tile_engine/ops/gemm/gemm_full_benchmark.py b/projects/composablekernel/tile_engine/ops/gemm/gemm_full_benchmark.py index f6bdecfae56c..9a9db49d0a66 100644 --- a/projects/composablekernel/tile_engine/ops/gemm/gemm_full_benchmark.py +++ b/projects/composablekernel/tile_engine/ops/gemm/gemm_full_benchmark.py @@ -62,10 +62,17 @@ VARIANT_CONFIGS = { "gemm_universal": "configs", "gemm_multi_d": "gemm_multi_d/configs", + "gemm_multi_abd": "gemm_multi_abd/configs", "gemm_preshuffle": "gemm_preshuffle/configs", "grouped_gemm": "grouped_gemm/configs", } DEFAULT_VARIANT = "gemm_universal" + +# Bridge variant string (expand_sweep / GemmKernelConfig.variant) per benchmark +# variant. Anything not listed goes through the regular "standard" path. +BRIDGE_VARIANT = { + "gemm_multi_abd": "multi_abd", +} CI_CONFIG_NAME = "default_ci_config.json" EXAMPLE_PROBLEMS_NAME = "example_problems.json" @@ -79,7 +86,9 @@ SUPPORTED_DTYPES = ("fp16", "bf16") # Row-major C only: ck_tile's universal GEMM rejects column-major C at build. -SUPPORTED_LAYOUTS = ("rcr", "rrr", "crr", "ccr") +# The 4-char codes (rcrr, ...) are the multi_abd A,B,E,D layouts; TE gemm_multi_abd +# only supports rcrr today. +SUPPORTED_LAYOUTS = ("rcr", "rrr", "crr", "ccr", "rcrr", "rrrr", "crrr", "ccrr") def detect_devices(): @@ -350,10 +359,23 @@ def main(): print(f" Variant: {args.variant}") print(f" Configs: {', '.join(config_paths)}") + bridge_variant = BRIDGE_VARIANT.get(args.variant, "standard") + # Multi-ABD needs the 4-char (A,B,E,D) layout; if the user left the 3-char + # default in place, extend it (D defaults to the C/E layout). + sweep_layout = args.layout + if bridge_variant == "multi_abd" and len(sweep_layout) == 3: + sweep_layout = sweep_layout + sweep_layout[2] + all_configs = [] for cfg_path in config_paths: all_configs.extend( - expand_sweep(cfg_path, args.arch, dtype=args.dtype, layout=args.layout) + expand_sweep( + cfg_path, + args.arch, + dtype=args.dtype, + layout=sweep_layout, + variant=bridge_variant, + ) ) if args.max_kernels > 0: diff --git a/projects/composablekernel/tile_engine/ops/gemm/run_one_gemm_kernel.py b/projects/composablekernel/tile_engine/ops/gemm/run_one_gemm_kernel.py index 54e138198a96..f1347eb89f20 100644 --- a/projects/composablekernel/tile_engine/ops/gemm/run_one_gemm_kernel.py +++ b/projects/composablekernel/tile_engine/ops/gemm/run_one_gemm_kernel.py @@ -34,10 +34,15 @@ if p and p not in sys.path: sys.path.insert(0, p) -from gemm_utils import GemmProblem, GpuGemmRunner # noqa: E402 +from gemm_utils import GemmProblem, GpuGemmRunner, GpuMultiABDRunner # noqa: E402 import numpy as np # noqa: E402 +def _is_multi_abd(kernel_name): + """A multi_abd kernel name carries the ``_multiabd_`` suffix (see codegen).""" + return "_multiabd_" in (kernel_name or "") + + def _run_one(idx, so_path, prob_dict, kernel_name, verify=False, verify_tol=2e-2): """Run a single kernel and emit its result as one JSON line. @@ -49,21 +54,30 @@ def _run_one(idx, so_path, prob_dict, kernel_name, verify=False, verify_tol=2e-2 try: problem = GemmProblem.from_dict(prob_dict) - # Cache host matrices per shape so batch mode doesn't regenerate huge inputs per kernel. - cache = getattr(_run_one, "_ab_cache", {}) - key = (problem.M, problem.N, problem.K) - if key not in cache: - rng = np.random.RandomState(42) - cache[key] = ( - (rng.randn(problem.M, problem.K) * 0.1).astype(np.float32), - (rng.randn(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 = GpuGemmRunner(lib_path=so_path) - result = runner.run(A, B, problem) + # Multi-ABD has a divergent (array-pointer) ABI and its own runner, which + # generates its A/B/D operands internally (no shared A,B). Route by name. + if _is_multi_abd(kernel_name): + # CRITICAL: load the library ONLY inside this subprocess. + runner = GpuMultiABDRunner(lib_path=so_path) + result = runner.run(problem) + A = B = None + else: + # Cache host matrices per shape so batch mode doesn't regenerate huge + # inputs per kernel. + cache = getattr(_run_one, "_ab_cache", {}) + key = (problem.M, problem.N, problem.K) + if key not in cache: + rng = np.random.RandomState(42) + cache[key] = ( + (rng.randn(problem.M, problem.K) * 0.1).astype(np.float32), + (rng.randn(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 = GpuGemmRunner(lib_path=so_path) + result = runner.run(A, B, problem) if result.success: non_zero = ( @@ -79,7 +93,12 @@ def _run_one(idx, so_path, prob_dict, kernel_name, verify=False, verify_tol=2e-2 "non_zero": non_zero, "kernel": kernel_name, } - if verify: + # Multi-ABD's reference is a multi-tensor sum with per-group element- + # wise ops whose operands the runner generates internally, so the + # simple A@B check does not apply; correctness for multi_abd is + # covered by the Old-TE profiler's own --verify path. Only run the + # numpy reference for the single-A/single-B standard path. + if verify and A is not None and B is not None: ref = A.astype(np.float32) @ B.astype(np.float32) got = result.output.astype(np.float32) denom = float(np.max(np.abs(ref))) or 1.0 From 58ce1478a6e3e5d326afbb408605c1201d6ae447 Mon Sep 17 00:00:00 2001 From: Muhammed Emin Ozturk <3836908+ozturkosu@users.noreply.github.com> Date: Sat, 11 Jul 2026 04:21:56 -0400 Subject: [PATCH 2/8] Address multi_abd bridge review: real numeric verification, driven ops, valid CI tiles Round-2 fixes for the gemm_multi_abd TileEngine->Dispatcher bridge: - B1: verify the BRIDGE launch path numerically. GpuMultiABDRunner now owns a numpy reference E = CDE( AB(As) @ BB(Bs), {Ds} ) that mirrors ck_tile::reference_gemm_multiple_abd byte-for-byte (A/B groups combined element-wise into one matrix, single GEMM, then the CDE op folds the D tensors) and reports max_rel on the result; run_one_gemm_kernel routes --verify through it. This also uncovered + fixed a real launch bug: the GemmMultiABDKernel does NOT derive strides from a 0 sentinel, so the runner now passes explicit leading strides for A/B/D/E (get_default_stride semantics) instead of 0, which had collapsed the whole output onto row 0. - B2: drive the element-wise ops + tensor counts. expand_sweep reads a multi_abd_config block (lists) and iterates the op/count combos, and gemm_full_benchmark exposes --multi-abd-{a,b,cde}-op / --multi-abd-num-{a,b,d} CLI overrides, so non-PassThrough kernels are actually generated and verified. - B3: default_ci_config.json warp tile 32x32x64 -> 32x32x16 (32x32x64 is not in the gfx942 fp16 XDL allow-list; CI now yields 16 kernels, was 0). - N1: comment the fp16_fp16_fp16 vs fp16_fp16_fp32 arch-filter key mismatch in ctypes_utils as a latent bridge-wide parity narrowing. - N2: derive layout + per-group ops from the config object (passed via the worker payload) instead of parsing the kernel name with a silent rcrr default. - Ensure build/examples exists before hipcc -o so the self-contained multi_abd .so builds without a prior full CMake configure. dtype x layout unchanged: exactly fp16 x rcrr (Old-TE parity set). Validated on gfx942 (MI300X, GPU0): default 2/2/2 PassThrough max_rel=2.9e-4, CDE=MultiDAdd max_rel=5.7e-4, A-op=MultiDAdd max_rel=4.1e-4 (all << fp16 tol); CI config yields 16 kernels; standard-variant expand_sweep regression clean. --- .../dispatcher/python/ctypes_utils.py | 12 + .../dispatcher/python/gemm_utils.py | 295 ++++++++++++++++-- .../ops/gemm/gemm_full_benchmark.py | 54 +++- .../configs/default_ci_config.json | 2 +- .../ops/gemm/run_one_gemm_kernel.py | 61 +++- 5 files changed, 379 insertions(+), 45 deletions(-) diff --git a/projects/composablekernel/dispatcher/python/ctypes_utils.py b/projects/composablekernel/dispatcher/python/ctypes_utils.py index cc94ede685c9..cf78be469eef 100644 --- a/projects/composablekernel/dispatcher/python/ctypes_utils.py +++ b/projects/composablekernel/dispatcher/python/ctypes_utils.py @@ -176,6 +176,18 @@ def get_arch_filter_data() -> Dict[str, Any]: "gfx942": [[1, 4, 1], [2, 2, 1], [4, 1, 1]], "gfx90a": [[1, 4, 1], [2, 2, 1], [4, 1, 1]], }, + # NOTE (N1, latent parity risk): this table keys the accumulator as + # ``fp16_fp16_fp16`` and lists only [[16,16,16],[32,32,16]], whereas + # the codegen arch_filter (dispatcher/codegen/arch_filter.py) keys + # fp16 GEMMs as ``fp16_fp16_fp32`` and additionally permits + # [32,32,8], [16,16,32], [4,64,16], [64,4,16]. Because this validator + # runs *before* codegen in expand_sweep, it silently NARROWS the + # accepted fp16 warp tiles to the intersection ([16,16,16],[32,32,16]) + # -- a bridge-wide (not multi_abd-specific) parity gap that can drop + # otherwise-valid Old-TE tiles. Left as-is here (the broad fix belongs + # with the shared bridge validator); tracked so it is not mistaken for + # a multi_abd regression. The multi_abd CI/default tiles all live in + # the intersection, so they are unaffected. "warp_tile_combos": { "gfx942": {"fp16_fp16_fp16": [[16, 16, 16], [32, 32, 16]]}, "gfx90a": {"fp16_fp16_fp16": [[16, 16, 16], [32, 32, 16]]}, diff --git a/projects/composablekernel/dispatcher/python/gemm_utils.py b/projects/composablekernel/dispatcher/python/gemm_utils.py index ed5a5f614946..36e38a71d1a0 100644 --- a/projects/composablekernel/dispatcher/python/gemm_utils.py +++ b/projects/composablekernel/dispatcher/python/gemm_utils.py @@ -314,6 +314,11 @@ class GemmResult: status: int tflops: float kernel_name: str + # Optional numeric-verification metric: global relative error of the kernel + # output vs a numpy reference (max|out-ref|/max|ref|). None when the runner + # did not compute a reference. Multi-ABD populates this in-runner because it + # generates its own A/B/D operands internally (see GpuMultiABDRunner.run). + max_rel: Optional[float] = None @property def success(self) -> bool: @@ -755,6 +760,10 @@ def run( elem_b: int, elem_d: int, elem_e: int, + stride_as: Optional[List[int]] = None, + stride_bs: Optional[List[int]] = None, + stride_ds: Optional[List[int]] = None, + stride_e: int = 0, ) -> Tuple[int, float]: def _ptr_array(arrays): arr = (ctypes.c_void_p * max(len(arrays), 1))() @@ -762,21 +771,35 @@ def _ptr_array(arrays): arr[i] = a.ctypes.data_as(ctypes.c_void_p) return arr + def _i64_array(vals): + # The GemmMultiABDKernel does NOT derive strides from a 0 sentinel + # (unlike some CK host helpers); it passes them straight to the + # UniversalGemm kernel args. We must therefore supply the explicit + # leading strides (see GpuMultiABDRunner.run, which mirrors the + # Old-TE profiler's get_default_stride). + if not vals: + return ctypes.POINTER(ctypes.c_int64)() + arr = (ctypes.c_int64 * len(vals))() + for i, v in enumerate(vals): + arr[i] = int(v) + return arr + as_ptrs = _ptr_array(as_arrays) bs_ptrs = _ptr_array(bs_arrays) ds_ptrs = _ptr_array(ds_arrays) - # Strides <= 0 tell the kernel to use its default layout-derived stride. - null_i64 = ctypes.POINTER(ctypes.c_int64)() + stride_as_arr = _i64_array(stride_as) + stride_bs_arr = _i64_array(stride_bs) + stride_ds_arr = _i64_array(stride_ds) time_ms = ctypes.c_float(0.0) status = self._lib.dispatcher_run_multi_abd( as_ptrs, bs_ptrs, ds_ptrs, e_array.ctypes.data_as(ctypes.c_void_p), - null_i64, - null_i64, - null_i64, - 0, + stride_as_arr, + stride_bs_arr, + stride_ds_arr, + int(stride_e), elem_a, elem_b, elem_d, @@ -795,6 +818,72 @@ def cleanup(self) -> None: self._lib.dispatcher_cleanup() +# Multi-ABD per-group element-wise CDE ops, as numpy reductions over +# (acc, D0, D1, ...). These mirror ck_tile::element_wise (see +# unary_element_wise_operation.hpp) EXACTLY so the numpy reference matches the +# device epilogue: +# PassThrough : E = C (D tensors ignored) +# MultiDAdd : E = C + D0 + D1 + ... +# MultiDMultiply : E = C * D0 * D1 * ... +# AddScale : E = scale * (C + D0 + D1 + ...) (default scale = 1.0; +# note AddScale folds *all* its arguments including C, so C +# participates in the sum -- see struct AddScale) +def _cde_reference(op: str, acc: np.ndarray, ds: List[np.ndarray]) -> np.ndarray: + """Apply the CDE element-wise op to the fp32 accumulator + D tensors.""" + acc = acc.astype(np.float32) + ds32 = [d.astype(np.float32) for d in ds] + if op == "PassThrough": + return acc + if op == "MultiDAdd": + out = acc.copy() + for d in ds32: + out = out + d + return out + if op == "MultiDMultiply": + out = acc.copy() + for d in ds32: + out = out * d + return out + if op == "AddScale": + # AddScale starts at 0 and folds every argument (C included). + out = acc.copy() + for d in ds32: + out = out + d + return out # scale defaults to 1.0 + raise ValueError(f"Unsupported CDE element-wise op for reference: {op}") + + +def _ab_reference(op: str, group: List[np.ndarray]) -> np.ndarray: + """Combine a group of A (or B) tensors into a single matrix via the op. + + Mirrors reference_gemm_multiple_abd's A/B pre-pass, which applies the group + element-wise op across the tuple of tensors element-by-element: + PassThrough : first tensor only (op(y, x0, x1...) assigns y = x0) + MultiDAdd : sum of all tensors + MultiDMultiply : product of all tensors + AddScale : scale * sum of all tensors (scale defaults to 1.0) + """ + g32 = [t.astype(np.float32) for t in group] + if op == "PassThrough": + return g32[0] + if op == "MultiDAdd": + out = g32[0].copy() + for t in g32[1:]: + out = out + t + return out + if op == "AddScale": + out = np.zeros_like(g32[0]) + for t in g32: + out = out + t + return out # scale defaults to 1.0 + if op == "MultiDMultiply": + out = g32[0].copy() + for t in g32[1:]: + out = out * t + return out + raise ValueError(f"Unsupported A/B element-wise op for reference: {op}") + + class GpuMultiABDRunner: """High-level multi-ABD runner: construct from a .so, call run(problem). @@ -802,48 +891,122 @@ class GpuMultiABDRunner: gemm_multi_abd op). The runner builds the host operand buffers in the kernel's element dtype + layout and hands raw pointer arrays to the .so, which owns GPU memory. fp16 is the only supported multi-abd dtype. + + Numeric verification (B1): the runner generates all A/B/D operands itself, + so it also owns the numpy reference. When ``verify`` is set on ``run`` it + computes ``E = CDE( AB(As) @ BB(Bs), {Ds} )`` -- byte-for-byte mirroring + ck_tile::reference_gemm_multiple_abd (A/B groups combined element-wise via + their ops into one matrix, single GEMM, then the CDE op folds the D tensors) + -- and reports max_rel = max|E_gpu - E_ref| / max|E_ref| on the result. + + Layout and per-group element-wise ops / tensor counts are taken from the + supplied config object when available (N2); the kernel name is used only as + a fallback so the runner never silently guesses a wrong layout. """ - def __init__(self, lib_path: Path): + def __init__( + self, + lib_path: Path, + layout4: Optional[str] = None, + a_elementwise_op: Optional[str] = None, + b_elementwise_op: Optional[str] = None, + cde_elementwise_op: Optional[str] = None, + ): self.lib = MultiABDDispatcherLib(lib_path) if not self.lib.initialize(): raise RuntimeError(f"Failed to initialize multi_abd .so: {lib_path}") self._kernel_name = self.lib.kernel_name self._num_a, self._num_b, self._num_d = self.lib.tensor_counts + # N2: prefer the layout / ops derived from the config object. Only fall + # back to parsing the kernel name (which is deterministic from config) + # when the caller did not supply them -- no silent "rcrr" default. + self._layout4 = layout4 or self._parse_layout4() + a_op, b_op, cde_op = self._parse_ops() + self._a_op = a_elementwise_op or a_op + self._b_op = b_elementwise_op or b_op + self._cde_op = cde_elementwise_op or cde_op + @property def kernel_name(self) -> str: return self._kernel_name - def run(self, problem: GemmProblem, seed: int = 0) -> GemmResult: + def _parse_layout4(self) -> str: + """Fallback: 4-char (A,B,E,D) layout from ``gemm___...``.""" + parts = self._kernel_name.split("_") + if len(parts) > 2 and len(parts[2]) == 4 and set(parts[2]) <= {"r", "c"}: + return parts[2] + raise ValueError( + f"Cannot derive multi_abd layout from kernel name {self._kernel_name!r}; " + "pass layout4 from the config object instead" + ) + + def _parse_ops(self) -> Tuple[str, str, str]: + """Fallback: parse the three element-wise ops from the kernel name. + + Name suffix is ``..._multiabd_a_b_d___``. + """ + marker = "_multiabd_" + idx = self._kernel_name.find(marker) + if idx >= 0: + tail = self._kernel_name[idx + len(marker) :].split("_") + # tail == [aNA, bNB, dND, Aop, Bop, CDEop] + if len(tail) >= 6: + return tail[3], tail[4], tail[5] + return "PassThrough", "PassThrough", "PassThrough" + + def run( + self, + problem: GemmProblem, + seed: int = 0, + verify: bool = False, + verify_tol: float = 2e-2, + ) -> GemmResult: M, N, K = problem.M, problem.N, problem.K dtype = _dtype_from_kernel_name(self._kernel_name) - # Kernel name for multi_abd is gemm__<4char-layout>_...; take the - # A/B/D layout chars (E is the C position, index 2). - parts = self._kernel_name.split("_") - layout4 = parts[2] if len(parts) > 2 else "rcrr" + layout4 = self._layout4 la, lb = layout4[0], layout4[1] ld = layout4[3] if len(layout4) >= 4 else "r" rng = np.random.default_rng(seed) + # Logical (row-major, M-major) operands used for the numpy reference, and + # the physically-laid-out contiguous buffers handed to the .so. The .so + # interprets each buffer per the compiled layout, so a column-major + # operand is stored transposed but represents the same logical matrix. def _mk(rows, cols, layout_char, n_tensors, lo, hi): - outs = [] + logical, physical = [], [] for _ in range(n_tensors): x = rng.uniform(lo, hi, size=(rows, cols)).astype(np.float32) - x_lay = x if layout_char == "r" else x.T - if dtype == "fp16": - outs.append(np.ascontiguousarray(x_lay, dtype=np.float16)) - else: + if dtype != "fp16": raise ValueError(f"multi_abd runner supports fp16 only, got {dtype}") - return outs - - as_arrays = _mk(M, K, la, self._num_a, -5.0, 5.0) - bs_arrays = _mk(K, N, lb, self._num_b, -5.0, 5.0) - ds_arrays = _mk(M, N, ld, self._num_d, -1.0, 1.0) + x16 = x.astype(np.float16) + logical.append(x16) + x_lay = x16 if layout_char == "r" else x16.T + physical.append(np.ascontiguousarray(x_lay, dtype=np.float16)) + return logical, physical + + as_logical, as_arrays = _mk(M, K, la, self._num_a, -5.0, 5.0) + bs_logical, bs_arrays = _mk(K, N, lb, self._num_b, -5.0, 5.0) + ds_logical, ds_arrays = _mk(M, N, ld, self._num_d, -1.0, 1.0) elem = _ELEM_BYTES.get(dtype, 2) e_array = np.zeros((M, N), dtype=np.float16) + # Explicit leading strides -- the GemmMultiABDKernel does NOT derive them + # from a zero sentinel (it forwards them straight to the UniversalGemm + # kernel args), so a 0 stride collapses the whole output onto row 0. + # Mirror the Old-TE profiler's get_default_stride(rows, cols, is_row): + # row-major -> #cols, col-major -> #rows. + def _lead_stride(rows, cols, layout_char): + return cols if layout_char == "r" else rows + + stride_as = [_lead_stride(M, K, la)] * self._num_a + stride_bs = [_lead_stride(K, N, lb)] * self._num_b + stride_ds = [_lead_stride(M, N, ld)] * self._num_d + # E is the C position (index 2) of the 4-char layout. + le = layout4[2] if len(layout4) >= 3 else "r" + stride_e = _lead_stride(M, N, le) + status, time_ms = self.lib.run( as_arrays, bs_arrays, @@ -856,14 +1019,33 @@ def _mk(rows, cols, layout_char, n_tensors, lo, hi): elem, elem, elem, + stride_as=stride_as, + stride_bs=stride_bs, + stride_ds=stride_ds, + stride_e=stride_e, ) tflops = (problem.flops / (time_ms * 1e-3)) / 1e12 if time_ms > 0 else 0.0 + + max_rel = None + if verify and status == 0: + # A/B groups are combined element-wise into a single matrix, then a + # single GEMM, then the CDE op folds the D tensors -- exactly + # reference_gemm_multiple_abd. Compute in fp32. + a_m_k = _ab_reference(self._a_op, as_logical) + b_k_n = _ab_reference(self._b_op, bs_logical) + acc = a_m_k @ b_k_n + ref = _cde_reference(self._cde_op, acc, ds_logical).astype(np.float32) + got = e_array.astype(np.float32) + denom = float(np.max(np.abs(ref))) or 1.0 + max_rel = float(np.max(np.abs(got - ref)) / denom) + return GemmResult( output=e_array, time_ms=time_ms, status=status, tflops=tflops, kernel_name=self._kernel_name, + max_rel=max_rel, ) @@ -963,6 +1145,11 @@ def _build_compile_jobs( static_lib = build_dir / "libck_tile_dispatcher.a" lib_path = build_dir / "examples" / f"lib{config.name}.so" + # Multi-ABD is self-contained (registry-bypass, no static lib) so it can be + # built without a prior full CMake configure of the dispatcher; that CMake + # step is what normally creates build/examples. Ensure the output directory + # exists so the hipcc -o path is always writable (harmless if it already is). + lib_path.parent.mkdir(parents=True, exist_ok=True) obj_file = lib_path.with_suffix(".o") compile_cmd = [ @@ -1193,6 +1380,7 @@ def expand_sweep( a_elementwise_op: str = "PassThrough", b_elementwise_op: str = "PassThrough", cde_elementwise_op: str = "PassThrough", + mabd_cli_overrides: Optional[Dict[str, Any]] = None, ) -> List[GemmKernelConfig]: """Expand a Tile Engine GEMM JSON sweep config into GemmKernelConfig list. @@ -1237,6 +1425,56 @@ def expand_sweep( # Multi-ABD carries a 4th (D) layout char; default D to C's layout otherwise. ld = layout[3] if (variant == "multi_abd" and len(layout) >= 4) else lc + # Multi-ABD (B2): the tensor counts and per-group element-wise ops are a real + # part of the swept configuration -- distinct ops produce distinct kernels + # (different epilogue math). Read them from an optional ``multi_abd_config`` + # block in the TE config JSON (lists of values), falling back to the scalar + # kwargs (which default to the Old-TE 2/2/2 all-PassThrough combo). This is + # what lets the driver actually generate + verify a non-PassThrough kernel. + # + # Allowed ops mirror the Old-TE gemm_multi_abd instance builder: + # {PassThrough, AddScale, MultiDMultiply, MultiDAdd}. + if variant == "multi_abd": + mabd = dict(cfg.get("multi_abd_config", {}) or {}) + # CLI overrides win over both the config block and the scalar defaults. + if mabd_cli_overrides: + mabd.update(mabd_cli_overrides) + + def _as_list(v, default): + if v is None: + return list(default) + return list(v) if isinstance(v, (list, tuple)) else [v] + + na_list = _as_list(mabd.get("num_a_tensors"), [num_a_tensors]) + nb_list = _as_list(mabd.get("num_b_tensors"), [num_b_tensors]) + nd_list = _as_list(mabd.get("num_d_tensors"), [num_d_tensors]) + a_ops = _as_list(mabd.get("a_elementwise_op"), [a_elementwise_op]) + b_ops = _as_list(mabd.get("b_elementwise_op"), [b_elementwise_op]) + cde_ops = _as_list(mabd.get("cde_elementwise_op"), [cde_elementwise_op]) + _ALLOWED_MABD_OPS = {"PassThrough", "AddScale", "MultiDMultiply", "MultiDAdd"} + for op in (*a_ops, *b_ops, *cde_ops): + if op not in _ALLOWED_MABD_OPS: + raise ValueError( + f"Invalid multi_abd element-wise op {op!r}; " + f"valid: {sorted(_ALLOWED_MABD_OPS)}" + ) + mabd_combos = list( + itertools.product(na_list, nb_list, nd_list, a_ops, b_ops, cde_ops) + ) + else: + # Non-multi_abd variants: a single, inert combo carrying the scalar + # kwargs (unused by those code paths' names). + mabd_combos = [ + ( + num_a_tensors, + num_b_tensors, + num_d_tensors, + a_elementwise_op, + b_elementwise_op, + cde_elementwise_op, + ) + ] + configs: List[GemmKernelConfig] = [] seen: set = set() for ( @@ -1274,6 +1512,7 @@ def expand_sweep( pad_ks, persistents, ): + for (m_na, m_nb, m_nd, m_aop, m_bop, m_cdeop) in mabd_combos: c = GemmKernelConfig( dtype_a=dtype, dtype_b=dtype, @@ -1300,12 +1539,12 @@ def expand_sweep( persistent=bool(persist), gfx_arch=arch, variant=variant, - num_a_tensors=num_a_tensors, - num_b_tensors=num_b_tensors, - num_d_tensors=num_d_tensors, - a_elementwise_op=a_elementwise_op, - b_elementwise_op=b_elementwise_op, - cde_elementwise_op=cde_elementwise_op, + num_a_tensors=m_na, + num_b_tensors=m_nb, + num_d_tensors=m_nd, + a_elementwise_op=m_aop, + b_elementwise_op=m_bop, + cde_elementwise_op=m_cdeop, layout_d=_LAYOUT_WORD[ld], ) if c.name in seen: diff --git a/projects/composablekernel/tile_engine/ops/gemm/gemm_full_benchmark.py b/projects/composablekernel/tile_engine/ops/gemm/gemm_full_benchmark.py index 9a9db49d0a66..d222ab3440c2 100644 --- a/projects/composablekernel/tile_engine/ops/gemm/gemm_full_benchmark.py +++ b/projects/composablekernel/tile_engine/ops/gemm/gemm_full_benchmark.py @@ -179,10 +179,19 @@ def _run_batch_on_device(device_id, unit, args, worker_path, base_env): prob_idx, prob_dict, batch = unit M, N, K = 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 - ] + def _item(cfg, lib): + it = {"so_path": str(lib), "problem": prob_dict, "kernel_name": cfg.name} + # N2: carry the multi_abd layout + per-group element-wise ops from the + # config object so the worker's numpy reference uses the real ops/layout + # instead of parsing (and defaulting) the kernel name. + if getattr(cfg, "variant", "standard") == "multi_abd": + it["layout4"] = cfg.layout4 + it["a_elementwise_op"] = cfg.a_elementwise_op + it["b_elementwise_op"] = cfg.b_elementwise_op + it["cde_elementwise_op"] = cfg.cde_elementwise_op + return it + + items = [_item(cfg, lib) for _, cfg, lib in batch] payload = json.dumps( {"items": items, "verify": args.verify, "verify_tol": args.verify_tol} ) @@ -311,6 +320,23 @@ def main(): choices=SUPPORTED_LAYOUTS, help=f"A/B/C layout (supported: {', '.join(SUPPORTED_LAYOUTS)})", ) + # Multi-ABD only: per-group element-wise ops + tensor counts. These override + # (are merged into) any multi_abd_config block in the config JSON so a single + # non-PassThrough kernel can be driven + verified without editing configs. + # Valid ops mirror the Old-TE gemm_multi_abd instance builder. + _MABD_OPS = ("PassThrough", "AddScale", "MultiDMultiply", "MultiDAdd") + parser.add_argument("--multi-abd-num-a", type=int, default=None, + help="multi_abd: number of A tensors (default: config/2)") + parser.add_argument("--multi-abd-num-b", type=int, default=None, + help="multi_abd: number of B tensors (default: config/2)") + parser.add_argument("--multi-abd-num-d", type=int, default=None, + help="multi_abd: number of D tensors (default: config/2)") + parser.add_argument("--multi-abd-a-op", default=None, choices=_MABD_OPS, + help="multi_abd: A element-wise op (default: config/PassThrough)") + parser.add_argument("--multi-abd-b-op", default=None, choices=_MABD_OPS, + help="multi_abd: B element-wise op (default: config/PassThrough)") + parser.add_argument("--multi-abd-cde-op", default=None, choices=_MABD_OPS, + help="multi_abd: CDE element-wise op (default: config/PassThrough)") parser.add_argument("--problems", default=None, help="JSON file of M,N,K problems") parser.add_argument("--csv", type=str, default="gemm_results.csv") parser.add_argument("--workers", type=int, default=8, help="Parallel build workers") @@ -366,6 +392,24 @@ def main(): if bridge_variant == "multi_abd" and len(sweep_layout) == 3: sweep_layout = sweep_layout + sweep_layout[2] + # Multi-ABD element-wise ops / tensor counts: CLI overrides win over the + # config; otherwise expand_sweep falls back to any multi_abd_config block in + # the JSON and finally to the Old-TE 2/2/2 all-PassThrough default. + mabd_kwargs = {} + if bridge_variant == "multi_abd": + if args.multi_abd_num_a is not None: + mabd_kwargs["num_a_tensors"] = args.multi_abd_num_a + if args.multi_abd_num_b is not None: + mabd_kwargs["num_b_tensors"] = args.multi_abd_num_b + if args.multi_abd_num_d is not None: + mabd_kwargs["num_d_tensors"] = args.multi_abd_num_d + if args.multi_abd_a_op is not None: + mabd_kwargs["a_elementwise_op"] = args.multi_abd_a_op + if args.multi_abd_b_op is not None: + mabd_kwargs["b_elementwise_op"] = args.multi_abd_b_op + if args.multi_abd_cde_op is not None: + mabd_kwargs["cde_elementwise_op"] = args.multi_abd_cde_op + all_configs = [] for cfg_path in config_paths: all_configs.extend( @@ -375,6 +419,8 @@ def main(): dtype=args.dtype, layout=sweep_layout, variant=bridge_variant, + mabd_cli_overrides=(mabd_kwargs or None), + **mabd_kwargs, ) ) diff --git a/projects/composablekernel/tile_engine/ops/gemm/gemm_multi_abd/configs/default_ci_config.json b/projects/composablekernel/tile_engine/ops/gemm/gemm_multi_abd/configs/default_ci_config.json index 05d4f171cc51..99902697f822 100644 --- a/projects/composablekernel/tile_engine/ops/gemm/gemm_multi_abd/configs/default_ci_config.json +++ b/projects/composablekernel/tile_engine/ops/gemm/gemm_multi_abd/configs/default_ci_config.json @@ -42,7 +42,7 @@ }, "warp_tile_k": { "values": [ - 64 + 16 ] } }, diff --git a/projects/composablekernel/tile_engine/ops/gemm/run_one_gemm_kernel.py b/projects/composablekernel/tile_engine/ops/gemm/run_one_gemm_kernel.py index f1347eb89f20..0ec9b5d93e29 100644 --- a/projects/composablekernel/tile_engine/ops/gemm/run_one_gemm_kernel.py +++ b/projects/composablekernel/tile_engine/ops/gemm/run_one_gemm_kernel.py @@ -43,13 +43,30 @@ def _is_multi_abd(kernel_name): return "_multiabd_" in (kernel_name or "") -def _run_one(idx, so_path, prob_dict, kernel_name, verify=False, verify_tol=2e-2): +def _run_one( + idx, + so_path, + prob_dict, + kernel_name, + verify=False, + verify_tol=2e-2, + layout4=None, + a_op=None, + b_op=None, + cde_op=None, +): """Run a single kernel and emit its result as one JSON line. When ``verify`` is set, the kernel output is checked against an fp32 numpy - reference (``A @ B``) using the global relative metric - ``max|out - ref| / max|ref|``; the emitted ``verified`` field then reflects - correctness, not just liveness (``non_zero``). + reference using the global relative metric ``max|out - ref| / max|ref|``; + the emitted ``verified`` field then reflects correctness, not just liveness + (``non_zero``). + + Standard path: reference is ``A @ B``. Multi-ABD path (B1): the reference is + the full ``E = CDE( AB(As) @ BB(Bs), {Ds} )`` computed inside + GpuMultiABDRunner (it owns the operands); the runner returns ``max_rel`` on + the result. ``layout4`` and the per-group element-wise ops come from the + config object (N2) so the runner never guesses the layout from the name. """ try: problem = GemmProblem.from_dict(prob_dict) @@ -58,8 +75,17 @@ def _run_one(idx, so_path, prob_dict, kernel_name, verify=False, verify_tol=2e-2 # generates its A/B/D operands internally (no shared A,B). Route by name. if _is_multi_abd(kernel_name): # CRITICAL: load the library ONLY inside this subprocess. - runner = GpuMultiABDRunner(lib_path=so_path) - result = runner.run(problem) + runner = GpuMultiABDRunner( + lib_path=so_path, + layout4=layout4, + a_elementwise_op=a_op, + b_elementwise_op=b_op, + cde_elementwise_op=cde_op, + ) + # The multi_abd runner computes its own numpy reference (it owns the + # operands), so verification flows through the runner, not the A@B + # check below. + result = runner.run(problem, verify=verify, verify_tol=verify_tol) A = B = None else: # Cache host matrices per shape so batch mode doesn't regenerate huge @@ -93,12 +119,15 @@ def _run_one(idx, so_path, prob_dict, kernel_name, verify=False, verify_tol=2e-2 "non_zero": non_zero, "kernel": kernel_name, } - # Multi-ABD's reference is a multi-tensor sum with per-group element- - # wise ops whose operands the runner generates internally, so the - # simple A@B check does not apply; correctness for multi_abd is - # covered by the Old-TE profiler's own --verify path. Only run the - # numpy reference for the single-A/single-B standard path. - if verify and A is not None and B is not None: + if verify and _is_multi_abd(kernel_name): + # Multi-ABD verifies against its own full multi-tensor reference + # (E = CDE(AB(As) @ BB(Bs), {Ds})) computed inside the runner, + # exercising the BRIDGE launch path (dispatcher_run_multi_abd -> + # SelectedKernel::launch), not just the Old-TE --verify path. + if result.max_rel is not None: + out["max_rel"] = result.max_rel + out["verified"] = bool(result.max_rel <= verify_tol) + elif verify and A is not None and B is not None: ref = A.astype(np.float32) @ B.astype(np.float32) got = result.output.astype(np.float32) denom = float(np.max(np.abs(ref))) or 1.0 @@ -151,6 +180,10 @@ def main(): item.get("kernel_name", "unknown"), verify=verify, verify_tol=verify_tol, + layout4=item.get("layout4"), + a_op=item.get("a_elementwise_op"), + b_op=item.get("b_elementwise_op"), + cde_op=item.get("cde_elementwise_op"), ) else: _run_one( @@ -160,6 +193,10 @@ def main(): d.get("kernel_name", "unknown"), verify=verify, verify_tol=verify_tol, + layout4=d.get("layout4"), + a_op=d.get("a_elementwise_op"), + b_op=d.get("b_elementwise_op"), + cde_op=d.get("cde_elementwise_op"), ) From e6a61a85dd6d89f840d1641eeb1e3519aec0d5b2 Mon Sep 17 00:00:00 2001 From: Muhammed Emin Ozturk <3836908+ozturkosu@users.noreply.github.com> Date: Sun, 12 Jul 2026 18:15:23 -0400 Subject: [PATCH 3/8] test: add multi_abd bridge unit test --- .../dispatcher/tests/test_multi_abd_bridge.py | 191 ++++++++++++++++++ 1 file changed, 191 insertions(+) create mode 100644 projects/composablekernel/dispatcher/tests/test_multi_abd_bridge.py diff --git a/projects/composablekernel/dispatcher/tests/test_multi_abd_bridge.py b/projects/composablekernel/dispatcher/tests/test_multi_abd_bridge.py new file mode 100644 index 000000000000..bb7a18b0462b --- /dev/null +++ b/projects/composablekernel/dispatcher/tests/test_multi_abd_bridge.py @@ -0,0 +1,191 @@ +#!/usr/bin/env python3 + +# Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +# SPDX-License-Identifier: MIT + +"""CPU-only unit tests for the multi_abd GEMM bridge. + +The Tile Engine -> Dispatcher multi_abd bridge relies on one hard invariant: +``GemmKernelConfig.name`` must reproduce, byte-for-byte, the kernel stem that +``unified_gemm_codegen.py`` bakes into the generated kernel (and that the .so +reports at runtime). For multi_abd that stem carries the 4-char (A,B,E,D) +layout plus a ``_multiabd_a{na}_b{nb}_d{nd}_{aop}_{bop}_{cdeop}`` suffix, so +distinct tensor counts / element-wise ops can never collapse onto one kernel. + +These tests exercise only pure host-side logic (name generation, the codegen +JSON projection, and the shipped configs/*.json). No GPU, hipcc, or build is +required, so the suite runs green in CPU-only CI. + +Run: python3 -m pytest tests/test_multi_abd_bridge.py -v +""" + +import json +import sys +import unittest +from pathlib import Path + +SCRIPT_DIR = Path(__file__).parent.resolve() +DISPATCHER_DIR = SCRIPT_DIR.parent +REPO_ROOT = DISPATCHER_DIR.parent +sys.path.insert(0, str(DISPATCHER_DIR / "python")) + +from gemm_utils import ( # noqa: E402 + GemmKernelConfig, + _output_dtype, + _dtype_from_kernel_name, +) + +# The shipped multi_abd sweep configs the bridge codegens from. +_CONFIG_DIR = ( + REPO_ROOT + / "tile_engine" + / "ops" + / "gemm" + / "gemm_multi_abd" + / "configs" +) + + +def _make_config(**overrides): + """A canonical multi_abd config; overrides tweak individual fields.""" + kw = dict( + dtype_a="fp16", + dtype_b="fp16", + dtype_c="fp16", + dtype_acc="fp32", + layout_a="row", + layout_b="col", + layout_c="row", + layout_d="row", + variant="multi_abd", + num_a_tensors=2, + num_b_tensors=2, + num_d_tensors=2, + ) + kw.update(overrides) + return GemmKernelConfig(**kw) + + +class TestMultiAbdName(unittest.TestCase): + """The multi_abd kernel-name contract (the byte-parity invariant).""" + + def test_name_carries_multiabd_suffix(self): + cfg = _make_config() + name = cfg.name + # 4-char (A,B,E,D) layout, not the 3-char C layout. + self.assertIn("_rcrr_", name) + # multiabd tensor-count + op suffix, exactly as codegen emits it. + self.assertTrue( + name.endswith( + "_multiabd_a2_b2_d2_PassThrough_PassThrough_PassThrough" + ), + name, + ) + + def test_full_stem_is_stable(self): + # Pin the entire stem so any drift in the naming scheme is caught. + cfg = _make_config( + 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, + ) + expected = ( + "gemm_fp16_rcrr_compv4_cshuffle_intrawave" + "_True_True_True_False" + "_128x128x32_2x2x1_32x32x16" + "_multiabd_a2_b2_d2_PassThrough_PassThrough_PassThrough" + ) + self.assertEqual(cfg.name, expected) + + def test_tensor_counts_change_the_name(self): + # Distinct tensor counts must not collapse onto one kernel name. + base = _make_config().name + more = _make_config(num_a_tensors=3, num_d_tensors=1).name + self.assertNotEqual(base, more) + self.assertIn("_a3_", more) + self.assertIn("_d1_", more) + + def test_elementwise_ops_change_the_name(self): + base = _make_config().name + scaled = _make_config(cde_elementwise_op="AddScale").name + self.assertNotEqual(base, scaled) + self.assertTrue(scaled.endswith("_PassThrough_PassThrough_AddScale")) + + def test_dtype_recovers_from_name(self): + # The runner reads the input dtype straight out of the compiled .so + # name, so every dtype the bridge builds must round-trip. + 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) + + def test_name_carries_four_char_layout(self): + # multi_abd uses the 4-char (A,B,E,D) layout code in the stem; the D + # char must reflect layout_d independently of the C layout. + for la, lb, lc, ld in ( + ("row", "col", "row", "row"), + ("row", "row", "row", "col"), + ("col", "col", "row", "row"), + ): + cfg = _make_config( + layout_a=la, layout_b=lb, layout_c=lc, layout_d=ld, + ) + self.assertIn(f"_{cfg.layout4}_", cfg.name) + + +class TestMultiAbdCodegenJson(unittest.TestCase): + """The codegen JSON projection must carry the multi_abd block.""" + + def test_codegen_json_has_multi_abd_block(self): + cfg = _make_config( + num_a_tensors=2, num_b_tensors=2, num_d_tensors=2, + cde_elementwise_op="AddScale", + ) + j = cfg.to_codegen_json() + self.assertIn("multi_abd_config", j) + mabd = j["multi_abd_config"] + self.assertEqual(mabd["num_a_tensors"], 2) + self.assertEqual(mabd["num_b_tensors"], 2) + self.assertEqual(mabd["num_d_tensors"], 2) + self.assertEqual(mabd["cde_elementwise_op"], "AddScale") + + 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 TestMultiAbdShippedConfigs(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 multi_abd 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}") + + +if __name__ == "__main__": + unittest.main() From 18a1a9f9dd04604195febe8f385ca5229f04b2aa Mon Sep 17 00:00:00 2001 From: Muhammed Emin Ozturk <3836908+ozturkosu@users.noreply.github.com> Date: Tue, 21 Jul 2026 14:12:54 -0400 Subject: [PATCH 4/8] =?UTF-8?q?fix(ck-tile):=20multi=5Fabd=20GEMM=20bridge?= =?UTF-8?q?=20=E2=80=94=20require=20GPU=20arch=20via=20get=5Farch,=20do=20?= =?UTF-8?q?not=20default=20to=20gfx942?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- .../dispatcher/bindings/ctypes/gemm_multi_abd_ctypes_lib.cpp | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/projects/composablekernel/dispatcher/bindings/ctypes/gemm_multi_abd_ctypes_lib.cpp b/projects/composablekernel/dispatcher/bindings/ctypes/gemm_multi_abd_ctypes_lib.cpp index 582cdf6c4e9d..3a518971d7d4 100644 --- a/projects/composablekernel/dispatcher/bindings/ctypes/gemm_multi_abd_ctypes_lib.cpp +++ b/projects/composablekernel/dispatcher/bindings/ctypes/gemm_multi_abd_ctypes_lib.cpp @@ -46,9 +46,10 @@ // KERNEL_NAME -- the byte-exact runtime kernel name // NumATensors / NumBTensors / NumDTensors -- tensor counts baked into the type -// GPU architecture - can be overridden via -DGFX_ARCH="gfx90a" at compile time +// GPU architecture - must be supplied at compile time via -DGFX_ARCH=. #ifndef GFX_ARCH -#define GFX_ARCH "gfx942" +#error \ + "GFX_ARCH must be defined at compile time (pass -DGFX_ARCH=); do not default to a specific GPU architecture." #endif // The force-included header exports the tensor counts into the global namespace From 90022e01fc0af8cafbb06ddecdf208542a63cd95 Mon Sep 17 00:00:00 2001 From: Muhammed Emin Ozturk <3836908+ozturkosu@users.noreply.github.com> Date: Tue, 21 Jul 2026 20:22:44 -0400 Subject: [PATCH 5/8] =?UTF-8?q?fix(ck-tile):=20multi=5Fabd=20bridge=20?= =?UTF-8?q?=E2=80=94=20address=20Copilot=20review=20(deterministic=20CMake?= =?UTF-8?q?=20glob=20sort,=20validate=20num=5Fa/b/d=5Ftensors,=20rcrr-only?= =?UTF-8?q?=20layout=20guard)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- .../dispatcher/bindings/ctypes/CMakeLists.txt | 3 +++ .../dispatcher/python/gemm_utils.py | 15 +++++++++++++++ .../tile_engine/ops/gemm/gemm_full_benchmark.py | 6 ++++++ 3 files changed, 24 insertions(+) diff --git a/projects/composablekernel/dispatcher/bindings/ctypes/CMakeLists.txt b/projects/composablekernel/dispatcher/bindings/ctypes/CMakeLists.txt index 0c765c5ca72e..9f535bb9bf14 100644 --- a/projects/composablekernel/dispatcher/bindings/ctypes/CMakeLists.txt +++ b/projects/composablekernel/dispatcher/bindings/ctypes/CMakeLists.txt @@ -84,6 +84,9 @@ endif() file(GLOB GEMM_MULTI_ABD_KERNEL_HEADERS "${CMAKE_BINARY_DIR}/generated_kernels/gemm_*_multiabd_*.hpp") if(GEMM_MULTI_ABD_KERNEL_HEADERS) + # Sort for deterministic selection: file(GLOB) order is unspecified, so + # list(GET 0) could otherwise pick an arbitrary header across configure runs. + list(SORT GEMM_MULTI_ABD_KERNEL_HEADERS) list(GET GEMM_MULTI_ABD_KERNEL_HEADERS 0 GEMM_MULTI_ABD_KERNEL_HEADER) message(STATUS "Found GEMM Multi-ABD kernel for ctypes lib: ${GEMM_MULTI_ABD_KERNEL_HEADER}") diff --git a/projects/composablekernel/dispatcher/python/gemm_utils.py b/projects/composablekernel/dispatcher/python/gemm_utils.py index 80933822d58d..93e551a6adb8 100644 --- a/projects/composablekernel/dispatcher/python/gemm_utils.py +++ b/projects/composablekernel/dispatcher/python/gemm_utils.py @@ -1811,6 +1811,21 @@ def _as_list(v, default): na_list = _as_list(mabd.get("num_a_tensors"), [num_a_tensors]) nb_list = _as_list(mabd.get("num_b_tensors"), [num_b_tensors]) nd_list = _as_list(mabd.get("num_d_tensors"), [num_d_tensors]) + # CK's GemmKernelMultiABD requires >=1 A and B tensors and + # DsLayout::size() > 0 (num_d_tensors >= 1); a 0 or non-integer count + # otherwise fails later with a cryptic tuple-size-0 compile error, so + # reject it here with a clear message. + for _label, _vals in ( + ("num_a_tensors", na_list), + ("num_b_tensors", nb_list), + ("num_d_tensors", nd_list), + ): + for _v in _vals: + if not isinstance(_v, int) or _v < 1: + raise ValueError( + f"multi_abd {_label} must be a positive integer (>= 1), " + f"got {_v!r}" + ) a_ops = _as_list(mabd.get("a_elementwise_op"), [a_elementwise_op]) b_ops = _as_list(mabd.get("b_elementwise_op"), [b_elementwise_op]) cde_ops = _as_list(mabd.get("cde_elementwise_op"), [cde_elementwise_op]) diff --git a/projects/composablekernel/tile_engine/ops/gemm/gemm_full_benchmark.py b/projects/composablekernel/tile_engine/ops/gemm/gemm_full_benchmark.py index d222ab3440c2..0cf56b8292e7 100644 --- a/projects/composablekernel/tile_engine/ops/gemm/gemm_full_benchmark.py +++ b/projects/composablekernel/tile_engine/ops/gemm/gemm_full_benchmark.py @@ -391,6 +391,12 @@ def main(): sweep_layout = args.layout if bridge_variant == "multi_abd" and len(sweep_layout) == 3: sweep_layout = sweep_layout + sweep_layout[2] + # multi_abd supports only the 'rcrr' layout today; reject anything else up + # front instead of silently building an unsupported/divergent kernel. + if bridge_variant == "multi_abd" and sweep_layout != "rcrr": + raise SystemExit( + f"multi_abd supports only the 'rcrr' layout today, got {sweep_layout!r}" + ) # Multi-ABD element-wise ops / tensor counts: CLI overrides win over the # config; otherwise expand_sweep falls back to any multi_abd_config block in From ffc81e99ab53016030397c013f12fcee71c1de3e Mon Sep 17 00:00:00 2001 From: Muhammed Emin Ozturk <3836908+ozturkosu@users.noreply.github.com> Date: Tue, 21 Jul 2026 20:51:37 -0400 Subject: [PATCH 6/8] =?UTF-8?q?test(ck-tile):=20multi=5Fabd=20bridge=20?= =?UTF-8?q?=E2=80=94=20add=20GPU-free=20C-ABI=20marshalling=20tests=20(arg?= =?UTF-8?q?types,=20pointer/count=20packing,=20arg=20order,=20status=20for?= =?UTF-8?q?warding)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- .../dispatcher/tests/test_multi_abd_bridge.py | 71 +++++++++++++++++++ 1 file changed, 71 insertions(+) diff --git a/projects/composablekernel/dispatcher/tests/test_multi_abd_bridge.py b/projects/composablekernel/dispatcher/tests/test_multi_abd_bridge.py index bb7a18b0462b..9f503a1fb5f0 100644 --- a/projects/composablekernel/dispatcher/tests/test_multi_abd_bridge.py +++ b/projects/composablekernel/dispatcher/tests/test_multi_abd_bridge.py @@ -19,10 +19,14 @@ Run: python3 -m pytest tests/test_multi_abd_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 @@ -31,6 +35,7 @@ from gemm_utils import ( # noqa: E402 GemmKernelConfig, + MultiABDDispatcherLib, _output_dtype, _dtype_from_kernel_name, ) @@ -187,5 +192,71 @@ def test_configs_are_valid_sweeps(self): self.assertIn(key, tr, f"{path.name} missing {key}") +class TestMultiAbdAbiMarshalling(unittest.TestCase): + """GPU-free tests of the C-ABI marshalling in ``MultiABDDispatcherLib``. + + The ``.so`` is mocked (``ctypes.CDLL`` patched), so no generated kernel or + GPU is needed. These cover the highest-risk-but-previously-untested surface: + the declared ABI, the pointer/stride array packing, the positional argument + order, and status forwarding. + """ + + def _make_lib(self, status=0): + fake = mock.MagicMock() + fake.dispatcher_run_multi_abd.return_value = status + with mock.patch("ctypes.CDLL", return_value=fake): + lib = MultiABDDispatcherLib(Path("/nonexistent/multi_abd.so")) + return lib, fake + + def test_run_abi_argtypes(self): + # The declared argtypes must match, in count and key types, the + # positional arguments run() actually passes (19). + _, fake = self._make_lib() + argt = fake.dispatcher_run_multi_abd.argtypes + self.assertEqual(len(argt), 19) + for i in range(3): # A/B/D host-pointer arrays + self.assertEqual(argt[i], ctypes.POINTER(ctypes.c_void_p)) + self.assertEqual(argt[3], ctypes.c_void_p) # E host pointer + for i in (4, 5, 6): # A/B/D leading-stride arrays + self.assertEqual(argt[i], ctypes.POINTER(ctypes.c_int64)) + self.assertEqual(argt[-1], ctypes.POINTER(ctypes.c_float)) # time_ms + + def test_run_marshals_pointer_arrays_counts_and_order(self): + lib, fake = self._make_lib(status=0) + na, nb, nd = 2, 2, 2 + as_arrays = [np.zeros(4, np.float16) for _ in range(na)] + bs_arrays = [np.zeros(4, np.float16) for _ in range(nb)] + ds_arrays = [np.zeros(4, np.float16) for _ in range(nd)] + e = np.zeros(4, np.float16) + status, _ = lib.run( + as_arrays, bs_arrays, ds_arrays, e, + M=2, N=2, K=2, elem_a=2, elem_b=2, elem_d=2, elem_e=2, + stride_as=[2, 2], stride_bs=[2, 2], stride_ds=[2, 2], stride_e=2, + ) + self.assertEqual(status, 0) + fake.dispatcher_run_multi_abd.assert_called_once() + args = fake.dispatcher_run_multi_abd.call_args[0] + self.assertEqual(len(args), 19) + # A/B/D host-pointer arrays carry exactly one entry per tensor. + self.assertEqual(len(args[0]), na) + self.assertEqual(len(args[1]), nb) + self.assertEqual(len(args[2]), nd) + # Tensor counts occupy positions 12/13/14; M/N/K positions 15/16/17. + self.assertEqual((args[12], args[13], args[14]), (na, nb, nd)) + self.assertEqual((args[15], args[16], args[17]), (2, 2, 2)) + + def test_run_forwards_nonzero_status(self): + # A thin shim must surface the C error code (e.g. -3 tensor-count + # mismatch) verbatim rather than swallowing it. + lib, fake = self._make_lib(status=-3) + a = [np.zeros(4, np.float16)] + status, _ = lib.run( + a, a, a, np.zeros(4, np.float16), + M=1, N=1, K=1, elem_a=2, elem_b=2, elem_d=2, elem_e=2, + stride_as=[1], stride_bs=[1], stride_ds=[1], stride_e=1, + ) + self.assertEqual(status, -3) + + if __name__ == "__main__": unittest.main() From 8b28afdb8a5e61c15a54934dd0b9021b1f50d18f Mon Sep 17 00:00:00 2001 From: Muhammed Emin Ozturk <3836908+ozturkosu@users.noreply.github.com> Date: Tue, 21 Jul 2026 20:54:19 -0400 Subject: [PATCH 7/8] =?UTF-8?q?fix(ck-tile):=20multi=5Fabd=20bridge=20hard?= =?UTF-8?q?ening=20=E2=80=94=20explicit=20size=5Ft=20byte-size=20casts,=20?= =?UTF-8?q?32-bit=20index/stride=20overflow=20guard,=20=5Fparse=5Fops=20ra?= =?UTF-8?q?ises=20on=20unparseable=20name?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- .../ctypes/gemm_multi_abd_ctypes_lib.cpp | 51 +++++++++++++++++-- .../dispatcher/python/gemm_utils.py | 10 +++- 2 files changed, 55 insertions(+), 6 deletions(-) diff --git a/projects/composablekernel/dispatcher/bindings/ctypes/gemm_multi_abd_ctypes_lib.cpp b/projects/composablekernel/dispatcher/bindings/ctypes/gemm_multi_abd_ctypes_lib.cpp index 3a518971d7d4..af7b6f77caa2 100644 --- a/projects/composablekernel/dispatcher/bindings/ctypes/gemm_multi_abd_ctypes_lib.cpp +++ b/projects/composablekernel/dispatcher/bindings/ctypes/gemm_multi_abd_ctypes_lib.cpp @@ -222,10 +222,16 @@ int dispatcher_run_multi_abd(const void** as_hosts, } } - const size_t a_bytes = static_cast(M) * K * elem_a; - const size_t b_bytes = static_cast(K) * N * elem_b; - const size_t d_bytes = static_cast(M) * N * elem_d; - const size_t e_bytes = static_cast(M) * N * elem_e; + // Cast every factor to size_t so the products are computed in 64-bit + // unsigned arithmetic (no reliance on operand-promotion order). + const size_t a_bytes = + static_cast(M) * static_cast(K) * static_cast(elem_a); + const size_t b_bytes = + static_cast(K) * static_cast(N) * static_cast(elem_b); + const size_t d_bytes = + static_cast(M) * static_cast(N) * static_cast(elem_d); + const size_t e_bytes = + static_cast(M) * static_cast(N) * static_cast(elem_e); std::vector a_dev(kNumA, nullptr), b_dev(kNumB, nullptr), d_dev(kNumD, nullptr); void* e_dev = nullptr; @@ -304,7 +310,42 @@ int dispatcher_run_multi_abd(const void** as_hosts, std::array str_bs{}; std::array str_ds{}; - // Strides validated non-null and strictly positive above, so pack directly. + // ck_tile::index_t is 32-bit: reject any dimension or stride that would not + // fit before the static_casts below, so an out-of-range value fails loudly + // instead of silently truncating. + { + auto fits_i32 = [](int64_t v) { return v <= static_cast(INT32_MAX); }; + if(!fits_i32(M) || !fits_i32(N) || !fits_i32(K) || !fits_i32(stride_e)) + { + std::cerr << "dispatcher_run_multi_abd: M/N/K or stride_e exceeds the 32-bit " + "index range\n"; + return -1; + } + for(ck_tile::index_t i = 0; i < kNumA; ++i) + if(!fits_i32(stride_as[i])) + { + std::cerr << "dispatcher_run_multi_abd: stride_as exceeds the 32-bit index " + "range\n"; + return -1; + } + for(ck_tile::index_t i = 0; i < kNumB; ++i) + if(!fits_i32(stride_bs[i])) + { + std::cerr << "dispatcher_run_multi_abd: stride_bs exceeds the 32-bit index " + "range\n"; + return -1; + } + for(ck_tile::index_t i = 0; i < kNumD; ++i) + if(!fits_i32(stride_ds[i])) + { + std::cerr << "dispatcher_run_multi_abd: stride_ds exceeds the 32-bit index " + "range\n"; + return -1; + } + } + + // Strides validated non-null, strictly positive, and 32-bit-safe above, so + // pack directly. for(ck_tile::index_t i = 0; i < kNumA; ++i) { as[i] = a_dev[i]; diff --git a/projects/composablekernel/dispatcher/python/gemm_utils.py b/projects/composablekernel/dispatcher/python/gemm_utils.py index 93e551a6adb8..d5d867a8bd88 100644 --- a/projects/composablekernel/dispatcher/python/gemm_utils.py +++ b/projects/composablekernel/dispatcher/python/gemm_utils.py @@ -1249,7 +1249,15 @@ def _parse_ops(self) -> Tuple[str, str, str]: # tail == [aNA, bNB, dND, Aop, Bop, CDEop] if len(tail) >= 6: return tail[3], tail[4], tail[5] - return "PassThrough", "PassThrough", "PassThrough" + # Raise rather than silently defaulting to PassThrough (which would yield + # a wrong numeric reference): the config is the source of truth and the + # name is deterministic from it, so an unparseable name is a real error. + # Mirrors _parse_layout4, which also raises. + raise ValueError( + f"cannot parse multi_abd element-wise ops from kernel name " + f"{self._kernel_name!r}; expected a " + f"'..._multiabd_a_b_d___' suffix" + ) def run( self, From 1e4cba2707d67096e95382235e4e4c6305fc0a78 Mon Sep 17 00:00:00 2001 From: muozturk Date: Wed, 22 Jul 2026 22:51:56 -0400 Subject: [PATCH 8/8] fix(ck-tile): gate multi_abd bridge sweep on Old-TE warp-config constraint expand_sweep emitted the full wave_m x wave_n x wave_k cartesian product but never applied Old-TE's validate_warp_configuration gate (WARP_SUPPORTED_COMBINATIONS[arch], which allows only [1,4,1]/[2,2,1]/[4,1,1] on CDNA). The bridge therefore generated warps-per-block triples (product 1/2/16) that Old-TE never emits, diverging the two instance sets (gfx950 whole-config: bridge 5308 vs Old-TE 4752 buildable, overlap only 1920). Add _warp_config_supported() (reusing the canonical table from gemm_validation_utils with a hardcoded fallback) and gate the emit loop on it, right after the tile-divisibility check, so the bridge and Old-TE produce the same warp configurations for a given config+arch. --- .../dispatcher/python/gemm_utils.py | 37 +++++++++++++++++++ 1 file changed, 37 insertions(+) diff --git a/projects/composablekernel/dispatcher/python/gemm_utils.py b/projects/composablekernel/dispatcher/python/gemm_utils.py index 5d7451eb885c..14cac7fc117c 100644 --- a/projects/composablekernel/dispatcher/python/gemm_utils.py +++ b/projects/composablekernel/dispatcher/python/gemm_utils.py @@ -1891,6 +1891,39 @@ def _is_power_of_two(x: int) -> bool: return x > 0 and (x & (x - 1)) == 0 +# --- 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]. +# expand_sweep must apply the SAME gate or the bridge emits wave configs Old-TE +# never generates (product != 4 on CDNA), diverging the two instance sets. +_WARP_SUPPORTED_COMBINATIONS_FALLBACK = { + "gfx90a": [[1, 4, 1], [2, 2, 1], [4, 1, 1]], + "gfx942": [[1, 4, 1], [2, 2, 1], [4, 1, 1]], + "gfx950": [[1, 4, 1], [2, 2, 1], [4, 1, 1]], + "gfx1201": [[2, 4, 1], [1, 8, 1], [8, 1, 1], [4, 2, 1]], +} + + +def _warp_supported_table(): + """Canonical WARP_SUPPORTED_COMBINATIONS from gemm_validation_utils, with a + hardcoded fallback so the bridge never silently skips the gate.""" + try: + from gemm_validation_utils import WARP_SUPPORTED_COMBINATIONS as _t + return _t + except Exception: # pragma: no cover - fallback keeps the gate active + return _WARP_SUPPORTED_COMBINATIONS_FALLBACK + + +def _warp_config_supported(wave_m: int, wave_n: int, wave_k: int, arch: str) -> bool: + """True iff [wave_m, wave_n, wave_k] is an allowed warps-per-block triple + for ``arch`` (mirrors Old-TE validate_warp_configuration). Unknown arch => + permissive (matches Old-TE's log-and-allow behavior).""" + allowed = _warp_supported_table().get(arch) + if not allowed: + return True + return [wave_m, wave_n, wave_k] in allowed + + def expand_sweep( config_path: str, arch: Optional[str] = None, @@ -2125,6 +2158,10 @@ def _as_list(v, default): n_div = wn * wtn if m_div <= 0 or n_div <= 0 or tm % m_div != 0 or tn % n_div != 0: continue + # Parity gate: only emit warps-per-block triples Old-TE allows + # (WARP_SUPPORTED_COMBINATIONS[arch]); see _warp_config_supported. + if not _warp_config_supported(wm, wn, wk, arch): + continue if epi == "cshuffle" and not _cshuffle_store_ok( tm // m_div, tn // n_div, wtm, wtn ):