Skip to content

Commit 08b7a71

Browse files
maerhartmaerhart
authored
[mlir][bufferization] Define a pipeline for buffer deallocation (#66352)
Since ownership based 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).
1 parent bdac5de commit 08b7a71

File tree

15 files changed

+155
-3
lines changed

15 files changed

+155
-3
lines changed

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

Lines changed: 50 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,50 @@
1+
//===- Passes.h - Bufferization pipeline entry points -----------*- C++ -*-===//
2+
//
3+
// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
4+
// See https://llvm.org/LICENSE.txt for license information.
5+
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
6+
//
7+
//===----------------------------------------------------------------------===//
8+
//
9+
// This header file defines prototypes of all bufferization pipelines.
10+
//
11+
//===----------------------------------------------------------------------===//
12+
13+
#ifndef MLIR_DIALECT_BUFFERIZATION_PIPELINES_PASSES_H
14+
#define MLIR_DIALECT_BUFFERIZATION_PIPELINES_PASSES_H
15+
16+
#include "mlir/Pass/PassOptions.h"
17+
18+
namespace mlir {
19+
namespace bufferization {
20+
21+
/// Options for the buffer deallocation pipeline.
22+
struct BufferDeallocationPipelineOptions
23+
: public PassPipelineOptions<BufferDeallocationPipelineOptions> {
24+
PassOptions::Option<bool> privateFunctionDynamicOwnership{
25+
*this, "private-function-dynamic-ownership",
26+
llvm::cl::desc(
27+
"Allows to add additional arguments to private functions to "
28+
"dynamically pass ownership of memrefs to callees. This can enable "
29+
"earlier deallocations."),
30+
llvm::cl::init(false)};
31+
};
32+
33+
//===----------------------------------------------------------------------===//
34+
// Building and Registering.
35+
//===----------------------------------------------------------------------===//
36+
37+
/// Adds the buffer deallocation pipeline to the `OpPassManager`. This
38+
/// is the standard pipeline for deallocating the MemRefs introduced by the
39+
/// One-Shot bufferization pass.
40+
void buildBufferDeallocationPipeline(
41+
OpPassManager &pm, const BufferDeallocationPipelineOptions &options);
42+
43+
/// Registers all pipelines for the `bufferization` dialect. Currently,
44+
/// this includes only the "buffer-deallocation-pipeline".
45+
void registerBufferizationPipelines();
46+
47+
} // namespace bufferization
48+
} // namespace mlir
49+
50+
#endif // MLIR_DIALECT_BUFFERIZATION_PIPELINES_PASSES_H

‎mlir/include/mlir/Dialect/Bufferization/Transforms/Passes.h

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -30,7 +30,8 @@ std::unique_ptr<Pass> createBufferDeallocationPass();
3030

3131
/// Creates an instance of the OwnershipBasedBufferDeallocation pass to free all
3232
/// allocated buffers.
33-
std::unique_ptr<Pass> createOwnershipBasedBufferDeallocationPass();
33+
std::unique_ptr<Pass> createOwnershipBasedBufferDeallocationPass(
34+
bool privateFuncDynamicOwnership = false);
3435

3536
/// Creates a pass that optimizes `bufferization.dealloc` operations. For
3637
/// example, it reduces the number of alias checks needed at runtime using

‎mlir/include/mlir/InitAllPasses.h

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -20,6 +20,7 @@
2020
#include "mlir/Dialect/Arith/Transforms/Passes.h"
2121
#include "mlir/Dialect/ArmSME/Transforms/Passes.h"
2222
#include "mlir/Dialect/Async/Passes.h"
23+
#include "mlir/Dialect/Bufferization/Pipelines/Passes.h"
2324
#include "mlir/Dialect/Bufferization/Transforms/Passes.h"
2425
#include "mlir/Dialect/Func/Transforms/Passes.h"
2526
#include "mlir/Dialect/GPU/Transforms/Passes.h"
@@ -81,6 +82,7 @@ inline void registerAllPasses() {
8182
arm_sme::registerArmSMEPasses();
8283

8384
// Dialect pipelines
85+
bufferization::registerBufferizationPipelines();
8486
sparse_tensor::registerSparseTensorPipelines();
8587
}
8688

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

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,4 @@
11
add_subdirectory(IR)
2+
add_subdirectory(Pipelines)
23
add_subdirectory(TransformOps)
34
add_subdirectory(Transforms)

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

