[AutoDiff] make diff_wit_fns instead of fn_refs #28451
Merged
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
marcrasi
force-pushed
the
marcrasi-use-diff-witness
branch
from
38824cf to
4036d53
Compare
|
@swift-ci please test tensorflow |
|
@swift-ci Please test tensorflow linux |
rxwei
reviewed
Nov 23, 2019
lib/SILOptimizer/Mandatory/Differentiation/DerivativeLookup.cpp
Outdated
Show resolved
Hide resolved
dan-zheng
approved these changes
Nov 23, 2019
| const auto *minimalAttr = getMinimalASTDifferentiableAttr( | ||
| originalAFD, parameterIndices, minimalParameterIndices); | ||
|
|
||
| // TODO(TF-835): This will also need to search all `@differentiating` |
There was a problem hiding this comment.
Great future facing note!
lib/SILOptimizer/Mandatory/Differentiation/DerivativeLookup.cpp
Outdated
Show resolved
Hide resolved
lib/SILOptimizer/Mandatory/Differentiation/DerivativeLookup.cpp
Outdated
Show resolved
Hide resolved
|
I figured out why the tests are failing. I pushed a fix that makes it slightly better but does not fix it. I have to sign off now, feel free to do more stuff to this. The reason the tests are failing is that inlinable functions with differentiability_witness_function instructions are getting inlined into the module, but the corresponding witness tables are not getting inlined into the module. So deserialization tries to look up the witness table, does not find it, and dies. |
|
I had a bit of time to investigate the failure today, so I added a simple testcase that triggers it that should make it pretty clear what's happening. In the testcase, the Okay, that's all I'll be doing for now. |
Thanks for the clear explanation! I think I understand the SIL serialization issues and I'll investigate now. Edit: I merged the fix #28463 to |
dan-zheng
reviewed
Nov 24, 2019
test/AutoDiff/sil_differentiability_witness_reference_serialization.sil
Outdated
Show resolved
Hide resolved
TF-993 tracks further organization of differentiation code.
|
@swift-ci Please test tensorflow |
rxwei
approved these changes
Nov 25, 2019
Shall we investigate the (compile-time and/or runtime) performance regressions before merging this PR? I suspect generating |
|
I'll spend a bit of time measuring the compile time and runtime performance impact of this. I'll report back on the PR before merging and we can decide whether to go ahead. |
Runtime performanceMicrobenchmarkWhen compiled with swift-models BenchmarkNo noticeable difference when running this benchmark. Internal Google modelAn internal Google model is affected, with about the same amount of performance regression as the microbenchmark. DecisionI will implement devirtualization before merging this patch because we don't want to check in such a big performance regression without being sure that we have a solution. |
|
A swift-apis failure also remains after the latest fixes, so I will investigate that too: |
|
@swift-ci please test tensorflow |
1 similar comment
|
@swift-ci please test tensorflow |
|
@swift-ci please test tensorflow |
|
@swift-ci please test tensorflow macos |
|
Updates! tldr: Only one remaining blocker, a swift-apis compile failure, which I am investigating now. Runtime performance #28480 has fully fixed the performance regression in the microbenchmark posted above and in the internal google model. Compiletime performance There's a bit of slowdown as measured by running the AD compiler test suite. None of it seems egregious, so I will defer investigation. Filed TF-1013 for future investigation. You can see the measurements at https://gist.github.com/marcrasi/456d09d15a14f9af092db2e7b24186bb. swift-apis compile failure |
|
Okay, I have uploaded a fix that seems to work. I think it's not very proper, but if all the tests pass then it's probably good enough for now. Gonna retrigger the swift-apis tests. @swift-ci please test tensorflow |
|
@swift-ci please test tensorflow |
|
@swift-ci please test tensorflow macos |
1 similar comment
|
@swift-ci please test tensorflow macos |
|
swift-apis tests have gotten farther, but are still failing. The failure is currently a bit mysterious to me. Compilation succeeds, a lot of tests succeed, and then suddenly it says "Exited with signal code 11," with no other explanatory message. Investigating now. Full test output: |
|
Investigation update for the curious. This reproduces the segfault: The backtrace from lldb is The only SIL differences that this PR causes on this code are:
It's pretty surprising that this affects memory management. Current ideas of what might be going on:
Next, I'm going to carefully read through the SIL to see if I see anything suspicious in it. |
|
I think I solved it. I uploaded a commit, and I'm going to run all the tests now. This is the problem: The comment is false. LoadableByAddress can change the type of a function in a different module because pre-LoadableByAddress SIL references the pre-LoadableByAddress versions of external functions. So I have fixed it by making LoadableByAddress work on external differentiability witnesses. I have also added a cross-module LoadableByAddress test. While writing the test, I realized that we didn't have print/parse/serialization/deserialization for differentiability_witness_functions with explicit types, so I added all that too. |
|
@swift-ci please test tensorflow |
marcrasi
force-pushed
the
marcrasi-use-diff-witness
branch
from
fef6423 to
5f93d60
Compare
|
@swift-ci please test tensorflow |
|
This is really awesome. Thank you Marc! |
dan-zheng
reviewed
Nov 28, 2019
Nice! #28501 to |
|
@swift-ci Please test tensorflow |
|
Merging now! Verified that the extended test suite ( |
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.
Main change
The main point of this PR is to make the differentiation pass look up witnesses and make
differentiable_witness_functioninstructions, instead of looking up attributes and makingfunction_refs.Related changes
Some fixes were required to make this work properly:
LinkEntity::SecondaryPointerneeded to be set tonullptr. Otherwise it gets garbage that messes up equality comparisons.LoadableByAddressto handledifferentiable_witness_function.I used this PR as an opportunity to move all "derivative lookup" logic into its own file. I think this makes it easier to read and also it'll make incremental compilation faster.
What I'm doing next
I'm going to run toolchain build and swift-apis tests on this because it's a big change.
I'll be online from ~9pm to ~10pm tonight to fix any CI failures or address any comments. After that, I'll be out until Monday, so feel free to do whatever you want with this if you want to extend it over the weekend.
Next steps
This does not yet eliminate
SILDifferentiableAttr, which is still used for at least:SILDifferentiableAttr, and then the jvp/vjp from that gets copied to the witness.I think it'll be quite easy for a follow-up PR to remove these things and delete
SILDifferentiableAttr, and that's what I think the next PR should do.After that, all the typechecking and SILGen simplifications that we want should be unblocked. And then after that we should be able to implement cross-file and cross-module retrodiff and start using it in the stdlib.
I plan not to work on any of this until Monday, so feel free to start doing stuff on it if you want.