[Coroutines] Optimized coroutine elision based on reachability

[junparser]

Summary

For now, CoroElide pass can not handle case like:

coroA (cond){
   if (cond)
     co_await coroB
  else 
     co_await coroC
co_return
}

However, CoroElide should worked for this.

In this patch, If dominate relation is not satisfied, we iterate all paths from coro.begin to normal terminators to check whether there is any pathes from coro.begin to Terminators which not pass through any of the coro.destroys. If not, record the coro.begin.
As for the paths reach to normal terminators through suspend point, they have been skipped during DFS.

TestPlan: check-llvm, check-clang, cppcoro

[modocache]

Thank you for this patch, I'll give it a proper review tonight. But in the meantime, I wonder if you've tried this with any of the sample programs for this bug report, https://bugs.llvm.org/show_bug.cgi?id=40656 ?

[junparser]

In D75440#1901867, @modocache wrote:

Thank you for this patch, I'll give it a proper review tonight. But in the meantime, I wonder if you've tried this with any of the sample programs for this bug report, https://bugs.llvm.org/show_bug.cgi?id=40656 ?

yep, I have tried the two samples, same behavior as clang6.

[modocache]

yep, I have tried the two samples, same behavior as clang6.

Phenomenal!! Thanks for checking. Please feel free to address the nit-pick comments only if you wish. I'll wait to accept until I better understand the rationale for checking the number of cases in a suspend switch statement, if that's alright.

llvm/lib/Transforms/Coroutines/CoroElide.cpp

33

A nit-pick, but: it's spelled "switches" with an "e". So I would rename this CoroSuspendSwitches, or if that's too long of a variable name for you, CoroSwitches. Also, this could be marked const to signal these won't be "lowered" or transformed in any way: SmallPtrSet<const SwitchInst *, 4>.

147

Another nit-pick: this can be const CoroBeginInst * (with a few modifications below, like SmallVector<const BasicBlock *, 32> Worklist;). I think it'd be a small improvement because it'd make clear that this function just analyzes paths through the IR, it doesn't modify it.

172

Yet another nit-pick, sorry: maybe it's just me but I would find this much clearer with --Limit; as its own statement on its own line. Just my opinion, though, please ignore it if you disagree.

259

Typo: "donot" should be "do not" -- but also, this is the first instance in which these coroutine passes check for the number of cases in a switch statement. Could you expand on this comment? I'm afraid I don't follow why coro.suspend with switches between resume/cleanup should be part of the path analysis, but others should not.

[junparser]

In D75440#1904205, @modocache wrote:

yep, I have tried the two samples, same behavior as clang6.

Phenomenal!! Thanks for checking. Please feel free to address the nit-pick comments only if you wish. I'll wait to accept until I better understand the rationale for checking the number of cases in a suspend switch statement, if that's alright.

Thanks for review the patch. I'll update the patch as suggestion.

llvm/lib/Transforms/Coroutines/CoroElide.cpp

259

According to https://llvm.org/docs/Coroutines.html, when FE create coro.suspend with switches, it should be :

%0 = call i8 @llvm.coro.suspend(token none, i1 false)
switch i8 %0, label %suspend [i8 0, label %resume
                              i8 1, label %cleanup]

The behavior of default dest is suspend the coroutine and return to caller which means this a escape path to normal terminator.

so for case like:

coroA  () {
     co_await coroB();
}

statement co_await coroB() is expanded by FE, and there is at least one coro.suspend between coro.begin and coro.destroy of coroB.

Skip the default dest of coro.suspend switches is reasonable since contents of coroutine frame doesn't change outside the coroutine body.

Btw, most of time the switches can not be simplified except final.suspend which we do not care (real terminator right?).

[modocache]

llvm/lib/Transforms/Coroutines/CoroElide.cpp

259

Ah, I see! Very cool, thanks for the explanation, that helped me understand. (I think the comment could use an update based on what you wrote above.)

[junparser]

update the patch as suggestion

