[Bug] tt.spyre_tensor_layout accepts a stick size the dtype cannot have, and most lit fixtures use one

[fabianlim]

[Bug] tt.spyre_tensor_layout accepts a stick size the dtype cannot have, and most lit fixtures use one

Problem

A stick is a 128-byte aligned memory chunk, so the number of elements in one
stick is 128 / sizeof(element) — 32 for f32, 64 for f16. From
the Spyre architecture doc:

A stick is a 128-byte aligned memory chunk, which works out to 64 elements at
fp16. The size matches the natural granularity of data transfers between LPDDR5
device memory and the per-core LX scratchpad.

BYTES_IN_STICK = 128 there; STICK_BYTES = 128 in third_party/spyre/test/utils.py
here. Nothing enforces the relationship between that constant and an annotation's
stick size.

SpyreTensorLayoutOp::verify() checks that the three coordinate arrays have
equal length, that phys_src is in range, that phys_op is a known code, that
phys_arg is positive for floordiv/mod, and that a logical dim spans at most
one floordiv and one mod. It never relates phys_arg on the mod dim to the
descriptor's element type.

RewriteDescriptorLayout cannot catch it either: it reads the stick size out of
the annotation rather than deriving it.

// Classify.cpp
d.stickSize = coords.physBlock[d.lane];

So an f32 descriptor annotated with a stick of 64 elements — 256 bytes, twice the
hardware unit — is accepted, and the pass produces internally consistent IR from
an impossible layout. There is no diagnostic at any stage.

Where the wrong value is already in the tree

Twelve test/Conversion fixtures annotate f32 descriptors with phys_arg = 64:

Fixture	occurrences
matmul	21
invalid-ktir	12
absence	9
advanced, loop-rescale	6 each
elementwise-chain, reduce	5 each
invalid	4
chained-transpose, gather	2 each
idempotent, shared-idx	1 each

The f16 fixtures are right — batch-matmul, index-domain,
parallel-multistick all use 64 on f16, which is correct.

Two fixtures show the inconsistency directly rather than uniformly:

• sink-transpose uses phys_arg = 32 on f32 descriptors, which is correct.
• matmul uses both 32 and 64 on f32, in the same file
  (matmul_parallel_scatter_output at 32, the rest at 64).

So the codebase already disagrees with itself about the same dtype.

reduce additionally uses phys_arg = 96 on f32 (middle_axis and
leading_axis). 96 is not 128 / itemsize for any supported element type, so it
is not a stick size at all.

The pytest fixtures do it correctly and never hardcode — they derive from the
signature:

# fixtures/matmul/meta.py
"A_LAYOUT": [[(1, "floordiv", _SS("a_ptr")), 0, (1, "mod", _SS("a_ptr"))]],

with _SS = functools.partial(sticksize, _SIG_SPYRE) and

# test/utils.py
STICK_BYTES = 128          # 32 for fp32, 64 for fp16
def sticksize(signature, key):
    return STICK_BYTES // np.dtype(np_dtype(signature, key)).itemsize

Why it matters

The lit fixtures are the reference for what a correct annotation looks like. A
kernel author copying one gets a layout the hardware cannot express, and no stage
of the compiler says so — the failure surfaces later as wrong addresses or wrong
numbers, or not at all in a structural test.

It also means the physical shapes those fixtures pin are not the shapes a real
f32 kernel produces. [64, 128] stick-on-N is [2, 64, 64] at stick 64 but
[4, 64, 32] at stick 32, so the loop trip counts, slice extents and stick
factors the fixtures exercise are all off by a factor of two for f32.

Suggested fix

Two parts, and the second forces the first to land with it.

1. Validate in SpyreTensorLayoutOp::verify() that phys_arg on a mod dim
   equals 128 / itemsize for the descriptor's element type. The descriptor type
   is on the op, so the element type is available. This is the part that stops the
   next occurrence.

2. Correct the twelve fixtures. Changing 64 to 32 on f32 changes every
   physical shape those files pin, so the CHECK blocks regenerate — nine of them
   are positional %[[VAL_N]] chains where this is mechanical but bulky. Adding
   the check in (1) turns all twelve red, so they have to be fixed in the same
   change or the check lands disabled.

Worth deciding as part of this: whether reduce's 96 is a deliberate
non-hardware value used to exercise a code path, in which case the fixture needs a
comment saying so and the verifier needs a way to express that intent, or whether
it is simply wrong too.

Notes

third_party/spyre/docs/spyre-tensor-layouts.md states the derivation and warns
that the lit fixtures use 64 on f32 as a test convenience. That note documents the
discrepancy rather than resolving it; if (1) and (2) land, the warning should be
removed rather than kept.