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   bool IsPrivateAndUnprocessed =
60       F.hasFnAttribute(CORO_PRESPLIT_ATTR) && F.hasLocalLinkage();
61 
62   for (Instruction &I : llvm::make_early_inc_range(instructions(F))) {
63     if (auto *II = dyn_cast<IntrinsicInst>(&I)) {
64       switch (II->getIntrinsicID()) {
65       default:
66         continue;
67       case Intrinsic::coro_begin:
68         II->replaceAllUsesWith(II->getArgOperand(1));
69         break;
70       case Intrinsic::coro_free:
71         II->replaceAllUsesWith(II->getArgOperand(1));
72         break;
73       case Intrinsic::coro_alloc:
74         II->replaceAllUsesWith(ConstantInt::getTrue(Context));
75         break;
76       case Intrinsic::coro_async_resume:
77         II->replaceAllUsesWith(
78             ConstantPointerNull::get(cast<PointerType>(I.getType())));
79         break;
80       case Intrinsic::coro_id:
81       case Intrinsic::coro_id_retcon:
82       case Intrinsic::coro_id_retcon_once:
83       case Intrinsic::coro_id_async:
84         II->replaceAllUsesWith(ConstantTokenNone::get(Context));
85         break;
86       case Intrinsic::coro_subfn_addr:
87         lowerSubFn(Builder, cast<CoroSubFnInst>(II));
88         break;
89       case Intrinsic::coro_end:
90       case Intrinsic::coro_suspend_retcon:
91         if (IsPrivateAndUnprocessed) {
92           II->replaceAllUsesWith(UndefValue::get(II->getType()));
93         } else
94           continue;
95         break;
96       case Intrinsic::coro_async_size_replace:
97         auto *Target = cast<ConstantStruct>(
98             cast<GlobalVariable>(II->getArgOperand(0)->stripPointerCasts())
99                 ->getInitializer());
100         auto *Source = cast<ConstantStruct>(
101             cast<GlobalVariable>(II->getArgOperand(1)->stripPointerCasts())
102                 ->getInitializer());
103         auto *TargetSize = Target->getOperand(1);
104         auto *SourceSize = Source->getOperand(1);
105         if (TargetSize->isElementWiseEqual(SourceSize)) {
106           break;
107         }
108         auto *TargetRelativeFunOffset = Target->getOperand(0);
109         auto *NewFuncPtrStruct = ConstantStruct::get(
110             Target->getType(), TargetRelativeFunOffset, SourceSize);
111         Target->replaceAllUsesWith(NewFuncPtrStruct);
112         break;
113       }
114       II->eraseFromParent();
115       Changed = true;
116     }
117   }
118 
119   if (Changed) {
120     // After replacement were made we can cleanup the function body a little.
121     simplifyCFG(F);
122   }
123 
124   return Changed;
125 }
126 
127 static bool declaresCoroCleanupIntrinsics(const Module &M) {
128   return coro::declaresIntrinsics(
129       M, {"llvm.coro.alloc", "llvm.coro.begin", "llvm.coro.subfn.addr",
130           "llvm.coro.free", "llvm.coro.id", "llvm.coro.id.retcon",
131           "llvm.coro.id.retcon.once", "llvm.coro.async.size.replace",
132           "llvm.coro.async.resume"});
133 }
134 
135 PreservedAnalyses CoroCleanupPass::run(Function &F,
136                                        FunctionAnalysisManager &AM) {
137   auto &M = *F.getParent();
138   if (!declaresCoroCleanupIntrinsics(M) ||
139       !Lowerer(M).lowerRemainingCoroIntrinsics(F))
140     return PreservedAnalyses::all();
141 
142   return PreservedAnalyses::none();
143 }
144