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

This reverts commit 6f58fd4.

#26406 introduced TF-697;
reverting it until the issue is fixed.

Add testcase for TF-697 to prevent future regressions.

Author: dan-zheng

lib/SILOptimizer/Mandatory/Differentiation.cpp

@@ -217,19 +217,6 @@ 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();
@@ -2863,17 +2850,11 @@ 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 pullback's generic parameter types.
+    // will know the original function's generic parameter types.
     Lowering::GenericContextScope genericContextScope(
-        module.Types, pbGenericSig);
-
-    auto *pbGenericEnv = pbGenericSig
-        ? pbGenericSig->createGenericEnvironment()
-        : nullptr;
+        module.Types, origTy->getGenericSignature());
 
     // Given a type, returns its formal SIL parameter info.
     auto getTangentParameterInfoForOriginalResult = [&](
@@ -2965,6 +2946,10 @@ 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,
@@ -3286,7 +3271,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();
@@ -3536,6 +3521,12 @@ 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();
@@ -3562,13 +3553,6 @@ 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;
@@ -5991,7 +5975,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 JVP's generic parameter types.
+  // will know the VJP's generic parameter types.
   Lowering::GenericContextScope genericContextScope(
       module.Types, jvpGenericSig);
 

test/AutoDiff/autodiff_indirect_diagnostics.swift

@@ -12,9 +12,7 @@ func generic<T: Differentiable & FloatingPoint>(_ x: T) -> T {
 }
 _ = gradient(at: 1.0, in: generic)
 
-//===----------------------------------------------------------------------===//
-// Unmet generic requirements
-//===----------------------------------------------------------------------===//
+// Test unmet generic requirements.
 
 @differentiable(
   vjp: vjpWeirdExtraRequirements
@@ -91,6 +89,10 @@ 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) }
 
@@ -104,4 +106,5 @@ 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)) }
+let _: @differentiable (Float) -> TF_691<Float> = { x in id(TF_691(x)) }
+*/

test/AutoDiff/generics.swift

@@ -247,4 +247,33 @@ public func TF_688<Scalar: Differentiable>(
   reduction(x)
 }
 
+// TF-697: Test generic requirements of generated AD associated function.
+protocol TF_697_Module: Differentiable where AllDifferentiableVariables == TangentVector {
+    associatedtype Input
+    associatedtype Output: Differentiable
+
+    @differentiable(wrt: self)
+    func callModule(_ input: Input) -> Output
+}
+protocol TF_697_Layer: TF_697_Module where Input: Differentiable {
+    @differentiable
+    func callLayer(_ input: Input) -> Output
+}
+struct TF_697_Sequential<Layer1: TF_697_Module, Layer2: TF_697_Layer>: TF_697_Module
+    where Layer1.Output == Layer2.Input {
+    var layer1: Layer1
+    var layer2: Layer2
+
+    @differentiable(wrt: self)
+    func callModule(_ input: Layer1.Input) -> Layer2.Output {
+        layer2.callLayer(layer1.callModule(input))
+    }
+}
+extension TF_697_Sequential: TF_697_Layer where Layer1: TF_697_Layer {
+    @differentiable
+    func callLayer(_ input: Layer1.Input) -> Layer2.Output {
+        layer2.callLayer(layer1.callLayer(input))
+    }
+}
+
 // TODO: add more tests.