[Bug] `tt.broadcast` is not re-derived against the physical rank, blocking layouts on any reduce-then-broadcast kernel

[fabianlim]

Problem

RewriteDescriptorLayout physicalizes a loaded tensor to its stick rank, but a
tt.broadcast that must be combined with it keeps its logical target shape.
The consuming elementwise op then fails its own operand-type constraint:

error: 'arith.subf' op requires the same type for all operands and results

This is the reason the softmax fixture carries no tt.spyre_tensor_layout
variant (fixtures/softmax/meta.py:68-72). It blocks the whole
reduce → broadcast → elementwise shape, which is the core of softmax,
layernorm, and every other normalization kernel.

Reproduction

softmax_single_tile (fixtures/softmax/kernel.py:27) with a stick-on-N
annotation on both descriptors, BLOCK_SIZE=128, f16 (stick = 64), so
[1, 128] -> [128/64, 1, 64] = [2, 1, 64]:

tl.spyre_tensor_layout(in_desc,  [(1, "floordiv", 64), 0, (1, "mod", 64)])
tl.spyre_tensor_layout(out_desc, [(1, "floordiv", 64), 0, (1, "mod", 64)])
row      = in_desc.load([row_idx, 0])
row_max  = tl.max(row, axis=1)
num      = tl.exp(row - row_max)      # <-- fails here

Fails inside RewriteDescriptorLayout in the default pipeline.

Root cause

The relevant TTIR, as it reaches the pass (after lower-compute-ops):

%52 = "tt.broadcast"(%51) : (tensor<1x1xf32>) -> tensor<1x128xf32>
%53 = "arith.subf"(%50, %52) : (tensor<1x128xf32>, tensor<1x128xf32>)
                             -> tensor<1x128xf32>

%50 is the physicalized row; %52 is the broadcast row-max.

tt.broadcast encodes its target shape in its result type, not as a value.
Physicalizing %50 to tensor<2x1x64xf32> therefore requires recomputing that
target — and retypeChain (RewriteDescriptorLayout.cpp:596-635) has no
machinery for it:

• It walks forward over getUsers() and retypes each single-result op's result
  from operand 0 (:624-631). Reaching arith.subf via %50, it retypes
  the result to rank 3, but leaves operand %52 alone.
• It enqueues only op->getResult(0).getUsers() (:632-633). %52 is a sibling
  operand of arith.subf, not a successor, so the walk never visits it — and its
  producer %51 traces back to the reduce, not to the physicalized load, so no
  forward walk from the load can ever reach it.
• grep -n broadcast over RewriteDescriptorLayout.cpp returns zero hits.
  The pass has no concept of broadcast semantics.

Result: arith.subf ends up with operand 0 at rank 3 and operand 1 at rank 2, and
its verifier rejects it.

This is the same structural bug class as the inter-tile identity issue:
retypeChain propagates strictly forward along operand 0, so any op needing two
shape-coupled operands reconciled has one of them left stale. Softmax and
inter-tile are two instances; they need the same underlying capability, and the
fix should probably be designed once.

The broadcast target alone is enough to break it

It is tempting to assume the hard part is the reduce: when the layout splits the
reduced axis, that axis is spread across two physical dims (d0 = stick index,
d2 = stick lane) of [2, 1, 64], so a correct reduce must fold over {d0, d2}
— the multi-dim-reduce capability the normalize-to-generic discussion issue argues
for.

But that is not the blocker. Splitting the non-reduced axis fails
identically. Stick-on-M with a [128, 64] block ([M,N] -> [M//64, M%64, N],
so the reduced axis N stays intact as a single dim), reduce over axis 1 with
keep_dims=True:

error: 'arith.subf' op requires the same type for all operands and results

Same diagnostic, same site. Since the reduced axis is untouched here, the reduce's
dimensions need no rewriting at all — the only thing wrong is the broadcast's
stale target shape.

That narrows the issue usefully:

• Minimum fix is re-deriving the broadcast target against the physical rank.
  This alone unblocks stick-on-non-reduced-axis, which is a real, testable subset.
• Multi-dim reduce is an additional requirement, needed only when the layout
  splits the reduced axis. It is not a prerequisite for the first subset.

So a stick-on-M softmax variant is the natural first milestone, and it does not
depend on the normalize-to-generic work.

Scope

Blocks layout annotations on any kernel where a reduction result is broadcast back
against a physicalized tensor:

• all three softmax variants (single_tile, multi_tile, 2pass) — each does
  row - row_max and num / denom
• layernorm / RMSnorm-shaped kernels generally
• any masked-max or normalize step

reduce and matmul are unaffected: their reduce results are stored, never
broadcast back against a physical-rank operand — which is why those fixtures do
have working spyre_stick variants.

Failure quality

Loud, and caught by an MLIR op verifier rather than producing wrong numbers. The
diagnostic names arith.subf and not spyre_tensor_layout, so the connection to
the layout annotation is not obvious from the message alone — but it does fail.

Proposed direction

Not proposing a specific patch — the reduce/broadcast coupling above means this is
a design question, not a local edit. Two shapes to consider:

1. Re-derive the broadcast target (unblocks the non-reduced-axis subset). When
   retypeChain retypes an elementwise op's result to physical rank, reconcile its
   other operands: for a tt.broadcast, recompute the result type against the
   physical shape instead of leaving the logical one. Per the section above this is
   sufficient on its own for stick-on-non-reduced-axis, and is the smallest change
   that yields a landable softmax variant.
2. Add multi-dim reduce (unblocks splitting the reduced axis). Only needed on
   top of (1) when the layout splits the axis being reduced. Shares the
   dimensions = [d0, d2] work described in the normalize-to-generic issue.
3. Or do it all in generic form. If reductions/contractions are normalized to
   linalg.generic before this pass, a broadcast becomes an indexing_map that
   drops the stick dims rather than a materialized shape — no target to re-derive,
   and multi-dim reduce falls out of iterator_types for free. Best long-term
   shape, and it would fix the inter-tile identity case in the same stroke.

Interim behavior is acceptable (loud verifier failure), so this is a capability gap
rather than a correctness emergency — unlike the ragged-stick and dispatch-boundary
issues.

Test coverage gap

• No lit case covers a layout-annotated reduce-then-broadcast kernel.
• fixtures/softmax/meta.py has no spyre_stick variant; the blocker is
  documented in the module comment (:68-72) but was not tracked in an issue
  until now.
• With a fix: a positive lit case pinning the physicalized broadcast + multi-dim
  reduce, and a numerical spyre_stick softmax variant.

Related

• The inter-tile identity issue — same retypeChain forward-only root cause,
  different stale operand. Consider fixing the class, not the instances.
• The normalize-to-generic discussion issue — option 2 above; its multi-dim
  reduce work is a prerequisite for the stick-on-N case.
• The tracking issue — coverage gap list.