[Frontend] bilinear interpolate emits invalid MLIR: gather value used across sibling loop regions (and its scratchpad is overwritten)

[YWHyuk]

Summary

On feature/togsim-cpp-trace, a bilinear F.interpolate emits invalid MLIR: a value defined
inside one gather's affine.for body is referenced from a later, sibling affine.for body.
Codegen aborts while the Python passes re-parse the kernel:

mlir._mlir_libs._site_initialize.<locals>.MLIRError: Unable to parse module assembly:
error: "-":186:42: use of undeclared SSA value name

Minimal repro

import torch
import torch.nn.functional as F

x = torch.randn(1, 8, 8, 8).to("npu:0")
f = torch.compile(dynamic=False)(
    lambda t: F.interpolate(t, scale_factor=2, mode="bilinear", align_corners=False)
)
f(x)   # -> MLIRError: use of undeclared SSA value name

What does and does not trigger it

case	result
x[i] (single gather)	OK
x[i0] - x[i1] (two gathers + sub)	OK
x[i0] + x[i1]	OK
F.interpolate(..., mode="nearest")	OK
F.interpolate(..., mode="bilinear")	FAIL

So two indirect gathers alone are fine. It needs the bilinear pattern, where several gathers of the
same tensor are blended and one gather's loaded value is consumed after a later gather has run.

Mechanism

Two things are broken at once. From the generated kernel:

119:  "togsim.transfer"(%in_ptr0, ..., %spad4, ..., %spad3) {indirect = true}   // gather #1 -> spad4
120:  affine.for %compute_idx = 0 to 64 step 16 {
122:      %tmp57  = affine.vector_load %spad4[...]                              // value A
        ...
      }                                                                        // <- scope closes

182:  "togsim.transfer"(%in_ptr0, ..., %spad4, ..., %spad7) {indirect = true}   // gather #2 -> SAME spad4
183:  affine.for %compute_idx = 0 to 64 step 16 {
185:      %tmp113 = affine.vector_load %spad4[...]                              // value B
186:      %tmp114 = arith.subf %tmp57, %tmp113                                  // A - B
      }

1. SSA scope violation. %tmp57 is defined in the first loop's region and used in a sibling
   region, so the module does not parse.
2. Even with scoping fixed, the data would be wrong. Both indirect MVINs target the same
   %spad4 (only the index buffer differs: %spad3 vs %spad7), so value A has already been
   overwritten by gather Scalar #2 by the time arith.subf runs.

A correct lowering has to either give each gather its own scratchpad buffer (and hoist the loaded
value so it stays live), or fuse the gather loops so both loads happen in one region.

Impact

Blocks HuggingFace SegFormer (issue #254). Running SegFormer on this branch, 10 of 1170 generated
kernels fail to parse; every one of them has >= 2 indirect MVINs writing the same scratchpad. The
decode head's bilinear upsample is the source.

Note this is not the error originally reported in #254 (NotImplementedError: Not supporting format),
which no longer reproduces. SegFormer now clears codegen for 145 kernels and all their gem5 runs
(~197 s with the functional Spike pass off), then dies here.

Related but distinct: #284 (the trace path does not model indirect access). This issue is about the
frontend emitting unparseable MLIR in the first place.

Environment

• branch feature/togsim-cpp-trace
• fails at PyTorchSimFrontend/mlir/passes/__init__.py:78 (Module.parse) via
  extension_codecache.load -> run_python_passes, i.e. before gem5/Spike, so it is pure frontend
  codegen and independent of the simulator binaries.

[YWHyuk]

Root cause

Two coupled defects. Both stem from gathers of the same tensor aliasing one scratchpad.

D1 — the scratchpad cache key ignores which index buffer the gather used

MLIRKernel.load() (mlir_codegen_backend.py) first factors the indirect part of the index out
into offset_desc:

index, offset_desc = self.convert_indirect_indexing(index)
...
sram_var, sram_index_var = self.get_scratchpad_buffer(dtype, name, local_tile_desc, index)

and allocate_sram_buffer caches the buffer by (dram_name, str(raw_index)):

if str(raw_index) not in self.global_vars_dict[dram_name]:
    ...
self.global_vars_dict[dram_name][str(raw_index)] = new_name

After convert_indirect_indexing, the residual raw_index for x[i0] and x[i1] is the same
string. offset_desc — the thing that actually distinguishes the two gathers — is not part of the
key. So both gathers get the same %spadN, each togsim.transfer {indirect = true} overwrites the
previous gather's data, and the emitted load text (affine.vector_load %spadN[...]) is identical.

D2 — values are referenced across step (region) boundaries without bridging

push_step() opens a new affine.for compute region:

def push_step(self):
    # New load->compute->store step; later emits land here, steps bridge via spad
    ...
    self.cse = self.cse.clone()  # share name counter, fresh dedup cache (region-safe)

The dedup cache is made region-safe, but a CSEVariable produced in an earlier step is still
printed verbatim when a later step consumes it. The backend already knows where that value came from
— load() records self.spad_buffer_dict[str(out)] = [sram_var, ...] — so it could re-emit the load
in the consuming region. It does not. The "steps bridge via spad" intent is not implemented for this
path.

How the two defects produce two very different symptoms

(a) No step boundary between the gathers -> silent wrong answer.
Both loads land in one region, identical load text, CSE collapses them:

"togsim.transfer"(..., %spad1, ..., %alloc1) {indirect = true}   // gather #1 -> spad1
"togsim.transfer"(..., %spad1, ..., %alloc3) {indirect = true}   // gather #2 -> SAME spad1
affine.for %compute_idx ... {
    %tmp3 = affine.vector_load %spad1[...]
    %tmp7 = arith.subf %tmp3, %tmp3                              // A - A  ==  0

Verified numerically — x[i0] - x[i1] returns all zeros on npu:

golden : [0.756, 1.246, 0.784, -1.844, ...]
npu    : [0.0, 0.0, 0.0, 0.0, ...]

Filed separately as a correctness bug, since a user hits it with no interpolate involved.

(b) A step boundary between the gathers -> invalid MLIR (this issue).
Bilinear interpolate needs its own MVIN per gather, so push_step() runs between them. The later
region then references %tmp57 from the closed region -> use of undeclared SSA value name.

Why fixing only one defect is not enough

• D1 alone (give each gather its own spad): the bilinear case still references %tmp57 from a
  closed region -> still a parse error.
• D2 alone (re-load the value in the consuming region): the reload reads %spad4, which now
  holds gather Scalar #2's data -> still wrong (B - B == 0).

Fix sketch

1. Make the indirect index descriptor part of the scratchpad cache key, so each live gather owns its
   buffer. The load texts then differ (%spad4 vs %spad9) and the CSE distinguishes them naturally.
2. Bridge cross-step values: when a later step consumes a CSEVariable produced in an earlier step,
   re-emit the load from spad_buffer_dict[var] inside the consuming region.

Worth checking SPAD pressure afterwards: bilinear keeps up to 4 gathers live.