[mlir][tensor][bufferize] Bufferize inserts into equivalent tensors in-place

[springerm]

Summary

Inserting a tensor into an equivalent tensor is a no-op after bufferization. No alloc is needed.

mlir/include/mlir/Dialect/Bufferization/IR/BufferizableOpInterface.h

public:
   ///   |          |          |          |
   ///   +----------+----------+----------+
   ///              |
   ///              1
   ///
   /// In the above example, Values with a star satisfy the condition. When
   /// starting the traversal from Value 1, the resulting SetVector is:
   /// { 2, 7, 8, 5 }
-  SetVector<Value> findValueInReverseUseDefChain(
-      Value value, llvm::function_ref<bool(Value)> condition) const;
+  ///
+  /// If `followEquivalentOnly` is set, only equivalent OpOperands are selected.
+  SetVector<Value>
+  findValueInReverseUseDefChain(Value value,
+                                llvm::function_ref<bool(Value)> condition,
+                                bool followEquivalentOnly = false) const;
 
   /// Find the Values of the last preceding write of a given Value.
   ///
   /// Note: Unknown ops are handled conservatively and assumed to be writes.
   /// Furthermore, BlockArguments are also assumed to be writes. There is no
   /// analysis across block boundaries.
   ///
   /// Note: When reaching an end of the reverse SSA use-def chain, that value

mlir/lib/Dialect/Bufferization/IR/BufferizableOpInterface.cpp

