[AutoDiff] Use @derivative attribute to register tgmath derivatives.

[dan-zheng]

• Compile stdlib with -enable-experimental-cross-file-derivative-registration.
  • This is necessary for using @derivative with tgmath functions, since
    some original functions (e.g. 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.

[dan-zheng]

@swift-ci Please test tensorflow

[rxwei]

Neat!

[dan-zheng]

@swift-ci Please test tensorflow

[dan-zheng]

Note: it'd be nicer to register a derivative for a foreign function that doesn't already have a registered derivative in the stdlib (e.g. tan). Let's revisit this test if issues arise with @derivative(of: tan).

[dan-zheng]

Some context below.

Swift tgmath functions are interesting because they call and alias Clang-imported functions. For example, func tan(_: Float) -> Float in Darwin/Glibc aliases the C-imported tan function.

// Clang-imported foreign function, `public_external` linkage:
sil [serializable] [readnone] [clang tan] @tan : $@convention(c) (Double) -> Double

// Clang-imported non-foreign function:
// @nonobjc tan(_:)
sil shared [serializable] [readnone] @$sSo3tanyS2dFTO : $@convention(thin) (Double) -> Double

// Wrapper Swift function in the Darwin/Glibc module:
sil [transparent] [serialized] @$s6Darwin3tanys7Float80VADF : $@convention(thin) (Float80) -> Float80

One idea is to use module-qualified names in @derivative attributes to always register a derivative for Swift functions:

@derivative(Darwin.tan)
func _vjpTan(_ x: Float) -> (value: Float, pullback: (Float) -> Float) { ... }

However, this doesn't work because some tgmath function references are SILGen'd as function_ref to the Clang-imported foreign function, not the wrapper Swift function in the Darwin module. I'm not exactly sure why this occurs:

import Darwin
// Foreign: sil [serializable] [readnone] [clang tan] @tan : $@convention(c) (Double) -> Double
_ = tan(Double(3))
// Non-foreign: sil [transparent] [serialized] @$s6Darwin3tanyS2fF : $@convention(thin) (Float) -> Float
_ = tan(Float(3))

In any case, registering derivatives for the actual underlying function seems more robust/future-proof than registering derivatives for any given wrapper function, because:

• There exist multiple wrapper functions, e.g. @$sSo3tanyS2dFTO and @$s6Darwin3tanys7Float80VADF.
• Any wrapper function can be differentiated if the underlying function has a registered derivative.

---

I refactored derivative registration support for imported functions to a separate PR #29016. I then rebased this PR on top of #29016: differentiability witnesses are emitted for all underlying tgmath functions.

[rxwei]

Looks like clang-imported symbols are in the __C module. And we definitely do want to allow users to register derivatives for a module-qualified C function. What if we make any module-qualified C function derivative also imply a derivative registration for the clang-imported C function in the __C module?

[dan-zheng]

Looks like clang-imported symbols are in the __C module. And we definitely do want to allow users to register derivatives for a module-qualified C function. What if we make any module-qualified C function derivative also imply a derivative registration for the clang-imported C function in the __C module?

Thanks for the idea! I actually found __C module functions to be problematic (TF-1087).

Example __C module function: mangled name comes from SILDeclRef(originalAFD) with isForeign = false.

$ swift-demangle sSo3tanyS2dFTO
$sSo3tanyS2dFTO ---> @nonobjc __C.tan(Swift.Double) -> Swift.Double

1. References to the __C module don't seem to be supported in Swift:

import Darwin
_ = __C.tan

test.swift:2:5: error: use of unresolved identifier '__C'
_ = __C.tan
    ^~~

I'm not sure how to trigger SILGen for __C module function references, e.g. sSo3tanyS2dFTO.

2. __C module functions have shared linkage and aren't naturally SILGen'd, so forcibly SILGen'ing them causes an assertion error (TF-1087):

--- a/lib/SILGen/SILGen.cpp
+++ b/lib/SILGen/SILGen.cpp
@@ -802,6 +802,12 @@ void SILGenModule::postEmitFunction(SILDeclRef constant,
                               derivativeGenSig);
         emitDifferentiabilityWitness(origAFD, origFn, config, jvp, vjp,
                                      derivAttr);
+        if (requiresForeignEntryPoint(origAFD)) {
+          auto origDeclRef = SILDeclRef(origAFD);
+          auto *origFn = getFunction(origDeclRef, NotForDefinition);
+          llvm::errs() << "Test non-foreign function (in the `__C` module):\n";
+          origFn->dump();
+        }
       }
     };
     if (auto *accessor = dyn_cast<AccessorDecl>(AFD))

$ ninja swift-stdlib
Test non-foreign function (in the `__C` module):
// @nonobjc tan(_:)
sil shared [serializable] [readnone] @$sSo3tanyS2dFTO : $@convention(thin) (Double) -> Double

...

SIL verification failed: external declaration of internal SILFunction not allowed: F->isAvailableExternally()In function:
// @nonobjc tan(_:)
sil shared [serializable] [readnone] @$sSo3tanyS2dFTO : $@convention(thin) (Double) -> Double

I believe the body of a __C function is a simple wrapper around the raw Clang-imported @convention(c) function, so __C functions can be differentiated if derivatives are registered for the Clang-imported functions.

[dan-zheng]

@swift-ci Please test tensorflow