[mlir][bufferization] Add deallocation option to remove existing dealloc operations, add option to specify the kind of alloc operations to consider

Author: maerhart

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

@@ -10,6 +10,8 @@
 #define MLIR_DIALECT_BUFFERIZATION_IR_BUFFERDEALLOCATIONOPINTERFACE_H_
 
 #include "mlir/Analysis/Liveness.h"
+#include "mlir/Dialect/Bufferization/IR/Bufferization.h"
+#include "mlir/Dialect/MemRef/IR/MemRef.h"
 #include "mlir/IR/Operation.h"
 #include "mlir/IR/SymbolTable.h"
 #include "mlir/Support/LLVM.h"
@@ -92,6 +94,9 @@ class Ownership {
 
 /// Options for BufferDeallocationOpInterface-based buffer deallocation.
 struct DeallocationOptions {
+  using DetectionFn = std::function<bool(Operation *)>;
+  using ReplaceDeallocFn = std::function<FailureOr<ValueRange>(Operation *)>;
+
   // A pass option indicating whether private functions should be modified to
   // pass the ownership of MemRef values instead of adhering to the function
   // boundary ABI.
@@ -106,6 +111,48 @@ struct DeallocationOptions {
   /// to, an error will already be emitted at compile time. This cannot be
   /// changed with this option.
   bool verifyFunctionBoundaryABI = true;
+
+  /// Given an allocation side-effect on the passed operation, determine whether
+  /// this allocation operation is of relevance (i.e., should assign ownership
+  /// to the allocated value). If it is determined to not be relevant,
+  /// ownership will be set to 'false', i.e., it will be leaked. This is useful
+  /// to support deallocation of multiple different kinds of allocation ops.
+  DetectionFn isRelevantAllocOp = [](Operation *op) {
+    return isa<memref::MemRefDialect, bufferization::BufferizationDialect>(
+        op->getDialect());
+  };
+
+  /// Given a free side-effect on the passed operation, determine whether this
+  /// deallocation operation is of relevance (i.e., should be removed if the
+  /// `removeExistingDeallocations` option is enabled or otherwise an error
+  /// should be emitted because existing deallocation operations are not
+  /// supported without that flag). If it is determined to not be relevant,
+  /// the operation will be ignored. This is useful to support deallocation of
+  /// multiple different kinds of allocation ops where deallocations for some of
+  /// them are already present in the IR.
+  DetectionFn isRelevantDeallocOp = [](Operation *op) {
+    return isa<memref::MemRefDialect, bufferization::BufferizationDialect>(
+        op->getDialect());
+  };
+
+  /// When enabled, remove deallocation operations determined to be relevant
+  /// according to `isRelevantDeallocOp`. If the operation has result values,
+  /// `getDeallocReplacement` will be called to determine the SSA values that
+  /// should be used as replacements.
+  bool removeExistingDeallocations = false;
+
+  /// Provides SSA values for deallocation operations when
+  /// `removeExistingDeallocations` is enabled. May return a failure when the
+  /// given deallocation operation is not supported (e.g., because no
+  /// replacement for a result value can be determined). A failure will directly
+  /// lead to a failure emitted by the deallocation pass.
+  ReplaceDeallocFn getDeallocReplacement =
+      [](Operation *op) -> FailureOr<ValueRange> {
+    if (isa<memref::DeallocOp>(op))
+      return ValueRange{};
+    // ReallocOp has to be expanded before running the dealloc pass.
+    return failure();
+  };
 };
 
 /// This class collects all the state that we need to perform the buffer

mlir/include/mlir/Dialect/Bufferization/Pipelines/Passes.h

@@ -40,6 +40,12 @@ struct BufferDeallocationPipelineOptions
           "statically that the ABI is not adhered to, an error will already be "
           "emitted at compile time. This cannot be changed with this option."),
       llvm::cl::init(true)};
