[EXTERNAL] Check the direct-to-LDS swizzle against the chunk a warp writes - #556
bogdan-petkovic wants to merge 2 commits into
Conversation
There was a problem hiding this comment.
Copilot review overview
🟡 Changes recommended
Both lowering guards validate logical shapes instead of the allocation shapes used for memdesc subviews.
Review effort: Balanced
Findings: 1
What changed in this PR
Updates direct-to-LDS swizzle validation to use each warp’s written chunk and adds lowering safeguards and regression coverage.
Changes:
- Adds shared chunk-based swizzle validation.
- Falls back to stream copies for pinned layouts.
- Adds and updates clamp and lowering tests.
| File | Description |
|---|---|
triton-patches/triton-patch-content.txt |
Documents the downstream patch. |
triton-patches/patch-direct-to-lds-swizzle-chunk.patch |
Records the vendored Triton changes. |
mlir/test/Dialect/Rock/direct-to-lds-swizzle-clamp-pinned.mlir |
Tests pinned-layout fallback. |
mlir/test/Dialect/Rock/direct-to-lds-swizzle-clamp-load-vec.mlir |
Updates expected clamp behavior. |
mlir/test/Dialect/Rock/direct-to-lds-swizzle-chunk-guard.mlir |
Tests unsafe-swizzle diagnostics. |
external/triton/third_party/amd/lib/TritonAMDGPUTransforms/LowerLoops.cpp |
Applies chunk-based clamping and fallback. |
external/triton/third_party/amd/lib/TritonAMDGPUToLLVM/Utility.h |
Declares the validation helper. |
external/triton/third_party/amd/lib/TritonAMDGPUToLLVM/Utility.cpp |
Implements chunk validation. |
external/triton/third_party/amd/lib/TritonAMDGPUToLLVM/LoadStoreOpToLLVM.cpp |
Guards both direct-to-LDS lowerings. |
external/triton/test/TritonGPU/amd/amd-pipeline-shared-layout-async-copy-gfx9.mlir |
Updates the vendored clamp expectation. |
💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
![rocmlir-pr-reviewer[bot]](https://avatars.githubusercontent.com/u/7440191?s=60&v=4){kind=small}
There was a problem hiding this comment.
Verdict: COMMENT · Findings: 3 (0 Critical, 1 Major, 2 Minor)
Scope
Downstream-only change to the vendored Triton tree: the GFX9 direct-to-LDS swizzle legality check moves from walking the lanes of the pre-CoalesceAsyncCopy load layout to a layout-only test (isSwizzleInsideDirectToLdsChunk) against the aligned chunk a warp writes. swizzlesInsideWarp (clamp side, LowerLoops.cpp) and both direct-to-LDS lowerings (LoadStoreOpToLLVM.cpp) now share that predicate; the lowerings hard-error instead of silently shuffling pointers from outside the warp. createStreamOps falls back to a stream copy when the clamp would change an encoding pinned by a memdesc user. Three rocMLIR lit tests (two new, one updated), one vendored test expectation, plus the patch record and index entry.
Findings
external/triton/third_party/amd/lib/TritonAMDGPUToLLVM/Utility.cpp:902— the XOR/low-bits correctness argument requireschunkSizeto be a power of two; add an assert (Major).external/triton/third_party/amd/lib/TritonAMDGPUToLLVM/LoadStoreOpToLLVM.cpp:809— the guard and its diagnostic string are duplicated verbatim in both conversion patterns (Minor).external/triton/third_party/amd/lib/TritonAMDGPUTransforms/LowerLoops.cpp:666— thepinnedpredicate matches any memdesc-producing user regardless of its encoding (Minor).
Notes
Process checks all pass: the vendored-tree commit 704a0dda carries the [EXTERNAL] prefix; triton-patches/patch-direct-to-lds-swizzle-chunk.patch matches that commit's diff hunk-for-hunk (paths rebased to the subtree root); the triton-patch-content.txt entry is placed after patch-direct-to-lds-swizzle-check-vec.patch and documents the stacking, and the reverse-apply order the bump guide uses (sort -r) happens to unwind chunk before check-vec, which is the required order. Both new .mlir files carry the first-party AMD/Apache-with-LLVM-exception header; the negative test uses -split-input-file with --verify-diagnostics.
No rocMLIR back-port note is needed — every touched path is in the rocmlirTriton-only list (external/triton/, triton-patches/) or is lit coverage for Triton-only passes under mlir/test/Dialect/Rock/.
Design note (not a finding): the new lowering guard converts what was previously a silent wrong-result miscompile into a compile-time error. The clamp side computes the chunk from the pre-coalescing vec, so a case where CoalesceAsyncCopy narrows further than min(loadVec, getNumConsecutiveInOut()) models would surface as that hard error rather than bad numerics. The PR reports the error never fires across 16748 compiled kernels, which is good evidence, but it is worth keeping in mind as the failure mode if a future Triton bump changes coalescing.
CI status
No failed or cancelled checks. Jenkins, Build and Test, MIGraphX, Code coverage and ml-ci-internal were still pending at review time, so the functional signal for this change is not yet in.
f53ac2e to
e135c91
Compare
There was a problem hiding this comment.
Verdict: APPROVE -- submitted as COMMENT (automated reviews are advisory) · New findings: 0 (0 Critical, 0 Major, 0 Minor)
Scope
Re-review of the direct-to-LDS swizzle chunk check. The PR replaces the lane-walk in swizzlesInsideWarp with a new shared-layout-only predicate isSwizzleInsideDirectToLdsChunk (external/triton/third_party/amd/lib/TritonAMDGPUToLLVM/Utility.cpp:900), wires it into both direct-to-LDS lowerings as a hard error, and stream-copies loads whose shared layout is pinned by a memdesc user. Two vendored-tree commits are correctly split: the external/triton/ edits land under [EXTERNAL] Check the direct-to-LDS swizzle against the chunk a warp writes, and the patch record plus lit tests land in a separate first-party commit.
Findings
No blocking issues found. Both actionable findings from the previous round are addressed in the current head:
- The power-of-two precondition on
chunkSizeis now asserted atUtility.cpp:903with a descriptive message, and the requirement is documented in theUtility.h:126doc comment. - The duplicated nine-line guard is hoisted into
DirectToLdsLoadConversionBase::verifySwizzleInsideDirectToLdsChunk(LoadStoreOpToLLVM.cpp:297) and called from bothBufferLoadToLocalOpConversionandAsyncCopyGlobalToLocalOpConversion, so the diagnostic string exists once.
The third thread (the pinned predicate in LowerLoops.cpp:664) is closed on the author's explanation that getSharedEncIfAllUsersAreDotEnc only yields an encoding when every memdesc user agrees on perPhase/maxPhase/order, so a full-encoding comparison would wrongly unpin the vec-only-differs case.
Notes
- Spot-checked that
chunkSizecannot be zero at either lowering call site:canLoadDirectToLDSsetsvectorSize = contigand rejects widths the target cannot write, sovecis a non-zero power of two by the time the guard runs. InswizzlesInsideWarpthevec == 0early return covers the clamp path. triton-patches/patch-direct-to-lds-swizzle-chunk.patchmatches the[EXTERNAL]commit diff hunk-for-hunk across all five vendored files, and thetriton-patch-content.txtentry documents the stacking onpatch-direct-to-lds-swizzle-check-vec.patch. The bump guide's reverse-apply order (ls | sort -r) puts...-chunk.patchahead of...-check-vec.patch, which is the required newest-first order for this stack.- New lit files under
mlir/test/Dialect/Rock/carry the first-partyCopyright Advanced Micro Devices, Inc./Apache-2.0 WITH LLVM-exceptionheader, and coverage includes both negative cases (direct-to-lds-swizzle-chunk-guard.mlir) and the pinned-layout fallback. - rocMLIR back-port check does not apply: every touched path is in the rocmlirTriton-only set (
external/triton/,triton-patches/) or is a lit test exercising Triton-only passes that do not exist upstream.
CI status
No failing or cancelled checks; py-checks and detect pass, the reviewer checks are still in progress.
| Stacks on patch-direct-to-lds-swizzle-check-vec.patch and must be applied | ||
| after it. | ||
|
|
||
| Not filed upstream yet; upstream PR 11295 carries the same check. |
…rites swizzlesInsideWarp walked the lanes of warp 0 in the layout the load has in LowerLoops. CoalesceAsyncCopy rewrites that layout whenever it narrows the load, so the walk measured a layout that is never lowered and overstated how far the lane shuffle moves. It clamped swizzles that stay inside the warp, down to maxPhase 1 for the load in loop-pipeline-hip.mlir whose local_alloc keeps maxPhase 8, and the pipeliner then could not predicate that local_alloc. Without scatter support each direct-to-LDS load of a warp writes one aligned chunk of lanes * vec consecutive elements, whatever layout CoalesceAsyncCopy picks, and the lane shuffle stays inside the warp exactly when the swizzle keeps every element inside its chunk. isSwizzleInsideDirectToLdsChunk checks that on the shared layout alone, through the basis offsets of the map from unswizzled to swizzled offsets. The clamp uses it with the load vector it already computes, and both direct-to-LDS lowerings use it with their final vector and emit an error instead of shuffling pointers from outside the warp. When the clamp would change a layout that a memdesc user pins, the pipeliner cannot stage a conversion between the two, so createStreamOps stream copies such a load instead. async_copy_swizzle_clamped_to_warp in amd-pipeline-shared-layout-async-copy-gfx9.mlir now expects maxPhase 8. Its coalesced copy writes two whole rows per warp, the same lowered copy as async_copy_swizzle_stays_in_warp.
Records the change as triton-patches/patch-direct-to-lds-swizzle-chunk.patch, stacked on patch-direct-to-lds-swizzle-check-vec.patch, with its entry in triton-patch-content.txt. direct-to-lds-swizzle-clamp-load-vec.mlir now expects maxPhase 8 in all four cases. Each tile loads one f32 per lane, so a warp writes half of a 128-wide row, and in three of them CoalesceAsyncCopy narrows the layout, which the old check read as a longer shuffle. direct-to-lds-swizzle-clamp-pinned.mlir covers the stream copy fallback for a pinned layout, and direct-to-lds-swizzle-chunk-guard.mlir covers the error from both lowerings.
e135c91 to
afd7c7d
Compare
Motivation
#548 made the swizzle clamp from triton#11295 measure the direct to LDS lane shuffle with the load's vector, but the check still walks the lanes of the layout the load has in LowerLoops. CoalesceAsyncCopy rewrites that layout whenever it narrows the load, so the check looks at a layout that is never lowered. It reports shuffles that leave the warp when they do not, and clamps more than it has to.
In
TritonGPU/loop-pipeline-hip.mlir, from the vendored suite that #494 turns on, this clamps a gfx942 load from maxPhase 8 to 1 while itslocal_allockeeps 8, and the pipeliner stops with "Fatal pipeliner error". The test passed before #548.Technical Details
Without scatter support each direct to LDS load of a warp writes one aligned chunk of lanes times vec consecutive elements, whatever layout CoalesceAsyncCopy picks, because
canLoadDirectToLDSonly accepts layouts that write that way. The swizzle is applied by exchanging source pointers between the lanes, so it works exactly when the swizzle keeps every element inside its chunk. The newisSwizzleInsideDirectToLdsChunkchecks that on the shared layout alone. It builds the map from unswizzled to swizzled offsets and, since that map is linear over XOR, checks that no basis offset moves above the low log2(chunk) bits.swizzlesInsideWarpkeeps the load vector from #548 and calls the new function instead of walking lanes. The halving loop inclampSwizzleForDirectToLdsis unchanged.Both direct to LDS lowerings,
BufferLoadToLocalOpConversionandAsyncCopyGlobalToLocalOpConversion, now call the same function with their final vector and emit an error instead of shuffling pointers from outside the warp. Until now the lowering only assumed this, and a miss showed up as wrong results at runtime. The error is explicit becauseamdg.buffer_load_to_localis not marked illegal in our pinned conversion, so a bare failure would only surface later in LLVM translation.A load whose shared layout is pinned by a memdesc user, a
local_allocfor example, cannot be clamped, because the pipeliner cannot stage a conversion to the pinned layout. When the clamp would change such a layout,createStreamOpsnow stream copies the load in the pinned layout.async_copy_swizzle_clamped_to_warpinamd-pipeline-shared-layout-async-copy-gfx9.mlirnow expects maxPhase 8. CoalesceAsyncCopy narrows its copy to 32-bit loads, so each warp writes two whole rows, and it lowers to the same copy asasync_copy_swizzle_stays_in_warp, which already expects 8. The case keeps its name to stay close to upstream.The change is recorded as
triton-patches/patch-direct-to-lds-swizzle-chunk.patch, stacked onpatch-direct-to-lds-swizzle-check-vec.patch, with its entry intriton-patch-content.txt. It is not filed upstream yet, triton#11295 carries the same check.Test Plan
Lit, locally, the full vendored Triton suite and
check-rocmlir, plus the new and changed tests run against develop's code to make sure they fail there.An instrumented build computed, inside the lowering, the lanes every copy actually reads and the largest maxPhase that stays inside the warp on the layout the lowering gets. I compiled 16748 rocMLIR gemms with develop and with this change, on gfx950 (f32, f16, bf16, fp8 and i8, nine values of k, all four transposes, the heuristic config, the quick tuning space and the full one for f32 k 769 and 770) and on gfx942 with async copy forced (f32, f16 and bf16).
On gfx950 I ran every changed kernel with validation, compared the compiled output of every unchanged kernel byte for byte, and benchmarked the changed kernels with
rocmlir-tuning-driver.Test Result
The vendored suite fails only the four tests #494 fixes, and
loop-pipeline-hip.mlirpasses.check-rocmlirfails only theperf-scriptstests, which need pandas on my machine. The three new or changed rocMLIR tests and the changed vendored test fail with develop's code and pass with this change.9892 of the 16748 kernels have swizzled direct to LDS copies. With this change no copy leaves its warp, and every clamp lands on the largest maxPhase that stays inside. Develop also keeps every copy inside, but clamps 2029 kernels further than needed. Only those 2029 change, always to a larger maxPhase and with the same load width. The other 12449 that compile give byte identical output, and the compile failures are the same 2892 (LDS limit) in both. The error and the stream copy fallback never trigger on this corpus.
All 861 changed gfx950 kernels that compile pass validation. Performance moves both ways. The six changed kernels on the default heuristic config come out 1.1% faster on geometric mean, from 6% faster to 4.6% slower. Taking the best config of the quick space for 12 f32 problems, the best config is the same with and without the change, the geometric mean is 0.7% slower, and the non-transposed k 770 and 774 problems are 4 to 5% slower. Turning the swizzle off altogether when the full one does not fit was 16% slower on the default kernels, so keeping the largest maxPhase that stays inside the warp is the better rule.
Submission Checklist