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

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