I have a function with a for-loop that I like to turn multi- or single-threaded by a separate argument. For example:
function foo(cond)
if cond
@inbounds Threads.@threads for n = 1:N
# computations
end
else
@inbounds @simd for n = 1:N
# same computations as above
end
end
end
To prevent duplicating code I wrote a macro (note that I’m very new to the world of macros):
# use threaded for loop macro
macro usethreads(cond::Bool, expr)
@assert expr.head == :for
if cond
quote
@inbounds Threads.@threads for $(esc(expr.args[1].args[1])) = $(esc(expr.args[1].args[2]))
$(esc(expr.args[2]))
end
end
else
quote
@inbounds @simd for $(expr.args[1].args[1]) = $(esc(expr.args[1].args[2]))
$(esc(expr.args[2]))
end
end
end
end
@usethreads true for i = 1:3
println(Threads.threadid())
end
@usethreads false for i = 1:3
println(Threads.threadid())
end
My first question is if there exist any better option or preferred way to do this than a macro?
My second question is if this macro is just rubbish? For example, I could not use esc at the @inbounds @simd line in the same way as I did in the @inbounds Threads.@threads line:
...
@inbounds @simd for $(esc(expr.args[1].args[1])) = $(esc(expr.args[1].args[2]))
...
I get ERROR: LoadError: Base.SimdLoop.SimdError("simd loop index must be a symbol"). If macros is the way to go, any help polishing my first attempt is very appreciated
This won’t work. Threads.@threads creates a closure, which effectively “protects” everything inside from the @inbounds.
You should do something more like
Threads.@threads for $(esc(expr.args[1].args[1])) = $(esc(expr.args[1].args[2]))
@inbounds begin
$(esc(expr.args[2]))
end
end
EDIT:
My preferred approach to macros is to esc everything, and gensym all the symbols I add manually.
Except for very simple macros, I prefer handling hygiene manually. I find it much simpler and more intuitive.
That said, for the @simd, you’re not inserting symbols of your own, so just esc the entire expression.
You can interpolate the modules if you’re worried about someone breaking your macro by defining their own @inbounds and @simd in the scope they’re trying to use @usethreads:
julia> macro usethreads(cond::Bool, expr)
@assert expr.head == :for
if cond
quote
Threads.@threads for $(esc(expr.args[1].args[1])) = $(esc(expr.args[1].args[2]))
@inbounds begin
$(esc(expr.args[2]))
end
end
end
else
quote
$Base.@inbounds $Base.@simd for $(expr.args[1].args[1]) = $(expr.args[1].args[2])
$(expr.args[2])
end
end |> esc
end
end
@usethreads (macro with 1 method)
julia> @usethreads false for i = 1:3
println(Threads.threadid())
end
1
1
1
Indeed, @threads macro can trivially be implemented as a function which composes well unlike a macro
function foreach_threaded_if(f, usethreads, xs)
if usethreads
Threads.@threads for x in xs
# @inline f(x) # requires Julia 1.8 or later
f(x)
end
else
for x in xs
# @inline f(x) # requires Julia 1.8 or later
f(x)
end
end
end
foreach_threaded_if(true, 1:3) do x
println(Threads.threadid())
end
foreach_threaded_if(false, 1:3) do x
println(Threads.threadid())
end
See also Folds.foreach that accepts an executor as an optional argument to switch single-thread and multi-thread implementations (using SequentialEx and ThreadedEx, respectively). Use FLoops.jl macro if you prefer a for loop syntax that supports the executor.