[Coroutines] CoroElide enhancement

[junparser]

Summary

According to https://reviews.llvm.org/D43242, coroelide pass only considers llvm.coro.destroy that dominate non-exceptional return. However, this cannot be met if current function is coroutine function because of suspend point.
Consider case from http://www.open-std.org/jtc1/sc22/wg21/docs/papers/2018/p0978r0.html :

task<> big_task() { // lifetime of `subtask` is fully enclosed in its caller
  co_await subtask(); // coroutine frame allocation elided 
}

such cases should be optimized.

this patch does two changes in shouldElide function:

1. Consider final coro.suspend as real non-exceptional terminator when function is coroutine, since final suspend is executed with or without exception at last.
2. When function body is wrapped by try-catch, block of coro.suspend may has multiple predecessors : one normal path and the others are exceptional path, and they do not dominate each other. If all the predecessors dominate coro.destroys that reference same coro.begin, we consider condition is met.
3. Keep normal function as usual

Test Plan:
Compile https://godbolt.org/g/mCKfnr and confirm it is still optimized to a single instruction, 'return 1190'.
Compile and run https://godbolt.org/z/AbdRw4, there is no "new" and "delete" output
Compile and run https://godbolt.org/z/gjYSi5 , there is no "new" and "delete" output
check-llvm

[junparser]

kindly ping~

[junparser]

@modocache kindly ping~

[junparser]

In D71663#1798419, @junparser wrote:

@modocache kindly ping~

@modocache any comments？

[junparser]

gental ping～

[modocache]

Sorry this review took so long, and thank you for addressing the regression I introduced in D43242!

This patch no longer cleanly applies due to my changes in D71900, but I believe I was able to rebase it correctly in P8200. Based on that I think this looks like a great improvement, thank you!

