[AutoDiff] Change @differentiable syntax to use parameter names. (#22000)

* [AutoDiff] Change `@differentiable` syntax to use parameter names.

Use parameter names instead of parameter indices in `@differentiable` attribute `wrt` list.
The old syntax used parameter indices with a prefix dot, e.g. `wrt: (.0, .1)`.
The new syntax is more readable and understandable when good parameter names are chosen.

New syntax:
```
@differentiable(wrt: (self, input))
func applied(to input: Float) -> Output { ... }
```

The internal, lowered representation of parameter indices is unchanged. This patch affects
only AST attribute parsing/print.
Update `wrt` clause usages in stdlib and tests.

Todos:
- Support parameter name `self` in `wrt` clause, probably via escape backticks: ``self``.
- Low priority: add fixits/specific diagnostics for old parameter index style.

* Address comments from @rxwei.

- Don't allocate auxiliary vector.
- Update parameter index related diagnostics.

* Support single `wrt` parameter without parentheses.

Single `wrt` parameters don't require parenthesis.
The following are now valid:
```
@differentiable(wrt: x)
@differentiable(wrt: self)
```

Multiple parameters still require parentheses: `@differentiable(wrt: x, y)`.

* Rewrite stdlib usages of `@differentiable`, update libSyntax support.

Rewrite stdlib usages of `@differentiable` with one differentiation parameter:
`wrt: (self)` -> `wrt: self`.

Note: printing still always prints parentheses around the differentiation
parameter list.

* Fix `test/Syntax/round_trip_parse_gen.swift`.

Author: dan-zheng

include/swift/AST/AutoDiff.h

@@ -25,34 +25,34 @@ namespace swift {
 
 class ParsedAutoDiffParameter {
 public:
-  enum class Kind { Index, Self };
+  enum class Kind { Named, Self };
 
 private:
   SourceLoc Loc;
   Kind Kind;
   union Value {
-    struct { unsigned Index; }; // Index
+    struct { Identifier Name; }; // Index
     struct {};                  // Self
-    Value(unsigned index) : Index(index) {}
+    Value(Identifier name) : Name(name) {}
     Value() {}
   } V;
 
 public:
   ParsedAutoDiffParameter(SourceLoc loc, enum Kind kind, Value value)
     : Loc(loc), Kind(kind), V(value) {}
 
-  static ParsedAutoDiffParameter getIndexParameter(SourceLoc loc,
-                                                   unsigned index) {
-    return { loc, Kind::Index, index };
+  static ParsedAutoDiffParameter getNamedParameter(SourceLoc loc,
+                                                   Identifier name) {
+    return { loc, Kind::Named, name };
   }
 
   static ParsedAutoDiffParameter getSelfParameter(SourceLoc loc) {
     return { loc, Kind::Self, {} };
   }
 
-  unsigned getIndex() const {
-    assert(Kind == Kind::Index);
-    return V.Index;
+  Identifier getName() const {
+    assert(Kind == Kind::Named);
+    return V.Name;
   }
 
   enum Kind getKind() const {
@@ -64,8 +64,8 @@ class ParsedAutoDiffParameter {
   }
 
   bool isEqual(const ParsedAutoDiffParameter &other) const {
-    if (getKind() == other.getKind() && getKind() == Kind::Index)
-      return getIndex() == other.getIndex();
+    if (getKind() == other.getKind() && getKind() == Kind::Named)
+      return getName() == other.getName();
     return getKind() == other.getKind() && getKind() == Kind::Self;
   }
 };

include/swift/AST/DiagnosticsParse.def

@@ -1466,17 +1466,11 @@ ERROR(attr_implements_expected_member_name,PointsToFirstBadToken,
 ERROR(attr_differentiable_expected_function_name,PointsToFirstBadToken,
       "expected a qualified %0 function", (StringRef))
 ERROR(attr_differentiable_expected_parameter_list,PointsToFirstBadToken,
-      "expected a list of parameters to differentiate with respect to, e.g. (.0, w, b)", ())
+      "expected a list of parameters to differentiate with respect to", ())
 ERROR(attr_differentiable_use_wrt_not_withrespectto,none,
       "use 'wrt:' to specify parameters to differentiate with respect to", ())
 ERROR(attr_differentiable_expected_parameter,PointsToFirstBadToken,
-      "expected a parameter, which can be the index of a function parameter with a leading dot (e.g. '.0'), or 'self'", ())
-ERROR(attr_differentiable_parameter_index_must_be_positive_integer,PointsToFirstBadToken,
-      "parameter index must be a positive integer", ())
-ERROR(attr_differentiable_nondifferentiable_function,none,
-      "function '%0' is not differentiable", (StringRef))
-ERROR(attr_differentiable_duplicate_config_option,PointsToFirstBadToken,
-      "duplicated gradient configuration option '%0'", (StringRef))
+      "expected a parameter, which can be a function parameter name or 'self'", ())
 ERROR(attr_differentiable_missing_label,PointsToFirstBadToken,
       "missing label '%0:' in '@differentiable' attribute", (StringRef))
 ERROR(attr_differentiable_expected_colon_after_label,PointsToFirstBadToken,

include/swift/AST/DiagnosticsSema.def

@@ -2665,10 +2665,10 @@ ERROR(differentiable_attr_overload_not_found,none,
       "%0 does not have expected type %1", (DeclName, Type))
 ERROR(differentiable_attr_wrt_self_must_be_first,none,
       "'self' parameter must come first in the parameter list", ())
-ERROR(differentiable_attr_wrt_indices_must_be_ascending,none,
-      "parameter indices must be ascending", ())
-ERROR(differentiable_attr_wrt_index_out_of_bounds,none,
-      "parameter index out of bounds", ())
+ERROR(differentiable_attr_wrt_names_not_original_order,none,
+      "parameter names must be specified in original order", ())
+ERROR(differentiable_attr_wrt_name_unknown,none,
+      "unknown parameter name %0", (Identifier))
 ERROR(differentiable_attr_wrt_self_instance_method_only,none,
       "'self' parameter is only applicable to instance methods", ())
 ERROR(differentiable_attr_cannot_diff_wrt_objects_or_existentials,none,

lib/AST/Attr.cpp

@@ -21,6 +21,8 @@
 #include "swift/AST/GenericEnvironment.h"
 #include "swift/AST/Module.h"
 #include "swift/AST/Types.h"
+// SWIFT_ENABLE_TENSORFLOW
+#include "swift/AST/ParameterList.h"
 #include "swift/Basic/Defer.h"
 #include "llvm/ADT/SmallString.h"
 #include "llvm/Support/raw_ostream.h"
@@ -580,16 +582,16 @@ bool DeclAttribute::printImpl(ASTPrinter &Printer, const PrintOptions &Options,
         if (isProperty || (isMethod && index == indices->parameters.size() - 1))
           Printer << "self";
         else
-          Printer << "." << index;
+          Printer << original->getParameters()->get(index)->getName().str();
       }, [&] { Printer << ", "; });
       Printer << ")";
     } else if (!parsedParams.empty()) {
       printCommaIfNecessary();
       Printer << "wrt: (";
       interleave(parsedParams, [&](const ParsedAutoDiffParameter &param) {
         switch (param.getKind()) {
-        case ParsedAutoDiffParameter::Kind::Index:
-          Printer << '.' << param.getIndex();
+        case ParsedAutoDiffParameter::Kind::Named:
+          Printer << '.' << param.getName();
           break;
         case ParsedAutoDiffParameter::Kind::Self:
           Printer << "self";

lib/Parse/ParseDecl.cpp

@@ -911,37 +911,29 @@ bool Parser::parseDifferentiableAttributeArguments(
     return errorAndSkipToEnd();
   }
   if (Tok.is(tok::identifier) && Tok.getText() == "wrt") {
-    SyntaxParsingContext DiffParamsContext(
-        SyntaxContext, SyntaxKind::DifferentiableAttributeDiffParams);
+    SyntaxParsingContext DiffParamsClauseContext(
+        SyntaxContext, SyntaxKind::DifferentiableAttributeDiffParamsClause);
     consumeToken(tok::identifier);
     if (!consumeIf(tok::colon)) {
       diagnose(Tok, diag::attr_differentiable_expected_colon_after_label,
                "wrt");
       return errorAndSkipToEnd();
     }
-    SourceLoc leftLoc;
-    if (parseToken(tok::l_paren, leftLoc,
-                   diag::attr_differentiable_expected_parameter_list)) {
-      return errorAndSkipToEnd();
-    }
 
     // Function that parses a parameter into `params`. Returns true if error
     // occurred.
-    auto parseParam = [&]() -> bool {
+    auto parseParam = [&](bool parseTrailingComma = true) -> bool {
       SyntaxParsingContext DiffParamContext(
           SyntaxContext, SyntaxKind::DifferentiableAttributeDiffParam);
       SourceLoc paramLoc;
       switch (Tok.getKind()) {
-      case tok::period_prefix: {
-        SyntaxParsingContext IndexParamContext(
-            SyntaxContext, SyntaxKind::DifferentiationIndexParam);
-        consumeToken(tok::period_prefix);
-        unsigned index;
-        if (parseUnsignedInteger(index, paramLoc,
-                                 diag::attr_differentiable_expected_parameter))
+      case tok::identifier: {
+        Identifier paramName;
+        if (parseIdentifier(paramName, paramLoc,
+                            diag::attr_differentiable_expected_parameter))
           return true;
-        params.push_back(
-          ParsedAutoDiffParameter::getIndexParameter(paramLoc, index));
+        params.push_back(ParsedAutoDiffParameter::getNamedParameter(
+            paramLoc, paramName));
         break;
       }
       case tok::kw_self: {
@@ -953,24 +945,34 @@ bool Parser::parseDifferentiableAttributeArguments(
         diagnose(Tok, diag::attr_differentiable_expected_parameter);
         return true;
       }
-      if (Tok.isNot(tok::r_paren))
+      if (parseTrailingComma && Tok.isNot(tok::r_paren))
         return parseToken(tok::comma, diag::attr_expected_comma, AttrName,
                           /*isDeclModifier=*/false);
       return false;
     };
 
-    // Parse first parameter. At least one is required.
-    if (parseParam())
-      return errorAndSkipToEnd(2);
-    // Parse remaining parameters until ')'.
-    while (Tok.isNot(tok::r_paren))
+    // Parse opening '(' of the parameter list.
+    if (Tok.is(tok::l_paren)) {
+      SyntaxParsingContext DiffParamsContext(
+          SyntaxContext, SyntaxKind::DifferentiableAttributeDiffParams);
+      consumeToken(tok::l_paren);
+      // Parse first parameter. At least one is required.
       if (parseParam())
         return errorAndSkipToEnd(2);
-
-    SyntaxContext->collectNodesInPlace(
-        SyntaxKind::DifferentiableAttributeDiffParamList);
-    // Parse closing ')' of the parameter list.
-    consumeToken(tok::r_paren);
+      // Parse remaining parameters until ')'.
+      while (Tok.isNot(tok::r_paren))
+        if (parseParam())
+          return errorAndSkipToEnd(2);
+      SyntaxContext->collectNodesInPlace(
+          SyntaxKind::DifferentiableAttributeDiffParamList);
+      // Parse closing ')' of the parameter list.
+      consumeToken(tok::r_paren);
+    }
+    // If no opening '(' for parameter list, parse a single parameter.
+    else {
+      if (parseParam(/*parseTrailingComma*/ false))
+        return errorAndSkipToEnd();
+    }
     // If no trailing comma or 'where' clause, terminate parsing arguments.
     if (Tok.isNot(tok::comma) && Tok.isNot(tok::kw_where))
       return false;

lib/Sema/TypeCheckAttr.cpp

@@ -2380,17 +2380,22 @@ void AttributeChecker::visitDifferentiableAttr(DifferentiableAttr *attr) {
     for (unsigned i : indices(parsedWrtParams)) {
       auto paramLoc = parsedWrtParams[i].getLoc();
       switch (parsedWrtParams[i].getKind()) {
-      case ParsedAutoDiffParameter::Kind::Index: {
-        unsigned index = parsedWrtParams[i].getIndex();
-        if ((int)index <= lastIndex) {
-          TC.diagnose(paramLoc,
-                      diag::differentiable_attr_wrt_indices_must_be_ascending);
+      case ParsedAutoDiffParameter::Kind::Named: {
+        auto nameIter =
+            llvm::find_if(originalParams.getArray(), [&](ParamDecl *param) {
+              return param->getName() == parsedWrtParams[i].getName();
+            });
+        // Parameter name must exist.
+        if (nameIter == originalParams.end()) {
+          TC.diagnose(paramLoc, diag::differentiable_attr_wrt_name_unknown,
+                      parsedWrtParams[i].getName());
           return;
         }
-        // Parameter index cannot exceed bounds.
-        if (index >= originalParams.size()) {
+        // Parameter names must be specified in the original order.
+        unsigned index = std::distance(originalParams.begin(), nameIter);
+        if ((int)index <= lastIndex) {
           TC.diagnose(paramLoc,
-                      diag::differentiable_attr_wrt_index_out_of_bounds);
+                      diag::differentiable_attr_wrt_names_not_original_order);
           return;
         }
         autoDiffParameterIndicesBuilder.setParameter(index);

stdlib/public/TensorFlow/Gradients.swift

@@ -617,7 +617,7 @@ extension Tensor where Scalar : Differentiable & FloatingPoint {
   /// TensorFlow builtin conv2d gradient helper for the input.
   @inlinable
   @differentiable(
-    wrt: (.1, .2),
+    wrt: (filter, backpropOutput),
     vjp: _vjpTFConv2DBackpropInput(_:_:_:_:_:)
   )
   func _TFConv2DBackpropInput(
@@ -638,7 +638,7 @@ extension Tensor where Scalar : Differentiable & FloatingPoint {
   /// TensorFlow builtin conv2d gradient helper for the filter.
   @inlinable
   @differentiable(
-    wrt: (.0, .2),
+    wrt: (input, backpropOutput),
     vjp: _vjpTFConv2DBackpropFilter(_:_:_:_:_:)
   )
   func _TFConv2DBackpropFilter(