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>(Val: II->getArgOperand(i: 0)->stripPointerCasts())
144 ->getInitializer());
145 auto *Source = cast<ConstantStruct>(
146 Val: cast<GlobalVariable>(Val: II->getArgOperand(i: 1)->stripPointerCasts())
147 ->getInitializer());
148 auto *TargetSize = Target->getOperand(i_nocapture: 1);
149 auto *SourceSize = Source->getOperand(i_nocapture: 1);
150 if (TargetSize->isElementWiseEqual(Y: SourceSize)) {
151 break;
152 }
153 auto *TargetRelativeFunOffset = Target->getOperand(i_nocapture: 0);
154 auto *NewFuncPtrStruct = ConstantStruct::get(
155 T: Target->getType(), Vs: TargetRelativeFunOffset, Vs: SourceSize);
156 Target->replaceAllUsesWith(V: NewFuncPtrStruct);
157 break;
158 }
159 DeadInsts.insert(Ptr: II);
160 Changed = true;
161 }
162 }
163
164 for (auto *I : DeadInsts)
165 I->eraseFromParent();
166 return Changed;
167}
168
169void Lowerer::lowerCoroNoop(IntrinsicInst *II) {
170 if (!NoopCoro) {
171 LLVMContext &C = Builder.getContext();
172 Module &M = *II->getModule();
173
174 // Create a noop.frame struct type.
175 auto *FnTy = FunctionType::get(Result: Type::getVoidTy(C), Params: Builder.getPtrTy(AddrSpace: 0),
176 /*isVarArg=*/false);
177 auto *FnPtrTy = Builder.getPtrTy(AddrSpace: 0);
178 StructType *FrameTy =
179 StructType::create(Elements: {FnPtrTy, FnPtrTy}, Name: "NoopCoro.Frame");
180
181 // Create a Noop function that does nothing.
182 Function *NoopFn = Function::createWithDefaultAttr(
183 Ty: FnTy, Linkage: GlobalValue::LinkageTypes::InternalLinkage,
184 AddrSpace: M.getDataLayout().getProgramAddressSpace(), N: "__NoopCoro_ResumeDestroy",
185 M: &M);
186 NoopFn->setCallingConv(CallingConv::Fast);
187 buildDebugInfoForNoopResumeDestroyFunc(NoopFn);
188 auto *Entry = BasicBlock::Create(Context&: C, Name: "entry", Parent: NoopFn);
189 ReturnInst::Create(C, InsertAtEnd: Entry);
190
191 // Create a constant struct for the frame.
192 Constant *Values[] = {NoopFn, NoopFn};
193 Constant *NoopCoroConst = ConstantStruct::get(T: FrameTy, V: Values);
194 NoopCoro = new GlobalVariable(
195 M, NoopCoroConst->getType(), /*isConstant=*/true,
196 GlobalVariable::PrivateLinkage, NoopCoroConst, "NoopCoro.Frame.Const");
197 cast<GlobalVariable>(Val: NoopCoro)->setNoSanitizeMetadata();
198 }
199
200 Builder.SetInsertPoint(II);
201 auto *NoopCoroVoidPtr = Builder.CreateBitCast(V: NoopCoro, DestTy: Int8Ptr);
202 II->replaceAllUsesWith(V: NoopCoroVoidPtr);
203}
204
205void NoopCoroElider::run(IntrinsicInst *II) {
206 visitPtr(I&: *II);
207
208 Worklist.clear();
209 VisitedUses.clear();
210}
211
212void NoopCoroElider::visitCallBase(CallBase &CB) {
213 auto *V = U->get();
214 bool ResumeOrDestroy = V == CB.getCalledOperand();
215 if (ResumeOrDestroy) {
216 [[maybe_unused]] bool Success = tryEraseCallInvoke(I: &CB);
217 assert(Success && "Unexpected CallBase");
218
219 auto AboutToDeleteCallback = [this](Value *V) {
220 eraseFromWorklist(I: cast<Instruction>(Val: V));
221 };
222 RecursivelyDeleteTriviallyDeadInstructions(V, TLI: nullptr, MSSAU: nullptr,
223 AboutToDeleteCallback);
224 }
225}
226
227void NoopCoroElider::visitIntrinsicInst(IntrinsicInst &II) {
228 if (auto *SubFn = dyn_cast<CoroSubFnInst>(Val: &II)) {
229 auto *User = SubFn->getUniqueUndroppableUser();
230 assert(User && "Broken module");
231 if (!tryEraseCallInvoke(I: cast<Instruction>(Val: User)))
232 return;
233 SubFn->eraseFromParent();
234 }
235}
236
237bool NoopCoroElider::tryEraseCallInvoke(Instruction *I) {
238 if (auto *Call = dyn_cast<CallInst>(Val: I)) {
239 eraseFromWorklist(I: Call);
240 Call->eraseFromParent();
241 return true;
242 }
243
244 if (auto *II = dyn_cast<InvokeInst>(Val: I)) {
245 Builder.SetInsertPoint(II);
246 Builder.CreateBr(Dest: II->getNormalDest());
247 eraseFromWorklist(I: II);
248 II->getUnwindDest()->removePredecessor(Pred: II->getParent());
249 II->eraseFromParent();
250 return true;
251 }
252 return false;
253}
254
255void NoopCoroElider::eraseFromWorklist(Instruction *I) {
256 erase_if(C&: Worklist, P: [I](UseToVisit &U) {
257 return I == U.UseAndIsOffsetKnown.getPointer()->getUser();
258 });
259}
260
261static bool declaresCoroCleanupIntrinsics(const Module &M) {
262 return coro::declaresIntrinsics(
263 M, List: {Intrinsic::coro_alloc, Intrinsic::coro_begin,
264 Intrinsic::coro_subfn_addr, Intrinsic::coro_free,
265 Intrinsic::coro_dead, Intrinsic::coro_id, Intrinsic::coro_id_retcon,
266 Intrinsic::coro_id_async, Intrinsic::coro_id_retcon_once,
267 Intrinsic::coro_noop, Intrinsic::coro_async_size_replace,
268 Intrinsic::coro_async_resume, Intrinsic::coro_begin_custom_abi});
269}
270
271PreservedAnalyses CoroCleanupPass::run(Module &M,
272 ModuleAnalysisManager &MAM) {
273 if (!declaresCoroCleanupIntrinsics(M))
274 return PreservedAnalyses::all();
275
276 FunctionAnalysisManager &FAM =
277 MAM.getResult<FunctionAnalysisManagerModuleProxy>(IR&: M).getManager();
278
279 FunctionPassManager FPM;
280 FPM.addPass(Pass: SimplifyCFGPass());
281
282 PreservedAnalyses FuncPA;
283 FuncPA.preserveSet<CFGAnalyses>();
284
285 Lowerer L(M);
286 for (auto &F : M) {
287 if (L.lower(F)) {
288 FAM.invalidate(IR&: F, PA: FuncPA);
289 FPM.run(IR&: F, AM&: FAM);
290 }
291 }
292
293 return PreservedAnalyses::none();
294}
295