[AutoDiff] Speed up compile time by 3.5x by caching derivative function types.

include/swift/AST/ASTContext.h

@@ -117,6 +117,7 @@ namespace swift {
   // SWIFT_ENABLE_TENSORFLOW
   struct AutoDiffConfig;
   struct AutoDiffDerivativeFunctionKind;
+  struct SILAutoDiffDerivativeFunctionKey;
   class DerivativeAttr;
   class DifferentiableAttr;
   // SWIFT_ENABLE_TENSORFLOW END
@@ -293,6 +294,10 @@ class ASTContext final {
 
   /// Cached mapping from types to their associated tangent spaces.
   llvm::DenseMap<Type, Optional<TangentSpace>> AutoDiffTangentSpaces;
+
+  /// A cache of derivative function types per configuration.
+  llvm::DenseMap<SILAutoDiffDerivativeFunctionKey, CanSILFunctionType>
+      SILAutoDiffDerivativeFunctions;
 
   /// Cache of `@differentiable` attributes keyed by parameter indices. Used to
   /// diagnose duplicate `@differentiable` attributes for the same key.

include/swift/AST/AutoDiff.h

@@ -465,6 +465,15 @@ inline llvm::raw_ostream &operator<<(llvm::raw_ostream &s,
   return s;
 }
 
+struct SILAutoDiffDerivativeFunctionKey {
+  SILFunctionType *originalType;
+  IndexSubset *parameterIndices;
+  IndexSubset *resultIndices;
+  AutoDiffDerivativeFunctionKind kind;
+  CanGenericSignature derivativeFnGenSig;
+  bool isReabstractionThunk;
+};
+
 /// In conjunction with the original function declaration, identifies an
 /// autodiff derivative function.
 ///
@@ -599,7 +608,60 @@ template <> struct DenseMapInfo<SILAutoDiffIndices> {
   }
 };
 
-} // end namespace llvm
+using swift::SILAutoDiffDerivativeFunctionKey;
+using swift::SILFunctionType;
+using swift::IndexSubset;
+using swift::AutoDiffDerivativeFunctionKind;
+using swift::GenericSignature;
+using swift::CanGenericSignature;
+
+template <> struct DenseMapInfo<SILAutoDiffDerivativeFunctionKey> {
+
+  static bool isEqual(const SILAutoDiffDerivativeFunctionKey lhs,
+                      const SILAutoDiffDerivativeFunctionKey rhs) {
+    return lhs.originalType == rhs.originalType &&
+        lhs.parameterIndices == rhs.parameterIndices &&
+        lhs.resultIndices == rhs.resultIndices &&
+        lhs.kind == rhs.kind &&
+        lhs.derivativeFnGenSig == rhs.derivativeFnGenSig &&
+        lhs.isReabstractionThunk == rhs.isReabstractionThunk;
+  }
+
+  static inline SILAutoDiffDerivativeFunctionKey getEmptyKey() {
+    return {
+        DenseMapInfo<SILFunctionType *>::getEmptyKey(),
+        DenseMapInfo<IndexSubset *>::getEmptyKey(),
+        DenseMapInfo<IndexSubset *>::getEmptyKey(),
+        AutoDiffDerivativeFunctionKind::innerty(
+            DenseMapInfo<unsigned>::getEmptyKey()),
+        CanGenericSignature(DenseMapInfo<GenericSignature>::getEmptyKey()),
+        (bool)DenseMapInfo<unsigned>::getEmptyKey()};
+  }
+
+  static inline SILAutoDiffDerivativeFunctionKey getTombstoneKey() {
+    return {
+        DenseMapInfo<SILFunctionType *>::getTombstoneKey(),
+        DenseMapInfo<IndexSubset *>::getTombstoneKey(),
+        DenseMapInfo<IndexSubset *>::getTombstoneKey(),
+        AutoDiffDerivativeFunctionKind::innerty(
+            DenseMapInfo<unsigned>::getTombstoneKey()),
+        CanGenericSignature(DenseMapInfo<GenericSignature>::getTombstoneKey()),
+        (bool)DenseMapInfo<unsigned>::getTombstoneKey()};
+  }
+
+  static unsigned getHashValue(const SILAutoDiffDerivativeFunctionKey &Val) {
+    return hash_combine(
+        DenseMapInfo<SILFunctionType *>::getHashValue(Val.originalType),
+        DenseMapInfo<IndexSubset *>::getHashValue(Val.parameterIndices),
+        DenseMapInfo<IndexSubset *>::getHashValue(Val.resultIndices),
+        DenseMapInfo<unsigned>::getHashValue((unsigned)Val.kind.rawValue),
+        DenseMapInfo<GenericSignature>::getHashValue(Val.derivativeFnGenSig),
+        DenseMapInfo<unsigned>::getHashValue(
+            (unsigned)Val.isReabstractionThunk));
+  }
+};
+
+} // namespace llvm
 // SWIFT_ENABLE_TENSORFLOW END
 
 #endif // SWIFT_AST_AUTODIFF_H

