Use GPUToolbox's cooperative_wait for synchronization

[maleadt]

CUDA.jl, oneAPI.jl and OpenCL.jl wait for the GPU on a worker thread, so that the calling task doesn't block its thread. That implementation now lives in GPUToolbox as cooperative_wait (JuliaGPU/GPUToolbox.jl#23). This PR switches AMDGPU.jl over to it for streams, events and HIP.device_synchronize().

Until now, nonblocking stream synchronization enqueued a host function with hipLaunchHostFunc that woke the waiting task through libuv. Because that host function might never run if something went wrong, every synchronization also spawned a second task that polled the stream on a 1 s timer. cooperative_wait polls the stream first, exactly like before, and then hands a plain hipStreamSynchronize to a small pool of worker threads while the calling task waits on an event. Errors come back from hipStreamSynchronize itself, so the watchdog goes away. This also fixes a few other things:

• Host functions are slow and get in the way of the stream. CLR puts a blocking marker behind the callback, so later work on the stream waits for the host function too, and the wakeup then goes through the libuv event loop. On CUDA, the equivalent cuLaunchHostFunc path added 20–75 µs over a blocking wait, against 2–4 µs for handing the wait to a worker (measured with CUDA.jl on an RTX 5080).

• Event synchronization blocked the thread. synchronize(::HIPEvent) polled for a while and then called hipEventSynchronize on the calling thread. It now waits on a worker too.

• HIP.device_synchronize() blocked the thread, and could deadlock with hostcalls. The host side of a hostcall is a task. If a kernel waits for a hostcall whose task can only run on the thread blocked in hipDeviceSynchronize, nothing makes progress. For example, with julia --threads=1:

  hc = AMDGPU.Device.HostCallHolder(Float32, Tuple{Float32}; continuous=true) do x
      x + 42f0
  end
  kernel!(y, hc) = (y[1] = AMDGPU.Device.hostcall!(hc, y[1]); nothing)
  y = ROCArray(Float32[0])
  @roc kernel!(y, hc)
  AMDGPU.HIP.device_synchronize()   # hangs on main, returns with this PR

• Every synchronization called device! from the watchdog task, which also sets the global default device.

The workers call HIP directly rather than through the checked wrappers, because those switch to the context of the task making the call. They select the caller's device with hipSetDevice first, since the null stream and hipDeviceSynchronize depend on it. blocking=true and the nonblocking_synchronization preference work as before. Finalizers, which can't switch tasks, keep the context they selected and wait on the calling thread.

