[llvm] [mlir][bufferization] Define a pipeline for buffer deallocation (PR #66352)
Martin Erhart via llvm-commits
llvm-commits at lists.llvm.org
Thu Sep 14 07:22:50 PDT 2023
https://github.com/maerhart updated https://github.com/llvm/llvm-project/pull/66352:
>From f5c014d8cc6729e6dc5ba733e99212f9409563ea Mon Sep 17 00:00:00 2001
From: Martin Erhart <merhart at google.com>
Date: Tue, 12 Sep 2023 15:23:22 +0000
Subject: [PATCH] [mlir][bufferization] Define a pipeline for buffer
deallocation
Since buffer deallocation requires a few passes to be run in a somewhat fixed
sequence, it makes sense to have a pipeline for convenience (and to reduce the
number of transform ops to represent default deallocation).
Was already reviewd in Differential Revision: https://reviews.llvm.org/D159432
---
.../Dialect/Bufferization/Pipelines/Passes.h | 50 +++++++++++++++++++
.../Dialect/Bufferization/Transforms/Passes.h | 3 +-
mlir/include/mlir/InitAllPasses.h | 2 +
mlir/lib/Dialect/Bufferization/CMakeLists.txt | 1 +
.../Pipelines/BufferizationPipelines.cpp | 46 +++++++++++++++++
.../Bufferization/Pipelines/CMakeLists.txt | 13 +++++
.../OwnershipBasedBufferDeallocation.cpp | 11 +++-
.../dealloc-branchop-interface.mlir | 2 +
.../dealloc-callop-interface.mlir | 3 ++
.../dealloc-existing-deallocs.mlir | 2 +
.../dealloc-function-boundaries.mlir | 3 ++
.../dealloc-memoryeffect-interface.mlir | 2 +
.../dealloc-region-branchop-interface.mlir | 2 +
.../dealloc-subviews.mlir | 2 +
.../llvm-project-overlay/mlir/BUILD.bazel | 16 ++++++
15 files changed, 155 insertions(+), 3 deletions(-)
create mode 100644 mlir/include/mlir/Dialect/Bufferization/Pipelines/Passes.h
create mode 100644 mlir/lib/Dialect/Bufferization/Pipelines/BufferizationPipelines.cpp
create mode 100644 mlir/lib/Dialect/Bufferization/Pipelines/CMakeLists.txt
diff --git a/mlir/include/mlir/Dialect/Bufferization/Pipelines/Passes.h b/mlir/include/mlir/Dialect/Bufferization/Pipelines/Passes.h
new file mode 100644
index 000000000000000..7acacb763cd2c18
--- /dev/null
+++ b/mlir/include/mlir/Dialect/Bufferization/Pipelines/Passes.h
@@ -0,0 +1,50 @@
+//===- Passes.h - Bufferization pipeline entry points -----------*- 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
+//
+//===----------------------------------------------------------------------===//
+//
+// This header file defines prototypes of all bufferization pipelines.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef MLIR_DIALECT_BUFFERIZATION_PIPELINES_PASSES_H
+#define MLIR_DIALECT_BUFFERIZATION_PIPELINES_PASSES_H
+
+#include "mlir/Pass/PassOptions.h"
+
+namespace mlir {
+namespace bufferization {
+
+/// Options for the buffer deallocation pipeline.
+struct BufferDeallocationPipelineOptions
+ : public PassPipelineOptions<BufferDeallocationPipelineOptions> {
+ PassOptions::Option<bool> privateFunctionDynamicOwnership{
+ *this, "private-function-dynamic-ownership",
+ llvm::cl::desc(
+ "Allows to add additional arguments to private functions to "
+ "dynamically pass ownership of memrefs to callees. This can enable "
+ "earlier deallocations."),
+ llvm::cl::init(false)};
+};
+
+//===----------------------------------------------------------------------===//
+// Building and Registering.
+//===----------------------------------------------------------------------===//
+
+/// Adds the buffer deallocation pipeline to the `OpPassManager`. This
+/// is the standard pipeline for deallocating the MemRefs introduced by the
+/// One-Shot bufferization pass.
+void buildBufferDeallocationPipeline(
+ OpPassManager &pm, const BufferDeallocationPipelineOptions &options);
+
+/// Registers all pipelines for the `bufferization` dialect. Currently,
+/// this includes only the "buffer-deallocation-pipeline".
+void registerBufferizationPipelines();
+
+} // namespace bufferization
+} // namespace mlir
+
+#endif // MLIR_DIALECT_BUFFERIZATION_PIPELINES_PASSES_H
diff --git a/mlir/include/mlir/Dialect/Bufferization/Transforms/Passes.h b/mlir/include/mlir/Dialect/Bufferization/Transforms/Passes.h
index 23eed02a15d4801..92520eb13da6875 100644
--- a/mlir/include/mlir/Dialect/Bufferization/Transforms/Passes.h
+++ b/mlir/include/mlir/Dialect/Bufferization/Transforms/Passes.h
@@ -30,7 +30,8 @@ std::unique_ptr<Pass> createBufferDeallocationPass();
/// Creates an instance of the OwnershipBasedBufferDeallocation pass to free all
/// allocated buffers.
-std::unique_ptr<Pass> createOwnershipBasedBufferDeallocationPass();
+std::unique_ptr<Pass> createOwnershipBasedBufferDeallocationPass(
+ bool privateFuncDynamicOwnership = false);
/// Creates a pass that optimizes `bufferization.dealloc` operations. For
/// example, it reduces the number of alias checks needed at runtime using
diff --git a/mlir/include/mlir/InitAllPasses.h b/mlir/include/mlir/InitAllPasses.h
index f7271737c66d1cb..8a45da7d1b982f1 100644
--- a/mlir/include/mlir/InitAllPasses.h
+++ b/mlir/include/mlir/InitAllPasses.h
@@ -20,6 +20,7 @@
#include "mlir/Dialect/Arith/Transforms/Passes.h"
#include "mlir/Dialect/ArmSME/Transforms/Passes.h"
#include "mlir/Dialect/Async/Passes.h"
+#include "mlir/Dialect/Bufferization/Pipelines/Passes.h"
#include "mlir/Dialect/Bufferization/Transforms/Passes.h"
#include "mlir/Dialect/Func/Transforms/Passes.h"
#include "mlir/Dialect/GPU/Transforms/Passes.h"
@@ -81,6 +82,7 @@ inline void registerAllPasses() {
arm_sme::registerArmSMEPasses();
// Dialect pipelines
+ bufferization::registerBufferizationPipelines();
sparse_tensor::registerSparseTensorPipelines();
}
diff --git a/mlir/lib/Dialect/Bufferization/CMakeLists.txt b/mlir/lib/Dialect/Bufferization/CMakeLists.txt
index 660deb21479d297..215ec562c9818c8 100644
--- a/mlir/lib/Dialect/Bufferization/CMakeLists.txt
+++ b/mlir/lib/Dialect/Bufferization/CMakeLists.txt
@@ -1,3 +1,4 @@
add_subdirectory(IR)
+add_subdirectory(Pipelines)
add_subdirectory(TransformOps)
add_subdirectory(Transforms)
diff --git a/mlir/lib/Dialect/Bufferization/Pipelines/BufferizationPipelines.cpp b/mlir/lib/Dialect/Bufferization/Pipelines/BufferizationPipelines.cpp
new file mode 100644
index 000000000000000..b2a60feb9a7f011
--- /dev/null
+++ b/mlir/lib/Dialect/Bufferization/Pipelines/BufferizationPipelines.cpp
@@ -0,0 +1,46 @@
+//===- BufferizationPipelines.cpp - Pipelines for bufferization -----------===//
+//
+// 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/Pipelines/Passes.h"
+
+#include "mlir/Dialect/Bufferization/Transforms/Passes.h"
+#include "mlir/Dialect/Func/IR/FuncOps.h"
+#include "mlir/Dialect/MemRef/Transforms/Passes.h"
+#include "mlir/Pass/PassManager.h"
+#include "mlir/Transforms/Passes.h"
+
+//===----------------------------------------------------------------------===//
+// Pipeline implementation.
+//===----------------------------------------------------------------------===//
+
+void mlir::bufferization::buildBufferDeallocationPipeline(
+ OpPassManager &pm, const BufferDeallocationPipelineOptions &options) {
+ pm.addNestedPass<func::FuncOp>(
+ memref::createExpandReallocPass(/*emitDeallocs=*/false));
+ pm.addNestedPass<func::FuncOp>(createCanonicalizerPass());
+ pm.addNestedPass<func::FuncOp>(createOwnershipBasedBufferDeallocationPass(
+ options.privateFunctionDynamicOwnership.getValue()));
+ pm.addNestedPass<func::FuncOp>(createCanonicalizerPass());
+ pm.addNestedPass<func::FuncOp>(createBufferDeallocationSimplificationPass());
+ pm.addPass(createLowerDeallocationsPass());
+ pm.addNestedPass<func::FuncOp>(createCSEPass());
+ pm.addNestedPass<func::FuncOp>(createCanonicalizerPass());
+}
+
+//===----------------------------------------------------------------------===//
+// Pipeline registration.
+//===----------------------------------------------------------------------===//
+
+void mlir::bufferization::registerBufferizationPipelines() {
+ PassPipelineRegistration<BufferDeallocationPipelineOptions>(
+ "buffer-deallocation-pipeline",
+ "The default pipeline for automatically inserting deallocation "
+ "operations after one-shot bufferization. Deallocation operations "
+ "(except `memref.realloc`) may not be present already.",
+ buildBufferDeallocationPipeline);
+}
diff --git a/mlir/lib/Dialect/Bufferization/Pipelines/CMakeLists.txt b/mlir/lib/Dialect/Bufferization/Pipelines/CMakeLists.txt
new file mode 100644
index 000000000000000..6e8dab64ba6b935
--- /dev/null
+++ b/mlir/lib/Dialect/Bufferization/Pipelines/CMakeLists.txt
@@ -0,0 +1,13 @@
+add_mlir_dialect_library(MLIRBufferizationPipelines
+ BufferizationPipelines.cpp
+
+ ADDITIONAL_HEADER_DIRS
+ ${MLIR_MAIN_INCLUDE_DIR}/mlir/Dialect/Bufferization
+
+ LINK_LIBS PUBLIC
+ MLIRBufferizationTransforms
+ MLIRMemRefTransforms
+ MLIRFuncDialect
+ MLIRPass
+ MLIRTransforms
+)
diff --git a/mlir/lib/Dialect/Bufferization/Transforms/OwnershipBasedBufferDeallocation.cpp b/mlir/lib/Dialect/Bufferization/Transforms/OwnershipBasedBufferDeallocation.cpp
index 02fb4d3c42fa521..43ba11cf132cb92 100644
--- a/mlir/lib/Dialect/Bufferization/Transforms/OwnershipBasedBufferDeallocation.cpp
+++ b/mlir/lib/Dialect/Bufferization/Transforms/OwnershipBasedBufferDeallocation.cpp
@@ -994,6 +994,11 @@ namespace {
struct OwnershipBasedBufferDeallocationPass
: public bufferization::impl::OwnershipBasedBufferDeallocationBase<
OwnershipBasedBufferDeallocationPass> {
+ OwnershipBasedBufferDeallocationPass() = default;
+ OwnershipBasedBufferDeallocationPass(bool privateFuncDynamicOwnership)
+ : OwnershipBasedBufferDeallocationPass() {
+ this->privateFuncDynamicOwnership.setValue(privateFuncDynamicOwnership);
+ }
void runOnOperation() override {
func::FuncOp func = getOperation();
if (func.isExternal())
@@ -1025,6 +1030,8 @@ LogicalResult bufferization::deallocateBuffersOwnershipBased(
//===----------------------------------------------------------------------===//
std::unique_ptr<Pass>
-mlir::bufferization::createOwnershipBasedBufferDeallocationPass() {
- return std::make_unique<OwnershipBasedBufferDeallocationPass>();
+mlir::bufferization::createOwnershipBasedBufferDeallocationPass(
+ bool privateFuncDynamicOwnership) {
+ return std::make_unique<OwnershipBasedBufferDeallocationPass>(
+ privateFuncDynamicOwnership);
}
diff --git a/mlir/test/Dialect/Bufferization/Transforms/OwnershipBasedBufferDeallocation/dealloc-branchop-interface.mlir b/mlir/test/Dialect/Bufferization/Transforms/OwnershipBasedBufferDeallocation/dealloc-branchop-interface.mlir
index 23a628cc2b83d99..3ae0529ab7d7466 100644
--- a/mlir/test/Dialect/Bufferization/Transforms/OwnershipBasedBufferDeallocation/dealloc-branchop-interface.mlir
+++ b/mlir/test/Dialect/Bufferization/Transforms/OwnershipBasedBufferDeallocation/dealloc-branchop-interface.mlir
@@ -2,6 +2,8 @@
// RUN: -buffer-deallocation-simplification -split-input-file %s | FileCheck %s
// RUN: mlir-opt -verify-diagnostics -ownership-based-buffer-deallocation=private-function-dynamic-ownership=true -split-input-file %s > /dev/null
+// RUN: mlir-opt %s -buffer-deallocation-pipeline --split-input-file > /dev/null
+
// Test Case:
// bb0
// / \
diff --git a/mlir/test/Dialect/Bufferization/Transforms/OwnershipBasedBufferDeallocation/dealloc-callop-interface.mlir b/mlir/test/Dialect/Bufferization/Transforms/OwnershipBasedBufferDeallocation/dealloc-callop-interface.mlir
index bff06d4499938df..389840c9cf77c9a 100644
--- a/mlir/test/Dialect/Bufferization/Transforms/OwnershipBasedBufferDeallocation/dealloc-callop-interface.mlir
+++ b/mlir/test/Dialect/Bufferization/Transforms/OwnershipBasedBufferDeallocation/dealloc-callop-interface.mlir
@@ -3,6 +3,9 @@
// RUN: mlir-opt -verify-diagnostics -ownership-based-buffer-deallocation=private-function-dynamic-ownership=true \
// RUN: --buffer-deallocation-simplification -split-input-file %s | FileCheck %s --check-prefix=CHECK-DYNAMIC
+// RUN: mlir-opt %s -buffer-deallocation-pipeline --split-input-file > /dev/null
+// RUN: mlir-opt %s -buffer-deallocation-pipeline=private-function-dynamic-ownership --split-input-file > /dev/null
+
func.func private @f(%arg0: memref<f64>) -> memref<f64> {
return %arg0 : memref<f64>
}
diff --git a/mlir/test/Dialect/Bufferization/Transforms/OwnershipBasedBufferDeallocation/dealloc-existing-deallocs.mlir b/mlir/test/Dialect/Bufferization/Transforms/OwnershipBasedBufferDeallocation/dealloc-existing-deallocs.mlir
index bf4eabd31a81241..7014746e348d25e 100644
--- a/mlir/test/Dialect/Bufferization/Transforms/OwnershipBasedBufferDeallocation/dealloc-existing-deallocs.mlir
+++ b/mlir/test/Dialect/Bufferization/Transforms/OwnershipBasedBufferDeallocation/dealloc-existing-deallocs.mlir
@@ -1,6 +1,8 @@
// RUN: mlir-opt -verify-diagnostics -expand-realloc=emit-deallocs=false -ownership-based-buffer-deallocation \
// RUN: --buffer-deallocation-simplification -split-input-file %s | FileCheck %s
+// RUN: mlir-opt %s -buffer-deallocation-pipeline --split-input-file > /dev/null
+
func.func @auto_dealloc() {
%c10 = arith.constant 10 : index
%c100 = arith.constant 100 : index
diff --git a/mlir/test/Dialect/Bufferization/Transforms/OwnershipBasedBufferDeallocation/dealloc-function-boundaries.mlir b/mlir/test/Dialect/Bufferization/Transforms/OwnershipBasedBufferDeallocation/dealloc-function-boundaries.mlir
index 44f3e20c5009309..13c55d0289880ef 100644
--- a/mlir/test/Dialect/Bufferization/Transforms/OwnershipBasedBufferDeallocation/dealloc-function-boundaries.mlir
+++ b/mlir/test/Dialect/Bufferization/Transforms/OwnershipBasedBufferDeallocation/dealloc-function-boundaries.mlir
@@ -3,6 +3,9 @@
// RUN: mlir-opt --allow-unregistered-dialect -verify-diagnostics -ownership-based-buffer-deallocation=private-function-dynamic-ownership=true \
// RUN: --buffer-deallocation-simplification -split-input-file %s | FileCheck %s --check-prefix=CHECK-DYNAMIC
+// RUN: mlir-opt %s -buffer-deallocation-pipeline --split-input-file > /dev/null
+// RUN: mlir-opt %s -buffer-deallocation-pipeline=private-function-dynamic-ownership --split-input-file > /dev/null
+
// Test Case: Existing AllocOp with no users.
// BufferDeallocation expected behavior: It should insert a DeallocOp right
// before ReturnOp.
diff --git a/mlir/test/Dialect/Bufferization/Transforms/OwnershipBasedBufferDeallocation/dealloc-memoryeffect-interface.mlir b/mlir/test/Dialect/Bufferization/Transforms/OwnershipBasedBufferDeallocation/dealloc-memoryeffect-interface.mlir
index 460e37aa03059ff..44cf16385603e07 100644
--- a/mlir/test/Dialect/Bufferization/Transforms/OwnershipBasedBufferDeallocation/dealloc-memoryeffect-interface.mlir
+++ b/mlir/test/Dialect/Bufferization/Transforms/OwnershipBasedBufferDeallocation/dealloc-memoryeffect-interface.mlir
@@ -2,6 +2,8 @@
// RUN: --buffer-deallocation-simplification -split-input-file %s | FileCheck %s
// RUN: mlir-opt -verify-diagnostics -ownership-based-buffer-deallocation=private-function-dynamic-ownership=true -split-input-file %s > /dev/null
+// RUN: mlir-opt %s -buffer-deallocation-pipeline --split-input-file > /dev/null
+
// Test Case: Dead operations in a single block.
// BufferDeallocation expected behavior: It only inserts the two missing
// DeallocOps after the last BufferBasedOp.
diff --git a/mlir/test/Dialect/Bufferization/Transforms/OwnershipBasedBufferDeallocation/dealloc-region-branchop-interface.mlir b/mlir/test/Dialect/Bufferization/Transforms/OwnershipBasedBufferDeallocation/dealloc-region-branchop-interface.mlir
index 66449aa2ffdb60e..dc372749fc074be 100644
--- a/mlir/test/Dialect/Bufferization/Transforms/OwnershipBasedBufferDeallocation/dealloc-region-branchop-interface.mlir
+++ b/mlir/test/Dialect/Bufferization/Transforms/OwnershipBasedBufferDeallocation/dealloc-region-branchop-interface.mlir
@@ -2,6 +2,8 @@
// RUN: --buffer-deallocation-simplification -split-input-file %s | FileCheck %s
// RUN: mlir-opt -allow-unregistered-dialect -verify-diagnostics -ownership-based-buffer-deallocation=private-function-dynamic-ownership=true -split-input-file %s > /dev/null
+// RUN: mlir-opt %s -buffer-deallocation-pipeline --split-input-file --verify-diagnostics > /dev/null
+
// Test Case: Nested regions - This test defines a BufferBasedOp inside the
// region of a RegionBufferBasedOp.
// BufferDeallocation expected behavior: The AllocOp for the BufferBasedOp
diff --git a/mlir/test/Dialect/Bufferization/Transforms/OwnershipBasedBufferDeallocation/dealloc-subviews.mlir b/mlir/test/Dialect/Bufferization/Transforms/OwnershipBasedBufferDeallocation/dealloc-subviews.mlir
index 666b3b08995e8ad..35523319de1548e 100644
--- a/mlir/test/Dialect/Bufferization/Transforms/OwnershipBasedBufferDeallocation/dealloc-subviews.mlir
+++ b/mlir/test/Dialect/Bufferization/Transforms/OwnershipBasedBufferDeallocation/dealloc-subviews.mlir
@@ -2,6 +2,8 @@
// RUN: --buffer-deallocation-simplification -split-input-file %s | FileCheck %s
// RUN: mlir-opt -verify-diagnostics -ownership-based-buffer-deallocation=private-function-dynamic-ownership=true -split-input-file %s > /dev/null
+// RUN: mlir-opt %s -buffer-deallocation-pipeline --split-input-file > /dev/null
+
// CHECK-LABEL: func @subview
func.func @subview(%arg0 : index, %arg1 : index, %arg2 : memref<?x?xf32>) {
%0 = memref.alloc() : memref<64x4xf32, strided<[4, 1], offset: 0>>
diff --git a/utils/bazel/llvm-project-overlay/mlir/BUILD.bazel b/utils/bazel/llvm-project-overlay/mlir/BUILD.bazel
index 9cdbe29587fed62..9bea555f701757c 100644
--- a/utils/bazel/llvm-project-overlay/mlir/BUILD.bazel
+++ b/utils/bazel/llvm-project-overlay/mlir/BUILD.bazel
@@ -8249,6 +8249,7 @@ cc_library(
":AsyncToLLVM",
":AsyncTransforms",
":BufferizationDialect",
+ ":BufferizationPipelines",
":BufferizationTransformOps",
":BufferizationTransforms",
":CastInterfaces",
@@ -12214,6 +12215,21 @@ cc_library(
],
)
+cc_library(
+ name = "BufferizationPipelines",
+ srcs = glob(["lib/Dialect/Bufferization/Pipelines/*.cpp"]),
+ hdrs = ["include/mlir/Dialect/Bufferization/Pipelines/Passes.h"],
+ includes = ["include"],
+ deps = [
+ ":BufferizationToMemRef",
+ ":BufferizationTransforms",
+ ":FuncDialect",
+ ":MemRefTransforms",
+ ":Pass",
+ ":Transforms",
+ ],
+)
+
td_library(
name = "DLTIDialectTdFiles",
srcs = [
More information about the llvm-commits
mailing list