Repository navigation
fix(rocke): fwd perf regressions after switch to aot - #13109
Conversation
✅ 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. |
There was a problem hiding this comment.
Copilot review overview
🟡 Changes recommended
Unresolved correctness, validation-consistency, and regression-testing issues block approval.
Review effort: Balanced
Findings: 2
Open (7)
Exclude failed kernels from timing and pruning · New Regenerate stale representative IR golden hashes · New Add small-group spatial plan selection tests · New Add numerical tests for streaming depthwise dgrad · New Add numerical tests for 7x7 stride-2 depthwise tails · New Add large-filter spatial correctness coverage · New Align depthwise filter limits across Python and C++ · New
What changed in this PR
Addresses convolution performance regressions after the switch to ahead-of-time compilation, with depthwise kernel optimizations and lower benchmark launch overhead.
Changes:
- Adds prefetched depthwise pipelines and runtime spatial tiling.
- Introduces reusable launches, cross-run early stopping, and filtered cache inspection.
- Extends spatial cache registration and launcher, ABI, and parity tests.
Paths below are relative to dnn-providers/hip-kernel-provider/rocke/.
| File | Description |
|---|---|
platform/tests/test_rocke.py |
Tests reusable launch binding. |
platform/tests/instances/differential/golden/llvm_gfx_all.json |
Updates depthwise IR baselines. |
platform/python/rocke/runtime/launcher.py |
Adds prepacked launch binding. |
platform/python/rocke/runtime/hip_module.py |
Prepares reusable HIP launch arguments. |
platform/cpp/instances/common/conv_direct_grouped_build_depthwise.cpp |
Adds prefetching and a separate tail loop. |
platform/cpp/instances/common/conv_direct_grouped_build_depthwise_spatial.cpp |
Pipelines loads and computes runtime spatial tiling. |
platform/cpp/include/rocke/instance_conv_direct_grouped_internal.h |
Expands depthwise filter storage. |
library/tests/test_early_stop.py |
Tests timing record and seed handoff. |
library/tests/test_conv_abi.py |
Checks spatial group-count invariance. |
library/tests/parity/conv_direct_grouped_emit.py |
Adds a larger-filter Python parity case. |
library/tests/parity/conv_direct_grouped_emit.c |
Adds the matching C parity case. |
library/kernels/common/conv_direct_grouped.py |
Optimizes depthwise forward and gradient pipelines. |
library/benchmarks/common/kernel_sweep.py |
Filters cache descriptions by direction. |
library/benchmarks/common/early_stop.py |
Records and seeds per-case timings. |
library/benchmarks/common/direct_kernel_sweep.py |
Registers spatial depthwise cache builds. |
library/benchmarks/common/benchmark_implicit_gemm_conv.py |
Uses bound launches and seeded timing. |
library/benchmarks/common/benchmark_direct_conv.py |
Integrates bound launches and timing records. |
library/benchmarks/common/benchmark_conv_compare.py |
Streams output and shares direct timing seeds. |
💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
aledudek
left a comment
There was a problem hiding this comment.
LGTM - added a few minor comments
|
Cancelling the remaining gfx950 jobs on this run so the TheRock submodule bump PRs can be prioritized on the gfx950 queue. This change does not target gfx950, and gfx950 was the only work still outstanding. Jobs that already finished are unchanged. Re-run failed jobs if you need gfx950 results before merge. |
…arkocot/fwd-aot-regression-fixes
…arkocot/fwd-aot-regression-fixes
The differential golden had gone stale: 627 of its 1240 LLVM IR entries and 537 of its 1206 canonical IR entries no longer matched the tree, because nothing ran the check. Recorded afresh in the new format, with the Python reference emitters, from a tree whose two-engine gate is green for all 82 families: - ll at llvm20, llvm22 and llvm23: 82 families, 1518 configs per flavor; 1615 at llvm23, where gfx1250_lowering has configs that need LLVM 23 and are refused below it. - ir_canonical: 81 families, 1598 configs (one family'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 magic division sum emitted without signed-overflow poison (#13802): 215 configs in each of the four sections, 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 per section in attention_strided_kv, gfx942_attention_tiled_3d and gfx950_attention_tiled_3d; - wave_all on wave32 (#13896): 18 gfx1151_wmma_fmha_fwd configs per LLVM flavor, canonical IR unchanged; - new configs: the barrier_fence family (#13575, 10) and one conv_implicit_gemm_dgrad config (#13782). 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>
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>
…13974) ## Motivation The differential golden (`tests/instances/differential/golden/llvm_gfx_all.json`) is the only tripwire for an output change in about 55 parity families that the representative golden does not cover (norms, reduce, transpose, fmha, wmma GEMMs, flatmm, mx_gemm, ...). The byte-identity gate cannot see a change made the same way in both engines; this file can. But nothing ran its check, and on develop 627 of its 1240 LLVM IR entries and 537 of its 1206 canonical IR entries no longer matched. The checker itself made it easy to go stale: - Entries carried no LLVM flavor, so blesses from hosts of different flavors mixed and no flavor matched the file. - A re-bless limited with `--only` replaced the whole section; one did, and the pasted-back copy left two repeated JSON keys whose older copies silently won. - A check limited with `--only` reported every other family missing. - The check needed the C++ build although it compares Python reference hashes only. ISSUE ID : AICK-2654 ## Technical Details Three commits. 1. **The checker** (`fix`) - New `golden_store.py` owns the file (schema 2): LLVM IR hashes in one section per flavor (`llvm20`, `llvm22`, `llvm23`), canonical IR in one section (produced before lowering; identical across flavors, checked). Loading refuses repeated keys and the old format. A filtered record replaces only the families it ran; a filtered check compares only those. - `run_diff.py` pins the flavor for every golden run (`--flavor`, else `ROCKE_LLVM_FLAVOR`, else the host's) and exports it to both emitters. - `--python-only` records or checks from the Python reference emitters alone: no C++ engine, any flavor on any host, a list of flavors in one run. It produces the same hashes as the two-engine run (all 82 families, checked). - The canonical hash has one definition (`python_batch_worker.canonical_digest`), shared by the batch worker, which now computes it in parallel, and `run_diff.py`. - `run_diff.py` finds the rocke and library packages in an installed test tree too. 2. **The re-bless** (`test`, data only): recorded afresh at llvm20, llvm22, llvm23 and canonical IR from a tree whose two-engine gate is green for all 82 families (re-recorded after rebasing onto develop with #13802, whose magic-division change moved 215 configs in 19 families in every section). Every mismatch that could be traced is an intended code generation change that did not re-bless this file (#8844, #9583, #9133, #10112, #12710, #9220, #13109); 441 entries were already stale at the start of the history available and cannot be attributed; no unexplained change was found after it. The commit message has the breakdown. 3. **CI** (`test`): two quick-tier CTest entries run the Python-only check from the installed test tree, `rocke_differential_golden_ll` (all three flavors) and `rocke_differential_golden_canonical`. The install ships only what that check needs from `tests/instances/differential`; the rest still needs the C++ build and stays out. The golden README and the troubleshooting entry describe the check and re-bless commands. No emitter or IR changes. ## Test Plan - `test_golden_store.py`: repeated keys refused by name; the old format refused with the way out; flavor sections separate; a full record replaces, a filtered record merges; a filtered check ignores other families; new configs alone pass. - The existing differential tests (gate verdict, scaled MFMA batch) still pass. - Python-only vs two-engine record at llvm20: identical hashes for all 82 families; canonical IR identical at llvm20 and llvm23. - The check passes on the re-blessed tree and fails (1058 changed configs) when the lowerer's output is altered. - Installed tree (provider layout, staged install): both CTest entries pass, and fail when the installed lowerer's output is altered. - Byte identity at llvm20, llvm22, llvm23; ctest; full platform suite. ## Test Result - All of the above pass. - Installed CTest time: about 15 s (LLVM IR, three flavors) and 13 s (canonical IR). - Byte identity green at llvm20, llvm22 and llvm23; ctest passes; full platform suite passes with no failures; pre-commit passes. ## Submission Checklist - [ ] Look over the contributing guidelines at https://github.com/ROCm/TheRock/blob/main/GOVERNANCE.md#pull-requests. 🤖 Generated with [Claude Code](https://claude.com/claude-code) --------- Co-authored-by: Claude <noreply@anthropic.com>


Motivation
Fix perf regression after switch to AOT.
Technical Details
Improvements for direct and implicit conv
Test Plan
test_conv_*_correctness
Test Result
pass
Submission Checklist
JIRA ID: AICK-2509