[AutoDiff] Closure specialization: specialize branch tracing enums #83926

kovdan01 · 2025-08-26T21:21:16Z

This patch contains part of the changes intended to resolve #68944.

Closure info gathering logic.
Branch tracing enum specialization logic.
Specialization of branch tracing enum basic block arguments in VJP.
Specialization of branch tracing enum payload basic block arguments in pullback.

kovdan01 · 2025-08-26T21:22:58Z

Tagging @JaapWijnen

JaapWijnen · 2025-09-12T13:04:20Z

Is this still meant to be marked as a draft?

kovdan01 · 2025-09-12T21:06:34Z

Is this still meant to be marked as a draft?

@asl Could you please ensure that there are no blocking comments from your side before we un-draft this?

kovdan01 · 2025-09-18T09:35:41Z

@swift-ci please test

eeckstein · 2025-09-22T12:06:31Z

Why is it required to implement this in C++ and not in Swift directly?

kovdan01 · 2025-09-23T16:35:59Z

Why is it required to implement this in C++ and not in Swift directly?

@eeckstein Theoretically, we can move some part of implementation which is now in C++ to Swift, but this would require us to create missing bridges for EnumDecl, EnumElementDecl, ParamDecl, GenericParamList, etc. While we have some bridges for some of these objects, there are many ones which are needed for this patch and which are not present currently. I can implement them just for this patch, but I'm not sure if it's worth it since it would not be used anywhere else (at least at this point). And the amount of new bridges that we need seems non-trivial, so I decided just to use C++ for all this logic.

Also note that as for things related to demangling, I do not see any existing bridges which would allow us to move logic which relies on demangling to Swift. That's not a very big chunk of the patch though.

And one more thing: the "core" function of the patch is autodiffSpecializeBranchTracingEnum which is heavily inspired by LinearMapInfo::createBranchingTraceDecl. Since the latter one is implemented in C++, I suppose that it makes sense to have the new function implemented in C++ as well.

To sum up: I might try to move some of the logic to Swift, but I'll need to define many missing bridges if we decide to do so. And anyway, some logic would remain implemented in C++.

@eeckstein Do you want me to try to do such a shift from C++ to Swift or we can keep things "as is" for now?

eeckstein · 2025-09-24T12:38:49Z

It is the well known conflict of "do something quickly" and "do it right". At some day we want all this code to be in Swift, anyway. So, implementing it in Swift now (including the required briding code) is not wasted time but investment in a better infrastructure for the future.
However it is your call and I can only give you my recommendation.

kovdan01 · 2025-09-24T13:00:33Z

It is the well known conflict of "do something quickly" and "do it right". At some day we want all this code to be in Swift, anyway. So, implementing it in Swift now (including the required briding code) is not wasted time but investment in a better infrastructure for the future.
However it is your call and I can only give you my recommendation.

@eeckstein OK, thanks for a detailed explanation! This makes perfect sense, I was just not sure if implementing all these missing bridges would be worth it. Since the answer is yes, I'm fine taking this approach. Some logic will probably remain implemented in C++ as for now (especially the demangling-related stuff), but I'll try to move as much C++ logic to Swift as I can.

I'll probably submit bridging implementation as a separate PR to avoid mixing the autodiff-specific and generic changes, which should make reviews faster and git history clearer.

In swiftlang#83926, part of the changes resolving swiftlang#68944 is submitted. The AutoDiff closure specialization optimizer pass relies on several AST-level bridges not implemented yet. This patch introduces these missing bridges.

eeckstein · 2025-10-06T08:26:58Z

@kovdan01 Sorry that my PR #84704 interferes with your work here. It's a big simplification of the pass and it should also make your work easier. If you have some questions about the new pass design, just let me know.

Also, can you add a simple SIL example in a comment describing what this PR is about?

In swiftlang#83926, part of the changes resolving swiftlang#68944 is submitted. The AutoDiff closure specialization optimizer pass relies on several AST-level bridges not implemented yet. This patch introduces these missing bridges.

