Skip to main content

rlx_driver/
stream.rs

1// RLX — versatile ML compiler + runtime.
2// Copyright (C) 2026 Eugene Hauptmann, Nataliya Kosmyna.
3//
4// This program is free software: you can redistribute it and/or modify
5// it under the terms of the GNU General Public License as published by
6// the Free Software Foundation, version 3.
7//
8// This program is distributed in the hope that it will be useful,
9// but WITHOUT ANY WARRANTY; without even the implied warranty of
10// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
11// GNU General Public License for more details.
12//
13// You should have received a copy of the GNU General Public License
14// along with this program. If not, see <https://www.gnu.org/licenses/>.
15
16//! Command-stream abstraction.
17//!
18//! Every GPU-shaped backend has the same shape: enqueue work, submit, wait.
19//!   - Metal: `MTLCommandBuffer` + `commit` + `waitUntilCompleted`
20//!   - CUDA: `cudaStream_t` + `cudaStreamSynchronize`
21//!   - ROCm: `hipStream_t` + `hipStreamSynchronize`
22//!   - wgpu: `CommandEncoder.finish()` → `Queue.submit()` → `Device.poll(Wait)`
23//!   - WASM (single-threaded): no-op (work runs synchronously)
24//!
25//! Hoisting this into one trait means:
26//!   - the runtime can drive *any* backend via the same submit-and-wait API
27//!   - new backends only need a thin command-stream impl
28//!   - test infrastructure works against the trait, not per-backend types
29
30/// Per-backend command stream.
31///
32/// Implementations are free to be no-ops on synchronous backends (host CPU,
33/// WASM): `submit` runs work eagerly, `wait` returns immediately.
34pub trait CommandStream {
35    /// Submit any pending work to the device (non-blocking).
36    fn submit(&mut self);
37
38    /// Block until all submitted work has completed.
39    fn wait(&mut self);
40
41    /// Convenience: `submit` followed by `wait`. Backends may override
42    /// for a fused fast path (e.g. Metal's `commit` + `waitUntilCompleted`).
43    fn submit_and_wait(&mut self) {
44        self.submit();
45        self.wait();
46    }
47}
48
49/// Default implementation for synchronous backends — work has already
50/// happened by the time `submit` is called.
51pub struct SyncStream;
52
53impl CommandStream for SyncStream {
54    fn submit(&mut self) {}
55    fn wait(&mut self) {}
56}