[TMTensor] Bufferize ins operands at the op that reads them

[nmalimban]

tm-tensor-bufferize is a partial bufferizer: the TMTensor ops it rewrites
become pure memref ops, whose reads One-Shot Bufferize cannot see (its
analysis is defined over tensor operands). Each tensor->buffer crossing was
left where the conversion driver put it -- a bufferization.to_buffer at the
operand's definition. That is the only read One-Shot sees, and since it
happens before any later write, One-Shot finds no read-after-write conflict
and may fold two values onto one buffer.

convert-torch-to-tmtensor produces exactly this shape for
torch.aten.scatter_reduce.two with include_self=false: the reduction
identity (a linalg.fill) and the update values (a linalg.generic) share
one tensor.empty init, and are then read by two tm_tensor.scatters. With
the two writers collapsed, reduce='sum' returns 2x the expected value at
every scatter target, and the other reduce modes are wrong in other ways
(e.g. amax returns the identity, -inf).

This does not show up in the e2e suite, for two independent reasons:
refbackend runs One-Shot with copy-before-write, which skips the
analysis, and its linalg-fuse-elementwise-ops dissolves the shared init
before bufferization runs (any later cse brings it back).

Example

The shape convert-torch-to-tmtensor emits (reduced from
test/Dialect/TMTensor/bufferize-one-shot-hazard.mlir). One init, two
writers, two readers:

%init  = tensor.empty() : tensor<2xf32>
%zeros = linalg.fill ins(%zero : f32) outs(%init : tensor<2xf32>) -> tensor<2xf32>
%vals  = linalg.generic ... ins(%src : tensor<2xf32>) outs(%init : tensor<2xf32>) {...}
%cleared = tm_tensor.scatter ... ins(%zeros, %indices ...) outs(%self ...)    // plant 0.0
%result  = tm_tensor.scatter ... ins(%vals, %indices ...) outs(%cleared ...)  // add values

Without the fix, -tm-tensor-bufferize states the read of %zeros at
its definition, before the linalg.generic writes the same init:

%3 = linalg.fill ins(%cst : f32) outs(%2 : tensor<2xf32>) -> tensor<2xf32>
%4 = bufferization.to_buffer %3 : tensor<2xf32> to memref<2xf32>   // read of %zeros
%5 = linalg.generic ... ins(%arg1 : tensor<2xf32>) outs(%2 : tensor<2xf32>) {...}
%6 = bufferization.to_buffer %5 : tensor<2xf32> to memref<2xf32>
...
tm_tensor.scatter ... ins(%4, %1 : memref<2xf32>, memref<2x1xi32>) outs(%alloc ...)

%4 is the only read of %zeros that One-Shot can see, and it comes before
the generic, so One-Shot sees no conflict and bufferizes both writers onto
the same buffer. The generic overwrites the zeros, and both loops then read
the values:

%alloc_1 = memref.alloc() {alignment = 64 : i64} : memref<2xf32>
linalg.fill ins(%cst : f32) outs(%alloc_1 : memref<2xf32>)
linalg.generic ... ins(%0 ...) outs(%alloc_1 : memref<2xf32>) {...}   // overwrites the zeros
scf.for ... {   // clear loop
  %4 = memref.load %alloc_1[%arg3] : memref<2xf32>                    // reads values, not 0.0
  ...
scf.for ... {   // add loop
  %4 = memref.load %alloc_1[%arg3] : memref<2xf32>                    // same buffer

With the fix, the read is stated at the scatter that performs it, after
both writers:

%1 = linalg.fill ins(%cst : f32) outs(%0 : tensor<2xf32>) -> tensor<2xf32>
%2 = linalg.generic ... ins(%arg1 : tensor<2xf32>) outs(%0 : tensor<2xf32>) {...}
...
%4 = bufferization.to_buffer %1 read_only : tensor<2xf32> to memref<2xf32>
%5 = bufferization.to_buffer %arg0 read_only : tensor<2x1xi32> to memref<2x1xi32>
tm_tensor.scatter ... ins(%4, %5 : memref<2xf32>, memref<2x1xi32>) outs(%alloc ...)

One-Shot now sees a read of %zeros after a conflicting write, so it gives
the generic a buffer of its own:

%alloc   = memref.alloc() {alignment = 64 : i64} : memref<2xf32>
%alloc_0 = memref.alloc() {alignment = 64 : i64} : memref<2xf32>
linalg.fill ins(%cst : f32) outs(%alloc_0 : memref<2xf32>)
linalg.generic ... ins(%1 ...) outs(%alloc : memref<2xf32>) {...}
scf.for ... {   // clear loop
  %3 = memref.load %alloc_0[%arg3] : memref<2xf32>                    // 0.0
  ...
scf.for ... {   // add loop
  %3 = memref.load %alloc[%arg3] : memref<2xf32>                      // values

(Final IR is from -tm-tensor-bufferize -tm-tensor-to-loops | mlir-opt --one-shot-bufferize="allow-unknown-ops" -canonicalize.)

Change

Each crossing is now stated inside the pattern, at the TMTensor op that
reads the operand, and built from the original tensor operand rather than
from the adaptor's converted value:

• ins operands: bufferization.to_buffer %operand read_only. At the op,
  this is a read-after-write against any intervening write, so One-Shot
  finds the conflict. read_only is truthful: TMTensor ops never write
  ins operands.
• outs operands read by the payload: memref.alloc +
  materialize_in_destination from the tensor operand, replacing
  cloneMemref on the converted buffer.
• outs operands the payload does not read, with dynamic shapes: sizes
  are taken from the tensor operand (DPS guarantees they match).

materializeToTensor now marks its to_tensor restrict. Without it,
One-Shot Analysis rejects the pass's output outright, so nothing that
analyses (rather than copying defensively) could consume it. restrict
holds: every buffer reaching it is a fresh memref.alloc.

Cost

On main, One-Shot Analysis rejects this pass's output (the to_tensor
lacks restrict), so there is no analysed baseline to compare against.
The "without the fix" column below is the cost if One-Shot could run its
analysis on the old output -- i.e. the old crossings plus only the
restrict change. Counts are memref.copy after
tm-tensor-bufferize,tm-tensor-to-loops then
one-shot-bufferize{bufferize-function-boundaries function-boundary-type-conversion=identity-layout-map}:

op	without the fix	with the fix
scan inclusive / exclusive / dynamic	0 / 1 / 0	0 / 1 / 0
scatter update / add	1 / 1	1 / 1
sort	2	2
topk	2	2
attention	0	0
scatter_reduce hazard	2 (wrong result)	2 (correct)

The fix adds no copies. On the hazard IR, One-Shot resolves the conflict by
giving the clobbering linalg.generic a fresh outs buffer. That costs no
copy, because the generic's payload does not read its outs.

Tests

• bufferize-scatter-reduce-numeric.mlir: lowers the real torch op, runs
  plain --one-shot-bufferize, executes the result under mlir-runner,
  and checks the printed values for sum and amax, which fail in the two
  different ways described above. Guarded by REQUIRES: mlir-runner.
• bufferize-one-shot-hazard.mlir: checks at the IR level that the two
  scatter loops read distinct buffers.
• bufferize.mlir: CHECKs updated, plus @scan_1d_dynamic for the
  dynamic write-only outs path.

All three fail with Bufferize.cpp reverted. test/lit.cfg.py and
test/CMakeLists.txt gain the mlir-opt / mlir-runner tools and
dependencies these tests need.

🤖 Generated with Claude Code

[nmalimban]

Hi @sahas3 @hariprasadravi, I am requesting reviews via this comment since I do not have Triage permissions for this repo. Thanks!