[AutoDiff] Support custom derivatives for @_alwaysEmitIntoClient functions

Fixes #54445

Author: kovdan01

include/swift/AST/DiagnosticsSema.def

@@ -4290,6 +4290,9 @@ NOTE(derivative_attr_fix_access,none,
      "mark the derivative function as "
      "'%select{private|fileprivate|internal|package|@usableFromInline|@usableFromInline}0' "
      "to match the original function", (AccessLevel))
+ERROR(derivative_attr_always_emit_into_client_mismatch,none,
+      "either both or none of derivative and original function must have "
+      "@alwaysEmitIntoClient attribute", ())
 ERROR(derivative_attr_static_method_mismatch_original,none,
       "unexpected derivative function declaration; "
       "%0 requires the derivative function %1 to be %select{an instance|a 'static'}2 method",

lib/SIL/IR/Linker.cpp

@@ -159,9 +159,29 @@ void SILLinkerVisitor::maybeAddFunctionToWorklist(
   // HiddenExternal linkage when they are declarations, then they
   // become Shared after the body has been deserialized.
   // So try deserializing HiddenExternal functions too.
-  if (linkage == SILLinkage::HiddenExternal)
-    return deserializeAndPushToWorklist(F);
-  
+  if (linkage == SILLinkage::HiddenExternal) {
+    deserializeAndPushToWorklist(F);
+    if (!F->markedAsAlwaysEmitIntoClient())
+      return;
+    for (SILDifferentiabilityWitness &witness :
+         F->getModule().getDifferentiabilityWitnesses()) {
+      if (witness.getOriginalFunction() != F)
+        continue;
+      SILDifferentiabilityWitness *loadedWitness =
+          F->getModule().getSILLoader()->lookupDifferentiabilityWitness(
+              witness.getKey());
+      if (loadedWitness == nullptr)
+        continue;
+      assert(loadedWitness == &witness);
+      SILFunction *jvp = loadedWitness->getJVP();
+      SILFunction *vjp = loadedWitness->getVJP();
+      assert(jvp && vjp);
+      deserializeAndPushToWorklist(jvp);
+      deserializeAndPushToWorklist(vjp);
+    }
+    return;
+  }
+
   // Update the linkage of the function in case it's different in the serialized
   // SIL than derived from the AST. This can be the case with cross-module-
   // optimizations.

lib/SILGen/SILGen.cpp

@@ -1401,13 +1401,15 @@ void SILGenModule::emitDifferentiabilityWitness(
   if (!diffWitness) {
     // Differentiability witnesses have the same linkage as the original
     // function, stripping external.
-    auto linkage = stripExternalFromLinkage(originalFunction->getLinkage());
+    auto linkage =
+        originalFunction->markedAsAlwaysEmitIntoClient()
+            ? SILLinkage::PublicNonABI
+            : stripExternalFromLinkage(originalFunction->getLinkage());
     diffWitness = SILDifferentiabilityWitness::createDefinition(
         M, linkage, originalFunction, diffKind, silConfig.parameterIndices,
         silConfig.resultIndices, config.derivativeGenericSignature,
         /*jvp*/ nullptr, /*vjp*/ nullptr,
-        /*isSerialized*/ hasPublicVisibility(originalFunction->getLinkage()),
-        attr);
+        /*isSerialized*/ hasPublicVisibility(linkage), attr);
   }
 
   // Set derivative function in differentiability witness.

lib/SILGen/SILGenPoly.cpp

@@ -6282,10 +6282,19 @@ SILFunction *SILGenModule::getOrCreateCustomDerivativeThunk(
   SILGenFunctionBuilder fb(*this);
   // Derivative thunks have the same linkage as the original function, stripping
   // external.
-  auto linkage = stripExternalFromLinkage(originalFn->getLinkage());
+  auto linkage = originalFn->markedAsAlwaysEmitIntoClient()
+                     ? SILLinkage::PublicNonABI
+                     : stripExternalFromLinkage(originalFn->getLinkage());
+
+  auto serializedKind = customDerivativeFn->getSerializedKind();
+  // See comment for an identical if statement in
+  // DifferentiationTransformer::canonicalizeDifferentiabilityWitness.
+  if (originalFn->getLinkage() == SILLinkage::HiddenExternal &&
+      !originalFn->markedAsAlwaysEmitIntoClient())
+    serializedKind = IsNotSerialized;
+
   auto *thunk = fb.getOrCreateFunction(
-      loc, name, linkage, thunkFnTy, IsBare, IsNotTransparent,
-      customDerivativeFn->getSerializedKind(),
+      loc, name, linkage, thunkFnTy, IsBare, IsNotTransparent, serializedKind,
       customDerivativeFn->isDynamicallyReplaceable(),
       customDerivativeFn->isDistributed(),
       customDerivativeFn->isRuntimeAccessible(),

lib/SILOptimizer/Differentiation/Common.cpp

@@ -538,9 +538,11 @@ SILDifferentiabilityWitness *getOrCreateMinimalASTDifferentiabilityWitness(
          "definitions with explicit differentiable attributes");
 
   return SILDifferentiabilityWitness::createDeclaration(
-      module, SILLinkage::PublicExternal, original, kind,
-      minimalConfig->parameterIndices, minimalConfig->resultIndices,
-      minimalConfig->derivativeGenericSignature);
+      module,
+      original->markedAsAlwaysEmitIntoClient() ? SILLinkage::PublicNonABI
+                                               : SILLinkage::PublicExternal,
+      original, kind, minimalConfig->parameterIndices,
+      minimalConfig->resultIndices, minimalConfig->derivativeGenericSignature);
 }
 
 } // end namespace autodiff

lib/SILOptimizer/Mandatory/Differentiation.cpp

@@ -914,7 +914,8 @@ bool DifferentiationTransformer::canonicalizeDifferentiabilityWitness(
   // flag. Important exception here hidden_external functions as they are
   // serializable but corresponding hidden ones would be not and the SIL
   // verifier will fail. Patch `serializeFunctions` for this case.
-  if (orig->getLinkage() == SILLinkage::HiddenExternal)
+  if (orig->getLinkage() == SILLinkage::HiddenExternal &&
+      !orig->markedAsAlwaysEmitIntoClient())
     serializeFunctions = IsNotSerialized;
 
   // If the JVP doesn't exist, need to synthesize it.

lib/Sema/TypeCheckAttr.cpp

@@ -6743,6 +6743,13 @@ static bool typeCheckDerivativeAttr(DerivativeAttr *attr) {
     return true;
   }
 
+  if (originalAFD->getAttrs().hasAttribute<AlwaysEmitIntoClientAttr>() !=
+      derivative->getAttrs().hasAttribute<AlwaysEmitIntoClientAttr>()) {
+    diags.diagnose(derivative->getLoc(),
+                   diag::derivative_attr_always_emit_into_client_mismatch);
+    return true;
+  }
+
   // Get the resolved differentiability parameter indices.
   auto *resolvedDiffParamIndices = attr->getParameterIndices();
 

stdlib/public/Differentiation/SIMDDifferentiation.swift.gyb

@@ -405,9 +405,6 @@ where
   }
 }
 
