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/Function.h"
12#include "llvm/IR/IRBuilder.h"
13#include "llvm/IR/InstIterator.h"
14#include "llvm/IR/Module.h"
15#include "llvm/Transforms/Coroutines/CoroShape.h"
16
17using namespace llvm;
18
19#define DEBUG_TYPE "coro-early"
20
21namespace {
22// Created on demand if the coro-early pass has work to do.
23class Lowerer : public coro::LowererBase {
24 IRBuilder<> Builder;
25 PointerType *const AnyResumeFnPtrTy;
26
27 void lowerResumeOrDestroy(CallBase &CB, CoroSubFnInst::ResumeKind);
28 void lowerCoroPromise(CoroPromiseInst *Intrin);
29 void lowerCoroDone(IntrinsicInst *II);
30 void hidePromiseAlloca(CoroIdInst *CoroId, CoroBeginInst *CoroBegin);
31
32public:
33 Lowerer(Module &M)
34 : LowererBase(M), Builder(Context),
35 AnyResumeFnPtrTy(PointerType::getUnqual(C&: Context)) {}
36 void lowerEarlyIntrinsics(Function &F);
37};
38} // namespace
39
40// Replace a direct call to coro.resume or coro.destroy with an indirect call to
41// an address returned by coro.subfn.addr intrinsic. This is done so that
42// CGPassManager recognizes devirtualization when CoroElide pass replaces a call
43// to coro.subfn.addr with an appropriate function address.
44void Lowerer::lowerResumeOrDestroy(CallBase &CB,
45 CoroSubFnInst::ResumeKind Index) {
46 Value *ResumeAddr = makeSubFnCall(Arg: CB.getArgOperand(i: 0), Index, InsertPt: &CB);
47 CB.setCalledOperand(ResumeAddr);
48 CB.setCallingConv(CallingConv::Fast);
49}
50
51// Coroutine promise field is always at the fixed offset from the beginning of
52// the coroutine frame. i8* coro.promise(i8*, i1 from) intrinsic adds an offset
53// to a passed pointer to move from coroutine frame to coroutine promise and
54// vice versa. Since we don't know exactly which coroutine frame it is, we build
55// a coroutine frame mock up starting with two function pointers, followed by a
56// properly aligned coroutine promise field.
57// TODO: Handle the case when coroutine promise alloca has align override.
58void Lowerer::lowerCoroPromise(CoroPromiseInst *Intrin) {
59 Value *Operand = Intrin->getArgOperand(i: 0);
60 Align Alignment = Intrin->getAlignment();
61 Type *Int8Ty = Builder.getInt8Ty();
62
63 auto *SampleStruct =
64 StructType::get(Context, Elements: {AnyResumeFnPtrTy, AnyResumeFnPtrTy, Int8Ty});
65 const DataLayout &DL = TheModule.getDataLayout();
66 int64_t Offset = alignTo(
67 Size: DL.getStructLayout(Ty: SampleStruct)->getElementOffset(Idx: 2), A: Alignment);
68 if (Intrin->isFromPromise())
69 Offset = -Offset;
70
71 Builder.SetInsertPoint(Intrin);
72 Value *Replacement =
73 Builder.CreateConstInBoundsGEP1_32(Ty: Int8Ty, Ptr: Operand, Idx0: Offset);
74
75 Intrin->replaceAllUsesWith(V: Replacement);
76 Intrin->eraseFromParent();
77}
78
79// When a coroutine reaches final suspend point, it zeros out ResumeFnAddr in
80// the coroutine frame (it is UB to resume from a final suspend point).
81// The llvm.coro.done intrinsic is used to check whether a coroutine is
82// suspended at the final suspend point or not.
83void Lowerer::lowerCoroDone(IntrinsicInst *II) {
84 Value *Operand = II->getArgOperand(i: 0);
85
86 // ResumeFnAddr is the first pointer sized element of the coroutine frame.
87 static_assert(coro::Shape::SwitchFieldIndex::Resume == 0,
88 "resume function not at offset zero");
89 auto *FrameTy = Int8Ptr;
90
91 Builder.SetInsertPoint(II);
92 auto *Load = Builder.CreateLoad(Ty: FrameTy, Ptr: Operand);
93 auto *Cond = Builder.CreateICmpEQ(LHS: Load, RHS: NullPtr);
94
95 II->replaceAllUsesWith(V: Cond);
96 II->eraseFromParent();
97}
98
99// Later middle-end passes will assume promise alloca dead after coroutine
100// suspend, leading to misoptimizations. We hide promise alloca using
101// coro.promise and will lower it back to alloca at CoroSplit.
102void Lowerer::hidePromiseAlloca(CoroIdInst *CoroId, CoroBeginInst *CoroBegin) {
103 auto *PA = CoroId->getPromise();
104 if (!PA || !CoroBegin)
105 return;
106 Builder.SetInsertPoint(*CoroBegin->getInsertionPointAfterDef());
107
108 auto *Alignment = Builder.getInt32(C: PA->getAlign().value());
109 auto *FromPromise = Builder.getInt1(V: false);
110 SmallVector<Value *, 3> Arg{CoroBegin, Alignment, FromPromise};
111 auto *PI = Builder.CreateIntrinsic(
112 RetTy: Builder.getPtrTy(), ID: Intrinsic::coro_promise, Args: Arg, FMFSource: {}, Name: "promise.addr");
113 PI->setCannotDuplicate();
114 // Remove lifetime markers, as these are only allowed on allocas.
115 for (User *U : make_early_inc_range(Range: PA->users())) {
116 auto *I = cast<Instruction>(Val: U);
117 if (I->isLifetimeStartOrEnd())
118 I->eraseFromParent();
119 }
120 PA->replaceUsesWithIf(New: PI, ShouldReplace: [CoroId](Use &U) {
121 bool IsBitcast = U == U.getUser()->stripPointerCasts();
122 bool IsCoroId = U.getUser() == CoroId;
123 return !IsBitcast && !IsCoroId;
124 });
125}
126
127// Prior to CoroSplit, calls to coro.begin needs to be marked as NoDuplicate,
128// as CoroSplit assumes there is exactly one coro.begin. After CoroSplit,
129// NoDuplicate attribute will be removed from coro.begin otherwise, it will
130// interfere with inlining.
131static void setCannotDuplicate(CoroIdInst *CoroId) {
132 for (User *U : CoroId->users())
133 if (auto *CB = dyn_cast<CoroBeginInst>(Val: U))
134 CB->setCannotDuplicate();
135}
136
137void Lowerer::lowerEarlyIntrinsics(Function &F) {
138 CoroIdInst *CoroId = nullptr;
139 CoroBeginInst *CoroBegin = nullptr;
140 SmallVector<CoroFreeInst *, 4> CoroFrees;
141 bool HasCoroSuspend = false;
142 for (Instruction &I : llvm::make_early_inc_range(Range: instructions(F))) {
143 auto *CB = dyn_cast<CallBase>(Val: &I);
144 if (!CB)
145 continue;
146
147 switch (CB->getIntrinsicID()) {
148 default:
149 continue;
150 case Intrinsic::coro_begin:
151 case Intrinsic::coro_begin_custom_abi:
152 if (CoroBegin)
153 report_fatal_error(
154 reason: "coroutine should have exactly one defining @llvm.coro.begin");
155 CoroBegin = cast<CoroBeginInst>(Val: &I);
156 break;
157 case Intrinsic::coro_free:
158 CoroFrees.push_back(Elt: cast<CoroFreeInst>(Val: &I));
159 break;
160 case Intrinsic::coro_suspend:
161 // Make sure that final suspend point is not duplicated as CoroSplit
162 // pass expects that there is at most one final suspend point.
163 if (cast<CoroSuspendInst>(Val: &I)->isFinal())
164 CB->setCannotDuplicate();
165 HasCoroSuspend = true;
166 break;
167 case Intrinsic::coro_end_async:
168 case Intrinsic::coro_end:
169 // Make sure that fallthrough coro.end is not duplicated as CoroSplit
170 // pass expects that there is at most one fallthrough coro.end.
171 if (cast<AnyCoroEndInst>(Val: &I)->isFallthrough())
172 CB->setCannotDuplicate();
173 break;
174 case Intrinsic::coro_id:
175 if (auto *CII = cast<CoroIdInst>(Val: &I)) {
176 if (CII->getInfo().isPreSplit()) {
177 assert(F.isPresplitCoroutine() &&
178 "The frontend uses Switch-Resumed ABI should emit "
179 "\"presplitcoroutine\" attribute for the coroutine.");
180 setCannotDuplicate(CII);
181 CII->setCoroutineSelf();
182 CoroId = cast<CoroIdInst>(Val: &I);
183 }
184 }
185 break;
186 case Intrinsic::coro_id_retcon:
187 case Intrinsic::coro_id_retcon_once:
188 case Intrinsic::coro_id_async:
189 F.setPresplitCoroutine();
190 break;
191 case Intrinsic::coro_resume:
192 lowerResumeOrDestroy(CB&: *CB, Index: CoroSubFnInst::ResumeIndex);
193 break;
194 case Intrinsic::coro_destroy:
195 lowerResumeOrDestroy(CB&: *CB, Index: CoroSubFnInst::DestroyIndex);
196 break;
197 case Intrinsic::coro_promise:
198 lowerCoroPromise(Intrin: cast<CoroPromiseInst>(Val: &I));
199 break;
200 case Intrinsic::coro_done:
201 lowerCoroDone(II: cast<IntrinsicInst>(Val: &I));
202 break;
203 }
204 }
205
206 if (CoroId) {
207 // Make sure that all CoroFree reference the coro.id intrinsic.
208 // Token type is not exposed through coroutine C/C++ builtins to plain C, so
209 // we allow specifying none and fixing it up here.
210 for (CoroFreeInst *CF : CoroFrees)
211 CF->setArgOperand(i: 0, v: CoroId);
212
213 hidePromiseAlloca(CoroId, CoroBegin);
214 }
215
216 // Coroutine suspention could potentially lead to any argument modified
217 // outside of the function, hence arguments should not have noalias
218 // attributes.
219 if (HasCoroSuspend)
220 for (Argument &A : F.args())
221 if (A.hasNoAliasAttr())
222 A.removeAttr(Kind: Attribute::NoAlias);
223}
224
225static bool declaresCoroEarlyIntrinsics(const Module &M) {
226 // coro_suspend omitted as it is overloaded.
227 return coro::declaresIntrinsics(
228 M, List: {Intrinsic::coro_id, Intrinsic::coro_id_retcon,
229 Intrinsic::coro_id_retcon_once, Intrinsic::coro_id_async,
230 Intrinsic::coro_destroy, Intrinsic::coro_done, Intrinsic::coro_end,
231 Intrinsic::coro_end_async, Intrinsic::coro_free,
232 Intrinsic::coro_promise, Intrinsic::coro_resume});
233}
234
235PreservedAnalyses CoroEarlyPass::run(Module &M, ModuleAnalysisManager &) {
236 if (!declaresCoroEarlyIntrinsics(M))
237 return PreservedAnalyses::all();
238
239 Lowerer L(M);
240 for (auto &F : M)
241 L.lowerEarlyIntrinsics(F);
242
243 PreservedAnalyses PA;
244 PA.preserveSet<CFGAnalyses>();
245 return PA;
246}
247