rlx_driver/arena.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//! Backend-agnostic arena trait — the contract every backend's memory
6//! plan obeys.
7//!
8//! Lifted from CpuExecutable / MetalExecutable's previously duplicated
9//! arena helpers. Each new backend (CUDA, ROCm, wgpu, WASM, TPU) implements
10//! this trait once and gets:
11//! - typed input feed (`f32 → arena_dtype`)
12//! - typed output read (`arena_dtype → f32`)
13//! - per-node byte offset resolution
14//!
15//! The trait deliberately exposes raw pointers / byte offsets rather than
16//! Rust slices so the same implementation works for host-resident memory
17//! (CPU/WASM), unified memory (Apple Silicon Metal/MPSGraph), and
18//! discrete-VRAM backends (CUDA/ROCm) where reading involves a copy.
19
20use rlx_ir::{DType, NodeId};
21
22/// Per-backend arena interface.
23///
24/// All concrete arenas — `rlx-cpu::Arena`, `rlx-metal::Arena`, future
25/// `rlx-cuda::Arena`, `rlx-wgpu::Arena` — implement this trait so the
26/// runtime can drive them uniformly. The actual byte layout is owned
27/// by the backend; we only require offset-based access.
28pub trait DeviceArena {
29 /// Byte offset of `id`'s buffer slot in the arena. `usize::MAX` for
30 /// nodes that don't have an arena slot (e.g. fused-away intermediates).
31 fn byte_offset(&self, id: NodeId) -> usize;
32
33 /// True if `id` has a real arena slot.
34 fn has_buffer(&self, id: NodeId) -> bool;
35
36 /// Total arena size in bytes.
37 fn size_bytes(&self) -> usize;
38
39 /// Write a host-side `f32` slice into `id`'s slot, casting to `dtype`
40 /// if necessary. Truncates to the buffer's capacity (no panic on overflow).
41 ///
42 /// For discrete-memory backends this involves a host→device copy; for
43 /// unified-memory backends (Apple Silicon, integrated GPUs) it's a
44 /// direct write.
45 fn write_input_f32(&mut self, id: NodeId, dtype: DType, data: &[f32]);
46
47 /// Read `id`'s slot as a host-side `Vec<f32>`, casting from `dtype` if
48 /// necessary. The number of elements is determined by the backend
49 /// based on the memory plan (typically `shape.num_elements()`).
50 fn read_output_f32(&self, id: NodeId, dtype: DType, n_elements: usize) -> Vec<f32>;
51}
52
53/// Helper: cast f32 input to bytes of `dtype` and write to `dst_ptr`.
54/// Used by every CPU-resident-arena backend. GPU backends can call this
55/// after staging into a host buffer, then upload.
56///
57/// Currently supports F32 / F64 / F16 / BF16 / C64 / C128. Other dtypes
58/// fall through to F32.
59pub unsafe fn write_typed_from_f32(dst_ptr: *mut u8, dtype: DType, src: &[f32], max_elems: usize) {
60 let n = src.len().min(max_elems);
61 match dtype {
62 DType::F64 => unsafe {
63 // F64 slots are 8 B/elem; widen the f32 input. (Values carry
64 // f32 precision — this is the f32 entry path; `run_typed`'s
65 // `all_f64` branch is the full-precision route.)
66 let dst = dst_ptr as *mut f64;
67 for i in 0..n {
68 *dst.add(i) = src[i] as f64;
69 }
70 },
71 DType::F16 => unsafe {
72 let dst = dst_ptr as *mut half::f16;
73 for i in 0..n {
74 *dst.add(i) = half::f16::from_f32(src[i]);
75 }
76 },
77 DType::BF16 => unsafe {
78 let dst = dst_ptr as *mut half::bf16;
79 for i in 0..n {
80 *dst.add(i) = half::bf16::from_f32(src[i]);
81 }
82 },
83 DType::C64 => unsafe {
84 // Interleaved [re, im, re, im, ...]; `max_elems` is complex count.
85 let dst = dst_ptr as *mut f32;
86 let n = src.len().min(max_elems.saturating_mul(2));
87 std::ptr::copy_nonoverlapping(src.as_ptr(), dst, n);
88 },
89 DType::C128 => unsafe {
90 // Interleaved [re, im, ...] f64 pairs (16 B/elem); `max_elems`
91 // is complex count. This is the f32 entry path: widen each
92 // interleaved f32 lane to f64 (values carry f32 precision).
93 let dst = dst_ptr as *mut f64;
94 let n = src.len().min(max_elems.saturating_mul(2));
95 for i in 0..n {
96 *dst.add(i) = src[i] as f64;
97 }
98 },
99 _ => unsafe {
100 let dst = dst_ptr as *mut f32;
101 std::ptr::copy_nonoverlapping(src.as_ptr(), dst, n);
102 },
103 }
104}
105
106/// Helper: read `n_elems` of `dtype` from `src_ptr`, returning `Vec<f32>`.
107pub unsafe fn read_typed_to_f32(src_ptr: *const u8, dtype: DType, n_elems: usize) -> Vec<f32> {
108 match dtype {
109 DType::F64 => {
110 // F64 slots are 8 B/elem; narrow to f32 for the f32 read path.
111 let mut out = Vec::with_capacity(n_elems);
112 unsafe {
113 let src = src_ptr as *const f64;
114 for i in 0..n_elems {
115 out.push(*src.add(i) as f32);
116 }
117 }
118 out
119 }
120 DType::F16 => {
121 let mut out = Vec::with_capacity(n_elems);
122 unsafe {
123 let src = src_ptr as *const half::f16;
124 for i in 0..n_elems {
125 out.push((*src.add(i)).to_f32());
126 }
127 }
128 out
129 }
130 DType::BF16 => {
131 let mut out = Vec::with_capacity(n_elems);
132 unsafe {
133 let src = src_ptr as *const half::bf16;
134 for i in 0..n_elems {
135 out.push((*src.add(i)).to_f32());
136 }
137 }
138 out
139 }
140 DType::C64 => unsafe {
141 // Interleaved [re, im, re, im, ...]; `n_elems` is complex count.
142 let src = src_ptr as *const f32;
143 std::slice::from_raw_parts(src, n_elems.saturating_mul(2)).to_vec()
144 },
145 DType::C128 => unsafe {
146 // Interleaved [re, im, ...] f64 pairs (16 B/elem); `n_elems` is
147 // complex count. f32 read path: narrow each f64 lane to f32.
148 let src = src_ptr as *const f64;
149 let lanes = n_elems.saturating_mul(2);
150 let mut out = Vec::with_capacity(lanes);
151 for i in 0..lanes {
152 out.push(*src.add(i) as f32);
153 }
154 out
155 },
156 _ => unsafe {
157 let src = src_ptr as *const f32;
158 std::slice::from_raw_parts(src, n_elems).to_vec()
159 },
160 }
161}