add threads static to aid thread API evolution

Author: JeffBezanson

NEWS.md

@@ -97,6 +97,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

@@ -82,8 +82,8 @@ function _threadsfor(iter,lbody)
             end
         end
         end
-        if threadid() != 1
-            # only thread 1 can enter/exit _threadedregion
+        if threadid() != 1 || ccall(:jl_in_threaded_region, Cint, ()) != 0
+            # only use threads when called from thread 1, outside @threads
             Base.invokelatest(threadsfor_fun, true)
         else
             threading_run(threadsfor_fun)
@@ -93,31 +93,45 @@ 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. If called from inside another
+`@threads` loop, or from a thread other than 1, then all iterations run on the
+current thread.
+
+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
+    elseif na == 1
+        sched = :static
+        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 sched != :static
+        throw(ArgumentError("unsupported schedule argument in @threads"))
+    end
+    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])
 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,16 @@ 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();]
+
 try
     @macroexpand @threads(for i = 1:10, j = 1:10; end)
 catch ex