-// FIXME(TF-1103): Derivative registration does not yet support
-// `@_alwaysEmitIntoClient` original functions like `SIMD.sum()`.
-/*
 extension SIMD
 where
   Self: Differentiable,
@@ -417,6 +414,7 @@ where
   TangentVector == Self
 {
   @inlinable
+  @_alwaysEmitIntoClient
   @derivative(of: sum)
   func _vjpSum() -> (
     value: Scalar, pullback: (Scalar.TangentVector) -> TangentVector
@@ -425,14 +423,14 @@ where
   }
 
   @inlinable
+  @_alwaysEmitIntoClient
   @derivative(of: sum)
   func _jvpSum() -> (
     value: Scalar, differential: (TangentVector) -> Scalar.TangentVector
   ) {
     return (sum(), { v in Scalar.TangentVector(v.sum()) })
   }
 }
-*/
 
 extension SIMD
 where

test/AutoDiff/SILGen/nil_coalescing.swift

@@ -1,8 +1,9 @@
-// RUN: %target-swift-frontend -Xllvm -sil-print-types -emit-sil -verify %s | %FileCheck %s
+/// Note: -primary-file prevents non_abi->shared linkage change in `removeSerializedFlagFromAllFunctions`
+// RUN: %target-swift-frontend -Xllvm -sil-print-types -emit-sil -verify -primary-file %s | %FileCheck %s
 
 import _Differentiation
 
-// CHECK: sil @test_nil_coalescing
+// CHECK: sil non_abi @test_nil_coalescing
 // CHECK: bb0(%{{.*}} : $*T, %[[ARG_OPT:.*]] : $*Optional<T>, %[[ARG_PB:.*]] :
 // CHECK:    $@noescape @callee_guaranteed @substituted <τ_0_0> () -> (@out τ_0_0, @error any Error) for <T>):
 // CHECK: %[[ALLOC_OPT:.*]] = alloc_stack [lexical] $Optional<T>
@@ -15,7 +16,7 @@ import _Differentiation
 //
 @_silgen_name("test_nil_coalescing")
 @derivative(of: ??)
