[AutoDiff upstream] Add TBDGen for AutoDiff symbols.

Emit symbols for `@differentiable` and `@derivative` declaration attributes:
- Differentiability witness symbols.
- Derivative function (JVP/VJP) symbols.
- Linear map (differential/pullback) symbols.

Add TBDGen test.

Author: dan-zheng

lib/TBDGen/TBDGen.cpp

@@ -502,6 +502,95 @@ void TBDGenVisitor::addConformances(DeclContext *DC) {
   }
 }
 
+void TBDGenVisitor::addAutoDiffLinearMapFunction(AbstractFunctionDecl *original,
+                                                 AutoDiffConfig config,
+                                                 AutoDiffLinearMapKind kind) {
+  auto &ctx = original->getASTContext();
+  auto declRef =
+      SILDeclRef(original).asForeign(requiresForeignEntryPoint(original));
+
+  if (!declRef.isSerialized())
+    return;
+  // Linear maps are public only when the original function is serialized.
+  if (!declRef.isSerialized())
+    return;
+  // Differential functions are emitted only when forward-mode is enabled.
+  if (kind == AutoDiffLinearMapKind::Differential &&
+      !ctx.LangOpts.EnableExperimentalForwardModeDifferentiation)
+    return;
+  auto *loweredParamIndices = autodiff::getLoweredParameterIndices(
+      config.parameterIndices,
+      original->getInterfaceType()->castTo<AnyFunctionType>());
+  Mangle::ASTMangler mangler;
+  AutoDiffConfig silConfig{loweredParamIndices, config.resultIndices,
+                           config.derivativeGenericSignature};
+  std::string linearMapName =
+      mangler.mangleAutoDiffLinearMapHelper(declRef.mangle(), kind, silConfig);
+  addSymbol(linearMapName);
+}
+
+void TBDGenVisitor::addAutoDiffDerivativeFunction(
+    AbstractFunctionDecl *original, IndexSubset *parameterIndices,
+    GenericSignature derivativeGenericSignature,
+    AutoDiffDerivativeFunctionKind kind) {
+  auto *assocFnId = AutoDiffDerivativeFunctionIdentifier::get(
+      kind, parameterIndices, derivativeGenericSignature,
+      original->getASTContext());
+  auto declRef =
+      SILDeclRef(original).asForeign(requiresForeignEntryPoint(original));
+  addSymbol(declRef.asAutoDiffDerivativeFunction(assocFnId));
+}
+
+void TBDGenVisitor::addDifferentiabilityWitness(
+    AbstractFunctionDecl *original, IndexSubset *astParameterIndices,
+    IndexSubset *resultIndices, GenericSignature derivativeGenericSignature) {
+  bool foreign = requiresForeignEntryPoint(original);
+  auto declRef = SILDeclRef(original).asForeign(foreign);
+
+  // Skip symbol emission for original functions that do not have public
+  // linkage. Exclude original functions that require a foreign entry point with
+  // `public_external` linkage.
+  auto originalLinkage = declRef.getLinkage(ForDefinition);
+  if (foreign)
+    originalLinkage = stripExternalFromLinkage(originalLinkage);
+  if (originalLinkage != SILLinkage::Public)
+    return;
+
+  auto *silParamIndices = autodiff::getLoweredParameterIndices(
+      astParameterIndices,
+      original->getInterfaceType()->castTo<AnyFunctionType>());
+
+  auto originalMangledName = declRef.mangle();
+  AutoDiffConfig config{silParamIndices, resultIndices,
+                        derivativeGenericSignature};
+  SILDifferentiabilityWitnessKey key(originalMangledName, config);
+
+  Mangle::ASTMangler mangler;
+  auto mangledName = mangler.mangleSILDifferentiabilityWitnessKey(key);
+  addSymbol(mangledName);
+}
+
+void TBDGenVisitor::addDerivativeConfiguration(AbstractFunctionDecl *original,
+                                               AutoDiffConfig config) {
+  auto inserted = AddedDerivatives.insert({original, config});
+  if (!inserted.second)
+    return;
+
+  addAutoDiffLinearMapFunction(original, config,
+                               AutoDiffLinearMapKind::Differential);
+  addAutoDiffLinearMapFunction(original, config,
+                               AutoDiffLinearMapKind::Pullback);
+  addAutoDiffDerivativeFunction(original, config.parameterIndices,
+                                config.derivativeGenericSignature,
+                                AutoDiffDerivativeFunctionKind::JVP);
+  addAutoDiffDerivativeFunction(original, config.parameterIndices,
+                                config.derivativeGenericSignature,
+                                AutoDiffDerivativeFunctionKind::VJP);
+  addDifferentiabilityWitness(original, config.parameterIndices,
+                              config.resultIndices,
+                              config.derivativeGenericSignature);
+}
+
 /// Determine whether dynamic replacement should be emitted for the allocator or
 /// the initializer given a decl.
 /// The rule is that structs and convenience init of classes emit a
@@ -565,6 +654,22 @@ void TBDGenVisitor::visitAbstractFunctionDecl(AbstractFunctionDecl *AFD) {
     addSymbol(SILDeclRef(AFD).asForeign());
   }
 
+  // Add derivative function symbols.
+  for (const auto *differentiableAttr :
+       AFD->getAttrs().getAttributes<DifferentiableAttr>())
+    addDerivativeConfiguration(
+        AFD,
+        AutoDiffConfig(differentiableAttr->getParameterIndices(),
+                       IndexSubset::get(AFD->getASTContext(), 1, {0}),
+                       differentiableAttr->getDerivativeGenericSignature()));
+  for (const auto *derivativeAttr :
+       AFD->getAttrs().getAttributes<DerivativeAttr>())
+    addDerivativeConfiguration(
+        derivativeAttr->getOriginalFunction(),
+        AutoDiffConfig(derivativeAttr->getParameterIndices(),
+                       IndexSubset::get(AFD->getASTContext(), 1, {0}),
+                       AFD->getGenericSignature()));
+
   visitDefaultArguments(AFD, AFD->getParameters());
 }
 
