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

lib/Dialect/TMTensor/Transforms/Bufferize.cpp

@@ -32,18 +32,8 @@ namespace mlir::torch::TMTensor {
 #define GEN_PASS_DEF_TMTENSORBUFFERIZE
 #include "torch-mlir-dialects/Dialect/TMTensor/Transforms/Passes.h.inc"
 
-static Value cloneMemref(Location loc, Value memref, OpBuilder &b) {
-  auto memrefType = cast<MemRefType>(memref.getType());
-  auto alloc =
-      memref::AllocOp::create(b, loc, memref::getMixedSizes(b, loc, memref),
-                              memrefType.getElementType());
-  memref::CopyOp::create(b, loc, memref, alloc);
-  return alloc;
-}
-
 static LogicalResult
 allocateBuffersForResults(Location loc, TMTensorOp tmtensorOp,
-                          ValueRange outputs,
                           SmallVectorImpl<Value> &resultBuffers, OpBuilder &b) {
   // Lazily compute loopRanges.
   SmallVector<Range, 4> loopRanges;
@@ -62,12 +52,23 @@ allocateBuffersForResults(Location loc, TMTensorOp tmtensorOp,
     }
     auto tensorShape = tensorType.getShape();
     auto memrefType = MemRefType::get(tensorShape, tensorType.getElementType());
-    Value resultTensor = outputs[resultIndex];
 
     // Clone output buffers whose value is actually used.
     OpOperand *tiedOpOperand = tmtensorOp.getOutputOperand(resultIndex);
     if (tmtensorOp.payloadUsesValueFromOperand(tiedOpOperand)) {
-      resultBuffers.push_back(cloneMemref(loc, resultTensor, b));
+      // The op updates this operand in place, so the buffer it writes has to
+      // start out holding the operand's incoming value. Copy from the *tensor*
+      // operand, not from the buffer the conversion driver handed us: a copy
+      // with a tensor source is itself the read of that tensor, so the read
+      // lands here, at the op, instead of at the operand's definition -- the
+      // same reason the `ins` crossings below are built here.
+      Value alloc = memref::AllocOp::create(
+          b, loc, tensor::getMixedSizes(b, loc, tiedOpOperand->get()),
+          memrefType.getElementType());
+      bufferization::MaterializeInDestinationOp::create(
+          b, loc, /*result=*/TypeRange{}, tiedOpOperand->get(), alloc,
+          /*restrict=*/true, /*writable=*/true);
+      resultBuffers.push_back(alloc);
       continue;
     }
 
@@ -77,8 +78,13 @@ allocateBuffersForResults(Location loc, TMTensorOp tmtensorOp,
       continue;
     }
 
+    // The payload does not read this output, so the buffer needs the operand's
+    // shape but not its contents. Take the dynamic sizes from the *tensor*
+    // operand; the destination-passing-style contract guarantees they match the
+    // result's -- "Init operands and their tied OpResults have the same type.
+    // Dynamic dimension sizes also match at runtime."
     resultBuffers.push_back(memref::AllocOp::create(
-        b, loc, memref::getMixedSizes(b, loc, resultTensor),
+        b, loc, tensor::getMixedSizes(b, loc, tiedOpOperand->get()),
         memrefType.getElementType()));
   }
   return success();
@@ -108,16 +114,37 @@ class BufferizeAnyTMTensorOp : public OpInterfaceConversionPattern<TMTensorOp> {
     Location loc = op.getLoc();
     SmallVector<Value, 2> newOutputBuffers;
 
-    SmallVector<Value> outputs(operands.begin() + op.getNumInputs(),
-                               operands.end());
-    if (failed(allocateBuffersForResults(loc, op, outputs, newOutputBuffers,
-                                         rewriter))) {
+    if (failed(
+            allocateBuffersForResults(loc, op, newOutputBuffers, rewriter))) {
       return op.emitOpError()
              << "Failed to allocate buffers for tensor results.";
     }
 
-    SmallVector<Value> inputs(operands.begin(),
-                              operands.begin() + op.getNumInputs());
+    // Take each tensor `ins` operand's buffer here, at the op that reads it,
+    // instead of using the operand the conversion driver converted for us. The
+    // driver pins a target materialization to the operand's *definition*, which
+    // leaves the read invisible to a later full bufferization: this pass is
+    // partial, so the TMTensor op it feeds reads raw memory, and One-Shot
+    // Bufferize analyses tensor operands only. Any write between the definition
+    // and this op may then be bufferized in place over the very buffer we were
+    // handed, silently changing the value the op observes. Stating the crossing
+    // here instead puts the read at the program point where it happens, so
+    // One-Shot sees the read-after-write and preserves the value. `read_only`
+    // is truthful -- a TMTensor op never writes an `ins` operand -- and keeps
+    // the crossing from counting as a write as well.
+    SmallVector<Value> inputs;
+    for (OpOperand &opOperand :
+         op->getOpOperands().take_front(op.getNumInputs())) {
+      auto tensorType = dyn_cast<RankedTensorType>(opOperand.get().getType());
+      if (!tensorType) {
+        inputs.push_back(operands[opOperand.getOperandNumber()]);
+        continue;
+      }
+      inputs.push_back(bufferization::ToBufferOp::create(
+          rewriter, loc,
+          MemRefType::get(tensorType.getShape(), tensorType.getElementType()),
+          opOperand.get(), /*read_only=*/true));
+    }
     createTMTensorOpOnBuffers(rewriter, op, inputs, newOutputBuffers);
     // Replace the results of the old op with the new output buffers.
     rewriter.replaceOp(op, newOutputBuffers);
@@ -131,7 +158,8 @@ static Value materializeToTensor(OpBuilder &builder, TensorType type,
                                  ValueRange inputs, Location loc) {
   assert(inputs.size() == 1);
   assert(isa<BaseMemRefType>(inputs[0].getType()));
-  return bufferization::ToTensorOp::create(builder, loc, type, inputs[0]);
+  return bufferization::ToTensorOp::create(builder, loc, type, inputs[0],
+                                           /*restrict=*/true);
 }
 
 /// Converts TMTensor operations that work on tensor-type operands or results to
@@ -188,8 +216,9 @@ struct TMTensorBufferizePass
     });
 
     // Mark all Standard operations legal.
-    target.addLegalDialect<arith::ArithDialect, func::FuncDialect,
-                           memref::MemRefDialect, tensor::TensorDialect>();
+    target.addLegalDialect<
+        arith::ArithDialect, bufferization::BufferizationDialect,
+        func::FuncDialect, memref::MemRefDialect, tensor::TensorDialect>();
 
     // Mark all TMTensor operations illegal as long as they work on tensors.
     auto isLegalOperation = [&](Operation *op) {

test/CMakeLists.txt

@@ -16,6 +16,13 @@ set(TORCH_MLIR_TEST_DEPENDS
         TorchMLIRPythonModules
         torch-mlir-opt
         torch-mlir-capi-torch-test
+        # Needed by tests whose RUN lines reach for upstream tooling:
+        # --one-shot-bufferize is not registered in torch-mlir-opt, and the
+        # numeric regression tests execute the lowered IR.
+        mlir-opt
+        mlir-runner
+        mlir_runner_utils
+        mlir_c_runner_utils
         )
 
 add_lit_testsuite(check-torch-mlir "Running the torch-mlir regression tests"

test/Dialect/TMTensor/bufferize-one-shot-hazard.mlir

@@ -0,0 +1,89 @@
+// RUN: torch-mlir-opt %s -tm-tensor-bufferize -tm-tensor-to-loops \
+// RUN:   | mlir-opt --one-shot-bufferize="allow-unknown-ops" \
+// RUN:   | FileCheck %s
+
+// The IR shape `convert-torch-to-tmtensor` emits for
+// `torch.aten.scatter_reduce.two` with `include_self=false`: one `tensor.empty`
+// init with two writers -- the reduction identity and the update values -- read
+// by two separate `tm_tensor.scatter`s, a "clear" pass and an "add" pass.
+//
+// The property pinned here: the two loops must read from two *different*
+// buffers, and the buffer the clear loop reads must be the one filled with 0.0.
+// See bufferize-scatter-reduce-numeric.mlir for the same bug at the level of
+// values.
+
+// CHECK-LABEL:   func.func @scatter_reduce_sum_no_include_self(
+// CHECK-SAME:                            %[[INDICES:.*]]: tensor<2x1xi32>,
+// CHECK-SAME:                            %[[SRC:.*]]: tensor<2xf32>,
+// CHECK-SAME:                            %[[SELF:.*]]: tensor<4xf32>
+// CHECK-DAG:       %[[SELF_BUF:.*]] = bufferization.to_buffer %[[SELF]]
+// CHECK-DAG:       %[[SRC_BUF:.*]] = bufferization.to_buffer %[[SRC]]
+// CHECK-DAG:       %[[ZERO:.*]] = arith.constant 0.000000e+00 : f32
+
+// The two writers of the shared init get separate buffers: neither can be
+// written in place, because the scatter that reads the other one conflicts.
+// CHECK:           %[[VALS:.*]] = memref.alloc(){{.*}} : memref<2xf32>
+// CHECK:           %[[ZEROS:.*]] = memref.alloc(){{.*}} : memref<2xf32>
+// CHECK:           linalg.fill ins(%[[ZERO]] : f32) outs(%[[ZEROS]] : memref<2xf32>)
+// CHECK:           linalg.generic {{.*}} ins(%[[SRC_BUF]]{{.*}}) outs(%[[VALS]] : memref<2xf32>)
+
+// The clear loop reads the zeros and overwrites its targets.
+// CHECK:           %[[CLEARED:.*]] = memref.alloc() : memref<4xf32>
+// CHECK:           memref.copy %[[SELF_BUF]], %[[CLEARED]]
+// CHECK:           %[[IDX_BUF:.*]] = bufferization.to_buffer %[[INDICES]] read_only
+// CHECK:           scf.for
+// CHECK:             %[[U0:.*]] = memref.load %[[ZEROS]]
+// CHECK:             memref.load %[[IDX_BUF]]
+// CHECK:             memref.store %[[U0]], %[[CLEARED]]
+// CHECK:           }
+
+// The add loop reads the values -- a *different* buffer.
+// CHECK:           %[[RESULT:.*]] = memref.alloc() : memref<4xf32>
+// CHECK:           %[[RESULT_TENSOR:.*]] = bufferization.to_tensor %[[RESULT]] restrict
+// CHECK:           memref.copy %[[CLEARED]], %[[RESULT]]
+// CHECK:           %[[IDX_BUF_2:.*]] = bufferization.to_buffer %[[INDICES]] read_only
+// CHECK:           scf.for
+// CHECK:             %[[U1:.*]] = memref.load %[[VALS]]
+// CHECK:             memref.load %[[IDX_BUF_2]]
+// CHECK:             %[[O:.*]] = memref.load %[[RESULT]]
+// CHECK:             %[[SUM:.*]] = arith.addf %[[U1]], %[[O]] : f32
+// CHECK:             memref.store %[[SUM]], %[[RESULT]]
+// CHECK:           }
+// CHECK:           return %[[RESULT_TENSOR]] : tensor<4xf32>
+func.func @scatter_reduce_sum_no_include_self(
+    %indices: tensor<2x1xi32>, %src: tensor<2xf32>,
+    %self: tensor<4xf32>) -> tensor<4xf32> {
+  %zero = arith.constant 0.000000e+00 : f32
+  %init = tensor.empty() : tensor<2xf32>
+
+  // Writer 1: the reduction identity.
+  %zeros = linalg.fill ins(%zero : f32) outs(%init : tensor<2xf32>) -> tensor<2xf32>
+
+  // Writer 2: the update values, sharing the SAME init tensor.
+  %vals = linalg.generic {
+      indexing_maps = [affine_map<(d0) -> (d0)>, affine_map<(d0) -> (d0)>],
+      iterator_types = ["parallel"]}
+      ins(%src : tensor<2xf32>) outs(%init : tensor<2xf32>) {
+  ^bb0(%in: f32, %out: f32):
+    linalg.yield %in : f32
+  } -> tensor<2xf32>
+
+  // Clear pass: plant the identity at every scatter target.
+  %cleared = tm_tensor.scatter {dimension_map = array<i64: 0>} unique_indices(true)
+      ins(%zeros, %indices : tensor<2xf32>, tensor<2x1xi32>)
+      outs(%self : tensor<4xf32>) {
+  ^bb0(%update: f32, %orig: f32):
+    tm_tensor.yield %update : f32
+  } -> tensor<4xf32>
+
+  // Add pass: accumulate the values on top.
+  %result = tm_tensor.scatter {dimension_map = array<i64: 0>} unique_indices(false)
+      ins(%vals, %indices : tensor<2xf32>, tensor<2x1xi32>)
+      outs(%cleared : tensor<4xf32>) {
+  ^bb0(%update: f32, %orig: f32):
+    %sum = arith.addf %update, %orig : f32
+    tm_tensor.yield %sum : f32
+  } -> tensor<4xf32>
+
+  return %result : tensor<4xf32>
+}

test/Dialect/TMTensor/bufferize-scatter-reduce-numeric.mlir

@@ -0,0 +1,145 @@
+// REQUIRES: mlir-runner
+// RUN: torch-mlir-opt %s \
+// RUN:   --pass-pipeline='builtin.module(func.func(convert-torch-to-tmtensor,convert-torch-to-linalg,canonicalize,cse,torch-finalizing-backend-type-conversion,tm-tensor-bufferize,tm-tensor-to-loops))' \
+// RUN:   | mlir-opt --one-shot-bufferize="bufferize-function-boundaries function-boundary-type-conversion=identity-layout-map" \
+// RUN:       --convert-linalg-to-loops --expand-strided-metadata --lower-affine \
+// RUN:       --convert-scf-to-cf --finalize-memref-to-llvm --convert-func-to-llvm \
+// RUN:       --convert-arith-to-llvm --convert-cf-to-llvm --reconcile-unrealized-casts \
+// RUN:   | mlir-runner -e main --entry-point-result=void \
+// RUN:       --shared-libs=%mlir_lib_dir/libmlir_runner_utils%shlibext \
+// RUN:       --shared-libs=%mlir_lib_dir/libmlir_c_runner_utils%shlibext \
+// RUN:   | FileCheck %s
+
+// Numeric regression test for `torch.aten.scatter_reduce.two` with
+// `include_self=false`, starting from the actual torch op and *executing* the
+// result. This is the only level at which the bug this file guards against is
+// observable as a value rather than as a shape of IR.
+//
+// The pipeline deliberately finishes bufferization with plain
+// `--one-shot-bufferize`, so that One-Shot's analysis actually runs. Adding
+// `copy-before-write` -- as `refbackend`'s pipeline does -- would skip the
+// analysis and copy defensively, which hides the bug.
+//
+// An e2e test cannot stand in for this file. `ScatterReduceFloatSumModule`
+// already *is* this op with `include_self=false`, and it passes on a build
+// without the fix, for two independent reasons. First, `refbackend`'s pipeline
+// passes `copy-before-write`, which skips One-Shot's analysis altogether.
+// Second, it would pass even with the analysis enabled, because the pipeline
+// opens with `linalg-fuse-elementwise-ops`, whose `RemoveOutsDependency`
+// pattern rewrites every `outs` operand the payload does not read to a fresh
+// `tensor.empty` -- dissolving the shared init described below before
+// bufferization is ever consulted.
+//
+// Neither defence is one this pass can rely on. `copy-before-write` is opt-in,
+// and the `RemoveOutsDependency` rewrite is undone by any `cse` between fusion
+// and bufferization: the fresh `tensor.empty` is structurally identical to the
+// one the `linalg.fill` below already writes into, so CSE merges the two and
+// the shared init comes straight back. What keeps the e2e suite green is
+// therefore an accident of pass ordering, not anything the suite can be
+// configured to expose -- which is why this file pins the behaviour at the
+// level of values instead.
+//
+// The hazard `tm-tensor-bufferize` is handed: `convert-torch-to-tmtensor`
+// computes the reduction identity (a `linalg.fill`) and the update values (a
+// `linalg.generic`) into the *same* `tensor.empty` init, then runs two
+// `tm_tensor.scatter`s -- a "clear" pass that plants the identity at each target
+// and a "reduce" pass that combines the updates in. One init, several writers.
+//
+// `cse` in the pipeline above is load-bearing, not tidying: the lowering emits
+// two structurally identical `tensor.empty` + `linalg.fill` pairs, and it is CSE
+// that merges them into the single init the two scatters then share. Drop it and
+// this test passes even on a build without the fix.
+//
+// If the scatters' reads are stated at their operands' definitions rather than
+// at the scatters, One-Shot Bufferize sees no conflict between those writers and
+// folds them onto one buffer, so one of the two values is lost. See
+// bufferize-one-shot-hazard.mlir for the IR-level property.
+//
+// Both modes below use the same inputs:
+//
+//   self  = [10 20; 30 40; 50 60; 70 80]
+//   index = [[0 1]; [0 3]]      src = [[1 2]; [3 4]]
+//   scatter_reduce(dim=0, include_self=false)
+//
+// Cell (0,0) receives both 1 and 3; (1,1) receives 2; (3,1) receives 4; every
+// other cell keeps its `self` value.
+//
+// Two reduce modes are checked because the collapse destroys a different value
+// in each, and a fix that recovered only one shape would still be wrong:
+//
+//   sum  -- the identity is 0.0, which is also the constant the `linalg.generic`
+//           init is filled with, so CSE merges the two fills into one that
+//           *precedes* the generic. The surviving value is the updates, which
+//           the clear pass then scatters instead of zeros and the add pass adds
+//           a second time: 2x at every singly-targeted cell. (At (0,0) it is
+//           last-write plus full sum, 3 + (1+3) = 7, not 8.)
+//
+//   amax -- the identity is -inf, a different constant, so its fill is a third
+//           writer of the shared init emitted *after* the generic. That fill is
+//           what survives, destroying the updates, and both passes read -inf.
+//
+//   mode   correct                        collapsed
+//   sum    [4 20; 30 2; 50 60; 70 4]      [7 20; 30 4; 50 60; 70 8]
+//   amax   [3 20; 30 2; 50 60; 70 4]      [-inf 20; 30 -inf; 50 60; 70 -inf]
+//
+// CHECK:      Unranked Memref
+// CHECK-NEXT: [4, 20]
+// CHECK-NEXT: [30, 2]
+// CHECK-NEXT: [50, 60]
+// CHECK-NEXT: [70, 4]
+// CHECK:      Unranked Memref
+// CHECK-NEXT: [3, 20]
+// CHECK-NEXT: [30, 2]
+// CHECK-NEXT: [50, 60]
+// CHECK-NEXT: [70, 4]
+
+func.func private @printMemrefF32(tensor<*xf32>)
+
+// The `torch_c` casts wrap the torch op inside a builtin-tensor signature so
+// that @main can call it directly and the whole test stays in one file. They
+// are erased by `torch-finalizing-backend-type-conversion`.
+func.func @scatter_reduce_sum(%self_b: tensor<4x2xf32>, %index_b: tensor<2x2xi64>,
+                              %src_b: tensor<2x2xf32>) -> tensor<4x2xf32> {
+  %self = torch_c.from_builtin_tensor %self_b : tensor<4x2xf32> -> !torch.vtensor<[4,2],f32>
+  %index = torch_c.from_builtin_tensor %index_b : tensor<2x2xi64> -> !torch.vtensor<[2,2],si64>
+  %src = torch_c.from_builtin_tensor %src_b : tensor<2x2xf32> -> !torch.vtensor<[2,2],f32>
+  %dim = torch.constant.int 0
+  %reduce = torch.constant.str "sum"
+  %include_self = torch.constant.bool false
+  %0 = torch.aten.scatter_reduce.two %self, %dim, %index, %src, %reduce, %include_self :
+      !torch.vtensor<[4,2],f32>, !torch.int, !torch.vtensor<[2,2],si64>,
+      !torch.vtensor<[2,2],f32>, !torch.str, !torch.bool -> !torch.vtensor<[4,2],f32>
+  %r = torch_c.to_builtin_tensor %0 : !torch.vtensor<[4,2],f32> -> tensor<4x2xf32>
+  return %r : tensor<4x2xf32>
+}
+
+func.func @scatter_reduce_amax(%self_b: tensor<4x2xf32>, %index_b: tensor<2x2xi64>,
+                               %src_b: tensor<2x2xf32>) -> tensor<4x2xf32> {
+  %self = torch_c.from_builtin_tensor %self_b : tensor<4x2xf32> -> !torch.vtensor<[4,2],f32>
+  %index = torch_c.from_builtin_tensor %index_b : tensor<2x2xi64> -> !torch.vtensor<[2,2],si64>
+  %src = torch_c.from_builtin_tensor %src_b : tensor<2x2xf32> -> !torch.vtensor<[2,2],f32>
+  %dim = torch.constant.int 0
+  %reduce = torch.constant.str "amax"
+  %include_self = torch.constant.bool false
+  %0 = torch.aten.scatter_reduce.two %self, %dim, %index, %src, %reduce, %include_self :
+      !torch.vtensor<[4,2],f32>, !torch.int, !torch.vtensor<[2,2],si64>,
+      !torch.vtensor<[2,2],f32>, !torch.str, !torch.bool -> !torch.vtensor<[4,2],f32>
+  %r = torch_c.to_builtin_tensor %0 : !torch.vtensor<[4,2],f32> -> tensor<4x2xf32>
+  return %r : tensor<4x2xf32>
+}
+
+func.func @main() {
+  %self = arith.constant dense<[[10.0, 20.0], [30.0, 40.0], [50.0, 60.0], [70.0, 80.0]]> : tensor<4x2xf32>
+  %index = arith.constant dense<[[0, 1], [0, 3]]> : tensor<2x2xi64>
+  %src = arith.constant dense<[[1.0, 2.0], [3.0, 4.0]]> : tensor<2x2xf32>
+  %sum = call @scatter_reduce_sum(%self, %index, %src)
+      : (tensor<4x2xf32>, tensor<2x2xi64>, tensor<2x2xf32>) -> tensor<4x2xf32>
+  %sum_unranked = tensor.cast %sum : tensor<4x2xf32> to tensor<*xf32>
+  call @printMemrefF32(%sum_unranked) : (tensor<*xf32>) -> ()
+
+  %amax = call @scatter_reduce_amax(%self, %index, %src)
+      : (tensor<4x2xf32>, tensor<2x2xi64>, tensor<2x2xf32>) -> tensor<4x2xf32>
+  %amax_unranked = tensor.cast %amax : tensor<4x2xf32> to tensor<*xf32>
+  call @printMemrefF32(%amax_unranked) : (tensor<*xf32>) -> ()
+  return
+}