1 //===- Coroutines.cpp -----------------------------------------------------===//
2 //
3 //                     The LLVM Compiler Infrastructure
4 //
5 // This file is distributed under the University of Illinois Open Source
6 // License. See LICENSE.TXT for details.
7 //
8 //===----------------------------------------------------------------------===//
9 //
10 // This file implements the common infrastructure for Coroutine Passes.
11 //
12 //===----------------------------------------------------------------------===//
13 
14 #include "CoroInstr.h"
15 #include "CoroInternal.h"
16 #include "llvm/ADT/SmallVector.h"
17 #include "llvm/ADT/StringRef.h"
18 #include "llvm/Analysis/CallGraph.h"
19 #include "llvm/Analysis/CallGraphSCCPass.h"
20 #include "llvm/IR/Attributes.h"
21 #include "llvm/IR/CallSite.h"
22 #include "llvm/IR/Constants.h"
23 #include "llvm/IR/DerivedTypes.h"
24 #include "llvm/IR/Function.h"
25 #include "llvm/IR/InstIterator.h"
26 #include "llvm/IR/Instructions.h"
27 #include "llvm/IR/IntrinsicInst.h"
28 #include "llvm/IR/Intrinsics.h"
29 #include "llvm/IR/LegacyPassManager.h"
30 #include "llvm/IR/Module.h"
31 #include "llvm/IR/Type.h"
32 #include "llvm/Support/Casting.h"
33 #include "llvm/Support/ErrorHandling.h"
34 #include "llvm/Transforms/Coroutines.h"
35 #include "llvm/Transforms/IPO.h"
36 #include "llvm/Transforms/IPO/PassManagerBuilder.h"
37 #include "llvm/Transforms/Utils/Local.h"
38 #include <cassert>
39 #include <cstddef>
40 #include <utility>
41 
42 using namespace llvm;
43 
44 void llvm::initializeCoroutines(PassRegistry &Registry) {
45   initializeCoroEarlyPass(Registry);
46   initializeCoroSplitPass(Registry);
47   initializeCoroElidePass(Registry);
48   initializeCoroCleanupPass(Registry);
49 }
50 
51 static void addCoroutineOpt0Passes(const PassManagerBuilder &Builder,
52                                    legacy::PassManagerBase &PM) {
53   PM.add(createCoroSplitPass());
54   PM.add(createCoroElidePass());
55 
56   PM.add(createBarrierNoopPass());
57   PM.add(createCoroCleanupPass());
58 }
59 
60 static void addCoroutineEarlyPasses(const PassManagerBuilder &Builder,
61                                     legacy::PassManagerBase &PM) {
62   PM.add(createCoroEarlyPass());
63 }
64 
65 static void addCoroutineScalarOptimizerPasses(const PassManagerBuilder &Builder,
66                                               legacy::PassManagerBase &PM) {
67   PM.add(createCoroElidePass());
68 }
69 
70 static void addCoroutineSCCPasses(const PassManagerBuilder &Builder,
71                                   legacy::PassManagerBase &PM) {
72   PM.add(createCoroSplitPass());
73 }
74 
75 static void addCoroutineOptimizerLastPasses(const PassManagerBuilder &Builder,
76                                             legacy::PassManagerBase &PM) {
77   PM.add(createCoroCleanupPass());
78 }
79 
80 void llvm::addCoroutinePassesToExtensionPoints(PassManagerBuilder &Builder) {
81   Builder.addExtension(PassManagerBuilder::EP_EarlyAsPossible,
82                        addCoroutineEarlyPasses);
83   Builder.addExtension(PassManagerBuilder::EP_EnabledOnOptLevel0,
84                        addCoroutineOpt0Passes);
85   Builder.addExtension(PassManagerBuilder::EP_CGSCCOptimizerLate,
86                        addCoroutineSCCPasses);
87   Builder.addExtension(PassManagerBuilder::EP_ScalarOptimizerLate,
88                        addCoroutineScalarOptimizerPasses);
89   Builder.addExtension(PassManagerBuilder::EP_OptimizerLast,
90                        addCoroutineOptimizerLastPasses);
91 }
92 
93 // Construct the lowerer base class and initialize its members.
94 coro::LowererBase::LowererBase(Module &M)
95     : TheModule(M), Context(M.getContext()),
96       Int8Ptr(Type::getInt8PtrTy(Context)),
97       ResumeFnType(FunctionType::get(Type::getVoidTy(Context), Int8Ptr,
98                                      /*isVarArg=*/false)),
99       NullPtr(ConstantPointerNull::get(Int8Ptr)) {}
100 
101 // Creates a sequence of instructions to obtain a resume function address using
102 // llvm.coro.subfn.addr. It generates the following sequence:
103 //
104 //    call i8* @llvm.coro.subfn.addr(i8* %Arg, i8 %index)
105 //    bitcast i8* %2 to void(i8*)*
106 
107 Value *coro::LowererBase::makeSubFnCall(Value *Arg, int Index,
108                                         Instruction *InsertPt) {
109   auto *IndexVal = ConstantInt::get(Type::getInt8Ty(Context), Index);
110   auto *Fn = Intrinsic::getDeclaration(&TheModule, Intrinsic::coro_subfn_addr);
111 
112   assert(Index >= CoroSubFnInst::IndexFirst &&
113          Index < CoroSubFnInst::IndexLast &&
114          "makeSubFnCall: Index value out of range");
115   auto *Call = CallInst::Create(Fn, {Arg, IndexVal}, "", InsertPt);
116 
117   auto *Bitcast =
118       new BitCastInst(Call, ResumeFnType->getPointerTo(), "", InsertPt);
119   return Bitcast;
120 }
121 
122 #ifndef NDEBUG
123 static bool isCoroutineIntrinsicName(StringRef Name) {
124   // NOTE: Must be sorted!
125   static const char *const CoroIntrinsics[] = {
126       "llvm.coro.alloc",   "llvm.coro.begin",   "llvm.coro.destroy",
127       "llvm.coro.done",    "llvm.coro.end",     "llvm.coro.frame",
128       "llvm.coro.free",    "llvm.coro.id",      "llvm.coro.param",
129       "llvm.coro.promise", "llvm.coro.resume",  "llvm.coro.save",
130       "llvm.coro.size",    "llvm.coro.subfn.addr", "llvm.coro.suspend",
131   };
132   return Intrinsic::lookupLLVMIntrinsicByName(CoroIntrinsics, Name) != -1;
133 }
134 #endif
135 
136 // Verifies if a module has named values listed. Also, in debug mode verifies
137 // that names are intrinsic names.
138 bool coro::declaresIntrinsics(Module &M,
139                               std::initializer_list<StringRef> List) {
140   for (StringRef Name : List) {
141     assert(isCoroutineIntrinsicName(Name) && "not a coroutine intrinsic");
142     if (M.getNamedValue(Name))
143       return true;
144   }
145 
146   return false;
147 }
148 
149 // Replace all coro.frees associated with the provided CoroId either with 'null'
150 // if Elide is true and with its frame parameter otherwise.
151 void coro::replaceCoroFree(CoroIdInst *CoroId, bool Elide) {
152   SmallVector<CoroFreeInst *, 4> CoroFrees;
153   for (User *U : CoroId->users())
154     if (auto CF = dyn_cast<CoroFreeInst>(U))
155       CoroFrees.push_back(CF);
156 
157   if (CoroFrees.empty())
158     return;
159 
160   Value *Replacement =
161       Elide ? ConstantPointerNull::get(Type::getInt8PtrTy(CoroId->getContext()))
162             : CoroFrees.front()->getFrame();
163 
164   for (CoroFreeInst *CF : CoroFrees) {
165     CF->replaceAllUsesWith(Replacement);
166     CF->eraseFromParent();
167   }
168 }
169 
170 // FIXME: This code is stolen from CallGraph::addToCallGraph(Function *F), which
171 // happens to be private. It is better for this functionality exposed by the
172 // CallGraph.
173 static void buildCGN(CallGraph &CG, CallGraphNode *Node) {
174   Function *F = Node->getFunction();
175 
176   // Look for calls by this function.
177   for (Instruction &I : instructions(F))
178     if (CallSite CS = CallSite(cast<Value>(&I))) {
179       const Function *Callee = CS.getCalledFunction();
180       if (!Callee || !Intrinsic::isLeaf(Callee->getIntrinsicID()))
181         // Indirect calls of intrinsics are not allowed so no need to check.
182         // We can be more precise here by using TargetArg returned by
183         // Intrinsic::isLeaf.
184         Node->addCalledFunction(CS, CG.getCallsExternalNode());
185       else if (!Callee->isIntrinsic())
186         Node->addCalledFunction(CS, CG.getOrInsertFunction(Callee));
187     }
188 }
189 
190 // Rebuild CGN after we extracted parts of the code from ParentFunc into
191 // NewFuncs. Builds CGNs for the NewFuncs and adds them to the current SCC.
192 void coro::updateCallGraph(Function &ParentFunc, ArrayRef<Function *> NewFuncs,
193                            CallGraph &CG, CallGraphSCC &SCC) {
194   // Rebuild CGN from scratch for the ParentFunc
195   auto *ParentNode = CG[&ParentFunc];
196   ParentNode->removeAllCalledFunctions();
197   buildCGN(CG, ParentNode);
198 
199   SmallVector<CallGraphNode *, 8> Nodes(SCC.begin(), SCC.end());
200 
201   for (Function *F : NewFuncs) {
202     CallGraphNode *Callee = CG.getOrInsertFunction(F);
203     Nodes.push_back(Callee);
204     buildCGN(CG, Callee);
205   }
206 
207   SCC.initialize(Nodes);
208 }
209 
210 static void clear(coro::Shape &Shape) {
211   Shape.CoroBegin = nullptr;
212   Shape.CoroEnds.clear();
213   Shape.CoroSizes.clear();
214   Shape.CoroSuspends.clear();
215 
216   Shape.FrameTy = nullptr;
217   Shape.FramePtr = nullptr;
218   Shape.AllocaSpillBlock = nullptr;
219   Shape.ResumeSwitch = nullptr;
220   Shape.PromiseAlloca = nullptr;
221   Shape.HasFinalSuspend = false;
222 }
223 
224 static CoroSaveInst *createCoroSave(CoroBeginInst *CoroBegin,
225                                     CoroSuspendInst *SuspendInst) {
226   Module *M = SuspendInst->getModule();
227   auto *Fn = Intrinsic::getDeclaration(M, Intrinsic::coro_save);
228   auto *SaveInst =
229       cast<CoroSaveInst>(CallInst::Create(Fn, CoroBegin, "", SuspendInst));
230   assert(!SuspendInst->getCoroSave());
231   SuspendInst->setArgOperand(0, SaveInst);
232   return SaveInst;
233 }
234 
235 // Collect "interesting" coroutine intrinsics.
236 void coro::Shape::buildFrom(Function &F) {
237   size_t FinalSuspendIndex = 0;
238   clear(*this);
239   SmallVector<CoroFrameInst *, 8> CoroFrames;
240   SmallVector<CoroSaveInst *, 2> UnusedCoroSaves;
241 
242   for (Instruction &I : instructions(F)) {
243     if (auto II = dyn_cast<IntrinsicInst>(&I)) {
244       switch (II->getIntrinsicID()) {
245       default:
246         continue;
247       case Intrinsic::coro_size:
248         CoroSizes.push_back(cast<CoroSizeInst>(II));
249         break;
250       case Intrinsic::coro_frame:
251         CoroFrames.push_back(cast<CoroFrameInst>(II));
252         break;
253       case Intrinsic::coro_save:
254         // After optimizations, coro_suspends using this coro_save might have
255         // been removed, remember orphaned coro_saves to remove them later.
256         if (II->use_empty())
257           UnusedCoroSaves.push_back(cast<CoroSaveInst>(II));
258         break;
259       case Intrinsic::coro_suspend:
260         CoroSuspends.push_back(cast<CoroSuspendInst>(II));
261         if (CoroSuspends.back()->isFinal()) {
262           if (HasFinalSuspend)
263             report_fatal_error(
264               "Only one suspend point can be marked as final");
265           HasFinalSuspend = true;
266           FinalSuspendIndex = CoroSuspends.size() - 1;
267         }
268         break;
269       case Intrinsic::coro_begin: {
270         auto CB = cast<CoroBeginInst>(II);
271         if (CB->getId()->getInfo().isPreSplit()) {
272           if (CoroBegin)
273             report_fatal_error(
274                 "coroutine should have exactly one defining @llvm.coro.begin");
275           CB->addAttribute(AttributeList::ReturnIndex, Attribute::NonNull);
276           CB->addAttribute(AttributeList::ReturnIndex, Attribute::NoAlias);
277           CB->removeAttribute(AttributeList::FunctionIndex,
278                               Attribute::NoDuplicate);
279           CoroBegin = CB;
280         }
281         break;
282       }
283       case Intrinsic::coro_end:
284         CoroEnds.push_back(cast<CoroEndInst>(II));
285         if (CoroEnds.back()->isFallthrough()) {
286           // Make sure that the fallthrough coro.end is the first element in the
287           // CoroEnds vector.
288           if (CoroEnds.size() > 1) {
289             if (CoroEnds.front()->isFallthrough())
290               report_fatal_error(
291                   "Only one coro.end can be marked as fallthrough");
292             std::swap(CoroEnds.front(), CoroEnds.back());
293           }
294         }
295         break;
296       }
297     }
298   }
299 
300   // If for some reason, we were not able to find coro.begin, bailout.
301   if (!CoroBegin) {
302     // Replace coro.frame which are supposed to be lowered to the result of
303     // coro.begin with undef.
304     auto *Undef = UndefValue::get(Type::getInt8PtrTy(F.getContext()));
305     for (CoroFrameInst *CF : CoroFrames) {
306       CF->replaceAllUsesWith(Undef);
307       CF->eraseFromParent();
308     }
309 
310     // Replace all coro.suspend with undef and remove related coro.saves if
311     // present.
312     for (CoroSuspendInst *CS : CoroSuspends) {
313       CS->replaceAllUsesWith(UndefValue::get(CS->getType()));
314       CS->eraseFromParent();
315       if (auto *CoroSave = CS->getCoroSave())
316         CoroSave->eraseFromParent();
317     }
318 
319     // Replace all coro.ends with unreachable instruction.
320     for (CoroEndInst *CE : CoroEnds)
321       changeToUnreachable(CE, /*UseLLVMTrap=*/false);
322 
323     return;
324   }
325 
326   // The coro.free intrinsic is always lowered to the result of coro.begin.
327   for (CoroFrameInst *CF : CoroFrames) {
328     CF->replaceAllUsesWith(CoroBegin);
329     CF->eraseFromParent();
330   }
331 
332   // Canonicalize coro.suspend by inserting a coro.save if needed.
333   for (CoroSuspendInst *CS : CoroSuspends)
334     if (!CS->getCoroSave())
335       createCoroSave(CoroBegin, CS);
336 
337   // Move final suspend to be the last element in the CoroSuspends vector.
338   if (HasFinalSuspend &&
339       FinalSuspendIndex != CoroSuspends.size() - 1)
340     std::swap(CoroSuspends[FinalSuspendIndex], CoroSuspends.back());
341 
342   // Remove orphaned coro.saves.
343   for (CoroSaveInst *CoroSave : UnusedCoroSaves)
344     CoroSave->eraseFromParent();
345 }
346