[AutoDiff] Add JVPs for tgmath functions. (#29290)

Add JVPs for tgmath functions except `remainder(_:_:)` and `fmod(_:_:)`.

JVP/VJP functions for unary tgmath functions are identical, but there is no way
to register a derivative function as both a JVP/VJP of the same original
function. TF-1111 tracks this; the issue will naturally resolve when linear
functions and transposition are done and VJPs are removed.

Partially resolves TF-1108.

Author: dan-zheng

stdlib/public/Platform/tgmath_derivatives.swift.gyb

@@ -11,14 +11,17 @@
 // See https://swift.org/CONTRIBUTORS.txt for the list of Swift project authors
 //
 //===----------------------------------------------------------------------===//
+// This file defines derivatives for tgmath functions.
+//===----------------------------------------------------------------------===//
 
 @usableFromInline
-@derivative(of: sqrt)
-func _vjpSqrt<T: FloatingPoint & Differentiable> (
-  _ x: T
-) -> (value: T, pullback: (T) -> T) where T == T.TangentVector {
-  let value = sqrt(x)
-  return (value, { v in v / (2 * value) })
+@derivative(of: fma)
+func _jvpFma<T: FloatingPoint & Differentiable> (
+  _ x: T,
+  _ y: T,
+  _ z: T
+) -> (value: T, differential: (T, T, T) -> T) where T == T.TangentVector {
+  return (fma(x, y, z), { (dx, dy, dz) in dx * y + dy * x + dz })
 }
 
 @usableFromInline
@@ -31,6 +34,18 @@ func _vjpFma<T: FloatingPoint & Differentiable> (
   return (fma(x, y, z), { v in (v * y, v * x, v) })
 }
 
+@usableFromInline
+@derivative(of: remainder)
+func _jvpRemainder<T: FloatingPoint & Differentiable> (
+  _ x: T,
+  _ y: T
+) -> (value: T, differential: (T, T) -> T) where T == T.TangentVector {
+  fatalError("""
+    Unimplemented JVP for 'remainder(_:)'. \
+    https://bugs.swift.org/browse/TF-1108 tracks this issue
+    """)
+}
+
 @usableFromInline
 @derivative(of: remainder)
 func _vjpRemainder<T: FloatingPoint & Differentiable> (
@@ -40,6 +55,18 @@ func _vjpRemainder<T: FloatingPoint & Differentiable> (
   return (remainder(x, y), { v in (v, -v * ((x / y).rounded(.toNearestOrEven))) })
 }
 
+@usableFromInline
+@derivative(of: fmod)
+func _jvpFmod<T: FloatingPoint & Differentiable> (
+  _ x: T,
+  _ y: T
+) -> (value: T, differential: (T, T) -> T) where T == T.TangentVector {
+  fatalError("""
+    Unimplemented JVP for 'fmod(_:)'. \
+    https://bugs.swift.org/browse/TF-1108 tracks this issue
+    """)
+}
+
 @usableFromInline
 @derivative(of: fmod)
 func _vjpFmod<T: FloatingPoint & Differentiable> (
@@ -49,173 +76,188 @@ func _vjpFmod<T: FloatingPoint & Differentiable> (
   return (fmod(x, y), { v in (v, -v * ((x / y).rounded(.towardZero))) })
 }
 
+%for derivative_kind in ['jvp', 'vjp']:
+%  linear_map_kind = 'differential' if derivative_kind == 'jvp' else 'pullback'
+@usableFromInline
+@derivative(of: sqrt)
+func _${derivative_kind}Sqrt<T: FloatingPoint & Differentiable> (
+  _ x: T
+) -> (value: T, ${linear_map_kind}: (T) -> T) where T == T.TangentVector {
+  let value = sqrt(x)
+  return (value, { v in v / (2 * value) })
+}
+
 @usableFromInline
 @derivative(of: ceil)
-func _vjpCeil<T: FloatingPoint & Differentiable> (
+func _${derivative_kind}Ceil<T: FloatingPoint & Differentiable> (
   _ x: T
-) -> (value: T, pullback: (T) -> T) where T == T.TangentVector {
+) -> (value: T, ${linear_map_kind}: (T) -> T) where T == T.TangentVector {
   return (ceil(x), { v in 0 })
 }
 
 @usableFromInline
 @derivative(of: floor)
-func _vjpFloor<T: FloatingPoint & Differentiable> (
+func _${derivative_kind}Floor<T: FloatingPoint & Differentiable> (
   _ x: T
-) -> (value: T, pullback: (T) -> T) where T == T.TangentVector {
+) -> (value: T, ${linear_map_kind}: (T) -> T) where T == T.TangentVector {
   return (floor(x), { v in 0 })
 }
 
 @usableFromInline
 @derivative(of: round)
-func _vjpRound<T: FloatingPoint & Differentiable> (
+func _${derivative_kind}Round<T: FloatingPoint & Differentiable> (
   _ x: T
-) -> (value: T, pullback: (T) -> T) where T == T.TangentVector {
+) -> (value: T, ${linear_map_kind}: (T) -> T) where T == T.TangentVector {
   return (round(x), { v in 0 })
 }
 
 @usableFromInline
 @derivative(of: trunc)
-func _vjpTrunc<T: FloatingPoint & Differentiable> (
+func _${derivative_kind}Trunc<T: FloatingPoint & Differentiable> (
   _ x: T
-) -> (value: T, pullback: (T) -> T) where T == T.TangentVector {
+) -> (value: T, ${linear_map_kind}: (T) -> T) where T == T.TangentVector {
   return (trunc(x), { v in 0 })
 }
+%end # for derivative_kind in ['jvp', 'vjp']:
 
-%for T in ['Float', 'Double', 'Float80']:
-%   if T == 'Float80':
+%for derivative_kind in ['jvp', 'vjp']:
+%  linear_map_kind = 'differential' if derivative_kind == 'jvp' else 'pullback'
+%  for T in ['Float', 'Double', 'Float80']:
+%    if T == 'Float80':
 #if !(os(Windows) || os(Android)) && (arch(i386) || arch(x86_64))
-%   end
+%    end
 @inlinable
 @derivative(of: exp)
-func _vjpExp(_ x: ${T}) -> (value: ${T}, pullback: (${T}) -> ${T}) {
+func _${derivative_kind}Exp(_ x: ${T}) -> (value: ${T}, ${linear_map_kind}: (${T}) -> ${T}) {
   let value = exp(x)
   return (value, { v in value * v })
 }
 
 @inlinable
 @derivative(of: exp2)
-func _vjpExp2(_ x: ${T}) -> (value: ${T}, pullback: (${T}) -> ${T}) {
+func _${derivative_kind}Exp2(_ x: ${T}) -> (value: ${T}, ${linear_map_kind}: (${T}) -> ${T}) {
   let value = exp2(x)
   return (value, { v in v * ${T}(M_LN2) * value })
 }
 
 @inlinable
 @derivative(of: log)
-func _vjpLog(_ x: ${T}) -> (value: ${T}, pullback: (${T}) -> ${T}) {
+func _${derivative_kind}Log(_ x: ${T}) -> (value: ${T}, ${linear_map_kind}: (${T}) -> ${T}) {
   return (log(x), { v in v / x })
 }
 
 @inlinable
 @derivative(of: log10)
-func _vjpLog10(_ x: ${T}) -> (value: ${T}, pullback: (${T}) -> ${T}) {
+func _${derivative_kind}Log10(_ x: ${T}) -> (value: ${T}, ${linear_map_kind}: (${T}) -> ${T}) {
   return (log10(x), { v in v * ${T}(M_LOG10E) / x })
 }
 
 @inlinable
 @derivative(of: log2)
-func _vjpLog2(_ x: ${T}) -> (value: ${T}, pullback: (${T}) -> ${T}) {
+func _${derivative_kind}Log2(_ x: ${T}) -> (value: ${T}, ${linear_map_kind}: (${T}) -> ${T}) {
   return (log2(x), { v in v / (${T}(M_LN2) * x) })
 }
 
 @inlinable
 @derivative(of: sin)
-func _vjpSin(_ x: ${T}) -> (value: ${T}, pullback: (${T}) -> ${T}) {
+func _${derivative_kind}Sin(_ x: ${T}) -> (value: ${T}, ${linear_map_kind}: (${T}) -> ${T}) {
   return (sin(x), { v in v * cos(x) })
 }
 
 @inlinable
 @derivative(of: cos)
-func _vjpCos(_ x: ${T}) -> (value: ${T}, pullback: (${T}) -> ${T}) {
+func _${derivative_kind}Cos(_ x: ${T}) -> (value: ${T}, ${linear_map_kind}: (${T}) -> ${T}) {
   return (cos(x), { v in -v * sin(x) })
 }
 
 @inlinable
 @derivative(of: tan)
-func _vjpTan(_ x: ${T}) -> (value: ${T}, pullback: (${T}) -> ${T}) {
+func _${derivative_kind}Tan(_ x: ${T}) -> (value: ${T}, ${linear_map_kind}: (${T}) -> ${T}) {
   let value = tan(x)
   return (value, { v in v * (1 + value * value) })
 }
 
 @inlinable
 @derivative(of: asin)
-func _vjpAsin(_ x: ${T}) -> (value: ${T}, pullback: (${T}) -> ${T}) {
+func _${derivative_kind}Asin(_ x: ${T}) -> (value: ${T}, ${linear_map_kind}: (${T}) -> ${T}) {
   return (asin(x), { v in v / sqrt(1 - x * x) })
 }
 
 @inlinable
 @derivative(of: acos)
-func _vjpAcos(_ x: ${T}) -> (value: ${T}, pullback: (${T}) -> ${T}) {
+func _${derivative_kind}Acos(_ x: ${T}) -> (value: ${T}, ${linear_map_kind}: (${T}) -> ${T}) {
   return (acos(x), { v in -v / sqrt(1 - x * x) })
 }
 
 @inlinable
 @derivative(of: atan)
-func _vjpAtan(_ x: ${T}) -> (value: ${T}, pullback: (${T}) -> ${T}) {
+func _${derivative_kind}Atan(_ x: ${T}) -> (value: ${T}, ${linear_map_kind}: (${T}) -> ${T}) {
   return (atan(x), { v in v / (1 + x * x) })
 }
 
 @inlinable
 @derivative(of: sinh)
-func _vjpSinh(_ x: ${T}) -> (value: ${T}, pullback: (${T}) -> ${T}) {
+func _${derivative_kind}Sinh(_ x: ${T}) -> (value: ${T}, ${linear_map_kind}: (${T}) -> ${T}) {
   return (sinh(x), { v in v * cosh(x) })
 }
 
 @inlinable
 @derivative(of: cosh)
-func _vjpCosh(_ x: ${T}) -> (value: ${T}, pullback: (${T}) -> ${T}) {
+func _${derivative_kind}Cosh(_ x: ${T}) -> (value: ${T}, ${linear_map_kind}: (${T}) -> ${T}) {
   return (cosh(x), { v in v * sinh(x) })
 }
 
 @inlinable
 @derivative(of: tanh)
-func _vjpTanh(_ x: ${T}) -> (value: ${T}, pullback: (${T}) -> ${T}) {
+func _${derivative_kind}Tanh(_ x: ${T}) -> (value: ${T}, ${linear_map_kind}: (${T}) -> ${T}) {
   let value = tanh(x)
   return (value, { v in v * (1 - value * value) })
 }
 
 @inlinable
 @derivative(of: asinh)
-func _vjpAsinh(_ x: ${T}) -> (value: ${T}, pullback: (${T}) -> ${T}) {
+func _${derivative_kind}Asinh(_ x: ${T}) -> (value: ${T}, ${linear_map_kind}: (${T}) -> ${T}) {
   return (asinh(x), { v in v / sqrt(1 + x * x) })
 }
 
 @inlinable
 @derivative(of: acosh)
-func _vjpAcosh(_ x: ${T}) -> (value: ${T}, pullback: (${T}) -> ${T}) {
+func _${derivative_kind}Acosh(_ x: ${T}) -> (value: ${T}, ${linear_map_kind}: (${T}) -> ${T}) {
   return (acosh(x), { v in v / sqrt(x * x - 1) })
 }
 
 @inlinable
 @derivative(of: atanh)
-func _vjpAtanh(_ x: ${T}) -> (value: ${T}, pullback: (${T}) -> ${T}) {
+func _${derivative_kind}Atanh(_ x: ${T}) -> (value: ${T}, ${linear_map_kind}: (${T}) -> ${T}) {
   return (atanh(x), { v in v / (1 - x * x) })
 }
 
 @inlinable
 @derivative(of: expm1)
-func _vjpExpm1(_ x: ${T}) -> (value: ${T}, pullback: (${T}) -> ${T}) {
+func _${derivative_kind}Expm1(_ x: ${T}) -> (value: ${T}, ${linear_map_kind}: (${T}) -> ${T}) {
   return (expm1(x), { v in exp(x) * v })
 }
 
 @inlinable
 @derivative(of: log1p)
-func _vjpLog1p(_ x: ${T}) -> (value: ${T}, pullback: (${T}) -> ${T}) {
+func _${derivative_kind}Log1p(_ x: ${T}) -> (value: ${T}, ${linear_map_kind}: (${T}) -> ${T}) {
   return (log1p(x), { v in v / (x + 1) })
 }
 
 @inlinable
 @derivative(of: erf)
-func _vjpErf(_ x: ${T}) -> (value: ${T}, pullback: (${T}) -> ${T}) {
+func _${derivative_kind}Erf(_ x: ${T}) -> (value: ${T}, ${linear_map_kind}: (${T}) -> ${T}) {
   return (erf(x), { v in v * ${T}(M_2_SQRTPI) * exp(-x * x) })
 }
 
 @inlinable
 @derivative(of: erfc)
-func _vjpErfc(_ x: ${T}) -> (value: ${T}, pullback: (${T}) -> ${T}) {
+func _${derivative_kind}Erfc(_ x: ${T}) -> (value: ${T}, ${linear_map_kind}: (${T}) -> ${T}) {
   return (erfc(x), { v in v * -${T}(M_2_SQRTPI) * exp(-x * x) })
 }
 
-%   if T == 'Float80':
+%    if T == 'Float80':
 #endif
-%   end
-%end
+%    end # if T == 'Float80':
+%  end # for T in ['Float', 'Double', 'Float80']:
+%end # for derivative_kind in ['jvp', 'vjp']:

test/AutoDiff/downstream/forward_mode_runtime.swift

@@ -1,4 +1,4 @@
-// RUN: %target_run_simple_swift_forward_mode_differentiation
+// RUN: %target-run-simple-swift-forward-mode-differentiation
 // REQUIRES: executable_test
 
 import StdlibUnittest

test/AutoDiff/downstream/nonvaried_result.swift

@@ -1,7 +1,7 @@
 // RUN: %target-swift-frontend -Xllvm -sil-print-after=differentiation %s -emit-sil -o /dev/null 2>&1 | %FileCheck %s
 // RUN: %target-run-simple-swift
 // TODO: Test forward-mode differentiation when it supports control flow.
-// UN: %target_run_simple_swift_forward_mode_differentiation
+// UN: %target-run-simple-swift-forward-mode-differentiation
 // REQUIRES: executable_test
 
 // Test differentiation edge case: functions with non-varied results.

test/AutoDiff/downstream/simple_math.swift

@@ -1,6 +1,6 @@
 // RUN: %target-run-simple-swift
 // NOTE(TF-813): verify that enabling forward-mode does not affect reverse-mode.
-// RUN: %target_run_simple_swift_forward_mode_differentiation
+// RUN: %target-run-simple-swift-forward-mode-differentiation
 // RUN: %target-swift-frontend -Xllvm -sil-print-after=differentiation %s -emit-sil -o /dev/null 2>&1 | %FileCheck %s
 // REQUIRES: executable_test