rlx_driver/
handle.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//! Persistent buffer handles — state that survives across forward passes.
6//!
7//! For inference of stateful models (KV-cache, beam search) and for
8//! training (gradient accumulators, optimizer state), the runtime needs
9//! buffers that persist beyond a single `compiled.run()`. The arena
10//! is rebuilt every compile, so it can't carry state.
11//!
12//! `BufferHandle` is an opaque, stable identifier the user creates once
13//! and binds at compile time. The backend allocates a separate "handles"
14//! region (independent of the arena) and routes reads/writes there.
15//!
16//! Workflow:
17//!
18//! ```rust,ignore
19//! let kv_cache = BufferHandle::new("kv", &[batch, max_seq, num_heads, head_dim]);
20//! let session = Session::new(Device::Metal);
21//! let mut compiled = session.compile_with(graph, &CompileOptions::new()
22//!     .bind_handle(&kv_cache));
23//!
24//! for token in tokens {
25//!     compiled.bind_handle("kv", &kv_cache_data); // initial value
26//!     let logits = compiled.run(&[("token", &[token])]);
27//!     kv_cache_data = compiled.read_handle("kv").unwrap();
28//! }
29//! ```
30
31use rlx_ir::{DType, Shape};
32
33/// External, persistent buffer reference. Created once, bound at compile,
34/// carried across many `compiled.run()` invocations.
35#[derive(Debug, Clone)]
36pub struct BufferHandle {
37    pub name: String,
38    pub shape: Shape,
39}
40
41impl BufferHandle {
42    pub fn new(name: impl Into<String>, dims: &[usize], dtype: DType) -> Self {
43        Self {
44            name: name.into(),
45            shape: Shape::new(dims, dtype),
46        }
47    }
48
49    /// Total byte size of this handle.
50    pub fn byte_size(&self) -> usize {
51        self.shape.size_bytes().unwrap_or(0)
52    }
53
54    /// Number of elements (for slicing as &[f32] etc.).
55    pub fn num_elements(&self) -> usize {
56        self.shape.num_elements().unwrap_or(0)
57    }
58}
rlx_driver/handle.rs

rlx_driver/
handle.rs
