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:
Lowerer(Module & M)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.
lowerResumeOrDestroy(CallBase & CB,CoroSubFnInst::ResumeKind Index)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.
lowerCoroPromise(CoroPromiseInst * Intrin)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.
lowerCoroDone(IntrinsicInst * II)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
lowerCoroNoop(IntrinsicInst * II)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.
setCannotDuplicate(CoroIdInst * CoroId)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
lowerEarlyIntrinsics(Function & F)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 (auto IB = inst_begin(F), IE = inst_end(F); IB != IE;) {
154 Instruction &I = *IB++;
155 if (auto *CB = dyn_cast<CallBase>(&I)) {
156 switch (CB->getIntrinsicID()) {
157 default:
158 continue;
159 case Intrinsic::coro_free:
160 CoroFrees.push_back(cast<CoroFreeInst>(&I));
161 break;
162 case Intrinsic::coro_suspend:
163 // Make sure that final suspend point is not duplicated as CoroSplit
164 // pass expects that there is at most one final suspend point.
165 if (cast<CoroSuspendInst>(&I)->isFinal())
166 CB->setCannotDuplicate();
167 HasCoroSuspend = true;
168 break;
169 case Intrinsic::coro_end_async:
170 case Intrinsic::coro_end:
171 // Make sure that fallthrough coro.end is not duplicated as CoroSplit
172 // pass expects that there is at most one fallthrough coro.end.
173 if (cast<AnyCoroEndInst>(&I)->isFallthrough())
174 CB->setCannotDuplicate();
175 break;
176 case Intrinsic::coro_noop:
177 lowerCoroNoop(cast<IntrinsicInst>(&I));
178 break;
179 case Intrinsic::coro_id:
180 // Mark a function that comes out of the frontend that has a coro.id
181 // with a coroutine attribute.
182 if (auto *CII = cast<CoroIdInst>(&I)) {
183 if (CII->getInfo().isPreSplit()) {
184 F.addFnAttr(CORO_PRESPLIT_ATTR, UNPREPARED_FOR_SPLIT);
185 setCannotDuplicate(CII);
186 CII->setCoroutineSelf();
187 CoroId = cast<CoroIdInst>(&I);
188 }
189 }
190 break;
191 case Intrinsic::coro_id_retcon:
192 case Intrinsic::coro_id_retcon_once:
193 case Intrinsic::coro_id_async:
194 F.addFnAttr(CORO_PRESPLIT_ATTR, PREPARED_FOR_SPLIT);
195 break;
196 case Intrinsic::coro_resume:
197 lowerResumeOrDestroy(*CB, CoroSubFnInst::ResumeIndex);
198 break;
199 case Intrinsic::coro_destroy:
200 lowerResumeOrDestroy(*CB, CoroSubFnInst::DestroyIndex);
201 break;
202 case Intrinsic::coro_promise:
203 lowerCoroPromise(cast<CoroPromiseInst>(&I));
204 break;
205 case Intrinsic::coro_done:
206 lowerCoroDone(cast<IntrinsicInst>(&I));
207 break;
208 }
209 Changed = true;
210 }
211 }
212 // Make sure that all CoroFree reference the coro.id intrinsic.
213 // Token type is not exposed through coroutine C/C++ builtins to plain C, so
214 // we allow specifying none and fixing it up here.
215 if (CoroId)
216 for (CoroFreeInst *CF : CoroFrees)
217 CF->setArgOperand(0, CoroId);
218 // Coroutine suspention could potentially lead to any argument modified
219 // outside of the function, hence arguments should not have noalias
220 // attributes.
221 if (HasCoroSuspend)
222 for (Argument &A : F.args())
223 if (A.hasNoAliasAttr())
224 A.removeAttr(Attribute::NoAlias);
225 return Changed;
226 }
227
declaresCoroEarlyIntrinsics(const Module & M)228 static bool declaresCoroEarlyIntrinsics(const Module &M) {
229 return coro::declaresIntrinsics(
230 M, {"llvm.coro.id", "llvm.coro.id.retcon", "llvm.coro.id.retcon.once",
231 "llvm.coro.id.async", "llvm.coro.destroy", "llvm.coro.done",
232 "llvm.coro.end", "llvm.coro.end.async", "llvm.coro.noop",
233 "llvm.coro.free", "llvm.coro.promise", "llvm.coro.resume",
234 "llvm.coro.suspend"});
235 }
236
run(Function & F,FunctionAnalysisManager &)237 PreservedAnalyses CoroEarlyPass::run(Function &F, FunctionAnalysisManager &) {
238 Module &M = *F.getParent();
239 if (!declaresCoroEarlyIntrinsics(M) || !Lowerer(M).lowerEarlyIntrinsics(F))
240 return PreservedAnalyses::all();
241
242 PreservedAnalyses PA;
243 PA.preserveSet<CFGAnalyses>();
244 return PA;
245 }
246
247 namespace {
248
249 struct CoroEarlyLegacy : public FunctionPass {
250 static char ID; // Pass identification, replacement for typeid.
CoroEarlyLegacy__anon054474440211::CoroEarlyLegacy251 CoroEarlyLegacy() : FunctionPass(ID) {
252 initializeCoroEarlyLegacyPass(*PassRegistry::getPassRegistry());
253 }
254
255 std::unique_ptr<Lowerer> L;
256
257 // This pass has work to do only if we find intrinsics we are going to lower
258 // in the module.
doInitialization__anon054474440211::CoroEarlyLegacy259 bool doInitialization(Module &M) override {
260 if (declaresCoroEarlyIntrinsics(M))
261 L = std::make_unique<Lowerer>(M);
262 return false;
263 }
264
runOnFunction__anon054474440211::CoroEarlyLegacy265 bool runOnFunction(Function &F) override {
266 if (!L)
267 return false;
268
269 return L->lowerEarlyIntrinsics(F);
270 }
271
getAnalysisUsage__anon054474440211::CoroEarlyLegacy272 void getAnalysisUsage(AnalysisUsage &AU) const override {
273 AU.setPreservesCFG();
274 }
getPassName__anon054474440211::CoroEarlyLegacy275 StringRef getPassName() const override {
276 return "Lower early coroutine intrinsics";
277 }
278 };
279 }
280
281 char CoroEarlyLegacy::ID = 0;
282 INITIALIZE_PASS(CoroEarlyLegacy, "coro-early",
283 "Lower early coroutine intrinsics", false, false)
284
createCoroEarlyLegacyPass()285 Pass *llvm::createCoroEarlyLegacyPass() { return new CoroEarlyLegacy(); }
286