[TF] TF-496: Conditionally conform SIMD types to Differentiable #24786

bartchr808 · 2019-05-14T22:49:03Z

Implements this JIRA ticket which makes SIMD vectors differentiable.

Also here is a link to the Swift Forums where some of this was discussed. See the forum post for any context regarding the changes.

There is also the change to the -=/+= operator. As background, we need SIMD to conform to the AdditiveArithmetic protocol which already provides a conditional definition of the -= and += operator. So to avoid the "dual default conditional definition protocol" issue, I removed the default implementation in SIMD for now on our branch, but we will fix it when we pitch it to get it merged into master.

Also, I had to remove @_alwaysEmitIntoClient from sum() due to a bug (filed as TF-545).

And finally, there are custom VJPs that I defined that don't actually need to be defined, as the body of the original methods define methods that I already marked as differentiable. But the bug TF-547 is blocking me on that.

bartchr808 · 2019-05-14T22:54:51Z

Forgot to remove the zero property as mentioned in the forum....will push change soon.

stdlib/public/core/SIMDVectorTypes.swift.gyb

bartchr808 · 2019-05-15T00:12:41Z

@stephentyrone I'm not sure what you and the the Swift Core team thinks as to whether this is a bug fix or a new feature that would require a proposal for, but I just realized that we could create a separate PR just for the -=/+= part and merge that into master.

Additionally, for the +=/-= fix of my PR, which do you think would be better (or possibly even an alternate idea if you have one):

My current PR where I removed the operators from the SIMD protocol, keeping the one in AdditiveArithmetic
The complete reverse where we keep the SIMD protocol conditional definition, and remove the conditional definition in AdditiveArithmetic

In either case the SIMD implementation would be preferred, it just depends whether in the protocol or the various structs

rxwei · 2019-05-16T05:27:27Z

stdlib/public/core/SIMDVectorTypes.swift.gyb

