[AIROCMLIR-1292] [EXTERNAL] Do not stagger CDNA4 padded layouts over rows that do not exist

[bogdan-petkovic]

Motivation

The weekly tuning stage aborts the compiler on gfx950. Three consecutive runs of PR #361 died the same way, always on gemm f16 m=384 k=3072 n=768 with the tuner emitted config mPerBlock=1, nPerBlock=16, kPerBlock=2048, about a minute into the gemm stage. Because a failing chip aborts the rest of the matrix, this single config was enough to stop the whole weekly and hide every other chip's results.

Technical Details

composePaddedLayoutForAsyncCopyCDNA4 in third_party/amd/lib/TritonAMDGPUTransforms/Utility.cpp computes warpSize / contigLanes with integer division in two places. Once a row of the tile spans more than a warp of vectors, that is contigLanes > warpSize, the quotient truncates to zero and two things break. requiredRows collapses to zero, so xWayConflicts always reports one and the "too few rows to stagger" path never triggers. And the useBestWrap guard degenerates to nonContigDim >= 0, which is always true, so wrap is forced to 16 even for a tile with a single row. Bases are then added for rows 1, 2, 4 and 8 that the tile does not have, the resulting linear layout broadcasts in the offset dimension, and PaddedSharedEncodingAttr::verify rejects it. In an assertions build that aborts the compiler rather than falling back. The failing call chain is PipelinePass to expandLoops to lowerLoops to lowerLoop to getSharedEncIfAllUsersAreDotEnc to composePaddedLayout to composePaddedLayoutForAsyncCopyCDNA4.

The fix is the one suggested in the upstream review of triton-lang/triton#11890: contigLanes is clamped to [1, warpSize]. A warp then counts as covering one row when a row is wider than a warp, so requiredRows and the useBestWrap guard keep their original form and work again, and a row narrower than a vector can no longer make contigLanes zero. The [EXTERNAL] commit is the upstream change plus its lit test, recorded as triton-patches/patch11890.patch; the downstream copy of the test only adds the MIT header this repo requires.

Test Plan

• A sweep of 70 gemm configs on the weekly's failing problem for gfx950, f16 and f32 with mPerBlock 1 to 64 and kPerBlock 128 to 2048, compiled through rocmlir-driver -c with and without the change, comparing the generated code by hashing the lowered module, which embeds the HSACO.
• The full tier1 gemm quick tuning space for gfx950, compiled and compared the same way.
• Triton's AMD lit tests with and without the change.
• The new lit test amd-pipeline-padded-layout-wide-row-gfx950.mlir and the new E2E test mlir/test/fusion/pr-e2e/rock-gemm-cdna4-padded-layout-wide-row.mlir.

Test Result

• In the 70-config sweep, compiler aborts go from 8 to 0, and the 55 configs that compiled before produce bit identical code.
• On today's develop the tier1 gemm quick space no longer contains the failing configs; all 1578 of its gfx950 pairs compile with and without the change and produce bit identical code.
• Triton's AMD lit tests give the same results with and without the change.
• The lit test aborts triton-opt without the change and passes with it. The E2E test runs the weekly's config plus a 2-row and a 4-row tile. All three abort the compiler without the change, and with it the test passes on the gfx950 CI node.

Submission Checklist

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

[pabloantoniom]

This makes sense to me, it looks like a genuine bug in Triton upstream...Could you create a PR in https://github.com/triton-lang/triton ? They should have more context and they might be able to give more valuable feedback

[bogdan-petkovic]

This makes sense to me, it looks like a genuine bug in Triton upstream...Could you create a PR in https://github.com/triton-lang/triton ? They should have more context and they might be able to give more valuable feedback

Sure, I've opened the upstream PR here triton-lang/triton#11890

[erizheng-amd]

I can confirm this fixed the recent gemm tuning failures observed on gfx950

[Copilot]

Copilot review overview

🟡 Changes recommended

Add wide-row regression coverage and correct the patch index placement.

Get a fresh assessment by requesting another Copilot review.

Review effort: Balanced
Findings: 1 · 1

Open (2)

