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