[Pallas:MGPU] Support transposing the swizzled dimension in `UnswizzleRef`. #41854

Closed
copybara-service[bot] opened 11:28am on October 21, 2026 wanted to merge 226 Ξ” into google/jax main from
test_1003935315

Pull Request Overview

  • Opened on October 21, 2026
  • Status Closed
  • Commit count 1 with first commit October 21, 2026

Total Delta

226 Total Diff Delta

Open Days

Open 17 weekdays

Test Delta

92 Diff Delta in Test Files
Breakdown by Phase

How long has this pull request spent in each phase of its lifecycle?

Data pending calculation for pull request

Author avatar

[Pallas:MGPU] Support transposing the swizzled dimension in `UnswizzleRef`.

[Pallas:MGPU] Support transposing the swizzled dimension in UnswizzleRef.

Previously, UnswizzleRef assumed the swizzled dimension was always the minormost dimension (-1) and rejected any TransposeTransform that moved it. Because of this restriction, _get_lowering_rule and _swap_lowering_rule had to manually pop UnswizzleRef before calling _handle_transforms: when a tiled and swizzled SMEM reference is transposed, commuting TransposeTransform past UntilingTransform also transposes the 2D tile dimensions, which moves the swizzled dimension away from -1.

To allow UnswizzleRef to flow through _handle_transforms uniformly with UntilingTransform (in a followup change), UnswizzleRef needs to track which dimension is swizzled (swizzled_dim) across TransposeTransform and NDIndexer rather than assuming it is always -1, while enforcing that operations requiring the swizzled dimension to remain minormost (undo and commute_reshape) or un-sliced (commute_ndindexer) check swizzled_dim.

PR was closed without comments.