[AutoDiff] TF-508: Temporary Hack Fix (#25115)

[JIRA Ticket TF-508](https://bugs.swift.org/browse/TF-508)

This ticket should not be closed, as this PR is a temporary hack to fix the problem and unblock me on `SIMD` vectors.

The problem has to do with not being able to resolve `Self` for the actual underlying struct type. For example, in the reproducer in the JIRA ticket, it doesn't know about `MyStruct`. So here, we take a look at the first type, which we are assuming is `Self`, but more specfically the struct which lets us successfully check the requirements.

Author: bartchr808

lib/SILOptimizer/Mandatory/Differentiation.cpp

@@ -786,7 +786,7 @@ void DifferentiationInvoker::print(llvm::raw_ostream &os) const {
 // subsitution map and in the given module.
 static bool checkRequirementsSatisfied(
     ArrayRef<Requirement> requirements, SubstitutionMap substMap,
-    ModuleDecl *swiftModule) {
+    SILFunction *original, ModuleDecl *swiftModule) {
   if (requirements.empty())
     return true;
   // Jointly iterate through associated function requirements/conformances.
@@ -831,10 +831,27 @@ static bool checkRequirementsSatisfied(
           continue;
         }
       }
+      Optional<ProtocolConformanceRef> conformance;
+      // FIXME terrible hack: If LHS is a dependent member, try to resolve it
+      // using the `Self` type, assuming the first replacement type is `Self`.
+      if (auto depMemType = firstType->getAs<DependentMemberType>()) {
+        if (original->hasSelfParam() &&
+            !substMap.getReplacementTypes().empty()) {
+          if (auto substType = depMemType->substBaseType(
+                  substMap.getReplacementTypes().front(),
+                  LookUpConformanceInModule(swiftModule))) {
+            firstType = substType;
+          }
+        }
+        conformance = swiftModule->conformsToProtocol(firstType, protocol);
+      }
       // Otherwise, try to look up conformance in substitution maps.
-      auto isConformanceMet = substMap.lookupConformance(
-          firstType->getCanonicalType(), protocol);
-      if (!isConformanceMet)
+      else {
+        conformance = substMap.lookupConformance(
+            firstType->getCanonicalType(), protocol);
+      }
+        
+      if (!conformance)
         unsatisfiedRequirements.push_back(req);
       continue;
     }
@@ -2111,7 +2128,7 @@ emitAssociatedFunctionReference(
     // TODO(TF-482): Change `lookupMinimalDifferentiableAttr`.
     if (!checkRequirementsSatisfied(
             minimalAttr->getRequirements(),
-            substMap, context.getModule().getSwiftModule())) {
+            substMap, originalFn, context.getModule().getSwiftModule())) {
       context.emitNondifferentiabilityError(original, invoker,
           diag::autodiff_function_assoc_func_requirements_unmet);
       return None;

test/AutoDiff/generics.swift

@@ -57,25 +57,75 @@ struct SupervisedTrainer<Model : Layer> {
 }
 
 // Tests TF-440.
-struct TF_440_Input<Input: Differentiable, State: Differentiable>: Differentiable {
-    var input: Input
-    var state: State
+struct TF_440_Input<Input: Differentiable, State: Differentiable>
+  : Differentiable {
+  var input: Input
+  var state: State
 }
 struct TF_440<T : Differentiable> {
-    @differentiable
-    func applied(to input: TF_440_Input<Float, Float>) -> Float {
-        return input.state
-    }
-
-    @differentiable
-    func applied(to input: TF_440_Input<T, Float>) -> Float {
-        return input.state
-    }
-
-    @differentiable
-    func applied(to input: TF_440_Input<T, Float>) -> T {
-        return input.input
-    }
+  @differentiable
+  func applied(to input: TF_440_Input<Float, Float>) -> Float {
+    return input.state
+  }
+
+  @differentiable
+  func applied(to input: TF_440_Input<T, Float>) -> Float {
+    return input.state
+  }
+
+  @differentiable
+  func applied(to input: TF_440_Input<T, Float>) -> T {
+    return input.input
+  }
+}
+
+// Tests TF-508
+protocol TF_508_Protocol {
+  associatedtype Scalar: BinaryFloatingPoint
+}
+
+extension TF_508_Protocol {
+  @differentiable(vjp: _vjpAdd(lhs:rhs:)
+    where Self : Differentiable,
+          Scalar : Differentiable,
+          Self.TangentVector : TF_508_Protocol)
+  static func +(lhs: Self, rhs: Self) -> Self {
+    return lhs
+  }
+  
+  static var zero: Self {
+    fatalError()
+  }
+  
+  static func - (lhs: Self, rhs: Self) -> Self {
+    fatalError()
+  }
+}
+
+extension TF_508_Protocol
+  where Self : Differentiable,
+        Scalar : Differentiable,
+        TangentVector : TF_508_Protocol {
+  static func _vjpAdd(lhs: Self, rhs: Self)
+    -> (Self, (TangentVector) -> (TangentVector, TangentVector)) {
+    return (lhs, { ($0, $0) })
+  }
+}
+
+struct TF_508_Struct<Scalar: BinaryFloatingPoint>
+  : TF_508_Protocol & AdditiveArithmetic {}
+
+extension TF_508_Struct : Differentiable
+  where Scalar : Differentiable {
+  typealias TangentVector = TF_508_Struct
+  typealias AllDifferentiableVariables = TF_508_Struct
+}
+
+let TF_508_inst = TF_508_Struct<Float>()
+func TF_508_func(x: TF_508_Struct<Float>, y: TF_508_Struct<Float>)
+  -> TF_508_Struct<Float> {
+  return x + y
 }
+let TF_508_bp = pullback(at: TF_508_inst, TF_508_inst, in: TF_508_func)
 
 // TODO: add more tests.