Merge pull request JuliaLang#12182 from eschnett/fast-complex-fastmath

Implement explicit mixed complex/real arithmetic functions
armgong · Jul 18, 2015 · 2539fcf · 2539fcf
2 parents f9048e3 + 9a7f9c6
commit 2539fcf
Show file tree

Hide file tree

Showing 2 changed files with 64 additions and 26 deletions.
diff --git a/base/fastmath.jl b/base/fastmath.jl
@@ -175,27 +175,49 @@ issubnormal_fast(x) = false
 ComplexTypes = Union{Complex64, Complex128}
 
 @fastmath begin
-    div_fast{T<:FloatTypes}(x::Complex{T}, y::T) =
-        Complex{T}(real(x)/y, imag(x)/y)
-
     abs_fast{T<:ComplexTypes}(x::T) = hypot(real(x), imag(x))
     abs2_fast{T<:ComplexTypes}(x::T) = real(x)*real(x) + imag(x)*imag(x)
-    add_fast{T<:ComplexTypes}(x::T, y::T) = T(real(x)+real(y), imag(x)+imag(y))
     conj_fast{T<:ComplexTypes}(x::T) = T(real(x), -imag(x))
-    function div_fast{T<:ComplexTypes}(x::T, y::T)
-        T(real(x)*real(y) + imag(x)*imag(y),
-           imag(x)*real(y) - real(x)*imag(y)) / abs2(y)
-    end
     inv_fast{T<:ComplexTypes}(x::T) = conj(x) / abs2(x)
     sign_fast{T<:ComplexTypes}(x::T) = x / abs(x)
-    sub_fast{T<:ComplexTypes}(x::T, y::T) = T(real(x)-real(y), imag(x)-imag(y))
-    function mul_fast{T<:ComplexTypes}(x::T, y::T)
+
+    add_fast{T<:ComplexTypes}(x::T, y::T) =
+        T(real(x)+real(y), imag(x)+imag(y))
+    add_fast{T<:FloatTypes}(x::Complex{T}, b::T) =
+        Complex{T}(real(x)+b, imag(x))
+    add_fast{T<:FloatTypes}(a::T, y::Complex{T}) =
+        Complex{T}(a+real(y), imag(y))
+
+    sub_fast{T<:ComplexTypes}(x::T, y::T) =
+        T(real(x)-real(y), imag(x)-imag(y))
+    sub_fast{T<:FloatTypes}(x::Complex{T}, b::T) =
+        Complex{T}(real(x)-b, imag(x))
+    sub_fast{T<:FloatTypes}(a::T, y::Complex{T}) =
+        Complex{T}(a-real(y), -imag(y))
+
+    mul_fast{T<:ComplexTypes}(x::T, y::T) =
         T(real(x)*real(y) - imag(x)*imag(y),
-           real(x)*imag(y) + imag(x)*real(y))
-    end
+          real(x)*imag(y) + imag(x)*real(y))
+    mul_fast{T<:FloatTypes}(x::Complex{T}, b::T) =
+        Complex{T}(real(x)*b, imag(x)*b)
+    mul_fast{T<:FloatTypes}(a::T, y::Complex{T}) =
+        Complex{T}(a*real(y), a*imag(y))
+
+    div_fast{T<:ComplexTypes}(x::T, y::T) =
+        T(real(x)*real(y) + imag(x)*imag(y),
+          imag(x)*real(y) - real(x)*imag(y)) / abs2(y)
+    div_fast{T<:FloatTypes}(x::Complex{T}, b::T) =
+        Complex{T}(real(x)/b, imag(x)/b)
+    div_fast{T<:FloatTypes}(a::T, y::Complex{T}) =
+        Complex{T}(a*real(y), -a*imag(y)) / abs2(y)
 
     eq_fast{T<:ComplexTypes}(x::T, y::T) =
         (real(x)==real(y)) & (imag(x)==imag(y))
+    eq_fast{T<:FloatTypes}(x::Complex{T}, b::T) =
+        (real(x)==b) & (imag(x)==T(0))
+    eq_fast{T<:FloatTypes}(a::T, y::Complex{T}) =
+        (a==real(y)) & (T(0)==imag(y))
+
     ne_fast{T<:ComplexTypes}(x::T, y::T) = !(x==y)
 end
 

diff --git a/test/fastmath.jl b/test/fastmath.jl
@@ -98,6 +98,7 @@ end
 
 # math functions
 
+# real arithmetic
 for T in (Float32, Float64, BigFloat)
     half = 1/convert(T, 2)
     third = 1/convert(T, 3)
@@ -132,21 +133,10 @@ for T in (Float32, Float64, BigFloat)
     end
 end
 
-for T in (Float32, Float64, BigFloat)
-    CT = Complex{T}
-    half = (1+1im)/convert(T, 2) # complex
-    third = 1/convert(T, 3)      # real
-
-    @test isapprox((@fastmath third^3), third^3)
-
-    @test isapprox((@fastmath half/third), half/third)
-    @test isapprox((@fastmath half^3), half^3)
-    @test isapprox((@fastmath cis(third)), cis(third))
-end
-
+# complex arithmetic
 for T in (Complex64, Complex128, Complex{BigFloat})
-    half = (1+1im)/convert(T, 2)
-    third = (1-1im)/convert(T, 3)
+    half = (1+1im)/T(2)
+    third = (1-1im)/T(3)
 
     for f in (:+, :-, :abs, :abs2, :conj, :inv, :sign,
               :acos, :acosh, :asin, :asinh, :atan, :atanh, :cos,
@@ -163,6 +153,32 @@ for T in (Complex64, Complex128, Complex{BigFloat})
     end
 end
 
+# mixed real/complex arithmetic
+for T in (Float32, Float64, BigFloat)
+    CT = Complex{T}
+    half = 1/T(2)
+    third = 1/T(3)
+    chalf = (1+1im)/CT(2)
+    cthird = (1-1im)/CT(3)
+
+    for f in (:+, :-, :*, :/, :(==), :!=, :^)
+        @test isapprox((@eval @fastmath $f($chalf, $third)),
+                       (@eval $f($chalf, $third)))
+        @test isapprox((@eval @fastmath $f($half, $cthird)),
+                       (@eval $f($half, $cthird)))
+        @test isapprox((@eval @fastmath $f($cthird, $half)),
+                       (@eval $f($cthird, $half)))
+        @test isapprox((@eval @fastmath $f($third, $chalf)),
+                       (@eval $f($third, $chalf)))
+    end
+
+    @test isapprox((@fastmath third^3), third^3)
+
+    @test isapprox((@fastmath chalf/third), chalf/third)
+    @test isapprox((@fastmath chalf^3), chalf^3)
+    @test isapprox((@fastmath cis(third)), cis(third))
+end
+
 # issue #10544
 let a = ones(2,2), b = ones(2,2)
     @test isapprox((@fastmath a[1] += 2.0), b[1] += 2.0)