KernelAbstractions.jl icon indicating copy to clipboard operation
KernelAbstractions.jl copied to clipboard

`KernelAbstractions.@spawn backend ...` for easy Julia task based concurrency

Open vchuravy opened this issue 2 years ago • 0 comments

In discussion with @maleadt, we came up with the idea to solve user code like:

kernel(...)
@sync for i in 1:4
  @spawn begin
      kernel(...) # can race with previous launch on parent stream
   end
end
kernel(...) # can race with spawned kernels

by introducing a KernelAbstractions.@spawn backend construct that takes care of creating the synchronization edges between the Julia task and the GPU work on the the child&parent stream.

An sketch of an implementation is below.

struct KAClosure{F, Backend, Token}
    f::F
    backend::Backend
    token::Token
    function KATask(f::F, backend::Backend) where {F, Backend}
        token = create_token(backend, underlying_stream(backend))
        new(f, backend, token)
    end
end

function (task::KAClosure{F})() where F
    device_barrier!(underlying_stream(task.backend), task.token)
    f()
end

struct KATask{Backend}
    task::Task
    backend::Backend
end

function Base.wait(task::KATask{<:CUDABackend})
    stream = task.task.tls[:cuda]
    device_barrier!(underlying_stream(task.backend),
                    create_token(task.backend, stream))
end

# @spawn backend begin
# end

vchuravy avatar Feb 28 '23 15:02 vchuravy