llvm/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
   9 #include "llvm/Transforms/Coroutines/CoroCleanup.h"
  10 #include "CoroInternal.h"
  11 #include "llvm/IR/Function.h"
  12 #include "llvm/IR/IRBuilder.h"
  13 #include "llvm/IR/InstIterator.h"
  14 #include "llvm/IR/Module.h"
  15 #include "llvm/IR/PassManager.h"
  16 #include "llvm/Transforms/Scalar/SimplifyCFG.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 lower(Function &F);
  28 };
  29 }
  30
  31 static void lowerSubFn(IRBuilder<> &Builder, CoroSubFnInst *SubFn) {
  32   Builder.SetInsertPoint(SubFn);
  33   Value *FramePtr = SubFn->getFrame();
  34   int Index = SubFn->getIndex();
  35
  36   auto *FrameTy = StructType::get(SubFn->getContext(),
  37                                   {Builder.getPtrTy(), Builder.getPtrTy()});
  38
  39   Builder.SetInsertPoint(SubFn);
  40   auto *Gep = Builder.CreateConstInBoundsGEP2_32(FrameTy, FramePtr, 0, Index);
  41   auto *Load = Builder.CreateLoad(FrameTy->getElementType(Index), Gep);
  42
  43   SubFn->replaceAllUsesWith(Load);
  44 }
  45
  46 bool Lowerer::lower(Function &F) {
  47   bool IsPrivateAndUnprocessed = F.isPresplitCoroutine() && F.hasLocalLinkage();
  48   bool Changed = false;
  49
  50   for (Instruction &I : llvm::make_early_inc_range(instructions(F))) {
  51     if (auto *II = dyn_cast<IntrinsicInst>(&I)) {
  52       switch (II->getIntrinsicID()) {
  53       default:
  54         continue;
  55       case Intrinsic::coro_begin:
  56       case Intrinsic::coro_begin_custom_abi:
  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(PoisonValue::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           "llvm.coro.begin.custom.abi"});
 118 }
 119
 120 PreservedAnalyses CoroCleanupPass::run(Module &M,
 121                                        ModuleAnalysisManager &MAM) {
 122   if (!declaresCoroCleanupIntrinsics(M))
 123     return PreservedAnalyses::all();
 124
 125   FunctionAnalysisManager &FAM =
 126       MAM.getResult<FunctionAnalysisManagerModuleProxy>(M).getManager();
 127
 128   FunctionPassManager FPM;
 129   FPM.addPass(SimplifyCFGPass());
 130
 131   PreservedAnalyses FuncPA;
 132   FuncPA.preserveSet<CFGAnalyses>();
 133
 134   Lowerer L(M);
 135   for (auto &F : M) {
 136     if (L.lower(F)) {
 137       FAM.invalidate(F, FuncPA);
 138       FPM.run(F, FAM);
 139     }
 140   }
 141
 142   return PreservedAnalyses::none();
 143 }