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 auto *FrameTy = Int8Ptr;
88
89 Builder.SetInsertPoint(II);
90 auto *Load = Builder.CreateLoad(Ty: FrameTy, Ptr: Operand);
91 auto *Cond = Builder.CreateICmpEQ(LHS: Load, RHS: NullPtr);
92
93 II->replaceAllUsesWith(V: Cond);
94 II->eraseFromParent();
95}
96
97// Later middle-end passes will assume promise alloca dead after coroutine
98// suspend, leading to misoptimizations. We hide promise alloca using
99// coro.promise and will lower it back to alloca at CoroSplit.
100void Lowerer::hidePromiseAlloca(CoroIdInst *CoroId, CoroBeginInst *CoroBegin) {
101 auto *PA = CoroId->getPromise();
102 if (!PA || !CoroBegin)
103 return;
104 Builder.SetInsertPoint(*CoroBegin->getInsertionPointAfterDef());
105
106 auto *Alignment = Builder.getInt32(C: PA->getAlign().value());
107 auto *FromPromise = Builder.getInt1(V: false);
108 SmallVector<Value *, 3> Arg{CoroBegin, Alignment, FromPromise};
109 auto *PI = Builder.CreateIntrinsic(
110 RetTy: Builder.getPtrTy(), ID: Intrinsic::coro_promise, Args: Arg, FMFSource: {}, Name: "promise.addr");
111 PI->setCannotDuplicate();
112 // Remove lifetime markers, as these are only allowed on allocas.
113 for (User *U : make_early_inc_range(Range: PA->users())) {
114 auto *I = cast<Instruction>(Val: U);
115 if (I->isLifetimeStartOrEnd())
116 I->eraseFromParent();
117 }
118 PA->replaceUsesWithIf(New: PI, ShouldReplace: [CoroId](Use &U) {
119 bool IsBitcast = U == U.getUser()->stripPointerCasts();
120 bool IsCoroId = U.getUser() == CoroId;
121 return !IsBitcast && !IsCoroId;
122 });
123}
124
125// Prior to CoroSplit, calls to coro.begin needs to be marked as NoDuplicate,
126// as CoroSplit assumes there is exactly one coro.begin. After CoroSplit,
127// NoDuplicate attribute will be removed from coro.begin otherwise, it will
128// interfere with inlining.
129static void setCannotDuplicate(CoroIdInst *CoroId) {
130 for (User *U : CoroId->users())
131 if (auto *CB = dyn_cast<CoroBeginInst>(Val: U))
132 CB->setCannotDuplicate();
133}
134
135void Lowerer::lowerEarlyIntrinsics(Function &F) {
136 CoroIdInst *CoroId = nullptr;
137 CoroBeginInst *CoroBegin = nullptr;
138 SmallVector<CoroFreeInst *, 4> CoroFrees;
139 bool HasCoroSuspend = false;
140 for (Instruction &I : llvm::make_early_inc_range(Range: instructions(F))) {
141 auto *CB = dyn_cast<CallBase>(Val: &I);
142 if (!CB)
143 continue;
144
145 switch (CB->getIntrinsicID()) {
146 default:
147 continue;
148 case Intrinsic::coro_begin:
149 case Intrinsic::coro_begin_custom_abi:
150 if (CoroBegin)
151 report_fatal_error(
152 reason: "coroutine should have exactly one defining @llvm.coro.begin");
153 CoroBegin = cast<CoroBeginInst>(Val: &I);
154 break;
155 case Intrinsic::coro_free:
156 CoroFrees.push_back(Elt: cast<CoroFreeInst>(Val: &I));
157 break;
158 case Intrinsic::coro_suspend:
159 // Make sure that final suspend point is not duplicated as CoroSplit
160 // pass expects that there is at most one final suspend point.
161 if (cast<CoroSuspendInst>(Val: &I)->isFinal())
162 CB->setCannotDuplicate();
163 HasCoroSuspend = true;
164 break;
165 case Intrinsic::coro_end_async:
166 case Intrinsic::coro_end:
167 // Make sure that fallthrough coro.end is not duplicated as CoroSplit
168 // pass expects that there is at most one fallthrough coro.end.
169 if (cast<AnyCoroEndInst>(Val: &I)->isFallthrough())
170 CB->setCannotDuplicate();
171 break;
172 case Intrinsic::coro_id:
173 if (auto *CII = cast<CoroIdInst>(Val: &I)) {
174 if (CII->getInfo().isPreSplit()) {
175 assert(F.isPresplitCoroutine() &&
176 "The frontend uses Switch-Resumed ABI should emit "
177 "\"presplitcoroutine\" attribute for the coroutine.");
178 setCannotDuplicate(CII);
179 CII->setCoroutineSelf();
180 CoroId = cast<CoroIdInst>(Val: &I);
181 }
182 }
183 break;
184 case Intrinsic::coro_id_retcon:
185 case Intrinsic::coro_id_retcon_once:
186 case Intrinsic::coro_id_async:
187 F.setPresplitCoroutine();
188 break;
189 case Intrinsic::coro_resume:
190 lowerResumeOrDestroy(CB&: *CB, Index: CoroSubFnInst::ResumeIndex);
191 break;
192 case Intrinsic::coro_destroy:
193 lowerResumeOrDestroy(CB&: *CB, Index: CoroSubFnInst::DestroyIndex);
194 break;
195 case Intrinsic::coro_promise:
196 lowerCoroPromise(Intrin: cast<CoroPromiseInst>(Val: &I));
197 break;
198 case Intrinsic::coro_done:
199 lowerCoroDone(II: cast<IntrinsicInst>(Val: &I));
200 break;
201 }
202 }
203
204 if (CoroId) {
205 // Make sure that all CoroFree reference the coro.id intrinsic.
206 // Token type is not exposed through coroutine C/C++ builtins to plain C, so
207 // we allow specifying none and fixing it up here.
208 for (CoroFreeInst *CF : CoroFrees)
209 CF->setArgOperand(i: 0, v: CoroId);
210
211 hidePromiseAlloca(CoroId, CoroBegin);
212 }
213
214 // Coroutine suspention could potentially lead to any argument modified
215 // outside of the function, hence arguments should not have noalias
216 // attributes.
217 if (HasCoroSuspend)
218 for (Argument &A : F.args())
219 if (A.hasNoAliasAttr())
220 A.removeAttr(Kind: Attribute::NoAlias);
221}
222
223static bool declaresCoroEarlyIntrinsics(const Module &M) {
224 // coro_suspend omitted as it is overloaded.
225 return coro::declaresIntrinsics(
226 M, List: {Intrinsic::coro_id, Intrinsic::coro_id_retcon,
227 Intrinsic::coro_id_retcon_once, Intrinsic::coro_id_async,
228 Intrinsic::coro_destroy, Intrinsic::coro_done, Intrinsic::coro_end,
229 Intrinsic::coro_end_async, Intrinsic::coro_free,
230 Intrinsic::coro_promise, Intrinsic::coro_resume});
231}
232
233PreservedAnalyses CoroEarlyPass::run(Module &M, ModuleAnalysisManager &) {
234 if (!declaresCoroEarlyIntrinsics(M))
235 return PreservedAnalyses::all();
236
237 Lowerer L(M);
238 for (auto &F : M)
239 L.lowerEarlyIntrinsics(F);
240
241 PreservedAnalyses PA;
242 PA.preserveSet<CFGAnalyses>();
243 return PA;
244}
245