1 //===- AsyncRuntimeRefCountingOpt.cpp - Async Ref Counting --------------===//
5 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
7 //===----------------------------------------------------------------------===//
11 //===----------------------------------------------------------------------===//
25 #define DEBUG_TYPE "async-ref-counting"
73   auto updateBlockUsersInfo = [&](Operation *user) {  in optimizeReferenceCounting()  argument
74     BlockUsersInfo &info = blockUsers[user->getBlock()];  in optimizeReferenceCounting()
75     info.users.push_back(user);  in optimizeReferenceCounting()
77     if (auto addRef = dyn_cast<RuntimeAddRefOp>(user))  in optimizeReferenceCounting()
79     if (auto dropRef = dyn_cast<RuntimeDropRefOp>(user))  in optimizeReferenceCounting()
83   for (Operation *user : value.getUsers()) {  in optimizeReferenceCounting()
84     while (user->getParentRegion() != definingRegion) {  in optimizeReferenceCounting()
85       updateBlockUsersInfo(user);  in optimizeReferenceCounting()
86       user = user->getParentOp();  in optimizeReferenceCounting()
87       assert(user != nullptr && "value user lies outside of the value region");  in optimizeReferenceCounting()
90     updateBlockUsersInfo(user);  in optimizeReferenceCounting()
94   auto preprocessBlockUsersInfo = [](BlockUsersInfo &info) -> BlockUsersInfo & {  in optimizeReferenceCounting()
95     auto isBeforeInBlock = [](Operation *a, Operation *b) -> bool {  in optimizeReferenceCounting()
96       return a->isBeforeInBlock(b);  in optimizeReferenceCounting()
100     llvm::sort(info.users, [&](Operation *a, Operation *b) -> bool {  in optimizeReferenceCounting()
116             dropRef->isBeforeInBlock(addRef.getOperation()))  in optimizeReferenceCounting()
139         for (Operation *user : info.users) {  in optimizeReferenceCounting()
140           // `user` operation lies after `addRef` ...  in optimizeReferenceCounting()
141           if (user == addRef || user->isBeforeInBlock(addRef))  in optimizeReferenceCounting()
144           if (user == dropRef || dropRef->isBeforeInBlock(user))  in optimizeReferenceCounting()
147           // Find the first function call user of the reference counted value.  in optimizeReferenceCounting()
148           Operation *functionCall = dyn_cast<func::CallOp>(user);  in optimizeReferenceCounting()
151                functionCall->isBeforeInBlock(firstFunctionCallUser))) {  in optimizeReferenceCounting()
156           // Find the last regular user of the reference counted value.  in optimizeReferenceCounting()
159                lastNonFunctionCallUser->isBeforeInBlock(user))) {  in optimizeReferenceCounting()
160             lastNonFunctionCallUser = user;  in optimizeReferenceCounting()
165         // Non function call user after the function call user of the reference  in optimizeReferenceCounting()
168             firstFunctionCallUser->isBeforeInBlock(lastNonFunctionCallUser))  in optimizeReferenceCounting()
178         if (emplaced.second) // successfully cancelled `add_ref` <-> `drop_ref`  in optimizeReferenceCounting()
197   WalkResult blockWalk = op->walk([&](Block *block) -> WalkResult {  in runOnOperation()
198     for (BlockArgument arg : block->getArguments())  in runOnOperation()
210   WalkResult opWalk = op->walk([&](Operation *op) -> WalkResult {  in runOnOperation()
211     for (unsigned i = 0; i < op->getNumResults(); ++i)  in runOnOperation()
212       if (isRefCounted(op->getResultTypes()[i]))  in runOnOperation()
213         if (failed(optimizeReferenceCounting(op->getResult(i), cancellable)))  in runOnOperation()
227   // Erase all cancellable `add_ref <-> drop_ref` operation pairs.  in runOnOperation()
229     kv.first->erase();  in runOnOperation()
230     kv.second->erase();  in runOnOperation()