[AutoDiff] [SIL] Tweak 'differentiable_function' syntax.

docs/SIL.rst

@@ -5611,28 +5611,29 @@ differentiable_function
                       sil-differentiable-function-derivative-functions-clause?
 
   sil-differentiable-function-parameter-indices ::=
-      '[' 'wrt' [0-9]+ (' ' [0-9]+)* ']'
+      '[' 'parameters' [0-9]+ (' ' [0-9]+)* ']'
   sil-differentiable-derivative-functions-clause ::=
-      'with' '{' sil-value ':' sil-type ',' sil-value ':' sil-type '}'
+      'with_derivative'
+      '{' sil-value ':' sil-type ',' sil-value ':' sil-type '}'
 
-  differentiable_function [wrt 0] %0 : $(T) -> T \
-    with {%1 : $(T) -> (T, (T) -> T), %2 : $(T) -> (T, (T) -> T)}
+  differentiable_function [parameters 0] %0 : $(T) -> T \
+    with_derivative {%1 : $(T) -> (T, (T) -> T), %2 : $(T) -> (T, (T) -> T)}
 
 Bundles a function with its derivative functions into a ``@differentiable``
 function. There are two derivative functions: a Jacobian-vector products (JVP)
 function and a vector-Jacobian products (VJP) function.
 
-``[wrt ...]`` specifies parameter indices that the original function is
+``[parameters ...]`` specifies parameter indices that the original function is
 differentiable with respect to. When not specified, it defaults to all
 parameters.
 
-A ``with`` clause specifies the differentiation functions associated
-with the original function. When a ``with`` clause is not specified, the first
-operand will be differentiated to produce derivative functions, and a ``with``
-clause will be added to the instruction.
+A ``with_derivative`` clause specifies the differentiation functions associated
+with the original function. When a ``with_derivative`` clause is not specified,
+the first operand will be differentiated to produce derivative functions, and a
+``with_derivative`` clause will be added to the instruction.
 
-In raw SIL, it is optional to provide a derivative function ``with`` clause.
-In canonical SIL, a ``with`` clause is mandatory.
+In raw SIL, it is optional to provide a derivative function ``with_derivative``
+clause. In canonical SIL, a ``with_derivative`` clause is mandatory.
 
 
 linear_function

lib/ParseSIL/ParseSIL.cpp

