Skip to content

Commit 7e6588d

Browse files
aschwaighoferaschwaighofer
committed
Update closure lifetime fixup to use partial_apply [stack]
1 parent cac7fda commit 7e6588d

File tree

2 files changed

+185
-20
lines changed

2 files changed

+185
-20
lines changed

‎lib/SILOptimizer/Mandatory/ClosureLifetimeFixup.cpp

Lines changed: 161 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,7 @@
2121
#include "swift/SILOptimizer/PassManager/Transforms.h"
2222
#include "swift/SILOptimizer/Utils/CFG.h"
2323
#include "swift/SILOptimizer/Utils/Local.h"
24+
#include "swift/SILOptimizer/Utils/StackNesting.h"
2425

2526
#include "llvm/Support/CommandLine.h"
2627

@@ -186,9 +187,145 @@ static SILInstruction *lookThroughRebastractionUsers(
186187
SingleNonDebugNonRefCountUser, Memoized));
187188
}
188189

190+
/// Insert a mark_dependence for any non-trivial argument of a partial_apply.
191+
static SILValue insertMarkDependenceForCapturedArguments(PartialApplyInst *PAI,
192+
SILBuilder &B) {
193+
SILValue curr(PAI);
194+
// Mark dependence on all non-trivial arguments.
195+
for (auto &arg : PAI->getArgumentOperands()) {
196+
if (arg.get()->getType().isTrivial(PAI->getModule()))
197+
continue;
198+
curr = B.createMarkDependence(PAI->getLoc(), curr, arg.get());
199+
}
200+
201+
return curr;
202+
}
203+
204+
/// Rewrite a partial_apply convert_escape_to_noescape sequence with a single
205+
/// apply/try_apply user to a partial_apply [stack] terminated with a
206+
/// dealloc_stack placed after the apply.
207+
///
208+
/// %p = partial_apply %f(%a, %b)
209+
/// %ne = convert_escape_to_noescape %p
210+
/// apply %f2(%p)
211+
/// destroy_value %p
212+
///
213+
/// =>
214+
///
215+
/// %p = partial_apply [stack] %f(%a, %b)
216+
/// %md = mark_dependence %p on %a
217+
/// %md2 = mark_dependence %md on %b
218+
/// apply %f2(%md2)
219+
/// dealloc_stack %p
220+
/// destroy_value %a
221+
/// destroy_value %b
222+
///
223+
/// Note: If the rewrite succeeded we have inserted a dealloc_stack. This
224+
/// dealloc_stack still needs to be balanced with other dealloc_stacks i.e the
225+
/// caller needs to use the StackNesting utility to update the dealloc_stack
226+
/// nesting.
227+
static bool tryRewriteToPartialApplyStack(
228+
SILLocation &loc, PartialApplyInst *origPA,
229+
ConvertEscapeToNoEscapeInst *cvt, SILInstruction *singleApplyUser,
230+
SILBasicBlock::iterator &advanceIfDelete,
231+
llvm::DenseMap<SILInstruction *, SILInstruction *> &Memoized) {
232+
233+
auto *convertOrPartialApply = cast<SingleValueInstruction>(origPA);
234+
if (cvt->getOperand() != origPA)
235+
convertOrPartialApply = cast<ConvertFunctionInst>(cvt->getOperand());
236+
237+
// Whenever we delete an instruction advance the iterator and remove the
238+
// instruction from the memoized map.
239+
auto saveDeleteInst = [&](SILInstruction *I) {
240+
if (&*advanceIfDelete == I)
241+
advanceIfDelete++;
242+
Memoized.erase(I);
243+
I->eraseFromParent();
244+
};
245+
246+
// Look for a single non ref count user of the partial_apply.
247+
SmallVector<SILInstruction *, 8> refCountInsts;
248+
SILInstruction *singleNonDebugNonRefCountUser = nullptr;
249+
for (auto *Usr : getNonDebugUses(convertOrPartialApply)) {
250+
auto *I = Usr->getUser();
251+
if (onlyAffectsRefCount(I)) {
252+
refCountInsts.push_back(I);
253+
continue;
254+
}
255+
if (singleNonDebugNonRefCountUser)
256+
return false;
257+
singleNonDebugNonRefCountUser = I;
258+
}
259+
260+
SILBuilderWithScope B(cvt);
261+
262+
// The convert_escape_to_noescape is the only user of the partial_apply.
263+
// Convert to a partial_apply [stack].
264+
SmallVector<SILValue, 8> args;
265+
for (auto &arg : origPA->getArgumentOperands())
266+
args.push_back(arg.get());
267+
auto newPA = B.createPartialApply(
268+
origPA->getLoc(), origPA->getCallee(), origPA->getSubstitutionMap(), args,
269+
origPA->getType().getAs<SILFunctionType>()->getCalleeConvention(),
270+
PartialApplyInst::OnStackKind::OnStack);
271+
272+
// Insert mark_dependence for any non-trivial operand to the partial_apply.
273+
auto closure = insertMarkDependenceForCapturedArguments(newPA, B);
274+
275+
// Optionally, replace the convert_function instruction.
276+
if (auto *convert = dyn_cast<ConvertFunctionInst>(convertOrPartialApply)) {
277+
auto origTy = convert->getType().castTo<SILFunctionType>();
278+
auto origWithNoEscape = SILType::getPrimitiveObjectType(
279+
origTy->getWithExtInfo(origTy->getExtInfo().withNoEscape()));
280+
closure = B.createConvertFunction(convert->getLoc(), closure, origWithNoEscape, false);
281+
convert->replaceAllUsesWith(closure);
282+
}
283+
284+
// Replace the convert_escape_to_noescape uses with the new
285+
// partial_apply [stack].
286+
cvt->replaceAllUsesWith(closure);
287+
saveDeleteInst(cvt);
288+
289+
// Delete the ref count operations on the original partial_apply.
290+
for (auto *refInst : refCountInsts)
291+
saveDeleteInst(refInst);
292+
convertOrPartialApply->replaceAllUsesWith(newPA);
293+
if (convertOrPartialApply != origPA)
294+
saveDeleteInst(convertOrPartialApply);
295+
saveDeleteInst(origPA);
296+
297+
// Insert destroys of arguments after the apply and the dealloc_stack.
298+
if (auto *Apply = dyn_cast<ApplyInst>(singleApplyUser)) {
299+
auto InsertPt = std::next(SILBasicBlock::iterator(Apply));
300+
SILBuilderWithScope B3(InsertPt);
301+
B3.createDeallocStack(loc, newPA);
302+
insertDestroyOfCapturedArguments(newPA, B3);
303+
} else if (auto *Try = dyn_cast<TryApplyInst>(singleApplyUser)) {
304+
for (auto *SuccBB : Try->getSuccessorBlocks()) {
305+
SILBuilderWithScope B3(SuccBB->begin());
306+
B3.createDeallocStack(loc, newPA);
307+
insertDestroyOfCapturedArguments(newPA, B3);
308+
}
309+
} else {
310+
llvm_unreachable("Unknown FullApplySite instruction kind");
311+
}
312+
return true;
313+
}
314+
315+
static SILValue skipConvert(SILValue v) {
316+
auto cvt = dyn_cast<ConvertFunctionInst>(v);
317+
if (!cvt)
318+
return v;
319+
auto *pa = dyn_cast<PartialApplyInst>(cvt->getOperand());
320+
if (!pa || !pa->hasOneUse())
321+
return v;
322+
return pa;
323+
}
324+
189325
static bool tryExtendLifetimeToLastUse(
190326
ConvertEscapeToNoEscapeInst *Cvt,
191-
llvm::DenseMap<SILInstruction *, SILInstruction *> &Memoized) {
327+
llvm::DenseMap<SILInstruction *, SILInstruction *> &Memoized,
328+
SILBasicBlock::iterator &AdvanceIfDelete) {
192329
// If there is a single user that is an apply this is simple: extend the
193330
// lifetime of the operand until after the apply.
194331
auto SingleUser = lookThroughRebastractionUsers(Cvt, Memoized);
@@ -203,6 +340,12 @@ static bool tryExtendLifetimeToLastUse(
203340
}
204341

205342
auto loc = RegularLocation::getAutoGeneratedLocation();
343+
auto origPA = dyn_cast<PartialApplyInst>(skipConvert(Cvt->getOperand()));
344+
if (origPA && tryRewriteToPartialApplyStack(
345+
loc, origPA, Cvt, SingleApplyUser.getInstruction(),
346+
AdvanceIfDelete,
347+
Memoized))
348+
return true;
206349

207350
// Insert a copy at the convert_escape_to_noescape [not_guaranteed] and
208351
// change the instruction to the guaranteed form.
@@ -473,7 +616,7 @@ static bool fixupCopyBlockWithoutEscaping(CopyBlockWithoutEscapingInst *CB) {
473616
return true;
474617
}
475618

476-
static bool fixupClosureLifetimes(SILFunction &Fn) {
619+
static bool fixupClosureLifetimes(SILFunction &Fn, bool &checkStackNesting) {
477620
bool Changed = false;
478621

479622
// tryExtendLifetimeToLastUse uses a cache of recursive instruction use
@@ -505,8 +648,9 @@ static bool fixupClosureLifetimes(SILFunction &Fn) {
505648
continue;
506649
}
507650

508-
if (tryExtendLifetimeToLastUse(Cvt, MemoizedQueries)) {
651+
if (tryExtendLifetimeToLastUse(Cvt, MemoizedQueries, I)) {
509652
Changed |= true;
653+
checkStackNesting = true;
510654
continue;
511655
}
512656

@@ -537,8 +681,20 @@ class ClosureLifetimeFixup : public SILFunctionTransform {
537681

538682
// Fixup convert_escape_to_noescape [not_guaranteed] and
539683
// copy_block_without_escaping instructions.
540-
if (fixupClosureLifetimes(*getFunction()))
541-
invalidateAnalysis(SILAnalysis::InvalidationKind::FunctionBody);
684+
685+
bool checkStackNesting = false;
686+
bool modifiedCFG = false;
687+
if (fixupClosureLifetimes(*getFunction(), checkStackNesting)) {
688+
if (checkStackNesting){
689+
StackNesting SN;
690+
modifiedCFG =
691+
SN.correctStackNesting(getFunction()) == StackNesting::Changes::CFG;
692+
}
693+
if (modifiedCFG)
694+
invalidateAnalysis(SILAnalysis::InvalidationKind::FunctionBody);
695+
else
696+
invalidateAnalysis(SILAnalysis::InvalidationKind::CallsAndInstructions);
697+
}
542698
LLVM_DEBUG(getFunction()->verify());
543699

544700
}

‎test/SILOptimizer/closure_lifetime_fixup.swift

Lines changed: 24 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -17,14 +17,13 @@ public class C {
1717
// CHECK: bb0([[ARG:%.*]] : $C):
1818
// CHECK: [[F:%.*]] = function_ref @$s22closure_lifetime_fixup10testSimple1cyAA1CC_tFyyXEfU_
1919
// CHECK-NEXT: strong_retain [[ARG]] : $C
20-
// CHECK-NEXT: [[PA:%.*]] = partial_apply [callee_guaranteed] [[F]]([[ARG]]) : $@convention(thin) (@guaranteed C) -> ()
21-
// CHECK-NEXT: strong_retain [[PA]] : $@callee_guaranteed () -> ()
22-
// CHECK-NEXT: [[CVT:%.*]] = convert_escape_to_noescape [[PA]] : $@callee_guaranteed () -> () to $@noescape @callee_guaranteed () -> ()
20+
// CHECK-NEXT: [[PA:%.*]] = partial_apply [callee_guaranteed] [on_stack] [[F]]([[ARG]]) : $@convention(thin) (@guaranteed C) -> ()
21+
// CHECK-NEXT: [[CL:%.*]] = mark_dependence [[PA]] : $@noescape @callee_guaranteed () -> () on [[ARG]] : $C
2322
// CHECK-NEXT: // function_ref use_closure(_:)
2423
// CHECK-NEXT: [[F2:%.*]] = function_ref @$s22closure_lifetime_fixup04use_A0yyyyXEF : $@convention(thin) (@noescape @callee_guaranteed () -> ()) -> ()
25-
// CHECK-NEXT: apply [[F2]]([[CVT]]) : $@convention(thin) (@noescape @callee_guaranteed () -> ()) -> ()
26-
// CHECK-NEXT: strong_release [[PA]] : $@callee_guaranteed () -> ()
27-
// CHECK-NEXT: strong_release [[PA]] : $@callee_guaranteed () -> ()
24+
// CHECK-NEXT: apply [[F2]]([[CL]]) : $@convention(thin) (@noescape @callee_guaranteed () -> ()) -> ()
25+
// CHECK-NEXT: dealloc_stack [[PA]] : $@noescape @callee_guaranteed () -> ()
26+
// CHECK-NEXT: strong_release [[ARG]] : $C
2827
// CHECK-NEXT: tuple ()
2928
// CHECK-NEXT: return {{.*}} : $()
3029
public func testSimple(c: C) {
@@ -35,19 +34,17 @@ public func testSimple(c: C) {
3534
// CHECK:bb0([[ARG:%.*]] : $C):
3635
// CHECK: [[F:%.*]] = function_ref @$s22closure_lifetime_fixup11testGeneric1cyAA1CC_tFSiyXEfU_ : $@convention(thin) (@guaranteed C) -> Int
3736
// CHECK-NEXT: strong_retain [[ARG]] : $C
38-
// CHECK-NEXT: [[PA:%.*]] = partial_apply [callee_guaranteed] [[F]]([[ARG]]) : $@convention(thin) (@guaranteed C) -> Int
39-
// CHECK-NEXT: strong_retain [[PA]] : $@callee_guaranteed () -> Int
40-
// CHECK-NEXT: [[CVT:%.*]] = convert_escape_to_noescape [[PA]] : $@callee_guaranteed () -> Int to $@noescape @callee_guaranteed () -> Int
37+
// CHECK-NEXT: [[PA:%.*]] = partial_apply [callee_guaranteed] [on_stack] [[F]]([[ARG]]) : $@convention(thin) (@guaranteed C) -> Int
38+
// CHECK-NEXT: [[MD:%.*]] = mark_dependence [[PA]] : $@noescape @callee_guaranteed () -> Int on [[ARG]] : $C
4139
// CHECK-NEXT: // function_ref thunk for @callee_guaranteed () -> (@unowned Int)
4240
// CHECK-NEXT: [[F:%.*]] = function_ref @$sSiIgd_SiIegr_TR : $@convention(thin) (@noescape @callee_guaranteed () -> Int) -> @out Int
43-
// CHECK-NEXT: [[PA2:%.*]] = partial_apply [callee_guaranteed] [[F]]([[CVT]]) : $@convention(thin) (@noescape @callee_guaranteed () -> Int) -> @out Int
44-
// CHECK-NEXT: [[CVT2:%.*]] = convert_escape_to_noescape [[PA2]] : $@callee_guaranteed () -> @out Int to $@noescape @callee_guaranteed () -> @out Int
45-
// CHECK-NEXT: strong_release [[PA]] : $@callee_guaranteed () -> Int
41+
// CHECK-NEXT: [[PA2:%.*]] = partial_apply [callee_guaranteed] [on_stack] [[F]]([[MD]]) : $@convention(thin) (@noescape @callee_guaranteed () -> Int) -> @out Int
4642
// CHECK-NEXT: // function_ref use_closureGeneric<A>(_:)
4743
// CHECK-NEXT: [[F2:%.*]] = function_ref @$s22closure_lifetime_fixup04use_A7GenericyyxyXElF : $@convention(thin) <τ_0_0> (@noescape @callee_guaranteed () -> @out τ_0_0) -> ()
48-
// CHECK-NEXT: %12 = apply [[F2]]<Int>([[CVT2]]) : $@convention(thin) <τ_0_0> (@noescape @callee_guaranteed () -> @out τ_0_0) -> ()
49-
// CHECK-NEXT: strong_release [[PA2]] : $@callee_guaranteed () -> @out Int
50-
// CHECK-NEXT: strong_release [[PA]] : $@callee_guaranteed () -> Int
44+
// CHECK-NEXT: apply [[F2]]<Int>([[PA2]]) : $@convention(thin) <τ_0_0> (@noescape @callee_guaranteed () -> @out τ_0_0) -> ()
45+
// CHECK-NEXT: dealloc_stack [[PA2]]
46+
// CHECK-NEXT: dealloc_stack [[PA]]
47+
// CHECK-NEXT: strong_release [[ARG]]
5148
// CHECK-NEXT: tuple ()
5249
// CHECK-NEXT: return {{.*}} : $()
5350

@@ -90,3 +87,15 @@ public func dontCrash<In, Out>(test: Bool, body: @escaping ((In) -> Out, In) ->
9087
}
9188
return result
9289
}
90+
91+
// CHECK-LABEL: sil @$s22closure_lifetime_fixup28to_stack_of_convert_function1pySvSg_tF
92+
// CHECK: [[PA:%.*]] = partial_apply [callee_guaranteed] [on_stack]
93+
// CHECK: [[MD:%.*]] = mark_dependence [[PA]]
94+
// CHECK: [[CVT:%.*]] = convert_function [[MD]]
95+
// CHECK: [[REABSTRACT:%.*]] = function_ref @$sSvSSs5Error_pIgyozo_SvSSsAA_pIegnrzo_TR
96+
// CHECK: [[PA2:%.*]] = partial_apply [callee_guaranteed] [on_stack] [[REABSTRACT]]([[CVT]])
97+
// CHECK: [[MAP:%.*]] = function_ref @$sSq3mapyqd__Sgqd__xKXEKlF
98+
// CHECK: try_apply [[MAP]]<UnsafeMutableRawPointer, String>({{.*}}, [[PA2]], {{.*}})
99+
public func to_stack_of_convert_function(p: UnsafeMutableRawPointer?) {
100+
_ = p.map(String.init(describing:))
101+
}

0 commit comments

Comments
 (0)