Skip to content

[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

Merged
merged 1 commit into from
Aug 6, 2019

Conversation

dan-zheng
Copy link
Contributor

@dan-zheng dan-zheng commented Aug 1, 2019

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.


Reproducer:

struct TF_698 : Differentiable & AdditiveArithmetic {
  var x: Float
  init(_ x: Float) {
    self.x = x
  }

  @differentiable(jvp: jvpMultiplied, vjp: vjpMultiplied)
  func multiplied(with other: Self) -> Self {
    return Self(x: x * other.x)
  }

  // Helper function to shorten JVP/VJP definitions.
  static func *(lhs: Self, rhs: Self) -> Self {
    return lhs.multiplied(with: rhs)
  }

  func jvpMultiplied(with other: Self) -> (Self, (Self, Self) -> Self) {
    return (multiplied(with: other), { dself, dother in
      self * dother + other * dself
    })
  }

  func vjpMultiplied(with other: Self) -> (Self, (Self) -> (Self, Self)) {
    return (multiplied(with: other), { v in
      (v * other, v * self)
    })
  }
}
let v = TF_698(1)
print(pullback(at: TF_698(3), TF_698(10)) { x, y in x.multiplied(with: y) }(v))
// Before:           (TF_698(3), TF_698(10))
// After (expected): (TF_698(10), TF_698(3))

Fixes derivatives of Tensor.concatenated(with:) and Tensor.++: tensorflow/swift-apis#403

@dan-zheng dan-zheng added the tensorflow This is for "tensorflow" branch PRs. label Aug 1, 2019
@dan-zheng dan-zheng requested a review from rxwei August 1, 2019 04:08
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.
@dan-zheng
Copy link
Contributor Author

@swift-ci Please test tensorflow

SILFunction *getOrCreateAutoDiffLinearMapThunk(
AutoDiffAssociatedFunctionKind assocFnKind,
ManagedValue getThunkedAutoDiffLinearMap(
ManagedValue linearMap, AutoDiffAssociatedFunctionKind assocFnKind,
Copy link
Contributor Author

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:

Suggested change
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.

Copy link
Contributor

@rxwei rxwei left a 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());
Copy link
Contributor

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).

Copy link
Contributor Author

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.)

@dan-zheng
Copy link
Contributor Author

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.

I have a WIP patch to enable ownership in SILGen thunks. Is it okay for this patch land first to resolve TF-698?

@rxwei
Copy link
Contributor

rxwei commented Aug 1, 2019

Yep, resolving correctness is more important. Could you please run a few manual memory tests through Instruments before merging?

@dan-zheng
Copy link
Contributor Author

dan-zheng commented Aug 1, 2019

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 partial_apply leaks, like in #26384, and why the current tests are ineffective - thanks!)

@dan-zheng
Copy link
Contributor Author

My patch enabling ownership for AD linear map thunks is ready for review.
I'll merge this PR first (fixing correctness but not leaks) to separate orthogonal changes.

@dan-zheng dan-zheng merged commit 2f89554 into swiftlang:tensorflow Aug 6, 2019
@dan-zheng dan-zheng deleted the TF-698 branch August 6, 2019 00:51
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
tensorflow This is for "tensorflow" branch PRs.
Projects
None yet
Development

Successfully merging this pull request may close these issues.

2 participants