Merge pull request #18947 from rjmccall/coroutine-optimizer-fixes

More SIL optimizer fixes for coroutines.

Author: rjmccall

lib/SIL/SILVerifier.cpp

@@ -4309,16 +4309,22 @@ class SILVerifier : public SILVerifierBase<SILVerifier> {
   void verifyEpilogBlocks(SILFunction *F) {
     bool FoundReturnBlock = false;
     bool FoundThrowBlock = false;
+    bool FoundUnwindBlock = false;
     for (auto &BB : *F) {
       if (isa<ReturnInst>(BB.getTerminator())) {
         require(!FoundReturnBlock,
                 "more than one return block in function");
         FoundReturnBlock = true;
-      }
-      if (isa<ThrowInst>(BB.getTerminator())) {
+      } else if (isa<ThrowInst>(BB.getTerminator())) {
         require(!FoundThrowBlock,
                 "more than one throw block in function");
         FoundThrowBlock = true;
+      } else if (isa<UnwindInst>(BB.getTerminator())) {
+        require(!FoundUnwindBlock,
+                "more than one unwind block in function");
+        FoundUnwindBlock = true;
+      } else {
+        assert(!BB.getTerminator()->isFunctionExiting());
       }
     }
   }

lib/SILOptimizer/Analysis/CFG.cpp

@@ -72,7 +72,7 @@ findAllNonFailureExitBBs(SILFunction *F,
       continue;
 
     // A return inst is always a non-failure exit bb.
-    if (isa<ReturnInst>(TI) || isa<ThrowInst>(TI)) {
+    if (TI->isFunctionExiting()) {
       BBs.push_back(&BB);
       continue;
     }

lib/SILOptimizer/IPO/ClosureSpecializer.cpp

@@ -787,18 +787,9 @@ void ClosureSpecCloner::populateCloned() {
       TermInst *TI = OpBB->getTerminator();
       auto Loc = CleanupLocation::get(NewClosure->getLoc());
 
-      // If we have a return/throw, we place the release right before it so we know
+      // If we have an exit, we place the release right before it so we know
       // that it will be executed at the end of the epilogue.
-      if (isa<ReturnInst>(TI)) {
-        Builder.setInsertionPoint(TI);
-        if (ClosureHasRefSemantics)
-          Builder.createReleaseValue(Loc, SILValue(NewClosure),
-                                     Builder.getDefaultAtomicity());
-        for (auto PAI : NeedsRelease)
-          Builder.createReleaseValue(Loc, SILValue(PAI),
-                                     Builder.getDefaultAtomicity());
-        continue;
-      } else if (isa<ThrowInst>(TI)) {
+      if (TI->isFunctionExiting()) {
         Builder.setInsertionPoint(TI);
         if (ClosureHasRefSemantics)
           Builder.createReleaseValue(Loc, SILValue(NewClosure),

lib/SILOptimizer/Transforms/ARCCodeMotion.cpp

@@ -306,7 +306,8 @@ class RetainCodeMotionContext : public CodeMotionContext {
     // end, this function is called many times.
     //
     // These terminator instructions block.
-    if (isa<ReturnInst>(II) || isa<ThrowInst>(II) || isa<UnreachableInst>(II))
+    if (isa<ReturnInst>(II) || isa<ThrowInst>(II) || isa<UnwindInst>(II) ||
+        isa<UnreachableInst>(II))
       return true;
     // Identical RC root blocks code motion, we will be able to move this retain
     // further once we move the blocking retain.

lib/SILOptimizer/Transforms/SimplifyCFG.cpp

@@ -262,12 +262,11 @@ static SILValue getTerminatorCondition(TermInst *Term) {
 /// Is this basic block jump threadable.
 static bool isThreadableBlock(SILBasicBlock *BB,
                               SmallPtrSetImpl<SILBasicBlock *> &LoopHeaders) {
-  if (isa<ReturnInst>(BB->getTerminator()))
-    return false;
+  auto TI = BB->getTerminator();
 
-  // We know how to handle cond_br and switch_enum .
-  if (!isa<CondBranchInst>(BB->getTerminator()) &&
-      !isa<SwitchEnumInst>(BB->getTerminator()))
+  // We know how to handle cond_br and switch_enum.
+  if (!isa<CondBranchInst>(TI) &&
+      !isa<SwitchEnumInst>(TI))
     return false;
 
   if (LoopHeaders.count(BB))
@@ -279,7 +278,7 @@ static bool isThreadableBlock(SILBasicBlock *BB,
       return false;
 
     // Don't jumpthread function calls.
-    if (isa<ApplyInst>(Inst))
+    if (FullApplySite::isa(&Inst))
       return false;
 
     // Only thread 'small blocks'.
@@ -975,7 +974,7 @@ bool SimplifyCFG::tryJumpThreading(BranchInst *BI) {
   auto *SrcBB = BI->getParent();
   // If the destination block ends with a return, we don't want to duplicate it.
   // We want to maintain the canonical form of a single return where possible.
-  if (isa<ReturnInst>(DestBB->getTerminator()))
+  if (DestBB->getTerminator()->isFunctionExiting())
     return false;
 
   // We need to update SSA if a value duplicated is used outside of the
@@ -2398,7 +2397,7 @@ static bool shouldTailDuplicate(SILBasicBlock &Block) {
   unsigned Cost = 0;
   bool SawRelease = false;
 
-  if (isa<ReturnInst>(Block.getTerminator()))
+  if (Block.getTerminator()->isFunctionExiting())
     return false;
 
   if (Block.getSinglePredecessorBlock())

lib/SILOptimizer/Utils/GenericCloner.cpp

@@ -159,7 +159,7 @@ void GenericCloner::populateCloned() {
         getBuilder().createReturn(RI->getLoc(), ReturnValue);
         continue;
       }
-    } else if (isa<ThrowInst>(OrigTermInst)) {
+    } else if (OrigTermInst->isFunctionExiting()) {
       for (AllocStackInst *ASI : reverse(AllocStacks)) {
         getBuilder().createDeallocStack(ASI->getLoc(), ASI);
       }

test/SIL/Parser/coroutines.sil

@@ -43,9 +43,12 @@ bb2:
   return %r : $()
 
 bb3:
-  unwind
+  br bb5
 
 bb4:
+  br bb5
+
+bb5:
   unwind
 }
 

test/SILOptimizer/specialize_reabstraction.sil

@@ -79,3 +79,55 @@ bb0(%0 : $Val<Bool>, %1 : $Val<Int>):
   %8 = apply %7(%1) : $@callee_owned (Val<Int>) -> Val<(Bool, Int)>
   return %8 : $Val<(Bool, Int)>
 }
+
+// CHECK-LABEL: sil shared @$S9coroutineSb_Tg5 : $@yield_once @convention(thin) (Bool) -> @yields @inout Bool {
+// CHECK:       bb0(%0 : $Bool):
+// CHECK-NEXT:    [[TEMP:%.*]] = alloc_stack $Bool
+// CHECK-NEXT:    store %0 to [[TEMP]] : $*Bool
+// CHECK-NEXT:    yield [[TEMP]] : $*Bool, resume bb1, unwind bb2
+// CHECK:       bb1:
+// CHECK-NEXT:    destroy_addr [[TEMP]] : $*Bool
+// CHECK-NEXT:    [[RV:%.*]] = tuple ()
+// CHECK-NEXT:    dealloc_stack [[TEMP]] : $*Bool
+// CHECK-NEXT:    return [[RV]] : $()
+// CHECK:       bb2:
+// CHECK-NEXT:    destroy_addr [[TEMP]] : $*Bool
+// CHECK-NEXT:    dealloc_stack [[TEMP]] : $*Bool
+// CHECK-NEXT:    unwind
+// CHECK-NEXT:  }
+sil @coroutine : $@yield_once @convention(thin) <T> (@in T) -> @yields @inout T {
+bb0(%0 : $*T):
+  yield %0 : $*T, resume bb1, unwind bb2
+bb1:
+  destroy_addr %0 : $*T
+  %rv = tuple ()
+  return %rv : $()
+bb2:
+  destroy_addr %0 : $*T
+  unwind
+}
+
+// CHECK-LABEL: @test_coroutine : $@convention(thin) (Bool) -> () {
+// CHECK:       bb0(%0 : $Bool):
+// CHECK-NEXT:    [[TEMP:%.*]] = alloc_stack $Bool
+// CHECK-NEXT:    store %0 to [[TEMP]] : $*Bool
+// CHECK-NEXT:    // function_ref
+// CHECK-NEXT:    [[CORO:%.*]] = function_ref @$S9coroutineSb_Tg5 : $@yield_once @convention(thin) (Bool) -> @yields @inout Bool
+// CHECK-NEXT:    [[LOAD:%.*]] = load [[TEMP]] : $*Bool
+// CHECK-NEXT:    ([[ADDR:%.*]], [[TOKEN:%.*]]) = begin_apply [[CORO]]([[LOAD]])
+// CHECK-NEXT:    end_apply [[TOKEN]]
+// CHECK-NEXT:    dealloc_stack [[TEMP]] : $*Bool
+// CHECK-NEXT:    [[RV:%.*]] = tuple ()
+// CHECK-NEXT:    return [[RV]] : $()
+// CHECK-NEXT:  }
+sil @test_coroutine : $@convention(thin) (Bool) -> () {
+bb0(%0 : $Bool):
+  %coro = function_ref @coroutine : $@yield_once @convention(thin) <T> (@in T) -> @yields @inout T
+  %temp = alloc_stack $Bool
+  store %0 to %temp : $*Bool
+  (%addr, %token) = begin_apply %coro<Bool>(%temp) : $@yield_once @convention(thin) <T> (@in T) -> @yields @inout T
+  end_apply %token
+  dealloc_stack %temp : $*Bool
+  %rv = tuple ()
+  return %rv : $()
+}