[mlir][bufferize] Provide default BufferizableOpInterface impl for destination style ops

[springerm]

Summary

tensor.insert and tensor.insert_slice (as destination style ops) do no longer need to implement the entire BufferizableOpInterface.

Depends On D136346

[pifon2a]

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

237–241

why is it included in the default implementation of bufferization interface? Shouldn't it be outside of default implementation for bufferization interface. Otherwise, continuing this logic, I could also check here if (auto pad = dyn_cast<PadOp>(..)) and then write code for PadOp bufferization here.

[springerm]

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

237–241

What do you mean by "outside"?

All DPS ops have this property. So by putting this logic into the default implementation, there is one less method to implement.

Most BufferizableOpInterface impls of DPS ops will become real simple. E.g., take a look at tensor.insert (InsertOpInterface in this revision). Only the bufferize method must be implemented.

[pifon2a]

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

237–241

I mean by "outside" not in the definition of BufferizableOpInterface. I would prefer if we had an external model for DPS Ops similar to what we had in Linalg/Transforms/BufferizableOpInterfaceImpl.cpp for all Linalg ops.

[springerm]

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

237–241

Unfortunately only ops can implement interfaces. Interfaces can currently not implement other interfaces.

[pifon2a]

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

237–241

We can template it w.r.t. operation type, can't we? Then the users would just register an external model for bufferizable interface for DPS ops. Yes, it will be 1 registration per op, but IMO it is cleaner than writing it in td.

[springerm]

address comments

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

237–241

We discussed offline. Instead of putting the default impl into the interface defaultImplementation, we provide a templatized external model struct with the default implementations.

[pifon2a]

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

711

nit: methods

[springerm]

rebase

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

 //===- BufferizableOpInterface.h - Bufferizable Ops -------------*- C++ -*-===//
 //
 // 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
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef MLIR_DIALECT_BUFFERIZATION_IR_BUFFERIZABLEOPINTERFACE_H_
 #define MLIR_DIALECT_BUFFERIZATION_IR_BUFFERIZABLEOPINTERFACE_H_
 
 #include "mlir/IR/Operation.h"
 #include "mlir/IR/PatternMatch.h"
