[AutoDiff] Plumb witness derivative generic signatures through SILGen.

When SILGenWitnessTable creates a decl ref for the witness of a derivative function requirement, it is using the requirement's derivative generic signature in the resulting witness decl ref. This is wrong because the witness may have a different derivative generic signature than the requirement, leading to a crash. This bug was never discovered because GSB's dark magic made it "just work", until requirement machine.

The fix is to store the matched witness derivative generic signature in `Witness` during type checking, and during witness table generation, use the witness' generic signature to create a witness decl ref.

Resolves rdar://84716758, rdar://84213107 and rdar://84987079.

Author: rxwei

include/swift/AST/Witness.h

@@ -94,6 +94,9 @@ class Witness {
     ConcreteDeclRef declRef;
     GenericEnvironment *syntheticEnvironment;
     SubstitutionMap reqToSyntheticEnvSubs;
+    /// The derivative generic signature, when the requirement is a derivative
+    /// function.
+    GenericSignature derivativeGenSig;
   };
 
   llvm::PointerUnion<ValueDecl *, StoredWitness *> storage;
@@ -124,7 +127,8 @@ class Witness {
   static Witness forDeserialized(ValueDecl *decl,
                                  SubstitutionMap substitutions) {
     // TODO: It's probably a good idea to have a separate 'deserialized' bit.
-    return Witness(decl, substitutions, nullptr, SubstitutionMap());
+    return Witness(
+        decl, substitutions, nullptr, SubstitutionMap(), CanGenericSignature());
   }
 
   /// Create a witness that requires substitutions.
@@ -138,10 +142,14 @@ class Witness {
   ///
   /// \param reqToSyntheticEnvSubs The mapping from the interface types of the
   /// requirement into the interface types of the synthetic environment.
+  ///
+  /// \param derivativeGenSig The derivative generic signature, when the
+  /// requirement is a derivative function.
   Witness(ValueDecl *decl,
           SubstitutionMap substitutions,
           GenericEnvironment *syntheticEnv,
-          SubstitutionMap reqToSyntheticEnvSubs);
+          SubstitutionMap reqToSyntheticEnvSubs,
+          GenericSignature derivativeGenSig);
 
   /// Retrieve the witness declaration reference, which includes the
   /// substitutions needed to use the witness from the synthetic environment
@@ -183,6 +191,13 @@ class Witness {
     return {};
   }
 
+  /// Retrieve the derivative generic signature.
+  GenericSignature getDerivativeGenericSignature() const {
+    if (auto *storedWitness = storage.dyn_cast<StoredWitness *>())
+      return storedWitness->derivativeGenSig;
+    return GenericSignature();
+  }
+
   SWIFT_DEBUG_DUMP;
 
   void dump(llvm::raw_ostream &out) const;

lib/AST/ProtocolConformance.cpp

@@ -42,7 +42,8 @@ using namespace swift;
 
 Witness::Witness(ValueDecl *decl, SubstitutionMap substitutions,
                  GenericEnvironment *syntheticEnv,
-                 SubstitutionMap reqToSynthesizedEnvSubs) {
+                 SubstitutionMap reqToSynthesizedEnvSubs,
+                 GenericSignature derivativeGenSig) {
   if (!syntheticEnv && substitutions.empty() &&
       reqToSynthesizedEnvSubs.empty()) {
     storage = decl;
@@ -53,7 +54,8 @@ Witness::Witness(ValueDecl *decl, SubstitutionMap substitutions,
   auto declRef = ConcreteDeclRef(decl, substitutions);
   auto storedMem = ctx.Allocate(sizeof(StoredWitness), alignof(StoredWitness));
   auto stored = new (storedMem) StoredWitness{declRef, syntheticEnv,
-                                              reqToSynthesizedEnvSubs};
+                                              reqToSynthesizedEnvSubs,
+                                              derivativeGenSig};
 
   storage = stored;
 }
@@ -892,7 +894,8 @@ NormalProtocolConformance::getWitnessUncached(ValueDecl *requirement) const {
 }
 
 Witness SelfProtocolConformance::getWitness(ValueDecl *requirement) const {
-  return Witness(requirement, SubstitutionMap(), nullptr, SubstitutionMap());
+  return Witness(requirement, SubstitutionMap(), nullptr, SubstitutionMap(),
+                 GenericSignature());
 }
 
 ConcreteDeclRef

lib/SIL/IR/TypeLowering.cpp

@@ -2631,7 +2631,8 @@ CanAnyFunctionType TypeConverter::makeConstantInterfaceType(SILDeclRef c) {
         makeConstantInterfaceType(c.asAutoDiffOriginalFunction());
     auto *derivativeFnTy = originalFnTy->getAutoDiffDerivativeFunctionType(
         derivativeId->getParameterIndices(), derivativeId->getKind(),
-        LookUpConformanceInModule(&M));
+        LookUpConformanceInModule(&M),
+        derivativeId->getDerivativeGenericSignature());
     return cast<AnyFunctionType>(derivativeFnTy->getCanonicalType());
   }
 

lib/SILGen/SILGenType.cpp

@@ -416,7 +416,7 @@ template<typename T> class SILGenWitnessTable : public SILWitnessVisitor<T> {
     if (!reqAccessor) {
       if (auto witness = asDerived().getWitness(reqDecl)) {
         return addMethodImplementation(
-            requirementRef, requirementRef.withDecl(witness.getDecl()),
+            requirementRef, getWitnessRef(requirementRef, witness),
             witness);
       }
 
@@ -444,7 +444,8 @@ template<typename T> class SILGenWitnessTable : public SILWitnessVisitor<T> {
       witnessStorage->getSynthesizedAccessor(reqAccessor->getAccessorKind());
 
     return addMethodImplementation(
-        requirementRef, requirementRef.withDecl(witnessAccessor), witness);
+        requirementRef, getWitnessRef(requirementRef, witnessAccessor),
+        witness);
   }
 
 private:
@@ -458,6 +459,21 @@ template<typename T> class SILGenWitnessTable : public SILWitnessVisitor<T> {
     asDerived().addMethodImplementation(requirementRef, witnessRef,
                                         isFree, witness);
   }
+
+  SILDeclRef getWitnessRef(SILDeclRef requirementRef, Witness witness) {
+    auto witnessRef = requirementRef.withDecl(witness.getDecl());
+    // If the requirement/witness is a derivative function, we need to
+    // substitute the witness's derivative generic signature in its derivative
+    // function identifier.
+    if (requirementRef.isAutoDiffDerivativeFunction()) {
+      auto *reqrRerivativeId = requirementRef.getDerivativeFunctionIdentifier();
+      auto *witnessDerivativeId = AutoDiffDerivativeFunctionIdentifier::get(
+          reqrRerivativeId->getKind(), reqrRerivativeId->getParameterIndices(),
+          witness.getDerivativeGenericSignature(), witnessRef.getASTContext());
+      witnessRef = witnessRef.asAutoDiffDerivativeFunction(witnessDerivativeId);
+    }
+    return witnessRef;
+  }
 };
 
 static IsSerialized_t isConformanceSerialized(RootProtocolConformance *conf) {

lib/Sema/TypeCheckProtocol.cpp

@@ -98,7 +98,8 @@ struct swift::RequirementCheck {
 swift::Witness RequirementMatch::getWitness(ASTContext &ctx) const {
   auto syntheticEnv = ReqEnv->getSyntheticEnvironment();
   return swift::Witness(this->Witness, WitnessSubstitutions,
-                        syntheticEnv, ReqEnv->getRequirementToSyntheticMap());
+                        syntheticEnv, ReqEnv->getRequirementToSyntheticMap(),
+                        DerivativeGenSig);
 }
 
 AssociatedTypeDecl *
@@ -306,17 +307,16 @@ static ValueDecl *getStandinForAccessor(AbstractStorageDecl *witness,
 /// Given a witness, a requirement, and an existing `RequirementMatch` result,
 /// check if the requirement's `@differentiable` attributes are met by the
 /// witness.
-/// - If requirement's `@differentiable` attributes are met, or if `result` is
-///   not viable, returns `result`.
+/// - If `result` is not viable, do nothing.
+/// - If requirement's `@differentiable` attributes are met, update `result`
+///   with the matched derivative generic signature.
 /// - Otherwise, returns a "missing `@differentiable` attribute"
 ///   `RequirementMatch`.
-// Note: the `result` argument is only necessary for using
-// `RequirementMatch::WitnessSubstitutions`.
-static RequirementMatch
+static void
 matchWitnessDifferentiableAttr(DeclContext *dc, ValueDecl *req,
-                               ValueDecl *witness, RequirementMatch result) {
+                               ValueDecl *witness, RequirementMatch &result) {
   if (!result.isViable())
-    return result;
+    return;
 
   // Get the requirement and witness attributes.
   const auto &reqAttrs = req->getAttrs();
@@ -377,6 +377,8 @@ matchWitnessDifferentiableAttr(DeclContext *dc, ValueDecl *req,
       if (witnessConfig.parameterIndices ==
           reqDiffAttr->getParameterIndices()) {
         foundExactConfig = true;
+        // Store the matched witness derivative generic signature.
+        result.DerivativeGenSig = witnessConfig.derivativeGenericSignature;
         break;
       }
       if (witnessConfig.parameterIndices->isSupersetOf(
@@ -407,12 +409,12 @@ matchWitnessDifferentiableAttr(DeclContext *dc, ValueDecl *req,
         // FIXME(TF-1014): `@differentiable` attribute diagnostic does not
         // appear if associated type inference is involved.
         if (auto *vdWitness = dyn_cast<VarDecl>(witness)) {
-          return RequirementMatch(
+          result = RequirementMatch(
               getStandinForAccessor(vdWitness, AccessorKind::Get),
               MatchKind::MissingDifferentiableAttr, reqDiffAttr);
         } else {
-          return RequirementMatch(witness, MatchKind::MissingDifferentiableAttr,
-                                  reqDiffAttr);
+          result = RequirementMatch(
+              witness, MatchKind::MissingDifferentiableAttr, reqDiffAttr);
         }
       }
 
@@ -461,6 +463,8 @@ matchWitnessDifferentiableAttr(DeclContext *dc, ValueDecl *req,
           witnessAFD->addDerivativeFunctionConfiguration(
               {newAttr->getParameterIndices(), resultIndices,
                newAttr->getDerivativeGenericSignature()});
+          // Store the witness derivative generic signature.
+          result.DerivativeGenSig = newAttr->getDerivativeGenericSignature();
         }
       }
       if (!success) {
@@ -475,17 +479,16 @@ matchWitnessDifferentiableAttr(DeclContext *dc, ValueDecl *req,
         // FIXME(TF-1014): `@differentiable` attribute diagnostic does not
         // appear if associated type inference is involved.
         if (auto *vdWitness = dyn_cast<VarDecl>(witness)) {
-          return RequirementMatch(
+          result = RequirementMatch(
               getStandinForAccessor(vdWitness, AccessorKind::Get),
               MatchKind::MissingDifferentiableAttr, reqDiffAttr);
         } else {
-          return RequirementMatch(witness, MatchKind::MissingDifferentiableAttr,
-                                  reqDiffAttr);
+          result = RequirementMatch(
+              witness, MatchKind::MissingDifferentiableAttr, reqDiffAttr);
         }
       }
     }
   }
-  return result;
 }
 
 /// A property or subscript witness must have the same or fewer
@@ -817,7 +820,7 @@ swift::matchWitness(
   auto result = finalize(anyRenaming, optionalAdjustments);
   // Check if the requirement's `@differentiable` attributes are satisfied by
   // the witness.
-  result = matchWitnessDifferentiableAttr(dc, req, witness, result);
+  matchWitnessDifferentiableAttr(dc, req, witness, result);
   return result;
 }
 

lib/Sema/TypeCheckProtocol.h

@@ -435,23 +435,27 @@ struct RequirementMatch {
   RequirementMatch(ValueDecl *witness, MatchKind kind,
                    Type witnessType,
                    Optional<RequirementEnvironment> env = None,
-                   ArrayRef<OptionalAdjustment> optionalAdjustments = {})
+                   ArrayRef<OptionalAdjustment> optionalAdjustments = {},
+                   GenericSignature derivativeGenSig = GenericSignature())
     : Witness(witness), Kind(kind), WitnessType(witnessType),
       ReqEnv(std::move(env)),
       OptionalAdjustments(optionalAdjustments.begin(),
-                          optionalAdjustments.end())
+                          optionalAdjustments.end()),
+      DerivativeGenSig(derivativeGenSig)
   {
     assert(hasWitnessType() == !witnessType.isNull() &&
            "Should (or should not) have witness type");
   }
 
   RequirementMatch(ValueDecl *witness, MatchKind kind, Requirement requirement,
                    Optional<RequirementEnvironment> env = None,
-                   ArrayRef<OptionalAdjustment> optionalAdjustments = {})
+                   ArrayRef<OptionalAdjustment> optionalAdjustments = {},
+                   GenericSignature derivativeGenSig = GenericSignature())
       : Witness(witness), Kind(kind), WitnessType(requirement.getFirstType()),
         MissingRequirement(requirement), ReqEnv(std::move(env)),
         OptionalAdjustments(optionalAdjustments.begin(),
-                            optionalAdjustments.end()) {
+                            optionalAdjustments.end()),
+        DerivativeGenSig(derivativeGenSig) {
     assert(hasWitnessType() && hasRequirement() &&
            "Should have witness type and requirement");
   }
@@ -481,6 +485,9 @@ struct RequirementMatch {
   /// environment.
   SubstitutionMap WitnessSubstitutions;
 
+  /// The matched derivative generic signature.
+  GenericSignature DerivativeGenSig;
+
   /// Determine whether this match is viable.
   bool isViable() const {
     switch(Kind) {

test/AutoDiff/compiler_crashers_fixed/rdar83894546-missing-generic-requirement.swift

@@ -0,0 +1,44 @@
+// RUN: %target-build-swift %s
+
+import _Differentiation
+
+public protocol Layer: Differentiable {
+  associatedtype Input: Differentiable
+  associatedtype Output: Differentiable
+  @differentiable(reverse, wrt: (self, input))
+  @differentiable(reverse, wrt: input)
+  func callAsFunction(_ input: Input) -> Output
+}
+
+// Test for explicitly declared `@differentiable` attributes.
+public class Function1<Input: Differentiable, Output: Differentiable>: Layer {
+  @noDerivative public let body: @differentiable(reverse) (Input) -> Output
+
+  public init(_ body: @escaping @differentiable(reverse) (Input) -> Output) {
+    self.body = body
+  }
+
+  @differentiable(reverse, wrt: (self, input))
+  @differentiable(reverse, wrt: input)
+  public func callAsFunction(_ input: Input) -> Output {
+    body(input)
+  }
+
+  @differentiable(reverse, wrt: x where T: Differentiable)
+  public func foo<T>(x: T) -> T {
+    x
+  }
+}
+
+// Test for implicitly generated `@differentiable` attributes.
+class Function2<Input: Differentiable, Output: Differentiable>: Layer {
+  @noDerivative let body: @differentiable(reverse) (Input) -> Output
+
+  init(_ body: @escaping @differentiable(reverse) (Input) -> Output) {
+    self.body = body
+  }
+
+  func callAsFunction(_ input: Input) -> Output {
+    body(input)
+  }
+}