lib/Transforms/Coroutines/CoroCleanup.cpp

   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 // This pass lowers all remaining coroutine intrinsics.
   9 //===----------------------------------------------------------------------===//
  10
  11 #include "CoroInternal.h"
  12 #include "llvm/IR/IRBuilder.h"
  13 #include "llvm/IR/InstIterator.h"
  14 #include "llvm/IR/LegacyPassManager.h"
  15 #include "llvm/Pass.h"
  16 #include "llvm/Transforms/Scalar.h"
  17
  18 using namespace llvm;
  19
  20 #define DEBUG_TYPE "coro-cleanup"
  21
  22 namespace {
  23 // Created on demand if CoroCleanup pass has work to do.
  24 struct Lowerer : coro::LowererBase {
  25   IRBuilder<> Builder;
  26   Lowerer(Module &M) : LowererBase(M), Builder(Context) {}
  27   bool lowerRemainingCoroIntrinsics(Function &F);
  28 };
  29 }
  30
  31 static void simplifyCFG(Function &F) {
  32   llvm::legacy::FunctionPassManager FPM(F.getParent());
  33   FPM.add(createCFGSimplificationPass());
  34
  35   FPM.doInitialization();
  36   FPM.run(F);
  37   FPM.doFinalization();
  38 }
  39
  40 static void lowerSubFn(IRBuilder<> &Builder, CoroSubFnInst *SubFn) {
  41   Builder.SetInsertPoint(SubFn);
  42   Value *FrameRaw = SubFn->getFrame();
  43   int Index = SubFn->getIndex();
  44
  45   auto *FrameTy = StructType::get(
  46       SubFn->getContext(), {Builder.getInt8PtrTy(), Builder.getInt8PtrTy()});
  47   PointerType *FramePtrTy = FrameTy->getPointerTo();
  48
  49   Builder.SetInsertPoint(SubFn);
  50   auto *FramePtr = Builder.CreateBitCast(FrameRaw, FramePtrTy);
  51   auto *Gep = Builder.CreateConstInBoundsGEP2_32(FrameTy, FramePtr, 0, Index);
  52   auto *Load = Builder.CreateLoad(FrameTy->getElementType(Index), Gep);
  53
  54   SubFn->replaceAllUsesWith(Load);
  55 }
  56
  57 bool Lowerer::lowerRemainingCoroIntrinsics(Function &F) {
  58   bool Changed = false;
  59
  60   for (auto IB = inst_begin(F), E = inst_end(F); IB != E;) {
  61     Instruction &I = *IB++;
  62     if (auto *II = dyn_cast<IntrinsicInst>(&I)) {
  63       switch (II->getIntrinsicID()) {
  64       default:
  65         continue;
  66       case Intrinsic::coro_begin:
  67         II->replaceAllUsesWith(II->getArgOperand(1));
  68         break;
  69       case Intrinsic::coro_free:
  70         II->replaceAllUsesWith(II->getArgOperand(1));
  71         break;
  72       case Intrinsic::coro_alloc:
  73         II->replaceAllUsesWith(ConstantInt::getTrue(Context));
  74         break;
  75       case Intrinsic::coro_id:
  76       case Intrinsic::coro_id_retcon:
  77       case Intrinsic::coro_id_retcon_once:
  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   return Changed;
  94 }
  95
  96 //===----------------------------------------------------------------------===//
  97 //                              Top Level Driver
  98 //===----------------------------------------------------------------------===//
  99
 100 namespace {
 101
 102 struct CoroCleanup : FunctionPass {
 103   static char ID; // Pass identification, replacement for typeid
 104
 105   CoroCleanup() : FunctionPass(ID) {
 106     initializeCoroCleanupPass(*PassRegistry::getPassRegistry());
 107   }
 108
 109   std::unique_ptr<Lowerer> L;
 110
 111   // This pass has work to do only if we find intrinsics we are going to lower
 112   // in the module.
 113   bool doInitialization(Module &M) override {
 114     if (coro::declaresIntrinsics(M, {"llvm.coro.alloc", "llvm.coro.begin",
 115                                      "llvm.coro.subfn.addr", "llvm.coro.free",
 116                                      "llvm.coro.id", "llvm.coro.id.retcon",
 117                                      "llvm.coro.id.retcon.once"}))
 118       L = std::make_unique<Lowerer>(M);
 119     return false;
 120   }
 121
 122   bool runOnFunction(Function &F) override {
 123     if (L)
 124       return L->lowerRemainingCoroIntrinsics(F);
 125     return false;
 126   }
 127   void getAnalysisUsage(AnalysisUsage &AU) const override {
 128     if (!L)
 129       AU.setPreservesAll();
 130   }
 131   StringRef getPassName() const override { return "Coroutine Cleanup"; }
 132 };
 133 }
 134
 135 char CoroCleanup::ID = 0;
 136 INITIALIZE_PASS(CoroCleanup, "coro-cleanup",
 137                 "Lower all coroutine related intrinsics", false, false)
 138
 139 Pass *llvm::createCoroCleanupPass() { return new CoroCleanup(); }