[AutoDiff] Initial Forward Mode AD Support

stdlib/private/DifferentiationUnittest/DifferentiationUnittest.swift

@@ -65,12 +65,12 @@ public struct Tracked<T> {
   }
   private var handle: Box
 
-  @differentiable(vjp: _vjpInit where T : Differentiable, T == T.TangentVector)
+  @differentiable(jvp: _jvpInit, vjp: _vjpInit where T : Differentiable, T == T.TangentVector)
   public init(_ value: T) {
     self.handle = Box(value)
   }
 
-  @differentiable(vjp: _vjpValue where T : Differentiable, T == T.TangentVector)
+  @differentiable(jvp: _jvpValue, vjp: _vjpValue where T : Differentiable, T == T.TangentVector)
   public var value: T {
     get { handle.value }
     set { handle.value = newValue }
@@ -177,10 +177,21 @@ extension Tracked where T : Differentiable, T == T.TangentVector {
     return (Tracked(value), { v in v.value })
   }
 
+  @usableFromInline
+  internal static func _jvpInit(_ value: T)
+      -> (value: Self, differential: (T.TangentVector) -> (Self.TangentVector)) {
+    return (Tracked(value), { v in Tracked(v) })
+  }
+
   @usableFromInline
   internal func _vjpValue() -> (T, (T.TangentVector) -> Self.TangentVector) {
     return (value, { v in Tracked(v) })
   }
+
+  @usableFromInline
+  internal func _jvpValue() -> (T, (Self.TangentVector) -> T.TangentVector) {
+    return (value, { v in v.value })
+  }
 }
 
 extension Tracked where T : Differentiable, T == T.TangentVector {
@@ -197,6 +208,20 @@ extension Tracked where T : Differentiable, T == T.TangentVector {
       -> (value: Self, pullback: (Self) -> (Self, Self)) {
     return (lhs - rhs, { v in (v, .zero - v) })
   }
+
+  @usableFromInline
+  @differentiating(+)
+  internal static func _vjpAdd(lhs: Self, rhs: Self)
+      -> (value: Self, differential: (Self, Self) -> (Self)) {
+    return (lhs + rhs, { (dx, dy) in dx + dy })
+  }
+
+  @usableFromInline
+  @differentiating(-)
+  internal static func _vjpSubtract(lhs: Self, rhs: Self)
+      -> (value: Self, differential: (Self, Self) -> (Self)) {
+    return (lhs - rhs, { (dx, dy) in dx - dy })
+  }
 }
 
 extension Tracked where T : Differentiable & SignedNumeric, T == T.Magnitude,
@@ -207,6 +232,13 @@ extension Tracked where T : Differentiable & SignedNumeric, T == T.Magnitude,
       -> (value: Self, pullback: (Self) -> (Self, Self)) {
     return (lhs * rhs, { v in (v * rhs, v * lhs) })
   }
+
+  @usableFromInline
+  @differentiating(*)
+  internal static func _vjpMultiply(lhs: Self, rhs: Self)
+      -> (value: Self, differential: (Self, Self) -> (Self)) {
+    return (lhs * rhs, { (dx, dy) in dx * rhs + dy * lhs })
+  }
 }
 
 extension Tracked where T : Differentiable & FloatingPoint, T == T.TangentVector {
@@ -216,6 +248,13 @@ extension Tracked where T : Differentiable & FloatingPoint, T == T.TangentVector
       -> (value: Self, pullback: (Self) -> (Self, Self)) {
     return (lhs / rhs, { v in (v / rhs, -lhs / (rhs * rhs) * v) })
   }
+
+  @usableFromInline
+  @differentiating(/)
+  internal static func _vjpDivide(lhs: Self, rhs: Self)
+      -> (value: Self, differential: (Self, Self) -> (Self)) {
+    return (lhs / rhs, { (dx, dy) in dx / rhs - lhs / (rhs * rhs) * dy })
+  }
 }
 
 // Differential operators for `Tracked<Float>`.

stdlib/public/core/AutoDiff.swift

