-
Notifications
You must be signed in to change notification settings - Fork 10.5k
[AutoDiff] Fix SILGen JVP/VJP thunking bug. #26448
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
Conversation
SILGen JVP/VJP thunks for methods now always perform self reordering. Add leak checking test. Previously, self reordering was not performed if actual JVP/VJP type matched expected JVP/VJP type, which is incorrect. Gardening included: - Change `SILFunction *SILGenFunction::getOrCreateAutoDiffLinearMapThunk` to `ManagedValue SILGenFunction::getThunkedAutoDiffLinearMap`. The new helper is more ergonomic and directly returns thunked values. - Use ad-hoc `_vtable_entry_thunk` suffix for JVP/VJP vtable entry thunks. - Add todo comments for TF-685: principled AD thunk mangling. Resolves TF-698.
@swift-ci Please test tensorflow |
SILFunction *getOrCreateAutoDiffLinearMapThunk( | ||
AutoDiffAssociatedFunctionKind assocFnKind, | ||
ManagedValue getThunkedAutoDiffLinearMap( | ||
ManagedValue linearMap, AutoDiffAssociatedFunctionKind assocFnKind, |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Use the recently added AutoDiffLinearMapEnum
:
ManagedValue linearMap, AutoDiffAssociatedFunctionKind assocFnKind, | |
ManagedValue linearMap, AutoDiffLinearMapEnum linearMapKind, |
Will address in a follow-up to unblock progress.
This PR fixes latent bug in nightly toolchains.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
The leak checking tests you added are valid, but they are almost never effective for the thunks you added. They do not test whether you leaked any partial_apply’s. You can either convert all of these thunks to use ownership (which should be relatively easy), or manually test these cases before you merge. Making these thunks use ownership is a must in the short term.
@@ -3746,8 +3761,8 @@ SILGenModule::getOrCreateAutoDiffAssociatedFunctionThunk( | |||
|
|||
SmallVector<SILValue, 8> directResults; | |||
extractAllElements(apply, thunkSGF.B, directResults); | |||
auto linearMap = directResults.back(); | |||
auto linearMapFnType = linearMap->getType().castTo<SILFunctionType>(); | |||
auto linearMap = ManagedValue::forBorrowedObjectRValue(directResults.back()); |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
You aren’t using any SILGen memory management in this thunk. This should be ‘ManagedValue::forUnmanaged’ or something close (I forgot the exact name).
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Will address shortly in a follow-up.
(I plan to enable ownership for AD associated function thunks soon.)
I have a WIP patch to enable ownership in SILGen thunks. Is it okay for this patch land first to resolve TF-698? |
Yep, resolving correctness is more important. Could you please run a few manual memory tests through Instruments before merging? |
I'm AFK now but I can use Instruments to test leaks in a few hours. (I understand the need to test |
My patch enabling ownership for AD linear map thunks is ready for review. |
SILGen JVP/VJP thunks for methods now always perform self reordering.
Add leak checking test.
Previously, self reordering was not performed if actual JVP/VJP type
matched expected JVP/VJP type, which is incorrect.
Gardening included:
SILFunction *SILGenFunction::getOrCreateAutoDiffLinearMapThunk
toManagedValue SILGenFunction::getThunkedAutoDiffLinearMap
. The new helperis more ergonomic and directly returns thunked values.
_vtable_entry_thunk
suffix for JVP/VJP vtable entry thunks.Resolves TF-698.
Reproducer:
Fixes derivatives of
Tensor.concatenated(with:)
andTensor.++
: tensorflow/swift-apis#403