diff --git a/Project.toml b/Project.toml index d19d0d1e..ed89a9d4 100644 --- a/Project.toml +++ b/Project.toml @@ -11,7 +11,7 @@ StableTasks = "91464d47-22a1-43fe-8b7f-2d57ee82463f" [compat] BangBang = "0.4" ChunkSplitters = "2.1" -StableTasks = "0.1.2" +StableTasks = "0.1.4" julia = "1.6" [extras] diff --git a/README.md b/README.md index b888aa8d..7ffb8da9 100644 --- a/README.md +++ b/README.md @@ -12,14 +12,14 @@ Unlike most JuliaFolds2 packages, it is not built off of Rather, OhMyThreads is meant to be a simpler, more maintainable, and more accessible alternative to packages like [ThreadsX.jl](https://github.com/tkf/ThreadsX.jl) or [Folds.jl](https://github.com/JuliaFolds2/Folds.jl). -OhMyThreads.jl re-exports the very useful function `chunks` from +OhMyThreads.jl re-exports the function `chunks` from [ChunkSplitters.jl](https://github.com/JuliaFolds2/ChunkSplitters.jl), and provides the following functions:
tmapreduce

``` -tmapreduce(f, op, A::AbstractArray; +tmapreduce(f, op, A::AbstractArray...; [init], nchunks::Int = nthreads(), split::Symbol = :batch, @@ -63,7 +63,7 @@ ____________________________

``` -treducemap(op, f, A::AbstractArray; +treducemap(op, f, A::AbstractArray...; [init], nchunks::Int = nthreads(), split::Symbol = :batch, @@ -107,7 +107,8 @@ ____________________________

``` -treduce(op, A::AbstractArray; [init], +treduce(op, A::AbstractArray...; + [init], nchunks::Int = nthreads(), split::Symbol = :batch, schedule::Symbol =:dynamic, @@ -150,7 +151,7 @@ ____________________________

``` -tmap(f, [OutputElementType], A::AbstractArray; +tmap(f, [OutputElementType], A::AbstractArray...; nchunks::Int = nthreads(), split::Symbol = :batch, schedule::Symbol =:dynamic) @@ -174,7 +175,7 @@ ____________________________

``` -tmap!(f, out, A::AbstractArray; +tmap!(f, out, A::AbstractArray...; nchunks::Int = nthreads(), split::Symbol = :batch, schedule::Symbol =:dynamic) @@ -198,7 +199,7 @@ ____________________________

``` -tforeach(f, A::AbstractArray; +tforeach(f, A::AbstractArray...; nchunks::Int = nthreads(), split::Symbol = :batch, schedule::Symbol =:dynamic) :: Nothing diff --git a/src/OhMyThreads.jl b/src/OhMyThreads.jl index 166070a5..80005b5a 100644 --- a/src/OhMyThreads.jl +++ b/src/OhMyThreads.jl @@ -7,7 +7,7 @@ using ChunkSplitters: chunks export chunks, treduce, tmapreduce, treducemap, tmap, tmap!, tforeach, tcollect """ - tmapreduce(f, op, A::AbstractArray; + tmapreduce(f, op, A::AbstractArray...; [init], nchunks::Int = nthreads(), split::Symbol = :batch, @@ -42,7 +42,7 @@ needed if you are using a `:static` schedule, since the `:dynamic` schedule is u function tmapreduce end """ - treducemap(op, f, A::AbstractArray; + treducemap(op, f, A::AbstractArray...; [init], nchunks::Int = nthreads(), split::Symbol = :batch, @@ -79,7 +79,8 @@ function treducemap end """ - treduce(op, A::AbstractArray; [init], + treduce(op, A::AbstractArray...; + [init], nchunks::Int = nthreads(), split::Symbol = :batch, schedule::Symbol =:dynamic, @@ -114,7 +115,7 @@ needed if you are using a `:static` schedule, since the `:dynamic` schedule is u function treduce end """ - tforeach(f, A::AbstractArray; + tforeach(f, A::AbstractArray...; nchunks::Int = nthreads(), split::Symbol = :batch, schedule::Symbol =:dynamic) :: Nothing @@ -134,7 +135,7 @@ A multithreaded function like `Base.foreach`. Apply `f` to each element of `A` o function tforeach end """ - tmap(f, [OutputElementType], A::AbstractArray; + tmap(f, [OutputElementType], A::AbstractArray...; nchunks::Int = nthreads(), split::Symbol = :batch, schedule::Symbol =:dynamic) @@ -154,7 +155,7 @@ fewer allocations than the version where `OutputElementType` is not specified. function tmap end """ - tmap!(f, out, A::AbstractArray; + tmap!(f, out, A::AbstractArray...; nchunks::Int = nthreads(), split::Symbol = :batch, schedule::Symbol =:dynamic) diff --git a/src/implementation.jl b/src/implementation.jl index 56d786ec..a2a6ee0c 100644 --- a/src/implementation.jl +++ b/src/implementation.jl @@ -9,77 +9,96 @@ using Base.Threads: nthreads, @threads using BangBang: BangBang, append!! -function tmapreduce(f, op, A; +function tmapreduce(f, op, Arrs...; nchunks::Int=nthreads(), split::Symbol=:batch, schedule::Symbol=:dynamic, outputtype::Type=Any, kwargs...) - if schedule === :dynamic - _tmapreduce(f, op, A, outputtype, nchunks, split; kwargs...) + if schedule === :dynamic + _tmapreduce(f, op, Arrs, outputtype, nchunks, split, :default; kwargs...) + elseif schedule === :interactive + _tmapreduce(f, op, Arrs, outputtype, nchunks, split, :interactive; kwargs...) elseif schedule === :static - _tmapreduce_static(f, op, A, outputtype, nchunks, split; kwargs...) + _tmapreduce_static(f, op, Arrs, outputtype, nchunks, split; kwargs...) else schedule_err(schedule) end end @noinline schedule_err(s) = error(ArgumentError("Invalid schedule option: $s, expected :dynamic or :static.")) -treducemap(op, f, A; kwargs...) = tmapreduce(f, op, A; kwargs...) +treducemap(op, f, A...; kwargs...) = tmapreduce(f, op, A...; kwargs...) -function _tmapreduce(f, op, A, ::Type{OutputType}, nchunks, split=:batch; kwargs...)::OutputType where {OutputType} - tasks = map(chunks(A; n=nchunks, split)) do inds - @spawn mapreduce(f, op, @view(A[inds]); kwargs...) +function _tmapreduce(f, op, Arrs, ::Type{OutputType}, nchunks, split, schedule; kwargs...)::OutputType where {OutputType} + check_all_have_same_indices(Arrs) + tasks = map(chunks(first(Arrs); n=nchunks, split)) do inds + args = map(A -> A[inds], Arrs) + @spawn schedule mapreduce(f, op, args...; kwargs...) end mapreduce(fetch, op, tasks) end -function _tmapreduce_static(f, op, A, ::Type{OutputType}, nchunks, split; kwargs...) where {OutputType} +function _tmapreduce_static(f, op, Arrs, ::Type{OutputType}, nchunks, split; kwargs...) where {OutputType} nt = nthreads() + check_all_have_same_indices(Arrs) if nchunks > nt # We could implement strategies, like round-robin, in the future throw(ArgumentError("We currently only support `nchunks <= nthreads()` for static scheduling.")) end - tasks = map(enumerate(chunks(A; n=nchunks, split))) do (c, inds) + tasks = map(enumerate(chunks(first(Arrs); n=nchunks, split))) do (c, inds) tid = @inbounds nthtid(c) - @spawnat tid mapreduce(f, op, @view(A[inds]); kwargs...) + args = map(A -> A[inds], Arrs) + @spawnat tid mapreduce(f, op, args...; kwargs...) end mapreduce(fetch, op, tasks) end + +check_all_have_same_indices(Arrs) = let A = first(Arrs), Arrs = Arrs[2:end] + if !all(B -> eachindex(A) == eachindex(B), Arrs) + error("The indices of the input arrays must match the indices of the output array.") + end +end + #------------------------------------------------------------- -function treduce(op, A; kwargs...) - tmapreduce(identity, op, A; kwargs...) +function treduce(op, A...; kwargs...) + tmapreduce(identity, op, A...; kwargs...) end #------------------------------------------------------------- -function tforeach(f, A::AbstractArray; kwargs...)::Nothing - tmapreduce(f, (l, r) -> l, A; kwargs..., init=nothing, outputtype=Nothing) +function tforeach(f, A...; kwargs...)::Nothing + tmapreduce(f, (l, r) -> l, A...; kwargs..., init=nothing, outputtype=Nothing) end #------------------------------------------------------------- -function tmap(f, ::Type{T}, A::AbstractArray; kwargs...) where {T} - tmap!(f, similar(A, T), A; kwargs...) +function tmap(f, ::Type{T}, A::AbstractArray, _Arrs::AbstractArray...; kwargs...) where {T} + Arrs = (A, _Arrs...) + tmap!(f, similar(A, T), Arrs...; kwargs...) end -function tmap(f, A; nchunks::Int= 2*nthreads(), kwargs...) +function tmap(f, A::AbstractArray, _Arrs::AbstractArray...; nchunks::Int=nthreads(), kwargs...) + Arrs = (A, _Arrs...) + check_all_have_same_indices(Arrs) the_chunks = collect(chunks(A; n=nchunks)) # It's vital that we force split=:batch here because we're not doing a commutative operation! v = tmapreduce(append!!, the_chunks; kwargs..., nchunks, split=:batch) do inds - map(f, @view A[inds]) + args = map(A -> @view(A[inds]), Arrs) + map(f, args...) end reshape(v, size(A)...) end -@propagate_inbounds function tmap!(f, out, A::AbstractArray; kwargs...) - @boundscheck eachindex(out) == eachindex(A) || error("The indices of the input array must match the indices of the output array.") +@propagate_inbounds function tmap!(f, out, A::AbstractArray, _Arrs::AbstractArray...; kwargs...) + Arrs = (A, _Arrs...) + @boundscheck check_all_have_same_indices((out, Arrs...)) # It's vital that we force split=:batch here because we're not doing a commutative operation! - tforeach(eachindex(A); kwargs..., split=:batch) do i - fAi = f(@inbounds A[i]) - out[i] = fAi + tforeach(eachindex(out); kwargs..., split=:batch) do i + args = map(A -> @inbounds(A[i]), Arrs) + res = f(args...) + out[i] = res end out end diff --git a/test/runtests.jl b/test/runtests.jl index c5b8d3af..f82c2140 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -1,43 +1,43 @@ using Test, OhMyThreads @testset "Basics" begin - for (~, f, op, itr) ∈ [ - (isapprox, sin, +, rand(ComplexF64, 10, 10)), - (isapprox, cos, max, 1:100000), - (==, round, vcat, randn(1000)), - (==, last, *, [1=>"a", 2=>"b", 3=>"c", 4=>"d", 5=>"e"]) + for (~, f, op, itrs) ∈ [ + (isapprox, sin∘*, +, (rand(ComplexF64, 10, 10), rand(-10:10, 10, 10))), + (isapprox, cos, max, (1:100000,)), + (==, round, vcat, (randn(1000),)), + (==, last, *, ([1=>"a", 2=>"b", 3=>"c", 4=>"d", 5=>"e"],)) ] - @testset for schedule ∈ (:static, :dynamic,) + @testset for schedule ∈ (:static, :dynamic, :interactive) @testset for split ∈ (:batch, :scatter) if split == :scatter # scatter only works for commutative operators if op ∈ (vcat, *) continue end end - for nchunks ∈ (1, 2, 6, 10, 100) + for nchunks ∈ (1, 2, 6, 10) if schedule == :static && nchunks > Threads.nthreads() continue end kwargs = (; schedule, split, nchunks) - mapreduce_f_op_itr = mapreduce(f, op, itr) - @test tmapreduce(f, op, itr; kwargs...) ~ mapreduce_f_op_itr - @test treducemap(op, f, itr; kwargs...) ~ mapreduce_f_op_itr - @test treduce(op, f.(itr); kwargs...) ~ mapreduce_f_op_itr + mapreduce_f_op_itr = mapreduce(f, op, itrs...) + @test tmapreduce(f, op, itrs...; kwargs...) ~ mapreduce_f_op_itr + @test treducemap(op, f, itrs...; kwargs...) ~ mapreduce_f_op_itr + @test treduce(op, f.(itrs...); kwargs...) ~ mapreduce_f_op_itr - map_f_itr = map(f, itr) - @test all(tmap(f, Any, itr; kwargs...) .~ map_f_itr) - @test all(tcollect(Any, (f(x) for x in itr); kwargs...) .~ map_f_itr) - @test all(tcollect(Any, f.(itr); kwargs...) .~ map_f_itr) + map_f_itr = map(f, itrs...) + @test all(tmap(f, Any, itrs...; kwargs...) .~ map_f_itr) + @test all(tcollect(Any, (f(x...) for x in collect(zip(itrs...))); kwargs...) .~ map_f_itr) + @test all(tcollect(Any, f.(itrs...); kwargs...) .~ map_f_itr) - @test tmap(f, itr; kwargs...) ~ map_f_itr - @test tcollect((f(x) for x in itr); kwargs...) ~ map_f_itr - @test tcollect(f.(itr); kwargs...) ~ map_f_itr + @test tmap(f, itrs...; kwargs...) ~ map_f_itr + @test tcollect((f(x...) for x in collect(zip(itrs...))); kwargs...) ~ map_f_itr + @test tcollect(f.(itrs...); kwargs...) ~ map_f_itr - RT = Core.Compiler.return_type(f, Tuple{eltype(itr)}) + RT = Core.Compiler.return_type(f, Tuple{eltype.(itrs)...}) - @test tmap(f, RT, itr; kwargs...) ~ map_f_itr - @test tcollect(RT, (f(x) for x in itr); kwargs...) ~ map_f_itr - @test tcollect(RT, f.(itr); kwargs...) ~ map_f_itr + @test tmap(f, RT, itrs...; kwargs...) ~ map_f_itr + @test tcollect(RT, (f(x...) for x in collect(zip(itrs...))); kwargs...) ~ map_f_itr + @test tcollect(RT, f.(itrs...); kwargs...) ~ map_f_itr end end end