[AutoDiff] Constrain wrt parameters to conform to Differentiable. (#26426)

Constrain all wrt parameters to conform to `Differentiable` when computing
AD associated function generic signatures.

This fixes crashes when differentiating generic original functions that
do not constrain parameters to be `Differentiable`, e.g. an unconstrained
identity function.

Gardening included:
- Remove unused `isSerialized` flag from
  `SILGenModule::getOrCreateAutoDiffAssociatedFunctionThunk`.
- Rename `whereClauseGenericSignature` in SIL to
  `associatedFunctionGenericSignature`.
  - The generic signature does not necessarily come from the `where`
    clause of a `[differentiable]` attribute.

Resolves TF-691 and TF-697.

Author: dan-zheng

include/swift/AST/Types.h

@@ -4228,11 +4228,11 @@ class SILFunctionType final : public TypeBase, public llvm::FoldingSetNode,
       AutoDiffIndexSubset *parameterIndices, unsigned resultIndex,
       unsigned differentiationOrder, AutoDiffAssociatedFunctionKind kind,
       SILModule &module, LookupConformanceFn lookupConformance,
-      CanGenericSignature whereClauseGenericSignature = nullptr);
+      CanGenericSignature associatedFunctionGenericSignature = nullptr);
 
   /// Returns a bit vector that specifices which parameters you can
   /// differentiate with respect to for this differentiable function type. (e.g.
-  /// which parameters are not @nondiff). The function type must be
+  /// which parameters are not `@nondiff`). The function type must be
   /// differentiable.
   AutoDiffIndexSubset *getDifferentiationParameterIndices();
 

lib/SIL/SILFunctionType.cpp

@@ -22,6 +22,8 @@
 #include "swift/AST/DiagnosticsSIL.h"
 #include "swift/AST/ForeignInfo.h"
 #include "swift/AST/GenericEnvironment.h"
+// SWIFT_ENABLE_TENSORFLOW
+#include "swift/AST/GenericSignatureBuilder.h"
 #include "swift/AST/Module.h"
 #include "swift/AST/ProtocolConformance.h"
 #include "swift/SIL/SILModule.h"
@@ -148,22 +150,68 @@ CanSILFunctionType SILFunctionType::getWithoutDifferentiability() {
                               getOptionalErrorResult(), getASTContext());
 }
 
