[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

215–219

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

215–219

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

215–219

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

215–219

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

[pifon2a]

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

215–219

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

215–219

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

712

nit: methods

[springerm]

rebase

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

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=*/[{
           // 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=*/[{
           // 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.

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

+//===- DstBufferizableOpInterfaceImpl.h - Dst Op Bufferization --*- 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_DSTBUFFERIZABLEOPINTERFACEIMPL_H_
+#define MLIR_DIALECT_BUFFERIZATION_IR_DSTBUFFERIZABLEOPINTERFACEIMPL_H_
+
+#include "mlir/Dialect/Bufferization/IR/BufferizableOpInterface.h"
+#include "mlir/Interfaces/DestinationStyleOpInterface.h"
+
+namespace mlir {
+namespace bufferization {
+
+/// Bufferizable ops that implement the DestinationStyleOpInterface can use this
+/// external model base class. It provides default implementations for various
+/// required interface methods.
+template <typename ConcreteModel, typename ConcreteOp>
+struct DstBufferizableOpInterfaceExternalModel
+    : public BufferizableOpInterface::ExternalModel<ConcreteModel, ConcreteOp> {
+  bool bufferizesToMemoryRead(Operation *op, OpOperand &opOperand,
+                              const AnalysisState &state) const {
+    // All inputs and outputs bufferize to a memory read.
+    assert(isa<DestinationStyleOpInterface>(op) &&
+           "expected that op implements DestinationStyleOpInterface");
+    return true;
+  }
+
+  bool bufferizesToMemoryWrite(Operation *op, OpOperand &opOperand,
+                               const AnalysisState &state) const {
+    // Only outputs bufferize to a memory write.
+    auto dstOp = cast<DestinationStyleOpInterface>(op);
+    return dstOp.isOutput(&opOperand);
+  }
+
+  SmallVector<OpResult> getAliasingOpResult(Operation *op, OpOperand &opOperand,
+                                            const AnalysisState &state) const {
+    // Output operands alias with their respective tied OpResults.
+    auto dstOp = cast<DestinationStyleOpInterface>(op);
+    if (dstOp.isOutput(&opOperand))
+      return {dstOp.getTiedOpResult(&opOperand)};
+    return {};
+  }
+
+  BufferRelation bufferRelation(Operation *op, OpResult opResult,
+                                const AnalysisState &state) const {
+    assert(isa<DestinationStyleOpInterface>(op) &&
+           "expected that op implements DestinationStyleOpInterface");
+    return BufferRelation::Equivalent;
+  }
+};
+
+} // namespace bufferization
+} // namespace mlir
+
+#endif // MLIR_DIALECT_BUFFERIZATION_IR_DSTBUFFERIZABLEOPINTERFACEIMPL_H_

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

 //===- BufferizableOpInterfaceImpl.cpp - Impl. of BufferizableOpInterface -===//
 //
 // 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/Tensor/Transforms/BufferizableOpInterfaceImpl.h"
 #include "mlir/Dialect/Affine/IR/AffineOps.h"
 #include "mlir/Dialect/Arith/IR/Arith.h"
 #include "mlir/Dialect/Bufferization/IR/BufferizableOpInterface.h"
 #include "mlir/Dialect/Bufferization/IR/Bufferization.h"
+#include "mlir/Dialect/Bufferization/IR/DstBufferizableOpInterfaceImpl.h"
 #include "mlir/Dialect/MemRef/IR/MemRef.h"
 #include "mlir/Dialect/SCF/IR/SCF.h"
 #include "mlir/Dialect/Tensor/IR/Tensor.h"
 #include "mlir/Dialect/Utils/StaticValueUtils.h"
 #include "mlir/IR/Dialect.h"
 #include "mlir/IR/Operation.h"
 
 using namespace mlir;
LogicalResult bufferize(Operation *op, RewriterBase &rewriter,
 
     replaceOpWithBufferizedValues(rewriter, op, buffer);
 
     return success();
   }
 };
 
 /// Bufferization of tensor.insert. Replace with memref.store.
+///
+/// Note: DstBufferizableOpInterfaceExternalModel provides many default method
+/// implementations for DestinationStyle ops.
 struct InsertOpInterface
-    : public BufferizableOpInterface::ExternalModel<InsertOpInterface,
+    : public DstBufferizableOpInterfaceExternalModel<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: DstBufferizableOpInterfaceExternalModel provides many default method
+/// implementations for DestinationStyle ops.
 struct InsertSliceOpInterface
-    : public BufferizableOpInterface::ExternalModel<InsertSliceOpInterface,
+    : public DstBufferizableOpInterfaceExternalModel<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

cc_library(
     srcs = [
         "lib/Dialect/Bufferization/IR/BufferizableOpInterface.cpp",
         "lib/Dialect/Bufferization/IR/BufferizationDialect.cpp",
         "lib/Dialect/Bufferization/IR/BufferizationOps.cpp",
     ],
     hdrs = [
         "include/mlir/Dialect/Bufferization/IR/BufferizableOpInterface.h",
         "include/mlir/Dialect/Bufferization/IR/Bufferization.h",
+        "include/mlir/Dialect/Bufferization/IR/DstBufferizableOpInterfaceImpl.h",
     ],
     includes = ["include"],
     deps = [
         ":AffineDialect",
         ":AllocationOpInterface",
         ":ArithDialect",
         ":BufferizableOpInterfaceIncGen",
         ":BufferizationBaseIncGen",
         ":BufferizationOpsIncGen",
         ":ControlFlowInterfaces",
         ":CopyOpInterface",
+        ":DestinationStyleOpInterface",
         ":FuncDialect",
         ":IR",
         ":InferTypeOpInterface",
         ":MemRefDialect",
         ":SparseTensorDialect",
         ":Support",
         ":TensorDialect",
         "//llvm:Support",