Merge pull request #73859 from nate-chandler/lifetime-completion/20240523/1

[LifetimeCompletion] Mark lifetimes in dead-end loops.

Author: nate-chandler

SwiftCompilerSources/Sources/Optimizer/Utilities/OwnershipLiveness.swift

@@ -64,8 +64,11 @@ func computeLinearLiveness(for definingValue: Value, _ context: Context)
   var range = InstructionRange(for: definingValue, context)
 
   // Compute liveness.
-  definingValue.lookThroughBorrowedFromUser.uses.endingLifetime.forEach {
-    range.insert($0.instruction)
+ for use in definingValue.lookThroughBorrowedFromUser.uses {
+    let instruction = use.instruction
+    if use.endsLifetime || instruction is ExtendLifetimeInst {
+      range.insert(instruction)
+    }
   }
   return range
 }

SwiftCompilerSources/Sources/SIL/Instruction.swift

@@ -514,6 +514,8 @@ final public class DestroyAddrInst : Instruction, UnaryInstruction {
 
 final public class EndLifetimeInst : Instruction, UnaryInstruction {}
 
+final public class ExtendLifetimeInst : Instruction, UnaryInstruction {}
+
 final public class InjectEnumAddrInst : Instruction, UnaryInstruction, EnumInstruction {
   public var `enum`: Value { operand.value }
   public var caseIndex: Int { bridged.InjectEnumAddrInst_caseIndex() }

SwiftCompilerSources/Sources/SIL/Registration.swift

@@ -78,6 +78,7 @@ public func registerSILClasses() {
   register(DestroyValueInst.self)
   register(DestroyAddrInst.self)
   register(EndLifetimeInst.self)
+  register(ExtendLifetimeInst.self)
   register(StrongCopyUnownedValueInst.self)
   register(StrongCopyUnmanagedValueInst.self)
   register(StrongCopyWeakValueInst.self)

docs/SIL.rst

@@ -4682,6 +4682,22 @@ The instruction accepts an object or address type.
 This instruction is valid only in OSSA and is lowered to a no-op when lowering
 to non-OSSA.
 
+extend_lifetime
+```````````````
+
+::
+
+   sil-instruction ::= 'extend_lifetime' sil-operand
+
+   // Indicate that %0's linear lifetime extends to this point
+   extend_lifetime %0 : $X
+
+Indicates that a value's linear lifetime extends to this point.  Inserted by
+OSSALifetimeCompletion(AvailabilityBoundary) in order to provide the invariant
+that a value is either consumed OR has an `extend_lifetime` user on all paths
+and furthermore that all uses are within the boundary defined by that set of
+instructions (the consumes and the `extend_lifetime`s).
+
 assign
 ``````
 ::

include/swift/SIL/OSSALifetimeCompletion.h

@@ -120,10 +120,17 @@ class OSSALifetimeCompletion {
     DoNotAllowLeaks = false,
   };
 
-  static void visitUnreachableLifetimeEnds(
+  enum class LifetimeEnd : uint8_t {
+    /// The lifetime ends at the boundary.
+    Boundary,
+    /// The lifetime "ends" in a loop.
+    Loop,
+  };
+
+  static void visitAvailabilityBoundary(
       SILValue value, AllowLeaks_t allowLeaks,
       const SSAPrunedLiveness &liveness,
-      llvm::function_ref<void(SILInstruction *)> visit);
+      llvm::function_ref<void(SILInstruction *, LifetimeEnd end)> visit);
 
 protected:
   bool analyzeAndUpdateLifetime(SILValue value, Boundary boundary);

include/swift/SIL/OwnershipLiveness.h

@@ -206,9 +206,19 @@ class LinearLiveness : public OSSALiveness {
   friend LinearLivenessVisitor;
 
 public:
-  LinearLiveness(SILValue def);
+  /// Whether extend_lifetime instructions should be added to the boundary.
+  /// Used to verify extend_lifetime instructions.
+  enum IncludeExtensions_t {
+    DoNotIncludeExtensions = false,
+    IncludeExtensions = true,
+  };
+  LinearLiveness(SILValue def,
+                 IncludeExtensions_t includeExtensions = IncludeExtensions);
 
   void compute();
+
+private:
+  const IncludeExtensions_t includeExtensions;
 };
 
 // Internal implementation

include/swift/SIL/OwnershipUseVisitor.h

@@ -213,6 +213,13 @@ bool OwnershipUseVisitor<Impl>::visitLifetimeEndingUses(SILValue ssaDef) {
 template <typename Impl>
 bool OwnershipUseVisitor<Impl>::visitConsumes(SILValue ssaDef) {
   for (Operand *use : ssaDef->getUses()) {
+    // extend_lifetime instructions are non-consuming but need to be visited
+    // because together with consuming uses they enclose all users of the value.
+    if (isa<ExtendLifetimeInst>(use->getUser())) {
+      if (!handleUsePoint(use, UseLifetimeConstraint::NonLifetimeEnding))
+        return false;
+      continue;
+    }
     if (use->isConsuming()) {
       if (PhiOperand(use) && !asImpl().handleOwnedPhi(use))
         return false;
@@ -245,6 +252,9 @@ bool OwnershipUseVisitor<Impl>::visitOuterBorrowScopeEnd(Operand *borrowEnd) {
 
     return handleUsePoint(borrowEnd, UseLifetimeConstraint::LifetimeEnding);
 
+  case OperandOwnership::InstantaneousUse:
+    assert(isa<ExtendLifetimeInst>(borrowEnd->getUser()));
+    return handleUsePoint(borrowEnd, UseLifetimeConstraint::NonLifetimeEnding);
   default:
     llvm_unreachable("expected borrow scope end");
   }

include/swift/SIL/PrunedLiveness.h

@@ -391,22 +391,20 @@ class PrunedLiveness {
       Ending,
       // The instruction doesn't use the value.
       NonUse,
-    };
-    Value value;
+    } value;
 
     LifetimeEnding(Value value) : value(value) {}
-    explicit LifetimeEnding(bool lifetimeEnding)
-        : value(lifetimeEnding ? Value::Ending : Value::NonEnding) {}
     operator Value() const { return value; }
+
+    static LifetimeEnding forUse(bool lifetimeEnding) {
+      return lifetimeEnding ? Value::Ending : Value::NonEnding;
+    }
+    bool isEnding() const { return value == Value::Ending; }
+
     LifetimeEnding meet(LifetimeEnding const other) const {
       return value < other.value ? *this : other;
     }
     void meetInPlace(LifetimeEnding const other) { *this = meet(other); }
-    bool isEnding() const { return value == Value::Ending; }
-
-    static LifetimeEnding NonUse() { return {Value::NonUse}; };
-    static LifetimeEnding Ending() { return {Value::Ending}; };
-    static LifetimeEnding NonEnding() { return {Value::NonEnding}; };
   };
 
 protected:
@@ -467,8 +465,8 @@ class PrunedLiveness {
     auto useIter = users.find(user);
     if (useIter == users.end())
       return NonUser;
-    return useIter->second == LifetimeEnding::Ending() ? LifetimeEndingUse
-                                                       : NonLifetimeEndingUse;
+    return useIter->second.isEnding() ? LifetimeEndingUse
+                                      : NonLifetimeEndingUse;
   }
 
   using ConstUserRange =

include/swift/SIL/SILBuilder.h

@@ -2292,6 +2292,11 @@ class SILBuilder {
                       EndLifetimeInst(getSILDebugLocation(Loc), Operand));
   }
 
+  ExtendLifetimeInst *createExtendLifetime(SILLocation Loc, SILValue Operand) {
+    return insert(new (getModule())
+                      ExtendLifetimeInst(getSILDebugLocation(Loc), Operand));
+  }
+
   UncheckedOwnershipConversionInst *
   createUncheckedOwnershipConversion(SILLocation Loc, SILValue Operand,
                                      ValueOwnershipKind Kind) {

include/swift/SIL/SILCloner.h

@@ -3019,6 +3019,18 @@ void SILCloner<ImplClass>::visitEndLifetimeInst(EndLifetimeInst *Inst) {
                                            getOpValue(Inst->getOperand())));
 }
 
+template <typename ImplClass>
+void SILCloner<ImplClass>::visitExtendLifetimeInst(ExtendLifetimeInst *Inst) {
+  getBuilder().setCurrentDebugScope(getOpScope(Inst->getDebugScope()));
+
+  if (!getBuilder().hasOwnership())
+    return;
+
+  recordClonedInstruction(
+      Inst, getBuilder().createExtendLifetime(getOpLocation(Inst->getLoc()),
+                                              getOpValue(Inst->getOperand())));
+}
+
 template <typename ImplClass>
 void SILCloner<ImplClass>::visitUncheckedOwnershipConversionInst(
     UncheckedOwnershipConversionInst *Inst) {