Merge pull request #30087 from gottesmm/pr-7d5f0d45e4677ebbd92aeb1e261ab388a4ca5756

[ownership] Add simple support for concatenating together simple live ranges

Author: gottesmm

include/swift/SIL/OwnershipUtils.h

@@ -218,6 +218,10 @@ bool isGuaranteedForwardingValue(SILValue value);
 /// forward guaranteed ownership.
 bool isOwnedForwardingValueKind(SILNodeKind kind);
 
+/// Does this SILInstruction 'forward' owned ownership, but may not be able to
+/// forward guaranteed ownership.
+bool isOwnedForwardingInstruction(SILInstruction *inst);
+
 struct BorrowScopeOperandKind {
   enum Kind {
     BeginBorrow,

include/swift/SIL/SILValue.h

@@ -272,9 +272,22 @@ class ValueBase : public SILNode, public SILAllocated<ValueBase> {
   /// otherwise.
   inline Operand *getSingleUse() const;
 
+  /// Returns .some(single user) if this value is non-trivial, we are in ossa,
+  /// and it has a single consuming user. Returns .none otherwise.
+  inline Operand *getSingleConsumingUse() const;
+
   template <class T>
   inline T *getSingleUserOfType() const;
 
+  template <class T> inline T *getSingleConsumingUserOfType() const;
+
+  /// Returns true if this operand has exactly two.
+  ///
+  /// This is useful if one has found a predefined set of 2 unique users and
+  /// wants to check if there are any other users without iterating over the
+  /// entire use list.
+  inline bool hasTwoUses() const;
+
   /// Helper struct for DowncastUserFilterRange
   struct UseToUser;
 
@@ -647,6 +660,10 @@ class Operand {
   /// Returns true if this operand acts as a use that consumes its associated
   /// value.
   bool isConsumingUse() const {
+    // Type dependent uses can never be consuming and do not have valid
+    // ownership maps since they do not participate in the ownership system.
+    if (isTypeDependent())
+      return false;
     auto map = getOwnershipKindMap();
     auto constraint = map.getLifetimeConstraint(get().getOwnershipKind());
     return constraint == UseLifetimeConstraint::MustBeInvalidated;
@@ -748,17 +765,48 @@ inline Operand *ValueBase::getSingleUse() const {
   return Op;
 }
 
+inline Operand *ValueBase::getSingleConsumingUse() const {
+  Operand *result = nullptr;
+  for (auto *op : getUses()) {
+    if (op->isConsumingUse()) {
+      if (result) {
+        return nullptr;
+      }
+      result = op;
+    }
+  }
+  return result;
+}
+
+inline bool ValueBase::hasTwoUses() const {
+  auto iter = use_begin(), end = use_end();
+  for (unsigned i = 0; i < 2; ++i) {
+    if (iter == end)
+      return false;
+    ++iter;
+  }
+  return iter == end;
+}
+
 template <class T>
 inline T *ValueBase::getSingleUserOfType() const {
-  T *Result = nullptr;
-  for (auto *Op : getUses()) {
-    if (auto *Tmp = dyn_cast<T>(Op->getUser())) {
-      if (Result)
+  T *result = nullptr;
+  for (auto *op : getUses()) {
+    if (auto *tmp = dyn_cast<T>(op->getUser())) {
+      if (result)
         return nullptr;
-      Result = Tmp;
+      result = tmp;
     }
   }
-  return Result;
+  return result;
+}
+
+template <class T> inline T *ValueBase::getSingleConsumingUserOfType() const {
+  auto *op = getSingleConsumingUse();
+  if (!op)
+    return nullptr;
+
+  return dyn_cast<T>(op->getUser());
 }
 
 struct ValueBase::UseToUser {

lib/SIL/OwnershipUtils.cpp

@@ -73,6 +73,16 @@ bool swift::isOwnedForwardingValueKind(SILNodeKind kind) {
   }
 }
 
+bool swift::isOwnedForwardingInstruction(SILInstruction *inst) {
+  auto kind = inst->getKind();
+  switch (kind) {
+  case SILInstructionKind::BranchInst:
+    return true;
+  default:
+    return isOwnershipForwardingValueKind(SILNodeKind(kind));
+  }
+}
+
 bool swift::isGuaranteedForwardingValue(SILValue value) {
   // If we have an argument from a transforming terminator, we can forward
   // guaranteed.

lib/SILOptimizer/Transforms/SemanticARCOpts.cpp

@@ -59,7 +59,6 @@ class LiveRange {
 
 public:
   LiveRange(SILValue value);
-
   LiveRange(const LiveRange &) = delete;
   LiveRange &operator=(const LiveRange &) = delete;
 
@@ -454,6 +453,7 @@ struct SemanticARCOptVisitor
 
   bool performGuaranteedCopyValueOptimization(CopyValueInst *cvi);
   bool eliminateDeadLiveRangeCopyValue(CopyValueInst *cvi);
+  bool tryJoiningCopyValueLiveRangeWithOperand(CopyValueInst *cvi);
 };
 
 } // end anonymous namespace
@@ -824,13 +824,184 @@ bool SemanticARCOptVisitor::eliminateDeadLiveRangeCopyValue(CopyValueInst *cvi)
   return true;
 }
 
+// Handle simple checking where we do not need to form live ranges and visit a
+// bunch of instructions.
+static bool canSafelyJoinSimpleRange(SILValue cviOperand,
+                                     DestroyValueInst *cviOperandDestroy,
+                                     CopyValueInst *cvi) {
+  // We only handle cases where our copy_value has a single consuming use that
+  // is not a forwarding use. We need to use the LiveRange functionality to
+  // guarantee correctness in the presence of forwarding uses.
+  //
+  // NOTE: This use may be any type of consuming use and may not be a
+  // destroy_value.
+  auto *cviConsumer = cvi->getSingleConsumingUse();
+  if (!cviConsumer || isOwnedForwardingInstruction(cviConsumer->getUser())) {
+    return false;
+  }
+
+  // Ok, we may be able to eliminate this. The main thing we need to be careful
+  // of here is that if the destroy_value is /after/ the consuming use of the
+  // operand of copy_value, we may have normal uses of the copy_value's operand
+  // that would become use-after-frees since we would be shrinking the lifetime
+  // of the object potentially. Consider the following SIL:
+  //
+  //   %0 = ...
+  //   %1 = copy_value %0
+  //   apply %cviConsumer(%1)
+  //   apply %guaranteedUser(%0)
+  //   destroy_value %0
+  //
+  // Easily, if we were to eliminate the copy_value, destroy_value, the object's
+  // lifetime could potentially be shrunk before guaranteedUser is executed,
+  // causing guaranteedUser to be a use-after-free.
+  //
+  // As an extra wrinkle, until all interior pointer constructs (e.x.:
+  // project_box) are guaranteed to be guaranted by a begin_borrow, we can not
+  // in general safely shrink lifetimes. So even if we think we can prove that
+  // all non-consuming uses of %0 are before apply %cviConsumer, we may miss
+  // implicit uses that are not guarded yet by a begin_borrow, resulting in
+  // use-after-frees.
+  //
+  // With that in mind, we only handle cases today where we can prove that
+  // destroy_value is strictly before the consuming use of the operand. This
+  // guarantees that we are not shrinking the lifetime of the underlying object.
+  //
+  // First we handle the simple case: where the cviConsumer is a return inst. In
+  // such a case, we know for sure that cviConsumer post-dominates the
+  // destroy_value.
+  auto cviConsumerIter = cviConsumer->getUser()->getIterator();
+  if (isa<ReturnInst>(cviConsumerIter)) {
+    return true;
+  }
+
+  // Then see if our cviConsumer is in the same block as a return inst and the
+  // destroy_value is not. In that case, we know that the cviConsumer must
+  // post-dominate the destroy_value.
+  auto *cviConsumingBlock = cviConsumerIter->getParent();
+  if (isa<ReturnInst>(cviConsumingBlock->getTerminator()) &&
+      cviConsumingBlock != cviOperandDestroy->getParent()) {
+    return true;
+  }
+
+  // Otherwise, we only support joining live ranges where the cvi and the cvi's
+  // operand's destroy are in the same block with the destroy_value of cvi
+  // operand needing to be strictly after the copy_value. This analysis can be
+  // made significantly stronger by using LiveRanges, but this is simple for
+  // now.
+  auto cviOperandDestroyIter = cviOperandDestroy->getIterator();
+  if (cviConsumingBlock != cviOperandDestroyIter->getParent()) {
+    return false;
+  }
+
+  // TODO: This should really be llvm::find, but for some reason, the templates
+  // do not match up given the current state of the iterators. This impl works
+  // in a pinch though.
+  return llvm::any_of(
+      llvm::make_range(cviOperandDestroyIter,
+                       cviOperandDestroyIter->getParent()->end()),
+      [&](const SILInstruction &val) { return &*cviConsumerIter == &val; });
+}
+
+// # The Problem We Are Solving
+//
+// The main idea here is that we are trying to eliminate the simplest, easiest
+// form of live range joining. Consider the following SIL:
+//
+//   ```
+//   %cviOperand = ...                // @owned value
+//   %cvi = copy_value %cviOperand    // copy of @owned value
+//   ...
+//   destroy_value %cviOperandDestroy // destruction of @owned value
+//   ...
+//   apply %consumingUser(%cvi)       // destruction of copy of @owned value
+//   ```
+//
+// We want to reduce reference count traffic by eliminating the middle
+// copy/destroy yielding:
+//
+//   ```
+//   %cviOperand = ...                // @owned value
+//   // *eliminated copy_value*
+//   ...
+//   // *eliminated destroy_value*
+//   ...
+//   apply %consumingUser(%cviOperand)       // destruction of copy of @owned
+//   value
+//   ```
+//
+// # Safety
+//
+// In order to do this safely, we need to take the union of the two objects
+// lifetimes since we are only joining lifetimes. This ensures that we can rely
+// on SILGen's correctness on inserting safe lifetimes. To keep this simple
+// today we only optimize if the destroy_value and consuming user are in the
+// same block and the consuming user is later in the block than the
+// destroy_value.
+//
+// DISCUSSION: The reason why we do not shrink lifetimes today is that all
+// interior pointers (e.x. project_box) are properly guarded by
+// begin_borrow. Because of that we can not shrink lifetimes and instead rely on
+// SILGen's correctness.
+bool SemanticARCOptVisitor::tryJoiningCopyValueLiveRangeWithOperand(
+    CopyValueInst *cvi) {
+  // First do a quick check if our operand is owned. If it is not owned, we can
+  // not join live ranges.
+  SILValue operand = cvi->getOperand();
+  if (operand.getOwnershipKind() != ValueOwnershipKind::Owned) {
+    return false;
+  }
+
+  // Then check if our operand has a single destroy_value. If it does and that
+  // destroy_value is strictly before the consumer of our copy_value in the same
+  // block as the consumer of said copy_value then we can always join the live
+  // ranges.
+  //
+  // Example:
+  //
+  //   ```
+  //   %1 = copy_value %0
+  //   ...
+  //   destroy_value %0
+  //   apply %consumingUser(%1)
+  //   ```
+  // ->
+  //
+  //   ```
+  //   apply %consumingUser(%0)
+  //   ```
+  //
+  // DISCUSSION: We need to ensure that the consuming use of the copy_value is
+  // strictly after the destroy_value to ensure that we do not shrink the live
+  // range of the operand if the operand has any normal uses beyond our copy
+  // value. Otherwise, we could have normal uses /after/ the consuming use of
+  // our copy_value.
+  if (auto *dvi = operand->getSingleConsumingUserOfType<DestroyValueInst>()) {
+    if (canSafelyJoinSimpleRange(operand, dvi, cvi)) {
+      eraseInstruction(dvi);
+      eraseAndRAUWSingleValueInstruction(cvi, operand);
+      NumEliminatedInsts += 2;
+      return true;
+    }
+  }
+
+  // Otherwise, we couldn't handle this case, so return false.
+  return false;
+}
+
 bool SemanticARCOptVisitor::visitCopyValueInst(CopyValueInst *cvi) {
   // If our copy value inst has only destroy_value users, it is a dead live
   // range. Try to eliminate them.
   if (eliminateDeadLiveRangeCopyValue(cvi)) {
     return true;
   }
 
+  // Then see if copy_value operand's lifetime ends after our copy_value via a
+  // destroy_value. If so, we can join their lifetimes.
+  if (tryJoiningCopyValueLiveRangeWithOperand(cvi)) {
+    return true;
+  }
+
   // Then try to perform the guaranteed copy value optimization.
   if (performGuaranteedCopyValueOptimization(cvi)) {
     return true;