Home | History | Annotate | Line # | Download | only in Coroutines 
      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       case Intrinsic::coro_async_size_replace:
     84         auto *Target = cast<ConstantStruct>(
     85             cast<GlobalVariable>(II->getArgOperand(0)->stripPointerCasts())
     86                 ->getInitializer());
     87         auto *Source = cast<ConstantStruct>(
     88             cast<GlobalVariable>(II->getArgOperand(1)->stripPointerCasts())
     89                 ->getInitializer());
     90         auto *TargetSize = Target->getOperand(1);
     91         auto *SourceSize = Source->getOperand(1);
     92         if (TargetSize->isElementWiseEqual(SourceSize)) {
     93           break;
     94         }
     95         auto *TargetRelativeFunOffset = Target->getOperand(0);
     96         auto *NewFuncPtrStruct = ConstantStruct::get(
     97             Target->getType(), TargetRelativeFunOffset, SourceSize);
     98         Target->replaceAllUsesWith(NewFuncPtrStruct);
     99         break;
    100       }
    101       II->eraseFromParent();
    102       Changed = true;
    103     }
    104   }
    105 
    106   if (Changed) {
    107     // After replacement were made we can cleanup the function body a little.
    108     simplifyCFG(F);
    109   }
    110 
    111   return Changed;
    112 }
    113 
    114 static bool declaresCoroCleanupIntrinsics(const Module &M) {
    115   return coro::declaresIntrinsics(
    116       M, {"llvm.coro.alloc", "llvm.coro.begin", "llvm.coro.subfn.addr",
    117           "llvm.coro.free", "llvm.coro.id", "llvm.coro.id.retcon",
    118           "llvm.coro.id.retcon.once", "llvm.coro.async.size.replace"});
    119 }
    120 
    121 PreservedAnalyses CoroCleanupPass::run(Function &F,
    122                                        FunctionAnalysisManager &AM) {
    123   auto &M = *F.getParent();
    124   if (!declaresCoroCleanupIntrinsics(M) ||
    125       !Lowerer(M).lowerRemainingCoroIntrinsics(F))
    126     return PreservedAnalyses::all();
    127 
    128   return PreservedAnalyses::none();
    129 }
    130 
    131 namespace {
    132 
    133 struct CoroCleanupLegacy : FunctionPass {
    134   static char ID; // Pass identification, replacement for typeid
    135 
    136   CoroCleanupLegacy() : FunctionPass(ID) {
    137     initializeCoroCleanupLegacyPass(*PassRegistry::getPassRegistry());
    138   }
    139 
    140   std::unique_ptr<Lowerer> L;
    141 
    142   // This pass has work to do only if we find intrinsics we are going to lower
    143   // in the module.
    144   bool doInitialization(Module &M) override {
    145     if (declaresCoroCleanupIntrinsics(M))
    146       L = std::make_unique<Lowerer>(M);
    147     return false;
    148   }
    149 
    150   bool runOnFunction(Function &F) override {
    151     if (L)
    152       return L->lowerRemainingCoroIntrinsics(F);
    153     return false;
    154   }
    155   void getAnalysisUsage(AnalysisUsage &AU) const override {
    156     if (!L)
    157       AU.setPreservesAll();
    158   }
    159   StringRef getPassName() const override { return "Coroutine Cleanup"; }
    160 };
    161 }
    162 
    163 char CoroCleanupLegacy::ID = 0;
    164 INITIALIZE_PASS(CoroCleanupLegacy, "coro-cleanup",
    165                 "Lower all coroutine related intrinsics", false, false)
    166 
    167 Pass *llvm::createCoroCleanupLegacyPass() { return new CoroCleanupLegacy(); }
    168