Add utilities to support OSSA update after running SSAUpdater.

This directly adds support to BasicBlockCloner for updating OSSA.

It also adds a much more general-purpose GuaranteedPhiBorrowFixup
utility which can be used for more complicated SSA updates, in which
multiple phis need to be created. More generally, it handles adding
nested borrow scopes for guaranteed phis even when that phi is used by
other guaranteed phis.

Author: atrick

include/swift/SILOptimizer/Utils/BasicBlockOptUtils.h

@@ -83,6 +83,19 @@ bool rotateLoop(SILLoop *loop, DominanceInfo *domInfo, SILLoopInfo *loopInfo,
                 bool rotateSingleBlockLoops, SILBasicBlock *upToBB,
                 bool shouldVerify);
 
+//===----------------------------------------------------------------------===//
+//                             BasicBlock Cloning
+//===----------------------------------------------------------------------===//
+
+/// Return true if the \p termInst can be cloned. If termInst produces a
+/// guaranteed result, then it must be possible to create a nested borrow scope
+/// for that result when the cloner generates a guaranteed phi.
+///
+/// Note: if all the terminators uses will also be cloned, then a guaranteed phi
+/// won't be necessary. This is one of the reasons that cloning a region is much
+/// better than cloning a single block at a time.
+bool canCloneTerminator(TermInst *termInst);
+
 /// Sink address projections to their out-of-block uses. This is
 /// required after cloning a block and before calling
 /// updateSSAAfterCloning to avoid address-type phis.
@@ -173,6 +186,7 @@ class BasicBlockCloner : public SILCloner<BasicBlockCloner> {
       if (!canCloneInstruction(&inst))
         return false;
     }
+    canCloneTerminator(origBB->getTerminator());
     return true;
   }
 
@@ -221,17 +235,34 @@ class BasicBlockCloner : public SILCloner<BasicBlockCloner> {
     bi->eraseFromParent();
   }
 
+  /// Create phis and maintain OSSA invariants.
+  ///
+  /// Note: This must be called after calling cloneBlock or cloneBranchTarget,
+  /// before using any OSSA utilities.
+  ///
+  /// The client may perform arbitrary branch fixups and dead block removal
+  /// after cloning and before calling this.
+  ///
+  /// WARNING: If client converts terminator results to phis (e.g. replaces a
+  /// switch_enum with a branch), then it must call this before performing that
+  /// transformation, or fix the OSSA representation of that value itself.
+  void updateOSSAAfterCloning();
+
   /// Get the newly cloned block corresponding to `origBB`.
   SILBasicBlock *getNewBB() {
     return remapBasicBlock(origBB);
   }
 
   bool wasCloned() { return isBlockCloned(origBB); }
 
-  /// Helper function to perform SSA updates after calling cloneBranchTarget.
-  void updateSSAAfterCloning();
-
 protected:
+  /// Helper function to perform SSA updates used by updateOSSAAfterCloning.
+  void updateSSAAfterCloning(SmallVectorImpl<SILPhiArgument *> &newPhis);
+
+  /// Given a terminator result, either from the original or the cloned block,
+  /// update OSSA for any phis created for the result during edge splitting.
+  void updateOSSATerminatorResult(SILPhiArgument *termResult);
+
   // MARK: CRTP overrides.
 
   /// Override getMappedValue to allow values defined outside the block to be

include/swift/SILOptimizer/Utils/OwnershipOptUtils.h

@@ -37,6 +37,19 @@ inline bool requiresOSSACleanup(SILValue v) {
 // Defined in BasicBlockUtils.h
 struct JointPostDominanceSetComputer;
 
+/// Given a new phi that may use a guaranteed value, create nested borrow scopes
+/// for its incoming operands and end_borrows that cover the phi's extended
+/// borrow scope, which transitively includes any phis that use this phi.
+///
+/// Returns true if any changes were made.
+///
+/// Note: \p newPhi itself might not have Guaranteed ownership. A phi that
+/// converts Guaranteed to None ownership still needs nested borrows.
+///
+/// Note: This may be called on partially invalid OSSA form, where multiple
+/// newly created phis do not yet have a borrow scope.
+bool createBorrowScopeForPhiOperands(SILPhiArgument *newPhi);
+
 /// A struct that contains context shared in between different operation +
 /// "ownership fixup" utilities. Please do not put actual methods on this, it is
 /// meant to be composed with.

lib/SILOptimizer/Transforms/SimplifyCFG.cpp

@@ -382,7 +382,7 @@ bool SimplifyCFG::threadEdge(const ThreadInfo &ti) {
           dyn_cast<BranchInst>(ThreadedSuccessorBlock->getTerminator())) {
     simplifyBranchBlock(branchInst);
   }
-  Cloner.updateSSAAfterCloning();
+  Cloner.updateOSSAAfterCloning();
   return true;
 }
 
