[coroutines] Part 4[ab]: Coroutine Devirtualization: Lower coro.resume and coro.destroy.

This is the forth patch in the coroutine series. CoroEaly pass now lowers coro.resume
and coro.destroy intrinsics by replacing them with an indirect call to an address
returned by coro.subfn.addr intrinsic. This is done so that CGPassManager recognizes
devirtualization when CoroElide replaces a call to coro.subfn.addr with an appropriate
function address.

Patch by Gor Nishanov!

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

git-svn-id: https://llvm.org/svn/llvm-project/llvm/trunk@277765 91177308-0d34-0410-b5e6-96231b3b80d8

include/llvm/IR/CallSite.h

@@ -109,6 +109,17 @@ public:
     *getCallee() = V;
   }
 
+  /// Return the intrinsic ID of the intrinsic called by this CallSite,
+  /// or Intrinsic::not_intrinsic if the called function is not an
+  /// intrinsic, or if this CallSite is an indirect call.
+  Intrinsic::ID getIntrinsicID() const {
+    if (auto *F = getCalledFunction())
+      return F->getIntrinsicID();
+    // Don't use Intrinsic::not_intrinsic, as it will require pulling
+    // Intrinsics.h into every header that uses CallSite.
+    return static_cast<Intrinsic::ID>(0);
+  }
+
   /// isCallee - Determine whether the passed iterator points to the
   /// callee operand's Use.
   bool isCallee(Value::const_user_iterator UI) const {

lib/IR/Verifier.cpp

@@ -3681,11 +3681,13 @@ void Verifier::visitInstruction(Instruction &I) {
       Assert(
           !F->isIntrinsic() || isa<CallInst>(I) ||
               F->getIntrinsicID() == Intrinsic::donothing ||
+              F->getIntrinsicID() == Intrinsic::coro_resume ||
+              F->getIntrinsicID() == Intrinsic::coro_destroy ||
               F->getIntrinsicID() == Intrinsic::experimental_patchpoint_void ||
               F->getIntrinsicID() == Intrinsic::experimental_patchpoint_i64 ||
               F->getIntrinsicID() == Intrinsic::experimental_gc_statepoint,
-          "Cannot invoke an intrinsic other than donothing, patchpoint or "
-          "statepoint",
+          "Cannot invoke an intrinsic other than donothing, patchpoint, "
+          "statepoint, coro_resume or coro_destroy",
           &I);
       Assert(F->getParent() == &M, "Referencing function in another module!",
              &I, &M, F, F->getParent());

lib/Transforms/Coroutines/CoroEarly.cpp

@@ -12,12 +12,70 @@
 //===----------------------------------------------------------------------===//
 
 #include "CoroInternal.h"
+#include "llvm/IR/CallSite.h"
+#include "llvm/IR/InstIterator.h"
+#include "llvm/IR/Module.h"
 #include "llvm/Pass.h"
 
 using namespace llvm;
 
 #define DEBUG_TYPE "coro-early"
 
+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.
+void Lowerer::lowerResumeOrDestroy(CallSite CS,
+                                   CoroSubFnInst::ResumeKind Index) {
+  Value *ResumeAddr =
+      makeSubFnCall(CS.getArgOperand(0), Index, CS.getInstruction());
+  CS.setCalledFunction(ResumeAddr);
+  CS.setCallingConv(CallingConv::Fast);
+}
+
+bool Lowerer::lowerEarlyIntrinsics(Function &F) {
+  bool Changed = false;
+  for (auto IB = inst_begin(F), IE = inst_end(F); IB != IE;) {
+    Instruction &I = *IB++;
+    if (auto CS = CallSite(&I)) {
+      switch (CS.getIntrinsicID()) {
+      default:
+        continue;
+      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
 //===----------------------------------------------------------------------===//
@@ -25,12 +83,25 @@ using namespace llvm;
 namespace {
 
 struct CoroEarly : public FunctionPass {
-  static char ID; // Pass identification, replacement for typeid
+  static char ID; // Pass identification, replacement for typeid.
   CoroEarly() : FunctionPass(ID) {}
 
-  bool runOnFunction(Function &F) override { return false; }
+  std::unique_ptr<Lowerer> L;
+
+  bool doInitialization(Module &M) override {
+    L = Lowerer::createIfNeeded(M);
+    return false;
+  }
+
+  bool runOnFunction(Function &F) override {
+    if (!L)
+      return false;
+
+    return L->lowerEarlyIntrinsics(F);
+  }
+
   void getAnalysisUsage(AnalysisUsage &AU) const override {
-    AU.setPreservesAll();
+    AU.setPreservesCFG();
   }
 };
 

lib/Transforms/Coroutines/CoroInstr.h

@@ -0,0 +1,64 @@
+//===-- CoroInstr.h - Coroutine Intrinsics Instruction Wrappers -*- C++ -*-===//
+//
+//                     The LLVM Compiler Infrastructure
+//
+// This file is distributed under the University of Illinois Open Source
+// License. See LICENSE.TXT for details.
+//
+//===----------------------------------------------------------------------===//
+// This file defines classes that make it really easy to deal with intrinsic
+// functions with the isa/dyncast family of functions.  In particular, this
+// allows you to do things like:
+//
+//     if (auto *SF = dyn_cast<CoroSubFnInst>(Inst))
+//        ... SF->getFrame() ... SF->getAlloc() ...
+//
+// 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
+// has state or virtual methods, which is an important part of this gross/neat
+// hack working.
+//
+// The helpful comment above is borrowed from llvm/IntrinsicInst.h, we keep
+// coroutine intrinsic wrappers here since they are only used by the passes in
+// the Coroutine library.
+//===----------------------------------------------------------------------===//
+
+#include "llvm/IR/GlobalVariable.h"
+#include "llvm/IR/IntrinsicInst.h"
+
+namespace llvm {
+
+/// This class represents the llvm.coro.subfn.addr instruction.
+class LLVM_LIBRARY_VISIBILITY CoroSubFnInst : public IntrinsicInst {
+  enum { FrameArg, IndexArg };
+
+public:
+  enum ResumeKind {
+    ResumeIndex,
+    DestroyIndex,
+    IndexLast,
+    IndexFirst = ResumeIndex
+  };
+
+  Value *getFrame() const { return getArgOperand(FrameArg); }
+  ResumeKind getIndex() const {
+    int64_t Index = getRawIndex()->getValue().getSExtValue();
+    assert(Index >= IndexFirst && Index < IndexLast &&
+           "unexpected CoroSubFnInst index argument");
+    return static_cast<ResumeKind>(Index);
+  }
+
+  ConstantInt *getRawIndex() const {
+    return cast<ConstantInt>(getArgOperand(IndexArg));
+  }
+
+  // Methods to support type inquiry through isa, cast, and dyn_cast:
+  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));
+  }
+};
+
+} // End namespace llvm.

lib/Transforms/Coroutines/CoroInternal.h

@@ -12,10 +12,14 @@
 #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 &);
@@ -23,6 +27,20 @@ void initializeCoroSplitPass(PassRegistry &);
 void initializeCoroElidePass(PassRegistry &);
 void initializeCoroCleanupPass(PassRegistry &);
 
-}
+namespace coro {
+
+// 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

@@ -66,3 +66,67 @@ void llvm::addCoroutinePassesToExtensionPoints(PassManagerBuilder &Builder) {
   Builder.addExtension(PassManagerBuilder::EP_OptimizerLast,
                        addCoroutineOptimizerLastPasses);
 }
+
+// Construct the lowerer base class and initialize its members.
+coro::LowererBase::LowererBase(Module &M)
+    : TheModule(M), Context(M.getContext()),
+      ResumeFnType(FunctionType::get(Type::getVoidTy(Context),
+                                     Type::getInt8PtrTy(Context),
+                                     /*isVarArg=*/false)) {}
+
+// Creates a sequence of instructions to obtain a resume function address using
+// llvm.coro.subfn.addr. It generates the following sequence:
+//
+//    call i8* @llvm.coro.subfn.addr(i8* %Arg, i8 %index)
+//    bitcast i8* %2 to void(i8*)*
+
+Value *coro::LowererBase::makeSubFnCall(Value *Arg, int Index,
+                                        Instruction *InsertPt) {
+  auto *IndexVal = ConstantInt::get(Type::getInt8Ty(Context), Index);
+  auto *Fn = Intrinsic::getDeclaration(&TheModule, Intrinsic::coro_subfn_addr);
+
+  assert(Index >= CoroSubFnInst::IndexFirst &&
+         Index < CoroSubFnInst::IndexLast &&
+         "makeSubFnCall: Index value out of range");
+  auto *Call = CallInst::Create(Fn, {Arg, IndexVal}, "", InsertPt);
+
+  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.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) {
+
+  for (StringRef Name : List) {
+    assert(isCoroutineIntrinsicName(Name) && "not a coroutine intrinsic");
+    if (M.getNamedValue(Name))
+      return true;
+  }
+
+  return false;
+}

test/Transforms/Coroutines/coro-early.ll

@@ -0,0 +1,41 @@
+; Tests that CoroEarly pass correctly lowers coro.resume and coro.destroy
+; intrinsics.
+; RUN: opt < %s -S -coro-early | FileCheck %s
+
+; CHECK-LABEL: @callResume
+define void @callResume(i8* %hdl) {
+; CHECK-NEXT: entry
+entry:
+; CHECK-NEXT: %0 = call i8* @llvm.coro.subfn.addr(i8* %hdl, i8 0)
+; CHECK-NEXT: %1 = bitcast i8* %0 to void (i8*)*
+; CHECK-NEXT: call fastcc void %1(i8* %hdl)
+  call void @llvm.coro.resume(i8* %hdl)
+
+; CHECK-NEXT: %2 = call i8* @llvm.coro.subfn.addr(i8* %hdl, i8 1)
+; CHECK-NEXT: %3 = bitcast i8* %2 to void (i8*)*
+; CHECK-NEXT: call fastcc void %3(i8* %hdl)
+  call void @llvm.coro.destroy(i8* %hdl)
+
+  ret void
+; CHECK-NEXT: ret void
+}
+
+; CHECK-LABEL: @eh
+define void @eh(i8* %hdl) personality i8* null {
+; CHECK-NEXT: entry
+entry:
+;  CHECK-NEXT: %0 = call i8* @llvm.coro.subfn.addr(i8* %hdl, i8 0)
+;  CHECK-NEXT: %1 = bitcast i8* %0 to void (i8*)*
+;  CHECK-NEXT: invoke fastcc void %1(i8* %hdl)
+  invoke void @llvm.coro.resume(i8* %hdl)
+          to label %cont unwind label %ehcleanup
+cont:
+  ret void
+
+ehcleanup:
+  %0 = cleanuppad within none []
+  cleanupret from %0 unwind to caller
+}
+
+declare void @llvm.coro.resume(i8*)
+declare void @llvm.coro.destroy(i8*)

test/Verifier/invoke.ll

@@ -46,7 +46,7 @@ contb:
 
 define i8 @f2() personality i8* bitcast (i32 (...)* @__gxx_personality_v0 to i8*) {
 entry:
-; CHECK: Cannot invoke an intrinsic other than donothing, patchpoint or statepoint
+; CHECK: Cannot invoke an intrinsic other than donothing, patchpoint, statepoint, coro_resume or coro_destroy
   invoke void @llvm.trap()
   to label %cont unwind label %lpad