Merge branch 'tensorflow' of github.com:apple/swift into sil-differentiability-witness

Author: dan-zheng

docs/SIL.rst

@@ -5609,38 +5609,34 @@ differentiable_function
                       sil-differentiable-function-parameter-indices?
                       sil-differentiable-function-order?
                       sil-value ':' sil-type
-                      sil-differentiable-function-associated-functions-clause?
+                      sil-differentiable-function-derivative-functions-clause?
                       
   sil-differentiable-function-parameter-indices ::=
       '[' 'wrt' [0-9]+ (',', [0-9]+)* ']'
   sil-differentiable-function-order ::= '[' 'order' [0-9]+ ']'
-  sil-differentiable-associated-functions-clause ::=
-      'with' sil-differentiable-associated-function-list
-      (',' sil-differentiable-associated-function-list)*
-  sil-differentiable-function-associated-function-list ::=
+  sil-differentiable-derivative-functions-clause ::=
+      'with' sil-differentiable-derivative-function-list
+      (',' sil-differentiable-derivative-function-list)*
+  sil-differentiable-function-derivative-function-list ::=
       '{' sil-value ',' sil-value '}'
 
-  differentiable_function [wrt 0] [order 1] %0 : $(T) -> T \
+  differentiable_function [wrt 0] %0 : $(T) -> T \
     with {%1 : $(T) -> (T, (T) -> T), %2 : $(T) -> (T, (T) -> T)}
 
-Bundles a function with its associated differentiation functions up to a
-specified differentiation order into an ``@differentiable`` function. There are
-two associated functions per differentiation order: a Jacobian-vector products
-(JVP) function and a vector-Jacobian products (VJP) function.
+Bundles a function with its derivative functions into a ``@differentiable``
+function. There are two derivative functions: a Jacobian-vector products (JVP)
+function and a vector-Jacobian products (VJP) function.
 
 ``[wrt ...]`` specifies parameter indices that the original function is
 differentiable with respect to. When not specified, it defaults to all
 parameters.
 
-``[order ...]`` specifies the maximum differentiation order for the resulting
-function. The number of lists of associated functions is equal to the order.
-
 A ``with`` clause specifies the differentiation functions associated
 with the original function. When a ``with`` clause is not specified, the first
-operand will be differentiated to produce associated functions, and a ``with``
+operand will be differentiated to produce derivative functions, and a ``with``
 clause will be added to the instruction.
 
-In raw SIL, it is optional to provide an associated function ``with`` clause.
+In raw SIL, it is optional to provide a derivative function ``with`` clause.
 In canonical SIL, a ``with`` clause is mandatory.
 
 
@@ -5660,12 +5656,12 @@ differentiable_function_extract
   sil-differentiable-function-differentiation-order ::= '[' 'order' [0-9]+ ']'
 
   differentiable_function_extract [original] %0 : $@differentiable (T) -> T
-  differentiable_function_extract [jvp] [order 1] %0 : $@differentiable (T) -> T
-  differentiable_function_extract [vjp] [order 1] %0 : $@differentiable (T) -> T
+  differentiable_function_extract [jvp] %0 : $@differentiable (T) -> T
+  differentiable_function_extract [vjp] %0 : $@differentiable (T) -> T
 
-Extracts the original function or an associated function from the given
-``@differentiable`` function at a specific differentiation order. It must be
-provided with an extractee: ``[original]``, ``[jvp]`` or ``[vjp]``.
+Extracts the original function or a derivative function from the given
+``@differentiable`` function. It must be provided with an extractee:
+``[original]``, ``[jvp]`` or ``[vjp]``.
 
 
 Assertion configuration

include/swift/AST/ASTContext.h

