1 //===- CoroEarly.cpp - Coroutine Early Function Pass ----------------------===//
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 // This pass lowers coroutine intrinsics that hide the details of the exact
10 // calling convention for coroutine resume and destroy functions and details of
11 // the structure of the coroutine frame.
12 //===----------------------------------------------------------------------===//
13 
14 #include "CoroInternal.h"
15 #include "llvm/IR/CallSite.h"
16 #include "llvm/IR/IRBuilder.h"
17 #include "llvm/IR/InstIterator.h"
18 #include "llvm/IR/Module.h"
19 #include "llvm/Pass.h"
20 
21 using namespace llvm;
22 
23 #define DEBUG_TYPE "coro-early"
24 
25 namespace {
26 // Created on demand if CoroEarly pass has work to do.
27 class Lowerer : public coro::LowererBase {
28   IRBuilder<> Builder;
29   PointerType *const AnyResumeFnPtrTy;
30 
31   void lowerResumeOrDestroy(CallSite CS, CoroSubFnInst::ResumeKind);
32   void lowerCoroPromise(CoroPromiseInst *Intrin);
33   void lowerCoroDone(IntrinsicInst *II);
34 
35 public:
36   Lowerer(Module &M)
37       : LowererBase(M), Builder(Context),
38         AnyResumeFnPtrTy(FunctionType::get(Type::getVoidTy(Context), Int8Ptr,
39                                            /*isVarArg=*/false)
40                              ->getPointerTo()) {}
41   bool lowerEarlyIntrinsics(Function &F);
42 };
43 }
44 
45 // Replace a direct call to coro.resume or coro.destroy with an indirect call to
46 // an address returned by coro.subfn.addr intrinsic. This is done so that
47 // CGPassManager recognizes devirtualization when CoroElide pass replaces a call
48 // to coro.subfn.addr with an appropriate function address.
49 void Lowerer::lowerResumeOrDestroy(CallSite CS,
50                                    CoroSubFnInst::ResumeKind Index) {
51   Value *ResumeAddr =
52       makeSubFnCall(CS.getArgOperand(0), Index, CS.getInstruction());
53   CS.setCalledFunction(ResumeAddr);
54   CS.setCallingConv(CallingConv::Fast);
55 }
56 
57 // Coroutine promise field is always at the fixed offset from the beginning of
58 // the coroutine frame. i8* coro.promise(i8*, i1 from) intrinsic adds an offset
59 // to a passed pointer to move from coroutine frame to coroutine promise and
60 // vice versa. Since we don't know exactly which coroutine frame it is, we build
61 // a coroutine frame mock up starting with two function pointers, followed by a
62 // properly aligned coroutine promise field.
63 // TODO: Handle the case when coroutine promise alloca has align override.
64 void Lowerer::lowerCoroPromise(CoroPromiseInst *Intrin) {
65   Value *Operand = Intrin->getArgOperand(0);
66   unsigned Alignement = Intrin->getAlignment();
67   Type *Int8Ty = Builder.getInt8Ty();
68 
69   auto *SampleStruct =
70       StructType::get(Context, {AnyResumeFnPtrTy, AnyResumeFnPtrTy, Int8Ty});
71   const DataLayout &DL = TheModule.getDataLayout();
72   int64_t Offset = alignTo(
73       DL.getStructLayout(SampleStruct)->getElementOffset(2), Alignement);
74   if (Intrin->isFromPromise())
75     Offset = -Offset;
76 
77   Builder.SetInsertPoint(Intrin);
78   Value *Replacement =
79       Builder.CreateConstInBoundsGEP1_32(Int8Ty, Operand, Offset);
80 
81   Intrin->replaceAllUsesWith(Replacement);
82   Intrin->eraseFromParent();
83 }
84 
85 // When a coroutine reaches final suspend point, it zeros out ResumeFnAddr in
86 // the coroutine frame (it is UB to resume from a final suspend point).
87 // The llvm.coro.done intrinsic is used to check whether a coroutine is
88 // suspended at the final suspend point or not.
89 void Lowerer::lowerCoroDone(IntrinsicInst *II) {
90   Value *Operand = II->getArgOperand(0);
91 
92   // ResumeFnAddr is the first pointer sized element of the coroutine frame.
93   auto *FrameTy = Int8Ptr;
94   PointerType *FramePtrTy = FrameTy->getPointerTo();
95 
96   Builder.SetInsertPoint(II);
97   auto *BCI = Builder.CreateBitCast(Operand, FramePtrTy);
98   auto *Gep = Builder.CreateConstInBoundsGEP1_32(FrameTy, BCI, 0);
99   auto *Load = Builder.CreateLoad(Gep);
100   auto *Cond = Builder.CreateICmpEQ(Load, NullPtr);
101 
102   II->replaceAllUsesWith(Cond);
103   II->eraseFromParent();
104 }
105 
106 // Prior to CoroSplit, calls to coro.begin needs to be marked as NoDuplicate,
107 // as CoroSplit assumes there is exactly one coro.begin. After CoroSplit,
108 // NoDuplicate attribute will be removed from coro.begin otherwise, it will
109 // interfere with inlining.
110 static void setCannotDuplicate(CoroIdInst *CoroId) {
111   for (User *U : CoroId->users())
112     if (auto *CB = dyn_cast<CoroBeginInst>(U))
113       CB->setCannotDuplicate();
114 }
115 
116 bool Lowerer::lowerEarlyIntrinsics(Function &F) {
117   bool Changed = false;
118   CoroIdInst *CoroId = nullptr;
119   SmallVector<CoroFreeInst *, 4> CoroFrees;
120   for (auto IB = inst_begin(F), IE = inst_end(F); IB != IE;) {
121     Instruction &I = *IB++;
122     if (auto CS = CallSite(&I)) {
123       switch (CS.getIntrinsicID()) {
124       default:
125         continue;
126       case Intrinsic::coro_free:
127         CoroFrees.push_back(cast<CoroFreeInst>(&I));
128         break;
129       case Intrinsic::coro_suspend:
130         // Make sure that final suspend point is not duplicated as CoroSplit
131         // pass expects that there is at most one final suspend point.
132         if (cast<CoroSuspendInst>(&I)->isFinal())
133           CS.setCannotDuplicate();
134         break;
135       case Intrinsic::coro_end:
136         // Make sure that fallthrough coro.end is not duplicated as CoroSplit
137         // pass expects that there is at most one fallthrough coro.end.
138         if (cast<CoroEndInst>(&I)->isFallthrough())
139           CS.setCannotDuplicate();
140         break;
141       case Intrinsic::coro_id:
142         // Mark a function that comes out of the frontend that has a coro.id
143         // with a coroutine attribute.
144         if (auto *CII = cast<CoroIdInst>(&I)) {
145           if (CII->getInfo().isPreSplit()) {
146             F.addFnAttr(CORO_PRESPLIT_ATTR, UNPREPARED_FOR_SPLIT);
147             setCannotDuplicate(CII);
148             CII->setCoroutineSelf();
149             CoroId = cast<CoroIdInst>(&I);
150           }
151         }
152         break;
153       case Intrinsic::coro_resume:
154         lowerResumeOrDestroy(CS, CoroSubFnInst::ResumeIndex);
155         break;
156       case Intrinsic::coro_destroy:
157         lowerResumeOrDestroy(CS, CoroSubFnInst::DestroyIndex);
158         break;
159       case Intrinsic::coro_promise:
160         lowerCoroPromise(cast<CoroPromiseInst>(&I));
161         break;
162       case Intrinsic::coro_done:
163         lowerCoroDone(cast<IntrinsicInst>(&I));
164         break;
165       }
166       Changed = true;
167     }
168   }
169   // Make sure that all CoroFree reference the coro.id intrinsic.
170   // Token type is not exposed through coroutine C/C++ builtins to plain C, so
171   // we allow specifying none and fixing it up here.
172   if (CoroId)
173     for (CoroFreeInst *CF : CoroFrees)
174       CF->setArgOperand(0, CoroId);
175   return Changed;
176 }
177 
178 //===----------------------------------------------------------------------===//
179 //                              Top Level Driver
180 //===----------------------------------------------------------------------===//
181 
182 namespace {
183 
184 struct CoroEarly : public FunctionPass {
185   static char ID; // Pass identification, replacement for typeid.
186   CoroEarly() : FunctionPass(ID) {}
187 
188   std::unique_ptr<Lowerer> L;
189 
190   // This pass has work to do only if we find intrinsics we are going to lower
191   // in the module.
192   bool doInitialization(Module &M) override {
193     if (coro::declaresIntrinsics(M, {"llvm.coro.id", "llvm.coro.destroy",
194                                      "llvm.coro.done", "llvm.coro.end",
195                                      "llvm.coro.free", "llvm.coro.promise",
196                                      "llvm.coro.resume", "llvm.coro.suspend"}))
197       L = llvm::make_unique<Lowerer>(M);
198     return false;
199   }
200 
201   bool runOnFunction(Function &F) override {
202     if (!L)
203       return false;
204 
205     return L->lowerEarlyIntrinsics(F);
206   }
207 
208   void getAnalysisUsage(AnalysisUsage &AU) const override {
209     AU.setPreservesCFG();
210   }
211 };
212 }
213 
214 char CoroEarly::ID = 0;
215 INITIALIZE_PASS(CoroEarly, "coro-early", "Lower early coroutine intrinsics",
216                 false, false)
217 
218 Pass *llvm::createCoroEarlyPass() { return new CoroEarly(); }
219