fix(rocke): close conv perf gaps after the switch to AOT by jakpiase · Pull Request #13320 · ROCm/rocm-libraries

jakpiase · 2026-10-08T15:52:54Z

Motivation

After #12783 (JIT → AOT), many conv shapes got slower. In some cases the AOT cache had no kernel at all. The causes:

The cache left out kernels the JIT sweep used to pick.
The shape-generic kernels compute index math with runtime divides inside the K loop.

Technical Details

Cache coverage

Direct-conv wgrad is now cached. DirectConvWgradSpec.runtime_channels makes one binary serve any channel count.
Direct MFMA dgrad for dense problems is now cached, through runtime-channel modes of DirectConvSpec and its transpose/reorganize helpers. The C++ engine gets ports of these.
Implicit-GEMM wgrad and dgrad entries now use the K-outer LDS layout, the same as the JIT sweep picks.
Non-split wgrad (split_k=1) is cached, and split-K degrees 256/512 are tried at launch.
13x13 and 31x31 depthwise forward filters are cached.

New kernel loop options (off by default, so existing IR is unchanged)

WgradConvSpec.prefetch_addr: computes the next tile's addresses while the current loads are in flight.
DgradConvSpec.uniform_k_decode: moves the per-element divides out of the K loop.
Both are turned on by deterministic, per-geometry selection rules.

Benchmark tooling

Fixed a device-memory leak in --run-from-cache (two-stage workspaces and dgrad sub-GEMM tables).
--split-k-prune now works on the cache path.
Early stop drops hopeless kernels after a single launch.
The compare script no longer skips direct conv for wgrad.

The follow-up commit addresses self-review findings: C++ bounds and validation, dgrad cache coverage, early-stop robustness, and extra tests.

Cache size grows about 16% per dtype. Most of it is the non-split wgrad entries.

Test Plan

check_byte_identity.py: llvm20 and llvm22. New parity configs cover each new mode.
platform/tests/run_all.py with ctest.
CPU: test_conv_abi, test_direct_kernel_cache, test_early_stop, test_conv_wgrad_prefetch_addr, test_conv_dgrad_uniform_k_decode, layering/signature tests.
GPU (gfx950): test_conv_wgrad_correctness, test_conv_dgrad_correctness, test_direct_conv_correctness, test_conv_fwd_correctness.
--run-from-cache --verify on sampled wgrad, dgrad and depthwise shapes.

Test Result

Byte identity is GREEN for all 73 families on both LLVM flavors.
ctest passes 14/14.
All the listed tests pass.
Every new kernel mode passes --verify against the torch reference.
run_all.py reports two failures that also occur on the base commit, unrelated to this change: the relative-path guard on platform/CMakeLists.txt, and test_parity_matrix for tf32 on gfx942.

Performance (relative only).

How it was measured:

Setup: gfx950, NHWC, bf16 and fp16. Shapes are 297 forward, wgrad and dgrad problems from three model shape lists.
Baseline: best JIT result before feat(rocke): switch conv kernels to AOT #12783.
New side: best kernel from this branch's AOT cache, measured with --run-from-cache at 3 warmup / 10 iterations.
Correctness: the new kernel modes passed --verify against a reference. The other cached kernels were not re-verified in the timing runs.

Results:

wgrad shapes went from typically 2–3× slower than the JIT baseline (worst case about 6×) to parity on most measured configurations.
dgrad shapes that were 1.1–1.3× slower are about level with the baseline.
The shape-generic kernel loops alone got 1.05–1.3× faster on the measured configurations.

Caveats:

At this iteration count some rows varied by up to 1.6× between runs.
A few remaining gaps are outside this PR:
- Direct depthwise kernels are slower in the JIT path too, which is being addressed separately.
- dgrad with stride > 2, and direct dgrad with an odd number of 16-channel dY tiles (80, 112, …), are not yet cached.

JIRA-ID : AICK-2519

The AOT kernel cache left out kernels the JIT sweep used to pick, and the shape-generic kernels carry runtime index math on the K-loop critical path. Cache coverage: - Direct-conv wgrad is now cached. DirectConvWgradSpec.runtime_channels reads cpg/kpg from kernargs, so one binary serves any channel count. - Dense (and any cpg/kpg) direct MFMA dgrad is now cached through runtime-channel modes of DirectConvSpec and its transpose/reorganize helpers. The C++ engine gains ports of the generic direct kernel and the dgrad weight transforms. - Implicit-GEMM wgrad entries use the K-outer LDS layout picked by default_lds_k_outer; dgrad uses it for vec_b > 1 binaries. - The non-split (split_k=1) wgrad kernel is cached (except group-merged), and split-K degrees 256/512 are tried at launch. - Large depthwise forward filters (13x13, 31x31; stride 1, same padding) are cached. Kernel loops (both default off, so existing IR is unchanged): - WgradConvSpec.prefetch_addr computes the next tile's load offsets while the current loads are in flight; enabled by default_prefetch_addr. - DgradConvSpec.uniform_k_decode carries the tile's filter tap as uniform loop state when kpg is a multiple of tile_k, hoisting the per-element divides out of the K loop; enabled by default_uniform_k_decode. Benchmark tooling: - --run-from-cache frees two-stage workspaces and dgrad sub-GEMM tables as it goes; they leaked per (kernel, degree) and could exhaust device memory. - --split-k-prune applies to --run-from-cache. - EarlyStop probes one launch first and skips kernels far past the stop bound without the rest of their warmup, also before --early-stop-after. - benchmark_conv_compare.py no longer skips direct conv for wgrad. Python and C++ engines emit byte-identical IR for every family (llvm20 and llvm22); parity configs cover the new modes. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