Lines changed: 46 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,46 @@
1+
//===- BufferizationPipelines.cpp - Pipelines for bufferization -----------===//
2+
//
3+
// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
4+
// See https://llvm.org/LICENSE.txt for license information.
5+
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
6+
//
7+
//===----------------------------------------------------------------------===//
8+
9+
#include "mlir/Dialect/Bufferization/Pipelines/Passes.h"
10+
11+
#include "mlir/Dialect/Bufferization/Transforms/Passes.h"
12+
#include "mlir/Dialect/Func/IR/FuncOps.h"
13+
#include "mlir/Dialect/MemRef/Transforms/Passes.h"
14+
#include "mlir/Pass/PassManager.h"
15+
#include "mlir/Transforms/Passes.h"
16+
17+
//===----------------------------------------------------------------------===//
18+
// Pipeline implementation.
19+
//===----------------------------------------------------------------------===//
20+
21+
void mlir::bufferization::buildBufferDeallocationPipeline(
22+
OpPassManager &pm, const BufferDeallocationPipelineOptions &options) {
23+
pm.addNestedPass<func::FuncOp>(
24+
memref::createExpandReallocPass(/*emitDeallocs=*/false));
25+
pm.addNestedPass<func::FuncOp>(createCanonicalizerPass());
26+
pm.addNestedPass<func::FuncOp>(createOwnershipBasedBufferDeallocationPass(
27+
options.privateFunctionDynamicOwnership.getValue()));
28+
pm.addNestedPass<func::FuncOp>(createCanonicalizerPass());
29+
pm.addNestedPass<func::FuncOp>(createBufferDeallocationSimplificationPass());
30+
pm.addPass(createLowerDeallocationsPass());
31+
pm.addNestedPass<func::FuncOp>(createCSEPass());
32+
pm.addNestedPass<func::FuncOp>(createCanonicalizerPass());
33+
}
34+
35+
//===----------------------------------------------------------------------===//
36+
// Pipeline registration.
37+
//===----------------------------------------------------------------------===//
38+
39+
void mlir::bufferization::registerBufferizationPipelines() {
40+
PassPipelineRegistration<BufferDeallocationPipelineOptions>(
41+
"buffer-deallocation-pipeline",
42+
"The default pipeline for automatically inserting deallocation "
43+
"operations after one-shot bufferization. Deallocation operations "
44+
"(except `memref.realloc`) may not be present already.",
45+
buildBufferDeallocationPipeline);
46+
}

‎mlir/lib/Dialect/Bufferization/Pipelines/CMakeLists.txt

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,13 @@
1+
add_mlir_dialect_library(MLIRBufferizationPipelines
2+
BufferizationPipelines.cpp
3+
4+
ADDITIONAL_HEADER_DIRS
5+
${MLIR_MAIN_INCLUDE_DIR}/mlir/Dialect/Bufferization
6+
7+
LINK_LIBS PUBLIC
8+
MLIRBufferizationTransforms
9+
MLIRMemRefTransforms
10+
MLIRFuncDialect
11+
MLIRPass
12+
MLIRTransforms
13+
)

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