@@ -1149,7 +1149,7 @@ bool SimplifyCFG::tryJumpThreading(BranchInst *BI) {
   // Duplicate the destination block into this one, rewriting uses of the BBArgs
   // to use the branch arguments as we go.
   Cloner.cloneBranchTarget(BI);
-  Cloner.updateSSAAfterCloning();
+  Cloner.updateOSSAAfterCloning();
 
   // Once all the instructions are copied, we can nuke BI itself.  We also add
   // the threaded and edge block to the worklist now that they (likely) can be
@@ -3027,7 +3027,7 @@ bool SimplifyCFG::tailDuplicateObjCMethodCallSuccessorBlocks() {
       continue;
 
     Cloner.cloneBranchTarget(Branch);
-    Cloner.updateSSAAfterCloning();
+    Cloner.updateOSSAAfterCloning();
 
     Changed = true;
     // Simplify the cloned block and continue tail duplicating through its new

lib/SILOptimizer/Utils/BasicBlockOptUtils.cpp

@@ -13,6 +13,7 @@
 #include "swift/SILOptimizer/Utils/BasicBlockOptUtils.h"
 #include "swift/SILOptimizer/Utils/CFGOptUtils.h"
 #include "swift/SILOptimizer/Utils/InstOptUtils.h"
+#include "swift/SILOptimizer/Utils/OwnershipOptUtils.h"
 #include "swift/SILOptimizer/Utils/SILSSAUpdater.h"
 
 using namespace swift;
@@ -87,7 +88,99 @@ bool swift::removeUnreachableBlocks(SILFunction &f) {
   return changed;
 }
 
-void BasicBlockCloner::updateSSAAfterCloning() {
+//===----------------------------------------------------------------------===//
+//                             BasicBlock Cloning
+//===----------------------------------------------------------------------===//
+
+// Return true if a guaranteed terminator result can be borrowed such that the
+// nested borrow scope covers all its uses.
+static bool canBorrowGuaranteedResult(SILValue guaranteedResult) {
+  if (guaranteedResult.getOwnershipKind() != OwnershipKind::Guaranteed) {
+    // Either this terminator forwards an owned value, or it is some legal
+    // conversion to a non-guaranteed value. Either way, not interesting.
+    return true;
+  }
+  SmallVector<Operand *, 16> usePoints;
+  return findInnerTransitiveGuaranteedUses(guaranteedResult, usePoints);
+}
+
+bool swift::canCloneTerminator(TermInst *termInst) {
+  // TODO: this is an awkward way to check for guaranteed terminator results.
+  for (Operand &oper : termInst->getAllOperands()) {
+    if (oper.getOperandOwnership() != OperandOwnership::ForwardingBorrow)
+      continue;
+
+    if (!ForwardingOperand(&oper).visitForwardedValues(
+          [&](SILValue termResult) {
+            return canBorrowGuaranteedResult(termResult);
+          })) {
+        return false;
+    }
+  }
+  return true;
+}
+
+/// Given a terminator result, either from the original or the cloned block,
+/// update OSSA for any phis created for the result during edge splitting.
+void BasicBlockCloner::updateOSSATerminatorResult(SILPhiArgument *termResult) {
+  assert(termResult->isTerminatorResult() && "precondition");
+
+  // If the terminator result is used by a phi, then it is invalid OSSA
+  // which was created by edge splitting.
+  for (Operand *termUse : termResult->getUses()) {
+    if (auto phiOper = PhiOperand(termUse)) {
+      createBorrowScopeForPhiOperands(phiOper.getValue());
+    }
+  }
+}
+
+// Cloning does not invalidate ownership lifetime. When it clones values, it
+// also either clones the consumes, or creates the necessary phis that consume
+// the new values on all paths.  However, cloning may create new phis of
+// inner guaranteed values. Since phis are reborrows, they are only allowed to
+// use BorrowedValues. Therefore, we must create nested borrow scopes for any
+// new phis whose arguments aren't BorrowedValues. Note that other newly created
+// phis are themselves BorrowedValues, so only one level of nested borrow is
+// needed per value, per new phi that the value reaches.
+void BasicBlockCloner::updateOSSAAfterCloning() {
+  SmallVector<SILPhiArgument *, 4> updateSSAPhis;
+  if (!origBB->getParent()->hasOwnership()) {
+    updateSSAAfterCloning(updateSSAPhis);
+    return;
+  }
+
+  // If the original basic block has terminator results, then all phis in the
+  // exit blocks are new phis that used to be terminator results.
+  //
+  // Create nested borrow scopes for terminator results that were converted to
+  // phis during edge splitting. This is simpler to check before SSA update.
+  //
+  // This assumes that the phis introduced by update-SSA below cannot be users
+  // of the phis that were created in exitBBs during block cloning. Otherwise
+  // borrowPhiArguments would handle them twice.
+  auto *termInst = origBB->getTerminator();
+  // FIXME: cond_br args should not exist in OSSA
+  if (!isa<BranchInst>(termInst) && !isa<CondBranchInst>(termInst)) {
+    // Update all of the terminator results.
+    for (auto *succBB : origBB->getSuccessorBlocks()) {
+      for (SILArgument *termResult : succBB->getArguments()) {
+        updateOSSATerminatorResult(cast<SILPhiArgument>(termResult));
+      }
+    }
+  }
+
+  // Update SSA form before calling OSSA update utilities to maintain a layering
+  // of SIL invariants.
+  updateSSAAfterCloning(updateSSAPhis);
+
+  // Create nested borrow scopes for phis created during SSA update.
+  for (auto *phi : updateSSAPhis) {
+    createBorrowScopeForPhiOperands(phi);
+  }
+}
+
+void BasicBlockCloner::updateSSAAfterCloning(
+    SmallVectorImpl<SILPhiArgument *> &newPhis) {
   // All instructions should have been checked by canCloneInstruction. But we
   // still need to check the arguments.
   for (auto arg : origBB->getArguments()) {
@@ -98,7 +191,7 @@ void BasicBlockCloner::updateSSAAfterCloning() {
   if (!needsSSAUpdate)
     return;
 
-  SILSSAUpdater ssaUpdater;
+  SILSSAUpdater ssaUpdater(&newPhis);
   for (auto availValPair : availVals) {
     auto inst = availValPair.first;
     if (inst->use_empty())

lib/SILOptimizer/Utils/CheckedCastBrJumpThreading.cpp

@@ -673,7 +673,7 @@ void CheckedCastBrJumpThreading::optimizeFunction() {
     edit->modifyCFGForSuccessPreds(Cloner);
 
     if (Cloner.wasCloned()) {
-      Cloner.updateSSAAfterCloning();
+      Cloner.updateOSSAAfterCloning();
 
       if (!Cloner.getNewBB()->pred_empty())
         BlocksForWorklist.push_back(Cloner.getNewBB());

lib/SILOptimizer/Utils/OwnershipOptUtils.cpp

@@ -1225,3 +1225,159 @@ SILBasicBlock::iterator OwnershipReplaceSingleUseHelper::perform() {
   SingleUseReplacementUtility utility{use, newValue, *ctx};
   return utility.perform();
 }
+
+//===----------------------------------------------------------------------===//
+//                      createBorrowScopeForPhiOperands
+//===----------------------------------------------------------------------===//
+
+/// Given a phi that has been newly created or converted from terminator
+/// results, check for inner guaranteed operands (which do not introduce a
+/// borrow scope). This is invalid OSSA because the phi is a reborrow, and all
+/// borrow-scope-ending instructions must directly use the BorrowedValue that
+/// introduces the scope.
+///
+/// Create nested borrow scopes for its operands.
+///
+/// Transitively follow its phi uses.
+///
+/// Create end_borrows at all points that cover the inner uses.
+///
+/// The client must check canCloneTerminator() first to make sure that the
+/// search for transitive uses does not encouter a PointerEscape.
+class GuaranteedPhiBorrowFixup {
+  // A phi in mustConvertPhis has already been determined to be part of this
+  // new nested borrow scope.
+  SmallSetVector<SILPhiArgument *, 8> mustConvertPhis;
+
+  // Phi operands that are already within the new nested borrow scope.
+  llvm::SmallDenseSet<PhiOperand, 8> nestedPhiOperands;
+
+public:
+  /// Return true if an extended nested borrow scope was created.
+  bool createExtendedNestedBorrowScope(SILPhiArgument *newPhi);
+
+protected:
+  bool phiOperandNeedsBorrow(Operand *operand) {
+    SILValue inVal = operand->get();
+    if (inVal.getOwnershipKind() != OwnershipKind::Guaranteed) {
+      assert(inVal.getOwnershipKind() == OwnershipKind::None);
+      return false;
+    }
+    // This operand needs a nested borrow if inVal is not a BorrowedValue.
+    return !bool(BorrowedValue(inVal));
+  }
+
+  void borrowPhiOperand(Operand *oper) {
+    // Begin the borrow just before the branch.
+    SILInstruction *borrowPoint = oper->getUser();
+    auto loc = RegularLocation::getAutoGeneratedLocation(borrowPoint->getLoc());
+    auto *borrow =
+        SILBuilderWithScope(borrowPoint).createBeginBorrow(loc, oper->get());
+    oper->set(borrow);
+  }
+
+  EndBorrowInst *createEndBorrow(SILValue guaranteedValue,
+                                 SILBasicBlock::iterator borrowPoint) {
+    auto loc = borrowPoint->getLoc();
+    return SILBuilderWithScope(borrowPoint)
+        .createEndBorrow(loc, guaranteedValue);
+  }
+
+  void insertEndBorrowsAndFindPhis(SILPhiArgument *phi);
+};
+
+void GuaranteedPhiBorrowFixup::insertEndBorrowsAndFindPhis(
+    SILPhiArgument *phi) {
+  // Scope ending instructions are only needed for nontrivial results.
+  if (phi->getOwnershipKind() != OwnershipKind::Guaranteed) {
+    assert(phi->getOwnershipKind() == OwnershipKind::None);
+    return;
+  }
+  SmallVector<Operand *, 16> usePoints;
+  bool result = findInnerTransitiveGuaranteedUses(phi, usePoints);
+  assert(result && "should be checked by canCloneTerminator");
+  (void)result;
+
+  // Add usePoints to a set for phi membership checking.
+  //
+  // FIXME: consider integrating with ValueLifetimeBoundary instead.
+  SmallPtrSet<Operand *, 16> useSet(usePoints.begin(), usePoints.end());
+
+  auto phiUsers = llvm::map_range(usePoints, ValueBase::UseToUser());
+  ValueLifetimeAnalysis lifetimeAnalysis(phi, phiUsers);
+  ValueLifetimeBoundary boundary;
+  lifetimeAnalysis.computeLifetimeBoundary(boundary);
+
+  for (auto *boundaryEdge : boundary.boundaryEdges) {
+    createEndBorrow(phi, boundaryEdge->begin());
+  }
+
+  for (SILInstruction *lastUser : boundary.lastUsers) {
+    // If the last use is a branch, transitively process the phi.
+    if (isa<BranchInst>(lastUser)) {
+      for (Operand &oper : lastUser->getAllOperands()) {
+        if (!useSet.count(&oper))
+          continue;
+
+        PhiOperand phiOper(&oper);
+        nestedPhiOperands.insert(phiOper);
+        mustConvertPhis.insert(phiOper.getValue());
+        continue;
+      }
+    }
+    // If the last user is a terminator, add the successors as boundary edges.
+    if (isa<TermInst>(lastUser)) {
+      for (auto *succBB : lastUser->getParent()->getSuccessorBlocks()) {
+        // succBB cannot already be in boundaryEdges. It has a
+        // single predecessor with liveness ending at the terminator, which
+        // means it was not live into any successor blocks.
+        createEndBorrow(phi, succBB->begin());
+      }
+      continue;
+    }
+    // Otherwise, just plop down an end_borrow after the last use.
+    createEndBorrow(phi, std::next(lastUser->getIterator()));
+  }
+};
+
+// For each phi that transitively uses an inner guaranteed value, create nested
+// borrow scopes so that it is a well-formed reborrow.
+bool GuaranteedPhiBorrowFixup::
+createExtendedNestedBorrowScope(SILPhiArgument *newPhi) {
+  // Determine if this new phi needs a nested borrow scope. If so, seed the
+  // Visit phi operands, returning false as soon as one needs a borrow.
+  if (!newPhi->visitIncomingPhiOperands(
+        [&](Operand *op) { return !phiOperandNeedsBorrow(op); })) {
+    mustConvertPhis.insert(newPhi);
+  }
+  if (mustConvertPhis.empty())
+    return false;
+
+  // mustConvertPhis grows in this loop.
+  for (unsigned mustConvertIdx = 0; mustConvertIdx < mustConvertPhis.size();
+         ++mustConvertIdx) {
+    SILPhiArgument *phi = mustConvertPhis[mustConvertIdx];
+    insertEndBorrowsAndFindPhis(phi);
+  }
+  // To handle recursive phis, first discover all phis before attempting to
+  // borrow any phi operands.
+  for (SILPhiArgument *phi : mustConvertPhis) {
+    phi->visitIncomingPhiOperands([&](Operand *op) {
+      if (!nestedPhiOperands.count(op))
+        borrowPhiOperand(op);
+      return true;
+    });
+  }
+  return true;
+}
+
+// Note: \p newPhi itself might not have Guaranteed ownership. A phi that
+// converts Guaranteed to None ownership still needs nested borrows.
+//
+// Note: This may be called on partially invalid OSSA form, where multiple
+// newly created phis do not yet have a borrow scope. The implementation
+// assumes that this API will eventually be called for all such new phis until
+// OSSA is fully valid.
+bool swift::createBorrowScopeForPhiOperands(SILPhiArgument *newPhi) {
+  return GuaranteedPhiBorrowFixup().createExtendedNestedBorrowScope(newPhi);
+}