- wgrad prefetch_addr (C++): size the K-loop iteration arguments for every loader plan (accumulators plus one offset per staged A/B vector). Scalar loads on a 128x128x64 tile overflowed the accumulator cap, so the C++ engine rejected a spec the Python engine builds. Parity config 38 (C = K = 3 on that tile) pins it. - Direct conv (C++): stop building after a failed validation and when the LDS loader passes exceed chunk_meta[], instead of dividing by an unvalidated wave count or writing past the array. - Runtime-channel MFMA dgrad cache: add a 2-way split of the 16-channel atoms to the shared dgrad grid, so even atom counts 4 does not divide get a binary; document the odd counts that still have none. - EarlyStop: the one-launch hopeless probe also needs an absolute excess over the bound (a cold first launch next to a tiny best), is skipped for kernels that failed verification, and its report names the probe. - --run-from-cache: report when depthwise filtering leaves no candidate, free the per-case buffers on the verify-stop path; refresh the --split-k, --split-k-prune and kernel_sweep docs for the new degrees, the non-split wgrad kernel and the K-outer/prefetch/uniform-decode policies. - Docstrings for the prefetch_addr / uniform_k_decode policies state the rationale instead of measured outcomes; runbook gains a prefetch_addr row. - Tests: dgrad uniform_k_decode split-K past gemm_k, tilde stride-2 split-K, grouped fallback (K a tile multiple, kpg not) and grouped split-K; pinned K-outer/prefetch expectations for the AOT job generators; EarlyStop absolute-excess and failed-kernel cases. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

therock-pr-bot · 2026-10-08T15:56:33Z

✅ 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-10-08T16:11:45Z

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

Copilot

🟡 Changes recommended

The C++ kernel-name append can write out of bounds, and new validation paths silently normalize malformed dtype input.

3 open findings

Reject missing dtype before constructing kernel identity · New
Prevent snprintf length overflow before append pointer arithmetic · New
Remove absolute performance cutoff from user-facing help · New

What changed in this PR

Closes convolution AOT performance gaps by expanding cache coverage, adding optimized K-loop modes, and improving benchmark reliability.

Changes:

Adds runtime-channel direct wgrad/dgrad cache support.
Introduces wgrad address prefetch and dgrad uniform K decoding.
Improves split-K coverage, pruning, early stopping, cleanup, and validation tests.