+  PassOptions::Option<bool> removeExistingDeallocations{
+      *this, "remove-existing-deallocations",
+      llvm::cl::desc("Removes all pre-existing memref.dealloc operations and "
+                     "insert all deallocations according to the buffer "
+                     "deallocation rules."),
+      llvm::cl::init(false)};
 
   /// Convert this BufferDeallocationPipelineOptions struct to a
   /// DeallocationOptions struct to be passed to the

mlir/include/mlir/Dialect/Bufferization/Transforms/Passes.td

@@ -233,6 +233,11 @@ def OwnershipBasedBufferDeallocation : Pass<
            "If it can be determined statically that the ABI is not adhered "
            "to, an error will already be emitted at compile time. This cannot "
            "be changed with this option.">,
+    Option<"removeExistingDeallocations", "remove-existing-deallocations",
+           "bool", /*default=*/"false",
+           "Remove already existing MemRef deallocation operations and let the "
+           "deallocation pass insert the deallocation operations according to "
+           "its rules.">,
   ];
   let constructor = "mlir::bufferization::createOwnershipBasedBufferDeallocationPass()";
 

mlir/lib/Dialect/Bufferization/Pipelines/BufferizationPipelines.cpp

@@ -27,6 +27,7 @@ BufferDeallocationPipelineOptions::asDeallocationOptions() const {
   DeallocationOptions opts;
   opts.privateFuncDynamicOwnership = privateFunctionDynamicOwnership.getValue();
   opts.verifyFunctionBoundaryABI = verifyFunctionBoundaryABI.getValue();
+  opts.removeExistingDeallocations = removeExistingDeallocations.getValue();
   return opts;
 }
 

mlir/lib/Dialect/Bufferization/Transforms/OwnershipBasedBufferDeallocation.cpp

