[mlir][bufferize] Do not copy buffers with undefined contents

[springerm]

Summary

Buffers with undefined contents (e.g., the result of an init_tensor) are no longer copied.

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

public:
   SetVector<Value> findLastPrecedingWrite(Value value) const;
 
   /// Return `true` if the given OpResult has been decided to bufferize inplace.
   virtual bool isInPlace(OpOperand &opOperand) const = 0;
 
   /// Return true if `v1` and `v2` bufferize to equivalent buffers.
   virtual bool areEquivalentBufferizedValues(Value v1, Value v2) const = 0;
 
+  /// Return `true` if the given tensor has undefined contents.
+  virtual bool hasUndefinedContents(OpOperand *opOperand) const = 0;
+
   /// Return true if the given tensor (or an aliasing tensor) is yielded from
   /// the containing block. Also include all aliasing tensors in the same block.
   ///
   /// Note: In the absence of an analysis, an implementation may return true for
   /// any given tensor.
   virtual bool isTensorYielded(Value tensor) const = 0;
 
   /// Return `true` if the given dialect state exists.
public:
   virtual ~AlwaysCopyAnalysisState() = default;
 
   /// Return `true` if the given OpResult has been decided to bufferize inplace.
   bool isInPlace(OpOperand &opOperand) const override;
 
   /// Return true if `v1` and `v2` bufferize to equivalent buffers.
   bool areEquivalentBufferizedValues(Value v1, Value v2) const override;
 
+  /// Return `true` if the given tensor has undefined contents.
+  bool hasUndefinedContents(OpOperand *opOperand) const override;
+
   /// Return true if the given tensor (or an aliasing tensor) is yielded from
   /// the containing block. Also include all aliasing tensors in the same block.
   bool isTensorYielded(Value tensor) const override;
 };
 
 /// BufferizationState provides helper functions for performing bufferization
 /// rewrites and handling memref buffers.
 struct BufferizationState {

mlir/include/mlir/Dialect/Bufferization/Transforms/OneShotAnalysis.h

public:
   BufferizationAliasInfo &getAliasInfo() { return aliasInfo; }
 
   /// Return `true` if the given OpResult has been decided to bufferize inplace.
   bool isInPlace(OpOperand &opOperand) const override;
 
   /// Return true if `v1` and `v2` bufferize to equivalent buffers.
   bool areEquivalentBufferizedValues(Value v1, Value v2) const override;
 
+  /// Return `true` if the given tensor has undefined contents.
+  bool hasUndefinedContents(OpOperand *opOperand) const override;
+
   /// Return true if the given tensor (or an aliasing tensor) is yielded from
   /// the containing block. Also include all aliasing tensors in the same block.
   bool isTensorYielded(Value tensor) const override;
 
+  /// Find all tensor values in the given operation that have undefined contents
+  /// and store them in `undefinedTensorUses`.
+  void gatherUndefinedTensorUses(Operation *op);
+
   /// Find all tensors that are yielded/returned from a block and store them in
   /// `yieldedTensors`. Also include all aliasing tensors in the same block.
   void gatherYieldedTensors(Operation *op);
 
 private:
   /// `aliasInfo` keeps track of aliasing and equivalent values. Only internal
   /// functions and `runOneShotBufferize` may access this object.
   BufferizationAliasInfo aliasInfo;
 
   /// A set of all tensors (and maybe aliasing tensors) that yielded from a
   /// block.
   DenseSet<Value> yieldedTensors;
+
+  /// A set of uses of tensors that have undefined contents.
+  DenseSet<OpOperand *> undefinedTensorUses;
 };
 
 /// Analyze `op` and its nested ops. Bufferization decisions are stored in
 /// `state`.
 LogicalResult analyzeOp(Operation *op, OneShotAnalysisState &state);
 
 /// Run One-Shot Bufferize on the given op: Analysis + Bufferization
 LogicalResult runOneShotBufferize(Operation *op,
                                   const OneShotBufferizationOptions &options);
 
 } // namespace bufferization
 } // namespace mlir
 
 #endif // MLIR_DIALECT_BUFFERIZATION_TRANSFORMS_ONESHOTANALYSIS_H

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

