KernelInterface

Project.toml

@@ -11,6 +11,7 @@ GPUArrays = "0c68f7d7-f131-5f86-a1c3-88cf8149b2d7"
 GPUCompiler = "61eb1bfa-7361-4325-ad38-22787b887f55"
 GPUToolbox = "096a3bc2-3ced-46d0-87f4-dd12716f4bfc"
 KernelAbstractions = "63c18a36-062a-441e-b654-da1e3ab1ce7c"
+KernelInterface = "4ee993da-d684-4d17-a7dd-4e58e78d92bf"
 LLVM = "929cbde3-209d-540e-8aea-75f648917ca0"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 OpenCL_jll = "6cb37087-e8b6-5417-8430-1f242f1e46e4"
@@ -27,6 +28,7 @@ StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
 spirv2clc_jll = "f0274c0c-8c8a-59f1-85b7-f7d60330c5fb"
 
 [sources]
+KernelInterface = {url = "https://github.com/JuliaGPU/KernelAbstractions.jl", rev = "tb/ki-0.3", subdir = "lib/KernelInterface"}
 SPIRVIntrinsics = {path = "lib/intrinsics"}
 
 [compat]
@@ -35,6 +37,7 @@ GPUArrays = "11.2.1"
 GPUCompiler = "2.9"
 GPUToolbox = "3.1"
 KernelAbstractions = "0.9.38"
+KernelInterface = "0.3"
 LLVM = "9.6"
 LinearAlgebra = "1"
 OpenCL_jll = "=2024.10.24"
@@ -44,7 +47,7 @@ Random = "1"
 Random123 = "1.7.1"
 RandomNumbers = "1.6.0"
 Reexport = "1"
-SPIRVIntrinsics = "1.1"
+SPIRVIntrinsics = "1.1.3"
 SPIRV_LLVM_Backend_jll = "23"
 SPIRV_Tools_jll = "2025.1"
 StaticArrays = "1"

lib/cl/CL.jl

@@ -20,6 +20,7 @@ include("device.jl")
 include("context.jl")
 include("cmdqueue.jl")
 include("event.jl")
+include("synchronization.jl")
 include("memory.jl")
 include("program.jl")
 include("kernel.jl")

lib/cl/cmdqueue.jl

@@ -58,8 +58,10 @@ end
 
 # `check_exceptions` surfaces device-side exceptions thrown by kernels that ran on the
 # queue; disable it where throwing is not an option (e.g. finalizers), in which case the
-# exception remains pending until the next check.
-function finish(q::CmdQueue; check_exceptions::Bool=true)
+# exception remains pending until the next check. `blocking=false` waits cooperatively,
+# letting other tasks run, which is not possible in finalizers either.
+function finish(q::CmdQueue; check_exceptions::Bool=true, blocking::Bool=true)
+    blocking || wait_cooperatively(q)
     OpenCL.check_exceptions(q; rethrow=check_exceptions)
     return q
 end

lib/cl/synchronization.jl

