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

[GorNishanov]

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)

[majnemer]

lib/Transforms/Coroutines/CoroElide.cpp

67–75

This can do a lot of unnecessary work because there can be lots of users of the Constant.

Instead, could we use replaceAndRecursivelySimplify? It is quite powerful.

I imagine something like:

for (CoroSubFnInst *I : Users)
  replaceAndRecursivelySimplify(I, Value);

107–110

Any reason why you need the constant folder here? I'd have used ConstantExpr::getExtractValue.

lib/Transforms/Coroutines/CoroInstr.h

105

It seems like the natural thing to do would be to mirror the logic for OutlinedParts.

If you wish to cast<> here, we need a verifier check to ensure that the GlobalVariable's initializer is either a ConstantArray or a ConstantStruct.

[GorNishanov]

1. removed useless continue in CoroEarly
2. coro-elide.ll no longer runs coro-early pass
3. add coro.begin sanity tester to verifier.cpp
4. switch to using replaceAndRecursivelySimplify
5. switch to ConstantExpr::extractValue

[mehdi_amini]

lib/Transforms/Coroutines/CoroInstr.h

76

frontend I guess?

84

Constant array? Array of constants? I'm not sure what you wanted here but "array constant" doesn't seem right to me.

[GorNishanov]

Fixed typo + clarified comment in CoroInstr.h

[majnemer]

LGTM w/ nits addressed.

lib/Transforms/Coroutines/CoroInstr.h

96–98

Can this ever fail? I thought the verifier checks to make sure that we have a GlobalVariable operand.

102–103

Could we have a ConstantStruct test case as well?

test/Transforms/Coroutines/coro-elide.ll

31

I'd recommend ending all these function names with an open paren like so:

; CHECK-LABEL: @callResume(

This way we will not confuse FileCheck when we add @callResumeComplicated()

[GorNishanov]

CHECK-LABEL tweak per David's suggestion.

[GorNishanov]

lib/Transforms/Coroutines/CoroInstr.h

96–98

Yes. There are three legal alternative that are described in the comment above and checked by the verifier.

i8* null
ConstantStruct
ConstantArray

When coro.begin comes fresh out of the frontend, info parameter is null.

102–103

Sure. In one of the patches working on CoroSplit, there will be a test with ConstantStruct.

include/llvm/IR/Intrinsics.td

def int_coro_done : Intrinsic<[llvm_i1_ty], [llvm_ptr_ty],
                               [IntrArgMemOnly, ReadOnly<0>, NoCapture<0>]>;
 def int_coro_promise : Intrinsic<[llvm_ptr_ty],
                                  [llvm_ptr_ty, llvm_i32_ty, llvm_i1_ty],
                                  [IntrNoMem, NoCapture<0>]>;
 
 // 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 --------------------------===//
 //
 def int_flt_rounds : Intrinsic<[llvm_i32_ty]>,
                      GCCBuiltin<"__builtin_flt_rounds">;
 def int_trap : Intrinsic<[], [], [IntrNoReturn]>,
                GCCBuiltin<"__builtin_trap">;

lib/IR/Verifier.cpp

void Verifier::visitIntrinsicCallSite(Intrinsic::ID ID, CallSite CS) {
   // or are local to *this* function.
   for (Value *V : CS.args()) 
     if (auto *MD = dyn_cast<MetadataAsValue>(V))
       visitMetadataAsValue(*MD, CS.getCaller());
 
   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)),
            "is_zero_undef argument of bit counting intrinsics must be a "
            "constant int",
            CS);
     break;
   case Intrinsic::dbg_declare: // llvm.dbg.declare

lib/Transforms/Coroutines/CoroEarly.cpp

 
 namespace {
 // Created on demand if CoroEarly pass has work to do.
 class Lowerer : public coro::LowererBase {
   void lowerResumeOrDestroy(CallSite CS, CoroSubFnInst::ResumeKind);
 
 public:
   Lowerer(Module &M) : LowererBase(M) {}
-  static std::unique_ptr<Lowerer> createIfNeeded(Module &M);
   bool lowerEarlyIntrinsics(Function &F);
 };
 }
 
 // Replace a direct call to coro.resume or coro.destroy with an indirect call to
 // an address returned by coro.subfn.addr intrinsic. This is done so that
 // CGPassManager recognizes devirtualization when CoroElide pass replaces a call
 // to coro.subfn.addr with an appropriate function address.
