contrib/llvm/lib/Transforms/Coroutines/Coroutines.cpp

   1 //===- Coroutines.cpp -----------------------------------------------------===//
   2 //
   3 //                     The LLVM Compiler Infrastructure
   4 //
   5 // This file is distributed under the University of Illinois Open Source
   6 // License. See LICENSE.TXT for details.
   7 //
   8 //===----------------------------------------------------------------------===//
   9 //
  10 // This file implements the common infrastructure for Coroutine Passes.
  11 //
  12 //===----------------------------------------------------------------------===//
  13
  14 #include "llvm/Transforms/Coroutines.h"
  15 #include "llvm-c/Transforms/Coroutines.h"
  16 #include "CoroInstr.h"
  17 #include "CoroInternal.h"
  18 #include "llvm/ADT/SmallVector.h"
  19 #include "llvm/ADT/StringRef.h"
  20 #include "llvm/Analysis/CallGraph.h"
  21 #include "llvm/Analysis/CallGraphSCCPass.h"
  22 #include "llvm/Transforms/Utils/Local.h"
  23 #include "llvm/IR/Attributes.h"
  24 #include "llvm/IR/CallSite.h"
  25 #include "llvm/IR/Constants.h"
  26 #include "llvm/IR/DerivedTypes.h"
  27 #include "llvm/IR/Function.h"
  28 #include "llvm/IR/InstIterator.h"
  29 #include "llvm/IR/Instructions.h"
  30 #include "llvm/IR/IntrinsicInst.h"
  31 #include "llvm/IR/Intrinsics.h"
  32 #include "llvm/IR/LegacyPassManager.h"
  33 #include "llvm/IR/Module.h"
  34 #include "llvm/IR/Type.h"
  35 #include "llvm/Support/Casting.h"
  36 #include "llvm/Support/ErrorHandling.h"
  37 #include "llvm/Transforms/IPO.h"
  38 #include "llvm/Transforms/IPO/PassManagerBuilder.h"
  39 #include <cassert>
  40 #include <cstddef>
  41 #include <utility>
  42
  43 using namespace llvm;
  44
  45 void llvm::initializeCoroutines(PassRegistry &Registry) {
  46   initializeCoroEarlyPass(Registry);
  47   initializeCoroSplitPass(Registry);
  48   initializeCoroElidePass(Registry);
  49   initializeCoroCleanupPass(Registry);
  50 }
  51
  52 static void addCoroutineOpt0Passes(const PassManagerBuilder &Builder,
  53                                    legacy::PassManagerBase &PM) {
  54   PM.add(createCoroSplitPass());
  55   PM.add(createCoroElidePass());
  56
  57   PM.add(createBarrierNoopPass());
  58   PM.add(createCoroCleanupPass());
  59 }
  60
  61 static void addCoroutineEarlyPasses(const PassManagerBuilder &Builder,
  62                                     legacy::PassManagerBase &PM) {
  63   PM.add(createCoroEarlyPass());
  64 }
  65
  66 static void addCoroutineScalarOptimizerPasses(const PassManagerBuilder &Builder,
  67                                               legacy::PassManagerBase &PM) {
  68   PM.add(createCoroElidePass());
  69 }
  70
  71 static void addCoroutineSCCPasses(const PassManagerBuilder &Builder,
  72                                   legacy::PassManagerBase &PM) {
  73   PM.add(createCoroSplitPass());
  74 }
  75
  76 static void addCoroutineOptimizerLastPasses(const PassManagerBuilder &Builder,
  77                                             legacy::PassManagerBase &PM) {
  78   PM.add(createCoroCleanupPass());
  79 }
  80
  81 void llvm::addCoroutinePassesToExtensionPoints(PassManagerBuilder &Builder) {
  82   Builder.addExtension(PassManagerBuilder::EP_EarlyAsPossible,
  83                        addCoroutineEarlyPasses);
  84   Builder.addExtension(PassManagerBuilder::EP_EnabledOnOptLevel0,
  85                        addCoroutineOpt0Passes);
  86   Builder.addExtension(PassManagerBuilder::EP_CGSCCOptimizerLate,
  87                        addCoroutineSCCPasses);
  88   Builder.addExtension(PassManagerBuilder::EP_ScalarOptimizerLate,
  89                        addCoroutineScalarOptimizerPasses);
  90   Builder.addExtension(PassManagerBuilder::EP_OptimizerLast,
  91                        addCoroutineOptimizerLastPasses);
  92 }
  93
  94 // Construct the lowerer base class and initialize its members.
  95 coro::LowererBase::LowererBase(Module &M)
  96     : TheModule(M), Context(M.getContext()),
  97       Int8Ptr(Type::getInt8PtrTy(Context)),
  98       ResumeFnType(FunctionType::get(Type::getVoidTy(Context), Int8Ptr,
  99                                      /*isVarArg=*/false)),
 100       NullPtr(ConstantPointerNull::get(Int8Ptr)) {}
 101
 102 // Creates a sequence of instructions to obtain a resume function address using
 103 // llvm.coro.subfn.addr. It generates the following sequence:
 104 //
 105 //    call i8* @llvm.coro.subfn.addr(i8* %Arg, i8 %index)
 106 //    bitcast i8* %2 to void(i8*)*
 107
 108 Value *coro::LowererBase::makeSubFnCall(Value *Arg, int Index,
 109                                         Instruction *InsertPt) {
 110   auto *IndexVal = ConstantInt::get(Type::getInt8Ty(Context), Index);
 111   auto *Fn = Intrinsic::getDeclaration(&TheModule, Intrinsic::coro_subfn_addr);
 112
 113   assert(Index >= CoroSubFnInst::IndexFirst &&
 114          Index < CoroSubFnInst::IndexLast &&
 115          "makeSubFnCall: Index value out of range");
 116   auto *Call = CallInst::Create(Fn, {Arg, IndexVal}, "", InsertPt);
 117
 118   auto *Bitcast =
 119       new BitCastInst(Call, ResumeFnType->getPointerTo(), "", InsertPt);
 120   return Bitcast;
 121 }
 122
 123 #ifndef NDEBUG
 124 static bool isCoroutineIntrinsicName(StringRef Name) {
 125   // NOTE: Must be sorted!
 126   static const char *const CoroIntrinsics[] = {
 127       "llvm.coro.alloc",   "llvm.coro.begin",   "llvm.coro.destroy",
 128       "llvm.coro.done",    "llvm.coro.end",     "llvm.coro.frame",
 129       "llvm.coro.free",    "llvm.coro.id",      "llvm.coro.noop",
 130       "llvm.coro.param",   "llvm.coro.promise", "llvm.coro.resume",
 131       "llvm.coro.save",    "llvm.coro.size",    "llvm.coro.subfn.addr",
 132       "llvm.coro.suspend",
 133   };
 134   return Intrinsic::lookupLLVMIntrinsicByName(CoroIntrinsics, Name) != -1;
 135 }
 136 #endif
 137
 138 // Verifies if a module has named values listed. Also, in debug mode verifies
 139 // that names are intrinsic names.
 140 bool coro::declaresIntrinsics(Module &M,
 141                               std::initializer_list<StringRef> List) {
 142   for (StringRef Name : List) {
 143     assert(isCoroutineIntrinsicName(Name) && "not a coroutine intrinsic");
 144     if (M.getNamedValue(Name))
 145       return true;
 146   }
 147
 148   return false;
 149 }
 150
 151 // Replace all coro.frees associated with the provided CoroId either with 'null'
 152 // if Elide is true and with its frame parameter otherwise.
 153 void coro::replaceCoroFree(CoroIdInst *CoroId, bool Elide) {
 154   SmallVector<CoroFreeInst *, 4> CoroFrees;
 155   for (User *U : CoroId->users())
 156     if (auto CF = dyn_cast<CoroFreeInst>(U))
 157       CoroFrees.push_back(CF);
 158
 159   if (CoroFrees.empty())
 160     return;
 161
 162   Value *Replacement =
 163       Elide ? ConstantPointerNull::get(Type::getInt8PtrTy(CoroId->getContext()))
 164             : CoroFrees.front()->getFrame();
 165
 166   for (CoroFreeInst *CF : CoroFrees) {
 167     CF->replaceAllUsesWith(Replacement);
 168     CF->eraseFromParent();
 169   }
 170 }
 171
 172 // FIXME: This code is stolen from CallGraph::addToCallGraph(Function *F), which
 173 // happens to be private. It is better for this functionality exposed by the
 174 // CallGraph.
 175 static void buildCGN(CallGraph &CG, CallGraphNode *Node) {
 176   Function *F = Node->getFunction();
 177
 178   // Look for calls by this function.
 179   for (Instruction &I : instructions(F))
 180     if (CallSite CS = CallSite(cast<Value>(&I))) {
 181       const Function *Callee = CS.getCalledFunction();
 182       if (!Callee || !Intrinsic::isLeaf(Callee->getIntrinsicID()))
 183         // Indirect calls of intrinsics are not allowed so no need to check.
 184         // We can be more precise here by using TargetArg returned by
 185         // Intrinsic::isLeaf.
 186         Node->addCalledFunction(CS, CG.getCallsExternalNode());
 187       else if (!Callee->isIntrinsic())
 188         Node->addCalledFunction(CS, CG.getOrInsertFunction(Callee));
 189     }
 190 }
 191
 192 // Rebuild CGN after we extracted parts of the code from ParentFunc into
 193 // NewFuncs. Builds CGNs for the NewFuncs and adds them to the current SCC.
 194 void coro::updateCallGraph(Function &ParentFunc, ArrayRef<Function *> NewFuncs,
 195                            CallGraph &CG, CallGraphSCC &SCC) {
 196   // Rebuild CGN from scratch for the ParentFunc
 197   auto *ParentNode = CG[&ParentFunc];
 198   ParentNode->removeAllCalledFunctions();
 199   buildCGN(CG, ParentNode);
 200
 201   SmallVector<CallGraphNode *, 8> Nodes(SCC.begin(), SCC.end());
 202
 203   for (Function *F : NewFuncs) {
 204     CallGraphNode *Callee = CG.getOrInsertFunction(F);
 205     Nodes.push_back(Callee);
 206     buildCGN(CG, Callee);
 207   }
 208
 209   SCC.initialize(Nodes);
 210 }
 211
 212 static void clear(coro::Shape &Shape) {
 213   Shape.CoroBegin = nullptr;
 214   Shape.CoroEnds.clear();
 215   Shape.CoroSizes.clear();
 216   Shape.CoroSuspends.clear();
 217
 218   Shape.FrameTy = nullptr;
 219   Shape.FramePtr = nullptr;
 220   Shape.AllocaSpillBlock = nullptr;
 221   Shape.ResumeSwitch = nullptr;
 222   Shape.PromiseAlloca = nullptr;
 223   Shape.HasFinalSuspend = false;
 224 }
 225
 226 static CoroSaveInst *createCoroSave(CoroBeginInst *CoroBegin,
 227                                     CoroSuspendInst *SuspendInst) {
 228   Module *M = SuspendInst->getModule();
 229   auto *Fn = Intrinsic::getDeclaration(M, Intrinsic::coro_save);
 230   auto *SaveInst =
 231       cast<CoroSaveInst>(CallInst::Create(Fn, CoroBegin, "", SuspendInst));
 232   assert(!SuspendInst->getCoroSave());
 233   SuspendInst->setArgOperand(0, SaveInst);
 234   return SaveInst;
 235 }
 236
 237 // Collect "interesting" coroutine intrinsics.
 238 void coro::Shape::buildFrom(Function &F) {
 239   size_t FinalSuspendIndex = 0;
 240   clear(*this);
 241   SmallVector<CoroFrameInst *, 8> CoroFrames;
 242   SmallVector<CoroSaveInst *, 2> UnusedCoroSaves;
 243
 244   for (Instruction &I : instructions(F)) {
 245     if (auto II = dyn_cast<IntrinsicInst>(&I)) {
 246       switch (II->getIntrinsicID()) {
 247       default:
 248         continue;
 249       case Intrinsic::coro_size:
 250         CoroSizes.push_back(cast<CoroSizeInst>(II));
 251         break;
 252       case Intrinsic::coro_frame:
 253         CoroFrames.push_back(cast<CoroFrameInst>(II));
 254         break;
 255       case Intrinsic::coro_save:
 256         // After optimizations, coro_suspends using this coro_save might have
 257         // been removed, remember orphaned coro_saves to remove them later.
 258         if (II->use_empty())
 259           UnusedCoroSaves.push_back(cast<CoroSaveInst>(II));
 260         break;
 261       case Intrinsic::coro_suspend:
 262         CoroSuspends.push_back(cast<CoroSuspendInst>(II));
 263         if (CoroSuspends.back()->isFinal()) {
 264           if (HasFinalSuspend)
 265             report_fatal_error(
 266               "Only one suspend point can be marked as final");
 267           HasFinalSuspend = true;
 268           FinalSuspendIndex = CoroSuspends.size() - 1;
 269         }
 270         break;
 271       case Intrinsic::coro_begin: {
 272         auto CB = cast<CoroBeginInst>(II);
 273         if (CB->getId()->getInfo().isPreSplit()) {
 274           if (CoroBegin)
 275             report_fatal_error(
 276                 "coroutine should have exactly one defining @llvm.coro.begin");
 277           CB->addAttribute(AttributeList::ReturnIndex, Attribute::NonNull);
 278           CB->addAttribute(AttributeList::ReturnIndex, Attribute::NoAlias);
 279           CB->removeAttribute(AttributeList::FunctionIndex,
 280                               Attribute::NoDuplicate);
 281           CoroBegin = CB;
 282         }
 283         break;
 284       }
 285       case Intrinsic::coro_end:
 286         CoroEnds.push_back(cast<CoroEndInst>(II));
 287         if (CoroEnds.back()->isFallthrough()) {
 288           // Make sure that the fallthrough coro.end is the first element in the
 289           // CoroEnds vector.
 290           if (CoroEnds.size() > 1) {
 291             if (CoroEnds.front()->isFallthrough())
 292               report_fatal_error(
 293                   "Only one coro.end can be marked as fallthrough");
 294             std::swap(CoroEnds.front(), CoroEnds.back());
 295           }
 296         }
 297         break;
 298       }
 299     }
 300   }
 301
 302   // If for some reason, we were not able to find coro.begin, bailout.
 303   if (!CoroBegin) {
 304     // Replace coro.frame which are supposed to be lowered to the result of
 305     // coro.begin with undef.
 306     auto *Undef = UndefValue::get(Type::getInt8PtrTy(F.getContext()));
 307     for (CoroFrameInst *CF : CoroFrames) {
 308       CF->replaceAllUsesWith(Undef);
 309       CF->eraseFromParent();
 310     }
 311
 312     // Replace all coro.suspend with undef and remove related coro.saves if
 313     // present.
 314     for (CoroSuspendInst *CS : CoroSuspends) {
 315       CS->replaceAllUsesWith(UndefValue::get(CS->getType()));
 316       CS->eraseFromParent();
 317       if (auto *CoroSave = CS->getCoroSave())
 318         CoroSave->eraseFromParent();
 319     }
 320
 321     // Replace all coro.ends with unreachable instruction.
 322     for (CoroEndInst *CE : CoroEnds)
 323       changeToUnreachable(CE, /*UseLLVMTrap=*/false);
 324
 325     return;
 326   }
 327
 328   // The coro.free intrinsic is always lowered to the result of coro.begin.
 329   for (CoroFrameInst *CF : CoroFrames) {
 330     CF->replaceAllUsesWith(CoroBegin);
 331     CF->eraseFromParent();
 332   }
 333
 334   // Canonicalize coro.suspend by inserting a coro.save if needed.
 335   for (CoroSuspendInst *CS : CoroSuspends)
 336     if (!CS->getCoroSave())
 337       createCoroSave(CoroBegin, CS);
 338
 339   // Move final suspend to be the last element in the CoroSuspends vector.
 340   if (HasFinalSuspend &&
 341       FinalSuspendIndex != CoroSuspends.size() - 1)
 342     std::swap(CoroSuspends[FinalSuspendIndex], CoroSuspends.back());
 343
 344   // Remove orphaned coro.saves.
 345   for (CoroSaveInst *CoroSave : UnusedCoroSaves)
 346     CoroSave->eraseFromParent();
 347 }
 348
 349 void LLVMAddCoroEarlyPass(LLVMPassManagerRef PM) {
 350   unwrap(PM)->add(createCoroEarlyPass());
 351 }
 352
 353 void LLVMAddCoroSplitPass(LLVMPassManagerRef PM) {
 354   unwrap(PM)->add(createCoroSplitPass());
 355 }
 356
 357 void LLVMAddCoroElidePass(LLVMPassManagerRef PM) {
 358   unwrap(PM)->add(createCoroElidePass());
 359 }
 360
 361 void LLVMAddCoroCleanupPass(LLVMPassManagerRef PM) {
 362   unwrap(PM)->add(createCoroCleanupPass());
 363 }