Merge pull request #36015 from rxwei/74380324-mangle-diff-witness-keys

Author: swift-ci

docs/ABI/Mangling.rst

@@ -167,6 +167,8 @@ Globals
   global ::= protocol-conformance protocol 'Wb' // base protocol witness table accessor
   global ::= type protocol-conformance 'Wl' // lazy protocol witness table accessor
 
+  global ::= global generic-signature? 'WJ' DIFFERENTIABILITY-KIND INDEX-SUBSET 'p' INDEX-SUBSET 'r' // differentiability witness
+
   global ::= type 'WV'                   // value witness table
   global ::= entity 'Wvd'                // field offset
   global ::= entity 'WC'                 // resilient enum tag index

docs/SIL.rst

@@ -1590,12 +1590,15 @@ Differentiability Witnesses
   sil-differentiability-witness ::=
       'sil_differentiability_witness'
       sil-linkage?
+      '[' differentiability-kind ']'
       '[' 'parameters' sil-differentiability-witness-function-index-list ']'
       '[' 'results' sil-differentiability-witness-function-index-list ']'
       generic-parameter-clause?
       sil-function-name ':' sil-type
       sil-differentiability-witness-body?
 
+  differentiability-kind ::= 'forward' | 'reverse' | 'normal' | 'linear'
+
   sil-differentiability-witness-body ::=
       '{' sil-differentiability-witness-entry?
           sil-differentiability-witness-entry? '}'
@@ -1625,7 +1628,7 @@ based on the key.
 
 ::
 