-@usableFromInline
+@_alwaysEmitIntoClient
 func nilCoalescing<T: Differentiable>(optional: T?, defaultValue: @autoclosure () throws -> T)
  rethrows -> (value: T, pullback: (T.TangentVector) -> Optional<T>.TangentVector)
 {

test/AutoDiff/Sema/derivative_attr_type_checking.swift

@@ -1062,6 +1062,15 @@ func _internal_original_inlinable_derivative(_ x: Float) -> (value: Float, pullb
   fatalError()
 }
 
+func internal_original_alwaysemitintoclient_derivative_error(_ x: Float) -> Float { x }
+@_alwaysEmitIntoClient
+@derivative(of: internal_original_alwaysemitintoclient_derivative_error)
+// expected-error @+1 {{either both or none of derivative and original function must have @alwaysEmitIntoClient attribute}}
+func _internal_original_alwaysemitintoclient_derivative_error(_ x: Float) -> (value: Float, pullback: (Float) -> Float) {
+  fatalError()
+}
+
+@_alwaysEmitIntoClient
 func internal_original_alwaysemitintoclient_derivative(_ x: Float) -> Float { x }
 @_alwaysEmitIntoClient
 @derivative(of: internal_original_alwaysemitintoclient_derivative)
@@ -1084,6 +1093,15 @@ package func _package_original_inlinable_derivative(_ x: Float) -> (value: Float
   fatalError()
 }
 
+@_alwaysEmitIntoClient
+package func package_original_alwaysemitintoclient_derivative_error(_ x: Float) -> Float { x }
+@derivative(of: package_original_alwaysemitintoclient_derivative_error)
+// expected-error @+1 {{either both or none of derivative and original function must have @alwaysEmitIntoClient attribute}}
+package func _package_original_alwaysemitintoclient_derivative_error(_ x: Float) -> (value: Float, pullback: (Float) -> Float) {
+  fatalError()
+}
+
+@_alwaysEmitIntoClient
 package func package_original_alwaysemitintoclient_derivative(_ x: Float) -> Float { x }
 @_alwaysEmitIntoClient
 @derivative(of: package_original_alwaysemitintoclient_derivative)

test/AutoDiff/stdlib/simd.swift

@@ -19,9 +19,6 @@ SIMDTests.test("init(repeating:)") {
   expectEqual(8, pb1(g))
 }
 
-// FIXME(TF-1103): Derivative registration does not yet support
-// `@_alwaysEmitIntoClient` original functions.
-/*
 SIMDTests.test("Sum") {
   let a = SIMD4<Float>(1, 2, 3, 4)
 
@@ -32,7 +29,6 @@ SIMDTests.test("Sum") {
   expectEqual(10, val1)
   expectEqual(SIMD4<Float>(3, 3, 3, 3), pb1(3))
 }
-*/
 
 SIMDTests.test("Identity") {
   let a = SIMD4<Float>(1, 2, 3, 4)
@@ -289,9 +285,6 @@ SIMDTests.test("Generics") {
   expectEqual(SIMD3<Double>(5, 10, 15), val4)
   expectEqual((SIMD3<Double>(5, 5, 5), 6), pb4(g))
 
-  // FIXME(TF-1103): Derivative registration does not yet support
-  // `@_alwaysEmitIntoClient` original functions like `SIMD.sum()`.
-  /*
   func testSum<Scalar, SIMDType: SIMD>(x: SIMDType) -> Scalar
     where SIMDType.Scalar == Scalar,
           SIMDType : Differentiable,
@@ -304,7 +297,6 @@ SIMDTests.test("Generics") {
   let (val5, pb5) = valueWithPullback(at: a, of: simd3Sum)
   expectEqual(6, val5)
   expectEqual(SIMD3<Double>(7, 7, 7), pb5(7))
-  */
 }
 
 runAllTests()

test/AutoDiff/validation-test/Inputs/always_emit_into_client/MultiFileModule/file1.swift

@@ -0,0 +1,4 @@
+@_alwaysEmitIntoClient
+public func f(_ x: Float) -> Float {
+  x
+}

test/AutoDiff/validation-test/Inputs/always_emit_into_client/MultiFileModule/file2.swift

@@ -0,0 +1,7 @@
+import _Differentiation
+
+@derivative(of: f)
+@_alwaysEmitIntoClient
+public func df(_ x: Float) -> (value: Float, pullback: (Float) -> Float) {
+  (x, { 42 * $0 })
+}

test/AutoDiff/validation-test/Inputs/always_emit_into_client/SingleFileModule/file.swift

@@ -0,0 +1,12 @@
+import _Differentiation
+
+@_alwaysEmitIntoClient
+public func f(_ x: Float) -> Float {
+  x
+}
+
+@derivative(of: f)
+@_alwaysEmitIntoClient
+public func df(_ x: Float) -> (value: Float, pullback: (Float) -> Float) {
+  (x, { 42 * $0 })
+}

test/AutoDiff/validation-test/always_emit_into_client/Inputs/MultiFileModule/file1.swift

@@ -0,0 +1,4 @@
+@_alwaysEmitIntoClient
+public func f(_ x: Float) -> Float {
+  x
+}

test/AutoDiff/validation-test/always_emit_into_client/Inputs/MultiFileModule/file2.swift

@@ -0,0 +1,7 @@
+import _Differentiation
+
+@derivative(of: f)
+@_alwaysEmitIntoClient
+public func df(_ x: Float) -> (value: Float, pullback: (Float) -> Float) {
+  (x, { 42 * $0 })
+}

test/AutoDiff/validation-test/always_emit_into_client/Inputs/MultiModule/file1.swift

@@ -0,0 +1,4 @@
+@_alwaysEmitIntoClient
+public func f(_ x: Float) -> Float {
+  x
+}

test/AutoDiff/validation-test/always_emit_into_client/Inputs/MultiModule/file2.swift

@@ -0,0 +1,8 @@
+import MultiModule1
+import _Differentiation
+
+@derivative(of: f)
+@_alwaysEmitIntoClient
+public func df(_ x: Float) -> (value: Float, pullback: (Float) -> Float) {
+  (x, { 42 * $0 })
+}

test/AutoDiff/validation-test/always_emit_into_client/Inputs/MultiModuleProtocol/file1.swift

@@ -0,0 +1,16 @@
+import _Differentiation
+
+public protocol Protocol {
+  var x : Float {get set}
+  init()
+}
+
+extension Protocol {
+  public init(_ val: Float) {
+    self.init()
+    x = val
+  }
+
+  @_alwaysEmitIntoClient
+  public func sum() -> Float { x }
+}

test/AutoDiff/validation-test/always_emit_into_client/Inputs/MultiModuleProtocol/file2.swift

@@ -0,0 +1,20 @@
+import MultiModuleProtocol1
+import _Differentiation
+
+extension Protocol where Self: Differentiable, Self.TangentVector == Self {
+  @_alwaysEmitIntoClient
+  @derivative(of: sum)
+  public func _vjpSum() -> (
+    value: Float, pullback: (Float) -> Self.TangentVector
+  ) {
+    (value: self.x, pullback: { Self.TangentVector(42 * $0) })
+  }
+
+  @_alwaysEmitIntoClient
+  @derivative(of: sum)
+  public func _jvpSum() -> (
+    value: Float, differential: (Self.TangentVector) -> Float
+  ) {
+    (value: self.x, differential: { 42 * $0.x })
+  }
+}

test/AutoDiff/validation-test/always_emit_into_client/Inputs/MultiModuleProtocol/file3.swift

@@ -0,0 +1,24 @@
+import MultiModuleProtocol1
+import MultiModuleProtocol2
+import _Differentiation
+
+public struct Struct : Protocol {
+  private var _x : Float
+  public  var x : Float {
+    get { _x }
+    set { _x = newValue }
+  }
+  public init() { _x = 0 }
+}
+
+extension Struct : AdditiveArithmetic {
+  public static func +(lhs: Self, rhs: Self) -> Self { Self(lhs.x + rhs.x) }
+  public static func -(lhs: Self, rhs: Self) -> Self { Self(lhs.x - rhs.x) }
+  public static func +=(a: inout Self, b: Self) { a.x = a.x + b.x }
+  public static func -=(a: inout Self, b: Self) { a.x = a.x - b.x }
+  public static var zero: Self { Self(0) }
+}
+
+extension Struct : Differentiable {
+  public  typealias TangentVector = Self
+}

test/AutoDiff/validation-test/always_emit_into_client/Inputs/MultiModuleStruct/file1.swift

@@ -0,0 +1,23 @@
+public struct Struct {
+  public var x : Float
+  public typealias TangentVector = Self
+  public init() { x = 0 }
+}
+
+extension Struct {
+  public init(_ val: Float) {
+    self.init()
+    x = val
+  }
+
+  @_alwaysEmitIntoClient
+  public func sum() -> Float { x }
+}
+
+extension Struct : AdditiveArithmetic {
+  public static func +(lhs: Self, rhs: Self) -> Self { lhs }
+  public static func -(lhs: Self, rhs: Self) -> Self { lhs }
+  public static func +=(a: inout Self, b: Self) { a }
+  public static func -=(a: inout Self, b: Self) { a }
+  public static var zero: Self { Self(0) }
+}