BufferizationState::getBuffer(RewriterBase &rewriter, OpOperand &opOperand,
   // is not yielded and deallocs are enabled in general.
   bool dealloc = llvm::none_of(aliasingOpResults, [&](Value v) {
     return getAnalysisState().isTensorYielded(v);
   });
   FailureOr<Value> resultBuffer = createAlloc(
       rewriter, loc, operandBuffer, dealloc && getOptions().createDeallocs);
   if (failed(resultBuffer))
     return failure();
-  // Do not copy if the last preceding writes of `operand` are ops that do
-  // not write (skipping ops that merely create aliases). E.g., InitTensorOp.
-  // Note: If `findLastPrecedingWrite` reaches the end of the reverse SSA
-  // use-def chain, it returns that value, regardless of whether it is a
-  // memory write or not.
-  SetVector<Value> lastWrites = analysisState.findLastPrecedingWrite(operand);
-  if (llvm::none_of(lastWrites, [&](Value lastWrite) {
-        if (auto bufferizableOp = options.dynCastBufferizableOp(lastWrite))
-          return bufferizableOp.isMemoryWrite(lastWrite.cast<OpResult>(),
-                                              analysisState);
-        return true;
-      }))
+  // Do not copy the buffer if its contents are undefined.
+  if (analysisState.hasUndefinedContents(&opOperand))
     return resultBuffer;
   // Do not copy if the copied data is never read.
   if (!aliasingOpResults.empty() &&
       !analysisState.bufferizesToMemoryRead(opOperand) &&
       llvm::none_of(aliasingOpResults, [&](OpResult opResult) {
         return analysisState.isValueRead(opResult);
       }))
     return resultBuffer;

 /// Return true if `v1` and `v2` bufferize to equivalent buffers.
 bool AlwaysCopyAnalysisState::areEquivalentBufferizedValues(Value v1,
                                                             Value v2) const {
   // There is no analysis, so we do not know if the values are equivalent. The
   // conservative answer is "false".
   return false;
 }
 
+/// Return `true` if the given tensor has undefined contents.
+bool AlwaysCopyAnalysisState::hasUndefinedContents(OpOperand *opOperand) const {
+  // There is no analysis, so the conservative answer is "false".
+  return false;
+}
+
 /// Return true if the given tensor (or an aliasing tensor) is yielded from
 /// the containing block. Also include all aliasing tensors in the same block.
 bool AlwaysCopyAnalysisState::isTensorYielded(Value tensor) const {
   // There is no analysis, so conservatively answer "true".
   return true;
 }
 
 //===----------------------------------------------------------------------===//

mlir/lib/Dialect/Bufferization/Transforms/OneShotAnalysis.cpp

for (OpOperand &returnValOperand : returnOp->getOpOperands()) {
           yieldedTensors.insert(v);
       });
     }
 
     return WalkResult::advance();
   });
 }
 
+void OneShotAnalysisState::gatherUndefinedTensorUses(Operation *op) {
+  op->walk([&](Operation *op) {
+    // Skip unknown ops.
+    auto bufferizableOp = getOptions().dynCastBufferizableOp(op);
+    if (!bufferizableOp)
+      return WalkResult::skip();
+
+    // Check all tensor OpResults.
+    for (OpResult opResult : op->getOpResults()) {
+      if (!opResult.getType().isa<TensorType>())
+        continue;
+
+      // If there is no preceding memory write, the tensor contents are
+      // undefined.
+      // Note: If `findLastPrecedingWrite` reaches the end of the reverse SSA
+      // use-def chain, it returns that value, regardless of whether it is a
+      // memory write or not.
+      SetVector<Value> lastWrites = findLastPrecedingWrite(opResult);
+      bool isUndefined = llvm::none_of(lastWrites, [&](Value lastWrite) {
+        if (auto bufferizableOp = getOptions().dynCastBufferizableOp(lastWrite))
+          return bufferizableOp.isMemoryWrite(lastWrite.cast<OpResult>(),
+                                              *this);
+        return true;
+      });
+      if (isUndefined)
+        for (OpOperand &use : opResult.getUses())
+          undefinedTensorUses.insert(&use);
+    }
+
+    return WalkResult::advance();
+  });
+}
+
+bool OneShotAnalysisState::hasUndefinedContents(OpOperand *opOperand) const {
+  return undefinedTensorUses.contains(opOperand);
+}
+
 bool OneShotAnalysisState::isTensorYielded(Value tensor) const {
   return yieldedTensors.contains(tensor);
 }
 
 //===----------------------------------------------------------------------===//
 // Bufferization-specific alias analysis.
 //===----------------------------------------------------------------------===//
 
LogicalResult bufferization::analyzeOp(Operation *op,
 
   bool failedAnalysis = false;
   if (!options.allowReturnAllocs) {
     SmallVector<Operation *> newOps;
     failedAnalysis |=
         failed(assertDestinationPassingStyle(op, state, aliasInfo, newOps));
   }
 
-  // Gather all yielded tensors.
+  // Gather some extra analysis data.
   state.gatherYieldedTensors(op);
+  state.gatherUndefinedTensorUses(op);
 
   // Analysis verification: After setting up alias/equivalence sets, each op
   // can check for expected invariants/limitations and fail the analysis if
   // necessary.
   op->walk([&](Operation *op) {
     if (BufferizableOpInterface bufferizableOp =
             options.dynCastBufferizableOp(op))
       failedAnalysis |= failed(bufferizableOp.verifyAnalysis(state));

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

%3 = scf.for %arg3 = %c0 to %c32 step %c8 iter_args(%arg4 = %1) -> (tensor<?x1x6x8xf32>) {
     %4 = affine.apply affine_map<(d0) -> (d0 ceildiv 8)>(%arg3)
     %5 = tensor.insert_slice %2 into %arg4[%4,0, 0, 0] [1, 1, 6, 8] [1, 1, 1, 1] :
       tensor<1x6x8xf32> into tensor<?x1x6x8xf32>
     scf.yield %5 : tensor<?x1x6x8xf32>
   }
   return %3 : tensor<?x1x6x8xf32>
 }
 
+// -----
+
+// CHECK-LABEL: func @do_not_copy_init_tensors(
+func.func @do_not_copy_init_tensors(%f1: f32, %f2: f32, %idx: index)
+  -> (tensor<5xf32>, tensor<5xf32>)
+{
+  // CHECK: memref.alloc
+  // CHECK: memref.alloc
+  // CHECK-NOT: copy
+  // CHECK: memref.store
+  // CHECK: memref.store
+  %0 = linalg.init_tensor [5] : tensor<5xf32>
+  %1 = tensor.insert %f1 into %0[%idx] : tensor<5xf32>
+  %2 = tensor.insert %f2 into %0[%idx] : tensor<5xf32>
+  return %1, %2 : tensor<5xf32>, tensor<5xf32>
+}