Support minsize setting; don't spawn tasks when only 1 chunk (#145)

* Support `minsize` setting; don't spawn tasks when only 1 chunk

* fix multiple chunks detection

* handle `minsize=nothing` case

* typo

* don't rely on upstream changes to ChunkSplitters

* make downgrade CI happy

* rename `minsize` -> `minchunksize`

* print nicer

* don't show minimum chunk size for FixedSize schedulers

Author: MasonProtter

CHANGELOG.md

@@ -1,6 +1,11 @@
 OhMyThreads.jl Changelog
 =========================
 
+Version 0.8.2
+------------
+- ![Feature][badge-feature] Added a `minchunksize` chunking argument for schedulers, so that they can specify a lower bound on the size of chunks which are worth parallelizing. For example, `treduce(+, 1:10; minchunksize=100)` will run serially, but `treduce(+, 1:1000000; minchunksize=100)` will be parallelized.
+- ![Enhancement][badge-enhancement] Operations on collections with only one 'chunk' no longer spawn an unnecessary task. That means operations like `treduce(+, 1:10; minchunksize=100)` will have less overhead.
+
 Version 0.8.1
 ------------
 - ![Feature][badge-feature] Added a `@localize` macro which turns `@localize x y expr` into `let x=x, y=y; expr end` ([#142][gh-pr-142])

Project.toml

@@ -19,7 +19,7 @@ MarkdownExt = "Markdown"
 [compat]
 Aqua = "0.8"
 BangBang = "0.3.40, 0.4"
-ChunkSplitters = "3"
+ChunkSplitters = "3.1"
 Markdown = "1"
 ScopedValues = "1.3"
 StableTasks = "0.1.5"

src/implementation.jl

@@ -7,7 +7,7 @@ using OhMyThreads: Scheduler,
                    DynamicScheduler, StaticScheduler, GreedyScheduler,
                    SerialScheduler
 using OhMyThreads.Schedulers: chunking_enabled,
-                              nchunks, chunksize, chunksplit, has_chunksplit,
+                              nchunks, chunksize, chunksplit, minchunksize, has_chunksplit,
                               chunking_mode, ChunkingMode, NoChunking,
                               FixedSize, FixedCount, scheduler_from_symbol, NotGiven,
                               isgiven
@@ -25,14 +25,17 @@ function _index_chunks(sched, arg)
     C = chunking_mode(sched)
     @assert C != NoChunking
     if C == FixedCount
+        msz = isnothing(minchunksize(sched)) ? nothing : min(minchunksize(sched), length(arg))
         index_chunks(arg;
             n = nchunks(sched),
-            split = chunksplit(sched))::IndexChunks{
+            split = chunksplit(sched),
+            minsize = msz)::IndexChunks{
             typeof(arg), ChunkSplitters.Internals.FixedCount}
     elseif C == FixedSize
         index_chunks(arg;
             size = chunksize(sched),
-            split = chunksplit(sched))::IndexChunks{
+            split = chunksplit(sched),
+            minsize = minchunksize(sched))::IndexChunks{
             typeof(arg), ChunkSplitters.Internals.FixedSize}
     end
 end
@@ -67,6 +70,23 @@ function _check_chunks_incompatible_kwargs(; kwargs...)
     return nothing
 end
 
+function has_multiple_chunks(scheduler, coll)
+    C = chunking_mode(scheduler)
+    if C == NoChunking || coll isa Union{AbstractChunks, ChunkSplitters.Internals.Enumerate}
+        length(coll) > 1
+    elseif C == FixedCount
+        if isnothing(minchunksize(scheduler))
+            mcs = 1
+        else
+            mcs = max(min(minchunksize(scheduler), length(coll)), 1)
+        end
+        min(length(coll) ÷ mcs, nchunks(scheduler)) > 1
+    elseif C == FixedSize
+        length(coll) ÷ chunksize(scheduler) > 1
+    end
+end
+
+
 function tmapreduce(f, op, Arrs...;
         scheduler::MaybeScheduler = NotGiven(),
         outputtype::Type = Any,
@@ -79,7 +99,7 @@ function tmapreduce(f, op, Arrs...;
     if A isa AbstractChunks || A isa ChunkSplitters.Internals.Enumerate
         _check_chunks_incompatible_kwargs(; kwargs...)
     end
-    if _scheduler isa SerialScheduler || isempty(first(Arrs))
+    if _scheduler isa SerialScheduler || !has_multiple_chunks(_scheduler, first(Arrs))
         # empty input collection → align with Base.mapreduce behavior
         mapreduce(f, op, Arrs...; mapreduce_kwargs...)
     else
@@ -107,11 +127,12 @@ function _tmapreduce(f,
     throw_if_boxed_captures(f, op)
     if chunking_enabled(scheduler)
         tasks = map(_index_chunks(scheduler, first(Arrs))) do inds
+
             args = map(A -> view(A, inds), Arrs)
             # Note, calling `promise_task_local` here is only safe because we're assuming that
             # Base.mapreduce isn't going to magically try to do multithreading on us...
             @spawn threadpool mapreduce(promise_task_local(f), promise_task_local(op),
-                args...; $mapreduce_kwargs...)
+                                        args...; $mapreduce_kwargs...)
         end
         mapreduce(fetch, promise_task_local(op), tasks)
     else

src/schedulers.jl

@@ -85,13 +85,15 @@ struct ChunkingArgs{C, S <: Split}
     n::Int
     size::Int
     split::S
+    minsize::Union{Int, Nothing}
 end
-ChunkingArgs(::Type{NoChunking}) = ChunkingArgs{NoChunking, NoSplit}(-1, -1, NoSplit())
+ChunkingArgs(::Type{NoChunking}) = ChunkingArgs{NoChunking, NoSplit}(-1, -1, NoSplit(), nothing)
 function ChunkingArgs(
         Sched::Type{<:Scheduler},
         n::MaybeInteger,
         size::MaybeInteger,
         split::Union{Symbol, Split};
+        minsize=nothing,
         chunking
 )
     chunking || return ChunkingArgs(NoChunking)
@@ -106,7 +108,7 @@ function ChunkingArgs(
 
     chunking_mode = size > 0 ? FixedSize : FixedCount
     split = _parse_split(split)
-    result = ChunkingArgs{chunking_mode, typeof(split)}(n, size, split)
+    result = ChunkingArgs{chunking_mode, typeof(split)}(n, size, split, minsize)
 
     # argument names in error messages are those of the scheduler constructor instead
     # of ChunkingArgs because the user should not be aware of the ChunkingArgs type
@@ -124,14 +126,20 @@ chunking_mode(::ChunkingArgs{C}) where {C} = C
 has_n(ca::ChunkingArgs) = ca.n > 0
 has_size(ca::ChunkingArgs) = ca.size > 0
 has_split(::ChunkingArgs{C, S}) where {C, S} = S !== NoSplit
+has_minsize(ca::ChunkingArgs) = !isnothing(ca.minsize)
 chunking_enabled(ca::ChunkingArgs) = chunking_mode(ca) != NoChunking
 
 _chunkingstr(ca::ChunkingArgs{NoChunking}) = "none"
 function _chunkingstr(ca::ChunkingArgs{FixedCount})
-    return "fixed count ($(ca.n)), split :$(_splitid(ca.split))"
+    str = "fixed count ($(ca.n)), split :$(_splitid(ca.split))"
+    if has_minsize(ca)
+        str = str * ", minimum chunk size  $(ca.minsize)"
+    end
+    str
 end
 function _chunkingstr(ca::ChunkingArgs{FixedSize})
-    return "fixed size ($(ca.size)), split :$(_splitid(ca.split))"
+    str = "fixed size ($(ca.size)), split :$(_splitid(ca.split))"
+    str
 end
 
 # Link between a scheduler and its chunking arguments
@@ -142,10 +150,12 @@ chunking_args(::Scheduler) = ChunkingArgs(NoChunking)
 nchunks(sched::Scheduler) = chunking_args(sched).n
 chunksize(sched::Scheduler) = chunking_args(sched).size
 chunksplit(sched::Scheduler) = chunking_args(sched).split
+minchunksize(sched::Scheduler) = chunking_args(sched).minsize
 
 has_nchunks(sched::Scheduler) = has_n(chunking_args(sched))
 has_chunksize(sched::Scheduler) = has_size(chunking_args(sched))
 has_chunksplit(sched::Scheduler) = has_split(chunking_args(sched))
+has_minchunksize(sched::Scheduler) = has_minsize(chunking_args(sched))
 
 chunking_mode(sched::Scheduler) = chunking_mode(chunking_args(sched))
 chunking_enabled(sched::Scheduler) = chunking_enabled(chunking_args(sched))
@@ -180,6 +190,8 @@ with other multithreaded code.
 - `chunksize::Integer` (default not set)
     * Specifies the desired chunk size (instead of the number of chunks).
     * The options `chunksize` and `nchunks`/`ntasks` are **mutually exclusive** (only one may be a positive integer).
+- `minchunksize::Union{Integer, Nothing}` (default `nothing`)
+    * Sets a lower bound on the size of chunks. This argument takes priority over `nchunks`, so `treduce(+, 1:10; nchunks=10, minchunksize=5)` will only operate on `2` chunks for example.
 - `split::Union{Symbol, OhMyThreads.Split}` (default `OhMyThreads.Consecutive()`):
     * Determines how the collection is divided into chunks (if chunking=true). By default, each chunk consists of contiguous elements and order is maintained.
     * See [ChunkSplitters.jl](https://github.com/JuliaFolds2/ChunkSplitters.jl) for more details and available options. We also allow users to pass `:consecutive` in place of `Consecutive()`, and `:roundrobin` in place of `RoundRobin()`
@@ -209,14 +221,15 @@ function DynamicScheduler(;
         ntasks::MaybeInteger = NotGiven(), # "alias" for nchunks
         chunksize::MaybeInteger = NotGiven(),
         chunking::Bool = true,
-        split::Union{Split, Symbol} = Consecutive())
+        split::Union{Split, Symbol} = Consecutive(),
+        minchunksize::Union{Nothing, Int}=nothing)
     if isgiven(ntasks)
         if isgiven(nchunks)
             throw(ArgumentError("For the dynamic scheduler, nchunks and ntasks are aliases and only one may be provided"))
         end
         nchunks = ntasks
     end
-    ca = ChunkingArgs(DynamicScheduler, nchunks, chunksize, split; chunking)
+    ca = ChunkingArgs(DynamicScheduler, nchunks, chunksize, split; chunking, minsize=minchunksize)
     return DynamicScheduler(threadpool, ca)
 end
 from_symbol(::Val{:dynamic}) = DynamicScheduler
@@ -250,6 +263,8 @@ Isn't well composable with other multithreaded code though.
 - `chunksize::Integer` (default not set)
     * Specifies the desired chunk size (instead of the number of chunks).
     * The options `chunksize` and `nchunks`/`ntasks` are **mutually exclusive** (only one may be non-zero).
+- `minchunksize::Union{Integer, Nothing}` (default `nothing`)
+    * Sets a lower bound on the size of chunks. This argument takes priority over `nchunks`, so `treduce(+, 1:10; nchunks=10, minchunksize=5)` will only operate on `2` chunks for example.
 - `chunking::Bool` (default `true`):
     * Controls whether input elements are grouped into chunks (`true`) or not (`false`).
     * For `chunking=false`, the arguments `nchunks`/`ntasks`, `chunksize`, and `split` are ignored and input elements are regarded as "chunks" as is. Hence, there will be one parallel task spawned per input element. Note that, depending on the input, this **might spawn many(!) tasks** and can be costly!
@@ -267,14 +282,15 @@ function StaticScheduler(;
         ntasks::MaybeInteger = NotGiven(), # "alias" for nchunks
         chunksize::MaybeInteger = NotGiven(),
         chunking::Bool = true,
-        split::Union{Split, Symbol} = Consecutive())
+        split::Union{Split, Symbol} = Consecutive(),
+        minchunksize::Union{Nothing, Int} = nothing)
     if isgiven(ntasks)
         if isgiven(nchunks)
             throw(ArgumentError("For the static scheduler, nchunks and ntasks are aliases and only one may be provided"))
         end
         nchunks = ntasks
     end
-    ca = ChunkingArgs(StaticScheduler, nchunks, chunksize, split; chunking)
+    ca = ChunkingArgs(StaticScheduler, nchunks, chunksize, split; chunking, minsize=minchunksize)
     return StaticScheduler(ca)
 end
 from_symbol(::Val{:static}) = StaticScheduler
@@ -315,6 +331,8 @@ some additional overhead.
 - `chunksize::Integer` (default not set)
     * Specifies the desired chunk size (instead of the number of chunks).
     * The options `chunksize` and `nchunks` are **mutually exclusive** (only one may be a positive integer).
+- `minchunksize::Union{Integer, Nothing}` (default `nothing`)
+    * Sets a lower bound on the size of chunks. This argument takes priority over `nchunks`, so `treduce(+, 1:10; nchunks=10, minchunksize=5)` will only operate on `2` chunks for example.
 - `split::Union{Symbol, OhMyThreads.Split}` (default `OhMyThreads.RoundRobin()`):
     * Determines how the collection is divided into chunks (if chunking=true).
     * See [ChunkSplitters.jl](https://github.com/JuliaFolds2/ChunkSplitters.jl) for more details and available options. We also allow users to pass `:consecutive` in place of `Consecutive()`, and `:roundrobin` in place of `RoundRobin()`
@@ -334,11 +352,12 @@ function GreedyScheduler(;
         nchunks::MaybeInteger = NotGiven(),
         chunksize::MaybeInteger = NotGiven(),
         chunking::Bool = false,
-        split::Union{Split, Symbol} = RoundRobin())
+        split::Union{Split, Symbol} = RoundRobin(),
+        minchunksize::Union{Nothing, Int} = nothing)
     if isgiven(nchunks) || isgiven(chunksize)
         chunking = true
     end
-    ca = ChunkingArgs(GreedyScheduler, nchunks, chunksize, split; chunking)
+    ca = ChunkingArgs(GreedyScheduler, nchunks, chunksize, split; chunking, minsize=minchunksize)
     return GreedyScheduler(ntasks, ca)
 end
 from_symbol(::Val{:greedy}) = GreedyScheduler

test/runtests.jl

@@ -28,45 +28,46 @@ ChunkedGreedy(; kwargs...) = GreedyScheduler(; kwargs...)
                 SerialScheduler, ChunkedGreedy)
                 @testset for split in (Consecutive(), RoundRobin(), :consecutive, :roundrobin)
                     for nchunks in (1, 2, 6)
-                        if sched == GreedyScheduler
-                            scheduler = sched(; ntasks = nchunks)
-                        elseif sched == DynamicScheduler{OhMyThreads.Schedulers.NoChunking}
-                            scheduler = DynamicScheduler(; chunking = false)
-                        elseif sched == SerialScheduler
-                            scheduler = SerialScheduler()
-                        else
-                            scheduler = sched(; nchunks, split)
+                        for minchunksize ∈ (nothing, 1, 3)
+                            if sched == GreedyScheduler
+                                scheduler = sched(; ntasks = nchunks, minchunksize)
+                            elseif sched == DynamicScheduler{OhMyThreads.Schedulers.NoChunking}
+                                scheduler = DynamicScheduler(; chunking = false)
+                            elseif sched == SerialScheduler
+                                scheduler = SerialScheduler()
+                            else
+                                scheduler = sched(; nchunks, split, minchunksize)
+                            end
+                            kwargs = (; scheduler)
+                            if (split in (RoundRobin(), :roundrobin) ||
+                                sched ∈ (GreedyScheduler, ChunkedGreedy)) || op ∉ (vcat, *)
+                                # scatter and greedy only works for commutative operators!
+                            else
+                                mapreduce_f_op_itr = mapreduce(f, op, itrs...)
+                                @test tmapreduce(f, op, itrs...; init, kwargs...) ~ mapreduce_f_op_itr
+                                @test treducemap(op, f, itrs...; init, kwargs...) ~ mapreduce_f_op_itr
+                                @test treduce(op, f.(itrs...); init, kwargs...) ~ mapreduce_f_op_itr
+                            end
+
+                            split in (RoundRobin(), :roundrobin) && continue
+                            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)
+
+                            RT = Core.Compiler.return_type(f, Tuple{eltype.(itrs)...})
+
+                            @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
+
+                            if sched ∉ (GreedyScheduler, ChunkedGreedy)
+                                @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
+                            end
                         end
-
-                        kwargs = (; scheduler)
-                        if (split in (RoundRobin(), :roundrobin) ||
-                            sched ∈ (GreedyScheduler, ChunkedGreedy)) || op ∉ (vcat, *)
-                            # scatter and greedy only works for commutative operators!
-                        else
-                            mapreduce_f_op_itr = mapreduce(f, op, itrs...)
-                            @test tmapreduce(f, op, itrs...; init, kwargs...) ~ mapreduce_f_op_itr
-                            @test treducemap(op, f, itrs...; init, kwargs...) ~ mapreduce_f_op_itr
-                            @test treduce(op, f.(itrs...); init, kwargs...) ~ mapreduce_f_op_itr
-                        end
-
-                        split in (RoundRobin(), :roundrobin) && continue
-                        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)
-
-                        RT = Core.Compiler.return_type(f, Tuple{eltype.(itrs)...})
-
-                        @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
-
-                        if sched ∉ (GreedyScheduler, ChunkedGreedy)
-                            @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
-                        end
-                    end
+                    end 
                 end
             end
         end
@@ -185,6 +186,10 @@ end;
         @set chunking = false
         i
     end) |> isnothing
+    @test @tasks(for i in 1:4
+        @set minchunksize=2
+        i
+    end) |> isnothing
     @test_throws ArgumentError @tasks(for i in 1:3
         @set scheduler = DynamicScheduler()
         @set chunking = false
@@ -386,7 +391,9 @@ end;
     @test tmapreduce(sin, +, 1:10000; split = RoundRobin()) ≈ res_tmr
     @test tmapreduce(sin, +, 1:10000; chunksize = 2) ≈ res_tmr
     @test tmapreduce(sin, +, 1:10000; chunking = false) ≈ res_tmr
-
+    @test tmapreduce(sin, +, 1:10000; minchunksize=10) ≈ res_tmr
+    @test tmapreduce(sin, +, 1:10; minchunksize=10) == mapreduce(sin, +, 1:10)
+    
     # scheduler isa Scheduler
     @test tmapreduce(sin, +, 1:10000; scheduler = StaticScheduler()) ≈ res_tmr
     @test_throws ArgumentError tmapreduce(