Add KernelAbstractions.@spawn and record_event/wait_event by vchuravy · Pull Request #750 · JuliaGPU/KernelAbstractions.jl

vchuravy · 2026-09-07T17:32:02Z

Summary

Adds KernelAbstractions.@spawn backend expr, which runs expr on a new Julia task while keeping the work queued on backend ordered between the two tasks. It encodes the discipline users otherwise have to write by hand:

synchronize(backend)
Threads.@spawn begin
    ...
    synchronize(backend)
end |> wait

The macro's protocol:

The spawning task calls record_event(backend).
The new task selects the spawning task's device with device!, then calls wait_event(backend, event).
After expr returns, the task calls synchronize(backend), so wait(task) / fetch(task) imply that all of the task's device work has completed.

An optional first argument is forwarded to Threads.@spawn as the threadpool (@spawn :interactive backend expr). The macro is not exported, to avoid clashing with Threads.@spawn.

Backend opt-in

Two new optional functions in KernelInterface:

record_event(backend) defaults to a full synchronize returning nothing, which is always correct.
wait_event(::Backend, ::Nothing) is a no-op.

A backend with task-local streams (CUDA.jl, AMDGPU.jl, ...) can override record_event to record an event on the current stream without blocking, and implement wait_event for that event type as a stream wait, turning the protocol into

event = record()
@spawn begin
    wait(event)
    ...
    synchronize()
end |> wait

Backends with a single global queue (POCL) need no changes.

Docs

synchronize docstring now recommends a cooperative, non-blocking implementation, since a blocking one serializes otherwise independent spawned tasks.
API and KernelInterface pages list the new functions; the implementer notes gain a "Task-local queues and @spawn" section; the quickstart's task-programming section now recommends @spawn.

Test plan

New Spawn testsuite entry (test/spawn.jl): ordering after the parent's queued work, visibility after wait, device pinning, threadpool forms, single evaluation of the backend expression, error propagation, many concurrent tasks. Passes on the CPU backend.
KernelInterface standalone tests cover the record_event/wait_event fallbacks.
Full CPU testsuite: all pass except a pre-existing Private error caused by my local GPUCompiler dev checkout lacking alloca (unrelated to this change).
CI on GPU backends (they use the default synchronize-based fallback until they opt in).

🤖 Generated with Claude Code

https://claude.ai/code/session_01YQWT6DHjjRoh1YaAsFEUN5

github-actions · 2026-09-07T17:51:01Z

Benchmark Results

Show table

	main	`54c0185`...	main / `54c0185`...
