Tail CSE in `_make_ktir` merges an author-written store/load pair into one access tile, aborting the scheduler

[fabianlim]

The tail CSE in _make_ktir merges an author-written store/load pair into one access tile

What happens

_make_ktir ends with add_canonicalizer + add_cse. On a kernel where the
author writes the HBM round-trip themselves — an intermediate declared as a
pointer argument with its own tl.make_tensor_descriptor, its own
tl.spyre_tensor_layout, an explicit store and an explicit load — that CSE
merges the store's and the load's ktdp.construct_access_tile into one
value, and the two tensor.empty ops into one. dbo-opt's
ComputeGroupExtraction then aborts.

The smallest case is out = sqrt(exp(x)) with exp(x) round-tripped through a
scratch pointer: two computes, one round-trip.

Before CSE (canonicalized _make_ktir output, abridged)

Two access tiles against the scratch view %2 — %16 for the store, %21 for
the load — and two tensor.empty, one per compute:

%4  = arith.muli %0, %c128 : index
%5  = arith.divsi %4, %c32 : index
%6  = arith.muli %0, %c128 : index
%7  = arith.remsi %6, %c32 : index
%8  = ktdp.construct_access_tile %1[%5, %7] {...} : memref<4x32xf32> -> !ktdp.access_tile<4x32xindex>
%9  = ktdp.load %8 : <4x32xindex> -> tensor<4x32xf32>
%10 = tensor.empty() : tensor<4x32xf32>
%11 = linalg.generic ... ins(%9 ...) outs(%10 ...) { math.exp } -> tensor<4x32xf32>
%12 = arith.muli %0, %c128 : index
%13 = arith.divsi %12, %c32 : index
%14 = arith.muli %0, %c128 : index
%15 = arith.remsi %14, %c32 : index
%16 = ktdp.construct_access_tile %2[%13, %15] {...} : memref<4x32xf32> -> !ktdp.access_tile<4x32xindex>
ktdp.store %11, %16 : tensor<4x32xf32>, <4x32xindex>          // store side
%17 = arith.muli %0, %c128 : index
%18 = arith.divsi %17, %c32 : index
%19 = arith.muli %0, %c128 : index
%20 = arith.remsi %19, %c32 : index
%21 = ktdp.construct_access_tile %2[%18, %20] {...} : memref<4x32xf32> -> !ktdp.access_tile<4x32xindex>
%22 = ktdp.load %21 : <4x32xindex> -> tensor<4x32xf32>          // load side
%23 = tensor.empty() : tensor<4x32xf32>
%24 = linalg.generic ... ins(%22 ...) outs(%23 ...) { math.sqrt } -> tensor<4x32xf32>
%29 = ktdp.construct_access_tile %3[%26, %28] {...}
ktdp.store %24, %29 : tensor<4x32xf32>, <4x32xindex>

After CSE

One access tile %11 on both sides of the fence — the store writes it and the
ktdp.load reads it — and one tensor.empty %9 as the outs of both
computes:

%4  = arith.muli %0, %c128 : index
%5  = arith.divsi %4, %c32 : index
%6  = arith.remsi %4, %c32 : index
%7  = ktdp.construct_access_tile %1[%5, %6] {...}
%8  = ktdp.load %7 : <4x32xindex> -> tensor<4x32xf32>
%9  = tensor.empty() : tensor<4x32xf32>
%10 = linalg.generic ... ins(%8 ...) outs(%9 ...) { math.exp } -> tensor<4x32xf32>
%11 = ktdp.construct_access_tile %2[%5, %6] {...}
ktdp.store %10, %11 : tensor<4x32xf32>, <4x32xindex>            // store side
%12 = ktdp.load %11 : <4x32xindex> -> tensor<4x32xf32>          // load side, SAME tile
%13 = linalg.generic ... ins(%12 ...) outs(%9 ...) { math.sqrt } -> tensor<4x32xf32>
%14 = ktdp.construct_access_tile %3[%5, %6] {...}
ktdp.store %13, %14 : tensor<4x32xf32>, <4x32xindex>

The merge order matters

The two construct_access_tile ops are not structurally equal in the input:
%16 takes [%13, %15] and %21 takes [%18, %20], four distinct SSA values.
What makes them equal is that CSE merges the index arithmetic first — the
duplicated muli / divsi / remsi triples collapse onto %4 / %5 / %6
within the same walk — and only then are both tiles %2[%5, %6] and mergeable.
This happens with --cse alone; the preceding canonicalizer is not needed.

The two assertions

