[AutoDiff] DifferentiabilityWitnessFunctionInst stores pointer to witness (#27919)

We want the `DifferentiabilityWitnessFunctionInst` to store a pointer to the witness because this makes the instruction data structure smaller and simpler, and makes misuse (e.g. referencing a witness that's not declared in the SIL module) harder. 

This PR mostly preserves the existing SIL syntax & rules. The changes are:
* `sil_differentiability_witness` declarations are now required to come before `differentiability_witness_function` instructions that reference them, so that we don't have to add "forward declaration" logic for them.
* `sil_differentiability_witness` constraints now look like `<τ_0_0 where τ_0_0 : _Differentiable>` instead of `[where T : _Differentiable]`. This allows me to share parsing logic between the `sil_differentiability_witness` constraints and the `differentiability_witness_function`, which simplifies the code and also avoids a problem [1].

This PR changes the `DifferentiabilityWitnessFunctionInst` serialization format to store the witness mangled name rather than the original function name plus indices and generic constraints.

Detailed outline of changes:
* Datastructure & utility changes:
  * `DifferentiabilityWitnessFunctionInst` now contains a `SILDifferentiabilityWitness*` instead of a pointer to original function, indices, and generic signature. (credit to @dan-zheng, I just copied these changes from his commit)
  * Added functions in `SILModule` for looking up `SILDifferentiabilityWitness*` because parsing and deserialization need to find the right `SILDifferentiabilityWitness*` to put in the inst.
* SIL printing changes:
  * Print `sil_differentiability_witness` declarations before function declarations
  * Print `<T where T ...>` instead of `[where T ...]` in `sil_differentiability_witness` declarations.
* SIL parsing changes:
  * Factor out a common `parseSILDifferentiabilityWitnessConfigAndFunction`, and use this for parsing both `sil_differentiability_witness` and `differentiability_witness_function`. This simplifies the code and avoids the problem [1].
  * When parsing a `differentiability_witness_function`, look for the `SILDifferentiabilityWitness*` in the module.
  * Bugfix in ParseStmt so that it can parse `sil_differentiability_witness` declarations that come at the beginning of the file.
* Serialization changes:
  * Deleted `SILInstDifferentiabilityWitnessFunctionLayout` from the serialization format, because the `DifferentiabilityWitnessFunctionInst` now fits into the `SILOneOperandLayout`.
  * Changed serialization to serialize `DifferentiabilityWitnessFunctionInst` a `SILOneOperandLayout` with the mangled witness name.
  * Changed deserialization to deserialize `DifferentiabilityWitnessFunctionInst` by reading the name and deserializing the corresponding witness.
  * Changed deserialization to treat declaration-only witnesses as "fully deserialized". Otherwise it tries to deserialize them twice and the module doesn't like having duplicate witnesses.

[1] The avoided problem: The previous code for parsing `sil_differentiability_witness` constraints of the form `[where T : _Differentiable]` fails to handle the case where the original function is declaration-only, because the code assumes that there is a generic environment, and declaration-only functions do not have a generic environment. Rather than fixing this code, it seemed better to delete it and use the existing code that parses `differentiability_witness_function` constraints, which does not suffer from this problem.

Author: marcrasi

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))) {