@@ -2922,17 +2922,17 @@ bool SILParser::parseSILInstruction(SILBuilder &B) {
 
   // SWIFT_ENABLE_TENSORFLOW
   case SILInstructionKind::DifferentiableFunctionInst: {
-    // e.g. differentiable_function [wrt 0 1 2] %0 : $T
+    // e.g. differentiable_function [parameters 0 1 2] %0 : $T
     //
-    // e.g. differentiable_function [wrt 0 1 2] %0 : $T with
+    // e.g. differentiable_function [parameters 0 1 2] %0 : $T with_derivative
     //      {%1 : $T, %2 : $T}
     //        ^ jvp    ^ vjp
     SourceLoc lastLoc;
     SmallVector<unsigned, 8> parameterIndices;
-    // Parse optional `[wrt <integer_literal>...]`
+    // Parse optional `[parameters <integer_literal>...]`
     if (P.Tok.is(tok::l_square) &&
         P.peekToken().is(tok::identifier) &&
-        P.peekToken().getText() == "wrt") {
+        P.peekToken().getText() == "parameters") {
       P.consumeToken(tok::l_square);
       P.consumeToken(tok::identifier);
       // Parse indices.
@@ -2960,8 +2960,9 @@ bool SILParser::parseSILInstruction(SILBuilder &B) {
       return true;
     }
     Optional<std::pair<SILValue, SILValue>> derivativeFunctions = None;
-    // Parse an optional operand list `with { <operand> , <operand> }`.
-    if (P.Tok.is(tok::identifier) && P.Tok.getText() == "with") {
+    // Parse an optional operand list
+    //   `with_derivative { <operand> , <operand> }`.
+    if (P.Tok.is(tok::identifier) && P.Tok.getText() == "with_derivative") {
       P.consumeToken(tok::identifier);
       // Parse derivative function values as an operand list.
       // FIXME(rxwei): Change this to *not* require a type signature once

lib/SIL/SILPrinter.cpp

@@ -1163,14 +1163,14 @@ class SILPrinter : public SILInstructionVisitor<SILPrinter> {
   // SWIFT_ENABLE_TENSORFLOW
   void visitDifferentiableFunctionInst(DifferentiableFunctionInst *dfi) {
     if (!dfi->getParameterIndices()->isEmpty()) {
-      *this << "[wrt";
+      *this << "[parameters";
       for (auto i : dfi->getParameterIndices()->getIndices())
         *this << ' ' << i;
       *this << "] ";
     }
     *this << getIDAndType(dfi->getOriginalFunction());
     if (dfi->hasDerivativeFunctions()) {
-      *this << " with ";
+      *this << " with_derivative ";
       *this << '{' << getIDAndType(dfi->getJVPFunction()) << ", "
             << getIDAndType(dfi->getVJPFunction()) << '}';
     }

test/AutoDiff/differentiable_function_inst.sil

@@ -17,20 +17,20 @@ sil @examplemethod : $@convention(method) (Float, Float, Float) -> Float
 sil @test : $@convention(thin) () -> () {
 bb0:
   %0 = function_ref @examplefunc : $@convention(thin) (Float, Float, Float) -> Float
-  %1 = differentiable_function [wrt 0 1 2] %0 : $@convention(thin) (Float, Float, Float) -> Float
+  %1 = differentiable_function [parameters 0 1 2] %0 : $@convention(thin) (Float, Float, Float) -> Float
 
   // CHECK: %2 = differentiable_function_extract [vjp] %1 : $@differentiable @convention(thin) (Float, Float, Float) -> Float
   %2 = differentiable_function_extract [vjp] %1 : $@differentiable @convention(thin) (Float, Float, Float) -> Float
-  %3 = differentiable_function [wrt 0] %0 : $@convention(thin) (Float, Float, Float) -> Float
+  %3 = differentiable_function [parameters 0] %0 : $@convention(thin) (Float, Float, Float) -> Float
 
   // CHECK: %4 = differentiable_function_extract [vjp] %3 : $@differentiable @convention(thin) (Float, @nondiff Float, @nondiff Float) -> Float
   %4 = differentiable_function_extract [vjp] %3 : $@differentiable @convention(thin) (Float, @nondiff Float, @nondiff Float) -> Float
   %5 = function_ref @examplemethod : $@convention(method) (Float, Float, Float) -> Float
-  %6 = differentiable_function [wrt 0 1 2] %5 : $@convention(method) (Float, Float, Float) -> Float
+  %6 = differentiable_function [parameters 0 1 2] %5 : $@convention(method) (Float, Float, Float) -> Float
 
   // CHECK: %7 = differentiable_function_extract [vjp] %6 : $@differentiable @convention(method) (Float, Float, Float) -> Float
   %7 = differentiable_function_extract [vjp] %6 : $@differentiable @convention(method) (Float, Float, Float) -> Float
-  %8 = differentiable_function [wrt 0] %5 : $@convention(method) (Float, Float, Float) -> Float
+  %8 = differentiable_function [parameters 0] %5 : $@convention(method) (Float, Float, Float) -> Float
 
   // CHECK: %9 = differentiable_function_extract [vjp] %8 : $@differentiable @convention(method) (Float, @nondiff Float, @nondiff Float) -> Float
   %9 = differentiable_function_extract [vjp] %8 : $@differentiable @convention(method) (Float, @nondiff Float, @nondiff Float) -> Float
@@ -68,19 +68,19 @@ bb0(%0 : $Float):
 sil @make_diff_func : $@convention(thin) () -> @differentiable @convention(thin) (Float) -> Float {
 bb0:
   %orig = function_ref @foo : $@convention(thin) (Float) -> Float
-  %undiffedFunc = differentiable_function [wrt 0] %orig : $@convention(thin) (Float) -> Float
+  %undiffedFunc = differentiable_function [parameters 0] %orig : $@convention(thin) (Float) -> Float
   %vjp = function_ref @foo_vjp : $@convention(thin) (Float) -> (Float, @owned @callee_guaranteed (Float) -> Float)
-  %diffFunc = differentiable_function [wrt 0] %orig : $@convention(thin) (Float) -> Float with {undef : $@convention(thin) (Float) -> (Float, @owned @callee_guaranteed (Float) -> Float), %vjp : $@convention(thin) (Float) -> (Float, @owned @callee_guaranteed (Float) -> Float)}
+  %diffFunc = differentiable_function [parameters 0] %orig : $@convention(thin) (Float) -> Float with_derivative {undef : $@convention(thin) (Float) -> (Float, @owned @callee_guaranteed (Float) -> Float), %vjp : $@convention(thin) (Float) -> (Float, @owned @callee_guaranteed (Float) -> Float)}
   %extractedVJP = differentiable_function_extract [vjp] %diffFunc : $@differentiable @convention(thin) (Float) -> Float
   %extractedOriginal = differentiable_function_extract [original] %diffFunc : $@differentiable @convention(thin) (Float) -> Float
   return %undiffedFunc : $@differentiable @convention(thin) (Float) -> Float
 }
 
 // CHECK-LABEL: @make_diff_func : $@convention(thin) () -> @differentiable @convention(thin) (Float) -> Float
 // CHECK:   [[FOO:%.*]] = function_ref @foo : $@convention(thin) (Float) -> Float
-// CHECK:   [[UNDIFFED_FOO:%.*]] = differentiable_function [wrt 0] [[FOO]] : $@convention(thin) (Float) -> Float
+// CHECK:   [[UNDIFFED_FOO:%.*]] = differentiable_function [parameters 0] [[FOO]] : $@convention(thin) (Float) -> Float
 // CHECK:   [[FOO_VJP:%.*]] = function_ref @foo_vjp : $@convention(thin) (Float) -> (Float, @owned @callee_guaranteed (Float) -> Float)
-// CHECK:   [[DIFFED_FOO:%.*]] = differentiable_function [wrt 0] [[FOO]] : $@convention(thin) (Float) -> Float with {undef : $@convention(thin) (Float) -> (Float, @owned @callee_guaranteed (Float) -> Float), [[FOO_VJP]] : $@convention(thin) (Float) -> (Float, @owned @callee_guaranteed (Float) -> Float)}
+// CHECK:   [[DIFFED_FOO:%.*]] = differentiable_function [parameters 0] [[FOO]] : $@convention(thin) (Float) -> Float with_derivative {undef : $@convention(thin) (Float) -> (Float, @owned @callee_guaranteed (Float) -> Float), [[FOO_VJP]] : $@convention(thin) (Float) -> (Float, @owned @callee_guaranteed (Float) -> Float)}
 // CHECK:   [[EXTRACTED_VJP:%.*]] = differentiable_function_extract [vjp] [[DIFFED_FOO]] : $@differentiable @convention(thin) (Float) -> Float
 // CHECK:   [[EXTRACTED_ORIG:%.*]] = differentiable_function_extract [original] [[DIFFED_FOO]] : $@differentiable @convention(thin) (Float) -> Float
 // CHECK:   return [[UNDIFFED_FOO]] : $@differentiable @convention(thin) (Float) -> Float

test/AutoDiff/differentiable_function_inst_irgen.sil

@@ -36,7 +36,7 @@ sil @make_diff_func : $@convention(thin) () -> (@convention(thin) (Float) -> Flo
 bb0:
   %orig = function_ref @foo : $@convention(thin) (Float) -> Float
   %vjp = function_ref @foo_vjp : $@convention(thin) (Float) -> (Float, @owned @callee_guaranteed (Float) -> Float)
-  %diffFunc = differentiable_function [wrt 0] %orig : $@convention(thin) (Float) -> Float with {undef : $@convention(thin) (Float) -> (Float, @owned @callee_guaranteed (Float) -> Float), %vjp : $@convention(thin) (Float) -> (Float, @owned @callee_guaranteed (Float) -> Float)}
+  %diffFunc = differentiable_function [parameters 0] %orig : $@convention(thin) (Float) -> Float with_derivative {undef : $@convention(thin) (Float) -> (Float, @owned @callee_guaranteed (Float) -> Float), %vjp : $@convention(thin) (Float) -> (Float, @owned @callee_guaranteed (Float) -> Float)}
   %extractedOrig = differentiable_function_extract [original] %diffFunc : $@differentiable @convention(thin) (Float) -> Float
   %extractedVJP = differentiable_function_extract [vjp] %diffFunc : $@differentiable @convention(thin) (Float) -> Float
   %tuple = tuple (%extractedOrig : $@convention(thin) (Float) -> Float, %extractedVJP : $@convention(thin) (Float) -> (Float, @owned @callee_guaranteed (Float) -> Float))

test/AutoDiff/differentiable_function_silgen.swift

@@ -75,7 +75,7 @@ func apply() {
 // CHECK-SILGEN-LABEL: @{{.*}}apply{{.*}}
 // CHECK-SILGEN:       [[ORIG:%.*]] = function_ref @{{.*}}thin{{.*}} : $@convention(thin) (Float) -> Float
 // CHECK-SILGEN-NEXT:  [[ORIG_THICK:%.*]] = thin_to_thick_function [[ORIG]] : $@convention(thin) (Float) -> Float to $@callee_guaranteed (Float) -> Float
-// CHECK-SILGEN-NEXT:  [[DIFFED:%.*]] = differentiable_function [wrt 0] [[ORIG_THICK]] : $@callee_guaranteed (Float) -> Float
+// CHECK-SILGEN-NEXT:  [[DIFFED:%.*]] = differentiable_function [parameters 0] [[ORIG_THICK]] : $@callee_guaranteed (Float) -> Float
 // CHECK-SILGEN:       [[ORIG:%.*]] = function_ref @{{.*}}thin{{.*}} : $@convention(thin) (Float) -> Float
 // CHECK-SILGEN-NEXT:  [[ORIG_THICK:%.*]] = thin_to_thick_function [[ORIG]] : $@convention(thin) (Float) -> Float to $@callee_guaranteed (Float) -> Float
 // CHECK-SILGEN-NEXT:  [[LIN:%.*]] = linear_function [parameters 0] [[ORIG_THICK]] : $@callee_guaranteed (Float) -> Float
@@ -110,6 +110,6 @@ func appliesReabstraction(_ f: @escaping @differentiable (Float) -> Float) {
 // CHECK-SILGEN:   [[VJP_COPY:%.*]] = copy_value [[VJP]] : $@callee_guaranteed (Float) -> (Float, @owned @callee_guaranteed (Float) -> Float)
 // CHECK-SILGEN:   [[REABS_VJP:%.*]] = function_ref @$sS4fIegyd_Iegydo_S4fIegnr_Iegnro_TR : $@convention(thin) (@in_guaranteed Float, @guaranteed @callee_guaranteed (Float) -> (Float, @owned @callee_guaranteed (Float) -> Float)) -> (@out Float, @owned @callee_guaranteed (@in_guaranteed Float) -> @out Float)
 // CHECK-SILGEN:   [[NEW_VJP:%.*]] = partial_apply [callee_guaranteed] [[REABS_VJP]]([[VJP_COPY]]) : $@convention(thin) (@in_guaranteed Float, @guaranteed @callee_guaranteed (Float) -> (Float, @owned @callee_guaranteed (Float) -> Float)) -> (@out Float, @owned @callee_guaranteed (@in_guaranteed Float) -> @out Float)
-// CHECK-SILGEN:   [[NEW_DIFF_FUNC:%.*]] = differentiable_function [wrt 0] [[NEW_ORIG]] : $@callee_guaranteed (@in_guaranteed Float) -> @out Float with {[[NEW_JVP]] : $@callee_guaranteed (@in_guaranteed Float) -> (@out Float, @owned @callee_guaranteed (@in_guaranteed Float) -> @out Float), [[NEW_VJP]] : $@callee_guaranteed (@in_guaranteed Float) -> (@out Float, @owned @callee_guaranteed (@in_guaranteed Float) -> @out Float)}
+// CHECK-SILGEN:   [[NEW_DIFF_FUNC:%.*]] = differentiable_function [parameters 0] [[NEW_ORIG]] : $@callee_guaranteed (@in_guaranteed Float) -> @out Float with_derivative {[[NEW_JVP]] : $@callee_guaranteed (@in_guaranteed Float) -> (@out Float, @owned @callee_guaranteed (@in_guaranteed Float) -> @out Float), [[NEW_VJP]] : $@callee_guaranteed (@in_guaranteed Float) -> (@out Float, @owned @callee_guaranteed (@in_guaranteed Float) -> @out Float)}
 // CHECK-SILGEN:   [[DIFF_API:%.*]] = function_ref @${{.*}}pullback{{.*}}at{{.*}} : $@convention(thin) <τ_0_0, τ_0_1 where τ_0_0 : _Differentiable, τ_0_1 : _Differentiable> (@in_guaranteed τ_0_0, @guaranteed @differentiable @callee_guaranteed (@in_guaranteed τ_0_0) -> @out τ_0_1) -> @owned @callee_guaranteed (@in_guaranteed τ_0_1.TangentVector) -> @out τ_0_0.TangentVector
 // CHECK-SILGEN:   apply [[DIFF_API]]<Float, Float>({{.*}}, [[NEW_DIFF_FUNC]]) : $@convention(thin) <τ_0_0, τ_0_1 where τ_0_0 : _Differentiable, τ_0_1 : _Differentiable> (@in_guaranteed τ_0_0, @guaranteed @differentiable @callee_guaranteed (@in_guaranteed τ_0_0) -> @out τ_0_1) -> @owned @callee_guaranteed (@in_guaranteed τ_0_1.TangentVector) -> @out τ_0_0.TangentVector