@@ -816,15 +816,28 @@ FailureOr<Operation *>
 BufferDeallocation::handleInterface(MemoryEffectOpInterface op) {
   auto *block = op->getBlock();
 
-  for (auto operand : llvm::make_filter_range(op->getOperands(), isMemref))
-    if (op.getEffectOnValue<MemoryEffects::Free>(operand).has_value())
+  for (auto operand : llvm::make_filter_range(op->getOperands(), isMemref)) {
+    if (op.getEffectOnValue<MemoryEffects::Free>(operand).has_value() &&
+        options.isRelevantDeallocOp(op)) {
+      if (auto repl = options.getDeallocReplacement(op);
+          succeeded(repl) && options.removeExistingDeallocations) {
+        op->replaceAllUsesWith(repl.value());
+        op.erase();
+        return FailureOr<Operation *>(nullptr);
+      }
+
       return op->emitError(
           "memory free side-effect on MemRef value not supported!");
+    }
+  }
 
   OpBuilder builder = OpBuilder::atBlockBegin(block);
   for (auto res : llvm::make_filter_range(op->getResults(), isMemref)) {
     auto allocEffect = op.getEffectOnValue<MemoryEffects::Allocate>(res);
     if (allocEffect.has_value()) {
+      // Assuming that an alloc effect is interpreted as MUST and not MAY.
+      state.resetOwnerships(res, block);
+
       if (isa<SideEffects::AutomaticAllocationScopeResource>(
               allocEffect->getResource())) {
         // Make sure that the ownership of auto-managed allocations is set to
@@ -839,8 +852,15 @@ BufferDeallocation::handleInterface(MemoryEffectOpInterface op) {
         continue;
       }
 
-      state.updateOwnership(res, buildBoolValue(builder, op.getLoc(), true));
-      state.addMemrefToDeallocate(res, block);
+      if (options.isRelevantAllocOp(op)) {
+        state.updateOwnership(res, buildBoolValue(builder, op.getLoc(), true));
+        state.addMemrefToDeallocate(res, block);
+        continue;
+      }
+
+      // Alloc operations from other dialects are expected to have matching
+      // deallocation operations inserted by another pass.
+      state.updateOwnership(res, buildBoolValue(builder, op.getLoc(), false));
     }
   }
 
@@ -943,16 +963,18 @@ struct OwnershipBasedBufferDeallocationPass
     : public bufferization::impl::OwnershipBasedBufferDeallocationBase<
           OwnershipBasedBufferDeallocationPass> {
   OwnershipBasedBufferDeallocationPass() = default;
-  OwnershipBasedBufferDeallocationPass(const DeallocationOptions &options)
-      : OwnershipBasedBufferDeallocationPass() {
+  OwnershipBasedBufferDeallocationPass(const DeallocationOptions &options) {
     privateFuncDynamicOwnership.setValue(options.privateFuncDynamicOwnership);
     verifyFunctionBoundaryABI.setValue(options.verifyFunctionBoundaryABI);
+    removeExistingDeallocations.setValue(options.removeExistingDeallocations);
   }
   void runOnOperation() override {
     DeallocationOptions options;
     options.privateFuncDynamicOwnership =
         privateFuncDynamicOwnership.getValue();
     options.verifyFunctionBoundaryABI = verifyFunctionBoundaryABI.getValue();
+    options.removeExistingDeallocations =
+        removeExistingDeallocations.getValue();
 
     auto status = getOperation()->walk([&](func::FuncOp func) {
       if (func.isExternal())

mlir/test/Dialect/Bufferization/Transforms/OwnershipBasedBufferDeallocation/dealloc-mixed-allocations.mlir

@@ -0,0 +1,33 @@
+// RUN: mlir-opt %s --test-ownership-based-buffer-deallocation -split-input-file | FileCheck %s
+
+func.func @mixed_allocations(%cond: i1) -> (memref<f32>, !gpu.async.token) {
+  %a1 = memref.alloc() : memref<f32>
+  %a2 = gpu.alloc() : memref<f32>
+  %0 = arith.select %cond, %a1, %a2 : memref<f32>
+  %token = gpu.dealloc async [] %a2 : memref<f32>
+  memref.dealloc %a1 : memref<f32>
+  return %0, %token : memref<f32>, !gpu.async.token
+}
+
+// CHECK: [[A1:%.+]] = memref.alloc(
+// CHECK: [[A2:%.+]] = gpu.alloc
+// CHECK: [[SELECT:%.+]] = arith.select {{.*}}, [[A1]], [[A2]]
+// CHECK: [[TOKEN:%.+]] = gpu.wait async
+// CHECK: [[A1_BASE:%.+]],{{.*}} = memref.extract_strided_metadata [[A1]]
+// CHECK: [[A1_PTR:%.+]] = memref.extract_aligned_pointer_as_index [[A1_BASE]]
+// CHECK: [[SELECT_PTR:%.+]] = memref.extract_aligned_pointer_as_index [[SELECT]]
+// CHECK: [[ALIAS0:%.+]] = arith.cmpi ne, [[A1_PTR]], [[SELECT_PTR]]
+// CHECK: [[COND0:%.+]] = arith.andi [[ALIAS0]], %true
+// CHECK: scf.if [[COND0]] {
+// CHECK:   memref.dealloc [[A1_BASE]]
+// CHECK: }
+// CHECK: [[A2_BASE:%.+]],{{.*}} = memref.extract_strided_metadata [[A2]]
+// CHECK: [[A2_PTR:%.+]] = memref.extract_aligned_pointer_as_index [[A2_BASE]]
+// CHECK: [[SELECT_PTR:%.+]] = memref.extract_aligned_pointer_as_index [[SELECT]]
+// CHECK: [[ALIAS1:%.+]] = arith.cmpi ne, [[A2_PTR]], [[SELECT_PTR]]
+// CHECK: [[COND1:%.+]] = arith.andi [[ALIAS1]], %true
+// CHECK: scf.if [[COND1]] {
+// TODO: add pass option to lower-deallocation to insert gpu.dealloc here
+// CHECK:   memref.dealloc [[A2_BASE]]
+// CHECK: }
+// CHECK: return [[SELECT]], [[TOKEN]]

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

@@ -5,8 +5,12 @@ add_mlir_library(MLIRBufferizationTestPasses
   EXCLUDE_FROM_LIBMLIR
 
   LINK_LIBS PUBLIC
+  MLIRArithDialect
   MLIRBufferizationDialect
   MLIRBufferizationTransforms
+  MLIRFuncDialect
+  MLIRGPUDialect
+  MLIRSCFDialect
   MLIRIR
   MLIRPass
 )

mlir/test/lib/Dialect/Bufferization/TestOwnershipBasedBufferDeallocation.cpp

@@ -0,0 +1,128 @@
+//===- TestOwnershipBasedBufferDeallocation.cpp -----------------*- 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
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/Dialect/Bufferization/IR/BufferDeallocationOpInterface.h"
+#include "mlir/Dialect/Bufferization/IR/Bufferization.h"
+#include "mlir/Dialect/Bufferization/Transforms/Passes.h"
+#include "mlir/Dialect/Func/IR/FuncOps.h"
+#include "mlir/Dialect/GPU/IR/GPUDialect.h"
+#include "mlir/Dialect/SCF/IR/SCF.h"
+#include "mlir/Pass/Pass.h"
+#include "mlir/Transforms/DialectConversion.h"
+
+using namespace mlir;
+
+namespace {
+/// This pass runs the ownership based deallocation pass once for `memref.alloc`
+/// operations, then lowers the `bufferization.dealloc` operations, and
+/// afterwards runs the deallocation pass again for `gpu.alloc` operations and
+/// lowers the inserted `bufferization.dealloc` operations again to the
+/// corresponding deallocation operations.
+struct TestOwnershipBasedBufferDeallocationPass
+    : public PassWrapper<TestOwnershipBasedBufferDeallocationPass,
+                         OperationPass<ModuleOp>> {
+  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(
+      TestOwnershipBasedBufferDeallocationPass)
+
+  TestOwnershipBasedBufferDeallocationPass() = default;
+  TestOwnershipBasedBufferDeallocationPass(
+      const TestOwnershipBasedBufferDeallocationPass &pass)
+      : TestOwnershipBasedBufferDeallocationPass() {}
+
+  void getDependentDialects(DialectRegistry &registry) const override {
+    registry.insert<bufferization::BufferizationDialect, memref::MemRefDialect,
+                    scf::SCFDialect, func::FuncDialect, arith::ArithDialect>();
+  }
+  StringRef getArgument() const final {
+    return "test-ownership-based-buffer-deallocation";
+  }
+  StringRef getDescription() const final {
+    return "Module pass to test the Ownership-based Buffer Deallocation pass";
+  }
+
+  void runOnOperation() override {
+    ModuleOp module = getOperation();
+
+    // Build the library function for the lowering of `bufferization.dealloc`.
+    OpBuilder builder = OpBuilder::atBlockBegin(module.getBody());
+    SymbolTable symbolTable(module);
+    func::FuncOp helper = bufferization::buildDeallocationLibraryFunction(
+        builder, module.getLoc(), symbolTable);
+
+    RewritePatternSet patterns(module->getContext());
+    bufferization::populateBufferizationDeallocLoweringPattern(patterns,
+                                                               helper);
+    FrozenRewritePatternSet frozenPatterns(std::move(patterns));
+
+    WalkResult result = getOperation()->walk([&](FunctionOpInterface funcOp) {
+      // Deallocate the `memref.alloc` operations.
+      bufferization::DeallocationOptions options;
+      options.removeExistingDeallocations = true;
+      if (failed(
+              bufferization::deallocateBuffersOwnershipBased(funcOp, options)))
+        return WalkResult::interrupt();
+
+      // Lower the inserted `bufferization.dealloc` operations.
+      ConversionTarget target(getContext());
+      target.addLegalDialect<memref::MemRefDialect, arith::ArithDialect,
+                             scf::SCFDialect, func::FuncDialect>();
+      target.addIllegalOp<bufferization::DeallocOp>();
+
+      if (failed(applyPartialConversion(funcOp, target, frozenPatterns)))
+        return WalkResult::interrupt();
+
+      // Deallocate the `gpu.alloc` operations.
+      options.isRelevantAllocOp = [](Operation *op) {
+        return isa<gpu::GPUDialect>(op->getDialect());
+      };
+      options.isRelevantDeallocOp = [](Operation *op) {
+        return isa<gpu::GPUDialect>(op->getDialect());
+      };
+      options.getDeallocReplacement =
+          [](Operation *op) -> FailureOr<ValueRange> {
+        if (auto gpuDealloc = dyn_cast<gpu::DeallocOp>(op)) {
+          if (gpuDealloc.getAsyncToken()) {
+            OpBuilder builder(op);
+            ValueRange token =
+                builder
+                    .create<gpu::WaitOp>(
+                        op->getLoc(),
+                        gpu::AsyncTokenType::get(builder.getContext()),
+                        ValueRange{})
+                    .getResults();
+            return token;
+          }
+          return ValueRange{};
+        }
+        return failure();
+      };
+      if (failed(
+              bufferization::deallocateBuffersOwnershipBased(funcOp, options)))
+        return WalkResult::interrupt();
+
+      // Lower the `bufferization.dealloc` operations inserted in the second
+      // deallocation run.
+      // TODO: they are currently also lowered to memref.dealloc, we need to
+      // add pass options to the lowering pass that allow us to select the
+      // dealloc operation to be inserted.
+      if (failed(applyPartialConversion(funcOp, target, frozenPatterns)))
+        return WalkResult::interrupt();
+
+      return WalkResult::advance();
+    });
+    if (result.wasInterrupted())
+      signalPassFailure();
+  }
+};
+} // namespace
+
+namespace mlir::test {
+void registerTestOwnershipBasedBufferDeallocationPass() {
+  PassRegistration<TestOwnershipBasedBufferDeallocationPass>();
+}
+} // namespace mlir::test

mlir/tools/mlir-opt/mlir-opt.cpp

@@ -121,6 +121,7 @@ void registerTestMemRefStrideCalculation();
 void registerTestNextAccessPass();
 void registerTestOneToNTypeConversionPass();
 void registerTestOpaqueLoc();
+void registerTestOwnershipBasedBufferDeallocationPass();
 void registerTestPadFusion();
 void registerTestPDLByteCodePass();
 void registerTestPDLLPasses();
@@ -241,6 +242,7 @@ void registerTestPasses() {
   mlir::test::registerTestNextAccessPass();
   mlir::test::registerTestOneToNTypeConversionPass();
   mlir::test::registerTestOpaqueLoc();
+  mlir::test::registerTestOwnershipBasedBufferDeallocationPass();
   mlir::test::registerTestPadFusion();
   mlir::test::registerTestPDLByteCodePass();
   mlir::test::registerTestPDLLPasses();

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

@@ -3,9 +3,9 @@
 # SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 
 load("@bazel_skylib//rules:expand_template.bzl", "expand_template")
+load("//llvm:lit_test.bzl", "package_path")
 load("//mlir:build_defs.bzl", "if_cuda_available")
 load("//mlir:tblgen.bzl", "gentbl_cc_library", "td_library")
-load("//llvm:lit_test.bzl", "package_path")
 
 package(
     default_visibility = ["//visibility:public"],
@@ -843,10 +843,15 @@ cc_library(
     defines = ["MLIR_CUDA_CONVERSIONS_ENABLED"],
     includes = ["lib/Dialect/Test"],
     deps = [
+        "//mlir:ArithDialect",
         "//mlir:BufferizationDialect",
         "//mlir:BufferizationTransforms",
+        "//mlir:FuncDialect",
+        "//mlir:GPUDialect",
         "//mlir:IR",
         "//mlir:Pass",
+        "//mlir:SCFDialect",
+        "//mlir:Transforms",
     ],
 )