I am writing a routine to store a matrix that contains the values of f(x, z) = \frac{1}{\sqrt{1-x^2}} \times \text{real}\left(\frac{\sqrt{z-1}\sqrt{z+1}}{z-x}\right) for \sim 1000 real values x \in [-1, 1] and \sim 50 complex numbers z.
This computation is a critical part of my code, do you have ideas of how I can improve it?
We cannot simplify \sqrt{z-1}\sqrt{z+1} since z is complex. Otherwise, the branch cut of the square root is no longer on the real axis between -1 and 1.
using LinearAlgebra
using BenchmarkTools
function test(x::Array{Float64,1}, z::Array{ComplexF64,1})
A = zeros(size(x,1), size(z, 1))
@inbounds for (j, zj) in enumerate(z)
bj = √(zj-1.0)*√(zj+1.0)
for (i, xi) in enumerate(x)
ci = (1.0 - xi^2)^(-0.5)
A[i,j] = ci*real(bj/(zj-xi))
end
end
return A
end
x = 2*rand(1000) .- 1
z = randn(50) + im*randn(50)
@btime test($x, $z)
989.748 μs (2 allocations: 390.70 KiB)
The computation arises in a potential flow problem in fluid mechanics. I am interested in the induced effect of a collection of singularities at the locations {z_j} on an infinitely-thin plate materizalized by the interval [-1,1]. Here f(x_i,z_j) computes the induce effect of a point vortex at location z_j at a point x_i along the plate. This computation is used in the forecast step of an ensemble Kalman filter. Therefore, I need to compute this matrix for about 300 (number of ensemble members) different sets of {z_j} values. The function test performs this computation for a distributed set of {x_i} values for only one set of {z_j} values. The set of {x_i} values is always the same. I hope this helps
LoopVectorization produces a huge speedup on my machine. Here is the contents of the file testit.jl:
using LinearAlgebra
using BenchmarkTools
using LoopVectorization
function test(x::Array{Float64,1}, z::Array{ComplexF64,1})
A = zeros(size(x,1), size(z, 1))
@inbounds for (j, zj) in enumerate(z)
bj = √(zj-1.0)*√(zj+1.0)
for (i, xi) in enumerate(x)
ci = (1.0 - xi^2)^(-0.5)
A[i,j] = ci*real(bj/(zj-xi))
end
end
return A
end
function test2(x, z)
a = zeros(length(x),length(z))
b = [√(t-1)*√(t+1) for t in z]
c = [1/sqrt(1 - t^2) for t in x]
@avx for i in axes(a,1), j in axes(a,2)
a[i,j] = c[i] * real(b[j]/(z[j]-x[i]))
end
a
end
x = 2*rand(1000) .- 1
z = randn(50) + im*randn(50)
versioninfo()
a = @btime test($x,$z)
a2 = @btime test2($x,$z)
@show norm(a-a2)
I can confirm this. You even get a little better if you do calculations in single precicion.
using BenchmarkTools
using LoopVectorization
function test(x::Vector{T}, z::Vector{Complex{T}}) where T<:Real
A = zeros(T,length(x), length(z))
b = map(z->√(z - one(T))*√(z + one(T)),z)
c = map(x->(one(T) - x^2)^(-0.5),x)
@avx for j in eachindex(z)
for i in eachindex(x)
A[i,j] = c[i]*real(b[j]/(z[j]-x[i]))
end
end
return A
end
T = Float32
x = 2*rand(T,1000) .- 1
z = randn(Complex{T},50)
It boosts from 70 μs to 62 μs on my machine, which is not the fastest. That is if you do not need the numerical accuracy of double precision.
For whatever it is worth, it is interesting that this function can be written without loops with adjoints. The @fastmath macro speeds up significantely but not as much as Tullio does:
function test4(x, z)
b = @. √(z-1)*√(z+1)
c = @. 1.0/sqrt(1 - x^2)
@fastmath a = @. c*real(b'/(z' - x)) # using adjoints
return a
end
In that case, you could pre-compute the 1/\sqrt{1-x_i^2} values, and pre-allocate the A matrix.
As others have pointed out, @fastmath helps a lot — the reason is that complex-number division normally requires a bunch of costly checks to eliminate spurious overflows and similar conditions that can arise for extreme inputs (e.g. x[j] = 1e300), which you can safely eliminate if you know that you always have moderate-sized inputs and never divide by zero.
One can also use threads, but probably you are better off using threads at a higher level (e.g. to process different ensemble instances in parallel), so I’m only looking at the single-threaded performance here:
using BenchmarkTools
function test!(A, x, c, z)
@fastmath @inbounds for j in eachindex(z)
zj = z[j]
bj = √(zj-1)*√(zj+1)
for i in eachindex(x)
A[i,j] = c[i] * real(bj / (zj - x[i]))
end
end
return A
end
x = 2*rand(1000) .- 1
z = randn(50) + im*randn(50)
A = zeros(size(x,1), size(z, 1))
c = @. inv(sqrt(1 - x^2))
which gives:
julia> @btime test($x, $z); # your original code
1.335 ms (2 allocations: 390.70 KiB)
julia> @btime test!($A, $x, $c, $z); # new code
34.047 μs (0 allocations: 0 bytes)
for about 40× speedup.
(I kept the explicit loops in order to continue to compute bj = -√(zj-1)*√(zj+1) outside the inner loop without storing it in an additional array.)
I tried @avx from LoopVectorization.jl and it unfortunately gives me an error… it works if I write out real(bj / (zj - x[i])) explicitly in terms of real and imaginary parts, but in that case I get negligible speedup from @avx.
Thank you very much for your help. This is definitly much more efficient.
Is eachindex preferable to enumerate for the performance?
Is there something that prevent me from using @fastmath in my future scripts everty time I am not dealing with very large entries?
If you want to be sure that @fastmath does what you want, please check out this page, where it is documented, which kind of optimizations @fastmath makes and what kind of assumptions.
eachindex is more generic, enumerate does not give indexes, it gives 1 with the first element, 2 with the second, and so on; if your Array has more than one dimension, or custom indexes (start at 0 for example) then your code will break. eachindex will always give valid indexes, and for N-dimensional arrays, it will give the indexes in the most efficient order for linear access. If you want something that both give you valid indexes and the value then look at pairs documentation and use it instead of enumerate.
I wouldn’t say that. enumerate is generic, works also for non-standard indexed arrays, and (without testing) possibly for non-indexed collections (I should check this).
It just doesn’t do what many people think it does. It doesn’t return indices. It’s unfortunate that enumerate is so widely known, and pairs isn’t, because I suspect it’s very often misused.
I wonder if a linter would be able to catch this. If output from enumerate is used for indexing, it might suggest pairs instead.
Fair enough. I was saying that eachindex is more generic for the purpose of having both index and element, but maybe it would be best to say that eachindex is the correct choice for that purpose and enumerate is designed for another purpose (easily adding an iteration counter on a loop).