Both are in
dataflow-scheduler/lib/Conversion/frontend/KTIRToScheduleIR/ComputeGroupExtraction.cpp,
and they are distinct failures from distinct merges. Measured by feeding dbo-opt
three versions of the same module:

Both merges (i.e. plain --cse) — the shared access tile transitively
collapses every access tile into one equivalence class, so the walk reaches a
store with no load-established class leader:

dbo-opt: .../ComputeGroupExtraction.cpp:570: void ComputeGroupExtractionPass::processBlockForExtraction(mlir::Block*, llvm::EquivalenceClasses<mlir::Value>&): Assertion `current_class_leader && "StoreOp found before any LoadOp"' failed.

Only the two tensor.empty merged, access tiles left alone — a separate
break, so fixing the tile alone would not be enough:

dbo-opt: .../ComputeGroupExtraction.cpp:472: void ComputeGroupExtractionPass::extractComputeGroup(mlir::ModuleOp, mlir::ModuleOp, const ComputeGroup&): Assertion `op->use_empty() && "Operation should have no uses left"' failed.

Neither merged — dbo-opt exits 0 and emits SpyreCode.

Why this cannot be fixed by exempting the ops

Both merged op kinds are Pure, so CSE is within its rights:

• ktdp.construct_access_tile is declared
  Ktdp_Op<"construct_access_tile", [AttrSizedOperandSegments, Pure]>
  (ktir-mlir-frontend, include/ktir/Dialect/KTDP/KTDP.td).
• tensor.empty is Pure upstream.

Upstream's CSE offers no trait, attribute or interface hook to exempt an
operation — it tests isMemoryEffectFree and OperationEquivalence, with no
opt-out — and --cse takes no options. So there is no way to keep the pass and
tell it that these ops are owned per compute group.

The structural point

This is the reason the issue exists, rather than the individual assertions.

HbmRoundtrip (#158) does the same thing an author does — it puts an
intermediate through HBM — and escapes this problem only by position: it runs
after the CSE and re-materializes its own per-group access tiles and
tensor.empty ops, so nothing merges them afterwards. Position is not available
to an author. Their ops exist before CSE by construction, because they came in
from the Triton source.

So per-group ownership of an access tile and of a compute's outs is a property
the scheduler requires and the pipeline does not maintain — it only happens to
hold for the passes that run late enough. Any pass, or any author, producing
these ops earlier hits this.

Measured cost of removing the CSE

On fabianlim/declared-buffers-doc plus six author-written declared-buffer
fixtures (elementwise__1d_device_dag_buffers and the five pooled variants), on
hardware:

	result
tail CSE present	6 failed, 191 passed, 2 xfailed (182 s)
tail CSE removed	197 passed, 2 xfailed (190 s)

No test that passes with the CSE fails without it. The only outcome difference
is the six new device launches, which flip from the StoreOp found before any LoadOp abort to passing and matching the NumPy oracle.

The cost is IR size — the redundant index arithmetic the CSE used to fold now
reaches dbo-opt, which accepts it:

kernel	ops with CSE	ops without
elementwise__1d_device[fp32, add]	21	30
softmax[M=1024]	47	58
elementwise__2d[M=512]	30	34

Wall-clock on the suite moved 182 s → 190 s, which is within run-to-run noise
for a device suite.

Note on a second CSE

_make_spyrecode also runs add_canonicalizer + add_cse, inside the
materialize_base_addresses branch. That branch is skipped in symbolic-argument
mode, which is the default (BUNDLE_SYMBOLIC_ARGS defaults to "1"), so it does
not fire on the path the tests take — but an address-bound compile of the same
kernel would hit the identical problem. Any resolution has to cover both sites.

Options, not a recommendation

Removing the tail CSE is what the accompanying commit does, and it is not
obviously wrong: nothing in the suite depends on it, and the canonicalizer
already handles the folding that motivated it (muli x, 1, cast chains). But it
is also not obviously right — it gives up a cheap cleanup for every kernel to
accommodate one op's ownership rule.

Recorded so the alternatives are visible rather than lost:

1. Leave the CSE out (current state). Cheapest; costs IR size everywhere.
2. Make the KTDP ops that must be owned per compute group not Pure, or give
   them a side effect on their memory view. Correct in principle; changes the
   dialect contract, and affects every consumer of KTIR.
3. Re-materialize per-group access tiles and tensor.empty in the spyrecode
   stage, whatever their provenance — i.e. do for author-written round-trips what
   HbmRoundtrip does for inferred ones, so ownership stops depending on which
   pass produced the op.
4. Have the scheduler tolerate a shared access tile / shared outs rather than
   assert.

The purpose of this issue is to record the problem and its options, not to argue
for one.