[CK_TILE] Stream-K bridge: merge deep-core #8094 codegen into the bridge

Port #8136's Tile-Engine->Dispatcher Stream-K bridge onto the rewritten
deep-core #8094 engine (KernelKey reduction fields, KernelInstance workspace
virtuals, StreamK backend, Dispatcher-owned reduction workspace, registry +
validation driver). 3-way merge over the shared stream_k ancestor; only the
streamk launch emitter in unified_gemm_codegen.py and 03_streamk_gemm_driver.cpp
conflicted -- both resolved to the deep-core side:
- codegen now emits the struct-scope Sk* kernel type + GetWorkSpaceSize +
  IsSupported, keeps the 2-arg internal-workspace launch the bridge ctypes lib
  calls, and adds the 3-arg dispatcher-owned-workspace launch.
- driver takes deep-core's stoll parse + apple-to-apple timing + validate cold shot.

Bridge ctypes lib still bypasses the registry and calls the 2-arg launch directly,
so the bridge runs the exact deep-core kernels. Codegen smoke: atomic/linear/tree
+ regular gemm all generate cleanly (0 failed).

Author: ozturkosu

projects/composablekernel/dispatcher/codegen/unified_gemm_codegen.py

@@ -774,30 +774,43 @@ def _launch_function_streamk(self, config: KernelConfig) -> str:
             "tree": "Tree",
         }[config.reduction_strategy]
         return f"""
-    static float launch(const ck_tile::StreamKHostArgs& args, const stream_config& stream) {{
-        constexpr auto scheduler = {self.tm.SCHEDULER_TO_CK[config.trait.scheduler]};
-        constexpr auto ReductionStrategy = ck_tile::StreamKReductionStrategy::{reduction_ck};
-
-        using GemmUniversalTraits = TileGemmUniversalTraits<kPadM, kPadN, kPadK, DoubleSmemBuffer,
-                                            ALayout, BLayout, CLayout, TransposeC,
-                                            UseStructuredSparsity, UsePersistentKernel,
-                                            NumWaveGroups, Preshuffle>;
-
-        using UniversalGemmProblem = UniversalGemmPipelineProblem<
-            ADataType, BDataType, AccDataType, TileShape, GemmUniversalTraits, scheduler>;
-
-        using GemmPipeline = {self.tm.PIPELINE_TO_CK[config.trait.pipeline]}<UniversalGemmProblem>;
-        {self._epilogue_code(config)}
+    // ---- 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<kPadM, kPadN, kPadK, DoubleSmemBuffer,
+                                        ALayout, BLayout, CLayout, TransposeC,
+                                        UseStructuredSparsity, UsePersistentKernel,
+                                        NumWaveGroups, Preshuffle>;
+    using SkUniversalGemmProblem = UniversalGemmPipelineProblem<
+        ADataType, BDataType, AccDataType, TileShape, SkGemmUniversalTraits, SkScheduler>;
+    using SkGemmPipeline = {self.tm.PIPELINE_TO_CK[config.trait.pipeline]}<SkUniversalGemmProblem>;
+    {self._epilogue_code(config)}
+    using SkStreamKTilePartitioner =
+        ck_tile::StreamKTilePartitioner<TileShape, SkReductionStrategy, UsePersistentKernel>;
+    using StreamKGemmKernel =
+        ck_tile::StreamKKernel<SkStreamKTilePartitioner, SkGemmPipeline, GemmEpilogue>;
+
+    // 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);
+    }}
 
-        using StreamKTilePartitioner =
-            ck_tile::StreamKTilePartitioner<TileShape, ReductionStrategy, UsePersistentKernel>;
-        using StreamKGemmKernel =
-            ck_tile::StreamKKernel<StreamKTilePartitioner, GemmPipeline, GemmEpilogue>;
+    // 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);
-
-        // Allocate the reduction workspace INTERNALLY (dispatcher idiom), not via
-        // an external pointer the way Tile Engine does.
         const auto ws_size = StreamKGemmKernel::GetWorkSpaceSize(kargs);
         ck_tile::DeviceMem workspace_dev(ws_size);
         workspace_dev.SetZero();
@@ -809,42 +822,66 @@ def _launch_function_streamk(self, config: KernelConfig) -> str:
 
         const dim3 grids = StreamKGemmKernel::GridSize(kargs.tile_partitioner);
         const dim3 blocks = StreamKGemmKernel::BlockSize();
-        constexpr int kBlockPerCu = {config.k_block_per_cu};
 
-        // Atomic reduction accumulates into C, so reset it before every run.
-        // Zero only the used MxN region honoring the leading dimension
-        // (stride_E): a flat M*N memset would skip elements when C has a padded
-        // stride and corrupt the atomic accumulation. hipMemset2DAsync handles
-        // the pitched 2D extent and its status is checked, not discarded.
+        // Atomic reduction accumulates into C, so reset buffers before each run.
         auto reset_data_buffers = [&]() {{
-            if constexpr (ReductionStrategy == ck_tile::StreamKReductionStrategy::Atomic) {{
-                constexpr size_t c_elem = sizeof(CDataType);
-                const size_t c_pitch = static_cast<size_t>(args.stride_E) * c_elem;
-                hipError_t memset_status;
-                if constexpr (std::is_same_v<CLayout, ck_tile::tensor_layout::gemm::RowMajor>) {{
-                    // Row-major C: M rows of N elements, row pitch = stride_E.
-                    memset_status = hipMemset2DAsync(args.e_ptr, c_pitch, 0,
-                        static_cast<size_t>(args.N) * c_elem,
-                        static_cast<size_t>(args.M), stream.stream_id_);
-                }} else {{
-                    // Column-major C: N columns of M elements, column pitch = stride_E.
-                    memset_status = hipMemset2DAsync(args.e_ptr, c_pitch, 0,
-                        static_cast<size_t>(args.M) * c_elem,
-                        static_cast<size_t>(args.N), stream.stream_id_);
-                }}
-                if(memset_status != hipSuccess) {{
-                    throw std::runtime_error(
-                        std::string("hipMemset2DAsync failed to zero C: ") +
-                        hipGetErrorString(memset_status));
-                }}
+            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<void()> preprocess = reset_data_buffers;
 
         float ave_time = launch_kernel_time_mask(stream, preprocess,
-            make_kernel<kBlockPerCu>(StreamKGemmKernel{{}}, grids, blocks, 0, kargs));
+            make_kernel<SkBlockPerCu>(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<void()> preprocess = reset_data_buffers;
+
+        float ave_time = launch_kernel_time_mask(stream, preprocess,
+            make_kernel<SkBlockPerCu>(StreamKGemmKernel{{}}, grids, blocks, 0, kargs));
         return ave_time;
     }}"""
 
