[mlir][bufferization][NFC] Put inplacability conflict resolution in op interface

[springerm]

Summary

The TensorCopyInsertion pass resolves out-of-place bufferization decisions by inserting explicit bufferization.alloc_tensor ops. This change moves that functionality into a new BufferizableOpInterface method, so that it can be overridden by op implementations. Some op bufferizations must insert additional alloc_tensor ops to make sure that certain aliasing invariants are not violated (e.g., scf::ForOp). This will be addressed in a subsequent change.

Depends On D126573

[springerm]

update

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

let methods = [
         /*defaultImplementation=*/[{
           // Does not have to be implemented for ops without tensor OpResults
           // that have an aliasing OpOperand.
           llvm_unreachable("bufferRelation not implemented");
         }]
       >,
       InterfaceMethod<
         /*desc=*/[{
+          Resolve all inplacability conflicts by inserting explicit
+          `bufferization.alloc_tensor` ops. Examples of inplacability conflicts
+          are read-after-write conflicts or writes into non-writable buffers.
+
+          This method should rewrite the IR in such a way that for each tensor
+          OpOperand t, buffer(t) can be directly used when during bufferization.
+          The bufferization does no longer have to care about inplacability
+          conflicts.
+
+          This method can query analysis information from the given analysis
+          state.
+        }],
+        /*retType=*/"LogicalResult",
+        /*methodName=*/"resolveConflicts",
+        /*args=*/(ins "RewriterBase &":$rewriter,
+                      "const AnalysisState &":$state),
+        /*methodBody=*/"",
+        /*defaultImplementation=*/[{
+          auto bufferizableOp =
+              cast<BufferizableOpInterface>($_op.getOperation());
+          return bufferizableOp.resolveTensorOpOperandConflicts(
+              rewriter, state);
+        }]
+      >,
+      InterfaceMethod<
+        /*desc=*/[{
           Bufferize this op, i.e., rewrite it into a memref-based equivalent.
           Buffers of tensor SSA values can be retrieved via `state.getBuffer`.
           Uses of tensor results of the existing tensor op can be replaced with
           `replaceOpWithBufferizedValues` or `replaceOpWithNewBufferizedOp`.
           These two functions automatically handle the tensor-to-memref type
           conversion.
 
           The implementation of this method must be consistent with the
let methods = [
         /*methodBody=*/"",
         /*defaultImplementation=*/[{
           return success();
         }]
       >
   ];
 
   let extraClassDeclaration = [{
+    /// Resolve out-of-place tensor OpOperands with explicit allocations in the
+    /// form of `bufferization.alloc_tensor` ops.
+    LogicalResult resolveTensorOpOperandConflicts(
+        RewriterBase &rewriter, const AnalysisState &state);
+
     /// Return `true` if the given OpOperand creates an alias but does neither
     /// read nor write. This implies that `bufferizesToMemoryRead` and
     /// `bufferizesToMemoryWrite` must return `false`. This method will never
     /// be called on OpOperands that do not have a tensor type.
     ///
     /// Examples of such ops are `tensor.extract_slice` and `tensor.cast`.
     bool bufferizesToAliasOnly(OpOperand &opOperand,
                                const AnalysisState &state) {

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

 #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"
 
+//===----------------------------------------------------------------------===//
+// BufferizableOpInterface
+//===----------------------------------------------------------------------===//
+
 namespace mlir {
 namespace bufferization {
 
 #include "mlir/Dialect/Bufferization/IR/BufferizableOpInterface.cpp.inc"
 
 } // namespace bufferization
 } // namespace mlir
 
 #define DEBUG_TYPE "bufferizable-op-interface"
 #define DBGS() (llvm::dbgs() << '[' << DEBUG_TYPE << "] ")
 #define LDBG(X) LLVM_DEBUG(DBGS() << (X))
 
 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;
 
+LogicalResult BufferizableOpInterface::resolveTensorOpOperandConflicts(
+    RewriterBase &rewriter, const AnalysisState &state) {
+  Operation *op = getOperation();
+  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");
+    auto tensorType = operandType.dyn_cast<RankedTensorType>();
+    if (!tensorType)
+      continue;
+    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);
+    rewriter.updateRootInPlace(op, [&]() { opOperand.set(copy); });
+  }
+  return success();
+}
+
 //===----------------------------------------------------------------------===//
 // OpFilter
 //===----------------------------------------------------------------------===//
 
 bool OpFilter::isOpAllowed(Operation *op) const {
   // All other ops: Allow/disallow according to filter.
   bool isAllowed = !hasAllowRule();
   for (const Entry &entry : entries) {

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

if (options.testAnalysisOnly)
     return success();
 
   return insertTensorCopies(op, state);
 }
 
 LogicalResult
 mlir::bufferization::insertTensorCopies(Operation *op,
                                         const AnalysisState &state) {
-  OpBuilder builder(op->getContext());
+  IRRewriter rewriter(op->getContext());
   WalkResult result = op->walk([&](Operation *op) {
     auto bufferizableOp = state.getOptions().dynCastBufferizableOp(op);
     if (!bufferizableOp)
       return WalkResult::skip();
 
     // Find AllocTensorOps without an `escape` attribute and add the attribute
     // based on analysis results.
     if (auto allocTensorOp = dyn_cast<AllocTensorOp>(op)) {
       if (allocTensorOp.escape())
         return WalkResult::advance();
       bool escape = state.isTensorYielded(allocTensorOp.result());
-      allocTensorOp.escapeAttr(builder.getBoolAttr(escape));
+      allocTensorOp.escapeAttr(rewriter.getBoolAttr(escape));
       return WalkResult::advance();
     }
 
-    // Find out-of-place tensor OpOperands and resolve them with an explicit
-    // tensor copy in the form of an AllocTensorOp.
-    builder.setInsertionPoint(op);
-    for (OpOperand &opOperand : op->getOpOperands()) {
-      if (opOperand.get().getType().isa<UnrankedTensorType>()) {
-        op->emitError("copies of unranked tensors are not supported");
-        return WalkResult::interrupt();
-      }
-      auto tensorType = opOperand.get().getType().dyn_cast<RankedTensorType>();
-      if (!tensorType)
-        continue;
-      if (state.isInPlace(opOperand))
-        continue;
-      SmallVector<OpResult> aliasingOpResults =
-          state.getAliasingOpResult(opOperand);
-      bool escape = llvm::any_of(
-          aliasingOpResults, [&](Value v) { return state.isTensorYielded(v); });
-      Value copy = builder.create<AllocTensorOp>(
-          op->getLoc(), tensorType, ValueRange(), opOperand.get(), escape);
-      opOperand.set(copy);
-    }
+    // Find inplacability conflicts and resolve them. (Typically with explicit
+    // tensor copies in the form of AllocTensorOps.)
+    rewriter.setInsertionPoint(op);
+    if (failed(bufferizableOp.resolveConflicts(rewriter, state)))
 
     return WalkResult::advance();
   });
 
   return failure(result.wasInterrupted());
 }
 
 namespace {