Add AutoDiffConfig and use in SILDifferentiabilityWitnessKey.

Author: dan-zheng

include/swift/AST/AutoDiff.h

@@ -208,6 +208,41 @@ struct AutoDiffDerivativeFunctionKind {
   }
 };
 
+/// Identifies an autodiff derivative function configuration:
+/// - Parameter indices.
+/// - Result indices.
+/// - Derivative generic signature (optional).
+// TODO(TF-893): Use `AutoDiffConfig` in `AutoDiffDerivativeFunctionIdentifier`
+// to avoid duplication.
+class AutoDiffConfig : public llvm::FoldingSetNode {
+  IndexSubset *const parameterIndices;
+  IndexSubset *const resultIndices;
+  GenericSignature *derivativeGenericSignature;
+
+  AutoDiffConfig(IndexSubset *parameterIndices, IndexSubset *resultIndices,
+                 GenericSignature *derivativeGenericSignature)
+      : parameterIndices(parameterIndices), resultIndices(resultIndices),
+        derivativeGenericSignature(derivativeGenericSignature) {}
+
+public:
+  IndexSubset *getParameterIndices() const { return parameterIndices; }
+  IndexSubset *getResultIndices() const { return resultIndices; }
+  GenericSignature *getDerivativeGenericSignature() const {
+    return derivativeGenericSignature;
+  }
+
+  static AutoDiffConfig *get(IndexSubset *parameterIndices,
+                             IndexSubset *resultIndices,
+                             GenericSignature *derivativeGenericSignature,
+                             ASTContext &C);
+
+  void Profile(llvm::FoldingSetNodeID &ID) {
+    ID.AddPointer(parameterIndices);
+    ID.AddPointer(resultIndices);
+    ID.AddPointer(derivativeGenericSignature);
+  }
+};
+
 /// In conjunction with the original function declaration, identifies an
 /// autodiff derivative function.
 ///
