[AutoDiff] Fix transpose type checking. (#28950)

- TF-1060: fix `@transpose` typing rules for instance methods wrt `self`.
  - The transpose of an instance method wrt `self` is now a static method
    in the same type context.
  - Nice new invariant: transpose functions are always declared in the same type
    context as the original function.
- TF-997: support `@transpose` for initializer original declarations.

Add TF-1063 negative test: `@transpose` type-checking crash for static methods.

Co-authored-by: Dan Zheng <[email protected]>

Author: bartchr808

include/swift/AST/DiagnosticsSema.def

@@ -3043,9 +3043,12 @@ ERROR(transpose_attr_overload_not_found,none,
 ERROR(transpose_attr_cannot_use_named_wrt_params,none,
       "cannot use named 'wrt' parameters in '@transpose(of:)' attribute, found "
       "%0", (Identifier))
-ERROR(transpose_attr_result_value_not_differentiable,none,
-      "'@transpose(of:)' attribute requires original function result %0 to "
-      "conform to 'Differentiable'", (Type))
+ERROR(transpose_func_wrt_self_must_be_static,none,
+     "the transpose of an instance method must be a 'static' method in the "
+     "same type when 'self' is a linearity parameter", ())
+NOTE(transpose_func_wrt_self_self_type_mismatch_note,none,
+     "the transpose is declared in %0 but the original function is declared in "
+     "%1", (Type, Type))
 
 // Automatic differentiation attributes
 ERROR(autodiff_attr_original_decl_invalid_kind,none,

lib/AST/Attr.cpp

@@ -372,51 +372,60 @@ static void printShortFormAvailable(ArrayRef<const DeclAttribute *> Attrs,
   Printer.printNewline();
 }
 
-/// Printing style for a differentiation parameter in a `wrt:` differentiation
-/// parameters clause. Used for printing `@differentiable`, `@derivative`, and
-/// `@transpose` attributes.
-enum class DifferentiationParameterPrintingStyle {
-  /// Print parameter by name.
+/// The kind of a differentiation parameter in a `wrt:` differentiation
+/// parameters clause: differentiability or linearity. Used for printing
+/// `@differentiable`, `@derivative`, and `@transpose` attributes.
+enum class DifferentiationParameterKind {
+  /// A differentiability parameter, printed by name.
   /// Used for `@differentiable` and `@derivative` attribute.
-  Name,
-  /// Print parameter by index.
+  Differentiability,
+  /// A linearity parameter, printed by index.
   /// Used for `@transpose` attribute.
-  Index
+  Linearity
 };
 
 /// Returns the differentiation parameters clause string for the given function,
-/// parameter indices, parsed parameters, . Use the parameter indices if
-/// specified; otherwise, use the parsed parameters.
+/// parameter indices, parsed parameters, and differentiation parameter kind.
+/// Use the parameter indices if specified; otherwise, use the parsed
+/// parameters.
 static std::string getDifferentiationParametersClauseString(
-    const AbstractFunctionDecl *function, IndexSubset *paramIndices,
+    const AbstractFunctionDecl *function, IndexSubset *parameterIndices,
     ArrayRef<ParsedAutoDiffParameter> parsedParams,
-    DifferentiationParameterPrintingStyle style) {
+    DifferentiationParameterKind parameterKind) {
   assert(function);
   bool isInstanceMethod = function->isInstanceMember();
+  bool isStaticMethod = function->isStatic();
   std::string result;
   llvm::raw_string_ostream printer(result);
 
   // Use the parameter indices, if specified.
-  if (paramIndices) {
-    auto parameters = paramIndices->getBitVector();
+  if (parameterIndices) {
+    auto parameters = parameterIndices->getBitVector();
     auto parameterCount = parameters.count();
     printer << "wrt: ";
     if (parameterCount > 1)
       printer << '(';
     // Check if differentiating wrt `self`. If so, manually print it first.
-    if (isInstanceMethod && parameters.test(parameters.size() - 1)) {
+    bool isWrtSelf =
+        (isInstanceMethod ||
+         (isStaticMethod &&
+          parameterKind == DifferentiationParameterKind::Linearity)) &&
+        parameters.test(parameters.size() - 1);
+    if (isWrtSelf) {
       parameters.reset(parameters.size() - 1);
       printer << "self";
       if (parameters.any())
         printer << ", ";
     }
     // Print remaining differentiation parameters.
     interleave(parameters.set_bits(), [&](unsigned index) {
-      switch (style) {
-      case DifferentiationParameterPrintingStyle::Name:
+      switch (parameterKind) {
+      // Print differentiability parameters by name.
+      case DifferentiationParameterKind::Differentiability:
         printer << function->getParameters()->get(index)->getName().str();
         break;
-      case DifferentiationParameterPrintingStyle::Index:
+      // Print linearity parameters by index.
+      case DifferentiationParameterKind::Linearity:
         printer << index;
         break;
       }
@@ -493,7 +502,7 @@ static void printDifferentiableAttrArguments(
   if (!omitWrtClause) {
     auto diffParamsString = getDifferentiationParametersClauseString(
         original, attr->getParameterIndices(), attr->getParsedParameters(),
-        DifferentiationParameterPrintingStyle::Name);
+        DifferentiationParameterKind::Differentiability);
     // Check whether differentiation parameter clause is empty.
     // Handles edge case where resolved parameter indices are unset and
     // parsed parameters are empty. This case should never trigger for
@@ -933,7 +942,7 @@ bool DeclAttribute::printImpl(ASTPrinter &Printer, const PrintOptions &Options,
     auto *derivative = cast<AbstractFunctionDecl>(D);
     auto diffParamsString = getDifferentiationParametersClauseString(
         derivative, attr->getParameterIndices(), attr->getParsedParameters(),
-        DifferentiationParameterPrintingStyle::Name);
+        DifferentiationParameterKind::Differentiability);
     if (!diffParamsString.empty())
       Printer << ", " << diffParamsString;
     Printer << ')';
@@ -948,7 +957,7 @@ bool DeclAttribute::printImpl(ASTPrinter &Printer, const PrintOptions &Options,
     auto *transpose = cast<AbstractFunctionDecl>(D);
     auto transParamsString = getDifferentiationParametersClauseString(
         transpose, attr->getParameterIndices(), attr->getParsedParameters(),
-        DifferentiationParameterPrintingStyle::Index);
+        DifferentiationParameterKind::Linearity);
     if (!transParamsString.empty())
       Printer << ", " << transParamsString;
     Printer << ')';

lib/AST/Type.cpp

@@ -5075,7 +5075,7 @@ AnyFunctionType *AnyFunctionType::getTransposeOriginalFunctionType(
   unsigned transposeParamsIndex = 0;
   bool isCurried = getResult()->is<AnyFunctionType>();
 
-  // Get the original function's result.
+  // Get the transpose function's parameters and result type.
   auto transposeParams = getParams();
   auto transposeResult = getResult();
   if (isCurried) {
@@ -5084,79 +5084,81 @@ AnyFunctionType *AnyFunctionType::getTransposeOriginalFunctionType(
     transposeResult = methodType->getResult();
   }
 
-  Type originalResult;
-  if (isCurried) {
-    // If it's curried, then the first parameter in the curried type, which is
-    // the 'Self' type, is the original result (no matter if we are
-    // transposing wrt self or not).
-    originalResult = getParams().front().getPlainType();
-  } else {
-    // If it's not curried, the last parameter, the tangent, is always the
-    // original result type as we require the last parameter of the transposing
-    // function to be the original result.
-    originalResult = transposeParams.back().getPlainType();
-  }
-  assert(originalResult);
+  // Get the original function's result type.
+  // The original result type is always equal to the type of the last
+  // parameter of the transpose function type.
+  auto originalResult = transposeParams.back().getPlainType();
 
+  // Get transposed result types.
+  // The transpose function result type may be a singular type or a tuple type.
   SmallVector<TupleTypeElt, 4> transposeResultTypes;
-  // Return type of transpose function can be a singular type or a tuple type.
   if (auto transposeResultTupleType = transposeResult->getAs<TupleType>()) {
     transposeResultTypes.append(transposeResultTupleType->getElements().begin(),
                                 transposeResultTupleType->getElements().end());
   } else {
     transposeResultTypes.push_back(transposeResult);
   }
-  assert(!transposeResultTypes.empty());
 
-  // If the function is curried and is transposing wrt 'self', then grab
-  // the type from the result list (guaranteed to be the first since 'self'
-  // is first in wrt list) and remove it. If it is still curried but not
-  // transposing wrt 'self', then the 'Self' type is the first parameter
-  // in the method.
+  // Get the `Self` type, if the transpose function type is curried.
+  // - If `self` is a linearity parameter, use the first transpose result type.
+  // - Otherwise, use the first transpose parameter type.
   unsigned transposeResultTypesIndex = 0;
   Type selfType;
   if (isCurried && wrtSelf) {
     selfType = transposeResultTypes.front().getType();
     transposeResultTypesIndex++;
   } else if (isCurried) {
-    selfType = transposeParams.front().getPlainType();
-    transposeParamsIndex++;
+    selfType = getParams().front().getPlainType();
   }
 
+  // Get the original function's parameters.
   SmallVector<AnyFunctionType::Param, 8> originalParams;
+  // The number of original parameters is equal to the sum of:
+  // - The number of original non-transposed parameters.
+  //   - This is the number of transpose parameters minus one. All transpose
+  //     parameters come from the original function, except the last parameter
+  //     (the transposed original result).
+  // - The number of original transposed parameters.
+  //   - This is the number of linearity parameters.
   unsigned originalParameterCount =
-      transposeParams.size() + wrtParamIndices->getNumIndices() - 1;
+      transposeParams.size() - 1 + wrtParamIndices->getNumIndices();
+  // Iterate over all original parameter indices.
   for (auto i : range(originalParameterCount)) {
-    // Need to check if it is the 'self' param since we handle it differently
-    // above.
-    bool lookingAtSelf = (i == wrtParamIndices->getCapacity() - 1) && wrtSelf;
-    if (wrtParamIndices->contains(i) && !lookingAtSelf) {
-      // If in wrt list, the item in the result tuple must be a parameter in the
-      // original function.
+    // Skip `self` parameter if `self` is a linearity parameter.
+    // The `self` is handled specially later to form a curried function type.
+    bool isSelfParameterAndWrtSelf =
+        wrtSelf && i == wrtParamIndices->getCapacity() - 1;
+    if (isSelfParameterAndWrtSelf)
+      continue;
+    // If `i` is a linearity parameter index, the next original parameter is
+    // the next transpose result.
+    if (wrtParamIndices->contains(i)) {
       auto resultType =
-          transposeResultTypes[transposeResultTypesIndex].getType();
+          transposeResultTypes[transposeResultTypesIndex++].getType();
       originalParams.push_back(AnyFunctionType::Param(resultType));
-      transposeResultTypesIndex++;
-    } else {
-      // Else if not in the wrt list, the parameter in the transposing function
-      // is a parameter in the original function.
-      originalParams.push_back(transposeParams[transposeParamsIndex]);
-      transposeParamsIndex++;
+    }
+    // Otherwise, the next original parameter is the next transpose parameter.
+    else {
+      originalParams.push_back(transposeParams[transposeParamsIndex++]);
     }
   }
 
+  // Compute the original function type.
   AnyFunctionType *originalType;
+  // If the transpose type is curried, the original function type is:
+  // `(Self) -> (<original parameters>) -> <original result>`.
   if (isCurried) {
-    assert(selfType);
-    // If curried, wrap the function into the 'Self' type to get a method.
+    assert(selfType && "`Self` type should be resolved");
     originalType = makeFunctionType(originalParams, originalResult, nullptr);
     originalType = makeFunctionType(AnyFunctionType::Param(selfType),
                                     originalType, getOptGenericSignature());
-  } else {
+  }
+  // Otherwise, the original function type is simply:
+  // `(<original parameters>) -> <original result>`.
+  else {
     originalType = makeFunctionType(originalParams, originalResult,
                                     getOptGenericSignature());
   }
-  assert(originalType);
   return originalType;
 }