Add libSyntax support for #chainableGradient. (#18255)

* Add #chainableGradient differential operator for seedable gradient

* Handle ExprKind::ChainableGradient.

* Add libSyntax support for #chainableGradient.

* Rename GradientExpr to ReverseAutoDiffExpr for consistency with the class hierarchy.

Author: rxwei

lib/Parse/ParseExpr.cpp

@@ -3665,7 +3665,8 @@ ParserResult<Expr> Parser::parseExprTypeOf() {
 ///     '.' [0-9]+
 ///
 ParserResult<Expr> Parser::parseExprGradientBody(ExprKind kind) {
-  SyntaxParsingContext GradientContext(SyntaxContext, SyntaxKind::GradientExpr);
+  SyntaxParsingContext RADEContext(SyntaxContext,
+                                   SyntaxKind::ReverseAutoDiffExpr);
 
   assert(Tok.isAny(tok::pound_gradient, tok::pound_valueAndGradient,
                    tok::pound_chainableGradient));
@@ -3727,7 +3728,7 @@ ParserResult<Expr> Parser::parseExprGradientBody(ExprKind kind) {
     // Function that parses one parameter.
     auto parseParam = [&]() -> bool {
       SyntaxParsingContext DiffParamContext(
-          SyntaxContext, SyntaxKind::GradientExprDiffParam);
+          SyntaxContext, SyntaxKind::ReverseAutoDiffExprParam);
       SourceLoc paramLoc;
       switch (Tok.getKind()) {
       case tok::period_prefix: {
@@ -3757,7 +3758,8 @@ ParserResult<Expr> Parser::parseExprGradientBody(ExprKind kind) {
     while (Tok.isNot(tok::r_paren))
       if (parseParam())
         return errorAndSkipToEnd();
-    SyntaxContext->collectNodesInPlace(SyntaxKind::GradientExprParamList);
+    SyntaxContext->collectNodesInPlace(
+        SyntaxKind::ReverseAutoDiffExprParamList);
   }
   // Parse the closing ')'.
   if (parseToken(tok::r_paren, rParenLoc, diag::expr_expected_rparen, exprName))

test/Syntax/Outputs/round_trip_parse_gen.swift.withkinds

@@ -561,11 +561,12 @@ func bar<GenericParameterClause><<GenericParameter>T : <SimpleTypeIdentifier>Num
 func bar<FunctionSignature><ParameterClause>(<FunctionParameter>_ x: <SimpleTypeIdentifier>Float</SimpleTypeIdentifier>, </FunctionParameter><FunctionParameter>_: <SimpleTypeIdentifier>Float</SimpleTypeIdentifier></FunctionParameter>) </ParameterClause><ReturnClause>-> <SimpleTypeIdentifier>Float </SimpleTypeIdentifier></ReturnClause></FunctionSignature><CodeBlock>{ <ReturnStmt>return <IntegerLiteralExpr>1 </IntegerLiteralExpr></ReturnStmt>}</CodeBlock></FunctionDecl><FunctionDecl><Attribute>
 
 @differentiable(<DifferentiableAttributeArguments>reverse, <DifferentiableAttributeDiffParams>wrt: (<DifferentiableAttributeDiffParam>self, </DifferentiableAttributeDiffParam><DifferentiableAttributeDiffParam><DifferentiationIndexParam>.0</DifferentiationIndexParam>, </DifferentiableAttributeDiffParam><DifferentiableAttributeDiffParam><DifferentiationIndexParam>.1</DifferentiationIndexParam></DifferentiableAttributeDiffParam>), </DifferentiableAttributeDiffParams><DifferentiableAttributeFuncSpecifier>primal: bar, </DifferentiableAttributeFuncSpecifier><DifferentiableAttributeFuncSpecifier>adjoint: foo<DeclNameArguments>(<DeclNameArgument>_:</DeclNameArgument><DeclNameArgument>_:</DeclNameArgument>) </DeclNameArguments></DifferentiableAttributeFuncSpecifier><GenericWhereClause>where <ConformanceRequirement><SimpleTypeIdentifier>T </SimpleTypeIdentifier>: <SimpleTypeIdentifier>FloatingPoint</SimpleTypeIdentifier></ConformanceRequirement></GenericWhereClause></DifferentiableAttributeArguments>)</Attribute>
-func bar<GenericParameterClause><<GenericParameter>T : <SimpleTypeIdentifier>Numeric</SimpleTypeIdentifier></GenericParameter>></GenericParameterClause><FunctionSignature><ParameterClause>(<FunctionParameter>_ x: <SimpleTypeIdentifier>T</SimpleTypeIdentifier>, </FunctionParameter><FunctionParameter>_: <SimpleTypeIdentifier>T</SimpleTypeIdentifier></FunctionParameter>) </ParameterClause><ReturnClause>-> <SimpleTypeIdentifier>T </SimpleTypeIdentifier></ReturnClause></FunctionSignature><CodeBlock>{ <ReturnStmt>return <IntegerLiteralExpr>1 </IntegerLiteralExpr></ReturnStmt>}</CodeBlock></FunctionDecl><GradientExpr>
+func bar<GenericParameterClause><<GenericParameter>T : <SimpleTypeIdentifier>Numeric</SimpleTypeIdentifier></GenericParameter>></GenericParameterClause><FunctionSignature><ParameterClause>(<FunctionParameter>_ x: <SimpleTypeIdentifier>T</SimpleTypeIdentifier>, </FunctionParameter><FunctionParameter>_: <SimpleTypeIdentifier>T</SimpleTypeIdentifier></FunctionParameter>) </ParameterClause><ReturnClause>-> <SimpleTypeIdentifier>T </SimpleTypeIdentifier></ReturnClause></FunctionSignature><CodeBlock>{ <ReturnStmt>return <IntegerLiteralExpr>1 </IntegerLiteralExpr></ReturnStmt>}</CodeBlock></FunctionDecl><ReverseAutoDiffExpr>
 
-#gradient(<IdentifierExpr>foo</IdentifierExpr>)</GradientExpr><GradientExpr>
-#gradient(<IdentifierExpr>foo</IdentifierExpr>, wrt: <GradientExprDiffParam><DifferentiationIndexParam>.0</DifferentiationIndexParam>, </GradientExprDiffParam><GradientExprDiffParam><DifferentiationIndexParam>.1</DifferentiationIndexParam></GradientExprDiffParam>)</GradientExpr><GradientExpr>
-#valueAndGradient(<IdentifierExpr>foo</IdentifierExpr>, wrt: <GradientExprDiffParam><DifferentiationIndexParam>.0</DifferentiationIndexParam>, </GradientExprDiffParam><GradientExprDiffParam><DifferentiationIndexParam>.1</DifferentiationIndexParam></GradientExprDiffParam>)</GradientExpr><AdjointExpr>
+#gradient(<IdentifierExpr>foo</IdentifierExpr>)</ReverseAutoDiffExpr><ReverseAutoDiffExpr>
+#gradient(<IdentifierExpr>foo</IdentifierExpr>, wrt: <ReverseAutoDiffExprParam><DifferentiationIndexParam>.0</DifferentiationIndexParam>, </ReverseAutoDiffExprParam><ReverseAutoDiffExprParam><DifferentiationIndexParam>.1</DifferentiationIndexParam></ReverseAutoDiffExprParam>)</ReverseAutoDiffExpr><ReverseAutoDiffExpr>
+#chainableGradient(<IdentifierExpr>foo</IdentifierExpr>, wrt: <ReverseAutoDiffExprParam><DifferentiationIndexParam>.0</DifferentiationIndexParam>, </ReverseAutoDiffExprParam><ReverseAutoDiffExprParam><DifferentiationIndexParam>.1</DifferentiationIndexParam></ReverseAutoDiffExprParam>)</ReverseAutoDiffExpr><ReverseAutoDiffExpr>
+#valueAndGradient(<IdentifierExpr>foo</IdentifierExpr>, wrt: <ReverseAutoDiffExprParam><DifferentiationIndexParam>.0</DifferentiationIndexParam>, </ReverseAutoDiffExprParam><ReverseAutoDiffExprParam><DifferentiationIndexParam>.1</DifferentiationIndexParam></ReverseAutoDiffExprParam>)</ReverseAutoDiffExpr><AdjointExpr>
 
 #adjoint(+)</AdjointExpr><AdjointExpr>
 #adjoint(foo<DeclNameArguments>(<DeclNameArgument>_:</DeclNameArgument><DeclNameArgument>_:</DeclNameArgument>)</DeclNameArguments>)</AdjointExpr><AdjointExpr>

test/Syntax/round_trip_parse_gen.swift

@@ -565,6 +565,7 @@ func bar<T : Numeric>(_ x: T, _: T) -> T { return 1 }
 
 #gradient(foo)
 #gradient(foo, wrt: .0, .1)
+#chainableGradient(foo, wrt: .0, .1)
 #valueAndGradient(foo, wrt: .0, .1)
 
 #adjoint(+)

utils/gyb_syntax_support/ExprNodes.py

@@ -571,14 +571,14 @@
          ]),
 
     # SWIFT_ENABLE_TENSORFLOW
-    # Expression generalizing #gradient and #valueAndGradient.
-    # e.g. #gradient(foo(_:_:), wrt: .0, .1)
-    Node('GradientExpr', kind='Expr',
+    # e.g. "#gradient(foo(_:_:), wrt: .0, .1)"
+    Node('ReverseAutoDiffExpr', kind='Expr',
          traits=['Parenthesized'],
          children=[
              Child('Identifier', kind='Token',
                    token_choices=[
                        'PoundGradientToken',
+                       'PoundChainableGradientToken',
                        'PoundValueAndGradientToken',
                    ]),
              Child('LeftParen', kind='LeftParenToken'),
@@ -587,19 +587,19 @@
              Child('WrtLabel', kind='IdentifierToken', is_optional=True,
                    text_choices=['wrt']),
              Child('Colon', kind='ColonToken', is_optional=True),
-             Child('DiffParams', kind='GradientExprParamList',
+             Child('DiffParams', kind='ReverseAutoDiffExprParamList',
                    is_optional=True),
              Child('RightParen', kind='RightParenToken'),
          ]),
 
-    # gradient-expr-diff-param-list ->
+    # reverse-autodiff-expr-param-list ->
     #     gradient-expr-diff-param gradient-expr-diff-param-list?
-    Node('GradientExprParamList', kind='SyntaxCollection',
-         element='GradientExprDiffParam'),
+    Node('ReverseAutoDiffExprParamList', kind='SyntaxCollection',
+         element='ReverseAutoDiffExprParam'),
 
-    # gradient-expr-diff-param ->
+    # reverse-autodiff-expr-param ->
     #     differentiation-index-param ','?
-    Node('GradientExprDiffParam', kind='Syntax',
+    Node('ReverseAutoDiffExprParam', kind='Syntax',
          description='''
          A differentiation parameter: a period followed by an unsigned integer \
          (e.g. `.0`).

utils/gyb_syntax_support/Token.py

@@ -125,6 +125,8 @@ def __init__(self, name, text):
           is_keyword=True),
     Token('PoundGradient', 'pound_gradient', text='#gradient',
           is_keyword=True),
+    Token('PoundChainableGradient', 'pound_chainableGradient',
+          text='#chainableGradient', is_keyword=True),
     Token('PoundValueAndGradient', 'pound_valueAndGradient',
           text='#valueAndGradient', is_keyword=True),
     Token('PoundAdjoint', 'pound_adjoint', text='#adjoint',