Skip to content

[AutoDiff upstream] Add SIL differentiability witnesses. #29623

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

Merged

Conversation

dan-zheng
Copy link
Contributor

SIL differentiability witnesses are a new top-level SIL construct mapping
"original" SIL functions to derivative SIL functions.

SIL differentiability witnesses have the following components:

  • "Original" SILFunction.
  • SIL linkage.
  • Differentiability parameter indices (IndexSubset).
  • Differentiability result indices (IndexSubset).
  • Derivative GenericSignature representing differentiability generic requirements (optional).
  • JVP derivative SILFunction (optional).
  • VJP derivative SILFunction (optional).
  • "Is serialized?" bit.
sil_differentiability_witness hidden [parameters 0] [results 0] <T where T : Differentiable> @id : $@convention(thin) (T) -> T {
  jvp: @id_jvp : $@convention(thin) (T) -> (T, @owned @callee_guaranteed (T.TangentVector) -> T.TangentVector)
  vjp: @id_vjp : $@convention(thin) (T) -> (T, @owned @callee_guaranteed (T.TangentVector) -> T.TangentVector)
}

This patch adds the SILDifferentiabilityWitness data structure, with
documentation, parsing, and printing.

Resolves TF-911.


Todos:

  • TF-1136: upstream SILDifferentiabilityWitness serialization.
  • TF-1137: upstream SILDifferentiabilityWitness verification.
  • TF-1138: upstream SILDifferentiabilityWitness SILGen from
    @differentiable and @derivative attributes.
  • TF-20: robust mangling for SILDifferentiabilityWitness names.

SIL differentiability witnesses are a new top-level SIL construct mapping
"original" SIL functions to derivative SIL functions.

SIL differentiability witnesses have the following components:
- "Original" `SILFunction`.
- SIL linkage.
- Differentiability parameter indices (`IndexSubset`).
- Differentiability result indices (`IndexSubset`).
- Derivative `GenericSignature` representing differentiability generic
  requirements (optional).
- JVP derivative `SILFunction` (optional).
- VJP derivative `SILFunction` (optional).
- "Is serialized?" bit.

This patch adds the `SILDifferentiabilityWitness` data structure, with
documentation, parsing, and printing.

Resolves TF-911.

Todos:
- TF-1136: upstream `SILDifferentiabilityWitness` serialization.
- TF-1137: upstream `SILDifferentiabilityWitness` verification.
- TF-1138: upstream `SILDifferentiabilityWitness` SILGen from
  `@differentiable` and `@derivative` attributes.
- TF-20: robust mangling for `SILDifferentiabilityWitness` names.
@dan-zheng dan-zheng requested review from rxwei and marcrasi February 4, 2020 01:07
@@ -2787,6 +2818,8 @@ void SILModule::print(SILPrintContext &PrintCtx, ModuleDecl *M,
}

printSILGlobals(PrintCtx, getSILGlobalList());
printSILDifferentiabilityWitnesses(PrintCtx,
Copy link
Contributor Author

Choose a reason for hiding this comment

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

Note: printing SIL differentiability witnesses before SIL functions may be slightly controversial.

This decision significantly simplifies SIL parsing: it allows us to avoid adding support for diff. witness forward declarations when parsing differentiability_witness_function instructions in SIL functions.

Other longer table-like SIL constructs are printed after SIL functions so that functions appear at the top. Diff. witnesses aren't as long as other table-like SIL constructs (up to four lines), so they don't take up too much space.

@dan-zheng dan-zheng requested a review from jckarter February 4, 2020 01:47
@dan-zheng
Copy link
Contributor Author

dan-zheng commented Feb 4, 2020

@jckarter: would you like to review this PR, or suggest a reviewer? I'm asking you since you seem like a SIL code owner.

This PR is part of our effort to upstream differentiable programming. It adds a new top-level SIL construct called "differentiability witnesses" and doesn't affect much existing code.

If there's no activity, I'll merge within a few days to unblock progress. Happy to address feedback at any time!

@dan-zheng
Copy link
Contributor Author

@swift-ci Please smoke test

@dan-zheng
Copy link
Contributor Author

@swift-ci Please test

@swift-ci
Copy link
Contributor

swift-ci commented Feb 4, 2020

Build failed
Swift Test OS X Platform
Git Sha - 9f7ca52

@jckarter
Copy link
Contributor

jckarter commented Feb 4, 2020

Looks good to me. If you're just adding data structures to support autodiff that have minimal interaction with the rest of the compiler like this, I would say don't wait for me.

@dan-zheng
Copy link
Contributor Author

Looks good to me. If you're just adding data structures to support autodiff that have minimal interaction with the rest of the compiler like this, I would say don't wait for me.

Got it, thank you!

If you don't mind, I'll continue to add you as a reviewer on SIL autodiff upstreaming patches, unless you have some other preference.

@dan-zheng
Copy link
Contributor Author

The macOS Python/python-lint test failure seems unrelated to this patch, so I'll merge now.

@dan-zheng dan-zheng merged commit 849bd62 into swiftlang:master Feb 4, 2020
@dan-zheng dan-zheng deleted the autodiff-upstream-sil-diff-witness branch February 4, 2020 20:53
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

4 participants