File	Description
`platform/python/rocke/helpers/loads.py`	Adds split address-planning/load helpers.
`platform/dsl_docs/optimization/optimization_runbook.md`	Documents new loop modes.
`platform/cpp/instances/common/conv_implicit_gemm_wgrad.cpp`	Ports address-prefetch wgrad loop.
`platform/cpp/instances/common/conv_implicit_gemm_dgrad.cpp`	Ports uniform dgrad K decoding.
`platform/cpp/instances/common/conv_direct_grouped_spec_and_validate.cpp`	Adds generic direct-conv specification support.
`platform/cpp/instances/common/conv_direct_grouped_build_wgrad.cpp`	Adds runtime-channel wgrad lowering.
`platform/cpp/include/rocke/instance_conv_implicit_gemm_wgrad.h`	Exposes the prefetch flag.
`platform/cpp/include/rocke/instance_conv_implicit_gemm_dgrad.h`	Exposes the uniform-decode flag.
`platform/cpp/include/rocke/instance_conv_direct_grouped.h`	Declares runtime-channel direct kernels.
`platform/cpp/include/rocke/instance_conv_direct_grouped_internal.h`	Declares dynamic KRSC descriptors.
`platform/cpp/include/rocke/helper_rocke.helpers.loads.h`	Declares split loader APIs.
`platform/cpp/helpers/loads.cpp`	Implements split loader APIs.
`platform/cpp/bindings/rocke_engine.cpp`	Binds the wgrad prefetch field.
`library/tests/test_early_stop.py`	Tests one-launch early stopping.
`library/tests/test_direct_kernel_cache.py`	Tests expanded direct-cache coverage.
`library/tests/test_direct_conv_correctness.py`	Tests runtime-channel direct kernels.
`library/tests/test_conv_wgrad_prefetch_addr.py`	Tests prefetch selection and IR.
`library/tests/test_conv_wgrad_correctness.py`	Adds prefetch numerical coverage.
`library/tests/test_conv_dgrad_uniform_k_decode.py`	Tests uniform decode selection and IR.
`library/tests/test_conv_dgrad_correctness.py`	Adds uniform-decode numerical coverage.
`library/tests/test_conv_abi.py`	Tests runtime-channel ABI invariance.
`library/tests/parity/conv_implicit_gemm_wgrad_emit.py`	Expands Python wgrad parity cases.
`library/tests/parity/conv_implicit_gemm_wgrad_emit.c`	Expands C wgrad parity cases.
`library/tests/parity/conv_implicit_gemm_dgrad_emit.py`	Expands Python dgrad parity cases.
`library/tests/parity/conv_implicit_gemm_dgrad_emit.c`	Expands C dgrad parity cases.
`library/tests/parity/conv_direct_grouped_emit.py`	Adds Python direct-kernel parity cases.
`library/tests/parity/conv_direct_grouped_emit.c`	Adds C direct-kernel parity cases.
`library/kernels/common/README_conv_implicit_gemm_dgrad.md`	Documents uniform K decoding.
`library/kernels/common/conv_implicit_gemm_wgrad.py`	Implements address-prefetch selection and lowering.
`library/kernels/common/conv_implicit_gemm_dgrad.py`	Implements uniform K decoding.
`library/benchmarks/common/kernel_sweep.py`	Expands cached kernels and policies.
`library/benchmarks/common/kernel_cache.py`	Adds new codegen fields to identities.
`library/benchmarks/common/early_stop.py`	Adds hopeless-kernel probing.
`library/benchmarks/common/benchmark_implicit_gemm_conv.py`	Improves split-K cache benchmarking.
`library/benchmarks/common/benchmark_direct_conv.py`	Adds cached direct-wgrad benchmarking.
`library/benchmarks/common/benchmark_conv_compare.py`	Enables cached direct-wgrad comparisons.

🧠 Review effort: Balanced

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

…ke-aot-wgrad-perf Conflicts resolved: - kernel_cache.py: take develop's lossless label() (kernel_label.encode); kernel_label gains baseline entries and pfa / ukd tags for the prefetch_addr and uniform_k_decode identity fields, and the tests that checked short_label() use label(). - conv_direct_grouped_emit.{c,py}: develop keeps indices 57..70 (tiled, dilated non-grouped, must-reject); this branch's runtime-channel wgrad configs move from 57..60 to 71..74 and the DirectConvSpec / dgrad weight-transform configs from 61..80 to 75..94. - conv_direct_grouped_spec_and_validate.cpp: keep both new sections (DirectDepthwiseTiledSpec from develop, DirectConvSpec from this branch). - direct_kernel_sweep.py: keep both the large depthwise filters and develop's tiled / dilated capability rows. Merge fixes: - direct_job_for_identity: pass knobs to probe_problem (the runtime-channel MFMA dgrad probe needs them) and find a helper's pipeline row among both MFMA dgrad variants, not only the baked one. - test_large_depthwise_filters_get_the_tiled_kernel: at stride 1 the column-streamed kernel now competes with the tiled one. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

- DirectConvSpec kernel_name / validate / is_valid_spec reject a missing problem dtype instead of reading it as fp16, so a malformed spec cannot take a valid fp16 spec's kernel identity. The generic and weight-transform builders pass the dtype through unchanged; rocke_b_io_ir_type already rejects NULL with the Python error text. - rocke_dgrad_conv_spec_kernel_name: check snprintf's wanted length before offsetting into the buffer, and check every flag append, so a name that fills the buffer returns ROCKE_ERR_VALUE instead of writing out of bounds. - --early-stop-after help: describe the hopeless-kernel cutoff qualitatively instead of printing its thresholds. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

therock-pr-bot · 2026-10-09T12:21:50Z

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.

Copilot

🟡 Changes recommended

Grouped C++ dgrad decoding is incorrect, and scratch reclamation can race pending launches after timing failures.

4 open findings

Drain stream before freeing buffers after launch exceptions · New
Handle grouped channel offsets in C++ uniform-k decode · New
Prevent snprintf length overflow before append pointer arithmetic
Reject missing dtype before constructing kernel identity

1 resolved since last review

Remove absolute performance cutoff from user-facing help

🧠 Review effort: Balanced

