feat(rocke): split-K forward implicit-GEMM convolution on gfx1250

dnn-providers/hip-kernel-provider/rocke/library/benchmarks/common/kernel_sweep.py

@@ -250,6 +250,7 @@ def enumerate_jobs(
             split_ks=split_ks,
             streamk_modes=streamk_modes,
             streamk_reductions=streamk_reductions,
+            llvm_flavor=flavor,
         )
         log(f"  {total} candidates to check with {jobs} process(es)")
 
@@ -540,7 +541,15 @@ def compile_all(
     ("grouped", lambda i: "yes" if i.grouped else "no"),
     (
         "split-K",
-        lambda i: ("two-stage" if i.two_stage else "atomic") if i.split_k > 1 else "no",
+        lambda i: (
+            (
+                "cluster"
+                if i.split_k_cluster
+                else "two-stage" if i.needs_stage2 else "atomic"
+            )
+            if i.split_k > 1
+            else "no"
+        ),
     ),
     ("group_merge", lambda i: i.group_merge),
     (
@@ -844,6 +853,18 @@ def _launch_values_for(direction, problem, identity, ptrs, sizes, extras):
 
     tm, tn, tk = identity.tile_m, identity.tile_n, identity.tile_k
     if direction == "fwd":
+        if "ws_ptr" in extras:
+            # Split-K forward: the degree and the workspace, with the slice
+            # width rounded to the binary's tile_k.
+            return ConvArgs.from_problem(
+                problem, tile_m=tm, tile_n=tn, tile_k=tk
+            ).to_launch_values(
+                *ptrs,
+                *sizes,
+                split_k=extras["split_k"],
+                ws_ptr=extras["ws_ptr"],
+                ws_bytes=extras["ws_bytes"],
+            )
         return ConvArgs.from_problem(problem, tile_m=tm, tile_n=tn).to_launch_values(
             *ptrs, *sizes
         )

dnn-providers/hip-kernel-provider/rocke/library/kernels/common/conv_abi.py

@@ -109,8 +109,13 @@ def _grid() -> List[ArgSpec]:
     return [("p_num_pid_m", "i32"), ("p_num_pid_n", "i32")]
 
 
-def _fwd_arg_names(*, is_3d: bool = False) -> List[ArgSpec]:
-    """Forward implicit-GEMM conv: ``A`` (N[D]HWC), ``B`` (K[Z]YXC), ``D`` (N[D]HWK)."""
+def _fwd_arg_names(*, is_3d: bool = False, split_k: bool = False) -> List[ArgSpec]:
+    """Forward implicit-GEMM conv: ``A`` (N[D]HWC), ``B`` (K[Z]YXC), ``D`` (N[D]HWK).
+
+    ``split_k`` is the split-K binary's ABI: the fp32 workspace its slices
+    store into, then the slice width and degree (wgrad's ``ks`` /
+    ``ks_count``).
+    """
     out: List[ArgSpec] = [
         ("A", "a"),
         ("B", "b"),
@@ -140,6 +145,9 @@ def _fwd_arg_names(*, is_3d: bool = False) -> List[ArgSpec]:
     # k -> ([z,] y, x, c): divisors cpg, X[, Y].
     out += _magic("k_", (["Y"] if is_3d else []) + ["X", "cpg"])
     out += _grid()
+    if split_k:
+        out += [("ws_ptr", "f32*"), ("ws_bytes", "i32")]
+        out += [("ks", "i32"), ("ks_count", "i32")]
     return out
 
 
@@ -313,17 +321,18 @@ def conv_arg_names(
     callers that already dispatch on it (the manifest emitter, the grouped
     dispatcher, the AOT cache) stop branching.
 
-    ``two_stage`` belongs to wgrad's deterministic epilogue and is rejected
-    elsewhere rather than ignored: silently accepting it would hand back a
-    signature that does not match any kernel the flag can build.
+    ``two_stage`` is wgrad's deterministic epilogue, and on fwd the split-K
+    binary (fp32 workspace slabs + a Stage 2 reduce). dgrad has neither, so it
+    is rejected there rather than ignored: silently accepting it would hand
+    back a signature that does not match any kernel the flag can build.
 
     ``streamk`` (wgrad only) is the stream-K reduction, or ``None`` for a
     kernel without stream-K. Stream-K partitions K itself, so it excludes the
     two-stage split-K variant.
     """
-    if two_stage and direction != "wgrad":
+    if two_stage and direction not in ("wgrad", "fwd"):
         raise ValueError(
-            f"two_stage is a wgrad-only variant (got direction={direction!r})"
+            f"two_stage is a wgrad / fwd split-K variant (got direction={direction!r})"
         )
     if streamk is not None:
         if direction != "wgrad":
@@ -338,7 +347,7 @@ def conv_arg_names(
         if two_stage:
             raise ValueError("streamk and two_stage are exclusive: both partition K_wg")
     if direction == "fwd":
-        return _fwd_arg_names(is_3d=is_3d)
+        return _fwd_arg_names(is_3d=is_3d, split_k=two_stage)
     if direction == "wgrad":
         return _wgrad_arg_names(is_3d=is_3d, two_stage=two_stage, streamk=streamk)
     if direction == "dgrad":

dnn-providers/hip-kernel-provider/rocke/library/kernels/common/conv_args.py

@@ -422,7 +422,8 @@ class ConvArgs:
     # mismatched by accident.
     tile_m: int = 0
     tile_n: int = 0
-    # wgrad only: the granularity the split-K slice width is rounded to.
+    # wgrad, and the fwd split-K binary: the granularity the split-K slice
+    # width is rounded to.
     tile_k: int = 0
 
     def __post_init__(self) -> None:
@@ -622,11 +623,11 @@ def _named_values(self) -> Dict[str, int]:
 
     def _variant_values(self, split_k: int) -> Dict[str, int]:
         """Scalars whose meaning depends on the variant, not just the name."""
-        if self.algorithm == "direct" or self.direction != "wgrad":
+        if self.algorithm == "direct" or self.direction == "dgrad":
             if split_k != 1:
                 raise ValueError(
-                    f"split_k={split_k}: only implicit-GEMM wgrad splits the "
-                    f"reduction ({self.algorithm} {self.direction} has no "
+                    f"split_k={split_k}: only implicit-GEMM wgrad and fwd split "
+                    f"the reduction ({self.algorithm} {self.direction} has no "
                     f"split-K arguments)"
                 )
         if self.algorithm == "direct":
@@ -638,10 +639,21 @@ def _variant_values(self, split_k: int) -> Dict[str, int]:
             # The tail past wg_K reads zero through the descriptor's bounds,
             # which is what lets all ``split_k`` slices share one uniform
             # bound computation.
-            values["ks"] = ceil_div(self.gemm_k, self.tile_k * split_k) * self.tile_k
-            values["ks_count"] = split_k
+            values.update(self._split_values(split_k))
         return values
 
+    def _split_values(self, split_k: int) -> Dict[str, int]:
+        """``ks`` / ``ks_count``: the split-K slice width and degree."""
+        if self.tile_k <= 0:
+            raise ValueError(
+                "a split-K launch needs the kernel's tile_k: the slice width is "
+                "rounded to it"
+            )
+        return {
+            "ks": ceil_div(self.gemm_k, self.tile_k * split_k) * self.tile_k,
+            "ks_count": split_k,
+        }
+
     def _divisor(self, name: str) -> int:
         g = self.geom
         return {
@@ -758,6 +770,14 @@ def to_launch_values(
         supplied = {k: v for k, v in supplied.items() if v is not None}
         named = self._named_values()
         named.update(self._variant_values(split_k))
+        if self.direction == "fwd" and "ks" in names:
+            # The split-K forward binary (the ABI has its workspace pair).
+            named.update(self._split_values(split_k))
+        elif split_k != 1 and self.direction == "fwd":
+            raise ValueError(
+                f"split_k={split_k}: a forward split-K launch needs the "
+                f"workspace (ws_ptr / ws_bytes)"
+            )
         if streamk is not None:
             named.update(self._streamk_values(streamk))
         for name in names[6:]:

dnn-providers/hip-kernel-provider/rocke/library/kernels/common/conv_fwd_split_k_reduce.py

@@ -0,0 +1,158 @@
+# Copyright (c) Advanced Micro Devices, Inc., or its affiliates.
+# SPDX-License-Identifier: MIT
+
+"""Stage 2 of the split-K forward convolution: fold the slabs and cast.
+
+Stage 1 (``ImplicitGemmConvSpec`` with ``split_k > 1``) stores each K slice's
+fp32 partial output into its own slab of a ``[ks_count][M][K]`` workspace --
+plain stores, no atomics, every element written exactly once. This kernel sums
+the ``ks_count`` slabs element-wise in a fixed order and converts the result to
+``dtype_d`` into the output tensor ``D`` (N[D]HWK, the same packed layout as one
+slab). The result is deterministic.
+
+The degree is a kernarg, so one binary per (``dtype_d``, ``vec``) serves every
+split. Each thread owns ``vec`` consecutive output elements: one vector load per
+slab, ``vec`` adds, one vector store. ``vec`` must divide the element count
+(``M * K``), which a K divisible by ``vec`` guarantees.
+
+Kernel signature::
+
+    ws_ptr   : f32 global ptr, readonly    -- workspace [ks_count * elems]
+    d_ptr    : dtype_d global ptr          -- output [elems]
+    elems    : i32 -- elements per slab (M * K)
+    ws_bytes : i32 -- workspace byte size (ABI boundary)
+    d_bytes  : i32 -- output byte size (ABI boundary)
+    ks_count : i32 -- number of slabs (the split-K degree)
+
+Grid: ``(ceil(elems / (vec * block_size)), 1, 1)``; block ``(block_size, 1, 1)``.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Tuple
+
+from rocke.core.ir import F32, I32, IRBuilder, KernelDef, PtrType
+from rocke.helpers.io import io_ir_type, pack_f32_to, store_scalar_from_f32
+from rocke.helpers.spec import SignatureBuilder, ceil_div, kernel_name_join
+
+
+@dataclass(frozen=True)
+class FwdSplitKReduceSpec:
+    """One Stage 2 binary: the output dtype and elements per thread.
+
+    ``degree`` > 0 specialises the fold for exactly that many slabs: the slab
+    loop is unrolled, so all of a thread's loads issue back to back instead of
+    one per dependent loop trip (``ks_count`` must then equal it). 0 keeps the
+    runtime loop over ``ks_count``.
+    """
+
+    dtype_d: str = "bf16"
+    vec: int = 4
+    block_size: int = 256
+    degree: int = 0
+    name: str = "conv_fwd_splitk_reduce"
+
+    def kernel_name(self) -> str:
+        parts = [f"v{self.vec}", f"b{self.block_size}"]
+        if self.degree:
+            parts.append(f"s{self.degree}")
+        return kernel_name_join(self.name, *parts, self.dtype_d)
+
+    def validate(self) -> None:
+        if self.vec not in (1, 2, 4):
+            raise ValueError(f"vec must be 1, 2 or 4 (got {self.vec})")
+        if self.dtype_d not in ("fp16", "bf16"):
+            raise ValueError(f"dtype_d must be fp16 or bf16 (got {self.dtype_d})")
+        if self.degree < 0:
+            raise ValueError(f"degree must be >= 0 (got {self.degree})")
+
+
+def build_conv_fwd_split_k_reduce(
+    spec: FwdSplitKReduceSpec, *, arch: str = "gfx950"  # noqa: ARG001
+) -> KernelDef:
+    spec.validate()
+    vec = spec.vec
+    b = IRBuilder(spec.kernel_name())
+    b.kernel.attrs["max_workgroup_size"] = spec.block_size
+
+    ws_ptr = b.param(
+        "ws_ptr", PtrType(F32, "global"), noalias=True, readonly=True, align=16
+    )
+    d_ptr = b.param(
+        "d_ptr",
+        PtrType(io_ir_type(spec.dtype_d), "global"),
+        noalias=True,
+        writeonly=True,
+        align=16,
+    )
+    elems = b.param("elems", I32)
+    _ws_bytes = b.param("ws_bytes", I32)  # noqa: F841 -- ABI boundary
+    _d_bytes = b.param("d_bytes", I32)  # noqa: F841 -- ABI boundary
+    ks_count = b.param("ks_count", I32)
+
+    gid = b.add(b.mul(b.block_id_x(), b.const_i32(spec.block_size)), b.thread_id_x())
+    idx = b.mul(gid, b.const_i32(vec))
+
+    def _load(off):
+        if vec == 1:
+            return [b.global_load_f32(ws_ptr, off)]
+        v = b.global_load_vN(ws_ptr, off, F32, vec)
+        return [b.vec_extract(v, i) for i in range(vec)]
+
+    with b.scf_if(b.cmp_lt(idx, elems)):
+        if spec.degree:
+            # Unrolled: every slab's load is issued before the first add.
+            slabs = [
+                _load(idx if s == 0 else b.add(b.mul(b.const_i32(s), elems), idx))
+                for s in range(spec.degree)
+            ]
+            total = slabs[0]
+            for part in slabs[1:]:
+                total = [b.fadd(a, p) for a, p in zip(total, part)]
+        else:
+            zero = b.const_f32(0.0)
+            init = [(f"acc{i}", zero) for i in range(vec)]
+            loop = b.scf_for_iter(
+                b.const_i32(0), ks_count, b.const_i32(1), init, iv_name="slab"
+            )
+            with loop as (s, accs):
+                parts = _load(b.add(b.mul(s, elems), idx))
+                b.scf_yield(*[b.fadd(a, p) for a, p in zip(accs, parts)])
+            total = list(loop.results)
+        if vec == 1:
+            store_scalar_from_f32(b, d_ptr, idx, total[0], dtype=spec.dtype_d)
+        else:
+            b.global_store_vN(
+                d_ptr, idx, pack_f32_to(b, total, dtype=spec.dtype_d), vec
+            )
+    return b.kernel
+
+
+def fwd_split_k_reduce_grid(
+    spec: FwdSplitKReduceSpec, elems: int
+) -> Tuple[int, int, int]:
+    return (max(1, ceil_div(elems, spec.vec * spec.block_size)), 1, 1)
+
+
+def fwd_split_k_reduce_signature(spec: FwdSplitKReduceSpec) -> list:
+    return (
+        SignatureBuilder()
+        .ptr("ws_ptr", "fp32")
+        .ptr("d_ptr", spec.dtype_d)
+        .scalar("elems", "i32")
+        .scalar("ws_bytes", "i32")
+        .scalar("d_bytes", "i32")
+        .scalar("ks_count", "i32")
+        .build()
+    )
+
+
+def fwd_split_k_reduce_vec(elems: int) -> int:
+    """The widest ``vec`` that divides ``elems``."""
+    return next(v for v in (4, 2, 1) if elems % v == 0)
+
+
+def fwd_split_k_workspace_nbytes(M: int, K: int, ks_count: int) -> int:
+    """Bytes of the ``[ks_count][M][K]`` fp32 workspace."""
+    return 4 * ks_count * M * K