Merge pull request #78015 from nate-chandler/cherrypick/release/6.1/rdar139842132

6.1: [OSSACanonicalizeOwned] Record traversed defs and don't traverse copies of guaranteed values.

Author: nate-chandler

include/swift/SILOptimizer/Utils/CanonicalizeOSSALifetime.h

@@ -96,8 +96,8 @@
 #ifndef SWIFT_SILOPTIMIZER_UTILS_CANONICALOSSALIFETIME_H
 #define SWIFT_SILOPTIMIZER_UTILS_CANONICALOSSALIFETIME_H
 
-#include "swift/Basic/GraphNodeWorklist.h"
 #include "swift/Basic/SmallPtrSetVector.h"
+#include "swift/Basic/TaggedUnion.h"
 #include "swift/SIL/PrunedLiveness.h"
 #include "swift/SIL/SILInstruction.h"
 #include "swift/SILOptimizer/Analysis/DeadEndBlocksAnalysis.h"
@@ -279,8 +279,71 @@ class CanonicalizeOSSALifetime final {
   /// outside the pruned liveness at the time it is discovered.
   llvm::SmallPtrSet<DebugValueInst *, 8> debugValues;
 
-  /// Visited set for general def-use traversal that prevents revisiting values.
-  GraphNodeWorklist<SILValue, 8> defUseWorklist;
+  class Def {
+    struct Root {
+      SILValue value;
+    };
+    struct Copy {
+      CopyValueInst *cvi;
+    };
+    struct BorrowedFrom {
+      BorrowedFromInst *bfi;
+    };
+    struct Reborrow {
+      SILArgument *argument;
+    };
+    using Payload = TaggedUnion<Root, Copy, BorrowedFrom, Reborrow>;
+    Payload payload;
+    Def(Payload payload) : payload(payload) {}
+
+  public:
+    enum class Kind {
+      Root,
+      Copy,
+      BorrowedFrom,
+      Reborrow,
+    };
+    Kind getKind() const {
+      if (payload.isa<Root>()) {
+        return Kind::Root;
+      }
+      if (payload.isa<Copy>()) {
+        return Kind::Copy;
+      }
+      if (payload.isa<BorrowedFrom>()) {
+        return Kind::BorrowedFrom;
+      }
+      assert(payload.isa<Reborrow>());
+      return Kind::Reborrow;
+    }
+    operator Kind() const { return getKind(); }
+    bool operator==(Def rhs) const {
+      return getKind() == rhs.getKind() && getValue() == rhs.getValue();
+    }
+    static Def root(SILValue value) { return {Root{value}}; }
+    static Def copy(CopyValueInst *cvi) { return {Copy{cvi}}; }
+    static Def borrowedFrom(BorrowedFromInst *bfi) {
+      return {BorrowedFrom{bfi}};
+    }
+    static Def reborrow(SILArgument *argument) { return {Reborrow{argument}}; }
+    SILValue getValue() const {
+      switch (*this) {
+      case Kind::Root:
+        return payload.get<Root>().value;
+      case Kind::Copy:
+        return payload.get<Copy>().cvi;
+      case Kind::BorrowedFrom:
+        return payload.get<BorrowedFrom>().bfi;
+      case Kind::Reborrow:
+        return payload.get<Reborrow>().argument;
+      }
+      llvm_unreachable("covered switch");
+    }
+  };
+  friend llvm::DenseMapInfo<Def>;
+
+  /// The defs derived from currentDef whose uses are added to liveness.
+  SmallVector<Def, 8> discoveredDefs;
 
   /// The blocks that were discovered by PrunedLiveness.
   SmallVector<SILBasicBlock *, 32> discoveredBlocks;

lib/SIL/Utils/OSSALifetimeCompletion.cpp

@@ -454,6 +454,7 @@ bool OSSALifetimeCompletion::analyzeAndUpdateLifetime(SILValue value,
 namespace swift::test {
 // Arguments:
 // - SILValue: value
+// - string: either "liveness" or "availability"
 // Dumps:
 // - function
 static FunctionTest OSSALifetimeCompletionTest(

lib/SILOptimizer/Mandatory/ConsumeOperatorCopyableValuesChecker.cpp

@@ -15,6 +15,7 @@
 #include "swift/AST/DiagnosticsSIL.h"
 #include "swift/Basic/Assertions.h"
 #include "swift/Basic/Defer.h"
+#include "swift/Basic/GraphNodeWorklist.h"
 #include "swift/SIL/BasicBlockBits.h"
 #include "swift/SIL/BasicBlockDatastructures.h"
 #include "swift/SIL/DebugUtils.h"

lib/SILOptimizer/Utils/CanonicalizeOSSALifetime.cpp

@@ -65,8 +65,8 @@
 
 #define DEBUG_TYPE "copy-propagation"
 
-#include "swift/Basic/Assertions.h"
 #include "swift/SILOptimizer/Utils/CanonicalizeOSSALifetime.h"
+#include "swift/Basic/Assertions.h"
 #include "swift/SIL/InstructionUtils.h"
 #include "swift/SIL/NodeDatastructures.h"
 #include "swift/SIL/OSSALifetimeCompletion.h"
@@ -132,17 +132,29 @@ static bool isDestroyOfCopyOf(SILInstruction *instruction, SILValue def) {
 bool CanonicalizeOSSALifetime::computeCanonicalLiveness() {
   LLVM_DEBUG(llvm::dbgs() << "Computing canonical liveness from:\n";
              getCurrentDef()->print(llvm::dbgs()));
-  defUseWorklist.initialize(getCurrentDef());
+  SmallVector<unsigned, 8> indexWorklist;
+  ValueSet visitedDefs(getCurrentDef()->getFunction());
+  auto addDefToWorklist = [&](Def def) {
+    if (!visitedDefs.insert(def.getValue()))
+      return;
+    discoveredDefs.push_back(def);
+    indexWorklist.push_back(discoveredDefs.size() - 1);
+  };
+  discoveredDefs.clear();
+  addDefToWorklist(Def::root(getCurrentDef()));
   // Only the first level of reborrows need to be consider. All nested inner
   // adjacent reborrows and phis are encapsulated within their lifetimes.
   SILPhiArgument *arg;
   if ((arg = dyn_cast<SILPhiArgument>(getCurrentDef())) && arg->isPhi()) {
     visitInnerAdjacentPhis(arg, [&](SILArgument *reborrow) {
-      defUseWorklist.insert(reborrow);
+      addDefToWorklist(Def::reborrow(reborrow));
       return true;
     });
   }
-  while (SILValue value = defUseWorklist.pop()) {
+  while (!indexWorklist.empty()) {
+    auto index = indexWorklist.pop_back_val();
+    auto def = discoveredDefs[index];
+    auto value = def.getValue();
     LLVM_DEBUG(llvm::dbgs() << "  Uses of value:\n";
                value->print(llvm::dbgs()));
 
@@ -153,11 +165,20 @@ bool CanonicalizeOSSALifetime::computeCanonicalLiveness() {
       auto *user = use->getUser();
       // Recurse through copies.
       if (auto *copy = dyn_cast<CopyValueInst>(user)) {
-        defUseWorklist.insert(copy);
+        // Don't recurse through copies of borrowed-froms or reborrows.
+        switch (def) {
+        case Def::Kind::Root:
+        case Def::Kind::Copy:
+          addDefToWorklist(Def::copy(copy));
+          break;
+        case Def::Kind::Reborrow:
+        case Def::Kind::BorrowedFrom:
+          break;
+        }
         continue;
       }
       if (auto *bfi = dyn_cast<BorrowedFromInst>(user)) {
-        defUseWorklist.insert(bfi);
+        addDefToWorklist(Def::borrowedFrom(bfi));
         continue;
       }
       // Handle debug_value instructions separately.
@@ -244,7 +265,7 @@ bool CanonicalizeOSSALifetime::computeCanonicalLiveness() {
         // This branch reborrows a guaranteed phi whose lifetime is dependent on
         // currentDef.  Uses of the reborrowing phi extend liveness.
         auto *reborrow = PhiOperand(use).getValue();
-        defUseWorklist.insert(reborrow);
+        addDefToWorklist(Def::reborrow(reborrow));
         break;
       }
     }
@@ -1148,20 +1169,17 @@ void CanonicalizeOSSALifetime::rewriteCopies(
     SmallVectorImpl<DestroyValueInst *> const &newDestroys) {
   assert(getCurrentDef()->getOwnershipKind() == OwnershipKind::Owned);
 
+  // Shadow discoveredDefs in order to constrain its uses.
+  const auto &discoveredDefs = this->discoveredDefs;
+
   InstructionSetVector instsToDelete(getCurrentDef()->getFunction());
-  defUseWorklist.clear();
 
   // Visit each operand in the def-use chain.
   //
   // Return true if the operand can use the current definition. Return false if
   // it requires a copy.
   auto visitUse = [&](Operand *use) {
     auto *user = use->getUser();
-    // Recurse through copies.
-    if (auto *copy = dyn_cast<CopyValueInst>(user)) {
-      defUseWorklist.insert(copy);
-      return true;
-    }
     if (destroys.contains(user)) {
       auto *destroy = cast<DestroyValueInst>(user);
       // If this destroy was marked as a final destroy, ignore it; otherwise,
@@ -1195,37 +1213,54 @@ void CanonicalizeOSSALifetime::rewriteCopies(
   };
 
   // Perform a def-use traversal, visiting each use operand.
-  for (auto useIter = getCurrentDef()->use_begin(),
-         endIter = getCurrentDef()->use_end(); useIter != endIter;) {
-    Operand *use = *useIter++;
-    if (!visitUse(use)) {
-      copyLiveUse(use, getCallbacks());
-    }
-  }
-  while (SILValue value = defUseWorklist.pop()) {
-    CopyValueInst *srcCopy = cast<CopyValueInst>(value);
-    // Recurse through copies while replacing their uses.
-    Operand *reusedCopyOp = nullptr;
-    for (auto useIter = srcCopy->use_begin(); useIter != srcCopy->use_end();) {
-      Operand *use = *useIter++;
-      if (!visitUse(use)) {
-        if (!reusedCopyOp && srcCopy->getParent() == use->getParentBlock()) {
-          reusedCopyOp = use;
-        } else {
+  for (auto def : discoveredDefs) {
+    switch (def) {
+    case Def::Kind::BorrowedFrom:
+    case Def::Kind::Reborrow:
+      // Direct uses of these defs never need to be rewritten.  Being guaranteed
+      // values, none of their direct uses consume an owned value.
+      assert(def.getValue()->getOwnershipKind() == OwnershipKind::Guaranteed);
+      break;
+    case Def::Kind::Root: {
+      SILValue value = def.getValue();
+      for (auto useIter = value->use_begin(), endIter = value->use_end();
+           useIter != endIter;) {
+        Operand *use = *useIter++;
+        if (!visitUse(use)) {
           copyLiveUse(use, getCallbacks());
         }
       }
+      break;
     }
-    if (!(reusedCopyOp && srcCopy->hasOneUse())) {
-      getCallbacks().replaceValueUsesWith(srcCopy, srcCopy->getOperand());
-      if (reusedCopyOp) {
-        reusedCopyOp->set(srcCopy);
-      } else {
-        if (instsToDelete.insert(srcCopy)) {
-          LLVM_DEBUG(llvm::dbgs() << "  Removing " << *srcCopy);
-          ++NumCopiesAndMovesEliminated;
+    case Def::Kind::Copy: {
+      SILValue value = def.getValue();
+      CopyValueInst *srcCopy = cast<CopyValueInst>(value);
+      // Recurse through copies while replacing their uses.
+      Operand *reusedCopyOp = nullptr;
+      for (auto useIter = srcCopy->use_begin();
+           useIter != srcCopy->use_end();) {
+        Operand *use = *useIter++;
+        if (!visitUse(use)) {
+          if (!reusedCopyOp && srcCopy->getParent() == use->getParentBlock()) {
+            reusedCopyOp = use;
+          } else {
+            copyLiveUse(use, getCallbacks());
+          }
         }
       }
+      if (!(reusedCopyOp && srcCopy->hasOneUse())) {
+        getCallbacks().replaceValueUsesWith(srcCopy, srcCopy->getOperand());
+        if (reusedCopyOp) {
+          reusedCopyOp->set(srcCopy);
+        } else {
+          if (instsToDelete.insert(srcCopy)) {
+            LLVM_DEBUG(llvm::dbgs() << "  Removing " << *srcCopy);
+            ++NumCopiesAndMovesEliminated;
+          }
+        }
+      }
+      break;
+    }
     }
   }
   assert(!consumes.hasUnclaimedConsumes());

test/SILOptimizer/canonicalize_ossa_lifetime_unit.sil

@@ -789,3 +789,71 @@ entry(%c : @owned $C):
   specify_test "canonicalize_ossa_lifetime true false true %c"
   unreachable
 }
+
+// CHECK-LABEL: begin running test {{.*}} on consume_copy_of_borrowed_from
+// CHECK-LABEL: sil [ossa] @consume_copy_of_borrowed_from : {{.*}} {
+// The copy can't be eliminated because it's a copy of a borrowed-from.  Such
+// copies are not rewritten.
+// CHECK:         copy_value
+// CHECK-LABEL: } // end sil function 'consume_copy_of_borrowed_from'
+// CHECK-LABEL: end running test {{.*}} on consume_copy_of_borrowed_from
+sil [ossa] @consume_copy_of_borrowed_from : $@convention(thin) (@owned C) -> () {
+bb0(%instance : @owned $C):
+  %borrow = begin_borrow %instance : $C
+  br bb1(%instance : $C, %borrow : $C)
+
+bb1(%instance2 : @owned $C, %reborrowed : @reborrow $C):
+  specify_test "canonicalize_ossa_lifetime true false true %instance2"
+  %reborrow = borrowed %reborrowed : $C from (%instance2 : $C)
+  cond_br undef, bb2, bb3
+
+bb2:
+  end_borrow %reborrow : $C
+  destroy_value %instance2 : $C
+  %retval = tuple ()
+  return %retval : $()
+
+bb3:
+  %copy = copy_value %reborrow : $C
+  %takeC = function_ref @takeC : $@convention(thin) (@owned C) -> ()
+  apply %takeC(%copy) : $@convention(thin) (@owned C) -> ()
+  unreachable
+}
+
+// CHECK-LABEL: begin running test {{.*}} on consume_copy_of_borrowed_from_2
+// CHECK-LABEL: sil [ossa] @consume_copy_of_borrowed_from_2 : {{.*}} {
+// The copy can't be eliminated because it's a copy of a borrowed-from.  Such
+// copies are not rewritten.
+// CHECK:         copy_value
+// CHECK-LABEL: } // end sil function 'consume_copy_of_borrowed_from_2'
+// CHECK-LABEL: end running test {{.*}} on consume_copy_of_borrowed_from_2
+sil [ossa] @consume_copy_of_borrowed_from_2 : $@convention(thin) (@owned C) -> () {
+bb0(%instance : @owned $C):
+  %borrow = begin_borrow %instance : $C
+  br bb1(%instance : $C, %borrow : $C)
+
+bb1(%instance2 : @owned $C, %reborrowed : @reborrow $C):
+  specify_test "canonicalize_ossa_lifetime true false true %instance2"
+  %reborrow = borrowed %reborrowed : $C from (%instance2 : $C)
+  cond_br undef, bb2, bb3
+
+bb2:
+  end_borrow %reborrow : $C
+  destroy_value %instance2 : $C
+  br exit
+
+bb3:
+  %copy = copy_value %reborrow : $C
+  end_borrow %reborrow : $C
+  destroy_value %instance2 : $C
+  br bb4
+
+bb4:
+  %takeC = function_ref @takeC : $@convention(thin) (@owned C) -> ()
+  apply %takeC(%copy) : $@convention(thin) (@owned C) -> ()
+  br exit
+
+exit:
+  %retval = tuple ()
+  return %retval : $()
+}

test/SILOptimizer/copy_propagation_opaque.sil

@@ -242,7 +242,7 @@ bb(%0 : @owned $T):
 // CHECK-ONONE-NEXT:  [[CP:%.*]] = copy_value [[INSTANCE]] : $T
 // CHECK-NEXT:  // function_ref takeMultipleOwned
 // CHECK-NEXT:  function_ref @takeMultipleOwned : $@convention(thin) <τ_0_0> (@in τ_0_0, @in τ_0_0) -> ()
-// CHECK-NEXT:  apply %{{.*}}<T>([[INSTANCE]], [[CP]]) : $@convention(thin) <τ_0_0> (@in τ_0_0, @in τ_0_0) -> ()
+// CHECK-NEXT:  apply %{{.*}}<T>([[CP]], [[INSTANCE]]) : $@convention(thin) <τ_0_0> (@in τ_0_0, @in τ_0_0) -> ()
 // CHECK-ONONE-NEXT:  destroy_value
 // CHECK-NEXT:  tuple ()
 // CHECK-NEXT:  return