Clarify what @spawn orders, for users and for backends

docs/src/implementations.md

@@ -16,49 +16,40 @@ 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. Likewise,
+`synchronize` **should not** wait for work on other queues that the current queue is not
+ordered after.
+
+Otherwise, a spawned task's first use of an array shared with its parent waits for work the
+parent queued after `@spawn`, reducing overlap. Neither recommendation guarantees that the
+work of different tasks runs concurrently.
 
 
 ## Moving data with `adapt`

lib/KernelInterface/src/host.jl

@@ -36,22 +36,19 @@ 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
-    [notes for backend implementations](@ref implementations_notes) for why.
+    Backend implementations **must** implement this function cooperatively, yielding to
+    the Julia scheduler while waiting rather than blocking inside a driver call. See the
+    [notes for backend implementations](@ref implementations_notes) for why, and for what
+    it should not wait for.
 """
 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 so far, and
+return a handle for [`wait_event`](@ref). Work queued later is not captured, and recording
+need not wait for the captured work to complete. The handle is only meant for `wait_event`.
 
 !!! note
     The default implementation calls [`synchronize`](@ref) and returns `nothing`.
@@ -68,28 +65,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
 

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,40 +32,40 @@ 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)`:
+A task started with plain `Threads.@spawn` runs on the backend's default device, not on
+the device of the task that started it. `@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
     mul2_kernel(backend, 64)(B, ndrange = length(B))
 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`"