[Coroutine] Collect CoroBegin if all of terminators are dominated by one coro.destroy

[ChuanqiXu]

Summary

The original logic seems to be we could collecting a CoroBegin if one of the terminators could be dominated by one of coro.destroy, which doesn't make sense.
This patch rewrites the logics to collect CoroBegin if all of terminators are dominated by one coro.destroy. If there is no such coro.destroy, we would call hasEscapePath to evaluate if we should collect it.

Test Plan: check-llvm

[lxfind]

Does it require that all terminators to be dominated by the same coro.destroy?
Reading the comments in the original code, I think the condition is that for each terminator it's dominated by a coro.destroy, but it doesn't require all terminators to be dominated by the same coro.destroy.
The original code does seem to be buggy, but the fix may not be correct. Perhaps in the next statement, the check on hasEscapePath and ReferencedCoroBegins.count is problematic?

[ChuanqiXu]

In D100614#2697468, @lxfind wrote:

Does it require that all terminators to be dominated by the same coro.destroy?
Reading the comments in the original code, I think the condition is that for each terminator it's dominated by a coro.destroy, but it doesn't require all terminators to be dominated by the same coro.destroy.
The original code does seem to be buggy, but the fix may not be correct. Perhaps in the next statement, the check on hasEscapePath and ReferencedCoroBegins.count is problematic?

The codes I changed means if all terminators are dominated by one coro.destroy we can collect the corresponding coro.begin directly. Then if there is no coro.destroy to dominate all terminators, we could only search in hasEscapePath. The key idea here is hasEscapePath is relatively slow, so we want to avoid to call it as much as possible. The check on ReferencedCoroBegins.count and hasEscapePath just did this.

[ChuanqiXu]

In another words, we could use hasEscapePath to handle all the cases if we don't care about compile-time.

[lxfind]

I see. LGTM. Thanks for fixing. Could you also add comments to the llvm::all_of search with what you just described?

[ychen]

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

97

Missing a CHECK?

[ChuanqiXu]

Address the comments.

[ChuanqiXu]

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

97

Thanks for reminding. It is legacy comments before previously. Here I use print(i32 2) to judge this coroutine is elided.

llvm/lib/Transforms/Coroutines/CoroElide.cpp

for (BasicBlock &B : *F) {
       if (TI->getNumSuccessors() == 0 && !TI->isExceptionalTerminator() &&
           !isa<UnreachableInst>(TI))
         Terminators.insert(&B);
     }
 
   // Filter out the coro.destroy that lie along exceptional paths.
   SmallPtrSet<CoroBeginInst *, 8> ReferencedCoroBegins;
   for (auto &It : DestroyAddr) {
+    // If there is any coro.destroy dominates all of the terminators for the
+    // coro.begin, we could know the corresponding coro.begin wouldn't escape.
     for (Instruction *DA : It.second) {
-      for (BasicBlock *TI : Terminators) {
-        if (DT.dominates(DA, TI->getTerminator())) {
-          ReferencedCoroBegins.insert(It.first);
-          break;
-        }
-      }
-    }
+      if (llvm::all_of(Terminators, [&](auto *TI) {
+            return DT.dominates(DA, TI->getTerminator());
+          })) {
 
     // Whether there is any paths from coro.begin to Terminators which not pass
     // through any of the coro.destroys.
+    //
+    // hasEscapePath is relatively slow, so we avoid to run it as much as
+    // possible.
     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.

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

 }
 
 ; 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]
+; cleanup part of the coroutine
+define fastcc void @f.cleanup(i8*) {
+  tail call void @print(i32 2)
+  ret void
+}
+
[e-mail] = internal constant [3 x void (i8*)*] [void (i8*)* @f.resume,
+                                                   void (i8*)* @f.destroy,
+                                                   void (i8*)* @f.cleanup]
 
 ; a coroutine start function
 define i8* @f() {
 entry:
   %id = call token @llvm.coro.id(i32 0, i8* null,
                           i8* bitcast (i8*()* @f to i8*),
-                          i8* bitcast ([2 x void (i8*)*]* @f.resumers to i8*))
+                          i8* bitcast ([3 x void (i8*)*]* @f.resumers to i8*))
+  %alloc = call i1 @llvm.coro.alloc(token %id)
   %hdl = call i8* @llvm.coro.begin(token %id, i8* null)
   ret i8* %hdl
 }
 
 ; CHECK-LABEL: @callResume(
 define void @callResume() {
 entry:
   %hdl = call i8* @f()
 
 ; CHECK: 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)
+; CHECK-NEXT: call void @print(i32 2)
   %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: @callResumeMultiRet(
+define void @callResumeMultiRet(i1 %b) {
+entry:
+  %hdl = call i8* @f()
+; CHECK: %alloc.i = call i1 @llvm.coro.alloc
+; CHECK: 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)
+  br i1 %b, label %destroy, label %ret
+
+destroy:
+; CHECK: 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)
+  ret void
+
+ret:
+  ret void
+}
+
+; CHECK-LABEL: @callResumeMultiRetDommmed(
+define void @callResumeMultiRetDommmed(i1 %b) {
+entry:
+  %hdl = call i8* @f()
+; CHECK-NOT: %alloc.i = call i1 @llvm.coro.alloc
+; CHECK: 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: call void @print(i32 2)
+  %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 i1 %b, label %destroy, label %ret
+
+destroy:
+  ret void
+
+ret:
+  ret void
+}
+
 ; CHECK-LABEL: @eh(
 define void @eh() personality i8* null {
 entry:
   %hdl = call i8* @f()
 
 ; CHECK: call void @print(i32 0)
   %0 = call i8* @llvm.coro.subfn.addr(i8* %hdl, i8 0)
   %1 = bitcast i8* %0 to void (i8*)*

 ; CHECK: ret void
   ret void
 }
 
 declare token @llvm.coro.id(i32, i8*, i8*, i8*)
 declare i8* @llvm.coro.begin(token, i8*)
 declare i8* @llvm.coro.frame()
 declare i8* @llvm.coro.subfn.addr(i8*, i8)
+declare i1 @llvm.coro.alloc(token)