[mlir][gpu][bufferization] Implement BufferDeallocationOpInterface for gpu.terminator (#66880)

This is necessary to support deallocation of IR with gpu.launch
operations because it does not implement the RegionBranchOpInterface.
Implementing the interface would require it to support regions with
unstructured control flow and produced arguments/results.

10 files changed, 149 insertions, 53 deletions

diff --git a/mlir/include/mlir/Dialect/Bufferization/IR/BufferDeallocationOpInterface.h b/mlir/include/mlir/Dialect/Bufferization/IR/BufferDeallocationOpInterface.h
index 7ac4592d..752a4a2 100644
--- a/mlir/include/mlir/Dialect/Bufferization/IR/BufferDeallocationOpInterface.h
+++ b/mlir/include/mlir/Dialect/Bufferization/IR/BufferDeallocationOpInterface.h
@@ -205,6 +205,23 @@ private:
   Liveness liveness;
 };
 
+namespace deallocation_impl {
+/// Insert a `bufferization.dealloc` operation right before `op` which has to be
+/// a terminator without any successors. Note that it is not required to have
+/// the ReturnLike trait attached. The MemRef values in the `operands` argument
+/// will be added to the list of retained values and their updated ownership
+/// values will be appended to the `updatedOperandOwnerships` list. `op` is not
+/// modified in any way. Returns failure if at least one of the MemRefs to
+/// deallocate does not have 'Unique' ownership (likely as a result of an
+/// incorrect implementation of the `process` or
+/// `materializeUniqueOwnershipForMemref` interface method) or the original
+/// `op`.
+FailureOr<Operation *>
+insertDeallocOpForReturnLike(DeallocationState &state, Operation *op,
+                             ValueRange operands,
+                             SmallVectorImpl<Value> &updatedOperandOwnerships);
+} // namespace deallocation_impl
+
 } // namespace bufferization
 } // namespace mlir
 
diff --git a/mlir/include/mlir/Dialect/GPU/Transforms/BufferDeallocationOpInterfaceImpl.h b/mlir/include/mlir/Dialect/GPU/Transforms/BufferDeallocationOpInterfaceImpl.h
new file mode 100644
index 0000000..16cf969
--- /dev/null
+++ b/mlir/include/mlir/Dialect/GPU/Transforms/BufferDeallocationOpInterfaceImpl.h
@@ -0,0 +1,22 @@
+//===- BufferDeallocationOpInterfaceImpl.h ----------------------*- 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_GPU_TRANSFORMS_BUFFERDEALLOCATIONOPINTERFACEIMPL_H
+#define MLIR_DIALECT_GPU_TRANSFORMS_BUFFERDEALLOCATIONOPINTERFACEIMPL_H
+
+namespace mlir {
+
+class DialectRegistry;
+
+namespace gpu {
+void registerBufferDeallocationOpInterfaceExternalModels(
+    DialectRegistry &registry);
+} // namespace gpu
+} // namespace mlir
+
+#endif // MLIR_DIALECT_GPU_TRANSFORMS_BUFFERDEALLOCATIONOPINTERFACEIMPL_H
diff --git a/mlir/include/mlir/InitAllDialects.h b/mlir/include/mlir/InitAllDialects.h
index 5b2b1ed..8a085d9 100644
--- a/mlir/include/mlir/InitAllDialects.h
+++ b/mlir/include/mlir/InitAllDialects.h
@@ -36,6 +36,7 @@
 #include "mlir/Dialect/EmitC/IR/EmitC.h"
 #include "mlir/Dialect/Func/IR/FuncOps.h"
 #include "mlir/Dialect/GPU/IR/GPUDialect.h"
+#include "mlir/Dialect/GPU/Transforms/BufferDeallocationOpInterfaceImpl.h"
 #include "mlir/Dialect/IRDL/IR/IRDL.h"
 #include "mlir/Dialect/Index/IR/IndexDialect.h"
 #include "mlir/Dialect/LLVMIR/LLVMDialect.h"
