Merge pull request #60121 from nate-chandler/rdar94346482

[SILOpt] Handle reborrows of owned phis in CanonicalizeOSSALifetime.

Author: nate-chandler

include/swift/SIL/OwnershipUtils.h

@@ -1217,6 +1217,14 @@ void findTransitiveReborrowBaseValuePairs(
     BorrowingOperand initialScopeOperand, SILValue origBaseValue,
     function_ref<void(SILPhiArgument *, SILValue)> visitReborrowBaseValuePair);
 
+/// Visit the phis in the same block as \p phi which are reborrows of a borrow
+/// of one of the values reaching \p phi.
+///
+/// If the visitor returns false, stops visiting and returns false.  Otherwise,
+/// returns true.
+bool visitAdjacentReborrowsOfPhi(SILPhiArgument *phi,
+                                 function_ref<bool(SILPhiArgument *)> visitor);
+
 /// Given a begin of a borrow scope, visit all end_borrow users of the borrow or
 /// its reborrows.
 void visitTransitiveEndBorrows(

include/swift/SIL/SILArgument.h

@@ -274,6 +274,13 @@ class SILPhiArgument : public SILArgument {
   /// If visitor returns false, iteration is stopped and we return false.
   bool visitIncomingPhiOperands(function_ref<bool(Operand *)> visitor) const;
 
+  /// Visit incoming phi operands and the argument into which they are incoming;
+  /// if an operand's value is itself a phi, visit that phi's operands.
+  ///
+  /// Returns false when called on a non-phi and when the visitor returns false.
+  bool visitTransitiveIncomingPhiOperands(
+      function_ref<bool(SILPhiArgument *, Operand *)> visitor);
+
   /// Returns true if we were able to find a single terminator operand value for
   /// each predecessor of this arguments basic block. The found values are
   /// stored in OutArray.

lib/SIL/IR/SILArgument.cpp

@@ -10,12 +10,13 @@
 //
 //===----------------------------------------------------------------------===//
 
-#include "llvm/ADT/STLExtras.h"
-#include "swift/SIL/SILBasicBlock.h"
 #include "swift/SIL/SILArgument.h"
+#include "swift/Basic/GraphNodeWorklist.h"
+#include "swift/SIL/SILBasicBlock.h"
 #include "swift/SIL/SILFunction.h"
 #include "swift/SIL/SILInstruction.h"
 #include "swift/SIL/SILModule.h"
+#include "llvm/ADT/STLExtras.h"
 
 using namespace swift;
 
@@ -226,6 +227,31 @@ bool SILPhiArgument::getIncomingPhiValues(
   return true;
 }
 
+bool SILPhiArgument::visitTransitiveIncomingPhiOperands(
+    function_ref<bool(SILPhiArgument *, Operand *)> visitor) {
+  if (!isPhi())
+    return false;
+
+  GraphNodeWorklist<SILPhiArgument *, 4> worklist;
+  worklist.initialize(this);
+
+  while (auto *argument = worklist.pop()) {
+    llvm::SmallVector<Operand *> operands;
+    argument->getIncomingPhiOperands(operands);
+
+    for (auto *operand : operands) {
+      SILPhiArgument *forwarded;
+      if ((forwarded = dyn_cast<SILPhiArgument>(operand->get())) &&
+          forwarded->isPhi()) {
+        worklist.insert(forwarded);
+      }
+      if (!visitor(argument, operand))
+        return false;
+    }
+  }
+  return true;
+}
+
 static SILValue
 getSingleTerminatorOperandForPred(const SILBasicBlock *parentBlock,
                                   const SILBasicBlock *predBlock,

lib/SIL/Utils/OwnershipUtils.cpp

@@ -1586,6 +1586,223 @@ void swift::findTransitiveReborrowBaseValuePairs(
   }
 }
 
+/// Visit the phis in the same block as \p phi which are reborrows of a borrow
+/// of one of the values reaching \p phi.
+///
+/// If the visitor returns false, stops visiting and returns false.  Otherwise,
+/// returns true.
+///
+///
+/// When an owned value is passed as a phi argument, it is consumed.  So any
+/// open scope borrowing that owned value must be ended no later than in that
+/// branch instruction.  Either such a borrow scope is ended beforehand
+///     %lifetime = begin_borrow %value
+///     ...
+///     end_borrow %lifetime <-- borrow scope ended here
+///     br block(%value)        <-- before consume
+/// or the borrow scope is ended in the same instruction as the owned value is
+/// consumed
+///     %lifetime = begin_borrow %value
+///     ...
+///     end_borrow %lifetime
+///     br block(%value, %lifetime)         <-- borrow scope ended here
+///                                         <-- in same instruction as consume
+/// In particular, the following is invalid
+///         %lifetime = begin_borrow %value
+///         ...
+///         br block(%value)
+///     block(%value_2 : @owned):
+///         end_borrow %lifetime
+///         destroy_value %value_2
+/// because %lifetime was guaranteed by %value but value is consumed at
+/// `br two`.
+///
+/// Similarly, when a guaranteed value is passed as a phi argument, its borrow
+/// scope ends and a new borrow scope is begun.  And so any open nested borrow
+/// of the original outer borrow must be ended no later than in that branch
+/// instruction.
+///
+///
+/// Given an phi argument
+///     block(..., %value : @owned, ...)
+/// this function finds the adjacent reborrow phis
+///     block(..., %lifetime : @guaranteed, ..., %value : @owned, ...)
+///                ^^^^^^^^^^^^^^^^^^^^^^^
+/// one of whose reaching values is a borrow of a reaching value of %value.
+///
+/// Finding these is more complicated than merely looking for guaranteed
+/// operands adjacent to the incoming operands to phi and which are borrows of
+/// the value consumed there.  The reason is that they might not be borrows of
+/// that incoming value _directly_ but rather reborrows of some other reborrow
+/// if the incoming value is itself a phi argument:
+///         %lifetime = begin_borrow %value
+///         br one(%value, %lifetime)
+///     one(%value_1 : @owned, %lifetime_1 : @guaranteed)
+///         br two(%value_1, %lifetime_1)
+///     two(%value_2 : @owned, %lifetime_2 : @guaranteed)
+///         end_borrow %lifetime_2
+///         destroy_value %value_2
+///
+/// When called with %value_2, \p visitor is invoked with both:
+///     two(%value_2 : @owned, %lifetime_2 : @guaranteed)
+///                            ^^^^^^^^^^^^^^^^^^^^^^^^^
+bool swift::visitAdjacentReborrowsOfPhi(
+    SILPhiArgument *phi, function_ref<bool(SILPhiArgument *)> visitor) {
+  assert(phi->isPhi());
+
+  // First, collect all the values that reach \p phi, that is:
+  // - operands to the phi
+  // - operands to phis which are operands to the phi
+  // - and so forth.
+  SmallPtrSet<SILValue, 8> reachingValues;
+  // At the same time, record all the phis in \p phi's phi web: the phis which
+  // are transitively operands to \p phi.  This is the subset of \p
+  // reachingValues that are phis.
+  SmallVector<SILPhiArgument *, 4> phis;
+  phi->visitTransitiveIncomingPhiOperands(
+      [&](auto *phi, auto *operand) -> bool {
+        phis.push_back(phi);
+        reachingValues.insert(phi);
+        reachingValues.insert(operand->get());
+        return true;
+      });
+
+  // Second, find all the guaranteed phis one of whose operands _could_ (by
+  // dint of being adjacent to a phi in the phi web with the appropriate
+  // ownership and type) be a reborrow of a reaching value of \p phi.
+  SmallVector<SILPhiArgument *, 4> candidates;
+  for (auto *phi : phis) {
+    SILBasicBlock *block = phi->getParentBlock();
+    for (auto *uncastAdjacent : block->getArguments()) {
+      auto *adjacent = cast<SILPhiArgument>(uncastAdjacent);
+      if (adjacent == phi)
+        continue;
+      if (adjacent->getType() != phi->getType())
+        continue;
+      if (adjacent->getOwnershipKind() != OwnershipKind::Guaranteed)
+        continue;
+      candidates.push_back(adjacent);
+    }
+  }
+
+  // Finally, look through \p candidates to find those one of whose incoming
+  // operands either
+  // (1) borrow one of reaching values of \p phi
+  // or (2) is itself a guaranteed phi which does so.
+  // Because we may discover a reborrow R1 of type (1) after visiting another
+  // R2 of type (2) which reborrows R1, we need to iterate to a fixed point.
+  //
+  // Record all the phis that we see which are borrows or reborrows of a
+  // reaching value \p so that we can check for case (2) above.
+  //
+  // Visit those phis which are both reborrows of a reaching value AND are in
+  // the same block as \phi.
+  //
+  // For example, given
+  //
+  //         %lifetime = begin_borrow %value
+  //         br one(%value, %lifetime)
+  //     one(%value_1 : @owned, %lifetime_1 : @guaranteed)
+  //         br two(%value_1, %lifetime_1)
+  //     two(%value_2 : @owned, %lifetime_2 : @guaranteed)
+  //         end_borrow %lifetime_2
+  //         destroy_value %value_2
+  //
+  // when visiting the reborrow phis adjacent to %value_2, The following steps
+  // would be taken:
+  //
+  // (1) Look at the first candidate:
+  //   two(%value_2 : @owned, %lifetime_2 : @guaranteed)
+  //                          ^^^^^^^^^^^^^^^^^^^^^^^^^
+  // but see that its one incoming value
+  //   br two(%value_1, %lifetime_1)
+  //                    ^^^^^^^^^^^
+  // although a phi argument itself, is not known (yet!) to be a reborrow phi.
+  // So the first candidate is NOT (yet!) added to reborrowPhis.
+  //
+  // (2) Look at the second candidate:
+  //     one(%value_1 : @owned, %lifetime_1 : @guaranteed)
+  //                            ^^^^^^^^^^^^^^^^^^^^^^^^^
+  // and see that one of its incoming values
+  //   br one(%value, %lifetime)
+  //                  ^^^^^^^^^
+  // is a borrow
+  //   %lifetime = begin_borrow %value
+  // of %value, one of the values reaching %value_2.
+  // So the second candidate IS added to reborrowPhis.
+  // AND changed is set to true, so we will repeat the outer loop.
+  // But this candidate is not adjacent to our phi %value_2, so it is not
+  // visited.
+  //
+  // (4.5) Changed is true: repeat the outer loop.  Set changed to false.
+  //
+  // (3) Look at the first candidate:
+  //   two(%value_2 : @owned, %lifetime_2 : @guaranteed)
+  //                          ^^^^^^^^^^^^^^^^^^^^^^^^^
+  // and see that one of its incoming values
+  //   br two(%value_1, %lifetime_1)
+  //                    ^^^^^^^^^^^
+  // is itself a phi
+  //   one(%value_1 : @owned, %lifetime_1 : @guaranteed)
+  //                          ^^^^^^^^^^^^^^^^^^^^^^^^^
+  // which was added to reborrowPhis in (2).
+  // So the first candidate IS added to reborrowPhis.
+  // AND changed is set to true.
+  // ALSO, see that the first candidate IS adjacent to our phi %value_2, so our
+  // visitor is invoked with the first candidate.
+  //
+  // (4) Look at the second candidate.
+  // See that it is already a member of reborrowPhis.
+  //
+  // (4.5) Changed is true: repeat the outer loop.  Set changed to false.
+  //
+  // (5) Look at the first candidate.
+  // See that it is already a member of reborrowPhis.
+  //
+  // (6) Look at the second candidate.
+  // See that it is already a member of reborrowPhis.
+  //
+  // (6.5) Changed is false: exit the outer loop.
+  bool changed = false;
+  SmallSetVector<SILPhiArgument *, 4> reborrowPhis;
+  do {
+    changed = false;
+    for (auto *candidate : candidates) {
+      if (reborrowPhis.contains(candidate))
+        continue;
+      auto success = candidate->visitIncomingPhiOperands([&](auto *operand) {
+        // If the value being reborrowed is itself a reborrow of a value
+        // reaching \p phi, then visit it.
+        SILPhiArgument *forwarded;
+        if ((forwarded = dyn_cast<SILPhiArgument>(operand->get()))) {
+          if (!reborrowPhis.contains(forwarded))
+            return true;
+          changed = true;
+          reborrowPhis.insert(candidate);
+          if (candidate->getParentBlock() == phi->getParentBlock())
+            return visitor(candidate);
+          return true;
+        }
+        BeginBorrowInst *bbi;
+        if (!(bbi = dyn_cast<BeginBorrowInst>(operand->get())))
+          return true;
+        auto borrowee = bbi->getOperand();
+        if (!reachingValues.contains(borrowee))
+          return true;
+        changed = true;
+        reborrowPhis.insert(candidate);
+        if (candidate->getParentBlock() == phi->getParentBlock())
+          return visitor(candidate);
+        return true;
+      });
+      if (!success)
+        return false;
+    }
+  } while (changed);
+
+  return true;
+}
+
 void swift::visitTransitiveEndBorrows(
     SILValue value,
     function_ref<void(EndBorrowInst *)> visitEndBorrow) {

lib/SILOptimizer/Utils/CanonicalOSSALifetime.cpp

@@ -114,6 +114,13 @@ void swift::copyLiveUse(Operand *use, InstModCallbacks &instModCallbacks) {
 bool CanonicalizeOSSALifetime::computeCanonicalLiveness() {
   defUseWorklist.initialize(currentDef);
   while (SILValue value = defUseWorklist.pop()) {
+    SILPhiArgument *arg;
+    if ((arg = dyn_cast<SILPhiArgument>(value)) && arg->isPhi()) {
+      visitAdjacentReborrowsOfPhi(arg, [&](SILPhiArgument *reborrow) {
+        defUseWorklist.insert(reborrow);
+        return true;
+      });
+    }
     for (Operand *use : value->getUses()) {
       auto *user = use->getUser();
 
@@ -168,8 +175,33 @@ bool CanonicalizeOSSALifetime::computeCanonicalLiveness() {
       case OperandOwnership::InteriorPointer:
       case OperandOwnership::ForwardingBorrow:
       case OperandOwnership::EndBorrow:
+        // Guaranteed values are considered uses of the value when the value is
+        // an owned phi and the guaranteed values are adjacent reborrow phis or
+        // reborrow of such.
+        liveness.updateForUse(user, /*lifetimeEnding*/ false);
+        break;
       case OperandOwnership::Reborrow:
-        llvm_unreachable("operand kind cannot take an owned value");
+        BranchInst *branch;
+        if (!(branch = dyn_cast<BranchInst>(user))) {
+          // Non-phi reborrows (tuples, etc) never end the lifetime of the owned
+          // value.
+          liveness.updateForUse(user, /*lifetimeEnding*/ false);
+          defUseWorklist.insert(cast<SingleValueInstruction>(user));
+          break;
+        }
+        if (is_contained(user->getOperandValues(), currentDef)) {
+          // An adjacent phi consumes the value being reborrowed.  Although this
+          // use doesn't end the lifetime, this user does.
+          liveness.updateForUse(user, /*lifetimeEnding*/ true);
+          break;
+        }
+        // No adjacent phi consumes the value.  This use is not lifetime ending.
+        liveness.updateForUse(user, /*lifetimeEnding*/ false);
+        // This branch reborrows a guaranteed phi whose lifetime is dependent on
+        // currentDef.  Uses of the reborrowing phi extend liveness.
+        auto *reborrow = branch->getArgForOperand(use);
+        defUseWorklist.insert(reborrow);
+        break;
       }
     }
   }