[AutoDiff] Directly SILGen @derivative attributes to diff witnesses.

Previously, `@derivative` attribute type-checking created implicit
`@differentiable` attributes on the original declaration. This was a
longstanding hack powering `@derivative` attribute derivative registration.

swiftlang#28608 made these changes:
- Derivative function configurations (from `@differentiable` and `@derivative`
  attributes) are serialized in modules and are loaded from imported modules.
- The differentiation transform uses these derivative function configurations
  for derivative function lookup instead of `@differentiable` attributes.

Now, `@derivative` attributes are directly lowered to differentiability
witnesses during SILGen, and implicit `@differentiable` attribute generation
is removed.

Type-checking changes:
- "Overlapping" `@differentiable` and `@derivative` attributes (for the same
  original declaration and parameter indices) are now disallowed. They
  semantically conflict because the first "requests derivative generation"
  while the second "registers a derivative".
  - "Overlapping" `@differentiable` and `@derivative` attributes are allowed
    for protocol requirements. Requirement `@differentiable` attributes
    mean "add JVP/VJP witness table entries" - not "request derivative
    generation", because there is no function body.
  - Note that relaxing the "overlapping" condition to consider derivative
    generic signatures is possible after derivative generic signature mangling
    for derivative functions: TF-680.

Resolves TF-835.

Unblocks TF-1021: lifting the "same-file derivative registration only"
limitation in `@derivative` attribute type-checking. This should be possible
without much work, but needs testing!

Exposes TF-1040: `@differentiable` attribute limitations for class methods.
Exposes TF-1041: untested protocol requirement `@differentiable` attribute
type-checking logic.

Author: dan-zheng

include/swift/AST/ASTContext.h

