[Coroutines] Part 5: Add CGSCC restart trigger

[GorNishanov]

Summary

CoroSplit pass processes the coroutine twice. First, it lets it go through
complete IPO optimization pipeline as a single function. It forces restart
of the pipeline by inserting an indirect call to an empty function "coro.devirt.trigger"
which is devirtualized by CoroElide pass that triggers a restart of the pipeline by CGPassManager.
(In later patches, when CoroSplit pass sees the same coroutine the second time, it splits it up,
adds coroutine subfunctions to the SCC to be processed by IPO pipeline.)

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 (https://reviews.llvm.org/D23229)
5.Add CGSCC restart trigger + tests. <= we are here
6.Add coroutine heap elision + tests.
7.Add the rest of the logic (split into more patches)

[mehdi_amini]

This seems "hacky", but it is maybe because I don't grasp the complexity of what is achieved here.

[majnemer]

lib/Transforms/Coroutines/CoroSplit.cpp

23

Would you consider prepare instead of makeReady?

31–34

Any reason why you don't just use EntryBlock.getFirstInsertionPt()?

67

auto *?

70

Ditto.

74

I've updated this API to be a little more reasonable, just pass in Nodes.

106

Can this be auto *?

[GorNishanov]

In D23234#507839, @mehdi_amini wrote:

This seems "hacky", but it is maybe because I don't grasp the complexity of what is achieved here.

I considered three other less hacky alternatives to force a restart of CGSCC pipeline. They all had an unfortunate property of slowing down (slightly) compilation even if the module has no coroutines at all. I found the hacky alternative offered by this patch tolerable because it:

a) does not affect non-coroutines, thus, no impact on compilation of regular functions
b) we already have devirtualization logic for coro.resume and coro.destroy and we utilize it as a restart trigger

Less hacky alternatives:

Option 1: make CGPassManager a little bit of coroutine aware.
http://reviews.llvm.org/D21569
I dislike this approach as I’d like to keep Coroutine specific logic in coroutine passes

Option 2: Generalization runOnSCC(SCC, bool& restartRequired) – now any pass can request a restart
http://reviews.llvm.org/D21570
Not sure about this one either, it is very “in your face”. On the other hand, it is very straightforward way of specifying the desired behavior compare to:

Option 3: Add virtual bool CallGraphSCCPass::restartRequested() const; CGManager will call it after each call to runOnSCC to check if restart is required

http://reviews.llvm.org/D21572
Well, I am not sure about this one either. It is slightly more convoluted way to do option 2.
On the plus side, this is not a breaking change, and also, forcing people to work a little harder to demand restart, may give them a chance to think and realize that may be they don’t need to rerun the pipeline at all.

[GorNishanov]

Addressed code review feedback:

1. s/makeReadyForSplit/prepareForSplit
2. no longer skipping over DbgInst in the entry block
3. auto * galore
4. updated to match new interface to SCC.initialize

[GorNishanov]

lib/Transforms/Coroutines/CoroSplit.cpp

31–34

getFirstInsertionPt does not skip over AllocaInst. Some passes, Reg2Mem, Inliner, for example, assume that entry block starts with (possibly empty) sequence of AllocaInst.

I don't have to skip over DbgInfoIntrinsic though. I can make this loop to be simply:

while (isa<AllocaInst>(*It)) ++It;

[mehdi_amini]

I'm still struggling to figure why you need to force the restart, can you point to a document that explains it?
(I think it should be explained in the code somehow as well).

[GorNishanov]

In D23234#507893, @mehdi_amini wrote:

I'm still struggling to figure why you need to force the restart, can you point to a document that explains it?
(I think it should be explained in the code somehow as well).

Would this comment address your concerns? I can add it to the beginning of CoroSplit.cpp

// We present a coroutine to an LLVM as an ordinary function with suspension
// points marked up with intrinsics. We let the optimizer party on the coroutine
// as a single function for as long as possible. Shortly before the coroutine is
// eligible to be inlined into its callers, we split up the coroutine into parts
// corresponding to an initial, resume and destroy invocations of the coroutine,
// adds them to the current SCC and restart the IPO pipeline to optimize the
// coroutine subfunctions we extracted before proceeding to the caller of the
// coroutine.

P.S.

This was discussed in RFC: Coroutine Optimization Passes: http://lists.llvm.org/pipermail/llvm-dev/2016-July/102337.html .

[GorNishanov]

Added a comment to CoroSplit.

[GorNishanov]

uploaded with more context (-U999999)