bool AnalysisState::isValueRead(Value value) const {
   return false;
 }
 
 // Starting from `value`, follow the use-def chain in reverse, always selecting
 // the aliasing OpOperands. Find and return Values for which `condition`
 // evaluates to true. OpOperands of such matching Values are not traversed any
 // further.
 llvm::SetVector<Value> AnalysisState::findValueInReverseUseDefChain(
-    Value value, llvm::function_ref<bool(Value)> condition) const {
+    Value value, llvm::function_ref<bool(Value)> condition,
+    bool followEquivalentOnly) const {
   llvm::SetVector<Value> result, workingSet;
   workingSet.insert(value);
 
   while (!workingSet.empty()) {
     Value value = workingSet.pop_back_val();
     if (condition(value) || value.isa<BlockArgument>()) {
       result.insert(value);
       continue;
     }
 
     OpResult opResult = value.cast<OpResult>();
+    BufferizableOpInterface bufferizableOp =
+        options.dynCastBufferizableOp(opResult.getDefiningOp());
     SmallVector<OpOperand *> opOperands = getAliasingOpOperand(opResult);
-    if (opOperands.empty() || !options.isOpAllowed(value.getDefiningOp())) {
+
+    // Stop iterating in either one of these cases:
+    // * The current op is not bufferizable or excluded in the filter.
+    // * There are no OpOperands to follow.
+    // * There is an OpOperand, but it is not an equivalent tensor (only if
+    //   `followEquivalentOnly` is set).
+    if (!bufferizableOp || opOperands.empty() ||
+        (followEquivalentOnly &&
+         bufferizableOp.bufferRelation(opResult, *this) !=
+             BufferRelation::Equivalent)) {
       result.insert(value);
       continue;
     }
 
     for (OpOperand *o : opOperands)
       workingSet.insert(o->get());
   }
 

mlir/lib/Dialect/Tensor/Transforms/BufferizableOpInterfaceImpl.cpp

BufferRelation bufferRelation(Operation *op, OpResult opResult,
                                 const AnalysisState &state) const {
     return BufferRelation::Equivalent;
   }
 };
 
 /// Return true if the (ExtractSliceOp, InsertSliceOp) pair match (i.e.
 /// equivalent operand / result and same offset/sizes/strides specification).
 template <typename OpTy>
-static bool areEquivalentExtractSliceOps(const AnalysisState &state,
-                                         ExtractSliceOp extractSliceOp,
-                                         OpTy insertSliceOp) {
+static bool areEquivalentSlices(const AnalysisState &state,
   if (!extractSliceOp || !insertSliceOp)
     return false;
   if (extractSliceOp != insertSliceOp &&
       !state.areEquivalentBufferizedValues(extractSliceOp.getSource(),
                                            insertSliceOp.getDest()))
     return false;
   if (!sameOffsetsSizesAndStrides(extractSliceOp, insertSliceOp,
                                   isEqualConstantIntOrValue))
     return false;
   return true;
 }
 
-/// Return true if `value` is originating from an ExtractSliceOp that matches
-/// the given InsertSliceOp.
+/// Return true if `value` is originating from the InsertSliceOp's destination
+/// or an ExtractSliceOp that matches the given InsertSliceOp.
 template <typename OpTy>
-static bool hasMatchingExtractSliceOp(const AnalysisState &state, Value value,
-                                      OpTy insertSliceOp) {
-  auto condition = [&](Value val) {
+static bool matchesInsertDestination(const AnalysisState &state, Value value,
+  // Look for matching slices.
+  auto matchesSlice = [&](Value val) {
     if (auto extractSliceOp = val.getDefiningOp<ExtractSliceOp>())
-      if (areEquivalentExtractSliceOps(state, extractSliceOp, insertSliceOp))
+      if (areEquivalentSlices(state, extractSliceOp, insertSliceOp))
         return true;
     return false;
   };
+  if (llvm::all_of(state.findValueInReverseUseDefChain(value, matchesSlice),
+                   matchesSlice))
+    return true;
 
-  return llvm::all_of(state.findValueInReverseUseDefChain(value, condition),
-                      condition);
+  // Look for equivalent values.
+  auto isEquivalent = [&](Value val) {
+    return state.areEquivalentBufferizedValues(val, insertSliceOp.getDest());
+  };
+  if (llvm::all_of(state.findValueInReverseUseDefChain(
+                       value, isEquivalent, /*followEquivalentOnly=*/true),
+                   isEquivalent))
+    return true;
+  return false;
 }
 
 template <typename OpTy>
 static bool isNotConflictingInsertSliceLikeOp(Operation *op, OpOperand *uRead,
                                               OpOperand *uConflictingWrite,
                                               const AnalysisState &state) {
   Operation *readingOp = uRead->getOwner();
   Operation *conflictingWritingOp = uConflictingWrite->getOwner();
 
   // Special rules for matching ExtractSliceOp/InsertSliceOp pairs. If
   // uRead is an InsertSliceOp...
   if (auto insertSliceOp = dyn_cast<OpTy>(readingOp)) {
     // As an example, consider the following IR.
     //
     // %0 = tensor.extract_slice %t[%a, %b][%c, %d][1, 1] {inplace = [true] }
     // %1 = linalg.fill %cst, %0 {inplace= [true] }
     // %2 = tensor.insert_slice %1 into %t[%a, %b][%c, %d][1, 1]
     //     {inplace= [true] }
 
     // TODO: Use insertSliceOp.getDestOpOperand etc. when available.
     if (uRead == &insertSliceOp->getOpOperand(1) /*dest*/ &&
-        hasMatchingExtractSliceOp(state, uConflictingWrite->get(),
-                                  insertSliceOp))
+        matchesInsertDestination(state, uConflictingWrite->get(),
       // Case 1: The main insight is that InsertSliceOp reads only part of
       // the destination tensor. The overwritten area is not read. If
       // uConflictingWrite writes into exactly the memory location that is
       // being read by uRead, this is not a conflict.
       //
       // In the above example:
       // uRead             = OpOperand 1 (%t) of tensor.insert_slice
       // uConflictingWrite = OpOperand 1 (%0) of linalg.fill
       //
       // The read of %t does not conflict with the write of the FillOp
       // (same aliases!) because the area that the FillOp operates on is
       // exactly the one that is *not* read via %t.
       return true;
 
     if (uRead == &insertSliceOp->getOpOperand(0) /*source*/ &&
         uConflictingWrite == &insertSliceOp->getOpOperand(1) /*dest*/ &&
-        hasMatchingExtractSliceOp(state, uRead->get(), insertSliceOp))
+        matchesInsertDestination(state, uRead->get(), insertSliceOp))
       // Case 2: The read of the source tensor and the write to the dest
       // tensor via an InsertSliceOp is not a conflict if the read is
       // reading exactly that part of an equivalent tensor that the
       // InsertSliceOp is writing.
       //
       // In the above example:
       // uRead             = OpOperand 0 (%1) of tensor.insert_slice
       // uConflictingWrite = OpOperand 1 (%t) of tensor.insert_slice
if (auto insertSliceOp = dyn_cast<OpTy>(conflictingWritingOp))
     // lastWrite         = %1
     //
     // This is not a conflict because the InsertSliceOp overwrites the
     // memory segment of %1 with the exact same data. (Effectively, there
     // is no memory write here.)
     if (uConflictingWrite == &insertSliceOp->getOpOperand(1) /*dest*/ &&
         state.areEquivalentBufferizedValues(uRead->get(),
                                             insertSliceOp.getSource()) &&
-        hasMatchingExtractSliceOp(state, insertSliceOp.getSource(),
-                                  insertSliceOp))
+        matchesInsertDestination(state, insertSliceOp.getSource(),
       return true;
 
   return false;
 }
 
 /// Bufferization of tensor.insert_slice. Replace with a memory copy. Under
 /// certain circumstances, this op can also be a no-op.
 struct InsertSliceOpInterface

mlir/test/Dialect/Tensor/one-shot-bufferize.mlir

func.func @insert_slice_fun_not_inplace(
   %r0 = tensor.insert_slice %t into %A[0][4][1] : tensor<4xf32> into tensor<?xf32>
 
   //     CHECK: return %{{.*}} : memref<?xf32>
   return %r0: tensor<?xf32>
 }
 
 // -----
 
-// CHECK-LABEL: func @tensor_cast_not_in_place(
-//  CHECK-SAME:     %[[A:.*]]: memref<?xf32{{.*}}>, %[[B:.*]]: memref<?xf32{{.*}}>
-//       CHECK:   %[[alloc:.*]] = memref.alloc
-//       CHECK:   memref.copy %[[A]], %[[alloc]]
+// CHECK-LABEL: func @tensor_cast_in_place(
+//  CHECK-SAME:     %[[A:.*]]: memref<?xf32{{.*}}>
 //       CHECK:   %[[subview:.*]] = memref.subview %[[A]][{{.*}}] [4] [1] : {{.*}} to memref<4xf32
-//       CHECK:   memref.copy %[[alloc]], %[[subview]]
-func.func @tensor_cast_not_in_place(
-    %A : tensor<?xf32> {bufferization.writable = true},
-    %B : tensor<?xf32> {bufferization.writable = false}, %idx: index)
+//       CHECK:   memref.copy %[[A]], %[[subview]]
+func.func @tensor_cast_in_place(
+    %A : tensor<?xf32> {bufferization.writable = true}, %idx: index)
   -> (tensor<?xf32>)
 {
   %r0 = tensor.cast %A : tensor<?xf32> to tensor<4xf32>
   %r1 = tensor.insert_slice %r0 into %A[%idx][4][1] : tensor<4xf32> into tensor<?xf32>
   return %r1 : tensor<?xf32>
 }
 
 // -----
func.func @dealloc_pad_buffer(%t1: tensor<?x10xindex>, %l2: index, %h1: index,
   %0 = tensor.pad %t1 low[5, %l2] high[%h1, %h2] {
   ^bb0(%arg0: index, %arg1: index):
     %m = arith.muli %arg0, %arg1 : index
     tensor.yield %m : index
   } : tensor<?x10xindex> to tensor<?x?xindex>
   %r = tensor.extract %0[%idx, %idx] : tensor<?x?xindex>
   return %r : index
 }
+
+// -----
+
+// CHECK-LABEL: func @insert_equivalent_tensor
+func.func @insert_equivalent_tensor(%t: tensor<10xf32>) -> tensor<10xf32> {
+  // CHECK-NOT: memref.alloc
+  %cst = arith.constant 4.200000e+01 : f32
+  // CHECK: linalg.fill
+  %0 = linalg.fill ins(%cst : f32) outs(%t : tensor<10xf32>) -> tensor<10xf32>
+  // CHECK-NOT: memref.copy
+  %1 = tensor.insert_slice %0 into %t[0][10][1] : tensor<10xf32> into tensor<10xf32>
+  return %1 : tensor<10xf32>
+}