@@ -112,8 +112,10 @@ namespace swift {
   class IndexSubset;
   // SWIFT_ENABLE_TENSORFLOW
   struct AutoDiffConfig;
-  class VectorSpace;
+  struct AutoDiffDerivativeFunctionKind;
+  class DerivativeAttr;
   class DifferentiableAttr;
+  class VectorSpace;
   // SWIFT_ENABLE_TENSORFLOW END
 
   enum class KnownProtocolKind : uint8_t;
@@ -290,11 +292,26 @@ class ASTContext final {
   /// Cache of autodiff-associated vector spaces.
   llvm::DenseMap<Type, Optional<VectorSpace>> AutoDiffVectorSpaces;
 
-  /// Cache of `@differentiable` attributes keyed by parameter indices. This
-  /// helps us diagnose multiple `@differentiable`s that are with respect to the
-  /// same set of parameters.
+  /// Cache of `@differentiable` attributes keyed by parameter indices. Used to
+  /// diagnose duplicate `@differentiable` attributes for the same key.
+  // NOTE(TF-680): relaxing the uniqueness condition to use derivative generic
+  // signature as a key is possible. It requires derivative generic signature
+  // mangling to avoid name collisions for SIL derivative functions with the
+  // same parameter indices but different derivative generic signatures.
   llvm::DenseMap<std::pair<Decl *, IndexSubset *>, DifferentiableAttr *>
       DifferentiableAttrs;
+
+  /// Cache of `@derivative` attributes keyed by parameter indices and
+  /// derivative function kind. Used to diagnose duplicate `@derivative`
+  /// attributes for the same key.
+  // NOTE(TF-680): relaxing the uniqueness condition to use derivative generic
+  // signature as a key is possible. It requires derivative generic signature
+  // mangling to avoid name collisions for SIL derivative functions with the
+  // same parameter indices but different derivative generic signatures.
+  llvm::DenseMap<
+      std::tuple<Decl *, IndexSubset *, AutoDiffDerivativeFunctionKind>,
+      DerivativeAttr *>
+      DerivativeAttrs;
   // SWIFT_ENABLE_TENSORFLOW END
 
 private:

include/swift/AST/Attr.h

@@ -1947,6 +1947,8 @@ class DerivativeAttr final
   unsigned NumParsedParameters = 0;
   /// The differentiation parameters' indices, resolved by the type checker.
   IndexSubset *ParameterIndices = nullptr;
+  /// The derivative function kind (JVP or VJP), resolved by the type checker.
+  Optional<AutoDiffDerivativeFunctionKind> Kind = None;
 
   explicit DerivativeAttr(bool implicit, SourceLoc atLoc, SourceRange baseRange,
                           DeclNameWithLoc original,
@@ -1975,6 +1977,12 @@ class DerivativeAttr final
     OriginalFunction = decl;
   }
 
+  AutoDiffDerivativeFunctionKind getDerivativeKind() const {
+    assert(Kind && "Derivative function kind has not yet been resolved");
+    return *Kind;
+  }
+  void setDerivativeKind(AutoDiffDerivativeFunctionKind kind) { Kind = kind; }
+
   /// The parsed differentiation parameters, i.e. the list of parameters
   /// specified in 'wrt:'.
   ArrayRef<ParsedAutoDiffParameter> getParsedParameters() const {

include/swift/AST/AutoDiff.h

@@ -306,6 +306,7 @@ struct AutoDiffConfig {
 class AutoDiffDerivativeFunctionIdentifier : public llvm::FoldingSetNode {
   const AutoDiffDerivativeFunctionKind kind;
   IndexSubset *const parameterIndices;
+  // TODO(TF-680): Mangle derivative generic signature requirements as well.
 
   AutoDiffDerivativeFunctionIdentifier(
       AutoDiffDerivativeFunctionKind kind, IndexSubset *parameterIndices) :
@@ -508,6 +509,27 @@ template<> struct DenseMapInfo<AutoDiffConfig> {
   }
 };
 
+template<> struct DenseMapInfo<AutoDiffDerivativeFunctionKind> {
+  static AutoDiffDerivativeFunctionKind getEmptyKey() {
+    return static_cast<AutoDiffDerivativeFunctionKind::innerty>(
+        DenseMapInfo<unsigned>::getEmptyKey());
+  }
+
+  static AutoDiffDerivativeFunctionKind getTombstoneKey() {
+    return static_cast<AutoDiffDerivativeFunctionKind::innerty>(
+        DenseMapInfo<unsigned>::getTombstoneKey());
+  }
+
+  static unsigned getHashValue(const AutoDiffDerivativeFunctionKind &Val) {
+    return DenseMapInfo<unsigned>::getHashValue(Val);
+  }
+
+  static bool isEqual(const AutoDiffDerivativeFunctionKind &LHS,
+                      const AutoDiffDerivativeFunctionKind &RHS) {
+    return LHS == RHS;
+  }
+};
+
 template<> struct DenseMapInfo<SILAutoDiffIndices> {
   static SILAutoDiffIndices getEmptyKey() {
     return { DenseMapInfo<unsigned>::getEmptyKey(), nullptr };

lib/SILGen/SILGen.cpp

@@ -778,6 +778,26 @@ void SILGenModule::postEmitFunction(SILDeclRef constant,
                               diffAttr->getDerivativeGenericSignature());
         emitDifferentiabilityWitness(AFD, F, config, jvp, vjp, diffAttr);
       }
+      for (auto *derivAttr : Attrs.getAttributes<DerivativeAttr>()) {
+        SILFunction *jvp = nullptr;
+        SILFunction *vjp = nullptr;
+        switch (derivAttr->getDerivativeKind()) {
+        case AutoDiffDerivativeFunctionKind::JVP:
+          jvp = F;
+          break;
+        case AutoDiffDerivativeFunctionKind::VJP:
+          vjp = F;
+          break;
+        }
+        auto *origAFD = derivAttr->getOriginalFunction();
+        auto *origFn = getFunction(SILDeclRef(origAFD), NotForDefinition);
+        auto derivativeGenSig = AFD->getGenericSignature();
+        auto *resultIndices = IndexSubset::get(getASTContext(), 1, {0});
+        AutoDiffConfig config(derivAttr->getParameterIndices(), resultIndices,
+                              derivativeGenSig);
+        emitDifferentiabilityWitness(origAFD, origFn, config, jvp, vjp,
+                                     derivAttr);
+      }
     };
     if (auto *accessor = dyn_cast<AccessorDecl>(AFD))
       if (accessor->isGetter())
@@ -790,21 +810,22 @@ void SILGenModule::postEmitFunction(SILDeclRef constant,
 void SILGenModule::emitDifferentiabilityWitness(
     AbstractFunctionDecl *originalAFD, SILFunction *originalFunction,
     const AutoDiffConfig &config, SILFunction *jvp, SILFunction *vjp,
-    const DeclAttribute *diffAttr) {
+    const DeclAttribute *attr) {
+  assert(isa<DifferentiableAttr>(attr) || isa<DerivativeAttr>(attr));
   auto *origFnType = originalAFD->getInterfaceType()->castTo<AnyFunctionType>();
   auto origSilFnType = originalFunction->getLoweredFunctionType();
-  auto *loweredParamIndices = autodiff::getLoweredParameterIndices(
-      config.parameterIndices, origFnType);
+  auto *silParamIndices =
+      autodiff::getLoweredParameterIndices(config.parameterIndices, origFnType);
   // NOTE(TF-893): Extending capacity is necessary when `origSilFnType` has
   // parameters corresponding to captured variables. These parameters do not
   // appear in the type of `origFnType`.
   // TODO: If posssible, change `autodiff::getLoweredParameterIndices` to
   // take `CaptureInfo` into account.
-  if (origSilFnType->getNumParameters() > loweredParamIndices->getCapacity())
-    loweredParamIndices = loweredParamIndices->extendingCapacity(
+  if (origSilFnType->getNumParameters() > silParamIndices->getCapacity())
+    silParamIndices = silParamIndices->extendingCapacity(
         getASTContext(), origSilFnType->getNumParameters());
   // TODO(TF-913): Replace usages of `SILAutoDiffIndices` with `AutoDiffConfig`.
-  SILAutoDiffIndices indices(/*source*/ 0, loweredParamIndices);
+  SILAutoDiffIndices indices(/*source*/ 0, silParamIndices);
 
   // Self reordering thunk is necessary if wrt at least two parameters,
   // including self.
@@ -818,14 +839,22 @@ void SILGenModule::emitDifferentiabilityWitness(
   };
   bool reorderSelf = shouldReorderSelf();
 
-  // Create new SIL differentiability witness.
+  // Get or create new SIL differentiability witness.
+  // Differentiability witness already exists when there are two `@derivative`
+  // attributes (JVP and VJP) for the same derivative function configuration.
   // Witness JVP and VJP are set below.
-  auto *diffWitness = SILDifferentiabilityWitness::createDefinition(
-      M, originalFunction->getLinkage(), originalFunction, loweredParamIndices,
-      config.resultIndices, config.derivativeGenericSignature,
-      /*jvp*/ nullptr, /*vjp*/ nullptr,
-      /*isSerialized*/ hasPublicVisibility(originalFunction->getLinkage()),
-      diffAttr);
+  AutoDiffConfig silConfig(silParamIndices, config.resultIndices,
+                           config.derivativeGenericSignature);
+  SILDifferentiabilityWitnessKey key{originalFunction->getName(), silConfig};
+  auto *diffWitness = M.lookUpDifferentiabilityWitness(key);
+  if (!diffWitness) {
+    diffWitness = SILDifferentiabilityWitness::createDefinition(
+        M, originalFunction->getLinkage(), originalFunction,
+        silConfig.parameterIndices, silConfig.resultIndices,
+        config.derivativeGenericSignature, /*jvp*/ nullptr, /*vjp*/ nullptr,
+        /*isSerialized*/ hasPublicVisibility(originalFunction->getLinkage()),
+        attr);
+  }
 
   // Set derivative function in differentiability witness.
   auto setDerivativeInDifferentiabilityWitness =

lib/Sema/TypeCheckAttr.cpp

@@ -3696,6 +3696,7 @@ void AttributeChecker::visitDerivativeAttr(DerivativeAttr *attr) {
     attr->setInvalid();
     return;
   }
+  attr->setDerivativeKind(kind);
   // `value: R` result tuple element must conform to `Differentiable`.
   auto diffableProto = Ctx.getProtocol(KnownProtocolKind::Differentiable);
   auto valueResultType = valueResultElt.getType();
@@ -3918,74 +3919,35 @@ void AttributeChecker::visitDerivativeAttr(DerivativeAttr *attr) {
     return;
   }
 
-  // Try to find a `@differentiable` attribute on the original function with the
-  // same differentiation parameters.
-  DifferentiableAttr *da = nullptr;
-  for (auto *cda : originalAFD->getAttrs().getAttributes<DifferentiableAttr>())
-    if (checkedWrtParamIndices == cda->getParameterIndices())
-      da = const_cast<DifferentiableAttr *>(cda);
-  // If the original function does not have a `@differentiable` attribute with
-  // the same differentiation parameters, create one.
-  // TODO(TF-835): Lower `@derivative` attributes directly to SIL
-  // differentiability witnesses during SILGen instead of generating implicit
-  // `@differentiable` attributes.
-  if (!da) {
-    da = DifferentiableAttr::create(
-        originalAFD, /*implicit*/ true, attr->AtLoc, attr->getRange(),
-        /*linear*/ false, checkedWrtParamIndices, /*jvp*/ None,
-        /*vjp*/ None, derivative->getGenericSignature());
-    switch (kind) {
-    case AutoDiffDerivativeFunctionKind::JVP:
-      da->setJVPFunction(derivative);
-      break;
-    case AutoDiffDerivativeFunctionKind::VJP:
-      da->setVJPFunction(derivative);
-      break;
-    }
-    auto insertion = Ctx.DifferentiableAttrs.try_emplace(
-        {originalAFD, checkedWrtParamIndices}, da);
-    // Valid `@differentiable` attributes are uniqued by their parameter
-    // indices. Reject duplicate attributes for the same decl and parameter
-    // indices pair.
-    if (!insertion.second && insertion.first->getSecond() != da) {
-      diagnoseAndRemoveAttr(da, diag::differentiable_attr_duplicate);
-      diagnose(insertion.first->getSecond()->getLocation(),
-               diag::differentiable_attr_duplicate_note);
-      return;
-    }
-    originalAFD->getAttrs().add(da);
-    return;
+  // Diagnose if there exists a `@differentiable` attribute with the same
+  // parameter indices as the `@derivative` function.
+  // NOTE(TF-680): relaxing this limitation requires derivative generic
+  // signature mangling to avoid name collisions for SIL derivative functions
+  // with the same parameter indices but different derivative generic
+  // signatures.
+  auto *diffAttr =
+      Ctx.DifferentiableAttrs.lookup({originalAFD, checkedWrtParamIndices});
+  bool isOriginalAFDProtocolRequirement =
+      isa<ProtocolDecl>(originalAFD->getDeclContext()) &&
+      originalAFD->isProtocolRequirement();
+  if (diffAttr && !isOriginalAFDProtocolRequirement) {
+    diagnoseAndRemoveAttr(attr,
+                          diag::derivative_attr_original_already_has_derivative,
+                          originalAFD->getFullName());
+    diagnose(diffAttr->getLocation(), diag::differentiable_attr_duplicate_note);
   }
-  // If the original function has a `@differentiable` attribute with the same
-  // differentiation parameters, check if the `@differentiable` attribute
-  // already has a different registered derivative. If so, emit an error on the
-  // `@derivative` attribute. Otherwise, register the derivative in the
-  // `@differentiable` attribute.
-  switch (kind) {
-  case AutoDiffDerivativeFunctionKind::JVP:
-    // If there's a different registered derivative, emit an error.
-    if ((da->getJVP() &&
-         da->getJVP()->Name.getBaseName() != derivative->getBaseName()) ||
-        (da->getJVPFunction() && da->getJVPFunction() != derivative)) {
-      diagnoseAndRemoveAttr(
-          attr, diag::derivative_attr_original_already_has_derivative,
-          originalAFD->getFullName());
-      return;
-    }
-    da->setJVPFunction(derivative);
-    break;
-  case AutoDiffDerivativeFunctionKind::VJP:
-    // If there's a different registered derivative, emit an error.
-    if ((da->getVJP() &&
-         da->getVJP()->Name.getBaseName() != derivative->getBaseName()) ||
-        (da->getVJPFunction() && da->getVJPFunction() != derivative)) {
-      diagnoseAndRemoveAttr(
-          attr, diag::derivative_attr_original_already_has_derivative,
-          originalAFD->getFullName());
-      return;
-    }
-    da->setVJPFunction(derivative);
-    break;
+
+  // Valid `@derivative` attributes are uniqued by original function and
+  // parameter indices. Reject duplicate attributes.
+  auto insertion = Ctx.DerivativeAttrs.try_emplace(
+      {originalAFD, checkedWrtParamIndices, kind}, attr);
+  if (!insertion.second) {
+    diagnoseAndRemoveAttr(attr,
+                          diag::derivative_attr_original_already_has_derivative,
+                          originalAFD->getFullName());
+    diagnose(insertion.first->getSecond()->getLocation(),
+             diag::differentiable_attr_duplicate_note);
+    return;
   }
 
   // Register derivative function configuration.

lib/Sema/TypeCheckProtocol.cpp

@@ -533,55 +533,61 @@ swift::matchWitness(
   // SWIFT_ENABLE_TENSORFLOW
   auto result = finalize(anyRenaming, optionalAdjustments);
   if (result.isViable()) {
-    // '@differentiable' attributes must match completely. If there exists a
-    // '@differentiable' attribute with a superset of the "wrt" parameters of
-    // a requirement, then an '@differentiable' attribute is added
-    // automatically.
+    // For all `@differentiable` attributes of the protocol requirement, check
+    // that the witness has a derivative configuration with exactly the same
+    // parameter indices, or one with "superset" parameter indices. If there
+    // exists a witness derivative configuration with "superset" parameter
+    // indices, create an implicit `@differentiable` attribute for the witness
+    // with the exact parameter indices from the requirement `@differentiable`
+    // attribute.
     ASTContext &ctx = witness->getASTContext();
-    auto witnessDiffAttrs = witnessAttrs
-        .getAttributes<DifferentiableAttr, /*AllowInvalid*/ true>();
+    auto *witnessAFD = dyn_cast<AbstractFunctionDecl>(witness);
+    if (auto *witnessASD = dyn_cast<AbstractStorageDecl>(witness))
+      witnessAFD = witnessASD->getAccessor(AccessorKind::Get);
     for (auto *reqDiffAttr : reqAttrs.getAttributes<DifferentiableAttr>()) {
-      // TODO(TF-482): Also check whether generic requirements are the same.
-      bool reqDiffAttrMatch = llvm::any_of(
-          witnessDiffAttrs, [&](const DifferentiableAttr *witnessDiffAttr) {
-            return witnessDiffAttr->getParameterIndices() &&
-                   reqDiffAttr->getParameterIndices() &&
-                   witnessDiffAttr->parametersMatch(*reqDiffAttr);
-          });
-      bool reqDiffAttrSupersetMatch = llvm::any_of(
-          witnessDiffAttrs, [&](const DifferentiableAttr *witnessDiffAttr) {
-            return witnessDiffAttr->getParameterIndices() &&
-                   reqDiffAttr->getParameterIndices() &&
-                   witnessDiffAttr->getParameterIndices()
-                     ->isSupersetOf(reqDiffAttr->getParameterIndices());
-          });
-      if (!reqDiffAttrMatch) {
-        auto implicitDiffAttr = false;
-        if (reqDiffAttrSupersetMatch) {
-          auto *witnessAFD = cast<AbstractFunctionDecl>(witness);
+      bool foundExactAttr = false;
+      bool foundSupersetAttr = false;
+      for (auto witnessConfig :
+           witnessAFD->getDerivativeFunctionConfigurations()) {
+        if (witnessConfig.parameterIndices ==
+            reqDiffAttr->getParameterIndices())
+          foundExactAttr = true;
+        if (witnessConfig.parameterIndices->isSupersetOf(
+                reqDiffAttr->getParameterIndices()))
+          foundSupersetAttr = true;
+      }
+      if (!foundExactAttr) {
+        bool success = false;
+        if (foundSupersetAttr) {
+          // If the witness has a "superset" derivative configuration, create an
+          // implicit `@differentiable` attribute with the exact requirement
+          // `@differentiable` attribute parameter indices.
+          // TODO(TF-1041): Investigate why this logic is necessary. When
+          // "implicit `@differentiable` attribute" logic is removed, core
+          // stdlib compilation succeeds and AutoDiff tests pass, but TensorFlow
+          // compilation crashes. An AutoDiff reproducer test should be added.
           auto *newAttr = DifferentiableAttr::create(
               witnessAFD, /*implicit*/ true, reqDiffAttr->AtLoc,
               reqDiffAttr->getRange(), reqDiffAttr->isLinear(),
               reqDiffAttr->getParameterIndices(), /*jvp*/ None,
               /*vjp*/ None, reqDiffAttr->getDerivativeGenericSignature());
           auto insertion = ctx.DifferentiableAttrs.try_emplace(
               {witnessAFD, newAttr->getParameterIndices()}, newAttr);
-          // Register derivative function configuration.
-          auto *resultIndices = IndexSubset::get(ctx, 1, {0});
-          witnessAFD->addDerivativeFunctionConfiguration(
-              {newAttr->getParameterIndices(), resultIndices,
-               newAttr->getDerivativeGenericSignature()});
-          // Valid `@differentiable` attributes are uniqued by their parameter
-          // indices. Reject duplicate attributes for the same decl and parameter
-          // indices pair.
+          // Valid `@differentiable` attributes are uniqued by original function
+          // and parameter indices. Reject duplicate attributes.
           if (!insertion.second) {
             newAttr->setInvalid();
           } else {
             witness->getAttrs().add(newAttr);
-            implicitDiffAttr = true;
+            // Register derivative function configuration.
+            auto *resultIndices = IndexSubset::get(ctx, 1, {0});
+            witnessAFD->addDerivativeFunctionConfiguration(
+                {newAttr->getParameterIndices(), resultIndices,
+                 newAttr->getDerivativeGenericSignature()});
+            success = true;
           }
         }
-        if (!implicitDiffAttr) {
+        if (!success) {
           if (auto *vdWitness = dyn_cast<VarDecl>(witness))
             return RequirementMatch(
                 getStandinForAccessor(vdWitness, AccessorKind::Get),