Task id as a local variable

[fredrikekre]

I think it would be useful to give the user access to the task id (i.e. 1:ntasks:). My motivation for this is code that have this pattern (which is pretty common):

for i in 1:N
    @tasks for j in 1:N
        @local cache = computation_with_heavy_allocs()
    end
end

where there currently isn't a great way to hoist the task local cache allocation out of the outer loop. If I can observe the task id I can do something like

caches = [computation_with_heavy_allocs() for _ in 1:ntasks]
for i in 1:N
    @tasks for j in 1:N
        @local cache = caches[OhMyThreads.taskid()] # or @taskid
    end
end

Is there a better way to achieve this?

[MasonProtter]

Yeah, your use-case here seems to be one of the biggest things that OhMyThreads has been missing.

To be honest, the main way I have been recommending that people deal with this is by just using ChunkSplitters.jl like so:

using ChunkSplitters
using Base.Threads: @spawn
caches = [computation_with_heavy_allocs() for _ in 1:ntasks]
for i in 1:N
    chunk = chunks(1:N; n=ntasks)
    @sync for (idx, chunk) in enumerate(chunks(1:N; n=ntasks))
        @spawn begin
            cache = caches[idx]
            for i in chunk
                # computations using cache
            end
        end
    end
end

but I really do want OhMyThreads to be able to address this use-case and provide a more compact syntax.

The problem I have is that I'm hesitant to expose something like a task_index here because I worry it'll just create some new analogue of the problems that happen today with threadid() and I'm not confident in my ability to forsee how those problems could arise...

[fredrikekre]

Yea I understand that worry. Migrating code from threadid is part of the motivation for this. Previously it was so "easy" for people to allocate up front but with the current OhMyThreads API you kind of have to do it locally for each threaded loop.