[PrunedLiveness] Add extendToNonUse.

nate-chandler · nate-chandler · commit c1716c9dfef0 · 2023-09-26T11:45:22.000-07:00
And use it in lifetime extension/maximization.

The new member function differs from updateForUse in that it doesn't
overwrite the old value for lifetime ending associated with the
instruction (calling updateForUse with lifetimeEnding=false overwrites
the flag set by a previous call with lifetimeEnding=true because if an
instruction both consumes and doesn't consume a copy-extended value, the
value must be live after the instruction).
diff --git a/include/swift/SIL/PrunedLiveness.h b/include/swift/SIL/PrunedLiveness.h
@@ -353,6 +353,61 @@ struct LiveRangeSummary {
 /// necessarily include liveness up to destroy_value or end_borrow
 /// instructions.
 class PrunedLiveness {
+public:
+  /// A tristate describing how an instruction uses the value:
+  /// - non-ending: the value's lifetime does not end at the instruction;
+  ///               results from updateForUse(, lifetimeEnding: false)
+  /// - ending: the value's lifetime ends at the instruction; results from
+  ///           updateForUse(, lifetimeEnding: true)
+  /// - non-use: the instruction doesn't use the value but liveness was extended
+  ///            to it; results from extendToNonUse
+  ///
+  /// If an instruction is added to liveness with multiple LifetimeEnding
+  /// instances, the stored instance needs to be updated appropriately, taking
+  /// into account the instances that were already seen.
+  ///
+  /// For example, if Nonuse is seen after Ending, it is ignored: that the
+  /// instruction ends the lifetime takes priority: the instruction doesn't end
+  /// the lifetime of the value any less because liveness was extended to it.
+  ///
+  /// Similarly, if Ending is seen after NonEnding, it is ignored: that the
+  /// instruction ends the lifetime of a copy of the value can't change the fact
+  /// that liveness already extends _beyond_ the instruction because it was
+  /// already recognized as a non-ending use.
+  ///
+  /// This relationship of "overriding" is captured by the order of the cases in
+  /// LifetimeEnding::Value and which case overrides another is computed by
+  /// taking the meet--i.e. the lesser of the two cases overrides.
+  ///
+  /// Note: Taking the meet may not be appropriate for branch instructions
+  ///       which may need to be recognized as lifetime-ending when they have
+  ///       as uses both a reborrow and a guaranteed phi.
+  struct LifetimeEnding {
+    enum class Value {
+      // The instruction doesn't consume the value.
+      NonEnding,
+      // The instruction consumes the value.
+      Ending,
+      // The instruction doesn't use the value.
+      NonUse,
+    };
+    Value value;
+
+    LifetimeEnding(Value value) : value(value) {}
+    explicit LifetimeEnding(bool lifetimeEnding)
+        : value(lifetimeEnding ? Value::Ending : Value::NonEnding) {}
+    operator Value() const { return value; }
+    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:
   PrunedLiveBlocks liveBlocks;
 
@@ -365,7 +420,7 @@ class PrunedLiveness {
   // they may be the last use in the block.
   //
   // Non-lifetime-ending within a LiveOut block are uninteresting.
-  llvm::SmallMapVector<SILInstruction *, bool, 8> users;
+  llvm::SmallMapVector<SILInstruction *, LifetimeEnding, 8> users;
 
 public:
   PrunedLiveness(SILFunction *function,
@@ -411,11 +466,12 @@ class PrunedLiveness {
     auto useIter = users.find(user);
     if (useIter == users.end())
       return NonUser;
-    return useIter->second ? LifetimeEndingUse : NonLifetimeEndingUse;
+    return useIter->second == LifetimeEnding::Ending() ? LifetimeEndingUse
+                                                       : NonLifetimeEndingUse;
   }
 
   using ConstUserRange =
-      iterator_range<const std::pair<SILInstruction *, bool> *>;
+      iterator_range<const std::pair<SILInstruction *, LifetimeEnding> *>;
   ConstUserRange getAllUsers() const {
     return llvm::make_range(users.begin(), users.end());
   }
@@ -425,49 +481,53 @@ class PrunedLiveness {
   /// IDE.
   struct RangeIterationHelpers {
     struct MapFunctor {
-      SILInstruction *
-      operator()(const std::pair<SILInstruction *, bool> &pair) const {
+      SILInstruction *operator()(
+          const std::pair<SILInstruction *, LifetimeEnding> &pair) const {
         // Strip off the const to ease use with other APIs.
         return const_cast<SILInstruction *>(pair.first);
       }
     };
 
-    struct LifetimeEnding {
+    struct IsLifetimeEnding {
       struct FilterFunctor {
-        bool operator()(const std::pair<SILInstruction *, bool> &pair) const {
-          return pair.second;
+        bool operator()(
+            const std::pair<SILInstruction *, LifetimeEnding> &pair) const {
+          return pair.second.isEnding();
         }
       };
 
       using MapFilterIter = llvm::mapped_iterator<
-          llvm::filter_iterator<const std::pair<SILInstruction *, bool> *,
-                                FilterFunctor>,
+          llvm::filter_iterator<
+              const std::pair<SILInstruction *, LifetimeEnding> *,
+              FilterFunctor>,
           MapFunctor>;
     };
 
     struct NonLifetimeEnding {
       struct FilterFunctor {
-        bool operator()(const std::pair<SILInstruction *, bool> &pair) const {
-          return !pair.second;
+        bool operator()(
+            const std::pair<SILInstruction *, LifetimeEnding> &pair) const {
+          return !pair.second.isEnding();
         }
       };
 
       using MapFilterIter = llvm::mapped_iterator<
-          llvm::filter_iterator<const std::pair<SILInstruction *, bool> *,
-                                FilterFunctor>,
+          llvm::filter_iterator<
+              const std::pair<SILInstruction *, LifetimeEnding> *,
+              FilterFunctor>,
           MapFunctor>;
     };
   };
   using LifetimeEndingUserRange = llvm::iterator_range<
-      RangeIterationHelpers::LifetimeEnding::MapFilterIter>;
+      RangeIterationHelpers::IsLifetimeEnding::MapFilterIter>;
 
   /// Return a range consisting of the current set of consuming users fed into
   /// this PrunedLiveness instance.
   LifetimeEndingUserRange getLifetimeEndingUsers() const {
     return map_range(
         llvm::make_filter_range(
             getAllUsers(),
-            RangeIterationHelpers::LifetimeEnding::FilterFunctor()),
+            RangeIterationHelpers::IsLifetimeEnding::FilterFunctor()),
         RangeIterationHelpers::MapFunctor());
   }
 
@@ -552,6 +612,8 @@ class PrunedLiveRange : public PrunedLiveness {
                                            ValueSet &visited,
                                            SILValue value);
 
+  void updateForUse(SILInstruction *user, LifetimeEnding lifetimeEnding);
+
 public:
   /// For flexibility, \p lifetimeEnding is provided by the
   /// caller. PrunedLiveness makes no assumptions about the def-use
@@ -560,6 +622,12 @@ class PrunedLiveRange : public PrunedLiveness {
   /// live range vs. a nested borrow scope within the extended live range.
   void updateForUse(SILInstruction *user, bool lifetimeEnding);
 
+  /// Adds \p inst which doesn't use the def to liveness.
+  ///
+  /// Different from calling updateForUse because it never overrides the value
+  /// \p lifetimeEnding stored for \p inst.
+  void extendToNonUse(SILInstruction *inst);
+
   /// Updates the liveness for a whole borrow scope, beginning at \p op.
   /// Returns false if this cannot be done. This assumes that nested OSSA
   /// lifetimes are complete.
diff --git a/lib/SIL/Utils/PrunedLiveness.cpp b/lib/SIL/Utils/PrunedLiveness.cpp
@@ -85,10 +85,17 @@ void PrunedLiveBlocks::dump() const {
 void PrunedLiveness::print(llvm::raw_ostream &OS) const {
   liveBlocks.print(OS);
   for (auto &userAndIsLifetimeEnding : users) {
-    if (userAndIsLifetimeEnding.second)
+    switch (userAndIsLifetimeEnding.second) {
+    case LifetimeEnding::Value::NonUse:
+      OS << "non-user: ";
+      break;
+    case LifetimeEnding::Value::Ending:
       OS << "lifetime-ending user: ";
-    else
+      break;
+    case LifetimeEnding::Value::NonEnding:
       OS << "regular user: ";
+      break;
+    }
     userAndIsLifetimeEnding.first->print(OS);
   }
 }
@@ -185,8 +192,9 @@ static FunctionTest
 //===----------------------------------------------------------------------===//
 
 template <typename LivenessWithDefs>
-void PrunedLiveRange<LivenessWithDefs>::updateForUse(SILInstruction *user,
-                                                     bool lifetimeEnding) {
+void PrunedLiveRange<LivenessWithDefs>::updateForUse(
+    SILInstruction *user,
+    PrunedLiveRange<LivenessWithDefs>::LifetimeEnding lifetimeEnding) {
   liveBlocks.updateForUse(user, asImpl().isUserBeforeDef(user));
 
   // Note that a user may use the current value from multiple operands. If any
@@ -202,7 +210,17 @@ void PrunedLiveRange<LivenessWithDefs>::updateForUse(SILInstruction *user,
   // argument must be copied.
   auto iterAndSuccess = users.insert({user, lifetimeEnding});
   if (!iterAndSuccess.second)
-    iterAndSuccess.first->second &= lifetimeEnding;
+    iterAndSuccess.first->second.meetInPlace(lifetimeEnding);
+}
+template <typename LivenessWithDefs>
+void PrunedLiveRange<LivenessWithDefs>::updateForUse(SILInstruction *user,
+                                                     bool lifetimeEnding) {
+  updateForUse(user, LifetimeEnding(lifetimeEnding));
+}
+
+template <typename LivenessWithDefs>
+void PrunedLiveRange<LivenessWithDefs>::extendToNonUse(SILInstruction *inst) {
+  updateForUse(inst, LifetimeEnding::NonUse());
 }
 
 template <typename LivenessWithDefs>
@@ -257,7 +275,7 @@ template <typename LivenessWithDefs>
 void PrunedLiveRange<LivenessWithDefs>::extendAcrossLiveness(
     PrunedLiveness &otherLiveness) {
   // update this liveness for all the interesting users in otherLiveness.
-  for (std::pair<SILInstruction *, bool> userAndEnd :
+  for (std::pair<SILInstruction *, LifetimeEnding> userAndEnd :
        otherLiveness.getAllUsers()) {
     updateForUse(userAndEnd.first, userAndEnd.second);
   }
@@ -371,6 +389,64 @@ static FunctionTest SSALivenessTest("ssa_liveness", [](auto &function,
   liveness.computeBoundary(boundary);
   boundary.print(llvm::outs());
 });
+
+// Arguments:
+// - SILValue: def whose pruned liveness will be calculated
+// - the string "uses:"
+// - variadic list of live-range user instructions
+// Dumps:
+// -
+static FunctionTest SSAUseLivenessTest("ssa_use_liveness", [](auto &function,
+                                                              auto &arguments,
+                                                              auto &test) {
+  auto value = arguments.takeValue();
+  SmallVector<SILBasicBlock *, 8> discoveredBlocks;
+  SSAPrunedLiveness liveness(&function, &discoveredBlocks);
+  liveness.initializeDef(value);
+
+  auto argument = arguments.takeArgument();
+  if (cast<StringArgument>(argument).getValue() != "uses:") {
+    llvm::report_fatal_error("test specification expects the 'uses:' label\n");
+  }
+
+  while (arguments.hasUntaken()) {
+    auto *inst = arguments.takeInstruction();
+    auto kindString = arguments.takeString();
+    enum Kind {
+      NonUse,
+      Ending,
+      NonEnding,
+    };
+    auto kind = llvm::StringSwitch<llvm::Optional<Kind>>(kindString)
+                    .Case("non-use", Kind::NonUse)
+                    .Case("ending", Kind::Ending)
+                    .Case("non-ending", Kind::NonEnding)
+                    .Default(llvm::None);
+    if (!kind.has_value()) {
+      llvm::errs() << "Unknown kind: " << kindString << "\n";
+      llvm::report_fatal_error("Bad user kind.  Value must be one of "
+                               "'non-use', 'ending', 'non-ending'");
+    }
+    switch (kind.value()) {
+    case Kind::NonUse:
+      liveness.extendToNonUse(inst);
+      break;
+    case Kind::Ending:
+      liveness.updateForUse(inst, /*lifetimeEnding*/ true);
+      break;
+    case Kind::NonEnding:
+      liveness.updateForUse(inst, /*lifetimeEnding*/ false);
+      break;
+    }
+  }
+
+  liveness.print(llvm::outs());
+
+  PrunedLivenessBoundary boundary;
+  liveness.computeBoundary(boundary);
+  boundary.print(llvm::outs());
+});
+
 } // end namespace swift::test
 
 template <typename LivenessWithDefs>
diff --git a/lib/SILOptimizer/Utils/CanonicalizeOSSALifetime.cpp b/lib/SILOptimizer/Utils/CanonicalizeOSSALifetime.cpp
@@ -262,21 +262,10 @@ void CanonicalizeOSSALifetime::extendLivenessToDeinitBarriers() {
   OSSALifetimeCompletion::visitUnreachableLifetimeEnds(
       getCurrentDef(), completeLiveness, [&](auto *unreachable) {
         recordConsumingUser(unreachable);
-        if (auto *previous = unreachable->getPreviousInstruction()) {
-          if (liveness->isInterestingUser(previous) ==
-              PrunedLiveness::IsInterestingUser::NonUser) {
-            liveness->updateForUse(previous, /*lifetimeEnding=*/false);
-          }
-          return;
-        }
-        for (auto *predecessor :
-             unreachable->getParent()->getPredecessorBlocks()) {
-          auto *previous = &predecessor->back();
-          if (liveness->isInterestingUser(previous) ==
-              PrunedLiveness::IsInterestingUser::NonUser) {
-            liveness->updateForUse(previous, /*lifetimeEnding=*/false);
-          }
-        }
+        unreachable->visitPriorInstructions([&](auto *inst) {
+          liveness->extendToNonUse(inst);
+          return true;
+        });
       });
 
   auto *def = getCurrentDef()->getDefiningInstruction();
@@ -648,14 +637,10 @@ void CanonicalizeOSSALifetime::extendUnconsumedLiveness(
         continue;
       // 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);
-        }
-      }
+      predecessor->getTerminator()->visitPriorInstructions([&](auto *inst) {
+        liveness->extendToNonUse(inst);
+        return true;
+      });
     }
   }
 
diff --git a/test/SILOptimizer/pruned_liveness_boundary.sil b/test/SILOptimizer/pruned_liveness_boundary.sil
