[AutoDiff] SILGen derivative function thunks must not be transparent. (#27752)

SILGen generates thunks for derivative functions registered via
`@differentiable` and `@differentiating` attributes.

The thunks have a canonical, consistent naming (currently based on
the original function name and parameter indices) that TBDGen can also
generate via `SILDeclRef::mangle`.

These SILGen derivative function thunks must not be transparent; otherwise,
the functions will be inlined during MandatoryInlining and will not be
exposed publicly.

This bug was discovered as a DeadFunctionElimination crash:
`function_ref` instructions referencing transparent derivative thunks
had their referenced functions set to null by MandatoryInling and became
invalid.

Author: dan-zheng

lib/SILGen/SILGenPoly.cpp

@@ -3695,6 +3695,8 @@ SILGenModule::getOrCreateAutoDiffDerivativeFunctionThunk(
   SILGenFunctionBuilder fb(*this);
   auto linkage = autodiff::getAutoDiffDerivativeFunctionLinkage(
       original->getLinkage(), /*isDerivativeFnExported*/ true);
+  // This thunk is publicly exposed and cannot be transparent.
+  // TODO(TF-925): Mark the thunks as "always inline" for optimization.
   auto *thunk = fb.getOrCreateFunction(
       loc, name, linkage, origDerivativeFnType, IsBare, IsNotTransparent,
       derivativeFn->isSerialized(), derivativeFn->isDynamicallyReplaceable(),

lib/SILGen/SILGenThunk.cpp

@@ -86,8 +86,10 @@ SILGenModule::getOrCreateAutoDiffThunk(SILDeclRef derivativeFnDeclRef,
   auto linkage = autodiff::getAutoDiffDerivativeFunctionLinkage(
       originalLinkage, /*isDerivativeFnExported*/ true);
   auto name = derivativeFnDeclRef.mangle();
+  // This thunk is publicly exposed and cannot be transparent.
+  // TODO(TF-925): Mark the thunks as "always inline" for optimization.
   auto *thunk = builder.getOrCreateFunction(
-      derivativeFnDecl, name, linkage, derivativeFnTy, IsBare, IsTransparent,
+      derivativeFnDecl, name, linkage, derivativeFnTy, IsBare, IsNotTransparent,
       derivativeFnDeclRef.isSerialized(), IsNotDynamic, ProfileCounter(),
       IsThunk);
   if (!thunk->empty())

test/AutoDiff/silgen_thunking/main.swift

@@ -10,6 +10,17 @@
 import StdlibUnittest
 import DifferentiationUnittest
 
+// Verify that SILGen derivative thunks are never `[transparent]`.
+@differentiable(vjp: vjpNoReabstraction)
+func noReabstraction<T: Differentiable>(_ x: T) -> T {
+  return x
+}
+func vjpNoReabstraction<T: Differentiable>(_ x: T) -> (T, (T.TangentVector) -> T.TangentVector) {
+  return (x, { $0 })
+}
+// Find the non-`[transparent]` SILGen thunk.
+// CHECK-LABEL: sil hidden [thunk] [ossa] @AD__$s4main15noReabstractionyxxs15_DifferentiableRzlF__vjp_src_0_wrt_0 : $@convention(thin) <τ_0_0 where τ_0_0 : _Differentiable> (@in_guaranteed τ_0_0) -> (@out τ_0_0, @owned @callee_guaranteed (@in_guaranteed τ_0_0.TangentVector) -> @out τ_0_0.TangentVector)
+
 var DerivativeSILGenThunkTests = TestSuite("DerivativeSILGenThunks")
 
 // TF-619: Test cross-module import of `@differentiable` methods with

test/AutoDiff/tbdgen.swift

@@ -12,6 +12,13 @@
 @differentiable public func publicDiffable(_ x: Float, _ y: Float) -> Float { return x }
 @differentiable(wrt: (x)) public func publicDiffableWRT(_ x: Float, _ y: Float) -> Float { return x }
 
+// Tests SILGen derivative "forwarding thunk" (no derivative reabstraction/self-reordering).
+@differentiable(vjp: publicNoDerivativeReabstractionVJP)
+public func publicNoDerivativeReabstraction<T: Differentiable>(_ x: T) -> T { return x }
+public func publicNoDerivativeReabstractionVJP<T: Differentiable>(_ x: T) -> (T, (T.TangentVector) -> T.TangentVector) {
+  return (x, { $0 })
+}
+
 @differentiable internal func internalDiffable(_ x: Float, _ y: Float) -> Float { return x }
 @differentiable(wrt: (x)) internal func internalDiffableWRT(_ x: Float, _ y: Float) -> Float { return x }