Physicalize linalg.broadcast in RewriteDescriptorLayout

[kcirred]

Summary

RewriteDescriptorLayout could not physicalize a linalg.broadcast, and
that single gap stopped every softmax- and layernorm-shaped kernel from
carrying a tt.spyre_tensor_layout annotation. Both failed with a
same-type error on the arith.subf that consumes the broadcast --
pointing at the user's subtract rather than at anything about layouts.
Both now compile. Addresses #91.

Four propagation rules were placeholders returning failure(), so
neither half of Phase 2A could cross a broadcast or the size-1 reshape
that LowerComputeOps leaves in front of it. They are implemented here:

• BroadcastRequirement::induce projects a requirement onto the axes the
  broadcast carries, dropping the ones it adds. This is what lets a
  requirement reach the reduce behind a broadcast, which a reduce needs
  before it will take the physical output space.
• ReshapeRequirement::induce and ReshapePropagation::propagate cross a
  reshape that only adds or drops size-1 axes, and keep declining when a
  reassociation regroups a real axis. The narrow case is exactly what
  tt.expand_dims followed by tt.broadcast produces.
• BroadcastPropagation::propagate recomputes the target shape against
  the physical rank, and RewriteBroadcastPattern renumbers dimensions
  and rebuilds the outs init to match.

Three match predicates were declining whenever an op's tensor operands
disagreed on a shape. Mid-pass they routinely do: Phase 1 physicalizes
loads and stops, so an op with one load-fed operand and one not is
guaranteed to see a mismatch -- the state the pass exists to resolve, not
evidence the op is unknown. The test was also vacuous on a single-operand
op, so a reshape or a broadcast passed it as though it preserved shape.
All three now exclude shape-changing ops by kind, through one shared
isShapeChangingOp.

Layernorm needed one more piece. Its divisor is the reduction length,
which arrives as a splat constant with no operands, so the operand-driven
forward walk cannot reach it and no propagation rule can claim it.
seedSplatOperands rebuilds such a constant at the physical shape from
the consumer side. Only a splat qualifies: a general dense<[...]>
constant has per-element data whose stick-tiled placement is a real
question, and guessing it would satisfy every verifier while computing
wrong numbers.

Known limitations, each pinned by a test that fails when it is fixed:

• A layout that splits an axis the broadcast carries declines rather
  than physicalizing. The input would need retyping too, which a rule
  returning only a result type cannot express; repairing it at the
  consuming op is separate work.
• A tensor.reshape in the chain still blocks the layout. It carries its
  target shape as a runtime operand, so there is no reassociation to
  project through, even though the result type here is static.
• Only a compile-time splat is seeded. A splat of a runtime scalar stays a
  linalg.fill, which seedSplatOperands does not match.

Scope

• Spyre backend or KTIR lowering
• Tests or fixtures
• Build, packaging, or CI
• Documentation or repository metadata
• Upstream Triton sync or compatibility

Testing

• I ran the relevant tests, or explained why they were not run.
• I added or updated tests when behavior changed.

Commands and results:

Checklist

• I initialized or updated submodules when needed.
• I checked public-facing text for private links or internal-only references.
• If this touches upstream Triton files outside third_party/spyre, I explained why in the summary.

[kcirred]

#129

[fabianlim]

these tests are confusing tro me because I dont actually see the splat in the below mlir

[fabianlim]

this kind of logic is repeated many times, and in general these functions are quite long, so I wonder if we should make some shared helpers.

[fabianlim]

The primary structural finding

2A resolves a result from one physical operand. 2B requires all operands
to agree. Those are different notions of "resolved", and most of the
problems below live in the gap between them.

Concretely, on the layernorm divide:

%x   = ktdp.load %tile : <64x2x64xindex> -> tensor<64x2x64xf32>   // physical (Phase 1)
%cst = arith.constant dense<1.28e2> : tensor<64x128xf32>          // logical, zero operands
%y   = arith.divf %x, %cst : tensor<64x128xf32>

