[AutoDiff] Mark SILGen derivative function thunks as "always inline". (#27754)

SILGen derivative function thunks must be publicly exposed and cannot be
`[transparent]`. Instead, mark them as "always inline" for optimization.

Author: dan-zheng

lib/SILGen/SILGenPoly.cpp

@@ -3696,12 +3696,13 @@ SILGenModule::getOrCreateAutoDiffDerivativeReabstractionThunk(
   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.
+  // Instead, mark it as "always inline" for optimization.
   auto *thunk = fb.getOrCreateFunction(
       loc, name, linkage, origDerivativeFnType, IsBare, IsNotTransparent,
       derivativeFn->isSerialized(), derivativeFn->isDynamicallyReplaceable(),
       derivativeFn->getEntryCount(), derivativeFn->isThunk(),
       derivativeFn->getClassSubclassScope());
+  thunk->setInlineStrategy(AlwaysInline);
   if (!thunk->empty())
     return thunk;
   thunk->setGenericEnvironment(thunkGenericEnv);

lib/SILGen/SILGenThunk.cpp

@@ -87,11 +87,12 @@ SILGenModule::getOrCreateAutoDiffDerivativeForwardingThunk(
       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.
+  // Instead, mark it as "always inline" for optimization.
   auto *thunk = builder.getOrCreateFunction(
       derivativeFnDecl, name, linkage, derivativeFnTy, IsBare, IsNotTransparent,
       derivativeFnDeclRef.isSerialized(), IsNotDynamic, ProfileCounter(),
       IsThunk);
+  thunk->setInlineStrategy(AlwaysInline);
   if (!thunk->empty())
     return thunk;
 

test/AutoDiff/silgen_thunking/main.swift

@@ -19,7 +19,7 @@ func vjpNoReabstraction<T: Differentiable>(_ x: T) -> (T, (T.TangentVector) -> T
   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)
+// CHECK-LABEL: sil hidden [thunk] [always_inline] [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")
 
@@ -51,7 +51,7 @@ struct SelfReordering : Differentiable & AdditiveArithmetic {
     return (value, { v in (Self(1), Self(2), Self(3)) })
   }
 
-// CHECK-LABEL: sil hidden [ossa] @AD__$s4main14SelfReorderingV20threeParameterMethodyA2C_ACtF__jvp_src_0_wrt_0_1_2 : $@convention(method) (@guaranteed SelfReordering, @guaranteed SelfReordering, @guaranteed SelfReordering) -> (@owned SelfReordering, @owned @callee_guaranteed (@guaranteed SelfReordering, @guaranteed SelfReordering, @guaranteed SelfReordering) -> @owned SelfReordering)
+// CHECK-LABEL: sil hidden [always_inline] [ossa] @AD__$s4main14SelfReorderingV20threeParameterMethodyA2C_ACtF__jvp_src_0_wrt_0_1_2 : $@convention(method) (@guaranteed SelfReordering, @guaranteed SelfReordering, @guaranteed SelfReordering) -> (@owned SelfReordering, @owned @callee_guaranteed (@guaranteed SelfReordering, @guaranteed SelfReordering, @guaranteed SelfReordering) -> @owned SelfReordering)
 // CHECK: bb0([[X:%.*]] : @guaranteed $SelfReordering, [[Y:%.*]] : @guaranteed $SelfReordering, [[SELF:%.*]] : @guaranteed $SelfReordering):
 // CHECK: [[JVP:%.*]] = function_ref @$s4main14SelfReorderingV23jvpThreeParameterMethodyAC_A2C_A2CtctAC_ACtF
 // CHECK: [[JVP_RESULT:%.*]] = apply [[JVP]]([[X]], [[Y]], [[SELF]])
@@ -66,7 +66,7 @@ struct SelfReordering : Differentiable & AdditiveArithmetic {
 // CHECK: [[DF_RESULT:%.*]] = apply [[DF]]([[DSELF]], [[DX]], [[DY]])
 // CHECK: return [[DF_RESULT]]
 
-// CHECK-LABEL: sil hidden [ossa] @AD__$s4main14SelfReorderingV20threeParameterMethodyA2C_ACtF__vjp_src_0_wrt_0_1_2 : $@convention(method) (@guaranteed SelfReordering, @guaranteed SelfReordering, @guaranteed SelfReordering) -> (@owned SelfReordering, @owned @callee_guaranteed (@guaranteed SelfReordering) -> (@owned SelfReordering, @owned SelfReordering, @owned SelfReordering))
+// CHECK-LABEL: sil hidden [always_inline] [ossa] @AD__$s4main14SelfReorderingV20threeParameterMethodyA2C_ACtF__vjp_src_0_wrt_0_1_2 : $@convention(method) (@guaranteed SelfReordering, @guaranteed SelfReordering, @guaranteed SelfReordering) -> (@owned SelfReordering, @owned @callee_guaranteed (@guaranteed SelfReordering) -> (@owned SelfReordering, @owned SelfReordering, @owned SelfReordering))
 // CHECK: bb0([[X:%.*]] : @guaranteed $SelfReordering, [[Y:%.*]] : @guaranteed $SelfReordering, [[SELF:%.*]] : @guaranteed $SelfReordering):
 // CHECK: [[VJP:%.*]] = function_ref @$s4main14SelfReorderingV23vjpThreeParameterMethodyAC_AC_A2CtACctAC_ACtF
 // CHECK: [[VJP_RESULT:%.*]] = apply [[VJP]]([[X]], [[Y]], [[SELF]])
@@ -115,7 +115,7 @@ where Dummy: Differentiable & ExpressibleByIntegerLiteral {
     return (value, { v in (v, 2.0, 3.0) })
   }
 
-// CHECK-LABEL: sil hidden [ossa] @AD__$s4main21SelfReorderingGenericV20threeParameterMethodyACyxGqd___qd_0_ts15_DifferentiableRd__sAFRd_0_s25ExpressibleByFloatLiteral13TangentVectorRpd__sAgHRpd_0_r0_lF__jvp_src_0_wrt_0_1_2 : $@convention(method) <τ_0_0 where τ_0_0 : ExpressibleByIntegerLiteral, τ_0_0 : _Differentiable><τ_1_0, τ_1_1 where τ_1_0 : _Differentiable, τ_1_1 : _Differentiable, τ_1_0.TangentVector : ExpressibleByFloatLiteral, τ_1_1.TangentVector : ExpressibleByFloatLiteral> (@in_guaranteed τ_1_0, @in_guaranteed τ_1_1, @in_guaranteed SelfReorderingGeneric<τ_0_0>) -> (@out SelfReorderingGeneric<τ_0_0>, @owned @callee_guaranteed (@in_guaranteed τ_1_0.TangentVector, @in_guaranteed τ_1_1.TangentVector, @in_guaranteed SelfReorderingGeneric<τ_0_0>.TangentVector) -> @out SelfReorderingGeneric<τ_0_0>.TangentVector) {
+// CHECK-LABEL: sil hidden [always_inline] [ossa] @AD__$s4main21SelfReorderingGenericV20threeParameterMethodyACyxGqd___qd_0_ts15_DifferentiableRd__sAFRd_0_s25ExpressibleByFloatLiteral13TangentVectorRpd__sAgHRpd_0_r0_lF__jvp_src_0_wrt_0_1_2 : $@convention(method) <τ_0_0 where τ_0_0 : ExpressibleByIntegerLiteral, τ_0_0 : _Differentiable><τ_1_0, τ_1_1 where τ_1_0 : _Differentiable, τ_1_1 : _Differentiable, τ_1_0.TangentVector : ExpressibleByFloatLiteral, τ_1_1.TangentVector : ExpressibleByFloatLiteral> (@in_guaranteed τ_1_0, @in_guaranteed τ_1_1, @in_guaranteed SelfReorderingGeneric<τ_0_0>) -> (@out SelfReorderingGeneric<τ_0_0>, @owned @callee_guaranteed (@in_guaranteed τ_1_0.TangentVector, @in_guaranteed τ_1_1.TangentVector, @in_guaranteed SelfReorderingGeneric<τ_0_0>.TangentVector) -> @out SelfReorderingGeneric<τ_0_0>.TangentVector) {
 // CHECK: bb0([[JVP_RESULT:%.*]] : $*SelfReorderingGeneric<τ_0_0>, [[X:%.*]] : $*τ_1_0, [[Y:%.*]] : $*τ_1_1, [[SELF:%.*]] : $*SelfReorderingGeneric<τ_0_0>):
 // CHECK: [[JVP:%.*]] = function_ref @$s4main21SelfReorderingGenericV23jvpThreeParameterMethodyACyxG_AC13TangentVectorVyx_GAH_AFQyd__AFQyd_0_tctqd___qd_0_ts15_DifferentiableRd__sAKRd_0_s25ExpressibleByFloatLiteralAIRQsAlJRQr0_lF
 // CHECK: [[DF:%.*]] = apply [[JVP]]<τ_0_0, τ_1_0, τ_1_1>([[JVP_RESULT]], [[X]], [[Y]], [[SELF]])
@@ -129,7 +129,7 @@ where Dummy: Differentiable & ExpressibleByIntegerLiteral {
 // CHECK: [[VOID:%.*]] = tuple ()
 // CHECK: return [[VOID]]
 
-// CHECK-LABEL: sil hidden [ossa] @AD__$s4main21SelfReorderingGenericV20threeParameterMethodyACyxGqd___qd_0_ts15_DifferentiableRd__sAFRd_0_s25ExpressibleByFloatLiteral13TangentVectorRpd__sAgHRpd_0_r0_lF__vjp_src_0_wrt_0_1_2 : $@convention(method) <τ_0_0 where τ_0_0 : ExpressibleByIntegerLiteral, τ_0_0 : _Differentiable><τ_1_0, τ_1_1 where τ_1_0 : _Differentiable, τ_1_1 : _Differentiable, τ_1_0.TangentVector : ExpressibleByFloatLiteral, τ_1_1.TangentVector : ExpressibleByFloatLiteral> (@in_guaranteed τ_1_0, @in_guaranteed τ_1_1, @in_guaranteed SelfReorderingGeneric<τ_0_0>) -> (@out SelfReorderingGeneric<τ_0_0>, @owned @callee_guaranteed (@in_guaranteed SelfReorderingGeneric<τ_0_0>.TangentVector) -> (@out τ_1_0.TangentVector, @out τ_1_1.TangentVector, @out SelfReorderingGeneric<τ_0_0>.TangentVector)) {
+// CHECK-LABEL: sil hidden [always_inline] [ossa] @AD__$s4main21SelfReorderingGenericV20threeParameterMethodyACyxGqd___qd_0_ts15_DifferentiableRd__sAFRd_0_s25ExpressibleByFloatLiteral13TangentVectorRpd__sAgHRpd_0_r0_lF__vjp_src_0_wrt_0_1_2 : $@convention(method) <τ_0_0 where τ_0_0 : ExpressibleByIntegerLiteral, τ_0_0 : _Differentiable><τ_1_0, τ_1_1 where τ_1_0 : _Differentiable, τ_1_1 : _Differentiable, τ_1_0.TangentVector : ExpressibleByFloatLiteral, τ_1_1.TangentVector : ExpressibleByFloatLiteral> (@in_guaranteed τ_1_0, @in_guaranteed τ_1_1, @in_guaranteed SelfReorderingGeneric<τ_0_0>) -> (@out SelfReorderingGeneric<τ_0_0>, @owned @callee_guaranteed (@in_guaranteed SelfReorderingGeneric<τ_0_0>.TangentVector) -> (@out τ_1_0.TangentVector, @out τ_1_1.TangentVector, @out SelfReorderingGeneric<τ_0_0>.TangentVector)) {
 // CHECK: bb0([[VJP_RESULT:%.*]] : $*SelfReorderingGeneric<τ_0_0>, [[X:%.*]] : $*τ_1_0, [[Y:%.*]] : $*τ_1_1, [[SELF:%.*]] : $*SelfReorderingGeneric<τ_0_0>):
 // CHECK: [[VJP:%.*]] = function_ref @$s4main21SelfReorderingGenericV23vjpThreeParameterMethodyACyxG_AC13TangentVectorVyx_G_AFQyd__AFQyd_0_tAHctqd___qd_0_ts15_DifferentiableRd__sAKRd_0_s25ExpressibleByFloatLiteralAIRQsAlJRQr0_lF
 // CHECK: [[PB:%.*]] = apply [[VJP]]<τ_0_0, τ_1_0, τ_1_1>([[VJP_RESULT]], [[X]], [[Y]], [[SELF]])