[AutoDiff] Revamp derivative type calculation using llvm::Expected. (#31727)

Create `DerivativeFunctionTypeError` representing all potential derivative
function type calculation errors.

Make `AnyFunctionType::getAutoDiffDerivativeFunctionLinearMapType` return
`llvm::Expected<AnyFunctionType *>`. This is much safer and caller-friendly
than performing assertions.

Delete hacks in `@differentiable` and `@derivative` attribute type-checking
for verifying that `Differentiable.TangentVector` type witnesses are valid:
this is no longer necessary.

Robust fix for TF-521: invalid `Differentiable` conformances during
`@derivative` attribute type-checking.

Resolves SR-12793: bad interaction between `@differentiable` and `@derivative`
attribute type-checking and `Differentiable` derived conformances.

Author: dan-zheng

include/swift/AST/AutoDiff.h

@@ -27,6 +27,7 @@
 #include "swift/Basic/Range.h"
 #include "swift/Basic/SourceLoc.h"
 #include "llvm/ADT/StringExtras.h"
+#include "llvm/Support/Error.h"
 
 namespace swift {
 
@@ -388,6 +389,68 @@ class TangentSpace {
   NominalTypeDecl *getNominal() const;
 };
 
+/// A derivative function type calculation error.
+class DerivativeFunctionTypeError
+    : public llvm::ErrorInfo<DerivativeFunctionTypeError> {
+public:
+  enum class Kind {
+    NoSemanticResults,
+    MultipleSemanticResults,
+    NonDifferentiableParameters,
+    NonDifferentiableResult
+  };
+
+  static const char ID;
+  /// The original function type.
+  AnyFunctionType *functionType;
+  /// The error kind.
+  Kind kind;
+
+private:
+  union Value {
+    IndexSubset *indices;
+    Type type;
+    Value(IndexSubset *indices) : indices(indices) {}
+    Value(Type type) : type(type) {}
+    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);
+  };
+
+  explicit DerivativeFunctionTypeError(AnyFunctionType *functionType, Kind kind,
+                                       Type nonDiffResultType)
+      : functionType(functionType), kind(kind), value(nonDiffResultType) {
+    assert(kind == Kind::NonDifferentiableResult);
+  };
+
+  IndexSubset *getNonDifferentiableParameterIndices() const {
+    assert(kind == Kind::NonDifferentiableParameters);
+    return value.indices;
+  }
+
+  Type getNonDifferentiableResultType() const {
+    assert(kind == Kind::NonDifferentiableResult);
+    return value.type;
+  }
+
+  void log(raw_ostream &OS) const override;
+
+  std::error_code convertToErrorCode() const override {
+    return llvm::inconvertibleErrorCode();
+  }
+};
+
 /// The key type used for uniquing `SILDifferentiabilityWitness` in
 /// `SILModule`: original function name, parameter indices, result indices, and
 /// derivative generic signature.

include/swift/AST/Types.h

@@ -39,6 +39,7 @@
 #include "llvm/ADT/PointerEmbeddedInt.h"
 #include "llvm/ADT/PointerUnion.h"
 #include "llvm/ADT/SmallBitVector.h"
+#include "llvm/Support/Error.h"
 #include "llvm/Support/ErrorHandling.h"
 #include "llvm/Support/TrailingObjects.h"
 
@@ -3350,7 +3351,7 @@ class AnyFunctionType : public TypeBase {
   /// first. `makeSelfParamFirst` should be true when working with user-facing
   /// derivative function types, e.g. when type-checking `@differentiable` and
   /// `@derivative` attributes.
-  AnyFunctionType *getAutoDiffDerivativeFunctionLinearMapType(
+  llvm::Expected<AnyFunctionType *> getAutoDiffDerivativeFunctionLinearMapType(
       IndexSubset *parameterIndices, AutoDiffLinearMapKind kind,
       LookupConformanceFn lookupConformance, bool makeSelfParamFirst = false);
 

lib/AST/AutoDiff.cpp

@@ -388,3 +388,26 @@ NominalTypeDecl *TangentSpace::getNominal() const {
   assert(isTangentVector());
   return getTangentVector()->getNominalOrBoundGenericNominal();
 }
+
+const char DerivativeFunctionTypeError::ID = '\0';
+
+void DerivativeFunctionTypeError::log(raw_ostream &OS) const {
+  OS << "original function type '";
+  functionType->print(OS);
+  OS << "' ";
+  switch (kind) {
+  case Kind::NoSemanticResults:
+    OS << "has no semantic results ('Void' result)";
+    break;
+  case Kind::MultipleSemanticResults:
+    OS << "has multiple semantic results";
+    break;
+  case Kind::NonDifferentiableParameters:
+    OS << "has non-differentiable parameters: ";
+    value.indices->print(OS);
+    break;
+  case Kind::NonDifferentiableResult:
+    OS << "has non-differentiable result: " << value.type;
+    break;
+  }
+}

lib/AST/Type.cpp

@@ -5136,9 +5136,11 @@ AnyFunctionType *AnyFunctionType::getAutoDiffDerivativeFunctionType(
 
   auto originalResult = curryLevels.back()->getResult();
 
-  Type linearMapType = getAutoDiffDerivativeFunctionLinearMapType(
+  auto linearMapTypeExpected = getAutoDiffDerivativeFunctionLinearMapType(
       parameterIndices, kind.getLinearMapKind(), lookupConformance,
       makeSelfParamFirst);
+  assert(linearMapTypeExpected && "Linear map type is invalid");
+  Type linearMapType = linearMapTypeExpected.get();
 
   // Build the full derivative function type: `(T...) -> (R, LinearMapType)`.
   SmallVector<TupleTypeElt, 2> retElts;
@@ -5165,7 +5167,8 @@ AnyFunctionType *AnyFunctionType::getAutoDiffDerivativeFunctionType(
   return derivativeFunctionType;
 }
 
-AnyFunctionType *AnyFunctionType::getAutoDiffDerivativeFunctionLinearMapType(
+llvm::Expected<AnyFunctionType *>
+AnyFunctionType::getAutoDiffDerivativeFunctionLinearMapType(
     IndexSubset *parameterIndices, AutoDiffLinearMapKind kind,
     LookupConformanceFn lookupConformance, bool makeSelfParamFirst) {
   assert(!parameterIndices->isEmpty() &&
@@ -5180,15 +5183,27 @@ AnyFunctionType *AnyFunctionType::getAutoDiffDerivativeFunctionLinearMapType(
   // Get the original semantic result type.
   SmallVector<AutoDiffSemanticFunctionResultType, 1> originalResults;
   autodiff::getFunctionSemanticResultTypes(this, originalResults);
-  assert(originalResults.size() == 1 &&
-         "Only functions with one semantic result are currently supported");
+  // Error if no original semantic results.
+  if (originalResults.empty())
+    return llvm::make_error<DerivativeFunctionTypeError>(
+        this, DerivativeFunctionTypeError::Kind::NoSemanticResults);
+  // Error if multiple original semantic results.
+  // TODO(TF-1250): Support functions with multiple semantic results.
+  if (originalResults.size() > 1)
+    return llvm::make_error<DerivativeFunctionTypeError>(
+        this, DerivativeFunctionTypeError::Kind::MultipleSemanticResults);
   auto originalResult = originalResults.front();
   auto originalResultType = originalResult.type;
 
   // Get the original semantic result type's `TangentVector` associated type.
   auto resultTan =
       originalResultType->getAutoDiffTangentSpace(lookupConformance);
-  assert(resultTan && "Original result has no tangent space?");
+  // Error if original semantic result has no tangent space.
+  if (!resultTan) {
+    return llvm::make_error<DerivativeFunctionTypeError>(
+        this, DerivativeFunctionTypeError::Kind::NonDifferentiableResult,
+        originalResultType);
+  }
   auto resultTanType = resultTan->getType();
 
   // Compute the result linear map function type.
@@ -5213,7 +5228,13 @@ AnyFunctionType *AnyFunctionType::getAutoDiffDerivativeFunctionLinearMapType(
     for (auto diffParam : diffParams) {
       auto paramType = diffParam.getPlainType();
       auto paramTan = paramType->getAutoDiffTangentSpace(lookupConformance);
-      assert(paramTan && "Parameter has no tangent space?");
+      // Error if paraneter has no tangent space.
+      if (!paramTan) {
+        return llvm::make_error<DerivativeFunctionTypeError>(
+            this,
+            DerivativeFunctionTypeError::Kind::NonDifferentiableParameters,
+            parameterIndices);
+      }
       differentialParams.push_back(AnyFunctionType::Param(
           paramTan->getType(), Identifier(), diffParam.getParameterFlags()));
       if (diffParam.isInOut())
@@ -5243,7 +5264,13 @@ AnyFunctionType *AnyFunctionType::getAutoDiffDerivativeFunctionLinearMapType(
     for (auto diffParam : diffParams) {
       auto paramType = diffParam.getPlainType();
       auto paramTan = paramType->getAutoDiffTangentSpace(lookupConformance);
-      assert(paramTan && "Parameter has no tangent space?");
+      // Error if paraneter has no tangent space.
+      if (!paramTan) {
+        return llvm::make_error<DerivativeFunctionTypeError>(
+            this,
+            DerivativeFunctionTypeError::Kind::NonDifferentiableParameters,
+            parameterIndices);
+      }
       if (diffParam.isInOut()) {
         hasInoutDiffParameter = true;
         continue;

lib/Sema/TypeCheckAttr.cpp

@@ -3375,42 +3375,20 @@ DynamicallyReplacedDeclRequest::evaluate(Evaluator &evaluator,
   return nullptr;
 }
 
-/// If the given type conforms to `Differentiable` in the given context, returns
-/// the `ProtocolConformanceRef`. Otherwise, returns an invalid
-/// `ProtocolConformanceRef`.
-///
-/// This helper verifies that the `TangentVector` type witness is valid, in case
-/// the conformance has not been fully checked and the type witness cannot be
-/// resolved.
-static ProtocolConformanceRef getDifferentiableConformance(Type type,
-                                                           DeclContext *DC) {
-  auto &ctx = type->getASTContext();
-  auto *differentiableProto =
-      ctx.getProtocol(KnownProtocolKind::Differentiable);
-  auto conf =
-      TypeChecker::conformsToProtocol(type, differentiableProto, DC);
-  if (!conf)
-    return ProtocolConformanceRef();
-  // Try to get the `TangentVector` type witness, in case the conformance has
-  // not been fully checked.
-  Type tanType = conf.getTypeWitnessByName(type, ctx.Id_TangentVector);
-  if (tanType.isNull() || tanType->hasError())
-    return ProtocolConformanceRef();
-  return conf;
-};
-
 /// Returns true if the given type conforms to `Differentiable` in the given
 /// contxt. If `tangentVectorEqualsSelf` is true, also check whether the given
 /// type satisfies `TangentVector == Self`.
 static bool conformsToDifferentiable(Type type, DeclContext *DC,
                                      bool tangentVectorEqualsSelf = false) {
-  auto conf = getDifferentiableConformance(type, DC);
+  auto &ctx = type->getASTContext();
+  auto *differentiableProto =
+      ctx.getProtocol(KnownProtocolKind::Differentiable);
+  auto conf = TypeChecker::conformsToProtocol(type, differentiableProto, DC);
   if (conf.isInvalid())
     return false;
   if (!tangentVectorEqualsSelf)
     return true;
-  auto &ctx = type->getASTContext();
-  Type tanType = conf.getTypeWitnessByName(type, ctx.Id_TangentVector);
+  auto tanType = conf.getTypeWitnessByName(type, ctx.Id_TangentVector);
   return type->isEqual(tanType);
 };
 
@@ -4602,67 +4580,81 @@ static bool typeCheckDerivativeAttr(ASTContext &Ctx, Decl *D,
   // Set the resolved differentiability parameter indices in the attribute.
   attr->setParameterIndices(resolvedDiffParamIndices);
 
-  // Get the original semantic result.
-  llvm::SmallVector<AutoDiffSemanticFunctionResultType, 1> originalResults;
-  autodiff::getFunctionSemanticResultTypes(
-      originalFnType, originalResults,
-      derivative->getGenericEnvironmentOfContext());
-  // Check that original function has at least one semantic result, i.e.
-  // that the original semantic result type is not `Void`.
-  if (originalResults.empty()) {
-    diags
-        .diagnose(attr->getLocation(), diag::autodiff_attr_original_void_result,
-                  derivative->getName())
-        .highlight(attr->getOriginalFunctionName().Loc.getSourceRange());
-    attr->setInvalid();
-    return true;
-  }
-  // Check that original function does not have multiple semantic results.
-  if (originalResults.size() > 1) {
-    diags
-        .diagnose(attr->getLocation(),
-                  diag::autodiff_attr_original_multiple_semantic_results)
-        .highlight(attr->getOriginalFunctionName().Loc.getSourceRange());
-    attr->setInvalid();
-    return true;
-  }
-  auto originalResult = originalResults.front();
-  auto originalResultType = originalResult.type;
-  // Check that the original semantic result conforms to `Differentiable`.
-  auto valueResultConf = getDifferentiableConformance(
-      originalResultType, derivative->getDeclContext());
-  if (!valueResultConf) {
-    diags.diagnose(attr->getLocation(),
-                   diag::derivative_attr_result_value_not_differentiable,
-                   valueResultElt.getType());
+  // Compute the expected differential/pullback type.
+  auto expectedLinearMapTypeOrError =
+      originalFnType->getAutoDiffDerivativeFunctionLinearMapType(
+          resolvedDiffParamIndices, kind.getLinearMapKind(), lookupConformance,
+          /*makeSelfParamFirst*/ true);
+
+  // Helper for diagnosing derivative function type errors.
+  auto errorHandler = [&](const DerivativeFunctionTypeError &error) {
+    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);
+      }
+      return;
+    }
+    case DerivativeFunctionTypeError::Kind::NonDifferentiableResult:
+      auto originalResultType = error.getNonDifferentiableResultType();
+      diags.diagnose(attr->getLocation(),
+                     diag::differentiable_attr_result_not_differentiable,
+                     originalResultType);
+      attr->setInvalid();
+      return;
+    }
+  };
+  // Diagnose any derivative function type errors.
+  if (!expectedLinearMapTypeOrError) {
+    auto error = expectedLinearMapTypeOrError.takeError();
+    handleAllErrors(std::move(error), errorHandler);
     return true;
   }
-
-  // Compute the actual differential/pullback type that we use for comparison
-  // with the expected type. We must canonicalize the derivative interface type
-  // before extracting the differential/pullback type from it, so that the
-  // derivative interface type generic signature is available for simplifying
-  // types.
+  Type expectedLinearMapType = expectedLinearMapTypeOrError.get();
+  if (expectedLinearMapType->hasTypeParameter())
+    expectedLinearMapType =
+        derivative->mapTypeIntoContext(expectedLinearMapType);
+  if (expectedLinearMapType->hasArchetype())
+    expectedLinearMapType = expectedLinearMapType->mapTypeOutOfContext();
+
+  // Compute the actual differential/pullback type for comparison with the
+  // expected type. We must canonicalize the derivative interface type before
+  // extracting the differential/pullback type from it so that types are
+  // simplified via the canonical generic signature.
   CanType canActualResultType = derivativeInterfaceType->getCanonicalType();
   while (isa<AnyFunctionType>(canActualResultType)) {
     canActualResultType =
         cast<AnyFunctionType>(canActualResultType).getResult();
   }
-  CanType actualFuncEltType =
+  CanType actualLinearMapType =
       cast<TupleType>(canActualResultType).getElementType(1);
 
-  // Compute expected differential/pullback type.
-  Type expectedFuncEltType =
-      originalFnType->getAutoDiffDerivativeFunctionLinearMapType(
-          resolvedDiffParamIndices, kind.getLinearMapKind(), lookupConformance,
-          /*makeSelfParamFirst*/ true);
-  if (expectedFuncEltType->hasTypeParameter())
-    expectedFuncEltType = derivative->mapTypeIntoContext(expectedFuncEltType);
-  if (expectedFuncEltType->hasArchetype())
-    expectedFuncEltType = expectedFuncEltType->mapTypeOutOfContext();
-
   // Check if differential/pullback type matches expected type.
-  if (!actualFuncEltType->isEqual(expectedFuncEltType)) {
+  if (!actualLinearMapType->isEqual(expectedLinearMapType)) {
     // Emit differential/pullback type mismatch error on attribute.
     diags.diagnose(attr->getLocation(),
                    diag::derivative_attr_result_func_type_mismatch,
@@ -4675,7 +4667,7 @@ static bool typeCheckDerivativeAttr(ASTContext &Ctx, Decl *D,
     diags
         .diagnose(funcEltTypeRepr->getStartLoc(),
                   diag::derivative_attr_result_func_type_mismatch_note,
-                  funcResultElt.getName(), expectedFuncEltType)
+                  funcResultElt.getName(), expectedLinearMapType)
         .highlight(funcEltTypeRepr->getSourceRange());
     // Emit note showing original function location, if possible.
     if (originalAFD->getLoc().isValid())