+  where Scalar : Differentiable & BinaryFloatingPoint,
+        Scalar.CotangentVector : BinaryFloatingPoint {
+  public func _vjpSubscript(index: Int) ->
+  (Scalar, (Scalar.CotangentVector) -> CotangentVector)


Indent by 2 and put the -> on this line.

Suggested change

(Scalar, (Scalar.CotangentVector) -> CotangentVector)

-> (Scalar, (Scalar.CotangentVector) -> CotangentVector)

To verify:

public func _vjpSubscript(index: Int) -> (Scalar, (Scalar.CotangentVector) -> CotangentVector)

Yes. VJP functions should not be public though. When the original function is public, a VJP function should be internal. When it's internal, if the original function has @inlinable, the VJP function should be @inlinable, otherwise the VJP function should be @usableFromInline.

Ahh okay. But for let's say _vjpSum, the original function (in the protocol) is marked as public and @inlinable. How should _vjpSum be marked? Looking at Tensor, its _vjpSum is just @inlinable in a regular extension

If the original function is @inlinable, the VJP should also be @inlinable. In all cases where the original function is public, the VJP should be internal.

But for SIMD.sum(), it's not @inlinable, so the VJP should be @usableFromInline.

Ahh okay thanks!

stdlib/public/core/SIMDVectorTypes.swift.gyb

rxwei · 2019-05-16T05:29:40Z

stdlib/public/core/SIMDVectorTypes.swift.gyb

+  public func _vjpSubscript(index: Int) ->
+  (Scalar, (Scalar.CotangentVector) -> CotangentVector)
+  {
+    return (self[index], { (v: Scalar.CotangentVector) in


I don't think the type signature is needed here. Did the compiler suggest otherwise?

Suggested change

return (self[index], { (v: Scalar.CotangentVector) in

return (self[index], { v in

At some point it was asking me for it either here in the struct or in the SIMD protocol where it gave me an error along the lines of cannot convert return expression of type '(Self, (_) -> _)' to return type '(Self, (Self.CotangentVector) -> (Self.Scalar, Self.CotangentVector))', but I think this ended up being fixed when I corrected the protocol refinements. So I actually can remove these all now!

stdlib/public/core/SIMDVector.swift

stdlib/public/core/SIMDVectorTypes.swift.gyb

stdlib/public/core/SIMDVector.swift

bartchr808 · 2019-05-16T21:43:36Z

Seems like this may have a bug, filed it here.

TODO: - Determine whether sum() needs to be differentiated - See what other inits I can differentiate - Remove some of the custom VJPs now that init(repeating:) is differentiable

bartchr808 · 2019-06-02T21:56:15Z

In case you see that I pushed commits to this PR, Ignore re-reviewing this PR for now while I clean up some stuff. Pushed it so it's accessible from desktop and laptop for me. Still need to do three things:

Determine whether sum() needs to be differentiated
See what other inits I can differentiate
Remove some of the custom VJPs now that init(repeating:) is differentiable

bartchr808 · 2019-06-02T22:21:14Z

stdlib/public/core/SIMDVector.swift

@@ -843,6 +874,13 @@ extension SIMD where Scalar: FloatingPoint {

  /// Returns the sum of the scalars in the vector.
  @_alwaysEmitIntoClient


Need to look into this, but I think this attribute is what is preventing me from differentiating the sum()

bartchr808 · 2019-06-03T20:02:51Z

stdlib/public/core/SIMDVector.swift

  public static func /(lhs: Self, rhs: Scalar) -> Self {
    return lhs / Self(repeating: rhs)
  }

-  @_transparent
-  public static func +=(lhs: inout Self, rhs: Self) {


In order to make SIMD conform to AdditiveArithmetic, I had to remove both of these implementations as two protocols cannot both have default implementations. Let me know if this is okay, or if there is something else we should do.

bartchr808 · 2019-06-03T20:04:56Z

@swift-ci please test tensorflow

bartchr808 · 2019-06-03T20:06:06Z

@swift-ci please test tensorflow

bartchr808 · 2019-06-03T21:07:21Z

TODO: getting assertion failure if I want to eliminate redundant custom VJPs: https://bugs.swift.org/browse/TF-547

dan-zheng

LGTM!

stdlib/public/core/SIMDVector.swift

dan-zheng · 2019-06-13T16:52:07Z

stdlib/public/core/SIMDVectorTypes.swift.gyb

+  -> (Scalar, (Scalar.TangentVector) -> TangentVector) {
+    return (self[index], { v in
+      var zeros = Self.zero
+      zeros[index] = Scalar(v)


Is the Scalar initializer call here necessary?

Suggested change

zeros[index] = Scalar(v)

zeros[index] =v

v is of type Scalar.TangentVector while zeros is of type SIMD3<Scalar>, and since the only constraint on Scalar.TangentVector is Scalar.TangentVector : BinaryFloatingPoint which Scalar also conforms to, I have to cast it or else I I get the error cannot assign value of type 'Scalar.TangentVector' to type 'Scalar'. I could add the more constraining requirement of Scalar.TangentVector == Scalar, but I think doing this cast is better. Let me know what you think!

cc @rxwei regarding whether to add Scalar.TangentVector == Scalar constraint.
I agree that not adding the constraint is better.

test/AutoDiff/SIMD.swift

bartchr808 · 2019-06-13T22:50:41Z

@swift-ci please test tensorflow

dan-zheng

Generic SIMD differentiation tests look great!

bartchr808 · 2019-06-14T15:24:36Z

@dan-zheng macOS tests are still going 16 hours later, is this a known problem currently? Or is this the linux passing good enough to merge in?

rxwei

Great work Bart!

bartchr808 · 2019-06-14T20:57:38Z

Thanks @dan-zheng and @rxwei ! 😄

bartchr808 added 6 commits May 13, 2019 14:23

Initial broken version.

7fb96b1

Add VectorView.

d3182a3

WIP while waiting for swift forums answer.

aa23187

WIP: removed from AdditiveArithmetic, kept in CIMD.

fe38362

Remove += from SIMD protocol, move it to Struct SIMDn.

5f109a0

Have SIMDn conform to Differentiable.

304026b

bartchr808 requested review from stephentyrone and rxwei May 14, 2019 22:49

Forgot to remove zero, already implemented in SIMDVector.swift

b75db94

dan-zheng added the tensorflow This is for "tensorflow" branch PRs. label May 14, 2019

rxwei reviewed May 14, 2019

View reviewed changes

stdlib/public/core/SIMDVectorTypes.swift.gyb Outdated Show resolved Hide resolved

PR fix: spacing.

d67acef

bartchr808 changed the title ~~TF-496: Conditionally conform SIMD types to Differentiable~~ [TF] TF-496: Conditionally conform SIMD types to Differentiable May 14, 2019

bartchr808 added 4 commits May 14, 2019 22:36

Start making + differentiable

17f52bf

Get + and - differentiable.

031f384

WIP

83e7f75

Make all necessary operators differentiable.

0268fc6

bartchr808 force-pushed the TF-496-conform-SIMD-types-differentiable branch from 32e963a to 0268fc6 Compare May 16, 2019 05:15

rxwei reviewed May 16, 2019

View reviewed changes

bartchr808 added 2 commits May 16, 2019 10:38

Address 1st batch of PR comments.

72d13e5

Merge branch 'tensorflow' into TF-496-conform-SIMD-types-differentiable

91fd99b

rxwei reviewed May 16, 2019

View reviewed changes

stdlib/public/core/SIMDVectorTypes.swift.gyb Outdated Show resolved Hide resolved

rxwei reviewed May 16, 2019

View reviewed changes

stdlib/public/core/SIMDVectorTypes.swift.gyb Show resolved Hide resolved

stdlib/public/core/SIMDVectorTypes.swift.gyb Outdated Show resolved Hide resolved

stdlib/public/core/SIMDVector.swift Show resolved Hide resolved

bartchr808 added 3 commits May 16, 2019 11:58

Fix due to changes in PR swiftlang#24825.

e057686

Address 2nd batch of PR comments.

260f2fb

Make _vjpSubscript internal and inline.

c4de3c9

bartchr808 removed the request for review from stephentyrone May 16, 2019 21:43

Fix bugs in code and add tests

99b8408

TODO: - Determine whether sum() needs to be differentiated - See what other inits I can differentiate - Remove some of the custom VJPs now that init(repeating:) is differentiable

bartchr808 commented Jun 2, 2019

View reviewed changes

bartchr808 added 2 commits June 2, 2019 18:10

Get sum() differentiable.

ffdb4c4

Cleanup and additional init(repeating:) test.

c8df93f

bartchr808 commented Jun 3, 2019

View reviewed changes

bartchr808 force-pushed the TF-496-conform-SIMD-types-differentiable branch 2 times, most recently from d387338 to b6b8fe8 Compare June 3, 2019 20:04

bartchr808 force-pushed the TF-496-conform-SIMD-types-differentiable branch from b6b8fe8 to 8693e01 Compare June 4, 2019 17:58

Fix up AdditiveArithmetic conformance.

2821b8f

bartchr808 force-pushed the TF-496-conform-SIMD-types-differentiable branch from 8693e01 to 2821b8f Compare June 4, 2019 17:59

Merge branch 'tensorflow' into TF-496-conform-SIMD-types-differentiable

bc32cf9

bartchr808 force-pushed the TF-496-conform-SIMD-types-differentiable branch from d7c29b5 to 8d495dd Compare June 13, 2019 15:47

Code cleanup and additonal tests.

7abb1c3

bartchr808 force-pushed the TF-496-conform-SIMD-types-differentiable branch from 8d495dd to 7abb1c3 Compare June 13, 2019 15:51

White space.

0adb6eb

bartchr808 requested review from rxwei and dan-zheng June 13, 2019 15:57

dan-zheng approved these changes Jun 13, 2019

View reviewed changes

Add generic tests.

d73f3da

dan-zheng approved these changes Jun 13, 2019

View reviewed changes

rxwei approved these changes Jun 14, 2019

View reviewed changes

bartchr808 merged commit 3417a08 into swiftlang:tensorflow Jun 14, 2019

bartchr808 deleted the TF-496-conform-SIMD-types-differentiable branch June 14, 2019 21:05

	(Scalar, (Scalar.CotangentVector) -> CotangentVector)
	-> (Scalar, (Scalar.CotangentVector) -> CotangentVector)

	return (self[index], { (v: Scalar.CotangentVector) in
	return (self[index], { v in

		@@ -843,6 +874,13 @@ extension SIMD where Scalar: FloatingPoint {

		/// Returns the sum of the scalars in the vector.
		@_alwaysEmitIntoClient

[TF] TF-496: Conditionally conform SIMD types to Differentiable #24786

[TF] TF-496: Conditionally conform SIMD types to Differentiable #24786

Uh oh!

Conversation

bartchr808 commented May 14, 2019 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

bartchr808 commented May 14, 2019 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Uh oh!

bartchr808 commented May 15, 2019 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

bartchr808 commented May 16, 2019 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

bartchr808 commented Jun 2, 2019

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

bartchr808 commented Jun 3, 2019

Uh oh!

bartchr808 commented Jun 3, 2019

Uh oh!

bartchr808 commented Jun 3, 2019

Uh oh!

dan-zheng left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

bartchr808 Jun 13, 2019 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

bartchr808 commented Jun 13, 2019

Uh oh!

dan-zheng left a comment

Choose a reason for hiding this comment

Uh oh!

bartchr808 commented Jun 14, 2019

Uh oh!

rxwei left a comment

bartchr808 commented May 14, 2019 •

edited

Loading

bartchr808 commented May 14, 2019 •

edited

Loading

bartchr808 commented May 15, 2019 •

edited

Loading

bartchr808 commented May 16, 2019 •

edited

Loading

bartchr808 Jun 13, 2019 •

edited

Loading