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

   1 //===- CoroSplit.cpp - Converts a coroutine into a state machine ----------===//
   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 // This pass builds the coroutine frame and outlines resume and destroy parts
  10 // of the coroutine into separate functions.
  11 //
  12 // We present a coroutine to an LLVM as an ordinary function with suspension
  13 // points marked up with intrinsics. We let the optimizer party on the coroutine
  14 // as a single function for as long as possible. Shortly before the coroutine is
  15 // eligible to be inlined into its callers, we split up the coroutine into parts
  16 // corresponding to an initial, resume and destroy invocations of the coroutine,
  17 // add them to the current SCC and restart the IPO pipeline to optimize the
  18 // coroutine subfunctions we extracted before proceeding to the caller of the
  19 // coroutine.
  20 //===----------------------------------------------------------------------===//
  21
  22 #include "CoroInternal.h"
  23 #include "llvm/Analysis/CallGraphSCCPass.h"
  24 #include "llvm/IR/DebugInfoMetadata.h"
  25 #include "llvm/IR/IRBuilder.h"
  26 #include "llvm/IR/LegacyPassManager.h"
  27 #include "llvm/IR/Verifier.h"
  28 #include "llvm/Transforms/Scalar.h"
  29 #include "llvm/Transforms/Utils/Cloning.h"
  30 #include "llvm/Transforms/Utils/Local.h"
  31 #include "llvm/Transforms/Utils/ValueMapper.h"
  32
  33 using namespace llvm;
  34
  35 #define DEBUG_TYPE "coro-split"
  36
  37 // Create an entry block for a resume function with a switch that will jump to
  38 // suspend points.
  39 static BasicBlock *createResumeEntryBlock(Function &F, coro::Shape &Shape) {
  40   LLVMContext &C = F.getContext();
  41
  42   // resume.entry:
  43   //  %index.addr = getelementptr inbounds %f.Frame, %f.Frame* %FramePtr, i32 0,
  44   //  i32 2
  45   //  % index = load i32, i32* %index.addr
  46   //  switch i32 %index, label %unreachable [
  47   //    i32 0, label %resume.0
  48   //    i32 1, label %resume.1
  49   //    ...
  50   //  ]
  51
  52   auto *NewEntry = BasicBlock::Create(C, "resume.entry", &F);
  53   auto *UnreachBB = BasicBlock::Create(C, "unreachable", &F);
  54
  55   IRBuilder<> Builder(NewEntry);
  56   auto *FramePtr = Shape.FramePtr;
  57   auto *FrameTy = Shape.FrameTy;
  58   auto *GepIndex = Builder.CreateConstInBoundsGEP2_32(
  59       FrameTy, FramePtr, 0, coro::Shape::IndexField, "index.addr");
  60   auto *Index = Builder.CreateLoad(GepIndex, "index");
  61   auto *Switch =
  62       Builder.CreateSwitch(Index, UnreachBB, Shape.CoroSuspends.size());
  63   Shape.ResumeSwitch = Switch;
  64
  65   size_t SuspendIndex = 0;
  66   for (CoroSuspendInst *S : Shape.CoroSuspends) {
  67     ConstantInt *IndexVal = Shape.getIndex(SuspendIndex);
  68
  69     // Replace CoroSave with a store to Index:
  70     //    %index.addr = getelementptr %f.frame... (index field number)
  71     //    store i32 0, i32* %index.addr1
  72     auto *Save = S->getCoroSave();
  73     Builder.SetInsertPoint(Save);
  74     if (S->isFinal()) {
  75       // Final suspend point is represented by storing zero in ResumeFnAddr.
  76       auto *GepIndex = Builder.CreateConstInBoundsGEP2_32(FrameTy, FramePtr, 0,
  77                                                           0, "ResumeFn.addr");
  78       auto *NullPtr = ConstantPointerNull::get(cast<PointerType>(
  79           cast<PointerType>(GepIndex->getType())->getElementType()));
  80       Builder.CreateStore(NullPtr, GepIndex);
  81     } else {
  82       auto *GepIndex = Builder.CreateConstInBoundsGEP2_32(
  83           FrameTy, FramePtr, 0, coro::Shape::IndexField, "index.addr");
  84       Builder.CreateStore(IndexVal, GepIndex);
  85     }
  86     Save->replaceAllUsesWith(ConstantTokenNone::get(C));
  87     Save->eraseFromParent();
  88
  89     // Split block before and after coro.suspend and add a jump from an entry
  90     // switch:
  91     //
  92     //  whateverBB:
  93     //    whatever
  94     //    %0 = call i8 @llvm.coro.suspend(token none, i1 false)
  95     //    switch i8 %0, label %suspend[i8 0, label %resume
  96     //                                 i8 1, label %cleanup]
  97     // becomes:
  98     //
  99     //  whateverBB:
 100     //     whatever
 101     //     br label %resume.0.landing
 102     //
 103     //  resume.0: ; <--- jump from the switch in the resume.entry
 104     //     %0 = tail call i8 @llvm.coro.suspend(token none, i1 false)
 105     //     br label %resume.0.landing
 106     //
 107     //  resume.0.landing:
 108     //     %1 = phi i8[-1, %whateverBB], [%0, %resume.0]
 109     //     switch i8 % 1, label %suspend [i8 0, label %resume
 110     //                                    i8 1, label %cleanup]
 111
 112     auto *SuspendBB = S->getParent();
 113     auto *ResumeBB =
 114         SuspendBB->splitBasicBlock(S, "resume." + Twine(SuspendIndex));
 115     auto *LandingBB = ResumeBB->splitBasicBlock(
 116         S->getNextNode(), ResumeBB->getName() + Twine(".landing"));
 117     Switch->addCase(IndexVal, ResumeBB);
 118
 119     cast<BranchInst>(SuspendBB->getTerminator())->setSuccessor(0, LandingBB);
 120     auto *PN = PHINode::Create(Builder.getInt8Ty(), 2, "", &LandingBB->front());
 121     S->replaceAllUsesWith(PN);
 122     PN->addIncoming(Builder.getInt8(-1), SuspendBB);
 123     PN->addIncoming(S, ResumeBB);
 124
 125     ++SuspendIndex;
 126   }
 127
 128   Builder.SetInsertPoint(UnreachBB);
 129   Builder.CreateUnreachable();
 130
 131   return NewEntry;
 132 }
 133
 134 // In Resumers, we replace fallthrough coro.end with ret void and delete the
 135 // rest of the block.
 136 static void replaceFallthroughCoroEnd(IntrinsicInst *End,
 137                                       ValueToValueMapTy &VMap) {
 138   auto *NewE = cast<IntrinsicInst>(VMap[End]);
 139   ReturnInst::Create(NewE->getContext(), nullptr, NewE);
 140
 141   // Remove the rest of the block, by splitting it into an unreachable block.
 142   auto *BB = NewE->getParent();
 143   BB->splitBasicBlock(NewE);
 144   BB->getTerminator()->eraseFromParent();
 145 }
 146
 147 // Rewrite final suspend point handling. We do not use suspend index to
 148 // represent the final suspend point. Instead we zero-out ResumeFnAddr in the
 149 // coroutine frame, since it is undefined behavior to resume a coroutine
 150 // suspended at the final suspend point. Thus, in the resume function, we can
 151 // simply remove the last case (when coro::Shape is built, the final suspend
 152 // point (if present) is always the last element of CoroSuspends array).
 153 // In the destroy function, we add a code sequence to check if ResumeFnAddress
 154 // is Null, and if so, jump to the appropriate label to handle cleanup from the
 155 // final suspend point.
 156 static void handleFinalSuspend(IRBuilder<> &Builder, Value *FramePtr,
 157                                coro::Shape &Shape, SwitchInst *Switch,
 158                                bool IsDestroy) {
 159   assert(Shape.HasFinalSuspend);
 160   auto FinalCase = --Switch->case_end();
 161   BasicBlock *ResumeBB = FinalCase.getCaseSuccessor();
 162   Switch->removeCase(FinalCase);
 163   if (IsDestroy) {
 164     BasicBlock *OldSwitchBB = Switch->getParent();
 165     auto *NewSwitchBB = OldSwitchBB->splitBasicBlock(Switch, "Switch");
 166     Builder.SetInsertPoint(OldSwitchBB->getTerminator());
 167     auto *GepIndex = Builder.CreateConstInBoundsGEP2_32(Shape.FrameTy, FramePtr,
 168                                                         0, 0, "ResumeFn.addr");
 169     auto *Load = Builder.CreateLoad(GepIndex);
 170     auto *NullPtr =
 171         ConstantPointerNull::get(cast<PointerType>(Load->getType()));
 172     auto *Cond = Builder.CreateICmpEQ(Load, NullPtr);
 173     Builder.CreateCondBr(Cond, ResumeBB, NewSwitchBB);
 174     OldSwitchBB->getTerminator()->eraseFromParent();
 175   }
 176 }
 177
 178 // Create a resume clone by cloning the body of the original function, setting
 179 // new entry block and replacing coro.suspend an appropriate value to force
 180 // resume or cleanup pass for every suspend point.
 181 static Function *createClone(Function &F, Twine Suffix, coro::Shape &Shape,
 182                              BasicBlock *ResumeEntry, int8_t FnIndex) {
 183   Module *M = F.getParent();
 184   auto *FrameTy = Shape.FrameTy;
 185   auto *FnPtrTy = cast<PointerType>(FrameTy->getElementType(0));
 186   auto *FnTy = cast<FunctionType>(FnPtrTy->getElementType());
 187
 188   Function *NewF =
 189       Function::Create(FnTy, GlobalValue::LinkageTypes::InternalLinkage,
 190                        F.getName() + Suffix, M);
 191   NewF->addAttribute(1, Attribute::NonNull);
 192   NewF->addAttribute(1, Attribute::NoAlias);
 193
 194   ValueToValueMapTy VMap;
 195   // Replace all args with undefs. The buildCoroutineFrame algorithm already
 196   // rewritten access to the args that occurs after suspend points with loads
 197   // and stores to/from the coroutine frame.
 198   for (Argument &A : F.getArgumentList())
 199     VMap[&A] = UndefValue::get(A.getType());
 200
 201   SmallVector<ReturnInst *, 4> Returns;
 202
 203   if (DISubprogram *SP = F.getSubprogram()) {
 204     // If we have debug info, add mapping for the metadata nodes that should not
 205     // be cloned by CloneFunctionInfo.
 206     auto &MD = VMap.MD();
 207     MD[SP->getUnit()].reset(SP->getUnit());
 208     MD[SP->getType()].reset(SP->getType());
 209     MD[SP->getFile()].reset(SP->getFile());
 210   }
 211   CloneFunctionInto(NewF, &F, VMap, /*ModuleLevelChanges=*/true, Returns);
 212
 213   // Remove old returns.
 214   for (ReturnInst *Return : Returns)
 215     changeToUnreachable(Return, /*UseLLVMTrap=*/false);
 216
 217   // Remove old return attributes.
 218   NewF->removeAttributes(
 219       AttributeSet::ReturnIndex,
 220       AttributeSet::get(
 221           NewF->getContext(), AttributeSet::ReturnIndex,
 222           AttributeFuncs::typeIncompatible(NewF->getReturnType())));
 223
 224   // Make AllocaSpillBlock the new entry block.
 225   auto *SwitchBB = cast<BasicBlock>(VMap[ResumeEntry]);
 226   auto *Entry = cast<BasicBlock>(VMap[Shape.AllocaSpillBlock]);
 227   Entry->moveBefore(&NewF->getEntryBlock());
 228   Entry->getTerminator()->eraseFromParent();
 229   BranchInst::Create(SwitchBB, Entry);
 230   Entry->setName("entry" + Suffix);
 231
 232   // Clear all predecessors of the new entry block.
 233   auto *Switch = cast<SwitchInst>(VMap[Shape.ResumeSwitch]);
 234   Entry->replaceAllUsesWith(Switch->getDefaultDest());
 235
 236   IRBuilder<> Builder(&NewF->getEntryBlock().front());
 237
 238   // Remap frame pointer.
 239   Argument *NewFramePtr = &NewF->getArgumentList().front();
 240   Value *OldFramePtr = cast<Value>(VMap[Shape.FramePtr]);
 241   NewFramePtr->takeName(OldFramePtr);
 242   OldFramePtr->replaceAllUsesWith(NewFramePtr);
 243
 244   // Remap vFrame pointer.
 245   auto *NewVFrame = Builder.CreateBitCast(
 246       NewFramePtr, Type::getInt8PtrTy(Builder.getContext()), "vFrame");
 247   Value *OldVFrame = cast<Value>(VMap[Shape.CoroBegin]);
 248   OldVFrame->replaceAllUsesWith(NewVFrame);
 249
 250   // Rewrite final suspend handling as it is not done via switch (allows to
 251   // remove final case from the switch, since it is undefined behavior to resume
 252   // the coroutine suspended at the final suspend point.
 253   if (Shape.HasFinalSuspend) {
 254     auto *Switch = cast<SwitchInst>(VMap[Shape.ResumeSwitch]);
 255     bool IsDestroy = FnIndex != 0;
 256     handleFinalSuspend(Builder, NewFramePtr, Shape, Switch, IsDestroy);
 257   }
 258
 259   // Replace coro suspend with the appropriate resume index.
 260   // Replacing coro.suspend with (0) will result in control flow proceeding to
 261   // a resume label associated with a suspend point, replacing it with (1) will
 262   // result in control flow proceeding to a cleanup label associated with this
 263   // suspend point.
 264   auto *NewValue = Builder.getInt8(FnIndex ? 1 : 0);
 265   for (CoroSuspendInst *CS : Shape.CoroSuspends) {
 266     auto *MappedCS = cast<CoroSuspendInst>(VMap[CS]);
 267     MappedCS->replaceAllUsesWith(NewValue);
 268     MappedCS->eraseFromParent();
 269   }
 270
 271   // Remove coro.end intrinsics.
 272   replaceFallthroughCoroEnd(Shape.CoroEnds.front(), VMap);
 273   // FIXME: coming in upcoming patches:
 274   // replaceUnwindCoroEnds(Shape.CoroEnds, VMap);
 275
 276   // Eliminate coro.free from the clones, replacing it with 'null' in cleanup,
 277   // to suppress deallocation code.
 278   coro::replaceCoroFree(cast<CoroIdInst>(VMap[Shape.CoroBegin->getId()]),
 279                         /*Elide=*/FnIndex == 2);
 280
 281   NewF->setCallingConv(CallingConv::Fast);
 282
 283   return NewF;
 284 }
 285
 286 static void removeCoroEnds(coro::Shape &Shape) {
 287   for (CoroEndInst *CE : Shape.CoroEnds)
 288     CE->eraseFromParent();
 289 }
 290
 291 static void replaceFrameSize(coro::Shape &Shape) {
 292   if (Shape.CoroSizes.empty())
 293     return;
 294
 295   // In the same function all coro.sizes should have the same result type.
 296   auto *SizeIntrin = Shape.CoroSizes.back();
 297   Module *M = SizeIntrin->getModule();
 298   const DataLayout &DL = M->getDataLayout();
 299   auto Size = DL.getTypeAllocSize(Shape.FrameTy);
 300   auto *SizeConstant = ConstantInt::get(SizeIntrin->getType(), Size);
 301
 302   for (CoroSizeInst *CS : Shape.CoroSizes) {
 303     CS->replaceAllUsesWith(SizeConstant);
 304     CS->eraseFromParent();
 305   }
 306 }
 307
 308 // Create a global constant array containing pointers to functions provided and
 309 // set Info parameter of CoroBegin to point at this constant. Example:
 310 //
 311 //   @f.resumers = internal constant [2 x void(%f.frame*)*]
 312 //                    [void(%f.frame*)* @f.resume, void(%f.frame*)* @f.destroy]
 313 //   define void @f() {
 314 //     ...
 315 //     call i8* @llvm.coro.begin(i8* null, i32 0, i8* null,
 316 //                    i8* bitcast([2 x void(%f.frame*)*] * @f.resumers to i8*))
 317 //
 318 // Assumes that all the functions have the same signature.
 319 static void setCoroInfo(Function &F, CoroBeginInst *CoroBegin,
 320                         std::initializer_list<Function *> Fns) {
 321
 322   SmallVector<Constant *, 4> Args(Fns.begin(), Fns.end());
 323   assert(!Args.empty());
 324   Function *Part = *Fns.begin();
 325   Module *M = Part->getParent();
 326   auto *ArrTy = ArrayType::get(Part->getType(), Args.size());
 327
 328   auto *ConstVal = ConstantArray::get(ArrTy, Args);
 329   auto *GV = new GlobalVariable(*M, ConstVal->getType(), /*isConstant=*/true,
 330                                 GlobalVariable::PrivateLinkage, ConstVal,
 331                                 F.getName() + Twine(".resumers"));
 332
 333   // Update coro.begin instruction to refer to this constant.
 334   LLVMContext &C = F.getContext();
 335   auto *BC = ConstantExpr::getPointerCast(GV, Type::getInt8PtrTy(C));
 336   CoroBegin->getId()->setInfo(BC);
 337 }
 338
 339 // Store addresses of Resume/Destroy/Cleanup functions in the coroutine frame.
 340 static void updateCoroFrame(coro::Shape &Shape, Function *ResumeFn,
 341                             Function *DestroyFn, Function *CleanupFn) {
 342
 343   IRBuilder<> Builder(Shape.FramePtr->getNextNode());
 344   auto *ResumeAddr = Builder.CreateConstInBoundsGEP2_32(
 345       Shape.FrameTy, Shape.FramePtr, 0, coro::Shape::ResumeField,
 346       "resume.addr");
 347   Builder.CreateStore(ResumeFn, ResumeAddr);
 348
 349   Value *DestroyOrCleanupFn = DestroyFn;
 350
 351   CoroIdInst *CoroId = Shape.CoroBegin->getId();
 352   if (CoroAllocInst *CA = CoroId->getCoroAlloc()) {
 353     // If there is a CoroAlloc and it returns false (meaning we elide the
 354     // allocation, use CleanupFn instead of DestroyFn).
 355     DestroyOrCleanupFn = Builder.CreateSelect(CA, DestroyFn, CleanupFn);
 356   }
 357
 358   auto *DestroyAddr = Builder.CreateConstInBoundsGEP2_32(
 359       Shape.FrameTy, Shape.FramePtr, 0, coro::Shape::DestroyField,
 360       "destroy.addr");
 361   Builder.CreateStore(DestroyOrCleanupFn, DestroyAddr);
 362 }
 363
 364 static void postSplitCleanup(Function &F) {
 365   removeUnreachableBlocks(F);
 366   llvm::legacy::FunctionPassManager FPM(F.getParent());
 367
 368   FPM.add(createVerifierPass());
 369   FPM.add(createSCCPPass());
 370   FPM.add(createCFGSimplificationPass());
 371   FPM.add(createEarlyCSEPass());
 372   FPM.add(createCFGSimplificationPass());
 373
 374   FPM.doInitialization();
 375   FPM.run(F);
 376   FPM.doFinalization();
 377 }
 378
 379 // Coroutine has no suspend points. Remove heap allocation for the coroutine
 380 // frame if possible.
 381 static void handleNoSuspendCoroutine(CoroBeginInst *CoroBegin, Type *FrameTy) {
 382   auto *CoroId = CoroBegin->getId();
 383   auto *AllocInst = CoroId->getCoroAlloc();
 384   coro::replaceCoroFree(CoroId, /*Elide=*/AllocInst != nullptr);
 385   if (AllocInst) {
 386     IRBuilder<> Builder(AllocInst);
 387     // FIXME: Need to handle overaligned members.
 388     auto *Frame = Builder.CreateAlloca(FrameTy);
 389     auto *VFrame = Builder.CreateBitCast(Frame, Builder.getInt8PtrTy());
 390     AllocInst->replaceAllUsesWith(Builder.getFalse());
 391     AllocInst->eraseFromParent();
 392     CoroBegin->replaceAllUsesWith(VFrame);
 393   } else {
 394     CoroBegin->replaceAllUsesWith(CoroBegin->getMem());
 395   }
 396   CoroBegin->eraseFromParent();
 397 }
 398
 399 // look for a very simple pattern
 400 //    coro.save
 401 //    no other calls
 402 //    resume or destroy call
 403 //    coro.suspend
 404 //
 405 // If there are other calls between coro.save and coro.suspend, they can
 406 // potentially resume or destroy the coroutine, so it is unsafe to eliminate a
 407 // suspend point.
 408 static bool simplifySuspendPoint(CoroSuspendInst *Suspend,
 409                                  CoroBeginInst *CoroBegin) {
 410   auto *Save = Suspend->getCoroSave();
 411   auto *BB = Suspend->getParent();
 412   if (BB != Save->getParent())
 413     return false;
 414
 415   CallSite SingleCallSite;
 416
 417   // Check that we have only one CallSite.
 418   for (Instruction *I = Save->getNextNode(); I != Suspend;
 419        I = I->getNextNode()) {
 420     if (isa<CoroFrameInst>(I))
 421       continue;
 422     if (isa<CoroSubFnInst>(I))
 423       continue;
 424     if (CallSite CS = CallSite(I)) {
 425       if (SingleCallSite)
 426         return false;
 427       else
 428         SingleCallSite = CS;
 429     }
 430   }
 431   auto *CallInstr = SingleCallSite.getInstruction();
 432   if (!CallInstr)
 433     return false;
 434
 435   auto *Callee = SingleCallSite.getCalledValue()->stripPointerCasts();
 436
 437   // See if the callsite is for resumption or destruction of the coroutine.
 438   auto *SubFn = dyn_cast<CoroSubFnInst>(Callee);
 439   if (!SubFn)
 440     return false;
 441
 442   // Does not refer to the current coroutine, we cannot do anything with it.
 443   if (SubFn->getFrame() != CoroBegin)
 444     return false;
 445
 446   // Replace llvm.coro.suspend with the value that results in resumption over
 447   // the resume or cleanup path.
 448   Suspend->replaceAllUsesWith(SubFn->getRawIndex());
 449   Suspend->eraseFromParent();
 450   Save->eraseFromParent();
 451
 452   // No longer need a call to coro.resume or coro.destroy.
 453   CallInstr->eraseFromParent();
 454
 455   if (SubFn->user_empty())
 456     SubFn->eraseFromParent();
 457
 458   return true;
 459 }
 460
 461 // Remove suspend points that are simplified.
 462 static void simplifySuspendPoints(coro::Shape &Shape) {
 463   auto &S = Shape.CoroSuspends;
 464   size_t I = 0, N = S.size();
 465   if (N == 0)
 466     return;
 467   for (;;) {
 468     if (simplifySuspendPoint(S[I], Shape.CoroBegin)) {
 469       if (--N == I)
 470         break;
 471       std::swap(S[I], S[N]);
 472       continue;
 473     }
 474     if (++I == N)
 475       break;
 476   }
 477   S.resize(N);
 478 }
 479
 480 static void splitCoroutine(Function &F, CallGraph &CG, CallGraphSCC &SCC) {
 481   coro::Shape Shape(F);
 482   if (!Shape.CoroBegin)
 483     return;
 484
 485   simplifySuspendPoints(Shape);
 486   buildCoroutineFrame(F, Shape);
 487   replaceFrameSize(Shape);
 488
 489   // If there are no suspend points, no split required, just remove
 490   // the allocation and deallocation blocks, they are not needed.
 491   if (Shape.CoroSuspends.empty()) {
 492     handleNoSuspendCoroutine(Shape.CoroBegin, Shape.FrameTy);
 493     removeCoroEnds(Shape);
 494     postSplitCleanup(F);
 495     coro::updateCallGraph(F, {}, CG, SCC);
 496     return;
 497   }
 498
 499   auto *ResumeEntry = createResumeEntryBlock(F, Shape);
 500   auto ResumeClone = createClone(F, ".resume", Shape, ResumeEntry, 0);
 501   auto DestroyClone = createClone(F, ".destroy", Shape, ResumeEntry, 1);
 502   auto CleanupClone = createClone(F, ".cleanup", Shape, ResumeEntry, 2);
 503
 504   // We no longer need coro.end in F.
 505   removeCoroEnds(Shape);
 506
 507   postSplitCleanup(F);
 508   postSplitCleanup(*ResumeClone);
 509   postSplitCleanup(*DestroyClone);
 510   postSplitCleanup(*CleanupClone);
 511
 512   // Store addresses resume/destroy/cleanup functions in the coroutine frame.
 513   updateCoroFrame(Shape, ResumeClone, DestroyClone, CleanupClone);
 514
 515   // Create a constant array referring to resume/destroy/clone functions pointed
 516   // by the last argument of @llvm.coro.info, so that CoroElide pass can
 517   // determined correct function to call.
 518   setCoroInfo(F, Shape.CoroBegin, {ResumeClone, DestroyClone, CleanupClone});
 519
 520   // Update call graph and add the functions we created to the SCC.
 521   coro::updateCallGraph(F, {ResumeClone, DestroyClone, CleanupClone}, CG, SCC);
 522 }
 523
 524 // When we see the coroutine the first time, we insert an indirect call to a
 525 // devirt trigger function and mark the coroutine that it is now ready for
 526 // split.
 527 static void prepareForSplit(Function &F, CallGraph &CG) {
 528   Module &M = *F.getParent();
 529 #ifndef NDEBUG
 530   Function *DevirtFn = M.getFunction(CORO_DEVIRT_TRIGGER_FN);
 531   assert(DevirtFn && "coro.devirt.trigger function not found");
 532 #endif
 533
 534   F.addFnAttr(CORO_PRESPLIT_ATTR, PREPARED_FOR_SPLIT);
 535
 536   // Insert an indirect call sequence that will be devirtualized by CoroElide
 537   // pass:
 538   //    %0 = call i8* @llvm.coro.subfn.addr(i8* null, i8 -1)
 539   //    %1 = bitcast i8* %0 to void(i8*)*
 540   //    call void %1(i8* null)
 541   coro::LowererBase Lowerer(M);
 542   Instruction *InsertPt = F.getEntryBlock().getTerminator();
 543   auto *Null = ConstantPointerNull::get(Type::getInt8PtrTy(F.getContext()));
 544   auto *DevirtFnAddr =
 545       Lowerer.makeSubFnCall(Null, CoroSubFnInst::RestartTrigger, InsertPt);
 546   auto *IndirectCall = CallInst::Create(DevirtFnAddr, Null, "", InsertPt);
 547
 548   // Update CG graph with an indirect call we just added.
 549   CG[&F]->addCalledFunction(IndirectCall, CG.getCallsExternalNode());
 550 }
 551
 552 // Make sure that there is a devirtualization trigger function that CoroSplit
 553 // pass uses the force restart CGSCC pipeline. If devirt trigger function is not
 554 // found, we will create one and add it to the current SCC.
 555 static void createDevirtTriggerFunc(CallGraph &CG, CallGraphSCC &SCC) {
 556   Module &M = CG.getModule();
 557   if (M.getFunction(CORO_DEVIRT_TRIGGER_FN))
 558     return;
 559
 560   LLVMContext &C = M.getContext();
 561   auto *FnTy = FunctionType::get(Type::getVoidTy(C), Type::getInt8PtrTy(C),
 562                                  /*IsVarArgs=*/false);
 563   Function *DevirtFn =
 564       Function::Create(FnTy, GlobalValue::LinkageTypes::PrivateLinkage,
 565                        CORO_DEVIRT_TRIGGER_FN, &M);
 566   DevirtFn->addFnAttr(Attribute::AlwaysInline);
 567   auto *Entry = BasicBlock::Create(C, "entry", DevirtFn);
 568   ReturnInst::Create(C, Entry);
 569
 570   auto *Node = CG.getOrInsertFunction(DevirtFn);
 571
 572   SmallVector<CallGraphNode *, 8> Nodes(SCC.begin(), SCC.end());
 573   Nodes.push_back(Node);
 574   SCC.initialize(Nodes);
 575 }
 576
 577 //===----------------------------------------------------------------------===//
 578 //                              Top Level Driver
 579 //===----------------------------------------------------------------------===//
 580
 581 namespace {
 582
 583 struct CoroSplit : public CallGraphSCCPass {
 584   static char ID; // Pass identification, replacement for typeid
 585   CoroSplit() : CallGraphSCCPass(ID) {}
 586
 587   bool Run = false;
 588
 589   // A coroutine is identified by the presence of coro.begin intrinsic, if
 590   // we don't have any, this pass has nothing to do.
 591   bool doInitialization(CallGraph &CG) override {
 592     Run = coro::declaresIntrinsics(CG.getModule(), {"llvm.coro.begin"});
 593     return CallGraphSCCPass::doInitialization(CG);
 594   }
 595
 596   bool runOnSCC(CallGraphSCC &SCC) override {
 597     if (!Run)
 598       return false;
 599
 600     // Find coroutines for processing.
 601     SmallVector<Function *, 4> Coroutines;
 602     for (CallGraphNode *CGN : SCC)
 603       if (auto *F = CGN->getFunction())
 604         if (F->hasFnAttribute(CORO_PRESPLIT_ATTR))
 605           Coroutines.push_back(F);
 606
 607     if (Coroutines.empty())
 608       return false;
 609
 610     CallGraph &CG = getAnalysis<CallGraphWrapperPass>().getCallGraph();
 611     createDevirtTriggerFunc(CG, SCC);
 612
 613     for (Function *F : Coroutines) {
 614       Attribute Attr = F->getFnAttribute(CORO_PRESPLIT_ATTR);
 615       StringRef Value = Attr.getValueAsString();
 616       DEBUG(dbgs() << "CoroSplit: Processing coroutine '" << F->getName()
 617                    << "' state: " << Value << "\n");
 618       if (Value == UNPREPARED_FOR_SPLIT) {
 619         prepareForSplit(*F, CG);
 620         continue;
 621       }
 622       F->removeFnAttr(CORO_PRESPLIT_ATTR);
 623       splitCoroutine(*F, CG, SCC);
 624     }
 625     return true;
 626   }
 627
 628   void getAnalysisUsage(AnalysisUsage &AU) const override {
 629     CallGraphSCCPass::getAnalysisUsage(AU);
 630   }
 631 };
 632 }
 633
 634 char CoroSplit::ID = 0;
 635 INITIALIZE_PASS(
 636     CoroSplit, "coro-split",
 637     "Split coroutine into a set of functions driving its state machine", false,
 638     false)
 639
 640 Pass *llvm::createCoroSplitPass() { return new CoroSplit(); }