@@ -0,0 +1,106 @@
+# cooperative synchronization
+#
+# Like CUDA.jl, busy-wait briefly for short operations, then block in the driver on a
+# separate thread, so that the waiting task yields to the Julia scheduler, and the thread
+# it runs on can do other work and take part in garbage collection.
+
+using GPUToolbox: @gcsafe_ccall
+
+# whether the command associated with `evt` has completed. errors count as completion:
+# waiting for the event reports them.
+function isdone(evt::AbstractEvent)
+    status = Ref{Cint}()
+    clGetEventInfo(evt, CL_EVENT_COMMAND_EXECUTION_STATUS, sizeof(Cint), status, C_NULL)
+    return status[] <= CL_COMPLETE
+end
+
+# before waiting on another thread, which has some overhead, busy-wait for the event,
+# initially without even yielding to other tasks. returns whether the event completed.
+function spinning_wait(evt::AbstractEvent)
+    isdone(evt) && return true
+    for spins in 1:256
+        if spins <= 32
+            ccall(:jl_cpu_pause, Cvoid, ())
+            # allow the GC to run while we're spinning
+            ccall(:jl_gc_safepoint, Cvoid, ())
+        else
+            yield()
+        end
+        isdone(evt) && return true
+    end
+    return false
+end
+
+struct SyncRequest
+    event::AbstractEvent
+    status::Base.RefValue{cl_int}
+    done::Base.Event
+end
+
+const MAX_SYNC_THREADS = 4
+const sync_channels = Vector{Channel{SyncRequest}}(undef, MAX_SYNC_THREADS)
+const sync_channel_cursor = Threads.Atomic{UInt32}(1)
+const sync_channel_lock = ReentrantLock()
+
+# runs on a thread of its own, waiting for the events it's sent
+function synchronization_worker(data::Ptr{Cvoid})
+    chan = sync_channels[Int(data)]
+    while true
+        req = take!(chan)
+        GC.@preserve req begin
+            id = Ref(req.event.id)
+            req.status[] = @gcsafe_ccall libopencl.clWaitForEvents(1::cl_uint,
+                                                                   id::Ptr{cl_event})::cl_int
+        end
+        notify(req.done)
+    end
+end
+
+@noinline function sync_channel(i::Int)
+    @lock sync_channel_lock begin
+        isassigned(sync_channels, i) && return sync_channels[i]
+        chan = Channel{SyncRequest}(Inf)
+        sync_channels[i] = chan
+
+        # we don't know the size of uv_thread_t, so reserve enough space
+        tid = Ref{NTuple{32, UInt8}}(ntuple(_ -> 0x00, 32))
+        cb = @cfunction(synchronization_worker, Cvoid, (Ptr{Cvoid},))
+        err = @ccall uv_thread_create(tid::Ptr{Cvoid}, cb::Ptr{Cvoid}, Ptr{Cvoid}(i)::Ptr{Cvoid})::Cint
+        err == 0 || Base.uv_error("uv_thread_create", err)
+        err = @ccall uv_thread_detach(tid::Ptr{Cvoid})::Cint
+        err == 0 || Base.uv_error("uv_thread_detach", err)
+        return chan
+    end
+end
+
+# wait for `evt` on a worker thread, while the calling task yields
+function nonblocking_wait(evt::AbstractEvent)
+    # sticky per task, so that a task keeps using the same worker, while concurrent tasks
+    # spread over the workers
+    i = get!(task_local_storage(), :CLSyncChannel) do
+        mod1(Int(Threads.atomic_add!(sync_channel_cursor, UInt32(1))), MAX_SYNC_THREADS)
+    end::Int
+    chan = isassigned(sync_channels, i) ? sync_channels[i] : sync_channel(i)
+
+    req = SyncRequest(evt, Ref{cl_int}(CL_SUCCESS), Base.Event())
+    put!(chan, req)
+    wait(req.done)
+    req.status[] == CL_SUCCESS || throw(CLError(req.status[]))
+    return
+end
+
+"""
+    cl.wait_cooperatively(q::CmdQueue)
+
+Wait for all commands queued on `q` to complete, while letting other tasks run. This does
+not check for errors or device-side exceptions, `cl.finish` does.
+"""
+function wait_cooperatively(q::CmdQueue)
+    evt = Ref{cl_event}()
+    clEnqueueMarkerWithWaitList(q, 0, C_NULL, evt)
+    marker = Event(evt[])
+    # the marker only completes once the queue has been submitted to the device
+    clFlush(q)
+    spinning_wait(marker) || nonblocking_wait(marker)
+    return
+end

src/OpenCL.jl

@@ -12,6 +12,8 @@ using Preferences
 
 import KernelAbstractions: KernelAbstractions
 
+import KernelInterface
+
 using Core: LLVMPtr
 
 # library wrappers
@@ -49,7 +51,12 @@ include("mapreduce.jl")
 include("gpuarrays.jl")
 include("random.jl")
 
-include("OpenCLKernels.jl")
+include("OpenCLKernelsOld.jl")
 import .OpenCLKernels: OpenCLBackend
 export OpenCLBackend
+
+# KernelInterface - NOT PUBLIC. Use KernelInterface.get_backend on an CLArray to get the backend
+include("OpenCLKernels.jl")
+import .OpenCLInterface
+
 end