@@ -479,15 +479,15 @@ public func pullback<T, U, V, R>(
 
 @inlinable
 public func derivative<T: FloatingPoint, R>(
-  at x: T, in f: @escaping @differentiable (T) -> R
+  at x: T, in f: @differentiable (T) -> R
 ) ->  R.TangentVector
   where T.TangentVector == T {
   return differential(at: x, in: f)(T(1))
 }
 
 @inlinable
 public func derivative<T: FloatingPoint, U: FloatingPoint, R>(
-  at x: T, _ y: U, in f: @escaping @differentiable (T, U) -> R
+  at x: T, _ y: U, in f: @differentiable (T, U) -> R
 ) -> R.TangentVector
   where T.TangentVector == T,
         U.TangentVector == U {
@@ -496,7 +496,7 @@ public func derivative<T: FloatingPoint, U: FloatingPoint, R>(
 
 @inlinable
 public func derivative<T: FloatingPoint, U: FloatingPoint, V: FloatingPoint, R>(
-  at x: T, _ y: U, _ z: V, in f: @escaping @differentiable (T, U, V) -> R
+  at x: T, _ y: U, _ z: V, in f: @differentiable (T, U, V) -> R
 ) -> R.TangentVector
   where T.TangentVector == T,
         U.TangentVector == U,
@@ -995,3 +995,14 @@ public extension Array where Element: Differentiable {
     return (value: values, pullback: pullback)
   }
 }
+
+//===----------------------------------------------------------------------===//
+// JVP Diagnostics
+//===----------------------------------------------------------------------===//
+@_silgen_name("_printJVPErrorAndExit")
+public func _printJVPErrorAndExit() -> Never {
+    fatalError("""
+        JVP does not exist. Differential-first differentiation APIs are \
+        experimental and should not be used.
+        """)
+}

stdlib/public/core/FloatingPointTypes.swift.gyb

@@ -1634,6 +1634,15 @@ extension ${Self} {
     -> (value: ${Self}, pullback: (${Self}) -> ${Self}) {
     return (-x, { v in -v })
   }
+
+  @usableFromInline
+  @_transparent
+  // SWIFT_ENABLE_TENSORFLOW
+  @differentiating(-)
+  static func _jvpNegate(x: ${Self})
+  -> (value: ${Self}, differential: (${Self}) -> ${Self}) {
+    return (-x, { dx in -dx })
+  }
 }
 
 //===----------------------------------------------------------------------===//
@@ -1797,6 +1806,15 @@ extension ${Self} {
     return (lhs + rhs, { v in (v, v) })
   }
 
+  @inlinable // FIXME(sil-serialize-all)
+  @_transparent
+  @differentiating(+)
+  static func _jvpAdd(
+    lhs: ${Self}, rhs: ${Self}
+  ) -> (value: ${Self}, differential: (${Self}, ${Self}) -> ${Self}) {
+    return (lhs + rhs, { (dlhs, drhs) in dlhs + drhs })
+  }
+
   @inlinable // FIXME(sil-serialize-all)
   @_transparent
   @differentiating(-)
@@ -1806,6 +1824,15 @@ extension ${Self} {
     return (lhs - rhs, { v in (v, -v) })
   }
 
+  @inlinable // FIXME(sil-serialize-all)
+  @_transparent
+  @differentiating(-)
+  static func _jvpSubtract(
+    lhs: ${Self}, rhs: ${Self}
+  ) -> (value: ${Self}, differential: (${Self}, ${Self}) -> ${Self}) {
+    return (lhs - rhs, { (dlhs, drhs) in dlhs - drhs })
+  }
+
   @inlinable // FIXME(sil-serialize-all)
   @_transparent
   @differentiating(*)
@@ -1815,6 +1842,15 @@ extension ${Self} {
     return (lhs * rhs, { v in (rhs * v, lhs * v) })
   }
 
+  @inlinable // FIXME(sil-serialize-all)
+  @_transparent
+  @differentiating(*)
+  static func _jvpMultiply(
+    lhs: ${Self}, rhs: ${Self}
+  ) -> (value: ${Self}, differential: (${Self}, ${Self}) -> ${Self}) {
+    return (lhs * rhs, { (dlhs, drhs) in lhs * drhs + rhs * dlhs })
+  }
+
   @inlinable // FIXME(sil-serialize-all)
   @_transparent
   @differentiating(/)
@@ -1823,6 +1859,15 @@ extension ${Self} {
   ) -> (value: ${Self}, pullback: (${Self}) -> (${Self}, ${Self})) {
     return (lhs / rhs, { v in (v / rhs, -lhs / (rhs * rhs) * v) })
   }
+
+  @inlinable // FIXME(sil-serialize-all)
+  @_transparent
+  @differentiating(/)
+  static func _jvpDivide(
+    lhs: ${Self}, rhs: ${Self}
+  ) -> (value: ${Self}, differential: (${Self}, ${Self}) -> ${Self}) {
+    return (lhs / rhs, { (dlhs, drhs) in dlhs / rhs - lhs / (rhs * rhs) * drhs })
+  }
 }
 
 //===----------------------------------------------------------------------===//

test/AutoDiff/autodiff_diagnostics.swift

@@ -277,7 +277,6 @@ func activeInoutArg(_ x: Float) -> Float {
 // expected-error @+1 {{function is not differentiable}}
 _ = pullback(at: .zero, in: activeInoutArg(_:))
 
-
 func activeInoutArgTuple(_ x: Float) -> Float {
   var tuple = (x, x)
   // expected-note @+1 {{cannot differentiate through 'inout' arguments}}

test/AutoDiff/differentiable_attr_silgen_cross_module.swift

@@ -24,4 +24,3 @@ _ = pullback(at: Wrapper(1)) { x in x + x * x }
 // CHECK-SILGEN-NEXT: sil [differentiable source 0 wrt 0, 1 vjp @AD__$s39differentiable_attr_silgen_other_module7WrapperV1poiyA2C_ACtFZ__vjp_src_0_wrt_0_1] @$s39differentiable_attr_silgen_other_module7WrapperV1poiyA2C_ACtFZ : $@convention(method) (Wrapper, Wrapper, @thin Wrapper.Type) -> Wrapper
 // CHECK-SIL-LABEL: // static Wrapper.+ infix(_:_:)
 // CHECK-SIL-NEXT: sil [differentiable source 0 wrt 0, 1 jvp @AD__$s39differentiable_attr_silgen_other_module7WrapperV1poiyA2C_ACtFZ__jvp_src_0_wrt_0_1 vjp @AD__$s39differentiable_attr_silgen_other_module7WrapperV1poiyA2C_ACtFZ__vjp_src_0_wrt_0_1] @$s39differentiable_attr_silgen_other_module7WrapperV1poiyA2C_ACtFZ : $@convention(method) (Wrapper, Wrapper, @thin Wrapper.Type) -> Wrapper
-

test/AutoDiff/forward_mode_diagnostics.swift

@@ -0,0 +1,95 @@
+// RUN: %target-swift-frontend -Xllvm -run-jvp-generation -emit-sil -verify %s
+
+// TODO: move these tests back into `autodiff_diagnostics.swift` once
+// forward mode reaches feature parity with reverse mode.
+
+//===----------------------------------------------------------------------===//
+// Basic function
+//===----------------------------------------------------------------------===//
+
+func one_to_one_0(_ x: Float) -> Float {
+  return x + 2
+}
+
+_ = derivative(at: 0, in: one_to_one_0) // okay!
+
+//===----------------------------------------------------------------------===//
+// Function composition
+//===----------------------------------------------------------------------===//
+
+func base(_ x: Float) -> Float {
+  // expected-error @+2 2 {{expression is not differentiable}}
+  // expected-note @+1 2 {{cannot differentiate through a non-differentiable result; do you want to use 'withoutDerivative(at:)'?}}
+  return Float(Int(x))
+}
+
+// TODO: Fix nested differentiation diagnostics. Need to fix indirect differentiation invokers.
+func nested(_ x: Float) -> Float {
+  // xpected-note @+1 {{when differentiating this function call}}
+  return base(x)
+}
+
+func middle(_ x: Float) -> Float {
+  // xpected-note @+1 {{when differentiating this function call}}
+  return nested(x)
+}
+
+func middle2(_ x: Float) -> Float {
+  // xpected-note @+1 {{when differentiating this function call}}
+  return middle(x)
+}
+
+func func_to_diff(_ x: Float) -> Float {
+  // xpected-note @+1 {{expression is not differentiable}}
+  return middle2(x)
+}
+
+func calls_diff_of_nested(_ x: Float) -> Float {
+  // xpected-error @+1 {{function is not differentiable}}
+  return derivative(at: x, in: func_to_diff)
+}
+
+//===----------------------------------------------------------------------===//
+// Inout arguments
+//===----------------------------------------------------------------------===//
+
+func activeInoutArg(_ x: Float) -> Float {
+  var a = x
+  // expected-note @+1 {{cannot differentiate through 'inout' arguments}}
+  a += x
+  return a
+}
+// expected-error @+1 {{function is not differentiable}}
+_ = differential(at: .zero, in: activeInoutArg(_:))
+
+func activeInoutArgTuple(_ x: Float) -> Float {
+  var tuple = (x, x)
+  // expected-note @+1 {{cannot differentiate through 'inout' arguments}}
+  tuple.0 *= x
+  return x * tuple.0
+}
+// expected-error @+1 {{function is not differentiable}}
+_ = differential(at: .zero, in: activeInoutArgTuple(_:))
+
+//===----------------------------------------------------------------------===//
+// Non-varied results
+//===----------------------------------------------------------------------===//
+
+func one() -> Float {
+  return 1
+}
+@differentiable
+func nonVariedResult(_ x: Float) -> Float {
+  // expected-warning @+1 2 {{result does not depend on differentiation arguments and will always have a zero derivative; do you want to use 'withoutDerivative(at:)'?}} {{10-10=withoutDerivative(at:}}
+  return one()
+}
+
+//===----------------------------------------------------------------------===//
+// Subset parameters
+//===----------------------------------------------------------------------===//
+
+func nondiff(_ f: @differentiable (Float, @nondiff Float) -> Float) -> Float {
+  // expected-note @+2 {{cannot differentiate with respect to a '@nondiff' parameter}}
+  // expected-error @+1 {{function is not differentiable}}
+  return derivative(at: 2, 3) { (x, y) in f(x * x, y) }
+}