[AutoDiff] Clean up derivative type calculation.

Remove all assertions from `AnyFunctionType::getAutoDiffDerivativeFunctionLinearMapType`.
All error cases are represented by `DerivativeFunctionTypeError` now.

Fix `DerivativeFunctionTypeError` error payloads and improve error case naming.

Author: dan-zheng

include/swift/AST/AutoDiff.h

@@ -394,9 +394,16 @@ class DerivativeFunctionTypeError
     : public llvm::ErrorInfo<DerivativeFunctionTypeError> {
 public:
   enum class Kind {
+    /// Original function type has no semantic results.
     NoSemanticResults,
+    /// Original function type has multiple semantic results.
+    // TODO(TF-1250): Support function types with multiple semantic results.
     MultipleSemanticResults,
-    NonDifferentiableParameters,
+    /// Differentiability parmeter indices are empty.
+    NoDifferentiabilityParameters,
+    /// A differentiability parameter does not conform to `Differentiable`.
+    NonDifferentiableDifferentiabilityParameter,
+    /// The original result type does not conform to `Differentiable`.
     NonDifferentiableResult
   };
 
@@ -406,42 +413,35 @@ class DerivativeFunctionTypeError
   /// The error kind.
   Kind kind;
 
+  /// The type and index of a differentiability parameter or result.
+  using TypeAndIndex = std::pair<Type, unsigned>;
+
 private:
   union Value {
-    IndexSubset *indices;
-    Type type;
-    Value(IndexSubset *indices) : indices(indices) {}
-    Value(Type type) : type(type) {}
+    TypeAndIndex typeAndIndex;
+    Value(TypeAndIndex typeAndIndex) : typeAndIndex(typeAndIndex) {}
     Value() {}
   } value;
 
 public:
   explicit DerivativeFunctionTypeError(AnyFunctionType *functionType, Kind kind)
       : functionType(functionType), kind(kind), value(Value()) {
     assert(kind == Kind::NoSemanticResults ||
-           kind == Kind::MultipleSemanticResults);
-  };
-
-  explicit DerivativeFunctionTypeError(AnyFunctionType *functionType, Kind kind,
-                                       IndexSubset *nonDiffParameterIndices)
-      : functionType(functionType), kind(kind), value(nonDiffParameterIndices) {
-    assert(kind == Kind::NonDifferentiableParameters);
+           kind == Kind::MultipleSemanticResults ||
+           kind == Kind::NoDifferentiabilityParameters);
   };
 
   explicit DerivativeFunctionTypeError(AnyFunctionType *functionType, Kind kind,
-                                       Type nonDiffResultType)
-      : functionType(functionType), kind(kind), value(nonDiffResultType) {
-    assert(kind == Kind::NonDifferentiableResult);
+                                       TypeAndIndex nonDiffTypeAndIndex)
+      : functionType(functionType), kind(kind), value(nonDiffTypeAndIndex) {
+    assert(kind == Kind::NonDifferentiableDifferentiabilityParameter ||
+           kind == Kind::NonDifferentiableResult);
   };
 
-  IndexSubset *getNonDifferentiableParameterIndices() const {
-    assert(kind == Kind::NonDifferentiableParameters);
-    return value.indices;
-  }
-
-  Type getNonDifferentiableResultType() const {
-    assert(kind == Kind::NonDifferentiableResult);
-    return value.type;
+  TypeAndIndex getNonDifferentiableTypeAndIndex() const {
+    assert(kind == Kind::NonDifferentiableDifferentiabilityParameter ||
+           kind == Kind::NonDifferentiableResult);
+    return value.typeAndIndex;
   }
 
   void log(raw_ostream &OS) const override;

lib/AST/AutoDiff.cpp

@@ -402,12 +402,20 @@ void DerivativeFunctionTypeError::log(raw_ostream &OS) const {
   case Kind::MultipleSemanticResults:
     OS << "has multiple semantic results";
     break;
-  case Kind::NonDifferentiableParameters:
-    OS << "has non-differentiable parameters: ";
-    value.indices->print(OS);
+  case Kind::NoDifferentiabilityParameters:
+    OS << "has no differentiability parameters";
     break;
-  case Kind::NonDifferentiableResult:
-    OS << "has non-differentiable result: " << value.type;
+  case Kind::NonDifferentiableDifferentiabilityParameter: {
+    auto nonDiffParam = getNonDifferentiableTypeAndIndex();
+    OS << "has non-differentiable differentiability parameter "
+       << nonDiffParam.second << ": " << nonDiffParam.first;
     break;
   }
+  case Kind::NonDifferentiableResult: {
+    auto nonDiffResult = getNonDifferentiableTypeAndIndex();
+    OS << "has non-differentiable result " << nonDiffResult.second << ": "
+       << nonDiffResult.first;
+    break;
+  }
+  }
 }

