Express `element_arrangement`in Triton kernel

[airin711]

Background

tt.spyre_tensor_layout (introduced in #74, "RewriteDescriptorLayout:
physicalize tensor descriptors for stick-tiled device layouts") lets a
Triton-Spyre kernel annotate a tensor descriptor with a physical coordinate
decomposition — three parallel arrays phys_src/phys_op/phys_arg
describing, per physical dim, which logical dim it derives from and whether
via identity, floordiv, or mod. This is the Triton-DSL analog of
TensorArg.device_coordinates in torch-spyre's OpSpec IR — it says where
each element physically lives.

What tt.spyre_tensor_layout does not express and what Triton's ordinary
tensor type system (tl.float16, tl.bfloat16, etc.) also does not express
is how elements are arranged/packed within a stick for certain hardware
conversions. torch-spyre tracks this separately as element_arrangement
(ElementArrangement enum: STANDARD, DL16_TO_FP32, FP32_TO_DL16,
QFP8CH, QFP8WT, EXX2) on SpyreTensorLayout/TensorArg. Certain
fp16/bf16↔fp32 hardware conversions produce a "staggered" physical layout,
elements interleaved/reordered within a stick rather than laid out in simple
row-major order, and this is orthogonal to dtype: it's about intra-stick
element ordering, not the element's storage format. A Triton kernel that
needs to consume or produce a staggered tensor currently has no way to
declare that.

Here is the example OpSpec using the element_arrangement field:

KERNEL_NAME = "sdsc_fused__to_copy_0"
POOL_SIZE = 0
BUNDLE_SYMBOLIC_ARGS = True
INPUTS_PT = os.path.splitext(os.path.abspath(__file__))[0] + ".inputs.pt"

# Host (shape, dtype) per kernel arg, in arg_index order.
SHAPES = [
    ([4, 128], torch.float16),
    ([4, 128], torch.float32)
]

# Exact device layout each arg had when the real graph ran.
LAYOUTS = [
    SpyreTensorLayout(
        device_size=[2, 4, 64],
        stride_map=[64, 128, 1],
        device_dtype=DataFormats.SEN169_FP16,
        element_arrangement=ElementArrangement.STANDARD,
    ),
    SpyreTensorLayout(
        device_size=[4, 4, 32],
        stride_map=[32, 128, 1],
        device_dtype=DataFormats.IEEE_FP32,
        element_arrangement=ElementArrangement.DL16_TO_FP32,
    ),
]

ops = [
    OpSpec(
        op='dl16tofp32',
        is_reduction=False,
        iteration_space={sympify('c0'): (sympify('4'), 4), sympify('c1'): (sympify('128'), 2)},
        op_info={},
        symbolic_dim_bounds={},
        debug_handle=DebugHandle(id=4224585593578919848, source=SourceLoc(file='staggered_ea.py', start_line=6, start_col=0, end_line=None, end_col=None), aten_op='aten._to_copy.default', ir_chain=('convert_element_type', 'buf0'), fused_from=(), transform_history=()),
        args=[
            TensorArg(
                is_input=True, arg_index=0, device_dtype=DataFormats.SEN169_FP16,
                device_size=[2, 4, 64],
                device_coordinates=[sympify('floor(c1/64)'), sympify('c0'), sympify('Mod(c1, 64)')],
                allocation={'hbm': 0},
            ),
            TensorArg(
                is_input=False, arg_index=1, device_dtype=DataFormats.IEEE_FP32,
                device_size=[4, 4, 32],
                device_coordinates=[sympify('floor(c1/32)'), sympify('c0'), sympify('Mod(c1, 32)')],
                allocation={'hbm': 1},
                element_arrangement=ElementArrangement.DL16_TO_FP32,
            ),
        ]
    ),
]

Proposal

Extend tt.spyre_tensor_layout (and its tl.spyre_tensor_layout frontend
builtin) with an element_arrangement attribute, applying to the whole
descriptor (not per-dimension, since staggering is not a coordinate
operation like phys_src/phys_op/phys_arg):

tl.spyre_tensor_layout(desc, layout, element_arrangement="DL16_TO_FP32")

lowering to an element_arrangement attribute on tt.spyre_tensor_layout
alongside the existing phys_src/phys_op/phys_arg arrays. Defaults to
STANDARD/omitted when not specified, matching the "omit if default"
convention already used throughout torch-spyre's own OpSpec serialization
(torch_spyre/_inductor/spyre_kernel.py).

Why this matters

Without this, any Triton-Spyre kernel that touches a staggered fp16↔fp32
conversion has no way to describe that at the descriptor level — the
information exists in torch-spyre's OpSpec IR today but has no Triton-DSL
counterpart, making it impossible to hand-write or lower to a Triton kernel
that participates in one of these conversions without out-of-band knowledge.

References

• RewriteDescriptorLayout: physicalize tensor descriptors for stick-tiled device layouts #74 — introduces tt.spyre_tensor_layout, tl.spyre_tensor_layout, and
  RewriteDescriptorLayout (the mechanism this issue extends)
• Add the Spyre driver skeleton so kernel[grid](...) reaches the backend #106 — example of current LAYOUT=[...] usage in driver-launch fixtures
• torch_spyre/csrc/spyre_tensor_impl.h — ElementArrangement definition
• torch_spyre/_inductor/spyre_kernel.py — existing "omit if default"
  serialization convention this proposal mirrors

cc. @fabianlim

[fabianlim]

@airin711 adding it to spyre_tensor_layout seems reasonable, but the OpSpec indicates a conversion. I can see the following issues

• the suitability of the adding the EA hint to spyre_tensor_layout only makes sense if it only depends on the tensor-desc. Is this specific only to inputs of the function (e.g., coming from the HBM or another core's LX) or it can also be an internal SSA.
• if it can apply internally in the function, then we cannot tag it only to a desc. We need to be able to tag it to an internal variable.
• we could consider generalizing spyre tensor layout to also tag to other tt.types other than desc, but this is a bigger arch descision and change.
• also the danger of tagging to a variable, is that it is position dependent, what if the EA needs to be applied to a copy of the tensor.

[airin711]

@fabianlim Thanks for the comments. Walking through it changed the shape of the proposal a bit.

On your first three points (input-only vs. internal, tagging a variable, generalizing to other tt.types):

element_arrangement should be tagged on the descriptor, not on an arbitrary SSA value or a generalized tt.types attribute — but "descriptor" already covers more than "function input." tt.make_tensor_descriptor takes an arbitrary base pointer, so an internally-allocated buffer gets its own descriptor the same way an incoming argument does; tagging that descriptor is enough to cover the "internal, but materialized" case without needing to tag SSA values directly. That also addresses the position-dependence/copy concern: a descriptor is anchored to a specific allocation, so it doesn't have the fragility a tag on a compiler-manipulable SSA value would (surviving CSE, copy-propagation, etc.). So we don't think this needs the bigger "generalize spyre_tensor_layout to other tt.types" change — descriptor scope is sufficient for what a Triton kernel author actually writes.

Where it isn't sufficient: SFP-to-SFP fusion without materializing.

This is the case that doesn't reduce to a descriptor problem. If the Triton→KTIR path allows multiple SFP ops to chain directly on a converted value before any write back to LX/HBM (e.g. fp32todl16 feeding an add that consumes two still-staggered operands), there's no descriptor anywhere in that chain to hang the tag on.

Proposal, split across the two layers where it actually lives:

1. Triton side (this issue, tt.spyre_tensor_layout) — add element_arrangement as a descriptor-level attribute, as originally proposed. This is what a kernel author needs: declaring that a tensor they load/store crosses a memory boundary in a non-standard arrangement. Scope stays exactly as narrow as before.
2. KTIR side (a separate proposal, against ktir-mlir-frontend) — represent EA as a type-level property of df16 values, the same way df16 itself already exists as a type rather than a side-channel attribute. That gives the fusion/scheduling pass that chains SFP ops a real type to verify against (operands to add must agree on arrangement; the result's arrangement follows from the op's semantics) instead of an inferred, easily-dropped tag.

What do you think?

[fabianlim]

@airin711

After thinking for awhile I think the simplest is to do the following

• if we put a hint in tt.spyre_tensor_layout that is type dependent, e.g., fp32todl16, then we will fail the lowering if the user passes in a type that contradict this hint
• for consistency, it is best to have both input / out descs to have hints, so that we can gaurantee that the output will have the expected arrangements.
• if a tensor_descriptor has an arrangement hint, then during the lowering, we can do an analysis, and if there are ops that are not arrangement compatible, then we have to insert some special ops (e.g., like a spyreop) to change the arrangement. This will be done automatically during lowering, like how we propagate phy shapes when handling spyre tensor layouts.

In general, the hint should not be a something-to-something hint. the hint should be a state. Like a is expected to be an fp32 with a particular arrangement.

[airin711]

Thank you, @fabianlim .
I completely agree with the idea that the hint should represent a state rather than a transformation.

I do have one question, though: what exactly do you mean by type-dependent in this context?
Some element arrangements are tied to a specific data type, while others are not.
For example, fp32todl16 implies that the data has previously undergone an fp32todl16 conversion and therefore has the df16 dtype. However, arrangements such as standard or exx2 are not necessarily associated with a particular data type in the same way.

[fabianlim]

@airin711 I tried to consider the 4 possible scenarios with the RMS Norm. While this document needs to be refined it gives us a starting point to discuss

RMSNorm x element arrangement: 4 scenarios

RMSNorm forward:

def forward(self, x):                      # x: [mb, hidden], hidden == stick dim
    input_dtype = x.dtype
    x = x.to(torch.float32)                 # upcast
    variance = x.pow(2).mean(-1, keepdim=True)
    x = x * torch.rsqrt(variance + self.eps)
    return self.weight * x.to(input_dtype)  # downcast

Terms follow the FP32 Element Arrangement RFC:
standard = sequential stick order. staggered = all values correct, but
within-stick position doesn't match logical order (produced by a dl16tofp32/fp32todl16
cast; RFC's DL16_TO_FP32/FP32_TO_DL16). Safe rule: an op is safe on staggered input
iff it never consults within-stick position — unary pointwise, full-dim-over-stick
reduction, and broadcast-with-scalar binary ops are all safe; combining a staggered
operand with a standard (or differently-staggered) operand, without broadcast, is not.

We vary x's dtype and incoming arrangement at kernel entry. weight is always a
fresh HBM load, i.e. always standard.

"Arrangement" vs "rearrangement" — these are not the same thing. Every scenario
below changes x's arrangement (standard↔staggered) — that's just what
dl16tofp32/fp32todl16 do, and it's expected. "Rearrangement" (the RFC's term)
means something extra: an additional restickify/relayout op inserted only to fix an
otherwise-illegal EA combination, beyond whatever casts the source code already
contains. The question each scenario answers is whether the casts RMSNorm's own code
already has are enough to close every stagger, or whether an extra op is needed on
top of them.

1. input = fp16, standard (the common case)

step	op	arrangement
entry	—	fp16, standard
x.to(fp32)	dl16tofp32	opens: fp32, staggered
pow/mean/add/rsqrt/mul	pointwise + reduction	stays staggered (all safe: unary, full-stick reduction, scalar-broadcast)
x.to(fp16)	fp32todl16	closes: fp16, standard (free un-stagger — restoration)
weight * x	mul	both standard — safe

x's arrangement does change twice here (standard→staggered→standard) — that's the
normal, expected behavior of the two casts. No extra rearrangement op is needed
beyond those two casts, which the source already has. Confirmed live in
rmsnorm_fused/: exactly 8 sdscs, dl16tofp32→(mul,mean,add,rsqrt,mul all at
stick=32)→fp32todl16→mul — no 9th restickify/identity op appears anywhere.

2. input = fp16, staggered (x is itself the unconsumed output of an earlier fp32→fp16 downcast)

step	op	arrangement
entry	—	fp16, staggered
x.to(fp32)	dl16tofp32	input already staggered → this upcast closes the old bracket instead of opening one: fp32, standard
pow/mean/add/rsqrt/mul	pointwise + reduction	stays standard (trivially safe)
x.to(fp16)	fp32todl16	input standard → opens a new bracket: fp16, staggered
weight * x	mul	weight standard vs x staggered, no broadcast → illegal

Rearrangement needed before the final weight * x. The closing downcast produces
a fresh stagger that nothing downstream absorbs (the next consumer is a plain mul,
not another upcast) — un-stagger x (or weight) before this multiply.

3. input = fp32, standard (module already runs in fp32)

step	op	arrangement
entry	—	fp32, standard
x.to(fp32)	(traced away — same-dtype .to() produces no graph node at all)	fp32, standard
pow/mean/add/rsqrt/mul	pointwise + reduction	stays standard
x.to(fp32)	(traced away — same-dtype .to() produces no graph node at all)	fp32, standard
weight * x	mul	both standard — safe

x's arrangement never changes at all here — it's standard on entry and stays
standard throughout, since dtype never changes so no cast op is ever emitted.
No rearrangement needed (trivially: there's nothing to fix, and no EA machinery
is even touched).

4. input = fp32, staggered (x arrives already staggered from an earlier, un-closed fp16→fp32 upcast upstream)

step	op	arrangement
entry	—	fp32, staggered
x.to(fp32)	(traced away — same-dtype .to(), no node exists to absorb the incoming stagger)	fp32, staggered
pow	pointwise	stays staggered (safe, unary)
mean(-1)	full-stick reduction	stays staggered (safe, order-irrelevant — reduced dim == stick dim)
add eps	binary, scalar broadcast	stays staggered (safe)
rsqrt	pointwise	stays staggered (safe, unary)
mul x*rsqrt	binary, broadcast across stick	stays staggered (safe)
x.to(fp32)	(traced away — same-dtype .to(), no node exists to restore it)	fp32, staggered
weight * x	mul	weight standard vs x staggered, no broadcast → illegal

Rearrangement needed before the final weight * x — same illegal shape as
scenario 2, but for a different reason: there is no dl16tofp32/fp32todl16 inside
this kernel at all (dtype never changes), so there's no cast to open or close a
bracket. Verified on an fp32-native RMSNorm compile: fx_graph_readable.py shows
zero convert_element_type/_to_copy nodes for either .to() call — they're
erased at Dynamo/AOTAutograd trace time, before torch-spyre's backend (and its
EA case-1..4 restoration logic in propagate_layouts.py) ever sees the graph. This
is the RFC's noted gap: an eager .to(fp32) that's a no-op on dtype leaves an
incoming stagger completely unmanaged, because there is no op left to manage it.

Summary

#	input dtype	input EA	x's arrangement changes inside the kernel?	extra rearrangement op needed?	where
1	fp16	standard	yes (standard→staggered→standard, via the 2 casts already in the code)	no	—
2	fp16	staggered	yes (staggered→standard→staggered, via the same 2 casts)	yes	final weight * x
3	fp32	standard	no (no cast ever runs)	no	—
4	fp32	staggered	no (no cast ever runs — dtype never changes)	yes	final weight * x

So "arrangement changes" and "needs an extra rearrangement op" are independent —
scenario 1 changes arrangement twice and needs nothing extra; scenario 4 never
changes arrangement at all and still needs a fix. Both failure cases (#2, #4) end
up in the same illegal shape (weight-standard vs x-staggered, no broadcast, at
the last op) but for different root causes: #2 is a fresh stagger opened by the
closing downcast with no consumer left to close it; #4 is an incoming stagger that
this kernel's .to() calls never touch because the dtype doesn't change.