From 746a9e2fa72f847efded87eb05e3c7a80f6e2875 Mon Sep 17 00:00:00 2001 From: Jakub Piasecki Date: Sun, 1 Nov 2026 17:49:12 +0000 Subject: [PATCH 01/13] feat(rocke): async-to-LDS K loop and cshuffle epilogue for gfx1250 WMMA wgrad The WMMA wgrad K loop was single-buffered with no overlap: load, barrier, WMMA, barrier. On gfx1250, async_dma now runs a one-barrier ping-pong over global_load_async_to_lds (CoalescedTileLoader.load_async_to_lds): - it reuses the synchronous K-outer loaders for the slot partition and keeps the K-outer row pad of 8, since the copy writes a per-lane LDS address (the gfx9 buffer_load_lds path keeps its pad of 0); - 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, which the wgrad descriptors rely on for the reduction tail and the free-axis overhang, so the loader ANDs byte_offset < num_bytes into each chunk's valid predicate and zero-fills failing chunks with an LDS store; - it needs lds_k_outer and a 4-, 8- or 16-byte chunk on both operands. Admit epilogue="cshuffle" on WMMA wgrad. Both wgrad cshuffle epilogues (direct store and split-K atomic_store) read the WMMA fragment through CShuffleEpilogue.from_grid_op, the path the forward and dgrad WMMA epilogues already use, so a 16-bit split-K partial is written as genuine <2 x dtype> pairs instead of one zero-partner packed atomic per element. Like MFMA's, that LDS pairing needs an even store vector, which both engines' validators now require for WMMA cshuffle split-K (the default WMMA epilogue still ignores it). pipeline stays "mem", and unroll_k / prefetch_addr / chiplet_swizzle stay rejected on WMMA, as before. Both engines are mirrored (rocke_conv_emit_kloop_async_global_lds and rocke_coalesced_tile_loader_load_async_to_lds in C++). Parity configs 55-62 cover the async loop (fp32 dW split-K, partial N tiles, 2x2 warps), the C=3 async reject, and the cshuffle epilogue (direct fp16, split-K bf16, with async at C=24, the gfx1201 reject, and the odd store-vector reject). Config 53 (gfx1250 bf16 split-K with cshuffle) is now admitted by both engines. Two new representative-IR cases (async, cshuffle) are blessed; no existing hash changes. Co-Authored-By: Claude Opus 5.5 --- .../common/README_conv_implicit_gemm_wgrad.md | 37 ++- .../common/conv_implicit_gemm_wgrad.py | 244 +++++++++++++++--- .../library/kernels/common/conv_split_k.py | 8 +- .../parity/conv_implicit_gemm_wgrad_emit.c | 88 +++++++ .../parity/conv_implicit_gemm_wgrad_emit.py | 113 +++++++- .../tests/test_conv_wgrad_correctness.py | 164 +++++++++++- .../rocke/platform/cpp/helpers/loads.cpp | 65 +++++ .../rocke/helper_rocke.helpers.loads.h | 19 ++ .../instance_conv_implicit_gemm_internal.h | 10 + .../conv_implicit_gemm_conv_kloop_drivers.cpp | 78 ++++++ .../conv_implicit_gemm_conv_load_phase.cpp | 26 ++ .../common/conv_implicit_gemm_wgrad.cpp | 173 +++++++++++-- .../dsl_docs/instances/convolution.md | 5 +- .../platform/python/rocke/helpers/loads.py | 61 +++++ .../python/rocke/instances/SUPPORT_MATRIX.md | 2 +- .../conv_wgrad.json | 64 +++++ .../differential/golden/llvm_gfx_all.json | 30 +++ .../instances/rocke_ir_parity_harness.py | 45 ++++ 18 files changed, 1155 insertions(+), 77 deletions(-) diff --git a/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/README_conv_implicit_gemm_wgrad.md b/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/README_conv_implicit_gemm_wgrad.md index bf815c57347a..58ab29545dec 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/README_conv_implicit_gemm_wgrad.md +++ b/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/README_conv_implicit_gemm_wgrad.md @@ -54,13 +54,46 @@ Other wave32 arches keep `split_k=1`, and no WMMA arch has the 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) into the descriptor's `valid`, + and zero-fills failing chunks with an LDS store. + - 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 diff --git a/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/conv_implicit_gemm_wgrad.py b/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/conv_implicit_gemm_wgrad.py index e3f2424279af..0fed277ba131 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/conv_implicit_gemm_wgrad.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/conv_implicit_gemm_wgrad.py @@ -179,6 +179,7 @@ def wgrad_atomic_epilogue_available( dtype_d: str, vector_size_c: Optional[int] = None, wave_size: int = 64, + epilogue: str = "default", ) -> Tuple[bool, str]: """Can split-K reach this problem through the *packed atomic* epilogue? @@ -195,9 +196,10 @@ def wgrad_atomic_epilogue_available( not dword-aligned; and * an even store-vector width -- the epilogue emits ``sv/2`` pairs per thread, so ``sv == 1`` (which is what a depthwise ``cpg == 1`` yields) - produces no pair at all. MFMA (wave64) only: the WMMA epilogue - (``wave_size == 32``) issues one pair per accumulator element, value in - its half and zero in the partner, and never reads the store width. + produces no pair at all. Not for the WMMA default epilogue + (``wave_size == 32``): it issues one pair per accumulator element, value + in its half and zero in the partner, and never reads the store width. + The WMMA ``cshuffle`` epilogue pairs from LDS like MFMA's and needs it. gfx9 has no scalar 16-bit atomic add to fall back on, so when this returns False the only ways to keep ``split_k > 1`` are the two-stage f32-workspace @@ -207,7 +209,8 @@ def wgrad_atomic_epilogue_available( if dtype_d not in ("bf16", "fp16"): return True, "ok" wgN = _wg_N(p) - sv = wgrad_atomic_store_vec(p, dtype_d, vector_size_c) if wave_size != 32 else 2 + wmma_default = wave_size == 32 and epilogue != "cshuffle" + sv = 2 if wmma_default else wgrad_atomic_store_vec(p, dtype_d, vector_size_c) if wgN % 2 != 0 or sv % 2 != 0: return False, ( f"split_k atomic with dtype_d={dtype_d!r} requires an even dW row " @@ -835,7 +838,11 @@ def validate(self) -> None: # cpg=3, X=2) is perfectly safe, and rejecting it needlessly pushed # depthwise-adjacent shapes off split-K entirely. _ok, _why = wgrad_atomic_epilogue_available( - self.problem, self.data.dtype_d, self.vector_size_c, self.wave_size + self.problem, + self.data.dtype_d, + self.vector_size_c, + self.wave_size, + self.epilogue, ) if not _ok: raise ValueError(_why) @@ -1347,6 +1354,86 @@ def wgrad_wmma_split_k_available(arch: str) -> bool: return ArchTarget.from_gfx(arch).memory.has_global_float_atomic_add +def _kouter_pad(spec: "WgradConvSpec") -> int: + """Row pad (elements) of the K-outer LDS tile. + + 8 spreads the transpose read's row groups across banks (see the builder). + The gfx9 direct load (``buffer_load_lds``) deposits lane-contiguous packed + bytes and cannot skip a pad, so its async path uses 0. The gfx1250 async + copy writes a per-lane LDS address and keeps the bank-friendly pad. + """ + return 0 if spec.async_dma and spec.wave_size == 64 else 8 + + +def _wgrad_free_axis_vec(chan: int, dtype: str, override: Optional[int] = None) -> int: + """Auto free-axis load width, clamped by an explicit spec override. + + ``vector_size_a`` / ``vector_size_b`` are a cap, not a replacement: a wider + request than the channel run allows would read across a stride-1 boundary + (wrong data), so the divisibility rule still wins. The forward and dgrad + instances honour the same two fields; wgrad was the only conv family that + declared and ignored them. + """ + widths = (8, 4, 2, 1) if dtype != "fp32" else (4, 2, 1) + auto = next(v for v in widths if chan % v == 0) + return auto if override is None else min(auto, override) + + +def _wmma_async_vecs(spec: "WgradConvSpec") -> Tuple[int, int]: + """The (A, B) K-outer chunk widths the gfx1250 async loader issues. + + The same choose_vec the synchronous K-outer loader makes, so the async and + sync paths load the identical tile partition. + """ + a_vec = CoalescedTileLoader.choose_vec( + tile_rows=spec.tile_k, + tile_cols=spec.tile_m, + block_size=spec.block_size, + max_vec=_wgrad_free_axis_vec(spec.problem.kpg, spec.data.dtype_a), + vector_axis="col", + ) + b_vec = CoalescedTileLoader.choose_vec( + tile_rows=spec.tile_k, + tile_cols=spec.tile_n, + block_size=spec.block_size, + max_vec=_wgrad_free_axis_vec(spec.problem.cpg, spec.data.dtype_b), + vector_axis="col", + ) + return a_vec, b_vec + + +def _wmma_async_available(spec: "WgradConvSpec", arch: str) -> Tuple[bool, str]: + """Whether a wave32 spec can take the gfx1250 async global->LDS K loop. + + gfx1250 has no gfx9 ``buffer_load_lds``; its direct load is the per-lane + ``global_load_async_to_lds_b{32,64,128}``, which moves 4, 8 or 16 bytes. A + 16-bit chunk therefore needs at least two elements, i.e. an even channel + run on both operands. The copy carries no buffer-resource range check, so + out-of-range chunks are zero-filled with a plain LDS store instead. + """ + from rocke.core.arch import ArchTarget + + if not ArchTarget.from_gfx(arch).memory.has_async_global_lds: + return ( + False, + f"WMMA wgrad async_dma needs global_load_async_to_lds (not on {arch})", + ) + if not spec.lds_k_outer: + return False, "WMMA wgrad async_dma requires lds_k_outer=True" + try: + vecs = _wmma_async_vecs(spec) + except ValueError as e: + return False, str(e) + elem_bytes = 4 if spec.data.dtype_a == "fp32" else 2 + for side, v in zip(("A", "B"), vecs): + if v * elem_bytes not in (4, 8, 16): + return False, ( + f"WMMA wgrad async_dma needs a 4/8/16-byte {side} chunk " + f"(global_load_async_to_lds width); got {v} x {elem_bytes} B" + ) + return True, "ok" + + def wgrad_two_stage_available(arch: str) -> bool: """Whether ``arch`` has the two-stage (f32 scratch) split-K wgrad path. @@ -1859,7 +1946,11 @@ def is_valid_wgrad_spec(spec: WgradConvSpec, arch: str = "gfx950") -> Tuple[bool # cpg, and it does not apply when the epilogue is the f32 scratch atomic. if _needs_atomic: _ok, _why = wgrad_atomic_epilogue_available( - spec.problem, spec.data.dtype_d, spec.vector_size_c, spec.wave_size + spec.problem, + spec.data.dtype_d, + spec.vector_size_c, + spec.wave_size, + spec.epilogue, ) if not _ok: return False, _why @@ -1867,10 +1958,10 @@ def is_valid_wgrad_spec(spec: WgradConvSpec, arch: str = "gfx950") -> Tuple[bool # pairs at the scattered MFMA layout, so it needs cshuffle's contiguous # pairs. Two cases are not on it. At split_k == 1 the epilogue is a direct # store, and under two_stage it is an f32 scratch atomic; neither emits - # packed atomics, so 'default' is fine. WMMA rejects cshuffle outright, so - # 'default' is its only epilogue; its split-K (admitted above only on - # wgrad_wmma_split_k_available arches) goes through the WMMA epilogue, which forms - # each packed pair itself. + # packed atomics, so 'default' is fine. WMMA is exempt too: its 'default' + # split-K (admitted above only on wgrad_wmma_split_k_available arches) goes + # through the WMMA epilogue, which forms each packed pair itself with a zero + # partner, and its 'cshuffle' split-K pairs through LDS exactly as on MFMA. # (_effective_two_stage is computed above, with the split-K validity gates.) if ( _is_atomic @@ -2005,7 +2096,7 @@ def is_valid_wgrad_spec(spec: WgradConvSpec, arch: str = "gfx950") -> Tuple[bool # whenever tile_k > tile_mn -- for 32x32x64 that is 1 KB per spec -- so a # spec that overflows the cap passes validation and fails later at # smem_alloc. Keep this in sync with the _KOUTER_PAD block in the builder. - _KOUTER_PAD = 0 if spec.async_dma else 8 + _KOUTER_PAD = _kouter_pad(spec) _a_shape = (spec.tile_k, spec.tile_m + _KOUTER_PAD) _b_shape = (spec.tile_k, spec.tile_n + _KOUTER_PAD) else: @@ -2058,19 +2149,25 @@ def is_valid_wgrad_spec(spec: WgradConvSpec, arch: str = "gfx950") -> Tuple[bool f"WMMA wgrad supports only the 'mem' pipeline " f"(got {spec.pipeline!r}) on {arch}" ) - if spec.epilogue != "default": + # cshuffle reads the accumulators through op.c_layout(), the same WMMA + # path the forward and dgrad cshuffle epilogues use. + if spec.epilogue not in ("default", "cshuffle"): return False, ( - f"WMMA wgrad supports only the 'default' epilogue " + f"WMMA wgrad supports only the 'default' and 'cshuffle' epilogues " f"(got {spec.epilogue!r}) on {arch}" ) for flag, label in ( - (spec.async_dma, "async_dma"), (spec.unroll_k, "unroll_k"), (spec.chiplet_swizzle, "chiplet_swizzle"), (spec.prefetch_addr, "prefetch_addr"), ): if flag: return False, f"WMMA wgrad does not support {label} on {arch}" + # async_dma is the gfx1250 global_load_async_to_lds K loop. + if spec.async_dma: + _async_ok, _async_why = _wmma_async_available(spec, arch) + if not _async_ok: + return False, _async_why return True, "ok" @@ -2438,7 +2535,7 @@ def _declare_wgrad_ptr(name: str, kind: str) -> Value: # Direct load deposits lane-contiguous *packed* bytes and cannot skip a # row pad, so the async path must use a pad of 0. The transpose read is # insensitive to the row stride, so dropping the pad costs nothing here. - _KOUTER_PAD = 0 if spec.async_dma else 8 + _KOUTER_PAD = _kouter_pad(spec) a_kouter_stride = block_m + _KOUTER_PAD b_kouter_stride = block_n + _KOUTER_PAD _a_shape = (block_k, a_kouter_stride) @@ -2497,21 +2594,13 @@ def _declare_wgrad_ptr(name: str, kind: str) -> Value: axis_a = axis_b = "col" load_vec_a = 1 load_vec_b = 1 + _free_axis_vec = _wgrad_free_axis_vec + # gfx1250 async: the per-lane global_load_async_to_lds K loop. It reuses the + # synchronous K-outer loaders for the slot partition and the descriptors, + # and only swaps the global->register->LDS hop for one async copy. + _wmma_async = spec.async_dma and op.family == "wmma" if not spec.async_dma: - def _free_axis_vec(chan: int, dtype: str, override: "int | None" = None) -> int: - """Auto free-axis load width, clamped by an explicit spec override. - - ``vector_size_a`` / ``vector_size_b`` are a cap, not a replacement: - a wider request than the channel run allows would read across a - stride-1 boundary (wrong data), so the divisibility rule still - wins. The forward and dgrad instances honour the same two fields; - wgrad was the only conv family that declared and ignored them. - """ - widths = (8, 4, 2, 1) if dtype != "fp32" else (4, 2, 1) - auto = next(v for v in widths if chan % v == 0) - return auto if override is None else min(auto, override) - va = CoalescedTileLoader.choose_vec( tile_rows=block_m, tile_cols=block_k, @@ -2609,7 +2698,7 @@ def x_descriptor(b_: IRBuilder, row: Value, col: Value): return X_desc.offset(b_, m=m_val, k=k_val, group=group_v) return X_desc.offset(b_, m=m_val, k=k_val) - if spec.async_dma: + if spec.async_dma and not _wmma_async: # K-outer: rows are K_wg, columns are the free axis. A chunk is then a # run along the free axis at one k_wg, which is contiguous in global for # both operands -- exactly what the intrinsic requires. @@ -2697,8 +2786,10 @@ def x_descriptor(b_: IRBuilder, row: Value, col: Value): vector_axis=b_axis, ) + # The gfx1250 async loop keeps the plain pipeline's (hint-free) policy; the + # MFMA async loop keeps its own. schedule = SchedulePolicy.for_pipeline( - "async_dma" if spec.async_dma else spec.pipeline + "async_dma" if spec.async_dma and not _wmma_async else spec.pipeline ) schedule.emit_prologue(b) @@ -2720,6 +2811,25 @@ def _x_kouter(b_, row, col): else: a_desc_fn, b_desc_fn = dy_descriptor, x_descriptor + if _wmma_async: + a_sync_loader.load_async_to_lds( + b, + tid=tid, + smem_dst=A_dst, + descriptor=a_desc_fn, + ptr=dY, + num_bytes=dY_bytes, + ) + b_sync_loader.load_async_to_lds( + b, + tid=tid, + smem_dst=B_dst, + descriptor=b_desc_fn, + ptr=X, + num_bytes=X_bytes, + ) + return + if spec.async_dma: from rocke.core.ir import CACHE_STREAM @@ -3318,6 +3428,49 @@ def _tile_offsets(k_off: Value) -> List[Value]: b.sync() b.scf_yield(*new_accs, *_next) final_accs = list(for_op.results)[:_n_acc] + elif _wmma_async: + # gfx1250 async ping-pong, one barrier per tile. Each phase waits for its + # own tile's copies (asynccnt 0: nothing else is in flight then), and the + # barrier both publishes that tile -- the copies and the zero-fill + # stores of every wave -- and retires the ds_reads of the tile computed + # out of the other buffer one phase earlier. Only then is the next tile + # issued into that buffer, and its copies stream while this tile's + # WMMAs run. The body is unrolled twice so each phase binds a build-time + # buffer; an odd tile count makes the last phase B compute a tile past + # _k_upper, redirected to the zero-filled wg_K as the unroll_k loop does. + c_2block_k = b.const_i32(2 * block_k) + + def _next_load(k_next: Value) -> Value: + if _k_zero_fill is None: + return k_next + return b.select(b.cmp_lt(k_next, _k_upper), k_next, _k_zero_fill) + + emit_load_phase(k_lo, A_smem, B_smem) + for_op_as = b.scf_for_iter(k_lo, _k_upper, c_2block_k, accs, iv_name="k_async") + with for_op_as as (k_async, iter_accs): + k_odd = b.add(k_async, c_block_k) + k_nxt_pair = b.add(k_async, c_2block_k) + + # Phase A: compute tile k out of buf0 while tile k+1 streams into buf1. + b.s_wait_asynccnt(0) + b.sync_lds_only() + emit_load_phase(_next_load(k_odd), A_smem2, B_smem2) + k_off_capture[0] = k_async + accs_a = emit_mfma_phase(A_smem, B_smem, list(iter_accs)) + + # Phase B: the buffers swap roles. + b.s_wait_asynccnt(0) + b.sync_lds_only() + emit_load_phase(_next_load(k_nxt_pair), A_smem, B_smem) + k_off_capture[0] = k_odd + accs_b = emit_mfma_phase(A_smem2, B_smem2, accs_a) + + b.scf_yield(*accs_b) + # The last phase left one tile in flight; drain it and the zero-fill + # stores before the epilogue reuses LDS. + b.s_wait_asynccnt(0) + b.sync() + final_accs = list(for_op_as.results) elif not spec.async_dma: # mem/compv3/compv4/basic: scf_for_iter with dynamic _k_upper (= p_wg_K or k_hi) for_op = b.scf_for_iter(k_lo, _k_upper, c_block_k, accs, iv_name="k0") @@ -3381,7 +3534,7 @@ def compute_dyn_wg(k_offset_val, buf_pair, state): params, gm_group=group_v, ) - elif _is_split_k and op.family == "wmma": + elif _is_split_k and op.family == "wmma" and spec.epilogue != "cshuffle": # WMMA split-K: atomic-add via the WMMA C-fragment layout (fp32/bf16/fp16). _emit_wgrad_split_k_epilogue_wmma( b, @@ -3398,12 +3551,14 @@ def compute_dyn_wg(k_offset_val, buf_pair, state): group=group_v, ) elif _is_split_k and spec.epilogue == "cshuffle": - # MFMA split-K + cshuffle: scatter to LDS then paired pk_atomic from smem. - # Avoids the zero-fill workaround of the direct split-K path for bf16/fp16. + # Split-K + cshuffle (MFMA or WMMA): scatter to LDS then paired + # pk_atomic from smem. Avoids the zero-fill workaround of the direct + # split-K path for bf16/fp16. _emit_wgrad_split_k_cshuffle_epilogue( b, spec, atom, + op, final_accs, grid, dW, @@ -3429,7 +3584,7 @@ def compute_dyn_wg(k_offset_val, buf_pair, state): ) elif spec.epilogue == "cshuffle": _emit_wgrad_cshuffle_epilogue( - b, spec, final_accs, grid, dw_rsrc, params, group=group_v + b, spec, op, final_accs, grid, dw_rsrc, params, group=group_v ) elif op.family == "wmma": _emit_wgrad_direct_epilogue_wmma( @@ -3742,10 +3897,22 @@ def _emit_single_packed_atomic(c_m: Value, c_n: Value, val_f32: Value) -> None: _emit_single_packed_atomic(c_m, c_n, b.vec_extract(acc, i)) +def _cshuffle_epilogue(op, atom, grid, kwargs: dict) -> CShuffleEpilogue: + """The cshuffle epilogue for this accumulator layout. + + WMMA reads the fragment through ``op.c_layout()`` (``from_grid_op``, as the + forward and dgrad WMMA cshuffle epilogues do); MFMA through the atom. + """ + if op.family == "wmma": + return CShuffleEpilogue.from_grid_op(op=op, grid=grid, **kwargs) + return CShuffleEpilogue.from_grid(atom=atom, grid=grid, **kwargs) + + def _emit_wgrad_split_k_cshuffle_epilogue( b: IRBuilder, spec: WgradConvSpec, - atom: MfmaAtom, + atom: Optional[MfmaAtom], + op, accs: Sequence[Value], grid, dW: Value, @@ -3816,9 +3983,7 @@ def _emit_wgrad_split_k_cshuffle_epilogue( eff_grid = grid wg_M_v = params["p_wg_M"] - CShuffleEpilogue.from_grid( - atom=atom, grid=eff_grid, **_cshuffle_kwargs - ).atomic_store( + _cshuffle_epilogue(op, atom, eff_grid, _cshuffle_kwargs).atomic_store( b, accs=accs, dw_ptr=dW, @@ -3969,6 +4134,7 @@ def _emit_wgrad_direct_epilogue_wmma( def _emit_wgrad_cshuffle_epilogue( b: IRBuilder, spec: WgradConvSpec, + op, accs: Sequence[Value], grid: WarpGrid, dw_rsrc: Value, @@ -4027,7 +4193,7 @@ def dw_addr(b_: IRBuilder, m_val: Value, n_val: Value): _vc_C, _vc_K, spec.data.dtype_d, split_k=spec.split_k ) _cshuffle_kwargs["max_store_vec"] = vec_c - CShuffleEpilogue.from_grid(atom=spec.atom, grid=grid, **_cshuffle_kwargs).store( + _cshuffle_epilogue(op, spec.atom, grid, _cshuffle_kwargs).store( b, accs=accs, addr_fn=dw_addr, diff --git a/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/conv_split_k.py b/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/conv_split_k.py index d1ce5a500759..0ace8899caaf 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/conv_split_k.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/conv_split_k.py @@ -203,10 +203,12 @@ def _resolve_wgrad( # drop an off-diagonal group pair, so only the workspace carries the mask. # On MFMA the packed 16-bit atomic is emitted by the cshuffle epilogue only # (the direct-store one scatters the MFMA layout and has no pairs to add). - # WMMA (wave32) has only the default epilogue, and its split-K epilogue - # forms each packed pair itself. + # On WMMA (wave32) the default split-K epilogue forms each packed pair + # itself; the cshuffle one pairs from LDS as on MFMA. wave_size = ArchTarget.from_gfx(arch).wave_size - layout_ok, atomic_why = wgrad_atomic_epilogue_available(p, dtype, None, wave_size) + layout_ok, atomic_why = wgrad_atomic_epilogue_available( + p, dtype, None, wave_size, epilogue + ) if gm > 1: atomic_why = "a merged tile reduces through the workspace only" elif wave_size == 32 and not wgrad_wmma_split_k_available(arch): diff --git a/dnn-providers/hip-kernel-provider/rocke/library/tests/parity/conv_implicit_gemm_wgrad_emit.c b/dnn-providers/hip-kernel-provider/rocke/library/tests/parity/conv_implicit_gemm_wgrad_emit.c index ca6fdc62cd72..a408008b86ff 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/tests/parity/conv_implicit_gemm_wgrad_emit.c +++ b/dnn-providers/hip-kernel-provider/rocke/library/tests/parity/conv_implicit_gemm_wgrad_emit.c @@ -412,6 +412,94 @@ static int make_base_cfg(int idx, rocke_implicit_gemm_conv_wgrad_spec_t* spec, c spec->split_k = 4; *arch = idx == 52 ? "gfx1201" : idx == 53 ? "gfx1250" : "gfx1151"; return 0; + case 55: + case 56: + case 57: + { + /* gfx1250 WMMA async K loop. Must match the Python emitter's 55-57. */ + const bool big = idx == 56; + if(idx == 57) + spec->problem = rocke_conv_problem_default(4, 12, 12, 3, 32, 3, 3); + else if(idx == 55) + spec->problem = rocke_conv_problem_default(4, 12, 12, 24, 32, 3, 3); + else + spec->problem = rocke_conv_problem_default(8, 56, 56, 64, 64, 3, 3); + spec->problem.pH = 1; + spec->problem.pW = 1; + spec->dtype_a = "fp16"; + spec->dtype_b = "fp16"; + spec->dtype_d = "fp32"; + spec->tile_m = big ? 64 : 32; + spec->tile_n = big ? 64 : 32; + spec->tile_k = big ? 64 : 32; + spec->warp_m = big ? 2 : 1; + spec->warp_n = big ? 2 : 1; + spec->warp_tile_m = 16; + spec->warp_tile_n = 16; + spec->warp_tile_k = 32; + spec->wave_size = 32; + spec->pipeline = "mem"; + spec->epilogue = "default"; + spec->split_k = big ? 4 : 2; + spec->lds_k_outer = true; + spec->async_dma = true; + *arch = "gfx1250"; + return 0; + } + case 58: + case 59: + case 60: + case 61: + { + /* WMMA cshuffle epilogue. Must match the Python emitter's 58-61. */ + const bool gfx1250 = idx != 61; + if(idx == 60) + spec->problem = rocke_conv_problem_default(4, 12, 12, 24, 32, 3, 3); + else + spec->problem = rocke_conv_problem_default(8, 56, 56, 64, 64, 3, 3); + spec->problem.pH = 1; + spec->problem.pW = 1; + spec->dtype_a = "fp16"; + spec->dtype_b = "fp16"; + spec->dtype_d = idx == 59 ? "bf16" : "fp16"; + spec->tile_m = 64; + spec->tile_n = 64; + spec->tile_k = 64; + spec->warp_m = 2; + spec->warp_n = 2; + spec->warp_tile_m = 16; + spec->warp_tile_n = 16; + spec->warp_tile_k = gfx1250 ? 32 : 16; + spec->wave_size = 32; + spec->pipeline = "mem"; + spec->epilogue = "cshuffle"; + spec->split_k = idx == 58 ? 1 : (idx == 60 ? 2 : 4); + spec->lds_k_outer = gfx1250; + spec->async_dma = idx == 60; + *arch = gfx1250 ? "gfx1250" : "gfx1201"; + return 0; + } + case 62: + /* WMMA cshuffle split-K with an odd store vector. Must match the + * Python emitter's 62 (both engines reject). */ + spec->problem = rocke_conv_problem_default(4, 12, 12, 3, 32, 2, 2); + spec->dtype_a = "fp16"; + spec->dtype_b = "fp16"; + spec->dtype_d = "fp16"; + spec->tile_m = 32; + spec->tile_n = 32; + spec->tile_k = 32; + spec->warp_m = 1; + spec->warp_n = 1; + spec->warp_tile_m = 16; + spec->warp_tile_n = 16; + spec->warp_tile_k = 32; + spec->wave_size = 32; + spec->pipeline = "mem"; + spec->epilogue = "cshuffle"; + spec->split_k = 4; + *arch = "gfx1250"; + return 0; default: return -1; } diff --git a/dnn-providers/hip-kernel-provider/rocke/library/tests/parity/conv_implicit_gemm_wgrad_emit.py b/dnn-providers/hip-kernel-provider/rocke/library/tests/parity/conv_implicit_gemm_wgrad_emit.py index 24aa0a3fd12a..78c746e165bd 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/tests/parity/conv_implicit_gemm_wgrad_emit.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/tests/parity/conv_implicit_gemm_wgrad_emit.py @@ -52,9 +52,24 @@ # 49 -- gfx1250 WMMA split-K=4, bf16 dW (zero-fill packed bf16 atomic), M-outer # 50 -- gfx1250 WMMA split-K=4, fp32 dW (scalar atomic), K-outer # 51 -- gfx1250 WMMA split-K=2, fp16 dW, C=24 (partial N tiles) -# 52-54 -- WMMA split-K both engines must reject: gfx1201 bf16 dW (arch not -# enabled), gfx1250 bf16 dW with the cshuffle epilogue (outside the -# WMMA wgrad subset), gfx1151 fp32 dW (arch not enabled, any dtype) +# 52-54 -- WMMA split-K: gfx1201 bf16 dW (arch not enabled) and gfx1151 fp32 +# dW (arch not enabled, any dtype), which both engines must reject, and +# gfx1250 bf16 dW with the cshuffle epilogue (M-outer), which both admit +# since the WMMA cshuffle epilogue +# 55 -- gfx1250 async_dma (global_load_async_to_lds ping-pong), K-outer, +# split-K=2 fp32 dW, C=24 (partial N tiles: zero-filled chunks) and odd +# slices (zero-fill redirect) +# 56 -- gfx1250 async_dma, K-outer, 64x64x64 2x2 warps, split-K=4 fp32 dW +# 57 -- both engines must reject: gfx1250 async_dma with C=3 (a one-element +# chunk is narrower than the copy) +# 58-60 -- gfx1250 cshuffle (WMMA fragment through LDS): direct fp16 dW, +# split-K bf16 dW (genuine paired atomics), and with async_dma under +# split-K at C=24 (partial N tiles) +# 61 -- both engines must reject: cshuffle split-K on gfx1201 (WMMA split-K +# is enabled only on gfx1250) +# 62 -- both engines must reject: gfx1250 cshuffle split-K fp16 dW whose +# store vector is odd (C=3 with a 2x2 filter: wg_N=12 is even, but +# the LDS pairing has no partner) # (Grouped wgrad reaches the C++ engine only through stream-K, and its tile # loader has no elem_dtype, so no case loads bf16 operands. The stream-K # partition and CTA pool are kernargs, so no case varies them. async_dma is @@ -70,7 +85,7 @@ # 107 -- streamk with unroll_k (must raise ValueError) # 108 -- streamk atomic reduction into fp16 dW (must raise ValueError) # (These illustrate the validator contract. The C emitter defines only cases -# 0-54, so run_diff.py stops at the shared END before reaching 100+; these +# 0-62, so run_diff.py stops at the shared END before reaching 100+; these # configs are not exercised by the differential gate.) from kernels.common.conv_implicit_gemm_wgrad import ( WgradConvSpec, @@ -694,6 +709,96 @@ def _spec(idx: int): arch, ) + if 55 <= idx <= 57: + # gfx1250 WMMA K loop: the global_load_async_to_lds ping-pong. + from kernels.common._conv_implicit_gemm_common import ConvDataSpec + + p = ConvProblem(N=8, Hi=56, Wi=56, C=64, K=64, Y=3, X=3, pH=1, pW=1) + p_small = ConvProblem(N=4, Hi=12, Wi=12, C=24, K=32, Y=3, X=3, pH=1, pW=1) + p_c3 = ConvProblem(N=4, Hi=12, Wi=12, C=3, K=32, Y=3, X=3, pH=1, pW=1) + problem, dtype_d, split_k, big, extra, arch = { + 55: (p_small, "fp32", 2, False, {"async_dma": True}, "gfx1250"), + 56: (p, "fp32", 4, True, {"async_dma": True}, "gfx1250"), + 57: (p_c3, "fp32", 2, False, {"async_dma": True}, "gfx1250"), + }[idx] + wave32 = arch == "gfx1250" + kw = dict( + problem=problem, + data=ConvDataSpec(dtype_a="fp16", dtype_b="fp16", dtype_d=dtype_d), + tile_m=64 if big else 32, + tile_n=64 if big else 32, + tile_k=64 if big else 32, + warp_m=2 if big else 1, + warp_n=2 if big else 1, + warp_tile_m=16, + warp_tile_n=16, + warp_tile_k=32 if wave32 else 16, + wave_size=32 if wave32 else 64, + pipeline="mem", + epilogue="default", + split_k=split_k, + lds_k_outer=wave32, + ) + kw.update(extra) + return WgradConvSpec(**kw), arch + + if 58 <= idx <= 61: + # WMMA cshuffle epilogue. + from kernels.common._conv_implicit_gemm_common import ConvDataSpec + + p64 = ConvProblem(N=8, Hi=56, Wi=56, C=64, K=64, Y=3, X=3, pH=1, pW=1) + p24 = ConvProblem(N=4, Hi=12, Wi=12, C=24, K=32, Y=3, X=3, pH=1, pW=1) + problem, dtype_d, split_k, extra, arch = { + 58: (p64, "fp16", 1, {}, "gfx1250"), + 59: (p64, "bf16", 4, {}, "gfx1250"), + 60: (p24, "fp16", 2, {"async_dma": True}, "gfx1250"), + 61: (p64, "fp16", 4, {}, "gfx1201"), + }[idx] + gfx1250 = arch == "gfx1250" + kw = dict( + problem=problem, + data=ConvDataSpec(dtype_a="fp16", dtype_b="fp16", dtype_d=dtype_d), + tile_m=64, + tile_n=64, + tile_k=64, + warp_m=2, + warp_n=2, + warp_tile_m=16, + warp_tile_n=16, + warp_tile_k=32 if gfx1250 else 16, + wave_size=32, + pipeline="mem", + epilogue="cshuffle", + split_k=split_k, + lds_k_outer=gfx1250, + ) + kw.update(extra) + return WgradConvSpec(**kw), arch + + if idx == 62: + # WMMA cshuffle split-K needs an even store vector, as on MFMA. + from kernels.common._conv_implicit_gemm_common import ConvDataSpec + + return ( + WgradConvSpec( + problem=ConvProblem(N=4, Hi=12, Wi=12, C=3, K=32, Y=2, X=2), + data=ConvDataSpec(dtype_a="fp16", dtype_b="fp16", dtype_d="fp16"), + tile_m=32, + tile_n=32, + tile_k=32, + warp_m=1, + warp_n=1, + warp_tile_m=16, + warp_tile_n=16, + warp_tile_k=32, + wave_size=32, + pipeline="mem", + epilogue="cshuffle", + split_k=4, + ), + "gfx1250", + ) + # ---------------------------------------------------------------- # Negative cases: these specs must be REJECTED by the validator. # The harness (run_emit) expects a ValueError / SystemExit when diff --git a/dnn-providers/hip-kernel-provider/rocke/library/tests/test_conv_wgrad_correctness.py b/dnn-providers/hip-kernel-provider/rocke/library/tests/test_conv_wgrad_correctness.py index 7b5803b5c974..8211ed6e3a66 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/tests/test_conv_wgrad_correctness.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/tests/test_conv_wgrad_correctness.py @@ -2308,14 +2308,32 @@ def test_16bit_split_k_atomic_arch_scope(self): ok, why = is_valid_wgrad_spec(other, arch) self.assertFalse(ok, f"{dtype} split-K on {arch} must stay rejected") self.assertIn("enabled only on", why) - # And no WMMA epilogue but 'default' is admitted for split-K. - ok, _ = is_valid_wgrad_spec( + # The cshuffle epilogue is admitted for WMMA split-K too: it pairs + # through LDS from the WMMA fragment (CShuffleEpilogue.from_grid_op). + ok, why = is_valid_wgrad_spec( self._gfx1250_spec(data=data, split_k=4, epilogue="cshuffle"), "gfx1250", ) - self.assertFalse( - ok, f"{dtype} cshuffle split-K on gfx1250 must be rejected" + self.assertTrue(ok, f"{dtype} cshuffle split-K on gfx1250: {why}") + + def test_wmma_cshuffle_split_k_needs_even_store_vec(self): + # The cshuffle epilogue pairs from LDS on WMMA as on MFMA, so unlike the + # default WMMA epilogue it needs an even store vector. + from kernels.common._conv_implicit_gemm_common import ConvDataSpec + from kernels.common.conv_implicit_gemm_wgrad import is_valid_wgrad_spec + + for dtype in ("bf16", "fp16"): + data = ConvDataSpec(dtype_a=dtype, dtype_b=dtype, dtype_d=dtype) + for sv, want in ((1, False), (2, True), (8, True)): + spec = self._gfx1250_spec( + data=data, split_k=4, epilogue="cshuffle", vector_size_c=sv + ) + ok, why = is_valid_wgrad_spec(spec, "gfx1250") + self.assertEqual(ok, want, f"{dtype} sv={sv}: {why}") + ok, _ = is_valid_wgrad_spec( + self._gfx1250_spec(data=data, split_k=4, vector_size_c=1), "gfx1250" ) + self.assertTrue(ok, "the default WMMA epilogue ignores the store vector") def test_wmma_split_k_ignores_store_vec(self): # The AOT cache grid pins vector_size_c=1 for the default epilogue. The @@ -2458,11 +2476,13 @@ class TestWgradDoubleBufferedSplitKTail(unittest.TestCase): ) def _sweep(self, **knobs) -> None: - # Both loops are MFMA-only (WMMA wgrad rejects async_dma and - # unroll_k). On MFMA every case must build and run: a skip (an - # invalid spec) is a failure, not a quiet pass. - if not _IS_MFMA: - self.skipTest(f"{knobs} wgrad is MFMA-only; running on {GPU_ARCH}") + # Both loops exist on MFMA; gfx1250 WMMA has async_dma only (the + # global_load_async_to_lds ping-pong) and rejects unroll_k. Every case + # must build and run: a skip (an invalid spec) is a failure, not a + # quiet pass. + gfx1250_async = GPU_ARCH == "gfx1250" and knobs.get("async_dma", False) + if not (_IS_MFMA or gfx1250_async): + self.skipTest(f"{knobs} wgrad is not built on {GPU_ARCH}") for shape in self._CASES: for dtype in _DTYPES: with self.subTest(shape=shape.id, dtype=dtype): @@ -2499,6 +2519,132 @@ def _assert_case_ran(test, ok: bool, why: str) -> None: test.assertFalse(why.startswith("skip"), f"case was skipped rather than run: {why}") +@unittest.skipIf(_skip_reason(), _skip_reason()) +class TestWgradWmmaCShuffle(unittest.TestCase): + """gfx1250 dW numerics for the WMMA cshuffle epilogue. + + The tile is staged through LDS from the WMMA C layout (op.c_layout()) and + written as wide stores, or as genuine packed atomic pairs under split-K. A + wrong lane-to-element mapping shows up only in dW, so every case runs + against the CPU reference: direct store and split-K, fp16 and bf16, partial + M/N tiles (C=24), pointwise, stride 2, grouped, and the async K loop. + """ + + _CASES = ( + _Shape( + "3x3_N2H14W14C64K64", N=2, Hi=14, Wi=14, C=64, K=64, Y=3, X=3, pH=1, pW=1 + ), + _Shape( + "3x3_C24_N2H12W12C24K32", + N=2, + Hi=12, + Wi=12, + C=24, + K=32, + Y=3, + X=3, + pH=1, + pW=1, + ), + _Shape("1x1_N2H14W14C32K32", N=2, Hi=14, Wi=14, C=32, K=32, Y=1, X=1), + _Shape( + "3x3s2_N2H15W15C32K48", + N=2, + Hi=15, + Wi=15, + C=32, + K=48, + Y=3, + X=3, + sH=2, + sW=2, + pH=1, + pW=1, + ), + _Shape( + "3x3_g2_N2H10W10C32K32", + N=2, + Hi=10, + Wi=10, + C=32, + K=32, + Y=3, + X=3, + pH=1, + pW=1, + groups=2, + ), + ) + + def _sweep(self, **knobs) -> None: + if GPU_ARCH != "gfx1250": + self.skipTest( + f"the WMMA cshuffle epilogue is gfx1250-tested; on {GPU_ARCH}" + ) + from kernels.common._conv_implicit_gemm_common import ConvProblem + from kernels.common.conv_implicit_gemm import ConvDataSpec + from kernels.common.conv_implicit_gemm_wgrad import WgradConvSpec + + for shape in self._CASES: + p = ConvProblem( + N=shape.N, + Hi=shape.Hi, + Wi=shape.Wi, + C=shape.C, + K=shape.K, + Y=shape.Y, + X=shape.X, + sH=shape.sH, + sW=shape.sW, + pH=shape.pH, + pW=shape.pW, + groups=shape.groups, + ) + for dtype in _DTYPES: + for split_k in (1, 2): + for tm, tn, tk, wm, wn in ((32, 64, 32, 2, 2), (64, 32, 64, 2, 1)): + spec = WgradConvSpec( + problem=p, + data=ConvDataSpec( + dtype_a=dtype, dtype_b=dtype, dtype_d=dtype + ), + tile_m=tm, + tile_n=tn, + tile_k=tk, + warp_m=wm, + warp_n=wn, + warp_tile_m=16, + warp_tile_n=16, + warp_tile_k=32, + wave_size=32, + pipeline="mem", + epilogue="cshuffle", + split_k=split_k, + lds_k_outer=True, + **knobs, + ) + with self.subTest( + shape=shape.id, dtype=dtype, sk=split_k, tile=(tm, tn, tk) + ): + ok, why = _run_one( + GPU_ARCH, + shape, + dtype, + "mem", + "cshuffle", + spec_override=spec, + ) + _assert_case_ran( + self, ok, f"{shape.id} {dtype} {knobs}: {why}" + ) + + def test_cshuffle(self): + self._sweep() + + def test_cshuffle_async(self): + self._sweep(async_dma=True) + + @unittest.skipUnless(not _SKIP_REASON, _SKIP_REASON or "needs CDNA GPU + torch") class TestWgradGroupMergeNumerics(unittest.TestCase): """GPU numerics for group merging, over the two-stage path. diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/helpers/loads.cpp b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/helpers/loads.cpp index 15aa90df6076..c9b765652bb9 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/helpers/loads.cpp +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/helpers/loads.cpp @@ -632,6 +632,71 @@ void rocke_coalesced_tile_loader_issue_slots(rocke_ir_builder_t* b, staged->count = count; } +void rocke_coalesced_tile_loader_load_async_to_lds(rocke_ir_builder_t* b, + const rocke_coalesced_tile_loader_t* self, + rocke_value_t* tid, + rocke_value_t* smem_dst, + rocke_loads_descriptor_fn descriptor, + void* descriptor_user, + rocke_value_t* ptr, + rocke_value_t* num_bytes, + int coherency) +{ + if(b != NULL && b->status != ROCKE_OK) + return; + if(self == NULL || descriptor == NULL || ptr == NULL || smem_dst == NULL || num_bytes == NULL) + { + rocke_i_set_err(b, ROCKE_ERR_VALUE, "load_async_to_lds: NULL argument"); + return; + } + if(self->vector_axis_row) + { + rocke_i_set_err(b, + ROCKE_ERR_VALUE, + "CoalescedTileLoader.load_async_to_lds supports only vector_axis='col' " + "(the copy cannot transpose on store)"); + return; + } + /* This port loads 16-bit elements (see slot_offsets). */ + const int width = self->load_vec * 2; + if(width != 4 && width != 8 && width != 16) + { + rocke_i_set_err(b, + ROCKE_ERR_VALUE, + "global_load_async_to_lds moves 4, 8 or 16 bytes per lane; a " + "%d x 2 B chunk is %d B", + self->load_vec, + width); + return; + } + + /* Python: zero = b.zero_vec(self.elem_dtype, self.load_vec) */ + rocke_value_t* zero = rocke_b_zero_vec(b, rocke_f16(), self->load_vec); + rocke_value_t* c_elem_bytes = rocke_b_const_i32(b, 2); + rocke_ctl_staged_t slots; + rocke_coalesced_tile_loader_plan_slots(b, self, tid, &slots); + for(int e = 0; e < slots.count; ++e) + { + rocke_value_t* idx[2]; + idx[0] = slots.vecs[e].row; + idx[1] = slots.vecs[e].col; + /* Python: off_elems, valid = descriptor(b, row, col) */ + rocke_value_t* valid = NULL; + rocke_value_t* off_elems = descriptor(b, idx[0], idx[1], &valid, descriptor_user); + /* Python: in_bounds = b.cmp_lt(b.mul(off_elems, c_elem_bytes), num_bytes) */ + rocke_value_t* off_bytes = rocke_b_mul(b, off_elems, c_elem_bytes); + rocke_value_t* in_bounds = rocke_b_cmp_lt(b, off_bytes, num_bytes); + valid = (valid == NULL) ? in_bounds : rocke_b_land(b, valid, in_bounds); + rocke_if_else_t ie = rocke_b_scf_if_else(b, valid); + rocke_b_region_enter(b, ie.then_region); + rocke_b_global_load_async_to_lds(b, ptr, off_elems, smem_dst, idx, 2, width, coherency, 0); + rocke_b_region_leave(b); + rocke_b_region_enter(b, ie.else_region); + rocke_b_smem_store_vN_f16(b, smem_dst, idx, 2, zero, self->load_vec); + rocke_b_region_leave(b); + } +} + /* ============================================================================ * AsyncTileLoader / AsyncTileLoaderSlot * ========================================================================== */ diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/helper_rocke.helpers.loads.h b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/helper_rocke.helpers.loads.h index 02f9d541aac2..a1cb87ed47c8 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/helper_rocke.helpers.loads.h +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/helper_rocke.helpers.loads.h @@ -256,6 +256,25 @@ void rocke_coalesced_tile_loader_issue_slots(rocke_ir_builder_t* b, rocke_value_t* const* offsets, rocke_ctl_staged_t* staged); +/* CoalescedTileLoader.load_async_to_lds() -- gfx1250 async variant of load(). + * + * Same slot partition and descriptor contract as load() in vector_axis="col" + * mode, but each in-range chunk is one per-lane global_load_async_to_lds of + * load_vec 16-bit elements from `ptr` to smem_dst[row, col]. The copy has no + * buffer-resource range check, so the one the resource would make (byte offset + * below `num_bytes`) is ANDed into the descriptor's valid; a chunk failing + * either is zero-filled with an ordinary LDS store instead. The caller drains + * the copies with s_wait_asynccnt (and the zero stores with an LDS barrier). */ +void rocke_coalesced_tile_loader_load_async_to_lds(rocke_ir_builder_t* b, + const rocke_coalesced_tile_loader_t* self, + rocke_value_t* tid, + rocke_value_t* smem_dst, + rocke_loads_descriptor_fn descriptor, + void* descriptor_user, + rocke_value_t* ptr, + rocke_value_t* num_bytes, + int coherency); + /* ============================================================================ * AsyncTileLoader / AsyncTileLoaderSlot * ========================================================================== */ diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/instance_conv_implicit_gemm_internal.h b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/instance_conv_implicit_gemm_internal.h index a576b518eccf..05c854a2a95a 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/instance_conv_implicit_gemm_internal.h +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/instance_conv_implicit_gemm_internal.h @@ -466,6 +466,12 @@ typedef struct rocke_conv_build_ctx * pipeline="wavelet": a_wavelet_loader/b_wavelet_loader are * CoalescedTileLoader sized to num_load_waves * wave_size threads. */ bool async_dma; /* spec.async_dma (cached) */ + /* gfx1250 async K loop (wgrad only): the sync loaders' slot partition with + * each chunk moved by global_load_async_to_lds from a_ptr / b_ptr instead of + * a buffer load. Zeroed by the ctx memset, so fwd/dgrad are unaffected. */ + bool async_global_lds; + rocke_value_t* a_ptr; + rocke_value_t* b_ptr; /* spec.lds_k_outer (wgrad only). Zeroed by the ctx memset, so the forward * and dgrad drivers keep their exact previous behaviour. */ bool lds_k_outer; @@ -758,6 +764,10 @@ void rocke_conv_emit_kloop_simple(rocke_conv_build_ctx_t* ctx); * its run_ping_pong sequencing inline against ctx until that port lands.) */ void rocke_conv_emit_kloop_async(rocke_conv_build_ctx_t* ctx); +/* gfx1250 async branch (wgrad only, ctx->async_global_lds): one-barrier + * ping-pong over global_load_async_to_lds copies drained by s_wait_asynccnt. */ +void rocke_conv_emit_kloop_async_global_lds(rocke_conv_build_ctx_t* ctx); + /* pipeline=="wavelet" branch (Python lines 1539-1797): dedicated load-wave / * math-wave split. Requires ctx->have_wavelet_loaders and the wavelet_* fields * to be populated by rocke_conv_build_ctx_init. Dispatches the WMMA diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/common/conv_implicit_gemm_conv_kloop_drivers.cpp b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/common/conv_implicit_gemm_conv_kloop_drivers.cpp index eac50da69cfb..223b1ad9d542 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/common/conv_implicit_gemm_conv_kloop_drivers.cpp +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/common/conv_implicit_gemm_conv_kloop_drivers.cpp @@ -140,6 +140,84 @@ void rocke_conv_emit_kloop_unroll(rocke_conv_build_ctx_t* ctx) rocke_conv_set_final_accs(ctx, for_op.op->results, for_op.op->num_results); } +/* ===================================================================== * + * rocke_conv_emit_kloop_async_global_lds (wgrad ``elif _wmma_async:``) + * + * gfx1250 async ping-pong with one barrier per tile. Each phase waits for its + * own tile's global_load_async_to_lds copies (asynccnt 0), the LDS barrier + * publishes that tile and retires the ds_reads of the other buffer, and only + * then is the next tile issued into the other buffer, streaming while this + * tile's WMMAs run. Unrolled twice so each phase binds a build-time buffer. + * ===================================================================== */ +void rocke_conv_emit_kloop_async_global_lds(rocke_conv_build_ctx_t* ctx) +{ + rocke_ir_builder_t* b = ctx->b; + int num_accs = ctx->num_accs; + int i; + + rocke_iter_arg_t iter_args[ROCKE_CONV_MAX_ACCS]; + rocke_value_t* iter_vars[ROCKE_CONV_MAX_ACCS]; + rocke_value_t* accs_a[ROCKE_CONV_MAX_ACCS]; + rocke_value_t* accs_b[ROCKE_CONV_MAX_ACCS]; + rocke_value_t* k_lo = (ctx->kloop_k_lo != NULL) ? ctx->kloop_k_lo : ctx->c0; + rocke_value_t* k_hi = (ctx->kloop_k_hi != NULL) ? ctx->kloop_k_hi : ctx->c_K_gemm; + rocke_value_t* c_2block_k = rocke_b_const_i32(b, 2 * ctx->block_k); + rocke_for_t for_op; + + rocke_conv_emit_load_phase(ctx, k_lo, ctx->A_smem, ctx->B_smem); + + for(i = 0; i < num_accs; ++i) + { + iter_args[i].name = ctx->acc_names[i]; + iter_args[i].init = ctx->acc_inits[i]; + } + for_op = rocke_b_scf_for_iter(b, + k_lo, + k_hi, + c_2block_k, + iter_args, + num_accs, + "k_async", + /*unroll=*/false, + /*elide_trailing_barrier=*/true); + for(i = 0; i < for_op.num_iter_vars; ++i) + iter_vars[i] = for_op.iter_vars[i]; + + rocke_b_region_enter(b, for_op.body); + { + rocke_value_t* k_odd = rocke_b_add(b, for_op.iv, ctx->c_block_k); + rocke_value_t* k_nxt_pair = rocke_b_add(b, for_op.iv, c_2block_k); + rocke_value_t* k_load; + + /* Phase A: compute tile k out of buf0 while tile k+1 streams into buf1. */ + rocke_b_s_wait_asynccnt(b, 0); + rocke_b_sync_lds_only(b); + k_load = rocke_conv_kloop_guard_prefetch(ctx, k_odd, k_hi); + rocke_conv_emit_load_phase(ctx, k_load, ctx->A_smem2, ctx->B_smem2); + ctx->k_off_capture = for_op.iv; + rocke_conv_emit_mfma_phase( + ctx, ctx->A_smem, ctx->B_smem, iter_vars, for_op.num_iter_vars, accs_a); + + /* Phase B: the buffers swap roles. */ + rocke_b_s_wait_asynccnt(b, 0); + rocke_b_sync_lds_only(b); + k_load = rocke_conv_kloop_guard_prefetch(ctx, k_nxt_pair, k_hi); + rocke_conv_emit_load_phase(ctx, k_load, ctx->A_smem, ctx->B_smem); + ctx->k_off_capture = k_odd; + rocke_conv_emit_mfma_phase( + ctx, ctx->A_smem2, ctx->B_smem2, accs_a, for_op.num_iter_vars, accs_b); + + rocke_b_scf_yield(b, accs_b, for_op.num_iter_vars); + } + rocke_b_region_leave(b); + + /* The last phase left one tile in flight; drain it and the zero-fill + * stores before the epilogue reuses LDS. */ + rocke_b_s_wait_asynccnt(b, 0); + rocke_b_sync(b); + rocke_conv_set_final_accs(ctx, for_op.op->results, for_op.op->num_results); +} + /* ===================================================================== * * rocke_conv_emit_kloop_simple (Python lines 1311-1319) * diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/common/conv_implicit_gemm_conv_load_phase.cpp b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/common/conv_implicit_gemm_conv_load_phase.cpp index c53dc3ac148f..84d6c227836d 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/common/conv_implicit_gemm_conv_load_phase.cpp +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/common/conv_implicit_gemm_conv_load_phase.cpp @@ -514,6 +514,32 @@ void rocke_conv_emit_load_phase(rocke_conv_build_ctx_t* ctx, /* k_off_capture[0] = k_off */ ctx->k_off_capture = k_off; + if(ctx->async_global_lds) + { + /* gfx1250: Python a_sync_loader / b_sync_loader.load_async_to_lds with + * the default coherency (the X tile is re-read across N tiles and + * filter taps, so it is not streamed past the cache). */ + rocke_coalesced_tile_loader_load_async_to_lds(b, + &ctx->a_sync_loader, + ctx->tid, + A_dst, + ctx->a_descriptor_fn, + ctx, + ctx->a_ptr, + ctx->a_buf_rsrc.num_bytes, + 0); + rocke_coalesced_tile_loader_load_async_to_lds(b, + &ctx->b_sync_loader, + ctx->tid, + B_dst, + ctx->b_descriptor_fn, + ctx, + ctx->b_ptr, + ctx->b_buf_rsrc.num_bytes, + 0); + return; + } + if(spec->async_dma) { /* CACHE_STREAM (SLC=1): each K-tile is consumed exactly once by the diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/common/conv_implicit_gemm_wgrad.cpp b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/common/conv_implicit_gemm_wgrad.cpp index 3fe9366baaad..885636a1eae4 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/common/conv_implicit_gemm_wgrad.cpp +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/common/conv_implicit_gemm_wgrad.cpp @@ -600,6 +600,76 @@ static bool wgrad_group_merge_available(const rocke_implicit_gemm_conv_wgrad_spe #undef GM_REJECT } +/* Python _wgrad_free_axis_vec(chan, dtype) with no override. */ +static int wgrad_free_axis_vec(int chan, const char* dtype) +{ + const bool f32 = dtype != NULL && strcmp(dtype, "fp32") == 0; + for(int v = f32 ? 4 : 8; v > 1; v /= 2) + if(chan % v == 0) + return v; + return 1; +} + +/* Mirrors Python _wmma_async_available: whether a wave32 spec can take the + * gfx1250 async global->LDS K loop. global_load_async_to_lds moves 4, 8 or 16 + * bytes per lane, so a 16-bit chunk needs at least two elements on both + * operands. */ +static bool wgrad_wmma_async_available(const rocke_implicit_gemm_conv_wgrad_spec_t* s, + const char* arch, + char* reason, + size_t reason_cap) +{ + const rocke_archtarget_t* t = rocke_archtarget_from_gfx(arch); + if(t == NULL || !t->memory.has_async_global_lds) + { + if(reason && reason_cap) + snprintf(reason, + reason_cap, + "WMMA wgrad async_dma needs global_load_async_to_lds (not on %s)", + arch); + return false; + } + if(!s->lds_k_outer) + { + if(reason && reason_cap) + snprintf(reason, reason_cap, "WMMA wgrad async_dma requires lds_k_outer=True"); + return false; + } + const int threads = s->warp_m * s->warp_n * s->wave_size; + const int cap_a = wgrad_free_axis_vec(rocke_conv_problem_kpg(&s->problem), s->dtype_a); + const int cap_b = wgrad_free_axis_vec(rocke_conv_problem_cpg(&s->problem), s->dtype_b); + int va = 1; + int vb = 1; + if(rocke_coalesced_tile_loader_choose_vec_axis(s->tile_k, s->tile_m, threads, cap_a, false, &va) + != ROCKE_OK + || rocke_coalesced_tile_loader_choose_vec_axis( + s->tile_k, s->tile_n, threads, cap_b, false, &vb) + != ROCKE_OK) + { + if(reason && reason_cap) + snprintf(reason, reason_cap, "no usable load_vec for the K-outer tile geometry"); + return false; + } + const int eb = (s->dtype_a != NULL && strcmp(s->dtype_a, "fp32") == 0) ? 4 : 2; + const int widths[2] = {va * eb, vb * eb}; + for(int i = 0; i < 2; ++i) + { + if(widths[i] != 4 && widths[i] != 8 && widths[i] != 16) + { + if(reason && reason_cap) + snprintf(reason, + reason_cap, + "WMMA wgrad async_dma needs a 4/8/16-byte %s chunk " + "(global_load_async_to_lds width); got %d x %d B", + i == 0 ? "A" : "B", + i == 0 ? va : vb, + eb); + return false; + } + } + return true; +} + bool rocke_implicit_gemm_conv_wgrad_is_valid_spec(const rocke_implicit_gemm_conv_wgrad_spec_t* s, const char* arch, char* reason, @@ -840,11 +910,13 @@ bool rocke_implicit_gemm_conv_wgrad_is_valid_spec(const rocke_implicit_gemm_conv * then rejects -- the same admits/build split this gate exists to * close. store_vec mirrors default_vector_sizes(..., split_k=1): * widest of 8/4/2/1 dividing the channel run (per-group when grouped). - * MFMA only: the WMMA epilogue (wave32) issues one pair per + * Not for the WMMA default epilogue (wave32): it issues one pair per * accumulator element with a zero partner and never reads the store - * width, so only the row parity applies there. */ + * width, so only the row parity applies there. The WMMA cshuffle + * epilogue pairs from LDS as MFMA does and needs the even width. */ + const bool cshuffle_epi = s->epilogue && strcmp(s->epilogue, "cshuffle") == 0; int store_vec; - if(is_wmma_family) + if(is_wmma_family && !cshuffle_epi) { store_vec = 2; } @@ -1117,10 +1189,56 @@ bool rocke_implicit_gemm_conv_wgrad_is_valid_spec(const rocke_implicit_gemm_conv fwd.dtype_d = s->dtype_d; /* Use a dummy problem that keeps K_gemm/M positive for the LDS calc. */ fwd.problem = rocke_conv_problem_default(1, 8, 8, 16, 16, 1, 1); + /* The forward WMMA subset has no async K loop, which wgrad's WMMA path has + * on gfx1250, so hand it the synchronous equivalent and charge the + * double-buffered LDS below, the way Python is_valid_wgrad_spec does. */ + const bool is_wave32 = s->wave_size == 32; + if(is_wave32) + fwd.async_dma = false; if(!rocke_implicit_gemm_conv_is_valid_spec(&fwd, arch, reason, reason_cap)) return false; + if(is_wave32 && s->async_dma) + { + /* Python: the A/B tile charged twice for the buffer-alternating loops, + * K-outer (tile_k, tile_mn + _kouter_pad) or the M-outer layout rows. */ + const char* da = s->dtype_a ? s->dtype_a : "fp16"; + const long eb = strcmp(da, "fp32") == 0 ? 4 : 2; + long ab; + if(s->lds_k_outer) + { + const long pad = ROCKE_WGRAD_KOUTER_PAD; /* wave32: never the gfx9 pad 0 */ + ab = ((long)s->tile_k * (s->tile_m + pad) + (long)s->tile_k * (s->tile_n + pad)) * eb; + } + else + { + rocke_conv_lds_layout_t lay; + char lay_reason[256]; + if(!rocke_implicit_gemm_conv_spec_effective_lds_layout( + &fwd, &lay, lay_reason, sizeof(lay_reason))) + { + if(reason && reason_cap) + snprintf(reason, reason_cap, "%s", lay_reason); + return false; + } + ab = ((long)s->tile_m * lay.row_stride + (long)s->tile_n * lay.row_stride) * eb; + } + const rocke_archtarget_t* t32 = rocke_archtarget_from_gfx(arch); + if(!rocke_archtarget_fits_lds(t32, 2 * ab)) + { + if(reason && reason_cap) + snprintf(reason, + reason_cap, + "LDS budget %ld bytes (A/B=x2 %ld, C=0, aliased) > %d cap on %s", + 2 * ab, + ab, + t32->lds_capacity_bytes, + arch); + return false; + } + } + /* The forward validator has pinned wave_size to the arch, so wave32 is the * WMMA family here. Its WMMA subset is the forward conv's, which is wider * than wgrad's (it also admits cshuffle, wavelet and 16x16x4); mirror the @@ -1153,13 +1271,16 @@ bool rocke_implicit_gemm_conv_wgrad_is_valid_spec(const rocke_implicit_gemm_conv arch); return false; } + /* cshuffle reads the accumulators through op.c_layout(), the WMMA path + * the forward and dgrad cshuffle epilogues use. */ const char* epilogue = s->epilogue ? s->epilogue : "default"; - if(strcmp(epilogue, "default") != 0) + if(strcmp(epilogue, "default") != 0 && strcmp(epilogue, "cshuffle") != 0) { if(reason && reason_cap) snprintf(reason, reason_cap, - "WMMA wgrad supports only the 'default' epilogue (got '%s') on %s", + "WMMA wgrad supports only the 'default' and 'cshuffle' epilogues " + "(got '%s') on %s", epilogue, arch); return false; @@ -1171,6 +1292,8 @@ bool rocke_implicit_gemm_conv_wgrad_is_valid_spec(const rocke_implicit_gemm_conv reason, reason_cap, "WMMA wgrad does not support chiplet_swizzle on %s", arch); return false; } + if(s->async_dma && !wgrad_wmma_async_available(s, arch, reason, reason_cap)) + return false; } /* Mirrors the Python WMMA flag gate. */ @@ -2195,8 +2318,10 @@ static void wgrad_emit_split_k_cshuffle_epilogue(rocke_ir_builder_t* b, vec_c = 1; } + /* WMMA reads the fragment through op.c_layout() (Python _cshuffle_epilogue). */ rocke_cshuffle_epilogue_t cepi - = rocke_cshuffle_epilogue_from_grid(ctx->atom, &ctx->grid, vec_c); + = ctx->is_wmma ? rocke_cshuffle_epilogue_from_grid_op(ctx->op, &ctx->grid, vec_c) + : rocke_cshuffle_epilogue_from_grid(ctx->atom, &ctx->grid, vec_c); cepi.out_dtype = dtype_d; rocke_cshuffle_epilogue_atomic_store(b, @@ -2781,9 +2906,12 @@ static bool wgrad_build_ctx_init(rocke_conv_build_ctx_t* ctx, if(spec->lds_k_outer) { a_sh[0] = ctx->block_k; - /* Direct load writes packed lane-contiguous bytes and cannot skip a - * row pad, so the async path uses a pad of 0. */ - const int kpad = spec->async_dma ? 0 : ROCKE_WGRAD_KOUTER_PAD; + /* The gfx9 direct load writes packed lane-contiguous bytes and + * cannot skip a row pad, so its async path uses a pad of 0; the + * gfx1250 async copy writes a per-lane address and keeps the pad. + * Mirrors Python _kouter_pad. */ + const int kpad + = (spec->async_dma && spec->wave_size == 64) ? 0 : ROCKE_WGRAD_KOUTER_PAD; a_sh[1] = ctx->block_m + kpad; b_sh[0] = ctx->block_k; b_sh[1] = ctx->block_n + kpad; @@ -2859,7 +2987,10 @@ static bool wgrad_build_ctx_init(rocke_conv_build_ctx_t* ctx, * (swapped coordinates on the K-outer tile) so it does not fall back to the * forward descriptor. */ ctx->a_descriptor_fn = spec->lds_k_outer ? wgrad_dy_descriptor_kouter : wgrad_dy_descriptor; - if(ctx->async_dma) + /* gfx1250 async (Python _wmma_async): reuses the synchronous K-outer + * loaders for the slot partition and swaps only the copy. */ + ctx->async_global_lds = spec->async_dma && ctx->is_wmma; + if(ctx->async_dma && !ctx->async_global_lds) { /* K-outer: rows are K_wg, columns are the free axis, so a chunk is a run * along the free axis at one k_wg -- contiguous in global for both @@ -2992,8 +3123,8 @@ static bool wgrad_build_ctx_init(rocke_conv_build_ctx_t* ctx, /* Schedule -- only compute the policy here; the caller emits the prologue * AFTER the buffer resources so the SSA order matches Python: * buffer_rsrc(dY/X/dW) -> schedule.emit_prologue(b) -> k-loop. */ - ctx->schedule - = rocke_schedule_policy_for_pipeline(b, ctx->async_dma ? "async_dma" : spec->pipeline); + ctx->schedule = rocke_schedule_policy_for_pipeline( + b, (ctx->async_dma && !ctx->async_global_lds) ? "async_dma" : spec->pipeline); return rocke_ir_builder_ok(b); } @@ -3923,6 +4054,9 @@ rocke_kernel_def_t* rocke_build_implicit_gemm_conv_wgrad( ctx.a_rsrc = a_rsrc.rsrc; ctx.b_rsrc = b_rsrc.rsrc; ctx.d_rsrc = d_rsrc.rsrc; + /* Plain global pointers for the gfx1250 async copy (no IR emitted). */ + ctx.a_ptr = dY; + ctx.b_ptr = X; /* Emit schedule prologue AFTER buffer resources -- mirrors Python ordering: * make_buffer_resource(dY/X/dW) then schedule.emit_prologue(b). */ @@ -3974,6 +4108,8 @@ rocke_kernel_def_t* rocke_build_implicit_gemm_conv_wgrad( rocke_conv_emit_kloop_unroll(&ctx); else if(spec->prefetch_addr) wgrad_emit_kloop_prefetch_addr(&ctx); + else if(ctx.async_global_lds) + rocke_conv_emit_kloop_async_global_lds(&ctx); else if(!spec->async_dma) rocke_conv_emit_kloop_simple(&ctx); else @@ -4006,11 +4142,13 @@ rocke_kernel_def_t* rocke_build_implicit_gemm_conv_wgrad( for(int i = 0; i < ctx.num_final_accs; ++i) ctx.final_accs[i] = post_accs[i]; - /* Same order as Python: the WMMA epilogue is chosen before the - * epilogue knob, since WMMA has no cshuffle. */ - if(ctx.is_wmma) + /* Same order as Python: WMMA takes its own atomic epilogue unless the + * cshuffle epilogue was asked for, which pairs through LDS on both + * families. */ + const bool cshuffle_epi = spec->epilogue && strcmp(spec->epilogue, "cshuffle") == 0; + if(ctx.is_wmma && !cshuffle_epi) wgrad_emit_split_k_epilogue_wmma(b, &ctx, spec, dW); - else if(spec->epilogue && strcmp(spec->epilogue, "cshuffle") == 0) + else if(cshuffle_epi) wgrad_emit_split_k_cshuffle_epilogue(b, &ctx, spec, dW, wg_M, wg_N); else wgrad_emit_split_k_epilogue_f32(b, &ctx, spec, dW, wg_M, wg_N); @@ -4066,7 +4204,8 @@ rocke_kernel_def_t* rocke_build_implicit_gemm_conv_wgrad( vec_c = 1; } rocke_cshuffle_epilogue_t cepi - = rocke_cshuffle_epilogue_from_grid(ctx.atom, &ctx.grid, vec_c); + = ctx.is_wmma ? rocke_cshuffle_epilogue_from_grid_op(ctx.op, &ctx.grid, vec_c) + : rocke_cshuffle_epilogue_from_grid(ctx.atom, &ctx.grid, vec_c); cepi.out_dtype = dtype_d; if(ctx.is_pointwise) { diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/instances/convolution.md b/dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/instances/convolution.md index 99d8929ed7df..6ef4dc6bda54 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/instances/convolution.md +++ b/dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/instances/convolution.md @@ -128,7 +128,7 @@ class ImplicitGemmConvSpec: `lds_k_pad=None` lets the kernel pick: `+8` on sync paths when `block_k >= 16`, `+0` on async paths. Override only for sweep experiments. -`lds_k_outer` is a backward-only field (`WgradConvSpec`, `DgradConvSpec`); the forward spec above has no counterpart. It stores the LDS tile K-outer as `T[k][mn]` and transposes inside the MFMA operand fetch instead of transposing on store. Its row stride does not come from `lds_k_pad`: it is fixed at `_KOUTER_PAD` — 8, or 0 on the wgrad async path — and wgrad rejects an explicit `lds_k_pad` under `lds_k_outer` rather than silently ignoring it. dgrad keeps its A tile M-outer, so `lds_k_pad` still applies there; dgrad instead rejects an explicit `lds_layout` or `async_dma`. Do not set the field by hand: ask `WgradConvSpec.default_lds_k_outer(...)` / `DgradConvSpec.default_lds_k_outer(...)`. The backward sections below cover the layout, the pad, and the gate. +`lds_k_outer` is a backward-only field (`WgradConvSpec`, `DgradConvSpec`); the forward spec above has no counterpart. It stores the LDS tile K-outer as `T[k][mn]` and transposes inside the MFMA operand fetch instead of transposing on store. Its row stride does not come from `lds_k_pad`: it is fixed at `_KOUTER_PAD` — 8, or 0 on the wave64 wgrad async path — and wgrad rejects an explicit `lds_k_pad` under `lds_k_outer` rather than silently ignoring it. dgrad keeps its A tile M-outer, so `lds_k_pad` still applies there; dgrad instead rejects an explicit `lds_layout` or `async_dma`. Do not set the field by hand: ask `WgradConvSpec.default_lds_k_outer(...)` / `DgradConvSpec.default_lds_k_outer(...)`. The backward sections below cover the layout, the pad, and the gate. The example bake-off (`example/ck_tile/dsl/08_bake_off_implicit_gemm`) uses: @@ -365,7 +365,7 @@ The removed writes were bank-degenerate by construction. In `"row"` mode the dec The K-outer row stride must **not** be a multiple of the 32-dword LDS bank period, or the transpose read degenerates. The row-walking lane term differs per regime but the condition does not: on wave64 it is `((l % 16) // 4)`, four row-groups; on wave32 it is `(l % 8)`, eight. Either contributes zero bank spread whenever `(stride_elems * 2 / 4) % 32 == 0`. A pad of 8 makes a 64-wide tile 36 dwords (`36 % 32 == 4`), spreading the row-groups across banks (`0, 4, 8, 12` on wave64; `0, 4, ..., 28` on wave32), and keeps rows 16-byte aligned for the `b128` store side. -The two directions carry the constant differently. dgrad fixes it at 8 as a module constant shared by the builder and `is_valid_dgrad_spec` so the charged and allocated shapes cannot drift (mirrored as `ROCKE_DGRAD_KOUTER_PAD`). wgrad derives it per build as `0 if spec.async_dma else 8`: the direct global→LDS load deposits packed lane-contiguous bytes and cannot skip a row pad. That is not free on the read side — a pad of 0 puts the row stride back at a whole multiple of the bank period, which is the degenerate case above, so the async leg gives up exactly the bank spread K-outer was introduced to obtain. `async_dma` trades read-side bank spread for the write-side and prefetch win of the direct path, which is why it is a swept axis rather than a deduced one. +The two directions carry the constant differently. dgrad fixes it at 8 as a module constant shared by the builder and `is_valid_dgrad_spec` so the charged and allocated shapes cannot drift (mirrored as `ROCKE_DGRAD_KOUTER_PAD`). wgrad derives it per build (`_kouter_pad`) as 0 on the wave64 async path and 8 otherwise: the gfx9 direct global→LDS load (`buffer_load_lds`) deposits packed lane-contiguous bytes and cannot skip a row pad. The gfx1250 async copy (`global_load_async_to_lds`) writes a per-lane LDS address, so the wave32 async path keeps the pad of 8. That is not free on the read side — a pad of 0 puts the row stride back at a whole multiple of the bank period, which is the degenerate case above, so the async leg gives up exactly the bank spread K-outer was introduced to obtain. `async_dma` trades read-side bank spread for the write-side and prefetch win of the direct path, which is why it is a swept axis rather than a deduced one. ### Fragment Length Is Per-Atom @@ -463,6 +463,7 @@ M-outer default for dgrad, and matches the fp32 reference for wgrad. - `pipeline="basic"`: now uses a bounded runtime `scf.for` K-loop (load → sync → MFMA → sync per tile), identical to the `"mem"` path. No K-trip-count limit. dgrad has no `basic` branch. - `async_dma` is a swept axis on the wgrad sweep driver, not a run-level flag. Unlike `lds_k_outer` it is not deducible from `(arch, spec)`: it removes the register staging of the tile, but it also forces the K-outer row pad to 0 and coarsens the load-width ladder to the widths the intrinsic accepts, and both terms are functions of tile width and channel run, which are themselves sweep axes. +- WMMA wgrad (gfx1250): `async_dma` is a ping-pong over `global_load_async_to_lds` (`CoalescedTileLoader.load_async_to_lds`, one barrier per tile, drained by `s_wait_asynccnt`). It reuses the synchronous K-outer loaders' slot partition, needs `lds_k_outer` and a 4-, 8- or 16-byte chunk on both operands, and ANDs the buffer-range check (`byte offset < num_bytes`) into each chunk's `valid`, zero-filling failing chunks with an LDS store. `epilogue="cshuffle"` is admitted on WMMA too: both cshuffle epilogues (direct store and split-K atomic) read the WMMA fragment through `CShuffleEpilogue.from_grid_op`. `pipeline` stays `mem`; `unroll_k`, `prefetch_addr` and `chiplet_swizzle` stay rejected on WMMA. - No K loop is unrolled at build time any more: `async_dma` runs `SoftwarePipeline.run_ping_pong_dynamic` and `unroll_k` a hand-rolled 2x ping-pong, both over the runtime extent, so the former `_MAX_UNROLLED_K_ITERS` / `ROCKE_MAX_UNROLLED_K_ITERS` cap is gone. - wgrad `split_k`: `1` = no split (direct store); any value `> 1` builds the one split-K kernel (name tag `spk`), whose degree is the `ks_count` / `ks` kernargs chosen at launch. There is no `split_k=0` "runtime degree" encoding. - Removed: the `--lds-k-outer`, `--lds-k-pad` and `--dtype-d` CLI flags and the `ROCKE_WGRAD_LDS_K_OUTER` env override. All replaced by deduction or dropped. diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/loads.py b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/loads.py index fc242fae4a5a..a3e5f6e336f4 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/loads.py +++ b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/loads.py @@ -506,6 +506,67 @@ def load( v = b.global_load_vN(ptr, off_elems, dtype, self.load_vec) self._store_tile(b, smem_dst, row, col, v) + def load_async_to_lds( + self, + b: IRBuilder, + *, + tid: Value, + smem_dst: Value, + descriptor: DescriptorFn, + ptr: Value, + num_bytes: Value, + coherency: int = 0, + ) -> None: + """gfx1250 async variant of :meth:`load` (``global_load_async_to_lds``). + + Same slot partition and descriptor contract as :meth:`load` in + ``vector_axis="col"`` mode, but each in-range chunk is one per-lane + async global->LDS copy of ``load_vec`` elements to ``smem_dst[row, col]`` + instead of a register round trip. The copy reads through a plain global + pointer with no buffer-resource range check, so the check the resource + would make -- byte offset below ``num_bytes`` -- is ANDed into the + descriptor's ``valid``. Descriptors lean on that check: a reduction + tail past the tensor, or a free-axis overhang at its last element, is + reported valid and only the resource turns it into zero. A chunk that + fails either test is zero-filled with an ordinary LDS store. The caller + drains the copies with ``s_wait_asynccnt`` (and the zero stores with an + LDS barrier) before reading the tile. + """ + if self.vector_axis != "col": + raise ValueError( + "CoalescedTileLoader.load_async_to_lds supports only vector_axis='col' " + "(the copy cannot transpose on store)" + ) + elem_bytes = _ELEM_BYTES.get(self.elem_dtype.name) + if elem_bytes is None: + raise ValueError( + f"CoalescedTileLoader: unsupported elem_dtype {self.elem_dtype.name!r}" + ) + width = self.load_vec * elem_bytes + if width not in (4, 8, 16): + raise ValueError( + f"global_load_async_to_lds moves 4, 8 or 16 bytes per lane; a " + f"{self.load_vec} x {elem_bytes} B chunk is {width} B" + ) + zero = b.zero_vec(self.elem_dtype, self.load_vec) + c_elem_bytes = b.const_i32(elem_bytes) + for row, col in self.plan_slots(b, tid=tid): + off_elems, valid = descriptor(b, row, col) + in_bounds = b.cmp_lt(b.mul(off_elems, c_elem_bytes), num_bytes) + valid = in_bounds if valid is None else b.land(valid, in_bounds) + with b.scf_if_else(valid) as (in_range, out_of_range): + with in_range: + b.global_load_async_to_lds( + ptr, + off_elems, + smem_dst, + [row, col], + width_bytes=width, + coherency=coherency, + ) + with out_of_range: + b.smem_store_vN(smem_dst, [row, col], zero, self.load_vec) + def fetch( self, b: IRBuilder, diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/instances/SUPPORT_MATRIX.md b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/instances/SUPPORT_MATRIX.md index 13b84f476ab0..1409ecefd95d 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/instances/SUPPORT_MATRIX.md +++ b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/instances/SUPPORT_MATRIX.md @@ -109,7 +109,7 @@ be smaller than 2 GiB. | `direct_depthwise_dgrad` (`DirectDepthwiseDgradStreamSpec`) | ✅ | ✅ | ❌ | `cpg=kpg=1`, any stride and padding; row-streaming scalar FMA at the target's wave size (wave64, or wave32 on gfx1250); `cpl` 2/4/8 channels per lane with packed `<2 x f32>` FMAs (`groups % cpl == 0`, `cpl*KH*KW <= 200`, `cpl*block_w*KH <= 192`). In the AOT direct cache; not yet wired into dispatch | | `direct_depthwise_tiled_dgrad` (`DirectDepthwiseTiledSpec(dgrad=True)`) | ✅ | ✅ | ❌ | `cpg=kpg=1`, stride 1, the forward tiled kernel with the filter rotated by 180 degrees; filters up to 64 on a side incl. non-square (51x5 / 5x51), per-axis padding (`pad_w`) with `PAD <= KH-1`, `PAD_W <= KW-1`. Served by the direct-conv AOT cache, not the dispatcher | | `conv_implicit_gemm_dgrad` stride-1 knobs | ✅ | ✅ | ❌ | folded sub-GEMM record (stride capability), `tap_outer_k` (wave64 MFMA, `kpg % tile_k == 0`), `dy_halo` / `dy_halo_setprio` / `dy_halo_kouter_pad` / `halo_max_wo` (on the tap-outer path, `Ho == Hi`, `Wo == Wi`; the pad needs `lds_k_outer`, gfx950); `xcd_tile_order` remaps only on gfx950. AOT grid: stride-1 tap-outer binaries next to the tilde ones, halo sub-grid on gfx950; not yet wired into dispatch | -| `conv_implicit_gemm_wgrad` | ✅ | ✅ | ✅ | gfx1151/gfx1201/gfx1250: WMMA, `mem`+`default`; split-K (WMMA packed atomic, no two-stage) on gfx1250 only, automatic for bf16/fp32 and explicit-only for fp16, `split_k=1` on gfx1151/gfx1201 | +| `conv_implicit_gemm_wgrad` | ✅ | ✅ | ✅ | gfx1151/gfx1201/gfx1250: WMMA, `mem` with the `default` or `cshuffle` epilogue (gfx1250 adds the `async_dma` K loop over `global_load_async_to_lds`); split-K (WMMA packed atomic, no two-stage) on gfx1250 only, automatic for bf16/fp32 and explicit-only for fp16, `split_k=1` on gfx1151/gfx1201 | | `conv_implicit_gemm_wgrad` stream-K | ✅ | ✅ | ❌ | CK iteration-balanced partition; `linear`/`tree` deterministic fixups need the cross-workgroup flag ops | | `conv_implicit_gemm` split-K forward (`split_k > 1`) | ❌ | ❌ | ❌ | gfx1250 only: WMMA, `mem` or `wavelet` pipeline, 2-D, not pointwise, identity `acc_epilogue`; no `async_dma`, `unroll_k`, `chiplet_swizzle` or `group_merge`. Each slice stores an fp32 partial tile into its slab of a `[ks_count][M][K]` workspace and `conv_fwd_split_k_reduce` folds the slabs into D; the degree is the `ks_count` kernarg, and `split_k = -1` resolves it (`resolve_conv_split_k`; 1 where splitting does not pay). Built into the AOT kernel cache; the dispatcher does not select it | | `conv_implicit_gemm` cluster split-K (`split_k_cluster`) | ❌ | ❌ | ❌ | gfx1250 only (llvm23, workgroup-cluster launch): `mem` pipeline, `split_k_cluster` in 2..15, `tile_n % 4 == 0`, K and kpg multiples of 4 (checked at build and on every launched problem). The slices of one output tile form a (1, 1, S) workgroup cluster and fold through a cluster barrier, so there is no Stage 2 launch; one binary per S. Built into the AOT kernel cache; the dispatcher does not select it | diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/tests/golden/rocke_representative_ir_sha256/conv_wgrad.json b/dnn-providers/hip-kernel-provider/rocke/platform/tests/golden/rocke_representative_ir_sha256/conv_wgrad.json index 28d269461b91..05af6e5d5761 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/tests/golden/rocke_representative_ir_sha256/conv_wgrad.json +++ b/dnn-providers/hip-kernel-provider/rocke/platform/tests/golden/rocke_representative_ir_sha256/conv_wgrad.json @@ -27,6 +27,22 @@ "sha256": "4cb4af882d9ba2405ac67c3d5c0cf7ddb6a893bed24d6c2496f7ba542187767c", "status": "ok" }, + "conv_wgrad/gfx1250/n1h8c16k32r3_spk4_async": { + "arch": "gfx1250", + "bytes": 85370, + "family": "conv_wgrad", + "kernel_name": "irhash_wgrad_1250_spk4_async_N1H8W8C16_K32Y3X3_t32x32x32_w2x1_a16x16x32_mem_default_async_kouter_spk", + "sha256": "94bb64f47a0c3398035207801b87b35e5ebd85895e382892f495f3c4302dfb31", + "status": "ok" + }, + "conv_wgrad/gfx1250/n1h8c16k32r3_spk4_cshuffle": { + "arch": "gfx1250", + "bytes": 54947, + "family": "conv_wgrad", + "kernel_name": "irhash_wgrad_1250_spk4_cshuffle_N1H8W8C16_K32Y3X3_t32x32x32_w2x1_a16x16x32_mem_cshuffle_spk", + "sha256": "d6709d4327c8ccf1074a9bc864afedd7fef77891fe804886ea1df4160e9e1c07", + "status": "ok" + }, "conv_wgrad/gfx1250/n1h8c64k64r3_g4": { "arch": "gfx1250", "bytes": 66522, @@ -208,6 +224,22 @@ "sha256": "f6bb79207ed082a136a1b1aa7d33730d33be31f729607ed99126d04e74b095ee", "status": "ok" }, + "conv_wgrad/gfx1250/n1h8c16k32r3_spk4_async": { + "arch": "gfx1250", + "bytes": 85509, + "family": "conv_wgrad", + "kernel_name": "irhash_wgrad_1250_spk4_async_N1H8W8C16_K32Y3X3_t32x32x32_w2x1_a16x16x32_mem_default_async_kouter_spk", + "sha256": "e7a41f7a1884f347467739385ef4979ad5f90a94fa99387f6ef07f4b6bbd1ceb", + "status": "ok" + }, + "conv_wgrad/gfx1250/n1h8c16k32r3_spk4_cshuffle": { + "arch": "gfx1250", + "bytes": 55086, + "family": "conv_wgrad", + "kernel_name": "irhash_wgrad_1250_spk4_cshuffle_N1H8W8C16_K32Y3X3_t32x32x32_w2x1_a16x16x32_mem_cshuffle_spk", + "sha256": "cce59f1ea889507c62d4abe9c141e5e6a889cb8dbb7c816730ba1362e6dede69", + "status": "ok" + }, "conv_wgrad/gfx1250/n1h8c64k64r3_g4": { "arch": "gfx1250", "bytes": 66661, @@ -389,6 +421,22 @@ "sha256": "599ca1e5afea44a50511e99fae46d7ecc37ee9c9c11bd5960aa6b04d3e15ce17", "status": "ok" }, + "conv_wgrad/gfx1250/n1h8c16k32r3_spk4_async": { + "arch": "gfx1250", + "bytes": 85573, + "family": "conv_wgrad", + "kernel_name": "irhash_wgrad_1250_spk4_async_N1H8W8C16_K32Y3X3_t32x32x32_w2x1_a16x16x32_mem_default_async_kouter_spk", + "sha256": "00734689dde7756b661433f5a481b149a5b505d7418d156853288ea0f317c8a8", + "status": "ok" + }, + "conv_wgrad/gfx1250/n1h8c16k32r3_spk4_cshuffle": { + "arch": "gfx1250", + "bytes": 55150, + "family": "conv_wgrad", + "kernel_name": "irhash_wgrad_1250_spk4_cshuffle_N1H8W8C16_K32Y3X3_t32x32x32_w2x1_a16x16x32_mem_cshuffle_spk", + "sha256": "0d4ff4ddc79acf33334ae99b5d2f2da8afd6199f35b0638c3e4d448d45d6a087", + "status": "ok" + }, "conv_wgrad/gfx1250/n1h8c64k64r3_g4": { "arch": "gfx1250", "bytes": 66725, @@ -570,6 +618,22 @@ "sha256": "599ca1e5afea44a50511e99fae46d7ecc37ee9c9c11bd5960aa6b04d3e15ce17", "status": "ok" }, + "conv_wgrad/gfx1250/n1h8c16k32r3_spk4_async": { + "arch": "gfx1250", + "bytes": 85573, + "family": "conv_wgrad", + "kernel_name": "irhash_wgrad_1250_spk4_async_N1H8W8C16_K32Y3X3_t32x32x32_w2x1_a16x16x32_mem_default_async_kouter_spk", + "sha256": "00734689dde7756b661433f5a481b149a5b505d7418d156853288ea0f317c8a8", + "status": "ok" + }, + "conv_wgrad/gfx1250/n1h8c16k32r3_spk4_cshuffle": { + "arch": "gfx1250", + "bytes": 55150, + "family": "conv_wgrad", + "kernel_name": "irhash_wgrad_1250_spk4_cshuffle_N1H8W8C16_K32Y3X3_t32x32x32_w2x1_a16x16x32_mem_cshuffle_spk", + "sha256": "0d4ff4ddc79acf33334ae99b5d2f2da8afd6199f35b0638c3e4d448d45d6a087", + "status": "ok" + }, "conv_wgrad/gfx1250/n1h8c64k64r3_g4": { "arch": "gfx1250", "bytes": 66725, diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/tests/instances/differential/golden/llvm_gfx_all.json b/dnn-providers/hip-kernel-provider/rocke/platform/tests/instances/differential/golden/llvm_gfx_all.json index 4d314fa06b45..9c1372fc5e3a 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/tests/instances/differential/golden/llvm_gfx_all.json +++ b/dnn-providers/hip-kernel-provider/rocke/platform/tests/instances/differential/golden/llvm_gfx_all.json @@ -562,7 +562,13 @@ "44": "e947ae1c080472d24798cab77fdb22be23224485379cb15e1abdcd9f8672d679", "48": "1a8c8001b691eba5e091506a1410b22363bf9fd48010921a1042d715f5b2779c", "5": "a8044100dfbe89a48ecf9af672bc80353812d0bfca61eaa7928bd0ef88f077f9", + "53": "51f6686426cc43a0130ddde2511b0aedc5e5636607b9ea523744998770ea0664", + "55": "e778edf08e2ed9999a24b747652c4d0fd3a3ea64dd68d922c82abb206bfe98e3", + "56": "87bff7628af8465e95e1dd9ffa12b05f5089c838a00ab986db749c777406eee5", + "58": "7fe2785d761f330bdc80afd6156e79f426ee197f61ec3bed99d2107c04fa90bd", + "59": "5ebbd1600763884bb810b302df64677a5fef6f222bfec3fac9c470a4ad20ed5d", "6": "779030de0f57eb186f48efc6b92827294b7aa1095c2a8ee96d2edae002cbaa1d", + "60": "1fda68fce4ddbfd4d1ab27127e91474868de275962bdb26e3be025ca21cb2d85", "7": "1106d4ce71314f3c532427fed8ceb00855f1cf1dd7f5117b89234cf326583087", "8": "82d8923b31be6a07c4371fbe40b33a2ea7452027299ce0993e903cf55ad092d8" }, @@ -2429,7 +2435,13 @@ "44": "8835b927ea9575fe7c32635a5c04b0cb6524ec242821f44f8e93fcb5a40cfebd", "48": "ee78fea15745bbe7cd6f54557c1c0e94b86c907758e13a1167881df99ae92b6f", "5": "9d32ae9f91e43f62e072658dba4bd2ae5f3dc07467ef6e9e83a1f067f1a638db", + "53": "2361cf65c6a7b3277143bc06f14ec547e626d361d6447d5247d29908321da28b", + "55": "845dd21a850297763ad037d18dcfdb89bb403142039d77d26d4b2d337d761012", + "56": "d22b27471808ec88c11cd416b5f1c0214b395332c65514ab2fc2e7c05499b42d", + "58": "0fd30468be138821abaf30432c005cc82999425fff3020f0746ff78c199f1a68", + "59": "46dd718d4f966e5b0659ebe0f8b2f6fde2ca6eac6035fd247c0b7f44df648408", "6": "b2605886350eb7b793616938e9179ad480084552a3533959f34918085a34484c", + "60": "e1f8d89e18ad5bb1b9e918a2df43fca9283adf375de2488619a56861a2b0cd3b", "7": "d653088ddce4d716c3b9e20fabe5592541b52ba5e800786a86b660763c25a596", "8": "889d9e3617a8eb2e30d29db4ef8da0aeac52c9fa0d92bc53541947cf85926c08" }, @@ -4198,7 +4210,13 @@ "44": "bb0541f7a5814241e9998ccba12441e518d93210881fb4d78faa34e99a7414c1", "48": "529267374fd88af7746495caff9f9fb19c8e3bdac84cab09af9e946c5e429b9b", "5": "cb1aab374b1c257ee44efdb653a51cedfef95f2dbc27c959a66215db38e0db35", + "53": "335c1050cc3ecc5c2288292ca4fa9b735f3702bec64305e7e7b1a764a6301031", + "55": "d4f321c49e9be50002075621fef0fd6e1c5d3085adc584be4a3a317582eaaf79", + "56": "6a7c3fec1ae26dd7faf162315dcbd9fdd92becdd4704ef81dbc1ed117444e5ca", + "58": "95aabfecf90f2021a781953eed801d3b14d86c4c5844c365dc6edd83416a0512", + "59": "654c0bfc412129aad45be434cb4e382fab2739ad94ffe5d607db6924e5f77b00", "6": "f3b4e8cf7ab611637137babc7994172995f76262274b0a63ca3b6999ac62ef04", + "60": "b3e910d01f76ec19f85babd5c26ed1debd157233c2560492711967b843de603a", "7": "e12021bd6f8923472a7595394ecde03a55dcd07d63d714b9adaea3efa79df887", "8": "2cd90bffada0eaec9a90caf5d5b76f2f213e4e884e3f413b13e79eef3071ec82" }, @@ -5997,7 +6015,13 @@ "44": "28b3e5675e42d727cfcc5c4706e5b8d432f6d635ae68370960063468c5777683", "48": "5ea399799a9a1af9e69a9dae94c3194e66f277472b290592869ba2deae88acab", "5": "0826f9183484b30541afcda4322fac4e5fec28374efb8995a373b9d227a97016", + "53": "bc38f32697b74c7208cd101293556b123c54cfe96b4469fc0562c23052b10fcb", + "55": "0c59e7ef5797f03be197533617fcb598f79e7df6759ca8fb3bf875a53e395731", + "56": "acf39c44252e4063825560e706c92070ec2532daca3f1072453e4b9bb747681a", + "58": "02b4d2adce33a29c4a426bc3c4617387e340a2e19b429e986468c6064dac2759", + "59": "fbcf10500498604862078885c4482d5cbe67269770db5929f39fa0d6919d3c89", "6": "2726776342b4ffd858fce5951d771a50515d7c143e07656eb6235e10fd08bd32", + "60": "916897d1f0ed7b3a09a499a4b6f5128ee271fc7bf4e1acbfa8fb2980b00d9ac9", "7": "44a58b32e3d95d8492b23193d701925de3dd9de8cddf8a2c006b6fdf5de87cf8", "8": "af0df509c9e1bfa9bcb32621a677b633c9bd7c0b0a22012ee20acd2894b44df8" }, @@ -8051,7 +8075,13 @@ "5": "0826f9183484b30541afcda4322fac4e5fec28374efb8995a373b9d227a97016", "50": "b54d4d875110a0dbd2b0c42b5844114122cd9f74b075e03ada23422cc62c185f", "51": "84826a6d0c1344658e348d3010e06274415a4badc09580b3bcf9c4d5ef102e73", + "53": "bc38f32697b74c7208cd101293556b123c54cfe96b4469fc0562c23052b10fcb", + "55": "0c59e7ef5797f03be197533617fcb598f79e7df6759ca8fb3bf875a53e395731", + "56": "acf39c44252e4063825560e706c92070ec2532daca3f1072453e4b9bb747681a", + "58": "02b4d2adce33a29c4a426bc3c4617387e340a2e19b429e986468c6064dac2759", + "59": "fbcf10500498604862078885c4482d5cbe67269770db5929f39fa0d6919d3c89", "6": "2726776342b4ffd858fce5951d771a50515d7c143e07656eb6235e10fd08bd32", + "60": "916897d1f0ed7b3a09a499a4b6f5128ee271fc7bf4e1acbfa8fb2980b00d9ac9", "7": "44a58b32e3d95d8492b23193d701925de3dd9de8cddf8a2c006b6fdf5de87cf8", "8": "af0df509c9e1bfa9bcb32621a677b633c9bd7c0b0a22012ee20acd2894b44df8" }, diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/tests/instances/rocke_ir_parity_harness.py b/dnn-providers/hip-kernel-provider/rocke/platform/tests/instances/rocke_ir_parity_harness.py index c8be1742c44f..1866e85cad4d 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/tests/instances/rocke_ir_parity_harness.py +++ b/dnn-providers/hip-kernel-provider/rocke/platform/tests/instances/rocke_ir_parity_harness.py @@ -196,6 +196,7 @@ def build_conv_wgrad( streamk="off", streamk_reduction="linear", streamk_ctas=-1, + **spec_extra, ): def _build(): from kernels.common.conv_implicit_gemm_wgrad import ( @@ -228,6 +229,7 @@ def _build(): streamk=streamk, streamk_reduction=streamk_reduction, streamk_ctas=streamk_ctas, + **spec_extra, ) return build_implicit_gemm_conv_wgrad(spec, arch=arch) @@ -2810,6 +2812,49 @@ def add(family, case_id, arch, build): dtype_d="bf16", ), ) + # WMMA split-K over the gfx1250 async K loop (global_load_async_to_lds). + add( + "conv_wgrad", + "conv_wgrad/gfx1250/n1h8c16k32r3_spk4_async", + "gfx1250", + build_conv_wgrad( + "irhash_wgrad_1250_spk4_async", + "gfx1250", + wgrad1, + wave_size=32, + wtm=16, + wtn=16, + wtk=32, + tile_m=32, + tile_n=32, + tile_k=32, + split_k=4, + dtype_d="bf16", + lds_k_outer=True, + async_dma=True, + ), + ) + # WMMA split-K through the cshuffle epilogue (paired atomics from LDS). + add( + "conv_wgrad", + "conv_wgrad/gfx1250/n1h8c16k32r3_spk4_cshuffle", + "gfx1250", + build_conv_wgrad( + "irhash_wgrad_1250_spk4_cshuffle", + "gfx1250", + wgrad1, + wave_size=32, + wtm=16, + wtn=16, + wtk=32, + tile_m=32, + tile_n=32, + tile_k=32, + split_k=4, + dtype_d="fp16", + epilogue="cshuffle", + ), + ) # Dgrad: backward-data implicit-GEMM. Problem args match the forward conv # set; atoms chosen to be valid on each arch. Stride=2 exercises the tiled From 5e4782116ae8f7f16dbf481ff2fd7f0c9c736d35 Mon Sep 17 00:00:00 2001 From: Jakub Piasecki Date: Sun, 1 Nov 2026 17:59:02 +0000 Subject: [PATCH 02/13] fix(rocke): lower the packed fp16 atomic add as a monotonic atomicrmw The legacy llvm.amdgcn.global.atomic.fadd.v2f16 intrinsic is lowered as a fenced atomic: a global_wb / global_inv pair (buffer_wbl2 / buffer_inv on gfx950) and full counter drains around every single add, which serialised the split-K epilogues. Emit the generic agent-scope monotonic atomicrmw fadd <2 x half> with the device-local memory-model metadata, as the packed bf16 add already does; the backend selects the same global_atomic_pk_add_f16 without the fences. Mirrored in the C++ engine (lower_llvm/mem.cpp). The intrinsic declaration has no other user, so it is dropped from both engines' declaration tables and from the intrinsic arch-domain columns. The HIP lowering keeps __builtin_amdgcn_global_atomic_fadd_v2f16, which clang already emits unfenced. Every kernel with a packed fp16 atomic changes its IR: the fp16 split-K wgrad configs and the fp16 atomic dgrad epilogue configs. The parity golden and the representative-IR golden are re-blessed for exactly those keys. Co-Authored-By: Claude Opus 5.5 --- .../platform/cpp/core/lower_llvm/data.cpp | 3 - .../platform/cpp/core/lower_llvm/mem.cpp | 19 ++++-- .../data/intrinsic_arch_domain.llvm20.json | 36 ---------- .../data/intrinsic_arch_domain.llvm22.json | 35 ---------- .../data/intrinsic_arch_domain.llvm23.json | 35 ---------- .../data/intrinsic_arch_domain.llvm24.json | 35 ---------- .../rocke/platform/python/rocke/core/ir.py | 5 +- .../platform/python/rocke/core/lower_llvm.py | 26 ++++---- .../conv_wgrad.json | 64 +++++++++--------- .../differential/golden/llvm_gfx_all.json | 66 +++++++++---------- 10 files changed, 96 insertions(+), 228 deletions(-) diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/lower_llvm/data.cpp b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/lower_llvm/data.cpp index 4f32ee9d47be..451a87b912b1 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/lower_llvm/data.cpp +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/lower_llvm/data.cpp @@ -237,9 +237,6 @@ const rocke_ll_decl_t ROCKE_LL_INTRINSIC_DECLS[] = { {"update.dpp.i32", "declare i32 @llvm.amdgcn.update.dpp.i32(i32, i32, i32 immarg, i32 immarg, i32 immarg, i1 " "immarg)"}, - {"global.atomic.fadd.v2f16", - "declare <2 x half> @llvm.amdgcn.global.atomic.fadd.v2f16.p1(" - "ptr addrspace(1), <2 x half>)"}, {"mbcnt.lo", "declare i32 @llvm.amdgcn.mbcnt.lo(i32, i32)"}, {"mbcnt.hi", "declare i32 @llvm.amdgcn.mbcnt.hi(i32, i32)"}, {"ds.read.tr16.b64", "declare <4 x i16> @llvm.amdgcn.ds.read.tr16.b64(ptr addrspace(3))"}, diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/lower_llvm/mem.cpp b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/lower_llvm/mem.cpp index b4a59605c3f1..15ace4abb44e 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/lower_llvm/mem.cpp +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/lower_llvm/mem.cpp @@ -393,20 +393,29 @@ static void op_memref_global_atomic_add_pk_bf16(rocke_lower_t* L, const rocke_op static void op_memref_global_atomic_add_pk_f16(rocke_lower_t* L, const rocke_op_t* op) { + /* Mirror of Python _op_memref_global_atomic_add_pk_f16: a generic + * atomicrmw fadd <2 x half> (agent scope, monotonic) with the AMDGPU + * memory-model metadata, as the packed bf16 add does. The backend selects + * the native global_atomic_pk_add_f16. The legacy + * llvm.amdgcn.global.atomic.fadd.v2f16 intrinsic it replaces is lowered on + * gfx1250 as a fenced atomic (global_wb / global_inv and counter drains + * around every add). */ const rocke_value_t* ptr = op->operands[0]; const rocke_value_t* idx = op->operands[1]; const rocke_value_t* val = op->operands[2]; - const char* gep; - rocke_ll_need(L, "global.atomic.fadd.v2f16"); - gep = rocke_ll_fresh(L, "gep"); + const char* gep = rocke_ll_fresh(L, "gep"); rocke_ll_emitf(L, " %s = getelementptr inbounds half, ptr addrspace(1) %s, i32 %s", gep, rocke_ll_operand(L, ptr), rocke_ll_operand(L, idx)); + L->needs_fp_atomic_md = true; rocke_ll_emitf(L, - " %s = call <2 x half> @llvm.amdgcn.global.atomic.fadd.v2f16.p1(" - "ptr addrspace(1) %s, <2 x half> %s)", + " %s = atomicrmw fadd ptr addrspace(1) %s, <2 x half> %s " + "syncscope(\"agent\") monotonic" + ", !amdgpu.no.fine.grained.memory !1" + ", !amdgpu.no.remote.memory !1" + ", !amdgpu.ignore.denormal.mode !1", ll_res(op), gep, rocke_ll_operand(L, val)); diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/core/arch/data/intrinsic_arch_domain.llvm20.json b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/core/arch/data/intrinsic_arch_domain.llvm20.json index 6ae3214f30fe..825d3b664bd9 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/core/arch/data/intrinsic_arch_domain.llvm20.json +++ b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/core/arch/data/intrinsic_arch_domain.llvm20.json @@ -42,7 +42,6 @@ "fmuladd.v8bf16": "llvm.fmuladd.v8bf16", "fmuladd.v8f16": "llvm.fmuladd.v8f16", "fmuladd.v8f32": "llvm.fmuladd.v8f32", - "global.atomic.fadd.v2f16": "(auto-upgraded)", "global.load.lds": "llvm.amdgcn.global.load.lds", "global.load.tr.b128.v8bf16": "llvm.amdgcn.global.load.tr.b128.v8bf16", "global.load.tr.b128.v8f16": "llvm.amdgcn.global.load.tr.b128.v8f16", @@ -2728,41 +2727,6 @@ "verified_on": "llvm20" } }, - "global.atomic.fadd.v2f16": { - "gfx11-generic": { - "status": "ok", - "verified_on": "llvm20" - }, - "gfx1151": { - "status": "ok", - "verified_on": "llvm20" - }, - "gfx12-generic": { - "status": "ok", - "verified_on": "llvm20" - }, - "gfx1201": { - "status": "ok", - "verified_on": "llvm20" - }, - "gfx1250": { - "evidence": "invalid target ID", - "status": "target_unsupported", - "verified_on": "llvm20" - }, - "gfx90a": { - "status": "ok", - "verified_on": "llvm20" - }, - "gfx942": { - "status": "ok", - "verified_on": "llvm20" - }, - "gfx950": { - "status": "ok", - "verified_on": "llvm20" - } - }, "global.load.async.to.lds.b128": { "gfx11-generic": { "evidence": "not an intrinsic in llvm20", diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/core/arch/data/intrinsic_arch_domain.llvm22.json b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/core/arch/data/intrinsic_arch_domain.llvm22.json index f0808d6a705b..1aff8864f74b 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/core/arch/data/intrinsic_arch_domain.llvm22.json +++ b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/core/arch/data/intrinsic_arch_domain.llvm22.json @@ -63,7 +63,6 @@ "fmuladd.v8bf16": "llvm.fmuladd.v8bf16", "fmuladd.v8f16": "llvm.fmuladd.v8f16", "fmuladd.v8f32": "llvm.fmuladd.v8f32", - "global.atomic.fadd.v2f16": "(auto-upgraded)", "global.load.async.to.lds.b128": "llvm.amdgcn.global.load.async.to.lds.b128", "global.load.async.to.lds.b32": "llvm.amdgcn.global.load.async.to.lds.b32", "global.load.async.to.lds.b64": "llvm.amdgcn.global.load.async.to.lds.b64", @@ -2645,40 +2644,6 @@ "verified_on": "llvm22" } }, - "global.atomic.fadd.v2f16": { - "gfx11-generic": { - "status": "ok", - "verified_on": "llvm22" - }, - "gfx1151": { - "status": "ok", - "verified_on": "llvm22" - }, - "gfx12-generic": { - "status": "ok", - "verified_on": "llvm22" - }, - "gfx1201": { - "status": "ok", - "verified_on": "llvm22" - }, - "gfx1250": { - "status": "ok", - "verified_on": "llvm22" - }, - "gfx90a": { - "status": "ok", - "verified_on": "llvm22" - }, - "gfx942": { - "status": "ok", - "verified_on": "llvm22" - }, - "gfx950": { - "status": "ok", - "verified_on": "llvm22" - } - }, "global.load.async.to.lds.b128": { "gfx11-generic": { "evidence": "fatal error: error in backend: Cannot select: intrinsic %llvm.amdgcn.global.load.async.to.lds.b128", diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/core/arch/data/intrinsic_arch_domain.llvm23.json b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/core/arch/data/intrinsic_arch_domain.llvm23.json index 03dee4e0cfbc..6f47e060f1b0 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/core/arch/data/intrinsic_arch_domain.llvm23.json +++ b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/core/arch/data/intrinsic_arch_domain.llvm23.json @@ -68,7 +68,6 @@ "fmuladd.v8bf16": "llvm.fmuladd.v8bf16", "fmuladd.v8f16": "llvm.fmuladd.v8f16", "fmuladd.v8f32": "llvm.fmuladd.v8f32", - "global.atomic.fadd.v2f16": "(auto-upgraded)", "global.load.async.to.lds.b128": "llvm.amdgcn.global.load.async.to.lds.b128", "global.load.async.to.lds.b32": "llvm.amdgcn.global.load.async.to.lds.b32", "global.load.async.to.lds.b64": "llvm.amdgcn.global.load.async.to.lds.b64", @@ -2623,40 +2622,6 @@ "verified_on": "llvm23" } }, - "global.atomic.fadd.v2f16": { - "gfx11-generic": { - "status": "ok", - "verified_on": "llvm23" - }, - "gfx1151": { - "status": "ok", - "verified_on": "llvm23" - }, - "gfx12-generic": { - "status": "ok", - "verified_on": "llvm23" - }, - "gfx1201": { - "status": "ok", - "verified_on": "llvm23" - }, - "gfx1250": { - "status": "ok", - "verified_on": "llvm23" - }, - "gfx90a": { - "status": "ok", - "verified_on": "llvm23" - }, - "gfx942": { - "status": "ok", - "verified_on": "llvm23" - }, - "gfx950": { - "status": "ok", - "verified_on": "llvm23" - } - }, "global.load.async.to.lds.b128": { "gfx11-generic": { "evidence": "fatal error: error in backend: Cannot select: intrinsic %llvm.amdgcn.global.load.async.to.lds.b128", diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/core/arch/data/intrinsic_arch_domain.llvm24.json b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/core/arch/data/intrinsic_arch_domain.llvm24.json index ff00d64bdfe8..c32a44b2064b 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/core/arch/data/intrinsic_arch_domain.llvm24.json +++ b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/core/arch/data/intrinsic_arch_domain.llvm24.json @@ -68,7 +68,6 @@ "fmuladd.v8bf16": "llvm.fmuladd.v8bf16", "fmuladd.v8f16": "llvm.fmuladd.v8f16", "fmuladd.v8f32": "llvm.fmuladd.v8f32", - "global.atomic.fadd.v2f16": "(auto-upgraded)", "global.load.async.to.lds.b128": "llvm.amdgcn.global.load.async.to.lds.b128", "global.load.async.to.lds.b32": "llvm.amdgcn.global.load.async.to.lds.b32", "global.load.async.to.lds.b64": "llvm.amdgcn.global.load.async.to.lds.b64", @@ -2623,40 +2622,6 @@ "verified_on": "llvm24" } }, - "global.atomic.fadd.v2f16": { - "gfx11-generic": { - "status": "ok", - "verified_on": "llvm24" - }, - "gfx1151": { - "status": "ok", - "verified_on": "llvm24" - }, - "gfx12-generic": { - "status": "ok", - "verified_on": "llvm24" - }, - "gfx1201": { - "status": "ok", - "verified_on": "llvm24" - }, - "gfx1250": { - "status": "ok", - "verified_on": "llvm24" - }, - "gfx90a": { - "status": "ok", - "verified_on": "llvm24" - }, - "gfx942": { - "status": "ok", - "verified_on": "llvm24" - }, - "gfx950": { - "status": "ok", - "verified_on": "llvm24" - } - }, "global.load.async.to.lds.b128": { "gfx11-generic": { "evidence": "fatal error: error in backend: Cannot select: intrinsic %llvm.amdgcn.global.load.async.to.lds.b128", diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/core/ir.py b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/core/ir.py index a1b2d47d3265..b052dfad04ae 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/core/ir.py +++ b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/core/ir.py @@ -2433,8 +2433,9 @@ def global_atomic_add_pk_f16( ) -> Value: """Packed-fp16 atomic add: two fp16 lanes per transaction. - Lowers to AMDGPU's ``llvm.amdgcn.global.atomic.fadd.v2f16`` - intrinsic (gfx940+); returns the pre-add value. ``value`` + Lowers to a generic ``atomicrmw fadd <2 x half>`` (agent scope) + with the AMDGPU memory-model metadata, which the backend selects as + ``global_atomic_pk_add_f16``; returns the pre-add value. ``value`` must be a ``<2 x f16>`` vector and the pointer must reach into an fp16 buffer with an even element index. """ diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/core/lower_llvm.py b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/core/lower_llvm.py index c17453ef6681..23644a20de99 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/core/lower_llvm.py +++ b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/core/lower_llvm.py @@ -581,11 +581,6 @@ def _resolve_llvm_flavor() -> str: "declare i32 @llvm.amdgcn.update.dpp.i32(" "i32, i32, i32 immarg, i32 immarg, i32 immarg, i1 immarg)" ), - # Packed fp16 atomic add (gfx940+). Two fp16 lanes per atomic transaction. - "global.atomic.fadd.v2f16": ( - "declare <2 x half> @llvm.amdgcn.global.atomic.fadd.v2f16.p1(" - "ptr addrspace(1), <2 x half>)" - ), "mbcnt.lo": ("declare i32 @llvm.amdgcn.mbcnt.lo(i32, i32)"), "mbcnt.hi": ("declare i32 @llvm.amdgcn.mbcnt.hi(i32, i32)"), "ds.read.tr16.b64": ( @@ -3574,21 +3569,28 @@ def _op_memref_global_atomic_add_pk_bf16(self, op: Op) -> None: def _op_memref_global_atomic_add_pk_f16(self, op: Op) -> None: """Lower the packed-fp16 atomic add to its AMDGCN intrinsic. - Mirrors ``_op_memref_global_atomic_add_pk_bf16`` for fp16. - GEPs into the fp16 buffer at ``idx`` and calls - ``llvm.amdgcn.global.atomic.fadd.v2f16.p1`` (gfx940+). + Mirrors ``_op_memref_global_atomic_add_pk_bf16`` for fp16: a generic + ``atomicrmw fadd <2 x half>`` (agent scope, monotonic) with the + device-local memory-model metadata, which the backend selects as the + native ``global_atomic_pk_add_f16``. The legacy + ``llvm.amdgcn.global.atomic.fadd.v2f16`` intrinsic it replaces is + lowered on gfx1250 as a fenced atomic: a ``global_wb`` / ``global_inv`` + pair and full counter drains around every single add, which + serialised the split-K epilogues. """ ptr, idx, val = op.operands - self._need("global.atomic.fadd.v2f16") gep = self._fresh("gep") self._current().emit( f" {gep} = getelementptr inbounds half, ptr addrspace(1) " f"{self._operand(ptr)}, i32 {self._operand(idx)}" ) + self._needs_fp_atomic_md = True self._current().emit( - f" {op.result.name} = call <2 x half> " - f"@llvm.amdgcn.global.atomic.fadd.v2f16.p1(" - f"ptr addrspace(1) {gep}, <2 x half> {self._operand(val)})" + f" {op.result.name} = atomicrmw fadd ptr addrspace(1) {gep}, " + f'<2 x half> {self._operand(val)} syncscope("agent") monotonic' + ", !amdgpu.no.fine.grained.memory !1" + ", !amdgpu.no.remote.memory !1" + ", !amdgpu.ignore.denormal.mode !1" ) def _op_memref_global_load_vN(self, op: Op) -> None: diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/tests/golden/rocke_representative_ir_sha256/conv_wgrad.json b/dnn-providers/hip-kernel-provider/rocke/platform/tests/golden/rocke_representative_ir_sha256/conv_wgrad.json index 05af6e5d5761..3a536167ad9f 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/tests/golden/rocke_representative_ir_sha256/conv_wgrad.json +++ b/dnn-providers/hip-kernel-provider/rocke/platform/tests/golden/rocke_representative_ir_sha256/conv_wgrad.json @@ -37,10 +37,10 @@ }, "conv_wgrad/gfx1250/n1h8c16k32r3_spk4_cshuffle": { "arch": "gfx1250", - "bytes": 54947, + "bytes": 55531, "family": "conv_wgrad", "kernel_name": "irhash_wgrad_1250_spk4_cshuffle_N1H8W8C16_K32Y3X3_t32x32x32_w2x1_a16x16x32_mem_cshuffle_spk", - "sha256": "d6709d4327c8ccf1074a9bc864afedd7fef77891fe804886ea1df4160e9e1c07", + "sha256": "d5edf46c7b910114fad2f49bfce7d464f2d04a66a5e3c064c337ad643fad40ea", "status": "ok" }, "conv_wgrad/gfx1250/n1h8c64k64r3_g4": { @@ -101,10 +101,10 @@ }, "conv_wgrad/gfx942/n1h8c64k64r3_g4_spk4": { "arch": "gfx942", - "bytes": 26896, + "bytes": 27148, "family": "conv_wgrad", "kernel_name": "irhash_wgrad_942_g4_spk4_N1H8W8C64_K64Y3X3G4_t64x32x16_w2x2_a16x16x16_mem_cshuffle_spk", - "sha256": "b679eee7b5aea368d4e7996ba4d5a95e3f652baecc95529dfa36a662b1cd9f87", + "sha256": "c4127c11b80f3fa0d4a2bdc7e52c4dc352c790f7ffd5eff6266053ab0ffefdda", "status": "ok" }, "conv_wgrad/gfx950/n1h8c16k32r3": { @@ -149,10 +149,10 @@ }, "conv_wgrad/gfx950/n1h8c16k32r3_spk4": { "arch": "gfx950", - "bytes": 26145, + "bytes": 26397, "family": "conv_wgrad", "kernel_name": "irhash_wgrad_950_spk4_N1H8W8C16_K32Y3X3_t64x32x16_w2x2_a16x16x16_mem_cshuffle_spk", - "sha256": "f88eb29eee495f0fa9ea2e0e1d3e543902a238589e117588cec1b0fbd2ffcf94", + "sha256": "0815ecf5b28c6c8560baa60ea0443ca0eaf3c5c10aea0ec3374cdd56c995574f", "status": "ok" }, "conv_wgrad/gfx950/n1h8c64k64r3_g4": { @@ -181,10 +181,10 @@ }, "conv_wgrad/gfx950/n1h8c64k64r3_g4_spk4": { "arch": "gfx950", - "bytes": 38655, + "bytes": 39239, "family": "conv_wgrad", "kernel_name": "irhash_wgrad_950_g4_spk4_N1H8W8C64_K64Y3X3G4_t64x64x32_w2x2_a32x32x16_mem_cshuffle_spk", - "sha256": "6e19701dff1ca74eefcaa75027a91e42eef9aa5b5c081b00987e5e0901de2a51", + "sha256": "b8e704266dc1d712907553c4928ffcc4d77fcd644dbd0abb0f2b3b47a1955805", "status": "ok" }, "conv_wgrad/gfx950/n2h16c32k32r1": { @@ -234,10 +234,10 @@ }, "conv_wgrad/gfx1250/n1h8c16k32r3_spk4_cshuffle": { "arch": "gfx1250", - "bytes": 55086, + "bytes": 55670, "family": "conv_wgrad", "kernel_name": "irhash_wgrad_1250_spk4_cshuffle_N1H8W8C16_K32Y3X3_t32x32x32_w2x1_a16x16x32_mem_cshuffle_spk", - "sha256": "cce59f1ea889507c62d4abe9c141e5e6a889cb8dbb7c816730ba1362e6dede69", + "sha256": "fbca7fe69c8677120991634e159227a1de11682f37daaaa3b00cbbf442d289a9", "status": "ok" }, "conv_wgrad/gfx1250/n1h8c64k64r3_g4": { @@ -298,10 +298,10 @@ }, "conv_wgrad/gfx942/n1h8c64k64r3_g4_spk4": { "arch": "gfx942", - "bytes": 27035, + "bytes": 27287, "family": "conv_wgrad", "kernel_name": "irhash_wgrad_942_g4_spk4_N1H8W8C64_K64Y3X3G4_t64x32x16_w2x2_a16x16x16_mem_cshuffle_spk", - "sha256": "7411ee05b8fc22b908f5910efb04b368a0471acf8c49a527507c833e536feb1c", + "sha256": "adf01950dc43c298c8c3962c2cbbba9296107d419d8688507956459361fbf780", "status": "ok" }, "conv_wgrad/gfx950/n1h8c16k32r3": { @@ -346,10 +346,10 @@ }, "conv_wgrad/gfx950/n1h8c16k32r3_spk4": { "arch": "gfx950", - "bytes": 26284, + "bytes": 26536, "family": "conv_wgrad", "kernel_name": "irhash_wgrad_950_spk4_N1H8W8C16_K32Y3X3_t64x32x16_w2x2_a16x16x16_mem_cshuffle_spk", - "sha256": "da9cc8fd173e6bee98f0c3f9ea45d1b44283230d165ebedc25466431882a034f", + "sha256": "abbb90f4b8cd8bbdce8a45f3ba49d5555973d95f2e16c31eeb916d8388db7c66", "status": "ok" }, "conv_wgrad/gfx950/n1h8c64k64r3_g4": { @@ -378,10 +378,10 @@ }, "conv_wgrad/gfx950/n1h8c64k64r3_g4_spk4": { "arch": "gfx950", - "bytes": 38794, + "bytes": 39378, "family": "conv_wgrad", "kernel_name": "irhash_wgrad_950_g4_spk4_N1H8W8C64_K64Y3X3G4_t64x64x32_w2x2_a32x32x16_mem_cshuffle_spk", - "sha256": "baaf64e74b63f526c6981224bdcfcfae0b8b3e3dac755eac97977964943945a3", + "sha256": "9960fd1b9d77bbca947b56a519bfedaeee34f49c8c5347518e5803f145d9b192", "status": "ok" }, "conv_wgrad/gfx950/n2h16c32k32r1": { @@ -431,10 +431,10 @@ }, "conv_wgrad/gfx1250/n1h8c16k32r3_spk4_cshuffle": { "arch": "gfx1250", - "bytes": 55150, + "bytes": 55734, "family": "conv_wgrad", "kernel_name": "irhash_wgrad_1250_spk4_cshuffle_N1H8W8C16_K32Y3X3_t32x32x32_w2x1_a16x16x32_mem_cshuffle_spk", - "sha256": "0d4ff4ddc79acf33334ae99b5d2f2da8afd6199f35b0638c3e4d448d45d6a087", + "sha256": "bc450909719641a857aca3d5d35096545aa83e065ac70c43ac16ec77e0781127", "status": "ok" }, "conv_wgrad/gfx1250/n1h8c64k64r3_g4": { @@ -495,10 +495,10 @@ }, "conv_wgrad/gfx942/n1h8c64k64r3_g4_spk4": { "arch": "gfx942", - "bytes": 27099, + "bytes": 27351, "family": "conv_wgrad", "kernel_name": "irhash_wgrad_942_g4_spk4_N1H8W8C64_K64Y3X3G4_t64x32x16_w2x2_a16x16x16_mem_cshuffle_spk", - "sha256": "004a3b6223e994093c87f281d71a3e166de55601d1f4f8854a9514371872ec14", + "sha256": "9e4586909ee7fe822a07a66edc01f23722fe5c443364255f36a89a1e2594de8a", "status": "ok" }, "conv_wgrad/gfx950/n1h8c16k32r3": { @@ -543,10 +543,10 @@ }, "conv_wgrad/gfx950/n1h8c16k32r3_spk4": { "arch": "gfx950", - "bytes": 26348, + "bytes": 26600, "family": "conv_wgrad", "kernel_name": "irhash_wgrad_950_spk4_N1H8W8C16_K32Y3X3_t64x32x16_w2x2_a16x16x16_mem_cshuffle_spk", - "sha256": "c090de9315a4d8a019dc6fabb172961345779dc8cc4b9c8cba1d040d99c5ac2f", + "sha256": "2f582f32c5dd795e2962a6ba75706d65dcba891d72811a9c8b5839df5b6aeed8", "status": "ok" }, "conv_wgrad/gfx950/n1h8c64k64r3_g4": { @@ -575,10 +575,10 @@ }, "conv_wgrad/gfx950/n1h8c64k64r3_g4_spk4": { "arch": "gfx950", - "bytes": 38858, + "bytes": 39442, "family": "conv_wgrad", "kernel_name": "irhash_wgrad_950_g4_spk4_N1H8W8C64_K64Y3X3G4_t64x64x32_w2x2_a32x32x16_mem_cshuffle_spk", - "sha256": "81bbd79a8447cf36bb163a127d12aede645368914b75f2e289ad5af3a9476c41", + "sha256": "6af4e940a085c42c53285a37c9cf4dada3851adb558fc87a98478736fd293241", "status": "ok" }, "conv_wgrad/gfx950/n2h16c32k32r1": { @@ -628,10 +628,10 @@ }, "conv_wgrad/gfx1250/n1h8c16k32r3_spk4_cshuffle": { "arch": "gfx1250", - "bytes": 55150, + "bytes": 55734, "family": "conv_wgrad", "kernel_name": "irhash_wgrad_1250_spk4_cshuffle_N1H8W8C16_K32Y3X3_t32x32x32_w2x1_a16x16x32_mem_cshuffle_spk", - "sha256": "0d4ff4ddc79acf33334ae99b5d2f2da8afd6199f35b0638c3e4d448d45d6a087", + "sha256": "bc450909719641a857aca3d5d35096545aa83e065ac70c43ac16ec77e0781127", "status": "ok" }, "conv_wgrad/gfx1250/n1h8c64k64r3_g4": { @@ -692,10 +692,10 @@ }, "conv_wgrad/gfx942/n1h8c64k64r3_g4_spk4": { "arch": "gfx942", - "bytes": 27099, + "bytes": 27351, "family": "conv_wgrad", "kernel_name": "irhash_wgrad_942_g4_spk4_N1H8W8C64_K64Y3X3G4_t64x32x16_w2x2_a16x16x16_mem_cshuffle_spk", - "sha256": "004a3b6223e994093c87f281d71a3e166de55601d1f4f8854a9514371872ec14", + "sha256": "9e4586909ee7fe822a07a66edc01f23722fe5c443364255f36a89a1e2594de8a", "status": "ok" }, "conv_wgrad/gfx950/n1h8c16k32r3": { @@ -740,10 +740,10 @@ }, "conv_wgrad/gfx950/n1h8c16k32r3_spk4": { "arch": "gfx950", - "bytes": 26348, + "bytes": 26600, "family": "conv_wgrad", "kernel_name": "irhash_wgrad_950_spk4_N1H8W8C16_K32Y3X3_t64x32x16_w2x2_a16x16x16_mem_cshuffle_spk", - "sha256": "c090de9315a4d8a019dc6fabb172961345779dc8cc4b9c8cba1d040d99c5ac2f", + "sha256": "2f582f32c5dd795e2962a6ba75706d65dcba891d72811a9c8b5839df5b6aeed8", "status": "ok" }, "conv_wgrad/gfx950/n1h8c64k64r3_g4": { @@ -772,10 +772,10 @@ }, "conv_wgrad/gfx950/n1h8c64k64r3_g4_spk4": { "arch": "gfx950", - "bytes": 38858, + "bytes": 39442, "family": "conv_wgrad", "kernel_name": "irhash_wgrad_950_g4_spk4_N1H8W8C64_K64Y3X3G4_t64x64x32_w2x2_a32x32x16_mem_cshuffle_spk", - "sha256": "81bbd79a8447cf36bb163a127d12aede645368914b75f2e289ad5af3a9476c41", + "sha256": "6af4e940a085c42c53285a37c9cf4dada3851adb558fc87a98478736fd293241", "status": "ok" }, "conv_wgrad/gfx950/n2h16c32k32r1": { diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/tests/instances/differential/golden/llvm_gfx_all.json b/dnn-providers/hip-kernel-provider/rocke/platform/tests/instances/differential/golden/llvm_gfx_all.json index 9c1372fc5e3a..1f4d19a8903c 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/tests/instances/differential/golden/llvm_gfx_all.json +++ b/dnn-providers/hip-kernel-provider/rocke/platform/tests/instances/differential/golden/llvm_gfx_all.json @@ -2353,7 +2353,7 @@ "13": "294128a0d12bf923feb75c23acc85a7e24bb7e43b9593575b9d7e38273d21c30", "15": "159c38c95cfe15b2e7783df25bee3627fc7a37ae7a3055155726f607a4b87901", "16": "575e820e57d00dcbd194bdf5b8c2879bd56ab5f2d35caec8a249e44a8e08a6c4", - "17": "c96041495be3f29b29f7a265b89cc63e061a1328a800a77fdff52caa38f78704", + "17": "02e3e60b594e023500c6ef2b0a6aa2e06b29eceaac3f6505210e156e7e3cab36", "18": "d82232e225b1e6fdedb51c7edb27e4c102d48a2bc7b2a2470187c68c6ac4812a", "19": "95f4c0fb3ffee944b8420498546cbbba1de9888ceef074a8186204a2d962d1bb", "2": "3d3da00794d8e8bfdf6d57b2b76086ec9224f3f076c294ab36792dbf523a7f8b", @@ -2363,8 +2363,8 @@ "24": "f84e004f1fd0dbcff1429d6f1c1f73dc8dac7643afbb2253c1e43634d6563ae9", "25": "583761e177e08834bbb5804d98291e41f4d06adbd3a82de793fd9fed1920f98b", "26": "618b36ce1f2566f04f385189959132c57a772b36d297992da4fe9ae51dfa7026", - "27": "7be37a207df5996ee604d307fd9df07410121129707f6132095ef0e2f0aab44d", - "28": "289eab1ad2060ccef66ac3ae5b6c0001b986ab4254e05b592fb01293f274fc97", + "27": "2b916770c06998f1439f386d1a6939ce7441f06336ad8d6bbb8646294f52457d", + "28": "04b20878792bf2cd9d9edce2cc0b9e6e991b7fec55618063205425a9ad2ed9fb", "29": "b10dc5a3316cf34eb6c2ba7adbbcf75db63ab773e1b2f3d04352ca345b616046", "30": "d004ce3365d9f655fe5ab1220c9edaf854327ebbc3a37ff342c04d623a548d71", "31": "d13aa5694cff1d8a954f0e31e198c2a0959918a22e277dcfb32f9aee25e71826", @@ -2376,7 +2376,7 @@ "4": "0068fff7f0d32fd8bdc5d6d10e059e3ec4d758bd882fd452b5742d31cd9d3207", "40": "a5a9429b46ade9d0b3ab710cbbda0ade52b56b25721d5c797275378c236ed39a", "41": "0562b1bdfd37f0ed4a978dfe64c5392d9d1e11df6a7b4d26739cccfd79b5ba1c", - "42": "3a31edd453ad5f7e4c9dc671aca3854e6a1a6a984a86aa90e628402e221ab0c1", + "42": "7570423b64dfc4881253cedda5bc430fb5ec77c1c17b87ae6aa6be31715655be", "43": "dc83ac08c3a41ae77d22ec26741f8ef548409d06aa0b5194ed751ffe30dadc90", "44": "4341c60765ad89f2574e433894226781f13b4e910984301125e7ec9e2d3e0cc2", "45": "ce50d815e84adff84a2a248d068f1e5f3e53522b679071f3e524855a118e2da0", @@ -2389,7 +2389,7 @@ "53": "5c2b64c5c13123d967aae65e448d9161b01e3441e46791a38e8f49dd749f50f1", "54": "edd965e19df8883fd5bfb4cff2a17886e41415725371374271cde0e1e746e0ba", "6": "ccb519f64f32f70e0672c09ce1e220c65bb0defe9dbce21e77ebdb07e302fa8b", - "7": "d8bebe2a5f9442f831666b5517ffaa7e436300442d73b58ad534020e4af13f89", + "7": "6a9ce44b92dd56970078680dae7f6b8adaaca18cdbf819a64676d196472f47c9", "8": "3bd26276378cf6242a2e78a23167421fa5d9a84b1df6ce96ed1c7999d5bea97c", "9": "ba3dd4497bb0519a4fc215232fc6788910a5f9b2b2da3959842802a0bca64d53" }, @@ -2406,7 +2406,7 @@ "17": "16bda865ba27b681b78967d718fc7d436c3fe410e49ce5b084e70b2765334dbd", "18": "6eb404062da7a49ff4b8068b350c6451e2b61146b5f68c6becfd12658dd42576", "19": "86688c89323f3f61065b0d4eb05c567e0c2b82c7d560b721124707c8586d61ae", - "20": "a12970b388558f8bdfcc02e355a96ac17cca32c8c98dbdd1fba34355dd6a210a", + "20": "13c387e65f9055a9a87986a154359de8e794380d370d92f21007acc48f8f5419", "21": "ec1f118e2e26b0e3b2ba9843135b0f43579be437ea44e40f8fdbb754dbb3e48d", "22": "1bea7e86014f1e74cee0a91d531569206ce6e1248ee2849e5584843ae59638b7", "23": "150a8c24e93555e726615b1a55e0685efe2c6810a02feecee9ed7f2d355b8c49", @@ -2423,7 +2423,7 @@ "33": "c8534322a81b336183061022fcf5fe0bbd0f156f4a5af0242f695b819fa41d15", "34": "237f8bc95caf1f4f0bfb32d89bbc9b55de5dad0853561a14a7e7ad8bbfae6cc8", "35": "9695fffd8fb578532a1484e6d9879ed94fd868428dc0880dc68e19ed2cf6de9c", - "36": "b14cf3d0b5c51226445f2d72d32d549bed4342daedab9879550a9a72a8f4ff2f", + "36": "dca82f294e424ec83860a30922c4ba1d72a5f99d6633f85b9d45e45b6c300ac1", "37": "10b39fedda5627a9445f95f7f4b6a8772d5bb64239c3cf702dcc05aabf259fc3", "38": "36566a1f072f41a1c80ba5d433d58f26eebaaf3f3063362ca150d3d2128ade5e", "39": "db53587e5e47f83283bec6b26072576b9b8f7525b5beebb02bc53b5efd5fc963", @@ -2441,7 +2441,7 @@ "58": "0fd30468be138821abaf30432c005cc82999425fff3020f0746ff78c199f1a68", "59": "46dd718d4f966e5b0659ebe0f8b2f6fde2ca6eac6035fd247c0b7f44df648408", "6": "b2605886350eb7b793616938e9179ad480084552a3533959f34918085a34484c", - "60": "e1f8d89e18ad5bb1b9e918a2df43fca9283adf375de2488619a56861a2b0cd3b", + "60": "55e677df75468f06ced7b24d600f48aa934f93a5c46963966d7a438383e05c59", "7": "d653088ddce4d716c3b9e20fabe5592541b52ba5e800786a86b660763c25a596", "8": "889d9e3617a8eb2e30d29db4ef8da0aeac52c9fa0d92bc53541947cf85926c08" }, @@ -4128,7 +4128,7 @@ "13": "4718aa9667f38e44cdcb145098c612b58880aaaabfac0b6f12b08087691b2c62", "15": "34603cc15bed77746bf7971dfb56a76c3c50a6baa2a07db3b8762ef464d2634c", "16": "77023b6b8a5ad5a799ba3ad0865d3de71f8e0e04447cd939bde535e4c76ab6c2", - "17": "4ae08237e04006485a919f2f577d4468aa3f99a71cb443f5796b89e3a8385526", + "17": "749c1cbbf5a4177442bb4c529662536e89bfd0968ce98ab6ad7d19a95a74e8bf", "18": "4ae3aebde852418c1d7dbdc866b5a648dfb9fac29a3c0f88263126cf225d5d15", "19": "ee1c2a0c50d1f01e5d0ffb3ef8d0ad2123150bec50839576c581108f92894bc7", "2": "12f31176d7d6634d1dba25c10a883b9e4291532d03e82e1e0308289d54167289", @@ -4138,8 +4138,8 @@ "24": "000b4c2573e376e9d98d0a5d0f917535f72c7fc5636b3db9ac3e6819366cc155", "25": "8173176c27cdfe0857efa8d819d34c8927da938581128b5657cefbfb28085e03", "26": "f693188c58e41d60203416b5a154dad959fa2b028aae78a4b1f3038be98eb905", - "27": "724c60183b73cbb65cc38f8cae8072df2f144368f2e86f7998aa4314f851a971", - "28": "d265b9e8ca38912b39ab903479ed6849f7e62e70ddefcb6c19038178f17067f1", + "27": "be6fd03794ed8b219d779d414b86bb3d39413dbaf342351fac19bfd7c1d8dc3b", + "28": "b750889a3e974b4ef14798752c535dc7098ce5cc31837d8d0217c15e78e37537", "29": "8c5cef30949fe162ad9cf8a76dcf61966ed49f48433ceb2dcda036b25cf26ac2", "30": "8b287fd7599f5a2773f2417ce7ef72ce3fff5d0cfbad129e5b9556cef490b66b", "31": "7d68058922f594ef74b3139a27f13816964e80bae1cefddf826a7199dff19ed4", @@ -4151,7 +4151,7 @@ "4": "1cf24baefccbd02adcb68f7f0b7ac6555dbe62523e8fee31281f54be49e35c5d", "40": "1817445314803697c4c6f4edc653debd454182261a41f7cf32415634ff6acaf6", "41": "63f9692e7476e40dfb3131ffd787dde702a665668ecefbc4d158367b158f5ffd", - "42": "d26b720383fbfa0e122d07c43cefe2c16ed5e8a01287da6e23c3d112d38950ef", + "42": "16044150a20db2467960685e56c3cce23a14ede8da641f506c1b1245d7f8f9bf", "43": "750d3deaedba7cfebfa1888f52c5f4a591eabb2ea797bd01223e99f2ba47c276", "44": "27c8b89de4f615c9977212ef1b4d8cdd7a68868a018a7979b2bd80024397e0e8", "45": "48560a51a8f519bbbbde3d47f1fa9e16a0db99e2d0e619e1072908e1111a0284", @@ -4164,7 +4164,7 @@ "53": "bfc732240858b6e80e45a64705703f9dc4d254c99657d32a932534c5fbcc4829", "54": "b14d621c3ddbfa30adfe2f19eb4fa6ac3dc57bbf17325ef673429547085072c5", "6": "6be1307c9c3fae4cade18018eb5c4844de596208590644b0d4dece1b2179f84a", - "7": "bcf2572bd61c9e2c1592adefd2bebbe06bb49e6f1556ce16a4d3cce00bbbd84e", + "7": "108a8290848c924fd31bf6c29e120d3a6c21d6f4e4d131c2c2502bbd5a67055a", "8": "d001f59ab48d690329d0b0365fd341b98d0ac826c13048d259a1d8eecf0cbc33", "9": "56615523785949e79185f7c6ba83fe9048ccc6fb6994293142de799ee348f9d3" }, @@ -4181,7 +4181,7 @@ "17": "9fef612e03e351580dd65a37fec9c2fe199b80f897d7f4ab513c5a53b35f64f8", "18": "845529ad02c3c581dadd872782d8f37708d3fa1d41163d6c7497c164c91214ff", "19": "f8341e9741d24b7c186e0c7c4daf5a68cc7110ed82148dfa767499128c1d0a35", - "20": "1e764006b3487b6dbcc05c1506c9eb669775809ec1c8ad1f7238a8e64d283c7f", + "20": "df65d4e4f8627688a660fc7d55ca9fc59ae669eed36df01d1baccaae9001123a", "21": "4e7944845415149a7ac9f665a0acdd20861ea4735f86a405357261473eaea387", "22": "06202a7998a26beeec528c8da532c69ebcf955510d5cbf93044970044b09bca4", "23": "4fd0a9d596283bdce3a9a55e061427eb4e9b993cb088b18bbec1235d24974007", @@ -4198,7 +4198,7 @@ "33": "5a7182ddbed0ba909cdad6da62c6109c2f2e9e4d84314c8c4ed57716a6a45e58", "34": "d4eb7804d910b122ece8e21663b5fec90a2faa5ba2ea656f390914513312ff44", "35": "cfbb77a3cec29b148a54cf677028196953efbba5893108988f6d5f5e18be86d7", - "36": "58092d04a0838e0032fcb1b8bd3016dde6a2576c64be5f8ddb57e2343cfb7839", + "36": "122b28d8b61f385ca67f2a64d9998a6c56875315d739ef20fb346abad90b1cc5", "37": "cd8da3081f3f8bad5f95e8124f9f9ad15f1f80921980eeab4e0e9683fae24571", "38": "2c888e4af5de99cb61ef77c2e6a9fdc3e8f21629c5cb79ec8aa176ea6f547419", "39": "66954dc0bd33f6dbbe0239dc8b5cdb37660f714d181f5ce2ec353106bd520fdd", @@ -4216,7 +4216,7 @@ "58": "95aabfecf90f2021a781953eed801d3b14d86c4c5844c365dc6edd83416a0512", "59": "654c0bfc412129aad45be434cb4e382fab2739ad94ffe5d607db6924e5f77b00", "6": "f3b4e8cf7ab611637137babc7994172995f76262274b0a63ca3b6999ac62ef04", - "60": "b3e910d01f76ec19f85babd5c26ed1debd157233c2560492711967b843de603a", + "60": "cafe652716f3f73fc0747ca5a6932007dcc21baa840e1367de312a93bbfdf4b1", "7": "e12021bd6f8923472a7595394ecde03a55dcd07d63d714b9adaea3efa79df887", "8": "2cd90bffada0eaec9a90caf5d5b76f2f213e4e884e3f413b13e79eef3071ec82" }, @@ -5933,7 +5933,7 @@ "13": "c4ef396ae940c4fdb26503f54756e85e60149d01d7cdfb73692da89bb2aca6ad", "15": "ffca81c20a6601e73bea48cf39a5279c1924e43928f81b159a40535e9f375503", "16": "0554832490d8c52f9c0cbd277c2a9292e4719696ec240210c6a384e89fd7410e", - "17": "0429b42176dc5784b8d3070103c7ae8231b9c8a35b478f708151c9f11ac3d0a0", + "17": "1ccbdf3e6c060ef2bc6fdb7930fdc9c266a3ba7e8acec35fe5ecca79c791cd66", "18": "1e64d9bb5b95352ab7bcc07eb5bf665ce521a37610ba3aa84f55f8afc825eda0", "19": "bd49133df28132892d75337112346dc851ca69089844dba935934439fadee185", "2": "c7aac42a9187c8b0c236460edc66f3cc7c6a5fe6bdbf51db302015864a9f4b75", @@ -5943,8 +5943,8 @@ "24": "cfc0bcb59e6187f87781711767ef61ef814f820a1926e6d6c88c6fb93a9b709d", "25": "c2f75a78ed8683372883b5c444455498cf84620fa6f284383ab3ea6d3167ac54", "26": "c3497f863152e01d88c656d446be40c94e3fcef90fed8fc24f2f5aec149aa61d", - "27": "47dfbe66c83d9de75df9d4252b30efbff432443b2ddc24387a488e31820b7db3", - "28": "bbdac24de1eccfd19cc06c7e966d44b3762524c63cdedbbdacbe9e0bb54895fa", + "27": "302a6e03ef807f27575c0f11619c03bb3a0be4588bb01f21bb083dfa73620d48", + "28": "970c3baf37817f9807d87c2c700bfe82615352f7f6e0885d9de57f03cc4f6601", "29": "9dcd9d95074aa655f2060d535ca7948eb7b05ae98998cab3be52fc407ca82048", "30": "0448cb69991b6e8ce01917cb83e89f02f43ab67e2b8d4dd7b63bb2e3fb4f1d5e", "31": "d2feb19ecc3a093beb4301905153497d98fa3bb1d2fb94920a9fbf2c92368112", @@ -5956,7 +5956,7 @@ "4": "11323ce1e5d8050df2b054d2a30986a696b5a67b17634407f7c053dfd078e076", "40": "9cb97aa3b43bdcecde831a14e325479dbf65fa53aa0b10924356bced359f0359", "41": "941f92518325f4eb3ac4bc871e812d1a5094d859a49e3ad2c4330525e686e3fe", - "42": "176cfb3b8717f8cbd727fdbe972ba927b7131d9126eaf1d34519ade3d9210140", + "42": "1eda5471fc3b66ad0bcfb88a39637c2df045da7a603fac08f283e37c3fc8d899", "43": "a3bb30ed428af721f349ad5d11a629feb5707632e6cd4cfb3d137f41dc460980", "44": "f95d1dff265e1f3ef10d622f69c471ba5d61a31daf0f6def2be2c854fd4da9af", "45": "9565bdfd21055a60a6fb8138986e2897f67d059ec941f02a8fba321527be281f", @@ -5969,7 +5969,7 @@ "53": "9cf2dae14c17530ab736153bfa3fdd5161dab0d011ca32faec7e70684a0b750c", "54": "ddc492f0f87e02c75b36863002554901a9ceb8867d6a32dbeaf1485fee81cb82", "6": "92ed5cc85215747725b6fdd1d458acaec92353a86be1977cc86c118eae077624", - "7": "e6d2b91d83452726233d67461d06705fdfa7bdab4563645d446270c68e4e097c", + "7": "a5fbbaab259b7669d8a1764d9de4caa611c93de6ab49ad91876e428cc98669e3", "8": "578f1d62afa50c173aae1f7af3839968a0adb864cbcf25e8143cb6712199d174", "9": "360081d8b3cf6d0a467ecbcbb36c57ea3bcdc255e9059712a5c8926374ab1306" }, @@ -5986,7 +5986,7 @@ "17": "208a4bfdd1f0d274697c79cfa22446d5f7cc234bfb4d9a5dfd597a669849b12f", "18": "73e41b9ba7b27bd6b80a834ea4d21a8715224092d8a5feec0a258089db83302a", "19": "f6923fb55bfb3e624b01f2de98b5987f5e2a82e3b79a13b29de55352fb67c9bc", - "20": "46075ade7e2a86ed1b3f16f205f0abf706eb4edfbe3616ce73889e9c9ad1909f", + "20": "7cdf478cedbf9ce5ac3b4565e06b28c104e7693af9afc16819cfdfcfbe3a167f", "21": "e9bae61d9bcc34a9fe5ef21223b3ff48c3f100d06f63c55745b8c949942b3489", "22": "c5e3032f5d6f031e897844bfc7dea0b8bc06431712e3e3a6e934f66b05789d91", "23": "d38a227c730eaa23db8dd4dd9bbfe64b368cb6818fd9d798a768f802d0c88337", @@ -6003,7 +6003,7 @@ "33": "487b4879dddfbe2e0a2288f5586a8449ceb3ee31645c51de0d9e93ffeb21a8b2", "34": "c6369561f0b3d8a45b4e791666835f0b97befee7a34a3687d61b22c386ffe1b0", "35": "0eb2bc1ec01ece751ddd8eef7c26dcf50647eadc7d50d5b8bb2dc0e1184d8a8f", - "36": "f8c3360f71f9da12da26329efa8f2b5a29f0500828d131c85b3c0f82d87a7eea", + "36": "662e13ba92a2f432d358e40ecdb0165fc779fe42e1b69060519ca2136de58ec5", "37": "b6861c4ebff53913f2dd613365cbf47d32b34829bfe1f92f43fb360d190902fa", "38": "ee1b5b1d4ea68e6ea2aa159c1f2426350c0022e8f6c6a90d9f3e6bdb96abeb03", "39": "d9474047aa9ca5912c5223045dd62ac9a8f17d51c621935e1a7b0556abcec1e8", @@ -6021,7 +6021,7 @@ "58": "02b4d2adce33a29c4a426bc3c4617387e340a2e19b429e986468c6064dac2759", "59": "fbcf10500498604862078885c4482d5cbe67269770db5929f39fa0d6919d3c89", "6": "2726776342b4ffd858fce5951d771a50515d7c143e07656eb6235e10fd08bd32", - "60": "916897d1f0ed7b3a09a499a4b6f5128ee271fc7bf4e1acbfa8fb2980b00d9ac9", + "60": "33ef9a12ea16c82a1e985e2ad05b557de385f4def104b8b4abd694e6989ff42e", "7": "44a58b32e3d95d8492b23193d701925de3dd9de8cddf8a2c006b6fdf5de87cf8", "8": "af0df509c9e1bfa9bcb32621a677b633c9bd7c0b0a22012ee20acd2894b44df8" }, @@ -7990,7 +7990,7 @@ "13": "c4ef396ae940c4fdb26503f54756e85e60149d01d7cdfb73692da89bb2aca6ad", "15": "ffca81c20a6601e73bea48cf39a5279c1924e43928f81b159a40535e9f375503", "16": "0554832490d8c52f9c0cbd277c2a9292e4719696ec240210c6a384e89fd7410e", - "17": "0429b42176dc5784b8d3070103c7ae8231b9c8a35b478f708151c9f11ac3d0a0", + "17": "1ccbdf3e6c060ef2bc6fdb7930fdc9c266a3ba7e8acec35fe5ecca79c791cd66", "18": "1e64d9bb5b95352ab7bcc07eb5bf665ce521a37610ba3aa84f55f8afc825eda0", "19": "bd49133df28132892d75337112346dc851ca69089844dba935934439fadee185", "2": "c7aac42a9187c8b0c236460edc66f3cc7c6a5fe6bdbf51db302015864a9f4b75", @@ -8000,8 +8000,8 @@ "24": "cfc0bcb59e6187f87781711767ef61ef814f820a1926e6d6c88c6fb93a9b709d", "25": "c2f75a78ed8683372883b5c444455498cf84620fa6f284383ab3ea6d3167ac54", "26": "c3497f863152e01d88c656d446be40c94e3fcef90fed8fc24f2f5aec149aa61d", - "27": "47dfbe66c83d9de75df9d4252b30efbff432443b2ddc24387a488e31820b7db3", - "28": "bbdac24de1eccfd19cc06c7e966d44b3762524c63cdedbbdacbe9e0bb54895fa", + "27": "302a6e03ef807f27575c0f11619c03bb3a0be4588bb01f21bb083dfa73620d48", + "28": "970c3baf37817f9807d87c2c700bfe82615352f7f6e0885d9de57f03cc4f6601", "29": "9dcd9d95074aa655f2060d535ca7948eb7b05ae98998cab3be52fc407ca82048", "30": "0448cb69991b6e8ce01917cb83e89f02f43ab67e2b8d4dd7b63bb2e3fb4f1d5e", "31": "d2feb19ecc3a093beb4301905153497d98fa3bb1d2fb94920a9fbf2c92368112", @@ -8013,7 +8013,7 @@ "4": "11323ce1e5d8050df2b054d2a30986a696b5a67b17634407f7c053dfd078e076", "40": "9cb97aa3b43bdcecde831a14e325479dbf65fa53aa0b10924356bced359f0359", "41": "941f92518325f4eb3ac4bc871e812d1a5094d859a49e3ad2c4330525e686e3fe", - "42": "176cfb3b8717f8cbd727fdbe972ba927b7131d9126eaf1d34519ade3d9210140", + "42": "1eda5471fc3b66ad0bcfb88a39637c2df045da7a603fac08f283e37c3fc8d899", "43": "a3bb30ed428af721f349ad5d11a629feb5707632e6cd4cfb3d137f41dc460980", "44": "f95d1dff265e1f3ef10d622f69c471ba5d61a31daf0f6def2be2c854fd4da9af", "45": "9565bdfd21055a60a6fb8138986e2897f67d059ec941f02a8fba321527be281f", @@ -8026,7 +8026,7 @@ "53": "9cf2dae14c17530ab736153bfa3fdd5161dab0d011ca32faec7e70684a0b750c", "54": "ddc492f0f87e02c75b36863002554901a9ceb8867d6a32dbeaf1485fee81cb82", "6": "92ed5cc85215747725b6fdd1d458acaec92353a86be1977cc86c118eae077624", - "7": "e6d2b91d83452726233d67461d06705fdfa7bdab4563645d446270c68e4e097c", + "7": "a5fbbaab259b7669d8a1764d9de4caa611c93de6ab49ad91876e428cc98669e3", "8": "578f1d62afa50c173aae1f7af3839968a0adb864cbcf25e8143cb6712199d174", "9": "360081d8b3cf6d0a467ecbcbb36c57ea3bcdc255e9059712a5c8926374ab1306" }, @@ -8043,7 +8043,7 @@ "17": "208a4bfdd1f0d274697c79cfa22446d5f7cc234bfb4d9a5dfd597a669849b12f", "18": "73e41b9ba7b27bd6b80a834ea4d21a8715224092d8a5feec0a258089db83302a", "19": "f6923fb55bfb3e624b01f2de98b5987f5e2a82e3b79a13b29de55352fb67c9bc", - "20": "46075ade7e2a86ed1b3f16f205f0abf706eb4edfbe3616ce73889e9c9ad1909f", + "20": "7cdf478cedbf9ce5ac3b4565e06b28c104e7693af9afc16819cfdfcfbe3a167f", "21": "e9bae61d9bcc34a9fe5ef21223b3ff48c3f100d06f63c55745b8c949942b3489", "22": "c5e3032f5d6f031e897844bfc7dea0b8bc06431712e3e3a6e934f66b05789d91", "23": "d38a227c730eaa23db8dd4dd9bbfe64b368cb6818fd9d798a768f802d0c88337", @@ -8060,7 +8060,7 @@ "33": "487b4879dddfbe2e0a2288f5586a8449ceb3ee31645c51de0d9e93ffeb21a8b2", "34": "c6369561f0b3d8a45b4e791666835f0b97befee7a34a3687d61b22c386ffe1b0", "35": "0eb2bc1ec01ece751ddd8eef7c26dcf50647eadc7d50d5b8bb2dc0e1184d8a8f", - "36": "f8c3360f71f9da12da26329efa8f2b5a29f0500828d131c85b3c0f82d87a7eea", + "36": "662e13ba92a2f432d358e40ecdb0165fc779fe42e1b69060519ca2136de58ec5", "37": "b6861c4ebff53913f2dd613365cbf47d32b34829bfe1f92f43fb360d190902fa", "38": "ee1b5b1d4ea68e6ea2aa159c1f2426350c0022e8f6c6a90d9f3e6bdb96abeb03", "39": "d9474047aa9ca5912c5223045dd62ac9a8f17d51c621935e1a7b0556abcec1e8", @@ -8074,14 +8074,14 @@ "49": "b1f8b4a6d3beaeb39c5be3594188f313cfec564e6f3893eb29a9271b8e945bb8", "5": "0826f9183484b30541afcda4322fac4e5fec28374efb8995a373b9d227a97016", "50": "b54d4d875110a0dbd2b0c42b5844114122cd9f74b075e03ada23422cc62c185f", - "51": "84826a6d0c1344658e348d3010e06274415a4badc09580b3bcf9c4d5ef102e73", + "51": "b367d89061b3bfb68f209d667d122cb3aa406996af41f508c91900415c24b049", "53": "bc38f32697b74c7208cd101293556b123c54cfe96b4469fc0562c23052b10fcb", "55": "0c59e7ef5797f03be197533617fcb598f79e7df6759ca8fb3bf875a53e395731", "56": "acf39c44252e4063825560e706c92070ec2532daca3f1072453e4b9bb747681a", "58": "02b4d2adce33a29c4a426bc3c4617387e340a2e19b429e986468c6064dac2759", "59": "fbcf10500498604862078885c4482d5cbe67269770db5929f39fa0d6919d3c89", "6": "2726776342b4ffd858fce5951d771a50515d7c143e07656eb6235e10fd08bd32", - "60": "916897d1f0ed7b3a09a499a4b6f5128ee271fc7bf4e1acbfa8fb2980b00d9ac9", + "60": "33ef9a12ea16c82a1e985e2ad05b557de385f4def104b8b4abd694e6989ff42e", "7": "44a58b32e3d95d8492b23193d701925de3dd9de8cddf8a2c006b6fdf5de87cf8", "8": "af0df509c9e1bfa9bcb32621a677b633c9bd7c0b0a22012ee20acd2894b44df8" }, From 9c875777206d1e7e5313a22eae646be5d1d1b9e8 Mon Sep 17 00:00:00 2001 From: Jakub Piasecki Date: Sun, 1 Nov 2026 18:00:47 +0000 Subject: [PATCH 03/13] feat(rocke): split fp16 wgrad automatically on gfx1250 again Give gfx1250 fp16 dW its own wgrad split policy: fill to 1024 CTAs (2048 grouped) with K slices down to 512, up to degree 128. bf16 keeps its accuracy-bound policy and fp32 the general one. This reverses the "fp16 cap 1" of the gfx1250 split-K policy. That cap was a symptom of the fenced packed fp16 atomic: with the atomic lowered as a monotonic atomicrmw (previous commit), deeper fp16 splits pay off. The policy was fitted on a degree sweep of the cached split-K binaries, and no fp16 degree failed the 5e-2 verification there. Python only: the C++ engine has no split-policy counterpart (the degree is a kernarg resolved by the host). Update the policy tests and the gfx1250 GPU split-K resolver test, whose fp16 cases now run at the resolved degree. Co-Authored-By: Claude Opus 5.5 --- .../common/README_conv_implicit_gemm_wgrad.md | 14 ++++----- .../rocke/library/tests/test_conv_split_k.py | 4 --- .../tests/test_conv_wgrad_correctness.py | 10 +++--- .../platform/python/rocke/helpers/split_k.py | 31 +++++++++++++++---- .../python/rocke/instances/SUPPORT_MATRIX.md | 2 +- .../dispatch_tests/gemm/test_split_k_wgrad.py | 24 ++++++++++++-- 6 files changed, 59 insertions(+), 26 deletions(-) diff --git a/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/README_conv_implicit_gemm_wgrad.md b/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/README_conv_implicit_gemm_wgrad.md index 58ab29545dec..7a5598bb835a 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/README_conv_implicit_gemm_wgrad.md +++ b/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/README_conv_implicit_gemm_wgrad.md @@ -102,13 +102,13 @@ dW rows, and the extra memory transactions outweigh the halved atomic count. 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 128. 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; no fp16 degree failed the 5e-2 + verification there. - 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 diff --git a/dnn-providers/hip-kernel-provider/rocke/library/tests/test_conv_split_k.py b/dnn-providers/hip-kernel-provider/rocke/library/tests/test_conv_split_k.py index 03a5d35b6c38..813086dd7475 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/tests/test_conv_split_k.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/tests/test_conv_split_k.py @@ -628,7 +628,6 @@ def test_explicit_atomic_that_cannot_be_built_raises(): def test_wmma_auto_without_a_reduction_runs_unsplit(dtype): # gfx1250 16-bit dW with an odd row (wg_N = 27): no packed pair, and WMMA # has no workspace. An auto request runs unsplit; a forced degree raises. - # fp16 is never split automatically, so only bf16 asks for a split here. p = ConvProblem(N=64, Hi=64, Wi=64, C=3, K=16, Y=3, X=3, pH=1, pW=1) kw = { "arch": "gfx1250", @@ -646,9 +645,6 @@ def test_wmma_auto_without_a_reduction_runs_unsplit(dtype): resolve_conv_split_k("wgrad", p, split_k=8, **kw) with pytest.raises(ValueError, match="packed atomic"): resolve_conv_split_k("wgrad", p, override=8, **kw) - if dtype == "fp16": - assert r.requested == 1 - return assert r.requested > 1 # the policy did ask for a split with pytest.raises(ValueError, match="packed atomic"): resolve_conv_split_k("wgrad", p, allow_unsplit=False, **kw) diff --git a/dnn-providers/hip-kernel-provider/rocke/library/tests/test_conv_wgrad_correctness.py b/dnn-providers/hip-kernel-provider/rocke/library/tests/test_conv_wgrad_correctness.py index 8211ed6e3a66..fabcb9c509b2 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/tests/test_conv_wgrad_correctness.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/tests/test_conv_wgrad_correctness.py @@ -834,9 +834,9 @@ def test_wide_tile_spec_at_degree_cap_gfx1250(self): These are the 2x2-warp tiles, and the bf16 problems are sized so the degree lands on its cap. The last bf16 problem (small channels, deep reduction) asks for 128 before the cap; all-positive partials miss the - tolerance there. fp16 is not split automatically (its cap is 1), so its - cases check that the resolver launches no split-K binary, then run the - binary at an explicit degree, as a caller that opts in would. The inputs are all-positive -- same-sign partials + tolerance there. fp16 takes its own policy (fill to 1024 CTAs, 2048 + grouped, slices down to 512, up to 128); its cases run at the degree the + resolver picks. The inputs are all-positive -- same-sign partials do not cancel their 16-bit rounding -- so this also bounds the accuracy at the cap. """ @@ -861,9 +861,9 @@ def test_wide_tile_spec_at_degree_cap_gfx1250(self): # expected degree); -1 is the resolver's pick. cases = [ (32, 32, 64, 96, 1, "bf16", (32, 64, 64), -1, 12), - (32, 32, 64, 128, 1, "fp16", (64, 64, 64), 8, 8), + (32, 32, 64, 128, 1, "fp16", (64, 64, 64), -1, 56), (32, 32, 64, 64, 4, "bf16", (32, 32, 32), -1, 12), - (32, 32, 64, 64, 4, "fp16", (32, 32, 32), 8, 8), + (32, 32, 64, 64, 4, "fp16", (32, 32, 32), -1, 64), (64, 64, 16, 16, 1, "bf16", (32, 64, 64), -1, 12), ] ran = 0 diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/split_k.py b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/split_k.py index 9eda8cdc95f7..3ff4c8f814ee 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/split_k.py +++ b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/split_k.py @@ -133,17 +133,32 @@ def atomic_rounding_degree_cap(fraction_bits: int, tol: float = WGRAD_TOLERANCE) # inputs then exceed that tolerance (at 16, one 1.0 partial followed by fifteen # 2**-8 ones rounds back to 1.0, a 5.5% error). # -# fp16 is not split automatically (cap 1). Its packed atomic does not scale -# like bf16's here: the best degree ranges from 16 on a tiny dW down to no -# split at all, where every degree is slower than split_k=1 on a mid-size -# dW, and no single cap avoids that loss. An explicit split_k still uses it. +# fp16 has its own policy (_WGRAD_SPLIT_POLICY_BY_DTYPE below). _WGRAD_SPLIT_POLICY: dict[str, WgradSplitPolicy] = { "gfx1250": WgradSplitPolicy( target_ctas=2048, grouped_target_ctas=4096, min_slice_k=2048, max_degree=128, - max_degree_by_dtype=(("bf16", atomic_rounding_degree_cap(7)), ("fp16", 1)), + max_degree_by_dtype=(("bf16", atomic_rounding_degree_cap(7)),), + ), +} + +# Per-dtype overrides of _WGRAD_SPLIT_POLICY, keyed by (arch, dW dtype). +# +# gfx1250 fp16: fill to 1024 CTAs (2048 grouped) with K slices as shallow as +# 512, up to degree 128. Fitted on a sample of wgrad shapes with every cached +# split-K binary timed at every degree 2..128 and verified, once the packed +# fp16 atomic was no longer fenced (it is an atomicrmw now). fp16 was not split +# automatically before (cap 1): that cap was a symptom of the fenced atomic. +# No fp16 degree failed the 5e-2 verification on those shapes, so the packed +# fp16 atomic's rounding is not the binding limit there. +_WGRAD_SPLIT_POLICY_BY_DTYPE: dict[tuple[str, str], WgradSplitPolicy] = { + ("gfx1250", "fp16"): WgradSplitPolicy( + target_ctas=1024, + grouped_target_ctas=2048, + min_slice_k=512, + max_degree=128, ), } @@ -437,7 +452,11 @@ def select_split_k_wgrad( if base_grid <= 0: return SplitKDecision(1, base_grid, max_capacity, "empty grid") - policy = _WGRAD_SPLIT_POLICY.get(arch) if dtype_d is not None else None + policy = ( + _WGRAD_SPLIT_POLICY_BY_DTYPE.get((arch, dtype_d), _WGRAD_SPLIT_POLICY.get(arch)) + if dtype_d is not None + else None + ) if policy is not None: if isinstance(groups, bool) or not isinstance(groups, int) or groups < 1: raise ValueError(f"groups must be a positive integer (got {groups!r})") diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/instances/SUPPORT_MATRIX.md b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/instances/SUPPORT_MATRIX.md index 1409ecefd95d..951f3b89d730 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/instances/SUPPORT_MATRIX.md +++ b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/instances/SUPPORT_MATRIX.md @@ -109,7 +109,7 @@ be smaller than 2 GiB. | `direct_depthwise_dgrad` (`DirectDepthwiseDgradStreamSpec`) | ✅ | ✅ | ❌ | `cpg=kpg=1`, any stride and padding; row-streaming scalar FMA at the target's wave size (wave64, or wave32 on gfx1250); `cpl` 2/4/8 channels per lane with packed `<2 x f32>` FMAs (`groups % cpl == 0`, `cpl*KH*KW <= 200`, `cpl*block_w*KH <= 192`). In the AOT direct cache; not yet wired into dispatch | | `direct_depthwise_tiled_dgrad` (`DirectDepthwiseTiledSpec(dgrad=True)`) | ✅ | ✅ | ❌ | `cpg=kpg=1`, stride 1, the forward tiled kernel with the filter rotated by 180 degrees; filters up to 64 on a side incl. non-square (51x5 / 5x51), per-axis padding (`pad_w`) with `PAD <= KH-1`, `PAD_W <= KW-1`. Served by the direct-conv AOT cache, not the dispatcher | | `conv_implicit_gemm_dgrad` stride-1 knobs | ✅ | ✅ | ❌ | folded sub-GEMM record (stride capability), `tap_outer_k` (wave64 MFMA, `kpg % tile_k == 0`), `dy_halo` / `dy_halo_setprio` / `dy_halo_kouter_pad` / `halo_max_wo` (on the tap-outer path, `Ho == Hi`, `Wo == Wi`; the pad needs `lds_k_outer`, gfx950); `xcd_tile_order` remaps only on gfx950. AOT grid: stride-1 tap-outer binaries next to the tilde ones, halo sub-grid on gfx950; not yet wired into dispatch | -| `conv_implicit_gemm_wgrad` | ✅ | ✅ | ✅ | gfx1151/gfx1201/gfx1250: WMMA, `mem` with the `default` or `cshuffle` epilogue (gfx1250 adds the `async_dma` K loop over `global_load_async_to_lds`); split-K (WMMA packed atomic, no two-stage) on gfx1250 only, automatic for bf16/fp32 and explicit-only for fp16, `split_k=1` on gfx1151/gfx1201 | +| `conv_implicit_gemm_wgrad` | ✅ | ✅ | ✅ | gfx1151/gfx1201/gfx1250: WMMA, `mem` with the `default` or `cshuffle` epilogue (gfx1250 adds the `async_dma` K loop over `global_load_async_to_lds`); split-K (WMMA packed atomic, no two-stage) on gfx1250 only, automatic for bf16/fp16/fp32, `split_k=1` on gfx1151/gfx1201 | | `conv_implicit_gemm_wgrad` stream-K | ✅ | ✅ | ❌ | CK iteration-balanced partition; `linear`/`tree` deterministic fixups need the cross-workgroup flag ops | | `conv_implicit_gemm` split-K forward (`split_k > 1`) | ❌ | ❌ | ❌ | gfx1250 only: WMMA, `mem` or `wavelet` pipeline, 2-D, not pointwise, identity `acc_epilogue`; no `async_dma`, `unroll_k`, `chiplet_swizzle` or `group_merge`. Each slice stores an fp32 partial tile into its slab of a `[ks_count][M][K]` workspace and `conv_fwd_split_k_reduce` folds the slabs into D; the degree is the `ks_count` kernarg, and `split_k = -1` resolves it (`resolve_conv_split_k`; 1 where splitting does not pay). Built into the AOT kernel cache; the dispatcher does not select it | | `conv_implicit_gemm` cluster split-K (`split_k_cluster`) | ❌ | ❌ | ❌ | gfx1250 only (llvm23, workgroup-cluster launch): `mem` pipeline, `split_k_cluster` in 2..15, `tile_n % 4 == 0`, K and kpg multiples of 4 (checked at build and on every launched problem). The slices of one output tile form a (1, 1, S) workgroup cluster and fold through a cluster barrier, so there is no Stage 2 launch; one binary per S. Built into the AOT kernel cache; the dispatcher does not select it | diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/tests/dispatch/dispatch_tests/gemm/test_split_k_wgrad.py b/dnn-providers/hip-kernel-provider/rocke/platform/tests/dispatch/dispatch_tests/gemm/test_split_k_wgrad.py index 17868fcbae28..472ecac292a3 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/tests/dispatch/dispatch_tests/gemm/test_split_k_wgrad.py +++ b/dnn-providers/hip-kernel-provider/rocke/platform/tests/dispatch/dispatch_tests/gemm/test_split_k_wgrad.py @@ -234,13 +234,31 @@ def test_slice_depth_term_binds(self): def test_dtype_caps(self): # A 4-tile grid on a long reduction asks for far more than any cap: - # each 16-bit dtype takes its own cap, and fp32 the general one. fp16's - # cap of 1 keeps it unsplit even here, where the demand is highest. + # bf16 takes its accuracy cap, and fp16 and fp32 the general one. args = (24, 9 * 24, 42 * 240 * 320) self.assertEqual(self._pick(*args, dtype_d="bf16").split_k, 12) - self.assertEqual(self._pick(*args, dtype_d="fp16").split_k, 1) + self.assertEqual(self._pick(*args, dtype_d="fp16").split_k, 128) self.assertEqual(self._pick(*args, dtype_d="fp32").split_k, 128) + def test_fp16_policy(self): + # fp16 has its own gfx1250 policy: fill to 1024 CTAs (2048 grouped), + # K slices down to 512, up to 128. + # Fill: base grid 2 x 18 = 36 tiles -> floor(1024 / 36) = 28 (fp32 + # fills to 2048 -> 56). + d = self._pick(128, 9 * 128, 42 * 120 * 160, dtype_d="fp16") + self.assertEqual(d.split_k, 28, d.reason) + # Slice depth: a 12288-deep reduction allows 12288 / 512 = 24 slices + # (fp32's 2048-element floor allows only 6). + d = self._pick(64, 9 * 64, 12288, dtype_d="fp16") + self.assertEqual(d.split_k, 24, d.reason) + self.assertEqual(self._pick(64, 9 * 64, 12288, dtype_d="fp32").split_k, 6) + # Grouped: 4 groups of a 1 x 3 tile grid -> floor(2048 / 12) = 170, + # capped at 128. + d = self._pick(16, 9 * 16, 42 * 240 * 320, tile=64, groups=4, dtype_d="fp16") + self.assertEqual(d.split_k, 128, d.reason) + d = self._pick(64, 9 * 64, 42 * 240 * 320, tile=64, groups=8, dtype_d="fp16") + self.assertEqual(d.split_k, 28, d.reason) + def test_bf16_degree_stays_within_tolerance(self): # N=64, 64x64, C=K=16, 3x3: wg_M=16, wg_N=144, wg_K=262144. Fill and # slice depth both ask for 128, where all-positive bf16 partials exceed From 2f7afdcb515c74fa29a498a32af0d5e2e0128c60 Mon Sep 17 00:00:00 2001 From: Jakub Piasecki Date: Sun, 1 Nov 2026 18:01:18 +0000 Subject: [PATCH 04/13] fix(rocke): rank async-DMA wgrad cache entries by their recorded chunk widths An async-DMA binary records no A / B load width (the copy picks its chunk), so the widest-vector filter of the from-cache sweep added None and aborted. Rank such an entry by the chunk widths recorded in its identity (async_chunk_a / async_chunk_b), as a synchronous kernel with the same load widths; fall back to 8 only when none is recorded. Ranking every async binary as 8 wide let it outrank, and so filter out, the widest synchronous kernels of a problem whose channel run caps the loads below 8. Benchmark / run side only; no kernel IR change. Co-Authored-By: Claude Opus 5.5 --- .../common/benchmark_implicit_gemm_conv.py | 19 +++++++++++++------ 1 file changed, 13 insertions(+), 6 deletions(-) diff --git a/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/common/benchmark_implicit_gemm_conv.py b/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/common/benchmark_implicit_gemm_conv.py index 3f07bc91471f..7a1bb5b28aa9 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/common/benchmark_implicit_gemm_conv.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/common/benchmark_implicit_gemm_conv.py @@ -5542,12 +5542,19 @@ def kind(ident): ) 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, as a synchronous kernel + # with the same load widths. (Ranking it as 8 wide let it outrank, and + # so filter out, the widest synchronous kernels of a problem whose + # channel run caps the loads below 8.) + 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 From 537bf561069bd08b3db3c622322c6eb72e3be634 Mon Sep 17 00:00:00 2001 From: Jakub Piasecki Date: Sun, 1 Nov 2026 18:06:42 +0000 Subject: [PATCH 05/13] feat(rocke): two-stage (f32 workspace) split-K wgrad on gfx1250 WMMA The two-stage scratch-store epilogue decoded the accumulator only through the MFMA atom, so wave32 had no two-stage path: a 16-bit dW with an odd row (a 3-channel stem) could not be split at all, and bf16 split-K stopped at the packed atomic's accuracy cap. Decode the WMMA C fragment through op.c_layout() in the scratch epilogue, as the WMMA atomic epilogue does, and enable two-stage on wave32 targets whose catalog marks float global atomics as validated (gfx1250; the f32 scratch adds are global float atomics). The workspace reduction accumulates in fp32, so the split policy resolves its degree with the fp32 cap. From the cache on wave32, bf16 also runs the two-stage kernels by default (--two-stage auto), next to the packed-atomic ones. The MFMA path's IR is unchanged. Mirrored in the C++ engine (availability gate and the WMMA scratch epilogue). Parity configs 63-66: fp16 dW with an odd dW row (C=3), bf16 dW on 64x64x64 2x2 warps K-outer, fp16 dW over the async K loop with one scratch slab, and the gfx1151 reject. One new representative-IR case. CPU tests cover the availability predicate and the gfx1151 reject; a new gfx1250 GPU test runs two-stage on an odd-dW-row stem and a 2x2 stride-2 downsample, fp16 / bf16, plain, K-outer and async K loops. Co-Authored-By: Claude Opus 5.5 --- .../common/benchmark_implicit_gemm_conv.py | 13 ++- .../common/README_conv_implicit_gemm_wgrad.md | 22 +++- .../common/conv_implicit_gemm_wgrad.py | 76 +++++++----- .../library/kernels/common/conv_split_k.py | 17 ++- .../parity/conv_implicit_gemm_wgrad_emit.c | 40 +++++++ .../parity/conv_implicit_gemm_wgrad_emit.py | 48 +++++++- .../tests/test_conv_wgrad_correctness.py | 109 ++++++++++++++---- .../common/conv_implicit_gemm_wgrad.cpp | 89 +++++++------- .../dsl_docs/instances/convolution.md | 2 +- .../python/rocke/instances/SUPPORT_MATRIX.md | 2 +- .../conv_wgrad_two_stage.json | 32 +++++ .../differential/golden/llvm_gfx_all.json | 15 +++ .../instances/rocke_ir_parity_harness.py | 21 ++++ 13 files changed, 374 insertions(+), 112 deletions(-) diff --git a/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/common/benchmark_implicit_gemm_conv.py b/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/common/benchmark_implicit_gemm_conv.py index 7a1bb5b28aa9..ac004758fb1b 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/common/benchmark_implicit_gemm_conv.py +++ b/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( diff --git a/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/README_conv_implicit_gemm_wgrad.md b/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/README_conv_implicit_gemm_wgrad.md index 7a5598bb835a..470eb7d1be8e 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/README_conv_implicit_gemm_wgrad.md +++ b/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/README_conv_implicit_gemm_wgrad.md @@ -45,10 +45,12 @@ 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. --- @@ -114,8 +116,16 @@ dW rows, and the extra memory transactions outweigh the halved atomic count. 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`) diff --git a/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/conv_implicit_gemm_wgrad.py b/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/conv_implicit_gemm_wgrad.py index 0fed277ba131..78396a6f8cf4 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/conv_implicit_gemm_wgrad.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/conv_implicit_gemm_wgrad.py @@ -1437,12 +1437,18 @@ def _wmma_async_available(spec: "WgradConvSpec", arch: str) -> Tuple[bool, str]: def wgrad_two_stage_available(arch: str) -> bool: """Whether ``arch`` has the two-stage (f32 scratch) split-K wgrad path. - The scratch-store epilogue decodes the C fragment through the MFMA atom, so - wave32 (WMMA) targets have none. is_valid_wgrad_spec rejects a two-stage - spec wherever this is False, so a caller choosing between the split paths - can ask the same question before it builds. + The scratch-store epilogue decodes the C fragment through the MFMA atom on + wave64 and through the WMMA C layout (``op.c_layout()``) on wave32. Its + f32 scratch adds are global float atomics, so a wave32 target needs the + same catalog flag as the WMMA split-K atomics + (:func:`wgrad_wmma_split_k_available`: gfx1250). is_valid_wgrad_spec + rejects a two-stage spec wherever this is False, so a caller choosing + between the split paths can ask the same question before it builds. """ - return _wgrad_mma_family(arch) != "wmma" + family = _wgrad_mma_family(arch) + if family == "wmma": + return wgrad_wmma_split_k_available(arch) + return family == "mma" _STREAMK_MODES = ("off", "dp_sk", "persistent") @@ -1915,18 +1921,16 @@ def is_valid_wgrad_spec(spec: WgradConvSpec, arch: str = "gfx950") -> Tuple[bool # Mirrors the builder's `_is_two_stage`: two-stage applies whenever the # reduction is split. See the matching comment in WgradConvSpec.validate(). _effective_two_stage = spec.two_stage and _is_atomic - # The two-stage scratch-atomic epilogue is MFMA-only. The packed *atomic* - # epilogue does have a WMMA variant (_emit_wgrad_split_k_epilogue_wmma), so - # split-K itself is fine on wave32 -- but _emit_wgrad_workspace_store_epilogue - # calls c_warp_params(atom), and `atom` is None on the WMMA path. The - # epilogue dispatch tests _is_two_stage BEFORE the wmma branch, so a - # two-stage wave32 spec reaches the MFMA-only emitter and dies with an - # AttributeError rather than a validation error. Reject it here, where every - # pre-filter (dispatch support(), the sweep drivers, benchmarks) can see it. + # The two-stage scratch-atomic epilogue decodes the WMMA C fragment through + # op.c_layout() and the MFMA one through the atom; its f32 scratch adds need + # the arch's float global atomics on wave32 (wgrad_two_stage_available). + # Reject it here, where every pre-filter (dispatch support(), the sweep + # drivers, benchmarks) can see it. if _effective_two_stage and not wgrad_two_stage_available(arch): return False, ( - f"two-stage wgrad is CDNA-only (got family 'wmma' on " - f"{arch}); the scratch-atomic epilogue has no WMMA variant" + f"two-stage wgrad on WMMA needs validated float global atomics " + f"(has_global_float_atomic_add; got {arch}): the scratch epilogue " + f"adds f32 partials atomically" ) # WMMA split-K is enabled for every dW dtype where the arch catalog marks # the float global atomics as run on a device. @@ -3533,6 +3537,7 @@ def compute_dyn_wg(k_offset_val, buf_pair, state): c_per_lane, params, gm_group=group_v, + **({"op": op} if op.family == "wmma" else {}), ) elif _is_split_k and op.family == "wmma" and spec.epilogue != "cshuffle": # WMMA split-K: atomic-add via the WMMA C-fragment layout (fp32/bf16/fp16). @@ -4219,9 +4224,14 @@ def _emit_wgrad_workspace_store_epilogue( params: Dict[str, Value], gm_group: Optional[Value] = None, replica_sel: Optional[Value] = None, + op=None, ) -> None: """Two-stage Stage 1 epilogue: f32 atomic-add into the scratch accumulator. + ``op`` (WMMA only): the accumulator is a WMMA C fragment, decoded through + ``op.c_layout()`` as :func:`_emit_wgrad_split_k_epilogue_wmma` does; the MFMA + path decodes it through ``atom``. + Identical coordinate computation to :func:`_emit_wgrad_split_k_epilogue`, but the add lands in an f32 scratch buffer instead of in ``dW``: @@ -4316,23 +4326,29 @@ def _replica_src() -> Value: block_warp_n_off = b.add(block_n_off, warp_n_off) # Decode C-fragment layout (same as atomic epilogue). - from rocke.helpers.atoms import c_warp_params, make_c_warp_dstr_encoding - from rocke.helpers.distribution import make_static_tile_distribution - - _, __, kc_m1, kc_nlane = c_warp_params(atom) - c_dist = make_static_tile_distribution(make_c_warp_dstr_encoding(atom)) - c_nlane = b.const_i32(kc_nlane) - n_in_atom = b.mod(lane, c_nlane) - m_blk = b.div(lane, c_nlane) - p_lane = [m_blk, n_in_atom] - rows: List[Value] = [] cols: List[Value] = [] - for i in range(c_per_lane): - ys = [b.const_i32(i // kc_m1), b.const_i32(i % kc_m1)] - x_row, x_col = c_dist.calculate_x(b, ys=ys, ps=[p_lane]) - rows.append(x_row) - cols.append(x_col) + if op is not None and op.family == "wmma": + c_map = op.c_layout() + for i in range(c_per_lane): + x_row, x_col = c_map.coord(b, lane, i) + rows.append(x_row) + cols.append(x_col) + else: + from rocke.helpers.atoms import c_warp_params, make_c_warp_dstr_encoding + from rocke.helpers.distribution import make_static_tile_distribution + + _, __, kc_m1, kc_nlane = c_warp_params(atom) + c_dist = make_static_tile_distribution(make_c_warp_dstr_encoding(atom)) + c_nlane = b.const_i32(kc_nlane) + n_in_atom = b.mod(lane, c_nlane) + m_blk = b.div(lane, c_nlane) + p_lane = [m_blk, n_in_atom] + for i in range(c_per_lane): + ys = [b.const_i32(i // kc_m1), b.const_i32(i % kc_m1)] + x_row, x_col = c_dist.calculate_x(b, ys=ys, ps=[p_lane]) + rows.append(x_row) + cols.append(x_col) flat = 0 for mi in range(mfmas_m): diff --git a/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/conv_split_k.py b/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/conv_split_k.py index 0ace8899caaf..a183d1bc5a70 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/conv_split_k.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/conv_split_k.py @@ -24,8 +24,9 @@ (:func:`wgrad_atomic_epilogue_available`) and accumulates rounding error with the degree (:data:`ATOMIC_16BIT_MAX_SPLIT_K`); a merged tile reduces through the workspace only; the workspace must fit the i32 ``ws_bytes`` kernarg; - WMMA (wave32) arches have no workspace path, and their packed atomic needs - neither cshuffle nor an even store vector; + on WMMA (wave32) the workspace path needs the arch's float global atomics + (gfx1250), and the packed atomic needs neither cshuffle nor an even store + vector; * the gridDim.z budget (``z = groups * split_k``); * for the forward split (the two-stage WMMA kernel, built on :data:`FWD_SPLIT_K_ARCHS` only): its GEMM ``N*[Do*]Ho*Wo x kpg x @@ -222,7 +223,8 @@ def _resolve_wgrad( atomic_why = "the default epilogue emits no packed 16-bit atomic" atomic_ok = atomic_why == "ok" # Deliberately unmerged: the workspace is R copies of the true dW. WMMA - # arches have no two-stage path, so no workspace at all. + # arches without float global atomics have no two-stage path, so no + # workspace at all. ws_ok = ( wgrad_two_stage_available(arch) and wgrad_workspace_nbytes(p, ws_replicas) <= MAX_WGRAD_WS_BYTES @@ -294,12 +296,15 @@ def _rule(requested: int, source: str, *, atomic_cap: bool) -> ConvSplitK: groups=grid_groups, block_size=block_size, # Selects the arch's WgradSplitPolicy (gfx1250) where it has one. - dtype_d=dtype, + # The workspace reduction accumulates in fp32 whatever dW's dtype, + # so the 16-bit packed-atomic accuracy caps do not apply to it. + dtype_d="fp32" if reduction == "workspace" else dtype, ) if reduction == "atomic" and not atomic_ok and allow_unsplit: # An auto degree the packed atomic cannot serve runs unsplit, as - # on the planner path (e.g. a WMMA 16-bit dW with an odd row: no - # packed pair and no workspace). A forced degree still raises. + # on the planner path (e.g. a WMMA 16-bit dW with an odd row and + # the reduction fixed to the atomic: no packed pair). A forced + # degree still raises. return ConvSplitK( 1, False, diff --git a/dnn-providers/hip-kernel-provider/rocke/library/tests/parity/conv_implicit_gemm_wgrad_emit.c b/dnn-providers/hip-kernel-provider/rocke/library/tests/parity/conv_implicit_gemm_wgrad_emit.c index a408008b86ff..3ad7b81f92ce 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/tests/parity/conv_implicit_gemm_wgrad_emit.c +++ b/dnn-providers/hip-kernel-provider/rocke/library/tests/parity/conv_implicit_gemm_wgrad_emit.c @@ -500,6 +500,46 @@ static int make_base_cfg(int idx, rocke_implicit_gemm_conv_wgrad_spec_t* spec, c spec->split_k = 4; *arch = "gfx1250"; return 0; + case 63: + case 64: + case 65: + case 66: + { + /* Two-stage (f32 workspace) split-K on WMMA. Must match the Python + * emitter's 63-66. */ + const bool gfx1250 = idx != 66; + const bool big = idx == 64; + if(idx == 63) + spec->problem = rocke_conv_problem_default(4, 12, 12, 3, 32, 3, 3); + else if(idx == 65) + spec->problem = rocke_conv_problem_default(4, 12, 12, 24, 32, 3, 3); + else + spec->problem = rocke_conv_problem_default(8, 56, 56, 64, 64, 3, 3); + spec->problem.pH = 1; + spec->problem.pW = 1; + spec->dtype_a = "fp16"; + spec->dtype_b = "fp16"; + spec->dtype_d = idx == 64 ? "bf16" : "fp16"; + spec->tile_m = big ? 64 : 32; + spec->tile_n = big ? 64 : 32; + spec->tile_k = big ? 64 : 32; + spec->warp_m = big ? 2 : 1; + spec->warp_n = big ? 2 : 1; + spec->warp_tile_m = 16; + spec->warp_tile_n = 16; + spec->warp_tile_k = gfx1250 ? 32 : 16; + spec->wave_size = 32; + spec->pipeline = "mem"; + spec->epilogue = "default"; + spec->split_k = 4; + spec->two_stage = true; + spec->lds_k_outer = idx == 64 || idx == 65; + spec->async_dma = idx == 65; + if(idx == 65) + spec->ws_replicas = 1; + *arch = gfx1250 ? "gfx1250" : "gfx1151"; + return 0; + } default: return -1; } diff --git a/dnn-providers/hip-kernel-provider/rocke/library/tests/parity/conv_implicit_gemm_wgrad_emit.py b/dnn-providers/hip-kernel-provider/rocke/library/tests/parity/conv_implicit_gemm_wgrad_emit.py index 78c746e165bd..57347e7c6dc5 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/tests/parity/conv_implicit_gemm_wgrad_emit.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/tests/parity/conv_implicit_gemm_wgrad_emit.py @@ -70,6 +70,12 @@ # 62 -- both engines must reject: gfx1250 cshuffle split-K fp16 dW whose # store vector is odd (C=3 with a 2x2 filter: wg_N=12 is even, but # the LDS pairing has no partner) +# 63-65 -- gfx1250 two-stage (f32 workspace) split-K on WMMA, the scratch +# epilogue decoding the WMMA C fragment: fp16 dW with an odd dW row +# (C=3, wg_N=27, M-outer), bf16 dW 64x64x64 2x2 warps K-outer, and fp16 +# dW over the async K loop with one scratch slab (ws_replicas=1) at C=24 +# 66 -- both engines must reject: two-stage on gfx1151 (no validated float +# global atomics for the scratch adds) # (Grouped wgrad reaches the C++ engine only through stream-K, and its tile # loader has no elem_dtype, so no case loads bf16 operands. The stream-K # partition and CTA pool are kernargs, so no case varies them. async_dma is @@ -85,7 +91,7 @@ # 107 -- streamk with unroll_k (must raise ValueError) # 108 -- streamk atomic reduction into fp16 dW (must raise ValueError) # (These illustrate the validator contract. The C emitter defines only cases -# 0-62, so run_diff.py stops at the shared END before reaching 100+; these +# 0-66, so run_diff.py stops at the shared END before reaching 100+; these # configs are not exercised by the differential gate.) from kernels.common.conv_implicit_gemm_wgrad import ( WgradConvSpec, @@ -799,6 +805,46 @@ def _spec(idx: int): "gfx1250", ) + if 63 <= idx <= 66: + # Two-stage (f32 workspace) split-K on WMMA. + from kernels.common._conv_implicit_gemm_common import ConvDataSpec + + p64 = ConvProblem(N=8, Hi=56, Wi=56, C=64, K=64, Y=3, X=3, pH=1, pW=1) + p24 = ConvProblem(N=4, Hi=12, Wi=12, C=24, K=32, Y=3, X=3, pH=1, pW=1) + p_c3 = ConvProblem(N=4, Hi=12, Wi=12, C=3, K=32, Y=3, X=3, pH=1, pW=1) + problem, dtype_d, big, extra, arch = { + 63: (p_c3, "fp16", False, {}, "gfx1250"), + 64: (p64, "bf16", True, {"lds_k_outer": True}, "gfx1250"), + 65: ( + p24, + "fp16", + False, + {"lds_k_outer": True, "async_dma": True, "ws_replicas": 1}, + "gfx1250", + ), + 66: (p64, "fp16", False, {}, "gfx1151"), + }[idx] + gfx1250 = arch == "gfx1250" + kw = dict( + problem=problem, + data=ConvDataSpec(dtype_a="fp16", dtype_b="fp16", dtype_d=dtype_d), + tile_m=64 if big else 32, + tile_n=64 if big else 32, + tile_k=64 if big else 32, + warp_m=2 if big else 1, + warp_n=2 if big else 1, + warp_tile_m=16, + warp_tile_n=16, + warp_tile_k=32 if gfx1250 else 16, + wave_size=32, + pipeline="mem", + epilogue="default", + split_k=4, + two_stage=True, + ) + kw.update(extra) + return WgradConvSpec(**kw), arch + # ---------------------------------------------------------------- # Negative cases: these specs must be REJECTED by the validator. # The harness (run_emit) expects a ValueError / SystemExit when diff --git a/dnn-providers/hip-kernel-provider/rocke/library/tests/test_conv_wgrad_correctness.py b/dnn-providers/hip-kernel-provider/rocke/library/tests/test_conv_wgrad_correctness.py index fabcb9c509b2..32a4f5fd0c2e 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/tests/test_conv_wgrad_correctness.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/tests/test_conv_wgrad_correctness.py @@ -2199,14 +2199,14 @@ def test_two_stage_with_split_k_1_rejected_by_predicate(self): self.assertIn("two_stage", why) -class TestWgradTwoStageIsCdnaOnly(unittest.TestCase): - """Two-stage wgrad must be rejected on WMMA rather than crashing the builder. - - ``_emit_wgrad_workspace_store_epilogue`` is MFMA-only -- it calls - ``c_warp_params(atom)`` and ``atom`` is None on wave32 -- and the epilogue - dispatch tests ``_is_two_stage`` before the WMMA branch. Without a validator - gate a two-stage wave32 spec reaches that emitter and dies with an - ``AttributeError``, which no caller pre-filters against. +class TestWgradTwoStageOnWmma(unittest.TestCase): + """Two-stage wgrad on WMMA: built on gfx1250, rejected cleanly elsewhere. + + ``_emit_wgrad_workspace_store_epilogue`` decodes the WMMA C fragment through + ``op.c_layout()`` (the MFMA one through the atom). Its f32 scratch adds are + global float atomics, so a wave32 arch without validated float atomics + (gfx1151) must get a validation error, never an exception out of the + emitter. """ def _gfx1250_spec(self, **kw): @@ -2216,9 +2216,10 @@ def _gfx1250_spec(self, **kw): ) from kernels.common.conv_implicit_gemm_wgrad import WgradConvSpec + dtype_d = kw.pop("dtype_d", "fp32") base = dict( problem=ConvProblem(N=8, Hi=56, Wi=56, C=64, K=64, Y=3, X=3, pH=1, pW=1), - data=ConvDataSpec(dtype_a="fp16", dtype_b="fp16", dtype_d="fp32"), + data=ConvDataSpec(dtype_a="fp16", dtype_b="fp16", dtype_d=dtype_d), tile_m=32, tile_n=32, tile_k=32, @@ -2234,28 +2235,36 @@ def _gfx1250_spec(self, **kw): base.update(kw) return WgradConvSpec(**base) - def test_two_stage_rejected_on_wmma(self): + def test_two_stage_admitted_on_gfx1250(self): from kernels.common.conv_implicit_gemm_wgrad import ( + build_implicit_gemm_conv_wgrad, is_valid_wgrad_spec, + wgrad_two_stage_available, ) - ok, why = is_valid_wgrad_spec( - self._gfx1250_spec(split_k=4, two_stage=True), "gfx1250" - ) - self.assertFalse(ok, "two-stage on WMMA must be rejected") - self.assertIn("CDNA", why) + self.assertTrue(wgrad_two_stage_available("gfx1250")) + for dtype_d in ("fp32", "fp16", "bf16"): + spec = self._gfx1250_spec(split_k=4, two_stage=True, dtype_d=dtype_d) + ok, why = is_valid_wgrad_spec(spec, "gfx1250") + self.assertTrue(ok, f"two-stage on gfx1250 WMMA ({dtype_d}): {why}") + build_implicit_gemm_conv_wgrad(spec, arch="gfx1250") - def test_two_stage_build_raises_value_error_not_attribute_error(self): - # The failure mode that matters: a clean ValueError a caller can handle, - # never an AttributeError out of the epilogue emitter. + def test_two_stage_rejected_without_float_atomics(self): + # gfx1151 has no validated float global atomics: a clean ValueError a + # caller can handle, never an exception out of the epilogue emitter. from kernels.common.conv_implicit_gemm_wgrad import ( build_implicit_gemm_conv_wgrad, + is_valid_wgrad_spec, + wgrad_two_stage_available, ) + self.assertFalse(wgrad_two_stage_available("gfx1151")) + spec = self._gfx1250_spec(split_k=4, two_stage=True, warp_tile_k=16) + ok, why = is_valid_wgrad_spec(spec, "gfx1151") + self.assertFalse(ok) + self.assertIn("has_global_float_atomic_add", why) with self.assertRaises(ValueError): - build_implicit_gemm_conv_wgrad( - self._gfx1250_spec(split_k=4, two_stage=True), arch="gfx1250" - ) + build_implicit_gemm_conv_wgrad(spec, arch="gfx1151") def test_split_k_atomic_still_valid_on_wmma(self): # The gate is two-stage-specific: the packed atomic epilogue DOES have a @@ -2275,9 +2284,9 @@ def test_two_stage_available_predicate(self): wgrad_two_stage_available, ) - for arch in ("gfx942", "gfx950"): + for arch in ("gfx942", "gfx950", "gfx1250"): self.assertTrue(wgrad_two_stage_available(arch), arch) - for arch in ("gfx1250", "gfx1151", "gfx1201"): + for arch in ("gfx1151", "gfx1201"): self.assertFalse(wgrad_two_stage_available(arch), arch) def test_16bit_split_k_atomic_arch_scope(self): @@ -2645,6 +2654,60 @@ def test_cshuffle_async(self): self._sweep(async_dma=True) +@unittest.skipIf(_skip_reason(), _skip_reason()) +class TestWgradTwoStageWmmaNumerics(unittest.TestCase): + """gfx1250 dW numerics for two-stage (f32 workspace) split-K on WMMA. + + The Stage 1 scratch epilogue decodes the WMMA C fragment through + op.c_layout(); a wrong lane-to-element mapping only shows in dW. Covers an + odd dW row (a 3-channel stem, wg_N = 27, which the packed atomic cannot + serve) and a 2x2 stride-2 downsample, fp16 and bf16, with the plain and the + async K loop. + """ + + _CASES = ( + _Shape( + "stem_C3_N4H32W32C3K32", N=4, Hi=32, Wi=32, C=3, K=32, Y=3, X=3, pH=1, pW=1 + ), + _Shape( + "2x2s2_N4H32W32C32K64", + N=4, + Hi=32, + Wi=32, + C=32, + K=64, + Y=2, + X=2, + sH=2, + sW=2, + ), + ) + + def _sweep(self, **knobs) -> None: + if GPU_ARCH != "gfx1250": + self.skipTest(f"two-stage WMMA wgrad is gfx1250-only; on {GPU_ARCH}") + for shape in self._CASES: + if knobs.get("async_dma") and shape.C % 2: + continue # the async copy needs an even channel run + for dtype in _DTYPES: + for split_k in (4, 16): + with self.subTest(shape=shape.id, dtype=dtype, sk=split_k): + ok, why = _check_two_stage( + shape, dtype, "mem", split_k, seed=7, **knobs + ) + self.assertTrue(ok, f"{shape.id} {dtype} {knobs}: {why}") + self.assertFalse(why.startswith("split_k=1"), why) + + def test_two_stage(self): + self._sweep() + + def test_two_stage_kouter(self): + self._sweep(lds_k_outer=True) + + def test_two_stage_async(self): + self._sweep(lds_k_outer=True, async_dma=True) + + @unittest.skipUnless(not _SKIP_REASON, _SKIP_REASON or "needs CDNA GPU + torch") class TestWgradGroupMergeNumerics(unittest.TestCase): """GPU numerics for group merging, over the two-stage path. diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/common/conv_implicit_gemm_wgrad.cpp b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/common/conv_implicit_gemm_wgrad.cpp index 885636a1eae4..ff5202676027 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/common/conv_implicit_gemm_wgrad.cpp +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/common/conv_implicit_gemm_wgrad.cpp @@ -848,18 +848,9 @@ bool rocke_implicit_gemm_conv_wgrad_is_valid_spec(const rocke_implicit_gemm_conv const bool is_wmma_family = target->wave_size == 32; if(sk > 1) { - /* The two-stage scratch-atomic epilogue is MFMA-only: it decodes the C - * fragment through ctx->atom, which is NULL on the WMMA path. */ - if(s->two_stage && is_wmma_family) - { - if(reason && reason_cap) - snprintf(reason, - reason_cap, - "two-stage wgrad is CDNA-only (got family 'wmma' on %s); the " - "scratch-atomic epilogue has no WMMA variant", - arch); - return false; - } + /* The two-stage scratch-atomic epilogue decodes the C fragment through + * ctx->atom on MFMA and through op.c_layout() on WMMA, so every wave32 and + * wave64 target has it (Python wgrad_two_stage_available). */ /* WMMA split-K is enabled for every dW dtype where the arch catalog marks * the float global atomics as run on a device. Mirrors Python * wgrad_wmma_split_k_available. */ @@ -1911,39 +1902,53 @@ static void wgrad_emit_workspace_store_epilogue(rocke_ir_builder_t* b, rocke_value_t* block_warp_m_off = rocke_b_add(b, ctx->block_m_off_v, warp_m_off); rocke_value_t* block_warp_n_off = rocke_b_add(b, ctx->block_n_off_v, warp_n_off); - /* 4. c_warp_params: compile-time lookup, no IR ops -- Python: c_warp_params(atom) */ - int m0, m_lane, m1, n_lane; - if(rocke_b_c_warp_params(b, atom, &m0, &m_lane, &m1, &n_lane) != ROCKE_OK) - return; - - /* 5. c_dist -- Python: c_dist = make_static_tile_distribution(make_c_warp_dstr_encoding) */ - rocke_tile_distribution_encoding_t* enc = rocke_make_c_warp_dstr_encoding(b, atom); - if(enc == NULL) - return; - rocke_tile_distribution_t* c_dist = rocke_make_static_tile_distribution(b, enc); - if(c_dist == NULL) - return; - - /* 6. c_nlane, n_in_atom, m_blk -- Python ordering */ - rocke_value_t* c_nlane_v = rocke_b_const_i32(b, n_lane); - rocke_value_t* n_in_atom = rocke_b_mod(b, ctx->lane, c_nlane_v); - rocke_value_t* m_blk = rocke_b_div(b, ctx->lane, c_nlane_v); - rocke_value_t* p_lane_subs[2] = {m_blk, n_in_atom}; - rocke_value_t* const* p_arr[1] = {p_lane_subs}; - int p_counts[1] = {2}; - - /* 7. Pre-compute per-slot (row, col) within the atom */ + /* 4-7. Per-slot (row, col) within the atom. WMMA (Python op is not None): + * the C fragment is decoded through op.c_layout(), as the WMMA atomic + * epilogue does. MFMA: through the atom's C-warp distribution. */ rocke_value_t* slot_rows[ROCKE_CONV_MAX_ACCS * 4]; rocke_value_t* slot_cols[ROCKE_CONV_MAX_ACCS * 4]; - for(int i = 0; i < c_per_lane; ++i) + if(ctx->is_wmma) { - rocke_value_t* y0 = rocke_b_const_i32(b, i / m1); - rocke_value_t* y1 = rocke_b_const_i32(b, i % m1); - rocke_value_t* ys[2] = {y0, y1}; - rocke_value_t* out_x[2] = {NULL, NULL}; - rocke_tile_distribution_calculate_x(b, c_dist, ys, 2, p_arr, p_counts, 1, out_x, 2); - slot_rows[i] = out_x[0]; - slot_cols[i] = out_x[1]; + const rocke_arch_layout_map_t* c_map = rocke_mmaop_c_layout(ctx->op, b); + if(c_map == NULL) + return; + for(int i = 0; i < c_per_lane; ++i) + rocke_arch_layout_map_coord(c_map, b, ctx->lane, i, &slot_rows[i], &slot_cols[i]); + } + else + { + /* 4. c_warp_params: compile-time lookup, no IR ops -- Python: c_warp_params(atom) */ + int m0, m_lane, m1, n_lane; + if(rocke_b_c_warp_params(b, atom, &m0, &m_lane, &m1, &n_lane) != ROCKE_OK) + return; + + /* 5. c_dist -- Python: c_dist = make_static_tile_distribution(make_c_warp_dstr_encoding) */ + rocke_tile_distribution_encoding_t* enc = rocke_make_c_warp_dstr_encoding(b, atom); + if(enc == NULL) + return; + rocke_tile_distribution_t* c_dist = rocke_make_static_tile_distribution(b, enc); + if(c_dist == NULL) + return; + + /* 6. c_nlane, n_in_atom, m_blk -- Python ordering */ + rocke_value_t* c_nlane_v = rocke_b_const_i32(b, n_lane); + rocke_value_t* n_in_atom = rocke_b_mod(b, ctx->lane, c_nlane_v); + rocke_value_t* m_blk = rocke_b_div(b, ctx->lane, c_nlane_v); + rocke_value_t* p_lane_subs[2] = {m_blk, n_in_atom}; + rocke_value_t* const* p_arr[1] = {p_lane_subs}; + int p_counts[1] = {2}; + + /* 7. Pre-compute per-slot (row, col) within the atom */ + for(int i = 0; i < c_per_lane; ++i) + { + rocke_value_t* y0 = rocke_b_const_i32(b, i / m1); + rocke_value_t* y1 = rocke_b_const_i32(b, i % m1); + rocke_value_t* ys[2] = {y0, y1}; + rocke_value_t* out_x[2] = {NULL, NULL}; + rocke_tile_distribution_calculate_x(b, c_dist, ys, 2, p_arr, p_counts, 1, out_x, 2); + slot_rows[i] = out_x[0]; + slot_cols[i] = out_x[1]; + } } /* 8. Main mi/ni/i loop */ diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/instances/convolution.md b/dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/instances/convolution.md index 6ef4dc6bda54..e7954618af69 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/instances/convolution.md +++ b/dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/instances/convolution.md @@ -463,7 +463,7 @@ M-outer default for dgrad, and matches the fp32 reference for wgrad. - `pipeline="basic"`: now uses a bounded runtime `scf.for` K-loop (load → sync → MFMA → sync per tile), identical to the `"mem"` path. No K-trip-count limit. dgrad has no `basic` branch. - `async_dma` is a swept axis on the wgrad sweep driver, not a run-level flag. Unlike `lds_k_outer` it is not deducible from `(arch, spec)`: it removes the register staging of the tile, but it also forces the K-outer row pad to 0 and coarsens the load-width ladder to the widths the intrinsic accepts, and both terms are functions of tile width and channel run, which are themselves sweep axes. -- WMMA wgrad (gfx1250): `async_dma` is a ping-pong over `global_load_async_to_lds` (`CoalescedTileLoader.load_async_to_lds`, one barrier per tile, drained by `s_wait_asynccnt`). It reuses the synchronous K-outer loaders' slot partition, needs `lds_k_outer` and a 4-, 8- or 16-byte chunk on both operands, and ANDs the buffer-range check (`byte offset < num_bytes`) into each chunk's `valid`, zero-filling failing chunks with an LDS store. `epilogue="cshuffle"` is admitted on WMMA too: both cshuffle epilogues (direct store and split-K atomic) read the WMMA fragment through `CShuffleEpilogue.from_grid_op`. `pipeline` stays `mem`; `unroll_k`, `prefetch_addr` and `chiplet_swizzle` stay rejected on WMMA. +- WMMA wgrad (gfx1250): `async_dma` is a ping-pong over `global_load_async_to_lds` (`CoalescedTileLoader.load_async_to_lds`, one barrier per tile, drained by `s_wait_asynccnt`). It reuses the synchronous K-outer loaders' slot partition, needs `lds_k_outer` and a 4-, 8- or 16-byte chunk on both operands, and ANDs the buffer-range check (`byte offset < num_bytes`) into each chunk's `valid`, zero-filling failing chunks with an LDS store. `epilogue="cshuffle"` is admitted on WMMA too: both cshuffle epilogues (direct store and split-K atomic) read the WMMA fragment through `CShuffleEpilogue.from_grid_op`. Two-stage (f32 workspace) split-K runs on gfx1250 as well: the Stage 1 scratch epilogue decodes the WMMA fragment through `op.c_layout()` (`wgrad_two_stage_available` needs the arch's float global atomics on wave32). `pipeline` stays `mem`; `unroll_k`, `prefetch_addr` and `chiplet_swizzle` stay rejected on WMMA. - No K loop is unrolled at build time any more: `async_dma` runs `SoftwarePipeline.run_ping_pong_dynamic` and `unroll_k` a hand-rolled 2x ping-pong, both over the runtime extent, so the former `_MAX_UNROLLED_K_ITERS` / `ROCKE_MAX_UNROLLED_K_ITERS` cap is gone. - wgrad `split_k`: `1` = no split (direct store); any value `> 1` builds the one split-K kernel (name tag `spk`), whose degree is the `ks_count` / `ks` kernargs chosen at launch. There is no `split_k=0` "runtime degree" encoding. - Removed: the `--lds-k-outer`, `--lds-k-pad` and `--dtype-d` CLI flags and the `ROCKE_WGRAD_LDS_K_OUTER` env override. All replaced by deduction or dropped. diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/instances/SUPPORT_MATRIX.md b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/instances/SUPPORT_MATRIX.md index 951f3b89d730..ca7f82c1ded1 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/instances/SUPPORT_MATRIX.md +++ b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/instances/SUPPORT_MATRIX.md @@ -109,7 +109,7 @@ be smaller than 2 GiB. | `direct_depthwise_dgrad` (`DirectDepthwiseDgradStreamSpec`) | ✅ | ✅ | ❌ | `cpg=kpg=1`, any stride and padding; row-streaming scalar FMA at the target's wave size (wave64, or wave32 on gfx1250); `cpl` 2/4/8 channels per lane with packed `<2 x f32>` FMAs (`groups % cpl == 0`, `cpl*KH*KW <= 200`, `cpl*block_w*KH <= 192`). In the AOT direct cache; not yet wired into dispatch | | `direct_depthwise_tiled_dgrad` (`DirectDepthwiseTiledSpec(dgrad=True)`) | ✅ | ✅ | ❌ | `cpg=kpg=1`, stride 1, the forward tiled kernel with the filter rotated by 180 degrees; filters up to 64 on a side incl. non-square (51x5 / 5x51), per-axis padding (`pad_w`) with `PAD <= KH-1`, `PAD_W <= KW-1`. Served by the direct-conv AOT cache, not the dispatcher | | `conv_implicit_gemm_dgrad` stride-1 knobs | ✅ | ✅ | ❌ | folded sub-GEMM record (stride capability), `tap_outer_k` (wave64 MFMA, `kpg % tile_k == 0`), `dy_halo` / `dy_halo_setprio` / `dy_halo_kouter_pad` / `halo_max_wo` (on the tap-outer path, `Ho == Hi`, `Wo == Wi`; the pad needs `lds_k_outer`, gfx950); `xcd_tile_order` remaps only on gfx950. AOT grid: stride-1 tap-outer binaries next to the tilde ones, halo sub-grid on gfx950; not yet wired into dispatch | -| `conv_implicit_gemm_wgrad` | ✅ | ✅ | ✅ | gfx1151/gfx1201/gfx1250: WMMA, `mem` with the `default` or `cshuffle` epilogue (gfx1250 adds the `async_dma` K loop over `global_load_async_to_lds`); split-K (WMMA packed atomic, no two-stage) on gfx1250 only, automatic for bf16/fp16/fp32, `split_k=1` on gfx1151/gfx1201 | +| `conv_implicit_gemm_wgrad` | ✅ | ✅ | ✅ | gfx1151/gfx1201/gfx1250: WMMA, `mem` with the `default` or `cshuffle` epilogue (gfx1250 adds the `async_dma` K loop over `global_load_async_to_lds`); split-K (WMMA packed atomic, cshuffle, or two-stage f32 workspace) on gfx1250 only, automatic for bf16/fp16/fp32, `split_k=1` on gfx1151/gfx1201 | | `conv_implicit_gemm_wgrad` stream-K | ✅ | ✅ | ❌ | CK iteration-balanced partition; `linear`/`tree` deterministic fixups need the cross-workgroup flag ops | | `conv_implicit_gemm` split-K forward (`split_k > 1`) | ❌ | ❌ | ❌ | gfx1250 only: WMMA, `mem` or `wavelet` pipeline, 2-D, not pointwise, identity `acc_epilogue`; no `async_dma`, `unroll_k`, `chiplet_swizzle` or `group_merge`. Each slice stores an fp32 partial tile into its slab of a `[ks_count][M][K]` workspace and `conv_fwd_split_k_reduce` folds the slabs into D; the degree is the `ks_count` kernarg, and `split_k = -1` resolves it (`resolve_conv_split_k`; 1 where splitting does not pay). Built into the AOT kernel cache; the dispatcher does not select it | | `conv_implicit_gemm` cluster split-K (`split_k_cluster`) | ❌ | ❌ | ❌ | gfx1250 only (llvm23, workgroup-cluster launch): `mem` pipeline, `split_k_cluster` in 2..15, `tile_n % 4 == 0`, K and kpg multiples of 4 (checked at build and on every launched problem). The slices of one output tile form a (1, 1, S) workgroup cluster and fold through a cluster barrier, so there is no Stage 2 launch; one binary per S. Built into the AOT kernel cache; the dispatcher does not select it | diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/tests/golden/rocke_representative_ir_sha256/conv_wgrad_two_stage.json b/dnn-providers/hip-kernel-provider/rocke/platform/tests/golden/rocke_representative_ir_sha256/conv_wgrad_two_stage.json index e624942aa789..89a8f24bf3a9 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/tests/golden/rocke_representative_ir_sha256/conv_wgrad_two_stage.json +++ b/dnn-providers/hip-kernel-provider/rocke/platform/tests/golden/rocke_representative_ir_sha256/conv_wgrad_two_stage.json @@ -3,6 +3,14 @@ "flavors": { "llvm20": { "cases": { + "conv_wgrad_two_stage/gfx1250/n1h8c16k32r3_spk4": { + "arch": "gfx1250", + "bytes": 46072, + "family": "conv_wgrad_two_stage", + "kernel_name": "irhash_wgrad_1250_twostage_spk4_N1H8W8C16_K32Y3X3_t32x32x32_w2x1_a16x16x32_mem_default_spk_twostage_wsr8", + "sha256": "fad682169b024e1652a409d9b7b2fa0816cff7f07ae1e7c676888ae19515c8c1", + "status": "ok" + }, "conv_wgrad_two_stage/gfx942/n1h8c16k32r3_spk4": { "arch": "gfx942", "bytes": 22397, @@ -24,6 +32,14 @@ }, "llvm22": { "cases": { + "conv_wgrad_two_stage/gfx1250/n1h8c16k32r3_spk4": { + "arch": "gfx1250", + "bytes": 46211, + "family": "conv_wgrad_two_stage", + "kernel_name": "irhash_wgrad_1250_twostage_spk4_N1H8W8C16_K32Y3X3_t32x32x32_w2x1_a16x16x32_mem_default_spk_twostage_wsr8", + "sha256": "0708e8b4e717907b5ea6b8bb9a4473ba4d88e5bd730a1b0107b3dae25a4fa822", + "status": "ok" + }, "conv_wgrad_two_stage/gfx942/n1h8c16k32r3_spk4": { "arch": "gfx942", "bytes": 22530, @@ -45,6 +61,14 @@ }, "llvm23": { "cases": { + "conv_wgrad_two_stage/gfx1250/n1h8c16k32r3_spk4": { + "arch": "gfx1250", + "bytes": 46275, + "family": "conv_wgrad_two_stage", + "kernel_name": "irhash_wgrad_1250_twostage_spk4_N1H8W8C16_K32Y3X3_t32x32x32_w2x1_a16x16x32_mem_default_spk_twostage_wsr8", + "sha256": "cc9e10ca972fea299383ae990b99b6ea786b949d3fb6618de71a0b40c7986951", + "status": "ok" + }, "conv_wgrad_two_stage/gfx942/n1h8c16k32r3_spk4": { "arch": "gfx942", "bytes": 22594, @@ -66,6 +90,14 @@ }, "llvm24": { "cases": { + "conv_wgrad_two_stage/gfx1250/n1h8c16k32r3_spk4": { + "arch": "gfx1250", + "bytes": 46275, + "family": "conv_wgrad_two_stage", + "kernel_name": "irhash_wgrad_1250_twostage_spk4_N1H8W8C16_K32Y3X3_t32x32x32_w2x1_a16x16x32_mem_default_spk_twostage_wsr8", + "sha256": "cc9e10ca972fea299383ae990b99b6ea786b949d3fb6618de71a0b40c7986951", + "status": "ok" + }, "conv_wgrad_two_stage/gfx942/n1h8c16k32r3_spk4": { "arch": "gfx942", "bytes": 22594, diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/tests/instances/differential/golden/llvm_gfx_all.json b/dnn-providers/hip-kernel-provider/rocke/platform/tests/instances/differential/golden/llvm_gfx_all.json index 1f4d19a8903c..4d86bc6b817d 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/tests/instances/differential/golden/llvm_gfx_all.json +++ b/dnn-providers/hip-kernel-provider/rocke/platform/tests/instances/differential/golden/llvm_gfx_all.json @@ -569,6 +569,9 @@ "59": "5ebbd1600763884bb810b302df64677a5fef6f222bfec3fac9c470a4ad20ed5d", "6": "779030de0f57eb186f48efc6b92827294b7aa1095c2a8ee96d2edae002cbaa1d", "60": "1fda68fce4ddbfd4d1ab27127e91474868de275962bdb26e3be025ca21cb2d85", + "63": "297b72544a13aa0b519b4f248a4ce655d39e493e730ae3f25adf7429b535bcb8", + "64": "0ea52c6b7172c309a13d3aadfff7f232f2f31f251b90d80ebdc9c4e80754b1f8", + "65": "d9e8f970e68f7d3c95b6d0edad0601481f50ca1f057b774ee8d40fa782cb58a4", "7": "1106d4ce71314f3c532427fed8ceb00855f1cf1dd7f5117b89234cf326583087", "8": "82d8923b31be6a07c4371fbe40b33a2ea7452027299ce0993e903cf55ad092d8" }, @@ -2442,6 +2445,9 @@ "59": "46dd718d4f966e5b0659ebe0f8b2f6fde2ca6eac6035fd247c0b7f44df648408", "6": "b2605886350eb7b793616938e9179ad480084552a3533959f34918085a34484c", "60": "55e677df75468f06ced7b24d600f48aa934f93a5c46963966d7a438383e05c59", + "63": "595ccf6f304c516b4bbfa25ab1b0714069494ed913a96de267265f720e941d3f", + "64": "bb76cde6ee2149851980c297b858c581cd115cef526902134f613937aa0f65f8", + "65": "11930bd06e3c6f3b486d7d2bfa4b409562f82a6af038bfea5eb8539eda991085", "7": "d653088ddce4d716c3b9e20fabe5592541b52ba5e800786a86b660763c25a596", "8": "889d9e3617a8eb2e30d29db4ef8da0aeac52c9fa0d92bc53541947cf85926c08" }, @@ -4217,6 +4223,9 @@ "59": "654c0bfc412129aad45be434cb4e382fab2739ad94ffe5d607db6924e5f77b00", "6": "f3b4e8cf7ab611637137babc7994172995f76262274b0a63ca3b6999ac62ef04", "60": "cafe652716f3f73fc0747ca5a6932007dcc21baa840e1367de312a93bbfdf4b1", + "63": "961ae6b9afbfcea5088451d9d356ee9d21cf208c39c0acd95ef29d3e0bca1e08", + "64": "1fc144e9617d38a78751c4452ebbe78aa0bcc06fd4a5a3e0f90eb8cf0213b268", + "65": "6bc8acde7c696b4ec316a439952e4425364c55f936802d77798494025fcf43d2", "7": "e12021bd6f8923472a7595394ecde03a55dcd07d63d714b9adaea3efa79df887", "8": "2cd90bffada0eaec9a90caf5d5b76f2f213e4e884e3f413b13e79eef3071ec82" }, @@ -6022,6 +6031,9 @@ "59": "fbcf10500498604862078885c4482d5cbe67269770db5929f39fa0d6919d3c89", "6": "2726776342b4ffd858fce5951d771a50515d7c143e07656eb6235e10fd08bd32", "60": "33ef9a12ea16c82a1e985e2ad05b557de385f4def104b8b4abd694e6989ff42e", + "63": "9b06d34d2d4381e75b0d249249e90cea4b09d0a3409328dd835eeb256546faae", + "64": "861ee261e315bb0202240ac717921ef9fb69b24ea6949194d5190778261d4706", + "65": "a664126829bd3246ff2d5a4db943d691e6c49d04c65fec57a78fee8218289608", "7": "44a58b32e3d95d8492b23193d701925de3dd9de8cddf8a2c006b6fdf5de87cf8", "8": "af0df509c9e1bfa9bcb32621a677b633c9bd7c0b0a22012ee20acd2894b44df8" }, @@ -8082,6 +8094,9 @@ "59": "fbcf10500498604862078885c4482d5cbe67269770db5929f39fa0d6919d3c89", "6": "2726776342b4ffd858fce5951d771a50515d7c143e07656eb6235e10fd08bd32", "60": "33ef9a12ea16c82a1e985e2ad05b557de385f4def104b8b4abd694e6989ff42e", + "63": "9b06d34d2d4381e75b0d249249e90cea4b09d0a3409328dd835eeb256546faae", + "64": "861ee261e315bb0202240ac717921ef9fb69b24ea6949194d5190778261d4706", + "65": "a664126829bd3246ff2d5a4db943d691e6c49d04c65fec57a78fee8218289608", "7": "44a58b32e3d95d8492b23193d701925de3dd9de8cddf8a2c006b6fdf5de87cf8", "8": "af0df509c9e1bfa9bcb32621a677b633c9bd7c0b0a22012ee20acd2894b44df8" }, diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/tests/instances/rocke_ir_parity_harness.py b/dnn-providers/hip-kernel-provider/rocke/platform/tests/instances/rocke_ir_parity_harness.py index 1866e85cad4d..42a3b776ca36 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/tests/instances/rocke_ir_parity_harness.py +++ b/dnn-providers/hip-kernel-provider/rocke/platform/tests/instances/rocke_ir_parity_harness.py @@ -3142,6 +3142,27 @@ def add(family, case_id, arch, build): ), ) + # gfx1250 WMMA: the scratch epilogue decodes the WMMA C fragment. + add( + "conv_wgrad_two_stage", + "conv_wgrad_two_stage/gfx1250/n1h8c16k32r3_spk4", + "gfx1250", + build_conv_wgrad( + "irhash_wgrad_1250_twostage_spk4", + "gfx1250", + wgrad1, + wave_size=32, + wtm=16, + wtn=16, + wtk=32, + tile_m=32, + tile_n=32, + tile_k=32, + split_k=4, + two_stage=True, + ), + ) + # Workspace reduce Stage 2: reads f32 partial sums written by Stage 1 and # reduces along split_k, emitting the result as dtype_d. # wg_M=32, wg_N=72 (3*3*8): small 3x3 filter footprint, fp16 output. From 363f522711202cb31909a855b707c3c5bb3b21fc Mon Sep 17 00:00:00 2001 From: Jakub Piasecki Date: Sun, 1 Nov 2026 18:11:08 +0000 Subject: [PATCH 06/13] feat(rocke): AOT-cache the gfx1250 WMMA wgrad cshuffle, two-stage and async variants On wave32 (WMMA) the implicit-GEMM wgrad cache keeps every kernel it built before: the synchronous default-epilogue kernels, unsplit and packed-atomic split-K, on the whole tile ladder. On top of that, on 19 (tile_m, tile_n, tile_k, warp_m, warp_n) geometries (CACHE_WGRAD_WMMA_GEOMETRIES), it adds: - the cshuffle epilogue, unsplit and packed-atomic split-K; - two-stage (f32 workspace, ws_replicas = 8) split-K on the default epilogue, which the two-stage availability on gfx1250 brings in; - the async_dma K loop (global_load_async_to_lds) next to the plain loop for every non-two-stage configuration, built once per geometry with A / B widths left to the copy and its chunk widths recorded in the identity (async_chunk_a / async_chunk_b). Each cover geometry was the fastest of at least one reference wgrad layer, and together they hold the fastest measured time of every layer. count_jobs follows the same rule (and now counts the async loop), and a CPU test pins it: the default grid is unchanged and the new variants stay on the cover. The gfx950 grid is unchanged. The gfx1250 GPU split-K resolver test takes its 2x2-warp tiles from the cover and only the synchronous packed-atomic binaries. Co-Authored-By: Claude Opus 5.5 --- .../rocke/library/kernels/conv/aot_jobs.py | 166 +++++++++++++++--- .../rocke/library/tests/test_conv_abi.py | 29 +++ .../tests/test_conv_wgrad_correctness.py | 17 +- 3 files changed, 185 insertions(+), 27 deletions(-) diff --git a/dnn-providers/hip-kernel-provider/rocke/library/kernels/conv/aot_jobs.py b/dnn-providers/hip-kernel-provider/rocke/library/kernels/conv/aot_jobs.py index b7785bfaeda0..67455bf1322a 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/kernels/conv/aot_jobs.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/kernels/conv/aot_jobs.py @@ -61,7 +61,14 @@ two-stage only, load widths derived by the builder), served only to depthwise problems whose merged GEMM fits the tile. ``lds_k_outer`` and ``prefetch_addr`` follow ``WgradConvSpec.default_lds_k_outer`` / - ``default_prefetch_addr`` per geometry. + ``default_prefetch_addr`` per geometry. On wave32 (WMMA, gfx1250) the + synchronous default-epilogue kernels (split_k = 1 and atomic split-K) are + built on the whole grid; the cshuffle epilogue, two-stage split-K and the + ``async_dma`` loop (global_load_async_to_lds) only on the + ``CACHE_WGRAD_WMMA_GEOMETRIES`` cover. The async loop is built next to the + plain K loop of each non-two-stage configuration, once per geometry with + the A / B chunk widths the copy picks recorded in the identity + (``async_chunk_a`` / ``_b``). Opt-in (``streamk_modes``; ``--compile-all --streamk``): stream-K kernels, default epilogue only, ``split_k`` = 1, c = 1. At the widest vector width W of the dtype (a = W; b in {W, 1}): one per launch mode @@ -177,6 +184,39 @@ # which have no factor of 9 to chase, and widening all three directions costs # several times the disk for no measured gain. CACHE_WGRAD_EXTRA_TILE_MN: Tuple[int, ...] = (144, 192) + +# ---- wave32 (WMMA) wgrad geometries --------------------------------------- +# The WMMA wgrad variants beyond the synchronous default epilogue (cshuffle, +# two-stage split-K, the async K loop) are built on these (tile_m, tile_n, +# tile_k, warp_m, warp_n) only; the default-epilogue kernels keep the whole +# ladder above. Profiling the wgrad layers of +# the reference corpora over a pruned grid (tile_m 32..256 bounded by kpg, +# tile_n 64..256, tile_k 32/64, 4 or 8 waves), with the cached kernels +# (split_k = 1 default epilogue, split-K cshuffle), each of these was the +# fastest geometry of at least one layer, and together they hold the fastest +# measured time of every layer. The rest of the ladder never won one; the JIT +# sweep still walks it. +CACHE_WGRAD_WMMA_GEOMETRIES: Tuple[Tuple[int, int, int, int, int], ...] = ( + (32, 64, 32, 2, 4), + (32, 64, 64, 2, 2), + (32, 64, 64, 2, 4), + (32, 128, 32, 2, 2), + (32, 128, 64, 1, 4), + (32, 128, 64, 1, 8), + (32, 256, 32, 1, 8), + (32, 256, 64, 1, 4), + (64, 64, 64, 2, 4), + (64, 64, 64, 4, 2), + (64, 128, 64, 2, 2), + (64, 128, 64, 2, 4), + (128, 64, 64, 4, 2), + (128, 128, 32, 1, 4), + (128, 128, 32, 2, 2), + (128, 128, 64, 2, 2), + (128, 128, 64, 4, 2), + (256, 64, 64, 4, 2), + (256, 128, 64, 4, 2), +) CACHE_WGRAD_EXTRA_WARP_MN: Tuple[int, ...] = (3, 6) # Cap on the f32 accumulator registers each lane holds: # (tile_m / warp_m) * (tile_n / warp_n) / wave_size. Above 256 the @@ -505,6 +545,60 @@ def _fwd_k_loops(pipeline: str) -> List[Tuple[bool, bool]]: return k_loops +def _wgrad_k_loops(wave_size: int, pipeline: str, two_stage: bool) -> List[dict]: + """K-loop variants built per wgrad configuration. + + The base loop always. On wave32 (WMMA) next to "mem", except under + two-stage: the gfx1250 ``async_dma`` loop (global_load_async_to_lds). + """ + loops: List[dict] = [{}] + if wave_size == 32 and pipeline == "mem" and not two_stage: + loops.append({"async_dma": True}) + return loops + + +def _wgrad_wmma_cover(mma_family: str, geometry: tuple) -> bool: + """Whether a wgrad geometry gets the WMMA variants beyond the plain loop. + + True off WMMA (nothing is restricted there) and, on WMMA, for the + :data:`CACHE_WGRAD_WMMA_GEOMETRIES` cover. Outside the cover a WMMA + geometry builds only the synchronous default-epilogue kernels. + """ + return mma_family != "wmma" or tuple(geometry[:5]) in CACHE_WGRAD_WMMA_GEOMETRIES + + +def _wgrad_async_chunks(cfg: dict, dtypes, caps: dict, arch: str) -> Optional[dict]: + """The gfx1250 async wgrad loaders' chunk widths, for the identity. + + ``{}`` for a synchronous loop; ``None`` when the async loop cannot serve + this configuration (the job is skipped). Taken from the builder's own + width choice (_wmma_async_vecs) on the probe problem, as the fwd chunks. + """ + if not cfg.get("async_dma"): + return {} + from kernels.common._conv_implicit_gemm_common import ConvDataSpec + from kernels.common.conv_implicit_gemm_wgrad import ( + WgradConvSpec, + _wmma_async_available, + _wmma_async_vecs, + ) + + da, db, dd = dtypes + try: + spec = WgradConvSpec( + problem=_probe_problem("wgrad", caps), + data=ConvDataSpec(dtype_a=da, dtype_b=db, dtype_d=dd), + **cfg, + ) + ok, _ = _wmma_async_available(spec, arch) + if not ok: + return None + a_vec, b_vec = _wmma_async_vecs(spec) + except (ValueError, TypeError): + return None + return dict(async_chunk_a=a_vec, async_chunk_b=b_vec) + + def _wgrad_two_stages(split_k: int) -> Tuple[bool, ...]: """Two-stage variants for a split-K degree. @@ -1074,6 +1168,11 @@ def _wgrad_jobs( ): if not _aot_pipeline("wgrad", pipeline): continue + in_cover = _wgrad_wmma_cover( + mma_family, (tile_m, tile_n, tile_k, warp_m, warp_n) + ) + if epilogue != "default" and not in_cover: + continue # The K-outer LDS tile, picked per geometry by the predicate the JIT # sweep and dispatch use. It is a pure function of arch/dtype/atom, so # a shape-generic binary can carry it; leaving it at the spec default @@ -1112,6 +1211,8 @@ def _wgrad_jobs( for two_stage, caps in itertools.product( _wgrad_two_stages(split_k), _layout_caps() ): + if two_stage and not in_cover: + continue grouped = caps["grouped"] cfg = dict( tile_m=tile_m, @@ -1137,23 +1238,39 @@ def _wgrad_jobs( # Pinned in the spec and the identity alike, so the host # sizes the scratch and Stage 2 from what was compiled. cfg["ws_replicas"] = _DEFAULT_WS_REPLICAS - cfg["prefetch_addr"] = _wgrad_prefetch_addr(arch, cfg) - yield BuildJob( - identity=KernelIdentity( - arch=arch, - direction="wgrad", - algorithm="implicit_gemm", - llvm_flavor=flavor, - dtype_a=da, - dtype_b=db, - dtype_d=dd, - grouped=grouped, - **cfg, - ), - direction="wgrad", - spec_kwargs=cfg, - caps=caps, + k_loops = ( + _wgrad_k_loops(wave_size, pipeline, two_stage) if in_cover else [{}] ) + for k_loop in k_loops: + kcfg = dict(cfg, **k_loop) + if k_loop.get("async_dma"): + # The async copy picks its own chunk width (recorded + # below), so A / B widths are left unset: built once + # per geometry, on the widest-vector iteration. + if (vec_a, vec_bc) != (_VECS[-1], _VECS[-1]): + continue + kcfg.update(vector_size_a=None, vector_size_b=None) + kcfg["prefetch_addr"] = _wgrad_prefetch_addr(arch, kcfg) + chunks = _wgrad_async_chunks(kcfg, (da, db, dd), caps, arch) + if chunks is None: + continue + yield BuildJob( + identity=KernelIdentity( + arch=arch, + direction="wgrad", + algorithm="implicit_gemm", + llvm_flavor=flavor, + dtype_a=da, + dtype_b=db, + dtype_d=dd, + grouped=grouped, + **kcfg, + **chunks, + ), + direction="wgrad", + spec_kwargs=kcfg, + caps=caps, + ) # Stream-K kernels (opt-in). Default epilogue only, so the store is # scalar and B takes its own width; split_k stays 1 (stream-K owns the # K partition). The workspace reduction reuses the two-stage scratch, @@ -1727,9 +1844,18 @@ def count_jobs( if fwd_cluster and _fwd_split_cluster_geometry(geo): total += len(CACHE_FWD_SPLIT_K_CLUSTER_DEGREES) elif direction == "wgrad": - total += ( - vecs * sum(len(_wgrad_two_stages(sk)) for sk in split_ks) * layouts - ) + in_cover = _wgrad_wmma_cover(mma_family, geo) + if geo[7] != "default" and not in_cover: + continue + if in_cover: + stages = sum(len(_wgrad_two_stages(sk)) for sk in split_ks) + # The async loop: once per non-two-stage configuration. + asyncs = len(split_ks) * ( + len(_wgrad_k_loops(wave_size, pipeline, False)) - 1 + ) + else: + stages, asyncs = len(split_ks), 0 + total += vecs * stages * layouts + asyncs * layouts total += merged total += layouts * len( _wgrad_streamk_builds( diff --git a/dnn-providers/hip-kernel-provider/rocke/library/tests/test_conv_abi.py b/dnn-providers/hip-kernel-provider/rocke/library/tests/test_conv_abi.py index ef2221fb08cd..8270755f44bf 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/tests/test_conv_abi.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/tests/test_conv_abi.py @@ -1863,6 +1863,35 @@ def test_cache_async_chunk_must_divide_cpg(): assert narrow_arch == {"async_chunk_a": 2, "async_chunk_b": 2} +def test_wmma_wgrad_keeps_the_default_grid_and_covers_the_new_variants(): + """gfx1250 wgrad: every geometry keeps its synchronous default-epilogue + kernels; cshuffle, two-stage and the async loop come only on the cover. + + Walks the real (unshrunk) wgrad generator and :func:`aj.count_jobs`, which + must agree on gfx1250 too. + """ + from rocke.core.arch import ArchTarget + + target = ArchTarget.from_gfx("gfx1250") + split_ks = aj.CACHE_WGRAD_SPLIT_KS + geos = set() + walked = 0 + for job in aj._wgrad_jobs("gfx1250", "bf16", 32, "wmma", target, split_ks): + walked += 1 + i = job.identity + geo = (i.tile_m, i.tile_n, i.tile_k, i.warp_m, i.warp_n) + geos.add(geo) + # The group-merged depthwise candidates are develop's, walked on every + # geometry as before (WMMA rejects them at validation). + new_variant = i.epilogue != "default" or i.two_stage or i.async_dma + if new_variant and not i.group_merge: + assert geo in aj.CACHE_WGRAD_WMMA_GEOMETRIES, i + assert set(aj.CACHE_WGRAD_WMMA_GEOMETRIES) < geos + assert walked == aj.count_jobs( + dtype="bf16", target=target, directions=["wgrad"], split_ks=split_ks + ) + + @pytest.mark.parametrize( "directions", [("fwd",), ("wgrad",), ("dgrad",), ("fwd", "wgrad", "dgrad")] ) diff --git a/dnn-providers/hip-kernel-provider/rocke/library/tests/test_conv_wgrad_correctness.py b/dnn-providers/hip-kernel-provider/rocke/library/tests/test_conv_wgrad_correctness.py index 32a4f5fd0c2e..9fd8d300316d 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/tests/test_conv_wgrad_correctness.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/tests/test_conv_wgrad_correctness.py @@ -831,7 +831,7 @@ def test_wide_tile_spec_at_degree_cap_gfx1250(self): at that degree rather than at the cache's ``split_k=2`` marker; the degree is a kernarg, so the IR is the same. - These are the 2x2-warp tiles, and the bf16 problems are sized so the + These are 2x2-warp tiles of the gfx1250 cover, and the bf16 problems are sized so the degree lands on its cap. The last bf16 problem (small channels, deep reduction) asks for 128 before the cap; all-positive partials miss the tolerance there. fp16 takes its own policy (fill to 1024 CTAs, 2048 @@ -861,9 +861,9 @@ def test_wide_tile_spec_at_degree_cap_gfx1250(self): # expected degree); -1 is the resolver's pick. cases = [ (32, 32, 64, 96, 1, "bf16", (32, 64, 64), -1, 12), - (32, 32, 64, 128, 1, "fp16", (64, 64, 64), -1, 56), - (32, 32, 64, 64, 4, "bf16", (32, 32, 32), -1, 12), - (32, 32, 64, 64, 4, "fp16", (32, 32, 32), -1, 64), + (32, 32, 64, 128, 1, "fp16", (64, 128, 64), -1, 64), + (32, 32, 64, 64, 4, "bf16", (32, 128, 32), -1, 12), + (32, 32, 64, 64, 4, "fp16", (32, 128, 32), -1, 64), (64, 64, 16, 16, 1, "bf16", (32, 64, 64), -1, 12), ] ran = 0 @@ -879,14 +879,17 @@ def test_wide_tile_spec_at_degree_cap_gfx1250(self): and g[5] == 16 and g[6:8] == ("mem", "default") ] - # The cache's split-K binaries for this geometry that would be - # built (spec validity) and offered for this problem (lookup). + # The cache's packed-atomic split-K binaries for this geometry + # (synchronous K loop) that would be built (spec validity) and + # offered for this problem (lookup). jobs = [ j for j in aot_jobs._direction_jobs( "wgrad", GPU_ARCH, dtype, target, (2,), 1, geometries=geometries ) - if aot_jobs._spec_is_valid(j, GPU_ARCH, dtype) + if not j.identity.two_stage + and not j.identity.async_dma + and aot_jobs._spec_is_valid(j, GPU_ARCH, dtype) and identity_supports_problem(j.identity, problem)[0] ] self.assertTrue(jobs, f"no cached split-K kernel serves {problem}") From d02b38751ef83bdc7abf8a30787d4599c2889938 Mon Sep 17 00:00:00 2001 From: Jakub Piasecki Date: Sun, 1 Nov 2026 20:23:01 +0000 Subject: [PATCH 07/13] fix(rocke): gate the wgrad async K loop on global_load_async_to_lds in the job count count_jobs counted the async_dma wgrad variants on every wave32 target, but _wgrad_jobs drops them on gfx1151 and gfx1201, which have no global_load_async_to_lds. The count then exceeded what the walk can yield. _wgrad_k_loops now takes the target and offers the async loop only where target.memory.has_async_global_lds holds; both the walk and the count use it. Test: count_jobs equals the wgrad walk on gfx1250, gfx1151 and gfx1201, and only gfx1250 yields async jobs. Co-Authored-By: Claude Opus 5.5 --- .../rocke/library/kernels/conv/aot_jobs.py | 22 ++++++++++++++----- .../rocke/library/tests/test_conv_abi.py | 22 +++++++++++++++++++ 2 files changed, 39 insertions(+), 5 deletions(-) diff --git a/dnn-providers/hip-kernel-provider/rocke/library/kernels/conv/aot_jobs.py b/dnn-providers/hip-kernel-provider/rocke/library/kernels/conv/aot_jobs.py index 67455bf1322a..8a27bd84a54c 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/kernels/conv/aot_jobs.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/kernels/conv/aot_jobs.py @@ -545,14 +545,24 @@ def _fwd_k_loops(pipeline: str) -> List[Tuple[bool, bool]]: return k_loops -def _wgrad_k_loops(wave_size: int, pipeline: str, two_stage: bool) -> List[dict]: +def _wgrad_k_loops( + wave_size: int, pipeline: str, two_stage: bool, target +) -> List[dict]: """K-loop variants built per wgrad configuration. The base loop always. On wave32 (WMMA) next to "mem", except under - two-stage: the gfx1250 ``async_dma`` loop (global_load_async_to_lds). + two-stage, and only on a target with ``global_load_async_to_lds`` + (gfx1250): the ``async_dma`` loop. :func:`_wgrad_jobs` and + :func:`count_jobs` both take the variants from here, so the count stays + the number of jobs the walk yields. """ loops: List[dict] = [{}] - if wave_size == 32 and pipeline == "mem" and not two_stage: + if ( + wave_size == 32 + and pipeline == "mem" + and not two_stage + and target.memory.has_async_global_lds + ): loops.append({"async_dma": True}) return loops @@ -1239,7 +1249,9 @@ def _wgrad_jobs( # sizes the scratch and Stage 2 from what was compiled. cfg["ws_replicas"] = _DEFAULT_WS_REPLICAS k_loops = ( - _wgrad_k_loops(wave_size, pipeline, two_stage) if in_cover else [{}] + _wgrad_k_loops(wave_size, pipeline, two_stage, target) + if in_cover + else [{}] ) for k_loop in k_loops: kcfg = dict(cfg, **k_loop) @@ -1851,7 +1863,7 @@ def count_jobs( stages = sum(len(_wgrad_two_stages(sk)) for sk in split_ks) # The async loop: once per non-two-stage configuration. asyncs = len(split_ks) * ( - len(_wgrad_k_loops(wave_size, pipeline, False)) - 1 + len(_wgrad_k_loops(wave_size, pipeline, False, target)) - 1 ) else: stages, asyncs = len(split_ks), 0 diff --git a/dnn-providers/hip-kernel-provider/rocke/library/tests/test_conv_abi.py b/dnn-providers/hip-kernel-provider/rocke/library/tests/test_conv_abi.py index 8270755f44bf..43bbdde64da2 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/tests/test_conv_abi.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/tests/test_conv_abi.py @@ -1892,6 +1892,28 @@ def test_wmma_wgrad_keeps_the_default_grid_and_covers_the_new_variants(): ) +@pytest.mark.parametrize("gfx", ["gfx1250", "gfx1151", "gfx1201"]) +def test_wmma_wgrad_count_matches_the_walk(gfx): + """:func:`aj.count_jobs` counts exactly the wgrad jobs the generator + yields on every wave32 target. The async loop is built only where the + target has ``global_load_async_to_lds`` (gfx1250); gfx1151 / gfx1201 get + none, and the count must not include them either. + """ + from rocke.core.arch import ArchTarget + + target = ArchTarget.from_gfx(gfx) + split_ks = aj.CACHE_WGRAD_SPLIT_KS + jobs = list(aj._wgrad_jobs(gfx, "fp16", 32, "wmma", target, split_ks)) + n_async = sum(1 for job in jobs if job.identity.async_dma) + if target.memory.has_async_global_lds: + assert n_async > 0 + else: + assert n_async == 0 + assert len(jobs) == aj.count_jobs( + dtype="fp16", target=target, directions=["wgrad"], split_ks=split_ks + ) + + @pytest.mark.parametrize( "directions", [("fwd",), ("wgrad",), ("dgrad",), ("fwd", "wgrad", "dgrad")] ) From 4d5e20069044ca1293b70a7936a94dcac9eb26ce Mon Sep 17 00:00:00 2001 From: Jakub Piasecki Date: Sun, 1 Nov 2026 20:24:38 +0000 Subject: [PATCH 08/13] fix(rocke): prune async-DMA cache kernels apart from synchronous ones The from-cache sweep keeps only the widest vector copy of each kernel kind. An async-DMA wgrad binary records the chunk widths its copies use, while a synchronous binary records its width caps, which its loader can narrow at build time. On the gfx1250 cover geometry 32x64x32 (2x4 warps) the async loop records A=4 / B=8 and the synchronous copy A=8 / B=8, so with both in one kind every async kernel was dropped before it was benchmarked. The kind now includes async_dma, so each loop is pruned to its own widest copy. Benchmark / run side only; no kernel IR change. Test: on that geometry the async cshuffle kernel survives the pruning and the synchronous kernels are still cut to their widest copy. Co-Authored-By: Claude Opus 5.5 --- .../common/benchmark_implicit_gemm_conv.py | 19 ++++---- .../rocke/library/tests/test_conv_abi.py | 45 +++++++++++++++++++ 2 files changed, 56 insertions(+), 8 deletions(-) diff --git a/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/common/benchmark_implicit_gemm_conv.py b/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/common/benchmark_implicit_gemm_conv.py index ac004758fb1b..1ba72624c49e 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/common/benchmark_implicit_gemm_conv.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/common/benchmark_implicit_gemm_conv.py @@ -5527,10 +5527,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, @@ -5548,14 +5551,14 @@ def kind(ident): ident.group_merge, ident.streamk, ident.streamk_reduction, + bool(getattr(ident, "async_dma", False)), ) def rank(ident): # An async-DMA binary leaves A / B unset and records the chunk widths - # its copies use instead; it ranks by those, as a synchronous kernel - # with the same load widths. (Ranking it as 8 wide let it outrank, and - # so filter out, the widest synchronous kernels of a problem whose - # channel run caps the loads below 8.) + # 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: diff --git a/dnn-providers/hip-kernel-provider/rocke/library/tests/test_conv_abi.py b/dnn-providers/hip-kernel-provider/rocke/library/tests/test_conv_abi.py index 43bbdde64da2..3d780dd0e92b 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/tests/test_conv_abi.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/tests/test_conv_abi.py @@ -1914,6 +1914,51 @@ def test_wmma_wgrad_count_matches_the_walk(gfx): ) +def test_widest_vector_pruning_keeps_the_async_wgrad_kernels(): + """The from-cache sweep's widest-vector pruning ranks async-DMA kernels + apart from synchronous ones. + + On the gfx1250 cover geometry 32x64x32 (2x4 warps) the async loop records + A=4 / B=8 chunks, while a synchronous binary records the A=8 / B=8 caps + (its loader narrows A at build time). Ranked together, every async kernel + lost to the synchronous ones and was dropped before it was benchmarked. + """ + from rocke.core.arch import ArchTarget + + from benchmarks.common.benchmark_implicit_gemm_conv import ( + _widest_vector_sizes, + ) + + target = ArchTarget.from_gfx("gfx1250") + idents = [ + job.identity + for job in aj._wgrad_jobs("gfx1250", "fp16", 32, "wmma", target, (1,)) + if ( + job.identity.tile_m, + job.identity.tile_n, + job.identity.tile_k, + job.identity.warp_m, + job.identity.warp_n, + ) + == (32, 64, 32, 2, 4) + and job.identity.grouped + ] + async_ids = [i for i in idents if i.async_dma] + assert async_ids and all( + (i.async_chunk_a, i.async_chunk_b) == (4, 8) for i in async_ids + ) + kept = [c[0] for c in _widest_vector_sizes([(i, None) for i in idents])] + kept_async = [i for i in kept if i.async_dma] + kept_sync = [i for i in kept if not i.async_dma] + # The async kernels survive (the cshuffle one; a default-epilogue kernel + # is dropped where a cshuffle one stores wider, async or not) ... + assert [i.epilogue for i in kept_async] == ["cshuffle"] + # ... and the synchronous ones are still pruned to their widest copy. + assert {(i.vector_size_a, i.vector_size_b, i.vector_size_c) for i in kept_sync} == { + (8, 8, 8) + } + + @pytest.mark.parametrize( "directions", [("fwd",), ("wgrad",), ("dgrad",), ("fwd", "wgrad", "dgrad")] ) From 6e49bdaa7257df8ae52422a194763e054434305e Mon Sep 17 00:00:00 2001 From: Jakub Piasecki Date: Sun, 1 Nov 2026 20:27:03 +0000 Subject: [PATCH 09/13] fix(rocke): cap the automatic gfx1250 fp16 wgrad split-K at the atomic's accuracy bound The gfx1250 fp16 wgrad split policy went up to degree 128. Each packed fp16 atomic add rounds the running sum in dW, so valid inputs can miss the 5e-2 tolerance at that degree: with one 1.0 partial followed by 127 partials of 2**-11 the fp16 sum stays at 1.0, a 5.8% error. The policy now caps fp16 at atomic_rounding_degree_cap(10) = 97, the same accuracy bound bf16 uses with 7 fraction bits. The cap binds only the packed-atomic reduction: the f32 workspace degree is already resolved as fp32 dW and keeps the general cap of 128. Explicit and overridden degrees are unchanged. No kernel IR or AOT job changes; an automatic fp16 dW degree on gfx1250 that would exceed 97 is now 97. Tests: the worst-case fp16 partials exceed the tolerance at 128 and stay within it at the automatic degree (the pointwise N=64, 32x32, C=K=16 shape on a 64x128x64 tile); the policy and resolver expectations move from 128 to 97, and the workspace reduction of the same kind of shape still resolves to 128. Co-Authored-By: Claude Opus 5.5 --- .../common/README_conv_implicit_gemm_wgrad.md | 10 +++-- .../rocke/library/tests/test_conv_split_k.py | 21 ++++++++++ .../platform/python/rocke/helpers/split_k.py | 18 ++++++--- .../dispatch_tests/gemm/test_split_k_wgrad.py | 40 ++++++++++++++++--- 4 files changed, 75 insertions(+), 14 deletions(-) diff --git a/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/README_conv_implicit_gemm_wgrad.md b/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/README_conv_implicit_gemm_wgrad.md index 470eb7d1be8e..daee53342dd0 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/README_conv_implicit_gemm_wgrad.md +++ b/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/README_conv_implicit_gemm_wgrad.md @@ -105,12 +105,16 @@ dW rows, and the extra memory transactions outweigh the halved atomic count. (`atomic_rounding_degree_cap`); higher degrees miss it on all-positive inputs. fp32 takes the general cap of 128. - fp16 has its own policy: fill to 1024 CTAs (2048 grouped) with K slices - down to 512, up to degree 128. Earlier fp16 was not split automatically + 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; no fp16 degree failed the 5e-2 - verification there. + 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 diff --git a/dnn-providers/hip-kernel-provider/rocke/library/tests/test_conv_split_k.py b/dnn-providers/hip-kernel-provider/rocke/library/tests/test_conv_split_k.py index 813086dd7475..7f28cdcef259 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/tests/test_conv_split_k.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/tests/test_conv_split_k.py @@ -650,6 +650,27 @@ def test_wmma_auto_without_a_reduction_runs_unsplit(dtype): resolve_conv_split_k("wgrad", p, allow_unsplit=False, **kw) +def test_wmma_fp16_atomic_degree_is_capped_for_accuracy(): + # gfx1250 fp16 dW, pointwise C=K=16 over a deep reduction (N*Ho*Wo = + # 262144) on one tile: the fill and slice terms ask for far more than the + # packed fp16 atomic's accuracy cap (97, see atomic_rounding_degree_cap), + # so the automatic atomic degree stops there. The f32 workspace reduction + # resolves as fp32 dW and keeps the general cap of 128. + p = ConvProblem(N=64, Hi=64, Wi=64, C=16, K=16, Y=1, X=1) + kw = { + "arch": "gfx1250", + "dtype": "fp16", + "tile_m": 64, + "tile_n": 128, + "tile_k": 64, + "block_size": 256, + } + atomic = resolve_conv_split_k("wgrad", p, reduction="atomic", **kw) + assert (atomic.split_k, atomic.two_stage) == (97, False) + ws = resolve_conv_split_k("wgrad", p, reduction="workspace", **kw) + assert (ws.split_k, ws.two_stage) == (128, True) + + @pytest.mark.parametrize("arch", ["gfx1151", "gfx1201"]) def test_wmma_auto_without_float_atomics_runs_unsplit(arch): # A wave32 arch whose catalog has no has_global_float_atomic_add: the dW diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/split_k.py b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/split_k.py index 3ff4c8f814ee..fb7dbf2a081e 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/split_k.py +++ b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/split_k.py @@ -147,18 +147,24 @@ def atomic_rounding_degree_cap(fraction_bits: int, tol: float = WGRAD_TOLERANCE) # Per-dtype overrides of _WGRAD_SPLIT_POLICY, keyed by (arch, dW dtype). # # gfx1250 fp16: fill to 1024 CTAs (2048 grouped) with K slices as shallow as -# 512, up to degree 128. Fitted on a sample of wgrad shapes with every cached -# split-K binary timed at every degree 2..128 and verified, once the packed -# fp16 atomic was no longer fenced (it is an atomicrmw now). fp16 was not split -# automatically before (cap 1): that cap was a symptom of the fenced atomic. -# No fp16 degree failed the 5e-2 verification on those shapes, so the packed -# fp16 atomic's rounding is not the binding limit there. +# 512. Fitted on a sample of wgrad shapes with every cached split-K binary +# timed at every degree 2..128 and verified, once the packed fp16 atomic was +# no longer fenced (it is an atomicrmw now). fp16 was not split automatically +# before (cap 1): that cap was a symptom of the fenced atomic. +# +# The degree is capped by the packed fp16 atomic's rounding, as bf16's is: fp16 +# has 10 fraction bits (unit roundoff 2**-11), so 97 same-sign partials stay +# within the 5e-2 tolerance. Beyond it valid inputs can miss it even where the +# fitted shapes did not (at 128, 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 atomic +# reduction: the f32 workspace (two-stage) degree is resolved as fp32 dW. _WGRAD_SPLIT_POLICY_BY_DTYPE: dict[tuple[str, str], WgradSplitPolicy] = { ("gfx1250", "fp16"): WgradSplitPolicy( target_ctas=1024, grouped_target_ctas=2048, min_slice_k=512, max_degree=128, + max_degree_by_dtype=(("fp16", atomic_rounding_degree_cap(10)),), ), } diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/tests/dispatch/dispatch_tests/gemm/test_split_k_wgrad.py b/dnn-providers/hip-kernel-provider/rocke/platform/tests/dispatch/dispatch_tests/gemm/test_split_k_wgrad.py index 472ecac292a3..2f9cdc1ebf1f 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/tests/dispatch/dispatch_tests/gemm/test_split_k_wgrad.py +++ b/dnn-providers/hip-kernel-provider/rocke/platform/tests/dispatch/dispatch_tests/gemm/test_split_k_wgrad.py @@ -234,15 +234,15 @@ def test_slice_depth_term_binds(self): def test_dtype_caps(self): # A 4-tile grid on a long reduction asks for far more than any cap: - # bf16 takes its accuracy cap, and fp16 and fp32 the general one. + # each 16-bit dtype takes its accuracy cap, and fp32 the general one. args = (24, 9 * 24, 42 * 240 * 320) self.assertEqual(self._pick(*args, dtype_d="bf16").split_k, 12) - self.assertEqual(self._pick(*args, dtype_d="fp16").split_k, 128) + self.assertEqual(self._pick(*args, dtype_d="fp16").split_k, 97) self.assertEqual(self._pick(*args, dtype_d="fp32").split_k, 128) def test_fp16_policy(self): # fp16 has its own gfx1250 policy: fill to 1024 CTAs (2048 grouped), - # K slices down to 512, up to 128. + # K slices down to 512, up to its accuracy cap of 97. # Fill: base grid 2 x 18 = 36 tiles -> floor(1024 / 36) = 28 (fp32 # fills to 2048 -> 56). d = self._pick(128, 9 * 128, 42 * 120 * 160, dtype_d="fp16") @@ -253,9 +253,9 @@ def test_fp16_policy(self): self.assertEqual(d.split_k, 24, d.reason) self.assertEqual(self._pick(64, 9 * 64, 12288, dtype_d="fp32").split_k, 6) # Grouped: 4 groups of a 1 x 3 tile grid -> floor(2048 / 12) = 170, - # capped at 128. + # capped at 97. d = self._pick(16, 9 * 16, 42 * 240 * 320, tile=64, groups=4, dtype_d="fp16") - self.assertEqual(d.split_k, 128, d.reason) + self.assertEqual(d.split_k, 97, d.reason) d = self._pick(64, 9 * 64, 42 * 240 * 320, tile=64, groups=8, dtype_d="fp16") self.assertEqual(d.split_k, 28, d.reason) @@ -294,6 +294,36 @@ def rel_err(degree): cap = self._pick(24, 9 * 24, 42 * 240 * 320, dtype_d="bf16").split_k self.assertLessEqual(rel_err(cap), 5e-2) + def test_fp16_cap_holds_on_worst_case_partials(self): + # Pointwise wgrad, N=64, 32x32, C=K=16 on a 64x128x64 tile: wg_M=16, + # wg_N=16, wg_K=65536, one tile. Fill and slice depth both ask for + # 128. fp16's unit roundoff is 2**-11 (10 fraction bits): one 1.0 + # partial followed by (degree - 1) partials of 2**-11 -- half an ulp at + # 1.0 each -- rounds back to 1.0 in a round-to-nearest-even fp16 sum. + # At 128 that loses 127/2048 (about 5.8%); the automatic degree must + # stay within the tolerance. + def rel_err(degree): + parts = [1.0] + [2.0**-11] * (degree - 1) + acc = 0.0 + for p in parts: + acc = struct.unpack(" Date: Sun, 1 Nov 2026 20:32:24 +0000 Subject: [PATCH 10/13] fix(rocke): unsigned range check and tile-typed zero-fill in the async LDS loader CoalescedTileLoader.load_async_to_lds makes the buffer-resource range check itself, since global_load_async_to_lds has none. Two fixes, in both engines: - The check compared the i32 byte offset with num_bytes signed. A padded reduction tail near 2 GiB can give the offset 0x80000000, which is negative as an i32 and so passed the check, issuing an out-of-bounds copy instead of a zero-fill. The compare is now unsigned (icmp ult), as the resource's is. - The C++ mirror zero-filled out-of-range chunks with an fp16 vector whatever the operands, while Python uses the loader's element type, so bf16-input async kernels diverged between the engines. The C++ loader carries no element type, so the zero now takes the destination LDS tile's (f16 or bf16; anything else is an error, as this port loads 16-bit elements). New parity config 67 (case 55 with bf16 operands) covers the bf16 zero-fill in both engines. The async wgrad goldens move with the compare: configs 55, 56, 60 and 65 and the new 67 in the differential golden (re-recording the family also adds the gfx1250 configs 49-51, which only the llvm24 section had recorded), and the async case in the representative conv_wgrad shard. Test: test_coalesced_load_async_to_lds.py checks the unsigned compare and the element type of the zero store for fp16 and bf16 tiles. Co-Authored-By: Claude Opus 5.5 --- .../common/README_conv_implicit_gemm_wgrad.md | 6 +- .../parity/conv_implicit_gemm_wgrad_emit.c | 10 +-- .../parity/conv_implicit_gemm_wgrad_emit.py | 13 ++-- .../rocke/platform/cpp/helpers/loads.cpp | 26 ++++++-- .../rocke/helper_rocke.helpers.loads.h | 5 +- .../dsl_docs/instances/convolution.md | 2 +- .../platform/python/rocke/helpers/loads.py | 7 +- .../core/test_coalesced_load_async_to_lds.py | 64 +++++++++++++++++++ .../conv_wgrad.json | 16 ++--- .../differential/golden/llvm_gfx_all.json | 57 +++++++++++------ 10 files changed, 158 insertions(+), 48 deletions(-) create mode 100644 dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_coalesced_load_async_to_lds.py diff --git a/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/README_conv_implicit_gemm_wgrad.md b/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/README_conv_implicit_gemm_wgrad.md index daee53342dd0..ed393a2e0888 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/README_conv_implicit_gemm_wgrad.md +++ b/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/README_conv_implicit_gemm_wgrad.md @@ -70,8 +70,10 @@ barrier. `async_dma` now covers it on gfx1250, mirrored in the C++ engine. - 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) into the descriptor's `valid`, - and zero-fills failing chunks with an LDS store. + 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. diff --git a/dnn-providers/hip-kernel-provider/rocke/library/tests/parity/conv_implicit_gemm_wgrad_emit.c b/dnn-providers/hip-kernel-provider/rocke/library/tests/parity/conv_implicit_gemm_wgrad_emit.c index 3ad7b81f92ce..54a035cf8e19 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/tests/parity/conv_implicit_gemm_wgrad_emit.c +++ b/dnn-providers/hip-kernel-provider/rocke/library/tests/parity/conv_implicit_gemm_wgrad_emit.c @@ -415,19 +415,21 @@ static int make_base_cfg(int idx, rocke_implicit_gemm_conv_wgrad_spec_t* spec, c case 55: case 56: case 57: + case 67: { - /* gfx1250 WMMA async K loop. Must match the Python emitter's 55-57. */ + /* gfx1250 WMMA async K loop. Must match the Python emitter's 55-57 + * and 67 (55 with bf16 operands). */ const bool big = idx == 56; if(idx == 57) spec->problem = rocke_conv_problem_default(4, 12, 12, 3, 32, 3, 3); - else if(idx == 55) + else if(idx == 55 || idx == 67) spec->problem = rocke_conv_problem_default(4, 12, 12, 24, 32, 3, 3); else spec->problem = rocke_conv_problem_default(8, 56, 56, 64, 64, 3, 3); spec->problem.pH = 1; spec->problem.pW = 1; - spec->dtype_a = "fp16"; - spec->dtype_b = "fp16"; + spec->dtype_a = idx == 67 ? "bf16" : "fp16"; + spec->dtype_b = idx == 67 ? "bf16" : "fp16"; spec->dtype_d = "fp32"; spec->tile_m = big ? 64 : 32; spec->tile_n = big ? 64 : 32; diff --git a/dnn-providers/hip-kernel-provider/rocke/library/tests/parity/conv_implicit_gemm_wgrad_emit.py b/dnn-providers/hip-kernel-provider/rocke/library/tests/parity/conv_implicit_gemm_wgrad_emit.py index 57347e7c6dc5..fcb102fdbbb5 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/tests/parity/conv_implicit_gemm_wgrad_emit.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/tests/parity/conv_implicit_gemm_wgrad_emit.py @@ -76,8 +76,11 @@ # dW over the async K loop with one scratch slab (ws_replicas=1) at C=24 # 66 -- both engines must reject: two-stage on gfx1151 (no validated float # global atomics for the scratch adds) +# 67 -- gfx1250 async_dma with bf16 operands: case 55 with bf16 A/B, so the +# out-of-range chunks zero-fill a bf16 LDS tile with bf16 zeros # (Grouped wgrad reaches the C++ engine only through stream-K, and its tile -# loader has no elem_dtype, so no case loads bf16 operands. The stream-K +# loader has no elem_dtype, so no other case loads bf16 operands: the +# async loop's copy carries no element type, only its zero-fill does. The stream-K # partition and CTA pool are kernargs, so no case varies them. async_dma is # rejected with stream-K, and so is prefetch_addr.) # @@ -91,7 +94,7 @@ # 107 -- streamk with unroll_k (must raise ValueError) # 108 -- streamk atomic reduction into fp16 dW (must raise ValueError) # (These illustrate the validator contract. The C emitter defines only cases -# 0-66, so run_diff.py stops at the shared END before reaching 100+; these +# 0-67, so run_diff.py stops at the shared END before reaching 100+; these # configs are not exercised by the differential gate.) from kernels.common.conv_implicit_gemm_wgrad import ( WgradConvSpec, @@ -715,7 +718,7 @@ def _spec(idx: int): arch, ) - if 55 <= idx <= 57: + if 55 <= idx <= 57 or idx == 67: # gfx1250 WMMA K loop: the global_load_async_to_lds ping-pong. from kernels.common._conv_implicit_gemm_common import ConvDataSpec @@ -726,11 +729,13 @@ def _spec(idx: int): 55: (p_small, "fp32", 2, False, {"async_dma": True}, "gfx1250"), 56: (p, "fp32", 4, True, {"async_dma": True}, "gfx1250"), 57: (p_c3, "fp32", 2, False, {"async_dma": True}, "gfx1250"), + 67: (p_small, "fp32", 2, False, {"async_dma": True}, "gfx1250"), }[idx] wave32 = arch == "gfx1250" + dtype_ab = "bf16" if idx == 67 else "fp16" kw = dict( problem=problem, - data=ConvDataSpec(dtype_a="fp16", dtype_b="fp16", dtype_d=dtype_d), + data=ConvDataSpec(dtype_a=dtype_ab, dtype_b=dtype_ab, dtype_d=dtype_d), tile_m=64 if big else 32, tile_n=64 if big else 32, tile_k=64 if big else 32, diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/helpers/loads.cpp b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/helpers/loads.cpp index c9b765652bb9..5f3574b74bca 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/helpers/loads.cpp +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/helpers/loads.cpp @@ -10,6 +10,8 @@ */ #include "rocke/helper_rocke.helpers.loads.h" +#include + #include "rocke/ir_internal.h" /* rocke_i_set_err */ /* ============================================================================ @@ -670,8 +672,22 @@ void rocke_coalesced_tile_loader_load_async_to_lds(rocke_ir_builder_t* b, return; } - /* Python: zero = b.zero_vec(self.elem_dtype, self.load_vec) */ - rocke_value_t* zero = rocke_b_zero_vec(b, rocke_f16(), self->load_vec); + /* Python: zero = b.zero_vec(self.elem_dtype, self.load_vec). The C loader + * carries no elem_dtype, so the zero takes the destination LDS tile's + * element type (f16 or bf16: the 16-bit elements this port loads). */ + const rocke_type_t* lds_elem + = (smem_dst->type != NULL && smem_dst->type->kind == ROCKE_TYPE_SMEM) ? smem_dst->type->elem + : NULL; + if(lds_elem == NULL + || (strcmp(lds_elem->name, "f16") != 0 && strcmp(lds_elem->name, "bf16") != 0)) + { + rocke_i_set_err(b, + ROCKE_ERR_VALUE, + "load_async_to_lds: the LDS tile must hold f16 or bf16 (got %s)", + lds_elem != NULL ? lds_elem->name : "a non-LDS value"); + return; + } + rocke_value_t* zero = rocke_b_zero_vec(b, lds_elem, self->load_vec); rocke_value_t* c_elem_bytes = rocke_b_const_i32(b, 2); rocke_ctl_staged_t slots; rocke_coalesced_tile_loader_plan_slots(b, self, tid, &slots); @@ -683,9 +699,11 @@ void rocke_coalesced_tile_loader_load_async_to_lds(rocke_ir_builder_t* b, /* Python: off_elems, valid = descriptor(b, row, col) */ rocke_value_t* valid = NULL; rocke_value_t* off_elems = descriptor(b, idx[0], idx[1], &valid, descriptor_user); - /* Python: in_bounds = b.cmp_lt(b.mul(off_elems, c_elem_bytes), num_bytes) */ + /* Python: in_bounds = b.cmp_ult(b.mul(off_elems, c_elem_bytes), num_bytes) + * -- unsigned, as the buffer resource's range check: a byte offset + * at or past 2 GiB (negative as i32) is out of range, not below. */ rocke_value_t* off_bytes = rocke_b_mul(b, off_elems, c_elem_bytes); - rocke_value_t* in_bounds = rocke_b_cmp_lt(b, off_bytes, num_bytes); + rocke_value_t* in_bounds = rocke_b_cmp_ult(b, off_bytes, num_bytes); valid = (valid == NULL) ? in_bounds : rocke_b_land(b, valid, in_bounds); rocke_if_else_t ie = rocke_b_scf_if_else(b, valid); rocke_b_region_enter(b, ie.then_region); diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/helper_rocke.helpers.loads.h b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/helper_rocke.helpers.loads.h index a1cb87ed47c8..1d207ea6da4d 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/helper_rocke.helpers.loads.h +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/helper_rocke.helpers.loads.h @@ -262,8 +262,9 @@ void rocke_coalesced_tile_loader_issue_slots(rocke_ir_builder_t* b, * mode, but each in-range chunk is one per-lane global_load_async_to_lds of * load_vec 16-bit elements from `ptr` to smem_dst[row, col]. The copy has no * buffer-resource range check, so the one the resource would make (byte offset - * below `num_bytes`) is ANDed into the descriptor's valid; a chunk failing - * either is zero-filled with an ordinary LDS store instead. The caller drains + * below `num_bytes`, compared unsigned) is ANDed into the descriptor's valid; a + * chunk failing either is zero-filled with an ordinary LDS store of the + * destination tile's element type (f16 or bf16) instead. The caller drains * the copies with s_wait_asynccnt (and the zero stores with an LDS barrier). */ void rocke_coalesced_tile_loader_load_async_to_lds(rocke_ir_builder_t* b, const rocke_coalesced_tile_loader_t* self, diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/instances/convolution.md b/dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/instances/convolution.md index e7954618af69..f0fed9bab7e1 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/instances/convolution.md +++ b/dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/instances/convolution.md @@ -463,7 +463,7 @@ M-outer default for dgrad, and matches the fp32 reference for wgrad. - `pipeline="basic"`: now uses a bounded runtime `scf.for` K-loop (load → sync → MFMA → sync per tile), identical to the `"mem"` path. No K-trip-count limit. dgrad has no `basic` branch. - `async_dma` is a swept axis on the wgrad sweep driver, not a run-level flag. Unlike `lds_k_outer` it is not deducible from `(arch, spec)`: it removes the register staging of the tile, but it also forces the K-outer row pad to 0 and coarsens the load-width ladder to the widths the intrinsic accepts, and both terms are functions of tile width and channel run, which are themselves sweep axes. -- WMMA wgrad (gfx1250): `async_dma` is a ping-pong over `global_load_async_to_lds` (`CoalescedTileLoader.load_async_to_lds`, one barrier per tile, drained by `s_wait_asynccnt`). It reuses the synchronous K-outer loaders' slot partition, needs `lds_k_outer` and a 4-, 8- or 16-byte chunk on both operands, and ANDs the buffer-range check (`byte offset < num_bytes`) into each chunk's `valid`, zero-filling failing chunks with an LDS store. `epilogue="cshuffle"` is admitted on WMMA too: both cshuffle epilogues (direct store and split-K atomic) read the WMMA fragment through `CShuffleEpilogue.from_grid_op`. Two-stage (f32 workspace) split-K runs on gfx1250 as well: the Stage 1 scratch epilogue decodes the WMMA fragment through `op.c_layout()` (`wgrad_two_stage_available` needs the arch's float global atomics on wave32). `pipeline` stays `mem`; `unroll_k`, `prefetch_addr` and `chiplet_swizzle` stay rejected on WMMA. +- WMMA wgrad (gfx1250): `async_dma` is a ping-pong over `global_load_async_to_lds` (`CoalescedTileLoader.load_async_to_lds`, one barrier per tile, drained by `s_wait_asynccnt`). It reuses the synchronous K-outer loaders' slot partition, needs `lds_k_outer` and a 4-, 8- or 16-byte chunk on both operands, and ANDs the buffer-range check (`byte offset < num_bytes`, unsigned as the resource compares it) into each chunk's `valid`, zero-filling failing chunks with an LDS store of the tile's element type. `epilogue="cshuffle"` is admitted on WMMA too: both cshuffle epilogues (direct store and split-K atomic) read the WMMA fragment through `CShuffleEpilogue.from_grid_op`. Two-stage (f32 workspace) split-K runs on gfx1250 as well: the Stage 1 scratch epilogue decodes the WMMA fragment through `op.c_layout()` (`wgrad_two_stage_available` needs the arch's float global atomics on wave32). `pipeline` stays `mem`; `unroll_k`, `prefetch_addr` and `chiplet_swizzle` stay rejected on WMMA. - No K loop is unrolled at build time any more: `async_dma` runs `SoftwarePipeline.run_ping_pong_dynamic` and `unroll_k` a hand-rolled 2x ping-pong, both over the runtime extent, so the former `_MAX_UNROLLED_K_ITERS` / `ROCKE_MAX_UNROLLED_K_ITERS` cap is gone. - wgrad `split_k`: `1` = no split (direct store); any value `> 1` builds the one split-K kernel (name tag `spk`), whose degree is the `ks_count` / `ks` kernargs chosen at launch. There is no `split_k=0` "runtime degree" encoding. - Removed: the `--lds-k-outer`, `--lds-k-pad` and `--dtype-d` CLI flags and the `ROCKE_WGRAD_LDS_K_OUTER` env override. All replaced by deduction or dropped. diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/loads.py b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/loads.py index a3e5f6e336f4..acc2594bbb73 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/loads.py +++ b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/loads.py @@ -524,8 +524,9 @@ def load_async_to_lds( async global->LDS copy of ``load_vec`` elements to ``smem_dst[row, col]`` instead of a register round trip. The copy reads through a plain global pointer with no buffer-resource range check, so the check the resource - would make -- byte offset below ``num_bytes`` -- is ANDed into the - descriptor's ``valid``. Descriptors lean on that check: a reduction + would make -- byte offset below ``num_bytes``, compared unsigned as the + resource does, so an offset at or past 2 GiB (negative as i32) is out + of range -- is ANDed into the descriptor's ``valid``. Descriptors lean on that check: a reduction tail past the tensor, or a free-axis overhang at its last element, is reported valid and only the resource turns it into zero. A chunk that fails either test is zero-filled with an ordinary LDS store. The caller @@ -552,7 +553,7 @@ def load_async_to_lds( c_elem_bytes = b.const_i32(elem_bytes) for row, col in self.plan_slots(b, tid=tid): off_elems, valid = descriptor(b, row, col) - in_bounds = b.cmp_lt(b.mul(off_elems, c_elem_bytes), num_bytes) + in_bounds = b.cmp_ult(b.mul(off_elems, c_elem_bytes), num_bytes) valid = in_bounds if valid is None else b.land(valid, in_bounds) with b.scf_if_else(valid) as (in_range, out_of_range): with in_range: diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_coalesced_load_async_to_lds.py b/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_coalesced_load_async_to_lds.py new file mode 100644 index 000000000000..8422507d3bf9 --- /dev/null +++ b/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_coalesced_load_async_to_lds.py @@ -0,0 +1,64 @@ +# Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +# SPDX-License-Identifier: MIT + +"""LLVM-text checks for ``CoalescedTileLoader.load_async_to_lds`` (gfx1250). + +The async global->LDS copy has no buffer-resource range check, so the loader +makes the one the resource would make itself -- byte offset below +``num_bytes`` -- and zero-fills a failing chunk with an LDS store. + +* The compare must be unsigned, as the resource's is: a byte offset at or past + 2 GiB is negative as an i32, and a signed compare would read it as in range + and issue an out-of-bounds copy instead of zero-filling. +* The zero-fill must have the LDS tile's element type (bf16 zeros into a bf16 + tile), as the C++ engine's mirror does. +""" + +from __future__ import annotations + +import re + +import pytest + +from rocke.core.ir import BF16, F16, I32, IRBuilder, PtrType +from rocke.core.lower_llvm import _lower_kernel_to_llvm_python +from rocke.helpers.loads import CoalescedTileLoader + + +def _build(dtype): + b = IRBuilder("load_async_to_lds_check") + src = b.param("src", PtrType(dtype, "global"), readonly=True, align=16) + num_bytes = b.param("num_bytes", I32) + tile = b.smem_alloc(dtype, [16, 32], name_hint="tile") + loader = CoalescedTileLoader.from_tile( + tile_rows=16, tile_cols=32, block_size=64, elem_dtype=dtype + ) + c32 = b.const_i32(32) + + def descriptor(b, row, col): + return b.add(b.mul(row, c32), col), None + + loader.load_async_to_lds( + b, + tid=b.thread_id_x(), + smem_dst=tile, + descriptor=descriptor, + ptr=src, + num_bytes=num_bytes, + ) + b.ret() + return b.kernel, loader.load_vec + + +@pytest.mark.parametrize("dtype, llvm_elem", [(F16, "half"), (BF16, "bfloat")]) +def test_range_check_is_unsigned_and_zero_fill_matches_the_tile(dtype, llvm_elem): + kernel, vec = _build(dtype) + ll = _lower_kernel_to_llvm_python(kernel, arch="gfx1250", llvm_flavor="llvm23") + assert "global.load.async.to.lds" in ll + # The range check: unsigned against num_bytes, never signed. + assert re.search(r"icmp ult i32 %[\w.]+, %num_bytes", ll), ll + assert not re.search(r"icmp s(lt|le|gt|ge) i32 %[\w.]+, %num_bytes", ll), ll + # The out-of-range chunks store zeros of the tile's element type. + assert re.search(rf"store <{vec} x {llvm_elem}> %cz", ll), ll + other = "bfloat" if llvm_elem == "half" else "half" + assert f"x {other}>" not in ll diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/tests/golden/rocke_representative_ir_sha256/conv_wgrad.json b/dnn-providers/hip-kernel-provider/rocke/platform/tests/golden/rocke_representative_ir_sha256/conv_wgrad.json index 3a536167ad9f..33e2aa3664bc 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/tests/golden/rocke_representative_ir_sha256/conv_wgrad.json +++ b/dnn-providers/hip-kernel-provider/rocke/platform/tests/golden/rocke_representative_ir_sha256/conv_wgrad.json @@ -29,10 +29,10 @@ }, "conv_wgrad/gfx1250/n1h8c16k32r3_spk4_async": { "arch": "gfx1250", - "bytes": 85370, + "bytes": 85394, "family": "conv_wgrad", "kernel_name": "irhash_wgrad_1250_spk4_async_N1H8W8C16_K32Y3X3_t32x32x32_w2x1_a16x16x32_mem_default_async_kouter_spk", - "sha256": "94bb64f47a0c3398035207801b87b35e5ebd85895e382892f495f3c4302dfb31", + "sha256": "c5a89a3a45b00c66cc832e9d10b557cf08f2fa94e0020a9fe46b4051e23ab7af", "status": "ok" }, "conv_wgrad/gfx1250/n1h8c16k32r3_spk4_cshuffle": { @@ -226,10 +226,10 @@ }, "conv_wgrad/gfx1250/n1h8c16k32r3_spk4_async": { "arch": "gfx1250", - "bytes": 85509, + "bytes": 85533, "family": "conv_wgrad", "kernel_name": "irhash_wgrad_1250_spk4_async_N1H8W8C16_K32Y3X3_t32x32x32_w2x1_a16x16x32_mem_default_async_kouter_spk", - "sha256": "e7a41f7a1884f347467739385ef4979ad5f90a94fa99387f6ef07f4b6bbd1ceb", + "sha256": "46c56b4b7b427f0f8873667d44b2f79b14127e4a3277b5f7e649fdba62a34eb5", "status": "ok" }, "conv_wgrad/gfx1250/n1h8c16k32r3_spk4_cshuffle": { @@ -423,10 +423,10 @@ }, "conv_wgrad/gfx1250/n1h8c16k32r3_spk4_async": { "arch": "gfx1250", - "bytes": 85573, + "bytes": 85597, "family": "conv_wgrad", "kernel_name": "irhash_wgrad_1250_spk4_async_N1H8W8C16_K32Y3X3_t32x32x32_w2x1_a16x16x32_mem_default_async_kouter_spk", - "sha256": "00734689dde7756b661433f5a481b149a5b505d7418d156853288ea0f317c8a8", + "sha256": "c8f08c9f12ded6e37f68be6f4691edf4b792edc15430fe836b1a3d7cbcbd7e09", "status": "ok" }, "conv_wgrad/gfx1250/n1h8c16k32r3_spk4_cshuffle": { @@ -620,10 +620,10 @@ }, "conv_wgrad/gfx1250/n1h8c16k32r3_spk4_async": { "arch": "gfx1250", - "bytes": 85573, + "bytes": 85597, "family": "conv_wgrad", "kernel_name": "irhash_wgrad_1250_spk4_async_N1H8W8C16_K32Y3X3_t32x32x32_w2x1_a16x16x32_mem_default_async_kouter_spk", - "sha256": "00734689dde7756b661433f5a481b149a5b505d7418d156853288ea0f317c8a8", + "sha256": "c8f08c9f12ded6e37f68be6f4691edf4b792edc15430fe836b1a3d7cbcbd7e09", "status": "ok" }, "conv_wgrad/gfx1250/n1h8c16k32r3_spk4_cshuffle": { diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/tests/instances/differential/golden/llvm_gfx_all.json b/dnn-providers/hip-kernel-provider/rocke/platform/tests/instances/differential/golden/llvm_gfx_all.json index 4d86bc6b817d..596b6e5e0968 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/tests/instances/differential/golden/llvm_gfx_all.json +++ b/dnn-providers/hip-kernel-provider/rocke/platform/tests/instances/differential/golden/llvm_gfx_all.json @@ -561,17 +561,21 @@ "43": "fd29174e03ef906b5db50f1f6c1973d22f365d40181cca476bf89801b506511b", "44": "e947ae1c080472d24798cab77fdb22be23224485379cb15e1abdcd9f8672d679", "48": "1a8c8001b691eba5e091506a1410b22363bf9fd48010921a1042d715f5b2779c", + "49": "34d805064b0901de42cdf21868aacb5458d131c77af3af7bd25b722625ce80c9", "5": "a8044100dfbe89a48ecf9af672bc80353812d0bfca61eaa7928bd0ef88f077f9", + "50": "1b96b8376012e6cecfa067c2f92ed163509c29beb4bd93a49b47b015accfdcff", + "51": "5596229f42429f0f3f3dffaa880d7eafda8b1ed858c7ebf37514d75b9738f371", "53": "51f6686426cc43a0130ddde2511b0aedc5e5636607b9ea523744998770ea0664", - "55": "e778edf08e2ed9999a24b747652c4d0fd3a3ea64dd68d922c82abb206bfe98e3", - "56": "87bff7628af8465e95e1dd9ffa12b05f5089c838a00ab986db749c777406eee5", + "55": "87b996fcaf8249d43d1883c24ddfd9f9fc2092d7efe5d94534a38a34fac9dc39", + "56": "ff5279193d763e5e68374e5fa9508afcef163e873b9fc6eed1d9430bb0e2bdd6", "58": "7fe2785d761f330bdc80afd6156e79f426ee197f61ec3bed99d2107c04fa90bd", "59": "5ebbd1600763884bb810b302df64677a5fef6f222bfec3fac9c470a4ad20ed5d", "6": "779030de0f57eb186f48efc6b92827294b7aa1095c2a8ee96d2edae002cbaa1d", - "60": "1fda68fce4ddbfd4d1ab27127e91474868de275962bdb26e3be025ca21cb2d85", + "60": "deee12c3d967a97148a661df5ecff241e5b51a113f1cc1caea39e923a7a134a8", "63": "297b72544a13aa0b519b4f248a4ce655d39e493e730ae3f25adf7429b535bcb8", "64": "0ea52c6b7172c309a13d3aadfff7f232f2f31f251b90d80ebdc9c4e80754b1f8", - "65": "d9e8f970e68f7d3c95b6d0edad0601481f50ca1f057b774ee8d40fa782cb58a4", + "65": "d0dee69c6ef1f69eeaf9250d4c0044cf4fd9713d9ee65c6e916f48453170363d", + "67": "f9bc06dc3bcc5f12696b8742d2bd64abbc170b256b89893706d7c31153b69649", "7": "1106d4ce71314f3c532427fed8ceb00855f1cf1dd7f5117b89234cf326583087", "8": "82d8923b31be6a07c4371fbe40b33a2ea7452027299ce0993e903cf55ad092d8" }, @@ -2437,17 +2441,21 @@ "43": "ba9dee6bc90d45f078ad1a0c422a9d8b170516560e92948b23731b3b2c007ddd", "44": "8835b927ea9575fe7c32635a5c04b0cb6524ec242821f44f8e93fcb5a40cfebd", "48": "ee78fea15745bbe7cd6f54557c1c0e94b86c907758e13a1167881df99ae92b6f", + "49": "c40e91cd2bb20d17c22e5861d257c979d626212510b12a0ed09361047c9c87c3", "5": "9d32ae9f91e43f62e072658dba4bd2ae5f3dc07467ef6e9e83a1f067f1a638db", + "50": "72a1fc459286d0cf8fc21b80359991c46fd2eb073faadcb1e5bbf4466d95644c", + "51": "8f2b55e58aca5f412ae2b2eb2892cfa75694be9fb00bcd14d7f3fd6ce23aee01", "53": "2361cf65c6a7b3277143bc06f14ec547e626d361d6447d5247d29908321da28b", - "55": "845dd21a850297763ad037d18dcfdb89bb403142039d77d26d4b2d337d761012", - "56": "d22b27471808ec88c11cd416b5f1c0214b395332c65514ab2fc2e7c05499b42d", + "55": "702222a1502c1836332e0363eac64b36b3df63627c51bc5ad2722cd6fdb35365", + "56": "81cae308a2f57f704cc0247d4d6cd9b973f60cf55aab0dc27bc9213b5d82cfde", "58": "0fd30468be138821abaf30432c005cc82999425fff3020f0746ff78c199f1a68", "59": "46dd718d4f966e5b0659ebe0f8b2f6fde2ca6eac6035fd247c0b7f44df648408", "6": "b2605886350eb7b793616938e9179ad480084552a3533959f34918085a34484c", - "60": "55e677df75468f06ced7b24d600f48aa934f93a5c46963966d7a438383e05c59", + "60": "cea4005d8a76d071991aa4723f478dfba12c3f5fa20f85ce2abd51af43cbb125", "63": "595ccf6f304c516b4bbfa25ab1b0714069494ed913a96de267265f720e941d3f", "64": "bb76cde6ee2149851980c297b858c581cd115cef526902134f613937aa0f65f8", - "65": "11930bd06e3c6f3b486d7d2bfa4b409562f82a6af038bfea5eb8539eda991085", + "65": "a79cd079bd3671c921a244740b4a3e94594e600c4a84d59075ad2346e56e6917", + "67": "2e0365dba91768754f0fe4fda140370b218e3c65b3385d869fe9e2955a3070d5", "7": "d653088ddce4d716c3b9e20fabe5592541b52ba5e800786a86b660763c25a596", "8": "889d9e3617a8eb2e30d29db4ef8da0aeac52c9fa0d92bc53541947cf85926c08" }, @@ -4215,17 +4223,21 @@ "43": "e3ec7cd29bfa9035d7eeb20353e261d0926b0b3359f547346013b4ce595c9c20", "44": "bb0541f7a5814241e9998ccba12441e518d93210881fb4d78faa34e99a7414c1", "48": "529267374fd88af7746495caff9f9fb19c8e3bdac84cab09af9e946c5e429b9b", + "49": "cfd83e34c301ff1deaf627b0b485300a196d1465b890efc6463e91f139ad176d", "5": "cb1aab374b1c257ee44efdb653a51cedfef95f2dbc27c959a66215db38e0db35", + "50": "c00a088e299bc0f5c525ce304a6f15f3cd98e309ac2e772f7eec444e838df8e1", + "51": "9efcd7edbefecc214015655e5c0f776d9510dc76462f4746db19fc082779f1c4", "53": "335c1050cc3ecc5c2288292ca4fa9b735f3702bec64305e7e7b1a764a6301031", - "55": "d4f321c49e9be50002075621fef0fd6e1c5d3085adc584be4a3a317582eaaf79", - "56": "6a7c3fec1ae26dd7faf162315dcbd9fdd92becdd4704ef81dbc1ed117444e5ca", + "55": "139f26e559d8e16bec9ac03f81fe3edcdcfd840af10dd5b2e643489eb8e9322d", + "56": "b4914cef56dee2a37d9dd16e1adf62ed47fda683a83bb31319c996c7c340d86c", "58": "95aabfecf90f2021a781953eed801d3b14d86c4c5844c365dc6edd83416a0512", "59": "654c0bfc412129aad45be434cb4e382fab2739ad94ffe5d607db6924e5f77b00", "6": "f3b4e8cf7ab611637137babc7994172995f76262274b0a63ca3b6999ac62ef04", - "60": "cafe652716f3f73fc0747ca5a6932007dcc21baa840e1367de312a93bbfdf4b1", + "60": "5600f2de386d5b75ac061f5e1baa8de0a41f8690edd193e1f35e50a25a8a3d61", "63": "961ae6b9afbfcea5088451d9d356ee9d21cf208c39c0acd95ef29d3e0bca1e08", "64": "1fc144e9617d38a78751c4452ebbe78aa0bcc06fd4a5a3e0f90eb8cf0213b268", - "65": "6bc8acde7c696b4ec316a439952e4425364c55f936802d77798494025fcf43d2", + "65": "379d276960a73e2da5b2c995088d3cbbd18f1bde2c47feeddd93602910d9418d", + "67": "881028a810289a2e95263f6e9317588742868c19826dc9d9ac49606a71fdc5c2", "7": "e12021bd6f8923472a7595394ecde03a55dcd07d63d714b9adaea3efa79df887", "8": "2cd90bffada0eaec9a90caf5d5b76f2f213e4e884e3f413b13e79eef3071ec82" }, @@ -6023,17 +6035,21 @@ "43": "dee7bac88cf82e6a5baf10626a042f3b4bf7098d54149761936d92fa642d5a10", "44": "28b3e5675e42d727cfcc5c4706e5b8d432f6d635ae68370960063468c5777683", "48": "5ea399799a9a1af9e69a9dae94c3194e66f277472b290592869ba2deae88acab", + "49": "b1f8b4a6d3beaeb39c5be3594188f313cfec564e6f3893eb29a9271b8e945bb8", "5": "0826f9183484b30541afcda4322fac4e5fec28374efb8995a373b9d227a97016", + "50": "b54d4d875110a0dbd2b0c42b5844114122cd9f74b075e03ada23422cc62c185f", + "51": "b367d89061b3bfb68f209d667d122cb3aa406996af41f508c91900415c24b049", "53": "bc38f32697b74c7208cd101293556b123c54cfe96b4469fc0562c23052b10fcb", - "55": "0c59e7ef5797f03be197533617fcb598f79e7df6759ca8fb3bf875a53e395731", - "56": "acf39c44252e4063825560e706c92070ec2532daca3f1072453e4b9bb747681a", + "55": "9c861cea11bff6bc2d13c8d40c5c45c6d5516db96a65856ea84301a655ce02ec", + "56": "de4fff1feda2f0462815d6fcfbe310d16c647999534e1c9a20bcf12e5c482501", "58": "02b4d2adce33a29c4a426bc3c4617387e340a2e19b429e986468c6064dac2759", "59": "fbcf10500498604862078885c4482d5cbe67269770db5929f39fa0d6919d3c89", "6": "2726776342b4ffd858fce5951d771a50515d7c143e07656eb6235e10fd08bd32", - "60": "33ef9a12ea16c82a1e985e2ad05b557de385f4def104b8b4abd694e6989ff42e", + "60": "acc8c8d0abce808bcd9359d4b8882c07907a651b7becfc134e9e82d924dcee26", "63": "9b06d34d2d4381e75b0d249249e90cea4b09d0a3409328dd835eeb256546faae", "64": "861ee261e315bb0202240ac717921ef9fb69b24ea6949194d5190778261d4706", - "65": "a664126829bd3246ff2d5a4db943d691e6c49d04c65fec57a78fee8218289608", + "65": "88e233291d3674ce2dd072faaed3eb546ea10d8e036f741a06c76531d19570a4", + "67": "79f10f6395545d6b2356fc4d1d4d02e6beb39a4ffe92f5d4060dc8abe0f2e67a", "7": "44a58b32e3d95d8492b23193d701925de3dd9de8cddf8a2c006b6fdf5de87cf8", "8": "af0df509c9e1bfa9bcb32621a677b633c9bd7c0b0a22012ee20acd2894b44df8" }, @@ -8088,15 +8104,16 @@ "50": "b54d4d875110a0dbd2b0c42b5844114122cd9f74b075e03ada23422cc62c185f", "51": "b367d89061b3bfb68f209d667d122cb3aa406996af41f508c91900415c24b049", "53": "bc38f32697b74c7208cd101293556b123c54cfe96b4469fc0562c23052b10fcb", - "55": "0c59e7ef5797f03be197533617fcb598f79e7df6759ca8fb3bf875a53e395731", - "56": "acf39c44252e4063825560e706c92070ec2532daca3f1072453e4b9bb747681a", + "55": "9c861cea11bff6bc2d13c8d40c5c45c6d5516db96a65856ea84301a655ce02ec", + "56": "de4fff1feda2f0462815d6fcfbe310d16c647999534e1c9a20bcf12e5c482501", "58": "02b4d2adce33a29c4a426bc3c4617387e340a2e19b429e986468c6064dac2759", "59": "fbcf10500498604862078885c4482d5cbe67269770db5929f39fa0d6919d3c89", "6": "2726776342b4ffd858fce5951d771a50515d7c143e07656eb6235e10fd08bd32", - "60": "33ef9a12ea16c82a1e985e2ad05b557de385f4def104b8b4abd694e6989ff42e", + "60": "acc8c8d0abce808bcd9359d4b8882c07907a651b7becfc134e9e82d924dcee26", "63": "9b06d34d2d4381e75b0d249249e90cea4b09d0a3409328dd835eeb256546faae", "64": "861ee261e315bb0202240ac717921ef9fb69b24ea6949194d5190778261d4706", - "65": "a664126829bd3246ff2d5a4db943d691e6c49d04c65fec57a78fee8218289608", + "65": "88e233291d3674ce2dd072faaed3eb546ea10d8e036f741a06c76531d19570a4", + "67": "79f10f6395545d6b2356fc4d1d4d02e6beb39a4ffe92f5d4060dc8abe0f2e67a", "7": "44a58b32e3d95d8492b23193d701925de3dd9de8cddf8a2c006b6fdf5de87cf8", "8": "af0df509c9e1bfa9bcb32621a677b633c9bd7c0b0a22012ee20acd2894b44df8" }, From 480216eb5854974c7f0624a0b057452708d1ced9 Mon Sep 17 00:00:00 2001 From: Jakub Piasecki Date: Tue, 3 Nov 2026 22:39:20 +0000 Subject: [PATCH 11/13] fix(rocke): reject explicit wgrad A/B load-width caps on the async K loop in C++ Python is_valid_wgrad_spec rejects an explicit vector_size_a/b when async_dma is set, because the async K loop derives its chunk widths from the channel counts alone, and range- and divisibility-checks the caps otherwise. The C++ validator had none of these checks, so a C spec with a cap on the async path passed validation and the builder ignored it. Port the Python loop into rocke_implicit_gemm_conv_wgrad_is_valid_spec at the same point in the gate order, with the same reason strings: reject a cap under async_dma, then require 1 <= cap <= 8 (4 for fp32 A) and that the cap divides the operand's stride-1 channel run (kpg for A, cpg for B). Tests: rocke_wgrad_streamk_host pins the Python reasons for the async rejection (A and B), the range and divisibility rejections, and the specs both engines accept; parity config 68 (config 55 with vector_size_a=4) is a both-engines-reject case. Co-Authored-By: Claude Opus 5.5 --- .../parity/conv_implicit_gemm_wgrad_emit.c | 13 ++- .../parity/conv_implicit_gemm_wgrad_emit.py | 15 ++- .../common/conv_implicit_gemm_wgrad.cpp | 66 +++++++++++++ .../tests/instances/wgrad_streamk_host.cpp | 98 +++++++++++++++++++ 4 files changed, 187 insertions(+), 5 deletions(-) diff --git a/dnn-providers/hip-kernel-provider/rocke/library/tests/parity/conv_implicit_gemm_wgrad_emit.c b/dnn-providers/hip-kernel-provider/rocke/library/tests/parity/conv_implicit_gemm_wgrad_emit.c index 54a035cf8e19..9c44f7169afd 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/tests/parity/conv_implicit_gemm_wgrad_emit.c +++ b/dnn-providers/hip-kernel-provider/rocke/library/tests/parity/conv_implicit_gemm_wgrad_emit.c @@ -416,13 +416,15 @@ static int make_base_cfg(int idx, rocke_implicit_gemm_conv_wgrad_spec_t* spec, c case 56: case 57: case 67: + case 68: { - /* gfx1250 WMMA async K loop. Must match the Python emitter's 55-57 - * and 67 (55 with bf16 operands). */ + /* gfx1250 WMMA async K loop. Must match the Python emitter's 55-57, + * 67 (55 with bf16 operands) and 68 (55 with an explicit vector_size_a, + * which both engines must reject). */ const bool big = idx == 56; if(idx == 57) spec->problem = rocke_conv_problem_default(4, 12, 12, 3, 32, 3, 3); - else if(idx == 55 || idx == 67) + else if(idx == 55 || idx == 67 || idx == 68) spec->problem = rocke_conv_problem_default(4, 12, 12, 24, 32, 3, 3); else spec->problem = rocke_conv_problem_default(8, 56, 56, 64, 64, 3, 3); @@ -445,6 +447,11 @@ static int make_base_cfg(int idx, rocke_implicit_gemm_conv_wgrad_spec_t* spec, c spec->split_k = big ? 4 : 2; spec->lds_k_outer = true; spec->async_dma = true; + if(idx == 68) + { + spec->has_vector_size_a = true; + spec->vector_size_a = 4; + } *arch = "gfx1250"; return 0; } diff --git a/dnn-providers/hip-kernel-provider/rocke/library/tests/parity/conv_implicit_gemm_wgrad_emit.py b/dnn-providers/hip-kernel-provider/rocke/library/tests/parity/conv_implicit_gemm_wgrad_emit.py index fcb102fdbbb5..39bd1894b2ff 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/tests/parity/conv_implicit_gemm_wgrad_emit.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/tests/parity/conv_implicit_gemm_wgrad_emit.py @@ -78,6 +78,9 @@ # global atomics for the scratch adds) # 67 -- gfx1250 async_dma with bf16 operands: case 55 with bf16 A/B, so the # out-of-range chunks zero-fill a bf16 LDS tile with bf16 zeros +# 68 -- both engines must reject: case 55 with an explicit vector_size_a +# (the async K loop derives its chunk widths from the channel counts, +# so an explicit A/B load-width cap is rejected rather than ignored) # (Grouped wgrad reaches the C++ engine only through stream-K, and its tile # loader has no elem_dtype, so no other case loads bf16 operands: the # async loop's copy carries no element type, only its zero-fill does. The stream-K @@ -94,7 +97,7 @@ # 107 -- streamk with unroll_k (must raise ValueError) # 108 -- streamk atomic reduction into fp16 dW (must raise ValueError) # (These illustrate the validator contract. The C emitter defines only cases -# 0-67, so run_diff.py stops at the shared END before reaching 100+; these +# 0-68, so run_diff.py stops at the shared END before reaching 100+; these # configs are not exercised by the differential gate.) from kernels.common.conv_implicit_gemm_wgrad import ( WgradConvSpec, @@ -718,7 +721,7 @@ def _spec(idx: int): arch, ) - if 55 <= idx <= 57 or idx == 67: + if 55 <= idx <= 57 or idx in (67, 68): # gfx1250 WMMA K loop: the global_load_async_to_lds ping-pong. from kernels.common._conv_implicit_gemm_common import ConvDataSpec @@ -730,6 +733,14 @@ def _spec(idx: int): 56: (p, "fp32", 4, True, {"async_dma": True}, "gfx1250"), 57: (p_c3, "fp32", 2, False, {"async_dma": True}, "gfx1250"), 67: (p_small, "fp32", 2, False, {"async_dma": True}, "gfx1250"), + 68: ( + p_small, + "fp32", + 2, + False, + {"async_dma": True, "vector_size_a": 4}, + "gfx1250", + ), }[idx] wave32 = arch == "gfx1250" dtype_ab = "bf16" if idx == 67 else "fp16" diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/common/conv_implicit_gemm_wgrad.cpp b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/common/conv_implicit_gemm_wgrad.cpp index ff5202676027..3041e560a3d2 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/common/conv_implicit_gemm_wgrad.cpp +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/common/conv_implicit_gemm_wgrad.cpp @@ -1050,6 +1050,72 @@ bool rocke_implicit_gemm_conv_wgrad_is_valid_spec(const rocke_implicit_gemm_conv return false; } + /* The A/B load widths are a cap on the free-axis auto-selection, so an + * inadmissible request is rejected rather than silently downgraded. The + * async K loop derives its chunk widths from the channel counts alone + * (wgrad_wmma_async_available and the gfx9 async loaders), so an explicit + * cap there is rejected rather than ignored. Mirrors the vector_size_a/b + * loop of Python is_valid_wgrad_spec, in its order (the range uses + * dtype_a for both operands, as Python does). */ + { + const char* da = s->dtype_a ? s->dtype_a : "fp16"; + const int vec_cap = strcmp(da, "fp32") == 0 ? 4 : 8; + const struct + { + const char* name; + bool has; + int v; + int chan; + } caps[2] = { + {"vector_size_a", + s->has_vector_size_a, + s->vector_size_a, + rocke_conv_problem_kpg(&s->problem)}, + {"vector_size_b", + s->has_vector_size_b, + s->vector_size_b, + rocke_conv_problem_cpg(&s->problem)}, + }; + for(const auto& c : caps) + { + if(!c.has) + continue; + if(s->async_dma) + { + if(reason && reason_cap) + snprintf(reason, + reason_cap, + "%s is not honoured on the async_dma path (the direct " + "global->LDS intrinsic derives its own width); leave it None", + c.name); + return false; + } + if(c.v < 1 || c.v > vec_cap) + { + if(reason && reason_cap) + snprintf(reason, + reason_cap, + "%s=%d out of range 1..%d for this dtype", + c.name, + c.v, + vec_cap); + return false; + } + if(c.chan % c.v) + { + if(reason && reason_cap) + snprintf(reason, + reason_cap, + "%s=%d must divide the stride-1 channel run (%d); a wider " + "vector would read across the contiguous boundary", + c.name, + c.v, + c.chan); + return false; + } + } + } + if(s->pipeline && strcmp(s->pipeline, "basic") == 0 && s->async_dma) { if(reason && reason_cap) diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/tests/instances/wgrad_streamk_host.cpp b/dnn-providers/hip-kernel-provider/rocke/platform/tests/instances/wgrad_streamk_host.cpp index 4504ba05bb72..b49aebca182a 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/tests/instances/wgrad_streamk_host.cpp +++ b/dnn-providers/hip-kernel-provider/rocke/platform/tests/instances/wgrad_streamk_host.cpp @@ -453,6 +453,103 @@ void check_wmma_split_k_validator() CHECK(strstr(why, "store_vec=1") != NULL); } +/* Explicit A/B load-width caps (Python is_valid_wgrad_spec's vector_size_a/b + * loop): rejected on the async K loop, whose chunk widths come from the + * channel counts alone, and range/divisibility-checked elsewhere. Reasons + * verbatim from Python. */ +void check_vector_cap_validator() +{ + char why[512]; + auto valid = [&](const rocke_implicit_gemm_conv_wgrad_spec_t& s, const char* arch) { + why[0] = '\0'; + return rocke_implicit_gemm_conv_wgrad_is_valid_spec(&s, arch, why, sizeof(why)); + }; + auto rejects = + [&](const rocke_implicit_gemm_conv_wgrad_spec_t& s, const char* arch, const char* reason) { + const bool ok = valid(s, arch); + if(ok || strcmp(why, reason) != 0) + { + fprintf(stderr, " want: %s\n got: %s\n", reason, ok ? "(accepted)" : why); + return false; + } + return true; + }; + /* The parity emitter's case 55: gfx1250 async K loop, C=24, fp32 dW. */ + auto async_spec = []() { + rocke_conv_problem_t p = rocke_conv_problem_default(4, 12, 12, 24, 32, 3, 3); + p.pH = 1; + p.pW = 1; + rocke_implicit_gemm_conv_wgrad_spec_t s = rocke_implicit_gemm_conv_wgrad_spec_default(); + s.problem = p; + s.dtype_a = "fp16"; + s.dtype_b = "fp16"; + s.dtype_d = "fp32"; + s.tile_m = 32; + s.tile_n = 32; + s.tile_k = 32; + s.warp_m = 1; + s.warp_n = 1; + s.warp_tile_m = 16; + s.warp_tile_n = 16; + s.warp_tile_k = 32; + s.wave_size = 32; + s.split_k = 2; + s.lds_k_outer = true; + s.async_dma = true; + return s; + }; + + rocke_implicit_gemm_conv_wgrad_spec_t s = async_spec(); + CHECK(valid(s, "gfx1250")); + s.has_vector_size_a = true; + s.vector_size_a = 4; + CHECK(rejects(s, + "gfx1250", + "vector_size_a is not honoured on the async_dma path (the direct " + "global->LDS intrinsic derives its own width); leave it None")); + /* Even a cap the auto width already meets is rejected, as in Python. */ + s.vector_size_a = 1; + CHECK(rejects(s, + "gfx1250", + "vector_size_a is not honoured on the async_dma path (the direct " + "global->LDS intrinsic derives its own width); leave it None")); + s = async_spec(); + s.has_vector_size_b = true; + s.vector_size_b = 8; + CHECK(rejects(s, + "gfx1250", + "vector_size_b is not honoured on the async_dma path (the direct " + "global->LDS intrinsic derives its own width); leave it None")); + /* The synchronous K-outer loop takes the caps. */ + s = async_spec(); + s.async_dma = false; + s.has_vector_size_a = true; + s.vector_size_a = 4; + s.has_vector_size_b = true; + s.vector_size_b = 2; + CHECK(valid(s, "gfx1250")); + + /* MFMA, no async: range and divisibility of the stride-1 channel run. */ + s = spec(dense()); + s.streamk = "off"; + CHECK(valid(s, "gfx950")); + s.has_vector_size_a = true; + s.vector_size_a = 4; + CHECK(valid(s, "gfx950")); + s.vector_size_a = 0; + CHECK(rejects(s, "gfx950", "vector_size_a=0 out of range 1..8 for this dtype")); + s.vector_size_a = -1; + CHECK(rejects(s, "gfx950", "vector_size_a=-1 out of range 1..8 for this dtype")); + s = spec(dense()); + s.streamk = "off"; + s.has_vector_size_b = true; + s.vector_size_b = 3; + CHECK(rejects(s, + "gfx950", + "vector_size_b=3 must divide the stride-1 channel run (32); a wider " + "vector would read across the contiguous boundary")); +} + } // namespace int main() @@ -463,6 +560,7 @@ int main() check_invalid_inputs(); check_validator(); check_wmma_split_k_validator(); + check_vector_cap_validator(); if(g_failures) { fprintf(stderr, "rocke_wgrad_streamk_host: %d failure(s)\n", g_failures); From 8ac37b6bc87592bfee7699112eff89235ce621c2 Mon Sep 17 00:00:00 2001 From: Jakub Piasecki Date: Wed, 4 Nov 2026 13:16:01 +0000 Subject: [PATCH 12/13] fix(rocke): apply explicit wgrad A/B load-width caps in the C++ builder Python caps the synchronous free-axis load widths of the implicit-GEMM wgrad with vector_size_a / vector_size_b (_wgrad_free_axis_vec takes the smaller of the auto width and the cap). The C++ builder computed the auto width inline and ignored both fields, so a capped spec built different IR in the two engines, and the pybind binding dropped both fields, so rocke_engine neither built the capped kernel nor ran the cap checks the C++ validator already has. - wgrad_free_axis_vec takes the cap, as Python does, and the synchronous loader setup (M-outer and K-outer) uses it in place of its inline copy. The async availability check still passes no cap, as _wmma_async_vecs does; both validators reject a cap there. - conv_wgrad_build_spec reads vector_size_a and vector_size_b. Tests: parity configs 69-70 (config 55 on the synchronous K loop with caps 4/2, K-outer and M-outer); both fail against the previous C++ builder. A new gfx1250 K-outer harness case with the same caps. test_conv_wgrad_streamk checks through the binding that the capped kernels lower byte-identically and that the validator reasons agree on the async, range and divisibility rejects. Golden: new keys only. Co-Authored-By: Claude Opus 5.5 --- .../parity/conv_implicit_gemm_wgrad_emit.c | 20 +++-- .../parity/conv_implicit_gemm_wgrad_emit.py | 24 +++++- .../library/tests/test_conv_wgrad_streamk.py | 77 +++++++++++++++++++ .../platform/cpp/bindings/rocke_engine.cpp | 16 +++- .../common/conv_implicit_gemm_wgrad.cpp | 52 ++++++------- .../conv_wgrad.json | 32 ++++++++ .../differential/golden/llvm_gfx_all.json | 10 +++ .../instances/rocke_ir_parity_harness.py | 24 ++++++ 8 files changed, 217 insertions(+), 38 deletions(-) diff --git a/dnn-providers/hip-kernel-provider/rocke/library/tests/parity/conv_implicit_gemm_wgrad_emit.c b/dnn-providers/hip-kernel-provider/rocke/library/tests/parity/conv_implicit_gemm_wgrad_emit.c index 9c44f7169afd..2ef95dd66e64 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/tests/parity/conv_implicit_gemm_wgrad_emit.c +++ b/dnn-providers/hip-kernel-provider/rocke/library/tests/parity/conv_implicit_gemm_wgrad_emit.c @@ -417,14 +417,17 @@ static int make_base_cfg(int idx, rocke_implicit_gemm_conv_wgrad_spec_t* spec, c case 57: case 67: case 68: + case 69: + case 70: { /* gfx1250 WMMA async K loop. Must match the Python emitter's 55-57, - * 67 (55 with bf16 operands) and 68 (55 with an explicit vector_size_a, - * which both engines must reject). */ + * 67 (55 with bf16 operands), 68 (55 with an explicit vector_size_a, + * which both engines must reject) and 69-70 (55 on the synchronous + * K loop with A/B load-width caps 4/2, K-outer and M-outer). */ const bool big = idx == 56; if(idx == 57) spec->problem = rocke_conv_problem_default(4, 12, 12, 3, 32, 3, 3); - else if(idx == 55 || idx == 67 || idx == 68) + else if(idx == 55 || idx >= 67) spec->problem = rocke_conv_problem_default(4, 12, 12, 24, 32, 3, 3); else spec->problem = rocke_conv_problem_default(8, 56, 56, 64, 64, 3, 3); @@ -445,13 +448,18 @@ static int make_base_cfg(int idx, rocke_implicit_gemm_conv_wgrad_spec_t* spec, c spec->pipeline = "mem"; spec->epilogue = "default"; spec->split_k = big ? 4 : 2; - spec->lds_k_outer = true; - spec->async_dma = true; - if(idx == 68) + spec->lds_k_outer = idx != 70; + spec->async_dma = idx < 69; + if(idx >= 68) { spec->has_vector_size_a = true; spec->vector_size_a = 4; } + if(idx >= 69) + { + spec->has_vector_size_b = true; + spec->vector_size_b = 2; + } *arch = "gfx1250"; return 0; } diff --git a/dnn-providers/hip-kernel-provider/rocke/library/tests/parity/conv_implicit_gemm_wgrad_emit.py b/dnn-providers/hip-kernel-provider/rocke/library/tests/parity/conv_implicit_gemm_wgrad_emit.py index 39bd1894b2ff..7c89ef857915 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/tests/parity/conv_implicit_gemm_wgrad_emit.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/tests/parity/conv_implicit_gemm_wgrad_emit.py @@ -81,6 +81,10 @@ # 68 -- both engines must reject: case 55 with an explicit vector_size_a # (the async K loop derives its chunk widths from the channel counts, # so an explicit A/B load-width cap is rejected rather than ignored) +# 69-70 -- gfx1250 synchronous K loop with explicit A/B load-width caps +# (vector_size_a=4, vector_size_b=2 under the auto 8/8): case 55 +# without async_dma, K-outer (69, vector_axis="col") and M-outer +# (70, the transpose-on-store vector_axis="row" loader) # (Grouped wgrad reaches the C++ engine only through stream-K, and its tile # loader has no elem_dtype, so no other case loads bf16 operands: the # async loop's copy carries no element type, only its zero-fill does. The stream-K @@ -97,7 +101,7 @@ # 107 -- streamk with unroll_k (must raise ValueError) # 108 -- streamk atomic reduction into fp16 dW (must raise ValueError) # (These illustrate the validator contract. The C emitter defines only cases -# 0-68, so run_diff.py stops at the shared END before reaching 100+; these +# 0-70, so run_diff.py stops at the shared END before reaching 100+; these # configs are not exercised by the differential gate.) from kernels.common.conv_implicit_gemm_wgrad import ( WgradConvSpec, @@ -721,7 +725,7 @@ def _spec(idx: int): arch, ) - if 55 <= idx <= 57 or idx in (67, 68): + if 55 <= idx <= 57 or idx in (67, 68, 69, 70): # gfx1250 WMMA K loop: the global_load_async_to_lds ping-pong. from kernels.common._conv_implicit_gemm_common import ConvDataSpec @@ -741,6 +745,22 @@ def _spec(idx: int): {"async_dma": True, "vector_size_a": 4}, "gfx1250", ), + 69: ( + p_small, + "fp32", + 2, + False, + {"vector_size_a": 4, "vector_size_b": 2}, + "gfx1250", + ), + 70: ( + p_small, + "fp32", + 2, + False, + {"vector_size_a": 4, "vector_size_b": 2, "lds_k_outer": False}, + "gfx1250", + ), }[idx] wave32 = arch == "gfx1250" dtype_ab = "bf16" if idx == 67 else "fp16" diff --git a/dnn-providers/hip-kernel-provider/rocke/library/tests/test_conv_wgrad_streamk.py b/dnn-providers/hip-kernel-provider/rocke/library/tests/test_conv_wgrad_streamk.py index 9e6a17908a77..596b65f0eae0 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/tests/test_conv_wgrad_streamk.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/tests/test_conv_wgrad_streamk.py @@ -714,5 +714,82 @@ def test_cpp_builds_auto_pool(self): self.assertEqual(eng.conv_wgrad_lower_llvm(_spec_dict(spec), "gfx950"), py_ll) +class TestLoadWidthCapsDualEngine(unittest.TestCase): + """Explicit vector_size_a / vector_size_b through the ``rocke_engine`` binding. + + The caps narrow the synchronous free-axis loads, so both engines must + build the capped kernel and reject the same inadmissible caps. + """ + + # The parity emitter's case 55 shape: kpg=32 and cpg=24 both auto to 8. + _P = ConvProblem(N=4, Hi=12, Wi=12, C=24, K=32, Y=3, X=3, pH=1, pW=1) + + def _wmma(self, **kw) -> WgradConvSpec: + base = dict( + problem=self._P, + data=ConvDataSpec(dtype_a="fp16", dtype_b="fp16", dtype_d="fp32"), + tile_m=32, + tile_n=32, + tile_k=32, + warp_m=1, + warp_n=1, + warp_tile_m=16, + warp_tile_n=16, + warp_tile_k=32, + wave_size=32, + split_k=2, + lds_k_outer=True, + ) + base.update(kw) + return WgradConvSpec(**base) + + def test_capped_sync_kernel_is_byte_identical(self): + from rocke.core.lower_llvm import _lower_kernel_to_llvm_python + + eng = _engine() + cases = [ + ("gfx1250_kouter", self._wmma(), "gfx1250"), + ("gfx1250_mouter", self._wmma(lds_k_outer=False), "gfx1250"), + ("gfx950_streamk", _spec(), "gfx950"), + ] + for label, base, arch in cases: + for caps in ((4, 2), (4, None), (None, 2)): + spec = replace(base, vector_size_a=caps[0], vector_size_b=caps[1]) + with self.subTest(label, caps=caps): + self.assertEqual(is_valid_wgrad_spec(spec, arch), (True, "ok")) + py_ll = _lower_kernel_to_llvm_python( + build_implicit_gemm_conv_wgrad(spec, arch=arch), arch=arch + ) + uncapped = _lower_kernel_to_llvm_python( + build_implicit_gemm_conv_wgrad( + replace(base, vector_size_a=None, vector_size_b=None), + arch=arch, + ), + arch=arch, + ) + self.assertNotEqual(py_ll, uncapped) + self.assertEqual( + eng.conv_wgrad_lower_llvm(_spec_dict(spec), arch), py_ll + ) + + def test_cpp_validity_agrees(self): + eng = _engine() + cases = [ + ("ok_sync", self._wmma(vector_size_a=4, vector_size_b=2)), + ("async_a", self._wmma(async_dma=True, vector_size_a=4)), + ("async_b", self._wmma(async_dma=True, vector_size_b=2)), + ("a_zero", self._wmma(vector_size_a=0)), + ("b_over_cap", self._wmma(vector_size_b=16)), + ("a_not_dividing", self._wmma(vector_size_a=3)), + ("b_not_dividing", self._wmma(vector_size_b=5)), + ] + for label, spec in cases: + want = is_valid_wgrad_spec(spec, "gfx1250") + with self.subTest(label): + self.assertEqual(label.startswith("ok_"), want[0], want[1]) + got = eng.conv_wgrad_is_valid_spec(_spec_dict(spec), "gfx1250") + self.assertEqual(tuple(got), want) + + if __name__ == "__main__": unittest.main(verbosity=2) diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/bindings/rocke_engine.cpp b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/bindings/rocke_engine.cpp index 4b1d679c0ae8..645dd2beabf8 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/bindings/rocke_engine.cpp +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/bindings/rocke_engine.cpp @@ -3161,13 +3161,25 @@ rocke_implicit_gemm_conv_wgrad_spec_t conv_wgrad_build_spec(const py::dict& d, s.chiplet_chunk_size = dict_int(d, "chiplet_chunk_size", s.chiplet_chunk_size); /* Optional ints: absent or None keeps Python's None (has_* = false). * Only the fields the C builder or validator honours are read: the C wgrad - * builder does not apply lds_k_pad or vector_size_a/b, so reading them - * would tag the kernel name (e.g. "pad") without changing its body. */ + * builder does not apply lds_k_pad, so reading it would tag the kernel + * name ("pad") without changing its body. vector_size_a/b cap the sync + * free-axis load widths (and are rejected on the async K loop); dropped, + * the engine would build the uncapped kernel and skip those checks. */ if(d.contains("waves_per_eu") && !d["waves_per_eu"].is_none()) { s.has_waves_per_eu = true; s.waves_per_eu = d["waves_per_eu"].cast(); } + if(d.contains("vector_size_a") && !d["vector_size_a"].is_none()) + { + s.has_vector_size_a = true; + s.vector_size_a = d["vector_size_a"].cast(); + } + if(d.contains("vector_size_b") && !d["vector_size_b"].is_none()) + { + s.has_vector_size_b = true; + s.vector_size_b = d["vector_size_b"].cast(); + } if(d.contains("vector_size_c") && !d["vector_size_c"].is_none()) { s.has_vector_size_c = true; diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/common/conv_implicit_gemm_wgrad.cpp b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/common/conv_implicit_gemm_wgrad.cpp index 3041e560a3d2..fe7651d724ac 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/common/conv_implicit_gemm_wgrad.cpp +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/common/conv_implicit_gemm_wgrad.cpp @@ -600,14 +600,16 @@ static bool wgrad_group_merge_available(const rocke_implicit_gemm_conv_wgrad_spe #undef GM_REJECT } -/* Python _wgrad_free_axis_vec(chan, dtype) with no override. */ -static int wgrad_free_axis_vec(int chan, const char* dtype) +/* Python _wgrad_free_axis_vec(chan, dtype, override): the widest of {8,4,2,1} + * ({4,2,1} for fp32) dividing the channel run, clamped by an explicit + * vector_size_a / vector_size_b. The override is a cap, not a replacement. */ +static int wgrad_free_axis_vec(int chan, const char* dtype, bool has_override, int override_vec) { const bool f32 = dtype != NULL && strcmp(dtype, "fp32") == 0; - for(int v = f32 ? 4 : 8; v > 1; v /= 2) - if(chan % v == 0) - return v; - return 1; + int v = f32 ? 4 : 8; + while(v > 1 && chan % v != 0) + v /= 2; + return (has_override && override_vec < v) ? override_vec : v; } /* Mirrors Python _wmma_async_available: whether a wave32 spec can take the @@ -636,8 +638,12 @@ static bool wgrad_wmma_async_available(const rocke_implicit_gemm_conv_wgrad_spec return false; } const int threads = s->warp_m * s->warp_n * s->wave_size; - const int cap_a = wgrad_free_axis_vec(rocke_conv_problem_kpg(&s->problem), s->dtype_a); - const int cap_b = wgrad_free_axis_vec(rocke_conv_problem_cpg(&s->problem), s->dtype_b); + /* Python _wmma_async_vecs passes no override: the async K loop does not + * honour vector_size_a/b (the validator rejects them there). */ + const int cap_a + = wgrad_free_axis_vec(rocke_conv_problem_kpg(&s->problem), s->dtype_a, false, 0); + const int cap_b + = wgrad_free_axis_vec(rocke_conv_problem_cpg(&s->problem), s->dtype_b, false, 0); int va = 1; int vb = 1; if(rocke_coalesced_tile_loader_choose_vec_axis(s->tile_k, s->tile_m, threads, cap_a, false, &va) @@ -3101,26 +3107,16 @@ static bool wgrad_build_ctx_init(rocke_conv_build_ctx_t* ctx, /* Every sync path (MFMA and WMMA); the async branch is handled above. */ { /* free-axis width: widest of {8,4,2,1} ({4,2,1} for fp32) dividing - * the channel count -- Python _free_axis_vec(kpg/cpg, dtype). */ - int kpg = rocke_conv_problem_kpg(&spec->problem); - int cpg = rocke_conv_problem_cpg(&spec->problem); - bool a_fp32 = (spec->dtype_a != NULL && strcmp(spec->dtype_a, "fp32") == 0); - bool b_fp32 = (spec->dtype_b != NULL && strcmp(spec->dtype_b, "fp32") == 0); - int cap_a = 1; - int cap_b = 1; - int wa, wb; - for(wa = a_fp32 ? 4 : 8; wa >= 1; wa /= 2) - if(kpg % wa == 0) - { - cap_a = wa; - break; - } - for(wb = b_fp32 ? 4 : 8; wb >= 1; wb /= 2) - if(cpg % wb == 0) - { - cap_b = wb; - break; - } + * the channel count, clamped by an explicit vector_size_a/b -- + * Python _free_axis_vec(kpg/cpg, dtype, spec.vector_size_a/b). */ + const int cap_a = wgrad_free_axis_vec(rocke_conv_problem_kpg(&spec->problem), + spec->dtype_a, + spec->has_vector_size_a, + spec->vector_size_a); + const int cap_b = wgrad_free_axis_vec(rocke_conv_problem_cpg(&spec->problem), + spec->dtype_b, + spec->has_vector_size_b, + spec->vector_size_b); /* Mirror the Python choose_vec, which RAISES ValueError when the tile * geometry admits no usable width (not even 1): fail fast here too so * a bad geometry surfaces as a clear error instead of silently diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/tests/golden/rocke_representative_ir_sha256/conv_wgrad.json b/dnn-providers/hip-kernel-provider/rocke/platform/tests/golden/rocke_representative_ir_sha256/conv_wgrad.json index 33e2aa3664bc..44614a74fa57 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/tests/golden/rocke_representative_ir_sha256/conv_wgrad.json +++ b/dnn-providers/hip-kernel-provider/rocke/platform/tests/golden/rocke_representative_ir_sha256/conv_wgrad.json @@ -43,6 +43,14 @@ "sha256": "d5edf46c7b910114fad2f49bfce7d464f2d04a66a5e3c064c337ad643fad40ea", "status": "ok" }, + "conv_wgrad/gfx1250/n1h8c16k32r3_spk4_kouter_vec4x2": { + "arch": "gfx1250", + "bytes": 70109, + "family": "conv_wgrad", + "kernel_name": "irhash_wgrad_1250_spk4_kouter_vec4x2_N1H8W8C16_K32Y3X3_t32x32x32_w2x1_a16x16x32_mem_default_kouter_spk", + "sha256": "8158ef708355f1aaada52a997d9e512679b2ef533102c83cb27f683846b12e30", + "status": "ok" + }, "conv_wgrad/gfx1250/n1h8c64k64r3_g4": { "arch": "gfx1250", "bytes": 66522, @@ -240,6 +248,14 @@ "sha256": "fbca7fe69c8677120991634e159227a1de11682f37daaaa3b00cbbf442d289a9", "status": "ok" }, + "conv_wgrad/gfx1250/n1h8c16k32r3_spk4_kouter_vec4x2": { + "arch": "gfx1250", + "bytes": 70248, + "family": "conv_wgrad", + "kernel_name": "irhash_wgrad_1250_spk4_kouter_vec4x2_N1H8W8C16_K32Y3X3_t32x32x32_w2x1_a16x16x32_mem_default_kouter_spk", + "sha256": "0a1b0fd05bc151eff083b235a31707b017f73d496379301f8e3dc5f69884272f", + "status": "ok" + }, "conv_wgrad/gfx1250/n1h8c64k64r3_g4": { "arch": "gfx1250", "bytes": 66661, @@ -437,6 +453,14 @@ "sha256": "bc450909719641a857aca3d5d35096545aa83e065ac70c43ac16ec77e0781127", "status": "ok" }, + "conv_wgrad/gfx1250/n1h8c16k32r3_spk4_kouter_vec4x2": { + "arch": "gfx1250", + "bytes": 70312, + "family": "conv_wgrad", + "kernel_name": "irhash_wgrad_1250_spk4_kouter_vec4x2_N1H8W8C16_K32Y3X3_t32x32x32_w2x1_a16x16x32_mem_default_kouter_spk", + "sha256": "e3514b46ba44f72d72cec6d11c14da6fd963be6789ee17134c7d259e30940c23", + "status": "ok" + }, "conv_wgrad/gfx1250/n1h8c64k64r3_g4": { "arch": "gfx1250", "bytes": 66725, @@ -634,6 +658,14 @@ "sha256": "bc450909719641a857aca3d5d35096545aa83e065ac70c43ac16ec77e0781127", "status": "ok" }, + "conv_wgrad/gfx1250/n1h8c16k32r3_spk4_kouter_vec4x2": { + "arch": "gfx1250", + "bytes": 70312, + "family": "conv_wgrad", + "kernel_name": "irhash_wgrad_1250_spk4_kouter_vec4x2_N1H8W8C16_K32Y3X3_t32x32x32_w2x1_a16x16x32_mem_default_kouter_spk", + "sha256": "e3514b46ba44f72d72cec6d11c14da6fd963be6789ee17134c7d259e30940c23", + "status": "ok" + }, "conv_wgrad/gfx1250/n1h8c64k64r3_g4": { "arch": "gfx1250", "bytes": 66725, diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/tests/instances/differential/golden/llvm_gfx_all.json b/dnn-providers/hip-kernel-provider/rocke/platform/tests/instances/differential/golden/llvm_gfx_all.json index 12b1774d9219..6813aff12660 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/tests/instances/differential/golden/llvm_gfx_all.json +++ b/dnn-providers/hip-kernel-provider/rocke/platform/tests/instances/differential/golden/llvm_gfx_all.json @@ -576,7 +576,9 @@ "64": "0ea52c6b7172c309a13d3aadfff7f232f2f31f251b90d80ebdc9c4e80754b1f8", "65": "d0dee69c6ef1f69eeaf9250d4c0044cf4fd9713d9ee65c6e916f48453170363d", "67": "f9bc06dc3bcc5f12696b8742d2bd64abbc170b256b89893706d7c31153b69649", + "69": "d926b9ba827147742087493fe29e5efe0bf8ac1126cb702793d3751c17b8a787", "7": "1106d4ce71314f3c532427fed8ceb00855f1cf1dd7f5117b89234cf326583087", + "70": "0cba53f2979d05f0368e8d72e89b39af3cd8b605872202b179a723bca71186f1", "8": "82d8923b31be6a07c4371fbe40b33a2ea7452027299ce0993e903cf55ad092d8" }, "conv_wgrad_workspace_reduce": { @@ -2456,7 +2458,9 @@ "64": "bb76cde6ee2149851980c297b858c581cd115cef526902134f613937aa0f65f8", "65": "a79cd079bd3671c921a244740b4a3e94594e600c4a84d59075ad2346e56e6917", "67": "2e0365dba91768754f0fe4fda140370b218e3c65b3385d869fe9e2955a3070d5", + "69": "f0753cf45846473ab80e98adf23c3805ef4c9354170ca9b1ce235fcce3ec9687", "7": "d653088ddce4d716c3b9e20fabe5592541b52ba5e800786a86b660763c25a596", + "70": "c30b77e00298c9769491cddfd407908798da5a63671019c3303ad38770223862", "8": "889d9e3617a8eb2e30d29db4ef8da0aeac52c9fa0d92bc53541947cf85926c08" }, "conv_wgrad_workspace_reduce": { @@ -4238,7 +4242,9 @@ "64": "1fc144e9617d38a78751c4452ebbe78aa0bcc06fd4a5a3e0f90eb8cf0213b268", "65": "379d276960a73e2da5b2c995088d3cbbd18f1bde2c47feeddd93602910d9418d", "67": "881028a810289a2e95263f6e9317588742868c19826dc9d9ac49606a71fdc5c2", + "69": "e9c9c979c9066e813cc5dd61e7a3604e52fa125147f182885467683788343129", "7": "e12021bd6f8923472a7595394ecde03a55dcd07d63d714b9adaea3efa79df887", + "70": "572073b022f0d0d0f36bd13f16ad18a23ab5a754307e2899b57d4d523b46b07f", "8": "2cd90bffada0eaec9a90caf5d5b76f2f213e4e884e3f413b13e79eef3071ec82" }, "conv_wgrad_workspace_reduce": { @@ -6050,7 +6056,9 @@ "64": "861ee261e315bb0202240ac717921ef9fb69b24ea6949194d5190778261d4706", "65": "88e233291d3674ce2dd072faaed3eb546ea10d8e036f741a06c76531d19570a4", "67": "79f10f6395545d6b2356fc4d1d4d02e6beb39a4ffe92f5d4060dc8abe0f2e67a", + "69": "1f70534411918acb3b05879c5ab3d30a2ae686ac3c7c337fa7c0fae5fd3bef2d", "7": "44a58b32e3d95d8492b23193d701925de3dd9de8cddf8a2c006b6fdf5de87cf8", + "70": "b94ad8371942d595842a888a2c867d1a3802274eaf763bee1ae6e02e5a10af16", "8": "af0df509c9e1bfa9bcb32621a677b633c9bd7c0b0a22012ee20acd2894b44df8" }, "conv_wgrad_workspace_reduce": { @@ -8114,7 +8122,9 @@ "64": "861ee261e315bb0202240ac717921ef9fb69b24ea6949194d5190778261d4706", "65": "88e233291d3674ce2dd072faaed3eb546ea10d8e036f741a06c76531d19570a4", "67": "79f10f6395545d6b2356fc4d1d4d02e6beb39a4ffe92f5d4060dc8abe0f2e67a", + "69": "1f70534411918acb3b05879c5ab3d30a2ae686ac3c7c337fa7c0fae5fd3bef2d", "7": "44a58b32e3d95d8492b23193d701925de3dd9de8cddf8a2c006b6fdf5de87cf8", + "70": "b94ad8371942d595842a888a2c867d1a3802274eaf763bee1ae6e02e5a10af16", "8": "af0df509c9e1bfa9bcb32621a677b633c9bd7c0b0a22012ee20acd2894b44df8" }, "conv_wgrad_workspace_reduce": { diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/tests/instances/rocke_ir_parity_harness.py b/dnn-providers/hip-kernel-provider/rocke/platform/tests/instances/rocke_ir_parity_harness.py index 247447efaab4..d803635479aa 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/tests/instances/rocke_ir_parity_harness.py +++ b/dnn-providers/hip-kernel-provider/rocke/platform/tests/instances/rocke_ir_parity_harness.py @@ -2889,6 +2889,30 @@ def add(family, case_id, arch, build): async_dma=True, ), ) + # The synchronous K-outer loop with explicit A/B load-width caps (4 and 2 + # under the auto 8/8), which both engines apply to the free-axis loads. + add( + "conv_wgrad", + "conv_wgrad/gfx1250/n1h8c16k32r3_spk4_kouter_vec4x2", + "gfx1250", + build_conv_wgrad( + "irhash_wgrad_1250_spk4_kouter_vec4x2", + "gfx1250", + wgrad1, + wave_size=32, + wtm=16, + wtn=16, + wtk=32, + tile_m=32, + tile_n=32, + tile_k=32, + split_k=4, + dtype_d="bf16", + lds_k_outer=True, + vector_size_a=4, + vector_size_b=2, + ), + ) # WMMA split-K through the cshuffle epilogue (paired atomics from LDS). add( "conv_wgrad", From b02e4e28018e4833a7af35a645d900eaae695dc0 Mon Sep 17 00:00:00 2001 From: Jakub Piasecki Date: Wed, 4 Nov 2026 19:46:03 +0000 Subject: [PATCH 13/13] test(rocke): pin the smaller of the auto width and the wgrad A/B cap A cap wider than the auto free-axis width (cap 6 at kpg 12, auto 4) must leave the width at 4. The new binding test checks that both engines lower that spec to the uncapped kernel's IR, on the K-outer and M-outer loops. Also drop the _free_axis_vec alias of _wgrad_free_axis_vec in the builder and name _wgrad_free_axis_vec in the C++ comment that cited the alias. Co-Authored-By: Claude Opus 5.5 --- .../common/conv_implicit_gemm_wgrad.py | 13 +++++---- .../library/tests/test_conv_wgrad_streamk.py | 27 +++++++++++++++++++ .../common/conv_implicit_gemm_wgrad.cpp | 2 +- 3 files changed, 36 insertions(+), 6 deletions(-) diff --git a/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/conv_implicit_gemm_wgrad.py b/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/conv_implicit_gemm_wgrad.py index 78396a6f8cf4..1c21cfa707d1 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/conv_implicit_gemm_wgrad.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/conv_implicit_gemm_wgrad.py @@ -2598,7 +2598,6 @@ def _declare_wgrad_ptr(name: str, kind: str) -> Value: axis_a = axis_b = "col" load_vec_a = 1 load_vec_b = 1 - _free_axis_vec = _wgrad_free_axis_vec # gfx1250 async: the per-lane global_load_async_to_lds K loop. It reuses the # synchronous K-outer loaders for the slot partition and the descriptors, # and only swaps the global->register->LDS hop for one async copy. @@ -2609,14 +2608,18 @@ def _declare_wgrad_ptr(name: str, kind: str) -> Value: tile_rows=block_m, tile_cols=block_k, block_size=threads, - max_vec=_free_axis_vec(p_load.kpg, spec.data.dtype_a, spec.vector_size_a), + max_vec=_wgrad_free_axis_vec( + p_load.kpg, spec.data.dtype_a, spec.vector_size_a + ), vector_axis="row", ) vb = CoalescedTileLoader.choose_vec( tile_rows=block_n, tile_cols=block_k, block_size=threads, - max_vec=_free_axis_vec(p_load.cpg, spec.data.dtype_b, spec.vector_size_b), + max_vec=_wgrad_free_axis_vec( + p_load.cpg, spec.data.dtype_b, spec.vector_size_b + ), vector_axis="row", ) if va > 1: @@ -2753,7 +2756,7 @@ def x_descriptor(b_: IRBuilder, row: Value, col: Value): tile_rows=a_tile_rows, tile_cols=a_tile_cols, block_size=threads, - max_vec=_free_axis_vec( + max_vec=_wgrad_free_axis_vec( p_load.kpg, spec.data.dtype_a, spec.vector_size_a ), vector_axis="col", @@ -2762,7 +2765,7 @@ def x_descriptor(b_: IRBuilder, row: Value, col: Value): tile_rows=b_tile_rows, tile_cols=b_tile_cols, block_size=threads, - max_vec=_free_axis_vec( + max_vec=_wgrad_free_axis_vec( p_load.cpg, spec.data.dtype_b, spec.vector_size_b ), vector_axis="col", diff --git a/dnn-providers/hip-kernel-provider/rocke/library/tests/test_conv_wgrad_streamk.py b/dnn-providers/hip-kernel-provider/rocke/library/tests/test_conv_wgrad_streamk.py index 596b65f0eae0..93854bf8a0c7 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/tests/test_conv_wgrad_streamk.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/tests/test_conv_wgrad_streamk.py @@ -772,6 +772,33 @@ def test_capped_sync_kernel_is_byte_identical(self): eng.conv_wgrad_lower_llvm(_spec_dict(spec), arch), py_ll ) + def test_cap_wider_than_auto_width_keeps_auto_width(self): + # kpg=12 autos to 4; a cap of 6 is admissible (it divides 12) but + # wider, so the load width stays 4 (the smaller of the two), not 6. + from rocke.core.lower_llvm import _lower_kernel_to_llvm_python + + eng = _engine() + base = self._wmma(problem=replace(self._P, K=12)) + for label, b in ( + ("kouter", base), + ("mouter", replace(base, lds_k_outer=False)), + ): + spec = replace(b, vector_size_a=6) + with self.subTest(label): + self.assertEqual(is_valid_wgrad_spec(spec, "gfx1250"), (True, "ok")) + py_ll = _lower_kernel_to_llvm_python( + build_implicit_gemm_conv_wgrad(spec, arch="gfx1250"), + arch="gfx1250", + ) + auto_ll = _lower_kernel_to_llvm_python( + build_implicit_gemm_conv_wgrad(b, arch="gfx1250"), + arch="gfx1250", + ) + self.assertEqual(py_ll, auto_ll) + self.assertEqual( + eng.conv_wgrad_lower_llvm(_spec_dict(spec), "gfx1250"), py_ll + ) + def test_cpp_validity_agrees(self): eng = _engine() cases = [ diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/common/conv_implicit_gemm_wgrad.cpp b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/common/conv_implicit_gemm_wgrad.cpp index fe7651d724ac..d5491a37064e 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/common/conv_implicit_gemm_wgrad.cpp +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/common/conv_implicit_gemm_wgrad.cpp @@ -3108,7 +3108,7 @@ static bool wgrad_build_ctx_init(rocke_conv_build_ctx_t* ctx, { /* free-axis width: widest of {8,4,2,1} ({4,2,1} for fp32) dividing * the channel count, clamped by an explicit vector_size_a/b -- - * Python _free_axis_vec(kpg/cpg, dtype, spec.vector_size_a/b). */ + * Python _wgrad_free_axis_vec(kpg/cpg, dtype, spec.vector_size_a/b). */ const int cap_a = wgrad_free_axis_vec(rocke_conv_problem_kpg(&spec->problem), spec->dtype_a, spec->has_vector_size_a,