From 73864f93fcd4534dc2df99ae52c75236e64e2982 Mon Sep 17 00:00:00 2001 From: efremale Date: Mon, 13 Jul 2020 20:53:01 +0200 Subject: [PATCH 1/3] [AutoDiff][stdlib] Add JVPs to SIMDDifferentiation.swift.gyb --- .../SIMDDifferentiation.swift.gyb | 172 ++++++++++++++++++ .../validation-test/forward_mode.swift | 166 +++++++++++++++++ 2 files changed, 338 insertions(+) diff --git a/stdlib/public/Differentiation/SIMDDifferentiation.swift.gyb b/stdlib/public/Differentiation/SIMDDifferentiation.swift.gyb index a60ab8dd461ed..e4110340ad696 100644 --- a/stdlib/public/Differentiation/SIMDDifferentiation.swift.gyb +++ b/stdlib/public/Differentiation/SIMDDifferentiation.swift.gyb @@ -59,6 +59,16 @@ where return zeros }) } + + @inlinable + @derivative(of: subscript(_:)) + internal func _jvpSubscript(index: Int) + -> (value: Scalar, differential: (TangentVector) -> Scalar.TangentVector) + { + return (self[index], { v in + return .init(v[index]) + }) + } } %end @@ -82,6 +92,18 @@ where }) } + @inlinable + @derivative(of: +) + static func _jvpAdd(lhs: Self, rhs: Self) + -> ( + value: Self, differential: (TangentVector, TangentVector) -> TangentVector + ) + { + return (lhs + rhs, { ltan, rtan in + return ltan + rtan + }) + } + @inlinable @derivative(of: -) static func _vjpSubtract(lhs: Self, rhs: Self) @@ -94,6 +116,18 @@ where }) } + @inlinable + @derivative(of: -) + static func _jvpSubtract(lhs: Self, rhs: Self) + -> ( + value: Self, differential: (TangentVector, TangentVector) -> TangentVector + ) + { + return (lhs - rhs, { ltan, rtan in + return ltan - rtan + }) + } + @inlinable @derivative(of: -) static func _vjpNegate(rhs: Self) @@ -103,6 +137,16 @@ where return -v }) } + + @inlinable + @derivative(of: -) + static func _jvpNegate(rhs: Self) + -> (value: Self, differential: (TangentVector) -> (TangentVector)) + { + return (-rhs, { v in + return -v + }) + } } extension SIMD @@ -124,6 +168,18 @@ where }) } + @inlinable + @derivative(of: *) + static func _jvpMultiply(lhs: Self, rhs: Self) + -> ( + value: Self, differential: (TangentVector, TangentVector) -> TangentVector + ) + { + return (lhs * rhs, { ltan, rtan in + return lhs * ltan + rtan * rhs + }) + } + @inlinable @derivative(of: /) static func _vjpDivide(lhs: Self, rhs: Self) @@ -135,6 +191,18 @@ where (v / rhs, -lhs / (rhs * rhs) * v) }) } + + @inlinable + @derivative(of: /) + static func _jvpDivide(lhs: Self, rhs: Self) + -> ( + value: Self, differential: (TangentVector, TangentVector) -> TangentVector + ) + { + return ( lhs / rhs, { ltan, rtan in + (ltan * rhs - lhs * rtan) / (rhs * rhs) + }) + } } extension SIMD @@ -156,6 +224,17 @@ where }) } + @inlinable + @derivative(of: +) + static func _jvpAdd(lhs: Scalar, rhs: Self) -> ( + value: Self, + differential: (Scalar.TangentVector, TangentVector) -> TangentVector + ) { + return (lhs + rhs, { ltan, rtan in + return ltan + rtan + }) + } + @inlinable @derivative(of: -) static func _vjpSubtract(lhs: Scalar, rhs: Self) -> ( @@ -167,6 +246,17 @@ where }) } + @inlinable + @derivative(of: -) + static func _jvpSubtract(lhs: Scalar, rhs: Self) -> ( + value: Self, + differential: (Scalar.TangentVector, TangentVector) -> TangentVector + ) { + return (lhs - rhs, { ltan, rtan in + return ltan - rtan + }) + } + @inlinable @derivative(of: +) static func _vjpAdd(lhs: Self, rhs: Scalar) -> ( @@ -178,6 +268,17 @@ where }) } + @inlinable + @derivative(of: +) + static func _jvpAdd(lhs: Self, rhs: Scalar) -> ( + value: Self, + differential: (TangentVector, Scalar.TangentVector) -> TangentVector + ) { + return (lhs + rhs, { ltan, rtan in + return ltan + rtan + }) + } + @inlinable @derivative(of: -) static func _vjpSubtract(lhs: Self, rhs: Scalar) -> ( @@ -188,6 +289,17 @@ where return (v, -v.sum()) }) } + + @inlinable + @derivative(of: -) + static func _jvpSubtract(lhs: Self, rhs: Scalar) -> ( + value: Self, + differential: (TangentVector, Scalar.TangentVector) -> TangentVector + ) { + return (lhs - rhs, { ltan, rtan in + return ltan - rtan + }) + } } extension SIMD @@ -209,6 +321,17 @@ where }) } + @inlinable + @derivative(of: *) + static func _jvpMultiply(lhs: Self, rhs: Scalar) -> ( + value: Self, + differential: (TangentVector, Scalar.TangentVector) -> TangentVector + ) { + return (lhs * rhs, { ltan, rtan in + return lhs * rtan + ltan * rhs + }) + } + @inlinable @derivative(of: /) static func _vjpDivide(lhs: Self, rhs: Scalar) -> ( @@ -220,6 +343,17 @@ where }) } + @inlinable + @derivative(of: /) + static func _jvpDivide(lhs: Self, rhs: Scalar) -> ( + value: Self, + differential: (TangentVector, Scalar.TangentVector) -> TangentVector + ) { + return (lhs / rhs, { ltan, rtan in + (ltan * rhs - lhs * rtan) / (rhs * rhs) + }) + } + @inlinable @derivative(of: *) static func _vjpMultiply(lhs: Scalar, rhs: Self) -> ( @@ -231,6 +365,17 @@ where }) } + @inlinable + @derivative(of: *) + static func _jvpMultiply(lhs: Scalar, rhs: Self) -> ( + value: Self, + differential: (Scalar.TangentVector, TangentVector) -> TangentVector + ) { + return (lhs * rhs, { ltan, rtan in + return lhs * rtan + ltan * rhs + }) + } + @inlinable @derivative(of: /) static func _vjpDivide(lhs: Scalar, rhs: Self) -> ( @@ -241,6 +386,17 @@ where ((v / rhs).sum(), -lhs / (rhs * rhs) * v) }) } + + @inlinable + @derivative(of: /) + static func _jvpDivide(lhs: Scalar, rhs: Self) -> ( + value: Self, + differential: (Scalar.TangentVector, TangentVector) -> TangentVector + ) { + return (lhs / rhs, { ltan, rtan in + (ltan * rhs - lhs * rtan) / (rhs * rhs) + }) + } } // FIXME(TF-1103): Derivative registration does not yet support @@ -261,6 +417,14 @@ where ) { return (sum(), { v in Self(repeating: Scalar(v)) }) } + + @inlinable + @derivative(of: sum) + func _jvpSum() -> ( + value: Scalar, differential (TangentVector) -> Scalar.TangentVector + ) { + return (sum(), { v in v.sum() } + } } */ @@ -279,4 +443,12 @@ where { return (Self(repeating: value), { v in v.sum() }) } + + @inlinable + @derivative(of: init(repeating:)) + static func _jvpInit(repeating value: Scalar) + -> (value: Self, differential: (Scalar.TangentVector) -> TangentVector) + { + return (Self(repeating: value), { v in Self(repeating: v) }) + } } diff --git a/test/AutoDiff/validation-test/forward_mode.swift b/test/AutoDiff/validation-test/forward_mode.swift index e43740c1244a5..5474bd0f6d834 100644 --- a/test/AutoDiff/validation-test/forward_mode.swift +++ b/test/AutoDiff/validation-test/forward_mode.swift @@ -1427,4 +1427,170 @@ ForwardModeTests.test("Array.differentiableReduce") { expectEqual(2 * x[0] + 2 * x[1] + 2 * x[2], differential(at: x, in: sumOfSquaresReduce)(tan)) } +//===----------------------------------------------------------------------===// +// SIMD methods from SIMDDifferentiation.swift.gyb +// Tests replicate reverse mode tests from test/AutoDiff/stdlib/simd.swift +//===----------------------------------------------------------------------===// + +ForwardModeTests.test("init(repeating:)") { + func foo1(x: Float) -> SIMD4 { + return SIMD4(repeating: 2 * x) + } + let (val1, df1) = valueWithDifferential(at: 5, in: foo1) + expectEqual(SIMD4(10, 10, 10, 10), val1) + expectEqual(SIMD4(6, 6, 6, 6), df1(3)) +} + +ForwardModeTests.test("Identity") { + let a = SIMD4(1, 2, 3, 4) + let g = SIMD4(1, 1, 1, 1) + + func foo1(x: SIMD4) -> SIMD4 { + return x + } + let (val1, df1) = valueWithDifferential(at: a, in: foo1) + expectEqual(a, val1) + expectEqual(g, df1(.init(g))) +} + +ForwardModeTests.test("Negate") { + let a = SIMD4(1, 2, 3, 4) + let g = SIMD4(1, 1, 1, 1) + + func foo1(x: SIMD4) -> SIMD4 { + return -x + } + let (val1, df1) = valueWithDifferential(at: a, in: foo1) + expectEqual(-a, val1) + expectEqual(-g, df1(.init(g))) +} + +ForwardModeTests.test("subscript") { + let a = SIMD4(1, 2, 3, 4) + + func foo1(x: SIMD4) -> Float { + return x[3] + } + + let (val1, df1) = valueWithDifferential(at: a, in: foo1) + expectEqual(4, val1) + expectEqual(4, df1(a)) +} + +ForwardModeTests.test("Addition") { + let a = SIMD4(1, 2, 3, 4) + let g = SIMD4(1, 1, 1, 1) + + // SIMD + SIMD + func foo1(x: SIMD4, y: SIMD4) -> SIMD4 { + return x + y + } + let (val1, df1) = valueWithDifferential(at: a, a, in: foo1) + expectEqual(SIMD4(2, 4, 6, 8), val1) + expectEqual(a + g, df1(a, g)) + + // SIMD + Scalar + func foo2(x: SIMD4, y: Float) -> SIMD4 { + return x + y + } + let (val2, df2) = valueWithDifferential(at: a, 5, in: foo2) + expectEqual(SIMD4(6, 7, 8, 9), val2) + expectEqual(g + 1, df2(g, 1)) + + // Scalar + SIMD + func foo3(x: SIMD4, y: Float) -> SIMD4 { + return y + x + } + let (val3, df3) = valueWithDifferential(at: a, 5, in: foo3) + expectEqual(SIMD4(6, 7, 8, 9), val3) + expectEqual(2 + g, df3(g, 2)) +} + +ForwardModeTests.test("Subtraction") { + let a = SIMD4(1, 2, 3, 4) + let g = SIMD4(1, 1, 1, 1) + + // SIMD - SIMD + func foo1(x: SIMD4, y: SIMD4) -> SIMD4 { + return x - y + } + let (val1, df1) = valueWithDifferential(at: a, a, in: foo1) + expectEqual(SIMD4(0, 0, 0, 0), val1) + expectEqual(g - a, df1(g, a)) + + // SIMD - Scalar + func foo2(x: SIMD4, y: Float) -> SIMD4 { + return x - y + } + let (val2, df2) = valueWithDifferential(at: a, 5, in: foo2) + expectEqual(SIMD4(-4, -3, -2, -1), val2) + expectEqual(g - 1, df2(g, 1)) + + // Scalar - SIMD + func foo3(x: SIMD4, y: Float) -> SIMD4 { + return y - x + } + let (val3, df3) = valueWithDifferential(at: a, 5, in: foo3) + expectEqual(SIMD4(4, 3, 2, 1), val3) + expectEqual(2 - g, df3(g, 2)) +} + +ForwardModeTests.test("Multiplication") { + let a = SIMD4(1, 2, 3, 4) + let g = SIMD4(1, 1, 1, 1) + + // SIMD * SIMD + func foo1(x: SIMD4, y: SIMD4) -> SIMD4 { + return x * y + } + let (val1, df1) = valueWithDifferential(at: a, a, in: foo1) + expectEqual(a * a, val1) + expectEqual(a * g + g * a, df1(g, g)) + + // SIMD * Scalar + func foo2(x: SIMD4, y: Float) -> SIMD4 { + return x * y + } + let (val2, df2) = valueWithDifferential(at: a, 5, in: foo2) + expectEqual(a * 5, val2) + expectEqual(a * 2 + g * 5, df2(g, 2)) + + // Scalar * SIMD + func foo3(x: SIMD4, y: Float) -> SIMD4 { + return y * x + } + let (val3, df3) = valueWithDifferential(at: a, 5, in: foo3) + expectEqual(a * 5, val3) + expectEqual(a * 3 + g * 5, df3(g, 3)) +} + +ForwardModeTests.test("Division") { + let a = SIMD4(1, 2, 3, 4) + let g = SIMD4(1, 1, 1, 1) + + // SIMD / SIMD + func foo1(x: SIMD4, y: SIMD4) -> SIMD4 { + return x / y + } + let (val1, df1) = valueWithDifferential(at: a, a, in: foo1) + expectEqual(a / a, val1) + expectEqual((g * a - a * g) / (a * a)/* == 0 */, df1(g, g)) + + // SIMD / Scalar + func foo2(x: SIMD4, y: Float) -> SIMD4 { + return x / y + } + let (val2, df2) = valueWithDifferential(at: a, 5, in: foo2) + expectEqual(a / 5, val2) + expectEqual((g * 5 - a * 2) / (5 * 5), df2(g, 2)) + + // Scalar / SIMD + func foo3(x: Float, y: SIMD4) -> SIMD4 { + return x / y + } + let (val3, df3) = valueWithDifferential(at: 5, a, in: foo3) + expectEqual(5 / a, val3) + expectEqual((3 * a - 5 * g) / (a * a), df3(3, g)) +} + runAllTests() From 2f9d0b8a09244e7779016d85137757ffa1a80d3e Mon Sep 17 00:00:00 2001 From: efremale Date: Mon, 13 Jul 2020 21:50:08 +0200 Subject: [PATCH 2/3] Add failing test --- .../validation-test/forward_mode.swift | 70 +++++++++++++++++++ 1 file changed, 70 insertions(+) diff --git a/test/AutoDiff/validation-test/forward_mode.swift b/test/AutoDiff/validation-test/forward_mode.swift index 5474bd0f6d834..68549abf7fcc6 100644 --- a/test/AutoDiff/validation-test/forward_mode.swift +++ b/test/AutoDiff/validation-test/forward_mode.swift @@ -1593,4 +1593,74 @@ ForwardModeTests.test("Division") { expectEqual((3 * a - 5 * g) / (a * a), df3(3, g)) } +// FIXME(SR-13210): Fix forward-mode SIL verification error. +/* ForwardModeTests.test("Generics") { + let a = SIMD3(1, 2, 3) + let g = SIMD3(1, 1, 1) + + func testInit(x: Scalar) -> SIMDType + where SIMDType.Scalar == Scalar, + SIMDType : Differentiable, + Scalar : BinaryFloatingPoint & Differentiable, + SIMDType.TangentVector == SIMDType, + Scalar.TangentVector == Scalar { + return SIMDType.init(repeating: x) + } + func simd3Init(x: Double) -> SIMD3 { testInit(x: x) } + let (val1, df1) = valueWithDifferential(at: 10, in: simd3Init) + expectEqual(SIMD3(10, 10, 10), val1) + expectEqual(SIMD3(5, 5, 5), df1(5)) + + // SIMDType + SIMDType + func testAddition(lhs: SIMDType, rhs: SIMDType) + -> SIMDType + where SIMDType.Scalar == Scalar, + SIMDType : Differentiable, + SIMDType.TangentVector : SIMD, + Scalar : BinaryFloatingPoint, + SIMDType.TangentVector.Scalar : BinaryFloatingPoint { + return lhs + rhs + } + func simd3Add(lhs: SIMD3, rhs: SIMD3) -> SIMD3 { + return testAddition(lhs: lhs, rhs: rhs) + } + let (val2, df2) = valueWithDifferential(at: a, a, in: simd3Add) + expectEqual(SIMD3(2, 4, 6), val2) + expectEqual(g + a, df2(g, a)) + + // Scalar - SIMDType + func testSubtraction(lhs: Scalar, rhs: SIMDType) + -> SIMDType + where SIMDType.Scalar == Scalar, + SIMDType : Differentiable, + Scalar : BinaryFloatingPoint & Differentiable, + SIMDType.TangentVector == SIMDType, + Scalar.TangentVector == Scalar { + return lhs - rhs + } + func simd3Subtract(lhs: Double, rhs: SIMD3) -> SIMD3 { + return testSubtraction(lhs: lhs, rhs: rhs) + } + let (val3, df3) = valueWithDifferential(at: 5, a, in: simd3Subtract) + expectEqual(SIMD3(4, 3, 2), val3) + expectEqual(2 - g, df3(2, g)) + + // SIMDType * Scalar + func testMultipication(lhs: SIMDType, rhs: Scalar) + -> SIMDType + where SIMDType.Scalar == Scalar, + SIMDType : Differentiable, + Scalar : BinaryFloatingPoint & Differentiable, + SIMDType.TangentVector == SIMDType, + Scalar.TangentVector == Scalar { + return lhs * rhs + } + func simd3Multiply(lhs: SIMD3, rhs: Double) -> SIMD3 { + return testMultipication(lhs: lhs, rhs: rhs) + } + let (val4, df4) = valueWithDifferential(at: a, 5, in: simd3Multiply) + expectEqual(SIMD3(5, 10, 15), val4) + expectEqual(a * 3 + g * 5 , df4(g, 3)) +} */ + runAllTests() From 98c836023bc51519e67594149f56c8f073866185 Mon Sep 17 00:00:00 2001 From: Dan Zheng Date: Mon, 13 Jul 2020 14:05:23 -0700 Subject: [PATCH 3/3] Narrow down SR-13210 failing test. --- test/AutoDiff/validation-test/forward_mode.swift | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/test/AutoDiff/validation-test/forward_mode.swift b/test/AutoDiff/validation-test/forward_mode.swift index 68549abf7fcc6..b75cc1645d59e 100644 --- a/test/AutoDiff/validation-test/forward_mode.swift +++ b/test/AutoDiff/validation-test/forward_mode.swift @@ -1593,11 +1593,12 @@ ForwardModeTests.test("Division") { expectEqual((3 * a - 5 * g) / (a * a), df3(3, g)) } -// FIXME(SR-13210): Fix forward-mode SIL verification error. -/* ForwardModeTests.test("Generics") { +ForwardModeTests.test("Generics") { let a = SIMD3(1, 2, 3) let g = SIMD3(1, 1, 1) + // FIXME(SR-13210): Fix forward-mode SIL verification error. + /* func testInit(x: Scalar) -> SIMDType where SIMDType.Scalar == Scalar, SIMDType : Differentiable, @@ -1610,6 +1611,7 @@ ForwardModeTests.test("Division") { let (val1, df1) = valueWithDifferential(at: 10, in: simd3Init) expectEqual(SIMD3(10, 10, 10), val1) expectEqual(SIMD3(5, 5, 5), df1(5)) + */ // SIMDType + SIMDType func testAddition(lhs: SIMDType, rhs: SIMDType) @@ -1661,6 +1663,6 @@ ForwardModeTests.test("Division") { let (val4, df4) = valueWithDifferential(at: a, 5, in: simd3Multiply) expectEqual(SIMD3(5, 10, 15), val4) expectEqual(a * 3 + g * 5 , df4(g, 3)) -} */ +} runAllTests()