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)
8across four memory backends: the platform's native zero-copy GPU buffer, POSIX shared memory, the system
9heap, and an OpenGL Pixel Buffer Object. The crate defines traits and structures for creating, reshaping,
10and mapping tensors into memory.
11
12## Examples
13```rust
14use edgefirst_tensor::{Error, Tensor, TensorMemory, TensorTrait};
15# fn main() -> Result<(), Error> {
16let tensor = Tensor::<f32>::new(&[2, 3, 4], Some(TensorMemory::Mem), Some("test_tensor"))?;
17assert_eq!(tensor.memory(), TensorMemory::Mem);
18assert_eq!(tensor.name(), "test_tensor");
19# Ok(())
20# }
21```
22
23## Overview
24The main structures and traits provided by the `edgefirst_tensor` crate are [`TensorTrait`] and
25[`TensorMapTrait`], which define the behavior of Tensors and their memory mappings, respectively.
26The [`Tensor<T>`] struct wraps a backend-specific storage with optional image format metadata
27([`PixelFormat`]), while the [`TensorMap`] enum provides access to the underlying data. The
28[`TensorDyn`] type-erased enum wraps `Tensor<T>` for runtime element-type dispatch.
29
30[`TensorMemory::Dma`] is one variant with three implementations behind it — a Linux DMA-heap DMA-BUF,
31a macOS/iOS `IOSurfaceRef`, or an Android `AHardwareBuffer`. Callers ask for `Dma` and get whichever
32the platform provides, so portable code never branches on the mechanism.
33
34## Image tensors
35Images go through [`Tensor::image`] (or [`Tensor::image_desc`] for the full-featured request) rather
36than [`Tensor::new`]. Two things set them apart:
37
38- They take a required [`CpuAccess`] declaration. Hardware access — GPU, NPU, ISP, codec — is always
39 implied; what the CPU intends is the one thing the allocator cannot guess. Mapping beyond the
40 declaration is best-effort and always counted by [`unplanned_cpu_access_count`].
41- DMA-backed images carry a 64-byte-aligned row stride in every layout, because Mali and Vivante
42 reject an EGLImage import at an unaligned pitch. Read the real pitch from
43 [`Tensor::effective_row_stride`]; don't assume `width * bpp`.
44 */
45#[cfg(target_os = "android")]
46mod ahardwarebuffer;
47// Pure AHardwareBuffer layout logic (format table, descriptor geometry,
48// overflow-checked shape math) — cfg-free so it compiles and unit-tests
49// on every host; the android module above consumes it.
50#[allow(dead_code)]
51mod ahardwarebuffer_layout;
52pub mod colorimetry;
53pub mod covguard;
54mod cuda;
55#[cfg(target_os = "linux")]
56mod dma;
57#[cfg(target_os = "linux")]
58mod dmabuf;
59mod error;
60mod format;
61#[cfg(any(target_os = "macos", target_os = "ios"))]
62mod iosurface;
63mod mem;
64mod pbo;
65#[cfg(unix)]
66mod shm;
67mod tensor_dyn;
68pub use colorimetry::{
69 ColorEncoding, ColorRange, ColorSpace, ColorTransfer, Colorimetry, MatrixWeights, RangeScaling,
70};
71
72/// Retained constructor: installs the coverage flush-on-abort handler for this
73/// crate's instrumented test binary. See `covguard`. Only present under
74/// coverage on Linux (`.init_array` is ELF-only; the i.MX flush is Linux-only).
75#[cfg(all(coverage, target_os = "linux"))]
76#[used]
77#[link_section = ".init_array"]
78static __EDGEFIRST_COV_INSTALL: extern "C" fn() = {
79 extern "C" fn ctor() {
80 crate::covguard::install();
81 }
82 ctor
83};
84
85// Backing tensor/map types are internal implementation details: callers
86// allocate `Tensor<T>` / `TensorDyn` and map them, never naming the per-memory
87// backing types directly. They are `pub(crate)` so they stay nameable for the
88// `TensorStorage` / `TensorMap` enums without leaking into the public API.
89// Exceptions kept public: `Pbo*` is a GL extension point implemented by the
90// image crate, and `image_iosurface_layout` is a public helper.
91#[cfg(target_os = "android")]
92pub use crate::ahardwarebuffer::image_ahardwarebuffer_layout;
93#[cfg(target_os = "android")]
94pub(crate) use crate::ahardwarebuffer::{AHardwareBufferMap, AHardwareBufferTensor};
95#[cfg(target_os = "linux")]
96pub(crate) use crate::dma::{DmaMap, DmaTensor};
97#[cfg(any(target_os = "macos", target_os = "ios"))]
98pub use crate::iosurface::image_iosurface_layout;
99#[cfg(any(target_os = "macos", target_os = "ios"))]
100pub(crate) use crate::iosurface::{IoSurfaceMap, IoSurfaceTensor};
101pub(crate) use crate::mem::{MemMap, MemTensor};
102pub use crate::pbo::{PboMap, PboMapping, PboOps, PboTensor};
103#[cfg(unix)]
104pub(crate) use crate::shm::{ShmMap, ShmTensor};
105pub use cuda::{
106 gl_map_resource, gl_register_buffer, gl_unmap_resource, gl_unregister_resource,
107 is_cuda_available, memcpy_device_to_host, stream_create, stream_destroy, stream_synchronize,
108 CudaGlOps, CudaHandle, CudaMap, CudaStream,
109};
110pub use error::{Error, Result};
111pub use format::{ChromaLayout, PixelFormat, PixelLayout};
112use num_traits::Num;
113use serde::{Deserialize, Serialize};
114#[cfg(unix)]
115use std::os::fd::OwnedFd;
116use std::{
117 fmt,
118 ops::{Deref, DerefMut},
119 sync::{
120 atomic::{AtomicU64, Ordering},
121 Arc, Weak,
122 },
123};
124pub use tensor_dyn::TensorDyn;
125
126/// Opaque keep-alive handle for a foreign-memory tensor (see
127/// [`Tensor::from_foreign`] / [`TensorDyn::from_foreign_ptr`]).
128///
129/// The HAL borrows the foreign buffer without owning it; this handle co-owns
130/// the *source* so the borrowed memory stays valid for the tensor's life. Its
131/// `Drop` releases the source — e.g. a small struct that calls `cudaFreeHost`,
132/// or a `Py<PyAny>` that decrements a NumPy array's refcount. Wrapping it in an
133/// `Arc` (then boxing each clone) makes the release fire exactly once, after
134/// the last sharing tensor/view/map drops, regardless of drop order.
135pub type ForeignOwner = Box<dyn std::any::Any + Send + Sync>;
136
137/// Re-export of `half::f16` so downstream crates can write
138/// `Tensor::<edgefirst_tensor::f16>::from_iosurface(…)` without
139/// adding `half` to their own dependency list. The version stays in
140/// lockstep with the `half` workspace dep.
141pub use half::f16;
142
143// =============================================================================
144// RGBA16F packed-layout geometry — single source of truth
145//
146// A `PlanarRgb` [3,H,W] or `PlanarRgba` [4,H,W] f16 tensor is represented
147// on the GPU as an RGBA16F surface (the only float format accepted by the
148// ANGLE IOSurface extension). Four contiguous f16 elements are packed into
149// each 8-byte RGBA16F texel, yielding a `(W/4, C*H)` surface.
150//
151// All call sites that need these dimensions must use `packed_rgba16f_layout`
152// so the rule lives in exactly one place. Currently consumed by:
153// - `crates/tensor/src/iosurface.rs` `new_image` (macOS IOSurface alloc)
154// - `crates/image/src/gl/iosurface_import.rs` (macOS GL IOSurface import)
155// - `crates/image/src/gl/processor/float.rs` (Linux GL float render — PBO
156// readback and DMA-BUF, also via the `dma_f16_packed_layout` wrapper)
157// =============================================================================
158
159/// Geometry of the RGBA16F-packed surface backing a planar F16 image tensor.
160///
161/// ANGLE only supports one float `(type, internal_format)` pair for IOSurface
162/// import: `(GL_HALF_FLOAT, GL_RGBA)` = RGBA16F (8 bytes/texel). To map a
163/// `[C, H, W]` f16 planar tensor onto such a surface, 4 contiguous f16
164/// elements are packed into each RGBA16F texel, yielding a surface of
165/// `(W/4, C*H)` texels at 8 bytes/texel. The byte stream is identical to a
166/// (nonexistent) R16F `(W, C*H)` surface and can be consumed as `&[f16]`
167/// with shape `[1, C, H, W]` without rearrangement.
168///
169/// Obtain via [`packed_rgba16f_layout`] — never construct directly.
170#[derive(Debug, Clone, Copy, PartialEq, Eq)]
171pub struct PackedRgba16fLayout {
172 /// Surface width in texels (`width / 4`).
173 pub surface_w: usize,
174 /// Surface height in texels (`planes * height`).
175 pub surface_h: usize,
176 /// Bytes per RGBA16F texel (always 8).
177 pub bytes_per_texel: usize,
178 /// Row pitch in bytes (`surface_w * 8`).
179 pub pitch: usize,
180}
181
182/// Canonical geometry for the RGBA16F-packed surface backing a planar F16
183/// image tensor.
184///
185/// Returns `Some(layout)` only when **all** of the following hold:
186///
187/// - `dtype == DType::F16`
188/// - `format` is `PixelFormat::PlanarRgb` (3 planes) or
189/// `PixelFormat::PlanarRgba` (4 planes)
190/// - `width % 4 == 0`
191///
192/// Returns `None` for any other `(format, dtype)` combination, misaligned
193/// width, or when the surface geometry would overflow `usize` — callers
194/// must fall back to a non-packed path or return a context-appropriate
195/// error.
196///
197/// # Examples
198///
199/// ```rust
200/// use edgefirst_tensor::{packed_rgba16f_layout, PixelFormat, DType};
201///
202/// let layout = packed_rgba16f_layout(PixelFormat::PlanarRgb, DType::F16, 640, 480).unwrap();
203/// assert_eq!(layout.surface_w, 160);
204/// assert_eq!(layout.surface_h, 1440);
205/// assert_eq!(layout.bytes_per_texel, 8);
206/// assert_eq!(layout.pitch, 1280);
207/// ```
208pub fn packed_rgba16f_layout(
209 format: PixelFormat,
210 dtype: DType,
211 width: usize,
212 height: usize,
213) -> Option<PackedRgba16fLayout> {
214 if dtype != DType::F16 {
215 return None;
216 }
217 let planes: usize = match format {
218 PixelFormat::PlanarRgb => 3,
219 PixelFormat::PlanarRgba => 4,
220 _ => return None,
221 };
222 if !width.is_multiple_of(4) {
223 return None;
224 }
225 let surface_w = width / 4;
226 // Checked arithmetic: a degenerate (height, width) could otherwise wrap
227 // and yield an under-sized layout, which downstream allocators trust for
228 // GPU/CPU buffer sizing. Overflow → None (handled like any other
229 // unsupported geometry).
230 let surface_h = planes.checked_mul(height)?;
231 let bytes_per_texel = 8;
232 let pitch = surface_w.checked_mul(bytes_per_texel)?;
233 Some(PackedRgba16fLayout {
234 surface_w,
235 surface_h,
236 bytes_per_texel,
237 pitch,
238 })
239}
240
241/// Geometry of the RGBA8888-packed surface backing a packed RGB u8/i8 image
242/// tensor.
243///
244/// GPUs have no 3-channel renderable format, so the GL engine's two-pass
245/// packed-RGB shader writes the tight `[H, W, 3]` byte stream into an
246/// RGBA8888 surface: each texel carries 4 consecutive RGB bytes, giving a
247/// `(W*3/4, H)` surface at 4 bytes/texel whose rows are byte-identical to
248/// tight RGB — consumable flat as `[H, W, 3]` with no rearrangement (the
249/// u8/i8 analog of [`packed_rgba16f_layout`]; i8 shares the layout since
250/// INT8 quantization is a per-byte `^0x80` bias, not a format change).
251///
252/// Returns `Some(layout)` only when `width % 4 == 0` (so `W*3` bytes divide
253/// into whole texels) and the geometry does not overflow `usize`.
254///
255/// # Examples
256///
257/// ```rust
258/// use edgefirst_tensor::packed_rgb888_layout;
259///
260/// let layout = packed_rgb888_layout(640, 480).unwrap();
261/// assert_eq!(layout.surface_w, 480); // 640*3/4
262/// assert_eq!(layout.surface_h, 480);
263/// assert_eq!(layout.bytes_per_texel, 4);
264/// assert_eq!(layout.pitch, 1920); // 640*3
265/// assert!(packed_rgb888_layout(641, 480).is_none());
266/// ```
267pub fn packed_rgb888_layout(width: usize, height: usize) -> Option<PackedRgb888Layout> {
268 if !width.is_multiple_of(4) {
269 return None;
270 }
271 let row_bytes = width.checked_mul(3)?;
272 let surface_w = row_bytes / 4;
273 Some(PackedRgb888Layout {
274 surface_w,
275 surface_h: height,
276 bytes_per_texel: 4,
277 pitch: row_bytes,
278 })
279}
280
281/// Geometry of the RGBA8888 surface backing a packed RGB u8/i8 image tensor.
282///
283/// Obtain via [`packed_rgb888_layout`] — never construct directly.
284#[derive(Debug, Clone, Copy, PartialEq, Eq)]
285pub struct PackedRgb888Layout {
286 /// Surface width in texels (`width * 3 / 4`).
287 pub surface_w: usize,
288 /// Surface height in texels (`height`).
289 pub surface_h: usize,
290 /// Bytes per RGBA8888 texel (always 4).
291 pub bytes_per_texel: usize,
292 /// Row pitch in bytes (`surface_w * 4` = `width * 3`).
293 pub pitch: usize,
294}
295
296/// Per-plane DMA-BUF descriptor for external buffer import.
297///
298/// Owns a duplicated file descriptor plus optional stride and offset metadata.
299/// The fd is duplicated eagerly in [`new()`](Self::new) so that a bad fd is
300/// caught immediately. `import_image` consumes the descriptor and takes
301/// ownership of the duped fd — no further cleanup is needed by the caller.
302///
303/// # Examples
304///
305/// ```rust,no_run
306/// use edgefirst_tensor::PlaneDescriptor;
307/// use std::os::fd::BorrowedFd;
308///
309/// // SAFETY: fd 42 is hypothetical; real code must pass a valid fd.
310/// let pd = unsafe { PlaneDescriptor::new(BorrowedFd::borrow_raw(42)) }
311/// .unwrap()
312/// .with_stride(2048)
313/// .with_offset(0);
314/// ```
315#[cfg(unix)]
316pub struct PlaneDescriptor {
317 fd: OwnedFd,
318 stride: Option<usize>,
319 offset: Option<usize>,
320}
321
322#[cfg(unix)]
323impl PlaneDescriptor {
324 /// Create a new plane descriptor by duplicating the given file descriptor.
325 ///
326 /// The fd is duped immediately — a bad fd fails here rather than inside
327 /// `import_image`. The caller retains ownership of the original fd.
328 ///
329 /// # Errors
330 ///
331 /// Returns an error if the `dup()` syscall fails (e.g. invalid fd or
332 /// fd limit reached).
333 pub fn new(fd: std::os::fd::BorrowedFd<'_>) -> Result<Self> {
334 let owned = fd.try_clone_to_owned()?;
335 Ok(Self {
336 fd: owned,
337 stride: None,
338 offset: None,
339 })
340 }
341
342 /// Set the row stride in bytes (consuming builder).
343 pub fn with_stride(mut self, stride: usize) -> Self {
344 self.stride = Some(stride);
345 self
346 }
347
348 /// Set the plane offset in bytes (consuming builder).
349 pub fn with_offset(mut self, offset: usize) -> Self {
350 self.offset = Some(offset);
351 self
352 }
353
354 /// Consume the descriptor and return the owned file descriptor.
355 pub fn into_fd(self) -> OwnedFd {
356 self.fd
357 }
358
359 /// Row stride in bytes, if set.
360 pub fn stride(&self) -> Option<usize> {
361 self.stride
362 }
363
364 /// Plane offset in bytes, if set.
365 pub fn offset(&self) -> Option<usize> {
366 self.offset
367 }
368}
369
370/// A rectangular sub-region of a tensor's leading spatial frame, in pixels.
371///
372/// `Region` is the single rectangle type in the workspace: the argument to
373/// [`Tensor::view`], the source sampling window in the image crate's `Crop`,
374/// and the geometry the image backend lowers to a `glViewport` (a destination
375/// tile) or a sampling rectangle (a source). Coordinates are pixel/element
376/// units of the leading spatial axes; byte addressing is derived from the
377/// parent's row stride, not stored here.
378#[derive(Debug, Clone, Copy, PartialEq, Eq)]
379pub struct Region {
380 pub x: usize,
381 pub y: usize,
382 pub width: usize,
383 pub height: usize,
384}
385
386impl Region {
387 /// Create a region at `(x, y)` spanning `width` × `height` pixels.
388 pub fn new(x: usize, y: usize, width: usize, height: usize) -> Self {
389 Self {
390 x,
391 y,
392 width,
393 height,
394 }
395 }
396
397 /// True when the region lies fully within a `width` × `height` frame.
398 pub fn fits_within(&self, width: usize, height: usize) -> bool {
399 self.x.saturating_add(self.width) <= width && self.y.saturating_add(self.height) <= height
400 }
401}
402
403/// The parent image a [`view`](Tensor::view)/[`batch`](Tensor::batch) sub-region
404/// was carved from, snapshotted at the time the view was created.
405///
406/// A view shares the parent's `BufferIdentity` and addresses a sub-rectangle of
407/// it. The GL backend keys its EGLImage import on the **parent** geometry (so all
408/// sibling views of one buffer collapse to a single import) and renders each
409/// view as a `glViewport`+`glScissor` ROI at `(x, y, width, height)` within that
410/// parent — the view is render state, never a distinct import. `parent_width`/
411/// `parent_height` are the parent's logical pixel dimensions; `x`/`y` are this
412/// view's top-left origin within the parent (pixels). Nested views compose:
413/// the snapshot always names the **root** parent, with offsets accumulated.
414#[derive(Debug, Clone, Copy, PartialEq, Eq)]
415pub struct ViewOrigin {
416 /// Logical width of the root parent image, in pixels.
417 pub parent_width: usize,
418 /// Logical height of the root parent image, in pixels. For a `batch(n)` view
419 /// of an `[N, H, W, C]` tensor this is `N * H` (the tiles stack vertically in
420 /// the shared buffer).
421 pub parent_height: usize,
422 /// The parent's row stride in **bytes**. The GL backend keys its EGLImage
423 /// import/cache and pitch on this — NOT on the view's own `row_stride`, which
424 /// a single-row view sets tight (for map-span safety). Using the parent
425 /// stride keeps the import pitch parent-consistent so single-row and
426 /// multi-row sibling views collapse onto the same parent import.
427 pub parent_row_stride: usize,
428 /// This view's top-left x origin within the root parent, in pixels.
429 pub x: usize,
430 /// This view's top-left y origin within the root parent, in pixels.
431 pub y: usize,
432}
433
434/// Element type discriminant for runtime type identification.
435#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
436#[repr(u8)]
437#[non_exhaustive]
438pub enum DType {
439 U8,
440 I8,
441 U16,
442 I16,
443 U32,
444 I32,
445 U64,
446 I64,
447 F16,
448 F32,
449 F64,
450}
451
452impl DType {
453 /// Size of one element in bytes.
454 pub const fn size(&self) -> usize {
455 match self {
456 Self::U8 | Self::I8 => 1,
457 Self::U16 | Self::I16 | Self::F16 => 2,
458 Self::U32 | Self::I32 | Self::F32 => 4,
459 Self::U64 | Self::I64 | Self::F64 => 8,
460 }
461 }
462
463 /// Short type name (e.g., "u8", "f32", "f16").
464 pub const fn name(&self) -> &'static str {
465 match self {
466 Self::U8 => "u8",
467 Self::I8 => "i8",
468 Self::U16 => "u16",
469 Self::I16 => "i16",
470 Self::U32 => "u32",
471 Self::I32 => "i32",
472 Self::U64 => "u64",
473 Self::I64 => "i64",
474 Self::F16 => "f16",
475 Self::F32 => "f32",
476 Self::F64 => "f64",
477 }
478 }
479}
480
481impl fmt::Display for DType {
482 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
483 f.write_str(self.name())
484 }
485}
486
487/// Map a static numeric type `T` to its `DType` discriminant, returning
488/// `None` for types that do not have a `DType` representation (e.g.
489/// user-defined wrappers in tests).
490///
491/// Runtime dtype of a static `Tensor<T>` element type — a safe `TypeId`
492/// allowlist of HAL's primitive numeric types (`Some` for `u8`..`i64` /
493/// `f16` / `f32` / `f64`, `None` otherwise). Used by the macOS IOSurface
494/// image constructors (FourCC / pixel-format lookup) and by the `Mem`
495/// backing's `alloc_zeroed` fast path (these types' `T::zero()` is the
496/// all-zeros bit pattern), so it is compiled on every target.
497pub(crate) fn dtype_of<T: 'static>() -> Option<DType> {
498 use std::any::TypeId;
499 let id = TypeId::of::<T>();
500 if id == TypeId::of::<u8>() {
501 Some(DType::U8)
502 } else if id == TypeId::of::<i8>() {
503 Some(DType::I8)
504 } else if id == TypeId::of::<u16>() {
505 Some(DType::U16)
506 } else if id == TypeId::of::<i16>() {
507 Some(DType::I16)
508 } else if id == TypeId::of::<u32>() {
509 Some(DType::U32)
510 } else if id == TypeId::of::<i32>() {
511 Some(DType::I32)
512 } else if id == TypeId::of::<u64>() {
513 Some(DType::U64)
514 } else if id == TypeId::of::<i64>() {
515 Some(DType::I64)
516 } else if id == TypeId::of::<half::f16>() {
517 Some(DType::F16)
518 } else if id == TypeId::of::<f32>() {
519 Some(DType::F32)
520 } else if id == TypeId::of::<f64>() {
521 Some(DType::F64)
522 } else {
523 None
524 }
525}
526
527// =============================================================================
528// Quantization metadata — type-gated to integer element types via sealed
529// `IntegerType` trait. Accessors on `Tensor<T>` only compile when `T` is
530// an integer type; calling them on `Tensor<f32>` / `Tensor<f16>` etc. is a
531// compile error, not a runtime one.
532// =============================================================================
533
534mod sealed {
535 pub trait Sealed {}
536 impl Sealed for u8 {}
537 impl Sealed for i8 {}
538 impl Sealed for u16 {}
539 impl Sealed for i16 {}
540 impl Sealed for u32 {}
541 impl Sealed for i32 {}
542 impl Sealed for u64 {}
543 impl Sealed for i64 {}
544 // Deliberately NOT implemented for f16 / f32 / f64.
545}
546
547/// Integer element types that may carry quantization metadata.
548///
549/// Sealed trait: implemented for `u8`, `i8`, `u16`, `i16`, `u32`, `i32`,
550/// `u64`, `i64`. Cannot be implemented downstream. Float element types
551/// (`half::f16`, `f32`, `f64`) are explicitly excluded — quantization
552/// metadata does not apply to float tensors per the edgefirst.json spec.
553pub trait IntegerType: sealed::Sealed {}
554impl IntegerType for u8 {}
555impl IntegerType for i8 {}
556impl IntegerType for u16 {}
557impl IntegerType for i16 {}
558impl IntegerType for u32 {}
559impl IntegerType for i32 {}
560impl IntegerType for u64 {}
561impl IntegerType for i64 {}
562
563/// Quantization parameters for an integer tensor.
564///
565/// Covers all four modes the edgefirst.json spec defines:
566///
567/// | Mode | `scale.len()` | `zero_point` | `axis` |
568/// |---|---|---|---|
569/// | Per-tensor symmetric | 1 | `None` | `None` |
570/// | Per-tensor asymmetric | 1 | `Some(len == 1)` | `None` |
571/// | Per-channel symmetric | >1 | `None` | `Some(c)` |
572/// | Per-channel asymmetric | >1 | `Some(len == scale.len())` | `Some(c)` |
573///
574/// The quantized storage type is carried on the parent [`Tensor<T>`]; this
575/// struct does not duplicate it. Construct via the four named constructors
576/// (the only public entry points); direct field mutation is not allowed so
577/// invalid combinations cannot be represented.
578///
579/// Dequantization formula:
580///
581/// ```text
582/// real_value = scale[c] × (quantized_value[c] - zero_point[c])
583/// ```
584///
585/// where `c` is the channel index (always `0` for per-tensor).
586#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
587pub struct Quantization {
588 /// Per-tensor: `vec![scale]`. Per-channel: `vec![scale_0, scale_1, ...]`.
589 #[serde(deserialize_with = "deserialize_scalar_or_vec_f32")]
590 scale: Vec<f32>,
591
592 /// `None` means symmetric (zero-point is 0). `Some(vec)` must have the
593 /// same length as `scale`.
594 #[serde(
595 default,
596 deserialize_with = "deserialize_opt_scalar_or_vec_i32",
597 skip_serializing_if = "Option::is_none"
598 )]
599 zero_point: Option<Vec<i32>>,
600
601 /// Channel axis for per-channel quantization. `Some(_)` iff
602 /// `scale.len() > 1`. Validated against the parent tensor's shape at
603 /// `set_quantization()` time.
604 #[serde(default, skip_serializing_if = "Option::is_none")]
605 axis: Option<usize>,
606}
607
608/// Semantic mode discriminant for hot-path kernel dispatch.
609///
610/// Obtain via [`Quantization::mode`] once at kernel entry; never inside a
611/// pixel-level loop. The enum is borrow-based so the hot kernel receives
612/// the scales / zero-points as slices without reallocation.
613#[derive(Debug, Clone, Copy)]
614pub enum QuantMode<'a> {
615 PerTensorSymmetric {
616 scale: f32,
617 },
618 PerTensor {
619 scale: f32,
620 zero_point: i32,
621 },
622 PerChannelSymmetric {
623 scales: &'a [f32],
624 axis: usize,
625 },
626 PerChannel {
627 scales: &'a [f32],
628 zero_points: &'a [i32],
629 axis: usize,
630 },
631}
632
633impl Quantization {
634 /// Per-tensor symmetric (zero_point = 0).
635 pub fn per_tensor_symmetric(scale: f32) -> Self {
636 Self {
637 scale: vec![scale],
638 zero_point: None,
639 axis: None,
640 }
641 }
642
643 /// Per-tensor asymmetric — the most common runtime shape.
644 pub fn per_tensor(scale: f32, zero_point: i32) -> Self {
645 Self {
646 scale: vec![scale],
647 zero_point: Some(vec![zero_point]),
648 axis: None,
649 }
650 }
651
652 /// Per-channel symmetric. Errors on empty `scales`.
653 pub fn per_channel_symmetric(scales: Vec<f32>, axis: usize) -> Result<Self> {
654 if scales.is_empty() {
655 return Err(Error::QuantizationInvalid {
656 field: "scale.len",
657 expected: "non-empty per-channel scales".to_string(),
658 got: "length 0".to_string(),
659 });
660 }
661 Ok(Self {
662 scale: scales,
663 zero_point: None,
664 axis: Some(axis),
665 })
666 }
667
668 /// Per-channel asymmetric. Errors on length mismatch between `scales`
669 /// and `zero_points`, or empty arrays.
670 pub fn per_channel(scales: Vec<f32>, zero_points: Vec<i32>, axis: usize) -> Result<Self> {
671 if scales.is_empty() {
672 return Err(Error::QuantizationInvalid {
673 field: "scale.len",
674 expected: "non-empty per-channel scales".to_string(),
675 got: "length 0".to_string(),
676 });
677 }
678 if scales.len() != zero_points.len() {
679 return Err(Error::QuantizationInvalid {
680 field: "zero_point.len",
681 expected: format!("length matches scale ({})", scales.len()),
682 got: format!("length {}", zero_points.len()),
683 });
684 }
685 Ok(Self {
686 scale: scales,
687 zero_point: Some(zero_points),
688 axis: Some(axis),
689 })
690 }
691
692 /// Borrow-based dispatch view. Match once at kernel entry.
693 pub fn mode(&self) -> QuantMode<'_> {
694 match (self.scale.len(), self.zero_point.as_deref(), self.axis) {
695 (1, None, _) => QuantMode::PerTensorSymmetric {
696 scale: self.scale[0],
697 },
698 (1, Some(zps), _) => QuantMode::PerTensor {
699 scale: self.scale[0],
700 zero_point: zps.first().copied().unwrap_or(0),
701 },
702 (_, None, Some(axis)) => QuantMode::PerChannelSymmetric {
703 scales: &self.scale,
704 axis,
705 },
706 (_, Some(zps), Some(axis)) => QuantMode::PerChannel {
707 scales: &self.scale,
708 zero_points: zps,
709 axis,
710 },
711 // The `validate()` path prevents constructing a
712 // per-channel Quantization without an axis, so the remaining
713 // pattern is unreachable in practice. Fall back to
714 // per-tensor symmetric using scale[0] to avoid panicking in
715 // release; debug builds assert.
716 _ => {
717 debug_assert!(
718 false,
719 "Quantization::mode: per-channel without axis is unreachable"
720 );
721 QuantMode::PerTensorSymmetric {
722 scale: self.scale.first().copied().unwrap_or(1.0),
723 }
724 }
725 }
726 }
727
728 /// Returns `true` for per-tensor quantization (`scale.len() == 1`).
729 pub fn is_per_tensor(&self) -> bool {
730 self.scale.len() == 1
731 }
732
733 /// Returns `true` for per-channel quantization (`scale.len() > 1`).
734 pub fn is_per_channel(&self) -> bool {
735 self.scale.len() > 1
736 }
737
738 /// Returns `true` for symmetric quantization (no zero-point, or
739 /// zero-point vector of all zeros).
740 pub fn is_symmetric(&self) -> bool {
741 match &self.zero_point {
742 None => true,
743 Some(zps) => zps.iter().all(|&z| z == 0),
744 }
745 }
746
747 /// Borrow the scale array. Length 1 for per-tensor; `num_channels` for
748 /// per-channel.
749 pub fn scale(&self) -> &[f32] {
750 &self.scale
751 }
752
753 /// Borrow the zero-point array. `None` for symmetric.
754 pub fn zero_point(&self) -> Option<&[i32]> {
755 self.zero_point.as_deref()
756 }
757
758 /// Channel axis for per-channel quantization. `None` for per-tensor.
759 pub fn axis(&self) -> Option<usize> {
760 self.axis
761 }
762
763 /// Validate against a target tensor shape. Runs in
764 /// `Tensor::set_quantization()`. Catches:
765 /// - empty `scale` (reject — must declare at least one factor)
766 /// - `zero_point` length inconsistent with `scale` (reject —
767 /// per-tensor must have len 1, per-channel must match `scale.len`)
768 /// - `axis >= shape.len()` (axis out of range)
769 /// - `scale.len() != shape[axis]` for per-channel
770 /// - per-channel without axis (reject)
771 /// - per-tensor with redundant axis (reject)
772 pub(crate) fn validate(&self, shape: &[usize]) -> Result<()> {
773 // `Quantization` is `Deserialize`, so malformed JSON like
774 // `{"scale": [], "zero_point": []}` could otherwise produce an
775 // ill-defined value that confuses `mode()` selection and the
776 // per-channel kernels' indexing.
777 if self.scale.is_empty() {
778 return Err(Error::QuantizationInvalid {
779 field: "scale.len",
780 expected: ">= 1".to_string(),
781 got: "0".to_string(),
782 });
783 }
784 if let Some(zps) = self.zero_point.as_ref() {
785 // Per-tensor: scale.len() == 1 and zero_point.len() must == 1.
786 // Per-channel: zero_point.len() must == scale.len().
787 let expected = if self.scale.len() == 1 {
788 1
789 } else {
790 self.scale.len()
791 };
792 if zps.len() != expected {
793 return Err(Error::QuantizationInvalid {
794 field: "zero_point.len",
795 expected: format!(
796 "{expected} (matching {})",
797 if self.scale.len() == 1 {
798 "per-tensor scale"
799 } else {
800 "per-channel scale.len"
801 }
802 ),
803 got: format!("length {}", zps.len()),
804 });
805 }
806 }
807
808 match (self.scale.len(), self.axis) {
809 (1, None) => Ok(()),
810 (1, Some(_)) => Err(Error::QuantizationInvalid {
811 field: "per_tensor_redundant_axis",
812 expected: "axis=None for per-tensor quantization".to_string(),
813 got: format!("axis={:?}", self.axis),
814 }),
815 (_, None) => Err(Error::QuantizationInvalid {
816 field: "per_channel_requires_axis",
817 expected: format!(
818 "axis=Some(_) for per-channel quantization (scale.len={})",
819 self.scale.len()
820 ),
821 got: "axis=None".to_string(),
822 }),
823 (n, Some(axis)) => {
824 if axis >= shape.len() {
825 return Err(Error::QuantizationInvalid {
826 field: "axis",
827 expected: format!("axis < tensor rank ({})", shape.len()),
828 got: format!("axis={axis}"),
829 });
830 }
831 if shape[axis] != n {
832 return Err(Error::QuantizationInvalid {
833 field: "scale.len",
834 expected: format!("length matches shape[{axis}] ({})", shape[axis]),
835 got: format!("length {n}"),
836 });
837 }
838 Ok(())
839 }
840 }
841 }
842}
843
844impl From<(f32, i32)> for Quantization {
845 /// Convenience construction from a `(scale, zero_point)` tuple. Matches
846 /// the legacy `QuantTuple` / `Quantization::new` calling convention so
847 /// existing `(0.1, -128).into()` sites keep working.
848 fn from((scale, zero_point): (f32, i32)) -> Self {
849 Self::per_tensor(scale, zero_point)
850 }
851}
852
853fn deserialize_scalar_or_vec_f32<'de, D: serde::Deserializer<'de>>(
854 de: D,
855) -> std::result::Result<Vec<f32>, D::Error> {
856 use serde::de::{self, Visitor};
857 struct V;
858 impl<'de> Visitor<'de> for V {
859 type Value = Vec<f32>;
860 fn expecting(&self, f: &mut fmt::Formatter) -> fmt::Result {
861 f.write_str("f32 or array of f32")
862 }
863 fn visit_f64<E: de::Error>(self, v: f64) -> std::result::Result<Self::Value, E> {
864 Ok(vec![v as f32])
865 }
866 #[allow(clippy::cast_possible_truncation)]
867 fn visit_i64<E: de::Error>(self, v: i64) -> std::result::Result<Self::Value, E> {
868 Ok(vec![v as f32])
869 }
870 #[allow(clippy::cast_possible_truncation, clippy::cast_precision_loss)]
871 fn visit_u64<E: de::Error>(self, v: u64) -> std::result::Result<Self::Value, E> {
872 Ok(vec![v as f32])
873 }
874 fn visit_seq<A: de::SeqAccess<'de>>(
875 self,
876 mut seq: A,
877 ) -> std::result::Result<Self::Value, A::Error> {
878 let mut out = Vec::with_capacity(seq.size_hint().unwrap_or(1));
879 while let Some(x) = seq.next_element::<f32>()? {
880 out.push(x);
881 }
882 Ok(out)
883 }
884 }
885 de.deserialize_any(V)
886}
887
888fn deserialize_opt_scalar_or_vec_i32<'de, D: serde::Deserializer<'de>>(
889 de: D,
890) -> std::result::Result<Option<Vec<i32>>, D::Error> {
891 use serde::de::{self, Visitor};
892 struct V;
893 impl<'de> Visitor<'de> for V {
894 type Value = Option<Vec<i32>>;
895 fn expecting(&self, f: &mut fmt::Formatter) -> fmt::Result {
896 f.write_str("null, i32, or array of i32")
897 }
898 fn visit_none<E: de::Error>(self) -> std::result::Result<Self::Value, E> {
899 Ok(None)
900 }
901 fn visit_unit<E: de::Error>(self) -> std::result::Result<Self::Value, E> {
902 Ok(None)
903 }
904 fn visit_some<D2: serde::Deserializer<'de>>(
905 self,
906 de: D2,
907 ) -> std::result::Result<Self::Value, D2::Error> {
908 struct Inner;
909 impl<'de> Visitor<'de> for Inner {
910 type Value = Vec<i32>;
911 fn expecting(&self, f: &mut fmt::Formatter) -> fmt::Result {
912 f.write_str("i32 or array of i32")
913 }
914 #[allow(clippy::cast_possible_truncation)]
915 fn visit_i64<E: de::Error>(self, v: i64) -> std::result::Result<Self::Value, E> {
916 Ok(vec![v as i32])
917 }
918 #[allow(clippy::cast_possible_truncation, clippy::cast_possible_wrap)]
919 fn visit_u64<E: de::Error>(self, v: u64) -> std::result::Result<Self::Value, E> {
920 Ok(vec![v as i32])
921 }
922 fn visit_seq<A: de::SeqAccess<'de>>(
923 self,
924 mut seq: A,
925 ) -> std::result::Result<Self::Value, A::Error> {
926 let mut out = Vec::with_capacity(seq.size_hint().unwrap_or(1));
927 while let Some(x) = seq.next_element::<i32>()? {
928 out.push(x);
929 }
930 Ok(out)
931 }
932 }
933 de.deserialize_any(Inner).map(Some)
934 }
935 #[allow(clippy::cast_possible_truncation)]
936 fn visit_i64<E: de::Error>(self, v: i64) -> std::result::Result<Self::Value, E> {
937 Ok(Some(vec![v as i32]))
938 }
939 #[allow(clippy::cast_possible_truncation, clippy::cast_possible_wrap)]
940 fn visit_u64<E: de::Error>(self, v: u64) -> std::result::Result<Self::Value, E> {
941 Ok(Some(vec![v as i32]))
942 }
943 fn visit_seq<A: de::SeqAccess<'de>>(
944 self,
945 mut seq: A,
946 ) -> std::result::Result<Self::Value, A::Error> {
947 let mut out = Vec::with_capacity(seq.size_hint().unwrap_or(1));
948 while let Some(x) = seq.next_element::<i32>()? {
949 out.push(x);
950 }
951 Ok(Some(out))
952 }
953 }
954 de.deserialize_option(V)
955}
956
957/// Monotonic counter for buffer identity IDs.
958static NEXT_BUFFER_ID: AtomicU64 = AtomicU64::new(1);
959
960/// Count of tensor maps requested beyond the buffer's declared
961/// [`CpuAccess`] (including any CPU map of a `CpuAccess::None` buffer).
962/// Undeclared CPU access is a pipeline smell: it forfeits layout
963/// optimizations (write-combined mappings, tile compression) and may be
964/// slow or refused on Android. Read via [`unplanned_cpu_access_count`];
965/// each offending buffer also logs one warning.
966static UNPLANNED_CPU_ACCESS: AtomicU64 = AtomicU64::new(0);
967
968/// Number of tensor maps that exceeded the buffer's declared
969/// [`CpuAccess`] since process start. A pipeline that declares its CPU
970/// access correctly holds this flat; see [`CpuAccess`] for the contract.
971pub fn unplanned_cpu_access_count() -> u64 {
972 UNPLANNED_CPU_ACCESS.load(Ordering::Relaxed)
973}
974
975/// Record an unplanned CPU access on `identity_id`, warning once per
976/// buffer (a steady-state pipeline maps the same buffer every frame — a
977/// per-map warn would flood the log; repeats count silently).
978pub(crate) fn note_unplanned_cpu_access(identity_id: u64, backend: &str, detail: &str) {
979 UNPLANNED_CPU_ACCESS.fetch_add(1, Ordering::Relaxed);
980 static WARNED: std::sync::OnceLock<std::sync::Mutex<std::collections::HashSet<u64>>> =
981 std::sync::OnceLock::new();
982 let warned = WARNED.get_or_init(Default::default);
983 if warned.lock().is_ok_and(|mut s| s.insert(identity_id)) {
984 log::warn!(
985 "unplanned CPU access on {backend} buffer (identity {identity_id}): {detail} — \
986 declare the intent at allocation (CpuAccess::Read/Write/ReadWrite) to make \
987 this a planned, optimized mapping"
988 );
989 }
990}
991
992/// Uniform guard for mutable access through a read-only map. Every
993/// `TensorMap` backend calls this from `as_mut_slice`/`deref_mut` so a
994/// `map_read()` misuse fails identically on all platforms instead of
995/// passing on tolerant ones and exploding only on Android.
996#[inline]
997pub(crate) fn assert_map_writable(writable: bool, backend: &str) {
998 assert!(
999 writable,
1000 "{backend} map is read-only (obtained via map_read()/CpuAccess::Read) — \
1001 use map_mut() or map_write() for mutable access"
1002 );
1003}
1004
1005/// Declared CPU involvement for an image tensor, chosen at allocation.
1006///
1007/// The HAL assumes buffers are produced and consumed by hardware (ISP,
1008/// codec, GPU, NPU) — hardware access needs no declaration. CPU access is
1009/// the opt-in: it selects the CPU usage/mapping mode at allocation
1010/// (write-combined for `Write`, cached for `Read`) and, on Android, pins
1011/// the layout linear (vendor tile compression requires `None`).
1012///
1013/// Mapping beyond the declared access is best-effort, never silent: it
1014/// may be refused ([`Error::NotImplemented`]) or take a slow path, and it
1015/// always increments [`unplanned_cpu_access_count`] with a once-per-buffer
1016/// warning.
1017#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
1018pub enum CpuAccess {
1019 /// Hardware-only buffer (the default): no CPU mapping declared.
1020 /// Compression-eligible on platforms with vendor tile layouts.
1021 #[default]
1022 None,
1023 /// CPU reads (verification, CPU consumers) — cached mapping.
1024 Read,
1025 /// CPU writes (decode targets) — write-combined mapping where the
1026 /// platform supports it; reading through a `Write` map is undeclared.
1027 Write,
1028 /// CPU reads and writes — the pre-CpuAccess implicit behavior.
1029 ReadWrite,
1030}
1031
1032impl CpuAccess {
1033 /// Whether this declaration includes CPU reads.
1034 pub fn reads(self) -> bool {
1035 matches!(self, CpuAccess::Read | CpuAccess::ReadWrite)
1036 }
1037
1038 /// Whether this declaration includes CPU writes.
1039 pub fn writes(self) -> bool {
1040 matches!(self, CpuAccess::Write | CpuAccess::ReadWrite)
1041 }
1042
1043 /// Whether this declaration covers `requested` (every direction the
1044 /// request needs is declared).
1045 pub fn covers(self, requested: CpuAccess) -> bool {
1046 (!requested.reads() || self.reads()) && (!requested.writes() || self.writes())
1047 }
1048}
1049
1050/// Requested tile-compression behavior for an image allocation (set via
1051/// [`ImageDesc::with_compression`]).
1052///
1053/// Vendor GPUs store textures in proprietary compressed tile layouts
1054/// (UBWC, AFBC, PVRIC, DCC) that cut memory bandwidth; eligibility
1055/// requires a hardware-only buffer ([`CpuAccess::None`]) because CPU
1056/// mapping pins the layout linear. The request records best knowledge on
1057/// the tensor ([`Tensor::compression`]); it never changes what bytes a
1058/// consumer sees through the GPU/NPU.
1059#[non_exhaustive]
1060#[derive(Debug, Clone, Copy, PartialEq, Eq)]
1061pub enum Compression {
1062 /// Let the platform use its native scheme when the format is
1063 /// eligible; otherwise allocate linear and count the fallback
1064 /// ([`compression_fallback_count`]). The right default for pipelines
1065 /// that want the bandwidth win without portability failures.
1066 Any,
1067 /// Require one specific scheme: allocation fails with
1068 /// [`Error::InvalidArgument`] when the device's scheme differs and
1069 /// [`Error::NotImplemented`] on platforms without vendor tile
1070 /// compression. For consumers whose ABI names a layout (e.g. a
1071 /// QNN context binary declaring UBWC inputs).
1072 Scheme(CompressionScheme),
1073}
1074
1075/// Vendor tile-compression schemes the HAL recognizes and records.
1076///
1077/// Unrecognized platforms record `None` (linear) — there is deliberately
1078/// no `Unknown` variant; a scheme is only recorded when the vendor is
1079/// positively identified.
1080#[non_exhaustive]
1081#[derive(Debug, Clone, Copy, PartialEq, Eq)]
1082pub enum CompressionScheme {
1083 /// Qualcomm Adreno Universal Bandwidth Compression.
1084 Ubwc,
1085 /// Arm Mali/Immortalis Framebuffer Compression.
1086 Afbc,
1087 /// Imagination PowerVR Image Compression (Google Tensor G5+).
1088 Pvric,
1089 /// Samsung Xclipse (AMD RDNA) Delta Color Compression.
1090 Dcc,
1091}
1092
1093/// Count of image allocations that requested [`Compression::Any`] but
1094/// resolved to a linear layout (ineligible format/dtype, unrecognized
1095/// vendor, or a platform without vendor tile compression). Read via
1096/// [`compression_fallback_count`].
1097static COMPRESSION_FALLBACKS: AtomicU64 = AtomicU64::new(0);
1098
1099/// Number of [`Compression::Any`] requests since process start that
1100/// resolved to a linear layout instead of a vendor tile scheme. A
1101/// steady-state pipeline holds this flat after warmup; growth means an
1102/// allocation path keeps requesting compression it never gets.
1103pub fn compression_fallback_count() -> u64 {
1104 COMPRESSION_FALLBACKS.load(Ordering::Relaxed)
1105}
1106
1107/// Record a `Compression::Any` request resolving linear.
1108pub(crate) fn note_compression_fallback(detail: &str) {
1109 COMPRESSION_FALLBACKS.fetch_add(1, Ordering::Relaxed);
1110 log::debug!("compression request fell back to linear: {detail}");
1111}
1112
1113/// Whether this platform can allocate `(format, dtype)` images in a
1114/// vendor tile-compressed layout. `true` requires an Android build, an
1115/// eligible format (RGBA8888 `u8`/`i8` initially), and a positively
1116/// identified GPU vendor; everywhere else the answer is `false` and
1117/// [`Compression::Any`] requests fall back to linear.
1118pub fn compression_support(format: PixelFormat, dtype: DType) -> bool {
1119 #[cfg(target_os = "android")]
1120 {
1121 crate::ahardwarebuffer_layout::compression_eligible(format, dtype)
1122 && crate::ahardwarebuffer::device_compression_scheme().is_some()
1123 }
1124 #[cfg(not(target_os = "android"))]
1125 {
1126 let _ = (format, dtype);
1127 false
1128 }
1129}
1130
1131/// Declarative image-allocation request — the full-featured front door
1132/// for image tensors ([`TensorDyn::image_desc`] and the image crate's
1133/// `ImageProcessor::create_image_desc`).
1134///
1135/// The classic constructors (`image`, `create_image`) cover the common
1136/// cases; the desc carries the optional requests — today the
1137/// [`Compression`] request — without another constructor-parameter
1138/// sweep. Fields are private and the builders consume/return by value,
1139/// so future options are non-breaking.
1140///
1141/// ```
1142/// use edgefirst_tensor::{Compression, CpuAccess, DType, ImageDesc, PixelFormat};
1143/// let desc = ImageDesc::new(640, 640, PixelFormat::Rgba, DType::U8)
1144/// .with_access(CpuAccess::None)
1145/// .with_compression(Compression::Any);
1146/// ```
1147#[derive(Debug, Clone)]
1148pub struct ImageDesc {
1149 width: usize,
1150 height: usize,
1151 format: PixelFormat,
1152 dtype: DType,
1153 memory: Option<TensorMemory>,
1154 access: CpuAccess,
1155 compression: Option<Compression>,
1156}
1157
1158impl ImageDesc {
1159 /// A new image request: auto-selected memory, [`CpuAccess::None`]
1160 /// (hardware-only), no compression request.
1161 pub fn new(width: usize, height: usize, format: PixelFormat, dtype: DType) -> Self {
1162 Self {
1163 width,
1164 height,
1165 format,
1166 dtype,
1167 memory: None,
1168 access: CpuAccess::None,
1169 compression: None,
1170 }
1171 }
1172
1173 /// Request a specific memory backing (`None` = auto-select).
1174 pub fn with_memory(mut self, memory: Option<TensorMemory>) -> Self {
1175 self.memory = memory;
1176 self
1177 }
1178
1179 /// Declare the CPU access (see [`CpuAccess`]). Any declaration other
1180 /// than `None` makes a compression request invalid.
1181 pub fn with_access(mut self, access: CpuAccess) -> Self {
1182 self.access = access;
1183 self
1184 }
1185
1186 /// Request a tile-compressed layout (see [`Compression`]).
1187 pub fn with_compression(mut self, compression: Compression) -> Self {
1188 self.compression = Some(compression);
1189 self
1190 }
1191
1192 /// Requested width in pixels.
1193 pub fn width(&self) -> usize {
1194 self.width
1195 }
1196
1197 /// Requested height in pixels.
1198 pub fn height(&self) -> usize {
1199 self.height
1200 }
1201
1202 /// Requested pixel format.
1203 pub fn format(&self) -> PixelFormat {
1204 self.format
1205 }
1206
1207 /// Requested element type.
1208 pub fn dtype(&self) -> DType {
1209 self.dtype
1210 }
1211
1212 /// Requested memory backing (`None` = auto-select).
1213 pub fn memory(&self) -> Option<TensorMemory> {
1214 self.memory
1215 }
1216
1217 /// Declared CPU access.
1218 pub fn access(&self) -> CpuAccess {
1219 self.access
1220 }
1221
1222 /// The compression request, if any.
1223 pub fn compression(&self) -> Option<Compression> {
1224 self.compression
1225 }
1226}
1227
1228/// Unique identity for a tensor's underlying buffer.
1229///
1230/// Created fresh on every buffer allocation or import. The `id` is a monotonic
1231/// u64 used as a cache key. The `guard` is an `Arc<()>` whose weak references
1232/// allow downstream caches to detect when the buffer has been dropped.
1233#[derive(Debug, Clone)]
1234pub struct BufferIdentity {
1235 id: u64,
1236 guard: Arc<()>,
1237}
1238
1239impl BufferIdentity {
1240 /// Create a new unique buffer identity.
1241 pub fn new() -> Self {
1242 Self {
1243 id: NEXT_BUFFER_ID.fetch_add(1, Ordering::Relaxed),
1244 guard: Arc::new(()),
1245 }
1246 }
1247
1248 /// Unique identifier for this buffer. Changes when the buffer changes.
1249 pub fn id(&self) -> u64 {
1250 self.id
1251 }
1252
1253 /// Returns a weak reference to the buffer guard. Goes dead when the
1254 /// owning Tensor is dropped (and no clones remain).
1255 pub fn weak(&self) -> Weak<()> {
1256 Arc::downgrade(&self.guard)
1257 }
1258
1259 /// Rebuild an identity from interned parts — crate-private: only the
1260 /// Android `AHardwareBuffer_getId` intern table may resurrect an
1261 /// existing identity (arbitrary construction would forge cache hits).
1262 // Only the Android AHardwareBuffer intern path uses these today.
1263 #[cfg_attr(not(target_os = "android"), allow(dead_code))]
1264 pub(crate) fn from_parts(id: u64, guard: Arc<()>) -> Self {
1265 Self { id, guard }
1266 }
1267
1268 /// The strong guard handle (for the intern table's mint path).
1269 #[cfg_attr(not(target_os = "android"), allow(dead_code))]
1270 pub(crate) fn guard_arc(&self) -> Arc<()> {
1271 Arc::clone(&self.guard)
1272 }
1273}
1274
1275impl Default for BufferIdentity {
1276 fn default() -> Self {
1277 Self::new()
1278 }
1279}
1280
1281#[cfg(target_os = "linux")]
1282use nix::sys::stat::{major, minor};
1283
1284/// Filesystem magic of the internal `dma_buf` mount, from
1285/// `include/uapi/linux/magic.h` (`DMA_BUF_MAGIC`, "DMAB").
1286///
1287/// This is stable UAPI and is the only reliable way to recognize a DMA-BUF
1288/// fd — see [`TensorStorage::from_fd`] for why `st_dev` cannot be used.
1289#[cfg(target_os = "linux")]
1290const DMA_BUF_MAGIC: u32 = 0x444d_4142;
1291
1292/// Filesystem magic of tmpfs, from `include/uapi/linux/magic.h`
1293/// (`TMPFS_MAGIC`). Covers both POSIX `shm_open` segments under `/dev/shm`
1294/// and anonymous `memfd_create` files.
1295#[cfg(target_os = "linux")]
1296const TMPFS_MAGIC: u32 = 0x0102_1994;
1297
1298/// Normalize a raw `fstatfs` `f_type` to the 32-bit unsigned magic that
1299/// `include/uapi/linux/magic.h` documents.
1300///
1301/// The width and signedness of `f_type` are target-dependent: `__fsword_t`
1302/// on Linux/gnu (`i64` on 64-bit, `i32` on 32-bit targets such as armv7,
1303/// i686 and aarch64-ilp32), `c_int` on uclibc, `c_ulong` on musl, and
1304/// `c_uint` on s390x. Where it is signed and 32 bits wide, widening it
1305/// sign-extends every magic with bit 31 set — hugetlbfs (`0x958458f6`),
1306/// btrfs (`0x9123683e`), f2fs (`0xf2f52010`) — which would make the value
1307/// reported by [`Error::UnknownBufferType`] impossible to find in the
1308/// header the docs point callers at.
1309///
1310/// Every magic is a 32-bit quantity, so truncating back to `u32` recovers
1311/// the true value from all four representations and loses nothing.
1312#[cfg(target_os = "linux")]
1313const fn fs_magic(raw: i64) -> u32 {
1314 raw as u32
1315}
1316
1317/// The operations every memory backend implements.
1318///
1319/// This is the seam that lets [`Tensor<T>`] hide which backend is in play.
1320/// The implementors are `DmaTensor` (Linux DMA-BUF), `IoSurfaceTensor`
1321/// (macOS/iOS), `AHardwareBufferTensor` (Android), `ShmTensor`, `MemTensor`,
1322/// and [`PboTensor`]; `TensorStorage<T>` dispatches to whichever is active.
1323/// Backend types are not public — allocate a [`Tensor<T>`] or [`TensorDyn`]
1324/// and call these methods through it.
1325///
1326/// Import the trait to get `shape`, `size`, `map`, `clone_fd`,
1327/// `buffer_identity`, and the zero-copy sub-region `view` on a tensor value.
1328/// `Tensor::view` and `Tensor::batch` route through [`TensorTrait::view`], so
1329/// each backend's identity-sharing rule lives in exactly one place.
1330///
1331/// [`TensorDyn`]: crate::TensorDyn
1332pub trait TensorTrait<T>: Send + Sync
1333where
1334 T: Num + Clone + fmt::Debug,
1335{
1336 /// Create a new tensor with the given shape and optional name. If no name
1337 /// is given, a random name will be generated.
1338 fn new(shape: &[usize], name: Option<&str>) -> Result<Self>
1339 where
1340 Self: Sized;
1341
1342 #[cfg(unix)]
1343 /// Import an existing buffer as a tensor, taking ownership of its file
1344 /// descriptor. The buffer is adopted in place — no bytes are copied.
1345 ///
1346 /// The backend is **detected**, not chosen: the fd already belongs to a
1347 /// buffer type, and this call's job is to recognize which. On Linux that
1348 /// is decided by the fd's filesystem magic, which is stable UAPI
1349 /// (`include/uapi/linux/magic.h`):
1350 ///
1351 /// | Filesystem | Magic | Resulting [`TensorMemory`] |
1352 /// |------------|-------|----------------------------|
1353 /// | `dma_buf` | `DMA_BUF_MAGIC` (`0x444d4142`) | [`TensorMemory::Dma`] |
1354 /// | `tmpfs` (`/dev/shm` **and** `memfd`) | `TMPFS_MAGIC` (`0x01021994`) | [`TensorMemory::Shm`] |
1355 /// | anything else | — | *rejected* — see Errors |
1356 ///
1357 /// Both supported types are identified **positively**. An unrecognized
1358 /// filesystem is an error, never a fallback to shared memory: the wrong
1359 /// branch does not fail loudly (a DMA-BUF is mmap-able, so it would
1360 /// import as a perfectly functional tensor that merely isn't DMA, and a
1361 /// pipe would import as a zero-length one), so guessing would trade a
1362 /// clear error here for silent loss of zero-copy far downstream.
1363 ///
1364 /// The device number is deliberately **not** consulted. `dma_buf` files
1365 /// live on an internal kernel mount whose `st_dev` comes from
1366 /// `get_anon_bdev()` — an IDA shared with every other anonymous
1367 /// pseudo-filesystem and allocated in boot order — so the minor a
1368 /// DMA-BUF lands on varies by kernel build and is not part of any ABI.
1369 ///
1370 /// On non-Linux Unix (macOS/iOS/Android) there is no fd-based DMA import;
1371 /// the fd is always adopted as [`TensorMemory::Shm`].
1372 ///
1373 /// # Arguments
1374 ///
1375 /// * `fd` - Owned descriptor for the buffer to import. Ownership
1376 /// transfers to the returned tensor and the fd is closed on drop; pass
1377 /// [`clone_fd`](TensorTrait::clone_fd) output to keep your own handle.
1378 /// * `shape` - Logical shape to interpret the buffer with. Must describe
1379 /// no more elements than the buffer holds.
1380 /// * `name` - Optional name; a random one is generated when `None`.
1381 ///
1382 /// # Returns
1383 ///
1384 /// A tensor sharing the imported buffer's memory, whose
1385 /// [`memory()`](TensorTrait::memory) reports the detected backend.
1386 ///
1387 /// # Errors
1388 ///
1389 /// * [`Error::UnknownBufferType`] - the fd is on a filesystem that is
1390 /// neither `dma_buf` nor `tmpfs`, so its buffer type cannot be
1391 /// determined. Carries the observed `fstatfs` magic as a `u32`,
1392 /// normalized so it can be looked up in `magic.h` directly on both
1393 /// 32- and 64-bit targets. Typical causes: a regular file, a pipe or
1394 /// socket, or a `MFD_HUGETLB` memfd (hugetlbfs, not tmpfs). Linux only.
1395 /// * [`Error::UnknownDeviceType`] - the fd's `st_dev` major is non-zero,
1396 /// i.e. it lives on a real block device rather than an anonymous or
1397 /// in-memory filesystem. Linux only.
1398 /// * [`Error::InvalidSize`] - `shape` is empty or describes zero
1399 /// elements.
1400 /// * [`Error::NixError`] - `fstat`, `fstatfs`, or `mmap` failed on the
1401 /// descriptor.
1402 ///
1403 /// # Examples
1404 ///
1405 /// ```no_run
1406 /// use edgefirst_tensor::{Tensor, TensorMemory, TensorTrait};
1407 ///
1408 /// # fn main() -> edgefirst_tensor::Result<()> {
1409 /// let src = Tensor::<u8>::new(&[480, 640, 3], Some(TensorMemory::Dma), None)?;
1410 ///
1411 /// // Round-tripping a DMA-BUF fd preserves the backend — the import is
1412 /// // still zero-copy, and still eligible for GPU/NPU paths.
1413 /// let imported = Tensor::<u8>::from_fd(src.clone_fd()?, src.shape(), None)?;
1414 /// assert_eq!(imported.memory(), TensorMemory::Dma);
1415 /// # Ok(())
1416 /// # }
1417 /// ```
1418 fn from_fd(fd: std::os::fd::OwnedFd, shape: &[usize], name: Option<&str>) -> Result<Self>
1419 where
1420 Self: Sized;
1421
1422 #[cfg(unix)]
1423 /// Clone the file descriptor associated with this tensor.
1424 fn clone_fd(&self) -> Result<std::os::fd::OwnedFd>;
1425
1426 /// Get the memory type of this tensor.
1427 fn memory(&self) -> TensorMemory;
1428
1429 /// Get the name of this tensor.
1430 fn name(&self) -> String;
1431
1432 /// Get the number of elements in this tensor.
1433 fn len(&self) -> usize {
1434 self.shape().iter().product()
1435 }
1436
1437 /// Check if the tensor is empty.
1438 fn is_empty(&self) -> bool {
1439 self.len() == 0
1440 }
1441
1442 /// Get the size in bytes of this tensor.
1443 fn size(&self) -> usize {
1444 self.len() * std::mem::size_of::<T>()
1445 }
1446
1447 /// Get the shape of this tensor.
1448 fn shape(&self) -> &[usize];
1449
1450 /// Reshape this tensor to the given shape. The total number of elements
1451 /// must remain the same.
1452 fn reshape(&mut self, shape: &[usize]) -> Result<()>;
1453
1454 /// Bytes of the underlying allocation (>= the current logical `size()`).
1455 /// Defaults to the logical size for storages without spare capacity.
1456 fn capacity_bytes(&self) -> usize {
1457 self.size()
1458 }
1459
1460 /// Set the logical shape to any shape whose byte size fits the allocation
1461 /// capacity, without the equal-size constraint of `reshape`.
1462 fn set_logical_shape(&mut self, shape: &[usize]) -> Result<()> {
1463 self.reshape(shape)
1464 }
1465
1466 /// Map the tensor into memory with the given access direction and
1467 /// return a TensorMap for accessing the data.
1468 ///
1469 /// `access` selects the platform mapping mode (read-only IOSurface
1470 /// lock, dma-buf sync direction, AHardwareBuffer lock usage) and the
1471 /// map's mutability: a map obtained with [`CpuAccess::Read`] rejects
1472 /// `as_mut_slice`. [`CpuAccess::None`] is not a mappable direction
1473 /// and returns [`Error::InvalidArgument`].
1474 ///
1475 /// Prefer the typed wrappers [`map_read`](Self::map_read) /
1476 /// [`map_write`](Self::map_write) / [`map_mut`](Self::map_mut).
1477 fn map_with(&self, access: CpuAccess) -> Result<TensorMap<T>>;
1478
1479 /// Map the tensor read-write (equivalent to
1480 /// `map_with(CpuAccess::ReadWrite)` — the historical `map()`
1481 /// behavior).
1482 fn map(&self) -> Result<TensorMap<T>> {
1483 self.map_with(CpuAccess::ReadWrite)
1484 }
1485
1486 /// Map the tensor for CPU reading only. The returned map rejects
1487 /// `as_mut_slice`; on macOS this takes the read-only IOSurface lock
1488 /// (skips the unlock flush), on Linux the dma-buf read-direction
1489 /// sync.
1490 fn map_read(&self) -> Result<TensorMap<T>> {
1491 self.map_with(CpuAccess::Read)
1492 }
1493
1494 /// Map the tensor for CPU writing (fill-only: reading through a
1495 /// write map may see write-combined memory — do not read the slice).
1496 fn map_write(&self) -> Result<TensorMap<T>> {
1497 self.map_with(CpuAccess::Write)
1498 }
1499
1500 /// Map the tensor read-write (alias of [`map`](Self::map) with the
1501 /// intent spelled out).
1502 fn map_mut(&self) -> Result<TensorMap<T>> {
1503 self.map_with(CpuAccess::ReadWrite)
1504 }
1505
1506 /// Get the buffer identity for cache keying and liveness tracking.
1507 fn buffer_identity(&self) -> &BufferIdentity;
1508
1509 /// Create a zero-copy sub-region view of this backing that shares the
1510 /// underlying allocation **and** [`BufferIdentity`].
1511 ///
1512 /// The window is `[offset_bytes, offset_bytes + shape.product() *
1513 /// size_of::<T>())` measured from this tensor's own logical start, so a
1514 /// sub-view of a sub-view composes by adding offsets. Sharing the parent's
1515 /// identity is the contract that lets identity-keyed caches (e.g. the GL
1516 /// EGLImage import cache) treat offset-distinct windows as one buffer rather
1517 /// than unrelated allocations — `view` must never mint a fresh identity.
1518 ///
1519 /// Defaults to [`Error::NotImplemented`]; every backend that supports
1520 /// sub-views overrides it (`Mem`, `Shm`, Linux DMA, macOS IOSurface, `Pbo`).
1521 fn view(&self, offset_bytes: usize, shape: &[usize]) -> Result<Self>
1522 where
1523 Self: Sized,
1524 {
1525 let _ = (offset_bytes, shape);
1526 Err(Error::NotImplemented(
1527 "view (zero-copy sub-region) is not supported for this tensor backend".to_owned(),
1528 ))
1529 }
1530}
1531
1532pub trait TensorMapTrait<T>
1533where
1534 T: Num + Clone + fmt::Debug,
1535{
1536 /// Get the shape of this tensor map.
1537 fn shape(&self) -> &[usize];
1538
1539 /// Unmap the tensor from memory.
1540 fn unmap(&mut self);
1541
1542 /// Get the number of elements in this tensor map.
1543 fn len(&self) -> usize {
1544 self.shape().iter().product()
1545 }
1546
1547 /// Check if the tensor map is empty.
1548 fn is_empty(&self) -> bool {
1549 self.len() == 0
1550 }
1551
1552 /// Get the size in bytes of this tensor map.
1553 fn size(&self) -> usize {
1554 self.len() * std::mem::size_of::<T>()
1555 }
1556
1557 /// Get a slice to the data in this tensor map.
1558 fn as_slice(&self) -> &[T];
1559
1560 /// Get a mutable slice to the data in this tensor map.
1561 fn as_mut_slice(&mut self) -> &mut [T];
1562
1563 #[cfg(feature = "ndarray")]
1564 /// Get an ndarray ArrayView of the tensor data.
1565 fn view(&'_ self) -> Result<ndarray::ArrayView<'_, T, ndarray::Dim<ndarray::IxDynImpl>>> {
1566 Ok(ndarray::ArrayView::from_shape(
1567 self.shape(),
1568 self.as_slice(),
1569 )?)
1570 }
1571
1572 #[cfg(feature = "ndarray")]
1573 /// Get an ndarray ArrayViewMut of the tensor data.
1574 fn view_mut(
1575 &'_ mut self,
1576 ) -> Result<ndarray::ArrayViewMut<'_, T, ndarray::Dim<ndarray::IxDynImpl>>> {
1577 let shape = self.shape().to_vec();
1578 Ok(ndarray::ArrayViewMut::from_shape(
1579 shape,
1580 self.as_mut_slice(),
1581 )?)
1582 }
1583}
1584
1585/// Which memory backend a tensor is (or should be) allocated from.
1586///
1587/// Pass `Some(..)` to a constructor to pin the backend, or `None` to
1588/// auto-select. Auto-selection differs by constructor: [`Tensor::new`] tries
1589/// `Dma` → `Shm` → `Mem`, while the image constructors try `Dma` → `Mem` and
1590/// skip `Shm`. `EDGEFIRST_TENSOR_FORCE_MEM=1` short-circuits either chain to
1591/// `Mem`, which is how tests run on hosts without DMA-heap permissions.
1592///
1593/// A pinned request has no fallback: if the backend cannot serve it, the
1594/// constructor fails rather than quietly giving you something slower.
1595/// [`TensorTrait::memory`] reports what a tensor actually ended up with.
1596#[derive(Debug, Clone, Copy, PartialEq, Eq)]
1597pub enum TensorMemory {
1598 /// Platform-native zero-copy GPU buffer.
1599 ///
1600 /// On Linux this is a DMA-BUF (`DmaTensor` in `crates/tensor/src/dma.rs`)
1601 /// allocated via the DRM/dma-heap subsystem. On macOS this is an
1602 /// IOSurface (`IoSurfaceTensor` in `crates/tensor/src/iosurface.rs`).
1603 /// Both fit into the same `TensorStorage::Dma` slot at the trait
1604 /// level — the public C API discriminant (`HalTensorMemory::Dma=1`)
1605 /// works on both platforms with no ABI break.
1606 ///
1607 /// Allows hardware-accelerated paths (OpenGL backend on Linux via
1608 /// `EGL_EXT_image_dma_buf_import`; macOS via
1609 /// `EGL_ANGLE_iosurface_client_buffer`). CPU access via `map()`
1610 /// incurs cache-coherency overhead on Linux DMA-BUF and is similar
1611 /// in cost on IOSurface; SHM/Mem are cheaper for CPU-only workloads.
1612 Dma,
1613 #[cfg(unix)]
1614 /// POSIX Shared Memory allocation. Suitable for inter-process
1615 /// communication, but not suitable for hardware acceleration.
1616 Shm,
1617
1618 /// Regular system memory allocation
1619 Mem,
1620
1621 /// OpenGL Pixel Buffer Object memory. Created by ImageProcessor
1622 /// when DMA-buf is unavailable but OpenGL is present.
1623 Pbo,
1624}
1625
1626impl From<TensorMemory> for String {
1627 fn from(memory: TensorMemory) -> Self {
1628 match memory {
1629 TensorMemory::Dma => "dma".to_owned(),
1630 #[cfg(unix)]
1631 TensorMemory::Shm => "shm".to_owned(),
1632 TensorMemory::Mem => "mem".to_owned(),
1633 TensorMemory::Pbo => "pbo".to_owned(),
1634 }
1635 }
1636}
1637
1638impl TryFrom<&str> for TensorMemory {
1639 type Error = Error;
1640
1641 fn try_from(s: &str) -> Result<Self> {
1642 match s {
1643 "dma" => Ok(TensorMemory::Dma),
1644 #[cfg(unix)]
1645 "shm" => Ok(TensorMemory::Shm),
1646 "mem" => Ok(TensorMemory::Mem),
1647 "pbo" => Ok(TensorMemory::Pbo),
1648 _ => Err(Error::InvalidMemoryType(s.to_owned())),
1649 }
1650 }
1651}
1652
1653#[derive(Debug)]
1654#[allow(dead_code)] // Variants are constructed by downstream crates via pub(crate) helpers
1655pub(crate) enum TensorStorage<T>
1656where
1657 T: Num + Clone + fmt::Debug + Send + Sync,
1658{
1659 /// Platform-native zero-copy GPU buffer. Inner type differs per
1660 /// target: `DmaTensor` on Linux (DMA-BUF fd), `IoSurfaceTensor` on
1661 /// macOS (CFRetained IOSurface). The shared variant name keeps the
1662 /// public `TensorMemory::Dma` discriminant stable across platforms.
1663 #[cfg(target_os = "linux")]
1664 Dma(DmaTensor<T>),
1665 #[cfg(any(target_os = "macos", target_os = "ios"))]
1666 Dma(IoSurfaceTensor<T>),
1667 #[cfg(target_os = "android")]
1668 Dma(AHardwareBufferTensor<T>),
1669 #[cfg(unix)]
1670 Shm(ShmTensor<T>),
1671 Mem(MemTensor<T>),
1672 Pbo(PboTensor<T>),
1673}
1674
1675impl<T> TensorStorage<T>
1676where
1677 T: Num + Clone + fmt::Debug + Send + Sync,
1678{
1679 /// The backing allocation's intrinsic physical row pitch in bytes, if it
1680 /// has one that is fixed independent of the logical shape. macOS IOSurface
1681 /// reports its 64-aligned `bytesPerRow`; other backings (Linux DMA, SHM,
1682 /// Mem, PBO) have no fixed pitch beyond the logical shape and return `None`.
1683 ///
1684 /// Used by `configure_image` to preserve the physical pitch when a reused
1685 /// pool tensor is reconfigured to a smaller logical image — so the decode
1686 /// writes rows at the surface's real stride and the GPU samples them with
1687 /// the same stride (the physical-grid / logical-ROI decoupling).
1688 pub(crate) fn backing_row_stride(&self) -> Option<usize> {
1689 match self {
1690 // Only genuine image-formatted IOSurfaces (height > 1) carry a real
1691 // per-row pitch; a generic byte-bag (height == 1) returns `None` so
1692 // `configure_image` does not adopt its whole-buffer "row" as a stride.
1693 #[cfg(any(target_os = "macos", target_os = "ios"))]
1694 TensorStorage::Dma(t) => t.image_backing_row_stride(),
1695 // Android AHardwareBuffer: same rule — only genuine 2D
1696 // image-formatted buffers (height > 1) carry a real row pitch.
1697 #[cfg(target_os = "android")]
1698 TensorStorage::Dma(t) => t.image_backing_row_stride(),
1699 _ => None,
1700 }
1701 }
1702
1703 /// Create a new tensor storage with the given shape, memory type, and
1704 /// optional name. If no name is given, a random name will be generated.
1705 /// If no memory type is given, the best available memory type will be
1706 /// chosen based on the platform and environment variables.
1707 fn new(shape: &[usize], memory: Option<TensorMemory>, name: Option<&str>) -> Result<Self> {
1708 match memory {
1709 #[cfg(target_os = "linux")]
1710 Some(TensorMemory::Dma) => {
1711 DmaTensor::<T>::new(shape, name).map(TensorStorage::Dma)
1712 }
1713 #[cfg(any(target_os = "macos", target_os = "ios"))]
1714 Some(TensorMemory::Dma) => {
1715 IoSurfaceTensor::<T>::new(shape, name).map(TensorStorage::Dma)
1716 }
1717 #[cfg(target_os = "android")]
1718 Some(TensorMemory::Dma) => {
1719 AHardwareBufferTensor::<T>::new(shape, name).map(TensorStorage::Dma)
1720 }
1721 #[cfg(not(any(
1722 target_os = "linux",
1723 target_os = "macos",
1724 target_os = "ios",
1725 target_os = "android"
1726 )))]
1727 Some(TensorMemory::Dma) => Err(crate::error::Error::NotImplemented(
1728 "TensorMemory::Dma is only available on Linux (DMA-BUF), macOS/iOS (IOSurface), \
1729 and Android (AHardwareBuffer)"
1730 .to_owned(),
1731 )),
1732 #[cfg(unix)]
1733 Some(TensorMemory::Shm) => {
1734 ShmTensor::<T>::new(shape, name).map(TensorStorage::Shm)
1735 }
1736 Some(TensorMemory::Mem) => {
1737 MemTensor::<T>::new(shape, name).map(TensorStorage::Mem)
1738 }
1739 Some(TensorMemory::Pbo) => Err(crate::error::Error::NotImplemented(
1740 "PboTensor cannot be created via Tensor::new() — use ImageProcessor::create_image()".to_owned(),
1741 )),
1742 None => {
1743 if std::env::var("EDGEFIRST_TENSOR_FORCE_MEM")
1744 .is_ok_and(|x| x != "0" && x.to_lowercase() != "false")
1745 {
1746 MemTensor::<T>::new(shape, name).map(TensorStorage::Mem)
1747 } else {
1748 // Auto-select priority: Dma > Mem. Shm is intentionally NOT
1749 // auto-selected — it offers no advantage over Mem for an
1750 // in-process tensor and Mem always succeeds, so Shm below Mem
1751 // is effectively never reached. Request Shm explicitly via
1752 // `TensorMemory::Shm` when cross-process sharing is needed.
1753 // (PBO sits between Dma and Mem but is GL-backed and created
1754 // only via `ImageProcessor::create_image`.)
1755 #[cfg(target_os = "linux")]
1756 {
1757 // Linux: Try DMA -> Mem
1758 match DmaTensor::<T>::new(shape, name) {
1759 Ok(tensor) => Ok(TensorStorage::Dma(tensor)),
1760 Err(_) => MemTensor::<T>::new(shape, name).map(TensorStorage::Mem),
1761 }
1762 }
1763 #[cfg(any(target_os = "macos", target_os = "ios"))]
1764 {
1765 // macOS/iOS: Try IOSurface -> Mem. IOSurface is the
1766 // GPU-shareable backend (zero-copy via ANGLE), filling the
1767 // same role as DMA-BUF on Linux.
1768 match IoSurfaceTensor::<T>::new(shape, name) {
1769 Ok(tensor) => Ok(TensorStorage::Dma(tensor)),
1770 Err(_) => MemTensor::<T>::new(shape, name).map(TensorStorage::Mem),
1771 }
1772 }
1773 #[cfg(target_os = "android")]
1774 {
1775 // Android: Mem. Unlike macOS (where a byte-bag
1776 // IOSurface is still GL-importable as R8) a generic
1777 // BLOB AHardwareBuffer cannot back a GPU texture, so
1778 // auto-selecting it buys nothing and costs plenty: a
1779 // gralloc allocator-HAL ioctl per allocation (orders
1780 // slower than malloc, ≥1 page + a dmabuf fd even for
1781 // tiny tensors) and a lock/unlock cache-maintenance
1782 // round trip on every map(). Zero-copy image tensors
1783 // come from `Tensor::image(..)`; callers that want a
1784 // BLOB (NNAPI handoff) request `TensorMemory::Dma`
1785 // explicitly.
1786 MemTensor::<T>::new(shape, name).map(TensorStorage::Mem)
1787 }
1788 #[cfg(all(
1789 unix,
1790 not(any(
1791 target_os = "linux",
1792 target_os = "macos",
1793 target_os = "ios",
1794 target_os = "android"
1795 ))
1796 ))]
1797 {
1798 // Other Unix (BSD): Mem only (no DMA; Shm is explicit-only)
1799 MemTensor::<T>::new(shape, name).map(TensorStorage::Mem)
1800 }
1801 #[cfg(not(unix))]
1802 {
1803 // Windows/other: Mem only
1804 MemTensor::<T>::new(shape, name).map(TensorStorage::Mem)
1805 }
1806 }
1807 }
1808 }
1809 }
1810
1811 /// Create a DMA-backed tensor storage with an explicit byte size that
1812 /// may exceed `shape.product() * sizeof(T)`. Used for image tensors
1813 /// with row-padded layouts (see `DmaTensor::new_with_byte_size`).
1814 ///
1815 /// This is intentionally DMA-only: padding is only meaningful for
1816 /// buffers that will be imported as GPU textures via EGLImage. PBO,
1817 /// Shm, and Mem storage doesn't benefit from pitch alignment and
1818 /// shouldn't pay the memory cost.
1819 #[cfg(target_os = "linux")]
1820 pub(crate) fn new_dma_with_byte_size(
1821 shape: &[usize],
1822 byte_size: usize,
1823 name: Option<&str>,
1824 ) -> Result<Self> {
1825 DmaTensor::<T>::new_with_byte_size(shape, byte_size, name).map(TensorStorage::Dma)
1826 }
1827
1828 // No non-Linux stub: the only caller (`Tensor::image_with_stride`)
1829 // returns `NotImplemented` directly on non-Linux without ever
1830 // reaching the storage layer, so defining a stub here would be
1831 // dead code and fail the `-D warnings` clippy gate on macOS CI.
1832
1833 /// Create a Mem-backed tensor storage with an explicit byte size that may
1834 /// exceed `shape.product() * sizeof(T)`. Used for image tensors with
1835 /// 64-byte-aligned row strides (see `MemTensor::with_capacity_bytes`).
1836 pub(crate) fn new_mem_with_byte_size(
1837 shape: &[usize],
1838 byte_size: usize,
1839 name: Option<&str>,
1840 ) -> Result<Self>
1841 where
1842 T: 'static,
1843 {
1844 MemTensor::<T>::with_capacity_bytes(shape, byte_size, name).map(TensorStorage::Mem)
1845 }
1846
1847 /// Create a Shm-backed tensor storage with an explicit byte size that may
1848 /// exceed `shape.product() * sizeof(T)`. Used for image tensors with
1849 /// 64-byte-aligned row strides (see `ShmTensor::new_with_byte_size`).
1850 #[cfg(unix)]
1851 pub(crate) fn new_shm_with_byte_size(
1852 shape: &[usize],
1853 byte_size: usize,
1854 name: Option<&str>,
1855 ) -> Result<Self> {
1856 ShmTensor::<T>::new_with_byte_size(shape, byte_size, name).map(TensorStorage::Shm)
1857 }
1858
1859 /// Allocate an image-formatted IOSurface-backed storage (macOS).
1860 ///
1861 /// Used by `Tensor::image()` when the caller requests
1862 /// `TensorMemory::Dma` and the format has an IOSurface FourCC
1863 /// mapping (YUYV, RGBA, BGRA today). Falls back to `new_with_byte_size`
1864 /// otherwise.
1865 #[cfg(any(target_os = "macos", target_os = "ios"))]
1866 pub(crate) fn new_image_iosurface(
1867 width: usize,
1868 height: usize,
1869 format: PixelFormat,
1870 dtype: DType,
1871 shape: &[usize],
1872 name: Option<&str>,
1873 ) -> Result<Self> {
1874 IoSurfaceTensor::<T>::new_image(width, height, format, dtype, shape, name)
1875 .map(TensorStorage::Dma)
1876 }
1877
1878 /// Allocate an image-formatted AHardwareBuffer-backed storage (Android).
1879 ///
1880 /// Used by `Tensor::image()` when the caller requests
1881 /// `TensorMemory::Dma` and the format has an AHardwareBuffer format
1882 /// mapping (RGBA8 and the RGBA16F float paths today). Falls back to
1883 /// `new_with_byte_size` otherwise.
1884 #[cfg(target_os = "android")]
1885 pub(crate) fn new_image_ahardwarebuffer(
1886 width: usize,
1887 height: usize,
1888 format: PixelFormat,
1889 dtype: DType,
1890 shape: &[usize],
1891 name: Option<&str>,
1892 access: CpuAccess,
1893 ) -> Result<Self> {
1894 AHardwareBufferTensor::<T>::new_image(width, height, format, dtype, shape, name, access)
1895 .map(TensorStorage::Dma)
1896 }
1897
1898 /// Create a new tensor storage using the given file descriptor, shape,
1899 /// and optional name.
1900 #[cfg(unix)]
1901 fn from_fd(fd: OwnedFd, shape: &[usize], name: Option<&str>) -> Result<Self> {
1902 #[cfg(target_os = "linux")]
1903 {
1904 use nix::sys::stat::fstat;
1905 use nix::sys::statfs::fstatfs;
1906
1907 let stat = fstat(&fd)?;
1908 let major = major(stat.st_dev);
1909 let minor = minor(stat.st_dev);
1910
1911 if major != 0 {
1912 // Dma and Shm tensors are expected to have major number 0
1913 return Err(Error::UnknownDeviceType(major, minor));
1914 }
1915
1916 // Classify by filesystem magic, never by the st_dev minor.
1917 //
1918 // dma_buf files live on an internal kernel mount (`dma_buf_mnt`,
1919 // created with `kern_mount`) whose superblock draws its device
1920 // number from `get_anon_bdev()` — an IDA allocated first-come,
1921 // first-served during boot, shared with every other anonymous
1922 // pseudo-filesystem (pipefs, sockfs, anon_inodefs, nsfs, bdev,
1923 // tracefs, ...). The minor a DMA-BUF lands on is therefore a
1924 // function of which pseudo-filesystems registered first on that
1925 // kernel build: it moves with kernel config, driver load order,
1926 // initramfs use, and kernel version. It is not part of any ABI.
1927 //
1928 // Observed values for a genuine DMA-BUF: 12 on x86 desktop, 8 on
1929 // the ADIS Verdin. An earlier revision hardcoded `9 | 10` here,
1930 // which silently imported real DMA-BUFs as shared memory
1931 // everywhere else.
1932 //
1933 // The filesystem magic, by contrast, is stable UAPI
1934 // (include/uapi/linux/magic.h). Both branches are identified
1935 // positively; anything else is genuinely unknown and is rejected
1936 // rather than guessed at. Falling back to SHM would "work" — a
1937 // DMA-BUF is mmap-able, and even a pipe imports as a zero-length
1938 // tensor — which is exactly why it must not be the default.
1939 let magic = fs_magic(fstatfs(&fd)?.filesystem_type().0 as i64);
1940
1941 log::debug!(
1942 "Creating tensor from fd: major={major}, minor={minor}, magic={magic:#010x}"
1943 );
1944
1945 match magic {
1946 DMA_BUF_MAGIC => DmaTensor::<T>::from_fd(fd, shape, name).map(TensorStorage::Dma),
1947 TMPFS_MAGIC => ShmTensor::<T>::from_fd(fd, shape, name).map(TensorStorage::Shm),
1948 other => Err(Error::UnknownBufferType(other)),
1949 }
1950 }
1951 #[cfg(all(unix, not(target_os = "linux")))]
1952 {
1953 // On macOS/iOS/BSD, always use SHM (no DMA-BUF fd import)
1954 ShmTensor::<T>::from_fd(fd, shape, name).map(TensorStorage::Shm)
1955 }
1956 }
1957}
1958
1959impl<T> TensorTrait<T> for TensorStorage<T>
1960where
1961 T: Num + Clone + fmt::Debug + Send + Sync,
1962{
1963 fn new(shape: &[usize], name: Option<&str>) -> Result<Self> {
1964 Self::new(shape, None, name)
1965 }
1966
1967 #[cfg(unix)]
1968 fn from_fd(fd: OwnedFd, shape: &[usize], name: Option<&str>) -> Result<Self> {
1969 Self::from_fd(fd, shape, name)
1970 }
1971
1972 #[cfg(unix)]
1973 fn clone_fd(&self) -> Result<OwnedFd> {
1974 match self {
1975 #[cfg(any(
1976 target_os = "linux",
1977 target_os = "macos",
1978 target_os = "ios",
1979 target_os = "android"
1980 ))]
1981 TensorStorage::Dma(t) => t.clone_fd(),
1982 TensorStorage::Shm(t) => t.clone_fd(),
1983 TensorStorage::Mem(t) => t.clone_fd(),
1984 TensorStorage::Pbo(t) => t.clone_fd(),
1985 }
1986 }
1987
1988 fn memory(&self) -> TensorMemory {
1989 match self {
1990 #[cfg(any(
1991 target_os = "linux",
1992 target_os = "macos",
1993 target_os = "ios",
1994 target_os = "android"
1995 ))]
1996 TensorStorage::Dma(_) => TensorMemory::Dma,
1997 #[cfg(unix)]
1998 TensorStorage::Shm(_) => TensorMemory::Shm,
1999 TensorStorage::Mem(_) => TensorMemory::Mem,
2000 TensorStorage::Pbo(_) => TensorMemory::Pbo,
2001 }
2002 }
2003
2004 fn name(&self) -> String {
2005 match self {
2006 #[cfg(any(
2007 target_os = "linux",
2008 target_os = "macos",
2009 target_os = "ios",
2010 target_os = "android"
2011 ))]
2012 TensorStorage::Dma(t) => t.name(),
2013 #[cfg(unix)]
2014 TensorStorage::Shm(t) => t.name(),
2015 TensorStorage::Mem(t) => t.name(),
2016 TensorStorage::Pbo(t) => t.name(),
2017 }
2018 }
2019
2020 fn shape(&self) -> &[usize] {
2021 match self {
2022 #[cfg(any(
2023 target_os = "linux",
2024 target_os = "macos",
2025 target_os = "ios",
2026 target_os = "android"
2027 ))]
2028 TensorStorage::Dma(t) => t.shape(),
2029 #[cfg(unix)]
2030 TensorStorage::Shm(t) => t.shape(),
2031 TensorStorage::Mem(t) => t.shape(),
2032 TensorStorage::Pbo(t) => t.shape(),
2033 }
2034 }
2035
2036 fn reshape(&mut self, shape: &[usize]) -> Result<()> {
2037 match self {
2038 #[cfg(any(
2039 target_os = "linux",
2040 target_os = "macos",
2041 target_os = "ios",
2042 target_os = "android"
2043 ))]
2044 TensorStorage::Dma(t) => t.reshape(shape),
2045 #[cfg(unix)]
2046 TensorStorage::Shm(t) => t.reshape(shape),
2047 TensorStorage::Mem(t) => t.reshape(shape),
2048 TensorStorage::Pbo(t) => t.reshape(shape),
2049 }
2050 }
2051
2052 fn capacity_bytes(&self) -> usize {
2053 match self {
2054 #[cfg(any(
2055 target_os = "linux",
2056 target_os = "macos",
2057 target_os = "ios",
2058 target_os = "android"
2059 ))]
2060 TensorStorage::Dma(t) => t.capacity_bytes(),
2061 #[cfg(unix)]
2062 TensorStorage::Shm(t) => t.capacity_bytes(),
2063 TensorStorage::Mem(t) => t.capacity_bytes(),
2064 TensorStorage::Pbo(t) => t.capacity_bytes(),
2065 }
2066 }
2067
2068 fn set_logical_shape(&mut self, shape: &[usize]) -> Result<()> {
2069 match self {
2070 #[cfg(any(
2071 target_os = "linux",
2072 target_os = "macos",
2073 target_os = "ios",
2074 target_os = "android"
2075 ))]
2076 TensorStorage::Dma(t) => t.set_logical_shape(shape),
2077 #[cfg(unix)]
2078 TensorStorage::Shm(t) => t.set_logical_shape(shape),
2079 TensorStorage::Mem(t) => t.set_logical_shape(shape),
2080 TensorStorage::Pbo(t) => t.set_logical_shape(shape),
2081 }
2082 }
2083
2084 fn map_with(&self, access: CpuAccess) -> Result<TensorMap<T>> {
2085 match self {
2086 #[cfg(any(
2087 target_os = "linux",
2088 target_os = "macos",
2089 target_os = "ios",
2090 target_os = "android"
2091 ))]
2092 TensorStorage::Dma(t) => t.map_with(access),
2093 #[cfg(unix)]
2094 TensorStorage::Shm(t) => t.map_with(access),
2095 TensorStorage::Mem(t) => t.map_with(access),
2096 TensorStorage::Pbo(t) => t.map_with(access),
2097 }
2098 }
2099
2100 fn buffer_identity(&self) -> &BufferIdentity {
2101 match self {
2102 #[cfg(any(
2103 target_os = "linux",
2104 target_os = "macos",
2105 target_os = "ios",
2106 target_os = "android"
2107 ))]
2108 TensorStorage::Dma(t) => t.buffer_identity(),
2109 #[cfg(unix)]
2110 TensorStorage::Shm(t) => t.buffer_identity(),
2111 TensorStorage::Mem(t) => t.buffer_identity(),
2112 TensorStorage::Pbo(t) => t.buffer_identity(),
2113 }
2114 }
2115
2116 /// Forward a sub-region view to the active backend, re-wrapping the
2117 /// backend's view (which shares the parent's allocation and
2118 /// [`BufferIdentity`]) back into the matching `TensorStorage` variant. Each
2119 /// backend's `view` is its own `TensorTrait::view` override; this single
2120 /// match is the only per-variant dispatch (`Tensor::subview` calls through
2121 /// here rather than matching the storage itself).
2122 fn view(&self, offset_bytes: usize, shape: &[usize]) -> Result<Self> {
2123 match self {
2124 #[cfg(any(
2125 target_os = "linux",
2126 target_os = "macos",
2127 target_os = "ios",
2128 target_os = "android"
2129 ))]
2130 TensorStorage::Dma(t) => t.view(offset_bytes, shape).map(TensorStorage::Dma),
2131 #[cfg(unix)]
2132 TensorStorage::Shm(t) => t.view(offset_bytes, shape).map(TensorStorage::Shm),
2133 TensorStorage::Mem(t) => t.view(offset_bytes, shape).map(TensorStorage::Mem),
2134 TensorStorage::Pbo(t) => t.view(offset_bytes, shape).map(TensorStorage::Pbo),
2135 }
2136 }
2137}
2138
2139/// Multi-backend tensor with optional image format metadata.
2140///
2141/// When `format` is `Some`, this tensor represents an image. Width, height,
2142/// and channels are derived from `shape` + `format`. When `format` is `None`,
2143/// this is a raw tensor (identical to the pre-refactoring behavior).
2144#[derive(Debug)]
2145pub struct Tensor<T>
2146where
2147 T: Num + Clone + fmt::Debug + Send + Sync,
2148{
2149 /// CUDA registration for this tensor, if any. Set after creation by
2150 /// the image crate once a PBO is registered with CUDA interop.
2151 ///
2152 /// MUST be declared before `storage`: CUDA must unregister the GL buffer
2153 /// before storage's Drop deletes it (cudaGraphicsUnregisterResource before
2154 /// glDeleteBuffers). Rust drops fields in declaration order.
2155 cuda: Option<crate::cuda::CudaHandle>,
2156 pub(crate) storage: TensorStorage<T>,
2157 format: Option<PixelFormat>,
2158 chroma: Option<Box<Tensor<T>>>,
2159 /// Row stride in bytes for externally allocated buffers with row padding.
2160 /// `None` means tightly packed (stride == width * bytes_per_pixel).
2161 row_stride: Option<usize>,
2162 /// Byte offset within the DMA-BUF where image data starts.
2163 /// `None` means offset 0 (data starts at the beginning of the buffer).
2164 plane_offset: Option<usize>,
2165 /// Quantization metadata for integer-typed tensors. Public access is
2166 /// gated by the `IntegerType` trait — `Tensor<f32>` etc. carry the
2167 /// field for layout uniformity but have no way to read or write it.
2168 pub(crate) quantization: Option<Quantization>,
2169 /// Optional colorimetry metadata. `None` = undefined; never auto-filled.
2170 colorimetry: Option<crate::Colorimetry>,
2171 /// Declared CPU access (see [`CpuAccess`]). Image constructors set it
2172 /// from their `access` parameter; non-image tensors (`new`, imports,
2173 /// numpy) default to `ReadWrite` — they are CPU-centric by nature.
2174 /// `map_with` counts requests beyond this declaration as unplanned.
2175 cpu_access: CpuAccess,
2176 /// Recorded vendor tile-compression scheme — best knowledge from
2177 /// allocation time (see [`Compression`]). `Some` only for Android
2178 /// hardware-only AHardwareBuffers whose allocation requested
2179 /// compression on an eligible format with a recognized vendor. When
2180 /// set, the row-stride accessors describe no meaningful linear
2181 /// layout. `configure_image` preserves it (physical layout, unlike
2182 /// colorimetry); views inherit it.
2183 compression: Option<CompressionScheme>,
2184 /// Parent-image snapshot when this tensor is a [`view`](Self::view)/
2185 /// [`batch`](Self::batch) sub-region; `None` for a whole tensor. Lets the GL
2186 /// backend key its import on the parent and render the view as a
2187 /// `glViewport`/`glScissor` ROI. See [`ViewOrigin`].
2188 view_origin: Option<ViewOrigin>,
2189}
2190
2191impl<T> Tensor<T>
2192where
2193 T: Num + Clone + fmt::Debug + Send + Sync,
2194{
2195 /// Wrap a TensorStorage in a Tensor with no image metadata.
2196 pub(crate) fn wrap(storage: TensorStorage<T>) -> Self {
2197 Self {
2198 storage,
2199 format: None,
2200 chroma: None,
2201 row_stride: None,
2202 plane_offset: None,
2203 quantization: None,
2204 cuda: None,
2205 colorimetry: None,
2206 cpu_access: CpuAccess::ReadWrite,
2207 compression: None,
2208 view_origin: None,
2209 }
2210 }
2211
2212 /// Construct a tensor from a row-major element slice + shape. Allocates a
2213 /// new buffer (`TensorMemory::Mem`) and memcpys the contents; caller
2214 /// retains ownership of the input slice.
2215 ///
2216 /// # Errors
2217 ///
2218 /// - [`Error::InvalidShape`] if `values.len() != shape.iter().product()`.
2219 /// - Propagates any allocation error from [`Self::new`].
2220 pub fn from_slice(values: &[T], shape: &[usize]) -> Result<Self>
2221 where
2222 T: Copy,
2223 {
2224 let expected: usize = shape.iter().product();
2225 if values.len() != expected {
2226 return Err(Error::InvalidShape(format!(
2227 "from_slice: values.len()={} but shape product={expected} (shape={shape:?})",
2228 values.len()
2229 )));
2230 }
2231 let t = Self::new(shape, Some(TensorMemory::Mem), None)?;
2232 {
2233 let mut m = t.map()?;
2234 m.as_mut_slice().copy_from_slice(values);
2235 }
2236 Ok(t)
2237 }
2238
2239 /// Wrap externally-owned memory as a tensor without copying. The tensor
2240 /// borrows `[ptr, ptr + shape.product() * size_of::<T>())` as
2241 /// [`TensorMemory::Mem`]; `owner`, when `Some`, co-owns the source so it
2242 /// outlives the tensor (and all derived views/maps). See [`ForeignOwner`].
2243 ///
2244 /// The canonical use is CUDA zero-copy: allocate host-coherent memory
2245 /// (`cudaHostAlloc`), wrap the host pointer here, and bind the matching
2246 /// device pointer to the inference engine — reads and writes hit the same
2247 /// physical buffer with no host copy. The identical primitive backs the
2248 /// Python `Tensor.from_numpy` zero-copy borrow (owner = the NumPy object).
2249 ///
2250 /// # Safety
2251 ///
2252 /// `ptr` must be non-null, aligned to `align_of::<T>()`, and valid for
2253 /// `shape.product()` elements of `T` for as long as the returned tensor —
2254 /// and every view/map sharing its backing — is alive. Pass an `owner` that
2255 /// co-owns the source to uphold that contract.
2256 ///
2257 /// # Errors
2258 ///
2259 /// [`Error::InvalidSize`] if `shape` is empty.
2260 pub unsafe fn from_foreign(
2261 ptr: *mut T,
2262 shape: &[usize],
2263 owner: Option<crate::ForeignOwner>,
2264 name: Option<&str>,
2265 ) -> Result<Self> {
2266 if shape.is_empty() {
2267 return Err(Error::InvalidSize(0));
2268 }
2269 if ptr.is_null() {
2270 return Err(Error::InvalidArgument(
2271 "from_foreign: ptr must be non-null".to_owned(),
2272 ));
2273 }
2274 shape
2275 .iter()
2276 .copied()
2277 .try_fold(1usize, |acc, dim| acc.checked_mul(dim))
2278 .ok_or_else(|| {
2279 Error::InvalidArgument(format!(
2280 "from_foreign: shape.product() overflows usize (shape={shape:?})"
2281 ))
2282 })?;
2283 let mem = MemTensor::<T>::from_foreign(ptr, shape, owner, name);
2284 Ok(Self::wrap(TensorStorage::Mem(mem)))
2285 }
2286
2287 /// Construct a tensor from a 3-D ndarray view. Respects strides — one
2288 /// copy in all cases; contiguous views take a memcpy fast path.
2289 ///
2290 /// Only available when the `ndarray` feature is enabled.
2291 #[cfg(feature = "ndarray")]
2292 pub fn from_arrayview3(view: ndarray::ArrayView3<'_, T>) -> Result<Self>
2293 where
2294 T: Copy,
2295 {
2296 let (h, w, c) = view.dim();
2297 let t = Self::new(&[h, w, c], Some(TensorMemory::Mem), None)?;
2298 {
2299 let mut m = t.map()?;
2300 let dst = m.as_mut_slice();
2301 if let Some(src) = view.as_slice() {
2302 dst.copy_from_slice(src);
2303 } else {
2304 for (d, &s) in dst.iter_mut().zip(view.iter()) {
2305 *d = s;
2306 }
2307 }
2308 }
2309 Ok(t)
2310 }
2311
2312 /// Create a new tensor with the given shape, memory type, and optional
2313 /// name. If no name is given, a random name will be generated. If no
2314 /// memory type is given, the best available memory type will be chosen
2315 /// based on the platform and environment variables.
2316 ///
2317 /// On Linux platforms, the order of preference is: Dma -> Shm -> Mem.
2318 /// On other Unix platforms (macOS), the order is: Shm -> Mem.
2319 /// On non-Unix platforms, only Mem is available.
2320 ///
2321 /// # Environment Variables
2322 /// - `EDGEFIRST_TENSOR_FORCE_MEM`: If set to a non-zero and non-false
2323 /// value, forces the use of regular system memory allocation
2324 /// (`TensorMemory::Mem`) regardless of platform capabilities.
2325 ///
2326 /// # Example
2327 /// ```rust
2328 /// use edgefirst_tensor::{Error, Tensor, TensorMemory, TensorTrait};
2329 /// # fn main() -> Result<(), Error> {
2330 /// let tensor = Tensor::<f32>::new(&[2, 3, 4], Some(TensorMemory::Mem), Some("test_tensor"))?;
2331 /// assert_eq!(tensor.memory(), TensorMemory::Mem);
2332 /// assert_eq!(tensor.name(), "test_tensor");
2333 /// # Ok(())
2334 /// # }
2335 /// ```
2336 pub fn new(shape: &[usize], memory: Option<TensorMemory>, name: Option<&str>) -> Result<Self> {
2337 let _span = tracing::trace_span!(
2338 "tensor.alloc",
2339 ?shape,
2340 memory = ?memory,
2341 dtype = std::any::type_name::<T>(),
2342 )
2343 .entered();
2344 #[cfg_attr(not(target_os = "linux"), allow(unused_mut))]
2345 let mut t = TensorStorage::new(shape, memory, name).map(Self::wrap)?;
2346 // Best-effort: attach a CUDA ExternalMemory handle for DMA tensors on
2347 // CUDA-capable hosts. Never blocks tensor creation on failure.
2348 // RUNTIME-UNVALIDATED: no CUDA+dma_heap test platform available; ABI
2349 // layout-asserted vs. CUDA 12.6 driver_types.h; mechanism proven by
2350 // gpu-probe O5 on Orin.
2351 #[cfg(target_os = "linux")]
2352 t.try_init_dma_cuda();
2353 Ok(t)
2354 }
2355
2356 /// Create an image tensor with the given format.
2357 /// Allocate an image tensor from a declarative request — the
2358 /// full-featured constructor behind [`Self::image`] and friends.
2359 ///
2360 /// Adds the [`Compression`] request to the classic parameters:
2361 ///
2362 /// - a request with any CPU access other than [`CpuAccess::None`] is
2363 /// [`Error::InvalidArgument`] (CPU mapping pins the layout linear);
2364 /// - [`Compression::Scheme`] fails with [`Error::NotImplemented`] on
2365 /// platforms without vendor tile compression, and with
2366 /// [`Error::InvalidArgument`] when the device's native scheme or
2367 /// the format eligibility doesn't match;
2368 /// - [`Compression::Any`] never fails for compression reasons: it
2369 /// records the scheme when the allocation is eligible and
2370 /// otherwise resolves linear, incrementing
2371 /// [`compression_fallback_count`].
2372 ///
2373 /// The recorded outcome is readable via [`Tensor::compression`].
2374 pub fn image_desc(desc: &ImageDesc) -> Result<Self>
2375 where
2376 T: 'static,
2377 {
2378 let Some(t_dtype) = dtype_of::<T>() else {
2379 return Err(Error::InvalidArgument(
2380 "image_desc: element type has no DType mapping".into(),
2381 ));
2382 };
2383 if t_dtype != desc.dtype {
2384 return Err(Error::InvalidArgument(format!(
2385 "image_desc: desc.dtype is {:?} but the tensor element type is {t_dtype:?}",
2386 desc.dtype
2387 )));
2388 }
2389
2390 // Compression-request guards. CPU access pins the layout linear,
2391 // so a request combined with any declared access is a
2392 // contradiction the caller should hear about immediately.
2393 if desc.compression.is_some() && desc.access != CpuAccess::None {
2394 return Err(Error::InvalidArgument(format!(
2395 "image_desc: a compression request requires CpuAccess::None (declared {:?}) — CPU mapping pins the layout linear",
2396 desc.access
2397 )));
2398 }
2399 if let Some(Compression::Scheme(requested)) = desc.compression {
2400 #[cfg(not(target_os = "android"))]
2401 {
2402 return Err(Error::NotImplemented(format!(
2403 "image_desc: Compression::Scheme({requested:?}) — no vendor tile compression on this platform (request Compression::Any for a portable fallback)"
2404 )));
2405 }
2406 #[cfg(target_os = "android")]
2407 {
2408 if matches!(
2409 desc.memory,
2410 Some(TensorMemory::Mem) | Some(TensorMemory::Shm) | Some(TensorMemory::Pbo)
2411 ) {
2412 return Err(Error::InvalidArgument(format!(
2413 "image_desc: Compression::Scheme({requested:?}) requires hardware memory (TensorMemory::Dma or auto-select), got {:?}",
2414 desc.memory
2415 )));
2416 }
2417 if !crate::ahardwarebuffer_layout::compression_eligible(desc.format, desc.dtype) {
2418 return Err(Error::InvalidArgument(format!(
2419 "image_desc: ({:?}, {:?}) is not compression-eligible (RGBA8888 u8/i8 initially)",
2420 desc.format, desc.dtype
2421 )));
2422 }
2423 let device = crate::ahardwarebuffer::device_compression_scheme();
2424 if device != Some(requested) {
2425 return Err(Error::InvalidArgument(format!(
2426 "image_desc: Compression::Scheme({requested:?}) but the device's native scheme is {device:?}"
2427 )));
2428 }
2429 }
2430 }
2431
2432 // A compression request implies a hardware pipeline, so auto
2433 // memory promotes to the platform's zero-copy allocation first
2434 // (`Tensor::image` only takes the AHardwareBuffer/IOSurface path
2435 // under an explicit Dma request). `Scheme` propagates the Dma
2436 // failure — the caller demanded a layout only that allocator can
2437 // produce; `Any` falls back to plain auto-select.
2438 #[allow(unused_mut)]
2439 let mut t = match (desc.memory, desc.compression) {
2440 (None, Some(request)) => {
2441 match Self::image(
2442 desc.width,
2443 desc.height,
2444 desc.format,
2445 Some(TensorMemory::Dma),
2446 desc.access,
2447 ) {
2448 Ok(t) => t,
2449 Err(e) if matches!(request, Compression::Scheme(_)) => return Err(e),
2450 Err(_) => Self::image(desc.width, desc.height, desc.format, None, desc.access)?,
2451 }
2452 }
2453 (memory, _) => Self::image(desc.width, desc.height, desc.format, memory, desc.access)?,
2454 };
2455
2456 // Record best knowledge / count fallbacks. Only an Android
2457 // hardware-only AHardwareBuffer allocation can actually hold a
2458 // vendor tile layout; everywhere else an Any request resolves
2459 // linear and is counted.
2460 if let Some(request) = desc.compression {
2461 #[cfg(target_os = "android")]
2462 {
2463 let eligible =
2464 crate::ahardwarebuffer_layout::compression_eligible(desc.format, desc.dtype);
2465 let scheme = crate::ahardwarebuffer::device_compression_scheme();
2466 let is_ahb = t.memory() == TensorMemory::Dma;
2467 match (request, scheme) {
2468 (_, Some(s)) if eligible && is_ahb => {
2469 t.set_compression_unchecked(Some(s));
2470 }
2471 (Compression::Scheme(requested), _) => {
2472 // Pre-validated eligible + scheme match, so the only
2473 // way here is the allocation resolving off-AHB.
2474 return Err(Error::InvalidOperation(format!(
2475 "image_desc: Compression::Scheme({requested:?}) requested but the allocation resolved to {:?} (not an AHardwareBuffer)",
2476 t.memory()
2477 )));
2478 }
2479 (Compression::Any, _) => {
2480 note_compression_fallback(&format!(
2481 "({:?}, {:?}) {}x{}: eligible={eligible}, scheme={scheme:?}, memory={:?}",
2482 desc.format,
2483 desc.dtype,
2484 desc.width,
2485 desc.height,
2486 t.memory()
2487 ));
2488 }
2489 }
2490 }
2491 #[cfg(not(target_os = "android"))]
2492 {
2493 debug_assert!(matches!(request, Compression::Any));
2494 let _ = request;
2495 note_compression_fallback(&format!(
2496 "({:?}, {:?}) {}x{}: no vendor tile compression on this platform",
2497 desc.format, desc.dtype, desc.width, desc.height
2498 ));
2499 }
2500 }
2501 Ok(t)
2502 }
2503
2504 /// Allocate an image tensor of `width` × `height` in `format`.
2505 ///
2506 /// This is the common-case image constructor. The shape follows from the
2507 /// format's layout — packed `[H, W, C]`, planar `[C, H, W]`, semi-planar
2508 /// `[H + chroma_rows, W]` — and the format is recorded on the tensor, so
2509 /// the image crate and the codec can work from it without out-of-band
2510 /// parameters.
2511 ///
2512 /// `memory` requests a backend, or `None` to auto-select. Auto-select for
2513 /// images is **DMA → Mem**: unlike [`Tensor::new`] it does not try SHM,
2514 /// which offers an image nothing that Mem doesn't (it is not
2515 /// GPU-importable, and Mem always succeeds). Ask for
2516 /// [`TensorMemory::Shm`] explicitly if you need it.
2517 ///
2518 /// `access` declares what the **CPU** intends to do with the buffer and is
2519 /// required; see [`CpuAccess`]. Hardware access — GPU, NPU, ISP, codec —
2520 /// is always implied and never declared. Prefer the narrowest declaration
2521 /// that fits: `None` keeps the allocation eligible for Android vendor tile
2522 /// compression, and anything else pins the layout linear.
2523 ///
2524 /// DMA-backed images always get a **64-byte-aligned row stride**, in every
2525 /// layout, because Mali and Vivante reject an EGLImage import at an
2526 /// unaligned pitch. Read the real pitch back with
2527 /// [`Tensor::effective_row_stride`] rather than assuming `width × bpp`.
2528 ///
2529 /// For a compression request or any other less common knob, build an
2530 /// [`ImageDesc`] and call [`Tensor::image_desc`].
2531 ///
2532 /// # Errors
2533 ///
2534 /// - [`Error::InvalidArgument`] if `width` × `height` is not a valid size
2535 /// for `format` (for example odd dimensions where the format's chroma
2536 /// subsampling forbids them), or, on macOS with an explicit
2537 /// [`TensorMemory::Dma`], if the format cannot be expressed as an
2538 /// image-formatted IOSurface at this width. The message names the
2539 /// aligned width to use.
2540 /// - Whatever the chosen backend returns if the allocation itself fails —
2541 /// for an explicit `memory` request there is no fallback.
2542 ///
2543 /// # Example
2544 ///
2545 /// ```rust
2546 /// use edgefirst_tensor::{CpuAccess, PixelFormat, Tensor, TensorMemory};
2547 ///
2548 /// # fn main() -> Result<(), edgefirst_tensor::Error> {
2549 /// // A camera frame the GPU converts and the CPU never touches.
2550 /// let frame = Tensor::<u8>::image(1920, 1080, PixelFormat::Nv12,
2551 /// Some(TensorMemory::Mem), CpuAccess::None)?;
2552 /// assert_eq!(frame.format(), Some(PixelFormat::Nv12));
2553 /// # Ok(())
2554 /// # }
2555 /// ```
2556 pub fn image(
2557 width: usize,
2558 height: usize,
2559 format: PixelFormat,
2560 memory: Option<TensorMemory>,
2561 access: CpuAccess,
2562 ) -> Result<Self>
2563 where
2564 T: 'static,
2565 {
2566 // Shape comes from the shared `PixelFormat::image_shape` helper (packed /
2567 // planar / semi-planar NV12·NV16). NV12 supports odd dimensions via the
2568 // `H + ceil(H/2)` combined-plane height.
2569 // The `T: 'static` bound is required by the macOS IOSurface path below.
2570 let shape = format.image_shape(width, height).ok_or_else(|| {
2571 Error::InvalidArgument(format!(
2572 "invalid dimensions {width}x{height} for format {format:?}"
2573 ))
2574 })?;
2575
2576 // macOS Dma path: allocate a format-aware IOSurface (FourCC +
2577 // 2D dimensions) so the GL backend can bind it via
2578 // `EGL_ANGLE_iosurface_client_buffer`. Without this, the IOSurface
2579 // would default to a generic byte buffer (FourCC 'L008') and
2580 // ANGLE would reject the import with `EGL_BAD_ATTRIBUTE`.
2581 //
2582 // Guard: IOSurface rounds `bytes_per_row` up to 64-byte alignment.
2583 // If the natural row pitch (`width * channels * sizeof(T)`) is not
2584 // already 64-byte aligned, the padded allocation cannot be mapped
2585 // as a contiguous packed tensor — CPU reads/writes would use the
2586 // wrong stride.
2587 //
2588 // Explicit-Dma contract: when the caller passes
2589 // `Some(TensorMemory::Dma)` they have asked for an
2590 // **image-formatted IOSurface**. Silently downgrading to the
2591 // generic 'L008' byte-bag when alignment fails buries the
2592 // mismatch — the caller only finds out hours later when ANGLE
2593 // (or any GL importer) rejects the bind with
2594 // `EGL_BAD_ATTRIBUTE`. Same anti-pattern bit us previously on
2595 // Mali GPUs with DMA-BUF padding. The right behaviour is to
2596 // fail loudly here with the alignment requirement spelled out
2597 // so the caller can either pick aligned dimensions, request
2598 // SHM/Mem explicitly, or pass `memory=None` for auto-select.
2599 #[cfg(any(target_os = "macos", target_os = "ios"))]
2600 if matches!(memory, Some(TensorMemory::Dma)) {
2601 // For planar formats the IOSurface stacks channels
2602 // vertically (channels * height rows), so the row stride is
2603 // single-channel width * sizeof(T). Packed formats keep the
2604 // natural width * channels * sizeof(T) stride.
2605 let natural_row_bytes = match format.layout() {
2606 PixelLayout::Planar => width * std::mem::size_of::<T>(),
2607 _ => width * format.channels() * std::mem::size_of::<T>(),
2608 };
2609 // A format with a real IOSurface FourCC (RGBA/BGRA/YUYV packed,
2610 // GREY/NV12/NV16/NV24 as R8) tolerates a non-64-aligned natural
2611 // pitch: the surface is allocated with its own 64-aligned
2612 // `bytes_per_row`, the tensor records that stride below, and a CPU
2613 // map iterates rows correctly via the strided-map path while the GL
2614 // import uses the surface's pitch directly — fully zero-copy.
2615 // Planar (the F16 RGBA16F packing) is consumed flat as
2616 // `[1, C, H, W]` with no stride, so it still requires an aligned
2617 // pitch; and formats without a FourCC would fall through to a
2618 // generic byte-bag GL can't bind. Both fail loudly rather than
2619 // silently downgrade.
2620 let has_image_fourcc = dtype_of::<T>()
2621 .and_then(|dt| crate::iosurface::image_iosurface_layout(format, dt))
2622 .is_some();
2623 let padded_ok = has_image_fourcc && format.layout() != PixelLayout::Planar;
2624 if !natural_row_bytes.is_multiple_of(64) && !padded_ok {
2625 let elem_size = std::mem::size_of::<T>();
2626 let per_pixel_bytes = match format.layout() {
2627 PixelLayout::Planar => elem_size.max(1),
2628 _ => format.channels().max(1) * elem_size.max(1),
2629 };
2630 // Compute the next 64-byte-aligned width by rounding the
2631 // natural row pitch up to the next multiple of 64 and
2632 // dividing back by per-pixel bytes. This handles every
2633 // `per_pixel_bytes` value correctly:
2634 //
2635 // * Divisors of 64 (1/2/4/8/16/32/64) → the suggestion
2636 // is always 64-byte aligned.
2637 // * Non-divisors of 64 (e.g. RGB u8 with 3 B/pixel) →
2638 // the next aligned row pitch may not be an integer
2639 // multiple of per_pixel_bytes (3 doesn't divide 64
2640 // in any way), so a "pad width to N" suggestion is
2641 // structurally impossible — omit the suggestion
2642 // instead of printing a wrong number.
2643 // * per_pixel_bytes > 64 → same situation, also
2644 // omitted; the previous formula divided by zero.
2645 //
2646 // The error always names the alignment requirement
2647 // verbatim and lists the two non-DMA alternatives so
2648 // the caller has at least one always-applicable fix.
2649 let aligned_row_bytes = natural_row_bytes.next_multiple_of(64);
2650 let pad_hint =
2651 if per_pixel_bytes > 0 && aligned_row_bytes.is_multiple_of(per_pixel_bytes) {
2652 let w = aligned_row_bytes / per_pixel_bytes;
2653 format!("Pad width to {w} (the next 64-byte-aligned stride), ")
2654 } else {
2655 String::new()
2656 };
2657 return Err(Error::InvalidArgument(format!(
2658 "Tensor::image: {format:?} {width}x{height} with element \
2659 size {elem_size} produces a {natural_row_bytes}-byte natural \
2660 row pitch, which is not 64-byte aligned. \
2661 IOSurface rounds bytes_per_row up to 64 bytes, so a \
2662 contiguous CPU map of this tensor would read garbage. \
2663 {pad_hint}pass memory=None to auto-fall-back to SHM, or \
2664 pass memory=Some(TensorMemory::Shm) or \
2665 Some(TensorMemory::Mem) explicitly."
2666 )));
2667 }
2668 // Alignment OK. Explicit-Dma contract: the caller asked for an
2669 // **image-formatted, GL-importable** IOSurface, so every failure
2670 // from here on is loud. The old behaviour fell through to the
2671 // generic 'L008' byte-bag `Some(other)` arm below, and the caller
2672 // only found out when the GL import rejected the bind with
2673 // `EGL_BAD_ATTRIBUTE` — the same silent-downgrade anti-pattern as
2674 // the alignment case above. (Semi-planar/Grey u8 map to 'L008'
2675 // *by design* in `image_iosurface_layout` — the R8-plane
2676 // representation the YUV shaders sample — so they return through
2677 // the mapped path and never reach these errors.)
2678 let dtype = dtype_of::<T>().ok_or_else(|| {
2679 Error::InvalidArgument(format!(
2680 "Tensor::image: element type {} has no DType, so no \
2681 image-formatted IOSurface exists. Pass memory=None or \
2682 Some(TensorMemory::Mem) for a CPU tensor.",
2683 std::any::type_name::<T>()
2684 ))
2685 })?;
2686 if crate::iosurface::image_iosurface_layout(format, dtype).is_none() {
2687 return Err(Error::InvalidArgument(format!(
2688 "Tensor::image: no zero-copy IOSurface mapping exists for \
2689 {format:?}/{dtype:?} on macOS/iOS (supported: \
2690 Rgba/Rgb @ U8/I8, Bgra/Yuyv/Grey/Nv12/Nv16/Nv24 @ U8, \
2691 Rgba/PlanarRgb/PlanarRgba @ F16). Pass memory=None to \
2692 auto-select, or Some(TensorMemory::Mem) explicitly, for \
2693 a CPU tensor."
2694 )));
2695 }
2696 // Packed RGB u8/i8 rides an RGBA8888 surface at (W*3/4, H) —
2697 // reject a width the texel packing cannot express up front
2698 // (mirrors the Android pre-guard).
2699 if format == PixelFormat::Rgb
2700 && matches!(dtype, DType::U8 | DType::I8)
2701 && packed_rgb888_layout(width, height).is_none()
2702 {
2703 return Err(Error::InvalidArgument(format!(
2704 "Tensor::image: Rgb {dtype:?} requires width%4==0 for the RGBA8888 \
2705 IOSurface packing (got width={width}). Pad the width, or pass \
2706 memory=Some(TensorMemory::Mem) for a CPU tensor."
2707 )));
2708 }
2709 let storage = TensorStorage::<T>::new_image_iosurface(
2710 width, height, format, dtype, &shape, None,
2711 )?;
2712 let mut t = Self::wrap(storage);
2713 t.format = Some(format);
2714 // IOSurface rounds `bytes_per_row` up to 64 bytes. When that
2715 // pitch exceeds the natural packed/planar row stride, record
2716 // it so CPU consumers iterate rows correctly (the GL import
2717 // already uses the surface's own pitch). For 64-aligned rows
2718 // — the common model-input case — the two match and no stride
2719 // is stored, leaving the flat mapping unchanged.
2720 if let TensorStorage::Dma(ref io) = t.storage {
2721 let bpr = io.bytes_per_row();
2722 if let Some(natural) = t.effective_row_stride() {
2723 if bpr > natural {
2724 t.set_row_stride_unchecked(bpr);
2725 }
2726 }
2727 }
2728 t.cpu_access = access;
2729 return Ok(t);
2730 }
2731
2732 // Android Dma path: allocate a format-aware AHardwareBuffer so the
2733 // GL backend can import it as an EGLImage
2734 // (`eglGetNativeClientBufferANDROID` → `eglCreateImageKHR`).
2735 //
2736 // Geometry: gralloc chooses the row pitch (`desc.stride`) at
2737 // allocation, and pads freely (validated on the Galaxy S26 Ultra,
2738 // where SnapAlloc pads the planar-F16 RGBA16F surface). A padded
2739 // pitch is recorded on the tensor so CPU maps iterate rows via the
2740 // strided-map path; the GL render uses the buffer's own pitch
2741 // through the EGLImage either way, so the GPU path stays fully
2742 // zero-copy. Consumers needing a FLAT layout (the future NPU
2743 // handoff's `[1, C, H, W]` contract) must check `row_stride()` and
2744 // repack when set — flatness is a per-device property here, unlike
2745 // macOS where IOSurface's 64-BYTE alignment keeps model-sized F16
2746 // surfaces naturally flat.
2747 #[cfg(target_os = "android")]
2748 if matches!(memory, Some(TensorMemory::Dma)) {
2749 // Explicit-Dma contract (mirrors the macOS block above): the
2750 // caller asked for an image-formatted, GL-importable
2751 // AHardwareBuffer, so an unmapped combination or a gralloc
2752 // refusal errors here instead of falling through to a BLOB
2753 // byte-bag the GL backend can never import — that downgrade
2754 // only surfaced as a silent per-frame CPU upload.
2755 let dtype = dtype_of::<T>().ok_or_else(|| {
2756 Error::InvalidArgument(format!(
2757 "Tensor::image: element type {} has no DType, so no \
2758 image-formatted AHardwareBuffer exists. Pass memory=None \
2759 or Some(TensorMemory::Mem) for a CPU tensor.",
2760 std::any::type_name::<T>()
2761 ))
2762 })?;
2763 // Planar F16 requires width % 4 == 0 for the RGBA16F
2764 // packing — reject up front with the requirement spelled
2765 // out (mirrors the macOS pre-allocation guard) instead of
2766 // falling through to a byte-bag GL cannot bind.
2767 if format.layout() == PixelLayout::Planar
2768 && dtype == DType::F16
2769 && packed_rgba16f_layout(format, dtype, width, height).is_none()
2770 {
2771 return Err(Error::InvalidArgument(format!(
2772 "Tensor::image: {format:?} F16 requires width%4==0 for the RGBA16F \
2773 AHardwareBuffer packing (got width={width}). Pad the width, or pass \
2774 memory=Some(TensorMemory::Mem) for a CPU tensor."
2775 )));
2776 }
2777 // Packed RGB u8/i8 rides an RGBA8888 surface at (W*3/4, H) —
2778 // the same whole-texel constraint as the F16 packing above.
2779 if format == PixelFormat::Rgb
2780 && matches!(dtype, DType::U8 | DType::I8)
2781 && packed_rgb888_layout(width, height).is_none()
2782 {
2783 return Err(Error::InvalidArgument(format!(
2784 "Tensor::image: Rgb {dtype:?} requires width%4==0 for the RGBA8888 \
2785 AHardwareBuffer packing (got width={width}). Pad the width, or pass \
2786 memory=Some(TensorMemory::Mem) for a CPU tensor."
2787 )));
2788 }
2789 if crate::ahardwarebuffer::image_ahardwarebuffer_layout(format, dtype).is_none() {
2790 return Err(Error::InvalidArgument(format!(
2791 "Tensor::image: no zero-copy AHardwareBuffer mapping exists \
2792 for {format:?}/{dtype:?} on Android (Grey/NV* need \
2793 R8-format buffers, API 29+; the HAL floor is API 26 — see \
2794 `image_ahardwarebuffer_layout`). Pass memory=None to \
2795 auto-select, or Some(TensorMemory::Mem) explicitly, for a \
2796 CPU tensor; camera NV12 stays zero-copy by wrapping the \
2797 camera's own AHardwareBuffer instead of allocating one."
2798 )));
2799 }
2800 let storage = TensorStorage::<T>::new_image_ahardwarebuffer(
2801 width, height, format, dtype, &shape, None, access,
2802 )?;
2803 let mut t = Self::wrap(storage);
2804 t.format = Some(format);
2805 if let TensorStorage::Dma(ref ahb) = t.storage {
2806 // gralloc chooses the row pitch (Qualcomm's
2807 // SnapAlloc pads e.g. the 160-px-wide RGBA16F
2808 // surface of a 640-wide planar F16 target).
2809 // When it exceeds the natural stride, record
2810 // it so CPU consumers iterate rows correctly
2811 // via `effective_row_stride()` — the GPU
2812 // renders through the EGLImage at the
2813 // buffer's real pitch regardless, so the
2814 // render stays fully zero-copy. Consumers
2815 // that need the FLAT `[1, C, H, W]` layout
2816 // (the future NPU handoff) must check
2817 // `row_stride()` is unset and repack — or
2818 // pick an aligned width — rather than assume
2819 // flatness (see the module docs).
2820 let bpr = ahb.bytes_per_row();
2821 if let Some(natural) = t.effective_row_stride() {
2822 if bpr > natural {
2823 log::debug!(
2824 "Tensor::image: gralloc padded the {format:?} \
2825 AHardwareBuffer pitch to {bpr} bytes (natural \
2826 {natural}); recording row stride"
2827 );
2828 t.set_row_stride_unchecked(bpr);
2829 }
2830 }
2831 }
2832 t.cpu_access = access;
2833 return Ok(t);
2834 }
2835
2836 // Compute the **64-byte-aligned** row stride for every image layout.
2837 //
2838 // Embedded GPUs reject `eglCreateImage` DMA-BUF imports whose row pitch
2839 // is not 64-byte aligned: Mali returns `EGL_BAD_ALLOC`, Vivante
2840 // `EGL_BAD_ACCESS`. This bit packed RGBA/RGB destinations at odd widths
2841 // AND at even non-multiple-of-16 widths (e.g. 321→1284, 322→1288 bytes —
2842 // neither divisible by 64), so an odd-source → RGBA convert failed on
2843 // imx95/imx8mp while succeeding on V3D/Tegra. Semi-planar already aligned
2844 // here; we now align packed and planar identically so every image()
2845 // allocation is GPU-importable regardless of width.
2846 //
2847 // The per-layout natural pitch and total row count:
2848 // * SemiPlanar `[total_h, width]` — pitch = even(width)·elem, rows = total_h
2849 // * Packed `[height, width, ch]` — pitch = width·ch·elem, rows = height
2850 // * Planar `[ch, height, width]` — pitch = width·elem, rows = ch·height
2851 // Allocation byte size = `aligned_stride · total_rows` (NOT the shape
2852 // product, which reflects only the logical width and under-allocates the
2853 // padding on odd / unaligned widths).
2854 let elem = std::mem::size_of::<T>();
2855 let channels = format.channels();
2856 let (natural_stride, total_rows) = match format.layout() {
2857 PixelLayout::SemiPlanar => (width.next_multiple_of(2) * elem, shape[0]),
2858 PixelLayout::Packed => (width * channels * elem, height),
2859 PixelLayout::Planar => (width * elem, channels * height),
2860 };
2861 let aligned_stride = natural_stride.next_multiple_of(64);
2862 let semi = format.layout() == PixelLayout::SemiPlanar;
2863
2864 // DMA buffers MUST carry a 64-aligned row pitch — Mali/Vivante reject a
2865 // DMA-BUF EGLImage whose pitch is not 64-aligned. Semi-planar also needs
2866 // the aligned pitch on every backend (its chroma-plane offset math
2867 // assumes it). Packed/planar on host-only memory (Mem/Shm) keep the
2868 // natural tight pitch so the many flat CPU consumers are unaffected.
2869 let host_stride = if semi { aligned_stride } else { natural_stride };
2870 let host_byte_size = host_stride * total_rows;
2871 #[cfg(target_os = "linux")]
2872 let dma_byte_size = aligned_stride * total_rows;
2873
2874 // `used_stride` is the actual row pitch of the storage created below.
2875 let (storage, used_stride) = match memory {
2876 #[cfg(target_os = "linux")]
2877 Some(TensorMemory::Dma) => (
2878 TensorStorage::<T>::new_dma_with_byte_size(&shape, dma_byte_size, None)?,
2879 aligned_stride,
2880 ),
2881 #[cfg(unix)]
2882 Some(TensorMemory::Shm) => (
2883 TensorStorage::<T>::new_shm_with_byte_size(&shape, host_byte_size, None)?,
2884 host_stride,
2885 ),
2886 Some(TensorMemory::Mem) => (
2887 TensorStorage::<T>::new_mem_with_byte_size(&shape, host_byte_size, None)?,
2888 host_stride,
2889 ),
2890 #[allow(unused_variables)]
2891 Some(other) => {
2892 // PBO and any future variants: fall through to standard new().
2893 return {
2894 let mut t = Self::new(&shape, Some(other), None)?;
2895 t.format = Some(format);
2896 t.cpu_access = access;
2897 Ok(t)
2898 };
2899 }
2900 None => {
2901 // Auto-select priority: DMA → Mem (DMA gets the 64-aligned
2902 // pitch; the Mem fallback keeps the tight host pitch, so the
2903 // recorded stride matches the storage used). Shm is NOT
2904 // auto-selected — it offers no advantage over Mem for an
2905 // in-process image and Mem always succeeds, so it sits below Mem
2906 // and is reached only via an explicit `TensorMemory::Shm`.
2907 #[cfg(target_os = "linux")]
2908 {
2909 match TensorStorage::<T>::new_dma_with_byte_size(&shape, dma_byte_size, None) {
2910 Ok(s) => (s, aligned_stride),
2911 Err(_) => (
2912 TensorStorage::<T>::new_mem_with_byte_size(
2913 &shape,
2914 host_byte_size,
2915 None,
2916 )?,
2917 host_stride,
2918 ),
2919 }
2920 }
2921 #[cfg(not(target_os = "linux"))]
2922 {
2923 (
2924 TensorStorage::<T>::new_mem_with_byte_size(&shape, host_byte_size, None)?,
2925 host_stride,
2926 )
2927 }
2928 }
2929 };
2930
2931 let mut t = Self::wrap(storage);
2932 t.format = Some(format);
2933 // Record the row stride when it exceeds the natural tight pitch (padding
2934 // is present — DMA packed/planar at an unaligned width, or always for
2935 // semi-planar), mirroring the IOSurface path above. Aligned-width and
2936 // host-only packed/planar images keep their flat layout with no explicit
2937 // stride; `effective_row_stride()` then falls back to the identical
2938 // computed pitch. When padding IS present, consumers must iterate rows by
2939 // `effective_row_stride()` to skip it.
2940 if semi || used_stride > natural_stride {
2941 t.set_row_stride_unchecked(used_stride);
2942 }
2943 debug_assert!(
2944 t.row_stride.is_some() || !semi,
2945 "image() must always set row_stride for semi-planar tensors"
2946 );
2947 t.cpu_access = access;
2948 #[cfg(target_os = "linux")]
2949 t.try_init_dma_cuda();
2950 Ok(t)
2951 }
2952
2953 /// Create a DMA-backed image tensor with an explicit row stride that
2954 /// may exceed the natural `width * channels * sizeof(T)` pitch.
2955 ///
2956 /// Used for image tensors that need GPU pitch alignment padding: the
2957 /// underlying DMA-BUF is sized to `row_stride * height` bytes, but
2958 /// the tensor's logical shape stays at `[height, width, channels]`.
2959 /// `width()` / `height()` / `shape()` continue to report the
2960 /// user-requested values; the padding is visible only via
2961 /// `row_stride()` / `effective_row_stride()` and is automatically
2962 /// propagated to the GL backend's EGLImage import so Mali Valhall
2963 /// accepts the buffer.
2964 ///
2965 /// # Supported formats
2966 ///
2967 /// Currently only **packed** pixel layouts (RGBA8, BGRA8, RGB888,
2968 /// Grey, etc.) are supported — the formats the GL backend uses as
2969 /// render destinations. Semi-planar formats (NV12, NV16) come from
2970 /// external allocators (camera capture, video decoders) and are
2971 /// imported via `TensorDyn::from_fd` + `set_row_stride`, which
2972 /// already supports padded strides.
2973 ///
2974 /// # Supported memory
2975 ///
2976 /// Currently only `TensorMemory::Dma` is supported. PBO and Mem
2977 /// storage don't go through EGLImage import so they don't need
2978 /// pitch alignment; if you pass any other memory type this returns
2979 /// `NotImplemented`. `None` (auto-select) is treated as `Dma`.
2980 ///
2981 /// # Errors
2982 ///
2983 /// - `InvalidArgument` if `row_stride_bytes < width * channels * sizeof(T)`
2984 /// (the requested stride would not fit a single row)
2985 /// - `NotImplemented` for non-packed formats or non-DMA memory
2986 /// - `IoError` if the DMA-heap allocation fails (propagated from
2987 /// `DmaTensor::new_with_byte_size`)
2988 pub fn image_with_stride(
2989 width: usize,
2990 height: usize,
2991 format: PixelFormat,
2992 row_stride_bytes: usize,
2993 memory: Option<TensorMemory>,
2994 access: CpuAccess,
2995 ) -> Result<Self> {
2996 #[cfg(not(target_os = "linux"))]
2997 let _ = access;
2998 // DMA backing (the only thing this constructor produces) is
2999 // Linux-only. On macOS/BSD/Windows the non-Linux block below is
3000 // the only compiled body and returns `NotImplemented` directly;
3001 // on Linux the non-Linux block is cfg-removed and the function
3002 // falls through to the real validation + allocation path. Each
3003 // target compiles exactly one of the two blocks, and the block
3004 // serves as the function's tail expression in both cases — so
3005 // neither needs an explicit `return` (avoids
3006 // `clippy::needless_return` on the macOS CI gate).
3007 #[cfg(not(target_os = "linux"))]
3008 {
3009 let _ = (width, height, format, row_stride_bytes, memory);
3010 Err(Error::NotImplemented(
3011 "image_with_stride requires DMA support (Linux only)".to_owned(),
3012 ))
3013 }
3014
3015 #[cfg(target_os = "linux")]
3016 {
3017 if format.layout() != PixelLayout::Packed {
3018 return Err(Error::NotImplemented(format!(
3019 "Tensor::image_with_stride only supports packed pixel layouts, got {format:?}"
3020 )));
3021 }
3022 let elem = std::mem::size_of::<T>();
3023 let min_stride = width
3024 .checked_mul(format.channels())
3025 .and_then(|p| p.checked_mul(elem))
3026 .ok_or_else(|| {
3027 Error::InvalidArgument(format!(
3028 "image_with_stride: width {width} × channels {} × sizeof::<T>={elem} \
3029 overflows usize",
3030 format.channels()
3031 ))
3032 })?;
3033 if row_stride_bytes < min_stride {
3034 return Err(Error::InvalidArgument(format!(
3035 "image_with_stride: row_stride {row_stride_bytes} < minimum {min_stride} \
3036 ({width} px × {} ch × {elem} B)",
3037 format.channels()
3038 )));
3039 }
3040 let total_byte_size = row_stride_bytes.checked_mul(height).ok_or_else(|| {
3041 Error::InvalidArgument(format!(
3042 "image_with_stride: row_stride {row_stride_bytes} × height {height} overflows usize"
3043 ))
3044 })?;
3045
3046 let shape = vec![height, width, format.channels()];
3047
3048 let storage = match memory {
3049 Some(TensorMemory::Dma) | None => {
3050 TensorStorage::<T>::new_dma_with_byte_size(&shape, total_byte_size, None)?
3051 }
3052 Some(other) => {
3053 return Err(Error::NotImplemented(format!(
3054 "image_with_stride: only TensorMemory::Dma is supported, got {other:?}"
3055 )));
3056 }
3057 };
3058
3059 let mut t = Self::wrap(storage);
3060 t.format = Some(format);
3061 t.row_stride = Some(row_stride_bytes);
3062 t.cpu_access = access;
3063 // Match new()/from_fd(): a DMA tensor must attempt CUDA external-
3064 // memory import so a strided DMA buffer is also zero-copy
3065 // CUDA-mappable (no-op when libcudart is absent).
3066 t.try_init_dma_cuda();
3067 Ok(t)
3068 }
3069 }
3070
3071 /// Attach format metadata to an existing tensor.
3072 ///
3073 /// # Arguments
3074 ///
3075 /// * `format` - The pixel format to attach
3076 ///
3077 /// # Returns
3078 ///
3079 /// `Ok(())` on success, with the format stored as metadata on the tensor.
3080 ///
3081 /// # Errors
3082 ///
3083 /// Returns `Error::InvalidShape` if the tensor shape is incompatible with
3084 /// the format's layout (packed expects `[H, W, C]`, planar expects
3085 /// `[C, H, W]`, semi-planar expects `[H*k, W]` with format-specific
3086 /// height constraints).
3087 pub fn set_format(&mut self, format: PixelFormat) -> Result<()> {
3088 let shape = self.shape();
3089 match format.layout() {
3090 PixelLayout::Packed => {
3091 if shape.len() != 3 || shape[2] != format.channels() {
3092 return Err(Error::InvalidShape(format!(
3093 "packed format {format:?} expects [H, W, {}], got {shape:?}",
3094 format.channels()
3095 )));
3096 }
3097 }
3098 PixelLayout::Planar => {
3099 if shape.len() != 3 || shape[0] != format.channels() {
3100 return Err(Error::InvalidShape(format!(
3101 "planar format {format:?} expects [{}, H, W], got {shape:?}",
3102 format.channels()
3103 )));
3104 }
3105 }
3106 PixelLayout::SemiPlanar => {
3107 if shape.len() != 2 {
3108 return Err(Error::InvalidShape(format!(
3109 "semi-planar format {format:?} expects [H*k, W], got {shape:?}"
3110 )));
3111 }
3112 match format {
3113 // Combined-plane height is `H + ceil(H/2)` (luma + chroma
3114 // rows). For even H that is `3H/2` (≡ 0 mod 3); for odd H it
3115 // is `(3H+1)/2` (≡ 2 mod 3). Only totals ≡ 1 mod 3 are
3116 // unreachable, so reject just those — odd-height NV12 is
3117 // valid (e.g. 725 rows for a 483-tall image).
3118 PixelFormat::Nv12 if shape[0] % 3 == 1 => {
3119 return Err(Error::InvalidShape(format!(
3120 "NV12 contiguous shape[0] must be H + ceil(H/2) for some height; \
3121 {} is unreachable (≡ 1 mod 3)",
3122 shape[0]
3123 )));
3124 }
3125 PixelFormat::Nv16 if !shape[0].is_multiple_of(2) => {
3126 return Err(Error::InvalidShape(format!(
3127 "NV16 contiguous shape[0] must be even, got {}",
3128 shape[0]
3129 )));
3130 }
3131 // NV24 (4:4:4): combined-plane height is 3H (Y + 2H chroma).
3132 PixelFormat::Nv24 if !shape[0].is_multiple_of(3) => {
3133 return Err(Error::InvalidShape(format!(
3134 "NV24 contiguous shape[0] must be a multiple of 3 (= 3H), got {}",
3135 shape[0]
3136 )));
3137 }
3138 _ => {}
3139 }
3140 }
3141 }
3142 // Clear stored stride/offset when format changes — they may be invalid
3143 // for the new format. Caller must re-set after changing format.
3144 if self.format != Some(format) {
3145 self.row_stride = None;
3146 self.plane_offset = None;
3147 match self.storage {
3148 TensorStorage::Mem(ref mut m) => m.set_offset(0),
3149 #[cfg(target_os = "linux")]
3150 TensorStorage::Dma(ref mut dma) => dma.mmap_offset = 0,
3151 _ => {}
3152 }
3153 }
3154 self.format = Some(format);
3155 Ok(())
3156 }
3157
3158 /// Set this tensor's logical dimensions and pixel format to a decoded
3159 /// image, reusing the existing allocation. The shape is derived from the
3160 /// format layout; fails with `Error::InsufficientCapacity` if the
3161 /// allocation cannot hold `width`×`height` in `format`, or
3162 /// `Error::InvalidArgument` if the dimensions are invalid for the format.
3163 ///
3164 /// For NV12/NV16/NV24 the buffer width is rounded up to even (a chroma-plane
3165 /// interleaving requirement); the true odd width is reported by the decoder
3166 /// in `ImageInfo` and trimmed by a `convert()` crop. See
3167 /// [`PixelFormat::image_shape`].
3168 ///
3169 /// When the backing has a fixed physical row pitch (an IOSurface's
3170 /// 64-aligned `bytesPerRow`) that exceeds the new format's natural row
3171 /// stride — i.e. a reused max-sized pool tensor reconfigured to a smaller
3172 /// image — the physical pitch is preserved as the tensor's `row_stride`.
3173 /// This keeps the **physical grid** (allocation stride/surface) fixed while
3174 /// the **logical ROI** (this image's W×H) changes, so the decode writes rows
3175 /// at the surface's real stride and the GPU samples them at the same stride.
3176 /// Exact-sized buffers (pitch == natural) stay tightly packed unchanged.
3177 pub fn configure_image(
3178 &mut self,
3179 width: usize,
3180 height: usize,
3181 format: PixelFormat,
3182 ) -> Result<()> {
3183 let shape = format.image_shape(width, height).ok_or_else(|| {
3184 Error::InvalidArgument(format!(
3185 "invalid dimensions {width}x{height} for format {format:?}"
3186 ))
3187 })?;
3188 // Capture the pre-existing row stride before `set_format` clears it.
3189 // For pool tensors that were allocated at a larger width (e.g. 1920-wide
3190 // pool decoding a 789-wide image), this preserves the backing pitch so
3191 // rows are still written at the correct physical stride.
3192 let prior_stride = self.row_stride;
3193
3194 self.storage.set_logical_shape(&shape)?;
3195 self.set_format(format)?; // clears any stale row_stride
3196
3197 // Restore the correct row pitch for the new geometry. A DMA buffer of
3198 // ANY layout, and a semi-planar buffer on any backing, MUST carry a
3199 // 64-byte-aligned pitch: Mali/Vivante reject a DMA-BUF EGLImage whose
3200 // row pitch is not 64-aligned (`EGL_BAD_ALLOC` / `EGL_BAD_ACCESS`), and
3201 // semi-planar chroma-offset math assumes it. This mirrors the pitch
3202 // `Tensor::image()` allocates, so a recycled pool buffer imports
3203 // identically to a fresh one — without it a `configure_image()`'d packed
3204 // buffer (e.g. Y800 96-wide → tight pitch 96) failed convert() on
3205 // imx95/imx8mp via the texture-upload fallback while the fresh oracle
3206 // (aligned 128) succeeded. Packed/planar on host-only memory (Mem/Shm)
3207 // keep the tight pitch so flat CPU consumers stay unaffected — matching
3208 // `image()`'s `host_stride` rule.
3209 let elem = std::mem::size_of::<T>();
3210 let channels = format.channels();
3211 let (min_stride, total_rows) = match format.layout() {
3212 PixelLayout::SemiPlanar => (width.next_multiple_of(2) * elem, shape[0]),
3213 PixelLayout::Packed => (width * channels * elem, height),
3214 PixelLayout::Planar => (width * elem, channels * height),
3215 };
3216 let needs_align = self.storage.memory() == TensorMemory::Dma
3217 || format.layout() == PixelLayout::SemiPlanar;
3218
3219 let active_stride = if let Some(pitch) = self.storage.backing_row_stride() {
3220 // macOS IOSurface: use the surface's native pitch.
3221 let natural = self.effective_row_stride().unwrap_or(0);
3222 if pitch > natural {
3223 self.set_row_stride_unchecked(pitch);
3224 pitch
3225 } else {
3226 natural
3227 }
3228 } else if needs_align {
3229 // Priority:
3230 // 1. Prior stride (pool reuse): if the pre-existing stride is
3231 // 64-aligned, >= this layout's minimum row, and fits the
3232 // allocation, keep it. This is the hot-loop reuse case (large
3233 // pool, small image).
3234 // 2. Compute a fresh 64-aligned stride for the current width.
3235 let aligned = min_stride.next_multiple_of(64);
3236 let capacity = self.storage.capacity_bytes();
3237
3238 let candidate = if let Some(ps) = prior_stride {
3239 if ps >= min_stride && ps % 64 == 0 && ps * total_rows <= capacity {
3240 ps
3241 } else {
3242 aligned
3243 }
3244 } else {
3245 aligned
3246 };
3247
3248 if candidate * total_rows <= capacity {
3249 self.set_row_stride_unchecked(candidate);
3250 candidate
3251 } else {
3252 // Shouldn't happen for legitimate pools, but don't crash.
3253 self.effective_row_stride().unwrap_or(0)
3254 }
3255 } else {
3256 self.effective_row_stride().unwrap_or(0)
3257 };
3258
3259 // Ensure the active stride fits the allocation. A pool reconfigured to a
3260 // wider image than its backing would silently SIGBUS on any subsequent
3261 // map/write — catch it here instead.
3262 if needs_align && active_stride > 0 {
3263 let needed = active_stride * total_rows;
3264 let capacity = self.storage.capacity_bytes();
3265 if needed > capacity {
3266 return Err(Error::InsufficientCapacity { needed, capacity });
3267 }
3268 }
3269 Ok(())
3270 }
3271
3272 /// Allocate an image tensor sized to hold up to `width`×`height` in
3273 /// `format`, reusable for any smaller image via `configure_image`.
3274 pub fn image_with_capacity(
3275 width: usize,
3276 height: usize,
3277 format: PixelFormat,
3278 memory: Option<TensorMemory>,
3279 access: CpuAccess,
3280 ) -> Result<Self>
3281 where
3282 T: 'static,
3283 {
3284 Self::image(width, height, format, memory, access)
3285 }
3286
3287 /// Pixel format (None if not an image).
3288 pub fn format(&self) -> Option<PixelFormat> {
3289 self.format
3290 }
3291
3292 /// Image width (None if not an image).
3293 pub fn width(&self) -> Option<usize> {
3294 let fmt = self.format?;
3295 let shape = self.shape();
3296 match fmt.layout() {
3297 PixelLayout::Packed => Some(shape[1]),
3298 PixelLayout::Planar => Some(shape[2]),
3299 PixelLayout::SemiPlanar => Some(shape[1]),
3300 }
3301 }
3302
3303 /// Image height (None if not an image).
3304 ///
3305 /// For semi-planar formats the combined-plane shape row count is divided
3306 /// by the format's luma-to-total ratio to recover logical height. This
3307 /// returns the exact logical height (including odd heights) only because
3308 /// the logical dimensions are tracked separately from the physical shape —
3309 /// `configure_image` stores the actual `(width, height)` in the format's
3310 /// `image_shape`, which round-trips losslessly via these accessors.
3311 pub fn height(&self) -> Option<usize> {
3312 let fmt = self.format?;
3313 let shape = self.shape();
3314 match fmt.layout() {
3315 PixelLayout::Packed => Some(shape[0]),
3316 PixelLayout::Planar => Some(shape[1]),
3317 PixelLayout::SemiPlanar => {
3318 if self.is_multiplane() {
3319 Some(shape[0])
3320 } else {
3321 match fmt {
3322 PixelFormat::Nv12 => Some(shape[0] * 2 / 3),
3323 PixelFormat::Nv16 => Some(shape[0] / 2),
3324 PixelFormat::Nv24 => Some(shape[0] / 3),
3325 _ => None,
3326 }
3327 }
3328 }
3329 }
3330 }
3331
3332 /// Create from separate Y and UV planes (multiplane NV12/NV16).
3333 pub fn from_planes(luma: Tensor<T>, chroma: Tensor<T>, format: PixelFormat) -> Result<Self> {
3334 if format.layout() != PixelLayout::SemiPlanar {
3335 return Err(Error::InvalidArgument(format!(
3336 "from_planes requires a semi-planar format, got {format:?}"
3337 )));
3338 }
3339 if chroma.format.is_some() || chroma.chroma.is_some() {
3340 return Err(Error::InvalidArgument(
3341 "chroma tensor must be a raw tensor (no format or chroma metadata)".into(),
3342 ));
3343 }
3344 let luma_shape = luma.shape();
3345 let chroma_shape = chroma.shape();
3346 if luma_shape.len() != 2 || chroma_shape.len() != 2 {
3347 return Err(Error::InvalidArgument(format!(
3348 "from_planes expects 2D shapes, got luma={luma_shape:?} chroma={chroma_shape:?}"
3349 )));
3350 }
3351 if luma_shape[1] != chroma_shape[1] {
3352 return Err(Error::InvalidArgument(format!(
3353 "luma width {} != chroma width {}",
3354 luma_shape[1], chroma_shape[1]
3355 )));
3356 }
3357 match format {
3358 PixelFormat::Nv12 => {
3359 if luma_shape[0] % 2 != 0 {
3360 return Err(Error::InvalidArgument(format!(
3361 "NV12 requires even luma height, got {}",
3362 luma_shape[0]
3363 )));
3364 }
3365 if chroma_shape[0] != luma_shape[0] / 2 {
3366 return Err(Error::InvalidArgument(format!(
3367 "NV12 chroma height {} != luma height / 2 ({})",
3368 chroma_shape[0],
3369 luma_shape[0] / 2
3370 )));
3371 }
3372 }
3373 PixelFormat::Nv16 => {
3374 if chroma_shape[0] != luma_shape[0] {
3375 return Err(Error::InvalidArgument(format!(
3376 "NV16 chroma height {} != luma height {}",
3377 chroma_shape[0], luma_shape[0]
3378 )));
3379 }
3380 }
3381 // NV24's chroma plane is full-resolution (2×-wide interleaved UV),
3382 // which the equal-width plane check above doesn't model. Multiplane
3383 // NV24 is unused (the JPEG decoder emits a contiguous NV24 buffer),
3384 // so it's not supported here yet.
3385 _ => {
3386 return Err(Error::InvalidArgument(format!(
3387 "from_planes only supports NV12 and NV16 (NV24 multiplane not yet \
3388 supported — use a contiguous NV24 tensor), got {format:?}"
3389 )));
3390 }
3391 }
3392
3393 Ok(Tensor {
3394 storage: luma.storage,
3395 format: Some(format),
3396 chroma: Some(Box::new(chroma)),
3397 row_stride: luma.row_stride,
3398 plane_offset: luma.plane_offset,
3399 quantization: luma.quantization,
3400 // A multiplane tensor spans two DMA-BUFs (luma + chroma); CUDA
3401 // external-memory import is per-fd, so there is no single device
3402 // pointer for the composite. Any CUDA handle the luma plane carried
3403 // is intentionally dropped — consumers needing CUDA access to
3404 // multiplane data must import each plane independently.
3405 cuda: None,
3406 colorimetry: luma.colorimetry,
3407 cpu_access: luma.cpu_access,
3408 compression: luma.compression,
3409 // A composed multiplane tensor is a whole image, not a sub-view.
3410 view_origin: None,
3411 })
3412 }
3413
3414 /// Whether this tensor uses separate plane allocations.
3415 pub fn is_multiplane(&self) -> bool {
3416 self.chroma.is_some()
3417 }
3418
3419 /// Access the chroma plane for multiplane semi-planar images.
3420 pub fn chroma(&self) -> Option<&Tensor<T>> {
3421 self.chroma.as_deref()
3422 }
3423
3424 /// Mutable access to the chroma plane for multiplane semi-planar images.
3425 pub fn chroma_mut(&mut self) -> Option<&mut Tensor<T>> {
3426 self.chroma.as_deref_mut()
3427 }
3428
3429 /// Row stride in bytes (`None` = tightly packed).
3430 pub fn row_stride(&self) -> Option<usize> {
3431 self.row_stride
3432 }
3433
3434 /// Effective row stride in bytes: the stored stride if set, otherwise the
3435 /// minimum stride computed from the format, width, and element size.
3436 /// Returns `None` only when no format is set and no explicit stride was
3437 /// stored via [`set_row_stride`](Self::set_row_stride).
3438 ///
3439 /// **GREY note:** `effective_row_stride()` for a GREY tensor returns the
3440 /// tight `width` bytes (no padding), which is what `normalize_to_numpy` and
3441 /// the CPU convert path expect. The codec's internal `native_row_stride`
3442 /// (64-byte-aligned) is used only during decoding and is not propagated to
3443 /// the tensor's stored stride, so callers reading via
3444 /// `effective_row_stride()` always see the tight value for GREY.
3445 pub fn effective_row_stride(&self) -> Option<usize> {
3446 if let Some(s) = self.row_stride {
3447 return Some(s);
3448 }
3449 let fmt = self.format?;
3450 let w = self.width()?;
3451 let elem = std::mem::size_of::<T>();
3452 Some(match fmt.layout() {
3453 PixelLayout::Packed => w * fmt.channels() * elem,
3454 PixelLayout::Planar => w * elem,
3455 // Semi-planar: minimum stride must cover the even width so the
3456 // interleaved chroma columns are byte-aligned on odd-width images.
3457 PixelLayout::SemiPlanar => w.next_multiple_of(2) * elem,
3458 })
3459 }
3460
3461 /// Copy the tensor's logical bytes into `dst`, compacting away any
3462 /// recorded row-stride padding.
3463 ///
3464 /// The flatness helper for NPU handoff: consumers that need a FLAT
3465 /// layout (e.g. `[1, C, H, W]` for NNAPI/LiteRT when the runtime cannot
3466 /// take a padded pitch) call this when [`row_stride`](Self::row_stride)
3467 /// is `Some` — on a tight tensor it degenerates to one memcpy, so it is
3468 /// safe to call unconditionally. `dst.len()` must equal the tight byte
3469 /// footprint (`shape` product × element size). Zero-copy consumers
3470 /// should prefer the buffer handle + [`Tensor::effective_row_stride`]
3471 /// (Self::effective_row_stride) and skip this copy entirely.
3472 pub fn copy_to_flat(&self, dst: &mut [u8]) -> Result<()> {
3473 let tight_bytes = crate::ahardwarebuffer_layout::checked_shape_bytes::<T>(self.shape())?;
3474 if dst.len() != tight_bytes {
3475 return Err(Error::InvalidArgument(format!(
3476 "copy_to_flat: dst is {} bytes but the tensor's tight \
3477 footprint is {tight_bytes} bytes (shape {:?})",
3478 dst.len(),
3479 self.shape()
3480 )));
3481 }
3482 let map = self.map()?;
3483 // SAFETY: T is a plain numeric type (crate-wide bound); viewing the
3484 // mapped elements as bytes is sound.
3485 let src: &[u8] = unsafe {
3486 std::slice::from_raw_parts(
3487 map.as_slice().as_ptr() as *const u8,
3488 std::mem::size_of_val(map.as_slice()),
3489 )
3490 };
3491 let Some(stride) = self.row_stride else {
3492 // Tight layout: the mapped window is exactly the logical bytes.
3493 let got = src.len().min(tight_bytes);
3494 if got < tight_bytes {
3495 return Err(Error::InvalidOperation(format!(
3496 "copy_to_flat: mapped {got} bytes < tight footprint {tight_bytes}"
3497 )));
3498 }
3499 dst.copy_from_slice(&src[..tight_bytes]);
3500 return Ok(());
3501 };
3502 // Strided: row count follows the strided-map convention (planar
3503 // stacks C planes of H rows; packed/semi-planar use shape[0]), and
3504 // the logical row is tight_bytes / rows for every layout.
3505 let rows = match self.format.map(|f| f.layout()) {
3506 Some(PixelLayout::Planar) => {
3507 let s = self.shape();
3508 if s.len() < 2 {
3509 return Err(Error::InvalidOperation(
3510 "copy_to_flat: strided planar tensor requires [C, H, W] shape".into(),
3511 ));
3512 }
3513 s[0].checked_mul(s[1]).ok_or_else(|| {
3514 Error::InvalidOperation(format!(
3515 "copy_to_flat: planar rows {} × {} overflows usize",
3516 s[0], s[1]
3517 ))
3518 })?
3519 }
3520 _ => *self.shape().first().ok_or_else(|| {
3521 Error::InvalidOperation("copy_to_flat: tensor has an empty shape".into())
3522 })?,
3523 };
3524 if rows == 0 || !tight_bytes.is_multiple_of(rows) {
3525 return Err(Error::InvalidOperation(format!(
3526 "copy_to_flat: tight footprint {tight_bytes} does not divide \
3527 into {rows} rows"
3528 )));
3529 }
3530 let row_bytes = tight_bytes / rows;
3531 let need = (rows - 1)
3532 .checked_mul(stride)
3533 .and_then(|b| b.checked_add(row_bytes))
3534 .ok_or_else(|| {
3535 Error::InvalidOperation(format!(
3536 "copy_to_flat: stride {stride} × rows {rows} overflows usize"
3537 ))
3538 })?;
3539 if src.len() < need {
3540 return Err(Error::InvalidOperation(format!(
3541 "copy_to_flat: mapped {} bytes but strided rows need {need}",
3542 src.len()
3543 )));
3544 }
3545 for r in 0..rows {
3546 dst[r * row_bytes..(r + 1) * row_bytes]
3547 .copy_from_slice(&src[r * stride..r * stride + row_bytes]);
3548 }
3549 Ok(())
3550 }
3551
3552 /// Set the row stride in bytes for externally allocated buffers with
3553 /// row padding (e.g. V4L2 or GStreamer allocators).
3554 ///
3555 /// The stride is propagated to the EGL DMA-BUF import attributes so
3556 /// the GPU interprets the padded buffer layout correctly. Must be
3557 /// called after [`set_format`](Self::set_format) and before the tensor
3558 /// is first passed to `ImageProcessor::convert`. The stored stride
3559 /// is cleared automatically if the pixel format is later changed.
3560 ///
3561 /// No stride-vs-buffer-size validation is performed because the
3562 /// backing allocation size is not reliably known: external DMA-BUFs
3563 /// may be over-allocated by the allocator, and internal tensors store
3564 /// a logical (unpadded) shape. An incorrect stride will be caught by
3565 /// the EGL driver at import time.
3566 ///
3567 /// # Arguments
3568 ///
3569 /// * `stride` - Row stride in bytes. Must be >= the minimum stride for
3570 /// the format (width * channels * sizeof(T) for packed,
3571 /// width * sizeof(T) for planar/semi-planar).
3572 ///
3573 /// # Errors
3574 ///
3575 /// * `InvalidArgument` if no pixel format is set on this tensor
3576 /// * `InvalidArgument` if `stride` is less than the minimum for the
3577 /// format and width
3578 pub fn set_row_stride(&mut self, stride: usize) -> Result<()> {
3579 let fmt = self.format.ok_or_else(|| {
3580 Error::InvalidArgument("cannot set row_stride without a pixel format".into())
3581 })?;
3582 let w = self.width().ok_or_else(|| {
3583 Error::InvalidArgument("cannot determine width for row_stride validation".into())
3584 })?;
3585 let elem = std::mem::size_of::<T>();
3586 let min_stride = match fmt.layout() {
3587 PixelLayout::Packed => w * fmt.channels() * elem,
3588 PixelLayout::Planar => w * elem,
3589 // Semi-planar: minimum must cover even width for chroma alignment.
3590 PixelLayout::SemiPlanar => w.next_multiple_of(2) * elem,
3591 };
3592 if stride < min_stride {
3593 return Err(Error::InvalidArgument(format!(
3594 "row_stride {stride} < minimum {min_stride} for {fmt:?} at width {w}"
3595 )));
3596 }
3597 self.row_stride = Some(stride);
3598 Ok(())
3599 }
3600
3601 /// Set the row stride without format validation.
3602 ///
3603 /// Use this for raw sub-tensors (e.g. chroma planes) that don't carry
3604 /// format metadata. The caller is responsible for ensuring the stride
3605 /// is valid.
3606 pub fn set_row_stride_unchecked(&mut self, stride: usize) {
3607 self.row_stride = Some(stride);
3608 }
3609
3610 /// Builder-style variant of [`set_row_stride`](Self::set_row_stride),
3611 /// consuming and returning `self`.
3612 ///
3613 /// # Errors
3614 ///
3615 /// Same conditions as [`set_row_stride`](Self::set_row_stride).
3616 pub fn with_row_stride(mut self, stride: usize) -> Result<Self> {
3617 self.set_row_stride(stride)?;
3618 Ok(self)
3619 }
3620
3621 /// Byte offset within the DMA-BUF where image data starts (`None` = 0).
3622 pub fn plane_offset(&self) -> Option<usize> {
3623 self.plane_offset
3624 }
3625
3626 /// The parent-image snapshot if this tensor is a [`view`](Self::view)/
3627 /// [`batch`](Self::batch) sub-region; `None` for a whole tensor. The GL
3628 /// backend keys its import on the parent geometry and renders this view as a
3629 /// `glViewport`/`glScissor` ROI at `(x, y, width, height)`. See [`ViewOrigin`].
3630 pub fn view_origin(&self) -> Option<ViewOrigin> {
3631 self.view_origin
3632 }
3633
3634 /// Set the byte offset within the DMA-BUF where image data starts.
3635 ///
3636 /// Propagated to `EGL_DMA_BUF_PLANE0_OFFSET_EXT` on GPU import.
3637 /// Unlike [`set_row_stride`](Self::set_row_stride), no format is required
3638 /// since the offset is format-independent.
3639 pub fn set_plane_offset(&mut self, offset: usize) {
3640 self.plane_offset = Some(offset);
3641 // The offset consulted by `map()` lives inside the storage variant.
3642 // Keep it in sync with the wrapper field for every backing that
3643 // honors it (DMA and Mem); see also the clear sites in `set_format`
3644 // and `reshape`.
3645 match self.storage {
3646 TensorStorage::Mem(ref mut m) => m.set_offset(offset),
3647 #[cfg(target_os = "linux")]
3648 TensorStorage::Dma(ref mut dma) => dma.mmap_offset = offset,
3649 _ => {}
3650 }
3651 }
3652
3653 /// Colorimetry metadata (`None` = undefined; never auto-filled).
3654 /// The CPU access declared for this tensor at allocation (see
3655 /// [`CpuAccess`]). Views share their parent's declaration.
3656 pub fn cpu_access(&self) -> CpuAccess {
3657 self.cpu_access
3658 }
3659
3660 /// Set the declared CPU access without re-allocating — crate-private:
3661 /// the declaration must reflect the underlying allocation's real
3662 /// capabilities (constructors and importers set it; arbitrary widening
3663 /// would defeat the contract).
3664 // Only the Android AHardwareBuffer importer derives a declaration from
3665 // an existing allocation's usage bits today.
3666 #[cfg_attr(not(target_os = "android"), allow(dead_code))]
3667 pub(crate) fn set_cpu_access_unchecked(&mut self, access: CpuAccess) {
3668 self.cpu_access = access;
3669 }
3670
3671 /// The vendor tile-compression scheme recorded at allocation, or
3672 /// `None` for a linear layout. `Some` means the pixels live in a
3673 /// proprietary tile order: the row-stride accessors describe no
3674 /// meaningful linear layout and CPU maps are best-effort (see
3675 /// [`Compression`]). Only Android hardware-only allocations that
3676 /// requested compression record a scheme.
3677 pub fn compression(&self) -> Option<CompressionScheme> {
3678 self.compression
3679 }
3680
3681 /// Record the compression scheme — crate-private: recording is an
3682 /// allocation-time fact ([`Tensor::image_desc`] sets it; arbitrary
3683 /// mutation would misdescribe the physical layout).
3684 // Only the Android allocation path records a scheme today.
3685 #[cfg_attr(not(target_os = "android"), allow(dead_code))]
3686 pub(crate) fn set_compression_unchecked(&mut self, scheme: Option<CompressionScheme>) {
3687 self.compression = scheme;
3688 }
3689
3690 pub fn colorimetry(&self) -> Option<crate::Colorimetry> {
3691 self.colorimetry
3692 }
3693
3694 /// Attach/clear colorimetry metadata.
3695 pub fn set_colorimetry(&mut self, c: Option<crate::Colorimetry>) {
3696 self.colorimetry = c;
3697 }
3698
3699 /// Builder-style colorimetry attach.
3700 pub fn with_colorimetry(mut self, c: crate::Colorimetry) -> Self {
3701 self.colorimetry = Some(c);
3702 self
3703 }
3704
3705 /// Create a zero-copy sub-region view of this tensor's backing buffer.
3706 ///
3707 /// The returned tensor shares this tensor's allocation (no copy) and maps
3708 /// the window `[offset_bytes, offset_bytes + shape.product()*size_of::<T>())`
3709 /// measured from this tensor's own logical start. N sub-views into one
3710 /// parent can be written independently, enabling batched assembly into a
3711 /// single buffer. Identical semantics across `Mem` (shared `Arc`) and
3712 /// `Dma` (shared fd) backings.
3713 ///
3714 /// # Disjointness
3715 ///
3716 /// Independent writes are sound *only* when the windows do not overlap. The
3717 /// shared backing uses interior mutability (`UnsafeCell` cells), so two
3718 /// sub-views whose byte ranges intersect alias the same cells: writing one
3719 /// while reading or writing the other is a data race and therefore
3720 /// **undefined behaviour**. The caller is responsible for keeping the
3721 /// windows disjoint; this method does not check for overlap.
3722 ///
3723 /// # Errors
3724 ///
3725 /// - [`Error::InvalidOperation`] if the backing is not `Mem` or `Dma`, or
3726 /// if `offset_bytes` is not a multiple of `align_of::<T>()`.
3727 /// - [`Error::InsufficientCapacity`] / [`Error::InvalidSize`] if the window
3728 /// exceeds the parent allocation.
3729 pub(crate) fn subview(&self, offset_bytes: usize, shape: &[usize]) -> Result<Tensor<T>> {
3730 // Offset is absolute into the backing allocation: a sub-view of a
3731 // sub-view composes by adding this tensor's own offset.
3732 let abs_offset = self
3733 .plane_offset
3734 .unwrap_or(0)
3735 .checked_add(offset_bytes)
3736 .ok_or(Error::InvalidSize(offset_bytes))?;
3737 // Every backend exposes `view(offset, shape)` via `TensorTrait`, sharing
3738 // the resource AND `BufferIdentity` (unlike `from_fd`/`from_surface`,
3739 // which mint a fresh identity). The GL backend keys the import on the
3740 // shared identity so offset-distinct sub-views of one buffer reuse a
3741 // single import and address their window via `glViewport`. `Mem`/`Shm`
3742 // share via the allocation `Arc` / a cloned fd; `Pbo` via the GL-buffer
3743 // `Arc`; Linux DMA-BUF / macOS IOSurface via the shared fd / CFRetain.
3744 // `TensorStorage::view` performs the one remaining per-variant dispatch.
3745 let mut t = Tensor::wrap(self.storage.view(offset_bytes, shape)?);
3746 // Inherit the parent's image metadata so the view is a ready-to-use
3747 // sub-image (e.g. a `convert()` destination). The offset is applied
3748 // LAST because `set_format` deliberately clears it — the offset is a
3749 // structural property of the sub-region, not format-dependent metadata.
3750 if let Some(fmt) = self.format {
3751 t.set_format(fmt)?;
3752 }
3753 if let Some(rs) = self.row_stride {
3754 t.set_row_stride_unchecked(rs);
3755 }
3756 t.quantization = self.quantization.clone();
3757 // A sub-region of an image carries the parent's colorimetry — it is the
3758 // same pixels, same color encoding. Inherit it like the other image
3759 // metadata above so a sub-view is a faithful convert() source/target.
3760 t.set_colorimetry(self.colorimetry);
3761 // The declared CPU access is a property of the underlying
3762 // allocation, so every view shares the parent's declaration.
3763 t.cpu_access = self.cpu_access;
3764 // Likewise the recorded compression scheme: the view shares the
3765 // parent's physical layout.
3766 t.compression = self.compression;
3767 if abs_offset > 0 {
3768 t.set_plane_offset(abs_offset);
3769 }
3770 Ok(t)
3771 }
3772
3773 /// Borrow batch element `n` of a batched tensor as a zero-copy view.
3774 ///
3775 /// A batched tensor prepends `N` as the leading dimension over the
3776 /// per-element image layout (`[N, H, W, C]` packed or `[N, C, H, W]`
3777 /// planar) — `N` is over the whole per-element block regardless of
3778 /// `HWC`/`CHW`. `batch(n)` returns element `n`: the contiguous per-element
3779 /// region at byte offset `n * element_size`, sharing the parent's
3780 /// `BufferIdentity` and inheriting its format / row stride / colorimetry.
3781 /// `batch(0)` on a tensor with `N == 1` is equivalent to the whole tensor.
3782 ///
3783 /// # Errors
3784 ///
3785 /// - [`Error::BatchIndexOutOfBounds`] if `n >= N`.
3786 /// - [`Error::InvalidShape`] if the tensor is not batched (a formatted
3787 /// tensor whose rank lacks the leading `N`, or an empty shape).
3788 pub fn batch(&self, n: usize) -> Result<Tensor<T>> {
3789 let shape = self.shape();
3790 // With a format we know the exact per-element rank, so a missing leading
3791 // `N` is a misuse we reject rather than silently treating a spatial dim
3792 // as the batch. Raw tensors take shape[0] as `N` by contract.
3793 if let Some(fmt) = self.format {
3794 let elem_rank = match fmt.layout() {
3795 PixelLayout::SemiPlanar => 2,
3796 _ => 3,
3797 };
3798 if shape.len() != elem_rank + 1 {
3799 return Err(Error::InvalidShape(format!(
3800 "batch(): tensor is not batched ({fmt:?} expects a leading N over a \
3801 {elem_rank}-D element, got shape {shape:?})"
3802 )));
3803 }
3804 }
3805 let batch = *shape
3806 .first()
3807 .ok_or_else(|| Error::InvalidShape("batch(): empty shape".into()))?;
3808 if n >= batch {
3809 return Err(Error::BatchIndexOutOfBounds { index: n, batch });
3810 }
3811 let elem_shape: Vec<usize> = shape[1..].to_vec();
3812 let elem_count: usize = elem_shape.iter().product();
3813 let elem_bytes = elem_count
3814 .checked_mul(std::mem::size_of::<T>())
3815 .ok_or(Error::InvalidSize(elem_count))?;
3816 let offset = n.checked_mul(elem_bytes).ok_or(Error::InvalidSize(n))?;
3817 // For a packed `[N, H, W, C]` tensor the N tiles stack vertically in the
3818 // shared buffer, so the GL import sees one `(W, N*H)` parent and each
3819 // tile is the row-band at `y = n*H`. Snapshot that parent so the backend
3820 // imports once and renders the tile via `glViewport`. Non-packed
3821 // (planar/semi-planar) batching keeps the per-slot path for now (planar
3822 // NCHW tiling is a separate step); raw tensors have no pixel geometry.
3823 let view_origin = match self.format.map(|f| f.layout()) {
3824 Some(PixelLayout::Packed) => {
3825 let tile_h = elem_shape[0];
3826 let tile_w = elem_shape[1];
3827 // Per-row pitch of the tall `(W, N*H)` parent — padded stride if
3828 // set, else the tight row width. The GL import keys on this.
3829 let bpp = elem_shape[2] * std::mem::size_of::<T>();
3830 let parent_stride = self.effective_row_stride().unwrap_or(tile_w * bpp);
3831 Some(self.compose_view_origin(tile_w, batch * tile_h, parent_stride, 0, n * tile_h))
3832 }
3833 _ => None,
3834 };
3835 let mut t = self.subview(offset, &elem_shape)?;
3836 t.view_origin = view_origin;
3837 Ok(t)
3838 }
3839
3840 /// Borrow a rectangular spatial sub-region of an image tensor as a
3841 /// zero-copy view — the **destination/source crop** primitive.
3842 ///
3843 /// `region` is in pixels of the image's leading frame. The returned view
3844 /// shares the parent's `BufferIdentity` and addresses the sub-rectangle by
3845 /// offset + the **parent's** row pitch (so each row lands at the correct
3846 /// columns). `convert(src, &mut dst.view(rect), …)` renders into that
3847 /// sub-rectangle of `dst`; a letterbox fit then clears the view and renders
3848 /// the aspect-preserved content into its inner region. `view`/`batch`/the
3849 /// whole tensor are the one coherent destination model — there is no
3850 /// separate `dst_rect`.
3851 ///
3852 /// # Known limitation: bottom-right multi-row views
3853 ///
3854 /// A multi-row view is mapped as `parent_pitch × rows` bytes **from its own
3855 /// offset**, so the window it asks for overruns the allocation by
3856 /// `region.x * bpp` whenever the view is both offset horizontally *and*
3857 /// reaches the parent's last row. Such a view constructs fine but fails at
3858 /// [`map`](TensorTrait::map) time with [`Error::InsufficientCapacity`] —
3859 /// loudly, never as an out-of-bounds access. For a tiled destination this is
3860 /// the bottom row of tiles at every column but the first. The single-row
3861 /// case is already special-cased below (it keeps its own tight stride); the
3862 /// general fix is to clamp the last row's exposure to the view's own row
3863 /// bytes, which is not yet implemented.
3864 ///
3865 /// # Errors
3866 ///
3867 /// - [`Error::RegionOutOfBounds`] if `region` exceeds the image bounds.
3868 /// - [`Error::InvalidOperation`] if the tensor is not a packed-format image
3869 /// (planar/semi-planar spatial sub-rects are not a single strided window;
3870 /// use [`batch`](Self::batch) for batched planar tensors).
3871 pub fn view(&self, region: Region) -> Result<Tensor<T>> {
3872 let fmt = self.format.ok_or_else(|| {
3873 Error::InvalidOperation("view() requires a formatted image tensor".into())
3874 })?;
3875 if fmt.layout() != PixelLayout::Packed {
3876 return Err(Error::InvalidOperation(format!(
3877 "view() supports packed formats only (got {fmt:?}); use batch(n) for batched \
3878 planar tensors"
3879 )));
3880 }
3881 let w = self
3882 .width()
3883 .ok_or_else(|| Error::InvalidOperation("view(): tensor has no image width".into()))?;
3884 let h = self
3885 .height()
3886 .ok_or_else(|| Error::InvalidOperation("view(): tensor has no image height".into()))?;
3887 if !region.fits_within(w, h) {
3888 return Err(Error::RegionOutOfBounds {
3889 region,
3890 bounds: (w, h),
3891 });
3892 }
3893 let elem = std::mem::size_of::<T>();
3894 let bpp = fmt.channels() * elem;
3895 let stride = self.effective_row_stride().unwrap_or(w * bpp);
3896 let offset = region
3897 .y
3898 .checked_mul(stride)
3899 .and_then(|yo| yo.checked_add(region.x.checked_mul(bpp)?))
3900 .ok_or(Error::InvalidSize(region.y))?;
3901 let sub_shape = fmt
3902 .image_shape(region.width, region.height)
3903 .ok_or_else(|| Error::InvalidShape(format!("view(): invalid shape for {fmt:?}")))?;
3904 let mut t = self.subview(offset, &sub_shape)?;
3905 // A multi-row sub-rect must advance rows by the PARENT pitch so each row
3906 // addresses the correct columns. A single-row view uses its own tight
3907 // stride — the parent pitch would make the strided `map()` (which exposes
3908 // `stride × rows`) expose a trailing row that runs past the buffer tail
3909 // for an offset (x>0 / bottom) view. The GL backend does NOT rely on this
3910 // (single-row-tight) `row_stride`: it reads the parent pitch from
3911 // `view_origin.parent_row_stride` so its import/cache pitch stays
3912 // parent-consistent for views of any height — see `ViewOrigin`.
3913 let view_stride = if region.height > 1 {
3914 stride
3915 } else {
3916 region.width * bpp
3917 };
3918 t.set_row_stride_unchecked(view_stride);
3919 // Snapshot the parent `(w, h, row_stride)` so the GL backend imports the
3920 // parent once (keyed on the parent pitch, not this view's possibly-tight
3921 // single-row stride) and renders this sub-rect as a `glViewport`/
3922 // `glScissor` ROI at `(region.x, region.y)`. Composes when viewing an
3923 // existing view.
3924 t.view_origin = Some(self.compose_view_origin(w, h, stride, region.x, region.y));
3925 Ok(t)
3926 }
3927
3928 /// Build the [`ViewOrigin`] for a new sub-region of `self`. When `self` is a
3929 /// whole tensor the snapshot names `self` as the parent; when `self` is
3930 /// already a view, the snapshot keeps the **root** parent and accumulates the
3931 /// local origin so nested views still resolve to one import.
3932 fn compose_view_origin(
3933 &self,
3934 parent_width: usize,
3935 parent_height: usize,
3936 parent_row_stride: usize,
3937 x: usize,
3938 y: usize,
3939 ) -> ViewOrigin {
3940 match self.view_origin {
3941 Some(root) => ViewOrigin {
3942 parent_width: root.parent_width,
3943 parent_height: root.parent_height,
3944 parent_row_stride: root.parent_row_stride,
3945 x: root.x.saturating_add(x),
3946 y: root.y.saturating_add(y),
3947 },
3948 None => ViewOrigin {
3949 parent_width,
3950 parent_height,
3951 parent_row_stride,
3952 x,
3953 y,
3954 },
3955 }
3956 }
3957
3958 /// Downcast to PBO tensor reference (for GL backends).
3959 pub fn as_pbo(&self) -> Option<&PboTensor<T>> {
3960 match &self.storage {
3961 TensorStorage::Pbo(p) => Some(p),
3962 _ => None,
3963 }
3964 }
3965
3966 /// Downcast to DMA tensor reference (for EGL import, G2D).
3967 #[cfg(target_os = "linux")]
3968 pub fn as_dma(&self) -> Option<&DmaTensor<T>> {
3969 match &self.storage {
3970 TensorStorage::Dma(d) => Some(d),
3971 _ => None,
3972 }
3973 }
3974
3975 /// Borrow the DMA-BUF file descriptor backing this tensor.
3976 ///
3977 /// # Returns
3978 ///
3979 /// A borrowed reference to the DMA-BUF file descriptor, tied to `self`'s
3980 /// lifetime.
3981 ///
3982 /// # Errors
3983 ///
3984 /// Returns `Error::NotImplemented` if the tensor is not DMA-backed.
3985 #[cfg(target_os = "linux")]
3986 pub fn dmabuf(&self) -> Result<std::os::fd::BorrowedFd<'_>> {
3987 use std::os::fd::AsFd;
3988 match &self.storage {
3989 TensorStorage::Dma(dma) => Ok(dma.fd.as_fd()),
3990 _ => Err(Error::NotImplemented(format!(
3991 "dmabuf requires DMA-backed tensor, got {:?}",
3992 self.storage.memory()
3993 ))),
3994 }
3995 }
3996
3997 /// Construct a Tensor from a PBO tensor (for GL backends that allocate PBOs).
3998 pub fn from_pbo(pbo: PboTensor<T>) -> Self {
3999 Self {
4000 storage: TensorStorage::Pbo(pbo),
4001 format: None,
4002 chroma: None,
4003 row_stride: None,
4004 plane_offset: None,
4005 quantization: None,
4006 cuda: None,
4007 colorimetry: None,
4008 cpu_access: CpuAccess::ReadWrite,
4009 compression: None,
4010 view_origin: None,
4011 }
4012 }
4013
4014 /// The CUDA registration for this tensor, if any (set at creation on CUDA devices).
4015 pub fn cuda(&self) -> Option<&crate::cuda::CudaHandle> {
4016 self.cuda.as_ref()
4017 }
4018
4019 /// Attach a CUDA handle (called by ImageProcessor::create_image after registering a PBO).
4020 pub fn set_cuda_handle(&mut self, h: crate::cuda::CudaHandle) {
4021 self.cuda = Some(h);
4022 }
4023
4024 /// Fast-fail CUDA map: None (no GL routing) when no handle; else map (PBO routes to the GL worker).
4025 ///
4026 /// Returns a scoped [`CudaMap`] guard holding the raw CUDA device pointer
4027 /// for the duration of the mapping. For GL-buffer-backed tensors the unmap is deferred until the
4028 /// guard drops, freeing the PBO for the next `convert()` call. When no CUDA handle is attached
4029 /// (the common case for plain `Mem`/`DMA` tensors without CUDA registration), returns `None`
4030 /// immediately — no GL routing, no allocation.
4031 ///
4032 /// # Example — zero-copy CUDA input with host fallback
4033 ///
4034 /// ```no_run
4035 /// use edgefirst_tensor::{Tensor, TensorMemory, TensorTrait};
4036 /// # fn feed_tensorrt(_dptr: *mut std::ffi::c_void, _bytes: usize) {}
4037 /// # fn demo(t: &Tensor<f32>) {
4038 /// // Try the zero-copy CUDA device pointer first.
4039 /// if let Some(cuda) = t.cuda_map() {
4040 /// feed_tensorrt(cuda.device_ptr(), cuda.len());
4041 /// // `cuda` (a CudaMap guard) unmaps when it goes out of scope, freeing
4042 /// // the GPU buffer for the next convert().
4043 /// } else {
4044 /// // Fall back to the host mapping when no CUDA handle is attached.
4045 /// let _host = t.map().expect("host map fallback must succeed");
4046 /// // `_host` is a TensorMap<f32> that derefs to &[f32].
4047 /// }
4048 /// # }
4049 /// ```
4050 pub fn cuda_map(&self) -> Option<crate::cuda::CudaMap<'_>> {
4051 self.cuda.as_ref()?.map()
4052 }
4053
4054 /// Attempt to attach a CUDA `ExternalMemory` handle for DMA-backed tensors.
4055 ///
4056 /// On a CUDA-capable host, imports the DMA-BUF fd via
4057 /// `cudaImportExternalMemory(OpaqueFd)` and maps it to a device pointer.
4058 /// Sets `self.cuda` to a persistent `ExternalMem` handle on success. No-op
4059 /// if CUDA is unavailable, the tensor is not DMA-backed, or a handle is
4060 /// already set. Import failure is silently ignored — the tensor remains
4061 /// usable without a CUDA handle.
4062 ///
4063 /// # RUNTIME-UNVALIDATED
4064 ///
4065 /// No test platform has both `/dev/dma_heap` and a CUDA device. ABI is
4066 /// layout-asserted vs. CUDA 12.6 `driver_types.h`; the mechanism is proven
4067 /// by gpu-probe O5 on Orin. Best-effort: tensor creation never fails here.
4068 #[cfg(target_os = "linux")]
4069 pub fn try_init_dma_cuda(&mut self) {
4070 // Fast-path: already imported, CUDA not available, or not a DMA tensor.
4071 if self.cuda.is_some() || !crate::cuda::is_cuda_available() {
4072 return;
4073 }
4074 let (raw_fd, buf_size) = match &self.storage {
4075 TensorStorage::Dma(dma) => {
4076 use std::os::fd::AsRawFd;
4077 (dma.fd.as_raw_fd(), dma.buf_size)
4078 }
4079 _ => return,
4080 };
4081 if let Some((ext, dptr)) = crate::cuda::import_dma_fd(raw_fd, buf_size) {
4082 self.cuda = Some(crate::cuda::CudaHandle::new_external(ext, dptr, buf_size));
4083 }
4084 }
4085}
4086
4087// Quantization accessors — type-gated to integer element types via the
4088// sealed `IntegerType` trait. Calling `.quantization()` on a `Tensor<f32>`
4089// produces a compile error, not a runtime one.
4090impl<T> Tensor<T>
4091where
4092 T: IntegerType + Num + Clone + fmt::Debug + Send + Sync,
4093{
4094 /// Quantization metadata for this tensor, if set.
4095 pub fn quantization(&self) -> Option<&Quantization> {
4096 self.quantization.as_ref()
4097 }
4098
4099 /// Attach quantization metadata to this tensor. Validates against the
4100 /// tensor's shape — returns [`Error::QuantizationInvalid`] on any
4101 /// inconsistency (mismatched scale/zp lengths, out-of-range axis, etc.).
4102 pub fn set_quantization(&mut self, q: Quantization) -> Result<()> {
4103 q.validate(self.shape())?;
4104 self.quantization = Some(q);
4105 Ok(())
4106 }
4107
4108 /// Builder-style variant of [`Self::set_quantization`]. Consumes `self`
4109 /// and returns `Result<Self>` — on success yields the tensor with the
4110 /// attached quantization; on validation failure returns
4111 /// [`Error::QuantizationInvalid`] and drops `self` (the tensor is not
4112 /// returned in the error arm).
4113 pub fn with_quantization(mut self, q: Quantization) -> Result<Self> {
4114 self.set_quantization(q)?;
4115 Ok(self)
4116 }
4117
4118 /// Clear any quantization metadata on this tensor.
4119 pub fn clear_quantization(&mut self) {
4120 self.quantization = None;
4121 }
4122}
4123
4124impl<T> TensorTrait<T> for Tensor<T>
4125where
4126 T: Num + Clone + fmt::Debug + Send + Sync,
4127{
4128 fn new(shape: &[usize], name: Option<&str>) -> Result<Self>
4129 where
4130 Self: Sized,
4131 {
4132 Self::new(shape, None, name)
4133 }
4134
4135 #[cfg(unix)]
4136 fn from_fd(fd: std::os::fd::OwnedFd, shape: &[usize], name: Option<&str>) -> Result<Self>
4137 where
4138 Self: Sized,
4139 {
4140 #[cfg_attr(not(target_os = "linux"), allow(unused_mut))]
4141 let mut t = Self::wrap(TensorStorage::from_fd(fd, shape, name)?);
4142 // Best-effort CUDA external memory import for DMA-backed tensors.
4143 // RUNTIME-UNVALIDATED: see try_init_dma_cuda().
4144 #[cfg(target_os = "linux")]
4145 t.try_init_dma_cuda();
4146 Ok(t)
4147 }
4148
4149 #[cfg(unix)]
4150 fn clone_fd(&self) -> Result<std::os::fd::OwnedFd> {
4151 self.storage.clone_fd()
4152 }
4153
4154 fn memory(&self) -> TensorMemory {
4155 self.storage.memory()
4156 }
4157
4158 fn name(&self) -> String {
4159 self.storage.name()
4160 }
4161
4162 fn shape(&self) -> &[usize] {
4163 self.storage.shape()
4164 }
4165
4166 fn reshape(&mut self, shape: &[usize]) -> Result<()> {
4167 if self.chroma.is_some() {
4168 return Err(Error::InvalidOperation(
4169 "cannot reshape a multiplane tensor — decompose planes first".into(),
4170 ));
4171 }
4172 self.storage.reshape(shape)?;
4173 self.format = None;
4174 self.row_stride = None;
4175 self.plane_offset = None;
4176 match self.storage {
4177 TensorStorage::Mem(ref mut m) => m.set_offset(0),
4178 #[cfg(target_os = "linux")]
4179 TensorStorage::Dma(ref mut dma) => dma.mmap_offset = 0,
4180 _ => {}
4181 }
4182 Ok(())
4183 }
4184
4185 fn map_with(&self, access: CpuAccess) -> Result<TensorMap<T>> {
4186 let _span = tracing::trace_span!(
4187 "tensor.map",
4188 memory = ?self.storage.memory(),
4189 ?access,
4190 )
4191 .entered();
4192 if access == CpuAccess::None {
4193 return Err(Error::InvalidArgument(
4194 "map_with(CpuAccess::None) is not a mappable direction — use \
4195 map_read()/map_write()/map_mut()"
4196 .into(),
4197 ));
4198 }
4199 // Declared-vs-requested telemetry (all platforms): mapping beyond
4200 // the allocation-time declaration is best-effort — tolerated where
4201 // the backing is CPU-mappable regardless (Mem/Shm/dma-buf/
4202 // IOSurface), refused by the Android backend for CpuAccess::None
4203 // buffers — but always loud and counted, never silent.
4204 if !self.cpu_access.covers(access) {
4205 note_unplanned_cpu_access(
4206 self.buffer_identity().id(),
4207 &format!("{:?}", self.storage.memory()),
4208 "map access exceeds the declared CpuAccess",
4209 );
4210 }
4211 // CPU mapping of a strided tensor exposes the full padded buffer
4212 // (`row_stride × rows`) so callers can iterate rows via
4213 // `effective_row_stride()` without running past the slice. This is sound
4214 // only when the HAL owns and can size-check the allocation:
4215 //
4216 // * Self-allocated Mem / Shm tensors (any platform) — the backing
4217 // `Vec` / shm segment is sized by `capacity_bytes()`, checked here.
4218 // * Self-allocated DMA tensors (Linux) — pitch padding from
4219 // `image_with_stride()`; checked against the DMA-BUF `buf_size`.
4220 //
4221 // * Self-allocated PBO tensors (any platform with GL) — the GL buffer
4222 // is sized by `capacity_bytes()` and may carry 64-byte row padding;
4223 // the JPEG decoder mmaps it and convert() reads it, both iterating
4224 // by `row_stride`. Checked against the PBO capacity below.
4225 //
4226 // Foreign DMA-BUFs (`from_fd()` + `set_row_stride()`, the V4L2 /
4227 // GStreamer case) and IOSurface are rejected: their layout comes from an
4228 // external allocator / GPU driver the HAL cannot validate for a strided
4229 // CPU view, and they are intended for the GPU path. (Earlier this
4230 // rejected *all* non-Linux strided maps with "DMA backing is Linux-only"
4231 // — that was an unimplemented path, not a platform limit; HAL-owned
4232 // Mem/Shm/PBO are trivially mappable and now are.)
4233 if let Some(stride) = self.row_stride {
4234 // Rows sit at `stride`-byte spacing. The row count is the first
4235 // shape dim for packed `[H, W, C]` and semi-planar `[H*k, W]`,
4236 // but planar `[C, H, W]` stacks C planes of H rows — its surface
4237 // row count is `C × H` (`shape[0]` alone would expose a 3-row
4238 // window and truncate the map; first hit by Android planar-F16
4239 // AHardwareBuffers, whose gralloc pads the pitch — macOS/Linux
4240 // planar pitches happen to be naturally aligned so no stride was
4241 // ever recorded there).
4242 let rows = match self.format.map(|f| f.layout()) {
4243 Some(PixelLayout::Planar) => {
4244 let s = self.shape();
4245 if s.len() < 2 {
4246 return Err(Error::InvalidOperation(
4247 "Tensor::map: strided planar mapping requires [C, H, W] shape".into(),
4248 ));
4249 }
4250 s[0].checked_mul(s[1]).ok_or_else(|| {
4251 Error::InvalidOperation(format!(
4252 "Tensor::map: planar rows {} × {} overflows usize",
4253 s[0], s[1]
4254 ))
4255 })?
4256 }
4257 _ => *self.shape().first().ok_or_else(|| {
4258 Error::InvalidOperation(
4259 "Tensor::map: strided mapping requires a non-empty shape".into(),
4260 )
4261 })?,
4262 };
4263 let total_bytes = stride.checked_mul(rows).ok_or_else(|| {
4264 Error::InvalidOperation(format!(
4265 "Tensor::map: row_stride {stride} × rows {rows} overflows usize"
4266 ))
4267 })?;
4268
4269 match &self.storage {
4270 #[cfg(target_os = "linux")]
4271 TensorStorage::Dma(dma) if !dma.is_imported => {
4272 // `set_row_stride()` only validates `stride >= min_stride`,
4273 // not that `stride × rows` fits the DMA-BUF, so re-check
4274 // here — mapping past `buf_size` would SIGBUS on access.
4275 let available_bytes = dma.buf_size.saturating_sub(dma.mmap_offset);
4276 if total_bytes > available_bytes {
4277 return Err(Error::InvalidOperation(format!(
4278 "Tensor::map: strided mapping needs {total_bytes} bytes \
4279 but DMA buffer only has {available_bytes} available \
4280 (buf_size={}, mmap_offset={}, stride={stride}, rows={rows}); \
4281 the row_stride was likely set larger than the original allocation",
4282 dma.buf_size, dma.mmap_offset
4283 )));
4284 }
4285 return dma
4286 .map_with_byte_size(total_bytes, access)
4287 .map(TensorMap::Dma);
4288 }
4289 TensorStorage::Mem(mem) => {
4290 let capacity = self.storage.capacity_bytes();
4291 if total_bytes > capacity {
4292 return Err(Error::InsufficientCapacity {
4293 needed: total_bytes,
4294 capacity,
4295 });
4296 }
4297 return mem.map_with_byte_size(total_bytes, access);
4298 }
4299 #[cfg(unix)]
4300 TensorStorage::Shm(shm) => {
4301 let capacity = self.storage.capacity_bytes();
4302 if total_bytes > capacity {
4303 return Err(Error::InsufficientCapacity {
4304 needed: total_bytes,
4305 capacity,
4306 });
4307 }
4308 return shm.map_with_byte_size(total_bytes, access);
4309 }
4310 // macOS/iOS: `TensorStorage::Dma` is the IOSurface. The lock yields
4311 // the full surface base address, and the row pitch
4312 // (`IOSurfaceGetBytesPerRow`) is known from the API for both
4313 // self-allocated and imported surfaces — unlike a foreign
4314 // DMA-BUF — so a strided CPU view is sound and zero-copy.
4315 #[cfg(any(target_os = "macos", target_os = "ios"))]
4316 TensorStorage::Dma(io) => {
4317 // A sub-view's window is `buf_size − view_offset`; the strided
4318 // span must fit the window, not the whole surface.
4319 let available = io.buf_size.saturating_sub(io.view_offset);
4320 if total_bytes > available {
4321 return Err(Error::InsufficientCapacity {
4322 needed: total_bytes,
4323 capacity: available,
4324 });
4325 }
4326 return io.map_with_byte_size(total_bytes, access);
4327 }
4328 // Android: `TensorStorage::Dma` is the AHardwareBuffer. The lock
4329 // yields the full buffer base address, and the row pitch is
4330 // known from the allocator-filled descriptor — so a strided CPU
4331 // view is sound and zero-copy, same as IOSurface.
4332 #[cfg(target_os = "android")]
4333 TensorStorage::Dma(ahb) => {
4334 // A sub-view's window is `buf_size − view_offset`; the strided
4335 // span must fit the window, not the whole buffer.
4336 let available = ahb.buf_size.saturating_sub(ahb.view_offset);
4337 if total_bytes > available {
4338 return Err(Error::InsufficientCapacity {
4339 needed: total_bytes,
4340 capacity: available,
4341 });
4342 }
4343 return ahb.map_with_byte_size(total_bytes, access);
4344 }
4345 TensorStorage::Pbo(pbo) => {
4346 // PBO: the GPU-side allocation may have a padded row stride
4347 // (e.g. 64-byte aligned). Expose the full padded buffer so a
4348 // CPU producer (JPEG decoder) and a strided convert source
4349 // can iterate rows via `effective_row_stride()` without
4350 // running past the slice — the logical `pbo.map()` view would
4351 // stop after `shape.product()` and lose bytes past row 0.
4352 // A sub-view's window is `capacity − view_offset`.
4353 let available = pbo.capacity_bytes().saturating_sub(pbo.view_offset);
4354 if total_bytes > available {
4355 return Err(Error::InsufficientCapacity {
4356 needed: total_bytes,
4357 capacity: available,
4358 });
4359 }
4360 return pbo.map_with_byte_size(total_bytes, access);
4361 }
4362 // Reachable on Linux for an IMPORTED DMA-BUF (the `Dma` arm above
4363 // is guarded `if !dma.is_imported`). On macOS/Windows every
4364 // storage variant is matched explicitly, so this catch-all is
4365 // unreachable there — allow it rather than cfg-gating per platform.
4366 #[allow(unreachable_patterns)]
4367 _ => {
4368 return Err(Error::InvalidOperation(
4369 "CPU mapping of strided tensors is supported only for HAL-allocated \
4370 Mem/Shm (any platform), self-allocated DMA (Linux), IOSurface \
4371 (macOS), and PBO; imported DMA-BUF without self-allocation is \
4372 GPU-path only"
4373 .into(),
4374 ));
4375 }
4376 }
4377 }
4378 // Offset tensors are supported for storages that apply the offset
4379 // inside their own `map()`: DMA (`DmaMap`/IOSurface adjust the mapped
4380 // base), Mem (`MemMap` adjusts the slice base), Shm (`ShmMap` adjusts
4381 // the slice base), and PBO (the staged copy starts at the offset). Every
4382 // self-allocated backing now carries a sub-region concept via `view`, so
4383 // a non-zero offset is honoured rather than rejected.
4384 if self.plane_offset.is_some_and(|o| o > 0) {
4385 let supported = matches!(self.storage, TensorStorage::Mem(_) | TensorStorage::Pbo(_));
4386 // macOS `Dma` is the IOSurface; Linux `Dma` is the DMA-BUF; Android
4387 // `Dma` is the AHardwareBuffer — all apply the offset in their map.
4388 // (`Dma` is the same variant name on each, hence one `cfg(any(...))`
4389 // arm rather than three.)
4390 #[cfg(any(
4391 target_os = "linux",
4392 target_os = "macos",
4393 target_os = "ios",
4394 target_os = "android"
4395 ))]
4396 let supported = supported || matches!(self.storage, TensorStorage::Dma(_));
4397 #[cfg(unix)]
4398 let supported = supported || matches!(self.storage, TensorStorage::Shm(_));
4399 if !supported {
4400 return Err(Error::InvalidOperation(
4401 "plane offset only supported for DMA, Mem, Shm, and PBO tensors".into(),
4402 ));
4403 }
4404 }
4405 self.storage.map_with(access)
4406 }
4407
4408 fn buffer_identity(&self) -> &BufferIdentity {
4409 self.storage.buffer_identity()
4410 }
4411}
4412
4413pub enum TensorMap<T>
4414where
4415 T: Num + Clone + fmt::Debug,
4416{
4417 #[cfg(target_os = "linux")]
4418 Dma(DmaMap<T>),
4419 #[cfg(any(target_os = "macos", target_os = "ios"))]
4420 IoSurface(IoSurfaceMap<T>),
4421 #[cfg(target_os = "android")]
4422 HardwareBuffer(AHardwareBufferMap<T>),
4423 #[cfg(unix)]
4424 Shm(ShmMap<T>),
4425 Mem(MemMap<T>),
4426 Pbo(PboMap<T>),
4427}
4428
4429impl<T> TensorMapTrait<T> for TensorMap<T>
4430where
4431 T: Num + Clone + fmt::Debug,
4432{
4433 fn shape(&self) -> &[usize] {
4434 match self {
4435 #[cfg(target_os = "linux")]
4436 TensorMap::Dma(map) => map.shape(),
4437 #[cfg(any(target_os = "macos", target_os = "ios"))]
4438 TensorMap::IoSurface(map) => map.shape(),
4439 #[cfg(target_os = "android")]
4440 TensorMap::HardwareBuffer(map) => map.shape(),
4441 #[cfg(unix)]
4442 TensorMap::Shm(map) => map.shape(),
4443 TensorMap::Mem(map) => map.shape(),
4444 TensorMap::Pbo(map) => map.shape(),
4445 }
4446 }
4447
4448 fn unmap(&mut self) {
4449 match self {
4450 #[cfg(target_os = "linux")]
4451 TensorMap::Dma(map) => map.unmap(),
4452 #[cfg(any(target_os = "macos", target_os = "ios"))]
4453 TensorMap::IoSurface(map) => map.unmap(),
4454 #[cfg(target_os = "android")]
4455 TensorMap::HardwareBuffer(map) => map.unmap(),
4456 #[cfg(unix)]
4457 TensorMap::Shm(map) => map.unmap(),
4458 TensorMap::Mem(map) => map.unmap(),
4459 TensorMap::Pbo(map) => map.unmap(),
4460 }
4461 }
4462
4463 fn as_slice(&self) -> &[T] {
4464 match self {
4465 #[cfg(target_os = "linux")]
4466 TensorMap::Dma(map) => map.as_slice(),
4467 #[cfg(any(target_os = "macos", target_os = "ios"))]
4468 TensorMap::IoSurface(map) => map.deref(),
4469 #[cfg(target_os = "android")]
4470 TensorMap::HardwareBuffer(map) => map.deref(),
4471 #[cfg(unix)]
4472 TensorMap::Shm(map) => map.as_slice(),
4473 TensorMap::Mem(map) => map.as_slice(),
4474 TensorMap::Pbo(map) => map.as_slice(),
4475 }
4476 }
4477
4478 fn as_mut_slice(&mut self) -> &mut [T] {
4479 match self {
4480 #[cfg(target_os = "linux")]
4481 TensorMap::Dma(map) => map.as_mut_slice(),
4482 #[cfg(any(target_os = "macos", target_os = "ios"))]
4483 TensorMap::IoSurface(map) => map.deref_mut(),
4484 #[cfg(target_os = "android")]
4485 TensorMap::HardwareBuffer(map) => map.deref_mut(),
4486 #[cfg(unix)]
4487 TensorMap::Shm(map) => map.as_mut_slice(),
4488 TensorMap::Mem(map) => map.as_mut_slice(),
4489 TensorMap::Pbo(map) => map.as_mut_slice(),
4490 }
4491 }
4492}
4493
4494impl<T> Deref for TensorMap<T>
4495where
4496 T: Num + Clone + fmt::Debug,
4497{
4498 type Target = [T];
4499
4500 fn deref(&self) -> &[T] {
4501 match self {
4502 #[cfg(target_os = "linux")]
4503 TensorMap::Dma(map) => map.deref(),
4504 #[cfg(any(target_os = "macos", target_os = "ios"))]
4505 TensorMap::IoSurface(map) => map.deref(),
4506 #[cfg(target_os = "android")]
4507 TensorMap::HardwareBuffer(map) => map.deref(),
4508 #[cfg(unix)]
4509 TensorMap::Shm(map) => map.deref(),
4510 TensorMap::Mem(map) => map.deref(),
4511 TensorMap::Pbo(map) => map.deref(),
4512 }
4513 }
4514}
4515
4516impl<T> DerefMut for TensorMap<T>
4517where
4518 T: Num + Clone + fmt::Debug,
4519{
4520 fn deref_mut(&mut self) -> &mut [T] {
4521 match self {
4522 #[cfg(target_os = "linux")]
4523 TensorMap::Dma(map) => map.deref_mut(),
4524 #[cfg(any(target_os = "macos", target_os = "ios"))]
4525 TensorMap::IoSurface(map) => map.deref_mut(),
4526 #[cfg(target_os = "android")]
4527 TensorMap::HardwareBuffer(map) => map.deref_mut(),
4528 #[cfg(unix)]
4529 TensorMap::Shm(map) => map.deref_mut(),
4530 TensorMap::Mem(map) => map.deref_mut(),
4531 TensorMap::Pbo(map) => map.deref_mut(),
4532 }
4533 }
4534}
4535
4536// ============================================================================
4537// Platform availability helpers
4538// ============================================================================
4539
4540/// Cached result of the Linux DMA-BUF availability probe.
4541#[cfg(target_os = "linux")]
4542static DMA_AVAILABLE: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
4543/// Cached result of the macOS/iOS IOSurface availability probe.
4544#[cfg(any(target_os = "macos", target_os = "ios"))]
4545static IOSURFACE_AVAILABLE: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
4546
4547/// Check if Linux DMA-BUF allocation is available on this system.
4548///
4549/// Linux-specific availability check (typically requires `/dev/dma_heap`
4550/// access — running as root or membership in a video/render group). For
4551/// portable code that wants "any zero-copy GPU buffer", use
4552/// [`is_gpu_buffer_available`] which also covers IOSurface on macOS.
4553///
4554/// This function caches its result after the first call.
4555#[cfg(target_os = "linux")]
4556pub fn is_dma_available() -> bool {
4557 *DMA_AVAILABLE.get_or_init(|| Tensor::<u8>::new(&[64], Some(TensorMemory::Dma), None).is_ok())
4558}
4559
4560/// Always returns `false` on non-Linux platforms.
4561#[cfg(not(target_os = "linux"))]
4562pub fn is_dma_available() -> bool {
4563 false
4564}
4565
4566/// Check if macOS/iOS IOSurface allocation is available on this system.
4567///
4568/// IOSurface is part of the macOS/iOS OS and is essentially always present;
4569/// this probe catches degraded scenarios such as memory pressure or
4570/// sandboxed contexts where `IOSurfaceCreate` fails. The result is
4571/// cached after the first call.
4572///
4573/// Always returns `false` on non-Apple platforms.
4574#[cfg(any(target_os = "macos", target_os = "ios"))]
4575pub fn is_iosurface_available() -> bool {
4576 *IOSURFACE_AVAILABLE.get_or_init(|| {
4577 // Probe via the same Dma path — on macOS/iOS this routes through
4578 // IoSurfaceTensor::new.
4579 Tensor::<u8>::new(&[64], Some(TensorMemory::Dma), None).is_ok()
4580 })
4581}
4582
4583#[cfg(not(any(target_os = "macos", target_os = "ios")))]
4584pub fn is_iosurface_available() -> bool {
4585 false
4586}
4587
4588/// Cached result of the Android AHardwareBuffer availability probe.
4589#[cfg(target_os = "android")]
4590static AHARDWAREBUFFER_AVAILABLE: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
4591
4592/// Check if Android AHardwareBuffer allocation is available on this system.
4593///
4594/// AHardwareBuffer is part of the Android OS (public NDK ABI since API
4595/// 26) and is essentially always present; this probe catches degraded
4596/// scenarios such as memory pressure or gralloc failures. The result is
4597/// cached after the first call.
4598#[cfg(target_os = "android")]
4599pub fn is_ahardwarebuffer_available() -> bool {
4600 *AHARDWAREBUFFER_AVAILABLE.get_or_init(|| {
4601 // Probe via the same Dma path — on Android this routes through
4602 // AHardwareBufferTensor::new.
4603 Tensor::<u8>::new(&[64], Some(TensorMemory::Dma), None).is_ok()
4604 })
4605}
4606
4607/// Always returns `false` on non-Android platforms.
4608#[cfg(not(target_os = "android"))]
4609pub fn is_ahardwarebuffer_available() -> bool {
4610 false
4611}
4612
4613/// Portable probe for the platform's native zero-copy GPU buffer
4614/// allocator (DMA-BUF on Linux, IOSurface on macOS/iOS, AHardwareBuffer on
4615/// Android). Returns `false` on
4616/// Windows and other platforms with no equivalent. Use this when writing
4617/// cross-platform code that cares whether the `Dma` tensor variant will
4618/// work, not which underlying mechanism is used.
4619pub fn is_gpu_buffer_available() -> bool {
4620 #[cfg(target_os = "linux")]
4621 {
4622 is_dma_available()
4623 }
4624 #[cfg(any(target_os = "macos", target_os = "ios"))]
4625 {
4626 is_iosurface_available()
4627 }
4628 #[cfg(target_os = "android")]
4629 {
4630 is_ahardwarebuffer_available()
4631 }
4632 #[cfg(not(any(
4633 target_os = "linux",
4634 target_os = "macos",
4635 target_os = "ios",
4636 target_os = "android"
4637 )))]
4638 {
4639 false
4640 }
4641}
4642
4643/// Check if POSIX shared memory allocation is available on this system.
4644///
4645/// Returns `true` on Unix systems (Linux, macOS, BSD) where POSIX shared memory
4646/// is supported. Always returns `false` on non-Unix platforms (Windows).
4647///
4648/// This function caches its result after the first call for efficiency.
4649#[cfg(unix)]
4650static SHM_AVAILABLE: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
4651
4652/// Check if POSIX shared memory allocation is available on this system.
4653#[cfg(unix)]
4654pub fn is_shm_available() -> bool {
4655 *SHM_AVAILABLE.get_or_init(|| Tensor::<u8>::new(&[64], Some(TensorMemory::Shm), None).is_ok())
4656}
4657
4658/// Check if POSIX shared memory allocation is available on this system.
4659///
4660/// Always returns `false` on non-Unix platforms since POSIX SHM is Unix-specific.
4661#[cfg(not(unix))]
4662pub fn is_shm_available() -> bool {
4663 false
4664}
4665
4666#[cfg(test)]
4667mod dtype_tests {
4668 use super::*;
4669
4670 #[test]
4671 fn dtype_size() {
4672 assert_eq!(DType::U8.size(), 1);
4673 assert_eq!(DType::I8.size(), 1);
4674 assert_eq!(DType::U16.size(), 2);
4675 assert_eq!(DType::I16.size(), 2);
4676 assert_eq!(DType::U32.size(), 4);
4677 assert_eq!(DType::I32.size(), 4);
4678 assert_eq!(DType::U64.size(), 8);
4679 assert_eq!(DType::I64.size(), 8);
4680 assert_eq!(DType::F16.size(), 2);
4681 assert_eq!(DType::F32.size(), 4);
4682 assert_eq!(DType::F64.size(), 8);
4683 }
4684
4685 #[test]
4686 fn dtype_name() {
4687 assert_eq!(DType::U8.name(), "u8");
4688 assert_eq!(DType::F16.name(), "f16");
4689 assert_eq!(DType::F32.name(), "f32");
4690 }
4691
4692 #[test]
4693 fn dtype_serde_roundtrip() {
4694 use serde_json;
4695 let dt = DType::F16;
4696 let json = serde_json::to_string(&dt).unwrap();
4697 let back: DType = serde_json::from_str(&json).unwrap();
4698 assert_eq!(dt, back);
4699 }
4700}
4701
4702#[cfg(test)]
4703mod image_tests {
4704 use super::*;
4705
4706 #[test]
4707 fn image_shape_per_layout() {
4708 assert_eq!(
4709 PixelFormat::Rgb.image_shape(640, 480),
4710 Some(vec![480, 640, 3])
4711 );
4712 assert_eq!(
4713 PixelFormat::Grey.image_shape(640, 480),
4714 Some(vec![480, 640, 1])
4715 );
4716 assert_eq!(
4717 PixelFormat::Nv12.image_shape(640, 480),
4718 Some(vec![720, 640])
4719 );
4720 // Odd height: combined-plane height is `481 + ceil(481/2)` = 481 + 241
4721 // = 722 rows. Logical height is recovered as `722 * 2 / 3` = 481.
4722 assert_eq!(
4723 PixelFormat::Nv12.image_shape(640, 481),
4724 Some(vec![722, 640])
4725 );
4726 // Odd width: shape carries the LOGICAL width (641).
4727 // The 64-aligned stride (>= 642) is stored separately on the Tensor.
4728 assert_eq!(
4729 PixelFormat::Nv12.image_shape(641, 480),
4730 Some(vec![720, 641])
4731 );
4732 // NV16 odd width: same — logical width in shape, stride separate.
4733 assert_eq!(
4734 PixelFormat::Nv16.image_shape(641, 480),
4735 Some(vec![960, 641])
4736 );
4737 assert_eq!(
4738 PixelFormat::PlanarRgb.image_shape(640, 480),
4739 Some(vec![3, 480, 640])
4740 );
4741 assert_eq!(
4742 PixelFormat::Nv16.image_shape(640, 480),
4743 Some(vec![960, 640])
4744 );
4745 }
4746
4747 #[test]
4748 fn raw_tensor_has_no_format() {
4749 let t = Tensor::<u8>::new(&[480, 640, 3], None, None).unwrap();
4750 assert!(t.format().is_none());
4751 assert!(t.width().is_none());
4752 assert!(t.height().is_none());
4753 assert!(!t.is_multiplane());
4754 assert!(t.chroma().is_none());
4755 }
4756
4757 #[test]
4758 fn image_tensor_packed() {
4759 let t = Tensor::<u8>::image(
4760 640,
4761 480,
4762 PixelFormat::Rgba,
4763 None,
4764 crate::CpuAccess::ReadWrite,
4765 )
4766 .unwrap();
4767 assert_eq!(t.format(), Some(PixelFormat::Rgba));
4768 assert_eq!(t.width(), Some(640));
4769 assert_eq!(t.height(), Some(480));
4770 assert_eq!(t.shape(), &[480, 640, 4]);
4771 assert!(!t.is_multiplane());
4772 }
4773
4774 #[test]
4775 fn image_tensor_planar() {
4776 let t = Tensor::<u8>::image(
4777 640,
4778 480,
4779 PixelFormat::PlanarRgb,
4780 None,
4781 crate::CpuAccess::ReadWrite,
4782 )
4783 .unwrap();
4784 assert_eq!(t.format(), Some(PixelFormat::PlanarRgb));
4785 assert_eq!(t.width(), Some(640));
4786 assert_eq!(t.height(), Some(480));
4787 assert_eq!(t.shape(), &[3, 480, 640]);
4788 }
4789
4790 #[test]
4791 #[cfg(target_os = "macos")]
4792 fn image_tensor_dma_non_aligned_packed_width_pads_zero_copy() {
4793 // RGBA u8 at width=4 → 4*4 = 16 bytes/row, not 64-byte aligned. RGBA has
4794 // a real IOSurface FourCC, so an explicit `Some(TensorMemory::Dma)`
4795 // request now allocates a padded image IOSurface (64-aligned
4796 // `bytes_per_row`) and records the stride — a fully zero-copy buffer GL
4797 // can bind and the CPU can map via the strided path. (Previously this
4798 // failed loudly to avoid an 'L008' byte-bag downgrade; with a real
4799 // FourCC surface that concern no longer applies.)
4800 let t = Tensor::<u8>::image(
4801 4,
4802 4,
4803 PixelFormat::Rgba,
4804 Some(TensorMemory::Dma),
4805 crate::CpuAccess::ReadWrite,
4806 )
4807 .expect("padded RGBA IOSurface should allocate");
4808 assert_eq!(t.format(), Some(PixelFormat::Rgba));
4809 assert_eq!(t.width(), Some(4));
4810 assert_eq!(t.height(), Some(4));
4811 let stride = t.effective_row_stride().expect("stride");
4812 assert_eq!(stride % 64, 0, "padded to 64-byte row alignment");
4813 assert!(stride >= 16);
4814 // A CPU map exposes the full padded surface for strided iteration.
4815 let m = t.map().expect("strided IOSurface map");
4816 assert_eq!(m.as_slice().len(), stride * 4);
4817 }
4818
4819 /// `per_pixel_bytes` that doesn't divide 64 evenly (e.g. RGB u8 with
4820 /// 3 B/pixel) makes a "Pad width to N" suggestion structurally
4821 /// impossible — there is no integer width whose `width * 3` is a
4822 /// multiple of 64. The error must still fire (no silent SHM
4823 /// fallback for explicit-DMA requests) and must spell out the
4824 /// alignment requirement; it just omits the misleading "pad to N"
4825 /// hint instead of printing a number whose row pitch still won't
4826 /// align.
4827 #[test]
4828 #[cfg(target_os = "macos")]
4829 fn image_tensor_dma_rejects_indivisible_pixel_pitch_without_pad_hint() {
4830 // Width=10 RGB f32 → 120 B/row, not 64-byte aligned, and (Rgb,
4831 // F32) has no IOSurface mapping so the padded-stride tolerance
4832 // does not apply. The next 64-multiple (128 B) isn't an integer
4833 // multiple of 12 B/pixel, so the "pad width to N" hint can't
4834 // produce a valid number and must be omitted. (Rgb u8 used to be
4835 // this test's subject but now has a real RGBA8888 mapping with
4836 // padded-stride tolerance — see the test below.)
4837 let err = Tensor::<f32>::image(
4838 10,
4839 10,
4840 PixelFormat::Rgb,
4841 Some(TensorMemory::Dma),
4842 crate::CpuAccess::ReadWrite,
4843 )
4844 .expect_err("RGB f32 with 12 B/pixel and non-aligned width must be rejected");
4845 match err {
4846 Error::InvalidArgument(msg) => {
4847 assert!(
4848 msg.contains("64-byte aligned"),
4849 "error must still name the alignment requirement: {msg}"
4850 );
4851 assert!(
4852 !msg.contains("Pad width"),
4853 "indivisible per-pixel pitch makes a width suggestion impossible; \
4854 hint must be omitted, got: {msg}"
4855 );
4856 assert!(
4857 msg.contains("memory=None") && msg.contains("TensorMemory::Mem"),
4858 "error must still list the always-applicable alternatives: {msg}"
4859 );
4860 }
4861 other => panic!("expected InvalidArgument, got {other:?}"),
4862 }
4863 }
4864
4865 #[test]
4866 #[cfg(target_os = "macos")]
4867 fn image_tensor_dma_packed_rgb_u8_contract() {
4868 // Packed RGB u8 @Dma is a designed RGBA8888 mapping at
4869 // (W*3/4, H) — the INT8 NPU input layout, shared with Android.
4870 // width%4 != 0 cannot form whole texels → loud InvalidArgument…
4871 let err = Tensor::<u8>::image(
4872 10,
4873 10,
4874 PixelFormat::Rgb,
4875 Some(TensorMemory::Dma),
4876 crate::CpuAccess::ReadWrite,
4877 )
4878 .expect_err("Rgb u8 width%4!=0 must be rejected");
4879 assert!(
4880 matches!(&err, Error::InvalidArgument(m) if m.contains("width%4==0")),
4881 "got {err:?}"
4882 );
4883 // …width%4 == 0 with a non-64-aligned pitch allocates PADDED
4884 // (36 B rows → 64 B surface pitch, recorded on the tensor)…
4885 let t = Tensor::<u8>::image(
4886 12,
4887 4,
4888 PixelFormat::Rgb,
4889 Some(TensorMemory::Dma),
4890 crate::CpuAccess::ReadWrite,
4891 )
4892 .expect("width 12 Rgb u8 must allocate padded");
4893 assert_eq!(t.memory(), TensorMemory::Dma);
4894 assert!(
4895 t.row_stride()
4896 .is_some_and(|s| s >= 64 && s.is_multiple_of(64)),
4897 "padded pitch must be recorded: {:?}",
4898 t.row_stride()
4899 );
4900 // …and the aligned model-input width stays flat (640*3 = 1920 is
4901 // 64-aligned → no recorded stride, the buffer IS [H, W, 3]).
4902 let t = Tensor::<u8>::image(
4903 640,
4904 8,
4905 PixelFormat::Rgb,
4906 Some(TensorMemory::Dma),
4907 crate::CpuAccess::ReadWrite,
4908 )
4909 .expect("width 640 Rgb u8 must allocate flat");
4910 assert_eq!(t.row_stride(), None);
4911 // I8 shares the layout (INT8 shader bias, not a format change).
4912 let t = Tensor::<i8>::image(
4913 640,
4914 8,
4915 PixelFormat::Rgb,
4916 Some(TensorMemory::Dma),
4917 crate::CpuAccess::ReadWrite,
4918 )
4919 .expect("Rgb i8 shares the RGBA8888 mapping");
4920 assert_eq!(t.memory(), TensorMemory::Dma);
4921 }
4922
4923 #[test]
4924 #[cfg(target_os = "macos")]
4925 fn image_tensor_dma_planar_f16_alignment() {
4926 // PlanarRgb F16 uses single-channel row pitch (width * 2 bytes).
4927 // Width=16 → 32 bytes/row (not aligned); width=32 → 64 bytes/row (aligned).
4928 let err = Tensor::<half::f16>::image(
4929 16,
4930 16,
4931 PixelFormat::PlanarRgb,
4932 Some(TensorMemory::Dma),
4933 crate::CpuAccess::ReadWrite,
4934 )
4935 .expect_err("width=16 PlanarRgb F16 is 32-byte row, must reject");
4936 assert!(matches!(err, Error::InvalidArgument(_)), "got {err:?}");
4937 // 32 wide should work.
4938 let t = Tensor::<half::f16>::image(
4939 32,
4940 8,
4941 PixelFormat::PlanarRgb,
4942 Some(TensorMemory::Dma),
4943 crate::CpuAccess::ReadWrite,
4944 )
4945 .expect("width=32 PlanarRgb F16 is 64-byte row, must succeed");
4946 assert_eq!(t.format(), Some(PixelFormat::PlanarRgb));
4947 }
4948
4949 #[test]
4950 fn image_tensor_semi_planar_contiguous() {
4951 let t = Tensor::<u8>::image(
4952 640,
4953 480,
4954 PixelFormat::Nv12,
4955 None,
4956 crate::CpuAccess::ReadWrite,
4957 )
4958 .unwrap();
4959 assert_eq!(t.format(), Some(PixelFormat::Nv12));
4960 assert_eq!(t.width(), Some(640));
4961 assert_eq!(t.height(), Some(480));
4962 // NV12: H*3/2 = 720
4963 assert_eq!(t.shape(), &[720, 640]);
4964 assert!(!t.is_multiplane());
4965 }
4966
4967 #[test]
4968 #[cfg(target_os = "linux")]
4969 fn image_tensor_with_stride_preserves_logical_width() {
4970 // Skip if DMA not available (e.g. sandboxed CI lacking dma_heap access).
4971 if !is_dma_available() {
4972 eprintln!("SKIPPED: DMA heap not available");
4973 return;
4974 }
4975 // 3004×1688 RGBA8: natural pitch 12016, padded to 12032 (64-aligned).
4976 let stride = 12032;
4977 let t = Tensor::<u8>::image_with_stride(
4978 3004,
4979 1688,
4980 PixelFormat::Rgba,
4981 stride,
4982 Some(TensorMemory::Dma),
4983 crate::CpuAccess::ReadWrite,
4984 )
4985 .unwrap();
4986 // Logical dimensions unchanged by padding — this is the contract.
4987 assert_eq!(t.width(), Some(3004));
4988 assert_eq!(t.height(), Some(1688));
4989 assert_eq!(t.shape(), &[1688, 3004, 4]);
4990 // Stride is carried separately and reports the padded pitch.
4991 assert_eq!(t.effective_row_stride(), Some(stride));
4992 // Buffer is sized to stride × height so the full padded layout fits,
4993 // and CPU map() works for self-allocated strided DMA tensors.
4994 use crate::TensorMapTrait;
4995 {
4996 let map = t.map().unwrap();
4997 assert!(
4998 map.as_slice().len() >= stride * 1688,
4999 "mapped buffer {} bytes < expected {}",
5000 map.as_slice().len(),
5001 stride * 1688
5002 );
5003 }
5004 // CPU write access works too — iterate rows using the padded stride,
5005 // touch only the active `width × bpp` region, verify it round-trips.
5006 {
5007 let mut map = t.map().unwrap();
5008 let slice = map.as_mut_slice();
5009 for y in 0..1688 {
5010 let row_start = y * stride;
5011 for x in 0..3004 {
5012 let p = row_start + x * 4;
5013 slice[p] = (y & 0xFF) as u8;
5014 slice[p + 1] = (x & 0xFF) as u8;
5015 slice[p + 2] = 0x42;
5016 slice[p + 3] = 0xFF;
5017 }
5018 }
5019 }
5020 {
5021 let map = t.map().unwrap();
5022 let slice = map.as_slice();
5023 // Sample a few pixels to confirm the round-trip.
5024 assert_eq!(slice[0], 0x00);
5025 assert_eq!(slice[1], 0x00);
5026 assert_eq!(slice[2], 0x42);
5027 assert_eq!(slice[3], 0xFF);
5028 let mid = 100 * stride + 50 * 4;
5029 assert_eq!(slice[mid], 100);
5030 assert_eq!(slice[mid + 1], 50);
5031 assert_eq!(slice[mid + 2], 0x42);
5032 }
5033 }
5034
5035 #[test]
5036 #[cfg(target_os = "linux")]
5037 fn image_tensor_with_stride_rejects_foreign_strided_map() {
5038 // A FOREIGN (imported via from_fd) DMA tensor with row_stride set
5039 // should still refuse CPU mapping — external allocator owns the
5040 // layout. This protects the V4L2 / GStreamer use case.
5041 //
5042 // We simulate a foreign import by wrapping our own allocation's
5043 // fd via `from_fd` and calling set_row_stride manually. The
5044 // `is_imported` flag on from_fd is true by construction.
5045 if !is_dma_available() {
5046 eprintln!("SKIPPED: DMA heap not available");
5047 return;
5048 }
5049 // Allocate a backing buffer large enough for a 320×240 BGRA8 image.
5050 let backing = Tensor::<u8>::new(&[240 * 320 * 4], Some(TensorMemory::Dma), None).unwrap();
5051 let fd = backing.clone_fd().unwrap();
5052 // Import it via from_fd — this marks is_imported=true.
5053 let shape = [240usize, 320, 4];
5054 let storage = TensorStorage::<u8>::from_fd(fd, &shape, None).unwrap();
5055 let mut t = Tensor::<u8>::wrap(storage);
5056 t.set_format(PixelFormat::Bgra).unwrap();
5057 t.set_row_stride(320 * 4).unwrap(); // natural, but still marks it as strided
5058 let err = t.map();
5059 assert!(
5060 matches!(err, Err(Error::InvalidOperation(_))),
5061 "foreign strided map should error"
5062 );
5063 }
5064
5065 #[test]
5066 #[cfg(target_os = "linux")]
5067 fn image_tensor_with_stride_map_rejects_tampered_stride() {
5068 // Round-3 PR feedback (C1): `set_row_stride` is public and only
5069 // validates `stride >= min_stride`, not that the new stride × height
5070 // fits the underlying buffer. A caller that tampers with the stride
5071 // after allocation must not be able to coerce `Tensor::map()` into
5072 // returning a slice larger than the backing mmap (that would be UB
5073 // in `DmaMap::as_slice`).
5074 if !is_dma_available() {
5075 eprintln!("SKIPPED: DMA heap not available");
5076 return;
5077 }
5078 // Allocate a 640×480 RGBA8 padded canvas (stride = 3072 = 768 px).
5079 // Backing buffer is 3072 × 480 = 1,474,560 bytes.
5080 let mut t = Tensor::<u8>::image_with_stride(
5081 640,
5082 480,
5083 PixelFormat::Rgba,
5084 3072,
5085 Some(TensorMemory::Dma),
5086 crate::CpuAccess::ReadWrite,
5087 )
5088 .unwrap();
5089 // Tamper: push the stride up to 4 × the original. This is >=
5090 // min_stride (2560), so `set_row_stride` accepts it.
5091 t.set_row_stride(12288).unwrap();
5092 // Map must now refuse — 12288 × 480 = 5,898,240 > 1,474,560.
5093 let err = t.map();
5094 assert!(
5095 matches!(err, Err(Error::InvalidOperation(_))),
5096 "map() with oversized stride must return InvalidOperation"
5097 );
5098 }
5099
5100 #[test]
5101 fn dma_tensor_new_with_byte_size_rejects_shape_overflow() {
5102 // Round-3 PR feedback (C3): shape.product() * sizeof(T) must use
5103 // checked arithmetic so a pathological shape can't wrap usize and
5104 // make the byte_size-vs-logical-size comparison incorrect.
5105 //
5106 // This test only exercises the overflow rejection path, which is
5107 // pure-Rust and doesn't touch dma_heap — safe to run on any target.
5108 #[cfg(target_os = "linux")]
5109 {
5110 let err = crate::dma::DmaTensor::<u64>::new_with_byte_size(
5111 &[usize::MAX, 2, 2],
5112 usize::MAX,
5113 None,
5114 );
5115 assert!(
5116 matches!(err, Err(Error::InvalidArgument(_))),
5117 "new_with_byte_size must detect shape.product() overflow"
5118 );
5119 }
5120 }
5121
5122 #[test]
5123 #[cfg(target_os = "linux")]
5124 fn image_tensor_with_stride_rejects_too_small_stride() {
5125 // 640×480 RGBA8 natural pitch = 2560, request 2400 → should error.
5126 let err = Tensor::<u8>::image_with_stride(
5127 640,
5128 480,
5129 PixelFormat::Rgba,
5130 2400,
5131 Some(TensorMemory::Dma),
5132 crate::CpuAccess::ReadWrite,
5133 );
5134 assert!(matches!(err, Err(Error::InvalidArgument(_))));
5135 }
5136
5137 #[test]
5138 #[cfg(target_os = "linux")]
5139 fn image_tensor_with_stride_rejects_non_packed() {
5140 // NV12 is SemiPlanar → not supported. (Linux-only because
5141 // `TensorMemory::Dma` itself is a Linux-only enum variant.)
5142 let err = Tensor::<u8>::image_with_stride(
5143 640,
5144 480,
5145 PixelFormat::Nv12,
5146 640,
5147 Some(TensorMemory::Dma),
5148 crate::CpuAccess::ReadWrite,
5149 );
5150 assert!(matches!(err, Err(Error::NotImplemented(_))));
5151 }
5152
5153 #[test]
5154 fn set_format_valid() {
5155 let mut t = Tensor::<u8>::new(&[480, 640, 3], None, None).unwrap();
5156 assert!(t.format().is_none());
5157 t.set_format(PixelFormat::Rgb).unwrap();
5158 assert_eq!(t.format(), Some(PixelFormat::Rgb));
5159 assert_eq!(t.width(), Some(640));
5160 assert_eq!(t.height(), Some(480));
5161 }
5162
5163 #[test]
5164 fn set_format_invalid_shape() {
5165 let mut t = Tensor::<u8>::new(&[480, 640, 4], None, None).unwrap();
5166 // RGB expects 3 channels, not 4
5167 let err = t.set_format(PixelFormat::Rgb);
5168 assert!(err.is_err());
5169 // Original tensor is unmodified
5170 assert!(t.format().is_none());
5171 }
5172
5173 #[test]
5174 fn reshape_clears_format() {
5175 let mut t = Tensor::<u8>::image(
5176 640,
5177 480,
5178 PixelFormat::Rgba,
5179 None,
5180 crate::CpuAccess::ReadWrite,
5181 )
5182 .unwrap();
5183 assert_eq!(t.format(), Some(PixelFormat::Rgba));
5184 // Reshape to flat — format cleared
5185 t.reshape(&[480 * 640 * 4]).unwrap();
5186 assert!(t.format().is_none());
5187 }
5188
5189 #[test]
5190 fn from_planes_nv12() {
5191 let y = Tensor::<u8>::new(&[480, 640], None, None).unwrap();
5192 let uv = Tensor::<u8>::new(&[240, 640], None, None).unwrap();
5193 let img = Tensor::from_planes(y, uv, PixelFormat::Nv12).unwrap();
5194 assert_eq!(img.format(), Some(PixelFormat::Nv12));
5195 assert!(img.is_multiplane());
5196 assert!(img.chroma().is_some());
5197 assert_eq!(img.width(), Some(640));
5198 assert_eq!(img.height(), Some(480));
5199 }
5200
5201 #[test]
5202 fn from_planes_rejects_non_semiplanar() {
5203 let y = Tensor::<u8>::new(&[480, 640], None, None).unwrap();
5204 let uv = Tensor::<u8>::new(&[240, 640], None, None).unwrap();
5205 let err = Tensor::from_planes(y, uv, PixelFormat::Rgb);
5206 assert!(err.is_err());
5207 }
5208
5209 #[test]
5210 fn reshape_multiplane_errors() {
5211 let y = Tensor::<u8>::new(&[480, 640], None, None).unwrap();
5212 let uv = Tensor::<u8>::new(&[240, 640], None, None).unwrap();
5213 let mut img = Tensor::from_planes(y, uv, PixelFormat::Nv12).unwrap();
5214 let err = img.reshape(&[480 * 640 + 240 * 640]);
5215 assert!(err.is_err());
5216 }
5217}
5218
5219#[cfg(test)]
5220mod compression_tests {
5221 use super::*;
5222
5223 #[test]
5224 fn desc_builder_roundtrips() {
5225 let desc = ImageDesc::new(640, 480, PixelFormat::Rgba, DType::U8)
5226 .with_memory(Some(TensorMemory::Mem))
5227 .with_access(CpuAccess::Read)
5228 .with_compression(Compression::Any);
5229 assert_eq!(desc.width(), 640);
5230 assert_eq!(desc.height(), 480);
5231 assert_eq!(desc.format(), PixelFormat::Rgba);
5232 assert_eq!(desc.dtype(), DType::U8);
5233 assert_eq!(desc.memory(), Some(TensorMemory::Mem));
5234 assert_eq!(desc.access(), CpuAccess::Read);
5235 assert_eq!(desc.compression(), Some(Compression::Any));
5236
5237 // Defaults: auto memory, hardware-only, no request.
5238 let plain = ImageDesc::new(2, 2, PixelFormat::Grey, DType::U8);
5239 assert_eq!(plain.memory(), None);
5240 assert_eq!(plain.access(), CpuAccess::None);
5241 assert_eq!(plain.compression(), None);
5242 }
5243
5244 #[test]
5245 fn desc_dtype_must_match_element_type() {
5246 let desc = ImageDesc::new(4, 4, PixelFormat::Rgba, DType::F32);
5247 match Tensor::<u8>::image_desc(&desc) {
5248 Err(Error::InvalidArgument(msg)) => assert!(msg.contains("dtype")),
5249 other => panic!("expected InvalidArgument, got {other:?}"),
5250 }
5251 }
5252
5253 #[test]
5254 fn compression_with_cpu_access_is_invalid() {
5255 let desc = ImageDesc::new(4, 4, PixelFormat::Rgba, DType::U8)
5256 .with_access(CpuAccess::ReadWrite)
5257 .with_compression(Compression::Any);
5258 match Tensor::<u8>::image_desc(&desc) {
5259 Err(Error::InvalidArgument(msg)) => assert!(msg.contains("CpuAccess::None")),
5260 other => panic!("expected InvalidArgument, got {other:?}"),
5261 }
5262 }
5263
5264 #[cfg(not(target_os = "android"))]
5265 #[test]
5266 fn scheme_request_off_android_is_not_implemented() {
5267 let desc = ImageDesc::new(4, 4, PixelFormat::Rgba, DType::U8)
5268 .with_compression(Compression::Scheme(CompressionScheme::Ubwc));
5269 match Tensor::<u8>::image_desc(&desc) {
5270 Err(Error::NotImplemented(msg)) => assert!(msg.contains("Ubwc")),
5271 other => panic!("expected NotImplemented, got {other:?}"),
5272 }
5273 }
5274
5275 #[cfg(not(target_os = "android"))]
5276 #[test]
5277 fn any_request_off_android_resolves_linear_and_counts() {
5278 let before = compression_fallback_count();
5279 let desc = ImageDesc::new(64, 64, PixelFormat::Rgba, DType::U8)
5280 .with_memory(Some(TensorMemory::Mem))
5281 .with_compression(Compression::Any);
5282 let t = Tensor::<u8>::image_desc(&desc).unwrap();
5283 assert_eq!(t.compression(), None);
5284 assert!(compression_fallback_count() > before);
5285 }
5286
5287 #[test]
5288 fn desc_without_request_matches_classic_constructor() {
5289 let desc = ImageDesc::new(32, 32, PixelFormat::Rgba, DType::U8)
5290 .with_memory(Some(TensorMemory::Mem))
5291 .with_access(CpuAccess::ReadWrite);
5292 let t = Tensor::<u8>::image_desc(&desc).unwrap();
5293 assert_eq!(t.compression(), None);
5294 assert_eq!(t.cpu_access(), CpuAccess::ReadWrite);
5295 assert_eq!(t.width(), Some(32));
5296 // Mappable exactly like the classic constructor's result.
5297 let m = t.map_read().unwrap();
5298 assert_eq!(m.as_slice().len(), 32 * 32 * 4);
5299 }
5300
5301 #[test]
5302 fn configure_image_preserves_compression_and_views_inherit() {
5303 // No host platform records a scheme, so emulate the recording to
5304 // pin the preserve/inherit semantics (the physical layout does
5305 // not change when the logical image is reconfigured).
5306 let mut t = Tensor::<u8>::image(
5307 64,
5308 64,
5309 PixelFormat::Rgba,
5310 Some(TensorMemory::Mem),
5311 CpuAccess::None,
5312 )
5313 .unwrap();
5314 t.compression = Some(CompressionScheme::Ubwc);
5315 t.configure_image(32, 32, PixelFormat::Rgba).unwrap();
5316 assert_eq!(t.compression(), Some(CompressionScheme::Ubwc));
5317 let view = t.subview(0, &[16, 32, 4]).unwrap();
5318 assert_eq!(view.compression(), Some(CompressionScheme::Ubwc));
5319 }
5320
5321 #[test]
5322 fn tensor_dyn_dispatches_desc_and_compression() {
5323 let desc = ImageDesc::new(16, 16, PixelFormat::Rgba, DType::U8)
5324 .with_memory(Some(TensorMemory::Mem))
5325 .with_access(CpuAccess::ReadWrite);
5326 let t = TensorDyn::image_desc(&desc).unwrap();
5327 assert_eq!(t.compression(), None);
5328 assert!(matches!(t, TensorDyn::U8(_)));
5329 }
5330}
5331
5332#[cfg(test)]
5333mod cpu_access_tests {
5334 use super::*;
5335
5336 #[test]
5337 fn covers_matrix() {
5338 use CpuAccess::*;
5339 // Every declaration covers a narrower or equal request…
5340 for a in [None, Read, Write, ReadWrite] {
5341 assert!(a.covers(None), "{a:?} must cover None");
5342 assert!(ReadWrite.covers(a), "ReadWrite must cover {a:?}");
5343 }
5344 assert!(Read.covers(Read));
5345 assert!(Write.covers(Write));
5346 // …and never a wider one.
5347 assert!(!None.covers(Read));
5348 assert!(!None.covers(Write));
5349 assert!(!Read.covers(Write));
5350 assert!(!Read.covers(ReadWrite));
5351 assert!(!Write.covers(Read));
5352 assert!(!Write.covers(ReadWrite));
5353 }
5354
5355 #[test]
5356 fn map_with_none_is_invalid() {
5357 let t = Tensor::<u8>::new(&[16], Some(TensorMemory::Mem), None).unwrap();
5358 match t.map_with(CpuAccess::None) {
5359 Err(Error::InvalidArgument(_)) => {}
5360 Err(other) => panic!("expected InvalidArgument, got {other:?}"),
5361 Ok(_) => panic!("map_with(CpuAccess::None) must not succeed"),
5362 }
5363 }
5364
5365 #[test]
5366 fn read_map_rejects_mutation_uniformly() {
5367 // Mem backend: map_read yields a working read view whose mutable
5368 // accessor panics (the uniform cross-backend contract).
5369 let t = Tensor::<u8>::new(&[8], Some(TensorMemory::Mem), None).unwrap();
5370 t.map_mut().unwrap().as_mut_slice().copy_from_slice(&[7; 8]);
5371 let ro = t.map_read().unwrap();
5372 assert_eq!(ro.as_slice(), &[7; 8]);
5373 drop(ro);
5374 let result = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
5375 let mut ro = t.map_read().unwrap();
5376 let _ = ro.as_mut_slice();
5377 }));
5378 assert!(result.is_err(), "as_mut_slice through map_read must panic");
5379 }
5380
5381 #[test]
5382 fn write_and_rw_maps_stay_mutable() {
5383 let t = Tensor::<u8>::new(&[4], Some(TensorMemory::Mem), None).unwrap();
5384 t.map_write()
5385 .unwrap()
5386 .as_mut_slice()
5387 .copy_from_slice(&[1; 4]);
5388 t.map().unwrap().as_mut_slice().copy_from_slice(&[2; 4]);
5389 assert_eq!(t.map_read().unwrap().as_slice(), &[2; 4]);
5390 }
5391
5392 /// The read-only IOSurface lock path: a `map_read` must observe data
5393 /// written through a prior read-write lock, and its unlock (which
5394 /// skips the cache flush) must not disturb subsequent reads.
5395 #[test]
5396 #[cfg(target_os = "macos")]
5397 fn iosurface_read_only_lock_roundtrip() {
5398 let Ok(t) = Tensor::<u8>::new(&[64], Some(TensorMemory::Dma), None) else {
5399 eprintln!("SKIPPED: IOSurface unavailable");
5400 return;
5401 };
5402 {
5403 let mut m = t.map_mut().unwrap();
5404 for (i, b) in m.as_mut_slice().iter_mut().enumerate() {
5405 *b = (i * 3) as u8;
5406 }
5407 }
5408 for _ in 0..2 {
5409 let ro = t.map_read().unwrap();
5410 for (i, b) in ro.as_slice().iter().enumerate() {
5411 assert_eq!(*b, (i * 3) as u8, "byte {i} through read-only lock");
5412 }
5413 }
5414 let result = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
5415 let mut ro = t.map_read().unwrap();
5416 let _ = ro.as_mut_slice();
5417 }));
5418 assert!(result.is_err(), "IOSurface read map must reject mutation");
5419 }
5420}
5421
5422#[cfg(test)]
5423mod tests {
5424 #[cfg(target_os = "linux")]
5425 use nix::unistd::{access, AccessFlags};
5426 #[cfg(target_os = "linux")]
5427 use std::io::Write as _;
5428 use std::sync::RwLock;
5429
5430 use super::*;
5431
5432 #[ctor::ctor(unsafe)]
5433 fn init() {
5434 env_logger::Builder::from_env(env_logger::Env::default().default_filter_or("info")).init();
5435 }
5436
5437 /// Macro to get the current function name for logging in tests.
5438 #[cfg(target_os = "linux")]
5439 macro_rules! function {
5440 () => {{
5441 fn f() {}
5442 fn type_name_of<T>(_: T) -> &'static str {
5443 std::any::type_name::<T>()
5444 }
5445 let name = type_name_of(f);
5446
5447 // Find and cut the rest of the path
5448 match &name[..name.len() - 3].rfind(':') {
5449 Some(pos) => &name[pos + 1..name.len() - 3],
5450 None => &name[..name.len() - 3],
5451 }
5452 }};
5453 }
5454
5455 #[test]
5456 #[cfg(target_os = "linux")]
5457 fn test_tensor() {
5458 let _lock = FD_LOCK.read().unwrap();
5459 let shape = vec![1];
5460 let tensor = DmaTensor::<f32>::new(&shape, Some("dma_tensor"));
5461 let dma_enabled = tensor.is_ok();
5462
5463 let tensor = Tensor::<f32>::new(&shape, None, None).expect("Failed to create tensor");
5464 // Auto-select priority is Dma > Mem; Shm is never auto-selected.
5465 match dma_enabled {
5466 true => assert_eq!(tensor.memory(), TensorMemory::Dma),
5467 false => assert_eq!(tensor.memory(), TensorMemory::Mem),
5468 }
5469 }
5470
5471 #[test]
5472 #[cfg(any(target_os = "macos", target_os = "ios"))]
5473 fn test_tensor() {
5474 let shape = vec![1];
5475 let tensor = Tensor::<f32>::new(&shape, None, None).expect("Failed to create tensor");
5476 // macOS/iOS auto-fallback chain: IOSurface (Dma) → Mem. Healthy systems
5477 // return Dma; Mem only appears under memory pressure or sandboxed
5478 // contexts where IOSurfaceCreate fails. Shm is never auto-selected.
5479 let m = tensor.memory();
5480 assert!(
5481 matches!(m, TensorMemory::Dma | TensorMemory::Mem),
5482 "Unexpected auto-fallback result on macOS/iOS: {m:?}"
5483 );
5484 }
5485
5486 #[test]
5487 #[cfg(all(
5488 unix,
5489 not(any(target_os = "linux", target_os = "macos", target_os = "ios"))
5490 ))]
5491 fn test_tensor() {
5492 let shape = vec![1];
5493 let tensor = Tensor::<f32>::new(&shape, None, None).expect("Failed to create tensor");
5494 // Other Unix (BSD): no DMA, so auto-selection is Mem (Shm is
5495 // explicit-only, never auto-selected).
5496 assert_eq!(tensor.memory(), TensorMemory::Mem);
5497 }
5498
5499 #[test]
5500 #[cfg(not(unix))]
5501 fn test_tensor() {
5502 let shape = vec![1];
5503 let tensor = Tensor::<f32>::new(&shape, None, None).expect("Failed to create tensor");
5504 assert_eq!(tensor.memory(), TensorMemory::Mem);
5505 }
5506
5507 #[test]
5508 #[cfg(target_os = "linux")]
5509 fn test_dma_tensor() {
5510 let _lock = FD_LOCK.read().unwrap();
5511 match access(
5512 "/dev/dma_heap/linux,cma",
5513 AccessFlags::R_OK | AccessFlags::W_OK,
5514 ) {
5515 Ok(_) => println!("/dev/dma_heap/linux,cma is available"),
5516 Err(_) => match access(
5517 "/dev/dma_heap/system",
5518 AccessFlags::R_OK | AccessFlags::W_OK,
5519 ) {
5520 Ok(_) => println!("/dev/dma_heap/system is available"),
5521 Err(e) => {
5522 writeln!(
5523 &mut std::io::stdout(),
5524 "[WARNING] DMA Heap is unavailable: {e}"
5525 )
5526 .unwrap();
5527 return;
5528 }
5529 },
5530 }
5531
5532 let shape = vec![2, 3, 4];
5533 let tensor =
5534 DmaTensor::<f32>::new(&shape, Some("test_tensor")).expect("Failed to create tensor");
5535
5536 const DUMMY_VALUE: f32 = 12.34;
5537
5538 assert_eq!(tensor.memory(), TensorMemory::Dma);
5539 assert_eq!(tensor.name(), "test_tensor");
5540 assert_eq!(tensor.shape(), &shape);
5541 assert_eq!(tensor.size(), 2 * 3 * 4 * std::mem::size_of::<f32>());
5542 assert_eq!(tensor.len(), 2 * 3 * 4);
5543
5544 {
5545 let mut tensor_map = tensor.map().expect("Failed to map DMA memory");
5546 tensor_map.fill(42.0);
5547 assert!(tensor_map.iter().all(|&x| x == 42.0));
5548 }
5549
5550 {
5551 let shared = Tensor::<f32>::from_fd(
5552 tensor
5553 .clone_fd()
5554 .expect("Failed to duplicate tensor file descriptor"),
5555 &shape,
5556 Some("test_tensor_shared"),
5557 )
5558 .expect("Failed to create tensor from fd");
5559
5560 assert_eq!(shared.memory(), TensorMemory::Dma);
5561 assert_eq!(shared.name(), "test_tensor_shared");
5562 assert_eq!(shared.shape(), &shape);
5563
5564 let mut tensor_map = shared.map().expect("Failed to map DMA memory from fd");
5565 tensor_map.fill(DUMMY_VALUE);
5566 assert!(tensor_map.iter().all(|&x| x == DUMMY_VALUE));
5567 }
5568
5569 {
5570 let tensor_map = tensor.map().expect("Failed to map DMA memory");
5571 assert!(tensor_map.iter().all(|&x| x == DUMMY_VALUE));
5572 }
5573
5574 let mut tensor = DmaTensor::<u8>::new(&shape, None).expect("Failed to create tensor");
5575 assert_eq!(tensor.shape(), &shape);
5576 let new_shape = vec![3, 4, 4];
5577 assert!(
5578 tensor.reshape(&new_shape).is_err(),
5579 "Reshape should fail due to size mismatch"
5580 );
5581 assert_eq!(tensor.shape(), &shape, "Shape should remain unchanged");
5582
5583 let new_shape = vec![2, 3, 4];
5584 tensor.reshape(&new_shape).expect("Reshape should succeed");
5585 assert_eq!(
5586 tensor.shape(),
5587 &new_shape,
5588 "Shape should be updated after successful reshape"
5589 );
5590
5591 {
5592 let mut tensor_map = tensor.map().expect("Failed to map DMA memory");
5593 tensor_map.fill(1);
5594 assert!(tensor_map.iter().all(|&x| x == 1));
5595 }
5596
5597 {
5598 let mut tensor_map = tensor.map().expect("Failed to map DMA memory");
5599 tensor_map[2] = 42;
5600 assert_eq!(tensor_map[1], 1, "Value at index 1 should be 1");
5601 assert_eq!(tensor_map[2], 42, "Value at index 2 should be 42");
5602 }
5603 }
5604
5605 #[test]
5606 #[cfg(unix)]
5607 fn test_shm_tensor() {
5608 let _lock = FD_LOCK.read().unwrap();
5609 let shape = vec![2, 3, 4];
5610 let tensor =
5611 ShmTensor::<f32>::new(&shape, Some("test_tensor")).expect("Failed to create tensor");
5612 assert_eq!(tensor.shape(), &shape);
5613 assert_eq!(tensor.size(), 2 * 3 * 4 * std::mem::size_of::<f32>());
5614 assert_eq!(tensor.name(), "test_tensor");
5615
5616 const DUMMY_VALUE: f32 = 12.34;
5617 {
5618 let mut tensor_map = tensor.map().expect("Failed to map shared memory");
5619 tensor_map.fill(42.0);
5620 assert!(tensor_map.iter().all(|&x| x == 42.0));
5621 }
5622
5623 {
5624 let shared = Tensor::<f32>::from_fd(
5625 tensor
5626 .clone_fd()
5627 .expect("Failed to duplicate tensor file descriptor"),
5628 &shape,
5629 Some("test_tensor_shared"),
5630 )
5631 .expect("Failed to create tensor from fd");
5632
5633 assert_eq!(shared.memory(), TensorMemory::Shm);
5634 assert_eq!(shared.name(), "test_tensor_shared");
5635 assert_eq!(shared.shape(), &shape);
5636
5637 let mut tensor_map = shared.map().expect("Failed to map shared memory from fd");
5638 tensor_map.fill(DUMMY_VALUE);
5639 assert!(tensor_map.iter().all(|&x| x == DUMMY_VALUE));
5640 }
5641
5642 {
5643 let tensor_map = tensor.map().expect("Failed to map shared memory");
5644 assert!(tensor_map.iter().all(|&x| x == DUMMY_VALUE));
5645 }
5646
5647 let mut tensor = ShmTensor::<u8>::new(&shape, None).expect("Failed to create tensor");
5648 assert_eq!(tensor.shape(), &shape);
5649 let new_shape = vec![3, 4, 4];
5650 assert!(
5651 tensor.reshape(&new_shape).is_err(),
5652 "Reshape should fail due to size mismatch"
5653 );
5654 assert_eq!(tensor.shape(), &shape, "Shape should remain unchanged");
5655
5656 let new_shape = vec![2, 3, 4];
5657 tensor.reshape(&new_shape).expect("Reshape should succeed");
5658 assert_eq!(
5659 tensor.shape(),
5660 &new_shape,
5661 "Shape should be updated after successful reshape"
5662 );
5663
5664 {
5665 let mut tensor_map = tensor.map().expect("Failed to map shared memory");
5666 tensor_map.fill(1);
5667 assert!(tensor_map.iter().all(|&x| x == 1));
5668 }
5669
5670 {
5671 let mut tensor_map = tensor.map().expect("Failed to map shared memory");
5672 tensor_map[2] = 42;
5673 assert_eq!(tensor_map[1], 1, "Value at index 1 should be 1");
5674 assert_eq!(tensor_map[2], 42, "Value at index 2 should be 42");
5675 }
5676 }
5677
5678 #[test]
5679 fn mem_subview_partitions_parent_buffer() {
5680 // One heap [2,4] u8 parent (8 bytes). Two [1,4] sub-views at byte
5681 // offsets 0 and 4 must share the parent allocation (zero-copy) and be
5682 // independently writable: view 0 owns bytes [0,4), view 1 owns [4,8).
5683 // Today this is impossible — heap offset is rejected and there is no
5684 // shared sub-view constructor.
5685 let parent = Tensor::<u8>::new(&[2, 4], Some(TensorMemory::Mem), None).unwrap();
5686 let view0 = parent.subview(0, &[1, 4]).expect("subview at offset 0");
5687 let view1 = parent.subview(4, &[1, 4]).expect("subview at offset 4");
5688
5689 view1
5690 .map()
5691 .unwrap()
5692 .as_mut_slice()
5693 .copy_from_slice(&[10, 20, 30, 40]);
5694 view0
5695 .map()
5696 .unwrap()
5697 .as_mut_slice()
5698 .copy_from_slice(&[1, 2, 3, 4]);
5699
5700 // Each view sees only its own window.
5701 assert_eq!(view0.map().unwrap().as_slice(), &[1, 2, 3, 4]);
5702 assert_eq!(view1.map().unwrap().as_slice(), &[10, 20, 30, 40]);
5703 // The parent buffer is correctly partitioned (shared, zero-copy).
5704 assert_eq!(
5705 parent.map().unwrap().as_slice(),
5706 &[1, 2, 3, 4, 10, 20, 30, 40]
5707 );
5708 }
5709
5710 #[test]
5711 fn batch_partitions_leading_dim() {
5712 // Raw [4,2,2,3] u8 batched tensor: 4 elements of 12 bytes each. batch(n)
5713 // yields element n at offset n*12, sharing the parent buffer (zero-copy).
5714 let parent = Tensor::<u8>::new(&[4, 2, 2, 3], Some(TensorMemory::Mem), None).unwrap();
5715 for i in 0..4u8 {
5716 let e = parent.batch(i as usize).expect("batch element");
5717 assert_eq!(e.shape(), &[2, 2, 3]);
5718 // A batch element shares the parent's BufferIdentity.
5719 assert_eq!(e.buffer_identity().id(), parent.buffer_identity().id());
5720 for b in e.map().unwrap().as_mut_slice() {
5721 *b = i + 1;
5722 }
5723 }
5724 // Each element occupies its own 12-byte band of the parent.
5725 let whole = parent.map().unwrap();
5726 let s = whole.as_slice();
5727 for i in 0..4usize {
5728 assert!(
5729 s[i * 12..(i + 1) * 12].iter().all(|&b| b == (i as u8 + 1)),
5730 "band {i} not partitioned: {:?}",
5731 &s[i * 12..(i + 1) * 12]
5732 );
5733 }
5734 }
5735
5736 #[test]
5737 fn view_origin_snapshots_parent_and_composes() {
5738 // view() on a whole image snapshots the parent dims + the view's origin.
5739 let parent = Tensor::<u8>::image(
5740 100,
5741 80,
5742 PixelFormat::Rgba,
5743 Some(TensorMemory::Mem),
5744 crate::CpuAccess::ReadWrite,
5745 )
5746 .unwrap();
5747 assert_eq!(
5748 parent.view_origin(),
5749 None,
5750 "whole tensor has no view_origin"
5751 );
5752 let v = parent.view(Region::new(10, 20, 30, 40)).unwrap();
5753 assert_eq!(
5754 v.view_origin(),
5755 Some(ViewOrigin {
5756 parent_width: 100,
5757 parent_height: 80,
5758 parent_row_stride: 100 * 4, // tight RGBA pitch
5759 x: 10,
5760 y: 20
5761 })
5762 );
5763 // A view of a view keeps the ROOT parent and accumulates the origin.
5764 let v2 = v.view(Region::new(5, 5, 10, 10)).unwrap();
5765 assert_eq!(
5766 v2.view_origin(),
5767 Some(ViewOrigin {
5768 parent_width: 100,
5769 parent_height: 80,
5770 parent_row_stride: 100 * 4,
5771 x: 15,
5772 y: 25
5773 }),
5774 "nested view composes onto the root parent"
5775 );
5776 }
5777
5778 #[test]
5779 fn view_origin_none_for_raw_batch() {
5780 // A raw (unformatted) batched tensor has no pixel geometry, so batch()
5781 // leaves view_origin None (the per-slot path, not the one-import pivot).
5782 let parent = Tensor::<u8>::new(&[4, 2, 2, 3], Some(TensorMemory::Mem), None).unwrap();
5783 assert_eq!(parent.batch(2).unwrap().view_origin(), None);
5784 }
5785
5786 #[test]
5787 fn batch_rejects_out_of_bounds_index() {
5788 let parent = Tensor::<u8>::new(&[4, 2, 2, 3], Some(TensorMemory::Mem), None).unwrap();
5789 match parent.batch(4) {
5790 Err(Error::BatchIndexOutOfBounds { index, batch }) => {
5791 assert_eq!((index, batch), (4, 4));
5792 }
5793 other => panic!("expected BatchIndexOutOfBounds, got {other:?}"),
5794 }
5795 }
5796
5797 #[test]
5798 fn batch_zero_on_unit_n_is_whole() {
5799 // N == 1: batch(0) is the whole per-element block at offset 0 (no plane_offset).
5800 let parent = Tensor::<u8>::new(&[1, 2, 2, 3], Some(TensorMemory::Mem), None).unwrap();
5801 let e = parent.batch(0).unwrap();
5802 assert_eq!(e.shape(), &[2, 2, 3]);
5803 assert_eq!(e.plane_offset(), None);
5804 assert_eq!(e.buffer_identity().id(), parent.buffer_identity().id());
5805 }
5806
5807 #[test]
5808 fn mem_subview_rejects_unaligned_offset() {
5809 // f32 has align 4; a byte offset of 2 cannot back a valid `*const f32`.
5810 let parent = Tensor::<f32>::new(&[8], Some(TensorMemory::Mem), None).unwrap();
5811 assert!(parent.subview(2, &[1]).is_err());
5812 // A correctly aligned offset is accepted.
5813 assert!(parent.subview(4, &[1]).is_ok());
5814 }
5815
5816 #[test]
5817 fn mem_subview_rejects_out_of_bounds() {
5818 let parent = Tensor::<u8>::new(&[8], Some(TensorMemory::Mem), None).unwrap();
5819 // offset 6 + 4 bytes = 10 exceeds the 8-byte allocation.
5820 assert!(parent.subview(6, &[4]).is_err());
5821 }
5822
5823 /// Regression guard for the `TensorTrait::view` promotion (R2): a `subview`
5824 /// must share the parent's `BufferIdentity` on **every** backend, not mint a
5825 /// fresh one. Identity-keyed caches (the GL EGLImage import) rely on this to
5826 /// treat offset-distinct windows of one buffer as a single import; a fresh
5827 /// identity would silently break that and regress zero-copy import reuse.
5828 ///
5829 /// Runs each backend that can be allocated without a GPU/GL context on the
5830 /// test host: `Mem` (always); `Shm` (when POSIX shm is available); the
5831 /// platform-native zero-copy buffer `Dma` (DMA-BUF on Linux / IOSurface on
5832 /// macOS, when available). `Pbo` shares its identity the same way (see
5833 /// `pbo.rs` `view`) but needs a live GL context, so it is exercised by the
5834 /// image-crate GL tests rather than here.
5835 #[test]
5836 fn subview_shares_buffer_identity_all_backends() {
5837 // u8 has align 1, so every byte offset is valid for the alignment check;
5838 // this isolates the identity-sharing contract from alignment concerns.
5839 let assert_shares = |memory: TensorMemory, label: &str| {
5840 let parent = Tensor::<u8>::new(&[64], Some(memory), None)
5841 .unwrap_or_else(|e| panic!("{label}: parent alloc failed: {e:?}"));
5842 let parent_id = parent.buffer_identity().id();
5843 // Two offset-distinct windows must both carry the parent's identity.
5844 let v0 = parent
5845 .subview(0, &[16])
5846 .unwrap_or_else(|e| panic!("{label}: subview(0) failed: {e:?}"));
5847 let v1 = parent
5848 .subview(16, &[16])
5849 .unwrap_or_else(|e| panic!("{label}: subview(16) failed: {e:?}"));
5850 assert_eq!(
5851 v0.buffer_identity().id(),
5852 parent_id,
5853 "{label}: subview(0) minted a fresh BufferIdentity"
5854 );
5855 assert_eq!(
5856 v1.buffer_identity().id(),
5857 parent_id,
5858 "{label}: subview(16) minted a fresh BufferIdentity"
5859 );
5860 };
5861
5862 assert_shares(TensorMemory::Mem, "Mem");
5863
5864 #[cfg(unix)]
5865 if crate::is_shm_available() {
5866 assert_shares(TensorMemory::Shm, "Shm");
5867 }
5868
5869 // Dma == DMA-BUF on Linux, IOSurface on macOS; same public variant.
5870 if crate::is_gpu_buffer_available() {
5871 assert_shares(TensorMemory::Dma, "Dma");
5872 }
5873 }
5874
5875 #[test]
5876 fn mem_subview_four_views_no_aliasing() {
5877 // One [4,3] f32 parent; four [1,3] views at 12-byte strides, each
5878 // written independently. Exercises a multi-byte element type (offsets
5879 // must stay element-aligned) and N-way zero-copy sharing.
5880 let parent = Tensor::<f32>::new(&[4, 3], Some(TensorMemory::Mem), None).unwrap();
5881 let frame = 3 * std::mem::size_of::<f32>();
5882 for i in 0..4 {
5883 let v = parent.subview(i * frame, &[1, 3]).unwrap();
5884 let val = i as f32 + 1.0;
5885 v.map()
5886 .unwrap()
5887 .as_mut_slice()
5888 .copy_from_slice(&[val, val, val]);
5889 }
5890 assert_eq!(
5891 parent.map().unwrap().as_slice(),
5892 &[1.0, 1.0, 1.0, 2.0, 2.0, 2.0, 3.0, 3.0, 3.0, 4.0, 4.0, 4.0]
5893 );
5894 }
5895
5896 #[test]
5897 fn mem_subview_inherits_format_and_row_stride() {
5898 // A sub-view is a ready-to-use sub-image: it inherits the parent's
5899 // pixel format and (crucially) its padded row stride, so a strided
5900 // parent yields strided windows. Set a stride wider than the tight row
5901 // to exercise the row_stride inheritance path specifically.
5902 let mut parent = Tensor::<u8>::image(
5903 100,
5904 100,
5905 PixelFormat::Rgba,
5906 Some(TensorMemory::Mem),
5907 crate::CpuAccess::ReadWrite,
5908 )
5909 .unwrap();
5910 parent.set_row_stride_unchecked(512); // padded stride (> 100*4)
5911 let view = parent.subview(4096, &[10, 10, 4]).unwrap();
5912 assert_eq!(view.format(), Some(PixelFormat::Rgba), "format inherited");
5913 assert_eq!(view.row_stride(), Some(512), "row_stride inherited");
5914 }
5915
5916 #[test]
5917 fn mem_strided_subview_maps_offset_and_byte_size() {
5918 // Integration of the sub-region offset (PR #89) and the strided-map
5919 // `byte_size_override` (PR #90): a strided sub-view exposes its full
5920 // padded window (`row_stride × rows`) starting at the view's byte
5921 // offset, mapped zero-copy into the parent.
5922 let parent = Tensor::<u8>::new(&[2048], Some(TensorMemory::Mem), None).unwrap();
5923 let mut view = parent.subview(128, &[8, 16]).unwrap(); // 8 rows × 16 @ off 128
5924 assert_eq!(view.plane_offset(), Some(128));
5925 view.set_row_stride_unchecked(32); // padded stride (> 16)
5926
5927 {
5928 let mut m = view.map().unwrap();
5929 let s = m.as_mut_slice();
5930 // Strided map exposes the padded window: stride(32) × rows(8) = 256.
5931 assert_eq!(
5932 s.len(),
5933 256,
5934 "strided map exposes the full padded byte window"
5935 );
5936 s[0] = 0xAA; // row 0, col 0
5937 s[32] = 0xBB; // row 1, col 0 (one stride in)
5938 }
5939
5940 // Zero-copy: the writes land in the parent at the view's offset.
5941 let p = parent.map().unwrap();
5942 let pb = p.as_slice();
5943 assert_eq!(pb[128], 0xAA, "row 0 writes at parent offset 128");
5944 assert_eq!(
5945 pb[128 + 32],
5946 0xBB,
5947 "row 1 writes at parent offset 128 + stride"
5948 );
5949 }
5950
5951 #[test]
5952 #[cfg(unix)]
5953 fn shm_subview_partitions_parent_buffer() {
5954 // Mirrors `mem_subview_partitions_parent_buffer` for Shm: one [2,4] u8
5955 // parent shared segment (8 bytes); two [1,4] sub-views at byte offsets 0
5956 // and 4 must share the segment (zero-copy, via cloned fd) and be
5957 // independently writable — view 0 owns [0,4), view 1 owns [4,8).
5958 if !crate::is_shm_available() {
5959 eprintln!("SKIPPED: shm not available");
5960 return;
5961 }
5962 let parent = Tensor::<u8>::new(&[2, 4], Some(TensorMemory::Shm), None).unwrap();
5963 let view0 = parent.subview(0, &[1, 4]).expect("shm subview at offset 0");
5964 let view1 = parent.subview(4, &[1, 4]).expect("shm subview at offset 4");
5965
5966 view1
5967 .map()
5968 .unwrap()
5969 .as_mut_slice()
5970 .copy_from_slice(&[10, 20, 30, 40]);
5971 view0
5972 .map()
5973 .unwrap()
5974 .as_mut_slice()
5975 .copy_from_slice(&[1, 2, 3, 4]);
5976
5977 assert_eq!(view0.map().unwrap().as_slice(), &[1, 2, 3, 4]);
5978 assert_eq!(view1.map().unwrap().as_slice(), &[10, 20, 30, 40]);
5979 // The parent sees the full partitioned segment (shared, zero-copy).
5980 assert_eq!(
5981 parent.map().unwrap().as_slice(),
5982 &[1, 2, 3, 4, 10, 20, 30, 40]
5983 );
5984 // A sub-view of a sub-view composes the offset.
5985 let nested = view1.subview(2, &[1, 2]).expect("nested shm subview");
5986 assert_eq!(nested.map().unwrap().as_slice(), &[30, 40]);
5987 }
5988
5989 #[test]
5990 #[cfg(unix)]
5991 fn shm_subview_rejects_unaligned_and_oob() {
5992 if !crate::is_shm_available() {
5993 eprintln!("SKIPPED: shm not available");
5994 return;
5995 }
5996 // f32 align 4: a 2-byte offset cannot back a valid `*const f32`.
5997 let parent = Tensor::<f32>::new(&[8], Some(TensorMemory::Shm), None).unwrap();
5998 assert!(parent.subview(2, &[1]).is_err());
5999 assert!(parent.subview(4, &[1]).is_ok());
6000 // Out of bounds: offset 6 + 4 bytes = 10 > 8-byte (u8) segment.
6001 let p2 = Tensor::<u8>::new(&[8], Some(TensorMemory::Shm), None).unwrap();
6002 assert!(p2.subview(6, &[4]).is_err());
6003 }
6004
6005 #[test]
6006 #[cfg(target_os = "linux")]
6007 fn dma_subview_matches_mem_subview() {
6008 // Serialize against the fd-leak tests: this test opens DMA fds (alloc +
6009 // clone_fd), which would otherwise perturb their fd counts.
6010 let _lock = FD_LOCK.read().unwrap();
6011 // Identical sub-view semantics across Dma (shared fd) and Mem (shared
6012 // Arc): same offsets → same logical windows → same partition.
6013 let dma = match Tensor::<u8>::new(&[8], Some(TensorMemory::Dma), None) {
6014 Ok(t) => t,
6015 Err(_) => {
6016 eprintln!("SKIPPED: DMA not available");
6017 return;
6018 }
6019 };
6020 let mem = Tensor::<u8>::new(&[8], Some(TensorMemory::Mem), None).unwrap();
6021 for parent in [&dma, &mem] {
6022 let v0 = parent.subview(0, &[4]).unwrap();
6023 let v1 = parent.subview(4, &[4]).unwrap();
6024 v0.map()
6025 .unwrap()
6026 .as_mut_slice()
6027 .copy_from_slice(&[1, 2, 3, 4]);
6028 v1.map()
6029 .unwrap()
6030 .as_mut_slice()
6031 .copy_from_slice(&[5, 6, 7, 8]);
6032 assert_eq!(parent.map().unwrap().as_slice(), &[1, 2, 3, 4, 5, 6, 7, 8]);
6033 }
6034 }
6035
6036 #[test]
6037 #[cfg(target_os = "linux")]
6038 fn dma_strided_subview_maps_padded_window() {
6039 // The strided-map path differs by backing: DMA maps through
6040 // `mmap_offset` + the `byte_size_override`, not the Mem `Arc` slice. A
6041 // padded sub-view of a DMA buffer must still expose its full
6042 // `row_stride × rows` window zero-copy at the view's offset (the GPU
6043 // batched-render-to-DMA case). Mirrors
6044 // `mem_strided_subview_maps_offset_and_byte_size` on a Dma parent.
6045 let _lock = FD_LOCK.read().unwrap();
6046 let parent = match Tensor::<u8>::new(&[2048], Some(TensorMemory::Dma), None) {
6047 Ok(t) => t,
6048 Err(_) => {
6049 eprintln!("SKIPPED: DMA not available");
6050 return;
6051 }
6052 };
6053 let mut view = parent.subview(128, &[8, 16]).unwrap();
6054 assert_eq!(view.plane_offset(), Some(128));
6055 view.set_row_stride_unchecked(32); // padded stride (> 16)
6056
6057 {
6058 let mut m = view.map().unwrap();
6059 let s = m.as_mut_slice();
6060 assert_eq!(s.len(), 256, "strided DMA map exposes stride(32) × rows(8)");
6061 s[0] = 0xAA; // row 0, col 0
6062 s[32] = 0xBB; // row 1, col 0 (one stride in)
6063 }
6064
6065 let p = parent.map().unwrap();
6066 let pb = p.as_slice();
6067 assert_eq!(pb[128], 0xAA, "row 0 writes at parent offset 128");
6068 assert_eq!(
6069 pb[128 + 32],
6070 0xBB,
6071 "row 1 writes at parent offset 128 + stride"
6072 );
6073 }
6074
6075 #[test]
6076 #[cfg(target_os = "linux")]
6077 fn view_single_row_snapshots_parent_stride() {
6078 // A single-row `view()` keeps a TIGHT `row_stride` for map-span safety,
6079 // but its `view_origin` snapshots the PARENT row stride — the GL backend
6080 // keys its EGLImage import/pitch on that snapshot (not the view's tight
6081 // stride), so single-row and multi-row sibling views collapse onto the
6082 // same parent import.
6083 let _lock = FD_LOCK.read().unwrap();
6084 // 8x4 RGBA with a padded 64-byte row stride (tight row = 8*4 = 32).
6085 let parent = match Tensor::<u8>::image_with_stride(
6086 8,
6087 4,
6088 PixelFormat::Rgba,
6089 64,
6090 Some(TensorMemory::Dma),
6091 crate::CpuAccess::ReadWrite,
6092 ) {
6093 Ok(t) => t,
6094 Err(_) => {
6095 eprintln!("SKIPPED: DMA not available");
6096 return;
6097 }
6098 };
6099 assert_eq!(parent.effective_row_stride(), Some(64));
6100 // Bottom row (y=3) at x>0 — the case the tight single-row stride guards.
6101 let row = parent.view(Region::new(2, 3, 4, 1)).unwrap();
6102 // The view's own stride is tight (4*4 = 16) so its strided map stays in
6103 // bounds; the GL-facing parent pitch (64) lives in `view_origin`.
6104 assert_eq!(row.effective_row_stride(), Some(16));
6105 let vo = row.view_origin().expect("a view carries a view_origin");
6106 assert_eq!(
6107 vo.parent_row_stride, 64,
6108 "GL keys/pitches a view on the parent stride, not its tight one"
6109 );
6110 // The tight stride keeps map() in-bounds for the bottom / x>0 single row.
6111 assert_eq!(row.map().unwrap().as_slice().len(), 16);
6112 }
6113
6114 #[test]
6115 fn test_mem_tensor() {
6116 let shape = vec![2, 3, 4];
6117 let tensor =
6118 MemTensor::<f32>::new(&shape, Some("test_tensor")).expect("Failed to create tensor");
6119 assert_eq!(tensor.shape(), &shape);
6120 assert_eq!(tensor.size(), 2 * 3 * 4 * std::mem::size_of::<f32>());
6121 assert_eq!(tensor.name(), "test_tensor");
6122
6123 {
6124 let mut tensor_map = tensor.map().expect("Failed to map memory");
6125 tensor_map.fill(42.0);
6126 assert!(tensor_map.iter().all(|&x| x == 42.0));
6127 }
6128
6129 let mut tensor = MemTensor::<u8>::new(&shape, None).expect("Failed to create tensor");
6130 assert_eq!(tensor.shape(), &shape);
6131 let new_shape = vec![3, 4, 4];
6132 assert!(
6133 tensor.reshape(&new_shape).is_err(),
6134 "Reshape should fail due to size mismatch"
6135 );
6136 assert_eq!(tensor.shape(), &shape, "Shape should remain unchanged");
6137
6138 let new_shape = vec![2, 3, 4];
6139 tensor.reshape(&new_shape).expect("Reshape should succeed");
6140 assert_eq!(
6141 tensor.shape(),
6142 &new_shape,
6143 "Shape should be updated after successful reshape"
6144 );
6145
6146 {
6147 let mut tensor_map = tensor.map().expect("Failed to map memory");
6148 tensor_map.fill(1);
6149 assert!(tensor_map.iter().all(|&x| x == 1));
6150 }
6151
6152 {
6153 let mut tensor_map = tensor.map().expect("Failed to map memory");
6154 tensor_map[2] = 42;
6155 assert_eq!(tensor_map[1], 1, "Value at index 1 should be 1");
6156 assert_eq!(tensor_map[2], 42, "Value at index 2 should be 42");
6157 }
6158 }
6159
6160 #[test]
6161 #[cfg(target_os = "linux")]
6162 fn test_dma_no_fd_leaks() {
6163 let _lock = FD_LOCK.write().unwrap();
6164 if !is_dma_available() {
6165 log::warn!(
6166 "SKIPPED: {} - DMA memory allocation not available (permission denied or no DMA-BUF support)",
6167 function!()
6168 );
6169 return;
6170 }
6171
6172 let proc = procfs::process::Process::myself()
6173 .expect("Failed to get current process using /proc/self");
6174
6175 let start_open_fds = proc
6176 .fd_count()
6177 .expect("Failed to get open file descriptor count");
6178
6179 for _ in 0..100 {
6180 let tensor = Tensor::<u8>::new(&[100, 100], Some(TensorMemory::Dma), None)
6181 .expect("Failed to create tensor");
6182 let mut map = tensor.map().unwrap();
6183 map.as_mut_slice().fill(233);
6184 }
6185
6186 let end_open_fds = proc
6187 .fd_count()
6188 .expect("Failed to get open file descriptor count");
6189
6190 assert_eq!(
6191 start_open_fds, end_open_fds,
6192 "File descriptor leak detected: {} -> {}",
6193 start_open_fds, end_open_fds
6194 );
6195 }
6196
6197 #[test]
6198 #[cfg(target_os = "linux")]
6199 fn test_dma_from_fd_no_fd_leaks() {
6200 let _lock = FD_LOCK.write().unwrap();
6201 if !is_dma_available() {
6202 log::warn!(
6203 "SKIPPED: {} - DMA memory allocation not available (permission denied or no DMA-BUF support)",
6204 function!()
6205 );
6206 return;
6207 }
6208
6209 let proc = procfs::process::Process::myself()
6210 .expect("Failed to get current process using /proc/self");
6211
6212 let start_open_fds = proc
6213 .fd_count()
6214 .expect("Failed to get open file descriptor count");
6215
6216 let orig = Tensor::<u8>::new(&[100, 100], Some(TensorMemory::Dma), None).unwrap();
6217
6218 for _ in 0..100 {
6219 let tensor =
6220 Tensor::<u8>::from_fd(orig.clone_fd().unwrap(), orig.shape(), None).unwrap();
6221 assert_eq!(
6222 tensor.memory(),
6223 TensorMemory::Dma,
6224 "DMA-BUF fd must import as Dma, not be silently downgraded"
6225 );
6226 let mut map = tensor.map().unwrap();
6227 map.as_mut_slice().fill(233);
6228 }
6229 drop(orig);
6230
6231 let end_open_fds = proc.fd_count().unwrap();
6232
6233 assert_eq!(
6234 start_open_fds, end_open_fds,
6235 "File descriptor leak detected: {} -> {}",
6236 start_open_fds, end_open_fds
6237 );
6238 }
6239
6240 /// A filesystem magic must report its true 32-bit value regardless of
6241 /// how wide, and how signed, `fstatfs`'s `f_type` is on the target.
6242 ///
6243 /// `fs_type_t` is `__fsword_t` on Linux/gnu — `i64` on 64-bit but `i32`
6244 /// on 32-bit (armv7, i686, aarch64-ilp32) — `c_int` on uclibc, and
6245 /// unsigned on musl and s390x. Widening a *signed* 32-bit `f_type`
6246 /// sign-extends every magic with bit 31 set, so a naive widening cast
6247 /// reports e.g. `0xffffffff958458f6` for hugetlbfs. The docs tell
6248 /// callers to look the reported value up in `include/uapi/linux/magic.h`,
6249 /// so a sign-extended value is actively misleading.
6250 ///
6251 /// This cannot be reproduced on a 64-bit host — the defect only exists
6252 /// where `f_type` is a signed 32-bit type — so the test drives the
6253 /// conversion directly with the value such a target would produce.
6254 #[test]
6255 #[cfg(target_os = "linux")]
6256 fn test_fs_magic_normalizes_sign_extended_values() {
6257 // Magics whose bit 31 is set. HUGETLBFS is the one that matters
6258 // most in practice: a MFD_HUGETLB memfd is the likeliest fd to land
6259 // in the UnknownBufferType arm.
6260 const HUGETLBFS_MAGIC: u32 = 0x9584_58f6;
6261 const F2FS_MAGIC: u32 = 0xf2f5_2010;
6262 const BTRFS_MAGIC: u32 = 0x9123_683e;
6263
6264 for magic in [HUGETLBFS_MAGIC, F2FS_MAGIC, BTRFS_MAGIC] {
6265 // 64-bit gnu: f_type is i64 and already holds the true value.
6266 assert_eq!(fs_magic(i64::from(magic)), magic);
6267
6268 // 32-bit gnu / uclibc: f_type is i32, so the value arrives
6269 // sign-extended once widened. It must still report its true
6270 // 32 bits.
6271 let sign_extended = i64::from(magic as i32);
6272 assert!(sign_extended < 0, "{magic:#x} should have bit 31 set");
6273 assert_eq!(
6274 fs_magic(sign_extended),
6275 magic,
6276 "{magic:#x} must survive a signed 32-bit f_type"
6277 );
6278 }
6279
6280 // The two magics we actually classify on are below 2^31, so they are
6281 // unaffected by signedness either way — this is why the DMA-vs-SHM
6282 // decision is correct on 32-bit even without this normalization.
6283 for magic in [DMA_BUF_MAGIC, TMPFS_MAGIC] {
6284 assert!(magic < 0x8000_0000);
6285 assert_eq!(fs_magic(i64::from(magic as i32)), magic);
6286 }
6287 }
6288
6289 /// A DMA-BUF fd must import as [`TensorMemory::Dma`].
6290 ///
6291 /// Regression test for the `st_dev` minor-number classifier, which
6292 /// hardcoded `9 | 10` as "this is DMA". Those minors come from
6293 /// `get_anon_bdev()` and are assigned first-come-first-served at boot,
6294 /// so they vary by kernel build and boot order — a real DMA-BUF is
6295 /// minor 12 on x86 desktop and minor 8 on the ADIS Verdin. Both fell
6296 /// into the `_` arm and imported as SHM, which "works" (a DMA-BUF is
6297 /// mmap-able) but silently forfeits zero-copy.
6298 #[test]
6299 #[cfg(target_os = "linux")]
6300 fn test_from_fd_dma_imports_as_dma() {
6301 let _lock = FD_LOCK.read().unwrap();
6302 if !is_dma_available() {
6303 log::warn!("SKIPPED: {} - DMA memory not available", function!());
6304 return;
6305 }
6306
6307 let orig = Tensor::<u8>::new(&[64, 64], Some(TensorMemory::Dma), None).unwrap();
6308 assert_eq!(orig.memory(), TensorMemory::Dma);
6309
6310 let imported = Tensor::<u8>::from_fd(orig.clone_fd().unwrap(), orig.shape(), None).unwrap();
6311
6312 assert_eq!(
6313 imported.memory(),
6314 TensorMemory::Dma,
6315 "a DMA-BUF fd must import as Dma"
6316 );
6317 }
6318
6319 /// A tmpfs/SHM fd must import as [`TensorMemory::Shm`].
6320 ///
6321 /// The companion to `test_from_fd_dma_imports_as_dma`: confirms the
6322 /// magic-based classifier identifies SHM positively rather than by
6323 /// falling through.
6324 #[test]
6325 #[cfg(target_os = "linux")]
6326 fn test_from_fd_shm_imports_as_shm() {
6327 let _lock = FD_LOCK.read().unwrap();
6328 if !is_shm_available() {
6329 log::warn!("SKIPPED: {} - SHM memory not available", function!());
6330 return;
6331 }
6332
6333 let orig = Tensor::<u8>::new(&[64, 64], Some(TensorMemory::Shm), None).unwrap();
6334 assert_eq!(orig.memory(), TensorMemory::Shm);
6335
6336 let imported = Tensor::<u8>::from_fd(orig.clone_fd().unwrap(), orig.shape(), None).unwrap();
6337
6338 assert_eq!(
6339 imported.memory(),
6340 TensorMemory::Shm,
6341 "a tmpfs fd must import as Shm"
6342 );
6343 }
6344
6345 /// An fd that is neither a DMA-BUF nor tmpfs must be rejected.
6346 ///
6347 /// A pipe is the convenient probe: it lives on `pipefs`, another
6348 /// `get_anon_bdev()` pseudo-filesystem, so it shares major 0 with the
6349 /// buffer types we do support and is only distinguishable by magic.
6350 /// Importing one as SHM is meaningless — `mmap` on a pipe fails — so
6351 /// the classifier must say "unknown" rather than guess.
6352 #[test]
6353 #[cfg(target_os = "linux")]
6354 fn test_from_fd_rejects_unknown_filesystem() {
6355 let _lock = FD_LOCK.read().unwrap();
6356
6357 let (read_end, _write_end) = nix::unistd::pipe().unwrap();
6358
6359 let result = Tensor::<u8>::from_fd(read_end, &[64], None);
6360
6361 match result {
6362 Err(Error::UnknownBufferType(magic)) => {
6363 // PIPEFS_MAGIC, from include/uapi/linux/magic.h
6364 assert_eq!(magic, 0x5049_5045, "expected PIPEFS_MAGIC");
6365 }
6366 other => panic!("expected UnknownBufferType for a pipe fd, got {other:?}"),
6367 }
6368 }
6369
6370 #[test]
6371 #[cfg(target_os = "linux")]
6372 fn test_shm_no_fd_leaks() {
6373 let _lock = FD_LOCK.write().unwrap();
6374 if !is_shm_available() {
6375 log::warn!(
6376 "SKIPPED: {} - SHM memory allocation not available (permission denied or no SHM support)",
6377 function!()
6378 );
6379 return;
6380 }
6381
6382 let proc = procfs::process::Process::myself()
6383 .expect("Failed to get current process using /proc/self");
6384
6385 let start_open_fds = proc
6386 .fd_count()
6387 .expect("Failed to get open file descriptor count");
6388
6389 for _ in 0..100 {
6390 let tensor = Tensor::<u8>::new(&[100, 100], Some(TensorMemory::Shm), None)
6391 .expect("Failed to create tensor");
6392 let mut map = tensor.map().unwrap();
6393 map.as_mut_slice().fill(233);
6394 }
6395
6396 let end_open_fds = proc
6397 .fd_count()
6398 .expect("Failed to get open file descriptor count");
6399
6400 assert_eq!(
6401 start_open_fds, end_open_fds,
6402 "File descriptor leak detected: {} -> {}",
6403 start_open_fds, end_open_fds
6404 );
6405 }
6406
6407 #[test]
6408 #[cfg(target_os = "linux")]
6409 fn test_shm_from_fd_no_fd_leaks() {
6410 let _lock = FD_LOCK.write().unwrap();
6411 if !is_shm_available() {
6412 log::warn!(
6413 "SKIPPED: {} - SHM memory allocation not available (permission denied or no SHM support)",
6414 function!()
6415 );
6416 return;
6417 }
6418
6419 let proc = procfs::process::Process::myself()
6420 .expect("Failed to get current process using /proc/self");
6421
6422 let start_open_fds = proc
6423 .fd_count()
6424 .expect("Failed to get open file descriptor count");
6425
6426 let orig = Tensor::<u8>::new(&[100, 100], Some(TensorMemory::Shm), None).unwrap();
6427
6428 for _ in 0..100 {
6429 let tensor =
6430 Tensor::<u8>::from_fd(orig.clone_fd().unwrap(), orig.shape(), None).unwrap();
6431 let mut map = tensor.map().unwrap();
6432 map.as_mut_slice().fill(233);
6433 }
6434 drop(orig);
6435
6436 let end_open_fds = proc.fd_count().unwrap();
6437
6438 assert_eq!(
6439 start_open_fds, end_open_fds,
6440 "File descriptor leak detected: {} -> {}",
6441 start_open_fds, end_open_fds
6442 );
6443 }
6444
6445 #[cfg(feature = "ndarray")]
6446 #[test]
6447 fn test_ndarray() {
6448 let _lock = FD_LOCK.read().unwrap();
6449 let shape = vec![2, 3, 4];
6450 let tensor = Tensor::<f32>::new(&shape, None, None).expect("Failed to create tensor");
6451
6452 let mut tensor_map = tensor.map().expect("Failed to map tensor memory");
6453 tensor_map.fill(1.0);
6454
6455 let view = tensor_map.view().expect("Failed to get ndarray view");
6456 assert_eq!(view.shape(), &[2, 3, 4]);
6457 assert!(view.iter().all(|&x| x == 1.0));
6458
6459 let mut view_mut = tensor_map
6460 .view_mut()
6461 .expect("Failed to get mutable ndarray view");
6462 view_mut[[0, 0, 0]] = 42.0;
6463 assert_eq!(view_mut[[0, 0, 0]], 42.0);
6464 assert_eq!(tensor_map[0], 42.0, "Value at index 0 should be 42");
6465 }
6466
6467 #[test]
6468 fn test_buffer_identity_unique() {
6469 let id1 = BufferIdentity::new();
6470 let id2 = BufferIdentity::new();
6471 assert_ne!(
6472 id1.id(),
6473 id2.id(),
6474 "Two identities should have different ids"
6475 );
6476 }
6477
6478 #[test]
6479 fn test_buffer_identity_clone_shares_guard() {
6480 let id1 = BufferIdentity::new();
6481 let weak = id1.weak();
6482 assert!(
6483 weak.upgrade().is_some(),
6484 "Weak should be alive while original exists"
6485 );
6486
6487 let id2 = id1.clone();
6488 assert_eq!(id1.id(), id2.id(), "Cloned identity should have same id");
6489
6490 drop(id1);
6491 assert!(
6492 weak.upgrade().is_some(),
6493 "Weak should still be alive (clone holds Arc)"
6494 );
6495
6496 drop(id2);
6497 assert!(
6498 weak.upgrade().is_none(),
6499 "Weak should be dead after all clones dropped"
6500 );
6501 }
6502
6503 #[test]
6504 fn test_tensor_buffer_identity() {
6505 let t1 = Tensor::<u8>::new(&[100], Some(TensorMemory::Mem), Some("t1")).unwrap();
6506 let t2 = Tensor::<u8>::new(&[100], Some(TensorMemory::Mem), Some("t2")).unwrap();
6507 assert_ne!(
6508 t1.buffer_identity().id(),
6509 t2.buffer_identity().id(),
6510 "Different tensors should have different buffer ids"
6511 );
6512 }
6513
6514 // ------------------------------------------------------------------------
6515 // Quantization — constructor validation + accessor correctness.
6516 // ------------------------------------------------------------------------
6517
6518 #[test]
6519 fn test_quantization_per_tensor_constructors() {
6520 let q = Quantization::per_tensor(0.1, -5);
6521 assert!(q.is_per_tensor());
6522 assert!(!q.is_per_channel());
6523 assert!(!q.is_symmetric());
6524 assert_eq!(q.scale(), &[0.1]);
6525 assert_eq!(q.zero_point(), Some(&[-5][..]));
6526
6527 let qs = Quantization::per_tensor_symmetric(0.05);
6528 assert!(qs.is_per_tensor());
6529 assert!(qs.is_symmetric());
6530 assert_eq!(qs.zero_point(), None);
6531 }
6532
6533 #[test]
6534 fn test_quantization_per_channel_constructors() {
6535 let q = Quantization::per_channel(vec![0.1, 0.2, 0.3], vec![0, -1, 1], 2).unwrap();
6536 assert!(q.is_per_channel());
6537 assert!(!q.is_symmetric());
6538 assert_eq!(q.axis(), Some(2));
6539 assert_eq!(q.scale().len(), 3);
6540
6541 let qs = Quantization::per_channel_symmetric(vec![0.054, 0.089, 0.195], 0).unwrap();
6542 assert!(qs.is_per_channel());
6543 assert!(qs.is_symmetric());
6544 assert_eq!(qs.axis(), Some(0));
6545 }
6546
6547 #[test]
6548 fn test_quantization_per_channel_length_mismatch_rejected() {
6549 // len(scales) != len(zero_points) → rejected at construction.
6550 let err = Quantization::per_channel(vec![0.1, 0.2], vec![0, 0, 0], 0).unwrap_err();
6551 assert!(matches!(err, Error::QuantizationInvalid { .. }));
6552 }
6553
6554 #[test]
6555 fn test_quantization_per_channel_empty_rejected() {
6556 let err = Quantization::per_channel_symmetric(vec![], 0).unwrap_err();
6557 assert!(matches!(err, Error::QuantizationInvalid { .. }));
6558 }
6559
6560 /// Constructors guard scale/zero_point length invariants, but
6561 /// `Quantization` is `Deserialize`, so malformed JSON (e.g. an
6562 /// empty `scale` array, or `zero_point` length that disagrees with
6563 /// `scale`) bypasses the constructor checks. `set_quantization`
6564 /// must reject these via `validate()` so they don't poison
6565 /// downstream `mode()` selection or per-channel kernel indexing.
6566 #[test]
6567 fn test_quantization_validate_rejects_malformed_deserialize() {
6568 let mut t = Tensor::<i8>::new(&[1, 1, 4], Some(TensorMemory::Mem), None).unwrap();
6569
6570 // Empty scale array: must be rejected.
6571 let q: Quantization = serde_json::from_str(r#"{"scale": []}"#).unwrap();
6572 assert!(matches!(
6573 t.set_quantization(q).unwrap_err(),
6574 Error::QuantizationInvalid { .. }
6575 ));
6576
6577 // Per-tensor with multi-element zero_point: must be rejected.
6578 let q: Quantization =
6579 serde_json::from_str(r#"{"scale": 0.1, "zero_point": [0, 0, 0]}"#).unwrap();
6580 assert!(matches!(
6581 t.set_quantization(q).unwrap_err(),
6582 Error::QuantizationInvalid { .. }
6583 ));
6584
6585 // Per-channel zero_point length != scale length: must be rejected.
6586 let q: Quantization = serde_json::from_str(
6587 r#"{"scale": [0.1, 0.2, 0.3, 0.4], "zero_point": [0, 0], "axis": 2}"#,
6588 )
6589 .unwrap();
6590 assert!(matches!(
6591 t.set_quantization(q).unwrap_err(),
6592 Error::QuantizationInvalid { .. }
6593 ));
6594 }
6595
6596 #[test]
6597 fn test_quantization_mode_dispatch() {
6598 let pt = Quantization::per_tensor(0.1, -5);
6599 assert!(matches!(
6600 pt.mode(),
6601 QuantMode::PerTensor { scale, zero_point } if scale == 0.1 && zero_point == -5
6602 ));
6603
6604 let pts = Quantization::per_tensor_symmetric(0.05);
6605 assert!(matches!(
6606 pts.mode(),
6607 QuantMode::PerTensorSymmetric { scale } if scale == 0.05
6608 ));
6609
6610 let pc = Quantization::per_channel(vec![0.1, 0.2], vec![0, -1], 2).unwrap();
6611 assert!(matches!(pc.mode(), QuantMode::PerChannel { axis: 2, .. }));
6612
6613 let pcs = Quantization::per_channel_symmetric(vec![0.1, 0.2], 0).unwrap();
6614 assert!(matches!(
6615 pcs.mode(),
6616 QuantMode::PerChannelSymmetric { axis: 0, .. }
6617 ));
6618 }
6619
6620 #[test]
6621 fn test_tensor_quantization_roundtrip_integer() {
6622 let mut t = Tensor::<i8>::new(&[2, 3, 4], Some(TensorMemory::Mem), None).unwrap();
6623 assert!(t.quantization().is_none());
6624 t.set_quantization(Quantization::per_tensor(0.1, -5))
6625 .unwrap();
6626 let q = t.quantization().unwrap();
6627 assert_eq!(q.scale(), &[0.1]);
6628 t.clear_quantization();
6629 assert!(t.quantization().is_none());
6630 }
6631
6632 #[test]
6633 fn test_tensor_with_quantization_builder() {
6634 let t = Tensor::<i8>::new(&[4, 4], Some(TensorMemory::Mem), None)
6635 .unwrap()
6636 .with_quantization(Quantization::per_tensor_symmetric(0.05))
6637 .unwrap();
6638 assert!(t.quantization().is_some());
6639 }
6640
6641 #[test]
6642 fn test_tensor_dyn_quantization_float_arm_returns_none() {
6643 let t = Tensor::<f32>::new(&[2, 2], Some(TensorMemory::Mem), None).unwrap();
6644 let td = TensorDyn::F32(t);
6645 assert!(td.quantization().is_none());
6646 }
6647
6648 #[test]
6649 fn test_tensor_dyn_set_quantization_float_arm_errors() {
6650 let t = Tensor::<f32>::new(&[2, 2], Some(TensorMemory::Mem), None).unwrap();
6651 let mut td = TensorDyn::F32(t);
6652 let err = td
6653 .set_quantization(Quantization::per_tensor(0.1, 0))
6654 .unwrap_err();
6655 // float path returns a QuantizationInvalid error.
6656 assert!(matches!(err, Error::QuantizationInvalid { .. }));
6657 }
6658
6659 /// Compile-time type gate — calling `Tensor::<f32>::quantization()` must
6660 /// fail to compile (the `IntegerType` trait bound is not satisfied by
6661 /// `f32`). This doctest anchors the invariant.
6662 ///
6663 /// ```compile_fail
6664 /// use edgefirst_tensor::{Tensor, TensorMemory};
6665 /// let t = Tensor::<f32>::new(&[2, 2], Some(TensorMemory::Mem), None).unwrap();
6666 /// let _ = t.quantization(); // compile error: f32 not IntegerType
6667 /// ```
6668 fn _compile_fail_doctest_anchor() {}
6669
6670 // Any test that cares about the fd count must grab it exclusively.
6671 // Any tests which modifies the fd count by opening or closing fds must grab it
6672 // shared.
6673 pub static FD_LOCK: RwLock<()> = RwLock::new(());
6674
6675 /// Test that DMA is NOT available on non-Linux platforms.
6676 /// This verifies the cross-platform behavior of is_dma_available().
6677 #[test]
6678 #[cfg(not(target_os = "linux"))]
6679 fn test_dma_not_available_on_non_linux() {
6680 assert!(
6681 !is_dma_available(),
6682 "DMA memory allocation should NOT be available on non-Linux platforms"
6683 );
6684 }
6685
6686 #[test]
6687 fn colorimetry_defaults_none_and_roundtrips_without_auto_fill() {
6688 use crate::{ColorEncoding, ColorRange, Colorimetry, PixelFormat, TensorMemory};
6689 let mut t = Tensor::<u8>::image(
6690 1280,
6691 720,
6692 PixelFormat::Nv12,
6693 Some(TensorMemory::Mem),
6694 crate::CpuAccess::ReadWrite,
6695 )
6696 .unwrap();
6697 assert_eq!(t.colorimetry(), None); // default undefined
6698 let c = Colorimetry::default()
6699 .with_encoding(ColorEncoding::Bt709)
6700 .with_range(ColorRange::Limited);
6701 t.set_colorimetry(Some(c));
6702 assert_eq!(t.colorimetry(), Some(c));
6703 // configure_image must NOT touch colorimetry
6704 t.configure_image(640, 480, PixelFormat::Grey).unwrap();
6705 assert_eq!(t.colorimetry(), Some(c));
6706 }
6707
6708 #[test]
6709 fn configure_image_within_capacity() {
6710 let mut t = Tensor::<u8>::image_with_capacity(
6711 640,
6712 480,
6713 PixelFormat::Rgb,
6714 None,
6715 crate::CpuAccess::ReadWrite,
6716 )
6717 .unwrap();
6718 t.configure_image(320, 240, PixelFormat::Nv12).unwrap();
6719 assert_eq!(t.format(), Some(PixelFormat::Nv12));
6720 assert_eq!(t.width(), Some(320));
6721 assert_eq!(t.height(), Some(240));
6722 assert_eq!(t.shape(), &[360, 320]); // 240*3/2
6723 }
6724
6725 #[test]
6726 fn configure_image_too_large_errors() {
6727 let mut t = Tensor::<u8>::image_with_capacity(
6728 64,
6729 64,
6730 PixelFormat::Grey,
6731 None,
6732 crate::CpuAccess::ReadWrite,
6733 )
6734 .unwrap();
6735 let err = t
6736 .configure_image(1920, 1080, PixelFormat::Nv12)
6737 .unwrap_err();
6738 assert!(matches!(err, Error::InsufficientCapacity { .. }));
6739 }
6740
6741 /// A reused max-sized IOSurface pool keeps its physical `bytesPerRow` when
6742 /// reconfigured to a smaller logical image (physical-grid / logical-ROI
6743 /// decoupling), instead of collapsing to the frame's natural row stride.
6744 #[test]
6745 #[cfg(target_os = "macos")]
6746 fn configure_image_preserves_iosurface_physical_stride() {
6747 // Pool: GREY/R8 IOSurface 100 wide → bytesPerRow padded to 128.
6748 let mut pool = Tensor::<u8>::image(
6749 100,
6750 64,
6751 PixelFormat::Grey,
6752 Some(TensorMemory::Dma),
6753 crate::CpuAccess::ReadWrite,
6754 )
6755 .unwrap();
6756 let pitch = pool.effective_row_stride().unwrap();
6757 assert!(
6758 pitch >= 128 && pitch.is_multiple_of(64),
6759 "padded bytesPerRow, got {pitch}"
6760 );
6761
6762 // Reconfigure to a smaller NV12 frame; the physical pitch must survive
6763 // (natural would be 32, but the surface stride is the 128-padded pitch).
6764 pool.configure_image(32, 16, PixelFormat::Nv12).unwrap();
6765 assert_eq!(pool.format(), Some(PixelFormat::Nv12));
6766 assert_eq!(pool.width(), Some(32));
6767 assert_eq!(pool.height(), Some(16));
6768 assert_eq!(
6769 pool.effective_row_stride(),
6770 Some(pitch),
6771 "configure_image must preserve the IOSurface physical bytesPerRow"
6772 );
6773
6774 // Reconfigure again to NV24 — pitch still preserved.
6775 pool.configure_image(32, 16, PixelFormat::Nv24).unwrap();
6776 assert_eq!(pool.effective_row_stride(), Some(pitch));
6777 }
6778
6779 /// `configure_image` on a Mem backing reconfigures to the format's
6780 /// **64-byte-aligned** row stride (the odd-dim contract: every image tensor
6781 /// carries a 64-aligned `row_stride`). For NV12 32×16 the minimum is
6782 /// `even(32)=32`, rounded up to the 64-byte alignment → 64. The capacity
6783 /// (64×64×4 RGBA = 16 KiB) easily holds the 24×64 = 1.5 KiB NV12 layout.
6784 #[test]
6785 fn configure_image_mem_aligns_stride() {
6786 let mut t = Tensor::<u8>::image_with_capacity(
6787 64,
6788 64,
6789 PixelFormat::Rgba,
6790 Some(TensorMemory::Mem),
6791 crate::CpuAccess::ReadWrite,
6792 )
6793 .unwrap();
6794 t.configure_image(32, 16, PixelFormat::Nv12).unwrap();
6795 let s = t.effective_row_stride().unwrap();
6796 assert_eq!(s % 64, 0, "stride must be 64-aligned");
6797 assert!(s >= 32, "stride must cover the even-width minimum");
6798 assert_eq!(s, 64);
6799 }
6800
6801 #[test]
6802 fn strided_mem_tensor_cpu_maps_full_padded_buffer() {
6803 // A packed RGBA image with row padding (GPU-pitch style): logical width
6804 // 8 px (32 B/row) but a 48-byte row stride. Over-allocate capacity (for
6805 // 16 px), narrow the logical width, then record the padded stride.
6806 // Previously `map()` rejected this on non-Linux with
6807 // "DMA backing is Linux-only"; HAL-owned Mem is now mappable.
6808 let mut t = Tensor::<u8>::image_with_capacity(
6809 16,
6810 3,
6811 PixelFormat::Rgba,
6812 Some(TensorMemory::Mem),
6813 crate::CpuAccess::ReadWrite,
6814 )
6815 .unwrap(); // capacity 3 × 16 × 4 = 192 B
6816 t.configure_image(8, 3, PixelFormat::Rgba).unwrap(); // logical [3, 8, 4] = 96 B
6817 t.set_row_stride(48).unwrap(); // padded stride (>= 32 B min)
6818
6819 let map = t.map().expect("strided Mem tensor should CPU-map");
6820 // Full padded buffer (stride 48 × 3 rows = 144 B), not the 96 B logical
6821 // view — callers iterate rows via `effective_row_stride()`.
6822 assert_eq!(map.as_slice().len(), 144);
6823 // Logical shape is still reported for shape-aware consumers.
6824 assert_eq!(map.shape(), &[3, 8, 4]);
6825 }
6826
6827 #[test]
6828 fn strided_mem_tensor_over_capacity_errors() {
6829 // Stride larger than the allocation: 64 B × 3 rows = 192 B > 96 B cap.
6830 let mut t = Tensor::<u8>::new(&[3, 8, 4], Some(TensorMemory::Mem), None).unwrap();
6831 t.set_format(PixelFormat::Rgba).unwrap();
6832 t.set_row_stride(64).unwrap();
6833 assert!(matches!(t.map(), Err(Error::InsufficientCapacity { .. })));
6834 }
6835
6836 /// Test that SHM memory allocation is available and usable on Unix systems.
6837 /// This is a basic functional test; Linux has additional FD leak tests using procfs.
6838 #[test]
6839 #[cfg(unix)]
6840 fn test_shm_available_and_usable() {
6841 assert!(
6842 is_shm_available(),
6843 "SHM memory allocation should be available on Unix systems"
6844 );
6845
6846 // Create a tensor with SHM backing
6847 let tensor = Tensor::<u8>::new(&[100, 100], Some(TensorMemory::Shm), None)
6848 .expect("Failed to create SHM tensor");
6849
6850 // Verify we can map and write to it
6851 let mut map = tensor.map().expect("Failed to map SHM tensor");
6852 map.as_mut_slice().fill(0xAB);
6853
6854 // Verify the data was written correctly
6855 assert!(
6856 map.as_slice().iter().all(|&b| b == 0xAB),
6857 "SHM tensor data should be writable and readable"
6858 );
6859 }
6860
6861 // =========================================================================
6862 // packed_rgba16f_layout — host-runnable geometry unit tests (TDD)
6863 // =========================================================================
6864
6865 #[test]
6866 fn packed_rgba16f_layout_planar_rgb_f16() {
6867 let layout =
6868 packed_rgba16f_layout(PixelFormat::PlanarRgb, DType::F16, 640, 640).expect("Some");
6869 assert_eq!(layout.surface_w, 160);
6870 assert_eq!(layout.surface_h, 1920);
6871 assert_eq!(layout.bytes_per_texel, 8);
6872 assert_eq!(layout.pitch, 1280);
6873 }
6874
6875 #[test]
6876 fn packed_rgba16f_layout_planar_rgba_f16() {
6877 let layout =
6878 packed_rgba16f_layout(PixelFormat::PlanarRgba, DType::F16, 640, 640).expect("Some");
6879 assert_eq!(layout.surface_w, 160);
6880 assert_eq!(layout.surface_h, 2560); // 4 planes
6881 assert_eq!(layout.bytes_per_texel, 8);
6882 assert_eq!(layout.pitch, 1280);
6883 }
6884
6885 #[test]
6886 fn packed_rgba16f_layout_rejects_misaligned() {
6887 assert!(packed_rgba16f_layout(PixelFormat::PlanarRgb, DType::F16, 642, 640).is_none());
6888 }
6889
6890 #[test]
6891 fn packed_rgba16f_layout_rejects_non_f16() {
6892 // Non-F16 dtype with planar RGB
6893 assert!(packed_rgba16f_layout(PixelFormat::PlanarRgb, DType::U8, 640, 640).is_none());
6894 // Non-planar format with F32
6895 assert!(packed_rgba16f_layout(PixelFormat::Rgb, DType::F32, 640, 640).is_none());
6896 // Packed Rgba with F16 is not a planar format → None
6897 assert!(packed_rgba16f_layout(PixelFormat::Rgba, DType::F16, 640, 640).is_none());
6898 }
6899
6900 #[test]
6901 fn cuda_map_fast_fails_to_none_without_handle() {
6902 let t = Tensor::<f32>::new(&[4], Some(TensorMemory::Mem), None).unwrap();
6903 assert!(t.cuda().is_none());
6904 assert!(t.cuda_map().is_none()); // pure local check, no GL routing
6905 }
6906
6907 #[test]
6908 fn cuda_returns_none_without_handle() {
6909 // A plain Mem-backed tensor has no CUDA handle attached.
6910 let t = Tensor::<f32>::new(&[2, 2], Some(TensorMemory::Mem), None).unwrap();
6911 assert!(t.cuda().is_none(), "no CUDA handle on a Mem tensor");
6912 assert!(t.cuda_map().is_none(), "fast-fail map → None");
6913 }
6914
6915 #[test]
6916 fn cuda_map_then_host_map_fallback() {
6917 // The documented client pattern: try cuda_map() first; when it is None
6918 // (no CUDA handle — the case for a plain Mem tensor), fall back to map().
6919 let t = Tensor::<f32>::new(&[2, 2], Some(TensorMemory::Mem), None).unwrap();
6920 // Bind to a named variable so the CudaMap guard (and its borrow of `t`)
6921 // is dropped at the end of this statement, before the else branch borrows `t` again.
6922 let cuda = t.cuda_map();
6923 if let Some(_c) = cuda {
6924 // On a CUDA-registered tensor we'd use the device ptr here.
6925 unreachable!("a Mem tensor has no CUDA handle");
6926 } else {
6927 let host = t.map().expect("host map fallback must succeed");
6928 // TensorMapTrait::len() returns the element count (not bytes).
6929 assert_eq!(host.len(), 4); // 2*2 f32 elements
6930 }
6931 }
6932
6933 // -------------------------------------------------------------------------
6934 // Tensor::from_foreign — public API tests at the Tensor<T> layer.
6935 //
6936 // The low-level MemTensor::from_foreign mechanics (owner-drop, view sharing)
6937 // are covered in mem.rs. These tests exercise the Tensor<T> guard paths
6938 // (null ptr, empty shape, size overflow) and the basic wrap+readback
6939 // contract, confirming the public unsafe API wires through correctly.
6940 // -------------------------------------------------------------------------
6941
6942 #[test]
6943 fn from_foreign_valid_wrap_and_readback() {
6944 // The canonical CUDA zero-copy shape: wrap a caller allocation as a
6945 // Mem tensor and verify the tensor reads the exact same bytes.
6946 let mut buf: Vec<f32> = vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0];
6947 let ptr = buf.as_mut_ptr();
6948 let t = unsafe { Tensor::<f32>::from_foreign(ptr, &[2, 3], None, Some("test_foreign")) }
6949 .expect("valid from_foreign must succeed");
6950 assert_eq!(t.shape(), &[2, 3]);
6951 assert_eq!(t.memory(), TensorMemory::Mem);
6952 assert_eq!(t.name(), "test_foreign");
6953 let m = t.map().unwrap();
6954 // The tensor is a zero-copy borrow — it sees the caller's data.
6955 assert_eq!(m.as_slice(), &[1.0_f32, 2.0, 3.0, 4.0, 5.0, 6.0]);
6956 }
6957
6958 #[test]
6959 fn from_foreign_write_visible_in_caller_allocation() {
6960 // Writes through the tensor's map land in the caller's buffer (zero-copy).
6961 let mut buf: Vec<u8> = vec![0u8; 6];
6962 let ptr = buf.as_mut_ptr();
6963 let t = unsafe { Tensor::<u8>::from_foreign(ptr, &[2, 3], None, None) }.unwrap();
6964 {
6965 let mut m = t.map().unwrap();
6966 m.as_mut_slice().copy_from_slice(&[10, 20, 30, 40, 50, 60]);
6967 }
6968 drop(t);
6969 // Mutations are visible in the original Vec — same physical buffer.
6970 assert_eq!(buf, vec![10, 20, 30, 40, 50, 60]);
6971 }
6972
6973 #[test]
6974 fn from_foreign_rejects_null_ptr() {
6975 let err = unsafe { Tensor::<u8>::from_foreign(std::ptr::null_mut(), &[4], None, None) }
6976 .unwrap_err();
6977 assert!(
6978 matches!(err, Error::InvalidArgument(ref m) if m.contains("non-null")),
6979 "expected InvalidArgument(non-null), got {err:?}"
6980 );
6981 }
6982
6983 #[test]
6984 fn from_foreign_rejects_empty_shape() {
6985 let mut dummy: u8 = 0;
6986 let err = unsafe { Tensor::<u8>::from_foreign(&mut dummy, &[], None, None) }.unwrap_err();
6987 assert!(
6988 matches!(err, Error::InvalidSize(0)),
6989 "expected InvalidSize(0) for empty shape, got {err:?}"
6990 );
6991 }
6992
6993 #[test]
6994 fn from_foreign_rejects_overflow_shape() {
6995 // Two dimensions whose product overflows usize — the overflow guard must
6996 // fire before any pointer arithmetic is attempted.
6997 let mut dummy: u8 = 0;
6998 let huge = [usize::MAX / 2 + 1, 2];
6999 let err = unsafe { Tensor::<u8>::from_foreign(&mut dummy, &huge, None, None) }.unwrap_err();
7000 assert!(
7001 matches!(err, Error::InvalidArgument(ref m) if m.contains("overflow")),
7002 "expected InvalidArgument(overflow), got {err:?}"
7003 );
7004 }
7005
7006 #[test]
7007 fn from_foreign_owner_keeps_allocation_alive() {
7008 // When `owner` is `Some`, dropping the Tensor must not free the backing
7009 // before the owner is also gone — the owner's Drop fires on last ref.
7010 use std::sync::atomic::{AtomicBool, Ordering};
7011 let flag = std::sync::Arc::new(AtomicBool::new(false));
7012 let flag2 = flag.clone();
7013 struct Guard(std::sync::Arc<AtomicBool>);
7014 impl Drop for Guard {
7015 fn drop(&mut self) {
7016 self.0.store(true, Ordering::SeqCst);
7017 }
7018 }
7019 let mut buf: Vec<u32> = vec![42u32; 4];
7020 let ptr = buf.as_mut_ptr();
7021 let owner: ForeignOwner = Box::new(Guard(flag2));
7022 let t = unsafe { Tensor::<u32>::from_foreign(ptr, &[4], Some(owner), None) }.unwrap();
7023 // Map co-owns the backing Arc; the owner must stay alive while the map lives.
7024 let m = t.map().unwrap();
7025 assert_eq!(m.as_slice()[0], 42);
7026 drop(t); // tensor dropped while map is still live
7027 assert!(
7028 !flag.load(Ordering::SeqCst),
7029 "owner must not drop while a map shares the backing"
7030 );
7031 drop(m);
7032 assert!(
7033 flag.load(Ordering::SeqCst),
7034 "owner Drop must fire when the last Arc reference is released"
7035 );
7036 }
7037}