[AIROCMLIR-1090] Clamp masked-out lane offsets in transforms-to-ptr

[bogdan-petkovic]

Motivation

resolve: https://amd-hub.atlassian.net/browse/AIROCMLIR-1090

Enabling the weekly CI stages for rocmlirTriton (#361) turned on the attention parameter sweep for the first time, which surfaced an attention configuration that crashes the GPU with an illegal memory access (hipErrorIllegalAddress). This PR fixes the underlying out-of-bounds memory access in the rock→Triton pointer lowering so masked-out lanes can never address out-of-bounds memory.

Technical Details

Root cause: in RockTransformsToPointerArithPass (mlir/lib/Dialect/Rock/Transforms/TransformsToPointerArith.cpp), validity-impacting transforms (e.g. Pad, introduced when a non-tile-aligned dimension is padded up to the block size) pass their out-of-bounds coordinate straight through to the computed offset and only record the out-of-boundsness in the mask tensor. On AMD, a masked tt.load is lowered as load-then-select (a buffer op over the full tt.pointer_range=32 / ~2 GB window, with the mask applied to the loaded value via arith.select, not to the memory access), so a masked-out lane whose offset is out of bounds still issues the access and can page-fault.

How it surfaced: a grouped-query-attention (GQA) + KV-cache flash-decoding attention config pads the first-GEMM M dimension (= numRepeats = numHeadsQ/numHeadsKV) up to MPerBlock; the masked padding rows of the pre-softmax scale/bias load then addressed past the end of the scale/bias buffer, faulting whenever that address fell on an unmapped page (which is why the failure looked non-deterministic across head_dim_v values — head_dim_v only shifts the allocation layout). Confirmed with rocgdb: the faulting instruction is a buffer_load from the scale buffer at an offset well past its end.

Fix: in TransformsToPtrRewritePattern, zero the offset of masked-out (invalid) lanes before forming the pointer, so every lane addresses in-bounds memory. The mask still discards the loaded value for loads, and stores are mask-predicated and therefore never write, so this is safe for both loads and stores. It is only emitted for non-trivial masks, so fully in-bounds memory ops keep their original offset chain (and its buffer-op offset analysis) untouched. Net change: 1 file, +23 / -7.

Test Plan

Reproduced and validated on real gfx90a (MI250) hardware, standalone (rocmlir-gen | rocmlir-driver --host-pipeline=highlevel | rocmlir-driver -c | rocm-run): the exact failing weekly config; a full head_dim_v sweep over the previously-faulting values; a GQA-ratio × MPerBlock matrix (the padded-M cases); f16/bf16/f32; and regressions for non-GQA, non-KV-cache attention and plain gemm. Root cause independently confirmed under rocgdb (faulting buffer_load address vs. buffer bounds) and by toggling scale/bias on/off.

Test Result

Before: the GQA + KV-cache attention config aborts with hipErrorIllegalAddress (GPU memory fault). After: all of the above pass with correct verification ([1 1 1]) and no local regressions; the previously-faulting head_dim_v values and GQA/MPerBlock combinations all pass. Full lit regression and the weekly attention sweep will be validated on CI.

rocMLIR back-port

Not applicable — this change is rocmlirTriton-only. RockTransformsToPointerArithPass and the rock.transforms_to_ptr op it lowers do not exist in ROCm/rocMLIR (verified against the rocMLIR develop tree: no TransformsToPointerArith.cpp, no TransformsToPtrOp, and no Triton ops emitted anywhere under mlir/lib/Dialect/Rock/). This pass is part of the rocmlirTriton-specific Triton bridge — it emits triton::SplatOp/triton::BroadcastOp that feed RockTensorToTritonPtr — and the fix specifically targets AMD's Triton masked-load lowering (load-then-select), which has no rocMLIR counterpart. No upstream back-port is required.

Submission Checklist

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

[Copilot]

Pull request overview

Fixes an out-of-bounds memory access risk in Rock → Triton pointer lowering by clamping masked-out lanes’ computed offsets to a safe in-bounds value before forming the pointer tensor. This addresses AMD’s masked-load lowering behavior (load-then-select), where masked-out lanes may still issue memory accesses.

Changes:

• Broadcast isValid into a per-lane mask tensor and use it to zero offsets for invalid lanes prior to pointer formation.
• Skip emitting the offset-clamp when the validity mask is trivially all-true (scalar constant).

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[rocmlir-pr-reviewer]

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

---

Scope

Single-file fix in mlir/lib/Dialect/Rock/Transforms/TransformsToPointerArith.cpp (+23/-7). In TransformsToPtrRewritePattern, the offset of masked-out (invalid) lanes is zeroed via arith.select before forming the pointer tensor, so masked lanes with out-of-bounds offsets (e.g. from a Pad transform) no longer address OOB memory. This fixes a hipErrorIllegalAddress in a GQA + KV-cache flash-decoding attention config, because AMD lowers a masked tt.load as load-then-select (the access still issues for masked lanes).

Findings

• TransformsToPointerArith.cpp:622 — the file is under mlir/lib/Dialect/Rock/ (a rocMLIR-shared path) and is not one of the enumerated Triton-only bridge files; the PR description has no rocMLIR back-port note. (Major)
• TransformsToPointerArith.cpp:637 — the change alters emitted IR (new arith.select for masked ops) but adds no Lit/FileCheck test asserting it. (Major)

Notes

The core logic is correct: arith.select(maskTensor, offset, zeroOffset) shapes are consistent (all shape); the trivial-mask guard correctly matches the initial constant-true isValid; and for masked lanes the pointer collapses to base (in-bounds), which is safe for loads (value discarded) and predicated stores (never written). The cast<RankedTensorType>(offset.getType()) is safe because baseAddrSplat is always a ranked splat and ensureCompatible keeps offset a matching tensor. Restricting the select to non-trivial masks preserves the original offset chain (and downstream buffer-op offset analysis) for fully in-bounds ops.

CI status

No non-self CI failures. py-checks and detect pass; review / copilot-pull-request-reviewer are in progress (the auto-review pipeline itself). The full Lit regression and weekly attention sweep are still pending per the PR description.

[Copilot]

Pull request overview

Copilot reviewed 2 out of 2 changed files in this pull request and generated 1 comment.

[rocmlir-pr-reviewer]

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

---

Scope

Fixes an out-of-bounds GPU memory access in RockTransformsToPointerArithPass: for non-trivial validity masks, masked-out lane offsets are now zeroed with an arith.select before the pointer add, so AMD's load-then-select masked-load lowering can never fault on a masked lane. Touches one source file and one lit test.

Findings

• Minor: mlir/test/Dialect/Rock/rock-transforms-to-pointer-arith.mlir:18 — the CHECK-NOT: arith.select guarding the trivial-mask path is scoped after the load, so it does not actually catch a clamp that would be emitted before it.

Notes

• The core fix is correct and well-scoped: trivialMask is only set when isValid is a constant true, so fully in-bounds ops keep their original offset chain (and downstream buffer-op offset analysis) untouched; the arith.select(mask, offset, zero) clamp is emitted only for non-trivial masks. Safe for both loads (value discarded by mask) and stores (mask-predicated).
• The positive assertion in @test_pad_mask correctly pins the emitted arith.select(mask, offset, zero) shape.
• rocMLIR back-port: the PR description now includes a ## rocMLIR back-port note stating this pass/op is rocmlirTriton-only; this satisfies the back-port check for the shared mlir/lib/Dialect/Rock/ path.

Re-review

Both prior findings (rocMLIR back-port note; missing FileCheck coverage) are resolved on this revision. Only one new Minor test-robustness point remains.

[rocmlir-pr-reviewer]

Verdict: APPROVE -- submitted as COMMENT (automated reviews are advisory)  ·  New findings: 0 (0 Critical, 0 Major, 0 Minor)

---

Scope

Fixes an out-of-bounds GPU memory access (hipErrorIllegalAddress) in RockTransformsToPointerArithPass. Validity-impacting transforms (e.g. Pad) previously passed an OOB coordinate straight through to the computed offset; since AMD lowers masked tt.load as load-then-select, a masked-out lane with an OOB offset still issued the access and could page-fault. The fix zeroes masked-out lane offsets before forming the pointer, gated to non-trivial masks so all-valid ops keep their original offset chain. Touches 1 source file and its lit test (+27 / -11).

Findings

No blocking issues found. The fix is correct: arith.select(mask, offset, 0) keeps valid lanes' offsets and zeroes invalid ones; the mask still discards the loaded value on loads and predicates stores, so it is safe for both. The trivial-mask fast path is preserved via a scalar arith.constant bool check, keeping downstream buffer-op offset analysis untouched for fully in-bounds ops.

Notes

Re-review: all three previously-flagged issues are resolved in the current head.

• FileCheck coverage added: @test_pad_mask asserts the arith.select(mask, offset, 0) clamp is emitted for a Pad-induced non-trivial mask; @test_transforms_to_ptr_load asserts (CHECK-NOT: arith.select) no clamp for the trivial mask.
• The CHECK-NOT: arith.select is now bounded between the mask tt.splat and rock.blockwise_load_ptr, so it scans the pointer-forming region where a regressed clamp would appear.
• rocMLIR back-port: the PR description explains RockTransformsToPointerArithPass / rock.transforms_to_ptr are rocmlirTriton-only (no equivalent in ROCm/rocMLIR), satisfying the back-port-note requirement.

Minor (non-blocking, no action needed): a tensor-splat all-true constant mask would not be recognized as trivial and would take the clamp path, emitting a redundant (but harmless and correct) arith.select. The common trivial case is a scalar i1 true, which is handled.

CI status

No non-self checks report fail/cancel. Jenkins and ml-ci-internal.amd.com are pending; the auto-review pipeline's own review check being in-progress is expected.

[Copilot]

Pull request overview

Copilot reviewed 2 out of 2 changed files in this pull request and generated no new comments.

[justinrosner]

Can you also add a new E2E test for the case where this failure was first observed?

[justinrosner]

Can you move buildMaskTensor out of the if/else? If it is generating the same thing in each case then we should be able to pull this out?

[justinrosner]

If we call it once it would also remove the need for it to be a lambda?

[dhernandez0]

n AMD, a masked tt.load is lowered as load-then-select (a buffer op over the full tt.pointer_range=32 / ~2 GB window, with the mask applied to the loaded value via arith.select, not to the memory access), so a masked-out lane whose offset is out of bounds still issues the access and can page-fault.

the buffer op is able to do this, why do we do load-then-select? this looks like a bug in triton

the faulting instruction is a buffer_load from the scale buffer at an offset well past its end.

Buffer ops are designed for this use case, to return 0 when the index is out of bounds. This looks like a triton bug that we should solve in triton rather than adding extra code on our side (which is not efficient because we will add extra assembly to mask pointers). Please, can you investigate why triton accesses masked out pointers?

[bogdan-petkovic]

n AMD, a masked tt.load is lowered as load-then-select (a buffer op over the full tt.pointer_range=32 / ~2 GB window, with the mask applied to the loaded value via arith.select, not to the memory access), so a masked-out lane whose offset is out of bounds still issues the access and can page-fault.

the buffer op is able to do this, why do we do load-then-select? this looks like a bug in triton

the faulting instruction is a buffer_load from the scale buffer at an offset well past its end.

Buffer ops are designed for this use case, to return 0 when the index is out of bounds. This looks like a triton bug that we should solve in triton rather than adding extra code on our side (which is not efficient because we will add extra assembly to mask pointers). Please, can you investigate why triton accesses masked out pointers?

You're right, this is a Triton bug and this PR is just a workaround. I dug into it, here is what's going on.

Buffer ops do return 0 on OOB, but only for the voffset against num_records, and num_records is hardcoded to ~2GB instead of the real buffer size. Masked lanes are kept safe by pushing their voffset past that, and that part works fine.

What breaks it is the soffset optimization (337). It moves the uniform part of the address into soffset, but the hardware only bounds-checks the voffset and then adds soffset without checking it, and the safety pass only checks the offset is non-negative, not in bounds. So for a masked load where the uniform part is out of bounds (the padded GQA case on the last workgroup) the OOB lands in soffset, skips the check and faults. The masking doesn't help because it only touches voffset.

I tested it on gfx90a. Drop the clamp and just skip the split annotation for that load and the fault is gone with the right result. Leave the split on with no clamp and it faults.

My plan is to fix it in Triton and not set split_soffset_safe on ops that have a real mask. Those keep the whole offset in voffset so the OOB lanes stay safe, and everything else keeps the optimization. This only turns the optimization off for masked ops, the common in-bounds loads still get it, so we keep the conv perf wins. What do you think about that fix?

[dhernandez0]

n AMD, a masked tt.load is lowered as load-then-select (a buffer op over the full tt.pointer_range=32 / ~2 GB window, with the mask applied to the loaded value via arith.select, not to the memory access), so a masked-out lane whose offset is out of bounds still issues the access and can page-fault.

the buffer op is able to do this, why do we do load-then-select? this looks like a bug in triton

the faulting instruction is a buffer_load from the scale buffer at an offset well past its end.

Buffer ops are designed for this use case, to return 0 when the index is out of bounds. This looks like a triton bug that we should solve in triton rather than adding extra code on our side (which is not efficient because we will add extra assembly to mask pointers). Please, can you investigate why triton accesses masked out pointers?

You're right, this is a Triton bug and this PR is just a workaround. I dug into it, here is what's going on.

Buffer ops do return 0 on OOB, but only for the voffset against num_records, and num_records is hardcoded to ~2GB instead of the real buffer size. Masked lanes are kept safe by pushing their voffset past that, and that part works fine.

What breaks it is the soffset optimization (337). It moves the uniform part of the address into soffset, but the hardware only bounds-checks the voffset and then adds soffset without checking it, and the safety pass only checks the offset is non-negative, not in bounds. So for a masked load where the uniform part is out of bounds (the padded GQA case on the last workgroup) the OOB lands in soffset, skips the check and faults. The masking doesn't help because it only touches voffset.

I tested it on gfx90a. Drop the clamp and just skip the split annotation for that load and the fault is gone with the right result. Leave the split on with no clamp and it faults.

My plan is to fix it in Triton and not set split_soffset_safe on ops that have a real mask. Those keep the whole offset in voffset so the OOB lanes stay safe, and everything else keeps the optimization. This only turns the optimization off for masked ops, the common in-bounds loads still get it, so we keep the conv perf wins. What do you think about that fix?

That pass was removed from triton temporarily until they fix the issues. We recovered it because it was good for perf and we didn't see any accuracy issues until now. I think either we fix it or remove the pass.

[justinrosner]

n AMD, a masked tt.load is lowered as load-then-select (a buffer op over the full tt.pointer_range=32 / ~2 GB window, with the mask applied to the loaded value via arith.select, not to the memory access), so a masked-out lane whose offset is out of bounds still issues the access and can page-fault.

the buffer op is able to do this, why do we do load-then-select? this looks like a bug in triton

the faulting instruction is a buffer_load from the scale buffer at an offset well past its end.

Buffer ops are designed for this use case, to return 0 when the index is out of bounds. This looks like a triton bug that we should solve in triton rather than adding extra code on our side (which is not efficient because we will add extra assembly to mask pointers). Please, can you investigate why triton accesses masked out pointers?

You're right, this is a Triton bug and this PR is just a workaround. I dug into it, here is what's going on.
Buffer ops do return 0 on OOB, but only for the voffset against num_records, and num_records is hardcoded to ~2GB instead of the real buffer size. Masked lanes are kept safe by pushing their voffset past that, and that part works fine.
What breaks it is the soffset optimization (337). It moves the uniform part of the address into soffset, but the hardware only bounds-checks the voffset and then adds soffset without checking it, and the safety pass only checks the offset is non-negative, not in bounds. So for a masked load where the uniform part is out of bounds (the padded GQA case on the last workgroup) the OOB lands in soffset, skips the check and faults. The masking doesn't help because it only touches voffset.
I tested it on gfx90a. Drop the clamp and just skip the split annotation for that load and the fault is gone with the right result. Leave the split on with no clamp and it faults.
My plan is to fix it in Triton and not set split_soffset_safe on ops that have a real mask. Those keep the whole offset in voffset so the OOB lanes stay safe, and everything else keeps the optimization. This only turns the optimization off for masked ops, the common in-bounds loads still get it, so we keep the conv perf wins. What do you think about that fix?

That pass was removed from triton temporarily until they fix the issues. We recovered it because it was good for perf and we didn't see any accuracy issues until now. I think either we fix it or remove the pass.

I would argue that since we are very deep into perf testing and really care about optimal performance right now, that we should keep the AnnotateBufferSoffset pass and come up with a targeted fix for the time being, knowing that eventually this will all be superseded by upstream Triton when they eventually rewrite the pass. Without this pass we saw significant performance drops.

@umangyadav @pabloantoniom @dhernandez0 do you guys have any strong opinions here?

[umangyadav]

Buffer ops do return 0 on OOB, but only for the voffset against num_records, and num_records is hardcoded to ~2GB instead of the real buffer size. Masked lanes are kept safe by pushing their voffset past that, and that part works fine.

I ran into this issue earlier here #347

Upstream rocMLIR tolerates that over-read: it lowers global loads to amdgpu.raw_buffer_load, whose descriptor carries num_records = the real allocation, so the hardware clamps out-of-bounds reads to zero (rock even marks such loads canReadOffEnd). Triton's AMD buffer lowering (BufferOpsEmitter) instead builds every descriptor with num_records = INT_MAX-1, which disables that hardware bound — so on gfx950 the over-read hits unmapped memory and faults intermittently.

It would be good to check if we can clamp the loop bounds as PR#347 and see if it can fix the issue for now

[umangyadav]

@umangyadav @pabloantoniom @dhernandez0 do you guys have any strong opinions here?

I think it would be good to check if we can fix loop bounds in our code first and if not then yes we can do some temporary workarounds

[bogdan-petkovic]

@umangyadav @pabloantoniom @dhernandez0 do you guys have any strong opinions here?

I think it would be good to check if we can fix loop bounds in our code first and if not then yes we can do some temporary workarounds

@umangyadav I pulled your #347 branch and it fixes our weekly fault. I think we should just reopen #347 and land it there, and I'll close mine.

[bogdan-petkovic]

Fix is in this PR: #347