Small optimizations on Scheduler creation (#148)

* make ChunkingArgs a bitstype

* make DynamicScheduler a bitstype

* dispatch symbol to its scheduler type by hand to improve performance

* fix failed tests

* inline tmapreduce

* don't run boxing tests on 1 thread

* revert make ChunkingArgs a bitstype

* revert dispatch symbol to its scheduler type

* nothing instead of sentinel values for chunking arguments

* update ChunkingArgs docstring

---------

Co-authored-by: Mason Protter <[email protected]>

Author: arnaud-ma

src/implementation.jl

@@ -1,43 +1,33 @@
 module Implementation
 
 import OhMyThreads: treduce, tmapreduce, treducemap, tforeach, tmap, tmap!, tcollect
-using OhMyThreads: @spawn, @spawnat, WithTaskLocals, promise_task_local, ChannelLike, allowing_boxed_captures
+using OhMyThreads: @spawn, @spawnat, WithTaskLocals, promise_task_local, ChannelLike,
+                   allowing_boxed_captures
 using OhMyThreads.Tools: nthtid
 using OhMyThreads: Scheduler,
                    DynamicScheduler, StaticScheduler, GreedyScheduler,
                    SerialScheduler
-using OhMyThreads.Schedulers: chunking_enabled,
+using OhMyThreads.Schedulers: chunksplitter_mode, chunking_enabled,
                               nchunks, chunksize, chunksplit, minchunksize, has_chunksplit,
+                              has_minchunksize, chunkingargs_to_kwargs,
                               chunking_mode, ChunkingMode, NoChunking,
                               FixedSize, FixedCount, scheduler_from_symbol, NotGiven,
-                              isgiven
+                              isgiven, threadpool as get_threadpool
 using Base: @propagate_inbounds
 using Base.Threads: nthreads, @threads
 using BangBang: append!!
 using ChunkSplitters: ChunkSplitters, index_chunks, Consecutive
 using ChunkSplitters.Internals: AbstractChunks, IndexChunks
 
-const MaybeScheduler = Union{NotGiven, Scheduler, Symbol}
+const MaybeScheduler = Union{NotGiven, Scheduler, Symbol, Val}
 
 include("macro_impl.jl")
 
-function _index_chunks(sched, arg)
+@inline 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),
-            minsize = msz)::IndexChunks{
-            typeof(arg), ChunkSplitters.Internals.FixedCount}
-    elseif C == FixedSize
-        index_chunks(arg;
-            size = chunksize(sched),
-            split = chunksplit(sched),
-            minsize = minchunksize(sched))::IndexChunks{
-            typeof(arg), ChunkSplitters.Internals.FixedSize}
-    end
+    @assert chunking_enabled(sched)
+    kwargs = chunkingargs_to_kwargs(sched, arg)
+    return index_chunks(arg; kwargs...)::IndexChunks{typeof(arg), chunksplitter_mode(C)}
 end
 
 function _scheduler_from_userinput(scheduler::MaybeScheduler; kwargs...)
@@ -75,19 +65,23 @@ function has_multiple_chunks(scheduler, coll)
     if C == NoChunking || coll isa Union{AbstractChunks, ChunkSplitters.Internals.Enumerate}
         length(coll) > 1
     elseif C == FixedCount
-        if isnothing(minchunksize(scheduler))
+        if !has_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
+    else
+        throw(ArgumentError("Unknown chunking mode: $C."))
     end
 end
 
