[AutoDiff] Support custom derivatives for @_alwaysEmitIntoClient functions

TODO:
- Multi file module (function and derivative in different files)
- Single file module: fix crash if we have function with
  `@_alwaysEmitIntoClient` and its derivative w/o this attribute

Fixes swiftlang#54445

Author: kovdan01

lib/SIL/IR/Linker.cpp

@@ -159,9 +159,21 @@ 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->getDeclRef().getLinkage(ForDefinition) != SILLinkage::PublicNonABI)
+      return;
+    for (SILDifferentiabilityWitness &witness :
+         Mod.getDifferentiabilityWitnesses()) {
+      if (witness.getOriginalFunction() != F)
+        continue;
+      Mod.getSILLoader()->lookupDifferentiabilityWitness(witness.getKey());
+      deserializeAndPushToWorklist(witness.getJVP());
+      deserializeAndPushToWorklist(witness.getVJP());
+    }
+    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/SILOptimizer/Differentiation/Common.cpp

@@ -537,10 +537,15 @@ SILDifferentiabilityWitness *getOrCreateMinimalASTDifferentiabilityWitness(
          "SILGen should create differentiability witnesses for all function "
          "definitions with explicit differentiable attributes");
 
+  bool isOrigAlwaysEmitIntoClient =
+      original->getDeclRef().getLinkage(ForDefinition) ==
+      SILLinkage::PublicNonABI;
   return SILDifferentiabilityWitness::createDeclaration(
-      module, SILLinkage::PublicExternal, original, kind,
-      minimalConfig->parameterIndices, minimalConfig->resultIndices,
-      minimalConfig->derivativeGenericSignature);
+      module,
+      isOrigAlwaysEmitIntoClient ? SILLinkage::PublicNonABI
+                                 : SILLinkage::PublicExternal,
+      original, kind, minimalConfig->parameterIndices,
+      minimalConfig->resultIndices, minimalConfig->derivativeGenericSignature);
 }
 
 } // end namespace autodiff

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_single_file.swift

@@ -0,0 +1,27 @@
+// REQUIRES: executable_test
+// RUN: %empty-directory(%t)
+
+// RUN: %target-build-swift-dylib(%t/%target-library-name(SingleFileModule)) %S/Inputs/always_emit_into_client/SingleFileModule/file.swift \
+// RUN:   -emit-module -emit-module-path %t/SingleFileModule.swiftmodule -module-name SingleFileModule
+// RUN: %target-build-swift -I%t -L%t %s -lm -lSingleFileModule -o %t/a.out %target-rpath(%t)
+// RUN: %target-run %t/a.out
+
+// RUN: %target-build-swift -I%t %s -emit-ir | %FileCheck %s
+
+import StdlibUnittest
+
+import SingleFileModule
+
+import _Differentiation
+
+var AlwaysEmitIntoClientTests = TestSuite("AlwaysEmitIntoClient")
+
+AlwaysEmitIntoClientTests.test("registration") {
+  expectEqual(42, gradient(at: 0, of: f))
+  expectEqual(42, gradient(at: 1, of: f))
+  expectEqual(42, gradient(at: 2, of: f))
+}
+
+runAllTests()
+
+// CHECK: @"16SingleFileModule1fyS2fFWJrSpSr" = weak_odr hidden global { ptr, ptr } { ptr @"$s16SingleFileModule1fyS2fFTJfSpSr", ptr @"$s16SingleFileModule1fyS2fFTJrSpSr" }, align 8