Part 4c: Coroutine Devirtualization: Devirtualize coro.resume and coro.destroy.

Summary:
This is the 4c patch of the coroutine series. CoroElide pass now checks if PostSplit coro.begin
is referenced by coro.subfn.addr intrinsics. If so replace coro.subfn.addrs with an appropriate coroutine
subfunction associated with that coro.begin.

Documentation and overview is here: http://llvm.org/docs/Coroutines.html.

Upstreaming sequence (rough plan)
1.Add documentation. (https://reviews.llvm.org/D22603)
2.Add coroutine intrinsics. (https://reviews.llvm.org/D22659)
3.Add empty coroutine passes. (https://reviews.llvm.org/D22847)
4.Add coroutine devirtualization + tests.
ab) Lower coro.resume and coro.destroy (https://reviews.llvm.org/D22998)
c) Do devirtualization <= we are here
5.Add CGSCC restart trigger + tests.
6.Add coroutine heap elision + tests.
7.Add the rest of the logic (split into more patches)

Reviewers: majnemer

Subscribers: mehdi_amini, llvm-commits

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

llvm-svn: 277908

6 files changed, 175 insertions, 35 deletions

diff --git a/llvm/lib/IR/Verifier.cpp b/llvm/lib/IR/Verifier.cpp
index 2b6bd3aa64f..05140403116 100644
--- a/llvm/lib/IR/Verifier.cpp
+++ b/llvm/lib/IR/Verifier.cpp
@@ -3837,6 +3837,20 @@ void Verifier::visitIntrinsicCallSite(Intrinsic::ID ID, CallSite CS) {
   switch (ID) {
   default:
     break;
+  case Intrinsic::coro_begin: {
+    auto *InfoArg = CS.getArgOperand(3)->stripPointerCasts();
+    if (isa<ConstantPointerNull>(InfoArg))
+      break;
+    auto *GV = dyn_cast<GlobalVariable>(InfoArg);
+    Assert(GV && GV->isConstant() && GV->hasDefinitiveInitializer(),
+      "info argument of llvm.coro.begin must refer to an initialized "
+      "constant");
+    Constant *Init = GV->getInitializer();
+    Assert(isa<ConstantStruct>(Init) || isa<ConstantArray>(Init),
+      "info argument of llvm.coro.begin must refer to either a struct or "
+      "an array");
+    break;
+  }
   case Intrinsic::ctlz:  // llvm.ctlz
   case Intrinsic::cttz:  // llvm.cttz
     Assert(isa<ConstantInt>(CS.getArgOperand(1)),
diff --git a/llvm/lib/Transforms/Coroutines/CoroEarly.cpp b/llvm/lib/Transforms/Coroutines/CoroEarly.cpp
index 786245ab2ac..40baca5a897 100644
--- a/llvm/lib/Transforms/Coroutines/CoroEarly.cpp
+++ b/llvm/lib/Transforms/Coroutines/CoroEarly.cpp
@@ -28,7 +28,6 @@ class Lowerer : public coro::LowererBase {
 
 public:
   Lowerer(Module &M) : LowererBase(M) {}
-  static std::unique_ptr<Lowerer> createIfNeeded(Module &M);
   bool lowerEarlyIntrinsics(Function &F);
 };
 }
@@ -61,21 +60,11 @@ bool Lowerer::lowerEarlyIntrinsics(Function &F) {
         break;
       }
       Changed = true;
-      continue;
     }
   }
   return Changed;
 }
 
-// This pass has work to do only if we find intrinsics we are going to lower in
-// the module.
-std::unique_ptr<Lowerer> Lowerer::createIfNeeded(Module &M) {
-  if (declaresIntrinsics(M, {"llvm.coro.resume", "llvm.coro.destroy"}))
-    return llvm::make_unique<Lowerer>(M);
-
-  return {};
-}
-
 //===----------------------------------------------------------------------===//
 //                              Top Level Driver
 //===----------------------------------------------------------------------===//
@@ -88,8 +77,11 @@ struct CoroEarly : public FunctionPass {
 
   std::unique_ptr<Lowerer> L;
 
+  // This pass has work to do only if we find intrinsics we are going to lower
+  // in the module.
   bool doInitialization(Module &M) override {
-    L = Lowerer::createIfNeeded(M);
+    if (coro::declaresIntrinsics(M, {"llvm.coro.resume", "llvm.coro.destroy"}))
+      L = llvm::make_unique<Lowerer>(M);
     return false;
   }
 
@@ -104,7 +96,6 @@ struct CoroEarly : public FunctionPass {
     AU.setPreservesCFG();
   }
 };
-
 }
 
 char CoroEarly::ID = 0;
