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

onnx_runtime_ep_api/
tensor.rs
