[AutoDiff] Fix syntax support for @derivative and @transpose. (#28761)

Add `QualifiedDeclName` syntax node representing qualified declaration names
with an optional base type.

Add syntax support for `@transpose` attribute with qualified declaration.

Fix trailing comma parsing for `DerivativeRegistrationAttributeArguments`
syntax.

Resolves TF-1009.
Enables TF-1058: `@derivative` attribute with qualified declaration.

Author: dan-zheng

lib/Parse/ParseDecl.cpp

@@ -877,7 +877,7 @@ static bool errorAndSkipUntilConsumeRightParen(Parser &P, StringRef attrName,
 bool Parser::parseDifferentiationParametersClause(
     SmallVectorImpl<ParsedAutoDiffParameter> &params, StringRef attrName) {
   SyntaxParsingContext DiffParamsClauseContext(
-       SyntaxContext, SyntaxKind::DifferentiationParamsClause);
+      SyntaxContext, SyntaxKind::DifferentiationParamsClause);
   consumeToken(tok::identifier);
   if (!consumeIf(tok::colon)) {
     diagnose(Tok, diag::expected_colon_after_label, "wrt");
@@ -1207,17 +1207,19 @@ ParserResult<DerivativeAttr> Parser::parseDerivativeAttribute(SourceLoc atLoc,
     }
     {
       // Parse the name of the function.
-      SyntaxParsingContext FuncDeclNameContext(SyntaxContext,
-                                               SyntaxKind::FunctionDeclName);
+      // TODO(TF-1058): Make `@derivative` attribute support qualified
+      // original declarations.
+      SyntaxParsingContext DeclNameContext(SyntaxContext,
+                                           SyntaxKind::QualifiedDeclName);
       // NOTE: Use `afterDot = true` and `allowDeinitAndSubscript = true` to
       // enable, e.g. `@derivative(of: init)` and `@derivative(of: subscript)`.
       original.Name = parseUnqualifiedDeclName(
           /*afterDot*/ true, original.Loc,
           diag::attr_derivative_expected_original_name, /*allowOperators*/ true,
           /*allowZeroArgCompoundNames*/ true, /*allowDeinitAndSubscript*/ true);
-      if (consumeIfTrailingComma())
-        return makeParserError();
     }
+    if (consumeIfTrailingComma())
+      return makeParserError();
     // Parse the optional 'wrt' differentiation parameters clause.
     if (isIdentifier(Tok, "wrt") &&
         parseDifferentiationParametersClause(params, AttrName))
@@ -1269,18 +1271,20 @@ Parser::parseDifferentiatingAttribute(SourceLoc atLoc, SourceLoc loc) {
         SyntaxKind::DeprecatedDerivativeRegistrationAttributeArguments);
     {
       // Parse the name of the function.
-      SyntaxParsingContext FuncDeclNameContext(
-          SyntaxContext, SyntaxKind::FunctionDeclName);
+      // TODO(TF-1058): Make `@differentiating` attribute support qualified
+      // original declarations.
+      SyntaxParsingContext DeclNameContext(SyntaxContext,
+                                           SyntaxKind::QualifiedDeclName);
       // NOTE: Use `afterDot = true` and `allowDeinitAndSubscript = true` to
       // enable, e.g. `@differentiating(init)` and
       // `@differentiating(subscript)`.
       original.Name = parseUnqualifiedDeclName(
           /*afterDot*/ true, original.Loc,
           diag::attr_derivative_expected_original_name, /*allowOperators*/ true,
           /*allowZeroArgCompoundNames*/ true, /*allowDeinitAndSubscript*/ true);
-      if (consumeIfTrailingComma())
-        return makeParserError();
     }
+    if (consumeIfTrailingComma())
+      return makeParserError();
     // Parse the optional 'wrt' differentiation parameters clause.
     if (isIdentifier(Tok, "wrt") &&
         parseDifferentiationParametersClause(params, AttrName))
@@ -1335,6 +1339,8 @@ static bool parseBaseTypeForQualifiedDeclName(Parser &P, TypeRepr *&baseType) {
 /// Returns true on error (if function decl name could not be parsed).
 bool parseQualifiedDeclName(Parser &P, Diag<> nameParseError,
                             TypeRepr *&baseType, DeclNameWithLoc &original) {
+  SyntaxParsingContext DeclNameContext(P.SyntaxContext,
+                                       SyntaxKind::QualifiedDeclName);
   if (parseBaseTypeForQualifiedDeclName(P, baseType))
     return true;
 
@@ -1399,18 +1405,13 @@ ParserResult<TransposeAttr> Parser::parseTransposeAttribute(SourceLoc atLoc,
     }
     {
       // Parse the optionally qualified function name.
-      // TODO(TF-1009): Fix syntax support for dot-separated qualified names.
-      // Currently, `SyntaxKind::FunctionDeclName` only supports unqualified
-      // names.
-      SyntaxParsingContext FuncDeclNameContext(SyntaxContext,
-                                               SyntaxKind::FunctionDeclName);
       if (parseQualifiedDeclName(*this,
                                  diag::attr_transpose_expected_original_name,
                                  baseType, original))
         return makeParserError();
-      if (consumeIfTrailingComma())
-        return makeParserError();
     }
+    if (consumeIfTrailingComma())
+      return makeParserError();
     // Parse the optional 'wrt' transposed parameters clause.
     if (Tok.is(tok::identifier) && Tok.getText() == "wrt" &&
         parseTransposedParametersClause(params, AttrName))

test/Syntax/Outputs/round_trip_parse_gen.swift.withkinds

@@ -601,32 +601,35 @@ func bar<FunctionSignature><ParameterClause>(<FunctionParameter>_ x: <SimpleType
 @differentiable(<DifferentiableAttributeArguments><DifferentiationParamsClause>wrt: <DifferentiationParams>(<DifferentiationParam>self, </DifferentiationParam><DifferentiationParam>x, </DifferentiationParam><DifferentiationParam>y</DifferentiationParam>)</DifferentiationParams></DifferentiationParamsClause>, <DifferentiableAttributeFuncSpecifier>jvp: <FunctionDeclName>bar</FunctionDeclName>, </DifferentiableAttributeFuncSpecifier><DifferentiableAttributeFuncSpecifier>vjp: <FunctionDeclName>foo<DeclNameArguments>(<DeclNameArgument>_:</DeclNameArgument><DeclNameArgument>_:</DeclNameArgument>) </DeclNameArguments></FunctionDeclName></DifferentiableAttributeFuncSpecifier><GenericWhereClause>where <GenericRequirement><ConformanceRequirement><SimpleTypeIdentifier>T </SimpleTypeIdentifier>: <SimpleTypeIdentifier>FloatingPoint</SimpleTypeIdentifier></ConformanceRequirement></GenericRequirement></GenericWhereClause></DifferentiableAttributeArguments>)</Attribute>
 func bar<GenericParameterClause><<GenericParameter>T : <SimpleTypeIdentifier>Numeric</SimpleTypeIdentifier></GenericParameter>></GenericParameterClause><FunctionSignature><ParameterClause>(<FunctionParameter>_ x: <SimpleTypeIdentifier>T</SimpleTypeIdentifier>, </FunctionParameter><FunctionParameter>y: <SimpleTypeIdentifier>T</SimpleTypeIdentifier></FunctionParameter>) </ParameterClause><ReturnClause>-> <SimpleTypeIdentifier>T </SimpleTypeIdentifier></ReturnClause></FunctionSignature><CodeBlock>{ <ReturnStmt>return <IntegerLiteralExpr>1 </IntegerLiteralExpr></ReturnStmt>}</CodeBlock></FunctionDecl><FunctionDecl><Attribute>
 
-@derivative(<DerivativeRegistrationAttributeArguments>of: <FunctionDeclName>-</FunctionDeclName></DerivativeRegistrationAttributeArguments>)</Attribute>
+@derivative(<DerivativeRegistrationAttributeArguments>of: <QualifiedDeclName>-</QualifiedDeclName></DerivativeRegistrationAttributeArguments>)</Attribute>
 func negateDerivative<FunctionSignature><ParameterClause>(<FunctionParameter>_ x: <SimpleTypeIdentifier>Float</SimpleTypeIdentifier></FunctionParameter>)</ParameterClause><ReturnClause>
     -> <TupleType>(<TupleTypeElement>value: <SimpleTypeIdentifier>Float</SimpleTypeIdentifier>, </TupleTypeElement><TupleTypeElement>pullback: <FunctionType>(<TupleTypeElement><SimpleTypeIdentifier>Float</SimpleTypeIdentifier></TupleTypeElement>) -> <SimpleTypeIdentifier>Float</SimpleTypeIdentifier></FunctionType></TupleTypeElement>) </TupleType></ReturnClause></FunctionSignature><CodeBlock>{<ReturnStmt>
   return <TupleExpr>(<TupleExprElement><PrefixOperatorExpr>-<IdentifierExpr>x</IdentifierExpr></PrefixOperatorExpr>, </TupleExprElement><TupleExprElement><ClosureExpr>{ <ClosureSignature><ClosureParam>v </ClosureParam>in </ClosureSignature><PrefixOperatorExpr>-<IdentifierExpr>v </IdentifierExpr></PrefixOperatorExpr>}</ClosureExpr></TupleExprElement>)</TupleExpr></ReturnStmt>
 }</CodeBlock></FunctionDecl><FunctionDecl><Attribute>
 
-@derivative(<DerivativeRegistrationAttributeArguments>of: <FunctionDeclName>baz<DeclNameArguments>(<DeclNameArgument>label:</DeclNameArgument><DeclNameArgument>_:</DeclNameArgument>)</DeclNameArguments></FunctionDeclName></DerivativeRegistrationAttributeArguments>)</Attribute>
+@derivative(<DerivativeRegistrationAttributeArguments>of: <QualifiedDeclName>baz<DeclNameArguments>(<DeclNameArgument>label:</DeclNameArgument><DeclNameArgument>_:</DeclNameArgument>)</DeclNameArguments></QualifiedDeclName>, <DifferentiationParamsClause>wrt: <DifferentiationParams>(<DifferentiationParam>x</DifferentiationParam>)</DifferentiationParams></DifferentiationParamsClause></DerivativeRegistrationAttributeArguments>)</Attribute>
 func bazDerivative<FunctionSignature><ParameterClause>(<FunctionParameter>_ x: <SimpleTypeIdentifier>Float</SimpleTypeIdentifier>, </FunctionParameter><FunctionParameter>y: <SimpleTypeIdentifier>Float</SimpleTypeIdentifier></FunctionParameter>)</ParameterClause><ReturnClause>
-    -> <TupleType>(<TupleTypeElement>value: <SimpleTypeIdentifier>Float</SimpleTypeIdentifier>, </TupleTypeElement><TupleTypeElement>pullback: <FunctionType>(<TupleTypeElement><SimpleTypeIdentifier>Float</SimpleTypeIdentifier></TupleTypeElement>) -> <TupleType>(<TupleTypeElement><SimpleTypeIdentifier>Float</SimpleTypeIdentifier>, </TupleTypeElement><TupleTypeElement><SimpleTypeIdentifier>Float</SimpleTypeIdentifier></TupleTypeElement>)</TupleType></FunctionType></TupleTypeElement>) </TupleType></ReturnClause></FunctionSignature><CodeBlock>{<ReturnStmt>
+    -> <TupleType>(<TupleTypeElement>value: <SimpleTypeIdentifier>Float</SimpleTypeIdentifier>, </TupleTypeElement><TupleTypeElement>pullback: <FunctionType>(<TupleTypeElement><SimpleTypeIdentifier>Float</SimpleTypeIdentifier></TupleTypeElement>) -> <SimpleTypeIdentifier>Float</SimpleTypeIdentifier></FunctionType></TupleTypeElement>) </TupleType></ReturnClause></FunctionSignature><CodeBlock>{<ReturnStmt>
   return <TupleExpr>(<TupleExprElement><IdentifierExpr>x</IdentifierExpr>, </TupleExprElement><TupleExprElement><ClosureExpr>{ <ClosureSignature><ClosureParam>v </ClosureParam>in </ClosureSignature><IdentifierExpr>v </IdentifierExpr>}</ClosureExpr></TupleExprElement>)</TupleExpr></ReturnStmt>
 }</CodeBlock></FunctionDecl><FunctionDecl><Attribute>
 
-@transpose(<DerivativeRegistrationAttributeArguments>of: <FunctionDeclName>+</FunctionDeclName></DerivativeRegistrationAttributeArguments>)</Attribute>
+@transpose(<DerivativeRegistrationAttributeArguments>of: <QualifiedDeclName>+</QualifiedDeclName></DerivativeRegistrationAttributeArguments>)</Attribute>
 func addTranspose<FunctionSignature><ParameterClause>(<FunctionParameter>_ v: <SimpleTypeIdentifier>Float</SimpleTypeIdentifier></FunctionParameter>) </ParameterClause><ReturnClause>-> <TupleType>(<TupleTypeElement><SimpleTypeIdentifier>Float</SimpleTypeIdentifier>, </TupleTypeElement><TupleTypeElement><SimpleTypeIdentifier>Float</SimpleTypeIdentifier></TupleTypeElement>) </TupleType></ReturnClause></FunctionSignature><CodeBlock>{<ReturnStmt>
   return <TupleExpr>(<TupleExprElement><IdentifierExpr>v</IdentifierExpr>, </TupleExprElement><TupleExprElement><IdentifierExpr>v</IdentifierExpr></TupleExprElement>)</TupleExpr></ReturnStmt>
 }</CodeBlock></FunctionDecl><FunctionDecl><Attribute>
 
-@differentiating(<DeprecatedDerivativeRegistrationAttributeArguments><FunctionDeclName>baz<DeclNameArguments>(<DeclNameArgument>label:</DeclNameArgument><DeclNameArgument>_:</DeclNameArgument>)</DeclNameArguments></FunctionDeclName></DeprecatedDerivativeRegistrationAttributeArguments>)</Attribute>
+@differentiating(<DeprecatedDerivativeRegistrationAttributeArguments><QualifiedDeclName>baz<DeclNameArguments>(<DeclNameArgument>label:</DeclNameArgument><DeclNameArgument>_:</DeclNameArgument>)</DeclNameArguments></QualifiedDeclName>, <DifferentiationParamsClause>wrt: <DifferentiationParams>(<DifferentiationParam>x</DifferentiationParam>)</DifferentiationParams></DifferentiationParamsClause></DeprecatedDerivativeRegistrationAttributeArguments>)</Attribute>
 func bazDerivative<FunctionSignature><ParameterClause>(<FunctionParameter>_ x: <SimpleTypeIdentifier>Float</SimpleTypeIdentifier>, </FunctionParameter><FunctionParameter>y: <SimpleTypeIdentifier>Float</SimpleTypeIdentifier></FunctionParameter>)</ParameterClause><ReturnClause>
-    -> <TupleType>(<TupleTypeElement>value: <SimpleTypeIdentifier>Float</SimpleTypeIdentifier>, </TupleTypeElement><TupleTypeElement>pullback: <FunctionType>(<TupleTypeElement><SimpleTypeIdentifier>Float</SimpleTypeIdentifier></TupleTypeElement>) -> <TupleType>(<TupleTypeElement><SimpleTypeIdentifier>Float</SimpleTypeIdentifier>, </TupleTypeElement><TupleTypeElement><SimpleTypeIdentifier>Float</SimpleTypeIdentifier></TupleTypeElement>)</TupleType></FunctionType></TupleTypeElement>) </TupleType></ReturnClause></FunctionSignature><CodeBlock>{<ReturnStmt>
+    -> <TupleType>(<TupleTypeElement>value: <SimpleTypeIdentifier>Float</SimpleTypeIdentifier>, </TupleTypeElement><TupleTypeElement>pullback: <FunctionType>(<TupleTypeElement><SimpleTypeIdentifier>Float</SimpleTypeIdentifier></TupleTypeElement>) -> <SimpleTypeIdentifier>Float</SimpleTypeIdentifier></FunctionType></TupleTypeElement>) </TupleType></ReturnClause></FunctionSignature><CodeBlock>{<ReturnStmt>
   return <TupleExpr>(<TupleExprElement><IdentifierExpr>x</IdentifierExpr>, </TupleExprElement><TupleExprElement><ClosureExpr>{ <ClosureSignature><ClosureParam>v </ClosureParam>in </ClosureSignature><IdentifierExpr>v </IdentifierExpr>}</ClosureExpr></TupleExprElement>)</TupleExpr></ReturnStmt>
-}</CodeBlock></FunctionDecl>
+}</CodeBlock></FunctionDecl><FunctionDecl><Attribute>
 
-// TODO(TF-1009): Add syntax support for dot-separated qualified names in
-// `@transpose(of:)` attributes.
-@transpose(of: <SimpleTypeIdentifier>Float</SimpleTypeIdentifier>.-)
+@transpose(<DerivativeRegistrationAttributeArguments>of: <QualifiedDeclName>-</QualifiedDeclName>, <DifferentiationParamsClause>wrt: <DifferentiationParams>(<DifferentiationParam>0, </DifferentiationParam><DifferentiationParam>1</DifferentiationParam>)</DifferentiationParams></DifferentiationParamsClause></DerivativeRegistrationAttributeArguments>)</Attribute>
+func subtractTranspose<FunctionSignature><ParameterClause>(<FunctionParameter>_ v: <SimpleTypeIdentifier>Float</SimpleTypeIdentifier></FunctionParameter>) </ParameterClause><ReturnClause>-> <TupleType>(<TupleTypeElement><SimpleTypeIdentifier>Float</SimpleTypeIdentifier>, </TupleTypeElement><TupleTypeElement><SimpleTypeIdentifier>Float</SimpleTypeIdentifier></TupleTypeElement>) </TupleType></ReturnClause></FunctionSignature><CodeBlock>{<ReturnStmt>
+  return <TupleExpr>(<TupleExprElement><IdentifierExpr>v</IdentifierExpr>, </TupleExprElement><TupleExprElement><PrefixOperatorExpr>-<IdentifierExpr>v</IdentifierExpr></PrefixOperatorExpr></TupleExprElement>)</TupleExpr></ReturnStmt>
+}</CodeBlock></FunctionDecl><FunctionDecl><Attribute>
+
+@transpose(<DerivativeRegistrationAttributeArguments>of: <QualifiedDeclName><SimpleTypeIdentifier>Float</SimpleTypeIdentifier>.-</QualifiedDeclName>, <DifferentiationParamsClause>wrt: <DifferentiationParam>0</DifferentiationParam></DifferentiationParamsClause></DerivativeRegistrationAttributeArguments>)</Attribute>
 func negateTranspose<FunctionSignature><ParameterClause>(<FunctionParameter>_ v: <SimpleTypeIdentifier>Float</SimpleTypeIdentifier></FunctionParameter>) </ParameterClause><ReturnClause>-> <SimpleTypeIdentifier>Float </SimpleTypeIdentifier></ReturnClause></FunctionSignature><CodeBlock>{<ReturnStmt>
   return <PrefixOperatorExpr>-<IdentifierExpr>v</IdentifierExpr></PrefixOperatorExpr></ReturnStmt>
-}</CodeBlock>
+}</CodeBlock></FunctionDecl>

test/Syntax/round_trip_parse_gen.swift

@@ -607,9 +607,9 @@ func negateDerivative(_ x: Float)
   return (-x, { v in -v })
 }
 
-@derivative(of: baz(label:_:))
+@derivative(of: baz(label:_:), wrt: (x))
 func bazDerivative(_ x: Float, y: Float)
-    -> (value: Float, pullback: (Float) -> (Float, Float)) {
+    -> (value: Float, pullback: (Float) -> Float) {
   return (x, { v in v })
 }
 
@@ -618,15 +618,18 @@ func addTranspose(_ v: Float) -> (Float, Float) {
   return (v, v)
 }
 
-@differentiating(baz(label:_:))
+@differentiating(baz(label:_:), wrt: (x))
 func bazDerivative(_ x: Float, y: Float)
-    -> (value: Float, pullback: (Float) -> (Float, Float)) {
+    -> (value: Float, pullback: (Float) -> Float) {
   return (x, { v in v })
 }
 
-// TODO(TF-1009): Add syntax support for dot-separated qualified names in
-// `@transpose(of:)` attributes.
-@transpose(of: Float.-)
+@transpose(of: -, wrt: (0, 1))
+func subtractTranspose(_ v: Float) -> (Float, Float) {
+  return (v, -v)
+}
+
+@transpose(of: Float.-, wrt: 0)
 func negateTranspose(_ v: Float) -> Float {
   return -v
 }