@@ -110,8 +110,8 @@ namespace swift {
   class VarDecl;
   class UnifiedStatsReporter;
   // SWIFT_ENABLE_TENSORFLOW
+  class AutoDiffIndexSubset;
   class VectorSpace;
-  class AutoDiffParameterIndices;
   class DifferentiableAttr;
 
   enum class KnownProtocolKind : uint8_t;
@@ -280,8 +280,9 @@ class ASTContext final {
   /// Cache of `@differentiable` attributes keyed by parameter indices. This
   /// helps us diagnose multiple `@differentiable`s that are with respect to the
   /// same set of parameters.
-  llvm::DenseMap<std::pair<Decl *, AutoDiffParameterIndices *>,
-                 DifferentiableAttr *> DifferentiableAttrs;
+  llvm::DenseMap<std::pair<Decl *, AutoDiffIndexSubset *>,
+                 DifferentiableAttr *>
+      DifferentiableAttrs;
 
 private:
   /// The current generation number, which reflects the number of

include/swift/AST/ASTMangler.h

@@ -155,12 +155,12 @@ class ASTMangler : public Mangler {
                                              ModuleDecl *Module);
 
   // SWIFT_ENABLE_TENSORFLOW
-  // Mangle the autodiff associated function (JVP/VJP) with the given:
+  // Mangle the derivative function (JVP/VJP) with the given:
   // - Mangled original function name.
-  // - Associated function kind.
+  // - Derivative function kind.
   // - Parameter/result indices.
-  std::string mangleAutoDiffAssociatedFunctionHelper(
-      StringRef name, AutoDiffAssociatedFunctionKind kind,
+  std::string mangleAutoDiffDerivativeFunctionHelper(
+      StringRef name, AutoDiffDerivativeFunctionKind kind,
       const SILAutoDiffIndices &indices);
 
   // SWIFT_ENABLE_TENSORFLOW

include/swift/AST/Attr.h

@@ -1552,10 +1552,10 @@ class DifferentiableAttr final
   /// specified.
   FuncDecl *VJPFunction = nullptr;
   /// The differentiation parameters' indices, resolved by the type checker.
-  AutoDiffParameterIndices *ParameterIndices = nullptr;
+  AutoDiffIndexSubset *ParameterIndices = nullptr;
   /// The trailing where clause (optional).
   TrailingWhereClause *WhereClause = nullptr;
-  /// The generic signature for autodiff associated functions. Resolved by the
+  /// The generic signature for autodiff derivative functions. Resolved by the
   /// type checker based on the original function's generic signature and the
   /// attribute's where clause requirements. This is set only if the attribute
   /// has a where clause.
@@ -1571,8 +1571,7 @@ class DifferentiableAttr final
 
   explicit DifferentiableAttr(ASTContext &context, bool implicit,
                               SourceLoc atLoc, SourceRange baseRange,
-                              bool linear,
-                              AutoDiffParameterIndices *indices,
+                              bool linear, AutoDiffIndexSubset *indices,
                               Optional<DeclNameWithLoc> jvp,
                               Optional<DeclNameWithLoc> vjp,
                               GenericSignature *derivativeGenericSignature);
@@ -1588,8 +1587,7 @@ class DifferentiableAttr final
 
   static DifferentiableAttr *create(ASTContext &context, bool implicit,
                                     SourceLoc atLoc, SourceRange baseRange,
-                                    bool linear,
-                                    AutoDiffParameterIndices *indices,
+                                    bool linear, AutoDiffIndexSubset *indices,
                                     Optional<DeclNameWithLoc> jvp,
                                     Optional<DeclNameWithLoc> vjp,
                                     GenericSignature *derivativeGenSig);
@@ -1604,10 +1602,10 @@ class DifferentiableAttr final
   /// registered VJP.
   Optional<DeclNameWithLoc> getVJP() const { return VJP; }
 
-  AutoDiffParameterIndices *getParameterIndices() const {
+  AutoDiffIndexSubset *getParameterIndices() const {
     return ParameterIndices;
   }
-  void setParameterIndices(AutoDiffParameterIndices *pi) {
+  void setParameterIndices(AutoDiffIndexSubset *pi) {
     ParameterIndices = pi;
   }
 
@@ -1642,7 +1640,7 @@ class DifferentiableAttr final
 
   bool parametersMatch(const DifferentiableAttr &other) const {
     assert(ParameterIndices && other.ParameterIndices);
-    return ParameterIndices->parameters == other.ParameterIndices->parameters;
+    return ParameterIndices == other.ParameterIndices;
   }
 
   /// Get the derivative generic environment for the given `@differentiable`
@@ -1652,10 +1650,10 @@ class DifferentiableAttr final
 
   // Print the attribute to the given stream.
   // If `omitWrtClause` is true, omit printing the `wrt:` clause.
-  // If `omitAssociatedFunctions` is true, omit printing associated functions.
+  // If `omitDerivativeFunctions` is true, omit printing derivative functions.
   void print(llvm::raw_ostream &OS, const Decl *D,
              bool omitWrtClause = false,
-             bool omitAssociatedFunctions = false) const;
+             bool omitDerivativeFunctions = false) const;
 
   static bool classof(const DeclAttribute *DA) {
     return DA->getKind() == DAK_Differentiable;
@@ -1683,7 +1681,7 @@ class DifferentiatingAttr final
   /// The number of parsed parameters specified in 'wrt:'.
   unsigned NumParsedParameters = 0;
   /// The differentiation parameters' indices, resolved by the type checker.
-  AutoDiffParameterIndices *ParameterIndices = nullptr;
+  AutoDiffIndexSubset *ParameterIndices = nullptr;
 
   explicit DifferentiatingAttr(ASTContext &context, bool implicit,
                                SourceLoc atLoc, SourceRange baseRange,
@@ -1693,7 +1691,7 @@ class DifferentiatingAttr final
   explicit DifferentiatingAttr(ASTContext &context, bool implicit,
                                SourceLoc atLoc, SourceRange baseRange,
                                DeclNameWithLoc original, bool linear,
-                               AutoDiffParameterIndices *indices);
+                               AutoDiffIndexSubset *indices);
 
 public:
   static DifferentiatingAttr *create(ASTContext &context, bool implicit,
@@ -1704,7 +1702,7 @@ class DifferentiatingAttr final
   static DifferentiatingAttr *create(ASTContext &context, bool implicit,
                                      SourceLoc atLoc, SourceRange baseRange,
                                      DeclNameWithLoc original, bool linear,
-                                     AutoDiffParameterIndices *indices);
+                                     AutoDiffIndexSubset *indices);
 
   DeclNameWithLoc getOriginal() const { return Original; }
 
@@ -1725,10 +1723,10 @@ class DifferentiatingAttr final
     return NumParsedParameters;
   }
 
-  AutoDiffParameterIndices *getParameterIndices() const {
+  AutoDiffIndexSubset *getParameterIndices() const {
     return ParameterIndices;
   }
-  void setParameterIndices(AutoDiffParameterIndices *pi) {
+  void setParameterIndices(AutoDiffIndexSubset *pi) {
     ParameterIndices = pi;
   }