In #83926, part of the changes resolving #68944 is submitted. The AutoDiff closure specialization optimizer pass relies on several bridges not implemented yet. This patch introduces these missing bridges. See the list of the changes below. **AST:** * Define `.asValueDecl` on each BridgedXXXDecl type that's also a ValueDecl. * Define `.asNominalTypeDecl` on each BridgedXXXDecl type that's also a NominalTypeDecl. * Define `.asGenericContext` on each BridgedXXXDecl type that's also a GenericContext. * `class BridgedSourceFile`: - Make nullable - `func addTopLevelDecl(_ decl: BridgedDecl)` * `class BridgedFileUnit`: - `func castToSourceFile() -> BridgedNullableSourceFile` * `class BridgedDecl`: - `func getDeclContext() -> BridgedDeclContext` * `class BridgedParamDecl`: - `func cloneWithoutType() -> BridgedParamDecl` - `func setInterfaceType(_ type : BridgedASTType)` * Define `BridgedDecl.setImplicit()` instead of `BridgedParamDecl.setImplicit()` * `class BridgedGenericContext`: - `setGenericSignature(_ genSig: BridgedGenericSignature)` * Change return type of `BridgedEnumDecl.createParsed(/*...*/)` from `BridgedNominalTypeDecl` to `BridgedEnumDecl` * `class BridgedValueDecl`: - `func setAccess(_ accessLevel: swift.AccessLevel)` * `class BridgedNominalTypeDecl`: - `func addMember(_ member: BridgedDecl)` * `class BridgedGenericTypeParamDecl`: - `func createImplicit(declContext: BridgedDeclContext, name: swift.Identifier, depth: UInt, index: UInt, paramKind: swift.GenericTypeParamKind)` * `class ValueDecl`: - `var baseIdentifier: swift.Identifier` * `class NominalTypeDecl`: - `var declaredInterfaceType: Type` * `class EnumElementDecl`: - `var parameterList: BridgedParameterList` - `var nameStr: StringRef` * `struct GenericSignature`: - `var canonicalSignature: CanGenericSignature` * `struct CanGenericSignature`: - `var isEmpty: Bool` - `var genericSignature: GenericSignature` * `struct Type`: - `func mapTypeOutOfContext() -> Type` - `func getReducedType(sig: GenericSignature) -> CanonicalType` - `func GenericTypeParam_getName() -> swift.Identifier` - `func GenericTypeParam_getDepth() -> UInt` - `func GenericTypeParam_getIndex() -> UInt` - `func GenericTypeParam_getParamKind() -> swift.GenericTypeParamKind` * `struct CanonicalType`: - `func SILFunctionType_getSubstGenericSignature() -> CanGenericSignature` - `func loweredType(in function: SIL.Function) -> SIL.Type` **SIL:** * `class Argument`: - `func replaceAllUsesWith(newArg: Argument)` * `class BasicBlock`: - `func insertPhiArgument(atPosition: Int, type: Type, ownership: Ownership, _ context: some MutatingContext) -> Argument` * `struct Builder`: - `func createTuple(elements: [Value]) -> TupleInst` * `protocol Context`: - `func getTupleType(elements: [AST.Type]) -> AST.Type` - `func getTupleTypeWithLabels(elements: [AST.Type], labels: [swift.Identifier]) -> AST.Type` * `class Function`: - `var sourceFile: BridgedNullableSourceFile` - `func mapTypeIntoContext(_ type: Type) -> Type` * `class PartialApplyInst`: - `var substitutionMap: SubstitutionMap` * `class SwitchEnumInst`: - `var numCases: Int` - `public func getSuccessorForDefault() -> BasicBlock?` * `Type`: - `var category: ValueCategory` - `func getEnumCasePayload(caseIdx: Int, function: Function) -> Type` - `func mapTypeOutOfContext() -> Type` - `static func getPrimitiveType(canType: CanonicalType, silValueCategory: ValueCategory) -> Type` * `struct EnumCase`: - `let enumElementDecl: EnumElementDecl` * `struct TupleElementArray`: - `func label(at index: Int) -> swift.Identifier` * Define `enum ValueCategory` with `address` and `object` elements

This patch contains part of the changes intended to resolve swiftlang#68944. 1. Closure info gathering logic. 2. Branch tracing enum specialization logic. 3. Specialization of branch tracing enum basic block arguments in VJP. 4. Specialization of branch tracing enum payload basic block arguments in pullback.

kovdan01 · 2025-10-06T15:09:26Z

