Define real, imag, conj, angle, and abs2 for Vec{N,<:Real} (

#127)
eschnett · Sep 18, 2024 · 33b78fc · 33b78fc
1 parent c332a03
commit 33b78fc
Show file tree

Hide file tree

Showing 2 changed files with 14 additions and 4 deletions.
diff --git a/src/simdvec.jl b/src/simdvec.jl
@@ -164,6 +164,7 @@ Base.FastMath.sub_fast(v::Vec{<:Any, <:FloatingTypes}) = Vec(Intrinsics.fneg(v.d
 Base.:~(v::Vec{N, T}) where {N, T<:IntegerTypes} = Vec(Intrinsics.xor(v.data, Vec{N, T}(-1).data))
 Base.:~(v::Vec{N, Bool}) where {N} = Vec(Intrinsics.xor(v.data, Vec{N, Bool}(true).data))
 Base.abs(v::Vec{N, T}) where {N, T} = Vec(vifelse(v < zero(T), -v, v))
+Base.abs2(v::Vec{N, T}) where {N, T} = v*v
 Base.:!(v1::Vec{N,Bool}) where {N} = ~v1
 Base.inv(v::Vec{N, T}) where {N, T<:FloatingTypes} = one(T) / v
 
@@ -492,3 +493,12 @@ end
 @inline function shufflevector(x::Vec{N, T}, y::Vec{N, T}, ::Val{I}) where {N, T, I}
     Vec(Intrinsics.shufflevector(x.data, y.data, Val(I)))
 end
+
+############################
+# Complex Number functions #
+############################
+
+@inline Base.real(v::Vec{N, T}) where {N, T<:ScalarTypes} = v
+@inline Base.conj(v::Vec{N, T}) where {N, T<:ScalarTypes} = v
+@inline Base.angle(v::Vec{N, T}) where {N, T <: FloatingTypes} = vifelse(signbit(v), Vec{N, T}(T(pi)), zero(v))
+@inline Base.imag(v::Vec{N, T}) where {N, T <: ScalarTypes} = zero(v)
diff --git a/test/runtests.jl b/test/runtests.jl
@@ -102,8 +102,8 @@ llvm_ir(f, args) = sprint(code_llvm, f, Base.typesof(args...))
         global const v8i32c = map(x->Int32(x*2), v8i32)
 
         notbool(x) = !(x>=typeof(x)(0))
-        for op in (~, +, -, abs, notbool, sign, signbit, count_ones, count_zeros,
-                   leading_ones, leading_zeros, trailing_ones, trailing_zeros)
+        for op in (~, +, -, abs, abs2, notbool, sign, signbit, count_ones, count_zeros,
+                   leading_ones, leading_zeros, conj, real, imag, trailing_ones, trailing_zeros)
             @test Tuple(op(V8I32(v8i32))) == map(op, v8i32)
         end
 
@@ -181,8 +181,8 @@ llvm_ir(f, args) = sprint(code_llvm, f, Base.typesof(args...))
         sqrtabs(x) = sqrt(abs(x))
         for op in (
                 +, -,
-                abs, ceil, inv, isfinite, isinf, isnan, issubnormal, floor, powi4,
-                round, sign, signbit, sqrtabs, trunc)
+                abs, abs2, angle, ceil, conj, inv, imag, isfinite, isinf, isnan, issubnormal, floor, powi4,
+                real, round, sign, signbit, sqrtabs, trunc)
             @test Tuple(op(V4F64(v4f64))) === map(op, v4f64)
         end
         function Base.isapprox(t1::Tuple,t2::Tuple)