Merge pull request #2440 from swiftwasm/main

[pull] swiftwasm from main

Author: pull[bot]

docs/ABI/Mangling.rst

@@ -229,6 +229,7 @@ types where the metadata itself has unknown layout.)
   global ::= entity generic-signature? type type* 'Tk' // key path setter
   global ::= type generic-signature 'TH' // key path equality
   global ::= type generic-signature 'Th' // key path hasher
+  global ::= global generic-signature? 'TJ' AUTODIFF-FUNCTION-KIND INDEX-SUBSET 'p' INDEX-SUBSET 'r' // autodiff function
 
   global ::= protocol 'TL'               // protocol requirements base descriptor
   global ::= assoc-type-name 'Tl'        // associated type descriptor
@@ -271,6 +272,16 @@ are always non-polymorphic ``<impl-function-type>`` types.
 ``<VALUE-WITNESS-KIND>`` differentiates the kinds of value
 witness functions for a type.
 
+::
+
+  AUTODIFF-FUNCTION-KIND ::= 'f'        // JVP (forward-mode derivative)
+  AUTODIFF-FUNCTION-KIND ::= 'r'        // VJP (reverse-mode derivative)
+  AUTODIFF-FUNCTION-KIND ::= 'd'        // differential
+  AUTODIFF-FUNCTION-KIND ::= 'p'        // pullback
+
+``<AUTODIFF-FUNCTION-KIND>`` differentiates the kinds of functions assocaited
+with a differentiable function used for differentiable programming.
+
 ::
 
   global ::= generic-signature? type 'WOy' // Outlined copy
@@ -1004,6 +1015,13 @@ Numbers and Indexes
 ``<INDEX>`` is a production for encoding numbers in contexts that can't
 end in a digit; it's optimized for encoding smaller numbers.
 
+::
+
+  INDEX-SUBSET ::= ('S' | 'U')+
+
+``<INDEX-SUBSET>`` is encoded like a bit vector and is optimized for encoding
+indices with a small upper bound.
+
 Function Specializations
 ~~~~~~~~~~~~~~~~~~~~~~~~
 

include/swift/AST/ASTMangler.h

@@ -166,23 +166,23 @@ class ASTMangler : public Mangler {
                                                    bool predefined);
 
   /// Mangle the derivative function (JVP/VJP) for the given:
-  /// - Mangled original function name.
+  /// - Mangled original function declaration.
   /// - Derivative function kind.
   /// - Derivative function configuration: parameter/result indices and
   ///   derivative generic signature.
   std::string
-  mangleAutoDiffDerivativeFunctionHelper(StringRef name,
-                                         AutoDiffDerivativeFunctionKind kind,
-                                         AutoDiffConfig config);
+  mangleAutoDiffDerivativeFunction(const AbstractFunctionDecl *originalAFD,
+                                   AutoDiffDerivativeFunctionKind kind,
+                                   AutoDiffConfig config);
 
   /// Mangle the linear map (differential/pullback) for the given:
-  /// - Mangled original function name.
+  /// - Mangled original function declaration.
   /// - Linear map kind.
   /// - Derivative function configuration: parameter/result indices and
   ///   derivative generic signature.
-  std::string mangleAutoDiffLinearMapHelper(StringRef name,
-                                            AutoDiffLinearMapKind kind,
-                                            AutoDiffConfig config);
+  std::string mangleAutoDiffLinearMap(const AbstractFunctionDecl *originalAFD,
+                                      AutoDiffLinearMapKind kind,
+                                      AutoDiffConfig config);
 
   /// Mangle the AutoDiff generated declaration for the given:
   /// - Generated declaration kind: linear map struct or branching trace enum.
