Move self-parameter-reordering type computation from SIL to AST.

Address review feedback from @marcrasi.

Computing self-parameter-reordered JVP/VJP types on AST function types
instead of SIL function types is safer and simpler.

Author: dan-zheng

include/swift/AST/AutoDiff.h

@@ -149,14 +149,17 @@ class AutoDiffParameterIndices : public llvm::FoldingSetNode {
   ///
   ///   functionType = (A, B) -> (C, D) -> R
   ///   if "A", "C", and "D" are in the set,
-  ///   ==> pushes {A, C, D} to `paramTypes`.
+  ///   ==> pushes {A, C, D} to `paramTypes` if `reverseCurryLevels` is false,
+  ///    or pushes {C, D, A} otherwise.
   ///
   ///   functionType = (Self) -> (A, B, C) -> R
   ///   if "Self" and "C" are in the set,
-  ///   ==> pushes {Self, C} to `paramTypes`.
+  ///   ==> pushes {Self, C} to `paramTypes` if `reverseCurryLevels` is false,
+  ///    or pushes {C, Self} otherwise.
   ///
   void getSubsetParameterTypes(AnyFunctionType *functionType,
-                               SmallVectorImpl<Type> &paramTypes) const;
+                               SmallVectorImpl<Type> &paramTypes,
+                               bool reverseCurryLevels = false) const;
 
   /// Returns a bitvector for the SILFunction parameters corresponding to the
   /// parameters in this set. In particular, this explodes tuples. For example,

include/swift/AST/Types.h

@@ -3072,14 +3072,24 @@ class AnyFunctionType : public TypeBase {
   /// Given `indices`, `differentiationOrder`, and `kind`, calculates the type
   /// of the corresponding autodiff associated function.
   ///
-  /// \note The original function type (`self`) need not be `@differentiable`,
-  /// and the resulting function will preserve all `ExtInfo` of the original
+  /// By default, if the original type has a self parameter list and parameter
+  /// indices include self, the computed associated function type will return a
+  /// linear map taking/returning self's tangent/cotangent *last* instead of
+  /// first, for consistency with SIL.
+  ///
+  /// If `makeSelfParamFirst` is true, self's tangent/cotangent is reordered to
+  /// appear first. This should be used during type-checking, e.g.
+  /// type-checking `@differentiable` and `@differentiating` attributes.
+  ///
+  /// \note The original function type (`self`) need not be `@differentiable`.
+  /// The resulting function will preserve all `ExtInfo` of the original
   /// function, including `@differentiable`.
   AnyFunctionType *getAutoDiffAssociatedFunctionType(
       AutoDiffParameterIndices *indices, unsigned resultIndex,
       unsigned differentiationOrder, AutoDiffAssociatedFunctionKind kind,
       LookupConformanceFn lookupConformance,
-      GenericSignature *whereClauseGenericSignature = nullptr);
+      GenericSignature *whereClauseGenericSignature = nullptr,
+      bool makeSelfParamFirst = false);
 
   /// Given the type of an autodiff associated function, returns the
   /// corresponding original function type.

lib/AST/AutoDiff.cpp

