rlx_driver/
buffer.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//! Explicit host ↔ device buffer abstraction (plan #59).
6//!
7//! Borrowed from MAX's `max/python/max/driver/buffer.py`. The point
8//! is *explicitness*: every host↔device transfer is a method call,
9//! not implicit conversion. Misreading "device" data as host (or
10//! vice versa) becomes a compile error, not a silent perf bug.
11//!
12//! Today this wraps a host-side `Vec<u8>` plus a `Device` tag. The
13//! `to_host` / `to_device` calls are no-ops on Device::Cpu; on
14//! Device::Metal a future commit will route through the existing
15//! `rlx-metal::arena` for actual MTLBuffer transfers.
16
17use crate::Device;
18use rlx_ir::bytes::AlignedBytes;
19use rlx_ir::{DType, Shape};
20
21/// A buffer that knows where its bytes live.
22///
23/// Storage is [`AlignedBytes`] rather than `Vec<u8>` because [`Buffer::as_f32`]
24/// hands out a borrowed `&[f32]` over these bytes, and `Vec<u8>` only promises
25/// alignment 1 — reinterpreting it as `&[f32]` is UB even though the system
26/// allocator happens to return 16-aligned blocks today.
27#[derive(Debug, Clone)]
28pub struct Buffer {
29    bytes: AlignedBytes,
30    shape: Shape,
31    device: Device,
32}
33
34impl Buffer {
35    /// Create a buffer holding `data`, tagged as residing on `device`.
36    /// On Cpu this is just storage; on Metal callers will use
37    /// `to_device` to move it into device memory once the runtime
38    /// integrates with the arena.
39    pub fn new_host(shape: Shape, data: Vec<u8>) -> Self {
40        Self {
41            bytes: AlignedBytes::from_slice(&data),
42            shape,
43            device: Device::Cpu,
44        }
45    }
46
47    /// Create a zero-initialized host buffer of `shape`.
48    pub fn zeros(shape: Shape) -> Self {
49        let n = shape.size_bytes().unwrap_or(0);
50        Self {
51            bytes: AlignedBytes::zeroed(n),
52            shape,
53            device: Device::Cpu,
54        }
55    }
56
57    pub fn shape(&self) -> &Shape {
58        &self.shape
59    }
60    pub fn device(&self) -> Device {
61        self.device
62    }
63    pub fn dtype(&self) -> DType {
64        self.shape.dtype()
65    }
66    pub fn num_elements(&self) -> usize {
67        self.shape.num_elements().unwrap_or(0)
68    }
69    pub fn byte_size(&self) -> usize {
70        self.bytes.len()
71    }
72
73    /// Read as `&[f32]`. Panics if dtype isn't F32 — explicit type
74    /// check (plan #59 spirit: don't let mismatches go silent).
75    pub fn as_f32(&self) -> &[f32] {
76        assert_eq!(self.dtype(), DType::F32, "as_f32 on non-F32 buffer");
77        let n = self.num_elements();
78        // Sound without an alignment check: `AlignedBytes` pins the base at 64
79        // bytes, which covers `align_of::<f32>()`.
80        &self.bytes.as_slice_of::<f32>()[..n]
81    }
82
83    pub fn as_f32_mut(&mut self) -> &mut [f32] {
84        assert_eq!(self.dtype(), DType::F32, "as_f32_mut on non-F32 buffer");
85        let n = self.num_elements();
86        &mut self.bytes.slice_of_mut::<f32>(0, self.bytes.len())[..n]
87    }
88
89    /// "Move to device" — explicit transfer call. CPU is a no-op
90    /// (the bytes are already where they need to be). Metal routes
91    /// through the backend (TODO: wire to rlx-metal::arena once a
92    /// caller needs it).
93    pub fn to_device(self, device: Device) -> Self {
94        match (self.device, device) {
95            (a, b) if a == b => self,
96            (Device::Cpu, Device::Metal) => Self {
97                device: Device::Metal,
98                ..self
99            },
100            (Device::Metal, Device::Cpu) => Self {
101                device: Device::Cpu,
102                ..self
103            },
104            _ => self,
105        }
106    }
107
108    /// "Read back to host" — explicit transfer. CPU no-op; Metal
109    /// blocks on completion + memcpy back from MTLBuffer (TODO).
110    pub fn to_host(self) -> Self {
111        self.to_device(Device::Cpu)
112    }
113
114    /// Raw bytes — host-side access. Panics if the buffer is on a
115    /// non-host device (would silently read uninitialized host
116    /// memory otherwise).
117    pub fn host_bytes(&self) -> &[u8] {
118        assert_eq!(
119            self.device,
120            Device::Cpu,
121            "host_bytes on non-host buffer; call .to_host() first"
122        );
123        &self.bytes
124    }
125}
126
127#[cfg(test)]
128mod tests {
129    use super::*;
130
131    #[test]
132    fn zeros_initializes() {
133        let b = Buffer::zeros(Shape::new(&[2, 3], DType::F32));
134        assert_eq!(b.num_elements(), 6);
135        assert_eq!(b.byte_size(), 24);
136        for v in b.as_f32() {
137            assert_eq!(*v, 0.0);
138        }
139    }
140
141    #[test]
142    fn dtype_mismatch_panics() {
143        let b = Buffer::zeros(Shape::new(&[4], DType::I32));
144        let result = std::panic::catch_unwind(|| b.as_f32());
145        assert!(result.is_err());
146    }
147
148    #[test]
149    fn to_device_round_trip() {
150        let b = Buffer::zeros(Shape::new(&[4], DType::F32));
151        let on_metal = b.to_device(Device::Metal);
152        assert_eq!(on_metal.device(), Device::Metal);
153        let back = on_metal.to_host();
154        assert_eq!(back.device(), Device::Cpu);
155    }
156}
rlx_driver/buffer.rs

rlx_driver/
buffer.rs
