feat(shard): express swizzled composed layouts as CuTe does - #174
Merged
Merged
Conversation
zhen8838
force-pushed
the
feat/shard-swizzle-layout
branch
from
September 21, 2026 08:21
e051166 to
1f87d3f
Compare
zhen8838
commented
Sep 21, 2026
zhen8838
force-pushed
the
feat/shard-swizzle-layout
branch
from
September 21, 2026 08:35
1f87d3f to
8985cd4
Compare
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
force-pushed
the
feat/shard-swizzle-layout
branch
from
September 21, 2026 08:44
8985cd4 to
0198303
Compare
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
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.
Why
LayoutandComposedLayoutcould only state an affine coord-to-offset mapping, so the XOR swizzle a WGMMA shared-memory operand needs (issue feat(shard): let a layout carry a non-affine transform #173) had nowhere to live.What
Swizzle(bits, base, shift)ports CuTe'sSwizzle<B,M,S>as a frozen functor thatComposedLayout.innermay carry. It is not aLayoutBase: its domain comes from the composition'souter.layout_algebragains CuTe's swizzle specializations:apply,cosize,coalesce, the involution inverses, andcomposition— includingcomposition(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 aSwizzleinouter.cute::make_composed_layout(cute::Swizzle<B,M,S>{}, cute::Int<offset>{}, ...), for both the layout type and the layout value.stride()on aComposedLayout) and projects a composed one by folding the instance's origin into the composition's offset instead of advancing the engine pointer.(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.innerwidens toLayoutBase | Swizzle | None;outer,ShardLayout.layoutandTensorType.layoutare unchanged, so a bareSwizzleis still not a tensor layout.docs/spec/shard.md§4 and the new §4.1 state it.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.Swizzlestates where an element lives, never which element a logical coordinate names:shape,domain_rank, the axis numbering aSplitreferences and the access relation an op states are read offouterand are the unswizzled ones.docs/spec/shard.md§4.1 states it.Risk
Transpose,ReshapeandSlicecarry the swizzle, since all three only relabel the domain and a window's start is a constant shift of the index thatoffsetalready holds. A register engine refuses it by name, as do a composedinnerthat is not aSwizzleand aSwizzleover aShardLayoutouter, 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 aSwizzleon the domain side, which this IR cannot represent.compositionand swizzled inverses have no caller insrc/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.