Add reborrow dependencies to DCE

A reborrow's base value may change if the base value is also passed as
an operand on a branch. Record reborrow dependencies so that we can mark
the base value as live, if the reborrow was also live. This ensures we
do not insert destroy_value on the base value prematurely before the
lifetime of the reborrow ends.

Author: meg-gupta

lib/SILOptimizer/Transforms/DeadCodeElimination.cpp

@@ -12,12 +12,13 @@
 
 #define DEBUG_TYPE "sil-dce"
 #include "swift/SIL/DebugUtils.h"
+#include "swift/SIL/OwnershipUtils.h"
 #include "swift/SIL/SILArgument.h"
 #include "swift/SIL/SILBasicBlock.h"
+#include "swift/SIL/SILBitfield.h"
 #include "swift/SIL/SILBuilder.h"
 #include "swift/SIL/SILFunction.h"
 #include "swift/SIL/SILUndef.h"
-#include "swift/SIL/SILBitfield.h"
 #include "swift/SILOptimizer/Analysis/DominanceAnalysis.h"
 #include "swift/SILOptimizer/PassManager/Passes.h"
 #include "swift/SILOptimizer/PassManager/Transforms.h"
@@ -118,6 +119,14 @@ class DCE {
   llvm::DenseMap<SILValue, SmallPtrSet<SILInstruction *, 4>>
       ReverseDependencies;
 
+  // reborrowDependencies tracks the dependency of a reborrowed phiArg with its
+  // renamed base value.
+  // A reborrowed phiArg may have a new base value, if it's original base value
+  // was also passed as a branch operand. The renamed base value should then be
+  // live if the reborrow phiArg was also live.
+  using BaseValueSet = SmallPtrSet<SILValue, 8>;
+  llvm::DenseMap<SILPhiArgument *, BaseValueSet> reborrowDependencies;
+
   /// Tracks if the pass changed branches.
   bool BranchesChanged = false;
   /// Tracks if the pass changed ApplyInsts.
@@ -128,6 +137,9 @@ class DCE {
   /// Record a reverse dependency from \p from to \p to meaning \p to is live
   /// if \p from is also live.
   void addReverseDependency(SILValue from, SILInstruction *to);
+  /// Starting from \p borrowInst find all reborrow dependency of its reborrows
+  /// with their renamed base values.
+  void findReborrowDependencies(BeginBorrowInst *borrowInst);
   bool removeDead();
 
   void computeLevelNumbers(PostDomTreeNode *root);
@@ -253,6 +265,10 @@ void DCE::markLive() {
         addReverseDependency(I.getOperand(0), &I);
         break;
       }
+      case SILInstructionKind::BeginBorrowInst: {
+        findReborrowDependencies(cast<BeginBorrowInst>(&I));
+        break;
+      }
       default:
         if (seemsUseful(&I))
           markInstructionLive(&I);
@@ -275,6 +291,21 @@ void DCE::addReverseDependency(SILValue from, SILInstruction *to) {
   ReverseDependencies[from].insert(to);
 }
 
+void DCE::findReborrowDependencies(BeginBorrowInst *borrowInst) {
+  LLVM_DEBUG(llvm::dbgs() << "Finding reborrow dependencies of " << borrowInst
+                          << "\n");
+  BorrowingOperand initialScopedOperand(&borrowInst->getOperandRef());
+  auto visitReborrowBaseValuePair = [&](SILPhiArgument *phiArg,
+                                        SILValue baseValue) {
+    reborrowDependencies[phiArg].insert(baseValue);
+  };
+  // Find all reborrow dependencies starting from \p borrowInst and populate
+  // them in reborrowDependencies
+  findTransitiveReborrowBaseValuePairs(initialScopedOperand,
+                                       borrowInst->getOperand(),
+                                       visitReborrowBaseValuePair);
+}
+
 // Mark as live the terminator argument at index ArgIndex in Pred that
 // targets Succ.
 void DCE::markTerminatorArgsLive(SILBasicBlock *Pred,
@@ -352,6 +383,12 @@ void DCE::propagateLiveBlockArgument(SILArgument *Arg) {
     markInstructionLive(depInst);
   }
 
+  if (auto *phi = dyn_cast<SILPhiArgument>(Arg)) {
+    for (auto depVal : reborrowDependencies.lookup(phi)) {
+      markValueLive(depVal);
+    }
+  }
+
   auto *Block = Arg->getParent();
   auto ArgIndex = Arg->getIndex();
 

test/SILOptimizer/dead_code_elimination_nontrivial_ossa.sil

@@ -339,3 +339,86 @@ bb4:
   return %res : $()
 }
 
+struct TestStruct {
+  var val:Klass
+  var index:Int
+}
+
+// CHECK-LABEL: sil [ossa] @dce_destructurenotfullydead :
+// CHECK-NOT: copy_value
+// CHECK-LABEL: } // end sil function 'dce_destructurenotfullydead'
+sil [ossa] @dce_destructurenotfullydead : $@convention(thin) (@owned TestStruct) -> Int {
+bb0(%0 : @owned $TestStruct):
+  %stk = alloc_stack $TestStruct
+  store %0 to [init] %stk : $*TestStruct
+  %copy = load [take] %stk : $*TestStruct
+  (%2, %3) = destructure_struct %copy : $TestStruct
+  %4 = struct $TestStruct (%2 : $Klass, %3 : $Int)
+  destroy_value %4 : $TestStruct
+  dealloc_stack %stk : $*TestStruct
+  return %3 : $Int
+}
+
+// CHECK-LABEL: sil [ossa] @dce_borrowlifetime1 :
+// CHECK: bb1([[ARG1:%.*]] : @owned $NonTrivialStruct, [[ARG2:%.*]] : @guaranteed $NonTrivialStruct):
+// CHECK-LABEL: } // end sil function 'dce_borrowlifetime1'
+sil [ossa] @dce_borrowlifetime1 : $@convention(thin) (@guaranteed NonTrivialStruct) -> @owned NonTrivialStruct {
+bb0(%0 : @guaranteed $NonTrivialStruct):
+  %copy = copy_value %0 : $NonTrivialStruct
+  %borrow = begin_borrow %copy : $NonTrivialStruct
+  br bb1(%copy : $NonTrivialStruct, %borrow : $NonTrivialStruct)
+
+bb1(%copy2 : @owned $NonTrivialStruct, %borrow2 : @guaranteed $NonTrivialStruct):
+  %newcopy = copy_value %borrow2 : $NonTrivialStruct
+  end_borrow %borrow2 : $NonTrivialStruct
+  destroy_value %copy2 : $NonTrivialStruct
+  return %newcopy : $NonTrivialStruct
+}
+
+// CHECK-LABEL: sil [ossa] @infinite_loop :
+// CHECK-NOT: copy_value
+// CHECK-LABEL: } // end sil function 'infinite_loop'
+sil [ossa] @infinite_loop : $@convention(thin) (@guaranteed NonTrivialStruct, @guaranteed NonTrivialStruct) -> () {
+bb0(%0 : @guaranteed $NonTrivialStruct, %1 : @guaranteed $NonTrivialStruct):
+  cond_br undef, bb1, bb4
+
+bb1:
+  %copy0 = copy_value %0 : $NonTrivialStruct
+  %borrow0 = begin_borrow %copy0 : $NonTrivialStruct
+  br bb3(%borrow0 : $NonTrivialStruct, %copy0 : $NonTrivialStruct)
+
+bb3(%newborrow : @guaranteed $NonTrivialStruct, %newowned : @owned $NonTrivialStruct):
+  br bb3(%newborrow : $NonTrivialStruct, %newowned : $NonTrivialStruct)
+
+bb4:
+  %ret = tuple ()
+  return %ret : $()
+}
+
+// CHECK-LABEL: sil [ossa] @dce_reborrow_with_different_basevalues :
+// CHECK: bb3([[ARG1:%.*]] : @guaranteed $NonTrivialStruct, [[ARG2:%.*]] : @owned $NonTrivialStruct, [[ARG3:%.*]] : @owned $NonTrivialStruct):
+// CHECK-LABEL: } // end sil function 'dce_reborrow_with_different_basevalues'
+sil [ossa] @dce_reborrow_with_different_basevalues : $@convention(thin) (@guaranteed NonTrivialStruct, @guaranteed NonTrivialStruct) -> @owned NonTrivialStruct {
+bb0(%0 : @guaranteed $NonTrivialStruct, %1 : @guaranteed $NonTrivialStruct):
+  cond_br undef, bb1, bb2
+
+bb1:
+  %copy0a = copy_value %0 : $NonTrivialStruct
+  %borrow0a = begin_borrow %copy0a : $NonTrivialStruct
+  %copy1a = copy_value %1 : $NonTrivialStruct
+  br bb3(%borrow0a : $NonTrivialStruct, %copy0a : $NonTrivialStruct, %copy1a : $NonTrivialStruct)
+
+bb2:
+  %copy1b = copy_value %1 : $NonTrivialStruct
+  %borrow0b = begin_borrow %copy1b : $NonTrivialStruct
+  %copy0b = copy_value %0 : $NonTrivialStruct
+  br bb3(%borrow0b : $NonTrivialStruct, %copy0b : $NonTrivialStruct, %copy1b : $NonTrivialStruct)
+
+bb3(%newborrow : @guaranteed $NonTrivialStruct, %newowned1 : @owned $NonTrivialStruct, %newowned2 : @owned $NonTrivialStruct):
+  %res = copy_value %newborrow : $NonTrivialStruct
+  end_borrow %newborrow : $NonTrivialStruct
+  destroy_value %newowned1 : $NonTrivialStruct
+  destroy_value %newowned2 : $NonTrivialStruct
+  return %res : $NonTrivialStruct
+}
+