Merge pull request swiftlang#76095 from nate-chandler/lifetime-completion/20240826/1

[LifetimeCompletion] Avoid instruction list walk.

Author: nate-chandler

include/swift/AST/SILOptions.h

@@ -188,10 +188,10 @@ class SILOptions {
   /// If set to true, compile with the SIL Opaque Values enabled.
   bool EnableSILOpaqueValues = false;
 
-  /// Require linear OSSA lifetimes after SILGen
-  bool OSSACompleteLifetimes = false;
+  /// Introduce linear OSSA lifetimes after SILGen
+  bool OSSACompleteLifetimes = true;
 
-  /// Verify linear OSSA lifetimes after SILGen
+  /// Verify linear OSSA lifetimes throughout OSSA pipeline.
   bool OSSAVerifyComplete = false;
 
   /// Enable pack metadata stack "promotion".

include/swift/SIL/PrunedLiveness.h

@@ -543,10 +543,37 @@ class PrunedLiveness {
         RangeIterationHelpers::MapFunctor());
   }
 
+  void visitUsers(llvm::function_ref<void(SILInstruction *, LifetimeEnding)>
+                      visitor) const {
+    for (auto &pair : users) {
+      visitor(pair.first, pair.second);
+    }
+  }
+
   void print(llvm::raw_ostream &OS) const;
   void dump() const;
 };
 
+/// Recording liveness boundary at some level of detail; see concrete subclasses
+/// PrunedLivenessBoundary and PrunedLivenessBlockBoundary.
+struct AnyPrunedLivenessBoundary {
+  virtual ~AnyPrunedLivenessBoundary() {}
+  /// Targets whose single predecessor has at least one non-boundary successor.
+  SmallVector<SILBasicBlock *, 8> boundaryEdges;
+
+  friend class SSAPrunedLiveness;
+  friend class MultiDefPrunedLiveness;
+
+private:
+  virtual void findBoundaryInSSADefBlock(SILNode *ssaDef,
+                                         const PrunedLiveness &liveness) = 0;
+  virtual void
+  findBoundaryInMultiDefBlock(SILBasicBlock *block, bool isLiveOut,
+                              const MultiDefPrunedLiveness &liveness) = 0;
+  virtual void findBoundaryInNonDefBlock(SILBasicBlock *block,
+                                         const PrunedLiveness &liveness) = 0;
+};
+
 /// Record the last use points and CFG edges that form the boundary of
 /// PrunedLiveness.
 ///
