-
Notifications
You must be signed in to change notification settings - Fork 10.5k
[AutoDiff] Add 'SILDifferentiableFunctionType'. #23482
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.
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
Closed
Closed
Conversation
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
marcrasi
reviewed
Mar 22, 2019
rxwei
added a commit
to rxwei/swift
that referenced
this pull request
Apr 8, 2019
…on types. This is to unblock `@differentiable` functions with `@nondiff` parameters. Once `SILDifferentiableFunctionType` in swiftlang#23482 lands, `@nondiff` should be nuked from SIL.
rxwei
added a commit
that referenced
this pull request
Apr 11, 2019
…on types. (#23854) This is to unblock `@differentiable` functions with `@nondiff` parameters. - Propagate `@nondiff` from AST to SIL. - Add `AnyFunctionType::getWithoutDifferentiability`, which drops all `@differentiable` and `@nondiff` attributes from a function type. - Use autodiff parameter indices from function types now that `@nondiff` has been propagated. - Replace currying logic from `SILFunctionType::getAssociatedFunctionType` with lightweight logic that handles methods, which is needed for differentiable protocol requirements. - Emit an error when a `@nondiff` parameter is being differentiated with respect to. - Fix `@nondiff` AST printing. Note: - Once `SILDifferentiableFunctionType` in #23482 lands, `@nondiff` should be nuked from SIL. - Before merging, pull from the `tensorflow` branch to make sure #23887 is merged. Resolves [TF-421](https://bugs.swift.org/browse/TF-421).
rxwei
added a commit
to rxwei/swift
that referenced
this pull request
May 11, 2019
…on types. (swiftlang#23854) This is to unblock `@differentiable` functions with `@nondiff` parameters. - Propagate `@nondiff` from AST to SIL. - Add `AnyFunctionType::getWithoutDifferentiability`, which drops all `@differentiable` and `@nondiff` attributes from a function type. - Use autodiff parameter indices from function types now that `@nondiff` has been propagated. - Replace currying logic from `SILFunctionType::getAssociatedFunctionType` with lightweight logic that handles methods, which is needed for differentiable protocol requirements. - Emit an error when a `@nondiff` parameter is being differentiated with respect to. - Fix `@nondiff` AST printing. Note: - Once `SILDifferentiableFunctionType` in swiftlang#23482 lands, `@nondiff` should be nuked from SIL. - Before merging, pull from the `tensorflow` branch to make sure swiftlang#23887 is merged. Resolves [TF-421](https://bugs.swift.org/browse/TF-421).
rxwei
added a commit
that referenced
this pull request
May 14, 2019
### Introduce `AutoDiffIndexSubset` This PR overhauls data structures used for parameter indices in the AutoDiff infrastructure in SIL. Previously, we used `llvm::SmallBitVector` to represent differentiation parameter indices in both AST and SIL. It was not efficient, and most importantly there's no way to put this in an instruction without causing memory leaks. This change replaces all uses of `llvm::SmallBitVector` in SIL AutoDiff code paths with a `ASTContext`-uniqued `AutoDiffIndexSubset *` where bits are stored as trailing objects. `AutoDiffIndexSubset` does not have "parameter indices" in its name because it is not only designed for parameter indices, but also for result indices as we move to supporting multi-result differentiation. `AutoDiffIndexSubset` has set operations like `isSubsetOf`, `isSupersetOf`, and `contains`, but it also has a special _capacity_ property. All differentiable function's parameter indices data should store the number of parameters as capacity, so that the differentiation transform won't need special logic to check whether an index is out of range. Another minor change is the module format layout of `SILDifferentiableAttr`. It used to store parameter indices as consecutive `bool` bits, but now stores numeric parameter indices directly for efficiency. It will be necessary to refactor or eliminate `AutoDiffParameterIndices` to make use of `AutoDiffIndexSubset`. `AutoDiffParameterIndices` is at the AST level, so it is not in the scope for this PR. Unblocks #23482. Partially resolves [TF-67](https://bugs.swift.org/browse/TF-67).
ca99bd1
to
bca46c7
Compare
e877ecc
to
ae75f5f
Compare
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
WIP
SILDifferentiableBundleType
is a SIL-only first-class type that represents a differentiable function. Previously, we used a@differentiable
attribute on normal functions and made functions with that attribute be lowered to a struct. However, that had a ton of problems (TF-11, TF-17, TF-171) so that we had to turn off LoadableByAddress.The right solution, which @marcrasi initially suggested, is to store the autodiff-associated function types into differentiable SIL function types. More specifically, we want to store the type of all parameters' and results' member types
TangentVector
andCotangentVector
, which are used in differentials and pullbacks, respectively.This PR is doing just that. A
SILDifferentiableFunctionType
contains the following information:2
.maxOrder
must be-1
.maxOrder
must be greater than0
.@convention(thin) (T) -> U
.@callee_guaranteed (T.TangentVector) -> U.TangentVector
.@callee_guaranteed (U.CotangentVector) -> T.CotangentVector
.Putting things together, a Swift order-1 differentiable function
@differentiable (T) -> U
in SIL looks like:A linear function (e.g. a differential or a pullback) looks like: