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/IR/IRBuilder.h"
12 #include "llvm/IR/InstIterator.h"
13 #include "llvm/IR/PassManager.h"
14 #include "llvm/IR/Function.h"
15 #include "llvm/Transforms/Scalar/SimplifyCFG.h"
16 
17 using namespace llvm;
18 
19 #define DEBUG_TYPE "coro-cleanup"
20 
21 namespace {
22 // Created on demand if CoroCleanup pass has work to do.
23 struct Lowerer : coro::LowererBase {
24   IRBuilder<> Builder;
25   Lowerer(Module &M) : LowererBase(M), Builder(Context) {}
26   bool lower(Function &F);
27 };
28 }
29 
30 static void lowerSubFn(IRBuilder<> &Builder, CoroSubFnInst *SubFn) {
31   Builder.SetInsertPoint(SubFn);
32   Value *FrameRaw = SubFn->getFrame();
33   int Index = SubFn->getIndex();
34 
35   auto *FrameTy = StructType::get(
36       SubFn->getContext(), {Builder.getInt8PtrTy(), Builder.getInt8PtrTy()});
37   PointerType *FramePtrTy = FrameTy->getPointerTo();
38 
39   Builder.SetInsertPoint(SubFn);
40   auto *FramePtr = Builder.CreateBitCast(FrameRaw, FramePtrTy);
41   auto *Gep = Builder.CreateConstInBoundsGEP2_32(FrameTy, FramePtr, 0, Index);
42   auto *Load = Builder.CreateLoad(FrameTy->getElementType(Index), Gep);
43 
44   SubFn->replaceAllUsesWith(Load);
45 }
46 
47 bool Lowerer::lower(Function &F) {
48   bool IsPrivateAndUnprocessed =
49       F.hasFnAttribute(CORO_PRESPLIT_ATTR) && F.hasLocalLinkage();
50   bool Changed = false;
51 
52   for (Instruction &I : llvm::make_early_inc_range(instructions(F))) {
53     if (auto *II = dyn_cast<IntrinsicInst>(&I)) {
54       switch (II->getIntrinsicID()) {
55       default:
56         continue;
57       case Intrinsic::coro_begin:
58         II->replaceAllUsesWith(II->getArgOperand(1));
59         break;
60       case Intrinsic::coro_free:
61         II->replaceAllUsesWith(II->getArgOperand(1));
62         break;
63       case Intrinsic::coro_alloc:
64         II->replaceAllUsesWith(ConstantInt::getTrue(Context));
65         break;
66       case Intrinsic::coro_async_resume:
67         II->replaceAllUsesWith(
68             ConstantPointerNull::get(cast<PointerType>(I.getType())));
69         break;
70       case Intrinsic::coro_id:
71       case Intrinsic::coro_id_retcon:
72       case Intrinsic::coro_id_retcon_once:
73       case Intrinsic::coro_id_async:
74         II->replaceAllUsesWith(ConstantTokenNone::get(Context));
75         break;
76       case Intrinsic::coro_subfn_addr:
77         lowerSubFn(Builder, cast<CoroSubFnInst>(II));
78         break;
79       case Intrinsic::coro_end:
80       case Intrinsic::coro_suspend_retcon:
81         if (IsPrivateAndUnprocessed) {
82           II->replaceAllUsesWith(UndefValue::get(II->getType()));
83         } else
84           continue;
85         break;
86       case Intrinsic::coro_async_size_replace:
87         auto *Target = cast<ConstantStruct>(
88             cast<GlobalVariable>(II->getArgOperand(0)->stripPointerCasts())
89                 ->getInitializer());
90         auto *Source = cast<ConstantStruct>(
91             cast<GlobalVariable>(II->getArgOperand(1)->stripPointerCasts())
92                 ->getInitializer());
93         auto *TargetSize = Target->getOperand(1);
94         auto *SourceSize = Source->getOperand(1);
95         if (TargetSize->isElementWiseEqual(SourceSize)) {
96           break;
97         }
98         auto *TargetRelativeFunOffset = Target->getOperand(0);
99         auto *NewFuncPtrStruct = ConstantStruct::get(
100             Target->getType(), TargetRelativeFunOffset, SourceSize);
101         Target->replaceAllUsesWith(NewFuncPtrStruct);
102         break;
103       }
104       II->eraseFromParent();
105       Changed = true;
106     }
107   }
108 
109   return Changed;
110 }
111 
112 static bool declaresCoroCleanupIntrinsics(const Module &M) {
113   return coro::declaresIntrinsics(
114       M, {"llvm.coro.alloc", "llvm.coro.begin", "llvm.coro.subfn.addr",
115           "llvm.coro.free", "llvm.coro.id", "llvm.coro.id.retcon",
116           "llvm.coro.id.retcon.once", "llvm.coro.async.size.replace",
117           "llvm.coro.async.resume"});
118 }
119 
120 PreservedAnalyses CoroCleanupPass::run(Module &M,
121                                        ModuleAnalysisManager &MAM) {
122   if (!declaresCoroCleanupIntrinsics(M))
123     return PreservedAnalyses::all();
124 
125   FunctionAnalysisManager &FAM =
126       MAM.getResult<FunctionAnalysisManagerModuleProxy>(M).getManager();
127 
128   FunctionPassManager FPM;
129   FPM.addPass(SimplifyCFGPass());
130 
131   Lowerer L(M);
132   for (auto &F : M)
133     if (L.lower(F))
134       FPM.run(F, FAM);
135 
136   return PreservedAnalyses::none();
137 }
138