SIGN IN SIGN UP

[Bugfix] Make TMA layouts region-aware to keep slices contiguous (#3089)

[Bugfix] Make TMA linear layouts region-aware to keep versioned slices contiguous

The inferred "linear" shared layout hoisted every 256-split quotient above
all dims, so a pipeline-versioned slice became two gapped chunks: LowerBulk's
bijection check fired on warp-specialized wide reductions, and the TMA
atomic-add path (which has no such check) silently read the other version.

MakeTmaLinearLayout now takes the copied region and orders modes as
[region-fixed dims, box quotients, box contents], so a fixed slice owns one
contiguous run of whole TMA boxes and the existing rest-loop machinery emits
one instruction per box over the true (unaligned) global shape. The
atomic-add path shares the helper and computes the region base offset through
the layout instead of hand-rolled row-major strides over the physical shape.

Co-authored-by: Claude Fable 5 <noreply@anthropic.com>
Y
Yongqi Zhuo committed
0ce2a7bacf01ab4622d4745ca48b004f1caa14aa
Parent: 728f2c8
Committed by GitHub <noreply@github.com> on 8/26/2026, 7:32:05 AM