[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