saxpy/default/Float32/1024	0.0507 ± 0.023 ms	0.0477 ± 0.025 ms	1.06 ± 0.73
saxpy/default/Float32/1048576	0.198 ± 0.016 ms	0.197 ± 0.019 ms	1.01 ± 0.13
saxpy/default/Float32/16384	0.0521 ± 0.02 ms	0.05 ± 0.021 ms	1.04 ± 0.59
saxpy/default/Float32/2048	0.0502 ± 0.022 ms	0.0478 ± 0.023 ms	1.05 ± 0.68
saxpy/default/Float32/256	0.0502 ± 0.023 ms	0.0464 ± 0.025 ms	1.08 ± 0.78
saxpy/default/Float32/262144	0.0909 ± 0.022 ms	0.0887 ± 0.023 ms	1.03 ± 0.37
saxpy/default/Float32/32768	0.0536 ± 0.022 ms	0.0492 ± 0.022 ms	1.09 ± 0.65
saxpy/default/Float32/4096	0.0502 ± 0.021 ms	0.0496 ± 0.022 ms	1.01 ± 0.62
saxpy/default/Float32/512	0.0506 ± 0.023 ms	0.0476 ± 0.025 ms	1.06 ± 0.75
saxpy/default/Float32/64	0.0503 ± 0.023 ms	0.0459 ± 0.026 ms	1.1 ± 0.8
saxpy/default/Float32/65536	0.0568 ± 0.022 ms	0.0524 ± 0.023 ms	1.08 ± 0.63
saxpy/default/Float64/1024	0.0499 ± 0.023 ms	0.0479 ± 0.024 ms	1.04 ± 0.71
saxpy/default/Float64/1048576	0.274 ± 0.02 ms	0.269 ± 0.024 ms	1.02 ± 0.12
saxpy/default/Float64/16384	0.0512 ± 0.021 ms	0.0483 ± 0.022 ms	1.06 ± 0.64
saxpy/default/Float64/2048	0.0507 ± 0.023 ms	0.0493 ± 0.024 ms	1.03 ± 0.67
saxpy/default/Float64/256	0.051 ± 0.023 ms	0.0474 ± 0.025 ms	1.08 ± 0.75
saxpy/default/Float64/262144	0.113 ± 0.02 ms	0.114 ± 0.02 ms	0.997 ± 0.24
saxpy/default/Float64/32768	0.0561 ± 0.021 ms	0.0538 ± 0.021 ms	1.04 ± 0.57
saxpy/default/Float64/4096	0.0497 ± 0.023 ms	0.0484 ± 0.022 ms	1.03 ± 0.67
saxpy/default/Float64/512	0.0504 ± 0.023 ms	0.0479 ± 0.025 ms	1.05 ± 0.74
saxpy/default/Float64/64	0.0506 ± 0.023 ms	0.0472 ± 0.025 ms	1.07 ± 0.75
saxpy/default/Float64/65536	0.0618 ± 0.021 ms	0.0575 ± 0.022 ms	1.07 ± 0.54
saxpy/static workgroup=(1024,)/Float32/1024	0.0494 ± 0.024 ms	0.0467 ± 0.024 ms	1.06 ± 0.75
saxpy/static workgroup=(1024,)/Float32/1048576	0.691 ± 0.013 ms	0.693 ± 0.014 ms	0.998 ± 0.027
saxpy/static workgroup=(1024,)/Float32/16384	0.0614 ± 0.021 ms	0.0609 ± 0.021 ms	1.01 ± 0.49
saxpy/static workgroup=(1024,)/Float32/2048	0.0512 ± 0.023 ms	0.0505 ± 0.024 ms	1.01 ± 0.65
saxpy/static workgroup=(1024,)/Float32/256	0.0491 ± 0.024 ms	0.0478 ± 0.024 ms	1.03 ± 0.73
saxpy/static workgroup=(1024,)/Float32/262144	0.214 ± 0.018 ms	0.215 ± 0.019 ms	0.998 ± 0.12
saxpy/static workgroup=(1024,)/Float32/32768	0.0685 ± 0.021 ms	0.0675 ± 0.022 ms	1.02 ± 0.46
saxpy/static workgroup=(1024,)/Float32/4096	0.0536 ± 0.02 ms	0.0525 ± 0.021 ms	1.02 ± 0.56
saxpy/static workgroup=(1024,)/Float32/512	0.0496 ± 0.024 ms	0.0476 ± 0.024 ms	1.04 ± 0.73
saxpy/static workgroup=(1024,)/Float32/64	0.0492 ± 0.024 ms	0.048 ± 0.024 ms	1.03 ± 0.71
saxpy/static workgroup=(1024,)/Float32/65536	0.0897 ± 0.021 ms	0.0878 ± 0.022 ms	1.02 ± 0.35
saxpy/static workgroup=(1024,)/Float64/1024	0.0484 ± 0.023 ms	0.0476 ± 0.023 ms	1.02 ± 0.7
saxpy/static workgroup=(1024,)/Float64/1048576	0.712 ± 0.02 ms	0.711 ± 0.025 ms	1 ± 0.045
saxpy/static workgroup=(1024,)/Float64/16384	0.059 ± 0.021 ms	0.058 ± 0.022 ms	1.02 ± 0.53
saxpy/static workgroup=(1024,)/Float64/2048	0.0519 ± 0.023 ms	0.0517 ± 0.023 ms	1 ± 0.63
saxpy/static workgroup=(1024,)/Float64/256	0.0488 ± 0.024 ms	0.0466 ± 0.024 ms	1.05 ± 0.74
saxpy/static workgroup=(1024,)/Float64/262144	0.215 ± 0.022 ms	0.214 ± 0.022 ms	1 ± 0.15
saxpy/static workgroup=(1024,)/Float64/32768	0.0699 ± 0.021 ms	0.0687 ± 0.022 ms	1.02 ± 0.45
saxpy/static workgroup=(1024,)/Float64/4096	0.0536 ± 0.02 ms	0.0516 ± 0.021 ms	1.04 ± 0.58
saxpy/static workgroup=(1024,)/Float64/512	0.0492 ± 0.025 ms	0.0462 ± 0.025 ms	1.07 ± 0.78
saxpy/static workgroup=(1024,)/Float64/64	0.0488 ± 0.024 ms	0.0462 ± 0.025 ms	1.05 ± 0.78
saxpy/static workgroup=(1024,)/Float64/65536	0.0919 ± 0.022 ms	0.0906 ± 0.022 ms	1.01 ± 0.34
time_to_load	0.771 ± 0.014 s	0.794 ± 0.0026 s	0.97 ± 0.018