+// Returns the canonical generic signature for an autodiff associated function
+// given an existing associated function generic signature. All differentiation
+// parameters are constrained to conform to `Differentiable`.
+static CanGenericSignature getAutoDiffAssociatedFunctionGenericSignature(
+    CanGenericSignature assocFnGenSig,
+    ArrayRef<SILParameterInfo> originalParameters,
+    AutoDiffIndexSubset *parameterIndices, SILModule &module) {
+  // If associated function has no
+  if (!assocFnGenSig)
+    return nullptr;
+  auto &ctx = module.getASTContext();
+  GenericSignatureBuilder builder(ctx);
+
+  // Add associated function generic signature.
+  builder.addGenericSignature(assocFnGenSig);
+  // Constrain all wrt parameters to conform to `Differentiable`.
+  auto source =
+      GenericSignatureBuilder::FloatingRequirementSource::forAbstract();
+  auto *diffableProto = ctx.getProtocol(KnownProtocolKind::Differentiable);
+  for (unsigned paramIdx : parameterIndices->getIndices()) {
+    auto paramType = originalParameters[paramIdx].getType();
+    Requirement req(RequirementKind::Conformance, paramType,
+                    diffableProto->getDeclaredType());
+    builder.addRequirement(req, source, module.getSwiftModule());
+  }
+  return std::move(builder)
+      .computeGenericSignature(SourceLoc(), /*allowConcreteGenericParams*/ true)
+      ->getCanonicalSignature();
+}
+
 CanSILFunctionType SILFunctionType::getAutoDiffAssociatedFunctionType(
     AutoDiffIndexSubset *parameterIndices, unsigned resultIndex,
     unsigned differentiationOrder, AutoDiffAssociatedFunctionKind kind,
     SILModule &module, LookupConformanceFn lookupConformance,
-    CanGenericSignature whereClauseGenSig) {
+    CanGenericSignature assocFnGenSig) {
   // JVP: (T...) -> ((R...),
   //                 (T.TangentVector...) -> (R.TangentVector...))
   // VJP: (T...) -> ((R...),
   //                 (R.TangentVector...) -> (T.TangentVector...))
 
   auto &ctx = getASTContext();
   auto &typeConverter = module.Types;
-  if (!whereClauseGenSig)
-    whereClauseGenSig = getGenericSignature();
-  Lowering::GenericContextScope genericContextScope(
-      module.Types, whereClauseGenSig);
+
+  // Helper function testing if we are differentiating wrt this index.
+  auto isWrtIndex = [&](unsigned index) -> bool {
+    return index < parameterIndices->getCapacity() &&
+        parameterIndices->contains(index);
+  };
+
+  // Calculate differentiation parameter infos.
+  SmallVector<SILParameterInfo, 4> wrtParams;
+  for (auto valueAndIndex : enumerate(getParameters()))
+    if (isWrtIndex(valueAndIndex.index()))
+      wrtParams.push_back(valueAndIndex.value());
+
+  // Get the canonical associated function generic signature.
+  if (!assocFnGenSig)
+    assocFnGenSig = getGenericSignature();
+  assocFnGenSig = getAutoDiffAssociatedFunctionGenericSignature(
+      assocFnGenSig, getParameters(), parameterIndices, module);
+  Lowering::GenericContextScope genericContextScope(module.Types,
+                                                    assocFnGenSig);
 
   // Given a type, returns its formal SIL parameter info.
   auto getTangentParameterInfoForOriginalResult = [&](
@@ -214,18 +262,6 @@ CanSILFunctionType SILFunctionType::getAutoDiffAssociatedFunctionType(
     return {tanType, conv};
   };
 
-  // Helper function testing if we are differentiating wrt this index.
-  auto isWrtIndex = [&](unsigned index) -> bool {
-    return index < parameterIndices->getCapacity() &&
-        parameterIndices->contains(index);
-  };
-
-  // Calculate differentiation parameter infos.
-  SmallVector<SILParameterInfo, 4> wrtParams;
-  for (auto valueAndIndex : enumerate(getParameters()))
-    if (isWrtIndex(valueAndIndex.index()))
-      wrtParams.push_back(valueAndIndex.value());
-
   CanSILFunctionType closureType;
   switch (kind) {
   case AutoDiffAssociatedFunctionKind::JVP: {
@@ -280,12 +316,12 @@ CanSILFunctionType SILFunctionType::getAutoDiffAssociatedFunctionType(
   newResults.reserve(getNumResults() + 1);
   for (auto &result : getResults()) {
     auto mappedResult = result.getWithType(
-        result.getType()->getCanonicalType(whereClauseGenSig));
+        result.getType()->getCanonicalType(assocFnGenSig));
     newResults.push_back(mappedResult);
   }
-  newResults.push_back({closureType->getCanonicalType(whereClauseGenSig),
+  newResults.push_back({closureType->getCanonicalType(assocFnGenSig),
                         ResultConvention::Owned});
-  return SILFunctionType::get(whereClauseGenSig, getExtInfo(),
+  return SILFunctionType::get(assocFnGenSig, getExtInfo(),
                               getCoroutineKind(), getCalleeConvention(),
                               getParameters(), getYields(), newResults,
                               getOptionalErrorResult(), ctx,

lib/SILGen/SILGen.cpp

@@ -816,8 +816,7 @@ void SILGenModule::postEmitFunction(SILDeclRef constant,
         auto *jvpFn = getFunction(SILDeclRef(jvpDecl), NotForDefinition);
         if (jvpFn->getLoweredFunctionType() != expectedJVPType) {
           jvpThunk = getOrCreateAutoDiffAssociatedFunctionThunk(
-              F, indices, jvpFn, AutoDiffAssociatedFunctionKind::JVP,
-              jvpFn->isSerialized());
+              F, indices, jvpFn, AutoDiffAssociatedFunctionKind::JVP);
         } else {
           auto *id = AutoDiffAssociatedFunctionIdentifier::get(
               AutoDiffAssociatedFunctionKind::JVP, /*differentiationOrder*/ 1,
@@ -836,8 +835,7 @@ void SILGenModule::postEmitFunction(SILDeclRef constant,
         auto *vjpFn = getFunction(SILDeclRef(vjpDecl), NotForDefinition);
         if (vjpFn->getLoweredFunctionType() != expectedVJPType) {
           vjpThunk = getOrCreateAutoDiffAssociatedFunctionThunk(
-              F, indices, vjpFn, AutoDiffAssociatedFunctionKind::VJP,
-              vjpFn->isSerialized());
+              F, indices, vjpFn, AutoDiffAssociatedFunctionKind::VJP);
         } else {
           auto *id = AutoDiffAssociatedFunctionIdentifier::get(
               AutoDiffAssociatedFunctionKind::VJP, /*differentiationOrder*/ 1,

lib/SILGen/SILGen.h

@@ -187,8 +187,7 @@ class LLVM_LIBRARY_VISIBILITY SILGenModule : public ASTVisitor<SILGenModule> {
   /// - The last result in the returned pullback.
   SILFunction *getOrCreateAutoDiffAssociatedFunctionThunk(
       SILFunction *original, SILAutoDiffIndices &indices,
-      SILFunction *assocFn, AutoDiffAssociatedFunctionKind assocFnKind,
-      IsSerialized_t isSerialized);
+      SILFunction *assocFn, AutoDiffAssociatedFunctionKind assocFnKind);
 
   /// Determine whether the given class has any instance variables that
   /// need to be destroyed.

lib/SILGen/SILGenPoly.cpp

@@ -3687,8 +3687,7 @@ static void forwardFunctionArgumentsConvertingOwnership(
 SILFunction *
 SILGenModule::getOrCreateAutoDiffAssociatedFunctionThunk(
     SILFunction *original, SILAutoDiffIndices &indices,
-    SILFunction *assocFn, AutoDiffAssociatedFunctionKind assocFnKind,
-    IsSerialized_t isSerialized) {
+    SILFunction *assocFn, AutoDiffAssociatedFunctionKind assocFnKind) {
   auto assocFnType = assocFn->getLoweredFunctionType();
 
   Mangle::ASTMangler mangler;

lib/SILOptimizer/Mandatory/Differentiation.cpp

@@ -199,26 +199,38 @@ static LoadOwnershipQualifier getBufferLOQ(Type type, SILFunction &fn) {
   return LoadOwnershipQualifier::Unqualified;
 }
 
-// Return the expected generic signature for autodiff associated functions given
-// a SILDifferentiableAttr. The expected generic signature is built from the
-// original generic signature and the attribute's requirements.
+// Returns the generic signature for an autodiff associated function given a
+// `SILDifferentiableAttr` and the original function. The associated function's
+// generic signature is built from the original function's generic signature and
+// the attribute's requirements. All differentiation parameters are constrained
+// to conform to `Differentiable`.
 static CanGenericSignature
 getAssociatedFunctionGenericSignature(SILDifferentiableAttr *attr,
                                       SILFunction *original) {
-  auto originalGenSig =
-      original->getLoweredFunctionType()->getGenericSignature();
+  auto originalFnTy = original->getLoweredFunctionType();
+  auto originalGenSig = originalFnTy->getGenericSignature();
   if (!originalGenSig)
     return nullptr;
-  GenericSignatureBuilder builder(original->getASTContext());
+  auto &ctx = original->getASTContext();
+  GenericSignatureBuilder builder(ctx);
   // Add original generic signature.
   builder.addGenericSignature(originalGenSig);
   // Add where clause requirements.
   auto source =
       GenericSignatureBuilder::FloatingRequirementSource::forAbstract();
   for (auto &req : attr->getRequirements())
     builder.addRequirement(req, source, original->getModule().getSwiftModule());
+  // Constrain all wrt parameters to conform to `Differentiable`.
+  auto *diffableProto = ctx.getProtocol(KnownProtocolKind::Differentiable);
+  auto paramIndexSet = attr->getIndices().parameters;
+  for (unsigned paramIdx : paramIndexSet->getIndices()) {
+    auto paramType = originalFnTy->getParameters()[paramIdx].getType();
+    Requirement req(RequirementKind::Conformance, paramType,
+                    diffableProto->getDeclaredType());
+    builder.addRequirement(req, source, original->getModule().getSwiftModule());
+  }
   return std::move(builder)
-      .computeGenericSignature(SourceLoc(), /*allowConcreteGenericParams=*/true)
+      .computeGenericSignature(SourceLoc(), /*allowConcreteGenericParams*/ true)
       ->getCanonicalSignature();
 }
 

test/AutoDiff/autodiff_indirect_diagnostics.swift

@@ -89,10 +89,6 @@ let _: @differentiable (Float) -> TF_687<Any> = { x in TF_687<Any>(x, dummy: x)
 // Add `Differentiable` conformance for generic wrt parameters
 //===----------------------------------------------------------------------===//
 
-// FIXME(TF-697): The tests below were fixed by
-// https://github.com/apple/swift/pull/26406, which was reverted because it
-// introduced TF-697.
-/*
 func id<T>(_ x: T) -> T { x }
 let _: @differentiable (Float) -> Float = { x in id(x) }
 
@@ -107,4 +103,3 @@ extension TF_691: Differentiable where Scalar: Differentiable {}
 func identity<T>(_ x: TF_691<T>) -> TF_691<T> { x }
 let _: @differentiable (Float) -> TF_691<Float> = { x in identity(TF_691(x)) }
 let _: @differentiable (Float) -> TF_691<Float> = { x in id(TF_691(x)) }
-*/