[AutoDiff] Serialize derivative function configurations per module. #28608
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
dan-zheng
commented
Dec 6, 2019
| }); | ||
| } | ||
|
|
||
| ArrayRef<AutoDiffConfig> |
There was a problem hiding this comment.
Note: even though @differentiable and @derivative attributes do not currently store result indices, I decided to use AutoDiffConfig (which contains result indices) to represent AST derivative function configurations. I also tried creating a new "ASTAutoDiffConfig" struct containing just parameter indices and derivative generic signature, but decided against the code duplication.
I think the direction is to eventually plumb result indices through the differentiation system (TF-1038), so preemptively using result indices (AutoDiffConfig) seems good.
dan-zheng
force-pushed
the
derivative-config-lookup
branch
2 times, most recently
from
d38c72b to
7ab66b9
Compare
`@differentiable` and `@derivative` attributes register derivatives for `AbstractFunctionDecl`s for a particular "derivative function configuration": parameter indices and dervative generic signature. To find `@derivative` functions registered in other Swift modules, derivative function configurations must be serialized per module. When configurations for a `AbstractFunctionDecl` are requested, all configurations from imported modules are deserialized. This module serialization technique has precedent: it is used for protocol conformances (e.g. extension declarations for a nominal type) and Obj-C members for a class type. Add `AbstractFunctionDecl::getDerivativeFunctionConfigurations` entry point for accessing derivative function configurations. Use `AbstractFunctionDecl::getDerivativeFunctionConfigurations` to implement `findMinimalDerivativeConfiguration` for canonical derivative function configuration lookup, replacing `getMinimalASTDifferentiableAttr`. Unblocks TF-815: lowering `@derivative` attributes directly to SIL differentiability witnesses without generating implicit `@differentiable` attributes.
dan-zheng
force-pushed
the
derivative-config-lookup
branch
from
7ab66b9 to
ceb8ce0
Compare
marcrasi
approved these changes
Dec 6, 2019
include/swift/SILOptimizer/Utils/Differentiation/DerivativeLookup.h
Outdated
Show resolved
Hide resolved
Change `findMinimalDerivativeConfiguration` to return `Optional<AutoDiffConfig>`.
|
@swift-ci Please test tensorflow |
|
Extended test suite passed. Merging to unblock progress. |
|
Very nice. Thanks for pushing this through :) |
dan-zheng
added a commit
to dan-zheng/swift
that referenced
this pull request
Dec 6, 2019
Previously, `@derivative` attribute type-checking created implicit `@differentiable` attributes on the original declaration. This was a longstanding hack powering `@derivative` attribute derivative registration. swiftlang#28608 made these changes: - Derivative function configurations (from `@differentiable` and `@derivative` attributes) are serialized in modules and are loaded from imported modules. - The differentiation transform uses these derivative function configurations for derivative function lookup instead of `@differentiable` attributes. Now, `@derivative` attributes are directly lowered to differentiability witnesses during SILGen, and implicit `@differentiable` attribute generation is removed. Type-checking changes: - "Overlapping" `@differentiable` and `@derivative` attributes (for the same original declaration and parameter indices) are now disallowed. They semantically conflict because the first "requests derivative generation" while the second "registers a derivative". - "Overlapping" `@differentiable` and `@derivative` attributes are allowed for protocol requirements. Requirement `@differentiable` attributes mean "add JVP/VJP witness table entries" - not "request derivative generation", because there is no function body. - Note that relaxing the "overlapping" condition to consider derivative generic signatures is possible after derivative generic signature mangling for derivative functions: TF-680. Resolves TF-835. Unblocks TF-1021: lifting the "same-file derivative registration only" limitation in `@derivative` attribute type-checking. This should be possible without much work, but needs testing! Exposes TF-1040: `@differentiable` attribute limitations for class methods. Exposes TF-1041: untested protocol requirement `@differentiable` attribute type-checking logic.
dan-zheng
added a commit
that referenced
this pull request
Dec 6, 2019
Previously, `@derivative` attribute type-checking created implicit `@differentiable` attributes on the original declaration. This was a longstanding hack powering `@derivative` attribute derivative registration. #28608 made these changes: - Derivative function configurations (from `@differentiable` and `@derivative` attributes) are serialized in modules and are loaded from imported modules. - The differentiation transform uses these derivative function configurations for derivative function lookup instead of `@differentiable` attributes. Now, `@derivative` attributes are directly lowered to differentiability witnesses during SILGen, and implicit `@differentiable` attribute generation is removed. Type-checking changes: - "Overlapping" `@differentiable` and `@derivative` attributes (for the same original declaration and parameter indices) are now disallowed. They semantically conflict because the first "requests derivative generation" while the second "registers a derivative". - "Overlapping" `@differentiable` and `@derivative` attributes are allowed for protocol requirements. Requirement `@differentiable` attributes mean "add JVP/VJP witness table entries" - not "request derivative generation", because there is no function body. - Note that relaxing the "overlapping" condition to consider derivative generic signatures is possible after derivative generic signature mangling for derivative functions: TF-680. Resolves TF-835. Unblocks TF-1021: lifting the "same-file derivative registration only" limitation in `@derivative` attribute type-checking. This should be possible without much work, but needs testing! Exposes TF-1040: `@differentiable` attribute limitations for class methods. Exposes TF-1041: untested protocol requirement `@differentiable` attribute type-checking logic.
dan-zheng
added a commit
that referenced
this pull request
Dec 10, 2019
#28621) Previously, `@derivative` attribute type-checking created implicit `@differentiable` attributes on the original declaration. This was a longstanding hack powering `@derivative` attribute derivative registration. #28608 made these changes: - Derivative function configurations (from `@differentiable` and `@derivative` attributes) are serialized in modules and are loaded from imported modules. - The differentiation transform uses these derivative function configurations for derivative function lookup instead of `@differentiable` attributes. Now, `@derivative` attributes are directly lowered to differentiability witnesses during SILGen, and implicit `@differentiable` attribute generation is removed. `@derivative` attributes are now serialized. Resolves TF-835. Unblocks TF-1021: lifting the "same-file derivative registration only" limitation in `@derivative` attribute type-checking. This should be trivially possible but requires more testing. Exposes TF-1037: crash due to `@differentiable` + `@derivative` attribute with different derivative generic signatures. Exposes TF-1040: `@differentiable` attribute limitations for class methods. Exposes TF-1041: untested protocol requirement `@differentiable` attribute type-checking logic. Tracks TF-1042: remove `ASTContext::{Differentiable,Derivative}Attrs` and use `AbstractFunctionDecl::getDerivativeFunctionConfigurations` to detect duplicate `@differentiable` + `@derivative` attributes.
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.
@differentiableand@derivativeattributes register derivatives forAbstractFunctionDecls for a particular "derivative function configuration":parameter indices and derivative generic signature.
To find
@derivativefunctions registered in other Swift modules, derivativefunction configurations must be serialized per module. When configurations for
a
AbstractFunctionDeclare requested, all configurations from importedmodules are deserialized. This module serialization technique has precedent: it
is used for protocol conformances (e.g. extension declarations for a nominal
type) and Obj-C members for a class type.
Add
AbstractFunctionDecl::getDerivativeFunctionConfigurationsentry pointfor accessing derivative function configurations.
Use
AbstractFunctionDecl::getDerivativeFunctionConfigurationstoimplement
findMinimalDerivativeConfigurationfor canonical derivativefunction configuration lookup, replacing
getMinimalASTDifferentiableAttr.Unblocks TF-835: lowering
@derivativeattributes directly to SILdifferentiability witnesses without generating implicit
@differentiableattributes.