feat(rocke/conv): add transpose load for bwd weight and bwd data by aledudek · Pull Request #11602 · ROCm/rocm-libraries

aledudek · 2026-09-02T12:15:46Z

ISSUE ID: AICK-1946

Add the K-outer LDS transpose load for backward-weight (wgrad) and
backward-data (dgrad) convolution, plus the direct global->LDS loader and a
basic K-loop on top of it.

Motivation

Improve rocke performance.

Technical Details

K-outer LDS tile: store A/B in global order and transpose on the operand
fetch with ds_read_tr16_b64, instead of transposing on store. Mirrored in
the C engine; ROCKE_WGRAD_KOUTER_PAD mirrors _KOUTER_PAD, and a "kouter"
kernel-name flag sits between "async" and "spk*".
Direct global->LDS load (async_dma) on the K-outer tile, with the
AsyncTileLoader chunk-addressing fix and contig_cols guards.
pipeline='basic': single LDS buffer, global-read/compute overlap.
Selection, not knobs. lds_k_outer is deduced by
WgradConvSpec.default_lds_k_outer (arch/dtype/atom/wave size), which both
dispatch and the sweep driver call. async_dma is a swept axis, since the
pad it forces and the load-width ladder it narrows both depend on sweep
axes. No --lds-k-outer, --lds-k-pad or --dtype-d flag, and no env override.
Output dtype is left to the two-stage work in feat(rocke): add two-stage deterministic backward-weight (wgrad) convolution #10571.
Validation. lds_k_outer is gated on gfx950 in both engines
(ds_read_tr16_b64 is CDNA4 only); an explicit lds_k_pad with lds_k_outer is
an error rather than a silent no-op; async specs run validate_for_async();
runtime split-K capability is decided per combo, so pipeline='basic' is no
longer dropped from a --split-k 0 sweep.

dgrad (B tile only)

Same K-outer tile and ds_read_tr16_b64 transpose read, applied to B (W,
KYXC). A (dY, NHWK) stays M-outer: its reduction axis is already stride-1,
so it has no transpose-on-store to remove and flipping it would destroy
global coalescing.
Selection, not a knob: DgradConvSpec.default_lds_k_outer, gated on
gfx950/MFMA/wave64/16-bit B and on an even channel run.
rocke_conv_tr_frag is shared between the compute phase and dgrad's fetch,
so there is one lane mapping rather than two copies.
Parity configs 11 and 12; config 12 covers the 16x16x16 atom, where the
per-lane fragment length is 4 rather than 8.

Parity: configs 11-14 cover K-outer, async and the 16x16x16 atom in both the
.py and .c emitters. They use cshuffle deliberately -- the validator rejects
16-bit output with epilogue="default", and a rejected config lands as
BOTH_REJECTED, which the gate counts as a pass.

Golden
One entry moves: conv/gfx950/n1h8c16k32r3, all three llvm flavors. Cause is
dropping "compv4" from the double_buffer expression in conv_implicit_gemm.py
-- compv4 shares the single-buffer K-loop, so the second A/B tile was a dead
allocation that only shrank the LDS budget and rejected valid specs. Both
allocs were already dead-stripped, so the IR is identical modulo SSA
renumbering, hence -4 bytes. No wgrad golden moves.

Test Plan

Byte-identity GREEN across all families at every llvm flavor.
Golden: only the entry above.
GPU numeric: test_lds_k_outer_matches_default and the split-K variants
against the torch oracle.
dgrad K-outer output is bitwise identical to the M-outer path across bf16
and fp16, the 32x32x16 and 16x16x16 atoms, strided (tilde) shapes, and a
non-tile-aligned K. split_k is checked against a torch reference instead,
because the atomic epilogue's accumulation order is not deterministic.

Submission Checklist

Look over the contributing guidelines at https://github.com/ROCm/ROCm/blob/develop/CONTRIBUTING.md#pull-requests.

therock-pr-bot · 2026-09-02T12:21:55Z

✅ All Checks Passed — Ready for Review

Check	Status	Details
📝 PR Description	✅ Pass	—
⛔ Forbidden Files	✅ Pass	—
🧪 Unit Test	✅ Pass	—
🔎 pre-commit	✅ Pass	—
🚫 Draft PR	🔜 To Be Enabled	—
🚩 Feature Flag	🔜 To Be Enabled	—
📊 Code Coverage	🔜 To Be Enabled	—
🤖 therock-pr-bot	✅ Pass	—

