Macro to make a nested loop with threads

General Usage
1

This is a reoccurring question, but with a twist that I can’t figure out (as this is my first time fiddling with macros).

I have gone through most of the threads here and on SO concerning using macro to conditionally use @threads macro on loops (and this one and this show couple of ways to get there).
But I was interested in going a step further and having something like this:

@alt_thread 4 for i in collection
    print(i)
end

and transforming it to a nested loop with threading on collection partitioned into given number of chunks. So in this case it would result in code analogous to:

@threads for j in collect(partition(collection, 4))
    for i in j
        print(i)
    end
end

It seems like it would be such a clean way to control if and how many threads are spawned for a given task.

So doing all steps without threading seem to work fine:

macro alt_thread(num, loop)
    i = esc(loop.args[1].args[1])
    iter = esc(loop.args[1].args[2])
    body = esc(loop.args[2])
    
    quote
        batch = collect(Iterators.partition($iter, $num))
        for j in batch
            for $i in j
                print("$j ")
                $body
            end
        end
    end
end

but trying to add threading breaks the macro.
When I do @macroexpand on my code I can see threading fails when called with less threads than available (it seems to be falling automatically to :static scheduling with errors as not composable in macro.
I thought about using directly _threadsfor but writing out the whole new nested body seemed way over my head for a first attempt.
Can anyone help me fix it (or tell me why it can’t be done, if that’s the case)?

Bonus points: I read the docs and watched a nice yt video about macro hygiene, but can anyone explain why I have to have $i in the quote since I overwrite it immediately? But without the dollar sign it says it is not defined.

Disclaimer: After trying to make it work for some time I have acknowledged that probably all of this is easily achievable with a custom function producing partition that could e.g. read preferences object to decide how many to make (or leaving it as one collection, which would in a sense turn off threading), but still would be nice to learn some more about macros.

1 Like