1 //===- CoroElide.cpp - Coroutine Frame Allocation Elision Pass ------------===//
3 // The LLVM Compiler Infrastructure
5 // This file is distributed under the University of Illinois Open Source
6 // License. See LICENSE.TXT for details.
8 //===----------------------------------------------------------------------===//
9 // This pass replaces dynamic allocation of coroutine frame with alloca and
10 // replaces calls to llvm.coro.resume and llvm.coro.destroy with direct calls
11 // to coroutine sub-functions.
12 //===----------------------------------------------------------------------===//
14 #include "CoroInternal.h"
15 #include "llvm/Analysis/AliasAnalysis.h"
16 #include "llvm/Analysis/InstructionSimplify.h"
17 #include "llvm/IR/InstIterator.h"
18 #include "llvm/Pass.h"
19 #include "llvm/Support/ErrorHandling.h"
23 #define DEBUG_TYPE "coro-elide"
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;
35 Lowerer(Module &M) : LowererBase(M) {}
37 void elideHeapAllocations(Function *F, Type *FrameTy, AAResults &AA);
38 bool shouldElide() const;
39 bool processCoroId(CoroIdInst *, AAResults &AA);
41 } // end anonymous namespace
43 // Go through the list of coro.subfn.addr intrinsics and replace them with the
45 static void replaceWithConstant(Constant *Value,
46 SmallVectorImpl<CoroSubFnInst *> &Users) {
50 // See if we need to bitcast the constant to match the type of the intrinsic
51 // being replaced. Note: All coro.subfn.addr intrinsics return the same type,
52 // so we only need to examine the type of the first one in the list.
53 Type *IntrTy = Users.front()->getType();
54 Type *ValueTy = Value->getType();
55 if (ValueTy != IntrTy) {
56 // May need to tweak the function type to match the type expected at the
58 assert(ValueTy->isPointerTy() && IntrTy->isPointerTy());
59 Value = ConstantExpr::getBitCast(Value, IntrTy);
62 // Now the value type matches the type of the intrinsic. Replace them all!
63 for (CoroSubFnInst *I : Users)
64 replaceAndRecursivelySimplify(I, Value);
67 // See if any operand of the call instruction references the coroutine frame.
68 static bool operandReferences(CallInst *CI, AllocaInst *Frame, AAResults &AA) {
69 for (Value *Op : CI->operand_values())
70 if (AA.alias(Op, Frame) != NoAlias)
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 static void removeTailCallAttribute(AllocaInst *Frame, AAResults &AA) {
79 Function &F = *Frame->getFunction();
80 MemoryLocation Mem(Frame);
81 for (Instruction &I : instructions(F))
82 if (auto *Call = dyn_cast<CallInst>(&I))
83 if (Call->isTailCall() && operandReferences(Call, Frame, AA)) {
84 // FIXME: If we ever hit this check. Evaluate whether it is more
85 // appropriate to retain musttail and allow the code to compile.
86 if (Call->isMustTailCall())
87 report_fatal_error("Call referring to the coroutine frame cannot be "
88 "marked as musttail");
89 Call->setTailCall(false);
93 // Given a resume function @f.resume(%f.frame* %frame), returns %f.frame type.
94 static Type *getFrameType(Function *Resume) {
95 auto *ArgType = Resume->arg_begin()->getType();
96 return cast<PointerType>(ArgType)->getElementType();
99 // Finds first non alloca instruction in the entry block of a function.
100 static Instruction *getFirstNonAllocaInTheEntryBlock(Function *F) {
101 for (Instruction &I : F->getEntryBlock())
102 if (!isa<AllocaInst>(&I))
104 llvm_unreachable("no terminator in the entry block");
107 // To elide heap allocations we need to suppress code blocks guarded by
108 // llvm.coro.alloc and llvm.coro.free instructions.
109 void Lowerer::elideHeapAllocations(Function *F, Type *FrameTy, AAResults &AA) {
110 LLVMContext &C = FrameTy->getContext();
112 getFirstNonAllocaInTheEntryBlock(CoroIds.front()->getFunction());
114 // Replacing llvm.coro.alloc with false will suppress dynamic
115 // allocation as it is expected for the frontend to generate the code that
118 // mem = coro.alloc(id) ? malloc(coro.size()) : 0;
119 // coro.begin(id, mem)
120 auto *False = ConstantInt::getFalse(C);
121 for (auto *CA : CoroAllocs) {
122 CA->replaceAllUsesWith(False);
123 CA->eraseFromParent();
126 // FIXME: Design how to transmit alignment information for every alloca that
127 // is spilled into the coroutine frame and recreate the alignment information
128 // here. Possibly we will need to do a mini SROA here and break the coroutine
129 // frame into individual AllocaInst recreating the original alignment.
130 const DataLayout &DL = F->getParent()->getDataLayout();
131 auto *Frame = new AllocaInst(FrameTy, DL.getAllocaAddrSpace(), "", InsertPt);
133 new BitCastInst(Frame, Type::getInt8PtrTy(C), "vFrame", InsertPt);
135 for (auto *CB : CoroBegins) {
136 CB->replaceAllUsesWith(FrameVoidPtr);
137 CB->eraseFromParent();
140 // Since now coroutine frame lives on the stack we need to make sure that
141 // any tail call referencing it, must be made non-tail call.
142 removeTailCallAttribute(Frame, AA);
145 bool Lowerer::shouldElide() const {
146 // If no CoroAllocs, we cannot suppress allocation, so elision is not
148 if (CoroAllocs.empty())
151 // Check that for every coro.begin there is a coro.destroy directly
152 // referencing the SSA value of that coro.begin. If the value escaped, then
153 // coro.destroy would have been referencing a memory location storing that
154 // value and not the virtual register.
156 SmallPtrSet<CoroBeginInst *, 8> ReferencedCoroBegins;
158 for (CoroSubFnInst *DA : DestroyAddr) {
159 if (auto *CB = dyn_cast<CoroBeginInst>(DA->getFrame()))
160 ReferencedCoroBegins.insert(CB);
165 // If size of the set is the same as total number of CoroBegins, means we
166 // found a coro.free or coro.destroy mentioning a coro.begin and we can
167 // perform heap elision.
168 return ReferencedCoroBegins.size() == CoroBegins.size();
171 bool Lowerer::processCoroId(CoroIdInst *CoroId, AAResults &AA) {
178 // Collect all coro.begin and coro.allocs associated with this coro.id.
179 for (User *U : CoroId->users()) {
180 if (auto *CB = dyn_cast<CoroBeginInst>(U))
181 CoroBegins.push_back(CB);
182 else if (auto *CA = dyn_cast<CoroAllocInst>(U))
183 CoroAllocs.push_back(CA);
184 else if (auto *CF = dyn_cast<CoroFreeInst>(U))
185 CoroFrees.push_back(CF);
188 // Collect all coro.subfn.addrs associated with coro.begin.
189 // Note, we only devirtualize the calls if their coro.subfn.addr refers to
190 // coro.begin directly. If we run into cases where this check is too
191 // conservative, we can consider relaxing the check.
192 for (CoroBeginInst *CB : CoroBegins) {
193 for (User *U : CB->users())
194 if (auto *II = dyn_cast<CoroSubFnInst>(U))
195 switch (II->getIndex()) {
196 case CoroSubFnInst::ResumeIndex:
197 ResumeAddr.push_back(II);
199 case CoroSubFnInst::DestroyIndex:
200 DestroyAddr.push_back(II);
203 llvm_unreachable("unexpected coro.subfn.addr constant");
207 // PostSplit coro.id refers to an array of subfunctions in its Info
209 ConstantArray *Resumers = CoroId->getInfo().Resumers;
210 assert(Resumers && "PostSplit coro.id Info argument must refer to an array"
211 "of coroutine subfunctions");
212 auto *ResumeAddrConstant =
213 ConstantExpr::getExtractValue(Resumers, CoroSubFnInst::ResumeIndex);
215 replaceWithConstant(ResumeAddrConstant, ResumeAddr);
217 bool ShouldElide = shouldElide();
219 auto *DestroyAddrConstant = ConstantExpr::getExtractValue(
221 ShouldElide ? CoroSubFnInst::CleanupIndex : CoroSubFnInst::DestroyIndex);
223 replaceWithConstant(DestroyAddrConstant, DestroyAddr);
226 auto *FrameTy = getFrameType(cast<Function>(ResumeAddrConstant));
227 elideHeapAllocations(CoroId->getFunction(), FrameTy, AA);
228 coro::replaceCoroFree(CoroId, /*Elide=*/true);
234 // See if there are any coro.subfn.addr instructions referring to coro.devirt
235 // trigger, if so, replace them with a direct call to devirt trigger function.
236 static bool replaceDevirtTrigger(Function &F) {
237 SmallVector<CoroSubFnInst *, 1> DevirtAddr;
238 for (auto &I : instructions(F))
239 if (auto *SubFn = dyn_cast<CoroSubFnInst>(&I))
240 if (SubFn->getIndex() == CoroSubFnInst::RestartTrigger)
241 DevirtAddr.push_back(SubFn);
243 if (DevirtAddr.empty())
246 Module &M = *F.getParent();
247 Function *DevirtFn = M.getFunction(CORO_DEVIRT_TRIGGER_FN);
248 assert(DevirtFn && "coro.devirt.fn not found");
249 replaceWithConstant(DevirtFn, DevirtAddr);
254 //===----------------------------------------------------------------------===//
256 //===----------------------------------------------------------------------===//
259 struct CoroElide : FunctionPass {
261 CoroElide() : FunctionPass(ID) {}
263 std::unique_ptr<Lowerer> L;
265 bool doInitialization(Module &M) override {
266 if (coro::declaresIntrinsics(M, {"llvm.coro.id"}))
267 L = llvm::make_unique<Lowerer>(M);
271 bool runOnFunction(Function &F) override {
275 bool Changed = false;
277 if (F.hasFnAttribute(CORO_PRESPLIT_ATTR))
278 Changed = replaceDevirtTrigger(F);
282 // Collect all PostSplit coro.ids.
283 for (auto &I : instructions(F))
284 if (auto *CII = dyn_cast<CoroIdInst>(&I))
285 if (CII->getInfo().isPostSplit())
286 // If it is the coroutine itself, don't touch it.
287 if (CII->getCoroutine() != CII->getFunction())
288 L->CoroIds.push_back(CII);
290 // If we did not find any coro.id, there is nothing to do.
291 if (L->CoroIds.empty())
294 AAResults &AA = getAnalysis<AAResultsWrapperPass>().getAAResults();
296 for (auto *CII : L->CoroIds)
297 Changed |= L->processCoroId(CII, AA);
301 void getAnalysisUsage(AnalysisUsage &AU) const override {
302 AU.addRequired<AAResultsWrapperPass>();
307 char CoroElide::ID = 0;
308 INITIALIZE_PASS_BEGIN(
309 CoroElide, "coro-elide",
310 "Coroutine frame allocation elision and indirect calls replacement", false,
312 INITIALIZE_PASS_DEPENDENCY(AAResultsWrapperPass)
314 CoroElide, "coro-elide",
315 "Coroutine frame allocation elision and indirect calls replacement", false,
318 Pass *llvm::createCoroElidePass() { return new CoroElide(); }