Grouped dgrad (groups > 1) landed in the Python builder only, so the C++ engine built every grouped spec -- uniform_k_decode or not -- with ungrouped addressing. Port the Python grouped terms throughout: - group index on blockIdx.y with g*kpg / g*cpg channel bases; - dY / W descriptors and the uniform-k decode loops decode by kpg, add the output-channel group base, and stride W by cpg (per-group packed KYXC); - every dX epilogue adds the input-channel group base; - load / store widths, sub-GEMM enumeration and dg_N / dg_K are per group; - the validator rejects depthwise, requires even cpg for split-K packed atomics, and rejects group counts that do not divide C and K. Two ungrouped divergences these configs exposed are fixed too: the tilde cshuffle store width was fixed at 8 instead of following cpg, and the tilde WMMA direct epilogue emitted ops in a different order than Python. Parity configs 21..33 cover grouped dgrad (MFMA and WMMA, stride 1 and 2, direct / cshuffle / atomic epilogues, uniform-k decode, K-outer, split-K, 1x1, and two must-reject specs). Ungrouped IR is unchanged. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

time_launches and launch_kernel's timing loop drained only on success, so when a launch raised, the caller's handler could free scratch (a two-stage workspace, a dgrad sub-GEMM table) that a queued kernel still used. Both now drain the device before the exception propagates; a failed drain does not mask the launch error and is attached to it as a note. Interrupts are not drained, as a hanging kernel would block the drain for good. The AOT conv sweeps also drain in their own handlers before scratch is reclaimed, and a launch error during --verify counts as a failed kernel instead of aborting the sweep. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

Copilot

🟡 Changes recommended

Timing-event cleanup can leak the second event and mask the original launch failure when destruction fails.

1 open finding

Preserve launch errors while independently cleaning up both events · New

4 resolved since last review

Handle grouped channel offsets in C++ uniform-k decode
Drain stream before freeing buffers after launch exceptions
Prevent snprintf length overflow before append pointer arithmetic
Reject missing dtype before constructing kernel identity

🧠 Review effort: Balanced

The wgrad sweep timed each split-K binary at a fixed list, _SPLIT_K_AUTO plus a hard-coded (512, 256). Derive the extra degrees instead: the degree select_split_k_wgrad picks for the tile config (what dispatch launches) and the powers of two above it up to 4x (_WGRAD_SPLIT_K_OVERSUBSCRIBE). The formula fills the device once (~2 CTAs per CU). Long-K, tiny-dW wgrads are latency-bound per CTA and run faster split further: on N42 C3->K24 480x640 the formula's 256 is 3.1x slower than 512 for the winning t32x16x64 binary, and on the g3 11x11 shape 168 is 7% slower. Both shapes pick 512 again; shapes with a small formula degree sweep the same degrees as before. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

Drop the extra launch-time degrees (512, 256 and the formula-derived oversubscription that replaced them) and go back to develop's _SPLIT_K_AUTO. Large-K shapes that win above 128 are a gap in the split-K auto-deduction (select_split_k_wgrad), which dispatch uses too; that is tracked separately instead of being patched in the benchmark. This reverts commit 4b8fe14. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

…ke-aot-wgrad-perf Conflicts with develop's stream-K wgrad (#13151) and dgrad AOT fixes (#13494), resolved by keeping both sides. Follow-on fixes for how the two interact: - wgrad: prefetch_addr is rejected with stream-K (Python and C), since stream-K never reaches the K loop it rewrites; the JIT sweep only sets it on non-stream-K combos. - wgrad parity emitters: the prefetch_addr configs move from 21-38 to 31-48, after develop's stream-K configs 21-30. - dgrad: the uniform_k_decode A decode uses the record's magic pairs, as develop's plain loop now does, and the C kernel's kpg divisor is c_kdiv (kpg when grouped). The AOT sweep keeps develop's lds_k_outer predicate (both sides had added an equivalent one). - run-from-cache: _cached_bind registers its scratch in live_extras, so the pruning/free logic keeps working; depthwise problems keep stream-K kernels; an early stop also frees the stream-K workspace. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

assistant-librarian · 2026-10-14T11:25:20Z

TheRock Submodule Bump Activity

Newest first

2026-10-14 — Included in TheRock via ROCm/TheRock#8986.

