POCL: wait for kernels cooperatively

Project.toml

@@ -10,6 +10,7 @@ projects = ["test", "lib/KernelInterface", "lib/KernelInterface/test"]
 Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"
 Atomix = "a9b6321e-bd34-4604-b9c9-b65b8de01458"
 GPUCompiler = "61eb1bfa-7361-4325-ad38-22787b887f55"
+GPUToolbox = "096a3bc2-3ced-46d0-87f4-dd12716f4bfc"
 KernelInterface = "4ee993da-d684-4d17-a7dd-4e58e78d92bf"
 LLVM = "929cbde3-209d-540e-8aea-75f648917ca0"
 MacroTools = "1914dd2f-81c6-5fcd-8719-6d5c9610ff09"
@@ -42,6 +43,7 @@ Adapt = "0.4, 1.0, 2.0, 3.0, 4"
 Atomix = "1.2.1"
 EnzymeCore = "0.7, 0.8.1"
 GPUCompiler = "2.7"
+GPUToolbox = "3.3.2"
 KernelInterface = "0.4"
 LLVM = "9.9"
 LinearAlgebra = "1.6"

src/pocl/backend.jl

@@ -154,13 +154,17 @@ end
 function KI.launch(obj::KI.Kernel{POCLBackend}, groups::Dims{3}, items::Dims{3}, args::Tuple)
     # the kernel only gets pointers to the arrays in `args` and captured by `f`, so keep
     # them alive until it completes. POCL launches synchronously, see the implementation
-    # note on `synchronize`
+    # note on `synchronize`. waiting for the kernel yields to other tasks, as `synchronize`
+    # should (see the documentation on its semantics).
     f = obj.kern.f
-    event = GC.@preserve f args begin
+    GC.@preserve f args begin
         event = POCL.launch_tuple(obj.kern.kernel, args; local_size = items, global_size = groups .* items)
-        wait(event)
+        try
+            wait(event)
+        finally
+            cl.clReleaseEvent(event)
+        end
     end
-    cl.clReleaseEvent(event)
     return nothing
 end
 

src/pocl/nanoOpenCL.jl

@@ -4,6 +4,8 @@ using ..POCL: platform, device, context, queue
 
 import pocl_standalone_jll
 
+using GPUToolbox: GPUToolbox, @gcsafe_ccall, cooperative_wait
+
 using Printf
 
 const libopencl = pocl_standalone_jll.libpocl
@@ -511,7 +513,7 @@ const cl_command_queue_properties = cl_bitfield
 const cl_event_info = cl_uint
 
 @checked function clGetPlatformIDs(num_entries, platforms, num_platforms)
-    @ccall libopencl.POclGetPlatformIDs(
+    @gcsafe_ccall libopencl.POclGetPlatformIDs(
         num_entries::cl_uint, platforms::Ptr{cl_platform_id},
         num_platforms::Ptr{cl_uint}
     )::cl_int
@@ -521,7 +523,7 @@ end
         platform, param_name, param_value_size, param_value,
         param_value_size_ret
     )