diff --git a/llvm/lib/Transforms/Coroutines/CoroElide.cpp b/llvm/lib/Transforms/Coroutines/CoroElide.cpp
index b21380ca008..dc8dad37931 100644
--- a/llvm/lib/Transforms/Coroutines/CoroElide.cpp
+++ b/llvm/lib/Transforms/Coroutines/CoroElide.cpp
@@ -13,6 +13,9 @@
 
 #include "CoroInternal.h"
 #include "llvm/Analysis/AliasAnalysis.h"
+#include "llvm/Analysis/InstructionSimplify.h"
+#include "llvm/IR/ConstantFolder.h"
+#include "llvm/IR/InstIterator.h"
 #include "llvm/Pass.h"
 
 using namespace llvm;
@@ -24,16 +27,104 @@ using namespace llvm;
 //===----------------------------------------------------------------------===//
 
 namespace {
-
 struct CoroElide : FunctionPass {
   static char ID;
   CoroElide() : FunctionPass(ID) {}
-  bool runOnFunction(Function &F) override { return false; }
+
+  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.setPreservesAll();
+    AU.setPreservesCFG();
   }
 };
+}
+
+// Go through the list of coro.subfn.addr intrinsics and replace them with the
+// provided constant.
+static void replaceWithConstant(Constant *Value,
+                                SmallVectorImpl<CoroSubFnInst *> &Users) {
+  if (Users.empty())
+    return;
+
+  // See if we need to bitcast the constant to match the type of the intrinsic
+  // being replaced. Note: All coro.subfn.addr intrinsics return the same type,
+  // so we only need to examine the type of the first one in the list.
+  Type *IntrTy = Users.front()->getType();
+  Type *ValueTy = Value->getType();
+  if (ValueTy != IntrTy) {
+    // May need to tweak the function type to match the type expected at the
+    // use site.
+    assert(ValueTy->isPointerTy() && IntrTy->isPointerTy());
+    Value = ConstantExpr::getBitCast(Value, IntrTy);
+  }
+
+  // Now the value type matches the type of the intrinsic. Replace them all!
+  for (CoroSubFnInst *I : Users)
+    replaceAndRecursivelySimplify(I, Value);
+}
+
+// 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) {
+  SmallVector<CoroSubFnInst *, 8> ResumeAddr;
+  SmallVector<CoroSubFnInst *, 8> DestroyAddr;
+
+  for (User *U : CoroBegin->users()) {
+    if (auto *II = dyn_cast<CoroSubFnInst>(U)) {
+      switch (II->getIndex()) {
+      case CoroSubFnInst::ResumeIndex:
+        ResumeAddr.push_back(II);
+        break;
+      case CoroSubFnInst::DestroyIndex:
+        DestroyAddr.push_back(II);
+        break;
+      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
+  // argument.
+  ConstantArray *Resumers = CoroBegin->getInfo().Resumers;
+  assert(Resumers && "PostSplit coro.begin Info argument must refer to an array"
+                     "of coroutine subfunctions");
+  auto *ResumeAddrConstant =
+      ConstantExpr::getExtractValue(Resumers, CoroSubFnInst::ResumeIndex);
+  auto *DestroyAddrConstant =
+      ConstantExpr::getExtractValue(Resumers, CoroSubFnInst::DestroyIndex);
+
+  replaceWithConstant(ResumeAddrConstant, ResumeAddr);
+  replaceWithConstant(DestroyAddrConstant, DestroyAddr);
+  return true;
+}
+
+bool CoroElide::runOnFunction(Function &F) {
+  // 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);
+
+  if (CoroBegins.empty())
+    return false;
+
+  bool Changed = false;
+
+  for (auto *CB : CoroBegins)
+    Changed |= replaceIndirectCalls(CB);
 
+  return Changed;
 }
 
 char CoroElide::ID = 0;
diff --git a/llvm/lib/Transforms/Coroutines/CoroInstr.h b/llvm/lib/Transforms/Coroutines/CoroInstr.h
index e7f99cfb4de..6f531f40d6d 100644
--- a/llvm/lib/Transforms/Coroutines/CoroInstr.h
+++ b/llvm/lib/Transforms/Coroutines/CoroInstr.h
@@ -61,4 +61,58 @@ public:
   }
 };
 
