[mlir] Move AllocationOpInterface to Bufferize/IR/AllocationOpInterface.td.

[pifon2a]

Summary

Remove it from op defs in MemRefOps.td and make it an external model.

This is the first PR of many that will move bufferization-related ops,
interfaces, passes to Dialect/Bufferize.

RFC: https://llvm.discourse.group/t/rfc-dialect-for-bufferization-related-ops/4712
It is still debated if the comprehensive bufferization has to be moved there as
well, so for now I am just moving the "gradual" bufferization.

[rriddle]

It wasn't clear to me that an agreement had been reached in the RFC to start populating anything here.

[pifon2a]

Rename the dialect to "bufferization".

[pifon2a]

In D114147#3139632, @rriddle wrote:

It wasn't clear to me that an agreement had been reached in the RFC to start populating anything here.

Could you please unblock it?

[herhut]

Just minor nits. Thank you!

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

43 ↗

(On Diff #388812)

This should not be a default implementation because it only applies to buffers allocated with memref::alloc. So it should be a specific implementation for that op.

For clone, I could envision it grows into a more generic clone, especially as it now is moving to the bufferize dialect and is no longer tied to memref. So there using it as default seems more reasonable, but we also need to weaken its type requirements then.

mlir/lib/Transforms/BufferDeallocation.cpp

620

Is this not generated somewhere?

[pifon2a]

Address Stephan's comments.

[pifon2a]

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

43 ↗

(On Diff #388812)

Removed default implementation for buildDealloc.

mlir/lib/Transforms/BufferDeallocation.cpp

620

Yes, it is. I should have used FallbackModel for that.

[mehdi_amini]

This broke the CMake buildbot FYI

mlir/include/mlir/Dialect/Bufferize/CMakeLists.txt

+add_subdirectory(IR)

mlir/include/mlir/Dialect/Bufferize/IR/AllocationOpInterface.h

+//===- AllocationOpInterface.h - Allocation op interface ------------------===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+//
+// This file implements the operation interface for allocation ops.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef MLIR_DIALECT_BUFFERIZE_IR_ALLOCATIONOPINTERFACE_H_
+#define MLIR_DIALECT_BUFFERIZE_IR_ALLOCATIONOPINTERFACE_H_
+
+#include "mlir/Dialect/MemRef/IR/MemRef.h"
+#include "mlir/IR/Builders.h"
+
+#include "mlir/Dialect/Bufferize/IR/AllocationOpInterface.h.inc"
+
+#endif // MLIR_DIALECT_BUFFERIZE_IR_ALLOCATIONOPINTERFACE_H_

mlir/include/mlir/Dialect/Bufferize/IR/AllocationOpInterface.td

+//===-- AllocationOpInterface.td - Allocation op interface -*- tablegen -*-===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+//
+// Defines the interface with allocation-related methods. It is used by the
+// buffer deallocation pass.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef ALLOCATION_OP_INTERFACE
+#define ALLOCATION_OP_INTERFACE
+
+include "mlir/IR/OpBase.td"
+
+//===----------------------------------------------------------------------===//
+// AllocationOpInterface
+//===----------------------------------------------------------------------===//
+
+def AllocationOpInterface : OpInterface<"AllocationOpInterface"> {
+  let description = [{
+    This interface provides general allocation-related methods that are
+    designed for allocation operations. For example, it offers the ability to
+    construct associated deallocation and clone operations that are compatible
+    with the current allocation operation.
+  }];
+  let cppNamespace = "::mlir::bufferize";
+
+  let methods = [
+    StaticInterfaceMethod<[{
+        Builds a deallocation operation using the provided builder and the
+        current allocation value (which refers to the current Op implementing
+        this interface). The allocation value is a result of the current
+        operation implementing this interface. If there is no compatible
+        deallocation operation, this method can return ::llvm::None.
+      }],
+      "::mlir::Optional<::mlir::Operation*>", "buildDealloc",
+      (ins "::mlir::OpBuilder&":$builder, "::mlir::Value":$alloc), [{}],
+      /*defaultImplementation=*/[{
+          return builder.create<memref::DeallocOp>(alloc.getLoc(), alloc)
+            .getOperation();
+      }]
+    >,
+    StaticInterfaceMethod<[{
+        Builds a clone operation using the provided builder and the current
+        allocation value (which refers to the current Op implementing this
+        interface). The allocation value is a result of the current operation
+        implementing this interface. If there is no compatible clone operation,
+        this method can return ::llvm::None.
+      }],
+      "::mlir::Optional<::mlir::Value>", "buildClone",
+      (ins "::mlir::OpBuilder&":$builder, "::mlir::Value":$alloc), [{}],
+      /*defaultImplementation=*/[{
+          return builder.create<memref::CloneOp>(alloc.getLoc(), alloc)
+            .getResult();
+      }]
+    >
+  ];
+}
+
+#endif  // ALLOCATION_OP_INTERFACE

mlir/include/mlir/Dialect/Bufferize/IR/CMakeLists.txt

+add_mlir_interface(AllocationOpInterface)

mlir/include/mlir/Dialect/CMakeLists.txt

 add_subdirectory(Affine)
 add_subdirectory(Arithmetic)
 add_subdirectory(Async)
 add_subdirectory(ArmNeon)
 add_subdirectory(ArmSVE)
 add_subdirectory(AMX)
+add_subdirectory(Bufferize)
 add_subdirectory(Complex)
 add_subdirectory(DLTI)
 add_subdirectory(EmitC)
 add_subdirectory(GPU)
 add_subdirectory(Math)
 add_subdirectory(Linalg)
 add_subdirectory(LLVMIR)
 add_subdirectory(MemRef)

mlir/include/mlir/Dialect/MemRef/IR/MemRefOps.td

def AssumeAlignmentOp : MemRef_Op<"assume_alignment"> {
 
   let assemblyFormat = "$memref `,` $alignment attr-dict `:` type($memref)";
 }
 
 //===----------------------------------------------------------------------===//
 // AllocOp
 //===----------------------------------------------------------------------===//
 
-def MemRef_AllocOp : AllocLikeOp<"alloc", DefaultResource, [
-    DeclareOpInterfaceMethods<AllocationOpInterface,
-      ["buildDealloc", "buildClone"]>]
-  > {
+def MemRef_AllocOp : AllocLikeOp<"alloc", DefaultResource, []> {
   let summary = "memory allocation operation";
   let description = [{
     The `alloc` operation allocates a region of memory, as specified by its
     memref type.
 
     Example:
 
     ```mlir

 
 //===----------------------------------------------------------------------===//
 // CloneOp
 //===----------------------------------------------------------------------===//
 
 def CloneOp : MemRef_Op<"clone", [
     CopyOpInterface,
     DeclareOpInterfaceMethods<MemoryEffectsOpInterface>,
-    DeclareOpInterfaceMethods<AllocationOpInterface,
-      ["buildDealloc", "buildClone"]>
   ]> {
   let builders = [
     OpBuilder<(ins "Value":$value), [{
       return build($_builder, $_state, value.getType(), value);
     }]>];
 
   let description = [{
     Clones the data in the input view into an implicitly defined output view.

mlir/include/mlir/Interfaces/SideEffectInterfaces.td

 //===----------------------------------------------------------------------===//
 
 #ifndef MLIR_INTERFACES_SIDEEFFECTS
 #define MLIR_INTERFACES_SIDEEFFECTS
 
 include "mlir/Interfaces/SideEffectInterfaceBase.td"
 
 //===----------------------------------------------------------------------===//
-// AllocationOpInterface
-//===----------------------------------------------------------------------===//
-
-def AllocationOpInterface : OpInterface<"AllocationOpInterface"> {
-  let description = [{
-    This interface provides general allocation-related methods that are
-    designed for allocation operations. For example, it offers the ability to
-    construct associated deallocation and clone operations that are compatible
-    with the current allocation operation.
-  }];
-  let cppNamespace = "::mlir";
-
-  let methods = [
-    StaticInterfaceMethod<[{
-        Builds a deallocation operation using the provided builder and the
-        current allocation value (which refers to the current Op implementing
-        this interface). The allocation value is a result of the current
-        operation implementing this interface. If there is no compatible
-        deallocation operation, this method can return ::llvm::None.
-      }],
-      "::mlir::Optional<::mlir::Operation*>", "buildDealloc",
-      (ins "::mlir::OpBuilder&":$opBuilder, "::mlir::Value":$alloc), [{}],
-      /*defaultImplementation=*/[{ return llvm::None; }]
-    >,
-    StaticInterfaceMethod<[{
-        Builds a clone operation using the provided builder and the current
-        allocation value (which refers to the current Op implementing this
-        interface). The allocation value is a result of the current operation
-        implementing this interface. If there is no compatible clone operation,
-        this method can return ::llvm::None.
-      }],
-      "::mlir::Optional<::mlir::Value>", "buildClone",
-      (ins "::mlir::OpBuilder&":$opBuilder, "::mlir::Value":$alloc), [{}],
-      /*defaultImplementation=*/[{ return llvm::None; }]
-    >
-  ];
-}
-
-//===----------------------------------------------------------------------===//
 // MemoryEffects
 //===----------------------------------------------------------------------===//
 
 // This def represents the definition for the memory effects interface. Users
 // should generally not use this directly, and should instead use
 // `MemoryEffects`.
 def MemoryEffectsOpInterface
     : EffectOpInterfaceBase<"MemoryEffectOpInterface",

mlir/lib/Dialect/Bufferize/CMakeLists.txt

+add_subdirectory(IR)

mlir/lib/Dialect/Bufferize/IR/AllocationOpInterface.cpp

+//===- AllocationOpInterface.cpp -  Allocation op interface ---------------===//
+//
+// 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/Bufferize/IR/AllocationOpInterface.h"
+#include "mlir/Dialect/Bufferize/IR/AllocationOpInterface.cpp.inc"

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

+add_mlir_library(MLIRAllocationOpInterface
+  AllocationOpInterface.cpp
+
+  ADDITIONAL_HEADER_DIRS
+  ${MLIR_MAIN_INCLUDE_DIR}/mlir/Dialect/Bufferize/IR
+
+  DEPENDS
+  MLIRAllocationOpInterfaceIncGen
+
+  LINK_LIBS PUBLIC
+  MLIRIR
+  )

mlir/lib/Dialect/CMakeLists.txt

 add_subdirectory(Affine)
 add_subdirectory(Arithmetic)
 add_subdirectory(ArmNeon)
 add_subdirectory(ArmSVE)
 add_subdirectory(Async)
 add_subdirectory(AMX)
+add_subdirectory(Bufferize)
 add_subdirectory(Complex)
 add_subdirectory(DLTI)
 add_subdirectory(EmitC)
 add_subdirectory(GPU)
 add_subdirectory(Linalg)
 add_subdirectory(LLVMIR)
 add_subdirectory(Math)
 add_subdirectory(MemRef)

mlir/lib/Dialect/MemRef/IR/MemRefOps.cpp

for (Operation *user : llvm::make_early_inc_range(alloc->getUsers()))
       rewriter.eraseOp(user);
 
     rewriter.eraseOp(alloc);
     return success();
   }
 };
 } // end anonymous namespace.
 
-Optional<Operation *> AllocOp::buildDealloc(OpBuilder &builder, Value alloc) {
-  return builder.create<memref::DeallocOp>(alloc.getLoc(), alloc)
-      .getOperation();
-}
-
-Optional<Value> AllocOp::buildClone(OpBuilder &builder, Value alloc) {
-  return builder.create<memref::CloneOp>(alloc.getLoc(), alloc).getResult();
-}
-
 void AllocOp::getCanonicalizationPatterns(RewritePatternSet &results,
                                           MLIRContext *context) {
   results.add<SimplifyAllocConst<AllocOp>, SimplifyDeadAlloc<AllocOp>>(context);
 }
 
 void AllocaOp::getCanonicalizationPatterns(RewritePatternSet &results,
                                            MLIRContext *context) {
   results.add<SimplifyAllocConst<AllocaOp>, SimplifyDeadAlloc<AllocaOp>>(
void CloneOp::getCanonicalizationPatterns(OwningRewritePatternList &results,
                                           MLIRContext *context) {
   results.insert<SimplifyClones>(context);
 }
 
 OpFoldResult CloneOp::fold(ArrayRef<Attribute> operands) {
   return succeeded(foldMemRefCast(*this)) ? getResult() : Value();
 }
 
-Optional<Operation *> CloneOp::buildDealloc(OpBuilder &builder, Value alloc) {
-  return builder.create<memref::DeallocOp>(alloc.getLoc(), alloc)
-      .getOperation();
-}
-
-Optional<Value> CloneOp::buildClone(OpBuilder &builder, Value alloc) {
-  return builder.create<memref::CloneOp>(alloc.getLoc(), alloc).getResult();
-}
-
 //===----------------------------------------------------------------------===//
 // DeallocOp
 //===----------------------------------------------------------------------===//
 
 LogicalResult DeallocOp::fold(ArrayRef<Attribute> cstOperands,
                               SmallVectorImpl<OpFoldResult> &results) {
   /// dealloc(memrefcast) -> dealloc
   return foldMemRefCast(*this);

mlir/lib/Transforms/BufferDeallocation.cpp

 // The current implementation does not support explicit-control-flow loops and
 // the resulting code will be invalid with respect to program semantics.
 // However, structured control-flow loops are fully supported. Furthermore, it
 // doesn't accept functions which return buffers already.
 //
 //===----------------------------------------------------------------------===//
 
 #include "PassDetail.h"
+#include "mlir/Dialect/Bufferize/IR/AllocationOpInterface.h"
 #include "mlir/Dialect/MemRef/IR/MemRef.h"
 #include "mlir/Dialect/StandardOps/IR/Ops.h"
 #include "mlir/IR/Operation.h"
 #include "mlir/Interfaces/ControlFlowInterfaces.h"
 #include "mlir/Interfaces/LoopLikeInterface.h"
 #include "mlir/Pass/Pass.h"
 #include "mlir/Transforms/BufferUtils.h"
 #include "mlir/Transforms/Passes.h"

 // BufferDeallocation
 //===----------------------------------------------------------------------===//
 
 /// The buffer deallocation transformation which ensures that all allocs in the
 /// program have a corresponding de-allocation. As a side-effect, it might also
 /// introduce clones that in turn leads to additional deallocations.
 class BufferDeallocation : public BufferPlacementTransformationBase {
 public:
-  using AliasAllocationMapT = llvm::DenseMap<Value, AllocationOpInterface>;
+  using AliasAllocationMapT =
+      llvm::DenseMap<Value, bufferize::AllocationOpInterface>;
 
   BufferDeallocation(Operation *op)
       : BufferPlacementTransformationBase(op), dominators(op),
         postDominators(op) {}
 
   /// Checks if all allocation operations either provide an already existing
   /// deallocation operation or implement the AllocationOpInterface. In
   /// addition, this method initializes the internal alias to
   /// AllocationOpInterface mapping in order to get compatible
   /// AllocationOpInterface implementations for aliases.
   LogicalResult prepare() {
     for (const BufferPlacementAllocs::AllocEntry &entry : allocs) {
       // Get the defining allocation operation.
       Value alloc = std::get<0>(entry);
-      auto allocationInterface = alloc.getDefiningOp<AllocationOpInterface>();
+      auto allocationInterface =
+          alloc.getDefiningOp<bufferize::AllocationOpInterface>();
       // If there is no existing deallocation operation and no implementation of
       // the AllocationOpInterface, we cannot apply the BufferDeallocation pass.
       if (!std::get<1>(entry) && !allocationInterface) {
         return alloc.getDefiningOp()->emitError(
             "Allocation is not deallocated explicitly nor does the operation "
             "implement the AllocationOpInterface.");
       }
 
private:
   /// Maps aliases to their source allocation interfaces (inverse mapping).
   AliasAllocationMapT aliasToAllocations;
 };
 
 //===----------------------------------------------------------------------===//
 // BufferDeallocationPass
 //===----------------------------------------------------------------------===//
 
+template <typename T>
+struct DefaultAllocationInterface
+    : public bufferize::AllocationOpInterface::ExternalModel<
+          DefaultAllocationInterface<T>, T> {};
+
 /// The actual buffer deallocation pass that inserts and moves dealloc nodes
 /// into the right positions. Furthermore, it inserts additional clones if
 /// necessary. It uses the algorithm described at the top of the file.
 struct BufferDeallocationPass : BufferDeallocationBase<BufferDeallocationPass> {
+  void getDependentDialects(DialectRegistry &registry) const override {
+    registry.insert<memref::MemRefDialect>();
+    registry.addOpInterface<memref::AllocOp,
+                            DefaultAllocationInterface<memref::AllocOp>>();
+    registry.addOpInterface<memref::CloneOp,
+                            DefaultAllocationInterface<memref::CloneOp>>();
+  }
 
   void runOnFunction() override {
     // Ensure that there are supported loops only.
     FuncOp func = getFunction();
+
     Backedges backedges(func);
     if (backedges.size()) {
       func.emitError("Only structured control-flow loops are supported.");
       return signalPassFailure();
     }
 
     // Check that the control flow structures are supported.
     if (!validateSupportedControlFlow(func.getRegion()))

mlir/lib/Transforms/CMakeLists.txt

add_mlir_library(MLIRTransforms
 
   DEPENDS
   MLIRStandardOpsIncGen
   MLIRTransformsPassIncGen
 
   LINK_LIBS PUBLIC
   MLIRAffine
   MLIRAnalysis
+  MLIRAllocationOpInterface
   MLIRCopyOpInterface
   MLIRLoopLikeInterface
   MLIRMemRef
   MLIRSCF
   MLIRPass
   MLIRSupport
   MLIRTransformUtils
   MLIRVector
   )

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

cc_library(
     srcs = glob([
         "lib/Transforms/*.cpp",
         "lib/Transforms/*.h",
     ]),
     hdrs = glob(["include/mlir/Transforms/*.h"]),
     includes = ["include"],
     deps = [
         ":Affine",
+        ":AllocationOpInterface",
         ":Analysis",
         ":ArithmeticDialect",
         ":ControlFlowInterfaces",
         ":CopyOpInterface",
         ":IR",
         ":LinalgOps",
         ":LoopLikeInterface",
         ":MemRefDialect",
deps = [
         ":TensorDialect",
         ":Transforms",
         ":VectorOps",
         "//llvm:Support",
     ],
 )
 
 td_library(
+    name = "AllocationOpInterfaceTdFiles",
+    srcs = ["include/mlir/Dialect/Bufferize/IR/AllocationOpInterface.td"],
+    includes = ["include"],
+    deps = [":OpBaseTdFiles"],
+)
+
+gentbl_cc_library(
+    name = "AllocationOpInterfaceIncGen",
+    strip_include_prefix = "include",
+    tbl_outs = [
+        (
+            ["-gen-op-interface-decls"],
+            "include/mlir/Dialect/Bufferize/IR/AllocationOpInterface.h.inc",
+        ),
+        (
+            ["-gen-op-interface-defs"],
+            "include/mlir/Dialect/Bufferize/IR/AllocationOpInterface.cpp.inc",
+        ),
+    ],
+    tblgen = ":mlir-tblgen",
+    td_file = "include/mlir/Dialect/Bufferize/IR/AllocationOpInterface.td",
+    deps = [":AllocationOpInterfaceTdFiles"],
+)
+
+cc_library(
+    name = "AllocationOpInterface",
+    srcs = ["lib/Dialect/Bufferize/IR/AllocationOpInterface.cpp"],
+    hdrs = ["include/mlir/Dialect/Bufferize/IR/AllocationOpInterface.h"],
+    includes = ["include"],
+    deps = [
+        ":AllocationOpInterfaceIncGen",
+        ":IR",
+        ":MemRefDialect",
+    ],
+)
+
+td_library(
     name = "DLTIDialectTdFiles",
     srcs = [
         "include/mlir/Dialect/DLTI/DLTI.td",
         "include/mlir/Dialect/DLTI/DLTIBase.td",
     ],
     includes = ["include"],
     deps = [":OpBaseTdFiles"],
 )