Repository: JuliaFolds2/OhMyThreads.jl
Pull request #148: Small optimizations on Scheduler creation
Author: login: arnaud-ma
Created: 2025-04-04T23:23:36Z | Merged: 2025-04-13T20:58:38Z

Now that we can completely skip the threading part when there is a single chunk, it would be good to get as close as possible to `mapreduce` (in performance) when we are in this case.

Base  `mapreduce`:
```julia
julia> @btime mapreduce(sin, +, $(1:10))
  65.331 ns (0 allocations: 0 bytes)
```

before:
```julia
julia> @btime tmapreduce(sin, +, $(1:10); minchunksize = 100)
  101.726 ns (3 allocations: 96 bytes)
```

With this PR:
```julia
julia> @btime tmapreduce(sin, +, $(1:10); minchunksize = 100)
  67.880 ns (0 allocations: 0 bytes)
```
This is done by changing the types to be "is bits" (in this case, i.e., not heap-allocated):

- `minsize` is negative instead of `nothing` in `ChunkingArgs`
- `threadpool` is a compile time symbol in `DynamicScheduler`. So `sched.threadpool` no longer works and is replaced by `threadpool(sched)`.

By doing some profiling, we can find that the small gap that remains with `mapreduce` is the call to `div` which allows to calculate the number of chunks. So it seems difficult to go further than that.


There is still an issue if we pass the scheduler as a symbol:
```julia
julia> @btime tmapreduce(sin, +, $(1:10); scheduler = :dynamic, minchunksize = 100)
  177.354 ns (5 allocations: 144 bytes)
```

This is because the symbol is a runtime variable. So to bind it to its associated type, there will necessarily be type instability. Sometimes Julia manages to do this by taking the symbol as a constant, but it seems rather random. There are several solutions:

```julia
julia> @btime tmapreduce(sin, +, $(1:10); scheduler = Val(:dynamic), minchunksize = 100)
  67.880 ns (0 allocations: 0 bytes)

julia> @btime tmapreduce(sin, +, $(1:10); scheduler = DynamicScheduler(minchunksize = 100))
  67.887 ns (0 allocations: 0 bytes)
```

One could go even further and implement the scheduler as a positional argument instead of a keyword one, i.e. `tmapreduce(sin, +, 1:10, DynamicScheduler(minchunksize = 100))`, but i am not really sure about that. 


Diff:
diff --git a/src/implementation.jl b/src/implementation.jl
index 967e90c..e491f98 100644
--- a/src/implementation.jl
+++ b/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,7 +65,7 @@ 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)
@@ -83,11 +73,15 @@ function has_multiple_chunks(scheduler, coll)
         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)
diff --git a/src/schedulers.jl b/src/schedulers.jl
index 4ae481c..339fb63 100644
--- a/src/schedulers.jl
+++ b/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
diff --git a/test/runtests.jl b/test/runtests.jl
index e1b7703..001749d 100644
--- a/test/runtests.jl
+++ b/test/runtests.jl
@@ -67,7 +67,7 @@ ChunkedGreedy(; kwargs...) = GreedyScheduler(; kwargs...)
                                 @test tcollect(f.(itrs...); kwargs...) ~ map_f_itr
                             end
                         end
-                    end 
+                    end
                 end
             end
         end
@@ -347,10 +347,10 @@ end;
                 nchunks = 2, chunksize = 4, chunking = false)) ==
                   OhMyThreads.Schedulers.NoChunking
             @test OhMyThreads.Schedulers.chunking_mode(sched(;
-                nchunks = -2, chunksize = -4, split = :whatever, chunking = false)) ==
+                nchunks = nothing, chunksize = nothing, split = :whatever, chunking = false)) ==
                   OhMyThreads.Schedulers.NoChunking
             @test OhMyThreads.Schedulers.chunking_enabled(sched(;
-                nchunks = -2, chunksize = -4, chunking = false)) == false
+                nchunks = nothing, chunksize = nothing, chunking = false)) == false
             @test OhMyThreads.Schedulers.chunking_enabled(sched(;
                 nchunks = 2, chunksize = 4, chunking = false)) == false
         else
@@ -369,8 +369,6 @@ end;
         @test OhMyThreads.Schedulers.chunking_enabled(sched(; chunksize = 2)) == true
         @test OhMyThreads.Schedulers.chunking_enabled(sched(; nchunks = 2)) == true
         @test_throws ArgumentError sched(; nchunks = 2, chunksize = 3)
-        @test_throws ArgumentError sched(; nchunks = 0, chunksize = 0)
-        @test_throws ArgumentError sched(; nchunks = -2, chunksize = -3)
         @test_throws ArgumentError sched(; nchunks = 2, split = :whatever)
 
         let scheduler = sched(; chunksize = 2, split = :batch)
