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

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.

Resolves TF-691.

Author: dan-zheng

lib/SILOptimizer/Mandatory/Differentiation.cpp

@@ -241,6 +241,19 @@ getAssociatedFunctionGenericSignature(SILDifferentiableAttr *attr,
       GenericSignatureBuilder::FloatingRequirementSource::forAbstract();
   for (auto &req : attr->getRequirements())
     builder.addRequirement(req, source, original->getModule().getSwiftModule());
+  // Constrain all wrt parameters to conform to `Differentiable`.
+  auto &ctx = original->getASTContext();
+  auto *diffableProto = ctx.getProtocol(KnownProtocolKind::Differentiable);
+  auto paramIndexSet = attr->getIndices().parameters;
+  for (unsigned paramIdx : paramIndexSet->getIndices()) {
+    if (!paramIndexSet->contains(paramIdx))
+      continue;
+    auto paramType =
+        original->getConventions().getSILArgumentType(paramIdx).getASTType();
+    Requirement req(RequirementKind::Conformance, paramType,
+                    diffableProto->getDeclaredType());
+    builder.addRequirement(req, source, original->getModule().getSwiftModule());
+  }
   return std::move(builder)
       .computeGenericSignature(SourceLoc(), /*allowConcreteGenericParams=*/true)
       ->getCanonicalSignature();
@@ -2874,11 +2887,17 @@ class VJPEmitter final
     auto origTy = original->getLoweredFunctionType();
     auto lookupConformance = LookUpConformanceInModule(module.getSwiftModule());
 
+    auto pbGenericSig = getAssociatedFunctionGenericSignature(attr, original);
+
     // RAII that pushes the original function's generic signature to
     // `module.Types` so that the calls `module.Types.getTypeLowering()` below
-    // will know the original function's generic parameter types.
+    // will know the pullback's generic parameter types.
     Lowering::GenericContextScope genericContextScope(
-        module.Types, origTy->getGenericSignature());
+        module.Types, pbGenericSig);
+
+    auto *pbGenericEnv = pbGenericSig
+        ? pbGenericSig->createGenericEnvironment()
+        : nullptr;
 
     // Given a type, returns its formal SIL parameter info.
     auto getTangentParameterInfoForOriginalResult = [&](
@@ -2970,10 +2989,6 @@ class VJPEmitter final
         mangler.mangleAutoDiffLinearMapHelper(
             original->getName(), AutoDiffLinearMapKind::Pullback,
             indices)).str();
-    auto pbGenericSig = getAssociatedFunctionGenericSignature(attr, original);
-    auto *pbGenericEnv = pbGenericSig
-        ? pbGenericSig->createGenericEnvironment()
-        : nullptr;
     auto pbType = SILFunctionType::get(
         pbGenericSig, origTy->getExtInfo(), origTy->getCoroutineKind(),
         origTy->getCalleeConvention(), pbParams, {}, adjResults, None,
@@ -3296,7 +3311,7 @@ class VJPEmitter final
     auto original = getOpValue(ai->getCallee());
     auto functionSource = original;
     SILValue vjpValue;
-    // If functionSource is a @differentiable function, just extract it.
+    // If `functionSource` is a `@differentiable` function, just extract it.
     auto originalFnTy = original->getType().castTo<SILFunctionType>();
     if (originalFnTy->isDifferentiable()) {
       auto paramIndices = originalFnTy->getDifferentiationParameterIndices();
@@ -3531,12 +3546,6 @@ class JVPEmitter final
     auto origTy = original->getLoweredFunctionType();
     auto lookupConformance = LookUpConformanceInModule(module.getSwiftModule());
 
-    // RAII that pushes the original function's generic signature to
-    // `module.Types` so that the calls `module.Types.getTypeLowering()` below
-    // will know the original function's generic parameter types.
-    Lowering::GenericContextScope genericContextScope(
-        module.Types, origTy->getGenericSignature());
-
     SmallVector<SILParameterInfo, 8> diffParams;
     SmallVector<SILResultInfo, 8> diffResults;
     auto origParams = origTy->getParameters();
@@ -3563,6 +3572,13 @@ class JVPEmitter final
             original->getName(), AutoDiffLinearMapKind::Differential,
             indices)).str();
     auto diffGenericSig = getAssociatedFunctionGenericSignature(attr, original);
+
+    // RAII that pushes the original function's generic signature to
+    // `module.Types` so that the calls `module.Types.getTypeLowering()` below
+    // will know the differential's generic parameter types.
+    Lowering::GenericContextScope genericContextScope(
+        module.Types, diffGenericSig);
+
     auto *diffGenericEnv = diffGenericSig
         ? diffGenericSig->createGenericEnvironment()
         : nullptr;
@@ -5985,7 +6001,7 @@ static SILFunction *createEmptyJVP(
 
   // RAII that pushes the original function's generic signature to
   // `module.Types` so that the calls `module.Types.getTypeLowering()` below
-  // will know the VJP's generic parameter types.
+  // will know the JVP's generic parameter types.
   Lowering::GenericContextScope genericContextScope(
       module.Types, jvpGenericSig);
 

test/AutoDiff/autodiff_indirect_diagnostics.swift

@@ -12,7 +12,9 @@ func generic<T: Differentiable & FloatingPoint>(_ x: T) -> T {
 }
 _ = gradient(at: 1.0, in: generic)
 
-// Test unmet generic requirements.
+//===----------------------------------------------------------------------===//
+// Unmet generic requirements
+//===----------------------------------------------------------------------===//
 
 @differentiable(
   vjp: vjpWeirdExtraRequirements
@@ -67,3 +69,22 @@ struct TF8Struct<Scalar> : TF8Proto where Scalar : FloatingPoint & Differentiabl
 }
 
 _ = gradient(at: 1.0, in: { x in x.squareRoot() })
+
+//===----------------------------------------------------------------------===//
+// Add `Differentiable` conformance for generic wrt parameters
+//===----------------------------------------------------------------------===//
+
+func id<T>(_ x: T) -> T { x }
+let _: @differentiable (Float) -> Float = { x in id(x) }
+
+struct TF_691<Scalar> {
+  var x: Scalar
+  init(_ x: Scalar) {
+    self.x = x
+  }
+}
+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)) }