edgefirst_tensor/lib.rs
1// SPDX-FileCopyrightText: Copyright 2025 Au-Zone Technologies
2// SPDX-License-Identifier: Apache-2.0
3
4/*!
5EdgeFirst HAL - Tensor Module
6
7The `edgefirst_tensor` crate provides a unified interface for managing multi-dimensional arrays (tensors)
8with support for different memory types, including Direct Memory Access (DMA), POSIX Shared Memory (Shm),
9and system memory. The crate defines traits and structures for creating, reshaping, and mapping tensors into memory.
10
11## Examples
12```rust
13use edgefirst_tensor::{Error, Tensor, TensorMemory, TensorTrait};
14# fn main() -> Result<(), Error> {
15let tensor = Tensor::<f32>::new(&[2, 3, 4], Some(TensorMemory::Mem), Some("test_tensor"))?;
16assert_eq!(tensor.memory(), TensorMemory::Mem);
17assert_eq!(tensor.name(), "test_tensor");
18# Ok(())
19# }
20```
21
22## Overview
23The main structures and traits provided by the `edgefirst_tensor` crate are `TensorTrait` and `TensorMapTrait`,
24which define the behavior of Tensors and their memory mappings, respectively.
25The `Tensor<T>` struct wraps a backend-specific storage with optional image format metadata (`PixelFormat`),
26while the `TensorMap` enum provides access to the underlying data. The `TensorDyn` type-erased enum
27wraps `Tensor<T>` for runtime element-type dispatch.
28 */
29pub mod colorimetry;
30pub mod covguard;
31mod cuda;
32#[cfg(target_os = "linux")]
33mod dma;
34#[cfg(target_os = "linux")]
35mod dmabuf;
36mod error;
37mod format;
38#[cfg(target_os = "macos")]
39mod iosurface;
40mod mem;
41mod pbo;
42#[cfg(unix)]
43mod shm;
44mod tensor_dyn;
45pub use colorimetry::{
46 ColorEncoding, ColorRange, ColorSpace, ColorTransfer, Colorimetry, MatrixWeights, RangeScaling,
47};
48
49/// Retained constructor: installs the coverage flush-on-abort handler for this
50/// crate's instrumented test binary. See `covguard`. Only present under
51/// coverage on Linux (`.init_array` is ELF-only; the i.MX flush is Linux-only).
52#[cfg(all(coverage, target_os = "linux"))]
53#[used]
54#[link_section = ".init_array"]
55static __EDGEFIRST_COV_INSTALL: extern "C" fn() = {
56 extern "C" fn ctor() {
57 crate::covguard::install();
58 }
59 ctor
60};
61
62// Backing tensor/map types are internal implementation details: callers
63// allocate `Tensor<T>` / `TensorDyn` and map them, never naming the per-memory
64// backing types directly. They are `pub(crate)` so they stay nameable for the
65// `TensorStorage` / `TensorMap` enums without leaking into the public API.
66// Exceptions kept public: `Pbo*` is a GL extension point implemented by the
67// image crate, and `image_iosurface_layout` is a public helper.
68#[cfg(target_os = "linux")]
69pub(crate) use crate::dma::{DmaMap, DmaTensor};
70#[cfg(target_os = "macos")]
71pub use crate::iosurface::image_iosurface_layout;
72#[cfg(target_os = "macos")]
73pub(crate) use crate::iosurface::{IoSurfaceMap, IoSurfaceTensor};
74pub(crate) use crate::mem::{MemMap, MemTensor};
75pub use crate::pbo::{PboMap, PboMapping, PboOps, PboTensor};
76#[cfg(unix)]
77pub(crate) use crate::shm::{ShmMap, ShmTensor};
78pub use cuda::{
79 gl_map_resource, gl_register_buffer, gl_unmap_resource, gl_unregister_resource,
80 is_cuda_available, memcpy_device_to_host, stream_create, stream_destroy, stream_synchronize,
81 CudaGlOps, CudaHandle, CudaMap, CudaStream,
82};
83pub use error::{Error, Result};
84pub use format::{ChromaLayout, PixelFormat, PixelLayout};
85use num_traits::Num;
86use serde::{Deserialize, Serialize};
87#[cfg(unix)]
88use std::os::fd::OwnedFd;
89use std::{
90 fmt,
91 ops::{Deref, DerefMut},
92 sync::{
93 atomic::{AtomicU64, Ordering},
94 Arc, Weak,
95 },
96};
97pub use tensor_dyn::TensorDyn;
98
99/// Opaque keep-alive handle for a foreign-memory tensor (see
100/// [`Tensor::from_foreign`] / [`TensorDyn::from_foreign_ptr`]).
101///
102/// The HAL borrows the foreign buffer without owning it; this handle co-owns
103/// the *source* so the borrowed memory stays valid for the tensor's life. Its
104/// `Drop` releases the source — e.g. a small struct that calls `cudaFreeHost`,
105/// or a `Py<PyAny>` that decrements a NumPy array's refcount. Wrapping it in an
106/// `Arc` (then boxing each clone) makes the release fire exactly once, after
107/// the last sharing tensor/view/map drops, regardless of drop order.
108pub type ForeignOwner = Box<dyn std::any::Any + Send + Sync>;
109
110/// Re-export of `half::f16` so downstream crates can write
111/// `Tensor::<edgefirst_tensor::f16>::from_iosurface(…)` without
112/// adding `half` to their own dependency list. The version stays in
113/// lockstep with the `half` workspace dep.
114pub use half::f16;
115
116// =============================================================================
117// RGBA16F packed-layout geometry — single source of truth
118//
119// A `PlanarRgb` [3,H,W] or `PlanarRgba` [4,H,W] f16 tensor is represented
120// on the GPU as an RGBA16F surface (the only float format accepted by the
121// ANGLE IOSurface extension). Four contiguous f16 elements are packed into
122// each 8-byte RGBA16F texel, yielding a `(W/4, C*H)` surface.
123//
124// All call sites that need these dimensions must use `packed_rgba16f_layout`
125// so the rule lives in exactly one place. Currently consumed by:
126// - `crates/tensor/src/iosurface.rs` `new_image` (macOS IOSurface alloc)
127// - `crates/image/src/gl/iosurface_import.rs` (macOS GL IOSurface import)
128// - `crates/image/src/gl/processor/float.rs` (Linux GL float render — PBO
129// readback and DMA-BUF, also via the `dma_f16_packed_layout` wrapper)
130// =============================================================================
131
132/// Geometry of the RGBA16F-packed surface backing a planar F16 image tensor.
133///
134/// ANGLE only supports one float `(type, internal_format)` pair for IOSurface
135/// import: `(GL_HALF_FLOAT, GL_RGBA)` = RGBA16F (8 bytes/texel). To map a
136/// `[C, H, W]` f16 planar tensor onto such a surface, 4 contiguous f16
137/// elements are packed into each RGBA16F texel, yielding a surface of
138/// `(W/4, C*H)` texels at 8 bytes/texel. The byte stream is identical to a
139/// (nonexistent) R16F `(W, C*H)` surface and can be consumed as `&[f16]`
140/// with shape `[1, C, H, W]` without rearrangement.
141///
142/// Obtain via [`packed_rgba16f_layout`] — never construct directly.
143#[derive(Debug, Clone, Copy, PartialEq, Eq)]
144pub struct PackedRgba16fLayout {
145 /// Surface width in texels (`width / 4`).
146 pub surface_w: usize,
147 /// Surface height in texels (`planes * height`).
148 pub surface_h: usize,
149 /// Bytes per RGBA16F texel (always 8).
150 pub bytes_per_texel: usize,
151 /// Row pitch in bytes (`surface_w * 8`).
152 pub pitch: usize,
153}
154
155/// Canonical geometry for the RGBA16F-packed surface backing a planar F16
156/// image tensor.
157///
158/// Returns `Some(layout)` only when **all** of the following hold:
159///
160/// - `dtype == DType::F16`
161/// - `format` is `PixelFormat::PlanarRgb` (3 planes) or
162/// `PixelFormat::PlanarRgba` (4 planes)
163/// - `width % 4 == 0`
164///
165/// Returns `None` for any other `(format, dtype)` combination, misaligned
166/// width, or when the surface geometry would overflow `usize` — callers
167/// must fall back to a non-packed path or return a context-appropriate
168/// error.
169///
170/// # Examples
171///
172/// ```rust
173/// use edgefirst_tensor::{packed_rgba16f_layout, PixelFormat, DType};
174///
175/// let layout = packed_rgba16f_layout(PixelFormat::PlanarRgb, DType::F16, 640, 480).unwrap();
176/// assert_eq!(layout.surface_w, 160);
177/// assert_eq!(layout.surface_h, 1440);
178/// assert_eq!(layout.bytes_per_texel, 8);
179/// assert_eq!(layout.pitch, 1280);
180/// ```
181pub fn packed_rgba16f_layout(
182 format: PixelFormat,
183 dtype: DType,
184 width: usize,
185 height: usize,
186) -> Option<PackedRgba16fLayout> {
187 if dtype != DType::F16 {
188 return None;
189 }
190 let planes: usize = match format {
191 PixelFormat::PlanarRgb => 3,
192 PixelFormat::PlanarRgba => 4,
193 _ => return None,
194 };
195 if !width.is_multiple_of(4) {
196 return None;
197 }
198 let surface_w = width / 4;
199 // Checked arithmetic: a degenerate (height, width) could otherwise wrap
200 // and yield an under-sized layout, which downstream allocators trust for
201 // GPU/CPU buffer sizing. Overflow → None (handled like any other
202 // unsupported geometry).
203 let surface_h = planes.checked_mul(height)?;
204 let bytes_per_texel = 8;
205 let pitch = surface_w.checked_mul(bytes_per_texel)?;
206 Some(PackedRgba16fLayout {
207 surface_w,
208 surface_h,
209 bytes_per_texel,
210 pitch,
211 })
212}
213
214/// Per-plane DMA-BUF descriptor for external buffer import.
215///
216/// Owns a duplicated file descriptor plus optional stride and offset metadata.
217/// The fd is duplicated eagerly in [`new()`](Self::new) so that a bad fd is
218/// caught immediately. `import_image` consumes the descriptor and takes
219/// ownership of the duped fd — no further cleanup is needed by the caller.
220///
221/// # Examples
222///
223/// ```rust,no_run
224/// use edgefirst_tensor::PlaneDescriptor;
225/// use std::os::fd::BorrowedFd;
226///
227/// // SAFETY: fd 42 is hypothetical; real code must pass a valid fd.
228/// let pd = unsafe { PlaneDescriptor::new(BorrowedFd::borrow_raw(42)) }
229/// .unwrap()
230/// .with_stride(2048)
231/// .with_offset(0);
232/// ```
233#[cfg(unix)]
234pub struct PlaneDescriptor {
235 fd: OwnedFd,
236 stride: Option<usize>,
237 offset: Option<usize>,
238}
239
240#[cfg(unix)]
241impl PlaneDescriptor {
242 /// Create a new plane descriptor by duplicating the given file descriptor.
243 ///
244 /// The fd is duped immediately — a bad fd fails here rather than inside
245 /// `import_image`. The caller retains ownership of the original fd.
246 ///
247 /// # Errors
248 ///
249 /// Returns an error if the `dup()` syscall fails (e.g. invalid fd or
250 /// fd limit reached).
251 pub fn new(fd: std::os::fd::BorrowedFd<'_>) -> Result<Self> {
252 let owned = fd.try_clone_to_owned()?;
253 Ok(Self {
254 fd: owned,
255 stride: None,
256 offset: None,
257 })
258 }
259
260 /// Set the row stride in bytes (consuming builder).
261 pub fn with_stride(mut self, stride: usize) -> Self {
262 self.stride = Some(stride);
263 self
264 }
265
266 /// Set the plane offset in bytes (consuming builder).
267 pub fn with_offset(mut self, offset: usize) -> Self {
268 self.offset = Some(offset);
269 self
270 }
271
272 /// Consume the descriptor and return the owned file descriptor.
273 pub fn into_fd(self) -> OwnedFd {
274 self.fd
275 }
276
277 /// Row stride in bytes, if set.
278 pub fn stride(&self) -> Option<usize> {
279 self.stride
280 }
281
282 /// Plane offset in bytes, if set.
283 pub fn offset(&self) -> Option<usize> {
284 self.offset
285 }
286}
287
288/// A rectangular sub-region of a tensor's leading spatial frame, in pixels.
289///
290/// `Region` is the single rectangle type in the workspace: the argument to
291/// [`Tensor::view`], the source sampling window in the image crate's `Crop`,
292/// and the geometry the image backend lowers to a `glViewport` (a destination
293/// tile) or a sampling rectangle (a source). Coordinates are pixel/element
294/// units of the leading spatial axes; byte addressing is derived from the
295/// parent's row stride, not stored here.
296#[derive(Debug, Clone, Copy, PartialEq, Eq)]
297pub struct Region {
298 pub x: usize,
299 pub y: usize,
300 pub width: usize,
301 pub height: usize,
302}
303
304impl Region {
305 /// Create a region at `(x, y)` spanning `width` × `height` pixels.
306 pub fn new(x: usize, y: usize, width: usize, height: usize) -> Self {
307 Self {
308 x,
309 y,
310 width,
311 height,
312 }
313 }
314
315 /// True when the region lies fully within a `width` × `height` frame.
316 pub fn fits_within(&self, width: usize, height: usize) -> bool {
317 self.x.saturating_add(self.width) <= width && self.y.saturating_add(self.height) <= height
318 }
319}
320
321/// The parent image a [`view`](Tensor::view)/[`batch`](Tensor::batch) sub-region
322/// was carved from, snapshotted at the time the view was created.
323///
324/// A view shares the parent's `BufferIdentity` and addresses a sub-rectangle of
325/// it. The GL backend keys its EGLImage import on the **parent** geometry (so all
326/// sibling views of one buffer collapse to a single import) and renders each
327/// view as a `glViewport`+`glScissor` ROI at `(x, y, width, height)` within that
328/// parent — the view is render state, never a distinct import. `parent_width`/
329/// `parent_height` are the parent's logical pixel dimensions; `x`/`y` are this
330/// view's top-left origin within the parent (pixels). Nested views compose:
331/// the snapshot always names the **root** parent, with offsets accumulated.
332#[derive(Debug, Clone, Copy, PartialEq, Eq)]
333pub struct ViewOrigin {
334 /// Logical width of the root parent image, in pixels.
335 pub parent_width: usize,
336 /// Logical height of the root parent image, in pixels. For a `batch(n)` view
337 /// of an `[N, H, W, C]` tensor this is `N * H` (the tiles stack vertically in
338 /// the shared buffer).
339 pub parent_height: usize,
340 /// The parent's row stride in **bytes**. The GL backend keys its EGLImage
341 /// import/cache and pitch on this — NOT on the view's own `row_stride`, which
342 /// a single-row view sets tight (for map-span safety). Using the parent
343 /// stride keeps the import pitch parent-consistent so single-row and
344 /// multi-row sibling views collapse onto the same parent import.
345 pub parent_row_stride: usize,
346 /// This view's top-left x origin within the root parent, in pixels.
347 pub x: usize,
348 /// This view's top-left y origin within the root parent, in pixels.
349 pub y: usize,
350}
351
352/// Element type discriminant for runtime type identification.
353#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
354#[repr(u8)]
355#[non_exhaustive]
356pub enum DType {
357 U8,
358 I8,
359 U16,
360 I16,
361 U32,
362 I32,
363 U64,
364 I64,
365 F16,
366 F32,
367 F64,
368}
369
370impl DType {
371 /// Size of one element in bytes.
372 pub const fn size(&self) -> usize {
373 match self {
374 Self::U8 | Self::I8 => 1,
375 Self::U16 | Self::I16 | Self::F16 => 2,
376 Self::U32 | Self::I32 | Self::F32 => 4,
377 Self::U64 | Self::I64 | Self::F64 => 8,
378 }
379 }
380
381 /// Short type name (e.g., "u8", "f32", "f16").
382 pub const fn name(&self) -> &'static str {
383 match self {
384 Self::U8 => "u8",
385 Self::I8 => "i8",
386 Self::U16 => "u16",
387 Self::I16 => "i16",
388 Self::U32 => "u32",
389 Self::I32 => "i32",
390 Self::U64 => "u64",
391 Self::I64 => "i64",
392 Self::F16 => "f16",
393 Self::F32 => "f32",
394 Self::F64 => "f64",
395 }
396 }
397}
398
399impl fmt::Display for DType {
400 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
401 f.write_str(self.name())
402 }
403}
404
405/// Map a static numeric type `T` to its `DType` discriminant, returning
406/// `None` for types that do not have a `DType` representation (e.g.
407/// user-defined wrappers in tests).
408///
409/// Runtime dtype of a static `Tensor<T>` element type — a safe `TypeId`
410/// allowlist of HAL's primitive numeric types (`Some` for `u8`..`i64` /
411/// `f16` / `f32` / `f64`, `None` otherwise). Used by the macOS IOSurface
412/// image constructors (FourCC / pixel-format lookup) and by the `Mem`
413/// backing's `alloc_zeroed` fast path (these types' `T::zero()` is the
414/// all-zeros bit pattern), so it is compiled on every target.
415pub(crate) fn dtype_of<T: 'static>() -> Option<DType> {
416 use std::any::TypeId;
417 let id = TypeId::of::<T>();
418 if id == TypeId::of::<u8>() {
419 Some(DType::U8)
420 } else if id == TypeId::of::<i8>() {
421 Some(DType::I8)
422 } else if id == TypeId::of::<u16>() {
423 Some(DType::U16)
424 } else if id == TypeId::of::<i16>() {
425 Some(DType::I16)
426 } else if id == TypeId::of::<u32>() {
427 Some(DType::U32)
428 } else if id == TypeId::of::<i32>() {
429 Some(DType::I32)
430 } else if id == TypeId::of::<u64>() {
431 Some(DType::U64)
432 } else if id == TypeId::of::<i64>() {
433 Some(DType::I64)
434 } else if id == TypeId::of::<half::f16>() {
435 Some(DType::F16)
436 } else if id == TypeId::of::<f32>() {
437 Some(DType::F32)
438 } else if id == TypeId::of::<f64>() {
439 Some(DType::F64)
440 } else {
441 None
442 }
443}
444
445// =============================================================================
446// Quantization metadata — type-gated to integer element types via sealed
447// `IntegerType` trait. Accessors on `Tensor<T>` only compile when `T` is
448// an integer type; calling them on `Tensor<f32>` / `Tensor<f16>` etc. is a
449// compile error, not a runtime one.
450// =============================================================================
451
452mod sealed {
453 pub trait Sealed {}
454 impl Sealed for u8 {}
455 impl Sealed for i8 {}
456 impl Sealed for u16 {}
457 impl Sealed for i16 {}
458 impl Sealed for u32 {}
459 impl Sealed for i32 {}
460 impl Sealed for u64 {}
461 impl Sealed for i64 {}
462 // Deliberately NOT implemented for f16 / f32 / f64.
463}
464
465/// Integer element types that may carry quantization metadata.
466///
467/// Sealed trait: implemented for `u8`, `i8`, `u16`, `i16`, `u32`, `i32`,
468/// `u64`, `i64`. Cannot be implemented downstream. Float element types
469/// (`half::f16`, `f32`, `f64`) are explicitly excluded — quantization
470/// metadata does not apply to float tensors per the edgefirst.json spec.
471pub trait IntegerType: sealed::Sealed {}
472impl IntegerType for u8 {}
473impl IntegerType for i8 {}
474impl IntegerType for u16 {}
475impl IntegerType for i16 {}
476impl IntegerType for u32 {}
477impl IntegerType for i32 {}
478impl IntegerType for u64 {}
479impl IntegerType for i64 {}
480
481/// Quantization parameters for an integer tensor.
482///
483/// Covers all four modes the edgefirst.json spec defines:
484///
485/// | Mode | `scale.len()` | `zero_point` | `axis` |
486/// |---|---|---|---|
487/// | Per-tensor symmetric | 1 | `None` | `None` |
488/// | Per-tensor asymmetric | 1 | `Some(len == 1)` | `None` |
489/// | Per-channel symmetric | >1 | `None` | `Some(c)` |
490/// | Per-channel asymmetric | >1 | `Some(len == scale.len())` | `Some(c)` |
491///
492/// The quantized storage type is carried on the parent [`Tensor<T>`]; this
493/// struct does not duplicate it. Construct via the four named constructors
494/// (the only public entry points); direct field mutation is not allowed so
495/// invalid combinations cannot be represented.
496///
497/// Dequantization formula:
498///
499/// ```text
500/// real_value = scale[c] × (quantized_value[c] - zero_point[c])
501/// ```
502///
503/// where `c` is the channel index (always `0` for per-tensor).
504#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
505pub struct Quantization {
506 /// Per-tensor: `vec![scale]`. Per-channel: `vec![scale_0, scale_1, ...]`.
507 #[serde(deserialize_with = "deserialize_scalar_or_vec_f32")]
508 scale: Vec<f32>,
509
510 /// `None` means symmetric (zero-point is 0). `Some(vec)` must have the
511 /// same length as `scale`.
512 #[serde(
513 default,
514 deserialize_with = "deserialize_opt_scalar_or_vec_i32",
515 skip_serializing_if = "Option::is_none"
516 )]
517 zero_point: Option<Vec<i32>>,
518
519 /// Channel axis for per-channel quantization. `Some(_)` iff
520 /// `scale.len() > 1`. Validated against the parent tensor's shape at
521 /// `set_quantization()` time.
522 #[serde(default, skip_serializing_if = "Option::is_none")]
523 axis: Option<usize>,
524}
525
526/// Semantic mode discriminant for hot-path kernel dispatch.
527///
528/// Obtain via [`Quantization::mode`] once at kernel entry; never inside a
529/// pixel-level loop. The enum is borrow-based so the hot kernel receives
530/// the scales / zero-points as slices without reallocation.
531#[derive(Debug, Clone, Copy)]
532pub enum QuantMode<'a> {
533 PerTensorSymmetric {
534 scale: f32,
535 },
536 PerTensor {
537 scale: f32,
538 zero_point: i32,
539 },
540 PerChannelSymmetric {
541 scales: &'a [f32],
542 axis: usize,
543 },
544 PerChannel {
545 scales: &'a [f32],
546 zero_points: &'a [i32],
547 axis: usize,
548 },
549}
550
551impl Quantization {
552 /// Per-tensor symmetric (zero_point = 0).
553 pub fn per_tensor_symmetric(scale: f32) -> Self {
554 Self {
555 scale: vec![scale],
556 zero_point: None,
557 axis: None,
558 }
559 }
560
561 /// Per-tensor asymmetric — the most common runtime shape.
562 pub fn per_tensor(scale: f32, zero_point: i32) -> Self {
563 Self {
564 scale: vec![scale],
565 zero_point: Some(vec![zero_point]),
566 axis: None,
567 }
568 }
569
570 /// Per-channel symmetric. Errors on empty `scales`.
571 pub fn per_channel_symmetric(scales: Vec<f32>, axis: usize) -> Result<Self> {
572 if scales.is_empty() {
573 return Err(Error::QuantizationInvalid {
574 field: "scale.len",
575 expected: "non-empty per-channel scales".to_string(),
576 got: "length 0".to_string(),
577 });
578 }
579 Ok(Self {
580 scale: scales,
581 zero_point: None,
582 axis: Some(axis),
583 })
584 }
585
586 /// Per-channel asymmetric. Errors on length mismatch between `scales`
587 /// and `zero_points`, or empty arrays.
588 pub fn per_channel(scales: Vec<f32>, zero_points: Vec<i32>, axis: usize) -> Result<Self> {
589 if scales.is_empty() {
590 return Err(Error::QuantizationInvalid {
591 field: "scale.len",
592 expected: "non-empty per-channel scales".to_string(),
593 got: "length 0".to_string(),
594 });
595 }
596 if scales.len() != zero_points.len() {
597 return Err(Error::QuantizationInvalid {
598 field: "zero_point.len",
599 expected: format!("length matches scale ({})", scales.len()),
600 got: format!("length {}", zero_points.len()),
601 });
602 }
603 Ok(Self {
604 scale: scales,
605 zero_point: Some(zero_points),
606 axis: Some(axis),
607 })
608 }
609
610 /// Borrow-based dispatch view. Match once at kernel entry.
611 pub fn mode(&self) -> QuantMode<'_> {
612 match (self.scale.len(), self.zero_point.as_deref(), self.axis) {
613 (1, None, _) => QuantMode::PerTensorSymmetric {
614 scale: self.scale[0],
615 },
616 (1, Some(zps), _) => QuantMode::PerTensor {
617 scale: self.scale[0],
618 zero_point: zps.first().copied().unwrap_or(0),
619 },
620 (_, None, Some(axis)) => QuantMode::PerChannelSymmetric {
621 scales: &self.scale,
622 axis,
623 },
624 (_, Some(zps), Some(axis)) => QuantMode::PerChannel {
625 scales: &self.scale,
626 zero_points: zps,
627 axis,
628 },
629 // The `validate()` path prevents constructing a
630 // per-channel Quantization without an axis, so the remaining
631 // pattern is unreachable in practice. Fall back to
632 // per-tensor symmetric using scale[0] to avoid panicking in
633 // release; debug builds assert.
634 _ => {
635 debug_assert!(
636 false,
637 "Quantization::mode: per-channel without axis is unreachable"
638 );
639 QuantMode::PerTensorSymmetric {
640 scale: self.scale.first().copied().unwrap_or(1.0),
641 }
642 }
643 }
644 }
645
646 /// Returns `true` for per-tensor quantization (`scale.len() == 1`).
647 pub fn is_per_tensor(&self) -> bool {
648 self.scale.len() == 1
649 }
650
651 /// Returns `true` for per-channel quantization (`scale.len() > 1`).
652 pub fn is_per_channel(&self) -> bool {
653 self.scale.len() > 1
654 }
655
656 /// Returns `true` for symmetric quantization (no zero-point, or
657 /// zero-point vector of all zeros).
658 pub fn is_symmetric(&self) -> bool {
659 match &self.zero_point {
660 None => true,
661 Some(zps) => zps.iter().all(|&z| z == 0),
662 }
663 }
664
665 /// Borrow the scale array. Length 1 for per-tensor; `num_channels` for
666 /// per-channel.
667 pub fn scale(&self) -> &[f32] {
668 &self.scale
669 }
670
671 /// Borrow the zero-point array. `None` for symmetric.
672 pub fn zero_point(&self) -> Option<&[i32]> {
673 self.zero_point.as_deref()
674 }
675
676 /// Channel axis for per-channel quantization. `None` for per-tensor.
677 pub fn axis(&self) -> Option<usize> {
678 self.axis
679 }
680
681 /// Validate against a target tensor shape. Runs in
682 /// `Tensor::set_quantization()`. Catches:
683 /// - empty `scale` (reject — must declare at least one factor)
684 /// - `zero_point` length inconsistent with `scale` (reject —
685 /// per-tensor must have len 1, per-channel must match `scale.len`)
686 /// - `axis >= shape.len()` (axis out of range)
687 /// - `scale.len() != shape[axis]` for per-channel
688 /// - per-channel without axis (reject)
689 /// - per-tensor with redundant axis (reject)
690 pub(crate) fn validate(&self, shape: &[usize]) -> Result<()> {
691 // `Quantization` is `Deserialize`, so malformed JSON like
692 // `{"scale": [], "zero_point": []}` could otherwise produce an
693 // ill-defined value that confuses `mode()` selection and the
694 // per-channel kernels' indexing.
695 if self.scale.is_empty() {
696 return Err(Error::QuantizationInvalid {
697 field: "scale.len",
698 expected: ">= 1".to_string(),
699 got: "0".to_string(),
700 });
701 }
702 if let Some(zps) = self.zero_point.as_ref() {
703 // Per-tensor: scale.len() == 1 and zero_point.len() must == 1.
704 // Per-channel: zero_point.len() must == scale.len().
705 let expected = if self.scale.len() == 1 {
706 1
707 } else {
708 self.scale.len()
709 };
710 if zps.len() != expected {
711 return Err(Error::QuantizationInvalid {
712 field: "zero_point.len",
713 expected: format!(
714 "{expected} (matching {})",
715 if self.scale.len() == 1 {
716 "per-tensor scale"
717 } else {
718 "per-channel scale.len"
719 }
720 ),
721 got: format!("length {}", zps.len()),
722 });
723 }
724 }
725
726 match (self.scale.len(), self.axis) {
727 (1, None) => Ok(()),
728 (1, Some(_)) => Err(Error::QuantizationInvalid {
729 field: "per_tensor_redundant_axis",
730 expected: "axis=None for per-tensor quantization".to_string(),
731 got: format!("axis={:?}", self.axis),
732 }),
733 (_, None) => Err(Error::QuantizationInvalid {
734 field: "per_channel_requires_axis",
735 expected: format!(
736 "axis=Some(_) for per-channel quantization (scale.len={})",
737 self.scale.len()
738 ),
739 got: "axis=None".to_string(),
740 }),
741 (n, Some(axis)) => {
742 if axis >= shape.len() {
743 return Err(Error::QuantizationInvalid {
744 field: "axis",
745 expected: format!("axis < tensor rank ({})", shape.len()),
746 got: format!("axis={axis}"),
747 });
748 }
749 if shape[axis] != n {
750 return Err(Error::QuantizationInvalid {
751 field: "scale.len",
752 expected: format!("length matches shape[{axis}] ({})", shape[axis]),
753 got: format!("length {n}"),
754 });
755 }
756 Ok(())
757 }
758 }
759 }
760}
761
762impl From<(f32, i32)> for Quantization {
763 /// Convenience construction from a `(scale, zero_point)` tuple. Matches
764 /// the legacy `QuantTuple` / `Quantization::new` calling convention so
765 /// existing `(0.1, -128).into()` sites keep working.
766 fn from((scale, zero_point): (f32, i32)) -> Self {
767 Self::per_tensor(scale, zero_point)
768 }
769}
770
771fn deserialize_scalar_or_vec_f32<'de, D: serde::Deserializer<'de>>(
772 de: D,
773) -> std::result::Result<Vec<f32>, D::Error> {
774 use serde::de::{self, Visitor};
775 struct V;
776 impl<'de> Visitor<'de> for V {
777 type Value = Vec<f32>;
778 fn expecting(&self, f: &mut fmt::Formatter) -> fmt::Result {
779 f.write_str("f32 or array of f32")
780 }
781 fn visit_f64<E: de::Error>(self, v: f64) -> std::result::Result<Self::Value, E> {
782 Ok(vec![v as f32])
783 }
784 #[allow(clippy::cast_possible_truncation)]
785 fn visit_i64<E: de::Error>(self, v: i64) -> std::result::Result<Self::Value, E> {
786 Ok(vec![v as f32])
787 }
788 #[allow(clippy::cast_possible_truncation, clippy::cast_precision_loss)]
789 fn visit_u64<E: de::Error>(self, v: u64) -> std::result::Result<Self::Value, E> {
790 Ok(vec![v as f32])
791 }
792 fn visit_seq<A: de::SeqAccess<'de>>(
793 self,
794 mut seq: A,
795 ) -> std::result::Result<Self::Value, A::Error> {
796 let mut out = Vec::with_capacity(seq.size_hint().unwrap_or(1));
797 while let Some(x) = seq.next_element::<f32>()? {
798 out.push(x);
799 }
800 Ok(out)
801 }
802 }
803 de.deserialize_any(V)
804}
805
806fn deserialize_opt_scalar_or_vec_i32<'de, D: serde::Deserializer<'de>>(
807 de: D,
808) -> std::result::Result<Option<Vec<i32>>, D::Error> {
809 use serde::de::{self, Visitor};
810 struct V;
811 impl<'de> Visitor<'de> for V {
812 type Value = Option<Vec<i32>>;
813 fn expecting(&self, f: &mut fmt::Formatter) -> fmt::Result {
814 f.write_str("null, i32, or array of i32")
815 }
816 fn visit_none<E: de::Error>(self) -> std::result::Result<Self::Value, E> {
817 Ok(None)
818 }
819 fn visit_unit<E: de::Error>(self) -> std::result::Result<Self::Value, E> {
820 Ok(None)
821 }
822 fn visit_some<D2: serde::Deserializer<'de>>(
823 self,
824 de: D2,
825 ) -> std::result::Result<Self::Value, D2::Error> {
826 struct Inner;
827 impl<'de> Visitor<'de> for Inner {
828 type Value = Vec<i32>;
829 fn expecting(&self, f: &mut fmt::Formatter) -> fmt::Result {
830 f.write_str("i32 or array of i32")
831 }
832 #[allow(clippy::cast_possible_truncation)]
833 fn visit_i64<E: de::Error>(self, v: i64) -> std::result::Result<Self::Value, E> {
834 Ok(vec![v as i32])
835 }
836 #[allow(clippy::cast_possible_truncation, clippy::cast_possible_wrap)]
837 fn visit_u64<E: de::Error>(self, v: u64) -> std::result::Result<Self::Value, E> {
838 Ok(vec![v as i32])
839 }
840 fn visit_seq<A: de::SeqAccess<'de>>(
841 self,
842 mut seq: A,
843 ) -> std::result::Result<Self::Value, A::Error> {
844 let mut out = Vec::with_capacity(seq.size_hint().unwrap_or(1));
845 while let Some(x) = seq.next_element::<i32>()? {
846 out.push(x);
847 }
848 Ok(out)
849 }
850 }
851 de.deserialize_any(Inner).map(Some)
852 }
853 #[allow(clippy::cast_possible_truncation)]
854 fn visit_i64<E: de::Error>(self, v: i64) -> std::result::Result<Self::Value, E> {
855 Ok(Some(vec![v as i32]))
856 }
857 #[allow(clippy::cast_possible_truncation, clippy::cast_possible_wrap)]
858 fn visit_u64<E: de::Error>(self, v: u64) -> std::result::Result<Self::Value, E> {
859 Ok(Some(vec![v as i32]))
860 }
861 fn visit_seq<A: de::SeqAccess<'de>>(
862 self,
863 mut seq: A,
864 ) -> std::result::Result<Self::Value, A::Error> {
865 let mut out = Vec::with_capacity(seq.size_hint().unwrap_or(1));
866 while let Some(x) = seq.next_element::<i32>()? {
867 out.push(x);
868 }
869 Ok(Some(out))
870 }
871 }
872 de.deserialize_option(V)
873}
874
875/// Monotonic counter for buffer identity IDs.
876static NEXT_BUFFER_ID: AtomicU64 = AtomicU64::new(1);
877
878/// Unique identity for a tensor's underlying buffer.
879///
880/// Created fresh on every buffer allocation or import. The `id` is a monotonic
881/// u64 used as a cache key. The `guard` is an `Arc<()>` whose weak references
882/// allow downstream caches to detect when the buffer has been dropped.
883#[derive(Debug, Clone)]
884pub struct BufferIdentity {
885 id: u64,
886 guard: Arc<()>,
887}
888
889impl BufferIdentity {
890 /// Create a new unique buffer identity.
891 pub fn new() -> Self {
892 Self {
893 id: NEXT_BUFFER_ID.fetch_add(1, Ordering::Relaxed),
894 guard: Arc::new(()),
895 }
896 }
897
898 /// Unique identifier for this buffer. Changes when the buffer changes.
899 pub fn id(&self) -> u64 {
900 self.id
901 }
902
903 /// Returns a weak reference to the buffer guard. Goes dead when the
904 /// owning Tensor is dropped (and no clones remain).
905 pub fn weak(&self) -> Weak<()> {
906 Arc::downgrade(&self.guard)
907 }
908}
909
910impl Default for BufferIdentity {
911 fn default() -> Self {
912 Self::new()
913 }
914}
915
916#[cfg(target_os = "linux")]
917use nix::sys::stat::{major, minor};
918
919pub trait TensorTrait<T>: Send + Sync
920where
921 T: Num + Clone + fmt::Debug,
922{
923 /// Create a new tensor with the given shape and optional name. If no name
924 /// is given, a random name will be generated.
925 fn new(shape: &[usize], name: Option<&str>) -> Result<Self>
926 where
927 Self: Sized;
928
929 #[cfg(unix)]
930 /// Create a new tensor using the given file descriptor, shape, and optional
931 /// name. If no name is given, a random name will be generated.
932 ///
933 /// On Linux: Inspects the fd to determine DMA vs SHM based on device major/minor.
934 /// On other Unix (macOS): Always creates SHM tensor.
935 fn from_fd(fd: std::os::fd::OwnedFd, shape: &[usize], name: Option<&str>) -> Result<Self>
936 where
937 Self: Sized;
938
939 #[cfg(unix)]
940 /// Clone the file descriptor associated with this tensor.
941 fn clone_fd(&self) -> Result<std::os::fd::OwnedFd>;
942
943 /// Get the memory type of this tensor.
944 fn memory(&self) -> TensorMemory;
945
946 /// Get the name of this tensor.
947 fn name(&self) -> String;
948
949 /// Get the number of elements in this tensor.
950 fn len(&self) -> usize {
951 self.shape().iter().product()
952 }
953
954 /// Check if the tensor is empty.
955 fn is_empty(&self) -> bool {
956 self.len() == 0
957 }
958
959 /// Get the size in bytes of this tensor.
960 fn size(&self) -> usize {
961 self.len() * std::mem::size_of::<T>()
962 }
963
964 /// Get the shape of this tensor.
965 fn shape(&self) -> &[usize];
966
967 /// Reshape this tensor to the given shape. The total number of elements
968 /// must remain the same.
969 fn reshape(&mut self, shape: &[usize]) -> Result<()>;
970
971 /// Bytes of the underlying allocation (>= the current logical `size()`).
972 /// Defaults to the logical size for storages without spare capacity.
973 fn capacity_bytes(&self) -> usize {
974 self.size()
975 }
976
977 /// Set the logical shape to any shape whose byte size fits the allocation
978 /// capacity, without the equal-size constraint of `reshape`.
979 fn set_logical_shape(&mut self, shape: &[usize]) -> Result<()> {
980 self.reshape(shape)
981 }
982
983 /// Map the tensor into memory and return a TensorMap for accessing the
984 /// data.
985 fn map(&self) -> Result<TensorMap<T>>;
986
987 /// Get the buffer identity for cache keying and liveness tracking.
988 fn buffer_identity(&self) -> &BufferIdentity;
989
990 /// Create a zero-copy sub-region view of this backing that shares the
991 /// underlying allocation **and** [`BufferIdentity`].
992 ///
993 /// The window is `[offset_bytes, offset_bytes + shape.product() *
994 /// size_of::<T>())` measured from this tensor's own logical start, so a
995 /// sub-view of a sub-view composes by adding offsets. Sharing the parent's
996 /// identity is the contract that lets identity-keyed caches (e.g. the GL
997 /// EGLImage import cache) treat offset-distinct windows as one buffer rather
998 /// than unrelated allocations — `view` must never mint a fresh identity.
999 ///
1000 /// Defaults to [`Error::NotImplemented`]; every backend that supports
1001 /// sub-views overrides it (`Mem`, `Shm`, Linux DMA, macOS IOSurface, `Pbo`).
1002 fn view(&self, offset_bytes: usize, shape: &[usize]) -> Result<Self>
1003 where
1004 Self: Sized,
1005 {
1006 let _ = (offset_bytes, shape);
1007 Err(Error::NotImplemented(
1008 "view (zero-copy sub-region) is not supported for this tensor backend".to_owned(),
1009 ))
1010 }
1011}
1012
1013pub trait TensorMapTrait<T>
1014where
1015 T: Num + Clone + fmt::Debug,
1016{
1017 /// Get the shape of this tensor map.
1018 fn shape(&self) -> &[usize];
1019
1020 /// Unmap the tensor from memory.
1021 fn unmap(&mut self);
1022
1023 /// Get the number of elements in this tensor map.
1024 fn len(&self) -> usize {
1025 self.shape().iter().product()
1026 }
1027
1028 /// Check if the tensor map is empty.
1029 fn is_empty(&self) -> bool {
1030 self.len() == 0
1031 }
1032
1033 /// Get the size in bytes of this tensor map.
1034 fn size(&self) -> usize {
1035 self.len() * std::mem::size_of::<T>()
1036 }
1037
1038 /// Get a slice to the data in this tensor map.
1039 fn as_slice(&self) -> &[T];
1040
1041 /// Get a mutable slice to the data in this tensor map.
1042 fn as_mut_slice(&mut self) -> &mut [T];
1043
1044 #[cfg(feature = "ndarray")]
1045 /// Get an ndarray ArrayView of the tensor data.
1046 fn view(&'_ self) -> Result<ndarray::ArrayView<'_, T, ndarray::Dim<ndarray::IxDynImpl>>> {
1047 Ok(ndarray::ArrayView::from_shape(
1048 self.shape(),
1049 self.as_slice(),
1050 )?)
1051 }
1052
1053 #[cfg(feature = "ndarray")]
1054 /// Get an ndarray ArrayViewMut of the tensor data.
1055 fn view_mut(
1056 &'_ mut self,
1057 ) -> Result<ndarray::ArrayViewMut<'_, T, ndarray::Dim<ndarray::IxDynImpl>>> {
1058 let shape = self.shape().to_vec();
1059 Ok(ndarray::ArrayViewMut::from_shape(
1060 shape,
1061 self.as_mut_slice(),
1062 )?)
1063 }
1064}
1065
1066#[derive(Debug, Clone, Copy, PartialEq, Eq)]
1067pub enum TensorMemory {
1068 /// Platform-native zero-copy GPU buffer.
1069 ///
1070 /// On Linux this is a DMA-BUF (`DmaTensor` in `crates/tensor/src/dma.rs`)
1071 /// allocated via the DRM/dma-heap subsystem. On macOS this is an
1072 /// IOSurface (`IoSurfaceTensor` in `crates/tensor/src/iosurface.rs`).
1073 /// Both fit into the same `TensorStorage::Dma` slot at the trait
1074 /// level — the public C API discriminant (`HalTensorMemory::Dma=1`)
1075 /// works on both platforms with no ABI break.
1076 ///
1077 /// Allows hardware-accelerated paths (OpenGL backend on Linux via
1078 /// `EGL_EXT_image_dma_buf_import`; macOS via
1079 /// `EGL_ANGLE_iosurface_client_buffer`). CPU access via `map()`
1080 /// incurs cache-coherency overhead on Linux DMA-BUF and is similar
1081 /// in cost on IOSurface; SHM/Mem are cheaper for CPU-only workloads.
1082 Dma,
1083 #[cfg(unix)]
1084 /// POSIX Shared Memory allocation. Suitable for inter-process
1085 /// communication, but not suitable for hardware acceleration.
1086 Shm,
1087
1088 /// Regular system memory allocation
1089 Mem,
1090
1091 /// OpenGL Pixel Buffer Object memory. Created by ImageProcessor
1092 /// when DMA-buf is unavailable but OpenGL is present.
1093 Pbo,
1094}
1095
1096impl From<TensorMemory> for String {
1097 fn from(memory: TensorMemory) -> Self {
1098 match memory {
1099 TensorMemory::Dma => "dma".to_owned(),
1100 #[cfg(unix)]
1101 TensorMemory::Shm => "shm".to_owned(),
1102 TensorMemory::Mem => "mem".to_owned(),
1103 TensorMemory::Pbo => "pbo".to_owned(),
1104 }
1105 }
1106}
1107
1108impl TryFrom<&str> for TensorMemory {
1109 type Error = Error;
1110
1111 fn try_from(s: &str) -> Result<Self> {
1112 match s {
1113 "dma" => Ok(TensorMemory::Dma),
1114 #[cfg(unix)]
1115 "shm" => Ok(TensorMemory::Shm),
1116 "mem" => Ok(TensorMemory::Mem),
1117 "pbo" => Ok(TensorMemory::Pbo),
1118 _ => Err(Error::InvalidMemoryType(s.to_owned())),
1119 }
1120 }
1121}
1122
1123#[derive(Debug)]
1124#[allow(dead_code)] // Variants are constructed by downstream crates via pub(crate) helpers
1125pub(crate) enum TensorStorage<T>
1126where
1127 T: Num + Clone + fmt::Debug + Send + Sync,
1128{
1129 /// Platform-native zero-copy GPU buffer. Inner type differs per
1130 /// target: `DmaTensor` on Linux (DMA-BUF fd), `IoSurfaceTensor` on
1131 /// macOS (CFRetained IOSurface). The shared variant name keeps the
1132 /// public `TensorMemory::Dma` discriminant stable across platforms.
1133 #[cfg(target_os = "linux")]
1134 Dma(DmaTensor<T>),
1135 #[cfg(target_os = "macos")]
1136 Dma(IoSurfaceTensor<T>),
1137 #[cfg(unix)]
1138 Shm(ShmTensor<T>),
1139 Mem(MemTensor<T>),
1140 Pbo(PboTensor<T>),
1141}
1142
1143impl<T> TensorStorage<T>
1144where
1145 T: Num + Clone + fmt::Debug + Send + Sync,
1146{
1147 /// The backing allocation's intrinsic physical row pitch in bytes, if it
1148 /// has one that is fixed independent of the logical shape. macOS IOSurface
1149 /// reports its 64-aligned `bytesPerRow`; other backings (Linux DMA, SHM,
1150 /// Mem, PBO) have no fixed pitch beyond the logical shape and return `None`.
1151 ///
1152 /// Used by `configure_image` to preserve the physical pitch when a reused
1153 /// pool tensor is reconfigured to a smaller logical image — so the decode
1154 /// writes rows at the surface's real stride and the GPU samples them with
1155 /// the same stride (the physical-grid / logical-ROI decoupling).
1156 pub(crate) fn backing_row_stride(&self) -> Option<usize> {
1157 match self {
1158 // Only genuine image-formatted IOSurfaces (height > 1) carry a real
1159 // per-row pitch; a generic byte-bag (height == 1) returns `None` so
1160 // `configure_image` does not adopt its whole-buffer "row" as a stride.
1161 #[cfg(target_os = "macos")]
1162 TensorStorage::Dma(t) => t.image_backing_row_stride(),
1163 _ => None,
1164 }
1165 }
1166
1167 /// Create a new tensor storage with the given shape, memory type, and
1168 /// optional name. If no name is given, a random name will be generated.
1169 /// If no memory type is given, the best available memory type will be
1170 /// chosen based on the platform and environment variables.
1171 fn new(shape: &[usize], memory: Option<TensorMemory>, name: Option<&str>) -> Result<Self> {
1172 match memory {
1173 #[cfg(target_os = "linux")]
1174 Some(TensorMemory::Dma) => {
1175 DmaTensor::<T>::new(shape, name).map(TensorStorage::Dma)
1176 }
1177 #[cfg(target_os = "macos")]
1178 Some(TensorMemory::Dma) => {
1179 IoSurfaceTensor::<T>::new(shape, name).map(TensorStorage::Dma)
1180 }
1181 #[cfg(not(any(target_os = "linux", target_os = "macos")))]
1182 Some(TensorMemory::Dma) => Err(crate::error::Error::NotImplemented(
1183 "TensorMemory::Dma is only available on Linux (DMA-BUF) and macOS (IOSurface)"
1184 .to_owned(),
1185 )),
1186 #[cfg(unix)]
1187 Some(TensorMemory::Shm) => {
1188 ShmTensor::<T>::new(shape, name).map(TensorStorage::Shm)
1189 }
1190 Some(TensorMemory::Mem) => {
1191 MemTensor::<T>::new(shape, name).map(TensorStorage::Mem)
1192 }
1193 Some(TensorMemory::Pbo) => Err(crate::error::Error::NotImplemented(
1194 "PboTensor cannot be created via Tensor::new() — use ImageProcessor::create_image()".to_owned(),
1195 )),
1196 None => {
1197 if std::env::var("EDGEFIRST_TENSOR_FORCE_MEM")
1198 .is_ok_and(|x| x != "0" && x.to_lowercase() != "false")
1199 {
1200 MemTensor::<T>::new(shape, name).map(TensorStorage::Mem)
1201 } else {
1202 // Auto-select priority: Dma > Mem. Shm is intentionally NOT
1203 // auto-selected — it offers no advantage over Mem for an
1204 // in-process tensor and Mem always succeeds, so Shm below Mem
1205 // is effectively never reached. Request Shm explicitly via
1206 // `TensorMemory::Shm` when cross-process sharing is needed.
1207 // (PBO sits between Dma and Mem but is GL-backed and created
1208 // only via `ImageProcessor::create_image`.)
1209 #[cfg(target_os = "linux")]
1210 {
1211 // Linux: Try DMA -> Mem
1212 match DmaTensor::<T>::new(shape, name) {
1213 Ok(tensor) => Ok(TensorStorage::Dma(tensor)),
1214 Err(_) => MemTensor::<T>::new(shape, name).map(TensorStorage::Mem),
1215 }
1216 }
1217 #[cfg(target_os = "macos")]
1218 {
1219 // macOS: Try IOSurface -> Mem. IOSurface is the
1220 // GPU-shareable backend (zero-copy via ANGLE), filling the
1221 // same role as DMA-BUF on Linux.
1222 match IoSurfaceTensor::<T>::new(shape, name) {
1223 Ok(tensor) => Ok(TensorStorage::Dma(tensor)),
1224 Err(_) => MemTensor::<T>::new(shape, name).map(TensorStorage::Mem),
1225 }
1226 }
1227 #[cfg(all(unix, not(any(target_os = "linux", target_os = "macos"))))]
1228 {
1229 // Other Unix (BSD): Mem only (no DMA; Shm is explicit-only)
1230 MemTensor::<T>::new(shape, name).map(TensorStorage::Mem)
1231 }
1232 #[cfg(not(unix))]
1233 {
1234 // Windows/other: Mem only
1235 MemTensor::<T>::new(shape, name).map(TensorStorage::Mem)
1236 }
1237 }
1238 }
1239 }
1240 }
1241
1242 /// Create a DMA-backed tensor storage with an explicit byte size that
1243 /// may exceed `shape.product() * sizeof(T)`. Used for image tensors
1244 /// with row-padded layouts (see `DmaTensor::new_with_byte_size`).
1245 ///
1246 /// This is intentionally DMA-only: padding is only meaningful for
1247 /// buffers that will be imported as GPU textures via EGLImage. PBO,
1248 /// Shm, and Mem storage doesn't benefit from pitch alignment and
1249 /// shouldn't pay the memory cost.
1250 #[cfg(target_os = "linux")]
1251 pub(crate) fn new_dma_with_byte_size(
1252 shape: &[usize],
1253 byte_size: usize,
1254 name: Option<&str>,
1255 ) -> Result<Self> {
1256 DmaTensor::<T>::new_with_byte_size(shape, byte_size, name).map(TensorStorage::Dma)
1257 }
1258
1259 // No non-Linux stub: the only caller (`Tensor::image_with_stride`)
1260 // returns `NotImplemented` directly on non-Linux without ever
1261 // reaching the storage layer, so defining a stub here would be
1262 // dead code and fail the `-D warnings` clippy gate on macOS CI.
1263
1264 /// Create a Mem-backed tensor storage with an explicit byte size that may
1265 /// exceed `shape.product() * sizeof(T)`. Used for image tensors with
1266 /// 64-byte-aligned row strides (see `MemTensor::with_capacity_bytes`).
1267 pub(crate) fn new_mem_with_byte_size(
1268 shape: &[usize],
1269 byte_size: usize,
1270 name: Option<&str>,
1271 ) -> Result<Self>
1272 where
1273 T: 'static,
1274 {
1275 MemTensor::<T>::with_capacity_bytes(shape, byte_size, name).map(TensorStorage::Mem)
1276 }
1277
1278 /// Create a Shm-backed tensor storage with an explicit byte size that may
1279 /// exceed `shape.product() * sizeof(T)`. Used for image tensors with
1280 /// 64-byte-aligned row strides (see `ShmTensor::new_with_byte_size`).
1281 #[cfg(unix)]
1282 pub(crate) fn new_shm_with_byte_size(
1283 shape: &[usize],
1284 byte_size: usize,
1285 name: Option<&str>,
1286 ) -> Result<Self> {
1287 ShmTensor::<T>::new_with_byte_size(shape, byte_size, name).map(TensorStorage::Shm)
1288 }
1289
1290 /// Allocate an image-formatted IOSurface-backed storage (macOS).
1291 ///
1292 /// Used by `Tensor::image()` when the caller requests
1293 /// `TensorMemory::Dma` and the format has an IOSurface FourCC
1294 /// mapping (YUYV, RGBA, BGRA today). Falls back to `new_with_byte_size`
1295 /// otherwise.
1296 #[cfg(target_os = "macos")]
1297 pub(crate) fn new_image_iosurface(
1298 width: usize,
1299 height: usize,
1300 format: PixelFormat,
1301 dtype: DType,
1302 shape: &[usize],
1303 name: Option<&str>,
1304 ) -> Result<Self> {
1305 IoSurfaceTensor::<T>::new_image(width, height, format, dtype, shape, name)
1306 .map(TensorStorage::Dma)
1307 }
1308
1309 /// Create a new tensor storage using the given file descriptor, shape,
1310 /// and optional name.
1311 #[cfg(unix)]
1312 fn from_fd(fd: OwnedFd, shape: &[usize], name: Option<&str>) -> Result<Self> {
1313 #[cfg(target_os = "linux")]
1314 {
1315 use nix::sys::stat::fstat;
1316
1317 let stat = fstat(&fd)?;
1318 let major = major(stat.st_dev);
1319 let minor = minor(stat.st_dev);
1320
1321 log::debug!("Creating tensor from fd: major={major}, minor={minor}");
1322
1323 if major != 0 {
1324 // Dma and Shm tensors are expected to have major number 0
1325 return Err(Error::UnknownDeviceType(major, minor));
1326 }
1327
1328 match minor {
1329 9 | 10 => {
1330 // minor number 9 & 10 indicates DMA memory
1331 DmaTensor::<T>::from_fd(fd, shape, name).map(TensorStorage::Dma)
1332 }
1333 _ => {
1334 // other minor numbers are assumed to be shared memory
1335 ShmTensor::<T>::from_fd(fd, shape, name).map(TensorStorage::Shm)
1336 }
1337 }
1338 }
1339 #[cfg(all(unix, not(target_os = "linux")))]
1340 {
1341 // On macOS/BSD, always use SHM (no DMA support)
1342 ShmTensor::<T>::from_fd(fd, shape, name).map(TensorStorage::Shm)
1343 }
1344 }
1345}
1346
1347impl<T> TensorTrait<T> for TensorStorage<T>
1348where
1349 T: Num + Clone + fmt::Debug + Send + Sync,
1350{
1351 fn new(shape: &[usize], name: Option<&str>) -> Result<Self> {
1352 Self::new(shape, None, name)
1353 }
1354
1355 #[cfg(unix)]
1356 fn from_fd(fd: OwnedFd, shape: &[usize], name: Option<&str>) -> Result<Self> {
1357 Self::from_fd(fd, shape, name)
1358 }
1359
1360 #[cfg(unix)]
1361 fn clone_fd(&self) -> Result<OwnedFd> {
1362 match self {
1363 TensorStorage::Dma(t) => t.clone_fd(),
1364 TensorStorage::Shm(t) => t.clone_fd(),
1365 TensorStorage::Mem(t) => t.clone_fd(),
1366 TensorStorage::Pbo(t) => t.clone_fd(),
1367 }
1368 }
1369
1370 fn memory(&self) -> TensorMemory {
1371 match self {
1372 #[cfg(any(target_os = "linux", target_os = "macos"))]
1373 TensorStorage::Dma(_) => TensorMemory::Dma,
1374 #[cfg(unix)]
1375 TensorStorage::Shm(_) => TensorMemory::Shm,
1376 TensorStorage::Mem(_) => TensorMemory::Mem,
1377 TensorStorage::Pbo(_) => TensorMemory::Pbo,
1378 }
1379 }
1380
1381 fn name(&self) -> String {
1382 match self {
1383 #[cfg(any(target_os = "linux", target_os = "macos"))]
1384 TensorStorage::Dma(t) => t.name(),
1385 #[cfg(unix)]
1386 TensorStorage::Shm(t) => t.name(),
1387 TensorStorage::Mem(t) => t.name(),
1388 TensorStorage::Pbo(t) => t.name(),
1389 }
1390 }
1391
1392 fn shape(&self) -> &[usize] {
1393 match self {
1394 #[cfg(any(target_os = "linux", target_os = "macos"))]
1395 TensorStorage::Dma(t) => t.shape(),
1396 #[cfg(unix)]
1397 TensorStorage::Shm(t) => t.shape(),
1398 TensorStorage::Mem(t) => t.shape(),
1399 TensorStorage::Pbo(t) => t.shape(),
1400 }
1401 }
1402
1403 fn reshape(&mut self, shape: &[usize]) -> Result<()> {
1404 match self {
1405 #[cfg(any(target_os = "linux", target_os = "macos"))]
1406 TensorStorage::Dma(t) => t.reshape(shape),
1407 #[cfg(unix)]
1408 TensorStorage::Shm(t) => t.reshape(shape),
1409 TensorStorage::Mem(t) => t.reshape(shape),
1410 TensorStorage::Pbo(t) => t.reshape(shape),
1411 }
1412 }
1413
1414 fn capacity_bytes(&self) -> usize {
1415 match self {
1416 #[cfg(any(target_os = "linux", target_os = "macos"))]
1417 TensorStorage::Dma(t) => t.capacity_bytes(),
1418 #[cfg(unix)]
1419 TensorStorage::Shm(t) => t.capacity_bytes(),
1420 TensorStorage::Mem(t) => t.capacity_bytes(),
1421 TensorStorage::Pbo(t) => t.capacity_bytes(),
1422 }
1423 }
1424
1425 fn set_logical_shape(&mut self, shape: &[usize]) -> Result<()> {
1426 match self {
1427 #[cfg(any(target_os = "linux", target_os = "macos"))]
1428 TensorStorage::Dma(t) => t.set_logical_shape(shape),
1429 #[cfg(unix)]
1430 TensorStorage::Shm(t) => t.set_logical_shape(shape),
1431 TensorStorage::Mem(t) => t.set_logical_shape(shape),
1432 TensorStorage::Pbo(t) => t.set_logical_shape(shape),
1433 }
1434 }
1435
1436 fn map(&self) -> Result<TensorMap<T>> {
1437 match self {
1438 #[cfg(any(target_os = "linux", target_os = "macos"))]
1439 TensorStorage::Dma(t) => t.map(),
1440 #[cfg(unix)]
1441 TensorStorage::Shm(t) => t.map(),
1442 TensorStorage::Mem(t) => t.map(),
1443 TensorStorage::Pbo(t) => t.map(),
1444 }
1445 }
1446
1447 fn buffer_identity(&self) -> &BufferIdentity {
1448 match self {
1449 #[cfg(any(target_os = "linux", target_os = "macos"))]
1450 TensorStorage::Dma(t) => t.buffer_identity(),
1451 #[cfg(unix)]
1452 TensorStorage::Shm(t) => t.buffer_identity(),
1453 TensorStorage::Mem(t) => t.buffer_identity(),
1454 TensorStorage::Pbo(t) => t.buffer_identity(),
1455 }
1456 }
1457
1458 /// Forward a sub-region view to the active backend, re-wrapping the
1459 /// backend's view (which shares the parent's allocation and
1460 /// [`BufferIdentity`]) back into the matching `TensorStorage` variant. Each
1461 /// backend's `view` is its own `TensorTrait::view` override; this single
1462 /// match is the only per-variant dispatch (`Tensor::subview` calls through
1463 /// here rather than matching the storage itself).
1464 fn view(&self, offset_bytes: usize, shape: &[usize]) -> Result<Self> {
1465 match self {
1466 #[cfg(any(target_os = "linux", target_os = "macos"))]
1467 TensorStorage::Dma(t) => t.view(offset_bytes, shape).map(TensorStorage::Dma),
1468 #[cfg(unix)]
1469 TensorStorage::Shm(t) => t.view(offset_bytes, shape).map(TensorStorage::Shm),
1470 TensorStorage::Mem(t) => t.view(offset_bytes, shape).map(TensorStorage::Mem),
1471 TensorStorage::Pbo(t) => t.view(offset_bytes, shape).map(TensorStorage::Pbo),
1472 }
1473 }
1474}
1475
1476/// Multi-backend tensor with optional image format metadata.
1477///
1478/// When `format` is `Some`, this tensor represents an image. Width, height,
1479/// and channels are derived from `shape` + `format`. When `format` is `None`,
1480/// this is a raw tensor (identical to the pre-refactoring behavior).
1481#[derive(Debug)]
1482pub struct Tensor<T>
1483where
1484 T: Num + Clone + fmt::Debug + Send + Sync,
1485{
1486 /// CUDA registration for this tensor, if any. Set after creation by
1487 /// the image crate once a PBO is registered with CUDA interop.
1488 ///
1489 /// MUST be declared before `storage`: CUDA must unregister the GL buffer
1490 /// before storage's Drop deletes it (cudaGraphicsUnregisterResource before
1491 /// glDeleteBuffers). Rust drops fields in declaration order.
1492 cuda: Option<crate::cuda::CudaHandle>,
1493 pub(crate) storage: TensorStorage<T>,
1494 format: Option<PixelFormat>,
1495 chroma: Option<Box<Tensor<T>>>,
1496 /// Row stride in bytes for externally allocated buffers with row padding.
1497 /// `None` means tightly packed (stride == width * bytes_per_pixel).
1498 row_stride: Option<usize>,
1499 /// Byte offset within the DMA-BUF where image data starts.
1500 /// `None` means offset 0 (data starts at the beginning of the buffer).
1501 plane_offset: Option<usize>,
1502 /// Quantization metadata for integer-typed tensors. Public access is
1503 /// gated by the `IntegerType` trait — `Tensor<f32>` etc. carry the
1504 /// field for layout uniformity but have no way to read or write it.
1505 pub(crate) quantization: Option<Quantization>,
1506 /// Optional colorimetry metadata. `None` = undefined; never auto-filled.
1507 colorimetry: Option<crate::Colorimetry>,
1508 /// Parent-image snapshot when this tensor is a [`view`](Self::view)/
1509 /// [`batch`](Self::batch) sub-region; `None` for a whole tensor. Lets the GL
1510 /// backend key its import on the parent and render the view as a
1511 /// `glViewport`/`glScissor` ROI. See [`ViewOrigin`].
1512 view_origin: Option<ViewOrigin>,
1513}
1514
1515impl<T> Tensor<T>
1516where
1517 T: Num + Clone + fmt::Debug + Send + Sync,
1518{
1519 /// Wrap a TensorStorage in a Tensor with no image metadata.
1520 pub(crate) fn wrap(storage: TensorStorage<T>) -> Self {
1521 Self {
1522 storage,
1523 format: None,
1524 chroma: None,
1525 row_stride: None,
1526 plane_offset: None,
1527 quantization: None,
1528 cuda: None,
1529 colorimetry: None,
1530 view_origin: None,
1531 }
1532 }
1533
1534 /// Construct a tensor from a row-major element slice + shape. Allocates a
1535 /// new buffer (`TensorMemory::Mem`) and memcpys the contents; caller
1536 /// retains ownership of the input slice.
1537 ///
1538 /// # Errors
1539 ///
1540 /// - [`Error::InvalidShape`] if `values.len() != shape.iter().product()`.
1541 /// - Propagates any allocation error from [`Self::new`].
1542 pub fn from_slice(values: &[T], shape: &[usize]) -> Result<Self>
1543 where
1544 T: Copy,
1545 {
1546 let expected: usize = shape.iter().product();
1547 if values.len() != expected {
1548 return Err(Error::InvalidShape(format!(
1549 "from_slice: values.len()={} but shape product={expected} (shape={shape:?})",
1550 values.len()
1551 )));
1552 }
1553 let t = Self::new(shape, Some(TensorMemory::Mem), None)?;
1554 {
1555 let mut m = t.map()?;
1556 m.as_mut_slice().copy_from_slice(values);
1557 }
1558 Ok(t)
1559 }
1560
1561 /// Wrap externally-owned memory as a tensor without copying. The tensor
1562 /// borrows `[ptr, ptr + shape.product() * size_of::<T>())` as
1563 /// [`TensorMemory::Mem`]; `owner`, when `Some`, co-owns the source so it
1564 /// outlives the tensor (and all derived views/maps). See [`ForeignOwner`].
1565 ///
1566 /// The canonical use is CUDA zero-copy: allocate host-coherent memory
1567 /// (`cudaHostAlloc`), wrap the host pointer here, and bind the matching
1568 /// device pointer to the inference engine — reads and writes hit the same
1569 /// physical buffer with no host copy. The identical primitive backs the
1570 /// Python `Tensor.from_numpy` zero-copy borrow (owner = the NumPy object).
1571 ///
1572 /// # Safety
1573 ///
1574 /// `ptr` must be non-null, aligned to `align_of::<T>()`, and valid for
1575 /// `shape.product()` elements of `T` for as long as the returned tensor —
1576 /// and every view/map sharing its backing — is alive. Pass an `owner` that
1577 /// co-owns the source to uphold that contract.
1578 ///
1579 /// # Errors
1580 ///
1581 /// [`Error::InvalidSize`] if `shape` is empty.
1582 pub unsafe fn from_foreign(
1583 ptr: *mut T,
1584 shape: &[usize],
1585 owner: Option<crate::ForeignOwner>,
1586 name: Option<&str>,
1587 ) -> Result<Self> {
1588 if shape.is_empty() {
1589 return Err(Error::InvalidSize(0));
1590 }
1591 if ptr.is_null() {
1592 return Err(Error::InvalidArgument(
1593 "from_foreign: ptr must be non-null".to_owned(),
1594 ));
1595 }
1596 shape
1597 .iter()
1598 .copied()
1599 .try_fold(1usize, |acc, dim| acc.checked_mul(dim))
1600 .ok_or_else(|| {
1601 Error::InvalidArgument(format!(
1602 "from_foreign: shape.product() overflows usize (shape={shape:?})"
1603 ))
1604 })?;
1605 let mem = MemTensor::<T>::from_foreign(ptr, shape, owner, name);
1606 Ok(Self::wrap(TensorStorage::Mem(mem)))
1607 }
1608
1609 /// Construct a tensor from a 3-D ndarray view. Respects strides — one
1610 /// copy in all cases; contiguous views take a memcpy fast path.
1611 ///
1612 /// Only available when the `ndarray` feature is enabled.
1613 #[cfg(feature = "ndarray")]
1614 pub fn from_arrayview3(view: ndarray::ArrayView3<'_, T>) -> Result<Self>
1615 where
1616 T: Copy,
1617 {
1618 let (h, w, c) = view.dim();
1619 let t = Self::new(&[h, w, c], Some(TensorMemory::Mem), None)?;
1620 {
1621 let mut m = t.map()?;
1622 let dst = m.as_mut_slice();
1623 if let Some(src) = view.as_slice() {
1624 dst.copy_from_slice(src);
1625 } else {
1626 for (d, &s) in dst.iter_mut().zip(view.iter()) {
1627 *d = s;
1628 }
1629 }
1630 }
1631 Ok(t)
1632 }
1633
1634 /// Create a new tensor with the given shape, memory type, and optional
1635 /// name. If no name is given, a random name will be generated. If no
1636 /// memory type is given, the best available memory type will be chosen
1637 /// based on the platform and environment variables.
1638 ///
1639 /// On Linux platforms, the order of preference is: Dma -> Shm -> Mem.
1640 /// On other Unix platforms (macOS), the order is: Shm -> Mem.
1641 /// On non-Unix platforms, only Mem is available.
1642 ///
1643 /// # Environment Variables
1644 /// - `EDGEFIRST_TENSOR_FORCE_MEM`: If set to a non-zero and non-false
1645 /// value, forces the use of regular system memory allocation
1646 /// (`TensorMemory::Mem`) regardless of platform capabilities.
1647 ///
1648 /// # Example
1649 /// ```rust
1650 /// use edgefirst_tensor::{Error, Tensor, TensorMemory, TensorTrait};
1651 /// # fn main() -> Result<(), Error> {
1652 /// let tensor = Tensor::<f32>::new(&[2, 3, 4], Some(TensorMemory::Mem), Some("test_tensor"))?;
1653 /// assert_eq!(tensor.memory(), TensorMemory::Mem);
1654 /// assert_eq!(tensor.name(), "test_tensor");
1655 /// # Ok(())
1656 /// # }
1657 /// ```
1658 pub fn new(shape: &[usize], memory: Option<TensorMemory>, name: Option<&str>) -> Result<Self> {
1659 let _span = tracing::trace_span!(
1660 "tensor.alloc",
1661 ?shape,
1662 memory = ?memory,
1663 dtype = std::any::type_name::<T>(),
1664 )
1665 .entered();
1666 #[cfg_attr(not(target_os = "linux"), allow(unused_mut))]
1667 let mut t = TensorStorage::new(shape, memory, name).map(Self::wrap)?;
1668 // Best-effort: attach a CUDA ExternalMemory handle for DMA tensors on
1669 // CUDA-capable hosts. Never blocks tensor creation on failure.
1670 // RUNTIME-UNVALIDATED: no CUDA+dma_heap test platform available; ABI
1671 // layout-asserted vs. CUDA 12.6 driver_types.h; mechanism proven by
1672 // gpu-probe O5 on Orin.
1673 #[cfg(target_os = "linux")]
1674 t.try_init_dma_cuda();
1675 Ok(t)
1676 }
1677
1678 /// Create an image tensor with the given format.
1679 pub fn image(
1680 width: usize,
1681 height: usize,
1682 format: PixelFormat,
1683 memory: Option<TensorMemory>,
1684 ) -> Result<Self>
1685 where
1686 T: 'static,
1687 {
1688 // Shape comes from the shared `PixelFormat::image_shape` helper (packed /
1689 // planar / semi-planar NV12·NV16). NV12 supports odd dimensions via the
1690 // `H + ceil(H/2)` combined-plane height.
1691 // The `T: 'static` bound is required by the macOS IOSurface path below.
1692 let shape = format.image_shape(width, height).ok_or_else(|| {
1693 Error::InvalidArgument(format!(
1694 "invalid dimensions {width}x{height} for format {format:?}"
1695 ))
1696 })?;
1697
1698 // macOS Dma path: allocate a format-aware IOSurface (FourCC +
1699 // 2D dimensions) so the GL backend can bind it via
1700 // `EGL_ANGLE_iosurface_client_buffer`. Without this, the IOSurface
1701 // would default to a generic byte buffer (FourCC 'L008') and
1702 // ANGLE would reject the import with `EGL_BAD_ATTRIBUTE`.
1703 //
1704 // Guard: IOSurface rounds `bytes_per_row` up to 64-byte alignment.
1705 // If the natural row pitch (`width * channels * sizeof(T)`) is not
1706 // already 64-byte aligned, the padded allocation cannot be mapped
1707 // as a contiguous packed tensor — CPU reads/writes would use the
1708 // wrong stride.
1709 //
1710 // Explicit-Dma contract: when the caller passes
1711 // `Some(TensorMemory::Dma)` they have asked for an
1712 // **image-formatted IOSurface**. Silently downgrading to the
1713 // generic 'L008' byte-bag when alignment fails buries the
1714 // mismatch — the caller only finds out hours later when ANGLE
1715 // (or any GL importer) rejects the bind with
1716 // `EGL_BAD_ATTRIBUTE`. Same anti-pattern bit us previously on
1717 // Mali GPUs with DMA-BUF padding. The right behaviour is to
1718 // fail loudly here with the alignment requirement spelled out
1719 // so the caller can either pick aligned dimensions, request
1720 // SHM/Mem explicitly, or pass `memory=None` for auto-select.
1721 #[cfg(target_os = "macos")]
1722 if matches!(memory, Some(TensorMemory::Dma)) {
1723 // For planar formats the IOSurface stacks channels
1724 // vertically (channels * height rows), so the row stride is
1725 // single-channel width * sizeof(T). Packed formats keep the
1726 // natural width * channels * sizeof(T) stride.
1727 let natural_row_bytes = match format.layout() {
1728 PixelLayout::Planar => width * std::mem::size_of::<T>(),
1729 _ => width * format.channels() * std::mem::size_of::<T>(),
1730 };
1731 // A format with a real IOSurface FourCC (RGBA/BGRA/YUYV packed,
1732 // GREY/NV12/NV16/NV24 as R8) tolerates a non-64-aligned natural
1733 // pitch: the surface is allocated with its own 64-aligned
1734 // `bytes_per_row`, the tensor records that stride below, and a CPU
1735 // map iterates rows correctly via the strided-map path while the GL
1736 // import uses the surface's pitch directly — fully zero-copy.
1737 // Planar (the F16 RGBA16F packing) is consumed flat as
1738 // `[1, C, H, W]` with no stride, so it still requires an aligned
1739 // pitch; and formats without a FourCC would fall through to a
1740 // generic byte-bag GL can't bind. Both fail loudly rather than
1741 // silently downgrade.
1742 let has_image_fourcc = dtype_of::<T>()
1743 .and_then(|dt| crate::iosurface::image_iosurface_layout(format, dt))
1744 .is_some();
1745 let padded_ok = has_image_fourcc && format.layout() != PixelLayout::Planar;
1746 if !natural_row_bytes.is_multiple_of(64) && !padded_ok {
1747 let elem_size = std::mem::size_of::<T>();
1748 let per_pixel_bytes = match format.layout() {
1749 PixelLayout::Planar => elem_size.max(1),
1750 _ => format.channels().max(1) * elem_size.max(1),
1751 };
1752 // Compute the next 64-byte-aligned width by rounding the
1753 // natural row pitch up to the next multiple of 64 and
1754 // dividing back by per-pixel bytes. This handles every
1755 // `per_pixel_bytes` value correctly:
1756 //
1757 // * Divisors of 64 (1/2/4/8/16/32/64) → the suggestion
1758 // is always 64-byte aligned.
1759 // * Non-divisors of 64 (e.g. RGB u8 with 3 B/pixel) →
1760 // the next aligned row pitch may not be an integer
1761 // multiple of per_pixel_bytes (3 doesn't divide 64
1762 // in any way), so a "pad width to N" suggestion is
1763 // structurally impossible — omit the suggestion
1764 // instead of printing a wrong number.
1765 // * per_pixel_bytes > 64 → same situation, also
1766 // omitted; the previous formula divided by zero.
1767 //
1768 // The error always names the alignment requirement
1769 // verbatim and lists the two non-DMA alternatives so
1770 // the caller has at least one always-applicable fix.
1771 let aligned_row_bytes = natural_row_bytes.next_multiple_of(64);
1772 let pad_hint =
1773 if per_pixel_bytes > 0 && aligned_row_bytes.is_multiple_of(per_pixel_bytes) {
1774 let w = aligned_row_bytes / per_pixel_bytes;
1775 format!("Pad width to {w} (the next 64-byte-aligned stride), ")
1776 } else {
1777 String::new()
1778 };
1779 return Err(Error::InvalidArgument(format!(
1780 "Tensor::image: {format:?} {width}x{height} with element \
1781 size {elem_size} produces a {natural_row_bytes}-byte natural \
1782 row pitch, which is not 64-byte aligned. \
1783 IOSurface rounds bytes_per_row up to 64 bytes, so a \
1784 contiguous CPU map of this tensor would read garbage. \
1785 {pad_hint}pass memory=None to auto-fall-back to SHM, or \
1786 pass memory=Some(TensorMemory::Shm) or \
1787 Some(TensorMemory::Mem) explicitly."
1788 )));
1789 }
1790 // Alignment OK. Try image-formatted IOSurface; on any
1791 // structural failure (unsupported (format, dtype) combo,
1792 // out-of-memory, etc.) fall through to SHM. dtype_of
1793 // returns None for non-standard numeric T used in tests —
1794 // those legitimately have no IOSurface FourCC mapping.
1795 if let Some(dtype) = dtype_of::<T>() {
1796 if let Ok(storage) = TensorStorage::<T>::new_image_iosurface(
1797 width, height, format, dtype, &shape, None,
1798 ) {
1799 let mut t = Self::wrap(storage);
1800 t.format = Some(format);
1801 // IOSurface rounds `bytes_per_row` up to 64 bytes. When that
1802 // pitch exceeds the natural packed/planar row stride, record
1803 // it so CPU consumers iterate rows correctly (the GL import
1804 // already uses the surface's own pitch). For 64-aligned rows
1805 // — the common model-input case — the two match and no stride
1806 // is stored, leaving the flat mapping unchanged.
1807 if let TensorStorage::Dma(ref io) = t.storage {
1808 let bpr = io.bytes_per_row();
1809 if let Some(natural) = t.effective_row_stride() {
1810 if bpr > natural {
1811 t.set_row_stride_unchecked(bpr);
1812 }
1813 }
1814 }
1815 return Ok(t);
1816 }
1817 }
1818 // Unsupported (format, dtype) on the IOSurface path —
1819 // e.g. PlanarRgb u8 has no IOSurface FourCC mapping today
1820 // (only F16 PlanarRgb is wired). Fall through to SHM/Mem
1821 // since the caller asked for Dma but no image-formatted DMA
1822 // storage exists for this combination.
1823 }
1824
1825 // Compute the **64-byte-aligned** row stride for every image layout.
1826 //
1827 // Embedded GPUs reject `eglCreateImage` DMA-BUF imports whose row pitch
1828 // is not 64-byte aligned: Mali returns `EGL_BAD_ALLOC`, Vivante
1829 // `EGL_BAD_ACCESS`. This bit packed RGBA/RGB destinations at odd widths
1830 // AND at even non-multiple-of-16 widths (e.g. 321→1284, 322→1288 bytes —
1831 // neither divisible by 64), so an odd-source → RGBA convert failed on
1832 // imx95/imx8mp while succeeding on V3D/Tegra. Semi-planar already aligned
1833 // here; we now align packed and planar identically so every image()
1834 // allocation is GPU-importable regardless of width.
1835 //
1836 // The per-layout natural pitch and total row count:
1837 // * SemiPlanar `[total_h, width]` — pitch = even(width)·elem, rows = total_h
1838 // * Packed `[height, width, ch]` — pitch = width·ch·elem, rows = height
1839 // * Planar `[ch, height, width]` — pitch = width·elem, rows = ch·height
1840 // Allocation byte size = `aligned_stride · total_rows` (NOT the shape
1841 // product, which reflects only the logical width and under-allocates the
1842 // padding on odd / unaligned widths).
1843 let elem = std::mem::size_of::<T>();
1844 let channels = format.channels();
1845 let (natural_stride, total_rows) = match format.layout() {
1846 PixelLayout::SemiPlanar => (width.next_multiple_of(2) * elem, shape[0]),
1847 PixelLayout::Packed => (width * channels * elem, height),
1848 PixelLayout::Planar => (width * elem, channels * height),
1849 };
1850 let aligned_stride = natural_stride.next_multiple_of(64);
1851 let semi = format.layout() == PixelLayout::SemiPlanar;
1852
1853 // DMA buffers MUST carry a 64-aligned row pitch — Mali/Vivante reject a
1854 // DMA-BUF EGLImage whose pitch is not 64-aligned. Semi-planar also needs
1855 // the aligned pitch on every backend (its chroma-plane offset math
1856 // assumes it). Packed/planar on host-only memory (Mem/Shm) keep the
1857 // natural tight pitch so the many flat CPU consumers are unaffected.
1858 let host_stride = if semi { aligned_stride } else { natural_stride };
1859 let host_byte_size = host_stride * total_rows;
1860 #[cfg(target_os = "linux")]
1861 let dma_byte_size = aligned_stride * total_rows;
1862
1863 // `used_stride` is the actual row pitch of the storage created below.
1864 let (storage, used_stride) = match memory {
1865 #[cfg(target_os = "linux")]
1866 Some(TensorMemory::Dma) => (
1867 TensorStorage::<T>::new_dma_with_byte_size(&shape, dma_byte_size, None)?,
1868 aligned_stride,
1869 ),
1870 #[cfg(unix)]
1871 Some(TensorMemory::Shm) => (
1872 TensorStorage::<T>::new_shm_with_byte_size(&shape, host_byte_size, None)?,
1873 host_stride,
1874 ),
1875 Some(TensorMemory::Mem) => (
1876 TensorStorage::<T>::new_mem_with_byte_size(&shape, host_byte_size, None)?,
1877 host_stride,
1878 ),
1879 #[allow(unused_variables)]
1880 Some(other) => {
1881 // PBO and any future variants: fall through to standard new().
1882 return {
1883 let mut t = Self::new(&shape, Some(other), None)?;
1884 t.format = Some(format);
1885 Ok(t)
1886 };
1887 }
1888 None => {
1889 // Auto-select priority: DMA → Mem (DMA gets the 64-aligned
1890 // pitch; the Mem fallback keeps the tight host pitch, so the
1891 // recorded stride matches the storage used). Shm is NOT
1892 // auto-selected — it offers no advantage over Mem for an
1893 // in-process image and Mem always succeeds, so it sits below Mem
1894 // and is reached only via an explicit `TensorMemory::Shm`.
1895 #[cfg(target_os = "linux")]
1896 {
1897 match TensorStorage::<T>::new_dma_with_byte_size(&shape, dma_byte_size, None) {
1898 Ok(s) => (s, aligned_stride),
1899 Err(_) => (
1900 TensorStorage::<T>::new_mem_with_byte_size(
1901 &shape,
1902 host_byte_size,
1903 None,
1904 )?,
1905 host_stride,
1906 ),
1907 }
1908 }
1909 #[cfg(not(target_os = "linux"))]
1910 {
1911 (
1912 TensorStorage::<T>::new_mem_with_byte_size(&shape, host_byte_size, None)?,
1913 host_stride,
1914 )
1915 }
1916 }
1917 };
1918
1919 let mut t = Self::wrap(storage);
1920 t.format = Some(format);
1921 // Record the row stride when it exceeds the natural tight pitch (padding
1922 // is present — DMA packed/planar at an unaligned width, or always for
1923 // semi-planar), mirroring the IOSurface path above. Aligned-width and
1924 // host-only packed/planar images keep their flat layout with no explicit
1925 // stride; `effective_row_stride()` then falls back to the identical
1926 // computed pitch. When padding IS present, consumers must iterate rows by
1927 // `effective_row_stride()` to skip it.
1928 if semi || used_stride > natural_stride {
1929 t.set_row_stride_unchecked(used_stride);
1930 }
1931 debug_assert!(
1932 t.row_stride.is_some() || !semi,
1933 "image() must always set row_stride for semi-planar tensors"
1934 );
1935 #[cfg(target_os = "linux")]
1936 t.try_init_dma_cuda();
1937 Ok(t)
1938 }
1939
1940 /// Create a DMA-backed image tensor with an explicit row stride that
1941 /// may exceed the natural `width * channels * sizeof(T)` pitch.
1942 ///
1943 /// Used for image tensors that need GPU pitch alignment padding: the
1944 /// underlying DMA-BUF is sized to `row_stride * height` bytes, but
1945 /// the tensor's logical shape stays at `[height, width, channels]`.
1946 /// `width()` / `height()` / `shape()` continue to report the
1947 /// user-requested values; the padding is visible only via
1948 /// `row_stride()` / `effective_row_stride()` and is automatically
1949 /// propagated to the GL backend's EGLImage import so Mali Valhall
1950 /// accepts the buffer.
1951 ///
1952 /// # Supported formats
1953 ///
1954 /// Currently only **packed** pixel layouts (RGBA8, BGRA8, RGB888,
1955 /// Grey, etc.) are supported — the formats the GL backend uses as
1956 /// render destinations. Semi-planar formats (NV12, NV16) come from
1957 /// external allocators (camera capture, video decoders) and are
1958 /// imported via `TensorDyn::from_fd` + `set_row_stride`, which
1959 /// already supports padded strides.
1960 ///
1961 /// # Supported memory
1962 ///
1963 /// Currently only `TensorMemory::Dma` is supported. PBO and Mem
1964 /// storage don't go through EGLImage import so they don't need
1965 /// pitch alignment; if you pass any other memory type this returns
1966 /// `NotImplemented`. `None` (auto-select) is treated as `Dma`.
1967 ///
1968 /// # Errors
1969 ///
1970 /// - `InvalidArgument` if `row_stride_bytes < width * channels * sizeof(T)`
1971 /// (the requested stride would not fit a single row)
1972 /// - `NotImplemented` for non-packed formats or non-DMA memory
1973 /// - `IoError` if the DMA-heap allocation fails (propagated from
1974 /// `DmaTensor::new_with_byte_size`)
1975 pub fn image_with_stride(
1976 width: usize,
1977 height: usize,
1978 format: PixelFormat,
1979 row_stride_bytes: usize,
1980 memory: Option<TensorMemory>,
1981 ) -> Result<Self> {
1982 // DMA backing (the only thing this constructor produces) is
1983 // Linux-only. On macOS/BSD/Windows the non-Linux block below is
1984 // the only compiled body and returns `NotImplemented` directly;
1985 // on Linux the non-Linux block is cfg-removed and the function
1986 // falls through to the real validation + allocation path. Each
1987 // target compiles exactly one of the two blocks, and the block
1988 // serves as the function's tail expression in both cases — so
1989 // neither needs an explicit `return` (avoids
1990 // `clippy::needless_return` on the macOS CI gate).
1991 #[cfg(not(target_os = "linux"))]
1992 {
1993 let _ = (width, height, format, row_stride_bytes, memory);
1994 Err(Error::NotImplemented(
1995 "image_with_stride requires DMA support (Linux only)".to_owned(),
1996 ))
1997 }
1998
1999 #[cfg(target_os = "linux")]
2000 {
2001 if format.layout() != PixelLayout::Packed {
2002 return Err(Error::NotImplemented(format!(
2003 "Tensor::image_with_stride only supports packed pixel layouts, got {format:?}"
2004 )));
2005 }
2006 let elem = std::mem::size_of::<T>();
2007 let min_stride = width
2008 .checked_mul(format.channels())
2009 .and_then(|p| p.checked_mul(elem))
2010 .ok_or_else(|| {
2011 Error::InvalidArgument(format!(
2012 "image_with_stride: width {width} × channels {} × sizeof::<T>={elem} \
2013 overflows usize",
2014 format.channels()
2015 ))
2016 })?;
2017 if row_stride_bytes < min_stride {
2018 return Err(Error::InvalidArgument(format!(
2019 "image_with_stride: row_stride {row_stride_bytes} < minimum {min_stride} \
2020 ({width} px × {} ch × {elem} B)",
2021 format.channels()
2022 )));
2023 }
2024 let total_byte_size = row_stride_bytes.checked_mul(height).ok_or_else(|| {
2025 Error::InvalidArgument(format!(
2026 "image_with_stride: row_stride {row_stride_bytes} × height {height} overflows usize"
2027 ))
2028 })?;
2029
2030 let shape = vec![height, width, format.channels()];
2031
2032 let storage = match memory {
2033 Some(TensorMemory::Dma) | None => {
2034 TensorStorage::<T>::new_dma_with_byte_size(&shape, total_byte_size, None)?
2035 }
2036 Some(other) => {
2037 return Err(Error::NotImplemented(format!(
2038 "image_with_stride: only TensorMemory::Dma is supported, got {other:?}"
2039 )));
2040 }
2041 };
2042
2043 let mut t = Self::wrap(storage);
2044 t.format = Some(format);
2045 t.row_stride = Some(row_stride_bytes);
2046 // Match new()/from_fd(): a DMA tensor must attempt CUDA external-
2047 // memory import so a strided DMA buffer is also zero-copy
2048 // CUDA-mappable (no-op when libcudart is absent).
2049 t.try_init_dma_cuda();
2050 Ok(t)
2051 }
2052 }
2053
2054 /// Attach format metadata to an existing tensor.
2055 ///
2056 /// # Arguments
2057 ///
2058 /// * `format` - The pixel format to attach
2059 ///
2060 /// # Returns
2061 ///
2062 /// `Ok(())` on success, with the format stored as metadata on the tensor.
2063 ///
2064 /// # Errors
2065 ///
2066 /// Returns `Error::InvalidShape` if the tensor shape is incompatible with
2067 /// the format's layout (packed expects `[H, W, C]`, planar expects
2068 /// `[C, H, W]`, semi-planar expects `[H*k, W]` with format-specific
2069 /// height constraints).
2070 pub fn set_format(&mut self, format: PixelFormat) -> Result<()> {
2071 let shape = self.shape();
2072 match format.layout() {
2073 PixelLayout::Packed => {
2074 if shape.len() != 3 || shape[2] != format.channels() {
2075 return Err(Error::InvalidShape(format!(
2076 "packed format {format:?} expects [H, W, {}], got {shape:?}",
2077 format.channels()
2078 )));
2079 }
2080 }
2081 PixelLayout::Planar => {
2082 if shape.len() != 3 || shape[0] != format.channels() {
2083 return Err(Error::InvalidShape(format!(
2084 "planar format {format:?} expects [{}, H, W], got {shape:?}",
2085 format.channels()
2086 )));
2087 }
2088 }
2089 PixelLayout::SemiPlanar => {
2090 if shape.len() != 2 {
2091 return Err(Error::InvalidShape(format!(
2092 "semi-planar format {format:?} expects [H*k, W], got {shape:?}"
2093 )));
2094 }
2095 match format {
2096 // Combined-plane height is `H + ceil(H/2)` (luma + chroma
2097 // rows). For even H that is `3H/2` (≡ 0 mod 3); for odd H it
2098 // is `(3H+1)/2` (≡ 2 mod 3). Only totals ≡ 1 mod 3 are
2099 // unreachable, so reject just those — odd-height NV12 is
2100 // valid (e.g. 725 rows for a 483-tall image).
2101 PixelFormat::Nv12 if shape[0] % 3 == 1 => {
2102 return Err(Error::InvalidShape(format!(
2103 "NV12 contiguous shape[0] must be H + ceil(H/2) for some height; \
2104 {} is unreachable (≡ 1 mod 3)",
2105 shape[0]
2106 )));
2107 }
2108 PixelFormat::Nv16 if !shape[0].is_multiple_of(2) => {
2109 return Err(Error::InvalidShape(format!(
2110 "NV16 contiguous shape[0] must be even, got {}",
2111 shape[0]
2112 )));
2113 }
2114 // NV24 (4:4:4): combined-plane height is 3H (Y + 2H chroma).
2115 PixelFormat::Nv24 if !shape[0].is_multiple_of(3) => {
2116 return Err(Error::InvalidShape(format!(
2117 "NV24 contiguous shape[0] must be a multiple of 3 (= 3H), got {}",
2118 shape[0]
2119 )));
2120 }
2121 _ => {}
2122 }
2123 }
2124 }
2125 // Clear stored stride/offset when format changes — they may be invalid
2126 // for the new format. Caller must re-set after changing format.
2127 if self.format != Some(format) {
2128 self.row_stride = None;
2129 self.plane_offset = None;
2130 match self.storage {
2131 TensorStorage::Mem(ref mut m) => m.set_offset(0),
2132 #[cfg(target_os = "linux")]
2133 TensorStorage::Dma(ref mut dma) => dma.mmap_offset = 0,
2134 _ => {}
2135 }
2136 }
2137 self.format = Some(format);
2138 Ok(())
2139 }
2140
2141 /// Set this tensor's logical dimensions and pixel format to a decoded
2142 /// image, reusing the existing allocation. The shape is derived from the
2143 /// format layout; fails with `Error::InsufficientCapacity` if the
2144 /// allocation cannot hold `width`×`height` in `format`, or
2145 /// `Error::InvalidArgument` if the dimensions are invalid for the format.
2146 ///
2147 /// For NV12/NV16/NV24 the buffer width is rounded up to even (a chroma-plane
2148 /// interleaving requirement); the true odd width is reported by the decoder
2149 /// in `ImageInfo` and trimmed by a `convert()` crop. See
2150 /// [`PixelFormat::image_shape`].
2151 ///
2152 /// When the backing has a fixed physical row pitch (an IOSurface's
2153 /// 64-aligned `bytesPerRow`) that exceeds the new format's natural row
2154 /// stride — i.e. a reused max-sized pool tensor reconfigured to a smaller
2155 /// image — the physical pitch is preserved as the tensor's `row_stride`.
2156 /// This keeps the **physical grid** (allocation stride/surface) fixed while
2157 /// the **logical ROI** (this image's W×H) changes, so the decode writes rows
2158 /// at the surface's real stride and the GPU samples them at the same stride.
2159 /// Exact-sized buffers (pitch == natural) stay tightly packed unchanged.
2160 pub fn configure_image(
2161 &mut self,
2162 width: usize,
2163 height: usize,
2164 format: PixelFormat,
2165 ) -> Result<()> {
2166 let shape = format.image_shape(width, height).ok_or_else(|| {
2167 Error::InvalidArgument(format!(
2168 "invalid dimensions {width}x{height} for format {format:?}"
2169 ))
2170 })?;
2171 // Capture the pre-existing row stride before `set_format` clears it.
2172 // For pool tensors that were allocated at a larger width (e.g. 1920-wide
2173 // pool decoding a 789-wide image), this preserves the backing pitch so
2174 // rows are still written at the correct physical stride.
2175 let prior_stride = self.row_stride;
2176
2177 self.storage.set_logical_shape(&shape)?;
2178 self.set_format(format)?; // clears any stale row_stride
2179
2180 // Restore the correct row pitch for the new geometry. A DMA buffer of
2181 // ANY layout, and a semi-planar buffer on any backing, MUST carry a
2182 // 64-byte-aligned pitch: Mali/Vivante reject a DMA-BUF EGLImage whose
2183 // row pitch is not 64-aligned (`EGL_BAD_ALLOC` / `EGL_BAD_ACCESS`), and
2184 // semi-planar chroma-offset math assumes it. This mirrors the pitch
2185 // `Tensor::image()` allocates, so a recycled pool buffer imports
2186 // identically to a fresh one — without it a `configure_image()`'d packed
2187 // buffer (e.g. Y800 96-wide → tight pitch 96) failed convert() on
2188 // imx95/imx8mp via the texture-upload fallback while the fresh oracle
2189 // (aligned 128) succeeded. Packed/planar on host-only memory (Mem/Shm)
2190 // keep the tight pitch so flat CPU consumers stay unaffected — matching
2191 // `image()`'s `host_stride` rule.
2192 let elem = std::mem::size_of::<T>();
2193 let channels = format.channels();
2194 let (min_stride, total_rows) = match format.layout() {
2195 PixelLayout::SemiPlanar => (width.next_multiple_of(2) * elem, shape[0]),
2196 PixelLayout::Packed => (width * channels * elem, height),
2197 PixelLayout::Planar => (width * elem, channels * height),
2198 };
2199 let needs_align = self.storage.memory() == TensorMemory::Dma
2200 || format.layout() == PixelLayout::SemiPlanar;
2201
2202 let active_stride = if let Some(pitch) = self.storage.backing_row_stride() {
2203 // macOS IOSurface: use the surface's native pitch.
2204 let natural = self.effective_row_stride().unwrap_or(0);
2205 if pitch > natural {
2206 self.set_row_stride_unchecked(pitch);
2207 pitch
2208 } else {
2209 natural
2210 }
2211 } else if needs_align {
2212 // Priority:
2213 // 1. Prior stride (pool reuse): if the pre-existing stride is
2214 // 64-aligned, >= this layout's minimum row, and fits the
2215 // allocation, keep it. This is the hot-loop reuse case (large
2216 // pool, small image).
2217 // 2. Compute a fresh 64-aligned stride for the current width.
2218 let aligned = min_stride.next_multiple_of(64);
2219 let capacity = self.storage.capacity_bytes();
2220
2221 let candidate = if let Some(ps) = prior_stride {
2222 if ps >= min_stride && ps % 64 == 0 && ps * total_rows <= capacity {
2223 ps
2224 } else {
2225 aligned
2226 }
2227 } else {
2228 aligned
2229 };
2230
2231 if candidate * total_rows <= capacity {
2232 self.set_row_stride_unchecked(candidate);
2233 candidate
2234 } else {
2235 // Shouldn't happen for legitimate pools, but don't crash.
2236 self.effective_row_stride().unwrap_or(0)
2237 }
2238 } else {
2239 self.effective_row_stride().unwrap_or(0)
2240 };
2241
2242 // Ensure the active stride fits the allocation. A pool reconfigured to a
2243 // wider image than its backing would silently SIGBUS on any subsequent
2244 // map/write — catch it here instead.
2245 if needs_align && active_stride > 0 {
2246 let needed = active_stride * total_rows;
2247 let capacity = self.storage.capacity_bytes();
2248 if needed > capacity {
2249 return Err(Error::InsufficientCapacity { needed, capacity });
2250 }
2251 }
2252 Ok(())
2253 }
2254
2255 /// Allocate an image tensor sized to hold up to `width`×`height` in
2256 /// `format`, reusable for any smaller image via `configure_image`.
2257 pub fn image_with_capacity(
2258 width: usize,
2259 height: usize,
2260 format: PixelFormat,
2261 memory: Option<TensorMemory>,
2262 ) -> Result<Self>
2263 where
2264 T: 'static,
2265 {
2266 Self::image(width, height, format, memory)
2267 }
2268
2269 /// Pixel format (None if not an image).
2270 pub fn format(&self) -> Option<PixelFormat> {
2271 self.format
2272 }
2273
2274 /// Image width (None if not an image).
2275 pub fn width(&self) -> Option<usize> {
2276 let fmt = self.format?;
2277 let shape = self.shape();
2278 match fmt.layout() {
2279 PixelLayout::Packed => Some(shape[1]),
2280 PixelLayout::Planar => Some(shape[2]),
2281 PixelLayout::SemiPlanar => Some(shape[1]),
2282 }
2283 }
2284
2285 /// Image height (None if not an image).
2286 ///
2287 /// For semi-planar formats the combined-plane shape row count is divided
2288 /// by the format's luma-to-total ratio to recover logical height. This
2289 /// returns the exact logical height (including odd heights) only because
2290 /// the logical dimensions are tracked separately from the physical shape —
2291 /// `configure_image` stores the actual `(width, height)` in the format's
2292 /// `image_shape`, which round-trips losslessly via these accessors.
2293 pub fn height(&self) -> Option<usize> {
2294 let fmt = self.format?;
2295 let shape = self.shape();
2296 match fmt.layout() {
2297 PixelLayout::Packed => Some(shape[0]),
2298 PixelLayout::Planar => Some(shape[1]),
2299 PixelLayout::SemiPlanar => {
2300 if self.is_multiplane() {
2301 Some(shape[0])
2302 } else {
2303 match fmt {
2304 PixelFormat::Nv12 => Some(shape[0] * 2 / 3),
2305 PixelFormat::Nv16 => Some(shape[0] / 2),
2306 PixelFormat::Nv24 => Some(shape[0] / 3),
2307 _ => None,
2308 }
2309 }
2310 }
2311 }
2312 }
2313
2314 /// Create from separate Y and UV planes (multiplane NV12/NV16).
2315 pub fn from_planes(luma: Tensor<T>, chroma: Tensor<T>, format: PixelFormat) -> Result<Self> {
2316 if format.layout() != PixelLayout::SemiPlanar {
2317 return Err(Error::InvalidArgument(format!(
2318 "from_planes requires a semi-planar format, got {format:?}"
2319 )));
2320 }
2321 if chroma.format.is_some() || chroma.chroma.is_some() {
2322 return Err(Error::InvalidArgument(
2323 "chroma tensor must be a raw tensor (no format or chroma metadata)".into(),
2324 ));
2325 }
2326 let luma_shape = luma.shape();
2327 let chroma_shape = chroma.shape();
2328 if luma_shape.len() != 2 || chroma_shape.len() != 2 {
2329 return Err(Error::InvalidArgument(format!(
2330 "from_planes expects 2D shapes, got luma={luma_shape:?} chroma={chroma_shape:?}"
2331 )));
2332 }
2333 if luma_shape[1] != chroma_shape[1] {
2334 return Err(Error::InvalidArgument(format!(
2335 "luma width {} != chroma width {}",
2336 luma_shape[1], chroma_shape[1]
2337 )));
2338 }
2339 match format {
2340 PixelFormat::Nv12 => {
2341 if luma_shape[0] % 2 != 0 {
2342 return Err(Error::InvalidArgument(format!(
2343 "NV12 requires even luma height, got {}",
2344 luma_shape[0]
2345 )));
2346 }
2347 if chroma_shape[0] != luma_shape[0] / 2 {
2348 return Err(Error::InvalidArgument(format!(
2349 "NV12 chroma height {} != luma height / 2 ({})",
2350 chroma_shape[0],
2351 luma_shape[0] / 2
2352 )));
2353 }
2354 }
2355 PixelFormat::Nv16 => {
2356 if chroma_shape[0] != luma_shape[0] {
2357 return Err(Error::InvalidArgument(format!(
2358 "NV16 chroma height {} != luma height {}",
2359 chroma_shape[0], luma_shape[0]
2360 )));
2361 }
2362 }
2363 // NV24's chroma plane is full-resolution (2×-wide interleaved UV),
2364 // which the equal-width plane check above doesn't model. Multiplane
2365 // NV24 is unused (the JPEG decoder emits a contiguous NV24 buffer),
2366 // so it's not supported here yet.
2367 _ => {
2368 return Err(Error::InvalidArgument(format!(
2369 "from_planes only supports NV12 and NV16 (NV24 multiplane not yet \
2370 supported — use a contiguous NV24 tensor), got {format:?}"
2371 )));
2372 }
2373 }
2374
2375 Ok(Tensor {
2376 storage: luma.storage,
2377 format: Some(format),
2378 chroma: Some(Box::new(chroma)),
2379 row_stride: luma.row_stride,
2380 plane_offset: luma.plane_offset,
2381 quantization: luma.quantization,
2382 // A multiplane tensor spans two DMA-BUFs (luma + chroma); CUDA
2383 // external-memory import is per-fd, so there is no single device
2384 // pointer for the composite. Any CUDA handle the luma plane carried
2385 // is intentionally dropped — consumers needing CUDA access to
2386 // multiplane data must import each plane independently.
2387 cuda: None,
2388 colorimetry: luma.colorimetry,
2389 // A composed multiplane tensor is a whole image, not a sub-view.
2390 view_origin: None,
2391 })
2392 }
2393
2394 /// Whether this tensor uses separate plane allocations.
2395 pub fn is_multiplane(&self) -> bool {
2396 self.chroma.is_some()
2397 }
2398
2399 /// Access the chroma plane for multiplane semi-planar images.
2400 pub fn chroma(&self) -> Option<&Tensor<T>> {
2401 self.chroma.as_deref()
2402 }
2403
2404 /// Mutable access to the chroma plane for multiplane semi-planar images.
2405 pub fn chroma_mut(&mut self) -> Option<&mut Tensor<T>> {
2406 self.chroma.as_deref_mut()
2407 }
2408
2409 /// Row stride in bytes (`None` = tightly packed).
2410 pub fn row_stride(&self) -> Option<usize> {
2411 self.row_stride
2412 }
2413
2414 /// Effective row stride in bytes: the stored stride if set, otherwise the
2415 /// minimum stride computed from the format, width, and element size.
2416 /// Returns `None` only when no format is set and no explicit stride was
2417 /// stored via [`set_row_stride`](Self::set_row_stride).
2418 ///
2419 /// **GREY note:** `effective_row_stride()` for a GREY tensor returns the
2420 /// tight `width` bytes (no padding), which is what `normalize_to_numpy` and
2421 /// the CPU convert path expect. The codec's internal `native_row_stride`
2422 /// (64-byte-aligned) is used only during decoding and is not propagated to
2423 /// the tensor's stored stride, so callers reading via
2424 /// `effective_row_stride()` always see the tight value for GREY.
2425 pub fn effective_row_stride(&self) -> Option<usize> {
2426 if let Some(s) = self.row_stride {
2427 return Some(s);
2428 }
2429 let fmt = self.format?;
2430 let w = self.width()?;
2431 let elem = std::mem::size_of::<T>();
2432 Some(match fmt.layout() {
2433 PixelLayout::Packed => w * fmt.channels() * elem,
2434 PixelLayout::Planar => w * elem,
2435 // Semi-planar: minimum stride must cover the even width so the
2436 // interleaved chroma columns are byte-aligned on odd-width images.
2437 PixelLayout::SemiPlanar => w.next_multiple_of(2) * elem,
2438 })
2439 }
2440
2441 /// Set the row stride in bytes for externally allocated buffers with
2442 /// row padding (e.g. V4L2 or GStreamer allocators).
2443 ///
2444 /// The stride is propagated to the EGL DMA-BUF import attributes so
2445 /// the GPU interprets the padded buffer layout correctly. Must be
2446 /// called after [`set_format`](Self::set_format) and before the tensor
2447 /// is first passed to [`ImageProcessor::convert`]. The stored stride
2448 /// is cleared automatically if the pixel format is later changed.
2449 ///
2450 /// No stride-vs-buffer-size validation is performed because the
2451 /// backing allocation size is not reliably known: external DMA-BUFs
2452 /// may be over-allocated by the allocator, and internal tensors store
2453 /// a logical (unpadded) shape. An incorrect stride will be caught by
2454 /// the EGL driver at import time.
2455 ///
2456 /// # Arguments
2457 ///
2458 /// * `stride` - Row stride in bytes. Must be >= the minimum stride for
2459 /// the format (width * channels * sizeof(T) for packed,
2460 /// width * sizeof(T) for planar/semi-planar).
2461 ///
2462 /// # Errors
2463 ///
2464 /// * `InvalidArgument` if no pixel format is set on this tensor
2465 /// * `InvalidArgument` if `stride` is less than the minimum for the
2466 /// format and width
2467 pub fn set_row_stride(&mut self, stride: usize) -> Result<()> {
2468 let fmt = self.format.ok_or_else(|| {
2469 Error::InvalidArgument("cannot set row_stride without a pixel format".into())
2470 })?;
2471 let w = self.width().ok_or_else(|| {
2472 Error::InvalidArgument("cannot determine width for row_stride validation".into())
2473 })?;
2474 let elem = std::mem::size_of::<T>();
2475 let min_stride = match fmt.layout() {
2476 PixelLayout::Packed => w * fmt.channels() * elem,
2477 PixelLayout::Planar => w * elem,
2478 // Semi-planar: minimum must cover even width for chroma alignment.
2479 PixelLayout::SemiPlanar => w.next_multiple_of(2) * elem,
2480 };
2481 if stride < min_stride {
2482 return Err(Error::InvalidArgument(format!(
2483 "row_stride {stride} < minimum {min_stride} for {fmt:?} at width {w}"
2484 )));
2485 }
2486 self.row_stride = Some(stride);
2487 Ok(())
2488 }
2489
2490 /// Set the row stride without format validation.
2491 ///
2492 /// Use this for raw sub-tensors (e.g. chroma planes) that don't carry
2493 /// format metadata. The caller is responsible for ensuring the stride
2494 /// is valid.
2495 pub fn set_row_stride_unchecked(&mut self, stride: usize) {
2496 self.row_stride = Some(stride);
2497 }
2498
2499 /// Builder-style variant of [`set_row_stride`](Self::set_row_stride),
2500 /// consuming and returning `self`.
2501 ///
2502 /// # Errors
2503 ///
2504 /// Same conditions as [`set_row_stride`](Self::set_row_stride).
2505 pub fn with_row_stride(mut self, stride: usize) -> Result<Self> {
2506 self.set_row_stride(stride)?;
2507 Ok(self)
2508 }
2509
2510 /// Byte offset within the DMA-BUF where image data starts (`None` = 0).
2511 pub fn plane_offset(&self) -> Option<usize> {
2512 self.plane_offset
2513 }
2514
2515 /// The parent-image snapshot if this tensor is a [`view`](Self::view)/
2516 /// [`batch`](Self::batch) sub-region; `None` for a whole tensor. The GL
2517 /// backend keys its import on the parent geometry and renders this view as a
2518 /// `glViewport`/`glScissor` ROI at `(x, y, width, height)`. See [`ViewOrigin`].
2519 pub fn view_origin(&self) -> Option<ViewOrigin> {
2520 self.view_origin
2521 }
2522
2523 /// Set the byte offset within the DMA-BUF where image data starts.
2524 ///
2525 /// Propagated to `EGL_DMA_BUF_PLANE0_OFFSET_EXT` on GPU import.
2526 /// Unlike [`set_row_stride`](Self::set_row_stride), no format is required
2527 /// since the offset is format-independent.
2528 pub fn set_plane_offset(&mut self, offset: usize) {
2529 self.plane_offset = Some(offset);
2530 // The offset consulted by `map()` lives inside the storage variant.
2531 // Keep it in sync with the wrapper field for every backing that
2532 // honors it (DMA and Mem); see also the clear sites in `set_format`
2533 // and `reshape`.
2534 match self.storage {
2535 TensorStorage::Mem(ref mut m) => m.set_offset(offset),
2536 #[cfg(target_os = "linux")]
2537 TensorStorage::Dma(ref mut dma) => dma.mmap_offset = offset,
2538 _ => {}
2539 }
2540 }
2541
2542 /// Colorimetry metadata (`None` = undefined; never auto-filled).
2543 pub fn colorimetry(&self) -> Option<crate::Colorimetry> {
2544 self.colorimetry
2545 }
2546
2547 /// Attach/clear colorimetry metadata.
2548 pub fn set_colorimetry(&mut self, c: Option<crate::Colorimetry>) {
2549 self.colorimetry = c;
2550 }
2551
2552 /// Builder-style colorimetry attach.
2553 pub fn with_colorimetry(mut self, c: crate::Colorimetry) -> Self {
2554 self.colorimetry = Some(c);
2555 self
2556 }
2557
2558 /// Create a zero-copy sub-region view of this tensor's backing buffer.
2559 ///
2560 /// The returned tensor shares this tensor's allocation (no copy) and maps
2561 /// the window `[offset_bytes, offset_bytes + shape.product()*size_of::<T>())`
2562 /// measured from this tensor's own logical start. N sub-views into one
2563 /// parent can be written independently, enabling batched assembly into a
2564 /// single buffer. Identical semantics across `Mem` (shared `Arc`) and
2565 /// `Dma` (shared fd) backings.
2566 ///
2567 /// # Disjointness
2568 ///
2569 /// Independent writes are sound *only* when the windows do not overlap. The
2570 /// shared backing uses interior mutability (`UnsafeCell` cells), so two
2571 /// sub-views whose byte ranges intersect alias the same cells: writing one
2572 /// while reading or writing the other is a data race and therefore
2573 /// **undefined behaviour**. The caller is responsible for keeping the
2574 /// windows disjoint; this method does not check for overlap.
2575 ///
2576 /// # Errors
2577 ///
2578 /// - [`Error::InvalidOperation`] if the backing is not `Mem` or `Dma`, or
2579 /// if `offset_bytes` is not a multiple of `align_of::<T>()`.
2580 /// - [`Error::InsufficientCapacity`] / [`Error::InvalidSize`] if the window
2581 /// exceeds the parent allocation.
2582 pub(crate) fn subview(&self, offset_bytes: usize, shape: &[usize]) -> Result<Tensor<T>> {
2583 // Offset is absolute into the backing allocation: a sub-view of a
2584 // sub-view composes by adding this tensor's own offset.
2585 let abs_offset = self
2586 .plane_offset
2587 .unwrap_or(0)
2588 .checked_add(offset_bytes)
2589 .ok_or(Error::InvalidSize(offset_bytes))?;
2590 // Every backend exposes `view(offset, shape)` via `TensorTrait`, sharing
2591 // the resource AND `BufferIdentity` (unlike `from_fd`/`from_surface`,
2592 // which mint a fresh identity). The GL backend keys the import on the
2593 // shared identity so offset-distinct sub-views of one buffer reuse a
2594 // single import and address their window via `glViewport`. `Mem`/`Shm`
2595 // share via the allocation `Arc` / a cloned fd; `Pbo` via the GL-buffer
2596 // `Arc`; Linux DMA-BUF / macOS IOSurface via the shared fd / CFRetain.
2597 // `TensorStorage::view` performs the one remaining per-variant dispatch.
2598 let mut t = Tensor::wrap(self.storage.view(offset_bytes, shape)?);
2599 // Inherit the parent's image metadata so the view is a ready-to-use
2600 // sub-image (e.g. a `convert()` destination). The offset is applied
2601 // LAST because `set_format` deliberately clears it — the offset is a
2602 // structural property of the sub-region, not format-dependent metadata.
2603 if let Some(fmt) = self.format {
2604 t.set_format(fmt)?;
2605 }
2606 if let Some(rs) = self.row_stride {
2607 t.set_row_stride_unchecked(rs);
2608 }
2609 t.quantization = self.quantization.clone();
2610 // A sub-region of an image carries the parent's colorimetry — it is the
2611 // same pixels, same color encoding. Inherit it like the other image
2612 // metadata above so a sub-view is a faithful convert() source/target.
2613 t.set_colorimetry(self.colorimetry);
2614 if abs_offset > 0 {
2615 t.set_plane_offset(abs_offset);
2616 }
2617 Ok(t)
2618 }
2619
2620 /// Borrow batch element `n` of a batched tensor as a zero-copy view.
2621 ///
2622 /// A batched tensor prepends `N` as the leading dimension over the
2623 /// per-element image layout (`[N, H, W, C]` packed or `[N, C, H, W]`
2624 /// planar) — `N` is over the whole per-element block regardless of
2625 /// `HWC`/`CHW`. `batch(n)` returns element `n`: the contiguous per-element
2626 /// region at byte offset `n * element_size`, sharing the parent's
2627 /// `BufferIdentity` and inheriting its format / row stride / colorimetry.
2628 /// `batch(0)` on a tensor with `N == 1` is equivalent to the whole tensor.
2629 ///
2630 /// # Errors
2631 ///
2632 /// - [`Error::BatchIndexOutOfBounds`] if `n >= N`.
2633 /// - [`Error::InvalidShape`] if the tensor is not batched (a formatted
2634 /// tensor whose rank lacks the leading `N`, or an empty shape).
2635 pub fn batch(&self, n: usize) -> Result<Tensor<T>> {
2636 let shape = self.shape();
2637 // With a format we know the exact per-element rank, so a missing leading
2638 // `N` is a misuse we reject rather than silently treating a spatial dim
2639 // as the batch. Raw tensors take shape[0] as `N` by contract.
2640 if let Some(fmt) = self.format {
2641 let elem_rank = match fmt.layout() {
2642 PixelLayout::SemiPlanar => 2,
2643 _ => 3,
2644 };
2645 if shape.len() != elem_rank + 1 {
2646 return Err(Error::InvalidShape(format!(
2647 "batch(): tensor is not batched ({fmt:?} expects a leading N over a \
2648 {elem_rank}-D element, got shape {shape:?})"
2649 )));
2650 }
2651 }
2652 let batch = *shape
2653 .first()
2654 .ok_or_else(|| Error::InvalidShape("batch(): empty shape".into()))?;
2655 if n >= batch {
2656 return Err(Error::BatchIndexOutOfBounds { index: n, batch });
2657 }
2658 let elem_shape: Vec<usize> = shape[1..].to_vec();
2659 let elem_count: usize = elem_shape.iter().product();
2660 let elem_bytes = elem_count
2661 .checked_mul(std::mem::size_of::<T>())
2662 .ok_or(Error::InvalidSize(elem_count))?;
2663 let offset = n.checked_mul(elem_bytes).ok_or(Error::InvalidSize(n))?;
2664 // For a packed `[N, H, W, C]` tensor the N tiles stack vertically in the
2665 // shared buffer, so the GL import sees one `(W, N*H)` parent and each
2666 // tile is the row-band at `y = n*H`. Snapshot that parent so the backend
2667 // imports once and renders the tile via `glViewport`. Non-packed
2668 // (planar/semi-planar) batching keeps the per-slot path for now (planar
2669 // NCHW tiling is a separate step); raw tensors have no pixel geometry.
2670 let view_origin = match self.format.map(|f| f.layout()) {
2671 Some(PixelLayout::Packed) => {
2672 let tile_h = elem_shape[0];
2673 let tile_w = elem_shape[1];
2674 // Per-row pitch of the tall `(W, N*H)` parent — padded stride if
2675 // set, else the tight row width. The GL import keys on this.
2676 let bpp = elem_shape[2] * std::mem::size_of::<T>();
2677 let parent_stride = self.effective_row_stride().unwrap_or(tile_w * bpp);
2678 Some(self.compose_view_origin(tile_w, batch * tile_h, parent_stride, 0, n * tile_h))
2679 }
2680 _ => None,
2681 };
2682 let mut t = self.subview(offset, &elem_shape)?;
2683 t.view_origin = view_origin;
2684 Ok(t)
2685 }
2686
2687 /// Borrow a rectangular spatial sub-region of an image tensor as a
2688 /// zero-copy view — the **destination/source crop** primitive.
2689 ///
2690 /// `region` is in pixels of the image's leading frame. The returned view
2691 /// shares the parent's `BufferIdentity` and addresses the sub-rectangle by
2692 /// offset + the **parent's** row pitch (so each row lands at the correct
2693 /// columns). `convert(src, &mut dst.view(rect), …)` renders into that
2694 /// sub-rectangle of `dst`; a letterbox fit then clears the view and renders
2695 /// the aspect-preserved content into its inner region. `view`/`batch`/the
2696 /// whole tensor are the one coherent destination model — there is no
2697 /// separate `dst_rect`.
2698 ///
2699 /// # Errors
2700 ///
2701 /// - [`Error::RegionOutOfBounds`] if `region` exceeds the image bounds.
2702 /// - [`Error::InvalidOperation`] if the tensor is not a packed-format image
2703 /// (planar/semi-planar spatial sub-rects are not a single strided window;
2704 /// use [`batch`](Self::batch) for batched planar tensors).
2705 pub fn view(&self, region: Region) -> Result<Tensor<T>> {
2706 let fmt = self.format.ok_or_else(|| {
2707 Error::InvalidOperation("view() requires a formatted image tensor".into())
2708 })?;
2709 if fmt.layout() != PixelLayout::Packed {
2710 return Err(Error::InvalidOperation(format!(
2711 "view() supports packed formats only (got {fmt:?}); use batch(n) for batched \
2712 planar tensors"
2713 )));
2714 }
2715 let w = self
2716 .width()
2717 .ok_or_else(|| Error::InvalidOperation("view(): tensor has no image width".into()))?;
2718 let h = self
2719 .height()
2720 .ok_or_else(|| Error::InvalidOperation("view(): tensor has no image height".into()))?;
2721 if !region.fits_within(w, h) {
2722 return Err(Error::RegionOutOfBounds {
2723 region,
2724 bounds: (w, h),
2725 });
2726 }
2727 let elem = std::mem::size_of::<T>();
2728 let bpp = fmt.channels() * elem;
2729 let stride = self.effective_row_stride().unwrap_or(w * bpp);
2730 let offset = region
2731 .y
2732 .checked_mul(stride)
2733 .and_then(|yo| yo.checked_add(region.x.checked_mul(bpp)?))
2734 .ok_or(Error::InvalidSize(region.y))?;
2735 let sub_shape = fmt
2736 .image_shape(region.width, region.height)
2737 .ok_or_else(|| Error::InvalidShape(format!("view(): invalid shape for {fmt:?}")))?;
2738 let mut t = self.subview(offset, &sub_shape)?;
2739 // A multi-row sub-rect must advance rows by the PARENT pitch so each row
2740 // addresses the correct columns. A single-row view uses its own tight
2741 // stride — the parent pitch would make the strided `map()` (which exposes
2742 // `stride × rows`) expose a trailing row that runs past the buffer tail
2743 // for an offset (x>0 / bottom) view. The GL backend does NOT rely on this
2744 // (single-row-tight) `row_stride`: it reads the parent pitch from
2745 // `view_origin.parent_row_stride` so its import/cache pitch stays
2746 // parent-consistent for views of any height — see `ViewOrigin`.
2747 let view_stride = if region.height > 1 {
2748 stride
2749 } else {
2750 region.width * bpp
2751 };
2752 t.set_row_stride_unchecked(view_stride);
2753 // Snapshot the parent `(w, h, row_stride)` so the GL backend imports the
2754 // parent once (keyed on the parent pitch, not this view's possibly-tight
2755 // single-row stride) and renders this sub-rect as a `glViewport`/
2756 // `glScissor` ROI at `(region.x, region.y)`. Composes when viewing an
2757 // existing view.
2758 t.view_origin = Some(self.compose_view_origin(w, h, stride, region.x, region.y));
2759 Ok(t)
2760 }
2761
2762 /// Build the [`ViewOrigin`] for a new sub-region of `self`. When `self` is a
2763 /// whole tensor the snapshot names `self` as the parent; when `self` is
2764 /// already a view, the snapshot keeps the **root** parent and accumulates the
2765 /// local origin so nested views still resolve to one import.
2766 fn compose_view_origin(
2767 &self,
2768 parent_width: usize,
2769 parent_height: usize,
2770 parent_row_stride: usize,
2771 x: usize,
2772 y: usize,
2773 ) -> ViewOrigin {
2774 match self.view_origin {
2775 Some(root) => ViewOrigin {
2776 parent_width: root.parent_width,
2777 parent_height: root.parent_height,
2778 parent_row_stride: root.parent_row_stride,
2779 x: root.x.saturating_add(x),
2780 y: root.y.saturating_add(y),
2781 },
2782 None => ViewOrigin {
2783 parent_width,
2784 parent_height,
2785 parent_row_stride,
2786 x,
2787 y,
2788 },
2789 }
2790 }
2791
2792 /// Downcast to PBO tensor reference (for GL backends).
2793 pub fn as_pbo(&self) -> Option<&PboTensor<T>> {
2794 match &self.storage {
2795 TensorStorage::Pbo(p) => Some(p),
2796 _ => None,
2797 }
2798 }
2799
2800 /// Downcast to DMA tensor reference (for EGL import, G2D).
2801 #[cfg(target_os = "linux")]
2802 pub fn as_dma(&self) -> Option<&DmaTensor<T>> {
2803 match &self.storage {
2804 TensorStorage::Dma(d) => Some(d),
2805 _ => None,
2806 }
2807 }
2808
2809 /// Borrow the DMA-BUF file descriptor backing this tensor.
2810 ///
2811 /// # Returns
2812 ///
2813 /// A borrowed reference to the DMA-BUF file descriptor, tied to `self`'s
2814 /// lifetime.
2815 ///
2816 /// # Errors
2817 ///
2818 /// Returns `Error::NotImplemented` if the tensor is not DMA-backed.
2819 #[cfg(target_os = "linux")]
2820 pub fn dmabuf(&self) -> Result<std::os::fd::BorrowedFd<'_>> {
2821 use std::os::fd::AsFd;
2822 match &self.storage {
2823 TensorStorage::Dma(dma) => Ok(dma.fd.as_fd()),
2824 _ => Err(Error::NotImplemented(format!(
2825 "dmabuf requires DMA-backed tensor, got {:?}",
2826 self.storage.memory()
2827 ))),
2828 }
2829 }
2830
2831 /// Construct a Tensor from a PBO tensor (for GL backends that allocate PBOs).
2832 pub fn from_pbo(pbo: PboTensor<T>) -> Self {
2833 Self {
2834 storage: TensorStorage::Pbo(pbo),
2835 format: None,
2836 chroma: None,
2837 row_stride: None,
2838 plane_offset: None,
2839 quantization: None,
2840 cuda: None,
2841 colorimetry: None,
2842 view_origin: None,
2843 }
2844 }
2845
2846 /// The CUDA registration for this tensor, if any (set at creation on CUDA devices).
2847 pub fn cuda(&self) -> Option<&crate::cuda::CudaHandle> {
2848 self.cuda.as_ref()
2849 }
2850
2851 /// Attach a CUDA handle (called by ImageProcessor::create_image after registering a PBO).
2852 pub fn set_cuda_handle(&mut self, h: crate::cuda::CudaHandle) {
2853 self.cuda = Some(h);
2854 }
2855
2856 /// Fast-fail CUDA map: None (no GL routing) when no handle; else map (PBO routes to the GL worker).
2857 ///
2858 /// Returns a scoped [`CudaMap`](crate::cuda::CudaMap) guard holding the raw CUDA device pointer
2859 /// for the duration of the mapping. For GL-buffer-backed tensors the unmap is deferred until the
2860 /// guard drops, freeing the PBO for the next `convert()` call. When no CUDA handle is attached
2861 /// (the common case for plain `Mem`/`DMA` tensors without CUDA registration), returns `None`
2862 /// immediately — no GL routing, no allocation.
2863 ///
2864 /// # Example — zero-copy CUDA input with host fallback
2865 ///
2866 /// ```no_run
2867 /// use edgefirst_tensor::{Tensor, TensorMemory, TensorTrait};
2868 /// # fn feed_tensorrt(_dptr: *mut std::ffi::c_void, _bytes: usize) {}
2869 /// # fn demo(t: &Tensor<f32>) {
2870 /// // Try the zero-copy CUDA device pointer first.
2871 /// if let Some(cuda) = t.cuda_map() {
2872 /// feed_tensorrt(cuda.device_ptr(), cuda.len());
2873 /// // `cuda` (a CudaMap guard) unmaps when it goes out of scope, freeing
2874 /// // the GPU buffer for the next convert().
2875 /// } else {
2876 /// // Fall back to the host mapping when no CUDA handle is attached.
2877 /// let _host = t.map().expect("host map fallback must succeed");
2878 /// // `_host` is a TensorMap<f32> that derefs to &[f32].
2879 /// }
2880 /// # }
2881 /// ```
2882 pub fn cuda_map(&self) -> Option<crate::cuda::CudaMap<'_>> {
2883 self.cuda.as_ref()?.map()
2884 }
2885
2886 /// Attempt to attach a CUDA `ExternalMemory` handle for DMA-backed tensors.
2887 ///
2888 /// On a CUDA-capable host, imports the DMA-BUF fd via
2889 /// `cudaImportExternalMemory(OpaqueFd)` and maps it to a device pointer.
2890 /// Sets `self.cuda` to a persistent `ExternalMem` handle on success. No-op
2891 /// if CUDA is unavailable, the tensor is not DMA-backed, or a handle is
2892 /// already set. Import failure is silently ignored — the tensor remains
2893 /// usable without a CUDA handle.
2894 ///
2895 /// # RUNTIME-UNVALIDATED
2896 ///
2897 /// No test platform has both `/dev/dma_heap` and a CUDA device. ABI is
2898 /// layout-asserted vs. CUDA 12.6 `driver_types.h`; the mechanism is proven
2899 /// by gpu-probe O5 on Orin. Best-effort: tensor creation never fails here.
2900 #[cfg(target_os = "linux")]
2901 pub fn try_init_dma_cuda(&mut self) {
2902 // Fast-path: already imported, CUDA not available, or not a DMA tensor.
2903 if self.cuda.is_some() || !crate::cuda::is_cuda_available() {
2904 return;
2905 }
2906 let (raw_fd, buf_size) = match &self.storage {
2907 TensorStorage::Dma(dma) => {
2908 use std::os::fd::AsRawFd;
2909 (dma.fd.as_raw_fd(), dma.buf_size)
2910 }
2911 _ => return,
2912 };
2913 if let Some((ext, dptr)) = crate::cuda::import_dma_fd(raw_fd, buf_size) {
2914 self.cuda = Some(crate::cuda::CudaHandle::new_external(ext, dptr, buf_size));
2915 }
2916 }
2917}
2918
2919// Quantization accessors — type-gated to integer element types via the
2920// sealed `IntegerType` trait. Calling `.quantization()` on a `Tensor<f32>`
2921// produces a compile error, not a runtime one.
2922impl<T> Tensor<T>
2923where
2924 T: IntegerType + Num + Clone + fmt::Debug + Send + Sync,
2925{
2926 /// Quantization metadata for this tensor, if set.
2927 pub fn quantization(&self) -> Option<&Quantization> {
2928 self.quantization.as_ref()
2929 }
2930
2931 /// Attach quantization metadata to this tensor. Validates against the
2932 /// tensor's shape — returns [`Error::QuantizationInvalid`] on any
2933 /// inconsistency (mismatched scale/zp lengths, out-of-range axis, etc.).
2934 pub fn set_quantization(&mut self, q: Quantization) -> Result<()> {
2935 q.validate(self.shape())?;
2936 self.quantization = Some(q);
2937 Ok(())
2938 }
2939
2940 /// Builder-style variant of [`Self::set_quantization`]. Consumes `self`
2941 /// and returns `Result<Self>` — on success yields the tensor with the
2942 /// attached quantization; on validation failure returns
2943 /// [`Error::QuantizationInvalid`] and drops `self` (the tensor is not
2944 /// returned in the error arm).
2945 pub fn with_quantization(mut self, q: Quantization) -> Result<Self> {
2946 self.set_quantization(q)?;
2947 Ok(self)
2948 }
2949
2950 /// Clear any quantization metadata on this tensor.
2951 pub fn clear_quantization(&mut self) {
2952 self.quantization = None;
2953 }
2954}
2955
2956impl<T> TensorTrait<T> for Tensor<T>
2957where
2958 T: Num + Clone + fmt::Debug + Send + Sync,
2959{
2960 fn new(shape: &[usize], name: Option<&str>) -> Result<Self>
2961 where
2962 Self: Sized,
2963 {
2964 Self::new(shape, None, name)
2965 }
2966
2967 #[cfg(unix)]
2968 fn from_fd(fd: std::os::fd::OwnedFd, shape: &[usize], name: Option<&str>) -> Result<Self>
2969 where
2970 Self: Sized,
2971 {
2972 #[cfg_attr(not(target_os = "linux"), allow(unused_mut))]
2973 let mut t = Self::wrap(TensorStorage::from_fd(fd, shape, name)?);
2974 // Best-effort CUDA external memory import for DMA-backed tensors.
2975 // RUNTIME-UNVALIDATED: see try_init_dma_cuda().
2976 #[cfg(target_os = "linux")]
2977 t.try_init_dma_cuda();
2978 Ok(t)
2979 }
2980
2981 #[cfg(unix)]
2982 fn clone_fd(&self) -> Result<std::os::fd::OwnedFd> {
2983 self.storage.clone_fd()
2984 }
2985
2986 fn memory(&self) -> TensorMemory {
2987 self.storage.memory()
2988 }
2989
2990 fn name(&self) -> String {
2991 self.storage.name()
2992 }
2993
2994 fn shape(&self) -> &[usize] {
2995 self.storage.shape()
2996 }
2997
2998 fn reshape(&mut self, shape: &[usize]) -> Result<()> {
2999 if self.chroma.is_some() {
3000 return Err(Error::InvalidOperation(
3001 "cannot reshape a multiplane tensor — decompose planes first".into(),
3002 ));
3003 }
3004 self.storage.reshape(shape)?;
3005 self.format = None;
3006 self.row_stride = None;
3007 self.plane_offset = None;
3008 match self.storage {
3009 TensorStorage::Mem(ref mut m) => m.set_offset(0),
3010 #[cfg(target_os = "linux")]
3011 TensorStorage::Dma(ref mut dma) => dma.mmap_offset = 0,
3012 _ => {}
3013 }
3014 Ok(())
3015 }
3016
3017 fn map(&self) -> Result<TensorMap<T>> {
3018 let _span = tracing::trace_span!(
3019 "tensor.map",
3020 memory = ?self.storage.memory(),
3021 )
3022 .entered();
3023 // CPU mapping of a strided tensor exposes the full padded buffer
3024 // (`row_stride × rows`) so callers can iterate rows via
3025 // `effective_row_stride()` without running past the slice. This is sound
3026 // only when the HAL owns and can size-check the allocation:
3027 //
3028 // * Self-allocated Mem / Shm tensors (any platform) — the backing
3029 // `Vec` / shm segment is sized by `capacity_bytes()`, checked here.
3030 // * Self-allocated DMA tensors (Linux) — pitch padding from
3031 // `image_with_stride()`; checked against the DMA-BUF `buf_size`.
3032 //
3033 // * Self-allocated PBO tensors (any platform with GL) — the GL buffer
3034 // is sized by `capacity_bytes()` and may carry 64-byte row padding;
3035 // the JPEG decoder mmaps it and convert() reads it, both iterating
3036 // by `row_stride`. Checked against the PBO capacity below.
3037 //
3038 // Foreign DMA-BUFs (`from_fd()` + `set_row_stride()`, the V4L2 /
3039 // GStreamer case) and IOSurface are rejected: their layout comes from an
3040 // external allocator / GPU driver the HAL cannot validate for a strided
3041 // CPU view, and they are intended for the GPU path. (Earlier this
3042 // rejected *all* non-Linux strided maps with "DMA backing is Linux-only"
3043 // — that was an unimplemented path, not a platform limit; HAL-owned
3044 // Mem/Shm/PBO are trivially mappable and now are.)
3045 if let Some(stride) = self.row_stride {
3046 // Rows sit at `stride`-byte spacing; the first shape dim is the row
3047 // count for packed `[H, W, C]` and semi-planar `[H*k, W]` alike.
3048 let rows = *self.shape().first().ok_or_else(|| {
3049 Error::InvalidOperation(
3050 "Tensor::map: strided mapping requires a non-empty shape".into(),
3051 )
3052 })?;
3053 let total_bytes = stride.checked_mul(rows).ok_or_else(|| {
3054 Error::InvalidOperation(format!(
3055 "Tensor::map: row_stride {stride} × rows {rows} overflows usize"
3056 ))
3057 })?;
3058
3059 match &self.storage {
3060 #[cfg(target_os = "linux")]
3061 TensorStorage::Dma(dma) if !dma.is_imported => {
3062 // `set_row_stride()` only validates `stride >= min_stride`,
3063 // not that `stride × rows` fits the DMA-BUF, so re-check
3064 // here — mapping past `buf_size` would SIGBUS on access.
3065 let available_bytes = dma.buf_size.saturating_sub(dma.mmap_offset);
3066 if total_bytes > available_bytes {
3067 return Err(Error::InvalidOperation(format!(
3068 "Tensor::map: strided mapping needs {total_bytes} bytes \
3069 but DMA buffer only has {available_bytes} available \
3070 (buf_size={}, mmap_offset={}, stride={stride}, rows={rows}); \
3071 the row_stride was likely set larger than the original allocation",
3072 dma.buf_size, dma.mmap_offset
3073 )));
3074 }
3075 return dma.map_with_byte_size(total_bytes).map(TensorMap::Dma);
3076 }
3077 TensorStorage::Mem(mem) => {
3078 let capacity = self.storage.capacity_bytes();
3079 if total_bytes > capacity {
3080 return Err(Error::InsufficientCapacity {
3081 needed: total_bytes,
3082 capacity,
3083 });
3084 }
3085 return mem.map_with_byte_size(total_bytes);
3086 }
3087 #[cfg(unix)]
3088 TensorStorage::Shm(shm) => {
3089 let capacity = self.storage.capacity_bytes();
3090 if total_bytes > capacity {
3091 return Err(Error::InsufficientCapacity {
3092 needed: total_bytes,
3093 capacity,
3094 });
3095 }
3096 return shm.map_with_byte_size(total_bytes);
3097 }
3098 // macOS: `TensorStorage::Dma` is the IOSurface. The lock yields
3099 // the full surface base address, and the row pitch
3100 // (`IOSurfaceGetBytesPerRow`) is known from the API for both
3101 // self-allocated and imported surfaces — unlike a foreign
3102 // DMA-BUF — so a strided CPU view is sound and zero-copy.
3103 #[cfg(target_os = "macos")]
3104 TensorStorage::Dma(io) => {
3105 // A sub-view's window is `buf_size − view_offset`; the strided
3106 // span must fit the window, not the whole surface.
3107 let available = io.buf_size.saturating_sub(io.view_offset);
3108 if total_bytes > available {
3109 return Err(Error::InsufficientCapacity {
3110 needed: total_bytes,
3111 capacity: available,
3112 });
3113 }
3114 return io.map_with_byte_size(total_bytes);
3115 }
3116 TensorStorage::Pbo(pbo) => {
3117 // PBO: the GPU-side allocation may have a padded row stride
3118 // (e.g. 64-byte aligned). Expose the full padded buffer so a
3119 // CPU producer (JPEG decoder) and a strided convert source
3120 // can iterate rows via `effective_row_stride()` without
3121 // running past the slice — the logical `pbo.map()` view would
3122 // stop after `shape.product()` and lose bytes past row 0.
3123 // A sub-view's window is `capacity − view_offset`.
3124 let available = pbo.capacity_bytes().saturating_sub(pbo.view_offset);
3125 if total_bytes > available {
3126 return Err(Error::InsufficientCapacity {
3127 needed: total_bytes,
3128 capacity: available,
3129 });
3130 }
3131 return pbo.map_with_byte_size(total_bytes);
3132 }
3133 // Reachable on Linux for an IMPORTED DMA-BUF (the `Dma` arm above
3134 // is guarded `if !dma.is_imported`). On macOS/Windows every
3135 // storage variant is matched explicitly, so this catch-all is
3136 // unreachable there — allow it rather than cfg-gating per platform.
3137 #[allow(unreachable_patterns)]
3138 _ => {
3139 return Err(Error::InvalidOperation(
3140 "CPU mapping of strided tensors is supported only for HAL-allocated \
3141 Mem/Shm (any platform), self-allocated DMA (Linux), IOSurface \
3142 (macOS), and PBO; imported DMA-BUF without self-allocation is \
3143 GPU-path only"
3144 .into(),
3145 ));
3146 }
3147 }
3148 }
3149 // Offset tensors are supported for storages that apply the offset
3150 // inside their own `map()`: DMA (`DmaMap`/IOSurface adjust the mapped
3151 // base), Mem (`MemMap` adjusts the slice base), Shm (`ShmMap` adjusts
3152 // the slice base), and PBO (the staged copy starts at the offset). Every
3153 // self-allocated backing now carries a sub-region concept via `view`, so
3154 // a non-zero offset is honoured rather than rejected.
3155 if self.plane_offset.is_some_and(|o| o > 0) {
3156 let supported = matches!(self.storage, TensorStorage::Mem(_) | TensorStorage::Pbo(_));
3157 // macOS `Dma` is the IOSurface; Linux `Dma` is the DMA-BUF — both
3158 // apply the offset in their map. (`Dma` is the same variant name on
3159 // both, hence one `cfg(any(...))` arm rather than two.)
3160 #[cfg(any(target_os = "linux", target_os = "macos"))]
3161 let supported = supported || matches!(self.storage, TensorStorage::Dma(_));
3162 #[cfg(unix)]
3163 let supported = supported || matches!(self.storage, TensorStorage::Shm(_));
3164 if !supported {
3165 return Err(Error::InvalidOperation(
3166 "plane offset only supported for DMA, Mem, Shm, and PBO tensors".into(),
3167 ));
3168 }
3169 }
3170 self.storage.map()
3171 }
3172
3173 fn buffer_identity(&self) -> &BufferIdentity {
3174 self.storage.buffer_identity()
3175 }
3176}
3177
3178pub enum TensorMap<T>
3179where
3180 T: Num + Clone + fmt::Debug,
3181{
3182 #[cfg(target_os = "linux")]
3183 Dma(DmaMap<T>),
3184 #[cfg(target_os = "macos")]
3185 IoSurface(IoSurfaceMap<T>),
3186 #[cfg(unix)]
3187 Shm(ShmMap<T>),
3188 Mem(MemMap<T>),
3189 Pbo(PboMap<T>),
3190}
3191
3192impl<T> TensorMapTrait<T> for TensorMap<T>
3193where
3194 T: Num + Clone + fmt::Debug,
3195{
3196 fn shape(&self) -> &[usize] {
3197 match self {
3198 #[cfg(target_os = "linux")]
3199 TensorMap::Dma(map) => map.shape(),
3200 #[cfg(target_os = "macos")]
3201 TensorMap::IoSurface(map) => map.shape(),
3202 #[cfg(unix)]
3203 TensorMap::Shm(map) => map.shape(),
3204 TensorMap::Mem(map) => map.shape(),
3205 TensorMap::Pbo(map) => map.shape(),
3206 }
3207 }
3208
3209 fn unmap(&mut self) {
3210 match self {
3211 #[cfg(target_os = "linux")]
3212 TensorMap::Dma(map) => map.unmap(),
3213 #[cfg(target_os = "macos")]
3214 TensorMap::IoSurface(map) => map.unmap(),
3215 #[cfg(unix)]
3216 TensorMap::Shm(map) => map.unmap(),
3217 TensorMap::Mem(map) => map.unmap(),
3218 TensorMap::Pbo(map) => map.unmap(),
3219 }
3220 }
3221
3222 fn as_slice(&self) -> &[T] {
3223 match self {
3224 #[cfg(target_os = "linux")]
3225 TensorMap::Dma(map) => map.as_slice(),
3226 #[cfg(target_os = "macos")]
3227 TensorMap::IoSurface(map) => map.deref(),
3228 #[cfg(unix)]
3229 TensorMap::Shm(map) => map.as_slice(),
3230 TensorMap::Mem(map) => map.as_slice(),
3231 TensorMap::Pbo(map) => map.as_slice(),
3232 }
3233 }
3234
3235 fn as_mut_slice(&mut self) -> &mut [T] {
3236 match self {
3237 #[cfg(target_os = "linux")]
3238 TensorMap::Dma(map) => map.as_mut_slice(),
3239 #[cfg(target_os = "macos")]
3240 TensorMap::IoSurface(map) => map.deref_mut(),
3241 #[cfg(unix)]
3242 TensorMap::Shm(map) => map.as_mut_slice(),
3243 TensorMap::Mem(map) => map.as_mut_slice(),
3244 TensorMap::Pbo(map) => map.as_mut_slice(),
3245 }
3246 }
3247}
3248
3249impl<T> Deref for TensorMap<T>
3250where
3251 T: Num + Clone + fmt::Debug,
3252{
3253 type Target = [T];
3254
3255 fn deref(&self) -> &[T] {
3256 match self {
3257 #[cfg(target_os = "linux")]
3258 TensorMap::Dma(map) => map.deref(),
3259 #[cfg(target_os = "macos")]
3260 TensorMap::IoSurface(map) => map.deref(),
3261 #[cfg(unix)]
3262 TensorMap::Shm(map) => map.deref(),
3263 TensorMap::Mem(map) => map.deref(),
3264 TensorMap::Pbo(map) => map.deref(),
3265 }
3266 }
3267}
3268
3269impl<T> DerefMut for TensorMap<T>
3270where
3271 T: Num + Clone + fmt::Debug,
3272{
3273 fn deref_mut(&mut self) -> &mut [T] {
3274 match self {
3275 #[cfg(target_os = "linux")]
3276 TensorMap::Dma(map) => map.deref_mut(),
3277 #[cfg(target_os = "macos")]
3278 TensorMap::IoSurface(map) => map.deref_mut(),
3279 #[cfg(unix)]
3280 TensorMap::Shm(map) => map.deref_mut(),
3281 TensorMap::Mem(map) => map.deref_mut(),
3282 TensorMap::Pbo(map) => map.deref_mut(),
3283 }
3284 }
3285}
3286
3287// ============================================================================
3288// Platform availability helpers
3289// ============================================================================
3290
3291/// Cached result of the Linux DMA-BUF availability probe.
3292#[cfg(target_os = "linux")]
3293static DMA_AVAILABLE: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
3294/// Cached result of the macOS IOSurface availability probe.
3295#[cfg(target_os = "macos")]
3296static IOSURFACE_AVAILABLE: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
3297
3298/// Check if Linux DMA-BUF allocation is available on this system.
3299///
3300/// Linux-specific availability check (typically requires `/dev/dma_heap`
3301/// access — running as root or membership in a video/render group). For
3302/// portable code that wants "any zero-copy GPU buffer", use
3303/// [`is_gpu_buffer_available`] which also covers IOSurface on macOS.
3304///
3305/// This function caches its result after the first call.
3306#[cfg(target_os = "linux")]
3307pub fn is_dma_available() -> bool {
3308 *DMA_AVAILABLE.get_or_init(|| Tensor::<u8>::new(&[64], Some(TensorMemory::Dma), None).is_ok())
3309}
3310
3311/// Always returns `false` on non-Linux platforms.
3312#[cfg(not(target_os = "linux"))]
3313pub fn is_dma_available() -> bool {
3314 false
3315}
3316
3317/// Check if macOS IOSurface allocation is available on this system.
3318///
3319/// IOSurface is part of the macOS OS and is essentially always present;
3320/// this probe catches degraded scenarios such as memory pressure or
3321/// sandboxed contexts where `IOSurfaceCreate` fails. The result is
3322/// cached after the first call.
3323///
3324/// Always returns `false` on non-macOS platforms.
3325#[cfg(target_os = "macos")]
3326pub fn is_iosurface_available() -> bool {
3327 *IOSURFACE_AVAILABLE.get_or_init(|| {
3328 // Probe via the same Dma path — on macOS this routes through
3329 // IoSurfaceTensor::new.
3330 Tensor::<u8>::new(&[64], Some(TensorMemory::Dma), None).is_ok()
3331 })
3332}
3333
3334#[cfg(not(target_os = "macos"))]
3335pub fn is_iosurface_available() -> bool {
3336 false
3337}
3338
3339/// Portable probe for the platform's native zero-copy GPU buffer
3340/// allocator (DMA-BUF on Linux, IOSurface on macOS). Returns `false` on
3341/// Windows and other platforms with no equivalent. Use this when writing
3342/// cross-platform code that cares whether the `Dma` tensor variant will
3343/// work, not which underlying mechanism is used.
3344pub fn is_gpu_buffer_available() -> bool {
3345 #[cfg(target_os = "linux")]
3346 {
3347 is_dma_available()
3348 }
3349 #[cfg(target_os = "macos")]
3350 {
3351 is_iosurface_available()
3352 }
3353 #[cfg(not(any(target_os = "linux", target_os = "macos")))]
3354 {
3355 false
3356 }
3357}
3358
3359/// Check if POSIX shared memory allocation is available on this system.
3360///
3361/// Returns `true` on Unix systems (Linux, macOS, BSD) where POSIX shared memory
3362/// is supported. Always returns `false` on non-Unix platforms (Windows).
3363///
3364/// This function caches its result after the first call for efficiency.
3365#[cfg(unix)]
3366static SHM_AVAILABLE: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
3367
3368/// Check if POSIX shared memory allocation is available on this system.
3369#[cfg(unix)]
3370pub fn is_shm_available() -> bool {
3371 *SHM_AVAILABLE.get_or_init(|| Tensor::<u8>::new(&[64], Some(TensorMemory::Shm), None).is_ok())
3372}
3373
3374/// Check if POSIX shared memory allocation is available on this system.
3375///
3376/// Always returns `false` on non-Unix platforms since POSIX SHM is Unix-specific.
3377#[cfg(not(unix))]
3378pub fn is_shm_available() -> bool {
3379 false
3380}
3381
3382#[cfg(test)]
3383mod dtype_tests {
3384 use super::*;
3385
3386 #[test]
3387 fn dtype_size() {
3388 assert_eq!(DType::U8.size(), 1);
3389 assert_eq!(DType::I8.size(), 1);
3390 assert_eq!(DType::U16.size(), 2);
3391 assert_eq!(DType::I16.size(), 2);
3392 assert_eq!(DType::U32.size(), 4);
3393 assert_eq!(DType::I32.size(), 4);
3394 assert_eq!(DType::U64.size(), 8);
3395 assert_eq!(DType::I64.size(), 8);
3396 assert_eq!(DType::F16.size(), 2);
3397 assert_eq!(DType::F32.size(), 4);
3398 assert_eq!(DType::F64.size(), 8);
3399 }
3400
3401 #[test]
3402 fn dtype_name() {
3403 assert_eq!(DType::U8.name(), "u8");
3404 assert_eq!(DType::F16.name(), "f16");
3405 assert_eq!(DType::F32.name(), "f32");
3406 }
3407
3408 #[test]
3409 fn dtype_serde_roundtrip() {
3410 use serde_json;
3411 let dt = DType::F16;
3412 let json = serde_json::to_string(&dt).unwrap();
3413 let back: DType = serde_json::from_str(&json).unwrap();
3414 assert_eq!(dt, back);
3415 }
3416}
3417
3418#[cfg(test)]
3419mod image_tests {
3420 use super::*;
3421
3422 #[test]
3423 fn image_shape_per_layout() {
3424 assert_eq!(
3425 PixelFormat::Rgb.image_shape(640, 480),
3426 Some(vec![480, 640, 3])
3427 );
3428 assert_eq!(
3429 PixelFormat::Grey.image_shape(640, 480),
3430 Some(vec![480, 640, 1])
3431 );
3432 assert_eq!(
3433 PixelFormat::Nv12.image_shape(640, 480),
3434 Some(vec![720, 640])
3435 );
3436 // Odd height: combined-plane height is `481 + ceil(481/2)` = 481 + 241
3437 // = 722 rows. Logical height is recovered as `722 * 2 / 3` = 481.
3438 assert_eq!(
3439 PixelFormat::Nv12.image_shape(640, 481),
3440 Some(vec![722, 640])
3441 );
3442 // Odd width: shape carries the LOGICAL width (641).
3443 // The 64-aligned stride (>= 642) is stored separately on the Tensor.
3444 assert_eq!(
3445 PixelFormat::Nv12.image_shape(641, 480),
3446 Some(vec![720, 641])
3447 );
3448 // NV16 odd width: same — logical width in shape, stride separate.
3449 assert_eq!(
3450 PixelFormat::Nv16.image_shape(641, 480),
3451 Some(vec![960, 641])
3452 );
3453 assert_eq!(
3454 PixelFormat::PlanarRgb.image_shape(640, 480),
3455 Some(vec![3, 480, 640])
3456 );
3457 assert_eq!(
3458 PixelFormat::Nv16.image_shape(640, 480),
3459 Some(vec![960, 640])
3460 );
3461 }
3462
3463 #[test]
3464 fn raw_tensor_has_no_format() {
3465 let t = Tensor::<u8>::new(&[480, 640, 3], None, None).unwrap();
3466 assert!(t.format().is_none());
3467 assert!(t.width().is_none());
3468 assert!(t.height().is_none());
3469 assert!(!t.is_multiplane());
3470 assert!(t.chroma().is_none());
3471 }
3472
3473 #[test]
3474 fn image_tensor_packed() {
3475 let t = Tensor::<u8>::image(640, 480, PixelFormat::Rgba, None).unwrap();
3476 assert_eq!(t.format(), Some(PixelFormat::Rgba));
3477 assert_eq!(t.width(), Some(640));
3478 assert_eq!(t.height(), Some(480));
3479 assert_eq!(t.shape(), &[480, 640, 4]);
3480 assert!(!t.is_multiplane());
3481 }
3482
3483 #[test]
3484 fn image_tensor_planar() {
3485 let t = Tensor::<u8>::image(640, 480, PixelFormat::PlanarRgb, None).unwrap();
3486 assert_eq!(t.format(), Some(PixelFormat::PlanarRgb));
3487 assert_eq!(t.width(), Some(640));
3488 assert_eq!(t.height(), Some(480));
3489 assert_eq!(t.shape(), &[3, 480, 640]);
3490 }
3491
3492 #[test]
3493 #[cfg(target_os = "macos")]
3494 fn image_tensor_dma_non_aligned_packed_width_pads_zero_copy() {
3495 // RGBA u8 at width=4 → 4*4 = 16 bytes/row, not 64-byte aligned. RGBA has
3496 // a real IOSurface FourCC, so an explicit `Some(TensorMemory::Dma)`
3497 // request now allocates a padded image IOSurface (64-aligned
3498 // `bytes_per_row`) and records the stride — a fully zero-copy buffer GL
3499 // can bind and the CPU can map via the strided path. (Previously this
3500 // failed loudly to avoid an 'L008' byte-bag downgrade; with a real
3501 // FourCC surface that concern no longer applies.)
3502 let t = Tensor::<u8>::image(4, 4, PixelFormat::Rgba, Some(TensorMemory::Dma))
3503 .expect("padded RGBA IOSurface should allocate");
3504 assert_eq!(t.format(), Some(PixelFormat::Rgba));
3505 assert_eq!(t.width(), Some(4));
3506 assert_eq!(t.height(), Some(4));
3507 let stride = t.effective_row_stride().expect("stride");
3508 assert_eq!(stride % 64, 0, "padded to 64-byte row alignment");
3509 assert!(stride >= 16);
3510 // A CPU map exposes the full padded surface for strided iteration.
3511 let m = t.map().expect("strided IOSurface map");
3512 assert_eq!(m.as_slice().len(), stride * 4);
3513 }
3514
3515 /// `per_pixel_bytes` that doesn't divide 64 evenly (e.g. RGB u8 with
3516 /// 3 B/pixel) makes a "Pad width to N" suggestion structurally
3517 /// impossible — there is no integer width whose `width * 3` is a
3518 /// multiple of 64. The error must still fire (no silent SHM
3519 /// fallback for explicit-DMA requests) and must spell out the
3520 /// alignment requirement; it just omits the misleading "pad to N"
3521 /// hint instead of printing a number whose row pitch still won't
3522 /// align.
3523 #[test]
3524 #[cfg(target_os = "macos")]
3525 fn image_tensor_dma_rejects_indivisible_pixel_pitch_without_pad_hint() {
3526 // Width=10 RGB u8 → 30 B/row, not 64-byte aligned. The next
3527 // 64-multiple (64 B) isn't an integer multiple of 3 B/pixel,
3528 // so the "pad width to N" hint can't produce a valid number
3529 // and must be omitted. (Width=640 happens to align — 640*3 =
3530 // 1920 = 30*64 — so don't pick that for this regression
3531 // guard.)
3532 let err = Tensor::<u8>::image(10, 10, PixelFormat::Rgb, Some(TensorMemory::Dma))
3533 .expect_err("RGB u8 with 3 B/pixel and non-aligned width must be rejected");
3534 match err {
3535 Error::InvalidArgument(msg) => {
3536 assert!(
3537 msg.contains("64-byte aligned"),
3538 "error must still name the alignment requirement: {msg}"
3539 );
3540 assert!(
3541 !msg.contains("Pad width"),
3542 "indivisible per-pixel pitch makes a width suggestion impossible; \
3543 hint must be omitted, got: {msg}"
3544 );
3545 assert!(
3546 msg.contains("memory=None") && msg.contains("TensorMemory::Mem"),
3547 "error must still list the always-applicable alternatives: {msg}"
3548 );
3549 }
3550 other => panic!("expected InvalidArgument, got {other:?}"),
3551 }
3552 }
3553
3554 #[test]
3555 #[cfg(target_os = "macos")]
3556 fn image_tensor_dma_planar_f16_alignment() {
3557 // PlanarRgb F16 uses single-channel row pitch (width * 2 bytes).
3558 // Width=16 → 32 bytes/row (not aligned); width=32 → 64 bytes/row (aligned).
3559 let err =
3560 Tensor::<half::f16>::image(16, 16, PixelFormat::PlanarRgb, Some(TensorMemory::Dma))
3561 .expect_err("width=16 PlanarRgb F16 is 32-byte row, must reject");
3562 assert!(matches!(err, Error::InvalidArgument(_)), "got {err:?}");
3563 // 32 wide should work.
3564 let t = Tensor::<half::f16>::image(32, 8, PixelFormat::PlanarRgb, Some(TensorMemory::Dma))
3565 .expect("width=32 PlanarRgb F16 is 64-byte row, must succeed");
3566 assert_eq!(t.format(), Some(PixelFormat::PlanarRgb));
3567 }
3568
3569 #[test]
3570 fn image_tensor_semi_planar_contiguous() {
3571 let t = Tensor::<u8>::image(640, 480, PixelFormat::Nv12, None).unwrap();
3572 assert_eq!(t.format(), Some(PixelFormat::Nv12));
3573 assert_eq!(t.width(), Some(640));
3574 assert_eq!(t.height(), Some(480));
3575 // NV12: H*3/2 = 720
3576 assert_eq!(t.shape(), &[720, 640]);
3577 assert!(!t.is_multiplane());
3578 }
3579
3580 #[test]
3581 #[cfg(target_os = "linux")]
3582 fn image_tensor_with_stride_preserves_logical_width() {
3583 // Skip if DMA not available (e.g. sandboxed CI lacking dma_heap access).
3584 if !is_dma_available() {
3585 eprintln!("SKIPPED: DMA heap not available");
3586 return;
3587 }
3588 // 3004×1688 RGBA8: natural pitch 12016, padded to 12032 (64-aligned).
3589 let stride = 12032;
3590 let t = Tensor::<u8>::image_with_stride(
3591 3004,
3592 1688,
3593 PixelFormat::Rgba,
3594 stride,
3595 Some(TensorMemory::Dma),
3596 )
3597 .unwrap();
3598 // Logical dimensions unchanged by padding — this is the contract.
3599 assert_eq!(t.width(), Some(3004));
3600 assert_eq!(t.height(), Some(1688));
3601 assert_eq!(t.shape(), &[1688, 3004, 4]);
3602 // Stride is carried separately and reports the padded pitch.
3603 assert_eq!(t.effective_row_stride(), Some(stride));
3604 // Buffer is sized to stride × height so the full padded layout fits,
3605 // and CPU map() works for self-allocated strided DMA tensors.
3606 use crate::TensorMapTrait;
3607 {
3608 let map = t.map().unwrap();
3609 assert!(
3610 map.as_slice().len() >= stride * 1688,
3611 "mapped buffer {} bytes < expected {}",
3612 map.as_slice().len(),
3613 stride * 1688
3614 );
3615 }
3616 // CPU write access works too — iterate rows using the padded stride,
3617 // touch only the active `width × bpp` region, verify it round-trips.
3618 {
3619 let mut map = t.map().unwrap();
3620 let slice = map.as_mut_slice();
3621 for y in 0..1688 {
3622 let row_start = y * stride;
3623 for x in 0..3004 {
3624 let p = row_start + x * 4;
3625 slice[p] = (y & 0xFF) as u8;
3626 slice[p + 1] = (x & 0xFF) as u8;
3627 slice[p + 2] = 0x42;
3628 slice[p + 3] = 0xFF;
3629 }
3630 }
3631 }
3632 {
3633 let map = t.map().unwrap();
3634 let slice = map.as_slice();
3635 // Sample a few pixels to confirm the round-trip.
3636 assert_eq!(slice[0], 0x00);
3637 assert_eq!(slice[1], 0x00);
3638 assert_eq!(slice[2], 0x42);
3639 assert_eq!(slice[3], 0xFF);
3640 let mid = 100 * stride + 50 * 4;
3641 assert_eq!(slice[mid], 100);
3642 assert_eq!(slice[mid + 1], 50);
3643 assert_eq!(slice[mid + 2], 0x42);
3644 }
3645 }
3646
3647 #[test]
3648 #[cfg(target_os = "linux")]
3649 fn image_tensor_with_stride_rejects_foreign_strided_map() {
3650 // A FOREIGN (imported via from_fd) DMA tensor with row_stride set
3651 // should still refuse CPU mapping — external allocator owns the
3652 // layout. This protects the V4L2 / GStreamer use case.
3653 //
3654 // We simulate a foreign import by wrapping our own allocation's
3655 // fd via `from_fd` and calling set_row_stride manually. The
3656 // `is_imported` flag on from_fd is true by construction.
3657 if !is_dma_available() {
3658 eprintln!("SKIPPED: DMA heap not available");
3659 return;
3660 }
3661 // Allocate a backing buffer large enough for a 320×240 BGRA8 image.
3662 let backing = Tensor::<u8>::new(&[240 * 320 * 4], Some(TensorMemory::Dma), None).unwrap();
3663 let fd = backing.clone_fd().unwrap();
3664 // Import it via from_fd — this marks is_imported=true.
3665 let shape = [240usize, 320, 4];
3666 let storage = TensorStorage::<u8>::from_fd(fd, &shape, None).unwrap();
3667 let mut t = Tensor::<u8>::wrap(storage);
3668 t.set_format(PixelFormat::Bgra).unwrap();
3669 t.set_row_stride(320 * 4).unwrap(); // natural, but still marks it as strided
3670 let err = t.map();
3671 assert!(
3672 matches!(err, Err(Error::InvalidOperation(_))),
3673 "foreign strided map should error"
3674 );
3675 }
3676
3677 #[test]
3678 #[cfg(target_os = "linux")]
3679 fn image_tensor_with_stride_map_rejects_tampered_stride() {
3680 // Round-3 PR feedback (C1): `set_row_stride` is public and only
3681 // validates `stride >= min_stride`, not that the new stride × height
3682 // fits the underlying buffer. A caller that tampers with the stride
3683 // after allocation must not be able to coerce `Tensor::map()` into
3684 // returning a slice larger than the backing mmap (that would be UB
3685 // in `DmaMap::as_slice`).
3686 if !is_dma_available() {
3687 eprintln!("SKIPPED: DMA heap not available");
3688 return;
3689 }
3690 // Allocate a 640×480 RGBA8 padded canvas (stride = 3072 = 768 px).
3691 // Backing buffer is 3072 × 480 = 1,474,560 bytes.
3692 let mut t = Tensor::<u8>::image_with_stride(
3693 640,
3694 480,
3695 PixelFormat::Rgba,
3696 3072,
3697 Some(TensorMemory::Dma),
3698 )
3699 .unwrap();
3700 // Tamper: push the stride up to 4 × the original. This is >=
3701 // min_stride (2560), so `set_row_stride` accepts it.
3702 t.set_row_stride(12288).unwrap();
3703 // Map must now refuse — 12288 × 480 = 5,898,240 > 1,474,560.
3704 let err = t.map();
3705 assert!(
3706 matches!(err, Err(Error::InvalidOperation(_))),
3707 "map() with oversized stride must return InvalidOperation"
3708 );
3709 }
3710
3711 #[test]
3712 fn dma_tensor_new_with_byte_size_rejects_shape_overflow() {
3713 // Round-3 PR feedback (C3): shape.product() * sizeof(T) must use
3714 // checked arithmetic so a pathological shape can't wrap usize and
3715 // make the byte_size-vs-logical-size comparison incorrect.
3716 //
3717 // This test only exercises the overflow rejection path, which is
3718 // pure-Rust and doesn't touch dma_heap — safe to run on any target.
3719 #[cfg(target_os = "linux")]
3720 {
3721 let err = crate::dma::DmaTensor::<u64>::new_with_byte_size(
3722 &[usize::MAX, 2, 2],
3723 usize::MAX,
3724 None,
3725 );
3726 assert!(
3727 matches!(err, Err(Error::InvalidArgument(_))),
3728 "new_with_byte_size must detect shape.product() overflow"
3729 );
3730 }
3731 }
3732
3733 #[test]
3734 #[cfg(target_os = "linux")]
3735 fn image_tensor_with_stride_rejects_too_small_stride() {
3736 // 640×480 RGBA8 natural pitch = 2560, request 2400 → should error.
3737 let err = Tensor::<u8>::image_with_stride(
3738 640,
3739 480,
3740 PixelFormat::Rgba,
3741 2400,
3742 Some(TensorMemory::Dma),
3743 );
3744 assert!(matches!(err, Err(Error::InvalidArgument(_))));
3745 }
3746
3747 #[test]
3748 #[cfg(target_os = "linux")]
3749 fn image_tensor_with_stride_rejects_non_packed() {
3750 // NV12 is SemiPlanar → not supported. (Linux-only because
3751 // `TensorMemory::Dma` itself is a Linux-only enum variant.)
3752 let err = Tensor::<u8>::image_with_stride(
3753 640,
3754 480,
3755 PixelFormat::Nv12,
3756 640,
3757 Some(TensorMemory::Dma),
3758 );
3759 assert!(matches!(err, Err(Error::NotImplemented(_))));
3760 }
3761
3762 #[test]
3763 fn set_format_valid() {
3764 let mut t = Tensor::<u8>::new(&[480, 640, 3], None, None).unwrap();
3765 assert!(t.format().is_none());
3766 t.set_format(PixelFormat::Rgb).unwrap();
3767 assert_eq!(t.format(), Some(PixelFormat::Rgb));
3768 assert_eq!(t.width(), Some(640));
3769 assert_eq!(t.height(), Some(480));
3770 }
3771
3772 #[test]
3773 fn set_format_invalid_shape() {
3774 let mut t = Tensor::<u8>::new(&[480, 640, 4], None, None).unwrap();
3775 // RGB expects 3 channels, not 4
3776 let err = t.set_format(PixelFormat::Rgb);
3777 assert!(err.is_err());
3778 // Original tensor is unmodified
3779 assert!(t.format().is_none());
3780 }
3781
3782 #[test]
3783 fn reshape_clears_format() {
3784 let mut t = Tensor::<u8>::image(640, 480, PixelFormat::Rgba, None).unwrap();
3785 assert_eq!(t.format(), Some(PixelFormat::Rgba));
3786 // Reshape to flat — format cleared
3787 t.reshape(&[480 * 640 * 4]).unwrap();
3788 assert!(t.format().is_none());
3789 }
3790
3791 #[test]
3792 fn from_planes_nv12() {
3793 let y = Tensor::<u8>::new(&[480, 640], None, None).unwrap();
3794 let uv = Tensor::<u8>::new(&[240, 640], None, None).unwrap();
3795 let img = Tensor::from_planes(y, uv, PixelFormat::Nv12).unwrap();
3796 assert_eq!(img.format(), Some(PixelFormat::Nv12));
3797 assert!(img.is_multiplane());
3798 assert!(img.chroma().is_some());
3799 assert_eq!(img.width(), Some(640));
3800 assert_eq!(img.height(), Some(480));
3801 }
3802
3803 #[test]
3804 fn from_planes_rejects_non_semiplanar() {
3805 let y = Tensor::<u8>::new(&[480, 640], None, None).unwrap();
3806 let uv = Tensor::<u8>::new(&[240, 640], None, None).unwrap();
3807 let err = Tensor::from_planes(y, uv, PixelFormat::Rgb);
3808 assert!(err.is_err());
3809 }
3810
3811 #[test]
3812 fn reshape_multiplane_errors() {
3813 let y = Tensor::<u8>::new(&[480, 640], None, None).unwrap();
3814 let uv = Tensor::<u8>::new(&[240, 640], None, None).unwrap();
3815 let mut img = Tensor::from_planes(y, uv, PixelFormat::Nv12).unwrap();
3816 let err = img.reshape(&[480 * 640 + 240 * 640]);
3817 assert!(err.is_err());
3818 }
3819}
3820
3821#[cfg(test)]
3822mod tests {
3823 #[cfg(target_os = "linux")]
3824 use nix::unistd::{access, AccessFlags};
3825 #[cfg(target_os = "linux")]
3826 use std::io::Write as _;
3827 use std::sync::RwLock;
3828
3829 use super::*;
3830
3831 #[ctor::ctor]
3832 fn init() {
3833 env_logger::Builder::from_env(env_logger::Env::default().default_filter_or("info")).init();
3834 }
3835
3836 /// Macro to get the current function name for logging in tests.
3837 #[cfg(target_os = "linux")]
3838 macro_rules! function {
3839 () => {{
3840 fn f() {}
3841 fn type_name_of<T>(_: T) -> &'static str {
3842 std::any::type_name::<T>()
3843 }
3844 let name = type_name_of(f);
3845
3846 // Find and cut the rest of the path
3847 match &name[..name.len() - 3].rfind(':') {
3848 Some(pos) => &name[pos + 1..name.len() - 3],
3849 None => &name[..name.len() - 3],
3850 }
3851 }};
3852 }
3853
3854 #[test]
3855 #[cfg(target_os = "linux")]
3856 fn test_tensor() {
3857 let _lock = FD_LOCK.read().unwrap();
3858 let shape = vec![1];
3859 let tensor = DmaTensor::<f32>::new(&shape, Some("dma_tensor"));
3860 let dma_enabled = tensor.is_ok();
3861
3862 let tensor = Tensor::<f32>::new(&shape, None, None).expect("Failed to create tensor");
3863 // Auto-select priority is Dma > Mem; Shm is never auto-selected.
3864 match dma_enabled {
3865 true => assert_eq!(tensor.memory(), TensorMemory::Dma),
3866 false => assert_eq!(tensor.memory(), TensorMemory::Mem),
3867 }
3868 }
3869
3870 #[test]
3871 #[cfg(target_os = "macos")]
3872 fn test_tensor() {
3873 let shape = vec![1];
3874 let tensor = Tensor::<f32>::new(&shape, None, None).expect("Failed to create tensor");
3875 // macOS auto-fallback chain: IOSurface (Dma) → Mem. Healthy systems
3876 // return Dma; Mem only appears under memory pressure or sandboxed
3877 // contexts where IOSurfaceCreate fails. Shm is never auto-selected.
3878 let m = tensor.memory();
3879 assert!(
3880 matches!(m, TensorMemory::Dma | TensorMemory::Mem),
3881 "Unexpected auto-fallback result on macOS: {m:?}"
3882 );
3883 }
3884
3885 #[test]
3886 #[cfg(all(unix, not(any(target_os = "linux", target_os = "macos"))))]
3887 fn test_tensor() {
3888 let shape = vec![1];
3889 let tensor = Tensor::<f32>::new(&shape, None, None).expect("Failed to create tensor");
3890 // Other Unix (BSD): no DMA, so auto-selection is Mem (Shm is
3891 // explicit-only, never auto-selected).
3892 assert_eq!(tensor.memory(), TensorMemory::Mem);
3893 }
3894
3895 #[test]
3896 #[cfg(not(unix))]
3897 fn test_tensor() {
3898 let shape = vec![1];
3899 let tensor = Tensor::<f32>::new(&shape, None, None).expect("Failed to create tensor");
3900 assert_eq!(tensor.memory(), TensorMemory::Mem);
3901 }
3902
3903 #[test]
3904 #[cfg(target_os = "linux")]
3905 fn test_dma_tensor() {
3906 let _lock = FD_LOCK.read().unwrap();
3907 match access(
3908 "/dev/dma_heap/linux,cma",
3909 AccessFlags::R_OK | AccessFlags::W_OK,
3910 ) {
3911 Ok(_) => println!("/dev/dma_heap/linux,cma is available"),
3912 Err(_) => match access(
3913 "/dev/dma_heap/system",
3914 AccessFlags::R_OK | AccessFlags::W_OK,
3915 ) {
3916 Ok(_) => println!("/dev/dma_heap/system is available"),
3917 Err(e) => {
3918 writeln!(
3919 &mut std::io::stdout(),
3920 "[WARNING] DMA Heap is unavailable: {e}"
3921 )
3922 .unwrap();
3923 return;
3924 }
3925 },
3926 }
3927
3928 let shape = vec![2, 3, 4];
3929 let tensor =
3930 DmaTensor::<f32>::new(&shape, Some("test_tensor")).expect("Failed to create tensor");
3931
3932 const DUMMY_VALUE: f32 = 12.34;
3933
3934 assert_eq!(tensor.memory(), TensorMemory::Dma);
3935 assert_eq!(tensor.name(), "test_tensor");
3936 assert_eq!(tensor.shape(), &shape);
3937 assert_eq!(tensor.size(), 2 * 3 * 4 * std::mem::size_of::<f32>());
3938 assert_eq!(tensor.len(), 2 * 3 * 4);
3939
3940 {
3941 let mut tensor_map = tensor.map().expect("Failed to map DMA memory");
3942 tensor_map.fill(42.0);
3943 assert!(tensor_map.iter().all(|&x| x == 42.0));
3944 }
3945
3946 {
3947 let shared = Tensor::<f32>::from_fd(
3948 tensor
3949 .clone_fd()
3950 .expect("Failed to duplicate tensor file descriptor"),
3951 &shape,
3952 Some("test_tensor_shared"),
3953 )
3954 .expect("Failed to create tensor from fd");
3955
3956 assert_eq!(shared.memory(), TensorMemory::Dma);
3957 assert_eq!(shared.name(), "test_tensor_shared");
3958 assert_eq!(shared.shape(), &shape);
3959
3960 let mut tensor_map = shared.map().expect("Failed to map DMA memory from fd");
3961 tensor_map.fill(DUMMY_VALUE);
3962 assert!(tensor_map.iter().all(|&x| x == DUMMY_VALUE));
3963 }
3964
3965 {
3966 let tensor_map = tensor.map().expect("Failed to map DMA memory");
3967 assert!(tensor_map.iter().all(|&x| x == DUMMY_VALUE));
3968 }
3969
3970 let mut tensor = DmaTensor::<u8>::new(&shape, None).expect("Failed to create tensor");
3971 assert_eq!(tensor.shape(), &shape);
3972 let new_shape = vec![3, 4, 4];
3973 assert!(
3974 tensor.reshape(&new_shape).is_err(),
3975 "Reshape should fail due to size mismatch"
3976 );
3977 assert_eq!(tensor.shape(), &shape, "Shape should remain unchanged");
3978
3979 let new_shape = vec![2, 3, 4];
3980 tensor.reshape(&new_shape).expect("Reshape should succeed");
3981 assert_eq!(
3982 tensor.shape(),
3983 &new_shape,
3984 "Shape should be updated after successful reshape"
3985 );
3986
3987 {
3988 let mut tensor_map = tensor.map().expect("Failed to map DMA memory");
3989 tensor_map.fill(1);
3990 assert!(tensor_map.iter().all(|&x| x == 1));
3991 }
3992
3993 {
3994 let mut tensor_map = tensor.map().expect("Failed to map DMA memory");
3995 tensor_map[2] = 42;
3996 assert_eq!(tensor_map[1], 1, "Value at index 1 should be 1");
3997 assert_eq!(tensor_map[2], 42, "Value at index 2 should be 42");
3998 }
3999 }
4000
4001 #[test]
4002 #[cfg(unix)]
4003 fn test_shm_tensor() {
4004 let _lock = FD_LOCK.read().unwrap();
4005 let shape = vec![2, 3, 4];
4006 let tensor =
4007 ShmTensor::<f32>::new(&shape, Some("test_tensor")).expect("Failed to create tensor");
4008 assert_eq!(tensor.shape(), &shape);
4009 assert_eq!(tensor.size(), 2 * 3 * 4 * std::mem::size_of::<f32>());
4010 assert_eq!(tensor.name(), "test_tensor");
4011
4012 const DUMMY_VALUE: f32 = 12.34;
4013 {
4014 let mut tensor_map = tensor.map().expect("Failed to map shared memory");
4015 tensor_map.fill(42.0);
4016 assert!(tensor_map.iter().all(|&x| x == 42.0));
4017 }
4018
4019 {
4020 let shared = Tensor::<f32>::from_fd(
4021 tensor
4022 .clone_fd()
4023 .expect("Failed to duplicate tensor file descriptor"),
4024 &shape,
4025 Some("test_tensor_shared"),
4026 )
4027 .expect("Failed to create tensor from fd");
4028
4029 assert_eq!(shared.memory(), TensorMemory::Shm);
4030 assert_eq!(shared.name(), "test_tensor_shared");
4031 assert_eq!(shared.shape(), &shape);
4032
4033 let mut tensor_map = shared.map().expect("Failed to map shared memory from fd");
4034 tensor_map.fill(DUMMY_VALUE);
4035 assert!(tensor_map.iter().all(|&x| x == DUMMY_VALUE));
4036 }
4037
4038 {
4039 let tensor_map = tensor.map().expect("Failed to map shared memory");
4040 assert!(tensor_map.iter().all(|&x| x == DUMMY_VALUE));
4041 }
4042
4043 let mut tensor = ShmTensor::<u8>::new(&shape, None).expect("Failed to create tensor");
4044 assert_eq!(tensor.shape(), &shape);
4045 let new_shape = vec![3, 4, 4];
4046 assert!(
4047 tensor.reshape(&new_shape).is_err(),
4048 "Reshape should fail due to size mismatch"
4049 );
4050 assert_eq!(tensor.shape(), &shape, "Shape should remain unchanged");
4051
4052 let new_shape = vec![2, 3, 4];
4053 tensor.reshape(&new_shape).expect("Reshape should succeed");
4054 assert_eq!(
4055 tensor.shape(),
4056 &new_shape,
4057 "Shape should be updated after successful reshape"
4058 );
4059
4060 {
4061 let mut tensor_map = tensor.map().expect("Failed to map shared memory");
4062 tensor_map.fill(1);
4063 assert!(tensor_map.iter().all(|&x| x == 1));
4064 }
4065
4066 {
4067 let mut tensor_map = tensor.map().expect("Failed to map shared memory");
4068 tensor_map[2] = 42;
4069 assert_eq!(tensor_map[1], 1, "Value at index 1 should be 1");
4070 assert_eq!(tensor_map[2], 42, "Value at index 2 should be 42");
4071 }
4072 }
4073
4074 #[test]
4075 fn mem_subview_partitions_parent_buffer() {
4076 // One heap [2,4] u8 parent (8 bytes). Two [1,4] sub-views at byte
4077 // offsets 0 and 4 must share the parent allocation (zero-copy) and be
4078 // independently writable: view 0 owns bytes [0,4), view 1 owns [4,8).
4079 // Today this is impossible — heap offset is rejected and there is no
4080 // shared sub-view constructor.
4081 let parent = Tensor::<u8>::new(&[2, 4], Some(TensorMemory::Mem), None).unwrap();
4082 let view0 = parent.subview(0, &[1, 4]).expect("subview at offset 0");
4083 let view1 = parent.subview(4, &[1, 4]).expect("subview at offset 4");
4084
4085 view1
4086 .map()
4087 .unwrap()
4088 .as_mut_slice()
4089 .copy_from_slice(&[10, 20, 30, 40]);
4090 view0
4091 .map()
4092 .unwrap()
4093 .as_mut_slice()
4094 .copy_from_slice(&[1, 2, 3, 4]);
4095
4096 // Each view sees only its own window.
4097 assert_eq!(view0.map().unwrap().as_slice(), &[1, 2, 3, 4]);
4098 assert_eq!(view1.map().unwrap().as_slice(), &[10, 20, 30, 40]);
4099 // The parent buffer is correctly partitioned (shared, zero-copy).
4100 assert_eq!(
4101 parent.map().unwrap().as_slice(),
4102 &[1, 2, 3, 4, 10, 20, 30, 40]
4103 );
4104 }
4105
4106 #[test]
4107 fn batch_partitions_leading_dim() {
4108 // Raw [4,2,2,3] u8 batched tensor: 4 elements of 12 bytes each. batch(n)
4109 // yields element n at offset n*12, sharing the parent buffer (zero-copy).
4110 let parent = Tensor::<u8>::new(&[4, 2, 2, 3], Some(TensorMemory::Mem), None).unwrap();
4111 for i in 0..4u8 {
4112 let e = parent.batch(i as usize).expect("batch element");
4113 assert_eq!(e.shape(), &[2, 2, 3]);
4114 // A batch element shares the parent's BufferIdentity.
4115 assert_eq!(e.buffer_identity().id(), parent.buffer_identity().id());
4116 for b in e.map().unwrap().as_mut_slice() {
4117 *b = i + 1;
4118 }
4119 }
4120 // Each element occupies its own 12-byte band of the parent.
4121 let whole = parent.map().unwrap();
4122 let s = whole.as_slice();
4123 for i in 0..4usize {
4124 assert!(
4125 s[i * 12..(i + 1) * 12].iter().all(|&b| b == (i as u8 + 1)),
4126 "band {i} not partitioned: {:?}",
4127 &s[i * 12..(i + 1) * 12]
4128 );
4129 }
4130 }
4131
4132 #[test]
4133 fn view_origin_snapshots_parent_and_composes() {
4134 // view() on a whole image snapshots the parent dims + the view's origin.
4135 let parent =
4136 Tensor::<u8>::image(100, 80, PixelFormat::Rgba, Some(TensorMemory::Mem)).unwrap();
4137 assert_eq!(
4138 parent.view_origin(),
4139 None,
4140 "whole tensor has no view_origin"
4141 );
4142 let v = parent.view(Region::new(10, 20, 30, 40)).unwrap();
4143 assert_eq!(
4144 v.view_origin(),
4145 Some(ViewOrigin {
4146 parent_width: 100,
4147 parent_height: 80,
4148 parent_row_stride: 100 * 4, // tight RGBA pitch
4149 x: 10,
4150 y: 20
4151 })
4152 );
4153 // A view of a view keeps the ROOT parent and accumulates the origin.
4154 let v2 = v.view(Region::new(5, 5, 10, 10)).unwrap();
4155 assert_eq!(
4156 v2.view_origin(),
4157 Some(ViewOrigin {
4158 parent_width: 100,
4159 parent_height: 80,
4160 parent_row_stride: 100 * 4,
4161 x: 15,
4162 y: 25
4163 }),
4164 "nested view composes onto the root parent"
4165 );
4166 }
4167
4168 #[test]
4169 fn view_origin_none_for_raw_batch() {
4170 // A raw (unformatted) batched tensor has no pixel geometry, so batch()
4171 // leaves view_origin None (the per-slot path, not the one-import pivot).
4172 let parent = Tensor::<u8>::new(&[4, 2, 2, 3], Some(TensorMemory::Mem), None).unwrap();
4173 assert_eq!(parent.batch(2).unwrap().view_origin(), None);
4174 }
4175
4176 #[test]
4177 fn batch_rejects_out_of_bounds_index() {
4178 let parent = Tensor::<u8>::new(&[4, 2, 2, 3], Some(TensorMemory::Mem), None).unwrap();
4179 match parent.batch(4) {
4180 Err(Error::BatchIndexOutOfBounds { index, batch }) => {
4181 assert_eq!((index, batch), (4, 4));
4182 }
4183 other => panic!("expected BatchIndexOutOfBounds, got {other:?}"),
4184 }
4185 }
4186
4187 #[test]
4188 fn batch_zero_on_unit_n_is_whole() {
4189 // N == 1: batch(0) is the whole per-element block at offset 0 (no plane_offset).
4190 let parent = Tensor::<u8>::new(&[1, 2, 2, 3], Some(TensorMemory::Mem), None).unwrap();
4191 let e = parent.batch(0).unwrap();
4192 assert_eq!(e.shape(), &[2, 2, 3]);
4193 assert_eq!(e.plane_offset(), None);
4194 assert_eq!(e.buffer_identity().id(), parent.buffer_identity().id());
4195 }
4196
4197 #[test]
4198 fn mem_subview_rejects_unaligned_offset() {
4199 // f32 has align 4; a byte offset of 2 cannot back a valid `*const f32`.
4200 let parent = Tensor::<f32>::new(&[8], Some(TensorMemory::Mem), None).unwrap();
4201 assert!(parent.subview(2, &[1]).is_err());
4202 // A correctly aligned offset is accepted.
4203 assert!(parent.subview(4, &[1]).is_ok());
4204 }
4205
4206 #[test]
4207 fn mem_subview_rejects_out_of_bounds() {
4208 let parent = Tensor::<u8>::new(&[8], Some(TensorMemory::Mem), None).unwrap();
4209 // offset 6 + 4 bytes = 10 exceeds the 8-byte allocation.
4210 assert!(parent.subview(6, &[4]).is_err());
4211 }
4212
4213 /// Regression guard for the `TensorTrait::view` promotion (R2): a `subview`
4214 /// must share the parent's `BufferIdentity` on **every** backend, not mint a
4215 /// fresh one. Identity-keyed caches (the GL EGLImage import) rely on this to
4216 /// treat offset-distinct windows of one buffer as a single import; a fresh
4217 /// identity would silently break that and regress zero-copy import reuse.
4218 ///
4219 /// Runs each backend that can be allocated without a GPU/GL context on the
4220 /// test host: `Mem` (always); `Shm` (when POSIX shm is available); the
4221 /// platform-native zero-copy buffer `Dma` (DMA-BUF on Linux / IOSurface on
4222 /// macOS, when available). `Pbo` shares its identity the same way (see
4223 /// `pbo.rs` `view`) but needs a live GL context, so it is exercised by the
4224 /// image-crate GL tests rather than here.
4225 #[test]
4226 fn subview_shares_buffer_identity_all_backends() {
4227 // u8 has align 1, so every byte offset is valid for the alignment check;
4228 // this isolates the identity-sharing contract from alignment concerns.
4229 let assert_shares = |memory: TensorMemory, label: &str| {
4230 let parent = Tensor::<u8>::new(&[64], Some(memory), None)
4231 .unwrap_or_else(|e| panic!("{label}: parent alloc failed: {e:?}"));
4232 let parent_id = parent.buffer_identity().id();
4233 // Two offset-distinct windows must both carry the parent's identity.
4234 let v0 = parent
4235 .subview(0, &[16])
4236 .unwrap_or_else(|e| panic!("{label}: subview(0) failed: {e:?}"));
4237 let v1 = parent
4238 .subview(16, &[16])
4239 .unwrap_or_else(|e| panic!("{label}: subview(16) failed: {e:?}"));
4240 assert_eq!(
4241 v0.buffer_identity().id(),
4242 parent_id,
4243 "{label}: subview(0) minted a fresh BufferIdentity"
4244 );
4245 assert_eq!(
4246 v1.buffer_identity().id(),
4247 parent_id,
4248 "{label}: subview(16) minted a fresh BufferIdentity"
4249 );
4250 };
4251
4252 assert_shares(TensorMemory::Mem, "Mem");
4253
4254 #[cfg(unix)]
4255 if crate::is_shm_available() {
4256 assert_shares(TensorMemory::Shm, "Shm");
4257 }
4258
4259 // Dma == DMA-BUF on Linux, IOSurface on macOS; same public variant.
4260 if crate::is_gpu_buffer_available() {
4261 assert_shares(TensorMemory::Dma, "Dma");
4262 }
4263 }
4264
4265 #[test]
4266 fn mem_subview_four_views_no_aliasing() {
4267 // One [4,3] f32 parent; four [1,3] views at 12-byte strides, each
4268 // written independently. Exercises a multi-byte element type (offsets
4269 // must stay element-aligned) and N-way zero-copy sharing.
4270 let parent = Tensor::<f32>::new(&[4, 3], Some(TensorMemory::Mem), None).unwrap();
4271 let frame = 3 * std::mem::size_of::<f32>();
4272 for i in 0..4 {
4273 let v = parent.subview(i * frame, &[1, 3]).unwrap();
4274 let val = i as f32 + 1.0;
4275 v.map()
4276 .unwrap()
4277 .as_mut_slice()
4278 .copy_from_slice(&[val, val, val]);
4279 }
4280 assert_eq!(
4281 parent.map().unwrap().as_slice(),
4282 &[1.0, 1.0, 1.0, 2.0, 2.0, 2.0, 3.0, 3.0, 3.0, 4.0, 4.0, 4.0]
4283 );
4284 }
4285
4286 #[test]
4287 fn mem_subview_inherits_format_and_row_stride() {
4288 // A sub-view is a ready-to-use sub-image: it inherits the parent's
4289 // pixel format and (crucially) its padded row stride, so a strided
4290 // parent yields strided windows. Set a stride wider than the tight row
4291 // to exercise the row_stride inheritance path specifically.
4292 let mut parent =
4293 Tensor::<u8>::image(100, 100, PixelFormat::Rgba, Some(TensorMemory::Mem)).unwrap();
4294 parent.set_row_stride_unchecked(512); // padded stride (> 100*4)
4295 let view = parent.subview(4096, &[10, 10, 4]).unwrap();
4296 assert_eq!(view.format(), Some(PixelFormat::Rgba), "format inherited");
4297 assert_eq!(view.row_stride(), Some(512), "row_stride inherited");
4298 }
4299
4300 #[test]
4301 fn mem_strided_subview_maps_offset_and_byte_size() {
4302 // Integration of the sub-region offset (PR #89) and the strided-map
4303 // `byte_size_override` (PR #90): a strided sub-view exposes its full
4304 // padded window (`row_stride × rows`) starting at the view's byte
4305 // offset, mapped zero-copy into the parent.
4306 let parent = Tensor::<u8>::new(&[2048], Some(TensorMemory::Mem), None).unwrap();
4307 let mut view = parent.subview(128, &[8, 16]).unwrap(); // 8 rows × 16 @ off 128
4308 assert_eq!(view.plane_offset(), Some(128));
4309 view.set_row_stride_unchecked(32); // padded stride (> 16)
4310
4311 {
4312 let mut m = view.map().unwrap();
4313 let s = m.as_mut_slice();
4314 // Strided map exposes the padded window: stride(32) × rows(8) = 256.
4315 assert_eq!(
4316 s.len(),
4317 256,
4318 "strided map exposes the full padded byte window"
4319 );
4320 s[0] = 0xAA; // row 0, col 0
4321 s[32] = 0xBB; // row 1, col 0 (one stride in)
4322 }
4323
4324 // Zero-copy: the writes land in the parent at the view's offset.
4325 let p = parent.map().unwrap();
4326 let pb = p.as_slice();
4327 assert_eq!(pb[128], 0xAA, "row 0 writes at parent offset 128");
4328 assert_eq!(
4329 pb[128 + 32],
4330 0xBB,
4331 "row 1 writes at parent offset 128 + stride"
4332 );
4333 }
4334
4335 #[test]
4336 #[cfg(unix)]
4337 fn shm_subview_partitions_parent_buffer() {
4338 // Mirrors `mem_subview_partitions_parent_buffer` for Shm: one [2,4] u8
4339 // parent shared segment (8 bytes); two [1,4] sub-views at byte offsets 0
4340 // and 4 must share the segment (zero-copy, via cloned fd) and be
4341 // independently writable — view 0 owns [0,4), view 1 owns [4,8).
4342 if !crate::is_shm_available() {
4343 eprintln!("SKIPPED: shm not available");
4344 return;
4345 }
4346 let parent = Tensor::<u8>::new(&[2, 4], Some(TensorMemory::Shm), None).unwrap();
4347 let view0 = parent.subview(0, &[1, 4]).expect("shm subview at offset 0");
4348 let view1 = parent.subview(4, &[1, 4]).expect("shm subview at offset 4");
4349
4350 view1
4351 .map()
4352 .unwrap()
4353 .as_mut_slice()
4354 .copy_from_slice(&[10, 20, 30, 40]);
4355 view0
4356 .map()
4357 .unwrap()
4358 .as_mut_slice()
4359 .copy_from_slice(&[1, 2, 3, 4]);
4360
4361 assert_eq!(view0.map().unwrap().as_slice(), &[1, 2, 3, 4]);
4362 assert_eq!(view1.map().unwrap().as_slice(), &[10, 20, 30, 40]);
4363 // The parent sees the full partitioned segment (shared, zero-copy).
4364 assert_eq!(
4365 parent.map().unwrap().as_slice(),
4366 &[1, 2, 3, 4, 10, 20, 30, 40]
4367 );
4368 // A sub-view of a sub-view composes the offset.
4369 let nested = view1.subview(2, &[1, 2]).expect("nested shm subview");
4370 assert_eq!(nested.map().unwrap().as_slice(), &[30, 40]);
4371 }
4372
4373 #[test]
4374 #[cfg(unix)]
4375 fn shm_subview_rejects_unaligned_and_oob() {
4376 if !crate::is_shm_available() {
4377 eprintln!("SKIPPED: shm not available");
4378 return;
4379 }
4380 // f32 align 4: a 2-byte offset cannot back a valid `*const f32`.
4381 let parent = Tensor::<f32>::new(&[8], Some(TensorMemory::Shm), None).unwrap();
4382 assert!(parent.subview(2, &[1]).is_err());
4383 assert!(parent.subview(4, &[1]).is_ok());
4384 // Out of bounds: offset 6 + 4 bytes = 10 > 8-byte (u8) segment.
4385 let p2 = Tensor::<u8>::new(&[8], Some(TensorMemory::Shm), None).unwrap();
4386 assert!(p2.subview(6, &[4]).is_err());
4387 }
4388
4389 #[test]
4390 #[cfg(target_os = "linux")]
4391 fn dma_subview_matches_mem_subview() {
4392 // Serialize against the fd-leak tests: this test opens DMA fds (alloc +
4393 // clone_fd), which would otherwise perturb their fd counts.
4394 let _lock = FD_LOCK.read().unwrap();
4395 // Identical sub-view semantics across Dma (shared fd) and Mem (shared
4396 // Arc): same offsets → same logical windows → same partition.
4397 let dma = match Tensor::<u8>::new(&[8], Some(TensorMemory::Dma), None) {
4398 Ok(t) => t,
4399 Err(_) => {
4400 eprintln!("SKIPPED: DMA not available");
4401 return;
4402 }
4403 };
4404 let mem = Tensor::<u8>::new(&[8], Some(TensorMemory::Mem), None).unwrap();
4405 for parent in [&dma, &mem] {
4406 let v0 = parent.subview(0, &[4]).unwrap();
4407 let v1 = parent.subview(4, &[4]).unwrap();
4408 v0.map()
4409 .unwrap()
4410 .as_mut_slice()
4411 .copy_from_slice(&[1, 2, 3, 4]);
4412 v1.map()
4413 .unwrap()
4414 .as_mut_slice()
4415 .copy_from_slice(&[5, 6, 7, 8]);
4416 assert_eq!(parent.map().unwrap().as_slice(), &[1, 2, 3, 4, 5, 6, 7, 8]);
4417 }
4418 }
4419
4420 #[test]
4421 #[cfg(target_os = "linux")]
4422 fn dma_strided_subview_maps_padded_window() {
4423 // The strided-map path differs by backing: DMA maps through
4424 // `mmap_offset` + the `byte_size_override`, not the Mem `Arc` slice. A
4425 // padded sub-view of a DMA buffer must still expose its full
4426 // `row_stride × rows` window zero-copy at the view's offset (the GPU
4427 // batched-render-to-DMA case). Mirrors
4428 // `mem_strided_subview_maps_offset_and_byte_size` on a Dma parent.
4429 let _lock = FD_LOCK.read().unwrap();
4430 let parent = match Tensor::<u8>::new(&[2048], Some(TensorMemory::Dma), None) {
4431 Ok(t) => t,
4432 Err(_) => {
4433 eprintln!("SKIPPED: DMA not available");
4434 return;
4435 }
4436 };
4437 let mut view = parent.subview(128, &[8, 16]).unwrap();
4438 assert_eq!(view.plane_offset(), Some(128));
4439 view.set_row_stride_unchecked(32); // padded stride (> 16)
4440
4441 {
4442 let mut m = view.map().unwrap();
4443 let s = m.as_mut_slice();
4444 assert_eq!(s.len(), 256, "strided DMA map exposes stride(32) × rows(8)");
4445 s[0] = 0xAA; // row 0, col 0
4446 s[32] = 0xBB; // row 1, col 0 (one stride in)
4447 }
4448
4449 let p = parent.map().unwrap();
4450 let pb = p.as_slice();
4451 assert_eq!(pb[128], 0xAA, "row 0 writes at parent offset 128");
4452 assert_eq!(
4453 pb[128 + 32],
4454 0xBB,
4455 "row 1 writes at parent offset 128 + stride"
4456 );
4457 }
4458
4459 #[test]
4460 #[cfg(target_os = "linux")]
4461 fn view_single_row_snapshots_parent_stride() {
4462 // A single-row `view()` keeps a TIGHT `row_stride` for map-span safety,
4463 // but its `view_origin` snapshots the PARENT row stride — the GL backend
4464 // keys its EGLImage import/pitch on that snapshot (not the view's tight
4465 // stride), so single-row and multi-row sibling views collapse onto the
4466 // same parent import.
4467 let _lock = FD_LOCK.read().unwrap();
4468 // 8x4 RGBA with a padded 64-byte row stride (tight row = 8*4 = 32).
4469 let parent = match Tensor::<u8>::image_with_stride(
4470 8,
4471 4,
4472 PixelFormat::Rgba,
4473 64,
4474 Some(TensorMemory::Dma),
4475 ) {
4476 Ok(t) => t,
4477 Err(_) => {
4478 eprintln!("SKIPPED: DMA not available");
4479 return;
4480 }
4481 };
4482 assert_eq!(parent.effective_row_stride(), Some(64));
4483 // Bottom row (y=3) at x>0 — the case the tight single-row stride guards.
4484 let row = parent.view(Region::new(2, 3, 4, 1)).unwrap();
4485 // The view's own stride is tight (4*4 = 16) so its strided map stays in
4486 // bounds; the GL-facing parent pitch (64) lives in `view_origin`.
4487 assert_eq!(row.effective_row_stride(), Some(16));
4488 let vo = row.view_origin().expect("a view carries a view_origin");
4489 assert_eq!(
4490 vo.parent_row_stride, 64,
4491 "GL keys/pitches a view on the parent stride, not its tight one"
4492 );
4493 // The tight stride keeps map() in-bounds for the bottom / x>0 single row.
4494 assert_eq!(row.map().unwrap().as_slice().len(), 16);
4495 }
4496
4497 #[test]
4498 fn test_mem_tensor() {
4499 let shape = vec![2, 3, 4];
4500 let tensor =
4501 MemTensor::<f32>::new(&shape, Some("test_tensor")).expect("Failed to create tensor");
4502 assert_eq!(tensor.shape(), &shape);
4503 assert_eq!(tensor.size(), 2 * 3 * 4 * std::mem::size_of::<f32>());
4504 assert_eq!(tensor.name(), "test_tensor");
4505
4506 {
4507 let mut tensor_map = tensor.map().expect("Failed to map memory");
4508 tensor_map.fill(42.0);
4509 assert!(tensor_map.iter().all(|&x| x == 42.0));
4510 }
4511
4512 let mut tensor = MemTensor::<u8>::new(&shape, None).expect("Failed to create tensor");
4513 assert_eq!(tensor.shape(), &shape);
4514 let new_shape = vec![3, 4, 4];
4515 assert!(
4516 tensor.reshape(&new_shape).is_err(),
4517 "Reshape should fail due to size mismatch"
4518 );
4519 assert_eq!(tensor.shape(), &shape, "Shape should remain unchanged");
4520
4521 let new_shape = vec![2, 3, 4];
4522 tensor.reshape(&new_shape).expect("Reshape should succeed");
4523 assert_eq!(
4524 tensor.shape(),
4525 &new_shape,
4526 "Shape should be updated after successful reshape"
4527 );
4528
4529 {
4530 let mut tensor_map = tensor.map().expect("Failed to map memory");
4531 tensor_map.fill(1);
4532 assert!(tensor_map.iter().all(|&x| x == 1));
4533 }
4534
4535 {
4536 let mut tensor_map = tensor.map().expect("Failed to map memory");
4537 tensor_map[2] = 42;
4538 assert_eq!(tensor_map[1], 1, "Value at index 1 should be 1");
4539 assert_eq!(tensor_map[2], 42, "Value at index 2 should be 42");
4540 }
4541 }
4542
4543 #[test]
4544 #[cfg(target_os = "linux")]
4545 fn test_dma_no_fd_leaks() {
4546 let _lock = FD_LOCK.write().unwrap();
4547 if !is_dma_available() {
4548 log::warn!(
4549 "SKIPPED: {} - DMA memory allocation not available (permission denied or no DMA-BUF support)",
4550 function!()
4551 );
4552 return;
4553 }
4554
4555 let proc = procfs::process::Process::myself()
4556 .expect("Failed to get current process using /proc/self");
4557
4558 let start_open_fds = proc
4559 .fd_count()
4560 .expect("Failed to get open file descriptor count");
4561
4562 for _ in 0..100 {
4563 let tensor = Tensor::<u8>::new(&[100, 100], Some(TensorMemory::Dma), None)
4564 .expect("Failed to create tensor");
4565 let mut map = tensor.map().unwrap();
4566 map.as_mut_slice().fill(233);
4567 }
4568
4569 let end_open_fds = proc
4570 .fd_count()
4571 .expect("Failed to get open file descriptor count");
4572
4573 assert_eq!(
4574 start_open_fds, end_open_fds,
4575 "File descriptor leak detected: {} -> {}",
4576 start_open_fds, end_open_fds
4577 );
4578 }
4579
4580 #[test]
4581 #[cfg(target_os = "linux")]
4582 fn test_dma_from_fd_no_fd_leaks() {
4583 let _lock = FD_LOCK.write().unwrap();
4584 if !is_dma_available() {
4585 log::warn!(
4586 "SKIPPED: {} - DMA memory allocation not available (permission denied or no DMA-BUF support)",
4587 function!()
4588 );
4589 return;
4590 }
4591
4592 let proc = procfs::process::Process::myself()
4593 .expect("Failed to get current process using /proc/self");
4594
4595 let start_open_fds = proc
4596 .fd_count()
4597 .expect("Failed to get open file descriptor count");
4598
4599 let orig = Tensor::<u8>::new(&[100, 100], Some(TensorMemory::Dma), None).unwrap();
4600
4601 for _ in 0..100 {
4602 let tensor =
4603 Tensor::<u8>::from_fd(orig.clone_fd().unwrap(), orig.shape(), None).unwrap();
4604 let mut map = tensor.map().unwrap();
4605 map.as_mut_slice().fill(233);
4606 }
4607 drop(orig);
4608
4609 let end_open_fds = proc.fd_count().unwrap();
4610
4611 assert_eq!(
4612 start_open_fds, end_open_fds,
4613 "File descriptor leak detected: {} -> {}",
4614 start_open_fds, end_open_fds
4615 );
4616 }
4617
4618 #[test]
4619 #[cfg(target_os = "linux")]
4620 fn test_shm_no_fd_leaks() {
4621 let _lock = FD_LOCK.write().unwrap();
4622 if !is_shm_available() {
4623 log::warn!(
4624 "SKIPPED: {} - SHM memory allocation not available (permission denied or no SHM support)",
4625 function!()
4626 );
4627 return;
4628 }
4629
4630 let proc = procfs::process::Process::myself()
4631 .expect("Failed to get current process using /proc/self");
4632
4633 let start_open_fds = proc
4634 .fd_count()
4635 .expect("Failed to get open file descriptor count");
4636
4637 for _ in 0..100 {
4638 let tensor = Tensor::<u8>::new(&[100, 100], Some(TensorMemory::Shm), None)
4639 .expect("Failed to create tensor");
4640 let mut map = tensor.map().unwrap();
4641 map.as_mut_slice().fill(233);
4642 }
4643
4644 let end_open_fds = proc
4645 .fd_count()
4646 .expect("Failed to get open file descriptor count");
4647
4648 assert_eq!(
4649 start_open_fds, end_open_fds,
4650 "File descriptor leak detected: {} -> {}",
4651 start_open_fds, end_open_fds
4652 );
4653 }
4654
4655 #[test]
4656 #[cfg(target_os = "linux")]
4657 fn test_shm_from_fd_no_fd_leaks() {
4658 let _lock = FD_LOCK.write().unwrap();
4659 if !is_shm_available() {
4660 log::warn!(
4661 "SKIPPED: {} - SHM memory allocation not available (permission denied or no SHM support)",
4662 function!()
4663 );
4664 return;
4665 }
4666
4667 let proc = procfs::process::Process::myself()
4668 .expect("Failed to get current process using /proc/self");
4669
4670 let start_open_fds = proc
4671 .fd_count()
4672 .expect("Failed to get open file descriptor count");
4673
4674 let orig = Tensor::<u8>::new(&[100, 100], Some(TensorMemory::Shm), None).unwrap();
4675
4676 for _ in 0..100 {
4677 let tensor =
4678 Tensor::<u8>::from_fd(orig.clone_fd().unwrap(), orig.shape(), None).unwrap();
4679 let mut map = tensor.map().unwrap();
4680 map.as_mut_slice().fill(233);
4681 }
4682 drop(orig);
4683
4684 let end_open_fds = proc.fd_count().unwrap();
4685
4686 assert_eq!(
4687 start_open_fds, end_open_fds,
4688 "File descriptor leak detected: {} -> {}",
4689 start_open_fds, end_open_fds
4690 );
4691 }
4692
4693 #[cfg(feature = "ndarray")]
4694 #[test]
4695 fn test_ndarray() {
4696 let _lock = FD_LOCK.read().unwrap();
4697 let shape = vec![2, 3, 4];
4698 let tensor = Tensor::<f32>::new(&shape, None, None).expect("Failed to create tensor");
4699
4700 let mut tensor_map = tensor.map().expect("Failed to map tensor memory");
4701 tensor_map.fill(1.0);
4702
4703 let view = tensor_map.view().expect("Failed to get ndarray view");
4704 assert_eq!(view.shape(), &[2, 3, 4]);
4705 assert!(view.iter().all(|&x| x == 1.0));
4706
4707 let mut view_mut = tensor_map
4708 .view_mut()
4709 .expect("Failed to get mutable ndarray view");
4710 view_mut[[0, 0, 0]] = 42.0;
4711 assert_eq!(view_mut[[0, 0, 0]], 42.0);
4712 assert_eq!(tensor_map[0], 42.0, "Value at index 0 should be 42");
4713 }
4714
4715 #[test]
4716 fn test_buffer_identity_unique() {
4717 let id1 = BufferIdentity::new();
4718 let id2 = BufferIdentity::new();
4719 assert_ne!(
4720 id1.id(),
4721 id2.id(),
4722 "Two identities should have different ids"
4723 );
4724 }
4725
4726 #[test]
4727 fn test_buffer_identity_clone_shares_guard() {
4728 let id1 = BufferIdentity::new();
4729 let weak = id1.weak();
4730 assert!(
4731 weak.upgrade().is_some(),
4732 "Weak should be alive while original exists"
4733 );
4734
4735 let id2 = id1.clone();
4736 assert_eq!(id1.id(), id2.id(), "Cloned identity should have same id");
4737
4738 drop(id1);
4739 assert!(
4740 weak.upgrade().is_some(),
4741 "Weak should still be alive (clone holds Arc)"
4742 );
4743
4744 drop(id2);
4745 assert!(
4746 weak.upgrade().is_none(),
4747 "Weak should be dead after all clones dropped"
4748 );
4749 }
4750
4751 #[test]
4752 fn test_tensor_buffer_identity() {
4753 let t1 = Tensor::<u8>::new(&[100], Some(TensorMemory::Mem), Some("t1")).unwrap();
4754 let t2 = Tensor::<u8>::new(&[100], Some(TensorMemory::Mem), Some("t2")).unwrap();
4755 assert_ne!(
4756 t1.buffer_identity().id(),
4757 t2.buffer_identity().id(),
4758 "Different tensors should have different buffer ids"
4759 );
4760 }
4761
4762 // ------------------------------------------------------------------------
4763 // Quantization — constructor validation + accessor correctness.
4764 // ------------------------------------------------------------------------
4765
4766 #[test]
4767 fn test_quantization_per_tensor_constructors() {
4768 let q = Quantization::per_tensor(0.1, -5);
4769 assert!(q.is_per_tensor());
4770 assert!(!q.is_per_channel());
4771 assert!(!q.is_symmetric());
4772 assert_eq!(q.scale(), &[0.1]);
4773 assert_eq!(q.zero_point(), Some(&[-5][..]));
4774
4775 let qs = Quantization::per_tensor_symmetric(0.05);
4776 assert!(qs.is_per_tensor());
4777 assert!(qs.is_symmetric());
4778 assert_eq!(qs.zero_point(), None);
4779 }
4780
4781 #[test]
4782 fn test_quantization_per_channel_constructors() {
4783 let q = Quantization::per_channel(vec![0.1, 0.2, 0.3], vec![0, -1, 1], 2).unwrap();
4784 assert!(q.is_per_channel());
4785 assert!(!q.is_symmetric());
4786 assert_eq!(q.axis(), Some(2));
4787 assert_eq!(q.scale().len(), 3);
4788
4789 let qs = Quantization::per_channel_symmetric(vec![0.054, 0.089, 0.195], 0).unwrap();
4790 assert!(qs.is_per_channel());
4791 assert!(qs.is_symmetric());
4792 assert_eq!(qs.axis(), Some(0));
4793 }
4794
4795 #[test]
4796 fn test_quantization_per_channel_length_mismatch_rejected() {
4797 // len(scales) != len(zero_points) → rejected at construction.
4798 let err = Quantization::per_channel(vec![0.1, 0.2], vec![0, 0, 0], 0).unwrap_err();
4799 assert!(matches!(err, Error::QuantizationInvalid { .. }));
4800 }
4801
4802 #[test]
4803 fn test_quantization_per_channel_empty_rejected() {
4804 let err = Quantization::per_channel_symmetric(vec![], 0).unwrap_err();
4805 assert!(matches!(err, Error::QuantizationInvalid { .. }));
4806 }
4807
4808 /// Constructors guard scale/zero_point length invariants, but
4809 /// `Quantization` is `Deserialize`, so malformed JSON (e.g. an
4810 /// empty `scale` array, or `zero_point` length that disagrees with
4811 /// `scale`) bypasses the constructor checks. `set_quantization`
4812 /// must reject these via `validate()` so they don't poison
4813 /// downstream `mode()` selection or per-channel kernel indexing.
4814 #[test]
4815 fn test_quantization_validate_rejects_malformed_deserialize() {
4816 let mut t = Tensor::<i8>::new(&[1, 1, 4], Some(TensorMemory::Mem), None).unwrap();
4817
4818 // Empty scale array: must be rejected.
4819 let q: Quantization = serde_json::from_str(r#"{"scale": []}"#).unwrap();
4820 assert!(matches!(
4821 t.set_quantization(q).unwrap_err(),
4822 Error::QuantizationInvalid { .. }
4823 ));
4824
4825 // Per-tensor with multi-element zero_point: must be rejected.
4826 let q: Quantization =
4827 serde_json::from_str(r#"{"scale": 0.1, "zero_point": [0, 0, 0]}"#).unwrap();
4828 assert!(matches!(
4829 t.set_quantization(q).unwrap_err(),
4830 Error::QuantizationInvalid { .. }
4831 ));
4832
4833 // Per-channel zero_point length != scale length: must be rejected.
4834 let q: Quantization = serde_json::from_str(
4835 r#"{"scale": [0.1, 0.2, 0.3, 0.4], "zero_point": [0, 0], "axis": 2}"#,
4836 )
4837 .unwrap();
4838 assert!(matches!(
4839 t.set_quantization(q).unwrap_err(),
4840 Error::QuantizationInvalid { .. }
4841 ));
4842 }
4843
4844 #[test]
4845 fn test_quantization_mode_dispatch() {
4846 let pt = Quantization::per_tensor(0.1, -5);
4847 assert!(matches!(
4848 pt.mode(),
4849 QuantMode::PerTensor { scale, zero_point } if scale == 0.1 && zero_point == -5
4850 ));
4851
4852 let pts = Quantization::per_tensor_symmetric(0.05);
4853 assert!(matches!(
4854 pts.mode(),
4855 QuantMode::PerTensorSymmetric { scale } if scale == 0.05
4856 ));
4857
4858 let pc = Quantization::per_channel(vec![0.1, 0.2], vec![0, -1], 2).unwrap();
4859 assert!(matches!(pc.mode(), QuantMode::PerChannel { axis: 2, .. }));
4860
4861 let pcs = Quantization::per_channel_symmetric(vec![0.1, 0.2], 0).unwrap();
4862 assert!(matches!(
4863 pcs.mode(),
4864 QuantMode::PerChannelSymmetric { axis: 0, .. }
4865 ));
4866 }
4867
4868 #[test]
4869 fn test_tensor_quantization_roundtrip_integer() {
4870 let mut t = Tensor::<i8>::new(&[2, 3, 4], Some(TensorMemory::Mem), None).unwrap();
4871 assert!(t.quantization().is_none());
4872 t.set_quantization(Quantization::per_tensor(0.1, -5))
4873 .unwrap();
4874 let q = t.quantization().unwrap();
4875 assert_eq!(q.scale(), &[0.1]);
4876 t.clear_quantization();
4877 assert!(t.quantization().is_none());
4878 }
4879
4880 #[test]
4881 fn test_tensor_with_quantization_builder() {
4882 let t = Tensor::<i8>::new(&[4, 4], Some(TensorMemory::Mem), None)
4883 .unwrap()
4884 .with_quantization(Quantization::per_tensor_symmetric(0.05))
4885 .unwrap();
4886 assert!(t.quantization().is_some());
4887 }
4888
4889 #[test]
4890 fn test_tensor_dyn_quantization_float_arm_returns_none() {
4891 let t = Tensor::<f32>::new(&[2, 2], Some(TensorMemory::Mem), None).unwrap();
4892 let td = TensorDyn::F32(t);
4893 assert!(td.quantization().is_none());
4894 }
4895
4896 #[test]
4897 fn test_tensor_dyn_set_quantization_float_arm_errors() {
4898 let t = Tensor::<f32>::new(&[2, 2], Some(TensorMemory::Mem), None).unwrap();
4899 let mut td = TensorDyn::F32(t);
4900 let err = td
4901 .set_quantization(Quantization::per_tensor(0.1, 0))
4902 .unwrap_err();
4903 // float path returns a QuantizationInvalid error.
4904 assert!(matches!(err, Error::QuantizationInvalid { .. }));
4905 }
4906
4907 /// Compile-time type gate — calling `Tensor::<f32>::quantization()` must
4908 /// fail to compile (the `IntegerType` trait bound is not satisfied by
4909 /// `f32`). This doctest anchors the invariant.
4910 ///
4911 /// ```compile_fail
4912 /// use edgefirst_tensor::{Tensor, TensorMemory};
4913 /// let t = Tensor::<f32>::new(&[2, 2], Some(TensorMemory::Mem), None).unwrap();
4914 /// let _ = t.quantization(); // compile error: f32 not IntegerType
4915 /// ```
4916 fn _compile_fail_doctest_anchor() {}
4917
4918 // Any test that cares about the fd count must grab it exclusively.
4919 // Any tests which modifies the fd count by opening or closing fds must grab it
4920 // shared.
4921 pub static FD_LOCK: RwLock<()> = RwLock::new(());
4922
4923 /// Test that DMA is NOT available on non-Linux platforms.
4924 /// This verifies the cross-platform behavior of is_dma_available().
4925 #[test]
4926 #[cfg(not(target_os = "linux"))]
4927 fn test_dma_not_available_on_non_linux() {
4928 assert!(
4929 !is_dma_available(),
4930 "DMA memory allocation should NOT be available on non-Linux platforms"
4931 );
4932 }
4933
4934 #[test]
4935 fn colorimetry_defaults_none_and_roundtrips_without_auto_fill() {
4936 use crate::{ColorEncoding, ColorRange, Colorimetry, PixelFormat, TensorMemory};
4937 let mut t =
4938 Tensor::<u8>::image(1280, 720, PixelFormat::Nv12, Some(TensorMemory::Mem)).unwrap();
4939 assert_eq!(t.colorimetry(), None); // default undefined
4940 let c = Colorimetry::default()
4941 .with_encoding(ColorEncoding::Bt709)
4942 .with_range(ColorRange::Limited);
4943 t.set_colorimetry(Some(c));
4944 assert_eq!(t.colorimetry(), Some(c));
4945 // configure_image must NOT touch colorimetry
4946 t.configure_image(640, 480, PixelFormat::Grey).unwrap();
4947 assert_eq!(t.colorimetry(), Some(c));
4948 }
4949
4950 #[test]
4951 fn configure_image_within_capacity() {
4952 let mut t = Tensor::<u8>::image_with_capacity(640, 480, PixelFormat::Rgb, None).unwrap();
4953 t.configure_image(320, 240, PixelFormat::Nv12).unwrap();
4954 assert_eq!(t.format(), Some(PixelFormat::Nv12));
4955 assert_eq!(t.width(), Some(320));
4956 assert_eq!(t.height(), Some(240));
4957 assert_eq!(t.shape(), &[360, 320]); // 240*3/2
4958 }
4959
4960 #[test]
4961 fn configure_image_too_large_errors() {
4962 let mut t = Tensor::<u8>::image_with_capacity(64, 64, PixelFormat::Grey, None).unwrap();
4963 let err = t
4964 .configure_image(1920, 1080, PixelFormat::Nv12)
4965 .unwrap_err();
4966 assert!(matches!(err, Error::InsufficientCapacity { .. }));
4967 }
4968
4969 /// A reused max-sized IOSurface pool keeps its physical `bytesPerRow` when
4970 /// reconfigured to a smaller logical image (physical-grid / logical-ROI
4971 /// decoupling), instead of collapsing to the frame's natural row stride.
4972 #[test]
4973 #[cfg(target_os = "macos")]
4974 fn configure_image_preserves_iosurface_physical_stride() {
4975 // Pool: GREY/R8 IOSurface 100 wide → bytesPerRow padded to 128.
4976 let mut pool =
4977 Tensor::<u8>::image(100, 64, PixelFormat::Grey, Some(TensorMemory::Dma)).unwrap();
4978 let pitch = pool.effective_row_stride().unwrap();
4979 assert!(
4980 pitch >= 128 && pitch.is_multiple_of(64),
4981 "padded bytesPerRow, got {pitch}"
4982 );
4983
4984 // Reconfigure to a smaller NV12 frame; the physical pitch must survive
4985 // (natural would be 32, but the surface stride is the 128-padded pitch).
4986 pool.configure_image(32, 16, PixelFormat::Nv12).unwrap();
4987 assert_eq!(pool.format(), Some(PixelFormat::Nv12));
4988 assert_eq!(pool.width(), Some(32));
4989 assert_eq!(pool.height(), Some(16));
4990 assert_eq!(
4991 pool.effective_row_stride(),
4992 Some(pitch),
4993 "configure_image must preserve the IOSurface physical bytesPerRow"
4994 );
4995
4996 // Reconfigure again to NV24 — pitch still preserved.
4997 pool.configure_image(32, 16, PixelFormat::Nv24).unwrap();
4998 assert_eq!(pool.effective_row_stride(), Some(pitch));
4999 }
5000
5001 /// `configure_image` on a Mem backing reconfigures to the format's
5002 /// **64-byte-aligned** row stride (the odd-dim contract: every image tensor
5003 /// carries a 64-aligned `row_stride`). For NV12 32×16 the minimum is
5004 /// `even(32)=32`, rounded up to the 64-byte alignment → 64. The capacity
5005 /// (64×64×4 RGBA = 16 KiB) easily holds the 24×64 = 1.5 KiB NV12 layout.
5006 #[test]
5007 fn configure_image_mem_aligns_stride() {
5008 let mut t =
5009 Tensor::<u8>::image_with_capacity(64, 64, PixelFormat::Rgba, Some(TensorMemory::Mem))
5010 .unwrap();
5011 t.configure_image(32, 16, PixelFormat::Nv12).unwrap();
5012 let s = t.effective_row_stride().unwrap();
5013 assert_eq!(s % 64, 0, "stride must be 64-aligned");
5014 assert!(s >= 32, "stride must cover the even-width minimum");
5015 assert_eq!(s, 64);
5016 }
5017
5018 #[test]
5019 fn strided_mem_tensor_cpu_maps_full_padded_buffer() {
5020 // A packed RGBA image with row padding (GPU-pitch style): logical width
5021 // 8 px (32 B/row) but a 48-byte row stride. Over-allocate capacity (for
5022 // 16 px), narrow the logical width, then record the padded stride.
5023 // Previously `map()` rejected this on non-Linux with
5024 // "DMA backing is Linux-only"; HAL-owned Mem is now mappable.
5025 let mut t =
5026 Tensor::<u8>::image_with_capacity(16, 3, PixelFormat::Rgba, Some(TensorMemory::Mem))
5027 .unwrap(); // capacity 3 × 16 × 4 = 192 B
5028 t.configure_image(8, 3, PixelFormat::Rgba).unwrap(); // logical [3, 8, 4] = 96 B
5029 t.set_row_stride(48).unwrap(); // padded stride (>= 32 B min)
5030
5031 let map = t.map().expect("strided Mem tensor should CPU-map");
5032 // Full padded buffer (stride 48 × 3 rows = 144 B), not the 96 B logical
5033 // view — callers iterate rows via `effective_row_stride()`.
5034 assert_eq!(map.as_slice().len(), 144);
5035 // Logical shape is still reported for shape-aware consumers.
5036 assert_eq!(map.shape(), &[3, 8, 4]);
5037 }
5038
5039 #[test]
5040 fn strided_mem_tensor_over_capacity_errors() {
5041 // Stride larger than the allocation: 64 B × 3 rows = 192 B > 96 B cap.
5042 let mut t = Tensor::<u8>::new(&[3, 8, 4], Some(TensorMemory::Mem), None).unwrap();
5043 t.set_format(PixelFormat::Rgba).unwrap();
5044 t.set_row_stride(64).unwrap();
5045 assert!(matches!(t.map(), Err(Error::InsufficientCapacity { .. })));
5046 }
5047
5048 /// Test that SHM memory allocation is available and usable on Unix systems.
5049 /// This is a basic functional test; Linux has additional FD leak tests using procfs.
5050 #[test]
5051 #[cfg(unix)]
5052 fn test_shm_available_and_usable() {
5053 assert!(
5054 is_shm_available(),
5055 "SHM memory allocation should be available on Unix systems"
5056 );
5057
5058 // Create a tensor with SHM backing
5059 let tensor = Tensor::<u8>::new(&[100, 100], Some(TensorMemory::Shm), None)
5060 .expect("Failed to create SHM tensor");
5061
5062 // Verify we can map and write to it
5063 let mut map = tensor.map().expect("Failed to map SHM tensor");
5064 map.as_mut_slice().fill(0xAB);
5065
5066 // Verify the data was written correctly
5067 assert!(
5068 map.as_slice().iter().all(|&b| b == 0xAB),
5069 "SHM tensor data should be writable and readable"
5070 );
5071 }
5072
5073 // =========================================================================
5074 // packed_rgba16f_layout — host-runnable geometry unit tests (TDD)
5075 // =========================================================================
5076
5077 #[test]
5078 fn packed_rgba16f_layout_planar_rgb_f16() {
5079 let layout =
5080 packed_rgba16f_layout(PixelFormat::PlanarRgb, DType::F16, 640, 640).expect("Some");
5081 assert_eq!(layout.surface_w, 160);
5082 assert_eq!(layout.surface_h, 1920);
5083 assert_eq!(layout.bytes_per_texel, 8);
5084 assert_eq!(layout.pitch, 1280);
5085 }
5086
5087 #[test]
5088 fn packed_rgba16f_layout_planar_rgba_f16() {
5089 let layout =
5090 packed_rgba16f_layout(PixelFormat::PlanarRgba, DType::F16, 640, 640).expect("Some");
5091 assert_eq!(layout.surface_w, 160);
5092 assert_eq!(layout.surface_h, 2560); // 4 planes
5093 assert_eq!(layout.bytes_per_texel, 8);
5094 assert_eq!(layout.pitch, 1280);
5095 }
5096
5097 #[test]
5098 fn packed_rgba16f_layout_rejects_misaligned() {
5099 assert!(packed_rgba16f_layout(PixelFormat::PlanarRgb, DType::F16, 642, 640).is_none());
5100 }
5101
5102 #[test]
5103 fn packed_rgba16f_layout_rejects_non_f16() {
5104 // Non-F16 dtype with planar RGB
5105 assert!(packed_rgba16f_layout(PixelFormat::PlanarRgb, DType::U8, 640, 640).is_none());
5106 // Non-planar format with F32
5107 assert!(packed_rgba16f_layout(PixelFormat::Rgb, DType::F32, 640, 640).is_none());
5108 // Packed Rgba with F16 is not a planar format → None
5109 assert!(packed_rgba16f_layout(PixelFormat::Rgba, DType::F16, 640, 640).is_none());
5110 }
5111
5112 #[test]
5113 fn cuda_map_fast_fails_to_none_without_handle() {
5114 let t = Tensor::<f32>::new(&[4], Some(TensorMemory::Mem), None).unwrap();
5115 assert!(t.cuda().is_none());
5116 assert!(t.cuda_map().is_none()); // pure local check, no GL routing
5117 }
5118
5119 #[test]
5120 fn cuda_returns_none_without_handle() {
5121 // A plain Mem-backed tensor has no CUDA handle attached.
5122 let t = Tensor::<f32>::new(&[2, 2], Some(TensorMemory::Mem), None).unwrap();
5123 assert!(t.cuda().is_none(), "no CUDA handle on a Mem tensor");
5124 assert!(t.cuda_map().is_none(), "fast-fail map → None");
5125 }
5126
5127 #[test]
5128 fn cuda_map_then_host_map_fallback() {
5129 // The documented client pattern: try cuda_map() first; when it is None
5130 // (no CUDA handle — the case for a plain Mem tensor), fall back to map().
5131 let t = Tensor::<f32>::new(&[2, 2], Some(TensorMemory::Mem), None).unwrap();
5132 // Bind to a named variable so the CudaMap guard (and its borrow of `t`)
5133 // is dropped at the end of this statement, before the else branch borrows `t` again.
5134 let cuda = t.cuda_map();
5135 if let Some(_c) = cuda {
5136 // On a CUDA-registered tensor we'd use the device ptr here.
5137 unreachable!("a Mem tensor has no CUDA handle");
5138 } else {
5139 let host = t.map().expect("host map fallback must succeed");
5140 // TensorMapTrait::len() returns the element count (not bytes).
5141 assert_eq!(host.len(), 4); // 2*2 f32 elements
5142 }
5143 }
5144
5145 // -------------------------------------------------------------------------
5146 // Tensor::from_foreign — public API tests at the Tensor<T> layer.
5147 //
5148 // The low-level MemTensor::from_foreign mechanics (owner-drop, view sharing)
5149 // are covered in mem.rs. These tests exercise the Tensor<T> guard paths
5150 // (null ptr, empty shape, size overflow) and the basic wrap+readback
5151 // contract, confirming the public unsafe API wires through correctly.
5152 // -------------------------------------------------------------------------
5153
5154 #[test]
5155 fn from_foreign_valid_wrap_and_readback() {
5156 // The canonical CUDA zero-copy shape: wrap a caller allocation as a
5157 // Mem tensor and verify the tensor reads the exact same bytes.
5158 let mut buf: Vec<f32> = vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0];
5159 let ptr = buf.as_mut_ptr();
5160 let t = unsafe { Tensor::<f32>::from_foreign(ptr, &[2, 3], None, Some("test_foreign")) }
5161 .expect("valid from_foreign must succeed");
5162 assert_eq!(t.shape(), &[2, 3]);
5163 assert_eq!(t.memory(), TensorMemory::Mem);
5164 assert_eq!(t.name(), "test_foreign");
5165 let m = t.map().unwrap();
5166 // The tensor is a zero-copy borrow — it sees the caller's data.
5167 assert_eq!(m.as_slice(), &[1.0_f32, 2.0, 3.0, 4.0, 5.0, 6.0]);
5168 }
5169
5170 #[test]
5171 fn from_foreign_write_visible_in_caller_allocation() {
5172 // Writes through the tensor's map land in the caller's buffer (zero-copy).
5173 let mut buf: Vec<u8> = vec![0u8; 6];
5174 let ptr = buf.as_mut_ptr();
5175 let t =
5176 unsafe { Tensor::<u8>::from_foreign(ptr, &[2, 3], None, None) }.unwrap();
5177 {
5178 let mut m = t.map().unwrap();
5179 m.as_mut_slice().copy_from_slice(&[10, 20, 30, 40, 50, 60]);
5180 }
5181 drop(t);
5182 // Mutations are visible in the original Vec — same physical buffer.
5183 assert_eq!(buf, vec![10, 20, 30, 40, 50, 60]);
5184 }
5185
5186 #[test]
5187 fn from_foreign_rejects_null_ptr() {
5188 let err = unsafe {
5189 Tensor::<u8>::from_foreign(std::ptr::null_mut(), &[4], None, None)
5190 }
5191 .unwrap_err();
5192 assert!(
5193 matches!(err, Error::InvalidArgument(ref m) if m.contains("non-null")),
5194 "expected InvalidArgument(non-null), got {err:?}"
5195 );
5196 }
5197
5198 #[test]
5199 fn from_foreign_rejects_empty_shape() {
5200 let mut dummy: u8 = 0;
5201 let err = unsafe { Tensor::<u8>::from_foreign(&mut dummy, &[], None, None) }.unwrap_err();
5202 assert!(
5203 matches!(err, Error::InvalidSize(0)),
5204 "expected InvalidSize(0) for empty shape, got {err:?}"
5205 );
5206 }
5207
5208 #[test]
5209 fn from_foreign_rejects_overflow_shape() {
5210 // Two dimensions whose product overflows usize — the overflow guard must
5211 // fire before any pointer arithmetic is attempted.
5212 let mut dummy: u8 = 0;
5213 let huge = [usize::MAX / 2 + 1, 2];
5214 let err =
5215 unsafe { Tensor::<u8>::from_foreign(&mut dummy, &huge, None, None) }.unwrap_err();
5216 assert!(
5217 matches!(err, Error::InvalidArgument(ref m) if m.contains("overflow")),
5218 "expected InvalidArgument(overflow), got {err:?}"
5219 );
5220 }
5221
5222 #[test]
5223 fn from_foreign_owner_keeps_allocation_alive() {
5224 // When `owner` is `Some`, dropping the Tensor must not free the backing
5225 // before the owner is also gone — the owner's Drop fires on last ref.
5226 use std::sync::atomic::{AtomicBool, Ordering};
5227 let flag = std::sync::Arc::new(AtomicBool::new(false));
5228 let flag2 = flag.clone();
5229 struct Guard(std::sync::Arc<AtomicBool>);
5230 impl Drop for Guard {
5231 fn drop(&mut self) {
5232 self.0.store(true, Ordering::SeqCst);
5233 }
5234 }
5235 let mut buf: Vec<u32> = vec![42u32; 4];
5236 let ptr = buf.as_mut_ptr();
5237 let owner: ForeignOwner = Box::new(Guard(flag2));
5238 let t = unsafe { Tensor::<u32>::from_foreign(ptr, &[4], Some(owner), None) }.unwrap();
5239 // Map co-owns the backing Arc; the owner must stay alive while the map lives.
5240 let m = t.map().unwrap();
5241 assert_eq!(m.as_slice()[0], 42);
5242 drop(t); // tensor dropped while map is still live
5243 assert!(
5244 !flag.load(Ordering::SeqCst),
5245 "owner must not drop while a map shares the backing"
5246 );
5247 drop(m);
5248 assert!(
5249 flag.load(Ordering::SeqCst),
5250 "owner Drop must fire when the last Arc reference is released"
5251 );
5252 }
5253}