Enable Mem2Reg for allocs with load_borrow (#59350)

This commit is contained in:
Meghana Gupta
2022-06-15 14:37:10 -07:00
committed by GitHub
parent 4859f89198
commit a0a0ebb558
2 changed files with 459 additions and 38 deletions

View File

@@ -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,12 +239,13 @@ 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) {
@@ -247,9 +253,9 @@ replaceLoad(LoadInst *li, SILValue newValue, AllocStackInst *asi,
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;
}