Use multithreading for rolling function

user_1117766 0
Viewed 39

I have some rolling function which evaluates an array according to a function,

rolling(x, func, window=10) = [func(x[1+i:window+i+1]) for i in 0:(length(x)-window-1)]

for instance for the mean:

julia> using Statistics

julia> rolling(1:10,mean,5)
5-element Array{Float64,1}:
 3.5
 4.5
 5.5
 6.5
 7.5

I wonder if there was an easy and fast and safe way to parallelize this function?

julia parallel-processing
1 Answers

Here it is. Note that for mean it could be done differently because of its rolling properties. Performance will also vary. For smaller loops sometimes @simd could be better (strongly depends on the actual scenario. Hence consider this a starting point.


using Statistics

rolling0(x, func, window=10) = [func(@view x[1+i:window+i+1]) for i in 0:(length(x)-window-1)]

function rolling(x, func, window=10)
    res = Vector{Float64}(undef, length(x)-window)
    Threads.@threads for i in 1:(length(x)-window)
        @inbounds res[i] = func(@view x[i:i+window]) 
    end
    res
end

z = rand(10000)
@assert rolling0(z,mean) == rolling(z,mean)
using BenchmarkTools

Tests:

julia> @btime rolling0($z, mean);
  104.601 μs (9994 allocations: 546.48 KiB)

julia> @btime rolling($z, mean);
  48.100 μs (10021 allocations: 549.88 KiB)
user_99577100
Related