@@ -144,6 +145,7 @@ inline void registerAllDialects(DialectRegistry &registry) {
   builtin::registerCastOpInterfaceExternalModels(registry);
   cf::registerBufferizableOpInterfaceExternalModels(registry);
   cf::registerBufferDeallocationOpInterfaceExternalModels(registry);
+  gpu::registerBufferDeallocationOpInterfaceExternalModels(registry);
   linalg::registerBufferizableOpInterfaceExternalModels(registry);
   linalg::registerTilingInterfaceExternalModels(registry);
   linalg::registerValueBoundsOpInterfaceExternalModels(registry);
diff --git a/mlir/lib/Dialect/Bufferization/IR/BufferDeallocationOpInterface.cpp b/mlir/lib/Dialect/Bufferization/IR/BufferDeallocationOpInterface.cpp
index 407d75e..8d21446 100644
--- a/mlir/lib/Dialect/Bufferization/IR/BufferDeallocationOpInterface.cpp
+++ b/mlir/lib/Dialect/Bufferization/IR/BufferDeallocationOpInterface.cpp
@@ -272,3 +272,44 @@ bool ValueComparator::operator()(const Value &lhs, const Value &rhs) const {
   assert(lhsRegion && "this should only happen if lhs == rhs");
   return false;
 }
+
+//===----------------------------------------------------------------------===//
+// Implementation utilities
+//===----------------------------------------------------------------------===//
+
+FailureOr<Operation *> deallocation_impl::insertDeallocOpForReturnLike(
+    DeallocationState &state, Operation *op, ValueRange operands,
+    SmallVectorImpl<Value> &updatedOperandOwnerships) {
+  assert(op->hasTrait<OpTrait::IsTerminator>() && "must be a terminator");
+  assert(!op->hasSuccessors() && "must not have any successors");
+  // Collect the values to deallocate and retain and use them to create the
+  // dealloc operation.
+  OpBuilder builder(op);
+  Block *block = op->getBlock();
+  SmallVector<Value> memrefs, conditions, toRetain;
+  if (failed(state.getMemrefsAndConditionsToDeallocate(
+          builder, op->getLoc(), block, memrefs, conditions)))
+    return failure();
+
+  state.getMemrefsToRetain(block, /*toBlock=*/nullptr, operands, toRetain);
+  if (memrefs.empty() && toRetain.empty())
+    return op;
+
+  auto deallocOp = builder.create<bufferization::DeallocOp>(
+      op->getLoc(), memrefs, conditions, toRetain);
+
+  // We want to replace the current ownership of the retained values with the
+  // result values of the dealloc operation as they are always unique.
+  state.resetOwnerships(deallocOp.getRetained(), block);
+  for (auto [retained, ownership] :
+       llvm::zip(deallocOp.getRetained(), deallocOp.getUpdatedConditions()))
+    state.updateOwnership(retained, ownership, block);
+
+  unsigned numMemrefOperands = llvm::count_if(operands, isMemref);
+  auto newOperandOwnerships =
+      deallocOp.getUpdatedConditions().take_front(numMemrefOperands);
+  updatedOperandOwnerships.append(newOperandOwnerships.begin(),
+                                  newOperandOwnerships.end());
+
+  return op;
+}
diff --git a/mlir/lib/Dialect/Bufferization/Transforms/OwnershipBasedBufferDeallocation.cpp b/mlir/lib/Dialect/Bufferization/Transforms/OwnershipBasedBufferDeallocation.cpp
index 09d3083..94a26f3 100644
--- a/mlir/lib/Dialect/Bufferization/Transforms/OwnershipBasedBufferDeallocation.cpp
+++ b/mlir/lib/Dialect/Bufferization/Transforms/OwnershipBasedBufferDeallocation.cpp
@@ -47,10 +47,6 @@ static Value buildBoolValue(OpBuilder &builder, Location loc, bool value) {
 
 static bool isMemref(Value v) { return v.getType().isa<BaseMemRefType>(); }
 
-static bool isMemrefOperand(OpOperand &operand) {
-  return isMemref(operand.get());
-}
-
 //===----------------------------------------------------------------------===//
 // Backedges analysis
 //===----------------------------------------------------------------------===//
@@ -917,35 +913,16 @@ BufferDeallocation::handleInterface(RegionBranchTerminatorOpInterface op) {
   MutableOperandRange operands =
       op.getMutableSuccessorOperands(RegionBranchPoint::parent());
 
-  // Collect the values to deallocate and retain and use them to create the
-  // dealloc operation.
-  Block *block = op->getBlock();
-  SmallVector<Value> memrefs, conditions, toRetain;
-  if (failed(state.getMemrefsAndConditionsToDeallocate(
-          builder, op.getLoc(), block, memrefs, conditions)))
-    return failure();
-
-  state.getMemrefsToRetain(block, nullptr, OperandRange(operands), toRetain);
-  if (memrefs.empty() && toRetain.empty())
-    return op.getOperation();
-
-  auto deallocOp = builder.create<bufferization::DeallocOp>(
-      op.getLoc(), memrefs, conditions, toRetain);
-
-  // We want to replace the current ownership of the retained values with the
-  // result values of the dealloc operation as they are always unique.
-  state.resetOwnerships(deallocOp.getRetained(), block);
-  for (auto [retained, ownership] :
-       llvm::zip(deallocOp.getRetained(), deallocOp.getUpdatedConditions()))
-    state.updateOwnership(retained, ownership, block);
+  SmallVector<Value> updatedOwnerships;
+  auto result = deallocation_impl::insertDeallocOpForReturnLike(
+      state, op, OperandRange(operands), updatedOwnerships);
+  if (failed(result) || !*result)
+    return result;
 
   // Add an additional operand for every MemRef for the ownership indicator.
   if (!funcWithoutDynamicOwnership) {
-    unsigned numMemRefs = llvm::count_if(operands, isMemrefOperand);
     SmallVector<Value> newOperands{OperandRange(operands)};
-    auto ownershipValues =
-        deallocOp.getUpdatedConditions().take_front(numMemRefs);
-    newOperands.append(ownershipValues.begin(), ownershipValues.end());
+    newOperands.append(updatedOwnerships.begin(), updatedOwnerships.end());
     operands.assign(newOperands);
   }
 
diff --git a/mlir/lib/Dialect/GPU/CMakeLists.txt b/mlir/lib/Dialect/GPU/CMakeLists.txt
index 6244132..324d5c13 100644
--- a/mlir/lib/Dialect/GPU/CMakeLists.txt
+++ b/mlir/lib/Dialect/GPU/CMakeLists.txt
@@ -48,6 +48,7 @@ add_mlir_dialect_library(MLIRGPUDialect
 add_mlir_dialect_library(MLIRGPUTransforms
   Transforms/AllReduceLowering.cpp
   Transforms/AsyncRegionRewriter.cpp
+  Transforms/BufferDeallocationOpInterfaceImpl.cpp
   Transforms/DecomposeMemrefs.cpp
   Transforms/GlobalIdRewriter.cpp
   Transforms/KernelOutlining.cpp
@@ -79,6 +80,7 @@ add_mlir_dialect_library(MLIRGPUTransforms
   MLIRAffineUtils
   MLIRArithDialect
   MLIRAsyncDialect
+  MLIRBufferizationDialect
   MLIRBuiltinToLLVMIRTranslation
   MLIRDataLayoutInterfaces
   MLIRExecutionEngineUtils
diff --git a/mlir/lib/Dialect/GPU/Transforms/BufferDeallocationOpInterfaceImpl.cpp b/mlir/lib/Dialect/GPU/Transforms/BufferDeallocationOpInterfaceImpl.cpp
new file mode 100644
index 0000000..6ccc0a2
--- /dev/null
+++ b/mlir/lib/Dialect/GPU/Transforms/BufferDeallocationOpInterfaceImpl.cpp
@@ -0,0 +1,37 @@
+//===- BufferDeallocationOpInterfaceImpl.cpp ------------------------------===//
+//
+// 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/GPU/Transforms/BufferDeallocationOpInterfaceImpl.h"
+#include "mlir/Dialect/Bufferization/IR/BufferDeallocationOpInterface.h"
+#include "mlir/Dialect/Bufferization/IR/Bufferization.h"
+#include "mlir/Dialect/GPU/IR/GPUDialect.h"
+
+using namespace mlir;
+using namespace mlir::bufferization;
+
+namespace {
+///
+struct GPUTerminatorOpInterface
+    : public BufferDeallocationOpInterface::ExternalModel<
+          GPUTerminatorOpInterface, gpu::TerminatorOp> {
+  FailureOr<Operation *> process(Operation *op, DeallocationState &state,
+                                 const DeallocationOptions &options) const {
+    SmallVector<Value> updatedOperandOwnerships;
+    return deallocation_impl::insertDeallocOpForReturnLike(
+        state, op, {}, updatedOperandOwnerships);
+  }
+};
+
+} // namespace
+
+void mlir::gpu::registerBufferDeallocationOpInterfaceExternalModels(
+    DialectRegistry &registry) {
+  registry.addExtension(+[](MLIRContext *ctx, GPUDialect *dialect) {
+    gpu::TerminatorOp::attachInterface<GPUTerminatorOpInterface>(*ctx);
+  });
+}
diff --git a/mlir/lib/Dialect/SCF/Transforms/BufferDeallocationOpInterfaceImpl.cpp b/mlir/lib/Dialect/SCF/Transforms/BufferDeallocationOpInterfaceImpl.cpp
index 88cb3e9..4ded8ba55 100644
--- a/mlir/lib/Dialect/SCF/Transforms/BufferDeallocationOpInterfaceImpl.cpp
+++ b/mlir/lib/Dialect/SCF/Transforms/BufferDeallocationOpInterfaceImpl.cpp
@@ -47,33 +47,12 @@ struct InParallelOpInterface
   FailureOr<Operation *> process(Operation *op, DeallocationState &state,
                                  const DeallocationOptions &options) const {
     auto inParallelOp = cast<scf::InParallelOp>(op);
-    OpBuilder builder(op);
     if (!inParallelOp.getBody()->empty())
       return op->emitError("only supported when nested region is empty");
 
-    // Collect the values to deallocate and retain and use them to create the
-    // dealloc operation.
-    Block *block = op->getBlock();
-    SmallVector<Value> memrefs, conditions, toRetain;
-    if (failed(state.getMemrefsAndConditionsToDeallocate(
-            builder, op->getLoc(), block, memrefs, conditions)))
-      return failure();
-
-    state.getMemrefsToRetain(block, /*toBlock=*/nullptr, {}, toRetain);
-    if (memrefs.empty() && toRetain.empty())
-      return op;
-
-    auto deallocOp = builder.create<bufferization::DeallocOp>(
-        op->getLoc(), memrefs, conditions, toRetain);
-
-    // We want to replace the current ownership of the retained values with the
-    // result values of the dealloc operation as they are always unique.
-    state.resetOwnerships(deallocOp.getRetained(), block);
-    for (auto [retained, ownership] :
-         llvm::zip(deallocOp.getRetained(), deallocOp.getUpdatedConditions()))
-      state.updateOwnership(retained, ownership, block);
-
-    return op;
+    SmallVector<Value> updatedOperandOwnership;
+    return deallocation_impl::insertDeallocOpForReturnLike(
+        state, op, {}, updatedOperandOwnership);
   }
 };
 
diff --git a/mlir/test/Dialect/GPU/bufferization-buffer-deallocation.mlir b/mlir/test/Dialect/GPU/bufferization-buffer-deallocation.mlir
new file mode 100644
index 0000000..2534996
--- /dev/null
+++ b/mlir/test/Dialect/GPU/bufferization-buffer-deallocation.mlir
@@ -0,0 +1,18 @@
+// RUN: mlir-opt %s -buffer-deallocation-pipeline --allow-unregistered-dialect | FileCheck %s
+
+func.func @gpu_launch() {
+  %c1 = arith.constant 1 : index
+  gpu.launch blocks(%arg0, %arg1, %arg2) in (%arg6 = %c1, %arg7 = %c1, %arg8 = %c1)
+    threads(%arg3, %arg4, %arg5) in (%arg9 = %c1, %arg10 = %c1, %arg11 = %c1) {
+    %alloc = memref.alloc() : memref<2xf32>
+    "test.memref_user"(%alloc) : (memref<2xf32>) -> ()
+    gpu.terminator
+  }
+  return
+}
+
+// CHECK-LABEL: func @gpu_launch
+//       CHECK:   gpu.launch
+//       CHECK:     [[ALLOC:%.+]] = memref.alloc(
+//       CHECK:     memref.dealloc [[ALLOC]]
+//       CHECK:     gpu.terminator
diff --git a/utils/bazel/llvm-project-overlay/mlir/BUILD.bazel b/utils/bazel/llvm-project-overlay/mlir/BUILD.bazel
index 7ae9b61..3c167ab 100644
--- a/utils/bazel/llvm-project-overlay/mlir/BUILD.bazel
+++ b/utils/bazel/llvm-project-overlay/mlir/BUILD.bazel
@@ -4927,6 +4927,7 @@ cc_library(
         ":ArithDialect",
         ":AsmParser",
         ":AsyncDialect",
+        ":BufferizationDialect",
         ":ControlFlowDialect",
         ":DLTIDialect",
         ":DialectUtils",