@kovdan01 Sorry that my PR #84704 interferes with your work here. It's a big simplification of the pass and it should also make your work easier. If you have some questions about the new pass design, just let me know.

@eeckstein Thanks! I'll take a look at the patch within a couple of days and let you know if there are any questions/suggestions from my side.

Given that I'll need to do a non-trivial rebase of this PR after your patch gets merged, I'll leave this PR as draft at the moment.

Also, can you add a simple SIL example in a comment describing what this PR is about?

Sure, I'll do that, thanks for suggestion! I just need to minimize an example to avoid cluttering the PR description with giant chunks of code :)

In swiftlang#83926, part of the changes resolving swiftlang#68944 is submitted. The AutoDiff closure specialization optimizer pass relies on several bridges not implemented yet. This patch introduces these missing bridges. See the list of the changes below. **AST:** * Define `.asValueDecl` on each BridgedXXXDecl type that's also a ValueDecl. * Define `.asNominalTypeDecl` on each BridgedXXXDecl type that's also a NominalTypeDecl. * Define `.asGenericContext` on each BridgedXXXDecl type that's also a GenericContext. * `class BridgedSourceFile`: - Make nullable - `func addTopLevelDecl(_ decl: BridgedDecl)` * `class BridgedFileUnit`: - `func castToSourceFile() -> BridgedNullableSourceFile` * `class BridgedDecl`: - `func getDeclContext() -> BridgedDeclContext` * `class BridgedParamDecl`: - `func cloneWithoutType() -> BridgedParamDecl` - `func setInterfaceType(_ type : BridgedASTType)` * Define `BridgedDecl.setImplicit()` instead of `BridgedParamDecl.setImplicit()` * `class BridgedGenericContext`: - `setGenericSignature(_ genSig: BridgedGenericSignature)` * Change return type of `BridgedEnumDecl.createParsed(/*...*/)` from `BridgedNominalTypeDecl` to `BridgedEnumDecl` * `class BridgedValueDecl`: - `func setAccess(_ accessLevel: swift.AccessLevel)` * `class BridgedNominalTypeDecl`: - `func addMember(_ member: BridgedDecl)` * `class BridgedGenericTypeParamDecl`: - `func createImplicit(declContext: BridgedDeclContext, name: swift.Identifier, depth: UInt, index: UInt, paramKind: swift.GenericTypeParamKind)` * `class ValueDecl`: - `var baseIdentifier: swift.Identifier` * `class NominalTypeDecl`: - `var declaredInterfaceType: Type` * `class EnumElementDecl`: - `var parameterList: BridgedParameterList` - `var nameStr: StringRef` * `struct GenericSignature`: - `var canonicalSignature: CanGenericSignature` * `struct CanGenericSignature`: - `var isEmpty: Bool` - `var genericSignature: GenericSignature` * `struct Type`: - `func mapTypeOutOfContext() -> Type` - `func getReducedType(sig: GenericSignature) -> CanonicalType` - `func GenericTypeParam_getName() -> swift.Identifier` - `func GenericTypeParam_getDepth() -> UInt` - `func GenericTypeParam_getIndex() -> UInt` - `func GenericTypeParam_getParamKind() -> swift.GenericTypeParamKind` * `struct CanonicalType`: - `func SILFunctionType_getSubstGenericSignature() -> CanGenericSignature` - `func loweredType(in function: SIL.Function) -> SIL.Type` **SIL:** * `class Argument`: - `func replaceAllUsesWith(newArg: Argument)` * `class BasicBlock`: - `func insertPhiArgument(atPosition: Int, type: Type, ownership: Ownership, _ context: some MutatingContext) -> Argument` * `struct Builder`: - `func createTuple(elements: [Value]) -> TupleInst` * `protocol Context`: - `func getTupleType(elements: [AST.Type]) -> AST.Type` - `func getTupleTypeWithLabels(elements: [AST.Type], labels: [swift.Identifier]) -> AST.Type` * `class Function`: - `var sourceFile: BridgedNullableSourceFile` - `func mapTypeIntoContext(_ type: Type) -> Type` * `class PartialApplyInst`: - `var substitutionMap: SubstitutionMap` * `class SwitchEnumInst`: - `var numCases: Int` - `public func getSuccessorForDefault() -> BasicBlock?` * `Type`: - `var category: ValueCategory` - `func getEnumCasePayload(caseIdx: Int, function: Function) -> Type` - `func mapTypeOutOfContext() -> Type` - `static func getPrimitiveType(canType: CanonicalType, silValueCategory: ValueCategory) -> Type` * `struct EnumCase`: - `let enumElementDecl: EnumElementDecl` * `struct TupleElementArray`: - `func label(at index: Int) -> swift.Identifier` * Define `enum ValueCategory` with `address` and `object` elements

