[semantic-arc-opts] Teach the guaranteed copy_value peephole how to handle instructions that can forward either owned or guaranteed ownership.

This fixes issues exposed by my turning off the Nominal Type RValue peephole. It
should give us some nice ARC wins as well potentially.

Author: gottesmm

include/swift/SIL/SILInstruction.h

@@ -505,6 +505,16 @@ class SILInstruction
     return isTypeDependentOperand(Op.getOperandNumber());
   }
 
+private:
+  /// Predicate used to filter OperandValueRange.
+  struct OperandToValue;
+
+public:
+  using OperandValueRange =
+      OptionalTransformRange<ArrayRef<Operand>, OperandToValue>;
+  OperandValueRange
+  getOperandValues(bool skipTypeDependentOperands = false) const;
+
   SILValue getOperand(unsigned Num) const {
     return getAllOperands()[Num].get();
   }
@@ -514,6 +524,7 @@ class SILInstruction
   }
 
   /// Return the list of results produced by this instruction.
+  bool hasResults() const { return !getResults().empty(); }
   SILInstructionResultArray getResults() const { return getResultsImpl(); }
   unsigned getNumResults() const { return getResults().size(); }
 
@@ -668,6 +679,27 @@ class SILInstruction
   static bool classof(const ValueBase *) = delete;
 };
 
+struct SILInstruction::OperandToValue {
+  const SILInstruction &i;
+  bool skipTypeDependentOps;
+
+  OperandToValue(const SILInstruction &i, bool skipTypeDependentOps)
+      : i(i), skipTypeDependentOps(skipTypeDependentOps) {}
+
+  Optional<SILValue> operator()(const Operand &use) const {
+    if (skipTypeDependentOps && i.isTypeDependentOperand(use))
+      return None;
+    return use.get();
+  }
+};
+
+inline auto
+SILInstruction::getOperandValues(bool skipTypeDependentOperands) const
+    -> OperandValueRange {
+  return OperandValueRange(getAllOperands(),
+                           OperandToValue(*this, skipTypeDependentOperands));
+}
+
 inline llvm::raw_ostream &operator<<(llvm::raw_ostream &OS,
                                      const SILInstruction &I) {
   I.print(OS);
@@ -801,6 +833,9 @@ class OwnershipForwardingSingleValueInst : public SingleValueInstruction {
 
 public:
   ValueOwnershipKind getOwnershipKind() const { return ownershipKind; }
+  void setOwnershipKind(ValueOwnershipKind newOwnershipKind) {
+    ownershipKind = newOwnershipKind;
+  }
 };
 
 /// A value base result of a multiple value instruction.
@@ -838,6 +873,11 @@ class MultipleValueInstructionResult : public ValueBase {
   /// This is stored in the bottom 3 bits of ValueBase's subclass data.
   ValueOwnershipKind getOwnershipKind() const;
 
+  /// Set the ownership kind assigned to this result.
+  ///
+  /// This is stored in SILNode in the subclass data.
+  void setOwnershipKind(ValueOwnershipKind Kind);
+
   static bool classof(const SILInstruction *) = delete;
   static bool classof(const SILUndef *) = delete;
   static bool classof(const SILArgument *) = delete;
@@ -851,11 +891,6 @@ class MultipleValueInstructionResult : public ValueBase {
   }
 
 protected:
-  /// Set the ownership kind assigned to this result.
-  ///
-  /// This is stored in SILNode in the subclass data.
-  void setOwnershipKind(ValueOwnershipKind Kind);
-
   /// Set the index of this result.
   void setIndex(unsigned NewIndex);
 };
@@ -4077,6 +4112,9 @@ class OwnershipForwardingConversionInst : public ConversionInst {
 
 public:
   ValueOwnershipKind getOwnershipKind() const { return ownershipKind; }
+  void setOwnershipKind(ValueOwnershipKind newOwnershipKind) {
+    ownershipKind = newOwnershipKind;
+  }
 };
 
 /// ConvertFunctionInst - Change the type of a function value without
@@ -5305,6 +5343,9 @@ class OwnershipForwardingSelectEnumInstBase : public SelectEnumInstBase {
 
 public:
   ValueOwnershipKind getOwnershipKind() const { return ownershipKind; }
+  void setOwnershipKind(ValueOwnershipKind newOwnershipKind) {
+    ownershipKind = newOwnershipKind;
+  }
 };
 
 /// Select one of a set of values based on the case of an enum.
@@ -7849,6 +7890,11 @@ inline void SILSuccessor::pred_iterator::cacheBasicBlock() {
   }
 }
 
+// Declared in SILValue.h
+inline bool Operand::isTypeDependent() const {
+  return getUser()->isTypeDependentOperand(*this);
+}
+
 } // end swift namespace
 
 //===----------------------------------------------------------------------===//