lib/AST/ASTContext.cpp

@@ -435,10 +435,6 @@ struct ASTContext::Implementation {
   llvm::FoldingSet<DeclName::CompoundDeclName> CompoundNames;
   llvm::DenseMap<UUID, OpenedArchetypeType *> OpenedExistentialArchetypes;
 
-  // SWIFT_ENABLE_TENSORFLOW
-  /// A cache of tangent spaces per type.
-  llvm::DenseMap<CanType, Optional<TangentSpace>> TangentSpaces;
-
   /// For uniquifying `IndexSubset` allocations.
   llvm::FoldingSet<IndexSubset> IndexSubsets;
 

lib/SIL/SILFunctionType.cpp

@@ -260,12 +260,23 @@ CanSILFunctionType SILFunctionType::getAutoDiffDerivativeFunctionType(
     AutoDiffDerivativeFunctionKind kind, TypeConverter &TC,
     LookupConformanceFn lookupConformance,
     CanGenericSignature derivativeFnGenSig, bool isReabstractionThunk) {
+  auto &ctx = getASTContext();
+  SILAutoDiffDerivativeFunctionKey key{
+    this, parameterIndices,
+    IndexSubset::get(ctx, getNumResults(), {resultIndex}),
+    kind, derivativeFnGenSig, isReabstractionThunk
+  };
+  auto insertion =
+      ctx.SILAutoDiffDerivativeFunctions.try_emplace(key, CanSILFunctionType());
+  auto &cachedResult = insertion.first->getSecond();
+  if (!insertion.second)
+    return cachedResult;
+
   // JVP: (T...) -> ((R...),
   //                 (T.TangentVector...) -> (R.TangentVector...))
   // VJP: (T...) -> ((R...),
   //                 (R.TangentVector...) -> (T.TangentVector...))
 
-  auto &ctx = getASTContext();
 
   // Helper function testing if we are differentiating wrt this index.
   auto isWrtIndex = [&](unsigned index) -> bool {
@@ -431,11 +442,12 @@ CanSILFunctionType SILFunctionType::getAutoDiffDerivativeFunctionType(
   auto extInfo = getExtInfo();
   if (getRepresentation() == SILFunctionTypeRepresentation::CFunctionPointer)
     extInfo = extInfo.withRepresentation(SILFunctionTypeRepresentation::Thin);
-  return SILFunctionType::get(canGenSig, extInfo, getCoroutineKind(),
-                              getCalleeConvention(), newParameters, getYields(),
-                              newResults, getOptionalErrorResult(),
-                              getSubstitutions(), isGenericSignatureImplied(),
-                              ctx, getWitnessMethodConformanceOrInvalid());
+  cachedResult = SILFunctionType::get(
+      canGenSig, extInfo, getCoroutineKind(), getCalleeConvention(),
+      newParameters, getYields(), newResults, getOptionalErrorResult(),
+      getSubstitutions(), isGenericSignatureImplied(), ctx,
+      getWitnessMethodConformanceOrInvalid());
+  return cachedResult;
 }
 
 CanSILFunctionType SILFunctionType::getAutoDiffTransposeFunctionType(