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 /// # Errors
3853 ///
3854 /// - [`Error::RegionOutOfBounds`] if `region` exceeds the image bounds.
3855 /// - [`Error::InvalidOperation`] if the tensor is not a packed-format image
3856 /// (planar/semi-planar spatial sub-rects are not a single strided window;
3857 /// use [`batch`](Self::batch) for batched planar tensors).
3858 pub fn view(&self, region: Region) -> Result<Tensor<T>> {
3859 let fmt = self.format.ok_or_else(|| {
3860 Error::InvalidOperation("view() requires a formatted image tensor".into())
3861 })?;
3862 if fmt.layout() != PixelLayout::Packed {
3863 return Err(Error::InvalidOperation(format!(
3864 "view() supports packed formats only (got {fmt:?}); use batch(n) for batched \
3865 planar tensors"
3866 )));
3867 }
3868 let w = self
3869 .width()
3870 .ok_or_else(|| Error::InvalidOperation("view(): tensor has no image width".into()))?;
3871 let h = self
3872 .height()
3873 .ok_or_else(|| Error::InvalidOperation("view(): tensor has no image height".into()))?;
3874 if !region.fits_within(w, h) {
3875 return Err(Error::RegionOutOfBounds {
3876 region,
3877 bounds: (w, h),
3878 });
3879 }
3880 let elem = std::mem::size_of::<T>();
3881 let bpp = fmt.channels() * elem;
3882 let stride = self.effective_row_stride().unwrap_or(w * bpp);
3883 let offset = region
3884 .y
3885 .checked_mul(stride)
3886 .and_then(|yo| yo.checked_add(region.x.checked_mul(bpp)?))
3887 .ok_or(Error::InvalidSize(region.y))?;
3888 let sub_shape = fmt
3889 .image_shape(region.width, region.height)
3890 .ok_or_else(|| Error::InvalidShape(format!("view(): invalid shape for {fmt:?}")))?;
3891 let mut t = self.subview(offset, &sub_shape)?;
3892 // A multi-row sub-rect must advance rows by the PARENT pitch so each row
3893 // addresses the correct columns. A single-row view uses its own tight
3894 // stride — the parent pitch would make the strided `map()` (which exposes
3895 // `stride × rows`) expose a trailing row that runs past the buffer tail
3896 // for an offset (x>0 / bottom) view. The GL backend does NOT rely on this
3897 // (single-row-tight) `row_stride`: it reads the parent pitch from
3898 // `view_origin.parent_row_stride` so its import/cache pitch stays
3899 // parent-consistent for views of any height — see `ViewOrigin`.
3900 let view_stride = if region.height > 1 {
3901 stride
3902 } else {
3903 region.width * bpp
3904 };
3905 t.set_row_stride_unchecked(view_stride);
3906 // Snapshot the parent `(w, h, row_stride)` so the GL backend imports the
3907 // parent once (keyed on the parent pitch, not this view's possibly-tight
3908 // single-row stride) and renders this sub-rect as a `glViewport`/
3909 // `glScissor` ROI at `(region.x, region.y)`. Composes when viewing an
3910 // existing view.
3911 t.view_origin = Some(self.compose_view_origin(w, h, stride, region.x, region.y));
3912 Ok(t)
3913 }
3914
3915 /// Build the [`ViewOrigin`] for a new sub-region of `self`. When `self` is a
3916 /// whole tensor the snapshot names `self` as the parent; when `self` is
3917 /// already a view, the snapshot keeps the **root** parent and accumulates the
3918 /// local origin so nested views still resolve to one import.
3919 fn compose_view_origin(
3920 &self,
3921 parent_width: usize,
3922 parent_height: usize,
3923 parent_row_stride: usize,
3924 x: usize,
3925 y: usize,
3926 ) -> ViewOrigin {
3927 match self.view_origin {
3928 Some(root) => ViewOrigin {
3929 parent_width: root.parent_width,
3930 parent_height: root.parent_height,
3931 parent_row_stride: root.parent_row_stride,
3932 x: root.x.saturating_add(x),
3933 y: root.y.saturating_add(y),
3934 },
3935 None => ViewOrigin {
3936 parent_width,
3937 parent_height,
3938 parent_row_stride,
3939 x,
3940 y,
3941 },
3942 }
3943 }
3944
3945 /// Downcast to PBO tensor reference (for GL backends).
3946 pub fn as_pbo(&self) -> Option<&PboTensor<T>> {
3947 match &self.storage {
3948 TensorStorage::Pbo(p) => Some(p),
3949 _ => None,
3950 }
3951 }
3952
3953 /// Downcast to DMA tensor reference (for EGL import, G2D).
3954 #[cfg(target_os = "linux")]
3955 pub fn as_dma(&self) -> Option<&DmaTensor<T>> {
3956 match &self.storage {
3957 TensorStorage::Dma(d) => Some(d),
3958 _ => None,
3959 }
3960 }
3961
3962 /// Borrow the DMA-BUF file descriptor backing this tensor.
3963 ///
3964 /// # Returns
3965 ///
3966 /// A borrowed reference to the DMA-BUF file descriptor, tied to `self`'s
3967 /// lifetime.
3968 ///
3969 /// # Errors
3970 ///
3971 /// Returns `Error::NotImplemented` if the tensor is not DMA-backed.
3972 #[cfg(target_os = "linux")]
3973 pub fn dmabuf(&self) -> Result<std::os::fd::BorrowedFd<'_>> {
3974 use std::os::fd::AsFd;
3975 match &self.storage {
3976 TensorStorage::Dma(dma) => Ok(dma.fd.as_fd()),
3977 _ => Err(Error::NotImplemented(format!(
3978 "dmabuf requires DMA-backed tensor, got {:?}",
3979 self.storage.memory()
3980 ))),
3981 }
3982 }
3983
3984 /// Construct a Tensor from a PBO tensor (for GL backends that allocate PBOs).
3985 pub fn from_pbo(pbo: PboTensor<T>) -> Self {
3986 Self {
3987 storage: TensorStorage::Pbo(pbo),
3988 format: None,
3989 chroma: None,
3990 row_stride: None,
3991 plane_offset: None,
3992 quantization: None,
3993 cuda: None,
3994 colorimetry: None,
3995 cpu_access: CpuAccess::ReadWrite,
3996 compression: None,
3997 view_origin: None,
3998 }
3999 }
4000
4001 /// The CUDA registration for this tensor, if any (set at creation on CUDA devices).
4002 pub fn cuda(&self) -> Option<&crate::cuda::CudaHandle> {
4003 self.cuda.as_ref()
4004 }
4005
4006 /// Attach a CUDA handle (called by ImageProcessor::create_image after registering a PBO).
4007 pub fn set_cuda_handle(&mut self, h: crate::cuda::CudaHandle) {
4008 self.cuda = Some(h);
4009 }
4010
4011 /// Fast-fail CUDA map: None (no GL routing) when no handle; else map (PBO routes to the GL worker).
4012 ///
4013 /// Returns a scoped [`CudaMap`] guard holding the raw CUDA device pointer
4014 /// for the duration of the mapping. For GL-buffer-backed tensors the unmap is deferred until the
4015 /// guard drops, freeing the PBO for the next `convert()` call. When no CUDA handle is attached
4016 /// (the common case for plain `Mem`/`DMA` tensors without CUDA registration), returns `None`
4017 /// immediately — no GL routing, no allocation.
4018 ///
4019 /// # Example — zero-copy CUDA input with host fallback
4020 ///
4021 /// ```no_run
4022 /// use edgefirst_tensor::{Tensor, TensorMemory, TensorTrait};
4023 /// # fn feed_tensorrt(_dptr: *mut std::ffi::c_void, _bytes: usize) {}
4024 /// # fn demo(t: &Tensor<f32>) {
4025 /// // Try the zero-copy CUDA device pointer first.
4026 /// if let Some(cuda) = t.cuda_map() {
4027 /// feed_tensorrt(cuda.device_ptr(), cuda.len());
4028 /// // `cuda` (a CudaMap guard) unmaps when it goes out of scope, freeing
4029 /// // the GPU buffer for the next convert().
4030 /// } else {
4031 /// // Fall back to the host mapping when no CUDA handle is attached.
4032 /// let _host = t.map().expect("host map fallback must succeed");
4033 /// // `_host` is a TensorMap<f32> that derefs to &[f32].
4034 /// }
4035 /// # }
4036 /// ```
4037 pub fn cuda_map(&self) -> Option<crate::cuda::CudaMap<'_>> {
4038 self.cuda.as_ref()?.map()
4039 }
4040
4041 /// Attempt to attach a CUDA `ExternalMemory` handle for DMA-backed tensors.
4042 ///
4043 /// On a CUDA-capable host, imports the DMA-BUF fd via
4044 /// `cudaImportExternalMemory(OpaqueFd)` and maps it to a device pointer.
4045 /// Sets `self.cuda` to a persistent `ExternalMem` handle on success. No-op
4046 /// if CUDA is unavailable, the tensor is not DMA-backed, or a handle is
4047 /// already set. Import failure is silently ignored — the tensor remains
4048 /// usable without a CUDA handle.
4049 ///
4050 /// # RUNTIME-UNVALIDATED
4051 ///
4052 /// No test platform has both `/dev/dma_heap` and a CUDA device. ABI is
4053 /// layout-asserted vs. CUDA 12.6 `driver_types.h`; the mechanism is proven
4054 /// by gpu-probe O5 on Orin. Best-effort: tensor creation never fails here.
4055 #[cfg(target_os = "linux")]
4056 pub fn try_init_dma_cuda(&mut self) {
4057 // Fast-path: already imported, CUDA not available, or not a DMA tensor.
4058 if self.cuda.is_some() || !crate::cuda::is_cuda_available() {
4059 return;
4060 }
4061 let (raw_fd, buf_size) = match &self.storage {
4062 TensorStorage::Dma(dma) => {
4063 use std::os::fd::AsRawFd;
4064 (dma.fd.as_raw_fd(), dma.buf_size)
4065 }
4066 _ => return,
4067 };
4068 if let Some((ext, dptr)) = crate::cuda::import_dma_fd(raw_fd, buf_size) {
4069 self.cuda = Some(crate::cuda::CudaHandle::new_external(ext, dptr, buf_size));
4070 }
4071 }
4072}
4073
4074// Quantization accessors — type-gated to integer element types via the
4075// sealed `IntegerType` trait. Calling `.quantization()` on a `Tensor<f32>`
4076// produces a compile error, not a runtime one.
4077impl<T> Tensor<T>
4078where
4079 T: IntegerType + Num + Clone + fmt::Debug + Send + Sync,
4080{
4081 /// Quantization metadata for this tensor, if set.
4082 pub fn quantization(&self) -> Option<&Quantization> {
4083 self.quantization.as_ref()
4084 }
4085
4086 /// Attach quantization metadata to this tensor. Validates against the
4087 /// tensor's shape — returns [`Error::QuantizationInvalid`] on any
4088 /// inconsistency (mismatched scale/zp lengths, out-of-range axis, etc.).
4089 pub fn set_quantization(&mut self, q: Quantization) -> Result<()> {
4090 q.validate(self.shape())?;
4091 self.quantization = Some(q);
4092 Ok(())
4093 }
4094
4095 /// Builder-style variant of [`Self::set_quantization`]. Consumes `self`
4096 /// and returns `Result<Self>` — on success yields the tensor with the
4097 /// attached quantization; on validation failure returns
4098 /// [`Error::QuantizationInvalid`] and drops `self` (the tensor is not
4099 /// returned in the error arm).
4100 pub fn with_quantization(mut self, q: Quantization) -> Result<Self> {
4101 self.set_quantization(q)?;
4102 Ok(self)
4103 }
4104
4105 /// Clear any quantization metadata on this tensor.
4106 pub fn clear_quantization(&mut self) {
4107 self.quantization = None;
4108 }
4109}
4110
4111impl<T> TensorTrait<T> for Tensor<T>
4112where
4113 T: Num + Clone + fmt::Debug + Send + Sync,
4114{
4115 fn new(shape: &[usize], name: Option<&str>) -> Result<Self>
4116 where
4117 Self: Sized,
4118 {
4119 Self::new(shape, None, name)
4120 }
4121
4122 #[cfg(unix)]
4123 fn from_fd(fd: std::os::fd::OwnedFd, shape: &[usize], name: Option<&str>) -> Result<Self>
4124 where
4125 Self: Sized,
4126 {
4127 #[cfg_attr(not(target_os = "linux"), allow(unused_mut))]
4128 let mut t = Self::wrap(TensorStorage::from_fd(fd, shape, name)?);
4129 // Best-effort CUDA external memory import for DMA-backed tensors.
4130 // RUNTIME-UNVALIDATED: see try_init_dma_cuda().
4131 #[cfg(target_os = "linux")]
4132 t.try_init_dma_cuda();
4133 Ok(t)
4134 }
4135
4136 #[cfg(unix)]
4137 fn clone_fd(&self) -> Result<std::os::fd::OwnedFd> {
4138 self.storage.clone_fd()
4139 }
4140
4141 fn memory(&self) -> TensorMemory {
4142 self.storage.memory()
4143 }
4144
4145 fn name(&self) -> String {
4146 self.storage.name()
4147 }
4148
4149 fn shape(&self) -> &[usize] {
4150 self.storage.shape()
4151 }
4152
4153 fn reshape(&mut self, shape: &[usize]) -> Result<()> {
4154 if self.chroma.is_some() {
4155 return Err(Error::InvalidOperation(
4156 "cannot reshape a multiplane tensor — decompose planes first".into(),
4157 ));
4158 }
4159 self.storage.reshape(shape)?;
4160 self.format = None;
4161 self.row_stride = None;
4162 self.plane_offset = None;
4163 match self.storage {
4164 TensorStorage::Mem(ref mut m) => m.set_offset(0),
4165 #[cfg(target_os = "linux")]
4166 TensorStorage::Dma(ref mut dma) => dma.mmap_offset = 0,
4167 _ => {}
4168 }
4169 Ok(())
4170 }
4171
4172 fn map_with(&self, access: CpuAccess) -> Result<TensorMap<T>> {
4173 let _span = tracing::trace_span!(
4174 "tensor.map",
4175 memory = ?self.storage.memory(),
4176 ?access,
4177 )
4178 .entered();
4179 if access == CpuAccess::None {
4180 return Err(Error::InvalidArgument(
4181 "map_with(CpuAccess::None) is not a mappable direction — use \
4182 map_read()/map_write()/map_mut()"
4183 .into(),
4184 ));
4185 }
4186 // Declared-vs-requested telemetry (all platforms): mapping beyond
4187 // the allocation-time declaration is best-effort — tolerated where
4188 // the backing is CPU-mappable regardless (Mem/Shm/dma-buf/
4189 // IOSurface), refused by the Android backend for CpuAccess::None
4190 // buffers — but always loud and counted, never silent.
4191 if !self.cpu_access.covers(access) {
4192 note_unplanned_cpu_access(
4193 self.buffer_identity().id(),
4194 &format!("{:?}", self.storage.memory()),
4195 "map access exceeds the declared CpuAccess",
4196 );
4197 }
4198 // CPU mapping of a strided tensor exposes the full padded buffer
4199 // (`row_stride × rows`) so callers can iterate rows via
4200 // `effective_row_stride()` without running past the slice. This is sound
4201 // only when the HAL owns and can size-check the allocation:
4202 //
4203 // * Self-allocated Mem / Shm tensors (any platform) — the backing
4204 // `Vec` / shm segment is sized by `capacity_bytes()`, checked here.
4205 // * Self-allocated DMA tensors (Linux) — pitch padding from
4206 // `image_with_stride()`; checked against the DMA-BUF `buf_size`.
4207 //
4208 // * Self-allocated PBO tensors (any platform with GL) — the GL buffer
4209 // is sized by `capacity_bytes()` and may carry 64-byte row padding;
4210 // the JPEG decoder mmaps it and convert() reads it, both iterating
4211 // by `row_stride`. Checked against the PBO capacity below.
4212 //
4213 // Foreign DMA-BUFs (`from_fd()` + `set_row_stride()`, the V4L2 /
4214 // GStreamer case) and IOSurface are rejected: their layout comes from an
4215 // external allocator / GPU driver the HAL cannot validate for a strided
4216 // CPU view, and they are intended for the GPU path. (Earlier this
4217 // rejected *all* non-Linux strided maps with "DMA backing is Linux-only"
4218 // — that was an unimplemented path, not a platform limit; HAL-owned
4219 // Mem/Shm/PBO are trivially mappable and now are.)
4220 if let Some(stride) = self.row_stride {
4221 // Rows sit at `stride`-byte spacing. The row count is the first
4222 // shape dim for packed `[H, W, C]` and semi-planar `[H*k, W]`,
4223 // but planar `[C, H, W]` stacks C planes of H rows — its surface
4224 // row count is `C × H` (`shape[0]` alone would expose a 3-row
4225 // window and truncate the map; first hit by Android planar-F16
4226 // AHardwareBuffers, whose gralloc pads the pitch — macOS/Linux
4227 // planar pitches happen to be naturally aligned so no stride was
4228 // ever recorded there).
4229 let rows = match self.format.map(|f| f.layout()) {
4230 Some(PixelLayout::Planar) => {
4231 let s = self.shape();
4232 if s.len() < 2 {
4233 return Err(Error::InvalidOperation(
4234 "Tensor::map: strided planar mapping requires [C, H, W] shape".into(),
4235 ));
4236 }
4237 s[0].checked_mul(s[1]).ok_or_else(|| {
4238 Error::InvalidOperation(format!(
4239 "Tensor::map: planar rows {} × {} overflows usize",
4240 s[0], s[1]
4241 ))
4242 })?
4243 }
4244 _ => *self.shape().first().ok_or_else(|| {
4245 Error::InvalidOperation(
4246 "Tensor::map: strided mapping requires a non-empty shape".into(),
4247 )
4248 })?,
4249 };
4250 let total_bytes = stride.checked_mul(rows).ok_or_else(|| {
4251 Error::InvalidOperation(format!(
4252 "Tensor::map: row_stride {stride} × rows {rows} overflows usize"
4253 ))
4254 })?;
4255
4256 match &self.storage {
4257 #[cfg(target_os = "linux")]
4258 TensorStorage::Dma(dma) if !dma.is_imported => {
4259 // `set_row_stride()` only validates `stride >= min_stride`,
4260 // not that `stride × rows` fits the DMA-BUF, so re-check
4261 // here — mapping past `buf_size` would SIGBUS on access.
4262 let available_bytes = dma.buf_size.saturating_sub(dma.mmap_offset);
4263 if total_bytes > available_bytes {
4264 return Err(Error::InvalidOperation(format!(
4265 "Tensor::map: strided mapping needs {total_bytes} bytes \
4266 but DMA buffer only has {available_bytes} available \
4267 (buf_size={}, mmap_offset={}, stride={stride}, rows={rows}); \
4268 the row_stride was likely set larger than the original allocation",
4269 dma.buf_size, dma.mmap_offset
4270 )));
4271 }
4272 return dma
4273 .map_with_byte_size(total_bytes, access)
4274 .map(TensorMap::Dma);
4275 }
4276 TensorStorage::Mem(mem) => {
4277 let capacity = self.storage.capacity_bytes();
4278 if total_bytes > capacity {
4279 return Err(Error::InsufficientCapacity {
4280 needed: total_bytes,
4281 capacity,
4282 });
4283 }
4284 return mem.map_with_byte_size(total_bytes, access);
4285 }
4286 #[cfg(unix)]
4287 TensorStorage::Shm(shm) => {
4288 let capacity = self.storage.capacity_bytes();
4289 if total_bytes > capacity {
4290 return Err(Error::InsufficientCapacity {
4291 needed: total_bytes,
4292 capacity,
4293 });
4294 }
4295 return shm.map_with_byte_size(total_bytes, access);
4296 }
4297 // macOS/iOS: `TensorStorage::Dma` is the IOSurface. The lock yields
4298 // the full surface base address, and the row pitch
4299 // (`IOSurfaceGetBytesPerRow`) is known from the API for both
4300 // self-allocated and imported surfaces — unlike a foreign
4301 // DMA-BUF — so a strided CPU view is sound and zero-copy.
4302 #[cfg(any(target_os = "macos", target_os = "ios"))]
4303 TensorStorage::Dma(io) => {
4304 // A sub-view's window is `buf_size − view_offset`; the strided
4305 // span must fit the window, not the whole surface.
4306 let available = io.buf_size.saturating_sub(io.view_offset);
4307 if total_bytes > available {
4308 return Err(Error::InsufficientCapacity {
4309 needed: total_bytes,
4310 capacity: available,
4311 });
4312 }
4313 return io.map_with_byte_size(total_bytes, access);
4314 }
4315 // Android: `TensorStorage::Dma` is the AHardwareBuffer. The lock
4316 // yields the full buffer base address, and the row pitch is
4317 // known from the allocator-filled descriptor — so a strided CPU
4318 // view is sound and zero-copy, same as IOSurface.
4319 #[cfg(target_os = "android")]
4320 TensorStorage::Dma(ahb) => {
4321 // A sub-view's window is `buf_size − view_offset`; the strided
4322 // span must fit the window, not the whole buffer.
4323 let available = ahb.buf_size.saturating_sub(ahb.view_offset);
4324 if total_bytes > available {
4325 return Err(Error::InsufficientCapacity {
4326 needed: total_bytes,
4327 capacity: available,
4328 });
4329 }
4330 return ahb.map_with_byte_size(total_bytes, access);
4331 }
4332 TensorStorage::Pbo(pbo) => {
4333 // PBO: the GPU-side allocation may have a padded row stride
4334 // (e.g. 64-byte aligned). Expose the full padded buffer so a
4335 // CPU producer (JPEG decoder) and a strided convert source
4336 // can iterate rows via `effective_row_stride()` without
4337 // running past the slice — the logical `pbo.map()` view would
4338 // stop after `shape.product()` and lose bytes past row 0.
4339 // A sub-view's window is `capacity − view_offset`.
4340 let available = pbo.capacity_bytes().saturating_sub(pbo.view_offset);
4341 if total_bytes > available {
4342 return Err(Error::InsufficientCapacity {
4343 needed: total_bytes,
4344 capacity: available,
4345 });
4346 }
4347 return pbo.map_with_byte_size(total_bytes, access);
4348 }
4349 // Reachable on Linux for an IMPORTED DMA-BUF (the `Dma` arm above
4350 // is guarded `if !dma.is_imported`). On macOS/Windows every
4351 // storage variant is matched explicitly, so this catch-all is
4352 // unreachable there — allow it rather than cfg-gating per platform.
4353 #[allow(unreachable_patterns)]
4354 _ => {
4355 return Err(Error::InvalidOperation(
4356 "CPU mapping of strided tensors is supported only for HAL-allocated \
4357 Mem/Shm (any platform), self-allocated DMA (Linux), IOSurface \
4358 (macOS), and PBO; imported DMA-BUF without self-allocation is \
4359 GPU-path only"
4360 .into(),
4361 ));
4362 }
4363 }
4364 }
4365 // Offset tensors are supported for storages that apply the offset
4366 // inside their own `map()`: DMA (`DmaMap`/IOSurface adjust the mapped
4367 // base), Mem (`MemMap` adjusts the slice base), Shm (`ShmMap` adjusts
4368 // the slice base), and PBO (the staged copy starts at the offset). Every
4369 // self-allocated backing now carries a sub-region concept via `view`, so
4370 // a non-zero offset is honoured rather than rejected.
4371 if self.plane_offset.is_some_and(|o| o > 0) {
4372 let supported = matches!(self.storage, TensorStorage::Mem(_) | TensorStorage::Pbo(_));
4373 // macOS `Dma` is the IOSurface; Linux `Dma` is the DMA-BUF; Android
4374 // `Dma` is the AHardwareBuffer — all apply the offset in their map.
4375 // (`Dma` is the same variant name on each, hence one `cfg(any(...))`
4376 // arm rather than three.)
4377 #[cfg(any(
4378 target_os = "linux",
4379 target_os = "macos",
4380 target_os = "ios",
4381 target_os = "android"
4382 ))]
4383 let supported = supported || matches!(self.storage, TensorStorage::Dma(_));
4384 #[cfg(unix)]
4385 let supported = supported || matches!(self.storage, TensorStorage::Shm(_));
4386 if !supported {
4387 return Err(Error::InvalidOperation(
4388 "plane offset only supported for DMA, Mem, Shm, and PBO tensors".into(),
4389 ));
4390 }
4391 }
4392 self.storage.map_with(access)
4393 }
4394
4395 fn buffer_identity(&self) -> &BufferIdentity {
4396 self.storage.buffer_identity()
4397 }
4398}
4399
4400pub enum TensorMap<T>
4401where
4402 T: Num + Clone + fmt::Debug,
4403{
4404 #[cfg(target_os = "linux")]
4405 Dma(DmaMap<T>),
4406 #[cfg(any(target_os = "macos", target_os = "ios"))]
4407 IoSurface(IoSurfaceMap<T>),
4408 #[cfg(target_os = "android")]
4409 HardwareBuffer(AHardwareBufferMap<T>),
4410 #[cfg(unix)]
4411 Shm(ShmMap<T>),
4412 Mem(MemMap<T>),
4413 Pbo(PboMap<T>),
4414}
4415
4416impl<T> TensorMapTrait<T> for TensorMap<T>
4417where
4418 T: Num + Clone + fmt::Debug,
4419{
4420 fn shape(&self) -> &[usize] {
4421 match self {
4422 #[cfg(target_os = "linux")]
4423 TensorMap::Dma(map) => map.shape(),
4424 #[cfg(any(target_os = "macos", target_os = "ios"))]
4425 TensorMap::IoSurface(map) => map.shape(),
4426 #[cfg(target_os = "android")]
4427 TensorMap::HardwareBuffer(map) => map.shape(),
4428 #[cfg(unix)]
4429 TensorMap::Shm(map) => map.shape(),
4430 TensorMap::Mem(map) => map.shape(),
4431 TensorMap::Pbo(map) => map.shape(),
4432 }
4433 }
4434
4435 fn unmap(&mut self) {
4436 match self {
4437 #[cfg(target_os = "linux")]
4438 TensorMap::Dma(map) => map.unmap(),
4439 #[cfg(any(target_os = "macos", target_os = "ios"))]
4440 TensorMap::IoSurface(map) => map.unmap(),
4441 #[cfg(target_os = "android")]
4442 TensorMap::HardwareBuffer(map) => map.unmap(),
4443 #[cfg(unix)]
4444 TensorMap::Shm(map) => map.unmap(),
4445 TensorMap::Mem(map) => map.unmap(),
4446 TensorMap::Pbo(map) => map.unmap(),
4447 }
4448 }
4449
4450 fn as_slice(&self) -> &[T] {
4451 match self {
4452 #[cfg(target_os = "linux")]
4453 TensorMap::Dma(map) => map.as_slice(),
4454 #[cfg(any(target_os = "macos", target_os = "ios"))]
4455 TensorMap::IoSurface(map) => map.deref(),
4456 #[cfg(target_os = "android")]
4457 TensorMap::HardwareBuffer(map) => map.deref(),
4458 #[cfg(unix)]
4459 TensorMap::Shm(map) => map.as_slice(),
4460 TensorMap::Mem(map) => map.as_slice(),
4461 TensorMap::Pbo(map) => map.as_slice(),
4462 }
4463 }
4464
4465 fn as_mut_slice(&mut self) -> &mut [T] {
4466 match self {
4467 #[cfg(target_os = "linux")]
4468 TensorMap::Dma(map) => map.as_mut_slice(),
4469 #[cfg(any(target_os = "macos", target_os = "ios"))]
4470 TensorMap::IoSurface(map) => map.deref_mut(),
4471 #[cfg(target_os = "android")]
4472 TensorMap::HardwareBuffer(map) => map.deref_mut(),
4473 #[cfg(unix)]
4474 TensorMap::Shm(map) => map.as_mut_slice(),
4475 TensorMap::Mem(map) => map.as_mut_slice(),
4476 TensorMap::Pbo(map) => map.as_mut_slice(),
4477 }
4478 }
4479}
4480
4481impl<T> Deref for TensorMap<T>
4482where
4483 T: Num + Clone + fmt::Debug,
4484{
4485 type Target = [T];
4486
4487 fn deref(&self) -> &[T] {
4488 match self {
4489 #[cfg(target_os = "linux")]
4490 TensorMap::Dma(map) => map.deref(),
4491 #[cfg(any(target_os = "macos", target_os = "ios"))]
4492 TensorMap::IoSurface(map) => map.deref(),
4493 #[cfg(target_os = "android")]
4494 TensorMap::HardwareBuffer(map) => map.deref(),
4495 #[cfg(unix)]
4496 TensorMap::Shm(map) => map.deref(),
4497 TensorMap::Mem(map) => map.deref(),
4498 TensorMap::Pbo(map) => map.deref(),
4499 }
4500 }
4501}
4502
4503impl<T> DerefMut for TensorMap<T>
4504where
4505 T: Num + Clone + fmt::Debug,
4506{
4507 fn deref_mut(&mut self) -> &mut [T] {
4508 match self {
4509 #[cfg(target_os = "linux")]
4510 TensorMap::Dma(map) => map.deref_mut(),
4511 #[cfg(any(target_os = "macos", target_os = "ios"))]
4512 TensorMap::IoSurface(map) => map.deref_mut(),
4513 #[cfg(target_os = "android")]
4514 TensorMap::HardwareBuffer(map) => map.deref_mut(),
4515 #[cfg(unix)]
4516 TensorMap::Shm(map) => map.deref_mut(),
4517 TensorMap::Mem(map) => map.deref_mut(),
4518 TensorMap::Pbo(map) => map.deref_mut(),
4519 }
4520 }
4521}
4522
4523// ============================================================================
4524// Platform availability helpers
4525// ============================================================================
4526
4527/// Cached result of the Linux DMA-BUF availability probe.
4528#[cfg(target_os = "linux")]
4529static DMA_AVAILABLE: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
4530/// Cached result of the macOS/iOS IOSurface availability probe.
4531#[cfg(any(target_os = "macos", target_os = "ios"))]
4532static IOSURFACE_AVAILABLE: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
4533
4534/// Check if Linux DMA-BUF allocation is available on this system.
4535///
4536/// Linux-specific availability check (typically requires `/dev/dma_heap`
4537/// access — running as root or membership in a video/render group). For
4538/// portable code that wants "any zero-copy GPU buffer", use
4539/// [`is_gpu_buffer_available`] which also covers IOSurface on macOS.
4540///
4541/// This function caches its result after the first call.
4542#[cfg(target_os = "linux")]
4543pub fn is_dma_available() -> bool {
4544 *DMA_AVAILABLE.get_or_init(|| Tensor::<u8>::new(&[64], Some(TensorMemory::Dma), None).is_ok())
4545}
4546
4547/// Always returns `false` on non-Linux platforms.
4548#[cfg(not(target_os = "linux"))]
4549pub fn is_dma_available() -> bool {
4550 false
4551}
4552
4553/// Check if macOS/iOS IOSurface allocation is available on this system.
4554///
4555/// IOSurface is part of the macOS/iOS OS and is essentially always present;
4556/// this probe catches degraded scenarios such as memory pressure or
4557/// sandboxed contexts where `IOSurfaceCreate` fails. The result is
4558/// cached after the first call.
4559///
4560/// Always returns `false` on non-Apple platforms.
4561#[cfg(any(target_os = "macos", target_os = "ios"))]
4562pub fn is_iosurface_available() -> bool {
4563 *IOSURFACE_AVAILABLE.get_or_init(|| {
4564 // Probe via the same Dma path — on macOS/iOS this routes through
4565 // IoSurfaceTensor::new.
4566 Tensor::<u8>::new(&[64], Some(TensorMemory::Dma), None).is_ok()
4567 })
4568}
4569
4570#[cfg(not(any(target_os = "macos", target_os = "ios")))]
4571pub fn is_iosurface_available() -> bool {
4572 false
4573}
4574
4575/// Cached result of the Android AHardwareBuffer availability probe.
4576#[cfg(target_os = "android")]
4577static AHARDWAREBUFFER_AVAILABLE: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
4578
4579/// Check if Android AHardwareBuffer allocation is available on this system.
4580///
4581/// AHardwareBuffer is part of the Android OS (public NDK ABI since API
4582/// 26) and is essentially always present; this probe catches degraded
4583/// scenarios such as memory pressure or gralloc failures. The result is
4584/// cached after the first call.
4585#[cfg(target_os = "android")]
4586pub fn is_ahardwarebuffer_available() -> bool {
4587 *AHARDWAREBUFFER_AVAILABLE.get_or_init(|| {
4588 // Probe via the same Dma path — on Android this routes through
4589 // AHardwareBufferTensor::new.
4590 Tensor::<u8>::new(&[64], Some(TensorMemory::Dma), None).is_ok()
4591 })
4592}
4593
4594/// Always returns `false` on non-Android platforms.
4595#[cfg(not(target_os = "android"))]
4596pub fn is_ahardwarebuffer_available() -> bool {
4597 false
4598}
4599
4600/// Portable probe for the platform's native zero-copy GPU buffer
4601/// allocator (DMA-BUF on Linux, IOSurface on macOS/iOS, AHardwareBuffer on
4602/// Android). Returns `false` on
4603/// Windows and other platforms with no equivalent. Use this when writing
4604/// cross-platform code that cares whether the `Dma` tensor variant will
4605/// work, not which underlying mechanism is used.
4606pub fn is_gpu_buffer_available() -> bool {
4607 #[cfg(target_os = "linux")]
4608 {
4609 is_dma_available()
4610 }
4611 #[cfg(any(target_os = "macos", target_os = "ios"))]
4612 {
4613 is_iosurface_available()
4614 }
4615 #[cfg(target_os = "android")]
4616 {
4617 is_ahardwarebuffer_available()
4618 }
4619 #[cfg(not(any(
4620 target_os = "linux",
4621 target_os = "macos",
4622 target_os = "ios",
4623 target_os = "android"
4624 )))]
4625 {
4626 false
4627 }
4628}
4629
4630/// Check if POSIX shared memory allocation is available on this system.
4631///
4632/// Returns `true` on Unix systems (Linux, macOS, BSD) where POSIX shared memory
4633/// is supported. Always returns `false` on non-Unix platforms (Windows).
4634///
4635/// This function caches its result after the first call for efficiency.
4636#[cfg(unix)]
4637static SHM_AVAILABLE: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
4638
4639/// Check if POSIX shared memory allocation is available on this system.
4640#[cfg(unix)]
4641pub fn is_shm_available() -> bool {
4642 *SHM_AVAILABLE.get_or_init(|| Tensor::<u8>::new(&[64], Some(TensorMemory::Shm), None).is_ok())
4643}
4644
4645/// Check if POSIX shared memory allocation is available on this system.
4646///
4647/// Always returns `false` on non-Unix platforms since POSIX SHM is Unix-specific.
4648#[cfg(not(unix))]
4649pub fn is_shm_available() -> bool {
4650 false
4651}
4652
4653#[cfg(test)]
4654mod dtype_tests {
4655 use super::*;
4656
4657 #[test]
4658 fn dtype_size() {
4659 assert_eq!(DType::U8.size(), 1);
4660 assert_eq!(DType::I8.size(), 1);
4661 assert_eq!(DType::U16.size(), 2);
4662 assert_eq!(DType::I16.size(), 2);
4663 assert_eq!(DType::U32.size(), 4);
4664 assert_eq!(DType::I32.size(), 4);
4665 assert_eq!(DType::U64.size(), 8);
4666 assert_eq!(DType::I64.size(), 8);
4667 assert_eq!(DType::F16.size(), 2);
4668 assert_eq!(DType::F32.size(), 4);
4669 assert_eq!(DType::F64.size(), 8);
4670 }
4671
4672 #[test]
4673 fn dtype_name() {
4674 assert_eq!(DType::U8.name(), "u8");
4675 assert_eq!(DType::F16.name(), "f16");
4676 assert_eq!(DType::F32.name(), "f32");
4677 }
4678
4679 #[test]
4680 fn dtype_serde_roundtrip() {
4681 use serde_json;
4682 let dt = DType::F16;
4683 let json = serde_json::to_string(&dt).unwrap();
4684 let back: DType = serde_json::from_str(&json).unwrap();
4685 assert_eq!(dt, back);
4686 }
4687}
4688
4689#[cfg(test)]
4690mod image_tests {
4691 use super::*;
4692
4693 #[test]
4694 fn image_shape_per_layout() {
4695 assert_eq!(
4696 PixelFormat::Rgb.image_shape(640, 480),
4697 Some(vec![480, 640, 3])
4698 );
4699 assert_eq!(
4700 PixelFormat::Grey.image_shape(640, 480),
4701 Some(vec![480, 640, 1])
4702 );
4703 assert_eq!(
4704 PixelFormat::Nv12.image_shape(640, 480),
4705 Some(vec![720, 640])
4706 );
4707 // Odd height: combined-plane height is `481 + ceil(481/2)` = 481 + 241
4708 // = 722 rows. Logical height is recovered as `722 * 2 / 3` = 481.
4709 assert_eq!(
4710 PixelFormat::Nv12.image_shape(640, 481),
4711 Some(vec![722, 640])
4712 );
4713 // Odd width: shape carries the LOGICAL width (641).
4714 // The 64-aligned stride (>= 642) is stored separately on the Tensor.
4715 assert_eq!(
4716 PixelFormat::Nv12.image_shape(641, 480),
4717 Some(vec![720, 641])
4718 );
4719 // NV16 odd width: same — logical width in shape, stride separate.
4720 assert_eq!(
4721 PixelFormat::Nv16.image_shape(641, 480),
4722 Some(vec![960, 641])
4723 );
4724 assert_eq!(
4725 PixelFormat::PlanarRgb.image_shape(640, 480),
4726 Some(vec![3, 480, 640])
4727 );
4728 assert_eq!(
4729 PixelFormat::Nv16.image_shape(640, 480),
4730 Some(vec![960, 640])
4731 );
4732 }
4733
4734 #[test]
4735 fn raw_tensor_has_no_format() {
4736 let t = Tensor::<u8>::new(&[480, 640, 3], None, None).unwrap();
4737 assert!(t.format().is_none());
4738 assert!(t.width().is_none());
4739 assert!(t.height().is_none());
4740 assert!(!t.is_multiplane());
4741 assert!(t.chroma().is_none());
4742 }
4743
4744 #[test]
4745 fn image_tensor_packed() {
4746 let t = Tensor::<u8>::image(
4747 640,
4748 480,
4749 PixelFormat::Rgba,
4750 None,
4751 crate::CpuAccess::ReadWrite,
4752 )
4753 .unwrap();
4754 assert_eq!(t.format(), Some(PixelFormat::Rgba));
4755 assert_eq!(t.width(), Some(640));
4756 assert_eq!(t.height(), Some(480));
4757 assert_eq!(t.shape(), &[480, 640, 4]);
4758 assert!(!t.is_multiplane());
4759 }
4760
4761 #[test]
4762 fn image_tensor_planar() {
4763 let t = Tensor::<u8>::image(
4764 640,
4765 480,
4766 PixelFormat::PlanarRgb,
4767 None,
4768 crate::CpuAccess::ReadWrite,
4769 )
4770 .unwrap();
4771 assert_eq!(t.format(), Some(PixelFormat::PlanarRgb));
4772 assert_eq!(t.width(), Some(640));
4773 assert_eq!(t.height(), Some(480));
4774 assert_eq!(t.shape(), &[3, 480, 640]);
4775 }
4776
4777 #[test]
4778 #[cfg(target_os = "macos")]
4779 fn image_tensor_dma_non_aligned_packed_width_pads_zero_copy() {
4780 // RGBA u8 at width=4 → 4*4 = 16 bytes/row, not 64-byte aligned. RGBA has
4781 // a real IOSurface FourCC, so an explicit `Some(TensorMemory::Dma)`
4782 // request now allocates a padded image IOSurface (64-aligned
4783 // `bytes_per_row`) and records the stride — a fully zero-copy buffer GL
4784 // can bind and the CPU can map via the strided path. (Previously this
4785 // failed loudly to avoid an 'L008' byte-bag downgrade; with a real
4786 // FourCC surface that concern no longer applies.)
4787 let t = Tensor::<u8>::image(
4788 4,
4789 4,
4790 PixelFormat::Rgba,
4791 Some(TensorMemory::Dma),
4792 crate::CpuAccess::ReadWrite,
4793 )
4794 .expect("padded RGBA IOSurface should allocate");
4795 assert_eq!(t.format(), Some(PixelFormat::Rgba));
4796 assert_eq!(t.width(), Some(4));
4797 assert_eq!(t.height(), Some(4));
4798 let stride = t.effective_row_stride().expect("stride");
4799 assert_eq!(stride % 64, 0, "padded to 64-byte row alignment");
4800 assert!(stride >= 16);
4801 // A CPU map exposes the full padded surface for strided iteration.
4802 let m = t.map().expect("strided IOSurface map");
4803 assert_eq!(m.as_slice().len(), stride * 4);
4804 }
4805
4806 /// `per_pixel_bytes` that doesn't divide 64 evenly (e.g. RGB u8 with
4807 /// 3 B/pixel) makes a "Pad width to N" suggestion structurally
4808 /// impossible — there is no integer width whose `width * 3` is a
4809 /// multiple of 64. The error must still fire (no silent SHM
4810 /// fallback for explicit-DMA requests) and must spell out the
4811 /// alignment requirement; it just omits the misleading "pad to N"
4812 /// hint instead of printing a number whose row pitch still won't
4813 /// align.
4814 #[test]
4815 #[cfg(target_os = "macos")]
4816 fn image_tensor_dma_rejects_indivisible_pixel_pitch_without_pad_hint() {
4817 // Width=10 RGB f32 → 120 B/row, not 64-byte aligned, and (Rgb,
4818 // F32) has no IOSurface mapping so the padded-stride tolerance
4819 // does not apply. The next 64-multiple (128 B) isn't an integer
4820 // multiple of 12 B/pixel, so the "pad width to N" hint can't
4821 // produce a valid number and must be omitted. (Rgb u8 used to be
4822 // this test's subject but now has a real RGBA8888 mapping with
4823 // padded-stride tolerance — see the test below.)
4824 let err = Tensor::<f32>::image(
4825 10,
4826 10,
4827 PixelFormat::Rgb,
4828 Some(TensorMemory::Dma),
4829 crate::CpuAccess::ReadWrite,
4830 )
4831 .expect_err("RGB f32 with 12 B/pixel and non-aligned width must be rejected");
4832 match err {
4833 Error::InvalidArgument(msg) => {
4834 assert!(
4835 msg.contains("64-byte aligned"),
4836 "error must still name the alignment requirement: {msg}"
4837 );
4838 assert!(
4839 !msg.contains("Pad width"),
4840 "indivisible per-pixel pitch makes a width suggestion impossible; \
4841 hint must be omitted, got: {msg}"
4842 );
4843 assert!(
4844 msg.contains("memory=None") && msg.contains("TensorMemory::Mem"),
4845 "error must still list the always-applicable alternatives: {msg}"
4846 );
4847 }
4848 other => panic!("expected InvalidArgument, got {other:?}"),
4849 }
4850 }
4851
4852 #[test]
4853 #[cfg(target_os = "macos")]
4854 fn image_tensor_dma_packed_rgb_u8_contract() {
4855 // Packed RGB u8 @Dma is a designed RGBA8888 mapping at
4856 // (W*3/4, H) — the INT8 NPU input layout, shared with Android.
4857 // width%4 != 0 cannot form whole texels → loud InvalidArgument…
4858 let err = Tensor::<u8>::image(
4859 10,
4860 10,
4861 PixelFormat::Rgb,
4862 Some(TensorMemory::Dma),
4863 crate::CpuAccess::ReadWrite,
4864 )
4865 .expect_err("Rgb u8 width%4!=0 must be rejected");
4866 assert!(
4867 matches!(&err, Error::InvalidArgument(m) if m.contains("width%4==0")),
4868 "got {err:?}"
4869 );
4870 // …width%4 == 0 with a non-64-aligned pitch allocates PADDED
4871 // (36 B rows → 64 B surface pitch, recorded on the tensor)…
4872 let t = Tensor::<u8>::image(
4873 12,
4874 4,
4875 PixelFormat::Rgb,
4876 Some(TensorMemory::Dma),
4877 crate::CpuAccess::ReadWrite,
4878 )
4879 .expect("width 12 Rgb u8 must allocate padded");
4880 assert_eq!(t.memory(), TensorMemory::Dma);
4881 assert!(
4882 t.row_stride()
4883 .is_some_and(|s| s >= 64 && s.is_multiple_of(64)),
4884 "padded pitch must be recorded: {:?}",
4885 t.row_stride()
4886 );
4887 // …and the aligned model-input width stays flat (640*3 = 1920 is
4888 // 64-aligned → no recorded stride, the buffer IS [H, W, 3]).
4889 let t = Tensor::<u8>::image(
4890 640,
4891 8,
4892 PixelFormat::Rgb,
4893 Some(TensorMemory::Dma),
4894 crate::CpuAccess::ReadWrite,
4895 )
4896 .expect("width 640 Rgb u8 must allocate flat");
4897 assert_eq!(t.row_stride(), None);
4898 // I8 shares the layout (INT8 shader bias, not a format change).
4899 let t = Tensor::<i8>::image(
4900 640,
4901 8,
4902 PixelFormat::Rgb,
4903 Some(TensorMemory::Dma),
4904 crate::CpuAccess::ReadWrite,
4905 )
4906 .expect("Rgb i8 shares the RGBA8888 mapping");
4907 assert_eq!(t.memory(), TensorMemory::Dma);
4908 }
4909
4910 #[test]
4911 #[cfg(target_os = "macos")]
4912 fn image_tensor_dma_planar_f16_alignment() {
4913 // PlanarRgb F16 uses single-channel row pitch (width * 2 bytes).
4914 // Width=16 → 32 bytes/row (not aligned); width=32 → 64 bytes/row (aligned).
4915 let err = Tensor::<half::f16>::image(
4916 16,
4917 16,
4918 PixelFormat::PlanarRgb,
4919 Some(TensorMemory::Dma),
4920 crate::CpuAccess::ReadWrite,
4921 )
4922 .expect_err("width=16 PlanarRgb F16 is 32-byte row, must reject");
4923 assert!(matches!(err, Error::InvalidArgument(_)), "got {err:?}");
4924 // 32 wide should work.
4925 let t = Tensor::<half::f16>::image(
4926 32,
4927 8,
4928 PixelFormat::PlanarRgb,
4929 Some(TensorMemory::Dma),
4930 crate::CpuAccess::ReadWrite,
4931 )
4932 .expect("width=32 PlanarRgb F16 is 64-byte row, must succeed");
4933 assert_eq!(t.format(), Some(PixelFormat::PlanarRgb));
4934 }
4935
4936 #[test]
4937 fn image_tensor_semi_planar_contiguous() {
4938 let t = Tensor::<u8>::image(
4939 640,
4940 480,
4941 PixelFormat::Nv12,
4942 None,
4943 crate::CpuAccess::ReadWrite,
4944 )
4945 .unwrap();
4946 assert_eq!(t.format(), Some(PixelFormat::Nv12));
4947 assert_eq!(t.width(), Some(640));
4948 assert_eq!(t.height(), Some(480));
4949 // NV12: H*3/2 = 720
4950 assert_eq!(t.shape(), &[720, 640]);
4951 assert!(!t.is_multiplane());
4952 }
4953
4954 #[test]
4955 #[cfg(target_os = "linux")]
4956 fn image_tensor_with_stride_preserves_logical_width() {
4957 // Skip if DMA not available (e.g. sandboxed CI lacking dma_heap access).
4958 if !is_dma_available() {
4959 eprintln!("SKIPPED: DMA heap not available");
4960 return;
4961 }
4962 // 3004×1688 RGBA8: natural pitch 12016, padded to 12032 (64-aligned).
4963 let stride = 12032;
4964 let t = Tensor::<u8>::image_with_stride(
4965 3004,
4966 1688,
4967 PixelFormat::Rgba,
4968 stride,
4969 Some(TensorMemory::Dma),
4970 crate::CpuAccess::ReadWrite,
4971 )
4972 .unwrap();
4973 // Logical dimensions unchanged by padding — this is the contract.
4974 assert_eq!(t.width(), Some(3004));
4975 assert_eq!(t.height(), Some(1688));
4976 assert_eq!(t.shape(), &[1688, 3004, 4]);
4977 // Stride is carried separately and reports the padded pitch.
4978 assert_eq!(t.effective_row_stride(), Some(stride));
4979 // Buffer is sized to stride × height so the full padded layout fits,
4980 // and CPU map() works for self-allocated strided DMA tensors.
4981 use crate::TensorMapTrait;
4982 {
4983 let map = t.map().unwrap();
4984 assert!(
4985 map.as_slice().len() >= stride * 1688,
4986 "mapped buffer {} bytes < expected {}",
4987 map.as_slice().len(),
4988 stride * 1688
4989 );
4990 }
4991 // CPU write access works too — iterate rows using the padded stride,
4992 // touch only the active `width × bpp` region, verify it round-trips.
4993 {
4994 let mut map = t.map().unwrap();
4995 let slice = map.as_mut_slice();
4996 for y in 0..1688 {
4997 let row_start = y * stride;
4998 for x in 0..3004 {
4999 let p = row_start + x * 4;
5000 slice[p] = (y & 0xFF) as u8;
5001 slice[p + 1] = (x & 0xFF) as u8;
5002 slice[p + 2] = 0x42;
5003 slice[p + 3] = 0xFF;
5004 }
5005 }
5006 }
5007 {
5008 let map = t.map().unwrap();
5009 let slice = map.as_slice();
5010 // Sample a few pixels to confirm the round-trip.
5011 assert_eq!(slice[0], 0x00);
5012 assert_eq!(slice[1], 0x00);
5013 assert_eq!(slice[2], 0x42);
5014 assert_eq!(slice[3], 0xFF);
5015 let mid = 100 * stride + 50 * 4;
5016 assert_eq!(slice[mid], 100);
5017 assert_eq!(slice[mid + 1], 50);
5018 assert_eq!(slice[mid + 2], 0x42);
5019 }
5020 }
5021
5022 #[test]
5023 #[cfg(target_os = "linux")]
5024 fn image_tensor_with_stride_rejects_foreign_strided_map() {
5025 // A FOREIGN (imported via from_fd) DMA tensor with row_stride set
5026 // should still refuse CPU mapping — external allocator owns the
5027 // layout. This protects the V4L2 / GStreamer use case.
5028 //
5029 // We simulate a foreign import by wrapping our own allocation's
5030 // fd via `from_fd` and calling set_row_stride manually. The
5031 // `is_imported` flag on from_fd is true by construction.
5032 if !is_dma_available() {
5033 eprintln!("SKIPPED: DMA heap not available");
5034 return;
5035 }
5036 // Allocate a backing buffer large enough for a 320×240 BGRA8 image.
5037 let backing = Tensor::<u8>::new(&[240 * 320 * 4], Some(TensorMemory::Dma), None).unwrap();
5038 let fd = backing.clone_fd().unwrap();
5039 // Import it via from_fd — this marks is_imported=true.
5040 let shape = [240usize, 320, 4];
5041 let storage = TensorStorage::<u8>::from_fd(fd, &shape, None).unwrap();
5042 let mut t = Tensor::<u8>::wrap(storage);
5043 t.set_format(PixelFormat::Bgra).unwrap();
5044 t.set_row_stride(320 * 4).unwrap(); // natural, but still marks it as strided
5045 let err = t.map();
5046 assert!(
5047 matches!(err, Err(Error::InvalidOperation(_))),
5048 "foreign strided map should error"
5049 );
5050 }
5051
5052 #[test]
5053 #[cfg(target_os = "linux")]
5054 fn image_tensor_with_stride_map_rejects_tampered_stride() {
5055 // Round-3 PR feedback (C1): `set_row_stride` is public and only
5056 // validates `stride >= min_stride`, not that the new stride × height
5057 // fits the underlying buffer. A caller that tampers with the stride
5058 // after allocation must not be able to coerce `Tensor::map()` into
5059 // returning a slice larger than the backing mmap (that would be UB
5060 // in `DmaMap::as_slice`).
5061 if !is_dma_available() {
5062 eprintln!("SKIPPED: DMA heap not available");
5063 return;
5064 }
5065 // Allocate a 640×480 RGBA8 padded canvas (stride = 3072 = 768 px).
5066 // Backing buffer is 3072 × 480 = 1,474,560 bytes.
5067 let mut t = Tensor::<u8>::image_with_stride(
5068 640,
5069 480,
5070 PixelFormat::Rgba,
5071 3072,
5072 Some(TensorMemory::Dma),
5073 crate::CpuAccess::ReadWrite,
5074 )
5075 .unwrap();
5076 // Tamper: push the stride up to 4 × the original. This is >=
5077 // min_stride (2560), so `set_row_stride` accepts it.
5078 t.set_row_stride(12288).unwrap();
5079 // Map must now refuse — 12288 × 480 = 5,898,240 > 1,474,560.
5080 let err = t.map();
5081 assert!(
5082 matches!(err, Err(Error::InvalidOperation(_))),
5083 "map() with oversized stride must return InvalidOperation"
5084 );
5085 }
5086
5087 #[test]
5088 fn dma_tensor_new_with_byte_size_rejects_shape_overflow() {
5089 // Round-3 PR feedback (C3): shape.product() * sizeof(T) must use
5090 // checked arithmetic so a pathological shape can't wrap usize and
5091 // make the byte_size-vs-logical-size comparison incorrect.
5092 //
5093 // This test only exercises the overflow rejection path, which is
5094 // pure-Rust and doesn't touch dma_heap — safe to run on any target.
5095 #[cfg(target_os = "linux")]
5096 {
5097 let err = crate::dma::DmaTensor::<u64>::new_with_byte_size(
5098 &[usize::MAX, 2, 2],
5099 usize::MAX,
5100 None,
5101 );
5102 assert!(
5103 matches!(err, Err(Error::InvalidArgument(_))),
5104 "new_with_byte_size must detect shape.product() overflow"
5105 );
5106 }
5107 }
5108
5109 #[test]
5110 #[cfg(target_os = "linux")]
5111 fn image_tensor_with_stride_rejects_too_small_stride() {
5112 // 640×480 RGBA8 natural pitch = 2560, request 2400 → should error.
5113 let err = Tensor::<u8>::image_with_stride(
5114 640,
5115 480,
5116 PixelFormat::Rgba,
5117 2400,
5118 Some(TensorMemory::Dma),
5119 crate::CpuAccess::ReadWrite,
5120 );
5121 assert!(matches!(err, Err(Error::InvalidArgument(_))));
5122 }
5123
5124 #[test]
5125 #[cfg(target_os = "linux")]
5126 fn image_tensor_with_stride_rejects_non_packed() {
5127 // NV12 is SemiPlanar → not supported. (Linux-only because
5128 // `TensorMemory::Dma` itself is a Linux-only enum variant.)
5129 let err = Tensor::<u8>::image_with_stride(
5130 640,
5131 480,
5132 PixelFormat::Nv12,
5133 640,
5134 Some(TensorMemory::Dma),
5135 crate::CpuAccess::ReadWrite,
5136 );
5137 assert!(matches!(err, Err(Error::NotImplemented(_))));
5138 }
5139
5140 #[test]
5141 fn set_format_valid() {
5142 let mut t = Tensor::<u8>::new(&[480, 640, 3], None, None).unwrap();
5143 assert!(t.format().is_none());
5144 t.set_format(PixelFormat::Rgb).unwrap();
5145 assert_eq!(t.format(), Some(PixelFormat::Rgb));
5146 assert_eq!(t.width(), Some(640));
5147 assert_eq!(t.height(), Some(480));
5148 }
5149
5150 #[test]
5151 fn set_format_invalid_shape() {
5152 let mut t = Tensor::<u8>::new(&[480, 640, 4], None, None).unwrap();
5153 // RGB expects 3 channels, not 4
5154 let err = t.set_format(PixelFormat::Rgb);
5155 assert!(err.is_err());
5156 // Original tensor is unmodified
5157 assert!(t.format().is_none());
5158 }
5159
5160 #[test]
5161 fn reshape_clears_format() {
5162 let mut t = Tensor::<u8>::image(
5163 640,
5164 480,
5165 PixelFormat::Rgba,
5166 None,
5167 crate::CpuAccess::ReadWrite,
5168 )
5169 .unwrap();
5170 assert_eq!(t.format(), Some(PixelFormat::Rgba));
5171 // Reshape to flat — format cleared
5172 t.reshape(&[480 * 640 * 4]).unwrap();
5173 assert!(t.format().is_none());
5174 }
5175
5176 #[test]
5177 fn from_planes_nv12() {
5178 let y = Tensor::<u8>::new(&[480, 640], None, None).unwrap();
5179 let uv = Tensor::<u8>::new(&[240, 640], None, None).unwrap();
5180 let img = Tensor::from_planes(y, uv, PixelFormat::Nv12).unwrap();
5181 assert_eq!(img.format(), Some(PixelFormat::Nv12));
5182 assert!(img.is_multiplane());
5183 assert!(img.chroma().is_some());
5184 assert_eq!(img.width(), Some(640));
5185 assert_eq!(img.height(), Some(480));
5186 }
5187
5188 #[test]
5189 fn from_planes_rejects_non_semiplanar() {
5190 let y = Tensor::<u8>::new(&[480, 640], None, None).unwrap();
5191 let uv = Tensor::<u8>::new(&[240, 640], None, None).unwrap();
5192 let err = Tensor::from_planes(y, uv, PixelFormat::Rgb);
5193 assert!(err.is_err());
5194 }
5195
5196 #[test]
5197 fn reshape_multiplane_errors() {
5198 let y = Tensor::<u8>::new(&[480, 640], None, None).unwrap();
5199 let uv = Tensor::<u8>::new(&[240, 640], None, None).unwrap();
5200 let mut img = Tensor::from_planes(y, uv, PixelFormat::Nv12).unwrap();
5201 let err = img.reshape(&[480 * 640 + 240 * 640]);
5202 assert!(err.is_err());
5203 }
5204}
5205
5206#[cfg(test)]
5207mod compression_tests {
5208 use super::*;
5209
5210 #[test]
5211 fn desc_builder_roundtrips() {
5212 let desc = ImageDesc::new(640, 480, PixelFormat::Rgba, DType::U8)
5213 .with_memory(Some(TensorMemory::Mem))
5214 .with_access(CpuAccess::Read)
5215 .with_compression(Compression::Any);
5216 assert_eq!(desc.width(), 640);
5217 assert_eq!(desc.height(), 480);
5218 assert_eq!(desc.format(), PixelFormat::Rgba);
5219 assert_eq!(desc.dtype(), DType::U8);
5220 assert_eq!(desc.memory(), Some(TensorMemory::Mem));
5221 assert_eq!(desc.access(), CpuAccess::Read);
5222 assert_eq!(desc.compression(), Some(Compression::Any));
5223
5224 // Defaults: auto memory, hardware-only, no request.
5225 let plain = ImageDesc::new(2, 2, PixelFormat::Grey, DType::U8);
5226 assert_eq!(plain.memory(), None);
5227 assert_eq!(plain.access(), CpuAccess::None);
5228 assert_eq!(plain.compression(), None);
5229 }
5230
5231 #[test]
5232 fn desc_dtype_must_match_element_type() {
5233 let desc = ImageDesc::new(4, 4, PixelFormat::Rgba, DType::F32);
5234 match Tensor::<u8>::image_desc(&desc) {
5235 Err(Error::InvalidArgument(msg)) => assert!(msg.contains("dtype")),
5236 other => panic!("expected InvalidArgument, got {other:?}"),
5237 }
5238 }
5239
5240 #[test]
5241 fn compression_with_cpu_access_is_invalid() {
5242 let desc = ImageDesc::new(4, 4, PixelFormat::Rgba, DType::U8)
5243 .with_access(CpuAccess::ReadWrite)
5244 .with_compression(Compression::Any);
5245 match Tensor::<u8>::image_desc(&desc) {
5246 Err(Error::InvalidArgument(msg)) => assert!(msg.contains("CpuAccess::None")),
5247 other => panic!("expected InvalidArgument, got {other:?}"),
5248 }
5249 }
5250
5251 #[cfg(not(target_os = "android"))]
5252 #[test]
5253 fn scheme_request_off_android_is_not_implemented() {
5254 let desc = ImageDesc::new(4, 4, PixelFormat::Rgba, DType::U8)
5255 .with_compression(Compression::Scheme(CompressionScheme::Ubwc));
5256 match Tensor::<u8>::image_desc(&desc) {
5257 Err(Error::NotImplemented(msg)) => assert!(msg.contains("Ubwc")),
5258 other => panic!("expected NotImplemented, got {other:?}"),
5259 }
5260 }
5261
5262 #[cfg(not(target_os = "android"))]
5263 #[test]
5264 fn any_request_off_android_resolves_linear_and_counts() {
5265 let before = compression_fallback_count();
5266 let desc = ImageDesc::new(64, 64, PixelFormat::Rgba, DType::U8)
5267 .with_memory(Some(TensorMemory::Mem))
5268 .with_compression(Compression::Any);
5269 let t = Tensor::<u8>::image_desc(&desc).unwrap();
5270 assert_eq!(t.compression(), None);
5271 assert!(compression_fallback_count() > before);
5272 }
5273
5274 #[test]
5275 fn desc_without_request_matches_classic_constructor() {
5276 let desc = ImageDesc::new(32, 32, PixelFormat::Rgba, DType::U8)
5277 .with_memory(Some(TensorMemory::Mem))
5278 .with_access(CpuAccess::ReadWrite);
5279 let t = Tensor::<u8>::image_desc(&desc).unwrap();
5280 assert_eq!(t.compression(), None);
5281 assert_eq!(t.cpu_access(), CpuAccess::ReadWrite);
5282 assert_eq!(t.width(), Some(32));
5283 // Mappable exactly like the classic constructor's result.
5284 let m = t.map_read().unwrap();
5285 assert_eq!(m.as_slice().len(), 32 * 32 * 4);
5286 }
5287
5288 #[test]
5289 fn configure_image_preserves_compression_and_views_inherit() {
5290 // No host platform records a scheme, so emulate the recording to
5291 // pin the preserve/inherit semantics (the physical layout does
5292 // not change when the logical image is reconfigured).
5293 let mut t = Tensor::<u8>::image(
5294 64,
5295 64,
5296 PixelFormat::Rgba,
5297 Some(TensorMemory::Mem),
5298 CpuAccess::None,
5299 )
5300 .unwrap();
5301 t.compression = Some(CompressionScheme::Ubwc);
5302 t.configure_image(32, 32, PixelFormat::Rgba).unwrap();
5303 assert_eq!(t.compression(), Some(CompressionScheme::Ubwc));
5304 let view = t.subview(0, &[16, 32, 4]).unwrap();
5305 assert_eq!(view.compression(), Some(CompressionScheme::Ubwc));
5306 }
5307
5308 #[test]
5309 fn tensor_dyn_dispatches_desc_and_compression() {
5310 let desc = ImageDesc::new(16, 16, PixelFormat::Rgba, DType::U8)
5311 .with_memory(Some(TensorMemory::Mem))
5312 .with_access(CpuAccess::ReadWrite);
5313 let t = TensorDyn::image_desc(&desc).unwrap();
5314 assert_eq!(t.compression(), None);
5315 assert!(matches!(t, TensorDyn::U8(_)));
5316 }
5317}
5318
5319#[cfg(test)]
5320mod cpu_access_tests {
5321 use super::*;
5322
5323 #[test]
5324 fn covers_matrix() {
5325 use CpuAccess::*;
5326 // Every declaration covers a narrower or equal request…
5327 for a in [None, Read, Write, ReadWrite] {
5328 assert!(a.covers(None), "{a:?} must cover None");
5329 assert!(ReadWrite.covers(a), "ReadWrite must cover {a:?}");
5330 }
5331 assert!(Read.covers(Read));
5332 assert!(Write.covers(Write));
5333 // …and never a wider one.
5334 assert!(!None.covers(Read));
5335 assert!(!None.covers(Write));
5336 assert!(!Read.covers(Write));
5337 assert!(!Read.covers(ReadWrite));
5338 assert!(!Write.covers(Read));
5339 assert!(!Write.covers(ReadWrite));
5340 }
5341
5342 #[test]
5343 fn map_with_none_is_invalid() {
5344 let t = Tensor::<u8>::new(&[16], Some(TensorMemory::Mem), None).unwrap();
5345 match t.map_with(CpuAccess::None) {
5346 Err(Error::InvalidArgument(_)) => {}
5347 Err(other) => panic!("expected InvalidArgument, got {other:?}"),
5348 Ok(_) => panic!("map_with(CpuAccess::None) must not succeed"),
5349 }
5350 }
5351
5352 #[test]
5353 fn read_map_rejects_mutation_uniformly() {
5354 // Mem backend: map_read yields a working read view whose mutable
5355 // accessor panics (the uniform cross-backend contract).
5356 let t = Tensor::<u8>::new(&[8], Some(TensorMemory::Mem), None).unwrap();
5357 t.map_mut().unwrap().as_mut_slice().copy_from_slice(&[7; 8]);
5358 let ro = t.map_read().unwrap();
5359 assert_eq!(ro.as_slice(), &[7; 8]);
5360 drop(ro);
5361 let result = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
5362 let mut ro = t.map_read().unwrap();
5363 let _ = ro.as_mut_slice();
5364 }));
5365 assert!(result.is_err(), "as_mut_slice through map_read must panic");
5366 }
5367
5368 #[test]
5369 fn write_and_rw_maps_stay_mutable() {
5370 let t = Tensor::<u8>::new(&[4], Some(TensorMemory::Mem), None).unwrap();
5371 t.map_write()
5372 .unwrap()
5373 .as_mut_slice()
5374 .copy_from_slice(&[1; 4]);
5375 t.map().unwrap().as_mut_slice().copy_from_slice(&[2; 4]);
5376 assert_eq!(t.map_read().unwrap().as_slice(), &[2; 4]);
5377 }
5378
5379 /// The read-only IOSurface lock path: a `map_read` must observe data
5380 /// written through a prior read-write lock, and its unlock (which
5381 /// skips the cache flush) must not disturb subsequent reads.
5382 #[test]
5383 #[cfg(target_os = "macos")]
5384 fn iosurface_read_only_lock_roundtrip() {
5385 let Ok(t) = Tensor::<u8>::new(&[64], Some(TensorMemory::Dma), None) else {
5386 eprintln!("SKIPPED: IOSurface unavailable");
5387 return;
5388 };
5389 {
5390 let mut m = t.map_mut().unwrap();
5391 for (i, b) in m.as_mut_slice().iter_mut().enumerate() {
5392 *b = (i * 3) as u8;
5393 }
5394 }
5395 for _ in 0..2 {
5396 let ro = t.map_read().unwrap();
5397 for (i, b) in ro.as_slice().iter().enumerate() {
5398 assert_eq!(*b, (i * 3) as u8, "byte {i} through read-only lock");
5399 }
5400 }
5401 let result = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
5402 let mut ro = t.map_read().unwrap();
5403 let _ = ro.as_mut_slice();
5404 }));
5405 assert!(result.is_err(), "IOSurface read map must reject mutation");
5406 }
5407}
5408
5409#[cfg(test)]
5410mod tests {
5411 #[cfg(target_os = "linux")]
5412 use nix::unistd::{access, AccessFlags};
5413 #[cfg(target_os = "linux")]
5414 use std::io::Write as _;
5415 use std::sync::RwLock;
5416
5417 use super::*;
5418
5419 #[ctor::ctor(unsafe)]
5420 fn init() {
5421 env_logger::Builder::from_env(env_logger::Env::default().default_filter_or("info")).init();
5422 }
5423
5424 /// Macro to get the current function name for logging in tests.
5425 #[cfg(target_os = "linux")]
5426 macro_rules! function {
5427 () => {{
5428 fn f() {}
5429 fn type_name_of<T>(_: T) -> &'static str {
5430 std::any::type_name::<T>()
5431 }
5432 let name = type_name_of(f);
5433
5434 // Find and cut the rest of the path
5435 match &name[..name.len() - 3].rfind(':') {
5436 Some(pos) => &name[pos + 1..name.len() - 3],
5437 None => &name[..name.len() - 3],
5438 }
5439 }};
5440 }
5441
5442 #[test]
5443 #[cfg(target_os = "linux")]
5444 fn test_tensor() {
5445 let _lock = FD_LOCK.read().unwrap();
5446 let shape = vec![1];
5447 let tensor = DmaTensor::<f32>::new(&shape, Some("dma_tensor"));
5448 let dma_enabled = tensor.is_ok();
5449
5450 let tensor = Tensor::<f32>::new(&shape, None, None).expect("Failed to create tensor");
5451 // Auto-select priority is Dma > Mem; Shm is never auto-selected.
5452 match dma_enabled {
5453 true => assert_eq!(tensor.memory(), TensorMemory::Dma),
5454 false => assert_eq!(tensor.memory(), TensorMemory::Mem),
5455 }
5456 }
5457
5458 #[test]
5459 #[cfg(any(target_os = "macos", target_os = "ios"))]
5460 fn test_tensor() {
5461 let shape = vec![1];
5462 let tensor = Tensor::<f32>::new(&shape, None, None).expect("Failed to create tensor");
5463 // macOS/iOS auto-fallback chain: IOSurface (Dma) → Mem. Healthy systems
5464 // return Dma; Mem only appears under memory pressure or sandboxed
5465 // contexts where IOSurfaceCreate fails. Shm is never auto-selected.
5466 let m = tensor.memory();
5467 assert!(
5468 matches!(m, TensorMemory::Dma | TensorMemory::Mem),
5469 "Unexpected auto-fallback result on macOS/iOS: {m:?}"
5470 );
5471 }
5472
5473 #[test]
5474 #[cfg(all(
5475 unix,
5476 not(any(target_os = "linux", target_os = "macos", target_os = "ios"))
5477 ))]
5478 fn test_tensor() {
5479 let shape = vec![1];
5480 let tensor = Tensor::<f32>::new(&shape, None, None).expect("Failed to create tensor");
5481 // Other Unix (BSD): no DMA, so auto-selection is Mem (Shm is
5482 // explicit-only, never auto-selected).
5483 assert_eq!(tensor.memory(), TensorMemory::Mem);
5484 }
5485
5486 #[test]
5487 #[cfg(not(unix))]
5488 fn test_tensor() {
5489 let shape = vec![1];
5490 let tensor = Tensor::<f32>::new(&shape, None, None).expect("Failed to create tensor");
5491 assert_eq!(tensor.memory(), TensorMemory::Mem);
5492 }
5493
5494 #[test]
5495 #[cfg(target_os = "linux")]
5496 fn test_dma_tensor() {
5497 let _lock = FD_LOCK.read().unwrap();
5498 match access(
5499 "/dev/dma_heap/linux,cma",
5500 AccessFlags::R_OK | AccessFlags::W_OK,
5501 ) {
5502 Ok(_) => println!("/dev/dma_heap/linux,cma is available"),
5503 Err(_) => match access(
5504 "/dev/dma_heap/system",
5505 AccessFlags::R_OK | AccessFlags::W_OK,
5506 ) {
5507 Ok(_) => println!("/dev/dma_heap/system is available"),
5508 Err(e) => {
5509 writeln!(
5510 &mut std::io::stdout(),
5511 "[WARNING] DMA Heap is unavailable: {e}"
5512 )
5513 .unwrap();
5514 return;
5515 }
5516 },
5517 }
5518
5519 let shape = vec![2, 3, 4];
5520 let tensor =
5521 DmaTensor::<f32>::new(&shape, Some("test_tensor")).expect("Failed to create tensor");
5522
5523 const DUMMY_VALUE: f32 = 12.34;
5524
5525 assert_eq!(tensor.memory(), TensorMemory::Dma);
5526 assert_eq!(tensor.name(), "test_tensor");
5527 assert_eq!(tensor.shape(), &shape);
5528 assert_eq!(tensor.size(), 2 * 3 * 4 * std::mem::size_of::<f32>());
5529 assert_eq!(tensor.len(), 2 * 3 * 4);
5530
5531 {
5532 let mut tensor_map = tensor.map().expect("Failed to map DMA memory");
5533 tensor_map.fill(42.0);
5534 assert!(tensor_map.iter().all(|&x| x == 42.0));
5535 }
5536
5537 {
5538 let shared = Tensor::<f32>::from_fd(
5539 tensor
5540 .clone_fd()
5541 .expect("Failed to duplicate tensor file descriptor"),
5542 &shape,
5543 Some("test_tensor_shared"),
5544 )
5545 .expect("Failed to create tensor from fd");
5546
5547 assert_eq!(shared.memory(), TensorMemory::Dma);
5548 assert_eq!(shared.name(), "test_tensor_shared");
5549 assert_eq!(shared.shape(), &shape);
5550
5551 let mut tensor_map = shared.map().expect("Failed to map DMA memory from fd");
5552 tensor_map.fill(DUMMY_VALUE);
5553 assert!(tensor_map.iter().all(|&x| x == DUMMY_VALUE));
5554 }
5555
5556 {
5557 let tensor_map = tensor.map().expect("Failed to map DMA memory");
5558 assert!(tensor_map.iter().all(|&x| x == DUMMY_VALUE));
5559 }
5560
5561 let mut tensor = DmaTensor::<u8>::new(&shape, None).expect("Failed to create tensor");
5562 assert_eq!(tensor.shape(), &shape);
5563 let new_shape = vec![3, 4, 4];
5564 assert!(
5565 tensor.reshape(&new_shape).is_err(),
5566 "Reshape should fail due to size mismatch"
5567 );
5568 assert_eq!(tensor.shape(), &shape, "Shape should remain unchanged");
5569
5570 let new_shape = vec![2, 3, 4];
5571 tensor.reshape(&new_shape).expect("Reshape should succeed");
5572 assert_eq!(
5573 tensor.shape(),
5574 &new_shape,
5575 "Shape should be updated after successful reshape"
5576 );
5577
5578 {
5579 let mut tensor_map = tensor.map().expect("Failed to map DMA memory");
5580 tensor_map.fill(1);
5581 assert!(tensor_map.iter().all(|&x| x == 1));
5582 }
5583
5584 {
5585 let mut tensor_map = tensor.map().expect("Failed to map DMA memory");
5586 tensor_map[2] = 42;
5587 assert_eq!(tensor_map[1], 1, "Value at index 1 should be 1");
5588 assert_eq!(tensor_map[2], 42, "Value at index 2 should be 42");
5589 }
5590 }
5591
5592 #[test]
5593 #[cfg(unix)]
5594 fn test_shm_tensor() {
5595 let _lock = FD_LOCK.read().unwrap();
5596 let shape = vec![2, 3, 4];
5597 let tensor =
5598 ShmTensor::<f32>::new(&shape, Some("test_tensor")).expect("Failed to create tensor");
5599 assert_eq!(tensor.shape(), &shape);
5600 assert_eq!(tensor.size(), 2 * 3 * 4 * std::mem::size_of::<f32>());
5601 assert_eq!(tensor.name(), "test_tensor");
5602
5603 const DUMMY_VALUE: f32 = 12.34;
5604 {
5605 let mut tensor_map = tensor.map().expect("Failed to map shared memory");
5606 tensor_map.fill(42.0);
5607 assert!(tensor_map.iter().all(|&x| x == 42.0));
5608 }
5609
5610 {
5611 let shared = Tensor::<f32>::from_fd(
5612 tensor
5613 .clone_fd()
5614 .expect("Failed to duplicate tensor file descriptor"),
5615 &shape,
5616 Some("test_tensor_shared"),
5617 )
5618 .expect("Failed to create tensor from fd");
5619
5620 assert_eq!(shared.memory(), TensorMemory::Shm);
5621 assert_eq!(shared.name(), "test_tensor_shared");
5622 assert_eq!(shared.shape(), &shape);
5623
5624 let mut tensor_map = shared.map().expect("Failed to map shared memory from fd");
5625 tensor_map.fill(DUMMY_VALUE);
5626 assert!(tensor_map.iter().all(|&x| x == DUMMY_VALUE));
5627 }
5628
5629 {
5630 let tensor_map = tensor.map().expect("Failed to map shared memory");
5631 assert!(tensor_map.iter().all(|&x| x == DUMMY_VALUE));
5632 }
5633
5634 let mut tensor = ShmTensor::<u8>::new(&shape, None).expect("Failed to create tensor");
5635 assert_eq!(tensor.shape(), &shape);
5636 let new_shape = vec![3, 4, 4];
5637 assert!(
5638 tensor.reshape(&new_shape).is_err(),
5639 "Reshape should fail due to size mismatch"
5640 );
5641 assert_eq!(tensor.shape(), &shape, "Shape should remain unchanged");
5642
5643 let new_shape = vec![2, 3, 4];
5644 tensor.reshape(&new_shape).expect("Reshape should succeed");
5645 assert_eq!(
5646 tensor.shape(),
5647 &new_shape,
5648 "Shape should be updated after successful reshape"
5649 );
5650
5651 {
5652 let mut tensor_map = tensor.map().expect("Failed to map shared memory");
5653 tensor_map.fill(1);
5654 assert!(tensor_map.iter().all(|&x| x == 1));
5655 }
5656
5657 {
5658 let mut tensor_map = tensor.map().expect("Failed to map shared memory");
5659 tensor_map[2] = 42;
5660 assert_eq!(tensor_map[1], 1, "Value at index 1 should be 1");
5661 assert_eq!(tensor_map[2], 42, "Value at index 2 should be 42");
5662 }
5663 }
5664
5665 #[test]
5666 fn mem_subview_partitions_parent_buffer() {
5667 // One heap [2,4] u8 parent (8 bytes). Two [1,4] sub-views at byte
5668 // offsets 0 and 4 must share the parent allocation (zero-copy) and be
5669 // independently writable: view 0 owns bytes [0,4), view 1 owns [4,8).
5670 // Today this is impossible — heap offset is rejected and there is no
5671 // shared sub-view constructor.
5672 let parent = Tensor::<u8>::new(&[2, 4], Some(TensorMemory::Mem), None).unwrap();
5673 let view0 = parent.subview(0, &[1, 4]).expect("subview at offset 0");
5674 let view1 = parent.subview(4, &[1, 4]).expect("subview at offset 4");
5675
5676 view1
5677 .map()
5678 .unwrap()
5679 .as_mut_slice()
5680 .copy_from_slice(&[10, 20, 30, 40]);
5681 view0
5682 .map()
5683 .unwrap()
5684 .as_mut_slice()
5685 .copy_from_slice(&[1, 2, 3, 4]);
5686
5687 // Each view sees only its own window.
5688 assert_eq!(view0.map().unwrap().as_slice(), &[1, 2, 3, 4]);
5689 assert_eq!(view1.map().unwrap().as_slice(), &[10, 20, 30, 40]);
5690 // The parent buffer is correctly partitioned (shared, zero-copy).
5691 assert_eq!(
5692 parent.map().unwrap().as_slice(),
5693 &[1, 2, 3, 4, 10, 20, 30, 40]
5694 );
5695 }
5696
5697 #[test]
5698 fn batch_partitions_leading_dim() {
5699 // Raw [4,2,2,3] u8 batched tensor: 4 elements of 12 bytes each. batch(n)
5700 // yields element n at offset n*12, sharing the parent buffer (zero-copy).
5701 let parent = Tensor::<u8>::new(&[4, 2, 2, 3], Some(TensorMemory::Mem), None).unwrap();
5702 for i in 0..4u8 {
5703 let e = parent.batch(i as usize).expect("batch element");
5704 assert_eq!(e.shape(), &[2, 2, 3]);
5705 // A batch element shares the parent's BufferIdentity.
5706 assert_eq!(e.buffer_identity().id(), parent.buffer_identity().id());
5707 for b in e.map().unwrap().as_mut_slice() {
5708 *b = i + 1;
5709 }
5710 }
5711 // Each element occupies its own 12-byte band of the parent.
5712 let whole = parent.map().unwrap();
5713 let s = whole.as_slice();
5714 for i in 0..4usize {
5715 assert!(
5716 s[i * 12..(i + 1) * 12].iter().all(|&b| b == (i as u8 + 1)),
5717 "band {i} not partitioned: {:?}",
5718 &s[i * 12..(i + 1) * 12]
5719 );
5720 }
5721 }
5722
5723 #[test]
5724 fn view_origin_snapshots_parent_and_composes() {
5725 // view() on a whole image snapshots the parent dims + the view's origin.
5726 let parent = Tensor::<u8>::image(
5727 100,
5728 80,
5729 PixelFormat::Rgba,
5730 Some(TensorMemory::Mem),
5731 crate::CpuAccess::ReadWrite,
5732 )
5733 .unwrap();
5734 assert_eq!(
5735 parent.view_origin(),
5736 None,
5737 "whole tensor has no view_origin"
5738 );
5739 let v = parent.view(Region::new(10, 20, 30, 40)).unwrap();
5740 assert_eq!(
5741 v.view_origin(),
5742 Some(ViewOrigin {
5743 parent_width: 100,
5744 parent_height: 80,
5745 parent_row_stride: 100 * 4, // tight RGBA pitch
5746 x: 10,
5747 y: 20
5748 })
5749 );
5750 // A view of a view keeps the ROOT parent and accumulates the origin.
5751 let v2 = v.view(Region::new(5, 5, 10, 10)).unwrap();
5752 assert_eq!(
5753 v2.view_origin(),
5754 Some(ViewOrigin {
5755 parent_width: 100,
5756 parent_height: 80,
5757 parent_row_stride: 100 * 4,
5758 x: 15,
5759 y: 25
5760 }),
5761 "nested view composes onto the root parent"
5762 );
5763 }
5764
5765 #[test]
5766 fn view_origin_none_for_raw_batch() {
5767 // A raw (unformatted) batched tensor has no pixel geometry, so batch()
5768 // leaves view_origin None (the per-slot path, not the one-import pivot).
5769 let parent = Tensor::<u8>::new(&[4, 2, 2, 3], Some(TensorMemory::Mem), None).unwrap();
5770 assert_eq!(parent.batch(2).unwrap().view_origin(), None);
5771 }
5772
5773 #[test]
5774 fn batch_rejects_out_of_bounds_index() {
5775 let parent = Tensor::<u8>::new(&[4, 2, 2, 3], Some(TensorMemory::Mem), None).unwrap();
5776 match parent.batch(4) {
5777 Err(Error::BatchIndexOutOfBounds { index, batch }) => {
5778 assert_eq!((index, batch), (4, 4));
5779 }
5780 other => panic!("expected BatchIndexOutOfBounds, got {other:?}"),
5781 }
5782 }
5783
5784 #[test]
5785 fn batch_zero_on_unit_n_is_whole() {
5786 // N == 1: batch(0) is the whole per-element block at offset 0 (no plane_offset).
5787 let parent = Tensor::<u8>::new(&[1, 2, 2, 3], Some(TensorMemory::Mem), None).unwrap();
5788 let e = parent.batch(0).unwrap();
5789 assert_eq!(e.shape(), &[2, 2, 3]);
5790 assert_eq!(e.plane_offset(), None);
5791 assert_eq!(e.buffer_identity().id(), parent.buffer_identity().id());
5792 }
5793
5794 #[test]
5795 fn mem_subview_rejects_unaligned_offset() {
5796 // f32 has align 4; a byte offset of 2 cannot back a valid `*const f32`.
5797 let parent = Tensor::<f32>::new(&[8], Some(TensorMemory::Mem), None).unwrap();
5798 assert!(parent.subview(2, &[1]).is_err());
5799 // A correctly aligned offset is accepted.
5800 assert!(parent.subview(4, &[1]).is_ok());
5801 }
5802
5803 #[test]
5804 fn mem_subview_rejects_out_of_bounds() {
5805 let parent = Tensor::<u8>::new(&[8], Some(TensorMemory::Mem), None).unwrap();
5806 // offset 6 + 4 bytes = 10 exceeds the 8-byte allocation.
5807 assert!(parent.subview(6, &[4]).is_err());
5808 }
5809
5810 /// Regression guard for the `TensorTrait::view` promotion (R2): a `subview`
5811 /// must share the parent's `BufferIdentity` on **every** backend, not mint a
5812 /// fresh one. Identity-keyed caches (the GL EGLImage import) rely on this to
5813 /// treat offset-distinct windows of one buffer as a single import; a fresh
5814 /// identity would silently break that and regress zero-copy import reuse.
5815 ///
5816 /// Runs each backend that can be allocated without a GPU/GL context on the
5817 /// test host: `Mem` (always); `Shm` (when POSIX shm is available); the
5818 /// platform-native zero-copy buffer `Dma` (DMA-BUF on Linux / IOSurface on
5819 /// macOS, when available). `Pbo` shares its identity the same way (see
5820 /// `pbo.rs` `view`) but needs a live GL context, so it is exercised by the
5821 /// image-crate GL tests rather than here.
5822 #[test]
5823 fn subview_shares_buffer_identity_all_backends() {
5824 // u8 has align 1, so every byte offset is valid for the alignment check;
5825 // this isolates the identity-sharing contract from alignment concerns.
5826 let assert_shares = |memory: TensorMemory, label: &str| {
5827 let parent = Tensor::<u8>::new(&[64], Some(memory), None)
5828 .unwrap_or_else(|e| panic!("{label}: parent alloc failed: {e:?}"));
5829 let parent_id = parent.buffer_identity().id();
5830 // Two offset-distinct windows must both carry the parent's identity.
5831 let v0 = parent
5832 .subview(0, &[16])
5833 .unwrap_or_else(|e| panic!("{label}: subview(0) failed: {e:?}"));
5834 let v1 = parent
5835 .subview(16, &[16])
5836 .unwrap_or_else(|e| panic!("{label}: subview(16) failed: {e:?}"));
5837 assert_eq!(
5838 v0.buffer_identity().id(),
5839 parent_id,
5840 "{label}: subview(0) minted a fresh BufferIdentity"
5841 );
5842 assert_eq!(
5843 v1.buffer_identity().id(),
5844 parent_id,
5845 "{label}: subview(16) minted a fresh BufferIdentity"
5846 );
5847 };
5848
5849 assert_shares(TensorMemory::Mem, "Mem");
5850
5851 #[cfg(unix)]
5852 if crate::is_shm_available() {
5853 assert_shares(TensorMemory::Shm, "Shm");
5854 }
5855
5856 // Dma == DMA-BUF on Linux, IOSurface on macOS; same public variant.
5857 if crate::is_gpu_buffer_available() {
5858 assert_shares(TensorMemory::Dma, "Dma");
5859 }
5860 }
5861
5862 #[test]
5863 fn mem_subview_four_views_no_aliasing() {
5864 // One [4,3] f32 parent; four [1,3] views at 12-byte strides, each
5865 // written independently. Exercises a multi-byte element type (offsets
5866 // must stay element-aligned) and N-way zero-copy sharing.
5867 let parent = Tensor::<f32>::new(&[4, 3], Some(TensorMemory::Mem), None).unwrap();
5868 let frame = 3 * std::mem::size_of::<f32>();
5869 for i in 0..4 {
5870 let v = parent.subview(i * frame, &[1, 3]).unwrap();
5871 let val = i as f32 + 1.0;
5872 v.map()
5873 .unwrap()
5874 .as_mut_slice()
5875 .copy_from_slice(&[val, val, val]);
5876 }
5877 assert_eq!(
5878 parent.map().unwrap().as_slice(),
5879 &[1.0, 1.0, 1.0, 2.0, 2.0, 2.0, 3.0, 3.0, 3.0, 4.0, 4.0, 4.0]
5880 );
5881 }
5882
5883 #[test]
5884 fn mem_subview_inherits_format_and_row_stride() {
5885 // A sub-view is a ready-to-use sub-image: it inherits the parent's
5886 // pixel format and (crucially) its padded row stride, so a strided
5887 // parent yields strided windows. Set a stride wider than the tight row
5888 // to exercise the row_stride inheritance path specifically.
5889 let mut parent = Tensor::<u8>::image(
5890 100,
5891 100,
5892 PixelFormat::Rgba,
5893 Some(TensorMemory::Mem),
5894 crate::CpuAccess::ReadWrite,
5895 )
5896 .unwrap();
5897 parent.set_row_stride_unchecked(512); // padded stride (> 100*4)
5898 let view = parent.subview(4096, &[10, 10, 4]).unwrap();
5899 assert_eq!(view.format(), Some(PixelFormat::Rgba), "format inherited");
5900 assert_eq!(view.row_stride(), Some(512), "row_stride inherited");
5901 }
5902
5903 #[test]
5904 fn mem_strided_subview_maps_offset_and_byte_size() {
5905 // Integration of the sub-region offset (PR #89) and the strided-map
5906 // `byte_size_override` (PR #90): a strided sub-view exposes its full
5907 // padded window (`row_stride × rows`) starting at the view's byte
5908 // offset, mapped zero-copy into the parent.
5909 let parent = Tensor::<u8>::new(&[2048], Some(TensorMemory::Mem), None).unwrap();
5910 let mut view = parent.subview(128, &[8, 16]).unwrap(); // 8 rows × 16 @ off 128
5911 assert_eq!(view.plane_offset(), Some(128));
5912 view.set_row_stride_unchecked(32); // padded stride (> 16)
5913
5914 {
5915 let mut m = view.map().unwrap();
5916 let s = m.as_mut_slice();
5917 // Strided map exposes the padded window: stride(32) × rows(8) = 256.
5918 assert_eq!(
5919 s.len(),
5920 256,
5921 "strided map exposes the full padded byte window"
5922 );
5923 s[0] = 0xAA; // row 0, col 0
5924 s[32] = 0xBB; // row 1, col 0 (one stride in)
5925 }
5926
5927 // Zero-copy: the writes land in the parent at the view's offset.
5928 let p = parent.map().unwrap();
5929 let pb = p.as_slice();
5930 assert_eq!(pb[128], 0xAA, "row 0 writes at parent offset 128");
5931 assert_eq!(
5932 pb[128 + 32],
5933 0xBB,
5934 "row 1 writes at parent offset 128 + stride"
5935 );
5936 }
5937
5938 #[test]
5939 #[cfg(unix)]
5940 fn shm_subview_partitions_parent_buffer() {
5941 // Mirrors `mem_subview_partitions_parent_buffer` for Shm: one [2,4] u8
5942 // parent shared segment (8 bytes); two [1,4] sub-views at byte offsets 0
5943 // and 4 must share the segment (zero-copy, via cloned fd) and be
5944 // independently writable — view 0 owns [0,4), view 1 owns [4,8).
5945 if !crate::is_shm_available() {
5946 eprintln!("SKIPPED: shm not available");
5947 return;
5948 }
5949 let parent = Tensor::<u8>::new(&[2, 4], Some(TensorMemory::Shm), None).unwrap();
5950 let view0 = parent.subview(0, &[1, 4]).expect("shm subview at offset 0");
5951 let view1 = parent.subview(4, &[1, 4]).expect("shm subview at offset 4");
5952
5953 view1
5954 .map()
5955 .unwrap()
5956 .as_mut_slice()
5957 .copy_from_slice(&[10, 20, 30, 40]);
5958 view0
5959 .map()
5960 .unwrap()
5961 .as_mut_slice()
5962 .copy_from_slice(&[1, 2, 3, 4]);
5963
5964 assert_eq!(view0.map().unwrap().as_slice(), &[1, 2, 3, 4]);
5965 assert_eq!(view1.map().unwrap().as_slice(), &[10, 20, 30, 40]);
5966 // The parent sees the full partitioned segment (shared, zero-copy).
5967 assert_eq!(
5968 parent.map().unwrap().as_slice(),
5969 &[1, 2, 3, 4, 10, 20, 30, 40]
5970 );
5971 // A sub-view of a sub-view composes the offset.
5972 let nested = view1.subview(2, &[1, 2]).expect("nested shm subview");
5973 assert_eq!(nested.map().unwrap().as_slice(), &[30, 40]);
5974 }
5975
5976 #[test]
5977 #[cfg(unix)]
5978 fn shm_subview_rejects_unaligned_and_oob() {
5979 if !crate::is_shm_available() {
5980 eprintln!("SKIPPED: shm not available");
5981 return;
5982 }
5983 // f32 align 4: a 2-byte offset cannot back a valid `*const f32`.
5984 let parent = Tensor::<f32>::new(&[8], Some(TensorMemory::Shm), None).unwrap();
5985 assert!(parent.subview(2, &[1]).is_err());
5986 assert!(parent.subview(4, &[1]).is_ok());
5987 // Out of bounds: offset 6 + 4 bytes = 10 > 8-byte (u8) segment.
5988 let p2 = Tensor::<u8>::new(&[8], Some(TensorMemory::Shm), None).unwrap();
5989 assert!(p2.subview(6, &[4]).is_err());
5990 }
5991
5992 #[test]
5993 #[cfg(target_os = "linux")]
5994 fn dma_subview_matches_mem_subview() {
5995 // Serialize against the fd-leak tests: this test opens DMA fds (alloc +
5996 // clone_fd), which would otherwise perturb their fd counts.
5997 let _lock = FD_LOCK.read().unwrap();
5998 // Identical sub-view semantics across Dma (shared fd) and Mem (shared
5999 // Arc): same offsets → same logical windows → same partition.
6000 let dma = match Tensor::<u8>::new(&[8], Some(TensorMemory::Dma), None) {
6001 Ok(t) => t,
6002 Err(_) => {
6003 eprintln!("SKIPPED: DMA not available");
6004 return;
6005 }
6006 };
6007 let mem = Tensor::<u8>::new(&[8], Some(TensorMemory::Mem), None).unwrap();
6008 for parent in [&dma, &mem] {
6009 let v0 = parent.subview(0, &[4]).unwrap();
6010 let v1 = parent.subview(4, &[4]).unwrap();
6011 v0.map()
6012 .unwrap()
6013 .as_mut_slice()
6014 .copy_from_slice(&[1, 2, 3, 4]);
6015 v1.map()
6016 .unwrap()
6017 .as_mut_slice()
6018 .copy_from_slice(&[5, 6, 7, 8]);
6019 assert_eq!(parent.map().unwrap().as_slice(), &[1, 2, 3, 4, 5, 6, 7, 8]);
6020 }
6021 }
6022
6023 #[test]
6024 #[cfg(target_os = "linux")]
6025 fn dma_strided_subview_maps_padded_window() {
6026 // The strided-map path differs by backing: DMA maps through
6027 // `mmap_offset` + the `byte_size_override`, not the Mem `Arc` slice. A
6028 // padded sub-view of a DMA buffer must still expose its full
6029 // `row_stride × rows` window zero-copy at the view's offset (the GPU
6030 // batched-render-to-DMA case). Mirrors
6031 // `mem_strided_subview_maps_offset_and_byte_size` on a Dma parent.
6032 let _lock = FD_LOCK.read().unwrap();
6033 let parent = match Tensor::<u8>::new(&[2048], Some(TensorMemory::Dma), None) {
6034 Ok(t) => t,
6035 Err(_) => {
6036 eprintln!("SKIPPED: DMA not available");
6037 return;
6038 }
6039 };
6040 let mut view = parent.subview(128, &[8, 16]).unwrap();
6041 assert_eq!(view.plane_offset(), Some(128));
6042 view.set_row_stride_unchecked(32); // padded stride (> 16)
6043
6044 {
6045 let mut m = view.map().unwrap();
6046 let s = m.as_mut_slice();
6047 assert_eq!(s.len(), 256, "strided DMA map exposes stride(32) × rows(8)");
6048 s[0] = 0xAA; // row 0, col 0
6049 s[32] = 0xBB; // row 1, col 0 (one stride in)
6050 }
6051
6052 let p = parent.map().unwrap();
6053 let pb = p.as_slice();
6054 assert_eq!(pb[128], 0xAA, "row 0 writes at parent offset 128");
6055 assert_eq!(
6056 pb[128 + 32],
6057 0xBB,
6058 "row 1 writes at parent offset 128 + stride"
6059 );
6060 }
6061
6062 #[test]
6063 #[cfg(target_os = "linux")]
6064 fn view_single_row_snapshots_parent_stride() {
6065 // A single-row `view()` keeps a TIGHT `row_stride` for map-span safety,
6066 // but its `view_origin` snapshots the PARENT row stride — the GL backend
6067 // keys its EGLImage import/pitch on that snapshot (not the view's tight
6068 // stride), so single-row and multi-row sibling views collapse onto the
6069 // same parent import.
6070 let _lock = FD_LOCK.read().unwrap();
6071 // 8x4 RGBA with a padded 64-byte row stride (tight row = 8*4 = 32).
6072 let parent = match Tensor::<u8>::image_with_stride(
6073 8,
6074 4,
6075 PixelFormat::Rgba,
6076 64,
6077 Some(TensorMemory::Dma),
6078 crate::CpuAccess::ReadWrite,
6079 ) {
6080 Ok(t) => t,
6081 Err(_) => {
6082 eprintln!("SKIPPED: DMA not available");
6083 return;
6084 }
6085 };
6086 assert_eq!(parent.effective_row_stride(), Some(64));
6087 // Bottom row (y=3) at x>0 — the case the tight single-row stride guards.
6088 let row = parent.view(Region::new(2, 3, 4, 1)).unwrap();
6089 // The view's own stride is tight (4*4 = 16) so its strided map stays in
6090 // bounds; the GL-facing parent pitch (64) lives in `view_origin`.
6091 assert_eq!(row.effective_row_stride(), Some(16));
6092 let vo = row.view_origin().expect("a view carries a view_origin");
6093 assert_eq!(
6094 vo.parent_row_stride, 64,
6095 "GL keys/pitches a view on the parent stride, not its tight one"
6096 );
6097 // The tight stride keeps map() in-bounds for the bottom / x>0 single row.
6098 assert_eq!(row.map().unwrap().as_slice().len(), 16);
6099 }
6100
6101 #[test]
6102 fn test_mem_tensor() {
6103 let shape = vec![2, 3, 4];
6104 let tensor =
6105 MemTensor::<f32>::new(&shape, Some("test_tensor")).expect("Failed to create tensor");
6106 assert_eq!(tensor.shape(), &shape);
6107 assert_eq!(tensor.size(), 2 * 3 * 4 * std::mem::size_of::<f32>());
6108 assert_eq!(tensor.name(), "test_tensor");
6109
6110 {
6111 let mut tensor_map = tensor.map().expect("Failed to map memory");
6112 tensor_map.fill(42.0);
6113 assert!(tensor_map.iter().all(|&x| x == 42.0));
6114 }
6115
6116 let mut tensor = MemTensor::<u8>::new(&shape, None).expect("Failed to create tensor");
6117 assert_eq!(tensor.shape(), &shape);
6118 let new_shape = vec![3, 4, 4];
6119 assert!(
6120 tensor.reshape(&new_shape).is_err(),
6121 "Reshape should fail due to size mismatch"
6122 );
6123 assert_eq!(tensor.shape(), &shape, "Shape should remain unchanged");
6124
6125 let new_shape = vec![2, 3, 4];
6126 tensor.reshape(&new_shape).expect("Reshape should succeed");
6127 assert_eq!(
6128 tensor.shape(),
6129 &new_shape,
6130 "Shape should be updated after successful reshape"
6131 );
6132
6133 {
6134 let mut tensor_map = tensor.map().expect("Failed to map memory");
6135 tensor_map.fill(1);
6136 assert!(tensor_map.iter().all(|&x| x == 1));
6137 }
6138
6139 {
6140 let mut tensor_map = tensor.map().expect("Failed to map memory");
6141 tensor_map[2] = 42;
6142 assert_eq!(tensor_map[1], 1, "Value at index 1 should be 1");
6143 assert_eq!(tensor_map[2], 42, "Value at index 2 should be 42");
6144 }
6145 }
6146
6147 #[test]
6148 #[cfg(target_os = "linux")]
6149 fn test_dma_no_fd_leaks() {
6150 let _lock = FD_LOCK.write().unwrap();
6151 if !is_dma_available() {
6152 log::warn!(
6153 "SKIPPED: {} - DMA memory allocation not available (permission denied or no DMA-BUF support)",
6154 function!()
6155 );
6156 return;
6157 }
6158
6159 let proc = procfs::process::Process::myself()
6160 .expect("Failed to get current process using /proc/self");
6161
6162 let start_open_fds = proc
6163 .fd_count()
6164 .expect("Failed to get open file descriptor count");
6165
6166 for _ in 0..100 {
6167 let tensor = Tensor::<u8>::new(&[100, 100], Some(TensorMemory::Dma), None)
6168 .expect("Failed to create tensor");
6169 let mut map = tensor.map().unwrap();
6170 map.as_mut_slice().fill(233);
6171 }
6172
6173 let end_open_fds = proc
6174 .fd_count()
6175 .expect("Failed to get open file descriptor count");
6176
6177 assert_eq!(
6178 start_open_fds, end_open_fds,
6179 "File descriptor leak detected: {} -> {}",
6180 start_open_fds, end_open_fds
6181 );
6182 }
6183
6184 #[test]
6185 #[cfg(target_os = "linux")]
6186 fn test_dma_from_fd_no_fd_leaks() {
6187 let _lock = FD_LOCK.write().unwrap();
6188 if !is_dma_available() {
6189 log::warn!(
6190 "SKIPPED: {} - DMA memory allocation not available (permission denied or no DMA-BUF support)",
6191 function!()
6192 );
6193 return;
6194 }
6195
6196 let proc = procfs::process::Process::myself()
6197 .expect("Failed to get current process using /proc/self");
6198
6199 let start_open_fds = proc
6200 .fd_count()
6201 .expect("Failed to get open file descriptor count");
6202
6203 let orig = Tensor::<u8>::new(&[100, 100], Some(TensorMemory::Dma), None).unwrap();
6204
6205 for _ in 0..100 {
6206 let tensor =
6207 Tensor::<u8>::from_fd(orig.clone_fd().unwrap(), orig.shape(), None).unwrap();
6208 assert_eq!(
6209 tensor.memory(),
6210 TensorMemory::Dma,
6211 "DMA-BUF fd must import as Dma, not be silently downgraded"
6212 );
6213 let mut map = tensor.map().unwrap();
6214 map.as_mut_slice().fill(233);
6215 }
6216 drop(orig);
6217
6218 let end_open_fds = proc.fd_count().unwrap();
6219
6220 assert_eq!(
6221 start_open_fds, end_open_fds,
6222 "File descriptor leak detected: {} -> {}",
6223 start_open_fds, end_open_fds
6224 );
6225 }
6226
6227 /// A filesystem magic must report its true 32-bit value regardless of
6228 /// how wide, and how signed, `fstatfs`'s `f_type` is on the target.
6229 ///
6230 /// `fs_type_t` is `__fsword_t` on Linux/gnu — `i64` on 64-bit but `i32`
6231 /// on 32-bit (armv7, i686, aarch64-ilp32) — `c_int` on uclibc, and
6232 /// unsigned on musl and s390x. Widening a *signed* 32-bit `f_type`
6233 /// sign-extends every magic with bit 31 set, so a naive widening cast
6234 /// reports e.g. `0xffffffff958458f6` for hugetlbfs. The docs tell
6235 /// callers to look the reported value up in `include/uapi/linux/magic.h`,
6236 /// so a sign-extended value is actively misleading.
6237 ///
6238 /// This cannot be reproduced on a 64-bit host — the defect only exists
6239 /// where `f_type` is a signed 32-bit type — so the test drives the
6240 /// conversion directly with the value such a target would produce.
6241 #[test]
6242 #[cfg(target_os = "linux")]
6243 fn test_fs_magic_normalizes_sign_extended_values() {
6244 // Magics whose bit 31 is set. HUGETLBFS is the one that matters
6245 // most in practice: a MFD_HUGETLB memfd is the likeliest fd to land
6246 // in the UnknownBufferType arm.
6247 const HUGETLBFS_MAGIC: u32 = 0x9584_58f6;
6248 const F2FS_MAGIC: u32 = 0xf2f5_2010;
6249 const BTRFS_MAGIC: u32 = 0x9123_683e;
6250
6251 for magic in [HUGETLBFS_MAGIC, F2FS_MAGIC, BTRFS_MAGIC] {
6252 // 64-bit gnu: f_type is i64 and already holds the true value.
6253 assert_eq!(fs_magic(i64::from(magic)), magic);
6254
6255 // 32-bit gnu / uclibc: f_type is i32, so the value arrives
6256 // sign-extended once widened. It must still report its true
6257 // 32 bits.
6258 let sign_extended = i64::from(magic as i32);
6259 assert!(sign_extended < 0, "{magic:#x} should have bit 31 set");
6260 assert_eq!(
6261 fs_magic(sign_extended),
6262 magic,
6263 "{magic:#x} must survive a signed 32-bit f_type"
6264 );
6265 }
6266
6267 // The two magics we actually classify on are below 2^31, so they are
6268 // unaffected by signedness either way — this is why the DMA-vs-SHM
6269 // decision is correct on 32-bit even without this normalization.
6270 for magic in [DMA_BUF_MAGIC, TMPFS_MAGIC] {
6271 assert!(magic < 0x8000_0000);
6272 assert_eq!(fs_magic(i64::from(magic as i32)), magic);
6273 }
6274 }
6275
6276 /// A DMA-BUF fd must import as [`TensorMemory::Dma`].
6277 ///
6278 /// Regression test for the `st_dev` minor-number classifier, which
6279 /// hardcoded `9 | 10` as "this is DMA". Those minors come from
6280 /// `get_anon_bdev()` and are assigned first-come-first-served at boot,
6281 /// so they vary by kernel build and boot order — a real DMA-BUF is
6282 /// minor 12 on x86 desktop and minor 8 on the ADIS Verdin. Both fell
6283 /// into the `_` arm and imported as SHM, which "works" (a DMA-BUF is
6284 /// mmap-able) but silently forfeits zero-copy.
6285 #[test]
6286 #[cfg(target_os = "linux")]
6287 fn test_from_fd_dma_imports_as_dma() {
6288 let _lock = FD_LOCK.read().unwrap();
6289 if !is_dma_available() {
6290 log::warn!("SKIPPED: {} - DMA memory not available", function!());
6291 return;
6292 }
6293
6294 let orig = Tensor::<u8>::new(&[64, 64], Some(TensorMemory::Dma), None).unwrap();
6295 assert_eq!(orig.memory(), TensorMemory::Dma);
6296
6297 let imported = Tensor::<u8>::from_fd(orig.clone_fd().unwrap(), orig.shape(), None).unwrap();
6298
6299 assert_eq!(
6300 imported.memory(),
6301 TensorMemory::Dma,
6302 "a DMA-BUF fd must import as Dma"
6303 );
6304 }
6305
6306 /// A tmpfs/SHM fd must import as [`TensorMemory::Shm`].
6307 ///
6308 /// The companion to `test_from_fd_dma_imports_as_dma`: confirms the
6309 /// magic-based classifier identifies SHM positively rather than by
6310 /// falling through.
6311 #[test]
6312 #[cfg(target_os = "linux")]
6313 fn test_from_fd_shm_imports_as_shm() {
6314 let _lock = FD_LOCK.read().unwrap();
6315 if !is_shm_available() {
6316 log::warn!("SKIPPED: {} - SHM memory not available", function!());
6317 return;
6318 }
6319
6320 let orig = Tensor::<u8>::new(&[64, 64], Some(TensorMemory::Shm), None).unwrap();
6321 assert_eq!(orig.memory(), TensorMemory::Shm);
6322
6323 let imported = Tensor::<u8>::from_fd(orig.clone_fd().unwrap(), orig.shape(), None).unwrap();
6324
6325 assert_eq!(
6326 imported.memory(),
6327 TensorMemory::Shm,
6328 "a tmpfs fd must import as Shm"
6329 );
6330 }
6331
6332 /// An fd that is neither a DMA-BUF nor tmpfs must be rejected.
6333 ///
6334 /// A pipe is the convenient probe: it lives on `pipefs`, another
6335 /// `get_anon_bdev()` pseudo-filesystem, so it shares major 0 with the
6336 /// buffer types we do support and is only distinguishable by magic.
6337 /// Importing one as SHM is meaningless — `mmap` on a pipe fails — so
6338 /// the classifier must say "unknown" rather than guess.
6339 #[test]
6340 #[cfg(target_os = "linux")]
6341 fn test_from_fd_rejects_unknown_filesystem() {
6342 let _lock = FD_LOCK.read().unwrap();
6343
6344 let (read_end, _write_end) = nix::unistd::pipe().unwrap();
6345
6346 let result = Tensor::<u8>::from_fd(read_end, &[64], None);
6347
6348 match result {
6349 Err(Error::UnknownBufferType(magic)) => {
6350 // PIPEFS_MAGIC, from include/uapi/linux/magic.h
6351 assert_eq!(magic, 0x5049_5045, "expected PIPEFS_MAGIC");
6352 }
6353 other => panic!("expected UnknownBufferType for a pipe fd, got {other:?}"),
6354 }
6355 }
6356
6357 #[test]
6358 #[cfg(target_os = "linux")]
6359 fn test_shm_no_fd_leaks() {
6360 let _lock = FD_LOCK.write().unwrap();
6361 if !is_shm_available() {
6362 log::warn!(
6363 "SKIPPED: {} - SHM memory allocation not available (permission denied or no SHM support)",
6364 function!()
6365 );
6366 return;
6367 }
6368
6369 let proc = procfs::process::Process::myself()
6370 .expect("Failed to get current process using /proc/self");
6371
6372 let start_open_fds = proc
6373 .fd_count()
6374 .expect("Failed to get open file descriptor count");
6375
6376 for _ in 0..100 {
6377 let tensor = Tensor::<u8>::new(&[100, 100], Some(TensorMemory::Shm), None)
6378 .expect("Failed to create tensor");
6379 let mut map = tensor.map().unwrap();
6380 map.as_mut_slice().fill(233);
6381 }
6382
6383 let end_open_fds = proc
6384 .fd_count()
6385 .expect("Failed to get open file descriptor count");
6386
6387 assert_eq!(
6388 start_open_fds, end_open_fds,
6389 "File descriptor leak detected: {} -> {}",
6390 start_open_fds, end_open_fds
6391 );
6392 }
6393
6394 #[test]
6395 #[cfg(target_os = "linux")]
6396 fn test_shm_from_fd_no_fd_leaks() {
6397 let _lock = FD_LOCK.write().unwrap();
6398 if !is_shm_available() {
6399 log::warn!(
6400 "SKIPPED: {} - SHM memory allocation not available (permission denied or no SHM support)",
6401 function!()
6402 );
6403 return;
6404 }
6405
6406 let proc = procfs::process::Process::myself()
6407 .expect("Failed to get current process using /proc/self");
6408
6409 let start_open_fds = proc
6410 .fd_count()
6411 .expect("Failed to get open file descriptor count");
6412
6413 let orig = Tensor::<u8>::new(&[100, 100], Some(TensorMemory::Shm), None).unwrap();
6414
6415 for _ in 0..100 {
6416 let tensor =
6417 Tensor::<u8>::from_fd(orig.clone_fd().unwrap(), orig.shape(), None).unwrap();
6418 let mut map = tensor.map().unwrap();
6419 map.as_mut_slice().fill(233);
6420 }
6421 drop(orig);
6422
6423 let end_open_fds = proc.fd_count().unwrap();
6424
6425 assert_eq!(
6426 start_open_fds, end_open_fds,
6427 "File descriptor leak detected: {} -> {}",
6428 start_open_fds, end_open_fds
6429 );
6430 }
6431
6432 #[cfg(feature = "ndarray")]
6433 #[test]
6434 fn test_ndarray() {
6435 let _lock = FD_LOCK.read().unwrap();
6436 let shape = vec![2, 3, 4];
6437 let tensor = Tensor::<f32>::new(&shape, None, None).expect("Failed to create tensor");
6438
6439 let mut tensor_map = tensor.map().expect("Failed to map tensor memory");
6440 tensor_map.fill(1.0);
6441
6442 let view = tensor_map.view().expect("Failed to get ndarray view");
6443 assert_eq!(view.shape(), &[2, 3, 4]);
6444 assert!(view.iter().all(|&x| x == 1.0));
6445
6446 let mut view_mut = tensor_map
6447 .view_mut()
6448 .expect("Failed to get mutable ndarray view");
6449 view_mut[[0, 0, 0]] = 42.0;
6450 assert_eq!(view_mut[[0, 0, 0]], 42.0);
6451 assert_eq!(tensor_map[0], 42.0, "Value at index 0 should be 42");
6452 }
6453
6454 #[test]
6455 fn test_buffer_identity_unique() {
6456 let id1 = BufferIdentity::new();
6457 let id2 = BufferIdentity::new();
6458 assert_ne!(
6459 id1.id(),
6460 id2.id(),
6461 "Two identities should have different ids"
6462 );
6463 }
6464
6465 #[test]
6466 fn test_buffer_identity_clone_shares_guard() {
6467 let id1 = BufferIdentity::new();
6468 let weak = id1.weak();
6469 assert!(
6470 weak.upgrade().is_some(),
6471 "Weak should be alive while original exists"
6472 );
6473
6474 let id2 = id1.clone();
6475 assert_eq!(id1.id(), id2.id(), "Cloned identity should have same id");
6476
6477 drop(id1);
6478 assert!(
6479 weak.upgrade().is_some(),
6480 "Weak should still be alive (clone holds Arc)"
6481 );
6482
6483 drop(id2);
6484 assert!(
6485 weak.upgrade().is_none(),
6486 "Weak should be dead after all clones dropped"
6487 );
6488 }
6489
6490 #[test]
6491 fn test_tensor_buffer_identity() {
6492 let t1 = Tensor::<u8>::new(&[100], Some(TensorMemory::Mem), Some("t1")).unwrap();
6493 let t2 = Tensor::<u8>::new(&[100], Some(TensorMemory::Mem), Some("t2")).unwrap();
6494 assert_ne!(
6495 t1.buffer_identity().id(),
6496 t2.buffer_identity().id(),
6497 "Different tensors should have different buffer ids"
6498 );
6499 }
6500
6501 // ------------------------------------------------------------------------
6502 // Quantization — constructor validation + accessor correctness.
6503 // ------------------------------------------------------------------------
6504
6505 #[test]
6506 fn test_quantization_per_tensor_constructors() {
6507 let q = Quantization::per_tensor(0.1, -5);
6508 assert!(q.is_per_tensor());
6509 assert!(!q.is_per_channel());
6510 assert!(!q.is_symmetric());
6511 assert_eq!(q.scale(), &[0.1]);
6512 assert_eq!(q.zero_point(), Some(&[-5][..]));
6513
6514 let qs = Quantization::per_tensor_symmetric(0.05);
6515 assert!(qs.is_per_tensor());
6516 assert!(qs.is_symmetric());
6517 assert_eq!(qs.zero_point(), None);
6518 }
6519
6520 #[test]
6521 fn test_quantization_per_channel_constructors() {
6522 let q = Quantization::per_channel(vec![0.1, 0.2, 0.3], vec![0, -1, 1], 2).unwrap();
6523 assert!(q.is_per_channel());
6524 assert!(!q.is_symmetric());
6525 assert_eq!(q.axis(), Some(2));
6526 assert_eq!(q.scale().len(), 3);
6527
6528 let qs = Quantization::per_channel_symmetric(vec![0.054, 0.089, 0.195], 0).unwrap();
6529 assert!(qs.is_per_channel());
6530 assert!(qs.is_symmetric());
6531 assert_eq!(qs.axis(), Some(0));
6532 }
6533
6534 #[test]
6535 fn test_quantization_per_channel_length_mismatch_rejected() {
6536 // len(scales) != len(zero_points) → rejected at construction.
6537 let err = Quantization::per_channel(vec![0.1, 0.2], vec![0, 0, 0], 0).unwrap_err();
6538 assert!(matches!(err, Error::QuantizationInvalid { .. }));
6539 }
6540
6541 #[test]
6542 fn test_quantization_per_channel_empty_rejected() {
6543 let err = Quantization::per_channel_symmetric(vec![], 0).unwrap_err();
6544 assert!(matches!(err, Error::QuantizationInvalid { .. }));
6545 }
6546
6547 /// Constructors guard scale/zero_point length invariants, but
6548 /// `Quantization` is `Deserialize`, so malformed JSON (e.g. an
6549 /// empty `scale` array, or `zero_point` length that disagrees with
6550 /// `scale`) bypasses the constructor checks. `set_quantization`
6551 /// must reject these via `validate()` so they don't poison
6552 /// downstream `mode()` selection or per-channel kernel indexing.
6553 #[test]
6554 fn test_quantization_validate_rejects_malformed_deserialize() {
6555 let mut t = Tensor::<i8>::new(&[1, 1, 4], Some(TensorMemory::Mem), None).unwrap();
6556
6557 // Empty scale array: must be rejected.
6558 let q: Quantization = serde_json::from_str(r#"{"scale": []}"#).unwrap();
6559 assert!(matches!(
6560 t.set_quantization(q).unwrap_err(),
6561 Error::QuantizationInvalid { .. }
6562 ));
6563
6564 // Per-tensor with multi-element zero_point: must be rejected.
6565 let q: Quantization =
6566 serde_json::from_str(r#"{"scale": 0.1, "zero_point": [0, 0, 0]}"#).unwrap();
6567 assert!(matches!(
6568 t.set_quantization(q).unwrap_err(),
6569 Error::QuantizationInvalid { .. }
6570 ));
6571
6572 // Per-channel zero_point length != scale length: must be rejected.
6573 let q: Quantization = serde_json::from_str(
6574 r#"{"scale": [0.1, 0.2, 0.3, 0.4], "zero_point": [0, 0], "axis": 2}"#,
6575 )
6576 .unwrap();
6577 assert!(matches!(
6578 t.set_quantization(q).unwrap_err(),
6579 Error::QuantizationInvalid { .. }
6580 ));
6581 }
6582
6583 #[test]
6584 fn test_quantization_mode_dispatch() {
6585 let pt = Quantization::per_tensor(0.1, -5);
6586 assert!(matches!(
6587 pt.mode(),
6588 QuantMode::PerTensor { scale, zero_point } if scale == 0.1 && zero_point == -5
6589 ));
6590
6591 let pts = Quantization::per_tensor_symmetric(0.05);
6592 assert!(matches!(
6593 pts.mode(),
6594 QuantMode::PerTensorSymmetric { scale } if scale == 0.05
6595 ));
6596
6597 let pc = Quantization::per_channel(vec![0.1, 0.2], vec![0, -1], 2).unwrap();
6598 assert!(matches!(pc.mode(), QuantMode::PerChannel { axis: 2, .. }));
6599
6600 let pcs = Quantization::per_channel_symmetric(vec![0.1, 0.2], 0).unwrap();
6601 assert!(matches!(
6602 pcs.mode(),
6603 QuantMode::PerChannelSymmetric { axis: 0, .. }
6604 ));
6605 }
6606
6607 #[test]
6608 fn test_tensor_quantization_roundtrip_integer() {
6609 let mut t = Tensor::<i8>::new(&[2, 3, 4], Some(TensorMemory::Mem), None).unwrap();
6610 assert!(t.quantization().is_none());
6611 t.set_quantization(Quantization::per_tensor(0.1, -5))
6612 .unwrap();
6613 let q = t.quantization().unwrap();
6614 assert_eq!(q.scale(), &[0.1]);
6615 t.clear_quantization();
6616 assert!(t.quantization().is_none());
6617 }
6618
6619 #[test]
6620 fn test_tensor_with_quantization_builder() {
6621 let t = Tensor::<i8>::new(&[4, 4], Some(TensorMemory::Mem), None)
6622 .unwrap()
6623 .with_quantization(Quantization::per_tensor_symmetric(0.05))
6624 .unwrap();
6625 assert!(t.quantization().is_some());
6626 }
6627
6628 #[test]
6629 fn test_tensor_dyn_quantization_float_arm_returns_none() {
6630 let t = Tensor::<f32>::new(&[2, 2], Some(TensorMemory::Mem), None).unwrap();
6631 let td = TensorDyn::F32(t);
6632 assert!(td.quantization().is_none());
6633 }
6634
6635 #[test]
6636 fn test_tensor_dyn_set_quantization_float_arm_errors() {
6637 let t = Tensor::<f32>::new(&[2, 2], Some(TensorMemory::Mem), None).unwrap();
6638 let mut td = TensorDyn::F32(t);
6639 let err = td
6640 .set_quantization(Quantization::per_tensor(0.1, 0))
6641 .unwrap_err();
6642 // float path returns a QuantizationInvalid error.
6643 assert!(matches!(err, Error::QuantizationInvalid { .. }));
6644 }
6645
6646 /// Compile-time type gate — calling `Tensor::<f32>::quantization()` must
6647 /// fail to compile (the `IntegerType` trait bound is not satisfied by
6648 /// `f32`). This doctest anchors the invariant.
6649 ///
6650 /// ```compile_fail
6651 /// use edgefirst_tensor::{Tensor, TensorMemory};
6652 /// let t = Tensor::<f32>::new(&[2, 2], Some(TensorMemory::Mem), None).unwrap();
6653 /// let _ = t.quantization(); // compile error: f32 not IntegerType
6654 /// ```
6655 fn _compile_fail_doctest_anchor() {}
6656
6657 // Any test that cares about the fd count must grab it exclusively.
6658 // Any tests which modifies the fd count by opening or closing fds must grab it
6659 // shared.
6660 pub static FD_LOCK: RwLock<()> = RwLock::new(());
6661
6662 /// Test that DMA is NOT available on non-Linux platforms.
6663 /// This verifies the cross-platform behavior of is_dma_available().
6664 #[test]
6665 #[cfg(not(target_os = "linux"))]
6666 fn test_dma_not_available_on_non_linux() {
6667 assert!(
6668 !is_dma_available(),
6669 "DMA memory allocation should NOT be available on non-Linux platforms"
6670 );
6671 }
6672
6673 #[test]
6674 fn colorimetry_defaults_none_and_roundtrips_without_auto_fill() {
6675 use crate::{ColorEncoding, ColorRange, Colorimetry, PixelFormat, TensorMemory};
6676 let mut t = Tensor::<u8>::image(
6677 1280,
6678 720,
6679 PixelFormat::Nv12,
6680 Some(TensorMemory::Mem),
6681 crate::CpuAccess::ReadWrite,
6682 )
6683 .unwrap();
6684 assert_eq!(t.colorimetry(), None); // default undefined
6685 let c = Colorimetry::default()
6686 .with_encoding(ColorEncoding::Bt709)
6687 .with_range(ColorRange::Limited);
6688 t.set_colorimetry(Some(c));
6689 assert_eq!(t.colorimetry(), Some(c));
6690 // configure_image must NOT touch colorimetry
6691 t.configure_image(640, 480, PixelFormat::Grey).unwrap();
6692 assert_eq!(t.colorimetry(), Some(c));
6693 }
6694
6695 #[test]
6696 fn configure_image_within_capacity() {
6697 let mut t = Tensor::<u8>::image_with_capacity(
6698 640,
6699 480,
6700 PixelFormat::Rgb,
6701 None,
6702 crate::CpuAccess::ReadWrite,
6703 )
6704 .unwrap();
6705 t.configure_image(320, 240, PixelFormat::Nv12).unwrap();
6706 assert_eq!(t.format(), Some(PixelFormat::Nv12));
6707 assert_eq!(t.width(), Some(320));
6708 assert_eq!(t.height(), Some(240));
6709 assert_eq!(t.shape(), &[360, 320]); // 240*3/2
6710 }
6711
6712 #[test]
6713 fn configure_image_too_large_errors() {
6714 let mut t = Tensor::<u8>::image_with_capacity(
6715 64,
6716 64,
6717 PixelFormat::Grey,
6718 None,
6719 crate::CpuAccess::ReadWrite,
6720 )
6721 .unwrap();
6722 let err = t
6723 .configure_image(1920, 1080, PixelFormat::Nv12)
6724 .unwrap_err();
6725 assert!(matches!(err, Error::InsufficientCapacity { .. }));
6726 }
6727
6728 /// A reused max-sized IOSurface pool keeps its physical `bytesPerRow` when
6729 /// reconfigured to a smaller logical image (physical-grid / logical-ROI
6730 /// decoupling), instead of collapsing to the frame's natural row stride.
6731 #[test]
6732 #[cfg(target_os = "macos")]
6733 fn configure_image_preserves_iosurface_physical_stride() {
6734 // Pool: GREY/R8 IOSurface 100 wide → bytesPerRow padded to 128.
6735 let mut pool = Tensor::<u8>::image(
6736 100,
6737 64,
6738 PixelFormat::Grey,
6739 Some(TensorMemory::Dma),
6740 crate::CpuAccess::ReadWrite,
6741 )
6742 .unwrap();
6743 let pitch = pool.effective_row_stride().unwrap();
6744 assert!(
6745 pitch >= 128 && pitch.is_multiple_of(64),
6746 "padded bytesPerRow, got {pitch}"
6747 );
6748
6749 // Reconfigure to a smaller NV12 frame; the physical pitch must survive
6750 // (natural would be 32, but the surface stride is the 128-padded pitch).
6751 pool.configure_image(32, 16, PixelFormat::Nv12).unwrap();
6752 assert_eq!(pool.format(), Some(PixelFormat::Nv12));
6753 assert_eq!(pool.width(), Some(32));
6754 assert_eq!(pool.height(), Some(16));
6755 assert_eq!(
6756 pool.effective_row_stride(),
6757 Some(pitch),
6758 "configure_image must preserve the IOSurface physical bytesPerRow"
6759 );
6760
6761 // Reconfigure again to NV24 — pitch still preserved.
6762 pool.configure_image(32, 16, PixelFormat::Nv24).unwrap();
6763 assert_eq!(pool.effective_row_stride(), Some(pitch));
6764 }
6765
6766 /// `configure_image` on a Mem backing reconfigures to the format's
6767 /// **64-byte-aligned** row stride (the odd-dim contract: every image tensor
6768 /// carries a 64-aligned `row_stride`). For NV12 32×16 the minimum is
6769 /// `even(32)=32`, rounded up to the 64-byte alignment → 64. The capacity
6770 /// (64×64×4 RGBA = 16 KiB) easily holds the 24×64 = 1.5 KiB NV12 layout.
6771 #[test]
6772 fn configure_image_mem_aligns_stride() {
6773 let mut t = Tensor::<u8>::image_with_capacity(
6774 64,
6775 64,
6776 PixelFormat::Rgba,
6777 Some(TensorMemory::Mem),
6778 crate::CpuAccess::ReadWrite,
6779 )
6780 .unwrap();
6781 t.configure_image(32, 16, PixelFormat::Nv12).unwrap();
6782 let s = t.effective_row_stride().unwrap();
6783 assert_eq!(s % 64, 0, "stride must be 64-aligned");
6784 assert!(s >= 32, "stride must cover the even-width minimum");
6785 assert_eq!(s, 64);
6786 }
6787
6788 #[test]
6789 fn strided_mem_tensor_cpu_maps_full_padded_buffer() {
6790 // A packed RGBA image with row padding (GPU-pitch style): logical width
6791 // 8 px (32 B/row) but a 48-byte row stride. Over-allocate capacity (for
6792 // 16 px), narrow the logical width, then record the padded stride.
6793 // Previously `map()` rejected this on non-Linux with
6794 // "DMA backing is Linux-only"; HAL-owned Mem is now mappable.
6795 let mut t = Tensor::<u8>::image_with_capacity(
6796 16,
6797 3,
6798 PixelFormat::Rgba,
6799 Some(TensorMemory::Mem),
6800 crate::CpuAccess::ReadWrite,
6801 )
6802 .unwrap(); // capacity 3 × 16 × 4 = 192 B
6803 t.configure_image(8, 3, PixelFormat::Rgba).unwrap(); // logical [3, 8, 4] = 96 B
6804 t.set_row_stride(48).unwrap(); // padded stride (>= 32 B min)
6805
6806 let map = t.map().expect("strided Mem tensor should CPU-map");
6807 // Full padded buffer (stride 48 × 3 rows = 144 B), not the 96 B logical
6808 // view — callers iterate rows via `effective_row_stride()`.
6809 assert_eq!(map.as_slice().len(), 144);
6810 // Logical shape is still reported for shape-aware consumers.
6811 assert_eq!(map.shape(), &[3, 8, 4]);
6812 }
6813
6814 #[test]
6815 fn strided_mem_tensor_over_capacity_errors() {
6816 // Stride larger than the allocation: 64 B × 3 rows = 192 B > 96 B cap.
6817 let mut t = Tensor::<u8>::new(&[3, 8, 4], Some(TensorMemory::Mem), None).unwrap();
6818 t.set_format(PixelFormat::Rgba).unwrap();
6819 t.set_row_stride(64).unwrap();
6820 assert!(matches!(t.map(), Err(Error::InsufficientCapacity { .. })));
6821 }
6822
6823 /// Test that SHM memory allocation is available and usable on Unix systems.
6824 /// This is a basic functional test; Linux has additional FD leak tests using procfs.
6825 #[test]
6826 #[cfg(unix)]
6827 fn test_shm_available_and_usable() {
6828 assert!(
6829 is_shm_available(),
6830 "SHM memory allocation should be available on Unix systems"
6831 );
6832
6833 // Create a tensor with SHM backing
6834 let tensor = Tensor::<u8>::new(&[100, 100], Some(TensorMemory::Shm), None)
6835 .expect("Failed to create SHM tensor");
6836
6837 // Verify we can map and write to it
6838 let mut map = tensor.map().expect("Failed to map SHM tensor");
6839 map.as_mut_slice().fill(0xAB);
6840
6841 // Verify the data was written correctly
6842 assert!(
6843 map.as_slice().iter().all(|&b| b == 0xAB),
6844 "SHM tensor data should be writable and readable"
6845 );
6846 }
6847
6848 // =========================================================================
6849 // packed_rgba16f_layout — host-runnable geometry unit tests (TDD)
6850 // =========================================================================
6851
6852 #[test]
6853 fn packed_rgba16f_layout_planar_rgb_f16() {
6854 let layout =
6855 packed_rgba16f_layout(PixelFormat::PlanarRgb, DType::F16, 640, 640).expect("Some");
6856 assert_eq!(layout.surface_w, 160);
6857 assert_eq!(layout.surface_h, 1920);
6858 assert_eq!(layout.bytes_per_texel, 8);
6859 assert_eq!(layout.pitch, 1280);
6860 }
6861
6862 #[test]
6863 fn packed_rgba16f_layout_planar_rgba_f16() {
6864 let layout =
6865 packed_rgba16f_layout(PixelFormat::PlanarRgba, DType::F16, 640, 640).expect("Some");
6866 assert_eq!(layout.surface_w, 160);
6867 assert_eq!(layout.surface_h, 2560); // 4 planes
6868 assert_eq!(layout.bytes_per_texel, 8);
6869 assert_eq!(layout.pitch, 1280);
6870 }
6871
6872 #[test]
6873 fn packed_rgba16f_layout_rejects_misaligned() {
6874 assert!(packed_rgba16f_layout(PixelFormat::PlanarRgb, DType::F16, 642, 640).is_none());
6875 }
6876
6877 #[test]
6878 fn packed_rgba16f_layout_rejects_non_f16() {
6879 // Non-F16 dtype with planar RGB
6880 assert!(packed_rgba16f_layout(PixelFormat::PlanarRgb, DType::U8, 640, 640).is_none());
6881 // Non-planar format with F32
6882 assert!(packed_rgba16f_layout(PixelFormat::Rgb, DType::F32, 640, 640).is_none());
6883 // Packed Rgba with F16 is not a planar format → None
6884 assert!(packed_rgba16f_layout(PixelFormat::Rgba, DType::F16, 640, 640).is_none());
6885 }
6886
6887 #[test]
6888 fn cuda_map_fast_fails_to_none_without_handle() {
6889 let t = Tensor::<f32>::new(&[4], Some(TensorMemory::Mem), None).unwrap();
6890 assert!(t.cuda().is_none());
6891 assert!(t.cuda_map().is_none()); // pure local check, no GL routing
6892 }
6893
6894 #[test]
6895 fn cuda_returns_none_without_handle() {
6896 // A plain Mem-backed tensor has no CUDA handle attached.
6897 let t = Tensor::<f32>::new(&[2, 2], Some(TensorMemory::Mem), None).unwrap();
6898 assert!(t.cuda().is_none(), "no CUDA handle on a Mem tensor");
6899 assert!(t.cuda_map().is_none(), "fast-fail map → None");
6900 }
6901
6902 #[test]
6903 fn cuda_map_then_host_map_fallback() {
6904 // The documented client pattern: try cuda_map() first; when it is None
6905 // (no CUDA handle — the case for a plain Mem tensor), fall back to map().
6906 let t = Tensor::<f32>::new(&[2, 2], Some(TensorMemory::Mem), None).unwrap();
6907 // Bind to a named variable so the CudaMap guard (and its borrow of `t`)
6908 // is dropped at the end of this statement, before the else branch borrows `t` again.
6909 let cuda = t.cuda_map();
6910 if let Some(_c) = cuda {
6911 // On a CUDA-registered tensor we'd use the device ptr here.
6912 unreachable!("a Mem tensor has no CUDA handle");
6913 } else {
6914 let host = t.map().expect("host map fallback must succeed");
6915 // TensorMapTrait::len() returns the element count (not bytes).
6916 assert_eq!(host.len(), 4); // 2*2 f32 elements
6917 }
6918 }
6919
6920 // -------------------------------------------------------------------------
6921 // Tensor::from_foreign — public API tests at the Tensor<T> layer.
6922 //
6923 // The low-level MemTensor::from_foreign mechanics (owner-drop, view sharing)
6924 // are covered in mem.rs. These tests exercise the Tensor<T> guard paths
6925 // (null ptr, empty shape, size overflow) and the basic wrap+readback
6926 // contract, confirming the public unsafe API wires through correctly.
6927 // -------------------------------------------------------------------------
6928
6929 #[test]
6930 fn from_foreign_valid_wrap_and_readback() {
6931 // The canonical CUDA zero-copy shape: wrap a caller allocation as a
6932 // Mem tensor and verify the tensor reads the exact same bytes.
6933 let mut buf: Vec<f32> = vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0];
6934 let ptr = buf.as_mut_ptr();
6935 let t = unsafe { Tensor::<f32>::from_foreign(ptr, &[2, 3], None, Some("test_foreign")) }
6936 .expect("valid from_foreign must succeed");
6937 assert_eq!(t.shape(), &[2, 3]);
6938 assert_eq!(t.memory(), TensorMemory::Mem);
6939 assert_eq!(t.name(), "test_foreign");
6940 let m = t.map().unwrap();
6941 // The tensor is a zero-copy borrow — it sees the caller's data.
6942 assert_eq!(m.as_slice(), &[1.0_f32, 2.0, 3.0, 4.0, 5.0, 6.0]);
6943 }
6944
6945 #[test]
6946 fn from_foreign_write_visible_in_caller_allocation() {
6947 // Writes through the tensor's map land in the caller's buffer (zero-copy).
6948 let mut buf: Vec<u8> = vec![0u8; 6];
6949 let ptr = buf.as_mut_ptr();
6950 let t = unsafe { Tensor::<u8>::from_foreign(ptr, &[2, 3], None, None) }.unwrap();
6951 {
6952 let mut m = t.map().unwrap();
6953 m.as_mut_slice().copy_from_slice(&[10, 20, 30, 40, 50, 60]);
6954 }
6955 drop(t);
6956 // Mutations are visible in the original Vec — same physical buffer.
6957 assert_eq!(buf, vec![10, 20, 30, 40, 50, 60]);
6958 }
6959
6960 #[test]
6961 fn from_foreign_rejects_null_ptr() {
6962 let err = unsafe { Tensor::<u8>::from_foreign(std::ptr::null_mut(), &[4], None, None) }
6963 .unwrap_err();
6964 assert!(
6965 matches!(err, Error::InvalidArgument(ref m) if m.contains("non-null")),
6966 "expected InvalidArgument(non-null), got {err:?}"
6967 );
6968 }
6969
6970 #[test]
6971 fn from_foreign_rejects_empty_shape() {
6972 let mut dummy: u8 = 0;
6973 let err = unsafe { Tensor::<u8>::from_foreign(&mut dummy, &[], None, None) }.unwrap_err();
6974 assert!(
6975 matches!(err, Error::InvalidSize(0)),
6976 "expected InvalidSize(0) for empty shape, got {err:?}"
6977 );
6978 }
6979
6980 #[test]
6981 fn from_foreign_rejects_overflow_shape() {
6982 // Two dimensions whose product overflows usize — the overflow guard must
6983 // fire before any pointer arithmetic is attempted.
6984 let mut dummy: u8 = 0;
6985 let huge = [usize::MAX / 2 + 1, 2];
6986 let err = unsafe { Tensor::<u8>::from_foreign(&mut dummy, &huge, None, None) }.unwrap_err();
6987 assert!(
6988 matches!(err, Error::InvalidArgument(ref m) if m.contains("overflow")),
6989 "expected InvalidArgument(overflow), got {err:?}"
6990 );
6991 }
6992
6993 #[test]
6994 fn from_foreign_owner_keeps_allocation_alive() {
6995 // When `owner` is `Some`, dropping the Tensor must not free the backing
6996 // before the owner is also gone — the owner's Drop fires on last ref.
6997 use std::sync::atomic::{AtomicBool, Ordering};
6998 let flag = std::sync::Arc::new(AtomicBool::new(false));
6999 let flag2 = flag.clone();
7000 struct Guard(std::sync::Arc<AtomicBool>);
7001 impl Drop for Guard {
7002 fn drop(&mut self) {
7003 self.0.store(true, Ordering::SeqCst);
7004 }
7005 }
7006 let mut buf: Vec<u32> = vec![42u32; 4];
7007 let ptr = buf.as_mut_ptr();
7008 let owner: ForeignOwner = Box::new(Guard(flag2));
7009 let t = unsafe { Tensor::<u32>::from_foreign(ptr, &[4], Some(owner), None) }.unwrap();
7010 // Map co-owns the backing Arc; the owner must stay alive while the map lives.
7011 let m = t.map().unwrap();
7012 assert_eq!(m.as_slice()[0], 42);
7013 drop(t); // tensor dropped while map is still live
7014 assert!(
7015 !flag.load(Ordering::SeqCst),
7016 "owner must not drop while a map shares the backing"
7017 );
7018 drop(m);
7019 assert!(
7020 flag.load(Ordering::SeqCst),
7021 "owner Drop must fire when the last Arc reference is released"
7022 );
7023 }
7024}