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/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 *FramePtr = SubFn->getFrame();
  33   int Index = SubFn->getIndex();
  34
  35   auto *FrameTy = StructType::get(SubFn->getContext(),
  36                                   {Builder.getPtrTy(), Builder.getPtrTy()});
  37
  38   Builder.SetInsertPoint(SubFn);
  39   auto *Gep = Builder.CreateConstInBoundsGEP2_32(FrameTy, FramePtr, 0, Index);
  40   auto *Load = Builder.CreateLoad(FrameTy->getElementType(Index), Gep);
  41
  42   SubFn->replaceAllUsesWith(Load);
  43 }
  44
  45 bool Lowerer::lower(Function &F) {
  46   bool IsPrivateAndUnprocessed = F.isPresplitCoroutine() && F.hasLocalLinkage();
  47   bool Changed = false;
  48
  49   for (Instruction &I : llvm::make_early_inc_range(instructions(F))) {
  50     if (auto *II = dyn_cast<IntrinsicInst>(&I)) {
  51       switch (II->getIntrinsicID()) {
  52       default:
  53         continue;
  54       case Intrinsic::coro_begin:
  55         II->replaceAllUsesWith(II->getArgOperand(1));
  56         break;
  57       case Intrinsic::coro_free:
  58         II->replaceAllUsesWith(II->getArgOperand(1));
  59         break;
  60       case Intrinsic::coro_alloc:
  61         II->replaceAllUsesWith(ConstantInt::getTrue(Context));
  62         break;
  63       case Intrinsic::coro_async_resume:
  64         II->replaceAllUsesWith(
  65             ConstantPointerNull::get(cast<PointerType>(I.getType())));
  66         break;
  67       case Intrinsic::coro_id:
  68       case Intrinsic::coro_id_retcon:
  69       case Intrinsic::coro_id_retcon_once:
  70       case Intrinsic::coro_id_async:
  71         II->replaceAllUsesWith(ConstantTokenNone::get(Context));
  72         break;
  73       case Intrinsic::coro_subfn_addr:
  74         lowerSubFn(Builder, cast<CoroSubFnInst>(II));
  75         break;
  76       case Intrinsic::coro_end:
  77       case Intrinsic::coro_suspend_retcon:
  78         if (IsPrivateAndUnprocessed) {
  79           II->replaceAllUsesWith(UndefValue::get(II->getType()));
  80         } else
  81           continue;
  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   return Changed;
 107 }
 108
 109 static bool declaresCoroCleanupIntrinsics(const Module &M) {
 110   return coro::declaresIntrinsics(
 111       M, {"llvm.coro.alloc", "llvm.coro.begin", "llvm.coro.subfn.addr",
 112           "llvm.coro.free", "llvm.coro.id", "llvm.coro.id.retcon",
 113           "llvm.coro.id.async", "llvm.coro.id.retcon.once",
 114           "llvm.coro.async.size.replace", "llvm.coro.async.resume"});
 115 }
 116
 117 PreservedAnalyses CoroCleanupPass::run(Module &M,
 118                                        ModuleAnalysisManager &MAM) {
 119   if (!declaresCoroCleanupIntrinsics(M))
 120     return PreservedAnalyses::all();
 121
 122   FunctionAnalysisManager &FAM =
 123       MAM.getResult<FunctionAnalysisManagerModuleProxy>(M).getManager();
 124
 125   FunctionPassManager FPM;
 126   FPM.addPass(SimplifyCFGPass());
 127
 128   PreservedAnalyses FuncPA;
 129   FuncPA.preserveSet<CFGAnalyses>();
 130
 131   Lowerer L(M);
 132   for (auto &F : M) {
 133     if (L.lower(F)) {
 134       FAM.invalidate(F, FuncPA);
 135       FPM.run(F, FAM);
 136     }
 137   }
 138
 139   return PreservedAnalyses::none();
 140 }