[mlir][bufferize] Avoid tensor copies when the data is not read

[springerm]

Summary

There are various shortcuts in BufferizationState::getBuffer that avoid a buffer copy when we just need an allocation (and no initialization). This change adds those shortcuts to the TensorCopyInsertion pass, so that getBuffer can be simplified in a subsequent change.

Depends On D126819

[springerm]

update

[aartbik]

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

354

nit: rephrase, old values not needed or something like this

[springerm]

address comments

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

public:
   /// Return true if `opOperand` bufferizes to a memory write. Return true` if
   /// the op is not bufferizable.
   bool bufferizesToMemoryWrite(OpOperand &opOperand) const;
 
   /// Return true if `opOperand` does neither read nor write but bufferizes to
   /// an alias. Return false if the op is not bufferizable.
   bool bufferizesToAliasOnly(OpOperand &opOperand) const;
 
+  /// Return true if a copy can always be avoided when allocating a new tensor
+  /// for the given OpOperand.
+  bool canOmitTensorCopy(OpOperand &opOperand) const;
+
   /// Return true if the given value is read by an op that bufferizes to a
   /// memory read. Also takes into account ops that create an alias but do not
   /// read by themselves (e.g., ExtractSliceOp).
   bool isValueRead(Value value) const;
 
   /// 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

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

 //===- BufferizableOpInterface.cpp - Bufferizable Ops  ---=----------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Dialect/Bufferization/IR/BufferizableOpInterface.h"
 #include "mlir/Dialect/Bufferization/IR/Bufferization.h"
 #include "mlir/Dialect/Func/IR/FuncOps.h"
 #include "mlir/Dialect/MemRef/IR/MemRef.h"
+#include "mlir/Dialect/Tensor/IR/Tensor.h"
 #include "mlir/IR/AsmState.h"
 #include "mlir/IR/BlockAndValueMapping.h"
 #include "mlir/IR/BuiltinOps.h"
 #include "mlir/IR/Operation.h"
 #include "mlir/IR/TypeUtilities.h"
 #include "mlir/IR/Value.h"
 #include "llvm/Support/Debug.h"
 

 using namespace mlir;
 using namespace bufferization;
 
 /// Attribute name used to mark region arguments that can be bufferized
 /// in-place during linalg comprehensive bufferization.
 constexpr const ::llvm::StringLiteral
     bufferization::BufferizableOpInterface::kInplaceableAttrName;
 
