Revamp serialization to enable lookup by key.

Use mangling to support string key lookup.

Author: dan-zheng

include/swift/AST/ASTMangler.h

@@ -155,23 +155,31 @@ class ASTMangler : public Mangler {
                                              ModuleDecl *Module);
 
   // SWIFT_ENABLE_TENSORFLOW
-  // Mangle the derivative function (JVP/VJP) with the given:
-  // - Mangled original function name.
-  // - Derivative function kind.
-  // - Parameter/result indices.
+  /// Mangle the derivative function (JVP/VJP) with the given:
+  /// - Mangled original function name.
+  /// - Derivative function kind.
+  /// - Parameter/result indices.
   std::string mangleAutoDiffDerivativeFunctionHelper(
       StringRef name, AutoDiffDerivativeFunctionKind kind,
       const SILAutoDiffIndices &indices);
 
-  // SWIFT_ENABLE_TENSORFLOW
-  // Mangle the autodiff linear map (differential/pullback) with the given:
-  // - Mangled original function name.
-  // - Linear map kind.
-  // - Parameter/result indices.
+  /// Mangle the autodiff linear map (differential/pullback) with the given:
+  /// - Mangled original function name.
+  /// - Linear map kind.
+  /// - Parameter/result indices.
   std::string mangleAutoDiffLinearMapHelper(
       StringRef name, AutoDiffLinearMapKind kind,
       const SILAutoDiffIndices &indices);
 