This requires GPUToolbox 3.3.1, which fixes two ways the workers could deadlock (JuliaGPU/GPUToolbox.jl#26). Both matter for the hostcall case above. hipDeviceSynchronize can't be polled, and before 3.3.1 such a wait had to wait for a free worker when all of them were busy, which never happens if the busy workers are waiting for kernels that need the hostcall to be served. And a worker could run a finalizer that blocks (e.g. freeing memory, which can wait for the device to go idle) before waking up the task it waited for. Now such waits get an extra worker, and workers never run finalizers.

This hasn't been run on a discrete AMD GPU yet, so CI is the real validation. I did compare both implementations on the integrated GPU of a Ryzen 9 9950X (gfx1036, run as gfx1030 with ROCm 7.2.4), in one process, alternating between them in blocks. The tables show the median (p90) time for a launch plus synchronize() of a busy-waiting kernel:

--threads=1:

kernel	blocking	before	after
0 µs	6.7 (7.3) µs	8.8 (12.1) µs	8.7 (11.6) µs
100 µs	112.2 (112.9) µs	117.4 (120.4) µs	113.8 (116.1) µs
1 ms	1018.5 (1022.5) µs	1027.4 (1057.9) µs	1021.2 (1037.9) µs
5 ms	5026.2 (5050.9) µs	5068.9 (5119.9) µs	5050.7 (5083.0) µs

--threads=4:

kernel	blocking	before	after
0 µs	6.8 (7.2) µs	7.0 (8.3) µs	7.0 (7.8) µs
100 µs	111.2 (111.6) µs	122.2 (128.6) µs	112.9 (114.6) µs
1 ms	1017.3 (1018.2) µs	1044.7 (1068.0) µs	1020.3 (1035.1) µs
5 ms	5025.0 (5047.6) µs	5084.9 (5169.0) µs	5044.6 (5083.1) µs

Short kernels finish while polling, so they aren't affected. For longer ones, the new version was never slower than the host-function path in any run, and usually had a much tighter tail. The machine was heavily loaded, so the 5 ms numbers varied between runs, but the ordering didn't change.

The trade-off is CPU usage while waiting. HIP busy-waits in hipStreamSynchronize, so the worker keeps a core busy, just like blocking=true does on the calling thread. On this machine, polling a stream with hipStreamQuery also kept one of HSA's runtime threads busy, both before and after this PR. As a result, a long wait used about 2 cores instead of about 1. HIP.synchronize(stream; spin=false) skips the polling and brings that back to 1 core, at the cost of a few µs on short kernels.

The new tests check that synchronize(stream) (including the null stream and spin=false), synchronize(event) and HIP.device_synchronize() don't block the thread, and that a hostcall gets served while device_synchronize() waits in a single-threaded process. They don't measure anything. A kernel spins until a flag in host memory is set, and another task on the same thread sets it after yielding many more times than the polling before a wait does, so the wait can only return if that task got to run in the meantime. If it never does, the kernel gives up after 20 s or so and records that it timed out, so a regression fails instead of hanging. That also makes the tests indifferent to the MI250 runner stalling test processes (individual tests there regularly take 5–30× longer than usual, also on main). On main, the event and device tests fail, and the hostcall test times out. With GPUToolbox 3.3.2, they pass on the integrated GPU mentioned above, also with --threads=1 and while the process is repeatedly stopped and resumed, and with nonblocking_synchronization=false they fail as expected.

[gbaraldi]

Tested on MI300A + MI250 (ROCm 7.2.4), 1 and 4 threads. The new tests pass (on main the event/device tests fail and the hostcall test hangs), multi-GPU and a 64-task stress run are OK, and latency is ≤ main with tighter tails.

1. Overflow threads. device_synchronize has no isdone, so each call made while the workers are busy gets a new overflow thread, and those threads are never torn down. 64 concurrent calls left 42 extra spinning threads. Allow one sync per device at a time, or poll an event on the null stream.
2. Ctrl-C during synchronize waits until the kernel finishes (12 s vs 0.5 s on main), so a hung kernel can't be interrupted. Document it or expose cancellable.
3. CPU use. Once any host function has been used in the process (e.g. #1116), polling with hipStreamQuery makes a CLR thread spin: 2 cores per wait instead of 1 (spin=false: 1).
4. Nits: spin=false costs 15–20 µs on short kernels; spin isn't documented or passed through AMDGPU.synchronize; the device_synchronize docstring still says "Blocks"; docs/src/api/system.md says nonblocking_synchronize (predates the PR).

[github-actions]

AMDGPU.jl Benchmarks

Details

Benchmark suite	Current: 93b690c	Previous: 4a144ad	Ratio
amdgpu/synchronization/context/device	1765 ns	555 ns	3.18
amdgpu/synchronization/stream/blocking	220 ns	222.5 ns	0.99
amdgpu/synchronization/stream/nonblocking	342.5 ns	305 ns	1.12
applications/bitonic_sort	1079366.25 ns	1090213.75 ns	0.99
applications/convolution	92228.5 ns	91664 ns	1.01
applications/floyd_warshall	9060402.75 ns	9044243.75 ns	1.00
applications/histogram	770639.75 ns	812457 ns	0.95
applications/prefix_sum	188204.75 ns	188707.75 ns	1.00
array/accumulate/Float32/1d	80026 ns	80296.25 ns	1.00
array/accumulate/Float32/dims=1	251087.75 ns	259326.25 ns	0.97
array/accumulate/Float32/dims=1L	103211.25 ns	101036.5 ns	1.02
array/accumulate/Float32/dims=2	105086.25 ns	103989 ns	1.01
array/accumulate/Float32/dims=2L	3008542.25 ns	2961769.25 ns	1.02
array/accumulate/Int64/1d	87111 ns	87134 ns	1.00
array/accumulate/Int64/dims=1	244520.5 ns	242691.25 ns	1.01
array/accumulate/Int64/dims=1L	114326.25 ns	112086.5 ns	1.02
array/accumulate/Int64/dims=2	77781 ns	103629 ns	0.75
array/accumulate/Int64/dims=2L	3365491.75 ns	3392515.5 ns	0.99
array/broadcast	29800.5 ns	29495.5 ns	1.01
array/construct	2272.5 ns	2045 ns	1.11
array/copy	40668 ns	41110.75 ns	0.99
array/copyto!/cpu_to_gpu	89146 ns	88973.75 ns	1.00
array/copyto!/gpu_to_cpu	89436 ns	89698.75 ns	1.00
array/copyto!/gpu_to_gpu	33813 ns	34610.5 ns	0.98
array/iteration/findall/bool	123309 ns	119711.75 ns	1.03
array/iteration/findall/int	111761.25 ns	112654.25 ns	0.99
array/iteration/findfirst/bool	150807 ns	153357.25 ns	0.98
array/iteration/findfirst/int	134116.75 ns	165282.5 ns	0.81
array/iteration/findmin/1d	108726.25 ns	104999 ns	1.04
array/iteration/findmin/2d	94266.25 ns	81738.5 ns	1.15
array/iteration/logical	191774.75 ns	183337.75 ns	1.05
array/iteration/scalar	302246.25 ns	309797 ns	0.98
array/permutedims/2d	58623.25 ns	37878.25 ns	1.55
array/permutedims/3d	47293 ns	57238.5 ns	0.83
array/permutedims/4d	62470.75 ns	62181 ns	1.00
array/random/rand/Float32	44630.5 ns	38143 ns	1.17
array/random/rand/Int64	43023 ns	37603 ns	1.14
array/random/rand!/Float32	41308 ns	42010.5 ns	0.98
array/random/rand!/Int64	37733 ns	36980.75 ns	1.02
array/random/randn/Float32	71225.75 ns	71413.75 ns	1.00
array/random/randn!/Float32	53065.75 ns	57343.25 ns	0.93
array/reductions/mapreduce/Float32/1d	109571.5 ns	110571.75 ns	0.99
array/reductions/mapreduce/Float32/dims=1	87588.5 ns	69198.5 ns	1.27
array/reductions/mapreduce/Float32/dims=1L	817452.25 ns	847172.5 ns	0.96
array/reductions/mapreduce/Float32/dims=2	93736 ns	93956.5 ns	1.00
array/reductions/mapreduce/Float32/dims=2L	134884 ns	135664.5 ns	0.99
array/reductions/mapreduce/Int64/1d	109564 ns	110764 ns	0.99
array/reductions/mapreduce/Int64/dims=1	88053.5 ns	88948.75 ns	0.99
array/reductions/mapreduce/Int64/dims=1L	822312.5 ns	840315 ns	0.98
array/reductions/mapreduce/Int64/dims=2	92178.75 ns	92799 ns	0.99
array/reductions/mapreduce/Int64/dims=2L	136594 ns	136762 ns	1.00
array/reductions/reduce/Float32/1d	109491.25 ns	110571.75 ns	0.99
array/reductions/reduce/Float32/dims=1	87701 ns	88131.25 ns	1.00
array/reductions/reduce/Float32/dims=1L	818402.5 ns	838072.5 ns	0.98
array/reductions/reduce/Float32/dims=2	93561.25 ns	93711.5 ns	1.00
array/reductions/reduce/Float32/dims=2L	135739.25 ns	134609.5 ns	1.01
array/reductions/reduce/Int64/1d	100638.75 ns	110611.75 ns	0.91
array/reductions/reduce/Int64/dims=1	88023.75 ns	88733.75 ns	0.99
array/reductions/reduce/Int64/dims=1L	820442.75 ns	839045 ns	0.98
array/reductions/reduce/Int64/dims=2	92026 ns	92376.25 ns	1.00
array/reductions/reduce/Int64/dims=2L	136396.5 ns	136572 ns	1.00
array/reverse/1d	42770.5 ns	41338.25 ns	1.03
array/reverse/1dL	52263.25 ns	68548.5 ns	0.76
array/reverse/1dL_inplace	53285.75 ns	54298.25 ns	0.98
array/reverse/1d_inplace	34500.5 ns	35448 ns	0.97
array/reverse/2d	45783 ns	46618 ns	0.98
array/reverse/2dL	89813.5 ns	90621.25 ns	0.99
array/reverse/2dL_inplace	64363.25 ns	65223.5 ns	0.99
array/reverse/2d_inplace	36430.5 ns	37133 ns	0.98
array/sorting/1d	321893.5 ns	337702.75 ns	0.95
gemm/tiled	2840073.25 ns	1966209 ns	1.44
gemm/tiled_unbounded	2834583.5 ns	1965436.5 ns	1.44
integration/byval/reference	38751 ns	38851 ns	1.00
integration/byval/slices=1	40561 ns	40371 ns	1.00
integration/byval/slices=2	158162 ns	145102 ns	1.09
integration/byval/slices=3	212273 ns	209573 ns	1.01
integration/volumerhs	4957783 ns	4879902 ns	1.02
kernel/indexing	30215.5 ns	30490.5 ns	0.99
kernel/indexing_checked	36098 ns	36730.75 ns	0.98
kernel/launch	1222.5 ns	1205 ns	1.01
kernel/rand	37858 ns	51168.25 ns	0.74
latency/import	1632619781 ns	1625449571 ns	1.00
latency/precompile	29299214001 ns	29273376579 ns	1.00
latency/ttfp	2192262246 ns	2179782468 ns	1.01
stencil/diffusion3d	1554549.75 ns	1600543.75 ns	0.97
stencil/diffusion3d_checked	1595055 ns	1636781.75 ns	0.97

This comment was automatically generated by workflow using github-action-benchmark.