Enable Mem2Reg for allocs with load_borrow (#59350)

Author: meg-gupta

lib/SILOptimizer/Transforms/SILMem2Reg.cpp

@@ -202,9 +202,14 @@ static void promoteDebugValueAddr(DebugValueInst *dvai, SILValue value,
 
 /// Returns true if \p I is a load which loads from \p ASI.
 static bool isLoadFromStack(SILInstruction *i, AllocStackInst *asi) {
-  if (!isa<LoadInst>(i))
+  if (!isa<LoadInst>(i) && !isa<LoadBorrowInst>(i))
     return false;
 
+  if (auto *lbi = dyn_cast<LoadBorrowInst>(i)) {
+    if (BorrowedValue(lbi).hasReborrow())
+      return false;
+  }
+
   // Skip struct and tuple address projections.
   ValueBase *op = i->getOperand(0);
   while (op != asi) {
@@ -218,9 +223,9 @@ static bool isLoadFromStack(SILInstruction *i, AllocStackInst *asi) {
 
 /// Collects all load instructions which (transitively) use \p I as address.
 static void collectLoads(SILInstruction *i,
-                         SmallVectorImpl<LoadInst *> &foundLoads) {
-  if (auto *load = dyn_cast<LoadInst>(i)) {
-    foundLoads.push_back(load);
+                         SmallVectorImpl<SILInstruction *> &foundLoads) {
+  if (isa<LoadInst>(i) || isa<LoadBorrowInst>(i)) {
+    foundLoads.push_back(i);
     return;
   }
   if (!isa<UncheckedAddrCastInst>(i) && !isa<StructElementAddrInst>(i) &&
@@ -234,22 +239,23 @@ static void collectLoads(SILInstruction *i,
 }
 
 static void
-replaceLoad(LoadInst *li, SILValue newValue, AllocStackInst *asi,
+replaceLoad(SILInstruction *inst, SILValue newValue, AllocStackInst *asi,
             SILBuilderContext &ctx, InstructionDeleter &deleter,
             SmallVectorImpl<SILInstruction *> &instructionsToDelete) {
+  assert(isa<LoadInst>(inst) || isa<LoadBorrowInst>(inst));
   ProjectionPath projections(newValue->getType());
-  SILValue op = li->getOperand();
-  SILBuilderWithScope builder(li, ctx);
+  SILValue op = inst->getOperand(0);
+  SILBuilderWithScope builder(inst, ctx);
   SILOptScope scope;
 
   while (op != asi) {
     assert(isa<UncheckedAddrCastInst>(op) || isa<StructElementAddrInst>(op) ||
            isa<TupleElementAddrInst>(op) &&
                "found instruction that should have been skipped in "
                "isLoadFromStack");
-    auto *inst = cast<SingleValueInstruction>(op);
-    projections.push_back(Projection(inst));
-    op = inst->getOperand(0);
+    auto *projInst = cast<SingleValueInstruction>(op);
+    projections.push_back(Projection(projInst));
+    op = projInst->getOperand(0);
   }
 
   for (const auto &proj : llvm::reverse(projections)) {
@@ -262,33 +268,51 @@ replaceLoad(LoadInst *li, SILValue newValue, AllocStackInst *asi,
     // to guaranteed!
     if (proj.getKind() == ProjectionKind::Struct ||
         proj.getKind() == ProjectionKind::Tuple) {
-      if (auto opVal = scope.borrowValue(li, newValue)) {
+      if (auto opVal = scope.borrowValue(inst, newValue)) {
         assert(*opVal != newValue &&
                "Valid value should be different from input value");
         newValue = *opVal;
       }
     }
     newValue =
-        proj.createObjectProjection(builder, li->getLoc(), newValue).get();
+        proj.createObjectProjection(builder, inst->getLoc(), newValue).get();
   }
 
-  op = li->getOperand();
+  op = inst->getOperand(0);
 
-  // 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(), newValue));
+  if (auto *lbi = dyn_cast<LoadBorrowInst>(inst)) {
+    if (shouldAddLexicalLifetime(asi)) {
+      assert(isa<BeginBorrowInst>(newValue));
+      SmallVector<SILInstruction *, 4> endBorrows;
+      for (auto *ebi : lbi->getUsersOfType<EndBorrowInst>()) {
+        endBorrows.push_back(ebi);
+      }
+      for (auto *ebi : endBorrows) {
+        ebi->eraseFromParent();
+      }
+      lbi->replaceAllUsesWith(newValue);
+    }
+    else {
+      auto *borrow = SILBuilderWithScope(lbi, ctx).createBeginBorrow(
+          lbi->getLoc(), newValue, asi->isLexical());
+      lbi->replaceAllUsesWith(borrow);
+    }
   } else {
-    assert(!asi->getFunction()->hasOwnership() ||
-           newValue.getOwnershipKind() != OwnershipKind::Guaranteed);
-    li->replaceAllUsesWith(newValue);
+    auto *li = cast<LoadInst>(inst);
+    // Replace users of the loaded value with `newValue`
+    // If we have a load [copy], replace the users with copy_value of `newValue`
+    if (li->getOwnershipQualifier() == LoadOwnershipQualifier::Copy) {
+      li->replaceAllUsesWith(builder.createCopyValue(li->getLoc(), newValue));
+    } else {
+      li->replaceAllUsesWith(newValue);
+    }
   }
 
   // Pop the scope so that we emit cleanups.
   std::move(scope).popAtEndOfScope(&*builder.getInsertionPoint());
 
   // Delete the load
-  prepareForDeletion(li, instructionsToDelete);
+  prepareForDeletion(inst, instructionsToDelete);
 
   while (op != asi && op->use_empty()) {
     assert(isa<UncheckedAddrCastInst>(op) || isa<StructElementAddrInst>(op) ||
@@ -577,6 +601,20 @@ StoreInst *StackAllocationPromoter::promoteAllocationInBlock(
 
     if (isLoadFromStack(inst, asi)) {
       assert(!runningVals || runningVals->isStorageValid);
+      if (auto *lbi = dyn_cast<LoadBorrowInst>(inst)) {
+        if (runningVals) {
+          if (shouldAddLexicalLifetime(asi)) {
+            replaceLoad(lbi, runningVals->value.borrow, asi, ctx,
+                        deleter, instructionsToDelete);
+          }
+          else {
+            replaceLoad(lbi, runningVals->value.replacement(asi), asi, ctx,
+                        deleter, instructionsToDelete);
+          }
+          ++NumInstRemoved;
+        }
+        continue;
+      }
       auto *li = cast<LoadInst>(inst);
       if (li->getOwnershipQualifier() == LoadOwnershipQualifier::Take) {
         if (shouldAddLexicalLifetime(asi)) {
@@ -597,12 +635,13 @@ StoreInst *StackAllocationPromoter::promoteAllocationInBlock(
         if (lastStoreInst)
           lastStoreInst->isStorageValid = false;
       }
+
       if (runningVals) {
         // 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: " << *li);
-        replaceLoad(li, runningVals->value.replacement(asi), asi, ctx, deleter,
-                    instructionsToDelete);
+        replaceLoad(inst, runningVals->value.replacement(asi), asi, ctx,
+                    deleter, instructionsToDelete);
         ++NumInstRemoved;
       } else if (li->getOperand() == asi &&
                  li->getOwnershipQualifier() != LoadOwnershipQualifier::Copy) {
@@ -789,10 +828,8 @@ StackAllocationPromoter::getLiveOutValues(BlockSetVector &phiBlocks,
       SILValue borrow = SILValue();
       SILValue copy = SILValue();
       if (shouldAddLexicalLifetime(asi)) {
-        auto *bbi = cast<BeginBorrowInst>(&*std::next(si->getIterator()));
-        borrow = bbi;
-        auto *cvi = cast<CopyValueInst>(bbi->getNextInstruction());
-        copy = cvi;
+        borrow = cast<BeginBorrowInst>(&*std::next(si->getIterator()));
+        copy = cast<CopyValueInst>(borrow->getNextInstruction());
       }
       LiveValues values = {stored, borrow, copy};
       return values;
@@ -953,7 +990,7 @@ void StackAllocationPromoter::propagateLiveness(
 void StackAllocationPromoter::fixBranchesAndUses(BlockSetVector &phiBlocks,
                                                  BlockSetVector &phiBlocksOut) {
   // First update uses of the value.
-  SmallVector<LoadInst *, 4> collectedLoads;
+  SmallVector<SILInstruction *, 4> collectedLoads;
 
   for (auto ui = asi->use_begin(), ue = asi->use_end(); ui != ue;) {
     auto *user = ui->getUser();
@@ -1421,6 +1458,13 @@ static bool isAddressForLoad(SILInstruction *load, SILBasicBlock *&singleBlock,
     return true;
   }
 
+  if (isa<LoadBorrowInst>(load)) {
+    if (involvesUntakableProjection) {
+      return false;
+    }
+    return true;
+  }
+
   if (!isa<UncheckedAddrCastInst>(load) && !isa<StructElementAddrInst>(load) &&
       !isa<TupleElementAddrInst>(load))
     return false;
@@ -1480,8 +1524,8 @@ static bool isDeadAddrProjection(SILInstruction *inst) {
 }
 
 /// Returns true if this AllocStacks is captured.
-/// Sets \p inSingleBlock to true if all uses of \p ASI are in a single block.
-static bool isCaptured(AllocStackInst *asi, bool &inSingleBlock) {
+/// Sets \p inSingleBlock to true if all uses of \p asi are in a single block.
+static bool isCaptured(AllocStackInst *asi, bool *inSingleBlock) {
   SILBasicBlock *singleBlock = asi->getParent();
 
   // For all users of the AllocStack instruction.
@@ -1518,7 +1562,7 @@ static bool isCaptured(AllocStackInst *asi, bool &inSingleBlock) {
   }
 
   // None of the users capture the AllocStack.
-  inSingleBlock = (singleBlock != nullptr);
+  *inSingleBlock = (singleBlock != nullptr);
   return false;
 }
 
@@ -1533,6 +1577,10 @@ bool MemoryToRegisters::isWriteOnlyAllocation(AllocStackInst *asi) {
       if (!isa<AllocStackInst>(si->getSrc()))
         continue;
 
+    if (auto *sbi = dyn_cast<StoreBorrowInst>(user))
+      if (!isa<AllocStackInst>(sbi->getSrc()))
+        continue;
+
     // Deallocation is also okay.
     if (isa<DeallocStackInst>(user))
       continue;
@@ -1557,7 +1605,6 @@ void MemoryToRegisters::removeSingleBlockAllocation(AllocStackInst *asi) {
   LLVM_DEBUG(llvm::dbgs() << "*** Promoting in-block: " << *asi);
 
   SILBasicBlock *parentBlock = asi->getParent();
-
   // The default value of the AllocStack is NULL because we don't have
   // uninitialized variables in Swift.
   Optional<StorageStateTracking<LiveValues>> runningVals;
@@ -1580,8 +1627,9 @@ void MemoryToRegisters::removeSingleBlockAllocation(AllocStackInst *asi) {
             /*isStorageValid=*/true};
       }
       assert(runningVals && runningVals->isStorageValid);
-      if (cast<LoadInst>(inst)->getOwnershipQualifier() ==
-          LoadOwnershipQualifier::Take) {
+      auto *loadInst = dyn_cast<LoadInst>(inst);
+      if (loadInst &&
+          loadInst->getOwnershipQualifier() == LoadOwnershipQualifier::Take) {
         if (shouldAddLexicalLifetime(asi)) {
           // End the lexical lifetime at a load [take].  The storage is no
           // longer keeping the value alive.
@@ -1590,8 +1638,8 @@ void MemoryToRegisters::removeSingleBlockAllocation(AllocStackInst *asi) {
         }
         runningVals->isStorageValid = false;
       }
-      replaceLoad(cast<LoadInst>(inst), runningVals->value.replacement(asi),
-                  asi, ctx, deleter, instructionsToDelete);
+      replaceLoad(inst, runningVals->value.replacement(asi), asi, ctx, deleter,
+                  instructionsToDelete);
       ++NumInstRemoved;
       continue;
     }
@@ -1656,7 +1704,7 @@ void MemoryToRegisters::removeSingleBlockAllocation(AllocStackInst *asi) {
     // Remove deallocation.
     if (auto *dsi = dyn_cast<DeallocStackInst>(inst)) {
       if (dsi->getOperand() == asi) {
-        deleter.forceDelete(inst);
+        deleter.forceDelete(dsi);
         NumInstRemoved++;
         // No need to continue scanning after deallocation.
         break;
@@ -1741,7 +1789,7 @@ bool MemoryToRegisters::promoteSingleAllocation(AllocStackInst *alloc) {
 
   // Don't handle captured AllocStacks.
   bool inSingleBlock = false;
-  if (isCaptured(alloc, inSingleBlock)) {
+  if (isCaptured(alloc, &inSingleBlock)) {
     ++NumAllocStackCaptured;
     return false;
   }