onnx_runtime_eager/
tensor.rs

1//! The owned, device-aware [`Tensor`] the eager engine dispatches over
2//! (`docs/EAGER.md` §3, §10.1).
3//!
4//! This mirrors the owned tensor in `onnx-runtime-session` (`src/tensor.rs`):
5//! a thin owner over an [`onnx_runtime_ep_api::DeviceBuffer`] plus the IR
6//! vocabulary ([`DataType`], [`TensorLayout`], shape). It is duplicated here
7//! rather than shared because the eager crate must not depend on the session
8//! crate (see `docs/EAGER.md` §12 crate layout). When the tensor type is hoisted
9//! into a shared `onnx-runtime-tensor` crate (the design's open question), both
10//! this crate and the session collapse onto it.
11//!
12//! ## The single `unsafe` seam
13//!
14//! A [`DeviceBuffer`] hands out only raw base pointers; reading/writing the
15//! bytes is `unsafe` and sound only on host-accessible devices. Every such
16//! access funnels through [`host_bytes`] / [`write_host`], which assert host
17//! accessibility, so the rest of the crate is safe Rust over the EP contract.
18//!
19//! DLPack / numpy interop (`docs/EAGER.md` §3) is DEFERRED — those are binding
20//! concerns handled by the (also DEFERRED) PyO3 layer (§11).
21
22use std::ffi::c_void;
23use std::sync::Arc;
24use std::sync::OnceLock;
25
26use onnx_runtime_ep_api::{DeviceBuffer, ExecutionProvider};
27use onnx_runtime_ep_cpu::CpuExecutionProvider;
28use onnx_runtime_ir::{DataType, DeviceId, TensorLayout};
29
30use crate::error::{EagerError, Result};
31
32/// A process-wide, already-initialized CPU execution provider used to back
33/// user-constructed [`Tensor`]s. Host `malloc`/`free` is global, so any
34/// `CpuExecutionProvider` can free any other's CPU allocation (mirrors
35/// `onnx-runtime-session`'s shared CPU EP).
36fn shared_cpu_ep() -> Arc<dyn ExecutionProvider> {
37    static EP: OnceLock<Arc<CpuExecutionProvider>> = OnceLock::new();
38    EP.get_or_init(|| {
39        let mut ep = CpuExecutionProvider::new();
40        // Pure-Rust CPU EP: `initialize` only flips a flag and never fails.
41        let _ = ep.initialize(&Default::default());
42        Arc::new(ep)
43    })
44    .clone()
45}
46
47/// Borrow the raw bytes of a host-accessible device buffer.
48///
49/// # Panics
50///
51/// Panics if `buffer` is not on a host-accessible device.
52fn host_bytes(buffer: &DeviceBuffer) -> &[u8] {
53    assert!(
54        buffer.device().is_host_accessible(),
55        "host_bytes on non-host device {:?}",
56        buffer.device()
57    );
58    if buffer.is_empty() {
59        return &[];
60    }
61    // SAFETY: host-accessible device (asserted) means `as_ptr` is a real,
62    // readable host address; the EP guarantees `len()` valid bytes behind it.
63    // The lifetime is tied to `&buffer`, so the slice cannot dangle.
64    unsafe { std::slice::from_raw_parts(buffer.as_ptr() as *const u8, buffer.len()) }
65}
66
67/// Copy `src` into the front of a host-accessible device buffer.
68fn write_host(buffer: &mut DeviceBuffer, src: &[u8]) -> Result<()> {
69    assert!(
70        buffer.device().is_host_accessible(),
71        "write_host on non-host device {:?}",
72        buffer.device()
73    );
74    if src.len() > buffer.len() {
75        return Err(EagerError::Kernel(
76            onnx_runtime_ep_api::EpError::KernelFailed(format!(
77                "write_host: source {} bytes exceeds buffer {} bytes",
78                src.len(),
79                buffer.len()
80            )),
81        ));
82    }
83    if src.is_empty() {
84        return Ok(());
85    }
86    let dst = buffer.as_mut_ptr() as *mut u8;
87    // SAFETY: host-accessible device (asserted); `dst` is a unique writable host
88    // pointer obtained via `&mut buffer` (no alias), with at least `src.len()`
89    // bytes of capacity (checked above). `src` is a distinct owned slice.
90    unsafe {
91        std::ptr::copy_nonoverlapping(src.as_ptr(), dst, src.len());
92    }
93    Ok(())
94}
95
96/// An owned, device-aware tensor (`docs/EAGER.md` §3).
97///
98/// Owns the [`DeviceBuffer`] that holds its elements and the EP that must free
99/// it. On Phase-1 CPU the buffer is a host allocation, so [`Tensor::as_bytes`]
100/// and the typed accessors read it directly.
101pub struct Tensor {
102    dtype: DataType,
103    shape: Vec<usize>,
104    layout: TensorLayout,
105    device: DeviceId,
106    /// `Some` while the tensor is live; taken by [`Drop`] to free exactly once.
107    buffer: Option<DeviceBuffer>,
108    /// The EP that allocated [`Tensor::buffer`] and must deallocate it.
109    allocator: Arc<dyn ExecutionProvider>,
110}
111
112impl Tensor {
113    /// Allocate a tensor from raw little-endian element bytes using `allocator`.
114    ///
115    /// `bytes` must hold exactly `storage_bytes(numel)` bytes for `dtype` and
116    /// `shape`.
117    pub fn from_raw_in(
118        allocator: Arc<dyn ExecutionProvider>,
119        dtype: DataType,
120        shape: Vec<usize>,
121        bytes: &[u8],
122    ) -> Result<Self> {
123        let numel: usize = shape.iter().product();
124        let expected = dtype.storage_bytes(numel);
125        if bytes.len() != expected {
126            return Err(EagerError::Kernel(
127                onnx_runtime_ep_api::EpError::KernelFailed(format!(
128                    "Tensor::from_raw_in: {} bytes for shape {shape:?} dtype {dtype:?}, expected {expected}",
129                    bytes.len()
130                )),
131            ));
132        }
133        let layout = TensorLayout::contiguous();
134        let mut buffer = allocator.allocate(expected.max(1), layout.alignment)?;
135        write_host(&mut buffer, bytes)?;
136        Ok(Self {
137            dtype,
138            shape,
139            layout,
140            device: buffer.device(),
141            buffer: Some(buffer),
142            allocator,
143        })
144    }
145
146    /// Allocate a zero-initialized tensor of `shape`/`dtype` using `allocator`.
147    ///
148    /// Used by dispatch to materialise output tensors before kernel execution
149    /// (`docs/EAGER.md` §10.1 step 6).
150    pub fn zeros_in(
151        allocator: Arc<dyn ExecutionProvider>,
152        dtype: DataType,
153        shape: Vec<usize>,
154    ) -> Result<Self> {
155        let numel: usize = shape.iter().product();
156        let bytes = vec![0u8; dtype.storage_bytes(numel)];
157        Self::from_raw_in(allocator, dtype, shape, &bytes)
158    }
159
160    /// Build a tensor from raw little-endian bytes on the shared CPU device.
161    pub fn from_raw(dtype: DataType, shape: Vec<usize>, bytes: &[u8]) -> Result<Self> {
162        Self::from_raw_in(shared_cpu_ep(), dtype, shape, bytes)
163    }
164
165    /// Build a zero-initialized tensor on the shared CPU device.
166    pub fn zeros(dtype: DataType, shape: Vec<usize>) -> Result<Self> {
167        Self::zeros_in(shared_cpu_ep(), dtype, shape)
168    }
169
170    /// Build an `f32` tensor from a dense row-major slice on the CPU device.
171    pub fn from_f32(shape: &[usize], data: &[f32]) -> Result<Self> {
172        let mut bytes = Vec::with_capacity(data.len() * 4);
173        for v in data {
174            bytes.extend_from_slice(&v.to_le_bytes());
175        }
176        Self::from_raw(DataType::Float32, shape.to_vec(), &bytes)
177    }
178
179    /// Build an `i64` tensor from a dense row-major slice on the CPU device.
180    pub fn from_i64(shape: &[usize], data: &[i64]) -> Result<Self> {
181        let mut bytes = Vec::with_capacity(data.len() * 8);
182        for v in data {
183            bytes.extend_from_slice(&v.to_le_bytes());
184        }
185        Self::from_raw(DataType::Int64, shape.to_vec(), &bytes)
186    }
187
188    /// The element type.
189    pub fn dtype(&self) -> DataType {
190        self.dtype
191    }
192
193    /// The logical shape (static dims).
194    pub fn shape(&self) -> &[usize] {
195        &self.shape
196    }
197
198    /// The physical layout (row-major contiguous for tensors this crate produces).
199    pub fn layout(&self) -> &TensorLayout {
200        &self.layout
201    }
202
203    /// The device this tensor lives on.
204    pub fn device(&self) -> DeviceId {
205        self.device
206    }
207
208    /// Number of elements.
209    pub fn numel(&self) -> usize {
210        self.shape.iter().product()
211    }
212
213    fn buffer(&self) -> &DeviceBuffer {
214        self.buffer
215            .as_ref()
216            .expect("Tensor buffer taken only in Drop")
217    }
218
219    /// Shared raw base pointer to the element storage. Safe to obtain;
220    /// dereferencing is `unsafe` and sound only within the owning EP's context.
221    pub(crate) fn device_ptr(&self) -> *const c_void {
222        self.buffer().as_ptr()
223    }
224
225    /// Unique raw base pointer to the element storage.
226    pub(crate) fn device_ptr_mut(&mut self) -> *mut c_void {
227        self.buffer
228            .as_mut()
229            .expect("Tensor buffer taken only in Drop")
230            .as_mut_ptr()
231    }
232
233    /// Borrow the raw little-endian element bytes (host tensors only).
234    pub fn as_bytes(&self) -> &[u8] {
235        let n = self.dtype.storage_bytes(self.numel());
236        &host_bytes(self.buffer())[..n]
237    }
238
239    /// Copy out the elements as `f32`. Panics if the dtype is not `Float32`.
240    pub fn to_vec_f32(&self) -> Vec<f32> {
241        assert_eq!(
242            self.dtype,
243            DataType::Float32,
244            "to_vec_f32 on non-f32 tensor"
245        );
246        self.as_bytes()
247            .chunks_exact(4)
248            .map(|c| f32::from_le_bytes([c[0], c[1], c[2], c[3]]))
249            .collect()
250    }
251
252    /// Copy out the elements as `i64`. Panics if the dtype is not `Int64`.
253    pub fn to_vec_i64(&self) -> Vec<i64> {
254        assert_eq!(self.dtype, DataType::Int64, "to_vec_i64 on non-i64 tensor");
255        self.as_bytes()
256            .chunks_exact(8)
257            .map(|c| i64::from_le_bytes(c.try_into().unwrap()))
258            .collect()
259    }
260}
261
262impl Clone for Tensor {
263    fn clone(&self) -> Self {
264        Self::from_raw_in(
265            self.allocator.clone(),
266            self.dtype,
267            self.shape.clone(),
268            self.as_bytes(),
269        )
270        .expect("Tensor::clone: re-allocation of identical bytes")
271    }
272}
273
274impl std::fmt::Debug for Tensor {
275    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
276        f.debug_struct("Tensor")
277            .field("dtype", &self.dtype)
278            .field("shape", &self.shape)
279            .field("device", &self.device)
280            .finish()
281    }
282}
283
284impl Drop for Tensor {
285    fn drop(&mut self) {
286        if let Some(buffer) = self.buffer.take() {
287            // `DeviceBuffer` has no `Drop`; the owning EP must free it exactly
288            // once (ep-api §4.4 invariant #2). A failed free leaks, never
289            // double-frees, so we swallow the error.
290            let _ = self.allocator.deallocate(buffer);
291        }
292    }
293}
onnx_runtime_eager/tensor.rs

onnx_runtime_eager/
tensor.rs
