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