@@ -241,9 +276,7 @@ class AutoDiffDerivativeFunctionIdentifier : public llvm::FoldingSetNode {
 /// The key type used for uniquing `SILDifferentiabilityWitness` in
 /// `SILModule`: original function name, parameter indices, result indices, and
 /// derivative generic signature.
-// TODO(TF-893): Unify with `AutoDiffDerivativeFunctionIdentifier`.
-using SILDifferentiabilityWitnessKey =
-    std::tuple<StringRef, IndexSubset *, IndexSubset *, GenericSignature *>;
+using SILDifferentiabilityWitnessKey = std::pair<StringRef, AutoDiffConfig *>;
 
 /// Automatic differentiation utility namespace.
 namespace autodiff {

include/swift/SIL/SILDifferentiabilityWitness.h

@@ -29,7 +29,6 @@
 #include "swift/AST/AutoDiff.h"
 #include "swift/AST/GenericSignature.h"
 #include "swift/SIL/SILAllocated.h"
-#include "swift/SIL/SILInstruction.h"
 #include "llvm/ADT/ilist_node.h"
 #include "llvm/ADT/ilist.h"
 
@@ -48,12 +47,9 @@ class SILDifferentiabilityWitness
   SILLinkage linkage;
   /// The original function.
   SILFunction *originalFunction;
-  /// The parameter indices.
-  IndexSubset *parameterIndices;
-  /// The result indices.
-  IndexSubset *resultIndices;
-  /// The derivative generic signature (optional).
-  GenericSignature *derivativeGenericSignature;
+  /// The autodiff configuration: parameter indices, result indices, and
+  /// derivative generic signature (optional).
+  AutoDiffConfig *autoDiffConfig;
   /// The JVP (Jacobian-vector products) derivative function.
   SILFunction *jvp;
   /// The VJP (vector-Jacobian products) derivative function.
@@ -62,6 +58,11 @@ class SILDifferentiabilityWitness
   /// devirtualization from another module.
   bool serialized;
 
+  static AutoDiffConfig *
+  getAutoDiffConfig(SILModule &module, IndexSubset *parameterIndices,
+                    IndexSubset *resultIndices,
+                    GenericSignature *derivativeGenSig);
+
   SILDifferentiabilityWitness(SILModule &module, SILLinkage linkage,
                               SILFunction *originalFunction,
                               IndexSubset *parameterIndices,
@@ -70,9 +71,9 @@ class SILDifferentiabilityWitness
                               SILFunction *jvp, SILFunction *vjp,
                               bool isSerialized)
     : module(module), linkage(linkage), originalFunction(originalFunction),
-      parameterIndices(parameterIndices), resultIndices(resultIndices),
-      derivativeGenericSignature(derivativeGenSig), jvp(jvp), vjp(vjp),
-      serialized(isSerialized) {}
+      autoDiffConfig(getAutoDiffConfig(
+          module, parameterIndices, resultIndices, derivativeGenSig)),
+      jvp(jvp), vjp(vjp), serialized(isSerialized) {}
 
 public:
   static SILDifferentiabilityWitness *create(
@@ -86,16 +87,29 @@ class SILDifferentiabilityWitness
   SILLinkage getLinkage() const { return linkage; }
   SILFunction *getOriginalFunction() const { return originalFunction; }
   IndexSubset *getParameterIndices() const {
-    return parameterIndices;
+    return autoDiffConfig->getParameterIndices();
   }
   IndexSubset *getResultIndices() const {
-    return resultIndices;
+    return autoDiffConfig->getResultIndices();
   }
   GenericSignature *getDerivativeGenericSignature() const {
-    return derivativeGenericSignature;
+    return autoDiffConfig->getDerivativeGenericSignature();
   }
   SILFunction *getJVP() const { return jvp; }
   SILFunction *getVJP() const { return vjp; }
+  SILFunction *getDerivative(AutoDiffDerivativeFunctionKind kind) const {
+    switch (kind) {
+    case AutoDiffDerivativeFunctionKind::JVP: return jvp;
+    case AutoDiffDerivativeFunctionKind::VJP: return vjp;
+    }
+  }
+  void setDerivative(AutoDiffDerivativeFunctionKind kind,
+                     SILFunction *derivative) {
+    switch (kind) {
+    case AutoDiffDerivativeFunctionKind::JVP: jvp = derivative; break;
+    case AutoDiffDerivativeFunctionKind::VJP: vjp = derivative; break;
+    }
+  }
   bool isSerialized() const { return serialized; }
 
   /// Verify that the differentiability witness is well-formed.

lib/AST/ASTContext.cpp

@@ -449,9 +449,13 @@ FOR_KNOWN_FOUNDATION_TYPES(CACHE_FOUNDATION_DECL)
   /// For uniquifying `IndexSubset` allocations.
   llvm::FoldingSet<IndexSubset> IndexSubsets;
 
+  /// For uniquifying `AutoDiffConfig` allocations.
+  llvm::FoldingSet<AutoDiffConfig> AutoDiffConfigs;
+
   /// For uniquifying `AutoDiffDerivativeFunctionIdentifier` allocations.
   llvm::FoldingSet<AutoDiffDerivativeFunctionIdentifier>
       AutoDiffDerivativeFunctionIdentifiers;
+  // SWIFT_ENABLE_TENSORFLOW END
 
   /// A cache of information about whether particular nominal types
   /// are representable in a foreign language.
@@ -4828,6 +4832,27 @@ IndexSubset::get(ASTContext &ctx, const SmallBitVector &indices) {
   return newNode;
 }
 
+AutoDiffConfig *AutoDiffConfig::get(
+    IndexSubset *parameterIndices, IndexSubset *resultIndices,
+    GenericSignature *derivativeGenericSignature, ASTContext &C) {
+  assert(parameterIndices);
+  assert(resultIndices);
+  auto &foldingSet = C.getImpl().AutoDiffConfigs;
+  llvm::FoldingSetNodeID id;
+  id.AddPointer(parameterIndices);
+  id.AddPointer(resultIndices);
+  id.AddPointer(derivativeGenericSignature);
+  void *insertPos;
+  auto *existing = foldingSet.FindNodeOrInsertPos(id, insertPos);
+  if (existing)
+    return existing;
+  void *buf = C.Allocate(sizeof(AutoDiffConfig), alignof(AutoDiffConfig));
+  auto *newNode = new (buf) AutoDiffConfig(
+      parameterIndices, resultIndices, derivativeGenericSignature);
+  foldingSet.InsertNode(newNode, insertPos);
+  return newNode;
+}
+
 AutoDiffDerivativeFunctionIdentifier *
 AutoDiffDerivativeFunctionIdentifier::get(
     AutoDiffDerivativeFunctionKind kind, IndexSubset *parameterIndices,

lib/AST/ASTMangler.cpp

@@ -431,10 +431,11 @@ std::string ASTMangler::mangleSILDifferentiabilityWitnessKey(
   // TODO(TF-20): Make the mangling scheme robust.
   beginManglingWithoutPrefix();
 
-  auto originalName = std::get<0>(key);
-  auto *parameterIndices = std::get<1>(key);
-  auto *resultIndices = std::get<2>(key);
-  auto *derivativeGenericSignature = std::get<3>(key);
+  auto originalName = key.first;
+  auto *autoDiffConfig = key.second;
+  auto *parameterIndices = autoDiffConfig->getParameterIndices();
+  auto *resultIndices = autoDiffConfig->getResultIndices();
+  auto *derivativeGenericSignature = autoDiffConfig->getDerivativeGenericSignature();
 
   Buffer << "AD__" << originalName << '_';
   Buffer << "P" << parameterIndices->getString();

lib/SIL/SILDifferentiabilityWitness.cpp

@@ -35,7 +35,13 @@ SILDifferentiabilityWitness *SILDifferentiabilityWitness::create(
   return diffWitness;
 }
 
+AutoDiffConfig *SILDifferentiabilityWitness::getAutoDiffConfig(
+    SILModule &module, IndexSubset *parameterIndices,
+    IndexSubset *resultIndices, GenericSignature *derivativeGenSig) {
+  return AutoDiffConfig::get(parameterIndices, resultIndices, derivativeGenSig,
+                             module.getASTContext());
+}
+
 SILDifferentiabilityWitnessKey SILDifferentiabilityWitness::getKey() const {
-  return std::make_tuple(originalFunction->getName(), parameterIndices,
-                         resultIndices, derivativeGenericSignature);
+  return std::make_pair(originalFunction->getName(), autoDiffConfig);
 }

lib/SIL/SILPrinter.cpp

@@ -3068,31 +3068,27 @@ void SILDifferentiabilityWitness::print(
   printLinkage(OS, linkage, ForDefinition);
   // [parameters 0 1 ...]
   OS << "[parameters ";
-  interleave(parameterIndices->getIndices(),
+  interleave(getParameterIndices()->getIndices(),
              [&](unsigned index) { OS << index; },
              [&] { OS << " "; });
   // [results 0 1 ...]
   OS << "] [results ";
-  interleave(resultIndices->getIndices(),
+  interleave(getResultIndices()->getIndices(),
              [&](unsigned index) { OS << index; },
              [&] { OS << " "; });
   OS << ']';
   // [where ...]
-  if (derivativeGenericSignature) {
-    // NOTE: This needs to be changed if there is no utility for parsing
-    // generic signatures. Idea: we could instead print the type of the original
-    // function substituted into this generic signature.
+  if (auto *derivativeGenSig = getDerivativeGenericSignature()) {
     ArrayRef<Requirement> requirements;
     SmallVector<Requirement, 4> requirementsScratch;
     auto *origGenEnv = originalFunction->getGenericEnvironment();
-    if (derivativeGenericSignature) {
+    if (derivativeGenSig) {
       if (origGenEnv) {
-        requirementsScratch =
-            derivativeGenericSignature->requirementsNotSatisfiedBy(
-                origGenEnv->getGenericSignature());
+        requirementsScratch = derivativeGenSig->requirementsNotSatisfiedBy(
+            origGenEnv->getGenericSignature());
         requirements = requirementsScratch;
       } else {
-        requirements = derivativeGenericSignature->getRequirements();
+        requirements = derivativeGenSig->getRequirements();
       }
     }
     if (!requirements.empty()) {

lib/SIL/SILVerifier.cpp

@@ -5365,7 +5365,7 @@ void SILDifferentiabilityWitness::verify(const SILModule &M) const {
     // TODO(TF-893): Change `SILFunctionType::getAutoDiffDerivativeFunctionType`
     // to accept result indices.
     auto expectedJVPType = origFnType->getAutoDiffDerivativeFunctionType(
-        getParameterIndices(), /*resultIndex*/ *resultIndices->begin(),
+        getParameterIndices(), /*resultIndex*/ *getResultIndices()->begin(),
         AutoDiffDerivativeFunctionKind::JVP, M.Types,
         LookUpConformanceInModule(M.getSwiftModule()), derivativeCanGenSig);
     SILVerifier(*jvp).requireSameType(
@@ -5377,7 +5377,7 @@ void SILDifferentiabilityWitness::verify(const SILModule &M) const {
     // TODO(TF-893): Change `SILFunctionType::getAutoDiffDerivativeFunctionType`
     // to result indices.
     auto expectedVJPType = origFnType->getAutoDiffDerivativeFunctionType(
-        getParameterIndices(), /*resultIndex*/ *resultIndices->begin(),
+        getParameterIndices(), /*resultIndex*/ *getResultIndices()->begin(),
         AutoDiffDerivativeFunctionKind::VJP, M.Types,
         LookUpConformanceInModule(M.getSwiftModule()), derivativeCanGenSig);
     SILVerifier(*vjp).requireSameType(