1//===- AsyncRuntimeRefCountingOpt.cpp - Async Ref Counting --------------===//
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// Optimize Async dialect reference counting operations.
10//
11//===----------------------------------------------------------------------===//
12
13#include "mlir/Dialect/Async/Passes.h"
14
15#include "mlir/Dialect/Async/IR/Async.h"
16#include "mlir/Dialect/Func/IR/FuncOps.h"
17#include "llvm/Support/Debug.h"
18
19namespace mlir {
20#define GEN_PASS_DEF_ASYNCRUNTIMEREFCOUNTINGOPTPASS
21#include "mlir/Dialect/Async/Passes.h.inc"
22} // namespace mlir
23
24#define DEBUG_TYPE "async-ref-counting"
25
26using namespace mlir;
27using namespace mlir::async;
28
29namespace {
30
31class AsyncRuntimeRefCountingOptPass
32 : public impl::AsyncRuntimeRefCountingOptPassBase<
33 AsyncRuntimeRefCountingOptPass> {
34public:
35 AsyncRuntimeRefCountingOptPass() = default;
36 void runOnOperation() override;
37
38private:
39 LogicalResult optimizeReferenceCounting(
40 Value value, llvm::SmallDenseMap<Operation *, Operation *> &cancellable);
41};
42
43} // namespace
44
45LogicalResult AsyncRuntimeRefCountingOptPass::optimizeReferenceCounting(
46 Value value, llvm::SmallDenseMap<Operation *, Operation *> &cancellable) {
47 Region *definingRegion = value.getParentRegion();
48
49 // Find all users of the `value` inside each block, including operations that
50 // do not use `value` directly, but have a direct use inside nested region(s).
51 //
52 // Example:
53 //
54 // ^bb1:
55 // %token = ...
56 // scf.if %cond {
57 // ^bb2:
58 // async.runtime.await %token : !async.token
59 // }
60 //
61 // %token has a use inside ^bb2 (`async.runtime.await`) and inside ^bb1
62 // (`scf.if`).
63
64 struct BlockUsersInfo {
65 llvm::SmallVector<RuntimeAddRefOp, 4> addRefs;
66 llvm::SmallVector<RuntimeDropRefOp, 4> dropRefs;
67 llvm::SmallVector<Operation *, 4> users;
68 };
69
70 llvm::DenseMap<Block *, BlockUsersInfo> blockUsers;
71
72 auto updateBlockUsersInfo = [&](Operation *user) {
73 BlockUsersInfo &info = blockUsers[user->getBlock()];
74 info.users.push_back(Elt: user);
75
76 if (auto addRef = dyn_cast<RuntimeAddRefOp>(Val: user))
77 info.addRefs.push_back(Elt: addRef);
78 if (auto dropRef = dyn_cast<RuntimeDropRefOp>(Val: user))
79 info.dropRefs.push_back(Elt: dropRef);
80 };
81
82 for (Operation *user : value.getUsers()) {
83 while (user->getParentRegion() != definingRegion) {
84 updateBlockUsersInfo(user);
85 user = user->getParentOp();
86 assert(user != nullptr && "value user lies outside of the value region");
87 }
88
89 updateBlockUsersInfo(user);
90 }
91
92 // Sort all operations found in the block.
93 auto preprocessBlockUsersInfo = [](BlockUsersInfo &info) -> BlockUsersInfo & {
94 auto isBeforeInBlock = [](Operation *a, Operation *b) -> bool {
95 return a->isBeforeInBlock(other: b);
96 };
97 llvm::sort(C&: info.addRefs, Comp: isBeforeInBlock);
98 llvm::sort(C&: info.dropRefs, Comp: isBeforeInBlock);
99 llvm::sort(C&: info.users, Comp: [&](Operation *a, Operation *b) -> bool {
100 return isBeforeInBlock(a, b);
101 });
102
103 return info;
104 };
105
106 // Find and erase matching pairs of `add_ref` / `drop_ref` operations in the
107 // blocks that modify the reference count of the `value`.
108 for (auto &kv : blockUsers) {
109 BlockUsersInfo &info = preprocessBlockUsersInfo(kv.second);
110
111 for (RuntimeAddRefOp addRef : info.addRefs) {
112 for (RuntimeDropRefOp dropRef : info.dropRefs) {
113 // `drop_ref` operation after the `add_ref` with matching count.
114 if (dropRef.getCount() != addRef.getCount() ||
115 dropRef->isBeforeInBlock(other: addRef.getOperation()))
116 continue;
117
118 // When reference counted value passed to a function as an argument,
119 // function takes ownership of +1 reference and it will drop it before
120 // returning.
121 //
122 // Example:
123 //
124 // %token = ... : !async.token
125 //
126 // async.runtime.add_ref %token {count = 1 : i64} : !async.token
127 // call @pass_token(%token: !async.token, ...)
128 //
129 // async.await %token : !async.token
130 // async.runtime.drop_ref %token {count = 1 : i64} : !async.token
131 //
132 // In this example if we'll cancel a pair of reference counting
133 // operations we might end up with a deallocated token when we'll
134 // reach `async.await` operation.
135 Operation *firstFunctionCallUser = nullptr;
136 Operation *lastNonFunctionCallUser = nullptr;
137
138 for (Operation *user : info.users) {
139 // `user` operation lies after `addRef` ...
140 if (user == addRef || user->isBeforeInBlock(other: addRef))
141 continue;
142 // ... and before `dropRef`.
143 if (user == dropRef || dropRef->isBeforeInBlock(other: user))
144 break;
145
146 // Find the first function call user of the reference counted value.
147 Operation *functionCall = dyn_cast<func::CallOp>(Val: user);
148 if (functionCall &&
149 (!firstFunctionCallUser ||
150 functionCall->isBeforeInBlock(other: firstFunctionCallUser))) {
151 firstFunctionCallUser = functionCall;
152 continue;
153 }
154
155 // Find the last regular user of the reference counted value.
156 if (!functionCall &&
157 (!lastNonFunctionCallUser ||
158 lastNonFunctionCallUser->isBeforeInBlock(other: user))) {
159 lastNonFunctionCallUser = user;
160 continue;
161 }
162 }
163
164 // Non function call user after the function call user of the reference
165 // counted value.
166 if (firstFunctionCallUser && lastNonFunctionCallUser &&
167 firstFunctionCallUser->isBeforeInBlock(other: lastNonFunctionCallUser))
168 continue;
169
170 // Try to cancel the pair of `add_ref` and `drop_ref` operations.
171 auto emplaced = cancellable.try_emplace(Key: dropRef.getOperation(),
172 Args: addRef.getOperation());
173
174 if (!emplaced.second) // `drop_ref` was already marked for removal
175 continue; // go to the next `drop_ref`
176
177 if (emplaced.second) // successfully cancelled `add_ref` <-> `drop_ref`
178 break; // go to the next `add_ref`
179 }
180 }
181 }
182
183 return success();
184}
185
186void AsyncRuntimeRefCountingOptPass::runOnOperation() {
187 Operation *op = getOperation();
188
189 // Mapping from `dropRef.getOperation()` to `addRef.getOperation()`.
190 //
191 // Find all cancellable pairs of operation and erase them in the end to keep
192 // all iterators valid while we are walking the function operations.
193 llvm::SmallDenseMap<Operation *, Operation *> cancellable;
194
195 // Optimize reference counting for values defined by block arguments.
196 WalkResult blockWalk = op->walk(callback: [&](Block *block) -> WalkResult {
197 for (BlockArgument arg : block->getArguments())
198 if (isRefCounted(type: arg.getType()))
199 if (failed(Result: optimizeReferenceCounting(value: arg, cancellable)))
200 return WalkResult::interrupt();
201
202 return WalkResult::advance();
203 });
204
205 if (blockWalk.wasInterrupted())
206 signalPassFailure();
207
208 // Optimize reference counting for values defined by operation results.
209 WalkResult opWalk = op->walk(callback: [&](Operation *op) -> WalkResult {
210 for (unsigned i = 0; i < op->getNumResults(); ++i)
211 if (isRefCounted(type: op->getResultTypes()[i]))
212 if (failed(Result: optimizeReferenceCounting(value: op->getResult(idx: i), cancellable)))
213 return WalkResult::interrupt();
214
215 return WalkResult::advance();
216 });
217
218 if (opWalk.wasInterrupted())
219 signalPassFailure();
220
221 LLVM_DEBUG({
222 llvm::dbgs() << "Found " << cancellable.size()
223 << " cancellable reference counting operations\n";
224 });
225
226 // Erase all cancellable `add_ref <-> drop_ref` operation pairs.
227 for (auto &kv : cancellable) {
228 kv.first->erase();
229 kv.second->erase();
230 }
231}
232

source code of mlir/lib/Dialect/Async/Transforms/AsyncRuntimeRefCountingOpt.cpp