[AutoDiff] Check derivative function @usableFromInline consistency.

include/swift/AST/DiagnosticsSema.def

@@ -3076,6 +3076,15 @@ NOTE(derivative_attr_fix_access,none,
      "mark the derivative function as "
      "'%select{private|fileprivate|internal|@usableFromInline|@usableFromInline}0' "
      "to match the original function", (AccessLevel))
+ERROR(derivative_attr_usable_from_inline_mismatch,none,
+      "non-'@usableFromInline' original function must not have a "
+      "'@usableFromInline' derivative function", ())
+NOTE(derivative_attr_fix_add_usable_from_inline,none,
+     "consider adding '@usableFromInline' to the original function %0",
+     (DeclName))
+NOTE(derivative_attr_fix_remove_usable_from_inline,none,
+     "consider removing '@usableFromInline' from the derivative function %0",
+     (DeclName))
 
 // @transpose
 ERROR(transpose_attr_invalid_linearity_parameter_or_result,none,

lib/Sema/TypeCheckAttr.cpp

@@ -4540,11 +4540,22 @@ static bool typeCheckDerivativeAttr(ASTContext &Ctx, Decl *D,
   }
   attr->setOriginalFunction(originalAFD);
 
+  bool usableFromInlineMismatch =
+      originalAFD->getFormalAccess() <= AccessLevel::Internal &&
+      derivative->getFormalAccess() <= AccessLevel::Internal &&
+      !originalAFD->isUsableFromInline() && derivative->isUsableFromInline();
+
+  bool sameAccessLevel =
+      originalAFD->getFormalAccess() == derivative->getFormalAccess();
+
   // Returns true if:
-  // - Original function and derivative function have the same access level.
+  // - Original function and derivative function have the same access level and
+  //   the same `@usableFromInline` status.
   // - Original function is public and derivative function is internal
   //   `@usableFromInline`. This is the only special case.
   auto compatibleAccessLevels = [&]() {
+    if (usableFromInlineMismatch)
+      return false;
     if (originalAFD->getFormalAccess() == derivative->getFormalAccess())
       return true;
     return originalAFD->getFormalAccess() == AccessLevel::Public &&
@@ -4553,6 +4564,23 @@ static bool typeCheckDerivativeAttr(ASTContext &Ctx, Decl *D,
 
   // Check access level compatibility for original and derivative functions.
   if (!compatibleAccessLevels()) {
+    // Diagnose if access levels match, but derivative is `@usableFromInline`
+    // while original is not.
+    if (sameAccessLevel && usableFromInlineMismatch) {
+      diags.diagnose(originalName.Loc,
+                     diag::derivative_attr_usable_from_inline_mismatch);
+      // Suggest adding `@usableFromInline` to original.
+      originalAFD
+          ->diagnose(diag::derivative_attr_fix_add_usable_from_inline,
+                     originalAFD->getName())
+          .fixItInsert(
+              originalAFD->getAttributeInsertionLoc(/*forModifier*/ false),
+              "@usableFromInline ");
+      // Suggest remove `@usableFromInline` from derivative.
+      derivative->diagnose(diag::derivative_attr_fix_remove_usable_from_inline,
+                           derivative->getName());
+      return true;
+    }
     auto originalAccess = originalAFD->getFormalAccess();
     auto derivativeAccess =
         derivative->getFormalAccessScope().accessLevelForDiagnostics();

test/AutoDiff/Sema/derivative_attr_type_checking.swift

@@ -922,3 +922,35 @@ func internal_original_fileprivate_derivative(_ x: Float) -> Float { x }
 fileprivate func _internal_original_fileprivate_derivative(_ x: Float) -> (value: Float, pullback: (Float) -> Float) {
   fatalError()
 }
+
+// MARK: - Original vs derivative `@usableFromInline` mismatch
+
+// expected-note @+1 {{consider adding '@usableFromInline' to the original function 'internal_original_usablefrominline_derivative'}}
+func internal_original_usablefrominline_derivative(_ x: Float) -> Float { x }
+@usableFromInline
+// expected-error @+1 {{non-'@usableFromInline' original function must not have a '@usableFromInline' derivative function}}
+@derivative(of: internal_original_usablefrominline_derivative)
+// expected-note @+1 {{consider removing '@usableFromInline' from the derivative function '_internal_original_usablefrominline_derivative'}}
+func _internal_original_usablefrominline_derivative(_ x: Float) -> (value: Float, pullback: (Float) -> Float) {
+  fatalError()
+}
+
+// expected-note @+1 {{consider adding '@usableFromInline' to the original function 'internal_original_inlinable_derivative'}}
+func internal_original_inlinable_derivative(_ x: Float) -> Float { x }
+@inlinable
+// expected-error @+1 {{non-'@usableFromInline' original function must not have a '@usableFromInline' derivative function}}
+@derivative(of: internal_original_inlinable_derivative)
+// expected-note @+1 {{consider removing '@usableFromInline' from the derivative function '_internal_original_inlinable_derivative'}}
+func _internal_original_inlinable_derivative(_ x: Float) -> (value: Float, pullback: (Float) -> Float) {
+  fatalError()
+}
+
+// expected-note @+1 {{consider adding '@usableFromInline' to the original function 'internal_original_alwaysemitintoclient_derivative'}}
+func internal_original_alwaysemitintoclient_derivative(_ x: Float) -> Float { x }
+@inlinable
+// expected-error @+1 {{non-'@usableFromInline' original function must not have a '@usableFromInline' derivative function}}
+@derivative(of: internal_original_alwaysemitintoclient_derivative)
+// expected-note @+1 {{consider removing '@usableFromInline' from the derivative function '_internal_original_alwaysemitintoclient_derivative'}}
+func _internal_original_alwaysemitintoclient_derivative(_ x: Float) -> (value: Float, pullback: (Float) -> Float) {
+  fatalError()
+}

test/AutoDiff/compiler_crashers_fixed/tf1160-derivative-usable-from-inline-mismatch.swift

@@ -0,0 +1,24 @@
+// RUN: %target-swift-frontend -c %s -verify
+// REQUIRES: asserts
+
+// TF-1160: Linker error for `@usableFromInline` derivative function but
+// non-`@usableFromInline` internal original function.
+
+import _Differentiation
+
+// expected-note @+1 {{consider adding '@usableFromInline' to the original function 'internalOriginal'}}
+func internalOriginal(_ x: Float) -> Float {
+  x
+}
+
+@usableFromInline
+// expected-error @+1 {{non-'@usableFromInline' original function must not have a '@usableFromInline' derivative function}}
+@derivative(of: internalOriginal)
+// expected-note @+1 {{consider removing '@usableFromInline' from the derivative function 'usableFromInlineDerivative'}}
+func usableFromInlineDerivative(_ x: Float) -> (value: Float, pullback: (Float) -> Float) {
+  (x, { $0 })
+}
+
+// Original error: type-checking passes but TBDGen is not consistent with IRGen.
+// <unknown>:0: error: symbol 'AD__$s4main16internalOriginalyS2fF__vjp_src_0_wrt_0' (AD__$s4main16internalOriginalyS2fF__vjp_src_0_wrt_0) is in generated IR file, but not in TBD file
+// <unknown>:0: error: please file a radar or open a bug on bugs.swift.org with this code, and add -Xfrontend -validate-tbd-against-ir=none to squash the errors