+#include "mlir/Interfaces/DestinationStyleOpInterface.h"
 #include "mlir/Support/LLVM.h"
 #include "llvm/ADT/SetVector.h"
 
 namespace mlir {
 class OpBuilder;
 
 namespace bufferization {
 

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

let methods = [
           interface) as reading OpOperands.
         }],
         /*retType=*/"bool",
         /*methodName=*/"bufferizesToMemoryRead",
         /*args=*/(ins "OpOperand &":$opOperand,
                       "const AnalysisState &":$state),
         /*methodBody=*/"",
         /*defaultImplementation=*/[{
+          if (isa<DestinationStyleOpInterface>($_op.getOperation())) {
+            // Default for all destination style ops: All Inputs and outputs
+            // bufferize to a memory read.
+            return true;
+          }
+
           // Does not have to be implemented for ops without tensor OpOperands.
           llvm_unreachable("bufferizesToMemoryRead not implemented");
          }]
       >,
       InterfaceMethod<
         /*desc=*/[{
           Return `true` if the given OpOperand bufferizes to a memory write.
 
let methods = [
           interface) as writing OpOperands.
         }],
         /*retType=*/"bool",
         /*methodName=*/"bufferizesToMemoryWrite",
         /*args=*/(ins "OpOperand &":$opOperand,
                       "const AnalysisState &":$state),
         /*methodBody=*/"",
         /*defaultImplementation=*/[{
+          if (auto dstOp =
+                  dyn_cast<DestinationStyleOpInterface>($_op.getOperation())) {
+            // Default for all destination style ops: Only outputs bufferize to
+            // a memory write.
+            return dstOp.isOutput(&opOperand);
+          }
+
           // Does not have to be implemented for ops without tensor OpOperands.
           // Does not have to be implemented for OpOperands that do not have an
           // aliasing OpResult.
           llvm_unreachable("bufferizesToMemoryWrite not implemented");
          }]
       >,
       InterfaceMethod<
           /*desc=*/[{
let methods = [
           return false;
         }]
       >,
       InterfaceMethod<
         /*desc=*/[{
           Return the OpResult that aliases with a given OpOperand when
           bufferized in-place. This method will never be called on OpOperands
           that do not have a tensor type.
+
+          Note: This method can return multiple OpResults, indicating that a
+          given OpOperand may at runtime alias with any (or multiple) of the
+          returned OpResults.
         }],
         /*retType=*/"SmallVector<OpResult>",
         /*methodName=*/"getAliasingOpResult",
         /*args=*/(ins "OpOperand &":$opOperand,
                       "const AnalysisState &":$state),
         /*methodBody=*/"",
         /*defaultImplementation=*/[{
+          if (auto dstOp =
+                  dyn_cast<DestinationStyleOpInterface>($_op.getOperation())) {
+            // Default for all destination style ops: Output operands alias with
+            // their respective tied OpResults.
+            if (dstOp.isOutput(&opOperand))
+              return { dstOp.getTiedOpResult(&opOperand) };
+            return {};
+          }
+
           // Does not have to be implemented for ops without tensor OpOperands.
           llvm_unreachable("getAliasingOpResult not implemented");
         }]
       >,
       InterfaceMethod<
         /*desc=*/[{
           Return the OpOperands that alias with a given OpResult when
           bufferized in-place. This method will never be called on OpResults
           that do not have a tensor type.
 
           By default, this method is the inverse of `getAliasingOpResult`. Ops
           with a region that yield values may want to override this method to
           return the OpOperands that are yielded by the terminator.
 
           Note: This method can return multiple OpOperands, indicating that the
-          given OpResult may at runtime alias with any of the OpOperands. This
-          is useful for branches and for ops such as `arith.select`.
+          given OpResult may at runtime alias with any (or multiple) of the
+          returned OpOperands. This can be useful for branches and for ops such
+          as `arith.select`.
         }],
         /*retType=*/"SmallVector<OpOperand *>",
         /*methodName=*/"getAliasingOpOperand",
         /*args=*/(ins "OpResult":$opResult,
                       "const AnalysisState &":$state),
         /*methodBody=*/"",
         /*defaultImplementation=*/[{
           assert(opResult.getType().isa<TensorType>() &&
let methods = [
           OpResult".
         }],
         /*retType=*/"BufferRelation",
         /*methodName=*/"bufferRelation",
         /*args=*/(ins "OpResult":$opResult,
                       "const AnalysisState &":$state),
         /*methodBody=*/"",
         /*defaultImplementation=*/[{
+          if (auto dstOp =
+                  dyn_cast<DestinationStyleOpInterface>($_op.getOperation())) {
+            // Default for all destination style ops: Equivalent buffers.
+            return BufferRelation::Equivalent;
+          }
+
           // 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

mlir/lib/Dialect/Bufferization/IR/CMakeLists.txt

 add_mlir_dialect_library(MLIRBufferizationDialect
   AllocationOpInterface.cpp
   BufferizableOpInterface.cpp
   BufferizationOps.cpp
   BufferizationDialect.cpp
 
   ADDITIONAL_HEADER_DIRS
   ${MLIR_MAIN_INCLUDE_DIR}/mlir/Dialect/Bufferization
 
   DEPENDS
   MLIRAllocationOpInterfaceIncGen
   MLIRBufferizationOpsIncGen
 
   LINK_LIBS PUBLIC
   MLIRAffineDialect
+  MLIRDestinationStyleOpInterface
   MLIRDialect
   MLIRFuncDialect
   MLIRIR
   MLIRSparseTensorDialect
   MLIRTensorDialect
   MLIRMemRefDialect
   )

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

LogicalResult bufferize(Operation *op, RewriterBase &rewriter,
 
     replaceOpWithBufferizedValues(rewriter, op, buffer);
 
     return success();
   }
 };
 
 /// Bufferization of tensor.insert. Replace with memref.store.
+///
+/// Note: InsertOp implements DestinationStyleOpInterface, so it is sufficient
+/// to implement only the `bufferize` interface method.
 struct InsertOpInterface
     : public BufferizableOpInterface::ExternalModel<InsertOpInterface,
                                                     tensor::InsertOp> {
-  bool bufferizesToMemoryRead(Operation *op, OpOperand &opOperand,
-                              const AnalysisState &state) const {
-    return true;
-  }
-
-  bool bufferizesToMemoryWrite(Operation *op, OpOperand &opOperand,
-                               const AnalysisState &state) const {
-    return true;
-  }
-
-  SmallVector<OpResult> getAliasingOpResult(Operation *op, OpOperand &opOperand,
-                                            const AnalysisState &state) const {
-    assert(&opOperand == &op->getOpOperand(1) /*dest*/ &&
-           "expected dest OpOperand");
-    return {op->getOpResult(0)};
-  }
-
-  SmallVector<OpOperand *>
-  getAliasingOpOperand(Operation *op, OpResult opResult,
-                       const AnalysisState &state) const {
-    return {&op->getOpOperand(1) /*dest*/};
-  }
-
   LogicalResult bufferize(Operation *op, RewriterBase &rewriter,
                           const BufferizationOptions &options) const {
     auto insertOp = cast<tensor::InsertOp>(op);
     FailureOr<Value> destMemref =
         getBuffer(rewriter, insertOp.getDest(), options);
     if (failed(destMemref))
       return failure();
     rewriter.create<memref::StoreOp>(insertOp.getLoc(), insertOp.getScalar(),
                                      *destMemref, insertOp.getIndices());
     replaceOpWithBufferizedValues(rewriter, op, *destMemref);
     return success();
   }
-
-  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 areEquivalentSlices(const AnalysisState &state,
                                 ExtractSliceOp extractSliceOp,
                                 OpTy insertSliceOp) {
if (uConflictingWrite == &insertSliceOp->getOpOperand(1) /*dest*/ &&
                                  insertSliceOp))
       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.
+///
+/// Note: InsertSliceOp implements DestinationStyleOpInterface for which the
+/// BufferizableOpInterface provides many default method implementations.
 struct InsertSliceOpInterface
     : public BufferizableOpInterface::ExternalModel<InsertSliceOpInterface,
                                                     tensor::InsertSliceOp> {
-  bool bufferizesToMemoryRead(Operation *op, OpOperand &opOperand,
-                              const AnalysisState &state) const {
-    return true;
-  }
-
-  bool bufferizesToMemoryWrite(Operation *op, OpOperand &opOperand,
-                               const AnalysisState &state) const {
-    return &opOperand == &op->getOpOperand(1) /*dest*/;
-  }
-
-  SmallVector<OpResult> getAliasingOpResult(Operation *op, OpOperand &opOperand,
-                                            const AnalysisState &state) const {
-    if (&opOperand == &op->getOpOperand(1) /*dest*/)
-      return {op->getResult(0)};
-    return {};
-  }
-
-  BufferRelation bufferRelation(Operation *op, OpResult opResult,
-                                const AnalysisState &state) const {
-    return BufferRelation::Equivalent;
-  }
-
   bool isNotConflicting(Operation *op, OpOperand *uRead,
                         OpOperand *uConflictingWrite,
                         const AnalysisState &state) const {
     return isNotConflictingInsertSliceLikeOp<tensor::InsertSliceOp>(
         op, uRead, uConflictingWrite, state);
   }
 
   LogicalResult bufferize(Operation *op, RewriterBase &rewriter,

utils/bazel/llvm-project-overlay/mlir/BUILD.bazel

deps = [
         ":AffineDialect",
         ":AllocationOpInterface",
         ":ArithDialect",
         ":BufferizableOpInterfaceIncGen",
         ":BufferizationBaseIncGen",
         ":BufferizationOpsIncGen",
         ":ControlFlowInterfaces",
         ":CopyOpInterface",
+        ":DestinationStyleOpInterface",
         ":FuncDialect",
         ":IR",
         ":InferTypeOpInterface",
         ":MemRefDialect",
         ":SparseTensorDialect",
         ":Support",
         ":TensorDialect",
         "//llvm:Support",