Skip to content

Commit 01654fd

Browse files
jkshtjjkshtj
authored
Merge pull request #71775 from jkshtj/main
[Autodiff] Adds part of the Autodiff specific closure-specialization optimization pass
2 parents e74cfb0 + c6330a7 commit 01654fd

38 files changed

+3682
-91
lines changed

‎SwiftCompilerSources/Sources/Optimizer/FunctionPasses/CMakeLists.txt

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -12,6 +12,7 @@ swift_compiler_sources(Optimizer
1212
AsyncDemotion.swift
1313
BooleanLiteralFolding.swift
1414
CleanupDebugSteps.swift
15+
ClosureSpecialization.swift
1516
ComputeEscapeEffects.swift
1617
ComputeSideEffects.swift
1718
DeadStoreElimination.swift

‎SwiftCompilerSources/Sources/Optimizer/FunctionPasses/ClosureSpecialization.swift

Lines changed: 720 additions & 0 deletions
Large diffs are not rendered by default.

‎SwiftCompilerSources/Sources/Optimizer/PassManager/PassRegistration.swift

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -93,6 +93,9 @@ private func registerSwiftPasses() {
9393
registerPass(lifetimeDependenceDiagnosticsPass, { lifetimeDependenceDiagnosticsPass.run($0) })
9494
registerPass(lifetimeDependenceInsertionPass, { lifetimeDependenceInsertionPass.run($0) })
9595
registerPass(lifetimeDependenceScopeFixupPass, { lifetimeDependenceScopeFixupPass.run($0) })
96+
registerPass(generalClosureSpecialization, { generalClosureSpecialization.run($0) })
97+
registerPass(autodiffClosureSpecialization, { autodiffClosureSpecialization.run($0) })
98+
9699
// Instruction passes
97100
registerForSILCombine(BeginCOWMutationInst.self, { run(BeginCOWMutationInst.self, $0) })
98101
registerForSILCombine(GlobalValueInst.self, { run(GlobalValueInst.self, $0) })

‎SwiftCompilerSources/Sources/Optimizer/Utilities/OptUtils.swift

Lines changed: 21 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -398,6 +398,23 @@ extension LoadInst {
398398
}
399399
}
400400

401+
extension PartialApplyInst {
402+
var isPartialApplyOfReabstractionThunk: Bool {
403+
// A partial_apply of a reabstraction thunk either has a single capture
404+
// (a function) or two captures (function and dynamic Self type).
405+
if self.numArguments == 1 || self.numArguments == 2,
406+
let fun = self.referencedFunction,
407+
fun.isReabstractionThunk,
408+
self.arguments[0].type.isFunction,
409+
self.arguments[0].type.isReferenceCounted(in: self.parentFunction) || self.callee.type.isThickFunction
410+
{
411+
return true
412+
}
413+
414+
return false
415+
}
416+
}
417+
401418
extension FunctionPassContext {
402419
/// Returns true if any blocks were removed.
403420
func removeDeadBlocks(in function: Function) -> Bool {
@@ -540,6 +557,10 @@ extension Function {
540557
}
541558
return nil
542559
}
560+
561+
var mayBindDynamicSelf: Bool {
562+
self.bridged.mayBindDynamicSelf()
563+
}
543564
}
544565

545566
extension FullApplySite {

‎SwiftCompilerSources/Sources/Optimizer/Utilities/Test.swift

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -164,7 +164,8 @@ public func registerOptimizerTests() {
164164
lifetimeDependenceUseTest,
165165
linearLivenessTest,
166166
parseTestSpecificationTest,
167-
variableIntroducerTest
167+
variableIntroducerTest,
168+
gatherCallSitesTest
168169
)
169170

170171
// Finally register the thunk they all call through.

‎SwiftCompilerSources/Sources/SIL/ApplySite.swift

Lines changed: 13 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -108,6 +108,10 @@ public protocol ApplySite : Instruction {
108108
extension ApplySite {
109109
public var callee: Value { operands[ApplyOperandConventions.calleeIndex].value }
110110

111+
public var hasSubstitutions: Bool {
112+
return substitutionMap.hasAnySubstitutableParams
113+
}
114+
111115
public var isAsync: Bool {
112116
return callee.type.isAsyncFunction
113117
}
@@ -126,7 +130,15 @@ extension ApplySite {
126130
return false
127131
}
128132

129-
/// Returns the subset of operands that are argument operands.
133+
public var isCalleeNoReturn: Bool {
134+
bridged.ApplySite_isCalleeNoReturn()
135+
}
136+
137+
public var isCalleeTrapNoReturn: Bool {
138+
referencedFunction?.isTrapNoReturn ?? false
139+
}
140+
141+
/// Returns the subset of operands which are argument operands.
130142
///
131143
/// This does not include the callee function operand.
132144
public var argumentOperands: OperandArray {

‎SwiftCompilerSources/Sources/SIL/BasicBlock.swift

Lines changed: 20 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -57,6 +57,26 @@ final public class BasicBlock : CustomStringConvertible, HasShortDescription, Ha
5757
successors.count == 1 ? successors[0] : nil
5858
}
5959

60+
/// All function exiting blocks except for ones with an `unreachable` terminator,
61+
/// not immediately preceded by an apply of a no-return function.
62+
public var isReachableExitBlock: Bool {
63+
switch terminator {
64+
case let termInst where termInst.isFunctionExiting:
65+
return true
66+
case is UnreachableInst:
67+
if let instBeforeUnreachable = terminator.previous,
68+
let ai = instBeforeUnreachable as? ApplyInst,
69+
ai.isCalleeNoReturn && !ai.isCalleeTrapNoReturn
70+
{
71+
return true
72+
}
73+
74+
return false
75+
default:
76+
return false
77+
}
78+
}
79+
6080
/// The index of the basic block in its function.
6181
/// This has O(n) complexity. Only use it for debugging
6282
public var index: Int {

‎SwiftCompilerSources/Sources/SIL/Function.swift

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -31,6 +31,12 @@ final public class Function : CustomStringConvertible, HasShortDescription, Hash
3131
hasher.combine(ObjectIdentifier(self))
3232
}
3333

34+
public var isTrapNoReturn: Bool { bridged.isTrapNoReturn() }
35+
36+
public var isAutodiffVJP: Bool { bridged.isAutodiffVJP() }
37+
38+
public var specializationLevel: Int { bridged.specializationLevel() }
39+
3440
public var hasOwnership: Bool { bridged.hasOwnership() }
3541

3642
public var hasLoweredAddresses: Bool { bridged.hasLoweredAddresses() }
@@ -59,6 +65,10 @@ final public class Function : CustomStringConvertible, HasShortDescription, Hash
5965
entryBlock.arguments.lazy.map { $0 as! FunctionArgument }
6066
}
6167

68+
public func argument(at index: Int) -> FunctionArgument {
69+
entryBlock.arguments[index] as! FunctionArgument
70+
}
71+
6272
/// All instructions of all blocks.
6373
public var instructions: LazySequence<FlattenSequence<LazyMapSequence<BasicBlockList, InstructionList>>> {
6474
blocks.lazy.flatMap { $0.instructions }
@@ -83,6 +93,8 @@ final public class Function : CustomStringConvertible, HasShortDescription, Hash
8393

8494
public var isAsync: Bool { bridged.isAsync() }
8595

96+
public var isReabstractionThunk: Bool { bridged.isReabstractionThunk() }
97+
8698
/// True if this is a `[global_init]` function.
8799
///
88100
/// Such a function is typically a global addressor which calls the global's

‎SwiftCompilerSources/Sources/SIL/Instruction.swift

Lines changed: 10 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -893,7 +893,16 @@ class ConvertFunctionInst : SingleValueInstruction, UnaryInstruction {
893893
}
894894

895895
final public
896-
class ThinToThickFunctionInst : SingleValueInstruction, UnaryInstruction {}
896+
class ThinToThickFunctionInst : SingleValueInstruction, UnaryInstruction {
897+
public var callee: Value { operand.value }
898+
899+
public var referencedFunction: Function? {
900+
if let fri = callee as? FunctionRefInst {
901+
return fri.referencedFunction
902+
}
903+
return nil
904+
}
905+
}
897906

898907
final public class ThickToObjCMetatypeInst : SingleValueInstruction {}
899908
final public class ObjCToThickMetatypeInst : SingleValueInstruction {}

‎SwiftCompilerSources/Sources/SIL/SubstitutionMap.swift

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -25,6 +25,8 @@ public struct SubstitutionMap {
2525

2626
public var isEmpty: Bool { bridged.isEmpty() }
2727

28+
public var hasAnySubstitutableParams: Bool { bridged.hasAnySubstitutableParams() }
29+
2830
public var replacementTypes: OptionalTypeArray {
2931
let types = BridgedTypeArray.fromReplacementTypes(bridged)
3032
return OptionalTypeArray(bridged: types)

‎SwiftCompilerSources/Sources/SIL/Type.swift

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -58,6 +58,7 @@ public struct Type : CustomStringConvertible, NoReflectionChildren {
5858
public var isMetatype: Bool { bridged.isMetatype() }
5959
public var isNoEscapeFunction: Bool { bridged.isNoEscapeFunction() }
6060
public var containsNoEscapeFunction: Bool { bridged.containsNoEscapeFunction() }
61+
public var isThickFunction: Bool { bridged.isThickFunction() }
6162
public var isAsyncFunction: Bool { bridged.isAsyncFunction() }
6263

6364
public var canBeClass: BridgedType.TraitResult { bridged.canBeClass() }

‎include/swift/AST/SILOptions.h

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -290,6 +290,10 @@ class SILOptions {
290290
/// Are we building in embedded Swift + -no-allocations?
291291
bool NoAllocations = false;
292292

293+
/// Should we use the experimental Swift based closure-specialization
294+
/// optimization pass instead of the existing C++ one.
295+
bool EnableExperimentalSwiftBasedClosureSpecialization = false;
296+
293297
/// The name of the file to which the backend should save optimization
294298
/// records.
295299
std::string OptRecordFile;

‎include/swift/Option/FrontendOptions.td

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -359,6 +359,10 @@ def enable_experimental_async_top_level :
359359
// HIDDEN FLAGS
360360
let Flags = [FrontendOption, NoDriverOption, HelpHidden] in {
361361

362+
def enable_experimental_swift_based_closure_specialization :
363+
Flag<["-"], "experimental-swift-based-closure-specialization">,
364+
HelpText<"Use the experimental Swift based closure-specialization optimization pass instead of the existing C++ one">;
365+
362366
def checked_async_objc_bridging : Joined<["-"], "checked-async-objc-bridging=">,
363367
HelpText<"Control whether checked continuations are used when bridging "
364368
"async calls from Swift to ObjC: 'on', 'off' ">;

‎include/swift/SIL/SILBridging.h

Lines changed: 10 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -347,6 +347,7 @@ struct BridgedType {
347347
BRIDGED_INLINE bool isMetatype() const;
348348
BRIDGED_INLINE bool isNoEscapeFunction() const;
349349
BRIDGED_INLINE bool containsNoEscapeFunction() const;
350+
BRIDGED_INLINE bool isThickFunction() const;
350351
BRIDGED_INLINE bool isAsyncFunction() const;
351352
BRIDGED_INLINE bool isEmpty(BridgedFunction f) const;
352353
BRIDGED_INLINE TraitResult canBeClass() const;
@@ -373,7 +374,8 @@ struct BridgedType {
373374
BRIDGED_INLINE bool isEndCaseIterator(EnumElementIterator i) const;
374375
SWIFT_IMPORT_UNSAFE BRIDGED_INLINE BridgedType getEnumCasePayload(EnumElementIterator i, BridgedFunction f) const;
375376
BRIDGED_INLINE SwiftInt getNumTupleElements() const;
376-
SWIFT_IMPORT_UNSAFE BRIDGED_INLINE BridgedType getTupleElementType(SwiftInt idx) const;
377+
SWIFT_IMPORT_UNSAFE BRIDGED_INLINE BridgedType
378+
getTupleElementType(SwiftInt idx) const;
377379
SWIFT_IMPORT_UNSAFE BRIDGED_INLINE BridgedType getFunctionTypeWithNoEscape(bool withNoEscape) const;
378380
};
379381

@@ -547,13 +549,15 @@ struct BridgedFunction {
547549
BRIDGED_INLINE bool isAvailableExternally() const;
548550
BRIDGED_INLINE bool isTransparent() const;
549551
BRIDGED_INLINE bool isAsync() const;
552+
BRIDGED_INLINE bool isReabstractionThunk() const;
550553
BRIDGED_INLINE bool isGlobalInitFunction() const;
551554
BRIDGED_INLINE bool isGlobalInitOnceFunction() const;
552555
BRIDGED_INLINE bool isDestructor() const;
553556
BRIDGED_INLINE bool isGeneric() const;
554557
BRIDGED_INLINE bool hasSemanticsAttr(BridgedStringRef attrName) const;
555558
BRIDGED_INLINE bool hasUnsafeNonEscapableResult() const;
556559
BRIDGED_INLINE bool hasResultDependsOnSelf() const;
560+
bool mayBindDynamicSelf() const;
557561
BRIDGED_INLINE EffectsKind getEffectAttribute() const;
558562
BRIDGED_INLINE PerformanceConstraints getPerformanceConstraints() const;
559563
BRIDGED_INLINE InlineStrategy getInlineStrategy() const;
@@ -566,6 +570,9 @@ struct BridgedFunction {
566570
BRIDGED_INLINE void setIsPerformanceConstraint(bool isPerfConstraint) const;
567571
BRIDGED_INLINE bool isResilientNominalDecl(BridgedNominalTypeDecl decl) const;
568572
BRIDGED_INLINE BridgedType getLoweredType(BridgedASTType type) const;
573+
bool isTrapNoReturn() const;
574+
bool isAutodiffVJP() const;
575+
SwiftInt specializationLevel() const;
569576

570577
enum class ParseEffectsMode {
571578
argumentEffectsFromSource,
@@ -658,6 +665,7 @@ struct BridgedSubstitutionMap {
658665

659666
BRIDGED_INLINE BridgedSubstitutionMap();
660667
BRIDGED_INLINE bool isEmpty() const;
668+
BRIDGED_INLINE bool hasAnySubstitutableParams() const;
661669
};
662670

663671
struct BridgedTypeArray {
@@ -917,6 +925,7 @@ struct BridgedInstruction {
917925
SWIFT_IMPORT_UNSAFE BRIDGED_INLINE BridgedSubstitutionMap ApplySite_getSubstitutionMap() const;
918926
SWIFT_IMPORT_UNSAFE BRIDGED_INLINE BridgedASTType ApplySite_getSubstitutedCalleeType() const;
919927
BRIDGED_INLINE SwiftInt ApplySite_getNumArguments() const;
928+
BRIDGED_INLINE bool ApplySite_isCalleeNoReturn() const;
920929
BRIDGED_INLINE SwiftInt FullApplySite_numIndirectResultArguments() const;
921930

922931
// =========================================================================//

‎include/swift/SIL/SILBridgingImpl.h

Lines changed: 16 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -262,6 +262,10 @@ bool BridgedType::containsNoEscapeFunction() const {
262262
return unbridged().containsNoEscapeFunction();
263263
}
264264

265+
bool BridgedType::isThickFunction() const {
266+
return unbridged().isThickFunction();
267+
}
268+
265269
bool BridgedType::isAsyncFunction() const {
266270
return unbridged().isAsyncFunction();
267271
}
@@ -530,6 +534,10 @@ bool BridgedSubstitutionMap::isEmpty() const {
530534
return unbridged().empty();
531535
}
532536

537+
bool BridgedSubstitutionMap::hasAnySubstitutableParams() const {
538+
return unbridged().hasAnySubstitutableParams();
539+
}
540+
533541
//===----------------------------------------------------------------------===//
534542
// BridgedLocation
535543
//===----------------------------------------------------------------------===//
@@ -635,6 +643,10 @@ bool BridgedFunction::isAsync() const {
635643
return getFunction()->isAsync();
636644
}
637645

646+
bool BridgedFunction::isReabstractionThunk() const {
647+
return getFunction()->isThunk() == swift::IsReabstractionThunk;
648+
}
649+
638650
bool BridgedFunction::isGlobalInitFunction() const {
639651
return getFunction()->isGlobalInit();
640652
}
@@ -1283,6 +1295,10 @@ SwiftInt BridgedInstruction::ApplySite_getNumArguments() const {
12831295
return swift::ApplySite(unbridged()).getNumArguments();
12841296
}
12851297

1298+
bool BridgedInstruction::ApplySite_isCalleeNoReturn() const {
1299+
return swift::ApplySite(unbridged()).isCalleeNoReturn();
1300+
}
1301+
12861302
SwiftInt BridgedInstruction::FullApplySite_numIndirectResultArguments() const {
12871303
auto fas = swift::FullApplySite(unbridged());
12881304
return fas.getNumIndirectSILResults();

‎include/swift/SIL/SILType.h

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -555,6 +555,13 @@ class SILType {
555555
// Handle whatever AST types are known to hold functions. Namely tuples.
556556
return ty->isNoEscape();
557557
}
558+
559+
bool isThickFunction() const {
560+
if (auto *fTy = getASTType()->getAs<SILFunctionType>()) {
561+
return fTy->getRepresentation() == SILFunctionType::Representation::Thick;
562+
}
563+
return false;
564+
}
558565

559566
bool isAsyncFunction() const {
560567
if (auto *fTy = getASTType()->getAs<SILFunctionType>()) {

‎include/swift/SILOptimizer/IPO/ClosureSpecializer.h

Lines changed: 34 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,34 @@
1+
//===-------------------------- ClosureSpecializer.h ------------------------------===//
2+
//
3+
// This source file is part of the Swift.org open source project
4+
//
5+
// Copyright (c) 2014 - 2019 Apple Inc. and the Swift project authors
6+
// Licensed under Apache License v2.0 with Runtime Library Exception
7+
//
8+
// See https://swift.org/LICENSE.txt for license information
9+
// See https://swift.org/CONTRIBUTORS.txt for the list of Swift project authors
10+
//
11+
//===-----------------------------------------------------------------------------===//
12+
#ifndef SWIFT_SILOPTIMIZER_CLOSURESPECIALIZER_H
13+
#define SWIFT_SILOPTIMIZER_CLOSURESPECIALIZER_H
14+
15+
#include "swift/SIL/SILFunction.h"
16+
17+
namespace swift {
18+
19+
/// If \p function is a function-signature specialization for a constant-
20+
/// propagated function argument, returns 1.
21+
/// If \p function is a specialization of such a specialization, returns 2.
22+
/// And so on.
23+
int getSpecializationLevel(SILFunction *f);
24+
25+
enum class AutoDiffFunctionComponent : char { JVP = 'f', VJP = 'r' };
26+
27+
/// Returns true if the function is the JVP or the VJP corresponding to
28+
/// a differentiable function.
29+
bool isDifferentiableFuncComponent(
30+
SILFunction *f,
31+
AutoDiffFunctionComponent component = AutoDiffFunctionComponent::VJP);
32+
33+
} // namespace swift
34+
#endif

‎include/swift/SILOptimizer/OptimizerBridging.h

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -239,6 +239,9 @@ struct BridgedPassContext {
239239
SWIFT_IMPORT_UNSAFE BridgedOwnedString mangleWithDeadArgs(const SwiftInt * _Nullable deadArgs,
240240
SwiftInt numDeadArgs,
241241
BridgedFunction function) const;
242+
SWIFT_IMPORT_UNSAFE BridgedOwnedString mangleWithClosureArgs(BridgedValueArray closureArgs,
243+
BridgedArrayRef closureArgIndices,
244+
BridgedFunction applySiteCallee) const;
242245

243246
SWIFT_IMPORT_UNSAFE BridgedGlobalVar createGlobalVariable(BridgedStringRef name, BridgedType type,
244247
bool isPrivate) const;

‎include/swift/SILOptimizer/OptimizerBridgingImpl.h

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -19,11 +19,12 @@
1919
#ifndef SWIFT_SILOPTIMIZER_OPTIMIZERBRIDGING_IMPL_H
2020
#define SWIFT_SILOPTIMIZER_OPTIMIZERBRIDGING_IMPL_H
2121

22-
#include "swift/SILOptimizer/OptimizerBridging.h"
22+
#include "swift/Demangling/Demangle.h"
2323
#include "swift/SILOptimizer/Analysis/AliasAnalysis.h"
2424
#include "swift/SILOptimizer/Analysis/BasicCalleeAnalysis.h"
2525
#include "swift/SILOptimizer/Analysis/DeadEndBlocksAnalysis.h"
2626
#include "swift/SILOptimizer/Analysis/DominanceAnalysis.h"
27+
#include "swift/SILOptimizer/OptimizerBridging.h"
2728
#include "swift/SILOptimizer/PassManager/PassManager.h"
2829
#include "swift/SILOptimizer/Utils/InstOptUtils.h"
2930

‎include/swift/SILOptimizer/PassManager/Passes.def

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -146,6 +146,13 @@ PASS(CapturePropagation, "capture-prop",
146146
"Captured Constant Propagation")
147147
PASS(ClosureSpecializer, "closure-specialize",
148148
"Closure Specialization on Constant Function Arguments")
149+
150+
// NOTE - ExperimentalSwiftBasedClosureSpecialization and AutodiffClosureSpecialization are a WIP
151+
SWIFT_FUNCTION_PASS(ExperimentalSwiftBasedClosureSpecialization, "experimental-swift-based-closure-specialization",
152+
"General closure-specialization pass written in Swift")
153+
SWIFT_FUNCTION_PASS(AutodiffClosureSpecialization, "autodiff-closure-specialization",
154+
"Autodiff specific closure-specialization pass")
155+
149156
PASS(ClosureLifetimeFixup, "closure-lifetime-fixup",
150157
"Closure Lifetime Fixup")
151158
PASS(CodeSinking, "code-sinking",

0 commit comments

Comments
 (0)