[AutoDiff upstream] Introduce @differentiable attribute to mark functions as differentiable. #27506

bgogul · 2019-10-02T23:57:09Z

This PR introduces @differentiable attribute to mark functions as differentiable. This PR only contains changes related to parsing the attribute. Type checking and other changes will be added in subsequent patches.

See https://github.com/apple/swift/pull/27506/files#diff-f3216f4188fd5ed34e1007e5a9c2490f for examples and tests for the new attribute.

saeta · 2019-10-03T01:41:15Z

@bgogul I think you might have forgotten to handle the switch at: https://github.com/apple/swift/pull/27506/files#diff-a99d62dc86d23321183dc044483cf2caR808

rxwei · 2019-10-03T08:49:47Z

I think the only thing we need to upstream in AutoDiff.h for parsing @differentiable is ParsedAutoDiffParameterIndices.

This PR introduces the `@transposing` attribute to mark functions as transposing other functions. This PR only contains changes related to parsing the attribute. Type checking and other changes will be added in subsequent patches. This work is related to the `@differentiable` attribute in #27506.

include/swift/AST/Attr.h

rxwei · 2019-10-15T17:34:46Z

@DougGregor could you review this patch when you get a chance?

rxwei · 2019-10-15T18:03:32Z

@swift-ci please smoke test

bgogul · 2019-10-15T21:28:12Z

@swift-ci please smoke test

bgogul · 2019-10-16T16:25:20Z

@swift-ci please clean smoke test

bgogul · 2019-10-16T18:12:42Z

@swift-ci please smoke test linux

bgogul · 2019-10-24T22:07:17Z

@DougGregor, hold off review on this one. I will try to split it further so that it is a slightly more easier to review.

…entally.

…erge.

bgogul · 2019-11-06T19:46:06Z

@swift-ci please smoke test

lattner · 2019-11-08T05:09:10Z

include/swift/AST/Attr.def