@@ -165,15 +165,16 @@ static void unwrapCurryLevels(AnyFunctionType *fnTy,
 ///   ==> pushes {A, C} to `paramTypes`.
 ///
 ///   functionType = (A, B) -> (C, D) -> R
-///   if "A", "C", and "D" are in the set,
-///   ==> pushes {A, C, D} to `paramTypes`.
+///   ==> pushes {A, C, D} to `paramTypes` if `reverseCurryLevels` is true,
+///    or pushes {C, D, A} otherwise.
 ///
 ///   functionType = (Self) -> (A, B, C) -> R
-///   if "Self" and "C" are in the set,
-///   ==> pushes {Self, C} to `paramTypes`.
+///   ==> pushes {Self, C} to `paramTypes` if `reverseCurryLevels` is true,
+///    or pushes {C, Self} otherwise.
 ///
 void AutoDiffParameterIndices::getSubsetParameterTypes(
-    AnyFunctionType *functionType, SmallVectorImpl<Type> &paramTypes) const {
+    AnyFunctionType *functionType, SmallVectorImpl<Type> &paramTypes,
+    bool reverseCurryLevels) const {
   SmallVector<AnyFunctionType *, 2> curryLevels;
   unwrapCurryLevels(functionType, curryLevels);
 
@@ -184,6 +185,13 @@ void AutoDiffParameterIndices::getSubsetParameterTypes(
     currentOffset += curryLevels[curryLevelIndex]->getNumParams();
   }
 
+  // If `reverseCurryLevels` is true, reverse the curry levels and offsets.
+  if (reverseCurryLevels) {
+    std::reverse(curryLevels.begin(), curryLevels.end());
+    std::reverse(curryLevelParameterIndexOffsets.begin(),
+                 curryLevelParameterIndexOffsets.end());
+  }
+
   for (unsigned curryLevelIndex : indices(curryLevels)) {
     auto *curryLevel = curryLevels[curryLevelIndex];
     unsigned parameterIndexOffset =

lib/AST/Type.cpp

@@ -4444,7 +4444,7 @@ AnyFunctionType *AnyFunctionType::getAutoDiffAssociatedFunctionType(
     AutoDiffParameterIndices *indices, unsigned resultIndex,
     unsigned differentiationOrder, AutoDiffAssociatedFunctionKind kind,
     LookupConformanceFn lookupConformance,
-    GenericSignature *whereClauseGenSig) {
+    GenericSignature *whereClauseGenSig, bool makeSelfParamFirst) {
   // JVP: (T...) -> ((R...),
   //                 (T.TangentVector...) -> (R.TangentVector...))
   // VJP: (T...) -> ((R...),
@@ -4460,7 +4460,8 @@ AnyFunctionType *AnyFunctionType::getAutoDiffAssociatedFunctionType(
   auto &ctx = getASTContext();
 
   SmallVector<Type, 8> wrtParamTypes;
-  indices->getSubsetParameterTypes(this, wrtParamTypes);
+  indices->getSubsetParameterTypes(
+      this, wrtParamTypes, /*reverseCurryLevels*/ !makeSelfParamFirst);
 
   // Unwrap curry levels. At most, two parameter lists are necessary, for
   // curried method types with a `(Self)` parameter list.

lib/SIL/SILFunctionType.cpp

@@ -1222,74 +1222,6 @@ static CanSILFunctionType getSILFunctionType(
     ? isPseudogeneric(*constant)
     : false;
 
-  // SWIFT_ENABLE_TENSORFLOW
-  // If constant is an autodiff associated function and is differentiable wrt
-  // self, reorder self so that it appears as:
-  // - The last parameter in the differential.
-  // - The last result in the pullback.
-  if (constant && constant->hasFuncDecl() &&
-      constant->autoDiffAssociatedFunctionIdentifier) {
-    auto *AFD = constant->getAbstractFunctionDecl();
-    auto *autoDiffFuncId = constant->autoDiffAssociatedFunctionIdentifier;
-    assert(results.size() == 2);
-    auto linearMapResultInfo = results.back();
-    auto linearMapType =
-        linearMapResultInfo.getSILStorageType().castTo<SILFunctionType>();
-    // Compute autodiff indices.
-    auto loweredParamIndices =
-        autoDiffFuncId->getParameterIndices()->getLowered(
-            M.getASTContext(),
-            AFD->getInterfaceType()->castTo<AnyFunctionType>());
-    SILAutoDiffIndices indices(/*source*/ 0, loweredParamIndices);
-    auto selfParamIndex = inputs.size() - 1;
-    // Reorder *only* if differentiable wrt to self.
-    if (indices.isWrtParameter(selfParamIndex)) {
-      switch (autoDiffFuncId->getKind()) {
-      // Reorder self as the first differential parameter.
-      case AutoDiffAssociatedFunctionKind::JVP: {
-        SmallVector<SILParameterInfo, 8> differentialParams;
-        differentialParams.append(linearMapType->getParameters().begin(),
-                                  linearMapType->getParameters().end());
-        std::rotate(differentialParams.begin(), differentialParams.begin() + 1,
-                    differentialParams.end());
-        auto differentialType = SILFunctionType::get(
-            linearMapType->getGenericSignature(), linearMapType->getExtInfo(),
-            linearMapType->getCoroutineKind(),
-            linearMapType->getCalleeConvention(), differentialParams,
-            linearMapType->getYields(), linearMapType->getResults(),
-            linearMapType->getOptionalErrorResult(), M.getASTContext());
-        results.pop_back();
-        auto newDifferentialResult = SILResultInfo(
-            differentialType->getCanonicalType(),
-            linearMapResultInfo.getConvention());
-        results.push_back(newDifferentialResult);
-        break;
-      }
-      // Reorder self as the last pullback result.
-      case AutoDiffAssociatedFunctionKind::VJP: {
-        SmallVector<SILResultInfo, 8> pullbackResults;
-        pullbackResults.append(linearMapType->getResults().begin(),
-                               linearMapType->getResults().end());
-        std::rotate(pullbackResults.begin(), pullbackResults.begin() + 1,
-                    pullbackResults.end());
-        auto pullbackType = SILFunctionType::get(
-            linearMapType->getGenericSignature(), linearMapType->getExtInfo(),
-            linearMapType->getCoroutineKind(),
-            linearMapType->getCalleeConvention(),
-            linearMapType->getParameters(), linearMapType->getYields(),
-            pullbackResults, linearMapType->getOptionalErrorResult(),
-            M.getASTContext());
-        auto newPullbackResult = SILResultInfo(
-            pullbackType->getCanonicalType(),
-            linearMapResultInfo.getConvention());
-        results.pop_back();
-        results.push_back(newPullbackResult);
-        break;
-      }
-      }
-    }
-  }
-
   // NOTE: SILFunctionType::ExtInfo doesn't track everything that
   // AnyFunctionType::ExtInfo tracks. For example: 'throws' or 'auto-closure'
   auto silExtInfo = SILFunctionType::ExtInfo()

lib/SILGen/SILGenPoly.cpp

@@ -4210,57 +4210,6 @@ void SILGenFunction::emitProtocolWitness(AbstractionPattern reqtOrigTy,
                                                 origWitnessFTy,
                                                 witnessKind, witnessSubs,
                                                 witnessParams, loc);
-  // SWIFT_ENABLE_TENSORFLOW
-  // Perform witness method thunking.
-  if (auto *autoDiffFuncId = witness.autoDiffAssociatedFunctionIdentifier) {
-    auto *AFD = witness.getAbstractFunctionDecl();
-    auto loweredParamIndices =
-        autoDiffFuncId->getParameterIndices()->getLowered(
-            this->getASTContext(),
-            AFD->getInterfaceType()->castTo<AnyFunctionType>());
-    SILAutoDiffIndices indices(/*source*/ 0, loweredParamIndices);
-    auto assocFnKind = autoDiffFuncId->getKind();
-
-    unsigned selfParamIndex = F.getLoweredFunctionType()->getNumParameters() - 1;
-    bool isWrtSelf = indices.isWrtParameter(selfParamIndex);
-    if (isWrtSelf && indices.parameters->getNumIndices() > 1) {
-      // Given the type of an autodiff associated method that is differentiable
-      // wrt self, return a version where self's tangent/cotangent is reordered
-      // in the returned linear map.
-      // This is done by computing the original function type and recomputing
-      // the corresponding associated function type.
-      auto getReorderedFunctionType =
-          [&](CanAnyFunctionType fnType) -> CanAnyFunctionType {
-        auto *fnOrigType = fnType->getAutoDiffOriginalFunctionType();
-        auto *newFnType = fnOrigType->getAutoDiffAssociatedFunctionType(
-            autoDiffFuncId->getParameterIndices(), /*resultIndex*/ 0,
-            /*differentiationOrder*/ 1, assocFnKind,
-            LookUpConformanceInModule(F.getModule().getSwiftModule()),
-            fnType->getOptGenericSignature());
-        CanGenericSignature canGenSig;
-        if (auto *genSig = newFnType->getOptGenericSignature())
-          canGenSig = genSig->getCanonicalSignature();
-        return CanAnyFunctionType::get(
-            canGenSig, newFnType->getParams(),
-            newFnType->getResult()->getCanonicalType(canGenSig));
-      };
-
-      // Given the abstraction pattern of an autodiff associated method that is
-      // differentiable wrt self, return a version where self's
-      // tangent/cotangent is reordered in the returned linear map.
-      auto getReorderedUpdatePattern =
-          [&](AbstractionPattern pattern) -> AbstractionPattern {
-        auto canFnTy = pattern.getAs<AnyFunctionType>();
-        canFnTy = getReorderedFunctionType(canFnTy);
-        return AbstractionPattern(pattern.getGenericSignature(), canFnTy);
-      };
-
-      witnessOrigTy = getReorderedUpdatePattern(witnessOrigTy);
-      reqtOrigTy = getReorderedUpdatePattern(reqtOrigTy);
-      witnessSubstTy = getReorderedFunctionType(witnessSubstTy);
-      reqtSubstTy = getReorderedFunctionType(reqtSubstTy);
-    }
-  }
 
   auto coroutineKind =
     witnessFnRef->getType().castTo<SILFunctionType>()->getCoroutineKind();

lib/Sema/TypeCheckAttr.cpp

@@ -3120,7 +3120,7 @@ void AttributeChecker::visitDifferentiableAttr(DifferentiableAttr *attr) {
         originalFnTy->getAutoDiffAssociatedFunctionType(
             checkedWrtParamIndices, /*resultIndex*/ 0,
             /*differentiationOrder*/ 1, AutoDiffAssociatedFunctionKind::JVP,
-            lookupConformance, whereClauseGenSig);
+            lookupConformance, whereClauseGenSig, /*makeSelfParamFirst*/ true);
 
     auto isValidJVP = [&](FuncDecl *jvpCandidate) {
       TC.validateDeclForNameLookup(jvpCandidate);
@@ -3146,7 +3146,7 @@ void AttributeChecker::visitDifferentiableAttr(DifferentiableAttr *attr) {
         originalFnTy->getAutoDiffAssociatedFunctionType(
             checkedWrtParamIndices, /*resultIndex*/ 0,
             /*differentiationOrder*/ 1, AutoDiffAssociatedFunctionKind::VJP,
-            lookupConformance, whereClauseGenSig);
+            lookupConformance, whereClauseGenSig, /*makeSelfParamFirst*/ true);
 
     auto isValidVJP = [&](FuncDecl *vjpCandidate) {
       TC.validateDeclForNameLookup(vjpCandidate);