Skip to content

[AutoDiff] Start supporting loadable types with address-only tangents. #32540

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

Conversation

dan-zheng
Copy link
Contributor

Previously, PullbackEmitter assumed that original values' value category
matches their TangentVector types' value category. This was problematic
for loadable types with address-only TangentVector types.

Now, PullbackEmitter starts to support differentiation of loadable types with
address-only TangentVector types. This patch focuses on supporting and testing
class types, more support can be added incrementally.

Resolves TF-1149.
Unblocks future _Differentiation module changes like TF-1267.


Example:

import _Differentiation

// Class<T>: loadable type
// Class<T>.TangentVector: address-only struct type

class Class<T: Differentiable>: Differentiable {
  @differentiable
  var stored: T

  init(_ stored: T) {
    self.stored = stored
  }

  // struct TangentVector: Differentiable, AdditiveArithmetic {
  //  var stored: T.TangentVector
  // }
}

Before: cryptic error due to artificial limitations (lack of compiler support).

$ swift tf-1149.swift
tf-1149.swift:7:4: error: function is not differentiable
  @differentiable
  ~^~~~~~~~~~~~~~
tf-1149.swift:8:7: note: when differentiating this function definition
  var stored: T
      ^
tf-1149.swift:8:7: note: cannot yet differentiate value whose type 'Class<T>' has a compile-time known size, but whose 'TangentVector' contains stored properties of unknown size; consider modifying 'Class<τ_0_0>.TangentVector' to use fewer generic parameters in stored properties
  var stored: T
      ^

After: no error.

Previously, PullbackEmitter assumed that original values' value category
matches their `TangentVector` types' value category. This was problematic
for loadable types with address-only `TangentVector` types.

Now, PullbackEmitter starts to support differentiation of loadable types with
address-only `TangentVector` types. This patch focuses on supporting and testing
class types, more support can be added incrementally.

Resolves TF-1149.
@dan-zheng dan-zheng requested review from rxwei and marcrasi June 25, 2020 00:13
@dan-zheng
Copy link
Contributor Author

@swift-ci Please smoke test

@@ -178,6 +178,21 @@ Optional<TangentSpace> PullbackEmitter::getTangentSpace(CanType type) {
LookUpConformanceInModule(getModule().getSwiftModule()));
}

SILValueCategory PullbackEmitter::getTangentValueCategory(SILValue v) {
// Quick check: if the value has an address type, the tangent value category
Copy link
Contributor

@rxwei rxwei Jun 25, 2020

Choose a reason for hiding this comment

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

This is not always true. Address only original types can totally have an non-address-only tangent, simplest example:

struct Foo<T>: Differentiable {
    var a: Float
    @noDerivative var b: T
}

Could you add a FIXME here with an appropriate bug number?

Copy link
Contributor Author

Choose a reason for hiding this comment

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

Address only original types can totally have an non-address-only tangent

This case is actually currently supported today. The differentiation transform always creates adjoint buffers for address-typed original values, regardless of whether the tangent type is loadable or address-only.

Offline, you pointed out that we should make "tangent value category" depend solely on whether the tangent type is loadable or address-only. For loadable tangent types, using symbolic adjoint values instead of concrete adjoint buffers would be more efficient.

I filed SR-13077 to track this improvement and left a comment in the code.

Clarify `PullbackEmitter::getTangentValueCategory`.
@dan-zheng
Copy link
Contributor Author

@swift-ci Please smoke test

@dan-zheng dan-zheng merged commit 46c2073 into swiftlang:master Jun 25, 2020
@dan-zheng dan-zheng deleted the autodiff-addressonly-tangentvector branch June 25, 2020 15:27
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

2 participants