kovdan01 requested a review from asl August 26, 2025 21:22

kovdan01 deleted the branch swiftlang:users/kovdan01/ast-bridges-for-autodiff-closure-spec August 27, 2025 06:12

kovdan01 closed this Aug 27, 2025

kovdan01 reopened this Aug 27, 2025

kovdan01 changed the base branch from users/kovdan01/autodiff-closure-specialization-auto-format to main August 27, 2025 06:14

kovdan01 force-pushed the autodiff-closure-specialization-multi-bb-specialize-branch-tracing-enum branch 2 times, most recently from 37f6ddb to 77191ca Compare August 27, 2025 22:30

kovdan01 mentioned this pull request Aug 28, 2025

[AutoDiff] Implement the closure optimization that is specialized towards the linear map tuples / enums produced by autodiff. #68944

Open

kovdan01 force-pushed the autodiff-closure-specialization-multi-bb-specialize-branch-tracing-enum branch 2 times, most recently from 5ba5d12 to f568182 Compare September 18, 2025 09:34

kovdan01 marked this pull request as ready for review September 18, 2025 09:35

kovdan01 requested review from eeckstein and jckarter as code owners September 18, 2025 09:35

kovdan01 marked this pull request as draft October 2, 2025 14:26

kovdan01 mentioned this pull request Oct 2, 2025

Bridging: Implement bridges required for ongoing AutoDiff changes #84648

Draft

kovdan01 force-pushed the autodiff-closure-specialization-multi-bb-specialize-branch-tracing-enum branch from f568182 to cf717ce Compare October 2, 2025 14:47

kovdan01 changed the base branch from main to users/kovdan01/ast-bridges-for-autodiff-closure-spec October 2, 2025 14:47

kovdan01 force-pushed the users/kovdan01/ast-bridges-for-autodiff-closure-spec branch from 10085d6 to b3d31c7 Compare October 6, 2025 14:58

kovdan01 force-pushed the autodiff-closure-specialization-multi-bb-specialize-branch-tracing-enum branch from cf717ce to fe18b09 Compare October 6, 2025 15:05

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

[AutoDiff] Closure specialization: specialize branch tracing enums #83926

[AutoDiff] Closure specialization: specialize branch tracing enums #83926

kovdan01 commented Aug 26, 2025 •

edited

Loading

Uh oh!

kovdan01 commented Aug 26, 2025

Uh oh!

JaapWijnen commented Sep 12, 2025

Uh oh!

kovdan01 commented Sep 12, 2025

Uh oh!

kovdan01 commented Sep 18, 2025

Uh oh!

eeckstein commented Sep 22, 2025

Uh oh!

kovdan01 commented Sep 23, 2025

Uh oh!

eeckstein commented Sep 24, 2025

Uh oh!

kovdan01 commented Sep 24, 2025

Uh oh!

eeckstein commented Oct 6, 2025

Uh oh!

kovdan01 commented Oct 6, 2025

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

[AutoDiff] Closure specialization: specialize branch tracing enums #83926

Are you sure you want to change the base?

[AutoDiff] Closure specialization: specialize branch tracing enums #83926

Conversation

kovdan01 commented Aug 26, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

kovdan01 commented Aug 26, 2025

Uh oh!

JaapWijnen commented Sep 12, 2025

Uh oh!

kovdan01 commented Sep 12, 2025

Uh oh!

kovdan01 commented Sep 18, 2025

Uh oh!

eeckstein commented Sep 22, 2025

Uh oh!

kovdan01 commented Sep 23, 2025

Uh oh!

eeckstein commented Sep 24, 2025

Uh oh!

kovdan01 commented Sep 24, 2025

Uh oh!

eeckstein commented Oct 6, 2025

Uh oh!

kovdan01 commented Oct 6, 2025

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

kovdan01 commented Aug 26, 2025 •

edited

Loading