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

Add utility to create end_lifetime at leaking blocks for an owned value

Author: meg-gupta

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/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

@@ -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())) {
@@ -330,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(),
@@ -346,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.
@@ -371,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) {
@@ -416,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 &&
@@ -461,11 +474,17 @@ void StackAllocationPromoter::addBlockArguments(BlockSetVector &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())
@@ -487,7 +506,7 @@ SILValue StackAllocationPromoter::getLiveOutValue(BlockSetVector &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(BlockSetVector &phiBlocks,
@@ -596,14 +615,22 @@ void StackAllocationPromoter::fixBranchesAndUses(BlockSetVector &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);
+      }
     }
   }
 }
@@ -1049,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)) {
@@ -1058,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/mem2reg.sil

@@ -466,11 +466,11 @@ bb0(%0 : $Optional<Klass>):
   return %4 : $()
 }
 
-// check no dead args are passed to bb3
+// dead args okay in non-ossa
 // CHECK-LABEL: sil @multi_basic_block_use_on_one_path :
 // CHECK-NOT: alloc_stack
-// CHECK: bb3:
-// CHECK-LABEL: } // end sil function 'multi_basic_block_use_on_one_path'
+// CHECK: br bb3(undef : $Klass)
+// CHECK: bb3([[dead_arg:%.*]]):
 sil @multi_basic_block_use_on_one_path : $@convention(thin) (@owned Klass) -> () {
 bb0(%0 : $Klass):
   %1 = alloc_stack $Klass