+  /// Mangle a SIL differentiability witness key.
+  /// - Mangled original function name.
+  /// - Parameter indices.
+  /// - Result indices.
+  /// - Derivative generic signature (optional).
+  std::string mangleSILDifferentiabilityWitnessKey(
+      SILDifferentiabilityWitnessKey key);
+  // SWIFT_ENABLE_TENSORFLOW END
+
   std::string mangleKeyPathGetterThunkHelper(const AbstractStorageDecl *property,
                                              GenericSignature *signature,
                                              CanType baseType,

include/swift/AST/AutoDiff.h

@@ -481,6 +481,14 @@ 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: Unify with `AutoDiffDerivativeFunctionIdentifier`.
+using SILDifferentiabilityWitnessKey =
+std::tuple<StringRef, AutoDiffIndexSubset *,
+           AutoDiffIndexSubset *, GenericSignature *>;
+
 /// Automatic differentiation utility namespace.
 namespace autodiff {
 /// Appends the subset's parameter's types to `result`, in the order in

include/swift/SIL/SILDifferentiabilityWitness.h

@@ -75,16 +75,7 @@ class SILDifferentiabilityWitness
       serialized(isSerialized) {}
 
 public:
-  /// The key type, used for uniquing `SILDifferentiabilityWitness` in
-  /// `SILModule`, original function, parameter indices, result indices, and
-  /// derivative generic signature.
-  using Key = std::tuple<const SILFunction *, AutoDiffIndexSubset *,
-                         AutoDiffIndexSubset *, GenericSignature *>;
-  Key getKey() {
-    return std::make_tuple(originalFunction, parameterIndices, resultIndices,
-                           derivativeGenericSignature);
-  }
-
+  SILDifferentiabilityWitnessKey getKey() const;
   SILModule &getModule() const { return module; }
   SILLinkage getLinkage() const { return linkage; }
   SILFunction *getOriginalFunction() const { return originalFunction; }

include/swift/SIL/SILModule.h

@@ -207,9 +207,8 @@ class SILModule {
   /// Lookup table for SIL differentiability witnesses from original functions.
   /// Indexed by key type: original function, parameter indices, result indices,
   /// and derivative generic signature.
-  llvm::DenseMap<SILDifferentiabilityWitness::Key,
-                 SILDifferentiabilityWitness *>
-  DifferentiabilityWitnessMap;
+  llvm::DenseMap<SILDifferentiabilityWitnessKey, SILDifferentiabilityWitness *>
+      DifferentiabilityWitnessMap;
 
   /// The list of SILDifferentiabilityWitnesses in the module.
   DifferentiabilityWitnessListType differentiabilityWitnesses;

include/swift/Serialization/SerializedSILLoader.h

@@ -13,6 +13,9 @@
 #ifndef SWIFT_SERIALIZATION_SILLOADER_H
 #define SWIFT_SERIALIZATION_SILLOADER_H
 
+// SWIFT_ENABLE_TENSORFLOW
+#include "swift/AST/AutoDiff.h"
+// SWIFT_ENABLE_TENSORFLOW END
 #include "swift/AST/Decl.h"
 #include "swift/AST/Identifier.h"
 #include "swift/SIL/Notifications.h"
@@ -32,6 +35,9 @@ class SILModule;
 class SILVTable;
 class SILWitnessTable;
 class SILDefaultWitnessTable;
+// SWIFT_ENABLE_TENSORFLOW
+class SILDifferentiabilityWitness;
+// SWIFT_ENABLE_TENSORFLOW END
 
 /// Maintains a list of SILDeserializer, one for each serialized modules
 /// in ASTContext. It provides lookupSILFunction that will perform lookup
@@ -64,6 +70,10 @@ class SerializedSILLoader {
   SILVTable *lookupVTable(const ClassDecl *C);
   SILWitnessTable *lookupWitnessTable(SILWitnessTable *C);
   SILDefaultWitnessTable *lookupDefaultWitnessTable(SILDefaultWitnessTable *C);
+  // SWIFT_ENABLE_TENSORFLOW
+  SILDifferentiabilityWitness *
+  lookupDifferentiabilityWitness(SILDifferentiabilityWitnessKey key);
+  // SWIFT_ENABLE_TENSORFLOW END
 
   /// Invalidate the cached entries for deserialized SILFunctions.
   void invalidateCaches();

lib/AST/ASTMangler.cpp

@@ -379,11 +379,12 @@ std::string ASTMangler::mangleReabstractionThunkHelper(
   return finalize();
 }
 
+// SWIFT_ENABLE_TENSORFLOW
 std::string ASTMangler::mangleAutoDiffDerivativeFunctionHelper(
     StringRef name, AutoDiffDerivativeFunctionKind kind,
     const SILAutoDiffIndices &indices) {
   // TODO(TF-20): Make the mangling scheme robust.
-  // TODO(TF-680): Mangle `@differentiable` atttribute requirements as well.
+  // TODO(TF-680): Mangle derivative generic signature as well.
   beginManglingWithoutPrefix();
 
   Buffer << "AD__" << name << '_';
@@ -406,7 +407,7 @@ std::string ASTMangler::mangleAutoDiffLinearMapHelper(
     StringRef name, AutoDiffLinearMapKind kind,
     const SILAutoDiffIndices &indices) {
   // TODO(TF-20): Make the mangling scheme robust.
-  // TODO(TF-680): Mangle `@differentiable` atttribute requirements as well.
+  // TODO(TF-680): Mangle derivative generic signature as well.
   beginManglingWithoutPrefix();
 
   Buffer << "AD__" << name << '_';
@@ -425,6 +426,29 @@ std::string ASTMangler::mangleAutoDiffLinearMapHelper(
   return result;
 }
 
+std::string ASTMangler::mangleSILDifferentiabilityWitnessKey(
+    SILDifferentiabilityWitnessKey key) {
+  // TODO(TF-20): Make the mangling scheme robust.
+  // TODO(TF-680): Mangle derivative generic signature as well.
+  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);
+
+  Buffer << "AD__" << originalName << '_';
+  Buffer << "P" << parameterIndices->getString();
+  Buffer << "R" << resultIndices->getString();
+  if (derivativeGenericSignature)
+    appendGenericSignature(derivativeGenericSignature);
+
+  auto result = Storage.str().str();
+  Storage.clear();
+  return result;
+}
+// SWIFT_ENABLE_TENSORFLOW END
+
 std::string ASTMangler::mangleTypeForDebugger(Type Ty, const DeclContext *DC) {
   PrettyStackTraceType prettyStackTrace(Ty->getASTContext(),
                                         "mangling type for debugger", Ty);

lib/SIL/SILDifferentiabilityWitness.cpp

@@ -34,3 +34,8 @@ SILDifferentiabilityWitness *SILDifferentiabilityWitness::create(
   module.getDifferentiabilityWitnessList().push_back(diffWitness);
   return diffWitness;
 }
+
+SILDifferentiabilityWitnessKey SILDifferentiabilityWitness::getKey() const {
+  return std::make_tuple(originalFunction->getName(), parameterIndices,
+                         resultIndices, derivativeGenericSignature);
+}