[AutoDiff upstream] Add SIL derivative function type caching. (#29953)

Upstream #29590:
cache `SILFunctionType::getAutoDiffDerivativeFunctionType` results.

Author: dan-zheng

include/swift/AST/ASTContext.h

@@ -114,6 +114,7 @@ namespace swift {
   class VarDecl;
   class UnifiedStatsReporter;
   class IndexSubset;
+  struct SILAutoDiffDerivativeFunctionKey;
 
   enum class KnownProtocolKind : uint8_t;
 
@@ -288,6 +289,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.
   llvm::DenseMap<std::pair<Decl *, IndexSubset *>, DifferentiableAttr *>

include/swift/AST/AutoDiff.h

@@ -113,6 +113,16 @@ struct AutoDiffConfig {
   SWIFT_DEBUG_DUMP;
 };
 
+/// Key for caching SIL derivative function types.
+struct SILAutoDiffDerivativeFunctionKey {
+  SILFunctionType *originalType;
+  IndexSubset *parameterIndices;
+  IndexSubset *resultIndices;
+  AutoDiffDerivativeFunctionKind kind;
+  CanGenericSignature derivativeFnGenSig;
+  bool isReabstractionThunk;
+};
+
 class ParsedAutoDiffParameter {
 public:
   enum class Kind { Named, Ordered, Self };
@@ -281,8 +291,11 @@ namespace llvm {
 
 using swift::AutoDiffConfig;
 using swift::AutoDiffDerivativeFunctionKind;
+using swift::CanGenericSignature;
 using swift::GenericSignature;
 using swift::IndexSubset;
+using swift::SILAutoDiffDerivativeFunctionKey;
+using swift::SILFunctionType;
 
 template <typename T> struct DenseMapInfo;
 
@@ -354,6 +367,50 @@ template <> struct DenseMapInfo<AutoDiffDerivativeFunctionKind> {
   }
 };
 
+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.rawValue == rhs.kind.rawValue &&
+           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));
+  }
+};
+
 } // end namespace llvm
 
 #endif // SWIFT_AST_AUTODIFF_H

lib/SIL/SILFunctionType.cpp

@@ -247,6 +247,15 @@ CanSILFunctionType SILFunctionType::getAutoDiffDerivativeFunctionType(
     LookupConformanceFn lookupConformance,
     CanGenericSignature derivativeFnGenSig, bool isReabstractionThunk) {
   auto &ctx = getASTContext();
+  auto resultIndices = IndexSubset::get(ctx, getNumResults(), {resultIndex});
+  SILAutoDiffDerivativeFunctionKey key{
+      this, parameterIndices,   resultIndices,
+      kind, derivativeFnGenSig, isReabstractionThunk};
+  auto insertion =
+      ctx.SILAutoDiffDerivativeFunctions.try_emplace(key, CanSILFunctionType());
+  auto &cachedResult = insertion.first->getSecond();
+  if (!insertion.second)
+    return cachedResult;
 
   // Returns true if `index` is a differentiability parameter index.
   auto isDiffParamIndex = [&](unsigned index) -> bool {
@@ -396,11 +405,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(