/Users/buildslave/jenkins/workspace/clang-stage2-coverage-R/llvm/lib/Transforms/Coroutines/CoroElide.cpp
Line | Count | Source (jump to first uncovered line) |
1 | | //===- CoroElide.cpp - Coroutine Frame Allocation Elision Pass ------------===// |
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 | | // This pass replaces dynamic allocation of coroutine frame with alloca and |
9 | | // replaces calls to llvm.coro.resume and llvm.coro.destroy with direct calls |
10 | | // to coroutine sub-functions. |
11 | | //===----------------------------------------------------------------------===// |
12 | | |
13 | | #include "CoroInternal.h" |
14 | | #include "llvm/Analysis/AliasAnalysis.h" |
15 | | #include "llvm/Analysis/InstructionSimplify.h" |
16 | | #include "llvm/IR/Dominators.h" |
17 | | #include "llvm/IR/InstIterator.h" |
18 | | #include "llvm/Pass.h" |
19 | | #include "llvm/Support/ErrorHandling.h" |
20 | | |
21 | | using namespace llvm; |
22 | | |
23 | | #define DEBUG_TYPE "coro-elide" |
24 | | |
25 | | namespace { |
26 | | // Created on demand if CoroElide pass has work to do. |
27 | | struct Lowerer : coro::LowererBase { |
28 | | SmallVector<CoroIdInst *, 4> CoroIds; |
29 | | SmallVector<CoroBeginInst *, 1> CoroBegins; |
30 | | SmallVector<CoroAllocInst *, 1> CoroAllocs; |
31 | | SmallVector<CoroSubFnInst *, 4> ResumeAddr; |
32 | | SmallVector<CoroSubFnInst *, 4> DestroyAddr; |
33 | | SmallVector<CoroFreeInst *, 1> CoroFrees; |
34 | | |
35 | 12 | Lowerer(Module &M) : LowererBase(M) {} |
36 | | |
37 | | void elideHeapAllocations(Function *F, Type *FrameTy, AAResults &AA); |
38 | | bool shouldElide(Function *F, DominatorTree &DT) const; |
39 | | bool processCoroId(CoroIdInst *, AAResults &AA, DominatorTree &DT); |
40 | | }; |
41 | | } // end anonymous namespace |
42 | | |
43 | | // Go through the list of coro.subfn.addr intrinsics and replace them with the |
44 | | // provided constant. |
45 | | static void replaceWithConstant(Constant *Value, |
46 | 56 | SmallVectorImpl<CoroSubFnInst *> &Users) { |
47 | 56 | if (Users.empty()) |
48 | 21 | return; |
49 | 35 | |
50 | 35 | // See if we need to bitcast the constant to match the type of the intrinsic |
51 | 35 | // being replaced. Note: All coro.subfn.addr intrinsics return the same type, |
52 | 35 | // so we only need to examine the type of the first one in the list. |
53 | 35 | Type *IntrTy = Users.front()->getType(); |
54 | 35 | Type *ValueTy = Value->getType(); |
55 | 35 | if (ValueTy != IntrTy) { |
56 | 35 | // May need to tweak the function type to match the type expected at the |
57 | 35 | // use site. |
58 | 35 | assert(ValueTy->isPointerTy() && IntrTy->isPointerTy()); |
59 | 35 | Value = ConstantExpr::getBitCast(Value, IntrTy); |
60 | 35 | } |
61 | 35 | |
62 | 35 | // Now the value type matches the type of the intrinsic. Replace them all! |
63 | 35 | for (CoroSubFnInst *I : Users) |
64 | 41 | replaceAndRecursivelySimplify(I, Value); |
65 | 35 | } |
66 | | |
67 | | // See if any operand of the call instruction references the coroutine frame. |
68 | 16 | static bool operandReferences(CallInst *CI, AllocaInst *Frame, AAResults &AA) { |
69 | 16 | for (Value *Op : CI->operand_values()) |
70 | 31 | if (AA.alias(Op, Frame) != NoAlias) |
71 | 7 | return true; |
72 | 16 | return false9 ; |
73 | 16 | } |
74 | | |
75 | | // Look for any tail calls referencing the coroutine frame and remove tail |
76 | | // attribute from them, since now coroutine frame resides on the stack and tail |
77 | | // call implies that the function does not references anything on the stack. |
78 | 3 | static void removeTailCallAttribute(AllocaInst *Frame, AAResults &AA) { |
79 | 3 | Function &F = *Frame->getFunction(); |
80 | 3 | for (Instruction &I : instructions(F)) |
81 | 75 | if (auto *Call = dyn_cast<CallInst>(&I)) |
82 | 22 | if (Call->isTailCall() && operandReferences(Call, Frame, AA)16 ) { |
83 | 7 | // FIXME: If we ever hit this check. Evaluate whether it is more |
84 | 7 | // appropriate to retain musttail and allow the code to compile. |
85 | 7 | if (Call->isMustTailCall()) |
86 | 0 | report_fatal_error("Call referring to the coroutine frame cannot be " |
87 | 0 | "marked as musttail"); |
88 | 7 | Call->setTailCall(false); |
89 | 7 | } |
90 | 3 | } |
91 | | |
92 | | // Given a resume function @f.resume(%f.frame* %frame), returns %f.frame type. |
93 | 3 | static Type *getFrameType(Function *Resume) { |
94 | 3 | auto *ArgType = Resume->arg_begin()->getType(); |
95 | 3 | return cast<PointerType>(ArgType)->getElementType(); |
96 | 3 | } |
97 | | |
98 | | // Finds first non alloca instruction in the entry block of a function. |
99 | 3 | static Instruction *getFirstNonAllocaInTheEntryBlock(Function *F) { |
100 | 3 | for (Instruction &I : F->getEntryBlock()) |
101 | 3 | if (!isa<AllocaInst>(&I)) |
102 | 3 | return &I; |
103 | 3 | llvm_unreachable0 ("no terminator in the entry block"); |
104 | 3 | } |
105 | | |
106 | | // To elide heap allocations we need to suppress code blocks guarded by |
107 | | // llvm.coro.alloc and llvm.coro.free instructions. |
108 | 3 | void Lowerer::elideHeapAllocations(Function *F, Type *FrameTy, AAResults &AA) { |
109 | 3 | LLVMContext &C = FrameTy->getContext(); |
110 | 3 | auto *InsertPt = |
111 | 3 | getFirstNonAllocaInTheEntryBlock(CoroIds.front()->getFunction()); |
112 | 3 | |
113 | 3 | // Replacing llvm.coro.alloc with false will suppress dynamic |
114 | 3 | // allocation as it is expected for the frontend to generate the code that |
115 | 3 | // looks like: |
116 | 3 | // id = coro.id(...) |
117 | 3 | // mem = coro.alloc(id) ? malloc(coro.size()) : 0; |
118 | 3 | // coro.begin(id, mem) |
119 | 3 | auto *False = ConstantInt::getFalse(C); |
120 | 3 | for (auto *CA : CoroAllocs) { |
121 | 3 | CA->replaceAllUsesWith(False); |
122 | 3 | CA->eraseFromParent(); |
123 | 3 | } |
124 | 3 | |
125 | 3 | // FIXME: Design how to transmit alignment information for every alloca that |
126 | 3 | // is spilled into the coroutine frame and recreate the alignment information |
127 | 3 | // here. Possibly we will need to do a mini SROA here and break the coroutine |
128 | 3 | // frame into individual AllocaInst recreating the original alignment. |
129 | 3 | const DataLayout &DL = F->getParent()->getDataLayout(); |
130 | 3 | auto *Frame = new AllocaInst(FrameTy, DL.getAllocaAddrSpace(), "", InsertPt); |
131 | 3 | auto *FrameVoidPtr = |
132 | 3 | new BitCastInst(Frame, Type::getInt8PtrTy(C), "vFrame", InsertPt); |
133 | 3 | |
134 | 3 | for (auto *CB : CoroBegins) { |
135 | 3 | CB->replaceAllUsesWith(FrameVoidPtr); |
136 | 3 | CB->eraseFromParent(); |
137 | 3 | } |
138 | 3 | |
139 | 3 | // Since now coroutine frame lives on the stack we need to make sure that |
140 | 3 | // any tail call referencing it, must be made non-tail call. |
141 | 3 | removeTailCallAttribute(Frame, AA); |
142 | 3 | } |
143 | | |
144 | 23 | bool Lowerer::shouldElide(Function *F, DominatorTree &DT) const { |
145 | 23 | // If no CoroAllocs, we cannot suppress allocation, so elision is not |
146 | 23 | // possible. |
147 | 23 | if (CoroAllocs.empty()) |
148 | 17 | return false; |
149 | 6 | |
150 | 6 | // Check that for every coro.begin there is a coro.destroy directly |
151 | 6 | // referencing the SSA value of that coro.begin along a non-exceptional path. |
152 | 6 | // If the value escaped, then coro.destroy would have been referencing a |
153 | 6 | // memory location storing that value and not the virtual register. |
154 | 6 | |
155 | 6 | // First gather all of the non-exceptional terminators for the function. |
156 | 6 | SmallPtrSet<Instruction *, 8> Terminators; |
157 | 29 | for (BasicBlock &B : *F) { |
158 | 29 | auto *TI = B.getTerminator(); |
159 | 29 | if (TI->getNumSuccessors() == 0 && !TI->isExceptionalTerminator()8 && |
160 | 29 | !isa<UnreachableInst>(TI)6 ) |
161 | 6 | Terminators.insert(TI); |
162 | 29 | } |
163 | 6 | |
164 | 6 | // Filter out the coro.destroy that lie along exceptional paths. |
165 | 6 | SmallPtrSet<CoroSubFnInst *, 4> DAs; |
166 | 6 | for (CoroSubFnInst *DA : DestroyAddr) { |
167 | 4 | for (Instruction *TI : Terminators) { |
168 | 4 | if (DT.dominates(DA, TI)) { |
169 | 3 | DAs.insert(DA); |
170 | 3 | break; |
171 | 3 | } |
172 | 4 | } |
173 | 4 | } |
174 | 6 | |
175 | 6 | // Find all the coro.begin referenced by coro.destroy along happy paths. |
176 | 6 | SmallPtrSet<CoroBeginInst *, 8> ReferencedCoroBegins; |
177 | 6 | for (CoroSubFnInst *DA : DAs) { |
178 | 3 | if (auto *CB = dyn_cast<CoroBeginInst>(DA->getFrame())) |
179 | 3 | ReferencedCoroBegins.insert(CB); |
180 | 0 | else |
181 | 0 | return false; |
182 | 3 | } |
183 | 6 | |
184 | 6 | // If size of the set is the same as total number of coro.begin, that means we |
185 | 6 | // found a coro.free or coro.destroy referencing each coro.begin, so we can |
186 | 6 | // perform heap elision. |
187 | 6 | return ReferencedCoroBegins.size() == CoroBegins.size(); |
188 | 6 | } |
189 | | |
190 | | bool Lowerer::processCoroId(CoroIdInst *CoroId, AAResults &AA, |
191 | 23 | DominatorTree &DT) { |
192 | 23 | CoroBegins.clear(); |
193 | 23 | CoroAllocs.clear(); |
194 | 23 | CoroFrees.clear(); |
195 | 23 | ResumeAddr.clear(); |
196 | 23 | DestroyAddr.clear(); |
197 | 23 | |
198 | 23 | // Collect all coro.begin and coro.allocs associated with this coro.id. |
199 | 31 | for (User *U : CoroId->users()) { |
200 | 31 | if (auto *CB = dyn_cast<CoroBeginInst>(U)) |
201 | 23 | CoroBegins.push_back(CB); |
202 | 8 | else if (auto *CA = dyn_cast<CoroAllocInst>(U)) |
203 | 6 | CoroAllocs.push_back(CA); |
204 | 2 | else if (auto *CF = dyn_cast<CoroFreeInst>(U)) |
205 | 2 | CoroFrees.push_back(CF); |
206 | 31 | } |
207 | 23 | |
208 | 23 | // Collect all coro.subfn.addrs associated with coro.begin. |
209 | 23 | // Note, we only devirtualize the calls if their coro.subfn.addr refers to |
210 | 23 | // coro.begin directly. If we run into cases where this check is too |
211 | 23 | // conservative, we can consider relaxing the check. |
212 | 23 | for (CoroBeginInst *CB : CoroBegins) { |
213 | 23 | for (User *U : CB->users()) |
214 | 119 | if (auto *II = dyn_cast<CoroSubFnInst>(U)) |
215 | 31 | switch (II->getIndex()) { |
216 | 31 | case CoroSubFnInst::ResumeIndex: |
217 | 20 | ResumeAddr.push_back(II); |
218 | 20 | break; |
219 | 31 | case CoroSubFnInst::DestroyIndex: |
220 | 11 | DestroyAddr.push_back(II); |
221 | 11 | break; |
222 | 31 | default: |
223 | 0 | llvm_unreachable("unexpected coro.subfn.addr constant"); |
224 | 31 | } |
225 | 23 | } |
226 | 23 | |
227 | 23 | // PostSplit coro.id refers to an array of subfunctions in its Info |
228 | 23 | // argument. |
229 | 23 | ConstantArray *Resumers = CoroId->getInfo().Resumers; |
230 | 23 | assert(Resumers && "PostSplit coro.id Info argument must refer to an array" |
231 | 23 | "of coroutine subfunctions"); |
232 | 23 | auto *ResumeAddrConstant = |
233 | 23 | ConstantExpr::getExtractValue(Resumers, CoroSubFnInst::ResumeIndex); |
234 | 23 | |
235 | 23 | replaceWithConstant(ResumeAddrConstant, ResumeAddr); |
236 | 23 | |
237 | 23 | bool ShouldElide = shouldElide(CoroId->getFunction(), DT); |
238 | 23 | |
239 | 23 | auto *DestroyAddrConstant = ConstantExpr::getExtractValue( |
240 | 23 | Resumers, |
241 | 23 | ShouldElide ? CoroSubFnInst::CleanupIndex3 : CoroSubFnInst::DestroyIndex20 ); |
242 | 23 | |
243 | 23 | replaceWithConstant(DestroyAddrConstant, DestroyAddr); |
244 | 23 | |
245 | 23 | if (ShouldElide) { |
246 | 3 | auto *FrameTy = getFrameType(cast<Function>(ResumeAddrConstant)); |
247 | 3 | elideHeapAllocations(CoroId->getFunction(), FrameTy, AA); |
248 | 3 | coro::replaceCoroFree(CoroId, /*Elide=*/true); |
249 | 3 | } |
250 | 23 | |
251 | 23 | return true; |
252 | 23 | } |
253 | | |
254 | | // See if there are any coro.subfn.addr instructions referring to coro.devirt |
255 | | // trigger, if so, replace them with a direct call to devirt trigger function. |
256 | 10 | static bool replaceDevirtTrigger(Function &F) { |
257 | 10 | SmallVector<CoroSubFnInst *, 1> DevirtAddr; |
258 | 10 | for (auto &I : instructions(F)) |
259 | 313 | if (auto *SubFn = dyn_cast<CoroSubFnInst>(&I)) |
260 | 10 | if (SubFn->getIndex() == CoroSubFnInst::RestartTrigger) |
261 | 10 | DevirtAddr.push_back(SubFn); |
262 | 10 | |
263 | 10 | if (DevirtAddr.empty()) |
264 | 0 | return false; |
265 | 10 | |
266 | 10 | Module &M = *F.getParent(); |
267 | 10 | Function *DevirtFn = M.getFunction(CORO_DEVIRT_TRIGGER_FN); |
268 | 10 | assert(DevirtFn && "coro.devirt.fn not found"); |
269 | 10 | replaceWithConstant(DevirtFn, DevirtAddr); |
270 | 10 | |
271 | 10 | return true; |
272 | 10 | } |
273 | | |
274 | | //===----------------------------------------------------------------------===// |
275 | | // Top Level Driver |
276 | | //===----------------------------------------------------------------------===// |
277 | | |
278 | | namespace { |
279 | | struct CoroElide : FunctionPass { |
280 | | static char ID; |
281 | 39 | CoroElide() : FunctionPass(ID) { |
282 | 39 | initializeCoroElidePass(*PassRegistry::getPassRegistry()); |
283 | 39 | } |
284 | | |
285 | | std::unique_ptr<Lowerer> L; |
286 | | |
287 | 39 | bool doInitialization(Module &M) override { |
288 | 39 | if (coro::declaresIntrinsics(M, {"llvm.coro.id"})) |
289 | 12 | L = llvm::make_unique<Lowerer>(M); |
290 | 39 | return false; |
291 | 39 | } |
292 | | |
293 | 250 | bool runOnFunction(Function &F) override { |
294 | 250 | if (!L) |
295 | 140 | return false; |
296 | 110 | |
297 | 110 | bool Changed = false; |
298 | 110 | |
299 | 110 | if (F.hasFnAttribute(CORO_PRESPLIT_ATTR)) |
300 | 10 | Changed = replaceDevirtTrigger(F); |
301 | 110 | |
302 | 110 | L->CoroIds.clear(); |
303 | 110 | |
304 | 110 | // Collect all PostSplit coro.ids. |
305 | 110 | for (auto &I : instructions(F)) |
306 | 1.63k | if (auto *CII = dyn_cast<CoroIdInst>(&I)) |
307 | 47 | if (CII->getInfo().isPostSplit()) |
308 | 36 | // If it is the coroutine itself, don't touch it. |
309 | 36 | if (CII->getCoroutine() != CII->getFunction()) |
310 | 23 | L->CoroIds.push_back(CII); |
311 | 110 | |
312 | 110 | // If we did not find any coro.id, there is nothing to do. |
313 | 110 | if (L->CoroIds.empty()) |
314 | 87 | return Changed; |
315 | 23 | |
316 | 23 | AAResults &AA = getAnalysis<AAResultsWrapperPass>().getAAResults(); |
317 | 23 | DominatorTree &DT = getAnalysis<DominatorTreeWrapperPass>().getDomTree(); |
318 | 23 | |
319 | 23 | for (auto *CII : L->CoroIds) |
320 | 23 | Changed |= L->processCoroId(CII, AA, DT); |
321 | 23 | |
322 | 23 | return Changed; |
323 | 23 | } |
324 | 39 | void getAnalysisUsage(AnalysisUsage &AU) const override { |
325 | 39 | AU.addRequired<AAResultsWrapperPass>(); |
326 | 39 | AU.addRequired<DominatorTreeWrapperPass>(); |
327 | 39 | } |
328 | 250 | StringRef getPassName() const override { return "Coroutine Elision"; } |
329 | | }; |
330 | | } |
331 | | |
332 | | char CoroElide::ID = 0; |
333 | 11.0k | INITIALIZE_PASS_BEGIN( |
334 | 11.0k | CoroElide, "coro-elide", |
335 | 11.0k | "Coroutine frame allocation elision and indirect calls replacement", false, |
336 | 11.0k | false) |
337 | 11.0k | INITIALIZE_PASS_DEPENDENCY(AAResultsWrapperPass) |
338 | 11.0k | INITIALIZE_PASS_END( |
339 | | CoroElide, "coro-elide", |
340 | | "Coroutine frame allocation elision and indirect calls replacement", false, |
341 | | false) |
342 | | |
343 | 36 | Pass *llvm::createCoroElidePass() { return new CoroElide(); } |