[AutoDiff upstream] Relax @differentiable for protocol witnesses. (#30629)

Previously, all witnesses of a `@differentiable` protocol requirement were
required to have the same attribute (or one with superset parameter indices).

However, this leads to many annotations on witnesses and is not ideal for
usability. `@differentiable` attributes are really only significant on
public witnesses, so that they are clearly `@differentiable` at a glance (in
source code, interface files, and API documentation), without looking through
protocol conformance hierarchies.

Now, only *public* witnesses of `@differentiable` protocol requirements are
required to have the same attribute (or one with superset parameter indices).
For less-visible witnesses, an implicit `@differentiable` attribute is created
with the same configuration as the requirement's.

Resolves TF-1117.
Upstreams #29771 from tensorflow branch.

Author: dan-zheng

include/swift/AST/Attr.h

@@ -1709,6 +1709,13 @@ class DifferentiableAttr final
   /// attribute's where clause requirements. This is set only if the attribute
   /// has a where clause.
   GenericSignature DerivativeGenericSignature;
+  /// The source location of the implicitly inherited protocol requirement
+  /// `@differentiable` attribute. Used for diagnostics, not serialized.
+  ///
+  /// This is set during conformance type-checking, only for implicit
+  /// `@differentiable` attributes created for non-public protocol witnesses of
+  /// protocol requirements with `@differentiable` attributes.
+  SourceLoc ImplicitlyInheritedDifferentiableAttrLocation;
 
   explicit DifferentiableAttr(bool implicit, SourceLoc atLoc,
                               SourceRange baseRange, bool linear,
@@ -1771,6 +1778,14 @@ class DifferentiableAttr final
     DerivativeGenericSignature = derivativeGenSig;
   }
 
+  SourceLoc getImplicitlyInheritedDifferentiableAttrLocation() const {
+    return ImplicitlyInheritedDifferentiableAttrLocation;
+  }
+  void getImplicitlyInheritedDifferentiableAttrLocation(SourceLoc loc) {
+    assert(isImplicit());
+    ImplicitlyInheritedDifferentiableAttrLocation = loc;
+  }
+
   /// Get the derivative generic environment for the given `@differentiable`
   /// attribute and original function.
   GenericEnvironment *

include/swift/AST/DiagnosticsSema.def

@@ -2958,6 +2958,12 @@ ERROR(overriding_decl_missing_differentiable_attr,none,
       "overriding declaration is missing attribute '%0'", (StringRef))
 NOTE(protocol_witness_missing_differentiable_attr,none,
      "candidate is missing attribute '%0'", (StringRef))
+NOTE(protocol_witness_missing_differentiable_attr_nonpublic_other_file,none,
+     "non-public %1 %2 must have explicit '%0' attribute to satisfy "
+     "requirement %3 %4 (in protocol %6) because it is declared in a different "
+     "file than the conformance of %5 to %6",
+     (StringRef, DescriptiveDeclKind, DeclName, DescriptiveDeclKind, DeclName,
+      Type, Type))
 
 // @derivative
 ERROR(derivative_attr_expected_result_tuple,none,

lib/Sema/TypeCheckProtocol.cpp

@@ -308,7 +308,8 @@ static ValueDecl *getStandinForAccessor(AbstractStorageDecl *witness,
 /// witness.
 /// - If requirement's `@differentiable` attributes are met, or if `result` is
 ///   not viable, returns `result`.
-/// - Otherwise, returns a `DifferentiableConflict` `RequirementMatch`.
+/// - Otherwise, returns a "missing `@differentiable` attribute"
+///   `RequirementMatch`.
 // Note: the `result` argument is only necessary for using
 // `RequirementMatch::WitnessSubstitutions`.
 static RequirementMatch
@@ -384,15 +385,50 @@ matchWitnessDifferentiableAttr(DeclContext *dc, ValueDecl *req,
     }
     if (!foundExactConfig) {
       bool success = false;
-      if (supersetConfig) {
-        // If the witness has a "superset" derivative configuration, create an
-        // implicit `@differentiable` attribute with the exact requirement
-        // `@differentiable` attribute parameter indices.
+      // If no exact witness derivative configuration was found, check
+      // conditions for creating an implicit witness `@differentiable` attribute
+      // with the exact derivative configuration:
+      // - If the witness has a "superset" derivative configuration.
+      // - If the witness is less than public and is declared in the same file
+      //   as the conformance.
+      //   - `@differentiable` attributes are really only significant for public
+      //     declarations: it improves usability to not require explicit
+      //     `@differentiable` attributes for less-visible declarations.
+      bool createImplicitWitnessAttribute =
+          supersetConfig || witness->getFormalAccess() < AccessLevel::Public;
+      // If the witness has less-than-public visibility and is declared in a
+      // different file than the conformance, produce an error.
+      if (!supersetConfig && witness->getFormalAccess() < AccessLevel::Public &&
+          dc->getModuleScopeContext() !=
+              witness->getDeclContext()->getModuleScopeContext()) {
+        // FIXME(TF-1014): `@differentiable` attribute diagnostic does not
+        // appear if associated type inference is involved.
+        if (auto *vdWitness = dyn_cast<VarDecl>(witness)) {
+          return RequirementMatch(
+              getStandinForAccessor(vdWitness, AccessorKind::Get),
+              MatchKind::MissingDifferentiableAttr, reqDiffAttr);
+        } else {
+          return RequirementMatch(witness, MatchKind::MissingDifferentiableAttr,
+                                  reqDiffAttr);
+        }
+      }
+      if (createImplicitWitnessAttribute) {
+        auto derivativeGenSig = witnessAFD->getGenericSignature();
+        if (supersetConfig)
+          derivativeGenSig = supersetConfig->derivativeGenericSignature;
+        // Use source location of the witness declaration as the source location
+        // of the implicit `@differentiable` attribute.
         auto *newAttr = DifferentiableAttr::create(
-            witnessAFD, /*implicit*/ true, reqDiffAttr->AtLoc,
-            reqDiffAttr->getRange(), reqDiffAttr->isLinear(),
-            reqDiffAttr->getParameterIndices(),
-            supersetConfig->derivativeGenericSignature);
+            witnessAFD, /*implicit*/ true, witness->getLoc(), witness->getLoc(),
+            reqDiffAttr->isLinear(), reqDiffAttr->getParameterIndices(),
+            derivativeGenSig);
+        // If the implicit attribute is inherited from a protocol requirement's
+        // attribute, store the protocol requirement attribute's location for
+        // use in diagnostics.
+        if (witness->getFormalAccess() < AccessLevel::Public) {
+          newAttr->getImplicitlyInheritedDifferentiableAttrLocation(
+              reqDiffAttr->getLocation());
+        }
         auto insertion = ctx.DifferentiableAttrs.try_emplace(
             {witnessAFD, newAttr->getParameterIndices()}, newAttr);
         // Valid `@differentiable` attributes are uniqued by original function
@@ -418,9 +454,9 @@ matchWitnessDifferentiableAttr(DeclContext *dc, ValueDecl *req,
         if (auto *vdWitness = dyn_cast<VarDecl>(witness)) {
           return RequirementMatch(
               getStandinForAccessor(vdWitness, AccessorKind::Get),
-              MatchKind::DifferentiableConflict, reqDiffAttr);
+              MatchKind::MissingDifferentiableAttr, reqDiffAttr);
         } else {
-          return RequirementMatch(witness, MatchKind::DifferentiableConflict,
+          return RequirementMatch(witness, MatchKind::MissingDifferentiableAttr,
                                   reqDiffAttr);
         }
       }
@@ -2318,14 +2354,15 @@ diagnoseMatch(ModuleDecl *module, NormalProtocolConformance *conformance,
   case MatchKind::NonObjC:
     diags.diagnose(match.Witness, diag::protocol_witness_not_objc);
     break;
-  case MatchKind::DifferentiableConflict: {
+  case MatchKind::MissingDifferentiableAttr: {
+    auto *witness = match.Witness;
     // Emit a note and fix-it showing the missing requirement `@differentiable`
     // attribute.
     auto *reqAttr = cast<DifferentiableAttr>(match.UnmetAttribute);
     assert(reqAttr);
     // Omit printing `wrt:` clause if attribute's differentiability
     // parameters match inferred differentiability parameters.
-    auto *original = cast<AbstractFunctionDecl>(match.Witness);
+    auto *original = cast<AbstractFunctionDecl>(witness);
     auto *whereClauseGenEnv =
         reqAttr->getDerivativeGenericEnvironment(original);
     auto *inferredParameters = TypeChecker::inferDifferentiabilityParameters(
@@ -2336,11 +2373,29 @@ diagnoseMatch(ModuleDecl *module, NormalProtocolConformance *conformance,
     llvm::raw_string_ostream os(reqDiffAttrString);
     reqAttr->print(os, req, omitWrtClause);
     os.flush();
-    diags
-        .diagnose(match.Witness,
-                  diag::protocol_witness_missing_differentiable_attr,
-                  reqDiffAttrString)
-        .fixItInsert(match.Witness->getStartLoc(), reqDiffAttrString + ' ');
+    // If the witness has less-than-public visibility and is declared in a
+    // different file than the conformance, emit a specialized diagnostic.
+    if (witness->getFormalAccess() < AccessLevel::Public &&
+        conformance->getDeclContext()->getModuleScopeContext() !=
+            witness->getDeclContext()->getModuleScopeContext()) {
+      diags
+          .diagnose(
+              witness,
+              diag::
+                  protocol_witness_missing_differentiable_attr_nonpublic_other_file,
+              reqDiffAttrString, witness->getDescriptiveKind(),
+              witness->getFullName(), req->getDescriptiveKind(),
+              req->getFullName(), conformance->getType(),
+              conformance->getProtocol()->getDeclaredInterfaceType())
+          .fixItInsert(match.Witness->getStartLoc(), reqDiffAttrString + ' ');
+    }
+    // Otherwise, emit a general "missing attribute" diagnostic.
+    else {
+      diags
+          .diagnose(witness, diag::protocol_witness_missing_differentiable_attr,
+                    reqDiffAttrString)
+          .fixItInsert(witness->getStartLoc(), reqDiffAttrString + ' ');
+    }
     break;
   }
   }

lib/Sema/TypeCheckProtocol.h

@@ -209,9 +209,8 @@ enum class MatchKind : uint8_t {
   /// The witness is explicitly @nonobjc but the requirement is @objc.
   NonObjC,
 
-  /// The witness does not have a `@differentiable` attribute satisfying one
-  /// from the requirement.
-  DifferentiableConflict,
+  /// The witness is missing a `@differentiable` attribute from the requirement.
+  MissingDifferentiableAttr,
 };
 
 /// Describes the kind of optional adjustment performed when
@@ -362,7 +361,7 @@ struct RequirementMatch {
       : Witness(witness), Kind(kind), WitnessType(), UnmetAttribute(attr),
         ReqEnv(None) {
     assert(!hasWitnessType() && "Should have witness type");
-    assert(UnmetAttribute);
+    assert(hasUnmetAttribute() && "Should have unmet attribute");
   }
 
   RequirementMatch(ValueDecl *witness, MatchKind kind,
@@ -437,7 +436,7 @@ struct RequirementMatch {
     case MatchKind::RethrowsConflict:
     case MatchKind::ThrowsConflict:
     case MatchKind::NonObjC:
-    case MatchKind::DifferentiableConflict:
+    case MatchKind::MissingDifferentiableAttr:
       return false;
     }
 
@@ -467,7 +466,7 @@ struct RequirementMatch {
     case MatchKind::RethrowsConflict:
     case MatchKind::ThrowsConflict:
     case MatchKind::NonObjC:
-    case MatchKind::DifferentiableConflict:
+    case MatchKind::MissingDifferentiableAttr:
       return false;
     }
 
@@ -478,7 +477,9 @@ struct RequirementMatch {
   bool hasRequirement() { return Kind == MatchKind::MissingRequirement; }
 
   /// Determine whether this requirement match has an unmet attribute.
-  bool hasUnmetAttribute() { return Kind == MatchKind::DifferentiableConflict; }
+  bool hasUnmetAttribute() {
+    return Kind == MatchKind::MissingDifferentiableAttr;
+  }
 
   swift::Witness getWitness(ASTContext &ctx) const;
 };