@@ -255,6 +255,8 @@ class ASTMangler : public Mangler {
 
   std::string mangleOpaqueTypeDecl(const ValueDecl *decl);
 
+  std::string mangleGenericSignature(const GenericSignature sig);
+
   enum SpecialContext {
     ObjCContext,
     ClangImporterContext,
@@ -427,6 +429,12 @@ class ASTMangler : public Mangler {
   void appendSymbolicReference(SymbolicReferent referent);
 
   void appendOpaqueDeclName(const OpaqueTypeDecl *opaqueDecl);
+
+  void beginManglingWithAutoDiffOriginalFunction(
+      const AbstractFunctionDecl *afd);
+  void appendAutoDiffFunctionParts(char functionKindCode,
+                                   AutoDiffConfig config);
+  void appendIndexSubset(IndexSubset *indexSubset);
 };
 
 } // end namespace Mangle

include/swift/AST/AutoDiff.h

@@ -26,6 +26,7 @@
 #include "swift/AST/TypeAlignments.h"
 #include "swift/Basic/Range.h"
 #include "swift/Basic/SourceLoc.h"
+#include "swift/Demangling/Demangle.h"
 #include "llvm/ADT/StringExtras.h"
 #include "llvm/Support/Error.h"
 
@@ -655,6 +656,17 @@ getDifferentiabilityWitnessGenericSignature(GenericSignature origGenSig,
 
 } // end namespace swift
 
+namespace swift {
+namespace Demangle {
+
+AutoDiffFunctionKind
+getAutoDiffFunctionKind(AutoDiffDerivativeFunctionKind kind);
+
+AutoDiffFunctionKind getAutoDiffFunctionKind(AutoDiffLinearMapKind kind);
+
+} // end namespace autodiff
+} // end namespace swift
+
 namespace llvm {
 
 using swift::AutoDiffConfig;

include/swift/Demangling/Demangle.h

@@ -117,6 +117,13 @@ enum class FunctionSigSpecializationParamKind : unsigned {
   ExistentialToGeneric = 1 << 10,
 };
 
+enum class AutoDiffFunctionKind : char {
+  JVP = 'f',
+  VJP = 'r',
+  Differential = 'd',
+  Pullback = 'p',
+};
+
 /// The pass that caused the specialization to occur. We use this to make sure
 /// that two passes that generate similar changes do not yield the same
 /// mangling. This currently cannot happen, so this is just a safety measure

include/swift/Demangling/DemangleNodes.def

@@ -308,6 +308,9 @@ NODE(CanonicalPrespecializedGenericTypeCachingOnceToken)
 
 // Added in Swift 5.5
 NODE(AsyncFunctionPointer)
+NODE(AutoDiffFunction)
+NODE(AutoDiffFunctionKind)
+NODE(IndexSubset)
 
 #undef CONTEXT_NODE
 #undef NODE

include/swift/Demangling/Demangler.h

@@ -569,6 +569,8 @@ class Demangler : public NodeFactory {
 
   NodePointer demangleTypeMangling();
   NodePointer demangleSymbolicReference(unsigned char rawKind);
+  NodePointer demangleAutoDiffFunctionKind();
+  NodePointer demangleIndexSubset();
 
   bool demangleBoundGenerics(Vector<NodePointer> &TypeListList,
                              NodePointer &RetroactiveConformances);

include/swift/SILOptimizer/Differentiation/JVPCloner.h

@@ -39,9 +39,8 @@ class JVPCloner final {
   ///
   /// The parent JVP cloner stores the original function and an empty
   /// to-be-generated pullback function.
-  explicit JVPCloner(ADContext &context, SILFunction *original,
-                     SILDifferentiabilityWitness *witness, SILFunction *jvp,
-                     DifferentiationInvoker invoker);
+  explicit JVPCloner(ADContext &context, SILDifferentiabilityWitness *witness,
+                     SILFunction *jvp, DifferentiationInvoker invoker);
   ~JVPCloner();
 
   /// Performs JVP generation on the empty JVP function. Returns true if any

include/swift/SILOptimizer/Differentiation/VJPCloner.h

@@ -39,9 +39,8 @@ class VJPCloner final {
   ///
   /// The parent VJP cloner stores the original function and an empty
   /// to-be-generated pullback function.
-  explicit VJPCloner(ADContext &context, SILFunction *original,
-                     SILDifferentiabilityWitness *witness, SILFunction *vjp,
-                     DifferentiationInvoker invoker);
+  explicit VJPCloner(ADContext &context, SILDifferentiabilityWitness *witness,
+                     SILFunction *vjp, DifferentiationInvoker invoker);
   ~VJPCloner();
 
   ADContext &getContext() const;

include/swift/SILOptimizer/Utils/DifferentiationMangler.h

@@ -0,0 +1,46 @@
+//===------- DifferentiationMangler.h --------- differentiation -*- C++ -*-===//
+//
+// This source file is part of the Swift.org open source project
+//
+// Copyright (c) 2014 - 2017 Apple Inc. and the Swift project authors
+// Licensed under Apache License v2.0 with Runtime Library Exception
+//
+// See https://swift.org/LICENSE.txt for license information
+// See https://swift.org/CONTRIBUTORS.txt for the list of Swift project authors
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef SWIFT_SIL_UTILS_DIFFERENTIATIONMANGLER_H
+#define SWIFT_SIL_UTILS_DIFFERENTIATIONMANGLER_H
+
+#include "swift/AST/ASTMangler.h"
+#include "swift/AST/AutoDiff.h"
+#include "swift/Basic/NullablePtr.h"
+#include "swift/Demangling/Demangler.h"
+#include "swift/SIL/SILFunction.h"
+
+namespace swift {
+namespace Mangle {
+
+/// A mangler for generated differentiation functions.
+class DifferentiationMangler : public ASTMangler {
+public:
+  DifferentiationMangler() {}
+  /// Returns the mangled name for a differentiation function of the given kind.
+  std::string mangle(SILFunction *originalFunction,
+                     Demangle::AutoDiffFunctionKind kind,
+                     AutoDiffConfig config);
+  /// Returns the mangled name for a derivative function of the given kind.
+  std::string mangleDerivativeFunction(SILFunction *originalFunction,
+                                       AutoDiffDerivativeFunctionKind kind,
+                                       AutoDiffConfig config);
+  /// Returns the mangled name for a linear map of the given kind.
+  std::string mangleLinearMap(SILFunction *originalFunction,
+                              AutoDiffLinearMapKind kind,
+                              AutoDiffConfig config);
+};
+
+} // end namespace Mangle
+} // end namespace swift
+
+#endif /* SWIFT_SIL_UTILS_DIFFERENTIATIONMANGLER_H */

include/swift/Sema/ConstraintLocator.h

@@ -762,6 +762,20 @@ class LocatorPathElt::ArgumentAttribute final : public StoredIntegerElement<1> {
   }
 };
 
+class LocatorPathElt::ConformanceRequirement final
+    : public StoredPointerElement<ProtocolDecl> {
+public:
+  ConformanceRequirement(ProtocolDecl *protocol)
+      : StoredPointerElement(PathElementKind::ConformanceRequirement,
+                             protocol) {}
+
+  ProtocolDecl *getRequirement() const { return getStoredPointer(); }
+
+  static bool classof(const LocatorPathElt *elt) {
+    return elt->getKind() == ConstraintLocator::ConformanceRequirement;
+  }
+};
+
 namespace details {
   template <typename CustomPathElement>
   class PathElement {