[AutoDiff] Remove differentiation order from AD-related instructions.

The differentiation order field in `differentiable_function` and `differentiable_function_extract` instructions is unsupported and will not be used by the current design. Quite a lot of dead code exists to try to handle `order`, but it is mostly incomplete and untested. This PR removes the differentiation order from the code base to simplify what we upstream to the 'master' branch.

Changes include:
* Remove `differentiationOrder` from `DifferentiableFunctionInst` and `DifferentiableFunctionExtractInst`.
* Make `DifferentiableFunctionInst::DifferentiableFunctionInst` take an optional pair of JVP and VJP instead of a variable-size array.
* Rename "associated functions" to "derivative functions" in `DifferentiableFunctionInst` to align better with [the design](https://forums.swift.org/t/differentiable-programming-mega-proposal/28547). Filed task [TF-882](https://bugs.swift.org/browse/TF-882) to track the renaming of all other occurrences of "associated functions".

Resolves [TF-880](https://bugs.swift.org/browse/TF-880).

Author: rxwei

docs/SIL.rst

@@ -5620,21 +5620,18 @@ differentiable_function
   sil-differentiable-function-associated-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 associated differentiation functions into a
+``@differentiable`` function. There are two associated 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``
@@ -5660,12 +5657,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]``.
+``@differentiable`` function. It must be provided with an extractee:
+``[original]``, ``[jvp]`` or ``[vjp]``.
 
 
 Assertion configuration

include/swift/AST/AutoDiff.h

@@ -458,29 +458,25 @@ struct AutoDiffAssociatedFunctionKind {
 /// compared by opaque pointer value.
 class AutoDiffAssociatedFunctionIdentifier : public llvm::FoldingSetNode {
   const AutoDiffAssociatedFunctionKind kind;
-  const unsigned differentiationOrder;
   AutoDiffIndexSubset *const parameterIndices;
 
   AutoDiffAssociatedFunctionIdentifier(
-      AutoDiffAssociatedFunctionKind kind, unsigned differentiationOrder,
+      AutoDiffAssociatedFunctionKind kind,
       AutoDiffIndexSubset *parameterIndices) :
-    kind(kind), differentiationOrder(differentiationOrder),
-    parameterIndices(parameterIndices) {}
+    kind(kind), parameterIndices(parameterIndices) {}
 
 public:
   AutoDiffAssociatedFunctionKind getKind() const { return kind; }
-  unsigned getDifferentiationOrder() const { return differentiationOrder; }
   AutoDiffIndexSubset *getParameterIndices() const {
     return parameterIndices;
   }
 
   static AutoDiffAssociatedFunctionIdentifier *get(
-      AutoDiffAssociatedFunctionKind kind, unsigned differentiationOrder,
+      AutoDiffAssociatedFunctionKind kind,
       AutoDiffIndexSubset *parameterIndices, ASTContext &C);
 
   void Profile(llvm::FoldingSetNodeID &ID) {
     ID.AddInteger(kind);
-    ID.AddInteger(differentiationOrder);
     ID.AddPointer(parameterIndices);
   }
 };
@@ -520,29 +516,12 @@ void getSubsetParameterTypes(AutoDiffIndexSubset *indices,
 AutoDiffIndexSubset *getLoweredParameterIndices(AutoDiffIndexSubset *indices,
                                                 AnyFunctionType *type);
 
-/// Returns the offset for an associated function at a specific differentiation
-/// order.
-/// This is used for both ordering in the `differentiable_function` instruction
-/// and ABI layout.
-///
-///                Order 1       Order 2     ...
-/// |----------| |-----|-----| |-----|-----| ...
-/// | Original | | JVP | VJP | | JVP | VJP | ...
-/// |----------| |-----|-----| |-----|-----| ...
-unsigned
-getOffsetForAutoDiffAssociatedFunction(unsigned order,
-                                       AutoDiffAssociatedFunctionKind kind);
-
-unsigned
-getNumAutoDiffAssociatedFunctions(unsigned differentiationOrder);
-
 /// Retrieve config from the function name of a variant of
 /// `Builtin.autodiffApply`, e.g. `Builtin.autodiffApply_jvp_arity2_order1`.
 /// Returns true if the function name is parsed successfully.
 bool getBuiltinAutoDiffApplyConfig(StringRef operationName,
                                    AutoDiffAssociatedFunctionKind &kind,
-                                   unsigned &arity, unsigned &order,
-                                   bool &rethrows);
+                                   unsigned &arity, bool &rethrows);
 
 /// Computes the correct linkage for an associated function given the linkage of
 /// the original function. If the original linkage is not external and

include/swift/AST/DiagnosticsParse.def

@@ -1594,10 +1594,6 @@ ERROR(sil_attr_differentiable_expected_source_index,PointsToFirstBadToken,
 // SIL autodiff
 ERROR(sil_inst_autodiff_attr_expected_rsquare,PointsToFirstBadToken,
       "expected ']' to complete the %0", (StringRef))
-ERROR(sil_inst_autodiff_expected_order,PointsToFirstBadToken,
-      "expected an unsigned integer indicating the differentiation order", ())
-ERROR(sil_inst_autodiff_expected_nonzero_order,PointsToFirstBadToken,
-      "expected a non-zero differentiation order", ())
 ERROR(sil_inst_autodiff_expected_parameter_index,PointsToFirstBadToken,
       "expected the index of a parameter to differentiate with respect to", ())
 ERROR(sil_inst_autodiff_operand_list_expected_lbrace,PointsToFirstBadToken,

include/swift/AST/Types.h

@@ -3099,8 +3099,8 @@ class AnyFunctionType : public TypeBase {
   }
 
   // SWIFT_ENABLE_TENSORFLOW
-  /// Given `indices`, `differentiationOrder`, and `kind`, calculates the type
-  /// of the corresponding autodiff associated function.
+  /// Given `indices` and `kind`, calculates the type of the corresponding
+  /// autodiff associated function.
   ///
   /// By default, if the original type has a self parameter list and parameter
   /// indices include self, the computed associated function type will return a
@@ -3116,7 +3116,7 @@ class AnyFunctionType : public TypeBase {
   /// function, including `@differentiable`.
   AnyFunctionType *getAutoDiffAssociatedFunctionType(
       AutoDiffIndexSubset *indices, unsigned resultIndex,
-      unsigned differentiationOrder, AutoDiffAssociatedFunctionKind kind,
+      AutoDiffAssociatedFunctionKind kind,
       LookupConformanceFn lookupConformance,
       GenericSignature *whereClauseGenericSignature = nullptr,
       bool makeSelfParamFirst = false);
@@ -4216,16 +4216,16 @@ class SILFunctionType final : public TypeBase, public llvm::FoldingSetNode,
 
   // SWIFT_ENABLE_TENSORFLOW
   CanSILFunctionType getWithDifferentiability(
-      unsigned differentiationOrder, AutoDiffIndexSubset *parameterIndices);
+      AutoDiffIndexSubset *parameterIndices);
 
   CanSILFunctionType getWithoutDifferentiability();
 
   /// Returns the type of a differentiation function that is associated with
   /// a function of this type.
   CanSILFunctionType getAutoDiffAssociatedFunctionType(
       AutoDiffIndexSubset *parameterIndices, unsigned resultIndex,
-      unsigned differentiationOrder, AutoDiffAssociatedFunctionKind kind,
-      Lowering::TypeConverter &TC, LookupConformanceFn lookupConformance,
+      AutoDiffAssociatedFunctionKind kind, Lowering::TypeConverter &TC,
+      LookupConformanceFn lookupConformance,
       CanGenericSignature associatedFunctionGenericSignature = nullptr);
 
   /// Returns a bit vector that specifices which parameters you can

include/swift/SIL/SILBuilder.h

@@ -511,28 +511,27 @@ class SILBuilder {
 
   /// SWIFT_ENABLE_TENSORFLOW
   DifferentiableFunctionInst *createDifferentiableFunction(
-      SILLocation loc, AutoDiffIndexSubset *parameterIndices,
-      unsigned differentiationOrder, SILValue original,
-      ArrayRef<SILValue> associatedFunctions = {}) {
+      SILLocation Loc, AutoDiffIndexSubset *ParameterIndices,
+      SILValue OriginalFunction,
+      Optional<std::pair<SILValue, SILValue>> JVPAndVJPFunctions = None) {
     return insert(DifferentiableFunctionInst::create(
-        getModule(), getSILDebugLocation(loc), parameterIndices,
-        differentiationOrder, original, associatedFunctions));
+        getModule(), getSILDebugLocation(Loc), ParameterIndices,
+        OriginalFunction, JVPAndVJPFunctions, hasOwnership()));
   }
 
   DifferentiableFunctionExtractInst *createDifferentiableFunctionExtract(
-      SILLocation loc, DifferentiableFunctionExtractee extractee,
-      unsigned differentiationOrder, SILValue theFunction) {
+      SILLocation Loc, DifferentiableFunctionExtractee Extractee,
+      SILValue TheFunction) {
     return insert(new (getModule()) DifferentiableFunctionExtractInst(
-        getModule(), getSILDebugLocation(loc), extractee, differentiationOrder,
-        theFunction));
+        getModule(), getSILDebugLocation(Loc), Extractee, TheFunction));
   }
 
   DifferentiableFunctionExtractInst *
-  createDifferentiableFunctionExtractOriginal(SILLocation loc,
-                                              SILValue theFunction) {
+  createDifferentiableFunctionExtractOriginal(SILLocation Loc,
+                                              SILValue TheFunction) {
     return insert(new (getModule()) DifferentiableFunctionExtractInst(
-        getModule(), getSILDebugLocation(loc),
-        DifferentiableFunctionExtractee::Original, 0, theFunction));
+        getModule(), getSILDebugLocation(Loc),
+        DifferentiableFunctionExtractee::Original, TheFunction));
   }
 
   BuiltinInst *createBuiltin(SILLocation Loc, Identifier Name, SILType ResultTy,

include/swift/SIL/SILCloner.h

@@ -970,15 +970,14 @@ template<typename ImplClass>
 void SILCloner<ImplClass>::visitDifferentiableFunctionInst(
     DifferentiableFunctionInst *Inst) {
   getBuilder().setCurrentDebugScope(getOpScope(Inst->getDebugScope()));
-  SmallVector<SILValue, 16> mappedAssocFns;
-  mappedAssocFns.reserve(Inst->getNumAssociatedFunctions());
-  for (auto &fn : Inst->getAssociatedFunctions())
-    mappedAssocFns.push_back(getOpValue(fn.get()));
+  Optional<std::pair<SILValue, SILValue>> assocFns = None;
+  if (Inst->hasDerivativeFunctions())
+    assocFns = std::make_pair(getOpValue(Inst->getJVPFunction()),
+                              getOpValue(Inst->getVJPFunction()));
   recordClonedInstruction(
       Inst, getBuilder().createDifferentiableFunction(
                 getOpLocation(Inst->getLoc()), Inst->getParameterIndices(),
-                Inst->getDifferentiationOrder(),
-                getOpValue(Inst->getOriginalFunction()), mappedAssocFns));
+                getOpValue(Inst->getOriginalFunction()), assocFns));
 }
 
 template<typename ImplClass>
