feat(rocke): gfx1250 direct wgrad kernels with two-stage reduction

dnn-providers/hip-kernel-provider/ROCKE_ENGINE_test_categories_external.yaml

@@ -61,6 +61,7 @@ test_categories:
       - "rocke_fp8_scale_conversion_gate"
       - "rocke_mx_scale_decode"
       - "rocke_buffer_rsrc_hip"
+      - "rocke_buffer_atomic_fadd"
       - "rocke_tiled_attention_2d_reentrancy"
       - "rocke_unified_attention_scalar_lower_to_llvm"
       - "rocke_python_import_smoke"
@@ -106,6 +107,7 @@ test_categories:
       - "rocke_fp8_scale_conversion_gate"
       - "rocke_mx_scale_decode"
       - "rocke_buffer_rsrc_hip"
+      - "rocke_buffer_atomic_fadd"
       - "rocke_tiled_attention_2d_reentrancy"
       - "rocke_unified_attention_scalar_lower_to_llvm"
       - "rocke_python_import_smoke"

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

@@ -35,6 +35,7 @@
 import dataclasses
 import itertools
 import os
+import struct
 import sys
 from dataclasses import dataclass
 from typing import List
@@ -1305,6 +1306,7 @@ def _run_wgrad_sweep(
     """Benchmark the direct wgrad kernel."""
     import torch
 
+    from rocke.core.arch import ArchTarget
     from kernels.common.conv_abi import conv_direct_args_signature
     from kernels.common.conv_args import ConvArgs
     from kernels.common.conv_direct_grouped import (
@@ -1356,12 +1358,15 @@ def _run_wgrad_sweep(
     _WAVES = sorted({w for w in _WAVES if w[1] <= n_c_tiles_1})
     _HPB = [30, 60, 120]
     _MK = [32]
+    # wave64 MFMA on gfx950, wave32 WMMA on gfx1250: the block-size cap and the
+    # spec both follow the arch.
+    _wave = ArchTarget.from_gfx(arch).wave_size
     combos = [
         (wk, wc, wq, hpb, mk)
         for wk, wc, wq in _WAVES
         for hpb in _HPB
         for mk in _MK
-        if wk * wc <= 16 and wk * wc * wq * 64 <= 1024
+        if wk * wc <= 16 and wk * wc * wq * _wave <= 1024
     ]
 
     print(
@@ -1380,6 +1385,7 @@ def _run_wgrad_sweep(
             waves_q=waves_q,
             ho_per_block=hpb,
             mfma_k=mk,
+            wave_size=_wave,
         )
         ok, _ = is_valid_wgrad_spec(spec, arch=arch)
         if not ok:
@@ -2253,10 +2259,15 @@ def bind(plan):
             return None
 
         # Each step packed once: a timed loop that packs per launch times the
-        # host, not a short kernel.
-        bound = [
-            launcher.bind(values, config=cfg) for launcher, values, cfg in launches
-        ]
+        # host, not a short kernel. A kernarg past its packed width (a byte
+        # extent over 2^31 on an i32 field) skips the plan, not the sweep.
+        try:
+            bound = [
+                launcher.bind(values, config=cfg) for launcher, values, cfg in launches
+            ]
+        except struct.error as e:
+            print(f"  [skip] {plan.label}: kernarg does not fit ({e})", flush=True)
+            return None
 
         def run(bound=bound):
             for launch in bound:
@@ -2269,6 +2280,10 @@ def verify(plan, run) -> bool:
         nonlocal rc
 
         rt.memset(ptrs[out_name], 0, Out_t.nbytes)
+        # Two-stage wgrad scratch accumulates too: zeroed again for the
+        # verified launch, as the target is.
+        for name in plan.zeroed:
+            rt.memset(ptrs[name], 0, sizes[name])
         run()
         synchronize_and_release(0)
         out_cpu = torch.empty_like(Out_t)
@@ -2466,6 +2481,7 @@ def _cache_dispatch(args, arch: str, cases: list) -> int:
             "direct_dgrad_helper",
             "direct_fwd_helper",
             "direct_wgrad",
+            "direct_wgrad_helper",
         ),
     )
     if rc:
@@ -2478,13 +2494,14 @@ def _cache_dispatch(args, arch: str, cases: list) -> int:
 
 def _batch_chunks(p, direction: str) -> List[int]:
     """Images per launch: the fewest balanced batch chunks that keep every
-    activation under 2 GiB (``[N]`` when one launch fits, and for wgrad).
+    activation under 2 GiB (``[N]`` when one launch fits).
 
     The chunks need not divide N; they differ by at most one image (N=31 that
-    fits 15 images per launch runs as 11 + 10 + 10).
+    fits 15 images per launch runs as 11 + 10 + 10). wgrad reduces over N, so
+    its chunks are not independent launches, but they can accumulate: every
+    chunk adds into the one dW, zeroed once, which only a kernel that adds
+    into dW can do (:func:`_accumulates_dw`).
     """
-    if direction == "wgrad":
-        return [p.N]
     elem = 2
     per_image = max(p.H * p.W * p.total_c, p.Ho * p.Wo * p.total_k) * elem
     per_launch = (_TENSOR_LIMIT - 1) // per_image
@@ -2502,7 +2519,9 @@ def _chunk_buffers(ptrs, sizes, inputs, out_name, chunks, i):
         return ptrs, sizes
     ptrs, sizes = dict(ptrs), dict(sizes)
     n_total, n0 = sum(chunks), sum(chunks[:i])
-    for name in [n for n in inputs if n != "w"] + [out_name]:
+    # wgrad's dW has no batch: every chunk adds into the whole of it.
+    batched = [out_name] if out_name != "dw" else []
+    for name in [n for n in inputs if n != "w"] + batched:
         per_image = sizes[name] // n_total
         ptrs[name] = ptrs[name] + n0 * per_image
         sizes[name] = chunks[i] * per_image
@@ -2537,6 +2556,31 @@ def _chunk_plans(plan_fn, p, chunks):
     return plans, rejected, per_chunk
 
 
+def _accumulates_dw(plan) -> bool:
+    """A one-step wgrad plan with no scratch: its kernel adds its tiles into
+    the zeroed fp32 dW (atomics), so its launches over batch chunks sum. A fold
+    or reduce step, or a scratch slab, writes dW instead."""
+    return len(plan.steps) == 1 and not plan.workspaces and not plan.zeroed
+
+
+def _wgrad_chunk_plans(plans, rejected, per_chunk, chunks):
+    """``(plans, rejected)`` for a wgrad problem run as ``chunks``: past one
+    chunk, only the plans whose every chunk accumulates into dW."""
+    if len(chunks) == 1:
+        return plans, rejected
+    keep = [pl for pl in plans if all(map(_accumulates_dw, per_chunk[pl.label]))]
+    rejected = list(rejected) + [
+        (
+            pl.identity,
+            f"a wgrad batch run as {len(chunks)} chunks needs a kernel that "
+            f"accumulates into dW",
+        )
+        for pl in plans
+        if pl not in keep
+    ]
+    return keep, rejected
+
+
 def _run_from_cache(args, arch: str, cases: list, cache, directions) -> int:
     """Benchmark every cached kernel that can run each case. Nothing is compiled.
 
@@ -2599,6 +2643,8 @@ def _u8(t):
                 p,
                 chunks,
             )
+            if direction == "wgrad":
+                plans, rejected = _wgrad_chunk_plans(plans, rejected, per_chunk, chunks)
             if not plans:
                 why = _no_plan_reason(rejected, dtype, variants)
                 print(
@@ -2665,6 +2711,10 @@ def _u8(t):
                         sizes[name] = max(sizes.get(name, 0), nbytes)
             for name in set(sizes) - set(ptrs):
                 ptrs[name] = rt.alloc(sizes[name])
+            # Two-stage wgrad scratch: zeroed like dW (once for timing, again
+            # before each verified run).
+            for name in {n for plan in plans for n in plan.zeroed}:
+                rt.memset(ptrs[name], 0, sizes[name])
 
             flop = float(p.flops)
             bytes_xfer = float(sum(t.nbytes for t in inputs.values()) + Out_t.nbytes)

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

@@ -4783,6 +4783,20 @@ def _u8(t):
                         flush=True,
                     )
                     continue
+            if direction == "wgrad" and len(batch_chunks) > 1:
+                # A batch too large for the kernels' i32 addressing runs as
+                # slices that accumulate into dW: only the kernels that add
+                # into dW can serve it.
+                candidates = [c for c in candidates if _wgrad_accumulates(c[0])]
+                if not candidates:
+                    print(
+                        f"Case {case_idx} {problem.short()} {case_dtype} wgrad: "
+                        f"no compatible kernel (the batch runs as "
+                        f"{len(batch_chunks)} slices, which needs a split-K "
+                        f"kernel that accumulates into dW)",
+                        flush=True,
+                    )
+                    continue
             if direction == "wgrad" and args.two_stage != "always":
                 # Same rule as the JIT sweep's --two-stage: by default the
                 # two-stage (scratch + reduce) kernels only run when C/groups
@@ -6090,9 +6104,14 @@ def _batch_chunks(direction, problem, dtype) -> list:
     sentinel is the i32 maximum, so no operand it is launched on may exceed
     2^31 - 1 bytes. fwd and dgrad compute every image independently, and N is
     a kernarg: a batch whose activations are larger runs as consecutive
-    launches of the same binary over equal slices of N, at offset pointers.
-    wgrad reduces over N, so it cannot be split that way; neither can a weight
-    tensor, which has no batch.
+    launches of the same binary over equal slices of N, at offset pointers. A
+    weight tensor, which has no batch, cannot be split.
+
+    wgrad reduces over N, so its slices are not independent launches -- but
+    they can accumulate: a kernel that adds into dW (the packed-atomic split-K
+    family), run over consecutive slices with dW zeroed once before the first,
+    computes the whole reduction. The caller offers only such kernels when
+    this returns more than one slice (see :func:`_wgrad_accumulates`).
 
     Raises ``ValueError`` for a problem no split can bring into range.
     """
@@ -6108,12 +6127,6 @@ def _batch_chunks(direction, problem, dtype) -> list:
         )
     if per_image * problem.N <= _INT32_MAX:
         return [(0, problem.N)]
-    if direction == "wgrad":
-        raise ValueError(
-            f"the activations ({per_image * problem.N} bytes) exceed the "
-            f"kernels' 2^31-byte operand addressing, and wgrad reduces over the "
-            f"batch, so it cannot run as batch slices"
-        )
     if per_image > _INT32_MAX:
         raise ValueError(
             f"one image's activations ({per_image} bytes) exceed the kernels' "
@@ -6124,6 +6137,17 @@ def _batch_chunks(direction, problem, dtype) -> list:
     return [(n0, min(step, problem.N - n0)) for n0 in range(0, problem.N, step)]
 
 
+def _wgrad_accumulates(ident) -> bool:
+    """Whether a cached wgrad binary adds into dW rather than storing it.
+
+    Only those can run a wgrad problem as batch slices (:func:`_batch_chunks`):
+    the packed-atomic split-K family accumulates into a dW zeroed once. The
+    non-split kernel stores dW, and two-stage and stream-K finish with a
+    reduce that overwrites it, so each slice would erase the previous ones.
+    """
+    return ident.split_k > 1 and not ident.two_stage and ident.streamk == "off"
+
+
 def _cached_bind_batch(batch_chunks, direction, problem, ident, *rest):
     """:func:`_cached_bind` over the batch slices of :func:`_batch_chunks`.
 
@@ -6135,9 +6159,28 @@ def _cached_bind_batch(batch_chunks, direction, problem, ident, *rest):
     if len(batch_chunks) == 1:
         return _cached_bind(direction, problem, ident, *rest)
     (meta, rt, launcher, case_dtype, split_k, ptrs, sizes, *tail) = rest
-    a_image, d_image = sizes[0] // problem.N, sizes[2] // problem.N
     runs = []
+    if direction == "wgrad":
+        # A (dY) and B (X) are the activations, D (dW) the reduction target:
+        # every slice adds into the whole dW, which only the first one zeroes.
+        if not _wgrad_accumulates(ident):
+            print(
+                f"  [skip] {ident.label()}: a wgrad problem split into "
+                f"{len(batch_chunks)} batch slices needs a kernel that "
+                f"accumulates into dW",
+                flush=True,
+            )
+            return None
+        a_image, b_image = sizes[0] // problem.N, sizes[1] // problem.N
+    else:
+        a_image, d_image = sizes[0] // problem.N, sizes[2] // problem.N
     for n0, n in batch_chunks:
+        if direction == "wgrad":
+            sl_ptrs = (ptrs[0] + n0 * a_image, ptrs[1] + n0 * b_image, ptrs[2])
+            sl_sizes = (n * a_image, n * b_image, sizes[2])
+        else:
+            sl_ptrs = (ptrs[0] + n0 * a_image, ptrs[1], ptrs[2] + n0 * d_image)
+            sl_sizes = (n * a_image, sizes[1], n * d_image)
         bound = _cached_bind(
             direction,
             dc_replace(problem, N=n),
@@ -6147,9 +6190,10 @@ def _cached_bind_batch(batch_chunks, direction, problem, ident, *rest):
             launcher,
             case_dtype,
             split_k,
-            (ptrs[0] + n0 * a_image, ptrs[1], ptrs[2] + n0 * d_image),
-            (n * a_image, sizes[1], n * d_image),
+            sl_ptrs,
+            sl_sizes,
             *tail,
+            zero_dw=n0 == 0,
         )
         if bound is None:
             return None
@@ -6178,9 +6222,14 @@ def _cached_bind(
     KernelLauncher,
     LaunchConfig,
     live_extras: list,
+    zero_dw: bool = True,
 ):
     """``(run, label)`` of a cached (non-stream-K) binary at ``split_k``.
 
+    ``zero_dw=False`` leaves out the dW clear of the packed-atomic split-K
+    wgrad launch: the later batch slices of a sliced wgrad problem add into
+    the dW the first slice zeroed (:func:`_cached_bind_batch`).
+
     ``ptrs`` / ``sizes`` are the A, B, D buffers. The device scratch the
     launch config allocates is appended to ``live_extras`` (freed by
     :func:`_free_extras`). Returns None, after saying why, when the binary
@@ -6228,7 +6277,8 @@ def _cached_bind(
     # queues ahead of this kernel's launches rather than draining for them.
     stream = resolve_stream(cfg.stream)
     if (
-        direction == "wgrad"
+        zero_dw
+        and direction == "wgrad"
         and ident.split_k > 1
         and not ident.two_stage
         and ident.streamk == "off"