Enable SILMem2Reg for OSSA

SILMem2Reg has roughly 2 central algorithms, removal of alloc_stack with uses in a
single block vs multiple blocks. While replacing loads and stores to the
stack location in each of the 2 algorithms, new handling of qualifiers
like [assign], [copy] and [take] which are new to OSSA are needed.

Also Disable SILMem2Reg when we see this pattern:
load [take] (struct_element_addr/tuple_element_addr %ASI)

Convert SILMem2Reg tests into ossa
And add new SILMem2Reg tests for non-trivial values.
Thanks to zoecarver for additional tests.

Author: meg-gupta

lib/SILOptimizer/Transforms/SILMem2Reg.cpp

@@ -183,23 +183,43 @@ class MemoryToRegisters {
 /// Returns true if \p I is an address of a LoadInst, skipping struct and
 /// tuple address projections. Sets \p singleBlock to null if the load (or
 /// it's address is not in \p singleBlock.
-static bool isAddressForLoad(SILInstruction *I, SILBasicBlock *&singleBlock) {
-  
-  if (isa<LoadInst>(I))
+/// This function looks for these patterns:
+/// 1. (load %ASI)
+/// 2. (load (struct_element_addr/tuple_element_addr/unchecked_addr_cast %ASI))
+static bool isAddressForLoad(SILInstruction *I, SILBasicBlock *&singleBlock,
+                             bool &hasGuaranteedOwnership) {
+
+  if (isa<LoadInst>(I)) {
+    // SILMem2Reg is disabled when we find:
+    // (load [take] (struct_element_addr/tuple_element_addr %ASI))
+    // struct_element_addr and tuple_element_addr are lowered into
+    // struct_extract and tuple_extract and these SIL instructions have a
+    // guaranteed ownership. For replacing load's users, we need an owned value.
+    // We will need a new copy and destroy of the running val placed after the
+    // last use. This is not implemented currently.
+    if (hasGuaranteedOwnership && cast<LoadInst>(I)->getOwnershipQualifier() ==
+                                      LoadOwnershipQualifier::Take) {
+      return false;
+    }
     return true;
+  }
 
   if (!isa<UncheckedAddrCastInst>(I) && !isa<StructElementAddrInst>(I) &&
       !isa<TupleElementAddrInst>(I))
     return false;
-  
+
+  if (isa<StructElementAddrInst>(I) || isa<TupleElementAddrInst>(I)) {
+    hasGuaranteedOwnership = true;
+  }
+
   // Recursively search for other (non-)loads in the instruction's uses.
   for (auto UI : cast<SingleValueInstruction>(I)->getUses()) {
     SILInstruction *II = UI->getUser();
     if (II->getParent() != singleBlock)
       singleBlock = nullptr;
-    
-    if (!isAddressForLoad(II, singleBlock))
-        return false;
+
+    if (!isAddressForLoad(II, singleBlock, hasGuaranteedOwnership))
+      return false;
   }
   return true;
 }
@@ -233,7 +253,8 @@ static bool isCaptured(AllocStackInst *ASI, bool &inSingleBlock) {
       singleBlock = nullptr;
 
     // Loads are okay.
-    if (isAddressForLoad(II, singleBlock))
+    bool hasGuaranteedOwnership = false;
+    if (isAddressForLoad(II, singleBlock, hasGuaranteedOwnership))
       continue;
 
     // We can store into an AllocStack (but not the pointer).
@@ -348,21 +369,54 @@ static void collectLoads(SILInstruction *I, SmallVectorImpl<LoadInst *> &Loads)
 static void replaceLoad(LoadInst *LI, SILValue val, AllocStackInst *ASI) {
   ProjectionPath projections(val->getType());
   SILValue op = LI->getOperand();
+  SILBuilderWithScope builder(LI);
+
   while (op != ASI) {
     assert(isa<UncheckedAddrCastInst>(op) || isa<StructElementAddrInst>(op) ||
            isa<TupleElementAddrInst>(op));
     auto *Inst = cast<SingleValueInstruction>(op);
     projections.push_back(Projection(Inst));
     op = Inst->getOperand(0);
   }
-  SILBuilder builder(LI);
+
+  SmallVector<SILValue, 4> borrowedVals;
   for (auto iter = projections.rbegin(); iter != projections.rend(); ++iter) {
     const Projection &projection = *iter;
+    assert(projection.getKind() == ProjectionKind::BitwiseCast ||
+           projection.getKind() == ProjectionKind::Struct ||
+           projection.getKind() == ProjectionKind::Tuple);
+
+    // struct_extract and tuple_extract expect guaranteed operand ownership
+    // non-trivial RunningVal is owned. Insert borrow operation to convert
+    if (projection.getKind() == ProjectionKind::Struct ||
+        projection.getKind() == ProjectionKind::Tuple) {
+      SILValue opVal = builder.emitBeginBorrowOperation(LI->getLoc(), val);
+      if (opVal != val) {
+        borrowedVals.push_back(opVal);
+        val = opVal;
+      }
+    }
     val = projection.createObjectProjection(builder, LI->getLoc(), val).get();
   }
+
   op = LI->getOperand();
-  LI->replaceAllUsesWith(val);
+  // Replace users of the loaded value with `val`
+  // If we have a load [copy], replace the users with copy_value of `val`
+  if (LI->getOwnershipQualifier() == LoadOwnershipQualifier::Copy) {
+    LI->replaceAllUsesWith(builder.createCopyValue(LI->getLoc(), val));
+  } else {
+    assert(!ASI->getFunction()->hasOwnership() ||
+           val.getOwnershipKind() != ValueOwnershipKind::Guaranteed);
+    LI->replaceAllUsesWith(val);
+  }
+
+  for (auto borrowedVal : borrowedVals) {
+    builder.emitEndBorrowOperation(LI->getLoc(), borrowedVal);
+  }
+
+  // Delete the load
   LI->eraseFromParent();
+
   while (op != ASI && op->use_empty()) {
     assert(isa<UncheckedAddrCastInst>(op) || isa<StructElementAddrInst>(op) ||
            isa<TupleElementAddrInst>(op));
@@ -399,6 +453,7 @@ StoreInst *
 StackAllocationPromoter::promoteAllocationInBlock(SILBasicBlock *BB) {
   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.
@@ -415,12 +470,16 @@ StackAllocationPromoter::promoteAllocationInBlock(SILBasicBlock *BB) {
         // If we are loading from the AllocStackInst and we already know the
         // content of the Alloca then use it.
         LLVM_DEBUG(llvm::dbgs() << "*** Promoting load: " << *Load);
-        
         replaceLoad(Load, RunningVal, ASI);
         ++NumInstRemoved;
-      } else if (Load->getOperand() == ASI) {
+      } else if (Load->getOperand() == ASI &&
+                 Load->getOwnershipQualifier() !=
+                     LoadOwnershipQualifier::Copy) {
         // If we don't know the content of the AllocStack then the loaded
         // value *is* the new value;
+        // Don't use result of load [copy] as a RunningVal, it necessitates
+        // additional logic for cleanup of consuming instructions of the result.
+        // StackAllocationPromoter::fixBranchesAndUses will later handle it.
         LLVM_DEBUG(llvm::dbgs() << "*** First load: " << *Load);
         RunningVal = Load;
       }
@@ -433,16 +492,51 @@ StackAllocationPromoter::promoteAllocationInBlock(SILBasicBlock *BB) {
       if (SI->getDest() != ASI)
         continue;
 
-      // The stored value is the new running value.
-      RunningVal = SI->getSrc();
+      // 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);
+          SILBuilderWithScope(SI).createDestroyValue(SI->getLoc(), RunningVal);
+        }
+      }
 
       // 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) {
-        ++NumInstRemoved;
-        LLVM_DEBUG(llvm::dbgs() << "*** Removing redundant store: "
-                                << *LastStore);
-        LastStore->eraseFromParent();
+        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();
+        }
       }
+
+      // The stored value is the new running value.
+      RunningVal = SI->getSrc();
+      // The current store is now the LastStore
       LastStore = SI;
       continue;
     }
@@ -466,6 +560,15 @@ StackAllocationPromoter::promoteAllocationInBlock(SILBasicBlock *BB) {
       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) {
@@ -516,6 +619,10 @@ void MemoryToRegisters::removeSingleBlockAllocation(AllocStackInst *ASI) {
     // value.
     if (auto *SI = dyn_cast<StoreInst>(Inst)) {
       if (SI->getDest() == ASI) {
+        if (SI->getOwnershipQualifier() == StoreOwnershipQualifier::Assign) {
+          assert(RunningVal);
+          SILBuilderWithScope(SI).createDestroyValue(SI->getLoc(), RunningVal);
+        }
         RunningVal = SI->getSrc();
         Inst->eraseFromParent();
         ++NumInstRemoved;
@@ -647,6 +754,21 @@ void StackAllocationPromoter::fixPhiPredBlock(BlockSet &PhiBlocks,
   TI->eraseFromParent();
 }
 
+static bool hasOnlyUndefIncomingValues(SILPhiArgument *phiArg) {
+  SmallVector<SILValue, 8> incomingValues;
+  phiArg->getIncomingPhiValues(incomingValues);
+  for (auto predArg : incomingValues) {
+    if (isa<SILUndef>(predArg))
+      continue;
+    if (isa<SILPhiArgument>(predArg) &&
+        hasOnlyUndefIncomingValues(cast<SILPhiArgument>(predArg))) {
+      continue;
+    }
+    return false;
+  }
+  return true;
+}
+
 void StackAllocationPromoter::fixBranchesAndUses(BlockSet &PhiBlocks) {
   // First update uses of the value.
   SmallVector<LoadInst *, 4> collectedLoads;
@@ -683,6 +805,16 @@ 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);
@@ -714,6 +846,22 @@ void StackAllocationPromoter::fixBranchesAndUses(BlockSet &PhiBlocks) {
       fixPhiPredBlock(PhiBlocks, Block, PBB);
     }
   }
+
+  // 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.
+  if (ASI->getFunction()->hasOwnership()) {
+    for (auto Block : PhiBlocks) {
+      auto *phiArg = cast<SILPhiArgument>(
+          Block->getArgument(Block->getNumArguments() - 1));
+      if (phiArg->getOwnershipKind() == ValueOwnershipKind::Owned &&
+          phiArg->use_empty()) {
+        assert(hasOnlyUndefIncomingValues(phiArg));
+        SILBuilderWithScope(&Block->front())
+            .createEndLifetime(Block->front().getLoc(), phiArg);
+      }
+    }
+  }
 }
 
 void StackAllocationPromoter::pruneAllocStackUsage() {
@@ -960,10 +1108,6 @@ class SILMem2Reg : public SILFunctionTransform {
   void run() override {
     SILFunction *F = getFunction();
 
-    // FIXME: We should be able to support ownership.
-    if (F->hasOwnership())
-      return;
-
     LLVM_DEBUG(llvm::dbgs() << "** Mem2Reg on function: " << F->getName()
                             << " **\n");
 

test/SILOptimizer/mem2reg.sil

@@ -449,12 +449,12 @@ bb0:
   return %16 : $((), ())
 }
 
-// CHECK-LABEL: sil @unchecked_ref_cast
+// CHECK-LABEL: sil @unchecked_bitwise_cast
 // CHECK-NOT: alloc_stack
 // CHECK: [[CAST:%.*]] = unchecked_bitwise_cast %0 : $Optional<Klass> to $Klass
 // CHECK: release_value [[CAST]]
 // CHECK: return
-sil @unchecked_ref_cast : $@convention(thin) (@owned Optional<Klass>) -> () {
+sil @unchecked_bitwise_cast : $@convention(thin) (@owned Optional<Klass>) -> () {
 bb0(%0 : $Optional<Klass>):
   %1 = alloc_stack $Optional<Klass>
   store %0 to %1 : $*Optional<Klass>

test/SILOptimizer/mem2reg_liveness_ossa.sil

@@ -0,0 +1,65 @@
+// RUN: %target-sil-opt -enable-sil-verify-all %s -mem2reg | %FileCheck %s
+
+import Builtin
+import Swift
+
+sil [ossa] @_Ts5printFT3valSi_T_ : $@convention(thin) (Int64) -> ()
+
+// CHECK-LABEL: sil [ossa] @liveness0 :
+// CHECK-NOT: alloc_stack
+sil [ossa] @liveness0 : $@convention(thin) (Int64) -> () {
+bb0(%0 : $Int64):
+  %1 = alloc_stack $Int64, var, name "x"
+  store %0 to [trivial] %1 : $*Int64
+  %3 = integer_literal $Builtin.Int64, 10
+  %5 = struct_extract %0 : $Int64, #Int64._value
+  %6 = builtin "cmp_eq_Int64"(%5 : $Builtin.Int64, %3 : $Builtin.Int64) : $Builtin.Int1
+  cond_br %6, bb1, bb5
+
+bb1:
+  %8 = alloc_stack $Int64, var, name "y"
+  %9 = integer_literal $Builtin.Int64, 20
+  %10 = struct $Int64 (%9 : $Builtin.Int64)
+  store %10 to [trivial] %8 : $*Int64
+  %12 = integer_literal $Builtin.Int64, 3
+  %14 = struct_extract %0 : $Int64, #Int64._value
+  %15 = builtin "cmp_sgt_Int64"(%14 : $Builtin.Int64, %12 : $Builtin.Int64) : $Builtin.Int1
+  cond_br %15, bb2, bb3
+
+bb2:
+  %17 = integer_literal $Builtin.Int64, 0
+  %18 = struct $Int64 (%17 : $Builtin.Int64)
+  store %18 to [trivial] %8 : $*Int64
+  br bb4
+
+bb3:
+  %21 = integer_literal $Builtin.Int64, 2
+  %22 = struct $Int64 (%21 : $Builtin.Int64)
+  store %22 to [trivial] %8 : $*Int64
+  br bb4
+
+bb4:
+  // function_ref
+  %25 = function_ref @_Ts5printFT3valSi_T_ : $@convention(thin) (Int64) -> ()
+  %26 = load [trivial] %8 : $*Int64
+  %27 = integer_literal $Builtin.Int64, 2
+  %28 = integer_literal $Builtin.Int1, -1
+  %30 = struct_extract %26 : $Int64, #Int64._value
+  %31 = builtin "sadd_with_overflow_Int64"(%30 : $Builtin.Int64, %27 : $Builtin.Int64, %28 : $Builtin.Int1) : $(Builtin.Int64, Builtin.Int1)
+  %32 = tuple_extract %31 : $(Builtin.Int64, Builtin.Int1), 0
+  %33 = tuple_extract %31 : $(Builtin.Int64, Builtin.Int1), 1
+  %34 = struct $Int64 (%32 : $Builtin.Int64)
+  cond_fail %33 : $Builtin.Int1
+  %36 = apply %25(%34) : $@convention(thin) (Int64) -> ()
+  dealloc_stack %8 : $*Int64
+  br bb5
+
+// We don't need a PHI node here because the value is dead!
+// CHECK: bb5:
+bb5:
+  dealloc_stack %1 : $*Int64
+  %40 = tuple ()
+  return %40 : $()
+}
+// CHECK-LABEL: } // end sil function 'liveness0'
+