## Motivation After ROCm#12783 (JIT → AOT), many conv shapes got slower. In some cases the AOT cache had no kernel at all. The causes: 1. The cache left out kernels the JIT sweep used to pick. 2. The shape-generic kernels compute index math with runtime divides inside the K loop. ## Technical Details **Cache coverage** - **Direct-conv wgrad** is now cached. `DirectConvWgradSpec.runtime_channels` makes one binary serve any channel count. - **Direct MFMA dgrad for dense problems** is now cached, through runtime-channel modes of `DirectConvSpec` and its transpose/reorganize helpers. The C++ engine gets ports of these. - **Implicit-GEMM wgrad and dgrad** entries now use the K-outer LDS layout, the same as the JIT sweep picks. - **Non-split wgrad** (split_k=1) is cached, and split-K degrees 256/512 are tried at launch. - **13x13 and 31x31 depthwise forward** filters are cached. **New kernel loop options** (off by default, so existing IR is unchanged) - `WgradConvSpec.prefetch_addr`: computes the next tile's addresses while the current loads are in flight. - `DgradConvSpec.uniform_k_decode`: moves the per-element divides out of the K loop. - Both are turned on by deterministic, per-geometry selection rules. **Benchmark tooling** - Fixed a device-memory leak in `--run-from-cache` (two-stage workspaces and dgrad sub-GEMM tables). - `--split-k-prune` now works on the cache path. - Early stop drops hopeless kernels after a single launch. - The compare script no longer skips direct conv for wgrad. The follow-up commit addresses self-review findings: C++ bounds and validation, dgrad cache coverage, early-stop robustness, and extra tests. Cache size grows about 16% per dtype. Most of it is the non-split wgrad entries. ## Test Plan - `check_byte_identity.py`: llvm20 and llvm22. New parity configs cover each new mode. - `platform/tests/run_all.py` with ctest. - CPU: `test_conv_abi`, `test_direct_kernel_cache`, `test_early_stop`, `test_conv_wgrad_prefetch_addr`, `test_conv_dgrad_uniform_k_decode`, layering/signature tests. - GPU (gfx950): `test_conv_wgrad_correctness`, `test_conv_dgrad_correctness`, `test_direct_conv_correctness`, `test_conv_fwd_correctness`. - `--run-from-cache --verify` on sampled wgrad, dgrad and depthwise shapes. ## Test Result - Byte identity is GREEN for all 73 families on both LLVM flavors. - ctest passes 14/14. - All the listed tests pass. - Every new kernel mode passes `--verify` against the torch reference. - `run_all.py` reports two failures that also occur on the base commit, unrelated to this change: the relative-path guard on `platform/CMakeLists.txt`, and `test_parity_matrix` for `tf32` on gfx942. **Performance** (relative only). How it was measured: - **Setup:** gfx950, NHWC, bf16 and fp16. Shapes are 297 forward, wgrad and dgrad problems from three model shape lists. - **Baseline:** best JIT result before ROCm#12783. - **New side:** best kernel from this branch's AOT cache, measured with `--run-from-cache` at 3 warmup / 10 iterations. - **Correctness:** the new kernel modes passed `--verify` against a reference. The other cached kernels were not re-verified in the timing runs. Results: - wgrad shapes went from typically 2–3× slower than the JIT baseline (worst case about 6×) to parity on most measured configurations. - dgrad shapes that were 1.1–1.3× slower are about level with the baseline. - The shape-generic kernel loops alone got 1.05–1.3× faster on the measured configurations. Caveats: - At this iteration count some rows varied by up to 1.6× between runs. - A few remaining gaps are outside this PR: - Direct depthwise kernels are slower in the JIT path too, which is being addressed separately. - dgrad with stride > 2, and direct dgrad with an odd number of 16-channel dY tiles (80, 112, …), are not yet cached. JIRA-ID : AICK-2519 --------- Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>

