Condition to use or not to use threaded for loop

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 :slight_smile:

1 Like
@inbounds Threads.@threads

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
1 Like

Thank you very very much! I also added an additional macro to enable me to send in a variable instead of a bool:

macro usethreads(cond, expr)
    quote
        if $(esc(cond))
            @usethreads true $(expr)
        else
            @usethreads false $(expr)
        end
    end
end

Nice! Would use Threads.@threadsif myself.

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.

5 Likes