KernelAbstractions.jl
KernelAbstractions.jl copied to clipboard
`KernelAbstractions.@spawn backend ...` for easy Julia task based concurrency
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