[AutoDiff] Be able to parse linear argument in @differentiable attribute (#25228)

- Adds a 'linear` argument to '@differentiable' attribute
- 'linear' argument is optional, but when added, needs to be the first argument in the attribute

Author: bartchr808

include/swift/AST/DiagnosticsParse.def

@@ -1498,7 +1498,8 @@ ERROR(attr_differentiable_use_wrt_not_withrespectto,none,
 ERROR(attr_differentiable_missing_label,PointsToFirstBadToken,
       "missing label '%0:' in '@differentiable' attribute", (StringRef))
 ERROR(attr_differentiable_expected_label,none,
-      "expected a function specifier label, e.g. 'wrt:', 'jvp:', or 'vjp:'", ())
+      "expected either 'wrt:' or a function specifier label, e.g. 'jvp:', "
+      "or 'vjp:'", ())
 
 // differentiating
 ERROR(attr_differentiating_expected_original_name,PointsToFirstBadToken,

include/swift/Parse/Parser.h

@@ -954,7 +954,7 @@ class Parser {
 
   /// Parse the arguments inside the @differentiable attribute.
   bool parseDifferentiableAttributeArguments(
-      SmallVectorImpl<ParsedAutoDiffParameter> &params,
+      bool &linear, SmallVectorImpl<ParsedAutoDiffParameter> &params,
       Optional<DeclNameWithLoc> &jvpSpec, Optional<DeclNameWithLoc> &vjpSpec,
       TrailingWhereClause *&whereClause);
 

lib/Parse/ParseDecl.cpp

@@ -834,11 +834,12 @@ Parser::parseDifferentiableAttribute(SourceLoc atLoc, SourceLoc loc) {
   Optional<DeclNameWithLoc> jvpSpec;
   Optional<DeclNameWithLoc> vjpSpec;
   TrailingWhereClause *whereClause = nullptr;
+  bool linear;
 
   // Parse '('.
   if (consumeIf(tok::l_paren, lParenLoc)) {
     // Parse @differentiable attribute arguments.
-    if (parseDifferentiableAttributeArguments(params, jvpSpec, vjpSpec,
+    if (parseDifferentiableAttributeArguments(linear, params, jvpSpec, vjpSpec,
                                               whereClause))
       return makeParserError();
     // Parse ')'.
@@ -932,7 +933,7 @@ bool Parser::parseDifferentiationParametersClause(
 }
 
 bool Parser::parseDifferentiableAttributeArguments(
-    SmallVectorImpl<ParsedAutoDiffParameter> &params,
+    bool &linear, SmallVectorImpl<ParsedAutoDiffParameter> &params,
     Optional<DeclNameWithLoc> &jvpSpec, Optional<DeclNameWithLoc> &vjpSpec,
     TrailingWhereClause *&whereClause) {
   StringRef AttrName = "differentiable";
@@ -956,8 +957,9 @@ bool Parser::parseDifferentiableAttributeArguments(
       diagnose(Tok, diag::unexpected_separator, ",");
       return true;
     }
-    // Check that token after comma is a function specifier label.
-    if (!Tok.is(tok::identifier) || !(Tok.getText() == "jvp" ||
+    // 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")) {
       diagnose(Tok, diag::attr_differentiable_expected_label);
       return true;
@@ -970,7 +972,19 @@ bool Parser::parseDifferentiableAttributeArguments(
   SyntaxParsingContext ContentContext(
       SyntaxContext, SyntaxKind::DifferentiableAttributeArguments);
 
-  // Parse optional differentiation parameters, starting with the 'wrt:' label.
+  // Parse optional differentiation parameters.
+  // Parse 'linear' label (optional).
+  linear = false;
+  if (Tok.is(tok::identifier) && Tok.getText() == "linear") {
+    linear = true;
+    consumeToken(tok::identifier);
+    // If no trailing comma or 'where' clause, terminate parsing arguments.
+    if (Tok.isNot(tok::comma) && Tok.isNot(tok::kw_where))
+      return false;
+    if (consumeIfTrailingComma())
+      return errorAndSkipToEnd();
+  }
+
   // If 'withRespectTo' is used, make the user change it to 'wrt'.
   if (Tok.is(tok::identifier) && Tok.getText() == "withRespectTo") {
     SourceRange withRespectToRange(Tok.getLoc(), peekToken().getLoc());

test/AutoDiff/differentiable_attr_parse.swift

@@ -57,24 +57,39 @@ public func squareRoot() -> Self {
   return lhs
 }
 
+@differentiable(linear) // okay
+func identity(_ x: Float) -> Float {
+  return x
+}
+
+@differentiable(linear, wrt: x) // okay
+func slope2(_ x: Float) -> Float {
+  return 2 * x
+}
+
+@differentiable(linear, wrt: x, vjp: const3) // okay
+func slope3(_ x: Float) -> Float {
+  return 3 * x
+}
+
 /// Bad
 
-@differentiable(3) // expected-error {{expected a function specifier label, e.g. 'wrt:', 'jvp:', or 'vjp:'}}
+@differentiable(3) // expected-error {{expected either 'wrt:' or a function specifier label, e.g. 'jvp:', or 'vjp:'}}
 func bar(_ x: Float, _: Float) -> Float {
   return 1 + x
 }
 
-@differentiable(foo(_:_:)) // expected-error {{expected a function specifier label, e.g. 'wrt:', 'jvp:', or 'vjp:'}}
+@differentiable(foo(_:_:)) // expected-error {{expected either 'wrt:' or a function specifier label, e.g. 'jvp:', or 'vjp:'}}
 func bar(_ x: Float, _: Float) -> Float {
   return 1 + x
 }
 
-@differentiable(vjp: foo(_:_:), 3) // expected-error {{expected a function specifier label, e.g. 'wrt:', 'jvp:', or 'vjp:'}}
+@differentiable(vjp: foo(_:_:), 3) // expected-error {{expected either 'wrt:' or a function specifier label, e.g. 'jvp:', or 'vjp:'}}
 func bar(_ x: Float, _: Float) -> Float {
   return 1 + x
 }
 
-@differentiable(wrt: (x), foo(_:_:)) // expected-error {{expected a function specifier label, e.g. 'wrt:', 'jvp:', or 'vjp:'}}
+@differentiable(wrt: (x), foo(_:_:)) // expected-error {{expected either 'wrt:' or a function specifier label, e.g. 'jvp:', or 'vjp:'}}
 func bar(_ x: Float, _: Float) -> Float {
   return 1 + x
 }
@@ -84,7 +99,7 @@ func bar(_ x: Float, _: Float) -> Float {
   return 1 + x
 }
 
-@differentiable(wrt: x, y) // expected-error {{expected a function specifier label, e.g. 'wrt:', 'jvp:', or 'vjp:'}}
+@differentiable(wrt: x, y) // expected-error {{expected either 'wrt:' or a function specifier label, e.g. 'jvp:', or 'vjp:'}}
 func bar(_ x: Float, _ y: Float) -> Float {
   return 1 + x
 }
@@ -99,7 +114,7 @@ func bar<T : Numeric>(_ x: T, _: T) -> T {
     return 1 + x
 }
 
-@differentiable(,) // expected-error {{expected a function specifier label, e.g. 'wrt:', 'jvp:', or 'vjp:'}}
+@differentiable(,) // expected-error {{expected either 'wrt:' or a function specifier label, e.g. 'jvp:', or 'vjp:'}}
 func bar(_ x: Float, _: Float) -> Float {
   return 1 + x
 }
@@ -113,3 +128,18 @@ func bar(_ x: Float, _: Float) -> Float {
 func bar<T : Numeric>(_ x: T, _: T) -> T {
     return 1 + x
 }
+
+@differentiable(wrt: x, linear) // expected-error {{expected either 'wrt:' or a function specifier label, e.g. 'jvp:', or 'vjp:'}}
+func slope4(_ x: Float) -> Float {
+  return 4 * x
+}
+
+@differentiable(wrt: x, linear, vjp: const5) // expected-error {{expected either 'wrt:' or a function specifier label, e.g. 'jvp:', or 'vjp:'}}
+func slope5(_ x: Float) -> Float {
+  return 5 * x
+}
+
+@differentiable(wrt: x, vjp: const6, linear) // expected-error {{expected either 'wrt:' or a function specifier label, e.g. 'jvp:', or 'vjp:'}}
+func slope5(_ x: Float) -> Float {
+  return 6 * x
+}