@@ -558,9 +585,8 @@ class PrunedLiveness {
 /// Each boundary edge is identified by its target block. The source of the edge
 /// is the target block's single predecessor which must have at least one other
 /// non-boundary successor.
-struct PrunedLivenessBoundary {
+struct PrunedLivenessBoundary : AnyPrunedLivenessBoundary {
   SmallVector<SILInstruction *, 8> lastUsers;
-  SmallVector<SILBasicBlock *, 8> boundaryEdges;
   SmallVector<SILNode *, 1> deadDefs;
 
   void clear() {
@@ -579,6 +605,39 @@ struct PrunedLivenessBoundary {
 
   void print(llvm::raw_ostream &OS) const;
   void dump() const;
+
+private:
+  void findBoundaryInSSADefBlock(SILNode *ssaDef,
+                                 const PrunedLiveness &liveness) override;
+  void
+  findBoundaryInMultiDefBlock(SILBasicBlock *block, bool isLiveOut,
+                              const MultiDefPrunedLiveness &liveness) override;
+  void findBoundaryInNonDefBlock(SILBasicBlock *block,
+                                 const PrunedLiveness &liveness) override;
+};
+
+/// Record the blocks which either contain last use points or are boundary edge
+/// targets.
+///
+/// Enables clients only interested in block-level details to avoid expensive
+/// and for-them wasteful instruction list iteration.
+struct PrunedLivenessBlockBoundary : AnyPrunedLivenessBoundary {
+  /// Blocks containing last users or dead defs.
+  SmallVector<SILBasicBlock *, 8> endBlocks;
+
+  void clear() {
+    endBlocks.clear();
+    boundaryEdges.clear();
+  }
+
+private:
+  void findBoundaryInSSADefBlock(SILNode *ssaDef,
+                                 const PrunedLiveness &liveness) override;
+  void
+  findBoundaryInMultiDefBlock(SILBasicBlock *block, bool isLiveOut,
+                              const MultiDefPrunedLiveness &liveness) override;
+  void findBoundaryInNonDefBlock(SILBasicBlock *block,
+                                 const PrunedLiveness &liveness) override;
 };
 
 /// PrunedLiveness with information about defs for computing the live range
@@ -689,7 +748,7 @@ class PrunedLiveRange : public PrunedLiveness {
   /// The computed boundary will completely post-dominate, including dead end
   /// paths. The client should query DeadEndBlocks to ignore those dead end
   /// paths.
-  void computeBoundary(PrunedLivenessBoundary &boundary) const;
+  void computeBoundary(AnyPrunedLivenessBoundary &boundary) const;
 
   /// Compute the boundary from a backward CFG traversal from a known set of
   /// jointly post-dominating blocks. Avoids the need to record an ordered list
@@ -765,7 +824,7 @@ class SSAPrunedLiveness : public PrunedLiveRange<SSAPrunedLiveness> {
 
   /// SSA implementation of computeBoundary.
   void findBoundariesInBlock(SILBasicBlock *block, bool isLiveOut,
-                             PrunedLivenessBoundary &boundary) const;
+                             AnyPrunedLivenessBoundary &boundary) const;
 
   /// Compute liveness for a single SSA definition. The lifetime-ending uses are
   /// also recorded--destroy_value or end_borrow.
@@ -853,9 +912,9 @@ class MultiDefPrunedLiveness : public PrunedLiveRange<MultiDefPrunedLiveness> {
 
   /// Multi-Def implementation of computeBoundary.
   void findBoundariesInBlock(SILBasicBlock *block, bool isLiveOut,
-                             PrunedLivenessBoundary &boundary) const;
+                             AnyPrunedLivenessBoundary &boundary) const;
 
-  /// Compute liveness for a all currently initialized definitions. The
+  /// Compute liveness for all currently initialized definitions. The
   /// lifetime-ending uses are also recorded--destroy_value or
   /// end_borrow. However destroy_values might not jointly-post dominate if
   /// dead-end blocks are present.
@@ -872,6 +931,9 @@ class MultiDefPrunedLiveness : public PrunedLiveRange<MultiDefPrunedLiveness> {
   /// also lack scope-ending instructions, so the liveness of their nested uses
   /// may be ignored.
   LiveRangeSummary computeSimple();
+
+  friend struct PrunedLivenessBoundary;
+  friend struct PrunedLivenessBlockBoundary;
 };
 
 //===----------------------------------------------------------------------===//

lib/SIL/Utils/OSSALifetimeCompletion.cpp

@@ -262,9 +262,18 @@ void AvailabilityBoundaryVisitor::visit(const SSAPrunedLiveness &liveness,
 void AvailabilityBoundaryVisitor::computeRegion(
     const SSAPrunedLiveness &liveness) {
   // (1) Compute the complete liveness boundary.
-  PrunedLivenessBoundary boundary;
+  PrunedLivenessBlockBoundary boundary;
   liveness.computeBoundary(boundary);
 
+  BasicBlockSet consumingBlocks(value->getFunction());
+
+  liveness.visitUsers(
+      [&consumingBlocks](auto *instruction, auto lifetimeEnding) {
+        if (lifetimeEnding.isEnding()) {
+          consumingBlocks.insert(instruction->getParent());
+        }
+      });
+
   // Used in the forward walk below (3).
   BasicBlockWorklist regionWorklist(value->getFunction());
 
@@ -284,18 +293,14 @@ void AvailabilityBoundaryVisitor::computeRegion(
     regionWorklist.push(block);
   };
 
-  for (SILInstruction *lastUser : boundary.lastUsers) {
-    if (liveness.isInterestingUser(lastUser)
-        != PrunedLiveness::LifetimeEndingUse) {
-      collect(lastUser->getParent());
+  for (auto *endBlock : boundary.endBlocks) {
+    if (!consumingBlocks.contains(endBlock)) {
+      collect(endBlock);
     }
   }
   for (SILBasicBlock *edge : boundary.boundaryEdges) {
     collect(edge);
   }
-  for (SILNode *deadDef : boundary.deadDefs) {
-    collect(deadDef->getParentBlock());
-  }
 
   // (3) Forward walk to find the region in which `value` might be available.
   while (auto *block = regionWorklist.pop()) {