Skip to content

[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
wants to merge 3 commits into from

Conversation

rxwei
Copy link
Contributor

@rxwei rxwei commented Mar 22, 2019

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 and CotangentVector, which are used in differentials and pullbacks, respectively.

This PR is doing just that. A SILDifferentiableFunctionType contains the following information:

  • The max differentiation order, e.g. 2.
  • The differentiability representation kind
    • Linear: Represents a linear map. maxOrder must be -1.
    • Normal: Represents a non-linear differentiable function. maxOrder must be greater than 0.
  • The original function type, e.g. @convention(thin) (T) -> U.
  • The differential type, e.g. @callee_guaranteed (T.TangentVector) -> U.TangentVector.
  • The pullback type, e.g. @callee_guaranteed (U.CotangentVector) -> T.CotangentVector.

Putting things together, a Swift order-1 differentiable function @differentiable (T) -> U in SIL looks like:

$@differentiable(1) {@convention(thin) (T) -> U, differential: @callee_guaranteed (T.TangentVector) -> U.TangentVector, pullback: @callee_guaranteed (U.CotangentVector) -> T.CotangentVector}

A linear function (e.g. a differential or a pullback) looks like:

$@differentiable(linear) {@convention(thin) (T.TangentVector) -> U.CotangentVector, transpose: @callee_guaranteed (U.CotangentVector) -> T.CotangentVector}

@rxwei rxwei requested a review from marcrasi March 22, 2019 06:06
@rxwei rxwei added the tensorflow This is for "tensorflow" branch PRs. label Apr 1, 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).
@rxwei rxwei force-pushed the diffable-bundle branch 3 times, most recently from ca99bd1 to bca46c7 Compare May 22, 2019 07:05
@rxwei rxwei force-pushed the diffable-bundle branch 6 times, most recently from e877ecc to ae75f5f Compare May 26, 2019 11:23
@rxwei rxwei force-pushed the diffable-bundle branch from ae75f5f to 871f916 Compare May 26, 2019 11:27
@rxwei rxwei closed this Aug 7, 2019
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
tensorflow This is for "tensorflow" branch PRs.
Projects
None yet
Development

Successfully merging this pull request may close these issues.

2 participants