• Add regression test for contigLanes exceeding warpSize · New
• Move gfx1310 patch entry into the filename index · New

What changed in this PR

Fixes CDNA4 padded-layout generation for rows wider than a warp, preventing gfx950 compiler aborts.

Changes:

• Replaces truncating divisions with scaled comparisons.
• Prevents staggering across nonexistent rows.
• Records the downstream Triton patch.

File	Review
triton-patches/​triton-patch-content.txt	Documents the patch; index entry is misplaced.
triton-patches/​patch11890.patch	Records the vendored Triton fix.
external/​triton/​third_party/​amd/​lib/​TritonAMDGPUTransforms/​Utility.cpp	Corrects layout guards; the triggering wide-row case lacks regression coverage.

---

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

[rocmlir-pr-reviewer]

Verdict: COMMENT  ·  Findings: 3 (0 Critical, 2 Major, 1 Minor)

---

Scope

Three files: a downstream fix to composePaddedLayoutForAsyncCopyCDNA4 in the vendored Triton tree (external/triton/third_party/amd/lib/TritonAMDGPUTransforms/Utility.cpp), the matching record triton-patches/patch11890.patch, and an entry in triton-patches/triton-patch-content.txt.

The code change itself looks correct. Both guards previously computed warpSize / contigLanes with integer division, which truncates to zero once a tile row spans more than a warp of vectors (contigLanes > warpSize); rescaling by contigLanes on the other side of the comparison is the right exact-arithmetic rewrite and is a no-op whenever contigLanes divides warpSize (it always does here, since contigDim and vecSize are powers of two). The extra nonContigDim >= bestWrap term is genuinely needed: the rescaled guard alone still admits nonContigDim == 4 at contigLanes == 256, which would stagger over 16 rows that do not exist.

Findings

• triton-patches/triton-patch-content.txt:36 — the new filename is inserted between the patch-gfx1310-isa-family.patch heading and its description body, duplicating the full entry added at the bottom of the file.
• external/triton/.../Utility.cpp:281 — no lit coverage added for the fixed path.
• external/triton/.../Utility.cpp:258 — the rewrite carries over a latent division by zero when contigLanes == 0.

Notes

