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/LegacyPassManager.h"
14 #include "llvm/Pass.h"
15 #include "llvm/Transforms/Scalar.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 lowerRemainingCoroIntrinsics(Function &F);
27 };
28 }
29 
30 static void simplifyCFG(Function &F) {
31   llvm::legacy::FunctionPassManager FPM(F.getParent());
32   FPM.add(createCFGSimplificationPass());
33 
34   FPM.doInitialization();
35   FPM.run(F);
36   FPM.doFinalization();
37 }
38 
39 static void lowerSubFn(IRBuilder<> &Builder, CoroSubFnInst *SubFn) {
40   Builder.SetInsertPoint(SubFn);
41   Value *FrameRaw = SubFn->getFrame();
42   int Index = SubFn->getIndex();
43 
44   auto *FrameTy = StructType::get(
45       SubFn->getContext(), {Builder.getInt8PtrTy(), Builder.getInt8PtrTy()});
46   PointerType *FramePtrTy = FrameTy->getPointerTo();
47 
48   Builder.SetInsertPoint(SubFn);
49   auto *FramePtr = Builder.CreateBitCast(FrameRaw, FramePtrTy);
50   auto *Gep = Builder.CreateConstInBoundsGEP2_32(FrameTy, FramePtr, 0, Index);
51   auto *Load = Builder.CreateLoad(FrameTy->getElementType(Index), Gep);
52 
53   SubFn->replaceAllUsesWith(Load);
54 }
55 
56 bool Lowerer::lowerRemainingCoroIntrinsics(Function &F) {
57   bool Changed = false;
58 
59   for (auto IB = inst_begin(F), E = inst_end(F); IB != E;) {
60     Instruction &I = *IB++;
61     if (auto *II = dyn_cast<IntrinsicInst>(&I)) {
62       switch (II->getIntrinsicID()) {
63       default:
64         continue;
65       case Intrinsic::coro_begin:
66         II->replaceAllUsesWith(II->getArgOperand(1));
67         break;
68       case Intrinsic::coro_free:
69         II->replaceAllUsesWith(II->getArgOperand(1));
70         break;
71       case Intrinsic::coro_alloc:
72         II->replaceAllUsesWith(ConstantInt::getTrue(Context));
73         break;
74       case Intrinsic::coro_id:
75       case Intrinsic::coro_id_retcon:
76       case Intrinsic::coro_id_retcon_once:
77       case Intrinsic::coro_id_async:
78         II->replaceAllUsesWith(ConstantTokenNone::get(Context));
79         break;
80       case Intrinsic::coro_subfn_addr:
81         lowerSubFn(Builder, cast<CoroSubFnInst>(II));
82         break;
83       }
84       II->eraseFromParent();
85       Changed = true;
86     }
87   }
88 
89   if (Changed) {
90     // After replacement were made we can cleanup the function body a little.
91     simplifyCFG(F);
92   }
93 
94   return Changed;
95 }
96 
97 static bool declaresCoroCleanupIntrinsics(const Module &M) {
98   return coro::declaresIntrinsics(M, {"llvm.coro.alloc", "llvm.coro.begin",
99                                       "llvm.coro.subfn.addr", "llvm.coro.free",
100                                       "llvm.coro.id", "llvm.coro.id.retcon",
101                                       "llvm.coro.id.retcon.once"});
102 }
103 
104 PreservedAnalyses CoroCleanupPass::run(Function &F,
105                                        FunctionAnalysisManager &AM) {
106   auto &M = *F.getParent();
107   if (!declaresCoroCleanupIntrinsics(M) ||
108       !Lowerer(M).lowerRemainingCoroIntrinsics(F))
109     return PreservedAnalyses::all();
110 
111   return PreservedAnalyses::none();
112 }
113 
114 namespace {
115 
116 struct CoroCleanupLegacy : FunctionPass {
117   static char ID; // Pass identification, replacement for typeid
118 
119   CoroCleanupLegacy() : FunctionPass(ID) {
120     initializeCoroCleanupLegacyPass(*PassRegistry::getPassRegistry());
121   }
122 
123   std::unique_ptr<Lowerer> L;
124 
125   // This pass has work to do only if we find intrinsics we are going to lower
126   // in the module.
127   bool doInitialization(Module &M) override {
128     if (declaresCoroCleanupIntrinsics(M))
129       L = std::make_unique<Lowerer>(M);
130     return false;
131   }
132 
133   bool runOnFunction(Function &F) override {
134     if (L)
135       return L->lowerRemainingCoroIntrinsics(F);
136     return false;
137   }
138   void getAnalysisUsage(AnalysisUsage &AU) const override {
139     if (!L)
140       AU.setPreservesAll();
141   }
142   StringRef getPassName() const override { return "Coroutine Cleanup"; }
143 };
144 }
145 
146 char CoroCleanupLegacy::ID = 0;
147 INITIALIZE_PASS(CoroCleanupLegacy, "coro-cleanup",
148                 "Lower all coroutine related intrinsics", false, false)
149 
150 Pass *llvm::createCoroCleanupLegacyPass() { return new CoroCleanupLegacy(); }
151