…m kernargs At stride 1 the tilde decomposition is the identity in both spatial axes (y_tilde = x_tilde = 1 whatever the dilation, pad, filter or grouping), so there is exactly one sub-GEMM and every field of its record is an existing kernarg or a constant: [0, num_pid_m, num_pid_n, dg_M, dg_K, Hi, Wi, pH, pW, Y, X, -dH, -dW, 1, 0, 1, 0, 1, 0, 1, 0, roundup(dg_K, tile_k * split_k)] DgradConvSpec.folds_sub_gemm_record (stride_h == stride_w == 1, read off the spec's problem; for an AOT binary that is its stride capability, which the probe and every served problem share) makes the builder bind those values instead of emitting the binary search over the record buffer and the 21 dependent record loads behind it; the unit strides and zero offsets fold into the address arithmetic, and local_flat drops the block_start subtraction. The magic-division table (#13320) is still read, from its one entry right after the one record: eight loads at constant offsets that depend on nothing. The fold never keys on a sub-GEMM count evaluated on a problem, and no problem extent is read at build time. ABI, identity, label, sweep and dispatch are unchanged; sub_gemm_buf and num_sub_gemms stay in the signature and the buffer stays mandatory (its magic table is read). The fold covers the JIT ungrouped pointwise path too: its record is the same function of the kernargs, so folding it is a straight substitution (the PR's exclusion was about constant trip counts, which runtime extents no longer give). Not a knob: the fold is the stride-1 code path, default on (waiver of the "new behaviour is a default-off knob" rule of KERNEL_AUTHORING.md: there is nothing to switch off -- the folded and the loaded record are equal by construction, test_dgrad_stride1_record_is_a_function_of_the_kernargs). Every stride-1 dgrad kernel changes IR; existing caches hold no stride-1 binaries (the AOT grid builds the tilde kernel only), so no cache entry goes stale. C++: the same in _build_tilde_dgrad (_dgrad_folds_sub_gemm_record), field for field and op for op. Parity (conv_implicit_gemm_dgrad): configs 0, 4 (1x1 s1, ungrouped pointwise) and 21-24 now emit the folded record in both engines; new 38 (s1 d2: folded record with the tilde epilogue), 39 (grouped 5x5 pad 0, K-outer + uniform_k_decode + cshuffle + explicit widths, the AOT shape), 40 (s1 split-K: folded gemm_k_padded), 41 (gfx942 grouped 1x3 pad (0,1)). Golden: changed conv_dgrad/{gfx90a,gfx942,gfx950,gfx1151}/n1h8c16k32r3_s1 and conv_dgrad/{gfx942,gfx950,gfx1151,gfx1250}/n1h8c64k64r3_g4 (stride-1 grouped) at llvm20/22/23; every *_s2 case and everything else unchanged. Tests: - test_conv_abi: the record-equivalence proof (packed record of the one sub-GEMM == the record rebuilt from launch values, over stride-1 problems x dilation 1/2 x pads x filters x groups x tiles x split-K); the IR reads only the eight magic-table entries and runs no search at stride 1 (grouped, dilated) while stride 2 still searches and loads; a stride-1 shape-invariance test that varies C/K, Y/X and pad. - test_conv_dgrad_correctness: TestConvDgradStride1Specializations, re-done for AOT: stride-1 binaries built once against the AOT probe (dilation 1 and 2) and the tilde binary (stride-2 probe), all launched through ConvArgs launch values on adversarial stride-1 shapes (odd H/W, N=1, 5x5 / 1x3 filters, pad 0/1/2 and asymmetric, partial M/N tiles, kpg not a multiple of tile_k, grouped, dilation 2) with dX poisoned: bit-identical to the tilde binary and within tolerance of the torch reference; plus split-K against the reference.

The AOT dgrad grid built the tilde kernel only (stride and dilation at run time), so the stride-1 record fold had no binary to live in. _dgrad_caps now yields a second capability per configuration, stride 1 / dilation 1 (CACHE_DGRAD_STRIDE1): its probe problem is stride 1 / dilation 1, so the binary folds the sub-GEMM record and takes the unit-stride epilogue, and its identity records stride_h/w = dilation_h/w = 1, so supports_problem (the existing exact match on non-zero baked fields) offers it to stride-1, dilation-1 problems only. The tilde binary of the same configuration keeps serving every stride, stride 1 included. count_jobs follows _dgrad_caps; job_for_identity round-trips both. K-outer: the plan's K-outer axis is not added. Since #13320 every dgrad geometry already takes the layout DgradConvSpec.default_lds_k_outer picks (K-outer wherever the arch, dtype, atom, pipeline and a B load wider than one element allow it), so K-outer dgrad binaries are in the grid; a (False, True) axis would only add M-outer twins of them. The K-outer fragment rule the builder enforces is in the validator since the validator-hardening commit. AOT catalog (D14), --compile-all --direction dgrad, valid identities: gfx950 bf16 and fp16 56064 -> 112128 each (count_jobs 85376 -> 170752), gfx942 bf16 and fp16 65208 -> 130416 each. Tests: test_aot_dgrad_builds_tilde_and_stride1_binaries (probes, identity fields, which problems each binary is offered, job_for_identity); test_describe_jobs_breaks_the_grid_down counts the stride x dilation axis. Checked: a --limit 240 gfx950 bf16 cache run with --run-from-cache --verify on stride-1 shapes (3x3 pad 0/1/2, 5x5 pad 2, grouped) passes every tilde and stride-1 binary; a dilation-2 shape is offered the tilde binaries only.

