[AutoDiff] Add #chainableGradient differential operator for seedable gradient (#18250)

rxwei · web-flow · commit 0c660553017d · 2018-07-26T03:54:28.000-07:00
* Add #chainableGradient differential operator for seedable gradient

* Handle ExprKind::ChainableGradient.
diff --git a/include/swift/AST/Expr.h b/include/swift/AST/Expr.h
@@ -3818,7 +3818,8 @@ class DynamicTypeExpr : public Expr {
 };
 
 /// SWIFT_ENABLE_TENSORFLOW
-/// Base class for #gradient and #valueAndGradient expressions.
+/// Base class for differential operators, such as `#gradient`,
+/// `#chainableGradient`, and `#valueAndGradient`.
 class ReverseAutoDiffExpr : public Expr {
 public:
   Expr *getOriginalExpr() const {
@@ -3894,6 +3895,36 @@ class GradientExpr : public ReverseAutoDiffExpr {
     : ReverseAutoDiffExpr(ExprKind::Gradient, loc, lParenLoc, originalExpr,
                           params, rParenLoc) {}
 };
+  
+/// Chainable gradient expression - An expression that produces the
+/// automatically differentiated function that computes the gradient (or
+/// vector-Jacobian products) with respect to specified parameters, taking an
+/// extra result-typed argument representing the seed, i.e. the backpropagated
+/// adjoint.
+/// Examples:
+///   #chainableGradient(baz)
+///   #chainableGradient(bar, wrt: .0, .1)
+///   #chainableGradient(foo(_:_:), wrt: .0)
+///
+class ChainableGradientExpr : public ReverseAutoDiffExpr {
+public:
+  static ChainableGradientExpr *create(ASTContext &ctx, SourceLoc loc,
+                                       SourceLoc lParenLoc, Expr *originalExpr,
+                                   ArrayRef<AutoDiffIndexParameter> parameters,
+                                       SourceLoc rParenLoc);
+  
+  static bool classof(const Expr *E) {
+    return E->getKind() == ExprKind::ChainableGradient;
+  }
+  
+private:
+  explicit ChainableGradientExpr(SourceLoc loc, SourceLoc lParenLoc,
+                                 Expr *originalExpr,
+                                 ArrayRef<AutoDiffIndexParameter> params,
+                                 SourceLoc rParenLoc)
+  : ReverseAutoDiffExpr(ExprKind::ChainableGradient, loc, lParenLoc,
+                        originalExpr, params, rParenLoc) {}
+};
 
 /// ValueAndGradient expression - An expression that produces an automatically
 /// differentiated function that returns the result of the original function and
diff --git a/include/swift/AST/ExprNodes.def b/include/swift/AST/ExprNodes.def
@@ -190,6 +190,7 @@ UNCHECKED_EXPR(KeyPathDot, Expr)
 // SWIFT_ENABLE_TENSORFLOW
 ABSTRACT_EXPR(ReverseAutoDiff, Expr)
   EXPR(Gradient, ReverseAutoDiffExpr)
+  EXPR(ChainableGradient, ReverseAutoDiffExpr)
   EXPR(ValueAndGradient, ReverseAutoDiffExpr)
 EXPR(Adjoint, Expr)
 EXPR(PoundAssert, Expr)
diff --git a/include/swift/Syntax/TokenKinds.def b/include/swift/Syntax/TokenKinds.def
@@ -286,6 +286,7 @@ POUND_OBJECT_LITERAL(tfop, "tfop",  ExpressibleByTensorFlowOp)
 
 // SWIFT_ENABLE_TENSORFLOW
 POUND_KEYWORD(gradient)
+POUND_KEYWORD(chainableGradient)
 POUND_KEYWORD(valueAndGradient)
 POUND_KEYWORD(adjoint)
 POUND_KEYWORD(assert)
diff --git a/lib/AST/ASTDumper.cpp b/lib/AST/ASTDumper.cpp
@@ -1857,6 +1857,11 @@ class PrintExpr : public ExprVisitor<PrintExpr> {
     printCommon(E, "gradient_expr");
     printReverseAutoDiffExpr(E);
   }
+  
+  void visitChainableGradientExpr(ChainableGradientExpr *E) {
+    printCommon(E, "chainable_gradient_expr");
+    printReverseAutoDiffExpr(E);
+  }
 
   void visitValueAndGradientExpr(ValueAndGradientExpr *E) {
     printCommon(E, "value_and_gradient_expr");
diff --git a/lib/AST/Expr.cpp b/lib/AST/Expr.cpp
@@ -358,6 +358,7 @@ ConcreteDeclRef Expr::getReferencedDecl() const {
   NO_REFERENCE(KeyPathDot);
   // SWIFT_ENABLE_TENSORFLOW
   NO_REFERENCE(Gradient);
+  NO_REFERENCE(ChainableGradient);
   NO_REFERENCE(ValueAndGradient);
   NO_REFERENCE(Adjoint);
   NO_REFERENCE(PoundAssert);
@@ -532,6 +533,7 @@ bool Expr::canAppendPostfixExpression(bool appendingPostfixOperator) const {
   case ExprKind::KeyPath:
   // SWIFT_ENABLE_TENSORFLOW
   case ExprKind::Gradient:
+  case ExprKind::ChainableGradient:
   case ExprKind::ValueAndGradient:
   case ExprKind::Adjoint:
     return true;
@@ -1142,21 +1144,34 @@ GradientExpr *GradientExpr::create(ASTContext &ctx, SourceLoc loc,
                                    SourceLoc rParenLoc) {
   unsigned numParams = parameters.size();
   unsigned size =
-    sizeof(GradientExpr) + numParams * sizeof(AutoDiffIndexParameter);
+      sizeof(GradientExpr) + numParams * sizeof(AutoDiffIndexParameter);
   void *memory = ctx.Allocate(size, alignof(GradientExpr));
   return new (memory) GradientExpr(loc, lParenLoc, originalExpr, parameters,
                                    rParenLoc);
 }
 
 
+ChainableGradientExpr *
+ChainableGradientExpr::create(ASTContext &ctx, SourceLoc loc,
+                              SourceLoc lParenLoc, Expr *originalExpr,
+                              ArrayRef<AutoDiffIndexParameter> parameters,
+                              SourceLoc rParenLoc) {
+  unsigned numParams = parameters.size();
+  unsigned size = sizeof(ChainableGradientExpr)
+      + numParams * sizeof(AutoDiffIndexParameter);
+  void *memory = ctx.Allocate(size, alignof(ChainableGradientExpr));
+  return new (memory) ChainableGradientExpr(loc, lParenLoc, originalExpr,
+                                            parameters, rParenLoc);
+}
+
 ValueAndGradientExpr *
 ValueAndGradientExpr::create(ASTContext &ctx, SourceLoc loc,
                              SourceLoc lParenLoc, Expr *originalExpr,
                              ArrayRef<AutoDiffIndexParameter> parameters,
                              SourceLoc rParenLoc) {
   unsigned numParams = parameters.size();
   unsigned size =
-    sizeof(ValueAndGradientExpr) + numParams * sizeof(AutoDiffIndexParameter);
+      sizeof(ValueAndGradientExpr) + numParams * sizeof(AutoDiffIndexParameter);
   void *memory = ctx.Allocate(size, alignof(ValueAndGradientExpr));
   return new (memory) ValueAndGradientExpr(loc, lParenLoc, originalExpr,
                                            parameters, rParenLoc);
diff --git a/lib/Parse/ParseExpr.cpp b/lib/Parse/ParseExpr.cpp
@@ -1764,6 +1764,10 @@ ParserResult<Expr> Parser::parseExprPrimary(Diag<> ID, bool isExprBasic) {
   case tok::pound_gradient:
     return parseExprGradientBody(ExprKind::Gradient);
     break;
+      
+  case tok::pound_chainableGradient:
+    return parseExprGradientBody(ExprKind::ChainableGradient);
+    break;
 
   case tok::pound_valueAndGradient:
     return parseExprGradientBody(ExprKind::ValueAndGradient);
@@ -3663,7 +3667,8 @@ ParserResult<Expr> Parser::parseExprTypeOf() {
 ParserResult<Expr> Parser::parseExprGradientBody(ExprKind kind) {
   SyntaxParsingContext GradientContext(SyntaxContext, SyntaxKind::GradientExpr);
 
-  assert(Tok.is(tok::pound_gradient) || Tok.is(tok::pound_valueAndGradient));
+  assert(Tok.isAny(tok::pound_gradient, tok::pound_valueAndGradient,
+                   tok::pound_chainableGradient));
   auto poundGradLoc = consumeToken();
   SourceLoc lParenLoc;
   SourceLoc rParenLoc;
@@ -3673,6 +3678,9 @@ ParserResult<Expr> Parser::parseExprGradientBody(ExprKind kind) {
   case ExprKind::Gradient:
     exprName = "#gradient";
     break;
+  case ExprKind::ChainableGradient:
+    exprName = "#chainableGradient";
+    break;
   case ExprKind::ValueAndGradient:
     exprName = "#valueAndGradient";
     break;
@@ -3764,6 +3772,11 @@ ParserResult<Expr> Parser::parseExprGradientBody(ExprKind kind) {
                                   originalFnParseResult.get(), params,
                                   rParenLoc);
     break;
+  case ExprKind::ChainableGradient:
+    result = ChainableGradientExpr::create(Context, poundGradLoc, lParenLoc,
+                                           originalFnParseResult.get(), params,
+                                           rParenLoc);
+    break;
   case ExprKind::ValueAndGradient:
     result = ValueAndGradientExpr::create(Context, poundGradLoc, lParenLoc,
                                           originalFnParseResult.get(), params,
diff --git a/lib/SILGen/SILGenExpr.cpp b/lib/SILGen/SILGenExpr.cpp
@@ -480,6 +480,7 @@ namespace {
                                               SGFContext C);
     // SWIFT_ENABLE_TENSORFLOW
     RValue visitGradientExpr(GradientExpr *E, SGFContext C);
+    RValue visitChainableGradientExpr(ChainableGradientExpr *E, SGFContext C);
     RValue visitValueAndGradientExpr(ValueAndGradientExpr *E, SGFContext C);
     RValue visitAdjointExpr(AdjointExpr *E, SGFContext C);
     RValue visitObjectLiteralExpr(ObjectLiteralExpr *E, SGFContext C);
@@ -2824,6 +2825,11 @@ visitGradientExpr(GradientExpr *E, SGFContext C) {
   return emitGradientInst(*this, C, E);
 }
 
+RValue RValueEmitter::
+visitChainableGradientExpr(ChainableGradientExpr *E, SGFContext C) {
+  return emitGradientInst(*this, C, E, SILGradientFlags::Seedable);
+}
+
 RValue RValueEmitter::
 visitValueAndGradientExpr(ValueAndGradientExpr *E, SGFContext C) {
   return emitGradientInst(*this, C, E, SILGradientFlags::PreservingResult);
diff --git a/lib/Sema/CSApply.cpp b/lib/Sema/CSApply.cpp
@@ -2439,6 +2439,10 @@ namespace {
     Expr *visitGradientExpr(GradientExpr *expr) {
       return handleReverseAutoDiffExpr(expr, /*preservingOriginalResult=*/false);
     }
+    
+    Expr *visitChainableGradientExpr(ChainableGradientExpr *expr) {
+      llvm_unreachable("Unhandled");
+    }
 
     Expr *visitValueAndGradientExpr(ValueAndGradientExpr *expr) {
       return handleReverseAutoDiffExpr(expr, /*preservingOriginalResult=*/true);
diff --git a/lib/Sema/CSDiag.cpp b/lib/Sema/CSDiag.cpp
@@ -1091,8 +1091,7 @@ class FailureDiagnosis :public ASTVisitor<FailureDiagnosis, /*exprresult*/bool>{
   bool diagnoseSubscriptErrors(SubscriptExpr *SE, bool performingSet);
 
   // SWIFT_ENABLE_TENSORFLOW
-  bool diagnoseReverseAutoDiffExpr(ReverseAutoDiffExpr *GE,
-                                   bool preservingPrimalResult);
+  bool diagnoseReverseAutoDiffExpr(ReverseAutoDiffExpr *GE);
 
   /// Diagnose the usage of 'subscript' instead of the operator when calling
   /// a subscript and offer a fixit if the inputs are compatible.
@@ -1123,8 +1122,7 @@ class FailureDiagnosis :public ASTVisitor<FailureDiagnosis, /*exprresult*/bool>{
   bool visitClosureExpr(ClosureExpr *CE);
   bool visitKeyPathExpr(KeyPathExpr *KPE);
   // SWIFT_ENABLE_TENSORFLOW
-  bool visitGradientExpr(GradientExpr *GE);
-  bool visitValueAndGradientExpr(ValueAndGradientExpr *GE);
+  bool visitReverseAutoDiffExpr(ReverseAutoDiffExpr *RADE);
   bool visitPoundAssertExpr(PoundAssertExpr *PAE);
 };
 } // end anonymous namespace
@@ -7471,11 +7469,10 @@ bool FailureDiagnosis::visitKeyPathExpr(KeyPathExpr *KPE) {
 
 // SWIFT_ENABLE_TENSORFLOW
 bool FailureDiagnosis::
-diagnoseReverseAutoDiffExpr(ReverseAutoDiffExpr *GE,
-                            bool preservingPrimalResult) {
+diagnoseReverseAutoDiffExpr(ReverseAutoDiffExpr *RADE) {
   // TODO: Sema diagnostics for gradient expressions could be improved by
   // diagnosing non-differentiable arguments/non-differentiable constraints.
-  auto gradType = CS.getType(GE);
+  auto gradType = CS.getType(RADE);
   auto gradFnType = gradType->getAs<AnyFunctionType>();
   assert(gradFnType && "Gradient expression should have function type.");
 
@@ -7486,20 +7483,16 @@ diagnoseReverseAutoDiffExpr(ReverseAutoDiffExpr *GE,
   // If gradient expression has a generic primal, then conversion to the
   // contextual type was not possible.
   if (gradType->hasTypeVariable()) {
-    diagnose(GE->getLoc(), diag::gradient_expr_incompatible_contextual_type,
+    diagnose(RADE->getLoc(), diag::gradient_expr_incompatible_contextual_type,
              contextualType);
     return true;
   }
 
   return false;
 }
 
-bool FailureDiagnosis::visitGradientExpr(GradientExpr *GE) {
-  return diagnoseReverseAutoDiffExpr(GE, /*preservingPrimalResult=*/false);
-}
-
-bool FailureDiagnosis::visitValueAndGradientExpr(ValueAndGradientExpr *GE) {
-  return diagnoseReverseAutoDiffExpr(GE, /*preservingPrimalResult=*/true);
+bool FailureDiagnosis::visitReverseAutoDiffExpr(ReverseAutoDiffExpr *RADE) {
+  return diagnoseReverseAutoDiffExpr(RADE);
 }
 
 bool FailureDiagnosis::visitPoundAssertExpr(PoundAssertExpr *PAE) {
diff --git a/lib/Sema/CSGen.cpp b/lib/Sema/CSGen.cpp
@@ -1391,6 +1391,10 @@ namespace {
     Type visitGradientExpr(GradientExpr *GE) {
       return handleReverseAutoDiffExpr(GE, /*preservingOriginalResult=*/false);
     }
+    
+    Type visitChainableGradientExpr(ChainableGradientExpr *CGE) {
+      llvm_unreachable("Unhandled");
+    }
 
     Type visitValueAndGradientExpr(ValueAndGradientExpr *VGE) {
       return handleReverseAutoDiffExpr(VGE, /*preservingOriginalResult=*/true);
diff --git a/test/AutoDiff/gradient_expr_parse.swift b/test/AutoDiff/gradient_expr_parse.swift
@@ -8,6 +8,8 @@
 #gradient(foo, wrt: .0, self) // expected-error {{expected a parameter, which must be }}
 #gradient(foo, wrt: .0, .1) // okay
 
+#chainableGradient(foo, wrt: .0, .1) // okay
+
 #valueAndGradient(foo, wrt: .0, .1) // okay
 
 #adjoint(foo(_:_:)) // okay
