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