@@ -988,7 +987,6 @@ visitDifferentiableFunctionExtractInst(DifferentiableFunctionExtractInst *Inst)
   recordClonedInstruction(
       Inst, getBuilder().createDifferentiableFunctionExtract(
                 getOpLocation(Inst->getLoc()), Inst->getExtractee(),
-                Inst->getDifferentiationOrder(),
                 getOpValue(Inst->getFunctionOperand())));
 }
 // SWIFT_ENABLE_TENSORFLOW END

include/swift/SIL/SILInstruction.h

@@ -7855,60 +7855,69 @@ class DifferentiableFunctionInst final :
 private:
   friend SILBuilder;
   /// Differentiation parameter indices.
-  AutoDiffIndexSubset *parameterIndices;
-  /// The order of differentiation.
-  unsigned differentiationOrder;
-  /// The number of operands. The first operand is always the original function.
-  /// The rest of operands determined by the order of differentiation and whether
-  /// this is the new AD model or the legacy reverse-mode AD model.
-  unsigned numOperands;
+  AutoDiffIndexSubset *ParameterIndices;
+  /// Indicates whether derivative functions (JVP/VJP) exist.
+  bool HasDerivativeFunctions;
 
-  DifferentiableFunctionInst(SILModule &module, SILDebugLocation debugLoc,
-                             AutoDiffIndexSubset *parameterIndices,
-                             unsigned differentiationOrder,
-                             SILValue originalFunction,
-                             ArrayRef<SILValue> associatedFunctions);
+  DifferentiableFunctionInst(
+      SILDebugLocation DebugLoc, AutoDiffIndexSubset *ParameterIndices,
+      SILValue OriginalFunction, ArrayRef<SILValue> DerivativeFunctions,
+      bool HasOwnership);
+
+  static SILType getDifferentiableFunctionType(
+      SILValue Original, AutoDiffIndexSubset *ParameterIndices);
+
+  static ValueOwnershipKind getMergedOwnershipKind(
+      SILValue Original, ArrayRef<SILValue> DerivativeFunctions);
 
 public:
   static DifferentiableFunctionInst *create(
-      SILModule &module, SILDebugLocation debugLoc,
-      AutoDiffIndexSubset *parameterIndices, unsigned differentiationOrder,
-      SILValue originalFunction, ArrayRef<SILValue> associatedFunctions);
-
-  static SILType getAutoDiffType(SILValue original,
-                                 unsigned differentiationOrder,
-                                 AutoDiffIndexSubset *parameterIndices);
+      SILModule &Module, SILDebugLocation DebugLoc,
+      AutoDiffIndexSubset *ParameterIndices, SILValue OriginalFunction,
+      Optional<std::pair<SILValue, SILValue>> VJPAndJVPFunctions,
+      bool HasOwnership);
 
   /// Returns the original function.
