add threads static to aid thread API evolution (#35646)

Author: JeffBezanson

NEWS.md

@@ -104,6 +104,9 @@ Command-line option changes
 
 Multi-threading changes
 -----------------------
+* `@threads` now allows an optional schedule argument. Use `@threads :static ...` to
+  ensure that the same schedule will be used as in past versions; the default schedule
+  is likely to change in the future.
 
 
 Build system changes

base/threadingconstructs.jl

@@ -38,7 +38,7 @@ function threading_run(func)
     end
 end
 
-function _threadsfor(iter,lbody)
+function _threadsfor(iter, lbody, schedule)
     lidx = iter.args[1]         # index
     range = iter.args[2]
     quote
@@ -82,9 +82,13 @@ function _threadsfor(iter,lbody)
             end
         end
         end
-        if threadid() != 1
-            # only thread 1 can enter/exit _threadedregion
-            Base.invokelatest(threadsfor_fun, true)
+        if 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
+              :(Base.invokelatest(threadsfor_fun, true))
+              end)
         else
             threading_run(threadsfor_fun)
         end
@@ -93,31 +97,50 @@ function _threadsfor(iter,lbody)
 end
 
 """
-    Threads.@threads
+    Threads.@threads [schedule] for ... end
 
-A macro to parallelize a for-loop to run with multiple threads. This spawns [`nthreads()`](@ref)
-number of threads, splits the iteration space amongst them, and iterates in parallel.
-A barrier is placed at the end of the loop which waits for all the threads to finish
-execution, and the loop returns.
+A macro to parallelize a `for` loop to run with multiple threads. Splits the iteration
+space among multiple tasks and runs those tasks on threads according to a scheduling
+policy.
+A barrier is placed at the end of the loop which waits for all tasks to finish
+execution.
+
+The `schedule` argument can be used to request a particular scheduling policy.
+The only currently supported value is `:static`, which creates one task per thread
+and divides the iterations equally among them. Specifying `:static` is an error
+if used from inside another `@threads` loop or from a thread other than 1.
+
+The default schedule (used when no `schedule` argument is present) is subject to change.
+
+!!! compat "Julia 1.5"
+    The `schedule` argument is available as of Julia 1.5.
 """
 macro threads(args...)
     na = length(args)
-    if na != 1
+    if na == 2
+        sched, ex = args
+        if sched isa QuoteNode
+            sched = sched.value
+        elseif sched isa Symbol
+            # for now only allow quoted symbols
+            sched = nothing
+        end
+        if sched !== :static
+            throw(ArgumentError("unsupported schedule argument in @threads"))
+        end
+    elseif na == 1
+        sched = :default
+        ex = args[1]
+    else
         throw(ArgumentError("wrong number of arguments in @threads"))
     end
-    ex = args[1]
-    if !isa(ex, Expr)
-        throw(ArgumentError("need an expression argument to @threads"))
+    if !(isa(ex, Expr) && ex.head === :for)
+        throw(ArgumentError("@threads requires a `for` loop expression"))
     end
-    if ex.head === :for
-        if ex.args[1] isa Expr && ex.args[1].head === :(=)
-            return _threadsfor(ex.args[1], ex.args[2])
-        else
-            throw(ArgumentError("nested outer loops are not currently supported by @threads"))
-        end
-    else
-        throw(ArgumentError("unrecognized argument to @threads"))
+    if !(ex.args[1] isa Expr && ex.args[1].head === :(=))
+        throw(ArgumentError("nested outer loops are not currently supported by @threads"))
     end
+    return _threadsfor(ex.args[1], ex.args[2], sched)
 end
 
 """

src/threading.c

@@ -481,6 +481,11 @@ void jl_start_threads(void)
 
 unsigned volatile _threadedregion; // HACK: keep track of whether it is safe to do IO
 
+JL_DLLEXPORT int jl_in_threaded_region(void)
+{
+    return _threadedregion != 0;
+}
+
 JL_DLLEXPORT void jl_enter_threaded_region(void)
 {
     _threadedregion += 1;

test/threads_exec.jl

@@ -729,6 +729,17 @@ let a = zeros(nthreads())
     @test a == [1:nthreads();]
 end
 
+# static schedule
+function _atthreads_static_schedule()
+    ids = zeros(Int, nthreads())
+    Threads.@threads :static for i = 1:nthreads()
+        ids[i] = Threads.threadid()
+    end
+    return ids
+end
+@test _atthreads_static_schedule() == [1:nthreads();]
+@test_throws TaskFailedException @threads for i = 1:1; _atthreads_static_schedule(); end
+
 try
     @macroexpand @threads(for i = 1:10, j = 1:10; end)
 catch ex