Skip to content

[temp-rvalue] Handle promoting alloc_stack that are destroyed via a load [take] in ossa. #30088

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
90 changes: 82 additions & 8 deletions lib/SILOptimizer/Transforms/TempRValueElimination.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -70,7 +70,9 @@ class TempRValueOptPass : public SILFunctionTransform {
checkNoSourceModification(CopyAddrInst *copyInst,
const SmallPtrSetImpl<SILInstruction *> &useInsts);

bool checkTempObjectDestroy(AllocStackInst *tempObj, CopyAddrInst *copyInst);
bool
checkTempObjectDestroy(AllocStackInst *tempObj, CopyAddrInst *copyInst,
ValueLifetimeAnalysis::Frontier &tempAddressFrontier);

bool tryOptimizeCopyIntoTemp(CopyAddrInst *copyInst);

Expand Down Expand Up @@ -189,6 +191,18 @@ bool TempRValueOptPass::collectLoads(
case SILInstructionKind::LoadBorrowInst:
// Loads are the end of the data flow chain. The users of the load can't
// access the temporary storage.
//
// That being said, if we see a load [take] here then we must have had a
// load [take] of a projection of our temporary stack location since we skip
// all the load [take] of the top level allocation in the caller of this
// function. So if we have such a load [take], we /must/ have a
// reinitialization or an alloc_stack that does not fit the pattern we are
// expecting from SILGen. Be conservative and return false.
if (auto *li = dyn_cast<LoadInst>(user)) {
if (li->getOwnershipQualifier() == LoadOwnershipQualifier::Take) {
return false;
}
}
loadInsts.insert(user);
return true;

Expand Down Expand Up @@ -250,8 +264,9 @@ bool TempRValueOptPass::checkNoSourceModification(
/// it is legal to destroy an in-memory object by loading the value and
/// releasing it. Rather than detecting unbalanced load releases, simply check
/// that tempObj is destroyed directly on all paths.
bool TempRValueOptPass::checkTempObjectDestroy(AllocStackInst *tempObj,
CopyAddrInst *copyInst) {
bool TempRValueOptPass::checkTempObjectDestroy(
AllocStackInst *tempObj, CopyAddrInst *copyInst,
ValueLifetimeAnalysis::Frontier &tempAddressFrontier) {
// If the original copy was a take, then replacing all uses cannot affect
// the lifetime.
if (copyInst->isTakeOfSrc())
Expand All @@ -275,7 +290,6 @@ bool TempRValueOptPass::checkTempObjectDestroy(AllocStackInst *tempObj,
}
// Find the boundary of tempObj's address lifetime, starting at copyInst.
ValueLifetimeAnalysis vla(copyInst, users);
ValueLifetimeAnalysis::Frontier tempAddressFrontier;
if (!vla.computeFrontier(tempAddressFrontier,
ValueLifetimeAnalysis::DontModifyCFG)) {
return false;
Expand All @@ -289,13 +303,19 @@ bool TempRValueOptPass::checkTempObjectDestroy(AllocStackInst *tempObj,
if (pos == frontierInst->getParent()->begin())
return false;

// Look for a known destroy point as described in the funciton level
// Look for a known destroy point as described in the function level
// comment. This whitelist can be expanded as more cases are handled in
// tryOptimizeCopyIntoTemp during copy replacement.
SILInstruction *lastUser = &*std::prev(pos);
if (isa<DestroyAddrInst>(lastUser))
continue;

if (auto *li = dyn_cast<LoadInst>(lastUser)) {
if (li->getOwnershipQualifier() == LoadOwnershipQualifier::Take) {
continue;
}
}

if (auto *cai = dyn_cast<CopyAddrInst>(lastUser)) {
assert(cai->getSrc() == tempObj && "collectLoads checks for writes");
assert(!copyInst->isTakeOfSrc() && "checked above");
Expand Down Expand Up @@ -334,6 +354,15 @@ bool TempRValueOptPass::tryOptimizeCopyIntoTemp(CopyAddrInst *copyInst) {
if (isa<DestroyAddrInst>(user) || isa<DeallocStackInst>(user))
continue;

// Same for load [take] on the top level temp object. SILGen always takes
// whole values from temporaries. If we have load [take] on projections from
// our base, we fail since those would be re-initializations.
if (auto *li = dyn_cast<LoadInst>(user)) {
if (li->getOwnershipQualifier() == LoadOwnershipQualifier::Take) {
continue;
}
}

if (!collectLoads(useOper, user, tempObj, copyInst->getSrc(), loadInsts))
return false;
}
Expand All @@ -342,7 +371,8 @@ bool TempRValueOptPass::tryOptimizeCopyIntoTemp(CopyAddrInst *copyInst) {
if (!checkNoSourceModification(copyInst, loadInsts))
return false;

if (!checkTempObjectDestroy(tempObj, copyInst))
ValueLifetimeAnalysis::Frontier tempAddressFrontier;
if (!checkTempObjectDestroy(tempObj, copyInst, tempAddressFrontier))
return false;

LLVM_DEBUG(llvm::dbgs() << " Success: replace temp" << *tempObj);
Expand All @@ -351,6 +381,10 @@ bool TempRValueOptPass::tryOptimizeCopyIntoTemp(CopyAddrInst *copyInst) {
// the source address. Note: we must not delete the original copyInst because
// it would crash the instruction iteration in run(). Instead the copyInst
// gets identical Src and Dest operands.
//
// NOTE: We delete instructions at the end to allow us to use
// tempAddressFrontier to insert compensating destroys for load [take].
SmallVector<SILInstruction *, 4> toDelete;
while (!tempObj->use_empty()) {
Operand *use = *tempObj->use_begin();
SILInstruction *user = use->getUser();
Expand All @@ -359,11 +393,13 @@ bool TempRValueOptPass::tryOptimizeCopyIntoTemp(CopyAddrInst *copyInst) {
if (copyInst->isTakeOfSrc()) {
use->set(copyInst->getSrc());
} else {
user->eraseFromParent();
user->dropAllReferences();
toDelete.push_back(user);
}
break;
case SILInstructionKind::DeallocStackInst:
user->eraseFromParent();
user->dropAllReferences();
toDelete.push_back(user);
break;
case SILInstructionKind::CopyAddrInst: {
auto *cai = cast<CopyAddrInst>(user);
Expand All @@ -375,6 +411,40 @@ bool TempRValueOptPass::tryOptimizeCopyIntoTemp(CopyAddrInst *copyInst) {
use->set(copyInst->getSrc());
break;
}
case SILInstructionKind::LoadInst: {
// If we do not have a load [take] or we have a load [take] and our
// copy_addr takes the source, just do the normal thing of setting the
// load to use the copyInst's source.
auto *li = cast<LoadInst>(user);
if (li->getOwnershipQualifier() != LoadOwnershipQualifier::Take ||
copyInst->isTakeOfSrc()) {
use->set(copyInst->getSrc());
break;
}

// Otherwise, since copy_addr is not taking src, we need to ensure that we
// insert a copy of our value. We do that by creating a load [copy] at the
// copy_addr inst and RAUWing the load [take] with that. We then insert
// destroy_value for the load [copy] at all points where we had destroys
// that are not the specific take that we were optimizing.
SILBuilderWithScope builder(copyInst);
SILValue newLoad = builder.emitLoadValueOperation(
copyInst->getLoc(), copyInst->getSrc(), LoadOwnershipQualifier::Copy);
for (auto *inst : tempAddressFrontier) {
assert(inst->getIterator() != inst->getParent()->begin() &&
"Should have caught this when checking destructor");
auto prevInst = std::prev(inst->getIterator());
if (&*prevInst == li)
continue;
SILBuilderWithScope builder(prevInst);
builder.emitDestroyValueOperation(prevInst->getLoc(), newLoad);
}
li->replaceAllUsesWith(newLoad);
li->dropAllReferences();
toDelete.push_back(li);
break;
}

// ASSUMPTION: no operations that may be handled by this default clause can
// destroy tempObj. This includes operations that load the value from memory
// and release it.
Expand All @@ -383,6 +453,10 @@ bool TempRValueOptPass::tryOptimizeCopyIntoTemp(CopyAddrInst *copyInst) {
break;
}
}

while (!toDelete.empty()) {
toDelete.pop_back_val()->eraseFromParent();
}
tempObj->eraseFromParent();
return true;
}
Expand Down
103 changes: 90 additions & 13 deletions test/SILOptimizer/temp_rvalue_opt_ossa.sil
Original file line number Diff line number Diff line change
Expand Up @@ -18,6 +18,8 @@ class Klass {}

sil @unknown : $@convention(thin) () -> ()

sil @use_gsbase_builtinnativeobject : $@convention(thin) (@guaranteed GS<Builtin.NativeObject>) -> ()

sil [ossa] @inguaranteed_user_without_result : $@convention(thin) (@in_guaranteed Klass) -> () {
bb0(%0 : $*Klass):
%9999 = tuple()
Expand Down Expand Up @@ -604,25 +606,17 @@ bb3:
br bb2
}

///////////////////////////////////////////////////////////////////////////////
// Test checkTempObjectDestroy
// <rdar://problem/56393491> Use-after free crashing an XCTest.

sil [ossa] @takeGuaranteedObj : $@convention(thin) (@guaranteed Builtin.NativeObject) -> ()

// Do not remove a copy that is released via a load (because
// TempRValueOpt is an address-based optimization does not know how to
// remove releases, and trying to do that would reduce to ARC
// Now that we support ossa, eliminate the alloc_stack and change the load
// [take] to a load [copy] in the process.
//
// optimization).
// CHECK-LABEL: sil [ossa] @copyWithLoadRelease : $@convention(thin) (@in_guaranteed Builtin.NativeObject) -> () {
// CHECK: bb0(%0 : $*Builtin.NativeObject):
// CHECK: [[STK:%.*]] = alloc_stack $Builtin.NativeObject
// CHECK: copy_addr %0 to [initialization] [[STK]] : $*Builtin.NativeObject
// CHECK: [[VAL:%.*]] = load [take] [[STK]] : $*Builtin.NativeObject
// CHECK-NOT: alloc_stack
// CHECK: [[VAL:%.*]] = load [copy] %0 : $*Builtin.NativeObject
// CHECK: apply %{{.*}}([[VAL]]) : $@convention(thin) (@guaranteed Builtin.NativeObject) -> ()
// CHECK: destroy_value [[VAL]] : $Builtin.NativeObject
// CHECK: dealloc_stack [[STK]] : $*Builtin.NativeObject
// CHECK-LABEL: } // end sil function 'copyWithLoadRelease'
sil [ossa] @copyWithLoadRelease : $@convention(thin) (@in_guaranteed Builtin.NativeObject) -> () {
bb0(%0 : $*Builtin.NativeObject):
Expand All @@ -637,7 +631,8 @@ bb0(%0 : $*Builtin.NativeObject):
return %v : $()
}

// Remove a copy that is released via a load as long as it was a copy [take].
// Remove a copy that is released via a load. Leave the load [take] alone since
// our copy_addr is taking from source.
//
// CHECK-LABEL: sil [ossa] @takeWithLoadRelease : $@convention(thin) (@in Builtin.NativeObject) -> () {
// CHECK: bb0(%0 : $*Builtin.NativeObject):
Expand All @@ -657,3 +652,85 @@ bb0(%0 : $*Builtin.NativeObject):
%v = tuple ()
return %v : $()
}

// Do not remove a copy that is released via a load of a projection. This is not
// the pattern from SILGen that we are targeting, so we reduce the state space by banning the pattern.
//
// CHECK-LABEL: sil [ossa] @takeWithLoadReleaseOfProjection : $@convention(thin) (@in GS<Builtin.NativeObject>) -> () {
// CHECK: alloc_stack
// CHECK: } // end sil function 'takeWithLoadReleaseOfProjection'
sil [ossa] @takeWithLoadReleaseOfProjection : $@convention(thin) (@in GS<Builtin.NativeObject>) -> () {
bb0(%0 : $*GS<Builtin.NativeObject>):
%stk = alloc_stack $GS<Builtin.NativeObject>
copy_addr [take] %0 to [initialization] %stk : $*GS<Builtin.NativeObject>
%proj = struct_element_addr %stk : $*GS<Builtin.NativeObject>, #GS._base
%obj = load [take] %proj : $*Builtin.NativeObject
%f = function_ref @takeGuaranteedObj : $@convention(thin) (@guaranteed Builtin.NativeObject) -> ()
%call = apply %f(%obj) : $@convention(thin) (@guaranteed Builtin.NativeObject) -> ()
destroy_value %obj : $Builtin.NativeObject
dealloc_stack %stk : $*GS<Builtin.NativeObject>
%v = tuple ()
return %v : $()
}

// Make sure that when we convert the load [take] to a load [copy], we hoist the
// load of src /before/ the destroy of %0.
// CHECK-LABEL: sil [ossa] @hoist_load_copy_to_src_copy_addr_site : $@convention(thin) (@in GS<Builtin.NativeObject>) -> @owned GS<Builtin.NativeObject> {
// CHECK: bb0([[ARG:%.*]] :
// CHECK: [[VAL:%.*]] = load [copy] [[ARG]]
// CHECK: destroy_addr [[ARG]]
// CHECK: apply {{%.*}}([[VAL]])
// CHECK: return [[VAL]]
// CHECK: } // end sil function 'hoist_load_copy_to_src_copy_addr_site'
sil [ossa] @hoist_load_copy_to_src_copy_addr_site : $@convention(thin) (@in GS<Builtin.NativeObject>) -> @owned GS<Builtin.NativeObject> {
bb0(%0 : $*GS<Builtin.NativeObject>):
%f = function_ref @use_gsbase_builtinnativeobject : $@convention(thin) (@guaranteed GS<Builtin.NativeObject>) -> ()
%stk = alloc_stack $GS<Builtin.NativeObject>
copy_addr %0 to [initialization] %stk : $*GS<Builtin.NativeObject>
destroy_addr %0 : $*GS<Builtin.NativeObject>
%obj = load [take] %stk : $*GS<Builtin.NativeObject>
dealloc_stack %stk : $*GS<Builtin.NativeObject>
apply %f(%obj) : $@convention(thin) (@guaranteed GS<Builtin.NativeObject>) -> ()
return %obj : $GS<Builtin.NativeObject>
}

// CHECK-LABEL: sil [ossa] @hoist_load_copy_to_src_copy_addr_site_two_takes : $@convention(thin) (@in GS<Builtin.NativeObject>) -> @owned GS<Builtin.NativeObject> {
// CHECK: bb0([[ARG:%.*]] :
// CHECK: [[COPY_1:%.*]] = load [copy] [[ARG]]
// CHECK: [[COPY_2:%.*]] = load [copy] [[ARG]]
// CHECK: destroy_addr [[ARG]]
// CHECK: cond_br undef, bb1, bb2
//
// CHECK: bb1:
// CHECK: destroy_value [[COPY_1]]
// CHECK: br bb3([[COPY_2]] :
//
// CHECK: bb2:
// CHECK: destroy_value [[COPY_2]]
// CHECK: br bb3([[COPY_1]] :
//
// CHECK: bb3([[RESULT:%.*]] : @owned
// CHECK: apply {{%.*}}([[RESULT]])
// CHECK: return [[RESULT]]
// CHECK: } // end sil function 'hoist_load_copy_to_src_copy_addr_site_two_takes'
sil [ossa] @hoist_load_copy_to_src_copy_addr_site_two_takes : $@convention(thin) (@in GS<Builtin.NativeObject>) -> @owned GS<Builtin.NativeObject> {
bb0(%0 : $*GS<Builtin.NativeObject>):
%f = function_ref @use_gsbase_builtinnativeobject : $@convention(thin) (@guaranteed GS<Builtin.NativeObject>) -> ()
%stk = alloc_stack $GS<Builtin.NativeObject>
copy_addr %0 to [initialization] %stk : $*GS<Builtin.NativeObject>
destroy_addr %0 : $*GS<Builtin.NativeObject>
cond_br undef, bb1, bb2

bb1:
%obj = load [take] %stk : $*GS<Builtin.NativeObject>
br bb3(%obj : $GS<Builtin.NativeObject>)

bb2:
%obj2 = load [take] %stk : $*GS<Builtin.NativeObject>
br bb3(%obj2 : $GS<Builtin.NativeObject>)

bb3(%obj3 : @owned $GS<Builtin.NativeObject>):
dealloc_stack %stk : $*GS<Builtin.NativeObject>
apply %f(%obj3) : $@convention(thin) (@guaranteed GS<Builtin.NativeObject>) -> ()
return %obj3 : $GS<Builtin.NativeObject>
}