llvm/lib/Transforms/Coroutines/CoroElide.cpp

 // Created on demand if the coro-elide pass has work to do.
 struct Lowerer : coro::LowererBase {
   SmallVector<CoroIdInst *, 4> CoroIds;
   SmallVector<CoroBeginInst *, 1> CoroBegins;
   SmallVector<CoroAllocInst *, 1> CoroAllocs;
   SmallVector<CoroSubFnInst *, 4> ResumeAddr;
   DenseMap<CoroBeginInst *, SmallVector<CoroSubFnInst *, 4>> DestroyAddr;
   SmallVector<CoroFreeInst *, 1> CoroFrees;
-  CoroSuspendInst *CoroFinalSuspend;
+  SmallPtrSet<const SwitchInst *, 4> CoroSuspendSwitches;
 
   Lowerer(Module &M) : LowererBase(M) {}
 
   void elideHeapAllocations(Function *F, Type *FrameTy, AAResults &AA);
   bool shouldElide(Function *F, DominatorTree &DT) const;
   void collectPostSplitCoroIds(Function *F);
   bool processCoroId(CoroIdInst *, AAResults &AA, DominatorTree &DT);
+  bool hasEscapePath(const CoroBeginInst *,
+                     const SmallPtrSetImpl<BasicBlock *> &) const;
 };
 } // end anonymous namespace
 
 // 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())
for (auto *CB : CoroBegins) {
     CB->eraseFromParent();
   }
 
   // Since now coroutine frame lives on the stack we need to make sure that
   // any tail call referencing it, must be made non-tail call.
   removeTailCallAttribute(Frame, AA);
 }
 
