Get sum() differentiable.

Author: bartchr808

lib/Sema/TypeCheckAttr.cpp

@@ -2628,7 +2628,8 @@ static FuncDecl *resolveAutoDiffAssociatedFunction(
   auto isABIPublic = [&](AbstractFunctionDecl *func) {
     return func->getFormalAccess() >= AccessLevel::Public ||
            func->getAttrs().hasAttribute<InlinableAttr>() ||
-           func->getAttrs().hasAttribute<UsableFromInlineAttr>();
+           func->getAttrs().hasAttribute<UsableFromInlineAttr>() ||
+           func->getAttrs().hasAttribute<AlwaysEmitIntoClientAttr>();
   };
 
   // If the original function is exported (i.e. it is public or

stdlib/public/core/SIMDVector.swift

@@ -873,14 +873,16 @@ extension SIMD where Scalar : FloatingPoint {
   }
 
   /// Returns the sum of the scalars in the vector.
-  @_alwaysEmitIntoClient
   // SWIFT_ENABLE_TENSORFLOW
-//  @differentiable(vjp: _vjpSum
-//    where Self : Differentiable,
-//          Self.TangentVector : SIMD,
-//          Scalar : BinaryFloatingPoint & Differentiable,
-//          Scalar.TangentVector : BinaryFloatingPoint,
-//          Self.TangentVector == Self)
+  @inlinable
+  // FIXME: TF-545 we want the sum() func to be marked as
+  // `@_alwaysEmitIntoClient` like before when we define the VJP
+  @differentiable(vjp: _vjpSum
+    where Self : Differentiable,
+          Self.TangentVector : SIMD,
+          Scalar : BinaryFloatingPoint & Differentiable,
+          Scalar.TangentVector : BinaryFloatingPoint,
+          Self.TangentVector == Self)
   public func sum() -> Scalar {
     // Implementation note: this eventually be defined to lower to either
     // llvm.experimental.vector.reduce.fadd or an explicit tree-sum. Open-
@@ -1650,17 +1652,17 @@ extension SIMD
   }
 }
 
-//extension SIMD
-//  where Self : Differentiable,
-//        TangentVector : SIMD,
-//        Scalar : BinaryFloatingPoint & Differentiable,
-//        Scalar.TangentVector : BinaryFloatingPoint,
-//        TangentVector == Self {
-//  @usableFromInline
-//  func _vjpSum() -> (Scalar, (Scalar.TangentVector) -> TangentVector) {
-//    return (sum(), { v in Self(repeating: Scalar(v)) })
-//  }
-//}
+extension SIMD
+  where Self : Differentiable,
+        TangentVector : SIMD,
+        Scalar : BinaryFloatingPoint & Differentiable,
+        Scalar.TangentVector : BinaryFloatingPoint,
+        TangentVector == Self {
+  @inlinable
+  func _vjpSum() -> (Scalar, (Scalar.TangentVector) -> TangentVector) {
+    return (sum(), { v in Self(repeating: Scalar(v)) })
+  }
+}
 
 extension SIMD
   where Self : Differentiable,

test/AutoDiff/SIMD.swift

@@ -30,15 +30,15 @@ SIMDTests.test("Negate") {
   expectEqual(-a, bp1(a))
 }
 
-//SIMDTests.test("Sum") {
-//  let a = SIMD4<Float>(1, 2, 3, 4)
-//
-//  let foo1 = { (x: SIMD4<Float>) -> Float in
-//    return x.sum()
-//  }
-//  let bp1 = pullback(at: a, in: foo1)
-//  expectEqual(SIMD4<Float>(3, 3, 3, 3), bp1(3))
-//}
+SIMDTests.test("Sum") {
+  let a = SIMD4<Float>(1, 2, 3, 4)
+
+  let foo1 = { (x: SIMD4<Float>) -> Float in
+    return x.sum()
+  }
+  let bp1 = pullback(at: a, in: foo1)
+  expectEqual(SIMD4<Float>(3, 3, 3, 3), bp1(3))
+}
 
 SIMDTests.test("Addition") {
   let a = SIMD4<Float>(1, 2, 3, 4)