🎉 All checks passed! This PR is ready for review.

📖 Need help? See the Policy FAQ for details on every check and how to fix failures.

🙋 Wish to Override Policy?

therock-pr-bot · 2026-09-02T12:21:59Z

Pre-commit check failed

⛔ pre-commit failed

Please run locally:

python -m pip install pre-commit
pre-commit install
pre-commit run --all-files --show-diff-on-failure

This repo uses .pre-commit-config.yaml.

therock-pr-bot · 2026-09-04T07:29:56Z

🎉 All checks passed! This PR is ready for review.

Copilot

🟡 Changes recommended

Wgrad async validation is incomplete (missing validate_for_async()), plus there’s a benchmark vec reporting bug and a PR-description mismatch vs the updated golden.

Once you've addressed the issues Copilot identified, you can request another Copilot review.

Pull request overview

This PR adds and validates a new wgrad “K-outer” LDS layout path (plus optional async global→LDS loads) to improve performance on gfx950 while preserving Python↔C++ byte-identity for representative parity configs, and it tightens async tile-loader correctness.

Changes:

Implement/port wgrad K-outer LDS storage with gfx950 transpose reads (ds_read_b64_tr_b16) and add parity configs to exercise it in both Python and C.
Fix and harden AsyncTileLoader chunk decoding/shape admissibility (including contiguous-column constraints) and add focused unit tests.
Extend tests/bench/dispatch knobs for wgrad (lds_k_outer, async_dma, lds_k_pad, output dtype) and document the new optimization knob.

File summaries

File	Description
dnn-providers/hip-kernel-provider/rocke/platform/tests/instances/test_conv_wgrad_correctness.py	Adds correctness coverage for `lds_k_outer`/async_dma variants and atom edge cases.
dnn-providers/hip-kernel-provider/rocke/platform/tests/instances/parity/conv_implicit_gemm_wgrad_emit.py	Adds parity configs to exercise K-outer and async_dma paths in the Python emitter.
dnn-providers/hip-kernel-provider/rocke/platform/tests/instances/parity/conv_implicit_gemm_wgrad_emit.c	Adds corresponding parity configs in the C emitter to preserve byte-identity gating coverage.
dnn-providers/hip-kernel-provider/rocke/platform/tests/golden/rocke_representative_ir_sha256.json	Updates representative IR hashes/sizes for affected kernels.
dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_async_tile_loader_decode.py	New unit tests for async tile-loader decode arithmetic/regressions (fast tier).
dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/instances/common/conv_implicit_gemm.py	Corrects double-buffering/LDS budgeting logic for `compv4` vs async/unroll.
dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/instances/common/conv_implicit_gemm_wgrad.py	Implements wgrad K-outer layout, transpose-read fragment feed, async_dma gating, and validation/policy updates.
dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/loads.py	Fixes async loader decode divisor and adds contig/run + whole-pass admissibility constraints.
dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/benchmark/benchmark_implicit_gemm_conv.py	Extends benchmark CLI/sweep axes for wgrad (pad/out dtype/kouter/async) and verification tolerance logic.
dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/optimization/optimization_runbook.md	Documents the new `lds_k_outer` tuning knob and selection policy.
dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/common/conv_implicit_gemm_wgrad.cpp	C++ engine port of K-outer path, validation gates, kernel naming flags, and K-loop plumbing.
dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/common/conv_implicit_gemm_conv_load_phase.cpp	Ensures async/basic global-read paths honor per-family descriptor overrides (needed for wgrad).
dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/common/conv_implicit_gemm_conv_kloop_drivers.cpp	Adds k-offset sequencing/iteration overrides for wgrad split-K and preserves SSA ordering.
dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/common/conv_implicit_gemm_conv_compute_phase.cpp	Adds transpose-read fragment feed for K-outer path (shared compute phase).
dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/common/conv_implicit_gemm_conv_build_glue.cpp	Aligns double-buffer selection and async loader contig_cols wiring with Python.
dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/instance_conv_implicit_gemm_wgrad.h	Adds public ABI fields/macros for `lds_k_outer` and basic-pipeline unroll cap.
dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/instance_conv_implicit_gemm_internal.h	Extends build ctx with K-outer/descriptor/k-offset/iter-count fields for parity.
dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/helper_rocke.helpers.loads.h	Updates C helper API for async loader contig_cols and decode divisor.
dnn-providers/hip-kernel-provider/rocke/platform/cpp/helpers/loads.cpp	Ports async loader decode fix + contig_cols/whole-pass admissibility constraints to C.
dnn-providers/hip-kernel-provider/rocke/platform/AGENTS.md	Documents env override for wgrad K-outer selection policy.
dnn-providers/hip-kernel-provider/rocke/library/dispatch/grouped_convolution.py	Propagates/auto-selects `lds_k_outer` for gfx950 wgrad dispatch candidates.
dnn-providers/hip-kernel-provider/.gitignore	Ignores Rocke benchmark CSV outputs.

