|
Uni20
Asynchronous tensor and runtime library.
|
Status: concrete host and CUDA promise types, promise-neutral suspended-task routing, capability-checked scheduler inheritance, explicit scheduler routing, optional CUDA device affinity, explicit device selection, and unified host/multi-device execution in both deterministic and oneTBB schedulers are implemented. Scoped process-wide CUDA resource initialization and global device-bound schedule(CudaTask&&, int) admission are also implemented; automatic storage-driven initial admission remains future work.
This note separates three related mechanisms: initial scheduler admission, rescheduling after suspension, and nested continuation routing. CUDA backend and resource policy are described in CUDA Kernel Dispatch and Device Scheduling.
Related documents:
Uni20 uses distinct public task and promise types over one shared promise implementation and one internal ownership representation:
The concrete type controls initial scheduler admission:
IAsyncScheduler::schedule(AsyncTask&&) admits ordinary host tasks;ICudaScheduler::schedule(CudaTask&&) admits a CUDA task without affinity and routes it through the scheduler's default device;ICudaScheduler::schedule(CudaTask&&, int) admits a CUDA task and binds an explicit device;Each concrete promise constructs its declared public task directly. After admission, schedulers and awaiters store and route the task as BasicTask. Both concrete task types consequently use the same tested implementation for:
TaskHandle always obtains its paired coroutine pointer, promise pointer, and domain tag from the original typed coroutine handle. It is non-owning and is valid only while some runtime ownership path keeps the coroutine frame alive. Code must not recreate a std::coroutine_handle<TaskPromiseBase> from a frame address: promise inheritance does not make that conversion valid. BasicTask carries the actual ownership claim wherever generic runtime code does not need the declared task kind.
This arrangement keeps initial admission and declared task identity type-safe without duplicating the coroutine runtime or adding virtual promise dispatch. The erased handle retains the concrete promise's domain so continuation routing can reject a direct cross-domain transfer. CUDA-aware code may narrow a CUDA-tagged promise to inspect its device. Neither property overrides the runtime scheduler binding.
IScheduler is the internal route for a task that has already been admitted. Its private reschedule(BasicTask&&) operation is callable only by task machinery. Public scheduler interfaces derive from it and add their typed initial-admission operations.
The shared promise records an IScheduler*:
BasicTask; the scheduler has no queued activation for it.BasicTask::reschedule() submits it through the recorded IScheduler.IScheduler* is a non-owning route, not scheduler membership. A public task may carry a route while remaining wholly owned by its AsyncTask or CudaTask. Debug registry entries likewise retain observational metadata rather than a coroutine ownership claim.
The TBB scheduler implementations accept initial and resumed activations with task_group::defer() followed by non-blocking task_arena::enqueue(). A publishing or completion-service thread therefore does not enter the target arena. Because defer registers the activation first, a run_all() sequenced after admission returns cannot miss it. Failure to allocate or publish this internal activation is fatal and reports scheduler/device context rather than returning an unowned coroutine handle.
The scheduler route is runtime continuation state, not a numerical backend tag. Tensor storage and backend dispatch still determine where the operands reside and which kernel is legal.
A suspended task owned by a BasicTask may replace its scheduler route. The old scheduler has no queued activation or coroutine membership to remove: only the task owner knows about the suspended frame. Route replacement first validates the task domain and any CUDA device affinity against the target scheduler, then stores the new non-owning route. The next resubmission belongs entirely to the new scheduler.
When a coroutine awaits another Uni20 task, an unbound child inherits the parent scheduler only when both tasks have the same domain. A same-domain CUDA child may also inherit its parent's device. Cross-domain nesting requires the child to have been explicitly admitted or bound to a scheduler before it is awaited. Entry then follows this rule:
| Prepared child route | Action |
|---|---|
| different scheduler | enqueue on the child scheduler |
| different task domain | enqueue through the recorded scheduler |
| same scheduler and domain, scheduler rejects direct transfer | enqueue through the recorded scheduler |
| same scheduler and domain, scheduler accepts direct transfer | transfer directly |
When the child finishes:
| Parent continuation route | Action |
|---|---|
| scheduler, domain, or scheduler-specific route differs | resubmit the continuation |
| scheduler accepts the same-domain route | transfer directly to the continuation |
Scheduler-pointer equality is therefore necessary but not sufficient for symmetric transfer. The same predicate governs nested task entry, a task returned by a forwarding awaiter, and final continuation return.
The common promise base and promise-neutral continuation handle make this work for all four parent/child combinations of AsyncTask and CudaTask. In particular, an AsyncTask parent may await a scheduler-unbound CudaTask. The child inherits the parent's scheduler only when that scheduler accepts the child's CUDA route and any established device affinity. The child runs in the CUDA domain and the parent later resumes in the host domain through the same unified scheduler. A host-only scheduler rejects the child route.
The parent does not become a CudaTask, and the child does not carry the parent onto the CUDA scheduler after completion.
DebugCudaScheduler is the deterministic first unified implementation. It derives from DebugScheduler, implements CUDA admission, and uses one runnable queue for ordinary and CUDA tasks. Ordinary activations run in the host domain. Each CUDA activation resolves an effective device from the affinity stored in CudaTaskPromise or the scheduler's default device. It establishes that device with cuda::ScopedDevice and restores the calling thread before the next activation. One scheduler can therefore drive host work and CUDA work for every visible device, including from get_wait().
Tests cover explicit and unbound CUDA admission, explicit device changes after execution starts, resumption after buffer suspension, multi-device tasks in one queue, explicit host/CUDA nesting in both directions, one get_wait() driving a complete host/CUDA/host continuation chain, same-device direct transfer, cross-device resubmission, exception propagation, cancellation, and calling-thread device restoration. A separate non-CUDA death test proves that cross-domain inheritance fails when the parent scheduler does not accept the child route.
TbbCudaScheduler is the parallel implementation. It extends TbbScheduler with one worker-only oneTBB arena per enrolled CUDA device. The host arena and all device arenas share one scheduler route, task group, pause state, wait state. The default constructor enrolls every visible device; tests may provide an explicit device set.
Each device arena has a task_scheduler_observer. The observer saves and selects that arena's device whenever a worker participates, then restores the previous device on exit. The application thread waits and helps through the host arena; it does not enter a CUDA arena merely to admit or wait for CUDA work. Same-device nested tasks may transfer directly while already inside the correct arena. Cross-domain or cross-device continuations are enqueued through the shared scheduler and routed to their required arena.
run_all() and get_wait() use the shared task group, so one call observes all currently admitted host and CUDA activations. They retain the same activation-quiescence meaning as TbbScheduler: a coroutine suspended on an external event may become runnable and submit a later activation.
Tests cover host, explicitly bound CUDA, and unbound CUDA admission through one scheduler; explicit migration between device arenas; synchronous waits that drive CUDA work; non-blocking resumption into a saturated device arena; concurrent worker device selection before and after suspension; one scheduler routing tasks across every visible device; cross-device continuation return; cross-domain exception and cancellation propagation; and calling-thread device restoration. Multi-device cases skip when fewer than two devices are visible.
CudaTaskPromise contains an optional CUDA device ordinal. An empty value means that the task has not established device affinity; it does not make the task invalid. Device-neutral code may execute in this state, and the scheduler uses its default CUDA device as the activation context without writing that default into the promise.
Device-sensitive code selects affinity in either of two ways:
co_await cuda::set_device(device) records or changes affinity while the coroutine is suspended. If the scheduler-established current CUDA device already matches, the coroutine continues immediately. Otherwise it is resubmitted through the unified scheduler so its next activation occurs in the selected device context. The candidate device is validated against the current scheduler before promise affinity changes.An unbound CUDA child nested under a CUDA parent inherits the parent's affinity when the parent has one. If both are unbound, both use the scheduler default and remain unbound. Direct transfer compares effective activation devices, not just the two optional promise values.
The device ordinal is CUDA-specific promise state. Generic task machinery sees only the immutable domain tag stored in TaskHandle; after checking that tag, CUDA-aware code may narrow the common promise pointer to CudaTaskPromise. Streams, provider handles, workspaces, and buffer leases remain operation-local RAII values rather than task-global state.
The global schedule(CudaTask&&, int) overload selects explicit initial device affinity and avoids a first-activation migration when the device is already known. Tensor operations whose device is only available after awaiting their operand epochs may schedule an unbound task and select the authoritative storage device inside the coroutine. Automatic pre-admission inference should remain above the shared coroutine machinery and must not inspect a pending Tensor value unsafely.
Device migration within one unified CUDA scheduler is implemented by cuda::set_device. It is legal only at that explicit suspension point. A task must not hold a stream, provider handle, workspace, or other device-local lease across the operation.
The awaiter uses cudaGetDevice() to inspect the execution context already established by DebugCudaScheduler or the oneTBB device-arena observer. This is not an independent placement decision: a scheduler that resumes a CUDA task without first establishing the correct CUDA device violates the scheduler contract. No second scheduler-specific TLS registry is required.
Cross-device selection is always enqueued through the unified scheduler. It does not call task_arena::execute() to enter another device arena synchronously. That oneTBB operation may wait when immediate arena admission is unavailable and may let the calling thread participate in the target arena; either behavior would make an ordinary CUDA scheduling point unexpectedly blocking and complicate continuation ownership. Same-effective-device nested CUDA tasks already use direct coroutine transfer without entering another arena.
The promise-neutral route machinery permits scheduler migration whenever a suspended task is exclusively owned outside a scheduler. A user-facing migration awaiter is not yet implemented; when added, it should be an explicit suspension operation such as:
The awaiter would validate the target scheduler, replace the route in the shared promise, and submit exactly one activation. It must not change the coroutine's concrete return type, move Tensor storage, or silently change operand placement.
Before adding this capability, define and test:
A CPU parent can enter a CUDA domain by awaiting a separately admitted CudaTask, which keeps the transition explicit and returns naturally to CPU control flow.
A task remembers which scheduler should run its next activation, but it does not own that scheduler or keep it alive. The scheduler must remain alive until every task routed through it has completed, or has been cancelled so that it can never become runnable again. The same rule applies to resource waiters and completion callbacks that may publish a later activation.
The reverse is also important: the scheduler does not own a coroutine merely because its promise remembers that scheduler. An initially suspended public task is owned by its task object, and an externally suspended task is owned by its awaiter. The scheduler sees either only a queued activation or no live activation at all.
Application schedulers should normally be long-lived runtime services, often lasting until process shutdown. A test may use a stack-local scheduler, but it must finish or cancel every routed task before leaving the scheduler's scope and ensure that no external callback can submit a later resumption.
run_all() is scheduler-local activation quiescence. It does not imply that an externally suspended task has completed, nor does it drain every scheduler in a multi-device process. In particular, waiting for the current oneTBB task group does not make it safe to destroy a scheduler while a coroutine remains suspended on an epoch, resource, or external event. Process-wide CUDA shutdown diagnoses live buffers, but scheduler shutdown still needs an explicit contract that diagnoses outstanding routed tasks and resource waiters.
The unified debug and TBB tests establish the task-domain route. The next checkpoints should add:
1.9.8