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