Review details

Suppressed comments (1)

dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/instances/common/conv_implicit_gemm_wgrad.py:551

WgradConvSpec.validate() no longer calls layout.validate_for_async() when async_dma is enabled, so specs can be constructed (and is_valid_wgrad_spec can return ok) even if a user-provided lds_layout has k_pad != 0 or a swizzle, which will only fail later during lowering. This makes validation inconsistent with conv/dgrad and can turn dispatch/tuning into a late build failure.

        layout = self.effective_lds_layout()
        if self.async_dma and self.lds_k_pad not in (None, 0):
            raise ValueError(
                "async_dma requires lds_k_pad to be 0/None because "
                "raw_ptr_buffer_load_lds writes a packed lane-contiguous tile"

Files reviewed: 22/22 changed files
Comments generated: 3
Review effort level: Lite

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

… it in dispatch Port the K-outer LDS + ds_read_b64_tr_b16 wgrad path to the C99 engine, add the parity config that proves the two engines agree, and select it from the gfx950 wgrad dispatch candidate so real callers get it.

ISSUE ID: AICK-XXXX Make async_dma legal and correct for wgrad, gated on lds_k_outer, and add the contig_cols guard that keeps a direct-load chunk inside one contiguous global run. The guard also fixes a live silent-wrongness bug in the forward conv async path. Direct load is available but NOT selected by default: it currently measures slower than the synchronous loader on the K-outer tile (see below). ## Motivation Direct load was rejected outright for wgrad because raw_ptr_buffer_load_lds maps contiguous-global to contiguous-LDS, and wgrad's reduction axis K_wg is strided in both operands. The K-outer tile removes that obstacle: a chunk is then a run along the free axis at a fixed (n, ho, wo, y, x), which is contiguous in global AND shares a single boundary-pad predicate, so one load with one predicate is correct. The packed layout the intrinsic forces is also free here, because the transpose read is insensitive to the LDS row stride. ## Technical Details Validator: the blanket wgrad async_dma rejection becomes "async_dma requires lds_k_outer", in both engines. Loader: on the K-outer path the async loader is built with the tile axes swapped (tile_rows = block_k, tile_cols = the free axis) and the existing K-outer descriptor adapters are shared with the sync path, so global addressing is unchanged. The LDS pad is forced to 0 for async, since the intrinsic writes packed lane-contiguous bytes and cannot skip a pad. New AsyncTileLoader.contig_cols: the length of the longest run of tile columns that is contiguous in global memory. choose_dwords now refuses any chunk width that exceeds it or does not divide it, and degrades to a narrower dwords value instead. wgrad passes kpg for dY (NHWK, dense over the output-channel slab) and cpg for X (NHWC, dense only within one filter position). This is a real bug fix beyond wgrad. The forward conv reduces over (y, x, c) where only c is stride-1, and its async path passed no such bound: for any cpg that is not a multiple of the chunk width the chunk straddled a filter position and fetched the wrong elements with no diagnostic. Verified on GPU: fwd async at C=4 and C=12 went from failing the torch oracle to passing, and C=64 is unchanged. C engine mirror, all NULL/0-defaulted so fwd and dgrad emit exactly what they emitted before: - contig_cols on the loader struct, choose_dwords and from_tile. - ctx->a_descriptor_fn: the async A slot had no hook and always called the forward descriptor, so a wgrad async kernel would have addressed dY wrongly. - ctx->kloop_k_lo: wgrad's k-loop offsets are add(k_lo, const) because the slice starts at the split-K lower bound, where the forward conv uses a bare const. - The build-path stub forward spec now carries async_dma and unroll_k. Without that the shared load phase and k-loop drivers never saw the flags, so both paths were silently unreachable from wgrad and fell through to a zeroed sync loader. Parity config 12 (K-outer + async_dma) covers the new path in both emitters. ## Test Plan Byte-identity GREEN - 69/69 families. The wgrad family is GREEN with config 11 (K-outer) and config 12 (K-outer + async_dma) both IDENTICAL, so the new path is genuinely compared between the engines rather than skipped. Golden unchanged. Suite - 727 passed, 1540 subtests, 80 skipped. GPU numeric - new test_lds_k_outer_async_dma covers fp16 and bf16 across every shape in the sweep, including the C=24 and C=8 shapes that exercise the contig_cols guard; new test_async_dma_requires_k_outer pins the rejection. Direct contiguity stress on GPU: C=4 and C=12 failed the oracle before the guard and pass after it; C=24 and C=64 pass throughout. Forward conv async: C=4 and C=12 FAIL -> PASS, C=64 unchanged. ## Test Result All of the above passed. Direct load is correct but slower than the synchronous loader on the K-outer tile, consistently across repeats and unchanged by unroll_k. The cause is visible in the counters: the async k-loop driver requires double-buffered LDS (single-buffering it is numerically wrong), which doubles the LDS footprint and costs occupancy, while VGPR usage goes up rather than down because the synchronous path was not register-bound to begin with. It does deliver the intended reduction in LDS instructions per MFMA; that saving simply does not pay for the occupancy loss here. It is therefore left available but off: async_dma still defaults to False and the gfx950 dispatch candidate does not set it. Landing it correctly removes a silently-wrong code path, fixes the forward-conv guard bug, and leaves the option open for shapes or pipelines where the trade lands differently. ## Submission Checklist - [ ] Look over the contributing guidelines at https://github.com/ROCm/ROCm/blob/develop/CONTRIBUTING.md#pull-requests.

ISSUE ID: AICK-XXXX Expose WgradConvSpec.async_dma on the benchmark driver so the direct global->LDS load path can be A/B'd against the register-staged loader without editing the sweep. Threaded through the picklable work tuple so it survives --jobs > 1. Requires --lds-k-outer; on its own every config is rejected by the validator, which is the intended behaviour.

…DS over-run ISSUE ID: AICK-XXXX Two silent-wrongness fixes in the direct-load path, both found by extending the parity harness to configurations it did not previously cover. ## Motivation The wgrad async parity config landed with split_k at its default of 1, and that was the only reason the byte-identity gate passed: at split_k > 1 the two engines emitted different kernels. Separately, the async loader could write past the LDS tile it was given for any geometry whose chunk count is not a whole multiple of the block size. Neither had test coverage. ## Technical Details Iteration count. The shared unrolled and async k-loop drivers sized the loop from the problem's full K_gemm. For wgrad that is the stub problem's UNSLICED wg_K, while the Python emitter sizes it from wg_K_padded()/split_k -- so the engines agreed only at split_k == 1, and the C engine emitted a full-range reduction under split-K atomics. New ctx->kloop_num_iters carries the sliced count; 0 keeps the previous derivation, so the forward conv and dgrad are untouched. Applied at all three driver sites (unroll, simple-unroll, async). LDS over-run. AsyncTileLoaderSlot.issue fires block_size chunks per pass, and chunks past chunks_total still issue with the OOB sentinel, which the hardware zero-fills. When chunks_total is not a multiple of block_size the final partial pass therefore writes past the tile's allocation and into the next smem_alloc. choose_dwords now requires whole passes and falls through to a narrower dwords value. Verified this loses no configurations: the affected tiles (32x48, 32x96, 64x48) degrade from dwords=4 to dwords=3 and land exactly inside their allocation. Note AsyncTileLoaderSlot.required_lds_bytes() has no callers, which is why nothing caught this. Parity config 12 now uses split_k=4 with fp32 output and the direct-store epilogue, so it actually gates the iteration-count fix. ## Test Plan Byte-identity GREEN - 69/69 families, with wgrad config 11 (K-outer) and config 12 (K-outer + async + split_k=4) both IDENTICAL. Config 12 at split_k=4 reproduced the divergence before the fix and is clean after it. Golden unchanged. Suite - 716 passed, 1526 subtests, 80 skipped. GPU numeric - wgrad async + K-outer at split_k=8 passes the torch oracle, as does the synchronous control. Forward conv async remains correct at C=4, C=12 and C=64 after the over-run guard. ## Known gap, not introduced here split_k > 1 combined with the cshuffle atomic epilogue diverges between the two engines. It reproduces with neither async nor K-outer involved, so it predates this work; it had simply never been byte-compared, because the existing split-K parity config uses the direct-store epilogue and the two cshuffle split-K configs are both rejected by the validator. It is worth its own fix: that combination is what the gfx950 dispatch candidate selects. ## Submission Checklist - [ ] Look over the contributing guidelines at https://github.com/ROCm/ROCm/blob/develop/CONTRIBUTING.md#pull-requests.

…pelines ISSUE ID: AICK-XXXX --split-k 0 means "sweep every split-K degree". For wgrad that is normally done with a single runtime-degree kernel: the degree rides a kernel argument and the degrees are swept at launch time, which avoids recompiling per degree. The async and unrolled pipelines cannot use that kernel -- they need a compile-time trip count to lay out the pipeline -- so every config was silently skipped and the sweep reported "0 compiled, 24000 skipped" with no reason. The sweep now compiles one kernel per concrete degree when the runtime-degree kernel is unavailable. Same coverage, more compiles; the non-async path is unchanged and still uses the single runtime kernel. Also reject split_k == 0 with async_dma/unroll_k in both validators, so the combination fails with its reason instead of raising a ValueError deep in the k-loop that the sweep drivers swallow into a silent skip. ## Test Plan --split-k 0 --lds-k-outer --async-dma now compiles configs where it previously compiled none; --split-k 0 --lds-k-outer alone is unchanged (still one runtime-degree kernel per config, same combo count). Byte-identity GREEN - 69/69 families; golden unchanged. Suite - 716 passed, 1526 subtests, 80 skipped. ## Submission Checklist - [ ] Look over the contributing guidelines at https://github.com/ROCm/ROCm/blob/develop/CONTRIBUTING.md#pull-requests.

…transpose read The K-outer transpose-read lane mapping stepped k by a hardcoded 8 between lane groups. That stride is the per-lane MFMA operand length: 8 for 32x32x16 and 16x16x32, but 4 for 16x16x16. The 4-element atom read past the end of its tile and produced NaN. Carry the operand length instead, in both engines. IR is unchanged for the 8-element atoms, so no golden moves. It shipped because every lds_k_outer test pinned _WARP_TILE_MN = 32 and both K-outer parity configs used 32x32x16. Added parity config 13 and test_lds_k_outer_atom_16x16x16 to cover the 4-element atom. Also: - WgradConvSpec.default_lds_k_outer selection policy + ROCKE_WGRAD_LDS_K_OUTER override. Policy only; the spec field still defaults to False. - benchmark --dtype-d, so dW can be wider than A/B. The packed 16-bit split-K atomic needs an even cpg; an fp32 dW uses a scalar atomic and does not. - _verify_kernel tolerance follows the compute dtype, not the storage dtype. - lds_k_outer added to the runbook Knob Catalog; new env flag documented.

…lds_k_pad Add the CK pipeline_basic wgrad K-loop: one LDS buffer, with the next tile's global read issued before the current tile's MFMA so VMEM latency hides behind compute. The sync loader is split into emit_global_read / emit_lds_write so the next tile waits in VGPRs instead of a second LDS tile. Also: - compv4 no longer forces double buffering. It shares the single-buffer scf.for body with mem/compv3, so the second A/B tile was a dead allocation that only shrank the LDS budget and rejected valid specs. - wgrad honours vector_size_a/vector_size_b, which it declared and ignored. They cap the free-axis width and are now rejected when inadmissible rather than silently downgraded. - lds_k_pad is tagged into the kernel name, so two pads no longer collide on one symbol and get the same cached binary. - split_k=0 is rejected for pipeline='basic' alongside async_dma/unroll_k; all three need a compile-time trip count. - basic is capped at 128 unrolled K iterations (build practicality). Mirrored in the C engine with parity coverage; golden IR hashes updated.

Deduce or delete the knobs the sweep should not expose, and restore a validation call dropped earlier on this branch. - Restore WgradConvSpec.validate()'s layout.validate_for_async() call, dropped in 73a270d and reachable again since c629128. Mirror it into is_valid_wgrad_spec so a bad layout returns a reason instead of a silent sweep skip, and add the regression test. - Gate lds_k_outer on the arch in both engines. It was accepted on gfx942 and emitted ds_read_tr16_b64, which is CDNA4 only. - Delete --lds-k-outer and ROCKE_WGRAD_LDS_K_OUTER. It is a pure function of arch/dtype/atom, so deduce it per combo via default_lds_k_outer, and point dispatch at the same predicate instead of its own divergent copy. - Delete --lds-k-pad. The K-outer row stride comes from _KOUTER_PAD, so an explicit pad changed the kernel name and the LDS budget without changing an emitted op; it is now an error in both engines. - Delete --dtype-d. The two-stage wgrad path in #10571 covers it. - Sweep async_dma instead of passing it: it is not deducible, since the pad it forces and the load-width ladder it narrows both depend on sweep axes. Pipeline is pinned on the async leg, where it is ignored. - Decide runtime split-K capability per combo. The run-level predicate tested only async_dma, so every pipeline='basic' combo was dropped from a --split-k 0 sweep without a reason string. - Rename test_async_tile_loader_decode.py; revert the .gitignore block. Byte-identity GREEN across all conv families. No golden hashes move.

…ed loops dgrad (B tile only): - Store the B (W, KYXC) LDS tile K-outer and transpose on the MFMA operand fetch with ds_read_tr16_b64, removing the transpose-on-store scatter. A (dY, NHWK) is left M-outer: it already has a stride-1 reduction axis, so its store is already wide and its read already conflict-free via lds_k_pad. - Deduced per combo by DgradConvSpec.default_lds_k_outer, not a flag. The gate is asymmetric (dtype_b/warp_tile_n) since only B flips, and carries cpg because the saving is zero when the channel run is odd. - Gated on gfx950/MFMA/wave64/16-bit B, and rejects an explicit lds_layout or async_dma rather than silently ignoring them. - Export conv_tr_frag as rocke_conv_tr_frag, taking lane and the hoisted constants instead of a build ctx, so the shared compute phase and dgrad's own fetch share one lane mapping. IR-neutral. - Fix dgrad's C kernel_name flag order to match Python (async, kouter, spk). It only diverged for a spec setting both, but kouter sits between them. - Parity configs 11 and 12 in both emitters; 12 pins the 16x16x16 atom, where the fragment length is 4 and hardcoding 8 reads past the tile end. - Six in-process A/B tests asserting bitwise equality with the M-outer path. split_k uses tolerance: the atomic epilogue is not order-deterministic and does not reproduce itself bitwise. wgrad: - Apply the K-iteration unroll cap to async_dma, not just pipeline='basic'. A deep reduction at a low split-K degree unrolled five figures of bodies into one kernel and exhausted host memory during the IR build. - Report async_dma in the benchmark's ranked table; it is a swept axis now, so the two legs were indistinguishable in the output.

yraparti · 2026-09-06T16:32:40Z

LGTM! along with dsl_docs, Please update any operator specific README/Algorithm.md to reflect the new knobs and the effect on instruction counts and performance.

Document pipeline + knobs | Fail. One Knob Catalog row for wgrad lds_k_outer. No dgrad row, no pipeline=basic, async_dma still described as pairing with compv4. -- | --

bartekxk

LGTM, one comment is that there is no gfx1250 support. Feel free to merge and create next PR with support for gfx1250

assistant-librarian · 2026-09-08T04:32:27Z

TheRock Submodule Bump Activity

Newest first

2026-09-08 — Included in TheRock via ROCm/TheRock#7975.

Extend the K-outer LDS path from #11602 to gfx1250 for both backward directions. gfx1250 (wave32 WMMA) - No new primitive: the IR op is shared and core/isa/backend.py already selects ds_load_tr16_b128 (.v8bf16/.v8f16). What differs is the lane mapping, and the distinction that matters is between the layout a lane must end up holding and the address it must supply. ds_load_tr16_b128 transposes an 8x8 element block within each group of 8 lanes: the 8 lanes each read 8 contiguous elements, and lane j of the group receives element j from all 8 runs. So to land the WMMA B layout (lane l holds column l % 16, K-half l // 16), the group addresses the 8-column block containing its target column and each lane supplies a different K row: col = mn_base + ((l % 16) // 8) * 8 row0 = k_base + (l // 16) * n + (l % 8) # read r at row0 + 8*r The read returns 8 per lane, so gfx1250's 16-element fragment is two reads. gfx950's ((l%16)//4) and (l%4)*4 terms come from the different grouping of ds_read_b64_tr_b16 under 64 lanes. - _tr_frag branches on wave_size; both formulas sit side by side, in both engines, with one const_i32(16) and one const_i32(8) bound and reused so SSA numbering does not drift. - Gating is two regimes, _LDS_K_OUTER_ARCH_WAVE = {gfx950: 64, gfx1250: 32}, each arch pinned to its wave size so a mismatched spec is rejected rather than emitting a formula the hardware lacks. The wave32 regime admits one atom, 16x16x32. - Wire the transpose read into the WMMA operand fetch in both engines; the MFMA path already had it. - wgrad parity config 15 covers the new path. Verified on gfx1250 - Numeric A/B green: K-outer is bitwise identical to the M-outer default for dgrad, and matches the fp32 reference for wgrad. - The lane map is confirmed by a per-lane probe of the intrinsic (fill LDS with an encoded (k, col), read back, decode): all 512 fragment slots land where the atom requires, with the ISA checked to confirm the compiler had not substituted a plain b128 load. - Byte-identity GREEN at llvm20 and llvm22. Tests - The K-outer numeric A/B is arch-aware: wave size and MMA family come from the target, and the gfx950 tiling is remapped to the one atom wave32 admits. Cases with no wave32 counterpart (16x16x16, split_k > 1) skip with an explicit reason. Asking for family="mma" on gfx1250 returns None and would silently skip every case, so the family must be derived. - Same for the wgrad K-outer sweep, which additionally has to request epilogue='default' (the only one WMMA wgrad accepts) and accumulate into an fp32 dW (see below). Since a test whose subTests all skip still reports *passed*, _assert_ran() makes an all-skipped sweep a hard failure rather than a false green. - CPU-only guard that the K-outer B fetch lowers to ds_load_tr16_b128 of <8 x T> from an addrspace(3) pointer, exactly twice per fragment. Docs - convolution.md: lds_k_outer in Spec Defaults, and a Backward Passes section. The file documented forward conv only and did not mention either backward direction. - optimization_runbook.md: the lds_k_outer row said "Wgrad only" and the lds_k_pad row told readers to sweep a deleted flag. Both fixed; rows added for the conv pipeline set and _MAX_UNROLLED_K_ITERS. - lds-optimization-rocke.md: new "Method 5: K-Outer Tile + Transpose Read", with the _KOUTER_PAD bank-spread derivation as a worked table. - All three record that byte-identity proves the two engines agree, never that a lane formula is right, and that a per-lane hardware probe settles a transpose-read map in one run. Note the atom-name collision: the gfx1250 WMMA 16x16x32 has n == 16, unlike the MFMA 16x16x32 (n == 8). - Drop four measured figures from convolution.md per AGENTS.md compliance; the levers and their ranking stay. Known gaps, not addressed here - No valid gfx1250 wgrad spec exists with fp16/bf16 dW: WMMA requires epilogue='default', the 16-bit-output rule requires 'cshuffle'. That rule is worded as a split_k-atomic constraint but is not guarded by _is_atomic, so it fires at split_k=1 and for M-outer too. - pipeline="wavelet" + lds_k_outer=True validates and builds, but the wavelet loader writes B M-outer while _tr_frag reads it transposed. - The dgrad b_per_lane % 4 arity guard is the wave64 condition; wave32 needs % 8. Moot at 16, latent otherwise.

…11762) ISSUE ID: AICK-2219 ## Motivation Extend the K-outer LDS path from #11602 to gfx1250 for both backward directions, and write the docs yraparti asked for. ## Technical Details gfx1250 (wave32 WMMA) - No new primitive: the IR op is shared and core/isa/backend.py already selects ds_load_tr16_b128 (.v8bf16/.v8f16). What differs is the lane mapping, and it is simpler. 32 lanes over a 16-wide atom edge give two groups that split K rather than the free axis, so lane l owns column l % 16 and K-half l // 16, and a fragment is a run of n K values at one column from (l // 16) * n. The read returns 8 per lane, so gfx1250's 16-element fragment is two reads. gfx950's ((l%16)//4) and (l%4)*4 terms exist only because 64 lanes over that edge create groups within the free axis. - _tr_frag branches on wave_size; both formulas sit side by side. - Gating is now two regimes, _LDS_K_OUTER_ARCH_WAVE = {gfx950: 64, gfx1250: 32}, with each arch pinned to its wave size so a mismatched spec is rejected rather than emitting a formula the hardware lacks. The wave32 regime admits one atom, 16x16x32. - Wire the transpose read into the WMMA operand fetch in both engines; the MFMA path already had it. - wgrad parity config 15 covers the new path. Docs - convolution.md: lds_k_outer in Spec Defaults, and a Backward Passes section. The file documented forward conv only and did not mention either backward direction. - optimization_runbook.md: the lds_k_outer row said "Wgrad only" and the lds_k_pad row told readers to sweep a deleted flag. Both fixed; rows added for the conv pipeline set and _MAX_UNROLLED_K_ITERS. - lds-optimization-rocke.md: new "Method 5: K-Outer Tile + Transpose Read", with the _KOUTER_PAD bank-spread derivation as a worked table. - Drop four measured figures from convolution.md per AGENTS.md compliance; the levers and their ranking stay. ## Submission Checklist - [ ] Look over the contributing guidelines at https://github.com/ROCm/ROCm/blob/develop/CONTRIBUTING.md#pull-requests.

