[mlir][linalg][bufferize] Support custom insertion point for buffer copies

[springerm]

Summary

By default, copies are inserted right before the tensor OpOperand use. With this change, bufferize implementation can change the insertion point. This is needed for some ops where it would be illegal to insert a copy right before the use.

Depends On D117224

mlir/include/mlir/Dialect/Linalg/ComprehensiveBufferize/BufferizableOpInterface.h

public:
   void createMemCpy(OpBuilder &b, Location loc, Value from, Value to) const;
 
   /// Return `true` if the given OpResult has been decided to bufferize inplace.
   bool isInPlace(OpOperand &opOperand) const;
 
   /// Return the buffer (memref) for a given OpOperand (tensor). Allocate
   /// a new buffer and copy over data from the existing buffer if out-of-place
   /// bufferization was decided.
-  FailureOr<Value> getBuffer(RewriterBase &rewriter, OpOperand &opOperand,
-                             bool forceInPlace = false) const;
+  FailureOr<Value>
+  getBuffer(RewriterBase &rewriter, OpOperand &opOperand,
+            bool forceInPlace = false,
+            Optional<Operation *> customCopyInsertionPoint = None) const;
 
   /// Return dialect-specific bufferization state.
   template <typename StateT>
   Optional<const StateT *> getDialectState(StringRef name) const {
     auto it = dialectState.find(name);
     if (it == dialectState.end())
       return None;
     return static_cast<const StateT *>(it->getSecond().get());

mlir/lib/Dialect/Linalg/ComprehensiveBufferize/BufferizableOpInterface.cpp

return rewriter.create<bufferization::ToMemrefOp>(tensor.getLoc(), memrefType,
                                                     tensor);
 }
 
 /// Return the result buffer (memref) for a given OpResult (tensor). Allocate
 /// a new buffer and copy over data from the existing buffer if out-of-place
 /// bufferization is necessary.
 FailureOr<Value>
 mlir::linalg::comprehensive_bufferize::BufferizationState::getBuffer(
-    RewriterBase &rewriter, OpOperand &opOperand, bool forceInPlace) const {
+    RewriterBase &rewriter, OpOperand &opOperand, bool forceInPlace,
+    Optional<Operation *> customCopyInsertionPoint) const {
   OpBuilder::InsertionGuard guard(rewriter);
   Operation *op = opOperand.getOwner();
   Location loc = op->getLoc();
   Value operand = opOperand.get();
   Value operandBuffer = lookupBuffer(rewriter, operand);
 
   if (forceInPlace || aliasInfo.isInPlace(opOperand))
     return operandBuffer;
mlir::linalg::comprehensive_bufferize::BufferizationState::getBuffer(
   OpResult aliasingOpResult = getAliasingOpResult(opOperand);
   if (aliasingOpResult && !bufferizesToMemoryRead(opOperand) &&
       !isValueRead(aliasingOpResult))
     return resultBuffer;
   // Do not copy if this op does not read the data, but writes it.
   if (bufferizesToMemoryWrite(opOperand) && !bufferizesToMemoryRead(opOperand))
     return resultBuffer;
 
+  if (customCopyInsertionPoint) {
+    rewriter.setInsertionPoint(*customCopyInsertionPoint);
+  } else {
     // The copy happens right before the op that is bufferized.
     rewriter.setInsertionPoint(op);
+  }
   createMemCpy(rewriter, loc, operandBuffer, *resultBuffer);
+
   return resultBuffer;
 }
 
 void mlir::linalg::comprehensive_bufferize::replaceOpWithBufferizedValues(
     RewriterBase &rewriter, Operation *op, ValueRange values) {
   OpBuilder::InsertionGuard g(rewriter);
 
   // Replace all OpResults with the given values.