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::{DType, Shape};
19
20/// A buffer that knows where its bytes live.
21#[derive(Debug, Clone)]
22pub struct Buffer {
23    bytes: Vec<u8>,
24    shape: Shape,
25    device: Device,
26}
27
28impl Buffer {
29    /// Create a buffer holding `data`, tagged as residing on `device`.
30    /// On Cpu this is just storage; on Metal callers will use
31    /// `to_device` to move it into device memory once the runtime
32    /// integrates with the arena.
33    pub fn new_host(shape: Shape, data: Vec<u8>) -> Self {
34        Self {
35            bytes: data,
36            shape,
37            device: Device::Cpu,
38        }
39    }
40
41    /// Create a zero-initialized host buffer of `shape`.
42    pub fn zeros(shape: Shape) -> Self {
43        let n = shape.size_bytes().unwrap_or(0);
44        Self {
45            bytes: vec![0u8; n],
46            shape,
47            device: Device::Cpu,
48        }
49    }
50
51    pub fn shape(&self) -> &Shape {
52        &self.shape
53    }
54    pub fn device(&self) -> Device {
55        self.device
56    }
57    pub fn dtype(&self) -> DType {
58        self.shape.dtype()
59    }
60    pub fn num_elements(&self) -> usize {
61        self.shape.num_elements().unwrap_or(0)
62    }
63    pub fn byte_size(&self) -> usize {
64        self.bytes.len()
65    }
66
67    /// Read as `&[f32]`. Panics if dtype isn't F32 — explicit type
68    /// check (plan #59 spirit: don't let mismatches go silent).
69    pub fn as_f32(&self) -> &[f32] {
70        assert_eq!(self.dtype(), DType::F32, "as_f32 on non-F32 buffer");
71        let n = self.num_elements();
72        unsafe { std::slice::from_raw_parts(self.bytes.as_ptr() as *const f32, n) }
73    }
74
75    pub fn as_f32_mut(&mut self) -> &mut [f32] {
76        assert_eq!(self.dtype(), DType::F32, "as_f32_mut on non-F32 buffer");
77        let n = self.num_elements();
78        unsafe { std::slice::from_raw_parts_mut(self.bytes.as_mut_ptr() as *mut f32, n) }
79    }
80
81    /// "Move to device" — explicit transfer call. CPU is a no-op
82    /// (the bytes are already where they need to be). Metal routes
83    /// through the backend (TODO: wire to rlx-metal::arena once a
84    /// caller needs it).
85    pub fn to_device(self, device: Device) -> Self {
86        match (self.device, device) {
87            (a, b) if a == b => self,
88            (Device::Cpu, Device::Metal) => Self {
89                device: Device::Metal,
90                ..self
91            },
92            (Device::Metal, Device::Cpu) => Self {
93                device: Device::Cpu,
94                ..self
95            },
96            _ => self,
97        }
98    }
99
100    /// "Read back to host" — explicit transfer. CPU no-op; Metal
101    /// blocks on completion + memcpy back from MTLBuffer (TODO).
102    pub fn to_host(self) -> Self {
103        self.to_device(Device::Cpu)
104    }
105
106    /// Raw bytes — host-side access. Panics if the buffer is on a
107    /// non-host device (would silently read uninitialized host
108    /// memory otherwise).
109    pub fn host_bytes(&self) -> &[u8] {
110        assert_eq!(
111            self.device,
112            Device::Cpu,
113            "host_bytes on non-host buffer; call .to_host() first"
114        );
115        &self.bytes
116    }
117}
118
119#[cfg(test)]
120mod tests {
121    use super::*;
122
123    #[test]
124    fn zeros_initializes() {
125        let b = Buffer::zeros(Shape::new(&[2, 3], DType::F32));
126        assert_eq!(b.num_elements(), 6);
127        assert_eq!(b.byte_size(), 24);
128        for v in b.as_f32() {
129            assert_eq!(*v, 0.0);
130        }
131    }
132
133    #[test]
134    fn dtype_mismatch_panics() {
135        let b = Buffer::zeros(Shape::new(&[4], DType::I32));
136        let result = std::panic::catch_unwind(|| b.as_f32());
137        assert!(result.is_err());
138    }
139
140    #[test]
141    fn to_device_round_trip() {
142        let b = Buffer::zeros(Shape::new(&[4], DType::F32));
143        let on_metal = b.to_device(Device::Metal);
144        assert_eq!(on_metal.device(), Device::Metal);
145        let back = on_metal.to_host();
146        assert_eq!(back.device(), Device::Cpu);
147    }
148}
rlx_driver/buffer.rs

rlx_driver/
buffer.rs
