Merge pull request #61126 from meg-gupta/fixdce

Don't shrink escaped borrow scopes in DCE

Author: meg-gupta

include/swift/SIL/OwnershipUtils.h

@@ -78,6 +78,8 @@ inline bool isForwardingConsume(SILValue value) {
   return canOpcodeForwardOwnedValues(value);
 }
 
+bool hasEscaped(BorrowedValue value);
+
 /// Find leaf "use points" of \p guaranteedValue that determine its lifetime
 /// requirement. Return true if no PointerEscape use was found.
 ///

lib/SIL/Utils/OwnershipUtils.cpp

@@ -25,6 +25,62 @@
 
 using namespace swift;
 
+bool swift::hasEscaped(BorrowedValue value) {
+  GraphNodeWorklist<Operand *, 8> worklist;
+  for (Operand *use : value->getUses()) {
+    if (use->getOperandOwnership() != OperandOwnership::NonUse)
+      worklist.insert(use);
+  }
+
+  while (Operand *op = worklist.pop()) {
+    switch (op->getOperandOwnership()) {
+    case OperandOwnership::NonUse:
+    case OperandOwnership::TrivialUse:
+    case OperandOwnership::ForwardingConsume:
+    case OperandOwnership::DestroyingConsume:
+      llvm_unreachable("this operand cannot handle an inner guaranteed use");
+
+    case OperandOwnership::ForwardingUnowned:
+    case OperandOwnership::PointerEscape:
+    case OperandOwnership::BitwiseEscape:
+      return true;
+
+    case OperandOwnership::Borrow:
+    case OperandOwnership::EndBorrow:
+    case OperandOwnership::InstantaneousUse:
+    case OperandOwnership::UnownedInstantaneousUse:
+    case OperandOwnership::InteriorPointer:
+      break;
+
+    case OperandOwnership::Reborrow: {
+      SILArgument *phi = cast<BranchInst>(op->getUser())
+                             ->getDestBB()
+                             ->getArgument(op->getOperandNumber());
+      for (auto *use : phi->getUses()) {
+        if (use->getOperandOwnership() != OperandOwnership::NonUse)
+          worklist.insert(use);
+      }
+      break;
+    }
+    case OperandOwnership::ForwardingBorrow: {
+      ForwardingOperand(op).visitForwardedValues([&](SILValue result) {
+        // Do not include transitive uses with 'none' ownership
+        if (result->getOwnershipKind() == OwnershipKind::None)
+          return true;
+        for (auto *resultUse : result->getUses()) {
+          if (resultUse->getOperandOwnership() != OperandOwnership::NonUse) {
+            worklist.insert(resultUse);
+          }
+        }
+        return true;
+      });
+      break;
+    }
+    }
+  }
+  return false;
+}
+
 bool swift::isValueAddressOrTrivial(SILValue v) {
   return v->getType().isAddress() ||
          v->getOwnershipKind() == OwnershipKind::None;

lib/SILOptimizer/Transforms/DeadCodeElimination.cpp

@@ -301,8 +301,15 @@ void DCE::markLive() {
           }
           continue;
         }
-        // If not populate reborrowDependencies for this borrow
+        // Populate reborrowDependencies for this borrow
         findReborrowDependencies(borrowInst);
+        if (hasEscaped(BorrowedValue(borrowInst))) {
+          // Visit all end_borrows and mark them live
+          visitTransitiveEndBorrows(borrowInst, [&](EndBorrowInst *endBorrow) {
+            markInstructionLive(endBorrow);
+          });
+          continue;
+        }
         break;
       }
       default:

test/SILOptimizer/dead_code_elimination_nontrivial_ossa.sil

@@ -972,3 +972,26 @@ bb0(%0 : $@sil_unmanaged AnyObject):
   %242 = tuple ()
   return %242 : $()
 }
+
+sil @foo : $@convention(thin) () -> ()
+
+// CHECK-LABEL: sil [ossa] @test_escape : {{.*}} {
+// CHECK: br bb1
+// CHECK: end_borrow
+// CHECK-LABEL: } // end sil function 'test_escape'
+sil [ossa] @test_escape : $@convention(thin) () -> () {
+bb0:
+  %13 = function_ref @foo : $@convention(thin) () -> ()
+  %14 = partial_apply [callee_guaranteed] %13() : $@convention(thin) () -> ()
+  %15 = convert_escape_to_noescape %14 : $@callee_guaranteed () -> () to $@noescape @callee_guaranteed () -> ()
+  %17 = begin_borrow %14 : $@callee_guaranteed () -> ()
+  %18 = mark_dependence %15 : $@noescape @callee_guaranteed () -> () on %17 : $@callee_guaranteed () -> ()
+  br bb1(%18 : $@noescape @callee_guaranteed () -> (), %17 : $@callee_guaranteed () -> (), %14 : $@callee_guaranteed () -> ())
+
+bb1(%21 : $@noescape @callee_guaranteed () -> (), %22 : @guaranteed $@callee_guaranteed () -> (), %23 : @owned $@callee_guaranteed () -> ()):
+  %27 = apply %21() : $@noescape @callee_guaranteed () -> ()
+  end_borrow %22 : $@callee_guaranteed () -> ()
+  destroy_value %23 : $@callee_guaranteed () -> ()
+  %28 = tuple ()
+  return %28: $()
+}