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}