Fix Mem2Reg for store [assign]

lib/SILOptimizer/Transforms/SILMem2Reg.cpp

@@ -495,46 +495,30 @@ StackAllocationPromoter::promoteAllocationInBlock(SILBasicBlock *BB) {
       if (SI->getDest() != ASI)
         continue;
 
-      // Special handling of entry block
-      // If we have a store [assign] in the first block, OSSA guarantees we can
-      // find the previous value stored in the stack location in RunningVal.
-      // Create destroy_value of the RunningVal.
-      // For all other blocks we may not know the previous value stored in the
-      // stack location. So we will create destroy_value in
-      // StackAllocationPromoter::fixBranchesAndUses, by getting the live-in
-      // value to the block.
-      if (BB->isEntry()) {
-        if (SI->getOwnershipQualifier() == StoreOwnershipQualifier::Assign) {
-          assert(RunningVal);
+      // If we see a store [assign], always convert it to a store [init]. This
+      // simplifies further processing.
+      if (SI->getOwnershipQualifier() == StoreOwnershipQualifier::Assign) {
+        if (RunningVal) {
           SILBuilderWithScope(SI).createDestroyValue(SI->getLoc(), RunningVal);
+        } else {
+          SILBuilderWithScope localBuilder(SI);
+          auto *newLoad = localBuilder.createLoad(SI->getLoc(), ASI,
+                                                  LoadOwnershipQualifier::Take);
+          localBuilder.createDestroyValue(SI->getLoc(), newLoad);
         }
+        SI->setOwnershipQualifier(StoreOwnershipQualifier::Init);
       }
 
       // If we met a store before this one, delete it.
-      // If the LastStore was a store with [assign], delete it only if we know
-      // the RunningValue to destroy. If not, it will be deleted in
-      // StackAllocationPromoter::fixBranchesAndUses.
       if (LastStore) {
-        if (LastStore->getOwnershipQualifier() ==
-            StoreOwnershipQualifier::Assign) {
-          if (RunningVal) {
-            // For entry block, we would have already created the destroy_value,
-            // skip it.
-            if (!BB->isEntry()) {
-              SILBuilderWithScope(LastStore).createDestroyValue(
-                  LastStore->getLoc(), RunningVal);
-            }
-            LLVM_DEBUG(llvm::dbgs()
-                       << "*** Removing redundant store: " << *LastStore);
-            ++NumInstRemoved;
-            LastStore->eraseFromParent();
-          }
-        } else {
-          LLVM_DEBUG(llvm::dbgs()
-                     << "*** Removing redundant store: " << *LastStore);
-          ++NumInstRemoved;
-          LastStore->eraseFromParent();
-        }
+        assert(LastStore->getOwnershipQualifier() !=
+                   StoreOwnershipQualifier::Assign &&
+               "store [assign] to the stack location should have been "
+               "transformed to a store [init]");
+        LLVM_DEBUG(llvm::dbgs()
+                   << "*** Removing redundant store: " << *LastStore);
+        ++NumInstRemoved;
+        LastStore->eraseFromParent();
       }
 
       // The stored value is the new running value.
@@ -578,7 +562,12 @@ StackAllocationPromoter::promoteAllocationInBlock(SILBasicBlock *BB) {
         break;
     }
   }
+
   if (LastStore) {
+    assert(LastStore->getOwnershipQualifier() !=
+               StoreOwnershipQualifier::Assign &&
+           "store [assign] to the stack location should have been "
+           "transformed to a store [init]");
     LLVM_DEBUG(llvm::dbgs() << "*** Finished promotion. Last store: "
                             << *LastStore);
   } else {
@@ -756,6 +745,7 @@ void StackAllocationPromoter::fixPhiPredBlock(BlockSet &PhiBlocks,
 void StackAllocationPromoter::fixBranchesAndUses(BlockSet &PhiBlocks) {
   // First update uses of the value.
   SmallVector<LoadInst *, 4> collectedLoads;
+
   for (auto UI = ASI->use_begin(), E = ASI->use_end(); UI != E;) {
     auto *Inst = UI->getUser();
     ++UI;
@@ -789,16 +779,6 @@ void StackAllocationPromoter::fixBranchesAndUses(BlockSet &PhiBlocks) {
     // on.
     SILBasicBlock *BB = Inst->getParent();
 
-    if (!BB->isEntry()) {
-      if (auto *SI = dyn_cast<StoreInst>(Inst)) {
-        if (SI->getOwnershipQualifier() == StoreOwnershipQualifier::Assign) {
-          SILValue Def = getLiveInValue(PhiBlocks, BB);
-          SILBuilderWithScope(SI).createDestroyValue(SI->getLoc(), Def);
-          continue;
-        }
-      }
-    }
-
     if (auto *DVAI = dyn_cast<DebugValueAddrInst>(Inst)) {
       // Replace DebugValueAddr with DebugValue.
       SILValue Def = getLiveInValue(PhiBlocks, BB);

test/SILOptimizer/mem2reg_ossa_nontrivial.sil

@@ -188,6 +188,84 @@ bb3:
   return %ret : $()
 }
 
+// CHECK-LABEL: sil [ossa] @multiple_store_vals6 :
+// CHECK-NOT: alloc_stack
+// CHECK-LABEL: } // end sil function 'multiple_store_vals6'
+sil [ossa] @multiple_store_vals6 : $@convention(thin) (@owned Klass, @owned Klass) -> () {
+bb0(%0 : @owned $Klass, %1 : @owned $Klass):
+  br bb1
+
+bb1:
+  %2 = alloc_stack $Klass
+  store %0 to [init] %2 : $*Klass
+  %3 = integer_literal $Builtin.Int1, 0
+  cond_fail %3 : $Builtin.Int1
+  store %1 to [assign] %2 : $*Klass
+  br bb2
+
+bb2:
+  %4 = function_ref @blackhole_spl : $@convention(thin) (@guaranteed Klass) -> ()
+  %5 = load [take] %2 : $*Klass
+  %6 = apply %4(%5) : $@convention(thin) (@guaranteed Klass) -> ()
+  destroy_value %5 : $Klass
+  dealloc_stack %2 : $*Klass
+  %7 = tuple ()
+  return %7 : $()
+}
+
+// CHECK-LABEL: sil [ossa] @multiple_store_vals7 :
+// CHECK-NOT: alloc_stack
+// CHECK-LABEL: } // end sil function 'multiple_store_vals7'
+sil [ossa] @multiple_store_vals7 : $@convention(thin) (@owned Klass, @owned Klass, @owned Klass) -> () {
+bb0(%0 : @owned $Klass, %1 : @owned $Klass, %2 : @owned $Klass):
+  %stk = alloc_stack $Klass
+  store %0 to [init] %stk : $*Klass
+  %3 = integer_literal $Builtin.Int1, 0
+  cond_fail %3 : $Builtin.Int1
+  br bb1
+
+bb1:
+  store %1 to [assign] %stk : $*Klass
+  store %2 to [assign] %stk : $*Klass
+  br bb2
+
+bb2:
+  %4 = function_ref @blackhole_spl : $@convention(thin) (@guaranteed Klass) -> ()
+  %5 = load [take] %stk : $*Klass
+  %6 = apply %4(%5) : $@convention(thin) (@guaranteed Klass) -> ()
+  destroy_value %5 : $Klass
+  dealloc_stack %stk : $*Klass
+  %7 = tuple ()
+  return %7 : $()
+}
+
+// CHECK-LABEL: sil [ossa] @multiple_store_vals8 :
+// CHECK-NOT: alloc_stack
+// CHECK-LABEL: } // end sil function 'multiple_store_vals8'
+sil [ossa] @multiple_store_vals8 : $@convention(thin) (@owned Klass, @owned Klass, @owned Klass) -> () {
+bb0(%0 : @owned $Klass, %1 : @owned $Klass, %2 : @owned $Klass):
+  %stk = alloc_stack $Klass
+  store %0 to [init] %stk : $*Klass
+  %3 = integer_literal $Builtin.Int1, 0
+  cond_fail %3 : $Builtin.Int1
+  br bb1
+
+bb1:
+  store %1 to [assign] %stk : $*Klass
+  destroy_addr %stk : $*Klass
+  store %2 to [init] %stk : $*Klass
+  br bb2
+
+bb2:
+  %4 = function_ref @blackhole_spl : $@convention(thin) (@guaranteed Klass) -> ()
+  %5 = load [take] %stk : $*Klass
+  %6 = apply %4(%5) : $@convention(thin) (@guaranteed Klass) -> ()
+  destroy_value %5 : $Klass
+  dealloc_stack %stk : $*Klass
+  %7 = tuple ()
+  return %7 : $()
+}
+
 // CHECK-LABEL: sil [ossa] @with_loads :
 // CHECK-NOT: alloc_stack
 // CHECK-LABEL: } // end sil function 'with_loads'
@@ -676,3 +754,4 @@ bbret(%new : @owned $Klass):
   dealloc_stack %stk1 : $*Klass
   return %new : $Klass
 }
+