onnx_runtime_ep_api/
tensor.rs

1//! Zero-copy device tensor views (§5.4) and their DLPack alignment (§5.3).
2//!
3//! These are thin, non-owning views over device-resident memory used at the
4//! kernel boundary. They intentionally use raw device pointers; the owning
5//! session/EP guarantees the backing memory outlives the view.
6//!
7//! ## DLPack correspondence (§5.3)
8//!
9//! A [`TensorView`] carries exactly the fields a DLPack `DLTensor` needs to be
10//! reconstructed, so import/export is a field-wise mapping rather than a copy:
11//!
12//! | `DLTensor`            | `TensorView`                |
13//! |----------------------|-----------------------------|
14//! | `data`               | [`TensorView::data`]        |
15//! | `byte_offset`        | [`TensorView::byte_offset`] |
16//! | `shape` / `ndim`     | [`TensorView::shape`]       |
17//! | `strides` (elements) | [`TensorView::strides`]     |
18//! | `device`             | [`TensorView::device`]      |
19//! | `dtype`              | [`TensorView::dtype`]       |
20//!
21//! DLPack keeps `byte_offset` *separate* from `data` (the base pointer is the
22//! allocation start; `byte_offset` selects the element origin) and permits
23//! negative strides. Both are representable here: strides are `i64` and the
24//! offset is applied lazily by [`TensorView::data_ptr`]. Use
25//! [`TensorView::validate`] to check the invariants of an imported view before
26//! handing it to a kernel.
27
28use onnx_runtime_ir::DataType;
29use onnx_runtime_ir::DeviceId;
30use std::marker::PhantomData;
31
32use crate::error::{EpError, Result};
33
34/// An opaque immutable device pointer (a host pointer for CPU tensors).
35#[derive(Clone, Copy, Debug)]
36pub struct DevicePtr(pub *const std::ffi::c_void);
37
38/// An opaque mutable device pointer.
39#[derive(Clone, Copy, Debug)]
40pub struct DevicePtrMut(pub *mut std::ffi::c_void);
41
42/// One validated read-only range in an ONNX external-data mmap.
43#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)]
44pub struct ExternalMmapRegion {
45    /// Identity of the live mapping that owns this range.
46    pub mapping_id: usize,
47    /// Absolute byte offset within the mapping.
48    pub offset: usize,
49    /// Length of the mapped tensor range in bytes.
50    pub len: usize,
51}
52
53/// Provenance relevant to kernels that can consume lazy external weights.
54#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
55pub enum TensorBacking {
56    #[default]
57    Opaque,
58    /// Read-only bytes directly alias an ONNX external-data mmap.
59    ExternalMmap(ExternalMmapRegion),
60}
61
62impl DevicePtr {
63    /// Reinterpret as a typed const pointer. Caller ensures the element type.
64    pub fn as_ptr<T>(self) -> *const T {
65        self.0 as *const T
66    }
67
68    /// Whether the underlying address is null.
69    pub fn is_null(self) -> bool {
70        self.0.is_null()
71    }
72}
73
74impl DevicePtrMut {
75    /// Reinterpret as a typed mutable pointer. Caller ensures the element type.
76    pub fn as_ptr<T>(self) -> *mut T {
77        self.0 as *mut T
78    }
79
80    /// Whether the underlying address is null.
81    pub fn is_null(self) -> bool {
82        self.0.is_null()
83    }
84}
85
86/// Shared invariant check for both view kinds (§5.3 DLPack import path).
87///
88/// Verifies the properties a kernel relies on: matching rank between `shape`
89/// and `strides`, a raw-view-representable dtype, and a `byte_offset` that keeps
90/// the element origin aligned to the element size. Storage bounds cannot be
91/// checked here — the view does not know its backing allocation size — so that
92/// remains the owning EP's responsibility.
93fn validate_view(
94    data_is_null: bool,
95    dtype: DataType,
96    shape: &[usize],
97    strides: &[i64],
98    byte_offset: usize,
99) -> Result<()> {
100    if data_is_null {
101        return Err(EpError::InvalidTensorView {
102            reason: "data pointer is null".into(),
103        });
104    }
105    if shape.len() != strides.len() {
106        return Err(EpError::InvalidTensorView {
107            reason: format!(
108                "rank mismatch: shape has {} dims but strides has {}",
109                shape.len(),
110                strides.len()
111            ),
112        });
113    }
114    if dtype == DataType::String {
115        return Err(EpError::InvalidTensorView {
116            reason: "String dtype has no fixed-width raw layout".into(),
117        });
118    }
119    // For fixed-width (non-sub-byte) types the element origin must land on an
120    // element boundary; sub-byte packed types are addressed per byte.
121    let esize = dtype.byte_size();
122    if esize > 1 && !byte_offset.is_multiple_of(esize) {
123        return Err(EpError::InvalidTensorView {
124            reason: format!("byte_offset {byte_offset} is not a multiple of element size {esize}"),
125        });
126    }
127    Ok(())
128}
129
130/// Immutable, non-owning view of a tensor on any device.
131#[derive(Clone, Copy)]
132pub struct TensorView<'a> {
133    pub data: DevicePtr,
134    pub dtype: DataType,
135    pub shape: &'a [usize],
136    /// Strides in **elements** (may be negative, matching DLPack).
137    pub strides: &'a [i64],
138    /// Offset in **bytes** of the element origin from `data` (DLPack semantics).
139    pub byte_offset: usize,
140    pub device: DeviceId,
141    pub backing: TensorBacking,
142    _marker: PhantomData<&'a ()>,
143}
144
145impl<'a> TensorView<'a> {
146    /// Construct a zero-offset view. `data` must remain valid for `'a`.
147    pub fn new(
148        data: DevicePtr,
149        dtype: DataType,
150        shape: &'a [usize],
151        strides: &'a [i64],
152        device: DeviceId,
153    ) -> Self {
154        Self {
155            data,
156            dtype,
157            shape,
158            strides,
159            byte_offset: 0,
160            device,
161            backing: TensorBacking::Opaque,
162            _marker: PhantomData,
163        }
164    }
165
166    /// Construct an **absent** view: the positional placeholder an executor
167    /// passes for an *omitted optional input* (an ONNX empty-string input name).
168    /// It carries a null pointer and an empty shape so it can never be read as a
169    /// real tensor — kernels test [`TensorView::is_absent`] before touching an
170    /// optional slot. This preserves positional arity so a later present input
171    /// (e.g. `Slice`'s `steps` when `axes` is omitted) is not misread as the
172    /// omitted one.
173    pub fn absent(dtype: DataType) -> Self {
174        Self {
175            data: DevicePtr(std::ptr::null()),
176            dtype,
177            shape: &[],
178            strides: &[],
179            byte_offset: 0,
180            device: DeviceId::cpu(),
181            backing: TensorBacking::Opaque,
182            _marker: PhantomData,
183        }
184    }
185
186    /// Whether this view is the [`TensorView::absent`] placeholder for an
187    /// omitted optional input (null backing pointer). Kernels with optional
188    /// inputs check this to distinguish "input not supplied" from a real,
189    /// possibly-empty, tensor.
190    pub fn is_absent(&self) -> bool {
191        self.data.is_null()
192    }
193
194    /// Set the DLPack-style byte offset of the element origin.
195    pub fn with_byte_offset(mut self, byte_offset: usize) -> Self {
196        self.byte_offset = byte_offset;
197        self
198    }
199
200    pub fn with_backing(mut self, backing: TensorBacking) -> Self {
201        self.backing = backing;
202        self
203    }
204
205    /// Check the view's invariants (rank, dtype, offset alignment). See
206    /// [`validate_view`]; call this on any view imported from DLPack before use.
207    pub fn validate(&self) -> Result<()> {
208        validate_view(
209            self.data.is_null(),
210            self.dtype,
211            self.shape,
212            self.strides,
213            self.byte_offset,
214        )
215    }
216
217    /// Whether the view is contiguous row-major.
218    pub fn is_contiguous(&self) -> bool {
219        onnx_runtime_ir::is_contiguous(self.shape, self.strides)
220    }
221
222    /// Number of elements.
223    pub fn numel(&self) -> usize {
224        self.shape.iter().product()
225    }
226
227    /// Logical element byte size (dense; ignores stride gaps). Uses
228    /// `storage_bytes` so sub-byte (int4/uint4) types are counted correctly.
229    pub fn byte_size(&self) -> usize {
230        self.dtype.storage_bytes(self.numel())
231    }
232
233    /// Typed const pointer to the element origin, applying `byte_offset`.
234    /// Computed with wrapping arithmetic (no deref) — safe to call.
235    pub fn data_ptr<T>(&self) -> *const T {
236        (self.data.0 as *const u8).wrapping_add(self.byte_offset) as *const T
237    }
238}
239
240/// Mutable, non-owning view of a tensor on any device.
241pub struct TensorMut<'a> {
242    pub data: DevicePtrMut,
243    pub dtype: DataType,
244    pub shape: &'a [usize],
245    /// Strides in **elements** (may be negative, matching DLPack).
246    pub strides: &'a [i64],
247    /// Offset in **bytes** of the element origin from `data` (DLPack semantics).
248    pub byte_offset: usize,
249    pub device: DeviceId,
250    _marker: PhantomData<&'a mut ()>,
251}
252
253impl<'a> TensorMut<'a> {
254    /// Construct a zero-offset mutable view. `data` must remain valid and
255    /// exclusively borrowed for `'a`.
256    pub fn new(
257        data: DevicePtrMut,
258        dtype: DataType,
259        shape: &'a [usize],
260        strides: &'a [i64],
261        device: DeviceId,
262    ) -> Self {
263        Self {
264            data,
265            dtype,
266            shape,
267            strides,
268            byte_offset: 0,
269            device,
270            _marker: PhantomData,
271        }
272    }
273
274    /// Set the DLPack-style byte offset of the element origin.
275    pub fn with_byte_offset(mut self, byte_offset: usize) -> Self {
276        self.byte_offset = byte_offset;
277        self
278    }
279
280    /// Check the view's invariants (rank, dtype, offset alignment).
281    pub fn validate(&self) -> Result<()> {
282        validate_view(
283            self.data.is_null(),
284            self.dtype,
285            self.shape,
286            self.strides,
287            self.byte_offset,
288        )
289    }
290
291    /// Whether the view is contiguous row-major.
292    pub fn is_contiguous(&self) -> bool {
293        onnx_runtime_ir::is_contiguous(self.shape, self.strides)
294    }
295
296    /// Number of elements.
297    pub fn numel(&self) -> usize {
298        self.shape.iter().product()
299    }
300
301    /// Logical element byte size (dense; ignores stride gaps). Uses
302    /// `storage_bytes` so sub-byte (int4/uint4) types are counted correctly.
303    pub fn byte_size(&self) -> usize {
304        self.dtype.storage_bytes(self.numel())
305    }
306
307    /// Typed mutable pointer to the element origin, applying `byte_offset`.
308    /// Computed with wrapping arithmetic (no deref) — safe to call.
309    pub fn data_ptr_mut<T>(&mut self) -> *mut T {
310        (self.data.0 as *mut u8).wrapping_add(self.byte_offset) as *mut T
311    }
312}
313
314#[cfg(test)]
315mod tests {
316    use super::*;
317    use onnx_runtime_ir::compute_contiguous_strides;
318
319    fn ptr(buf: &[u8]) -> DevicePtr {
320        DevicePtr(buf.as_ptr() as *const std::ffi::c_void)
321    }
322
323    #[test]
324    fn contiguous_view_roundtrips_invariants() {
325        let buf = vec![0u8; 6 * 4];
326        let shape = [2usize, 3];
327        let strides = compute_contiguous_strides(&shape);
328        let v = TensorView::new(ptr(&buf), DataType::Float32, &shape, &strides, DeviceId::cpu());
329        v.validate().unwrap();
330        assert!(v.is_contiguous());
331        assert_eq!(v.numel(), 6);
332        assert_eq!(v.byte_size(), 24);
333        assert_eq!(v.byte_offset, 0);
334    }
335
336    #[test]
337    fn strided_noncontiguous_view_is_representable() {
338        // A transposed [3,2] view over a [2,3] contiguous buffer: strides [1,3].
339        let buf = vec![0u8; 6 * 4];
340        let shape = [3usize, 2];
341        let strides = [1i64, 3];
342        let v = TensorView::new(ptr(&buf), DataType::Float32, &shape, &strides, DeviceId::cpu())
343            .with_byte_offset(4);
344        v.validate().unwrap();
345        assert!(!v.is_contiguous());
346        assert_eq!(v.shape, &[3, 2]);
347        assert_eq!(v.strides, &[1, 3]);
348        assert_eq!(v.byte_offset, 4);
349        // data_ptr applies the byte offset (one f32 in).
350        let base = buf.as_ptr() as usize;
351        assert_eq!(v.data_ptr::<f32>() as usize, base + 4);
352    }
353
354    #[test]
355    fn negative_strides_are_representable() {
356        // DLPack permits negative strides (reverse iteration).
357        let buf = vec![0u8; 4 * 4];
358        let shape = [4usize];
359        let strides = [-1i64];
360        let v = TensorView::new(ptr(&buf), DataType::Float32, &shape, &strides, DeviceId::cpu());
361        v.validate().unwrap();
362        assert!(!v.is_contiguous());
363    }
364
365    #[test]
366    fn validate_rejects_rank_mismatch() {
367        let buf = vec![0u8; 8];
368        let shape = [2usize, 2];
369        let strides = [1i64]; // wrong rank
370        let v = TensorView::new(ptr(&buf), DataType::Float32, &shape, &strides, DeviceId::cpu());
371        assert!(v.validate().is_err());
372    }
373
374    #[test]
375    fn validate_rejects_misaligned_offset_and_string_and_null() {
376        let buf = vec![0u8; 16];
377        let shape = [2usize];
378        let strides = [1i64];
379        // byte_offset not a multiple of f32 element size.
380        let bad_off =
381            TensorView::new(ptr(&buf), DataType::Float32, &shape, &strides, DeviceId::cpu())
382                .with_byte_offset(1);
383        assert!(bad_off.validate().is_err());
384        // String has no fixed-width raw layout.
385        let bad_dt =
386            TensorView::new(ptr(&buf), DataType::String, &shape, &strides, DeviceId::cpu());
387        assert!(bad_dt.validate().is_err());
388        // Null data pointer.
389        let bad_null = TensorView::new(
390            DevicePtr(std::ptr::null()),
391            DataType::Float32,
392            &shape,
393            &strides,
394            DeviceId::cpu(),
395        );
396        assert!(bad_null.validate().is_err());
397    }
398
399    #[test]
400    fn mut_view_offset_pointer() {
401        let mut buf = vec![0u8; 4 * 4];
402        let base = buf.as_ptr() as usize;
403        let shape = [2usize];
404        let strides = [1i64];
405        let mut v = TensorMut::new(
406            DevicePtrMut(buf.as_mut_ptr() as *mut std::ffi::c_void),
407            DataType::Float32,
408            &shape,
409            &strides,
410            DeviceId::cpu(),
411        )
412        .with_byte_offset(8);
413        v.validate().unwrap();
414        assert_eq!(v.data_ptr_mut::<f32>() as usize, base + 8);
415    }
416
417    #[test]
418    fn sub_byte_byte_size_uses_packing() {
419        let buf = vec![0u8; 4];
420        let shape = [5usize];
421        let strides = [1i64];
422        let v = TensorView::new(ptr(&buf), DataType::Int4, &shape, &strides, DeviceId::cpu());
423        // 5 int4 elements pack into 3 bytes.
424        assert_eq!(v.byte_size(), 3);
425        v.validate().unwrap();
426    }
427}
onnx_runtime_ep_api/tensor.rs

onnx_runtime_ep_api/
tensor.rs
