Skip to content

feat(shard): express swizzled composed layouts as CuTe does - #174

Merged
zhen8838 merged 1 commit into
tile-ai:mainfrom
zhen8838:feat/shard-swizzle-layout
Sep 21, 2026
Merged

zhen8838 merged 1 commit into
tile-ai:mainfrom
zhen8838:feat/shard-swizzle-layout

Conversation

@zhen8838

@zhen8838 zhen8838 commented Sep 21, 2026

Copy link
Copy Markdown
Collaborator

Why

What

  • Swizzle(bits, base, shift) ports CuTe's Swizzle<B,M,S> as a frozen functor that ComposedLayout.inner may carry. It is not a LayoutBase: its domain comes from the composition's outer.
  • layout_algebra gains CuTe's swizzle specializations: apply, cosize, coalesce, the involution inverses, and composition — including composition(Layout, Swizzle), which rebuilds the swizzle over the bits the layout leaves active so the answer is again a swizzle-inner composed layout rather than a Swizzle in outer.
  • CUDA codegen renders the composition unchanged, cute::make_composed_layout(cute::Swizzle<B,M,S>{}, cute::Int<offset>{}, ...), for both the layout type and the layout value.
  • The device runtime reads steps off the affine part of a shard layout (CuTe deletes stride() on a ComposedLayout) and projects a composed one by folding the instance's origin into the composition's offset instead of advancing the engine pointer.
  • One new test carries a (128, 4) f32 tile through swizzled SMEM: print/parse round-trip, codegen, nvcc, and the device, plus the emitted-layout assertion that running the kernel cannot make.

Contract

  • ComposedLayout.inner widens to LayoutBase | Swizzle | None; outer, ShardLayout.layout and TensorType.layout are unchanged, so a bare Swizzle is still not a tensor layout. docs/spec/shard.md §4 and the new §4.1 state it.
  • The device projection rule for a composed shard layout is new public behavior: docs/spec/runtime.md §2.4.1 states that the origin reaches the composition's offset and not the engine pointer, and that a composed layout does not take the whole-tensor shortcut.
  • A Swizzle states where an element lives, never which element a logical coordinate names: shape, domain_rank, the axis numbering a Split references and the access relation an op states are read off outer and are the unswizzled ones. docs/spec/shard.md §4.1 states it.
  • A swizzled layout is still refused as a mesh execution scope, unchanged.

Risk

  • Transpose, Reshape and Slice carry the swizzle, since all three only relabel the domain and a window's start is a constant shift of the index that offset already holds. A register engine refuses it by name, as do a composed inner that is not a Swizzle and a Swizzle over a ShardLayout outer, which the shard slice path would drop.
  • composition(Layout, Swizzle) and the inverses reject a non-zero composition offset: CuTe's rule for that case puts a Swizzle on the domain side, which this IR cannot represent.
  • The ported composition and swizzled inverses have no caller in src/ yet, so nothing exercises them; this PR ships them because the issue asks for CuTe's algebra, not because a code path needs them today.
  • No WGMMA instruction and no performance gate ship here; this is the layout the operand needs, not the instruction that consumes it.

@zhen8838
zhen8838 force-pushed the feat/shard-swizzle-layout branch from e051166 to 1f87d3f Compare September 21, 2026 08:21
Comment thread src/tilefoundry/ir/hir/tensor/slice.py Outdated
Comment thread src/tilefoundry/ir/types/shard/layout.py
@zhen8838
zhen8838 force-pushed the feat/shard-swizzle-layout branch from 1f87d3f to 8985cd4 Compare September 21, 2026 08:35
Port CuTe's Swizzle<B,M,S> as a frozen value type that ComposedLayout may
carry in `inner`, so a shared-memory tile can state the XOR swizzle a WGMMA
operand needs instead of only an affine stride rule. The layout algebra gains
CuTe's swizzle specializations -- apply, cosize, coalesce, the involution
inverses and the composition that canonicalizes a Swizzle back to the inner
side rather than into an `outer` that has no domain.

The CUDA path emits the composition unchanged, as
cute::make_composed_layout(cute::Swizzle<B,M,S>{}, cute::Int<offset>{}, ...),
for both the layout type and the layout value. Because CuTe deletes
stride() on a ComposedLayout, the device runtime now reads steps off the
affine part of a shard layout, and projects a composed one by folding the
instance's origin into the composition's offset instead of advancing the
engine pointer: the swizzle runs on the whole tensor's index, so pointer
arithmetic would hand two instances one address.

A Swizzle states where an element lives, never which element a logical
coordinate names, so shape, domain rank, the axis numbering a Split
references and the access relation an op states are all read off `outer` and
are the unswizzled ones. Transpose, Reshape and Slice therefore carry the
same Swizzle through, a window's start being a constant shift of the index
that `offset` already holds. A register engine refuses it by name rather
than rebuilding the layout from register strides.
@zhen8838
zhen8838 force-pushed the feat/shard-swizzle-layout branch from 8985cd4 to 0198303 Compare September 21, 2026 08:44
@zhen8838
zhen8838 merged commit 1631e23 into tile-ai:main Sep 21, 2026
1 check passed
@zhen8838
zhen8838 deleted the feat/shard-swizzle-layout branch September 21, 2026 09:44
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant