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