+DECL_ATTR(differentiable, Differentiable,
+  OnAccessor | OnConstructor | OnFunc | OnVar | OnSubscript | LongAttribute |
+  AllowMultipleAttributes |
+  ABIStableToAdd | ABIStableToRemove | APIStableToAdd | APIStableToRemove,


I believe this affects ABI, so I'd recommend removing the "ABIStableToAdd | ABIStableToRemove" tags, as well as APIStableToRemove which seems inappropriate.

It's not clear to me what these tags mean exactly, but I'll try the following configuration on tensorflow branch to see if things break:

ABIBreakingToAdd | ABIBreakingToRemove | APIStableToAdd | APIBreakingToRemove

Testing on tensorflow branch in #28148. The PR description has more context.

Replaced the attributes to be similar to tensorflow branch.

lattner

looks great with some changes requested, thanks!

lattner · 2019-11-08T05:10:30Z

include/swift/AST/AutoDiff.h

+#define SWIFT_AST_AUTODIFF_H
+
+#include "ASTContext.h"
+#include "llvm/ADT/SmallBitVector.h"


SmallBitVector.h looks unused.

Also, is ASTContext.h necessary here? Can you forward declare or include smaller headers? ASTContext.h includes the world.

Thanks for catching these. It turned out, I only needed IndexSubset.h instead of ASTContext.h

lattner · 2019-11-08T05:12:50Z

include/swift/AST/AutoDiff.h

+    return V.Ordered.Index;
+  }
+
+  enum Kind getKind() const {


nit: s/enum Kind/Kind/

lattner · 2019-11-08T05:17:42Z

lib/Parse/Parser.cpp

@@ -907,6 +907,18 @@ bool Parser::parseMatchingToken(tok K, SourceLoc &TokLoc, Diag<> ErrorDiag,
  return false;
 }

+bool Parser::parseUnsignedInteger(unsigned &Result, SourceLoc &Loc,


Why do you need this new method? DAK_Alignment handling in ParseDecl.cpp does this:

StringRef alignmentText = Tok.getText();
unsigned alignmentValue;
if (alignmentText.getAsInteger(0, alignmentValue)) {
diagnose(Loc, diag::alignment_must_be_positive_integer);
return false;
}

Won't this work for you at the callsite?

If you're trying to introduce a new generally useful utility, I'd recommend doing that as a separate patch which also migrates the existing callsites to it.

I found the code snippet similar to DAK_Alignment in a bunch of places. I agree it is better to introduce this in a separate PR. I will do so shortly.

There are also more uses of parseUnsignedInteger in the AD code base that has not been upstreamed yet.

lattner · 2019-11-08T05:20:27Z

lib/Parse/ParseDecl.cpp

+    // Check that token after comma is 'wrt:' or a function specifier label.
+    if (!Tok.is(tok::identifier) || !(Tok.getText() == "wrt" ||
+                                      Tok.getText() == "jvp" ||
+                                      Tok.getText() == "vjp")) {


Would it help to have a function like:

enum { wrt, jvp, vjp, invalid } classifyLabel(StringRef str);

function to keep all the classification logic in sync, and move the magic string parsing into a single place?

“jvp:” and “vjp:” will go away very soon with the almost-finished retroactive derivative registration feature (‘@differentiating’). So I think it is fine to leave them as is!

Good to know, but it is important to keep master consistent and following best practices. If these aren't important, then it would be fine to remove them from the patch. If they need to be in the patch, then please consider implementing them in a nice way :).

I'm not saying that classifyLabel is an appropriate thing, just saying that the rationale doesn't make sense.

Ok, sounds good!

I have added a bunch of helpers. PTAL.

(Although other code in master is using Tok.getText() to compare against literal strings.)

lattner · 2019-11-08T05:21:29Z

lib/Parse/ParseDecl.cpp

+                                 funcDiag, /*allowOperators=*/true,
+                                 /*allowZeroArgCompoundNames=*/true);
+    // If no trailing comma or 'where' clause, terminate parsing arguments.
+    if (Tok.isNot(tok::comma) && Tok.isNot(tok::kw_where))


Tok.isNot takes multiple arguments, please use that instead of && here and anywhere else this comes up.

lattner · 2019-11-08T05:24:19Z

test/AutoDiff/differentiable_attr_parse.swift

@@ -0,0 +1,180 @@
+// RUN: %target-swift-frontend -parse -verify %s
+
+/// Good


nice, thank you for the testcases!

bgogul

Thanks for the review, Chris!

bgogul · 2019-11-08T23:13:23Z

lib/Parse/Parser.cpp

@@ -907,6 +907,18 @@ bool Parser::parseMatchingToken(tok K, SourceLoc &TokLoc, Diag<> ErrorDiag,
  return false;
 }

+bool Parser::parseUnsignedInteger(unsigned &Result, SourceLoc &Loc,


I found the code snippet similar to DAK_Alignment in a bunch of places. I agree it is better to introduce this in a separate PR. I will do so shortly.

There are also more uses of parseUnsignedInteger in the AD code base that has not been upstreamed yet.

bgogul · 2019-11-08T23:13:46Z

include/swift/AST/AutoDiff.h

+    return V.Ordered.Index;
+  }
+
+  enum Kind getKind() const {


bgogul · 2019-11-08T23:14:33Z

include/swift/AST/AutoDiff.h

+#define SWIFT_AST_AUTODIFF_H
+
+#include "ASTContext.h"
+#include "llvm/ADT/SmallBitVector.h"


Thanks for catching these. It turned out, I only needed IndexSubset.h instead of ASTContext.h

bgogul · 2019-11-08T23:46:08Z

lib/Parse/ParseDecl.cpp

+    // Check that token after comma is 'wrt:' or a function specifier label.
+    if (!Tok.is(tok::identifier) || !(Tok.getText() == "wrt" ||
+                                      Tok.getText() == "jvp" ||
+                                      Tok.getText() == "vjp")) {


I have added a bunch of helpers. PTAL.

(Although other code in master is using Tok.getText() to compare against literal strings.)

bgogul · 2019-11-08T23:46:46Z

include/swift/AST/Attr.def

+DECL_ATTR(differentiable, Differentiable,
+  OnAccessor | OnConstructor | OnFunc | OnVar | OnSubscript | LongAttribute |
+  AllowMultipleAttributes |
+  ABIStableToAdd | ABIStableToRemove | APIStableToAdd | APIStableToRemove,


Replaced the attributes to be similar to tensorflow branch.

lattner · 2019-11-10T22:49:31Z

LGTM.

rxwei · 2019-11-10T23:57:15Z

@swift-ci please smoke test

rxwei · 2019-11-10T23:57:33Z

@swift-ci please smoke test

rxwei · 2019-11-11T10:26:21Z

@swift-ci please smoke test macOS

include/swift/AST/Attr.def

bgogul · 2019-11-11T19:47:13Z

@swift-ci please smoke test

gottesmm · 2019-11-12T14:10:28Z

This broke the build. Please when doing these larger sorts of changes use a full test.

gottesmm · 2019-11-12T14:10:35Z

See: https://ci.swift.org/job/swift-PR-osx/16810/console

dan-zheng · 2019-11-12T15:28:55Z

This broke the build. Please when doing these larger sorts of changes use a full test.

Sorry about that! We'll be sure to use @swift-ci Please test on all PRs in the future.
The error appears to be that SwiftSyntax sources need to be regenerated:

01:37:13 --- 5067 ----
01:37:13 diff -r -x '.*' --context=0 /var/folders/_8/79jmzf2142z2xydc_01btlx00000gn/T/tmpvgMewm/syntax_nodes/SyntaxNodes.swift /Users/buildnode/jenkins/workspace/swift-PR-osx/branch-master/swift-syntax/Sources/SwiftSyntax/gyb_generated/syntax_nodes/SyntaxNodes.swift
01:37:13 *** /var/folders/_8/79jmzf2142z2xydc_01btlx00000gn/T/tmpvgMewm/syntax_nodes/SyntaxNodes.swift	Tue Nov 12 01:37:13 2019
01:37:13 --- /Users/buildnode/jenkins/workspace/swift-PR-osx/branch-master/swift-syntax/Sources/SwiftSyntax/gyb_generated/syntax_nodes/SyntaxNodes.swift	Mon Nov 11 23:53:23 2019
01:37:13 ***************
01:37:13 *** 6904,7844 ****
01:37:13 - // MARK: - DifferentiableAttributeArgumentsSyntax
01:37:13 - 
01:37:13 - /// 
01:37:13 - /// The arguments for the `@differentiable` attribute: an optional          differentiation parameter list and associated functions.
01:37:13 - /// 
01:37:13 - public struct DifferentiableAttributeArgumentsSyntax: SyntaxProtocol, SyntaxHashable {
01:37:13 -   enum Cursor: Int {
01:37:13 -     case diffParams
01:37:13 -     case diffParamsComma
01:37:13 -     case maybePrimal
01:37:13 -     case maybeAdjoint
01:37:13 -     case maybeJVP
01:37:13 -     case maybeVJP
01:37:13 -     case whereClause
01:37:13 -   }
...
01:37:13 FAIL: Gyb-generated files committed to repository do not match generated ones. Please re-generate the gyb-files and recommit them.

I'm taking a look into the issue now. @gottesmm: would you prefer that we revert this PR in the meantime?

EDIT: swift-syntax fix in swiftlang/swift-syntax#178.

…attribute. Friend PR: swiftlang/swift#27506

Update gyb-generated files for `@differentiable` attribute. Friend PR: swiftlang/swift#27506

utils/gyb_syntax_support/AttributeNodes.py

atrick · 2019-11-12T16:39:41Z

@dan-zheng @gottesmm PR size has nothing to do with it. Any functional change that could conceivably be sensitive to build configuration needs to run normal PR testing (swift-ci test). I always run normal testing on a functional change, but a lot of people aren't doing this.

For a major functional change full PR testing would include benchmarks, SCK, and maybe even compiler performance tests.

beccadax · 2019-12-19T02:39:54Z

@dan-zheng I noticed today that this PR includes serialization for @differentiable, but there doesn't seem to be any deserialization code that uses DifferentiableDeclAttrLayout, even now. Is this intentional?

dan-zheng · 2019-12-19T03:10:57Z

@dan-zheng I noticed today that this PR includes serialization for @differentiable, but there doesn't seem to be any deserialization code that uses DifferentiableDeclAttrLayout, even now. Is this intentional?

Thanks for noticing!

@differentiable attribute (de)serialization is currently not fully done because it's ~blocked by @differentiable attribute type-checking, which hasn't yet been upstreamed to master from tensorflow branch. One of the serialized components (IndexSubset * representing parameter indices) is resolved by @differentiable attribute type-checking.

TF-836 tracks @differentiable attribute serialization. There's a NOTE comment in the source code referencing TF-836.

Small aside: @differentiable attribute has a bunch of components and not all of them are serialized. For example:

Not serialized
- Raw parsed parameter indices (ArrayRef<ParsedAutoDiffParameter>)†.
- Raw DeclNameRefs for JVP/VJP derivative function declarations†.
Serialized
- The resolved parameter indices (IndexSubset *). Resolving parameter indices is not particularly cheap, so serializing them seemed sensible.
- The resolved JVP/VJP derivative function declarations (FuncDecl *).

†: Serializing this is necessary to print deserialized @differentiable attributes exactly like how they were written in source code. We didn't prioritize full fidelity printing, so we left this unserialized.

It's worthwhile to discuss which attribute components should be serialized! I started a forum question for discussion.

Precedent: DynamicReplacementAttr is similar to DifferentiableAttr/DerivativeAttr/TransposeAttr in that it also stores a DeclNameRef. For DynamicReplacementAttr, all DeclNameRef components are serialized.

Update gyb-generated files for `@differentiable` attribute. Friend PR: swiftlang/swift#27506

saeta mentioned this pull request Oct 4, 2019

[AutoDiff upstream] Add @transpose attribute. #27545

Merged

dan-zheng reviewed Oct 4, 2019

View reviewed changes

include/swift/AST/Attr.h Outdated Show resolved Hide resolved

bgogul force-pushed the upstream_differentiable_attr branch from 12c07b4 to 0429dbd Compare October 4, 2019 22:59

bgogul force-pushed the upstream_differentiable_attr branch from 0429dbd to a9ac462 Compare October 14, 2019 17:25

bgogul changed the title ~~[WIP] [AutoDiff upstream] Introduce @differentiable attribute to mark functions as differentiable.~~ [AutoDiff upstream] Introduce @differentiable attribute to mark functions as differentiable. Oct 14, 2019

bgogul requested a review from DougGregor October 14, 2019 18:47

rxwei approved these changes Oct 15, 2019

View reviewed changes

bgogul force-pushed the upstream_differentiable_attr branch from 4a150b1 to f10459e Compare October 15, 2019 17:53

bgogul removed the request for review from DougGregor October 24, 2019 22:07

bgogul added 9 commits November 6, 2019 10:10

Code for parseUnsignedInteger.

937e57f

Upstream @differentiable attribute parsing.

374e64d

Serialization of @differentiable attribute.

46ff5f2

Pick up changes from IndexSubset.

3b8498d

Fixup serialization code with changes from IndexSubset.

1b8c088

Remove stale fixme.

da744bf

Update serialization codes.

3613168

Remove OpaqueReturnTypeOfAttributeArguments that got included accid…

c8be97d

…entally.

Remove parseTokenSyntax and parseMatchingTokenSyntax from incorrect m…

608d6db

…erge.

bgogul force-pushed the upstream_differentiable_attr branch from 51c32a8 to 608d6db Compare November 6, 2019 18:37

Reinstatate GenericRequirement in NodeSerializationCodes.py

9038971

dan-zheng requested a review from lattner November 7, 2019 21:38

lattner reviewed Nov 8, 2019

View reviewed changes

lattner approved these changes Nov 8, 2019

View reviewed changes

dan-zheng mentioned this pull request Nov 8, 2019

[AutoDiff] Update @differentiable attribute API/ABI status. #28148

Merged

bgogul added 4 commits November 8, 2019 15:37

ABI Stability Flags.

52bb03c

enum Kind -> Kind

b37ac12

Adds is(WRT|JVP|VJP)Identifier helper function.

95b09d7

Change Tok.isNot(a) && Tok.isNot(b) pattern to Tok.isNot(a, b)

644760d

bgogul commented Nov 8, 2019

View reviewed changes

dan-zheng reviewed Nov 11, 2019

View reviewed changes

include/swift/AST/Attr.def Outdated Show resolved Hide resolved

Make @differentiable ABIStableToAdd

764bf72

bgogul mentioned this pull request Nov 11, 2019

Pull in the upstream @differentiable attribute changes #28198

Merged

rxwei merged commit 5accda2 into swiftlang:master Nov 11, 2019

dan-zheng added a commit to dan-zheng/swift-syntax that referenced this pull request Nov 12, 2019

[AutoDiff upstream] Update gyb-generated files for @differentiable …

d24d271

…attribute. Friend PR: swiftlang/swift#27506

dan-zheng added a commit to dan-zheng/swift-syntax that referenced this pull request Nov 12, 2019

[AutoDiff upstream] Update gyb-generated files.

6d2c80e

Update gyb-generated files for `@differentiable` attribute. Friend PR: swiftlang/swift#27506

dan-zheng mentioned this pull request Nov 12, 2019

[AutoDiff upstream] Update gyb-generated files. swiftlang/swift-syntax#178

Merged

dan-zheng reviewed Nov 12, 2019

View reviewed changes

dan-zheng mentioned this pull request Nov 12, 2019

[AutoDiff upstream] Update @differentiable syntax. #28211

Merged

dan-zheng added a commit to dan-zheng/swift-syntax that referenced this pull request Apr 8, 2020

[AutoDiff upstream] Update gyb-generated files.

546b94d

Update gyb-generated files for `@differentiable` attribute. Friend PR: swiftlang/swift#27506

		@@ -0,0 +1,180 @@
		// RUN: %target-swift-frontend -parse -verify %s

		/// Good

[AutoDiff upstream] Introduce @differentiable attribute to mark functions as differentiable. #27506

[AutoDiff upstream] Introduce @differentiable attribute to mark functions as differentiable. #27506

Uh oh!

Conversation

bgogul commented Oct 2, 2019 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

saeta commented Oct 3, 2019

Uh oh!

rxwei commented Oct 3, 2019 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Uh oh!

rxwei commented Oct 15, 2019

Uh oh!

rxwei commented Oct 15, 2019

Uh oh!

bgogul commented Oct 15, 2019

Uh oh!

bgogul commented Oct 16, 2019

Uh oh!

bgogul commented Oct 16, 2019

Uh oh!

bgogul commented Oct 24, 2019

Uh oh!

bgogul commented Nov 6, 2019

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!

lattner left a comment

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!

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!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

bgogul left a comment

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!

lattner commented Nov 10, 2019

bgogul commented Oct 2, 2019 •

edited

Loading

rxwei commented Oct 3, 2019 •

edited

Loading

dan-zheng commented Nov 12, 2019 •

edited

Loading

dan-zheng commented Dec 19, 2019 •

edited

Loading