+/// Create an AllocTensorOp for the given shaped value. Only ranked tensors are
+/// supported at the moment. If `copy` is set, the shaped value is copied.
+/// Otherwise, a tensor with undefined contents is allocated.
+static Value allocateTensorForShapedValue(OpBuilder &b, Location loc,
+                                          Value shapedValue, bool escape,
+                                          bool copy = true) {
+  auto tensorType = shapedValue.getType().dyn_cast<RankedTensorType>();
+  assert(tensorType && "only RankedTensorType supported at the moment");
+  Value alloc;
+  if (!copy) {
+    // No copy needed: Just allocate.
+    SmallVector<Value> dynamicSizes;
+    for (int64_t i = 0; i < tensorType.getRank(); ++i)
+      if (tensorType.isDynamicDim(i))
+        dynamicSizes.push_back(b.create<tensor::DimOp>(loc, shapedValue, i));
+    alloc = b.create<AllocTensorOp>(loc, tensorType, dynamicSizes,
+                                    /*copy=*/Value(), escape);
+  } else {
+    // Allocate and copy.
+    alloc = b.create<AllocTensorOp>(loc, tensorType,
+                                    /*dynamicSizes=*/ValueRange(), shapedValue,
+                                    escape);
+  }
+  return alloc;
+}
+
 LogicalResult BufferizableOpInterface::resolveTensorOpOperandConflicts(
     RewriterBase &rewriter, const AnalysisState &state) {
   OpBuilder::InsertionGuard g(rewriter);
   Operation *op = getOperation();
   SmallVector<OpOperand *> outOfPlaceOpOperands;
+  DenseSet<OpOperand *> copiedOpOperands;
   SmallVector<OpResult> outOfPlaceOpResults;
+  DenseSet<OpResult> copiedOpResults;
 
   // Find all out-of-place OpOperands.
   for (OpOperand &opOperand : op->getOpOperands()) {
     Type operandType = opOperand.get().getType();
     if (!operandType.isa<TensorType>())
       continue;
     if (state.isInPlace(opOperand))
       continue;
     if (operandType.isa<UnrankedTensorType>())
       return op->emitError("copies of unranked tensors are not supported");
 
     SmallVector<OpResult> aliasingOpResults =
         state.getAliasingOpResult(opOperand);
     if (aliasingOpResults.size() == 1 &&
         !state.bufferizesToMemoryWrite(opOperand) &&
         state.getAliasingOpOperand(aliasingOpResults.front()).size() == 1) {
       // The op itself does not write but may create exactly one alias. Instead
       // of copying the OpOperand, copy the OpResult. The OpResult can sometimes
       // be smaller than the OpOperand (e.g., in the case of an extract_slice,
       // where the result is usually a smaller part of the source).
       outOfPlaceOpResults.push_back(aliasingOpResults.front());
+      if (!state.canOmitTensorCopy(opOperand))
+        copiedOpResults.insert(aliasingOpResults.front());
     } else {
       // In all other cases, make a copy of the OpOperand.
       outOfPlaceOpOperands.push_back(&opOperand);
+      if (!state.canOmitTensorCopy(opOperand))
+        copiedOpOperands.insert(&opOperand);
     }
   }
 
   // Insert copies of OpOperands.
   rewriter.setInsertionPoint(op);
   for (OpOperand *opOperand : outOfPlaceOpOperands) {
-    auto tensorType = opOperand->get().getType().cast<RankedTensorType>();
     SmallVector<OpResult> aliasingOpResults =
         state.getAliasingOpResult(*opOperand);
     bool escape = llvm::any_of(
         aliasingOpResults, [&](Value v) { return state.isTensorYielded(v); });
-    Value copy = rewriter.create<AllocTensorOp>(
-        op->getLoc(), tensorType, ValueRange(), opOperand->get(), escape);
+    Value copy = allocateTensorForShapedValue(
+        rewriter, op->getLoc(), opOperand->get(), escape,
+        copiedOpOperands.contains(opOperand));
     rewriter.updateRootInPlace(op, [&]() { opOperand->set(copy); });
   }
 
   // Insert copies of OpResults.
   rewriter.setInsertionPointAfter(op);
   for (OpResult opResult : outOfPlaceOpResults) {
-    auto tensorType = opResult.getType().cast<RankedTensorType>();
     bool escape = state.isTensorYielded(opResult);
-    Value copy = rewriter.create<AllocTensorOp>(op->getLoc(), tensorType,
-                                                ValueRange(), opResult, escape);
+    Value copy =
+        allocateTensorForShapedValue(rewriter, op->getLoc(), opResult, escape,
+                                     copiedOpResults.count(opResult));
     SmallVector<OpOperand *> uses = llvm::to_vector(llvm::map_range(
         opResult.getUses(), [](OpOperand &use) { return &use; }));
     for (OpOperand *use : uses) {
       // Do not update the alloc_tensor op that we just created.
       if (use->getOwner() != copy.getDefiningOp())
         rewriter.updateRootInPlace(use->getOwner(), [&]() { use->set(copy); });
     }
   }

 
 AnalysisState::AnalysisState(const BufferizationOptions &options)
     : options(options) {
   for (const BufferizationOptions::AnalysisStateInitFn &fn :
        options.stateInitializers)
     fn(*this);
 }
 
+bool AnalysisState::canOmitTensorCopy(OpOperand &opOperand) const {
+  // Do not copy if the tensor has undefined contents.
+  if (hasUndefinedContents(&opOperand))
+    return true;
+
+  // Do not copy if the buffer of the tensor is entirely overwritten (with
+  // values that do not depend on the old tensor).
+  if (bufferizesToMemoryWrite(opOperand) && !bufferizesToMemoryRead(opOperand))
+    return true;
+
+  // Do not copy if the tensor is never read.
+  SmallVector<OpResult> aliasingOpResults = getAliasingOpResult(opOperand);
+  if (!bufferizesToMemoryRead(opOperand) &&
+      llvm::none_of(aliasingOpResults,
+                    [&](OpResult opResult) { return isValueRead(opResult); }))
+    return true;
+
+  // Default: Cannot omit the copy.
+  return false;
+}
+
 // bufferization.to_memref is not allowed to change the rank.
 static void ensureToMemrefOpIsValid(Value tensor, Type memrefType) {
 #ifndef NDEBUG
   auto rankedTensorType = tensor.getType().dyn_cast<RankedTensorType>();
   assert((!rankedTensorType || memrefType.cast<MemRefType>().getRank() ==
                                    rankedTensorType.getRank()) &&
          "to_memref would be invalid: mismatching ranks");
 #endif

mlir/test/Dialect/Bufferization/Transforms/tensor-copy-insertion.mlir

 // CHECK-LABEL: func @return_alloc_tensor
 // CHECK-FUNC-LABEL: func @return_alloc_tensor
 func.func @return_alloc_tensor() -> (tensor<5xf32>) {
   // CHECK: bufferization.alloc_tensor() {escape = false} : tensor<5xf32>
   // CHECK-FUNC: bufferization.alloc_tensor() {escape = true} : tensor<5xf32>
   %0 = bufferization.alloc_tensor() : tensor<5xf32>
   return %0 : tensor<5xf32>
 }
+
+// -----
+
+// CHECK-LABEL: func @do_not_copy_undefined_tensor
+func.func @do_not_copy_undefined_tensor(%f: f32, %idx: index)
+  -> (tensor<5xf32>, tensor<5xf32>)
+{
+  // CHECK: bufferization.alloc_tensor() {escape = false} : tensor<5xf32>
+  // The second alloc_tensor should not have a copy operand.
+  // CHECK: bufferization.alloc_tensor() {escape = false} : tensor<5xf32>
+  %0 = bufferization.alloc_tensor() : tensor<5xf32>
+  %1 = tensor.insert %f into %0[%idx] : tensor<5xf32>
+  return %0, %1 : tensor<5xf32>, tensor<5xf32>
+}
+
+// -----
+
+// CHECK-LABEL: func @do_not_copy_when_overwritten
+func.func @do_not_copy_when_overwritten(%t: tensor<5xf32>, %f: f32)
+  -> (tensor<5xf32>, tensor<5xf32>)
+{
+  // CHECK: %[[alloc:.*]] = bufferization.alloc_tensor() {escape = false} : tensor<5xf32>
+  // CHECK: linalg.generic {{.*}} outs(%[[alloc]] : tensor<5xf32>)
+  %r = linalg.generic {
+    indexing_maps = [affine_map<(d0) -> (d0)>],
+    iterator_types = ["parallel"]}
+    outs(%t : tensor<5xf32>) {
+      ^bb0(%arg0 : f32) :
+        linalg.yield %f : f32
+    } -> tensor<5xf32>
+  return %t, %r : tensor<5xf32>, tensor<5xf32>
+}
+
+// -----
+
+// CHECK-LABEL: func @do_not_copy_when_result_not_read
+func.func @do_not_copy_when_result_not_read(%t: tensor<5xf32>, %f: f32)
+  -> (tensor<3xf32>)
+{
+  %0 = tensor.extract_slice %t[0][3][1] : tensor<5xf32> to tensor<3xf32>
+  // CHECK: %[[alloc:.*]] = bufferization.alloc_tensor() {escape = false} : tensor<3xf32>
+  // CHECK: linalg.generic {{.*}} outs(%[[alloc]] : tensor<3xf32>)
+  %r = linalg.generic {
+    indexing_maps = [affine_map<(d0) -> (d0)>],
+    iterator_types = ["parallel"]}
+    outs(%0 : tensor<3xf32>) {
+      ^bb0(%arg0 : f32) :
+        linalg.yield %f : f32
+    } -> tensor<3xf32>
+  return %r : tensor<3xf32>
+}