diff --git a/docs/src/implementations.md b/docs/src/implementations.md index 273b7beab..30f188e4a 100644 --- a/docs/src/implementations.md +++ b/docs/src/implementations.md @@ -16,49 +16,37 @@ thread instead of letting independent tasks run concurrently. ## Task-local queues and `KernelAbstractions.@spawn` -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. - -[`KernelAbstractions.@spawn`](@ref) hides this from users by following a fixed protocol, -which backends can support with two optional functions: - -- Before the new task is created, the spawning task calls - [`record_event`](@ref KernelAbstractions.record_event) on the backend. The default - implementation is a full [`synchronize`](@ref) returning `nothing`, which is always - correct. A backend with task-local queues **may** instead record an event on the - current task's queue and return it, so that the spawning task does not have to wait. -- The new task selects its device with [`device!`](@ref KernelAbstractions.device!) — the - spawning task's, or the one the user asked for with `@spawn backend device=id` — and then - calls [`wait_event`](@ref KernelAbstractions.wait_event) with the recorded handle. The - order matters: `wait_event` makes the queue of the *currently active* device wait, so the - device has to be selected first. 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. -- After the user's code returns, the new task calls [`synchronize`](@ref), so that - `wait(task)` in any other task implies that all work queued by the spawned task has - completed. - -A new Julia task does not inherit the device of the task that spawned it: backends keep the -active device in task-local state, which Julia does not copy into a child task, so the task -Backends with more than one device -**must** implement the device interface ([`device`](@ref KernelAbstractions.device), -[`ndevices`](@ref KernelAbstractions.ndevices), [`device!`](@ref KernelAbstractions.device!)) -for `@spawn` to run on the right device. - -`@spawn backend device=id` records the event on the spawning task's device but waits on -`id`, so a multi-device backend **must** accept an event recorded on a device other than the -one active in `wait_event`. A backend whose driver cannot **must** fall back -to waiting cooperatively, as [`synchronize`](@ref) does. - -Because `device!` selects the queue that `wait_event` acts on, the same two functions are -what lets users order work across a device switch they make themselves: - -```julia -event = KernelAbstractions.record_event(backend) -KernelAbstractions.device!(backend, 2) -KernelAbstractions.wait_event(backend, event) -``` +Backends **should** give each Julia task its own queue, so that work from different tasks +can execute concurrently. Separate queues do not by themselves order work. + +[`KernelAbstractions.@spawn`](@ref) orders it with this protocol: + +1. The spawning task calls [`record_event`](@ref KernelAbstractions.record_event). +2. The new task selects its device with [`device!`](@ref KernelAbstractions.device!), then + calls [`wait_event`](@ref KernelAbstractions.wait_event) before running the user's code. +3. If that code returns normally, the new task calls [`synchronize`](@ref), so that a + successful `wait(task)` implies its queued work has completed. + +The default `record_event` synchronizes and returns `nothing`, for which `wait_event` does +nothing. Backends can return an event instead, so that the spawning task doesn't wait; see +the docstrings of both functions for what that requires. Backends with more than one device +**must** implement [`device`](@ref KernelAbstractions.device), +[`ndevices`](@ref KernelAbstractions.ndevices) and +[`device!`](@ref KernelAbstractions.device!), and `wait_event` **must** accept an event +recorded on another device, since `@spawn backend device=id` records on the spawning task's +device. + +Backends that track which queue last used an array, and wait for that queue before using the +array on another one, **should** skip that wait when the current queue is already ordered +after the array's last use: through an event recorded on the previous queue after that use +and waited for by the current queue, or through a [`synchronize`](@ref) of the previous +queue that completed that use. In particular, they should not wait, on the host or on the +device, for work queued on the previous queue after that event or synchronization. Waits +needed to make memory accessible or to keep it alive still apply, and uses through a pointer +taken before that event or synchronization may be synchronized conservatively. This lets a +spawned task use arrays it shares with its parent while the parent keeps queuing work, which +is also why [`synchronize`](@ref) **must not** wait for work on other queues that the +current queue is not ordered after. ## Moving data with `adapt` diff --git a/lib/KernelInterface/src/host.jl b/lib/KernelInterface/src/host.jl index b769783cd..1c3ff0573 100644 --- a/lib/KernelInterface/src/host.jl +++ b/lib/KernelInterface/src/host.jl @@ -36,9 +36,10 @@ Block the calling task until all work it has queued on the active device of `bac completed. !!! note - Backend implementations **must** implement this function, and it **must** be - cooperative: it may not block inside a driver call, but has to yield to the Julia - scheduler while waiting. See the + Backend implementations **must** implement this function cooperatively, yielding to + the Julia scheduler while waiting rather than blocking inside a driver call. It + **must not** wait for work on other queues that the current queue is not ordered + after, e.g., by [`wait_event`](@ref). See the [notes for backend implementations](@ref implementations_notes) for why. """ function synchronize end @@ -46,12 +47,9 @@ function synchronize end """ record_event(backend::Backend) -Capture the work the calling task has queued on `backend`'s currently active device so -far, and return a handle that [`wait_event`](@ref) can use to order later work after it, -either from another task or from the same task after switching devices. - -The handle is only meaningful for the pair `record_event`/`wait_event`; do not use it for -anything else. +Capture the work the calling task has queued on `backend`'s active device before this +call, and return a handle for [`wait_event`](@ref). Recording need not wait for that work to +complete. The handle is only meant for `wait_event`. !!! note The default implementation calls [`synchronize`](@ref) and returns `nothing`. @@ -68,28 +66,29 @@ end """ wait_event(backend::Backend, event) -Order the work the calling task subsequently queues on `backend`'s currently active device -after the work captured by `event`, which was returned by [`record_event`](@ref). +Order the work the calling task subsequently queues on `backend`'s active device after the +work captured by `event`, which was returned by [`record_event`](@ref). Returning does not +mean that the captured work has completed. -The dependency is queue-ordered rather than task-ordered: it applies to the device that is -active when `wait_event` is called, and a later [`device!`](@ref) leaves the newly selected -device unordered with respect to `event`. Select the device first and wait afterwards: +The wait applies to the queue of the device that is active when `wait_event` is called; +switching devices adds no ordering. To order work across a device switch, select the device +first and wait afterwards: ```julia event = record_event(backend) # captures work on the current device device!(backend, 2) -wait_event(backend, event) # device 2 now waits for that work +wait_event(backend, event) # orders this task's work on device 2 after it ``` !!! note `wait_event(::Backend, ::Nothing)` is a no-op, matching the default `record_event`. - A backend that implements [`record_event`](@ref) **must** implement this for the event - type it returns, either by enqueuing a dependency on the current task's queue, or by - waiting cooperatively as [`synchronize`](@ref) does. A backend with more than one - device **must** also accept an `event` that was recorded on a different device, by - enqueuing the cross-device dependency if the driver supports one (CUDA's - `cuStreamWaitEvent` does) and by waiting cooperatively otherwise. See the - [notes for backend implementations](@ref implementations_notes). + A backend that returns another event type **must** implement `wait_event` for it, + either by adding a dependency to the current task's queue or by waiting cooperatively + as [`synchronize`](@ref) does. A backend with more than one device **must** accept an + event recorded on another device, waiting cooperatively if the driver cannot add a + cross-device dependency. See the + [notes for backend implementations](@ref implementations_notes) for how this ordering + should interact with implicit synchronization. """ wait_event(::Backend, ::Nothing) = nothing diff --git a/src/spawn.jl b/src/spawn.jl index 12a5b1904..14417426e 100644 --- a/src/spawn.jl +++ b/src/spawn.jl @@ -8,10 +8,10 @@ place of `Threads.@spawn` to launch kernels from a task. It guarantees that that argument is given; - the work the task queues on `backend` runs after the work the spawning task had queued on `backend` before calling `@spawn`; -- once `wait(task)` or `fetch(task)` returns, all work the task queued on `backend` has - completed, so its results may be used from any task. `fetch(task)` returns the value of - `expr`. If `expr` throws, the task is not synchronized: its queued work may still be - running when the exception surfaces. +- once `wait(task)` or `fetch(task)` returns successfully, all work the task queued on + `backend` has completed, so its results may be used from any task. `fetch(task)` + returns the value of `expr`. If `expr` throws, the task is not synchronized: its queued + work may still be running when the exception surfaces. Everything else works as for `Threads.@spawn`: the optional `threadpool` argument (`:default` or `:interactive`) is forwarded, `\$x` captures the value of `x` at spawn time, @@ -32,11 +32,9 @@ fetch(task) == 4 * length(A) # Choosing the device -Backends keep the active device in task-local state, and Julia does not copy that state -into a child task. A task started with plain `Threads.@spawn` therefore runs on the -backend's *default* device, whichever device the spawning task was using. `@spawn` selects -the device explicitly instead: by default the one active in the spawning task, or the one -named by `device`, a 1-based index into `1:ndevices(backend)`: +Which device a task started with plain `Threads.@spawn` uses depends on the backend, and +need not be the spawning task's. `@spawn` selects the spawning task's device, or the one +given by `device`, an index into `1:ndevices(backend)`: ```julia task = KernelAbstractions.@spawn backend device=2 begin @@ -44,28 +42,30 @@ task = KernelAbstractions.@spawn backend device=2 begin end ``` -The ordering guarantee holds across that switch: the task's work on `device` is still -ordered after the work the spawning task had queued on *its* device. Backends that support -more than one device implement this with a cross-device -[`wait_event`](@ref KernelAbstractions.wait_event). +The task's work on that device still runs after the spawning task's earlier work on its own +device. !!! note - `expr` should not rely on data that the spawning task queues *after* `@spawn` returns. - Order later work by waiting on the task, or by spawning again. + `@spawn` does not order the task's work against work the spawning task queues + afterwards. Order conflicting uses of shared data by waiting for the task, or by spawning a new + task after that work. !!! note - Prefer `device=` over calling [`device!`](@ref KernelAbstractions.device!) inside - `expr`. A `device!` in the body carries no ordering of its own, so work queued after it - is ordered neither against the spawning task nor against what the body queued before - the switch; you would have to bracket it with - [`record_event`](@ref KernelAbstractions.record_event) and - [`wait_event`](@ref KernelAbstractions.wait_event) yourself. + Queued work is ordered, but the spawning task's earlier work need not have completed + when `expr` starts. Before passing that work's results to a consumer outside that + ordering, e.g., an MPI call on a GPU buffer, call `synchronize(backend)` in `expr`, + which also waits for that work. Some backends synchronize implicitly when the buffer's + pointer is taken, but portable code should not rely on that. `@spawn` does not wait for + asynchronous operations outside the backend, like `MPI.Isend`. !!! note - If `expr` throws the state of the device and the internal queue is unspecified. + Prefer `device=` over calling [`device!`](@ref KernelAbstractions.device!) in `expr`. + To keep work ordered across a manual switch, call + [`record_event`](@ref KernelAbstractions.record_event) before it and + [`wait_event`](@ref KernelAbstractions.wait_event) after it, as shown for `wait_event`. + Backend authors: see the [notes for backend implementations](@ref implementations_notes) -for the protocol behind these guarantees, and for how to support it without a full -[`synchronize`](@ref). +for the protocol behind these guarantees. """ macro spawn(args...) usage = "@spawn expects `@spawn [threadpool] backend [device=id] expr`"