Do you have commit access now, or would you like me to commit this (once you've updated it to apply cleanly) on your behalf?

llvm/lib/Transforms/Coroutines/CoroElide.cpp

161

typo: "coroutine". Also, nit-pick comment, but: I would write this "Consider the final coro.suspend as the real terminator when the current function is a coroutine." But I might be misunderstanding what you mean by "real terminator" and "coroutine expanded function" here.

[junparser]

In D71663#1888271, @modocache wrote:

Sorry this review took so long, and thank you for addressing the regression I introduced in D43242!

This patch no longer cleanly applies due to my changes in D71900, but I believe I was able to rebase it correctly in P8200. Based on that I think this looks like a great improvement, thank you!

Do you have commit access now, or would you like me to commit this (once you've updated it to apply cleanly) on your behalf?

Yes , thanks for review the patch.
And I have got the commit access, however, how do I push the patch? Any hint for first time checkin ? Thank you again.

[junparser]

Rebase and update the patch as suggestion.

[junparser]

In D71663#1893326, @junparser wrote:

In D71663#1888271, @modocache wrote:

Sorry this review took so long, and thank you for addressing the regression I introduced in D43242!

This patch no longer cleanly applies due to my changes in D71900, but I believe I was able to rebase it correctly in P8200. Based on that I think this looks like a great improvement, thank you!

Do you have commit access now, or would you like me to commit this (once you've updated it to apply cleanly) on your behalf?

Yes , thanks for review the patch.
And I have got the commit access, however, how do I push the patch? Any hint for first time checkin ? Thank you again.

@modocache Thanks, I have committed the patch.

llvm/lib/Transforms/Coroutines/CoroElide.cpp

 //===- CoroElide.cpp - Coroutine Frame Allocation Elision Pass ------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #include "llvm/Transforms/Coroutines/CoroElide.h"
 #include "CoroInternal.h"
+#include "llvm/ADT/DenseMap.h"
 #include "llvm/Analysis/AliasAnalysis.h"
 #include "llvm/Analysis/InstructionSimplify.h"
 #include "llvm/IR/Dominators.h"
 #include "llvm/IR/InstIterator.h"
 #include "llvm/InitializePasses.h"
 #include "llvm/Pass.h"
 #include "llvm/Support/ErrorHandling.h"
 
 using namespace llvm;
 
 #define DEBUG_TYPE "coro-elide"
 
 namespace {
 // 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;
-  SmallVector<CoroSubFnInst *, 4> DestroyAddr;
+  DenseMap<CoroBeginInst *, SmallVector<CoroSubFnInst *, 4>> DestroyAddr;
   SmallVector<CoroFreeInst *, 1> CoroFrees;
+  CoroSuspendInst *CoroFinalSuspend;
 
   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 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 a coro.destroy directly
-  // referencing the SSA value of that coro.begin along a non-exceptional path.
+  // 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.
 
-  // First gather all of the non-exceptional terminators for the function.
   SmallPtrSet<Instruction *, 8> Terminators;
+  bool HasMultiPred = false;
+  // 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);
     }
+  }
 
   // Filter out the coro.destroy that lie along exceptional paths.
   SmallPtrSet<CoroSubFnInst *, 4> DAs;
-  for (CoroSubFnInst *DA : DestroyAddr) {
+  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);
           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.
-  SmallPtrSet<CoroBeginInst *, 8> ReferencedCoroBegins;
   for (CoroSubFnInst *DA : DAs) {
     if (auto *CB = dyn_cast<CoroBeginInst>(DA->getFrame()))
       ReferencedCoroBegins.insert(CB);
     else
       return false;
   }
 
   // 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();
-  for (auto &I : instructions(F))
+  CoroFinalSuspend = nullptr;
+  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);
+
+    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");
+      }
+  }
 }
 
 bool Lowerer::processCoroId(CoroIdInst *CoroId, AAResults &AA,
                             DominatorTree &DT) {
   CoroBegins.clear();
   CoroAllocs.clear();
   CoroFrees.clear();
   ResumeAddr.clear();
bool Lowerer::processCoroId(CoroIdInst *CoroId, AAResults &AA,
   for (CoroBeginInst *CB : CoroBegins) {
     for (User *U : CB->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);
+          DestroyAddr[CB].push_back(II);
           break;
         default:
           llvm_unreachable("unexpected coro.subfn.addr constant");
         }
   }
 
   // PostSplit coro.id refers to an array of subfunctions in its Info
   // argument.
   ConstantArray *Resumers = CoroId->getInfo().Resumers;
   assert(Resumers && "PostSplit coro.id Info argument must refer to an array"
                      "of coroutine subfunctions");
   auto *ResumeAddrConstant =
       ConstantExpr::getExtractValue(Resumers, CoroSubFnInst::ResumeIndex);
 
   replaceWithConstant(ResumeAddrConstant, ResumeAddr);
 
   bool ShouldElide = shouldElide(CoroId->getFunction(), DT);
 
   auto *DestroyAddrConstant = ConstantExpr::getExtractValue(
       Resumers,
       ShouldElide ? CoroSubFnInst::CleanupIndex : CoroSubFnInst::DestroyIndex);
 
-  replaceWithConstant(DestroyAddrConstant, DestroyAddr);
+  for (auto &It : DestroyAddr)
+    replaceWithConstant(DestroyAddrConstant, It.second);
 
   if (ShouldElide) {
     auto *FrameTy = getFrameType(cast<Function>(ResumeAddrConstant));
     elideHeapAllocations(CoroId->getFunction(), FrameTy, AA);
     coro::replaceCoroFree(CoroId, /*Elide=*/true);
   }
 
   return true;

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

; CHECK-NEXT: call fastcc void bitcast (void (%f.frame*)* @f.cleanup to void (i8*)*)(i8* %vFrame)
   %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: @callResume_with_coro_suspend_1(
+define void @callResume_with_coro_suspend_1() {
+entry:
+; CHECK: alloca %f.frame
+; CHECK-NOT: 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_with_coro_suspend_2(
+define void @callResume_with_coro_suspend_2() personality i8* null {
+entry:
+; CHECK: alloca %f.frame
+; CHECK-NOT: coro.begin
+; CHECK-NOT: CustomAlloc
+; CHECK: call void @may_throw()
+  %hdl = call i8* @f()
+
+  %0 = call token @llvm.coro.save(i8* %hdl)
+; CHECK: invoke fastcc void bitcast (void (%f.frame*)* @f.resume to void (i8*)*)(i8* %vFrame)
+  %1 = call i8* @llvm.coro.subfn.addr(i8* %hdl, i8 0)
+  %2 = bitcast i8* %1 to void (i8*)*
+  invoke fastcc void %2(i8* %hdl)
+    to label %invoke.cont1 unwind label %lpad
+
+; CHECK-LABEL: invoke.cont1:
+invoke.cont1:
+  %3 = call i8 @llvm.coro.suspend(token %0, i1 false)
+  switch i8 %3, label  %coro.ret [
+    i8 0, label %final.ready
+    i8 1, label %cleanups
+  ]
+
+; CHECK-LABEL: lpad:
+lpad:
+  %4 = landingpad { i8*, i32 }
+          catch i8* null
+; CHECK: call fastcc void bitcast (void (%f.frame*)* @f.cleanup to void (i8*)*)(i8* %vFrame)
+  %5 = call i8* @llvm.coro.subfn.addr(i8* %hdl, i8 1)
+  %6 = bitcast i8* %5 to void (i8*)*
+  call fastcc void %6(i8* %hdl)
+  br label %final.suspend
+
+; CHECK-LABEL: final.ready:
+final.ready:
+; CHECK-NEXT: call fastcc void bitcast (void (%f.frame*)* @f.cleanup to void (i8*)*)(i8* %vFrame)
+  %7 = call i8* @llvm.coro.subfn.addr(i8* %hdl, i8 1)
+  %8 = bitcast i8* %7 to void (i8*)*
+  call fastcc void %8(i8* %hdl)
+  br label %final.suspend
+
+; CHECK-LABEL: final.suspend:
+final.suspend:
+  %9 = call token @llvm.coro.save(i8* %hdl)
+  %10 = call i8 @llvm.coro.suspend(token %9, i1 true)
+  switch i8 %10, 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)
+  %11 = call i8* @llvm.coro.subfn.addr(i8* %hdl, i8 1)
+  %12 = bitcast i8* %11 to void (i8*)*
+  call fastcc void %12(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()
 ; CHECK: tail call void @bar(
   tail call void @bar(i8* %hdl)

 }
 
 declare token @llvm.coro.id(i32, i8*, i8*, i8*)
 declare i1 @llvm.coro.alloc(token)
 declare i8* @llvm.coro.free(token, i8*)
 declare i8* @llvm.coro.begin(token, i8*)
 declare i8* @llvm.coro.frame(token)
 declare i8* @llvm.coro.subfn.addr(i8*, i8)
+declare i8 @llvm.coro.suspend(token, i1)
+declare token @llvm.coro.save(i8*)