@@ -393,7 +391,7 @@ end;
     @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(
@@ -768,77 +766,79 @@ end
          └ Threadpool: default"""
 end
 
-@testset "Boxing detection and error" begin
-    let
-        f1() = tmap(1:10) do i
-            A = i
-            sleep(rand()/10)
-            A
-        end
-        f2() = tmap(1:10) do i
-            local A = i
-            sleep(rand()/10)
-            A
-        end
-
-        @test f1() == 1:10
-        @test f2() == 1:10
-    end
+if Threads.nthreads() > 1
+    @testset "Boxing detection and error" begin
+        let
+            f1() = tmap(1:10) do i
+                A = i
+                sleep(rand()/10)
+                A
+            end
+            f2() = tmap(1:10) do i
+                local A = i
+                sleep(rand()/10)
+                A
+            end
 
-    let
-        f1() = tmap(1:10) do i
-            A = i
-            sleep(rand()/10)
-            A
-        end
-        f2() = tmap(1:10) do i
-            local A = i
-            sleep(rand()/10)
-            A
+            @test f1() == 1:10
+            @test f2() == 1:10
         end
 
-        @test_throws BoxedVariableError f1()
-        @test f2() == 1:10
+        let
+            f1() = tmap(1:10) do i
+                A = i
+                sleep(rand()/10)
+                A
+            end
+            f2() = tmap(1:10) do i
+                local A = i
+                sleep(rand()/10)
+                A
+            end
 
-        A = 1 # Cause spooky action-at-a-distance by making A outer-local to the whole let block!
-    end
+            @test_throws BoxedVariableError f1()
+            @test f2() == 1:10
 
-    let
-        A = 1
-        f1() = tmap(1:10) do i
-            A = 1
+            A = 1 # Cause spooky action-at-a-distance by making A outer-local to the whole let block!
         end
-        @test_throws BoxedVariableError f1() == ones(10) # Throws even though the redefinition is 'harmless'
 
-        @allow_boxed_captures begin
-            f2() = tmap(1:10) do i
+        let
+            A = 1
+            f1() = tmap(1:10) do i
                 A = 1
             end
-            @test f2() == ones(10)
-        end
+            @test_throws BoxedVariableError f1() == ones(10) # Throws even though the redefinition is 'harmless'
 
-        # Can nest allow and disallow because they're scoped values!
-        function f3()
-            @disallow_boxed_captures begin
-                tmap(1:10) do i
-                A = 1
+            @allow_boxed_captures begin
+                f2() = tmap(1:10) do i
+                    A = 1
                 end
+                @test f2() == ones(10)
+            end
+
+            # Can nest allow and disallow because they're scoped values!
+            function f3()
+                @disallow_boxed_captures begin
+                    tmap(1:10) do i
+                    A = 1
+                    end
+                end
+            end
+            @allow_boxed_captures begin
+                @test_throws BoxedVariableError f3() == ones(10)
             end
         end
-        @allow_boxed_captures begin
-            @test_throws BoxedVariableError f3() == ones(10)
-        end
-    end
-    @testset "@localize" begin
-        A = 1
-        if false
-            A = 2
-        end
-        ## This stops A from being boxed!
-        v = @localize A tmap(1:2) do _
-            A
+        @testset "@localize" begin
+            A = 1
+            if false
+                A = 2
+            end
+            ## This stops A from being boxed!
+            v = @localize A tmap(1:2) do _
+                A
+            end
+            @test v == [1, 1]
         end
-        @test v == [1, 1]
     end
 end
 


Reviews:
user: login: MasonProtter
body: This looks good. I think I'd rather keep the `Union{Nothing, Int}` pattern rather than rely on `-1` sentinel values unless there's a good reason to do so though.

I don't think it's a huge deal either way though, so if you have a strong opinion I'm open to changing my mind. 
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-04-07T09:17:11Z
user: login: arnaud-ma
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-07T23:03:02Z
user: login: MasonProtter
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-04-08T07:17:00Z
user: login: MasonProtter
body: Nice, thanks for this!
state: APPROVED
author_association: MEMBER
submitted_at: 2025-04-13T20:58:29Z

Review comments:
diff_hunk: @@ -85,19 +97,23 @@ struct ChunkingArgs{C, S <: Split}
     n::Int
     size::Int
     split::S
-    minsize::Union{Int, Nothing}
+    minsize::Int
path: src/schedulers.jl
user: login: MasonProtter
body: I don't think this is necessary to eliminate the allocations, `Union{Int, Nothing}` is an `isbitsunion` type and so it can be handled efficiently without any allocations.

I think it's usually nicer to work with real types like `Nothing` rather than using sentinal values like `-1`. 
created_at: 2025-04-07T09:13:07Z
author_association: MEMBER
diff_hunk: @@ -45,12 +45,24 @@ abstract type Scheduler end
 
 from_symbol(::Val) = throw(ArgumentError("unkown scheduler symbol"))
 
-scheduler_from_symbol(s::Symbol; kwargs...) = scheduler_from_symbol(Val(s); kwargs...)
-function scheduler_from_symbol(v::Val; kwargs...)
+function scheduler_from_symbol(v::Val, ; kwargs...)
     sched = from_symbol(v)
     return sched(; kwargs...)
 end
 
+function scheduler_from_symbol(s::Symbol; kwargs...)
+    # we could simply use the last line
+    # but there is a high chance to have a runtime dispatch to pass from Symbol to Val
+    # so we do the dispatch by hand as much as possible
+    # There will always be a performance penalty because the output
+    # of this function is not known at compile time.
+    s == :dynamic && return DynamicScheduler(; kwargs...)
+    s == :static && return StaticScheduler(; kwargs...)
+    s == :greedy && return GreedyScheduler(; kwargs...)
+    s == :serial && return SerialScheduler(; kwargs...)
+    return scheduler_from_symbol(Val(s), kwargs...)
+end
path: src/schedulers.jl
user: login: MasonProtter
body: I think this is probably a bit more defensive than it needs to be. Have you tried just going straight to `scheduler_from_val` instead of doing the pre-defined `if-else` switching? I'm pretty sure it should constant propagate fine.
created_at: 2025-04-07T09:15:34Z
author_association: MEMBER
diff_hunk: @@ -85,19 +97,23 @@ struct ChunkingArgs{C, S <: Split}
     n::Int
     size::Int
     split::S
-    minsize::Union{Int, Nothing}
+    minsize::Int
path: src/schedulers.jl
user: login: arnaud-ma
body: You're right, I don't see any difference compared to `Union{Int, Nothing}`. Does this mean we can apply the same approach to all chunking arguments (except `split`)? This could simplify the code since `ChunkSplitters` also uses `nothing` as a default argument.
created_at: 2025-04-07T23:03:02Z
author_association: CONTRIBUTOR
diff_hunk: @@ -85,19 +97,23 @@ struct ChunkingArgs{C, S <: Split}
     n::Int
     size::Int
     split::S
-    minsize::Union{Int, Nothing}
+    minsize::Int
path: src/schedulers.jl
user: login: MasonProtter
body: Yep!
created_at: 2025-04-08T07:17:00Z
author_association: MEMBER

PR comments:
user: login: codecov-commenter
created_at: 2025-04-04T23:25:51Z
body: :warning: Please install the !['codecov app svg image'](https://github.com/codecov/engineering-team/assets/152432831/e90313f4-9d3a-4b63-8b54-cfe14e7ec20d) to ensure uploads and comments are reliably processed by Codecov.

## [Codecov](https://app.codecov.io/gh/JuliaFolds2/OhMyThreads.jl/pull/148?dropdown=coverage&src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=None) Report
:x: Patch coverage is `92.30769%` with `3 lines` in your changes missing coverage. Please review.
:white_check_mark: Project coverage is 88.29%. Comparing base ([`cba9127`](https://app.codecov.io/gh/JuliaFolds2/OhMyThreads.jl/commit/cba91272e424ebfceb2214153d3b2f67aaf90e3d?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=None)) to head ([`23a5e74`](https://app.codecov.io/gh/JuliaFolds2/OhMyThreads.jl/commit/23a5e745bf74bda2b888cfaa9dc74decae50cd47?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=None)).
:warning: Report is 150 commits behind head on master.

| [Files with missing lines](https://app.codecov.io/gh/JuliaFolds2/OhMyThreads.jl/pull/148?dropdown=coverage&src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=None) | Patch % | Lines |
|---|---|---|
| [src/schedulers.jl](https://app.codecov.io/gh/JuliaFolds2/OhMyThreads.jl/pull/148?src=pr&el=tree&filepath=src%2Fschedulers.jl&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=None#diff-c3JjL3NjaGVkdWxlcnMuamw=) | 92.85% | [2 Missing :warning: ](https://app.codecov.io/gh/JuliaFolds2/OhMyThreads.jl/pull/148?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=None) |
| [src/implementation.jl](https://app.codecov.io/gh/JuliaFolds2/OhMyThreads.jl/pull/148?src=pr&el=tree&filepath=src%2Fimplementation.jl&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=None#diff-c3JjL2ltcGxlbWVudGF0aW9uLmps) | 90.90% | [1 Missing :warning: ](https://app.codecov.io/gh/JuliaFolds2/OhMyThreads.jl/pull/148?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=None) |
:exclamation: Your organization needs to install the [Codecov GitHub app](https://github.com/apps/codecov/installations/select_target) to enable full functionality.

<details><summary>Additional details and impacted files</summary>



```diff
@@            Coverage Diff             @@
##           master     #148      +/-   ##
==========================================
- Coverage   90.24%   88.29%   -1.96%     
==========================================
  Files           3        8       +5     
  Lines          82      632     +550     
==========================================
+ Hits           74      558     +484     
- Misses          8       74      +66     
```
</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/JuliaFolds2/OhMyThreads.jl/pull/148?dropdown=coverage&src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=None).   
:loudspeaker: Have feedback on the report? [Share it here](https://about.codecov.io/codecov-pr-comment-feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=None).
<details><summary> :rocket: New features to boost your workflow: </summary>

- :snowflake: [Test Analytics](https://docs.codecov.com/docs/test-analytics): Detect flaky tests, report on failures, and find test suite problems.
</details>
author_association: NONE
user: login: arnaud-ma
created_at: 2025-04-08T00:49:01Z
body: I just realized that simply inlining tmapreduce makes everything constant and there is almost nothing left to do
author_association: CONTRIBUTOR