lib/AST/Type.cpp

@@ -5171,9 +5171,11 @@ llvm::Expected<AnyFunctionType *>
 AnyFunctionType::getAutoDiffDerivativeFunctionLinearMapType(
     IndexSubset *parameterIndices, AutoDiffLinearMapKind kind,
     LookupConformanceFn lookupConformance, bool makeSelfParamFirst) {
-  assert(!parameterIndices->isEmpty() &&
-         "Expected at least one differentiability parameter");
   auto &ctx = getASTContext();
+  // Error if differentiability parameter indices are empty.
+  if (parameterIndices->isEmpty())
+    return llvm::make_error<DerivativeFunctionTypeError>(
+        this, DerivativeFunctionTypeError::Kind::NoDifferentiabilityParameters);
 
   // Get differentiability parameters.
   SmallVector<AnyFunctionType::Param, 8> diffParams;
@@ -5202,7 +5204,7 @@ AnyFunctionType::getAutoDiffDerivativeFunctionLinearMapType(
   if (!resultTan) {
     return llvm::make_error<DerivativeFunctionTypeError>(
         this, DerivativeFunctionTypeError::Kind::NonDifferentiableResult,
-        originalResultType);
+        std::make_pair(originalResultType, /*index*/ 0));
   }
   auto resultTanType = resultTan->getType();
 
@@ -5225,15 +5227,17 @@ AnyFunctionType::getAutoDiffDerivativeFunctionLinearMapType(
     // - Differential: `(T0.Tan, inout T1.Tan, ...) -> Void`
     SmallVector<AnyFunctionType::Param, 4> differentialParams;
     bool hasInoutDiffParameter = false;
-    for (auto diffParam : diffParams) {
+    for (auto i : range(diffParams.size())) {
+      auto diffParam = diffParams[i];
       auto paramType = diffParam.getPlainType();
       auto paramTan = paramType->getAutoDiffTangentSpace(lookupConformance);
       // Error if paraneter has no tangent space.
       if (!paramTan) {
         return llvm::make_error<DerivativeFunctionTypeError>(
             this,
-            DerivativeFunctionTypeError::Kind::NonDifferentiableParameters,
-            parameterIndices);
+            DerivativeFunctionTypeError::Kind::
+                NonDifferentiableDifferentiabilityParameter,
+            std::make_pair(paramType, i));
       }
       differentialParams.push_back(AnyFunctionType::Param(
           paramTan->getType(), Identifier(), diffParam.getParameterFlags()));
@@ -5261,15 +5265,17 @@ AnyFunctionType::getAutoDiffDerivativeFunctionLinearMapType(
     // - Pullback: `(inout T1.Tan) -> (T0.Tan, ...)`
     SmallVector<TupleTypeElt, 4> pullbackResults;
     bool hasInoutDiffParameter = false;
-    for (auto diffParam : diffParams) {
+    for (auto i : range(diffParams.size())) {
+      auto diffParam = diffParams[i];
       auto paramType = diffParam.getPlainType();
       auto paramTan = paramType->getAutoDiffTangentSpace(lookupConformance);
       // Error if paraneter has no tangent space.
       if (!paramTan) {
         return llvm::make_error<DerivativeFunctionTypeError>(
             this,
-            DerivativeFunctionTypeError::Kind::NonDifferentiableParameters,
-            parameterIndices);
+            DerivativeFunctionTypeError::Kind::
+                NonDifferentiableDifferentiabilityParameter,
+            std::make_pair(paramType, i));
       }
       if (diffParam.isInOut()) {
         hasInoutDiffParameter = true;

lib/Sema/TypeCheckAttr.cpp

@@ -4570,14 +4570,9 @@ static bool typeCheckDerivativeAttr(ASTContext &Ctx, Decl *D,
   if (!resolvedDiffParamIndices)
     return true;
 
-  // Check if the differentiability parameter indices are valid.
-  if (checkDifferentiabilityParameters(
-          originalAFD, resolvedDiffParamIndices, originalFnType,
-          derivative->getGenericEnvironment(), derivative->getModuleContext(),
-          parsedDiffParams, attr->getLocation()))
-    return true;
-
   // Set the resolved differentiability parameter indices in the attribute.
+  // Differentiability parameter indices verification is done by
+  // `AnyFunctionType::getAutoDiffDerivativeFunctionLinearMapType` below.
   attr->setParameterIndices(resolvedDiffParamIndices);
 
   // Compute the expected differential/pullback type.
@@ -4588,43 +4583,39 @@ static bool typeCheckDerivativeAttr(ASTContext &Ctx, Decl *D,
 
   // Helper for diagnosing derivative function type errors.
   auto errorHandler = [&](const DerivativeFunctionTypeError &error) {
+    attr->setInvalid();
     switch (error.kind) {
     case DerivativeFunctionTypeError::Kind::NoSemanticResults:
       diags
           .diagnose(attr->getLocation(),
                     diag::autodiff_attr_original_multiple_semantic_results)
           .highlight(attr->getOriginalFunctionName().Loc.getSourceRange());
-      attr->setInvalid();
       return;
     case DerivativeFunctionTypeError::Kind::MultipleSemanticResults:
       diags
           .diagnose(attr->getLocation(),
                     diag::autodiff_attr_original_multiple_semantic_results)
           .highlight(attr->getOriginalFunctionName().Loc.getSourceRange());
-      attr->setInvalid();
       return;
-    case DerivativeFunctionTypeError::Kind::NonDifferentiableParameters: {
-      auto *nonDiffParamIndices = error.getNonDifferentiableParameterIndices();
-      SmallVector<AnyFunctionType::Param, 4> diffParams;
-      error.functionType->getSubsetParameters(resolvedDiffParamIndices,
-                                              diffParams);
-      for (unsigned i : range(diffParams.size())) {
-        if (!nonDiffParamIndices->contains(i))
-          continue;
-        SourceLoc loc = parsedDiffParams.empty() ? attr->getLocation()
-                                                 : parsedDiffParams[i].getLoc();
-        auto diffParamType = diffParams[i].getPlainType();
-        diags.diagnose(loc, diag::diff_params_clause_param_not_differentiable,
-                       diffParamType);
-      }
+    case DerivativeFunctionTypeError::Kind::NoDifferentiabilityParameters:
+      diags.diagnose(attr->getLocation(),
+                     diag::diff_params_clause_no_inferred_parameters);
+      return;
+    case DerivativeFunctionTypeError::Kind::
+        NonDifferentiableDifferentiabilityParameter: {
+      auto nonDiffParam = error.getNonDifferentiableTypeAndIndex();
+      SourceLoc loc = parsedDiffParams.empty()
+                          ? attr->getLocation()
+                          : parsedDiffParams[nonDiffParam.second].getLoc();
+      diags.diagnose(loc, diag::diff_params_clause_param_not_differentiable,
+                     nonDiffParam.first);
       return;
     }
     case DerivativeFunctionTypeError::Kind::NonDifferentiableResult:
-      auto originalResultType = error.getNonDifferentiableResultType();
+      auto nonDiffResult = error.getNonDifferentiableTypeAndIndex();
       diags.diagnose(attr->getLocation(),
                      diag::differentiable_attr_result_not_differentiable,
-                     originalResultType);
-      attr->setInvalid();
+                     nonDiffResult.first);
       return;
     }
   };