[AutoDiff] Serialize derivative function configurations per module.

`@differentiable` and `@derivative` attributes register derivatives for
`AbstractFunctionDecl`s for a particular "derivative function configuration":
parameter indices and dervative generic signature.

To find `@derivative` functions registered in other Swift modules, derivative
function configurations must be serialized per module. When configurations for
a `AbstractFunctionDecl` are requested, all configurations from imported
modules are deserialized. This module serialization technique has precedent: it
is used for protocol conformances (e.g. extension declarations for a nominal
type) and Obj-C members for a class type.

Add `AbstractFunctionDecl::getDerivativeFunctionConfigurations` entry point
for accessing derivative function configurations.

Use `AbstractFunctionDecl::getDerivativeFunctionConfigurations` to
implement `findMinimalDerivativeConfiguration` for canonical derivative
function configuration lookup, replacing `getMinimalASTDifferentiableAttr`.

Unblocks TF-815: lowering `@derivative` attributes directly to SIL
differentiability witnesses without generating implicit `@differentiable`
attributes.

Author: dan-zheng

include/swift/AST/ASTContext.h

@@ -110,8 +110,13 @@ namespace swift {
   class VarDecl;
   class UnifiedStatsReporter;
   class IndexSubset;
-  class VectorSpace;
+  // SWIFT_ENABLE_TENSORFLOW
+  struct AutoDiffConfig;
+  struct AutoDiffDerivativeFunctionKind;
+  class DerivativeAttr;
   class DifferentiableAttr;
+  class VectorSpace;
+  // SWIFT_ENABLE_TENSORFLOW END
 
   enum class KnownProtocolKind : uint8_t;
 
@@ -287,11 +292,18 @@ class ASTContext final {
   /// Cache of autodiff-associated vector spaces.
   llvm::DenseMap<Type, Optional<VectorSpace>> AutoDiffVectorSpaces;
 
-  /// Cache of `@differentiable` attributes keyed by parameter indices. This
-  /// helps us diagnose multiple `@differentiable`s that are with respect to the
-  /// same set of parameters.
+  /// Cache of `@differentiable` attributes keyed by parameter indices. Used to
+  /// diagnose duplicate `@differentiable` attributes for the same key.
   llvm::DenseMap<std::pair<Decl *, IndexSubset *>, DifferentiableAttr *>
       DifferentiableAttrs;
+
+  /// Cache of `@derivative` attributes keyed by parameter indices and
+  /// derivative function kind. Used to diagnose duplicate `@derivative`
+  /// attributes for the same key.
+  llvm::DenseMap<
+      std::tuple<Decl *, IndexSubset *, AutoDiffDerivativeFunctionKind>,
+      DerivativeAttr *>
+      DerivativeAttrs;
   // SWIFT_ENABLE_TENSORFLOW END
 
 private:
@@ -702,6 +714,21 @@ class ASTContext final {
                        unsigned previousGeneration,
                        llvm::TinyPtrVector<AbstractFunctionDecl *> &methods);
 
+  // SWIFT_ENABLE_TENSORFLOW
+  /// Load derivative function configurations for the given
+  /// AbstractFunctionDecl.
+  ///
+  /// \param originalAFD The declaration whose derivative function
+  /// configurations should be loaded.
+  ///
+  /// \param previousGeneration The previous generation number. The AST already
+  /// contains derivative function configurations loaded from any generation up
+  /// to and including this one.
+  void loadDerivativeFunctionConfigurations(
+      AbstractFunctionDecl *originalAFD, unsigned previousGeneration,
+      llvm::SetVector<AutoDiffConfig> &results);
+  // SWIFT_ENABLE_TENSORFLOW END
+
   /// Retrieve the Clang module loader for this ASTContext.
   ///
   /// If there is no Clang module loader, returns a null pointer.

include/swift/AST/Attr.h

@@ -1947,6 +1947,8 @@ class DerivativeAttr final
   unsigned NumParsedParameters = 0;
   /// The differentiation parameters' indices, resolved by the type checker.
   IndexSubset *ParameterIndices = nullptr;
+  /// The derivative function kind (JVP or VJP), resolved by the type checker.
+  Optional<AutoDiffDerivativeFunctionKind> Kind = None;
 
   explicit DerivativeAttr(bool implicit, SourceLoc atLoc, SourceRange baseRange,
                           DeclNameWithLoc original,
@@ -1975,6 +1977,12 @@ class DerivativeAttr final
     OriginalFunction = decl;
   }
 
+  AutoDiffDerivativeFunctionKind getDerivativeKind() const {
+    assert(Kind && "Derivative function kind has not yet been resolved");
+    return *Kind;
+  }
+  void setDerivativeKind(AutoDiffDerivativeFunctionKind kind) { Kind = kind; }
+
   /// The parsed differentiation parameters, i.e. the list of parameters
   /// specified in 'wrt:'.
   ArrayRef<ParsedAutoDiffParameter> getParsedParameters() const {

include/swift/AST/AutoDiff.h

@@ -306,6 +306,7 @@ struct AutoDiffConfig {
 class AutoDiffDerivativeFunctionIdentifier : public llvm::FoldingSetNode {
   const AutoDiffDerivativeFunctionKind kind;
   IndexSubset *const parameterIndices;
+  // TODO(TF-680): Mangle derivative generic signature requirements as well.
 
   AutoDiffDerivativeFunctionIdentifier(
       AutoDiffDerivativeFunctionKind kind, IndexSubset *parameterIndices) :
@@ -508,6 +509,27 @@ template<> struct DenseMapInfo<AutoDiffConfig> {
   }
 };
 
+template<> struct DenseMapInfo<AutoDiffDerivativeFunctionKind> {
+  static AutoDiffDerivativeFunctionKind getEmptyKey() {
+    return static_cast<AutoDiffDerivativeFunctionKind::innerty>(
+        DenseMapInfo<unsigned>::getEmptyKey());
+  }
+
+  static AutoDiffDerivativeFunctionKind getTombstoneKey() {
+    return static_cast<AutoDiffDerivativeFunctionKind::innerty>(
+        DenseMapInfo<unsigned>::getTombstoneKey());
+  }
+
+  static unsigned getHashValue(const AutoDiffDerivativeFunctionKind &Val) {
+    return DenseMapInfo<unsigned>::getHashValue(Val);
+  }
+
+  static bool isEqual(const AutoDiffDerivativeFunctionKind &LHS,
+                      const AutoDiffDerivativeFunctionKind &RHS) {
+    return LHS == RHS;
+  }
+};
+
 template<> struct DenseMapInfo<SILAutoDiffIndices> {
   static SILAutoDiffIndices getEmptyKey() {
     return { DenseMapInfo<unsigned>::getEmptyKey(), nullptr };

include/swift/AST/Decl.h

@@ -5694,6 +5694,25 @@ class AbstractFunctionDecl : public GenericContext, public ValueDecl {
 private:
   ParameterList *Params;
 
+// SWIFT_ENABLE_TENSORFLOW
+private:
+  /// The generation at which we last loaded derivative function configurations.
+  unsigned DerivativeFunctionConfigGeneration = 0;
+  /// Prepare to traverse the list of derivative function configurations.
+  void prepareDerivativeFunctionConfigurations();
+
+  /// A uniqued list of derivative function configurations.
+  struct DerivativeFunctionConfigurationList;
+  DerivativeFunctionConfigurationList *DerivativeFunctionConfigs = nullptr;
+
+public:
+  /// Get all derivative function configurations.
+  ArrayRef<AutoDiffConfig> getDerivativeFunctionConfigurations();
+
+  /// Add the given derivative function configuration.
+  void addDerivativeFunctionConfiguration(AutoDiffConfig config);
+// SWIFT_ENABLE_TENSORFLOW END
+
 protected:
   // If a function has a body at all, we have either a parsed body AST node or
   // we have saved the end location of the unparsed body.

include/swift/AST/ModuleLoader.h

@@ -34,6 +34,9 @@ class DependencyCollector;
 
 namespace swift {
 
+// SWIFT_ENABLE_TENSORFLOW
+struct AutoDiffConfig;
+// SWIFT_ENABLE_TENSORFLOW END
 class AbstractFunctionDecl;
 class ClangImporterOptions;
 class ClassDecl;
@@ -151,6 +154,25 @@ class ModuleLoader {
                  unsigned previousGeneration,
                  llvm::TinyPtrVector<AbstractFunctionDecl *> &methods) = 0;
 
+  // SWIFT_ENABLE_TENSORFLOW
+  /// Load derivative function configurations for the given
+  /// AbstractFunctionDecl.
+  ///
+  /// \param originalAFD The declaration whose derivative function
+  /// configurations should be loaded.
+  ///
+  /// \param previousGeneration The previous generation number. The AST already
+  /// contains derivative function configurations loaded from any generation up
+  /// to and including this one.
+  ///
+  /// \param results The result list of derivative function configurations.
+  /// This list will be extended with any methods found in subsequent
+  /// generations.
+  virtual void loadDerivativeFunctionConfigurations(
+      AbstractFunctionDecl *originalAFD, unsigned previousGeneration,
+      llvm::SetVector<AutoDiffConfig> &results){};
+  // SWIFT_ENABLE_TENSORFLOW END
+
   /// Verify all modules loaded by this loader.
   virtual void verifyAllModules() { }
 };

include/swift/AST/Types.h

@@ -4359,7 +4359,7 @@ class SILFunctionType final : public TypeBase, public llvm::FoldingSetNode,
   ///   is necessary for the differentiation transform to support reabstraction
   ///   thunk differentiation because the function argument is opaque and cannot
   ///   be differentiated. Instead, the argument is made `@differentiable` and
-  ///   reabstraction thunk JVP/VJP callers are reponsible for passing a
+  ///   reabstraction thunk JVP/VJP callers are responsible for passing a
   ///   `@differentiable` function.
   ///   - TODO(TF-1036): Investigate more efficient reabstraction thunk
   ///     derivative approaches. The last argument can simply be a

include/swift/SILOptimizer/Utils/Differentiation/DerivativeLookup.h

@@ -36,17 +36,26 @@ getExactDifferentiabilityWitness(SILModule &module, SILFunction *original,
                                  IndexSubset *parameterIndices,
                                  IndexSubset *resultIndices);
 
-/// Finds the "@differentiable" attribute on `original` whose parameter indices
-/// are a minimal superset of the specified parameter indices. Returns `nullptr`
-/// if no such attribute exists.
+/// Finds the derivative configuration (from `@differentiable` and
+/// `@derivative` attributes) for `original` whose parameter indices are a
+/// minimal superset of the specified AST parameter indices. Returns true if
+/// such a configuration is found.
 ///
 /// \param parameterIndices must be lowered to SIL.
-/// \param minimalParameterIndices is an output parameter that is set to the SIL
-/// indices of the minimal attribute, or to `nullptr` if no attribute exists.
-const DifferentiableAttr *
-getMinimalASTDifferentiableAttr(AbstractFunctionDecl *original,
-                                IndexSubset *parameterIndices,
-                                IndexSubset *&minimalParameterIndices);
+/// \param minimalASTParameterIndices is an output parameter that is set to the
+/// AST indices of the minimal configuration, or to `nullptr` if no such
+/// configuration exists.
+/// \param minimalSILParameterIndices is an output parameter that is set to the
+/// SIL indices of the minimal configuration, or to `nullptr` if no such
+/// configuration exists.
+/// \param derivativeGenericSignature is an output parameter that is set to the
+/// derivative generic signature of the minimal configuration, or the `nullptr`
+/// if no such configuration exists.
+bool findMinimalDerivativeConfiguration(
+    AbstractFunctionDecl *original, IndexSubset *parameterIndices,
+    IndexSubset *&minimalASTParameterIndices,
+    IndexSubset *&minimalSILParameterIndices,
+    GenericSignature &derivativeGenericSignature);
 
 /// Returns a differentiability witness for `original` whose parameter indices
 /// are a minimal superset of the specified parameter indices and whose result

include/swift/Serialization/SerializedModuleLoader.h

@@ -166,6 +166,12 @@ class SerializedModuleLoaderBase : public ModuleLoader {
                  unsigned previousGeneration,
                  llvm::TinyPtrVector<AbstractFunctionDecl *> &methods) override;
 
+  // SWIFT_ENABLE_TENSORFLOW
+  virtual void loadDerivativeFunctionConfigurations(
+      AbstractFunctionDecl *originalAFD, unsigned previousGeneration,
+      llvm::SetVector<AutoDiffConfig> &results) override;
+  // SWIFT_ENABLE_TENSORFLOW END
+
   virtual void verifyAllModules() override;
 };
 

lib/AST/ASTContext.cpp

@@ -1611,6 +1611,19 @@ void ASTContext::loadObjCMethods(
   }
 }
 
+// SWIFT_ENABLE_TENSORFLOW
+void ASTContext::loadDerivativeFunctionConfigurations(
+    AbstractFunctionDecl *originalAFD, unsigned previousGeneration,
+    llvm::SetVector<AutoDiffConfig> &results) {
+  PrettyStackTraceDecl stackTrace(
+      "loading derivative function configurations for", originalAFD);
+  for (auto &loader : getImpl().ModuleLoaders) {
+    loader->loadDerivativeFunctionConfigurations(originalAFD,
+                                                 previousGeneration, results);
+  }
+}
+// SWIFT_ENABLE_TENSORFLOW END
+
 void ASTContext::verifyAllLoadedModules() const {
 #ifndef NDEBUG
   FrontendStatsTracer tracer(Stats, "verify-all-loaded-modules");

lib/AST/Decl.cpp

@@ -6998,6 +6998,49 @@ StringRef AbstractFunctionDecl::getInlinableBodyText(
   return extractInlinableText(getASTContext().SourceMgr, body, scratch);
 }
 
+// SWIFT_ENABLE_TENSORFLOW
+/// A uniqued list of derivative function configurations.
+struct AbstractFunctionDecl::DerivativeFunctionConfigurationList
+    : public llvm::SetVector<AutoDiffConfig> {
+  // Necessary for `ASTContext` allocation.
+  void *operator new(
+      size_t bytes, ASTContext &ctx,
+      unsigned alignment = alignof(DerivativeFunctionConfigurationList)) {
+    return ctx.Allocate(bytes, alignment);
+  }
+};
+
+void AbstractFunctionDecl::prepareDerivativeFunctionConfigurations() {
+  if (DerivativeFunctionConfigs)
+    return;
+  auto &ctx = getASTContext();
+  DerivativeFunctionConfigs = new (ctx) DerivativeFunctionConfigurationList();
+  // Register an `ASTContext` cleanup calling the list destructor.
+  ctx.addCleanup([this]() {
+    this->DerivativeFunctionConfigs->~DerivativeFunctionConfigurationList();
+  });
+}
+
+ArrayRef<AutoDiffConfig>
+AbstractFunctionDecl::getDerivativeFunctionConfigurations() {
+  prepareDerivativeFunctionConfigurations();
+  auto &ctx = getASTContext();
+  if (ctx.getCurrentGeneration() > DerivativeFunctionConfigGeneration) {
+    unsigned previousGeneration = DerivativeFunctionConfigGeneration;
+    DerivativeFunctionConfigGeneration = ctx.getCurrentGeneration();
+    ctx.loadDerivativeFunctionConfigurations(this, previousGeneration,
+                                             *DerivativeFunctionConfigs);
+  }
+  return DerivativeFunctionConfigs->getArrayRef();
+}
+
+void AbstractFunctionDecl::addDerivativeFunctionConfiguration(
+    AutoDiffConfig config) {
+  prepareDerivativeFunctionConfigurations();
+  DerivativeFunctionConfigs->insert(config);
+}
+// SWIFT_ENABLE_TENSORFLOW END
+
 FuncDecl *FuncDecl::createImpl(ASTContext &Context,
                                SourceLoc StaticLoc,
                                StaticSpellingKind StaticSpelling,