Don't pass dead values into proactively added phis in SILMem2Reg

include/swift/SILOptimizer/Utils/InstOptUtils.h

@@ -678,6 +678,10 @@ makeNewValueAvailable(SILValue value, SILBasicBlock *inBlock);
 /// Warning: This does not properly cleanup an OSSA lifetime with a consuming
 /// use blocks inside a loop relative to \p value. The client must create
 /// separate copies for any uses within the loop.
+void destroyValueAtLeakingBlocks(SILValue value,
+                                 ArrayRef<SILBasicBlock *> userBBs);
+
+/// Given an ssa value \p value, create end_lifetime at leaking blocks
 void endLifetimeAtLeakingBlocks(SILValue value,
                                 ArrayRef<SILBasicBlock *> userBBs);
 

lib/SILOptimizer/Transforms/DeadCodeElimination.cpp

@@ -45,10 +45,11 @@ namespace {
 // FIXME: Reconcile the similarities between this and
 //        isInstructionTriviallyDead.
 static bool seemsUseful(SILInstruction *I) {
-  // Even though begin_access/destroy_value/copy_value have side-effects, they
-  // can be DCE'ed if they do not have useful dependencies/reverse dependencies
+  // Even though begin_access/destroy_value/copy_value/end_lifetime have
+  // side-effects, they can be DCE'ed if they do not have useful
+  // dependencies/reverse dependencies
   if (isa<BeginAccessInst>(I) || isa<CopyValueInst>(I) ||
-      isa<DestroyValueInst>(I))
+      isa<DestroyValueInst>(I) || isa<EndLifetimeInst>(I))
     return false;
 
   // A load [copy] is okay to be DCE'ed if there are no useful dependencies
@@ -268,7 +269,8 @@ void DCE::markLive() {
         break;
       }
       case SILInstructionKind::DestroyValueInst:
-      case SILInstructionKind::EndBorrowInst: {
+      case SILInstructionKind::EndBorrowInst:
+      case SILInstructionKind::EndLifetimeInst: {
         // The instruction is live only if it's operand value is also live
         addReverseDependency(I.getOperand(0), &I);
         break;

lib/SILOptimizer/Transforms/RedundantLoadElimination.cpp

@@ -1361,7 +1361,7 @@ SILValue RLEContext::computePredecessorLocationValue(SILBasicBlock *BB,
     for (auto use : phi->getUses()) {
       userBBs.push_back(use->getParentBlock());
     }
-    endLifetimeAtLeakingBlocks(phi, userBBs);
+    destroyValueAtLeakingBlocks(phi, userBBs);
   }
   return makeNewValueAvailable(Val, BB);
 }

lib/SILOptimizer/Transforms/SILMem2Reg.cpp

@@ -228,7 +228,7 @@ namespace {
 
 /// Promotes a single AllocStackInst into registers..
 class StackAllocationPromoter {
-  using BlockSet = BasicBlockSetVector;
+  using BlockSetVector = BasicBlockSetVector;
   using BlockToInstMap = llvm::DenseMap<SILBasicBlock *, SILInstruction *>;
 
   // Use a priority queue keyed on dominator tree level so that inserted nodes
@@ -261,14 +261,18 @@ class StackAllocationPromoter {
   /// AllocStackInst.
   BlockToInstMap lastStoreInBlock;
 
+  /// Records blocks in which the value in asi is moved out.
+  /// i.e There was a load [take] of the asi with no following store.
+  BasicBlockSet deadValueBlock;
 public:
   /// C'tor.
   StackAllocationPromoter(AllocStackInst *inputASI, DominanceInfo *inputDomInfo,
                           DomTreeLevelMap &inputDomTreeLevels,
                           SILBuilderContext &inputCtx,
                           InstructionDeleter &deleter)
       : asi(inputASI), dsi(nullptr), domInfo(inputDomInfo),
-        domTreeLevels(inputDomTreeLevels), ctx(inputCtx), deleter(deleter) {
+        domTreeLevels(inputDomTreeLevels), ctx(inputCtx), deleter(deleter),
+        deadValueBlock(inputASI->getFunction()) {
     // Scan the users in search of a deallocation instruction.
     for (auto *use : asi->getUses()) {
       if (auto *foundDealloc = dyn_cast<DeallocStackInst>(use->getUser())) {
@@ -291,26 +295,27 @@ class StackAllocationPromoter {
   void promoteAllocationToPhi();
 
   /// Replace the dummy nodes with new block arguments.
-  void addBlockArguments(BlockSet &phiBlocks);
+  void addBlockArguments(BlockSetVector &phiBlocks);
 
   /// Fix all of the branch instructions and the uses to use
   /// the AllocStack definitions (which include stores and Phis).
-  void fixBranchesAndUses(BlockSet &blocks);
+  void fixBranchesAndUses(BlockSetVector &blocks);
 
   /// update the branch instructions with the new Phi argument.
   /// The blocks in \p PhiBlocks are blocks that define a value, \p Dest is
   /// the branch destination, and \p Pred is the predecessors who's branch we
   /// modify.
-  void fixPhiPredBlock(BlockSet &phiBlocks, SILBasicBlock *dest,
+  void fixPhiPredBlock(BlockSetVector &phiBlocks, SILBasicBlock *dest,
                        SILBasicBlock *pred);
 
   /// Get the value for this AllocStack variable that is
   /// flowing out of StartBB.
-  SILValue getLiveOutValue(BlockSet &phiBlocks, SILBasicBlock *startBlock);
+  SILValue getLiveOutValue(BlockSetVector &phiBlocks,
+                           SILBasicBlock *startBlock);
 
   /// Get the value for this AllocStack variable that is
   /// flowing into BB.
-  SILValue getLiveInValue(BlockSet &phiBlocks, SILBasicBlock *block);
+  SILValue getLiveInValue(BlockSetVector &phiBlocks, SILBasicBlock *block);
 
   /// Prune AllocStacks usage in the function. Scan the function
   /// and remove in-block usage of the AllocStack. Leave only the first
@@ -329,12 +334,14 @@ class StackAllocationPromoter {
 StoreInst *StackAllocationPromoter::promoteAllocationInBlock(
     SILBasicBlock *blockPromotingWithin) {
   LLVM_DEBUG(llvm::dbgs() << "*** Promoting ASI in block: " << *asi);
-
   // RunningVal is the current value in the stack location.
   // We don't know the value of the alloca until we find the first store.
   SILValue runningVal = SILValue();
   // Keep track of the last StoreInst that we found.
   StoreInst *lastStore = nullptr;
+  // Keep track if value in asi becomes dead via load [take] or destroy_addr in
+  // this block
+  bool deadAsi = false;
 
   // For all instructions in the block.
   for (auto bbi = blockPromotingWithin->begin(),
@@ -345,6 +352,10 @@ StoreInst *StackAllocationPromoter::promoteAllocationInBlock(
 
     if (isLoadFromStack(inst, asi)) {
       auto *li = cast<LoadInst>(inst);
+      // The value in asi was moved out, set deadAsi to true
+      if (li->getOwnershipQualifier() == LoadOwnershipQualifier::Take) {
+        deadAsi = true;
+      }
       if (runningVal) {
         // If we are loading from the AllocStackInst and we already know the
         // content of the Alloca then use it.
@@ -370,6 +381,9 @@ StoreInst *StackAllocationPromoter::promoteAllocationInBlock(
       if (si->getDest() != asi)
         continue;
 
+      // Found a store to asi, set deadAsi to false.
+      deadAsi = false;
+
       // If we see a store [assign], always convert it to a store [init]. This
       // simplifies further processing.
       if (si->getOwnershipQualifier() == StoreOwnershipQualifier::Assign) {
@@ -415,28 +429,28 @@ StoreInst *StackAllocationPromoter::promoteAllocationInBlock(
 
     // Replace destroys with a release of the value.
     if (auto *dai = dyn_cast<DestroyAddrInst>(inst)) {
-      if (dai->getOperand() == asi && runningVal) {
-        replaceDestroy(dai, runningVal, ctx, deleter);
+      if (dai->getOperand() == asi) {
+        if (runningVal) {
+          replaceDestroy(dai, runningVal, ctx, deleter);
+        }
+        // The value in asi was destroyed, set deadAsi to true
+        deadAsi = true;
       }
       continue;
     }
 
-    if (auto *dvi = dyn_cast<DestroyValueInst>(inst)) {
-      if (dvi->getOperand() == runningVal) {
-        // Reset LastStore.
-        // So that we don't end up passing dead values as phi args in
-        // StackAllocationPromoter::fixBranchesAndUses
-        lastStore = nullptr;
-      }
-    }
-
     // Stop on deallocation.
     if (auto *dsi = dyn_cast<DeallocStackInst>(inst)) {
       if (dsi->getOperand() == asi)
         break;
     }
   }
 
+  // If the asi was dead in this block, insert into deadValueBlock in OSSA.
+  if (deadAsi && asi->getFunction()->hasOwnership() &&
+      !asi->getType().isTrivial(*asi->getFunction())) {
+    deadValueBlock.insert(blockPromotingWithin);
+  }
   if (lastStore) {
     assert(lastStore->getOwnershipQualifier() !=
                StoreOwnershipQualifier::Assign &&
@@ -450,21 +464,27 @@ StoreInst *StackAllocationPromoter::promoteAllocationInBlock(
   return lastStore;
 }
 
-void StackAllocationPromoter::addBlockArguments(BlockSet &phiBlocks) {
+void StackAllocationPromoter::addBlockArguments(BlockSetVector &phiBlocks) {
   LLVM_DEBUG(llvm::dbgs() << "*** Adding new block arguments.\n");
 
   for (auto *block : phiBlocks)
     block->createPhiArgument(asi->getElementType(), OwnershipKind::Owned);
 }
 
-SILValue StackAllocationPromoter::getLiveOutValue(BlockSet &phiBlocks,
+SILValue StackAllocationPromoter::getLiveOutValue(BlockSetVector &phiBlocks,
                                                   SILBasicBlock *startBlock) {
   LLVM_DEBUG(llvm::dbgs() << "*** Searching for a value definition.\n");
+  auto *func = asi->getFunction();
+
   // Walk the Dom tree in search of a defining value:
   for (DomTreeNode *domNode = domInfo->getNode(startBlock); domNode;
        domNode = domNode->getIDom()) {
     SILBasicBlock *domBlock = domNode->getBlock();
 
+    // If the value in asi is moved out, then return undef
+    if (deadValueBlock.contains(domBlock)) {
+      return SILUndef::get(asi->getElementType(), *func);
+    }
     // If there is a store (that must come after the phi), use its value.
     BlockToInstMap::iterator it = lastStoreInBlock.find(domBlock);
     if (it != lastStoreInBlock.end())
@@ -486,10 +506,10 @@ SILValue StackAllocationPromoter::getLiveOutValue(BlockSet &phiBlocks,
     LLVM_DEBUG(llvm::dbgs() << "*** Walking up the iDOM.\n");
   }
   LLVM_DEBUG(llvm::dbgs() << "*** Could not find a Def. Using Undef.\n");
-  return SILUndef::get(asi->getElementType(), *asi->getFunction());
+  return SILUndef::get(asi->getElementType(), *func);
 }
 
-SILValue StackAllocationPromoter::getLiveInValue(BlockSet &phiBlocks,
+SILValue StackAllocationPromoter::getLiveInValue(BlockSetVector &phiBlocks,
                                                  SILBasicBlock *block) {
   // First, check if there is a Phi value in the current block. We know that
   // our loads happen before stores, so we need to first check for Phi nodes
@@ -512,7 +532,7 @@ SILValue StackAllocationPromoter::getLiveInValue(BlockSet &phiBlocks,
   return getLiveOutValue(phiBlocks, iDom->getBlock());
 }
 
-void StackAllocationPromoter::fixPhiPredBlock(BlockSet &phiBlocks,
+void StackAllocationPromoter::fixPhiPredBlock(BlockSetVector &phiBlocks,
                                               SILBasicBlock *destBlock,
                                               SILBasicBlock *predBlock) {
   TermInst *ti = predBlock->getTerminator();
@@ -526,7 +546,7 @@ void StackAllocationPromoter::fixPhiPredBlock(BlockSet &phiBlocks,
   deleter.forceDelete(ti);
 }
 
-void StackAllocationPromoter::fixBranchesAndUses(BlockSet &phiBlocks) {
+void StackAllocationPromoter::fixBranchesAndUses(BlockSetVector &phiBlocks) {
   // First update uses of the value.
   SmallVector<LoadInst *, 4> collectedLoads;
 
@@ -595,21 +615,29 @@ void StackAllocationPromoter::fixBranchesAndUses(BlockSet &phiBlocks) {
     }
   }
 
-  // If the owned phi arg we added did not have any uses, create end_lifetime to
-  // end its lifetime. In asserts mode, make sure we have only undef incoming
-  // values for such phi args.
-  for (auto *block : phiBlocks) {
-    auto *phiArg =
-        cast<SILPhiArgument>(block->getArgument(block->getNumArguments() - 1));
-    if (phiArg->use_empty()) {
-      erasePhiArgument(block, block->getNumArguments() - 1);
+  // Fix ownership of newly created non-trivial phis
+  if (asi->getFunction()->hasOwnership() &&
+      !asi->getType().isTrivial(*asi->getFunction())) {
+    // Go over all the proactively added phis and create end_lifetime at
+    // leaking blocks
+    SmallVector<SILBasicBlock *, 4> consumingBlocks;
+    for (auto *block : phiBlocks) {
+      auto *phiArg = cast<SILPhiArgument>(
+          block->getArgument(block->getNumArguments() - 1));
+      if (phiArg->use_empty()) {
+        auto *insertPt = phiArg->getNextInstruction();
+        SILBuilderWithScope builder(insertPt);
+        builder.createEndLifetime(
+            RegularLocation::getAutoGeneratedLocation(insertPt->getLoc()),
+            phiArg);
+      }
     }
   }
 }
 
 void StackAllocationPromoter::pruneAllocStackUsage() {
   LLVM_DEBUG(llvm::dbgs() << "*** Pruning : " << *asi);
-  BlockSet functionBlocks(asi->getFunction());
+  BlockSetVector functionBlocks(asi->getFunction());
 
   // Insert all of the blocks that asi is live in.
   for (auto *use : asi->getUses())
@@ -630,7 +658,7 @@ void StackAllocationPromoter::promoteAllocationToPhi() {
   LLVM_DEBUG(llvm::dbgs() << "*** Placing Phis for : " << *asi);
 
   // A list of blocks that will require new Phi values.
-  BlockSet phiBlocks(asi->getFunction());
+  BlockSetVector phiBlocks(asi->getFunction());
 
   // The "piggy-bank" data-structure that we use for processing the dom-tree
   // bottom-up.
@@ -1048,6 +1076,10 @@ bool MemoryToRegisters::promoteSingleAllocation(AllocStackInst *alloc) {
   LLVM_DEBUG(llvm::dbgs() << "*** Memory to register looking at: " << *alloc);
   ++NumAllocStackFound;
 
+  if (f.hasOwnership() && alloc->hasDynamicLifetime()) {
+    return false;
+  }
+
   // Don't handle captured AllocStacks.
   bool inSingleBlock = false;
   if (isCaptured(alloc, inSingleBlock)) {
@@ -1057,9 +1089,8 @@ bool MemoryToRegisters::promoteSingleAllocation(AllocStackInst *alloc) {
 
   // Remove write-only AllocStacks.
   if (isWriteOnlyAllocation(alloc)) {
-    deleter.forceDeleteWithUsers(alloc);
-
     LLVM_DEBUG(llvm::dbgs() << "*** Deleting store-only AllocStack: "<< *alloc);
+    deleter.forceDeleteWithUsers(alloc);
     return true;
   }
 

lib/SILOptimizer/Utils/InstOptUtils.cpp

@@ -2082,6 +2082,25 @@ bool swift::tryEliminateOnlyOwnershipUsedForwardingInst(
 
 // The consuming use blocks are assumed either not to inside a loop relative to
 // \p value or they must have their own copies.
+void swift::destroyValueAtLeakingBlocks(SILValue value,
+                                        ArrayRef<SILBasicBlock *> uses) {
+  if (!value->getFunction()->hasOwnership())
+    return;
+
+  if (value.getOwnershipKind() != OwnershipKind::Owned)
+    return;
+
+  findJointPostDominatingSet(
+      value->getParentBlock(), uses, [&](SILBasicBlock *loopBlock) {},
+      [&](SILBasicBlock *postDomBlock) {
+        // Insert a destroy_value in the leaking block
+        auto front = postDomBlock->begin();
+        SILBuilderWithScope newBuilder(front);
+        newBuilder.createDestroyValue(
+            RegularLocation::getAutoGeneratedLocation(), value);
+      });
+}
+
 void swift::endLifetimeAtLeakingBlocks(SILValue value,
                                        ArrayRef<SILBasicBlock *> uses) {
   if (!value->getFunction()->hasOwnership())
@@ -2096,7 +2115,7 @@ void swift::endLifetimeAtLeakingBlocks(SILValue value,
         // Insert a destroy_value in the leaking block
         auto front = postDomBlock->begin();
         SILBuilderWithScope newBuilder(front);
-        newBuilder.createDestroyValue(
+        newBuilder.createEndLifetime(
             RegularLocation::getAutoGeneratedLocation(), value);
       });
 }

test/SILOptimizer/dead_code_elimination_nontrivial_ossa.sil

@@ -611,3 +611,44 @@ bb3:
   return %9999 : $()
 }
 
+// CHECK-LABEL: sil [ossa] @dce_deadendlifetime1 :
+// CHECK-NOT: end_lifetime
+// CHECK-LABEL: } // end sil function 'dce_deadendlifetime1'
+sil [ossa] @dce_deadendlifetime1 : $@convention(thin) () -> () {
+bb0:
+  cond_br undef, bb1, bb2
+
+bb1:
+  br bb3(undef : $Klass)
+
+bb2:
+  br bb3(undef : $Klass)
+
+bb3(%2 : @owned $Klass):
+  end_lifetime %2 : $Klass
+  %res = tuple ()
+  return %res : $()
+}
+
+// CHECK-LABEL: sil [ossa] @dce_deadendlifetime2 :
+// CHECK-NOT: end_lifetime
+// CHECK-LABEL: } // end sil function 'dce_deadendlifetime2'
+sil [ossa] @dce_deadendlifetime2 : $@convention(thin) () -> () {
+bb0:
+  cond_br undef, bb1, bb2
+
+bb1:
+  br bb3(undef : $Klass)
+
+bb2:
+  br bb3(undef : $Klass)
+
+bb3(%2 : @owned $Klass):
+  br bb4(%2 : $Klass)
+
+bb4(%3 : @owned $Klass):
+  end_lifetime %3 : $Klass
+  %res = tuple ()
+  return %res : $()
+}
+