[Coroutines]: Part6b: Add coro.id intrinsic.

Summary:
1. Make coroutine representation more robust against optimization that may duplicate instruction by introducing coro.id intrinsics that returns a token that will get fed into coro.alloc and coro.begin. Due to coro.id returning a token, it won't get duplicated and can be used as reliable indicator of coroutine identify when a particular coroutine call gets inlined.
2. Move last three arguments of coro.begin into coro.id as they will be shared if coro.begin will get duplicated.
3. doc + test + code updated to support the new intrinsic.

Reviewers: mehdi_amini, majnemer

Subscribers: mehdi_amini, llvm-commits

Differential Revision: https://reviews.llvm.org/D23412

llvm-svn: 278481

5 files changed, 191 insertions, 180 deletions

diff --git a/llvm/lib/Transforms/Coroutines/CoroEarly.cpp b/llvm/lib/Transforms/Coroutines/CoroEarly.cpp
index 89aa40b7277..5ebb2104a20 100644
--- a/llvm/lib/Transforms/Coroutines/CoroEarly.cpp
+++ b/llvm/lib/Transforms/Coroutines/CoroEarly.cpp
@@ -52,11 +52,11 @@ bool Lowerer::lowerEarlyIntrinsics(Function &F) {
       switch (CS.getIntrinsicID()) {
       default:
         continue;
-      case Intrinsic::coro_begin:
+      case Intrinsic::coro_id:
         // Mark a function that comes out of the frontend that has a coro.begin
         // with a coroutine attribute.
-        if (auto *CB = cast<CoroBeginInst>(&I)) {
-          if (CB->getInfo().isPreSplit())
+        if (auto *CII = cast<CoroIdInst>(&I)) {
+          if (CII->getInfo().isPreSplit())
             F.addFnAttr(CORO_PRESPLIT_ATTR, UNPREPARED_FOR_SPLIT);
         }
         break;
diff --git a/llvm/lib/Transforms/Coroutines/CoroElide.cpp b/llvm/lib/Transforms/Coroutines/CoroElide.cpp
index c4568119429..cf15608e0ff 100644
--- a/llvm/lib/Transforms/Coroutines/CoroElide.cpp
+++ b/llvm/lib/Transforms/Coroutines/CoroElide.cpp
@@ -22,46 +22,20 @@ using namespace llvm;
 
 #define DEBUG_TYPE "coro-elide"
 
-//===----------------------------------------------------------------------===//
-//                              Top Level Driver
-//===----------------------------------------------------------------------===//
-
-namespace {
-struct CoroElide : FunctionPass {
-  static char ID;
-  CoroElide() : FunctionPass(ID) {}
-
-  bool NeedsToRun = false;
-
-  bool doInitialization(Module &M) override {
-    NeedsToRun = coro::declaresIntrinsics(M, {"llvm.coro.begin"});
-    return false;
-  }
-
-  bool runOnFunction(Function &F) override;
-  void getAnalysisUsage(AnalysisUsage &AU) const override {
-    AU.addRequired<AAResultsWrapperPass>();
-    AU.setPreservesCFG();
-  }
+// Created on demand if CoroElide pass has work to do.
+struct Lowerer : coro::LowererBase {
+  SmallVector<CoroIdInst *, 4> CoroIds;
+  SmallVector<CoroBeginInst *, 1> CoroBegins;
+  SmallVector<CoroAllocInst *, 1> CoroAllocs;
+  SmallVector<CoroSubFnInst *, 4> ResumeAddr;
+  SmallVector<CoroSubFnInst *, 4> DestroyAddr;
+  SmallVector<CoroFreeInst *, 1> CoroFrees;
+
+  Lowerer(Module &M) : LowererBase(M) {}
+
+  void elideHeapAllocations(Function *F, Type *FrameTy, AAResults &AA);
+  bool processCoroId(CoroIdInst *, AAResults &AA);
 };
-}
-
-char CoroElide::ID = 0;
-INITIALIZE_PASS_BEGIN(
-    CoroElide, "coro-elide",
-    "Coroutine frame allocation elision and indirect calls replacement", false,
-    false)
-INITIALIZE_PASS_DEPENDENCY(AAResultsWrapperPass)
-INITIALIZE_PASS_END(
-    CoroElide, "coro-elide",
-    "Coroutine frame allocation elision and indirect calls replacement", false,
-    false)
-
-Pass *llvm::createCoroElidePass() { return new CoroElide(); }
-
-//===----------------------------------------------------------------------===//
-//                              Implementation
-//===----------------------------------------------------------------------===//
 
 // Go through the list of coro.subfn.addr intrinsics and replace them with the
 // provided constant.
@@ -129,10 +103,30 @@ static Instruction *getFirstNonAllocaInTheEntryBlock(Function *F) {
 
 // To elide heap allocations we need to suppress code blocks guarded by
 // llvm.coro.alloc and llvm.coro.free instructions.
-static void elideHeapAllocations(CoroBeginInst *CoroBegin, Type *FrameTy,
-                                 CoroAllocInst *AllocInst, AAResults &AA) {
-  LLVMContext &C = CoroBegin->getContext();
-  auto *InsertPt = getFirstNonAllocaInTheEntryBlock(CoroBegin->getFunction());
+void Lowerer::elideHeapAllocations(Function *F, Type *FrameTy, AAResults &AA) {
+  LLVMContext &C = FrameTy->getContext();
+  auto *InsertPt =
+      getFirstNonAllocaInTheEntryBlock(CoroIds.front()->getFunction());
+
+  // Replacing llvm.coro.alloc with false will suppress dynamic
+  // allocation as it is expected for the frontend to generate the code that
+  // looks like:
+  //   id = coro.id(...)
+  //   mem = coro.alloc(id) ? malloc(coro.size()) : 0;
+  //   coro.begin(id, mem)
+  auto *False = ConstantInt::getFalse(C);
+  for (auto *CA : CoroAllocs) {
+    CA->replaceAllUsesWith(False);
+    CA->eraseFromParent();
+  }
+
+  // To suppress deallocation code, we replace all llvm.coro.free intrinsics
+  // associated with this coro.begin with null constant.
+  auto *NullPtr = ConstantPointerNull::get(Type::getInt8PtrTy(C));
+  for (auto *CF : CoroFrees) {
+    CF->replaceAllUsesWith(NullPtr);
+    CF->eraseFromParent();
+  }
 
   // FIXME: Design how to transmit alignment information for every alloca that
   // is spilled into the coroutine frame and recreate the alignment information
@@ -142,38 +136,37 @@ static void elideHeapAllocations(CoroBeginInst *CoroBegin, Type *FrameTy,
   auto *FrameVoidPtr =
       new BitCastInst(Frame, Type::getInt8PtrTy(C), "vFrame", InsertPt);
 
-  // Replacing llvm.coro.alloc with non-null value will suppress dynamic
-  // allocation as it is expected for the frontend to generate the code that
-  // looks like:
-  //   mem = coro.alloc();
-  //   if (!mem) mem = malloc(coro.size());
-  //   coro.begin(mem, ...)
-  AllocInst->replaceAllUsesWith(FrameVoidPtr);
-  AllocInst->eraseFromParent();
-
-  // To suppress deallocation code, we replace all llvm.coro.free intrinsics
-  // associated with this coro.begin with null constant.
-  auto *NullPtr = ConstantPointerNull::get(Type::getInt8PtrTy(C));
-  coro::replaceAllCoroFrees(CoroBegin, NullPtr);
-  CoroBegin->lowerTo(FrameVoidPtr);
+  for (auto *CB : CoroBegins) {
+    CB->replaceAllUsesWith(FrameVoidPtr);
+    CB->eraseFromParent();
+  }
 
   // Since now coroutine frame lives on the stack we need to make sure that
   // any tail call referencing it, must be made non-tail call.
   removeTailCallAttribute(Frame, AA);
 }
 
-// See if there are any coro.subfn.addr intrinsics directly referencing
-// the coro.begin. If found, replace them with an appropriate coroutine
-// subfunction associated with that coro.begin.
-static bool replaceIndirectCalls(CoroBeginInst *CoroBegin, AAResults &AA) {
-  SmallVector<CoroSubFnInst *, 8> ResumeAddr;
-  SmallVector<CoroSubFnInst *, 8> DestroyAddr;
-
-  for (User *CF : CoroBegin->users()) {
-    assert(isa<CoroFrameInst>(CF) &&
-           "CoroBegin can be only used by coro.frame instructions");
-    for (User *U : CF->users()) {
-      if (auto *II = dyn_cast<CoroSubFnInst>(U)) {
+bool Lowerer::processCoroId(CoroIdInst *CoroId, AAResults &AA) {
+  CoroBegins.clear();
+  CoroAllocs.clear();
+  ResumeAddr.clear();
+  DestroyAddr.clear();
+
+  // Collect all coro.begin and coro.allocs associated with this coro.id.
+  for (User *U : CoroId->users()) {
+    if (auto *CB = dyn_cast<CoroBeginInst>(U))
+      CoroBegins.push_back(CB);
+    else if (auto *CA = dyn_cast<CoroAllocInst>(U))
+      CoroAllocs.push_back(CA);
+  }
+
+  // Collect all coro.subfn.addrs associated with coro.begin.
+  // Note, we only devirtualize the calls if their coro.subfn.addr refers to
+  // coro.begin directly. If we run into cases where this check is too
+  // conservative, we can consider relaxing the check.
+  for (CoroBeginInst *CB : CoroBegins) {
+    for (User *U : CB->users())
+      if (auto *II = dyn_cast<CoroSubFnInst>(U))
         switch (II->getIndex()) {
         case CoroSubFnInst::ResumeIndex:
           ResumeAddr.push_back(II);
@@ -184,19 +177,16 @@ static bool replaceIndirectCalls(CoroBeginInst *CoroBegin, AAResults &AA) {
         default:
           llvm_unreachable("unexpected coro.subfn.addr constant");
         }
-      }
-    }
   }
-  if (ResumeAddr.empty() && DestroyAddr.empty())
-    return false;
 
-  // PostSplit coro.begin refers to an array of subfunctions in its Info
+  // PostSplit coro.id refers to an array of subfunctions in its Info
   // argument.
-  ConstantArray *Resumers = CoroBegin->getInfo().Resumers;
-  assert(Resumers && "PostSplit coro.begin Info argument must refer to an array"
+  ConstantArray *Resumers = CoroId->getInfo().Resumers;
+  assert(Resumers && "PostSplit coro.id Info argument must refer to an array"
                      "of coroutine subfunctions");
   auto *ResumeAddrConstant =
       ConstantExpr::getExtractValue(Resumers, CoroSubFnInst::ResumeIndex);
+
   replaceWithConstant(ResumeAddrConstant, ResumeAddr);
 
   if (DestroyAddr.empty())
@@ -204,10 +194,12 @@ static bool replaceIndirectCalls(CoroBeginInst *CoroBegin, AAResults &AA) {
 
   auto *DestroyAddrConstant =
       ConstantExpr::getExtractValue(Resumers, CoroSubFnInst::DestroyIndex);
+
   replaceWithConstant(DestroyAddrConstant, DestroyAddr);
 
-  // If llvm.coro.begin refers to llvm.coro.alloc, we can elide the allocation.
-  if (auto *AllocInst = CoroBegin->getAlloc()) {
+  // If there is a coro.alloc that llvm.coro.id refers to, we have the ability
+  // to suppress dynamic allocation.
+  if (!CoroAllocs.empty()) {
     // FIXME: The check above is overly lax. It only checks for whether we have
     // an ability to elide heap allocations, not whether it is safe to do so.
     // We need to do something like:
@@ -216,9 +208,8 @@ static bool replaceIndirectCalls(CoroBeginInst *CoroBegin, AAResults &AA) {
     // then it is safe to elide heap allocation, since the lifetime of coroutine
     // is fully enclosed in its caller.
     auto *FrameTy = getFrameType(cast<Function>(ResumeAddrConstant));
-    elideHeapAllocations(CoroBegin, FrameTy, AllocInst, AA);
+    elideHeapAllocations(CoroId->getFunction(), FrameTy, AA);
   }
-
   return true;
 }
 
@@ -242,25 +233,69 @@ static bool replaceDevirtTrigger(Function &F) {
   return true;
 }
 
-bool CoroElide::runOnFunction(Function &F) {
-  bool Changed = false;
+//===----------------------------------------------------------------------===//
+//                              Top Level Driver
+//===----------------------------------------------------------------------===//
+
+namespace {
+struct CoroElide : FunctionPass {
+  static char ID;
+  CoroElide() : FunctionPass(ID) {}
+
+  std::unique_ptr<Lowerer> L;
 
-  if (F.hasFnAttribute(CORO_PRESPLIT_ATTR))
-    Changed = replaceDevirtTrigger(F);
+  bool doInitialization(Module &M) override {
+    if (coro::declaresIntrinsics(M, {"llvm.coro.id"}))
+      L = llvm::make_unique<Lowerer>(M);
+    return false;
+  }
 
-  // Collect all PostSplit coro.begins.
-  SmallVector<CoroBeginInst *, 4> CoroBegins;
-  for (auto &I : instructions(F))
-    if (auto *CB = dyn_cast<CoroBeginInst>(&I))
-      if (CB->getInfo().isPostSplit())
-        CoroBegins.push_back(CB);
+  bool runOnFunction(Function &F) override {
+    if (!L)
+      return false;
 
-  if (CoroBegins.empty())
-    return Changed;
+    bool Changed = false;
 
-  AAResults &AA = getAnalysis<AAResultsWrapperPass>().getAAResults();
-  for (auto *CB : CoroBegins)
-    Changed |= replaceIndirectCalls(CB, AA);
+    if (F.hasFnAttribute(CORO_PRESPLIT_ATTR))
+      Changed = replaceDevirtTrigger(F);
 
-  return Changed;
+    L->CoroIds.clear();
+    L->CoroFrees.clear();
+
+    // Collect all PostSplit coro.ids and all coro.free.
+    for (auto &I : instructions(F))
+      if (auto *CF = dyn_cast<CoroFreeInst>(&I))
+        L->CoroFrees.push_back(CF);
+      else if (auto *CII = dyn_cast<CoroIdInst>(&I))
+        if (CII->getInfo().isPostSplit())
+          L->CoroIds.push_back(CII);
+
+    // If we did not find any coro.id, there is nothing to do.
+    if (L->CoroIds.empty())
+      return Changed;
+
+    AAResults &AA = getAnalysis<AAResultsWrapperPass>().getAAResults();
+    for (auto *CII : L->CoroIds)
+      Changed |= L->processCoroId(CII, AA);
+
+    return Changed;
+  }
+  void getAnalysisUsage(AnalysisUsage &AU) const override {
+    AU.addRequired<AAResultsWrapperPass>();
+    AU.setPreservesCFG();
+  }
+};
 }
+
+char CoroElide::ID = 0;
+INITIALIZE_PASS_BEGIN(
+    CoroElide, "coro-elide",
+    "Coroutine frame allocation elision and indirect calls replacement", false,
+    false)
+INITIALIZE_PASS_DEPENDENCY(AAResultsWrapperPass)
+INITIALIZE_PASS_END(
+    CoroElide, "coro-elide",
+    "Coroutine frame allocation elision and indirect calls replacement", false,
+    false)
+
+Pass *llvm::createCoroElidePass() { return new CoroElide(); }
diff --git a/llvm/lib/Transforms/Coroutines/CoroInstr.h b/llvm/lib/Transforms/Coroutines/CoroInstr.h
index 1aa1f493edb..5de6a477391 100644
--- a/llvm/lib/Transforms/Coroutines/CoroInstr.h
+++ b/llvm/lib/Transforms/Coroutines/CoroInstr.h
@@ -11,7 +11,7 @@
 // allows you to do things like:
 //
 //     if (auto *SF = dyn_cast<CoroSubFnInst>(Inst))
-//        ... SF->getFrame() ... SF->getAlloc() ...
+//        ... SF->getFrame() ... 
 //
 // All intrinsic function calls are instances of the call instruction, so these
 // are all subclasses of the CallInst class.  Note that none of these classes
@@ -74,52 +74,11 @@ public:
   }
 };
 
-/// This represents the llvm.coro.frame instruction.
-class LLVM_LIBRARY_VISIBILITY CoroFrameInst : public IntrinsicInst {
-public:
-  // Methods to support type inquiry through isa, cast, and dyn_cast:
-  static inline bool classof(const IntrinsicInst *I) {
-    return I->getIntrinsicID() == Intrinsic::coro_frame;
-  }
-  static inline bool classof(const Value *V) {
-    return isa<IntrinsicInst>(V) && classof(cast<IntrinsicInst>(V));
-  }
-};
-
-/// This represents the llvm.coro.free instruction.
-class LLVM_LIBRARY_VISIBILITY CoroFreeInst : public IntrinsicInst {
-public:
-  // Methods to support type inquiry through isa, cast, and dyn_cast:
-  static inline bool classof(const IntrinsicInst *I) {
-    return I->getIntrinsicID() == Intrinsic::coro_free;
-  }
-  static inline bool classof(const Value *V) {
-    return isa<IntrinsicInst>(V) && classof(cast<IntrinsicInst>(V));
-  }
-};
-
-/// This class represents the llvm.coro.begin instruction.
-class LLVM_LIBRARY_VISIBILITY CoroBeginInst : public IntrinsicInst {
-  enum { MemArg, ElideArg, AlignArg, PromiseArg, InfoArg };
-
+/// This represents the llvm.coro.alloc instruction.
+class LLVM_LIBRARY_VISIBILITY CoroIdInst : public IntrinsicInst {
+  enum { AlignArg, PromiseArg, InfoArg };
 public:
-  CoroAllocInst *getAlloc() const {
-    if (auto *CAI = dyn_cast<CoroAllocInst>(
-            getArgOperand(ElideArg)->stripPointerCasts()))
-      return CAI;
-
-    return nullptr;
-  }
-
-  Value *getMem() const { return getArgOperand(MemArg); }
-
-  Constant *getRawInfo() const {
-    return cast<Constant>(getArgOperand(InfoArg)->stripPointerCasts());
-  }
-
-  void setInfo(Constant *C) { setArgOperand(InfoArg, C); }
-
-  // Info argument of coro.begin is
+  // Info argument of coro.id is
   //   fresh out of the frontend: null ;
   //   outlined                 : {Init, Return, Susp1, Susp2, ...} ;
   //   postsplit                : [resume, destroy, cleanup] ;
@@ -153,23 +112,57 @@ public:
     Result.Resumers = cast<ConstantArray>(Initializer);
     return Result;
   }
+  Constant *getRawInfo() const {
+    return cast<Constant>(getArgOperand(InfoArg)->stripPointerCasts());
+  }
+
+  void setInfo(Constant *C) { setArgOperand(InfoArg, C); }
+
+
+  // Methods to support type inquiry through isa, cast, and dyn_cast:
+  static inline bool classof(const IntrinsicInst *I) {
+    return I->getIntrinsicID() == Intrinsic::coro_id;
+  }
+  static inline bool classof(const Value *V) {
+    return isa<IntrinsicInst>(V) && classof(cast<IntrinsicInst>(V));
+  }
+};
+
+/// This represents the llvm.coro.frame instruction.
+class LLVM_LIBRARY_VISIBILITY CoroFrameInst : public IntrinsicInst {
+public:
+  // Methods to support type inquiry through isa, cast, and dyn_cast:
+  static inline bool classof(const IntrinsicInst *I) {
+    return I->getIntrinsicID() == Intrinsic::coro_frame;
+  }
+  static inline bool classof(const Value *V) {
+    return isa<IntrinsicInst>(V) && classof(cast<IntrinsicInst>(V));
+  }
+};
 
-  // Replaces all coro.frame intrinsics that are associated with this coro.begin
-  // to a replacement value and removes coro.begin and all of the coro.frame
-  // intrinsics.
-  void lowerTo(Value* Replacement) {
-    SmallVector<CoroFrameInst*, 4> FrameInsts;
-    for (auto *CF : this->users())
-      FrameInsts.push_back(cast<CoroFrameInst>(CF));
+/// This represents the llvm.coro.free instruction.
+class LLVM_LIBRARY_VISIBILITY CoroFreeInst : public IntrinsicInst {
+public:
+  // Methods to support type inquiry through isa, cast, and dyn_cast:
+  static inline bool classof(const IntrinsicInst *I) {
+    return I->getIntrinsicID() == Intrinsic::coro_free;
+  }
+  static inline bool classof(const Value *V) {
+    return isa<IntrinsicInst>(V) && classof(cast<IntrinsicInst>(V));
+  }
+};
 
-    for (auto *CF : FrameInsts) {
-      CF->replaceAllUsesWith(Replacement);
-      CF->eraseFromParent();
-    }
+/// This class represents the llvm.coro.begin instruction.
+class LLVM_LIBRARY_VISIBILITY CoroBeginInst : public IntrinsicInst {
+  enum { IdArg, MemArg };
 
-    this->eraseFromParent();
+public:
+  CoroIdInst *getId() const {
+    return cast<CoroIdInst>(getArgOperand(IdArg));
   }
 
+  Value *getMem() const { return getArgOperand(MemArg); }
+
   // Methods for support type inquiry through isa, cast, and dyn_cast:
   static inline bool classof(const IntrinsicInst *I) {
     return I->getIntrinsicID() == Intrinsic::coro_begin;
diff --git a/llvm/lib/Transforms/Coroutines/CoroInternal.h b/llvm/lib/Transforms/Coroutines/CoroInternal.h
index a057597194d..fd600493781 100644
--- a/llvm/lib/Transforms/Coroutines/CoroInternal.h
+++ b/llvm/lib/Transforms/Coroutines/CoroInternal.h
@@ -42,6 +42,7 @@ void initializeCoroCleanupPass(PassRegistry &);
 namespace coro {
 
 bool declaresIntrinsics(Module &M, std::initializer_list<StringRef>);
+void replaceAllCoroAllocs(CoroBeginInst *CB, bool Replacement);
 void replaceAllCoroFrees(CoroBeginInst *CB, Value *Replacement);
 
 // Keeps data and helper functions for lowering coroutine intrinsics.
diff --git a/llvm/lib/Transforms/Coroutines/Coroutines.cpp b/llvm/lib/Transforms/Coroutines/Coroutines.cpp
index 569ee9f2f40..cc92a95e935 100644
--- a/llvm/lib/Transforms/Coroutines/Coroutines.cpp
+++ b/llvm/lib/Transforms/Coroutines/Coroutines.cpp
@@ -99,11 +99,11 @@ Value *coro::LowererBase::makeSubFnCall(Value *Arg, int Index,
 static bool isCoroutineIntrinsicName(StringRef Name) {
   // NOTE: Must be sorted!
   static const char *const CoroIntrinsics[] = {
-      "llvm.coro.alloc",   "llvm.coro.begin", "llvm.coro.destroy",
-      "llvm.coro.done",    "llvm.coro.end",   "llvm.coro.frame",
-      "llvm.coro.free",    "llvm.coro.param", "llvm.coro.promise",
-      "llvm.coro.resume",  "llvm.coro.save",  "llvm.coro.size",
-      "llvm.coro.suspend",
+      "llvm.coro.alloc",   "llvm.coro.begin",   "llvm.coro.destroy",
+      "llvm.coro.done",    "llvm.coro.end",     "llvm.coro.frame",
+      "llvm.coro.free",    "llvm.coro.id",      "llvm.coro.param",
+      "llvm.coro.promise", "llvm.coro.resume",  "llvm.coro.save",
+      "llvm.coro.size",    "llvm.coro.suspend",
   };
   return Intrinsic::lookupLLVMIntrinsicByName(CoroIntrinsics, Name) != -1;
 }
@@ -122,21 +122,3 @@ bool coro::declaresIntrinsics(Module &M,
 
   return false;
 }
-
-// Find all llvm.coro.free instructions associated with the provided coro.begin
-// and replace them with the provided replacement value.
-void coro::replaceAllCoroFrees(CoroBeginInst *CB, Value *Replacement) {
-  SmallVector<CoroFreeInst *, 4> CoroFrees;
-  for (User *FramePtr: CB->users())
-    for (User *U : FramePtr->users())
-      if (auto *CF = dyn_cast<CoroFreeInst>(U))
-        CoroFrees.push_back(CF);
-
-  if (CoroFrees.empty())
-    return;
-
-  for (CoroFreeInst *CF : CoroFrees) {
-    CF->replaceAllUsesWith(Replacement);
-    CF->eraseFromParent();
-  }
-}