What 2A does. getPhysicalizedType calls findPhysicalTensorOperand, which
returns the first tensor operand already in roots/resolved — here %x —
and stops looking. ElementwisePropagation::propagate then takes the result
shape from that one operand:

info.type = RankedTensorType::get(srcTy.getShape(), resTy.getElementType());

No sibling is consulted, and after this PR ElementwisePropagation::match no
longer compares operand shapes either. So 2A records: %y is physical,
tensor<64x2x64xf32>.

What 2B does. RewriteElementwisePattern walks every tensor operand and
returns failure() the moment one disagrees. %cst is still 64x128, so 2B
declines the very op 2A just promised.

Why that is silent. verifyPhysicalTypeAgreement is one-directional:
2B's physicalValues ⊆ 2A's PhysicalTypeMap. It catches 2A under-claiming,
which the design doc explains at length ("2B reads absence from the map as
genuinely logical and commits on that basis"). 2A over-claiming is precisely
the direction it does not check. So the gap has no alarm on it.

Which reframes seedSplatOperands. It is not a missing propagation rule and
not a phase-placement slip. It exists to retroactively make 2A's over-claim
true — mutating IR until 2B's all-operands condition holds. That is why it must
run before the agreement check, and why its failure mode is mutate-then-decline
(finding 2): a repair aimed at a condition the pattern has not yet evaluated
cannot know whether the repair was warranted.

Two further consequences of the same gap:

• The PR removes the operand-shape-agreement check from
  ElementwisePropagation (2A) and keeps it in RewriteElementwisePattern
  (2B). Both changes are individually justified in the diff. Together they widen
  the gap and write it into the code as intended behavior.
• The codebase already has the right concept and does not apply it here:
  reconcileOperandSet — "the trip count of a stick loop is a property of the
  operand set, not of one operand, and that agreement is one question, answered
  once per op." Source ops get that. Elementwise ops do not.

The fix is to make one of the two notions authoritative. Since 2B's is the one
that must hold for the IR to verify, 2A should adopt it: resolve the operand
set, reconcile it, and only then record a result type — step 3 below.

Two catch-alls, and the permissive one shadows the strict one

getPhysicalizedType already has a strict catch-all, and its intent is
exactly right:

const PhysicalPropagationPattern *pattern = lookupPattern(defOp, patterns);
if (!pattern) {
  // No pattern is this op's rule. An untaught op stays visible here rather
  // than acquiring a default: if an op can appear on a physicalized chain it
  // needs a pattern, and guessing one would silently compute the wrong slice.
  LLVM_DEBUG(...);
  return failure();
}

Two things make it dead code:

1. ElementwisePropagation is registered last and matches almost anything,
   so ops never reach that branch. populatePhysicalPropagationPatterns'
   comment documents the ordering hazard in the other direction ("every
   named-op pattern must be asked before the structural elementwise rule") — the
   fragility was seen, but treated as an ordering constraint to maintain rather
   than as evidence that a residual rule should not exist.
2. "Stays visible" means LLVM_DEBUG — not a diagnostic, not a counter, not
   an error. In 2B it is worse: RewriteElementwisePattern is
   MatchAnyOpTypeTag() at benefit 1 with no error branch at all, and the
   design doc concedes "if every pattern declines, the worklist drains and the
   driver calls that a fixpoint."

So the escape hatch that should catch an untaught op exists and is unreachable.
The catch-all should be the thing that errors. Elementwise should be a named
rule like any other.

Consequence: a rule for arith.constant is reachable, and no seeding is needed

getPhysicalizedType already recurses on operands:

Value src = findPhysicalTensorOperand(defOp, roots, resolved);
if (!src) {
  // No operand is known physical yet; try to resolve one transitively.
  for (Value o : defOp->getOperands()) { ... getPhysicalizedType(o, ...) ... }
}

So getPhysicalizedType(%cst) can be called and would reach
lookupPattern(arith.constant). Registering a rule is enough — no new
machinery, and in particular no seeding, forward or backward.

But note the gate: the recursion runs only when no operand is already known
physical. In layernorm the load-fed operand resolves immediately, src is found,
and the loop never executes, so %cst is never queried. The constant is
reachable only in the case where nothing else is, which is backwards.

Ungate that, and the splat's shape comes from the operand set — it never needed
to come from the store's requirement.

Proposed order of work

1. Demote elementwise from catch-all. Trait-based allowlist —
   OpTrait::Elementwise / SameOperandsAndResultShape, plus
   linalg::isElementwise for linalg.generic. isShapeChangingOp's
   hand-maintained denylist deletes; finding 5 is fixed; arith.constant stops
   being incidentally excluded by sawTensorOperand and becomes visibly
   unhandled. Same inversion in 2B, and drop MatchAnyOpTypeTag.
2. Make the strict catch-all speak. Promote the null-pattern branch from
   LLVM_DEBUG to a real diagnostic naming the op; give 2B the branch it does
   not have.
3. 2A resolves the operand set, not one operand. Ungate the transitive
   recursion and add the reconciliation reconcileOperandSet already models.
   This is what makes 2A's and 2B's "resolved" the same predicate.
4. Then the splat is a one-line rule. SplatConstantPropagation: "I am
   layout-invariant; I take whatever the operand set agreed on." Recorded in 2A,
   so 2B retypes off the record. seedSplatOperands deletes, finding 2
   evaporates, verifyPhysicalTypeAgreement keeps full strength.

Land 1–2 as a measurement first. Count unhandled ops across the 149 registry
variants; do not fail yet. That is the pattern #132 used — the backward analysis
landed as a measurement that agreed with the forward walk, and was only made
authoritative afterwards. It de-risks the real hazard: today's silent
fallthrough-to-logical may be load-bearing for kernels that currently pass.
Promote to a hard error once the count is zero.

---

The reshape rules are working around IR that should not exist

ReshapeRequirement / ReshapePropagation, ~150 lines across two analyses,
exist to see through an expand_shape 1→1x1 immediately followed by
collapse_shape 1x1→1. That pair is a no-op. It exists only because
LowerComputeOps lowers tt.expand_dims and tt.broadcast independently
(rules A3 and A4) — the PR's own comments say exactly this.

Evidence that the rule is fitted rather than derived: the two directions test
different conditions under the same name.

• Backward: "no reassociation group of size > 1 contains a floordiv/mod dim" — a
  real statement about reassociation.
• Forward: srcTy == operandTy, i.e. "this value is not actually split, so
  nothing can go wrong." Not a reshape rule at all; a "nothing to do here"
  guard that holds only because a reduce result in softmax is rank-1 and unsplit.

Consequence: the forward rule cannot cross the stick-split reduce result that
#131 just added support for (the [2, 64] Physical output-axis-space case).
The two PRs' capabilities do not compose.

The fix is upstream, not here — see the next section.

Should a cleanup prepass run before this pass?

Yes for targeted reshape folding. No for general CSE — and the
distinction is load-bearing, not fastidiousness.

The ordering today. In _make_ktir, canonicalize and cse are added
after the entire core pipeline (_CORE_PIPELINE_PASSES, then
distribute_work). So LowerComputeOps → RewriteDescriptorLayout runs with no
cleanup between them, and the no-op expand/collapse pair is alive when the pass
sees it.

Targeted folding — do it. MLIR already canonicalizes collapse(expand(x))
for compatible reassociations, so a fold anchored after lower_compute_ops would
very likely delete the pair for free, retiring most of the reshape rules above.
The pipeline even has the insertion mechanism — options.required_fixes anchors
a pass after a named core pass — though it currently names KTDP passes, so a
common-pass or a small dedicated KTDP cleanup would need wiring.

General CSE — do not. It would change decisions this pass makes, in ways the
design doc explicitly relies on not happening:

Wider point. The no-op reshape pair and the shaped splat constant are two
instances of one thing: an earlier pass committed to a shape it did not need to,
and this pass pays to un-commit it. Folding handles one; steps 1–4 handle the
other. A third instance would mean the fix belongs in LowerComputeOps.