feat(rocke): async K loop, cshuffle, fp16 split-K and two-stage split-K for gfx1250 WMMA wgrad

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

@@ -1390,7 +1390,8 @@ def main() -> int:
             "two-stage deterministic wgrad pipeline (Stage1 GEMM → workspace → Stage2 reduce), "
             "also for --run-from-cache: "
             "auto = two-stage only when C/groups is odd (where the packed atomic "
-            "cannot be built); "
+            "cannot be built) -- and, from the cache on wave32, for bf16 (whose "
+            "packed atomic is degree-capped for accuracy); "
             "always = time the two-stage kernels alongside the packed-atomic ones "
             "(nothing is dropped), so both split-K reductions compete; "
             "never = skip two-stage entirely (default: auto)"
@@ -4790,7 +4791,15 @@ def _u8(t):
                 # "never" drops them outright; "always" keeps them alongside
                 # the packed-atomic ones, so both reductions compete.
                 odd = (problem.C // max(1, problem.groups)) % 2 != 0
-                if args.two_stage == "never" or not odd:
+                # On wave32 (gfx1250) bf16 also runs the two-stage kernels: the
+                # packed bf16 atomic is degree-capped for accuracy (12), and the
+                # f32 workspace is the accuracy-safe way past that cap.
+                from rocke.core.arch import ArchTarget
+
+                deep_bf16 = (
+                    case_dtype == "bf16" and ArchTarget.from_gfx(arch).wave_size == 32
+                )
+                if args.two_stage == "never" or not (odd or deep_bf16):
                     candidates = [c for c in candidates if not c[0].two_stage]
                 if not candidates:
                     print(
@@ -5512,10 +5521,13 @@ def _widest_vector_sizes(candidates):
     thread) only in narrower copies. Narrower access does not win, so only the
     kernels at the problem's widest combination are run.
 
-    "Widest" is taken per kernel kind -- direction, epilogue, split-K mode and
-    group merging -- because the kind caps some widths regardless of the
-    problem: the default epilogue stores scalar (c = 1), and a group-merged
-    kernel derives its widths in the builder (recorded as 0). Default-epilogue
+    "Widest" is taken per kernel kind -- direction, epilogue, split-K mode,
+    group merging and the K loop (synchronous or async-DMA) -- because the
+    kind caps some widths regardless of the problem: the default epilogue
+    stores scalar (c = 1), a group-merged kernel derives its widths in the
+    builder (recorded as 0), and an async-DMA kernel records the chunk widths
+    its copies use where a synchronous one records its width caps, which its
+    loader can narrow at build time, so the two do not compare. Default-epilogue
     kernels are dropped altogether where a kernel of the same kind can store
     wider than 1 for this problem.
     ``supports_problem`` has already dropped widths the problem cannot take,
@@ -5533,15 +5545,22 @@ def kind(ident):
             ident.group_merge,
             ident.streamk,
             ident.streamk_reduction,
+            bool(getattr(ident, "async_dma", False)),
         )
 
     def rank(ident):
-        return (
-            ident.vector_size_a + ident.vector_size_b + ident.vector_size_c,
-            ident.vector_size_a,
-            ident.vector_size_b,
-            ident.vector_size_c,
-        )
+        # An async-DMA binary leaves A / B unset and records the chunk widths
+        # its copies use instead; it ranks by those against the other async
+        # binaries of its kind (kind() keeps it apart from the synchronous
+        # kernels, whose recorded caps are not their effective load widths).
+        a = ident.vector_size_a
+        b = ident.vector_size_b
+        if a is None:
+            a = getattr(ident, "async_chunk_a", 0) or 8
+        if b is None:
+            b = getattr(ident, "async_chunk_b", 0) or 8
+        c = ident.vector_size_c or 0
+        return (a + b + c, a, b, c)
 
     # The default epilogue stores scalar (c = 1) straight from the MMA
     # accumulators. Where the same kind of kernel can store wider -- a

dnn-providers/hip-kernel-provider/rocke/library/kernels/common/README_conv_implicit_gemm_wgrad.md

@@ -45,22 +45,59 @@ On MFMA targets, `epilogue="cshuffle"` is **required** for bf16/fp16 output when
 `split_k > 1`. The direct-atomic path emits zero-filled packed atomics at the
 scattered MFMA layout; the cshuffle path produces genuinely contiguous adjacent
 pairs after the LDS shuffle, which is necessary for correct `<2 x dtype>` packed
-atomics. WMMA has no cshuffle epilogue; on arches whose `ArchTarget` sets
-`memory.has_global_float_atomic_add` (gfx1250) its split-K epilogue places each
-value in its half of the pair itself, so `default` is the split-K epilogue there.
-Other wave32 arches keep `split_k=1`, and no WMMA arch has the two-stage path.
+atomics. On WMMA arches whose `ArchTarget` sets
+`memory.has_global_float_atomic_add` (gfx1250) the `default` split-K epilogue
+places each value in its half of the pair itself, and `cshuffle` pairs through
+LDS from the WMMA fragment; the two-stage path decodes the WMMA fragment
+through `op.c_layout()`. Other wave32 arches keep `split_k=1` and have no
+two-stage path.
 
 ---
 
 ## Changelog
 