-  sil_differentiability_witness hidden [parameters 0] [results 0] <T where T : Differentiable> @id : $@convention(thin) (T) -> T {
+  sil_differentiability_witness hidden [normal] [parameters 0] [results 0] <T where T : Differentiable> @id : $@convention(thin) (T) -> T {
     jvp: @id_jvp : $@convention(thin) (T) -> (T, @owned @callee_guaranteed (T.TangentVector) -> T.TangentVector)
     vjp: @id_vjp : $@convention(thin) (T) -> (T, @owned @callee_guaranteed (T.TangentVector) -> T.TangentVector)
   }
@@ -7066,6 +7069,7 @@ differentiability_witness_function
   sil-instruction ::=
       'differentiability_witness_function'
       '[' sil-differentiability-witness-function-kind ']'
+      '[' differentiability-kind ']'
       '[' 'parameters' sil-differentiability-witness-function-index-list ']'
       '[' 'results' sil-differentiability-witness-function-index-list ']'
       generic-parameter-clause?
@@ -7074,7 +7078,7 @@ differentiability_witness_function
   sil-differentiability-witness-function-kind ::= 'jvp' | 'vjp' | 'transpose'
   sil-differentiability-witness-function-index-list ::= [0-9]+ (' ' [0-9]+)*
 
-  differentiability_witness_function [jvp] [parameters 0] [results 0] \
+  differentiability_witness_function [vjp] [reverse] [parameters 0] [results 0] \
     <T where T: Differentiable> @foo : $(T) -> T
 
 Looks up a differentiability witness function (JVP, VJP, or transpose) for
@@ -7086,6 +7090,7 @@ look up: ``[jvp]``, ``[vjp]``, or ``[transpose]``.
 The remaining components identify the SIL differentiability witness:
 
 - Original function name.
+- Differentiability kind.
 - Parameter indices.
 - Result indices.
 - Witness generic parameter clause (optional). When parsing SIL, the parsed

include/swift/AST/ASTMangler.h

@@ -204,6 +204,11 @@ class ASTMangler : public Mangler {
       CanType fromType, CanType toType, GenericSignature signature,
       AutoDiffLinearMapKind linearMapKind);
 
+  /// Mangle a SIL differentiability witness.
+  std::string mangleSILDifferentiabilityWitness(StringRef originalName,
+                                                DifferentiabilityKind kind,
+                                                AutoDiffConfig config);
+
   /// Mangle the AutoDiff generated declaration for the given:
   /// - Generated declaration kind: linear map struct or branching trace enum.
   /// - Mangled original function name.
@@ -217,14 +222,6 @@ class ASTMangler : public Mangler {
                                      AutoDiffLinearMapKind linearMapKind,
                                      AutoDiffConfig config);
 
-  /// Mangle a SIL differentiability witness key:
-  /// - Mangled original function name.
-  /// - Parameter indices.
-  /// - Result indices.
-  /// - Derivative generic signature (optional).
-  std::string
-  mangleSILDifferentiabilityWitnessKey(SILDifferentiabilityWitnessKey key);
-
   std::string mangleKeyPathGetterThunkHelper(const AbstractStorageDecl *property,
                                              GenericSignature signature,
                                              CanType baseType,

include/swift/AST/AutoDiff.h

@@ -195,6 +195,7 @@ struct AutoDiffConfig {
   IndexSubset *resultIndices;
   GenericSignature derivativeGenericSignature;
 
+  /*implicit*/ AutoDiffConfig() = default;
   /*implicit*/ AutoDiffConfig(
       IndexSubset *parameterIndices, IndexSubset *resultIndices,
       GenericSignature derivativeGenericSignature = GenericSignature())
@@ -545,10 +546,20 @@ struct TangentPropertyInfo {
 
 void simple_display(llvm::raw_ostream &OS, TangentPropertyInfo info);
 
-/// The key type used for uniquing `SILDifferentiabilityWitness` in
-/// `SILModule`: original function name, parameter indices, result indices, and
-/// derivative generic signature.
-using SILDifferentiabilityWitnessKey = std::pair<StringRef, AutoDiffConfig>;
+/// The key type used for uniquing `SILDifferentiabilityWitness` in `SILModule`.
+struct SILDifferentiabilityWitnessKey {
+  StringRef originalFunctionName;
+  DifferentiabilityKind kind;
+  AutoDiffConfig config;
+
+  void print(llvm::raw_ostream &s = llvm::outs()) const;
+};
+
+inline llvm::raw_ostream &operator<<(
+    llvm::raw_ostream &s, const SILDifferentiabilityWitnessKey &key) {
+  key.print(s);
+  return s;
+}
 
 /// Returns `true` iff differentiable programming is enabled.
 bool isDifferentiableProgrammingEnabled(SourceFile &SF);
@@ -676,6 +687,9 @@ getAutoDiffFunctionKind(AutoDiffDerivativeFunctionKind kind);
 
 AutoDiffFunctionKind getAutoDiffFunctionKind(AutoDiffLinearMapKind kind);
 
+MangledDifferentiabilityKind
+getMangledDifferentiabilityKind(DifferentiabilityKind kind);
+
 } // end namespace autodiff
 } // end namespace swift
 
@@ -688,6 +702,8 @@ using swift::GenericSignature;
 using swift::IndexSubset;
 using swift::SILAutoDiffDerivativeFunctionKey;
 using swift::SILFunctionType;
+using swift::DifferentiabilityKind;
+using swift::SILDifferentiabilityWitnessKey;
 
 template <typename T> struct DenseMapInfo;
 
@@ -760,8 +776,8 @@ template <> struct DenseMapInfo<AutoDiffDerivativeFunctionKind> {
 };
 
 template <> struct DenseMapInfo<SILAutoDiffDerivativeFunctionKey> {
-  static bool isEqual(const SILAutoDiffDerivativeFunctionKey lhs,
-                      const SILAutoDiffDerivativeFunctionKey rhs) {
+  static bool isEqual(const SILAutoDiffDerivativeFunctionKey &lhs,
+                      const SILAutoDiffDerivativeFunctionKey &rhs) {
     return lhs.originalType == rhs.originalType &&
            lhs.parameterIndices == rhs.parameterIndices &&
            lhs.resultIndices == rhs.resultIndices &&
@@ -803,6 +819,36 @@ template <> struct DenseMapInfo<SILAutoDiffDerivativeFunctionKey> {
   }
 };
 
+template <> struct DenseMapInfo<SILDifferentiabilityWitnessKey> {
+  static bool isEqual(const SILDifferentiabilityWitnessKey &lhs,
+                      const SILDifferentiabilityWitnessKey &rhs) {
+    return DenseMapInfo<StringRef>::isEqual(
+               lhs.originalFunctionName, rhs.originalFunctionName) &&
+           DenseMapInfo<unsigned>::isEqual(
+               (unsigned)lhs.kind, (unsigned)rhs.kind) &&
+           DenseMapInfo<AutoDiffConfig>::isEqual(lhs.config, rhs.config);
+  }
+
+  static inline SILDifferentiabilityWitnessKey getEmptyKey() {
+    return {DenseMapInfo<StringRef>::getEmptyKey(),
+            (DifferentiabilityKind)DenseMapInfo<unsigned>::getEmptyKey(),
+            DenseMapInfo<AutoDiffConfig>::getEmptyKey()};
+  }
+
+  static inline SILDifferentiabilityWitnessKey getTombstoneKey() {
+    return {DenseMapInfo<StringRef>::getTombstoneKey(),
+            (DifferentiabilityKind)DenseMapInfo<unsigned>::getTombstoneKey(),
+            DenseMapInfo<AutoDiffConfig>::getTombstoneKey()};
+  }
+
+  static unsigned getHashValue(const SILDifferentiabilityWitnessKey &val) {
+    return hash_combine(
+        DenseMapInfo<StringRef>::getHashValue(val.originalFunctionName),
+        DenseMapInfo<unsigned>::getHashValue((unsigned)val.kind),
+        DenseMapInfo<AutoDiffConfig>::getHashValue(val.config));
+  }
+};
+
 } // end namespace llvm
 
 #endif // SWIFT_AST_AUTODIFF_H

include/swift/AST/DiagnosticsParse.def

@@ -676,6 +676,9 @@ ERROR(sil_witness_protocol_conformance_not_found,none,
 // SIL differentiability witnesses
 ERROR(sil_diff_witness_expected_token,PointsToFirstBadToken,
       "expected '%0' in differentiability witness", (StringRef))
+ERROR(sil_diff_witness_unknown_kind,PointsToFirstBadToken,
+      "unknonwn differentiability kind '%0'; expected 'forward', 'reverse', "
+      "'normal', or 'linear'", (StringRef))
 ERROR(sil_diff_witness_serialized_declaration,none,
       "differentiability witness declaration should not be serialized", ())
 ERROR(sil_diff_witness_undefined,PointsToFirstBadToken,

include/swift/Demangling/DemangleNodes.def

@@ -313,6 +313,7 @@ NODE(AutoDiffFunctionKind)
 NODE(AutoDiffSelfReorderingReabstractionThunk)
 NODE(AutoDiffSubsetParametersThunk)
 NODE(AutoDiffDerivativeVTableThunk)
+NODE(DifferentiabilityWitness)
 NODE(IndexSubset)
 
 #undef CONTEXT_NODE

include/swift/Demangling/Demangler.h

@@ -573,6 +573,7 @@ class Demangler : public NodeFactory {
   NodePointer demangleAutoDiffFunctionKind();
   NodePointer demangleAutoDiffSubsetParametersThunk();
   NodePointer demangleAutoDiffSelfReorderingReabstractionThunk();
+  NodePointer demangleDifferentiabilityWitness();
   NodePointer demangleIndexSubset();
 
   bool demangleBoundGenerics(Vector<NodePointer> &TypeListList,

include/swift/IRGen/Linking.h

@@ -380,7 +380,8 @@ class LinkEntity {
     /// ProtocolConformance*.
     ProtocolWitnessTableLazyCacheVariable,
 
-    /// A SIL differentiability witness.
+    /// A SIL differentiability witness. The pointer is a
+    /// SILDifferentiabilityWitness*.
     DifferentiabilityWitness,
 
     // Everything following this is a type kind.

include/swift/SIL/SILDifferentiabilityWitness.h

@@ -47,6 +47,8 @@ class SILDifferentiabilityWitness
   SILLinkage Linkage;
   /// The original function.
   SILFunction *OriginalFunction;
+  /// The differentiability kind.
+  DifferentiabilityKind Kind;
   /// The derivative configuration: parameter indices, result indices, and
   /// derivative generic signature (optional). The derivative generic signature
   /// may contain same-type requirements such that all generic parameters are
@@ -69,27 +71,32 @@ class SILDifferentiabilityWitness
 
   SILDifferentiabilityWitness(
       SILModule &module, SILLinkage linkage, SILFunction *originalFunction,
-      IndexSubset *parameterIndices, IndexSubset *resultIndices,
-      GenericSignature derivativeGenSig, SILFunction *jvp, SILFunction *vjp,
-      bool isDeclaration, bool isSerialized, const DeclAttribute *attribute)
+      DifferentiabilityKind kind, IndexSubset *parameterIndices,
+      IndexSubset *resultIndices, GenericSignature derivativeGenSig,
+      SILFunction *jvp, SILFunction *vjp, bool isDeclaration, bool isSerialized,
+      const DeclAttribute *attribute)
       : Module(module), Linkage(linkage), OriginalFunction(originalFunction),
+        Kind(kind),
         Config(parameterIndices, resultIndices, derivativeGenSig.getPointer()),
         JVP(jvp), VJP(vjp), IsDeclaration(isDeclaration),
-        IsSerialized(isSerialized), Attribute(attribute) {}
+        IsSerialized(isSerialized), Attribute(attribute) {
+    assert(kind != DifferentiabilityKind::NonDifferentiable);
+  }
 
 public:
   static SILDifferentiabilityWitness *
   createDeclaration(SILModule &module, SILLinkage linkage,
-                    SILFunction *originalFunction,
+                    SILFunction *originalFunction, DifferentiabilityKind kind,
                     IndexSubset *parameterIndices, IndexSubset *resultIndices,
                     GenericSignature derivativeGenSig,
                     const DeclAttribute *attribute = nullptr);
 
   static SILDifferentiabilityWitness *createDefinition(
       SILModule &module, SILLinkage linkage, SILFunction *originalFunction,
-      IndexSubset *parameterIndices, IndexSubset *resultIndices,
-      GenericSignature derivativeGenSig, SILFunction *jvp, SILFunction *vjp,
-      bool isSerialized, const DeclAttribute *attribute = nullptr);
+      DifferentiabilityKind kind, IndexSubset *parameterIndices,
+      IndexSubset *resultIndices, GenericSignature derivativeGenSig,
+      SILFunction *jvp, SILFunction *vjp, bool isSerialized,
+      const DeclAttribute *attribute = nullptr);
 
   void convertToDefinition(SILFunction *jvp, SILFunction *vjp,
                            bool isSerialized);
@@ -98,6 +105,7 @@ class SILDifferentiabilityWitness
   SILModule &getModule() const { return Module; }
   SILLinkage getLinkage() const { return Linkage; }
   SILFunction *getOriginalFunction() const { return OriginalFunction; }
+  DifferentiabilityKind getKind() const { return Kind; }
   const AutoDiffConfig &getConfig() const { return Config; }
   IndexSubset *getParameterIndices() const { return Config.parameterIndices; }
   IndexSubset *getResultIndices() const { return Config.resultIndices; }

include/swift/SILOptimizer/Differentiation/Common.h

@@ -250,8 +250,8 @@ findMinimalDerivativeConfiguration(AbstractFunctionDecl *original,
 /// \param parameterIndices must be lowered to SIL.
 /// \param resultIndices must be lowered to SIL.
 SILDifferentiabilityWitness *getOrCreateMinimalASTDifferentiabilityWitness(
-    SILModule &module, SILFunction *original, IndexSubset *parameterIndices,
-    IndexSubset *resultIndices);
+    SILModule &module, SILFunction *original, DifferentiabilityKind kind,
+    IndexSubset *parameterIndices, IndexSubset *resultIndices);
 
 } // end namespace autodiff