[Mlir-commits] [mlir] [mlir][bufferization] Add deallocation option to remove existing dealloc operations, add option to specify the kind of alloc operations to consider (PR #67556)

[Matthias Springer]

================
@@ -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;
----------------
matthias-springer wrote:

Add comment that by default the deallocation options are configured to look for memref.alloc/memref.dealloc only.

https://github.com/llvm/llvm-project/pull/67556