Repository navigation
perf(rocke): direct wgrad kernels for dense, depthwise and small-C - #13990
Conversation
…yers on gfx950 Adds four direct weight-gradient kernel families to the direct-conv AOT sweep (--compile-all / --run-from-cache) and the LDS alias-scope lowering their DMA pipelines depend on. Kernels (library/kernels/common): - conv_direct_wgrad_tiled: dense and grouped 3x3, stride 1. 32x32 (or 32x16) wave tiles x all taps, rows DMAed into workgroup-shared LDS buffers with an XOR chunk swizzle (bank-conflict free transpose reads), dY register ring, persistent grid with XCD remap, fp32 atomic epilogue. unfold_n packs several images into one 32-wide strip for narrow feature maps; q_groups/pingpong run two wave groups half a row apart. - conv_direct_dw_wgrad: depthwise on the VALU. DwWgradLdsSpec streams X and dY rows through LDS rings (C % 8 == 0); DwWgradColsSpec puts lanes along output columns for small C (stem layers). Both store one partial slot per partition; DwWgradFoldSpec folds them into dW (no atomics, no zeroing -- same-address fp32 atomics serialize on a depthwise dW). - conv_direct_wgrad_gather: dense small-C, any stride. One GEMM per filter row with N = (s, c); the X operand is gathered from LDS, or transpose-read when S * C keeps lanes 8-byte aligned. Partial slots + fold. Sweep integration (library/kernels/conv, library/benchmarks/common): variants direct_wgrad_tiled / _dw / _dw_cols / _gather plus the fold helper (direction direct_wgrad_helper), their capability rows and knob grids, and two-step launch plans (pass into scratch, fold into dW). Lowering (platform, Python and C++ engines): kernel attr lds_alias_scopes gives every smem allocation an alias scope and tags raw.ptr.buffer.load.lds and ds.read.tr16.b64 with !alias.scope / !noalias (noalias only between allocations whose pool ranges do not overlap). Every allocation lives in one pooled LDS global, so without it the waitcnt pass drains all in-flight LDS DMA (vmcnt(0)) before each ds_read and multi-row prefetch collapses. Opt-in: no existing kernel's IR changes; check_byte_identity GREEN at llvm20 and llvm22, and both engines emit identical IR for the new kernels. Not yet included: C++ twins of the new builders, parity emitters and pytest coverage; the JIT --direction wgrad sweep still runs only DirectConvWgradSpec. Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
4bcbd47 to
948e602
Compare
✅ All Checks Passed — Ready for Review
📖 Need help? See the Policy FAQ for details on every check and how to fix failures. |
|
🎉 All checks passed! This PR is ready for review. |
The AOT grid never varied these, so each becomes the behaviour it always
had:
- conv_direct_wgrad_tiled: xcd_remap (the XCD block remap is always on)
and waves_per_eu (fixed at 2, the budget the 32x32 tile is sized for).
- DwWgradLdsSpec / WgradGatherSpec: waves_per_eu, and rows_per_seg = 0
("whole image"), which baked the probe problem's height into an AOT
binary; rows_per_seg now defaults to 32 / 16 and must be >= 1.
- DwWgradFoldSpec: unroll (fixed at 4) and dtype_d (dW is always fp32,
so the 16-bit store path goes).
The emitted IR is unchanged apart from kernel names, which lose the
xcd / u4 / fp32 tokens.
Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
… kernels Completes the direct wgrad kernels added in 948e602 with the pieces its description listed as missing. C++ engine (platform/cpp): ports of build_direct_wgrad_tiled (with pick_swizzle), build_dw_wgrad_{fold,cols,lds} and build_wgrad_gather in instances/common/conv_direct_{wgrad_tiled,dw_wgrad,wgrad_gather}.cpp, their public headers, and a private instance_conv_direct_wgrad_internal.hpp whose emission helpers take at most one emitting argument per call, so the SSA order follows Python's left-to-right evaluation. Byte identity: configs 95..116 in library/tests/parity/ conv_direct_grouped_emit.{py,c} cover the tiled kernel (wave tiles, unfold_n, row_unroll, q_groups with and without pingpong, 1x1 and 5x5), the fold, the small-C and LDS-ring depthwise kernels and the gather kernel (gathered and transpose-read X), plus two specs both engines reject. check_byte_identity GREEN at llvm20 and llvm22. Eight IR goldens conv_direct_wgrad/* added; no existing golden changes. Tests: TestDirectWgradAotKernelsCorrectness builds each kernel on its kernel-cache probe and launches it on a different shape (partial-slot kernels followed by the fold) against a torch reference, fp16 and bf16; CPU checks of the two-step plans in test_direct_kernel_cache.py. Fixes two tests the earlier commits broke: the small-C depthwise probe now picks a width the kernel takes, so its PAD 0 capability rows build (test_every_capability_row_is_built), and the new spec classes are listed in REQUIRED_FIELDS. Docs: convolution.md section, instance index, taxonomy and support matrix. Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
aledudek
left a comment
There was a problem hiding this comment.
Added a few comments regarding validation checks and tests, overall LGTM
…nd need Review fixes for the direct wgrad kernels, in both engines. - DwWgradColsSpec: the epilogue loops over outputs tid, tid + T, ... so KW * C > threads no longer leaves partial slots unwritten (a 9x9 C=8 64-thread spec had rel_err 0.62). One pass emits the old IR. New lds_bytes() (the reduction buffer) gated at 160 KiB; threads >= 64, strip >= 1, target_blocks >= 1. - DwWgradLdsSpec, WgradGatherSpec, DirectWgradTiledSpec: lds_bytes() counts the ring slots as allocated (whole DMA instructions of every wave, the gather's +16 tail, the reduction buffers); the builders take the slot sizes from the same helpers. Each gate rejects the reviewed over-budget cases (200,704 / 196,864 / 196,608 B of rings). - The LDS-DMA kernels require a known wave64 target with 16-byte buffer_load_lds: gfx942 (one dword) and gfx1151/1201/1250 (wave32) are rejected. The AOT grid admitted 288 gfx942 dw-LDS jobs per dtype; the gfx950 job list is unchanged. - Tiled rejects row_unroll < 0 (-3 ran no rows) and target_blocks < 1; gather checks wave_k16 >= 1 and waves_k in 1..16. - C++: a NULL dtype is rejected instead of defaulting to fp16. New C API: rocke_dw_wgrad_cols_lds_bytes, rocke_direct_wgrad_tiled_lds_bytes. Tests: parity configs 117..128 (the two-pass small-C epilogue, gather at PAD 0 with (1,3) / (2,3) waves, nine specs both engines reject with the same reason); test_direct_wgrad_specs.py (lds_bytes() == the built kernel's smem, boundary/arch/degenerate rejections); six GPU cases in TestDirectWgradAotKernelsCorrectness (target_blocks=2, 128 partitions, gather PAD 0 wave configs, the 9x9 C=8 t64 stem). check_byte_identity GREEN at llvm20 and llvm22 (129 configs); IR goldens unchanged. Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
Resolves conflicts in the LDS alias scopes feature (lower_llvm core/mem engines + header, both C++ and Python), where this branch and develop each independently implemented the same kernel-attr-gated optimization. Keeps develop's implementation (_lds_aa_*, pointer-keyed provenance, correctly excludes unused smem allocations, covers more call sites) and ports this branch's extra call sites (crosslane ds_read_tr16_b64, async_buffer_load_lds) onto it, removing the duplicate old API. Verified via check_byte_identity.py (GREEN, all families) and the lds_alias_scopes / mla_prefill_gfx942 test suites. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
Missed in the previous merge commit: the ds_read_tr16_b64 alias-scope call still used the old rocke_ll_lds_alias_suffix name that was removed when resolving the LDS alias scopes merge conflict in favor of develop's _lds_aa_* implementation. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
CongMa13
left a comment
There was a problem hiding this comment.
No new issues found. I agree with the following existing comments:
- Insufficient threads leave filter accumulators unwritten: Copilot (thread PRRT_kwDOOej8qc6uWSGN)
- Workspace footprint underestimates padded DMA allocation: Copilot (thread PRRT_kwDOOej8qc6uWSHC)
- Validator ignores architecture capabilities for LDS DMA: Copilot (thread PRRT_kwDOOej8qc6uWSHl)
- LDS footprint omits DMA rounding and ring-slot padding: Copilot (thread PRRT_kwDOOej8qc6uWSIK)
- Missing dtype incorrectly defaults to fp16: Copilot (thread PRRT_kwDOOej8qc6uWSIx)
- Add LDS footprint check to DwWgradColsSpec validator: aledudek (thread PRRT_kwDOOej8qc6uXV0Q)
- Add LDS footprint check to DirectWgradTiledSpec validator: aledudek (thread PRRT_kwDOOej8qc6uXWqf)
- Require row_unroll >= 0 and target_blocks >= 1: aledudek (thread PRRT_kwDOOej8qc6uXXHR)
- Add multi-segment accumulation and multi-iteration fold tests: aledudek (thread PRRT_kwDOOej8qc6uXXqA)
CongMa13
left a comment
There was a problem hiding this comment.
No new P0 or P1 issues were identified during this review. The changes are approved.
…arkocot/direct-f-opt # Conflicts: # dnn-providers/hip-kernel-provider/rocke/library/benchmarks/common/direct_kernel_sweep.py # dnn-providers/hip-kernel-provider/rocke/library/kernels/conv/direct_aot_jobs.py # dnn-providers/hip-kernel-provider/rocke/library/kernels/conv/direct_variants.py # dnn-providers/hip-kernel-provider/rocke/library/tests/parity/conv_direct_grouped_emit.c # dnn-providers/hip-kernel-provider/rocke/library/tests/parity/conv_direct_grouped_emit.py # dnn-providers/hip-kernel-provider/rocke/library/tests/test_builder_signature_contract.py
|
Automated review: Model - qwen3.6:35b-a3b-coding-mtp-q4_K_M, effort: max PR #13990 — direct wgrad kernels for dense, depthwise and small-C — 🔴 First-round: blocking — the review fixes all landed with tests and no existing kernel's IR changes, but the PR's own golden entries are stale and TheRock CI fails on all four hipkernelprovider shards
Utility: 🟢 Positive — four direct wgrad families (dense, depthwise, small-C) with validators that now match what they allocate, plus an opt-in alias-scope lowering that unblocks multi-row LDS prefetch without touching existing kernels. Legend: 🔴 blocking · 🟠 should-fix · 🟡 minor nit · 🟢 info/good · ⚪ retracted First-round review: a local qwen3.6 model made the independent first pass; the orchestrator (Claude Opus 5.5) verified it — every thread against the code and its new test, merge simulation, native build + byte identity, a full IR sweep against develop, the suites against develop, the golden at head and at the approved commit, and the CI logs — and owns this merge-worthiness call. It flags whether a closer look is needed; it is not an approval or a rejection. |
Resolve conflicts with #13120 (additive bias) and #13990 (wgrad golden): - dispatch/attention/bindings.py: keep both explicit-spec checks (LSE buffer presence, additive-bias refusal). - kernels/common/attention_unified.py: keep both imports; graph identity keys and keep-alive refs carry additive_bias (develop) and lse (this branch); lse is kept alive only when written. - rocke_representative_ir_sha256.json: develop's cases and summary plus this branch's 3 cases and its summary delta (gfx942 +2, gfx950 +1, ok_cases +3). - test_unified_lse_dispatch.py: re-pin gfx950's pre-LSE split-KV identity to develop's post-#13120 value; the LSE-off identity equals develop's.
…-cache-locality #13990 added direct wgrad cases to the representative IR golden next to this branch's hq_minor_swz cases. Re-blessed on the merged tree: every entry equals one side (the wgrad cases develop's, the rest this branch's).
TheRock Submodule Bump ActivityNewest first
|
The differential golden had gone stale, because nothing ran the check: 675 of its 1116 LLVM IR entries (compared at llvm20) and 608 of its 1089 canonical IR entries no longer match the tree. Recorded afresh in the new format, with the Python reference emitters, from a tree whose two-engine gate is green at llvm20, llvm22 and llvm23: - ll at llvm20 and llvm22: 82 families, 1474 configs per flavor. At llvm23: 83 families, 1578 configs, where gfx1250_lowering has configs that need LLVM 23 and are refused below it, and every conv_direct_dgrad_wmma config needs LLVM 23. - ir_canonical: 82 families, 1561 configs (conv_abi's emitter has no IR mode). Canonical IR was checked to be identical at llvm20 and llvm23. What the re-bless absorbs: every mismatch that could be traced is an intended code generation change that did not re-bless this file -- the shared LDS pool (#8844, about 425 configs), the ds.swizzle declaration (#9583, about 177), renamed values in the attention transpose reads (#9133, 57), the smaller LDS pool (#10112, about 39), the flat work group size of one-wave wave32 kernels (#12710, about 39), #9220 (6) and the conv changes of #13109 (6). 441 entries were already stale at the start of the history available here and cannot be attributed; no unexplained change was found after it. The two repeated keys (conv_direct_grouped, conv_wgrad_workspace_reduce), which hid their newer copies, are gone. It also takes in the code generation changes merged since then, each traced by recording the file at the commit and at its parent; the count is per section unless a flavor is named: - the magic division sum emitted without signed-overflow poison (#13802): 215 configs in 19 families -- attention_unified 6, attention_unified_stress 32, conv_direct_grouped 16, conv_implicit_gemm 26, conv_implicit_gemm_dgrad 31, conv_implicit_gemm_wgrad 42, deep_fused_conv_pool 2, gfx1201_deep_fused_conv_pool 5, gfx950_deep_fused_conv_pool 5, dynamic_helpers 2, fmha_splitkv_decode 6, img2col 5, moe_fused_mega 5, moe_gemm_fused 6, moe_sorting 6, pooling 4, sage_attention 4, sparse_attention 6, transpose 6; - attention decoding on gfx942 and gfx950 (#13492): 23 changed and 25 new configs in attention_strided_kv, gfx942_attention_tiled_3d and gfx950_attention_tiled_3d; - the f64 follow-ups (#13887): 3 changed and 1 new elementwise config, 5 changed reduce configs; - MX scale and fp6 decoding (#13885): 6 mx_gemm configs; - the gfx950 conv dgrad speed-up (#13075): 36 new conv_direct_dgrad_ext configs; 22 changed, 15 new and 2 removed conv_implicit_gemm_dgrad configs; 2 changed and 3 new conv_direct_grouped configs; 2 new dynamic_helpers and 3 new target_intrinsics configs; - removing the ds.read.tr16.b128 declaration (#14054): gfx1250_lowering, 100 changed and 2 removed configs at llvm23 and in canonical IR, 5 changed, 2 new and 4 removed at llvm20 and llvm22; - optional LSE output on unified split-KV attention (#13955): 18 gfx942_attention_tiled_3d and 16 gfx950_attention_tiled_3d configs per LLVM flavor, canonical IR unchanged; - the running softmax max in 3D split-KV (#14276): 14 attention_strided_kv, 18 gfx942_attention_tiled_3d and 16 gfx950_attention_tiled_3d configs; - depthwise 7x7 and grouped wgrad on batched MFMA (#14157): 7 changed, 10 new and 1 removed conv_direct_grouped configs; - the retirement of wmma_fmha_fwd (#13897): the gfx1151_wmma_fmha_fwd (157 configs) and gfx1201_wmma_fmha_fwd (12) families are gone; - new configs, every existing hash unchanged: the barrier_fence family (#13575, 10), one conv_implicit_gemm_dgrad config (#13782), 14 conv_direct_grouped configs for the split-C and grouped LDS-halo kernels (#13993), one config each in fmha_splitkv_decode, sage_attention and sparse_attention (#13876), 23 conv_direct_grouped configs for the direct wgrad kernels (#13990), 2 conv_implicit_gemm_dgrad configs for patch-embedding dgrad (#13988), 6 f64_mma configs (#13794), 5 conv_direct_grouped and 7 conv_implicit_gemm configs for depthwise and grouped conv forward (#13987), and the conv_direct_dgrad_wmma family (7 configs, llvm23 and canonical IR only, #14156). #13885, #13955 and #14276 updated the old file's entries for their changes; the new llvm20 and canonical IR sections reproduce every one of those hashes. Co-Authored-By: Claude <noreply@anthropic.com>
The differential golden had gone stale, because nothing ran the check: 675 of its 1116 LLVM IR entries (compared at llvm20) and 608 of its 1089 canonical IR entries no longer match the tree. Recorded afresh in the new format, with the Python reference emitters, from a tree whose two-engine gate is green at llvm20, llvm22 and llvm23: - ll at llvm20 and llvm22: 82 families, 1474 configs per flavor. At llvm23: 83 families, 1578 configs, where gfx1250_lowering has configs that need LLVM 23 and are refused below it, and every conv_direct_dgrad_wmma config needs LLVM 23. - ir_canonical: 82 families, 1561 configs (conv_abi's emitter has no IR mode). Canonical IR was checked to be identical at llvm20 and llvm23. What the re-bless absorbs: every mismatch that could be traced is an intended code generation change that did not re-bless this file -- the shared LDS pool (#8844, about 425 configs), the ds.swizzle declaration (#9583, about 177), renamed values in the attention transpose reads (#9133, 57), the smaller LDS pool (#10112, about 39), the flat work group size of one-wave wave32 kernels (#12710, about 39), #9220 (6) and the conv changes of #13109 (6). 441 entries were already stale at the start of the history available here and cannot be attributed; no unexplained change was found after it. The two repeated keys (conv_direct_grouped, conv_wgrad_workspace_reduce), which hid their newer copies, are gone. It also takes in the code generation changes merged since then, each traced by recording the file at the commit and at its parent; the count is per section unless a flavor is named: - the magic division sum emitted without signed-overflow poison (#13802): 215 configs in 19 families -- attention_unified 6, attention_unified_stress 32, conv_direct_grouped 16, conv_implicit_gemm 26, conv_implicit_gemm_dgrad 31, conv_implicit_gemm_wgrad 42, deep_fused_conv_pool 2, gfx1201_deep_fused_conv_pool 5, gfx950_deep_fused_conv_pool 5, dynamic_helpers 2, fmha_splitkv_decode 6, img2col 5, moe_fused_mega 5, moe_gemm_fused 6, moe_sorting 6, pooling 4, sage_attention 4, sparse_attention 6, transpose 6; - attention decoding on gfx942 and gfx950 (#13492): 23 changed and 25 new configs in attention_strided_kv, gfx942_attention_tiled_3d and gfx950_attention_tiled_3d; - the f64 follow-ups (#13887): 3 changed and 1 new elementwise config, 5 changed reduce configs; - MX scale and fp6 decoding (#13885): 6 mx_gemm configs; - the gfx950 conv dgrad speed-up (#13075): 36 new conv_direct_dgrad_ext configs; 22 changed, 15 new and 2 removed conv_implicit_gemm_dgrad configs; 2 changed and 3 new conv_direct_grouped configs; 2 new dynamic_helpers and 3 new target_intrinsics configs; - removing the ds.read.tr16.b128 declaration (#14054): gfx1250_lowering, 100 changed and 2 removed configs at llvm23 and in canonical IR, 5 changed, 2 new and 4 removed at llvm20 and llvm22; - optional LSE output on unified split-KV attention (#13955): 18 gfx942_attention_tiled_3d and 16 gfx950_attention_tiled_3d configs per LLVM flavor, canonical IR unchanged; - the running softmax max in 3D split-KV (#14276): 14 attention_strided_kv, 18 gfx942_attention_tiled_3d and 16 gfx950_attention_tiled_3d configs; - depthwise 7x7 and grouped wgrad on batched MFMA (#14157): 7 changed, 10 new and 1 removed conv_direct_grouped configs; - the retirement of wmma_fmha_fwd (#13897): the gfx1151_wmma_fmha_fwd (157 configs) and gfx1201_wmma_fmha_fwd (12) families are gone; - new configs, every existing hash unchanged: the barrier_fence family (#13575, 10), one conv_implicit_gemm_dgrad config (#13782), 14 conv_direct_grouped configs for the split-C and grouped LDS-halo kernels (#13993), one config each in fmha_splitkv_decode, sage_attention and sparse_attention (#13876), 23 conv_direct_grouped configs for the direct wgrad kernels (#13990), 2 conv_implicit_gemm_dgrad configs for patch-embedding dgrad (#13988), 6 f64_mma configs (#13794), 5 conv_direct_grouped and 7 conv_implicit_gemm configs for depthwise and grouped conv forward (#13987), and the conv_direct_dgrad_wmma family (7 configs, llvm23 and canonical IR only, #14156). #13885, #13955 and #14276 updated the old file's entries for their changes; the new llvm20 and canonical IR sections reproduce every one of those hashes. Co-Authored-By: Claude <noreply@anthropic.com>

Adds four direct weight-gradient kernel families to the direct-conv AOT sweep (--compile-all / --run-from-cache) and the LDS alias-scope lowering their DMA pipelines depend on.
Kernels (library/kernels/common):
Sweep integration (library/kernels/conv, library/benchmarks/common): variants direct_wgrad_tiled / _dw / _dw_cols / _gather plus the fold helper (direction direct_wgrad_helper), their capability rows and knob grids, and two-step launch plans (pass into scratch, fold into dW).
Lowering (platform, Python and C++ engines): kernel attr lds_alias_scopes gives every smem allocation an alias scope and tags raw.ptr.buffer.load.lds and ds.read.tr16.b64 with !alias.scope / !noalias (noalias only between allocations whose pool ranges do not overlap). Every allocation lives in one pooled LDS global, so without it the waitcnt pass drains all in-flight LDS DMA (vmcnt(0)) before each ds_read and multi-row prefetch collapses. Opt-in: no existing kernel's IR changes; check_byte_identity GREEN at llvm20 and llvm22, and both engines emit identical IR for the new kernels.
JIRA ID: AICK-2372