Skip to content

fix(rocke): fwd perf regressions after switch to aot - #13109

Merged
bartekxk merged 9 commits into
developfrom
users/barkocot/fwd-aot-regression-fixes
Oct 7, 2026
Merged

bartekxk merged 9 commits into
developfrom
users/barkocot/fwd-aot-regression-fixes

Conversation

@bartekxk

@bartekxk bartekxk commented Oct 5, 2026 •

Copy link
Copy Markdown
Contributor

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

@bartekxk bartekxk self-assigned this Oct 5, 2026
@bartekxk
bartekxk requested review from a team as code owners October 5, 2026 19:37
@therock-pr-bot

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

Copy link
Copy Markdown

✅ 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

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

Copy link
Copy Markdown

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

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Copilot review overview

🟡 Changes recommended

Unresolved correctness, validation-consistency, and regression-testing issues block approval.

Review effort: Balanced
Findings: 2 High severity · 5 Medium severity

Open (7)
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.

bartekxk and others added 2 commits October 6, 2026 06:01
Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>

@aledudek aledudek left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

LGTM - added a few minor comments

@JeniferC99

Copy link
Copy Markdown
Contributor

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.

@bartekxk
bartekxk enabled auto-merge (squash) October 6, 2026 23:10
@bartekxk bartekxk mentioned this pull request Oct 7, 2026
1 task done
@bartekxk
bartekxk merged commit 3ecaa65 into develop Oct 7, 2026
93 checks passed
@bartekxk
bartekxk deleted the users/barkocot/fwd-aot-regression-fixes branch October 7, 2026 12:08
cgmillette added a commit that referenced this pull request Oct 20, 2026
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>
cgmillette added a commit that referenced this pull request Oct 26, 2026
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>
cgmillette added a commit that referenced this pull request Oct 26, 2026
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>
cgmillette added a commit that referenced this pull request Oct 27, 2026
…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>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

4 participants