--lds-k-outer was removed in #11602/#11762 under "selection, not knobs"; re-adding it to the sweep driver contradicted convolution.md, the wgrad case study and lds-optimization-rocke.md, and let a sweep measure a kernel dispatch cannot select. The sweep driver now calls DgradConvSpec.default_lds_k_outer unconditionally, as wgrad does. run_one_dgrad.py gains --kouter auto (the new default), which asks the same predicate so a plain trace captures the kernel that ships; on/off remain for the in-process A/B the numeric tests use.

github-actions Bot added project: hip-kernel-provider rocke: platform Touches rocKE platform labels Sep 2, 2026

assistant-librarian Bot added the organization: ROCm label Sep 2, 2026

aledudek changed the title ~~feat(rocke/conv): mirror wgrad K-outer LDS in the C engine~~ feat(rocke/conv): add transpose load for bwd weight and bwd Data Sep 2, 2026

aledudek changed the title ~~feat(rocke/conv): add transpose load for bwd weight and bwd Data~~ feat(rocke/conv): add transpose load for bwd weight and bwd data Sep 2, 2026

github-actions Bot added the documentation label Sep 4, 2026

aledudek force-pushed the users/aledudek/rocke-kouter branch from 115fe12 to 177c245 Compare September 4, 2026 07:23