Lines changed: 9 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -994,6 +994,11 @@ namespace {
994994
struct OwnershipBasedBufferDeallocationPass
995995
: public bufferization::impl::OwnershipBasedBufferDeallocationBase<
996996
OwnershipBasedBufferDeallocationPass> {
997+
OwnershipBasedBufferDeallocationPass() = default;
998+
OwnershipBasedBufferDeallocationPass(bool privateFuncDynamicOwnership)
999+
: OwnershipBasedBufferDeallocationPass() {
1000+
this->privateFuncDynamicOwnership.setValue(privateFuncDynamicOwnership);
1001+
}
9971002
void runOnOperation() override {
9981003
func::FuncOp func = getOperation();
9991004
if (func.isExternal())
@@ -1025,6 +1030,8 @@ LogicalResult bufferization::deallocateBuffersOwnershipBased(
10251030
//===----------------------------------------------------------------------===//
10261031

10271032
std::unique_ptr<Pass>
1028-
mlir::bufferization::createOwnershipBasedBufferDeallocationPass() {
1029-
return std::make_unique<OwnershipBasedBufferDeallocationPass>();
1033+
mlir::bufferization::createOwnershipBasedBufferDeallocationPass(
1034+
bool privateFuncDynamicOwnership) {
1035+
return std::make_unique<OwnershipBasedBufferDeallocationPass>(
1036+
privateFuncDynamicOwnership);
10301037
}

‎mlir/test/Dialect/Bufferization/Transforms/OwnershipBasedBufferDeallocation/dealloc-branchop-interface.mlir

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,8 @@
22
// RUN: -buffer-deallocation-simplification -split-input-file %s | FileCheck %s
33
// RUN: mlir-opt -verify-diagnostics -ownership-based-buffer-deallocation=private-function-dynamic-ownership=true -split-input-file %s > /dev/null
44

5+
// RUN: mlir-opt %s -buffer-deallocation-pipeline --split-input-file > /dev/null
6+
57
// Test Case:
68
// bb0
79
// / \

‎mlir/test/Dialect/Bufferization/Transforms/OwnershipBasedBufferDeallocation/dealloc-callop-interface.mlir

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,9 @@
33
// RUN: mlir-opt -verify-diagnostics -ownership-based-buffer-deallocation=private-function-dynamic-ownership=true \
44
// RUN: --buffer-deallocation-simplification -split-input-file %s | FileCheck %s --check-prefix=CHECK-DYNAMIC
55

6+
// RUN: mlir-opt %s -buffer-deallocation-pipeline --split-input-file > /dev/null
7+
// RUN: mlir-opt %s -buffer-deallocation-pipeline=private-function-dynamic-ownership --split-input-file > /dev/null
8+
69
func.func private @f(%arg0: memref<f64>) -> memref<f64> {
710
return %arg0 : memref<f64>
811
}

‎mlir/test/Dialect/Bufferization/Transforms/OwnershipBasedBufferDeallocation/dealloc-existing-deallocs.mlir

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,8 @@
11
// RUN: mlir-opt -verify-diagnostics -expand-realloc=emit-deallocs=false -ownership-based-buffer-deallocation \
22
// RUN: --buffer-deallocation-simplification -split-input-file %s | FileCheck %s
33

4+
// RUN: mlir-opt %s -buffer-deallocation-pipeline --split-input-file > /dev/null
5+
46
func.func @auto_dealloc() {
57
%c10 = arith.constant 10 : index
68
%c100 = arith.constant 100 : index

‎mlir/test/Dialect/Bufferization/Transforms/OwnershipBasedBufferDeallocation/dealloc-function-boundaries.mlir

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,9 @@
33
// RUN: mlir-opt --allow-unregistered-dialect -verify-diagnostics -ownership-based-buffer-deallocation=private-function-dynamic-ownership=true \
44
// RUN: --buffer-deallocation-simplification -split-input-file %s | FileCheck %s --check-prefix=CHECK-DYNAMIC
55

6+
// RUN: mlir-opt %s -buffer-deallocation-pipeline --split-input-file > /dev/null
7+
// RUN: mlir-opt %s -buffer-deallocation-pipeline=private-function-dynamic-ownership --split-input-file > /dev/null
8+
69
// Test Case: Existing AllocOp with no users.
710
// BufferDeallocation expected behavior: It should insert a DeallocOp right
811
// before ReturnOp.

‎mlir/test/Dialect/Bufferization/Transforms/OwnershipBasedBufferDeallocation/dealloc-memoryeffect-interface.mlir

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,8 @@
22
// RUN: --buffer-deallocation-simplification -split-input-file %s | FileCheck %s
33
// RUN: mlir-opt -verify-diagnostics -ownership-based-buffer-deallocation=private-function-dynamic-ownership=true -split-input-file %s > /dev/null
44

5+
// RUN: mlir-opt %s -buffer-deallocation-pipeline --split-input-file > /dev/null
6+
57
// Test Case: Dead operations in a single block.
68
// BufferDeallocation expected behavior: It only inserts the two missing
79
// DeallocOps after the last BufferBasedOp.

‎mlir/test/Dialect/Bufferization/Transforms/OwnershipBasedBufferDeallocation/dealloc-region-branchop-interface.mlir

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,8 @@
22
// RUN: --buffer-deallocation-simplification -split-input-file %s | FileCheck %s
33
// RUN: mlir-opt -allow-unregistered-dialect -verify-diagnostics -ownership-based-buffer-deallocation=private-function-dynamic-ownership=true -split-input-file %s > /dev/null
44

5+
// RUN: mlir-opt %s -buffer-deallocation-pipeline --split-input-file --verify-diagnostics > /dev/null
6+
57
// Test Case: Nested regions - This test defines a BufferBasedOp inside the
68
// region of a RegionBufferBasedOp.
79
// BufferDeallocation expected behavior: The AllocOp for the BufferBasedOp

‎mlir/test/Dialect/Bufferization/Transforms/OwnershipBasedBufferDeallocation/dealloc-subviews.mlir

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,8 @@
22
// RUN: --buffer-deallocation-simplification -split-input-file %s | FileCheck %s
33
// RUN: mlir-opt -verify-diagnostics -ownership-based-buffer-deallocation=private-function-dynamic-ownership=true -split-input-file %s > /dev/null
44

5+
// RUN: mlir-opt %s -buffer-deallocation-pipeline --split-input-file > /dev/null
6+
57
// CHECK-LABEL: func @subview
68
func.func @subview(%arg0 : index, %arg1 : index, %arg2 : memref<?x?xf32>) {
79
%0 = memref.alloc() : memref<64x4xf32, strided<[4, 1], offset: 0>>

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

Lines changed: 16 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -8249,6 +8249,7 @@ cc_library(
82498249
":AsyncToLLVM",
82508250
":AsyncTransforms",
82518251
":BufferizationDialect",
8252+
":BufferizationPipelines",
82528253
":BufferizationTransformOps",
82538254
":BufferizationTransforms",
82548255
":CastInterfaces",
@@ -12214,6 +12215,21 @@ cc_library(
1221412215
],
1221512216
)
1221612217

12218+
cc_library(
12219+
name = "BufferizationPipelines",
12220+
srcs = glob(["lib/Dialect/Bufferization/Pipelines/*.cpp"]),
12221+
hdrs = ["include/mlir/Dialect/Bufferization/Pipelines/Passes.h"],
12222+
includes = ["include"],
12223+
deps = [
12224+
":BufferizationToMemRef",
12225+
":BufferizationTransforms",
12226+
":FuncDialect",
12227+
":MemRefTransforms",
12228+
":Pass",
12229+
":Transforms",
12230+
],
12231+
)
12232+
1221712233
td_library(
1221812234
name = "DLTIDialectTdFiles",
1221912235
srcs = [

0 commit comments

Comments
 (0)