Skip to content

Commit cb24a41

Browse files
aschwaighoferaschwaighofer
authored
Merge pull request #12764 from aschwaighofer/fix_mandatory_inliner_guaranteed
Fix of mandatory inliner's handling of unowned and guaranteed captures
2 parents 3e45eb7 + 2ec064a commit cb24a41

File tree

2 files changed

+154
-24
lines changed

2 files changed

+154
-24
lines changed

‎lib/SILOptimizer/Mandatory/MandatoryInlining.cpp

Lines changed: 54 additions & 24 deletions
Original file line numberDiff line numberDiff line change
@@ -83,19 +83,30 @@ class DeleteInstructionsHandler : public DeleteNotificationHandler {
8383
/// applications: namely, that an apply of a thick function consumes the callee
8484
/// and that the function implementing the closure consumes its capture
8585
/// arguments.
86-
static void fixupReferenceCounts(SILBasicBlock::iterator I,
87-
SILValue CalleeValue,
88-
SmallVectorImpl<SILValue> &CaptureArgs) {
86+
static void fixupReferenceCounts(
87+
SILBasicBlock::iterator I, SILValue CalleeValue,
88+
SmallVectorImpl<std::pair<SILValue, ParameterConvention>> &CaptureArgs,
89+
bool isCalleeGuaranteed) {
8990
// Add a copy of each non-address type capture argument to lifetime extend the
9091
// captured argument over the inlined function. This deals with the
9192
// possibility of the closure being destroyed by an earlier application and
9293
// thus cause the captured argument to be destroyed.
93-
for (auto &CaptureArg : CaptureArgs)
94-
if (!CaptureArg->getType().isAddress())
95-
createIncrementBefore(CaptureArg, &*I);
94+
for (auto &CaptureArg : CaptureArgs) {
95+
if (!CaptureArg.first->getType().isAddress() &&
96+
CaptureArg.second != ParameterConvention::Direct_Guaranteed &&
97+
CaptureArg.second != ParameterConvention::Direct_Unowned) {
98+
createIncrementBefore(CaptureArg.first, &*I);
99+
} else {
100+
// FIXME: What about indirectly owned parameters? The invocation of the
101+
// closure would perform an indirect copy which we should mimick here.
102+
assert(CaptureArg.second != ParameterConvention::Indirect_In &&
103+
"Missing indirect copy");
104+
}
105+
}
96106

97107
// Destroy the callee as the apply would have done.
98-
createDecrementBefore(CalleeValue, &*I);
108+
if (!isCalleeGuaranteed)
109+
createDecrementBefore(CalleeValue, &*I);
99110
}
100111

101112
static SILValue cleanupLoadedCalleeValue(SILValue CalleeValue, LoadInst *LI) {
@@ -159,9 +170,9 @@ static SILValue cleanupLoadedCalleeValue(SILValue CalleeValue, LoadInst *LI) {
159170

160171
/// \brief Removes instructions that create the callee value if they are no
161172
/// longer necessary after inlining.
162-
static void
163-
cleanupCalleeValue(SILValue CalleeValue, ArrayRef<SILValue> CaptureArgs,
164-
ArrayRef<SILValue> FullArgs) {
173+
static void cleanupCalleeValue(
174+
SILValue CalleeValue,
175+
ArrayRef<SILValue> FullArgs) {
165176
SmallVector<SILInstruction*, 16> InstsToDelete;
166177
for (SILValue V : FullArgs) {
167178
if (V != CalleeValue)
@@ -217,6 +228,22 @@ cleanupCalleeValue(SILValue CalleeValue, ArrayRef<SILValue> CaptureArgs,
217228
}
218229
}
219230

231+
static void collectPartiallyAppliedArguments(
232+
PartialApplyInst *PAI,
233+
SmallVectorImpl<std::pair<SILValue, ParameterConvention>> &CapturedArgs,
234+
SmallVectorImpl<SILValue> &FullArgs) {
235+
ApplySite Site(PAI);
236+
SILFunctionConventions CalleeConv(Site.getSubstCalleeType(),
237+
PAI->getModule());
238+
for (auto &Arg : PAI->getArgumentOperands()) {
239+
unsigned CalleeArgumentIndex = Site.getCalleeArgIndex(Arg);
240+
assert(CalleeArgumentIndex >= CalleeConv.getSILArgIndexOfFirstParam());
241+
auto ParamInfo = CalleeConv.getParamInfoForSILArg(CalleeArgumentIndex);
242+
CapturedArgs.push_back(std::make_pair(Arg.get(), ParamInfo.getConvention()));
243+
FullArgs.push_back(Arg.get());
244+
}
245+
}
246+
220247
/// \brief Returns the callee SILFunction called at a call site, in the case
221248
/// that the call is transparent (as in, both that the call is marked
222249
/// with the transparent flag and that callee function is actually transparently
@@ -225,12 +252,11 @@ cleanupCalleeValue(SILValue CalleeValue, ArrayRef<SILValue> CaptureArgs,
225252
///
226253
/// In the case that a non-null value is returned, FullArgs contains effective
227254
/// argument operands for the callee function.
228-
static SILFunction *getCalleeFunction(SILFunction *F, FullApplySite AI,
229-
bool &IsThick,
230-
SmallVectorImpl<SILValue> &CaptureArgs,
231-
SmallVectorImpl<SILValue> &FullArgs,
232-
PartialApplyInst *&PartialApply,
233-
SILModule::LinkingMode Mode) {
255+
static SILFunction *getCalleeFunction(
256+
SILFunction *F, FullApplySite AI, bool &IsThick,
257+
SmallVectorImpl<std::pair<SILValue, ParameterConvention>> &CaptureArgs,
258+
SmallVectorImpl<SILValue> &FullArgs, PartialApplyInst *&PartialApply,
259+
SILModule::LinkingMode Mode) {
234260
IsThick = false;
235261
PartialApply = nullptr;
236262
CaptureArgs.clear();
@@ -319,10 +345,9 @@ static SILFunction *getCalleeFunction(SILFunction *F, FullApplySite AI,
319345
// one "thin to thick function" instructions, since those are the patterns
320346
// generated when using auto closures.
321347
if (auto *PAI = dyn_cast<PartialApplyInst>(CalleeValue)) {
322-
for (const auto &Arg : PAI->getArguments()) {
323-
CaptureArgs.push_back(Arg);
324-
FullArgs.push_back(Arg);
325-
}
348+
349+
// Collect the applied arguments and their convention.
350+
collectPartiallyAppliedArguments(PAI, CaptureArgs, FullArgs);
326351

327352
CalleeValue = PAI->getCallee();
328353
IsThick = true;
@@ -440,7 +465,7 @@ runOnFunctionRecursively(SILFunction *F, FullApplySite AI,
440465
// during this call and recursive subcalls).
441466
CurrentInliningSet = SetFactory.add(CurrentInliningSet, F);
442467

443-
SmallVector<SILValue, 16> CaptureArgs;
468+
SmallVector<std::pair<SILValue, ParameterConvention>, 16> CaptureArgs;
444469
SmallVector<SILValue, 32> FullArgs;
445470

446471
for (auto BI = F->begin(), BE = F->end(); BI != BE; ++BI) {
@@ -465,6 +490,7 @@ runOnFunctionRecursively(SILFunction *F, FullApplySite AI,
465490
PartialApplyInst *PAI;
466491
SILFunction *CalleeFunction = getCalleeFunction(
467492
F, InnerAI, IsThick, CaptureArgs, FullArgs, PAI, Mode);
493+
468494
if (!CalleeFunction)
469495
continue;
470496

@@ -529,8 +555,12 @@ runOnFunctionRecursively(SILFunction *F, FullApplySite AI,
529555

530556
// If we intend to inline a thick function, then we need to balance the
531557
// reference counts for correctness.
532-
if (IsThick)
533-
fixupReferenceCounts(II, CalleeValue, CaptureArgs);
558+
if (IsThick) {
559+
bool IsCalleeGuaranteed =
560+
PAI &&
561+
PAI->getType().castTo<SILFunctionType>()->isCalleeGuaranteed();
562+
fixupReferenceCounts(II, CalleeValue, CaptureArgs, IsCalleeGuaranteed);
563+
}
534564

535565
// Decrement our iterator (carefully, to avoid going off the front) so it
536566
// is valid after inlining is done. Inlining deletes the apply, and can
@@ -551,7 +581,7 @@ runOnFunctionRecursively(SILFunction *F, FullApplySite AI,
551581

552582
// Now that the IR is correct, see if we can remove dead callee
553583
// computations (e.g. dead partial_apply closures).
554-
cleanupCalleeValue(CalleeValue, CaptureArgs, FullArgs);
584+
cleanupCalleeValue(CalleeValue, FullArgs);
555585

556586
// Reposition iterators possibly invalidated by mutation.
557587
BI = SILFunction::iterator(ApplyBlock);

‎test/SILOptimizer/mandatory_inlining.sil

Lines changed: 100 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1084,3 +1084,103 @@ bb11(%128 : $Error):
10841084
bb12(%result : $C):
10851085
return %result : $C
10861086
}
1087+
1088+
sil @use_c : $@convention(thin) (@guaranteed C) -> ()
1089+
sil @use_c_unowned : $@convention(thin) (C) -> ()
1090+
1091+
sil [transparent] @guaranteed_closure_func : $@convention(thin) (@guaranteed C) -> () {
1092+
bb0(%0 : $C):
1093+
%use = function_ref @use_c : $@convention(thin) (@guaranteed C) -> ()
1094+
apply %use(%0) : $@convention(thin) (@guaranteed C) -> ()
1095+
%t = tuple ()
1096+
return %t : $()
1097+
}
1098+
1099+
sil [transparent] @unowned_closure_func : $@convention(thin) (C) -> () {
1100+
bb0(%0 : $C):
1101+
%use = function_ref @use_c_unowned : $@convention(thin) (C) -> ()
1102+
apply %use(%0) : $@convention(thin) (C) -> ()
1103+
%t = tuple ()
1104+
return %t : $()
1105+
}
1106+
1107+
// CHECK-LABEL: sil @test_guaranteed_closure_capture
1108+
// CHECK: bb0([[ARG:%.*]] : $C):
1109+
// CHECK: strong_retain [[ARG]] : $C
1110+
// @guaranteed closure captures *don't* release the capture in the closure implementation.
1111+
// CHECK-NOT: strong_retain [[ARG]] : $C
1112+
// CHECK: [[F:%.*]] = function_ref @use_c
1113+
// CHECK: apply [[F]]([[ARG]])
1114+
// CHECK: strong_release [[ARG]]
1115+
// CHECK: return
1116+
sil @test_guaranteed_closure_capture : $@convention(thin) (@guaranteed C) -> () {
1117+
bb0(%0: $C):
1118+
strong_retain %0 : $C
1119+
%closure_fun = function_ref @guaranteed_closure_func : $@convention(thin) (@guaranteed C) -> ()
1120+
%closure = partial_apply [callee_guaranteed] %closure_fun(%0) : $@convention(thin) (@guaranteed C) -> ()
1121+
apply %closure() : $@callee_guaranteed () -> ()
1122+
strong_release %closure : $@callee_guaranteed () -> ()
1123+
%t = tuple ()
1124+
return %t : $()
1125+
}
1126+
1127+
// CHECK-LABEL: sil @test_guaranteed_closure_capture2
1128+
// CHECK: bb0([[ARG:%.*]] : $C):
1129+
// @guaranteed closure captures *don't* release the capture in the closure implementation.
1130+
// CHECK: strong_retain [[ARG]] : $C
1131+
// CHECK: strong_release [[ARG]] : $C
1132+
// CHECK: [[F:%.*]] = function_ref @use_c
1133+
// CHECK: apply [[F]]([[ARG]])
1134+
// CHECK-NOT: strong_release [[ARG]]
1135+
// CHECK: return
1136+
sil @test_guaranteed_closure_capture2 : $@convention(thin) (@guaranteed C) -> () {
1137+
bb0(%0: $C):
1138+
strong_retain %0 : $C
1139+
%closure_fun = function_ref @guaranteed_closure_func : $@convention(thin) (@guaranteed C) -> ()
1140+
%closure = partial_apply %closure_fun(%0) : $@convention(thin) (@guaranteed C) -> ()
1141+
apply %closure() : $@callee_owned () -> ()
1142+
%t = tuple ()
1143+
return %t : $()
1144+
}
1145+
1146+
// CHECK-LABEL: sil @test_unowned_closure_capture2
1147+
// CHECK: bb0([[ARG:%.*]] : $C):
1148+
// CHECK: strong_retain [[ARG]]
1149+
// CHECK: strong_release [[ARG]]
1150+
// CHECK-NOT: strong_release [[ARG]]
1151+
// CHECK: [[F:%.*]] = function_ref @use_c_unowned
1152+
// CHECK: apply [[F]]([[ARG]]) : $@convention(thin) (C) -> ()
1153+
// CHECK-NOT: strong_release [[ARG]]
1154+
// CHECK: return
1155+
sil @test_unowned_closure_capture2 : $@convention(thin) (@guaranteed C) -> () {
1156+
bb0(%0: $C):
1157+
strong_retain %0 : $C
1158+
%closure_fun = function_ref @unowned_closure_func : $@convention(thin) (C) -> ()
1159+
%closure = partial_apply %closure_fun(%0) : $@convention(thin) (C) -> ()
1160+
apply %closure() : $@callee_owned () -> ()
1161+
%t = tuple ()
1162+
return %t : $()
1163+
}
1164+
// CHECK-LABEL: sil @test_guaranteed_closure
1165+
// CHECK: bb0([[ARG:%.*]] : $C):
1166+
// CHECK: strong_retain [[ARG]] : $C
1167+
// CHECK: [[F:%.*]] = function_ref @guaranteed_closure_func
1168+
// CHECK: [[C:%.*]] = partial_apply [callee_guaranteed] [[F]]([[ARG]])
1169+
// @guaranteed closure captures *don't* release the capture in the closure implementation.
1170+
// CHECK-NOT: strong_retain [[ARG]] : $C
1171+
// @callee_guaranteed closures *don't* release the context on application.
1172+
// CHECK-NOT: strong_release [[C]]
1173+
// CHECK: [[F:%.*]] = function_ref @use_c
1174+
// CHECK: apply [[F]]([[ARG]])
1175+
// CHECK: return [[C]]
1176+
1177+
sil @test_guaranteed_closure : $@convention(thin) (@guaranteed C) -> @owned @callee_guaranteed () -> () {
1178+
bb0(%0: $C):
1179+
strong_retain %0 : $C
1180+
%closure_fun = function_ref @guaranteed_closure_func : $@convention(thin) (@guaranteed C) -> ()
1181+
%closure = partial_apply [callee_guaranteed] %closure_fun(%0) : $@convention(thin) (@guaranteed C) -> ()
1182+
// NOTE: No ``strong_retain %closure`` is needed here because the context is
1183+
// @callee_guaranteed.
1184+
apply %closure() : $@callee_guaranteed () -> ()
1185+
return %closure : $@callee_guaranteed () -> ()
1186+
}

0 commit comments

Comments
 (0)