if (auto CS = CallSite(&I)) {
       case Intrinsic::coro_resume:
         lowerResumeOrDestroy(CS, CoroSubFnInst::ResumeIndex);
         break;
       case Intrinsic::coro_destroy:
         lowerResumeOrDestroy(CS, CoroSubFnInst::DestroyIndex);
         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
 //===----------------------------------------------------------------------===//
 
 namespace {
 
 struct CoroEarly : public FunctionPass {
   static char ID; // Pass identification, replacement for typeid.
   CoroEarly() : FunctionPass(ID) {}
 
   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;
   }
 
   bool runOnFunction(Function &F) override {
     if (!L)
       return false;
 
     return L->lowerEarlyIntrinsics(F);
   }
 
   void getAnalysisUsage(AnalysisUsage &AU) const override {
     AU.setPreservesCFG();
   }
 };
-
 }
 
 char CoroEarly::ID = 0;
 INITIALIZE_PASS(CoroEarly, "coro-early", "Lower early coroutine intrinsics",
                 false, false)
 
 Pass *llvm::createCoroEarlyPass() { return new CoroEarly(); }

lib/Transforms/Coroutines/CoroElide.cpp

 //===- CoroElide.cpp - Coroutine Frame Allocation Elision Pass ------------===//
 //
 //                     The LLVM Compiler Infrastructure
 //
 // This file is distributed under the University of Illinois Open Source
 // License. See LICENSE.TXT for details.
 //
 //===----------------------------------------------------------------------===//
 // This pass replaces dynamic allocation of coroutine frame with alloca and
 // replaces calls to llvm.coro.resume and llvm.coro.destroy with direct calls
 // to coroutine sub-functions.
 //===----------------------------------------------------------------------===//
 
 #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;
 
 #define DEBUG_TYPE "coro-elide"
 
 //===----------------------------------------------------------------------===//
 //                              Top Level Driver
 //===----------------------------------------------------------------------===//
 
 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;
 INITIALIZE_PASS_BEGIN(
     CoroElide, "coro-elide",
     "Coroutine frame allocation elision and indirect calls replacement", false,
     false)
 INITIALIZE_PASS_END(
     CoroElide, "coro-elide",
     "Coroutine frame allocation elision and indirect calls replacement", false,
     false)
 
 Pass *llvm::createCoroElidePass() { return new CoroElide(); }

lib/Transforms/Coroutines/CoroInstr.h

public:
   static inline bool classof(const IntrinsicInst *I) {
     return I->getIntrinsicID() == Intrinsic::coro_subfn_addr;
   }
   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, 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 front: 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 constant
+  // refered 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

 #ifndef LLVM_LIB_TRANSFORMS_COROUTINES_COROINTERNAL_H
 #define LLVM_LIB_TRANSFORMS_COROUTINES_COROINTERNAL_H
 
 #include "CoroInstr.h"
 #include "llvm/Transforms/Coroutines.h"
 
 namespace llvm {
 
-class FunctionType;
-class LLVMContext;
-class Module;
 class PassRegistry;
 
 void initializeCoroEarlyPass(PassRegistry &);
 void initializeCoroSplitPass(PassRegistry &);
 void initializeCoroElidePass(PassRegistry &);
 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;
   LLVMContext &Context;
   FunctionType *const ResumeFnType;
 
   LowererBase(Module &M);
   Value *makeSubFnCall(Value *Arg, int Index, Instruction *InsertPt);
-  static bool declaresIntrinsics(Module &M, std::initializer_list<StringRef>);
 };
 
 } // End namespace coro.
 } // End namespace llvm
 
 #endif

lib/Transforms/Coroutines/Coroutines.cpp

auto *Bitcast =
       new BitCastInst(Call, ResumeFnType->getPointerTo(), "", InsertPt);
   return Bitcast;
 }
 
 #ifndef NDEBUG
 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.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;
 }
 #endif
 
 // 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");
     if (M.getNamedValue(Name))
       return true;
   }
 
   return false;
 }

test/Transforms/Coroutines/coro-elide.ll

+; 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
+}
+
[e-mail] = 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)