rework range partitioning to be independent of threadid

Author: IanButterworth

base/threadingconstructs.jl

@@ -22,83 +22,67 @@ See also: `BLAS.get_num_threads` and `BLAS.set_num_threads` in the
 """
 nthreads() = Int(unsafe_load(cglobal(:jl_n_threads, Cint)))
 
-function threading_run(func, static)
-    static && ccall(:jl_enter_threaded_region, Cvoid, ())
-    n = nthreads()
-    tasks = Vector{Task}(undef, n)
-    for i = 1:n
-        t = Task(func)
-        t.sticky = static
-        static && ccall(:jl_set_task_tid, Cint, (Any, Cint), t, i-1)
-        tasks[i] = t
-        schedule(t)
-    end
-    try
-        for i = 1:n
-            wait(tasks[i])
-        end
-    finally
-        static && ccall(:jl_exit_threaded_region, Cvoid, ())
-    end
-end
-
 function _threadsfor(iter, lbody, schedule)
     lidx = iter.args[1]         # index
     range = iter.args[2]
     quote
-        local threadsfor_fun
         let range = $(esc(range))
-        function threadsfor_fun(onethread=false)
-            r = range # Load into local variable
-            lenr = length(r)
-            # divide loop iterations among threads
-            if onethread
-                tid = 1
-                len, rem = lenr, 0
+            local range_len = length(range)
+            range_len == 0 && return
+            local ntasks = if $(schedule === :dynamic)
+                nthreads()
+            elseif threadid() != 1 || ccall(:jl_in_threaded_region, Cint, ()) != 0
+                $(if schedule === :static
+                    :(error("`@threads :static` can only be used from thread 1 and not nested"))
+                else
+                    # only use threads when called from thread 1, outside @threads :static
+                    1
+                end)
             else
-                tid = threadid()
-                len, rem = divrem(lenr, nthreads())
+                nthreads()
             end
-            # not enough iterations for all the threads?
-            if len == 0
-                if tid > rem
-                    return
+            local n_per_part = div(range_len, ntasks, RoundUp)
+            if ntasks == 1
+                # run locally as a simple loop
+                for $(esc(lidx)) in range
+                    $(esc(lbody))
                 end
-                len, rem = 1, 0
-            end
-            # compute this thread's iterations
-            f = firstindex(r) + ((tid-1) * len)
-            l = f + len - 1
-            # distribute remaining iterations evenly
-            if rem > 0
-                if tid <= rem
-                    f = f + (tid-1)
-                    l = l + tid
-                else
-                    f = f + rem
-                    l = l + rem
+            else
+                local parts = Iterators.partition(range, n_per_part)
+                local nparts = length(parts)
+                $(if schedule !== :dynamic
+                    :(ccall(:jl_enter_threaded_region, Cvoid, ()))
+                end)
+                try
+                    local tasks = Vector{Task}(undef, nparts)
+                    local tid = 0
+                    for part in parts
+                        local t = @task begin
+                            local $(esc(lidx))
+                            for $(esc(lidx)) in part
+                                $(esc(lbody))
+                            end
+                        end
+                        $(if schedule === :dynamic
+                            :(t.sticky = false)
+                        else
+                            :(t.sticky = true)
+                            :(ccall(:jl_set_task_tid, Cint, (Any, Cint), t, tid))
+                        end)
+                        tasks[tid + 1] = t
+                        schedule(t)
+                        tid += 1
+                    end
+                    for i = 1:nparts
+                        wait(tasks[i])
+                    end
+                finally
+                    $(if schedule !== :dynamic
+                        :(ccall(:jl_exit_threaded_region, Cvoid, ()))
+                    end)
                 end
             end
-            # run this thread's iterations
-            for i = f:l
-                local $(esc(lidx)) = @inbounds r[i]
-                $(esc(lbody))
-            end
-        end
-        end
-        if $(schedule === :dynamic)
-            threading_run(threadsfor_fun, false)
-        elseif threadid() != 1 || ccall(:jl_in_threaded_region, Cint, ()) != 0
-            $(if schedule === :static
-              :(error("`@threads :static` can only be used from thread 1 and not nested"))
-              else
-              # only use threads when called from thread 1, outside @threads :static
-              :(Base.invokelatest(threadsfor_fun, true))
-              end)
-        else
-            threading_run(threadsfor_fun, true)
         end
-        nothing
     end
 end
 

test/threads_exec.jl

@@ -733,15 +733,19 @@ end
 @test_throws TaskFailedException @threads for i = 1:1; _atthreads_static_schedule(); end
 
 # dynamic schedule
-function _atthreads_dynamic_schedule()
+function _atthreads_dynamic_schedule(n)
     inc = Threads.Atomic{Int}(0)
-    Threads.@threads :dynamic for _ = 1:nthreads()
+    flags = falses(n)
+    Threads.@threads :dynamic for i = 1:n
         Threads.atomic_add!(inc, 1)
+        flags[i] = true
     end
-    return inc[]
+    return inc[], flags
 end
-inc = _atthreads_dynamic_schedule()
-@test inc == nthreads()
+@test _atthreads_dynamic_schedule(nthreads()) == (nthreads(), trues(nthreads()))
+@test _atthreads_dynamic_schedule(1) == (1, trues(1))
+@test _atthreads_dynamic_schedule(10) == (10, trues(10))
+@test _atthreads_dynamic_schedule(nthreads() * 2) == (nthreads() * 2, trues(nthreads() * 2))
 
 # nested dynamic schedule
 function _atthreads_dynamic_dynamic_schedule()