Skip to content
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

[AutoDiff] Fix @differentiable attribute SILGen and serialization. #21837

Merged
merged 12 commits into from Jan 16, 2019
Merged

[AutoDiff] Fix @differentiable attribute SILGen and serialization. #21837

merged 12 commits into from Jan 16, 2019

Conversation

dan-zheng
Copy link
Collaborator

@dan-zheng dan-zheng commented Jan 14, 2019 •
edited

  • Propagate @differentiable attributes in SILFunctionBuilder::addFunctionAttributes.
    • Previously, this logic was missing, causing SIL [differentiable] attributes
      to be missing during SILGen.
    • This unblocks using VJP definitions in the stdlib (with some extra work).
  • Serialize AutoDiffParameterIndices for DifferentiableAttr and use it as the
    primary source for differentiation parameter info (instead of
    ParsedAutoDiffParameter).
  • Rework differentiation to never load SIL functions.
    • Previously, differentiation relied on explicit loading to get SIL
      [differentiable] attributes. That was a hack.
    • Now that @differentiable attribute is propagated and serialized,
      such loading is no longer necessary.
  • Tighten differentiation infrastructure to uphold the following invariants:
    • Differentiable functions' primal/adjoint are visible only in defining module.
    • Differentiable functions' JVP/VJPs are always visible to other modules.
      The differentiation pass guarantees that JVP/VJPs exist.

There are two TensorFlowRuntime test regressions in GPE mode.

  • Both are because used device set includes RUNTIME device in
    GraphFunctionDeviceInfo::finalizeUsedDevices().

Todos:

  • Replace adjoint definitions in stdlib with VJP definitions.
  • Eliminate primal/adjoint from @differentiable attribute.
    Use JVP/VJPs everywhere.

@dan-zheng
[AutoDiff] Fix SILGen for @differentiable attribute.
99fc5e7 
Propagate `@differentiable` attributes in `SILFunctionBuilder::addFunctionAttributes`.
- Previously, this logic was missing, causing SIL `[differentiable]` attributes
  to be missing during SILGen.
- This unblocks using VJP definitions in the stdlib (with some extra work).

Serialize `AutoDiffParameterIndices` for `DifferentiableAttr` and use it as the
primary source for differentiation parameter info (instead of
`ParsedAutoDiffParameter`).

Fix `AutoDiff` tests.

There are two `TensorFlowRuntime` test regressions in GPE mode.
- Both are because used device set includes RUNTIME device in
  `GraphFunctionDeviceInfo::finalizeUsedDevices()`.
@dan-zheng dan-zheng added the tensorflow This is for "tensorflow" branch PRs. label Jan 14, 2019
@dan-zheng dan-zheng requested a review from rxwei January 14, 2019 03:39
rxwei
rxwei approved these changes Jan 14, 2019
View reviewed changes
include/swift/AST/AutoDiff.h Outdated Show resolved Hide resolved
lib/SILOptimizer/Mandatory/TFDifferentiation.cpp Outdated Show resolved Hide resolved
lib/Sema/TypeCheckAttr.cpp Outdated Show resolved Hide resolved
lib/Serialization/Deserialization.cpp Outdated Show resolved Hide resolved
test/AutoDiff/derived_differentiable_properties.swift
@@ -7,7 +7,7 @@ public struct Foo : Differentiable {
}

// CHECK-AST-LABEL: public struct Foo : Differentiable {
// CHECK-AST: @sil_stored @differentiable()
// CHECK-AST: @sil_stored @differentiable(wrt: (self))
Copy link
Member

Choose a reason for hiding this comment

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

I love that things have a canonical printout now when you use checked parameter indices.

@dan-zheng dan-zheng force-pushed the fix-differentiable-attr-silgen branch 2 times, most recently from c17b517 to c765cc6 Compare January 15, 2019 01:42
@dan-zheng dan-zheng force-pushed the fix-differentiable-attr-silgen branch from c765cc6 to 30714d8 Compare January 15, 2019 01:54
@dan-zheng
Copy link
Collaborator Author

@swift-ci Please test tensorflow

@dan-zheng
Copy link
Collaborator Author

@swift-ci Please clean test tensorflow macOS

@dan-zheng
Copy link
Collaborator Author

@swift-ci Please test tensorflow Linux GPU

@dan-zheng dan-zheng force-pushed the fix-differentiable-attr-silgen branch 8 times, most recently from 189593d to 7f9effe Compare January 16, 2019 02:37
rxwei
rxwei approved these changes Jan 16, 2019
View reviewed changes
lib/Serialization/Serialization.cpp Outdated Show resolved Hide resolved
@dan-zheng
[AutoDiff] Fix @differentiable attribute SILGen and serialization.
a5d5ae1 
- Propagate `@differentiable` attributes in `SILFunctionBuilder::addFunctionAttributes`.
  - Previously, this code was missing, causing `@differentiable` AST
    attributes to not propagate during SILGen.
- Rework differentiation to never load SIL functions.
  - Previously, differentiation relied on explicit loading to get SIL
    `[differentiable]` attributes. That was a hack.
  - Now that `@differentiable` attribute is propagated and serialized,
    such loading is no longer necessary.
- Tighten differentiation infrastructure to uphold the following invariants:
  - Differentiable functions' primal/adjoint are visible only in defining module.
  - Differentiable functions' JVP/VJP are always visible to other modules.
    The differentiation pass guarantees that JVP/VJPs exist.

Todos:
- Replace adjoint definitions in stdlib with VJP definitions.
- Eliminate primal/adjoint from `@differentiable` attribute.
  Use JVP/VJPs everywhere.
@dan-zheng
Add cross-module serialization test.
0eb023d
@dan-zheng dan-zheng force-pushed the fix-differentiable-attr-silgen branch 2 times, most recently from 5d8016b to 0aca275 Compare January 16, 2019 03:20
@dan-zheng dan-zheng force-pushed the fix-differentiable-attr-silgen branch from 0aca275 to 61ccf59 Compare January 16, 2019 03:25
@dan-zheng
Fix SILLoader related issue.
2cc9518 
Calling `module.getSILLoader()` here is necessary to prevent crashes related to
`lookUpFunctionInWitnessTable`. This might be a bug in the `SILLoader`
implementation.
@dan-zheng
Fix serialization tests.
d16f208
@dan-zheng dan-zheng changed the title [AutoDiff] Fix SILGen for @differentiable attribute. [AutoDiff] Fix @differentiable attribute SILGen and serialization. Jan 16, 2019
@dan-zheng
Fix crasher during argument explosion optimization.
2a1a645 
`test/AutoDiff/simple_model.swift` failed during argument epxlosion optimization.
@dan-zheng
Copy link
Collaborator Author

@swift-ci Please test tensorflow

@dan-zheng dan-zheng merged commit 835f124 into apple:tensorflow Jan 16, 2019
@dan-zheng dan-zheng deleted the fix-differentiable-attr-silgen branch January 16, 2019 06:32
@dan-zheng
Copy link
Collaborator Author

This was joint work with @rxwei. 🤝

@rxwei rxwei mentioned this pull request Jan 16, 2019
Closed
@dan-zheng dan-zheng mentioned this pull request Jan 16, 2019
Closed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Reviewers

@rxwei rxwei rxwei approved these changes

Assignees
No one assigned
Labels
tensorflow This is for "tensorflow" branch PRs.
Projects
None yet
Milestone
No milestone
Development

Successfully merging this pull request may close these issues.

None yet
2 participants
@dan-zheng @rxwei