Brings in fwd group-merge (#12793), the dgrad uniform_k_decode / wgrad prefetch_addr / direct RTC-dgrad and wgrad work (#13320) and the MFMA dgrad preload spill guard (#13495). develop changed the AOT grid, cache identity and direct-conv run side in their old homes; this branch moved them, so the changes are ported into kernels/common/{conv_aot_jobs,conv_kernel_identity, conv_direct_aot_jobs,conv_direct_variants,conv_abi}.py and the new tests import from there. gfx950_dgrad.json: implicit-GEMM dgrad instances now carry uniform_k_decode per DgradConvSpec.default_uniform_k_decode, the identity the AOT generator builds them with. Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>

…m kernargs At stride 1 the tilde decomposition is the identity in both spatial axes (y_tilde = x_tilde = 1 whatever the dilation, pad, filter or grouping), so there is exactly one sub-GEMM and every field of its record is an existing kernarg or a constant: [0, num_pid_m, num_pid_n, dg_M, dg_K, Hi, Wi, pH, pW, Y, X, -dH, -dW, 1, 0, 1, 0, 1, 0, 1, 0, roundup(dg_K, tile_k * split_k)] DgradConvSpec.folds_sub_gemm_record (stride_h == stride_w == 1, read off the spec's problem; for an AOT binary that is its stride capability, which the probe and every served problem share) makes the builder bind those values instead of emitting the binary search over the record buffer and the 21 dependent record loads behind it; the unit strides and zero offsets fold into the address arithmetic, and local_flat drops the block_start subtraction. The magic-division table (#13320) is still read, from its one entry right after the one record: eight loads at constant offsets that depend on nothing. The fold never keys on a sub-GEMM count evaluated on a problem, and no problem extent is read at build time. ABI, identity, label, sweep and dispatch are unchanged; sub_gemm_buf and num_sub_gemms stay in the signature and the buffer stays mandatory (its magic table is read). The fold covers the JIT ungrouped pointwise path too: its record is the same function of the kernargs, so folding it is a straight substitution (the PR's exclusion was about constant trip counts, which runtime extents no longer give). Not a knob: the fold is the stride-1 code path, default on (waiver of the "new behaviour is a default-off knob" rule of KERNEL_AUTHORING.md: there is nothing to switch off -- the folded and the loaded record are equal by construction, test_dgrad_stride1_record_is_a_function_of_the_kernargs). Every stride-1 dgrad kernel changes IR; existing caches hold no stride-1 binaries (the AOT grid builds the tilde kernel only), so no cache entry goes stale. C++: the same in _build_tilde_dgrad (_dgrad_folds_sub_gemm_record), field for field and op for op. Parity (conv_implicit_gemm_dgrad): configs 0, 4 (1x1 s1, ungrouped pointwise) and 21-24 now emit the folded record in both engines; new 40 (s1 d2: folded record with the tilde epilogue), 41 (grouped 5x5 pad 0, K-outer + uniform_k_decode + cshuffle + explicit widths, the AOT shape), 42 (s1 split-K: folded gemm_k_padded), 43 (gfx942 grouped 1x3 pad (0,1)). Golden: changed conv_dgrad/{gfx90a,gfx942,gfx950,gfx1151}/n1h8c16k32r3_s1 and conv_dgrad/{gfx942,gfx950,gfx1151,gfx1250}/n1h8c64k64r3_g4 (stride-1 grouped) at llvm20/22/23; every *_s2 case and everything else unchanged. Tests: - test_conv_abi: the record-equivalence proof (packed record of the one sub-GEMM == the record rebuilt from launch values, over stride-1 problems x dilation 1/2 x pads x filters x groups x tiles x split-K); the IR reads only the eight magic-table entries and runs no search at stride 1 (grouped, dilated) while stride 2 still searches and loads; a stride-1 shape-invariance test that varies C/K, Y/X and pad. - test_conv_dgrad_correctness: TestConvDgradStride1Specializations, re-done for AOT: stride-1 binaries built once against the AOT probe (dilation 1 and 2) and the tilde binary (stride-2 probe), all launched through ConvArgs launch values on adversarial stride-1 shapes (odd H/W, N=1, 5x5 / 1x3 filters, pad 0/1/2 and asymmetric, partial M/N tiles, kpg not a multiple of tile_k, grouped, dilation 2) with dX poisoned: bit-identical to the tilde binary and within tolerance of the torch reference; plus split-K against the reference.

The AOT dgrad grid built the tilde kernel only (stride and dilation at run time), so the stride-1 record fold had no binary to live in. _dgrad_caps now yields a second capability per configuration, stride 1 / dilation 1 (CACHE_DGRAD_STRIDE1): its probe problem is stride 1 / dilation 1, so the binary folds the sub-GEMM record and takes the unit-stride epilogue, and its identity records stride_h/w = dilation_h/w = 1, so identity_supports_problem (the existing exact match on non-zero baked fields) offers it to stride-1, dilation-1 problems only. The tilde binary of the same configuration keeps serving every stride, stride 1 included. count_jobs follows _dgrad_caps; job_for_identity round-trips both. K-outer: the plan's K-outer axis is not added. Since #13320 every dgrad geometry already takes the layout DgradConvSpec.default_lds_k_outer picks (K-outer wherever the arch, dtype, atom, pipeline and a B load wider than one element allow it), so K-outer dgrad binaries are in the grid; a (False, True) axis would only add M-outer twins of them. The K-outer fragment rule the builder enforces is in the validator since the validator-hardening commit. AOT catalog, --compile-all --direction dgrad: one stride-1 identity per tilde configuration, so the dgrad slice of the grid doubles. Tests: test_aot_dgrad_builds_tilde_and_stride1_binaries (probes, identity fields, which problems each binary is offered, job_for_identity); test_describe_jobs_breaks_the_grid_down counts the stride x dilation axis. Checked: a --limit 240 gfx950 bf16 cache run with --run-from-cache --verify on stride-1 shapes (3x3 pad 0/1/2, 5x5 pad 2, grouped) passes every tilde and stride-1 binary; a dilation-2 shape is offered the tilde binaries only.

jakpiase and others added 2 commits October 8, 2026 14:22

jakpiase requested review from a team as code owners October 8, 2026 15:52

github-actions Bot added documentation project: hip-kernel-provider rocke: platform Touches rocKE platform labels Oct 8, 2026

assistant-librarian Bot added the organization: ROCm label Oct 8, 2026

jakpiase requested a balanced review from Copilot October 9, 2026 11:47

Copilot started reviewing on behalf of jakpiase October 9, 2026 11:47

View session

Copilot AI reviewed Oct 9, 2026

View reviewed changes

jakpiase and others added 2 commits October 9, 2026 12:16

jakpiase requested a balanced review from Copilot October 9, 2026 12:18

Copilot started reviewing on behalf of jakpiase October 9, 2026 12:19

View session

Copilot AI reviewed Oct 9, 2026

View reviewed changes

Comment thread ...roviders/hip-kernel-provider/rocke/library/benchmarks/common/benchmark_implicit_gemm_conv.py Resolved

Comment thread ...oviders/hip-kernel-provider/rocke/platform/cpp/instances/common/conv_implicit_gemm_dgrad.cpp Resolved

jakpiase and others added 3 commits October 9, 2026 13:23

style(rocke): clang-format the dgrad kernel-name append macro

45094b7

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

jakpiase requested a balanced review from Copilot October 9, 2026 13:35

Copilot started reviewing on behalf of jakpiase October 9, 2026 13:36

View session

Copilot AI reviewed Oct 9, 2026

View reviewed changes

Comment thread dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/runtime/launcher.py Resolved

bartekxk reviewed Oct 9, 2026

View reviewed changes

Comment thread ...roviders/hip-kernel-provider/rocke/library/benchmarks/common/benchmark_implicit_gemm_conv.py Resolved

jakpiase and others added 2 commits October 12, 2026 09:49

Merge branch 'develop' into users/jakpiase/rocke-aot-wgrad-perf

daecd0c

bartekxk approved these changes Oct 12, 2026

View reviewed changes

jakpiase and others added 5 commits October 12, 2026 17:07

Merge branch 'develop' into users/jakpiase/rocke-aot-wgrad-perf

aa7e991

Merge branch 'develop' into users/jakpiase/rocke-aot-wgrad-perf

66d77d8

Merge branch 'develop' into users/jakpiase/rocke-aot-wgrad-perf

0a97a60

Merge branch 'develop' into users/jakpiase/rocke-aot-wgrad-perf

0e8eecf

bartekxk enabled auto-merge (squash) October 13, 2026 14:55

bartekxk merged commit 9f79212 into develop Oct 13, 2026
97 checks passed

bartekxk deleted the users/jakpiase/rocke-aot-wgrad-perf branch October 13, 2026 15:27

Navigation Menu

Repository navigation

fix(rocke): close conv perf gaps after the switch to AOT - #13320

fix(rocke): close conv perf gaps after the switch to AOT#13320
bartekxk merged 15 commits into
developfrom
users/jakpiase/rocke-aot-wgrad-perf

jakpiase commented Oct 8, 2026

Uh oh!

therock-pr-bot Bot commented Oct 8, 2026 •

edited

Loading

Uh oh!

therock-pr-bot Bot commented Oct 8, 2026 •

edited

Loading

Uh oh!

Copilot AI left a comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

therock-pr-bot Bot commented Oct 9, 2026

Uh oh!

Copilot AI left a comment

Uh oh!

Uh oh!

Uh oh!

Copilot AI left a comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

assistant-librarian Bot commented Oct 14, 2026

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

Conversation

jakpiase commented Oct 8, 2026

Motivation

Technical Details

Test Plan

Test Result

Uh oh!

therock-pr-bot Bot commented Oct 8, 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 Oct 8, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

🟡 Changes recommended

Uh oh!

Uh oh!

Uh oh!

Uh oh!

therock-pr-bot Bot commented Oct 9, 2026

Pre-commit check failed

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

🟡 Changes recommended

Uh oh!

Uh oh!

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

🟡 Changes recommended

Uh oh!

Uh oh!

Uh oh!

Uh oh!

assistant-librarian Bot commented Oct 14, 2026

TheRock Submodule Bump Activity

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

therock-pr-bot Bot commented Oct 8, 2026 •

edited

Loading

therock-pr-bot Bot commented Oct 8, 2026 •

edited

Loading