1//===- CoroCleanup.cpp - Coroutine Cleanup 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/CoroCleanup.h"
10#include "CoroInternal.h"
11#include "llvm/Analysis/PtrUseVisitor.h"
12#include "llvm/IR/DIBuilder.h"
13#include "llvm/IR/Function.h"
14#include "llvm/IR/IRBuilder.h"
15#include "llvm/IR/InstIterator.h"
16#include "llvm/IR/Module.h"
17#include "llvm/IR/PassManager.h"
18#include "llvm/Transforms/Scalar/SimplifyCFG.h"
19#include "llvm/Transforms/Utils/Local.h"
20
21using namespace llvm;
22
23#define DEBUG_TYPE "coro-cleanup"
24
25namespace {
26// Created on demand if CoroCleanup pass has work to do.
27struct Lowerer : coro::LowererBase {
28 IRBuilder<> Builder;
29 Constant *NoopCoro = nullptr;
30
31 Lowerer(Module &M) : LowererBase(M), Builder(Context) {}
32 bool lower(Function &F);
33
34private:
35 void lowerCoroNoop(IntrinsicInst *II);
36};
37
38// Recursively walk and eliminate resume/destroy call on noop coro
39class NoopCoroElider : public PtrUseVisitor<NoopCoroElider> {
40 using Base = PtrUseVisitor<NoopCoroElider>;
41
42 IRBuilder<> Builder;
43
44public:
45 NoopCoroElider(const DataLayout &DL, LLVMContext &C) : Base(DL), Builder(C) {}
46
47 void run(IntrinsicInst *II);
48
49 void visitLoadInst(LoadInst &I) { enqueueUsers(I); }
50 void visitCallBase(CallBase &CB);
51 void visitIntrinsicInst(IntrinsicInst &II);
52
53private:
54 bool tryEraseCallInvoke(Instruction *I);
55 void eraseFromWorklist(Instruction *I);
56};
57}
58
59static void lowerSubFn(IRBuilder<> &Builder, CoroSubFnInst *SubFn) {
60 Builder.SetInsertPoint(SubFn);
61 Value *FramePtr = SubFn->getFrame();
62 int Index = SubFn->getIndex();
63
64 auto *FrameTy = StructType::get(Context&: SubFn->getContext(),
65 Elements: {Builder.getPtrTy(), Builder.getPtrTy()});
66
67 Builder.SetInsertPoint(SubFn);
68 auto *Gep = Builder.CreateConstInBoundsGEP2_32(Ty: FrameTy, Ptr: FramePtr, Idx0: 0, Idx1: Index);
69 auto *Load = Builder.CreateLoad(Ty: FrameTy->getElementType(N: Index), Ptr: Gep);
70
71 SubFn->replaceAllUsesWith(V: Load);
72}
73
74static void buildDebugInfoForNoopResumeDestroyFunc(Function *NoopFn) {
75 Module &M = *NoopFn->getParent();
76 if (M.debug_compile_units().empty())
77 return;
78
79 DICompileUnit *CU = *M.debug_compile_units_begin();
80 DIBuilder DB(M, /*AllowUnresolved*/ false, CU);
81 std::array<Metadata *, 2> Params{nullptr, nullptr};
82 auto *SubroutineType =
83 DB.createSubroutineType(ParameterTypes: DB.getOrCreateTypeArray(Elements: Params));
84 StringRef Name = NoopFn->getName();
85 auto *SP = DB.createFunction(
86 Scope: CU, /*Name=*/Name, /*LinkageName=*/Name, /*File=*/CU->getFile(),
87 /*LineNo=*/0, Ty: SubroutineType, /*ScopeLine=*/0, Flags: DINode::FlagArtificial,
88 SPFlags: DISubprogram::SPFlagDefinition);
89 NoopFn->setSubprogram(SP);
90 DB.finalize();
91}
92
93bool Lowerer::lower(Function &F) {
94 bool IsPrivateAndUnprocessed = F.isPresplitCoroutine() && F.hasLocalLinkage();
95 bool Changed = false;
96
97 NoopCoroElider NCE(F.getDataLayout(), F.getContext());
98 SmallPtrSet<Instruction *, 8> DeadInsts{};
99 for (Instruction &I : instructions(F)) {
100 if (auto *II = dyn_cast<IntrinsicInst>(Val: &I)) {
101 switch (II->getIntrinsicID()) {
102 default:
103 continue;
104 case Intrinsic::coro_begin:
105 case Intrinsic::coro_begin_custom_abi:
106 II->replaceAllUsesWith(V: II->getArgOperand(i: 1));
107 break;
108 case Intrinsic::coro_free:
109 II->replaceAllUsesWith(V: II->getArgOperand(i: 1));
110 break;
111 case Intrinsic::coro_dead:
112 break;
113 case Intrinsic::coro_alloc:
114 II->replaceAllUsesWith(V: ConstantInt::getTrue(Context));
115 break;
116 case Intrinsic::coro_async_resume:
117 II->replaceAllUsesWith(
118 V: ConstantPointerNull::get(T: cast<PointerType>(Val: I.getType())));
119 break;
120 case Intrinsic::coro_id:
121 case Intrinsic::coro_id_retcon:
122 case Intrinsic::coro_id_retcon_once:
123 case Intrinsic::coro_id_async:
124 II->replaceAllUsesWith(V: ConstantTokenNone::get(Context));
125 break;
126 case Intrinsic::coro_noop:
127 NCE.run(II);
128 if (!II->user_empty())
129 lowerCoroNoop(II);
130 break;
131 case Intrinsic::coro_subfn_addr:
132 lowerSubFn(Builder, SubFn: cast<CoroSubFnInst>(Val: II));
133 break;
134 case Intrinsic::coro_suspend_retcon:
135 case Intrinsic::coro_is_in_ramp:
136 if (IsPrivateAndUnprocessed) {
137 II->replaceAllUsesWith(V: PoisonValue::get(T: II->getType()));
138 } else
139 continue;
140 break;
141 case Intrinsic::coro_async_size_replace:
142 auto *Target = cast<ConstantStruct>(
143 Val: cast<GlobalVariable>(
144 Val: II->getArgOperand(i: 0)->stripPointerCastsAndAliases())
145 ->getInitializer());
146 auto *Source = cast<ConstantStruct>(
147 Val: cast<GlobalVariable>(
148 Val: II->getArgOperand(i: 1)->stripPointerCastsAndAliases())
149 ->getInitializer());
150 auto *TargetSize = Target->getOperand(i_nocapture: 1);
151 auto *SourceSize = Source->getOperand(i_nocapture: 1);
152 if (TargetSize->isElementWiseEqual(Y: SourceSize)) {
153 break;
154 }
155 auto *TargetRelativeFunOffset = Target->getOperand(i_nocapture: 0);
156 auto *NewFuncPtrStruct = ConstantStruct::get(
157 T: Target->getType(), Vs: TargetRelativeFunOffset, Vs: SourceSize);
158 Target->replaceAllUsesWith(V: NewFuncPtrStruct);
159 break;
160 }
161 DeadInsts.insert(Ptr: II);
162 Changed = true;
163 }
164 }
165
166 for (auto *I : DeadInsts)
167 I->eraseFromParent();
168 return Changed;
169}
170
171void Lowerer::lowerCoroNoop(IntrinsicInst *II) {
172 if (!NoopCoro) {
173 LLVMContext &C = Builder.getContext();
174 Module &M = *II->getModule();
175
176 // Create a noop.frame struct type.
177 auto *FnTy = FunctionType::get(Result: Type::getVoidTy(C), Params: Builder.getPtrTy(AddrSpace: 0),
178 /*isVarArg=*/false);
179 auto *FnPtrTy = Builder.getPtrTy(AddrSpace: 0);
180 StructType *FrameTy =
181 StructType::create(Elements: {FnPtrTy, FnPtrTy}, Name: "NoopCoro.Frame");
182
183 // Create a Noop function that does nothing.
184 Function *NoopFn = Function::createWithDefaultAttr(
185 Ty: FnTy, Linkage: GlobalValue::LinkageTypes::InternalLinkage,
186 AddrSpace: M.getDataLayout().getProgramAddressSpace(), N: "__NoopCoro_ResumeDestroy",
187 M: &M);
188 buildDebugInfoForNoopResumeDestroyFunc(NoopFn);
189 auto *Entry = BasicBlock::Create(Context&: C, Name: "entry", Parent: NoopFn);
190 ReturnInst::Create(C, InsertAtEnd: Entry);
191
192 // Create a constant struct for the frame.
193 Constant *Values[] = {NoopFn, NoopFn};
194 Constant *NoopCoroConst = ConstantStruct::get(T: FrameTy, V: Values);
195 NoopCoro = new GlobalVariable(
196 M, NoopCoroConst->getType(), /*isConstant=*/true,
197 GlobalVariable::PrivateLinkage, NoopCoroConst, "NoopCoro.Frame.Const");
198 cast<GlobalVariable>(Val: NoopCoro)->setNoSanitizeMetadata();
199 }
200
201 Builder.SetInsertPoint(II);
202 auto *NoopCoroVoidPtr = Builder.CreateBitCast(V: NoopCoro, DestTy: Int8Ptr);
203 II->replaceAllUsesWith(V: NoopCoroVoidPtr);
204}
205
206void NoopCoroElider::run(IntrinsicInst *II) {
207 visitPtr(I&: *II);
208
209 Worklist.clear();
210 VisitedUses.clear();
211}
212
213void NoopCoroElider::visitCallBase(CallBase &CB) {
214 auto *V = U->get();
215 bool ResumeOrDestroy = V == CB.getCalledOperand();
216 if (ResumeOrDestroy) {
217 [[maybe_unused]] bool Success = tryEraseCallInvoke(I: &CB);
218 assert(Success && "Unexpected CallBase");
219
220 auto AboutToDeleteCallback = [this](Value *V) {
221 eraseFromWorklist(I: cast<Instruction>(Val: V));
222 };
223 RecursivelyDeleteTriviallyDeadInstructions(V, TLI: nullptr, MSSAU: nullptr,
224 AboutToDeleteCallback);
225 }
226}
227
228void NoopCoroElider::visitIntrinsicInst(IntrinsicInst &II) {
229 if (auto *SubFn = dyn_cast<CoroSubFnInst>(Val: &II)) {
230 auto *User = SubFn->getUniqueUndroppableUser();
231 assert(User && "Broken module");
232 if (!tryEraseCallInvoke(I: cast<Instruction>(Val: User)))
233 return;
234 SubFn->eraseFromParent();
235 }
236}
237
238bool NoopCoroElider::tryEraseCallInvoke(Instruction *I) {
239 if (auto *Call = dyn_cast<CallInst>(Val: I)) {
240 eraseFromWorklist(I: Call);
241 Call->eraseFromParent();
242 return true;
243 }
244
245 if (auto *II = dyn_cast<InvokeInst>(Val: I)) {
246 Builder.SetInsertPoint(II);
247 Builder.CreateBr(Dest: II->getNormalDest());
248 eraseFromWorklist(I: II);
249 II->getUnwindDest()->removePredecessor(Pred: II->getParent());
250 II->eraseFromParent();
251 return true;
252 }
253 return false;
254}
255
256void NoopCoroElider::eraseFromWorklist(Instruction *I) {
257 erase_if(C&: Worklist, P: [I](UseToVisit &U) {
258 return I == U.UseAndIsOffsetKnown.getPointer()->getUser();
259 });
260}
261
262static bool declaresCoroCleanupIntrinsics(const Module &M) {
263 return coro::declaresIntrinsics(
264 M, List: {Intrinsic::coro_alloc, Intrinsic::coro_begin,
265 Intrinsic::coro_subfn_addr, Intrinsic::coro_free,
266 Intrinsic::coro_dead, Intrinsic::coro_id, Intrinsic::coro_id_retcon,
267 Intrinsic::coro_id_async, Intrinsic::coro_id_retcon_once,
268 Intrinsic::coro_noop, Intrinsic::coro_async_size_replace,
269 Intrinsic::coro_async_resume, Intrinsic::coro_begin_custom_abi});
270}
271
272PreservedAnalyses CoroCleanupPass::run(Module &M,
273 ModuleAnalysisManager &MAM) {
274 if (!declaresCoroCleanupIntrinsics(M))
275 return PreservedAnalyses::all();
276
277 FunctionAnalysisManager &FAM =
278 MAM.getResult<FunctionAnalysisManagerModuleProxy>(IR&: M).getManager();
279
280 FunctionPassManager FPM;
281 FPM.addPass(Pass: SimplifyCFGPass());
282
283 PreservedAnalyses FuncPA;
284 FuncPA.preserveSet<CFGAnalyses>();
285
286 Lowerer L(M);
287 for (auto &F : M) {
288 if (L.lower(F)) {
289 FAM.invalidate(IR&: F, PA: FuncPA);
290 FPM.run(IR&: F, AM&: FAM);
291 }
292 }
293
294 return PreservedAnalyses::none();
295}
296