+### Async K loop on WMMA (gfx1250)
+
+The WMMA K loop was single-buffered with no overlap: load, barrier, WMMA,
+barrier. `async_dma` now covers it on gfx1250, mirrored in the C++ engine.
+`pipeline` stays `mem`, and `unroll_k` / `prefetch_addr` stay rejected on WMMA.
+
+- `async_dma` on gfx1250 is a ping-pong over `global_load_async_to_lds`, because
+  gfx1250 has no gfx9 `buffer_load_lds`.
+  - It reuses the synchronous K-outer loaders for the slot partition, and keeps
+    the bank-friendly row pad of 8, since the copy writes a per-lane LDS
+    address.
+  - Each phase waits `s_wait_asynccnt 0` and takes one LDS barrier, then issues
+    the next tile while the current tile's WMMAs run.
+  - The copy has no buffer-resource range check. The loader therefore ANDs that
+    check (byte offset below the tensor size, compared unsigned as the
+    resource does, so an offset at or past 2 GiB is out of range) into the
+    descriptor's `valid`, and zero-fills failing chunks with an LDS store of
+    the tile's element type (fp16 or bf16).
+  - It needs a 4-, 8- or 16-byte chunk on both operands, so an odd channel run
+    is rejected.
+
+### cshuffle epilogue on WMMA (gfx1250)
+
+The WMMA C fragment gives each lane one dW column of 8 rows. The default
+epilogue therefore writes a 16-bit split-K partial as one zero-partner packed
+atomic per element, so every dword takes two atomics. `epilogue="cshuffle"` now
+admits WMMA. It stages the tile through LDS with `CShuffleEpilogue.from_grid_op`,
+the WMMA path the forward and dgrad epilogues already use, and then issues
+genuine `<2 x dtype>` pairs from contiguous runs: half the atomics, still
+coalesced across lanes.
+
+Swapping the WMMA operands instead was tried and dropped. That swap gives each
+lane a contiguous run, but it spreads the 32 lanes of one instruction over 32
+dW rows, and the extra memory transactions outweigh the halved atomic count.
+
 ### Split-K on gfx1250 (WMMA)
 
 - 16-bit and fp32 split-K are enabled on gfx1250 through the WMMA atomic
   epilogue (`_emit_wgrad_split_k_epilogue_wmma`, mirrored in the C++ engine).
   Both engines admit WMMA split-K only where the arch catalog sets
-  `memory.has_global_float_atomic_add`, and only inside the WMMA wgrad subset (`mem` pipeline, `default` epilogue, 16x16x16 or
-  16x16x32 atom).
+  `memory.has_global_float_atomic_add`, and only inside the WMMA wgrad subset
+  (`default` or `cshuffle` epilogue, 16x16x16 or 16x16x32 atom).
 - Degree (`WgradSplitPolicy` in `rocke/helpers/split_k.py`): the smaller of a
   CTA-fill term and a slice-depth term, under a per-dtype cap. Each packed
   16-bit atomic rounds the running sum in dW, so with same-sign partials the
@@ -69,20 +106,32 @@ Other wave32 arches keep `split_k=1`, and no WMMA arch has the two-stage path.
   degree that keeps that bound under the suite's 5e-2 tolerance
   (`atomic_rounding_degree_cap`); higher degrees miss it on all-positive
   inputs. fp32 takes the general cap of 128.
-- fp16 is not split automatically: its cap is 1, so `--split-k -1` (and every
-  auto caller) runs it unsplit. The fp16 packed atomic does not scale like
-  bf16's on gfx1250. In a degree sweep the best fp16 degree ranged from 16 on
-  a tiny dW down to none at all, where every degree was slower than
-  `split_k=1`. No single cap avoids that loss. The split-K binaries are still
-  built, and an explicit degree launches them for callers whose shape gains.
-  The two-stage f32 path is the planned fp16 reduction.
+- fp16 has its own policy: fill to 1024 CTAs (2048 grouped) with K slices
+  down to 512, up to degree 97. Earlier fp16 was not split automatically
+  (cap 1), because the packed fp16 atomic was lowered as a fenced atomic
+  (`global_wb` / `global_inv` around every add). It is now an agent-scope
+  monotonic `atomicrmw`, like bf16's, and the policy was refitted on a degree
+  sweep of the cached split-K binaries. The cap of 97 is the accuracy bound
+  for fp16's 10 fraction bits (unit roundoff `2**-11`), as bf16's 12 is for
+  its 7: no fp16 degree failed the 5e-2 verification on the sweep, but at 128
+  all-positive inputs can miss it (one 1.0 partial followed by 127 partials of
+  `2**-11` rounds back to 1.0, a 5.8% error). The cap binds only the packed
+  atomic; the f32 workspace (two-stage) degree is resolved as fp32 dW.
 - Reachability: as on gfx950, wgrad has no dispatcher instance list; it is
   served from the AOT cache. The cache grid (`kernels/conv/aot_jobs.py`) builds
   one split-K binary per geometry, and the degree is resolved at launch by
   `select_split_k_wgrad` with the dW dtype (`benchmark_implicit_gemm_conv.py
   --run-from-cache --split-k -1`). A 16-bit dW whose row length (`Z*Y*X*cpg`) is
-  odd cannot use the packed atomic, and WMMA has no two-stage path to fall back
-  on, so such problems run at `split_k=1`.
+  odd cannot use the packed atomic; it reaches split-K through the two-stage
+  kernels (below).
+- Two-stage (f32 workspace) split-K on gfx1250: the Stage 1 scratch epilogue
+  decodes the WMMA C fragment through `op.c_layout()`, as the WMMA atomic
+  epilogue does, mirrored in the C++ engine. The workspace reduction
+  accumulates in fp32, so the split policy resolves its degree with the fp32
+  cap (`conv_split_k`). The AOT grid builds the two-stage kernels
+  (`ws_replicas` = 8) next to the packed-atomic ones, and from the cache on
+  wave32 bf16 runs them by default too (`--two-stage auto`), since the bf16
+  packed atomic is capped at degree 12 for accuracy.
 - Results are not bit-reproducible across runs once `split_k > 1` (atomic
   ordering), and the caller must zero `dW` before launch.
 - Runbook mapping (`platform/dsl_docs/optimization/optimization_runbook.md`)