@@ -899,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<KernelStruct, "
+                "KernelStruct::ADataType, KernelStruct::BDataType, "
+                "KernelStruct::CDataType, KernelStruct::AccDataType>"
+                f'(key, "{kernel_name}");'
+            )
+        else:
+            ret_stmt = (
+                "return std::make_shared<backends::GeneratedKernelInstance<KernelStruct>>"
+                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 {{
@@ -955,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<backends::GeneratedKernelInstance<KernelStruct>>(key, "{kernel_name}");
+
+    {ret_stmt}
 }}
 
 }}}}}}

projects/composablekernel/dispatcher/examples/gemm/cpp/03_streamk_gemm_driver.cpp

@@ -55,18 +55,36 @@ static std::string get_opt(int argc, char** argv, const std::string& key, const
 
 int main(int argc, char** argv)
 {
-    // Parse with std::stoll (not std::stoi) before narrowing to index_t: stoi
-    // throws std::out_of_range past INT_MAX, which would reject large GEMM sizes.
-    const ck_tile::index_t M =
-        static_cast<ck_tile::index_t>(std::stoll(get_opt(argc, argv, "--m", "3840")));
-    const ck_tile::index_t N =
-        static_cast<ck_tile::index_t>(std::stoll(get_opt(argc, argv, "--n", "4096")));
-    const ck_tile::index_t K =
-        static_cast<ck_tile::index_t>(std::stoll(get_opt(argc, argv, "--k", "2048")));
-    const int warmup         = std::stoi(get_opt(argc, argv, "--warmup", "10"));
-    const int repeat         = std::stoi(get_opt(argc, argv, "--repeat", "50"));
+    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";
 
@@ -100,7 +118,8 @@ int main(int argc, char** argv)
                                   sB,
                                   sC};
 
-    const ck_tile::stream_config s{nullptr, true, /*log=*/0, warmup, repeat};
+    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;
@@ -121,9 +140,34 @@ int main(int argc, char** argv)
         ref.SetZero();
         ck_tile::reference_gemm<ADataType, BDataType, AccDataType, CDataType>(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<ADataType, CDataType, AccDataType>(K);
-        const auto atol =
-            ck_tile::get_absolute_threshold<ADataType, CDataType, AccDataType>(maxv, K);
+
+        // 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<ck_tile::index_t>(1, kargs.tile_partitioner.estimate_num_wgs_per_tile());
+
+        using ComputeType =
+            std::conditional_t<sizeof(ADataType) < sizeof(BDataType), ADataType, BDataType>;
+        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<ComputeType, CDataType, AccDataType>(k_per_split);
+        const auto atol_base = ck_tile::get_absolute_threshold<ComputeType, CDataType, AccDataType>(
+            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<CDataType, CDataType, CDataType>(kbatch);
+        const auto atol_split_k =
+            ck_tile::get_absolute_threshold<CDataType, CDataType, CDataType>(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";
     }