[AutoDiff] Use @derivative attribute to register tgmath derivatives. (#28933)

- Compile stdlib with `-enable-experimental-cross-file-derivative-registration`.
  - This is necessary for using `@derivative` with tgmath functions, since
    some original functions (e.g. the `tan(_: Double) -> Double`) are imported
    from Clang.
- Replace `@differentiable(jvp:vjp:)` with `@derivative` for tgmath functions.

Progress towards TF-1085: using `@derivative` attribute in the stdlib.

Author: dan-zheng

cmake/modules/AddSwift.cmake

@@ -1910,6 +1910,11 @@ function(add_swift_target_library name)
       endif()
     endif()
 
+    # SWIFT_ENABLE_TENSORFLOW
+    # NOTE(TF-1021): Enable cross-file derivative registration for stdlib.
+    list(APPEND swiftlib_swift_compile_flags_all
+         -Xllvm -enable-experimental-cross-file-derivative-registration)
+    # SWIFT_ENABLE_TENSORFLOW END
 
     # Collect architecture agnostic SDK linker flags
     set(swiftlib_link_flags_all ${SWIFTLIB_LINK_FLAGS})

stdlib/public/Platform/tgmath.swift.gyb

@@ -167,114 +167,136 @@ public func ldexp(_ x: ${T}, _ n : Int) -> ${T} {
 %   if T == 'Float80':
 #if !(os(Windows) || os(Android)) && (arch(i386) || arch(x86_64))
 %   end
-@usableFromInline
-func _vjpExp(_ x: ${T}) -> (${T}, (${T}) -> ${T}) {
+@inlinable
+@derivative(of: exp)
+func _vjpExp(_ x: ${T}) -> (value: ${T}, pullback: (${T}) -> ${T}) {
   let value = exp(x)
   return (value, { v in value * v })
 }
 
-@usableFromInline
-func _vjpExp2(_ x: ${T}) -> (${T}, (${T}) -> ${T}) {
+@inlinable
+@derivative(of: exp2)
+func _vjpExp2(_ x: ${T}) -> (value: ${T}, pullback: (${T}) -> ${T}) {
   let value = exp2(x)
   return (value, { v in v * ${T}(M_LN2) * value })
 }
 
-@usableFromInline
-func _vjpLog(_ x: ${T}) -> (${T}, (${T}) -> ${T}) {
+@inlinable
+@derivative(of: log)
+func _vjpLog(_ x: ${T}) -> (value: ${T}, pullback: (${T}) -> ${T}) {
   return (log(x), { v in v / x })
 }
 
-@usableFromInline
-func _vjpLog10(_ x: ${T}) -> (${T}, (${T}) -> ${T}) {
+@inlinable
+@derivative(of: log10)
+func _vjpLog10(_ x: ${T}) -> (value: ${T}, pullback: (${T}) -> ${T}) {
   return (log10(x), { v in v * ${T}(M_LOG10E) / x })
 }
 
-@usableFromInline
-func _vjpLog2(_ x: ${T}) -> (${T}, (${T}) -> ${T}) {
+@inlinable
+@derivative(of: log2)
+func _vjpLog2(_ x: ${T}) -> (value: ${T}, pullback: (${T}) -> ${T}) {
   return (log2(x), { v in v / (${T}(M_LN2) * x) })
 }
 
-@usableFromInline
-func _vjpSin(_ x: ${T}) -> (${T}, (${T}) -> ${T}) {
+@inlinable
+@derivative(of: sin)
+func _vjpSin(_ x: ${T}) -> (value: ${T}, pullback: (${T}) -> ${T}) {
   return (sin(x), { v in v * cos(x) })
 }
 
-@usableFromInline
-func _vjpCos(_ x: ${T}) -> (${T}, (${T}) -> ${T}) {
+@inlinable
+@derivative(of: cos)
+func _vjpCos(_ x: ${T}) -> (value: ${T}, pullback: (${T}) -> ${T}) {
   return (cos(x), { v in -v * sin(x) })
 }
 
-@usableFromInline
-func _vjpTan(_ x: ${T}) -> (${T}, (${T}) -> ${T}) {
+@inlinable
+@derivative(of: tan)
+func _vjpTan(_ x: ${T}) -> (value: ${T}, pullback: (${T}) -> ${T}) {
   let value = tan(x)
   return (value, { v in v * (1 + value * value) })
 }
 
-@usableFromInline
-func _vjpAsin(_ x: ${T}) -> (${T}, (${T}) -> ${T}) {
+@inlinable
+@derivative(of: asin)
+func _vjpAsin(_ x: ${T}) -> (value: ${T}, pullback: (${T}) -> ${T}) {
   return (asin(x), { v in v / sqrt(1 - x * x) })
 }
 
-@usableFromInline
-func _vjpAcos(_ x: ${T}) -> (${T}, (${T}) -> ${T}) {
+@inlinable
+@derivative(of: acos)
+func _vjpAcos(_ x: ${T}) -> (value: ${T}, pullback: (${T}) -> ${T}) {
   return (acos(x), { v in -v / sqrt(1 - x * x) })
 }
 
-@usableFromInline
-func _vjpAtan(_ x: ${T}) -> (${T}, (${T}) -> ${T}) {
+@inlinable
+@derivative(of: atan)
+func _vjpAtan(_ x: ${T}) -> (value: ${T}, pullback: (${T}) -> ${T}) {
   return (atan(x), { v in v / (1 + x * x) })
 }
 
-@usableFromInline
-func _vjpSinh(_ x: ${T}) -> (${T}, (${T}) -> ${T}) {
+@inlinable
+@derivative(of: sinh)
+func _vjpSinh(_ x: ${T}) -> (value: ${T}, pullback: (${T}) -> ${T}) {
   return (sinh(x), { v in v * cosh(x) })
 }
 
-@usableFromInline
-func _vjpCosh(_ x: ${T}) -> (${T}, (${T}) -> ${T}) {
+@inlinable
+@derivative(of: cosh)
+func _vjpCosh(_ x: ${T}) -> (value: ${T}, pullback: (${T}) -> ${T}) {
   return (cosh(x), { v in v * sinh(x) })
 }
 
-@usableFromInline
-func _vjpTanh(_ x: ${T}) -> (${T}, (${T}) -> ${T}) {
+@inlinable
+@derivative(of: tanh)
+func _vjpTanh(_ x: ${T}) -> (value: ${T}, pullback: (${T}) -> ${T}) {
   let value = tanh(x)
   return (value, { v in v * (1 - value * value) })
 }
 
-@usableFromInline
-func _vjpAsinh(_ x: ${T}) -> (${T}, (${T}) -> ${T}) {
+@inlinable
+@derivative(of: asinh)
+func _vjpAsinh(_ x: ${T}) -> (value: ${T}, pullback: (${T}) -> ${T}) {
   return (asinh(x), { v in v / sqrt(1 + x * x) })
 }
 
-@usableFromInline
-func _vjpAcosh(_ x: ${T}) -> (${T}, (${T}) -> ${T}) {
+@inlinable
+@derivative(of: acosh)
+func _vjpAcosh(_ x: ${T}) -> (value: ${T}, pullback: (${T}) -> ${T}) {
   return (acosh(x), { v in v / sqrt(x * x - 1) })
 }
 
-@usableFromInline
-func _vjpAtanh(_ x: ${T}) -> (${T}, (${T}) -> ${T}) {
+@inlinable
+@derivative(of: atanh)
+func _vjpAtanh(_ x: ${T}) -> (value: ${T}, pullback: (${T}) -> ${T}) {
   return (atanh(x), { v in v / (1 - x * x) })
 }
 
-@usableFromInline
-func _vjpExpm1(_ x: ${T}) -> (${T}, (${T}) -> ${T}) {
+@inlinable
+@derivative(of: expm1)
+func _vjpExpm1(_ x: ${T}) -> (value: ${T}, pullback: (${T}) -> ${T}) {
   return (expm1(x), { v in exp(x) * v })
 }
 
-@usableFromInline
-func _vjpLog1p(_ x: ${T}) -> (${T}, (${T}) -> ${T}) {
+@inlinable
+@derivative(of: log1p)
+func _vjpLog1p(_ x: ${T}) -> (value: ${T}, pullback: (${T}) -> ${T}) {
   return (log1p(x), { v in v / (x + 1) })
 }
 
-@usableFromInline
-func _vjpErf(_ x: ${T}) -> (${T}, (${T}) -> ${T}) {
+@inlinable
+@derivative(of: erf)
+func _vjpErf(_ x: ${T}) -> (value: ${T}, pullback: (${T}) -> ${T}) {
   return (erf(x), { v in v * ${T}(M_2_SQRTPI) * exp(-x * x) })
 }
 
-@usableFromInline
-func _vjpErfc(_ x: ${T}) -> (${T}, (${T}) -> ${T}) {
+@inlinable
+@derivative(of: erfc)
+func _vjpErfc(_ x: ${T}) -> (value: ${T}, pullback: (${T}) -> ${T}) {
   return (erfc(x), { v in v * -${T}(M_2_SQRTPI) * exp(-x * x) })
 }
+
 // SWIFT_ENABLE_TENSORFLOW END
 %   if T == 'Float80':
 #endif
@@ -398,7 +420,7 @@ def TypedBinaryFunctions():
 @_transparent
 // SWIFT_ENABLE_TENSORFLOW
 %   if ufunc in HasVJP:
-@differentiable(vjp: _vjp${ufunc.capitalize()}(_:))
+@differentiable
 %   end
 public func ${ufunc}(_ x: ${T}) -> ${T} {
   return ${T}.${ufunc}(x)
@@ -452,7 +474,7 @@ public func tgamma(_ x: Float80) -> Float80 {
 @_transparent
 // SWIFT_ENABLE_TENSORFLOW
 %   if ufunc in HasVJP:
-@differentiable(vjp: _vjp${ufunc.capitalize()}(_:))
+@differentiable
 %   end
 public func ${ufunc}(_ x: ${T}) -> ${T} {
   return ${T}.${ufunc}(x)
@@ -462,7 +484,7 @@ public func ${ufunc}(_ x: ${T}) -> ${T} {
 @_transparent
 // SWIFT_ENABLE_TENSORFLOW
 %   if ufunc in HasVJP:
-@differentiable(vjp: _vjp${ufunc.capitalize()}(_:))
+@differentiable
 %   end
 public func ${ufunc}(_ x: ${T}) -> ${T} {
   return x.rounded(.toNearestOrEven)

validation-test/ParseableInterface/verify_all_overlays.py

@@ -7,6 +7,11 @@
 # RUN: %{python} %s %target-os %target-cpu %platform-sdk-overlay-dir %t \
 # RUN:   %target-swift-frontend -build-module-from-parseable-interface \
 # RUN:     -Fsystem %sdk/System/Library/PrivateFrameworks/ \
+# SWIFT_ENABLE_TENSORFLOW
+# NOTE(TF-1097): Remove flag when retroactive derivative registration is enabled
+# by default.
+# RUN:     -Xllvm -enable-experimental-cross-file-derivative-registration \
+# SWIFT_ENABLE_TENSORFLOW END
 # RUN:     | sort > %t/failures.txt
 # RUN: grep '# %target-os:' %s > %t/filter.txt || true
 # RUN: test ! -e %t/failures.txt || \