I have just put up this small package that provides a chunks function to be used within threaded loops. It is something that I find useful and appears somewhat often here:
This is so basic that I have always the impression that it must be done already, perhaps even in Base as a type of Iterator. function.
Yet, for example, Iterators.partition is not what one wants, because the partition is not the most even possible:
Nice! By looking at IteratorTools I now implemented the splitting as an iterator (Julia is actually cool, isn’t it?).
It works like this:
julia> using ChunkSplitters
julia> x = rand(7);
julia> Threads.@threads for (range,ichunk) in chunks(x, 3, :batch)
@show (range, ichunk)
end
(range, ichunk) = (6:7, 3)
(range, ichunk) = (1:3, 1)
(range, ichunk) = (4:5, 2)
Such that we can do, slightly more cleanly:
julia> using ChunkSplitters
julia> function sum_parallel(f, x; nchunks=Threads.nthreads())
s = fill(zero(eltype(x)), nchunks)
Threads.@threads for (range, ichunk) in chunks(x, nchunks)
for i in range
s[ichunk] += f(x[i])
end
end
return sum(s)
end
sum_parallel (generic function with 1 methods)
julia> x = rand(10^7);
julia> Threads.nthreads()
12
julia> @btime sum(x -> log(x)^7, $x)
115.026 ms (0 allocations: 0 bytes)
-5.062317099586189e10
julia> @btime sum_parallel(x -> log(x)^7, $x; nchunks=128)
19.210 ms (74 allocations: 7.58 KiB)
-5.062317099585973e10
That performs nicely, for example in comparison with:
And I may be registering the package anyways, so I can experiment more with the functionality (also I don’t know when or even if the PR will be accepted).
Finally this package, although simple, turned out to be useful, and more people are using it. Thus, a 1.0.0 version was released recently, and now I’ve put up a small but decent docs page:
We released now a new version of ChunkSplitters.jl (v2.1.0) with a new and breaking interface (although the release itself is not breaking because the previous interface is still functional, although no longer documented):
Now, a typical use of chunks is of the form:
julia> x = rand(7);
julia> @threads for inds in chunks(x; n=3, split=:batch)
@show inds
end
inds = 1:1:3
inds = 4:1:5
inds = 6:1:7
and it returns only the indices of the array of each chunk. Additionally, the number of chunks is
set by a keyword parameter, n, and the type of splitting is set by the optional split keyword.
To retrieve the chunk index (to avoid using threadid()) enumerate should be used:
julia> @threads for (ichunk, inds) in enumerate(chunks(x; n=3))
@show ichunk, inds
end
(ichunk, inds) = (1, 1:1:3)
(ichunk, inds) = (2, 4:1:5)
(ichunk, inds) = (3, 6:1:7)
This last output is equivalent to the previous interface with a call of the form for (ichunk, inds) in chunks(x, 3). This last form is still supported, but will be deprecated in version 3.0 of the package.
This change makes many codes cleaner when they do not use the chunk index, and also is consistent with the general Julia syntax of using enumerate to retrieve the indices of the counter of the iteration.
Important: the new interface is currently the only documented one.
I was using the ranges returned by chunks as input to a constructor of iterators, this doesn’t seem to be possible with the new interface - so may be the old one still can be kept documented ? This is the code: