Repository: JuliaFolds2/OhMyThreads.jl (julia)

Issue:
Make `tmap(f, chunks(x; n))` work
Currently, one gets

```julia
julia> tmap(sin, chunks(x; n=1))
ERROR: MethodError: no method matching tmap(::typeof(sin), ::ChunkSplitters.Chunk{Vector{Float64}})

Closest candidates are:
  tmap(::Any, ::AbstractArray, AbstractArray...; scheduler, kwargs...)
   @ OhMyThreads ~/repos/OhMyThreads.jl/src/implementation.jl:120
  tmap(::Any, ::Type{T}, ::AbstractArray, ::AbstractArray...; kwargs...) where T
   @ OhMyThreads ~/repos/OhMyThreads.jl/src/implementation.jl:115

Stacktrace:
 [1] top-level scope
   @ REPL[9]:1
```

Pull request description:
* Basic support for `ChunkSplitters.Chunk` (#47)
* Deprecates `SpawnAllScheduler` in favor of `DynamicScheduler(; nchunks=0)`. I chose a deprecation here because I don't want to tag a breaking release just for this (given that we've only introduces the scheduler a few days ago, I don't think anyone is using it).
* Adds `StaticScheduler(; nchunks=0)` which for now only works if the input collection has `<= nthreads()` elements.

Closes #47

Patch:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index 3e28707d..5297e3a2 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -1,6 +1,14 @@
 OhMyThreads.jl Changelog
 =========================
 
+Version 0.4.2
+-------------
+
+- ![Feature][badge-feature] `DynamicScheduler` now supports `nchunks=0`, which turns off internal chunking entirely.
+- ![Deprecation][badge-deprecation] `SpawnAllScheduler` is now deprecated in favor of `DynamicScheduler(; nchunks=0)`.
+- ![Feature][badge-feature] Partial support for passing in a `ChunkSplitters.Chunk` when using `DynamicScheduler` (default). In this case, one should generally use `DynamicScheduler(; nchunks=0)`, i.e. turn off internal chunking.
+- ![Feature][badge-feature] `StaticScheduler` now supports `nchunks=0`, which turns off internal chunking entirely. Only works for input that has `<= nthreads()` elements.
+
 Version 0.4.1
 -------------
 
diff --git a/Project.toml b/Project.toml
index b4bbffc9..5a3bb3db 100644
--- a/Project.toml
+++ b/Project.toml
@@ -1,7 +1,7 @@
 name = "OhMyThreads"
 uuid = "67456a42-1dca-4109-a031-0a68de7e3ad5"
 authors = ["Mason Protter <mason.protter@icloud.com>"]
-version = "0.4.1"
+version = "0.4.2"
 
 [deps]
 BangBang = "198e06fe-97b7-11e9-32a5-e1d131e6ad66"
diff --git a/docs/src/refs/api.md b/docs/src/refs/api.md
index 6c5a8817..6cdd422e 100644
--- a/docs/src/refs/api.md
+++ b/docs/src/refs/api.md
@@ -28,7 +28,6 @@ Scheduler
 DynamicScheduler
 StaticScheduler
 GreedyScheduler
-SpawnAllScheduler
 ```
 
 ## Non-Exported
diff --git a/docs/src/translation.md b/docs/src/translation.md
index 435af9a2..efcfbef4 100644
--- a/docs/src/translation.md
+++ b/docs/src/translation.md
@@ -47,7 +47,7 @@ end
 
 ```julia
 # OhMyThreads
-tforeach(1:10; scheduler=SpawnAllScheduler()) do i
+tforeach(1:10; scheduler=DynamicScheduler(; nchunks=0)) do i
     println(i)
 end
 ```
diff --git a/src/OhMyThreads.jl b/src/OhMyThreads.jl
index f5571eb1..08cc697e 100644
--- a/src/OhMyThreads.jl
+++ b/src/OhMyThreads.jl
@@ -198,7 +198,7 @@ function tcollect end
 include("tools.jl")
 include("schedulers.jl")
 using .Schedulers: Scheduler,
-    DynamicScheduler, StaticScheduler, GreedyScheduler, SpawnAllScheduler
+    DynamicScheduler, StaticScheduler, GreedyScheduler, SpawnAllScheduler, chunking_enabled
 include("implementation.jl")
 
 export treduce, tmapreduce, treducemap, tmap, tmap!, tforeach, tcollect
diff --git a/src/implementation.jl b/src/implementation.jl
index d635b2e6..8dd03d4a 100644
--- a/src/implementation.jl
+++ b/src/implementation.jl
@@ -5,12 +5,21 @@ import OhMyThreads: treduce, tmapreduce, treducemap, tforeach, tmap, tmap!, tcol
 using OhMyThreads: StableTasks, chunks, @spawn, @spawnat
 using OhMyThreads.Tools: nthtid
 using OhMyThreads: Scheduler,
-    DynamicScheduler, StaticScheduler, GreedyScheduler, SpawnAllScheduler
+    chunking_enabled, DynamicScheduler, StaticScheduler, GreedyScheduler
 using Base: @propagate_inbounds
 using Base.Threads: nthreads, @threads
 
 using BangBang: BangBang, append!!
 
+using ChunkSplitters: ChunkSplitters
+
+function auto_disable_chunking_warning()
+    @warn("You passed in a `ChunkSplitters.Chunk` but also a scheduler that has "*
+          "chunking enabled. Will turn off internal chunking to proceed.\n"*
+          "To avoid this warning, try to turn off chunking (`nchunks=0`) "*
+          "or pass in `collect(chunks(...))`.")
+end
+
 function tmapreduce(f, op, Arrs...;
         scheduler::Scheduler = DynamicScheduler(),
         outputtype::Type = Any,
@@ -28,7 +37,7 @@ end
 
 treducemap(op, f, A...; kwargs...) = tmapreduce(f, op, A...; kwargs...)
 
-# DynamicScheduler
+# DynamicScheduler: AbstractArray/Generic
 function _tmapreduce(f,
         op,
         Arrs,
@@ -37,25 +46,85 @@ function _tmapreduce(f,
         mapreduce_kwargs)::OutputType where {OutputType}
     (; nchunks, split, threadpool) = scheduler
     check_all_have_same_indices(Arrs)
-    tasks = map(chunks(first(Arrs); n = nchunks, split)) do inds
-        args = map(A -> view(A, inds), Arrs)
-        @spawn threadpool mapreduce(f, op, args...; $mapreduce_kwargs...)
+    if chunking_enabled(scheduler)
+        tasks = map(chunks(first(Arrs); n = nchunks, split)) do inds
+            args = map(A -> view(A, inds), Arrs)
+            @spawn threadpool mapreduce(f, op, args...; $mapreduce_kwargs...)
+        end
+        mapreduce(fetch, op, tasks)
+    else
+        tasks = map(eachindex(first(Arrs))) do i
+            args = map(A -> @inbounds(A[i]), Arrs)
+            @spawn threadpool f(args...)
+        end
+        mapreduce(fetch, op, tasks; mapreduce_kwargs...)
     end
-    mapreduce(fetch, op, tasks)
 end
 
-# SpawnAllScheduler
+# DynamicScheduler: ChunkSplitters.Chunk
 function _tmapreduce(f,
         op,
-        Arrs,
+        Arrs::Tuple{ChunkSplitters.Chunk{Vector{Float64}}}, # we don't support multiple chunks for now
         ::Type{OutputType},
-        scheduler::SpawnAllScheduler,
+        scheduler::DynamicScheduler,
         mapreduce_kwargs)::OutputType where {OutputType}
-    (; threadpool) = scheduler
+    (; nchunks, split, threadpool) = scheduler
+    chunking_enabled(scheduler) && auto_disable_chunking_warning()
+    tasks = map(only(Arrs)) do idcs
+        @spawn threadpool f(idcs)
+    end
+    mapreduce(fetch, op, tasks; mapreduce_kwargs...)
+end
+
+# StaticScheduler: AbstractArray/Generic
+function _tmapreduce(f,
+        op,
+        Arrs,
+        ::Type{OutputType},
+        scheduler::StaticScheduler,
+        mapreduce_kwargs) where {OutputType}
+    (; nchunks, split) = scheduler
     check_all_have_same_indices(Arrs)
-    tasks = map(eachindex(first(Arrs))) do i
-        args = map(A -> @inbounds(A[i]), Arrs)
-        @spawn threadpool f(args...)
+    if chunking_enabled(scheduler)
+        n = min(nthreads(), nchunks) # We could implement strategies, like round-robin, in the future
+        tasks = map(enumerate(chunks(first(Arrs); n, split))) do (c, inds)
+            tid = @inbounds nthtid(c)
+            args = map(A -> view(A, inds), Arrs)
+            @spawnat tid mapreduce(f, op, args...; mapreduce_kwargs...)
+        end
+        mapreduce(fetch, op, tasks)
+    else
+        if length(first(Arrs)) > nthreads()
+            error("You have disabled chunking but provided an input with more then " *
+                  "`nthreads()` elements. This is not supported for `StaticScheduler`.")
+        end
+        n = min(nthreads(), nchunks) # We could implement strategies, like round-robin, in the future
+        tasks = map(enumerate(eachindex(first(Arrs)))) do (c, i)
+            tid = @inbounds nthtid(c)
+            args = map(A -> @inbounds(A[i]), Arrs)
+            @spawnat tid f(args...)
+        end
+        mapreduce(fetch, op, tasks; mapreduce_kwargs...)
+    end
+end
+
+# StaticScheduler: ChunkSplitters.Chunk
+function _tmapreduce(f,
+        op,
+        Arrs::Tuple{ChunkSplitters.Chunk{Vector{Float64}}}, # we don't support multiple chunks for now
+        ::Type{OutputType},
+        scheduler::StaticScheduler,
+        mapreduce_kwargs) where {OutputType}
+    chunking_enabled(scheduler) && auto_disable_chunking_warning()
+    check_all_have_same_indices(Arrs)
+    chnks = only(Arrs)
+    if length(chnks) > nthreads()
+        error("You provided a `ChunkSplitters.Chunk` with more than `nthreads()` chunks " *
+              "as input, which is not supported by the `StaticScheduler`.")
+    end
+    tasks = map(enumerate(chnks)) do (c, idcs)
+        tid = @inbounds nthtid(c)
+        @spawnat tid f(idcs)
     end
     mapreduce(fetch, op, tasks; mapreduce_kwargs...)
 end
@@ -89,24 +158,6 @@ function _tmapreduce(f,
     mapreduce(fetch, op, tasks; mapreduce_kwargs...)
 end
 
-# StaticScheduler
-function _tmapreduce(f,
-        op,
-        Arrs,
-        ::Type{OutputType},
-        scheduler::StaticScheduler,
-        mapreduce_kwargs) where {OutputType}
-    (; nchunks, split) = scheduler
-    check_all_have_same_indices(Arrs)
-    n = min(nthreads(), nchunks) # We could implement strategies, like round-robin, in the future
-    tasks = map(enumerate(chunks(first(Arrs); n, split))) do (c, inds)
-        tid = @inbounds nthtid(c)
-        args = map(A -> view(A, inds), Arrs)
-        @spawnat tid mapreduce(f, op, args...; mapreduce_kwargs...)
-    end
-    mapreduce(fetch, op, tasks)
-end
-
 function check_all_have_same_indices(Arrs)
     let A = first(Arrs), Arrs = Arrs[2:end]
         if !all(B -> eachindex(A) == eachindex(B), Arrs)
@@ -135,40 +186,83 @@ function tmap(f, ::Type{T}, A::AbstractArray, _Arrs::AbstractArray...; kwargs...
 end
 
 function tmap(f,
-        A::AbstractArray,
+        A::Union{AbstractArray, ChunkSplitters.Chunk},
         _Arrs::AbstractArray...;
         scheduler::Scheduler = DynamicScheduler(),
         kwargs...)
     if scheduler isa GreedyScheduler
         error("Greedy scheduler isn't supported with `tmap` unless you provide an `OutputElementType` argument, since the greedy schedule requires a commutative reducing operator.")
-    elseif hasfield(typeof(scheduler), :split) && scheduler.split != :batch
+    end
+    if chunking_enabled(scheduler) && hasfield(typeof(scheduler), :split) &&
+       scheduler.split != :batch
         error("Only `split == :batch` is supported because the parallel operation isn't commutative. (Scheduler: $scheduler)")
     end
+    if A isa ChunkSplitters.Chunk && chunking_enabled(scheduler)
+        auto_disable_chunking_warning()
+        if scheduler isa DynamicScheduler
+            scheduler = DynamicScheduler(; nchunks = 0, scheduler.threadpool)
+        elseif scheduler isa StaticScheduler
+            scheduler = StaticScheduler(; nchunks = 0)
+        else
+            error("Can't disable chunking for this scheduler?! Shouldn't be reached.",
+                scheduler)
+        end
+    end
 
     Arrs = (A, _Arrs...)
     check_all_have_same_indices(Arrs)
     _tmap(scheduler, f, A, _Arrs...; kwargs...)
 end
 
-# SpawnAllScheduler
-function _tmap(scheduler::SpawnAllScheduler,
+# w/o chunking (DynamicScheduler{false}): AbstractArray
+function _tmap(scheduler::DynamicScheduler{false},
         f,
         A::AbstractArray,
         _Arrs::AbstractArray...;
         kwargs...)
     (; threadpool) = scheduler
     Arrs = (A, _Arrs...)
-    ts = map(eachindex(A)) do i
+    tasks = map(eachindex(A)) do i
         @spawn threadpool begin
             args = map(A -> A[i], Arrs)
             f(args...)
         end
     end
-    v = map(fetch, ts)
+    v = map(fetch, tasks)
     reshape(v, size(A)...)
 end
 
-# All other schedulers
+# w/o chunking (DynamicScheduler{false}): ChunkSplitters.Chunk
+function _tmap(scheduler::DynamicScheduler{false},
+        f,
+        A::ChunkSplitters.Chunk,
+        _Arrs::AbstractArray...;
+        kwargs...)
+    (; threadpool) = scheduler
+    tasks = map(A) do idcs
+        @spawn threadpool f(idcs)
+    end
+    map(fetch, tasks)
+end
+
+# w/o chunking (StaticScheduler{false}): ChunkSplitters.Chunk
+function _tmap(scheduler::StaticScheduler{false},
+        f,
+        A::ChunkSplitters.Chunk,
+        _Arrs::AbstractArray...;
+        kwargs...)
+    if length(A) > nthreads()
+        error("You provided a `ChunkSplitters.Chunk` with more than `nthreads()` chunks " *
+              "as input, which is not supported by the `StaticScheduler`.")
+    end
+    tasks = map(enumerate(A)) do (c, idcs)
+        tid = @inbounds nthtid(c)
+        @spawnat tid f(idcs)
+    end
+    map(fetch, tasks)
+end
+
+# w/ chunking
 function _tmap(scheduler::Scheduler,
         f,
         A::AbstractArray,
diff --git a/src/schedulers.jl b/src/schedulers.jl
index eab3752c..2ac00197 100644
--- a/src/schedulers.jl
+++ b/src/schedulers.jl
@@ -8,7 +8,6 @@ Supertype for all available schedulers:
 * [`DynamicScheduler`](@ref): default dynamic scheduler
 * [`StaticScheduler`](@ref): low-overhead static scheduler
 * [`GreedyScheduler`](@ref): greedy load-balancing scheduler
-* [`SpawnAllScheduler`](@ref): `@spawn` one task per element
 """
 abstract type Scheduler end
 
@@ -23,10 +22,12 @@ with other multithreaded code.
 
 ## Keyword arguments:
 
-- `nchunks::Int` (default `2 * nthreads(threadpool)`):
+- `nchunks::Integer` (default `2 * nthreads(threadpool)`):
     * Determines the number of chunks (and thus also the number of parallel tasks).
     * Increasing `nchunks` can help with [load balancing](https://en.wikipedia.org/wiki/Load_balancing_(computing)), but at the expense of creating more overhead. For `nchunks <= nthreads()` there are not enough chunks for any load balancing.
     * Setting `nchunks < nthreads()` is an effective way to use only a subset of the available threads.
+    * Setting `nchunks = 0` turns off the internal chunking entirely (a task is spawned for each element). Note that, depending on the input, this scheduler **might spawn many(!) tasks** and can be
+    very costly!
 - `split::Symbol` (default `:batch`):
     * Determines how the collection is divided into chunks. By default, each chunk consists of contiguous elements.
     * See [ChunkSplitters.jl](https://github.com/JuliaFolds2/ChunkSplitters.jl) for more details and available options.
@@ -34,19 +35,23 @@ 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.
 """
-Base.@kwdef struct DynamicScheduler <: Scheduler
+Base.@kwdef struct DynamicScheduler{C} <: Scheduler
     threadpool::Symbol = :default
     nchunks::Int = 2 * nthreads(threadpool) # a multiple of nthreads to enable load balancing
     split::Symbol = :batch
 
-    function DynamicScheduler(threadpool::Symbol, nchunks::Int, split::Symbol)
+    function DynamicScheduler(threadpool::Symbol, nchunks::Integer, split::Symbol)
         threadpool in (:default, :interactive) ||
             throw(ArgumentError("threadpool must be either :default or :interactive"))
-        nchunks > 0 || throw(ArgumentError("nchunks must be a positive integer"))
-        new(threadpool, nchunks, split)
+        nchunks >= 0 ||
+            throw(ArgumentError("nchunks must be a positive integer (or zero)."))
+        C = !(nchunks == 0) # type parameter indicates whether chunking is enabled
+        new{C}(threadpool, nchunks, split)
     end
 end
 
+chunking_enabled(::DynamicScheduler{C}) where {C} = C
+
 """
 A static low-overhead scheduler. Divides the given collection into chunks and
 then spawns a task per chunk to perform the requested operation in parallel.
@@ -59,7 +64,7 @@ Isn't well composable with other multithreaded code though.
 
 ## Keyword arguments:
 
-- `nchunks::Int` (default `nthreads()`):
+- `nchunks::Integer` (default `nthreads()`):
     * Determines the number of chunks (and thus also the number of parallel tasks).
     * Setting `nchunks < nthreads()` is an effective way to use only a subset of the available threads.
     * Currently, `nchunks > nthreads()` **isn't officialy supported** but, for now, will fall back to `nchunks = nthreads()`.
@@ -67,16 +72,20 @@ Isn't well composable with other multithreaded code though.
     * Determines how the collection is divided into chunks. By default, each chunk consists of contiguous elements.
     * See [ChunkSplitters.jl](https://github.com/JuliaFolds2/ChunkSplitters.jl) for more details and available options.
 """
-Base.@kwdef struct StaticScheduler <: Scheduler
+Base.@kwdef struct StaticScheduler{C} <: Scheduler
     nchunks::Int = nthreads()
     split::Symbol = :batch
 
-    function StaticScheduler(nchunks::Int, split::Symbol)
-        nchunks > 0 || throw(ArgumentError("nchunks must be a positive integer"))
-        new(nchunks, split)
+    function StaticScheduler(nchunks::Integer, split::Symbol)
+        nchunks >= 0 ||
+            throw(ArgumentError("nchunks must be a positive integer (or zero)."))
+        C = !(nchunks == 0) # type parameter indicates whether chunking is enabled
+        new{C}(nchunks, split)
     end
 end
 
+chunking_enabled(::StaticScheduler{C}) where {C} = C
+
 """
 A greedy dynamic scheduler. The elements of the collection are first put into a `Channel`
 and then dynamic, non-sticky tasks are spawned to process channel content in parallel.
@@ -103,31 +112,10 @@ Base.@kwdef struct GreedyScheduler <: Scheduler
     end
 end
 
-"""
-A scheduler that spawns a task per element (i.e. there is no internal chunking) to perform
-the requested operation in parallel. The tasks are assigned to threads by Julia's dynamic
-scheduler and are non-sticky, that is, they can migrate between threads.
-
-Note that, depending on the input, this scheduler **might spawn many(!) tasks** and can be
-very costly!
-
-Essentially the same as `DynamicScheduler(; nchunks=nelements)`, but with a simpler,
-potentially more efficient implementation.
+chunking_enabled(::GreedyScheduler) = false
 
-## Keyword arguments:
-
-- `threadpool::Symbol` (default `:default`):
-    * 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.
-"""
-Base.@kwdef struct SpawnAllScheduler <: Scheduler
-    threadpool::Symbol = :default
-
-    function SpawnAllScheduler(threadpool::Symbol)
-        threadpool in (:default, :interactive) ||
-            throw(ArgumentError("threadpool must be either :default or :interactive"))
-        new(threadpool)
-    end
-end
+@deprecate SpawnAllScheduler(args...; kwargs...) DynamicScheduler(args...;
+    nchunks = 0,
+    kwargs...)
 
 end # module


Test patch:
diff --git a/test/runtests.jl b/test/runtests.jl
index 0db19233..fbc941bb 100644
--- a/test/runtests.jl
+++ b/test/runtests.jl
@@ -11,13 +11,13 @@ sets_to_test = [
 @testset "Basics" begin
     for (; ~, f, op, itrs, init) ∈ sets_to_test
         @testset "f=$f, op=$op, itrs::$(typeof(itrs))" begin
-            @testset for sched ∈ (StaticScheduler, DynamicScheduler, GreedyScheduler, SpawnAllScheduler)
+            @testset for sched ∈ (StaticScheduler, DynamicScheduler, GreedyScheduler, DynamicScheduler{false})
                 @testset for split ∈ (:batch, :scatter)
                     for nchunks ∈ (1, 2, 6)
                         if sched == GreedyScheduler
                             scheduler = sched(; ntasks=nchunks)
-                        elseif sched == SpawnAllScheduler
-                            scheduler = sched()
+                        elseif sched == DynamicScheduler{false}
+                            scheduler = DynamicScheduler(; nchunks=0)
                         else
                             scheduler = sched(; nchunks, split)
                         end
@@ -56,4 +56,18 @@ sets_to_test = [
     end
 end
 
+@testset "ChunkSplitters.Chunk" begin
+    x = rand(100)
+    chnks = OhMyThreads.chunks(x; n=Threads.nthreads())
+    for scheduler in (DynamicScheduler(; nchunks=0), StaticScheduler(; nchunks=0))
+        @testset "$scheduler" begin
+            @test tmap(x -> sin.(x), chnks; scheduler) ≈ map(x -> sin.(x), chnks)
+            @test tmapreduce(x -> sin.(x), vcat, chnks; scheduler) ≈ mapreduce(x -> sin.(x), vcat, chnks)
+            @test tcollect(chnks; scheduler) == collect(chnks)
+            @test treduce(vcat, chnks; scheduler) == reduce(vcat, chnks)
+            @test isnothing(tforeach(x -> sin.(x), chnks; scheduler))
+        end
+    end
+end
+
 # Todo way more testing, and easier tests to deal with