-    @ccall libopencl.POclGetPlatformInfo(
+    @gcsafe_ccall libopencl.POclGetPlatformInfo(
         platform::cl_platform_id,
         param_name::cl_platform_info,
         param_value_size::Csize_t, param_value::Ptr{Cvoid},
@@ -530,7 +532,7 @@ end
 end
 
 @checked function clGetDeviceIDs(platform, device_type, num_entries, devices, num_devices)
-    @ccall libopencl.POclGetDeviceIDs(
+    @gcsafe_ccall libopencl.POclGetDeviceIDs(
         platform::cl_platform_id, device_type::cl_device_type,
         num_entries::cl_uint, devices::Ptr{cl_device_id},
         num_devices::Ptr{cl_uint}
@@ -541,7 +543,7 @@ end
         device, param_name, param_value_size, param_value,
         param_value_size_ret
     )
-    @ccall libopencl.POclGetDeviceInfo(
+    @gcsafe_ccall libopencl.POclGetDeviceInfo(
         device::cl_device_id, param_name::cl_device_info,
         param_value_size::Csize_t, param_value::Ptr{Cvoid},
         param_value_size_ret::Ptr{Csize_t}
@@ -552,7 +554,7 @@ function clCreateContext(
         properties, num_devices, devices, pfn_notify, user_data,
         errcode_ret
     )
-    return @ccall libopencl.POclCreateContext(
+    return @gcsafe_ccall libopencl.POclCreateContext(
         properties::Ptr{cl_context_properties},
         num_devices::cl_uint, devices::Ptr{cl_device_id},
         pfn_notify::Ptr{Cvoid}, user_data::Ptr{Cvoid},
@@ -561,26 +563,26 @@ function clCreateContext(
 end
 
 @checked function clReleaseContext(context)
-    @ccall libopencl.POclReleaseContext(context::cl_context)::cl_int
+    @gcsafe_ccall libopencl.POclReleaseContext(context::cl_context)::cl_int
 end
 
 function clCreateProgramWithIL(context, il, length, errcode_ret)
-    return @ccall libopencl.POclCreateProgramWithIL(
+    return @gcsafe_ccall libopencl.POclCreateProgramWithIL(
         context::cl_context, il::Ptr{Cvoid},
         length::Csize_t,
         errcode_ret::Ptr{cl_int}
     )::cl_program
 end
 
 @checked function clReleaseProgram(program)
-    @ccall libopencl.POclReleaseProgram(program::cl_program)::cl_int
+    @gcsafe_ccall libopencl.POclReleaseProgram(program::cl_program)::cl_int
 end
 
 @checked function clBuildProgram(
         program, num_devices, device_list, options, pfn_notify,
         user_data
     )
-    @ccall libopencl.POclBuildProgram(
+    @gcsafe_ccall libopencl.POclBuildProgram(
         program::cl_program, num_devices::cl_uint,
         device_list::Ptr{cl_device_id}, options::Ptr{Cchar},
         pfn_notify::Ptr{Cvoid}, user_data::Ptr{Cvoid}
@@ -591,7 +593,7 @@ end
         program, param_name, param_value_size, param_value,
         param_value_size_ret
     )
-    @ccall libopencl.POclGetProgramInfo(
+    @gcsafe_ccall libopencl.POclGetProgramInfo(
         program::cl_program, param_name::cl_program_info,
         param_value_size::Csize_t, param_value::Ptr{Cvoid},
         param_value_size_ret::Ptr{Csize_t}
@@ -602,7 +604,7 @@ end
         program, device, param_name, param_value_size,
         param_value, param_value_size_ret
     )
-    @ccall libopencl.POclGetProgramBuildInfo(
+    @gcsafe_ccall libopencl.POclGetProgramBuildInfo(
         program::cl_program, device::cl_device_id,
         param_name::cl_program_build_info,
         param_value_size::Csize_t,
@@ -612,25 +614,25 @@ end
 end
 
 function clCreateKernel(program, kernel_name, errcode_ret)
-    return @ccall libopencl.POclCreateKernel(
+    return @gcsafe_ccall libopencl.POclCreateKernel(
         program::cl_program, kernel_name::Ptr{Cchar},
         errcode_ret::Ptr{cl_int}
     )::cl_kernel
 end
 
 @checked function clReleaseKernel(kernel)
-    @ccall libopencl.POclReleaseKernel(kernel::cl_kernel)::cl_int
+    @gcsafe_ccall libopencl.POclReleaseKernel(kernel::cl_kernel)::cl_int
 end
 
 @checked function clSetKernelArg(kernel, arg_index, arg_size, arg_value)
-    @ccall libopencl.POclSetKernelArg(
+    @gcsafe_ccall libopencl.POclSetKernelArg(
         kernel::cl_kernel, arg_index::cl_uint,
         arg_size::Csize_t, arg_value::Ptr{Cvoid}
     )::cl_int
 end
 
 @checked function clSetKernelArgSVMPointer(kernel, arg_index, arg_value)
-    @ccall libopencl.POclSetKernelArgSVMPointer(
+    @gcsafe_ccall libopencl.POclSetKernelArgSVMPointer(
         kernel::cl_kernel, arg_index::cl_uint,
         arg_value::Ptr{Cvoid}
     )::cl_int
@@ -640,7 +642,7 @@ end
         kernel, device, param_name, param_value_size,
         param_value, param_value_size_ret
     )
-    @ccall libopencl.POclGetKernelWorkGroupInfo(
+    @gcsafe_ccall libopencl.POclGetKernelWorkGroupInfo(
         kernel::cl_kernel, device::cl_device_id,
         param_name::cl_kernel_work_group_info,
         param_value_size::Csize_t,
@@ -654,7 +656,7 @@ end
         input_value, param_value_size, param_value,
         param_value_size_ret
     )
-    @ccall libopencl.POclGetKernelSubGroupInfo(
+    @gcsafe_ccall libopencl.POclGetKernelSubGroupInfo(
         kernel::cl_kernel, device::cl_device_id,
         param_name::cl_kernel_sub_group_info,
         input_value_size::Csize_t,
@@ -670,7 +672,7 @@ end
         command_queue, kernel, work_dim,
         global_work_size::NTuple{3, Csize_t}, local_work_size::NTuple{3, Csize_t}, event::Ref{cl_event}
     )
-    @ccall libopencl.POclEnqueueNDRangeKernel(
+    @gcsafe_ccall libopencl.POclEnqueueNDRangeKernel(
         command_queue::cl_command_queue,
         kernel::cl_kernel, work_dim::cl_uint,
         C_NULL::Ptr{Csize_t},
@@ -688,7 +690,7 @@ end
         local_work_size, num_events_in_wait_list,
         event_wait_list, event
     )
-    @ccall libopencl.POclEnqueueNDRangeKernel(
+    @gcsafe_ccall libopencl.POclEnqueueNDRangeKernel(
         command_queue::cl_command_queue,
         kernel::cl_kernel, work_dim::cl_uint,
         global_work_offset::Ptr{Csize_t},
@@ -701,38 +703,49 @@ end
 end
 
 function clCreateCommandQueue(context, device, properties, errcode_ret)
-    return @ccall libopencl.POclCreateCommandQueue(
+    return @gcsafe_ccall libopencl.POclCreateCommandQueue(
         context::cl_context, device::cl_device_id,
         properties::cl_command_queue_properties,
         errcode_ret::Ptr{cl_int}
     )::cl_command_queue
 end
 
 @checked function clReleaseCommandQueue(command_queue)
-    @ccall libopencl.POclReleaseCommandQueue(command_queue::cl_command_queue)::cl_int
+    @gcsafe_ccall libopencl.POclReleaseCommandQueue(command_queue::cl_command_queue)::cl_int
 end
 
 @checked function clFinish(command_queue)
-    @ccall libopencl.POclFinish(command_queue::cl_command_queue)::cl_int
+    @gcsafe_ccall libopencl.POclFinish(command_queue::cl_command_queue)::cl_int
+end
+
+@checked function clFlush(command_queue)
+    @gcsafe_ccall libopencl.POclFlush(command_queue::cl_command_queue)::cl_int
+end
+
+@checked function clSetEventCallback(event, command_exec_callback_type, pfn_notify, user_data)
+    @gcsafe_ccall libopencl.POclSetEventCallback(
+        event::cl_event, command_exec_callback_type::cl_int,
+        pfn_notify::Ptr{Cvoid}, user_data::Ptr{Cvoid}
+    )::cl_int
 end
 
 @checked function clWaitForEvents(num_events, event_list)
-    @ccall libopencl.POclWaitForEvents(num_events::cl_uint, event_list::Ptr{cl_event})::cl_int
+    @gcsafe_ccall libopencl.POclWaitForEvents(num_events::cl_uint, event_list::Ptr{cl_event})::cl_int
 end
 
 @checked function clGetEventInfo(
         event, param_name, param_value_size, param_value,
         param_value_size_ret
     )
-    @ccall libopencl.POclGetEventInfo(
+    @gcsafe_ccall libopencl.POclGetEventInfo(
         event::cl_event, param_name::cl_event_info,
         param_value_size::Csize_t, param_value::Ptr{Cvoid},
         param_value_size_ret::Ptr{Csize_t}
     )::cl_int
 end
 
 @checked function clReleaseEvent(event)
-    @ccall libopencl.POclReleaseEvent(event::cl_event)::cl_int
+    @gcsafe_ccall libopencl.POclReleaseEvent(event::cl_event)::cl_int
 end
 
 # Init
@@ -1272,7 +1285,7 @@ end
 
 function set_arg!(k::Kernel, idx::Integer, arg::T) where {T}
     # `Ref{T}` makes `ccall` pass a pointer to a stack copy of `arg`
-    err = @ccall libopencl.POclSetKernelArg(
+    err = @gcsafe_ccall libopencl.POclSetKernelArg(
         k::cl_kernel, cl_uint(idx - 1)::cl_uint, sizeof(T)::Csize_t, arg::Ref{T}
     )::cl_int
     if err == CL_INVALID_ARG_SIZE
@@ -1541,6 +1554,7 @@ struct Event
 end
 Base.unsafe_convert(::Type{cl_event}, e::Event) = e.id
 
+const CL_EVENT_COMMAND_QUEUE = 0x11d0
 const CL_EVENT_COMMAND_EXECUTION_STATUS = 0x11d3
 
 function Base.getproperty(evt::Event, s::Symbol)
@@ -1555,11 +1569,58 @@ function Base.getproperty(evt::Event, s::Symbol)
     end
 end
 
+const CL_COMPLETE = 0
 const CL_EXEC_STATUS_ERROR_FOR_EVENTS_IN_WAIT_LIST = -14
 
+# driver notification that a command has completed
+function notify_completion(::cl_event, ::Cint, payload::Ptr{Cvoid})
+    GPUToolbox.signal_completion(payload)
+    return
+end
+function subscribe_completion(evt, payload)
+    callback = @cfunction(notify_completion, Cvoid, (cl_event, Cint, Ptr{Cvoid}))
+    return clSetEventCallback(evt, CL_COMPLETE, callback, payload)
+end
+
+# commands have completed when their execution status is `CL_COMPLETE`, or negative when
+# they were terminated abnormally
+iscomplete(evt::Event) = evt.status <= CL_COMPLETE
+
+blocking_wait(evt::Event) = unchecked_clWaitForEvents(cl_uint(1), Ref(evt.id))
+
+# block the thread while waiting for commands instead, e.g., to make tests independent of
+# how a wait was performed
+const blocking_waits = Ref(false)
+
 function Base.wait(evt::Event)
-    evt_id = Ref(evt.id)
-    err = unchecked_clWaitForEvents(cl_uint(1), evt_id)
+    # wait without blocking the thread, so that other tasks can run in the meantime. after
+    # polling briefly, PoCL notifies us when the command completes: waking a worker thread
+    # to wait for it, or polling for longer, would compete with the command for the CPU
+    # cores it executes on.
+    #
+    # this cannot be interrupted, as kernels may be using memory that callers would release:
+    # an interrupt is only thrown once the kernel has completed (or waiting failed, in which
+    # case we block), but host memory still needs to be synchronized before unwinding.
+    if !blocking_waits[]
+        try
+            # commands only need to start executing once their queue has been flushed
+            queue = Ref{cl_command_queue}()
+            clGetEventInfo(evt, CL_EVENT_COMMAND_QUEUE, sizeof(cl_command_queue), queue, C_NULL)
+            clFlush(queue[])
+
+            cooperative_wait(
+                blocking_wait, evt; subscribe = subscribe_completion, isdone = iscomplete,
+                spin = 10.0e-6
+            )
+        catch
+            blocking_wait(evt)
+            rethrow()
+        end
+    end
+
+    # synchronize host memory and report errors (without blocking anymore, unless waiting
+    # by blocking)
+    err = unchecked_clWaitForEvents(cl_uint(1), Ref(evt.id))
     if err == CL_EXEC_STATUS_ERROR_FOR_EVENTS_IN_WAIT_LIST
         error("Kernel execution failed")
     elseif err != CL_SUCCESS

test/runtests.jl

@@ -95,7 +95,16 @@ end
     mod.launch_few(few, A)
     mod.launch_many(many, A)
     @test all(==(sum(1:40)), A)
-    @test @allocated(mod.launch_many(many, A)) <= @allocated(mod.launch_few(few, A))
+    # waiting for a kernel allocates when it involves a completion callback, which depends
+    # on how long the kernel takes, so measure launches that wait by blocking instead
+    allocated(launch, k, A) = @allocated launch(k, A)
+    POCL.cl.blocking_waits[] = true
+    try
+        allocated(mod.launch_many, many, A)
+        @test allocated(mod.launch_many, many, A) <= allocated(mod.launch_few, few, A)
+    finally
+        POCL.cl.blocking_waits[] = false
+    end
 end
 
 @testset "POCL compilation cache" begin
@@ -160,6 +169,32 @@ end
     @test all(==(2.0f0), A)
 end
 
+@kernel function busy_kernel!(A, n)
+    I = @index(Global)
+    acc = 0.0f0
+    for j in 1:n
+        acc += sin(Float32(j) + acc)
+    end
+    @inbounds A[I] = acc
+end
+
+# POCL launches wait for the kernel to finish, but let other tasks run in the meantime
+@testset "POCL cooperative launches" begin
+    A = zeros(Float32, 1024)
+    kernel = busy_kernel!(POCLBackend())
+    kernel(A, 1; ndrange = length(A))   # compile
+    n = 1000
+    while @elapsed(kernel(A, n; ndrange = length(A))) < 0.1
+        n *= 2
+    end
+
+    ran = Ref(false)
+    task = @async ran[] = true
+    kernel(A, n; ndrange = length(A))
+    @test ran[]
+    wait(task)
+end
+
 # not part of the shared testsuite: not every back-end supports bits-union arrays
 @testset "POCL zeros/ones of bits-union types" begin
     for T in (Union{Missing, Bool}, Union{Missing, Int32})