+/// This class represents the llvm.coro.begin instruction.
+class LLVM_LIBRARY_VISIBILITY CoroBeginInst : public IntrinsicInst {
+  enum { MemArg, AlignArg, PromiseArg, InfoArg };
+
+public:
+  Constant *getRawInfo() const {
+    return cast<Constant>(getArgOperand(InfoArg)->stripPointerCasts());
+  }
+
+  void setInfo(Constant *C) { setArgOperand(InfoArg, C); }
+
+  // Info argument of coro.begin is
+  //   fresh out of the frontend: null ;
+  //   outlined                 : {Init, Return, Susp1, Susp2, ...} ;
+  //   postsplit                : [resume, destroy, cleanup] ;
+  //
+  // If parts of the coroutine were outlined to protect against undesirable
+  // code motion, these functions will be stored in a struct literal referred to
+  // by the Info parameter. Note: this is only needed before coroutine is split.
+  //
+  // After coroutine is split, resume functions are stored in an array
+  // referred to by this parameter.
+
+  struct Info {
+    ConstantStruct *OutlinedParts = nullptr;
+    ConstantArray *Resumers = nullptr;
+
+    bool hasOutlinedParts() const { return OutlinedParts != nullptr; }
+    bool isPostSplit() const { return Resumers != nullptr; }
+  };
+  Info getInfo() const {
+    Info Result;
+    auto *GV = dyn_cast<GlobalVariable>(getRawInfo());
+    if (!GV)
+      return Result;
+
+    assert(GV->isConstant() && GV->hasDefinitiveInitializer());
+    Constant *Initializer = GV->getInitializer();
+    if ((Result.OutlinedParts = dyn_cast<ConstantStruct>(Initializer)))
+      return Result;
+
+    Result.Resumers = cast<ConstantArray>(Initializer);
+    return Result;
+  }
+
+  // Methods for support type inquiry through isa, cast, and dyn_cast:
+  static inline bool classof(const IntrinsicInst *I) {
+    return I->getIntrinsicID() == Intrinsic::coro_begin;
+  }
+  static inline bool classof(const Value *V) {
+    return isa<IntrinsicInst>(V) && classof(cast<IntrinsicInst>(V));
+  }
+};
+
 } // End namespace llvm.
diff --git a/llvm/lib/Transforms/Coroutines/CoroInternal.h b/llvm/lib/Transforms/Coroutines/CoroInternal.h
index 38963c46a00..95e988cdf4e 100644
--- a/llvm/lib/Transforms/Coroutines/CoroInternal.h
+++ b/llvm/lib/Transforms/Coroutines/CoroInternal.h
@@ -17,9 +17,6 @@
 
 namespace llvm {
 
-class FunctionType;
-class LLVMContext;
-class Module;
 class PassRegistry;
 
 void initializeCoroEarlyPass(PassRegistry &);
@@ -29,6 +26,8 @@ void initializeCoroCleanupPass(PassRegistry &);
 
 namespace coro {
 
+bool declaresIntrinsics(Module &M, std::initializer_list<StringRef>);
+
 // Keeps data and helper functions for lowering coroutine intrinsics.
 struct LowererBase {
   Module &TheModule;
@@ -37,7 +36,6 @@ struct LowererBase {
 
   LowererBase(Module &M);
   Value *makeSubFnCall(Value *Arg, int Index, Instruction *InsertPt);
-  static bool declaresIntrinsics(Module &M, std::initializer_list<StringRef>);
 };
 
 } // End namespace coro.
diff --git a/llvm/lib/Transforms/Coroutines/Coroutines.cpp b/llvm/lib/Transforms/Coroutines/Coroutines.cpp
index 328759ad2bb..f9661f3575c 100644
--- a/llvm/lib/Transforms/Coroutines/Coroutines.cpp
+++ b/llvm/lib/Transforms/Coroutines/Coroutines.cpp
@@ -99,19 +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.param", "llvm.coro.promise",
+      "llvm.coro.resume",  "llvm.coro.save",  "llvm.coro.size",
+      "llvm.coro.suspend",
   };
   return Intrinsic::lookupLLVMIntrinsicByName(CoroIntrinsics, Name) != -1;
 }
@@ -119,8 +111,8 @@ static bool isCoroutineIntrinsicName(StringRef Name) {
 
 // Verifies if a module has named values listed. Also, in debug mode verifies
 // that names are intrinsic names.
-bool coro::LowererBase::declaresIntrinsics(
-    Module &M, std::initializer_list<StringRef> List) {
+bool coro::declaresIntrinsics(Module &M,
+                              std::initializer_list<StringRef> List) {
 
   for (StringRef Name : List) {
     assert(isCoroutineIntrinsicName(Name) && "not a coroutine intrinsic");