Enable SIMD with Base.literal_pow (#332)

Author: tkf

.travis.yml

@@ -10,9 +10,9 @@ notifications:
 sudo: false
 script:
     - if [[ -a .git/shallow ]]; then git fetch --unshallow; fi
-    - julia -e 'Pkg.clone(pwd()); Pkg.build("ForwardDiff"); Pkg.test("ForwardDiff"; coverage=true)';
-    - julia -O3 -e 'include("test/SIMDTest.jl")';
+    - julia --color=yes -e 'Pkg.clone(pwd()); Pkg.build("ForwardDiff"); Pkg.test("ForwardDiff"; coverage=true)';
+    - julia --color=yes -O3 -e 'using Pkg; Pkg.add("StaticArrays"); include("test/SIMDTest.jl")';
 after_success:
-    - julia -e 'Pkg.add("Coverage"); using Coverage; Coveralls.submit(Coveralls.process_folder())'
-    - julia -e 'Pkg.add("Documenter")'
-    - julia -e 'include("docs/make.jl")'
+    - julia --color=yes -e 'Pkg.add("Coverage"); using Coverage; Coveralls.submit(Coveralls.process_folder())'
+    - julia --color=yes -e 'Pkg.add("Documenter")'
+    - julia --color=yes -e 'include("docs/make.jl")'

src/dual.jl

@@ -426,6 +426,18 @@ for f in (:(Base.:^), :(NaNMath.pow))
     end
 end
 
+@inline Base.literal_pow(::typeof(^), x::Dual{T}, ::Val{0}) where {T} =
+    Dual{T}(one(value(x)), zero(partials(x)))
+
+for y in 1:3
+    @eval @inline function Base.literal_pow(::typeof(^), x::Dual{T}, ::Val{$y}) where {T}
+        v = value(x)
+        expv = v^$y
+        deriv = $y * v^$(y - 1)
+        return Dual{T}(expv, deriv * partials(x))
+    end
+end
+
 # hypot #
 #-------#
 

test/DualTest.jl

@@ -460,9 +460,10 @@ end
     x1 = Dual{:t1}(x0, 1.0)
     x2 = Dual{:t2}(x1, 1.0)
     x3 = Dual{:t3}(x2, 1.0)
-    @test x3^2 === x3 * x3
-    @test x2^1 === x2
-    @test x1^0 === Dual{:t1}(1.0, 0.0)
+    pow = ^  # to call non-literal power
+    @test pow(x3, 2) === x3^2 === x3 * x3
+    @test pow(x2, 1) === x2^1 === x2
+    @test pow(x1, 0) === x1^0 === Dual{:t1}(1.0, 0.0)
 end
 
 end # module

test/SIMDTest.jl

@@ -3,6 +3,7 @@ module SIMDTest
 using Test
 using ForwardDiff: Dual, valtype
 using InteractiveUtils: code_llvm
+using StaticArrays: SVector
 
 const DUALS = (Dual(1., 2., 3., 4.),
                Dual(1., 2., 3., 4., 5.),
@@ -17,7 +18,7 @@ function simd_sum(x::Vector{T}) where T
     return s
 end
 
-for D in map(typeof, DUALS)
+@testset "SIMD $D" for D in map(typeof, DUALS)
     plus_bitcode = sprint(io -> code_llvm(io, +, (D, D)))
     @test occursin("fadd <4 x double>", plus_bitcode)
 
@@ -32,6 +33,9 @@ for D in map(typeof, DUALS)
     @test occursin(r"fadd \<.*?x double\>", div_bitcode)
     @test occursin(r"fmul \<.*?x double\>", div_bitcode)
 
+    pow_bitcode = sprint(io -> code_llvm(io, ^, (D, Int)))
+    @test occursin(r"fmul \<.*?x double\>", pow_bitcode)
+
     exp_bitcode = sprint(io -> code_llvm(io, ^, (D, D)))
     @test occursin(r"fadd \<.*?x double\>", exp_bitcode)
     if !(valtype(D) <: Dual)
@@ -44,4 +48,26 @@ for D in map(typeof, DUALS)
     end
 end
 
+# `pow2dot` is chosen so that `@code_llvm pow2dot(SVector(1:1.0:4...))`
+# generates code with SIMD instructions.
+# See:
+# https://github.com/JuliaDiff/ForwardDiff.jl/pull/332
+# https://github.com/JuliaDiff/ForwardDiff.jl/pull/331#issuecomment-406107260
+@inline pow2(x) = x^2
+pow2dot(xs) = pow2.(xs)
+
+# Nested dual such as `Dual(Dual(1., 2.), Dual(3., 4.))` only produces
+# "fmul <2 x double>" so it is excluded from the following test.
+const POW_DUALS = (Dual(1., 2.),
+                   Dual(1., 2., 3.),
+                   Dual(1., 2., 3., 4.),
+                   Dual(1., 2., 3., 4., 5.))
+
+@testset "SIMD square of $D" for D in map(typeof, POW_DUALS)
+    pow_bitcode = sprint(io -> code_llvm(io, pow2dot, (SVector{4, D},)))
+    @test occursin(r"(.*fmul \<4 x double\>){2}"s, pow_bitcode)
+    # "{2}" is for asserting that fmul has to appear at least twice:
+    # once for `.value` and once for `.partials`.
+end
+
 end # module