From fd6d29d10b3f11e37e3b0112e4f326d45d0d74a6 Mon Sep 17 00:00:00 2001
From: Meghana Gupta <meghanavgupta@gmail.com>
Date: Tue, 12 Nov 2024 12:44:00 -0800
Subject: [PATCH] Insert end_borrow for dead phi operands only when required

DCE inserts end_borrow at phi operands when a guaranteed phi becomes dead.
This should be done for reborrows which end the lifetime of the incoming value.
The existing check was not accurate and ended up inserting end_borrow for forwarded values as well.

Fixes rdar://139283745
---
 .../Transforms/DeadCodeElimination.cpp        | 32 ++++++++++---------
 .../dead_code_elimination_nontrivial_ossa.sil | 17 ++++++++++
 2 files changed, 34 insertions(+), 15 deletions(-)

diff --git a/lib/SILOptimizer/Transforms/DeadCodeElimination.cpp b/lib/SILOptimizer/Transforms/DeadCodeElimination.cpp
index f740c083196..736a318c9d7 100644
--- a/lib/SILOptimizer/Transforms/DeadCodeElimination.cpp
+++ b/lib/SILOptimizer/Transforms/DeadCodeElimination.cpp
@@ -187,9 +187,8 @@ class DCE {
                                 llvm::SmallPtrSetImpl<SILBasicBlock *> &);
   SILBasicBlock *nearestUsefulPostDominator(SILBasicBlock *Block);
   void replaceBranchWithJump(SILInstruction *Inst, SILBasicBlock *Block);
-  /// If \p value is live, insert a lifetime ending operation in ossa.
-  /// destroy_value for @owned value and end_borrow for a @guaranteed value.
-  void endLifetimeOfLiveValue(SILValue value, SILInstruction *insertPt);
+  /// Insert lifetime ending instruction if defining value of \p op is live.
+  void endLifetimeOfLiveValue(Operand *op, SILInstruction *insertPt);
 
 public:
   DCE(SILFunction *F, PostDominanceInfo *PDT)
@@ -585,7 +584,8 @@ void DCE::replaceBranchWithJump(SILInstruction *Inst, SILBasicBlock *Block) {
   (void)Branch;
 }
 
-void DCE::endLifetimeOfLiveValue(SILValue value, SILInstruction *insertPt) {
+void DCE::endLifetimeOfLiveValue(Operand *op, SILInstruction *insertPt) {
+  auto value = op->get();
   if (SILInstruction *inst = value->getDefiningInstruction()) {
     if (!LiveInstructions.contains(inst))
       return;
@@ -593,15 +593,15 @@ void DCE::endLifetimeOfLiveValue(SILValue value, SILInstruction *insertPt) {
     if (!LiveArguments.contains(arg))
       return;
   }
+
+  assert(op->isLifetimeEnding());
   SILBuilderWithScope builder(insertPt);
   if (value->getOwnershipKind() == OwnershipKind::Owned) {
-    builder.emitDestroyOperation(RegularLocation::getAutoGeneratedLocation(),
-                                 value);
+    builder.createDestroyValue(RegularLocation::getAutoGeneratedLocation(),
+                               value);
   }
-  BorrowedValue borrow(lookThroughBorrowedFromDef(value));
-  if (borrow && borrow.isLocalScope()) {
-    builder.emitEndBorrowOperation(RegularLocation::getAutoGeneratedLocation(),
-                                   value);
+  if (value->getOwnershipKind() == OwnershipKind::Guaranteed) {
+    builder.createEndBorrow(RegularLocation::getAutoGeneratedLocation(), value);
   }
 }
 
@@ -657,7 +657,7 @@ bool DCE::removeDead() {
       for (auto *pred : BB.getPredecessorBlocks()) {
         auto *predTerm = pred->getTerminator();
         SILInstruction *insertPt = predTerm;
-        if (phiArg->getOwnershipKind() == OwnershipKind::Guaranteed) {
+        if (phiArg->isReborrow()) {
           // If the phiArg is dead and had reborrow dependencies, its baseValue
           // may also have been dead and a destroy_value of its baseValue may
           // have been inserted before the pred's terminator. Make sure to
@@ -677,8 +677,10 @@ bool DCE::removeDead() {
             insertPt = &predInst;
           }
         }
-
-        endLifetimeOfLiveValue(phiArg->getIncomingPhiValue(pred), insertPt);
+        auto *predOp = phiArg->getIncomingPhiOperand(pred);
+        if (predOp->isLifetimeEnding()) {
+          endLifetimeOfLiveValue(predOp, insertPt);
+        }
       }
       erasePhiArgument(&BB, i, /*cleanupDeadPhiOps=*/true,
                        InstModCallbacks().onCreateNewInst(
@@ -702,7 +704,7 @@ bool DCE::removeDead() {
 
         for (auto &op : termInst->getAllOperands()) {
           if (op.isLifetimeEnding()) {
-            endLifetimeOfLiveValue(op.get(), termInst);
+            endLifetimeOfLiveValue(&op, termInst);
           }
         }
         LLVM_DEBUG(llvm::dbgs() << "Replacing branch: ");
@@ -725,7 +727,7 @@ bool DCE::removeDead() {
       if (F->hasOwnership()) {
         for (auto &Op : Inst->getAllOperands()) {
           if (Op.isLifetimeEnding()) {
-            endLifetimeOfLiveValue(Op.get(), Inst);
+            endLifetimeOfLiveValue(&Op, Inst);
           }
         }
       }
diff --git a/test/SILOptimizer/dead_code_elimination_nontrivial_ossa.sil b/test/SILOptimizer/dead_code_elimination_nontrivial_ossa.sil
index 3f7635a25cb..4c8a0bc74f7 100644
--- a/test/SILOptimizer/dead_code_elimination_nontrivial_ossa.sil
+++ b/test/SILOptimizer/dead_code_elimination_nontrivial_ossa.sil
@@ -1220,3 +1220,20 @@ bb0:
   return %t : $()
 }
 
+// CHECK-LABEL: sil [ossa] @dce_dead_guaranteedphi :
+// bb2:
+// CHECK-LABEL: } // end sil function 'dce_dead_guaranteedphi'
+sil [ossa] @dce_dead_guaranteedphi : $@convention(thin) (@owned NonTrivialStruct) -> () {
+bb0(%0 : @owned $NonTrivialStruct):
+  %b = begin_borrow %0 : $NonTrivialStruct
+  %func = function_ref @$use_nontrivialstruct2 : $@convention(thin) (@guaranteed NonTrivialStruct) -> ()
+  %funcres = apply %func(%b) : $@convention(thin) (@guaranteed NonTrivialStruct) -> ()
+  br bb2(%b : $NonTrivialStruct, %b : $NonTrivialStruct)
+
+bb2(%phi1 : @guaranteed $NonTrivialStruct, %phi2 : @guaranteed $NonTrivialStruct):
+  end_borrow %phi1 : $NonTrivialStruct
+  destroy_value %0 : $NonTrivialStruct
+  %9999 = tuple()
+  return %9999 : $()
+}
+