+bool Lowerer::hasEscapePath(const CoroBeginInst *CB,
+                            const SmallPtrSetImpl<BasicBlock *> &TIs) const {
+  const auto &It = DestroyAddr.find(CB);
+  assert(It != DestroyAddr.end());
+
+  // Limit the number of blocks we visit.
+  unsigned Limit = 32 * (1 + It->second.size());
+
+  SmallVector<const BasicBlock *, 32> Worklist;
+  Worklist.push_back(CB->getParent());
+
+  SmallPtrSet<const BasicBlock *, 32> Visited;
+  // Consider basicblock of coro.destroy as visited one, so that we
+  // skip the path pass through coro.destroy.
+  for (auto *DA : It->second)
+    Visited.insert(DA->getParent());
+
+  do {
+    const auto *BB = Worklist.pop_back_val();
+    if (!Visited.insert(BB).second)
+      continue;
+    if (TIs.count(BB))
+      return true;
+
+    // Conservatively say that there is potentially a path.
+    if (!--Limit)
+      return true;
+
+    auto TI = BB->getTerminator();
+    // Although the default dest of coro.suspend switches is suspend pointer
+    // which means a escape path to normal terminator, it is reasonable to skip
+    // it since coroutine frame doesn't change outside the coroutine body.
+    if (isa<SwitchInst>(TI) &&
+        CoroSuspendSwitches.count(cast<SwitchInst>(TI))) {
+      Worklist.push_back(cast<SwitchInst>(TI)->getSuccessor(1));
+      Worklist.push_back(cast<SwitchInst>(TI)->getSuccessor(2));
+    } else
+      Worklist.append(succ_begin(BB), succ_end(BB));
+
+  } while (!Worklist.empty());
+
+  // We have exhausted all possible paths and are certain that coro.begin can
+  // not reach to any of terminators.
+  return false;
+}
+
 bool Lowerer::shouldElide(Function *F, DominatorTree &DT) const {
   // If no CoroAllocs, we cannot suppress allocation, so elision is not
   // possible.
   if (CoroAllocs.empty())
     return false;
 
   // Check that for every coro.begin there is at least one coro.destroy directly
   // referencing the SSA value of that coro.begin along each
   // non-exceptional path.
   // If the value escaped, then coro.destroy would have been referencing a
   // memory location storing that value and not the virtual register.
 
-  SmallPtrSet<Instruction *, 8> Terminators;
-  bool HasMultiPred = false;
+  SmallPtrSet<BasicBlock *, 8> Terminators;
   // First gather all of the non-exceptional terminators for the function.
   // Consider the final coro.suspend as the real terminator when the current
   // function is a coroutine.
-  if (CoroFinalSuspend) {
-    // If block of final coro.suspend has more than one predecessor,
-    // then there is one resume path and the others are exceptional paths,
-    // consider these predecessors as terminators.
-    BasicBlock *FinalBB = CoroFinalSuspend->getParent();
-    if (FinalBB->hasNPredecessorsOrMore(2)) {
-      HasMultiPred = true;
-      for (auto *B : predecessors(FinalBB))
-        Terminators.insert(B->getTerminator());
-    } else
-      Terminators.insert(CoroFinalSuspend);
-  } else {
     for (BasicBlock &B : *F) {
       auto *TI = B.getTerminator();
       if (TI->getNumSuccessors() == 0 && !TI->isExceptionalTerminator() &&
           !isa<UnreachableInst>(TI))
-        Terminators.insert(TI);
-    }
+        Terminators.insert(&B);
     }
 
   // Filter out the coro.destroy that lie along exceptional paths.
-  SmallPtrSet<CoroSubFnInst *, 4> DAs;
-  SmallPtrSet<Instruction *, 2> TIs;
   SmallPtrSet<CoroBeginInst *, 8> ReferencedCoroBegins;
   for (auto &It : DestroyAddr) {
-    for (CoroSubFnInst *DA : It.second) {
-      for (Instruction *TI : Terminators) {
-        if (DT.dominates(DA, TI)) {
-          if (HasMultiPred)
-            TIs.insert(TI);
-          else
-            DAs.insert(DA);
+    for (Instruction *DA : It.second) {
+      for (BasicBlock *TI : Terminators) {
+        if (DT.dominates(DA, TI->getTerminator())) {
+          ReferencedCoroBegins.insert(It.first);
           break;
         }
       }
     }
-    // If all the predecessors dominate coro.destroys that reference same
-    // coro.begin, record the coro.begin
-    if (TIs.size() == Terminators.size()) {
-      ReferencedCoroBegins.insert(It.first);
-      TIs.clear();
-    }
-  }
 
-  // Find all the coro.begin referenced by coro.destroy along happy paths.
-  for (CoroSubFnInst *DA : DAs) {
-    if (auto *CB = dyn_cast<CoroBeginInst>(DA->getFrame()))
-      ReferencedCoroBegins.insert(CB);
-    else
-      return false;
+    // Whether there is any paths from coro.begin to Terminators which not pass
+    // through any of the coro.destroys.
+    if (!ReferencedCoroBegins.count(It.first) &&
+        !hasEscapePath(It.first, Terminators))
+      ReferencedCoroBegins.insert(It.first);
   }
 
   // If size of the set is the same as total number of coro.begin, that means we
   // found a coro.free or coro.destroy referencing each coro.begin, so we can
   // perform heap elision.
   return ReferencedCoroBegins.size() == CoroBegins.size();
 }
 
 void Lowerer::collectPostSplitCoroIds(Function *F) {
   CoroIds.clear();
-  CoroFinalSuspend = nullptr;
+  CoroSuspendSwitches.clear();
   for (auto &I : instructions(F)) {
     if (auto *CII = dyn_cast<CoroIdInst>(&I))
       if (CII->getInfo().isPostSplit())
         // If it is the coroutine itself, don't touch it.
         if (CII->getCoroutine() != CII->getFunction())
           CoroIds.push_back(CII);
 
+    // Consider case like:
+    // %0 = call i8 @llvm.coro.suspend(...)
+    // switch i8 %0, label %suspend [i8 0, label %resume
+    //                              i8 1, label %cleanup]
+    // and collect the SwitchInsts which are used by escape analysis later.
     if (auto *CSI = dyn_cast<CoroSuspendInst>(&I))
-      if (CSI->isFinal()) {
-        if (!CoroFinalSuspend)
-          CoroFinalSuspend = CSI;
-        else
-          report_fatal_error("Only one suspend point can be marked as final");
+      if (CSI->hasOneUse() && isa<SwitchInst>(CSI->use_begin()->getUser())) {
+        SwitchInst *SWI = cast<SwitchInst>(CSI->use_begin()->getUser());
+        if (SWI->getNumCases() == 2)
+          CoroSuspendSwitches.insert(SWI);
       }
   }
 }
 
 bool Lowerer::processCoroId(CoroIdInst *CoroId, AAResults &AA,
                             DominatorTree &DT) {
   CoroBegins.clear();
   CoroAllocs.clear();

llvm/test/Transforms/Coroutines/coro-heap-elide.ll

; CHECK-NEXT: call fastcc void bitcast (void (%f.frame*)* @f.cleanup to void (i8*)*)(i8* %vFrame)
   br label %coro.ret
 
 ; CHECK-LABEL: coro.ret:
 coro.ret:
 ; CHECK-NEXT: ret void
   ret void
 }
 
+; CHECK-LABEL: @callResume_with_coro_suspend_3(
+define void @callResume_with_coro_suspend_3(i8 %cond) {
+entry:
+; CHECK: alloca %f.frame
+  switch i8 %cond, label  %coro.ret [
+    i8 0, label %init.suspend
+    i8 1, label %coro.ret
+  ]
+
+init.suspend:
+; CHECK-NOT: llvm.coro.begin
+; CHECK-NOT: CustomAlloc
+; CHECK: call void @may_throw()
+  %hdl = call i8* @f()
+; CHECK-NEXT: call fastcc void bitcast (void (%f.frame*)* @f.resume to void (i8*)*)(i8* %vFrame)
+  %0 = call i8* @llvm.coro.subfn.addr(i8* %hdl, i8 0)
+  %1 = bitcast i8* %0 to void (i8*)*
+  call fastcc void %1(i8* %hdl)
+  %2 = call token @llvm.coro.save(i8* %hdl)
+  %3 = call i8 @llvm.coro.suspend(token %2, i1 false)
+  switch i8 %3, label  %coro.ret [
+    i8 0, label %final.suspend
+    i8 1, label %cleanups
+  ]
+
+; CHECK-LABEL: final.suspend:
+final.suspend:
+; CHECK-NEXT: call fastcc void bitcast (void (%f.frame*)* @f.cleanup to void (i8*)*)(i8* %vFrame)
+  %4 = call i8* @llvm.coro.subfn.addr(i8* %hdl, i8 1)
+  %5 = bitcast i8* %4 to void (i8*)*
+  call fastcc void %5(i8* %hdl)
+  %6 = call token @llvm.coro.save(i8* %hdl)
+  %7 = call i8 @llvm.coro.suspend(token %6, i1 true)
+  switch i8 %7, label  %coro.ret [
+    i8 0, label %coro.ret
+    i8 1, label %cleanups
+  ]
+
+; CHECK-LABEL: cleanups:
+cleanups:
+; CHECK-NEXT: call fastcc void bitcast (void (%f.frame*)* @f.cleanup to void (i8*)*)(i8* %vFrame)
+  %8 = call i8* @llvm.coro.subfn.addr(i8* %hdl, i8 1)
+  %9 = bitcast i8* %8 to void (i8*)*
+  call fastcc void %9(i8* %hdl)
+  br label %coro.ret
+
+; CHECK-LABEL: coro.ret:
+coro.ret:
+; CHECK-NEXT: ret void
+  ret void
+}
+
 
 
 ; CHECK-LABEL: @callResume_PR34897_no_elision(
 define void @callResume_PR34897_no_elision(i1 %cond) {
 ; CHECK-LABEL: entry:
 entry:
 ; CHECK: call i8* @CustomAlloc(
   %hdl = call i8* @f()
if.else:
   br label %return
 
 ; CHECK-LABEL: return:
 return:
 ; CHECK: ret void
   ret void
 }
 
+; CHECK-LABEL: @callResume_PR34897_elision(
+define void @callResume_PR34897_elision(i1 %cond) {
+; CHECK-LABEL: entry:
+entry:
+; CHECK: alloca %f.frame
+; CHECK: tail call void @bar(
+  tail call void @bar(i8* null)
+  br i1 %cond, label %if.then, label %if.else
+
+if.then:
+; CHECK-NOT: CustomAlloc
+; CHECK: call void @may_throw()
+  %hdl = call i8* @f()
+; CHECK: call void @bar(
+  tail call void @bar(i8* %hdl)
+; CHECK: call fastcc void bitcast (void (%f.frame*)* @f.resume to void (i8*)*)(i8*
+  %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 fastcc void bitcast (void (%f.frame*)* @f.cleanup to void (i8*)*)(i8*
+  %2 = call i8* @llvm.coro.subfn.addr(i8* %hdl, i8 1)
+  %3 = bitcast i8* %2 to void (i8*)*
+  call fastcc void %3(i8* %hdl)
+  br label %return
+
+if.else:
+  br label %return
+
+; CHECK-LABEL: return:
+return:
+; CHECK: ret void
+  ret void
+}
+
+
 ; a coroutine start function (cannot elide heap alloc, due to second argument to
 ; coro.begin not pointint to coro.alloc)
 define i8* @f_no_elision() personality i8* null {
 entry:
   %id = call token @llvm.coro.id(i32 0, i8* null,
                       i8* bitcast (i8*()* @f_no_elision to i8*),
                       i8* bitcast ([3 x void (%f.frame*)*]* @f.resumers to i8*))
   %alloc = call i8* @CustomAlloc(i32 4)