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

This reverts and fixes commit de18b7dee6.

mlir/include/mlir/Dialect/Bufferization/CMakeLists.txt

@@ -0,0 +1 @@
+add_subdirectory(IR)

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

@@ -0,0 +1,21 @@
+//===- 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_BUFFERIZATION_IR_ALLOCATIONOPINTERFACE_H_
+#define MLIR_DIALECT_BUFFERIZATION_IR_ALLOCATIONOPINTERFACE_H_
+
+#include "mlir/Dialect/MemRef/IR/MemRef.h"
+#include "mlir/IR/Builders.h"
+
+#include "mlir/Dialect/Bufferization/IR/AllocationOpInterface.h.inc"
+
+#endif // MLIR_DIALECT_BUFFERIZATION_IR_ALLOCATIONOPINTERFACE_H_

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

@@ -0,0 +1,61 @@
+//===-- 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::bufferization";
+
+  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 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&":$builder, "::mlir::Value":$alloc), [{}],
+      /*defaultImplementation=*/[{
+          return builder.create<memref::CloneOp>(alloc.getLoc(), alloc)
+            .getResult();
+      }]
+    >
+  ];
+}
+
+#endif  // ALLOCATION_OP_INTERFACE

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

@@ -0,0 +1 @@
+add_mlir_interface(AllocationOpInterface)

mlir/include/mlir/Dialect/CMakeLists.txt

@@ -4,6 +4,7 @@ add_subdirectory(Async)
 add_subdirectory(ArmNeon)
 add_subdirectory(ArmSVE)
 add_subdirectory(AMX)
+add_subdirectory(Bufferization)
 add_subdirectory(Complex)
 add_subdirectory(DLTI)
 add_subdirectory(EmitC)

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

@@ -120,10 +120,7 @@ def AssumeAlignmentOp : MemRef_Op<"assume_alignment"> {
 // 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
@@ -418,8 +415,6 @@ def MemRef_CastOp : MemRef_Op<"cast", [
 def CloneOp : MemRef_Op<"clone", [
     CopyOpInterface,
     DeclareOpInterfaceMethods<MemoryEffectsOpInterface>,
-    DeclareOpInterfaceMethods<AllocationOpInterface,
-      ["buildDealloc", "buildClone"]>
   ]> {
   let builders = [
     OpBuilder<(ins "Value":$value), [{

mlir/include/mlir/Interfaces/SideEffectInterfaces.td

@@ -16,45 +16,6 @@
 
 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
 //===----------------------------------------------------------------------===//

mlir/lib/Dialect/Bufferization/CMakeLists.txt

@@ -0,0 +1 @@
+add_subdirectory(IR)

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

@@ -0,0 +1,10 @@
+//===- 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/Bufferization/IR/AllocationOpInterface.h"
+#include "mlir/Dialect/Bufferization/IR/AllocationOpInterface.cpp.inc"

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

@@ -0,0 +1,12 @@
+add_mlir_library(MLIRAllocationOpInterface
+  AllocationOpInterface.cpp
+
+  ADDITIONAL_HEADER_DIRS
+  ${MLIR_MAIN_INCLUDE_DIR}/mlir/Dialect/Bufferization/IR
+
+  DEPENDS
+  MLIRAllocationOpInterfaceIncGen
+
+  LINK_LIBS PUBLIC
+  MLIRIR
+  )

mlir/lib/Dialect/CMakeLists.txt

@@ -4,6 +4,7 @@ add_subdirectory(ArmNeon)
 add_subdirectory(ArmSVE)
 add_subdirectory(Async)
 add_subdirectory(AMX)
+add_subdirectory(Bufferization)
 add_subdirectory(Complex)
 add_subdirectory(DLTI)
 add_subdirectory(EmitC)

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

@@ -196,15 +196,6 @@ struct SimplifyDeadAlloc : public OpRewritePattern<T> {
 };
 } // 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);
@@ -653,15 +644,6 @@ 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
 //===----------------------------------------------------------------------===//

mlir/lib/Transforms/BufferDeallocation.cpp

@@ -51,6 +51,8 @@
 //===----------------------------------------------------------------------===//
 
 #include "PassDetail.h"
+
+#include "mlir/Dialect/Bufferization/IR/AllocationOpInterface.h"
 #include "mlir/Dialect/MemRef/IR/MemRef.h"
 #include "mlir/Dialect/StandardOps/IR/Ops.h"
 #include "mlir/IR/Operation.h"
@@ -193,7 +195,8 @@ private:
 /// 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, bufferization::AllocationOpInterface>;
 
   BufferDeallocation(Operation *op)
       : BufferPlacementTransformationBase(op), dominators(op),
@@ -208,7 +211,8 @@ public:
     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<bufferization::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) {
@@ -614,10 +618,27 @@ private:
 // BufferDeallocationPass
 //===----------------------------------------------------------------------===//
 
+template <typename T>
+struct DefaultAllocationInterface
+    : public bufferization::AllocationOpInterface::ExternalModel<
+          DefaultAllocationInterface<T>, T> {
+  static Optional<Operation *> buildDealloc(OpBuilder &builder, Value alloc) {
+    return builder.create<memref::DeallocOp>(alloc.getLoc(), alloc)
+        .getOperation();
+  }
+};
+
 /// 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.

mlir/lib/Transforms/CMakeLists.txt

@@ -32,6 +32,7 @@ add_mlir_library(MLIRTransforms
   LINK_LIBS PUBLIC
   MLIRAffine
   MLIRAnalysis
+  MLIRAllocationOpInterface
   MLIRCopyOpInterface
   MLIRLoopLikeInterface
   MLIRMemRef

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

@@ -4432,6 +4432,7 @@ cc_library(
     includes = ["include"],
     deps = [
         ":Affine",
+        ":AllocationOpInterface",
         ":Analysis",
         ":ArithmeticDialect",
         ":ControlFlowInterfaces",
@@ -7489,6 +7490,43 @@ cc_library(
     ],
 )
 
+td_library(
+    name = "AllocationOpInterfaceTdFiles",
+    srcs = ["include/mlir/Dialect/Bufferization/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/Bufferization/IR/AllocationOpInterface.h.inc",
+        ),
+        (
+            ["-gen-op-interface-defs"],
+            "include/mlir/Dialect/Bufferization/IR/AllocationOpInterface.cpp.inc",
+        ),
+    ],
+    tblgen = ":mlir-tblgen",
+    td_file = "include/mlir/Dialect/Bufferization/IR/AllocationOpInterface.td",
+    deps = [":AllocationOpInterfaceTdFiles"],
+)
+
+cc_library(
+    name = "AllocationOpInterface",
+    srcs = ["lib/Dialect/Bufferization/IR/AllocationOpInterface.cpp"],
+    hdrs = ["include/mlir/Dialect/Bufferization/IR/AllocationOpInterface.h"],
+    includes = ["include"],
+    deps = [
+        ":AllocationOpInterfaceIncGen",
+        ":IR",
+        ":MemRefDialect",
+    ],
+)
+
 td_library(
     name = "DLTIDialectTdFiles",
     srcs = [