@@ -617,6 +722,15 @@ void TBDGenVisitor::visitAbstractStorageDecl(AbstractStorageDecl *ASD) {
   ASD->visitEmittedAccessors([&](AccessorDecl *accessor) {
     visitFuncDecl(accessor);
   });
+
+  // Add derivative function symbols.
+  for (const auto *differentiableAttr :
+       ASD->getAttrs().getAttributes<DifferentiableAttr>())
+    addDerivativeConfiguration(
+        ASD->getAccessor(AccessorKind::Get),
+        AutoDiffConfig(differentiableAttr->getParameterIndices(),
+                       IndexSubset::get(ASD->getASTContext(), 1, {0}),
+                       differentiableAttr->getDerivativeGenericSignature()));
 }
 
 void TBDGenVisitor::visitVarDecl(VarDecl *VD) {

lib/TBDGen/TBDGenVisitor.h

@@ -70,6 +70,14 @@ class TBDGenVisitor : public ASTVisitor<TBDGenVisitor> {
   ModuleDecl *SwiftModule;
   const TBDGenOptions &Opts;
 
+  /// A set of original function and derivative configuration pairs for which
+  /// derivative symbols have been emitted.
+  ///
+  /// Used to deduplicate derivative symbol emission for `@differentiable` and
+  /// `@derivative` attributes.
+  llvm::DenseSet<std::pair<AbstractFunctionDecl *, AutoDiffConfig>>
+      AddedDerivatives;
+
 private:
   std::vector<Decl*> DeclStack;
   std::unique_ptr<std::map<std::string, InstallNameStore>>
@@ -98,6 +106,34 @@ class TBDGenVisitor : public ASTVisitor<TBDGenVisitor> {
   void addAssociatedConformanceDescriptor(AssociatedConformance conformance);
   void addBaseConformanceDescriptor(BaseConformance conformance);
 
+  /// Adds the symbol for the linear map function of the given kind associated
+  /// with the given original function and derivative function configuration.
+  void addAutoDiffLinearMapFunction(AbstractFunctionDecl *original,
+                                    AutoDiffConfig config,
+                                    AutoDiffLinearMapKind kind);
+
+  /// Adds the symbol for the autodiff function of the given kind associated
+  /// with the given original function, parameter indices, and derivative
+  /// generic signature.
+  void
+  addAutoDiffDerivativeFunction(AbstractFunctionDecl *original,
+                                IndexSubset *parameterIndices,
+                                GenericSignature derivativeGenericSignature,
+                                AutoDiffDerivativeFunctionKind kind);
+
+  /// Adds the symbol for the differentiability witness associated with the
+  /// given original function, AST parameter indices, result indices, and
+  /// derivative generic signature.
+  void addDifferentiabilityWitness(AbstractFunctionDecl *original,
+                                   IndexSubset *astParameterIndices,
+                                   IndexSubset *resultIndices,
+                                   GenericSignature derivativeGenericSignature);
+
+  /// Adds symbols associated with the given original function and
+  /// derivative function configuration.
+  void addDerivativeConfiguration(AbstractFunctionDecl *original,
+                                  AutoDiffConfig config);
+
 public:
   TBDGenVisitor(llvm::MachO::InterfaceFile &symbols,
                 llvm::MachO::TargetList targets, StringSet *stringSymbols,

test/AutoDiff/TBD/derivative_symbols.swift

@@ -0,0 +1,56 @@
+// RUN: %target-swift-frontend -emit-ir -o/dev/null -parse-as-library -module-name test -validate-tbd-against-ir=all %s
+// RUN: %target-swift-frontend -emit-ir -o/dev/null -parse-as-library -module-name test -validate-tbd-against-ir=all %s -O
+// RUN: %target-swift-frontend -emit-ir -o/dev/null -parse-as-library -module-name test -validate-tbd-against-ir=missing %s -enable-testing
+// RUN: %target-swift-frontend -emit-ir -o/dev/null -parse-as-library -module-name test -validate-tbd-against-ir=missing %s -enable-testing -O
+
+import _Differentiation
+
+@differentiable
+public func topLevelDifferentiable(_ x: Float, _ y: Float) -> Float { x }
+
+public func topLevelHasDerivative<T: Differentiable>(_ x: T) -> T {
+  x
+}
+
+@derivative(of: topLevelHasDerivative)
+public func topLevelDerivative<T: Differentiable>(_ x: T) -> (
+  value: T, pullback: (T.TangentVector) -> T.TangentVector
+) {
+  fatalError()
+}
+
+struct Struct: Differentiable {
+  var stored: Float
+
+  // Test property.
+  @differentiable
+  public var property: Float {
+    stored
+  }
+
+  // Test initializer.
+  @differentiable
+  public init(_ x: Float) {
+    stored = x
+  }
+
+  // Test method.
+  public func method(x: Float, y: Float) -> Float { x }
+
+  @derivative(of: method)
+  public func jvpMethod(x: Float, y: Float) -> (
+    value: Float, differential: (TangentVector, Float, Float) -> Float
+  ) {
+    fatalError()
+  }
+
+  // Test subscript.
+  public subscript(x: Float) -> Float { x }
+
+  @derivative(of: subscript)
+  public func vjpSubscript(x: Float) -> (
+    value: Float, pullback: (Float) -> (TangentVector, Float)
+  ) {
+    fatalError()
+  }
+}