Skip to content

[ABI] [AutoDiff] Fix function differentiability kind metadata and mangling. #36601

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom
Merged
merged 1 commit into from
Mar 30, 2021
Merged

[ABI] [AutoDiff] Fix function differentiability kind metadata and mangling. #36601

merged 1 commit into from
Mar 30, 2021

Conversation

rxwei
Copy link
Contributor

@rxwei rxwei commented Mar 26, 2021 •
edited
Loading

  • Move differentiability kinds from target function type metadata to trailing objects so that we don't exhaust all remaining bits of function type metadata.
    • Differentiability kind is now stored in a tail-allocated word when function type flags say it's differentiable, located immediately after the normal function type metadata's contents (with proper alignment in between).
    • Add new runtime function swift_getFunctionTypeMetadataDifferentiable which handles differentiable function types.
  • Fix mangling of different differentiability kinds in function types. Mangle it like ConcurrentFunctionType so that we can drop special cases for escaping functions. 
    function-signature ::= params-type params-type async? sendable? throws? differentiable? // results and parameters
...
differentiable ::= 'jf'                    // @differentiable(_forward) on function type
differentiable ::= 'jr'                    // @differentiable(reverse) on function type
differentiable ::= 'jd'                    // @differentiable on function type
differentiable ::= 'jl'                    // @differentiable(_linear) on function type

Resolves rdar://75240064.

@rxwei rxwei force-pushed the 75240064-differentiability-metadata branch 3 times, most recently from 7e6607d to d617f0e Compare March 26, 2021 10:18
@rxwei rxwei requested a review from jckarter March 26, 2021 10:22
jckarter
jckarter reviewed Mar 26, 2021
View reviewed changes
include/swift/ABI/Metadata.h
}
return TargetFunctionMetadataDifferentiabilityKind<StoredSize>
::NonDifferentiable;
}
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

You could use swift/ABI/TrailingObjects.h to help with the trailing allocation management here. Where do you initialize the tail storage?

Copy link
Contributor Author

@rxwei rxwei Mar 26, 2021 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yeah, initially I wanted to use that but noticed that parameters and parameter flags are tail-allocated manually, so I followed the same approach. Would it be odd to mix both, no? The allocation size calculation is in FunctionCacheEntry::getExtraAllocationSize(); I aligned the previously calculated size (because there are a few 32-bit ints for parameter flags) and appended a differentiability word after it.

@rxwei rxwei force-pushed the 75240064-differentiability-metadata branch 4 times, most recently from 65286df to e6fe273 Compare March 27, 2021 08:52
@rxwei rxwei changed the title [WIP] Fix function differentiability kind metadata and mangling. [ABI] [AutoDiff] Fix function differentiability kind metadata and mangling. Mar 27, 2021
rxwei
rxwei commented Mar 27, 2021
View reviewed changes
lib/Demangling/OldDemangler.cpp
@@ -1785,6 +1798,15 @@ class OldDemangler {
if (async) {
block->addChild(Factory.createNode(Node::Kind::AsyncAnnotation), Factory);
}
if (concurrent) {
Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This part seemed missing sendable handling so I added it.

lib/IRGen/MetadataRequest.cpp Outdated
.withNumParameters(numParams)
.withConvention(metadataConvention)
.withAsync(type->isAsync())
.withConcurrent(type->isConcurrent())
Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This part seemed missing sendable handling so I added it.

@rxwei
Copy link
Contributor Author

rxwei commented Mar 27, 2021

@swift-ci please test

@rxwei rxwei marked this pull request as ready for review March 27, 2021 09:01
@swift-ci
Copy link
Contributor

Build failed
Swift Test OS X Platform
Git Sha - e6fe2734e655bd18f7150577bc2ac578f8955ca7

@rxwei
Fix function differentiability kind metadata and mangling.
f5282ed 
* Move differentiability kinds from target function type metadata to trailing objects so that we don't exhaust all remaining bits of function type metadata.
* Fix mangling of different differentiability kinds in function types. Mangle it like `ConcurrentFunctionType` so that we can drop special cases for escaping functions.

Resolves rdar://75240064.
@rxwei rxwei force-pushed the 75240064-differentiability-metadata branch from e6fe273 to f5282ed Compare March 28, 2021 07:38
@swiftlang swiftlang deleted a comment from swift-ci Mar 28, 2021
@swiftlang swiftlang deleted a comment from swift-ci Mar 28, 2021
@swiftlang swiftlang deleted a comment from swift-ci Mar 28, 2021
@rxwei
Copy link
Contributor Author

rxwei commented Mar 28, 2021

@swift-ci please test

jckarter
jckarter approved these changes Mar 30, 2021
View reviewed changes
Copy link
Contributor

@jckarter jckarter left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Looks good, thanks Richard!

@rxwei rxwei merged commit d997526 into swiftlang:main Mar 30, 2021
@rxwei rxwei deleted the 75240064-differentiability-metadata branch March 30, 2021 16:59
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Reviewers

@jckarter jckarter jckarter approved these changes

@dan-zheng dan-zheng dan-zheng approved these changes

@joshnewnham joshnewnham Awaiting requested review from joshnewnham

@DougGregor DougGregor Awaiting requested review from DougGregor

Assignees
No one assigned
Labels
None yet
Projects
None yet
Milestone
No milestone
Development

Successfully merging this pull request may close these issues.
4 participants
@rxwei @swift-ci @jckarter @dan-zheng