fix tf-628 (#25924)

Author: marcrasi

include/swift/SIL/SILWitnessVisitor.h

@@ -120,14 +120,16 @@ template <class T> class SILWitnessVisitor : public ASTVisitor<T> {
 
   void visitAbstractStorageDecl(AbstractStorageDecl *sd) {
     sd->visitOpaqueAccessors([&](AccessorDecl *accessor) {
-      if (SILDeclRef::requiresNewWitnessTableEntry(accessor))
-        asDerived().addMethod(SILDeclRef(accessor, SILDeclRef::Kind::Func));
+      // SWIFT_ENABLE_TENSORFLOW
+      addMethodAndAutoDiffAssociatedMethodsIfRequired(accessor,
+                                                      SILDeclRef::Kind::Func);
     });
   }
 
   void visitConstructorDecl(ConstructorDecl *cd) {
-    if (SILDeclRef::requiresNewWitnessTableEntry(cd))
-      asDerived().addMethod(SILDeclRef(cd, SILDeclRef::Kind::Allocator));
+    // SWIFT_ENABLE_TENSORFLOW
+    addMethodAndAutoDiffAssociatedMethodsIfRequired(
+        cd, SILDeclRef::Kind::Allocator);
   }
 
   void visitAccessorDecl(AccessorDecl *func) {
@@ -137,22 +139,8 @@ template <class T> class SILWitnessVisitor : public ASTVisitor<T> {
   void visitFuncDecl(FuncDecl *func) {
     assert(!isa<AccessorDecl>(func));
     // SWIFT_ENABLE_TENSORFLOW
-    if (!SILDeclRef::requiresNewWitnessTableEntry(func))
-      return;
-
-    auto funcDeclRef = SILDeclRef(func, SILDeclRef::Kind::Func);
-    asDerived().addMethod(funcDeclRef);
-
-    if (auto *DA = func->getAttrs().getAttribute<DifferentiableAttr>()) {
-      asDerived().addMethod(funcDeclRef.asAutoDiffAssociatedFunction(
-          AutoDiffAssociatedFunctionIdentifier::get(
-              AutoDiffAssociatedFunctionKind::JVP, /*differentiationOrder*/ 1,
-              DA->getParameterIndices(), func->getASTContext())));
-      asDerived().addMethod(funcDeclRef.asAutoDiffAssociatedFunction(
-          AutoDiffAssociatedFunctionIdentifier::get(
-              AutoDiffAssociatedFunctionKind::VJP, /*differentiationOrder*/ 1,
-              DA->getParameterIndices(), func->getASTContext())));
-    }
+    addMethodAndAutoDiffAssociatedMethodsIfRequired(func,
+                                                    SILDeclRef::Kind::Func);
   }
 
   void visitMissingMemberDecl(MissingMemberDecl *placeholder) {
@@ -179,6 +167,28 @@ template <class T> class SILWitnessVisitor : public ASTVisitor<T> {
   void visitPoundDiagnosticDecl(PoundDiagnosticDecl *pdd) {
     // We don't care about diagnostics at this stage.
   }
+
+// SWIFT_ENABLE_TENSORFLOW
+private:
+  void addMethodAndAutoDiffAssociatedMethodsIfRequired(
+      AbstractFunctionDecl *func, SILDeclRef::Kind kind) {
+    if (!SILDeclRef::requiresNewWitnessTableEntry(func))
+      return;
+
+    auto funcDeclRef = SILDeclRef(func, kind);
+    asDerived().addMethod(funcDeclRef);
+
+    if (auto *DA = func->getAttrs().getAttribute<DifferentiableAttr>()) {
+      asDerived().addMethod(funcDeclRef.asAutoDiffAssociatedFunction(
+          AutoDiffAssociatedFunctionIdentifier::get(
+              AutoDiffAssociatedFunctionKind::JVP, /*differentiationOrder*/ 1,
+              DA->getParameterIndices(), func->getASTContext())));
+      asDerived().addMethod(funcDeclRef.asAutoDiffAssociatedFunction(
+          AutoDiffAssociatedFunctionIdentifier::get(
+              AutoDiffAssociatedFunctionKind::VJP, /*differentiationOrder*/ 1,
+              DA->getParameterIndices(), func->getASTContext())));
+    }
+  }
 };
 
 } // end namespace swift

test/AutoDiff/protocol_requirement_autodiff.swift

@@ -4,12 +4,15 @@ import StdlibUnittest
 
 var ProtocolRequirementAutodiffTests = TestSuite("ProtocolRequirementAutodiff")
 
+// MARK: - Func requirements.
+
 protocol DiffReq : Differentiable {
   @differentiable(wrt: (self, x))
   func f(_ x: Float) -> Float
 }
 
 extension DiffReq where TangentVector : AdditiveArithmetic {
+  @inline(never)  // Prevent specialization, to test all witness code.
   func gradF(at x: Float) -> (Self.TangentVector, Float) {
     return (valueWithPullback(at: x) { s, x in s.f(x) }).1(1)
   }
@@ -53,7 +56,76 @@ extension Quadratic : VectorProtocol {
   }
 }
 
-// Test witness method SIL type computation.
+ProtocolRequirementAutodiffTests.test("func") {
+  expectEqual((Quadratic(0, 0, 1), 12), Quadratic(11, 12, 13).gradF(at: 0))
+  expectEqual((Quadratic(1, 1, 1), 2 * 11 + 12),
+              Quadratic(11, 12, 13).gradF(at: 1))
+  expectEqual((Quadratic(4, 2, 1), 2 * 11 * 2 + 12),
+              Quadratic(11, 12, 13).gradF(at: 2))
+}
+
+// MARK: Constructor, accessor, and subscript requirements.
+
+protocol FunctionsOfX: Differentiable {
+  @differentiable
+  init(x: Float)
+
+  @differentiable
+  var x: Float { get }
+
+  @differentiable
+  var y: Float { get }
+
+  @differentiable
+  var z: Float { get }
+
+  @differentiable
+  subscript() -> Float { get }
+}
+
+struct TestFunctionsOfX: FunctionsOfX {
+  @differentiable
+  init(x: Float) {
+    self.x = x
+    self.y = x * x
+  }
+
+  /// x = x
+  var x: Float
+
+  /// y = x * x
+  var y: Float
+
+  /// z = x * x + x
+  var z: Float {
+    return y + x
+  }
+
+  @differentiable
+  subscript() -> Float {
+    return z
+  }
+}
+
+@inline(never)  // Prevent specialization, to test all witness code.
+func derivatives<F: FunctionsOfX>(at x: Float, in: F.Type)
+  -> (Float, Float, Float, Float)
+{
+  let dxdx = gradient(at: x) { x in F(x: x).x }
+  let dydx = gradient(at: x) { x in F(x: x).y }
+  let dzdx = gradient(at: x) { x in F(x: x).z }
+  let dsubscriptdx = gradient(at: x) { x in F(x: x)[] }
+  return (dxdx, dydx, dzdx, dsubscriptdx)
+}
+
+ProtocolRequirementAutodiffTests.test("constructor, accessor, subscript") {
+  expectEqual(
+    derivatives(at: 2.0, in: TestFunctionsOfX.self),
+    (1.0, 4.0, 5.0, 5.0))
+}
+
+// MARK: - Test witness method SIL type computation.
+
 protocol P : Differentiable {
   @differentiable(wrt: (x, y))
   func foo(_ x: Float, _ y: Double) -> Float
@@ -65,12 +137,4 @@ struct S : P {
   }
 }
 
-ProtocolRequirementAutodiffTests.test("Trivial") {
-  expectEqual((Quadratic(0, 0, 1), 12), Quadratic(11, 12, 13).gradF(at: 0))
-  expectEqual((Quadratic(1, 1, 1), 2 * 11 + 12),
-              Quadratic(11, 12, 13).gradF(at: 1))
-  expectEqual((Quadratic(4, 2, 1), 2 * 11 * 2 + 12),
-              Quadratic(11, 12, 13).gradF(at: 2))
-}
-
 runAllTests()