diff --git a/projects/composablekernel/dispatcher/bindings/ctypes/gemm_ctypes_lib.cpp b/projects/composablekernel/dispatcher/bindings/ctypes/gemm_ctypes_lib.cpp index 85c0c2f2c13a..72d256302978 100644 --- a/projects/composablekernel/dispatcher/bindings/ctypes/gemm_ctypes_lib.cpp +++ b/projects/composablekernel/dispatcher/bindings/ctypes/gemm_ctypes_lib.cpp @@ -20,6 +20,7 @@ #include #include #include +#include #include "ck_tile/dispatcher/dispatcher.hpp" #include "ck_tile/dispatcher/registry.hpp" @@ -71,9 +72,20 @@ int dispatcher_initialize() key.signature.dtype_b = DataType::FP16; key.signature.dtype_c = DataType::FP16; key.signature.dtype_acc = DataType::FP32; - key.signature.layout_a = LayoutTag::RowMajor; - key.signature.layout_b = LayoutTag::ColMajor; - key.signature.layout_c = LayoutTag::RowMajor; + // Derive A/B/C layouts from the force-included kernel's own layout types + // instead of hardcoding rcr. The dispatcher's supports() gate is layout-aware + // (it only constrains a dimension that an operand's inner axis maps to), so a + // wrong key layout makes it reject valid problems -- e.g. a crr kernel does not + // gate K, but with a hardcoded rcr key supports() would apply rcr's K-gate and + // reject TileK=192 problems that Old-TE runs. ALayout/BLayout/CLayout are the + // global aliases exported by the kernel header under CK_TILE_SINGLE_KERNEL_INCLUDE. + using RowMajorLayout = ck_tile::tensor_layout::gemm::RowMajor; + key.signature.layout_a = + std::is_same_v ? LayoutTag::RowMajor : LayoutTag::ColMajor; + key.signature.layout_b = + std::is_same_v ? LayoutTag::RowMajor : LayoutTag::ColMajor; + key.signature.layout_c = + std::is_same_v ? LayoutTag::RowMajor : LayoutTag::ColMajor; key.signature.transpose_a = false; key.signature.transpose_b = false; key.signature.grouped = false; @@ -310,10 +322,40 @@ int dispatcher_run_gemm( } /** - * Get kernel information + * Get kernel information (legacy single-kernel ABI). + * + * Returns the compile-time KERNEL_NAME of the force-included kernel header. + * Kept for backward compatibility with one-kernel-per-.so callers. */ const char* dispatcher_get_kernel_name() { return KERNEL_NAME; } +/** + * Get the name of the kernel at a given registry index (multi-kernel ABI). + * + * Mirrors the conv/fmha ctypes libs: copies the index-th registered kernel's + * name into the caller-provided buffer so one .so can report a whole batch and + * be selected by name at runtime. Returns 0 on success, -1 on bad args or + * out-of-range index. + */ +int dispatcher_get_kernel_name_at(int index, char* buffer, int buffer_size) +{ + if(!buffer || buffer_size <= 0) + { + return -1; + } + + auto kernels = Registry::instance().get_all(); + if(index < 0 || index >= static_cast(kernels.size())) + { + return -1; + } + + std::string name = kernels[index]->get_name(); + std::strncpy(buffer, name.c_str(), static_cast(buffer_size) - 1); + buffer[buffer_size - 1] = '\0'; + return 0; +} + /** * Initialize dispatcher (alias) */ diff --git a/projects/composablekernel/dispatcher/bindings/ctypes/streamk_gemm_ctypes_lib.cpp b/projects/composablekernel/dispatcher/bindings/ctypes/streamk_gemm_ctypes_lib.cpp new file mode 100644 index 000000000000..52027b55b7d3 --- /dev/null +++ b/projects/composablekernel/dispatcher/bindings/ctypes/streamk_gemm_ctypes_lib.cpp @@ -0,0 +1,290 @@ +// Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +// SPDX-License-Identifier: MIT + +/** + * Stream-K GEMM Dispatcher ctypes Library + * + * Provides C API for Python ctypes integration for the STREAM-K GEMM variant. + * Kernel header included via -include at compile time. + * + * Stream-K is a single GEMM (one A/B/C, one M/N/K) like regular GEMM, so this + * lib keeps the exact same C ABI as gemm_ctypes_lib.cpp -- ``dispatcher_run_gemm`` + * takes host A/B/C and M/N/K. The difference is internal: the generated launch + * has a Stream-K-specific signature + * + * static float launch(const ck_tile::StreamKHostArgs& args, const stream_config& stream); + * + * which allocates the reduction workspace internally (DeviceMem) and uses the + * Atomic reduction strategy. The single-problem registry path + * (g_dispatcher->run / GemmHostArgs) and the generated_tile_backend wrapper both + * hard-code the plain GemmHostArgs launch, so this lib bypasses the registry and + * calls SelectedKernel::launch(args, stream) directly, reporting the kernel name + * from the compile-time KERNEL_NAME macro. + * + * Because the C ABI matches the regular lib, the Python side reuses + * GemmDispatcherLib / GpuGemmRunner unchanged -- only the .so internals differ. + * + * Usage from Python: + * lib = ctypes.CDLL("libdispatcher_streamk_gemm.so") + * lib.dispatcher_init() + * lib.dispatcher_run_gemm(...) + */ + +#include +#include +#include +#include +#include +#include +#include + +// Kernel header included via -include compiler flag (with CK_TILE_SINGLE_KERNEL_INCLUDE). +// Defines: ADataType, BDataType, CDataType, AccDataType, SelectedKernel, KERNEL_NAME +// and transitively brings in ck_tile::StreamKHostArgs and ck_tile::stream_config. + +// GPU architecture - can be overridden via -DGFX_ARCH="gfx90a" at compile time +#ifndef GFX_ARCH +#define GFX_ARCH "gfx942" +#endif + +static bool g_initialized = false; + +// Read an integer benchmark knob from the environment, falling back to +// `fallback` when unset or unparseable. +static int env_int(const char* name, int fallback) +{ + const char* v = std::getenv(name); + if(v == nullptr || *v == '\0') + return fallback; + char* end = nullptr; + const long out = std::strtol(v, &end, 10); + if(end == v) + return fallback; + return static_cast(out); +} + +// Read a boolean benchmark knob ("0"/"false"/"off", any case => false, else true). +static bool env_bool(const char* name, bool fallback) +{ + const char* v = std::getenv(name); + if(v == nullptr || *v == '\0') + return fallback; + std::string s(v); + for(char& c : s) + if(c >= 'A' && c <= 'Z') + c = static_cast(c - 'A' + 'a'); + return !(s == "0" || s == "false" || s == "off"); +} + +extern "C" { + +/** + * Initialize the stream-k GEMM library. + * + * The stream-k path does not use the dispatcher/registry (it launches the + * force-included kernel directly), so this is a lightweight no-op kept for ABI + * parity with the regular GEMM lib. Returns 0 on success. + */ +int dispatcher_initialize() +{ + g_initialized = true; + return 0; +} + +/** + * Initialize dispatcher (alias) + */ +int dispatcher_init() { return dispatcher_initialize(); } + +/** + * Run a Stream-K GEMM on GPU by launching the force-included kernel directly. + * + * hipMalloc A/B/C, copy A and B host->device, memset C (the Atomic reduction + * strategy accumulates into C, so it must start zeroed), build a + * ck_tile::StreamKHostArgs whose strides are derived from the kernel's actual + * ALayout/BLayout/CLayout (no layout hardcoding) and launch. The launch + * allocates the reduction workspace internally and resets C between timed + * iterations. C is then copied back. + * + * The host buffers must be laid out to match each operand's layout (the Python + * runner arranges A/B/C as RowMajor=C-contiguous, ColumnMajor=F-contiguous). + * + * Returns: 0 on success, -1 on HIP error / generic throw, -2 if the kernel + * reports the arguments are unsupported. + */ +int dispatcher_run_gemm( + const void* A, const void* B, void* C, int64_t M, int64_t N, int64_t K, float* time_ms) +{ + if(!g_initialized || !A || !B || !C || M <= 0 || N <= 0 || K <= 0) + { + return -1; + } + + const ADataType* A_host = static_cast(A); + const BDataType* B_host = static_cast(B); + CDataType* C_host = static_cast(C); + + ADataType* A_dev = nullptr; + BDataType* B_dev = nullptr; + CDataType* C_dev = nullptr; + + auto cleanup_gpu_mem = [&]() { + if(A_dev) + (void)hipFree(A_dev); + if(B_dev) + (void)hipFree(B_dev); + if(C_dev) + (void)hipFree(C_dev); + }; + + if(hipMalloc(&A_dev, M * K * sizeof(ADataType)) != hipSuccess) + { + cleanup_gpu_mem(); + return -1; + } + if(hipMalloc(&B_dev, K * N * sizeof(BDataType)) != hipSuccess) + { + cleanup_gpu_mem(); + return -1; + } + if(hipMalloc(&C_dev, M * N * sizeof(CDataType)) != hipSuccess) + { + cleanup_gpu_mem(); + return -1; + } + + if(hipMemcpy(A_dev, A_host, M * K * sizeof(ADataType), hipMemcpyHostToDevice) != hipSuccess) + { + cleanup_gpu_mem(); + return -1; + } + if(hipMemcpy(B_dev, B_host, K * N * sizeof(BDataType), hipMemcpyHostToDevice) != hipSuccess) + { + cleanup_gpu_mem(); + return -1; + } + if(hipMemset(C_dev, 0, M * N * sizeof(CDataType)) != hipSuccess) + { + cleanup_gpu_mem(); + return -1; + } + + // Strides are DERIVED from the kernel's actual layouts (ALayout/BLayout/CLayout + // come from the force-included generated header) -- nothing layout-specific is + // hardcoded, so every layout (rcr/rrr/ccr/crr/...) works. A RowMajor R x C + // matrix has leading dim C; a ColumnMajor one has leading dim R. + // A is M x K, B is K x N, C is M x N. + using RowMajor = ck_tile::tensor_layout::gemm::RowMajor; + const ck_tile::index_t lda = static_cast( + std::is_same_v ? K : M); + const ck_tile::index_t ldb = static_cast( + std::is_same_v ? N : K); + const ck_tile::index_t ldc = static_cast( + std::is_same_v ? N : M); + // k_batch is fixed to 1 inside StreamKHostArgs. + ck_tile::StreamKHostArgs args(static_cast(A_dev), + static_cast(B_dev), + static_cast(C_dev), + static_cast(M), + static_cast(N), + static_cast(K), + /*stride_A=*/lda, + /*stride_B=*/ldb, + /*stride_C=*/ldc); + + // Benchmark parameters. warmup/repeat default to old Tile Engine's values + // (warmup=50, repeat=100); a generous warmup keeps the GPU clock ramped, and + // 100 timed iterations give a stable median. These were the knobs behind the + // regular bridge's spurious "perf gap" (#8123): the old default of warmup=3/ + // repeat=10 measured a cold, un-ramped clock. Each knob is env-overridable so + // a caller can match another harness without recompiling. + // + // Divergence from the regular path (generated_tile_backend.hpp): flush_cache_ + // and rotating_count_ default OFF here. The Stream-K Atomic reduction + // accumulates into C, and the generated launch's launch_kernel_time_mask + // preprocess re-zeros only the original args.e_ptr -- rotating C across + // multiple buffers would leave the rotated copies un-zeroed and corrupt the + // accumulation. Leave rotating_count_=1 unless a caller knows the kernel + // re-zeros every rotated buffer. + ck_tile::stream_config stream_cfg; + stream_cfg.stream_id_ = nullptr; + stream_cfg.time_kernel_ = true; + stream_cfg.log_level_ = 0; + stream_cfg.cold_niters_ = env_int("CK_TILE_BENCH_WARMUP", 50); + stream_cfg.nrepeat_ = env_int("CK_TILE_BENCH_REPEAT", 100); + stream_cfg.is_gpu_timer_ = true; + stream_cfg.flush_cache_ = env_bool("CK_TILE_BENCH_FLUSH", false); + stream_cfg.rotating_count_ = env_int("CK_TILE_BENCH_ROTATING", 1); + + float exec_time = 0.0f; + try + { + exec_time = SelectedKernel::launch(args, stream_cfg); + } + catch(const std::exception& e) + { + cleanup_gpu_mem(); + if(std::string(e.what()).find("not supported") != std::string::npos) + { + if(time_ms) + { + *time_ms = -1.0f; + } + return -2; // Arguments not supported by this kernel + } + return -1; + } + + if(hipMemcpy(C_host, C_dev, M * N * sizeof(CDataType), hipMemcpyDeviceToHost) != hipSuccess) + { + cleanup_gpu_mem(); + return -1; + } + + if(time_ms) + { + *time_ms = exec_time; + } + + cleanup_gpu_mem(); + return 0; +} + +/** + * Get kernel information (legacy single-kernel ABI). + * + * Returns the compile-time KERNEL_NAME of the force-included kernel header. + */ +const char* dispatcher_get_kernel_name() { return KERNEL_NAME; } + +/** + * Get the name of the kernel at a given registry index (multi-kernel ABI). + * + * Each stream-k .so force-includes exactly one kernel header, so index 0 reports + * KERNEL_NAME and any other index is out of range. Mirrors the regular GEMM lib's + * name ABI so the Python bridge can use the same name-lookup path. + * Returns 0 on success, -1 on bad args or out-of-range index. + */ +int dispatcher_get_kernel_name_at(int index, char* buffer, int buffer_size) +{ + if(!buffer || buffer_size <= 0 || index != 0) + { + return -1; + } + + std::strncpy(buffer, KERNEL_NAME, static_cast(buffer_size) - 1); + buffer[buffer_size - 1] = '\0'; + return 0; +} + +/** + * Get the number of kernels in this .so (always 1 for the stream-k single-include lib). + */ +int dispatcher_get_kernel_count() { return 1; } + +/** + * Cleanup library resources (no-op; kept for ABI parity). + */ +void dispatcher_cleanup() { g_initialized = false; } + +} // extern "C" diff --git a/projects/composablekernel/dispatcher/codegen/arch_filter.py b/projects/composablekernel/dispatcher/codegen/arch_filter.py index 63dbee2dd762..41262ae7fa01 100644 --- a/projects/composablekernel/dispatcher/codegen/arch_filter.py +++ b/projects/composablekernel/dispatcher/codegen/arch_filter.py @@ -50,6 +50,7 @@ class OperatorType(Enum): GEMM = "gemm" GEMM_PRESHUFFLE = "gemm_preshuffle" GEMM_MULTI_D = "gemm_multi_d" + GEMM_STREAMK = "gemm_streamk" CONV_FWD = "conv_fwd" CONV_BWD_DATA = "conv_bwd_data" CONV_BWD_WEIGHT = "conv_bwd_weight" @@ -85,6 +86,14 @@ class OperatorType(Enum): "tile_n_alignment": 16, "tile_k_alignment": 8, }, + OperatorType.GEMM_STREAMK: { + "min_tile_m": 16, + "min_tile_n": 16, + "min_tile_k": 8, + "tile_m_alignment": 16, + "tile_n_alignment": 16, + "tile_k_alignment": 8, + }, OperatorType.CONV_FWD: { "min_tile_m": 1, # N dimension can be 1 "min_tile_n": 16, # K (output channels) should be reasonable diff --git a/projects/composablekernel/dispatcher/codegen/unified_gemm_codegen.py b/projects/composablekernel/dispatcher/codegen/unified_gemm_codegen.py index c0fb08aa4436..ea8ece8b53c9 100755 --- a/projects/composablekernel/dispatcher/codegen/unified_gemm_codegen.py +++ b/projects/composablekernel/dispatcher/codegen/unified_gemm_codegen.py @@ -198,6 +198,7 @@ class GemmVariant(Enum): STANDARD = "standard" PRESHUFFLE = "preshuffle" MULTI_D = "multi_d" + STREAM_K = "stream_k" # TileConfig imported from codegen_common @@ -223,6 +224,9 @@ class KernelConfig: elementwise_op: str = "PassThrough" num_d_tensors: int = 0 d_layout: str = "r" # Layout for D tensors (r=row, c=col) - same for all D tensors + # Stream-K reduction strategy: "atomic" (partials atomic-add into C), + # "linear", or "tree" (partials accumulate through a device workspace). + reduction_strategy: str = "atomic" # Fixed parameters block_size: int = 256 @@ -285,6 +289,11 @@ def key_name(self, datatype: str, layout: str) -> str: parts.append(f"nd{self.num_d_tensors}") parts.append(f"dly_{self.d_layout}") + # Stream-K variant: reduction strategy distinguishes otherwise-identical + # kernels (each strategy is a separate compiled binary). + if self.variant == GemmVariant.STREAM_K: + parts.append(f"redux_{self.reduction_strategy}") + # Occupancy parameters (only if non-default) if self.num_wave_groups != 1: parts.append(f"wg{self.num_wave_groups}") @@ -340,6 +349,12 @@ 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.STREAM_K: + name += "_streamk" + # Atomic keeps the bare "_streamk" suffix for name parity with the + # original single-strategy bridge; linear/tree are disambiguated. + if config.reduction_strategy != "atomic": + name += f"_{config.reduction_strategy}" return name @@ -393,6 +408,15 @@ def _header(self, kernel_name: str, config: KernelConfig) -> str: includes += """ #include "ck_tile/ops/gemm/pipeline/wp_pipeline_agmem_bgmem_creg_v2.hpp" #include "ck_tile/ops/gemm/pipeline/wp_pipeline_agmem_bgmem_creg_base_policy.hpp" +""" + + if config.variant == GemmVariant.STREAM_K: + includes += """ +#include +#include +#include "ck_tile/host/device_memory.hpp" +#include "ck_tile/ops/gemm/kernel/streamk_gemm/streamk_gemm_kernel.hpp" +#include "ck_tile/ops/gemm/kernel/streamk_gemm/streamk_gemm_tile_partitioner.hpp" """ return includes @@ -520,6 +544,9 @@ 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[self.tm.get_output_dtype(self.datatype)]}; using AccDataType = float; +using ALayout = {ns_name}::ALayout; +using BLayout = {ns_name}::BLayout; +using CLayout = {ns_name}::CLayout; #endif // CK_TILE_SINGLE_KERNEL_INCLUDE """ @@ -545,6 +572,8 @@ def _launch_function(self, config: KernelConfig) -> str: """Generate launch function""" if config.variant == GemmVariant.MULTI_D: return self._launch_function_multi_d(config) + if config.variant == GemmVariant.STREAM_K: + return self._launch_function_streamk(config) if config.preshuffle: return self._launch_function_preshuffle(config) return self._launch_function_standard(config) @@ -725,6 +754,137 @@ def _launch_function_multi_d(self, config: KernelConfig) -> str: return launch(multi_d_args, stream); }}""" + def _launch_function_streamk(self, config: KernelConfig) -> str: + """Generate launch function for Stream-K GEMM (the dispatcher way). + + Stream-K is a single GEMM that splits the K dimension across CUs and + reduces partial results through a device workspace. Unlike Tile Engine + (which takes an external workspace pointer), the dispatcher allocates the + workspace INTERNALLY via DeviceMem inside launch(args, stream). + + The reduction strategy is taken from the config (atomic/linear/tree). + Atomic: partial tiles atomic-add into C, so C is zeroed before every + kernel invocation. Linear/Tree: partials accumulate through the device + workspace, which is zeroed instead. Both are handled by the preprocess + callback passed to launch_kernel_time_mask. + """ + reduction_ck = { + "atomic": "Atomic", + "linear": "Linear", + "tree": "Tree", + }[config.reduction_strategy] + return f""" + // ---- Stream-K kernel type, hoisted to struct scope so the workspace API + // ---- (GetWorkSpaceSize + external-workspace launch) can reuse the same type. ---- + static constexpr auto SkScheduler = {self.tm.SCHEDULER_TO_CK[config.trait.scheduler]}; + static constexpr auto SkReductionStrategy = ck_tile::StreamKReductionStrategy::{reduction_ck}; + static constexpr int SkBlockPerCu = {config.k_block_per_cu}; + + using SkGemmUniversalTraits = TileGemmUniversalTraits; + using SkUniversalGemmProblem = UniversalGemmPipelineProblem< + ADataType, BDataType, AccDataType, TileShape, SkGemmUniversalTraits, SkScheduler>; + using SkGemmPipeline = {self.tm.PIPELINE_TO_CK[config.trait.pipeline]}; + {self._epilogue_code(config)} + using SkStreamKTilePartitioner = + ck_tile::StreamKTilePartitioner; + using StreamKGemmKernel = + ck_tile::StreamKKernel; + + // Device workspace (bytes) this kernel needs for `args`. 0 for Atomic; + // >0 for Linear/Tree. The Dispatcher uses this to size the buffer it owns. + static std::size_t GetWorkSpaceSize(const ck_tile::StreamKHostArgs& args) {{ + auto kargs = StreamKGemmKernel::MakeKernelArgs(args); + return StreamKGemmKernel::GetWorkSpaceSize(kargs); + }} + + // Whether the kernel can actually partition this problem (enough tiles across + // CUs). Lets the dispatcher's supports() reject too-small problems and fall + // back to a non-Stream-K kernel instead of throwing at launch. + static bool IsSupported(const ck_tile::StreamKHostArgs& args) {{ + return StreamKGemmKernel::IsSupportedArgument(StreamKGemmKernel::MakeKernelArgs(args)); + }} + + // Internal-workspace launch: allocates a fresh DeviceMem on every call. + // Kept unchanged for the bridge ctypes lib and the standalone 03 driver. + static float launch(const ck_tile::StreamKHostArgs& args, const stream_config& stream) {{ + auto kargs = StreamKGemmKernel::MakeKernelArgs(args); + const auto ws_size = StreamKGemmKernel::GetWorkSpaceSize(kargs); + ck_tile::DeviceMem workspace_dev(ws_size); + workspace_dev.SetZero(); + StreamKGemmKernel::SetWorkSpacePointer(kargs, workspace_dev.GetDeviceBuffer()); + + if (!StreamKGemmKernel::IsSupportedArgument(kargs)) {{ + throw std::runtime_error("Arguments not supported for stream-k kernel!"); + }} + + const dim3 grids = StreamKGemmKernel::GridSize(kargs.tile_partitioner); + const dim3 blocks = StreamKGemmKernel::BlockSize(); + + // Atomic reduction accumulates into C, so reset buffers before each run. + auto reset_data_buffers = [&]() {{ + if constexpr (SkReductionStrategy == ck_tile::StreamKReductionStrategy::Atomic) {{ + // Stride-aware: CLayout is row-major with stride_E elems/row, so a + // padded C is zeroed correctly (not just the contiguous M*N case). + (void)hipMemset2DAsync(args.e_ptr, + args.stride_E * sizeof(CDataType), + 0, + args.N * sizeof(CDataType), + args.M, + stream.stream_id_); + }} else {{ + workspace_dev.SetZero(); + }} + }}; + std::function preprocess = reset_data_buffers; + + float ave_time = launch_kernel_time_mask(stream, preprocess, + make_kernel(StreamKGemmKernel{{}}, grids, blocks, 0, kargs)); + return ave_time; + }} + + // External-workspace launch (PR-D): the Dispatcher owns and reuses the + // reduction buffer and passes it in. `workspace` may be null for Atomic + // (size 0). The per-iteration reset stays here because it needs CDataType + // and the reduction strategy, which the dtype-erased Dispatcher lacks. + static float launch(const ck_tile::StreamKHostArgs& args, const stream_config& stream, + void* workspace) {{ + auto kargs = StreamKGemmKernel::MakeKernelArgs(args); + const auto ws_size = StreamKGemmKernel::GetWorkSpaceSize(kargs); + if (workspace != nullptr) {{ + StreamKGemmKernel::SetWorkSpacePointer(kargs, workspace); + }} + + if (!StreamKGemmKernel::IsSupportedArgument(kargs)) {{ + throw std::runtime_error("Arguments not supported for stream-k kernel!"); + }} + + const dim3 grids = StreamKGemmKernel::GridSize(kargs.tile_partitioner); + const dim3 blocks = StreamKGemmKernel::BlockSize(); + + auto reset_data_buffers = [&]() {{ + if constexpr (SkReductionStrategy == ck_tile::StreamKReductionStrategy::Atomic) {{ + // Stride-aware: CLayout is row-major with stride_E elems/row, so a + // padded C is zeroed correctly (not just the contiguous M*N case). + (void)hipMemset2DAsync(args.e_ptr, + args.stride_E * sizeof(CDataType), + 0, + args.N * sizeof(CDataType), + args.M, + stream.stream_id_); + }} else {{ + (void)hipMemsetAsync(workspace, 0, ws_size, stream.stream_id_); + }} + }}; + std::function preprocess = reset_data_buffers; + + float ave_time = launch_kernel_time_mask(stream, preprocess, + make_kernel(StreamKGemmKernel{{}}, grids, blocks, 0, kargs)); + return ave_time; + }}""" + def _epilogue_code(self, config: KernelConfig) -> str: """Generate epilogue code""" if config.variant == GemmVariant.MULTI_D: @@ -776,12 +936,49 @@ def generate( output_dtype = self.tm.get_output_dtype(self.datatype) rel_path = kernel_path.relative_to(output_dir) + # Stream-K kernels need the Stream-K backend (StreamKHostArgs launch) and + # the SK key fields, so the registry can tell atomic/linear/tree apart and + # the right launch path compiles. All other variants use the regular backend. + is_streamk = config.variant == GemmVariant.STREAM_K + backend_inc = ( + "generated_tile_backend_streamk.hpp" + if is_streamk + else "generated_kernel_backend.hpp" + ) + + sk_fields = "" + if is_streamk: + rs = {"atomic": "Atomic", "linear": "Linear", "tree": "Tree"}[ + config.reduction_strategy + ] + ws = str(config.reduction_strategy != "atomic").lower() + sk_fields = f""" + key.algorithm.pad_m = {str(config.trait.pad_m).lower()}; + key.algorithm.pad_n = {str(config.trait.pad_n).lower()}; + key.algorithm.pad_k = {str(config.trait.pad_k).lower()}; + key.algorithm.streamk = true; + key.algorithm.reduction_strategy = ::ck_tile::dispatcher::ReductionStrategy::{rs}; + key.algorithm.workspace = {ws};""" + + if is_streamk: + ret_stmt = ( + "return backends::create_generated_streamk_kernel" + f'(key, "{kernel_name}");' + ) + else: + ret_stmt = ( + "return std::make_shared>" + f'(key, "{kernel_name}");' + ) + return f"""// SPDX-License-Identifier: MIT // Auto-generated dispatcher wrapper #pragma once #include "ck_tile/dispatcher.hpp" -#include "ck_tile/dispatcher/backends/generated_kernel_backend.hpp" +#include "ck_tile/dispatcher/backends/{backend_inc}" #include "{rel_path}" namespace ck_tile {{ @@ -832,11 +1029,11 @@ def generate( key.algorithm.persistent = {str(config.trait.persistent).lower()}; key.algorithm.preshuffle = {str(config.preshuffle).lower()}; key.algorithm.transpose_c = false; - key.algorithm.num_wave_groups = {config.num_wave_groups}; - + key.algorithm.num_wave_groups = {config.num_wave_groups};{sk_fields} + key.gfx_arch = gfx_arch; - - return std::make_shared>(key, "{kernel_name}"); + + {ret_stmt} }} }}}}}} @@ -941,6 +1138,10 @@ def _load_config(self, config_file: Optional[Path]) -> Dict: "elementwise_ops": ["MultiDAdd", "MultiDMultiply"], "num_d_tensors": [1, 2], }, + "streamk_config": { + # Each reduction strategy compiles to a separate kernel binary. + "reduction_strategy": ["atomic", "linear", "tree"], + }, } def generate_all(self, parallel: bool = True) -> Dict: @@ -1030,14 +1231,38 @@ def _get_configs_for_variant(self, variant: GemmVariant) -> List[KernelConfig]: trait_configs = self._get_trait_configs() for tile, trait in itertools.product(tile_configs, trait_configs): - # Perform variant-specific architecture validation + # Perform variant-specific architecture validation against the + # trait's ACTUAL pipeline/scheduler (not a hard-coded compv4). if self.arch_filter and HAS_ARCH_FILTER: - if not self._is_tile_arch_valid(tile, variant): + if not self._is_tile_arch_valid( + tile, + variant, + pipeline=trait.pipeline, + scheduler=trait.scheduler, + ): continue if variant == GemmVariant.STANDARD: configs.append(KernelConfig(tile=tile, trait=trait, variant=variant)) + elif variant == GemmVariant.STREAM_K: + # Stream-K reuses the standard trait space but requires the cshuffle + # epilogue (the only epilogue the stream-K kernel supports). Each + # reduction strategy (atomic/linear/tree) is a distinct compiled + # kernel, so we expand one config per requested strategy. + if trait.epilogue == "cshuffle": + streamk_cfg = self.config.get("streamk_config", {}) + strategies = streamk_cfg.get("reduction_strategy", ["atomic"]) + for reduction_strategy in strategies: + configs.append( + KernelConfig( + tile=tile, + trait=trait, + variant=variant, + reduction_strategy=reduction_strategy, + ) + ) + elif variant == GemmVariant.PRESHUFFLE: # Preshuffle needs specific pipeline (preshufflev2) and scheduler (default) # Skip configs that don't use preshuffle-compatible traits @@ -1105,9 +1330,21 @@ def _get_tile_configs(self) -> List[TileConfig]: rejected_count += 1 continue - # Architecture-specific validation + # Architecture-specific validation. This is a pre-filter run before + # tiles are paired with traits, so keep a tile if it is legal under + # ANY configured pipeline/scheduler; the precise per-trait check + # happens later in _get_configs_for_variant. Filtering here with a + # single hard-coded pipeline (compv4) wrongly dropped tiles that are + # legal under mem/compv3. if self.arch_filter and HAS_ARCH_FILTER: - if not self._is_tile_arch_valid(tile): + trait_cfg = self.config.get("trait_config", {}) + pipelines = trait_cfg.get("pipeline") or ["compv4"] + schedulers = trait_cfg.get("scheduler") or ["intrawave"] + if not any( + self._is_tile_arch_valid(tile, pipeline=pl, scheduler=sc) + for pl in pipelines + for sc in schedulers + ): rejected_count += 1 continue @@ -1119,13 +1356,23 @@ def _get_tile_configs(self) -> List[TileConfig]: return configs def _is_tile_arch_valid( - self, tile: TileConfig, variant: GemmVariant = None + self, + tile: TileConfig, + variant: GemmVariant = None, + pipeline: str = None, + scheduler: str = None, ) -> bool: """Check if tile configuration is valid for target architecture Args: tile: Tile configuration to validate variant: GEMM variant (affects operator-specific constraints) + pipeline: Trait pipeline to validate against. Pass the config's + actual pipeline -- omitting it falls back to ``compv4``, whose + MFMA constraints are stricter than ``mem``/``compv3`` and would + wrongly reject tiles that are legal under those pipelines. + scheduler: Trait scheduler to validate against (defaults to + ``intrawave`` for the same reason). """ if not self.arch_filter or not HAS_ARCH_FILTER: return True @@ -1146,14 +1393,17 @@ def _is_tile_arch_valid( # Map GEMM variant to operator type for validation operator = None - pipeline = "compv4" # Default - scheduler = "intrawave" # Default + if pipeline is None: + pipeline = "compv4" # Default (representative compute pipeline) + if scheduler is None: + scheduler = "intrawave" # Default if OperatorType is not None and variant is not None: variant_to_operator = { GemmVariant.STANDARD: OperatorType.GEMM, GemmVariant.PRESHUFFLE: OperatorType.GEMM_PRESHUFFLE, GemmVariant.MULTI_D: OperatorType.GEMM_MULTI_D, + GemmVariant.STREAM_K: OperatorType.GEMM_STREAMK, } operator = variant_to_operator.get(variant, OperatorType.GEMM) @@ -1403,7 +1653,7 @@ def main(): parser.add_argument( "--variants", nargs="+", - choices=["standard", "preshuffle", "multi_d"], + choices=["standard", "preshuffle", "multi_d", "stream_k"], default=["standard"], help="Variants to generate", ) diff --git a/projects/composablekernel/dispatcher/examples/gemm/cpp/03_streamk_gemm_driver.cpp b/projects/composablekernel/dispatcher/examples/gemm/cpp/03_streamk_gemm_driver.cpp new file mode 100644 index 000000000000..18271adb5ce3 --- /dev/null +++ b/projects/composablekernel/dispatcher/examples/gemm/cpp/03_streamk_gemm_driver.cpp @@ -0,0 +1,176 @@ +// Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +// SPDX-License-Identifier: MIT + +/** + * Minimal standalone stream-K GEMM driver (dispatcher way). + * + * Stream-K is a SINGLE GEMM that splits the K dimension across CUs and reduces + * the partial results through a device workspace. Like grouped GEMM it cannot + * ride the standard dispatcher.run(A,B,C,problem) path, so this driver includes + * a single generated stream-K kernel header (CK_TILE_SINGLE_KERNEL_INCLUDE) and + * calls SelectedKernel::launch(args, stream) directly with a single + * StreamKHostArgs -- the same 2-arg signature the dispatcher generates (the + * workspace is allocated INSIDE launch() via DeviceMem). It builds one A/B/C + * tensor, runs, and verifies against ck_tile::reference_gemm. + * + * Build (single-kernel include style): + * hipcc -std=c++17 --offload-arch=gfx942 -O3 \ + * -DCK_TILE_SINGLE_KERNEL_INCLUDE \ + * -I /include -I \ + * -include /_streamk.hpp \ + * 03_streamk_gemm_driver.cpp -o streamk_gemm_driver + */ + +#include + +#include +#include +#include +#include +#include + +#include "ck_tile/core.hpp" +#include "ck_tile/host.hpp" +#include "ck_tile/ops/gemm.hpp" + +// The generated stream-K kernel header is injected on the command line with +// -include and -DCK_TILE_SINGLE_KERNEL_INCLUDE. It exports into the global +// namespace: SelectedKernel, ADataType, BDataType, CDataType, AccDataType, +// ALayout, BLayout, CLayout, and KERNEL_NAME. + +template +static constexpr inline auto is_row_major(Layout) +{ + return ck_tile::bool_constant< + std::is_same_v, + ck_tile::tensor_layout::gemm::RowMajor>>{}; +} + +static std::string get_opt(int argc, char** argv, const std::string& key, const std::string& def) +{ + for(int i = 1; i < argc - 1; ++i) + if(key == argv[i]) return argv[i + 1]; + return def; +} + +int main(int argc, char** argv) +{ + const ck_tile::index_t M = std::stoll(get_opt(argc, argv, "--m", "3840")); + const ck_tile::index_t N = std::stoll(get_opt(argc, argv, "--n", "4096")); + const ck_tile::index_t K = std::stoll(get_opt(argc, argv, "--k", "2048")); + int warmup = std::stoi(get_opt(argc, argv, "--warmup", "50")); + int repeat = std::stoi(get_opt(argc, argv, "--repeat", "100")); + const bool validate = get_opt(argc, argv, "--validate", "1") != "0"; + + // Apple-to-apple with tile_engine: time the kernel with the SAME methodology the + // tile_engine benchmark uses (gemm_streamk_profiler.hpp) -- gpu timer and a + // cold-cache measurement that flushes the cache and rotates input buffers each + // iteration. tile_engine defaults: timer=true, flush_cache=true, rotating_count=1000. + // Without these the driver measured a warm-cache best case and over-reported TFlops, + // which is the entire source of the dispatcher-vs-TE "performance gap". + const bool gpu_timer = get_opt(argc, argv, "--timer", "1") != "0"; + bool flush_cache = get_opt(argc, argv, "--flush_cache", "1") != "0"; + int rotating_count = std::stoi(get_opt(argc, argv, "--rotating_count", "1000")); + + // Verification reads C back and compares against the reference for the known A/B. + // Rotating buffers and multi-repeat rotate/accumulate the output, so the C left on + // the device would not correspond to the reference inputs. tile_engine handles this + // with repeat_once_if_verify(); we mirror it -- a validating run times a single cold + // shot. Run a separate --validate 0 pass to collect apple-to-apple perf numbers. + if(validate) + { + warmup = 0; + repeat = 1; + flush_cache = false; + rotating_count = 1; + } + + std::cout << "Kernel: " << KERNEL_NAME << "\n"; + std::cout << "M=" << M << " N=" << N << " K=" << K << "\n"; + + const ck_tile::index_t sA = ck_tile::get_default_stride(M, K, 0, is_row_major(ALayout{})); + const ck_tile::index_t sB = ck_tile::get_default_stride(K, N, 0, is_row_major(BLayout{})); + const ck_tile::index_t sC = ck_tile::get_default_stride(M, N, 0, is_row_major(CLayout{})); + + ck_tile::HostTensor a_host( + ck_tile::host_tensor_descriptor(M, K, sA, is_row_major(ALayout{}))); + ck_tile::HostTensor b_host( + ck_tile::host_tensor_descriptor(K, N, sB, is_row_major(BLayout{}))); + ck_tile::HostTensor c_host( + ck_tile::host_tensor_descriptor(M, N, sC, is_row_major(CLayout{}))); + + ck_tile::FillUniformDistribution{-1.f, 1.f}(a_host); + ck_tile::FillUniformDistribution{-1.f, 1.f}(b_host); + c_host.SetZero(); + + ck_tile::DeviceMem a_dev(a_host); + ck_tile::DeviceMem b_dev(b_host); + ck_tile::DeviceMem c_dev(c_host); + c_dev.SetZero(); + + ck_tile::StreamKHostArgs args{a_dev.GetDeviceBuffer(), + b_dev.GetDeviceBuffer(), + c_dev.GetDeviceBuffer(), + M, + N, + K, + sA, + sB, + sC}; + + const ck_tile::stream_config s{ + nullptr, true, /*log=*/0, warmup, repeat, gpu_timer, flush_cache, rotating_count}; + float ave_time = SelectedKernel::launch(args, s); + + const std::size_t flop = std::size_t(2) * M * N * K; + const std::size_t bytes = sizeof(ADataType) * M * K + sizeof(BDataType) * K * N + + sizeof(CDataType) * M * N; + const float tflops = static_cast(flop) / 1.E9 / ave_time; + const float gbps = static_cast(bytes) / 1.E6 / ave_time; + std::cout << "Perf: " << std::setw(10) << ave_time << " ms, " << tflops << " TFlops, " << gbps + << " GB/s\n"; + + c_dev.FromDevice(c_host.data()); + + bool pass = true; + if(validate) + { + ck_tile::HostTensor ref( + ck_tile::host_tensor_descriptor(M, N, sC, is_row_major(CLayout{}))); + ref.SetZero(); + ck_tile::reference_gemm(a_host, b_host, ref); + const float maxv = *std::max_element(ref.mData.begin(), ref.mData.end()); + + // Stream-K splits K across CUs and reduces partials. Atomic reduction + // accumulates those partials directly into low-precision C, so the + // verification tolerance must account for the split-K accumulation + // error -- exactly as the tile_engine verifier does in + // tile_engine/include/utility/validation.hpp::calculate_rtol_atol. + // kbatch is the number of workgroups reducing into a single output + // tile, taken from the kernel's own tile partitioner so the driver and + // tile_engine agree on the split factor. + auto kargs = SelectedKernel::StreamKGemmKernel::MakeKernelArgs(args); + const ck_tile::index_t kbatch = + std::max(1, kargs.tile_partitioner.estimate_num_wgs_per_tile()); + + using ComputeType = + std::conditional_t; + const ck_tile::index_t k_per_split = ck_tile::integer_divide_ceil(K, kbatch); + // single-pass (per-split) tolerance + const auto rtol_base = + ck_tile::get_relative_threshold(k_per_split); + const auto atol_base = ck_tile::get_absolute_threshold( + maxv / kbatch, k_per_split); + // error contributed by reducing kbatch partials in low-precision C + const auto rtol_split_k = + ck_tile::get_relative_threshold(kbatch); + const auto atol_split_k = + ck_tile::get_absolute_threshold(maxv, kbatch); + const auto rtol = std::max(rtol_base, rtol_split_k); + const auto atol = std::max(atol_base, atol_split_k); + pass = ck_tile::check_err(c_host, ref, "streamk", rtol, atol); + std::cout << "Verification: " << (pass ? "PASS" : "FAIL") << "\n"; + } + + return pass ? 0 : 1; +} diff --git a/projects/composablekernel/dispatcher/examples/gemm/cpp/04_streamk_registry_driver.cpp b/projects/composablekernel/dispatcher/examples/gemm/cpp/04_streamk_registry_driver.cpp new file mode 100644 index 000000000000..bd8bc81bbc61 --- /dev/null +++ b/projects/composablekernel/dispatcher/examples/gemm/cpp/04_streamk_registry_driver.cpp @@ -0,0 +1,259 @@ +// Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +// SPDX-License-Identifier: MIT + +/** + * Stream-K GEMM driver through the Registry + Dispatcher (deep-core path). + * + * Unlike 03_streamk_gemm_driver.cpp (which calls SelectedKernel::launch() + * DIRECTLY, bypassing the dispatcher), this driver proves the full deep-core + * path that PR-A..PR-C built: + * + * Registry::register_kernel(GeneratedStreamKKernelInstance) + * -> Dispatcher::run(Problem.stream_k(Atomic)) + * -> Dispatcher::select_first_fit -> SK instance.supports() + * -> GeneratedStreamKKernelInstance::run -> SelectedKernel::launch() + * + * It registers ONE generated Stream-K kernel (force-included via + * -include / -DCK_TILE_SINGLE_KERNEL_INCLUDE), selects it through the registry + * by Problem::reduction_strategy, runs it, and verifies vs reference_gemm. + * + * Build (single-kernel include style): + * hipcc -std=c++17 --offload-arch=gfx942 -O3 \ + * -DCK_TILE_SINGLE_KERNEL_INCLUDE \ + * -I /include -I /dispatcher/include -I \ + * -include /_streamk.hpp \ + * 04_streamk_registry_driver.cpp -o streamk_registry_driver + */ + +#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/dispatcher/dispatcher.hpp" +#include "ck_tile/dispatcher/registry.hpp" +#include "ck_tile/dispatcher/backends/generated_tile_backend_streamk.hpp" + +// The generated stream-K kernel header is injected on the command line with +// -include and -DCK_TILE_SINGLE_KERNEL_INCLUDE. It exports into the global +// namespace: SelectedKernel, ADataType, BDataType, CDataType, AccDataType, +// ALayout, BLayout, CLayout, and KERNEL_NAME. + +#ifndef GFX_ARCH +#define GFX_ARCH "gfx942" +#endif + +using namespace ck_tile::dispatcher; +using namespace ck_tile::dispatcher::backends; +using Priority = ck_tile::dispatcher::Registry::Priority; + +template +static constexpr inline auto is_row_major(Layout) +{ + return ck_tile::bool_constant< + std::is_same_v, + ck_tile::tensor_layout::gemm::RowMajor>>{}; +} + +// Map a ck_tile element type to the dispatcher's DataType enum so the registry +// key reflects the kernel that was actually generated (fp16/bf16/fp8/bf8/...), +// instead of assuming fp16. Keeps the registry identifier and selection correct +// across every datatype the codegen emits. +template +static constexpr DataType dtype_enum_of() +{ + using U = ck_tile::remove_cvref_t; + if constexpr(std::is_same_v) + return DataType::FP16; + else if constexpr(std::is_same_v) + return DataType::BF16; + else if constexpr(std::is_same_v) + return DataType::FP8; + else if constexpr(std::is_same_v) + return DataType::BF8; + else if constexpr(std::is_same_v) + return DataType::INT8; + else if constexpr(std::is_same_v) + return DataType::FP32; + else + return DataType::UNKNOWN; +} + +template +static constexpr LayoutTag layout_tag_of() +{ + return std::is_same_v, + ck_tile::tensor_layout::gemm::RowMajor> + ? LayoutTag::RowMajor + : LayoutTag::ColMajor; +} + +static std::string get_opt(int argc, char** argv, const std::string& key, const std::string& def) +{ + for(int i = 1; i < argc - 1; ++i) + if(key == argv[i]) return argv[i + 1]; + return def; +} + +// Build the KernelKey for the force-included Stream-K kernel. Only the Stream-K +// axis (streamk + reduction_strategy) governs selection; the remaining fields +// are populated for a faithful encode_identifier()/registry entry. +static KernelKey make_streamk_key(ReductionStrategy strategy) +{ + KernelKey key; + key.signature.dtype_a = dtype_enum_of(); + key.signature.dtype_b = dtype_enum_of(); + key.signature.dtype_c = dtype_enum_of(); + key.signature.dtype_acc = dtype_enum_of(); + key.signature.layout_a = layout_tag_of(); + key.signature.layout_b = layout_tag_of(); + key.signature.layout_c = layout_tag_of(); + key.signature.transpose_a = false; + key.signature.transpose_b = false; + key.signature.grouped = false; + key.signature.split_k = 1; + key.signature.elementwise_op = "PassThrough"; + key.signature.num_d_tensors = 0; + key.signature.structured_sparsity = false; + + key.algorithm.tile_shape = {SelectedKernel::TileM, SelectedKernel::TileN, SelectedKernel::TileK}; + key.algorithm.warp_tile_shape = {static_cast(SelectedKernel::WarpTileM), + static_cast(SelectedKernel::WarpTileN), + static_cast(SelectedKernel::WarpTileK)}; + key.algorithm.wave_shape = {2, 2, 1}; + key.algorithm.pipeline = Pipeline::CompV3; + key.algorithm.scheduler = Scheduler::Intrawave; + key.algorithm.epilogue = Epilogue::CShuffle; + key.algorithm.block_size = 256; + key.algorithm.double_buffer = false; + key.algorithm.persistent = false; + key.algorithm.preshuffle = false; + key.algorithm.transpose_c = false; + key.algorithm.num_wave_groups = 1; + key.algorithm.pad_m = SelectedKernel::kPadM; + key.algorithm.pad_n = SelectedKernel::kPadN; + key.algorithm.pad_k = SelectedKernel::kPadK; + + // The Stream-K selection axis (the whole point of this path). + key.algorithm.streamk = true; + key.algorithm.reduction_strategy = strategy; + key.algorithm.workspace = (strategy != ReductionStrategy::Atomic); + + key.gfx_arch = GFX_ARCH; + return key; +} + +static ReductionStrategy parse_strategy(const std::string& s) +{ + if(s == "linear") return ReductionStrategy::Linear; + if(s == "tree") return ReductionStrategy::Tree; + return ReductionStrategy::Atomic; +} + +int main(int argc, char** argv) +{ + const ck_tile::index_t M = std::stoll(get_opt(argc, argv, "--m", "3840")); + const ck_tile::index_t N = std::stoll(get_opt(argc, argv, "--n", "4096")); + const ck_tile::index_t K = std::stoll(get_opt(argc, argv, "--k", "2048")); + const bool validate = get_opt(argc, argv, "--validate", "1") != "0"; + const ReductionStrategy strategy = + parse_strategy(get_opt(argc, argv, "--strategy", "atomic")); + + std::cout << "Kernel: " << KERNEL_NAME << "\n"; + std::cout << "M=" << M << " N=" << N << " K=" << K + << " strategy=" << to_string(strategy) << "\n"; + + // --- Register the kernel into the global registry --------------------------- + KernelKey key = make_streamk_key(strategy); + auto kernel = create_generated_streamk_kernel(key, KERNEL_NAME); + Registry::instance().clear(); + Registry::instance().register_kernel(kernel, Priority::High); + std::cout << "Registered kernels: " << Registry::instance().size() + << " identifier=" << key.encode_identifier() << "\n"; + + // --- Build the problem requesting THIS Stream-K strategy -------------------- + Problem problem(M, N, K); + problem.streamk = true; + problem.reduction_strategy = strategy; + + Dispatcher dispatcher; + auto selected = dispatcher.select_kernel(problem); + if(!selected) + { + std::cout << "Dispatcher selected NO kernel for the Stream-K problem -> FAIL\n"; + return 1; + } + std::cout << "Dispatcher selected: " << selected->get_name() << "\n"; + + // --- Tensors (rcr) --------------------------------------------------------- + const ck_tile::index_t sA = ck_tile::get_default_stride(M, K, 0, is_row_major(ALayout{})); + const ck_tile::index_t sB = ck_tile::get_default_stride(K, N, 0, is_row_major(BLayout{})); + const ck_tile::index_t sC = ck_tile::get_default_stride(M, N, 0, is_row_major(CLayout{})); + + ck_tile::HostTensor a_host( + ck_tile::host_tensor_descriptor(M, K, sA, is_row_major(ALayout{}))); + ck_tile::HostTensor b_host( + ck_tile::host_tensor_descriptor(K, N, sB, is_row_major(BLayout{}))); + ck_tile::HostTensor c_host( + ck_tile::host_tensor_descriptor(M, N, sC, is_row_major(CLayout{}))); + + ck_tile::FillUniformDistribution{-1.f, 1.f}(a_host); + ck_tile::FillUniformDistribution{-1.f, 1.f}(b_host); + c_host.SetZero(); + + ck_tile::DeviceMem a_dev(a_host); + ck_tile::DeviceMem b_dev(b_host); + ck_tile::DeviceMem c_dev(c_host); + c_dev.SetZero(); + + // --- Run through the dispatcher (registry -> Dispatcher::run -> SK backend) - + float ave_time = 0.f; + try + { + ave_time = dispatcher.run( + a_dev.GetDeviceBuffer(), b_dev.GetDeviceBuffer(), c_dev.GetDeviceBuffer(), problem); + } + catch(const std::exception& e) + { + std::cout << "Dispatcher::run threw: " << e.what() << " -> FAIL\n"; + return 1; + } + + const std::size_t flop = std::size_t(2) * M * N * K; + const std::size_t bytes = sizeof(ADataType) * M * K + sizeof(BDataType) * K * N + + sizeof(CDataType) * M * N; + const float tflops = static_cast(flop) / 1.E9 / ave_time; + const float gbps = static_cast(bytes) / 1.E6 / ave_time; + std::cout << "Perf: " << std::setw(10) << ave_time << " ms, " << tflops << " TFlops, " << gbps + << " GB/s\n"; + + c_dev.FromDevice(c_host.data()); + + bool pass = true; + if(validate) + { + ck_tile::HostTensor ref( + ck_tile::host_tensor_descriptor(M, N, sC, is_row_major(CLayout{}))); + ref.SetZero(); + ck_tile::reference_gemm(a_host, b_host, ref); + const float maxv = *std::max_element(ref.mData.begin(), ref.mData.end()); + const auto rtol = ck_tile::get_relative_threshold(K); + const auto atol = + ck_tile::get_absolute_threshold(maxv, K); + pass = ck_tile::check_err(c_host, ref, "streamk_registry", rtol, atol); + std::cout << "Verification: " << (pass ? "PASS" : "FAIL") << "\n"; + } + + return pass ? 0 : 1; +} diff --git a/projects/composablekernel/dispatcher/include/ck_tile/dispatcher/backends/generated_tile_backend.hpp b/projects/composablekernel/dispatcher/include/ck_tile/dispatcher/backends/generated_tile_backend.hpp index be22d94b3331..c724977f9613 100644 --- a/projects/composablekernel/dispatcher/include/ck_tile/dispatcher/backends/generated_tile_backend.hpp +++ b/projects/composablekernel/dispatcher/include/ck_tile/dispatcher/backends/generated_tile_backend.hpp @@ -12,6 +12,8 @@ #include #include #include +#include +#include namespace ck_tile { namespace dispatcher { @@ -50,26 +52,46 @@ class GeneratedTileKernelInstance : public KernelInstance bool supports(const Problem& problem) const override { - // Check dimension divisibility if padding not enabled + // Tile-divisibility gate, mirroring ck_tile::GemmKernel::IsSupportedArgument + // exactly. A dimension only needs to be a multiple of its tile size when an + // operand whose contiguous (inner) axis is that dimension participates AND + // padding for it is disabled. This is layout-dependent: + // + // layout RowMajor A -> inner axis K | layout ColMajor A -> inner axis M + // layout RowMajor B -> inner axis N | layout ColMajor B -> inner axis K + // layout RowMajor C -> inner axis N | layout ColMajor C -> inner axis M + // + // The old check blindly required M % TileM == 0 for every layout, which + // wrongly rejected e.g. rcr kernels (RowMajor A & C never gate M) on + // M-indivisible problems that Old-TE runs fine. Anything this lets through + // is still validated by the kernel's own IsSupportedArgument inside launch(), + // so the bridge stays a strict functional equivalent of Old-TE. constexpr bool pad_m = SelectedKernel::kPadM; constexpr bool pad_n = SelectedKernel::kPadN; constexpr bool pad_k = SelectedKernel::kPadK; - if(pad_m && pad_n && pad_k) - { - return true; // Padding enabled - supports any size - } - - // Check divisibility constexpr int tile_m = SelectedKernel::TileM; constexpr int tile_n = SelectedKernel::TileN; constexpr int tile_k = SelectedKernel::TileK; - if(!pad_m && problem.M % tile_m != 0) + const auto is_row = [](LayoutTag l) { return l == LayoutTag::RowMajor; }; + const bool row_a = is_row(key_.signature.layout_a); + const bool row_b = is_row(key_.signature.layout_b); + const bool row_c = is_row(key_.signature.layout_c); + + // Which problem dimensions are actually constrained for this layout combo. + const bool require_m = (!row_a) || (!row_c); // ColMajor A or C gate M + const bool require_n = row_b || row_c; // RowMajor B or C gate N + const bool require_k = row_a || (!row_b); // RowMajor A or ColMajor B gate K + + const std::int64_t k_grain = + static_cast(tile_k) * (problem.k_batch > 0 ? problem.k_batch : 1); + + if(require_m && !pad_m && problem.M % tile_m != 0) return false; - if(!pad_n && problem.N % tile_n != 0) + if(require_n && !pad_n && problem.N % tile_n != 0) return false; - if(!pad_k && problem.K % tile_k != 0) + if(require_k && !pad_k && problem.K % k_grain != 0) return false; return true; @@ -101,16 +123,30 @@ class GeneratedTileKernelInstance : public KernelInstance problem.N // stride_E/C (row-major C: stride = N) ); + // Benchmark parameters. Defaults mirror old Tile Engine's + // gemm_common.hpp (warmup=50, repeat=100, flush_cache=true, + // rotating_count=1000), and a generous warmup keeps the GPU clock + // ramped. NOTE: matching these knobs does NOT by itself make + // bridge-vs-old-TE numbers comparable -- the byte-identical kernel + // measures ~18-20% faster here than through old TE's *standalone + // benchmark binary* at e.g. 1024^3/compv4, purely because that + // separate process runs the kernel at a lower sustained SCLK (+ more + // memory-stall cycles), not because of any bench knob, compiler, or + // kernel difference (rocprof-confirmed). For an honest A/B, measure + // BOTH kernels through the SAME harness (build the old-TE kernel into a + // .so and run it via run_one_gemm_kernel.py) -- the gap then collapses + // to ~1%. Each knob is env-overridable so a caller can match another + // harness without recompiling. const bool bench = this->benchmarking_; ck_tile::stream_config stream_cfg; stream_cfg.stream_id_ = reinterpret_cast(stream); stream_cfg.time_kernel_ = bench; stream_cfg.log_level_ = 0; - stream_cfg.cold_niters_ = bench ? 5 : 0; - stream_cfg.nrepeat_ = bench ? 10 : 1; + stream_cfg.cold_niters_ = bench ? env_int("CK_TILE_BENCH_WARMUP", 50) : 0; + stream_cfg.nrepeat_ = bench ? env_int("CK_TILE_BENCH_REPEAT", 100) : 1; stream_cfg.is_gpu_timer_ = bench; - stream_cfg.flush_cache_ = false; - stream_cfg.rotating_count_ = 1; + stream_cfg.flush_cache_ = bench && env_bool("CK_TILE_BENCH_FLUSH", true); + stream_cfg.rotating_count_ = bench ? env_int("CK_TILE_BENCH_ROTATING", 1000) : 1; // Call the generated kernel's launch method return SelectedKernel::launch(args, stream_cfg); @@ -134,6 +170,33 @@ class GeneratedTileKernelInstance : public KernelInstance } private: + // Read an integer benchmark knob from the environment, falling back to + // `fallback` when unset or unparseable. + static int env_int(const char* name, int fallback) + { + const char* v = std::getenv(name); + if(v == nullptr || *v == '\0') + return fallback; + char* end = nullptr; + const long out = std::strtol(v, &end, 10); + if(end == v) + return fallback; + return static_cast(out); + } + + // Read a boolean benchmark knob ("0"/"false"/"off", any case => false, else true). + static bool env_bool(const char* name, bool fallback) + { + const char* v = std::getenv(name); + if(v == nullptr || *v == '\0') + return fallback; + std::string s(v); + for(char& c : s) + if(c >= 'A' && c <= 'Z') + c = static_cast(c - 'A' + 'a'); + return !(s == "0" || s == "false" || s == "off"); + } + KernelKey key_; std::string name_; }; diff --git a/projects/composablekernel/dispatcher/include/ck_tile/dispatcher/backends/generated_tile_backend_streamk.hpp b/projects/composablekernel/dispatcher/include/ck_tile/dispatcher/backends/generated_tile_backend_streamk.hpp new file mode 100644 index 000000000000..166fc4ca1d9e --- /dev/null +++ b/projects/composablekernel/dispatcher/include/ck_tile/dispatcher/backends/generated_tile_backend_streamk.hpp @@ -0,0 +1,212 @@ +// Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +// SPDX-License-Identifier: MIT + +#pragma once + +#include "ck_tile/dispatcher/kernel_instance.hpp" +#include "ck_tile/core.hpp" +#include "ck_tile/host.hpp" +#include "ck_tile/ops/gemm.hpp" +#include "ck_tile/ops/gemm/kernel/streamk_gemm/streamk_gemm_kernel.hpp" +#include +#include +#include + +namespace ck_tile { +namespace dispatcher { +namespace backends { + +/** + * Kernel-instance wrapper for unified_gemm_codegen.py Stream-K kernels. + * + * Counterpart of GeneratedTileKernelInstance (regular GEMM) for the Stream-K + * variant. The difference is the host-args type: Stream-K needs + * ck_tile::StreamKHostArgs (workspace pointer + reduction strategy), which is + * ABI-incompatible with the GemmHostArgs path -- this is exactly why Stream-K + * could not previously ride the registry. With this backend it can: the + * Dispatcher selects the instance by KernelKey (which now carries streamk + + * reduction_strategy) and calls run(). + * + * supports() gates on the requested reduction strategy so that the registry can + * hold atomic/linear/tree side by side and the Dispatcher's first-fit selection + * picks the one the caller asked for via Problem::reduction_strategy. + * + * NOTE (PR-C): the generated SelectedKernel::launch(StreamKHostArgs, stream) + * still owns the reduction workspace internally (DeviceMem) and does the + * per-iter reset. PR-D relocates workspace ownership + reset to Dispatcher::run() + * via get_workspace_size()/the workspace-aware run() overload. + */ +template +class GeneratedStreamKKernelInstance : public KernelInstance +{ + public: + using ADataType = ADataType_; + using BDataType = BDataType_; + using CDataType = CDataType_; + using AccDataType = AccDataType_; + using SelectedKernel = SelectedKernelType; + + GeneratedStreamKKernelInstance(const KernelKey& key, const std::string& name) + : key_(key), name_(name) + { + } + + const KernelKey& get_key() const override { return key_; } + + std::string get_name() const override { return name_; } + + /// Accept ONLY when the caller requested a Stream-K kernel with THIS + /// instance's reduction strategy. Lets atomic/linear/tree coexist in the + /// registry and be selected by Problem::reduction_strategy. + bool supports(const Problem& problem) const override + { + if(!problem.streamk) + return false; + if(problem.reduction_strategy != key_.algorithm.reduction_strategy) + return false; + + // Stream-K distributes K-iterations across workgroups; padding flags + // mirror the regular backend's divisibility guard. + constexpr bool pad_m = SelectedKernel::kPadM; + constexpr bool pad_n = SelectedKernel::kPadN; + constexpr bool pad_k = SelectedKernel::kPadK; + if(!pad_m && problem.M % SelectedKernel::TileM != 0) + return false; + if(!pad_n && problem.N % SelectedKernel::TileN != 0) + return false; + if(!pad_k && problem.K % SelectedKernel::TileK != 0) + return false; + + // Final feasibility: enough tiles to partition across CUs. Rejecting here + // (instead of throwing at launch) lets the dispatcher's first-fit fall back + // to a non-Stream-K kernel for too-small problems. + return SelectedKernel::IsSupported(make_args(problem)); + } + + /// Device workspace (bytes) needed for `problem`. 0 for Atomic; >0 for + /// Linear/Tree. The Dispatcher uses this to size the buffer it owns and then + /// passes that buffer to the workspace-aware run() below. + std::size_t get_workspace_size(const Problem& problem) const override + { + return SelectedKernel::GetWorkSpaceSize(make_args(problem)); + } + + /// No-workspace entry point: delegates to the workspace-aware overload with a + /// null buffer, so the generated launch() falls back to its internal + /// (self-allocating) path. Used when the caller does not own a workspace. + float run(const void* a_ptr, + const void* b_ptr, + void* c_ptr, + const void** d_ptrs, + const Problem& problem, + void* stream) const override + { + return run(a_ptr, b_ptr, c_ptr, d_ptrs, /*workspace=*/nullptr, problem, stream); + } + + /// Workspace-aware execution (PR-D). `workspace` is the Dispatcher-owned + /// reduction buffer (may be null for Atomic, which needs none). When non-null + /// the generated launch() binds it instead of allocating its own DeviceMem. + float run(const void* a_ptr, + const void* b_ptr, + void* c_ptr, + const void** d_ptrs, + void* workspace, + const Problem& problem, + void* stream) const override + { + (void)d_ptrs; // Not used for Stream-K GEMM + + auto args = make_args(problem, a_ptr, b_ptr, c_ptr); + + const bool bench = this->benchmarking_; + ck_tile::stream_config stream_cfg; + stream_cfg.stream_id_ = reinterpret_cast(stream); + stream_cfg.time_kernel_ = bench; + stream_cfg.log_level_ = 0; + stream_cfg.cold_niters_ = bench ? 5 : 0; + stream_cfg.nrepeat_ = bench ? 10 : 1; + stream_cfg.is_gpu_timer_ = bench; + stream_cfg.flush_cache_ = false; + stream_cfg.rotating_count_ = 1; // atomic accumulates into C; never rotate + + if(workspace != nullptr) + return SelectedKernel::launch(args, stream_cfg, workspace); + return SelectedKernel::launch(args, stream_cfg); + } + + bool validate(const void* a_ptr, + const void* b_ptr, + const void* c_ptr, + const void** d_ptrs, + const Problem& problem, + float tolerance) const override + { + (void)a_ptr; + (void)b_ptr; + (void)c_ptr; + (void)d_ptrs; + (void)problem; + (void)tolerance; + return true; // reference validation handled by the TE/driver harness + } + + private: + /// Build StreamKHostArgs for `problem`. Strides are DERIVED from the kernel's + /// actual layouts (the force-included single kernel exposes the global + /// ALayout/BLayout/CLayout), not hardcoded to rcr. k_batch is owned by the + /// Stream-K tile partitioner, not passed here. Pointers default to null for + /// sizing-only use (GetWorkSpaceSize). StreamKHostArgs uses ck_tile::index_t + /// (int32); cast from Problem's int64. + ck_tile::StreamKHostArgs make_args(const Problem& problem, + const void* a_ptr = nullptr, + const void* b_ptr = nullptr, + void* c_ptr = nullptr) const + { + using idx = ck_tile::index_t; + using RowMajor = ck_tile::tensor_layout::gemm::RowMajor; + // A is MxK, B is KxN, C is MxN; RowMajor RxC has leading dim C, else R. + const idx lda = static_cast( + std::is_same_v<::ALayout, RowMajor> ? problem.K : problem.M); + const idx ldb = static_cast( + std::is_same_v<::BLayout, RowMajor> ? problem.N : problem.K); + const idx ldc = static_cast( + std::is_same_v<::CLayout, RowMajor> ? problem.N : problem.M); + return ck_tile::StreamKHostArgs{a_ptr, + b_ptr, + c_ptr, + static_cast(problem.M), + static_cast(problem.N), + static_cast(problem.K), + lda, + ldb, + ldc}; + } + + KernelKey key_; + std::string name_; +}; + +/// Helper to create a Stream-K kernel-instance wrapper. +template +std::shared_ptr create_generated_streamk_kernel(const KernelKey& key, + const std::string& name) +{ + return std::make_shared>(key, name); +} + +} // namespace backends +} // namespace dispatcher +} // namespace ck_tile diff --git a/projects/composablekernel/dispatcher/include/ck_tile/dispatcher/backends/tile_backend.hpp b/projects/composablekernel/dispatcher/include/ck_tile/dispatcher/backends/tile_backend.hpp index a3a0b0468562..e709c00e153f 100644 --- a/projects/composablekernel/dispatcher/include/ck_tile/dispatcher/backends/tile_backend.hpp +++ b/projects/composablekernel/dispatcher/include/ck_tile/dispatcher/backends/tile_backend.hpp @@ -29,27 +29,37 @@ class TileKernelInstance : public KernelInstance bool supports(const Problem& problem) const override { - // Check dimension divisibility if padding not enabled + // Tile-divisibility gate, layout-aware to match + // ck_tile::GemmKernel::IsSupportedArgument (see generated_tile_backend.hpp + // for the full rationale). A dimension is only constrained when an operand + // whose inner axis is that dimension participates and its padding is off: + // RowMajor A->K, ColMajor A->M; RowMajor B->N, ColMajor B->K; + // RowMajor C->N, ColMajor C->M. constexpr bool pad_m = SelectedKernel::kPadM; constexpr bool pad_n = SelectedKernel::kPadN; constexpr bool pad_k = SelectedKernel::kPadK; - if(pad_m && pad_n && pad_k) - { - // Padding enabled - supports any size - return true; - } - - // Check divisibility constexpr int tile_m = SelectedKernel::TileM; constexpr int tile_n = SelectedKernel::TileN; constexpr int tile_k = SelectedKernel::TileK; - if(!pad_m && problem.M % tile_m != 0) + const auto is_row = [](LayoutTag l) { return l == LayoutTag::RowMajor; }; + const bool row_a = is_row(key_.signature.layout_a); + const bool row_b = is_row(key_.signature.layout_b); + const bool row_c = is_row(key_.signature.layout_c); + + const bool require_m = (!row_a) || (!row_c); + const bool require_n = row_b || row_c; + const bool require_k = row_a || (!row_b); + + const std::int64_t k_grain = + static_cast(tile_k) * (problem.k_batch > 0 ? problem.k_batch : 1); + + if(require_m && !pad_m && problem.M % tile_m != 0) return false; - if(!pad_n && problem.N % tile_n != 0) + if(require_n && !pad_n && problem.N % tile_n != 0) return false; - if(!pad_k && problem.K % tile_k != 0) + if(require_k && !pad_k && problem.K % k_grain != 0) return false; // Check shared memory budget if specified diff --git a/projects/composablekernel/dispatcher/include/ck_tile/dispatcher/dispatcher.hpp b/projects/composablekernel/dispatcher/include/ck_tile/dispatcher/dispatcher.hpp index d266d693daf8..7e4f532e3008 100644 --- a/projects/composablekernel/dispatcher/include/ck_tile/dispatcher/dispatcher.hpp +++ b/projects/composablekernel/dispatcher/include/ck_tile/dispatcher/dispatcher.hpp @@ -27,6 +27,7 @@ #include "ck_tile/dispatcher/kernel_instance.hpp" #include "ck_tile/dispatcher/problem.hpp" #include "ck_tile/dispatcher/registry.hpp" +#include #include #include #include @@ -56,6 +57,9 @@ class Dispatcher /// @param gfx_arch Target GPU architecture (e.g. "gfx950") explicit Dispatcher(Registry* registry = nullptr, const std::string& gfx_arch = ""); + /// Frees the dispatcher-owned Stream-K reduction workspace, if any. + ~Dispatcher(); + void set_arch(const std::string& arch) { gfx_arch_ = arch; } [[nodiscard]] const std::string& arch() const { return gfx_arch_; } @@ -149,6 +153,16 @@ class Dispatcher std::string gfx_arch_; bool benchmarking_ = true; + // Dispatcher-owned, grow-on-demand reduction workspace for Stream-K kernels + // (linear/tree). Sized via KernelInstance::get_workspace_size() and reused + // across calls so we don't hipMalloc/hipFree on the hot path. Held as a raw + // pointer to keep HIP/ck_tile out of this public header. + mutable void* workspace_ = nullptr; + mutable std::size_t workspace_bytes_ = 0; + + /// Ensure the owned workspace holds at least `bytes`, growing it if needed. + void ensure_workspace(std::size_t bytes) const; + /// Select kernel using first-fit strategy [[nodiscard]] KernelInstancePtr select_first_fit(const Problem& problem) const; diff --git a/projects/composablekernel/dispatcher/include/ck_tile/dispatcher/kernel_instance.hpp b/projects/composablekernel/dispatcher/include/ck_tile/dispatcher/kernel_instance.hpp index b6ef76e4f879..9b577ece3589 100644 --- a/projects/composablekernel/dispatcher/include/ck_tile/dispatcher/kernel_instance.hpp +++ b/projects/composablekernel/dispatcher/include/ck_tile/dispatcher/kernel_instance.hpp @@ -5,6 +5,7 @@ #include "ck_tile/dispatcher/kernel_key.hpp" #include "ck_tile/dispatcher/problem.hpp" +#include #include #include @@ -45,6 +46,30 @@ class KernelInstance const Problem& problem, void* stream = nullptr) const = 0; + /// Device workspace (in bytes) this kernel needs for `problem` (0 = none). + /// Non-zero only for Stream-K linear/tree reductions; the caller (Dispatcher) + /// sizes and owns the buffer and passes it to the workspace-aware run(). + [[nodiscard]] virtual std::size_t get_workspace_size(const Problem& problem) const + { + (void)problem; + return 0; + } + + /// Workspace-aware execution. Default forwards to the no-workspace run(), so + /// existing (non-Stream-K) kernels need no change; the Stream-K backend + /// overrides this to set the reduction workspace pointer before launch. + [[nodiscard]] virtual float run(const void* a_ptr, + const void* b_ptr, + void* c_ptr, + const void** d_ptrs, + void* workspace, + const Problem& problem, + void* stream = nullptr) const + { + (void)workspace; + return run(a_ptr, b_ptr, c_ptr, d_ptrs, problem, stream); + } + /// Validate kernel output against reference implementation /// @param a_ptr Pointer to matrix A (device memory) /// @param b_ptr Pointer to matrix B (device memory) diff --git a/projects/composablekernel/dispatcher/include/ck_tile/dispatcher/kernel_key.hpp b/projects/composablekernel/dispatcher/include/ck_tile/dispatcher/kernel_key.hpp index 24b20ecd9b8f..58a1f2b4e1bc 100644 --- a/projects/composablekernel/dispatcher/include/ck_tile/dispatcher/kernel_key.hpp +++ b/projects/composablekernel/dispatcher/include/ck_tile/dispatcher/kernel_key.hpp @@ -71,6 +71,30 @@ enum class Scheduler : std::uint8_t Interwave }; +/// Stream-K partial-sum reduction strategy. `None` = not a Stream-K kernel. +/// Mirrors ck_tile::StreamKReductionStrategy (Atomic/Linear/Tree). +enum class ReductionStrategy : std::uint8_t +{ + None = 0, + Atomic, + Linear, + Tree +}; + +/// Canonical lower-case name for a reduction strategy. Matches the codegen suffix +/// scheme (atomic -> "atomic", etc.) so callers/drivers share one spelling. +inline const char* to_string(ReductionStrategy r) +{ + switch(r) + { + case ReductionStrategy::Atomic: return "atomic"; + case ReductionStrategy::Linear: return "linear"; + case ReductionStrategy::Tree: return "tree"; + case ReductionStrategy::None: return "none"; + } + return "none"; +} + /// KernelKey: Compile-time kernel configuration metadata /// Organized into Signature (what operation) and Algorithm (how it's implemented) struct KernelKey @@ -146,6 +170,11 @@ struct KernelKey bool pad_m = true; // Support arbitrary M dimensions via padding bool pad_n = true; // Support arbitrary N dimensions via padding bool pad_k = true; // Support arbitrary K dimensions via padding + + // Stream-K (workgroup K-stream) parameters + bool streamk = false; // is a Stream-K kernel + ReductionStrategy reduction_strategy = ReductionStrategy::None; // atomic / linear / tree + bool workspace = false; // needs a device accumulation buffer (linear/tree) } algorithm; std::string gfx_arch; // e.g. "gfx942", "gfx90a", "gfx908" @@ -194,7 +223,10 @@ struct KernelKey algorithm.num_wave_groups, algorithm.pad_m, algorithm.pad_n, - algorithm.pad_k); + algorithm.pad_k, + algorithm.streamk, + algorithm.reduction_strategy, + algorithm.workspace); } /// Equality comparison @@ -436,6 +468,18 @@ inline std::string KernelKey::encode_identifier() const if(algorithm.preshuffle) oss << "_preshuffle"; + // Stream-K suffix -- must match unified_gemm_codegen.py KernelNaming.generate(): + // atomic -> "..._streamk" linear -> "..._streamk_linear" tree -> "..._streamk_tree" + // Guarded by algorithm.streamk so non-Stream-K identifiers stay byte-identical. + if(algorithm.streamk) + { + oss << "_streamk"; + if(algorithm.reduction_strategy == ReductionStrategy::Linear) + oss << "_linear"; + else if(algorithm.reduction_strategy == ReductionStrategy::Tree) + oss << "_tree"; + } + return oss.str(); } diff --git a/projects/composablekernel/dispatcher/include/ck_tile/dispatcher/problem.hpp b/projects/composablekernel/dispatcher/include/ck_tile/dispatcher/problem.hpp index 5bffb56b49ba..fedbc1d2cb27 100644 --- a/projects/composablekernel/dispatcher/include/ck_tile/dispatcher/problem.hpp +++ b/projects/composablekernel/dispatcher/include/ck_tile/dispatcher/problem.hpp @@ -7,6 +7,8 @@ #include #include +#include "ck_tile/dispatcher/kernel_key.hpp" // ReductionStrategy + namespace ck_tile { namespace dispatcher { @@ -58,6 +60,10 @@ struct Problem // Validation control bool enable_validation; // Enable output validation against reference + // Stream-K request: which reduction strategy the caller wants (None = non-Stream-K) + bool streamk = false; + ReductionStrategy reduction_strategy = ReductionStrategy::None; + /// Default constructor with sensible defaults Problem() : M(0), @@ -66,7 +72,9 @@ struct Problem k_batch(1), smem_budget(0), prefer_persistent(false), - enable_validation(false) + enable_validation(false), + streamk(false), + reduction_strategy(ReductionStrategy::None) { } @@ -78,7 +86,9 @@ struct Problem k_batch(1), smem_budget(0), prefer_persistent(false), - enable_validation(false) + enable_validation(false), + streamk(false), + reduction_strategy(ReductionStrategy::None) { } @@ -293,6 +303,14 @@ class ProblemBuilder return *this; } + /// Request a Stream-K kernel with a given reduction strategy + ProblemBuilder& stream_k(ReductionStrategy strategy = ReductionStrategy::Atomic) + { + problem_.streamk = true; + problem_.reduction_strategy = strategy; + return *this; + } + /// Build the Problem [[nodiscard]] Problem build() const { diff --git a/projects/composablekernel/dispatcher/parity_diag/regression/ab_efficient_sweep.py b/projects/composablekernel/dispatcher/parity_diag/regression/ab_efficient_sweep.py new file mode 100644 index 000000000000..c7a332e9b938 --- /dev/null +++ b/projects/composablekernel/dispatcher/parity_diag/regression/ab_efficient_sweep.py @@ -0,0 +1,163 @@ +#!/usr/bin/env python3 +"""Efficient A/B sweep: bridge .so vs Old-TE binary, all layouts + fp16/bf16. + +Faster successor to run_alllayout_sweep.py: the bridge side batches all shapes +for a stem into ONE run_one_gemm_kernel.py worker call (one Python+numpy+CDLL +startup per stem instead of one per measurement). Old-TE binaries are run once +per shape; their internal warmup=50/repeat=100 already yields a stable median, +matching the prior methodology. + +- Bridge .so : main worktree dispatcher/build/examples (built from the FIXED source). +- Old-TE bin : develop-parity worktree build/bin (develop branch), per user instruction. + +Writes allresult_fp16_bf16.csv with resume support (keyed on stem,shape). + +CSV fields: stem,pipeline,dtype,layout,shape,bridge_tflops,old_tflops,gap_pct, + bridge_verified,oldte_built +""" +import csv, json, os, re, subprocess, sys, time +from pathlib import Path + +ROOT = Path("/home/AMD/muozturk/New_project/rocm-libraries/projects/composablekernel") +DISP = ROOT / "dispatcher" +WORKER = ROOT / "tile_engine/ops/gemm/run_one_gemm_kernel.py" +SO_DIR = DISP / "build" / "examples" +GEN_DIR = DISP / "build" / "generated_kernels" +OLD_BIN_DIR = Path( + "/home/AMD/muozturk/New_project/rocm-libraries/.claude/worktrees" + "/develop-parity/projects/composablekernel/build/bin" +) +REG = DISP / "parity_diag" / "regression" +STEMS_FILE = REG / "stems_selected.txt" +CSV_OUT = REG / "allresult_fp16_bf16.csv" + +PYPATH = os.pathsep.join([str(DISP / "python"), str(ROOT / "tile_engine/ops/gemm")]) +DEVICE = os.environ.get("PARITY_DEVICE", "0") + +SHAPES = [(512, 512, 512), (1024, 1024, 1024), (2048, 2048, 2048), + (1024, 512, 256), (4096, 4096, 4096)] + +FIELDS = ["stem", "pipeline", "dtype", "layout", "shape", + "bridge_tflops", "old_tflops", "gap_pct", + "bridge_verified", "oldte_built"] + +_TFLOPS_RE = re.compile(r'"tflops\(TFlops\)":\s*([0-9.]+)') + + +def pipeline_of(stem): + for p in ("compv3", "compv4", "mem"): + if f"_{p}_" in stem: + return p + return "other" + + +def base_env(): + env = os.environ.copy() + env["HIP_VISIBLE_DEVICES"] = DEVICE + env["GEMM_PYPATH"] = PYPATH + env["LD_LIBRARY_PATH"] = "/opt/rocm/lib:" + env.get("LD_LIBRARY_PATH", "") + return env + + +def run_bridge_all(stem): + """One batched worker call over all SHAPES. Returns {shape_str: tflops|None}.""" + so = SO_DIR / f"libgemm_{stem}.so" + out = {f"{M}x{N}x{K}": None for (M, N, K) in SHAPES} + if not so.exists(): + return out + # Staleness guard: a .so older than its generated header was built from an + # obsolete codegen and must NOT be measured -- doing so reports phantom + # regressions (the big 256-tile gaps in allresult_fp16_bf16_2.csv were all + # stale binaries that recovered to parity on rebuild). Treat stale as missing. + hdr = GEN_DIR / f"gemm_{stem}.hpp" + if hdr.exists() and so.stat().st_mtime < hdr.stat().st_mtime: + print(f" STALE .so (older than header), skipping: {stem}", file=sys.stderr, flush=True) + return out + items = [{"so_path": str(so), "problem": {"M": M, "N": N, "K": K}, + "kernel_name": f"gemm_{stem}"} for (M, N, K) in SHAPES] + payload = json.dumps({"items": items, "verify": False}) + try: + p = subprocess.run([sys.executable, str(WORKER)], input=payload.encode(), + stdout=subprocess.PIPE, stderr=subprocess.DEVNULL, + env=base_env(), timeout=900) + except subprocess.TimeoutExpired: + return out + for line in p.stdout.decode().strip().splitlines(): + try: + d = json.loads(line) + except json.JSONDecodeError: + continue + idx = d.get("idx") + if isinstance(idx, int) and 0 <= idx < len(SHAPES) and d.get("ok"): + M, N, K = SHAPES[idx] + out[f"{M}x{N}x{K}"] = d.get("tflops") + return out + + +def run_oldte(stem, M, N, K): + binp = OLD_BIN_DIR / f"benchmark_gemm_universal_{stem}" + if not binp.exists(): + return None + try: + p = subprocess.run([str(binp), f"-m={M}", f"-n={N}", f"-k={K}", + "-warmup=50", "-repeat=100"], + stdout=subprocess.PIPE, stderr=subprocess.DEVNULL, + env=base_env(), timeout=300) + except subprocess.TimeoutExpired: + return None + m = _TFLOPS_RE.search(p.stdout.decode()) + return float(m.group(1)) if m else None + + +def main(): + stems = [l.strip() for l in STEMS_FILE.read_text().splitlines() if l.strip()] + total = len(stems) * len(SHAPES) + done = set() + if CSV_OUT.exists(): + with open(CSV_OUT) as f: + for row in csv.DictReader(f): + done.add((row["stem"], row["shape"])) + mode = "a" if done else "w" + print(f"stems={len(stems)} shapes={len(SHAPES)} total={total} resume={len(done)}", flush=True) + + t0 = time.time(); n = len(done) + with open(CSV_OUT, mode, newline="") as fh: + w = csv.DictWriter(fh, fieldnames=FIELDS) + if mode == "w": + w.writeheader() + for stem in stems: + shapes_todo = [(M, N, K) for (M, N, K) in SHAPES + if (stem, f"{M}x{N}x{K}") not in done] + if not shapes_todo: + continue + parts = stem.split("_") + dtype, layout = parts[0], parts[1] + pipeline = pipeline_of(stem) + oldte_built = (OLD_BIN_DIR / f"benchmark_gemm_universal_{stem}").exists() + + bridge = run_bridge_all(stem) + for (M, N, K) in shapes_todo: + shape = f"{M}x{N}x{K}" + bt = bridge.get(shape) + ot = run_oldte(stem, M, N, K) + if bt is not None and ot not in (None, 0): + gap = (bt - ot) / ot * 100.0 + else: + gap = float("nan") + w.writerow(dict( + stem=stem, pipeline=pipeline, dtype=dtype, layout=layout, shape=shape, + bridge_tflops=f"{bt:.4f}" if bt is not None else "nan", + old_tflops=f"{ot:.4f}" if ot is not None else "nan", + gap_pct=f"{gap:.4f}" if gap == gap else "nan", + bridge_verified="None", oldte_built=str(oldte_built))) + fh.flush() + n += 1 + el = time.time() - t0 + rate = (n - len(done)) / el if el > 0 else 0 + eta = (total - n) / rate / 3600 if rate > 0 else 0 + print(f"[{n}/{total}] {stem[:48]:48} rate={rate:.1f}/s ETA={eta:.1f}h", flush=True) + print(f"DONE rows={n} -> {CSV_OUT}", flush=True) + + +if __name__ == "__main__": + main() diff --git a/projects/composablekernel/dispatcher/parity_diag/regression/ab_same_harness.py b/projects/composablekernel/dispatcher/parity_diag/regression/ab_same_harness.py new file mode 100644 index 000000000000..919107174498 --- /dev/null +++ b/projects/composablekernel/dispatcher/parity_diag/regression/ab_same_harness.py @@ -0,0 +1,310 @@ +#!/usr/bin/env python3 +"""Apples-to-apples GEMM A/B: bridge kernel vs old-TE kernel, ONE harness. + +Why this exists +--------------- +The earlier sweep (allsweep6144rcrfp16.py) compared the bridge's dispatcher +measurement against old TE's *standalone benchmark binary* +(benchmark_gemm_universal_). That comparison is NOT apples-to-apples: +the device kernel is byte-identical, yet old TE's standalone binary reports +~18-20% lower TFLOPS at e.g. 1024^3 / compv4. rocprof shows the identical +kernel genuinely runs longer in that process -- ~+8% cycles plus a lower +sustained SCLK -- a power/clock + execution-environment artifact of that +binary, NOT a bridge speedup, compiler difference, or kernel difference. +(See diagnose.md sec.4.) + +This harness removes the artifact: it builds the OLD-TE kernel into a .so from +old TE's own generated header and runs BOTH the bridge kernel and the old-TE +kernel through the SAME worker (run_one_gemm_kernel.py). Measured this way the +gap collapses to ~1%, which is the honest result. + +The old-TE generated-header directory is derived per stem as +``///`` (e.g. fp16/rcr, bf16/crr), so a single +run covers every dtype/layout. Set OLD_TE_GEN to pin one explicit leaf dir for +all stems (legacy behavior); set OLD_TE_GEN_BASE to relocate the base. + +Usage: + python3 ab_same_harness.py # default kernel list + shapes + python3 ab_same_harness.py [...] # explicit stems + python3 ab_same_harness.py --stems-file F [--csv OUT] # sweep a stems file +""" +import argparse +import csv +import json +import os +import statistics +import subprocess +import sys +from pathlib import Path + +# composablekernel root: .../composablekernel/dispatcher/parity_diag/regression/ +ROOT = Path(__file__).resolve().parents[3] +DISP = ROOT / "dispatcher" +GEN = DISP / "build" / "generated_kernels" +SRC = DISP / "bindings" / "ctypes" / "gemm_ctypes_lib.cpp" +STATIC = DISP / "build" / "libck_tile_dispatcher.a" +BR_SO_DIR = DISP / "build" / "examples" +WORKER = ROOT / "tile_engine/ops/gemm/run_one_gemm_kernel.py" +# Base dir of old-TE generated single-kernel headers; the per-stem leaf +# (/) is appended in old_gen_dir(). Points at a sibling +# develop-parity worktree under the rocm-libraries root by default. +OLD_GEN_BASE = Path(os.environ.get( + "OLD_TE_GEN_BASE", + str(ROOT.parents[1] / ".claude/worktrees/develop-parity" + "/projects/composablekernel/build/tile_engine/ops/gemm/gemm_universal"), +)) +# Legacy explicit override: when set, this exact leaf dir is used for ALL stems. +OLD_GEN_PIN = os.environ.get("OLD_TE_GEN") +OUT = DISP / "parity_diag" / "regression" / "_ab_same_harness_build" +ARCH = os.environ.get("GFX_ARCH", "gfx942") +DEVICE = os.environ.get("PARITY_DEVICE", "0") +REPEATS = int(os.environ.get("AB_REPEATS", "3")) + +SHAPES = [(512, 512, 512), (1024, 1024, 1024), (2048, 2048, 2048), + (1024, 512, 256), (4096, 4096, 4096)] + +DEFAULT_STEMS = [ + "fp16_rcr_compv4_default_intrawave_False_False_False_False_64x128x64_2x2x1_32x32x16", + "fp16_rcr_compv4_cshuffle_intrawave_False_False_False_False_64x128x64_1x4x1_32x32x16", + "fp16_rcr_compv4_default_intrawave_False_False_False_False_128x128x64_4x1x1_32x32x16", +] + +PYPATH = os.pathsep.join([str(DISP / "python"), str(ROOT / "tile_engine/ops/gemm")]) + + +def old_gen_dir(stem: str) -> Path: + """Old-TE header dir for a stem: // (or the pinned dir). + + Stems are named ``__...`` (e.g. fp16_rcr_..., bf16_crr_...), + which is exactly the develop-parity gen-tree layout, so the leaf is derived + from the stem itself -- no per-layout hardcoding. + """ + if OLD_GEN_PIN: + return Path(OLD_GEN_PIN) + parts = stem.split("_") + dtype, layout = parts[0], parts[1] + return OLD_GEN_BASE / dtype / layout + + +def build_old_so(stem: str) -> Path | None: + """Compile old TE's generated kernel header into a bridge-loadable .so. + + Cached: if the .so already exists it is reused, so a parallel --build-only + pre-pass (CPU-bound hipcc) can be separated from the serial GPU measurement. + """ + hdr = old_gen_dir(stem) / f"gemm_universal_single_{stem}.hpp" + if not hdr.exists(): + return None + OUT.mkdir(parents=True, exist_ok=True) + obj = OUT / f"{stem}.o" + lib = OUT / f"libold_{stem}.so" + if lib.exists(): + return lib + common = [ + "-fPIC", "-O3", + f"-I{DISP / 'include'}", f"-I{ROOT / 'include'}", f"-I{ROOT}", f"-I{GEN}", + "-DCK_TILE_SINGLE_KERNEL_INCLUDE", f"-include{hdr}", "-D__HIP_PLATFORM_AMD__", + f"--offload-arch={ARCH}", f'-DGFX_ARCH="{ARCH}"', + # Match the bridge build's AMDGPU codegen flags (gemm_utils.py + # _build_compile_jobs / _TILE_ENGINE_CODEGEN_FLAGS), which are also what + # Tile Engine's own CMake passes. Without these the old-TE side is built + # with a *different* instruction schedule (notably -enable-post-misched + # defaults back on) and runs ~10-40% faster than real old-TE, making the + # bridge look regressed when it is actually at parity. Build BOTH sides + # identically so the A/B measures the kernel, not a flag asymmetry. + "-mllvm", "-enable-noalias-to-md-conversion=0", + "-mllvm", "--lsr-drop-solution=1", + "-mllvm", "-enable-post-misched=0", + "-mllvm", "-amdgpu-early-inline-all=true", + "-mllvm", "-amdgpu-function-calls=false", + "-fno-offload-uniform-block", + "-Wno-undefined-func-template", "-Wno-float-equal", + ] + cc = subprocess.run(["/opt/rocm/bin/hipcc", "-c", *common, str(SRC), "-o", str(obj)], + capture_output=True) + if cc.returncode != 0: + return None + ln = subprocess.run(["/opt/rocm/bin/hipcc", "-shared", "-fPIC", + f"--offload-arch={ARCH}", "--hip-link", + str(obj), str(STATIC), "-o", str(lib)], capture_output=True) + return lib if ln.returncode == 0 else None + + +def meas(so: Path, M: int, N: int, K: int) -> float | None: + """Median TFLOPS over REPEATS worker calls (each call does its own + warmup=50/repeat=100 internally). Median, not max, to match the sweep + methodology and stay robust to the occasional clock-warmup outlier.""" + if not so or not Path(so).exists(): + return None + payload = json.dumps({"so_path": str(so), "problem": {"M": M, "N": N, "K": K}, + "kernel_name": "x"}) + env = os.environ.copy() + env["HIP_VISIBLE_DEVICES"] = DEVICE + env["GEMM_PYPATH"] = PYPATH + env["LD_LIBRARY_PATH"] = "/opt/rocm/lib:" + env.get("LD_LIBRARY_PATH", "") + samples = [] + for _ in range(REPEATS): + p = subprocess.run([sys.executable, str(WORKER)], input=payload.encode(), + stdout=subprocess.PIPE, stderr=subprocess.DEVNULL, env=env) + for line in p.stdout.decode().splitlines(): + try: + d = json.loads(line) + except json.JSONDecodeError: + continue + if d.get("ok"): + samples.append(d["tflops"]) + return statistics.median(samples) if samples else None + + +def meas_all(so: Path) -> dict: + """Median TFLOPS per shape from REPEATS *batched* worker calls. + + One worker call measures ALL shapes (5x fewer python+numpy+CDLL startups + than per-shape meas()), which is the throughput lever for a full sweep on a + single GPU. Returns {shape_str: tflops|None}.""" + out = {f"{M}x{N}x{K}": None for (M, N, K) in SHAPES} + if not so or not Path(so).exists(): + return out + items = [{"so_path": str(so), "problem": {"M": M, "N": N, "K": K}, + "kernel_name": "x"} for (M, N, K) in SHAPES] + payload = json.dumps({"items": items, "verify": False}) + env = os.environ.copy() + env["HIP_VISIBLE_DEVICES"] = DEVICE + env["GEMM_PYPATH"] = PYPATH + env["LD_LIBRARY_PATH"] = "/opt/rocm/lib:" + env.get("LD_LIBRARY_PATH", "") + samples = {s: [] for s in out} + for _ in range(REPEATS): + p = subprocess.run([sys.executable, str(WORKER)], input=payload.encode(), + stdout=subprocess.PIPE, stderr=subprocess.DEVNULL, + env=env, timeout=900) + for line in p.stdout.decode().splitlines(): + try: + d = json.loads(line) + except json.JSONDecodeError: + continue + idx = d.get("idx") + if isinstance(idx, int) and 0 <= idx < len(SHAPES) and d.get("ok"): + M, N, K = SHAPES[idx] + samples[f"{M}x{N}x{K}"].append(d["tflops"]) + for s, xs in samples.items(): + if xs: + out[s] = statistics.median(xs) + return out + + +def pipeline_of(stem: str) -> str: + for p in ("compv3", "compv4", "mem"): + if f"_{p}_" in stem: + return p + return "other" + + +def main(): + ap = argparse.ArgumentParser(description=__doc__) + ap.add_argument("stems", nargs="*", help="kernel stems to A/B") + ap.add_argument("--stems-file", help="file with one stem per line") + ap.add_argument("--csv", help="write results to CSV (resume-aware)") + ap.add_argument("--build-only", action="store_true", + help="parallel-compile old-TE .so for all stems, then exit " + "(CPU pre-pass; GPU measurement reuses the cache)") + ap.add_argument("--jobs", type=int, default=min(os.cpu_count() or 8, 16), + help="parallel compile jobs for --build-only") + args = ap.parse_args() + + stems = list(args.stems) + if args.stems_file: + stems += [l.strip() for l in Path(args.stems_file).read_text().splitlines() + if l.strip()] + stems = stems or DEFAULT_STEMS + + # Parallel CPU pre-compile of every old-TE .so (no GPU touched). + if args.build_only: + from concurrent.futures import ProcessPoolExecutor, as_completed + ok = miss = fail = 0 + print(f"build-only: {len(stems)} stems, jobs={args.jobs}", flush=True) + with ProcessPoolExecutor(max_workers=args.jobs) as ex: + futs = {ex.submit(build_old_so, s): s for s in stems} + for i, fut in enumerate(as_completed(futs), 1): + try: + r = fut.result() + except Exception: + r = None + s = futs[fut] + if r is None: + # distinguish "no header" from "compile failed" + if (old_gen_dir(s) / f"gemm_universal_single_{s}.hpp").exists(): + fail += 1 + else: + miss += 1 + else: + ok += 1 + if i % 100 == 0: + print(f" [{i}/{len(stems)}] ok={ok} no_header={miss} fail={fail}", + flush=True) + print(f"build-only DONE: ok={ok} no_header={miss} fail={fail}", flush=True) + return + + # CSV sweep mode: same columns as the (now-corrected) sweep, resume-aware. + if args.csv: + fields = ["stem", "pipeline", "dtype", "layout", "shape", + "bridge_tflops", "old_tflops", "gap_pct", "oldte_built"] + out = Path(args.csv) + done = set() + if out.exists(): + with open(out) as f: + for row in csv.DictReader(f): + done.add((row["stem"], row["shape"])) + mode = "a" if done else "w" + print(f"stems={len(stems)} shapes={len(SHAPES)} resume={len(done)} -> {out}", + flush=True) + with open(out, mode, newline="") as fh: + w = csv.DictWriter(fh, fieldnames=fields) + if mode == "w": + w.writeheader() + for stem in stems: + todo = [(M, N, K) for (M, N, K) in SHAPES + if (stem, f"{M}x{N}x{K}") not in done] + if not todo: + continue + parts = stem.split("_") + dtype, layout = parts[0], parts[1] + old_so = build_old_so(stem) + br_so = BR_SO_DIR / f"libgemm_{stem}.so" + # Batched: one worker call per side covers all shapes. + bridge = meas_all(br_so) + old = meas_all(old_so) if old_so else {} + for (M, N, K) in todo: + shape = f"{M}x{N}x{K}" + b = bridge.get(shape) + o = old.get(shape) + gap = (b - o) / o * 100 if (b and o) else float("nan") + w.writerow(dict( + stem=stem, pipeline=pipeline_of(stem), dtype=dtype, + layout=layout, shape=shape, + bridge_tflops=f"{b:.4f}" if b is not None else "nan", + old_tflops=f"{o:.4f}" if o is not None else "nan", + gap_pct=f"{gap:.4f}" if gap == gap else "nan", + oldte_built=str(old_so is not None))) + fh.flush() + print(f" done {stem[:60]}", flush=True) + print(f"DONE -> {out}", flush=True) + return + + # Pretty-print mode. + print(f"{'shape':>14} {'bridge':>9} {'oldTE':>9} {'gap%':>7} kernel") + for stem in stems: + old_so = build_old_so(stem) + br_so = BR_SO_DIR / f"libgemm_{stem}.so" + if old_so is None: + print(f" [skip: no old-TE header] {stem}") + continue + for (M, N, K) in SHAPES: + b = meas(br_so, M, N, K) + o = meas(old_so, M, N, K) + gap = (b - o) / o * 100 if (b and o) else float("nan") + print(f"{f'{M}x{N}x{K}':>14} {b or float('nan'):9.2f} " + f"{o or float('nan'):9.2f} {gap:7.2f} {stem[:40]}") + + +if __name__ == "__main__": + main() diff --git a/projects/composablekernel/dispatcher/python/ctypes_utils.py b/projects/composablekernel/dispatcher/python/ctypes_utils.py index d719d1405e5d..8a919f4ffa02 100644 --- a/projects/composablekernel/dispatcher/python/ctypes_utils.py +++ b/projects/composablekernel/dispatcher/python/ctypes_utils.py @@ -1073,7 +1073,7 @@ def _generate_single_kernel_subprocess(args: dict) -> Tuple[bool, Optional[str], "--config", config_file, "--variants", - "standard", + args.get("variant", "standard"), ] res = subprocess.run(cmd, capture_output=True, text=True, timeout=300) @@ -1389,9 +1389,13 @@ class KernelConfig: gfx_arch: str = "gfx942" # GEMM variant (affects arch filter validation) - # "standard", "preshuffle", or "multi_d" + # "standard", "preshuffle", "multi_d", or "stream_k" variant: str = "standard" + # Stream-K reduction strategy ("atomic"/"linear"/"tree"); only used by the + # stream_k variant. Does not affect tile/arch validation. + reduction_strategy: str = "atomic" + @property def layout(self) -> str: """Get layout string (e.g., 'rcr' for row-col-row)""" diff --git a/projects/composablekernel/dispatcher/python/gemm_utils.py b/projects/composablekernel/dispatcher/python/gemm_utils.py new file mode 100644 index 000000000000..d05f4887030c --- /dev/null +++ b/projects/composablekernel/dispatcher/python/gemm_utils.py @@ -0,0 +1,1034 @@ +# Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +# SPDX-License-Identifier: MIT +""" +GEMM Tile Engine <-> Dispatcher bridge. + +This is the GEMM counterpart of ``grouped_conv_utils.py`` / ``fmha_utils.py``: +a single shared config dataclass (``GemmKernelConfig``) that Tile Engine imports +and hands back to the dispatcher. There is no translator between two +vocabularies -- both sides share the one object whose ``.name`` mirrors the +kernel identifier baked into the generated kernel header. + +Public surface (mirrors the grouped_conv bridge): + + GemmKernelConfig -- the shared contract dataclass + .name -- registry/runtime lookup key (byte-exact) + .to_codegen_json() -- feeds unified_gemm_codegen.py + GemmProblem -- a single (M, N, K) problem + setup_multiple_gemm_dispatchers -- codegen + hipcc -> .so paths (NO GPU) + GemmDispatcherLib -- thin ctypes ABI wrapper + GpuGemmRunner -- GPU memory + run + time (from a .so path) + expand_sweep -- TE JSON sweep config -> [GemmKernelConfig] + +The heavy lifting for codegen and compilation is reused from ``ctypes_utils`` +so there is a single source of truth for how a kernel header is produced and +how it is compiled into a ``.so``. +""" + +from __future__ import annotations + +import ctypes +import itertools +import json +import multiprocessing +from concurrent.futures import ProcessPoolExecutor, as_completed +from dataclasses import dataclass, field +from pathlib import Path +from typing import Any, Dict, List, Optional, Tuple + +import numpy as np + +# Reuse the proven codegen/compile leaf helpers from the dispatcher's own +# python layer. gemm_utils is a thin bridge on top of these. +import ctypes_utils as _cu + + +# ============================================================================ +# Layout / dtype helpers +# ============================================================================ + +_LAYOUT_CHAR = {"row": "r", "col": "c", "r": "r", "c": "c"} +_LAYOUT_WORD = {"r": "row", "c": "col"} + + +def _cap(flag: bool) -> str: + """Reproduce Python ``str(bool).capitalize()`` -> 'True' / 'False'.""" + return "True" if flag else "False" + + +# ============================================================================ +# The shared contract: GemmKernelConfig +# ============================================================================ + + +@dataclass +class GemmKernelConfig: + """The common config struct shared by Tile Engine and the Dispatcher. + + Naming convention (the "warp/wave trap" lives here, in ONE place): + * ``wave_m/n/k`` -- warps per block (C++ ``wave_shape``; TE "warp"). + * ``warp_tile_m/n/k`` -- MFMA instruction shape (C++ ``warp_tile_shape``; + TE "warp_tile"). + """ + + # --- Signature: what operation is computed ----------------------------- + dtype_a: str = "fp16" + dtype_b: str = "fp16" + dtype_c: str = "fp16" + dtype_acc: str = "fp32" + layout_a: str = "row" + layout_b: str = "col" + layout_c: str = "row" + + # --- Algorithm: how it is implemented ---------------------------------- + tile_m: int = 128 + tile_n: int = 128 + tile_k: int = 32 + wave_m: int = 2 + wave_n: int = 2 + wave_k: int = 1 + warp_tile_m: int = 32 + warp_tile_n: int = 32 + warp_tile_k: int = 16 + + pipeline: str = "compv4" + scheduler: str = "intrawave" + epilogue: str = "cshuffle" + + pad_m: bool = True + pad_n: bool = True + pad_k: bool = True + persistent: bool = False + + gfx_arch: str = "gfx942" + variant: str = "standard" + # Stream-K reduction strategy: "atomic" (default), "linear", or "tree". + # Only meaningful when variant == "stream_k". + reduction_strategy: str = "atomic" + + # ------------------------------------------------------------------ # + # Derived string fragments + # ------------------------------------------------------------------ # + @property + def layout(self) -> str: + """3-char layout string, e.g. 'rcr'.""" + return ( + _LAYOUT_CHAR[self.layout_a] + + _LAYOUT_CHAR[self.layout_b] + + _LAYOUT_CHAR[self.layout_c] + ) + + @property + def tile_str(self) -> str: + return f"{self.tile_m}x{self.tile_n}x{self.tile_k}" + + @property + def wave_str(self) -> str: + return f"{self.wave_m}x{self.wave_n}x{self.wave_k}" + + @property + def warp_tile_str(self) -> str: + return f"{self.warp_tile_m}x{self.warp_tile_n}x{self.warp_tile_k}" + + @property + def name(self) -> str: + """Registry / runtime lookup key. + + Reproduces, byte-for-byte, the ``KERNEL_NAME`` that + ``unified_gemm_codegen.py::KernelNaming.generate`` bakes into the + generated kernel header (and that the .so reports via + ``dispatcher_get_kernel_name``). This is the single thread tying + config -> codegen -> runtime together. + """ + name = ( + f"gemm_{self.dtype_a}_{self.layout}" + f"_{self.pipeline}_{self.epilogue}_{self.scheduler}" + f"_{_cap(self.pad_m)}_{_cap(self.pad_n)}_{_cap(self.pad_k)}" + f"_{_cap(self.persistent)}" + f"_{self.tile_str}_{self.wave_str}_{self.warp_tile_str}" + ) + if self.variant == "preshuffle": + name += "_preshuffle" + elif self.variant == "stream_k": + name += "_streamk" + # Atomic keeps the bare "_streamk" suffix (original parity); linear + # and tree are disambiguated, matching KernelNaming.generate. + if self.reduction_strategy != "atomic": + name += f"_{self.reduction_strategy}" + return name + + # ------------------------------------------------------------------ # + # Serialization + # ------------------------------------------------------------------ # + def to_codegen_json(self) -> Dict[str, Any]: + """Single-config JSON consumed by unified_gemm_codegen.py. + + Note the warp/wave mapping: the codegen calls the warps-per-block + triple ``warp_*`` and the MFMA triple ``warp_tile_*``. We translate + from dispatcher semantics here so the mapping cannot drift. + """ + cfg: Dict[str, Any] = { + "tile_config": { + "tile_m": [self.tile_m], + "tile_n": [self.tile_n], + "tile_k": [self.tile_k], + # dispatcher wave_* -> codegen warp_* (warps per block) + "warp_m": [self.wave_m], + "warp_n": [self.wave_n], + "warp_k": [self.wave_k], + # dispatcher warp_tile_* -> codegen warp_tile_* (MFMA shape) + "warp_tile_m": [self.warp_tile_m], + "warp_tile_n": [self.warp_tile_n], + "warp_tile_k": [self.warp_tile_k], + }, + "trait_config": { + "pipeline": [self.pipeline], + "epilogue": [self.epilogue], + "scheduler": [self.scheduler], + "pad_m": [self.pad_m], + "pad_n": [self.pad_n], + "pad_k": [self.pad_k], + "persistent": [self.persistent], + }, + } + # Pin the single reduction strategy so stream-K codegen emits exactly this + # kernel (the generator otherwise expands all strategies in its default). + if self.variant == "stream_k": + cfg["streamk_config"] = {"reduction_strategy": [self.reduction_strategy]} + return cfg + + def to_dict(self) -> Dict[str, Any]: + return { + "dtype_a": self.dtype_a, + "dtype_b": self.dtype_b, + "dtype_c": self.dtype_c, + "dtype_acc": self.dtype_acc, + "layout": self.layout, + "tile": [self.tile_m, self.tile_n, self.tile_k], + "wave": [self.wave_m, self.wave_n, self.wave_k], + "warp_tile": [self.warp_tile_m, self.warp_tile_n, self.warp_tile_k], + "pipeline": self.pipeline, + "scheduler": self.scheduler, + "epilogue": self.epilogue, + "pad": [self.pad_m, self.pad_n, self.pad_k], + "persistent": self.persistent, + "gfx_arch": self.gfx_arch, + "variant": self.variant, + "reduction_strategy": self.reduction_strategy, + "name": self.name, + } + + def to_ctypes_config(self) -> "_cu.KernelConfig": + """Convert to the ctypes_utils.KernelConfig used by the codegen/validate + helpers. ctypes_utils renames the MFMA triple ``warp_*`` (no _tile).""" + return _cu.KernelConfig( + dtype_a=self.dtype_a, + dtype_b=self.dtype_b, + dtype_c=self.dtype_c, + dtype_acc=self.dtype_acc, + layout_a=_LAYOUT_WORD[_LAYOUT_CHAR[self.layout_a]], + layout_b=_LAYOUT_WORD[_LAYOUT_CHAR[self.layout_b]], + layout_c=_LAYOUT_WORD[_LAYOUT_CHAR[self.layout_c]], + tile_m=self.tile_m, + tile_n=self.tile_n, + tile_k=self.tile_k, + wave_m=self.wave_m, + wave_n=self.wave_n, + wave_k=self.wave_k, + warp_m=self.warp_tile_m, + warp_n=self.warp_tile_n, + warp_k=self.warp_tile_k, + pipeline=self.pipeline, + scheduler=self.scheduler, + epilogue=self.epilogue, + pad_m=self.pad_m, + pad_n=self.pad_n, + pad_k=self.pad_k, + gfx_arch=self.gfx_arch, + variant=self.variant, + reduction_strategy=self.reduction_strategy, + ) + + +# ============================================================================ +# Problem +# ============================================================================ + + +@dataclass +class GemmProblem: + """A single GEMM problem: C[MxN] = A[MxK] @ B[KxN].""" + + M: int + N: int + K: int + + @property + def flops(self) -> float: + return 2.0 * self.M * self.N * self.K + + def to_dict(self) -> Dict[str, int]: + return {"M": self.M, "N": self.N, "K": self.K} + + @classmethod + def from_dict(cls, d: Dict[str, int]) -> "GemmProblem": + return cls(M=int(d["M"]), N=int(d["N"]), K=int(d["K"])) + + +@dataclass +class GemmResult: + output: np.ndarray + time_ms: float + status: int + tflops: float + kernel_name: str + + @property + def success(self) -> bool: + return self.status == 0 + + +# ============================================================================ +# ctypes ABI wrapper +# ============================================================================ + + +class GemmDispatcherLib: + """Thin ctypes wrapper around a compiled GEMM dispatcher .so. + + Supports both the legacy single-kernel ABI (``dispatcher_get_kernel_name``) + and the multi-kernel ABI (``dispatcher_get_kernel_name_at(index, buf, n)``) + so one .so can report a whole batch and be selected by name. + """ + + def __init__(self, so_path: Path): + self._path = Path(so_path) + self._lib = ctypes.CDLL(str(self._path)) + self._has_indexed = hasattr(self._lib, "dispatcher_get_kernel_name_at") + self._setup_functions() + + def _setup_functions(self) -> None: + lib = self._lib + + lib.dispatcher_initialize.argtypes = [] + lib.dispatcher_initialize.restype = ctypes.c_int + + lib.dispatcher_get_kernel_count.argtypes = [] + lib.dispatcher_get_kernel_count.restype = ctypes.c_int + + lib.dispatcher_get_kernel_name.argtypes = [] + lib.dispatcher_get_kernel_name.restype = ctypes.c_char_p + + if self._has_indexed: + lib.dispatcher_get_kernel_name_at.argtypes = [ + ctypes.c_int, + ctypes.c_char_p, + ctypes.c_int, + ] + lib.dispatcher_get_kernel_name_at.restype = ctypes.c_int + + lib.dispatcher_run_gemm.argtypes = [ + ctypes.c_void_p, # A (host) + ctypes.c_void_p, # B (host) + ctypes.c_void_p, # C (host) + ctypes.c_int64, # M + ctypes.c_int64, # N + ctypes.c_int64, # K + ctypes.POINTER(ctypes.c_float), # time_ms + ] + lib.dispatcher_run_gemm.restype = ctypes.c_int + + lib.dispatcher_cleanup.argtypes = [] + lib.dispatcher_cleanup.restype = None + + @property + def path(self) -> Path: + return self._path + + def initialize(self) -> bool: + return self._lib.dispatcher_initialize() == 0 + + def get_kernel_count(self) -> int: + return int(self._lib.dispatcher_get_kernel_count()) + + @property + def kernel_names(self) -> List[str]: + """List every kernel the .so exposes, by index when available.""" + if self._has_indexed: + names: List[str] = [] + count = self.get_kernel_count() + buf = ctypes.create_string_buffer(256) + for i in range(count): + if self._lib.dispatcher_get_kernel_name_at(i, buf, 256) == 0: + names.append(buf.value.decode("utf-8")) + if names: + return names + # Legacy single-kernel fallback. + raw = self._lib.dispatcher_get_kernel_name() + return [raw.decode("utf-8")] if raw else [] + + def run( + self, A: np.ndarray, B: np.ndarray, C: np.ndarray, M: int, N: int, K: int + ) -> Tuple[int, float]: + time_ms = ctypes.c_float(0.0) + status = self._lib.dispatcher_run_gemm( + A.ctypes.data_as(ctypes.c_void_p), + B.ctypes.data_as(ctypes.c_void_p), + C.ctypes.data_as(ctypes.c_void_p), + M, + N, + K, + ctypes.byref(time_ms), + ) + return status, time_ms.value + + def cleanup(self) -> None: + self._lib.dispatcher_cleanup() + + +# ============================================================================ +# GPU runner (constructed from a .so path; loaded only inside a worker) +# ============================================================================ + + +# --------------------------------------------------------------------------- +# fp8 (E4M3) / bf8 (E5M2) -- FNUZ ("NANOO") encoding used by gfx942/MI300. +# +# numpy has no native 8-bit float, and the C ABI only cares about the 1-byte +# memory layout (sizeof(fp8_t) == sizeof(bf8_t) == 1). We carry the value as a +# uint8 bit pattern. As with bf16, the DECODE is the load-bearing half: it must +# return the exact value the device's fp8_t/bf8_t represents for a byte, so the +# NumPy reference multiplies bit-for-bit what the GPU multiplies. The ENCODE only +# needs to land on the nearest representable byte. +# +# FNUZ format (gfx942): bias = 2^(exp_bits-1); the all-1s exponent is a normal +# number (no Inf), the sole NaN is the sign=1/exp=0/mant=0 byte (0x80), and there +# is no negative zero. gfx950/MI350 uses the OCP fp8 format instead; this codec +# targets the gfx942 default and the OCP path needs separate handling (the runner +# raises a clear error for fp8/bf8 on a non-gfx942 arch). +# --------------------------------------------------------------------------- + + +def _fnuz_decode_table(exp_bits: int, mant_bits: int) -> np.ndarray: + """Build the 256-entry byte -> fp32 value table for an 8-bit FNUZ float.""" + bias = (1 << (exp_bits - 1)) + mant_max = 1 << mant_bits + sign_shift = exp_bits + mant_bits + exp_mask = (1 << exp_bits) - 1 + table = np.zeros(256, dtype=np.float32) + for b in range(256): + sign = (b >> sign_shift) & 1 + exp = (b >> mant_bits) & exp_mask + mant = b & (mant_max - 1) + if exp == 0 and mant == 0: + # +0 (0x00); the negative-zero slot (0x80) is the lone NaN. + table[b] = np.float32(np.nan) if sign else np.float32(0.0) + continue + if exp == 0: + val = (mant / mant_max) * (2.0 ** (1 - bias)) # subnormal + else: + val = (1.0 + mant / mant_max) * (2.0 ** (exp - bias)) # normal + table[b] = np.float32(-val if sign else val) + return table + + +def _fnuz_encode(x: np.ndarray, exp_bits: int, mant_bits: int) -> np.ndarray: + """Encode fp32 -> nearest 8-bit FNUZ float, returned as a uint8 bit pattern. + + PRESERVES the input's memory order (C or F) so a column-major operand stays + column-major after encoding. + """ + table = _fnuz_decode_table(exp_bits, mant_bits) + sign_byte = np.uint8(1 << (exp_bits + mant_bits)) # 0x80 + + # Positive half (bytes 0..127) holds every non-negative magnitude, sorted. + # Compare in float64: for very large inputs the gap between the two top + # magnitudes is below fp32 resolution, which would tie and mis-saturate. + pos_mag = table[: int(sign_byte)].astype(np.float64) + order = np.argsort(pos_mag) + sorted_mag = pos_mag[order] + sorted_byte = order.astype(np.uint8) + + xf = np.asarray(x, dtype=np.float32) + if not (xf.flags["C_CONTIGUOUS"] or xf.flags["F_CONTIGUOUS"]): + xf = np.ascontiguousarray(xf) + ax = np.abs(xf).astype(np.float64) + # Both neighbours come from the raw insertion point: raw==size saturates to + # the top magnitude (lo==hi), raw==0 pins to zero, otherwise compare the two. + raw = np.searchsorted(sorted_mag, ax) + hi = np.clip(raw, 0, sorted_mag.size - 1) + lo = np.clip(raw - 1, 0, sorted_mag.size - 1) + pick_lo = np.abs(sorted_mag[lo] - ax) <= np.abs(sorted_mag[hi] - ax) + chosen = np.where(pick_lo, lo, hi) + out = sorted_byte[chosen] + + # Apply sign, but never the 0x80 (-0 == NaN) slot: zeros stay +0. + is_zero = sorted_mag[chosen] == 0 + out = np.where((xf < 0) & ~is_zero, out | sign_byte, out) + out = np.where(np.isnan(xf), sign_byte, out) # NaN inputs -> NaN byte + # np.where collapses memory order; restore the operand's contiguity. + out = out.astype(np.uint8) + return np.asfortranarray(out) if xf.flags["F_CONTIGUOUS"] else np.ascontiguousarray(out) + + +def _fp32_to_fp8_u8(x: np.ndarray) -> np.ndarray: + """Encode fp32 -> fp8 E4M3 (FNUZ) bit pattern in a uint8 array.""" + return _fnuz_encode(x, exp_bits=4, mant_bits=3) + + +def _fp8_u8_to_fp32(u8: np.ndarray) -> np.ndarray: + """Decode an fp8 E4M3 (FNUZ) bit pattern back to fp32.""" + return _fnuz_decode_table(4, 3)[u8.astype(np.intp)] + + +def _fp32_to_bf8_u8(x: np.ndarray) -> np.ndarray: + """Encode fp32 -> bf8 E5M2 (FNUZ) bit pattern in a uint8 array.""" + return _fnuz_encode(x, exp_bits=5, mant_bits=2) + + +def _bf8_u8_to_fp32(u8: np.ndarray) -> np.ndarray: + """Decode a bf8 E5M2 (FNUZ) bit pattern back to fp32.""" + return _fnuz_decode_table(5, 2)[u8.astype(np.intp)] + + +# Output (C) element dtype for an A/B element dtype, mirroring the codegen's +# CommonTypeMappings.get_output_dtype: fp8/bf8 accumulate into fp16, int8 into +# int32, everything else stores in its own dtype. +_OUTPUT_DTYPE = {"fp8": "fp16", "bf8": "fp16", "int8": "int32"} + + +def _output_dtype(dtype: str) -> str: + return _OUTPUT_DTYPE.get(dtype, dtype) + + +# numpy carrier dtype for each output (C) element type. fp8/bf8 -> fp16 store, +# int8 -> int32 accumulate, bf16 carried as raw uint16 bits. +_C_NP = {"fp16": np.float16, "bf16": np.uint16, "int32": np.int32} + + +def _detect_gpu_arch() -> Optional[str]: + """Best-effort detection of the active GPU's gcnArchName (e.g. 'gfx942'). + + Parses ``rocminfo`` for the first ``gfx*`` Name line. Returns ``None`` if it + cannot be determined; callers treat that as "cannot verify arch" rather than + a hard failure for non-fp8 dtypes. + """ + import re + import subprocess + + try: + out = subprocess.run( + ["rocminfo"], capture_output=True, text=True, timeout=30 + ).stdout + except Exception: + return None + m = re.search(r"^\s*Name:\s*(gfx[0-9a-fA-F]+)\s*$", out, re.MULTILINE) + return m.group(1) if m else None + + +class GpuGemmRunner: + """High-level runner: construct from a .so path, call run(A, B, problem). + + The GEMM ctypes ABI takes HOST pointers and manages GPU memory internally + (hipMalloc/hipMemcpy/hipFree), so this runner stays simple -- it hands + numpy arrays straight to the .so. + """ + + def __init__(self, lib_path: Path): + self.lib = GemmDispatcherLib(lib_path) + if not self.lib.initialize(): + raise RuntimeError(f"Failed to initialize dispatcher .so: {lib_path}") + names = self.lib.kernel_names + self._kernel_name = names[0] if names else "unknown" + # dtype and layout are encoded in the kernel name: gemm___... + # layout is the 3-char A/B/C major code (e.g. 'rcr'). Nothing layout- or + # dtype-specific is hardcoded -- both are read off the compiled kernel. + parts = self._kernel_name.split("_") + self._dtype = parts[1] if len(parts) > 1 else "fp16" + lay = parts[2] if len(parts) > 2 and len(parts[2]) == 3 else "rcr" + self._layout = lay if set(lay) <= {"r", "c"} else "rcr" + + @property + def kernel_name(self) -> str: + return self._kernel_name + + @staticmethod + def _bf16_encode(x: np.ndarray) -> np.ndarray: + """float -> bfloat16 bits (uint16), round-to-nearest-even, PRESERVING the + input's memory order (C or F) so column-major operands stay column-major. + ENCODE need only be nearest-representable; DECODE must be bit-exact to + device bf16_t so the numpy reference multiplies what the GPU does.""" + f = np.asarray(x, dtype=np.float32) + if not (f.flags["C_CONTIGUOUS"] or f.flags["F_CONTIGUOUS"]): + f = np.ascontiguousarray(f) + u = f.view(np.uint32) + rounded = (u + 0x7FFF + ((u >> 16) & 1)) >> 16 + return rounded.astype(np.uint16) + + @staticmethod + def _bf16_decode(u16: np.ndarray) -> np.ndarray: + return (u16.astype(np.uint32) << 16).view(np.float32) + + # fp8/bf8 codecs are bit-exact to the device fp8_t/bf8_t (FNUZ on gfx942); + # re-exposed as static methods so references (smoke test, run_one) can build + # decode(encode(x)) quantized inputs without reaching into module functions. + _fp8_encode = staticmethod(_fp32_to_fp8_u8) + _fp8_decode = staticmethod(_fp8_u8_to_fp32) + _bf8_encode = staticmethod(_fp32_to_bf8_u8) + _bf8_decode = staticmethod(_bf8_u8_to_fp32) + + def _check_arch_for_dtype(self) -> None: + """fp8/bf8 use the gfx942 FNUZ format. gfx950/MI350 uses OCP fp8, a + different bit layout, so refuse rather than silently mis-decode.""" + if self._dtype not in ("fp8", "bf8"): + return + arch = _detect_gpu_arch() + if arch is not None and arch != "gfx942": + raise RuntimeError( + f"fp8/bf8 bridge codec is FNUZ (gfx942/MI300) only; detected " + f"GPU arch {arch!r}. gfx950/MI350 uses OCP fp8 (different bit " + f"layout) -- an OCP codec is required for that arch." + ) + + def _to_buf(self, X: np.ndarray, major: str) -> np.ndarray: + """Lay out an operand in the order its layout implies: RowMajor -> + C-contiguous, ColumnMajor -> F-contiguous. The .so reads a flat buffer + with the matching stride, so the raw byte order is what matters. The + encode helpers (bf16/fp8/bf8) preserve that contiguity; int8/fp16 keep + the requested order via astype(order='K').""" + arr = np.ascontiguousarray(X) if major == "r" else np.asfortranarray(X) + if self._dtype == "bf16": + return self._bf16_encode(arr) + if self._dtype == "fp8": + return _fp32_to_fp8_u8(arr) + if self._dtype == "bf8": + return _fp32_to_bf8_u8(arr) + if self._dtype == "int8": + return arr.astype(np.int8, order="K") + return arr.astype(np.float16, order="K") + + def run( + self, A: np.ndarray, B: np.ndarray, problem: GemmProblem + ) -> GemmResult: + M, N, K = problem.M, problem.N, problem.K + self._check_arch_for_dtype() + + # Arrange A (MxK), B (KxN), C (MxN) per the kernel's actual layout. The + # ctypes ABI is void*+sizeof, so each dtype just needs the right bit + # pattern: bf16 -> uint16, fp8/bf8 -> uint8, int8 -> int8, fp16 -> fp16. + la, lb, lc = self._layout[0], self._layout[1], self._layout[2] + A_h = self._to_buf(A, la) + B_h = self._to_buf(B, lb) + + # The C buffer's element size must equal sizeof(CDataType): fp8/bf8 + # accumulate into fp16, int8 into int32, otherwise the input dtype (bf16 + # carried as raw uint16 bits). + out_dtype = _output_dtype(self._dtype) + cdt = _C_NP.get(out_dtype, np.float16) + C_h = np.zeros((M, N), dtype=cdt, order=("C" if lc == "r" else "F")) + + status, time_ms = self.lib.run(A_h, B_h, C_h, M, N, K) + + # Decode the output to a comparable numeric array. fp16/fp8/bf8 store fp16 + # (already comparable); int8 stores int32; only bf16 needs bit-decode. + out = self._bf16_decode(C_h) if out_dtype == "bf16" else C_h + tflops = (problem.flops / (time_ms * 1e-3)) / 1e12 if time_ms > 0 else 0.0 + return GemmResult( + output=out, + time_ms=time_ms, + status=status, + tflops=tflops, + kernel_name=self._kernel_name, + ) + + +# ============================================================================ +# Build API: codegen + hipcc -> .so paths (no GPU) +# ============================================================================ + + +def _ctypes_source_name(variant: str) -> str: + """Select the ctypes bridge .cpp for a variant. + + Variants whose launch ABI differs from the single-problem + ``dispatcher_run_gemm`` path need their own lib. Stream-K keeps the same + C ABI (single A/B/C, M/N/K) but its lib builds a ``StreamKHostArgs`` and + calls ``SelectedKernel::launch(args, stream)`` directly instead of routing + through the registry, so it gets a dedicated source. + """ + if variant == "stream_k": + return "streamk_gemm_ctypes_lib.cpp" + return "gemm_ctypes_lib.cpp" + + +def _build_compile_jobs( + config: GemmKernelConfig, header: Path +) -> Tuple[Dict[str, Any], Path]: + """Replicate the (validated) compile+link commands from ctypes_utils.""" + root = _cu.get_dispatcher_root() + ck_root = root.parent + build_dir = _cu.get_build_dir() + output_dir = _cu.get_generated_kernels_dir() + ctypes_source = ( + root / "bindings" / "ctypes" / _ctypes_source_name(config.variant) + ) + static_lib = build_dir / "libck_tile_dispatcher.a" + + lib_path = build_dir / "examples" / f"lib{config.name}.so" + obj_file = lib_path.with_suffix(".o") + # The Stream-K path skips the cmake build that would normally create this + # directory, so ensure it exists before hipcc writes the object/.so here. + lib_path.parent.mkdir(parents=True, exist_ok=True) + + # Per-variant compile flags. Stream-K must match Tile Engine's gemm_streamk + # build EXACTLY for a fair A/B. Ground truth is a TE streamk build's + # compile_commands.json (the -mllvm flags come from the composablekernel + # project-root add_compile_options, applied globally to the TE benchmark, NOT + # the per-target options): -std=c++20 -fno-offload-uniform-block + # -mllvm --lsr-drop-solution=1 -mllvm -enable-post-misched=0 + # -mllvm -amdgpu-early-inline-all=true -mllvm -amdgpu-function-calls=false + # --offload-compress. Note -enable-post-misched=0 is applied UNCONDITIONALLY by + # TE for streamk (not persistent-gated like the gemm_universal bridge). TE + # streamk does NOT use -enable-noalias-to-md-conversion=0. The non-streamk path + # keeps its existing flags unchanged (that is a #8479 concern). + is_streamk = getattr(config, "variant", "") == "stream_k" + variant_flags = ( + [ + "-std=c++20", + "-fno-offload-uniform-block", + "-mllvm", "--lsr-drop-solution=1", + "-mllvm", "-enable-post-misched=0", + "-mllvm", "-amdgpu-early-inline-all=true", + "-mllvm", "-amdgpu-function-calls=false", + "--offload-compress", + ] + if is_streamk + else ["-mllvm", "-enable-noalias-to-md-conversion=0"] + ) + compile_cmd = [ + "/opt/rocm/bin/hipcc", + "-c", + "-fPIC", + "-O3", + f"-I{root / 'include'}", + f"-I{ck_root / 'include'}", + f"-I{ck_root}", + f"-I{str(output_dir)}", + "-DCK_TILE_SINGLE_KERNEL_INCLUDE", + f"-include{header}", + "-D__HIP_PLATFORM_AMD__", + f"--offload-arch={config.gfx_arch}", + f'-DGFX_ARCH="{config.gfx_arch}"', + *variant_flags, + "-Wno-undefined-func-template", + "-Wno-float-equal", + str(ctypes_source), + "-o", + str(obj_file), + ] + # The Stream-K ctypes lib launches the force-included kernel directly and does + # NOT reference any registry/dispatcher symbols, so its .so does not need the + # dispatcher static lib. Skipping it removes the libck_tile_dispatcher.a build + # dependency for the Stream-K bridge. The regular path still links it. + link_cmd = [ + "/opt/rocm/bin/hipcc", + "-shared", + "-fPIC", + f"--offload-arch={config.gfx_arch}", + "--hip-link", + str(obj_file), + *([] if is_streamk else [str(static_lib)]), + "-o", + str(lib_path), + ] + job = {"compile_cmd": compile_cmd, "link_cmd": link_cmd, "lib_path": str(lib_path)} + return job, lib_path + + +def setup_multiple_gemm_dispatchers( + configs: List[GemmKernelConfig], + verbose: bool = True, + max_workers: Optional[int] = None, +) -> List[Optional[Path]]: + """Codegen + compile each config into its own .so. Returns .so paths. + + This is the build half of the bridge. It touches NO GPU -- pure CPU + codegen + hipcc, run massively in parallel -- and returns only ``Path`` + objects (``None`` for configs that failed to generate/compile), aligned to + the input order. Benchmarking happens later, in an isolated worker. + """ + import sys + + n = len(configs) + results: List[Optional[Path]] = [None] * n + if n == 0: + return results + + max_workers = max_workers or min(multiprocessing.cpu_count(), 8) + + # Dedupe identical configs by name; compile once, share the path. + first_index: Dict[str, int] = {} + unique: List[int] = [] + for i, c in enumerate(configs): + key = c.name + if key not in first_index: + first_index[key] = i + unique.append(i) + + codegen_script = _cu.get_codegen_path() + output_dir = _cu.get_generated_kernels_dir() + static_lib = _cu.get_build_dir() / "libck_tile_dispatcher.a" + # All configs in a sweep share one variant; route to the matching bridge lib. + ctypes_source = ( + _cu.get_dispatcher_root() + / "bindings" + / "ctypes" + / _ctypes_source_name(configs[0].variant) + ) + # Stream-K .so links only the force-included kernel (no registry/dispatcher + # symbols), so it does not need the dispatcher static lib; only the regular + # path requires it. + streamk_build = configs[0].variant == "stream_k" + if (not streamk_build and not static_lib.exists()) or not ctypes_source.exists(): + raise FileNotFoundError( + "Missing static lib or ctypes source required for compilation:\n" + f" {static_lib}\n {ctypes_source}\n" + "Build the dispatcher first (cmake + make)." + ) + + # -- Step 1: parallel codegen (one header per unique config) ----------- + codegen_args = [] + for i in unique: + c = configs[i] + codegen_args.append( + { + "index": i, + "python": sys.executable, + "codegen_script": str(codegen_script), + "output_dir": str(output_dir), + "dtype": c.dtype_a, + "layout": c.layout, + "gpu_target": c.gfx_arch, + "tile_config_json": c.to_codegen_json(), + "hpp_glob_pattern": f"{c.name}.hpp", + # Honor the config's variant so non-standard kernels are codegen'd + # as themselves; the kernel name (and thus hpp_glob_pattern) already + # carries the variant suffix, so a missing/standard value here would + # produce a header whose name never matches the requested pattern. + "variant": c.variant, + } + ) + + if verbose: + print( + f"[gemm-bridge] codegen: {len(codegen_args)} headers " + f"(workers={max_workers})..." + ) + + headers: Dict[int, Path] = {} + with ProcessPoolExecutor(max_workers=max_workers) as ex: + futs = { + ex.submit(_cu._generate_single_kernel_subprocess, a): a["index"] + for a in codegen_args + } + for fut in as_completed(futs): + i = futs[fut] + ok, hdr, err = fut.result() + if ok and hdr: + headers[i] = Path(hdr) + if verbose: + print(f" OK codegen [{i}] {configs[i].name}") + elif verbose: + print(f" FAIL codegen [{i}] {configs[i].name}: {err}") + + # -- Step 2: parallel compile + link ----------------------------------- + compile_jobs = [] + job_index: List[int] = [] + for i in unique: + hdr = headers.get(i) + if hdr is None: + continue + job, _ = _build_compile_jobs(configs[i], hdr) + compile_jobs.append(job) + job_index.append(i) + + if verbose and compile_jobs: + print( + f"[gemm-bridge] compile: {len(compile_jobs)} .so " + f"(workers={max_workers})..." + ) + + with ProcessPoolExecutor(max_workers=max_workers) as ex: + futs = { + ex.submit(_cu._run_hipcc_subprocess, job): job_index[j] + for j, job in enumerate(compile_jobs) + } + for fut in as_completed(futs): + i = futs[fut] + ok, lp, err = fut.result() + if ok and lp: + results[i] = Path(lp) + if verbose: + print(f" OK compile [{i}] {Path(lp).name}") + elif verbose: + print(f" FAIL compile [{i}] {configs[i].name}: {err}") + + # -- Fan the deduped result back out to every input index -------------- + for i, c in enumerate(configs): + if results[i] is None: + results[i] = results[first_index[c.name]] + + if verbose: + ok_count = sum(1 for r in results if r is not None) + print(f"[gemm-bridge] setup complete: {ok_count}/{n} configs -> .so") + + return results + + +# ============================================================================ +# TE sweep config expansion +# ============================================================================ + + +def _expand_range(entry: Dict[str, Any]) -> List[int]: + """Expand a tile_config entry: either {min,max,step} or {values:[...]}.""" + if "values" in entry: + return list(entry["values"]) + lo = int(entry["min"]) + hi = int(entry["max"]) + step = int(entry.get("step", 1)) + return list(range(lo, hi + 1, step)) + + +def _expand_values(entry: Optional[Dict[str, Any]], default: List[Any]) -> List[Any]: + if entry is None: + return list(default) + return list(entry.get("values", default)) + + +def expand_sweep( + config_path: str, + arch: str, + dtype: str = "fp16", + layout: str = "rcr", + variant: str = "standard", +) -> List[GemmKernelConfig]: + """Expand a Tile Engine GEMM JSON sweep config into GemmKernelConfig list. + + The TE config uses ``tile_config`` (ranges/value-lists for tile, warp and + warp_tile triples) and ``trait_config`` (value-lists for pipeline, + scheduler, epilogue, pad_*, persistent). Every valid combination becomes + one GemmKernelConfig. Invalid combinations are dropped via the dispatcher's + own validator, and duplicates (by .name) are collapsed. + + For Phase 1 the signature is fixed to fp16 / rcr. + """ + with open(config_path) as f: + cfg = json.load(f) + + tc = cfg.get("tile_config", {}) + tr = cfg.get("trait_config", {}) + + tile_ms = _expand_range(tc["tile_m"]) + tile_ns = _expand_range(tc["tile_n"]) + tile_ks = _expand_range(tc["tile_k"]) + wave_ms = _expand_range(tc["warp_m"]) # TE "warp" == wave count + wave_ns = _expand_range(tc["warp_n"]) + wave_ks = _expand_range(tc["warp_k"]) + wt_ms = _expand_range(tc["warp_tile_m"]) + wt_ns = _expand_range(tc["warp_tile_n"]) + wt_ks = _expand_range(tc["warp_tile_k"]) + + pipelines = _expand_values(tr.get("pipeline"), ["compv3"]) + schedulers = _expand_values(tr.get("scheduler"), ["intrawave"]) + epilogues = _expand_values(tr.get("epilogue"), ["cshuffle"]) + pad_ms = _expand_values(tr.get("pad_m"), [False]) + pad_ns = _expand_values(tr.get("pad_n"), [False]) + pad_ks = _expand_values(tr.get("pad_k"), [False]) + persistents = _expand_values(tr.get("persistent"), [False]) + + # Stream-K only: sweep reduction strategies (atomic/linear/tree). Other + # variants keep a single dummy value so the product is unaffected. + if variant == "stream_k": + sk = cfg.get("streamk_config", {}) + reductions = _expand_values(sk.get("reduction_strategy"), ["atomic"]) + else: + reductions = ["atomic"] + + la, lb, lc = layout[0], layout[1], layout[2] + + configs: List[GemmKernelConfig] = [] + seen: set = set() + for ( + tm, + tn, + tk, + wm, + wn, + wk, + wtm, + wtn, + wtk, + pipe, + sched, + epi, + pm, + pn, + pk, + persist, + redux, + ) in itertools.product( + tile_ms, + tile_ns, + tile_ks, + wave_ms, + wave_ns, + wave_ks, + wt_ms, + wt_ns, + wt_ks, + pipelines, + schedulers, + epilogues, + pad_ms, + pad_ns, + pad_ks, + persistents, + reductions, + ): + c = GemmKernelConfig( + dtype_a=dtype, + dtype_b=dtype, + dtype_c=dtype, + layout_a=_LAYOUT_WORD[la], + layout_b=_LAYOUT_WORD[lb], + layout_c=_LAYOUT_WORD[lc], + tile_m=tm, + tile_n=tn, + tile_k=tk, + wave_m=wm, + wave_n=wn, + wave_k=wk, + warp_tile_m=wtm, + warp_tile_n=wtn, + warp_tile_k=wtk, + pipeline=pipe, + scheduler=sched, + epilogue=epi, + pad_m=bool(pm), + pad_n=bool(pn), + pad_k=bool(pk), + persistent=bool(persist), + gfx_arch=arch, + variant=variant, + reduction_strategy=redux, + ) + if c.name in seen: + continue + val = _cu.validate_kernel_config(c.to_ctypes_config()) + if not val.is_valid: + continue + seen.add(c.name) + configs.append(c) + + return configs diff --git a/projects/composablekernel/dispatcher/src/dispatcher.cpp b/projects/composablekernel/dispatcher/src/dispatcher.cpp index 133485b2487c..10ea4efc8953 100644 --- a/projects/composablekernel/dispatcher/src/dispatcher.cpp +++ b/projects/composablekernel/dispatcher/src/dispatcher.cpp @@ -3,6 +3,7 @@ #include "ck_tile/dispatcher/dispatcher.hpp" #include "ck_tile/dispatcher/dispatcher_error.hpp" +#include #include #include @@ -17,6 +18,39 @@ Dispatcher::Dispatcher(Registry* registry, const std::string& gfx_arch) { } +Dispatcher::~Dispatcher() +{ + if(workspace_) + { + (void)hipFree(workspace_); + workspace_ = nullptr; + workspace_bytes_ = 0; + } +} + +void Dispatcher::ensure_workspace(std::size_t bytes) const +{ + if(bytes <= workspace_bytes_) + { + return; // current buffer is big enough (also covers bytes == 0) + } + + if(workspace_) + { + (void)hipFree(workspace_); + workspace_ = nullptr; + workspace_bytes_ = 0; + } + + if(hipMalloc(&workspace_, bytes) != hipSuccess) + { + workspace_ = nullptr; + workspace_bytes_ = 0; + throw std::runtime_error("Dispatcher: failed to allocate Stream-K reduction workspace"); + } + workspace_bytes_ = bytes; +} + void Dispatcher::set_heuristic(HeuristicFunction heuristic) { heuristic_ = heuristic; @@ -66,7 +100,17 @@ float Dispatcher::run_fused(const void* a_ptr, } kernel->set_benchmarking(benchmarking_); - return kernel->run(a_ptr, b_ptr, c_ptr, d_ptrs, problem, stream); + + // Size and own the reduction workspace (0 for non-Stream-K and for Atomic). + const std::size_t ws_bytes = kernel->get_workspace_size(problem); + void* workspace = nullptr; + if(ws_bytes > 0) + { + ensure_workspace(ws_bytes); + workspace = workspace_; + } + + return kernel->run(a_ptr, b_ptr, c_ptr, d_ptrs, workspace, problem, stream); } float Dispatcher::run_explicit(const std::string& kernel_id, @@ -92,7 +136,17 @@ float Dispatcher::run_explicit(const std::string& kernel_id, } kernel->set_benchmarking(benchmarking_); - return kernel->run(a_ptr, b_ptr, c_ptr, d_ptrs, problem, stream); + + // Size and own the reduction workspace (0 for non-Stream-K and for Atomic). + const std::size_t ws_bytes = kernel->get_workspace_size(problem); + void* workspace = nullptr; + if(ws_bytes > 0) + { + ensure_workspace(ws_bytes); + workspace = workspace_; + } + + return kernel->run(a_ptr, b_ptr, c_ptr, d_ptrs, workspace, problem, stream); } bool Dispatcher::validate(const void* a_ptr, diff --git a/projects/composablekernel/dispatcher/tests/CMakeLists.txt b/projects/composablekernel/dispatcher/tests/CMakeLists.txt index a18663f76d59..473c1a716ac0 100644 --- a/projects/composablekernel/dispatcher/tests/CMakeLists.txt +++ b/projects/composablekernel/dispatcher/tests/CMakeLists.txt @@ -126,6 +126,20 @@ set_tests_properties(dispatcher_test_fmha_parity PROPERTIES ENVIRONMENT "PYTHONPATH=${CMAKE_CURRENT_SOURCE_DIR}/../python:${CMAKE_CURRENT_SOURCE_DIR}/../codegen:${CMAKE_CURRENT_SOURCE_DIR}/../scripts" ) +# Stream-K deep-core registry test (requires GPU + hipcc; SKIPs otherwise) +add_test( + NAME dispatcher_test_streamk_registry + COMMAND ${Python3_EXECUTABLE} ${CMAKE_CURRENT_SOURCE_DIR}/test_streamk_registry.py + WORKING_DIRECTORY ${CMAKE_CURRENT_SOURCE_DIR}/.. +) + +set_tests_properties(dispatcher_test_streamk_registry PROPERTIES + LABELS "dispatcher;python;streamk;gpu" + TIMEOUT 900 + SKIP_RETURN_CODE 77 + ENVIRONMENT "PYTHONPATH=${CMAKE_CURRENT_SOURCE_DIR}/../python:${CMAKE_CURRENT_SOURCE_DIR}/../codegen:${CMAKE_CURRENT_SOURCE_DIR}/../scripts" +) + # Stress Test Script add_test( NAME dispatcher_stress_test diff --git a/projects/composablekernel/dispatcher/tests/test_streamk_registry.py b/projects/composablekernel/dispatcher/tests/test_streamk_registry.py new file mode 100644 index 000000000000..93924ea1950b --- /dev/null +++ b/projects/composablekernel/dispatcher/tests/test_streamk_registry.py @@ -0,0 +1,234 @@ +#!/usr/bin/env python3 + +# Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +# SPDX-License-Identifier: MIT + +""" +Stream-K deep-core registry test (requires a GPU + hipcc). + +Guards the deep-core path that lets Stream-K ride the registry like regular GEMM: +codegen -> generated SK wrapper -> Registry -> Dispatcher::run() (workspace alloc ++ strategy-aware reset) -> generated_tile_backend_streamk -> verify vs reference. + +Each reduction strategy (atomic/linear/tree) is a *distinct compiled kernel* +(SkReductionStrategy is a compile-time constexpr), so we generate all three from a +single tile config and build the 04 registry driver once per strategy, force- +including that strategy's header. For each we assert: + * the encode_identifier() suffix matches the strategy (..._streamk[_linear|_tree]) + * the Dispatcher selects that kernel by Problem::reduction_strategy + * the result verifies against the reference GEMM + +The test SKIPs (exit 77) when no GPU or no hipcc is available, so it is safe in +CPU-only CI; it only runs the heavy build+launch where a GPU is present. + +Usage: + python3 test_streamk_registry.py + python3 test_streamk_registry.py --arch gfx942 --m 3840 --n 4096 --k 2048 +""" + +import argparse +import json +import os +import shutil +import subprocess +import sys +import tempfile +from pathlib import Path + +DISPATCHER_DIR = Path(__file__).resolve().parent.parent +CK_DIR = DISPATCHER_DIR.parent +CODEGEN = DISPATCHER_DIR / "codegen" / "unified_gemm_codegen.py" +DRIVER = DISPATCHER_DIR / "examples" / "gemm" / "cpp" / "04_streamk_registry_driver.cpp" +REGISTRY_SRC = DISPATCHER_DIR / "src" / "registry.cpp" +DISPATCHER_SRC = DISPATCHER_DIR / "src" / "dispatcher.cpp" + +SKIP = 77 # ctest SKIP_RETURN_CODE + +# One tile config, all three reduction strategies. +TILE = "128x128x64_2x2x1_32x32x16" +TILE_CONFIG_JSON = json.dumps( + { + "tile_config": { + "tile_m": [128], "tile_n": [128], "tile_k": [64], + "warp_m": [2], "warp_n": [2], "warp_k": [1], + "warp_tile_m": [32], "warp_tile_n": [32], "warp_tile_k": [16], + "block_size": [256], + }, + "trait_config": { + "pipeline": ["compv3"], "epilogue": ["cshuffle"], "scheduler": ["intrawave"], + "pad_m": [False], "pad_n": [False], "pad_k": [False], "persistent": [False], + }, + "streamk_config": {"reduction_strategy": ["atomic", "linear", "tree"]}, + } +) + +# strategy -> (header variant suffix, expected encode_identifier suffix) +STRATEGIES = { + "atomic": ("streamk", "_streamk"), + "linear": ("streamk_linear", "_streamk_linear"), + "tree": ("streamk_tree", "_streamk_tree"), +} + +# Datatypes the Stream-K dispatcher codegen supports end-to-end. fp8/bf8 inputs +# accumulate in fp32 and write an fp16 C tensor (get_output_dtype), matching +# Tile Engine; the registry identifier keys on the input dtype (dtype_a), so the +# expected encode_identifier prefix is "{dtype}_rcr" for each. +DATATYPES = ["fp16", "bf16", "fp8", "bf8"] + + +def detect_arch(fallback=None): + try: + sys.path.insert(0, str(DISPATCHER_DIR / "python")) + from dispatcher_common import detect_gpu_arch # noqa: E402 + + return detect_gpu_arch() + except Exception: + out = shutil.which("rocminfo") + if out: + try: + txt = subprocess.run( + ["rocminfo"], capture_output=True, text=True, timeout=30 + ).stdout + for line in txt.splitlines(): + if "gfx" in line and "Name:" in line: + return line.split("gfx")[1].split()[0].join(["gfx", ""]) + except Exception: + pass + return fallback + + +def run(cmd, **kw): + return subprocess.run(cmd, capture_output=True, text=True, **kw) + + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument("--arch", default=None) + ap.add_argument("--m", type=int, default=3840) + ap.add_argument("--n", type=int, default=4096) + ap.add_argument("--k", type=int, default=2048) + args = ap.parse_args() + + hipcc = shutil.which("hipcc") + if not hipcc: + print("SKIP: hipcc not found") + return SKIP + + arch = args.arch or detect_arch() + if not arch: + print("SKIP: no GPU / could not detect gfx arch") + return SKIP + print(f"Stream-K registry test on {arch} @ {args.m}x{args.n}x{args.k}") + + inc = ["-I", str(CK_DIR / "include"), "-I", str(DISPATCHER_DIR / "include")] + + with tempfile.TemporaryDirectory(prefix="sk_reg_test_") as td: + # Build the dtype-independent core objects once (no force-include). + reg_o, disp_o = Path(td) / "registry.o", Path(td) / "dispatcher.o" + for src, obj in ((REGISTRY_SRC, reg_o), (DISPATCHER_SRC, disp_o)): + c = run( + [hipcc, "-std=c++17", f"--offload-arch={arch}", "-O3", *inc, + "-c", str(src), "-o", str(obj)], + timeout=900, + ) + if c.returncode != 0: + print(f"FAIL: compiling {src.name}\n" + c.stderr[-2000:]) + return 1 + + failures = [] + for dtype in DATATYPES: + failures += run_for_dtype( + dtype, td, arch, args, hipcc, inc, reg_o, disp_o + ) + + if failures: + print("\nSTREAM-K REGISTRY TEST FAILED:") + for f in failures: + print(" - " + f) + return 1 + + print( + "All Stream-K datatypes " + f"({', '.join(DATATYPES)}) registered, dispatched, and verified." + ) + return 0 + + +def run_for_dtype(dtype, td, arch, args, hipcc, inc, reg_o, disp_o): + """Generate + build + run all reduction strategies for one datatype. + + Returns a list of failure strings (empty on success).""" + failures = [] + gen = Path(td) / f"gen_{dtype}" + + # 1) generate all three strategy headers from one tile config + g = run( + [ + sys.executable, str(CODEGEN), + "--datatype", dtype, "--layout", "rcr", + "--gpu-target", arch, "--variants", "stream_k", + "--tile-config-json", TILE_CONFIG_JSON, + "--output-dir", str(gen), + ], + timeout=600, + ) + if g.returncode != 0: + return [f"{dtype}: codegen failed\n" + g.stderr[-2000:]] + + for strat, (variant, want_suffix) in STRATEGIES.items(): + tag = f"{dtype}/{strat}" + header = gen / ( + f"gemm_{dtype}_rcr_compv3_cshuffle_intrawave_" + f"False_False_False_False_{TILE}_{variant}.hpp" + ) + if not header.exists(): + failures.append(f"{tag}: generated header missing ({header.name})") + continue + + stem = f"{dtype}_{variant}" + drv_o, exe = Path(td) / f"d_{stem}.o", Path(td) / f"skreg_{stem}" + c = run( + [hipcc, "-std=c++17", f"--offload-arch={arch}", "-O3", + "-DCK_TILE_SINGLE_KERNEL_INCLUDE", f'-DGFX_ARCH="{arch}"', + *inc, "-I", str(gen), "-include", str(header), + "-c", str(DRIVER), "-o", str(drv_o)], + timeout=900, + ) + if c.returncode != 0: + failures.append(f"{tag}: driver compile failed\n{c.stderr[-1500:]}") + continue + l = run( + [hipcc, f"--offload-arch={arch}", str(drv_o), str(disp_o), + str(reg_o), "-o", str(exe)], + timeout=300, + ) + if l.returncode != 0: + failures.append(f"{tag}: link failed\n{l.stderr[-1500:]}") + continue + + r = run( + [str(exe), "--m", str(args.m), "--n", str(args.n), + "--k", str(args.k), "--strategy", strat, "--validate", "1"], + timeout=300, + ) + out = r.stdout + ok_verify = "Verification: PASS" in out + ok_suffix = f"identifier={dtype}_rcr" in out and want_suffix in out.split( + "identifier=" + )[1].split()[0] + if r.returncode != 0 or not ok_verify or not ok_suffix: + failures.append( + f"{tag}: rc={r.returncode} verify={ok_verify} " + f"suffix_ok={ok_suffix}\n{out[-800:]}{r.stderr[-400:]}" + ) + else: + tflops = next( + (ln for ln in out.splitlines() if "TFlops" in ln), "" + ).strip() + print(f" PASS {dtype:5s} {strat:6s} -> {want_suffix} | {tflops}") + + return failures + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/projects/composablekernel/test/ck_tile/CMakeLists.txt b/projects/composablekernel/test/ck_tile/CMakeLists.txt index 441356265cc6..3222475240f0 100644 --- a/projects/composablekernel/test/ck_tile/CMakeLists.txt +++ b/projects/composablekernel/test/ck_tile/CMakeLists.txt @@ -71,9 +71,6 @@ if(BUILD_CK_TILE_FMHA_TESTS) add_subdirectory(fmha) endif() if(BUILD_CK_TILE_ENGINE_TESTS) -# TODO: The Universal GEMM tile engine test will be either removed -# or moved to the appropriate location in future work. -# add_subdirectory(gemm_tile_engine) add_subdirectory(pooling_tile_engine) endif() add_subdirectory(pooling) diff --git a/projects/composablekernel/test/ck_tile/gemm_tile_engine/CMakeLists.txt b/projects/composablekernel/test/ck_tile/gemm_tile_engine/CMakeLists.txt deleted file mode 100644 index 374370f57076..000000000000 --- a/projects/composablekernel/test/ck_tile/gemm_tile_engine/CMakeLists.txt +++ /dev/null @@ -1,348 +0,0 @@ -# Copyright (c) Advanced Micro Devices, Inc., or its affiliates. -# SPDX-License-Identifier: MIT - -# ============================================================================ -# GEMM Tile Engine Unit Tests -# -# This CMake file creates unit tests for tile_engine generated GEMM kernels. -# It follows the exact same build patterns as tile_engine for consistency -# and reliability. Each kernel configuration gets its own test executable. -# ============================================================================ - -# Locate tile_engine GEMM scripts directory -set(TILE_ENGINE_GEMM_DIR "${PROJECT_SOURCE_DIR}/tile_engine/ops/gemm/gemm_universal") - -if(NOT EXISTS ${TILE_ENGINE_GEMM_DIR}) - message(WARNING "Tile engine directory not found: ${TILE_ENGINE_GEMM_DIR}") - return() -endif() - -# ============================================================================ -# create_individual_gemm_test_target -# -# Creates a single test executable for a specific kernel configuration. -# Mirrors tile_engine's create_individual_gemm_target function for consistency. -# -# Parameters: -# datatype - Data type (fp16, bf16, fp32, etc.) -# layout - Matrix layout (rcr, rrr, ccr, crr) -# config_name - Configuration file name without .json extension -# trait - Kernel trait combination string -# tile_config - Tile configuration parameters -# config_json - Full path to JSON configuration file -# ============================================================================ -function(create_individual_gemm_test_target datatype layout config_name trait tile_config config_json) - set(target_name "test_gemm_universal_tile_engine_${datatype}_${layout}_${config_name}_${trait}_${tile_config}") - set(working_path "${CMAKE_CURRENT_BINARY_DIR}/${datatype}/${layout}/${config_name}") - - # Generated header path (already created during cmake configuration) - set(test_header "${working_path}/gemm_universal_single_${datatype}_${layout}_${trait}_${tile_config}.hpp") - set(test_params_header "${working_path}/test_params.hpp") - - # Verify header exists (should have been generated during cmake configuration) - if(NOT EXISTS ${test_header}) - message(WARNING "Generated header not found: ${test_header}") - return() - endif() - - # Verify test parameters header exists - if(NOT EXISTS ${test_params_header}) - message(WARNING "Test parameters header not found: ${test_params_header}") - return() - endif() - - - # Create GTest executable for this kernel configuration - add_gtest_executable(${target_name} - ${CMAKE_CURRENT_SOURCE_DIR}/test_gemm_simple.cpp - ) - - # Configure GPU architectures for HIP compilation - set_property(TARGET ${target_name} PROPERTY HIP_ARCHITECTURES ${GEMM_TEST_GPU_TARGETS}) - - # Define preprocessor macros for generated header location and test parameters - target_compile_definitions(${target_name} PRIVATE - GEMM_SINGLE_INSTANCE_HPP="${test_header}" - GEMM_TEST_PARAMS_HPP="${test_params_header}" - ) - - # Include directories for headers and dependencies - target_include_directories(${target_name} PRIVATE - ${PROJECT_SOURCE_DIR}/include - ${PROJECT_BINARY_DIR}/include - ${PROJECT_SOURCE_DIR} # Root directory for tile_engine access - ${GTEST_INCLUDE_DIRS} - ) - - # Compiler options matching tile_engine requirements - target_compile_options(${target_name} PRIVATE - -Wno-undefined-func-template # Suppress template warnings - -Wno-float-equal # Allow floating point comparisons - --offload-compress # Enable GPU code compression - -include ${test_header} # Auto-include generated header - ) - - # Add FP8 format definitions for proper data type interpretation - if(CK_USE_OCP_FP8) - target_compile_options(${target_name} PRIVATE -DCK_TILE_USE_OCP_FP8) - endif() - - message(DEBUG " Created test target: ${target_name}") -endfunction() - -# ============================================================================ -# build_gemm_test_targets -# -# Builds all test targets for a specific datatype/layout/config combination. -# Uses tile_engine's two-step process: list kernels, then generate tests. -# -# Parameters: -# datatype - Data type (fp16, bf16, fp32, etc.) -# layout - Matrix layout (rcr, rrr, ccr, crr) -# config_name - Configuration file name without .json extension -# ============================================================================ -function(build_gemm_test_targets datatype layout config_name) - set(working_path "${CMAKE_CURRENT_BINARY_DIR}/${datatype}/${layout}/${config_name}") - - # Locate and validate configuration file - set(config_filename "${config_name}.json") - set(json_blob "${CMAKE_CURRENT_SOURCE_DIR}/configs/${config_filename}") - - if(NOT EXISTS ${json_blob}) - message(WARNING "Test config file not found: ${json_blob}") - return() - endif() - - # Prepare build directory for this configuration - file(MAKE_DIRECTORY ${working_path}) - - # STEP 1: Discovery phase - list all valid kernel configurations - execute_process( - COMMAND ${Python3_EXECUTABLE} -u ${TILE_ENGINE_GEMM_DIR}/gemm_universal_instance_builder.py - --working_path ${working_path} - --datatype ${datatype} - --layout ${layout} - --config_json ${json_blob} - --list_kernels - --gpu_target "${GEMM_TEST_GPU_TARGETS}" - WORKING_DIRECTORY ${TILE_ENGINE_GEMM_DIR} - RESULT_VARIABLE ret - OUTPUT_VARIABLE list_output - ERROR_VARIABLE list_error - ) - - if(NOT ret EQUAL 0) - message(WARNING "Failed to list kernels for ${datatype}_${layout}_${config_name}: ${list_error}") - return() - endif() - - # Verify kernel list file was generated - if(NOT EXISTS ${working_path}/gemm_kernel_list.txt) - message(DEBUG "No kernels found for ${datatype}_${layout}_${config_name} (validation filtered out all combinations)") - return() - endif() - - message(DEBUG "Building tests for ${datatype}_${layout}_${config_name}") - - # STEP 2a: Extract test parameters from config - set(test_params_file "${working_path}/test_params.hpp") - execute_process( - COMMAND ${Python3_EXECUTABLE} -u ${CMAKE_CURRENT_SOURCE_DIR}/extract_test_params.py - --config_file ${json_blob} - --output_file ${test_params_file} - WORKING_DIRECTORY ${CMAKE_CURRENT_SOURCE_DIR} - RESULT_VARIABLE extract_ret - OUTPUT_VARIABLE extract_output - ERROR_VARIABLE extract_error - ) - - if(NOT extract_ret EQUAL 0) - message(WARNING "Failed to extract test parameters for ${datatype}_${layout}: ${extract_error}") - return() - endif() - - # STEP 2b: Header generation phase - generate headers using --gen_single - message(STATUS " Generating headers using --gen_single...") - - file(STRINGS ${working_path}/gemm_kernel_list.txt kernel_lines) - set(gen_count 0) - - foreach(line IN LISTS kernel_lines) - # Parse kernel specification format: kernel_name|tile_config|trait_combo - string(REPLACE "|" ";" parts "${line}") - list(LENGTH parts parts_len) - if(parts_len EQUAL 3) - list(GET parts 0 kernel_name) - list(GET parts 1 tile_config) - list(GET parts 2 trait_combo) - - # Generate header using --gen_single - execute_process( - COMMAND ${Python3_EXECUTABLE} -u ${TILE_ENGINE_GEMM_DIR}/gemm_universal_instance_builder.py - --working_path ${working_path} - --gpu_target "${GEMM_TEST_GPU_TARGETS}" - --datatype ${datatype} - --layout ${layout} - --config_json ${json_blob} - --gen_single - --kernel_name "${kernel_name}" - --tile_config "${tile_config}" - --trait_combo "${trait_combo}" - WORKING_DIRECTORY ${TILE_ENGINE_GEMM_DIR} - RESULT_VARIABLE gen_ret - OUTPUT_VARIABLE gen_output - ERROR_VARIABLE gen_error - ) - - if(NOT gen_ret EQUAL 0) - message(WARNING "Failed to generate header for ${kernel_name}: ${gen_error}") - else() - math(EXPR gen_count "${gen_count} + 1") - endif() - endif() - endforeach() - - message(STATUS " Generated ${gen_count} headers for ${datatype}_${layout}") - - # STEP 3: Target creation phase - create test targets - message(STATUS " Creating test targets...") - file(STRINGS ${working_path}/gemm_kernel_list.txt kernel_lines) - set(test_count 0) - foreach(line IN LISTS kernel_lines) - # Parse kernel specification format: kernel_name|tile_config|trait_combo - string(REPLACE "|" ";" parts "${line}") - list(LENGTH parts parts_len) - if(parts_len EQUAL 3) - list(GET parts 0 kernel_name) - list(GET parts 1 tile_config) - list(GET parts 2 trait_combo) - - # Generate test target for this kernel configuration - create_individual_gemm_test_target("${datatype}" "${layout}" "${config_name}" "${trait_combo}" "${tile_config}" "${json_blob}") - math(EXPR test_count "${test_count} + 1") - endif() - endforeach() - message(STATUS " Created ${test_count} test targets for ${datatype}_${layout}") -endfunction()# ============================================================================ -# MAIN EXECUTION - Test Target Generation -# ============================================================================ - -message(STATUS "=== Starting GEMM Tile Engine Test Configuration ===") -message(STATUS "SUPPORTED_GPU_TARGETS: ${SUPPORTED_GPU_TARGETS}") - -# GPU architecture filtering - only build tests for supported architectures -set(GEMM_TEST_GPU_TARGETS "") -set(DESIRED_TARGETS "gfx90a;gfx942;gfx950;gfx1201;gfx12-generic") - -foreach(target IN LISTS SUPPORTED_GPU_TARGETS) - if(target IN_LIST DESIRED_TARGETS) - list(APPEND GEMM_TEST_GPU_TARGETS ${target}) - message(STATUS " Adding GPU target for tests: ${target}") - endif() -endforeach() - -# Early exit if no compatible GPU architectures are available -if(NOT GEMM_TEST_GPU_TARGETS) - message(WARNING "Skipping GEMM Tile Engine tests: No supported GPU targets (gfx90a, gfx942, gfx950, gfx1201, gfx12-generic) found in SUPPORTED_GPU_TARGETS: ${SUPPORTED_GPU_TARGETS}") - return() -endif() - -message(STATUS "Building GEMM tile engine tests for GPU targets: ${GEMM_TEST_GPU_TARGETS}") - - # Enable parallel compilation optimizations - # Set up job pools for better parallel compilation control - set_property(GLOBAL PROPERTY JOB_POOLS - compile_heavy=4 # Limit heavy compilations to prevent OOM - compile_normal=16 # Allow more parallel normal compilations - ) - - # Enable compiler cache if available and explicitly requested - # Disabled by default due to permission issues in CI environments - option(ENABLE_CCACHE_TESTS "Enable ccache for test compilation" OFF) - if(ENABLE_CCACHE_TESTS) - find_program(CCACHE_PROGRAM ccache) - if(CCACHE_PROGRAM) - set(CMAKE_CXX_COMPILER_LAUNCHER ${CCACHE_PROGRAM}) - message(STATUS "Using ccache for faster test compilation") - else() - message(WARNING "ccache requested but not found") - endif() - else() - message(STATUS "ccache disabled for tests (use -DENABLE_CCACHE_TESTS=ON to enable)") - endif() - -# ============================================================================ -# Test Configuration Matrix - Clean Focused Design -# ============================================================================ - -# All supported data types and layouts for comprehensive testing -# Note: fp64 not included (no MFMA hardware support) -set(TEST_DATATYPES "fp16;fp8;bf16;fp32") -set(TEST_LAYOUTS "rcr;rrr;ccr;crr") - -# ============================================================================ -# Test Target Generation - Datatype-Specific Categories -# ============================================================================ - -# 1. SMALL DATATYPES: Test optimized config for small data types (fp8, fp16, bf16) -# These data types can use larger warp tiles due to smaller memory footprint -set(SMALL_DATATYPE_CONFIG "small_datatype_config") -set(SMALL_DATATYPE_CONFIG_FILE "${CMAKE_CURRENT_SOURCE_DIR}/configs/${SMALL_DATATYPE_CONFIG}.json") -set(SMALL_DATATYPES "fp8;fp16;bf16") - -if(EXISTS ${SMALL_DATATYPE_CONFIG_FILE}) - message(STATUS "Processing small datatype config: ${SMALL_DATATYPE_CONFIG} (fp8, fp16, bf16)") - foreach(datatype IN LISTS SMALL_DATATYPES) - # fp8, fp16, bf16: testing all layouts (rcr, rrr, ccr, crr) - foreach(layout IN LISTS TEST_LAYOUTS) - build_gemm_test_targets("${datatype}" "${layout}" "${SMALL_DATATYPE_CONFIG}") - endforeach() - endforeach() -else() - message(WARNING "Small datatype config file not found: ${SMALL_DATATYPE_CONFIG_FILE}") -endif() - -# 2. PADDING COVERAGE: Test padding combinations with fixed fp16/rcr configuration -# This focuses on padding behavior (pad_m, pad_n, pad_k) -set(PADDING_CONFIG "padding_coverage_config") -set(PADDING_CONFIG_FILE "${CMAKE_CURRENT_SOURCE_DIR}/configs/${PADDING_CONFIG}.json") - -if(EXISTS ${PADDING_CONFIG_FILE}) - message(STATUS "Processing padding config: ${PADDING_CONFIG} (fp16/rcr only)") - build_gemm_test_targets("fp16" "rcr" "${PADDING_CONFIG}") -else() - message(WARNING "Padding config file not found: ${PADDING_CONFIG_FILE}") -endif() - -# 3. COVERAGE LEVEL: Quick or comprehensive testing -# Quick: ~144 kernels with multiple tile sizes and trait combinations -# Comprehensive: Several thousand kernels with extensive tile sizes, warp configurations, and all trait combinations -set(COVERAGE_LEVEL "quick" CACHE STRING "Coverage level: quick or comprehensive") -set_property(CACHE COVERAGE_LEVEL PROPERTY STRINGS "quick" "comprehensive") - -if(COVERAGE_LEVEL STREQUAL "quick") - set(COVERAGE_CONFIG "quick_coverage_config") - set(COVERAGE_DESC "Quick - approximately 144 kernels with trait combinations") -elseif(COVERAGE_LEVEL STREQUAL "comprehensive") - set(COVERAGE_CONFIG "comprehensive_coverage_config") - set(COVERAGE_DESC "Comprehensive - several thousand kernels with extensive tile and trait coverage") -else() - message(FATAL_ERROR "Invalid COVERAGE_LEVEL: ${COVERAGE_LEVEL}. Must be 'quick' or 'comprehensive'") -endif() - -set(COVERAGE_CONFIG_FILE "${CMAKE_CURRENT_SOURCE_DIR}/configs/${COVERAGE_CONFIG}.json") - -if(EXISTS ${COVERAGE_CONFIG_FILE}) - message(STATUS "Processing coverage config: ${COVERAGE_LEVEL} - ${COVERAGE_DESC}") - build_gemm_test_targets("fp16" "rcr" "${COVERAGE_CONFIG}") -else() - message(WARNING "Coverage config file not found: ${COVERAGE_CONFIG_FILE}") -endif() -# ============================================================================ - - -message(STATUS "GEMM tile engine tests configured with datatype-specific design:") -message(STATUS " - Small datatypes: fp8/fp16/bf16 (all layouts)") -message(STATUS " - Padding coverage with fp16/rcr") -message(STATUS " - Coverage level: ${COVERAGE_LEVEL} (~144 kernels quick, several thousand comprehensive)") -message(STATUS " Use -DCOVERAGE_LEVEL=comprehensive for extensive testing") diff --git a/projects/composablekernel/test/ck_tile/gemm_tile_engine/README.md b/projects/composablekernel/test/ck_tile/gemm_tile_engine/README.md deleted file mode 100644 index 87ce0c9fd05c..000000000000 --- a/projects/composablekernel/test/ck_tile/gemm_tile_engine/README.md +++ /dev/null @@ -1,85 +0,0 @@ -# GEMM Tile Engine Unit Tests - -## How It Works - -This unit test system integrates **tile_engine's kernel generation** into automated testing: - -1. **Uses tile_engine scripts directly**: Same Python scripts that generate tile_engine kernels -2. **JSON-based configuration**: Define test parameters in JSON files (like tile_engine) -3. **Build-time generation**: CMake calls tile_engine scripts to generate kernel headers -4. **Individual test executables**: Each kernel configuration becomes a separate test -5. **Tile_engine verification**: Uses exact same error thresholds and validation as tile_engine - -## Tile Engine Integration - -``` -JSON Config → tile_engine Python scripts → Generated Headers → Test Executables -``` - -- **`--list_kernels`**: Get available kernel configurations from JSON -- **`--gen_individual`**: Generate all kernel headers in parallel during CMake configuration -- **`--gen_single`**: Generate individual kernel header for each configuration -- **Same verification**: Uses tile_engine's adaptive error thresholds and reference calculations -- **Same patterns**: Follows tile_engine's tensor initialization, stride calculation, and kernel launching - -### Config-Specific Test Parameters - -Each test configuration can specify optimized problem sizes in its JSON file: -- **`test_params.problem_sizes`**: Array of `{m, n, k, split_k}` configurations -- **CMake extraction**: `extract_test_params.py` generates config-specific test parameter files -- **Build integration**: Each test target uses parameters appropriate for its kernel configuration -- **Optimized testing**: Different configs test different problem sizes that showcase their strengths - - -The key idea: **Unit tests that use tile_engine's exact kernel generation and verification methodology** instead of creating separate test infrastructure. - -## Test Configurations - -### 1. **Simple Test** (`simple_test_config.json`) -- **Purpose**: Basic functionality validation -- **Config**: 128x128x64, warp 2x2x1, warp_tile 16x16x16 -- **Traits**: compv3 + compv4 pipelines -- **Coverage**: ~2 kernels per datatype/layout - -### 2. **Small Datatype** (`small_datatype_config.json`) -- **Purpose**: Optimized for fp8/fp16/bf16 data types -- **Config**: 128x128x32, warp 2x2x1, warp_tile 32x32x16 -- **Traits**: compv3 pipeline only -- **Coverage**: All 4 layouts (rcr, rrr, ccr, crr) for fp8, fp16, bf16 - -### 3. **Padding Coverage** (`padding_coverage_config.json`) -- **Purpose**: Test padding behavior with all padding flags enabled -- **Config**: Fixed 64x64x32, warp 2x2x1, warp_tile 32x32x16 -- **Padding**: All enabled (pad_m=true, pad_n=true, pad_k=true) -- **Problem sizes**: Vector-aligned but not tile-aligned (104×104×56, 200×152×80, 152×200×64) -- **Coverage**: 1 kernel configuration testing padding with irregular sizes - -### 4. **Coverage Testing** (Quick or Comprehensive) -- **Purpose**: Comprehensive testing across tile sizes, warp configurations, and trait combinations -- **Quick** (`quick_coverage_config.json`): Approximately 144 kernels - - tile_m/n: [32, 64, 256], tile_k: [16, 32] - - warp config: 2×2×1, warp_tile 16×16×16 - - Traits: 3 pipelines × 2 epilogues × 2 schedulers (persistent=false only) - - Focused set testing trait combinations with multiple tile sizes -- **Comprehensive** (`comprehensive_coverage_config.json`): Several thousand kernels - - tile_m/n: [16-256 step 16] - - tile_k: [16, 32, 64] - - warp_m/n: [1, 2, 4], warp_tile_m/n: [16, 32], warp_tile_k: [16, 32] - - Traits: 3 pipelines × 2 epilogues × 2 schedulers × 2 persistent - - Extensive coverage across all tile sizes, warp configurations, and trait combinations - - Exact count varies based on validation filtering -- **Note**: Use CMake option `-DCOVERAGE_LEVEL=comprehensive` to enable comprehensive testing (default is quick) - -## Data Type Support -- ✅ **fp8, fp16, bf16**: Fully supported - all layouts (rcr, rrr, ccr, crr) -- ❌ **fp64**: Not supported (hardware MFMA limitation) -- ⏳ **fp32, bf8, pk-int4-t**: Not yet supported by gemm_instance_builder (will be added later) - -## Test Result Behavior - -Tests automatically handle unsupported configurations through runtime validation: -- **PASSED**: Kernel executed correctly with results within error thresholds ✅ -- **SKIPPED**: Kernel validation returned "Arguments not supported" (expected for certain problem sizes/configurations) ⚠️ -- **FAILED**: Actual error or incorrect computation results ❌ - -When a kernel's `IsSupportedArgument()` check fails (e.g., due to vector alignment requirements, dimension constraints, or padding limitations), the test is automatically skipped rather than failed. This allows comprehensive testing across various problem sizes while gracefully handling configurations that don't meet specific kernel requirements. diff --git a/projects/composablekernel/test/ck_tile/gemm_tile_engine/configs/comprehensive_coverage_config.json b/projects/composablekernel/test/ck_tile/gemm_tile_engine/configs/comprehensive_coverage_config.json deleted file mode 100644 index f2524e4a619d..000000000000 --- a/projects/composablekernel/test/ck_tile/gemm_tile_engine/configs/comprehensive_coverage_config.json +++ /dev/null @@ -1,37 +0,0 @@ -{ - "problem": { - "description": "Comprehensive coverage testing - extensive tile size coverage (16-256, step 16) with multiple warp configurations and all trait combinations. Several thousand kernels." - }, - "test_params": { - "problem_sizes": [ - {"m": 512, "n": 512, "k": 256, "split_k": 1}, - {"m": 1024, "n": 512, "k": 512, "split_k": 1}, - {"m": 512, "n": 1024, "k": 512, "split_k": 1}, - {"m": 1024, "n": 1024, "k": 256, "split_k": 1}, - {"m": 1024, "n": 1024, "k": 256, "split_k": 2}, - {"m": 1024, "n": 1024, "k": 256, "split_k": 4} - ] - }, - "tile_config": { - "tile_m": {"values": [16, 32, 48, 64, 80, 96, 112, 128, 144, 160, 176, 192, 208, 224, 240, 256]}, - "tile_n": {"values": [16, 32, 48, 64, 80, 96, 112, 128, 144, 160, 176, 192, 208, 224, 240, 256]}, - "tile_k": {"values": [16, 32, 64]}, - "warp_m": {"values": [1, 2, 4]}, - "warp_n": {"values": [1, 2, 4]}, - "warp_k": {"values": [1]}, - "warp_tile_m": {"values": [16, 32]}, - "warp_tile_n": {"values": [16, 32]}, - "warp_tile_k": {"values": [8, 16, 32, 64, 128]} - }, - "trait_config": { - "pipeline": {"values": ["mem", "compv3", "compv4"]}, - "epilogue": {"values": ["default", "cshuffle"]}, - "scheduler": {"values": ["intrawave", "interwave"]}, - "pad_m": {"values": [false]}, - "pad_n": {"values": [false]}, - "pad_k": {"values": [false]}, - "persistent": {"values": [true, false]} - }, - "k_block_per_cu": 1, - "permute_n": false -} diff --git a/projects/composablekernel/test/ck_tile/gemm_tile_engine/configs/large_datatype_config.json b/projects/composablekernel/test/ck_tile/gemm_tile_engine/configs/large_datatype_config.json deleted file mode 100644 index e9fcb6fb8007..000000000000 --- a/projects/composablekernel/test/ck_tile/gemm_tile_engine/configs/large_datatype_config.json +++ /dev/null @@ -1,34 +0,0 @@ -{ - "problem": { - "description": "Configuration optimized for large data types (fp32) with smaller warp tiles due to memory constraints" - }, - "test_params": { - "problem_sizes": [ - {"m": 512, "n": 512, "k": 128, "split_k": 1}, - {"m": 512, "n": 256, "k": 192, "split_k": 1}, - {"m": 256, "n": 384, "k": 192, "split_k": 1} - ] - }, - "tile_config": { - "tile_m": {"values": [256]}, - "tile_n": {"values": [128]}, - "tile_k": {"values": [32]}, - "warp_m": {"values": [2]}, - "warp_n": {"values": [2]}, - "warp_k": {"values": [1]}, - "warp_tile_m": {"values": [16]}, - "warp_tile_n": {"values": [16]}, - "warp_tile_k": {"values": [16]} - }, - "trait_config": { - "pipeline": {"values": ["compv3"]}, - "epilogue": {"values": ["default"]}, - "scheduler": {"values": ["intrawave"]}, - "pad_m": {"values": [false]}, - "pad_n": {"values": [false]}, - "pad_k": {"values": [false]}, - "persistent": {"values": [false]} - }, - "k_block_per_cu": 1, - "permute_n": false -} diff --git a/projects/composablekernel/test/ck_tile/gemm_tile_engine/configs/padding_coverage_config.json b/projects/composablekernel/test/ck_tile/gemm_tile_engine/configs/padding_coverage_config.json deleted file mode 100644 index 33bada839de5..000000000000 --- a/projects/composablekernel/test/ck_tile/gemm_tile_engine/configs/padding_coverage_config.json +++ /dev/null @@ -1,34 +0,0 @@ -{ - "problem": { - "description": "Padding coverage testing - fixed config with fp16/rcr, varying only padding combinations" - }, - "test_params": { - "problem_sizes": [ - {"m": 104, "n": 104, "k": 56, "split_k": 1}, - {"m": 200, "n": 152, "k": 80, "split_k": 1}, - {"m": 152, "n": 200, "k": 64, "split_k": 1} - ] - }, - "tile_config": { - "tile_m": {"values": [64]}, - "tile_n": {"values": [64]}, - "tile_k": {"values": [32]}, - "warp_m": {"values": [2]}, - "warp_n": {"values": [2]}, - "warp_k": {"values": [1]}, - "warp_tile_m": {"values": [32]}, - "warp_tile_n": {"values": [32]}, - "warp_tile_k": {"values": [16]} - }, - "trait_config": { - "pipeline": {"values": ["compv3"]}, - "epilogue": {"values": ["default"]}, - "scheduler": {"values": ["intrawave"]}, - "pad_m": {"values": [true]}, - "pad_n": {"values": [true]}, - "pad_k": {"values": [true]}, - "persistent": {"values": [false]} - }, - "k_block_per_cu": 1, - "permute_n": false -} diff --git a/projects/composablekernel/test/ck_tile/gemm_tile_engine/configs/quick_coverage_config.json b/projects/composablekernel/test/ck_tile/gemm_tile_engine/configs/quick_coverage_config.json deleted file mode 100644 index dcc6e99aee5a..000000000000 --- a/projects/composablekernel/test/ck_tile/gemm_tile_engine/configs/quick_coverage_config.json +++ /dev/null @@ -1,34 +0,0 @@ -{ - "problem": { - "description": "Quick coverage testing - tests multiple tile sizes with all trait combinations (pipelines, epilogues, schedulers). Approximately 144 kernels." - }, - "test_params": { - "problem_sizes": [ - {"m": 512, "n": 1024, "k": 512, "split_k": 1}, - {"m": 1024, "n": 1024, "k": 256, "split_k": 2}, - {"m": 1024, "n": 1024, "k": 256, "split_k": 4} - ] - }, - "tile_config": { - "tile_m": {"values": [32, 64, 256]}, - "tile_n": {"values": [32, 64, 256]}, - "tile_k": {"values": [16, 32]}, - "warp_m": {"values": [2]}, - "warp_n": {"values": [2]}, - "warp_k": {"values": [1]}, - "warp_tile_m": {"values": [16]}, - "warp_tile_n": {"values": [16]}, - "warp_tile_k": {"values": [16]} - }, - "trait_config": { - "pipeline": {"values": ["mem", "compv3", "compv4"]}, - "epilogue": {"values": ["default", "cshuffle"]}, - "scheduler": {"values": ["intrawave", "interwave"]}, - "pad_m": {"values": [false]}, - "pad_n": {"values": [false]}, - "pad_k": {"values": [false]}, - "persistent": {"values": [false]} - }, - "k_block_per_cu": 1, - "permute_n": false -} diff --git a/projects/composablekernel/test/ck_tile/gemm_tile_engine/configs/simple_test_config.json b/projects/composablekernel/test/ck_tile/gemm_tile_engine/configs/simple_test_config.json deleted file mode 100644 index 498ef9fa33a1..000000000000 --- a/projects/composablekernel/test/ck_tile/gemm_tile_engine/configs/simple_test_config.json +++ /dev/null @@ -1,34 +0,0 @@ -{ - "problem": { - "description": "Basic functionality validation with moderate problem sizes" - }, - "test_params": { - "problem_sizes": [ - {"m": 256, "n": 256, "k": 128, "split_k": 1}, - {"m": 512, "n": 256, "k": 256, "split_k": 1}, - {"m": 256, "n": 512, "k": 256, "split_k": 1} - ] - }, - "tile_config": { - "tile_m": {"values": [128]}, - "tile_n": {"values": [128]}, - "tile_k": {"values": [64]}, - "warp_m": {"values": [2]}, - "warp_n": {"values": [2]}, - "warp_k": {"values": [1]}, - "warp_tile_m": {"values": [16]}, - "warp_tile_n": {"values": [16]}, - "warp_tile_k": {"values": [16]} - }, - "trait_config": { - "pipeline": {"values": ["compv3", "compv4"]}, - "epilogue": {"values": ["default"]}, - "scheduler": {"values": ["intrawave"]}, - "pad_m": {"values": [false]}, - "pad_n": {"values": [false]}, - "pad_k": {"values": [false]}, - "persistent": {"values": [false]} - }, - "k_block_per_cu": 1, - "permute_n": false -} diff --git a/projects/composablekernel/test/ck_tile/gemm_tile_engine/configs/small_datatype_config.json b/projects/composablekernel/test/ck_tile/gemm_tile_engine/configs/small_datatype_config.json deleted file mode 100644 index d0d9f99a0cc7..000000000000 --- a/projects/composablekernel/test/ck_tile/gemm_tile_engine/configs/small_datatype_config.json +++ /dev/null @@ -1,35 +0,0 @@ -{ - "problem": { - "description": "Configuration optimized for small data types (fp8, fp16, bf16) with larger warp tiles" - }, - "test_params": { - "problem_sizes": [ - {"m": 512, "n": 512, "k": 256, "split_k": 1}, - {"m": 1024, "n": 512, "k": 512, "split_k": 1}, - {"m": 512, "n": 1024, "k": 512, "split_k": 1}, - {"m": 1024, "n": 1024, "k": 256, "split_k": 1} - ] - }, - "tile_config": { - "tile_m": {"values": [128]}, - "tile_n": {"values": [128]}, - "tile_k": {"values": [32]}, - "warp_m": {"values": [2]}, - "warp_n": {"values": [2]}, - "warp_k": {"values": [1]}, - "warp_tile_m": {"values": [32]}, - "warp_tile_n": {"values": [32]}, - "warp_tile_k": {"values": [16]} - }, - "trait_config": { - "pipeline": {"values": ["compv3"]}, - "epilogue": {"values": ["default"]}, - "scheduler": {"values": ["intrawave"]}, - "pad_m": {"values": [false]}, - "pad_n": {"values": [false]}, - "pad_k": {"values": [false]}, - "persistent": {"values": [false]} - }, - "k_block_per_cu": 1, - "permute_n": false -} diff --git a/projects/composablekernel/test/ck_tile/gemm_tile_engine/extract_test_params.py b/projects/composablekernel/test/ck_tile/gemm_tile_engine/extract_test_params.py deleted file mode 100644 index 48ec8dba8352..000000000000 --- a/projects/composablekernel/test/ck_tile/gemm_tile_engine/extract_test_params.py +++ /dev/null @@ -1,74 +0,0 @@ -#!/usr/bin/env python3 -# Copyright (c) Advanced Micro Devices, Inc., or its affiliates. -# SPDX-License-Identifier: MIT - - -import json -import argparse -import os -from pathlib import Path - - -def extract_test_params(config_file, output_file): - """Extract test parameters from config JSON and write to output file""" - - # Read config file - with open(config_file, "r") as f: - config = json.load(f) - - # Extract test parameters - test_params = [] - if "test_params" in config and "problem_sizes" in config["test_params"]: - test_params = config["test_params"]["problem_sizes"] - else: - # Default test parameters if none specified - test_params = [ - {"m": 256, "n": 256, "k": 128, "split_k": 1}, - {"m": 256, "n": 256, "k": 1024, "split_k": 1}, - {"m": 256, "n": 512, "k": 512, "split_k": 1}, - {"m": 512, "n": 256, "k": 512, "split_k": 1}, - ] - - # Write to output file in C++ format - output_dir = Path(output_file).parent - output_dir.mkdir(parents=True, exist_ok=True) - - with open(output_file, "w") as f: - f.write("// Generated test parameters for this configuration\n") - f.write("// This file is auto-generated during CMake configuration\n\n") - f.write("static const std::vector CONFIG_TEST_PARAMS = {\n") - - for i, params in enumerate(test_params): - comma = "," if i < len(test_params) - 1 else "" - f.write( - f" {{{params['m']}, {params['n']}, {params['k']}, {params['split_k']}}}{comma}\n" - ) - - f.write("};\n") - - print( - f"Extracted {len(test_params)} test parameters from {config_file} -> {output_file}" - ) - - -def main(): - parser = argparse.ArgumentParser( - description="Extract test parameters from config JSON" - ) - parser.add_argument("--config_file", required=True, help="Input config JSON file") - parser.add_argument( - "--output_file", required=True, help="Output test parameters file" - ) - - args = parser.parse_args() - - if not os.path.exists(args.config_file): - print(f"Error: Config file not found: {args.config_file}") - return 1 - - extract_test_params(args.config_file, args.output_file) - return 0 - - -if __name__ == "__main__": - exit(main()) diff --git a/projects/composablekernel/test/ck_tile/gemm_tile_engine/test_gemm_simple.cpp b/projects/composablekernel/test/ck_tile/gemm_tile_engine/test_gemm_simple.cpp deleted file mode 100644 index e44e8c4182ac..000000000000 --- a/projects/composablekernel/test/ck_tile/gemm_tile_engine/test_gemm_simple.cpp +++ /dev/null @@ -1,241 +0,0 @@ -// Copyright (c) Advanced Micro Devices, Inc., or its affiliates. -// SPDX-License-Identifier: MIT - -/** - * @file test_gemm_simple.cpp - * @brief Unit tests for GEMM kernels generated by gemm_instance_builder - * - * This test includes kernels generated during CMake configuration by - * gemm_instance_builder.py and tests them with problem sizes extracted - * from the corresponding JSON configuration files. - */ - -#include -#include - -#include "ck_tile/core.hpp" -#include "ck_tile/host.hpp" -#include "tile_engine/ops/gemm/gemm_common.hpp" - -// The kernel header is included via compile command line with -include flag -// It defines SelectedKernel struct, KERNEL_NAME, and tensor data types - -// Adaptive error threshold calculation matching tile_engine's implementation -template -auto calculate_rtol_atol(const ck_tile::index_t K, - const ck_tile::index_t kbatch, - const float max_accumulated_value) -{ - using ComputeType = - std::conditional_t; - // Calculate thresholds - const auto rtol = ck_tile::get_relative_threshold( - ck_tile::integer_divide_ceil(K, kbatch)); - const auto atol = ck_tile::get_absolute_threshold( - max_accumulated_value / kbatch, ck_tile::integer_divide_ceil(K, kbatch)); - // Calculate error due to split_k accumulation - const auto rtol_split_k = - ck_tile::get_relative_threshold(kbatch); - const auto atol_split_k = ck_tile::get_absolute_threshold( - max_accumulated_value, kbatch); - // Use higher threshold - return ck_tile::make_tuple(std::max(rtol, rtol_split_k), std::max(atol, atol_split_k)); -} - -/// @brief Function to compare the results of the device and host computations (from tile_engine) -template -bool compare_results(std::string instanceName, - ck_tile::index_t K, - ck_tile::index_t kbatch, - ck_tile::HostTensor& c_m_n_dev_result, - ck_tile::HostTensor& c_m_n_host_result) -{ - const float max_accumulated_value = - *std::max_element(c_m_n_host_result.mData.begin(), c_m_n_host_result.mData.end()); - const auto rtol_atol = calculate_rtol_atol( - K, kbatch, max_accumulated_value); - bool pass = ck_tile::check_err(c_m_n_dev_result, - c_m_n_host_result, - "Error: Incorrect results!", - rtol_atol.at(ck_tile::number<0>{}), - rtol_atol.at(ck_tile::number<1>{})); - - std::cout << "For " << instanceName << " Relative error threshold is " - << rtol_atol.at(ck_tile::number<0>{}) << " Absolute error threshold is " - << rtol_atol.at(ck_tile::number<1>{}) << std::endl; - std::cout << "The verification result is:" << (pass ? "correct" : "fail") << std::endl; - - return pass; -} - -// Test parameter structure for matrix dimensions and split_k values -struct GemmTestParams -{ - int m, n, k, split_k; -}; - -// Include config-specific test parameters (after GemmTestParams struct is defined) -#ifdef GEMM_TEST_PARAMS_HPP -#include GEMM_TEST_PARAMS_HPP -#endif - -class GemmTileEngineTest : public ::testing::TestWithParam -{ - protected: - void SetUp() override - { - auto params = GetParam(); - m_ = params.m; - n_ = params.n; - k_ = params.k; - split_k_ = params.split_k; - - // Calculate strides (following tile_engine pattern) - if constexpr(std::is_same_v) - { - stride_a_ = k_; - } - else - { - stride_a_ = m_; - } - - if constexpr(std::is_same_v) - { - stride_b_ = n_; - } - else - { - stride_b_ = k_; - } - - if constexpr(std::is_same_v) - { - stride_c_ = n_; - } - else - { - stride_c_ = m_; - } - } - - // Test dimensions - int m_, n_, k_, split_k_; - int stride_a_, stride_b_, stride_c_; -}; - -TEST_P(GemmTileEngineTest, BasicFunctionality) -{ - // Get tensor layouts from generated kernel - const ALayout layout_a = ALayout{}; - const BLayout layout_b = BLayout{}; - const CLayout layout_c = CLayout{}; - - // Use split_k from test parameters - int split_k = split_k_; - int stride_a_calc = ck_tile::get_default_stride(m_, k_, 0, is_row_major(layout_a)); - int stride_b_calc = ck_tile::get_default_stride(k_, n_, 0, is_row_major(layout_b)); - int stride_c_calc = ck_tile::get_default_stride(m_, n_, 0, is_row_major(layout_c)); - - // Create host tensors with proper descriptors - ck_tile::HostTensor a_m_k( - ck_tile::host_tensor_descriptor(m_, k_, stride_a_calc, is_row_major(layout_a))); - ck_tile::HostTensor b_k_n( - ck_tile::host_tensor_descriptor(k_, n_, stride_b_calc, is_row_major(layout_b))); - ck_tile::HostTensor c_m_n_dev_result( - ck_tile::host_tensor_descriptor(m_, n_, stride_c_calc, is_row_major(layout_c))); - ck_tile::HostTensor c_m_n_host_result( - ck_tile::host_tensor_descriptor(m_, n_, stride_c_calc, is_row_major(layout_c))); - - // Initialize input tensors with uniform random distribution [-1.0, 1.0] (matches tile_engine) - ck_tile::FillUniformDistribution{-1.f, 1.f}(a_m_k); - ck_tile::FillUniformDistribution{-1.f, 1.f}(b_k_n); - - // Allocate GPU device memory - ck_tile::DeviceMem a_m_k_dev_buf(a_m_k.get_element_space_size_in_bytes()); - ck_tile::DeviceMem b_k_n_dev_buf(b_k_n.get_element_space_size_in_bytes()); - ck_tile::DeviceMem c_m_n_dev_buf(c_m_n_dev_result.get_element_space_size_in_bytes()); - - // Copy data to device and zero output buffer - a_m_k_dev_buf.ToDevice(a_m_k.data()); - b_k_n_dev_buf.ToDevice(b_k_n.data()); - c_m_n_dev_buf.SetZero(); - c_m_n_dev_result.SetZero(); - - // Calculate reference result on host for verification - ck_tile::reference_gemm( - a_m_k, b_k_n, c_m_n_host_result); - - // Create GEMM kernel arguments - ck_tile::GemmHostArgs gemm_args(a_m_k_dev_buf.GetDeviceBuffer(), - b_k_n_dev_buf.GetDeviceBuffer(), - c_m_n_dev_buf.GetDeviceBuffer(), - split_k, - m_, - n_, - k_, - stride_a_calc, - stride_b_calc, - stride_c_calc); - - // Configure kernel execution for maximum speed (no timing, no debug output) - ck_tile::stream_config stream_config{nullptr, // stream - false, // time_kernel (disable timing for speed) - 0, // log_level (disable debug output) - 0, // n_warmup - 1, // n_repeat - false, // is_gpu_timer (unused when time_kernel=false) - false, // flush_cache - 1}; // rotating_count - - // Launch the generated kernel (no timing overhead for fastest execution) - try - { - SelectedKernel::launch(gemm_args, stream_config); - // Kernel launched successfully if no exception thrown - } - catch(const std::exception& e) - { - std::string error_msg(e.what()); - // If arguments not supported, skip the test (configuration validation failure, not a bug) - if(error_msg.find("Arguments not supported") != std::string::npos) - { - GTEST_SKIP() << "Configuration not supported: " << e.what(); - } - else - { - FAIL() << "Kernel launch failed: " << e.what(); - } - } - - // Copy result back from device - c_m_n_dev_buf.FromDevice(c_m_n_dev_result.data()); - - // Verify results using tile_engine's adaptive error thresholds - bool verification_passed = compare_results( - KERNEL_NAME, k_, split_k, c_m_n_dev_result, c_m_n_host_result); - - EXPECT_TRUE(verification_passed) << "GEMM result verification failed"; -} - -TEST_P(GemmTileEngineTest, KernelInfo) -{ - // Simple test to verify kernel information is available - EXPECT_TRUE(strlen(KERNEL_NAME) > 0) << "Kernel name should not be empty"; - - std::cout << "Testing kernel: " << KERNEL_NAME << std::endl; - std::cout << "Problem size: " << m_ << "x" << n_ << "x" << k_ << " with split_k=" << split_k_ - << std::endl; -} - -// Use config-specific test parameters (included via compile flags) -// CONFIG_TEST_PARAMS is defined in the auto-generated test_params.hpp file -INSTANTIATE_TEST_SUITE_P(GemmVerification, - GemmTileEngineTest, - ::testing::ValuesIn(CONFIG_TEST_PARAMS), - [](const ::testing::TestParamInfo& param_info) { - return std::to_string(param_info.param.m) + "x" + - std::to_string(param_info.param.n) + "x" + - std::to_string(param_info.param.k) + "_splitk" + - std::to_string(param_info.param.split_k); - }); diff --git a/projects/composablekernel/tile_engine/ops/gemm/CMakeLists.txt b/projects/composablekernel/tile_engine/ops/gemm/CMakeLists.txt index b50a6790105a..c7f6e48930ad 100644 --- a/projects/composablekernel/tile_engine/ops/gemm/CMakeLists.txt +++ b/projects/composablekernel/tile_engine/ops/gemm/CMakeLists.txt @@ -15,7 +15,7 @@ if(NOT "${TILE_ENGINE_SAMPLING_TIER}" STREQUAL "") if(_te_budget GREATER 0) # Detect active ops from their DATATYPE variables set(_active_ops "") - foreach(_op gemm_universal gemm_multi_d gemm_preshuffle grouped_gemm gemm_streamk batched_contraction batched_gemm gemm_multi_abd mx_gemm gemm_rowcolquant gemm_tensor_quant grouped_gemm_rowcolquant grouped_gemm_tensorquant) + foreach(_op gemm_multi_d gemm_preshuffle grouped_gemm gemm_streamk batched_contraction batched_gemm gemm_multi_abd mx_gemm gemm_rowcolquant gemm_tensor_quant grouped_gemm_rowcolquant grouped_gemm_tensorquant) string(TOUPPER ${_op} _OP_UPPER) if(NOT "${${_OP_UPPER}_DATATYPE}" STREQUAL "") list(APPEND _active_ops ${_op}) @@ -45,7 +45,7 @@ if(NOT "${TILE_ENGINE_SAMPLING_TIER}" STREQUAL "") message(STATUS "Sampling budget allocation:\n${_alloc_output}") # Read per-op allocations (only if not already overridden) - foreach(_op gemm_universal gemm_multi_d gemm_preshuffle grouped_gemm gemm_streamk batched_contraction batched_gemm gemm_multi_abd mx_gemm gemm_rowcolquant gemm_tensor_quant grouped_gemm_rowcolquant grouped_gemm_tensorquant) + foreach(_op gemm_multi_d gemm_preshuffle grouped_gemm gemm_streamk batched_contraction batched_gemm gemm_multi_abd mx_gemm gemm_rowcolquant gemm_tensor_quant grouped_gemm_rowcolquant grouped_gemm_tensorquant) string(TOUPPER ${_op} _OP_UPPER) if("${${_OP_UPPER}_MAX_INSTANCES}" STREQUAL "") if(EXISTS "${_alloc_dir}/${_op}_budget.txt") @@ -73,7 +73,6 @@ if(NOT "${TILE_ENGINE_SAMPLING_TIER}" STREQUAL "") endif() endif() -add_subdirectory(gemm_universal EXCLUDE_FROM_ALL) add_subdirectory(gemm_multi_d EXCLUDE_FROM_ALL) add_subdirectory(gemm_preshuffle EXCLUDE_FROM_ALL) add_subdirectory(grouped_gemm EXCLUDE_FROM_ALL) diff --git a/projects/composablekernel/tile_engine/ops/gemm/README.md b/projects/composablekernel/tile_engine/ops/gemm/README.md index 5e0bae70806d..781f113f7afb 100644 --- a/projects/composablekernel/tile_engine/ops/gemm/README.md +++ b/projects/composablekernel/tile_engine/ops/gemm/README.md @@ -6,6 +6,7 @@ The CK Tile Engine GEMM module provides a comprehensive system for generating, b ## Table of Contents +0. [Dispatcher Bridge Workflow](#dispatcher-bridge-workflow) 1. [Build System Architecture](#build-system-architecture) 2. [Build Instructions](#build-instructions) 3. [Running Benchmarks](#running-benchmarks) @@ -16,6 +17,179 @@ The CK Tile Engine GEMM module provides a comprehensive system for generating, b 8. [Troubleshooting](#troubleshooting) 9. [Performance Tips](#performance-tips) +## Dispatcher Bridge Workflow + +The **Dispatcher bridge** is the recommended path for sweeping and benchmarking +GEMM kernels. Instead of building monolithic or per-kernel executables through +CMake, Tile Engine expands a sweep config into shared `GemmKernelConfig` objects +and hands them to the Dispatcher, which codegens and compiles each into its own +`.so`. The kernel name produced by the bridge is byte-for-byte identical to the +codegen `KERNEL_NAME`, so the bridge runs exactly the same kernels the native +Tile Engine does — it only swaps the harness. + +### Scripts + +| Script | Role | +|---|---| +| `gemm_full_benchmark.py` | Driver: compile (Phase 1) → load problems (Phase 2) → benchmark across all visible GPUs (Phase 3). | +| `run_one_gemm_kernel.py` | Disposable worker: loads one `.so` in an isolated subprocess and times it. A GPU fault kills only the worker. | + +### Folder layout + +The bridged regular-GEMM path follows the same op-root convention as the merged +`fmha/` and `grouped_conv/` bridges — driver + worker + a flat `configs/` at the +op root: + +``` +gemm/ +├── gemm_full_benchmark.py # bridge driver (op root) +├── run_one_gemm_kernel.py # disposable per-kernel worker (op root) +├── configs/ # bridged gemm_universal sweep configs (flat) +├── gemm_instance_builder.py # shared generator for the non-bridged variants +├── gemm_benchmark.{py,hpp}, gemm_common.hpp, gemm_profiler.hpp # shared harness +├── gemm_multi_d/ gemm_preshuffle/ grouped_gemm/ # legacy variants +└── README.md +``` + +`configs/` ships example sweep configs: + +- `default_ci_config.json` — small CI-sized sweep (the driver's default when no + config is passed). +- `default_config.json` — full sweep. +- `user_provided_config.json` — scratch space for custom sweeps. +- `example_problems.json` — example M/N/K problem set (used when `--problems` + is omitted). + +> The JSON used by **nightly** tests is intended to drop into the same +> `configs/` directory and be selected with a positional config — no driver +> changes needed. + +The not-yet-bridged variants (`gemm_multi_d/`, `gemm_preshuffle/`, +`grouped_gemm/`) keep their own per-variant `configs/` directories; the driver +selects them with `--variant`. + +### Running + +```bash +cd tile_engine/ops/gemm + +# Default: gemm_universal variant, its CI sweep + example problems, +# auto-detect and use all visible GPUs. +python gemm_full_benchmark.py + +# Full sweep, fp16/rcr, restricted to 4 GPUs, custom output: +python gemm_full_benchmark.py --variant gemm_universal \ + configs/default_config.json \ + --dtype fp16 --layout rcr --devices 4 --csv gemm_results.csv + +# Specific GPU ids and a custom problem file: +python gemm_full_benchmark.py --devices 0,2,5 \ + --problems configs/example_problems.json + +# Correctness mode: check every kernel against an fp32 numpy reference. +python gemm_full_benchmark.py --verify --max-kernels 8 +``` + +### Liveness vs correctness (`--verify`) + +By default a measurement is reported `OK` purely on **liveness** — the kernel +ran and produced a non-zero output (`ZERO` otherwise). It is *not* a correctness +check: a numerically wrong but non-zero result still reads `OK`. Pass `--verify` +to have each worker compare its output against an fp32 numpy reference +(`A @ B`) using the global relative metric `max|out - ref| / max|ref|`. With +`--verify`, results read `VERIFY` (within `--verify-tol`, default `2e-2`) or +`MISMATCH` (counted as a failure), and the `max_rel` / `verified` columns are +populated in the CSV. This gives self-contained per-kernel confidence; the +broader numeric parity against native Tile Engine remains a separate task. + +### Multi-GPU parallelism + +Phase 3 fans the `(kernel × problem)` work out across **every visible GPU** in +parallel. One worker thread per device pulls batches from a shared queue and +spawns a disposable subprocess pinned with `HIP_VISIBLE_DEVICES`, so an N-GPU box +benchmarks roughly N× faster while keeping per-batch fault isolation. Devices are +auto-detected (`HIP_VISIBLE_DEVICES`, then `rocm-smi`/`amd-smi`); override with +`--devices`. This supersedes the serial-GPU design inherited from grouped_conv. + +### Supported surface + +| Axis | Supported | +|---|---| +| dtype | `fp16` (bf16 follows in #8190) | +| layout | `rcr` (rrr/crr/ccr follow in #8191; row-major C only — ck_tile rejects column-major C at build) | + +### Variant scope + +The bridge is **one shared, variant-aware driver** (`gemm_full_benchmark.py` + +`run_one_gemm_kernel.py`), not a per-variant copy of the driver. The bridged +regular-GEMM path (`gemm_universal`) uses the op-root `configs/`; `--variant` +selects a not-yet-bridged variant's own `configs/` subdirectory. + +What that means for this PR: + +- **Only `gemm_universal` is wired and validated through the bridge here.** It is + the foundation variant; the dispatcher codegen path is exercised and parity- + checked for it alone. +- The `gemm_multi_d/`, `gemm_preshuffle/`, and `grouped_gemm/` `configs/` + directories are **scaffolding** that follows the per-variant convention so the + layout is ready. `--variant` will select them, but the bridge does **not** yet + produce correct kernels for those variants on this PR — do not treat their + presence as working support. +- Grouped GEMM and stream-K go through **separate bridge efforts** (stream-K in + #8136, grouped GEMM on its own branch), not this PR. + +### Stream-K bridge + +Stream-K reuses the same bridge mechanics through its **own** driver/worker pair +(it keeps a distinct ctypes lib that bypasses the registry; see +`dispatcher/bindings/ctypes/streamk_gemm_ctypes_lib.cpp`): + +| Script | Role | +|---|---| +| `streamk_gemm_full_benchmark.py` | Stream-K driver — same 3-phase, multi-GPU, `--verify`/`--devices` flags as the regular driver; threads `variant="stream_k"`. Defaults to `gemm_streamk/configs/default_config.json`. | +| `run_one_streamk_gemm_kernel.py` | Disposable per-kernel worker (same fp16 host path as the regular worker; the ABI matches). | + +```bash +# Default Stream-K config, all visible GPUs, with correctness checking: +python streamk_gemm_full_benchmark.py --verify + +# Explicit config, 4 GPUs, custom output: +python streamk_gemm_full_benchmark.py gemm_streamk/configs/default_config.json \ + --devices 4 --csv streamk_gemm_results.csv +``` + +Notes specific to Stream-K: +- Kernel names carry a `_streamk` suffix (`GemmKernelConfig(variant="stream_k").name`). +- Supported surface here is **fp16 / rcr**, same as the regular bridge. +- The Atomic reduction does multiple fp16 atomic-adds per K-split, so it is noisier + than regular GEMM; `--verify` still uses the default `2e-2` gate (observed + `max_rel ≤ 2.5e-3`, well within it). +- Benchmark knobs in `streamk_gemm_ctypes_lib.cpp` default to warmup=50/repeat=100 + (env-overridable via `CK_TILE_BENCH_WARMUP`/`REPEAT`/`FLUSH`/`ROTATING`), matching + the regular path — **except** `rotating_count` defaults to 1, because the Atomic + preprocess only re-zeros the original C buffer and rotating C would corrupt the + accumulation. +- Tiny problems (e.g. `257³`) have too few tiles to partition across CUs; the kernel + reports them unsupported (status `-2`) and the bridge surfaces that gracefully. + +### Removal note + +The legacy regular-GEMM standalone build path has been **removed**, and the +`gemm_universal/` folder is gone entirely. The per-config benchmark generator and +driver (`gemm_universal_instance_builder.py`, `gemm_universal_benchmark.py`, +`gemm_universal_benchmark*.{cpp,hpp}`, and `gemm_universal/CMakeLists.txt`) no +longer exist; its sweep configs were promoted to the op-root `configs/` directory +(matching the `fmha/` and `grouped_conv/` bridge convention) and are consumed by +the bridge. Regular GEMM now runs exclusively through the Dispatcher bridge +workflow above (`gemm_full_benchmark.py` / `run_one_gemm_kernel.py`). The other +variants (`gemm_multi_d/`, `gemm_preshuffle/`, `grouped_gemm/`) still use the +shared `gemm_instance_builder.py` generator. + +The build-system, build-instruction, and benchmark-execution sections below +describe that removed standalone path and are retained only as historical +reference for the non-bridged variants; the `benchmark_gemm_universal_*` targets +they mention are no longer produced. + ## Build System Architecture ### Individual Kernel Compilation (New Approach) @@ -171,8 +345,13 @@ The system uses JSON configuration files to specify kernel parameters: ### Python Scripts -#### gemm_universal_instance_builder.py -**Purpose**: Main kernel instance generation script that creates C++ kernel implementations based on configuration files. +#### gemm_instance_builder.py +**Purpose**: Shared kernel instance generator used by the non-bridged variants +(`gemm_multi_d`, `gemm_preshuffle`, `grouped_gemm`). Creates C++ kernel +implementations based on configuration files. + +> The regular-GEMM subclass `gemm_universal/gemm_universal_instance_builder.py` +> has been removed; regular GEMM now goes through the Dispatcher bridge. **Key Features**: - Generates individual kernel header files for separate compilation @@ -180,16 +359,6 @@ The system uses JSON configuration files to specify kernel parameters: - Validates tile configurations for correctness - Creates CMake integration files -**Usage**: -```bash -python gemm_universal_instance_builder.py \ - --working_path ./generated \ - --datatype fp16 \ - --layout rcr \ - --config_json configs/user_provided_config.json \ - --gen_all_individual -``` - #### gemm_instance_builder_parallel.py **Purpose**: Parallel version of the instance builder for faster generation of multiple kernel configurations. @@ -225,14 +394,6 @@ python test_validation.py - Trait combination validation - Full tile configuration validation -#### gemm_universal_benchmark.py -**Purpose**: Python script for running and analyzing GEMM benchmarks. - -**Features**: -- Automated benchmark execution -- Performance data collection -- Result analysis and reporting - #### json_config.py **Purpose**: Configuration file parsing and management. diff --git a/projects/composablekernel/tile_engine/ops/gemm/gemm_universal/configs/default_ci_config.json b/projects/composablekernel/tile_engine/ops/gemm/configs/default_ci_config.json similarity index 98% rename from projects/composablekernel/tile_engine/ops/gemm/gemm_universal/configs/default_ci_config.json rename to projects/composablekernel/tile_engine/ops/gemm/configs/default_ci_config.json index 38376a410b01..a2b83334245e 100644 --- a/projects/composablekernel/tile_engine/ops/gemm/gemm_universal/configs/default_ci_config.json +++ b/projects/composablekernel/tile_engine/ops/gemm/configs/default_ci_config.json @@ -32,17 +32,17 @@ }, "warp_tile_m": { "values": [ - 16 + 32 ] }, "warp_tile_n": { "values": [ - 16 + 32 ] }, "warp_tile_k": { "values": [ - 32 + 16 ] } }, diff --git a/projects/composablekernel/tile_engine/ops/gemm/gemm_universal/configs/default_config.json b/projects/composablekernel/tile_engine/ops/gemm/configs/default_config.json similarity index 100% rename from projects/composablekernel/tile_engine/ops/gemm/gemm_universal/configs/default_config.json rename to projects/composablekernel/tile_engine/ops/gemm/configs/default_config.json diff --git a/projects/composablekernel/tile_engine/ops/gemm/configs/example_problems.json b/projects/composablekernel/tile_engine/ops/gemm/configs/example_problems.json new file mode 100644 index 000000000000..4be0c5a82379 --- /dev/null +++ b/projects/composablekernel/tile_engine/ops/gemm/configs/example_problems.json @@ -0,0 +1,9 @@ +{ + "problems": [ + {"M": 512, "N": 512, "K": 512}, + {"M": 1024, "N": 1024, "K": 1024}, + {"M": 2048, "N": 2048, "K": 2048}, + {"M": 1024, "N": 512, "K": 256}, + {"M": 4096, "N": 4096, "K": 4096} + ] +} diff --git a/projects/composablekernel/tile_engine/ops/gemm/gemm_universal/configs/user_provided_config.json b/projects/composablekernel/tile_engine/ops/gemm/configs/user_provided_config.json similarity index 100% rename from projects/composablekernel/tile_engine/ops/gemm/gemm_universal/configs/user_provided_config.json rename to projects/composablekernel/tile_engine/ops/gemm/configs/user_provided_config.json diff --git a/projects/composablekernel/tile_engine/ops/gemm/gemm_full_benchmark.py b/projects/composablekernel/tile_engine/ops/gemm/gemm_full_benchmark.py new file mode 100644 index 000000000000..43201a8cbabb --- /dev/null +++ b/projects/composablekernel/tile_engine/ops/gemm/gemm_full_benchmark.py @@ -0,0 +1,506 @@ +#!/usr/bin/env python3 +"""Full GEMM benchmark sweep driven through the Dispatcher bridge. + +Phases: + Phase 1: Compile all kernels (parallel, returns .so paths only -- no GPU) + Phase 2: Load problems (M, N, K shapes) + Phase 3: Benchmark via subprocess isolation, distributed across all visible + GPUs (one device-pinned worker per GPU, batched, fault-isolated) + +Tile Engine generates NO binaries here: it expands its sweep config into shared +``GemmKernelConfig`` objects and hands them to the dispatcher, which codegens + +compiles each into a .so. Each kernel runs in a disposable worker subprocess so +a GPU fault (or ctypes' inability to unload a .so) takes down only one worker. + +Unlike the serial-GPU design inherited from grouped_conv, Phase 3 here fans the +work out across every visible GPU in parallel: each device runs its own stream of +disposable worker subprocesses pinned with ``HIP_VISIBLE_DEVICES``, so an N-GPU +box benchmarks roughly N times faster while keeping per-batch fault isolation. + +Examples: + # Default: gemm_universal variant, its CI sweep config + example problems, + # auto-detect and use all visible GPUs. + python gemm_full_benchmark.py + + # Explicit variant + full sweep config on 4 GPUs: + python gemm_full_benchmark.py --variant gemm_universal \ + configs/default_config.json --devices 4 --csv out.csv + +When no config is given the driver uses the chosen variant's +``configs/default_ci_config.json`` (a small CI-sized sweep); +``configs/default_config.json`` is the full sweep, and the JSON used by nightly +tests is intended to drop into the same ``configs/`` directory. +""" + +import argparse +import csv +import json +import os +import queue +import re +import subprocess +import sys +import threading +import time +from pathlib import Path + +_THIS_DIR = Path(__file__).resolve().parent +_DISPATCHER_ROOT = _THIS_DIR.parents[2] / "dispatcher" +sys.path.insert(0, str(_DISPATCHER_ROOT / "python")) +sys.path.insert(0, str(_THIS_DIR)) + +from gemm_utils import setup_multiple_gemm_dispatchers, expand_sweep # noqa: E402 + +# Config layout. The bridged regular-GEMM path (gemm_universal) keeps its sweep +# configs in this op's flat ``configs/`` directory (matching the fmha/grouped_conv +# bridge convention): default_ci_config.json (small CI sweep), default_config.json +# (full sweep), user_provided_config.json, example_problems.json. The other, +# not-yet-bridged variants still live in their own per-variant ``configs/`` dirs; +# they are registered so ``--variant`` can select them once their bridge lands. +VARIANT_CONFIGS = { + "gemm_universal": "configs", + "gemm_multi_d": "gemm_multi_d/configs", + "gemm_preshuffle": "gemm_preshuffle/configs", + "grouped_gemm": "grouped_gemm/configs", +} +DEFAULT_VARIANT = "gemm_universal" +CI_CONFIG_NAME = "default_ci_config.json" +EXAMPLE_PROBLEMS_NAME = "example_problems.json" + +# Fallback problem set if a variant ships no example_problems.json. +DEFAULT_PROBLEMS = [ + {"M": 1024, "N": 1024, "K": 1024}, + {"M": 2048, "N": 2048, "K": 2048}, + {"M": 4096, "N": 4096, "K": 4096}, + {"M": 257, "N": 257, "K": 257}, +] + +# Foundation bridge surface: fp16/rcr only. bf16 and the rrr/crr/ccr layouts +# land in the follow-up stack (#8190 bf16, #8191 layouts) once the dispatcher +# host path supports them. +SUPPORTED_DTYPES = ("fp16",) +SUPPORTED_LAYOUTS = ("rcr",) + + +def detect_devices(): + """Return a list of visible GPU id strings (best-effort).""" + env = os.environ.get("HIP_VISIBLE_DEVICES") or os.environ.get( + "CUDA_VISIBLE_DEVICES" + ) + if env: + ids = [d.strip() for d in env.split(",") if d.strip() != ""] + if ids: + return ids + try: + out = subprocess.check_output( + ["rocm-smi", "--showid"], stderr=subprocess.DEVNULL, text=True + ) + ids = sorted(set(re.findall(r"GPU\[(\d+)\]", out)), key=int) + if ids: + return ids + except Exception: + pass + try: + out = subprocess.check_output( + ["amd-smi", "list"], stderr=subprocess.DEVNULL, text=True + ) + ids = re.findall(r"^GPU:\s*(\d+)", out, re.MULTILINE) + if ids: + return ids + except Exception: + pass + return ["0"] + + +def resolve_devices(spec): + """Resolve --devices into a concrete list of device id strings. + + spec is None (auto: all visible), an int count, or a comma-list of ids. + A bare digit is a *count*, not an id; to target one specific id use the + comma form, e.g. "5,". + """ + detected = detect_devices() + if spec is None: + return detected + spec = str(spec).strip() + if "," in spec: + return [s.strip() for s in spec.split(",") if s.strip() != ""] + if spec.isdigit(): + n = int(spec) + if n <= 0: + return detected + # Treat a bare integer as a device *count*: take the first n detected + # ids, falling back to a plain 0..n-1 range if detection under-reports. + # To target one specific device id, use the comma form (e.g. "5,"). + return detected[:n] if len(detected) >= n else [str(i) for i in range(n)] + return [spec] + + +def resolve_configs(args): + """Resolve positional configs -> concrete list of config paths.""" + if args.configs: + return args.configs + cfg = _THIS_DIR / VARIANT_CONFIGS[args.variant] / CI_CONFIG_NAME + return [str(cfg)] + + +def load_problems(path, variant): + if path: + with open(path) as f: + data = json.load(f) + return data["problems"] if isinstance(data, dict) else data + example = _THIS_DIR / VARIANT_CONFIGS[variant] / EXAMPLE_PROBLEMS_NAME + if example.exists(): + with open(example) as f: + data = json.load(f) + return data["problems"] if isinstance(data, dict) else data + return DEFAULT_PROBLEMS + + +def _run_batch_on_device(device_id, unit, args, worker_path, base_env): + """Run one (problem, kernel-batch) unit in a device-pinned subprocess. + + Returns (rows, lines, n_fail) where rows are dicts ready for the CSV writer, + lines are formatted strings to print, and n_fail counts failures. + """ + 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 + ] + payload = json.dumps( + {"items": items, "verify": args.verify, "verify_tol": args.verify_tol} + ) + + env = base_env.copy() + env["HIP_VISIBLE_DEVICES"] = str(device_id) + + rows, lines, n_fail = [], [], 0 + proc = None + try: + proc = subprocess.Popen( + [sys.executable, str(worker_path)], + stdin=subprocess.PIPE, + stdout=subprocess.PIPE, + stderr=subprocess.DEVNULL, + env=env, + ) + stdout_bytes, _ = proc.communicate( + input=payload.encode("utf-8"), + timeout=args.kernel_timeout * len(batch), + ) + + reported = set() + for line in stdout_bytes.decode("utf-8").strip().split("\n"): + if not line: + continue + try: + result = json.loads(line) + except json.JSONDecodeError: + lines.append(f" [gpu{device_id}] Warning: bad result line: {line[:50]}") + n_fail += 1 + continue + bidx = result.get("idx", 0) + _, cfg, _ = batch[bidx] + reported.add(bidx) + if result.get("ok", False): + status = "OK" if result.get("non_zero", 0) > 0 else "ZERO" + mismatch = False + if args.verify and "verified" in result: + if result["verified"]: + status = "VERIFY" + else: + status = "MISMATCH" + mismatch = True + extra = ( + f" rel={result['max_rel']:.2e}" if "max_rel" in result else "" + ) + lines.append( + f" [gpu{device_id}] {cfg.name:<58} {result['ms']:>10.3f} " + f"{result['tflops']:>10.2f} {status:>8}{extra}" + ) + rows.append( + { + "kernel": cfg.name, + "problem_idx": prob_idx, + "M": M, + "N": N, + "K": K, + "device": device_id, + "latency_ms": result["ms"], + "tflops": result["tflops"], + "non_zero": result.get("non_zero", 0), + "max_rel": result.get("max_rel", ""), + "verified": result.get("verified", ""), + } + ) + if mismatch: + n_fail += 1 + else: + lines.append(f" [gpu{device_id}] {cfg.name:<58} FAILED") + lines.append(f" Error: {result.get('error', 'unknown')[:100]}") + n_fail += 1 + + missing = set(range(len(batch))) - reported + if missing or proc.returncode != 0: + if proc.returncode != 0: + lines.append(f" [gpu{device_id}] worker exited code {proc.returncode}") + for idx in sorted(missing): + _, cfg, _ = batch[idx] + lines.append(f" [gpu{device_id}] {cfg.name:<58} MISSING (crash)") + n_fail += len(missing) + + except subprocess.TimeoutExpired: + lines.append(f" [gpu{device_id}] batch timeout ({len(batch)} kernels)") + try: + proc.kill() + proc.communicate(timeout=5) + except Exception: + pass + n_fail += len(batch) + except Exception as e: + lines.append(f" [gpu{device_id}] batch error: {e}") + try: + if proc and proc.poll() is None: + proc.kill() + except Exception: + pass + n_fail += len(batch) + + return rows, lines, n_fail + + +def main(): + parser = argparse.ArgumentParser(description="GEMM Benchmark Sweep (via Dispatcher)") + parser.add_argument( + "configs", + nargs="*", + help="TE sweep config JSON files (default: variant's default_ci_config.json)", + ) + parser.add_argument( + "--variant", + default=DEFAULT_VARIANT, + choices=tuple(VARIANT_CONFIGS), + help="GEMM variant (selects the configs/ directory)", + ) + parser.add_argument("--arch", default="gfx942") + parser.add_argument( + "--dtype", + default="fp16", + choices=SUPPORTED_DTYPES, + help=f"Input dtype (supported: {', '.join(SUPPORTED_DTYPES)})", + ) + parser.add_argument( + "--layout", + default="rcr", + choices=SUPPORTED_LAYOUTS, + help=f"A/B/C layout (supported: {', '.join(SUPPORTED_LAYOUTS)})", + ) + parser.add_argument("--problems", default=None, help="JSON file of 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") + parser.add_argument( + "--devices", + default=None, + help="GPUs to use: int count (e.g. 4) or comma-list of ids (e.g. 0,2,5); " + "for one specific id use the comma form (e.g. 5,) since a bare digit is " + "a count; default auto-detects all visible", + ) + parser.add_argument( + "--batch-size", + type=int, + default=20, + help="Kernels per subprocess (overhead vs fault isolation)", + ) + parser.add_argument( + "--kernel-timeout", type=int, default=30, help="Per-kernel timeout (s)" + ) + parser.add_argument( + "--max-kernels", type=int, default=0, help="Limit to first N kernels (0=all)" + ) + parser.add_argument( + "--verify", + action="store_true", + help="Check each kernel's output against an fp32 numpy reference " + "(global max|out-ref|/max|ref|); a mismatch counts as a failure", + ) + parser.add_argument( + "--verify-tol", + type=float, + default=2e-2, + help="Relative tolerance for --verify (default 2e-2, suits fp16)", + ) + args = parser.parse_args() + + config_paths = resolve_configs(args) + devices = resolve_devices(args.devices) + + # ======================================================================== + # Phase 1: Compile kernels (parallel, no GPU) + # ======================================================================== + print(f"\n{'=' * 80}") + print("Phase 1: Compile kernels") + print(f"{'=' * 80}") + print(f" Variant: {args.variant}") + print(f" Configs: {', '.join(config_paths)}") + + all_configs = [] + for cfg_path in config_paths: + all_configs.extend( + expand_sweep(cfg_path, args.arch, dtype=args.dtype, layout=args.layout) + ) + + if args.max_kernels > 0: + all_configs = all_configs[: args.max_kernels] + + print(f" Expanded configs: {len(all_configs)}") + print(f" Build workers: {args.workers}") + + t0 = time.perf_counter() + # CRITICAL: returns Path objects only, does NOT load any .so. + lib_paths = setup_multiple_gemm_dispatchers( + all_configs, verbose=True, max_workers=args.workers + ) + build_time = time.perf_counter() - t0 + + built_kernels = [ + (cfg, lib) for cfg, lib in zip(all_configs, lib_paths) if lib is not None + ] + + # Dedupe by .so path (distinct configs can map to the same physical kernel). + seen_libs = set() + unique_kernels = [] + duplicate_count = 0 + for cfg, lib in built_kernels: + lib_key = str(lib.resolve()) + if lib_key not in seen_libs: + seen_libs.add(lib_key) + unique_kernels.append((cfg, lib)) + else: + duplicate_count += 1 + built_kernels = unique_kernels + + print( + f"\n Built {len(all_configs)} configs -> {len(built_kernels)} unique kernels " + f"({duplicate_count} duplicates filtered) in {build_time:.0f}s" + ) + + if not built_kernels: + print(" ERROR: No kernels built successfully") + return 1 + + # ======================================================================== + # Phase 2: Load problems + # ======================================================================== + print(f"\n{'=' * 80}") + print("Phase 2: Load test problems") + print(f"{'=' * 80}") + + problems = load_problems(args.problems, args.variant) + print(f" Problems: {len(problems)}") + print( + f" Total measurements: {len(built_kernels)} x {len(problems)} = " + f"{len(built_kernels) * len(problems)}" + ) + + # ======================================================================== + # Phase 3: Benchmark across all visible GPUs (subprocess isolation, batched) + # ======================================================================== + print(f"\n{'=' * 80}") + print("Phase 3: Benchmark (multi-GPU, subprocess isolation, batched)") + print(f"{'=' * 80}") + print(f" Devices: {len(devices)} -> {', '.join(devices)}") + print(f" Batch size: {args.batch_size} kernels per subprocess") + print(f" Timeout: {args.kernel_timeout}s per kernel\n") + + csv_path = Path(args.csv) + csv_fields = [ + "kernel", + "problem_idx", + "M", + "N", + "K", + "device", + "latency_ms", + "tflops", + "non_zero", + "max_rel", + "verified", + ] + csv_file = open(csv_path, "w", newline="") + writer = csv.DictWriter(csv_file, fieldnames=csv_fields) + writer.writeheader() + + worker_path = _THIS_DIR / "run_one_gemm_kernel.py" + base_env = os.environ.copy() + base_env["GEMM_PYPATH"] = os.pathsep.join( + [str(_DISPATCHER_ROOT / "python"), str(_THIS_DIR)] + ) + + # Build a single work queue of (prob_idx, prob_dict, kernel-batch) units and + # fan them out across device-pinned worker threads. + work_q = queue.Queue() + for prob_idx, prob in enumerate(problems): + prob_dict = {"M": int(prob["M"]), "N": int(prob["N"]), "K": int(prob["K"])} + for start in range(0, len(built_kernels), args.batch_size): + end = min(start + args.batch_size, len(built_kernels)) + batch = [ + (start + j, cfg, lib) + for j, (cfg, lib) in enumerate(built_kernels[start:end]) + ] + work_q.put((prob_idx, prob_dict, batch)) + + io_lock = threading.Lock() + stats = {"measurements": 0, "failures": 0} + bench_t0 = time.perf_counter() + + def device_thread(device_id): + while True: + try: + unit = work_q.get_nowait() + except queue.Empty: + return + rows, lines, n_fail = _run_batch_on_device( + device_id, unit, args, worker_path, base_env + ) + with io_lock: + for ln in lines: + print(ln) + for row in rows: + writer.writerow(row) + csv_file.flush() + stats["measurements"] += len(rows) + stats["failures"] += n_fail + work_q.task_done() + + threads = [ + threading.Thread(target=device_thread, args=(d,), daemon=True) for d in devices + ] + for t in threads: + t.start() + for t in threads: + t.join() + + bench_time = time.perf_counter() - bench_t0 + csv_file.close() + + # ======================================================================== + # Summary + # ======================================================================== + print(f"\n{'=' * 80}") + print("BENCHMARK COMPLETE") + print(f"{'=' * 80}") + print(f" Build time: {build_time:.0f}s") + print(f" Benchmark time: {bench_time:.0f}s") + print(f" Total time: {build_time + bench_time:.0f}s") + print(f" Devices used: {len(devices)}") + print(f" Successful measurements: {stats['measurements']}") + print(f" Failed measurements: {stats['failures']}") + print(f" Output: {csv_path}") + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/projects/composablekernel/tile_engine/ops/gemm/gemm_streamk/configs/default_config.json b/projects/composablekernel/tile_engine/ops/gemm/gemm_streamk/configs/default_config.json new file mode 100644 index 000000000000..59f9a888a4d6 --- /dev/null +++ b/projects/composablekernel/tile_engine/ops/gemm/gemm_streamk/configs/default_config.json @@ -0,0 +1,89 @@ +{ + "tile_config": { + "tile_m": { + "values": [ + 128 + ] + }, + "tile_n": { + "values": [ + 128 + ] + }, + "tile_k": { + "values": [ + 32, + 64 + ] + }, + "warp_m": { + "values": [ + 2 + ] + }, + "warp_n": { + "values": [ + 2 + ] + }, + "warp_k": { + "values": [ + 1 + ] + }, + "warp_tile_m": { + "values": [ + 32 + ] + }, + "warp_tile_n": { + "values": [ + 32 + ] + }, + "warp_tile_k": { + "values": [ + 16 + ] + } + }, + "trait_config": { + "pipeline": { + "values": [ + "compv3", + "compv4" + ] + }, + "scheduler": { + "values": [ + "intrawave" + ] + }, + "epilogue": { + "values": [ + "cshuffle" + ] + }, + "pad_m": { + "values": [ + true + ] + }, + "pad_n": { + "values": [ + true + ] + }, + "pad_k": { + "values": [ + true + ] + }, + "persistent": { + "values": [ + false + ] + } + }, + "k_block_per_cu": 1 +} diff --git a/projects/composablekernel/tile_engine/ops/gemm/gemm_universal/CMakeLists.txt b/projects/composablekernel/tile_engine/ops/gemm/gemm_universal/CMakeLists.txt deleted file mode 100644 index e0624b7067b2..000000000000 --- a/projects/composablekernel/tile_engine/ops/gemm/gemm_universal/CMakeLists.txt +++ /dev/null @@ -1,338 +0,0 @@ -# Copyright (c) Advanced Micro Devices, Inc., or its affiliates. -# SPDX-License-Identifier: MIT - -set(GEMM_UNIVERSAL_DATATYPE "fp8;fp16" CACHE STRING "List of datatypes for GEMM Universal (semicolon-separated)") -set(GEMM_UNIVERSAL_LAYOUT "rcr;rrr;crr;ccr" CACHE STRING "List of layout for GEMM Universal (semicolon-separated)") -set(GEMM_UNIVERSAL_CONFIG_FILE "" CACHE STRING "Custom config file name (without path, must be in configs/ folder)") -set(GEMM_UNIVERSAL_MAX_INSTANCES "" CACHE STRING "Max kernel instances per (dtype, layout) combo (empty = no cap)") -option(ENABLE_CCACHE_GEMM_UNIVERSAL "Enable ccache for GEMM Universal ops compilation" OFF) - -# Store the directory path for use in functions -set(GEMM_UNIVERSAL_SOURCE_DIR ${CMAKE_CURRENT_LIST_DIR}) - -# Function to create individual GEMM Universal targets -function(create_individual_gemm_universal_target datatype layout trait tile_config config_json) - # Use the parent scope GEMM_UNIVERSAL_GPU_TARGETS_INDIVIDUAL variable - if(NOT GEMM_UNIVERSAL_GPU_TARGETS_INDIVIDUAL) - message(WARNING "Skipping individual GEMM Universal target ${datatype}_${layout}_${trait}_${tile_config}: No supported GPU targets") - return() - endif() - - # Parse tile configuration: format is tile_mxtile_nxtile_k_warp_mxwarp_nxwarp_k_warp_tile_mxwarp_tile_nxwarp_tile_k - # First split by underscore to get three groups - string(REPLACE "_" ";" config_groups ${tile_config}) - list(GET config_groups 0 tile_dims) # e.g., 256x256x32 - list(GET config_groups 1 warp_dims) # e.g., 4x1x1 - list(GET config_groups 2 warp_tile_dims) # e.g., 16x16x16 - - # Parse tile dimensions - string(REPLACE "x" ";" tile_parts ${tile_dims}) - list(GET tile_parts 0 tile_m) - list(GET tile_parts 1 tile_n) - list(GET tile_parts 2 tile_k) - - # Parse warp dimensions - string(REPLACE "x" ";" warp_parts ${warp_dims}) - list(GET warp_parts 0 warp_m) - list(GET warp_parts 1 warp_n) - list(GET warp_parts 2 warp_k) - - # Parse warp tile dimensions - string(REPLACE "x" ";" warp_tile_parts ${warp_tile_dims}) - list(GET warp_tile_parts 0 warp_tile_m) - list(GET warp_tile_parts 1 warp_tile_n) - list(GET warp_tile_parts 2 warp_tile_k) - - set(target_name "benchmark_gemm_universal_${datatype}_${layout}_${trait}_${tile_config}") - set(working_path "${CMAKE_CURRENT_BINARY_DIR}/${datatype}/${layout}") - - # Generate the single instance header for this kernel - set(instance_header "${working_path}/gemm_universal_single_${datatype}_${layout}_${trait}_${tile_config}.hpp") - - # Add custom command to generate the header file at build time - add_custom_command( - OUTPUT ${instance_header} - COMMAND ${Python3_EXECUTABLE} ${GEMM_UNIVERSAL_SOURCE_DIR}/gemm_universal_instance_builder.py - --working_path ${working_path} - --datatype ${datatype} - --layout ${layout} - --config_json ${config_json} - --gen_single - --kernel_name "gemm_universal_${datatype}_${layout}_${trait}_${tile_config}" - --tile_config "${tile_config}" - --trait_combo "${trait}" - --gpu_target "${GEMM_UNIVERSAL_GPU_TARGETS_INDIVIDUAL}" - DEPENDS ${GEMM_UNIVERSAL_SOURCE_DIR}/gemm_universal_instance_builder.py ${config_json} - COMMENT "Generating ${instance_header}" - ) - - # Create the executable - add_executable(${target_name} - EXCLUDE_FROM_ALL - ${GEMM_UNIVERSAL_SOURCE_DIR}/gemm_universal_benchmark_single.cpp - ${instance_header} - ) - - # Set GPU architectures - set_property(TARGET ${target_name} PROPERTY HIP_ARCHITECTURES ${GEMM_UNIVERSAL_GPU_TARGETS_INDIVIDUAL}) - - # Set compile definitions - target_compile_definitions(${target_name} PRIVATE - GEMM_UNIVERSAL_SINGLE_INSTANCE_HPP="${instance_header}" - ) - - # Include directories - target_include_directories(${target_name} PRIVATE - ${GEMM_UNIVERSAL_SOURCE_DIR} - ${working_path} - ) - - # Compile options - target_compile_options(${target_name} PRIVATE - -Wno-undefined-func-template - -Wno-float-equal - --offload-compress - -include ${instance_header} - ) - - # Add to collection targets - add_dependencies(benchmark_gemm_universal_all ${target_name}) - add_dependencies(benchmark_gemm_universal_${datatype} ${target_name}) - add_dependencies(benchmark_gemm_universal_${layout} ${target_name}) - add_dependencies(benchmark_gemm_universal_${datatype}_${layout} ${target_name}) - - # Add to trait-specific targets - string(REPLACE "_" ";" trait_parts ${trait}) - list(GET trait_parts 0 pipeline) - list(GET trait_parts 1 epilogue) - list(GET trait_parts 2 scheduler) - - add_dependencies(benchmark_gemm_universal_${pipeline}_pipeline ${target_name}) - add_dependencies(benchmark_gemm_universal_${epilogue}_epilogue ${target_name}) - add_dependencies(benchmark_gemm_universal_${scheduler}_scheduler ${target_name}) -endfunction() - -# Function to build individual GEMM Universal targets -function(build_individual_gemm_universal_targets datatype layout) - set(working_path "${CMAKE_CURRENT_BINARY_DIR}/${datatype}/${layout}") - - # Choose config file - # Priority order: - # 1. Environment variable GEMM_UNIVERSAL_CONFIG_FILE - # 2. CMake variable GEMM_UNIVERSAL_CONFIG_FILE - # 3. Default based on layout - - # Check environment variable first - if(DEFINED ENV{GEMM_UNIVERSAL_CONFIG_FILE} AND NOT "$ENV{GEMM_UNIVERSAL_CONFIG_FILE}" STREQUAL "") - set(config_filename "$ENV{GEMM_UNIVERSAL_CONFIG_FILE}") - set(json_blob "${CMAKE_CURRENT_LIST_DIR}/configs/${config_filename}") - message(VERBOSE " Using config from environment variable: ${config_filename}") - elseif(NOT "${GEMM_UNIVERSAL_CONFIG_FILE}" STREQUAL "") - # Use CMake variable if set - set(json_blob "${CMAKE_CURRENT_LIST_DIR}/configs/${GEMM_UNIVERSAL_CONFIG_FILE}") - message(VERBOSE " Using custom config: ${GEMM_UNIVERSAL_CONFIG_FILE}") - else() - # Use default config for all layouts - set(json_blob "${CMAKE_CURRENT_LIST_DIR}/configs/default_config.json") - message(VERBOSE " Using default config for layout ${layout}") - endif() - - # Check if config file exists - if(NOT EXISTS ${json_blob}) - message(FATAL_ERROR "Config file not found: ${json_blob}") - endif() - - # Determine number of workers for parallel generation - if(DEFINED ENV{CMAKE_BUILD_PARALLEL_LEVEL}) - set(num_workers $ENV{CMAKE_BUILD_PARALLEL_LEVEL}) - else() - # Use processor count but limit to avoid memory issues - cmake_host_system_information(RESULT num_cores QUERY NUMBER_OF_LOGICAL_CORES) - math(EXPR num_workers "${num_cores}") - if(num_workers GREATER 8) - set(num_workers 8) - endif() - endif() - - # Generate individual kernel files using parallel version - message(VERBOSE "Generating individual kernels for ${datatype} ${layout} using ${num_workers} workers...") - message(VERBOSE " Working path: ${working_path}") - message(VERBOSE " Config file: ${json_blob}") - message(VERBOSE " Python executable: ${Python3_EXECUTABLE}") - message(VERBOSE " Script path: ${CMAKE_CURRENT_LIST_DIR}/gemm_universal_instance_builder.py") - - # Create working directory first - file(MAKE_DIRECTORY ${working_path}) - - message(VERBOSE "COMMAND: ${Python3_EXECUTABLE} -u ${CMAKE_CURRENT_LIST_DIR}/gemm_universal_instance_builder.py - --working_path ${working_path} - --datatype ${datatype} - --layout ${layout} - --config_json ${json_blob} - --gpu_target ${GEMM_UNIVERSAL_GPU_TARGETS_INDIVIDUAL} - --list_kernels ") - - # Build optional args for instance builder - set(extra_list_args "") - if(NOT "${GEMM_UNIVERSAL_MAX_INSTANCES}" STREQUAL "") - list(APPEND extra_list_args --max-instances ${GEMM_UNIVERSAL_MAX_INSTANCES}) - endif() - if(NOT "${TILE_ENGINE_SAMPLING_TIER}" STREQUAL "") - list(APPEND extra_list_args --tier ${TILE_ENGINE_SAMPLING_TIER}) - list(APPEND extra_list_args --manifest-path ${working_path}) - endif() - if(NOT "${TILE_ENGINE_SAMPLING_SEED}" STREQUAL "") - list(APPEND extra_list_args --seed ${TILE_ENGINE_SAMPLING_SEED}) - endif() - - # First, just list the kernels (fast operation) - message(VERBOSE " Listing kernel configurations...") - execute_process( - COMMAND ${Python3_EXECUTABLE} -u ${CMAKE_CURRENT_LIST_DIR}/gemm_universal_instance_builder.py - --working_path ${working_path} - --datatype ${datatype} - --layout ${layout} - --config_json ${json_blob} - --gpu_target ${GEMM_UNIVERSAL_GPU_TARGETS_INDIVIDUAL} - --list_kernels - ${extra_list_args} - WORKING_DIRECTORY ${CMAKE_CURRENT_LIST_DIR} - RESULT_VARIABLE ret - OUTPUT_VARIABLE list_output - ERROR_VARIABLE list_error - ) - - if(NOT ret EQUAL 0) - message(FATAL_ERROR "Failed to list kernels for ${datatype} ${layout}: ${list_error}") - endif() - - # Read kernel count - if(EXISTS ${working_path}/gemm_universal_kernel_count.txt) - file(READ ${working_path}/gemm_universal_kernel_count.txt kernel_count) - string(STRIP "${kernel_count}" kernel_count) - message(VERBOSE " Found ${kernel_count} kernel configurations") - else() - message(FATAL_ERROR "Kernel count file not found") - endif() - - # Read kernel list and create targets - if(EXISTS ${working_path}/gemm_universal_kernel_list.txt) - file(STRINGS ${working_path}/gemm_universal_kernel_list.txt kernel_lines) - foreach(line IN LISTS kernel_lines) - # Parse line: kernel_name|tile_config|trait_combo - string(REPLACE "|" ";" parts "${line}") - list(GET parts 0 kernel_name) - list(GET parts 1 tile_config) - list(GET parts 2 trait_combo) - - # Create individual target - create_individual_gemm_universal_target("${datatype}" "${layout}" "${trait_combo}" "${tile_config}" "${json_blob}") - endforeach() - else() - message(FATAL_ERROR "Kernel list file not found") - endif() -endfunction() - -# Main build logic - Only individual builds supported -message(VERBOSE "=== Starting Tile Engine GEMM Universal Configuration ===") -message(VERBOSE "GEMM_UNIVERSAL_DATATYPE: ${GEMM_UNIVERSAL_DATATYPE}") -message(VERBOSE "GEMM_UNIVERSAL_LAYOUT: ${GEMM_UNIVERSAL_LAYOUT}") -message(VERBOSE "SUPPORTED_GPU_TARGETS: ${SUPPORTED_GPU_TARGETS}") - -# Filter GPU targets to only gfx90a, gfx942, gfx950, gfx1201 -set(GEMM_UNIVERSAL_GPU_TARGETS_INDIVIDUAL "") -set(DESIRED_TARGETS "gfx90a;gfx942;gfx950;gfx1201;gfx12-generic") - -foreach(target IN LISTS SUPPORTED_GPU_TARGETS) - if(target IN_LIST DESIRED_TARGETS) - list(APPEND GEMM_UNIVERSAL_GPU_TARGETS_INDIVIDUAL ${target}) - message(VERBOSE " Adding GPU target: ${target}") - endif() -endforeach() - -# Skip build if no matching targets found -if(NOT GEMM_UNIVERSAL_GPU_TARGETS_INDIVIDUAL) - message(WARNING "Skipping Tile Engine GEMM Universal build: No supported GPU targets (gfx90a, gfx942, gfx950, gfx1201, gfx12-generic) found in SUPPORTED_GPU_TARGETS: ${SUPPORTED_GPU_TARGETS}") -else() - message(VERBOSE "Building individual GEMM Universal targets for GPU targets: ${GEMM_UNIVERSAL_GPU_TARGETS_INDIVIDUAL}") - - # Enable parallel compilation optimizations - # Set up job pools for better parallel compilation control - set_property(GLOBAL PROPERTY JOB_POOLS - compile_heavy=4 # Limit heavy compilations to prevent OOM - compile_normal=16 # Allow more parallel normal compilations - ) - - # Enable compiler cache if available and explicitly requested - # Disabled by default due to permission issues in CI environments - if(ENABLE_CCACHE_GEMM_UNIVERSAL) - find_program(CCACHE_PROGRAM ccache) - if(CCACHE_PROGRAM) - set(CMAKE_CXX_COMPILER_LAUNCHER ${CCACHE_PROGRAM}) - message(VERBOSE "Using ccache for faster compilation") - else() - message(WARNING "ccache requested but not found") - endif() - else() - message(VERBOSE "ccache disabled for GEMM Universal ops (use -DENABLE_CCACHE_GEMM_UNIVERSAL=ON to enable)") - endif() - - # Create master collection targets - add_custom_target(benchmark_gemm_universal_all) - - # Create datatype collection targets - foreach(dt IN LISTS GEMM_UNIVERSAL_DATATYPE) - add_custom_target(benchmark_gemm_universal_${dt}) - endforeach() - - # Create layout collection targets - foreach(l IN LISTS GEMM_UNIVERSAL_LAYOUT) - add_custom_target(benchmark_gemm_universal_${l}) - endforeach() - - # Create combined collection targets - foreach(dt IN LISTS GEMM_UNIVERSAL_DATATYPE) - foreach(l IN LISTS GEMM_UNIVERSAL_LAYOUT) - add_custom_target(benchmark_gemm_universal_${dt}_${l}) - endforeach() - endforeach() - - # Create trait-based collection targets - # These are common trait components used across all GEMM Universal kernels - set(GEMM_UNIVERSAL_PIPELINES "mem;compv3;compv4") - set(GEMM_UNIVERSAL_EPILOGUES "default;cshuffle") - set(GEMM_UNIVERSAL_SCHEDULERS "intrawave;interwave") - - foreach(pipeline IN LISTS GEMM_UNIVERSAL_PIPELINES) - add_custom_target(benchmark_gemm_universal_${pipeline}_pipeline) - endforeach() - - foreach(epilogue IN LISTS GEMM_UNIVERSAL_EPILOGUES) - add_custom_target(benchmark_gemm_universal_${epilogue}_epilogue) - endforeach() - - foreach(scheduler IN LISTS GEMM_UNIVERSAL_SCHEDULERS) - add_custom_target(benchmark_gemm_universal_${scheduler}_scheduler) - endforeach() - - # Divide MAX_INSTANCES budget across all active (dtype, layout) combos so that - # sampling fires per-combo rather than being a single cap larger than any combo's - # feasible set (which would make sampling a no-op for most combos). - if(NOT "${GEMM_UNIVERSAL_MAX_INSTANCES}" STREQUAL "") - list(LENGTH GEMM_UNIVERSAL_DATATYPE _gu_n_dt) - list(LENGTH GEMM_UNIVERSAL_LAYOUT _gu_n_lay) - math(EXPR _gu_n_combos "${_gu_n_dt} * ${_gu_n_lay}") - if(_gu_n_combos GREATER 0) - math(EXPR GEMM_UNIVERSAL_MAX_INSTANCES - "${GEMM_UNIVERSAL_MAX_INSTANCES} / ${_gu_n_combos}") - message(STATUS " gemm_universal: per-combo budget = ${GEMM_UNIVERSAL_MAX_INSTANCES} (${_gu_n_combos} combos)") - endif() - endif() - - # Build individual targets for each datatype/layout combination - foreach(dt IN LISTS GEMM_UNIVERSAL_DATATYPE) - foreach(l IN LISTS GEMM_UNIVERSAL_LAYOUT) - build_individual_gemm_universal_targets(${dt} ${l}) - endforeach() - endforeach() -endif() diff --git a/projects/composablekernel/tile_engine/ops/gemm/gemm_universal/gemm_universal_benchmark.hpp b/projects/composablekernel/tile_engine/ops/gemm/gemm_universal/gemm_universal_benchmark.hpp deleted file mode 100644 index 23338a6cd008..000000000000 --- a/projects/composablekernel/tile_engine/ops/gemm/gemm_universal/gemm_universal_benchmark.hpp +++ /dev/null @@ -1,73 +0,0 @@ -// Copyright (c) Advanced Micro Devices, Inc., or its affiliates. -// SPDX-License-Identifier: MIT - -#pragma once - -#include -#include -#include -#include -#include - -#include "ck_tile/core.hpp" -#include "ck_tile/host.hpp" -#include "gemm/gemm_benchmark.hpp" - -#if __clang_major__ >= 23 -#pragma clang diagnostic push -#pragma clang diagnostic ignored "-Wlifetime-safety-intra-tu-suggestions" -#endif -// Data types and Layouts are defined by the generated kernel headers -// No hardcoded type definitions here to avoid conflicts - -/// @brief Function to get the kernel output with reference implementation on CPU/GPU -void gemm_host_reference(int verify, - ck_tile::HostTensor& a_m_k, - ck_tile::HostTensor& b_k_n, - ck_tile::HostTensor& c_m_n_host_result, - ck_tile::DeviceMem& a_m_k_dev_buf, - ck_tile::DeviceMem& b_k_n_dev_buf, - ck_tile::index_t M, - ck_tile::index_t N, - ck_tile::index_t K, - ck_tile::index_t stride_A, - ck_tile::index_t stride_B, - ck_tile::index_t stride_C) -{ - if(verify == 1) - { - c_m_n_host_result.SetZero(); - - ck_tile::reference_gemm( - a_m_k, b_k_n, c_m_n_host_result); - } - else if(verify == 2) - { - if constexpr(std::is_same_v) - { - // Restore input for B for gpu reference - b_k_n_dev_buf.ToDevice(b_k_n.data()); - } - - ck_tile::DeviceMem c_m_n_gpu_buf_ref(c_m_n_host_result.get_element_space_size_in_bytes()); - c_m_n_host_result.SetZero(); - c_m_n_gpu_buf_ref.SetZero(); - - ADataType* d_A = static_cast(a_m_k_dev_buf.GetDeviceBuffer()); - BDataType* d_B = static_cast(b_k_n_dev_buf.GetDeviceBuffer()); - CDataType* d_C = static_cast(c_m_n_gpu_buf_ref.GetDeviceBuffer()); - - ck_tile::reference_gemm_gpu(d_A, d_B, d_C, M, N, K, stride_A, stride_B, stride_C); - - c_m_n_gpu_buf_ref.FromDevice(c_m_n_host_result.data()); - } -} -#if __clang_major__ >= 23 -#pragma clang diagnostic pop -#endif diff --git a/projects/composablekernel/tile_engine/ops/gemm/gemm_universal/gemm_universal_benchmark.py b/projects/composablekernel/tile_engine/ops/gemm/gemm_universal/gemm_universal_benchmark.py deleted file mode 100755 index 73ba1261a849..000000000000 --- a/projects/composablekernel/tile_engine/ops/gemm/gemm_universal/gemm_universal_benchmark.py +++ /dev/null @@ -1,149 +0,0 @@ -#!/usr/bin/env python3 -# Copyright (c) Advanced Micro Devices, Inc., or its affiliates. -# SPDX-License-Identifier: MIT - -import os -import sys -import argparse -import time -import importlib.util - - -def _import_gemm_benchmark(): - """Import gemm benchmark from parent directory.""" - current_dir = os.path.dirname(os.path.abspath(__file__)) - parent_dir = os.path.dirname(current_dir) - - # Load the module dynamically - spec = importlib.util.spec_from_file_location( - "gemm_benchmark", - os.path.join(parent_dir, "gemm_benchmark.py"), - ) - gemm_benchmark_module = importlib.util.module_from_spec(spec) - spec.loader.exec_module(gemm_benchmark_module) - - return gemm_benchmark_module.GemmBenchmark - - -def _import_benchmark_utils(): - """Import benchmark utilities from commons directory.""" - current_dir = os.path.dirname(os.path.abspath(__file__)) - parent_dir = os.path.dirname(os.path.dirname(current_dir)) - - # Load the module dynamically - spec = importlib.util.spec_from_file_location( - "benchmark_utils", - os.path.join(parent_dir, "common", "benchmark_utils.py"), - ) - benchmark_utils = importlib.util.module_from_spec(spec) - spec.loader.exec_module(benchmark_utils) - - return benchmark_utils - - -GemmBenchmark = _import_gemm_benchmark() -benchmark_utils = _import_benchmark_utils() - - -class GemmUniversalBenchmark(GemmBenchmark): - def __init__(self, build_dir: str, verbose: bool = False): - super().__init__(build_dir, verbose, name="benchmark_gemm_universal_") - - -def main(): - parser = argparse.ArgumentParser( - description="Universal GEMM Kernel Benchmarking Tool" - ) - parser.add_argument( - "build_dir", help="Build directory containing kernel executables" - ) - parser.add_argument( - "--problem-sizes", - nargs="+", - default=["1024,1024,1024", "2048,2048,2048", "4096,4096,4096"], - help="Problem sizes as M,N,K tuples", - ) - parser.add_argument( - "--split-k", nargs="+", type=int, default=[1], help="Split-K values to test" - ) - parser.add_argument("--verify", action="store_true", help="Enable verification") - parser.add_argument( - "--csv", - default="gemm_universal_benchmark_results.csv", - help="CSV output filename", - ) - parser.add_argument( - "--best", default="best_kernels.txt", help="Best kernels output filename" - ) - parser.add_argument("--verbose", action="store_true", help="Verbose output") - parser.add_argument( - "--warmup", - type=int, - default=50, - help="Number of warmup iterations (default: 50)", - ) - parser.add_argument( - "--repeat", - type=int, - default=100, - help="Number of benchmark iterations (default: 100)", - ) - parser.add_argument( - "--flush-cache", - action="store_true", - default=True, - help="Enable cache flushing (default: True)", - ) - parser.add_argument( - "--rotating-count", - type=int, - default=1000, - help="Number of iterations to rotate cache (default: 1000)", - ) - parser.add_argument("--json", help="JSON output filename (optional)") - - args = parser.parse_args() - - # Parse problem sizes - problem_sizes = [] - for size_str in args.problem_sizes: - try: - m, n, k = map(int, size_str.split(",")) - problem_sizes.append((m, n, k)) - except ValueError: - print(f"Invalid problem size: {size_str}") - return 1 - - # Create benchmark instance - benchmark = GemmUniversalBenchmark(args.build_dir, verbose=args.verbose) - - # Run benchmark sweep - print("Starting Universal GEMM kernel benchmark sweep...") - start_time = time.time() - - best_kernels = benchmark.benchmark_sweep( - problem_sizes=problem_sizes, - split_k_values=args.split_k, - verify=args.verify, - warmup=args.warmup, - repeat=args.repeat, - flush_cache=args.flush_cache, - rotating_count=args.rotating_count, - ) - - elapsed_time = time.time() - start_time - print(f"\nBenchmark completed in {elapsed_time:.2f} seconds") - - # Export results - benchmark_utils.export_csv(benchmark.results, args.csv) - benchmark_utils.export_best_kernels(best_kernels, args.best) - - # Export JSON if requested - if args.json: - benchmark_utils.export_json(benchmark.results, args.json, best_kernels) - - return 0 - - -if __name__ == "__main__": - sys.exit(main()) diff --git a/projects/composablekernel/tile_engine/ops/gemm/gemm_universal/gemm_universal_benchmark_single.cpp b/projects/composablekernel/tile_engine/ops/gemm/gemm_universal/gemm_universal_benchmark_single.cpp deleted file mode 100644 index 9e73077e2895..000000000000 --- a/projects/composablekernel/tile_engine/ops/gemm/gemm_universal/gemm_universal_benchmark_single.cpp +++ /dev/null @@ -1,102 +0,0 @@ -// Copyright (c) Advanced Micro Devices, Inc., or its affiliates. -// SPDX-License-Identifier: MIT - -#include -#include -#include -#include -#include -#include -#include - -#include "ck_tile/core.hpp" -#include "ck_tile/host.hpp" -#include "gemm/gemm_common.hpp" -#include "gemm_universal_profiler.hpp" - -// The kernel header is included via the compile command line with -include flag -// It defines SelectedKernel struct and KERNEL_NAME - -void benchmark_single(const ck_tile::ArgParser& arg_parser) -{ - // Use DataTypeTraits to get the actual type names from the generated header - // The generated header defines ADataType, BDataType, AccDataType, CDataType - std::string dtype_a = ck_tile::DataTypeTraits::name; - std::string dtype_b = ck_tile::DataTypeTraits::name; - std::string dtype_acc = ck_tile::DataTypeTraits::name; - std::string dtype_c = ck_tile::DataTypeTraits::name; - - // Layout names from the layout types - std::string layout_a = ALayout::name; - std::string layout_b = BLayout::name; - std::string layout_c = CLayout::name; - - // Create GemmProblem struct - GemmProblem gemm_problem{arg_parser.get_int("split_k"), - arg_parser.get_int("m"), - arg_parser.get_int("n"), - arg_parser.get_int("k"), - arg_parser.get_int("stride_a"), - arg_parser.get_int("stride_b"), - arg_parser.get_int("stride_c"), - dtype_a, - dtype_b, - dtype_acc, - dtype_c, - layout_a, - layout_b, - layout_c, - arg_parser.get_bool("structured_sparsity")}; - - // Create Settings struct - Settings setting{arg_parser.get_int("warmup"), - arg_parser.get_int("repeat"), - arg_parser.get_bool("timer"), - arg_parser.get_int("verify"), - arg_parser.get_int("init"), - arg_parser.get_bool("log"), - arg_parser.get_str("csv_filename"), - arg_parser.get_bool("flush_cache"), - arg_parser.get_int("rotating_count"), - arg_parser.get_bool("json_output")}; - - // Get the profiler instance - auto& profiler = UniversalGemmProfiler::GemmProfiler::instance(setting); - - try - { - // Create a lambda that wraps the kernel launch - auto kernel_func = [](const ck_tile::GemmHostArgs& args, - const ck_tile::stream_config& stream) { - return SelectedKernel::launch(args, stream); - }; - - // Benchmark the kernel - profiler.benchmark(gemm_problem, kernel_func); - - // Select best instance based on metric - profiler.select_best_instance(static_cast(arg_parser.get_int("metric"))); - } - catch(const std::exception& e) - { - std::cerr << "Benchmark failed: " << e.what() << std::endl; - } -} - -int main(int argc, char* argv[]) -{ - try - { - auto [result, parser] = create_args(argc, argv); - if(!result) - return EXIT_FAILURE; - - benchmark_single(parser); - return 0; - } - catch(const std::exception& e) - { - std::cerr << "Error: " << e.what() << "\n"; - return EXIT_FAILURE; - } -} diff --git a/projects/composablekernel/tile_engine/ops/gemm/gemm_universal/gemm_universal_instance_builder.py b/projects/composablekernel/tile_engine/ops/gemm/gemm_universal/gemm_universal_instance_builder.py deleted file mode 100644 index 0d13584ca065..000000000000 --- a/projects/composablekernel/tile_engine/ops/gemm/gemm_universal/gemm_universal_instance_builder.py +++ /dev/null @@ -1,344 +0,0 @@ -# Copyright (c) Advanced Micro Devices, Inc., or its affiliates. -# SPDX-License-Identifier: MIT - -import os -import argparse -import importlib.util -import multiprocessing -import concurrent.futures - - -def _import_gemm_kernel_builder(): - """Import validation utilities from commons directory.""" - current_dir = os.path.dirname(os.path.abspath(__file__)) - parent_dir = os.path.dirname(current_dir) - - # Load the module dynamically - spec = importlib.util.spec_from_file_location( - "gemm_instance_builder", - os.path.join(parent_dir, "gemm_instance_builder.py"), - ) - gemm_builder_module = importlib.util.module_from_spec(spec) - spec.loader.exec_module(gemm_builder_module) - - return gemm_builder_module.GemmKernelBuilder - - -GemmKernelBuilder = _import_gemm_kernel_builder() - - -class GemmUniversalKernelBuilder(GemmKernelBuilder): - def __init__( - self, - kernel_name_prefix, - working_path, - gpu_target, - datatype, - layout, - config_json=None, - max_instances=None, - seed=None, - tier=None, - manifest_path=None, - ): - super().__init__( - kernel_name_prefix, - working_path, - gpu_target, - datatype, - layout, - config_json, - max_instances=max_instances, - seed=seed, - tier=tier, - manifest_path=manifest_path, - ) - - def _generate_all_individual(self, num_workers=None): - """Generate individual kernel files for separate compilation with parallel processing""" - if num_workers is None: - num_workers = min( - multiprocessing.cpu_count(), 8 - ) # Limit to avoid memory issues - - tile_configs = self._get_tile_configs() - trait_combos = self._generate_trait_combinations() - - # Prepare work items for parallel processing - work_items = [] - for tile_config in tile_configs: - for trait_combo in trait_combos: - work_items.append( - ( - tile_config, - trait_combo, - self.kernel_name_prefix, - self.working_path, - self.gpu_target, - self.datatype, - self.layout, - self.config_json, - ) - ) - - # Apply RFC-compliant sampling (Sobol + LHS + maximin) - if self.max_instances is not None and len(work_items) > self.max_instances: - kernel_dicts = [ - {"tile_config": item[0], "trait_combo": item[1], "_work_item": item} - for item in work_items - ] - sampled = self._apply_sampling(kernel_dicts) - work_items = [k["_work_item"] for k in sampled] - - print( - f"Generating {len(work_items)} individual kernel files using {num_workers} workers..." - ) - print(f" Tile configs: {len(tile_configs)}") - print(f" Trait combinations: {len(trait_combos)}") - print(f" Total kernels: {len(work_items)}") - - # Show first few work items for debugging - if work_items: - print(" First work item example:") - tile_config, trait_combo = work_items[0][:2] - print(f" Tile config: {tile_config}") - print(f" Trait combo: {trait_combo[:3]}") # Show first 3 traits - - # Process work items in parallel - kernel_list = [] - completed = 0 - - with concurrent.futures.ProcessPoolExecutor( - max_workers=num_workers - ) as executor: - # Submit all work items - print(f" Submitting {len(work_items)} tasks to executor...") - future_to_item = { - executor.submit(_generate_single_kernel_individual, item): item - for item in work_items - } - print(" All tasks submitted, waiting for completion...") - - # Collect results with progress reporting - for future in concurrent.futures.as_completed(future_to_item): - completed += 1 - if completed % 100 == 0 or completed == len(work_items): - print( - f" Progress: {completed}/{len(work_items)} kernels generated" - ) - try: - result = future.result() - if result: - kernel_list.append(result) - except Exception as exc: - item = future_to_item[future] - print(f"Kernel generation failed for {item}: {exc}") - - # Sort kernel list for consistent ordering - kernel_list.sort(key=lambda x: x[0]) # Sort by kernel name - - # Generate CMake include file for individual targets - self._generate_cmake_individual_targets(kernel_list) - - print( - f"Generated {len(kernel_list)} individual kernel files in {self.working_path}" - ) - - -def _generate_single_kernel_individual(work_item): - """Worker function to generate a single individual kernel file""" - ( - tile_config, - trait_combo, - kernel_name_prefix, - working_path, - gpu_target, - datatype, - layout, - config_json, - ) = work_item - - # Create a temporary builder instance for this worker - builder = GemmUniversalKernelBuilder( - kernel_name_prefix, working_path, gpu_target, datatype, layout, config_json - ) - - try: - kernel_name, instance_code = builder._generate_kernel_instance( - tile_config, trait_combo - ) - - # Create simplified filename without the "gemm_universal_" prefix - # Remove "gemm_universal_" from the beginning of kernel_name for the filename - simplified_name = kernel_name - if simplified_name.startswith("gemm_universal_"): - simplified_name = simplified_name[ - len(kernel_name_prefix) + 1 : - ] # Remove "gemm_universal" prefix - - # Write individual header file - header_file = working_path / f"gemm_universal_single_{simplified_name}.hpp" - with open(header_file, "w") as f: - f.write(instance_code) - - return (kernel_name, trait_combo, tile_config) - except Exception as e: - print(f"Error generating individual kernel: {e}") - return None - - -def main(): - parser = argparse.ArgumentParser( - description="GEMM Universal kernel instance builder with parallel support" - ) - parser.add_argument("--working_path", required=True, help="Working directory path") - parser.add_argument( - "--gpu_target", - required=True, - help="GPU target architecture", - ) - parser.add_argument( - "--datatype", - required=True, - choices=["fp16", "fp8", "bf16", "bf8"], - help="Data type", - ) - parser.add_argument( - "--layout", - required=True, - choices=["rcr", "rrr", "ccr", "crr"], - help="Matrix layout", - ) - parser.add_argument("--config_json", help="Configuration JSON file") - parser.add_argument( - "--num_workers", type=int, help="Number of parallel workers (default: auto)" - ) - parser.add_argument( - "--gen_all_individual", - action="store_true", - help="Generate individual kernel files", - ) - parser.add_argument( - "--gen_single", action="store_true", help="Generate a single kernel file" - ) - parser.add_argument("--kernel_name", help="Kernel name for single generation") - parser.add_argument( - "--tile_config", help="Tile configuration string for single generation" - ) - parser.add_argument( - "--trait_combo", help="Trait combination string for single generation" - ) - parser.add_argument( - "--list_kernels", - action="store_true", - help="List kernel configurations without generating files", - ) - parser.add_argument( - "--max-instances", - type=int, - default=None, - help="Cap on number of kernel instances per (dtype, layout) combo", - ) - parser.add_argument( - "--seed", - type=int, - default=None, - help="RNG seed for deterministic sampling; if omitted, derived from today's date", - ) - parser.add_argument( - "--tier", - default=None, - help="Sampling tier (daily/weekly)", - ) - parser.add_argument( - "--manifest-path", - default=None, - help="Directory for chosen_instances.json", - ) - - args = parser.parse_args() - - assert args.datatype in ["fp16", "bf16", "fp8", "bf8"], ( - f"Invalid datatype string: {args.datatype} (supported datatypes are [fp16, bf16, fp8, and bf8])" - ) - - layout_parts = args.layout.lower() - assert len(layout_parts) == 3, ( - f"Invalid layout string: {args.layout} (must be 3 characters like 'rcr' where r stands for row major and c stands for column major)" - ) - assert layout_parts[0] in ["r", "c"] and layout_parts[1] in ["r", "c"], ( - f"Invalid matrix_a layout : {layout_parts[0]} or matrix_b layout: {layout_parts[1]} (matrix_a and matrix_b must be either 'r' for row major or 'c' for column major)" - ) - assert layout_parts[2] == "r", ( - f"Invalid matrix_c layout: {layout_parts[2]} (must be 'r' only as currently we are supporting only row major)" - ) - - kernel_name_prefix = "gemm_universal" - builder = GemmUniversalKernelBuilder( - kernel_name_prefix, - args.working_path, - args.gpu_target, - args.datatype, - args.layout, - args.config_json, - max_instances=args.max_instances, - seed=args.seed, - tier=args.tier, - manifest_path=args.manifest_path, - ) - - if args.list_kernels: - builder._list_kernels() - elif args.gen_single: - # Generate a single kernel file input validation - if not args.kernel_name or not args.tile_config or not args.trait_combo: - parser.error( - "--gen_single requires --kernel_name, --tile_config, and --trait_combo" - ) - - # Parse tile config - tile_parts = args.tile_config.split("_") - tile_dims = tile_parts[0].split("x") - warp_dims = tile_parts[1].split("x") - warp_tile_dims = tile_parts[2].split("x") - - tile_config = { - "tile_m": int(tile_dims[0]), - "tile_n": int(tile_dims[1]), - "tile_k": int(tile_dims[2]), - "warp_m": int(warp_dims[0]), - "warp_n": int(warp_dims[1]), - "warp_k": int(warp_dims[2]), - "warp_tile_m": int(warp_tile_dims[0]), - "warp_tile_n": int(warp_tile_dims[1]), - "warp_tile_k": int(warp_tile_dims[2]), - } - - # Parse trait combo - trait_parts = args.trait_combo.split("_") - trait_combo = ( - trait_parts[0], # pipeline - trait_parts[1], # epilogue - trait_parts[2], # scheduler - trait_parts[3] == "True", # pad_m - trait_parts[4] == "True", # pad_n - trait_parts[5] == "True", # pad_k - trait_parts[6] == "True", # persistent - ) - - # Generate the kernel - builder._generate_kernel_instance( - tile_config, - trait_combo, - ) - elif args.gen_all_individual: - # Generate all individual kernel files - builder._generate_all_individual(args.num_workers) - else: - parser.error( - "Must specify one of: --list_kernels, --gen_all_individual, or --gen_single" - ) - - -if __name__ == "__main__": - main() diff --git a/projects/composablekernel/tile_engine/ops/gemm/gemm_universal/gemm_universal_profiler.hpp b/projects/composablekernel/tile_engine/ops/gemm/gemm_universal/gemm_universal_profiler.hpp deleted file mode 100644 index 6eb4266aae88..000000000000 --- a/projects/composablekernel/tile_engine/ops/gemm/gemm_universal/gemm_universal_profiler.hpp +++ /dev/null @@ -1,147 +0,0 @@ -// Copyright (c) Advanced Micro Devices, Inc., or its affiliates. -// SPDX-License-Identifier: MIT - -#pragma once - -#include -#include -#include - -#include "ck_tile/host/device_prop.hpp" -#include "ck_tile/ops/gemm.hpp" -#include "gemm/gemm_benchmark.hpp" -#include "gemm/gemm_profiler.hpp" -#include "gemm_universal_benchmark.hpp" - -class UniversalGemmProfiler - : public GemmProfiler -{ - public: - using BaseGemm = GemmProfiler; - using BaseGemm::benchmark; - - UniversalGemmProfiler(Settings setting) - : GemmProfiler(setting) - { - } - - void benchmark(GemmProblem& gemm_problem, - std::vector( - ck_tile::GemmHostArgs&, const ck_tile::stream_config&)>>& callables) override - { - const ALayout layout_a = ALayout{}; - const BLayout layout_b = BLayout{}; - const CLayout layout_c = CLayout{}; - - gemm_problem.stride_a_ = ck_tile::get_default_stride( - gemm_problem.m_, gemm_problem.k_, gemm_problem.stride_a_, is_row_major(layout_a)); - gemm_problem.stride_b_ = ck_tile::get_default_stride( - gemm_problem.k_, gemm_problem.n_, gemm_problem.stride_b_, is_row_major(layout_b)); - gemm_problem.stride_c_ = ck_tile::get_default_stride( - gemm_problem.m_, gemm_problem.n_, gemm_problem.stride_c_, is_row_major(layout_c)); - - ck_tile::HostTensor a_m_k(ck_tile::host_tensor_descriptor( - gemm_problem.m_, gemm_problem.k_, gemm_problem.stride_a_, is_row_major(layout_a))); - ck_tile::HostTensor b_k_n(ck_tile::host_tensor_descriptor( - gemm_problem.k_, gemm_problem.n_, gemm_problem.stride_b_, is_row_major(layout_b))); - ck_tile::HostTensor c_m_n_dev_result(ck_tile::host_tensor_descriptor( - gemm_problem.m_, gemm_problem.n_, gemm_problem.stride_c_, is_row_major(layout_c))); - - if(setting_.init_method == 0) - { - ck_tile::FillUniformDistribution{-1.f, 1.f}(a_m_k); - ck_tile::FillUniformDistribution{-1.f, 1.f}(b_k_n); - } - else if(setting_.init_method == 1) - { - ck_tile::FillMonotonicSeq{}(a_m_k); - ck_tile::FillMonotonicSeq{}(b_k_n); - } - else if(setting_.init_method == 2) - { - ck_tile::FillConstant{static_cast(1)}(a_m_k); - ck_tile::FillConstant{static_cast(1)}(b_k_n); - } - else - { - a_m_k.SetZero(); - b_k_n.SetZero(); - } - - if(gemm_problem.structured_sparsity_) - { - ck_tile::AdjustToStructuredSparsity{}(a_m_k); - } - - ck_tile::DeviceMem a_m_k_dev_buf(a_m_k.get_element_space_size_in_bytes()); - ck_tile::DeviceMem b_k_n_dev_buf(b_k_n.get_element_space_size_in_bytes()); - ck_tile::DeviceMem c_m_n_dev_buf(c_m_n_dev_result.get_element_space_size_in_bytes()); - - if constexpr(std::is_same_v) - { - // Permute vector pk_i4x4 data for device implementation - ck_tile::HostTensor b_k_n_dev = b_k_n; - // permute_tensor_b(b_k_n_dev); - ck_tile::permute_vectors_i4x4_b(b_k_n_dev); - b_k_n_dev_buf.ToDevice(b_k_n_dev.data()); - } - else - { - b_k_n_dev_buf.ToDevice(b_k_n.data()); - } - - a_m_k_dev_buf.ToDevice(a_m_k.data()); - c_m_n_dev_buf.SetZero(); - c_m_n_dev_result.SetZero(); - - ck_tile::GemmHostArgs gemm_args = { - a_m_k_dev_buf.GetDeviceBuffer(), - b_k_n_dev_buf.GetDeviceBuffer(), - c_m_n_dev_buf.GetDeviceBuffer(), - gemm_problem.split_k_, - gemm_problem.m_, - gemm_problem.n_, - gemm_problem.k_, - gemm_problem.stride_a_, - gemm_problem.stride_b_, - gemm_problem.stride_c_, - }; - - ck_tile::HostTensor c_m_n_host_result(ck_tile::host_tensor_descriptor( - gemm_problem.m_, gemm_problem.n_, gemm_problem.stride_c_, is_row_major(layout_c))); - - if(setting_.verify) - { - gemm_host_reference(setting_.verify, - a_m_k, - b_k_n, - c_m_n_host_result, - a_m_k_dev_buf, - b_k_n_dev_buf, - gemm_problem.m_, - gemm_problem.n_, - gemm_problem.k_, - gemm_problem.stride_a_, - gemm_problem.stride_b_, - gemm_problem.stride_c_); - } - - for(auto& callable : callables) - { - auto kernel_run_result = callable(gemm_args, - ck_tile::stream_config{nullptr, - true, - setting_.log, - setting_.n_warmup, - setting_.n_repeat, - setting_.is_gpu_timer, - setting_.flush_cache, - setting_.rotating_count}); - process_result(gemm_problem, - c_m_n_dev_buf, - c_m_n_host_result, - c_m_n_dev_result, - kernel_run_result); - } - } -}; 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 new file mode 100644 index 000000000000..fccc675222e6 --- /dev/null +++ b/projects/composablekernel/tile_engine/ops/gemm/run_one_gemm_kernel.py @@ -0,0 +1,138 @@ +#!/usr/bin/env python3 +"""Worker script for running GEMM kernels in an isolated subprocess. + +Mirrors grouped_conv's run_one_grouped_conv_kernel.py: +- Receives kernel config + problem via stdin as JSON +- Loads the .so library ONLY inside this subprocess +- Outputs timing results as JSON to stdout (one line per kernel, flushed) +- A GPU fault kills only this process; the parent driver can continue + +Input JSON format: + Single: {"so_path": "...", "problem": {"M":.., "N":.., "K":..}, "kernel_name": "..."} + Batch: {"items": [{"so_path": "...", "problem": {...}, "kernel_name": "..."}, ...]} + +Optional top-level keys ``verify`` (bool) and ``verify_tol`` (float) enable an +fp32 numpy reference check; when set, each OK result also carries ``verified`` +and ``max_rel``. + +Output JSON format (one line per kernel): + {"idx": 0, "ok": true, "ms": 0.123, "tflops": 456.7, "non_zero": 1, "kernel": "..."} + {"idx": 0, "ok": true, ..., "verified": true, "max_rel": 3.1e-4} # with --verify + {"idx": 1, "ok": false, "error": "...", "kernel": "..."} +""" + +import json +import os +import sys + +# Add dispatcher python paths from environment (os.pathsep-separated). +gemm_pypath = os.environ.get("GEMM_PYPATH", "") +if gemm_pypath: + for p in gemm_pypath.split(os.pathsep): + if p and p not in sys.path: + sys.path.insert(0, p) + +from gemm_utils import GemmProblem, GpuGemmRunner # noqa: E402 +import numpy as np # noqa: E402 + + +def _run_one(idx, so_path, prob_dict, kernel_name, verify=False, verify_tol=2e-2): + """Run a single 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``). + """ + try: + problem = GemmProblem.from_dict(prob_dict) + + np.random.seed(42) + A = (np.random.randn(problem.M, problem.K) * 0.1).astype(np.float16) + B = (np.random.randn(problem.K, problem.N) * 0.1).astype(np.float16) + + # 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 = ( + int(np.count_nonzero(result.output)) + if result.output is not None + else 0 + ) + out = { + "idx": idx, + "ok": True, + "ms": result.time_ms, + "tflops": result.tflops, + "non_zero": non_zero, + "kernel": kernel_name, + } + if verify: + ref = A.astype(np.float32) @ B.astype(np.float32) + got = result.output.astype(np.float32) + denom = float(np.max(np.abs(ref))) or 1.0 + max_rel = float(np.max(np.abs(got - ref)) / denom) + out["max_rel"] = max_rel + out["verified"] = bool(max_rel <= verify_tol) + print(json.dumps(out), flush=True) + else: + print( + json.dumps( + { + "idx": idx, + "ok": False, + "error": f"kernel returned status {result.status}", + "kernel": kernel_name, + } + ), + flush=True, + ) + + except Exception as e: + print( + json.dumps( + {"idx": idx, "ok": False, "error": str(e), "kernel": kernel_name} + ), + flush=True, + ) + + +def main(): + """Read JSON from stdin, run kernel(s), output results.""" + try: + d = json.loads(sys.stdin.buffer.read()) + except Exception as e: + print( + json.dumps({"idx": 0, "ok": False, "error": f"JSON parse error: {e}"}), + flush=True, + ) + sys.exit(1) + + verify = bool(d.get("verify", False)) + verify_tol = float(d.get("verify_tol", 2e-2)) + + if "items" in d: + for i, item in enumerate(d["items"]): + _run_one( + i, + item["so_path"], + item["problem"], + item.get("kernel_name", "unknown"), + verify=verify, + verify_tol=verify_tol, + ) + else: + _run_one( + 0, + d["so_path"], + d["problem"], + d.get("kernel_name", "unknown"), + verify=verify, + verify_tol=verify_tol, + ) + + +if __name__ == "__main__": + main() diff --git a/projects/composablekernel/tile_engine/ops/gemm/run_one_streamk_gemm_kernel.py b/projects/composablekernel/tile_engine/ops/gemm/run_one_streamk_gemm_kernel.py new file mode 100644 index 000000000000..2a7c8a338689 --- /dev/null +++ b/projects/composablekernel/tile_engine/ops/gemm/run_one_streamk_gemm_kernel.py @@ -0,0 +1,170 @@ +#!/usr/bin/env python3 +"""Worker script for running Stream-K GEMM kernels in an isolated subprocess. + +Stream-K is a single-problem GEMM with the same C ABI as regular GEMM, so this +worker is identical in shape to ``run_one_gemm_kernel.py`` and reuses +``GemmProblem`` / ``GpuGemmRunner`` unchanged -- the Stream-K specifics live +entirely inside the force-included kernel and ``streamk_gemm_ctypes_lib.cpp``. + +- Receives kernel config + problem via stdin as JSON +- Loads the .so library ONLY inside this subprocess +- Outputs timing results as JSON to stdout (one line per kernel, flushed) +- A GPU fault kills only this process; the parent driver can continue + +Input JSON format: + Single: {"so_path": "...", "problem": {"M":.., "N":.., "K":..}, "kernel_name": "..."} + Batch: {"items": [{"so_path": "...", "problem": {...}, "kernel_name": "..."}, ...]} + +Optional top-level keys ``verify`` (bool) and ``verify_tol`` (float) enable an +fp32 numpy reference check; when set, each OK result also carries ``verified`` +and ``max_rel``. Stream-K's Atomic reduction does multiple fp16 atomic-adds (one +per K-split partial), so it is inherently noisier than a single fp32->fp16 store; +the default gate tolerance (2e-2) is loose enough to pass while still catching +gross errors. + +Output JSON format (one line per kernel): + {"idx": 0, "ok": true, "ms": 0.123, "tflops": 456.7, "non_zero": 1, "kernel": "..."} + {"idx": 0, "ok": true, ..., "verified": true, "max_rel": 1.1e-3} # with --verify + {"idx": 1, "ok": false, "error": "...", "kernel": "..."} +""" + +import json +import os +import sys + +# Add dispatcher python paths from environment (os.pathsep-separated). +gemm_pypath = os.environ.get("GEMM_PYPATH", "") +if gemm_pypath: + for p in gemm_pypath.split(os.pathsep): + if p and p not in sys.path: + sys.path.insert(0, p) + +from gemm_utils import GemmProblem, GpuGemmRunner # noqa: E402 +import numpy as np # noqa: E402 + + +def _run_one(idx, so_path, prob_dict, kernel_name, verify=False, verify_tol=2e-2): + """Run a single 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``). + """ + try: + problem = GemmProblem.from_dict(prob_dict) + + np.random.seed(42) + A = (np.random.randn(problem.M, problem.K) * 0.1).astype(np.float32) + B = (np.random.randn(problem.K, problem.N) * 0.1).astype(np.float32) + + # CRITICAL: load the library ONLY inside this subprocess. The runner reads + # dtype + layout off the kernel name and arranges/encodes A/B accordingly. + runner = GpuGemmRunner(lib_path=so_path) + result = runner.run(A, B, problem) + + if result.success: + non_zero = ( + int(np.count_nonzero(result.output)) + if result.output is not None + else 0 + ) + out = { + "idx": idx, + "ok": True, + "ms": result.time_ms, + "tflops": result.tflops, + "non_zero": non_zero, + "kernel": kernel_name, + } + if verify: + # Reference uses the SAME quantized inputs the device sees, per the + # kernel's dtype (bf16/fp8/bf8 bit-quantization vs fp16), so the + # metric isolates compute error from input quantization. int8 is + # exact: the device multiplies the int8 values directly. + kdt = getattr(runner, "_dtype", "fp16") + if kdt == "bf16": + Aq = GpuGemmRunner._bf16_decode(GpuGemmRunner._bf16_encode(A)) + Bq = GpuGemmRunner._bf16_decode(GpuGemmRunner._bf16_encode(B)) + ref = Aq @ Bq + elif kdt == "fp8": + Aq = GpuGemmRunner._fp8_decode(GpuGemmRunner._fp8_encode(A)) + Bq = GpuGemmRunner._fp8_decode(GpuGemmRunner._fp8_encode(B)) + ref = Aq @ Bq + elif kdt == "bf8": + Aq = GpuGemmRunner._bf8_decode(GpuGemmRunner._bf8_encode(A)) + Bq = GpuGemmRunner._bf8_decode(GpuGemmRunner._bf8_encode(B)) + ref = Aq @ Bq + elif kdt == "int8": + Aq = A.astype(np.int8).astype(np.int32) + Bq = B.astype(np.int8).astype(np.int32) + ref = (Aq @ Bq).astype(np.float32) + else: + Aq = A.astype(np.float16).astype(np.float32) + Bq = B.astype(np.float16).astype(np.float32) + ref = Aq @ Bq + got = result.output.astype(np.float32) + denom = float(np.max(np.abs(ref))) or 1.0 + max_rel = float(np.max(np.abs(got - ref)) / denom) + out["max_rel"] = max_rel + out["verified"] = bool(max_rel <= verify_tol) + print(json.dumps(out), flush=True) + else: + print( + json.dumps( + { + "idx": idx, + "ok": False, + "error": f"kernel returned status {result.status}", + "kernel": kernel_name, + } + ), + flush=True, + ) + + except Exception as e: + print( + json.dumps( + {"idx": idx, "ok": False, "error": str(e), "kernel": kernel_name} + ), + flush=True, + ) + + +def main(): + """Read JSON from stdin, run kernel(s), output results.""" + try: + d = json.loads(sys.stdin.buffer.read()) + except Exception as e: + print( + json.dumps({"idx": 0, "ok": False, "error": f"JSON parse error: {e}"}), + flush=True, + ) + sys.exit(1) + + verify = bool(d.get("verify", False)) + verify_tol = float(d.get("verify_tol", 2e-2)) + + if "items" in d: + for i, item in enumerate(d["items"]): + _run_one( + i, + item["so_path"], + item["problem"], + item.get("kernel_name", "unknown"), + verify=verify, + verify_tol=verify_tol, + ) + else: + _run_one( + 0, + d["so_path"], + d["problem"], + d.get("kernel_name", "unknown"), + verify=verify, + verify_tol=verify_tol, + ) + + +if __name__ == "__main__": + main() diff --git a/projects/composablekernel/tile_engine/ops/gemm/streamk_gemm_full_benchmark.py b/projects/composablekernel/tile_engine/ops/gemm/streamk_gemm_full_benchmark.py new file mode 100644 index 000000000000..8f7ba89eb840 --- /dev/null +++ b/projects/composablekernel/tile_engine/ops/gemm/streamk_gemm_full_benchmark.py @@ -0,0 +1,492 @@ +#!/usr/bin/env python3 +"""Full Stream-K GEMM benchmark sweep driven through the Dispatcher bridge. + +Same 3-phase architecture as ``gemm_full_benchmark.py`` -- Stream-K is a +single-problem GEMM (one A/B/C, one M/N/K) with the *same* C ABI, so the only +bridge difference is ``variant="stream_k"`` threaded into ``expand_sweep`` (which +makes the dispatcher codegen the Stream-K launch and ``gemm_utils`` select +``streamk_gemm_ctypes_lib.cpp``). The GPU worker reuses ``GemmProblem`` / +``GpuGemmRunner`` unchanged. + + Phase 1: Compile all kernels (parallel, returns .so paths only -- no GPU) + Phase 2: Load problems (M, N, K shapes) + Phase 3: Benchmark via subprocess isolation, distributed across all visible + GPUs (one device-pinned worker per GPU, batched, fault-isolated) + +Like the regular bridge driver, Phase 3 fans the (kernel x problem) work out +across every visible GPU in parallel: each device runs its own stream of +disposable worker subprocesses pinned with ``HIP_VISIBLE_DEVICES``, so an N-GPU +box benchmarks roughly N times faster while keeping per-batch fault isolation. + +Examples: + # Default config (gemm_streamk/configs/default_config.json), all visible GPUs: + python streamk_gemm_full_benchmark.py + + # Explicit config on 4 GPUs with correctness checking: + python streamk_gemm_full_benchmark.py gemm_streamk/configs/default_config.json \ + --devices 4 --verify --csv streamk_gemm_results.csv +""" + +import argparse +import csv +import json +import os +import queue +import re +import subprocess +import sys +import threading +import time +from pathlib import Path + +_THIS_DIR = Path(__file__).resolve().parent +_DISPATCHER_ROOT = _THIS_DIR.parents[2] / "dispatcher" +sys.path.insert(0, str(_DISPATCHER_ROOT / "python")) +sys.path.insert(0, str(_THIS_DIR)) + +from gemm_utils import setup_multiple_gemm_dispatchers, expand_sweep # noqa: E402 + +# Stream-K is a single variant; its sweep configs live in gemm_streamk/configs/. +DEFAULT_CONFIG = _THIS_DIR / "gemm_streamk" / "configs" / "default_config.json" + +# Default problem set: squares plus a large-K skinny shape -- Stream-K's sweet +# spot is few output tiles with a long K reduction. Tiny problems (e.g. 257^3) +# have too few tiles to partition across CUs and the kernel reports them as +# unsupported (status -2), which the bridge surfaces gracefully. +DEFAULT_PROBLEMS = [ + {"M": 1024, "N": 1024, "K": 1024}, + {"M": 2048, "N": 2048, "K": 2048}, + {"M": 4096, "N": 4096, "K": 4096}, + {"M": 512, "N": 512, "K": 8192}, +] + +# Bridge surface for Stream-K. The dispatcher host path +# (streamk_gemm_ctypes_lib.cpp) derives strides from the kernel's layouts and the +# worker (run_one_streamk_gemm_kernel.py) reads dtype/layout off the kernel name, +# so all 4 A/B/C layouts are supported. dtypes cover fp16 + bf16 + fp8 + bf8 (the +# codecs the bridge runner implements); fp8/bf8 use the gfx942 FNUZ format and +# accumulate into fp16. int8 is left out: it is blocked at the ck_tile engine +# level, not the bridge -- the int8 kernel codegens but fails to COMPILE for +# every reduction strategy (atomic/linear/tree). warp_gemm_dispatcher has no +# Dispatcher specialization for the streamk +# CompV3 path, so WarpGemm resolves to `int` and the BlockUniversalGemmAsBsCr +# WarpGemm::kM/kN static_asserts fail. The runner keeps an int8 codec ready for +# when the engine adds that instantiation; this matches PR #8094 leaving int8 out. +SUPPORTED_DTYPES = ("fp16", "bf16", "fp8", "bf8") +SUPPORTED_LAYOUTS = ("rcr", "rrr", "ccr", "crr") + + +def detect_devices(): + """Return a list of visible GPU id strings (best-effort).""" + env = os.environ.get("HIP_VISIBLE_DEVICES") or os.environ.get( + "CUDA_VISIBLE_DEVICES" + ) + if env: + ids = [d.strip() for d in env.split(",") if d.strip() != ""] + if ids: + return ids + try: + out = subprocess.check_output( + ["rocm-smi", "--showid"], stderr=subprocess.DEVNULL, text=True + ) + ids = sorted(set(re.findall(r"GPU\[(\d+)\]", out)), key=int) + if ids: + return ids + except Exception: + pass + try: + out = subprocess.check_output( + ["amd-smi", "list"], stderr=subprocess.DEVNULL, text=True + ) + ids = re.findall(r"^GPU:\s*(\d+)", out, re.MULTILINE) + if ids: + return ids + except Exception: + pass + return ["0"] + + +def resolve_devices(spec): + """Resolve --devices into a concrete list of device id strings. + + spec is None (auto: all visible), an int count, or a comma-list of ids. + A bare digit is a *count*, not an id; to target one specific id use the + comma form, e.g. "5,". + """ + detected = detect_devices() + if spec is None: + return detected + spec = str(spec).strip() + if "," in spec: + return [s.strip() for s in spec.split(",") if s.strip() != ""] + if spec.isdigit(): + n = int(spec) + if n <= 0: + return detected + return detected[:n] if len(detected) >= n else [str(i) for i in range(n)] + return [spec] + + +def resolve_configs(args): + """Resolve positional configs -> concrete list of config paths.""" + if args.configs: + return args.configs + return [str(DEFAULT_CONFIG)] + + +def load_problems(path): + if not path: + return DEFAULT_PROBLEMS + with open(path) as f: + data = json.load(f) + # Accept either a bare list or {"problems": [...]}. + return data["problems"] if isinstance(data, dict) else data + + +def _run_batch_on_device(device_id, unit, args, worker_path, base_env): + """Run one (problem, kernel-batch) unit in a device-pinned subprocess. + + Returns (rows, lines, n_fail) where rows are dicts ready for the CSV writer, + lines are formatted strings to print, and n_fail counts failures. + """ + 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 + ] + payload = json.dumps( + {"items": items, "verify": args.verify, "verify_tol": args.verify_tol} + ) + + env = base_env.copy() + env["HIP_VISIBLE_DEVICES"] = str(device_id) + + rows, lines, n_fail = [], [], 0 + proc = None + try: + proc = subprocess.Popen( + [sys.executable, str(worker_path)], + stdin=subprocess.PIPE, + stdout=subprocess.PIPE, + stderr=subprocess.DEVNULL, + env=env, + ) + stdout_bytes, _ = proc.communicate( + input=payload.encode("utf-8"), + timeout=args.kernel_timeout * len(batch), + ) + + reported = set() + for line in stdout_bytes.decode("utf-8").strip().split("\n"): + if not line: + continue + try: + result = json.loads(line) + except json.JSONDecodeError: + lines.append(f" [gpu{device_id}] Warning: bad result line: {line[:50]}") + n_fail += 1 + continue + bidx = result.get("idx", 0) + _, cfg, _ = batch[bidx] + reported.add(bidx) + if result.get("ok", False): + status = "OK" if result.get("non_zero", 0) > 0 else "ZERO" + mismatch = False + if args.verify and "verified" in result: + if result["verified"]: + status = "VERIFY" + else: + status = "MISMATCH" + mismatch = True + extra = f" rel={result['max_rel']:.2e}" if "max_rel" in result else "" + lines.append( + f" [gpu{device_id}] {cfg.name:<58} {result['ms']:>10.3f} " + f"{result['tflops']:>10.2f} {status:>8}{extra}" + ) + rows.append( + { + "kernel": cfg.name, + "problem_idx": prob_idx, + "M": M, + "N": N, + "K": K, + "device": device_id, + "latency_ms": result["ms"], + "tflops": result["tflops"], + "non_zero": result.get("non_zero", 0), + "max_rel": result.get("max_rel", ""), + "verified": result.get("verified", ""), + } + ) + if mismatch: + n_fail += 1 + else: + lines.append(f" [gpu{device_id}] {cfg.name:<58} FAILED") + lines.append(f" Error: {result.get('error', 'unknown')[:100]}") + n_fail += 1 + + missing = set(range(len(batch))) - reported + if missing or proc.returncode != 0: + if proc.returncode != 0: + lines.append(f" [gpu{device_id}] worker exited code {proc.returncode}") + for idx in sorted(missing): + _, cfg, _ = batch[idx] + lines.append(f" [gpu{device_id}] {cfg.name:<58} MISSING (crash)") + n_fail += len(missing) + + except subprocess.TimeoutExpired: + lines.append(f" [gpu{device_id}] batch timeout ({len(batch)} kernels)") + try: + proc.kill() + proc.communicate(timeout=5) + except Exception: + pass + n_fail += len(batch) + except Exception as e: + lines.append(f" [gpu{device_id}] batch error: {e}") + try: + if proc and proc.poll() is None: + proc.kill() + except Exception: + pass + n_fail += len(batch) + + return rows, lines, n_fail + + +def main(): + parser = argparse.ArgumentParser( + description="Stream-K GEMM Benchmark Sweep (via Dispatcher)" + ) + parser.add_argument( + "configs", + nargs="*", + help="TE sweep config JSON files (default: gemm_streamk/configs/default_config.json)", + ) + parser.add_argument("--arch", default="gfx942") + parser.add_argument( + "--dtype", + default="fp16", + choices=SUPPORTED_DTYPES, + help=f"Input dtype (supported: {', '.join(SUPPORTED_DTYPES)})", + ) + parser.add_argument( + "--layout", + default="rcr", + choices=SUPPORTED_LAYOUTS, + help=f"A/B/C layout (supported: {', '.join(SUPPORTED_LAYOUTS)})", + ) + parser.add_argument("--problems", default=None, help="JSON file of M,N,K problems") + parser.add_argument("--csv", type=str, default="streamk_gemm_results.csv") + parser.add_argument("--workers", type=int, default=8, help="Parallel build workers") + parser.add_argument( + "--devices", + default=None, + help="GPUs to use: int count (e.g. 4) or comma-list of ids (e.g. 0,2,5); " + "for one specific id use the comma form (e.g. 5,) since a bare digit is " + "a count; default auto-detects all visible", + ) + parser.add_argument( + "--batch-size", + type=int, + default=20, + help="Kernels per subprocess (overhead vs fault isolation)", + ) + parser.add_argument( + "--kernel-timeout", type=int, default=30, help="Per-kernel timeout (s)" + ) + parser.add_argument( + "--max-kernels", type=int, default=0, help="Limit to first N kernels (0=all)" + ) + parser.add_argument( + "--verify", + action="store_true", + help="Check each kernel's output against an fp32 numpy reference " + "(global max|out-ref|/max|ref|); a mismatch counts as a failure", + ) + parser.add_argument( + "--verify-tol", + type=float, + default=2e-2, + help="Relative tolerance for --verify (default 2e-2; Stream-K's Atomic " + "reduction is noisier than regular GEMM but stays well under this)", + ) + args = parser.parse_args() + + config_paths = resolve_configs(args) + devices = resolve_devices(args.devices) + + # ======================================================================== + # Phase 1: Compile kernels (parallel, no GPU) + # ======================================================================== + print(f"\n{'=' * 80}") + print("Phase 1: Compile Stream-K kernels") + print(f"{'=' * 80}") + print(f" Configs: {', '.join(config_paths)}") + + all_configs = [] + for cfg_path in config_paths: + all_configs.extend( + expand_sweep( + cfg_path, + args.arch, + dtype=args.dtype, + layout=args.layout, + variant="stream_k", + ) + ) + + if args.max_kernels > 0: + all_configs = all_configs[: args.max_kernels] + + print(f" Expanded configs: {len(all_configs)}") + print(f" Build workers: {args.workers}") + + t0 = time.perf_counter() + # CRITICAL: returns Path objects only, does NOT load any .so. + lib_paths = setup_multiple_gemm_dispatchers( + all_configs, verbose=True, max_workers=args.workers + ) + build_time = time.perf_counter() - t0 + + built_kernels = [ + (cfg, lib) for cfg, lib in zip(all_configs, lib_paths) if lib is not None + ] + + # Dedupe by .so path (distinct configs can map to the same physical kernel). + seen_libs = set() + unique_kernels = [] + duplicate_count = 0 + for cfg, lib in built_kernels: + lib_key = str(lib.resolve()) + if lib_key not in seen_libs: + seen_libs.add(lib_key) + unique_kernels.append((cfg, lib)) + else: + duplicate_count += 1 + built_kernels = unique_kernels + + print( + f"\n Built {len(all_configs)} configs -> {len(built_kernels)} unique kernels " + f"({duplicate_count} duplicates filtered) in {build_time:.0f}s" + ) + + if not built_kernels: + print(" ERROR: No kernels built successfully") + return 1 + + # ======================================================================== + # Phase 2: Load problems + # ======================================================================== + print(f"\n{'=' * 80}") + print("Phase 2: Load test problems") + print(f"{'=' * 80}") + + problems = load_problems(args.problems) + print(f" Problems: {len(problems)}") + print( + f" Total measurements: {len(built_kernels)} x {len(problems)} = " + f"{len(built_kernels) * len(problems)}" + ) + + # ======================================================================== + # Phase 3: Benchmark across all visible GPUs (subprocess isolation, batched) + # ======================================================================== + print(f"\n{'=' * 80}") + print("Phase 3: Benchmark (multi-GPU, subprocess isolation, batched)") + print(f"{'=' * 80}") + print(f" Devices: {len(devices)} -> {', '.join(devices)}") + print(f" Batch size: {args.batch_size} kernels per subprocess") + print(f" Timeout: {args.kernel_timeout}s per kernel\n") + + csv_path = Path(args.csv) + csv_fields = [ + "kernel", + "problem_idx", + "M", + "N", + "K", + "device", + "latency_ms", + "tflops", + "non_zero", + "max_rel", + "verified", + ] + csv_file = open(csv_path, "w", newline="") + writer = csv.DictWriter(csv_file, fieldnames=csv_fields) + writer.writeheader() + + worker_path = _THIS_DIR / "run_one_streamk_gemm_kernel.py" + base_env = os.environ.copy() + base_env["GEMM_PYPATH"] = os.pathsep.join( + [str(_DISPATCHER_ROOT / "python"), str(_THIS_DIR)] + ) + + # Build a single work queue of (prob_idx, prob_dict, kernel-batch) units and + # fan them out across device-pinned worker threads. + work_q = queue.Queue() + for prob_idx, prob in enumerate(problems): + prob_dict = {"M": int(prob["M"]), "N": int(prob["N"]), "K": int(prob["K"])} + for start in range(0, len(built_kernels), args.batch_size): + end = min(start + args.batch_size, len(built_kernels)) + batch = [ + (start + j, cfg, lib) + for j, (cfg, lib) in enumerate(built_kernels[start:end]) + ] + work_q.put((prob_idx, prob_dict, batch)) + + io_lock = threading.Lock() + stats = {"measurements": 0, "failures": 0} + bench_t0 = time.perf_counter() + + def device_thread(device_id): + while True: + try: + unit = work_q.get_nowait() + except queue.Empty: + return + rows, lines, n_fail = _run_batch_on_device( + device_id, unit, args, worker_path, base_env + ) + with io_lock: + for ln in lines: + print(ln) + for row in rows: + writer.writerow(row) + csv_file.flush() + stats["measurements"] += len(rows) + stats["failures"] += n_fail + work_q.task_done() + + threads = [ + threading.Thread(target=device_thread, args=(d,), daemon=True) for d in devices + ] + for t in threads: + t.start() + for t in threads: + t.join() + + bench_time = time.perf_counter() - bench_t0 + csv_file.close() + + # ======================================================================== + # Summary + # ======================================================================== + print(f"\n{'=' * 80}") + print("BENCHMARK COMPLETE") + print(f"{'=' * 80}") + print(f" Build time: {build_time:.0f}s") + print(f" Benchmark time: {bench_time:.0f}s") + print(f" Total time: {build_time + bench_time:.0f}s") + print(f" Devices used: {len(devices)}") + print(f" Successful measurements: {stats['measurements']}") + print(f" Failed measurements: {stats['failures']}") + print(f" Output: {csv_path}") + return 0 + + +if __name__ == "__main__": + sys.exit(main())