[GorNishanov]

use terminator instruction of the entry block as an insertion point for the restart trigger function call.

[majnemer]

LGTM w/ nits addressed.

lib/Transforms/Coroutines/CoroElide.cpp

16–17

Could we remove this include?

lib/Transforms/Coroutines/CoroInternal.h

36

I'd name this define CORO_PRESPLIT_ATTR.

37–38

Maybe just PREPARED_FOR_SPLIT and UNPREPARED_FOR_SPLIT ?

lib/Transforms/Coroutines/CoroSplit.cpp

15

Are you using this? I cannot quite see where...

[GorNishanov]

Nits addressed

include/llvm/IR/Function.h

setAttributes(
                                  AttributeSet::FunctionIndex, Kind));
   }
   void addFnAttr(StringRef Kind, StringRef Value) {
     setAttributes(
       AttributeSets.addAttribute(getContext(),
                                  AttributeSet::FunctionIndex, Kind, Value));
   }
 
+  /// @brief Remove function attribute from this function.
+  void removeFnAttr(StringRef Kind) {
+    setAttributes(AttributeSets.removeAttribute(
+        getContext(), AttributeSet::FunctionIndex, Kind));
+  }
+
   /// Set the entry count for this function.
   void setEntryCount(uint64_t Count);
 
   /// Get the entry count for this function.
   Optional<uint64_t> getEntryCount() const;
 
   /// @brief Return true if the function has the attribute.
   bool hasFnAttribute(Attribute::AttrKind Kind) const {

lib/Transforms/Coroutines/CoroEarly.cpp

 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_begin:
+        // 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())
+            F.addFnAttr(CORO_PRESPLIT_ATTR, UNPREPARED_FOR_SPLIT);
+        }
+        break;
       case Intrinsic::coro_resume:
         lowerResumeOrDestroy(CS, CoroSubFnInst::ResumeIndex);
         break;
       case Intrinsic::coro_destroy:
         lowerResumeOrDestroy(CS, CoroSubFnInst::DestroyIndex);
         break;
       }
       Changed = true;
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 {
-    if (coro::declaresIntrinsics(M, {"llvm.coro.resume", "llvm.coro.destroy"}))
+    if (coro::declaresIntrinsics(
+            M, {"llvm.coro.begin", "llvm.coro.resume", "llvm.coro.destroy"}))
       L = llvm::make_unique<Lowerer>(M);
     return false;
   }
 
   bool runOnFunction(Function &F) override {
     if (!L)
       return false;
 

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
static bool replaceIndirectCalls(CoroBeginInst *CoroBegin) {
   auto *DestroyAddrConstant =
       ConstantExpr::getExtractValue(Resumers, CoroSubFnInst::DestroyIndex);
 
   replaceWithConstant(ResumeAddrConstant, ResumeAddr);
   replaceWithConstant(DestroyAddrConstant, DestroyAddr);
   return true;
 }
 
+// See if there are any coro.subfn.addr instructions referring to coro.devirt
+// trigger, if so, replace them with a direct call to devirt trigger function.
+static bool replaceDevirtTrigger(Function &F) {
+  SmallVector<CoroSubFnInst *, 1> DevirtAddr;
+  for (auto &I : instructions(F))
+    if (auto *SubFn = dyn_cast<CoroSubFnInst>(&I))
+      if (SubFn->getIndex() == CoroSubFnInst::RestartTrigger)
+        DevirtAddr.push_back(SubFn);
+
+  if (DevirtAddr.empty())
+    return false;
+
+  Module &M = *F.getParent();
+  Function *DevirtFn = M.getFunction(CORO_DEVIRT_TRIGGER_FN);
+  assert(DevirtFn && "coro.devirt.fn not found");
+  replaceWithConstant(DevirtFn, DevirtAddr);
+
+  return true;
+}
+
 bool CoroElide::runOnFunction(Function &F) {
+  bool Changed = false;
+
+  if (F.hasFnAttribute(CORO_PRESPLIT_ATTR))
+    Changed = replaceDevirtTrigger(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;
+    return Changed;
 
   for (auto *CB : CoroBegins)
     Changed |= replaceIndirectCalls(CB);
 
   return Changed;
 }
 
 char CoroElide::ID = 0;

lib/Transforms/Coroutines/CoroInstr.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 {
+    RestartTrigger = -1,
     ResumeIndex,
     DestroyIndex,
     IndexLast,
-    IndexFirst = ResumeIndex
+    IndexFirst = RestartTrigger
   };
 
   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);
public:
   // 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; }
+    bool isPreSplit() const { return !isPostSplit(); }
   };
   Info getInfo() const {
     Info Result;
     auto *GV = dyn_cast<GlobalVariable>(getRawInfo());
     if (!GV)
       return Result;
 
     assert(GV->isConstant() && GV->hasDefinitiveInitializer());

lib/Transforms/Coroutines/CoroInternal.h

 
 class PassRegistry;
 
 void initializeCoroEarlyPass(PassRegistry &);
 void initializeCoroSplitPass(PassRegistry &);
 void initializeCoroElidePass(PassRegistry &);
 void initializeCoroCleanupPass(PassRegistry &);
 
+// CoroEarly pass marks every function that has coro.begin with a string
+// attribute "coroutine.presplit"="0". CoroSplit pass processes the coroutine
+// twice. First, it lets it go through complete IPO optimization pipeline as a
+// single function. It forces restart of the pipeline by inserting an indirect
+// call to an empty function "coro.devirt.trigger" which is devirtualized by
+// CoroElide pass that triggers a restart of the pipeline by CGPassManager.
+// When CoroSplit pass sees the same coroutine the second time, it splits it up,
+// adds coroutine subfunctions to the SCC to be processed by IPO pipeline.
+
+#define CORO_PRESPLIT_ATTR "coroutine.presplit"
+#define UNPREPARED_FOR_SPLIT "0"
+#define PREPARED_FOR_SPLIT "1"
+
+#define CORO_DEVIRT_TRIGGER_FN "coro.devirt.trigger"
+
 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;

lib/Transforms/Coroutines/CoroSplit.cpp

 //===- CoroSplit.cpp - Converts a coroutine into a state machine ----------===//
 //
 //                     The LLVM Compiler Infrastructure
 //
 // This file is distributed under the University of Illinois Open Source
 // License. See LICENSE.TXT for details.
 //
 //===----------------------------------------------------------------------===//
 // This pass builds the coroutine frame and outlines resume and destroy parts
 // of the coroutine into separate functions.
 //===----------------------------------------------------------------------===//
 
 #include "CoroInternal.h"
 #include "llvm/Analysis/CallGraphSCCPass.h"
 
 using namespace llvm;
 
 #define DEBUG_TYPE "coro-split"
 
+// We present a coroutine to an LLVM as an ordinary function with suspension
+// points marked up with intrinsics. We let the optimizer party on the coroutine
+// as a single function for as long as possible. Shortly before the coroutine is
+// eligible to be inlined into its callers, we split up the coroutine into parts
+// corresponding to an initial, resume and destroy invocations of the coroutine,
+// add them to the current SCC and restart the IPO pipeline to optimize the
+// coroutine subfunctions we extracted before proceeding to the caller of the
+// coroutine.
+
+// When we see the coroutine the first time, we insert an indirect call to a
+// devirt trigger function and mark the coroutine that it is now ready for
+// split.
+static void prepareForSplit(Function &F, CallGraph &CG) {
+  Module &M = *F.getParent();
+  Function *DevirtFn = M.getFunction(CORO_DEVIRT_TRIGGER_FN);
+  assert(DevirtFn && "coro.devirt.trigger function not found");
+
+  F.addFnAttr(CORO_PRESPLIT_ATTR, PREPARED_FOR_SPLIT);
+
+  // Insert an indirect call sequence that will be devirtualized by CoroElide
+  // pass:
+  //    %0 = call i8* @llvm.coro.subfn.addr(i8* null, i8 -1)
+  //    %1 = bitcast i8* %0 to void(i8*)*
+  //    call void %1(i8* null)
+  coro::LowererBase Lowerer(M);
+  Instruction *InsertPt = F.getEntryBlock().getTerminator();
+  auto *Null = ConstantPointerNull::get(Type::getInt8PtrTy(F.getContext()));
+  auto *DevirtFnAddr =
+      Lowerer.makeSubFnCall(Null, CoroSubFnInst::RestartTrigger, InsertPt);
+  auto *IndirectCall = CallInst::Create(DevirtFnAddr, Null, "", InsertPt);
+
+  // Update CG graph with an indirect call we just added.
+  CG[&F]->addCalledFunction(IndirectCall, CG.getCallsExternalNode());
+}
+
+// Make sure that there is a devirtualization trigger function that CoroSplit
+// pass uses the force restart CGSCC pipeline. If devirt trigger function is not
+// found, we will create one and add it to the current SCC.
+static void createDevirtTriggerFunc(CallGraph &CG, CallGraphSCC &SCC) {
+  Module &M = CG.getModule();
+  if (M.getFunction(CORO_DEVIRT_TRIGGER_FN))
+    return;
+
+  LLVMContext &C = M.getContext();
+  auto *FnTy = FunctionType::get(Type::getVoidTy(C), Type::getInt8PtrTy(C),
+                                /*IsVarArgs=*/false);
+  Function *DevirtFn =
+      Function::Create(FnTy, GlobalValue::LinkageTypes::PrivateLinkage,
+                       CORO_DEVIRT_TRIGGER_FN, &M);
+  DevirtFn->addFnAttr(Attribute::AlwaysInline);
+  auto *Entry = BasicBlock::Create(C, "entry", DevirtFn);
+  ReturnInst::Create(C, Entry);
+
+  auto *Node = CG.getOrInsertFunction(DevirtFn);
+
+  SmallVector<CallGraphNode *, 8> Nodes(SCC.begin(), SCC.end());
+  Nodes.push_back(Node);
+  SCC.initialize(Nodes);
+}
+
 //===----------------------------------------------------------------------===//
 //                              Top Level Driver
 //===----------------------------------------------------------------------===//
 
 namespace {
 
 struct CoroSplit : public CallGraphSCCPass {
   static char ID; // Pass identification, replacement for typeid
   CoroSplit() : CallGraphSCCPass(ID) {}
 
-  bool runOnSCC(CallGraphSCC &SCC) override { return false; }
+  bool Run = false;
+
+  // A coroutine is identified by the presence of coro.begin intrinsic, if
+  // we don't have any, this pass has nothing to do.
+  bool doInitialization(CallGraph &CG) override {
+    Run = coro::declaresIntrinsics(CG.getModule(), {"llvm.coro.begin"});
+    return CallGraphSCCPass::doInitialization(CG);
+  }
+
+  bool runOnSCC(CallGraphSCC &SCC) override {
+    if (!Run)
+      return false;
+
+    // Find coroutines for processing.
+    SmallVector<Function *, 4> Coroutines;
+    for (CallGraphNode *CGN : SCC)
+      if (auto *F = CGN->getFunction())
+        if (F->hasFnAttribute(CORO_PRESPLIT_ATTR))
+          Coroutines.push_back(F);
+
+    if (Coroutines.empty())
+      return false;
+
+    CallGraph &CG = getAnalysis<CallGraphWrapperPass>().getCallGraph();
+    createDevirtTriggerFunc(CG, SCC);
+
+    for (Function *F : Coroutines) {
+      Attribute Attr = F->getFnAttribute(CORO_PRESPLIT_ATTR);
+      StringRef Value = Attr.getValueAsString();
+      DEBUG(dbgs() << "CoroSplit: Processing coroutine '" << F->getName()
+                   << "' state: " << Value << "\n");
+      if (Value == UNPREPARED_FOR_SPLIT) {
+        prepareForSplit(*F, CG);
+        continue;
+      }
+      F->removeFnAttr(CORO_PRESPLIT_ATTR);
+    }
+    return true;
+  }
+
   void getAnalysisUsage(AnalysisUsage &AU) const override {
     AU.setPreservesAll();
     CallGraphSCCPass::getAnalysisUsage(AU);
   }
 };
-
 }
 
 char CoroSplit::ID = 0;
 INITIALIZE_PASS(
     CoroSplit, "coro-split",
     "Split coroutine into a set of functions driving its state machine", false,
     false)
 
 Pass *llvm::createCoroSplitPass() { return new CoroSplit(); }

test/Transforms/Coroutines/restart-trigger.ll

+; Verifies that restart trigger forces IPO pipelines restart and the same
+; coroutine is looked at by CoroSplit pass twice.
+; RUN: opt < %s -S -O0 -enable-coroutines -debug-only=coro-split 2>&1 | FileCheck %s
+; RUN: opt < %s -S -O1 -enable-coroutines -debug-only=coro-split 2>&1 | FileCheck %s
+
+; CHECK:      CoroSplit: Processing coroutine 'f' state: 0
+; CHECK-NEXT: CoroSplit: Processing coroutine 'f' state: 1
+
+declare i8* @llvm.coro.begin(i8*, i32, i8*, i8*)
+
+; a coroutine start function
+define i8* @f() {
+entry:
+  %hdl = call i8* @llvm.coro.begin(i8* null, i32 0, i8* null, i8* null)
+  ret i8* %hdl
+}