1//===- OptimizeAllocationLiveness.cpp - impl. optimize allocation liveness pass
2//-===//
3//
4// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
5// See https://llvm.org/LICENSE.txt for license information.
6// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
7//
8//===----------------------------------------------------------------------===//
9//
10// This file implements a pass for optimizing allocation liveness.
11// The pass moves the deallocation operation after the last user of the
12// allocated buffer.
13//===----------------------------------------------------------------------===//
14
15#include "mlir/Dialect/Bufferization/Transforms/BufferViewFlowAnalysis.h"
16#include "mlir/Dialect/Bufferization/Transforms/Passes.h"
17#include "mlir/Dialect/Func/IR/FuncOps.h"
18#include "mlir/Dialect/MemRef/IR/MemRef.h"
19#include "mlir/IR/Operation.h"
20#include "mlir/IR/Value.h"
21#include "mlir/Interfaces/SideEffectInterfaces.h"
22#include "llvm/Support/Debug.h"
23#include "llvm/Support/ErrorHandling.h"
24
25#define DEBUG_TYPE "optimize-allocation-liveness"
26#define DBGS() (llvm::dbgs() << '[' << DEBUG_TYPE << "] ")
27#define LDBG(X) LLVM_DEBUG(DBGS() << X << "\n")
28
29namespace mlir {
30namespace bufferization {
31#define GEN_PASS_DEF_OPTIMIZEALLOCATIONLIVENESSPASS
32#include "mlir/Dialect/Bufferization/Transforms/Passes.h.inc"
33} // namespace bufferization
34} // namespace mlir
35
36using namespace mlir;
37
38namespace {
39
40//===----------------------------------------------------------------------===//
41// Helper functions
42//===----------------------------------------------------------------------===//
43
44/// Return true if `a` happens before `b`, i.e., `a` or one of its ancestors
45/// properly dominates `b` and `b` is not inside `a`.
46static bool happensBefore(Operation *a, Operation *b) {
47 do {
48 if (a->isProperAncestor(other: b))
49 return false;
50 if (Operation *bAncestor = a->getBlock()->findAncestorOpInBlock(op&: *b)) {
51 return a->isBeforeInBlock(other: bAncestor);
52 }
53 } while ((a = a->getParentOp()));
54 return false;
55}
56
57/// This method searches for a user of value that is a dealloc operation.
58/// If multiple users with free effect are found, return nullptr.
59Operation *findUserWithFreeSideEffect(Value value) {
60 Operation *freeOpUser = nullptr;
61 for (Operation *user : value.getUsers()) {
62 if (MemoryEffectOpInterface memEffectOp =
63 dyn_cast<MemoryEffectOpInterface>(user)) {
64 SmallVector<MemoryEffects::EffectInstance, 2> effects;
65 memEffectOp.getEffects(effects);
66
67 for (const auto &effect : effects) {
68 if (isa<MemoryEffects::Free>(effect.getEffect())) {
69 if (freeOpUser) {
70 LDBG("Multiple users with free effect found: " << *freeOpUser
71 << " and " << *user);
72 return nullptr;
73 }
74 freeOpUser = user;
75 }
76 }
77 }
78 }
79 return freeOpUser;
80}
81
82/// Checks if the given op allocates memory.
83static bool hasMemoryAllocEffect(MemoryEffectOpInterface memEffectOp) {
84 SmallVector<MemoryEffects::EffectInstance, 2> effects;
85 memEffectOp.getEffects(effects);
86 for (const auto &effect : effects) {
87 if (isa<MemoryEffects::Allocate>(effect.getEffect())) {
88 return true;
89 }
90 }
91 return false;
92}
93
94/// Extracts OpResult's with Allocate effects from given op
95static SmallVector<OpResult>
96collectAllocations(MemoryEffectOpInterface allocOp) {
97 SmallVector<MemoryEffects::EffectInstance> effects;
98 allocOp.getEffects(effects);
99 SmallVector<OpResult> allocResults;
100 for (const MemoryEffects::EffectInstance &it : effects)
101 if (isa<MemoryEffects::Allocate>(it.getEffect()))
102 if (auto val = it.getValue(); val && val.getDefiningOp() == allocOp)
103 allocResults.push_back(cast<OpResult>(val));
104 return allocResults;
105}
106
107struct OptimizeAllocationLiveness
108 : public bufferization::impl::OptimizeAllocationLivenessPassBase<
109 OptimizeAllocationLiveness> {
110public:
111 OptimizeAllocationLiveness() = default;
112
113 void runOnOperation() override {
114 func::FuncOp func = getOperation();
115
116 if (func.isExternal())
117 return;
118
119 BufferViewFlowAnalysis analysis = BufferViewFlowAnalysis(func);
120
121 func.walk([&](MemoryEffectOpInterface memEffectOp) -> WalkResult {
122 if (!hasMemoryAllocEffect(memEffectOp))
123 return WalkResult::advance();
124
125 auto allocOp = memEffectOp;
126 LDBG("Checking alloc op: " << allocOp);
127
128 SmallVector<OpResult> allocationResults = collectAllocations(allocOp);
129 // Multiple allocations from a single op are not considered here yet.
130 if (allocationResults.size() != 1)
131 return WalkResult::advance();
132
133 OpResult allocResult = allocationResults[0];
134 LDBG("On allocation result: " << allocResult);
135
136 auto *deallocOp = findUserWithFreeSideEffect(value: allocResult);
137 if (!deallocOp || (deallocOp->getBlock() != allocOp->getBlock())) {
138 // The pass handles allocations that have a single dealloc op in the
139 // same block. We also should not hoist the dealloc op out of
140 // conditionals.
141 return WalkResult::advance();
142 }
143
144 Operation *lastUser = nullptr;
145 const BufferViewFlowAnalysis::ValueSetT &deps =
146 analysis.resolve(value: allocResult);
147 for (auto dep : llvm::make_early_inc_range(deps)) {
148 for (auto *user : dep.getUsers()) {
149 // We are looking for a non dealloc op user.
150 // check if user is the dealloc op itself.
151 if (user == deallocOp)
152 continue;
153
154 // find the ancestor of user that is in the same block as the allocOp.
155 auto topUser = allocOp->getBlock()->findAncestorOpInBlock(*user);
156 if (!lastUser || happensBefore(lastUser, topUser)) {
157 lastUser = topUser;
158 }
159 }
160 }
161 if (lastUser == nullptr) {
162 return WalkResult::advance();
163 }
164 LDBG("Last user found: " << *lastUser);
165 assert(lastUser->getBlock() == allocOp->getBlock());
166 assert(lastUser->getBlock() == deallocOp->getBlock());
167 // Move the dealloc op after the last user.
168 deallocOp->moveAfter(lastUser);
169 LDBG("Moved dealloc op after: " << *lastUser);
170
171 return WalkResult::advance();
172 });
173 }
174};
175
176} // end anonymous namespace
177

source code of mlir/lib/Dialect/Bufferization/Transforms/OptimizeAllocationLiveness.cpp