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/DIBuilder.h"
12#include "llvm/IR/Function.h"
13#include "llvm/IR/IRBuilder.h"
14#include "llvm/IR/InstIterator.h"
15#include "llvm/IR/Module.h"
16#include "llvm/Transforms/Coroutines/CoroShape.h"
17
18using namespace llvm;
19
20#define DEBUG_TYPE "coro-early"
21
22namespace {
23// Created on demand if the coro-early pass has work to do.
24class Lowerer : public coro::LowererBase {
25 IRBuilder<> Builder;
26 PointerType *const AnyResumeFnPtrTy;
27 Constant *NoopCoro = nullptr;
28
29 void lowerResumeOrDestroy(CallBase &CB, CoroSubFnInst::ResumeKind);
30 void lowerCoroPromise(CoroPromiseInst *Intrin);
31 void lowerCoroDone(IntrinsicInst *II);
32 void lowerCoroNoop(IntrinsicInst *II);
33 void hidePromiseAlloca(CoroIdInst *CoroId, CoroBeginInst *CoroBegin);
34
35public:
36 Lowerer(Module &M)
37 : LowererBase(M), Builder(Context),
38 AnyResumeFnPtrTy(PointerType::getUnqual(C&: Context)) {}
39 void 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.
47void Lowerer::lowerResumeOrDestroy(CallBase &CB,
48 CoroSubFnInst::ResumeKind Index) {
49 Value *ResumeAddr = makeSubFnCall(Arg: CB.getArgOperand(i: 0), Index, InsertPt: &CB);
50 CB.setCalledOperand(ResumeAddr);
51 CB.setCallingConv(CallingConv::Fast);
52}
53
54// Coroutine promise field is always at the fixed offset from the beginning of
55// the coroutine frame. i8* coro.promise(i8*, i1 from) intrinsic adds an offset
56// to a passed pointer to move from coroutine frame to coroutine promise and
57// vice versa. Since we don't know exactly which coroutine frame it is, we build
58// a coroutine frame mock up starting with two function pointers, followed by a
59// properly aligned coroutine promise field.
60// TODO: Handle the case when coroutine promise alloca has align override.
61void Lowerer::lowerCoroPromise(CoroPromiseInst *Intrin) {
62 Value *Operand = Intrin->getArgOperand(i: 0);
63 Align Alignment = Intrin->getAlignment();
64 Type *Int8Ty = Builder.getInt8Ty();
65
66 auto *SampleStruct =
67 StructType::get(Context, Elements: {AnyResumeFnPtrTy, AnyResumeFnPtrTy, Int8Ty});
68 const DataLayout &DL = TheModule.getDataLayout();
69 int64_t Offset = alignTo(
70 Size: DL.getStructLayout(Ty: SampleStruct)->getElementOffset(Idx: 2), A: Alignment);
71 if (Intrin->isFromPromise())
72 Offset = -Offset;
73
74 Builder.SetInsertPoint(Intrin);
75 Value *Replacement =
76 Builder.CreateConstInBoundsGEP1_32(Ty: Int8Ty, Ptr: Operand, Idx0: Offset);
77
78 Intrin->replaceAllUsesWith(V: Replacement);
79 Intrin->eraseFromParent();
80}
81
82// When a coroutine reaches final suspend point, it zeros out ResumeFnAddr in
83// the coroutine frame (it is UB to resume from a final suspend point).
84// The llvm.coro.done intrinsic is used to check whether a coroutine is
85// suspended at the final suspend point or not.
86void Lowerer::lowerCoroDone(IntrinsicInst *II) {
87 Value *Operand = II->getArgOperand(i: 0);
88
89 // ResumeFnAddr is the first pointer sized element of the coroutine frame.
90 static_assert(coro::Shape::SwitchFieldIndex::Resume == 0,
91 "resume function not at offset zero");
92 auto *FrameTy = Int8Ptr;
93
94 Builder.SetInsertPoint(II);
95 auto *Load = Builder.CreateLoad(Ty: FrameTy, Ptr: Operand);
96 auto *Cond = Builder.CreateICmpEQ(LHS: Load, RHS: NullPtr);
97
98 II->replaceAllUsesWith(V: Cond);
99 II->eraseFromParent();
100}
101
102static void buildDebugInfoForNoopResumeDestroyFunc(Function *NoopFn) {
103 Module &M = *NoopFn->getParent();
104 if (M.debug_compile_units().empty())
105 return;
106
107 DICompileUnit *CU = *M.debug_compile_units_begin();
108 DIBuilder DB(M, /*AllowUnresolved*/ false, CU);
109 std::array<Metadata *, 2> Params{nullptr, nullptr};
110 auto *SubroutineType =
111 DB.createSubroutineType(ParameterTypes: DB.getOrCreateTypeArray(Elements: Params));
112 StringRef Name = NoopFn->getName();
113 auto *SP = DB.createFunction(
114 Scope: CU, /*Name=*/Name, /*LinkageName=*/Name, /*File=*/ CU->getFile(),
115 /*LineNo=*/0, Ty: SubroutineType, /*ScopeLine=*/0, Flags: DINode::FlagArtificial,
116 SPFlags: DISubprogram::SPFlagDefinition);
117 NoopFn->setSubprogram(SP);
118 DB.finalize();
119}
120
121void Lowerer::lowerCoroNoop(IntrinsicInst *II) {
122 if (!NoopCoro) {
123 LLVMContext &C = Builder.getContext();
124 Module &M = *II->getModule();
125
126 // Create a noop.frame struct type.
127 auto *FnTy = FunctionType::get(Result: Type::getVoidTy(C), Params: Builder.getPtrTy(AddrSpace: 0),
128 /*isVarArg=*/false);
129 auto *FnPtrTy = Builder.getPtrTy(AddrSpace: 0);
130 StructType *FrameTy =
131 StructType::create(Elements: {FnPtrTy, FnPtrTy}, Name: "NoopCoro.Frame");
132
133 // Create a Noop function that does nothing.
134 Function *NoopFn = Function::createWithDefaultAttr(
135 Ty: FnTy, Linkage: GlobalValue::LinkageTypes::PrivateLinkage,
136 AddrSpace: M.getDataLayout().getProgramAddressSpace(), N: "__NoopCoro_ResumeDestroy",
137 M: &M);
138 NoopFn->setCallingConv(CallingConv::Fast);
139 buildDebugInfoForNoopResumeDestroyFunc(NoopFn);
140 auto *Entry = BasicBlock::Create(Context&: C, Name: "entry", Parent: NoopFn);
141 ReturnInst::Create(C, InsertAtEnd: Entry);
142
143 // Create a constant struct for the frame.
144 Constant* Values[] = {NoopFn, NoopFn};
145 Constant* NoopCoroConst = ConstantStruct::get(T: FrameTy, V: Values);
146 NoopCoro = new GlobalVariable(M, NoopCoroConst->getType(), /*isConstant=*/true,
147 GlobalVariable::PrivateLinkage, NoopCoroConst,
148 "NoopCoro.Frame.Const");
149 cast<GlobalVariable>(Val: NoopCoro)->setNoSanitizeMetadata();
150 }
151
152 Builder.SetInsertPoint(II);
153 auto *NoopCoroVoidPtr = Builder.CreateBitCast(V: NoopCoro, DestTy: Int8Ptr);
154 II->replaceAllUsesWith(V: NoopCoroVoidPtr);
155 II->eraseFromParent();
156}
157
158// Later middle-end passes will assume promise alloca dead after coroutine
159// suspend, leading to misoptimizations. We hide promise alloca using
160// coro.promise and will lower it back to alloca at CoroSplit.
161void Lowerer::hidePromiseAlloca(CoroIdInst *CoroId, CoroBeginInst *CoroBegin) {
162 auto *PA = CoroId->getPromise();
163 if (!PA || !CoroBegin)
164 return;
165 Builder.SetInsertPoint(*CoroBegin->getInsertionPointAfterDef());
166
167 auto *Alignment = Builder.getInt32(C: PA->getAlign().value());
168 auto *FromPromise = Builder.getInt1(V: false);
169 SmallVector<Value *, 3> Arg{CoroBegin, Alignment, FromPromise};
170 auto *PI = Builder.CreateIntrinsic(
171 RetTy: Builder.getPtrTy(), ID: Intrinsic::coro_promise, Args: Arg, FMFSource: {}, Name: "promise.addr");
172 PI->setCannotDuplicate();
173 PA->replaceUsesWithIf(New: PI, ShouldReplace: [CoroId](Use &U) {
174 bool IsBitcast = U == U.getUser()->stripPointerCasts();
175 bool IsCoroId = U.getUser() == CoroId;
176 return !IsBitcast && !IsCoroId;
177 });
178}
179
180// Prior to CoroSplit, calls to coro.begin needs to be marked as NoDuplicate,
181// as CoroSplit assumes there is exactly one coro.begin. After CoroSplit,
182// NoDuplicate attribute will be removed from coro.begin otherwise, it will
183// interfere with inlining.
184static void setCannotDuplicate(CoroIdInst *CoroId) {
185 for (User *U : CoroId->users())
186 if (auto *CB = dyn_cast<CoroBeginInst>(Val: U))
187 CB->setCannotDuplicate();
188}
189
190void Lowerer::lowerEarlyIntrinsics(Function &F) {
191 CoroIdInst *CoroId = nullptr;
192 CoroBeginInst *CoroBegin = nullptr;
193 SmallVector<CoroFreeInst *, 4> CoroFrees;
194 bool HasCoroSuspend = false;
195 for (Instruction &I : llvm::make_early_inc_range(Range: instructions(F))) {
196 auto *CB = dyn_cast<CallBase>(Val: &I);
197 if (!CB)
198 continue;
199
200 switch (CB->getIntrinsicID()) {
201 default:
202 continue;
203 case Intrinsic::coro_begin:
204 case Intrinsic::coro_begin_custom_abi:
205 if (CoroBegin)
206 report_fatal_error(
207 reason: "coroutine should have exactly one defining @llvm.coro.begin");
208 CoroBegin = cast<CoroBeginInst>(Val: &I);
209 break;
210 case Intrinsic::coro_free:
211 CoroFrees.push_back(Elt: cast<CoroFreeInst>(Val: &I));
212 break;
213 case Intrinsic::coro_suspend:
214 // Make sure that final suspend point is not duplicated as CoroSplit
215 // pass expects that there is at most one final suspend point.
216 if (cast<CoroSuspendInst>(Val: &I)->isFinal())
217 CB->setCannotDuplicate();
218 HasCoroSuspend = true;
219 break;
220 case Intrinsic::coro_end_async:
221 case Intrinsic::coro_end:
222 // Make sure that fallthrough coro.end is not duplicated as CoroSplit
223 // pass expects that there is at most one fallthrough coro.end.
224 if (cast<AnyCoroEndInst>(Val: &I)->isFallthrough())
225 CB->setCannotDuplicate();
226 break;
227 case Intrinsic::coro_noop:
228 lowerCoroNoop(II: cast<IntrinsicInst>(Val: &I));
229 break;
230 case Intrinsic::coro_id:
231 if (auto *CII = cast<CoroIdInst>(Val: &I)) {
232 if (CII->getInfo().isPreSplit()) {
233 assert(F.isPresplitCoroutine() &&
234 "The frontend uses Switch-Resumed ABI should emit "
235 "\"presplitcoroutine\" attribute for the coroutine.");
236 setCannotDuplicate(CII);
237 CII->setCoroutineSelf();
238 CoroId = cast<CoroIdInst>(Val: &I);
239 }
240 }
241 break;
242 case Intrinsic::coro_id_retcon:
243 case Intrinsic::coro_id_retcon_once:
244 case Intrinsic::coro_id_async:
245 F.setPresplitCoroutine();
246 break;
247 case Intrinsic::coro_resume:
248 lowerResumeOrDestroy(CB&: *CB, Index: CoroSubFnInst::ResumeIndex);
249 break;
250 case Intrinsic::coro_destroy:
251 lowerResumeOrDestroy(CB&: *CB, Index: CoroSubFnInst::DestroyIndex);
252 break;
253 case Intrinsic::coro_promise:
254 lowerCoroPromise(Intrin: cast<CoroPromiseInst>(Val: &I));
255 break;
256 case Intrinsic::coro_done:
257 lowerCoroDone(II: cast<IntrinsicInst>(Val: &I));
258 break;
259 }
260 }
261
262 if (CoroId) {
263 // Make sure that all CoroFree reference the coro.id intrinsic.
264 // Token type is not exposed through coroutine C/C++ builtins to plain C, so
265 // we allow specifying none and fixing it up here.
266 for (CoroFreeInst *CF : CoroFrees)
267 CF->setArgOperand(i: 0, v: CoroId);
268
269 hidePromiseAlloca(CoroId, CoroBegin);
270 }
271
272 // Coroutine suspention could potentially lead to any argument modified
273 // outside of the function, hence arguments should not have noalias
274 // attributes.
275 if (HasCoroSuspend)
276 for (Argument &A : F.args())
277 if (A.hasNoAliasAttr())
278 A.removeAttr(Kind: Attribute::NoAlias);
279}
280
281static bool declaresCoroEarlyIntrinsics(const Module &M) {
282 // coro_suspend omitted as it is overloaded.
283 return coro::declaresIntrinsics(
284 M, List: {Intrinsic::coro_id, Intrinsic::coro_id_retcon,
285 Intrinsic::coro_id_retcon_once, Intrinsic::coro_id_async,
286 Intrinsic::coro_destroy, Intrinsic::coro_done, Intrinsic::coro_end,
287 Intrinsic::coro_end_async, Intrinsic::coro_noop, Intrinsic::coro_free,
288 Intrinsic::coro_promise, Intrinsic::coro_resume});
289}
290
291PreservedAnalyses CoroEarlyPass::run(Module &M, ModuleAnalysisManager &) {
292 if (!declaresCoroEarlyIntrinsics(M))
293 return PreservedAnalyses::all();
294
295 Lowerer L(M);
296 for (auto &F : M)
297 L.lowerEarlyIntrinsics(F);
298
299 PreservedAnalyses PA;
300 PA.preserveSet<CFGAnalyses>();
301 return PA;
302}
303