rlx_driver/
stream.rs

1// RLX — versatile ML compiler + runtime.
2// Copyright (C) 2026 Eugene Hauptmann, Nataliya Kosmyna.
3// SPDX-License-Identifier: MIT OR Apache-2.0
4
5//! Command-stream abstraction.
6//!
7//! Every GPU-shaped backend has the same shape: enqueue work, submit, wait.
8//!   - Metal: `MTLCommandBuffer` + `commit` + `waitUntilCompleted`
9//!   - CUDA: `cudaStream_t` + `cudaStreamSynchronize`
10//!   - ROCm: `hipStream_t` + `hipStreamSynchronize`
11//!   - wgpu: `CommandEncoder.finish()` → `Queue.submit()` → `Device.poll(Wait)`
12//!   - WASM (single-threaded): no-op (work runs synchronously)
13//!
14//! Hoisting this into one trait means:
15//!   - the runtime can drive *any* backend via the same submit-and-wait API
16//!   - new backends only need a thin command-stream impl
17//!   - test infrastructure works against the trait, not per-backend types
18
19/// Per-backend command stream.
20///
21/// Implementations are free to be no-ops on synchronous backends (host CPU,
22/// WASM): `submit` runs work eagerly, `wait` returns immediately.
23pub trait CommandStream {
24    /// Submit any pending work to the device (non-blocking).
25    fn submit(&mut self);
26
27    /// Block until all submitted work has completed.
28    fn wait(&mut self);
29
30    /// Convenience: `submit` followed by `wait`. Backends may override
31    /// for a fused fast path (e.g. Metal's `commit` + `waitUntilCompleted`).
32    fn submit_and_wait(&mut self) {
33        self.submit();
34        self.wait();
35    }
36}
37
38/// Default implementation for synchronous backends — work has already
39/// happened by the time `submit` is called.
40pub struct SyncStream;
41
42impl CommandStream for SyncStream {
43    fn submit(&mut self) {}
44    fn wait(&mut self) {}
45}
rlx_driver/stream.rs

rlx_driver/
stream.rs