Benchmark Plots

A plot of the benchmark results have been uploaded as an artifact to the workflow run for this PR.
Go to "Actions"->"Benchmark a pull request"->[the most recent run]->"Artifacts" (at the bottom).

vchuravy · 2026-09-07T19:04:44Z

+Backends are free to give each Julia task its own queue (stream), so that kernels
+launched from different tasks can execute concurrently. The price is that work queued
+from two tasks is not ordered with respect to each other, and that a task waiting on
+another task with `wait` learns nothing about the state of that task's queue.


Suggested change

Backends are free to give each Julia task its own queue (stream), so that kernels

launched from different tasks can execute concurrently. The price is that work queued

from two tasks is not ordered with respect to each other, and that a task waiting on

another task with `wait` learns nothing about the state of that task's queue.

Backends should give each Julia task its own queue/stream, so that kernels

launched from different tasks can execute concurrently. This implies that work queued

from two tasks is not ordered with respect to each other.

vchuravy · 2026-09-07T19:05:40Z

+- The new task first selects the spawning task's device with [`device!`](@ref KernelAbstractions.device!),
+  then calls [`wait_event`](@ref KernelAbstractions.wait_event) with the recorded handle.
+  A backend that overrides `record_event` **must** implement `wait_event` for its event
+  type, typically by making the current task's queue wait on the event.


Hm... do we need the device!?

I think wrong device being selected/synchronized is why the OpenCL CUDA tests fail

Not sure how to best handle this since we never really dealt with #631 (comment)

The lazy solution is to add in the docs that the device interface should be supported to use @spawn with a non-default device and disable those tests on the OpenCL CUDA runs

vchuravy · 2026-09-07T19:06:01Z

+A backend with a single, global queue needs no changes: the defaults are exactly the
+"synchronize before, synchronize after" discipline users would otherwise write by hand.


Suggested change

A backend with a single, global queue needs no changes: the defaults are exactly the

"synchronize before, synchronize after" discipline users would otherwise write by hand.

vchuravy · 2026-09-07T19:06:40Z

+Some backends give each Julia task its own queue, so kernels launched from two tasks are not
+ordered with respect to each other, and `wait(task)` on its own says nothing about whether
+the kernels that task launched have finished. Use [`KernelAbstractions.@spawn`](@ref) instead
+of `Threads.@spawn` to launch kernels from a task. It orders the new task's work after the
+work the spawning task has already queued, runs it on the same device, and synchronizes the
+backend before the task finishes, so that `wait(task)` and `fetch(task)` guarantee its
+results are ready:


Needs re-wording

vchuravy · 2026-09-08T17:17:38Z

@christiangnrd we should have this for KI 0.2 as well

@Spawn

`KernelAbstractions.@Spawn backend expr` runs `expr` on a new Julia task while keeping the work queued on `backend` ordered between the two tasks: 1. the spawning task calls `record_event(backend)`, 2. the new task selects the same device and calls `wait_event`, 3. after `expr` returns the task calls `synchronize(backend)`, so that `wait(task)`/`fetch(task)` imply all of its device work has completed. `record_event` defaults to a full `synchronize` returning `nothing`, and `wait_event(::Backend, ::Nothing)` is a no-op, so every backend gets the "synchronize before, synchronize after" discipline for free. Backends with task-local streams may opt in by recording an event instead and implementing `wait_event` for it as a stream wait. Also document that `synchronize` should be cooperative, since a blocking implementation would serialize otherwise independent spawned tasks. Assisted-by: Claude Code (Fable 5.1)

@Spawn

`Threads.@Spawn :interactive` runs the task in the `:default` pool when the interactive pool has no threads, so asserting `Threads.threadpool() === :interactive` fails on Julia 1.10 and 1.11, which do not start an interactive thread by default. Expect whichever pool Julia will actually use. Assisted-by: Claude Code (Opus 5)

christiangnrd · 2026-09-08T18:05:49Z

@christiangnrd we should have this for KI 0.2 as well

