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
9#include "llvm/Transforms/Coroutines/CoroElide.h"
10#include "CoroInternal.h"
11#include "llvm/ADT/DenseMap.h"
12#include "llvm/ADT/Statistic.h"
13#include "llvm/Analysis/AliasAnalysis.h"
14#include "llvm/Analysis/InstructionSimplify.h"
15#include "llvm/Analysis/OptimizationRemarkEmitter.h"
16#include "llvm/IR/Dominators.h"
17#include "llvm/IR/InstIterator.h"
18#include "llvm/Support/ErrorHandling.h"
19#include "llvm/Support/FileSystem.h"
20#include <optional>
21
22using namespace llvm;
23
24#define DEBUG_TYPE "coro-elide"
25
26STATISTIC(NumOfCoroElided, "The # of coroutine get elided.");
27
28#ifndef NDEBUG
29static cl::opt<std::string> CoroElideInfoOutputFilename(
30 "coro-elide-info-output-file", cl::value_desc("filename"),
31 cl::desc("File to record the coroutines got elided"), cl::Hidden);
32#endif
33
34namespace {
35// Created on demand if the coro-elide pass has work to do.
36class FunctionElideInfo {
37public:
38 FunctionElideInfo(Function *F) : ContainingFunction(F) {
39 this->collectPostSplitCoroIds();
40 }
41
42 bool hasCoroIds() const { return !CoroIds.empty(); }
43
44 const SmallVectorImpl<CoroIdInst *> &getCoroIds() const { return CoroIds; }
45
46private:
47 Function *ContainingFunction;
48 SmallVector<CoroIdInst *, 4> CoroIds;
49 // Used in canCoroBeginEscape to distinguish coro.suspend switchs.
50 SmallPtrSet<const SwitchInst *, 4> CoroSuspendSwitches;
51
52 void collectPostSplitCoroIds();
53 friend class CoroIdElider;
54};
55
56class CoroIdElider {
57public:
58 CoroIdElider(CoroIdInst *CoroId, FunctionElideInfo &FEI, AAResults &AA,
59 DominatorTree &DT, OptimizationRemarkEmitter &ORE);
60 void elideHeapAllocations(uint64_t FrameSize, Align FrameAlign);
61 bool lifetimeEligibleForElide() const;
62 bool attemptElide();
63 bool canCoroBeginEscape(const CoroBeginInst *,
64 const SmallPtrSetImpl<BasicBlock *> &) const;
65
66private:
67 CoroIdInst *CoroId;
68 FunctionElideInfo &FEI;
69 AAResults &AA;
70 DominatorTree &DT;
71 OptimizationRemarkEmitter &ORE;
72
73 SmallVector<CoroBeginInst *, 1> CoroBegins;
74 SmallVector<CoroAllocInst *, 1> CoroAllocs;
75 SmallVector<CoroSubFnInst *, 4> ResumeAddr;
76 SmallVector<CoroSubFnInst *, 4> DestroyAddr;
77 DenseMap<CoroBeginInst *, SmallVector<IntrinsicInst *, 4>> BeginDeadMap;
78};
79} // end anonymous namespace
80
81// Go through the list of coro.subfn.addr intrinsics and replace them with the
82// provided constant.
83static void replaceWithConstant(Constant *Value,
84 SmallVectorImpl<CoroSubFnInst *> &Users) {
85 for (CoroSubFnInst *I : Users)
86 replaceAndRecursivelySimplify(I, SimpleV: Value);
87}
88
89// See if any operand of the call instruction references the coroutine frame.
90static bool operandReferences(CallInst *CI, AllocaInst *Frame, AAResults &AA) {
91 for (Value *Op : CI->operand_values())
92 if (Op->getType()->isPointerTy() && !AA.isNoAlias(V1: Op, V2: Frame))
93 return true;
94 return false;
95}
96
97// Look for any tail calls referencing the coroutine frame and remove tail
98// attribute from them, since now coroutine frame resides on the stack and tail
99// call implies that the function does not references anything on the stack.
100// However if it's a musttail call, we cannot remove the tailcall attribute.
101// It's safe to keep it there as the musttail call is for symmetric transfer,
102// and by that point the frame should have been destroyed and hence not
103// interfering with operands.
104static void removeTailCallAttribute(AllocaInst *Frame, AAResults &AA) {
105 Function &F = *Frame->getFunction();
106 for (Instruction &I : instructions(F))
107 if (auto *Call = dyn_cast<CallInst>(Val: &I))
108 if (Call->isTailCall() && operandReferences(CI: Call, Frame, AA) &&
109 !Call->isMustTailCall())
110 Call->setTailCall(false);
111}
112
113// Given a resume function @f.resume(%f.frame* %frame), returns the size
114// and expected alignment of %f.frame type.
115static std::optional<std::pair<uint64_t, Align>>
116getFrameLayout(Function *Resume) {
117 // Pull information from the function attributes.
118 auto Size = Resume->getParamDereferenceableBytes(ArgNo: 0);
119 if (!Size)
120 return std::nullopt;
121 return std::make_pair(x&: Size, y: Resume->getParamAlign(ArgNo: 0).valueOrOne());
122}
123
124// Finds first non alloca instruction in the entry block of a function.
125static Instruction *getFirstNonAllocaInTheEntryBlock(Function *F) {
126 for (Instruction &I : F->getEntryBlock())
127 if (!isa<AllocaInst>(Val: &I))
128 return &I;
129 llvm_unreachable("no terminator in the entry block");
130}
131
132#ifndef NDEBUG
133static std::unique_ptr<raw_fd_ostream> getOrCreateLogFile() {
134 assert(!CoroElideInfoOutputFilename.empty() &&
135 "coro-elide-info-output-file shouldn't be empty");
136 std::error_code EC;
137 auto Result = std::make_unique<raw_fd_ostream>(CoroElideInfoOutputFilename,
138 EC, sys::fs::OF_Append);
139 if (!EC)
140 return Result;
141 llvm::errs() << "Error opening coro-elide-info-output-file '"
142 << CoroElideInfoOutputFilename << " for appending!\n";
143 return std::make_unique<raw_fd_ostream>(2, false); // stderr.
144}
145#endif
146
147void FunctionElideInfo::collectPostSplitCoroIds() {
148 for (auto &I : instructions(F: this->ContainingFunction)) {
149 if (auto *CII = dyn_cast<CoroIdInst>(Val: &I))
150 if (CII->getInfo().isPostSplit())
151 CoroIds.push_back(Elt: CII);
152
153 // Consider case like:
154 // %0 = call i8 @llvm.coro.suspend(...)
155 // switch i8 %0, label %suspend [i8 0, label %resume
156 // i8 1, label %cleanup]
157 // and collect the SwitchInsts which are used by escape analysis later.
158 if (auto *CSI = dyn_cast<CoroSuspendInst>(Val: &I))
159 if (CSI->hasOneUse() && isa<SwitchInst>(Val: CSI->use_begin()->getUser())) {
160 SwitchInst *SWI = cast<SwitchInst>(Val: CSI->use_begin()->getUser());
161 if (SWI->getNumCases() == 2)
162 CoroSuspendSwitches.insert(Ptr: SWI);
163 }
164 }
165}
166
167CoroIdElider::CoroIdElider(CoroIdInst *CoroId, FunctionElideInfo &FEI,
168 AAResults &AA, DominatorTree &DT,
169 OptimizationRemarkEmitter &ORE)
170 : CoroId(CoroId), FEI(FEI), AA(AA), DT(DT), ORE(ORE) {
171 // Collect all coro.begin and coro.allocs associated with this coro.id.
172 for (User *U : CoroId->users()) {
173 if (auto *CB = dyn_cast<CoroBeginInst>(Val: U))
174 CoroBegins.push_back(Elt: CB);
175 else if (auto *CA = dyn_cast<CoroAllocInst>(Val: U))
176 CoroAllocs.push_back(Elt: CA);
177 }
178
179 for (CoroBeginInst *CB : CoroBegins) {
180 for (User *U : CB->users()) {
181 auto &CoroDeads = BeginDeadMap[CB];
182 // Collect all coro.subfn.addrs associated with coro.begin.
183 // Note, we only devirtualize the calls if their coro.subfn.addr refers to
184 // coro.begin directly. If we run into cases where this check is too
185 // conservative, we can consider relaxing the check.
186 if (auto *II = dyn_cast<CoroSubFnInst>(Val: U)) {
187 switch (II->getIndex()) {
188 case CoroSubFnInst::ResumeIndex:
189 ResumeAddr.push_back(Elt: II);
190 break;
191 case CoroSubFnInst::DestroyIndex:
192 CoroDeads.push_back(Elt: II); // coro.destroy implies coro.dead
193 DestroyAddr.push_back(Elt: II);
194 break;
195 default:
196 llvm_unreachable("unexpected coro.subfn.addr constant");
197 }
198 } else if (auto *II = dyn_cast<CoroDeadInst>(Val: U))
199 CoroDeads.push_back(Elt: II);
200 }
201 }
202}
203
204// To elide heap allocations we need to suppress code blocks guarded by
205// llvm.coro.alloc and llvm.coro.free instructions.
206void CoroIdElider::elideHeapAllocations(uint64_t FrameSize, Align FrameAlign) {
207 LLVMContext &C = FEI.ContainingFunction->getContext();
208 BasicBlock::iterator InsertPt =
209 getFirstNonAllocaInTheEntryBlock(F: FEI.ContainingFunction)->getIterator();
210
211 // Replacing llvm.coro.alloc with false will suppress dynamic
212 // allocation as it is expected for the frontend to generate the code that
213 // looks like:
214 // id = coro.id(...)
215 // mem = coro.alloc(id) ? malloc(coro.size()) : 0;
216 // coro.begin(id, mem)
217 auto *False = ConstantInt::getFalse(Context&: C);
218 for (auto *CA : CoroAllocs) {
219 CA->replaceAllUsesWith(V: False);
220 CA->eraseFromParent();
221 }
222
223 // FIXME: Design how to transmit alignment information for every alloca that
224 // is spilled into the coroutine frame and recreate the alignment information
225 // here. Possibly we will need to do a mini SROA here and break the coroutine
226 // frame into individual AllocaInst recreating the original alignment.
227 const DataLayout &DL = FEI.ContainingFunction->getDataLayout();
228 auto FrameTy = ArrayType::get(ElementType: Type::getInt8Ty(C), NumElements: FrameSize);
229 auto *Frame = new AllocaInst(FrameTy, DL.getAllocaAddrSpace(), "", InsertPt);
230 Frame->setAlignment(FrameAlign);
231 auto *FrameVoidPtr =
232 new BitCastInst(Frame, PointerType::getUnqual(C), "vFrame", InsertPt);
233
234 for (auto *CB : CoroBegins) {
235 coro::elideCoroFree(FramePtr: CB);
236 CB->replaceAllUsesWith(V: FrameVoidPtr);
237 CB->eraseFromParent();
238 }
239
240 // Since now coroutine frame lives on the stack we need to make sure that
241 // any tail call referencing it, must be made non-tail call.
242 removeTailCallAttribute(Frame, AA);
243}
244
245bool CoroIdElider::canCoroBeginEscape(
246 const CoroBeginInst *CB, const SmallPtrSetImpl<BasicBlock *> &TIs) const {
247 const auto &It = BeginDeadMap.find(Val: CB);
248 assert(It != BeginDeadMap.end());
249
250 // Limit the number of blocks we visit.
251 unsigned Limit = 32 * (1 + It->second.size());
252
253 SmallVector<const BasicBlock *, 32> Worklist;
254 Worklist.push_back(Elt: CB->getParent());
255
256 SmallPtrSet<const BasicBlock *, 32> Visited;
257 // Consider basicblock of coro.dead/destroy as visited one, so that we
258 // skip the path pass through it.
259 for (auto *DA : It->second)
260 Visited.insert(Ptr: DA->getParent());
261
262 SmallPtrSet<const BasicBlock *, 32> EscapingBBs;
263 for (auto *U : CB->users()) {
264 // The use from coroutine intrinsics are not a problem.
265 if (isa<CoroFreeInst, CoroSubFnInst, CoroSaveInst>(Val: U))
266 continue;
267
268 // Think all other usages may be an escaping candidate conservatively.
269 //
270 // Note that the major user of switch ABI coroutine (the C++) will store
271 // resume.fn, destroy.fn and the index to the coroutine frame immediately.
272 // So the parent of the coro.begin in C++ will be always escaping.
273 // Then we can't get any performance benefits for C++ by improving the
274 // precision of the method.
275 //
276 // The reason why we still judge it is we want to make LLVM Coroutine in
277 // switch ABIs to be self contained as much as possible instead of a
278 // by-product of C++20 Coroutines.
279 EscapingBBs.insert(Ptr: cast<Instruction>(Val: U)->getParent());
280 }
281
282 bool PotentiallyEscaped = false;
283
284 do {
285 const auto *BB = Worklist.pop_back_val();
286 if (!Visited.insert(Ptr: BB).second)
287 continue;
288
289 // A Path insensitive marker to test whether the coro.begin escapes.
290 // It is intentional to make it path insensitive while it may not be
291 // precise since we don't want the process to be too slow.
292 PotentiallyEscaped |= EscapingBBs.count(Ptr: BB);
293
294 if (TIs.count(Ptr: BB)) {
295 if (isa<ReturnInst>(Val: BB->getTerminator()) || PotentiallyEscaped)
296 return true;
297
298 // If the function ends with the exceptional terminator, the memory used
299 // by the coroutine frame can be released by stack unwinding
300 // automatically. So we can think the coro.begin doesn't escape if it
301 // exits the function by exceptional terminator.
302
303 continue;
304 }
305
306 // Conservatively say that there is potentially a path.
307 if (!--Limit)
308 return true;
309
310 auto TI = BB->getTerminator();
311 // Although the default dest of coro.suspend switches is suspend pointer
312 // which means a escape path to normal terminator, it is reasonable to skip
313 // it since coroutine frame doesn't change outside the coroutine body.
314 if (isa<SwitchInst>(Val: TI) &&
315 FEI.CoroSuspendSwitches.count(Ptr: cast<SwitchInst>(Val: TI))) {
316 Worklist.push_back(Elt: cast<SwitchInst>(Val: TI)->getSuccessor(idx: 1));
317 Worklist.push_back(Elt: cast<SwitchInst>(Val: TI)->getSuccessor(idx: 2));
318 } else
319 Worklist.append(in_start: succ_begin(BB), in_end: succ_end(BB));
320
321 } while (!Worklist.empty());
322
323 // We have exhausted all possible paths and are certain that coro.begin can
324 // not reach to any of terminators.
325 return false;
326}
327
328bool CoroIdElider::lifetimeEligibleForElide() const {
329 // If no CoroAllocs, we cannot suppress allocation, so elision is not
330 // possible.
331 if (CoroAllocs.empty())
332 return false;
333
334 // Check that for every coro.begin there is at least one coro.dead/destroy
335 // directly referencing the SSA value of that coro.begin along each
336 // non-exceptional path.
337 //
338 // If the value escaped, then coro.dead/destroy would have been referencing a
339 // memory location storing that value and not the virtual register.
340
341 SmallPtrSet<BasicBlock *, 8> Terminators;
342 // First gather all of the terminators for the function.
343 // Consider the final coro.suspend as the real terminator when the current
344 // function is a coroutine.
345 for (BasicBlock &B : *FEI.ContainingFunction) {
346 auto *TI = B.getTerminator();
347
348 if (TI->getNumSuccessors() != 0 || isa<UnreachableInst>(Val: TI))
349 continue;
350
351 Terminators.insert(Ptr: &B);
352 }
353
354 // Filter out the coro.dead/destroy that lie along exceptional paths.
355 for (const auto *CB : CoroBegins) {
356 auto It = BeginDeadMap.find(Val: CB);
357 if (It == BeginDeadMap.end())
358 return false;
359
360 // If every terminators is dominated by coro.dead/destroy, we could know the
361 // corresponding coro.begin wouldn't escape.
362 auto DominatesTerminator = [&](auto *TI) {
363 return llvm::any_of(It->second, [&](auto *Destroy) {
364 return DT.dominates(Destroy, TI->getTerminator());
365 });
366 };
367
368 if (llvm::all_of(Range&: Terminators, P: DominatesTerminator))
369 continue;
370
371 // Otherwise canCoroBeginEscape would decide whether there is any paths from
372 // coro.begin to Terminators which not pass through any of the
373 // coro.dead/destroy. This is a slower analysis.
374 //
375 // canCoroBeginEscape is relatively slow, so we avoid to run it as much as
376 // possible.
377 if (canCoroBeginEscape(CB, TIs: Terminators))
378 return false;
379 }
380
381 // We have checked all CoroBegins and their paths to the terminators without
382 // finding disqualifying code patterns, so we can perform heap allocations.
383 return true;
384}
385
386bool CoroIdElider::attemptElide() {
387 // PostSplit coro.id refers to an array of subfunctions in its Info
388 // argument.
389 ConstantArray *Resumers = CoroId->getInfo().Resumers;
390 assert(Resumers && "PostSplit coro.id Info argument must refer to an array"
391 "of coroutine subfunctions");
392 auto *ResumeAddrConstant =
393 Resumers->getAggregateElement(Elt: CoroSubFnInst::ResumeIndex);
394
395 replaceWithConstant(Value: ResumeAddrConstant, Users&: ResumeAddr);
396
397 bool EligibleForElide = lifetimeEligibleForElide();
398
399 auto *DestroyAddrConstant = Resumers->getAggregateElement(
400 Elt: EligibleForElide ? CoroSubFnInst::CleanupIndex
401 : CoroSubFnInst::DestroyIndex);
402
403 replaceWithConstant(Value: DestroyAddrConstant, Users&: DestroyAddr);
404
405 auto FrameSizeAndAlign = getFrameLayout(Resume: cast<Function>(Val: ResumeAddrConstant));
406
407 auto CallerFunctionName = FEI.ContainingFunction->getName();
408 auto CalleeCoroutineName = CoroId->getCoroutine()->getName();
409
410 if (EligibleForElide && FrameSizeAndAlign) {
411 elideHeapAllocations(FrameSize: FrameSizeAndAlign->first, FrameAlign: FrameSizeAndAlign->second);
412 NumOfCoroElided++;
413
414#ifndef NDEBUG
415 if (!CoroElideInfoOutputFilename.empty())
416 *getOrCreateLogFile() << "Elide " << CalleeCoroutineName << " in "
417 << FEI.ContainingFunction->getName() << "\n";
418#endif
419
420 ORE.emit(RemarkBuilder: [&]() {
421 return OptimizationRemark(DEBUG_TYPE, "CoroElide", CoroId)
422 << "'" << ore::NV("callee", CalleeCoroutineName)
423 << "' elided in '" << ore::NV("caller", CallerFunctionName)
424 << "' (frame_size="
425 << ore::NV("frame_size", FrameSizeAndAlign->first) << ", align="
426 << ore::NV("align", FrameSizeAndAlign->second.value()) << ")";
427 });
428 } else {
429 ORE.emit(RemarkBuilder: [&]() {
430 auto Remark = OptimizationRemarkMissed(DEBUG_TYPE, "CoroElide", CoroId)
431 << "'" << ore::NV("callee", CalleeCoroutineName)
432 << "' not elided in '"
433 << ore::NV("caller", CallerFunctionName);
434
435 if (FrameSizeAndAlign)
436 return Remark << "' (frame_size="
437 << ore::NV("frame_size", FrameSizeAndAlign->first)
438 << ", align="
439 << ore::NV("align", FrameSizeAndAlign->second.value())
440 << ")";
441 else
442 return Remark << "' (frame_size=unknown, align=unknown)";
443 });
444 }
445
446 return true;
447}
448
449PreservedAnalyses CoroElidePass::run(Function &F, FunctionAnalysisManager &AM) {
450 auto &M = *F.getParent();
451 if (!coro::declaresIntrinsics(M, List: Intrinsic::coro_id))
452 return PreservedAnalyses::all();
453
454 FunctionElideInfo FEI{&F};
455 // Elide is not necessary if there's no coro.id within the function.
456 if (!FEI.hasCoroIds())
457 return PreservedAnalyses::all();
458
459 AAResults &AA = AM.getResult<AAManager>(IR&: F);
460 DominatorTree &DT = AM.getResult<DominatorTreeAnalysis>(IR&: F);
461 auto &ORE = AM.getResult<OptimizationRemarkEmitterAnalysis>(IR&: F);
462
463 bool Changed = false;
464 for (auto *CII : FEI.getCoroIds()) {
465 CoroIdElider CIE(CII, FEI, AA, DT, ORE);
466 Changed |= CIE.attemptElide();
467 }
468
469 return Changed ? PreservedAnalyses::none() : PreservedAnalyses::all();
470}
471