include/swift/SIL/SILValue.h

@@ -586,6 +586,11 @@ class Operand {
   SILInstruction *getUser() { return Owner; }
   const SILInstruction *getUser() const { return Owner; }
 
+  /// Return true if this operand is a type dependent operand.
+  ///
+  /// Implemented in SILInstruction.h
+  bool isTypeDependent() const;
+
   /// Return which operand this is in the operand list of the using instruction.
   unsigned getOperandNumber() const;
 

lib/SIL/OperandOwnership.cpp

@@ -719,6 +719,10 @@ OperandOwnershipKindClassifier::visitFullApply(FullApplySite apply) {
     return Map::allLive();
   }
 
+  // If we have a type dependent operand, return an empty map.
+  if (apply.getInstruction()->isTypeDependentOperand(op))
+    return Map();
+
   unsigned argIndex = apply.getCalleeArgIndex(op);
   auto conv = apply.getSubstCalleeConv();
   SILParameterInfo paramInfo = conv.getParamInfoForSILArg(argIndex);

lib/SILOptimizer/Mandatory/SemanticARCOpts.cpp

@@ -40,36 +40,73 @@ STATISTIC(NumLoadCopyConvertedToLoadBorrow,
 ///
 /// Semantically this implies that a value is never passed off as +1 to memory
 /// or another function implying it can be used everywhere at +0.
-static bool isConsumed(SILValue v,
-                       SmallVectorImpl<DestroyValueInst *> &destroys) {
+static bool isConsumed(
+    SILValue v, SmallVectorImpl<DestroyValueInst *> &destroys,
+    NullablePtr<SmallVectorImpl<SILInstruction *>> forwardingInsts = nullptr) {
   assert(v.getOwnershipKind() == ValueOwnershipKind::Owned);
-  return !all_of(v->getUses(), [&destroys](Operand *op) {
-    // We know that a copy_value produces an @owned value. Look
-    // through all of our uses and classify them as either
-    // invalidating or not invalidating. Make sure that all of the
-    // invalidating ones are destroy_value since otherwise the
-    // live_range is not complete.
+  SmallVector<Operand *, 32> worklist(v->use_begin(), v->use_end());
+  while (!worklist.empty()) {
+    auto *op = worklist.pop_back_val();
+
+    // Skip type dependent operands.
+    if (op->isTypeDependent())
+      continue;
+
+    auto *user = op->getUser();
+
+    // We know that a copy_value produces an @owned value. Look through all of
+    // our uses and classify them as either invalidating or not
+    // invalidating. Make sure that all of the invalidating ones are
+    // destroy_value since otherwise the live_range is not complete.
     auto map = op->getOwnershipKindMap();
     auto constraint = map.getLifetimeConstraint(ValueOwnershipKind::Owned);
     switch (constraint) {
     case UseLifetimeConstraint::MustBeInvalidated: {
       // See if we have a destroy value. If we don't we have an
       // unknown consumer. Return false, we need this live range.
-      auto *dvi = dyn_cast<DestroyValueInst>(op->getUser());
-      if (!dvi)
-        return false;
+      if (auto *dvi = dyn_cast<DestroyValueInst>(user)) {
+        destroys.push_back(dvi);
+        continue;
+      }
+
+      // Otherwise, see if we have a forwarding value that has a single
+      // non-trivial operand that can accept a guaranteed value. If so, at its
+      // users to the worklist and continue.
+      //
+      // DISCUSSION: For now we do not support forwarding instructions with
+      // multiple non-trivial arguments since we would need to optimize all of
+      // the non-trivial arguments at the same time.
+      //
+      // NOTE: Today we do not support TermInsts for simplicity... we /could/
+      // support it though if we need to.
+      if (forwardingInsts.isNonNull() && !isa<TermInst>(user) &&
+          isGuaranteedForwardingInst(user) &&
+          1 == count_if(user->getOperandValues(
+                            true /*ignore type dependent operands*/),
+                        [&](SILValue v) {
+                          return v.getOwnershipKind() ==
+                                 ValueOwnershipKind::Owned;
+                        })) {
+        forwardingInsts.get()->push_back(user);
+        for (SILValue v : user->getResults()) {
+          copy(v->getUses(), std::back_inserter(worklist));
+        }
+        continue;
+      }
 
-      // Otherwise, return true and stash this destroy value.
-      destroys.push_back(dvi);
+      // Otherwise be conservative and assume that we /may consume/ the value.
       return true;
     }
     case UseLifetimeConstraint::MustBeLive:
       // Ok, this constraint can take something owned as live. Lets
       // see if it can also take something that is guaranteed. If it
       // can not, then we bail.
-      return map.canAcceptKind(ValueOwnershipKind::Guaranteed);
+      map.canAcceptKind(ValueOwnershipKind::Guaranteed);
+      continue;
     }
-  });
+  }
+
+  return false;
 }
 
 //===----------------------------------------------------------------------===//
@@ -168,13 +205,14 @@ static bool performGuaranteedCopyValueOptimization(CopyValueInst *cvi) {
   if (!canHandleOperand(cvi->getOperand(), borrowIntroducers))
     return false;
 
-  // Then go over all of our uses. Find our destroying instructions
-  // and make sure all of them are destroy_value. For our
-  // non-destroying instructions, make sure that they accept a
-  // guaranteed value. After that, make sure that our destroys are
-  // within the lifetime of our borrowed values.
+  // Then go over all of our uses. Find our destroying instructions (ignoring
+  // forwarding instructions that can forward both owned and guaranteed) and
+  // make sure all of them are destroy_value. For our non-destroying
+  // instructions, make sure that they accept a guaranteed value. After that,
+  // make sure that our destroys are within the lifetime of our borrowed values.
   SmallVector<DestroyValueInst *, 16> destroys;
-  if (isConsumed(cvi, destroys))
+  SmallVector<SILInstruction *, 16> guaranteedForwardingInsts;
+  if (isConsumed(cvi, destroys, &guaranteedForwardingInsts))
     return false;
 
   // If we reached this point, then we know that all of our users can
@@ -189,15 +227,56 @@ static bool performGuaranteedCopyValueOptimization(CopyValueInst *cvi) {
   assert(all_of(borrowIntroducers,
                 [](SILValue v) { return isa<SILFunctionArgument>(v); }));
 
-  // Otherwise, we know that our copy_value/destroy_values are all
-  // completely within the guaranteed value scope.
+  // Otherwise, we know that our copy_value/destroy_values are all completely
+  // within the guaranteed value scope. First delete the destroys/copies.
   while (!destroys.empty()) {
     auto *dvi = destroys.pop_back_val();
     dvi->eraseFromParent();
     ++NumEliminatedInsts;
   }
   cvi->replaceAllUsesWith(cvi->getOperand());
   cvi->eraseFromParent();
+
+  // Then change all of our guaranteed forwarding insts to have guaranteed
+  // ownership kind instead of what ever they previously had (ignoring trivial
+  // results);
+  while (!guaranteedForwardingInsts.empty()) {
+    auto *i = guaranteedForwardingInsts.pop_back_val();
+
+    assert(i->hasResults());
+
+    for (SILValue result : i->getResults()) {
+      if (auto *svi = dyn_cast<OwnershipForwardingSingleValueInst>(result)) {
+        if (svi->getOwnershipKind() == ValueOwnershipKind::Owned) {
+          svi->setOwnershipKind(ValueOwnershipKind::Guaranteed);
+        }
+        continue;
+      }
+
+      if (auto *ofci = dyn_cast<OwnershipForwardingConversionInst>(result)) {
+        if (ofci->getOwnershipKind() == ValueOwnershipKind::Owned) {
+          ofci->setOwnershipKind(ValueOwnershipKind::Guaranteed);
+        }
+        continue;
+      }
+
+      if (auto *sei = dyn_cast<OwnershipForwardingSelectEnumInstBase>(result)) {
+        if (sei->getOwnershipKind() == ValueOwnershipKind::Owned) {
+          sei->setOwnershipKind(ValueOwnershipKind::Guaranteed);
+        }
+        continue;
+      }
+
+      if (auto *mvir = dyn_cast<MultipleValueInstructionResult>(result)) {
+        if (mvir->getOwnershipKind() == ValueOwnershipKind::Owned) {
+          mvir->setOwnershipKind(ValueOwnershipKind::Guaranteed);
+        }
+        continue;
+      }
+
+      llvm_unreachable("unhandled forwarding instruction?!");
+    }
+  }
   ++NumEliminatedInsts;
   return true;
 }
@@ -350,6 +429,11 @@ bool SemanticARCOptVisitor::visitLoadInst(LoadInst *li) {
   // load_borrow.
   auto *lbi =
       SILBuilderWithScope(li).createLoadBorrow(li->getLoc(), li->getOperand());
+
+  // Since we are looking through forwarding uses that can accept guaranteed
+  // parameters, we can have multiple destroy_value along the same path. We need
+  // to find the post-dominating block set of these destroy value to ensure that
+  // we do not insert multiple end_borrow.
   while (!destroyValues.empty()) {
     auto *dvi = destroyValues.pop_back_val();
     SILBuilderWithScope(dvi).createEndBorrow(dvi->getLoc(), lbi);