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 = F.isPresplitCoroutine() && F.hasLocalLinkage();
49   bool Changed = false;
50 
51   for (Instruction &I : llvm::make_early_inc_range(instructions(F))) {
52     if (auto *II = dyn_cast<IntrinsicInst>(&I)) {
53       switch (II->getIntrinsicID()) {
54       default:
55         continue;
56       case Intrinsic::coro_begin:
57         II->replaceAllUsesWith(II->getArgOperand(1));
58         break;
59       case Intrinsic::coro_free:
60         II->replaceAllUsesWith(II->getArgOperand(1));
61         break;
62       case Intrinsic::coro_alloc:
63         II->replaceAllUsesWith(ConstantInt::getTrue(Context));
64         break;
65       case Intrinsic::coro_async_resume:
66         II->replaceAllUsesWith(
67             ConstantPointerNull::get(cast<PointerType>(I.getType())));
68         break;
69       case Intrinsic::coro_id:
70       case Intrinsic::coro_id_retcon:
71       case Intrinsic::coro_id_retcon_once:
72       case Intrinsic::coro_id_async:
73         II->replaceAllUsesWith(ConstantTokenNone::get(Context));
74         break;
75       case Intrinsic::coro_subfn_addr:
76         lowerSubFn(Builder, cast<CoroSubFnInst>(II));
77         break;
78       case Intrinsic::coro_end:
79       case Intrinsic::coro_suspend_retcon:
80         if (IsPrivateAndUnprocessed) {
81           II->replaceAllUsesWith(UndefValue::get(II->getType()));
82         } else
83           continue;
84         break;
85       case Intrinsic::coro_async_size_replace:
86         auto *Target = cast<ConstantStruct>(
87             cast<GlobalVariable>(II->getArgOperand(0)->stripPointerCasts())
88                 ->getInitializer());
89         auto *Source = cast<ConstantStruct>(
90             cast<GlobalVariable>(II->getArgOperand(1)->stripPointerCasts())
91                 ->getInitializer());
92         auto *TargetSize = Target->getOperand(1);
93         auto *SourceSize = Source->getOperand(1);
94         if (TargetSize->isElementWiseEqual(SourceSize)) {
95           break;
96         }
97         auto *TargetRelativeFunOffset = Target->getOperand(0);
98         auto *NewFuncPtrStruct = ConstantStruct::get(
99             Target->getType(), TargetRelativeFunOffset, SourceSize);
100         Target->replaceAllUsesWith(NewFuncPtrStruct);
101         break;
102       }
103       II->eraseFromParent();
104       Changed = true;
105     }
106   }
107 
108   return Changed;
109 }
110 
111 static bool declaresCoroCleanupIntrinsics(const Module &M) {
112   return coro::declaresIntrinsics(
113       M, {"llvm.coro.alloc", "llvm.coro.begin", "llvm.coro.subfn.addr",
114           "llvm.coro.free", "llvm.coro.id", "llvm.coro.id.retcon",
115           "llvm.coro.id.async", "llvm.coro.id.retcon.once",
116           "llvm.coro.async.size.replace", "llvm.coro.async.resume"});
117 }
118 
119 PreservedAnalyses CoroCleanupPass::run(Module &M,
120                                        ModuleAnalysisManager &MAM) {
121   if (!declaresCoroCleanupIntrinsics(M))
122     return PreservedAnalyses::all();
123 
124   FunctionAnalysisManager &FAM =
125       MAM.getResult<FunctionAnalysisManagerModuleProxy>(M).getManager();
126 
127   FunctionPassManager FPM;
128   FPM.addPass(SimplifyCFGPass());
129 
130   Lowerer L(M);
131   for (auto &F : M)
132     if (L.lower(F))
133       FPM.run(F, FAM);
134 
135   return PreservedAnalyses::none();
136 }
137