Revamp struct extraction differentiation semantics.

- Remove `@_fieldwiseDifferentiable`.
- Remove `StructExtractDifferentiationStrategy`.
- Require `TangentVector` to have a member of the same name.

Author: rxwei

include/swift/AST/Attr.def

@@ -421,10 +421,8 @@ DECL_ATTR(differentiating, Differentiating,
 SIMPLE_DECL_ATTR(compilerEvaluable, CompilerEvaluable,
                  OnAccessor | OnFunc | OnConstructor | OnSubscript,
                  /* Not serialized */ 90)
-SIMPLE_DECL_ATTR(_fieldwiseDifferentiable, FieldwiseDifferentiable,
-                 OnNominalType | UserInaccessible, 91)
 SIMPLE_DECL_ATTR(noDerivative, NoDerivative,
-                 OnVar, 92)
+                 OnVar, 91)
 
 #undef TYPE_ATTR
 #undef DECL_ATTR_ALIAS

include/swift/AST/DiagnosticsSIL.def

@@ -447,6 +447,9 @@ NOTE(autodiff_opaque_function_not_differentiable,none,
      "opaque non-'@differentiable' function is not differentiable", ())
 NOTE(autodiff_property_not_differentiable,none,
      "property is not differentiable", ())
+NOTE(autodiff_stored_property_no_corresponding_tangent,none,
+     "property cannot be differentiated because '%0.TangentVector' does not "
+     "have a member named '%1'", (StringRef, StringRef))
 NOTE(autodiff_value_defined_here,none,
      "value defined here", ())
 NOTE(autodiff_when_differentiating_function_call,none,

include/swift/AST/DiagnosticsSema.def

@@ -2808,11 +2808,6 @@ ERROR(noderivative_only_on_stored_properties_in_differentiable_structs,none,
       "'@noDerivative' is only allowed on stored properties in structure types "
       "that declare a conformance to 'Differentiable'", ())
 
-// @_fieldwiseDifferentiable attribute
-ERROR(fieldwise_differentiable_only_on_differentiable_structs,none,
-      "'@_fieldwiseDifferentiable' is only allowed on structure types that "
-      "conform to 'Differentiable'", ())
-
 //------------------------------------------------------------------------------
 // MARK: Type Check Expressions
 //------------------------------------------------------------------------------

lib/SILOptimizer/Mandatory/Differentiation.cpp

@@ -690,8 +690,6 @@ getOrSynthesizeSingleAssociatedStruct(DerivedConformance &derived,
   auto *structDecl = new (C) StructDecl(SourceLoc(), id, SourceLoc(),
                                         /*Inherited*/ C.AllocateCopy(inherited),
                                         /*GenericParams*/ {}, parentDC);
-  structDecl->getAttrs().add(
-      new (C) FieldwiseDifferentiableAttr(/*implicit*/ true));
   structDecl->setImplicit();
   structDecl->copyFormalAccessFrom(nominal, /*sourceIsParentContext*/ true);
 
@@ -960,12 +958,6 @@ deriveDifferentiable_AssociatedStruct(DerivedConformance &derived,
     if (!getAssociatedType(member, parentDC, id))
       return nullptr;
 
-  // Since associated types will be derived, we make this struct a fieldwise
-  // differentiable type.
-  if (!nominal->getAttrs().hasAttribute<FieldwiseDifferentiableAttr>())
-    nominal->getAttrs().add(
-        new (C) FieldwiseDifferentiableAttr(/*implicit*/ true));
-
   // Prevent re-synthesis during repeated calls.
   // FIXME: Investigate why this is necessary to prevent duplicate synthesis.
   auto lookup = nominal->lookupDirect(id);

lib/Sema/DerivedConformanceDifferentiable.cpp

@@ -135,7 +135,6 @@ class AttributeEarlyChecker : public AttributeVisitor<AttributeEarlyChecker> {
   IGNORED_ATTR(Differentiable)
   IGNORED_ATTR(Differentiating)
   IGNORED_ATTR(CompilerEvaluable)
-  IGNORED_ATTR(FieldwiseDifferentiable)
   IGNORED_ATTR(NoDerivative)
 #undef IGNORED_ATTR
 
@@ -872,7 +871,6 @@ class AttributeChecker : public AttributeVisitor<AttributeChecker> {
   void visitDifferentiableAttr(DifferentiableAttr *attr);
   void visitDifferentiatingAttr(DifferentiatingAttr *attr);
   void visitCompilerEvaluableAttr(CompilerEvaluableAttr *attr);
-  void visitFieldwiseDifferentiableAttr(FieldwiseDifferentiableAttr *attr);
   void visitNoDerivativeAttr(NoDerivativeAttr *attr);
 };
 } // end anonymous namespace
@@ -3576,23 +3574,6 @@ void AttributeChecker::visitCompilerEvaluableAttr(CompilerEvaluableAttr *attr) {
   // TypeChecker::checkFunctionBodyCompilerEvaluable().
 }
 
-// SWIFT_ENABLE_TENSORFLOW
-void AttributeChecker::visitFieldwiseDifferentiableAttr(
-    FieldwiseDifferentiableAttr *attr) {
-  auto *structDecl = dyn_cast<StructDecl>(D);
-  if (!structDecl) {
-    diagnoseAndRemoveAttr(attr,
-        diag::fieldwise_differentiable_only_on_differentiable_structs);
-    return;
-  }
-  if (!conformsToDifferentiableInModule(
-          structDecl->getDeclaredInterfaceType(), D->getModuleContext())) {
-    diagnoseAndRemoveAttr(attr,
-        diag::fieldwise_differentiable_only_on_differentiable_structs);
-    return;
-  }
-}
-
 // SWIFT_ENABLE_TENSORFLOW
 void AttributeChecker::visitNoDerivativeAttr(NoDerivativeAttr *attr) {
   auto *vd = dyn_cast<VarDecl>(D);

lib/Sema/TypeCheckAttr.cpp

@@ -1303,7 +1303,6 @@ namespace  {
     UNINTERESTING_ATTR(Differentiable)
     UNINTERESTING_ATTR(Differentiating)
     UNINTERESTING_ATTR(CompilerEvaluable)
-    UNINTERESTING_ATTR(FieldwiseDifferentiable)
     UNINTERESTING_ATTR(NoDerivative)
 
     // These can't appear on overridable declarations.

lib/Sema/TypeCheckDeclOverride.cpp

@@ -29,15 +29,23 @@ struct S {
 }
 
 extension S : Differentiable, VectorNumeric {
+  struct TangentVector: Differentiable, VectorNumeric {
+    var dp: Float
+  }
+  typealias AllDifferentiableVariables = S
   static var zero: S { return S(p: 0) }
   typealias Scalar = Float
   static func + (lhs: S, rhs: S) -> S { return S(p: lhs.p + rhs.p) }
   static func - (lhs: S, rhs: S) -> S { return S(p: lhs.p - rhs.p) }
   static func * (lhs: Float, rhs: S) -> S { return S(p: lhs * rhs.p) }
 
-  typealias TangentVector = S
+  func moved(along direction: TangentVector) -> S {
+    return S(p: p + direction.dp)
+  }
 }
 
+// expected-error @+2 {{function is not differentiable}}
+// expected-note @+1 {{property cannot be differentiated because 'S.TangentVector' does not have a member named 'p'}}
 _ = gradient(at: S(p: 0)) { s in 2 * s.p }
 
 struct NoDerivativeProperty : Differentiable {

test/AutoDiff/autodiff_diagnostics.swift

@@ -6,11 +6,11 @@ public struct Foo : Differentiable {
   public var a: Float
 }
 
-// CHECK-AST-LABEL: @_fieldwiseDifferentiable public struct Foo : Differentiable {
+// CHECK-AST-LABEL: public struct Foo : Differentiable {
 // CHECK-AST:   @differentiable
 // CHECK-AST:   public var a: Float
 // CHECK-AST:   internal init(a: Float)
-// CHECK-AST:   @_fieldwiseDifferentiable public struct AllDifferentiableVariables
+// CHECK-AST:   public struct AllDifferentiableVariables
 // CHECK-AST:     public typealias AllDifferentiableVariables = Foo.AllDifferentiableVariables
 // CHECK-AST:     public typealias TangentVector = Foo.AllDifferentiableVariables
 // CHECK-AST:   public typealias TangentVector = Foo.AllDifferentiableVariables
@@ -25,7 +25,7 @@ let _: @differentiable (AdditiveTangentIsSelf) -> Float = { x in
   x.a + x.a
 }
 
-// CHECK-AST-LABEL: @_fieldwiseDifferentiable internal struct AdditiveTangentIsSelf : AdditiveArithmetic, Differentiable {
+// CHECK-AST-LABEL: internal struct AdditiveTangentIsSelf : AdditiveArithmetic, Differentiable {
 // CHECK-AST:         internal var a: Float
 // CHECK-AST:         internal init(a: Float)
 // CHECK-AST:         internal typealias TangentVector = AdditiveTangentIsSelf
@@ -36,11 +36,11 @@ struct TestNoDerivative : Differentiable {
   @noDerivative var technicallyDifferentiable: Float
 }
 
-// CHECK-AST-LABEL: @_fieldwiseDifferentiable internal struct TestNoDerivative : Differentiable {
+// CHECK-AST-LABEL: internal struct TestNoDerivative : Differentiable {
 // CHECK-AST:         var w: Float
 // CHECK-AST:         @noDerivative internal var technicallyDifferentiable: Float
 // CHECK-AST:         internal init(w: Float, technicallyDifferentiable: Float)
-// CHECK-AST:         @_fieldwiseDifferentiable internal struct AllDifferentiableVariables : Differentiable, AdditiveArithmetic, VectorNumeric
+// CHECK-AST:         internal struct AllDifferentiableVariables : Differentiable, AdditiveArithmetic, VectorNumeric
 // CHECK-AST:           internal typealias AllDifferentiableVariables = TestNoDerivative.AllDifferentiableVariables
 // CHECK-AST:           internal typealias TangentVector = TestNoDerivative.AllDifferentiableVariables
 // CHECK-AST:         internal typealias TangentVector = TestNoDerivative.AllDifferentiableVariables
@@ -50,11 +50,11 @@ struct TestKeyPathIterable : Differentiable, KeyPathIterable {
   @noDerivative var technicallyDifferentiable: Float
 }
 
-// CHECK-AST-LABEL: @_fieldwiseDifferentiable internal struct TestKeyPathIterable : Differentiable, KeyPathIterable {
+// CHECK-AST-LABEL: internal struct TestKeyPathIterable : Differentiable, KeyPathIterable {
 // CHECK-AST:         var w: Float
 // CHECK-AST:         @noDerivative internal var technicallyDifferentiable: Float
 // CHECK-AST:         internal init(w: Float, technicallyDifferentiable: Float)
-// CHECK-AST:         @_fieldwiseDifferentiable internal struct AllDifferentiableVariables : Differentiable, AdditiveArithmetic, KeyPathIterable, VectorNumeric
+// CHECK-AST:         internal struct AllDifferentiableVariables : Differentiable, AdditiveArithmetic, KeyPathIterable, VectorNumeric
 // CHECK-AST:           internal typealias AllDifferentiableVariables = TestKeyPathIterable.AllDifferentiableVariables
 // CHECK-AST:           internal typealias TangentVector = TestKeyPathIterable.AllDifferentiableVariables
 // CHECK-AST:         internal typealias TangentVector = TestKeyPathIterable.AllDifferentiableVariables
@@ -66,7 +66,7 @@ struct GenericTanMember<T : Differentiable> : Differentiable, AdditiveArithmetic
 // TODO(TF-316): Revisit after `Differentiable` derived conformances behavior is standardized.
 // `AllDifferentiableVariables` and `TangentVector` structs need not both be synthesized.
 
-// CHECK-AST-LABEL: @_fieldwiseDifferentiable internal struct GenericTanMember<T> : Differentiable, AdditiveArithmetic where T : Differentiable
+// CHECK-AST-LABEL: internal struct GenericTanMember<T> : Differentiable, AdditiveArithmetic where T : Differentiable
 // CHECK-AST:   internal var x: T.TangentVector
 // CHECK-AST:   internal init(x: T.TangentVector)
 // CHECK-AST:   internal typealias TangentVector = GenericTanMember<T>
@@ -81,7 +81,7 @@ public struct ConditionallyDifferentiable<T> {
 }
 extension ConditionallyDifferentiable : Differentiable where T : Differentiable {}
 
-// CHECK-AST-LABEL: @_fieldwiseDifferentiable public struct ConditionallyDifferentiable<T> {
+// CHECK-AST-LABEL: public struct ConditionallyDifferentiable<T> {
 // CHECK-AST:         @differentiable(wrt: self where T : Differentiable)
 // CHECK-AST:         public let x: T
 // CHECK-AST:         internal init(x: T)

test/AutoDiff/derived_differentiable_properties.swift

@@ -9,7 +9,7 @@ import StdlibUnittest
 var E2EDifferentiablePropertyTests = TestSuite("E2EDifferentiableProperty")
 
 struct TangentSpace : VectorNumeric {
-  let dx, dy: Float
+  let x, y: Float
 }
 
 extension TangentSpace : Differentiable {
@@ -26,18 +26,14 @@ struct Space {
   }
 
   func vjpX() -> (Float, (Float) -> TangentSpace) {
-    return (x, { v in TangentSpace(dx: v, dy: 0) } )
+    return (x, { v in TangentSpace(x: v, y: 0) } )
   }
 
   private let storedX: Float
-  
+
   @differentiable
   var y: Float
 
-  func vjpY() -> (Float, (Float) -> TangentSpace) {
-    return (y, { v in TangentSpace(dx: 0, dy: v) })
-  }
-
   init(x: Float, y: Float) {
     self.storedX = x
     self.y = y
@@ -47,23 +43,23 @@ struct Space {
 extension Space : Differentiable {
   typealias TangentVector = TangentSpace
   func moved(along: TangentSpace) -> Space {
-    return Space(x: x + along.dx, y: y + along.dy)
+    return Space(x: x + along.x, y: y + along.y)
   }
 }
 
 E2EDifferentiablePropertyTests.test("computed property") {
   let actualGrad = gradient(at: Space(x: 0, y: 0)) { (point: Space) -> Float in
     return 2 * point.x
   }
-  let expectedGrad = TangentSpace(dx: 2, dy: 0)
+  let expectedGrad = TangentSpace(x: 2, y: 0)
   expectEqual(expectedGrad, actualGrad)
 }
 
 E2EDifferentiablePropertyTests.test("stored property") {
   let actualGrad = gradient(at: Space(x: 0, y: 0)) { (point: Space) -> Float in
     return 3 * point.y
   }
-  let expectedGrad = TangentSpace(dx: 0, dy: 3)
+  let expectedGrad = TangentSpace(x: 0, y: 3)
   expectEqual(expectedGrad, actualGrad)
 }
 
@@ -85,7 +81,6 @@ E2EDifferentiablePropertyTests.test("generic stored property") {
   expectEqual(expectedGrad, actualGrad)
 }
 
-@_fieldwiseDifferentiable
 struct ProductSpaceSelfTangent : VectorNumeric {
   let x, y: Float
 }
@@ -110,7 +105,6 @@ extension ProductSpaceOtherTangentTangentSpace : Differentiable {
   typealias TangentVector = ProductSpaceOtherTangentTangentSpace
 }
 
-@_fieldwiseDifferentiable
 struct ProductSpaceOtherTangent {
   let x, y: Float
 }

test/AutoDiff/e2e_differentiable_property.swift

@@ -10,15 +10,13 @@ import Glibc
 
 var SeparateTangentTypeTests = TestSuite("SeparateTangentType")
 
-@_fieldwiseDifferentiable
 struct DifferentiableSubset : Differentiable {
   @differentiable(wrt: self)
   var w: Float
   @differentiable(wrt: self)
   var b: Float
   @noDerivative var flag: Bool
 
-  @_fieldwiseDifferentiable
   struct TangentVector : Differentiable, VectorNumeric {
     typealias TangentVector = DifferentiableSubset.TangentVector
     var w: Float
@@ -41,12 +39,10 @@ SeparateTangentTypeTests.test("Initialization") {
   expectEqual(pb(DifferentiableSubset.TangentVector.zero), DifferentiableSubset.TangentVector.zero)
 }
 
-// FIXME(SR-9602): If `TangentVector` is not marked
-// `@_fieldwiseProductSpace`, call the VJP of the memberwise initializer.
-// SeparateTangentTypeTests.test("SomeArithmetics") {
-//   let x = DifferentiableSubset(w: 0, b: 1, flag: false)
-//   let pb = pullback(at: x) { x in DifferentiableSubset(w: x.w * x.w, b: x.b * x.b, flag: true) }
-//   expectEqual(pb(DifferentiableSubset.TangentVector.zero), DifferentiableSubset.TangentVector.zero)
-// }
+SeparateTangentTypeTests.test("SomeArithmetics") {
+  let x = DifferentiableSubset(w: 0, b: 1, flag: false)
+  let pb = pullback(at: x) { x in DifferentiableSubset(w: x.w * x.w, b: x.b * x.b, flag: true) }
+  expectEqual(pb(DifferentiableSubset.TangentVector.zero), DifferentiableSubset.TangentVector.zero)
+}
 
 runAllTests()

test/AutoDiff/separate_cotangent_type.swift

@@ -11,7 +11,6 @@ protocol Proto : Differentiable {
   func function3(_ x: Float, _ y: Double) -> Double
 }
 
-@_fieldwiseDifferentiable
 struct S : Proto, VectorNumeric {
   static var zero: S { return S(p: 0) }
   typealias Scalar = Float