• Commit hygiene is correct: the vendored-tree commit carries the [EXTERNAL] prefix, the record commit is separate, and patch11890.patch matches the [EXTERNAL] commit diff hunk-for-hunk.
• No rocMLIR back-port note needed — both touched trees (external/triton/, triton-patches/) are on the rocmlirTriton-only path list.
• Follow-up (out of scope, and acknowledged in the patch record's "3 that remain"): wrap can still exceed nonContigDim after this change. With warpSize=64, contigLanes=256, wrap=8, nonContigDim=2, the rescaled conflict guard passes (512 >= 512), useBestWrap is correctly declined, and the "Add rows [0, wrap]" loop at Utility.cpp:303 still emits bases for rows 2 and 4. A wrap = std::min(wrap, nonContigDim) clamp (or a bail-out) before the bases loop would close that arm too.
• The PR description's Test Result section links to an internal Jenkins host. Per the checklist's Critical bullet on internal-only hyperlinks — which applies even though this repo is private today — please replace it with a plain-text job/run reference.

CI status

Nothing red. Jenkins, Build and Test, MIGraphX and Code coverage are still pending at the time of review; re-check before merge.

[dhernandez0]

can you add a lit test for this?

[dhernandez0]

also an e2e test (rocmlirTriton)

[bogdan-petkovic]

Added both, the lit test from the upstream PR, and mlir/test/fusion/pr-e2e/rock-gemm-cdna4-padded-layout-wide-row.mlir, which runs the weekly config plus 2-row and 4-row tiles end to end on gfx950

[dhernandez0]

looks good to me, it's equivalent for non-crash cases of the previous version. But I'm not familiar with this code.

[dhernandez0]

nit: I wonder if adding a comment about the non-scaled comparison would make the code easier to read.

[bogdan-petkovic]

After the upstream review the fix changed to a clamp of contigLanes to [1, warpSize], so both guards keep their original form, and the clamp has a short comment. It's still equivalent for everything that compiled before, 55 of 55 bit identical

[dhernandez0]

nit: nonContigDim >= std::max(bestWrap, warpSize * bestWrap / contigLanes),

[bogdan-petkovic]

That comparison is gone with the clamp, the guard is back to its original form

[bogdan-petkovic]

Updated to follow the upstream review on triton-lang/triton#11890 the fix is now a single clamp of contigLanes to [1, warpSize] instead of the scaled comparisons. In the 70-combination gfx950 sweep it takes aborts from 8 to 0 (the scaling left 3), and everything that compiled before produces bit identical code. Also added the upstream lit test and a gfx950 E2E test, and fixed the index entry

[Copilot]

Copilot review overview

🟡 Changes recommended

The Technical Details materially describe a different implementation and test state than the current diff.

Review effort: Balanced
Findings: 2

Open (2)

• Update PR documentation to match clamp-based implementation · New
• Move gfx1310 patch entry into the filename index

Resolved since last review (1)

• Add regression test for contigLanes exceeding warpSize

[rocmlir-pr-reviewer]

Verdict: COMMENT  ·  New findings: 1 (0 Critical, 1 Major, 0 Minor)

---

Scope

Fixes a gfx950 compiler abort in composePaddedLayoutForAsyncCopyCDNA4. When a tile row spans more vectors than a warp has lanes, contigDim / vecSize exceeded warpSize (and could truncate to 0), so requiredRows and the useBestWrap guard both degenerated and the pass emitted a broadcasting layout that PaddedSharedEncodingAttr::verify rejects. The head clamps contigLanes to [1, warpSize] at Utility.cpp:243 and leaves the two downstream guards untouched. Six files: the Triton fix plus its upstream lit test (one [EXTERNAL] commit), and the patch record, index entry, gfx950 E2E test, and its lit.local.cfg gate (one downstream commit).

Findings

One documentation-only Major: the PR description describes a different fix than the one implemented (Utility.cpp:243). No code change requested.

Re-review status

All three findings from the previous review are fixed:

• patch11890.patch is now at the end of the ordered apply list in triton-patches/triton-patch-content.txt, no longer nested inside another patch's section.
• Lit coverage was added (amd-pipeline-padded-layout-wide-row-gfx950.mlir) and is carried in the patch record, so a future subtree bump can't silently drop it.
• The contigLanes == 0 division hazard is closed by the clamp's lower bound.

Notes

I hand-verified the decline path for the three configs in the E2E test. For mPerBlock=1, kPerBlock=2048 (f16, kWidth=8): contigLanes clamps 256 to 64, wrap = 8, requiredRows = 8, and with nonContigDim = 1 the conflict estimate is 4, so the function returns {} instead of building an invalid layout. mPerBlock=2 and mPerBlock=4, kPerBlock=1024 likewise decline. Inputs already in [1, warpSize] are bit-identical, consistent with the "55 of 55 identical" sweep in the description.

std::clamp needs <algorithm>, which this file does not include directly; it already relies on std::min at line 39 through the same transitive path, so this adds no new build risk — worth an explicit include if you touch the file again.

Patch record matches the [EXTERNAL] commit diff exactly modulo the subtree prefix; commit prefixes and license headers (MIT under external/triton/, Apache-2.0 WITH LLVM-exception for the new first-party .mlir) are correct. No rocMLIR back-port note needed — the only non-Triton paths are under mlir/test/.

Minor confidentiality note: the Test Result section links a Jenkins run on an internal CI host. Consider quoting the result rather than linking, since external readers cannot reach it.

CI status

All non-self checks pass. Code coverage is pending and the auto-review check is in progress; neither is a failure.

[bogdan-petkovic]

@dhernandez0 @pabloantoniom The upstream PR was approved and merged by Alex triton-lang/triton@1725224a65ebc2dc7ca2ee1d84458344d93ee797. It is identical to the [EXTERNAL] commit here apart from the MIT header on the test, and the patch record now points at that merge commit

[pabloantoniom]

Thanks for the patience in the review and for submitting the fix upstream!