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

include/llvm/IR/Intrinsics.td

@@ -634,7 +634,7 @@ def int_coro_promise : Intrinsic<[llvm_ptr_ty],
 // Coroutine Lowering Intrinsics. Used internally by coroutine passes.
 
 def int_coro_subfn_addr : Intrinsic<[llvm_ptr_ty], [llvm_ptr_ty, llvm_i8_ty],
-                                    [IntrArgMemOnly, ReadOnly<0>,
+                                    [IntrReadMem, IntrArgMemOnly, ReadOnly<0>,
                                      NoCapture<0>]>;
 
 ///===-------------------------- Other Intrinsics --------------------------===//

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)),

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;

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;

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.

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.

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.alloc",   "llvm.coro.begin", "llvm.coro.destroy",
-    "llvm.coro.begin",
+      "llvm.coro.done",    "llvm.coro.end",   "llvm.coro.frame",
-    "llvm.coro.destroy",
+      "llvm.coro.free",    "llvm.coro.param", "llvm.coro.promise",
-    "llvm.coro.done",
+      "llvm.coro.resume",  "llvm.coro.save",  "llvm.coro.size",
-    "llvm.coro.end",
+      "llvm.coro.suspend",
-    "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(
+bool coro::declaresIntrinsics(Module &M,
-    Module &M, std::initializer_list<StringRef> List) {
+                              std::initializer_list<StringRef> List) {
 
   for (StringRef Name : List) {
     assert(isCoroutineIntrinsicName(Name) && "not a coroutine intrinsic");

test/Transforms/Coroutines/coro-early.ll

@@ -2,7 +2,7 @@
 ; intrinsics.
 ; RUN: opt < %s -S -coro-early | FileCheck %s
 
-; CHECK-LABEL: @callResume
+; CHECK-LABEL: @callResume(
 define void @callResume(i8* %hdl) {
 ; CHECK-NEXT: entry
 entry:
@@ -20,7 +20,7 @@ entry:
 ; CHECK-NEXT: ret void
 }
 
-; CHECK-LABEL: @eh
+; CHECK-LABEL: @eh(
 define void @eh(i8* %hdl) personality i8* null {
 ; CHECK-NEXT: entry
 entry:

test/Transforms/Coroutines/coro-elide.ll

@@ -0,0 +1,110 @@
+; Tests that the coro.destroy and coro.resume are devirtualized where possible,
+; SCC pipeline restarts and inlines the direct calls.
+; RUN: opt < %s -S -inline -coro-elide | FileCheck %s
+
+declare void @print(i32) nounwind
+
+; resume part of the coroutine
+define fastcc void @f.resume(i8*) {
+  tail call void @print(i32 0)
+  ret void
+}
+
+; destroy part of the coroutine
+define fastcc void @f.destroy(i8*) {
+  tail call void @print(i32 1)
+  ret void
+}
+
+@f.resumers = internal constant [2 x void (i8*)*] [void (i8*)* @f.resume,
+                                                   void (i8*)* @f.destroy]
+
+; a coroutine start function
+define i8* @f() {
+entry:
+  %hdl = call i8* @llvm.coro.begin(i8* null, i32 0, i8* null,
+                          i8* bitcast ([2 x void (i8*)*]* @f.resumers to i8*))
+  ret i8* %hdl
+}
+
+; CHECK-LABEL: @callResume(
+define void @callResume() {
+entry:
+; CHECK: call i8* @llvm.coro.begin
+  %hdl = call i8* @f()
+
+; CHECK-NEXT: call void @print(i32 0)
+  %0 = call i8* @llvm.coro.subfn.addr(i8* %hdl, i8 0)
+  %1 = bitcast i8* %0 to void (i8*)*
+  call fastcc void %1(i8* %hdl)
+
+; CHECK-NEXT: call void @print(i32 1)
+  %2 = call i8* @llvm.coro.subfn.addr(i8* %hdl, i8 1)
+  %3 = bitcast i8* %2 to void (i8*)*
+  call fastcc void %3(i8* %hdl)
+
+; CHECK-NEXT: ret void
+  ret void
+}
+
+; CHECK-LABEL: @eh(
+define void @eh() personality i8* null {
+entry:
+; CHECK: call i8* @llvm.coro.begin
+  %hdl = call i8* @f()
+
+; CHECK-NEXT: call void @print(i32 0)
+  %0 = call i8* @llvm.coro.subfn.addr(i8* %hdl, i8 0)
+  %1 = bitcast i8* %0 to void (i8*)*
+  invoke void %1(i8* %hdl)
+          to label %cont unwind label %ehcleanup
+cont:
+  ret void
+
+ehcleanup:
+  %tok = cleanuppad within none []
+  cleanupret from %tok unwind to caller
+}
+
+; CHECK-LABEL: @no_devirt_info_null(
+; no devirtualization here, since coro.begin info parameter is null
+define void @no_devirt_info_null() {
+entry:
+  %hdl = call i8* @llvm.coro.begin(i8* null, i32 0, i8* null, i8* null)
+
+; CHECK: call i8* @llvm.coro.subfn.addr(i8* %hdl, i8 0)
+  %0 = call i8* @llvm.coro.subfn.addr(i8* %hdl, i8 0)
+  %1 = bitcast i8* %0 to void (i8*)*
+  call fastcc void %1(i8* %hdl)
+
+; CHECK: call i8* @llvm.coro.subfn.addr(i8* %hdl, i8 1)
+  %2 = call i8* @llvm.coro.subfn.addr(i8* %hdl, i8 1)
+  %3 = bitcast i8* %2 to void (i8*)*
+  call fastcc void %3(i8* %hdl)
+
+; CHECK: ret void
+  ret void
+}
+
+; CHECK-LABEL: @no_devirt_no_begin(
+; no devirtualization here, since coro.begin is not visible
+define void @no_devirt_no_begin(i8* %hdl) {
+entry:
+
+; CHECK: call i8* @llvm.coro.subfn.addr(i8* %hdl, i8 0)
+  %0 = call i8* @llvm.coro.subfn.addr(i8* %hdl, i8 0)
+  %1 = bitcast i8* %0 to void (i8*)*
+  call fastcc void %1(i8* %hdl)
+
+; CHECK: call i8* @llvm.coro.subfn.addr(i8* %hdl, i8 1)
+  %2 = call i8* @llvm.coro.subfn.addr(i8* %hdl, i8 1)
+  %3 = bitcast i8* %2 to void (i8*)*
+  call fastcc void %3(i8* %hdl)
+
+; CHECK: ret void
+  ret void
+}
+
+
+declare i8* @llvm.coro.begin(i8*, i32, i8*, i8*)
+declare i8* @llvm.coro.subfn.addr(i8*, i8)