[AutoDiff] DifferentiabilityWitnessFunctionInst stores pointer to witness

include/swift/AST/DiagnosticsParse.def

@@ -691,6 +691,8 @@ ERROR(sil_diff_witness_expected_token,PointsToFirstBadToken,
       "expected '%0' in differentiability witness", (StringRef))
 ERROR(sil_diff_witness_serialized_declaration,none,
       "differentiability witness declaration should not be serialized", ())
+ERROR(sil_diff_witness_undefined,PointsToFirstBadToken,
+      "reference to undefined differentiability witness", ())
 
 // SIL Coverage Map
 ERROR(sil_coverage_func_not_found, none,

include/swift/SIL/SILBuilder.h

@@ -551,13 +551,10 @@ class SILBuilder {
 
   DifferentiabilityWitnessFunctionInst *
   createDifferentiabilityWitnessFunction(
-      SILLocation Loc, SILFunction *OriginalFunction,
-      DifferentiabilityWitnessFunctionKind WitnessKind,
-      IndexSubset *ParameterIndices, IndexSubset *ResultIndices,
-      GenericSignature WitnessGenericSignature) {
+      SILLocation Loc, DifferentiabilityWitnessFunctionKind WitnessKind,
+      SILDifferentiabilityWitness *Witness) {
     return insert(new (getModule()) DifferentiabilityWitnessFunctionInst(
-        getModule(), getSILDebugLocation(Loc), OriginalFunction, WitnessKind,
-        ParameterIndices, ResultIndices, WitnessGenericSignature.getPointer()));
+        getModule(), getSILDebugLocation(Loc), WitnessKind, Witness));
   }
   // SWIFT_ENABLE_TENSORFLOW END
 

include/swift/SIL/SILCloner.h

@@ -1018,9 +1018,8 @@ void SILCloner<ImplClass>::visitDifferentiabilityWitnessFunctionInst(
   getBuilder().setCurrentDebugScope(getOpScope(Inst->getDebugScope()));
   recordClonedInstruction(
       Inst, getBuilder().createDifferentiabilityWitnessFunction(
-                getOpLocation(Inst->getLoc()), Inst->getOriginalFunction(),
-                Inst->getWitnessKind(), Inst->getParameterIndices(),
-                Inst->getResultIndices(), Inst->getWitnessGenericSignature()));
+                getOpLocation(Inst->getLoc()), Inst->getWitnessKind(),
+                Inst->getWitness()));
 }
 // SWIFT_ENABLE_TENSORFLOW END
 

include/swift/SIL/SILInstruction.h

@@ -61,6 +61,9 @@ class SILBasicBlock;
 class SILBuilder;
 class SILDebugLocation;
 class SILDebugScope;
+// SWIFT_ENABLE_TENSORFLOW
+class SILDifferentiabilityWitness;
+// SWIFT_ENABLE_TENSORFLOW_END
 class SILFunction;
 class SILGlobalVariable;
 class SILInstructionResultArray;
@@ -8038,43 +8041,31 @@ class DifferentiabilityWitnessFunctionInst
           SingleValueInstruction> {
 private:
   friend SILBuilder;
-  /// The original function.
-  SILFunction *originalFunction;
   /// The differentiability witness function kind.
   DifferentiabilityWitnessFunctionKind witnessKind;
-  /// The autodiff config: parameter indices, result indices, and witness
-  /// derivative signature.
-  AutoDiffConfig config;
+  /// The referenced SIL differentiability witness.
+  SILDifferentiabilityWitness *witness;
 
   static SILType getDifferentiabilityWitnessType(
-      SILModule &module, SILFunction *originalFunction,
+      SILModule &module,
       DifferentiabilityWitnessFunctionKind witnessKind,
-      IndexSubset *parameterIndices, IndexSubset *resultIndices,
-      GenericSignature witnessGenericSignature);
+      SILDifferentiabilityWitness *witness);
 
 public:
   DifferentiabilityWitnessFunctionInst(
-      SILModule &module, SILDebugLocation loc, SILFunction *originalFunction,
+      SILModule &module, SILDebugLocation loc,
       DifferentiabilityWitnessFunctionKind witnessKind,
-      IndexSubset *parameterIndices, IndexSubset *resultIndices,
-      GenericSignature witnessGenericSignature);
+      SILDifferentiabilityWitness *witness);
 
   static DifferentiabilityWitnessFunctionInst *create(
-      SILModule &module, SILDebugLocation loc, SILFunction *originalFunction,
+      SILModule &module, SILDebugLocation loc,
       DifferentiabilityWitnessFunctionKind witnessKind,
-      IndexSubset *parameterIndices, IndexSubset *resultIndices,
-      GenericSignature witnessGenericSignature);
+      SILDifferentiabilityWitness *witness);
 
   DifferentiabilityWitnessFunctionKind getWitnessKind() const {
     return witnessKind;
   }
-  SILFunction *getOriginalFunction() const { return originalFunction; }
-  AutoDiffConfig const &getConfig() const { return config; }
-  IndexSubset *getParameterIndices() const { return config.parameterIndices; }
-  IndexSubset *getResultIndices() const { return config.resultIndices; }
-  GenericSignature getWitnessGenericSignature() const {
-    return config.derivativeGenericSignature;
-  }
+  SILDifferentiabilityWitness *getWitness() const { return witness; }
 
   ArrayRef<Operand> getAllOperands() const { return {}; }
   MutableArrayRef<Operand> getAllOperands() { return {}; }

include/swift/SIL/SILModule.h

@@ -204,11 +204,8 @@ class SILModule {
   DefaultWitnessTableListType defaultWitnessTables;
 
   // SWIFT_ENABLE_TENSORFLOW
-  /// 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<SILDifferentiabilityWitnessKey, SILDifferentiabilityWitness *>
-      DifferentiabilityWitnessMap;
+  /// Lookup table for SIL differentiability witnesses, keyed by mangled name.
+  llvm::StringMap<SILDifferentiabilityWitness *> DifferentiabilityWitnessMap;
 
   /// The list of SILDifferentiabilityWitnesses in the module.
   DifferentiabilityWitnessListType differentiabilityWitnesses;
@@ -609,6 +606,15 @@ class SILModule {
   /// hierarchy of \p Class.
   SILFunction *lookUpFunctionInVTable(ClassDecl *Class, SILDeclRef Member);
 
+  // SWIFT_ENABLE_TENSORFLOW
+  /// Look up the differentiability witness with the given name.
+  SILDifferentiabilityWitness *lookUpDifferentiabilityWitness(StringRef name);
+
+  /// Look up the differentiability witness corresponding to the given key.
+  SILDifferentiabilityWitness *
+  lookUpDifferentiabilityWitness(SILDifferentiabilityWitnessKey key);
+  // SWIFT_ENABLE_TENSORFLOW_END
+
   // Given a protocol, attempt to create a default witness table declaration
   // for it.
   SILDefaultWitnessTable *

lib/Parse/ParseStmt.cpp

@@ -326,6 +326,9 @@ ParserStatus Parser::parseBraceItems(SmallVectorImpl<ASTNode> &Entries,
          Tok.isNot(tok::kw_sil_global) &&
          Tok.isNot(tok::kw_sil_witness_table) &&
          Tok.isNot(tok::kw_sil_default_witness_table) &&
+         // SWIFT_ENABLE_TENSORFLOW
+         Tok.isNot(tok::kw_sil_differentiability_witness) &&
+         // SWIFT_ENABLE_TENSORFLOW_END
          Tok.isNot(tok::kw_sil_property) &&
          (isConditionalBlock ||
           !isTerminatorForBraceItemListKind(Kind, Entries))) {