Merge pull request #29918 from dan-zheng/derivative-attr-diagnostics

Author: swift-ci

include/swift/AST/DiagnosticsSema.def

@@ -2993,18 +2993,18 @@ NOTE(protocol_witness_missing_differentiable_attr,none,
 // @derivative
 ERROR(derivative_attr_expected_result_tuple,none,
       "'@derivative(of:)' attribute requires function to return a two-element "
-      "tuple of type '(value: T..., pullback: (U.TangentVector) -> T.TangentVector...)' "
-      "or '(value: T..., differential: (T.TangentVector...) -> U.TangentVector)'", ())
+      "tuple; first element must have label 'value:' and second element must "
+      "have label 'pullback:' or 'differential:'", ())
 ERROR(derivative_attr_invalid_result_tuple_value_label,none,
       "'@derivative(of:)' attribute requires function to return a two-element "
-      "tuple (first element must have label 'value:')", ())
+      "tuple; first element must have label 'value:'", ())
 ERROR(derivative_attr_invalid_result_tuple_func_label,none,
       "'@derivative(of:)' attribute requires function to return a two-element "
-      "tuple (second element must have label 'pullback:' or 'differential:')",
+      "tuple; second element must have label 'pullback:' or 'differential:'",
       ())
 ERROR(derivative_attr_result_value_not_differentiable,none,
       "'@derivative(of:)' attribute requires function to return a two-element "
-      "tuple (first element type %0 must conform to 'Differentiable')", (Type))
+      "tuple; first element type %0 must conform to 'Differentiable'", (Type))
 ERROR(derivative_attr_result_func_type_mismatch,none,
       "function result's %0 type does not match %1", (Identifier, DeclName))
 NOTE(derivative_attr_result_func_type_mismatch_note,none,

lib/Sema/TypeCheckAttr.cpp

@@ -3822,6 +3822,31 @@ bool resolveDifferentiableAttrDerivativeFunctions(
   return false;
 }
 
+/// Checks whether differentiable programming is enabled for the given
+/// differentiation-related attribute. Returns true on error.
+bool checkIfDifferentiableProgrammingEnabled(
+    ASTContext &ctx, DeclAttribute *attr) {
+  auto &diags = ctx.Diags;
+  // The experimental differentiable programming flag must be enabled.
+  if (!ctx.LangOpts.EnableExperimentalDifferentiableProgramming) {
+    diags
+        .diagnose(attr->getLocation(),
+                  diag::experimental_differentiable_programming_disabled)
+        .highlight(attr->getRangeWithAt());
+    return true;
+  }
+  // The `Differentiable` protocol must be available.
+  // If unavailable, the `_Differentiation` module should be imported.
+  if (!ctx.getProtocol(KnownProtocolKind::Differentiable)) {
+    diags
+        .diagnose(attr->getLocation(), diag::attr_used_without_required_module,
+                  attr, ctx.Id_Differentiation)
+        .highlight(attr->getRangeWithAt());
+    return true;
+  }
+  return false;
+}
+
 llvm::Expected<IndexSubset *> DifferentiableAttributeTypeCheckRequest::evaluate(
     Evaluator &evaluator, DifferentiableAttr *attr) const {
   // Skip type-checking for implicit `@differentiable` attributes. We currently
@@ -3838,21 +3863,8 @@ llvm::Expected<IndexSubset *> DifferentiableAttributeTypeCheckRequest::evaluate(
   auto &diags = ctx.Diags;
   // `@differentiable` attribute requires experimental differentiable
   // programming to be enabled.
-  if (!ctx.LangOpts.EnableExperimentalDifferentiableProgramming) {
-    diags
-        .diagnose(attr->getLocation(),
-                  diag::experimental_differentiable_programming_disabled)
-        .highlight(attr->getRangeWithAt());
-    return nullptr;
-  }
-  // The `Differentiable` protocol must be available.
-  if (!ctx.getProtocol(KnownProtocolKind::Differentiable)) {
-    diags
-        .diagnose(attr->getLocation(), diag::attr_used_without_required_module,
-                  attr, ctx.Id_Differentiation)
-        .highlight(attr->getRangeWithAt());
+  if (checkIfDifferentiableProgrammingEnabled(ctx, attr))
     return nullptr;
-  }
 
   // Derivative registration is disabled for `@differentiable(linear)`
   // attributes. Instead, use `@transpose` attribute to register transpose
@@ -4011,7 +4023,7 @@ void AttributeChecker::visitDifferentiableAttr(DifferentiableAttr *attr) {
   (void)attr->getParameterIndices();
 }
 
-/// Typechecks the given derivative attribute `attr` on decl `D`.
+/// Type-checks the given `@derivative` attribute `attr` on declaration `D`.
 ///
 /// Effects are:
 /// - Sets the original function and parameter indices on `attr`.
@@ -4021,19 +4033,13 @@ void AttributeChecker::visitDifferentiableAttr(DifferentiableAttr *attr) {
 /// \returns true on error, false on success.
 static bool typeCheckDerivativeAttr(ASTContext &Ctx, Decl *D,
                                     DerivativeAttr *attr) {
-  // Note: Implementation must be idempotent because it can get called multiple
+  // Note: Implementation must be idempotent because it may be called multiple
   // times for the same attribute.
-
   auto &diags = Ctx.Diags;
-
   // `@derivative` attribute requires experimental differentiable programming
   // to be enabled.
-  auto &ctx = D->getASTContext();
-  if (!ctx.LangOpts.EnableExperimentalDifferentiableProgramming) {
-    diags.diagnose(attr->getLocation(),
-                   diag::experimental_differentiable_programming_disabled);
+  if (checkIfDifferentiableProgrammingEnabled(Ctx, attr))
     return true;
-  }
   auto *derivative = cast<FuncDecl>(D);
   auto lookupConformance =
       LookUpConformanceInModule(D->getDeclContext()->getParentModule());
@@ -4075,19 +4081,6 @@ static bool typeCheckDerivativeAttr(ASTContext &Ctx, Decl *D,
     return true;
   }
   attr->setDerivativeKind(kind);
-  // `value: R` result tuple element must conform to `Differentiable`.
-  auto diffableProto = Ctx.getProtocol(KnownProtocolKind::Differentiable);
-  auto valueResultType = valueResultElt.getType();
-  if (valueResultType->hasTypeParameter())
-    valueResultType = derivative->mapTypeIntoContext(valueResultType);
-  auto valueResultConf = TypeChecker::conformsToProtocol(
-      valueResultType, diffableProto, derivative->getDeclContext(), None);
-  if (!valueResultConf) {
-    diags.diagnose(attr->getLocation(),
-                   diag::derivative_attr_result_value_not_differentiable,
-                   valueResultElt.getType());
-    return true;
-  }
 
   // Compute expected original function type and look up original function.
   auto *originalFnType =
@@ -4274,6 +4267,7 @@ static bool typeCheckDerivativeAttr(ASTContext &Ctx, Decl *D,
                                     diffParamTypes);
 
   // Get the differentiability parameters' `TangentVector` associated types.
+  auto *diffableProto = Ctx.getProtocol(KnownProtocolKind::Differentiable);
   auto diffParamTanTypes =
       map<SmallVector<TupleTypeElt, 4>>(diffParamTypes, [&](Type paramType) {
         if (paramType->hasTypeParameter())
@@ -4287,7 +4281,19 @@ static bool typeCheckDerivativeAttr(ASTContext &Ctx, Decl *D,
         return TupleTypeElt(paramAssocType);
       });
 
+  // `value: R` result tuple element must conform to `Differentiable`.
   // Get the `TangentVector` associated type of the `value:` result type.
+  auto valueResultType = valueResultElt.getType();
+  if (valueResultType->hasTypeParameter())
+    valueResultType = derivative->mapTypeIntoContext(valueResultType);
+  auto valueResultConf = TypeChecker::conformsToProtocol(
+      valueResultType, diffableProto, derivative->getDeclContext(), None);
+  if (!valueResultConf) {
+    diags.diagnose(attr->getLocation(),
+                   diag::derivative_attr_result_value_not_differentiable,
+                   valueResultElt.getType());
+    return true;
+  }
   auto resultTanType = valueResultConf.getTypeWitnessByName(
       valueResultType, Ctx.Id_TangentVector);
 

test/AutoDiff/Sema/derivative_attr_type_checking.swift

@@ -61,29 +61,45 @@ func vjpSubtractWrt1(x: Float, y: Float) -> (value: Float, pullback: (Float) ->
   return (x - y, { $0 })
 }
 
+// Test invalid original function.
+
+// expected-error @+1 {{use of unresolved identifier 'nonexistentFunction'}}
+@derivative(of: nonexistentFunction)
+func vjpOriginalFunctionNotFound(_ x: Float) -> (value: Float, pullback: (Float) -> Float) {
+  fatalError()
+}
+
+// Test `@derivative` attribute where `value:` result does not conform to `Differentiable`.
+// Invalid original function should be diagnosed first.
+// expected-error @+1 {{use of unresolved identifier 'nonexistentFunction'}}
+@derivative(of: nonexistentFunction)
+func vjpOriginalFunctionNotFound2(_ x: Float) -> (value: Int, pullback: (Float) -> Float) {
+  fatalError()
+}
+
 // Test incorrect `@derivative` declaration type.
 
 // expected-note @+1 {{'incorrectDerivativeType' defined here}}
 func incorrectDerivativeType(_ x: Float) -> Float {
   return x
 }
 
-// expected-error @+1 {{'@derivative(of:)' attribute requires function to return a two-element tuple of type '(value: T..., pullback: (U.TangentVector) -> T.TangentVector...)' or '(value: T..., differential: (T.TangentVector...) -> U.TangentVector)'}}
+// expected-error @+1 {{'@derivative(of:)' attribute requires function to return a two-element tuple; first element must have label 'value:' and second element must have label 'pullback:' or 'differential:'}}
 @derivative(of: incorrectDerivativeType)
 func jvpResultIncorrect(x: Float) -> Float {
   return x
 }
-// expected-error @+1 {{'@derivative(of:)' attribute requires function to return a two-element tuple (first element must have label 'value:'}}
+// expected-error @+1 {{'@derivative(of:)' attribute requires function to return a two-element tuple; first element must have label 'value:'}}
 @derivative(of: incorrectDerivativeType)
 func vjpResultIncorrectFirstLabel(x: Float) -> (Float, (Float) -> Float) {
   return (x, { $0 })
 }
-// expected-error @+1 {{'@derivative(of:)' attribute requires function to return a two-element tuple (second element must have label 'pullback:' or 'differential:')}}
+// expected-error @+1 {{'@derivative(of:)' attribute requires function to return a two-element tuple; second element must have label 'pullback:' or 'differential:'}}
 @derivative(of: incorrectDerivativeType)
 func vjpResultIncorrectSecondLabel(x: Float) -> (value: Float, (Float) -> Float) {
   return (x, { $0 })
 }
-// expected-error @+1 {{'@derivative(of:)' attribute requires function to return a two-element tuple (first element type 'Int' must conform to 'Differentiable')}}
+// expected-error @+1 {{could not find function 'incorrectDerivativeType' with expected type '(Int) -> Int'}}
 @derivative(of: incorrectDerivativeType)
 func vjpResultNotDifferentiable(x: Int) -> (
   value: Int, pullback: (Int) -> Int