From 1f3f830fc7aa36b1adba00c251b9c24a202c2af2 Mon Sep 17 00:00:00 2001
From: Jamie <2119834+jamieQ@users.noreply.github.com>
Date: Tue, 17 Jun 2025 15:09:26 -0500
Subject: [PATCH] [SILOptimizer]: slow OSSA lifetime canonicalization
 mitigation

OSSA lifetime canonicalization can take a very long time in certain
cases in which there are large basic blocks. to mitigate this, add logic
to skip walking the liveness boundary for extending liveness to dead
ends when there aren't any dead ends in the function.

Updates `DeadEndBlocks` with a new `isEmpty` method and cache to
determine if there are any dead-end blocks in a given function.
---
 include/swift/SIL/BasicBlockUtils.h           | 15 +++++
 .../Utils/CanonicalizeOSSALifetime.h          |  4 ++
 lib/SIL/Utils/BasicBlockUtils.cpp             | 10 ++++
 .../Utils/CanonicalizeOSSALifetime.cpp        |  2 +-
 test/SILOptimizer/dead_end_blocks.sil         | 60 +++++++++++++++++++
 5 files changed, 90 insertions(+), 1 deletion(-)

diff --git a/include/swift/SIL/BasicBlockUtils.h b/include/swift/SIL/BasicBlockUtils.h
index d443b1f5abf03..1e15bcb78706d 100644
--- a/include/swift/SIL/BasicBlockUtils.h
+++ b/include/swift/SIL/BasicBlockUtils.h
@@ -68,6 +68,10 @@ class DeadEndBlocks {
   const SILFunction *f;
   bool didComputeValue = false;
 
+  /// When non-null, indicates whether dead-end blocks are present
+  /// in the current function.
+  std::optional<bool> hasAnyDeadEnds = std::nullopt;
+
   void compute();
 
 public:
@@ -85,6 +89,17 @@ class DeadEndBlocks {
     return reachableBlocks.count(block) == 0;
   }
 
+  /// Returns true iff none of the function's blocks is a dead-end.
+  /// Note: The underlying value is lazily computed & cached.
+  bool isEmpty() {
+    if (!hasAnyDeadEnds.has_value()) {
+      hasAnyDeadEnds = llvm::any_of(
+          *f, [this](const SILBasicBlock &BB) { return isDeadEnd(&BB); });
+    }
+
+    return !hasAnyDeadEnds.value();
+  }
+
   /// Return true if this dead end blocks has computed its internal cache yet.
   ///
   /// Used to determine if we need to verify a DeadEndBlocks.
diff --git a/include/swift/SILOptimizer/Utils/CanonicalizeOSSALifetime.h b/include/swift/SILOptimizer/Utils/CanonicalizeOSSALifetime.h
index ca7a6dd9a7f71..f686492a70e06 100644
--- a/include/swift/SILOptimizer/Utils/CanonicalizeOSSALifetime.h
+++ b/include/swift/SILOptimizer/Utils/CanonicalizeOSSALifetime.h
@@ -475,6 +475,10 @@ class CanonicalizeOSSALifetime final {
     return !endingLifetimeAtExplicitEnds();
   }
 
+  bool hasAnyDeadEnds() const {
+    return !deadEndBlocksAnalysis->get(function)->isEmpty();
+  }
+
   bool respectsDeinitBarriers() const {
     if (!currentDef->isLexical())
       return false;
diff --git a/lib/SIL/Utils/BasicBlockUtils.cpp b/lib/SIL/Utils/BasicBlockUtils.cpp
index c0653ca3e4035..b375c94d31a64 100644
--- a/lib/SIL/Utils/BasicBlockUtils.cpp
+++ b/lib/SIL/Utils/BasicBlockUtils.cpp
@@ -446,6 +446,16 @@ static FunctionTest DeadEndBlocksTest("dead_end_blocks", [](auto &function,
   }
 #endif
 });
+
+// Arguments:
+// - none
+// Dumps:
+// - message
+static FunctionTest HasAnyDeadEndBlocksTest(
+    "has_any_dead_ends", [](auto &function, auto &arguments, auto &test) {
+      auto deb = test.getDeadEndBlocks();
+      llvm::outs() << (deb->isEmpty() ? "no dead ends\n" : "has dead ends\n");
+    });
 } // end namespace swift::test
 
 //===----------------------------------------------------------------------===//
diff --git a/lib/SILOptimizer/Utils/CanonicalizeOSSALifetime.cpp b/lib/SILOptimizer/Utils/CanonicalizeOSSALifetime.cpp
index 2a9f209fd5f43..52f85ed36178c 100644
--- a/lib/SILOptimizer/Utils/CanonicalizeOSSALifetime.cpp
+++ b/lib/SILOptimizer/Utils/CanonicalizeOSSALifetime.cpp
@@ -1406,7 +1406,7 @@ bool CanonicalizeOSSALifetime::computeLiveness() {
     clear();
     return false;
   }
-  if (respectsDeadEnds()) {
+  if (respectsDeadEnds() && hasAnyDeadEnds()) {
     if (respectsDeinitBarriers()) {
       extendLexicalLivenessToDeadEnds();
     }
diff --git a/test/SILOptimizer/dead_end_blocks.sil b/test/SILOptimizer/dead_end_blocks.sil
index 6bc6e24935956..4332fa66370b7 100644
--- a/test/SILOptimizer/dead_end_blocks.sil
+++ b/test/SILOptimizer/dead_end_blocks.sil
@@ -47,4 +47,64 @@ exit:
   return %retval : $()
 }
 
+// no dead ends - simple return
+// CHECK-LABEL: begin running test {{.*}} on simple_function: has_any_dead_ends
+// CHECK:         no dead ends
+// CHECK-LABEL: end running test {{.*}} on simple_function: has_any_dead_ends
+sil @simple_function : $@convention(thin) () -> () {
+entry:
+  specify_test "has_any_dead_ends"
+  %retval = tuple ()
+  return %retval : $()
+}
+
+// dead ends - unreachable blocks
+// CHECK-LABEL: begin running test {{.*}} on function_with_dead_ends: has_any_dead_ends
+// CHECK:         has dead ends
+// CHECK-LABEL: end running test {{.*}} on function_with_dead_ends: has_any_dead_ends
+sil @function_with_dead_ends : $@convention(thin) () -> () {
+entry:
+  specify_test "has_any_dead_ends"
+  cond_br undef, die, exit
+
+die:
+  unreachable
+
+exit:
+  %retval = tuple ()
+  return %retval : $()
+}
 
+// dead ends – infinite loop
+// CHECK-LABEL: begin running test {{.*}} on function_with_loop: has_any_dead_ends
+// CHECK:         has dead ends
+// CHECK-LABEL: end running test {{.*}} on function_with_loop: has_any_dead_ends
+sil @function_with_loop : $@convention(thin) () -> () {
+entry:
+  specify_test "has_any_dead_ends"
+  cond_br undef, exit, loop
+
+loop:
+  br loop
+
+exit:
+  %retval = tuple ()
+  return %retval : $()
+}
+
+// no dead ends – conditional branches but all paths return
+// CHECK-LABEL: begin running test {{.*}} on branching_no_dead_ends: has_any_dead_ends
+// CHECK:         no dead ends
+// CHECK-LABEL: end running test {{.*}} on branching_no_dead_ends: has_any_dead_ends
+sil @branching_no_dead_ends : $@convention(thin) () -> () {
+entry:
+  specify_test "has_any_dead_ends"
+  cond_br undef, then, else
+
+then:
+  br else
+
+else:
+  %retval2 = tuple ()
+  return %retval2 : $()
+}