Skip to content

Commit 78f5211

Browse files
swiftixswiftix
committed
[capture-propagation] Support generic partial_apply instructions
1 parent b3f5583 commit 78f5211

File tree

2 files changed

+164
-30
lines changed

2 files changed

+164
-30
lines changed

‎lib/SILOptimizer/IPO/CapturePropagation.cpp

Lines changed: 92 additions & 30 deletions
Original file line numberDiff line numberDiff line change
@@ -11,12 +11,15 @@
1111
//===----------------------------------------------------------------------===//
1212

1313
#define DEBUG_TYPE "capture-prop"
14+
#include "swift/AST/GenericEnvironment.h"
1415
#include "swift/SILOptimizer/PassManager/Passes.h"
16+
#include "swift/SILOptimizer/Utils/Generics.h"
1517
#include "swift/SILOptimizer/Utils/SpecializationMangler.h"
1618
#include "swift/Demangling/Demangle.h"
1719
#include "swift/SIL/Mangle.h"
1820
#include "swift/SIL/SILCloner.h"
1921
#include "swift/SIL/SILInstruction.h"
22+
#include "swift/SIL/TypeSubstCloner.h"
2023
#include "swift/SILOptimizer/Analysis/ColdBlockInfo.h"
2124
#include "swift/SILOptimizer/Analysis/DominanceAnalysis.h"
2225
#include "swift/SILOptimizer/PassManager/Transforms.h"
@@ -99,16 +102,17 @@ namespace {
99102
/// caller, so the cloned function will have a mix of locations from different
100103
/// functions.
101104
class CapturePropagationCloner
102-
: public SILClonerWithScopes<CapturePropagationCloner> {
103-
using SuperTy = SILClonerWithScopes<CapturePropagationCloner>;
105+
: public TypeSubstCloner<CapturePropagationCloner> {
106+
using SuperTy = TypeSubstCloner<CapturePropagationCloner>;
104107
friend class SILVisitor<CapturePropagationCloner>;
105108
friend class SILCloner<CapturePropagationCloner>;
106109

107110
SILFunction *OrigF;
108111
bool IsCloningConstant;
109112
public:
110-
CapturePropagationCloner(SILFunction *OrigF, SILFunction *NewF)
111-
: SuperTy(*NewF), OrigF(OrigF), IsCloningConstant(false) {}
113+
CapturePropagationCloner(SILFunction *OrigF, SILFunction *NewF,
114+
SubstitutionList Subs)
115+
: SuperTy(*NewF, *OrigF, Subs), OrigF(OrigF), IsCloningConstant(false) {}
112116

113117
void cloneBlocks(OperandValueArrayRef Args);
114118

@@ -219,6 +223,20 @@ void CapturePropagationCloner::cloneBlocks(
219223
}
220224
}
221225

226+
CanSILFunctionType getPartialApplyInterfaceResultType(PartialApplyInst *PAI) {
227+
SILFunction *OrigF = PAI->getReferencedFunction();
228+
// The new partial_apply will no longer take any arguments--they are all
229+
// expressed as literals. So its callee signature will be the same as its
230+
// return signature.
231+
auto FTy = PAI->getType().castTo<SILFunctionType>();
232+
CanGenericSignature CanGenericSig;
233+
assert(!PAI->hasSubstitutions() || !hasArchetypes(PAI->getSubstitutions()));
234+
FTy = cast<SILFunctionType>(
235+
OrigF->mapTypeOutOfContext(FTy)->getCanonicalType());
236+
auto NewFTy = FTy;
237+
return NewFTy;
238+
}
239+
222240
/// Given a partial_apply instruction, create a specialized callee by removing
223241
/// all constant arguments and adding constant literals to the specialized
224242
/// function body.
@@ -243,12 +261,16 @@ SILFunction *CapturePropagation::specializeConstClosure(PartialApplyInst *PAI,
243261
// The new partial_apply will no longer take any arguments--they are all
244262
// expressed as literals. So its callee signature will be the same as its
245263
// return signature.
246-
CanSILFunctionType NewFTy =
247-
Lowering::adjustFunctionType(PAI->getType().castTo<SILFunctionType>(),
248-
SILFunctionType::Representation::Thin);
264+
auto NewFTy = getPartialApplyInterfaceResultType(PAI);
265+
NewFTy = Lowering::adjustFunctionType(NewFTy,
266+
SILFunctionType::Representation::Thin);
267+
268+
GenericEnvironment *GenericEnv = nullptr;
269+
if (NewFTy->getGenericSignature())
270+
GenericEnv = OrigF->getGenericEnvironment();
249271
SILFunction *NewF = OrigF->getModule().createFunction(
250272
SILLinkage::Shared, Name, NewFTy,
251-
OrigF->getGenericEnvironment(), OrigF->getLocation(), OrigF->isBare(),
273+
GenericEnv, OrigF->getLocation(), OrigF->isBare(),
252274
OrigF->isTransparent(), Fragile, OrigF->isThunk(),
253275
OrigF->getClassVisibility(), OrigF->getInlineStrategy(),
254276
OrigF->getEffectsKind(),
@@ -259,18 +281,28 @@ SILFunction *CapturePropagation::specializeConstClosure(PartialApplyInst *PAI,
259281
DEBUG(llvm::dbgs() << " Specialize callee as ";
260282
NewF->printName(llvm::dbgs()); llvm::dbgs() << " " << NewFTy << "\n");
261283

262-
CapturePropagationCloner cloner(OrigF, NewF);
284+
DEBUG(if (PAI->hasSubstitutions()) {
285+
llvm::dbgs() << "CapturePropagation of generic partial_apply:\n";
286+
PAI->dumpInContext();
287+
});
288+
CapturePropagationCloner cloner(OrigF, NewF, PAI->getSubstitutions());
263289
cloner.cloneBlocks(PAI->getArguments());
264290
assert(OrigF->getDebugScope()->Parent != NewF->getDebugScope()->Parent);
265291
return NewF;
266292
}
267293

268294
void CapturePropagation::rewritePartialApply(PartialApplyInst *OrigPAI,
269295
SILFunction *SpecialF) {
296+
DEBUG(llvm::dbgs() << "\n Rewriting a partial apply:\n";
297+
OrigPAI->dumpInContext(); llvm::dbgs() << " with special function: "
298+
<< SpecialF->getName() << "\n";
299+
llvm::dbgs() << "\nThe function being rewritten is:\n";
300+
OrigPAI->getFunction()->dump());
301+
270302
SILBuilderWithScope Builder(OrigPAI);
271303
auto FuncRef = Builder.createFunctionRef(OrigPAI->getLoc(), SpecialF);
272-
auto *T2TF = Builder.createThinToThickFunction(OrigPAI->getLoc(),
273-
FuncRef, OrigPAI->getType());
304+
auto *T2TF = Builder.createThinToThickFunction(OrigPAI->getLoc(), FuncRef,
305+
OrigPAI->getType());
274306
OrigPAI->replaceAllUsesWith(T2TF);
275307
recursivelyDeleteTriviallyDeadInstructions(OrigPAI, true);
276308
DEBUG(llvm::dbgs() << " Rewrote caller:\n" << *T2TF);
@@ -311,12 +343,16 @@ static bool onlyContainsReturnOrThrowOfArg(SILBasicBlock *BB) {
311343

312344
/// Checks if \p Orig is a thunk which calls another function but without
313345
/// passing the trailing \p numDeadParams dead parameters.
314-
static SILFunction *getSpecializedWithDeadParams(SILFunction *Orig,
315-
int numDeadParams) {
346+
/// If a generic specialization was performed for a generic capture,
347+
/// GenericSpecialized contains a tuple:
348+
/// (new specialized function, old function)
349+
static SILFunction *getSpecializedWithDeadParams(
350+
PartialApplyInst *PAI, SILFunction *Orig, int numDeadParams,
351+
std::pair<SILFunction *, SILFunction *> &GenericSpecialized) {
316352
SILBasicBlock &EntryBB = *Orig->begin();
317353
unsigned NumArgs = EntryBB.getNumArguments();
318354
SILModule &M = Orig->getModule();
319-
355+
320356
// Check if all dead parameters have trivial types. We don't support non-
321357
// trivial types because it's very hard to find places where we can release
322358
// those parameters (as a replacement for the removed partial_apply).
@@ -328,20 +364,20 @@ static SILFunction *getSpecializedWithDeadParams(SILFunction *Orig,
328364
}
329365
SILFunction *Specialized = nullptr;
330366
SILValue RetValue;
331-
367+
332368
// Check all instruction of the entry block.
333369
for (SILInstruction &I : EntryBB) {
334370
if (auto FAS = FullApplySite::isa(&I)) {
335-
336371
// Check if this is the call of the specialized function.
337-
// As the original function is not generic, also the specialized function
338-
// must be not generic.
339-
if (FAS.hasSubstitutions())
372+
// If the original partial_apply didn't have substitutions,
373+
// also the specialized function must be not generic.
374+
if (!PAI->hasSubstitutions() && FAS.hasSubstitutions())
340375
return nullptr;
376+
341377
// Is it the only call?
342378
if (Specialized)
343379
return nullptr;
344-
380+
345381
Specialized = FAS.getReferencedFunction();
346382
if (!Specialized)
347383
return nullptr;
@@ -376,29 +412,54 @@ static SILFunction *getSpecializedWithDeadParams(SILFunction *Orig,
376412
if (I.mayHaveSideEffects() || isa<TermInst>(&I))
377413
return nullptr;
378414
}
415+
416+
GenericSpecialized = std::make_pair(nullptr, nullptr);
417+
418+
if (PAI->hasSubstitutions()) {
419+
if (Specialized->isExternalDeclaration())
420+
return nullptr;
421+
// Perform a generic specialization of the Specialized function.
422+
ReabstractionInfo ReInfo(ApplySite(), Specialized, PAI->getSubstitutions(),
423+
/* ConvertIndirectToDirect */ false);
424+
GenericFuncSpecializer FuncSpecializer(Specialized,
425+
ReInfo.getClonerParamSubstitutions(),
426+
Specialized->isFragile(), ReInfo);
427+
428+
SILFunction *GenericSpecializedFunc = FuncSpecializer.trySpecialization();
429+
if (!GenericSpecializedFunc)
430+
return nullptr;
431+
GenericSpecialized = std::make_pair(GenericSpecializedFunc, Specialized);
432+
return GenericSpecializedFunc;
433+
}
379434
return Specialized;
380435
}
381436

382437
bool CapturePropagation::optimizePartialApply(PartialApplyInst *PAI) {
383-
// Check if the partial_apply has generic substitutions.
384-
// FIXME: We could handle generic thunks if it's worthwhile.
385-
if (PAI->hasSubstitutions())
386-
return false;
387-
388438
SILFunction *SubstF = PAI->getReferencedFunction();
389439
if (!SubstF)
390440
return false;
391441
if (SubstF->isExternalDeclaration())
392442
return false;
393443

394-
assert(!SubstF->getLoweredFunctionType()->isPolymorphic() &&
395-
"cannot specialize generic partial apply");
444+
if (PAI->hasSubstitutions() && hasArchetypes(PAI->getSubstitutions())) {
445+
DEBUG(llvm::dbgs()
446+
<< "CapturePropagation: cannot handle partial specialization "
447+
"of partial_apply:\n";
448+
PAI->dumpInContext());
449+
return false;
450+
}
451+
396452

397453
// First possibility: Is it a partial_apply where all partially applied
398454
// arguments are dead?
399-
if (SILFunction *NewFunc = getSpecializedWithDeadParams(SubstF,
400-
PAI->getNumArguments())) {
455+
std::pair<SILFunction *, SILFunction *> GenericSpecialized;
456+
if (auto *NewFunc = getSpecializedWithDeadParams(
457+
PAI, SubstF, PAI->getNumArguments(), GenericSpecialized)) {
401458
rewritePartialApply(PAI, NewFunc);
459+
if (GenericSpecialized.first) {
460+
// Notify the pass manager about the new function.
461+
notifyAddFunction(GenericSpecialized.first, GenericSpecialized.second);
462+
}
402463
return true;
403464
}
404465

@@ -411,7 +472,8 @@ bool CapturePropagation::optimizePartialApply(PartialApplyInst *PAI) {
411472
return false;
412473

413474
DEBUG(llvm::dbgs() << "Specializing closure for constant arguments:\n"
414-
<< " " << SubstF->getName() << "\n" << *PAI);
475+
<< " " << SubstF->getName() << "\n"
476+
<< *PAI);
415477
++NumCapturesPropagated;
416478
SILFunction *NewF = specializeConstClosure(PAI, SubstF);
417479
rewritePartialApply(PAI, NewF);

‎test/SILOptimizer/capture_propagation.sil

Lines changed: 72 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -378,3 +378,75 @@ bb0:
378378
return %2 : $@callee_owned (Int32, Int32) -> (Bool, @error Error)
379379
}
380380

381+
// Test generic capture propagation
382+
383+
sil @_TFtest_generic_capture_propagation2_closure : $@convention(thin) <T> (Builtin.Int32, Builtin.FPIEEE32, Builtin.RawPointer, @in T) -> () {
384+
bb0(%0 : $Builtin.Int32, %1 : $Builtin.FPIEEE32, %2 : $Builtin.RawPointer, %3 : $*T):
385+
%9999 = tuple()
386+
return %9999 : $()
387+
}
388+
389+
// CHECK-LABEL: sil @test_generic_capture_propagation2_caller
390+
// CHECK: %[[CALLEE:[0-9]+]] = function_ref @test_generic_capture_propagation2_callee
391+
// CHECK: %[[FR:[0-9]+]] = function_ref @{{.*}}test_generic_capture_propagation2_thunk : $@convention(thin) () -> ()
392+
// CHECK: %[[CONVERTED:[0-9]+]] = thin_to_thick_function %[[FR]] : $@convention(thin) () -> () to $@callee_owned () -> ()
393+
// CHECK-NOT: partial_apply
394+
// CHECK: apply %[[CALLEE]](%[[CONVERTED]]) : $@convention(thin) (@owned @callee_owned () -> ()) -> ()
395+
// CHECL-NOT: partial_apply
396+
// CHECK: return
397+
sil @test_generic_capture_propagation2_caller : $@convention(thin) () -> () {
398+
%0 = integer_literal $Builtin.Int32, 0
399+
%1 = float_literal $Builtin.FPIEEE32, 0
400+
%2 = string_literal utf8 "123"
401+
%3 = global_addr @globalinit_33_06E7F1D906492AE070936A9B58CBAE1C_token8 : $*Builtin.Word
402+
%4 = function_ref @_TFtest_generic_capture_propagation2_closure : $@convention(thin) <T> (Builtin.Int32, Builtin.FPIEEE32, Builtin.RawPointer, @in T) -> ()
403+
%5 = thin_to_thick_function %4 : $@convention(thin) <T> (Builtin.Int32, Builtin.FPIEEE32, Builtin.RawPointer, @in T) -> () to $@callee_owned <T> (Builtin.Int32, Builtin.FPIEEE32, Builtin.RawPointer, @in T) -> ()
404+
%6 = function_ref @test_generic_capture_propagation2_callee : $@convention(thin) (@owned @callee_owned () -> ()) -> ()
405+
%7 = function_ref @test_generic_capture_propagation2_thunk : $@convention(thin) <T> (Builtin.Int32, Builtin.FPIEEE32, Builtin.RawPointer, @in T, @owned @callee_owned <T> (Builtin.Int32, Builtin.FPIEEE32, Builtin.RawPointer, @in T) -> ()) -> ()
406+
%8 = partial_apply %7<Builtin.Word>(%0, %1, %2, %3, %5) : $@convention(thin) <T> (Builtin.Int32, Builtin.FPIEEE32, Builtin.RawPointer, @in T, @owned @callee_owned <T> (Builtin.Int32, Builtin.FPIEEE32, Builtin.RawPointer, @in T) -> ()) -> ()
407+
apply %6(%8) : $@convention(thin) (@owned @callee_owned () -> ()) -> ()
408+
%9999 = tuple()
409+
return %9999 : $()
410+
}
411+
412+
sil shared @test_generic_capture_propagation2_thunk : $@convention(thin) <T> (Builtin.Int32, Builtin.FPIEEE32, Builtin.RawPointer, @in T, @owned @callee_owned <T> (Builtin.Int32, Builtin.FPIEEE32, Builtin.RawPointer, @in T) -> ()) -> () {
413+
bb0(%0 : $Builtin.Int32, %1 : $Builtin.FPIEEE32, %2 : $Builtin.RawPointer, %3 : $*T, %4 : $@callee_owned <T> (Builtin.Int32, Builtin.FPIEEE32, Builtin.RawPointer, @in T) -> ()):
414+
apply %4<T>(%0, %1, %2, %3) : $@callee_owned <T> (Builtin.Int32, Builtin.FPIEEE32, Builtin.RawPointer, @in T) -> ()
415+
%9999 = tuple()
416+
return %9999 : $()
417+
}
418+
419+
sil shared @test_generic_capture_propagation2_callee : $@convention(thin) (@owned @callee_owned () -> ()) -> () {
420+
bb0(%0 : $@callee_owned () -> ()):
421+
apply %0() : $@callee_owned () -> ()
422+
%9999 = tuple()
423+
return %9999 : $()
424+
}
425+
426+
// Test dead partial applied arguments when using generics
427+
428+
sil @specialized_generic_nonthrowing_closure : $@convention(thin) <T> (@in T, @in T) -> Bool {
429+
bb0(%0 : $*T, %1 : $*T):
430+
%10 = integer_literal $Builtin.Int1, -1
431+
%9999 = struct $Bool (%10 : $Builtin.Int1)
432+
return %9999 : $Bool
433+
}
434+
435+
sil @nonthrowing_generic_closure : $@convention(method) <T> (@in T, @in T, @thin T.Type) -> Bool {
436+
bb0(%0 : $*T, %1 : $*T, %2 : $@thin T.Type):
437+
%3 = function_ref @specialized_generic_nonthrowing_closure : $@convention(thin) <T> (@in T, @in T) -> Bool
438+
%4 = apply %3<T>(%0, %1) : $@convention(thin) <T> (@in T, @in T) -> Bool
439+
return %4 : $Bool
440+
}
441+
442+
// CHECK-LABEL: sil @return_generic_nonthrowing_closure
443+
// CHECK: [[F:%[0-9]+]] = function_ref @_TTSg5Vs5Int32__specialized_generic_nonthrowing_closure
444+
// CHECK: [[R:%[0-9]+]] = thin_to_thick_function [[F]]
445+
// CHECK: return [[R]]
446+
sil @return_generic_nonthrowing_closure : $@convention(thin) () -> @owned @callee_owned (@in Int32, @in Int32) -> Bool {
447+
bb0:
448+
%0 = metatype $@thin Int32.Type
449+
%1 = function_ref @nonthrowing_generic_closure : $@convention(method) <T> (@in T, @in T, @thin T.Type) -> Bool
450+
%2 = partial_apply %1<Int32>(%0) : $@convention(method) <T>(@in T, @in T, @thin T.Type) -> Bool
451+
return %2 : $@callee_owned (@in Int32, @in Int32) -> Bool
452+
}

0 commit comments

Comments
 (0)