Merge pull request #73368 from nate-chandler/lifetime-completion/20240501/1

[NFC] LifetimeCompletion: Clarified API and documentation.

Author: nate-chandler

include/swift/SIL/OSSALifetimeCompletion.h

@@ -50,22 +50,54 @@ class OSSALifetimeCompletion {
   OSSALifetimeCompletion(SILFunction *function, const DominanceInfo *domInfo)
     : domInfo(domInfo), completedValues(function) {}
 
+  // The kind of boundary at which to complete the lifetime.
+  //
+  // Liveness: "As early as possible."  Consume the value after the last
+  //           non-consuming uses.
+  // Availability: "As late as possible."  Consume the value in the last blocks
+  //               beyond the non-consuming uses in which the value has been
+  //               consumed on no incoming paths.
+  struct Boundary {
+    enum Value : uint8_t {
+      Liveness,
+      Availability,
+    };
+    Value value;
+
+    Boundary(Value value) : value(value){};
+    operator Value() const { return value; }
+
+    static std::optional<Boundary> getForcingLiveness(bool force) {
+      if (!force)
+        return {};
+      return {Liveness};
+    }
+
+    bool isLiveness() { return value == Liveness; }
+    bool isAvailable() { return !isLiveness(); }
+  };
+
   /// Insert a lifetime-ending instruction on every path to complete the OSSA
-  /// lifetime of \p value. Lifetime completion is only relevant for owned
+  /// lifetime of \p value along \p boundary.
+  ///
+  /// If \p boundary is not specified, the following boundary will be used:
+  ///   \p value is lexical -> Boundary::Availability
+  ///   \p value is non-lexical -> Boundary::Liveness
+  ///
+  /// Lifetime completion is only relevant for owned
   /// values or borrow introducers.
-  /// For lexical values lifetime is completed at unreachable instructions.
-  /// For non-lexical values lifetime is completed at the lifetime boundary.
-  /// When \p forceBoundaryCompletion is true, the client is able to guarantee
-  /// that lifetime completion of lexical values at the lifetime boundary is
-  /// sufficient.
-  /// Currently \p forceBoundaryCompletion is used by mem2reg and temprvalueopt
-  /// to complete lexical enum values on trivial paths.
+  ///
+  /// Currently \p boundary == {Boundary::Availability} is used by Mem2Reg and
+  /// TempRValueOpt and PredicatbleMemOpt to complete lexical enum values on
+  /// trivial paths.
+  ///
   /// Returns true if any new instructions were created to complete the
   /// lifetime.
   ///
   /// TODO: We also need to complete scoped addresses (e.g. store_borrow)!
   LifetimeCompletion
-  completeOSSALifetime(SILValue value, bool forceBoundaryCompletion = false) {
+  completeOSSALifetime(SILValue value,
+                       std::optional<Boundary> maybeBoundary = std::nullopt) {
     if (value->getOwnershipKind() == OwnershipKind::None)
       return LifetimeCompletion::NoLifetime;
 
@@ -80,7 +112,10 @@ class OSSALifetimeCompletion {
     if (!completedValues.insert(value))
       return LifetimeCompletion::AlreadyComplete;
 
-    return analyzeAndUpdateLifetime(value, forceBoundaryCompletion)
+    Boundary boundary = maybeBoundary.value_or(
+        value->isLexical() ? Boundary::Availability : Boundary::Liveness);
+
+    return analyzeAndUpdateLifetime(value, boundary)
                ? LifetimeCompletion::WasCompleted
                : LifetimeCompletion::AlreadyComplete;
   }
@@ -90,7 +125,7 @@ class OSSALifetimeCompletion {
       llvm::function_ref<void(SILInstruction *)> visit);
 
 protected:
-  bool analyzeAndUpdateLifetime(SILValue value, bool forceBoundaryCompletion);
+  bool analyzeAndUpdateLifetime(SILValue value, Boundary boundary);
 };
 
 //===----------------------------------------------------------------------===//

lib/SIL/Utils/OSSALifetimeCompletion.cpp

@@ -70,8 +70,8 @@ static SILInstruction *endOSSALifetime(SILValue value, SILBuilder &builder) {
   return builder.createEndBorrow(loc, value);
 }
 
-static bool endLifetimeAtBoundary(SILValue value,
-                                  const SSAPrunedLiveness &liveness) {
+static bool endLifetimeAtLivenessBoundary(SILValue value,
+                                          const SSAPrunedLiveness &liveness) {
   PrunedLivenessBoundary boundary;
   liveness.computeBoundary(boundary);
 
@@ -149,7 +149,7 @@ class VisitUnreachableLifetimeEnds {
                             llvm::function_ref<void(SILInstruction *)> visit);
 
   struct State {
-    enum class Value : uint8_t {
+    enum Value : uint8_t {
       Unavailable = 0,
       Available,
       Unknown,
@@ -161,10 +161,6 @@ class VisitUnreachableLifetimeEnds {
     State meet(State const other) const {
       return *this < other ? *this : other;
     }
-
-    static State Unavailable() { return {Value::Unavailable}; }
-    static State Available() { return {Value::Available}; }
-    static State Unknown() { return {Value::Unknown}; }
   };
 
   struct Result {
@@ -197,36 +193,47 @@ class VisitUnreachableLifetimeEnds {
 
 void VisitUnreachableLifetimeEnds::computeRegion(
     const SSAPrunedLiveness &liveness) {
-  // Find the non-lifetime-ending boundary of `value`.
+  // (1) Compute the complete liveness boundary.
   PrunedLivenessBoundary boundary;
   liveness.computeBoundary(boundary);
 
+  // Used in the forward walk below (3).
+  BasicBlockWorklist regionWorklist(value->getFunction());
+
+  // (2) Collect the non-lifetime-ending liveness boundary.  This is the
+  //     portion of `boundary` consisting of:
+  // - non-lifetime-ending instructions (their parent blocks)
+  // - boundary edges
+  // - dead defs (their parent blocks)
+  auto collect = [&](SILBasicBlock *block) {
+    // `region` consists of the non-lifetime-ending boundary and all its
+    // iterative successors.
+    region.insert(block);
+    // `starts` just consists of the blocks in the non-lifetime-ending
+    // boundary.
+    starts.insert(block);
+    // The forward walk begins from the non-lifetime-ending boundary.
+    regionWorklist.push(block);
+  };
+
   for (SILInstruction *lastUser : boundary.lastUsers) {
     if (liveness.isInterestingUser(lastUser)
         != PrunedLiveness::LifetimeEndingUse) {
-      region.insert(lastUser->getParent());
-      starts.insert(lastUser->getParent());
+      collect(lastUser->getParent());
     }
   }
   for (SILBasicBlock *edge : boundary.boundaryEdges) {
-    region.insert(edge);
-    starts.insert(edge);
+    collect(edge);
   }
   for (SILNode *deadDef : boundary.deadDefs) {
-    region.insert(deadDef->getParentBlock());
-    starts.insert(deadDef->getParentBlock());
+    collect(deadDef->getParentBlock());
   }
 
-  // Forward walk to find the region in which `value` might be available.
-  BasicBlockWorklist regionWorklist(value->getFunction());
-  // Start the forward walk from the non-lifetime-ending boundary.
-  for (auto *start : region) {
-    regionWorklist.push(start);
-  }
+  // (3) Forward walk to find the region in which `value` might be available.
   while (auto *block = regionWorklist.pop()) {
     if (block->succ_empty()) {
-      // This assert will fail unless there are already lifetime-ending
-      // instruction on all paths to normal function exits.
+      // This assert will fail unless there is already a lifetime-ending
+      // instruction on each path to normal function exits.
       assert(isa<UnreachableInst>(block->getTerminator()));
     }
     for (auto *successor : block->getSuccessorBlocks()) {
@@ -244,9 +251,9 @@ void VisitUnreachableLifetimeEnds::propagateAvailablity(Result &result) {
   // - start blocks are ::Available
   for (auto *block : region) {
     if (starts.contains(block))
-      result.setState(block, State::Available());
+      result.setState(block, State::Available);
     else
-      result.setState(block, State::Unknown());
+      result.setState(block, State::Unknown);
   }
 
   BasicBlockWorklist worklist(value->getFunction());
@@ -280,14 +287,14 @@ void VisitUnreachableLifetimeEnds::propagateAvailablity(Result &result) {
 void VisitUnreachableLifetimeEnds::visitAvailabilityBoundary(
     Result const &result, llvm::function_ref<void(SILInstruction *)> visit) {
   for (auto *block : region) {
-    auto available = result.getState(block) == State::Available();
+    auto available = result.getState(block) == State::Available;
     if (!available) {
       continue;
     }
     auto hasUnreachableSuccessor = [&]() {
       // Use a lambda to avoid checking if possible.
       return llvm::any_of(block->getSuccessorBlocks(), [&result](auto *block) {
-        return result.getState(block) == State::Unavailable();
+        return result.getState(block) == State::Unavailable;
       });
     };
     if (!block->succ_empty() && !hasUnreachableSuccessor()) {
@@ -315,8 +322,9 @@ void OSSALifetimeCompletion::visitUnreachableLifetimeEnds(
   visitor.visitAvailabilityBoundary(result, visit);
 }
 
-static bool endLifetimeAtUnreachableBlocks(SILValue value,
-                                           const SSAPrunedLiveness &liveness) {
+static bool
+endLifetimeAtAvailabilityBoundary(SILValue value,
+                                  const SSAPrunedLiveness &liveness) {
   bool changed = false;
   OSSALifetimeCompletion::visitUnreachableLifetimeEnds(
       value, liveness, [&](auto *unreachable) {
@@ -330,12 +338,8 @@ static bool endLifetimeAtUnreachableBlocks(SILValue value,
 /// End the lifetime of \p value at unreachable instructions.
 ///
 /// Returns true if any new instructions were created to complete the lifetime.
-///
-/// This is only meant to cleanup lifetimes that lead to dead-end blocks. After
-/// recursively completing all nested scopes, it then simply ends the lifetime
-/// at the Unreachable instruction.
-bool OSSALifetimeCompletion::analyzeAndUpdateLifetime(
-    SILValue value, bool forceBoundaryCompletion) {
+bool OSSALifetimeCompletion::analyzeAndUpdateLifetime(SILValue value,
+                                                      Boundary boundary) {
   // Called for inner borrows, inner adjacent reborrows, inner reborrows, and
   // scoped addresses.
   auto handleInnerScope = [this](SILValue innerBorrowedValue) {
@@ -345,10 +349,13 @@ bool OSSALifetimeCompletion::analyzeAndUpdateLifetime(
   liveness.compute(domInfo, handleInnerScope);
 
   bool changed = false;
-  if (value->isLexical() && !forceBoundaryCompletion) {
-    changed |= endLifetimeAtUnreachableBlocks(value, liveness.getLiveness());
-  } else {
-    changed |= endLifetimeAtBoundary(value, liveness.getLiveness());
+  switch (boundary) {
+  case Boundary::Availability:
+    changed |= endLifetimeAtAvailabilityBoundary(value, liveness.getLiveness());
+    break;
+  case Boundary::Liveness:
+    changed |= endLifetimeAtLivenessBoundary(value, liveness.getLiveness());
+    break;
   }
   // TODO: Rebuild outer adjacent phis on demand (SILGen does not currently
   // produce guaranteed phis). See FindEnclosingDefs &
@@ -367,9 +374,15 @@ static FunctionTest OSSALifetimeCompletionTest(
     "ossa-lifetime-completion",
     [](auto &function, auto &arguments, auto &test) {
       SILValue value = arguments.takeValue();
+      std::optional<OSSALifetimeCompletion::Boundary> kind = std::nullopt;
+      if (arguments.hasUntaken()) {
+        kind = arguments.takeBool()
+                   ? OSSALifetimeCompletion::Boundary::Liveness
+                   : OSSALifetimeCompletion::Boundary::Availability;
+      }
       llvm::outs() << "OSSA lifetime completion: " << value;
       OSSALifetimeCompletion completion(&function, /*domInfo*/ nullptr);
-      completion.completeOSSALifetime(value);
+      completion.completeOSSALifetime(value, kind);
       function.print(llvm::outs());
     });
 } // end namespace swift::test

lib/SILOptimizer/Mandatory/PredictableMemOpt.cpp

@@ -2673,10 +2673,11 @@ bool AllocOptimize::tryToRemoveDeadAllocation() {
     // Lexical enums can have incomplete lifetimes in non payload paths that
     // don't end in unreachable. Force their lifetime to end immediately after
     // the last use instead.
-    bool forceBoundaryCompletion = v->getType().isOrHasEnum();
+    auto boundary = OSSALifetimeCompletion::Boundary::getForcingLiveness(
+        v->getType().isOrHasEnum());
     LLVM_DEBUG(llvm::dbgs() << "Completing lifetime of: ");
     LLVM_DEBUG(v->dump());
-    completion.completeOSSALifetime(v, forceBoundaryCompletion);
+    completion.completeOSSALifetime(v, boundary);
   }
 
   return true;

lib/SILOptimizer/Transforms/SILMem2Reg.cpp

@@ -1754,7 +1754,8 @@ void StackAllocationPromoter::run(BasicBlockSetVector &livePhiBlocks) {
   for (auto it : valuesToComplete) {
     // Set forceBoundaryCompletion as true so that we complete at boundary for
     // lexical values as well.
-    completion.completeOSSALifetime(it, /* forceBoundaryCompletion */ true);
+    completion.completeOSSALifetime(it,
+                                    OSSALifetimeCompletion::Boundary::Liveness);
   }
 }
 

lib/SILOptimizer/Transforms/TempRValueElimination.cpp

@@ -964,7 +964,8 @@ void TempRValueOptPass::run() {
   // Call the utlity to complete ossa lifetime.
   OSSALifetimeCompletion completion(function, da->get(function));
   for (auto it : valuesToComplete) {
-    completion.completeOSSALifetime(it, /* forceBoundaryCompletion */ true);
+    completion.completeOSSALifetime(it,
+                                    OSSALifetimeCompletion::Boundary::Liveness);
   }
 }