-
-function tmapreduce(f, op, Arrs...;
+# we can inline this function because we use @noinline on the main function
+# it can save some time in cases where we do not hit the main function (e.g. when
+# fallback to mapreduce without any threading)
+@inline function tmapreduce(f, op, Arrs...;
         scheduler::MaybeScheduler = NotGiven(),
         outputtype::Type = Any,
         init = NotGiven(),
@@ -122,7 +116,7 @@ function _tmapreduce(f,
         ::Type{OutputType},
         scheduler::DynamicScheduler,
         mapreduce_kwargs)::OutputType where {OutputType}
-    (; threadpool) = scheduler
+    threadpool = get_threadpool(scheduler)
     check_all_have_same_indices(Arrs)
     throw_if_boxed_captures(f, op)
     if chunking_enabled(scheduler)
@@ -151,7 +145,7 @@ function _tmapreduce(f,
         ::Type{OutputType},
         scheduler::DynamicScheduler,
         mapreduce_kwargs)::OutputType where {OutputType, T}
-    (; threadpool) = scheduler
+    threadpool = get_threadpool(scheduler)
     throw_if_boxed_captures(f, op)
     tasks = map(only(Arrs)) do idcs
         @spawn threadpool promise_task_local(f)(idcs)
@@ -442,7 +436,7 @@ function tmap(f,
         if chunking_enabled(_scheduler)
             if _scheduler isa DynamicScheduler
                 _scheduler = DynamicScheduler(;
-                    threadpool = _scheduler.threadpool,
+                    threadpool = threadpool(_scheduler),
                     chunking = false)
             elseif _scheduler isa StaticScheduler
                 _scheduler = StaticScheduler(; chunking = false)
@@ -468,7 +462,7 @@ function _tmap(scheduler::DynamicScheduler{NoChunking},
         f,
         A::AbstractArray,
         _Arrs::AbstractArray...;)
-    (; threadpool) = scheduler
+    threadpool = get_threadpool(scheduler)
     Arrs = (A, _Arrs...)
     throw_if_boxed_captures(f)
     tasks = map(eachindex(A)) do i
@@ -486,7 +480,7 @@ function _tmap(scheduler::DynamicScheduler{NoChunking},
         f,
         A::Union{AbstractChunks, ChunkSplitters.Internals.Enumerate},
         _Arrs::AbstractArray...)
-    (; threadpool) = scheduler
+    threadpool = get_threadpool(scheduler)
     throw_if_boxed_captures(f)
     tasks = map(A) do idcs
         @spawn threadpool promise_task_local(f)(idcs)

src/schedulers.jl

@@ -1,7 +1,7 @@
 module Schedulers
 
 using Base.Threads: nthreads
-using ChunkSplitters: Split, Consecutive, RoundRobin
+using ChunkSplitters: Split, Consecutive, RoundRobin, ChunkSplitters
 
 # Used to indicate that a keyword argument has not been set by the user.
 # We don't use Nothing because nothing maybe sometimes be a valid user input (e.g. for init)
@@ -65,70 +65,61 @@ struct NoChunking <: ChunkingMode end
 struct FixedCount <: ChunkingMode end
 struct FixedSize <: ChunkingMode end
 
+chunksplitter_mode(::Type{FixedCount}) = ChunkSplitters.Internals.FixedCount
+chunksplitter_mode(::Type{FixedSize}) = ChunkSplitters.Internals.FixedSize
+
 """
-    ChunkingArgs{C, S <: Split}(n::Int, size::Int, split::S)
-    ChunkingArgs(Sched::Type{<:Scheduler}, n::MaybeInteger, size::MaybeInteger, split::Union{Symbol, Split}; chunking)
+    ChunkingArgs{C, S <: Split}(n::Union{Int, Nothing}, size::Union{Int, Nothing}, minsize::Union{Int, Nothing}, split::S)
+    ChunkingArgs(Sched::Type{<:Scheduler}; n = nothing, size = nothing, minsize = nothing, split::Union{Symbol, Split}; chunking)
 
 Stores all the information needed for chunking. The type parameter `C` is the chunking mode
-(`NoChunking`, `FixedSize`, or `FixedCount`).
-
-`MaybeInteger` arguments are arguments that can be `NotGiven`. If it is the case, the
-constructor automatically throws errors or gives defaults values while taking into account
-the kind of scheduler (provided by `Sched`, e.g. `DynamicScheduler`). The `chunking` keyword
-argument is a boolean and if true, everything is skipped and `C = NoChunking`.
+(`NoChunking`, `FixedSize`, or `FixedCount`). The `chunking` keyword argument is a boolean
+and if true, everything is skipped and `C = NoChunking`.
 
 Once the object is created, use the `has_fieldname(object)` function (e.g. `has_size(object)`)
-to know if the field is effectively used, since it is no longer
-`NotGiven` for type stability.
+to know if the field is effectively used.
 """
 struct ChunkingArgs{C, S <: Split}
-    n::Int
-    size::Int
-    split::S
+    n::Union{Int, Nothing}
+    size::Union{Int, Nothing}
     minsize::Union{Int, Nothing}
+    split::S
+end
+function ChunkingArgs(::Type{NoChunking})
+    ChunkingArgs{NoChunking, NoSplit}(nothing, nothing, nothing, NoSplit())
 end
-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,
+        Sched::Type{<:Scheduler};
+        n = nothing,
+        size = nothing,
+        minsize = nothing,
+        split::Union{Symbol, Split},
         chunking
 )
     chunking || return ChunkingArgs(NoChunking)
 
-    if !isgiven(n) && !isgiven(size)
+    if isnothing(n) && isnothing(size)
         n = default_nchunks(Sched)
-        size = -1
-    else
-        n = isgiven(n) ? n : -1
-        size = isgiven(size) ? size : -1
+    elseif !isnothing(n) && !isnothing(size)
+        throw(ArgumentError("nchunks and chunksize are mutually exclusive"))
     end
-
-    chunking_mode = size > 0 ? FixedSize : FixedCount
+    chunking_mode = isnothing(n) ? FixedSize : FixedCount
     split = _parse_split(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
-    # (e.g. `nchunks` instead of `n`)
-    if !(has_n(result) || has_size(result))
-        throw(ArgumentError("Either `nchunks` or `chunksize` must be a positive integer (or chunking=false)."))
-    end
-    if has_n(result) && has_size(result)
-        throw(ArgumentError("`nchunks` and `chunksize` are mutually exclusive and only one of them may be a positive integer"))
-    end
-    return result
+    return ChunkingArgs{chunking_mode, typeof(split)}(n, size, minsize, split)
 end
 
 chunking_mode(::ChunkingArgs{C}) where {C} = C
-has_n(ca::ChunkingArgs) = ca.n > 0
-has_size(ca::ChunkingArgs) = ca.size > 0
+has_n(ca::ChunkingArgs) = !isnothing(ca.n)
+has_size(ca::ChunkingArgs) = !isnothing(ca.size)
 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
 
+function chunkingargs_to_kwargs(ca::ChunkingArgs, arg)
+    minsize = !has_minsize(ca) ? nothing : min(ca.minsize, length(arg))
+    return (; ca.n, ca.size, minsize, ca.split)
+end
+
 _chunkingstr(ca::ChunkingArgs{NoChunking}) = "none"
 function _chunkingstr(ca::ChunkingArgs{FixedCount})
     str = "fixed count ($(ca.n)), split :$(_splitid(ca.split))"
@@ -157,6 +148,10 @@ 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))
 
+function chunkingargs_to_kwargs(sched::Scheduler, arg)
+    chunkingargs_to_kwargs(chunking_args(sched), arg)
+end
+
 chunking_mode(sched::Scheduler) = chunking_mode(chunking_args(sched))
 chunking_enabled(sched::Scheduler) = chunking_enabled(chunking_args(sched))
 _chunkingstr(sched::Scheduler) = _chunkingstr(chunking_args(sched))
@@ -203,43 +198,45 @@ with other multithreaded code.
     * Possible options are `:default` and `:interactive`.
     * The high-priority pool `:interactive` should be used very carefully since tasks on this threadpool should not be allowed to run for a long time without `yield`ing as it can interfere with [heartbeat](https://en.wikipedia.org/wiki/Heartbeat_(computing)) processes.
 """
-struct DynamicScheduler{C <: ChunkingMode, S <: Split} <: Scheduler
-    threadpool::Symbol
+struct DynamicScheduler{C <: ChunkingMode, S <: Split, threadpool} <: Scheduler
     chunking_args::ChunkingArgs{C, S}
 
     function DynamicScheduler(threadpool::Symbol, ca::ChunkingArgs)
         if !(threadpool in (:default, :interactive))
             throw(ArgumentError("threadpool must be either :default or :interactive"))
         end
-        new{chunking_mode(ca), typeof(ca.split)}(threadpool, ca)
+        new{chunking_mode(ca), typeof(ca.split), threadpool}(ca)
     end
 end
 
 function DynamicScheduler(;
         threadpool::Symbol = :default,
-        nchunks::MaybeInteger = NotGiven(),
-        ntasks::MaybeInteger = NotGiven(), # "alias" for nchunks
-        chunksize::MaybeInteger = NotGiven(),
-        chunking::Bool = true,
+        nchunks = nothing,
+        ntasks = nothing, # "alias" for nchunks
+        chunksize = nothing,
         split::Union{Split, Symbol} = Consecutive(),
-        minchunksize::Union{Nothing, Int}=nothing)
-    if isgiven(ntasks)
-        if isgiven(nchunks)
+        minchunksize = nothing,
+        chunking::Bool = true
+)
+    if !isnothing(ntasks)
+        if !isnothing(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, minsize=minchunksize)
+    ca = ChunkingArgs(DynamicScheduler;
+        n = nchunks, size = chunksize, minsize = minchunksize, split, chunking)
     return DynamicScheduler(threadpool, ca)
 end
 from_symbol(::Val{:dynamic}) = DynamicScheduler
 chunking_args(sched::DynamicScheduler) = sched.chunking_args
+threadpool(::DynamicScheduler{C, S, T}) where {C, S, T} = T
 
 function Base.show(io::IO, mime::MIME{Symbol("text/plain")}, s::DynamicScheduler)
     print(io, "DynamicScheduler", "\n")
     cstr = _chunkingstr(s.chunking_args)
     println(io, "├ Chunking: ", cstr)
-    print(io, "└ Threadpool: ", s.threadpool)
+    print(io, "└ Threadpool: ", threadpool(s))
 end
 
 """
@@ -278,19 +275,21 @@ struct StaticScheduler{C <: ChunkingMode, S <: Split} <: Scheduler
 end
 
 function StaticScheduler(;
-        nchunks::MaybeInteger = NotGiven(),
-        ntasks::MaybeInteger = NotGiven(), # "alias" for nchunks
-        chunksize::MaybeInteger = NotGiven(),
-        chunking::Bool = true,
+        nchunks = nothing,
+        ntasks = nothing, # "alias" for nchunks
+        chunksize = nothing,
+        minchunksize = nothing,
         split::Union{Split, Symbol} = Consecutive(),
-        minchunksize::Union{Nothing, Int} = nothing)
-    if isgiven(ntasks)
-        if isgiven(nchunks)
+        chunking::Bool = true
+)
+    if !isnothing(ntasks)
+        if !isnothing(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, minsize=minchunksize)
+    ca = ChunkingArgs(StaticScheduler;
+        n = nchunks, size = chunksize, minsize = minchunksize, split, chunking)
     return StaticScheduler(ca)
 end
 from_symbol(::Val{:static}) = StaticScheduler
@@ -349,15 +348,17 @@ end
 
 function GreedyScheduler(;
         ntasks::Integer = nthreads(),
-        nchunks::MaybeInteger = NotGiven(),
-        chunksize::MaybeInteger = NotGiven(),
-        chunking::Bool = false,
+        nchunks = nothing,
+        chunksize = nothing,
+        minchunksize = nothing,
         split::Union{Split, Symbol} = RoundRobin(),
-        minchunksize::Union{Nothing, Int} = nothing)
-    if isgiven(nchunks) || isgiven(chunksize)
+        chunking::Bool = false
+)
+    if !(isnothing(nchunks) && isnothing(chunksize))
         chunking = true
     end
-    ca = ChunkingArgs(GreedyScheduler, nchunks, chunksize, split; chunking, minsize=minchunksize)
+    ca = ChunkingArgs(GreedyScheduler;
+        n = nchunks, size = chunksize, minsize = minchunksize, split, chunking)
     return GreedyScheduler(ntasks, ca)
 end
 from_symbol(::Val{:greedy}) = GreedyScheduler