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    /// Set the DLPack-style byte offset of the element origin.
144    pub fn with_byte_offset(mut self, byte_offset: usize) -> Self {
145        self.byte_offset = byte_offset;
146        self
147    }
148
149    /// Check the view's invariants (rank, dtype, offset alignment). See
150    /// [`validate_view`]; call this on any view imported from DLPack before use.
151    pub fn validate(&self) -> Result<()> {
152        validate_view(
153            self.data.is_null(),
154            self.dtype,
155            self.shape,
156            self.strides,
157            self.byte_offset,
158        )
159    }
160
161    /// Whether the view is contiguous row-major.
162    pub fn is_contiguous(&self) -> bool {
163        onnx_runtime_ir::is_contiguous(self.shape, self.strides)
164    }
165
166    /// Number of elements.
167    pub fn numel(&self) -> usize {
168        self.shape.iter().product()
169    }
170
171    /// Logical element byte size (dense; ignores stride gaps). Uses
172    /// `storage_bytes` so sub-byte (int4/uint4) types are counted correctly.
173    pub fn byte_size(&self) -> usize {
174        self.dtype.storage_bytes(self.numel())
175    }
176
177    /// Typed const pointer to the element origin, applying `byte_offset`.
178    /// Computed with wrapping arithmetic (no deref) — safe to call.
179    pub fn data_ptr<T>(&self) -> *const T {
180        (self.data.0 as *const u8).wrapping_add(self.byte_offset) as *const T
181    }
182}
183
184/// Mutable, non-owning view of a tensor on any device.
185pub struct TensorMut<'a> {
186    pub data: DevicePtrMut,
187    pub dtype: DataType,
188    pub shape: &'a [usize],
189    /// Strides in **elements** (may be negative, matching DLPack).
190    pub strides: &'a [i64],
191    /// Offset in **bytes** of the element origin from `data` (DLPack semantics).
192    pub byte_offset: usize,
193    pub device: DeviceId,
194    _marker: PhantomData<&'a mut ()>,
195}
196
197impl<'a> TensorMut<'a> {
198    /// Construct a zero-offset mutable view. `data` must remain valid and
199    /// exclusively borrowed for `'a`.
200    pub fn new(
201        data: DevicePtrMut,
202        dtype: DataType,
203        shape: &'a [usize],
204        strides: &'a [i64],
205        device: DeviceId,
206    ) -> Self {
207        Self {
208            data,
209            dtype,
210            shape,
211            strides,
212            byte_offset: 0,
213            device,
214            _marker: PhantomData,
215        }
216    }
217
218    /// Set the DLPack-style byte offset of the element origin.
219    pub fn with_byte_offset(mut self, byte_offset: usize) -> Self {
220        self.byte_offset = byte_offset;
221        self
222    }
223
224    /// Check the view's invariants (rank, dtype, offset alignment).
225    pub fn validate(&self) -> Result<()> {
226        validate_view(
227            self.data.is_null(),
228            self.dtype,
229            self.shape,
230            self.strides,
231            self.byte_offset,
232        )
233    }
234
235    /// Whether the view is contiguous row-major.
236    pub fn is_contiguous(&self) -> bool {
237        onnx_runtime_ir::is_contiguous(self.shape, self.strides)
238    }
239
240    /// Number of elements.
241    pub fn numel(&self) -> usize {
242        self.shape.iter().product()
243    }
244
245    /// Logical element byte size (dense; ignores stride gaps). Uses
246    /// `storage_bytes` so sub-byte (int4/uint4) types are counted correctly.
247    pub fn byte_size(&self) -> usize {
248        self.dtype.storage_bytes(self.numel())
249    }
250
251    /// Typed mutable pointer to the element origin, applying `byte_offset`.
252    /// Computed with wrapping arithmetic (no deref) — safe to call.
253    pub fn data_ptr_mut<T>(&mut self) -> *mut T {
254        (self.data.0 as *mut u8).wrapping_add(self.byte_offset) as *mut T
255    }
256}
257
258#[cfg(test)]
259mod tests {
260    use super::*;
261    use onnx_runtime_ir::compute_contiguous_strides;
262
263    fn ptr(buf: &[u8]) -> DevicePtr {
264        DevicePtr(buf.as_ptr() as *const std::ffi::c_void)
265    }
266
267    #[test]
268    fn contiguous_view_roundtrips_invariants() {
269        let buf = vec![0u8; 6 * 4];
270        let shape = [2usize, 3];
271        let strides = compute_contiguous_strides(&shape);
272        let v = TensorView::new(ptr(&buf), DataType::Float32, &shape, &strides, DeviceId::cpu());
273        v.validate().unwrap();
274        assert!(v.is_contiguous());
275        assert_eq!(v.numel(), 6);
276        assert_eq!(v.byte_size(), 24);
277        assert_eq!(v.byte_offset, 0);
278    }
279
280    #[test]
281    fn strided_noncontiguous_view_is_representable() {
282        // A transposed [3,2] view over a [2,3] contiguous buffer: strides [1,3].
283        let buf = vec![0u8; 6 * 4];
284        let shape = [3usize, 2];
285        let strides = [1i64, 3];
286        let v = TensorView::new(ptr(&buf), DataType::Float32, &shape, &strides, DeviceId::cpu())
287            .with_byte_offset(4);
288        v.validate().unwrap();
289        assert!(!v.is_contiguous());
290        assert_eq!(v.shape, &[3, 2]);
291        assert_eq!(v.strides, &[1, 3]);
292        assert_eq!(v.byte_offset, 4);
293        // data_ptr applies the byte offset (one f32 in).
294        let base = buf.as_ptr() as usize;
295        assert_eq!(v.data_ptr::<f32>() as usize, base + 4);
296    }
297
298    #[test]
299    fn negative_strides_are_representable() {
300        // DLPack permits negative strides (reverse iteration).
301        let buf = vec![0u8; 4 * 4];
302        let shape = [4usize];
303        let strides = [-1i64];
304        let v = TensorView::new(ptr(&buf), DataType::Float32, &shape, &strides, DeviceId::cpu());
305        v.validate().unwrap();
306        assert!(!v.is_contiguous());
307    }
308
309    #[test]
310    fn validate_rejects_rank_mismatch() {
311        let buf = vec![0u8; 8];
312        let shape = [2usize, 2];
313        let strides = [1i64]; // wrong rank
314        let v = TensorView::new(ptr(&buf), DataType::Float32, &shape, &strides, DeviceId::cpu());
315        assert!(v.validate().is_err());
316    }
317
318    #[test]
319    fn validate_rejects_misaligned_offset_and_string_and_null() {
320        let buf = vec![0u8; 16];
321        let shape = [2usize];
322        let strides = [1i64];
323        // byte_offset not a multiple of f32 element size.
324        let bad_off =
325            TensorView::new(ptr(&buf), DataType::Float32, &shape, &strides, DeviceId::cpu())
326                .with_byte_offset(1);
327        assert!(bad_off.validate().is_err());
328        // String has no fixed-width raw layout.
329        let bad_dt =
330            TensorView::new(ptr(&buf), DataType::String, &shape, &strides, DeviceId::cpu());
331        assert!(bad_dt.validate().is_err());
332        // Null data pointer.
333        let bad_null = TensorView::new(
334            DevicePtr(std::ptr::null()),
335            DataType::Float32,
336            &shape,
337            &strides,
338            DeviceId::cpu(),
339        );
340        assert!(bad_null.validate().is_err());
341    }
342
343    #[test]
344    fn mut_view_offset_pointer() {
345        let mut buf = vec![0u8; 4 * 4];
346        let base = buf.as_ptr() as usize;
347        let shape = [2usize];
348        let strides = [1i64];
349        let mut v = TensorMut::new(
350            DevicePtrMut(buf.as_mut_ptr() as *mut std::ffi::c_void),
351            DataType::Float32,
352            &shape,
353            &strides,
354            DeviceId::cpu(),
355        )
356        .with_byte_offset(8);
357        v.validate().unwrap();
358        assert_eq!(v.data_ptr_mut::<f32>() as usize, base + 8);
359    }
360
361    #[test]
362    fn sub_byte_byte_size_uses_packing() {
363        let buf = vec![0u8; 4];
364        let shape = [5usize];
365        let strides = [1i64];
366        let v = TensorView::new(ptr(&buf), DataType::Int4, &shape, &strides, DeviceId::cpu());
367        // 5 int4 elements pack into 3 bytes.
368        assert_eq!(v.byte_size(), 3);
369        v.validate().unwrap();
370    }
371}
onnx_runtime_ep_api/tensor.rs

onnx_runtime_ep_api/
tensor.rs
