[CanonicalizeOSSALifetime] Extend Onone lifetimes.

To improve the debugging experience of values whose lifetimes are
canonicalized without compromising the semantics expressed in the source
language, when canonicalizing OSSA lifetimes at Onone, lengthen
lifetimes as much as possible without incurring copies that would be
eliminated at O.

rdar://99618502

Author: nate-chandler

include/swift/SILOptimizer/Utils/CanonicalizeOSSALifetime.h

@@ -205,6 +205,10 @@ class CanonicalizeOSSALifetime {
   /// liveness may be pruned during canonicalization.
   bool pruneDebugMode;
 
+  /// If true, lifetimes will not be shortened except when necessary to avoid
+  /// copies.
+  bool maximizeLifetime;
+
   /// If true and we are processing a value of move_only type, emit a diagnostic
   /// when-ever we need to insert a copy_value.
   std::function<void(Operand *)> moveOnlyCopyValueNotification;
@@ -240,6 +244,9 @@ class CanonicalizeOSSALifetime {
   /// Visited set for general def-use traversal that prevents revisiting values.
   GraphNodeWorklist<SILValue, 8> defUseWorklist;
 
+  /// The blocks that were discovered by PrunedLiveness.
+  SmallVector<SILBasicBlock *, 32> discoveredBlocks;
+
   /// Pruned liveness for the extended live range including copies. For this
   /// purpose, only consuming instructions are considered "lifetime
   /// ending". end_borrows do not end a liverange that may include owned copies.
@@ -248,7 +255,7 @@ class CanonicalizeOSSALifetime {
   /// The destroys of the value.  These are not uses, but need to be recorded so
   /// that we know when the last use in a consuming block is (without having to
   /// repeatedly do use-def walks from destroys).
-  SmallPtrSet<SILInstruction *, 8> destroys;
+  SmallPtrSetVector<SILInstruction *, 8> destroys;
 
   /// Information about consuming instructions discovered in this canonical OSSA
   /// lifetime.
@@ -286,15 +293,17 @@ class CanonicalizeOSSALifetime {
   }
 
   CanonicalizeOSSALifetime(
-      bool pruneDebugMode, NonLocalAccessBlockAnalysis *accessBlockAnalysis,
-      DominanceInfo *domTree, InstructionDeleter &deleter,
+      bool pruneDebugMode, bool maximizeLifetime,
+      NonLocalAccessBlockAnalysis *accessBlockAnalysis, DominanceInfo *domTree,
+      InstructionDeleter &deleter,
       std::function<void(Operand *)> moveOnlyCopyValueNotification = nullptr,
       std::function<void(Operand *)> moveOnlyFinalConsumingUse = nullptr)
-      : pruneDebugMode(pruneDebugMode),
+      : pruneDebugMode(pruneDebugMode), maximizeLifetime(maximizeLifetime),
         moveOnlyCopyValueNotification(moveOnlyCopyValueNotification),
         moveOnlyFinalConsumingUse(moveOnlyFinalConsumingUse),
         accessBlockAnalysis(accessBlockAnalysis), domTree(domTree),
-        deleter(deleter) {}
+        deleter(deleter),
+        liveness(maximizeLifetime ? &discoveredBlocks : nullptr) {}
 
   SILValue getCurrentDef() const { return liveness.getDef(); }
 
@@ -313,6 +322,7 @@ class CanonicalizeOSSALifetime {
     consumingBlocks.clear();
     debugValues.clear();
     liveness.clear();
+    discoveredBlocks.clear();
   }
 
   /// Top-Level API: rewrites copies and destroys within \p def's extended
@@ -345,6 +355,8 @@ class CanonicalizeOSSALifetime {
 
   void findExtendedBoundary(PrunedLivenessBoundary &boundary);
 
+  void extendUnconsumedLiveness(PrunedLivenessBoundary &boundary);
+
   void insertDestroysOnBoundary(PrunedLivenessBoundary &boundary);
 
   void rewriteCopies();

lib/SILOptimizer/Mandatory/MoveOnlyAddressChecker.cpp

@@ -746,7 +746,9 @@ struct MoveOnlyChecker {
           instToDelete->eraseFromParent();
         })),
         canonicalizer(
-            false /*pruneDebugMode*/, accessBlockAnalysis, domTree, deleter,
+            false /*pruneDebugMode*/,
+            !fn->shouldOptimize() /*maximizeLifetime*/, accessBlockAnalysis,
+            domTree, deleter,
             [&](Operand *use) { consumingUsesNeedingCopy.push_back(use); },
             [&](Operand *use) { finalConsumingUses.push_back(use); }) {}
 

lib/SILOptimizer/Mandatory/MoveOnlyObjectChecker.cpp

@@ -401,8 +401,9 @@ bool MoveOnlyChecker::check(NonLocalAccessBlockAnalysis *accessBlockAnalysis,
   };
 
   CanonicalizeOSSALifetime canonicalizer(
-      false /*pruneDebugMode*/, accessBlockAnalysis, domTree, deleter,
-      foundConsumingUseNeedingCopy, foundConsumingUseNotNeedingCopy);
+      false /*pruneDebugMode*/, !fn->shouldOptimize() /*maximizeLifetime*/,
+      accessBlockAnalysis, domTree, deleter, foundConsumingUseNeedingCopy,
+      foundConsumingUseNotNeedingCopy);
   auto moveIntroducers = llvm::makeArrayRef(moveIntroducersToProcess.begin(),
                                             moveIntroducersToProcess.end());
   SmallPtrSet<MarkMustCheckInst *, 4> valuesWithDiagnostics;

lib/SILOptimizer/SILCombiner/SILCombine.cpp

@@ -350,8 +350,10 @@ void SILCombiner::canonicalizeOSSALifetimes(SILInstruction *currentInst) {
   InstructionDeleter deleter(std::move(canonicalizeCallbacks));
 
   DominanceInfo *domTree = DA->get(&Builder.getFunction());
-  CanonicalizeOSSALifetime canonicalizer(false /*prune debug*/, NLABA, domTree,
-                                         deleter);
+  CanonicalizeOSSALifetime canonicalizer(
+      false /*prune debug*/,
+      !parentTransform->getFunction()->shouldOptimize() /*maximizeLifetime*/,
+      NLABA, domTree, deleter);
   CanonicalizeBorrowScope borrowCanonicalizer(deleter);
 
   while (!defsToCanonicalize.empty()) {

lib/SILOptimizer/Transforms/CopyPropagation.cpp

@@ -373,7 +373,7 @@ static bool sinkOwnedForward(SILInstruction *ownedForward,
 namespace {
 
 class CopyPropagation : public SILFunctionTransform {
-  /// True if debug_value instructions should be pruned.
+  /// If true, debug_value instructions should be pruned.
   bool pruneDebug;
   /// True if all values should be canonicalized.
   bool canonicalizeAll;
@@ -437,8 +437,9 @@ void CopyPropagation::run() {
 
   // canonicalizer performs all modifications through deleter's callbacks, so we
   // don't need to explicitly check for changes.
-  CanonicalizeOSSALifetime canonicalizer(pruneDebug, accessBlockAnalysis,
-                                         domTree, deleter);
+  CanonicalizeOSSALifetime canonicalizer(
+      pruneDebug, /*maximizeLifetime=*/!getFunction()->shouldOptimize(),
+      accessBlockAnalysis, domTree, deleter);
 
   // NOTE: We assume that the function is in reverse post order so visiting the
   //       blocks and pushing begin_borrows as we see them and then popping them

lib/SILOptimizer/Utils/CanonicalizeOSSALifetime.cpp

@@ -52,6 +52,9 @@
 ///
 /// 2. Liveness is extended out to original destroys to avoid spurious changes.
 ///
+/// 3. In the Onone mode, liveness is preserved to its previous extent whenever
+/// doing so doesn't incur extra copies compared to what is done in the O mode.
+///
 //===----------------------------------------------------------------------===//
 
 #define DEBUG_TYPE "copy-propagation"
@@ -64,6 +67,7 @@
 #include "swift/SILOptimizer/Utils/DebugOptUtils.h"
 #include "swift/SILOptimizer/Utils/InstructionDeleter.h"
 #include "swift/SILOptimizer/Utils/ValueLifetime.h"
+#include "llvm/ADT/STLExtras.h"
 #include "llvm/ADT/Statistic.h"
 
 using namespace swift;
@@ -542,6 +546,8 @@ class CanonicalizeOSSALifetimeBoundaryExtender {
 
 void CanonicalizeOSSALifetime::findExtendedBoundary(
     PrunedLivenessBoundary &boundary) {
+  assert(boundary.lastUsers.size() == 0 && boundary.boundaryEdges.size() == 0 &&
+         boundary.deadDefs.size() == 0);
   PrunedLivenessBoundary originalBoundary;
   liveness.computeBoundary(originalBoundary, consumingBlocks.getArrayRef());
 
@@ -611,6 +617,130 @@ void CanonicalizeOSSALifetime::insertDestroysOnBoundary(
   }
 }
 
+/// At -Onone, there are some conflicting goals:
+/// On the one hand: good debugging experience.
+/// (1) do not shorten value's lifetime
+/// On the other: demonstrate semantics.
+/// (2) consume value at same places it will be consumed at -O
+/// (3) ensure there are no more copies than there would be at -O
+///
+/// (2) and (3) are equivalent--extra (compared to -O) copies arise from failing
+/// to end lifetimes at consuming uses (which then need their own copies).
+///
+/// We achieve (2) and (3) always.  We achieve (1) where possible.
+///
+/// Conceptually, the strategy is the following:
+/// - Collect the blocks in which the value was live before canonicalization.
+///   These are the "original" live blocks (originalLiveBlocks).
+///   [Color these blocks green.]
+/// - From within that collection, collect the blocks which contain a _final_
+///   consuming, non-destroy use, and their successors.
+///   These are the "consumed" blocks (consumedAtExitBlocks).
+///   [Color these blocks red.]
+/// - Extend liveness down to the boundary between originalLiveBlocks and
+///   consumedAtExitBlocks blocks.
+///   [Extend liveness down to the boundary between green blocks and red.]
+/// - In particular, in regions of originalLiveBlocks which have no boundary
+///   with consumedAtExitBlocks, liveness should be extended to its original
+///   extent.
+///   [Extend liveness down to the boundary between green blocks and uncolored.]
+void CanonicalizeOSSALifetime::extendUnconsumedLiveness(
+    PrunedLivenessBoundary &boundary) {
+  auto currentDef = getCurrentDef();
+
+  // First, collect the blocks that were _originally_ live.  We can't use
+  // liveness here because it doesn't include blocks that occur before a
+  // destroy_value.
+  BasicBlockSet originalLiveBlocks(currentDef->getFunction());
+  {
+    // Some of the work here was already done by computeCanonicalLiveness.
+    // Specifically, it already discovered all blocks containing (transitive)
+    // uses and blocks that appear between them and the def.
+    //
+    // Seed the set with what it already discovered.
+    for (auto *discoveredBlock : liveness.getDiscoveredBlocks())
+      originalLiveBlocks.insert(discoveredBlock);
+
+    // Start the walk from the consuming blocks (which includes destroys as well
+    // as the other consuming uses).
+    BasicBlockWorklist worklist(currentDef->getFunction());
+    for (auto *consumingBlock : consumingBlocks) {
+      worklist.push(consumingBlock);
+    }
+
+    // Walk backwards from consuming blocks.
+    while (auto *block = worklist.pop()) {
+      originalLiveBlocks.insert(block);
+      for (auto *predecessor : block->getPredecessorBlocks()) {
+        // If the block was discovered by liveness, we already added it to the
+        // set.
+        if (originalLiveBlocks.contains(predecessor))
+          continue;
+        worklist.pushIfNotVisited(predecessor);
+      }
+    }
+  }
+
+  // Second, collect the blocks which occur after a _final_ consuming use.
+  BasicBlockSet consumedAtExitBlocks(currentDef->getFunction());
+  StackList<SILBasicBlock *> consumedAtEntryBlocks(currentDef->getFunction());
+  {
+    // Start the forward walk from blocks which contain _final_ non-destroy
+    // consumes. These are just the instructions on the boundary which aren't
+    // destroys.
+    BasicBlockWorklist worklist(currentDef->getFunction());
+    for (auto *instruction : boundary.lastUsers) {
+      if (dynCastToDestroyOf(instruction, getCurrentDef()))
+        continue;
+      if (liveness.isInterestingUser(instruction) !=
+          PrunedLiveness::IsInterestingUser::LifetimeEndingUse)
+        continue;
+      worklist.push(instruction->getParent());
+    }
+    while (auto *block = worklist.pop()) {
+      consumedAtExitBlocks.insert(block);
+      for (auto *successor : block->getSuccessorBlocks()) {
+        if (!originalLiveBlocks.contains(successor))
+          continue;
+        worklist.pushIfNotVisited(successor);
+        consumedAtEntryBlocks.push_back(successor);
+      }
+    }
+  }
+
+  // Third, find the blocks on the boundary between the originally-live blocks
+  // and the originally-live-but-consumed blocks.  These are new boundary
+  // blocks.
+  for (auto *block : consumedAtEntryBlocks) {
+    for (auto *predecessor : block->getPredecessorBlocks()) {
+      if (!consumedAtExitBlocks.contains(predecessor)) {
+        // Add "the instruction(s) before the terminator" of the predecessor to
+        // liveness.
+        if (auto *inst =
+                predecessor->getTerminator()->getPreviousInstruction()) {
+          liveness.updateForUse(inst, /*lifetimeEnding*/ false);
+        } else {
+          for (auto *grandPredecessor : predecessor->getPredecessorBlocks()) {
+            liveness.updateForUse(grandPredecessor->getTerminator(),
+                                  /*lifetimeEnding*/ false);
+          }
+        }
+      }
+    }
+  }
+
+  // Finally, preserve the destroys which weren't in the consumed region in
+  // place: hoisting such destroys would not avoid copies.
+  for (auto *destroy : destroys) {
+    auto *block = destroy->getParent();
+    // If the destroy is in a consumed block or a final consuming block,
+    // hoisting it would avoid a copy.
+    if (consumedAtExitBlocks.contains(block))
+      continue;
+    liveness.updateForUse(destroy, /*lifetimeEnding*/ true);
+  }
+}
+
 //===----------------------------------------------------------------------===//
 // MARK: Step 4. Rewrite copies and destroys
 //===----------------------------------------------------------------------===//
@@ -781,6 +911,11 @@ bool CanonicalizeOSSALifetime::canonicalizeValueLifetime(SILValue def) {
   // Step 2: compute boundary
   PrunedLivenessBoundary boundary;
   findExtendedBoundary(boundary);
+  if (maximizeLifetime) {
+    extendUnconsumedLiveness(boundary);
+    boundary.clear();
+    findExtendedBoundary(boundary);
+  }
   // Step 3: insert destroys and record consumes
   insertDestroysOnBoundary(boundary);
   // Step 4: rewrite copies and delete extra destroys