-  SILValue getOriginalFunction() const { return getAllOperands()[0].get(); }
+  SILValue getOriginalFunction() const { return getOperand(0); }
 
   /// Returns differentiation indices.
-  AutoDiffIndexSubset *getParameterIndices() const {
-    return parameterIndices;
-  }
+  AutoDiffIndexSubset *getParameterIndices() const { return ParameterIndices; }
 
-  /// Returns the differentiation order.
-  unsigned getDifferentiationOrder() const {
-    return differentiationOrder;
-  }
+  /// Returns true if derivative functions (JVP/VJP) exist.
+  bool hasDerivativeFunctions() const { return HasDerivativeFunctions; }
 
-  unsigned getNumAssociatedFunctions() const {
-    return numOperands - 1;
+  /// Returns the derivative functions, namely the JVP and VJP functions, if
+  /// they exist. Otherwise, return None.
+  Optional<std::pair<SILValue, SILValue>>
+  getOptionalDerivativeFunctionPair() const {
+    if (!HasDerivativeFunctions)
+      return None;
+    return std::make_pair(getOperand(1), getOperand(2));
   }
 
-  bool hasAssociatedFunctions() const {
-    return numOperands > 1;
+  ArrayRef<Operand> getDerivativeFunctionArray() const {
+    return getAllOperands().drop_front();
   }
 
-  ArrayRef<Operand> getAssociatedFunctions() const {
-    return getAllOperands().drop_front();
+  /// Returns the JVP function.
+  SILValue getJVPFunction() const {
+    assert(HasDerivativeFunctions);
+    return getOperand(1);
   }
 
-  std::pair<SILValue, SILValue>
-  getAssociatedFunctionPair(unsigned differentiationOrder) const;
+  /// Returns the VJP function.
+  SILValue getVJPFunction() const {
+    assert(HasDerivativeFunctions);
+    return getOperand(2);
+  }
 
-  SILValue getAssociatedFunction(unsigned differentiationOrder,
-                                 AutoDiffAssociatedFunctionKind kind) const;
+  /// Returns the derivative function (JVP or VJP) that matches the given kind.
+  SILValue getDerivativeFunction(AutoDiffAssociatedFunctionKind kind) const {
+    switch (kind) {
+    case AutoDiffAssociatedFunctionKind::JVP: return getJVPFunction();
+    case AutoDiffAssociatedFunctionKind::VJP: return getVJPFunction();
+    }
+  }
 };
 
 /// `differentiable_function_extract` - given an `@differentiable` function
@@ -7917,7 +7926,7 @@ class DifferentiableFunctionInst final :
 class DifferentiableFunctionExtractInst
     : public InstructionBase<
           SILInstructionKind::DifferentiableFunctionExtractInst,
-          OwnershipForwardingSingleValueInst> {
+          SingleValueInstruction> {
 public:
   struct Extractee {
     enum innerty : unsigned {
@@ -7939,46 +7948,28 @@ class DifferentiableFunctionExtractInst
 private:
   /// The extractee.
   Extractee extractee;
-  /// The differentiation order. A zero value is only legal when the extractee
-  /// is the original function, and it is a private representation only.
-  unsigned differentiationOrder;
   /// The list containing the `@differentiable` function operand.
   FixedOperandList<1> operands;
 
   static SILType
-  getExtracteeType(SILValue function, Extractee extractee,
-                   unsigned differentiationOrder, SILModule &module);
+  getExtracteeType(SILValue function, Extractee extractee, SILModule &module);
 
 public:
   explicit DifferentiableFunctionExtractInst(
       SILModule &module, SILDebugLocation debugLoc, Extractee extractee,
-      unsigned differentiationOrder, SILValue theFunction);
+      SILValue theFunction);
 
-  Extractee getExtractee() const {
-    return extractee;
-  }
+  Extractee getExtractee() const { return extractee; }
 
   AutoDiffAssociatedFunctionKind getAssociatedFunctionKind() const {
     auto kind = extractee.getExtracteeAsAssociatedFunction();
     assert(kind);
     return *kind;
   }
 
-  SILValue getFunctionOperand() const {
-    return operands[0].get();
-  }
-
-  unsigned getDifferentiationOrder() const {
-    return differentiationOrder;
-  }
-
-  ArrayRef<Operand> getAllOperands() const {
-    return operands.asArray();
-  }
-
-  MutableArrayRef<Operand> getAllOperands() {
-    return operands.asArray();
-  }
+  SILValue getFunctionOperand() const { return operands[0].get(); }
+  ArrayRef<Operand> getAllOperands() const { return operands.asArray(); }
+  MutableArrayRef<Operand> getAllOperands() { return operands.asArray(); }
 };
 
 typedef DifferentiableFunctionExtractInst::Extractee