[LoadStoreOpToLLVM] Fix packed-type OOB boundary in NaN-padded block loads (#6002) - #7828
Draft
NathanVoldman wants to merge 20 commits into
Draft
[LoadStoreOpToLLVM] Fix packed-type OOB boundary in NaN-padded block loads (#6002)#7828NathanVoldman wants to merge 20 commits into
NathanVoldman wants to merge 20 commits into
Conversation
NathanVoldman
changed the title
NathanVoldman
force-pushed
the
dev/nvoldman/fix-6002-block-io-nan-packed-boundary
branch
from
August 26, 2026 14:20
6e4f6e2 to
d7e7a49
Compare
Extends test_block_load_dpas_layout to cover float16 (opsPerChan=2, kWidth=2) in addition to float32. The float16 path exercises the VNNI B-operand packed boundary: the B descriptor has odd K dimension (N-1), so the last in-bounds K row is packed with the first OOB row in a single i32 word. With NaN padding, that in-bounds row must remain 1.0 and the OOB row must become NaN. Also replaces the zero-padding reference approach with explicit expected tensors (1.0 in-bounds, NaN at the border row/col). The zero-padding reference is unreliable for packed fp16 operands: hardware OOB checks at i32 granularity can cause the reference to read OOB elements as 1.0 instead of 0.0, making the zero->NaN conversion silently incorrect. The explicit expected tensors directly express the correct NaN-padding semantics and are independent of hardware zero-padding behaviour. Tested on BMG (Arc Pro B65): 18/18 passed (9 float32 + 9 float16). Co-Authored-By: Claude Sonnet 4.6 (1M context) <noreply@anthropic.com>
…loads (#6002) Hardware 2D block load instructions with elem_size_in_bits=32 (used for packed fp16, numPackedVals=2) perform OOB boundary checks at i32-word granularity. When base_width is not a multiple of the packed element size, the last in-bounds element shares a packed word with the first OOB element. On platforms with coarse-grained boundary protection (e.g. PVC Max 1100), the hardware zeroes the entire packed word, giving 0.0 for the in-bounds element instead of the correct value. The software NaN mask then propagates this incorrect 0.0 through the select operation. Fix: in Subgroup2DBlockLoadOpConversion::computeAddress, when pad_nan is active and numPackedVals > 1, apply the same pointer-shift + surface-widening pattern already used by Subgroup2DBlockLoadFromPtrOpConversion. Shifting addrElem back by offsetX elements and widening hwBaseWidth by the same byte count ensures that the sub-tile's packed words fall within the hardware's in-bounds region, while the software NaN mask (nanMaskElems) retains responsibility for correctly marking actual OOB elements as NaN. The fix is harmless on hardware with fine-grained OOB protection (BMG, PVC Max 1550) where the existing tests already pass. Co-Authored-By: Claude Sonnet 4.6 (1M context) <noreply@anthropic.com>
b.sub() returns a SubOp wrapper, not a Value. Passing it directly to b.gep caused a compilation error: no known conversion from 'mlir::LLVM::SubOp' to 'ArrayRef<GEPArg>' Store the result as Value (same pattern as line 3060) before the GEP. Co-Authored-By: Claude Sonnet 4.6 (1M context) <noreply@anthropic.com>
… block loads (#6002) The previous pointer-shift approach only widened the hardware surface for sub-tiles at offsetX > 0. For sub-tiles at offsetX=0 (covering all N columns in a single hardware instruction), offsetXBytes=0 made the fix a no-op, leaving base_width=30 unchanged. Replace with a round-up approach: unconditionally round base_width up to the nearest packed-element boundary (4 bytes for fp16). This ensures every packed word within the tile is in-bounds from the hardware's perspective regardless of sub-tile offset. Example: base_width=30 bytes (15 fp16) -> 32 bytes (8 packed i32 words). Hardware now sees packed word 7 (fp16 cols 14-15) as in-bounds instead of borderline OOB (floor(30/4)=7, so 7<7=False without the fix). Co-Authored-By: Claude Sonnet 4.6 (1M context) <noreply@anthropic.com>
… block loads (#6002) The previous fix used cfg.numPackedVals > 1 as the guard, but the hardware instruction for A (opIdx=0, kWidth=1) uses elem_size_in_bits=16 with numPackedVals=1 — so the condition was always false for A and the fix never applied. Root cause (confirmed on X1 cluster, PVC Max 1100): the hardware applies OOB boundary checks at i32 (4-byte) granularity for 2D block loads, regardless of elem_size_in_bits. For base_width=30 bytes (15 fp16), internal packed word 7 (fp16 cols 14-15) triggers a 7<7=False OOB check, zeroing both cols even though col 14 is in-bounds. Fix: whenever NaN-padding is active, unconditionally round base_width up to the nearest 4-byte boundary. For float32 or already-aligned shapes this is a no-op. For 15 fp16 cols (30 bytes) it rounds to 32, making the hardware check 7<8=True so col 14 is correctly loaded as 1.0. Co-Authored-By: Claude Sonnet 4.6 (1M context) <noreply@anthropic.com>
… block loads (#6002) Also round base_height up to nearest 2 rows to fix the VNNI K-row boundary case. For VNNI B operands, PVC Max 1100 also applies i32 granularity to the Y (row) direction: K row pairs (K=2i, K=2i+1) share one i32 unit, so base_height=15 causes 7<7=False OOB, zeroing K=14. Rounding base_height up to the nearest 2 (e.g. 15->16) ensures every K-row pair is in-bounds. This fixes the two remaining test failures: float16-256-64: B K-rows 62/63 (base_height=63 -> 64) float16-128-16: B K-rows 14/15 (base_height=15 -> 16) Co-Authored-By: Claude Sonnet 4.6 (1M context) <noreply@anthropic.com>
… block loads (#6002) Restrict base_height rounding to VNNI loads only. Unconditionally rounding hwBaseHeight for non-VNNI (A operand) loads was causing non-deterministic memory corruption at arbitrary positions on Max 1100 for specific shapes (e.g. M=256 with hwBaseHeight=256). The mechanism is unclear but the symptom is values like 2.0, -0.0, or NaN at in-bounds positions that should be 1.0. The A operand does not pair rows in the same way VNNI does, so the base_height i32-granularity issue does not apply — the Y boundary check is per-row for A. Co-Authored-By: Claude Sonnet 4.6 (1M context) <noreply@anthropic.com>
… block loads (#6002) Replace the unsafe base_width/height rounding approach (which corrupted memory when base_width > pitch) with a skip-and-fallback strategy in LowerTo2DBlockLoad. Root cause: PVC Max 1100 applies OOB checks at i32 (4-byte) granularity for 2D block loads. When base_width is not 4-byte aligned (fp16 odd-col boundary) or base_height is not 2-row aligned (VNNI odd-K boundary), the hardware zeroes an entire i32 word including its in-bounds element. Rounding up base_width to fix this is unsafe when base_width would exceed the row pitch, causing address corruption. Fix: in LowerTo2DBlockLoad::convertDescriptorLoadOp, when NaN-padding is requested and the inner shape is not i32-aligned (or outer for VNNI is not 2-aligned), skip the 2D block load conversion. The descriptor load op falls through to DescriptorLoadOpConversion which uses correct scalar (gather) loads. This trades performance for correctness on the boundary case (odd-dimension tensors with NaN padding). Co-Authored-By: Claude Sonnet 4.6 (1M context) <noreply@anthropic.com>
… block loads (#6002) Implement pitch-aware base_width rounding instead of unconditional scalar fallback for the column direction. The previous fix always fell back to scalar loads when base_width was not 4-byte aligned. This caused a performance regression on all hardware (BMG, Max 1550) for NaN-padded descriptor loads with odd fp16 column counts. The hardware constraint is: base_width % 4 == 0 AND pitch >= base_width. We can satisfy both by rounding base_width up — but only when pitch allows: rounded = ceil(base_width / 4) * 4 if pitch >= rounded: use 2D block load with rounded base_width (spec- compliant, correct on Max 1100 and all other HW) if pitch < rounded: scalar fallback (pitch too small to safely widen) For real-world matmuls where the tensor's row stride (pitch) >= the rounded boundary width, 2D block load is preserved. Scalar fallback only triggers for contiguous tensors where pitch == inner_shape_bytes (no padding room). The K-row direction (VNNI base_height) retains the conservative odd-fallback since there is no pitch-equivalent in the height direction. Co-Authored-By: Claude Sonnet 4.6 (1M context) <noreply@anthropic.com>
… block loads (#6002) Three changes: 1. LowerTo2DBlockLoad: skip non-constant-shape descriptors instead of falling back to scalar. When getFoldedConstantValue returns nullopt (runtime shape, e.g. function arguments), we cannot determine alignment statically — proceed with 2D block load as before. Only fall back when the shape IS a known constant and is non-4-byte-aligned with insufficient pitch. Fixes regression in existing test_block_load_dpas_layout (which uses runtime shapes + pad_nan and expects ttig.2d_block_load). 2. Lit test: test/TritonIntelGPU/LowerTo2DBlockLoad/ descriptor-load-nan-boundary.mlir — two cases: (a) shape=15, stride=15 (contiguous): pitch=30 < rounded=32 -> scalar fallback, tt.descriptor_load survives. (b) shape=15, stride=16 (padded): pitch=32 >= rounded=32 -> 2D block load emitted with arith.constant 32 as base_width. 3. Accuracy test: test_block_load_dpas_layout_pitch_rounding — uses stride=[N,1] (full tensor stride), so pitch = N*2 bytes. For N=16: (N-1)*2=30 not 4-aligned, rounded=32, pitch=32 >= 32 -> 2D block load path. Verifies col N-2=1.0 and col N-1=NaN on Max 1100, confirming that the rounded base_width correctly prevents coarse-grained OOB zeroing of the last valid column. Co-Authored-By: Claude Sonnet 4.6 (1M context) <noreply@anthropic.com>
…#6002) When base_width is rounded up to satisfy the 4-byte alignment constraint (pitch >= rounded), the NaN mask column bound must use the ORIGINAL inner shape rather than the rounded base_width / elemBytes. Previously the mask considered the extra rounded-up column as in-bounds, leaking the padding slot value instead of NaN. Fix: add optional nan_mask_col_bound attribute to ttig.2d_block_load. LowerTo2DBlockLoad sets it to the original innerSh when rounding occurs. LoadStoreOpToLLVM uses it in buildNaNMasks instead of udiv(baseWidth, elemBytes). Test: update test_block_load_dpas_layout_pitch_rounding to use a sentinel value (99.0) in the OOB padding slot so a leaked value is unmistakable. Update lit test to verify nan_mask_col_bound is emitted on the rounded path.
…col bound bug Single-warp DPAS layout (warpsPerCTA=[1,1]) ensures lane 15 owns col 15 = N-1. Uses sentinel value 99.0 in the OOB padding slot to make a leaked padding value unmistakably distinct from the expected NaN.
Output descriptor must use full [M, N] shape so row M-1 and col N-1 are actually written (previously [M-1, N-1] left those positions uninitialized).
…r the bug Previous design had two issues: 1. Constant M-1 caused the outer-shape scalar fallback to trigger (odd M-1 = 7) — use runtime %arg2 for M-1 so getFoldedConstantValue cannot fold it. 2. N=16 gave pitch=32 bytes < 64 byte minimum for 2D block load — use N=32 (pitch=64 bytes, the minimum) to ensure the 2D block load path is taken. With N=32: K_tile=16 x 2 reps, lane 15 covers col 31 = N-1 (OOB boundary). Bug: NaN mask bound = 32 (rounded) → col 31 in-bounds → sentinel 99.0 leaks. Fix: NaN mask bound = 31 (original innerSh) → col 31 OOB → NaN.
…ounding (#6002) Move base_width rounding from LowerTo2DBlockLoad.cpp to LoadStoreOpToLLVM.cpp. The nan_mask_col_bound attribute was needed only because the transform pass was rounding base_width too early (before the NaN mask was built). By keeping base_width at its original declared value in ttig.2d_block_load and applying the alignment rounding in LoadStoreOpToLLVM.cpp after the NaN mask is built, the mask automatically gets the correct original column bound. LowerTo2DBlockLoad.cpp: remove baseWidth rounding and nanMaskColBoundAttr. LoadStoreOpToLLVM.cpp: compute hwBaseWidth (rounded for HW) after NaN mask; use hwBaseWidth in computeAddress SubTileAddress return. TritonIntelGPUOps.td: remove nan_mask_col_bound attribute entirely.
…ounding (#6002) Remove getPadNan() guard and getFoldedConstantValue() check from hwBaseWidth rounding. The HW constraint base_width % max(4, elemSize) == 0 applies to all 2D block loads, not just pad_nan ones. Always emit the rounding as IR — for already-aligned constant values the ArithDialect canonicalizer folds it away.
NathanVoldman
force-pushed
the
dev/nvoldman/fix-6002-block-io-nan-packed-boundary
branch
from
August 27, 2026 08:04
479be66 to
d3993b1
Compare
…ounding (#6002) Fix pre-commit formatting (yapf/clang-format) and restore getPadNan() guard on hwBaseWidth rounding. Removing the guard caused unnecessary runtime arithmetic on non-pad_nan loads with runtime base_width, breaking the descriptor-load-block-2d.mlir lit test. Non-pad_nan loads are always aligned by LowerTo2DBlockLoad construction; the rounding is only needed for pad_nan boundary loads. The getFoldedConstantValue() check is still removed so runtime pad_nan base_width values are handled correctly.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
New contributor declaration
I am not making a trivial change, such as fixing a typo in a comment.
I have written a PR description following these
rules.
I have run
pre-commit run --from-ref origin/main --to-ref HEAD.Select one of the following.
/testforlittests/unittestfor C++ tests/python/testfor end-to-end testsFILL THIS IN.Select one of the following.
littests.littests I have added follow these best practices,including the "tests should be minimal" section. (Usually running Python code
and using the instructions it generates is not minimal.)