[AutoDiff] Add SIL differentiability witnesses. #27487

dan-zheng · 2019-10-02T07:47:37Z

SIL differentiability witnesses are a new top-level SIL construct mapping "original" SIL functions to derivatives. They will replace SIL function [differentiable] attributes, additionally enabling cross-module retroactive derivative registration.

SIL differentiability witnesses have the following components:

Original SILFunction.
Linkage.
Parameter indices (IndexSubset).
Result indices (IndexSubset).
Derivative generic signature (optional).
JVP function (optional).
VJP function (optional).

Example syntax:

sil_differentiability_witness hidden [parameters 0 1] [results 0] [where τ_0_0 : _Differentiable] @foo : $@convention(thin) <τ_0_0> (@in_guaranteed τ_0_0, Float) -> @out τ_0_0 {
  jvp: @AD__foo__jvp_src_0_wrt_0 : $@convention(thin) <τ_0_0 where τ_0_0 : _Differentiable> (@in_guaranteed τ_0_0, Float) -> (@out τ_0_0, @owned @callee_guaranteed (@in_guaranteed τ_0_0.TangentVector, Float) -> @out τ_0_0.TangentVector)
  vjp: @AD__foo__vjp_src_0_wrt_0 : $@convention(thin) <τ_0_0 where τ_0_0 : _Differentiable> (@in_guaranteed τ_0_0, Float) -> (@out τ_0_0, @owned @callee_guaranteed (@in_guaranteed τ_0_0.TangentVector) -> (@out τ_0_0.TangentVector, Float))
}

This patch adds the SILDifferentiabilityWitness data structure, along with parsing, printing, verification, and serialization (including lookup by key).

The master issue TF-866 tracks follow-up, including SILGen/IRGen/differentiation transform changes.

Early forum discussion.

- Printing compiles but may require changes to be parseable. Namely, if there is no suitable utility for parsing standalone generic signatures, changes are needed. - Parsing is a stub. - Note: it is difficult to test parsing/printing without generating SILDifferentiabilityWitness instances.

`parameters (0, 1, ...) results (0, 1, ...) where <...>`

…tiability-witness

Use mangling to support string key lookup.

…tiability-witness

rxwei

Because type signatures are long and sparse, I'd suggest bringing the configuration to the front of the declaration and using square brackets for them, for example:

sil_differentiability_witness hidden [parameters 0 1] [results 0] [where τ_0_0 : _Differentiable] @foo : $@convention(thin) <τ_0_0> (@in_guaranteed τ_0_0, Float) -> @out τ_0_0 {
  jvp: @AD__foo__jvp_src_0_wrt_0 : $@convention(thin) <τ_0_0 where τ_0_0 : _Differentiable> (@in_guaranteed τ_0_0, Float) -> (@out τ_0_0, @owned @callee_guaranteed (@in_guaranteed τ_0_0.TangentVector, Float) -> @out τ_0_0.TangentVector)
  vjp: @AD__foo__vjp_src_0_wrt_0 : $@convention(thin) <τ_0_0 where τ_0_0 : _Differentiable> (@in_guaranteed τ_0_0, Float) -> (@out τ_0_0, @owned @callee_guaranteed (@in_guaranteed τ_0_0.TangentVector) -> (@out τ_0_0.TangentVector, Float))
}

Parts of what SILDifferentiabilityWitnessKey stores could have been made a more general concept: AutoDiffConfig. How about defining a AutoDiffConfig in AutoDiff.h to store parameter indices, result indices, and the derivative generic signature? Then SILDifferentiabilityWitnessKey can just be defined as std::pair<StringRef, AutoDiffConfig>.

When the function name of a sil_differentiability_witness is demangleable, consider printing a descriptive comment above the witness declaration, for example:

// differentiability witness for closure #1 in foo<A>(_:_:) 
sil_differentiability_witness hidden [parameters 0 1] [results 0] [where τ_0_0 : _Differentiable] @someDemangleableName : $@convention(thin) <τ_0_0> (@in_guaranteed τ_0_0, Float) -> @out τ_0_0 {

include/swift/SIL/SILDifferentiabilityWitness.h

@foo

Print original function name in comment. ``` // differentiability witness for foo sil_differentiability_witness hidden [parameters 0 1] [results 0] [where τ_0_0 : _Differentiable] @foo ```

include/swift/AST/AutoDiff.h

The contents of `AutoDiffConfig` are all uniqued, so uniquing the product does not make sense.

Unserialized, to be used for diagnostics. Will revisit later when revamping the differentiation transform.

include/swift/AST/AutoDiff.h

include/swift/AST/DiagnosticsParse.def

lib/ParseSIL/ParseSIL.cpp

include/swift/AST/DiagnosticsParse.def

include/swift/SIL/SILDifferentiabilityWitness.h

lib/SIL/SILPrinter.cpp

lib/Serialization/DeserializeSIL.cpp

lib/Serialization/SerializeSIL.cpp

Manually verified parsing/printing. Chose not to add additional test for now to keep the test small.

rxwei

Are there serialization tests?

lib/SIL/SILPrinter.cpp

test/AutoDiff/sil_differentiability_witness_parse.sil

rxwei

The implementation LGTM, but I just wanted to make sure serialization is tested.

lib/ParseSIL/ParseSIL.cpp

test/AutoDiff/sil_differentiability_witness_parse.sil

Note: deserialization does not work when SIL differentiability witness references bodyless function declarations.

dan-zheng · 2019-10-13T09:01:03Z

@swift-ci Please test tensorflow

lib/SIL/SILVerifier.cpp

Verification assertion messages should appear on the differentiability witness.

dan-zheng · 2019-10-13T10:09:06Z

@swift-ci Please test tensorflow

dan-zheng added the tensorflow This is for "tensorflow" branch PRs. label Oct 2, 2019

dan-zheng requested a review from rxwei October 2, 2019 07:47

rxwei and others added 6 commits October 3, 2019 14:40

Diff witness

7e5d804

Rename SILDifferentiabilityWitness.h.

a5c1d13

Update SILDifferentiabilityWitness definition.

0ce7d91

Add SILDifferentiabilityWitness to SILModule.

471e660

Use improved syntax.

6b78684

`parameters (0, 1, ...) results (0, 1, ...) where <...>`

dan-zheng force-pushed the sil-differentiability-witness branch from 57e9634 to bac0181 Compare October 3, 2019 21:42

dan-zheng added 4 commits October 10, 2019 12:52

Merge branch 'tensorflow' of github.com:apple/swift into sil-differen…

ab304f1

…tiability-witness

Finish parsing/printing/serialization.

a2ae0f2

Revamp serialization to enable lookup by key.

419eea2

Use mangling to support string key lookup.

Add SIL verification.

b6cd1d7

dan-zheng force-pushed the sil-differentiability-witness branch from bac0181 to d2c6ab5 Compare October 11, 2019 05:07

dan-zheng added 2 commits October 10, 2019 22:10

Add miscellaneous todo comments.

573dd3e

Add round-trip parsing/printing test.

843b631

dan-zheng force-pushed the sil-differentiability-witness branch from d2c6ab5 to 843b631 Compare October 11, 2019 05:11

dan-zheng changed the title ~~[WIP] [AutoDiff] Add SIL differentiability witnesses.~~ [AutoDiff] Add SIL differentiability witnesses. Oct 11, 2019

dan-zheng marked this pull request as ready for review October 11, 2019 05:16

dan-zheng added 4 commits October 10, 2019 22:20

Clean up.

835f1c0

Add Swift source for parsing test.

aca1abb

Merge branch 'tensorflow' of github.com:apple/swift into sil-differen…

d75f6ce

…tiability-witness

AutoDiffIndexSubset -> IndexSubset

4187700

rxwei reviewed Oct 11, 2019

View reviewed changes

Minor fix.

873468f

rxwei requested a review from jckarter October 11, 2019 20:44

rxwei reviewed Oct 12, 2019

View reviewed changes

include/swift/SIL/SILDifferentiabilityWitness.h Show resolved Hide resolved

dan-zheng added 2 commits October 12, 2019 23:58

Update differentiability witness syntax.

aea64d3

Print original function name in comment. ``` // differentiability witness for foo sil_differentiability_witness hidden [parameters 0 1] [results 0] [where τ_0_0 : _Differentiable] @foo ```

Add AutoDiffConfig and use in SILDifferentiabilityWitnessKey.

bb93af3

rxwei reviewed Oct 13, 2019

View reviewed changes

include/swift/AST/AutoDiff.h Outdated Show resolved Hide resolved

dan-zheng added 2 commits October 13, 2019 01:16

Change AutoDiffConfig to a POD.

f240ed2

The contents of `AutoDiffConfig` are all uniqued, so uniquing the product does not make sense.

Add DeclAttribute * to SILDifferentiabilityWitness.

ad6b7aa

Unserialized, to be used for diagnostics. Will revisit later when revamping the differentiation transform.

dan-zheng force-pushed the sil-differentiability-witness branch from 4ddcd3f to ad6b7aa Compare October 13, 2019 01:16

Clean up.

7c63d03

dan-zheng requested a review from rxwei October 13, 2019 01:19

rxwei reviewed Oct 13, 2019

View reviewed changes

Address review feedback.

c3959ad

dan-zheng force-pushed the sil-differentiability-witness branch from 771143a to c3959ad Compare October 13, 2019 05:12

rxwei reviewed Oct 13, 2019

View reviewed changes

dan-zheng added 2 commits October 13, 2019 06:24

Address review feedback.

2073f86

Parse/print [serialized] flag.

27d7abc

Manually verified parsing/printing. Chose not to add additional test for now to keep the test small.

rxwei reviewed Oct 13, 2019

View reviewed changes

lib/SIL/SILPrinter.cpp Outdated Show resolved Hide resolved

test/AutoDiff/sil_differentiability_witness_parse.sil Outdated Show resolved Hide resolved

rxwei approved these changes Oct 13, 2019

View reviewed changes

rxwei reviewed Oct 13, 2019

View reviewed changes

lib/ParseSIL/ParseSIL.cpp Show resolved Hide resolved

dan-zheng force-pushed the sil-differentiability-witness branch from ef5f53c to 7a070f0 Compare October 13, 2019 07:35

rxwei reviewed Oct 13, 2019

View reviewed changes

test/AutoDiff/sil_differentiability_witness_parse.sil Outdated Show resolved Hide resolved

Add parsing/printing tests, address review feedback.

69209be

dan-zheng force-pushed the sil-differentiability-witness branch from 7a070f0 to 69209be Compare October 13, 2019 07:42

Fix serialization and add test.

df9cd49

Note: deserialization does not work when SIL differentiability witness references bodyless function declarations.

dan-zheng commented Oct 13, 2019

View reviewed changes

lib/SIL/SILVerifier.cpp Outdated Show resolved Hide resolved

Fix verification.

d673b70

Verification assertion messages should appear on the differentiability witness.

dan-zheng merged commit a5dc918 into swiftlang:tensorflow Oct 13, 2019

dan-zheng deleted the sil-differentiability-witness branch October 13, 2019 11:30

rxwei mentioned this pull request Oct 16, 2019

[AutoDiff upstream] Add the _Differentiation module. #27511

Merged

dan-zheng mentioned this pull request Oct 15, 2019

[SR-13660] SIL: differentiability witness data structure, parsing and serialization #53960

Closed

[AutoDiff] Add SIL differentiability witnesses. #27487

[AutoDiff] Add SIL differentiability witnesses. #27487

Uh oh!

Conversation

dan-zheng commented Oct 2, 2019 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

rxwei left a comment • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

rxwei left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

rxwei left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

dan-zheng commented Oct 13, 2019

Uh oh!

Uh oh!

dan-zheng commented Oct 13, 2019

Uh oh!

Uh oh!

dan-zheng commented Oct 2, 2019 •

edited

Loading

rxwei left a comment •

edited

Loading