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

   1 //===- CoroFrame.cpp - Builds and manipulates coroutine frame -------------===//
   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 file contains classes used to discover if for a particular value
  10 // there from sue to definition that crosses a suspend block.
  11 //
  12 // Using the information discovered we form a Coroutine Frame structure to
  13 // contain those values. All uses of those values are replaced with appropriate
  14 // GEP + load from the coroutine frame. At the point of the definition we spill
  15 // the value into the coroutine frame.
  16 //
  17 // TODO: pack values tightly using liveness info.
  18 //===----------------------------------------------------------------------===//
  19
  20 #include "CoroInternal.h"
  21 #include "llvm/ADT/BitVector.h"
  22 #include "llvm/Transforms/Utils/Local.h"
  23 #include "llvm/Config/llvm-config.h"
  24 #include "llvm/IR/CFG.h"
  25 #include "llvm/IR/Dominators.h"
  26 #include "llvm/IR/IRBuilder.h"
  27 #include "llvm/IR/InstIterator.h"
  28 #include "llvm/Support/Debug.h"
  29 #include "llvm/Support/MathExtras.h"
  30 #include "llvm/Support/circular_raw_ostream.h"
  31 #include "llvm/Transforms/Utils/BasicBlockUtils.h"
  32
  33 using namespace llvm;
  34
  35 // The "coro-suspend-crossing" flag is very noisy. There is another debug type,
  36 // "coro-frame", which results in leaner debug spew.
  37 #define DEBUG_TYPE "coro-suspend-crossing"
  38
  39 enum { SmallVectorThreshold = 32 };
  40
  41 // Provides two way mapping between the blocks and numbers.
  42 namespace {
  43 class BlockToIndexMapping {
  44   SmallVector<BasicBlock *, SmallVectorThreshold> V;
  45
  46 public:
  47   size_t size() const { return V.size(); }
  48
  49   BlockToIndexMapping(Function &F) {
  50     for (BasicBlock &BB : F)
  51       V.push_back(&BB);
  52     llvm::sort(V.begin(), V.end());
  53   }
  54
  55   size_t blockToIndex(BasicBlock *BB) const {
  56     auto *I = std::lower_bound(V.begin(), V.end(), BB);
  57     assert(I != V.end() && *I == BB && "BasicBlockNumberng: Unknown block");
  58     return I - V.begin();
  59   }
  60
  61   BasicBlock *indexToBlock(unsigned Index) const { return V[Index]; }
  62 };
  63 } // end anonymous namespace
  64
  65 // The SuspendCrossingInfo maintains data that allows to answer a question
  66 // whether given two BasicBlocks A and B there is a path from A to B that
  67 // passes through a suspend point.
  68 //
  69 // For every basic block 'i' it maintains a BlockData that consists of:
  70 //   Consumes:  a bit vector which contains a set of indices of blocks that can
  71 //              reach block 'i'
  72 //   Kills: a bit vector which contains a set of indices of blocks that can
  73 //          reach block 'i', but one of the path will cross a suspend point
  74 //   Suspend: a boolean indicating whether block 'i' contains a suspend point.
  75 //   End: a boolean indicating whether block 'i' contains a coro.end intrinsic.
  76 //
  77 namespace {
  78 struct SuspendCrossingInfo {
  79   BlockToIndexMapping Mapping;
  80
  81   struct BlockData {
  82     BitVector Consumes;
  83     BitVector Kills;
  84     bool Suspend = false;
  85     bool End = false;
  86   };
  87   SmallVector<BlockData, SmallVectorThreshold> Block;
  88
  89   iterator_range<succ_iterator> successors(BlockData const &BD) const {
  90     BasicBlock *BB = Mapping.indexToBlock(&BD - &Block[0]);
  91     return llvm::successors(BB);
  92   }
  93
  94   BlockData &getBlockData(BasicBlock *BB) {
  95     return Block[Mapping.blockToIndex(BB)];
  96   }
  97
  98   void dump() const;
  99   void dump(StringRef Label, BitVector const &BV) const;
 100
 101   SuspendCrossingInfo(Function &F, coro::Shape &Shape);
 102
 103   bool hasPathCrossingSuspendPoint(BasicBlock *DefBB, BasicBlock *UseBB) const {
 104     size_t const DefIndex = Mapping.blockToIndex(DefBB);
 105     size_t const UseIndex = Mapping.blockToIndex(UseBB);
 106
 107     assert(Block[UseIndex].Consumes[DefIndex] && "use must consume def");
 108     bool const Result = Block[UseIndex].Kills[DefIndex];
 109     LLVM_DEBUG(dbgs() << UseBB->getName() << " => " << DefBB->getName()
 110                       << " answer is " << Result << "\n");
 111     return Result;
 112   }
 113
 114   bool isDefinitionAcrossSuspend(BasicBlock *DefBB, User *U) const {
 115     auto *I = cast<Instruction>(U);
 116
 117     // We rewrote PHINodes, so that only the ones with exactly one incoming
 118     // value need to be analyzed.
 119     if (auto *PN = dyn_cast<PHINode>(I))
 120       if (PN->getNumIncomingValues() > 1)
 121         return false;
 122
 123     BasicBlock *UseBB = I->getParent();
 124     return hasPathCrossingSuspendPoint(DefBB, UseBB);
 125   }
 126
 127   bool isDefinitionAcrossSuspend(Argument &A, User *U) const {
 128     return isDefinitionAcrossSuspend(&A.getParent()->getEntryBlock(), U);
 129   }
 130
 131   bool isDefinitionAcrossSuspend(Instruction &I, User *U) const {
 132     return isDefinitionAcrossSuspend(I.getParent(), U);
 133   }
 134 };
 135 } // end anonymous namespace
 136
 137 #if !defined(NDEBUG) || defined(LLVM_ENABLE_DUMP)
 138 LLVM_DUMP_METHOD void SuspendCrossingInfo::dump(StringRef Label,
 139                                                 BitVector const &BV) const {
 140   dbgs() << Label << ":";
 141   for (size_t I = 0, N = BV.size(); I < N; ++I)
 142     if (BV[I])
 143       dbgs() << " " << Mapping.indexToBlock(I)->getName();
 144   dbgs() << "\n";
 145 }
 146
 147 LLVM_DUMP_METHOD void SuspendCrossingInfo::dump() const {
 148   for (size_t I = 0, N = Block.size(); I < N; ++I) {
 149     BasicBlock *const B = Mapping.indexToBlock(I);
 150     dbgs() << B->getName() << ":\n";
 151     dump("   Consumes", Block[I].Consumes);
 152     dump("      Kills", Block[I].Kills);
 153   }
 154   dbgs() << "\n";
 155 }
 156 #endif
 157
 158 SuspendCrossingInfo::SuspendCrossingInfo(Function &F, coro::Shape &Shape)
 159     : Mapping(F) {
 160   const size_t N = Mapping.size();
 161   Block.resize(N);
 162
 163   // Initialize every block so that it consumes itself
 164   for (size_t I = 0; I < N; ++I) {
 165     auto &B = Block[I];
 166     B.Consumes.resize(N);
 167     B.Kills.resize(N);
 168     B.Consumes.set(I);
 169   }
 170
 171   // Mark all CoroEnd Blocks. We do not propagate Kills beyond coro.ends as
 172   // the code beyond coro.end is reachable during initial invocation of the
 173   // coroutine.
 174   for (auto *CE : Shape.CoroEnds)
 175     getBlockData(CE->getParent()).End = true;
 176
 177   // Mark all suspend blocks and indicate that they kill everything they
 178   // consume. Note, that crossing coro.save also requires a spill, as any code
 179   // between coro.save and coro.suspend may resume the coroutine and all of the
 180   // state needs to be saved by that time.
 181   auto markSuspendBlock = [&](IntrinsicInst *BarrierInst) {
 182     BasicBlock *SuspendBlock = BarrierInst->getParent();
 183     auto &B = getBlockData(SuspendBlock);
 184     B.Suspend = true;
 185     B.Kills |= B.Consumes;
 186   };
 187   for (CoroSuspendInst *CSI : Shape.CoroSuspends) {
 188     markSuspendBlock(CSI);
 189     markSuspendBlock(CSI->getCoroSave());
 190   }
 191
 192   // Iterate propagating consumes and kills until they stop changing.
 193   int Iteration = 0;
 194   (void)Iteration;
 195
 196   bool Changed;
 197   do {
 198     LLVM_DEBUG(dbgs() << "iteration " << ++Iteration);
 199     LLVM_DEBUG(dbgs() << "==============\n");
 200
 201     Changed = false;
 202     for (size_t I = 0; I < N; ++I) {
 203       auto &B = Block[I];
 204       for (BasicBlock *SI : successors(B)) {
 205
 206         auto SuccNo = Mapping.blockToIndex(SI);
 207
 208         // Saved Consumes and Kills bitsets so that it is easy to see
 209         // if anything changed after propagation.
 210         auto &S = Block[SuccNo];
 211         auto SavedConsumes = S.Consumes;
 212         auto SavedKills = S.Kills;
 213
 214         // Propagate Kills and Consumes from block B into its successor S.
 215         S.Consumes |= B.Consumes;
 216         S.Kills |= B.Kills;
 217
 218         // If block B is a suspend block, it should propagate kills into the
 219         // its successor for every block B consumes.
 220         if (B.Suspend) {
 221           S.Kills |= B.Consumes;
 222         }
 223         if (S.Suspend) {
 224           // If block S is a suspend block, it should kill all of the blocks it
 225           // consumes.
 226           S.Kills |= S.Consumes;
 227         } else if (S.End) {
 228           // If block S is an end block, it should not propagate kills as the
 229           // blocks following coro.end() are reached during initial invocation
 230           // of the coroutine while all the data are still available on the
 231           // stack or in the registers.
 232           S.Kills.reset();
 233         } else {
 234           // This is reached when S block it not Suspend nor coro.end and it
 235           // need to make sure that it is not in the kill set.
 236           S.Kills.reset(SuccNo);
 237         }
 238
 239         // See if anything changed.
 240         Changed |= (S.Kills != SavedKills) || (S.Consumes != SavedConsumes);
 241
 242         if (S.Kills != SavedKills) {
 243           LLVM_DEBUG(dbgs() << "\nblock " << I << " follower " << SI->getName()
 244                             << "\n");
 245           LLVM_DEBUG(dump("S.Kills", S.Kills));
 246           LLVM_DEBUG(dump("SavedKills", SavedKills));
 247         }
 248         if (S.Consumes != SavedConsumes) {
 249           LLVM_DEBUG(dbgs() << "\nblock " << I << " follower " << SI << "\n");
 250           LLVM_DEBUG(dump("S.Consume", S.Consumes));
 251           LLVM_DEBUG(dump("SavedCons", SavedConsumes));
 252         }
 253       }
 254     }
 255   } while (Changed);
 256   LLVM_DEBUG(dump());
 257 }
 258
 259 #undef DEBUG_TYPE // "coro-suspend-crossing"
 260 #define DEBUG_TYPE "coro-frame"
 261
 262 // We build up the list of spills for every case where a use is separated
 263 // from the definition by a suspend point.
 264
 265 namespace {
 266 class Spill {
 267   Value *Def = nullptr;
 268   Instruction *User = nullptr;
 269   unsigned FieldNo = 0;
 270
 271 public:
 272   Spill(Value *Def, llvm::User *U) : Def(Def), User(cast<Instruction>(U)) {}
 273
 274   Value *def() const { return Def; }
 275   Instruction *user() const { return User; }
 276   BasicBlock *userBlock() const { return User->getParent(); }
 277
 278   // Note that field index is stored in the first SpillEntry for a particular
 279   // definition. Subsequent mentions of a defintion do not have fieldNo
 280   // assigned. This works out fine as the users of Spills capture the info about
 281   // the definition the first time they encounter it. Consider refactoring
 282   // SpillInfo into two arrays to normalize the spill representation.
 283   unsigned fieldIndex() const {
 284     assert(FieldNo && "Accessing unassigned field");
 285     return FieldNo;
 286   }
 287   void setFieldIndex(unsigned FieldNumber) {
 288     assert(!FieldNo && "Reassigning field number");
 289     FieldNo = FieldNumber;
 290   }
 291 };
 292 } // namespace
 293
 294 // Note that there may be more than one record with the same value of Def in
 295 // the SpillInfo vector.
 296 using SpillInfo = SmallVector<Spill, 8>;
 297
 298 #ifndef NDEBUG
 299 static void dump(StringRef Title, SpillInfo const &Spills) {
 300   dbgs() << "------------- " << Title << "--------------\n";
 301   Value *CurrentValue = nullptr;
 302   for (auto const &E : Spills) {
 303     if (CurrentValue != E.def()) {
 304       CurrentValue = E.def();
 305       CurrentValue->dump();
 306     }
 307     dbgs() << "   user: ";
 308     E.user()->dump();
 309   }
 310 }
 311 #endif
 312
 313 namespace {
 314 // We cannot rely solely on natural alignment of a type when building a
 315 // coroutine frame and if the alignment specified on the Alloca instruction
 316 // differs from the natural alignment of the alloca type we will need to insert
 317 // padding.
 318 struct PaddingCalculator {
 319   const DataLayout &DL;
 320   LLVMContext &Context;
 321   unsigned StructSize = 0;
 322
 323   PaddingCalculator(LLVMContext &Context, DataLayout const &DL)
 324       : DL(DL), Context(Context) {}
 325
 326   // Replicate the logic from IR/DataLayout.cpp to match field offset
 327   // computation for LLVM structs.
 328   void addType(Type *Ty) {
 329     unsigned TyAlign = DL.getABITypeAlignment(Ty);
 330     if ((StructSize & (TyAlign - 1)) != 0)
 331       StructSize = alignTo(StructSize, TyAlign);
 332
 333     StructSize += DL.getTypeAllocSize(Ty); // Consume space for this data item.
 334   }
 335
 336   void addTypes(SmallVectorImpl<Type *> const &Types) {
 337     for (auto *Ty : Types)
 338       addType(Ty);
 339   }
 340
 341   unsigned computePadding(Type *Ty, unsigned ForcedAlignment) {
 342     unsigned TyAlign = DL.getABITypeAlignment(Ty);
 343     auto Natural = alignTo(StructSize, TyAlign);
 344     auto Forced = alignTo(StructSize, ForcedAlignment);
 345
 346     // Return how many bytes of padding we need to insert.
 347     if (Natural != Forced)
 348       return std::max(Natural, Forced) - StructSize;
 349
 350     // Rely on natural alignment.
 351     return 0;
 352   }
 353
 354   // If padding required, return the padding field type to insert.
 355   ArrayType *getPaddingType(Type *Ty, unsigned ForcedAlignment) {
 356     if (auto Padding = computePadding(Ty, ForcedAlignment))
 357       return ArrayType::get(Type::getInt8Ty(Context), Padding);
 358
 359     return nullptr;
 360   }
 361 };
 362 } // namespace
 363
 364 // Build a struct that will keep state for an active coroutine.
 365 //   struct f.frame {
 366 //     ResumeFnTy ResumeFnAddr;
 367 //     ResumeFnTy DestroyFnAddr;
 368 //     int ResumeIndex;
 369 //     ... promise (if present) ...
 370 //     ... spills ...
 371 //   };
 372 static StructType *buildFrameType(Function &F, coro::Shape &Shape,
 373                                   SpillInfo &Spills) {
 374   LLVMContext &C = F.getContext();
 375   const DataLayout &DL = F.getParent()->getDataLayout();
 376   PaddingCalculator Padder(C, DL);
 377   SmallString<32> Name(F.getName());
 378   Name.append(".Frame");
 379   StructType *FrameTy = StructType::create(C, Name);
 380   auto *FramePtrTy = FrameTy->getPointerTo();
 381   auto *FnTy = FunctionType::get(Type::getVoidTy(C), FramePtrTy,
 382                                  /*IsVarArgs=*/false);
 383   auto *FnPtrTy = FnTy->getPointerTo();
 384
 385   // Figure out how wide should be an integer type storing the suspend index.
 386   unsigned IndexBits = std::max(1U, Log2_64_Ceil(Shape.CoroSuspends.size()));
 387   Type *PromiseType = Shape.PromiseAlloca
 388                           ? Shape.PromiseAlloca->getType()->getElementType()
 389                           : Type::getInt1Ty(C);
 390   SmallVector<Type *, 8> Types{FnPtrTy, FnPtrTy, PromiseType,
 391                                Type::getIntNTy(C, IndexBits)};
 392   Value *CurrentDef = nullptr;
 393
 394   Padder.addTypes(Types);
 395
 396   // Create an entry for every spilled value.
 397   for (auto &S : Spills) {
 398     if (CurrentDef == S.def())
 399       continue;
 400
 401     CurrentDef = S.def();
 402     // PromiseAlloca was already added to Types array earlier.
 403     if (CurrentDef == Shape.PromiseAlloca)
 404       continue;
 405
 406     Type *Ty = nullptr;
 407     if (auto *AI = dyn_cast<AllocaInst>(CurrentDef)) {
 408       Ty = AI->getAllocatedType();
 409       if (unsigned AllocaAlignment = AI->getAlignment()) {
 410         // If alignment is specified in alloca, see if we need to insert extra
 411         // padding.
 412         if (auto PaddingTy = Padder.getPaddingType(Ty, AllocaAlignment)) {
 413           Types.push_back(PaddingTy);
 414           Padder.addType(PaddingTy);
 415         }
 416       }
 417     } else {
 418       Ty = CurrentDef->getType();
 419     }
 420     S.setFieldIndex(Types.size());
 421     Types.push_back(Ty);
 422     Padder.addType(Ty);
 423   }
 424   FrameTy->setBody(Types);
 425
 426   return FrameTy;
 427 }
 428
 429 // We need to make room to insert a spill after initial PHIs, but before
 430 // catchswitch instruction. Placing it before violates the requirement that
 431 // catchswitch, like all other EHPads must be the first nonPHI in a block.
 432 //
 433 // Split away catchswitch into a separate block and insert in its place:
 434 //
 435 //   cleanuppad <InsertPt> cleanupret.
 436 //
 437 // cleanupret instruction will act as an insert point for the spill.
 438 static Instruction *splitBeforeCatchSwitch(CatchSwitchInst *CatchSwitch) {
 439   BasicBlock *CurrentBlock = CatchSwitch->getParent();
 440   BasicBlock *NewBlock = CurrentBlock->splitBasicBlock(CatchSwitch);
 441   CurrentBlock->getTerminator()->eraseFromParent();
 442
 443   auto *CleanupPad =
 444       CleanupPadInst::Create(CatchSwitch->getParentPad(), {}, "", CurrentBlock);
 445   auto *CleanupRet =
 446       CleanupReturnInst::Create(CleanupPad, NewBlock, CurrentBlock);
 447   return CleanupRet;
 448 }
 449
 450 // Replace all alloca and SSA values that are accessed across suspend points
 451 // with GetElementPointer from coroutine frame + loads and stores. Create an
 452 // AllocaSpillBB that will become the new entry block for the resume parts of
 453 // the coroutine:
 454 //
 455 //    %hdl = coro.begin(...)
 456 //    whatever
 457 //
 458 // becomes:
 459 //
 460 //    %hdl = coro.begin(...)
 461 //    %FramePtr = bitcast i8* hdl to %f.frame*
 462 //    br label %AllocaSpillBB
 463 //
 464 //  AllocaSpillBB:
 465 //    ; geps corresponding to allocas that were moved to coroutine frame
 466 //    br label PostSpill
 467 //
 468 //  PostSpill:
 469 //    whatever
 470 //
 471 //
 472 static Instruction *insertSpills(SpillInfo &Spills, coro::Shape &Shape) {
 473   auto *CB = Shape.CoroBegin;
 474   IRBuilder<> Builder(CB->getNextNode());
 475   PointerType *FramePtrTy = Shape.FrameTy->getPointerTo();
 476   auto *FramePtr =
 477       cast<Instruction>(Builder.CreateBitCast(CB, FramePtrTy, "FramePtr"));
 478   Type *FrameTy = FramePtrTy->getElementType();
 479
 480   Value *CurrentValue = nullptr;
 481   BasicBlock *CurrentBlock = nullptr;
 482   Value *CurrentReload = nullptr;
 483   unsigned Index = 0; // Proper field number will be read from field definition.
 484
 485   // We need to keep track of any allocas that need "spilling"
 486   // since they will live in the coroutine frame now, all access to them
 487   // need to be changed, not just the access across suspend points
 488   // we remember allocas and their indices to be handled once we processed
 489   // all the spills.
 490   SmallVector<std::pair<AllocaInst *, unsigned>, 4> Allocas;
 491   // Promise alloca (if present) has a fixed field number (Shape::PromiseField)
 492   if (Shape.PromiseAlloca)
 493     Allocas.emplace_back(Shape.PromiseAlloca, coro::Shape::PromiseField);
 494
 495   // Create a load instruction to reload the spilled value from the coroutine
 496   // frame.
 497   auto CreateReload = [&](Instruction *InsertBefore) {
 498     assert(Index && "accessing unassigned field number");
 499     Builder.SetInsertPoint(InsertBefore);
 500     auto *G = Builder.CreateConstInBoundsGEP2_32(FrameTy, FramePtr, 0, Index,
 501                                                  CurrentValue->getName() +
 502                                                      Twine(".reload.addr"));
 503     return isa<AllocaInst>(CurrentValue)
 504                ? G
 505                : Builder.CreateLoad(G,
 506                                     CurrentValue->getName() + Twine(".reload"));
 507   };
 508
 509   for (auto const &E : Spills) {
 510     // If we have not seen the value, generate a spill.
 511     if (CurrentValue != E.def()) {
 512       CurrentValue = E.def();
 513       CurrentBlock = nullptr;
 514       CurrentReload = nullptr;
 515
 516       Index = E.fieldIndex();
 517
 518       if (auto *AI = dyn_cast<AllocaInst>(CurrentValue)) {
 519         // Spilled AllocaInst will be replaced with GEP from the coroutine frame
 520         // there is no spill required.
 521         Allocas.emplace_back(AI, Index);
 522         if (!AI->isStaticAlloca())
 523           report_fatal_error("Coroutines cannot handle non static allocas yet");
 524       } else {
 525         // Otherwise, create a store instruction storing the value into the
 526         // coroutine frame.
 527
 528         Instruction *InsertPt = nullptr;
 529         if (isa<Argument>(CurrentValue)) {
 530           // For arguments, we will place the store instruction right after
 531           // the coroutine frame pointer instruction, i.e. bitcast of
 532           // coro.begin from i8* to %f.frame*.
 533           InsertPt = FramePtr->getNextNode();
 534         } else if (auto *II = dyn_cast<InvokeInst>(CurrentValue)) {
 535           // If we are spilling the result of the invoke instruction, split the
 536           // normal edge and insert the spill in the new block.
 537           auto NewBB = SplitEdge(II->getParent(), II->getNormalDest());
 538           InsertPt = NewBB->getTerminator();
 539         } else if (dyn_cast<PHINode>(CurrentValue)) {
 540           // Skip the PHINodes and EH pads instructions.
 541           BasicBlock *DefBlock = cast<Instruction>(E.def())->getParent();
 542           if (auto *CSI = dyn_cast<CatchSwitchInst>(DefBlock->getTerminator()))
 543             InsertPt = splitBeforeCatchSwitch(CSI);
 544           else
 545             InsertPt = &*DefBlock->getFirstInsertionPt();
 546         } else {
 547           // For all other values, the spill is placed immediately after
 548           // the definition.
 549           assert(!isa<TerminatorInst>(E.def()) && "unexpected terminator");
 550           InsertPt = cast<Instruction>(E.def())->getNextNode();
 551         }
 552
 553         Builder.SetInsertPoint(InsertPt);
 554         auto *G = Builder.CreateConstInBoundsGEP2_32(
 555             FrameTy, FramePtr, 0, Index,
 556             CurrentValue->getName() + Twine(".spill.addr"));
 557         Builder.CreateStore(CurrentValue, G);
 558       }
 559     }
 560
 561     // If we have not seen the use block, generate a reload in it.
 562     if (CurrentBlock != E.userBlock()) {
 563       CurrentBlock = E.userBlock();
 564       CurrentReload = CreateReload(&*CurrentBlock->getFirstInsertionPt());
 565     }
 566
 567     // If we have a single edge PHINode, remove it and replace it with a reload
 568     // from the coroutine frame. (We already took care of multi edge PHINodes
 569     // by rewriting them in the rewritePHIs function).
 570     if (auto *PN = dyn_cast<PHINode>(E.user())) {
 571       assert(PN->getNumIncomingValues() == 1 && "unexpected number of incoming "
 572                                                 "values in the PHINode");
 573       PN->replaceAllUsesWith(CurrentReload);
 574       PN->eraseFromParent();
 575       continue;
 576     }
 577
 578     // Replace all uses of CurrentValue in the current instruction with reload.
 579     E.user()->replaceUsesOfWith(CurrentValue, CurrentReload);
 580   }
 581
 582   BasicBlock *FramePtrBB = FramePtr->getParent();
 583   Shape.AllocaSpillBlock =
 584       FramePtrBB->splitBasicBlock(FramePtr->getNextNode(), "AllocaSpillBB");
 585   Shape.AllocaSpillBlock->splitBasicBlock(&Shape.AllocaSpillBlock->front(),
 586                                           "PostSpill");
 587
 588   Builder.SetInsertPoint(&Shape.AllocaSpillBlock->front());
 589   // If we found any allocas, replace all of their remaining uses with Geps.
 590   for (auto &P : Allocas) {
 591     auto *G =
 592         Builder.CreateConstInBoundsGEP2_32(FrameTy, FramePtr, 0, P.second);
 593     // We are not using ReplaceInstWithInst(P.first, cast<Instruction>(G)) here,
 594     // as we are changing location of the instruction.
 595     G->takeName(P.first);
 596     P.first->replaceAllUsesWith(G);
 597     P.first->eraseFromParent();
 598   }
 599   return FramePtr;
 600 }
 601
 602 // Sets the unwind edge of an instruction to a particular successor.
 603 static void setUnwindEdgeTo(TerminatorInst *TI, BasicBlock *Succ) {
 604   if (auto *II = dyn_cast<InvokeInst>(TI))
 605     II->setUnwindDest(Succ);
 606   else if (auto *CS = dyn_cast<CatchSwitchInst>(TI))
 607     CS->setUnwindDest(Succ);
 608   else if (auto *CR = dyn_cast<CleanupReturnInst>(TI))
 609     CR->setUnwindDest(Succ);
 610   else
 611     llvm_unreachable("unexpected terminator instruction");
 612 }
 613
 614 // Replaces all uses of OldPred with the NewPred block in all PHINodes in a
 615 // block.
 616 static void updatePhiNodes(BasicBlock *DestBB, BasicBlock *OldPred,
 617                            BasicBlock *NewPred,
 618                            PHINode *LandingPadReplacement) {
 619   unsigned BBIdx = 0;
 620   for (BasicBlock::iterator I = DestBB->begin(); isa<PHINode>(I); ++I) {
 621     PHINode *PN = cast<PHINode>(I);
 622
 623     // We manually update the LandingPadReplacement PHINode and it is the last
 624     // PHI Node. So, if we find it, we are done.
 625     if (LandingPadReplacement == PN)
 626       break;
 627
 628     // Reuse the previous value of BBIdx if it lines up.  In cases where we
 629     // have multiple phi nodes with *lots* of predecessors, this is a speed
 630     // win because we don't have to scan the PHI looking for TIBB.  This
 631     // happens because the BB list of PHI nodes are usually in the same
 632     // order.
 633     if (PN->getIncomingBlock(BBIdx) != OldPred)
 634       BBIdx = PN->getBasicBlockIndex(OldPred);
 635
 636     assert(BBIdx != (unsigned)-1 && "Invalid PHI Index!");
 637     PN->setIncomingBlock(BBIdx, NewPred);
 638   }
 639 }
 640
 641 // Uses SplitEdge unless the successor block is an EHPad, in which case do EH
 642 // specific handling.
 643 static BasicBlock *ehAwareSplitEdge(BasicBlock *BB, BasicBlock *Succ,
 644                                     LandingPadInst *OriginalPad,
 645                                     PHINode *LandingPadReplacement) {
 646   auto *PadInst = Succ->getFirstNonPHI();
 647   if (!LandingPadReplacement && !PadInst->isEHPad())
 648     return SplitEdge(BB, Succ);
 649
 650   auto *NewBB = BasicBlock::Create(BB->getContext(), "", BB->getParent(), Succ);
 651   setUnwindEdgeTo(BB->getTerminator(), NewBB);
 652   updatePhiNodes(Succ, BB, NewBB, LandingPadReplacement);
 653
 654   if (LandingPadReplacement) {
 655     auto *NewLP = OriginalPad->clone();
 656     auto *Terminator = BranchInst::Create(Succ, NewBB);
 657     NewLP->insertBefore(Terminator);
 658     LandingPadReplacement->addIncoming(NewLP, NewBB);
 659     return NewBB;
 660   }
 661   Value *ParentPad = nullptr;
 662   if (auto *FuncletPad = dyn_cast<FuncletPadInst>(PadInst))
 663     ParentPad = FuncletPad->getParentPad();
 664   else if (auto *CatchSwitch = dyn_cast<CatchSwitchInst>(PadInst))
 665     ParentPad = CatchSwitch->getParentPad();
 666   else
 667     llvm_unreachable("handling for other EHPads not implemented yet");
 668
 669   auto *NewCleanupPad = CleanupPadInst::Create(ParentPad, {}, "", NewBB);
 670   CleanupReturnInst::Create(NewCleanupPad, Succ, NewBB);
 671   return NewBB;
 672 }
 673
 674 static void rewritePHIs(BasicBlock &BB) {
 675   // For every incoming edge we will create a block holding all
 676   // incoming values in a single PHI nodes.
 677   //
 678   // loop:
 679   //    %n.val = phi i32[%n, %entry], [%inc, %loop]
 680   //
 681   // It will create:
 682   //
 683   // loop.from.entry:
 684   //    %n.loop.pre = phi i32 [%n, %entry]
 685   //    br %label loop
 686   // loop.from.loop:
 687   //    %inc.loop.pre = phi i32 [%inc, %loop]
 688   //    br %label loop
 689   //
 690   // After this rewrite, further analysis will ignore any phi nodes with more
 691   // than one incoming edge.
 692
 693   // TODO: Simplify PHINodes in the basic block to remove duplicate
 694   // predecessors.
 695
 696   LandingPadInst *LandingPad = nullptr;
 697   PHINode *ReplPHI = nullptr;
 698   if ((LandingPad = dyn_cast_or_null<LandingPadInst>(BB.getFirstNonPHI()))) {
 699     // ehAwareSplitEdge will clone the LandingPad in all the edge blocks.
 700     // We replace the original landing pad with a PHINode that will collect the
 701     // results from all of them.
 702     ReplPHI = PHINode::Create(LandingPad->getType(), 1, "", LandingPad);
 703     ReplPHI->takeName(LandingPad);
 704     LandingPad->replaceAllUsesWith(ReplPHI);
 705     // We will erase the original landing pad at the end of this function after
 706     // ehAwareSplitEdge cloned it in the transition blocks.
 707   }
 708
 709   SmallVector<BasicBlock *, 8> Preds(pred_begin(&BB), pred_end(&BB));
 710   for (BasicBlock *Pred : Preds) {
 711     auto *IncomingBB = ehAwareSplitEdge(Pred, &BB, LandingPad, ReplPHI);
 712     IncomingBB->setName(BB.getName() + Twine(".from.") + Pred->getName());
 713     auto *PN = cast<PHINode>(&BB.front());
 714     do {
 715       int Index = PN->getBasicBlockIndex(IncomingBB);
 716       Value *V = PN->getIncomingValue(Index);
 717       PHINode *InputV = PHINode::Create(
 718           V->getType(), 1, V->getName() + Twine(".") + BB.getName(),
 719           &IncomingBB->front());
 720       InputV->addIncoming(V, Pred);
 721       PN->setIncomingValue(Index, InputV);
 722       PN = dyn_cast<PHINode>(PN->getNextNode());
 723     } while (PN != ReplPHI); // ReplPHI is either null or the PHI that replaced
 724                              // the landing pad.
 725   }
 726
 727   if (LandingPad) {
 728     // Calls to ehAwareSplitEdge function cloned the original lading pad.
 729     // No longer need it.
 730     LandingPad->eraseFromParent();
 731   }
 732 }
 733
 734 static void rewritePHIs(Function &F) {
 735   SmallVector<BasicBlock *, 8> WorkList;
 736
 737   for (BasicBlock &BB : F)
 738     if (auto *PN = dyn_cast<PHINode>(&BB.front()))
 739       if (PN->getNumIncomingValues() > 1)
 740         WorkList.push_back(&BB);
 741
 742   for (BasicBlock *BB : WorkList)
 743     rewritePHIs(*BB);
 744 }
 745
 746 // Check for instructions that we can recreate on resume as opposed to spill
 747 // the result into a coroutine frame.
 748 static bool materializable(Instruction &V) {
 749   return isa<CastInst>(&V) || isa<GetElementPtrInst>(&V) ||
 750          isa<BinaryOperator>(&V) || isa<CmpInst>(&V) || isa<SelectInst>(&V);
 751 }
 752
 753 // Check for structural coroutine intrinsics that should not be spilled into
 754 // the coroutine frame.
 755 static bool isCoroutineStructureIntrinsic(Instruction &I) {
 756   return isa<CoroIdInst>(&I) || isa<CoroSaveInst>(&I) ||
 757          isa<CoroSuspendInst>(&I);
 758 }
 759
 760 // For every use of the value that is across suspend point, recreate that value
 761 // after a suspend point.
 762 static void rewriteMaterializableInstructions(IRBuilder<> &IRB,
 763                                               SpillInfo const &Spills) {
 764   BasicBlock *CurrentBlock = nullptr;
 765   Instruction *CurrentMaterialization = nullptr;
 766   Instruction *CurrentDef = nullptr;
 767
 768   for (auto const &E : Spills) {
 769     // If it is a new definition, update CurrentXXX variables.
 770     if (CurrentDef != E.def()) {
 771       CurrentDef = cast<Instruction>(E.def());
 772       CurrentBlock = nullptr;
 773       CurrentMaterialization = nullptr;
 774     }
 775
 776     // If we have not seen this block, materialize the value.
 777     if (CurrentBlock != E.userBlock()) {
 778       CurrentBlock = E.userBlock();
 779       CurrentMaterialization = cast<Instruction>(CurrentDef)->clone();
 780       CurrentMaterialization->setName(CurrentDef->getName());
 781       CurrentMaterialization->insertBefore(
 782           &*CurrentBlock->getFirstInsertionPt());
 783     }
 784
 785     if (auto *PN = dyn_cast<PHINode>(E.user())) {
 786       assert(PN->getNumIncomingValues() == 1 && "unexpected number of incoming "
 787                                                 "values in the PHINode");
 788       PN->replaceAllUsesWith(CurrentMaterialization);
 789       PN->eraseFromParent();
 790       continue;
 791     }
 792
 793     // Replace all uses of CurrentDef in the current instruction with the
 794     // CurrentMaterialization for the block.
 795     E.user()->replaceUsesOfWith(CurrentDef, CurrentMaterialization);
 796   }
 797 }
 798
 799 // Move early uses of spilled variable after CoroBegin.
 800 // For example, if a parameter had address taken, we may end up with the code
 801 // like:
 802 //        define @f(i32 %n) {
 803 //          %n.addr = alloca i32
 804 //          store %n, %n.addr
 805 //          ...
 806 //          call @coro.begin
 807 //    we need to move the store after coro.begin
 808 static void moveSpillUsesAfterCoroBegin(Function &F, SpillInfo const &Spills,
 809                                         CoroBeginInst *CoroBegin) {
 810   DominatorTree DT(F);
 811   SmallVector<Instruction *, 8> NeedsMoving;
 812
 813   Value *CurrentValue = nullptr;
 814
 815   for (auto const &E : Spills) {
 816     if (CurrentValue == E.def())
 817       continue;
 818
 819     CurrentValue = E.def();
 820
 821     for (User *U : CurrentValue->users()) {
 822       Instruction *I = cast<Instruction>(U);
 823       if (!DT.dominates(CoroBegin, I)) {
 824         LLVM_DEBUG(dbgs() << "will move: " << *I << "\n");
 825
 826         // TODO: Make this more robust. Currently if we run into a situation
 827         // where simple instruction move won't work we panic and
 828         // report_fatal_error.
 829         for (User *UI : I->users()) {
 830           if (!DT.dominates(CoroBegin, cast<Instruction>(UI)))
 831             report_fatal_error("cannot move instruction since its users are not"
 832                                " dominated by CoroBegin");
 833         }
 834
 835         NeedsMoving.push_back(I);
 836       }
 837     }
 838   }
 839
 840   Instruction *InsertPt = CoroBegin->getNextNode();
 841   for (Instruction *I : NeedsMoving)
 842     I->moveBefore(InsertPt);
 843 }
 844
 845 // Splits the block at a particular instruction unless it is the first
 846 // instruction in the block with a single predecessor.
 847 static BasicBlock *splitBlockIfNotFirst(Instruction *I, const Twine &Name) {
 848   auto *BB = I->getParent();
 849   if (&BB->front() == I) {
 850     if (BB->getSinglePredecessor()) {
 851       BB->setName(Name);
 852       return BB;
 853     }
 854   }
 855   return BB->splitBasicBlock(I, Name);
 856 }
 857
 858 // Split above and below a particular instruction so that it
 859 // will be all alone by itself in a block.
 860 static void splitAround(Instruction *I, const Twine &Name) {
 861   splitBlockIfNotFirst(I, Name);
 862   splitBlockIfNotFirst(I->getNextNode(), "After" + Name);
 863 }
 864
 865 void coro::buildCoroutineFrame(Function &F, Shape &Shape) {
 866   // Lower coro.dbg.declare to coro.dbg.value, since we are going to rewrite
 867   // access to local variables.
 868   LowerDbgDeclare(F);
 869
 870   Shape.PromiseAlloca = Shape.CoroBegin->getId()->getPromise();
 871   if (Shape.PromiseAlloca) {
 872     Shape.CoroBegin->getId()->clearPromise();
 873   }
 874
 875   // Make sure that all coro.save, coro.suspend and the fallthrough coro.end
 876   // intrinsics are in their own blocks to simplify the logic of building up
 877   // SuspendCrossing data.
 878   for (CoroSuspendInst *CSI : Shape.CoroSuspends) {
 879     splitAround(CSI->getCoroSave(), "CoroSave");
 880     splitAround(CSI, "CoroSuspend");
 881   }
 882
 883   // Put CoroEnds into their own blocks.
 884   for (CoroEndInst *CE : Shape.CoroEnds)
 885     splitAround(CE, "CoroEnd");
 886
 887   // Transforms multi-edge PHI Nodes, so that any value feeding into a PHI will
 888   // never has its definition separated from the PHI by the suspend point.
 889   rewritePHIs(F);
 890
 891   // Build suspend crossing info.
 892   SuspendCrossingInfo Checker(F, Shape);
 893
 894   IRBuilder<> Builder(F.getContext());
 895   SpillInfo Spills;
 896
 897   for (int Repeat = 0; Repeat < 4; ++Repeat) {
 898     // See if there are materializable instructions across suspend points.
 899     for (Instruction &I : instructions(F))
 900       if (materializable(I))
 901         for (User *U : I.users())
 902           if (Checker.isDefinitionAcrossSuspend(I, U))
 903             Spills.emplace_back(&I, U);
 904
 905     if (Spills.empty())
 906       break;
 907
 908     // Rewrite materializable instructions to be materialized at the use point.
 909     LLVM_DEBUG(dump("Materializations", Spills));
 910     rewriteMaterializableInstructions(Builder, Spills);
 911     Spills.clear();
 912   }
 913
 914   // Collect the spills for arguments and other not-materializable values.
 915   for (Argument &A : F.args())
 916     for (User *U : A.users())
 917       if (Checker.isDefinitionAcrossSuspend(A, U))
 918         Spills.emplace_back(&A, U);
 919
 920   for (Instruction &I : instructions(F)) {
 921     // Values returned from coroutine structure intrinsics should not be part
 922     // of the Coroutine Frame.
 923     if (isCoroutineStructureIntrinsic(I) || &I == Shape.CoroBegin)
 924       continue;
 925     // The Coroutine Promise always included into coroutine frame, no need to
 926     // check for suspend crossing.
 927     if (Shape.PromiseAlloca == &I)
 928       continue;
 929
 930     for (User *U : I.users())
 931       if (Checker.isDefinitionAcrossSuspend(I, U)) {
 932         // We cannot spill a token.
 933         if (I.getType()->isTokenTy())
 934           report_fatal_error(
 935               "token definition is separated from the use by a suspend point");
 936         Spills.emplace_back(&I, U);
 937       }
 938   }
 939   LLVM_DEBUG(dump("Spills", Spills));
 940   moveSpillUsesAfterCoroBegin(F, Spills, Shape.CoroBegin);
 941   Shape.FrameTy = buildFrameType(F, Shape, Spills);
 942   Shape.FramePtr = insertSpills(Spills, Shape);
 943 }