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