Sure! Is the idea that things you write your code assuming events are supported, and on backends where it isn't it just silently falls back to sequential operation?

vchuravy · 2026-09-08T18:06:47Z

Yeah the sequential ordering is the correct thing, and the events are an optimization.

christiangnrd · 2026-09-08T18:18:27Z

+!!! note "Cooperative synchronization"
+    Backend implementations **should** make `synchronize` cooperative rather than blocking.
+    That is, instead of blocking inside a driver call, it should poll or wait on a
+    completion signal while calling `yield` so that other Julia tasks can run in the
+    meantime. A blocking implementation stalls every task scheduled on the same thread,
+    which defeats overlapping kernels with host work or communication, and makes
+    [`KernelAbstractions.@spawn`](@ref)'s trailing `synchronize` serialize otherwise
+    independent tasks.


Doesn't the quickstart manual say they have to be cooperative?

christiangnrd

Claude seems to find it very important that someone reading any part of the documentation know that this defaults to synchronize and that if record_event is implemented, wait_event must also be implemented.

I think the docs should focus more on desired behaviour, and any implementation recommendations/guidelines should be tucked away and clearly marked as such since realistically it'll be one of us adding backend support and users probably won't care about implementation unless things break

christiangnrd · 2026-09-08T23:19:38Z

It'll probably end up in KI 0.2.0, but if I understand correctly, since the backend implementations are just an optimization (and optional), this can release whenever and backends just set compat to whatever version this is released in once they add support

Julia 1.10 ignores `[sources]`, so it resolves KernelInterface from the registry rather than from lib/KernelInterface. The "Dev KernelInterface" step compensates, but it was gated on `runner.os != 'Windows'`, copied from the neighboring runtest step; only the test run needs the de-escalated shell, not `Pkg.develop`. As a result the 1.10 Windows job picked up the registered KernelInterface v0.1.0, which lacks `record_event`, and KernelAbstractions failed to precompile with `UndefVarError: record_event not defined`. Run the step everywhere, under bash so the quoting works on Windows. Assisted-by: Claude Code (Opus 5)

`[sources]` is only supported from Julia 1.11 on. On 1.10 the entry pointing KernelInterface at lib/KernelInterface is silently ignored and Pkg resolves it from the registry, so the dev step is what keeps 1.10 testing this repo's copy. Record that where the step is, and cross-reference it from the OpenCL job, which dev's KernelInterface for the same reason. Assisted-by: Claude Code (Opus 5)

vchuravy requested review from christiangnrd and maleadt September 7, 2026 19:03

vchuravy commented Sep 7, 2026

View reviewed changes

vchuravy added this to the 0.10.0 milestone Sep 7, 2026

vchuravy force-pushed the vc/spawn branch from 634fe67 to 0be2934 Compare September 8, 2026 17:17

christiangnrd reviewed Sep 8, 2026

View reviewed changes

vchuravy added 2 commits September 9, 2026 16:09

		A backend with a single, global queue needs no changes: the defaults are exactly the
		"synchronize before, synchronize after" discipline users would otherwise write by hand.

Uh oh!

Conversation

vchuravy commented Sep 7, 2026

Summary

Backend opt-in

Docs

Test plan

Uh oh!

github-actions Bot commented Sep 7, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Benchmark Results

Benchmark Plots

Uh oh!

vchuravy Sep 7, 2026

Choose a reason for hiding this comment

Uh oh!

vchuravy Sep 7, 2026

Choose a reason for hiding this comment

Uh oh!

christiangnrd Sep 9, 2026

Choose a reason for hiding this comment

Uh oh!

christiangnrd Sep 9, 2026

Choose a reason for hiding this comment

Uh oh!

christiangnrd Sep 9, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

vchuravy Sep 7, 2026

Choose a reason for hiding this comment

Uh oh!

vchuravy Sep 7, 2026

Choose a reason for hiding this comment

Uh oh!

vchuravy commented Sep 8, 2026

Uh oh!

christiangnrd commented Sep 8, 2026

Uh oh!

vchuravy commented Sep 8, 2026

Uh oh!

christiangnrd Sep 8, 2026

Choose a reason for hiding this comment

Uh oh!

christiangnrd left a comment • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

christiangnrd commented Sep 8, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

github-actions Bot commented Sep 7, 2026 •

edited

Loading

christiangnrd Sep 9, 2026 •

edited

Loading

christiangnrd left a comment •

edited

Loading