aledudek marked this pull request as ready for review September 4, 2026 07:23

aledudek requested review from a team as code owners September 4, 2026 07:23

bartekxk reviewed Sep 4, 2026

View reviewed changes

bartekxk requested a lite review from Copilot September 4, 2026 10:22

Copilot started reviewing on behalf of bartekxk September 4, 2026 10:23

View session

Copilot AI reviewed Sep 4, 2026

View reviewed changes

aledudek changed the title ~~feat(rocke/conv): add transpose load for bwd weight and bwd data~~ feat(rocke/conv): add transpose load for bwd weight Sep 4, 2026

Aleksander Ledudek and others added 8 commits September 4, 2026 12:00

aledudek force-pushed the users/aledudek/rocke-kouter branch from 1164e82 to 96c9dba Compare September 4, 2026 17:00

aledudek changed the title ~~feat(rocke/conv): add transpose load for bwd weight~~ feat(rocke/conv): add transpose load for bwd weight and bwd data Sep 5, 2026

bartekxk approved these changes Sep 7, 2026

View reviewed changes

bartekxk merged commit d15aa0e into develop Sep 7, 2026
100 of 104 checks passed

bartekxk deleted the users/aledudek/rocke-kouter branch September 7, 2026 11:24

aledudek mentioned this pull request Sep 7, 2026

feat(rocke/conv): K-outer transpose load on gfx1250 + documentation #11762

Merged

1 task

Conversation

aledudek commented Sep 2, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Motivation

Technical Details

Test Plan

Submission Checklist

Uh oh!

therock-pr-bot Bot commented Sep 2, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

✅ All Checks Passed — Ready for Review

Uh oh!

therock-pr-bot Bot commented Sep 2, 2026

Pre-commit check failed

Uh oh!

therock-pr-bot Bot commented Sep 4, 2026

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

🟡 Changes recommended

Suppressed comments (1)

Uh oh!

Uh oh!

Uh oh!

Uh oh!

yraparti commented Sep 6, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

bartekxk left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

assistant-librarian Bot commented Sep 8, 2026

TheRock Submodule Bump Activity

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

aledudek commented Sep 2, 2026 •

edited

Loading

therock-pr-bot Bot commented Sep 2, 2026 •

edited

Loading

yraparti commented Sep 6, 2026 •

edited

Loading