Merge pull request #34276 from meg-gupta/mem2regossa

[ownership] Enable SILMem2Reg for OSSA

Author: meg-gupta

include/swift/SIL/SILBuilder.h

@@ -1125,6 +1125,7 @@ class SILBuilder {
 
   UncheckedValueCastInst *createUncheckedValueCast(SILLocation Loc, SILValue Op,
                                                    SILType Ty) {
+    assert(hasOwnership());
     return insert(UncheckedValueCastInst::create(
         getSILDebugLocation(Loc), Op, Ty, getFunction(), C.OpenedArchetypes));
   }

lib/SILOptimizer/SILCombiner/SILCombinerCastVisitors.cpp

@@ -472,6 +472,9 @@ visitUncheckedTrivialBitCastInst(UncheckedTrivialBitCastInst *UTBCI) {
 SILInstruction *
 SILCombiner::
 visitUncheckedBitwiseCastInst(UncheckedBitwiseCastInst *UBCI) {
+  if (UBCI->getFunction()->hasOwnership())
+    return nullptr;
+
   // (unchecked_bitwise_cast Y->Z (unchecked_bitwise_cast X->Y x))
   // OR (unchecked_trivial_cast Y->Z (unchecked_bitwise_cast X->Y x))
   //   ->

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).
@@ -302,8 +323,10 @@ promoteDebugValueAddr(DebugValueAddrInst *DVAI, SILValue Value, SILBuilder &B) {
   // Avoid inserting the same debug_value twice.
   for (Operand *Use : Value->getUses())
     if (auto *DVI = dyn_cast<DebugValueInst>(Use->getUser()))
-      if (*DVI->getVarInfo() == *DVAI->getVarInfo())
+      if (*DVI->getVarInfo() == *DVAI->getVarInfo()) {
+        DVAI->eraseFromParent();
         return;
+      }
   B.setInsertionPoint(DVAI);
   B.setCurrentDebugScope(DVAI->getDebugScope());
   B.createDebugValue(DVAI->getLoc(), Value, *DVAI->getVarInfo());
@@ -348,21 +371,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 +455,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 +472,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 +494,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,10 +562,23 @@ 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)
+      if (DSI->getOperand() == ASI) {
+        // Reset LastStore.
+        // So that we don't pass RunningVal as a phi arg beyond dealloc_stack
+        LastStore = nullptr;
         break;
+      }
     }
   }
   if (LastStore) {
@@ -512,6 +621,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;
@@ -643,6 +756,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;
@@ -679,6 +807,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);
@@ -710,6 +848,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() {
@@ -956,10 +1110,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>
@@ -465,3 +465,30 @@ bb0(%0 : $Optional<Klass>):
   %4 = tuple()
   return %4 : $()
 }
+
+// CHECK-LABEL: sil @multi_basic_block_use_on_one_path :
+// CHECK-NOT: alloc_stack
+// CHECK:bb2:
+// CHECK: br bb3(undef : $Klass)
+// CHECK-LABEL: } // end sil function 'multi_basic_block_use_on_one_path'
+sil @multi_basic_block_use_on_one_path : $@convention(thin) (@owned Klass) -> () {
+bb0(%0 : $Klass):
+  %1 = alloc_stack $Klass
+  cond_br undef, bb1, bb2
+
+bb1:
+  dealloc_stack %1 : $*Klass
+  strong_release %0 : $Klass
+  br bb3
+
+bb2:
+  store %0 to %1 : $*Klass
+  %7 = load %1 : $*Klass
+  dealloc_stack %1 : $*Klass
+  strong_release %7 : $Klass
+  br bb3
+
+bb3:
+  %11 = tuple ()
+  return %11 : $()
+}