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`]); mapping a tensor yields a [`HostView`], which implements
28[`TensorMapTrait`] and derefs to the underlying elements. The
29[`TensorDyn`] type-erased enum wraps `Tensor<T>` for runtime element-type dispatch.
30
31[`TensorMemory::DmaBuf`] is one variant with three implementations behind it — a Linux DMA-heap DMA-BUF,
32a macOS/iOS `IOSurfaceRef`, or an Android `AHardwareBuffer`. Callers ask for `Dma` and get whichever
33the platform provides, so portable code never branches on the mechanism.
34
35## Image tensors
36Images go through [`Tensor::image`] (or [`Tensor::image_desc`] for the full-featured request) rather
37than [`Tensor::new`]. Two things set them apart:
38
39- They take a required [`CpuAccess`] declaration. Hardware access — GPU, NPU, ISP, codec — is always
40 implied; what the CPU intends is the one thing the allocator cannot guess. Mapping beyond the
41 declaration is best-effort and always counted by [`unplanned_cpu_access_count`].
42- DMA-backed images carry a 64-byte-aligned row stride in every layout, because Mali and Vivante
43 reject an EGLImage import at an unaligned pitch. Read the real pitch from
44 [`Tensor::effective_row_stride`]; don't assume `width * bpp`.
45 */
46#[cfg(all(feature = "static", feature = "dynamic"))]
47compile_error!(
48 "edgefirst-tensor: `static` and `dynamic` are mutually exclusive. \
49 A consumer building the C-API leaves wants `--no-default-features \
50 --features dynamic`; everything else wants the default."
51);
52
53// The workspace declares `default-features = false` for this crate, so
54// nothing arrives here by accident anymore -- every consumer names a
55// backend explicitly. Before that change `static` was always a default and
56// this state was unreachable; now a manifest that forgot to ask for either
57// backend back builds an empty `tensor_dyn` module and fails on confusing
58// `E0432: unresolved import` errors instead of a message that names the
59// actual problem.
60#[cfg(not(any(feature = "static", feature = "dynamic")))]
61compile_error!(
62 "edgefirst-tensor: exactly one backend must be enabled. This crate has \
63 `default-features = false` at the workspace level, so every consumer \
64 must name one explicitly: `features = [\"static\"]` for Rust consumers, \
65 or `[\"dynamic\"]` for the C-API leaves."
66);
67
68// `static`-only: every constructor/accessor here is reached exclusively
69// through `TensorStorage`'s `AHardwareBuffer` variant (`impl<T> Tensor<T>`,
70// gated `#[cfg(feature = "static")]`); under `dynamic` nothing calls into
71// it, which clippy's dead-code lint on Android correctly flags without this
72// gate. Not `#[allow]`-ed: an `allow` here would hide the next real
73// dead-code signal in this module, the opposite of what a gate does.
74#[cfg(all(target_os = "android", feature = "static"))]
75mod ahardwarebuffer;
76// Pure AHardwareBuffer layout logic (format table, descriptor geometry,
77// overflow-checked shape math) — cfg-free so it compiles and unit-tests
78// on every host; the android module above consumes it.
79#[allow(dead_code)]
80mod ahardwarebuffer_layout;
81pub mod colorimetry;
82#[cfg(feature = "tracing")]
83pub mod trace;
84
85pub mod covguard;
86mod cuda;
87mod detect;
88// `static`-gated as well as Linux-gated. `DmaTensor` is storage, and only
89// the `static` backend has storage -- `TensorStorage` itself is already
90// `#[cfg(feature = "static")]`, and the `dynamic` backend reaches a
91// DMA-BUF's fd through `ef_tensor_plane_at` instead of ever constructing
92// one. Until task P2b's F3, this module was kept alive under `dynamic` by a
93// single always-`None` `Tensor::as_dma` whose signature named `DmaTensor`;
94// removing that stub revealed the whole module (the tensor, the sync
95// ioctls, the DRM PRIME helpers, the mmap owner) as dead code there --
96// twenty-six dead-code warnings' worth. It was always dead; one method
97// hid it.
98#[cfg(all(target_os = "linux", feature = "static"))]
99mod dma;
100#[cfg(all(target_os = "linux", feature = "static"))]
101mod dmabuf;
102mod error;
103mod format;
104// Layout table is backend-agnostic (pure DXGI-format math), same reasoning
105// as `iosurface_layout` below: compiled on every platform so it is
106// unit-tested on every CI lane. The allocator that consumes it
107// (`crate::d3d11`) is Windows-only.
108#[cfg(target_os = "windows")]
109pub mod d3d11;
110pub mod d3d11_layout;
111// Layout table is backend-agnostic (pure FourCC math). Storage
112// (`IoSurfaceTensor`) is `static`-only; `dynamic` reads IOSurface handles
113// through `ef_tensor_*`.
114#[cfg(any(target_os = "macos", target_os = "ios"))]
115mod iosurface_layout;
116// `static`-only, same reasoning as `ahardwarebuffer` above: reached only
117// through `TensorStorage`'s `IoSurface` variant.
118#[cfg(all(any(target_os = "macos", target_os = "ios"), feature = "static"))]
119mod iosurface;
120pub mod lens;
121// `static`-only: `MemTensor` backs `TensorStorage::Mem`, `impl<T> Tensor<T>`
122// only. `dynamic`'s host-memory allocation goes through `ef_tensor_new`/
123// `ef_tensor_builder_alloc` in `libedgefirst_tensor.so` instead -- there is
124// no local `MemTensor` under `dynamic` to be dead code, so the whole module
125// gates rather than an internal `#[allow]`.
126#[cfg(feature = "static")]
127mod mem;
128mod pbo;
129pub mod pin;
130pub mod protocol;
131// `static`-only, same reasoning as `mem` above: `ShmTensor` backs
132// `TensorStorage::Shm`, `impl<T> Tensor<T>` only.
133#[cfg(all(unix, feature = "static"))]
134mod shm;
135mod tensor_dyn;
136#[cfg(test)]
137mod test_support;
138pub mod view;
139mod vocabulary;
140pub use colorimetry::{
141 ColorEncoding, ColorRange, ColorSpace, ColorTransfer, Colorimetry, MatrixWeights, RangeScaling,
142};
143
144/// Retained constructor: installs the coverage flush-on-abort handler for this
145/// crate's instrumented test binary. See `covguard`. Only present under
146/// coverage on Linux (`.init_array` is ELF-only; the i.MX flush is Linux-only).
147#[cfg(all(coverage, target_os = "linux"))]
148#[used]
149#[link_section = ".init_array"]
150static __EDGEFIRST_COV_INSTALL: extern "C" fn() = {
151 extern "C" fn ctor() {
152 crate::covguard::install();
153 }
154 ctor
155};
156
157// Backing tensor/map types are internal implementation details: callers
158// allocate `Tensor<T>` / `TensorDyn` and map them, never naming the per-memory
159// backing types directly. They are `pub(crate)` so they stay nameable for the
160// `TensorStorage` / `TensorMap` enums without leaking into the public API.
161// Exceptions kept public: `Pbo*` is a GL extension point implemented by the
162// image crate, and `image_iosurface_layout` is a public helper.
163#[cfg(all(target_os = "android", feature = "static"))]
164pub use crate::ahardwarebuffer::image_ahardwarebuffer_layout;
165#[cfg(all(target_os = "android", feature = "static"))]
166pub(crate) use crate::ahardwarebuffer::AHardwareBufferTensor;
167#[cfg(all(target_os = "linux", feature = "static"))]
168pub(crate) use crate::dma::DmaTensor;
169#[cfg(all(any(target_os = "macos", target_os = "ios"), feature = "static"))]
170pub(crate) use crate::iosurface::IoSurfaceTensor;
171// The two geometry helpers are compiled for both backends: three of the four
172// surfaces that read a raw texture's geometry (the Python bindings, and the
173// protocol and blob imports) link `dynamic`, and only the C ABI links
174// `static`. See `crate::d3d11::geometry`.
175#[cfg(target_os = "windows")]
176pub use crate::d3d11::geometry::{d3d11_shared_handle_geometry, d3d11_texture_geometry};
177#[cfg(any(target_os = "macos", target_os = "ios"))]
178pub use crate::iosurface_layout::image_iosurface_layout;
179#[cfg(feature = "static")]
180pub(crate) use crate::mem::MemTensor;
181pub use crate::pbo::{
182 client_state_pbo_ops, read_pbo_vtable_parts, EfClientState, EfPboMapFn, EfPboUnmapFn,
183 PboClientParts, PboMapping, PboOps, PboOpsVtableParts, PboTensor,
184};
185#[cfg(all(unix, feature = "static"))]
186pub(crate) use crate::shm::ShmTensor;
187pub use cuda::{
188 client_state_cuda_ops, gl_map_resource, gl_register_buffer, gl_unmap_resource,
189 gl_unregister_resource, is_cuda_available, memcpy_device_to_host, memcpy_host_to_device,
190 runtime_path as cuda_runtime_path, stream_create, stream_destroy, stream_synchronize,
191 CudaGlOps, CudaHandle, CudaMap, CudaStream,
192};
193pub use detect::{unletter_norm, BoundingBox, DetectBox, ProtoData, ProtoLayout, Segmentation};
194pub use error::{Error, Result};
195// Static-backend-implementation code: backs `tensor-capi`'s own
196// `ef_tensor_export`/`ef_tensor_import` (the static cdylib's `TensorDyn`
197// serialization), the one caller in the whole workspace. It calls
198// `TensorDyn::new`/`from_fd` and reads fields (`.quantization()`,
199// `.colorimetry()`, `.pin_host()`, ...) that only the `static` backend's
200// richer `TensorDyn` surface has, so it never compiled under `dynamic` in
201// fact even though nothing gated it -- see `PRIMITIVE-INVENTORY.md`'s
202// "structural findings".
203#[cfg(feature = "static")]
204pub mod blob;
205#[cfg(feature = "static")]
206pub use blob::{
207 BlobError, BlobHeader, BlobPlane, BlobStrings, BlobView, Regions, TransportMode, HEADER_LEN,
208 PLANE_RECORD_LEN,
209};
210pub use format::{ChromaLayout, PixelFormat, PixelLayout, PlaneGeometry};
211pub use lens::Element;
212use num_traits::Num;
213pub use pin::HostPin;
214pub use protocol::{
215 dtype as tensor_dtype, dtype_to_dtype, format as tensor_format, format_from_code,
216 kind as tensor_kind, SendPtr, TensorDesc, ABI_VERSION,
217};
218use serde::{Deserialize, Serialize};
219#[cfg(unix)]
220use std::os::fd::OwnedFd;
221use std::{
222 fmt,
223 sync::{
224 atomic::{AtomicU64, Ordering},
225 Arc, Weak,
226 },
227};
228// `static`'s `Tensor<T>` is declared directly above, in this file; there is
229// nothing to re-export for it. `dynamic`'s lives in `tensor_dyn` (see that
230// module's docs) and is re-exported here under the same crate-root name so
231// both backends present the same `edgefirst_tensor::Tensor` to callers.
232#[cfg(feature = "dynamic")]
233pub use tensor_dyn::Tensor;
234pub use tensor_dyn::{Raw, TensorDyn};
235pub use view::HostView;
236pub use vocabulary::vocabulary_demo;
237
238/// Opaque keep-alive handle for a foreign-memory tensor (see
239/// [`Tensor::from_foreign`] / [`TensorDyn::from_foreign_ptr`]).
240///
241/// The HAL borrows the foreign buffer without owning it; this handle co-owns
242/// the *source* so the borrowed memory stays valid for the tensor's life. Its
243/// `Drop` releases the source — e.g. a small struct that calls `cudaFreeHost`,
244/// or a `Py<PyAny>` that decrements a NumPy array's refcount. Wrapping it in an
245/// `Arc` (then boxing each clone) makes the release fire exactly once, after
246/// the last sharing tensor/view/map drops, regardless of drop order.
247pub type ForeignOwner = Box<dyn std::any::Any + Send + Sync>;
248
249/// Re-export of `half::f16` so downstream crates can write
250/// `Tensor::<edgefirst_tensor::f16>::from_iosurface(…)` without
251/// adding `half` to their own dependency list. The version stays in
252/// lockstep with the `half` workspace dep.
253pub use half::f16;
254
255// =============================================================================
256// RGBA16F packed-layout geometry — single source of truth
257//
258// A `PlanarRgb` [3,H,W] or `PlanarRgba` [4,H,W] f16 tensor is represented
259// on the GPU as an RGBA16F surface (the only float format accepted by the
260// ANGLE IOSurface extension). Four contiguous f16 elements are packed into
261// each 8-byte RGBA16F texel, yielding a `(W/4, C*H)` surface.
262//
263// All call sites that need these dimensions must use `packed_rgba16f_layout`
264// so the rule lives in exactly one place. Currently consumed by:
265// - `crates/tensor/src/iosurface.rs` `new_image` (macOS IOSurface alloc)
266// - `crates/image/src/gl/iosurface_import.rs` (macOS GL IOSurface import)
267// - `crates/image/src/gl/processor/float.rs` (Linux GL float render — PBO
268// readback and DMA-BUF, also via the `dma_f16_packed_layout` wrapper)
269// =============================================================================
270
271/// Geometry of the RGBA16F-packed surface backing a planar F16 image tensor.
272///
273/// ANGLE only supports one float `(type, internal_format)` pair for IOSurface
274/// import: `(GL_HALF_FLOAT, GL_RGBA)` = RGBA16F (8 bytes/texel). To map a
275/// `[C, H, W]` f16 planar tensor onto such a surface, 4 contiguous f16
276/// elements are packed into each RGBA16F texel, yielding a surface of
277/// `(W/4, C*H)` texels at 8 bytes/texel. The byte stream is identical to a
278/// (nonexistent) R16F `(W, C*H)` surface and can be consumed as `&[f16]`
279/// with shape `[1, C, H, W]` without rearrangement.
280///
281/// Obtain via [`packed_rgba16f_layout`] — never construct directly.
282#[derive(Debug, Clone, Copy, PartialEq, Eq)]
283pub struct PackedRgba16fLayout {
284 /// Surface width in texels (`width / 4`).
285 pub surface_w: usize,
286 /// Surface height in texels (`planes * height`).
287 pub surface_h: usize,
288 /// Bytes per RGBA16F texel (always 8).
289 pub bytes_per_texel: usize,
290 /// Row pitch in bytes (`surface_w * 8`).
291 pub pitch: usize,
292}
293
294/// Canonical geometry for the RGBA16F-packed surface backing a planar F16
295/// image tensor.
296///
297/// Returns `Some(layout)` only when **all** of the following hold:
298///
299/// - `dtype == DType::F16`
300/// - `format` is `PixelFormat::PlanarRgb` (3 planes) or
301/// `PixelFormat::PlanarRgba` (4 planes)
302/// - `width % 4 == 0`
303///
304/// Returns `None` for any other `(format, dtype)` combination, misaligned
305/// width, or when the surface geometry would overflow `usize` — callers
306/// must fall back to a non-packed path or return a context-appropriate
307/// error.
308///
309/// # Examples
310///
311/// ```rust
312/// use edgefirst_tensor::{packed_rgba16f_layout, PixelFormat, DType};
313///
314/// let layout = packed_rgba16f_layout(PixelFormat::PlanarRgb, DType::F16, 640, 480).unwrap();
315/// assert_eq!(layout.surface_w, 160);
316/// assert_eq!(layout.surface_h, 1440);
317/// assert_eq!(layout.bytes_per_texel, 8);
318/// assert_eq!(layout.pitch, 1280);
319/// ```
320pub fn packed_rgba16f_layout(
321 format: PixelFormat,
322 dtype: DType,
323 width: usize,
324 height: usize,
325) -> Option<PackedRgba16fLayout> {
326 if dtype != DType::F16 {
327 return None;
328 }
329 let planes: usize = match format {
330 PixelFormat::PlanarRgb => 3,
331 PixelFormat::PlanarRgba => 4,
332 _ => return None,
333 };
334 if !width.is_multiple_of(4) {
335 return None;
336 }
337 let surface_w = width / 4;
338 // Checked arithmetic: a degenerate (height, width) could otherwise wrap
339 // and yield an under-sized layout, which downstream allocators trust for
340 // GPU/CPU buffer sizing. Overflow → None (handled like any other
341 // unsupported geometry).
342 let surface_h = planes.checked_mul(height)?;
343 let bytes_per_texel = 8;
344 let pitch = surface_w.checked_mul(bytes_per_texel)?;
345 Some(PackedRgba16fLayout {
346 surface_w,
347 surface_h,
348 bytes_per_texel,
349 pitch,
350 })
351}
352
353/// Geometry of the RGBA8888-packed surface backing a packed RGB u8/i8 image
354/// tensor.
355///
356/// GPUs have no 3-channel renderable format, so the GL engine's two-pass
357/// packed-RGB shader writes the tight `[H, W, 3]` byte stream into an
358/// RGBA8888 surface: each texel carries 4 consecutive RGB bytes, giving a
359/// `(W*3/4, H)` surface at 4 bytes/texel whose rows are byte-identical to
360/// tight RGB — consumable flat as `[H, W, 3]` with no rearrangement (the
361/// u8/i8 analog of [`packed_rgba16f_layout`]; i8 shares the layout since
362/// INT8 quantization is a per-byte `^0x80` bias, not a format change).
363///
364/// Returns `Some(layout)` only when `width % 4 == 0` (so `W*3` bytes divide
365/// into whole texels) and the geometry does not overflow `usize`.
366///
367/// # Examples
368///
369/// ```rust
370/// use edgefirst_tensor::packed_rgb888_layout;
371///
372/// let layout = packed_rgb888_layout(640, 480).unwrap();
373/// assert_eq!(layout.surface_w, 480); // 640*3/4
374/// assert_eq!(layout.surface_h, 480);
375/// assert_eq!(layout.bytes_per_texel, 4);
376/// assert_eq!(layout.pitch, 1920); // 640*3
377/// assert!(packed_rgb888_layout(641, 480).is_none());
378/// ```
379pub fn packed_rgb888_layout(width: usize, height: usize) -> Option<PackedRgb888Layout> {
380 if !width.is_multiple_of(4) {
381 return None;
382 }
383 let row_bytes = width.checked_mul(3)?;
384 let surface_w = row_bytes / 4;
385 Some(PackedRgb888Layout {
386 surface_w,
387 surface_h: height,
388 bytes_per_texel: 4,
389 pitch: row_bytes,
390 })
391}
392
393/// Geometry of the RGBA8888 surface backing a packed RGB u8/i8 image tensor.
394///
395/// Obtain via [`packed_rgb888_layout`] — never construct directly.
396#[derive(Debug, Clone, Copy, PartialEq, Eq)]
397pub struct PackedRgb888Layout {
398 /// Surface width in texels (`width * 3 / 4`).
399 pub surface_w: usize,
400 /// Surface height in texels (`height`).
401 pub surface_h: usize,
402 /// Bytes per RGBA8888 texel (always 4).
403 pub bytes_per_texel: usize,
404 /// Row pitch in bytes (`surface_w * 4` = `width * 3`).
405 pub pitch: usize,
406}
407
408/// The logical `(width, height)` an **allocation** shape describes for
409/// `format` -- the inverse of [`PixelFormat::allocation_shape`].
410///
411/// Used wherever a tensor arrives carrying a shape and a format but not its
412/// dimensions: the C constructors, the Python bindings and the cross-package
413/// protocol import.
414///
415/// # The input is the allocation shape, not the addressing shape
416///
417/// [`PixelFormat::allocation_shape`] is what [`Tensor::image`] allocates from
418/// and what `descriptor_pinned` emits from `shape()`, so that is what this
419/// inverts. [`PixelFormat::addressing_shape`] is a different, smaller grid,
420/// and for a semi-planar format the two collide: NV12's addressing shape is
421/// `[h, w]` and its allocation shape is `[combined_h, w]`, both rank 2. Handed
422/// an addressing shape this cannot tell -- NV12 `[480, 640]` reads as a
423/// combined height of 480 and answers `(640, 320)`, silently. Nothing in the
424/// shape distinguishes them, so callers must pass the allocation shape.
425///
426/// A D3D11 import is the one place this cannot bite: the constructors check
427/// the derived size against the texture description, so a mismatch fails with
428/// [`Error::InvalidArgument`] naming both. On Windows, prefer
429/// `d3d11_texture_geometry` and `d3d11_shared_handle_geometry`, which read the
430/// dimensions off the texture itself and leave the shape as a consistency
431/// check -- except for a semi-planar texture, whose width is its row stride
432/// rather than the image width, so those helpers take the width from the shape
433/// they are given and refuse to answer without one.
434///
435/// The rank is part of the format's contract, so a shape of the wrong rank is
436/// an error rather than an interpretation of whatever dimensions are present:
437///
438/// | Layout | Shape | Result |
439/// |--------|-------|--------|
440/// | Packed | `[h, w, c]` | `(shape[1], shape[0])` |
441/// | Planar | `[c, h, w]` | `(shape[2], shape[1])` |
442/// | SemiPlanar | `[combined_h, w]` | `(shape[1], image height)` |
443///
444/// A semi-planar shape carries the *combined* plane height, so the logical
445/// height is recovered through
446/// [`PixelFormat::image_height_from_combined`], which rejects a combined
447/// height no image height produces (721 for NV12, say) instead of rounding to
448/// a nearby one.
449///
450/// # Errors
451///
452/// [`Error::InvalidShape`] for a shape of the wrong rank, or a semi-planar
453/// combined height with no image height behind it.
454///
455/// # Example
456///
457/// ```
458/// use edgefirst_tensor::{image_dims_from_shape, PixelFormat};
459/// assert_eq!(image_dims_from_shape(PixelFormat::Rgba, &[480, 640, 4]).unwrap(), (640, 480));
460/// assert_eq!(image_dims_from_shape(PixelFormat::PlanarRgb, &[3, 480, 640]).unwrap(), (640, 480));
461/// assert_eq!(image_dims_from_shape(PixelFormat::Nv12, &[722, 640]).unwrap(), (640, 481));
462/// assert!(image_dims_from_shape(PixelFormat::Nv12, &[721, 640]).is_err());
463/// ```
464pub fn image_dims_from_shape(format: PixelFormat, shape: &[usize]) -> Result<(usize, usize)> {
465 let wrong_rank = |expected: usize| {
466 Error::InvalidShape(format!(
467 "image_dims_from_shape: {format:?} describes a rank-{expected} allocation \
468 shape, got rank {} ({shape:?})",
469 shape.len()
470 ))
471 };
472 match format.layout() {
473 PixelLayout::Packed => {
474 if shape.len() != 3 {
475 return Err(wrong_rank(3));
476 }
477 Ok((shape[1], shape[0]))
478 }
479 PixelLayout::Planar => {
480 if shape.len() != 3 {
481 return Err(wrong_rank(3));
482 }
483 Ok((shape[2], shape[1]))
484 }
485 PixelLayout::SemiPlanar => {
486 if shape.len() != 2 {
487 return Err(wrong_rank(2));
488 }
489 let height = format.image_height_from_combined(shape[0]).ok_or_else(|| {
490 Error::InvalidShape(format!(
491 "image_dims_from_shape: no {format:?} image height has combined plane \
492 height {} (shape {shape:?})",
493 shape[0]
494 ))
495 })?;
496 Ok((shape[1], height))
497 }
498 }
499}
500
501/// Per-plane DMA-BUF descriptor for external buffer import.
502///
503/// Owns a duplicated file descriptor plus optional stride and offset metadata.
504/// The fd is duplicated eagerly in [`new()`](Self::new) so that a bad fd is
505/// caught immediately. `import_image` consumes the descriptor and takes
506/// ownership of the duped fd — no further cleanup is needed by the caller.
507///
508/// # Examples
509///
510/// ```rust,no_run
511/// use edgefirst_tensor::PlaneDescriptor;
512/// use std::os::fd::BorrowedFd;
513///
514/// // SAFETY: fd 42 is hypothetical; real code must pass a valid fd.
515/// let pd = unsafe { PlaneDescriptor::new(BorrowedFd::borrow_raw(42)) }
516/// .unwrap()
517/// .with_stride(2048)
518/// .with_offset(0);
519/// ```
520#[cfg(unix)]
521pub struct PlaneDescriptor {
522 fd: OwnedFd,
523 stride: Option<usize>,
524 offset: Option<usize>,
525}
526
527#[cfg(unix)]
528impl PlaneDescriptor {
529 /// Create a new plane descriptor by duplicating the given file descriptor.
530 ///
531 /// The fd is duped immediately — a bad fd fails here rather than inside
532 /// `import_image`. The caller retains ownership of the original fd.
533 ///
534 /// # Errors
535 ///
536 /// Returns an error if the `dup()` syscall fails (e.g. invalid fd or
537 /// fd limit reached).
538 pub fn new(fd: std::os::fd::BorrowedFd<'_>) -> Result<Self> {
539 let owned = fd.try_clone_to_owned()?;
540 Ok(Self {
541 fd: owned,
542 stride: None,
543 offset: None,
544 })
545 }
546
547 /// Set the row stride in bytes (consuming builder).
548 pub fn with_stride(mut self, stride: usize) -> Self {
549 self.stride = Some(stride);
550 self
551 }
552
553 /// Set the plane offset in bytes (consuming builder).
554 pub fn with_offset(mut self, offset: usize) -> Self {
555 self.offset = Some(offset);
556 self
557 }
558
559 /// Consume the descriptor and return the owned file descriptor.
560 pub fn into_fd(self) -> OwnedFd {
561 self.fd
562 }
563
564 /// Row stride in bytes, if set.
565 pub fn stride(&self) -> Option<usize> {
566 self.stride
567 }
568
569 /// Plane offset in bytes, if set.
570 pub fn offset(&self) -> Option<usize> {
571 self.offset
572 }
573}
574
575/// A rectangular sub-region of a tensor's leading spatial frame, in pixels.
576///
577/// `Region` is the single rectangle type in the workspace: the argument to
578/// [`Tensor::view`], the source sampling window in the image crate's `Crop`,
579/// and the geometry the image backend lowers to a `glViewport` (a destination
580/// tile) or a sampling rectangle (a source). Coordinates are pixel/element
581/// units of the leading spatial axes; byte addressing is derived from the
582/// parent's row stride, not stored here.
583#[derive(Debug, Clone, Copy, PartialEq, Eq)]
584pub struct Region {
585 pub x: usize,
586 pub y: usize,
587 pub width: usize,
588 pub height: usize,
589}
590
591impl Region {
592 /// Create a region at `(x, y)` spanning `width` × `height` pixels.
593 pub fn new(x: usize, y: usize, width: usize, height: usize) -> Self {
594 Self {
595 x,
596 y,
597 width,
598 height,
599 }
600 }
601
602 /// True when the region lies fully within a `width` × `height` frame.
603 pub fn fits_within(&self, width: usize, height: usize) -> bool {
604 self.x.saturating_add(self.width) <= width && self.y.saturating_add(self.height) <= height
605 }
606}
607
608/// The parent image a [`view`](Tensor::view)/[`batch`](Tensor::batch) sub-region
609/// was carved from, snapshotted at the time the view was created.
610///
611/// A view shares the parent's `BufferIdentity` and addresses a sub-rectangle of
612/// it. The GL backend keys its EGLImage import on the **parent** geometry (so all
613/// sibling views of one buffer collapse to a single import) and renders each
614/// view as a `glViewport`+`glScissor` ROI at `(x, y, width, height)` within that
615/// parent — the view is render state, never a distinct import. `parent_width`/
616/// `parent_height` are the parent's logical pixel dimensions; `x`/`y` are this
617/// view's top-left origin within the parent (pixels). Nested views compose:
618/// the snapshot always names the **root** parent, with offsets accumulated.
619#[derive(Debug, Clone, Copy, PartialEq, Eq)]
620pub struct ViewOrigin {
621 /// Logical width of the root parent image, in pixels.
622 pub parent_width: usize,
623 /// Logical height of the root parent image, in pixels. For a `batch(n)` view
624 /// of an `[N, H, W, C]` tensor this is `N * H` (the tiles stack vertically in
625 /// the shared buffer).
626 pub parent_height: usize,
627 /// The parent's row stride in **bytes**. The GL backend keys its EGLImage
628 /// import/cache and pitch on this — NOT on the view's own `row_stride`, which
629 /// a single-row view sets tight (for map-span safety). Using the parent
630 /// stride keeps the import pitch parent-consistent so single-row and
631 /// multi-row sibling views collapse onto the same parent import.
632 pub parent_row_stride: usize,
633 /// This view's top-left x origin within the root parent, in pixels.
634 pub x: usize,
635 /// This view's top-left y origin within the root parent, in pixels.
636 pub y: usize,
637}
638
639// Declared through `ef_vocabulary!` (see `vocabulary.rs`): `.code()` is this
640// crate's wire numbering, mirrored by `protocol::dtype` and the C ABI's
641// `hal_dtype`. Both surfaces are generated from the same per-variant literal
642// below rather than re-typed by hand, and `dtype_codes_are_unchanged_by_the_macro`
643// in `tests/vocabulary.rs` pins them against the values already shipping.
644crate::ef_vocabulary! {
645 /// Element type discriminant for runtime type identification.
646 #[derive(Serialize, Deserialize)]
647 #[non_exhaustive]
648 pub enum DType {
649 U8 = 0, "u8", U8,
650 I8 = 1, "i8", I8,
651 U16 = 2, "u16", U16,
652 I16 = 3, "i16", I16,
653 U32 = 4, "u32", U32,
654 I32 = 5, "i32", I32,
655 U64 = 6, "u64", U64,
656 I64 = 7, "i64", I64,
657 F16 = 8, "f16", F16,
658 F32 = 9, "f32", F32,
659 F64 = 10, "f64", F64,
660 }
661 // `#[doc(hidden)]`: this module has to be `pub` for `protocol::dtype`'s
662 // `pub use` re-export below to compile, but it is emission plumbing, not
663 // a second documented API for the same eleven codes -- `protocol::dtype`
664 // is the canonical, documented path.
665 #[doc(hidden)]
666 pub mod dtype_wire;
667}
668
669impl DType {
670 /// Size of one element in bytes.
671 pub const fn size(&self) -> usize {
672 match self {
673 Self::U8 | Self::I8 => 1,
674 Self::U16 | Self::I16 | Self::F16 => 2,
675 Self::U32 | Self::I32 | Self::F32 => 4,
676 Self::U64 | Self::I64 | Self::F64 => 8,
677 }
678 }
679
680 /// Short type name (e.g., "u8", "f32", "f16"). Identical to
681 /// [`DType::as_str`]; kept as its own method for existing callers.
682 pub fn name(&self) -> &'static str {
683 self.as_str()
684 }
685}
686
687impl fmt::Display for DType {
688 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
689 f.write_str(self.name())
690 }
691}
692
693/// Map a static numeric type `T` to its `DType` discriminant, returning
694/// `None` for types that do not have a `DType` representation (e.g.
695/// user-defined wrappers in tests).
696///
697/// Runtime dtype of a static `Tensor<T>` element type — a safe `TypeId`
698/// allowlist of HAL's primitive numeric types (`Some` for `u8`..`i64` /
699/// `f16` / `f32` / `f64`, `None` otherwise). Used by the macOS IOSurface
700/// image constructors (FourCC / pixel-format lookup) and by the `Mem`
701/// backing's `alloc_zeroed` fast path (these types' `T::zero()` is the
702/// all-zeros bit pattern), so it is compiled on every target -- but only
703/// under `static`: every call site lives in `impl<T> Tensor<T>`, itself
704/// `#[cfg(feature = "static")]`.
705#[cfg(feature = "static")]
706pub(crate) fn dtype_of<T: 'static>() -> Option<DType> {
707 use std::any::TypeId;
708 let id = TypeId::of::<T>();
709 if id == TypeId::of::<u8>() {
710 Some(DType::U8)
711 } else if id == TypeId::of::<i8>() {
712 Some(DType::I8)
713 } else if id == TypeId::of::<u16>() {
714 Some(DType::U16)
715 } else if id == TypeId::of::<i16>() {
716 Some(DType::I16)
717 } else if id == TypeId::of::<u32>() {
718 Some(DType::U32)
719 } else if id == TypeId::of::<i32>() {
720 Some(DType::I32)
721 } else if id == TypeId::of::<u64>() {
722 Some(DType::U64)
723 } else if id == TypeId::of::<i64>() {
724 Some(DType::I64)
725 } else if id == TypeId::of::<half::f16>() {
726 Some(DType::F16)
727 } else if id == TypeId::of::<f32>() {
728 Some(DType::F32)
729 } else if id == TypeId::of::<f64>() {
730 Some(DType::F64)
731 } else {
732 None
733 }
734}
735
736// =============================================================================
737// Quantization metadata — type-gated to integer element types via sealed
738// `IntegerType` trait. Accessors on `Tensor<T>` only compile when `T` is
739// an integer type; calling them on `Tensor<f32>` / `Tensor<f16>` etc. is a
740// compile error, not a runtime one.
741// =============================================================================
742
743mod sealed {
744 pub trait Sealed {}
745 impl Sealed for u8 {}
746 impl Sealed for i8 {}
747 impl Sealed for u16 {}
748 impl Sealed for i16 {}
749 impl Sealed for u32 {}
750 impl Sealed for i32 {}
751 impl Sealed for u64 {}
752 impl Sealed for i64 {}
753 // f16 / f32 / f64 are sealed too (crate::lens::Element covers all
754 // eleven element types), but deliberately NOT given `IntegerType` below.
755 impl Sealed for half::f16 {}
756 impl Sealed for f32 {}
757 impl Sealed for f64 {}
758}
759
760/// Integer element types that may carry quantization metadata.
761///
762/// Sealed trait: implemented for `u8`, `i8`, `u16`, `i16`, `u32`, `i32`,
763/// `u64`, `i64`. Cannot be implemented downstream. Float element types
764/// (`half::f16`, `f32`, `f64`) are explicitly excluded — quantization
765/// metadata does not apply to float tensors per the edgefirst.json spec.
766pub trait IntegerType: sealed::Sealed {}
767impl IntegerType for u8 {}
768impl IntegerType for i8 {}
769impl IntegerType for u16 {}
770impl IntegerType for i16 {}
771impl IntegerType for u32 {}
772impl IntegerType for i32 {}
773impl IntegerType for u64 {}
774impl IntegerType for i64 {}
775
776/// Quantization parameters for an integer tensor.
777///
778/// Covers all four modes the edgefirst.json spec defines:
779///
780/// | Mode | `scale.len()` | `zero_point` | `axis` |
781/// |---|---|---|---|
782/// | Per-tensor symmetric | 1 | `None` | `None` |
783/// | Per-tensor asymmetric | 1 | `Some(len == 1)` | `None` |
784/// | Per-channel symmetric | >1 | `None` | `Some(c)` |
785/// | Per-channel asymmetric | >1 | `Some(len == scale.len())` | `Some(c)` |
786///
787/// The quantized storage type is carried on the parent [`Tensor<T>`]; this
788/// struct does not duplicate it. Construct via the four named constructors
789/// (the only public entry points); direct field mutation is not allowed so
790/// invalid combinations cannot be represented.
791///
792/// Dequantization formula:
793///
794/// ```text
795/// real_value = scale[c] × (quantized_value[c] - zero_point[c])
796/// ```
797///
798/// where `c` is the channel index (always `0` for per-tensor).
799#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
800pub struct Quantization {
801 /// Per-tensor: `vec![scale]`. Per-channel: `vec![scale_0, scale_1, ...]`.
802 #[serde(deserialize_with = "deserialize_scalar_or_vec_f32")]
803 scale: Vec<f32>,
804
805 /// `None` means symmetric (zero-point is 0). `Some(vec)` must have the
806 /// same length as `scale`.
807 #[serde(
808 default,
809 deserialize_with = "deserialize_opt_scalar_or_vec_i32",
810 skip_serializing_if = "Option::is_none"
811 )]
812 zero_point: Option<Vec<i32>>,
813
814 /// Channel axis for per-channel quantization. `Some(_)` iff
815 /// `scale.len() > 1`. Validated against the parent tensor's shape at
816 /// `set_quantization()` time.
817 #[serde(default, skip_serializing_if = "Option::is_none")]
818 axis: Option<usize>,
819}
820
821/// Semantic mode discriminant for hot-path kernel dispatch.
822///
823/// Obtain via [`Quantization::mode`] once at kernel entry; never inside a
824/// pixel-level loop. The enum is borrow-based so the hot kernel receives
825/// the scales / zero-points as slices without reallocation.
826#[derive(Debug, Clone, Copy)]
827pub enum QuantMode<'a> {
828 PerTensorSymmetric {
829 scale: f32,
830 },
831 PerTensor {
832 scale: f32,
833 zero_point: i32,
834 },
835 PerChannelSymmetric {
836 scales: &'a [f32],
837 axis: usize,
838 },
839 PerChannel {
840 scales: &'a [f32],
841 zero_points: &'a [i32],
842 axis: usize,
843 },
844}
845
846impl Quantization {
847 /// Per-tensor symmetric (zero_point = 0).
848 pub fn per_tensor_symmetric(scale: f32) -> Self {
849 Self {
850 scale: vec![scale],
851 zero_point: None,
852 axis: None,
853 }
854 }
855
856 /// Per-tensor asymmetric — the most common runtime shape.
857 pub fn per_tensor(scale: f32, zero_point: i32) -> Self {
858 Self {
859 scale: vec![scale],
860 zero_point: Some(vec![zero_point]),
861 axis: None,
862 }
863 }
864
865 /// Per-channel symmetric. Errors on empty `scales`.
866 pub fn per_channel_symmetric(scales: Vec<f32>, axis: usize) -> Result<Self> {
867 if scales.is_empty() {
868 return Err(Error::QuantizationInvalid {
869 field: "scale.len",
870 expected: "non-empty per-channel scales".to_string(),
871 got: "length 0".to_string(),
872 });
873 }
874 Ok(Self {
875 scale: scales,
876 zero_point: None,
877 axis: Some(axis),
878 })
879 }
880
881 /// Per-channel asymmetric. Errors on length mismatch between `scales`
882 /// and `zero_points`, or empty arrays.
883 pub fn per_channel(scales: Vec<f32>, zero_points: Vec<i32>, axis: usize) -> Result<Self> {
884 if scales.is_empty() {
885 return Err(Error::QuantizationInvalid {
886 field: "scale.len",
887 expected: "non-empty per-channel scales".to_string(),
888 got: "length 0".to_string(),
889 });
890 }
891 if scales.len() != zero_points.len() {
892 return Err(Error::QuantizationInvalid {
893 field: "zero_point.len",
894 expected: format!("length matches scale ({})", scales.len()),
895 got: format!("length {}", zero_points.len()),
896 });
897 }
898 Ok(Self {
899 scale: scales,
900 zero_point: Some(zero_points),
901 axis: Some(axis),
902 })
903 }
904
905 /// Borrow-based dispatch view. Match once at kernel entry.
906 pub fn mode(&self) -> QuantMode<'_> {
907 match (self.scale.len(), self.zero_point.as_deref(), self.axis) {
908 (1, None, _) => QuantMode::PerTensorSymmetric {
909 scale: self.scale[0],
910 },
911 (1, Some(zps), _) => QuantMode::PerTensor {
912 scale: self.scale[0],
913 zero_point: zps.first().copied().unwrap_or(0),
914 },
915 (_, None, Some(axis)) => QuantMode::PerChannelSymmetric {
916 scales: &self.scale,
917 axis,
918 },
919 (_, Some(zps), Some(axis)) => QuantMode::PerChannel {
920 scales: &self.scale,
921 zero_points: zps,
922 axis,
923 },
924 // The `validate()` path prevents constructing a
925 // per-channel Quantization without an axis, so the remaining
926 // pattern is unreachable in practice. Fall back to
927 // per-tensor symmetric using scale[0] to avoid panicking in
928 // release; debug builds assert.
929 _ => {
930 debug_assert!(
931 false,
932 "Quantization::mode: per-channel without axis is unreachable"
933 );
934 QuantMode::PerTensorSymmetric {
935 scale: self.scale.first().copied().unwrap_or(1.0),
936 }
937 }
938 }
939 }
940
941 /// Returns `true` for per-tensor quantization (`scale.len() == 1`).
942 pub fn is_per_tensor(&self) -> bool {
943 self.scale.len() == 1
944 }
945
946 /// Returns `true` for per-channel quantization (`scale.len() > 1`).
947 pub fn is_per_channel(&self) -> bool {
948 self.scale.len() > 1
949 }
950
951 /// Returns `true` for symmetric quantization (no zero-point, or
952 /// zero-point vector of all zeros).
953 pub fn is_symmetric(&self) -> bool {
954 match &self.zero_point {
955 None => true,
956 Some(zps) => zps.iter().all(|&z| z == 0),
957 }
958 }
959
960 /// Borrow the scale array. Length 1 for per-tensor; `num_channels` for
961 /// per-channel.
962 pub fn scale(&self) -> &[f32] {
963 &self.scale
964 }
965
966 /// Borrow the zero-point array. `None` for symmetric.
967 pub fn zero_point(&self) -> Option<&[i32]> {
968 self.zero_point.as_deref()
969 }
970
971 /// Channel axis for per-channel quantization. `None` for per-tensor.
972 pub fn axis(&self) -> Option<usize> {
973 self.axis
974 }
975
976 /// Validate against a target tensor shape. Runs in
977 /// `Tensor::set_quantization()`. Catches:
978 /// - empty `scale` (reject — must declare at least one factor)
979 /// - `zero_point` length inconsistent with `scale` (reject —
980 /// per-tensor must have len 1, per-channel must match `scale.len`)
981 /// - `axis >= shape.len()` (axis out of range)
982 /// - `scale.len() != shape[axis]` for per-channel
983 /// - per-channel without axis (reject)
984 /// - per-tensor with redundant axis (reject)
985 ///
986 /// `static`-only: `Tensor::set_quantization` is `impl<T> Tensor<T>`,
987 /// itself `#[cfg(feature = "static")]`. `dynamic`'s
988 /// `TensorDyn::set_quantization` (`dynamic_backend.rs`) sends the raw
989 /// scale/zero-point/axis over FFI and lets `libedgefirst_tensor.so`
990 /// validate them on its own (`static`) side instead of calling this
991 /// method locally.
992 #[cfg(feature = "static")]
993 pub(crate) fn validate(&self, shape: &[usize]) -> Result<()> {
994 // `Quantization` is `Deserialize`, so malformed JSON like
995 // `{"scale": [], "zero_point": []}` could otherwise produce an
996 // ill-defined value that confuses `mode()` selection and the
997 // per-channel kernels' indexing.
998 if self.scale.is_empty() {
999 return Err(Error::QuantizationInvalid {
1000 field: "scale.len",
1001 expected: ">= 1".to_string(),
1002 got: "0".to_string(),
1003 });
1004 }
1005 if let Some(zps) = self.zero_point.as_ref() {
1006 // Per-tensor: scale.len() == 1 and zero_point.len() must == 1.
1007 // Per-channel: zero_point.len() must == scale.len().
1008 let expected = if self.scale.len() == 1 {
1009 1
1010 } else {
1011 self.scale.len()
1012 };
1013 if zps.len() != expected {
1014 return Err(Error::QuantizationInvalid {
1015 field: "zero_point.len",
1016 expected: format!(
1017 "{expected} (matching {})",
1018 if self.scale.len() == 1 {
1019 "per-tensor scale"
1020 } else {
1021 "per-channel scale.len"
1022 }
1023 ),
1024 got: format!("length {}", zps.len()),
1025 });
1026 }
1027 }
1028
1029 match (self.scale.len(), self.axis) {
1030 (1, None) => Ok(()),
1031 (1, Some(_)) => Err(Error::QuantizationInvalid {
1032 field: "per_tensor_redundant_axis",
1033 expected: "axis=None for per-tensor quantization".to_string(),
1034 got: format!("axis={:?}", self.axis),
1035 }),
1036 (_, None) => Err(Error::QuantizationInvalid {
1037 field: "per_channel_requires_axis",
1038 expected: format!(
1039 "axis=Some(_) for per-channel quantization (scale.len={})",
1040 self.scale.len()
1041 ),
1042 got: "axis=None".to_string(),
1043 }),
1044 (n, Some(axis)) => {
1045 if axis >= shape.len() {
1046 return Err(Error::QuantizationInvalid {
1047 field: "axis",
1048 expected: format!("axis < tensor rank ({})", shape.len()),
1049 got: format!("axis={axis}"),
1050 });
1051 }
1052 if shape[axis] != n {
1053 return Err(Error::QuantizationInvalid {
1054 field: "scale.len",
1055 expected: format!("length matches shape[{axis}] ({})", shape[axis]),
1056 got: format!("length {n}"),
1057 });
1058 }
1059 Ok(())
1060 }
1061 }
1062 }
1063}
1064
1065impl From<(f32, i32)> for Quantization {
1066 /// Convenience construction from a `(scale, zero_point)` tuple. Matches
1067 /// the legacy `QuantTuple` / `Quantization::new` calling convention so
1068 /// existing `(0.1, -128).into()` sites keep working.
1069 fn from((scale, zero_point): (f32, i32)) -> Self {
1070 Self::per_tensor(scale, zero_point)
1071 }
1072}
1073
1074fn deserialize_scalar_or_vec_f32<'de, D: serde::Deserializer<'de>>(
1075 de: D,
1076) -> std::result::Result<Vec<f32>, D::Error> {
1077 use serde::de::{self, Visitor};
1078 struct V;
1079 impl<'de> Visitor<'de> for V {
1080 type Value = Vec<f32>;
1081 // Only serde's error text reads this; no value flows from it.
1082 #[cfg_attr(test, mutants::skip)]
1083 fn expecting(&self, f: &mut fmt::Formatter) -> fmt::Result {
1084 f.write_str("f32 or array of f32")
1085 }
1086 fn visit_f64<E: de::Error>(self, v: f64) -> std::result::Result<Self::Value, E> {
1087 Ok(vec![v as f32])
1088 }
1089 #[allow(clippy::cast_possible_truncation)]
1090 fn visit_i64<E: de::Error>(self, v: i64) -> std::result::Result<Self::Value, E> {
1091 Ok(vec![v as f32])
1092 }
1093 #[allow(clippy::cast_possible_truncation, clippy::cast_precision_loss)]
1094 fn visit_u64<E: de::Error>(self, v: u64) -> std::result::Result<Self::Value, E> {
1095 Ok(vec![v as f32])
1096 }
1097 fn visit_seq<A: de::SeqAccess<'de>>(
1098 self,
1099 mut seq: A,
1100 ) -> std::result::Result<Self::Value, A::Error> {
1101 let mut out = Vec::with_capacity(seq.size_hint().unwrap_or(1));
1102 while let Some(x) = seq.next_element::<f32>()? {
1103 out.push(x);
1104 }
1105 Ok(out)
1106 }
1107 }
1108 de.deserialize_any(V)
1109}
1110
1111fn deserialize_opt_scalar_or_vec_i32<'de, D: serde::Deserializer<'de>>(
1112 de: D,
1113) -> std::result::Result<Option<Vec<i32>>, D::Error> {
1114 use serde::de::{self, Visitor};
1115 struct V;
1116 impl<'de> Visitor<'de> for V {
1117 type Value = Option<Vec<i32>>;
1118 // Only serde's error text reads this; no value flows from it.
1119 #[cfg_attr(test, mutants::skip)]
1120 fn expecting(&self, f: &mut fmt::Formatter) -> fmt::Result {
1121 f.write_str("null, i32, or array of i32")
1122 }
1123 // `Ok(None)` is `Ok(Default::default())`, the only replacement.
1124 #[cfg_attr(test, mutants::skip)]
1125 fn visit_none<E: de::Error>(self) -> std::result::Result<Self::Value, E> {
1126 Ok(None)
1127 }
1128 // As `visit_none`.
1129 #[cfg_attr(test, mutants::skip)]
1130 fn visit_unit<E: de::Error>(self) -> std::result::Result<Self::Value, E> {
1131 Ok(None)
1132 }
1133 fn visit_some<D2: serde::Deserializer<'de>>(
1134 self,
1135 de: D2,
1136 ) -> std::result::Result<Self::Value, D2::Error> {
1137 struct Inner;
1138 impl<'de> Visitor<'de> for Inner {
1139 type Value = Vec<i32>;
1140 // Only serde's error text reads this; no value flows from it.
1141 #[cfg_attr(test, mutants::skip)]
1142 fn expecting(&self, f: &mut fmt::Formatter) -> fmt::Result {
1143 f.write_str("i32 or array of i32")
1144 }
1145 #[allow(clippy::cast_possible_truncation)]
1146 fn visit_i64<E: de::Error>(self, v: i64) -> std::result::Result<Self::Value, E> {
1147 Ok(vec![v as i32])
1148 }
1149 #[allow(clippy::cast_possible_truncation, clippy::cast_possible_wrap)]
1150 fn visit_u64<E: de::Error>(self, v: u64) -> std::result::Result<Self::Value, E> {
1151 Ok(vec![v as i32])
1152 }
1153 fn visit_seq<A: de::SeqAccess<'de>>(
1154 self,
1155 mut seq: A,
1156 ) -> std::result::Result<Self::Value, A::Error> {
1157 let mut out = Vec::with_capacity(seq.size_hint().unwrap_or(1));
1158 while let Some(x) = seq.next_element::<i32>()? {
1159 out.push(x);
1160 }
1161 Ok(out)
1162 }
1163 }
1164 de.deserialize_any(Inner).map(Some)
1165 }
1166 #[allow(clippy::cast_possible_truncation)]
1167 fn visit_i64<E: de::Error>(self, v: i64) -> std::result::Result<Self::Value, E> {
1168 Ok(Some(vec![v as i32]))
1169 }
1170 #[allow(clippy::cast_possible_truncation, clippy::cast_possible_wrap)]
1171 fn visit_u64<E: de::Error>(self, v: u64) -> std::result::Result<Self::Value, E> {
1172 Ok(Some(vec![v as i32]))
1173 }
1174 fn visit_seq<A: de::SeqAccess<'de>>(
1175 self,
1176 mut seq: A,
1177 ) -> std::result::Result<Self::Value, A::Error> {
1178 let mut out = Vec::with_capacity(seq.size_hint().unwrap_or(1));
1179 while let Some(x) = seq.next_element::<i32>()? {
1180 out.push(x);
1181 }
1182 Ok(Some(out))
1183 }
1184 }
1185 de.deserialize_option(V)
1186}
1187
1188/// Count of tensor maps requested beyond the buffer's declared
1189/// [`CpuAccess`] (including any CPU map of a `CpuAccess::None` buffer).
1190/// Undeclared CPU access is a pipeline smell: it forfeits layout
1191/// optimizations (write-combined mappings, tile compression) and may be
1192/// slow or refused on Android. Read via [`unplanned_cpu_access_count`];
1193/// each offending buffer also logs one warning.
1194static UNPLANNED_CPU_ACCESS: AtomicU64 = AtomicU64::new(0);
1195
1196/// Number of tensor maps that exceeded the buffer's declared
1197/// [`CpuAccess`] since process start. A pipeline that declares its CPU
1198/// access correctly holds this flat; see [`CpuAccess`] for the contract.
1199pub fn unplanned_cpu_access_count() -> u64 {
1200 UNPLANNED_CPU_ACCESS.load(Ordering::Relaxed)
1201}
1202
1203/// Record an unplanned CPU access on `identity_id`, warning once per
1204/// buffer (a steady-state pipeline maps the same buffer every frame — a
1205/// per-map warn would flood the log; repeats count silently).
1206///
1207/// `static`-only: the sole caller is `impl<T> TensorMapTrait<T> for
1208/// Tensor<T>` in `lib.rs`, itself `#[cfg(feature = "static")]`.
1209/// [`unplanned_cpu_access_count`] stays available under `dynamic` too --
1210/// it just always reads 0 there, honestly reflecting that this backend
1211/// produces no such telemetry.
1212#[cfg(feature = "static")]
1213pub(crate) fn note_unplanned_cpu_access(identity_id: u64, backend: &str, detail: &str) {
1214 UNPLANNED_CPU_ACCESS.fetch_add(1, Ordering::Relaxed);
1215 static WARNED: std::sync::OnceLock<std::sync::Mutex<std::collections::HashSet<u64>>> =
1216 std::sync::OnceLock::new();
1217 let warned = WARNED.get_or_init(Default::default);
1218 if warned.lock().is_ok_and(|mut s| s.insert(identity_id)) {
1219 log::warn!(
1220 "unplanned CPU access on {backend} buffer (identity {identity_id}): {detail} — \
1221 declare the intent at allocation (CpuAccess::Read/Write/ReadWrite) to make \
1222 this a planned, optimized mapping"
1223 );
1224 }
1225}
1226
1227/// Uniform guard for mutable access through a read-only map. Every
1228/// `TensorMap` backend calls this from `as_mut_slice`/`deref_mut` so a
1229/// `map_read()` misuse fails identically on all platforms instead of
1230/// passing on tolerant ones and exploding only on Android.
1231#[inline]
1232pub(crate) fn assert_map_writable(writable: bool, backend: &str) {
1233 assert!(
1234 writable,
1235 "{backend} map is read-only (obtained via map_read()/CpuAccess::Read) — \
1236 use map_mut() or map_write() for mutable access"
1237 );
1238}
1239
1240/// Declared CPU involvement for an image tensor, chosen at allocation.
1241///
1242/// The HAL assumes buffers are produced and consumed by hardware (ISP,
1243/// codec, GPU, NPU) — hardware access needs no declaration. CPU access is
1244/// the opt-in: it selects the CPU usage/mapping mode at allocation
1245/// (write-combined for `Write`, cached for `Read`) and, on Android, pins
1246/// the layout linear (vendor tile compression requires `None`).
1247///
1248/// Mapping beyond the declared access is best-effort, never silent: it
1249/// may be refused ([`Error::NotImplemented`]) or take a slow path, and it
1250/// always increments [`unplanned_cpu_access_count`] with a once-per-buffer
1251/// warning.
1252#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
1253pub enum CpuAccess {
1254 /// Hardware-only buffer (the default): no CPU mapping declared.
1255 /// Compression-eligible on platforms with vendor tile layouts.
1256 #[default]
1257 None,
1258 /// CPU reads (verification, CPU consumers) — cached mapping.
1259 Read,
1260 /// CPU writes (decode targets) — write-combined mapping where the
1261 /// platform supports it; reading through a `Write` map is undeclared.
1262 Write,
1263 /// CPU reads and writes — the pre-CpuAccess implicit behavior.
1264 ReadWrite,
1265}
1266
1267impl CpuAccess {
1268 /// Whether this declaration includes CPU reads.
1269 pub fn reads(self) -> bool {
1270 matches!(self, CpuAccess::Read | CpuAccess::ReadWrite)
1271 }
1272
1273 /// Whether this declaration includes CPU writes.
1274 pub fn writes(self) -> bool {
1275 matches!(self, CpuAccess::Write | CpuAccess::ReadWrite)
1276 }
1277
1278 /// Whether this declaration covers `requested` (every direction the
1279 /// request needs is declared).
1280 pub fn covers(self, requested: CpuAccess) -> bool {
1281 (!requested.reads() || self.reads()) && (!requested.writes() || self.writes())
1282 }
1283}
1284
1285/// Requested tile-compression behavior for an image allocation (set via
1286/// [`ImageDesc::with_compression`]).
1287///
1288/// Vendor GPUs store textures in proprietary compressed tile layouts
1289/// (UBWC, AFBC, PVRIC, DCC) that cut memory bandwidth; eligibility
1290/// requires a hardware-only buffer ([`CpuAccess::None`]) because CPU
1291/// mapping pins the layout linear. The request records best knowledge on
1292/// the tensor ([`Tensor::compression`]); it never changes what bytes a
1293/// consumer sees through the GPU/NPU.
1294#[non_exhaustive]
1295#[derive(Debug, Clone, Copy, PartialEq, Eq)]
1296pub enum Compression {
1297 /// Let the platform use its native scheme when the format is
1298 /// eligible; otherwise allocate linear and count the fallback
1299 /// ([`compression_fallback_count`]). The right default for pipelines
1300 /// that want the bandwidth win without portability failures.
1301 Any,
1302 /// Require one specific scheme: allocation fails with
1303 /// [`Error::InvalidArgument`] when the device's scheme differs and
1304 /// [`Error::NotImplemented`] on platforms without vendor tile
1305 /// compression. For consumers whose ABI names a layout (e.g. a
1306 /// QNN context binary declaring UBWC inputs).
1307 Scheme(CompressionScheme),
1308}
1309
1310// Declared through `ef_vocabulary!` (see `vocabulary.rs`). This vocabulary
1311// gained a wire numbering when `ef_tensor_compression` was added: the
1312// dynamic backend reads a scheme back across the C boundary, so the enum
1313// needs codes that cannot drift from the C enumerators. Hand-writing that
1314// mapping in `edgefirst-tensor-capi` was the alternative and it is strictly
1315// worse -- `#[non_exhaustive]` makes a cross-crate `match` require a
1316// wildcard arm, so a scheme added here would silently map to "linear"
1317// there. Generated from one declaration, it cannot.
1318//
1319// Code 0 is deliberately unassigned and reserved for "linear" -- the state
1320// this crate spells `Option::None`, which has no variant to carry a code of
1321// its own. The C side names it `EF_COMPRESSION_NONE`.
1322crate::ef_vocabulary! {
1323 /// Vendor tile-compression schemes the HAL recognizes and records.
1324 ///
1325 /// Unrecognized platforms record `None` (linear) — there is deliberately
1326 /// no `Unknown` variant; a scheme is only recorded when the vendor is
1327 /// positively identified.
1328 #[non_exhaustive]
1329 pub enum CompressionScheme {
1330 /// Qualcomm Adreno Universal Bandwidth Compression.
1331 Ubwc = 1, "ubwc", UBWC,
1332 /// Arm Mali/Immortalis Framebuffer Compression.
1333 Afbc = 2, "afbc", AFBC,
1334 /// Imagination PowerVR Image Compression (Google Tensor G5+).
1335 Pvric = 3, "pvric", PVRIC,
1336 /// Samsung Xclipse (AMD RDNA) Delta Color Compression.
1337 Dcc = 4, "dcc", DCC,
1338 }
1339 // `#[doc(hidden)]` for the same reason `dtype_wire` is: emission
1340 // plumbing for `const`-only consumers, not a second documented API.
1341 #[doc(hidden)]
1342 pub mod compression_wire;
1343}
1344
1345/// Count of image allocations that requested [`Compression::Any`] but
1346/// resolved to a linear layout (ineligible format/dtype, unrecognized
1347/// vendor, or a platform without vendor tile compression). Read via
1348/// [`compression_fallback_count`].
1349static COMPRESSION_FALLBACKS: AtomicU64 = AtomicU64::new(0);
1350
1351/// Number of [`Compression::Any`] requests since process start that
1352/// resolved to a linear layout instead of a vendor tile scheme. A
1353/// steady-state pipeline holds this flat after warmup; growth means an
1354/// allocation path keeps requesting compression it never gets.
1355pub fn compression_fallback_count() -> u64 {
1356 COMPRESSION_FALLBACKS.load(Ordering::Relaxed)
1357}
1358
1359/// Record a `Compression::Any` request resolving linear.
1360///
1361/// `static`-only: the sole caller is `Tensor::image_desc` (`impl<T>
1362/// Tensor<T>`), itself `#[cfg(feature = "static")]`.
1363/// [`compression_fallback_count`] stays available under `dynamic` too --
1364/// it just always reads 0 there.
1365#[cfg(feature = "static")]
1366pub(crate) fn note_compression_fallback(detail: &str) {
1367 COMPRESSION_FALLBACKS.fetch_add(1, Ordering::Relaxed);
1368 log::debug!("compression request fell back to linear: {detail}");
1369}
1370
1371/// Whether this platform can allocate `(format, dtype)` images in a
1372/// vendor tile-compressed layout. `true` requires an Android build, an
1373/// eligible format (RGBA8888 `u8`/`i8` initially), and a positively
1374/// identified GPU vendor; everywhere else the answer is `false` and
1375/// [`Compression::Any`] requests fall back to linear.
1376pub fn compression_support(format: PixelFormat, dtype: DType) -> bool {
1377 #[cfg(target_os = "android")]
1378 {
1379 crate::ahardwarebuffer_layout::compression_eligible(format, dtype)
1380 && crate::ahardwarebuffer::device_compression_scheme().is_some()
1381 }
1382 #[cfg(not(target_os = "android"))]
1383 {
1384 let _ = (format, dtype);
1385 false
1386 }
1387}
1388
1389/// Declarative image-allocation request — the full-featured front door
1390/// for image tensors ([`TensorDyn::image_desc`] and the image crate's
1391/// `ImageProcessor::create_image_desc`).
1392///
1393/// The classic constructors (`image`, `create_image`) cover the common
1394/// cases; the desc carries the optional requests — today the
1395/// [`Compression`] request — without another constructor-parameter
1396/// sweep. Fields are private and the builders consume/return by value,
1397/// so future options are non-breaking.
1398///
1399/// ```
1400/// use edgefirst_tensor::{Compression, CpuAccess, DType, ImageDesc, PixelFormat};
1401/// let desc = ImageDesc::new(640, 640, PixelFormat::Rgba, DType::U8)
1402/// .with_access(CpuAccess::None)
1403/// .with_compression(Compression::Any);
1404/// ```
1405#[derive(Debug, Clone)]
1406pub struct ImageDesc {
1407 width: usize,
1408 height: usize,
1409 format: PixelFormat,
1410 dtype: DType,
1411 memory: Option<TensorMemory>,
1412 access: CpuAccess,
1413 compression: Option<Compression>,
1414}
1415
1416impl ImageDesc {
1417 /// A new image request: auto-selected memory, [`CpuAccess::None`]
1418 /// (hardware-only), no compression request.
1419 pub fn new(width: usize, height: usize, format: PixelFormat, dtype: DType) -> Self {
1420 Self {
1421 width,
1422 height,
1423 format,
1424 dtype,
1425 memory: None,
1426 access: CpuAccess::None,
1427 compression: None,
1428 }
1429 }
1430
1431 /// Request a specific memory backing (`None` = auto-select).
1432 pub fn with_memory(mut self, memory: Option<TensorMemory>) -> Self {
1433 self.memory = memory;
1434 self
1435 }
1436
1437 /// Declare the CPU access (see [`CpuAccess`]). Any declaration other
1438 /// than `None` makes a compression request invalid.
1439 pub fn with_access(mut self, access: CpuAccess) -> Self {
1440 self.access = access;
1441 self
1442 }
1443
1444 /// Request a tile-compressed layout (see [`Compression`]).
1445 pub fn with_compression(mut self, compression: Compression) -> Self {
1446 self.compression = Some(compression);
1447 self
1448 }
1449
1450 /// Requested width in pixels.
1451 pub fn width(&self) -> usize {
1452 self.width
1453 }
1454
1455 /// Requested height in pixels.
1456 pub fn height(&self) -> usize {
1457 self.height
1458 }
1459
1460 /// Requested pixel format.
1461 pub fn format(&self) -> PixelFormat {
1462 self.format
1463 }
1464
1465 /// Requested element type.
1466 pub fn dtype(&self) -> DType {
1467 self.dtype
1468 }
1469
1470 /// Requested memory backing (`None` = auto-select).
1471 pub fn memory(&self) -> Option<TensorMemory> {
1472 self.memory
1473 }
1474
1475 /// Declared CPU access.
1476 pub fn access(&self) -> CpuAccess {
1477 self.access
1478 }
1479
1480 /// The compression request, if any.
1481 pub fn compression(&self) -> Option<Compression> {
1482 self.compression
1483 }
1484}
1485
1486/// What kind of system key an identity was derived from. Participates in the
1487/// id so that, say, inode 7 and IOSurfaceID 7 cannot collide.
1488///
1489/// Invariant every variant must uphold: the key passed to
1490/// [`BufferIdentity::derived`] is unique among live buffers of that kind
1491/// **that own an allocation**. There is no "no identity" case — where no
1492/// system-wide key exists a process-local one does (a pointer, a GL name),
1493/// and a process-local key is exactly as sound as a system-wide one:
1494/// pointer/name reuse after free is the same ABA hazard as inode reuse, and
1495/// is mitigated the same way, by the retaining cache (see the design spec's
1496/// Task 6).
1497///
1498/// The allocation qualifier is load-bearing, not a hedge. A zero-element
1499/// tensor (`[0, num_protos]`, the tracker's "no detections this frame"
1500/// sentinel — see [`crate::mem`]) owns no allocation, so its `Vec` yields the
1501/// dangling well-aligned pointer `align_of::<T>()`. Every zero-element
1502/// `Tensor<f32>` in a process therefore shares one identity, and can even
1503/// collide across kinds with a small integer key from another variant.
1504///
1505/// That is safe rather than tolerated: an identity exists to answer "are
1506/// these the same bytes?", and buffers with no bytes have nothing to alias.
1507/// Nor can they reach an import cache. There are four GPU import paths —
1508/// a dma-buf fd on Linux, an IOSurface ref on macOS, an `AHardwareBuffer` on
1509/// Android, and an `ID3D11Texture2D` on Windows — and each takes only its own
1510/// object, so a tensor with no allocation reaches none of them. The Windows
1511/// path additionally keys on the texture the import itself holds a reference
1512/// to, and `edgefirst-image` **refuses** any texture whose identity is not
1513/// `D3d11Texture` over that same pointer — in every build, and on every
1514/// import including one served from its cache
1515/// (`GlPlatform::validate_import_identity`). That refusal exists because the
1516/// rule was once only written
1517/// down: the dynamic backend handed the image crate `HostPtr` identities over
1518/// recyclable handle addresses, and converts silently rendered into dropped
1519/// textures. Do not "fix" the zero-element case by minting a distinct id per
1520/// empty tensor: that reintroduces a counter, which is the state this whole
1521/// design exists to remove.
1522#[derive(Debug, Clone, Copy, PartialEq, Eq)]
1523pub enum IdentityKind {
1524 /// `(st_dev, st_ino)` of a dma-buf or memfd. Survives `dup`, and is the
1525 /// same value in every process that holds the buffer.
1526 DmaBuf,
1527 Shm,
1528 /// `IOSurfaceGetID`.
1529 IoSurface,
1530 /// `AHardwareBuffer_getId`, API 31+.
1531 AHardwareBufferId,
1532 /// The `AHardwareBuffer*` pointer, pre-API-31 where no id exists.
1533 /// Process-local: unique among live buffers here, meaningless elsewhere.
1534 AHardwareBufferPtr,
1535 /// An `ID3D11Texture2D*`. Process-local, like `AHardwareBufferPtr`, and
1536 /// the same pointer for every kind of texture tensor: one the HAL
1537 /// allocated, one wrapped from a caller's pointer, and one opened from a
1538 /// shared NT handle all key on the texture object they hold.
1539 ///
1540 /// **Both backends derive it from the same pointer, so their ids are
1541 /// equal byte for byte.** The static backend reads the `ComPtr` it owns
1542 /// (`d3d11::texture::tex_key`); the dynamic backend reads the same
1543 /// pointer back across the ABI with `ef_tensor_d3d11_texture`
1544 /// (`TensorDyn::derive_identity`). That equality is the point, not a
1545 /// coincidence: [`BufferIdentity::derived`] exists so two independently
1546 /// linked copies of this crate name one buffer identically, and
1547 /// `edgefirst-image` links the dynamic backend inside every `-capi` leaf
1548 /// while its own tests link the static one.
1549 ///
1550 /// **An `ef_tensor` handle's address must never be the key.** It is not a
1551 /// buffer at all, nothing the cache holds keeps it alive, and this
1552 /// library's allocator hands a freed handle's address to the very next
1553 /// one it mints — measured at 400 texture handles over 4 distinct
1554 /// addresses. The dynamic backend did exactly that on Windows until
1555 /// 2026-09-04, and a convert whose destination key collided rendered into
1556 /// the dropped texture while returning success.
1557 ///
1558 /// A handle import therefore does **not** compare equal to the tensor it
1559 /// was exported from, nor to a second import of the same handle:
1560 /// `OpenSharedResource1` mints a fresh `ID3D11Texture2D` per open. Keying
1561 /// on a handle value would make those compare equal, and would be unsound:
1562 /// Windows recycles a closed handle value straight into the next
1563 /// `DuplicateHandle`, so a consumer that caches on this id past the
1564 /// producing tensor's lifetime -- as the image crate's EGLImage cache
1565 /// does -- could serve one texture's entry for a later, unrelated import.
1566 /// A COM object's address cannot be recycled that way: a cache that holds
1567 /// a reference to the texture keeps the address taken for as long as its
1568 /// entry can be found. The cost is a cache miss per re-import instead of
1569 /// a wrong image.
1570 D3d11Texture,
1571 /// A host mapping address. Process-local.
1572 HostPtr,
1573 /// A GL buffer name. Valid only in the context that created it.
1574 Pbo,
1575}
1576
1577/// Unique identity for a tensor's underlying buffer.
1578///
1579/// Created fresh on every buffer allocation or import. The `id` is a cache
1580/// key, `kind` records what system key (if any) it was derived from, and the
1581/// `guard` is an `Arc<()>` whose weak references allow downstream caches to
1582/// detect when the buffer has been dropped.
1583#[derive(Debug, Clone)]
1584pub struct BufferIdentity {
1585 id: u64,
1586 kind: IdentityKind,
1587 guard: Arc<()>,
1588}
1589
1590impl BufferIdentity {
1591 /// Derive an identity from a key the operating system already guarantees.
1592 ///
1593 /// Deliberately NOT a counter: two independently-linked copies of this
1594 /// crate must derive the same id for the same buffer, or a consumer's
1595 /// import cache misses every frame. See the design spec, "The core finding".
1596 pub fn derived(kind: IdentityKind, key: u64) -> Self {
1597 // Mix the discriminant into the high bits so different kinds with the
1598 // same numeric key never collide.
1599 let id = ((kind as u64) << 56) ^ key;
1600 Self {
1601 id,
1602 kind,
1603 guard: Arc::new(()),
1604 }
1605 }
1606
1607 /// Unique identifier for this buffer. Changes when the buffer changes.
1608 pub fn id(&self) -> u64 {
1609 self.id
1610 }
1611
1612 /// What kind of system key this identity was derived from.
1613 pub fn kind(&self) -> IdentityKind {
1614 self.kind
1615 }
1616
1617 /// Returns a weak reference to the buffer guard. Goes dead when the
1618 /// owning Tensor is dropped (and no clones remain).
1619 pub fn weak(&self) -> Weak<()> {
1620 Arc::downgrade(&self.guard)
1621 }
1622}
1623
1624// `static`-only from here through `fs_magic` below: every use site lives in
1625// `impl<T> TensorStorage<T>` (`from_fd`'s DMA-BUF-vs-tmpfs classification),
1626// itself `#[cfg(feature = "static")]`. `dynamic`'s `from_fd`
1627// (`dynamic_backend.rs`) drives `ef_tensor_builder_wrap` instead, which
1628// does this same classification inside `libedgefirst_tensor.so`.
1629#[cfg(all(target_os = "linux", feature = "static"))]
1630use nix::sys::stat::{major, minor};
1631
1632/// Filesystem magic of the internal `dma_buf` mount, from
1633/// `include/uapi/linux/magic.h` (`DMA_BUF_MAGIC`, "DMAB").
1634///
1635/// This is stable UAPI and is the only reliable way to recognize a DMA-BUF
1636/// fd — see [`TensorStorage::from_fd`] for why `st_dev` cannot be used.
1637#[cfg(all(target_os = "linux", feature = "static"))]
1638const DMA_BUF_MAGIC: u32 = 0x444d_4142;
1639
1640/// Filesystem magic of tmpfs, from `include/uapi/linux/magic.h`
1641/// (`TMPFS_MAGIC`). Covers both POSIX `shm_open` segments under `/dev/shm`
1642/// and anonymous `memfd_create` files.
1643#[cfg(all(target_os = "linux", feature = "static"))]
1644const TMPFS_MAGIC: u32 = 0x0102_1994;
1645
1646/// Normalize a raw `fstatfs` `f_type` to the 32-bit unsigned magic that
1647/// `include/uapi/linux/magic.h` documents.
1648///
1649/// The width and signedness of `f_type` are target-dependent: `__fsword_t`
1650/// on Linux/gnu (`i64` on 64-bit, `i32` on 32-bit targets such as armv7,
1651/// i686 and aarch64-ilp32), `c_int` on uclibc, `c_ulong` on musl, and
1652/// `c_uint` on s390x. Where it is signed and 32 bits wide, widening it
1653/// sign-extends every magic with bit 31 set — hugetlbfs (`0x958458f6`),
1654/// btrfs (`0x9123683e`), f2fs (`0xf2f52010`) — which would make the value
1655/// reported by [`Error::UnknownBufferType`] impossible to find in the
1656/// header the docs point callers at.
1657///
1658/// Every magic is a 32-bit quantity, so truncating back to `u32` recovers
1659/// the true value from all four representations and loses nothing.
1660#[cfg(all(target_os = "linux", feature = "static"))]
1661const fn fs_magic(raw: i64) -> u32 {
1662 raw as u32
1663}
1664
1665/// The operations every memory backend implements.
1666///
1667/// This is the seam that lets [`Tensor<T>`] hide which backend is in play.
1668/// The implementors are `DmaTensor` (Linux DMA-BUF), `IoSurfaceTensor`
1669/// (macOS/iOS), `AHardwareBufferTensor` (Android), `ShmTensor`, `MemTensor`,
1670/// and [`PboTensor`]; `TensorStorage<T>` dispatches to whichever is active.
1671/// Backend types are not public — allocate a [`Tensor<T>`] or [`TensorDyn`]
1672/// and call these methods through it.
1673///
1674/// Import the trait to get `shape`, `size`, `map`, `clone_fd`,
1675/// `buffer_identity`, and the zero-copy sub-region `view` on a tensor value.
1676/// `Tensor::view` and `Tensor::batch` route through [`TensorTrait::view`], so
1677/// each backend's identity-sharing rule lives in exactly one place.
1678///
1679/// [`TensorDyn`]: crate::TensorDyn
1680pub trait TensorTrait<T>: Send + Sync
1681where
1682 T: Num + Clone + fmt::Debug,
1683{
1684 /// Create a new tensor with the given shape and optional name. If no name
1685 /// is given, a random name will be generated.
1686 fn new(shape: &[usize], name: Option<&str>) -> Result<Self>
1687 where
1688 Self: Sized;
1689
1690 #[cfg(unix)]
1691 /// Import an existing buffer as a tensor, taking ownership of its file
1692 /// descriptor. The buffer is adopted in place — no bytes are copied.
1693 ///
1694 /// The backend is **detected**, not chosen: the fd already belongs to a
1695 /// buffer type, and this call's job is to recognize which. On Linux that
1696 /// is decided by the fd's filesystem magic, which is stable UAPI
1697 /// (`include/uapi/linux/magic.h`):
1698 ///
1699 /// | Filesystem | Magic | Resulting [`TensorMemory`] |
1700 /// |------------|-------|----------------------------|
1701 /// | `dma_buf` | `DMA_BUF_MAGIC` (`0x444d4142`) | [`TensorMemory::DmaBuf`] |
1702 /// | `tmpfs` (`/dev/shm` **and** `memfd`) | `TMPFS_MAGIC` (`0x01021994`) | [`TensorMemory::Shm`] |
1703 /// | anything else | — | *rejected* — see Errors |
1704 ///
1705 /// Both supported types are identified **positively**. An unrecognized
1706 /// filesystem is an error, never a fallback to shared memory: the wrong
1707 /// branch does not fail loudly (a DMA-BUF is mmap-able, so it would
1708 /// import as a perfectly functional tensor that merely isn't DMA, and a
1709 /// pipe would import as a zero-length one), so guessing would trade a
1710 /// clear error here for silent loss of zero-copy far downstream.
1711 ///
1712 /// The device number is deliberately **not** consulted. `dma_buf` files
1713 /// live on an internal kernel mount whose `st_dev` comes from
1714 /// `get_anon_bdev()` — an IDA shared with every other anonymous
1715 /// pseudo-filesystem and allocated in boot order — so the minor a
1716 /// DMA-BUF lands on varies by kernel build and is not part of any ABI.
1717 ///
1718 /// On non-Linux Unix (macOS/iOS/Android) there is no fd-based DMA import;
1719 /// the fd is always adopted as [`TensorMemory::Shm`].
1720 ///
1721 /// # Arguments
1722 ///
1723 /// * `fd` - Owned descriptor for the buffer to import. Ownership
1724 /// transfers to the returned tensor and the fd is closed on drop; pass
1725 /// [`clone_fd`](TensorTrait::clone_fd) output to keep your own handle.
1726 /// * `shape` - Logical shape to interpret the buffer with. Must describe
1727 /// no more elements than the buffer holds.
1728 /// * `name` - Optional name; a random one is generated when `None`.
1729 ///
1730 /// # Returns
1731 ///
1732 /// A tensor sharing the imported buffer's memory, whose
1733 /// [`memory()`](TensorTrait::memory) reports the detected backend.
1734 ///
1735 /// # Errors
1736 ///
1737 /// * [`Error::UnknownBufferType`] - the fd is on a filesystem that is
1738 /// neither `dma_buf` nor `tmpfs`, so its buffer type cannot be
1739 /// determined. Carries the observed `fstatfs` magic as a `u32`,
1740 /// normalized so it can be looked up in `magic.h` directly on both
1741 /// 32- and 64-bit targets. Typical causes: a regular file, a pipe or
1742 /// socket, or a `MFD_HUGETLB` memfd (hugetlbfs, not tmpfs). Linux only.
1743 /// * [`Error::UnknownDeviceType`] - the fd's `st_dev` major is non-zero,
1744 /// i.e. it lives on a real block device rather than an anonymous or
1745 /// in-memory filesystem. Linux only.
1746 /// * [`Error::InvalidSize`] - `shape` is empty or describes zero
1747 /// elements.
1748 /// * [`Error::NixError`] - `fstat`, `fstatfs`, or `mmap` failed on the
1749 /// descriptor.
1750 ///
1751 /// # Examples
1752 ///
1753 /// ```no_run
1754 /// use edgefirst_tensor::{Tensor, TensorMemory, TensorTrait};
1755 ///
1756 /// # fn main() -> edgefirst_tensor::Result<()> {
1757 /// let src = Tensor::<u8>::new(&[480, 640, 3], Some(TensorMemory::DmaBuf), None)?;
1758 ///
1759 /// // Round-tripping a DMA-BUF fd preserves the backend — the import is
1760 /// // still zero-copy, and still eligible for GPU/NPU paths.
1761 /// let imported = Tensor::<u8>::from_fd(src.clone_fd()?, src.shape(), None)?;
1762 /// assert_eq!(imported.memory(), TensorMemory::DmaBuf);
1763 /// # Ok(())
1764 /// # }
1765 /// ```
1766 fn from_fd(fd: std::os::fd::OwnedFd, shape: &[usize], name: Option<&str>) -> Result<Self>
1767 where
1768 Self: Sized;
1769
1770 #[cfg(unix)]
1771 /// Clone the file descriptor associated with this tensor.
1772 fn clone_fd(&self) -> Result<std::os::fd::OwnedFd>;
1773
1774 /// Get the memory type of this tensor.
1775 fn memory(&self) -> TensorMemory;
1776
1777 /// Get the name of this tensor.
1778 fn name(&self) -> String;
1779
1780 /// Get the number of elements in this tensor.
1781 fn len(&self) -> usize {
1782 self.shape().iter().product()
1783 }
1784
1785 /// Check if the tensor is empty.
1786 fn is_empty(&self) -> bool {
1787 self.len() == 0
1788 }
1789
1790 /// Get the size in bytes of this tensor.
1791 fn size(&self) -> usize {
1792 self.len() * std::mem::size_of::<T>()
1793 }
1794
1795 /// Get the shape of this tensor.
1796 fn shape(&self) -> &[usize];
1797
1798 /// Reshape this tensor to the given shape. The total number of elements
1799 /// must remain the same.
1800 fn reshape(&mut self, shape: &[usize]) -> Result<()>;
1801
1802 /// Bytes of the underlying allocation (>= the current logical `size()`).
1803 fn capacity_bytes(&self) -> usize;
1804
1805 /// Set the logical shape to any shape whose byte size fits the allocation
1806 /// capacity, without the equal-size constraint of `reshape`.
1807 fn set_logical_shape(&mut self, shape: &[usize]) -> Result<()>;
1808
1809 /// Map the tensor into memory with the given access direction and
1810 /// return a TensorMap for accessing the data.
1811 ///
1812 /// `access` selects the platform mapping mode (read-only IOSurface
1813 /// lock, dma-buf sync direction, AHardwareBuffer lock usage) and the
1814 /// map's mutability: a map obtained with [`CpuAccess::Read`] rejects
1815 /// `as_mut_slice`. [`CpuAccess::None`] is not a mappable direction
1816 /// and returns [`Error::InvalidArgument`].
1817 ///
1818 /// Prefer the typed wrappers [`map_read`](Self::map_read) /
1819 /// [`map_write`](Self::map_write) / [`map_mut`](Self::map_mut).
1820 /// Map for CPU access.
1821 ///
1822 /// Like [`pin_host`](Tensor::pin_host), `'a` bounds how long the backing
1823 /// allocation lives rather than borrowing the tensor: the returned view
1824 /// shares ownership through the pin's keepalive, so a map can coexist with
1825 /// a `&mut` on the same tensor.
1826 ///
1827 /// # Windows D3D11 textures: a write-only map of a partial window
1828 ///
1829 /// A [`CpuAccess::Write`] map is not refreshed from the texture first, so
1830 /// the bytes the caller leaves untouched are undefined -- the documented
1831 /// contract a whole-tensor map can meet and a sub-view cannot, since the
1832 /// window spans its own rows only.
1833 ///
1834 /// A tensor allocated [`CpuAccess::Write`] (no staging texture) publishes
1835 /// only the window's own rows on unmap, so a sub-view -- a
1836 /// [`view`](Tensor::view), or a shape narrowed with `set_logical_shape`
1837 /// -- is refused only when its window is not a whole, pitch-aligned row
1838 /// range (a plain-offset [`view`](Tensor::view) always is; an arbitrary
1839 /// byte offset from `set_plane_offset` might not be), which fails with
1840 /// [`Error::InvalidArgument`].
1841 ///
1842 /// A tensor allocated [`CpuAccess::Read`] or [`CpuAccess::ReadWrite`]
1843 /// (staging texture) has no such per-window publish: unmapping always
1844 /// copies the whole staging texture back, so a write-only sub-view is
1845 /// refused with [`Error::PartialWriteRequiresReadWrite`], naming
1846 /// [`CpuAccess::ReadWrite`] as the same window with the refresh that
1847 /// makes the untouched rows well-defined -- and, unlike the
1848 /// `CpuAccess::Write` case above, that retry always succeeds.
1849 fn map_with<'a>(&self, access: CpuAccess) -> Result<crate::view::HostView<'a, T>>
1850 where
1851 T: 'a;
1852
1853 /// Map the tensor read-write (equivalent to
1854 /// `map_with(CpuAccess::ReadWrite)` — the historical `map()`
1855 /// behavior).
1856 fn map<'a>(&self) -> Result<crate::view::HostView<'a, T>>
1857 where
1858 T: 'a,
1859 {
1860 self.map_with(CpuAccess::ReadWrite)
1861 }
1862
1863 /// Map the tensor for CPU reading only. The returned map rejects
1864 /// `as_mut_slice`; on macOS this takes the read-only IOSurface lock
1865 /// (skips the unlock flush), on Linux the dma-buf read-direction
1866 /// sync.
1867 fn map_read<'a>(&self) -> Result<crate::view::HostView<'a, T>>
1868 where
1869 T: 'a,
1870 {
1871 self.map_with(CpuAccess::Read)
1872 }
1873
1874 /// Map the tensor for CPU writing (fill-only: reading through a
1875 /// write map may see write-combined memory — do not read the slice).
1876 fn map_write<'a>(&self) -> Result<crate::view::HostView<'a, T>>
1877 where
1878 T: 'a,
1879 {
1880 self.map_with(CpuAccess::Write)
1881 }
1882
1883 /// Map the tensor read-write (alias of [`map`](Self::map) with the
1884 /// intent spelled out).
1885 fn map_mut<'a>(&self) -> Result<crate::view::HostView<'a, T>>
1886 where
1887 T: 'a,
1888 {
1889 self.map_with(CpuAccess::ReadWrite)
1890 }
1891
1892 /// Non-blocking [`map_with`](Self::map_with): returns
1893 /// `Err(Error::IoError(WouldBlock))` while a GPU copy the map depends on
1894 /// is still in flight, and makes progress on a retry.
1895 ///
1896 /// Backings with no such copy (host memory, shared memory, dma-buf,
1897 /// IOSurface, AHardwareBuffer, PBO) alias `map_with`, so a caller can use
1898 /// this unconditionally and only ever see `WouldBlock` where the wait is
1899 /// real -- today the Windows D3D11 texture's staging refresh.
1900 ///
1901 /// A retry loop must yield or sleep between attempts. On the WARP
1902 /// software adapter the CPU threads are the GPU, so a tight loop starves
1903 /// the very copy it is waiting for.
1904 ///
1905 /// [`CpuAccess::Write`] is refused on a Windows D3D11 texture tensor whose
1906 /// window is shorter than its backing, exactly as it is for
1907 /// [`map_with`](Self::map_with); see that method.
1908 fn try_map_with<'a>(&self, access: CpuAccess) -> Result<crate::view::HostView<'a, T>>
1909 where
1910 T: 'a,
1911 {
1912 self.map_with(access)
1913 }
1914
1915 /// Non-blocking [`map`](Self::map): read-write, and `WouldBlock` rather
1916 /// than a stall. See [`try_map_with`](Self::try_map_with).
1917 fn try_map<'a>(&self) -> Result<crate::view::HostView<'a, T>>
1918 where
1919 T: 'a,
1920 {
1921 self.try_map_with(CpuAccess::ReadWrite)
1922 }
1923
1924 /// Get the buffer identity for cache keying and liveness tracking.
1925 fn buffer_identity(&self) -> &BufferIdentity;
1926
1927 /// Create a zero-copy sub-region view of this backing that shares the
1928 /// underlying allocation **and** [`BufferIdentity`].
1929 ///
1930 /// The window is `[offset_bytes, offset_bytes + shape.product() *
1931 /// size_of::<T>())` measured from this tensor's own logical start, so a
1932 /// sub-view of a sub-view composes by adding offsets. Sharing the parent's
1933 /// identity is the contract that lets identity-keyed caches (e.g. the GL
1934 /// EGLImage import cache) treat offset-distinct windows as one buffer rather
1935 /// than unrelated allocations — `view` must never mint a fresh identity.
1936 ///
1937 /// Defaults to [`Error::NotImplemented`]; every backend that supports
1938 /// sub-views overrides it (`Mem`, `Shm`, Linux DMA, macOS IOSurface, `Pbo`).
1939 fn view(&self, offset_bytes: usize, shape: &[usize]) -> Result<Self>
1940 where
1941 Self: Sized,
1942 {
1943 let _ = (offset_bytes, shape);
1944 Err(Error::NotImplemented(
1945 "view (zero-copy sub-region) is not supported for this tensor backend".to_owned(),
1946 ))
1947 }
1948}
1949
1950pub trait TensorMapTrait<T>
1951where
1952 T: Num + Clone + fmt::Debug,
1953{
1954 /// Get the shape of this tensor map.
1955 fn shape(&self) -> &[usize];
1956
1957 /// Unmap the tensor from memory.
1958 fn unmap(&mut self);
1959
1960 /// Get the number of elements in this tensor map.
1961 fn len(&self) -> usize {
1962 self.shape().iter().product()
1963 }
1964
1965 /// Check if the tensor map is empty.
1966 fn is_empty(&self) -> bool {
1967 self.len() == 0
1968 }
1969
1970 /// Get the size in bytes of this tensor map.
1971 fn size(&self) -> usize {
1972 self.len() * std::mem::size_of::<T>()
1973 }
1974
1975 /// Get a slice to the data in this tensor map.
1976 fn as_slice(&self) -> &[T];
1977
1978 /// Get a mutable slice to the data in this tensor map.
1979 fn as_mut_slice(&mut self) -> &mut [T];
1980
1981 /// Whether this mapping permits writes -- i.e. whether the coherency
1982 /// bracket it holds will flush on release.
1983 ///
1984 /// The single source of truth for a binding that has to advertise
1985 /// mutability to a foreign consumer (the Python buffer protocol's
1986 /// `readonly` flag, say). Deriving that from the access a caller
1987 /// *requested* instead would leave the advertisement agreeing with the
1988 /// request while both disagreed with the mapping actually taken.
1989 fn is_writable(&self) -> bool;
1990
1991 #[cfg(feature = "ndarray")]
1992 /// Get an ndarray ArrayView of the tensor data.
1993 fn view(&'_ self) -> Result<ndarray::ArrayView<'_, T, ndarray::Dim<ndarray::IxDynImpl>>> {
1994 Ok(ndarray::ArrayView::from_shape(
1995 self.shape(),
1996 self.as_slice(),
1997 )?)
1998 }
1999
2000 #[cfg(feature = "ndarray")]
2001 /// Get an ndarray ArrayViewMut of the tensor data.
2002 fn view_mut(
2003 &'_ mut self,
2004 ) -> Result<ndarray::ArrayViewMut<'_, T, ndarray::Dim<ndarray::IxDynImpl>>> {
2005 let shape = self.shape().to_vec();
2006 Ok(ndarray::ArrayViewMut::from_shape(
2007 shape,
2008 self.as_mut_slice(),
2009 )?)
2010 }
2011}
2012
2013// Declared through `ef_vocabulary!` (see `vocabulary.rs`). Before this, the
2014// same four backings were numbered three different ways -- Rust's implicit
2015// discriminants (`Dma=0, Shm=1, Mem=2, Pbo=3`), the C ABI's
2016// `hal_tensor_memory` (`MEM=0, DMA=1, SHM=2, PBO=3`), and Python's
2017// `PyTensorMemory`, which had no explicit discriminants at all and a
2018// `#[cfg(unix)]` variant in the middle of the list, so its numbering
2019// depended on the target OS. The canonical assignment below is the one all
2020// of them converge on; `hal_tensor_memory` is the single explicitly-mapped
2021// outlier that keeps its own released values (see `to_hal_tensor_memory` in
2022// the capi crate).
2023//
2024// Note this is NOT `protocol::kind`, which is a genuinely different
2025// vocabulary and stays so: `kind` describes what a `TensorDesc`'s `handle`
2026// and `ptr` fields MEAN to an importer, and both `Mem` and `Shm` mean the
2027// same thing there (`kind::HOST`, host-addressable pointer, no handle). A
2028// vocabulary that collapses two variants into one cannot also be the
2029// vocabulary that distinguishes them. `kind_of()` in `protocol.rs` is the
2030// mapping between them.
2031crate::ef_vocabulary! {
2032 /// Which memory backend a tensor is (or should be) allocated from.
2033 ///
2034 /// Pass `Some(..)` to a constructor to pin the backend, or `None` to
2035 /// auto-select. Auto-selection differs by constructor: [`Tensor::new`]
2036 /// tries `DmaBuf` → `Mem`, while the image constructors try `DmaBuf` →
2037 /// `Mem` and skip `Shm`. `EDGEFIRST_TENSOR_FORCE_MEM=1` short-circuits
2038 /// either chain to `Mem`, which is how tests run on hosts without
2039 /// DMA-heap permissions.
2040 ///
2041 /// A pinned request has no fallback: if the backend cannot serve it, the
2042 /// constructor fails rather than quietly giving you something slower.
2043 /// [`TensorTrait::memory`] reports what a tensor actually ended up with.
2044 ///
2045 /// # Every variant exists on every platform
2046 ///
2047 /// This is a namespace of codes, not a list of what this build can
2048 /// allocate -- [`TensorMemory::is_available`] answers the latter, at
2049 /// runtime. `Shm` used to be `#[cfg(unix)]`, which shifted every
2050 /// variant declared after it on a non-unix build and made the
2051 /// numbering a compile-time detail of the host. It also made a
2052 /// cross-platform recording unreadable: a blob captured on Linux
2053 /// carries `"dmabuf"`, and a macOS consumer has to be able to NAME
2054 /// that backing in order to report "cannot materialise here" instead
2055 /// of failing with an unrecognised code.
2056 ///
2057 /// # `#[non_exhaustive]`
2058 ///
2059 /// Match this with a wildcard arm. The variant list is expected to
2060 /// keep growing -- `IoSurface` and `Cuda` were added here, Windows
2061 /// backings and a real CUDA device backing are named as coming -- and
2062 /// this attribute is what keeps a closed-source consumer's
2063 /// `match memory { .. }` compiling across those additions. It protects
2064 /// DOWNSTREAM, not this crate: inside `edgefirst-tensor` the matches
2065 /// stay exhaustive and compile-enforced. Same rule as [`DType`],
2066 /// [`Compression`] and [`PixelFormat`], which `TensorMemory` was the
2067 /// lone exception to.
2068 #[non_exhaustive]
2069 pub enum TensorMemory {
2070 /// Regular system memory allocation. Available everywhere; the
2071 /// floor every auto-selection chain ends on.
2072 Mem = 0, "mem", MEM,
2073 /// POSIX Shared Memory allocation. Suitable for inter-process
2074 /// communication, but not for hardware acceleration. Only
2075 /// allocatable on unix; nameable everywhere.
2076 Shm = 1, "shm", SHM,
2077 /// Platform-native zero-copy GPU buffer.
2078 ///
2079 /// On Linux this is a DMA-BUF (`DmaTensor` in
2080 /// `crates/tensor/src/dma.rs`) allocated via the DRM/dma-heap
2081 /// subsystem. On macOS/iOS it is currently also this variant that
2082 /// yields an IOSurface, and on Android an AHardwareBuffer: they
2083 /// share the `TensorStorage::Dma` slot at the trait level, and the
2084 /// public C API discriminant (`HAL_TENSOR_MEMORY_DMA = 1`) covers
2085 /// all three with no ABI break. `IoSurface` below is the code for
2086 /// naming that backing specifically; no backend reports it yet.
2087 ///
2088 /// Allows hardware-accelerated paths (OpenGL backend on Linux via
2089 /// `EGL_EXT_image_dma_buf_import`; macOS via
2090 /// `EGL_ANGLE_iosurface_client_buffer`). CPU access via `map()`
2091 /// incurs cache-coherency overhead on Linux DMA-BUF and is similar
2092 /// in cost on IOSurface; SHM/Mem are cheaper for CPU-only workloads.
2093 ///
2094 /// On Windows it is an `ID3D11Texture2D` on the process D3D11 device
2095 /// (`crates/tensor/src/d3d11/`), allocated only through
2096 /// [`Tensor::image`](crate::Tensor::image); CPU access goes through a
2097 /// staging copy. The first such allocation creates the process
2098 /// device, so a host supplying its own device must call
2099 /// `d3d11::use_external_device` before it -- and before
2100 /// [`crate::is_gpu_buffer_available`], which creates the device too.
2101 DmaBuf = 2, "dmabuf", DMABUF,
2102 /// Apple IOSurface, named specifically rather than through the
2103 /// portable `DmaBuf` spelling.
2104 ///
2105 /// Defined as a code so a descriptor or recording that names an
2106 /// IOSurface is parseable on every platform. **No backend produces
2107 /// or accepts it yet** -- macOS/iOS allocate and report `DmaBuf`.
2108 /// Pinning a request to it fails with `NotImplemented`.
2109 IoSurface = 3, "iosurface", IOSURFACE,
2110 /// OpenGL Pixel Buffer Object memory. Created by `ImageProcessor`
2111 /// when DMA-buf is unavailable but OpenGL is present -- never by
2112 /// [`Tensor::new`], which has no GL context to create it in.
2113 Pbo = 4, "pbo", PBO,
2114 /// CUDA device memory: a device pointer, not host-addressable.
2115 ///
2116 /// Defined as a code for the same reason as `IoSurface`. **No
2117 /// backend produces or accepts it yet**; pinning a request to it
2118 /// fails with `NotImplemented`.
2119 Cuda = 5, "cuda", CUDA,
2120 }
2121 // `#[doc(hidden)]`: `pub` so the const-only form is reachable by an FFI
2122 // or cbindgen consumer, but it is emission plumbing rather than a second
2123 // documented API -- `TensorMemory::code()` is the documented path, and
2124 // the two are generated from the same literal.
2125 #[doc(hidden)]
2126 pub mod tensor_memory_wire;
2127}
2128
2129impl TensorMemory {
2130 /// Whether this backing can actually be allocated in this process, now.
2131 ///
2132 /// A runtime question, not a compile-time one, which is why it is a
2133 /// method on the vocabulary rather than a `cfg` on the variant:
2134 /// `DmaBuf` needs `/dev/dma_heap` to exist *and* this process to have
2135 /// permission on it (a container, a user outside the `video`/`render`
2136 /// groups, and a stock desktop give three different answers on the same
2137 /// kernel), and `Shm` needs a writable `/dev/shm`. Each probe allocates
2138 /// once and caches.
2139 ///
2140 /// Reports what can be *allocated*, which is not the same as what can be
2141 /// *named*: every variant is nameable on every platform (that is the
2142 /// point of the vocabulary), and the ones this build cannot serve say so
2143 /// here rather than by being absent from the enum.
2144 ///
2145 /// Three variants answer `false` unconditionally today, for reasons
2146 /// worth stating rather than hiding behind a probe that cannot fail:
2147 ///
2148 /// * `IoSurface` -- no backend produces or accepts it yet; macOS/iOS
2149 /// serve IOSurface under the `DmaBuf` spelling, so `DmaBuf` is the
2150 /// variant that answers `true` there.
2151 /// * `Pbo` -- PBO tensors are created by `ImageProcessor::create_image`
2152 /// with a GL context current, never by [`Tensor::new`]. This crate
2153 /// holds no GL context and cannot see one; ask the image crate
2154 /// whether a GL backend is live.
2155 /// * `Cuda` -- there is no CUDA backing in this crate at all. Probing
2156 /// for `libcuda` would answer `true` on a Jetson while every
2157 /// allocation still failed, which is worse than answering `false`.
2158 ///
2159 /// One variant answers `true` for a narrower constructor than the rest:
2160 /// on Windows `DmaBuf` is an `ID3D11Texture2D`, which is image-formatted,
2161 /// so `true` there means [`Tensor::image`] can serve it and
2162 /// [`Tensor::new`] with `Some(DmaBuf)` still answers `NotImplemented`.
2163 /// [`is_gpu_buffer_available`](crate::is_gpu_buffer_available), which this
2164 /// forwards to, states the same split.
2165 pub fn is_available(self) -> bool {
2166 match self {
2167 // Always: the fallback every constructor chain ends on. If this
2168 // could fail, allocation would have no floor left.
2169 TensorMemory::Mem => true,
2170 TensorMemory::Shm => is_shm_available(),
2171 // Today `DmaBuf` is the portable spelling of "platform-native
2172 // zero-copy GPU buffer", so it is the portable probe that
2173 // answers for it -- true on macOS/iOS (IOSurface) and Android
2174 // (AHardwareBuffer) as well as Linux.
2175 TensorMemory::DmaBuf => is_gpu_buffer_available(),
2176 TensorMemory::IoSurface | TensorMemory::Pbo | TensorMemory::Cuda => false,
2177 }
2178 }
2179
2180 /// The error to return when a caller pins a request to a backing this
2181 /// build knows the name of but cannot serve.
2182 ///
2183 /// Deliberately not `unreachable!()`. `IoSurface` and `Cuda` are
2184 /// unreachable *today* only because no backend reports them yet; the
2185 /// day one does, an `unreachable!()` written now becomes a panic in
2186 /// library code reached from an ordinary caller request. A "this build
2187 /// cannot do that" error is the same information without the crash, and
2188 /// it is what every other unservable request here already returns.
2189 ///
2190 /// `static`-only: every call site is inside `impl<T> TensorStorage<T>`
2191 /// (`TensorStorage::new`), itself `#[cfg(feature = "static")]`.
2192 #[cfg(feature = "static")]
2193 pub(crate) fn unsupported_here(self) -> Error {
2194 let reason = match self {
2195 // Not reachable: `Mem` is servable on every target. Stated
2196 // rather than panicked on, for the reason above.
2197 TensorMemory::Mem => "Mem is servable on every platform, so this should not happen",
2198 TensorMemory::Shm => "POSIX shared memory is unix-only",
2199 TensorMemory::DmaBuf => {
2200 "this target has no platform GPU-buffer allocator (DMA-BUF is Linux, \
2201 IOSurface macOS/iOS, AHardwareBuffer Android, D3D11 textures on \
2202 Windows)"
2203 }
2204 TensorMemory::IoSurface => {
2205 "no backend produces or accepts IOSurface under its own code yet -- \
2206 macOS/iOS allocate and report TensorMemory::DmaBuf"
2207 }
2208 TensorMemory::Pbo => {
2209 "PBO tensors are created by ImageProcessor::create_image with a GL \
2210 context current, never by Tensor::new"
2211 }
2212 TensorMemory::Cuda => "there is no CUDA backing in this build",
2213 };
2214 Error::NotImplemented(format!(
2215 "TensorMemory::{self:?} is not supported by this build: {reason}"
2216 ))
2217 }
2218}
2219
2220impl From<TensorMemory> for String {
2221 /// The wire string, via [`TensorMemory::as_str`] -- one spelling, not a
2222 /// second one maintained alongside it.
2223 ///
2224 /// `DmaBuf` renders as `"dmabuf"`, not the `"dma"` this impl used to
2225 /// emit. `TryFrom<&str>` still accepts `"dma"` so anything that
2226 /// persisted the old spelling keeps parsing; the benchmark result keys
2227 /// in `.github/scripts` spell it `"dma"` but build that string
2228 /// themselves and never went through this impl.
2229 fn from(memory: TensorMemory) -> Self {
2230 memory.as_str().to_owned()
2231 }
2232}
2233
2234impl TryFrom<&str> for TensorMemory {
2235 type Error = Error;
2236
2237 /// Parses [`TensorMemory::as_str`]'s spelling, plus the legacy `"dma"`
2238 /// alias for `DmaBuf`.
2239 fn try_from(s: &str) -> Result<Self> {
2240 TensorMemory::from_str_code(s)
2241 .or(if s == "dma" {
2242 Some(TensorMemory::DmaBuf)
2243 } else {
2244 None
2245 })
2246 .ok_or_else(|| Error::InvalidMemoryType(s.to_owned()))
2247 }
2248}
2249
2250// Everything from here through the end of `impl<T> TensorTrait<T> for
2251// Tensor<T>` below is the `static` backend's storage-owning `Tensor<T>`:
2252// `TensorStorage<T>` (the per-memory-backend enum) and the `Tensor<T>`
2253// struct that wraps it. `dynamic`'s `Tensor<T>` is a different, much
2254// smaller type -- a typed lens over `TensorDyn` with no storage of its
2255// own -- defined in `tensor_dyn/dynamic_tensor.rs` and re-exported below.
2256#[cfg(feature = "static")]
2257#[derive(Debug)]
2258#[allow(dead_code)] // Variants are constructed by downstream crates via pub(crate) helpers
2259pub(crate) enum TensorStorage<T>
2260where
2261 T: Num + Clone + fmt::Debug + Send + Sync,
2262{
2263 /// Platform-native zero-copy GPU buffer. Inner type differs per
2264 /// target: `DmaTensor` on Linux (DMA-BUF fd), `IoSurfaceTensor` on
2265 /// macOS (CFRetained IOSurface). The shared variant name keeps the
2266 /// public `TensorMemory::DmaBuf` discriminant stable across platforms.
2267 #[cfg(target_os = "linux")]
2268 Dma(DmaTensor<T>),
2269 #[cfg(any(target_os = "macos", target_os = "ios"))]
2270 Dma(IoSurfaceTensor<T>),
2271 #[cfg(target_os = "android")]
2272 Dma(AHardwareBufferTensor<T>),
2273 #[cfg(target_os = "windows")]
2274 Dma(crate::d3d11::texture::D3d11TextureTensor<T>),
2275 #[cfg(unix)]
2276 Shm(ShmTensor<T>),
2277 Mem(MemTensor<T>),
2278 Pbo(PboTensor<T>),
2279}
2280
2281#[cfg(feature = "static")]
2282impl<T> TensorStorage<T>
2283where
2284 T: Num + Clone + fmt::Debug + Send + Sync,
2285{
2286 /// The backing allocation's intrinsic physical row pitch in bytes, if it
2287 /// has one that is fixed independent of the logical shape. macOS IOSurface
2288 /// reports its 64-aligned `bytesPerRow`; other backings (Linux DMA, SHM,
2289 /// Mem, PBO) have no fixed pitch beyond the logical shape and return `None`.
2290 ///
2291 /// Used by `configure_image` to preserve the physical pitch when a reused
2292 /// pool tensor is reconfigured to a smaller logical image — so the decode
2293 /// writes rows at the surface's real stride and the GPU samples them with
2294 /// the same stride (the physical-grid / logical-ROI decoupling).
2295 ///
2296 /// Compiled only where such a backing exists; on Linux and other unix
2297 /// targets it would always answer `None`.
2298 #[cfg(any(
2299 target_os = "macos",
2300 target_os = "ios",
2301 target_os = "android",
2302 target_os = "windows"
2303 ))]
2304 pub(crate) fn backing_row_stride(&self) -> Option<usize> {
2305 match self {
2306 // Only genuine image-formatted IOSurfaces (height > 1) carry a real
2307 // per-row pitch; a generic byte-bag (height == 1) returns `None` so
2308 // `configure_image` does not adopt its whole-buffer "row" as a stride.
2309 #[cfg(any(target_os = "macos", target_os = "ios"))]
2310 TensorStorage::Dma(t) => t.image_backing_row_stride(),
2311 // Android AHardwareBuffer: same rule — only genuine 2D
2312 // image-formatted buffers (height > 1) carry a real row pitch.
2313 #[cfg(target_os = "android")]
2314 TensorStorage::Dma(t) => t.image_backing_row_stride(),
2315 // Windows D3D11: the driver's staging row pitch. `None` when the
2316 // tensor has no staging (write-only or hardware-only), whose host
2317 // buffer has tight rows -- except for a semi-planar tensor, whose
2318 // texture is allocated as wide as that pitch for every access, so
2319 // it always has one to report.
2320 #[cfg(target_os = "windows")]
2321 TensorStorage::Dma(t) => t.image_backing_row_stride(),
2322 _ => None,
2323 }
2324 }
2325
2326 /// Create a new tensor storage with the given shape, memory type, and
2327 /// optional name. If no name is given, a random name will be generated.
2328 /// If no memory type is given, the best available memory type will be
2329 /// chosen based on the platform and environment variables.
2330 fn new(shape: &[usize], memory: Option<TensorMemory>, name: Option<&str>) -> Result<Self> {
2331 match memory {
2332 #[cfg(target_os = "linux")]
2333 Some(TensorMemory::DmaBuf) => {
2334 DmaTensor::<T>::new(shape, name).map(TensorStorage::Dma)
2335 }
2336 #[cfg(any(target_os = "macos", target_os = "ios"))]
2337 Some(TensorMemory::DmaBuf) => {
2338 IoSurfaceTensor::<T>::new(shape, name).map(TensorStorage::Dma)
2339 }
2340 #[cfg(target_os = "android")]
2341 Some(TensorMemory::DmaBuf) => {
2342 AHardwareBufferTensor::<T>::new(shape, name).map(TensorStorage::Dma)
2343 }
2344 // Windows has a DmaBuf backing, but it is a D3D11 texture: an image
2345 // layout, not a byte bag, so a shape alone cannot describe one.
2346 // Name the constructor that can, rather than quietly handing back
2347 // host memory the caller did not ask for.
2348 #[cfg(target_os = "windows")]
2349 Some(TensorMemory::DmaBuf) => Err(crate::error::Error::NotImplemented(
2350 "D3D11 texture tensors are image-formatted; use Tensor::image with \
2351 TensorMemory::DmaBuf"
2352 .to_owned(),
2353 )),
2354 #[cfg(not(any(
2355 target_os = "linux",
2356 target_os = "macos",
2357 target_os = "ios",
2358 target_os = "android",
2359 target_os = "windows"
2360 )))]
2361 Some(TensorMemory::DmaBuf) => Err(crate::error::Error::NotImplemented(
2362 "TensorMemory::DmaBuf is only available on Linux (DMA-BUF), macOS/iOS (IOSurface), \
2363 Android (AHardwareBuffer), and Windows (D3D11 textures, via Tensor::image)"
2364 .to_owned(),
2365 )),
2366 #[cfg(unix)]
2367 Some(TensorMemory::Shm) => {
2368 ShmTensor::<T>::new(shape, name).map(TensorStorage::Shm)
2369 }
2370 #[cfg(not(unix))]
2371 Some(m @ TensorMemory::Shm) => Err(m.unsupported_here()),
2372 Some(TensorMemory::Mem) => {
2373 MemTensor::<T>::new(shape, name).map(TensorStorage::Mem)
2374 }
2375 Some(TensorMemory::Pbo) => Err(crate::error::Error::NotImplemented(
2376 "PboTensor cannot be created via Tensor::new() — use ImageProcessor::create_image()".to_owned(),
2377 )),
2378 // Defined codes with no backing behind them yet. An error, not
2379 // an `unreachable!()` -- see `unsupported_here`.
2380 Some(m @ (TensorMemory::IoSurface | TensorMemory::Cuda)) => Err(m.unsupported_here()),
2381 None => {
2382 if std::env::var("EDGEFIRST_TENSOR_FORCE_MEM")
2383 .is_ok_and(|x| x != "0" && x.to_lowercase() != "false")
2384 {
2385 MemTensor::<T>::new(shape, name).map(TensorStorage::Mem)
2386 } else {
2387 // Auto-select priority: Dma > Mem. Shm is intentionally NOT
2388 // auto-selected — it offers no advantage over Mem for an
2389 // in-process tensor and Mem always succeeds, so Shm below Mem
2390 // is effectively never reached. Request Shm explicitly via
2391 // `TensorMemory::Shm` when cross-process sharing is needed.
2392 // (PBO sits between Dma and Mem but is GL-backed and created
2393 // only via `ImageProcessor::create_image`.)
2394 #[cfg(target_os = "linux")]
2395 {
2396 // Linux: Try DMA -> Mem
2397 match DmaTensor::<T>::new(shape, name) {
2398 Ok(tensor) => Ok(TensorStorage::Dma(tensor)),
2399 Err(_) => MemTensor::<T>::new(shape, name).map(TensorStorage::Mem),
2400 }
2401 }
2402 #[cfg(any(target_os = "macos", target_os = "ios"))]
2403 {
2404 // macOS/iOS: Try IOSurface -> Mem. IOSurface is the
2405 // GPU-shareable backend (zero-copy via ANGLE), filling the
2406 // same role as DMA-BUF on Linux.
2407 match IoSurfaceTensor::<T>::new(shape, name) {
2408 Ok(tensor) => Ok(TensorStorage::Dma(tensor)),
2409 Err(_) => MemTensor::<T>::new(shape, name).map(TensorStorage::Mem),
2410 }
2411 }
2412 #[cfg(target_os = "android")]
2413 {
2414 // Android: Mem. Unlike macOS (where a byte-bag
2415 // IOSurface is still GL-importable as R8) a generic
2416 // BLOB AHardwareBuffer cannot back a GPU texture, so
2417 // auto-selecting it buys nothing and costs plenty: a
2418 // gralloc allocator-HAL ioctl per allocation (orders
2419 // slower than malloc, ≥1 page + a dmabuf fd even for
2420 // tiny tensors) and a lock/unlock cache-maintenance
2421 // round trip on every map(). Zero-copy image tensors
2422 // come from `Tensor::image(..)`; callers that want a
2423 // BLOB (NNAPI handoff) request `TensorMemory::DmaBuf`
2424 // explicitly.
2425 MemTensor::<T>::new(shape, name).map(TensorStorage::Mem)
2426 }
2427 #[cfg(all(
2428 unix,
2429 not(any(
2430 target_os = "linux",
2431 target_os = "macos",
2432 target_os = "ios",
2433 target_os = "android"
2434 ))
2435 ))]
2436 {
2437 // Other Unix (BSD): Mem only (no DMA; Shm is explicit-only)
2438 MemTensor::<T>::new(shape, name).map(TensorStorage::Mem)
2439 }
2440 #[cfg(not(unix))]
2441 {
2442 // Windows/other: Mem only
2443 MemTensor::<T>::new(shape, name).map(TensorStorage::Mem)
2444 }
2445 }
2446 }
2447 }
2448 }
2449
2450 /// Create a DMA-backed tensor storage with an explicit byte size that
2451 /// may exceed `shape.product() * sizeof(T)`. Used for image tensors
2452 /// with row-padded layouts (see `DmaTensor::new_with_byte_size`).
2453 ///
2454 /// This is intentionally DMA-only: padding is only meaningful for
2455 /// buffers that will be imported as GPU textures via EGLImage. PBO,
2456 /// Shm, and Mem storage doesn't benefit from pitch alignment and
2457 /// shouldn't pay the memory cost.
2458 #[cfg(target_os = "linux")]
2459 pub(crate) fn new_dma_with_byte_size(
2460 shape: &[usize],
2461 byte_size: usize,
2462 name: Option<&str>,
2463 ) -> Result<Self> {
2464 DmaTensor::<T>::new_with_byte_size(shape, byte_size, name).map(TensorStorage::Dma)
2465 }
2466
2467 // No non-Linux stub: the only caller (`Tensor::image_with_stride`)
2468 // returns `NotImplemented` directly on non-Linux without ever
2469 // reaching the storage layer, so defining a stub here would be
2470 // dead code and fail the `-D warnings` clippy gate on macOS CI.
2471
2472 /// Create a Mem-backed tensor storage with an explicit byte size that may
2473 /// exceed `shape.product() * sizeof(T)`. Used for image tensors with
2474 /// 64-byte-aligned row strides (see `MemTensor::with_capacity_bytes`).
2475 pub(crate) fn new_mem_with_byte_size(
2476 shape: &[usize],
2477 byte_size: usize,
2478 name: Option<&str>,
2479 ) -> Result<Self>
2480 where
2481 T: 'static,
2482 {
2483 MemTensor::<T>::with_capacity_bytes(shape, byte_size, name).map(TensorStorage::Mem)
2484 }
2485
2486 /// Create a Shm-backed tensor storage with an explicit byte size that may
2487 /// exceed `shape.product() * sizeof(T)`. Used for image tensors with
2488 /// 64-byte-aligned row strides (see `ShmTensor::new_with_byte_size`).
2489 #[cfg(unix)]
2490 pub(crate) fn new_shm_with_byte_size(
2491 shape: &[usize],
2492 byte_size: usize,
2493 name: Option<&str>,
2494 ) -> Result<Self> {
2495 ShmTensor::<T>::new_with_byte_size(shape, byte_size, name).map(TensorStorage::Shm)
2496 }
2497
2498 /// Allocate an image-formatted IOSurface-backed storage (macOS).
2499 ///
2500 /// Used by `Tensor::image()` when the caller requests
2501 /// `TensorMemory::DmaBuf` and the format has an IOSurface FourCC
2502 /// mapping (YUYV, RGBA, BGRA today). Falls back to `new_with_byte_size`
2503 /// otherwise.
2504 #[cfg(any(target_os = "macos", target_os = "ios"))]
2505 pub(crate) fn new_image_iosurface(
2506 width: usize,
2507 height: usize,
2508 format: PixelFormat,
2509 dtype: DType,
2510 shape: &[usize],
2511 name: Option<&str>,
2512 ) -> Result<Self> {
2513 IoSurfaceTensor::<T>::new_image(width, height, format, dtype, shape, name)
2514 .map(TensorStorage::Dma)
2515 }
2516
2517 /// Allocate an image-formatted AHardwareBuffer-backed storage (Android).
2518 ///
2519 /// Used by `Tensor::image()` when the caller requests
2520 /// `TensorMemory::DmaBuf` and the format has an AHardwareBuffer format
2521 /// mapping (RGBA8 and the RGBA16F float paths today). Falls back to
2522 /// `new_with_byte_size` otherwise.
2523 #[cfg(target_os = "android")]
2524 pub(crate) fn new_image_ahardwarebuffer(
2525 width: usize,
2526 height: usize,
2527 format: PixelFormat,
2528 dtype: DType,
2529 shape: &[usize],
2530 name: Option<&str>,
2531 access: CpuAccess,
2532 ) -> Result<Self> {
2533 AHardwareBufferTensor::<T>::new_image(width, height, format, dtype, shape, name, access)
2534 .map(TensorStorage::Dma)
2535 }
2536
2537 /// Allocate an image-formatted D3D11 texture storage (Windows).
2538 ///
2539 /// Used by `Tensor::image()` when the caller requests
2540 /// `TensorMemory::DmaBuf`. Unlike the IOSurface and AHardwareBuffer
2541 /// constructors there is no byte-bag fallback: a format with no texture
2542 /// layout is an error, because a texture is the only thing this backing
2543 /// can be.
2544 #[cfg(target_os = "windows")]
2545 pub(crate) fn new_image_d3d11(
2546 width: usize,
2547 height: usize,
2548 format: PixelFormat,
2549 dtype: DType,
2550 shape: &[usize],
2551 name: Option<&str>,
2552 access: CpuAccess,
2553 ) -> Result<Self> {
2554 crate::d3d11::texture::D3d11TextureTensor::<T>::new_image(
2555 width, height, format, dtype, shape, name, access,
2556 )
2557 .map(TensorStorage::Dma)
2558 }
2559
2560 /// Create a new tensor storage using the given file descriptor, shape,
2561 /// and optional name.
2562 #[cfg(unix)]
2563 fn from_fd(fd: OwnedFd, shape: &[usize], name: Option<&str>) -> Result<Self> {
2564 #[cfg(target_os = "linux")]
2565 {
2566 use nix::sys::stat::fstat;
2567 use nix::sys::statfs::fstatfs;
2568
2569 let stat = fstat(&fd)?;
2570 let major = major(stat.st_dev);
2571 let minor = minor(stat.st_dev);
2572
2573 if major != 0 {
2574 // Dma and Shm tensors are expected to have major number 0
2575 return Err(Error::UnknownDeviceType(major, minor));
2576 }
2577
2578 // Classify by filesystem magic, never by the st_dev minor.
2579 //
2580 // dma_buf files live on an internal kernel mount (`dma_buf_mnt`,
2581 // created with `kern_mount`) whose superblock draws its device
2582 // number from `get_anon_bdev()` — an IDA allocated first-come,
2583 // first-served during boot, shared with every other anonymous
2584 // pseudo-filesystem (pipefs, sockfs, anon_inodefs, nsfs, bdev,
2585 // tracefs, ...). The minor a DMA-BUF lands on is therefore a
2586 // function of which pseudo-filesystems registered first on that
2587 // kernel build: it moves with kernel config, driver load order,
2588 // initramfs use, and kernel version. It is not part of any ABI.
2589 //
2590 // Observed values for a genuine DMA-BUF: 12 on x86 desktop, 8 on
2591 // the ADIS Verdin. An earlier revision hardcoded `9 | 10` here,
2592 // which silently imported real DMA-BUFs as shared memory
2593 // everywhere else.
2594 //
2595 // The filesystem magic, by contrast, is stable UAPI
2596 // (include/uapi/linux/magic.h). Both branches are identified
2597 // positively; anything else is genuinely unknown and is rejected
2598 // rather than guessed at. Falling back to SHM would "work" — a
2599 // DMA-BUF is mmap-able, and even a pipe imports as a zero-length
2600 // tensor — which is exactly why it must not be the default.
2601 let magic = fs_magic(fstatfs(&fd)?.filesystem_type().0 as i64);
2602
2603 log::debug!(
2604 "Creating tensor from fd: major={major}, minor={minor}, magic={magic:#010x}"
2605 );
2606
2607 match magic {
2608 DMA_BUF_MAGIC => DmaTensor::<T>::from_fd(fd, shape, name).map(TensorStorage::Dma),
2609 TMPFS_MAGIC => ShmTensor::<T>::from_fd(fd, shape, name).map(TensorStorage::Shm),
2610 other => Err(Error::UnknownBufferType(other)),
2611 }
2612 }
2613 #[cfg(all(unix, not(target_os = "linux")))]
2614 {
2615 // On macOS/iOS/BSD, always use SHM (no DMA-BUF fd import)
2616 ShmTensor::<T>::from_fd(fd, shape, name).map(TensorStorage::Shm)
2617 }
2618 }
2619}
2620
2621#[cfg(feature = "static")]
2622impl<T> TensorTrait<T> for TensorStorage<T>
2623where
2624 T: Num + Clone + fmt::Debug + Send + Sync,
2625{
2626 fn new(shape: &[usize], name: Option<&str>) -> Result<Self> {
2627 Self::new(shape, None, name)
2628 }
2629
2630 #[cfg(unix)]
2631 fn from_fd(fd: OwnedFd, shape: &[usize], name: Option<&str>) -> Result<Self> {
2632 Self::from_fd(fd, shape, name)
2633 }
2634
2635 #[cfg(unix)]
2636 fn clone_fd(&self) -> Result<OwnedFd> {
2637 match self {
2638 #[cfg(any(
2639 target_os = "linux",
2640 target_os = "macos",
2641 target_os = "ios",
2642 target_os = "android"
2643 ))]
2644 TensorStorage::Dma(t) => t.clone_fd(),
2645 TensorStorage::Shm(t) => t.clone_fd(),
2646 TensorStorage::Mem(t) => t.clone_fd(),
2647 TensorStorage::Pbo(t) => t.clone_fd(),
2648 }
2649 }
2650
2651 fn memory(&self) -> TensorMemory {
2652 match self {
2653 #[cfg(any(
2654 target_os = "linux",
2655 target_os = "macos",
2656 target_os = "ios",
2657 target_os = "android",
2658 target_os = "windows"
2659 ))]
2660 TensorStorage::Dma(_) => TensorMemory::DmaBuf,
2661 #[cfg(unix)]
2662 TensorStorage::Shm(_) => TensorMemory::Shm,
2663 TensorStorage::Mem(_) => TensorMemory::Mem,
2664 TensorStorage::Pbo(_) => TensorMemory::Pbo,
2665 }
2666 }
2667
2668 fn name(&self) -> String {
2669 match self {
2670 #[cfg(any(
2671 target_os = "linux",
2672 target_os = "macos",
2673 target_os = "ios",
2674 target_os = "android",
2675 target_os = "windows"
2676 ))]
2677 TensorStorage::Dma(t) => t.name(),
2678 #[cfg(unix)]
2679 TensorStorage::Shm(t) => t.name(),
2680 TensorStorage::Mem(t) => t.name(),
2681 TensorStorage::Pbo(t) => t.name(),
2682 }
2683 }
2684
2685 fn shape(&self) -> &[usize] {
2686 match self {
2687 #[cfg(any(
2688 target_os = "linux",
2689 target_os = "macos",
2690 target_os = "ios",
2691 target_os = "android",
2692 target_os = "windows"
2693 ))]
2694 TensorStorage::Dma(t) => t.shape(),
2695 #[cfg(unix)]
2696 TensorStorage::Shm(t) => t.shape(),
2697 TensorStorage::Mem(t) => t.shape(),
2698 TensorStorage::Pbo(t) => t.shape(),
2699 }
2700 }
2701
2702 fn reshape(&mut self, shape: &[usize]) -> Result<()> {
2703 match self {
2704 #[cfg(any(
2705 target_os = "linux",
2706 target_os = "macos",
2707 target_os = "ios",
2708 target_os = "android",
2709 target_os = "windows"
2710 ))]
2711 TensorStorage::Dma(t) => t.reshape(shape),
2712 #[cfg(unix)]
2713 TensorStorage::Shm(t) => t.reshape(shape),
2714 TensorStorage::Mem(t) => t.reshape(shape),
2715 TensorStorage::Pbo(t) => t.reshape(shape),
2716 }
2717 }
2718
2719 fn capacity_bytes(&self) -> usize {
2720 match self {
2721 #[cfg(any(
2722 target_os = "linux",
2723 target_os = "macos",
2724 target_os = "ios",
2725 target_os = "android",
2726 target_os = "windows"
2727 ))]
2728 TensorStorage::Dma(t) => t.capacity_bytes(),
2729 #[cfg(unix)]
2730 TensorStorage::Shm(t) => t.capacity_bytes(),
2731 TensorStorage::Mem(t) => t.capacity_bytes(),
2732 TensorStorage::Pbo(t) => t.capacity_bytes(),
2733 }
2734 }
2735
2736 fn set_logical_shape(&mut self, shape: &[usize]) -> Result<()> {
2737 match self {
2738 #[cfg(any(
2739 target_os = "linux",
2740 target_os = "macos",
2741 target_os = "ios",
2742 target_os = "android",
2743 target_os = "windows"
2744 ))]
2745 TensorStorage::Dma(t) => t.set_logical_shape(shape),
2746 #[cfg(unix)]
2747 TensorStorage::Shm(t) => t.set_logical_shape(shape),
2748 TensorStorage::Mem(t) => t.set_logical_shape(shape),
2749 TensorStorage::Pbo(t) => t.set_logical_shape(shape),
2750 }
2751 }
2752
2753 fn map_with<'a>(&self, access: CpuAccess) -> Result<crate::view::HostView<'a, T>>
2754 where
2755 T: 'a,
2756 {
2757 match self {
2758 #[cfg(any(
2759 target_os = "linux",
2760 target_os = "macos",
2761 target_os = "ios",
2762 target_os = "android",
2763 target_os = "windows"
2764 ))]
2765 TensorStorage::Dma(t) => t.map_with(access),
2766 #[cfg(unix)]
2767 TensorStorage::Shm(t) => t.map_with(access),
2768 TensorStorage::Mem(t) => t.map_with(access),
2769 TensorStorage::Pbo(t) => t.map_with(access),
2770 }
2771 }
2772
2773 fn try_map_with<'a>(&self, access: CpuAccess) -> Result<crate::view::HostView<'a, T>>
2774 where
2775 T: 'a,
2776 {
2777 match self {
2778 #[cfg(any(
2779 target_os = "linux",
2780 target_os = "macos",
2781 target_os = "ios",
2782 target_os = "android",
2783 target_os = "windows"
2784 ))]
2785 TensorStorage::Dma(t) => t.try_map_with(access),
2786 #[cfg(unix)]
2787 TensorStorage::Shm(t) => t.try_map_with(access),
2788 TensorStorage::Mem(t) => t.try_map_with(access),
2789 TensorStorage::Pbo(t) => t.try_map_with(access),
2790 }
2791 }
2792
2793 fn buffer_identity(&self) -> &BufferIdentity {
2794 match self {
2795 #[cfg(any(
2796 target_os = "linux",
2797 target_os = "macos",
2798 target_os = "ios",
2799 target_os = "android",
2800 target_os = "windows"
2801 ))]
2802 TensorStorage::Dma(t) => t.buffer_identity(),
2803 #[cfg(unix)]
2804 TensorStorage::Shm(t) => t.buffer_identity(),
2805 TensorStorage::Mem(t) => t.buffer_identity(),
2806 TensorStorage::Pbo(t) => t.buffer_identity(),
2807 }
2808 }
2809
2810 /// Forward a sub-region view to the active backend, re-wrapping the
2811 /// backend's view (which shares the parent's allocation and
2812 /// [`BufferIdentity`]) back into the matching `TensorStorage` variant. Each
2813 /// backend's `view` is its own `TensorTrait::view` override; this single
2814 /// match is the only per-variant dispatch (`Tensor::subview` calls through
2815 /// here rather than matching the storage itself).
2816 fn view(&self, offset_bytes: usize, shape: &[usize]) -> Result<Self> {
2817 match self {
2818 #[cfg(any(
2819 target_os = "linux",
2820 target_os = "macos",
2821 target_os = "ios",
2822 target_os = "android",
2823 target_os = "windows"
2824 ))]
2825 TensorStorage::Dma(t) => t.view(offset_bytes, shape).map(TensorStorage::Dma),
2826 #[cfg(unix)]
2827 TensorStorage::Shm(t) => t.view(offset_bytes, shape).map(TensorStorage::Shm),
2828 TensorStorage::Mem(t) => t.view(offset_bytes, shape).map(TensorStorage::Mem),
2829 TensorStorage::Pbo(t) => t.view(offset_bytes, shape).map(TensorStorage::Pbo),
2830 }
2831 }
2832}
2833
2834/// Multi-backend tensor with optional image format metadata.
2835///
2836/// When `format` is `Some`, this tensor represents an image. Width, height,
2837/// and channels are derived from `shape` + `format`. When `format` is `None`,
2838/// this is a raw tensor (identical to the pre-refactoring behavior).
2839///
2840/// This is the `static` backend's `Tensor<T>` -- it owns its storage. See
2841/// `tensor_dyn/dynamic_tensor.rs` for `dynamic`'s `Tensor<T>`, a typed lens
2842/// with no storage of its own.
2843#[cfg(feature = "static")]
2844#[derive(Debug)]
2845pub struct Tensor<T>
2846where
2847 T: Num + Clone + fmt::Debug + Send + Sync,
2848{
2849 /// CUDA registration for this tensor, if any. Set after creation by
2850 /// the image crate once a PBO is registered with CUDA interop.
2851 ///
2852 /// MUST be declared before `storage`: CUDA must unregister the GL buffer
2853 /// before storage's Drop deletes it (cudaGraphicsUnregisterResource before
2854 /// glDeleteBuffers). Rust drops fields in declaration order.
2855 cuda: Option<crate::cuda::CudaHandle>,
2856 pub(crate) storage: TensorStorage<T>,
2857 format: Option<PixelFormat>,
2858 chroma: Option<Box<Tensor<T>>>,
2859 /// Row stride in bytes for externally allocated buffers with row padding.
2860 /// `None` means tightly packed (stride == width * bytes_per_pixel).
2861 row_stride: Option<usize>,
2862 /// Byte offset within the DMA-BUF where image data starts.
2863 /// `None` means offset 0 (data starts at the beginning of the buffer).
2864 plane_offset: Option<usize>,
2865 /// Quantization metadata for integer-typed tensors. Public access is
2866 /// gated by the `IntegerType` trait — `Tensor<f32>` etc. carry the
2867 /// field for layout uniformity but have no way to read or write it.
2868 pub(crate) quantization: Option<Quantization>,
2869 /// Optional colorimetry metadata. `None` = undefined; never auto-filled.
2870 colorimetry: Option<crate::Colorimetry>,
2871 /// Declared CPU access (see [`CpuAccess`]). Image constructors set it
2872 /// from their `access` parameter; non-image tensors (`new`, imports,
2873 /// numpy) default to `ReadWrite` — they are CPU-centric by nature.
2874 /// `map_with` counts requests beyond this declaration as unplanned.
2875 cpu_access: CpuAccess,
2876 /// Recorded vendor tile-compression scheme — best knowledge from
2877 /// allocation time (see [`Compression`]). `Some` only for Android
2878 /// hardware-only AHardwareBuffers whose allocation requested
2879 /// compression on an eligible format with a recognized vendor. When
2880 /// set, the row-stride accessors describe no meaningful linear
2881 /// layout. `configure_image` preserves it (physical layout, unlike
2882 /// colorimetry); views inherit it.
2883 compression: Option<CompressionScheme>,
2884 /// Parent-image snapshot when this tensor is a [`view`](Self::view)/
2885 /// [`batch`](Self::batch) sub-region; `None` for a whole tensor. Lets the GL
2886 /// backend key its import on the parent and render the view as a
2887 /// `glViewport`/`glScissor` ROI. See [`ViewOrigin`].
2888 view_origin: Option<ViewOrigin>,
2889}
2890
2891#[cfg(feature = "static")]
2892impl<T> Tensor<T>
2893where
2894 T: Num + Clone + fmt::Debug + Send + Sync,
2895{
2896 /// Wrap a TensorStorage in a Tensor with no image metadata.
2897 pub(crate) fn wrap(storage: TensorStorage<T>) -> Self {
2898 Self {
2899 storage,
2900 format: None,
2901 chroma: None,
2902 row_stride: None,
2903 plane_offset: None,
2904 quantization: None,
2905 cuda: None,
2906 colorimetry: None,
2907 cpu_access: CpuAccess::ReadWrite,
2908 compression: None,
2909 view_origin: None,
2910 }
2911 }
2912
2913 /// Construct a tensor from a row-major element slice + shape. Allocates a
2914 /// new buffer (`TensorMemory::Mem`) and memcpys the contents; caller
2915 /// retains ownership of the input slice.
2916 ///
2917 /// # Errors
2918 ///
2919 /// - [`Error::InvalidShape`] if `values.len() != shape.iter().product()`.
2920 /// - Propagates any allocation error from [`Self::new`].
2921 pub fn from_slice(values: &[T], shape: &[usize]) -> Result<Self>
2922 where
2923 T: Copy,
2924 {
2925 let expected: usize = shape.iter().product();
2926 if values.len() != expected {
2927 return Err(Error::InvalidShape(format!(
2928 "from_slice: values.len()={} but shape product={expected} (shape={shape:?})",
2929 values.len()
2930 )));
2931 }
2932 let t = Self::new(shape, Some(TensorMemory::Mem), None)?;
2933 {
2934 let mut m = t.map()?;
2935 m.as_mut_slice().copy_from_slice(values);
2936 }
2937 Ok(t)
2938 }
2939
2940 /// Wrap externally-owned memory as a tensor without copying. The tensor
2941 /// borrows `[ptr, ptr + shape.product() * size_of::<T>())` as
2942 /// [`TensorMemory::Mem`]; `owner`, when `Some`, co-owns the source so it
2943 /// outlives the tensor (and all derived views/maps). See [`ForeignOwner`].
2944 ///
2945 /// The canonical use is CUDA zero-copy: allocate host-coherent memory
2946 /// (`cudaHostAlloc`), wrap the host pointer here, and bind the matching
2947 /// device pointer to the inference engine — reads and writes hit the same
2948 /// physical buffer with no host copy. The identical primitive backs the
2949 /// Python `Tensor.from_numpy` zero-copy borrow (owner = the NumPy object).
2950 ///
2951 /// # Safety
2952 ///
2953 /// `ptr` must be non-null, aligned to `align_of::<T>()`, and valid for
2954 /// `shape.product()` elements of `T` for as long as the returned tensor —
2955 /// and every view/map sharing its backing — is alive. Pass an `owner` that
2956 /// co-owns the source to uphold that contract.
2957 ///
2958 /// # Errors
2959 ///
2960 /// [`Error::InvalidSize`] if `shape` is empty.
2961 pub unsafe fn from_foreign(
2962 ptr: *mut T,
2963 shape: &[usize],
2964 owner: Option<crate::ForeignOwner>,
2965 name: Option<&str>,
2966 ) -> Result<Self> {
2967 // SAFETY: forwarded to `from_foreign_with_capacity`'s contract; `0`
2968 // clamps to the tight shape-derived footprint below, i.e. no extra
2969 // headroom -- identical to this function's own prior behaviour.
2970 unsafe { Self::from_foreign_with_capacity(ptr, shape, 0, owner, name) }
2971 }
2972
2973 /// [`Self::from_foreign`], additionally recording that the allocation
2974 /// behind `ptr` holds at least `capacity_bytes` bytes -- headroom beyond
2975 /// the tight `shape.product() * size_of::<T>()` footprint that a later
2976 /// [`configure_image`](Self::configure_image) / `set_logical_shape` can
2977 /// grow into without erroring.
2978 ///
2979 /// This is the consumer half of the cross-package capsule protocol's
2980 /// [`crate::protocol::TensorDesc::capacity`]: a producer's pool tensor
2981 /// (allocated for its largest expected image, holding a smaller one
2982 /// today) reports its real allocation size so the imported alias
2983 /// inherits the same headroom instead of being clamped to today's
2984 /// logical shape. `capacity_bytes` below the tight footprint is clamped
2985 /// up to it, never down -- the shape-implied floor always holds.
2986 ///
2987 /// # Safety
2988 ///
2989 /// Same as [`Self::from_foreign`], except `ptr` must be valid for
2990 /// `capacity_bytes` bytes (or the tight shape footprint, whichever is
2991 /// larger), not just the shape-implied footprint.
2992 ///
2993 /// # Errors
2994 ///
2995 /// [`Error::InvalidSize`] if `shape` is empty.
2996 pub unsafe fn from_foreign_with_capacity(
2997 ptr: *mut T,
2998 shape: &[usize],
2999 capacity_bytes: usize,
3000 owner: Option<crate::ForeignOwner>,
3001 name: Option<&str>,
3002 ) -> Result<Self> {
3003 if shape.is_empty() {
3004 return Err(Error::InvalidSize(0));
3005 }
3006 if ptr.is_null() {
3007 return Err(Error::InvalidArgument(
3008 "from_foreign: ptr must be non-null".to_owned(),
3009 ));
3010 }
3011 shape
3012 .iter()
3013 .copied()
3014 .try_fold(1usize, |acc, dim| acc.checked_mul(dim))
3015 .ok_or_else(|| {
3016 Error::InvalidArgument(format!(
3017 "from_foreign: shape.product() overflows usize (shape={shape:?})"
3018 ))
3019 })?;
3020 let capacity_elems = capacity_bytes / std::mem::size_of::<T>();
3021 // SAFETY: this function's own contract (see its doc comment) is
3022 // exactly `MemTensor::from_foreign_with_capacity`'s contract.
3023 let mem = unsafe {
3024 MemTensor::<T>::from_foreign_with_capacity(ptr, shape, capacity_elems, owner, name)
3025 };
3026 Ok(Self::wrap(TensorStorage::Mem(mem)))
3027 }
3028
3029 /// Pin a stable host address for this tensor's data.
3030 ///
3031 /// The returned [`HostPin`] carries **no borrow of the tensor**, so a
3032 /// caller can pin once at init and still pass `&mut self` to
3033 /// `ImageProcessor::convert` every frame. That borrow conflict is the
3034 /// blocker reported in [#134](https://github.com/EdgeFirstAI/hal/issues/134).
3035 ///
3036 /// The address stays valid until the pin is dropped, and survives
3037 /// re-layout: `configure_image` and `set_logical_shape` only rewrite the
3038 /// logical shape after a capacity check, and every backend's pin owns its
3039 /// own mapping, so nothing can move underneath it.
3040 ///
3041 /// A pin is therefore a **snapshot**, not a tracker: [`HostPin::len`]
3042 /// reports the extent at pin time and does not follow a later re-layout.
3043 /// That is the safe direction — the window always stays inside the
3044 /// allocation — but it means a caller who re-layouts should re-pin if it
3045 /// wants the new extent.
3046 ///
3047 /// The address is **not** a coherency guarantee: on non-coherent backends
3048 /// bracket CPU access with
3049 /// [`sync_for_cpu`](Self::sync_for_cpu) and
3050 /// [`sync_for_device`](Self::sync_for_device), which are no-ops where the
3051 /// backend does not require them.
3052 ///
3053 /// # Cost
3054 ///
3055 /// On `Mem` this is an `Arc` refcount bump — the backing allocation is
3056 /// never resized, so its base pointer is already stable. Other backends
3057 /// establish a mapping once and reuse it.
3058 ///
3059 /// # Errors
3060 ///
3061 /// - [`Error::InvalidArgument`] when `access` is [`CpuAccess::None`]: a
3062 /// pinned mapping is CPU access by definition, so this is an error
3063 /// rather than a warn-and-stage.
3064 /// - [`Error::NotImplemented`] for the two backends that **cannot** pin,
3065 /// rather than merely not having been wired yet: a PBO
3066 /// (`glMapBufferRange` establishes the address and the visibility
3067 /// together, and there is no `GL_MAP_PERSISTENT_BIT` in GLES) and an
3068 /// AHardwareBuffer (`AHardwareBuffer_lock` is what yields the address at
3069 /// all, so there is no address outside the lock). Both name
3070 /// `map()`/`map_with()` as the alternative. `Mem`, `Shm`, DMA-BUF,
3071 /// IOSurface and the Windows D3D11 texture all pin — the texture through
3072 /// a persistent host buffer the storage owns, which
3073 /// [`sync_for_cpu`](Self::sync_for_cpu) and
3074 /// [`sync_for_device`](Self::sync_for_device) move bytes between and the
3075 /// texture, rather than through the staging mapping (a pin outlives every
3076 /// map guard, so holding the staging texture mapped would freeze it).
3077 pub fn pin_host<'a>(&self, access: CpuAccess) -> Result<HostPin<'a>>
3078 where
3079 T: 'a,
3080 {
3081 if access == CpuAccess::None {
3082 return Err(Error::InvalidArgument(
3083 "pin_host: CpuAccess::None is not a valid pin; a pinned mapping \
3084 is CPU access by definition"
3085 .to_owned(),
3086 ));
3087 }
3088 // Backends yield everything addressable from the tensor's offset;
3089 // narrow to the logical extent so a consumer cannot mistake pitch
3090 // padding for data. A map guard keeps the wider window.
3091 let logical = self.len() * std::mem::size_of::<T>();
3092 let pin = match &self.storage {
3093 TensorStorage::Mem(m) => Ok(m.host_pin()),
3094 #[cfg(unix)]
3095 TensorStorage::Shm(s) => s.host_pin(),
3096 #[cfg(target_os = "linux")]
3097 TensorStorage::Dma(d) => d.host_pin(),
3098 #[cfg(any(target_os = "macos", target_os = "ios"))]
3099 TensorStorage::Dma(s) => s.host_pin(),
3100 // Windows D3D11: the pin is the storage's own persistent host
3101 // buffer, not the staging mapping -- see `D3d11TextureTensor::host_pin`.
3102 // The direction goes with it: a writable pin is what tells
3103 // `sync_for_device` the host buffer, not the staging copy, carries
3104 // the caller's writes.
3105 #[cfg(target_os = "windows")]
3106 TensorStorage::Dma(t) => t.host_pin(access),
3107 // AHardwareBuffer cannot support a pin at all, and that is a
3108 // platform contract rather than missing work: the NDK exposes no
3109 // way to obtain a CPU address outside AHardwareBuffer_lock, so a
3110 // "persistent" address would mean holding the lock for the pin's
3111 // lifetime -- which blocks the GPU from the buffer and makes any
3112 // further lock undefined. Both are the opposite of what a pin is
3113 // for. Say so, rather than implying a TODO.
3114 #[cfg(target_os = "android")]
3115 TensorStorage::Dma(_) => Err(Error::NotImplemented(
3116 "pin_host: AHardwareBuffer has no host address outside \
3117 AHardwareBuffer_lock, so it cannot hand out one that outlives \
3118 a guard. Use map()/map_with() for CPU access; for zero-copy \
3119 handoff pass the buffer itself."
3120 .to_owned(),
3121 )),
3122 // PBO is the same shape of constraint as AHardwareBuffer. The
3123 // message is a shared constant because `TensorDyn::pin_host`
3124 // (`tensor_dyn/dynamic_backend.rs`) must repeat this refusal --
3125 // the ABI has `ef_tensor_map` and no `ef_tensor_pin_host`, so the
3126 // library cannot tell a pin from a map. See
3127 // `crate::pbo::PIN_HOST_PBO_REFUSAL` for the reasoning at length.
3128 TensorStorage::Pbo(_) => Err(Error::NotImplemented(
3129 crate::pbo::PIN_HOST_PBO_REFUSAL.to_owned(),
3130 )),
3131 }?;
3132 Ok(pin.narrowed(logical))
3133 }
3134
3135 /// Make CPU reads/writes of this tensor coherent — the CPU is about to
3136 /// access the buffer.
3137 ///
3138 /// Corresponds to `DMA_BUF_SYNC_START`, and the name matches the Linux
3139 /// kernel DMA API for the same reason: `for_cpu` says who is about to
3140 /// touch the memory, which is what the caller actually knows.
3141 ///
3142 /// Pair with [`sync_for_device`](Self::sync_for_device).
3143 ///
3144 /// # Which backends owe maintenance
3145 ///
3146 /// | backend | behaviour |
3147 /// |---|---|
3148 /// | `Mem`, `Shm` | **no-op** — plain CPU memory, coherent by construction |
3149 /// | `Dma` (Linux DMA-BUF) | one `DMA_BUF_IOCTL_SYNC`, direction-tagged |
3150 /// | `Dma` (macOS/iOS IOSurface) | `IOSurfaceLock` / `IOSurfaceUnlock`, options matched to `access` |
3151 /// | `Dma` (Android AHardwareBuffer) | [`Error::NotImplemented`] — no sync independent of the lock |
3152 /// | `Dma` (Windows D3D11 texture) | refreshes the pinned host buffer from the texture through the staging copy, or uploads it back; a no-op until something pins |
3153 /// | `Pbo` | [`Error::NotImplemented`] — no coherency window independent of the map |
3154 ///
3155 /// Only the first row is a no-op, so "bracket unconditionally and pay
3156 /// nothing" holds for `Mem`/`Shm` alone. The two that report
3157 /// `NotImplemented` do so because their platform fuses addressing and
3158 /// coherency into one call — `AHardwareBuffer_lock` and
3159 /// `glMapBufferRange` each perform the maintenance *as* they hand out the
3160 /// address, so a separate bracket would have to nest inside a map guard
3161 /// that already owns the pair. Both name `map()`/`map_with()` as the
3162 /// alternative. They previously fell through a wildcard and answered
3163 /// `Ok(())`, which is worse than an error because the symptom is
3164 /// stale-but-plausible pixels rather than a failure. The match is
3165 /// exhaustive so a new backend cannot inherit that silence — it must state
3166 /// its coherency story to compile.
3167 ///
3168 /// # Cost
3169 ///
3170 /// `Mem`/`Shm`: nothing. `Dma` on Linux: one ioctl. IOSurface: one
3171 /// lock/unlock pair. Windows D3D11 texture: one `CopyResource` plus a
3172 /// staging map per direction, and nothing at all on a tensor no one has
3173 /// pinned.
3174 ///
3175 /// # Errors
3176 ///
3177 /// [`Error::InvalidArgument`] when `access` is [`CpuAccess::None`], which
3178 /// describes no access to make coherent.
3179 pub fn sync_for_cpu(&self, access: CpuAccess) -> Result<()> {
3180 Self::check_sync_access(access)?;
3181 match &self.storage {
3182 // Plain CPU memory: coherent by construction.
3183 TensorStorage::Mem(_) => Ok(()),
3184 #[cfg(unix)]
3185 TensorStorage::Shm(_) => Ok(()),
3186 #[cfg(target_os = "linux")]
3187 TensorStorage::Dma(d) => {
3188 // The direction tells the kernel which maintenance this access
3189 // needs: a read-only bracket skips the writeback at END, a
3190 // write-only one skips the invalidate at START. It MUST match
3191 // the direction passed to sync_for_device.
3192 crate::dmabuf::sync_access(&d.fd, true, access).map_err(Error::NixError)
3193 }
3194 #[cfg(any(target_os = "macos", target_os = "ios"))]
3195 TensorStorage::Dma(s) => s.sync_for_cpu(access),
3196 // Windows D3D11: refreshes the pinned host buffer from the texture
3197 // through the staging copy.
3198 #[cfg(target_os = "windows")]
3199 TensorStorage::Dma(t) => t.sync_for_cpu(access),
3200 #[cfg(target_os = "android")]
3201 TensorStorage::Dma(_) => Err(Error::NotImplemented(
3202 "sync_for_cpu: AHardwareBuffer has no sync independent of its \
3203 lock. AHardwareBuffer_lock both establishes the address and \
3204 performs the cache invalidate, and the NDK leaves locking an \
3205 already-locked buffer undefined, so a bracket on top of a map \
3206 guard would nest. Use map()/map_with(), which owns the pair."
3207 .to_owned(),
3208 )),
3209 TensorStorage::Pbo(_) => Err(Error::NotImplemented(
3210 "sync_for_cpu: a PBO has no coherency window independent of its \
3211 map -- glMapBufferRange establishes the address and the \
3212 visibility together, and glUnmapBuffer publishes the writes. \
3213 Use map()/map_with(), which owns the pair."
3214 .to_owned(),
3215 )),
3216 }
3217 }
3218
3219 /// Release the buffer back to the device — the CPU is done accessing it.
3220 ///
3221 /// Corresponds to `DMA_BUF_SYNC_END`. See
3222 /// [`sync_for_cpu`](Self::sync_for_cpu) for the pairing and cost.
3223 pub fn sync_for_device(&self, access: CpuAccess) -> Result<()> {
3224 Self::check_sync_access(access)?;
3225 match &self.storage {
3226 TensorStorage::Mem(_) => Ok(()),
3227 #[cfg(unix)]
3228 TensorStorage::Shm(_) => Ok(()),
3229 #[cfg(target_os = "linux")]
3230 TensorStorage::Dma(d) => {
3231 crate::dmabuf::sync_access(&d.fd, false, access).map_err(Error::NixError)
3232 }
3233 #[cfg(any(target_os = "macos", target_os = "ios"))]
3234 TensorStorage::Dma(s) => s.sync_for_device(access),
3235 // Windows D3D11: uploads the pinned host buffer to the texture.
3236 #[cfg(target_os = "windows")]
3237 TensorStorage::Dma(t) => t.sync_for_device(access),
3238 #[cfg(target_os = "android")]
3239 TensorStorage::Dma(_) => Err(Error::NotImplemented(
3240 "sync_for_device: AHardwareBuffer has no sync independent of \
3241 its lock — see sync_for_cpu. Use map()/map_with()."
3242 .to_owned(),
3243 )),
3244 TensorStorage::Pbo(_) => Err(Error::NotImplemented(
3245 "sync_for_device: a PBO has no coherency window independent of \
3246 its map -- see sync_for_cpu. Use map()/map_with()."
3247 .to_owned(),
3248 )),
3249 }
3250 }
3251
3252 fn check_sync_access(access: CpuAccess) -> Result<()> {
3253 if access == CpuAccess::None {
3254 return Err(Error::InvalidArgument(
3255 "sync_for_cpu/sync_for_device: CpuAccess::None describes no \
3256 access to synchronise"
3257 .to_owned(),
3258 ));
3259 }
3260 Ok(())
3261 }
3262
3263 /// Construct a tensor from a 3-D ndarray view. Respects strides — one
3264 /// copy in all cases; contiguous views take a memcpy fast path.
3265 ///
3266 /// Only available when the `ndarray` feature is enabled.
3267 #[cfg(feature = "ndarray")]
3268 pub fn from_arrayview3(view: ndarray::ArrayView3<'_, T>) -> Result<Self>
3269 where
3270 T: Copy,
3271 {
3272 let (h, w, c) = view.dim();
3273 let t = Self::new(&[h, w, c], Some(TensorMemory::Mem), None)?;
3274 {
3275 let mut m = t.map()?;
3276 let dst = m.as_mut_slice();
3277 if let Some(src) = view.as_slice() {
3278 dst.copy_from_slice(src);
3279 } else {
3280 for (d, &s) in dst.iter_mut().zip(view.iter()) {
3281 *d = s;
3282 }
3283 }
3284 }
3285 Ok(t)
3286 }
3287
3288 /// Create a new tensor with the given shape, memory type, and optional
3289 /// name. If no name is given, a random name will be generated. If no
3290 /// memory type is given, the best available memory type will be chosen
3291 /// based on the platform and environment variables.
3292 ///
3293 /// On Linux platforms, the order of preference is: Dma -> Shm -> Mem.
3294 /// On other Unix platforms (macOS), the order is: Shm -> Mem.
3295 /// On non-Unix platforms, only Mem is available.
3296 ///
3297 /// # Environment Variables
3298 /// - `EDGEFIRST_TENSOR_FORCE_MEM`: If set to a non-zero and non-false
3299 /// value, forces the use of regular system memory allocation
3300 /// (`TensorMemory::Mem`) regardless of platform capabilities.
3301 ///
3302 /// # Example
3303 /// ```rust
3304 /// use edgefirst_tensor::{Error, Tensor, TensorMemory, TensorTrait};
3305 /// # fn main() -> Result<(), Error> {
3306 /// let tensor = Tensor::<f32>::new(&[2, 3, 4], Some(TensorMemory::Mem), Some("test_tensor"))?;
3307 /// assert_eq!(tensor.memory(), TensorMemory::Mem);
3308 /// assert_eq!(tensor.name(), "test_tensor");
3309 /// # Ok(())
3310 /// # }
3311 /// ```
3312 pub fn new(shape: &[usize], memory: Option<TensorMemory>, name: Option<&str>) -> Result<Self> {
3313 let _span = tracing::trace_span!(
3314 "tensor.alloc",
3315 ?shape,
3316 memory = ?memory,
3317 dtype = std::any::type_name::<T>(),
3318 )
3319 .entered();
3320 #[cfg_attr(not(target_os = "linux"), allow(unused_mut))]
3321 let mut t = TensorStorage::new(shape, memory, name).map(Self::wrap)?;
3322 // Best-effort: attach a CUDA ExternalMemory handle for DMA tensors on
3323 // CUDA-capable hosts. Never blocks tensor creation on failure.
3324 // RUNTIME-UNVALIDATED: no CUDA+dma_heap test platform available; ABI
3325 // layout-asserted vs. CUDA 12.6 driver_types.h; mechanism proven by
3326 // gpu-probe O5 on Orin.
3327 #[cfg(target_os = "linux")]
3328 t.try_init_dma_cuda();
3329 Ok(t)
3330 }
3331
3332 /// Create an image tensor with the given format.
3333 /// Allocate an image tensor from a declarative request — the
3334 /// full-featured constructor behind [`Self::image`] and friends.
3335 ///
3336 /// Adds the [`Compression`] request to the classic parameters:
3337 ///
3338 /// - a request with any CPU access other than [`CpuAccess::None`] is
3339 /// [`Error::InvalidArgument`] (CPU mapping pins the layout linear);
3340 /// - [`Compression::Scheme`] fails with [`Error::NotImplemented`] on
3341 /// platforms without vendor tile compression, and with
3342 /// [`Error::InvalidArgument`] when the device's native scheme or
3343 /// the format eligibility doesn't match;
3344 /// - [`Compression::Any`] never fails for compression reasons: it
3345 /// records the scheme when the allocation is eligible and
3346 /// otherwise resolves linear, incrementing
3347 /// [`compression_fallback_count`].
3348 ///
3349 /// The recorded outcome is readable via [`Tensor::compression`].
3350 pub fn image_desc(desc: &ImageDesc) -> Result<Self>
3351 where
3352 T: 'static,
3353 {
3354 let Some(t_dtype) = dtype_of::<T>() else {
3355 return Err(Error::InvalidArgument(
3356 "image_desc: element type has no DType mapping".into(),
3357 ));
3358 };
3359 if t_dtype != desc.dtype {
3360 return Err(Error::InvalidArgument(format!(
3361 "image_desc: desc.dtype is {:?} but the tensor element type is {t_dtype:?}",
3362 desc.dtype
3363 )));
3364 }
3365
3366 // Compression-request guards. CPU access pins the layout linear,
3367 // so a request combined with any declared access is a
3368 // contradiction the caller should hear about immediately.
3369 if desc.compression.is_some() && desc.access != CpuAccess::None {
3370 return Err(Error::InvalidArgument(format!(
3371 "image_desc: a compression request requires CpuAccess::None (declared {:?}) — CPU mapping pins the layout linear",
3372 desc.access
3373 )));
3374 }
3375 if let Some(Compression::Scheme(requested)) = desc.compression {
3376 #[cfg(not(target_os = "android"))]
3377 {
3378 return Err(Error::NotImplemented(format!(
3379 "image_desc: Compression::Scheme({requested:?}) — no vendor tile compression on this platform (request Compression::Any for a portable fallback)"
3380 )));
3381 }
3382 #[cfg(target_os = "android")]
3383 {
3384 if matches!(
3385 desc.memory,
3386 Some(TensorMemory::Mem) | Some(TensorMemory::Shm) | Some(TensorMemory::Pbo)
3387 ) {
3388 return Err(Error::InvalidArgument(format!(
3389 "image_desc: Compression::Scheme({requested:?}) requires hardware memory (TensorMemory::DmaBuf or auto-select), got {:?}",
3390 desc.memory
3391 )));
3392 }
3393 if !crate::ahardwarebuffer_layout::compression_eligible(desc.format, desc.dtype) {
3394 return Err(Error::InvalidArgument(format!(
3395 "image_desc: ({:?}, {:?}) is not compression-eligible (RGBA8888 u8/i8 initially)",
3396 desc.format, desc.dtype
3397 )));
3398 }
3399 let device = crate::ahardwarebuffer::device_compression_scheme();
3400 if device != Some(requested) {
3401 return Err(Error::InvalidArgument(format!(
3402 "image_desc: Compression::Scheme({requested:?}) but the device's native scheme is {device:?}"
3403 )));
3404 }
3405 }
3406 }
3407
3408 // A compression request implies a hardware pipeline, so auto
3409 // memory promotes to the platform's zero-copy allocation first
3410 // (`Tensor::image` only takes the AHardwareBuffer/IOSurface path
3411 // under an explicit Dma request). `Scheme` propagates the Dma
3412 // failure — the caller demanded a layout only that allocator can
3413 // produce; `Any` falls back to plain auto-select.
3414 #[allow(unused_mut)]
3415 let mut t = match (desc.memory, desc.compression) {
3416 (None, Some(request)) => {
3417 match Self::image(
3418 desc.width,
3419 desc.height,
3420 desc.format,
3421 Some(TensorMemory::DmaBuf),
3422 desc.access,
3423 ) {
3424 Ok(t) => t,
3425 Err(e) if matches!(request, Compression::Scheme(_)) => return Err(e),
3426 Err(_) => Self::image(desc.width, desc.height, desc.format, None, desc.access)?,
3427 }
3428 }
3429 (memory, _) => Self::image(desc.width, desc.height, desc.format, memory, desc.access)?,
3430 };
3431
3432 // Record best knowledge / count fallbacks. Only an Android
3433 // hardware-only AHardwareBuffer allocation can actually hold a
3434 // vendor tile layout; everywhere else an Any request resolves
3435 // linear and is counted.
3436 if let Some(request) = desc.compression {
3437 #[cfg(target_os = "android")]
3438 {
3439 let eligible =
3440 crate::ahardwarebuffer_layout::compression_eligible(desc.format, desc.dtype);
3441 let scheme = crate::ahardwarebuffer::device_compression_scheme();
3442 let is_ahb = t.memory() == TensorMemory::DmaBuf;
3443 match (request, scheme) {
3444 (_, Some(s)) if eligible && is_ahb => {
3445 t.set_compression_unchecked(Some(s));
3446 }
3447 (Compression::Scheme(requested), _) => {
3448 // Pre-validated eligible + scheme match, so the only
3449 // way here is the allocation resolving off-AHB.
3450 return Err(Error::InvalidOperation(format!(
3451 "image_desc: Compression::Scheme({requested:?}) requested but the allocation resolved to {:?} (not an AHardwareBuffer)",
3452 t.memory()
3453 )));
3454 }
3455 (Compression::Any, _) => {
3456 note_compression_fallback(&format!(
3457 "({:?}, {:?}) {}x{}: eligible={eligible}, scheme={scheme:?}, memory={:?}",
3458 desc.format,
3459 desc.dtype,
3460 desc.width,
3461 desc.height,
3462 t.memory()
3463 ));
3464 }
3465 }
3466 }
3467 #[cfg(not(target_os = "android"))]
3468 {
3469 debug_assert!(matches!(request, Compression::Any));
3470 let _ = request;
3471 note_compression_fallback(&format!(
3472 "({:?}, {:?}) {}x{}: no vendor tile compression on this platform",
3473 desc.format, desc.dtype, desc.width, desc.height
3474 ));
3475 }
3476 }
3477 Ok(t)
3478 }
3479
3480 /// Allocate an image tensor of `width` × `height` in `format`.
3481 ///
3482 /// This is the common-case image constructor. The shape follows from the
3483 /// format's layout — packed `[H, W, C]`, planar `[C, H, W]`, semi-planar
3484 /// `[H + chroma_rows, W]` — and the format is recorded on the tensor, so
3485 /// the image crate and the codec can work from it without out-of-band
3486 /// parameters.
3487 ///
3488 /// `memory` requests a backend, or `None` to auto-select. Auto-select for
3489 /// images is **DMA → Mem**: unlike [`Tensor::new`] it does not try SHM,
3490 /// which offers an image nothing that Mem doesn't (it is not
3491 /// GPU-importable, and Mem always succeeds). Ask for
3492 /// [`TensorMemory::Shm`] explicitly if you need it.
3493 ///
3494 /// `access` declares what the **CPU** intends to do with the buffer and is
3495 /// required; see [`CpuAccess`]. Hardware access — GPU, NPU, ISP, codec —
3496 /// is always implied and never declared. Prefer the narrowest declaration
3497 /// that fits: `None` keeps the allocation eligible for Android vendor tile
3498 /// compression, and anything else pins the layout linear.
3499 ///
3500 /// DMA-backed images always get a **64-byte-aligned row stride**, in every
3501 /// layout, because Mali and Vivante reject an EGLImage import at an
3502 /// unaligned pitch. Read the real pitch back with
3503 /// [`Tensor::effective_row_stride`] rather than assuming `width × bpp`.
3504 ///
3505 /// For a compression request or any other less common knob, build an
3506 /// [`ImageDesc`] and call [`Tensor::image_desc`].
3507 ///
3508 /// # Errors
3509 ///
3510 /// - [`Error::InvalidArgument`] if `width` × `height` is not a valid size
3511 /// for `format` (for example odd dimensions where the format's chroma
3512 /// subsampling forbids them), or, on macOS with an explicit
3513 /// [`TensorMemory::DmaBuf`], if the format cannot be expressed as an
3514 /// image-formatted IOSurface at this width. The message names the
3515 /// aligned width to use.
3516 /// - Whatever the chosen backend returns if the allocation itself fails —
3517 /// for an explicit `memory` request there is no fallback.
3518 ///
3519 /// # Example
3520 ///
3521 /// ```rust
3522 /// use edgefirst_tensor::{CpuAccess, PixelFormat, Tensor, TensorMemory};
3523 ///
3524 /// # fn main() -> Result<(), edgefirst_tensor::Error> {
3525 /// // A camera frame the GPU converts and the CPU never touches.
3526 /// let frame = Tensor::<u8>::image(1920, 1080, PixelFormat::Nv12,
3527 /// Some(TensorMemory::Mem), CpuAccess::None)?;
3528 /// assert_eq!(frame.format(), Some(PixelFormat::Nv12));
3529 /// # Ok(())
3530 /// # }
3531 /// ```
3532 pub fn image(
3533 width: usize,
3534 height: usize,
3535 format: PixelFormat,
3536 memory: Option<TensorMemory>,
3537 access: CpuAccess,
3538 ) -> Result<Self>
3539 where
3540 T: 'static,
3541 {
3542 // Shape comes from the shared `PixelFormat::allocation_shape` helper (packed /
3543 // planar / semi-planar NV12·NV16). NV12 supports odd dimensions via the
3544 // `H + ceil(H/2)` combined-plane height.
3545 // The `T: 'static` bound is required by the macOS IOSurface path below.
3546 let shape = format.allocation_shape(width, height).ok_or_else(|| {
3547 Error::InvalidArgument(format!(
3548 "invalid dimensions {width}x{height} for format {format:?}"
3549 ))
3550 })?;
3551
3552 // macOS Dma path: allocate a format-aware IOSurface (FourCC +
3553 // 2D dimensions) so the GL backend can bind it via
3554 // `EGL_ANGLE_iosurface_client_buffer`. Without this, the IOSurface
3555 // would default to a generic byte buffer (FourCC 'L008') and
3556 // ANGLE would reject the import with `EGL_BAD_ATTRIBUTE`.
3557 //
3558 // Guard: IOSurface rounds `bytes_per_row` up to 64-byte alignment.
3559 // If the natural row pitch (`width * channels * sizeof(T)`) is not
3560 // already 64-byte aligned, the padded allocation cannot be mapped
3561 // as a contiguous packed tensor — CPU reads/writes would use the
3562 // wrong stride.
3563 //
3564 // Explicit-Dma contract: when the caller passes
3565 // `Some(TensorMemory::DmaBuf)` they have asked for an
3566 // **image-formatted IOSurface**. Silently downgrading to the
3567 // generic 'L008' byte-bag when alignment fails buries the
3568 // mismatch — the caller only finds out hours later when ANGLE
3569 // (or any GL importer) rejects the bind with
3570 // `EGL_BAD_ATTRIBUTE`. Same anti-pattern bit us previously on
3571 // Mali GPUs with DMA-BUF padding. The right behaviour is to
3572 // fail loudly here with the alignment requirement spelled out
3573 // so the caller can either pick aligned dimensions, request
3574 // SHM/Mem explicitly, or pass `memory=None` for auto-select.
3575 #[cfg(any(target_os = "macos", target_os = "ios"))]
3576 if matches!(memory, Some(TensorMemory::DmaBuf)) {
3577 // For planar formats the IOSurface stacks channels
3578 // vertically (channels * height rows), so the row stride is
3579 // single-channel width * sizeof(T). Packed formats keep the
3580 // natural width * channels * sizeof(T) stride.
3581 let natural_row_bytes = match format.layout() {
3582 PixelLayout::Planar => width * std::mem::size_of::<T>(),
3583 _ => width * format.channels() * std::mem::size_of::<T>(),
3584 };
3585 // A format with a real IOSurface FourCC (RGBA/BGRA/YUYV packed,
3586 // GREY/NV12/NV16/NV24 as R8) tolerates a non-64-aligned natural
3587 // pitch: the surface is allocated with its own 64-aligned
3588 // `bytes_per_row`, the tensor records that stride below, and a CPU
3589 // map iterates rows correctly via the strided-map path while the GL
3590 // import uses the surface's pitch directly — fully zero-copy.
3591 // Planar (the F16 RGBA16F packing) is consumed flat as
3592 // `[1, C, H, W]` with no stride, so it still requires an aligned
3593 // pitch; and formats without a FourCC would fall through to a
3594 // generic byte-bag GL can't bind. Both fail loudly rather than
3595 // silently downgrade.
3596 let has_image_fourcc = dtype_of::<T>()
3597 .and_then(|dt| crate::iosurface_layout::image_iosurface_layout(format, dt))
3598 .is_some();
3599 let padded_ok = has_image_fourcc && format.layout() != PixelLayout::Planar;
3600 if !natural_row_bytes.is_multiple_of(64) && !padded_ok {
3601 let elem_size = std::mem::size_of::<T>();
3602 let per_pixel_bytes = match format.layout() {
3603 PixelLayout::Planar => elem_size.max(1),
3604 _ => format.channels().max(1) * elem_size.max(1),
3605 };
3606 // Compute the next 64-byte-aligned width by rounding the
3607 // natural row pitch up to the next multiple of 64 and
3608 // dividing back by per-pixel bytes. This handles every
3609 // `per_pixel_bytes` value correctly:
3610 //
3611 // * Divisors of 64 (1/2/4/8/16/32/64) → the suggestion
3612 // is always 64-byte aligned.
3613 // * Non-divisors of 64 (e.g. RGB u8 with 3 B/pixel) →
3614 // the next aligned row pitch may not be an integer
3615 // multiple of per_pixel_bytes (3 doesn't divide 64
3616 // in any way), so a "pad width to N" suggestion is
3617 // structurally impossible — omit the suggestion
3618 // instead of printing a wrong number.
3619 // * per_pixel_bytes > 64 → same situation, also
3620 // omitted; the previous formula divided by zero.
3621 //
3622 // The error always names the alignment requirement
3623 // verbatim and lists the two non-DMA alternatives so
3624 // the caller has at least one always-applicable fix.
3625 let aligned_row_bytes = natural_row_bytes.next_multiple_of(64);
3626 let pad_hint =
3627 if per_pixel_bytes > 0 && aligned_row_bytes.is_multiple_of(per_pixel_bytes) {
3628 let w = aligned_row_bytes / per_pixel_bytes;
3629 format!("Pad width to {w} (the next 64-byte-aligned stride), ")
3630 } else {
3631 String::new()
3632 };
3633 return Err(Error::InvalidArgument(format!(
3634 "Tensor::image: {format:?} {width}x{height} with element \
3635 size {elem_size} produces a {natural_row_bytes}-byte natural \
3636 row pitch, which is not 64-byte aligned. \
3637 IOSurface rounds bytes_per_row up to 64 bytes, so a \
3638 contiguous CPU map of this tensor would read garbage. \
3639 {pad_hint}pass memory=None to auto-fall-back to SHM, or \
3640 pass memory=Some(TensorMemory::Shm) or \
3641 Some(TensorMemory::Mem) explicitly."
3642 )));
3643 }
3644 // Alignment OK. Explicit-Dma contract: the caller asked for an
3645 // **image-formatted, GL-importable** IOSurface, so every failure
3646 // from here on is loud. The old behaviour fell through to the
3647 // generic 'L008' byte-bag `Some(other)` arm below, and the caller
3648 // only found out when the GL import rejected the bind with
3649 // `EGL_BAD_ATTRIBUTE` — the same silent-downgrade anti-pattern as
3650 // the alignment case above. (Semi-planar/Grey u8 map to 'L008'
3651 // *by design* in `image_iosurface_layout` — the R8-plane
3652 // representation the YUV shaders sample — so they return through
3653 // the mapped path and never reach these errors.)
3654 let dtype = dtype_of::<T>().ok_or_else(|| {
3655 Error::InvalidArgument(format!(
3656 "Tensor::image: element type {} has no DType, so no \
3657 image-formatted IOSurface exists. Pass memory=None or \
3658 Some(TensorMemory::Mem) for a CPU tensor.",
3659 std::any::type_name::<T>()
3660 ))
3661 })?;
3662 if crate::iosurface_layout::image_iosurface_layout(format, dtype).is_none() {
3663 return Err(Error::InvalidArgument(format!(
3664 "Tensor::image: no zero-copy IOSurface mapping exists for \
3665 {format:?}/{dtype:?} on macOS/iOS (supported: \
3666 Rgba/Rgb @ U8/I8, Bgra/Yuyv/Grey/Nv12/Nv16/Nv24 @ U8, \
3667 Rgba/PlanarRgb/PlanarRgba @ F16). Pass memory=None to \
3668 auto-select, or Some(TensorMemory::Mem) explicitly, for \
3669 a CPU tensor."
3670 )));
3671 }
3672 // Packed RGB u8/i8 rides an RGBA8888 surface at (W*3/4, H) —
3673 // reject a width the texel packing cannot express up front
3674 // (mirrors the Android pre-guard).
3675 if format == PixelFormat::Rgb
3676 && matches!(dtype, DType::U8 | DType::I8)
3677 && packed_rgb888_layout(width, height).is_none()
3678 {
3679 return Err(Error::InvalidArgument(format!(
3680 "Tensor::image: Rgb {dtype:?} requires width%4==0 for the RGBA8888 \
3681 IOSurface packing (got width={width}). Pad the width, or pass \
3682 memory=Some(TensorMemory::Mem) for a CPU tensor."
3683 )));
3684 }
3685 let storage = TensorStorage::<T>::new_image_iosurface(
3686 width, height, format, dtype, &shape, None,
3687 )?;
3688 let mut t = Self::wrap(storage);
3689 t.format = Some(format);
3690 // IOSurface rounds `bytes_per_row` up to 64 bytes. When that
3691 // pitch exceeds the natural packed/planar row stride, record
3692 // it so CPU consumers iterate rows correctly (the GL import
3693 // already uses the surface's own pitch). For 64-aligned rows
3694 // — the common model-input case — the two match and no stride
3695 // is stored, leaving the flat mapping unchanged.
3696 if let TensorStorage::Dma(ref io) = t.storage {
3697 let bpr = io.bytes_per_row();
3698 if let Some(natural) = t.effective_row_stride() {
3699 if bpr > natural {
3700 t.set_row_stride_unchecked(bpr);
3701 }
3702 }
3703 }
3704 t.cpu_access = access;
3705 return Ok(t);
3706 }
3707
3708 // Android Dma path: allocate a format-aware AHardwareBuffer so the
3709 // GL backend can import it as an EGLImage
3710 // (`eglGetNativeClientBufferANDROID` → `eglCreateImageKHR`).
3711 //
3712 // Geometry: gralloc chooses the row pitch (`desc.stride`) at
3713 // allocation, and pads freely (validated on the Galaxy S26 Ultra,
3714 // where SnapAlloc pads the planar-F16 RGBA16F surface). A padded
3715 // pitch is recorded on the tensor so CPU maps iterate rows via the
3716 // strided-map path; the GL render uses the buffer's own pitch
3717 // through the EGLImage either way, so the GPU path stays fully
3718 // zero-copy. Consumers needing a FLAT layout (the future NPU
3719 // handoff's `[1, C, H, W]` contract) must check `row_stride()` and
3720 // repack when set — flatness is a per-device property here, unlike
3721 // macOS where IOSurface's 64-BYTE alignment keeps model-sized F16
3722 // surfaces naturally flat.
3723 #[cfg(target_os = "android")]
3724 if matches!(memory, Some(TensorMemory::DmaBuf)) {
3725 // Explicit-Dma contract (mirrors the macOS block above): the
3726 // caller asked for an image-formatted, GL-importable
3727 // AHardwareBuffer, so an unmapped combination or a gralloc
3728 // refusal errors here instead of falling through to a BLOB
3729 // byte-bag the GL backend can never import — that downgrade
3730 // only surfaced as a silent per-frame CPU upload.
3731 let dtype = dtype_of::<T>().ok_or_else(|| {
3732 Error::InvalidArgument(format!(
3733 "Tensor::image: element type {} has no DType, so no \
3734 image-formatted AHardwareBuffer exists. Pass memory=None \
3735 or Some(TensorMemory::Mem) for a CPU tensor.",
3736 std::any::type_name::<T>()
3737 ))
3738 })?;
3739 // Planar F16 requires width % 4 == 0 for the RGBA16F
3740 // packing — reject up front with the requirement spelled
3741 // out (mirrors the macOS pre-allocation guard) instead of
3742 // falling through to a byte-bag GL cannot bind.
3743 if format.layout() == PixelLayout::Planar
3744 && dtype == DType::F16
3745 && packed_rgba16f_layout(format, dtype, width, height).is_none()
3746 {
3747 return Err(Error::InvalidArgument(format!(
3748 "Tensor::image: {format:?} F16 requires width%4==0 for the RGBA16F \
3749 AHardwareBuffer packing (got width={width}). Pad the width, or pass \
3750 memory=Some(TensorMemory::Mem) for a CPU tensor."
3751 )));
3752 }
3753 // Packed RGB u8/i8 rides an RGBA8888 surface at (W*3/4, H) —
3754 // the same whole-texel constraint as the F16 packing above.
3755 if format == PixelFormat::Rgb
3756 && matches!(dtype, DType::U8 | DType::I8)
3757 && packed_rgb888_layout(width, height).is_none()
3758 {
3759 return Err(Error::InvalidArgument(format!(
3760 "Tensor::image: Rgb {dtype:?} requires width%4==0 for the RGBA8888 \
3761 AHardwareBuffer packing (got width={width}). Pad the width, or pass \
3762 memory=Some(TensorMemory::Mem) for a CPU tensor."
3763 )));
3764 }
3765 if crate::ahardwarebuffer::image_ahardwarebuffer_layout(format, dtype).is_none() {
3766 return Err(Error::InvalidArgument(format!(
3767 "Tensor::image: no zero-copy AHardwareBuffer mapping exists \
3768 for {format:?}/{dtype:?} on Android (Grey/NV* need \
3769 R8-format buffers, API 29+; the HAL floor is API 26 — see \
3770 `image_ahardwarebuffer_layout`). Pass memory=None to \
3771 auto-select, or Some(TensorMemory::Mem) explicitly, for a \
3772 CPU tensor; camera NV12 stays zero-copy by wrapping the \
3773 camera's own AHardwareBuffer instead of allocating one."
3774 )));
3775 }
3776 let storage = TensorStorage::<T>::new_image_ahardwarebuffer(
3777 width, height, format, dtype, &shape, None, access,
3778 )?;
3779 let mut t = Self::wrap(storage);
3780 t.format = Some(format);
3781 if let TensorStorage::Dma(ref ahb) = t.storage {
3782 // gralloc chooses the row pitch (Qualcomm's
3783 // SnapAlloc pads e.g. the 160-px-wide RGBA16F
3784 // surface of a 640-wide planar F16 target).
3785 // When it exceeds the natural stride, record
3786 // it so CPU consumers iterate rows correctly
3787 // via `effective_row_stride()` — the GPU
3788 // renders through the EGLImage at the
3789 // buffer's real pitch regardless, so the
3790 // render stays fully zero-copy. Consumers
3791 // that need the FLAT `[1, C, H, W]` layout
3792 // (the future NPU handoff) must check
3793 // `row_stride()` is unset and repack — or
3794 // pick an aligned width — rather than assume
3795 // flatness (see the module docs).
3796 let bpr = ahb.bytes_per_row();
3797 if let Some(natural) = t.effective_row_stride() {
3798 if bpr > natural {
3799 log::debug!(
3800 "Tensor::image: gralloc padded the {format:?} \
3801 AHardwareBuffer pitch to {bpr} bytes (natural \
3802 {natural}); recording row stride"
3803 );
3804 t.set_row_stride_unchecked(bpr);
3805 }
3806 }
3807 }
3808 t.cpu_access = access;
3809 return Ok(t);
3810 }
3811
3812 // Windows Dma path: allocate an `ID3D11Texture2D` on the process
3813 // device so the GL backend can bind it through ANGLE's D3D11 texture
3814 // import and a D3D12/CUDA consumer can open its shared NT handle.
3815 //
3816 // Explicit-Dma contract, as on macOS and Android: the caller asked for
3817 // a texture, so a format with no texture layout is an error here
3818 // rather than a silent fall-through to the host-memory arms below,
3819 // which would only surface much later as a failed GPU import.
3820 #[cfg(target_os = "windows")]
3821 if matches!(memory, Some(TensorMemory::DmaBuf)) {
3822 let dtype = dtype_of::<T>().ok_or_else(|| {
3823 Error::InvalidArgument(format!(
3824 "Tensor::image: element type {} has no DType, so no D3D11 texture \
3825 layout exists. Pass memory=None or Some(TensorMemory::Mem) for a \
3826 CPU tensor.",
3827 std::any::type_name::<T>()
3828 ))
3829 })?;
3830 let storage = TensorStorage::<T>::new_image_d3d11(
3831 width, height, format, dtype, &shape, None, access,
3832 )?;
3833 let mut t = Self::wrap(storage);
3834 t.format = Some(format);
3835 record_d3d11_row_stride(&mut t, format, "Tensor::image");
3836 t.cpu_access = access;
3837 // Best-effort CUDA import, as the Linux DMA arms do: a texture a
3838 // CUDA consumer can map is worth the one import at allocation.
3839 t.try_init_d3d11_cuda();
3840 return Ok(t);
3841 }
3842
3843 // Compute the **64-byte-aligned** row stride for every image layout.
3844 //
3845 // Embedded GPUs reject `eglCreateImage` DMA-BUF imports whose row pitch
3846 // is not 64-byte aligned: Mali returns `EGL_BAD_ALLOC`, Vivante
3847 // `EGL_BAD_ACCESS`. This bit packed RGBA/RGB destinations at odd widths
3848 // AND at even non-multiple-of-16 widths (e.g. 321→1284, 322→1288 bytes —
3849 // neither divisible by 64), so an odd-source → RGBA convert failed on
3850 // imx95/imx8mp while succeeding on V3D/Tegra. Semi-planar already aligned
3851 // here; we now align packed and planar identically so every image()
3852 // allocation is GPU-importable regardless of width.
3853 //
3854 // The per-layout natural pitch and total row count:
3855 // * SemiPlanar `[total_h, width]` — pitch = even(width)·elem, rows = total_h
3856 // * Packed `[height, width, ch]` — pitch = width·ch·elem, rows = height
3857 // * Planar `[ch, height, width]` — pitch = width·elem, rows = ch·height
3858 // Allocation byte size = `aligned_stride · total_rows` (NOT the shape
3859 // product, which reflects only the logical width and under-allocates the
3860 // padding on odd / unaligned widths).
3861 let elem = std::mem::size_of::<T>();
3862 let channels = format.channels();
3863 let (natural_stride, total_rows) = match format.layout() {
3864 PixelLayout::SemiPlanar => (width.next_multiple_of(2) * elem, shape[0]),
3865 PixelLayout::Packed => (width * channels * elem, height),
3866 PixelLayout::Planar => (width * elem, channels * height),
3867 };
3868 let aligned_stride = natural_stride.next_multiple_of(64);
3869 let semi = format.layout() == PixelLayout::SemiPlanar;
3870
3871 // DMA buffers MUST carry a 64-aligned row pitch — Mali/Vivante reject a
3872 // DMA-BUF EGLImage whose pitch is not 64-aligned. Semi-planar also needs
3873 // the aligned pitch on every backend (its chroma-plane offset math
3874 // assumes it). Packed/planar on host-only memory (Mem/Shm) keep the
3875 // natural tight pitch so the many flat CPU consumers are unaffected.
3876 let host_stride = if semi { aligned_stride } else { natural_stride };
3877 let host_byte_size = host_stride * total_rows;
3878 #[cfg(target_os = "linux")]
3879 let dma_byte_size = aligned_stride * total_rows;
3880
3881 // `used_stride` is the actual row pitch of the storage created below.
3882 let (storage, used_stride) = match memory {
3883 #[cfg(target_os = "linux")]
3884 Some(TensorMemory::DmaBuf) => (
3885 TensorStorage::<T>::new_dma_with_byte_size(&shape, dma_byte_size, None)?,
3886 aligned_stride,
3887 ),
3888 #[cfg(unix)]
3889 Some(TensorMemory::Shm) => (
3890 TensorStorage::<T>::new_shm_with_byte_size(&shape, host_byte_size, None)?,
3891 host_stride,
3892 ),
3893 Some(TensorMemory::Mem) => (
3894 TensorStorage::<T>::new_mem_with_byte_size(&shape, host_byte_size, None)?,
3895 host_stride,
3896 ),
3897 #[allow(unused_variables)]
3898 Some(other) => {
3899 // PBO and any future variants: fall through to standard new().
3900 return {
3901 let mut t = Self::new(&shape, Some(other), None)?;
3902 t.format = Some(format);
3903 t.cpu_access = access;
3904 Ok(t)
3905 };
3906 }
3907 None => {
3908 // Auto-select priority: DMA → Mem (DMA gets the 64-aligned
3909 // pitch; the Mem fallback keeps the tight host pitch, so the
3910 // recorded stride matches the storage used). Shm is NOT
3911 // auto-selected — it offers no advantage over Mem for an
3912 // in-process image and Mem always succeeds, so it sits below Mem
3913 // and is reached only via an explicit `TensorMemory::Shm`.
3914 #[cfg(target_os = "linux")]
3915 {
3916 match TensorStorage::<T>::new_dma_with_byte_size(&shape, dma_byte_size, None) {
3917 Ok(s) => (s, aligned_stride),
3918 Err(_) => (
3919 TensorStorage::<T>::new_mem_with_byte_size(
3920 &shape,
3921 host_byte_size,
3922 None,
3923 )?,
3924 host_stride,
3925 ),
3926 }
3927 }
3928 #[cfg(not(target_os = "linux"))]
3929 {
3930 (
3931 TensorStorage::<T>::new_mem_with_byte_size(&shape, host_byte_size, None)?,
3932 host_stride,
3933 )
3934 }
3935 }
3936 };
3937
3938 let mut t = Self::wrap(storage);
3939 t.format = Some(format);
3940 // Record the row stride when it exceeds the natural tight pitch (padding
3941 // is present — DMA packed/planar at an unaligned width, or always for
3942 // semi-planar), mirroring the IOSurface path above. Aligned-width and
3943 // host-only packed/planar images keep their flat layout with no explicit
3944 // stride; `effective_row_stride()` then falls back to the identical
3945 // computed pitch. When padding IS present, consumers must iterate rows by
3946 // `effective_row_stride()` to skip it.
3947 if semi || used_stride > natural_stride {
3948 t.set_row_stride_unchecked(used_stride);
3949 }
3950 debug_assert!(
3951 t.row_stride.is_some() || !semi,
3952 "image() must always set row_stride for semi-planar tensors"
3953 );
3954 t.cpu_access = access;
3955 #[cfg(target_os = "linux")]
3956 t.try_init_dma_cuda();
3957 Ok(t)
3958 }
3959
3960 /// Create a DMA-backed image tensor with an explicit row stride that
3961 /// may exceed the natural `width * channels * sizeof(T)` pitch.
3962 ///
3963 /// Used for image tensors that need GPU pitch alignment padding: the
3964 /// underlying DMA-BUF is sized to `row_stride * height` bytes, but
3965 /// the tensor's logical shape stays at `[height, width, channels]`.
3966 /// `width()` / `height()` / `shape()` continue to report the
3967 /// user-requested values; the padding is visible only via
3968 /// `row_stride()` / `effective_row_stride()` and is automatically
3969 /// propagated to the GL backend's EGLImage import so Mali Valhall
3970 /// accepts the buffer.
3971 ///
3972 /// # Supported formats
3973 ///
3974 /// Currently only **packed** pixel layouts (RGBA8, BGRA8, RGB888,
3975 /// Grey, etc.) are supported — the formats the GL backend uses as
3976 /// render destinations. Semi-planar formats (NV12, NV16) come from
3977 /// external allocators (camera capture, video decoders) and are
3978 /// imported via `TensorDyn::from_fd` + `set_row_stride`, which
3979 /// already supports padded strides.
3980 ///
3981 /// # Supported memory
3982 ///
3983 /// Currently only `TensorMemory::DmaBuf` is supported. PBO and Mem
3984 /// storage don't go through EGLImage import so they don't need
3985 /// pitch alignment; if you pass any other memory type this returns
3986 /// `NotImplemented`. `None` (auto-select) is treated as `Dma`.
3987 ///
3988 /// # Errors
3989 ///
3990 /// - `InvalidArgument` if `row_stride_bytes < width * channels * sizeof(T)`
3991 /// (the requested stride would not fit a single row)
3992 /// - `NotImplemented` for non-packed formats or non-DMA memory
3993 /// - `IoError` if the DMA-heap allocation fails (propagated from
3994 /// `DmaTensor::new_with_byte_size`)
3995 pub fn image_with_stride(
3996 width: usize,
3997 height: usize,
3998 format: PixelFormat,
3999 row_stride_bytes: usize,
4000 memory: Option<TensorMemory>,
4001 access: CpuAccess,
4002 ) -> Result<Self> {
4003 #[cfg(not(target_os = "linux"))]
4004 let _ = access;
4005 // DMA backing (the only thing this constructor produces) is
4006 // Linux-only. On macOS/BSD/Windows the non-Linux block below is
4007 // the only compiled body and returns `NotImplemented` directly;
4008 // on Linux the non-Linux block is cfg-removed and the function
4009 // falls through to the real validation + allocation path. Each
4010 // target compiles exactly one of the two blocks, and the block
4011 // serves as the function's tail expression in both cases — so
4012 // neither needs an explicit `return` (avoids
4013 // `clippy::needless_return` on the macOS CI gate).
4014 #[cfg(not(target_os = "linux"))]
4015 {
4016 let _ = (width, height, format, row_stride_bytes, memory);
4017 Err(Error::NotImplemented(
4018 "image_with_stride requires DMA support (Linux only)".to_owned(),
4019 ))
4020 }
4021
4022 #[cfg(target_os = "linux")]
4023 {
4024 if format.layout() != PixelLayout::Packed {
4025 return Err(Error::NotImplemented(format!(
4026 "Tensor::image_with_stride only supports packed pixel layouts, got {format:?}"
4027 )));
4028 }
4029 let elem = std::mem::size_of::<T>();
4030 let min_stride = width
4031 .checked_mul(format.channels())
4032 .and_then(|p| p.checked_mul(elem))
4033 .ok_or_else(|| {
4034 Error::InvalidArgument(format!(
4035 "image_with_stride: width {width} × channels {} × sizeof::<T>={elem} \
4036 overflows usize",
4037 format.channels()
4038 ))
4039 })?;
4040 if row_stride_bytes < min_stride {
4041 return Err(Error::InvalidArgument(format!(
4042 "image_with_stride: row_stride {row_stride_bytes} < minimum {min_stride} \
4043 ({width} px × {} ch × {elem} B)",
4044 format.channels()
4045 )));
4046 }
4047 let total_byte_size = row_stride_bytes.checked_mul(height).ok_or_else(|| {
4048 Error::InvalidArgument(format!(
4049 "image_with_stride: row_stride {row_stride_bytes} × height {height} overflows usize"
4050 ))
4051 })?;
4052
4053 let shape = vec![height, width, format.channels()];
4054
4055 let storage = match memory {
4056 Some(TensorMemory::DmaBuf) | None => {
4057 TensorStorage::<T>::new_dma_with_byte_size(&shape, total_byte_size, None)?
4058 }
4059 Some(other) => {
4060 return Err(Error::NotImplemented(format!(
4061 "image_with_stride: only TensorMemory::DmaBuf is supported, got {other:?}"
4062 )));
4063 }
4064 };
4065
4066 let mut t = Self::wrap(storage);
4067 t.format = Some(format);
4068 t.row_stride = Some(row_stride_bytes);
4069 t.cpu_access = access;
4070 // Match new()/from_fd(): a DMA tensor must attempt CUDA external-
4071 // memory import so a strided DMA buffer is also zero-copy
4072 // CUDA-mappable (no-op when libcudart is absent).
4073 t.try_init_dma_cuda();
4074 Ok(t)
4075 }
4076 }
4077
4078 /// Attach format metadata to an existing tensor.
4079 ///
4080 /// # Arguments
4081 ///
4082 /// * `format` - The pixel format to attach
4083 ///
4084 /// # Returns
4085 ///
4086 /// `Ok(())` on success, with the format stored as metadata on the tensor.
4087 ///
4088 /// # Errors
4089 ///
4090 /// Returns `Error::InvalidShape` if the tensor shape is incompatible with
4091 /// the format's layout (packed expects `[H, W, C]`, planar expects
4092 /// `[C, H, W]`, semi-planar expects `[H*k, W]` with format-specific
4093 /// height constraints).
4094 pub fn set_format(&mut self, format: PixelFormat) -> Result<()> {
4095 let shape = self.shape();
4096 match format.layout() {
4097 PixelLayout::Packed => {
4098 if shape.len() != 3 || shape[2] != format.channels() {
4099 return Err(Error::InvalidShape(format!(
4100 "packed format {format:?} expects [H, W, {}], got {shape:?}",
4101 format.channels()
4102 )));
4103 }
4104 }
4105 PixelLayout::Planar => {
4106 if shape.len() != 3 || shape[0] != format.channels() {
4107 return Err(Error::InvalidShape(format!(
4108 "planar format {format:?} expects [{}, H, W], got {shape:?}",
4109 format.channels()
4110 )));
4111 }
4112 }
4113 PixelLayout::SemiPlanar => {
4114 if shape.len() != 2 {
4115 return Err(Error::InvalidShape(format!(
4116 "semi-planar format {format:?} expects [H*k, W], got {shape:?}"
4117 )));
4118 }
4119 match format {
4120 // Combined-plane height is `H + ceil(H/2)` (luma + chroma
4121 // rows). For even H that is `3H/2` (≡ 0 mod 3); for odd H it
4122 // is `(3H+1)/2` (≡ 2 mod 3). Only totals ≡ 1 mod 3 are
4123 // unreachable, so reject just those — odd-height NV12 is
4124 // valid (e.g. 725 rows for a 483-tall image).
4125 PixelFormat::Nv12 if shape[0] % 3 == 1 => {
4126 return Err(Error::InvalidShape(format!(
4127 "NV12 contiguous shape[0] must be H + ceil(H/2) for some height; \
4128 {} is unreachable (≡ 1 mod 3)",
4129 shape[0]
4130 )));
4131 }
4132 PixelFormat::Nv16 if !shape[0].is_multiple_of(2) => {
4133 return Err(Error::InvalidShape(format!(
4134 "NV16 contiguous shape[0] must be even, got {}",
4135 shape[0]
4136 )));
4137 }
4138 // NV24 (4:4:4): combined-plane height is 3H (Y + 2H chroma).
4139 PixelFormat::Nv24 if !shape[0].is_multiple_of(3) => {
4140 return Err(Error::InvalidShape(format!(
4141 "NV24 contiguous shape[0] must be a multiple of 3 (= 3H), got {}",
4142 shape[0]
4143 )));
4144 }
4145 _ => {}
4146 }
4147 }
4148 }
4149 // Clear stored stride/offset when format changes — they may be invalid
4150 // for the new format. Caller must re-set after changing format.
4151 if self.format != Some(format) {
4152 self.row_stride = None;
4153 self.plane_offset = None;
4154 match self.storage {
4155 TensorStorage::Mem(ref mut m) => m.set_offset(0),
4156 #[cfg(target_os = "linux")]
4157 TensorStorage::Dma(ref mut dma) => dma.mmap_offset = 0,
4158 // Paired with the arm in `set_plane_offset`: a setter that
4159 // takes effect on this backing needs a clear that does too,
4160 // or the wrapper reports `None` while `map()` still starts at
4161 // the old offset -- a stale window rather than a lost one.
4162 // `reshape`'s clear site needs no such arm:
4163 // `IoSurfaceTensor::reshape` zeroes `view_offset` itself
4164 // before the match there runs.
4165 #[cfg(any(target_os = "macos", target_os = "ios"))]
4166 TensorStorage::Dma(ref mut io) => io.view_offset = 0,
4167 // The same pair for D3D11. Unlike `IoSurfaceTensor`,
4168 // `D3d11TextureTensor::reshape` leaves `view_offset` alone, so
4169 // `reshape`'s clear site carries this arm as well.
4170 #[cfg(target_os = "windows")]
4171 TensorStorage::Dma(ref mut tex) => tex.set_view_offset(0),
4172 // Paired with the arm in `set_plane_offset`: a setter that
4173 // takes effect on this backing needs a clear that does too.
4174 // `reshape`'s clear site needs no `Pbo` arm --
4175 // `PboTensor::reshape` zeroes `view_offset` itself before the
4176 // match there runs, exactly as `IoSurfaceTensor::reshape`
4177 // does. `set_logical_shape` deliberately does not, matching
4178 // `Mem`.
4179 TensorStorage::Pbo(ref mut p) => p.view_offset = 0,
4180 // Same pair, `Shm` side. `cfg(unix)` mirrors the `Shm`
4181 // variant's own gate on `TensorStorage`, not a per-OS
4182 // multiplex -- see the comment on the `set_plane_offset`
4183 // arm.
4184 #[cfg(unix)]
4185 TensorStorage::Shm(ref mut s) => s.set_offset(0),
4186 // No-op, pairing with `set_plane_offset`'s Android arm:
4187 // nothing writes `AHardwareBufferTensor::view_offset` there
4188 // either yet, so there is nothing here for a format change
4189 // to clear. See that arm's comment for the full accounting
4190 // of the still-incomplete gap.
4191 #[cfg(target_os = "android")]
4192 TensorStorage::Dma(_) => {}
4193 }
4194 }
4195 self.format = Some(format);
4196 Ok(())
4197 }
4198
4199 /// Set this tensor's logical dimensions and pixel format to a decoded
4200 /// image, reusing the existing allocation. The shape is derived from the
4201 /// format layout; fails with `Error::InsufficientCapacity` if the
4202 /// allocation cannot hold `width`×`height` in `format`, or
4203 /// `Error::InvalidArgument` if the dimensions are invalid for the format.
4204 ///
4205 /// For NV12/NV16/NV24 the buffer width is rounded up to even (a chroma-plane
4206 /// interleaving requirement); the true odd width is reported by the decoder
4207 /// in `ImageInfo` and trimmed by a `convert()` crop. See
4208 /// [`PixelFormat::allocation_shape`].
4209 ///
4210 /// When the backing has a fixed physical row pitch (an IOSurface's
4211 /// 64-aligned `bytesPerRow`) that exceeds the new format's natural row
4212 /// stride — i.e. a reused max-sized pool tensor reconfigured to a smaller
4213 /// image — the physical pitch is preserved as the tensor's `row_stride`.
4214 /// This keeps the **physical grid** (allocation stride/surface) fixed while
4215 /// the **logical ROI** (this image's W×H) changes, so the decode writes rows
4216 /// at the surface's real stride and the GPU samples them at the same stride.
4217 /// Exact-sized buffers (pitch == natural) stay tightly packed unchanged.
4218 pub fn configure_image(
4219 &mut self,
4220 width: usize,
4221 height: usize,
4222 format: PixelFormat,
4223 ) -> Result<()> {
4224 let shape = format.allocation_shape(width, height).ok_or_else(|| {
4225 Error::InvalidArgument(format!(
4226 "invalid dimensions {width}x{height} for format {format:?}"
4227 ))
4228 })?;
4229 // Capture the pre-existing row stride before `set_format` clears it.
4230 // For pool tensors that were allocated at a larger width (e.g. 1920-wide
4231 // pool decoding a 789-wide image), this preserves the backing pitch so
4232 // rows are still written at the correct physical stride.
4233 let prior_stride = self.row_stride;
4234
4235 self.storage.set_logical_shape(&shape)?;
4236 self.set_format(format)?;
4237 // `set_format` keeps the stride when the format is unchanged; the new
4238 // geometry's stride is chosen below, never inherited by accident.
4239 self.row_stride = None;
4240
4241 // Restore the correct row pitch for the new geometry. A DMA buffer of
4242 // ANY layout, and a semi-planar buffer on any backing, MUST carry a
4243 // 64-byte-aligned pitch: Mali/Vivante reject a DMA-BUF EGLImage whose
4244 // row pitch is not 64-aligned (`EGL_BAD_ALLOC` / `EGL_BAD_ACCESS`), and
4245 // semi-planar chroma-offset math assumes it. This mirrors the pitch
4246 // `Tensor::image()` allocates, so a recycled pool buffer imports
4247 // identically to a fresh one — without it a `configure_image()`'d packed
4248 // buffer (e.g. Y800 96-wide → tight pitch 96) failed convert() on
4249 // imx95/imx8mp via the texture-upload fallback while the fresh oracle
4250 // (aligned 128) succeeded. Packed/planar on host-only memory (Mem/Shm)
4251 // keep the tight pitch so flat CPU consumers stay unaffected — matching
4252 // `image()`'s `host_stride` rule.
4253 let elem = std::mem::size_of::<T>();
4254 let channels = format.channels();
4255 let (min_stride, total_rows) = match format.layout() {
4256 PixelLayout::SemiPlanar => (width.next_multiple_of(2) * elem, shape[0]),
4257 PixelLayout::Packed => (width * channels * elem, height),
4258 PixelLayout::Planar => (width * elem, channels * height),
4259 };
4260 let needs_align = self.storage.memory() == TensorMemory::DmaBuf
4261 || format.layout() == PixelLayout::SemiPlanar;
4262
4263 // A backing with its own fixed pitch (IOSurface, AHardwareBuffer, a
4264 // staged D3D11 texture) keeps it whenever it exceeds the new row.
4265 #[cfg(any(
4266 target_os = "macos",
4267 target_os = "ios",
4268 target_os = "android",
4269 target_os = "windows"
4270 ))]
4271 if let Some(pitch) = self.storage.backing_row_stride() {
4272 let active_stride = if pitch > min_stride {
4273 self.set_row_stride_unchecked(pitch);
4274 pitch
4275 } else {
4276 min_stride
4277 };
4278 if !needs_align {
4279 return Ok(());
4280 }
4281 return self.check_configured_stride(active_stride, total_rows);
4282 }
4283
4284 if !needs_align {
4285 return Ok(());
4286 }
4287 // Priority:
4288 // 1. Prior stride (pool reuse): if the pre-existing stride is
4289 // 64-aligned, >= this layout's minimum row, and fits the
4290 // allocation, keep it. This is the hot-loop reuse case (large
4291 // pool, small image).
4292 // 2. Compute a fresh 64-aligned stride for the current width.
4293 let aligned = min_stride.next_multiple_of(64);
4294 let capacity = self.storage.capacity_bytes();
4295
4296 let candidate = if let Some(ps) = prior_stride {
4297 if ps >= min_stride && ps % 64 == 0 && ps * total_rows <= capacity {
4298 ps
4299 } else {
4300 aligned
4301 }
4302 } else {
4303 aligned
4304 };
4305
4306 let active_stride = if candidate * total_rows <= capacity {
4307 self.set_row_stride_unchecked(candidate);
4308 candidate
4309 } else {
4310 // The tight row may still fit where the aligned one does not.
4311 min_stride
4312 };
4313 self.check_configured_stride(active_stride, total_rows)
4314 }
4315
4316 /// The last step of [`configure_image`](Self::configure_image) for an
4317 /// aligned layout: rows that do not fit the allocation are refused here,
4318 /// rather than left to fault on a later map or write.
4319 fn check_configured_stride(&self, active_stride: usize, total_rows: usize) -> Result<()> {
4320 let needed = active_stride * total_rows;
4321 let capacity = self.storage.capacity_bytes();
4322 if needed > capacity {
4323 return Err(Error::InsufficientCapacity { needed, capacity });
4324 }
4325 Ok(())
4326 }
4327
4328 /// Allocate an image tensor sized to hold up to `width`×`height` in
4329 /// `format`, reusable for any smaller image via `configure_image`.
4330 pub fn image_with_capacity(
4331 width: usize,
4332 height: usize,
4333 format: PixelFormat,
4334 memory: Option<TensorMemory>,
4335 access: CpuAccess,
4336 ) -> Result<Self>
4337 where
4338 T: 'static,
4339 {
4340 Self::image(width, height, format, memory, access)
4341 }
4342
4343 /// Pixel format (None if not an image).
4344 pub fn format(&self) -> Option<PixelFormat> {
4345 self.format
4346 }
4347
4348 /// The format and the trailing shape dims that describe one image: packed
4349 /// `[H, W, C]`, planar `[C, H, W]` or semi-planar `[H*k, W]`, optionally
4350 /// behind a leading batch `N` (see [`batch`](Self::batch)). `None` when
4351 /// unformatted or when the shape has any other rank --
4352 /// `set_logical_shape` keeps the format on a shape it no longer describes.
4353 fn image_dims(&self) -> Option<(PixelFormat, &[usize])> {
4354 let fmt = self.format?;
4355 let shape = self.shape();
4356 let rank = image_rank(fmt.layout());
4357 let batch_dims = shape.len().checked_sub(rank).filter(|&n| n <= 1)?;
4358 Some((fmt, &shape[batch_dims..]))
4359 }
4360
4361 /// Image width (None if not an image). For a batched `[N, ...]` tensor
4362 /// this is the width of one element.
4363 pub fn width(&self) -> Option<usize> {
4364 let (fmt, dims) = self.image_dims()?;
4365 match fmt.layout() {
4366 PixelLayout::Packed => Some(dims[1]),
4367 PixelLayout::Planar => Some(dims[2]),
4368 PixelLayout::SemiPlanar => Some(dims[1]),
4369 }
4370 }
4371
4372 /// Image height (None if not an image). For a batched `[N, ...]` tensor
4373 /// this is the height of one element.
4374 ///
4375 /// For semi-planar formats the combined-plane shape row count is divided
4376 /// by the format's luma-to-total ratio to recover logical height. This
4377 /// returns the exact logical height (including odd heights) only because
4378 /// the logical dimensions are tracked separately from the physical shape —
4379 /// `configure_image` stores the actual `(width, height)` in the format's
4380 /// `allocation_shape`, which round-trips losslessly via these accessors.
4381 pub fn height(&self) -> Option<usize> {
4382 let (fmt, dims) = self.image_dims()?;
4383 match fmt.layout() {
4384 PixelLayout::Packed => Some(dims[0]),
4385 PixelLayout::Planar => Some(dims[1]),
4386 PixelLayout::SemiPlanar => {
4387 if self.is_multiplane() {
4388 Some(dims[0])
4389 } else {
4390 match fmt {
4391 PixelFormat::Nv12 => Some(dims[0] * 2 / 3),
4392 PixelFormat::Nv16 => Some(dims[0] / 2),
4393 PixelFormat::Nv24 => Some(dims[0] / 3),
4394 _ => None,
4395 }
4396 }
4397 }
4398 }
4399 }
4400
4401 /// Create from separate Y and UV planes (multiplane NV12/NV16).
4402 pub fn from_planes(luma: Tensor<T>, chroma: Tensor<T>, format: PixelFormat) -> Result<Self> {
4403 if format.layout() != PixelLayout::SemiPlanar {
4404 return Err(Error::InvalidArgument(format!(
4405 "from_planes requires a semi-planar format, got {format:?}"
4406 )));
4407 }
4408 if chroma.format.is_some() || chroma.chroma.is_some() {
4409 return Err(Error::InvalidArgument(
4410 "chroma tensor must be a raw tensor (no format or chroma metadata)".into(),
4411 ));
4412 }
4413 let luma_shape = luma.shape();
4414 let chroma_shape = chroma.shape();
4415 if luma_shape.len() != 2 || chroma_shape.len() != 2 {
4416 return Err(Error::InvalidArgument(format!(
4417 "from_planes expects 2D shapes, got luma={luma_shape:?} chroma={chroma_shape:?}"
4418 )));
4419 }
4420 if luma_shape[1] != chroma_shape[1] {
4421 return Err(Error::InvalidArgument(format!(
4422 "luma width {} != chroma width {}",
4423 luma_shape[1], chroma_shape[1]
4424 )));
4425 }
4426 match format {
4427 PixelFormat::Nv12 => {
4428 if luma_shape[0] % 2 != 0 {
4429 return Err(Error::InvalidArgument(format!(
4430 "NV12 requires even luma height, got {}",
4431 luma_shape[0]
4432 )));
4433 }
4434 if chroma_shape[0] != luma_shape[0] / 2 {
4435 return Err(Error::InvalidArgument(format!(
4436 "NV12 chroma height {} != luma height / 2 ({})",
4437 chroma_shape[0],
4438 luma_shape[0] / 2
4439 )));
4440 }
4441 }
4442 PixelFormat::Nv16 => {
4443 if chroma_shape[0] != luma_shape[0] {
4444 return Err(Error::InvalidArgument(format!(
4445 "NV16 chroma height {} != luma height {}",
4446 chroma_shape[0], luma_shape[0]
4447 )));
4448 }
4449 }
4450 // NV24's chroma plane is full-resolution (2×-wide interleaved UV),
4451 // which the equal-width plane check above doesn't model. Multiplane
4452 // NV24 is unused (the JPEG decoder emits a contiguous NV24 buffer),
4453 // so it's not supported here yet.
4454 _ => {
4455 return Err(Error::InvalidArgument(format!(
4456 "from_planes only supports NV12 and NV16 (NV24 multiplane not yet \
4457 supported — use a contiguous NV24 tensor), got {format:?}"
4458 )));
4459 }
4460 }
4461
4462 Ok(Tensor {
4463 storage: luma.storage,
4464 format: Some(format),
4465 chroma: Some(Box::new(chroma)),
4466 row_stride: luma.row_stride,
4467 plane_offset: luma.plane_offset,
4468 quantization: luma.quantization,
4469 // A multiplane tensor spans two DMA-BUFs (luma + chroma); CUDA
4470 // external-memory import is per-fd, so there is no single device
4471 // pointer for the composite. Any CUDA handle the luma plane carried
4472 // is intentionally dropped — consumers needing CUDA access to
4473 // multiplane data must import each plane independently.
4474 cuda: None,
4475 colorimetry: luma.colorimetry,
4476 cpu_access: luma.cpu_access,
4477 compression: luma.compression,
4478 // A composed multiplane tensor is a whole image, not a sub-view.
4479 view_origin: None,
4480 })
4481 }
4482
4483 /// Whether this tensor uses separate plane allocations.
4484 pub fn is_multiplane(&self) -> bool {
4485 self.chroma.is_some()
4486 }
4487
4488 /// Access the chroma plane for multiplane semi-planar images.
4489 pub fn chroma(&self) -> Option<&Tensor<T>> {
4490 self.chroma.as_deref()
4491 }
4492
4493 /// Mutable access to the chroma plane for multiplane semi-planar images.
4494 pub fn chroma_mut(&mut self) -> Option<&mut Tensor<T>> {
4495 self.chroma.as_deref_mut()
4496 }
4497
4498 /// Independent tensor over this plane's allocation, copying geometry
4499 /// metadata. Used to answer `ef_tensor_chroma` from the handle without a
4500 /// client-side shadow field.
4501 pub(crate) fn clone_chroma_plane(&self) -> Result<Self> {
4502 // Every storage accepted by `from_planes` implements `view` as a
4503 // shared-backing lens. Using that common path keeps
4504 // `is_multiplane() => chroma().is_some()` true for SHM, PBO and
4505 // D3D11 too, instead of silently losing chroma when a hand-written
4506 // per-storage clone omits a variant.
4507 self.subview(0, self.shape())
4508 }
4509
4510 /// Row stride in bytes (`None` = tightly packed).
4511 pub fn row_stride(&self) -> Option<usize> {
4512 self.row_stride
4513 }
4514
4515 /// Effective row stride in bytes: the stored stride if set, otherwise the
4516 /// minimum stride computed from the format, width, and element size.
4517 /// Returns `None` only when no format is set and no explicit stride was
4518 /// stored via [`set_row_stride`](Self::set_row_stride).
4519 ///
4520 /// **GREY note:** `effective_row_stride()` for a GREY tensor returns the
4521 /// tight `width` bytes (no padding), which is what `normalize_to_numpy` and
4522 /// the CPU convert path expect. The codec's internal `native_row_stride`
4523 /// (64-byte-aligned) is used only during decoding and is not propagated to
4524 /// the tensor's stored stride, so callers reading via
4525 /// `effective_row_stride()` always see the tight value for GREY.
4526 pub fn effective_row_stride(&self) -> Option<usize> {
4527 if let Some(s) = self.row_stride {
4528 return Some(s);
4529 }
4530 let fmt = self.format?;
4531 let w = self.width()?;
4532 let elem = std::mem::size_of::<T>();
4533 Some(match fmt.layout() {
4534 PixelLayout::Packed => w * fmt.channels() * elem,
4535 PixelLayout::Planar => w * elem,
4536 // Semi-planar: minimum stride must cover the even width so the
4537 // interleaved chroma columns are byte-aligned on odd-width images.
4538 PixelLayout::SemiPlanar => w.next_multiple_of(2) * elem,
4539 })
4540 }
4541
4542 /// Copy the tensor's logical bytes into `dst`, compacting away any
4543 /// recorded row-stride padding.
4544 ///
4545 /// The flatness helper for NPU handoff: consumers that need a FLAT
4546 /// layout (e.g. `[1, C, H, W]` for NNAPI/LiteRT when the runtime cannot
4547 /// take a padded pitch) call this when [`row_stride`](Self::row_stride)
4548 /// is `Some` — on a tight tensor it degenerates to one memcpy, so it is
4549 /// safe to call unconditionally. `dst.len()` must equal the tight byte
4550 /// footprint (`shape` product × element size). Zero-copy consumers
4551 /// should prefer the buffer handle + [`Tensor::effective_row_stride`]
4552 /// (Self::effective_row_stride) and skip this copy entirely.
4553 pub fn copy_to_flat(&self, dst: &mut [u8]) -> Result<()> {
4554 let tight_bytes = crate::ahardwarebuffer_layout::checked_shape_bytes::<T>(self.shape())?;
4555 if dst.len() != tight_bytes {
4556 return Err(Error::InvalidArgument(format!(
4557 "copy_to_flat: dst is {} bytes but the tensor's tight \
4558 footprint is {tight_bytes} bytes (shape {:?})",
4559 dst.len(),
4560 self.shape()
4561 )));
4562 }
4563 // Read-only: this helper never writes back into the tensor, so a
4564 // read map suffices — and is required on backends (D3D11) that
4565 // reject a writable map on a tensor allocated CpuAccess::Read.
4566 let map = self.map_read()?;
4567 // SAFETY: T is a plain numeric type (crate-wide bound); viewing the
4568 // mapped elements as bytes is sound.
4569 let src: &[u8] = unsafe {
4570 std::slice::from_raw_parts(
4571 map.as_slice().as_ptr() as *const u8,
4572 std::mem::size_of_val(map.as_slice()),
4573 )
4574 };
4575 let Some(stride) = self.row_stride else {
4576 // Tight layout: the mapped window is exactly the logical bytes.
4577 let Some(tight) = src.get(..tight_bytes) else {
4578 return Err(Error::InvalidOperation(format!(
4579 "copy_to_flat: mapped {} bytes < tight footprint {tight_bytes}",
4580 src.len()
4581 )));
4582 };
4583 dst.copy_from_slice(tight);
4584 return Ok(());
4585 };
4586 // Strided: the rows the strided map spans, each holding one tight
4587 // logical row. `rows` is a factor of the element count, so it divides
4588 // the tight footprint exactly.
4589 let rows = strided_rows(self)?;
4590 if rows == 0 {
4591 return Err(Error::InvalidOperation(
4592 "copy_to_flat: a strided tensor with no rows has nothing to copy".into(),
4593 ));
4594 }
4595 let row_bytes = tight_bytes / rows;
4596 let need = (rows - 1)
4597 .checked_mul(stride)
4598 .and_then(|b| b.checked_add(row_bytes))
4599 .ok_or_else(|| {
4600 Error::InvalidOperation(format!(
4601 "copy_to_flat: stride {stride} × rows {rows} overflows usize"
4602 ))
4603 })?;
4604 if src.len() < need {
4605 return Err(Error::InvalidOperation(format!(
4606 "copy_to_flat: mapped {} bytes but strided rows need {need}",
4607 src.len()
4608 )));
4609 }
4610 for r in 0..rows {
4611 dst[r * row_bytes..(r + 1) * row_bytes]
4612 .copy_from_slice(&src[r * stride..r * stride + row_bytes]);
4613 }
4614 Ok(())
4615 }
4616
4617 /// Set the row stride in bytes for externally allocated buffers with
4618 /// row padding (e.g. V4L2 or GStreamer allocators).
4619 ///
4620 /// The stride is propagated to the EGL DMA-BUF import attributes so
4621 /// the GPU interprets the padded buffer layout correctly. Must be
4622 /// called after [`set_format`](Self::set_format) and before the tensor
4623 /// is first passed to `ImageProcessor::convert`. The stored stride
4624 /// is cleared automatically if the pixel format is later changed.
4625 ///
4626 /// No stride-vs-buffer-size validation is performed because the
4627 /// backing allocation size is not reliably known: external DMA-BUFs
4628 /// may be over-allocated by the allocator, and internal tensors store
4629 /// a logical (unpadded) shape. An incorrect stride will be caught by
4630 /// the EGL driver at import time.
4631 ///
4632 /// # Arguments
4633 ///
4634 /// * `stride` - Row stride in bytes. Must be >= the minimum stride for
4635 /// the format (width * channels * sizeof(T) for packed,
4636 /// width * sizeof(T) for planar/semi-planar).
4637 ///
4638 /// # Errors
4639 ///
4640 /// * `InvalidArgument` if no pixel format is set on this tensor
4641 /// * `InvalidArgument` if `stride` is less than the minimum for the
4642 /// format and width
4643 pub fn set_row_stride(&mut self, stride: usize) -> Result<()> {
4644 let fmt = self.format.ok_or_else(|| {
4645 Error::InvalidArgument("cannot set row_stride without a pixel format".into())
4646 })?;
4647 let w = self.width().ok_or_else(|| {
4648 Error::InvalidArgument("cannot determine width for row_stride validation".into())
4649 })?;
4650 let elem = std::mem::size_of::<T>();
4651 let min_stride = match fmt.layout() {
4652 PixelLayout::Packed => w * fmt.channels() * elem,
4653 PixelLayout::Planar => w * elem,
4654 // Semi-planar: minimum must cover even width for chroma alignment.
4655 PixelLayout::SemiPlanar => w.next_multiple_of(2) * elem,
4656 };
4657 if stride < min_stride {
4658 return Err(Error::InvalidArgument(format!(
4659 "row_stride {stride} < minimum {min_stride} for {fmt:?} at width {w}"
4660 )));
4661 }
4662 self.row_stride = Some(stride);
4663 Ok(())
4664 }
4665
4666 /// Set the row stride without format validation.
4667 ///
4668 /// Use this for raw sub-tensors (e.g. chroma planes) that don't carry
4669 /// format metadata. The caller is responsible for ensuring the stride
4670 /// is valid.
4671 pub fn set_row_stride_unchecked(&mut self, stride: usize) {
4672 self.row_stride = Some(stride);
4673 }
4674
4675 /// Builder-style variant of [`set_row_stride`](Self::set_row_stride),
4676 /// consuming and returning `self`.
4677 ///
4678 /// # Errors
4679 ///
4680 /// Same conditions as [`set_row_stride`](Self::set_row_stride).
4681 pub fn with_row_stride(mut self, stride: usize) -> Result<Self> {
4682 self.set_row_stride(stride)?;
4683 Ok(self)
4684 }
4685
4686 /// Byte offset within the DMA-BUF where image data starts (`None` = 0).
4687 pub fn plane_offset(&self) -> Option<usize> {
4688 self.plane_offset
4689 }
4690
4691 /// The parent-image snapshot if this tensor is a [`view`](Self::view)/
4692 /// [`batch`](Self::batch) sub-region; `None` for a whole tensor. The GL
4693 /// backend keys its import on the parent geometry and renders this view as a
4694 /// `glViewport`/`glScissor` ROI at `(x, y, width, height)`. See [`ViewOrigin`].
4695 pub fn view_origin(&self) -> Option<ViewOrigin> {
4696 self.view_origin
4697 }
4698
4699 /// Set the byte offset within the DMA-BUF where image data starts.
4700 ///
4701 /// Propagated to `EGL_DMA_BUF_PLANE0_OFFSET_EXT` on GPU import.
4702 /// Unlike [`set_row_stride`](Self::set_row_stride), no format is required
4703 /// since the offset is format-independent.
4704 pub fn set_plane_offset(&mut self, offset: usize) {
4705 self.plane_offset = Some(offset);
4706 // The offset consulted by `map()` lives inside the storage variant.
4707 // Keep it in sync with the wrapper field for every backing that
4708 // honors it; see also the clear sites in `set_format` and `reshape`.
4709 //
4710 // Note that `TensorStorage::Dma` is a cfg-multiplexed *name* rather
4711 // than one type -- `DmaTensor` on Linux, `IoSurfaceTensor` on
4712 // macOS/iOS, `AHardwareBufferTensor` on Android,
4713 // `D3d11TextureTensor` on Windows -- so each platform needs its own
4714 // arm below; a missing one would silently become a fall-through
4715 // rather than a compile error, which is how macOS went without one
4716 // for a while. The match is kept exhaustive (no `_ => {}`)
4717 // specifically so that hazard is a compile error, not a silent gap.
4718 match self.storage {
4719 TensorStorage::Mem(ref mut m) => m.set_offset(offset),
4720 #[cfg(target_os = "linux")]
4721 TensorStorage::Dma(ref mut dma) => dma.mmap_offset = offset,
4722 // The offset `IoSurfaceTensor::scoped_pin`/`map_inner` add to the
4723 // locked base address. Writing it here is what makes a
4724 // *reconstructed* view address its own sub-region instead of the
4725 // parent surface's origin; `IoSurfaceTensor::view` already set it
4726 // on the fresh-view path, which is why only the descriptor round
4727 // trip was ever wrong. Idempotent against `subview`, which sets
4728 // the same absolute value by both routes -- pinned by
4729 // `nested_iosurface_subviews_do_not_compound_their_plane_offset`.
4730 #[cfg(any(target_os = "macos", target_os = "ios"))]
4731 TensorStorage::Dma(ref mut io) => io.view_offset = offset,
4732 // The offset every `D3d11TextureTensor` pin adds to its host or
4733 // staging base, in backing (pitched) space like `row_stride()`.
4734 // `view()` bounds-checks the value it computes; this route does
4735 // not, so the pins refuse an offset past the backing instead
4736 // (`check_view_offset` in `d3d11/texture.rs`). Pinned by the
4737 // `d3d11` plane-offset tests in `crates/tensor/tests`.
4738 #[cfg(target_os = "windows")]
4739 TensorStorage::Dma(ref mut tex) => tex.set_view_offset(offset),
4740 // The offset `PboTensor::scoped_pin` adds to the mapped base.
4741 // Writing it here is what makes a *reconstructed* view address
4742 // its own sub-region instead of the parent buffer's origin;
4743 // `PboTensor::view` already set it on the fresh-view path, which
4744 // is why only the round trip was ever wrong -- the same shape
4745 // IOSurface had. Idempotent against `subview`, which sets the
4746 // same absolute value by both routes -- pinned by
4747 // `nested_pbo_subviews_do_not_compound_their_plane_offset`.
4748 //
4749 // Not `cfg`-gated, and deliberately so: unlike
4750 // `TensorStorage::Dma` -- a cfg-multiplexed *name* whose gated
4751 // arm silently becomes `_ => {}` on the other targets, the
4752 // mechanism behind this whole issue class -- `TensorStorage::Pbo`
4753 // is one type on every platform.
4754 TensorStorage::Pbo(ref mut p) => p.view_offset = offset,
4755 // The offset `ShmTensor::map_inner`/`host_pin` add to the mmap
4756 // base, exactly like `MemTensor::offset`. `cfg(unix)` here
4757 // mirrors `TensorStorage::Shm`'s own gate on the enum
4758 // definition -- it is not a per-OS multiplex the way
4759 // `TensorStorage::Dma` is, so this arm cannot silently become
4760 // `_ => {}` on a target where the `Shm` variant still exists.
4761 #[cfg(unix)]
4762 TensorStorage::Shm(ref mut s) => s.set_offset(offset),
4763 // STILL INCOMPLETE: `AHardwareBufferTensor` (Android) has a
4764 // `view_offset` that its own `map()` adds and its own `view()`
4765 // sets, but nothing writes it here yet, so restoring an offset
4766 // onto an already-built tensor is silently a no-op on this
4767 // backing -- a caller gets a map at the parent's origin. Not
4768 // reachable through `TensorDyn::import_descriptor` today
4769 // (`import_storage`'s DMABUF arm is `cfg(target_os = "linux")`),
4770 // so it is latent rather than silent. See
4771 // `edgefirst-python-common`'s `interop::apply_plane_offset` for
4772 // the full per-backing accounting. PBO and Shm were on this list
4773 // until Stage B of the client-side-state work gave them the arms
4774 // above.
4775 #[cfg(target_os = "android")]
4776 TensorStorage::Dma(_) => {}
4777 }
4778 }
4779
4780 /// Colorimetry metadata (`None` = undefined; never auto-filled).
4781 /// The CPU access declared for this tensor at allocation (see
4782 /// [`CpuAccess`]). Views share their parent's declaration.
4783 pub fn cpu_access(&self) -> CpuAccess {
4784 self.cpu_access
4785 }
4786
4787 /// Set the declared CPU access without re-allocating — crate-private:
4788 /// the declaration must reflect the underlying allocation's real
4789 /// capabilities (constructors and importers set it; arbitrary widening
4790 /// would defeat the contract).
4791 // Only the Android AHardwareBuffer importer derives a declaration from
4792 // an existing allocation's usage bits.
4793 #[cfg(target_os = "android")]
4794 pub(crate) fn set_cpu_access_unchecked(&mut self, access: CpuAccess) {
4795 self.cpu_access = access;
4796 }
4797
4798 /// The vendor tile-compression scheme recorded at allocation, or
4799 /// `None` for a linear layout. `Some` means the pixels live in a
4800 /// proprietary tile order: the row-stride accessors describe no
4801 /// meaningful linear layout and CPU maps are best-effort (see
4802 /// [`Compression`]). Only Android hardware-only allocations that
4803 /// requested compression record a scheme.
4804 pub fn compression(&self) -> Option<CompressionScheme> {
4805 self.compression
4806 }
4807
4808 /// Record the compression scheme — crate-private: recording is an
4809 /// allocation-time fact ([`Tensor::image_desc`] sets it; arbitrary
4810 /// mutation would misdescribe the physical layout).
4811 // Only the Android allocation path records a scheme.
4812 #[cfg(target_os = "android")]
4813 pub(crate) fn set_compression_unchecked(&mut self, scheme: Option<CompressionScheme>) {
4814 self.compression = scheme;
4815 }
4816
4817 pub fn colorimetry(&self) -> Option<crate::Colorimetry> {
4818 self.colorimetry
4819 }
4820
4821 /// Attach/clear colorimetry metadata.
4822 pub fn set_colorimetry(&mut self, c: Option<crate::Colorimetry>) {
4823 self.colorimetry = c;
4824 }
4825
4826 /// Builder-style colorimetry attach.
4827 pub fn with_colorimetry(mut self, c: crate::Colorimetry) -> Self {
4828 self.colorimetry = Some(c);
4829 self
4830 }
4831
4832 /// Create a zero-copy sub-region view of this tensor's backing buffer.
4833 ///
4834 /// The returned tensor shares this tensor's allocation (no copy) and maps
4835 /// the window `[offset_bytes, offset_bytes + shape.product()*size_of::<T>())`
4836 /// measured from this tensor's own logical start. N sub-views into one
4837 /// parent can be written independently, enabling batched assembly into a
4838 /// single buffer. Identical semantics across `Mem` (shared `Arc`) and
4839 /// `Dma` (shared fd) backings.
4840 ///
4841 /// # Disjointness
4842 ///
4843 /// Independent writes are sound *only* when the windows do not overlap. The
4844 /// shared backing uses interior mutability (`UnsafeCell` cells), so two
4845 /// sub-views whose byte ranges intersect alias the same cells: writing one
4846 /// while reading or writing the other is a data race and therefore
4847 /// **undefined behaviour**. The caller is responsible for keeping the
4848 /// windows disjoint; this method does not check for overlap.
4849 ///
4850 /// # Errors
4851 ///
4852 /// - [`Error::InvalidOperation`] if the backing is not `Mem` or `Dma`, or
4853 /// if `offset_bytes` is not a multiple of `align_of::<T>()`.
4854 /// - [`Error::InsufficientCapacity`] / [`Error::InvalidSize`] if the window
4855 /// exceeds the parent allocation.
4856 pub(crate) fn subview(&self, offset_bytes: usize, shape: &[usize]) -> Result<Tensor<T>> {
4857 // Offset is absolute into the backing allocation: a sub-view of a
4858 // sub-view composes by adding this tensor's own offset.
4859 let abs_offset = self
4860 .plane_offset
4861 .unwrap_or(0)
4862 .checked_add(offset_bytes)
4863 .ok_or(Error::InvalidSize(offset_bytes))?;
4864 // Every backend exposes `view(offset, shape)` via `TensorTrait`, sharing
4865 // the resource AND `BufferIdentity` (unlike `from_fd`/`from_surface`,
4866 // which mint a fresh identity). The GL backend keys the import on the
4867 // shared identity so offset-distinct sub-views of one buffer reuse a
4868 // single import and address their window via `glViewport`. `Mem`/`Shm`
4869 // share via the allocation `Arc` / a cloned fd; `Pbo` via the GL-buffer
4870 // `Arc`; Linux DMA-BUF / macOS IOSurface via the shared fd / CFRetain.
4871 // `TensorStorage::view` performs the one remaining per-variant dispatch.
4872 let mut t = Tensor::wrap(self.storage.view(offset_bytes, shape)?);
4873 // Inherit the parent's image metadata so the view is a ready-to-use
4874 // sub-image (e.g. a `convert()` destination). The offset is applied
4875 // LAST because `set_format` deliberately clears it — the offset is a
4876 // structural property of the sub-region, not format-dependent metadata.
4877 if let Some(fmt) = self.format {
4878 t.set_format(fmt)?;
4879 }
4880 if let Some(rs) = self.row_stride {
4881 t.set_row_stride_unchecked(rs);
4882 }
4883 t.quantization = self.quantization.clone();
4884 // A sub-region of an image carries the parent's colorimetry — it is the
4885 // same pixels, same color encoding. Inherit it like the other image
4886 // metadata above so a sub-view is a faithful convert() source/target.
4887 t.set_colorimetry(self.colorimetry);
4888 // The declared CPU access is a property of the underlying
4889 // allocation, so every view shares the parent's declaration.
4890 t.cpu_access = self.cpu_access;
4891 // Likewise the recorded compression scheme: the view shares the
4892 // parent's physical layout.
4893 t.compression = self.compression;
4894 if abs_offset > 0 {
4895 t.set_plane_offset(abs_offset);
4896 }
4897 Ok(t)
4898 }
4899
4900 /// Borrow batch element `n` of a batched tensor as a zero-copy view.
4901 ///
4902 /// A batched tensor prepends `N` as the leading dimension over the
4903 /// per-element image layout (`[N, H, W, C]` packed or `[N, C, H, W]`
4904 /// planar) — `N` is over the whole per-element block regardless of
4905 /// `HWC`/`CHW`. `batch(n)` returns element `n`: the contiguous per-element
4906 /// region at byte offset `n * element_size`, sharing the parent's
4907 /// `BufferIdentity` and inheriting its format / row stride / colorimetry.
4908 /// `batch(0)` on a tensor with `N == 1` is equivalent to the whole tensor.
4909 ///
4910 /// # Errors
4911 ///
4912 /// - [`Error::BatchIndexOutOfBounds`] if `n >= N`.
4913 /// - [`Error::InvalidShape`] if the tensor is not batched (a formatted
4914 /// tensor whose rank lacks the leading `N`, or an empty shape).
4915 pub fn batch(&self, n: usize) -> Result<Tensor<T>> {
4916 let shape = self.shape();
4917 // With a format we know the exact per-element rank, so a missing leading
4918 // `N` is a misuse we reject rather than silently treating a spatial dim
4919 // as the batch. Raw tensors take shape[0] as `N` by contract.
4920 if let Some(fmt) = self.format {
4921 let elem_rank = image_rank(fmt.layout());
4922 if shape.len() != elem_rank + 1 {
4923 return Err(Error::InvalidShape(format!(
4924 "batch(): tensor is not batched ({fmt:?} expects a leading N over a \
4925 {elem_rank}-D element, got shape {shape:?})"
4926 )));
4927 }
4928 }
4929 let batch = *shape
4930 .first()
4931 .ok_or_else(|| Error::InvalidShape("batch(): empty shape".into()))?;
4932 if n >= batch {
4933 return Err(Error::BatchIndexOutOfBounds { index: n, batch });
4934 }
4935 let elem_shape: Vec<usize> = shape[1..].to_vec();
4936 let elem_count: usize = elem_shape.iter().product();
4937 elem_count
4938 .checked_mul(std::mem::size_of::<T>())
4939 .ok_or(Error::InvalidSize(elem_count))?;
4940 // One element's footprint including row padding: the descriptor's
4941 // stride for `N`, so a padded batch is addressed the way it is mapped.
4942 let dims: Vec<u64> = shape.iter().map(|&d| d as u64).collect();
4943 let elem_bytes = crate::protocol::c_byte_strides(
4944 &dims,
4945 std::mem::size_of::<T>() as i64,
4946 self.format.map(|f| f.layout()),
4947 self.row_stride,
4948 )[0] as usize;
4949 let offset = n.checked_mul(elem_bytes).ok_or(Error::InvalidSize(n))?;
4950 // For a packed `[N, H, W, C]` tensor the N tiles stack vertically in the
4951 // shared buffer, so the GL import sees one `(W, N*H)` parent and each
4952 // tile is the row-band at `y = n*H`. Snapshot that parent so the backend
4953 // imports once and renders the tile via `glViewport`. Non-packed
4954 // (planar/semi-planar) batching keeps the per-slot path for now (planar
4955 // NCHW tiling is a separate step); raw tensors have no pixel geometry.
4956 let view_origin = match self.format.map(|f| f.layout()) {
4957 Some(PixelLayout::Packed) => {
4958 let tile_h = elem_shape[0];
4959 let tile_w = elem_shape[1];
4960 // Per-row pitch of the tall `(W, N*H)` parent — padded stride if
4961 // set, else the tight row width. The GL import keys on this.
4962 let parent_stride = self
4963 .effective_row_stride()
4964 .expect("a formatted tensor of its batched rank has a width");
4965 Some(self.compose_view_origin(tile_w, batch * tile_h, parent_stride, 0, n * tile_h))
4966 }
4967 _ => None,
4968 };
4969 let mut t = self.subview(offset, &elem_shape)?;
4970 t.view_origin = view_origin;
4971 Ok(t)
4972 }
4973
4974 /// Borrow a rectangular spatial sub-region of an image tensor as a
4975 /// zero-copy view — the **destination/source crop** primitive.
4976 ///
4977 /// `region` is in pixels of the image's leading frame. The returned view
4978 /// shares the parent's `BufferIdentity` and addresses the sub-rectangle by
4979 /// offset + the **parent's** row pitch (so each row lands at the correct
4980 /// columns). `convert(src, &mut dst.view(rect), …)` renders into that
4981 /// sub-rectangle of `dst`; a letterbox fit then clears the view and renders
4982 /// the aspect-preserved content into its inner region. `view`/`batch`/the
4983 /// whole tensor are the one coherent destination model — there is no
4984 /// separate `dst_rect`.
4985 ///
4986 /// # Known limitation: bottom-right multi-row views
4987 ///
4988 /// A multi-row view is mapped as `parent_pitch × rows` bytes **from its own
4989 /// offset**, so the window it asks for overruns the allocation by
4990 /// `region.x * bpp` whenever the view is both offset horizontally *and*
4991 /// reaches the parent's last row. Such a view constructs fine but fails at
4992 /// [`map`](TensorTrait::map) time with [`Error::InsufficientCapacity`] —
4993 /// loudly, never as an out-of-bounds access. For a tiled destination this is
4994 /// the bottom row of tiles at every column but the first. The single-row
4995 /// case is already special-cased below (it keeps its own tight stride); the
4996 /// general fix is to clamp the last row's exposure to the view's own row
4997 /// bytes, which is not yet implemented.
4998 ///
4999 /// # Errors
5000 ///
5001 /// - [`Error::RegionOutOfBounds`] if `region` exceeds the image bounds.
5002 /// - [`Error::InvalidOperation`] if the tensor is not a packed-format image
5003 /// (planar/semi-planar spatial sub-rects are not a single strided window;
5004 /// use [`batch`](Self::batch) for batched planar tensors).
5005 pub fn view(&self, region: Region) -> Result<Tensor<T>> {
5006 let fmt = self.format.ok_or_else(|| {
5007 Error::InvalidOperation("view() requires a formatted image tensor".into())
5008 })?;
5009 if fmt.layout() != PixelLayout::Packed {
5010 return Err(Error::InvalidOperation(format!(
5011 "view() supports packed formats only (got {fmt:?}); use batch(n) for batched \
5012 planar tensors"
5013 )));
5014 }
5015 let w = self
5016 .width()
5017 .ok_or_else(|| Error::InvalidOperation("view(): tensor has no image width".into()))?;
5018 let h = self
5019 .height()
5020 .ok_or_else(|| Error::InvalidOperation("view(): tensor has no image height".into()))?;
5021 if !region.fits_within(w, h) {
5022 return Err(Error::RegionOutOfBounds {
5023 region,
5024 bounds: (w, h),
5025 });
5026 }
5027 let elem = std::mem::size_of::<T>();
5028 let bpp = fmt.channels() * elem;
5029 let stride = self
5030 .effective_row_stride()
5031 .expect("a tensor with an image width has a row stride");
5032 let offset = region
5033 .y
5034 .checked_mul(stride)
5035 .and_then(|yo| yo.checked_add(region.x.checked_mul(bpp)?))
5036 .ok_or(Error::InvalidSize(region.y))?;
5037 let sub_shape = fmt
5038 .allocation_shape(region.width, region.height)
5039 .ok_or_else(|| Error::InvalidShape(format!("view(): invalid shape for {fmt:?}")))?;
5040 let mut t = self.subview(offset, &sub_shape)?;
5041 // A multi-row sub-rect must advance rows by the PARENT pitch so each row
5042 // addresses the correct columns. A single-row view uses its own tight
5043 // stride — the parent pitch would make the strided `map()` (which exposes
5044 // `stride × rows`) expose a trailing row that runs past the buffer tail
5045 // for an offset (x>0 / bottom) view. The GL backend does NOT rely on this
5046 // (single-row-tight) `row_stride`: it reads the parent pitch from
5047 // `view_origin.parent_row_stride` so its import/cache pitch stays
5048 // parent-consistent for views of any height — see `ViewOrigin`.
5049 let view_stride = if region.height > 1 {
5050 stride
5051 } else {
5052 region.width * bpp
5053 };
5054 t.set_row_stride_unchecked(view_stride);
5055 // Snapshot the parent `(w, h, row_stride)` so the GL backend imports the
5056 // parent once (keyed on the parent pitch, not this view's possibly-tight
5057 // single-row stride) and renders this sub-rect as a `glViewport`/
5058 // `glScissor` ROI at `(region.x, region.y)`. Composes when viewing an
5059 // existing view.
5060 t.view_origin = Some(self.compose_view_origin(w, h, stride, region.x, region.y));
5061 Ok(t)
5062 }
5063
5064 /// Build the [`ViewOrigin`] for a new sub-region of `self`. When `self` is a
5065 /// whole tensor the snapshot names `self` as the parent; when `self` is
5066 /// already a view, the snapshot keeps the **root** parent and accumulates the
5067 /// local origin so nested views still resolve to one import.
5068 fn compose_view_origin(
5069 &self,
5070 parent_width: usize,
5071 parent_height: usize,
5072 parent_row_stride: usize,
5073 x: usize,
5074 y: usize,
5075 ) -> ViewOrigin {
5076 match self.view_origin {
5077 Some(root) => ViewOrigin {
5078 parent_width: root.parent_width,
5079 parent_height: root.parent_height,
5080 parent_row_stride: root.parent_row_stride,
5081 x: root.x.saturating_add(x),
5082 y: root.y.saturating_add(y),
5083 },
5084 None => ViewOrigin {
5085 parent_width,
5086 parent_height,
5087 parent_row_stride,
5088 x,
5089 y,
5090 },
5091 }
5092 }
5093
5094 /// The wrapped `PboTensor`, for this module's own PBO accessors.
5095 ///
5096 /// Crate-private: `PboTensor<T>` is a `static`-side value with no
5097 /// `dynamic` counterpart to lend (there the buffer lives inside
5098 /// `libedgefirst_tensor.so`), so a public borrow of it could not have
5099 /// the same signature on both backends. Callers use
5100 /// [`Self::pbo_id`]/[`Self::pbo_is_mapped`], which are the whole surface
5101 /// `edgefirst-image` ever used.
5102 pub(crate) fn as_pbo(&self) -> Option<&PboTensor<T>> {
5103 match &self.storage {
5104 TensorStorage::Pbo(p) => Some(p),
5105 _ => None,
5106 }
5107 }
5108
5109 /// The GL buffer ID backing this tensor (for GL backends / the
5110 /// cross-package protocol's `handle`). `None` when the tensor is not
5111 /// PBO-backed.
5112 pub fn pbo_id(&self) -> Option<u32> {
5113 self.as_pbo().map(PboTensor::buffer_id)
5114 }
5115
5116 /// Whether this PBO currently holds -- or is establishing -- a CPU
5117 /// mapping. `None` when the tensor is not PBO-backed.
5118 ///
5119 /// The question a GL caller is really asking is "is this buffer free
5120 /// for GL operations?", so a map or unmap still in flight answers
5121 /// `Some(true)`: the GL side must not touch the buffer until the
5122 /// transition finishes. See [`PboTensor::is_mapped`].
5123 pub fn pbo_is_mapped(&self) -> Option<bool> {
5124 self.as_pbo().map(PboTensor::is_mapped)
5125 }
5126
5127 /// The C-ABI `PboOpsVtable` address backing this tensor (for the
5128 /// cross-package protocol's `ptr` under `kind::PBO`; see
5129 /// [`crate::protocol::TensorDesc::ptr`]'s doc comment). `None` when the
5130 /// tensor is not PBO-backed.
5131 pub fn pbo_vtable_ptr(&self) -> Option<*const std::ffi::c_void> {
5132 self.as_pbo()
5133 .map(|p| p.pbo_vtable() as *const _ as *const std::ffi::c_void)
5134 }
5135
5136 /// A type-erased keepalive that must stay alive for at least as long as
5137 /// [`Self::pbo_vtable_ptr`]'s address is used -- see
5138 /// [`PboTensor::pbo_keepalive`]'s own doc comment. `None` when the
5139 /// tensor is not PBO-backed.
5140 pub fn pbo_keepalive(&self) -> Option<Arc<dyn Send + Sync>> {
5141 self.as_pbo().map(PboTensor::pbo_keepalive)
5142 }
5143
5144 /// Downcast to DMA tensor reference (for EGL import, G2D).
5145 #[cfg(target_os = "linux")]
5146 pub fn as_dma(&self) -> Option<&DmaTensor<T>> {
5147 match &self.storage {
5148 TensorStorage::Dma(d) => Some(d),
5149 _ => None,
5150 }
5151 }
5152
5153 /// Borrow the DMA-BUF file descriptor backing this tensor.
5154 ///
5155 /// # Returns
5156 ///
5157 /// A borrowed reference to the DMA-BUF file descriptor, tied to `self`'s
5158 /// lifetime.
5159 ///
5160 /// # Errors
5161 ///
5162 /// Returns `Error::NotImplemented` if the tensor is not DMA-backed.
5163 #[cfg(target_os = "linux")]
5164 pub fn dmabuf(&self) -> Result<std::os::fd::BorrowedFd<'_>> {
5165 use std::os::fd::AsFd;
5166 match &self.storage {
5167 TensorStorage::Dma(dma) => Ok(dma.fd.as_fd()),
5168 _ => Err(Error::NotImplemented(format!(
5169 "dmabuf requires DMA-backed tensor, got {:?}",
5170 self.storage.memory()
5171 ))),
5172 }
5173 }
5174
5175 /// Construct a Tensor from a PBO tensor (for GL backends that allocate PBOs).
5176 ///
5177 /// Infallible in this backend (`Ok` always) -- kept as `Result` to match
5178 /// `dynamic`'s own `Tensor::from_pbo` (`dynamic_tensor.rs`), which
5179 /// genuinely can fail: it must allocate a real `ef_tensor_*` handle
5180 /// alongside `pbo` to carry shape/dtype/format metadata (this backend
5181 /// has nowhere else to put it), an allocation this backend never
5182 /// needed since it just wraps `pbo` into `TensorStorage::Pbo` directly.
5183 /// One shared, fallible signature lets `edgefirst-image`'s own call
5184 /// sites (`gl/threaded.rs`) stay backend-agnostic.
5185 pub fn from_pbo(pbo: PboTensor<T>) -> Result<Self> {
5186 Ok(Self {
5187 storage: TensorStorage::Pbo(pbo),
5188 format: None,
5189 chroma: None,
5190 row_stride: None,
5191 plane_offset: None,
5192 quantization: None,
5193 cuda: None,
5194 colorimetry: None,
5195 cpu_access: CpuAccess::ReadWrite,
5196 compression: None,
5197 view_origin: None,
5198 })
5199 }
5200
5201 /// The CUDA registration for this tensor, if any (set at creation on CUDA devices).
5202 pub fn cuda(&self) -> Option<&crate::cuda::CudaHandle> {
5203 self.cuda.as_ref()
5204 }
5205
5206 /// Whether a CUDA registration is attached to this tensor.
5207 pub fn is_cuda_attached(&self) -> bool {
5208 self.cuda.is_some()
5209 }
5210
5211 /// Attach a CUDA handle (called by ImageProcessor::create_image after registering a PBO).
5212 pub fn set_cuda_handle(&mut self, h: crate::cuda::CudaHandle) {
5213 self.cuda = Some(h);
5214 }
5215
5216 /// Fast-fail CUDA map: None (no GL routing) when no handle; else map (PBO routes to the GL worker).
5217 ///
5218 /// Returns a scoped [`CudaMap`] guard holding the raw CUDA device pointer
5219 /// for the duration of the mapping. For GL-buffer-backed tensors the unmap is deferred until the
5220 /// guard drops, freeing the PBO for the next `convert()` call. When no CUDA handle is attached
5221 /// (the common case for plain `Mem`/`DMA` tensors without CUDA registration), returns `None`
5222 /// immediately — no GL routing, no allocation.
5223 ///
5224 /// # Example — zero-copy CUDA input with host fallback
5225 ///
5226 /// ```no_run
5227 /// use edgefirst_tensor::{Tensor, TensorMemory, TensorTrait};
5228 /// # fn feed_tensorrt(_dptr: *mut std::ffi::c_void, _bytes: usize) {}
5229 /// # fn demo(t: &Tensor<f32>) {
5230 /// // Try the zero-copy CUDA device pointer first.
5231 /// if let Some(cuda) = t.cuda_map() {
5232 /// feed_tensorrt(cuda.device_ptr(), cuda.len());
5233 /// // `cuda` (a CudaMap guard) unmaps when it goes out of scope, freeing
5234 /// // the GPU buffer for the next convert().
5235 /// } else {
5236 /// // Fall back to the host mapping when no CUDA handle is attached.
5237 /// let _host = t.map().expect("host map fallback must succeed");
5238 /// // `_host` is a TensorMap<f32> that derefs to &[f32].
5239 /// }
5240 /// # }
5241 /// ```
5242 ///
5243 /// # D3D11 texture tensors: the mapping is tight rows
5244 ///
5245 /// On Windows the mapping is a linear copy of the texture laid out at
5246 /// `width * bytes_per_texel`, and `len()` is that times the row count --
5247 /// *not* [`row_stride`](Self::row_stride), which reports the D3D11 staging
5248 /// pitch a CPU map sees and is the larger number on a padded backing. A
5249 /// consumer that strides the device pointer by `row_stride()` reads past
5250 /// the end of the mapping.
5251 ///
5252 /// Semi-planar formats (NV12, NV16, NV24) are the exception: their texture
5253 /// is allocated as wide as its own row pitch, so the tight row *is* the
5254 /// pitch and the mapping equals `row_stride() x combined_height` bytes.
5255 /// Striding the device pointer by `row_stride()` is exactly right there.
5256 pub fn cuda_map(&self) -> Option<crate::cuda::CudaMap<'_>> {
5257 self.cuda.as_ref()?.map()
5258 }
5259
5260 /// Writable counterpart of [`cuda_map`](Self::cuda_map). See
5261 /// [`CudaHandle::map_mut`](crate::cuda::CudaHandle::map_mut) for which
5262 /// backings distinguish the two.
5263 ///
5264 /// The D3D11 arm is the one that does: the guard's release copies the
5265 /// buffer back into the texture. Its layout is [`cuda_map`](Self::cuda_map)'s
5266 /// -- tight rows of `width * bytes_per_texel`, not `row_stride()` -- and
5267 /// writing it at the wrong pitch scrambles the image rather than failing.
5268 pub fn cuda_map_mut(&self) -> Option<crate::cuda::CudaMap<'_>> {
5269 self.cuda.as_ref()?.map_mut()
5270 }
5271
5272 /// Attempt to attach a CUDA `ExternalMemory` handle for DMA-backed tensors.
5273 ///
5274 /// On a CUDA-capable host, imports the DMA-BUF fd via
5275 /// `cudaImportExternalMemory(OpaqueFd)` and maps it to a device pointer.
5276 /// Sets `self.cuda` to a persistent `ExternalMem` handle on success. No-op
5277 /// if CUDA is unavailable, the tensor is not DMA-backed, or a handle is
5278 /// already set. Import failure is silently ignored — the tensor remains
5279 /// usable without a CUDA handle.
5280 ///
5281 /// # RUNTIME-UNVALIDATED
5282 ///
5283 /// No test platform has both `/dev/dma_heap` and a CUDA device. ABI is
5284 /// layout-asserted vs. CUDA 12.6 `driver_types.h`; the mechanism is proven
5285 /// by gpu-probe O5 on Orin. Best-effort: tensor creation never fails here.
5286 #[cfg(target_os = "linux")]
5287 pub fn try_init_dma_cuda(&mut self) {
5288 // Fast-path: already imported, CUDA not available, or not a DMA tensor.
5289 if self.cuda.is_some() || !crate::cuda::is_cuda_available() {
5290 return;
5291 }
5292 let (raw_fd, buf_size) = match &self.storage {
5293 TensorStorage::Dma(dma) => {
5294 use std::os::fd::AsRawFd;
5295 (dma.fd.as_raw_fd(), dma.buf_size)
5296 }
5297 _ => return,
5298 };
5299 if let Some((ext, dptr)) = crate::cuda::import_dma_fd(raw_fd, buf_size) {
5300 self.cuda = Some(crate::cuda::CudaHandle::new_external(ext, dptr, buf_size));
5301 }
5302 }
5303
5304 /// Attempt to attach a CUDA external-memory handle for D3D11 texture
5305 /// tensors — the Windows counterpart of
5306 /// [`try_init_dma_cuda`](Self::try_init_dma_cuda).
5307 ///
5308 /// Imports the texture's shared NT handle via
5309 /// `cudaImportExternalMemory(D3D11Resource)` and the process device's
5310 /// fence as an external semaphore, so [`cuda_map`](Self::cuda_map) hands
5311 /// out a linear device buffer whose refresh is ordered after the D3D11
5312 /// work a producer recorded with [`set_gpu_write`](Self::set_gpu_write).
5313 ///
5314 /// No-op when a handle is already attached or the tensor is not
5315 /// texture-backed. Import failure is silently ignored: the WARP adapter
5316 /// has no CUDA device at all, and a texture whose import fails is still a
5317 /// valid tensor whose `cuda_map()` is `None`.
5318 #[cfg(target_os = "windows")]
5319 pub fn try_init_d3d11_cuda(&mut self) {
5320 if self.cuda.is_some() {
5321 return;
5322 }
5323 let TensorStorage::Dma(tex) = &self.storage else {
5324 return;
5325 };
5326 if let Some(h) = tex.try_init_cuda() {
5327 self.cuda = Some(h);
5328 }
5329 }
5330}
5331
5332/// Windows D3D11 texture accessors. No `'static` bound: none of these needs
5333/// `dtype_of::<T>()`.
5334#[cfg(all(feature = "static", target_os = "windows"))]
5335impl<T> Tensor<T>
5336where
5337 T: Num + Clone + fmt::Debug + Send + Sync,
5338{
5339 /// The borrowed `ID3D11Texture2D*` backing this tensor, valid while the
5340 /// tensor lives. `None` when the tensor is not texture-backed.
5341 pub fn d3d11_texture(&self) -> Option<*mut std::ffi::c_void> {
5342 match &self.storage {
5343 TensorStorage::Dma(t) => Some(t.texture_ptr()),
5344 _ => None,
5345 }
5346 }
5347
5348 /// The texture geometry the HAL chose for this image: DXGI format, texel
5349 /// dimensions and the matching GL internal format. `None` when the tensor
5350 /// is not texture-backed.
5351 pub fn d3d11_layout(&self) -> Option<crate::d3d11_layout::D3d11ImageLayout> {
5352 match &self.storage {
5353 TensorStorage::Dma(t) => Some(t.layout()),
5354 _ => None,
5355 }
5356 }
5357
5358 /// A duplicated NT handle the caller owns: open it on a D3D12 device,
5359 /// another D3D11 device, or CUDA, or duplicate it into another process.
5360 /// Closing it does not affect this tensor.
5361 pub fn d3d11_shared_handle(&self) -> Result<std::os::windows::io::OwnedHandle> {
5362 match &self.storage {
5363 TensorStorage::Dma(t) => t.shared_handle(),
5364 _ => Err(Error::NotImplemented(
5365 "not a D3D11 texture tensor".to_owned(),
5366 )),
5367 }
5368 }
5369
5370 /// The fence handle plus value a GPU consumer waits on before reading this
5371 /// texture, or `None` when no GPU write has been recorded on it.
5372 pub fn gpu_completion(&self) -> Result<Option<crate::d3d11::GpuCompletion>> {
5373 match &self.storage {
5374 TensorStorage::Dma(t) => t.gpu_completion(),
5375 _ => Err(Error::NotImplemented(
5376 "not a D3D11 texture tensor".to_owned(),
5377 )),
5378 }
5379 }
5380
5381 /// Record that GPU work writing this texture completes at `value` of the
5382 /// process device's fence. Monotonic: an older value never displaces a
5383 /// newer one, so recording out of order is safe rather than merely
5384 /// tolerated.
5385 ///
5386 /// Takes `&self`, not `&mut self`: the fence slot is an `AtomicU64` and
5387 /// this is a `fetch_max` on it. A producer records a write on the very
5388 /// tensor its consumers already hold and are reading, so demanding
5389 /// exclusive access here would forbid the one call this exists for.
5390 pub fn set_gpu_write(&self, value: u64) -> Result<()> {
5391 match &self.storage {
5392 TensorStorage::Dma(t) => {
5393 t.set_gpu_write(value);
5394 Ok(())
5395 }
5396 _ => Err(Error::NotImplemented(
5397 "not a D3D11 texture tensor".to_owned(),
5398 )),
5399 }
5400 }
5401
5402 /// The tensor's own NT handle value -- not a duplicate -- for descriptors
5403 /// whose consumer keeps the producing tensor alive.
5404 ///
5405 /// Borrowed rather than duplicated because the C exports (task A6) hand
5406 /// out owned duplicates, which a `Copy` descriptor cannot close; the
5407 /// cross-package protocol's `kind::D3D11_TEXTURE` reads this instead.
5408 pub(crate) fn d3d11_shared_handle_value(&self) -> Option<usize> {
5409 match &self.storage {
5410 TensorStorage::Dma(t) => Some(t.shared_handle_value()),
5411 _ => None,
5412 }
5413 }
5414}
5415
5416/// The recorded GPU write value, on every platform: the descriptor path and
5417/// the Python bindings forward it without a `cfg` of their own, and only
5418/// Windows can record one.
5419#[cfg(feature = "static")]
5420impl<T> Tensor<T>
5421where
5422 T: Num + Clone + fmt::Debug + Send + Sync,
5423{
5424 /// The fence value of the newest GPU write recorded on this tensor, or 0
5425 /// when there is none: the `value` of `gpu_completion` without the
5426 /// duplicated fence handle, for a consumer that already holds the process
5427 /// fence. 0 for a tensor that is not a D3D11 texture, and on every
5428 /// platform other than Windows.
5429 pub fn gpu_write_value(&self) -> u64 {
5430 #[cfg(target_os = "windows")]
5431 {
5432 if let TensorStorage::Dma(t) = &self.storage {
5433 return t.last_gpu_write();
5434 }
5435 }
5436 0
5437 }
5438}
5439
5440/// Records the D3D11 backing's row pitch on `t` when it exceeds the natural
5441/// stride for its format and width.
5442///
5443/// The staging pitch is the row stride CPU consumers see, and for a
5444/// semi-planar tensor it is also the texture's own width, which the GL
5445/// backend's Path B shader takes as its `tex_width`. Three constructors need
5446/// it identically -- `Tensor::image`'s Windows arm and the two wrap
5447/// constructors below -- and a wrapped texture is the case that makes it
5448/// load-bearing: without it such a tensor reports the natural `even(width)`
5449/// stride while its texture is the driver's pitch wide, and the shader samples
5450/// a padded texture at the image width.
5451///
5452/// `t.format` must already be set: the natural stride comes from it.
5453#[cfg(all(feature = "static", target_os = "windows"))]
5454fn record_d3d11_row_stride<T>(t: &mut Tensor<T>, format: PixelFormat, what: &str)
5455where
5456 T: Num + Clone + fmt::Debug + Send + Sync,
5457{
5458 let pitch = match &t.storage {
5459 TensorStorage::Dma(tex) => tex.image_backing_row_stride(),
5460 _ => None,
5461 };
5462 let (Some(pitch), Some(natural)) = (pitch, t.effective_row_stride()) else {
5463 return;
5464 };
5465 if pitch > natural {
5466 log::debug!(
5467 "{what}: the D3D11 backing pitch for the {format:?} texture is {pitch} bytes \
5468 (natural {natural}); recording row stride"
5469 );
5470 t.set_row_stride_unchecked(pitch);
5471 }
5472}
5473
5474/// The two texture-importing constructors. Separate from the accessors above
5475/// because `dtype_of::<T>()` needs `T: 'static`.
5476#[cfg(all(feature = "static", target_os = "windows"))]
5477impl<T> Tensor<T>
5478where
5479 T: Num + Clone + fmt::Debug + Send + Sync + 'static,
5480{
5481 /// Wrap an existing `ID3D11Texture2D` as a tensor. The texture is checked
5482 /// against the layout the HAL would have chosen for `format`/`width`/
5483 /// `height` and must live on the HAL's process device; ownership stays
5484 /// with the caller, since this takes its own reference.
5485 ///
5486 /// The backing's row pitch is recorded as the tensor's
5487 /// [`row_stride`](Self::row_stride) when it exceeds the natural stride, so
5488 /// a wrapped semi-planar texture -- whose width *is* the pitch -- reports
5489 /// the geometry its consumers address it by.
5490 ///
5491 /// # Safety
5492 ///
5493 /// `texture` must be null or a live `ID3D11Texture2D` created on the HAL
5494 /// device ([`crate::d3d11::device()`]).
5495 pub unsafe fn from_d3d11_texture(
5496 texture: *mut std::ffi::c_void,
5497 width: usize,
5498 height: usize,
5499 format: PixelFormat,
5500 access: CpuAccess,
5501 name: Option<&str>,
5502 ) -> Result<Self> {
5503 let dtype = dtype_of::<T>().ok_or_else(|| {
5504 Error::InvalidArgument(format!(
5505 "from_d3d11_texture: element type {} has no DType",
5506 std::any::type_name::<T>()
5507 ))
5508 })?;
5509 let shape = format.allocation_shape(width, height).ok_or_else(|| {
5510 Error::InvalidArgument(format!(
5511 "invalid dimensions {width}x{height} for format {format:?}"
5512 ))
5513 })?;
5514 // SAFETY: the caller guarantees `texture` is null or a live
5515 // `ID3D11Texture2D`; the storage takes its own reference.
5516 let storage = unsafe {
5517 crate::d3d11::texture::D3d11TextureTensor::<T>::from_d3d11_texture(
5518 texture, width, height, format, dtype, &shape, access, name,
5519 )
5520 }?;
5521 let mut t = Self::wrap(TensorStorage::Dma(storage));
5522 t.format = Some(format);
5523 record_d3d11_row_stride(&mut t, format, "from_d3d11_texture");
5524 t.cpu_access = access;
5525 t.try_init_d3d11_cuda();
5526 Ok(t)
5527 }
5528
5529 /// Open a shared texture by its NT handle on the HAL's device. When
5530 /// `completion` is given, its fence value is waited for on the immediate
5531 /// context, so a same-device reader needs no further ordering. Both
5532 /// handles stay owned by the caller: this duplicates what it keeps.
5533 ///
5534 /// Records the backing's row pitch as
5535 /// [`row_stride`](Self::row_stride) on the same rule as
5536 /// [`from_d3d11_texture`](Self::from_d3d11_texture).
5537 ///
5538 /// # Safety
5539 ///
5540 /// `handle` must be an NT shared handle of a D3D11 texture, valid in this
5541 /// process, and `completion`'s handle a shared fence handle.
5542 pub unsafe fn from_d3d11_shared_handle(
5543 handle: std::os::windows::io::RawHandle,
5544 width: usize,
5545 height: usize,
5546 format: PixelFormat,
5547 access: CpuAccess,
5548 completion: Option<(std::os::windows::io::RawHandle, u64)>,
5549 name: Option<&str>,
5550 ) -> Result<Self> {
5551 let dtype = dtype_of::<T>().ok_or_else(|| {
5552 Error::InvalidArgument(format!(
5553 "from_d3d11_shared_handle: element type {} has no DType",
5554 std::any::type_name::<T>()
5555 ))
5556 })?;
5557 let shape = format.allocation_shape(width, height).ok_or_else(|| {
5558 Error::InvalidArgument(format!(
5559 "invalid dimensions {width}x{height} for format {format:?}"
5560 ))
5561 })?;
5562 // SAFETY: the caller guarantees the handles are valid in this process.
5563 let storage = unsafe {
5564 crate::d3d11::texture::D3d11TextureTensor::<T>::from_d3d11_shared_handle(
5565 handle, width, height, format, dtype, &shape, access, completion, name,
5566 )
5567 }?;
5568 let mut t = Self::wrap(TensorStorage::Dma(storage));
5569 t.format = Some(format);
5570 record_d3d11_row_stride(&mut t, format, "from_d3d11_shared_handle");
5571 t.cpu_access = access;
5572 t.try_init_d3d11_cuda();
5573 Ok(t)
5574 }
5575}
5576
5577// Quantization accessors — type-gated to integer element types via the
5578// sealed `IntegerType` trait. Calling `.quantization()` on a `Tensor<f32>`
5579// produces a compile error, not a runtime one.
5580#[cfg(feature = "static")]
5581impl<T> Tensor<T>
5582where
5583 T: IntegerType + Num + Clone + fmt::Debug + Send + Sync,
5584{
5585 /// Quantization metadata for this tensor, if set.
5586 pub fn quantization(&self) -> Option<&Quantization> {
5587 self.quantization.as_ref()
5588 }
5589
5590 /// Attach quantization metadata to this tensor. Validates against the
5591 /// tensor's shape — returns [`Error::QuantizationInvalid`] on any
5592 /// inconsistency (mismatched scale/zp lengths, out-of-range axis, etc.).
5593 pub fn set_quantization(&mut self, q: Quantization) -> Result<()> {
5594 q.validate(self.shape())?;
5595 self.quantization = Some(q);
5596 Ok(())
5597 }
5598
5599 /// Builder-style variant of [`Self::set_quantization`]. Consumes `self`
5600 /// and returns `Result<Self>` — on success yields the tensor with the
5601 /// attached quantization; on validation failure returns
5602 /// [`Error::QuantizationInvalid`] and drops `self` (the tensor is not
5603 /// returned in the error arm).
5604 pub fn with_quantization(mut self, q: Quantization) -> Result<Self> {
5605 self.set_quantization(q)?;
5606 Ok(self)
5607 }
5608
5609 /// Clear any quantization metadata on this tensor.
5610 pub fn clear_quantization(&mut self) {
5611 self.quantization = None;
5612 }
5613}
5614
5615/// Rank of one image's shape in `layout`: `[H, W, C]` / `[C, H, W]` are 3-D,
5616/// semi-planar `[H*k, W]` is 2-D. A batched tensor has one more.
5617#[cfg(feature = "static")]
5618fn image_rank(layout: PixelLayout) -> usize {
5619 match layout {
5620 PixelLayout::SemiPlanar => 2,
5621 PixelLayout::Packed | PixelLayout::Planar => 3,
5622 }
5623}
5624
5625#[cfg(feature = "static")]
5626fn map_strided_host_view<'a, T>(
5627 tensor: &Tensor<T>,
5628 stride: usize,
5629 access: CpuAccess,
5630 non_blocking: bool,
5631) -> Result<crate::view::HostView<'a, T>>
5632where
5633 T: Num + Clone + fmt::Debug + Send + Sync + 'a,
5634{
5635 // Rows sit at `stride`-byte spacing; see `strided_rows` for the count.
5636 let rows = strided_rows(tensor)?;
5637 let total_bytes = stride.checked_mul(rows).ok_or_else(|| {
5638 Error::InvalidOperation(format!(
5639 "Tensor::map: row_stride {stride} × rows {rows} overflows usize"
5640 ))
5641 })?;
5642 map_strided_storage(tensor, stride, rows, total_bytes, access, non_blocking)
5643}
5644
5645/// Rows a strided map or flat copy of `tensor` spans at its row stride: a
5646/// raw tensor's leading dim, else every image row of every batch element --
5647/// `H` per packed image, `C * H` per planar one (C planes of H rows), and the
5648/// combined `H*k` per semi-planar one. No product here can overflow: each is
5649/// a factor of the shape's element count, which the allocation already holds.
5650#[cfg(feature = "static")]
5651fn strided_rows<T>(tensor: &Tensor<T>) -> Result<usize>
5652where
5653 T: Num + Clone + fmt::Debug + Send + Sync,
5654{
5655 let shape = tensor.shape();
5656 let Some(fmt) = tensor.format else {
5657 return shape.first().copied().ok_or_else(|| {
5658 Error::InvalidOperation("strided access requires a non-empty shape".into())
5659 });
5660 };
5661 let (_, dims) = tensor.image_dims().ok_or_else(|| {
5662 Error::InvalidOperation(format!(
5663 "strided access to a {fmt:?} tensor requires its image shape, optionally \
5664 batched; got {shape:?}"
5665 ))
5666 })?;
5667 let per_image = match fmt.layout() {
5668 PixelLayout::Planar => dims[0] * dims[1],
5669 PixelLayout::Packed | PixelLayout::SemiPlanar => dims[0],
5670 };
5671 let batch: usize = shape[..shape.len() - dims.len()].iter().product();
5672 Ok(batch * per_image)
5673}
5674
5675#[cfg(feature = "static")]
5676fn map_strided_storage<'a, T>(
5677 tensor: &Tensor<T>,
5678 _stride: usize,
5679 _rows: usize,
5680 total_bytes: usize,
5681 access: CpuAccess,
5682 _non_blocking: bool,
5683) -> Result<crate::view::HostView<'a, T>>
5684where
5685 T: Num + Clone + fmt::Debug + Send + Sync + 'a,
5686{
5687 match &tensor.storage {
5688 #[cfg(target_os = "linux")]
5689 TensorStorage::Dma(dma) => {
5690 // `set_row_stride()` only validates `stride >= min_stride`,
5691 // not that `stride × rows` fits the DMA-BUF, so re-check
5692 // here — mapping past `buf_size` would SIGBUS on access.
5693 // `buf_size` is real kernel truth either way, not the
5694 // importer's say-so: for a self-allocated tensor it is
5695 // the allocation size; for an imported one it is
5696 // `fstat`-derived when the kernel reports a usable size
5697 // (`DmaTensor::from_fd`), or the declared logical size in
5698 // the rare fallback -- either way an out-of-range mmap
5699 // against the real dma-buf is still rejected by the
5700 // kernel, so this bound stays sound regardless.
5701 let available_bytes = dma.buf_size.saturating_sub(dma.mmap_offset);
5702 if total_bytes > available_bytes {
5703 return Err(Error::InvalidOperation(format!(
5704 "Tensor::map: strided mapping needs {total_bytes} bytes \
5705 but DMA buffer only has {available_bytes} available \
5706 (buf_size={}, mmap_offset={}, stride={_stride}, rows={_rows}); \
5707 the row_stride was likely set larger than the original allocation",
5708 dma.buf_size, dma.mmap_offset
5709 )));
5710 }
5711 dma.map_with_byte_size(total_bytes, access)
5712 }
5713 TensorStorage::Mem(mem) => {
5714 let capacity = tensor.storage.capacity_bytes();
5715 if total_bytes > capacity {
5716 return Err(Error::InsufficientCapacity {
5717 needed: total_bytes,
5718 capacity,
5719 });
5720 }
5721 mem.map_with_byte_size(total_bytes, access)
5722 }
5723 #[cfg(unix)]
5724 TensorStorage::Shm(shm) => {
5725 let capacity = tensor.storage.capacity_bytes();
5726 if total_bytes > capacity {
5727 return Err(Error::InsufficientCapacity {
5728 needed: total_bytes,
5729 capacity,
5730 });
5731 }
5732 shm.map_with_byte_size(total_bytes, access)
5733 }
5734 // macOS/iOS: `TensorStorage::Dma` is the IOSurface. The lock yields
5735 // the full surface base address, and the row pitch
5736 // (`IOSurfaceGetBytesPerRow`) is known from the API for both
5737 // self-allocated and imported surfaces — unlike a foreign
5738 // DMA-BUF — so a strided CPU view is sound and zero-copy.
5739 #[cfg(any(target_os = "macos", target_os = "ios"))]
5740 TensorStorage::Dma(io) => {
5741 // A sub-view's window is `buf_size − view_offset`; the strided
5742 // span must fit the window, not the whole surface.
5743 let available = io.buf_size.saturating_sub(io.view_offset);
5744 if total_bytes > available {
5745 return Err(Error::InsufficientCapacity {
5746 needed: total_bytes,
5747 capacity: available,
5748 });
5749 }
5750 io.map_with_byte_size(total_bytes, access)
5751 }
5752 // Android: `TensorStorage::Dma` is the AHardwareBuffer. The lock
5753 // yields the full buffer base address, and the row pitch is
5754 // known from the allocator-filled descriptor — so a strided CPU
5755 // view is sound and zero-copy, same as IOSurface.
5756 #[cfg(target_os = "android")]
5757 TensorStorage::Dma(ahb) => {
5758 // A sub-view's window is `buf_size − view_offset`; the strided
5759 // span must fit the window, not the whole buffer.
5760 let available = ahb.buf_size.saturating_sub(ahb.view_offset);
5761 if total_bytes > available {
5762 return Err(Error::InsufficientCapacity {
5763 needed: total_bytes,
5764 capacity: available,
5765 });
5766 }
5767 ahb.map_with_byte_size(total_bytes, access)
5768 }
5769 // Windows: `TensorStorage::Dma` is the D3D11 texture, whose CPU side
5770 // is the staging copy at the driver's pitch -- the very pitch
5771 // `Tensor::image` recorded as this tensor's row stride. Its
5772 // `map_with_byte_size` bounds-checks the span against this tensor's
5773 // own window, as the IOSurface arm above does. This is also the only
5774 // backing with a real non-blocking map, so it is the only arm that
5775 // reads `_non_blocking`.
5776 #[cfg(target_os = "windows")]
5777 TensorStorage::Dma(tex) => {
5778 if _non_blocking {
5779 tex.try_map_with_byte_size(total_bytes, access)
5780 } else {
5781 tex.map_with_byte_size(total_bytes, access)
5782 }
5783 }
5784 TensorStorage::Pbo(pbo) => {
5785 // PBO: the GPU-side allocation may have a padded row stride
5786 // (e.g. 64-byte aligned). Expose the full padded buffer so a
5787 // CPU producer (JPEG decoder) and a strided convert source
5788 // can iterate rows via `effective_row_stride()` without
5789 // running past the slice — the logical `pbo.map()` view would
5790 // stop after `shape.product()` and lose bytes past row 0.
5791 // A sub-view's window is `capacity − view_offset`.
5792 let available = pbo.capacity_bytes().saturating_sub(pbo.view_offset);
5793 if total_bytes > available {
5794 return Err(Error::InsufficientCapacity {
5795 needed: total_bytes,
5796 capacity: available,
5797 });
5798 }
5799 pbo.map_with_byte_size(total_bytes, access)
5800 }
5801 // Every `TensorStorage` variant this build can construct is
5802 // matched explicitly above (Mem/Shm/Dma/Pbo, each cfg-gated
5803 // to the platforms it exists on) -- this catch-all exists
5804 // only so the match stays exhaustive across every
5805 // OS/cfg combination without a matching arm on every one of
5806 // them, not because any variant is meant to fall through to
5807 // it. `#[allow(unreachable_patterns)]` because on any given
5808 // platform it genuinely is unreachable.
5809 #[allow(unreachable_patterns)]
5810 _ => Err(Error::InvalidOperation(
5811 "CPU mapping of strided tensors is supported only for HAL-allocated \
5812 Mem/Shm (any platform), DMA (Linux, self-allocated or imported), \
5813 IOSurface (macOS), AHardwareBuffer (Android), D3D11 textures \
5814 (Windows), and PBO"
5815 .into(),
5816 )),
5817 }
5818}
5819
5820#[cfg(feature = "static")]
5821impl<T> Tensor<T>
5822where
5823 T: Num + Clone + fmt::Debug + Send + Sync,
5824{
5825 /// The shared body of [`TensorTrait::map_with`] and
5826 /// [`TensorTrait::try_map_with`]: the direction check, the
5827 /// declared-vs-requested telemetry, the strided and plane-offset routing,
5828 /// and finally the backing's own map. `non_blocking` chooses which of the
5829 /// backing's two maps that last step takes; everything before it is the
5830 /// same either way.
5831 fn map_impl<'a>(
5832 &self,
5833 access: CpuAccess,
5834 non_blocking: bool,
5835 ) -> Result<crate::view::HostView<'a, T>>
5836 where
5837 T: 'a,
5838 {
5839 let _span = tracing::trace_span!(
5840 "tensor.map",
5841 memory = ?self.storage.memory(),
5842 ?access,
5843 non_blocking,
5844 )
5845 .entered();
5846 if access == CpuAccess::None {
5847 return Err(Error::InvalidArgument(
5848 "map_with(CpuAccess::None) is not a mappable direction — use \
5849 map_read()/map_write()/map_mut()"
5850 .into(),
5851 ));
5852 }
5853 // Declared-vs-requested telemetry (all platforms): mapping beyond
5854 // the allocation-time declaration is best-effort — tolerated where
5855 // the backing is CPU-mappable regardless (Mem/Shm/dma-buf/
5856 // IOSurface), refused by the Android backend for CpuAccess::None
5857 // buffers — but always loud and counted, never silent.
5858 if !self.cpu_access.covers(access) {
5859 note_unplanned_cpu_access(
5860 self.buffer_identity().id(),
5861 &format!("{:?}", self.storage.memory()),
5862 "map access exceeds the declared CpuAccess",
5863 );
5864 }
5865 // CPU mapping of a strided tensor exposes the full padded buffer
5866 // (`row_stride × rows`) so callers can iterate rows via
5867 // `effective_row_stride()` without running past the slice. This is sound
5868 // only when the HAL owns and can size-check the allocation:
5869 //
5870 // * Self-allocated Mem / Shm tensors (any platform) — the backing
5871 // `Vec` / shm segment is sized by `capacity_bytes()`, checked here.
5872 // * Self-allocated DMA tensors (Linux) — pitch padding from
5873 // `image_with_stride()`; checked against the DMA-BUF `buf_size`.
5874 //
5875 // * Self-allocated PBO tensors (any platform with GL) — the GL buffer
5876 // is sized by `capacity_bytes()` and may carry 64-byte row padding;
5877 // the JPEG decoder mmaps it and convert() reads it, both iterating
5878 // by `row_stride`. Checked against the PBO capacity below.
5879 //
5880 // A strided **imported** DMA-BUF is accepted here too (Task 10b): the
5881 // `stride × rows <= buf_size` check just below is exactly what a
5882 // self-allocated Dma tensor relies on for its own soundness, and
5883 // `buf_size` for an imported fd is `fstat`-derived when the kernel
5884 // reports a usable size (`stat.st_size > 0` and at least the
5885 // declared shape's logical size; `DmaTensor::from_fd`) -- kernel
5886 // truth, not a caller's claim, in that case. The rare fallback (a
5887 // kernel that reports `st_size <= 0` for a DMA-BUF fd) uses the
5888 // declared logical size instead; that is not a hole here either, since
5889 // an out-of-range mmap against the real dma-buf is still rejected by
5890 // the kernel regardless of what this crate believed `buf_size` was.
5891 // So the check is equally load-bearing regardless of who allocated
5892 // the buffer. What
5893 // changed is only *which* callers can now reach this path: the
5894 // cross-package capsule protocol reconstructs an imported alias of a
5895 // HAL-owned Dma tensor and then calls `configure_image`, which
5896 // computes its stride with the exact same alignment rule a
5897 // self-allocated `Tensor::image()` used — so the value being
5898 // bounds-checked here is HAL-computed and format/width-derived, not
5899 // read from an external driver. Rejecting it unconditionally used to
5900 // make every foreign-imported GPU-backed destination fail any CPU
5901 // write path (the JPEG decoder's MCU writer among them) even when the
5902 // write was provably in-bounds. `set_row_stride()` was already
5903 // "trust the caller, bounds-check the map" for every *other* backing
5904 // (Mem/Shm/self-allocated Dma/PBO) — singling out imported Dma out
5905 // for a stricter standard was the inconsistency, not the trust model
5906 // itself. IOSurface keeps its own always-allowed strided path
5907 // unconditionally (see its arm below): a live `IOSurfaceRef` reports
5908 // its real `bytesPerRow` from the OS regardless of who allocated it,
5909 // so there was never anything to distrust there.
5910 if let Some(stride) = self.row_stride {
5911 return map_strided_host_view(self, stride, access, non_blocking);
5912 }
5913 // A non-zero plane offset needs no check here: every storage applies
5914 // its own offset inside its map -- Mem/Shm adjust the slice base, the
5915 // Linux DMA-BUF and Android AHardwareBuffer the mapped base, IOSurface
5916 // and D3D11 textures add their `view_offset`, and PBO starts the
5917 // staged copy at it.
5918 if non_blocking {
5919 self.storage.try_map_with(access)
5920 } else {
5921 self.storage.map_with(access)
5922 }
5923 }
5924}
5925
5926#[cfg(all(
5927 feature = "static",
5928 any(
5929 target_os = "macos",
5930 target_os = "ios",
5931 target_os = "android",
5932 target_os = "windows"
5933 )
5934))]
5935impl<T> Tensor<T>
5936where
5937 T: Num + Clone + fmt::Debug + Send + Sync,
5938{
5939 /// The pitched-backing half of [`TensorTrait::set_logical_shape`]: record
5940 /// the backing's own pitch as the row stride of the new `shape`.
5941 fn adopt_backing_pitch(&mut self, shape: &[usize]) {
5942 // Only a backing with a pitch of its own, and only while the shape
5943 // still has the rank the format's row is measured on: a caller may
5944 // flatten an image to one dimension and leave the format behind.
5945 let Some(pitch) = self.storage.backing_row_stride() else {
5946 return;
5947 };
5948 let Some((rank, rows)) = self.format.map(|f| match f.layout() {
5949 PixelLayout::Packed => (3, shape.first().copied()),
5950 PixelLayout::Planar => (3, shape.get(1).copied()),
5951 PixelLayout::SemiPlanar => (2, shape.first().copied()),
5952 }) else {
5953 return;
5954 };
5955 if rank != shape.len() {
5956 return;
5957 }
5958 // The natural stride of the new shape, not of the stride recorded for
5959 // the old one.
5960 let prior = self.row_stride.take();
5961 let Some(natural) = self.effective_row_stride() else {
5962 self.row_stride = prior;
5963 return;
5964 };
5965 if rows == Some(1) {
5966 self.set_row_stride_unchecked(natural);
5967 } else if pitch > natural {
5968 self.set_row_stride_unchecked(pitch);
5969 } else {
5970 self.row_stride = prior;
5971 }
5972 }
5973}
5974
5975#[cfg(feature = "static")]
5976impl<T> TensorTrait<T> for Tensor<T>
5977where
5978 T: Num + Clone + fmt::Debug + Send + Sync,
5979{
5980 fn new(shape: &[usize], name: Option<&str>) -> Result<Self>
5981 where
5982 Self: Sized,
5983 {
5984 Self::new(shape, None, name)
5985 }
5986
5987 #[cfg(unix)]
5988 fn from_fd(fd: std::os::fd::OwnedFd, shape: &[usize], name: Option<&str>) -> Result<Self>
5989 where
5990 Self: Sized,
5991 {
5992 #[cfg_attr(not(target_os = "linux"), allow(unused_mut))]
5993 let mut t = Self::wrap(TensorStorage::from_fd(fd, shape, name)?);
5994 // Best-effort CUDA external memory import for DMA-backed tensors.
5995 // RUNTIME-UNVALIDATED: see try_init_dma_cuda().
5996 #[cfg(target_os = "linux")]
5997 t.try_init_dma_cuda();
5998 Ok(t)
5999 }
6000
6001 #[cfg(unix)]
6002 fn clone_fd(&self) -> Result<std::os::fd::OwnedFd> {
6003 self.storage.clone_fd()
6004 }
6005
6006 fn memory(&self) -> TensorMemory {
6007 self.storage.memory()
6008 }
6009
6010 fn name(&self) -> String {
6011 self.storage.name()
6012 }
6013
6014 fn shape(&self) -> &[usize] {
6015 self.storage.shape()
6016 }
6017
6018 /// Bytes available in the underlying allocation, from this tensor's
6019 /// window start -- may exceed [`TensorTrait::size`] when the backing
6020 /// carries spare room beyond the current logical shape (a pitch-aligned
6021 /// image, or a pool tensor sized for a larger image than it holds
6022 /// today). See [`crate::protocol::TensorDesc::capacity`], the
6023 /// cross-package descriptor field this backs.
6024 fn capacity_bytes(&self) -> usize {
6025 self.storage.capacity_bytes()
6026 }
6027
6028 fn reshape(&mut self, shape: &[usize]) -> Result<()> {
6029 if self.chroma.is_some() {
6030 return Err(Error::InvalidOperation(
6031 "cannot reshape a multiplane tensor — decompose planes first".into(),
6032 ));
6033 }
6034 self.storage.reshape(shape)?;
6035 self.format = None;
6036 self.row_stride = None;
6037 self.plane_offset = None;
6038 match self.storage {
6039 TensorStorage::Mem(ref mut m) => m.set_offset(0),
6040 #[cfg(target_os = "linux")]
6041 TensorStorage::Dma(ref mut dma) => dma.mmap_offset = 0,
6042 // No-op: `IoSurfaceTensor::reshape` (iosurface.rs) zeroes its
6043 // own `view_offset` before this match runs.
6044 #[cfg(any(target_os = "macos", target_os = "ios"))]
6045 TensorStorage::Dma(_) => {}
6046 // `D3d11TextureTensor::reshape` does not self-zero, so the clear
6047 // that pairs with `set_plane_offset`'s D3D11 arm lives here.
6048 #[cfg(target_os = "windows")]
6049 TensorStorage::Dma(ref mut tex) => tex.set_view_offset(0),
6050 // No-op: `AHardwareBufferTensor::reshape` (ahardwarebuffer.rs)
6051 // zeroes its own `view_offset` before this match runs, the same
6052 // as `IoSurfaceTensor`. Unlike `set_plane_offset`'s Android arm,
6053 // this one is a genuine self-zeroing no-op, not a still-
6054 // incomplete gap: `view_offset` never survives a reshape here
6055 // regardless of whether `set_plane_offset` can restore it.
6056 #[cfg(target_os = "android")]
6057 TensorStorage::Dma(_) => {}
6058 // No-op: `PboTensor::reshape` (pbo.rs) zeroes its own
6059 // `view_offset` before this match runs, the same as
6060 // `IoSurfaceTensor`/`AHardwareBufferTensor` above -- which is
6061 // why, unlike `set_format`'s clear block, this match never
6062 // needed a `Pbo` arm that writes anything.
6063 TensorStorage::Pbo(_) => {}
6064 // Unlike `PboTensor::reshape`, `ShmTensor::reshape` (shm.rs)
6065 // only revalidates the element count and swaps `self.shape` --
6066 // it never touches `self.offset` -- so this clear site needs
6067 // the arm `set_format`'s does. `cfg(unix)` mirrors the `Shm`
6068 // variant's own gate, not a per-OS multiplex.
6069 #[cfg(unix)]
6070 TensorStorage::Shm(ref mut s) => s.set_offset(0),
6071 }
6072 Ok(())
6073 }
6074
6075 /// Set the logical shape to any shape whose bytes fit the allocation,
6076 /// without [`reshape`](TensorTrait::reshape)'s equal-count constraint.
6077 ///
6078 /// Forwards to the storage's capacity-aware implementation -- the same
6079 /// one [`Tensor::configure_image`] uses for pool reuse, a decode into an
6080 /// oversized reusable buffer reconfiguring it to a smaller image without
6081 /// reallocating -- so the two agree. The pixel format is kept.
6082 ///
6083 /// A narrower image inside a pitched backing (an IOSurface, an
6084 /// AHardwareBuffer, a D3D11 texture) still advances rows by the backing's
6085 /// pitch, so that pitch is recorded as the row stride whenever it exceeds
6086 /// the new shape's natural stride -- the same adoption `configure_image`
6087 /// makes. Without it a texture narrowed from 64 to 16 texels wide reports
6088 /// a 64-byte stride over rows that are 256 bytes apart. A single-row
6089 /// shape records its own tight stride instead, by [`view`](Self::view)'s
6090 /// rule: a strided `map()` spans `stride * rows`, and one pitched row
6091 /// runs past the backing when the window sits in the last row.
6092 fn set_logical_shape(&mut self, shape: &[usize]) -> Result<()> {
6093 self.storage.set_logical_shape(shape)?;
6094 #[cfg(any(
6095 target_os = "macos",
6096 target_os = "ios",
6097 target_os = "android",
6098 target_os = "windows"
6099 ))]
6100 self.adopt_backing_pitch(shape);
6101 Ok(())
6102 }
6103
6104 fn map_with<'a>(&self, access: CpuAccess) -> Result<crate::view::HostView<'a, T>>
6105 where
6106 T: 'a,
6107 {
6108 self.map_impl(access, false)
6109 }
6110
6111 fn try_map_with<'a>(&self, access: CpuAccess) -> Result<crate::view::HostView<'a, T>>
6112 where
6113 T: 'a,
6114 {
6115 self.map_impl(access, true)
6116 }
6117
6118 fn buffer_identity(&self) -> &BufferIdentity {
6119 self.storage.buffer_identity()
6120 }
6121}
6122
6123// ============================================================================
6124// Platform availability helpers
6125// ============================================================================
6126
6127/// Cached result of the Linux DMA-BUF availability probe.
6128#[cfg(target_os = "linux")]
6129static DMA_AVAILABLE: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
6130/// Cached result of the macOS/iOS IOSurface availability probe.
6131#[cfg(any(target_os = "macos", target_os = "ios"))]
6132static IOSURFACE_AVAILABLE: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
6133
6134/// Check if Linux DMA-BUF allocation is available on this system.
6135///
6136/// Linux-specific availability check (typically requires `/dev/dma_heap`
6137/// access — running as root or membership in a video/render group). For
6138/// portable code that wants "any zero-copy GPU buffer", use
6139/// [`is_gpu_buffer_available`] which also covers IOSurface on macOS.
6140///
6141/// This function caches its result after the first call.
6142#[cfg(target_os = "linux")]
6143pub fn is_dma_available() -> bool {
6144 *DMA_AVAILABLE
6145 .get_or_init(|| Tensor::<u8>::new(&[64], Some(TensorMemory::DmaBuf), None).is_ok())
6146}
6147
6148/// Always returns `false` on non-Linux platforms.
6149#[cfg(not(target_os = "linux"))]
6150pub fn is_dma_available() -> bool {
6151 false
6152}
6153
6154/// Check if macOS/iOS IOSurface allocation is available on this system.
6155///
6156/// IOSurface is part of the macOS/iOS OS and is essentially always present;
6157/// this probe catches degraded scenarios such as memory pressure or
6158/// sandboxed contexts where `IOSurfaceCreate` fails. The result is
6159/// cached after the first call.
6160///
6161/// Always returns `false` on non-Apple platforms.
6162#[cfg(any(target_os = "macos", target_os = "ios"))]
6163pub fn is_iosurface_available() -> bool {
6164 *IOSURFACE_AVAILABLE.get_or_init(|| {
6165 // Probe via the same Dma path — on macOS/iOS this routes through
6166 // IoSurfaceTensor::new.
6167 Tensor::<u8>::new(&[64], Some(TensorMemory::DmaBuf), None).is_ok()
6168 })
6169}
6170
6171#[cfg(not(any(target_os = "macos", target_os = "ios")))]
6172pub fn is_iosurface_available() -> bool {
6173 false
6174}
6175
6176/// Cached result of the Android AHardwareBuffer availability probe.
6177#[cfg(target_os = "android")]
6178static AHARDWAREBUFFER_AVAILABLE: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
6179
6180/// Check if Android AHardwareBuffer allocation is available on this system.
6181///
6182/// AHardwareBuffer is part of the Android OS (public NDK ABI since API
6183/// 26) and is essentially always present; this probe catches degraded
6184/// scenarios such as memory pressure or gralloc failures. The result is
6185/// cached after the first call.
6186#[cfg(target_os = "android")]
6187pub fn is_ahardwarebuffer_available() -> bool {
6188 *AHARDWAREBUFFER_AVAILABLE.get_or_init(|| {
6189 // Probe via the same Dma path — on Android this routes through
6190 // AHardwareBufferTensor::new.
6191 Tensor::<u8>::new(&[64], Some(TensorMemory::DmaBuf), None).is_ok()
6192 })
6193}
6194
6195/// Always returns `false` on non-Android platforms.
6196#[cfg(not(target_os = "android"))]
6197pub fn is_ahardwarebuffer_available() -> bool {
6198 false
6199}
6200
6201/// Portable probe for the platform's native zero-copy GPU buffer allocator
6202/// (DMA-BUF on Linux, IOSurface on macOS/iOS, AHardwareBuffer on Android,
6203/// D3D11 textures on Windows). Returns `false` on platforms with no
6204/// equivalent. Use this when writing cross-platform code that cares whether
6205/// the `Dma` tensor variant will work, not which underlying mechanism is used.
6206///
6207/// On Windows the `Dma` variant is image-formatted: `true` here means
6208/// [`Tensor::image`] can serve `TensorMemory::DmaBuf`, not [`Tensor::new`].
6209///
6210/// # Ordering on Windows
6211///
6212/// This probe **creates the process D3D11 device**, because the only honest
6213/// answer to "can this host allocate one" is to try. A host that must supply
6214/// its own device calls `d3d11::use_external_device` **before** this probe or
6215/// any allocation: the device is cached on first creation and installing one
6216/// afterwards is refused.
6217pub fn is_gpu_buffer_available() -> bool {
6218 #[cfg(target_os = "linux")]
6219 {
6220 is_dma_available()
6221 }
6222 #[cfg(any(target_os = "macos", target_os = "ios"))]
6223 {
6224 is_iosurface_available()
6225 }
6226 #[cfg(target_os = "android")]
6227 {
6228 is_ahardwarebuffer_available()
6229 }
6230 // No cached probe: `d3d11::device()` is itself a `OnceLock` over the
6231 // process device, so asking twice costs one atomic load.
6232 #[cfg(target_os = "windows")]
6233 {
6234 crate::d3d11::device().is_ok()
6235 }
6236 #[cfg(not(any(
6237 target_os = "linux",
6238 target_os = "macos",
6239 target_os = "ios",
6240 target_os = "android",
6241 target_os = "windows"
6242 )))]
6243 {
6244 false
6245 }
6246}
6247
6248/// Check if POSIX shared memory allocation is available on this system.
6249///
6250/// Returns `true` on Unix systems (Linux, macOS, BSD) where POSIX shared memory
6251/// is supported. Always returns `false` on non-Unix platforms (Windows).
6252///
6253/// This function caches its result after the first call for efficiency.
6254#[cfg(unix)]
6255static SHM_AVAILABLE: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
6256
6257/// Check if POSIX shared memory allocation is available on this system.
6258#[cfg(unix)]
6259pub fn is_shm_available() -> bool {
6260 *SHM_AVAILABLE.get_or_init(|| Tensor::<u8>::new(&[64], Some(TensorMemory::Shm), None).is_ok())
6261}
6262
6263/// Check if POSIX shared memory allocation is available on this system.
6264///
6265/// Always returns `false` on non-Unix platforms since POSIX SHM is Unix-specific.
6266#[cfg(not(unix))]
6267pub fn is_shm_available() -> bool {
6268 false
6269}
6270
6271#[cfg(test)]
6272mod dtype_tests {
6273 use super::*;
6274
6275 #[test]
6276 fn dtype_size() {
6277 assert_eq!(DType::U8.size(), 1);
6278 assert_eq!(DType::I8.size(), 1);
6279 assert_eq!(DType::U16.size(), 2);
6280 assert_eq!(DType::I16.size(), 2);
6281 assert_eq!(DType::U32.size(), 4);
6282 assert_eq!(DType::I32.size(), 4);
6283 assert_eq!(DType::U64.size(), 8);
6284 assert_eq!(DType::I64.size(), 8);
6285 assert_eq!(DType::F16.size(), 2);
6286 assert_eq!(DType::F32.size(), 4);
6287 assert_eq!(DType::F64.size(), 8);
6288 }
6289
6290 #[test]
6291 fn dtype_name() {
6292 assert_eq!(DType::U8.name(), "u8");
6293 assert_eq!(DType::F16.name(), "f16");
6294 assert_eq!(DType::F32.name(), "f32");
6295 }
6296
6297 #[test]
6298 fn dtype_serde_roundtrip() {
6299 use serde_json;
6300 let dt = DType::F16;
6301 let json = serde_json::to_string(&dt).unwrap();
6302 let back: DType = serde_json::from_str(&json).unwrap();
6303 assert_eq!(dt, back);
6304 }
6305}
6306
6307#[cfg(test)]
6308mod image_tests {
6309 use super::*;
6310
6311 #[test]
6312 fn image_dims_from_shape_inverts_allocation_shape() {
6313 for (format, w, h) in [
6314 (PixelFormat::Rgba, 640usize, 480usize),
6315 (PixelFormat::Rgb, 641, 3),
6316 (PixelFormat::Grey, 32, 32),
6317 (PixelFormat::PlanarRgb, 640, 480),
6318 (PixelFormat::PlanarRgba, 8, 5),
6319 (PixelFormat::Nv12, 640, 481),
6320 (PixelFormat::Nv16, 33, 17),
6321 (PixelFormat::Nv24, 16, 9),
6322 ] {
6323 let shape = format.allocation_shape(w, h).unwrap();
6324 assert_eq!(
6325 image_dims_from_shape(format, &shape).unwrap(),
6326 (w, h),
6327 "{format:?} {w}x{h} via {shape:?}"
6328 );
6329 }
6330 }
6331
6332 #[test]
6333 fn image_dims_from_shape_rejects_the_wrong_rank() {
6334 // Packed and planar want rank 3, semi-planar rank 2.
6335 for (format, shape) in [
6336 (PixelFormat::Rgba, vec![480usize, 640usize]),
6337 (PixelFormat::PlanarRgb, vec![3, 480]),
6338 (PixelFormat::Nv12, vec![720, 640, 1]),
6339 ] {
6340 let err = image_dims_from_shape(format, &shape).unwrap_err();
6341 assert!(
6342 matches!(err, Error::InvalidShape(ref m) if m.contains("rank")),
6343 "{format:?} {shape:?}: {err}"
6344 );
6345 }
6346 }
6347
6348 #[test]
6349 fn image_dims_from_shape_rejects_an_unproduced_combined_height() {
6350 // NV12 maps 480 -> 720 and 481 -> 722; nothing maps to 721, so the
6351 // shape is rejected rather than rounded to 480.
6352 let err = image_dims_from_shape(PixelFormat::Nv12, &[721, 640]).unwrap_err();
6353 assert!(
6354 matches!(err, Error::InvalidShape(ref m) if m.contains("combined plane height")),
6355 "{err}"
6356 );
6357 }
6358
6359 #[test]
6360 fn allocation_shape_per_layout() {
6361 assert_eq!(
6362 PixelFormat::Rgb.allocation_shape(640, 480),
6363 Some(vec![480, 640, 3])
6364 );
6365 assert_eq!(
6366 PixelFormat::Grey.allocation_shape(640, 480),
6367 Some(vec![480, 640, 1])
6368 );
6369 assert_eq!(
6370 PixelFormat::Nv12.allocation_shape(640, 480),
6371 Some(vec![720, 640])
6372 );
6373 // Odd height: combined-plane height is `481 + ceil(481/2)` = 481 + 241
6374 // = 722 rows. Logical height is recovered as `722 * 2 / 3` = 481.
6375 assert_eq!(
6376 PixelFormat::Nv12.allocation_shape(640, 481),
6377 Some(vec![722, 640])
6378 );
6379 // Odd width: shape carries the LOGICAL width (641).
6380 // The 64-aligned stride (>= 642) is stored separately on the Tensor.
6381 assert_eq!(
6382 PixelFormat::Nv12.allocation_shape(641, 480),
6383 Some(vec![720, 641])
6384 );
6385 // NV16 odd width: same — logical width in shape, stride separate.
6386 assert_eq!(
6387 PixelFormat::Nv16.allocation_shape(641, 480),
6388 Some(vec![960, 641])
6389 );
6390 assert_eq!(
6391 PixelFormat::PlanarRgb.allocation_shape(640, 480),
6392 Some(vec![3, 480, 640])
6393 );
6394 assert_eq!(
6395 PixelFormat::Nv16.allocation_shape(640, 480),
6396 Some(vec![960, 640])
6397 );
6398 }
6399
6400 #[test]
6401 fn raw_tensor_has_no_format() {
6402 let t = Tensor::<u8>::new(&[480, 640, 3], None, None).unwrap();
6403 assert!(t.format().is_none());
6404 assert!(t.width().is_none());
6405 assert!(t.height().is_none());
6406 assert!(!t.is_multiplane());
6407 assert!(t.chroma().is_none());
6408 }
6409
6410 #[test]
6411 fn image_tensor_packed() {
6412 let t = Tensor::<u8>::image(
6413 640,
6414 480,
6415 PixelFormat::Rgba,
6416 None,
6417 crate::CpuAccess::ReadWrite,
6418 )
6419 .unwrap();
6420 assert_eq!(t.format(), Some(PixelFormat::Rgba));
6421 assert_eq!(t.width(), Some(640));
6422 assert_eq!(t.height(), Some(480));
6423 assert_eq!(t.shape(), &[480, 640, 4]);
6424 assert!(!t.is_multiplane());
6425 }
6426
6427 #[test]
6428 fn image_tensor_planar() {
6429 let t = Tensor::<u8>::image(
6430 640,
6431 480,
6432 PixelFormat::PlanarRgb,
6433 None,
6434 crate::CpuAccess::ReadWrite,
6435 )
6436 .unwrap();
6437 assert_eq!(t.format(), Some(PixelFormat::PlanarRgb));
6438 assert_eq!(t.width(), Some(640));
6439 assert_eq!(t.height(), Some(480));
6440 assert_eq!(t.shape(), &[3, 480, 640]);
6441 }
6442
6443 #[test]
6444 #[cfg(target_os = "macos")]
6445 fn image_tensor_dma_non_aligned_packed_width_pads_zero_copy() {
6446 // Declares this test as an fd-opener; see FD_LOCK.
6447 let _lock = crate::tests::fd_lock_shared();
6448 // RGBA u8 at width=4 → 4*4 = 16 bytes/row, not 64-byte aligned. RGBA has
6449 // a real IOSurface FourCC, so an explicit `Some(TensorMemory::DmaBuf)`
6450 // request now allocates a padded image IOSurface (64-aligned
6451 // `bytes_per_row`) and records the stride — a fully zero-copy buffer GL
6452 // can bind and the CPU can map via the strided path. (Previously this
6453 // failed loudly to avoid an 'L008' byte-bag downgrade; with a real
6454 // FourCC surface that concern no longer applies.)
6455 let t = Tensor::<u8>::image(
6456 4,
6457 4,
6458 PixelFormat::Rgba,
6459 Some(TensorMemory::DmaBuf),
6460 crate::CpuAccess::ReadWrite,
6461 )
6462 .expect("padded RGBA IOSurface should allocate");
6463 assert_eq!(t.format(), Some(PixelFormat::Rgba));
6464 assert_eq!(t.width(), Some(4));
6465 assert_eq!(t.height(), Some(4));
6466 let stride = t.effective_row_stride().expect("stride");
6467 assert_eq!(stride % 64, 0, "padded to 64-byte row alignment");
6468 assert!(stride >= 16);
6469 // A CPU map exposes the full padded surface for strided iteration.
6470 let m = t.map().expect("strided IOSurface map");
6471 assert_eq!(m.as_slice().len(), stride * 4);
6472 }
6473
6474 /// `per_pixel_bytes` that doesn't divide 64 evenly (e.g. RGB u8 with
6475 /// 3 B/pixel) makes a "Pad width to N" suggestion structurally
6476 /// impossible — there is no integer width whose `width * 3` is a
6477 /// multiple of 64. The error must still fire (no silent SHM
6478 /// fallback for explicit-DMA requests) and must spell out the
6479 /// alignment requirement; it just omits the misleading "pad to N"
6480 /// hint instead of printing a number whose row pitch still won't
6481 /// align.
6482 #[test]
6483 #[cfg(target_os = "macos")]
6484 fn image_tensor_dma_rejects_indivisible_pixel_pitch_without_pad_hint() {
6485 // Declares this test as an fd-opener; see FD_LOCK.
6486 let _lock = crate::tests::fd_lock_shared();
6487 // Width=10 RGB f32 → 120 B/row, not 64-byte aligned, and (Rgb,
6488 // F32) has no IOSurface mapping so the padded-stride tolerance
6489 // does not apply. The next 64-multiple (128 B) isn't an integer
6490 // multiple of 12 B/pixel, so the "pad width to N" hint can't
6491 // produce a valid number and must be omitted. (Rgb u8 used to be
6492 // this test's subject but now has a real RGBA8888 mapping with
6493 // padded-stride tolerance — see the test below.)
6494 let err = Tensor::<f32>::image(
6495 10,
6496 10,
6497 PixelFormat::Rgb,
6498 Some(TensorMemory::DmaBuf),
6499 crate::CpuAccess::ReadWrite,
6500 )
6501 .expect_err("RGB f32 with 12 B/pixel and non-aligned width must be rejected");
6502 match err {
6503 Error::InvalidArgument(msg) => {
6504 assert!(
6505 msg.contains("64-byte aligned"),
6506 "error must still name the alignment requirement: {msg}"
6507 );
6508 assert!(
6509 !msg.contains("Pad width"),
6510 "indivisible per-pixel pitch makes a width suggestion impossible; \
6511 hint must be omitted, got: {msg}"
6512 );
6513 assert!(
6514 msg.contains("memory=None") && msg.contains("TensorMemory::Mem"),
6515 "error must still list the always-applicable alternatives: {msg}"
6516 );
6517 }
6518 other => panic!("expected InvalidArgument, got {other:?}"),
6519 }
6520 }
6521
6522 #[test]
6523 #[cfg(target_os = "macos")]
6524 fn image_tensor_dma_packed_rgb_u8_contract() {
6525 // Declares this test as an fd-opener; see FD_LOCK.
6526 let _lock = crate::tests::fd_lock_shared();
6527 // Packed RGB u8 @Dma is a designed RGBA8888 mapping at
6528 // (W*3/4, H) — the INT8 NPU input layout, shared with Android.
6529 // width%4 != 0 cannot form whole texels → loud InvalidArgument…
6530 let err = Tensor::<u8>::image(
6531 10,
6532 10,
6533 PixelFormat::Rgb,
6534 Some(TensorMemory::DmaBuf),
6535 crate::CpuAccess::ReadWrite,
6536 )
6537 .expect_err("Rgb u8 width%4!=0 must be rejected");
6538 assert!(
6539 matches!(&err, Error::InvalidArgument(m) if m.contains("width%4==0")),
6540 "got {err:?}"
6541 );
6542 // …width%4 == 0 with a non-64-aligned pitch allocates PADDED
6543 // (36 B rows → 64 B surface pitch, recorded on the tensor)…
6544 let t = Tensor::<u8>::image(
6545 12,
6546 4,
6547 PixelFormat::Rgb,
6548 Some(TensorMemory::DmaBuf),
6549 crate::CpuAccess::ReadWrite,
6550 )
6551 .expect("width 12 Rgb u8 must allocate padded");
6552 assert_eq!(t.memory(), TensorMemory::DmaBuf);
6553 assert!(
6554 t.row_stride()
6555 .is_some_and(|s| s >= 64 && s.is_multiple_of(64)),
6556 "padded pitch must be recorded: {:?}",
6557 t.row_stride()
6558 );
6559 // …and the aligned model-input width stays flat (640*3 = 1920 is
6560 // 64-aligned → no recorded stride, the buffer IS [H, W, 3]).
6561 let t = Tensor::<u8>::image(
6562 640,
6563 8,
6564 PixelFormat::Rgb,
6565 Some(TensorMemory::DmaBuf),
6566 crate::CpuAccess::ReadWrite,
6567 )
6568 .expect("width 640 Rgb u8 must allocate flat");
6569 assert_eq!(t.row_stride(), None);
6570 // I8 shares the layout (INT8 shader bias, not a format change).
6571 let t = Tensor::<i8>::image(
6572 640,
6573 8,
6574 PixelFormat::Rgb,
6575 Some(TensorMemory::DmaBuf),
6576 crate::CpuAccess::ReadWrite,
6577 )
6578 .expect("Rgb i8 shares the RGBA8888 mapping");
6579 assert_eq!(t.memory(), TensorMemory::DmaBuf);
6580 }
6581
6582 #[test]
6583 #[cfg(target_os = "macos")]
6584 fn image_tensor_dma_planar_f16_alignment() {
6585 // Declares this test as an fd-opener; see FD_LOCK.
6586 let _lock = crate::tests::fd_lock_shared();
6587 // PlanarRgb F16 uses single-channel row pitch (width * 2 bytes).
6588 // Width=16 → 32 bytes/row (not aligned); width=32 → 64 bytes/row (aligned).
6589 let err = Tensor::<half::f16>::image(
6590 16,
6591 16,
6592 PixelFormat::PlanarRgb,
6593 Some(TensorMemory::DmaBuf),
6594 crate::CpuAccess::ReadWrite,
6595 )
6596 .expect_err("width=16 PlanarRgb F16 is 32-byte row, must reject");
6597 assert!(matches!(err, Error::InvalidArgument(_)), "got {err:?}");
6598 // 32 wide should work.
6599 let t = Tensor::<half::f16>::image(
6600 32,
6601 8,
6602 PixelFormat::PlanarRgb,
6603 Some(TensorMemory::DmaBuf),
6604 crate::CpuAccess::ReadWrite,
6605 )
6606 .expect("width=32 PlanarRgb F16 is 64-byte row, must succeed");
6607 assert_eq!(t.format(), Some(PixelFormat::PlanarRgb));
6608 }
6609
6610 #[test]
6611 fn image_tensor_semi_planar_contiguous() {
6612 let t = Tensor::<u8>::image(
6613 640,
6614 480,
6615 PixelFormat::Nv12,
6616 None,
6617 crate::CpuAccess::ReadWrite,
6618 )
6619 .unwrap();
6620 assert_eq!(t.format(), Some(PixelFormat::Nv12));
6621 assert_eq!(t.width(), Some(640));
6622 assert_eq!(t.height(), Some(480));
6623 // NV12: H*3/2 = 720
6624 assert_eq!(t.shape(), &[720, 640]);
6625 assert!(!t.is_multiplane());
6626 }
6627
6628 #[test]
6629 #[cfg(target_os = "linux")]
6630 fn image_tensor_with_stride_preserves_logical_width() {
6631 // Declares this test as an fd-opener; see FD_LOCK.
6632 let _lock = crate::tests::fd_lock_shared();
6633 if !crate::test_support::dma_or_skip("image_tensor_with_stride_preserves_logical_width") {
6634 return;
6635 }
6636 // 3004×1688 RGBA8: natural pitch 12016, padded to 12032 (64-aligned).
6637 let stride = 12032;
6638 let t = Tensor::<u8>::image_with_stride(
6639 3004,
6640 1688,
6641 PixelFormat::Rgba,
6642 stride,
6643 Some(TensorMemory::DmaBuf),
6644 crate::CpuAccess::ReadWrite,
6645 )
6646 .unwrap();
6647 // Logical dimensions unchanged by padding — this is the contract.
6648 assert_eq!(t.width(), Some(3004));
6649 assert_eq!(t.height(), Some(1688));
6650 assert_eq!(t.shape(), &[1688, 3004, 4]);
6651 // Stride is carried separately and reports the padded pitch.
6652 assert_eq!(t.effective_row_stride(), Some(stride));
6653 // Buffer is sized to stride × height so the full padded layout fits,
6654 // and CPU map() works for self-allocated strided DMA tensors.
6655 use crate::TensorMapTrait;
6656 {
6657 let map = t.map().unwrap();
6658 assert!(
6659 map.as_slice().len() >= stride * 1688,
6660 "mapped buffer {} bytes < expected {}",
6661 map.as_slice().len(),
6662 stride * 1688
6663 );
6664 }
6665 // CPU write access works too — iterate rows using the padded stride,
6666 // touch only the active `width × bpp` region, verify it round-trips.
6667 {
6668 let mut map = t.map().unwrap();
6669 let slice = map.as_mut_slice();
6670 for y in 0..1688 {
6671 let row_start = y * stride;
6672 for x in 0..3004 {
6673 let p = row_start + x * 4;
6674 slice[p] = (y & 0xFF) as u8;
6675 slice[p + 1] = (x & 0xFF) as u8;
6676 slice[p + 2] = 0x42;
6677 slice[p + 3] = 0xFF;
6678 }
6679 }
6680 }
6681 {
6682 let map = t.map().unwrap();
6683 let slice = map.as_slice();
6684 // Sample a few pixels to confirm the round-trip.
6685 assert_eq!(slice[0], 0x00);
6686 assert_eq!(slice[1], 0x00);
6687 assert_eq!(slice[2], 0x42);
6688 assert_eq!(slice[3], 0xFF);
6689 let mid = 100 * stride + 50 * 4;
6690 assert_eq!(slice[mid], 100);
6691 assert_eq!(slice[mid + 1], 50);
6692 assert_eq!(slice[mid + 2], 0x42);
6693 }
6694 }
6695
6696 #[test]
6697 #[cfg(target_os = "linux")]
6698 fn image_tensor_foreign_strided_map_is_bounds_checked_not_refused() {
6699 // An imported (foreign) DMA-BUF with a row stride used to be refused
6700 // outright. That changed deliberately: the cross-package capsule
6701 // protocol reconstructs an imported alias of a HAL-owned Dma tensor,
6702 // and refusing every strided CPU map broke the flagship
6703 // create_image() -> decode_file_into() -> convert() pipeline on Linux.
6704 //
6705 // What replaced the refusal is a bounds check against `buf_size`,
6706 // which for an imported fd is `fstat`-derived kernel truth rather
6707 // than the caller's claim. So the property worth testing is no longer
6708 // "foreign strided maps fail" but "they succeed exactly when they
6709 // fit" -- which is what actually keeps the mapping sound.
6710 //
6711 // This test previously asserted the old refusal and went red on every
6712 // board with a dma-heap when the behaviour changed, unnoticed because
6713 // it needs both Linux and a dma-heap: it skips on macOS and on Orin
6714 // (nvmap, no CONFIG_DMABUF_HEAPS).
6715 if !crate::test_support::dma_or_skip(
6716 "image_tensor_foreign_strided_map_is_bounds_checked_not_refused",
6717 ) {
6718 return;
6719 }
6720 let backing =
6721 Tensor::<u8>::new(&[240 * 320 * 4], Some(TensorMemory::DmaBuf), None).unwrap();
6722
6723 // In bounds: the natural stride for the declared image fits the
6724 // allocation, so the map must succeed.
6725 {
6726 let fd = backing.clone_fd().unwrap();
6727 let shape = [240usize, 320, 4];
6728 let storage = TensorStorage::<u8>::from_fd(fd, &shape, None).unwrap();
6729 let mut t = Tensor::<u8>::wrap(storage);
6730 t.set_format(PixelFormat::Bgra).unwrap();
6731 t.set_row_stride(320 * 4).unwrap();
6732 assert!(
6733 t.map().is_ok(),
6734 "an in-bounds strided map of an imported DMA-BUF must be \
6735 accepted -- refusing it is what broke decode-into-GPU-tensor"
6736 );
6737 }
6738
6739 // Out of bounds: a stride the allocation cannot back must still be
6740 // refused, and the error must say so. This is the half that carries
6741 // the soundness.
6742 {
6743 let fd = backing.clone_fd().unwrap();
6744 let shape = [240usize, 320, 4];
6745 let storage = TensorStorage::<u8>::from_fd(fd, &shape, None).unwrap();
6746 let mut t = Tensor::<u8>::wrap(storage);
6747 t.set_format(PixelFormat::Bgra).unwrap();
6748 // Double the pitch: 240 rows x 2560 bytes is twice the allocation.
6749 t.set_row_stride(320 * 4 * 2).unwrap();
6750 let err = t.map();
6751 assert!(
6752 matches!(err, Err(Error::InvalidOperation(_))),
6753 "a strided map larger than the DMA allocation must be refused"
6754 );
6755 }
6756 }
6757
6758 #[test]
6759 #[cfg(target_os = "linux")]
6760 fn image_tensor_with_stride_map_rejects_tampered_stride() {
6761 // Declares this test as an fd-opener; see FD_LOCK.
6762 let _lock = crate::tests::fd_lock_shared();
6763 // Round-3 PR feedback (C1): `set_row_stride` is public and only
6764 // validates `stride >= min_stride`, not that the new stride × height
6765 // fits the underlying buffer. A caller that tampers with the stride
6766 // after allocation must not be able to coerce `Tensor::map()` into
6767 // returning a slice larger than the backing mmap (that would be UB
6768 // in `DmaMap::as_slice`).
6769 if !crate::test_support::dma_or_skip("image_tensor_with_stride_map_rejects_tampered_stride")
6770 {
6771 return;
6772 }
6773 // Allocate a 640×480 RGBA8 padded canvas (stride = 3072 = 768 px).
6774 // Backing buffer is 3072 × 480 = 1,474,560 bytes.
6775 let mut t = Tensor::<u8>::image_with_stride(
6776 640,
6777 480,
6778 PixelFormat::Rgba,
6779 3072,
6780 Some(TensorMemory::DmaBuf),
6781 crate::CpuAccess::ReadWrite,
6782 )
6783 .unwrap();
6784 // Tamper: push the stride up to 4 × the original. This is >=
6785 // min_stride (2560), so `set_row_stride` accepts it.
6786 t.set_row_stride(12288).unwrap();
6787 // Map must now refuse — 12288 × 480 = 5,898,240 > 1,474,560.
6788 let err = t.map();
6789 assert!(
6790 matches!(err, Err(Error::InvalidOperation(_))),
6791 "map() with oversized stride must return InvalidOperation"
6792 );
6793 }
6794
6795 #[test]
6796 fn dma_tensor_new_with_byte_size_rejects_shape_overflow() {
6797 // Round-3 PR feedback (C3): shape.product() * sizeof(T) must use
6798 // checked arithmetic so a pathological shape can't wrap usize and
6799 // make the byte_size-vs-logical-size comparison incorrect.
6800 //
6801 // This test only exercises the overflow rejection path, which is
6802 // pure-Rust and doesn't touch dma_heap — safe to run on any target.
6803 #[cfg(target_os = "linux")]
6804 {
6805 let err = crate::dma::DmaTensor::<u64>::new_with_byte_size(
6806 &[usize::MAX, 2, 2],
6807 usize::MAX,
6808 None,
6809 );
6810 assert!(
6811 matches!(err, Err(Error::InvalidArgument(_))),
6812 "new_with_byte_size must detect shape.product() overflow"
6813 );
6814 }
6815 }
6816
6817 #[test]
6818 #[cfg(target_os = "linux")]
6819 fn image_tensor_with_stride_rejects_too_small_stride() {
6820 // Declares this test as an fd-opener; see FD_LOCK.
6821 let _lock = crate::tests::fd_lock_shared();
6822 // 640×480 RGBA8 natural pitch = 2560, request 2400 → should error.
6823 let err = Tensor::<u8>::image_with_stride(
6824 640,
6825 480,
6826 PixelFormat::Rgba,
6827 2400,
6828 Some(TensorMemory::DmaBuf),
6829 crate::CpuAccess::ReadWrite,
6830 );
6831 assert!(matches!(err, Err(Error::InvalidArgument(_))));
6832 }
6833
6834 #[test]
6835 #[cfg(target_os = "linux")]
6836 fn image_tensor_with_stride_rejects_non_packed() {
6837 // Declares this test as an fd-opener; see FD_LOCK.
6838 let _lock = crate::tests::fd_lock_shared();
6839 // NV12 is SemiPlanar → not supported. (Linux-only because
6840 // `TensorMemory::DmaBuf` itself is a Linux-only enum variant.)
6841 let err = Tensor::<u8>::image_with_stride(
6842 640,
6843 480,
6844 PixelFormat::Nv12,
6845 640,
6846 Some(TensorMemory::DmaBuf),
6847 crate::CpuAccess::ReadWrite,
6848 );
6849 assert!(matches!(err, Err(Error::NotImplemented(_))));
6850 }
6851
6852 #[test]
6853 fn set_format_valid() {
6854 let mut t = Tensor::<u8>::new(&[480, 640, 3], None, None).unwrap();
6855 assert!(t.format().is_none());
6856 t.set_format(PixelFormat::Rgb).unwrap();
6857 assert_eq!(t.format(), Some(PixelFormat::Rgb));
6858 assert_eq!(t.width(), Some(640));
6859 assert_eq!(t.height(), Some(480));
6860 }
6861
6862 #[test]
6863 fn set_format_invalid_shape() {
6864 let mut t = Tensor::<u8>::new(&[480, 640, 4], None, None).unwrap();
6865 // RGB expects 3 channels, not 4
6866 let err = t.set_format(PixelFormat::Rgb);
6867 assert!(err.is_err());
6868 // Original tensor is unmodified
6869 assert!(t.format().is_none());
6870 }
6871
6872 #[test]
6873 fn reshape_clears_format() {
6874 let mut t = Tensor::<u8>::image(
6875 640,
6876 480,
6877 PixelFormat::Rgba,
6878 None,
6879 crate::CpuAccess::ReadWrite,
6880 )
6881 .unwrap();
6882 assert_eq!(t.format(), Some(PixelFormat::Rgba));
6883 // Reshape to flat — format cleared
6884 t.reshape(&[480 * 640 * 4]).unwrap();
6885 assert!(t.format().is_none());
6886 }
6887
6888 #[test]
6889 fn from_planes_nv12() {
6890 let y = Tensor::<u8>::new(&[480, 640], None, None).unwrap();
6891 let uv = Tensor::<u8>::new(&[240, 640], None, None).unwrap();
6892 let img = Tensor::from_planes(y, uv, PixelFormat::Nv12).unwrap();
6893 assert_eq!(img.format(), Some(PixelFormat::Nv12));
6894 assert!(img.is_multiplane());
6895 assert!(img.chroma().is_some());
6896 assert_eq!(img.width(), Some(640));
6897 assert_eq!(img.height(), Some(480));
6898 }
6899
6900 #[test]
6901 fn from_planes_rejects_non_semiplanar() {
6902 let y = Tensor::<u8>::new(&[480, 640], None, None).unwrap();
6903 let uv = Tensor::<u8>::new(&[240, 640], None, None).unwrap();
6904 let err = Tensor::from_planes(y, uv, PixelFormat::Rgb);
6905 assert!(err.is_err());
6906 }
6907
6908 #[test]
6909 fn reshape_multiplane_errors() {
6910 let y = Tensor::<u8>::new(&[480, 640], None, None).unwrap();
6911 let uv = Tensor::<u8>::new(&[240, 640], None, None).unwrap();
6912 let mut img = Tensor::from_planes(y, uv, PixelFormat::Nv12).unwrap();
6913 let err = img.reshape(&[480 * 640 + 240 * 640]);
6914 assert!(err.is_err());
6915 }
6916}
6917
6918#[cfg(test)]
6919mod compression_tests {
6920 use super::*;
6921
6922 #[test]
6923 fn desc_builder_roundtrips() {
6924 let desc = ImageDesc::new(640, 480, PixelFormat::Rgba, DType::U8)
6925 .with_memory(Some(TensorMemory::Mem))
6926 .with_access(CpuAccess::Read)
6927 .with_compression(Compression::Any);
6928 assert_eq!(desc.width(), 640);
6929 assert_eq!(desc.height(), 480);
6930 assert_eq!(desc.format(), PixelFormat::Rgba);
6931 assert_eq!(desc.dtype(), DType::U8);
6932 assert_eq!(desc.memory(), Some(TensorMemory::Mem));
6933 assert_eq!(desc.access(), CpuAccess::Read);
6934 assert_eq!(desc.compression(), Some(Compression::Any));
6935
6936 // Defaults: auto memory, hardware-only, no request.
6937 let plain = ImageDesc::new(2, 2, PixelFormat::Grey, DType::U8);
6938 assert_eq!(plain.memory(), None);
6939 assert_eq!(plain.access(), CpuAccess::None);
6940 assert_eq!(plain.compression(), None);
6941 }
6942
6943 #[test]
6944 fn desc_dtype_must_match_element_type() {
6945 let desc = ImageDesc::new(4, 4, PixelFormat::Rgba, DType::F32);
6946 match Tensor::<u8>::image_desc(&desc) {
6947 Err(Error::InvalidArgument(msg)) => assert!(msg.contains("dtype")),
6948 other => panic!("expected InvalidArgument, got {other:?}"),
6949 }
6950 }
6951
6952 #[test]
6953 fn compression_with_cpu_access_is_invalid() {
6954 let desc = ImageDesc::new(4, 4, PixelFormat::Rgba, DType::U8)
6955 .with_access(CpuAccess::ReadWrite)
6956 .with_compression(Compression::Any);
6957 match Tensor::<u8>::image_desc(&desc) {
6958 Err(Error::InvalidArgument(msg)) => assert!(msg.contains("CpuAccess::None")),
6959 other => panic!("expected InvalidArgument, got {other:?}"),
6960 }
6961 }
6962
6963 #[cfg(not(target_os = "android"))]
6964 #[test]
6965 fn scheme_request_off_android_is_not_implemented() {
6966 let desc = ImageDesc::new(4, 4, PixelFormat::Rgba, DType::U8)
6967 .with_compression(Compression::Scheme(CompressionScheme::Ubwc));
6968 match Tensor::<u8>::image_desc(&desc) {
6969 Err(Error::NotImplemented(msg)) => assert!(msg.contains("Ubwc")),
6970 other => panic!("expected NotImplemented, got {other:?}"),
6971 }
6972 }
6973
6974 #[cfg(not(target_os = "android"))]
6975 #[test]
6976 fn any_request_off_android_resolves_linear_and_counts() {
6977 let before = compression_fallback_count();
6978 let desc = ImageDesc::new(64, 64, PixelFormat::Rgba, DType::U8)
6979 .with_memory(Some(TensorMemory::Mem))
6980 .with_compression(Compression::Any);
6981 let t = Tensor::<u8>::image_desc(&desc).unwrap();
6982 assert_eq!(t.compression(), None);
6983 assert!(compression_fallback_count() > before);
6984 }
6985
6986 #[test]
6987 fn desc_without_request_matches_classic_constructor() {
6988 let desc = ImageDesc::new(32, 32, PixelFormat::Rgba, DType::U8)
6989 .with_memory(Some(TensorMemory::Mem))
6990 .with_access(CpuAccess::ReadWrite);
6991 let t = Tensor::<u8>::image_desc(&desc).unwrap();
6992 assert_eq!(t.compression(), None);
6993 assert_eq!(t.cpu_access(), CpuAccess::ReadWrite);
6994 assert_eq!(t.width(), Some(32));
6995 // Mappable exactly like the classic constructor's result.
6996 let m = t.map_read().unwrap();
6997 assert_eq!(m.as_slice().len(), 32 * 32 * 4);
6998 }
6999
7000 #[test]
7001 fn configure_image_preserves_compression_and_views_inherit() {
7002 // No host platform records a scheme, so emulate the recording to
7003 // pin the preserve/inherit semantics (the physical layout does
7004 // not change when the logical image is reconfigured).
7005 let mut t = Tensor::<u8>::image(
7006 64,
7007 64,
7008 PixelFormat::Rgba,
7009 Some(TensorMemory::Mem),
7010 CpuAccess::None,
7011 )
7012 .unwrap();
7013 t.compression = Some(CompressionScheme::Ubwc);
7014 t.configure_image(32, 32, PixelFormat::Rgba).unwrap();
7015 assert_eq!(t.compression(), Some(CompressionScheme::Ubwc));
7016 let view = t.subview(0, &[16, 32, 4]).unwrap();
7017 assert_eq!(view.compression(), Some(CompressionScheme::Ubwc));
7018 }
7019
7020 #[test]
7021 fn tensor_dyn_dispatches_desc_and_compression() {
7022 let desc = ImageDesc::new(16, 16, PixelFormat::Rgba, DType::U8)
7023 .with_memory(Some(TensorMemory::Mem))
7024 .with_access(CpuAccess::ReadWrite);
7025 let t = TensorDyn::image_desc(&desc).unwrap();
7026 assert_eq!(t.compression(), None);
7027 assert!(matches!(t, TensorDyn::U8(_)));
7028 }
7029}
7030
7031#[cfg(all(test, feature = "static"))]
7032mod vocabulary_and_probe_tests {
7033 use super::*;
7034
7035 #[test]
7036 fn dtype_displays_its_wire_name() {
7037 for (dt, name) in [
7038 (DType::U8, "u8"),
7039 (DType::I8, "i8"),
7040 (DType::U16, "u16"),
7041 (DType::I16, "i16"),
7042 (DType::U32, "u32"),
7043 (DType::I32, "i32"),
7044 (DType::U64, "u64"),
7045 (DType::I64, "i64"),
7046 (DType::F16, "f16"),
7047 (DType::F32, "f32"),
7048 (DType::F64, "f64"),
7049 ] {
7050 assert_eq!(dt.to_string(), name);
7051 }
7052 }
7053
7054 #[test]
7055 fn dtype_of_maps_every_element_type_and_nothing_else() {
7056 assert_eq!(dtype_of::<u8>(), Some(DType::U8));
7057 assert_eq!(dtype_of::<i8>(), Some(DType::I8));
7058 assert_eq!(dtype_of::<u16>(), Some(DType::U16));
7059 assert_eq!(dtype_of::<i16>(), Some(DType::I16));
7060 assert_eq!(dtype_of::<u32>(), Some(DType::U32));
7061 assert_eq!(dtype_of::<i32>(), Some(DType::I32));
7062 assert_eq!(dtype_of::<u64>(), Some(DType::U64));
7063 assert_eq!(dtype_of::<i64>(), Some(DType::I64));
7064 assert_eq!(dtype_of::<half::f16>(), Some(DType::F16));
7065 assert_eq!(dtype_of::<f32>(), Some(DType::F32));
7066 assert_eq!(dtype_of::<f64>(), Some(DType::F64));
7067 assert_eq!(dtype_of::<bool>(), None);
7068 }
7069
7070 #[test]
7071 fn tensor_memory_converts_to_and_from_its_wire_string() {
7072 assert_eq!(String::from(TensorMemory::DmaBuf), "dmabuf");
7073 assert_eq!(String::from(TensorMemory::Shm), "shm");
7074 assert_eq!(
7075 TensorMemory::try_from("dmabuf").unwrap(),
7076 TensorMemory::DmaBuf
7077 );
7078 assert_eq!(
7079 TensorMemory::try_from("dma").unwrap(),
7080 TensorMemory::DmaBuf,
7081 "the legacy spelling still parses"
7082 );
7083 assert_eq!(TensorMemory::try_from("mem").unwrap(), TensorMemory::Mem);
7084 match TensorMemory::try_from("bogus") {
7085 Err(Error::InvalidMemoryType(s)) => assert_eq!(s, "bogus"),
7086 other => panic!("expected InvalidMemoryType, got {other:?}"),
7087 }
7088 }
7089
7090 #[test]
7091 fn tensor_memory_is_available_matches_a_real_allocation() {
7092 let _lock = crate::tests::fd_lock_shared();
7093 assert!(TensorMemory::Mem.is_available());
7094 assert!(!TensorMemory::IoSurface.is_available());
7095 assert!(!TensorMemory::Pbo.is_available());
7096 assert!(!TensorMemory::Cuda.is_available());
7097 #[cfg(unix)]
7098 assert_eq!(
7099 TensorMemory::Shm.is_available(),
7100 Tensor::<u8>::new(&[64], Some(TensorMemory::Shm), None).is_ok()
7101 );
7102 #[cfg(not(target_os = "windows"))]
7103 assert_eq!(
7104 TensorMemory::DmaBuf.is_available(),
7105 Tensor::<u8>::new(&[64], Some(TensorMemory::DmaBuf), None).is_ok()
7106 );
7107 }
7108
7109 #[test]
7110 #[cfg(target_os = "linux")]
7111 fn linux_probes_report_what_linux_can_allocate() {
7112 let _lock = crate::tests::fd_lock_shared();
7113 let dma = DmaTensor::<u8>::new(&[64], None).is_ok();
7114 assert_eq!(is_dma_available(), dma);
7115 assert_eq!(is_gpu_buffer_available(), dma);
7116 assert!(!is_iosurface_available());
7117 assert!(!is_ahardwarebuffer_available());
7118 // Linux (unlike Android's bionic) always has POSIX shared memory.
7119 assert!(is_shm_available());
7120 }
7121}
7122
7123#[cfg(all(test, feature = "static"))]
7124mod quantization_serde_tests {
7125 use super::*;
7126
7127 fn parse(json: &str) -> Quantization {
7128 serde_json::from_str(json).unwrap_or_else(|e| panic!("{json}: {e}"))
7129 }
7130
7131 #[test]
7132 fn scale_accepts_every_json_number_form() {
7133 assert_eq!(parse(r#"{"scale": 0.5}"#).scale(), &[0.5]);
7134 assert_eq!(parse(r#"{"scale": 2}"#).scale(), &[2.0]);
7135 assert_eq!(parse(r#"{"scale": -2}"#).scale(), &[-2.0]);
7136 assert_eq!(parse(r#"{"scale": [0.25, 4]}"#).scale(), &[0.25, 4.0]);
7137 }
7138
7139 #[test]
7140 fn zero_point_accepts_null_scalar_and_array() {
7141 assert_eq!(
7142 parse(r#"{"scale": 0.1, "zero_point": -5}"#).zero_point(),
7143 Some(&[-5][..])
7144 );
7145 assert_eq!(
7146 parse(r#"{"scale": 0.1, "zero_point": 3}"#).zero_point(),
7147 Some(&[3][..])
7148 );
7149 assert_eq!(
7150 parse(r#"{"scale": [0.1, 0.2], "zero_point": [1, -2], "axis": 0}"#).zero_point(),
7151 Some(&[1, -2][..])
7152 );
7153 assert_eq!(
7154 parse(r#"{"scale": 0.1, "zero_point": null}"#).zero_point(),
7155 None
7156 );
7157 assert_eq!(parse(r#"{"scale": 0.1}"#).zero_point(), None);
7158 }
7159
7160 /// Deserializers whose `deserialize_option` forwards to
7161 /// `deserialize_any` (serde's value deserializers among them) reach the
7162 /// zero-point visitor's scalar and sequence arms directly, not through
7163 /// `visit_some`.
7164 #[test]
7165 fn zero_point_visitor_accepts_values_without_an_option_wrapper() {
7166 use serde::de::value::{
7167 Error as ValueError, I64Deserializer, SeqDeserializer, U64Deserializer,
7168 };
7169 assert_eq!(
7170 deserialize_opt_scalar_or_vec_i32(I64Deserializer::<ValueError>::new(-7)).unwrap(),
7171 Some(vec![-7])
7172 );
7173 assert_eq!(
7174 deserialize_opt_scalar_or_vec_i32(U64Deserializer::<ValueError>::new(9)).unwrap(),
7175 Some(vec![9])
7176 );
7177 let seq = SeqDeserializer::<_, ValueError>::new(vec![4i32, -4].into_iter());
7178 assert_eq!(
7179 deserialize_opt_scalar_or_vec_i32(seq).unwrap(),
7180 Some(vec![4, -4])
7181 );
7182 }
7183
7184 #[test]
7185 fn per_channel_is_not_per_tensor() {
7186 let q = Quantization::per_channel(vec![0.1, 0.2], vec![0, 1], 0).unwrap();
7187 assert!(!q.is_per_tensor());
7188 assert!(Quantization::per_tensor(0.1, 0).is_per_tensor());
7189 }
7190
7191 #[test]
7192 fn validate_accepts_matching_per_channel_and_names_the_failing_field() {
7193 let mut t = Tensor::<i8>::new(&[1, 1, 4], Some(TensorMemory::Mem), None).unwrap();
7194 let q = Quantization::per_channel(vec![0.1, 0.2, 0.3, 0.4], vec![0, 1, 2, 3], 2).unwrap();
7195 t.set_quantization(q)
7196 .expect("four channels on axis 2 of [1, 1, 4]");
7197 assert_eq!(
7198 t.quantization().and_then(Quantization::zero_point),
7199 Some(&[0, 1, 2, 3][..])
7200 );
7201
7202 let q = Quantization::per_channel(vec![0.1, 0.2, 0.3, 0.4], vec![0, 1, 2, 3], 3).unwrap();
7203 match t.set_quantization(q) {
7204 Err(Error::QuantizationInvalid { field, got, .. }) => {
7205 assert_eq!(field, "axis");
7206 assert_eq!(got, "axis=3");
7207 }
7208 other => panic!("axis 3 of a rank-3 tensor must be refused, got {other:?}"),
7209 }
7210
7211 let q = Quantization::per_channel_symmetric(vec![0.1, 0.2, 0.3, 0.4], 1).unwrap();
7212 match t.set_quantization(q) {
7213 Err(Error::QuantizationInvalid { field, got, .. }) => {
7214 assert_eq!(field, "scale.len");
7215 assert_eq!(got, "length 4");
7216 }
7217 other => panic!("four scales on a length-1 axis must be refused, got {other:?}"),
7218 }
7219 }
7220}
7221
7222#[cfg(all(test, feature = "static"))]
7223mod image_geometry_tests {
7224 use super::*;
7225
7226 fn mem_image<T>(width: usize, height: usize, format: PixelFormat) -> Tensor<T>
7227 where
7228 T: Num + Clone + fmt::Debug + Send + Sync + 'static,
7229 {
7230 Tensor::<T>::image(
7231 width,
7232 height,
7233 format,
7234 Some(TensorMemory::Mem),
7235 CpuAccess::ReadWrite,
7236 )
7237 .unwrap_or_else(|e| panic!("Mem {format:?} {width}x{height}: {e}"))
7238 }
7239
7240 fn raw_mem<T>(shape: &[usize]) -> Tensor<T>
7241 where
7242 T: Num + Clone + fmt::Debug + Send + Sync + 'static,
7243 {
7244 Tensor::<T>::new(shape, Some(TensorMemory::Mem), None).unwrap()
7245 }
7246
7247 #[test]
7248 fn semi_planar_image_stride_scales_with_the_element_size() {
7249 // u16 NV12 at width 100: 200 B rows, aligned to 256.
7250 let t = mem_image::<u16>(100, 64, PixelFormat::Nv12);
7251 assert_eq!(t.row_stride(), Some(256));
7252 assert_eq!(t.capacity_bytes(), 256 * 96);
7253 }
7254
7255 #[test]
7256 fn host_planar_image_is_tight_at_every_element_size() {
7257 let t = mem_image::<u8>(320, 240, PixelFormat::PlanarRgb);
7258 assert_eq!(t.row_stride(), None);
7259 assert_eq!(t.effective_row_stride(), Some(320));
7260 assert_eq!(t.capacity_bytes(), 320 * 240 * 3);
7261
7262 let t = mem_image::<f32>(100, 40, PixelFormat::PlanarRgb);
7263 assert_eq!(t.row_stride(), None);
7264 assert_eq!(t.effective_row_stride(), Some(400));
7265 assert_eq!(t.capacity_bytes(), 400 * 120);
7266 }
7267
7268 #[test]
7269 #[cfg(target_os = "linux")]
7270 fn dma_packed_image_records_the_aligned_pitch_only_when_padded() {
7271 let _lock = crate::tests::fd_lock_shared();
7272 if !crate::test_support::dma_or_skip(
7273 "dma_packed_image_records_the_aligned_pitch_only_when_padded",
7274 ) {
7275 return;
7276 }
7277 let t = Tensor::<u8>::image(
7278 100,
7279 10,
7280 PixelFormat::Rgba,
7281 Some(TensorMemory::DmaBuf),
7282 CpuAccess::ReadWrite,
7283 )
7284 .unwrap();
7285 assert_eq!(t.row_stride(), Some(448), "400 B rows pad to 448");
7286 let t = Tensor::<u8>::image(
7287 64,
7288 10,
7289 PixelFormat::Rgba,
7290 Some(TensorMemory::DmaBuf),
7291 CpuAccess::ReadWrite,
7292 )
7293 .unwrap();
7294 assert_eq!(t.row_stride(), None, "256 B rows are already aligned");
7295 }
7296
7297 #[test]
7298 #[cfg(target_os = "linux")]
7299 fn image_with_stride_validates_the_stride_before_the_memory() {
7300 match Tensor::<u8>::image_with_stride(
7301 640,
7302 480,
7303 PixelFormat::Rgba,
7304 2559,
7305 Some(TensorMemory::Mem),
7306 CpuAccess::ReadWrite,
7307 ) {
7308 Err(Error::InvalidArgument(msg)) => {
7309 assert!(msg.contains("2559 < minimum 2560"), "{msg}")
7310 }
7311 other => panic!("expected InvalidArgument, got {other:?}"),
7312 }
7313 match Tensor::<u8>::image_with_stride(
7314 640,
7315 480,
7316 PixelFormat::Rgba,
7317 2560,
7318 Some(TensorMemory::Mem),
7319 CpuAccess::ReadWrite,
7320 ) {
7321 Err(Error::NotImplemented(msg)) => assert!(msg.contains("only TensorMemory::DmaBuf")),
7322 other => panic!("a valid stride on Mem must reach the memory check, got {other:?}"),
7323 }
7324 }
7325
7326 #[test]
7327 fn set_format_checks_the_planar_and_semi_planar_shape() {
7328 let mut t = raw_mem::<u8>(&[480, 640, 3]);
7329 match t.set_format(PixelFormat::PlanarRgb) {
7330 Err(Error::InvalidShape(msg)) => assert!(msg.contains("[3, H, W]"), "{msg}"),
7331 other => panic!("[480, 640, 3] is not planar, got {other:?}"),
7332 }
7333 assert!(raw_mem::<u8>(&[3, 480, 640])
7334 .set_format(PixelFormat::PlanarRgb)
7335 .is_ok());
7336
7337 for (format, bad, good) in [
7338 (PixelFormat::Nv12, 721, [720usize, 722]),
7339 (PixelFormat::Nv16, 481, [480, 482]),
7340 (PixelFormat::Nv24, 1441, [1440, 1443]),
7341 ] {
7342 match raw_mem::<u8>(&[bad, 640]).set_format(format) {
7343 Err(Error::InvalidShape(msg)) => {
7344 assert!(msg.contains(&bad.to_string()), "{format:?}: {msg}")
7345 }
7346 other => panic!("{format:?} [{bad}, 640] must be refused, got {other:?}"),
7347 }
7348 for rows in good {
7349 raw_mem::<u8>(&[rows, 640])
7350 .set_format(format)
7351 .unwrap_or_else(|e| panic!("{format:?} [{rows}, 640]: {e}"));
7352 }
7353 }
7354 }
7355
7356 #[test]
7357 fn contiguous_semi_planar_height_divides_the_combined_rows() {
7358 assert_eq!(
7359 mem_image::<u8>(320, 240, PixelFormat::Nv24).height(),
7360 Some(240)
7361 );
7362 assert_eq!(
7363 mem_image::<u8>(320, 240, PixelFormat::Nv16).height(),
7364 Some(240)
7365 );
7366 assert_eq!(
7367 mem_image::<u8>(320, 241, PixelFormat::Nv12).height(),
7368 Some(241)
7369 );
7370 }
7371
7372 #[test]
7373 fn effective_row_stride_scales_with_the_element_size() {
7374 assert_eq!(
7375 mem_image::<f32>(10, 4, PixelFormat::Rgb).effective_row_stride(),
7376 Some(120)
7377 );
7378 assert_eq!(
7379 mem_image::<f32>(10, 4, PixelFormat::PlanarRgb).effective_row_stride(),
7380 Some(40)
7381 );
7382 let mut t = raw_mem::<u16>(&[360, 320]);
7383 t.set_format(PixelFormat::Nv12).unwrap();
7384 assert_eq!(t.effective_row_stride(), Some(640));
7385 let mut t = raw_mem::<u16>(&[360, 321]);
7386 t.set_format(PixelFormat::Nv12).unwrap();
7387 assert_eq!(
7388 t.effective_row_stride(),
7389 Some(644),
7390 "odd width rounds up to even"
7391 );
7392 }
7393
7394 #[test]
7395 fn set_row_stride_minimum_scales_with_the_element_size() {
7396 let refuse = |mut t: Tensor<f32>, stride: usize, min: usize| match t.set_row_stride(stride)
7397 {
7398 Err(Error::InvalidArgument(msg)) => {
7399 assert!(msg.contains(&format!("{stride} < minimum {min}")), "{msg}")
7400 }
7401 other => panic!("stride {stride} under {min} must be refused, got {other:?}"),
7402 };
7403 refuse(mem_image::<f32>(10, 4, PixelFormat::Rgb), 119, 120);
7404 refuse(mem_image::<f32>(10, 4, PixelFormat::PlanarRgb), 39, 40);
7405 let mut t = mem_image::<f32>(10, 4, PixelFormat::Rgb);
7406 t.set_row_stride(120).unwrap();
7407 assert_eq!(t.row_stride(), Some(120));
7408 let mut t = mem_image::<f32>(10, 4, PixelFormat::PlanarRgb);
7409 t.set_row_stride(40).unwrap();
7410 assert_eq!(t.row_stride(), Some(40));
7411
7412 let mut t = mem_image::<u16>(101, 64, PixelFormat::Nv12);
7413 match t.set_row_stride(203) {
7414 Err(Error::InvalidArgument(msg)) => assert!(msg.contains("203 < minimum 204"), "{msg}"),
7415 other => panic!("expected InvalidArgument, got {other:?}"),
7416 }
7417 t.set_row_stride(204).unwrap();
7418 }
7419
7420 #[test]
7421 fn configure_image_semi_planar_stride_scales_with_the_element_size() {
7422 let mut t = mem_image::<u16>(640, 480, PixelFormat::Rgba);
7423 t.configure_image(100, 64, PixelFormat::Nv12).unwrap();
7424 assert_eq!(t.row_stride(), Some(256));
7425 }
7426
7427 #[test]
7428 fn configure_image_keeps_host_packed_odd_widths_tight() {
7429 let mut t = mem_image::<u8>(640, 480, PixelFormat::Rgb);
7430 t.configure_image(321, 240, PixelFormat::Rgb).unwrap();
7431 assert_eq!(t.row_stride(), None);
7432 assert_eq!(t.effective_row_stride(), Some(963));
7433 }
7434
7435 #[test]
7436 fn configure_image_does_not_inherit_a_stride_under_the_same_format() {
7437 // Host packed: a stride set for one geometry must not survive into
7438 // the next, where it is below the new minimum row.
7439 let mut t = mem_image::<u8>(640, 480, PixelFormat::Rgb);
7440 t.configure_image(100, 100, PixelFormat::Rgb).unwrap();
7441 t.set_row_stride(400).unwrap();
7442 t.configure_image(300, 100, PixelFormat::Rgb).unwrap();
7443 assert_eq!(t.row_stride(), None);
7444 assert_eq!(t.effective_row_stride(), Some(900));
7445
7446 // Semi-planar whose aligned stride does not fit: the tight row does,
7447 // and must be what the tensor reports -- not the old 64.
7448 let mut t = mem_image::<u8>(64, 1000, PixelFormat::Nv12);
7449 assert_eq!(t.row_stride(), Some(64));
7450 t.configure_image(200, 300, PixelFormat::Nv12).unwrap();
7451 assert_eq!(t.row_stride(), None);
7452 assert_eq!(t.effective_row_stride(), Some(200));
7453 }
7454
7455 #[test]
7456 fn configure_image_pool_reuse_picks_the_stride_that_fits() {
7457 // A prior stride narrower than the new row is replaced by the
7458 // aligned one.
7459 let mut t = mem_image::<u8>(64, 1024, PixelFormat::Nv12);
7460 t.configure_image(200, 100, PixelFormat::Nv12).unwrap();
7461 assert_eq!(t.row_stride(), Some(256));
7462
7463 // An unaligned prior stride is replaced too.
7464 let mut t = mem_image::<u8>(64, 64, PixelFormat::Nv12);
7465 t.set_row_stride(100).unwrap();
7466 t.configure_image(64, 64, PixelFormat::Nv12).unwrap();
7467 assert_eq!(t.row_stride(), Some(64));
7468
7469 // A wide prior stride whose rows would overrun the allocation gives
7470 // way to the aligned stride of the new width.
7471 let mut t = mem_image::<u8>(640, 480, PixelFormat::Nv12);
7472 t.configure_image(320, 900, PixelFormat::Nv12).unwrap();
7473 assert_eq!(t.row_stride(), Some(320));
7474
7475 // Neither the aligned stride fits nor is it needed: the tight row
7476 // exactly fills a raw allocation.
7477 let mut t = raw_mem::<u8>(&[150 * 100]);
7478 t.configure_image(100, 100, PixelFormat::Nv12).unwrap();
7479 assert_eq!(t.row_stride(), None);
7480 assert_eq!(t.effective_row_stride(), Some(100));
7481 }
7482
7483 #[test]
7484 fn configure_image_refuses_rows_that_overrun_the_allocation() {
7485 // Odd width 99: the shape [150, 99] fits 14850 bytes, but NV12 rows
7486 // are at least the even width, 100 * 150 = 15000 bytes.
7487 let mut t = raw_mem::<u8>(&[150 * 99]);
7488 match t.configure_image(99, 100, PixelFormat::Nv12) {
7489 Err(Error::InsufficientCapacity { needed, capacity }) => {
7490 assert_eq!((needed, capacity), (15000, 14850))
7491 }
7492 other => panic!("expected InsufficientCapacity, got {other:?}"),
7493 }
7494 }
7495
7496 #[test]
7497 #[cfg(target_os = "linux")]
7498 fn configure_image_on_dma_aligns_packed_and_planar_rows() {
7499 let _lock = crate::tests::fd_lock_shared();
7500 if !crate::test_support::dma_or_skip("configure_image_on_dma_aligns_packed_and_planar_rows")
7501 {
7502 return;
7503 }
7504 // Packed u16 Rgb at width 100: 600 B rows align to 640.
7505 let mut t = Tensor::<u16>::image(
7506 640,
7507 480,
7508 PixelFormat::Rgba,
7509 Some(TensorMemory::DmaBuf),
7510 CpuAccess::ReadWrite,
7511 )
7512 .unwrap();
7513 t.configure_image(100, 50, PixelFormat::Rgb).unwrap();
7514 assert_eq!(t.row_stride(), Some(640));
7515
7516 // Planar u16 at width 100: 200 B rows would align to 256, but 150
7517 // rows of 256 B overrun the 32 KiB buffer while 150 rows of 200 B fit.
7518 let mut t = Tensor::<u16>::new(&[16 * 1024], Some(TensorMemory::DmaBuf), None).unwrap();
7519 assert_eq!(t.capacity_bytes(), 32 * 1024);
7520 t.configure_image(100, 50, PixelFormat::PlanarRgb).unwrap();
7521 assert_eq!(t.row_stride(), None);
7522 assert_eq!(t.effective_row_stride(), Some(200));
7523 t.configure_image(100, 40, PixelFormat::PlanarRgb).unwrap();
7524 assert_eq!(t.row_stride(), Some(256), "120 rows of 256 B fit");
7525 }
7526
7527 #[test]
7528 fn from_planes_requires_a_raw_two_dimensional_chroma() {
7529 let formatted_chroma = mem_image::<u8>(640, 240, PixelFormat::Grey);
7530 match Tensor::from_planes(
7531 raw_mem::<u8>(&[480, 640]),
7532 formatted_chroma,
7533 PixelFormat::Nv12,
7534 ) {
7535 Err(Error::InvalidArgument(msg)) => assert!(msg.contains("raw tensor"), "{msg}"),
7536 other => panic!("a formatted chroma plane must be refused, got {other:?}"),
7537 }
7538 match Tensor::from_planes(
7539 raw_mem::<u8>(&[480, 640, 1]),
7540 raw_mem::<u8>(&[240, 640]),
7541 PixelFormat::Nv12,
7542 ) {
7543 Err(Error::InvalidArgument(msg)) => assert!(msg.contains("2D shapes"), "{msg}"),
7544 other => panic!("a 3-D luma plane must be refused, got {other:?}"),
7545 }
7546 }
7547
7548 #[test]
7549 fn from_planes_nv16_needs_full_height_chroma() {
7550 let img = Tensor::from_planes(
7551 raw_mem::<u8>(&[480, 640]),
7552 raw_mem::<u8>(&[480, 640]),
7553 PixelFormat::Nv16,
7554 )
7555 .expect("NV16 chroma is full height");
7556 assert_eq!(img.format(), Some(PixelFormat::Nv16));
7557 assert_eq!(img.height(), Some(480));
7558
7559 match Tensor::from_planes(
7560 raw_mem::<u8>(&[480, 640]),
7561 raw_mem::<u8>(&[240, 640]),
7562 PixelFormat::Nv16,
7563 ) {
7564 Err(Error::InvalidArgument(msg)) => {
7565 assert!(msg.contains("NV16 chroma height"), "{msg}")
7566 }
7567 other => panic!("half-height NV16 chroma must be refused, got {other:?}"),
7568 }
7569 }
7570
7571 #[test]
7572 fn chroma_mut_writes_the_chroma_plane() {
7573 let mut img = Tensor::from_planes(
7574 raw_mem::<u8>(&[4, 8]),
7575 raw_mem::<u8>(&[2, 8]),
7576 PixelFormat::Nv12,
7577 )
7578 .unwrap();
7579 let chroma = img.chroma_mut().expect("a multiplane image has chroma");
7580 chroma.map().unwrap().as_mut_slice().fill(0x80);
7581 let chroma = img.chroma().unwrap();
7582 assert!(chroma.map().unwrap().as_slice().iter().all(|&b| b == 0x80));
7583 }
7584
7585 #[test]
7586 fn view_offset_and_stride_scale_with_the_element_size() {
7587 let t = mem_image::<u8>(100, 80, PixelFormat::Rgba);
7588 let v = t.view(Region::new(10, 20, 30, 40)).unwrap();
7589 assert_eq!(v.plane_offset(), Some(20 * 400 + 10 * 4));
7590 assert_eq!(v.row_stride(), Some(400));
7591
7592 let t = mem_image::<u16>(100, 80, PixelFormat::Rgba);
7593 let v = t.view(Region::new(10, 20, 30, 40)).unwrap();
7594 assert_eq!(v.plane_offset(), Some(20 * 800 + 10 * 8));
7595 assert_eq!(v.row_stride(), Some(800));
7596 assert_eq!(v.view_origin().map(|o| o.parent_row_stride), Some(800));
7597 }
7598
7599 #[test]
7600 fn single_row_view_keeps_a_tight_stride() {
7601 let t = mem_image::<u8>(8, 4, PixelFormat::Rgba);
7602 let row = t.view(Region::new(2, 3, 4, 1)).unwrap();
7603 assert_eq!(row.effective_row_stride(), Some(16));
7604 assert_eq!(row.map().unwrap().as_slice().len(), 16);
7605 assert_eq!(row.view_origin().map(|o| o.parent_row_stride), Some(32));
7606
7607 let two_rows = t.view(Region::new(2, 2, 4, 2)).unwrap();
7608 assert_eq!(two_rows.effective_row_stride(), Some(32));
7609 }
7610
7611 #[test]
7612 fn view_refuses_a_region_past_either_edge() {
7613 let t = mem_image::<u8>(64, 40, PixelFormat::Rgba);
7614 for region in [Region::new(0, 30, 10, 20), Region::new(60, 0, 10, 1)] {
7615 match t.view(region) {
7616 Err(Error::RegionOutOfBounds { region: r, bounds }) => {
7617 assert_eq!((r, bounds), (region, (64, 40)))
7618 }
7619 other => panic!("{region:?} must be refused, got {other:?}"),
7620 }
7621 }
7622 assert!(Region::new(0, 0, 64, 40).fits_within(64, 40));
7623 assert!(!Region::new(1, 0, 64, 40).fits_within(64, 40));
7624 assert!(!Region::new(0, 1, 64, 40).fits_within(64, 40));
7625 }
7626}
7627
7628#[cfg(all(test, feature = "static"))]
7629mod strided_access_tests {
7630 use super::*;
7631
7632 /// A patterned byte view of `t`'s whole mapping: byte `i` is `i % 251`.
7633 fn fill_pattern<T>(t: &Tensor<T>)
7634 where
7635 T: Num + Clone + fmt::Debug + Send + Sync,
7636 {
7637 let mut m = t.map().unwrap();
7638 let s = m.as_mut_slice();
7639 // SAFETY: plain numeric elements; every byte pattern is a valid value
7640 // for the integer and float types these tests use.
7641 let bytes = unsafe {
7642 std::slice::from_raw_parts_mut(s.as_mut_ptr() as *mut u8, std::mem::size_of_val(s))
7643 };
7644 for (i, b) in bytes.iter_mut().enumerate() {
7645 *b = (i % 251) as u8;
7646 }
7647 }
7648
7649 fn expected_rows(rows: usize, stride: usize, row_bytes: usize) -> Vec<u8> {
7650 (0..rows)
7651 .flat_map(|r| (r * stride..r * stride + row_bytes).map(|i| (i % 251) as u8))
7652 .collect()
7653 }
7654
7655 fn padded<T>(format: PixelFormat, stride: usize) -> Tensor<T>
7656 where
7657 T: Num + Clone + fmt::Debug + Send + Sync + 'static,
7658 {
7659 let mut t = Tensor::<T>::image_with_capacity(
7660 16,
7661 4,
7662 format,
7663 Some(TensorMemory::Mem),
7664 CpuAccess::ReadWrite,
7665 )
7666 .unwrap();
7667 t.configure_image(8, 4, format).unwrap();
7668 t.set_row_stride(stride).unwrap();
7669 t
7670 }
7671
7672 #[test]
7673 fn copy_to_flat_compacts_padded_packed_rows() {
7674 // 8 px * 3 B = 24 B rows at a 32 B pitch.
7675 let t = padded::<u8>(PixelFormat::Rgb, 32);
7676 fill_pattern(&t);
7677 let mut dst = vec![0u8; 96];
7678 t.copy_to_flat(&mut dst).unwrap();
7679 assert_eq!(dst, expected_rows(4, 32, 24));
7680 }
7681
7682 #[test]
7683 fn copy_to_flat_compacts_padded_planar_rows() {
7684 // f32 planar: 3 planes x 4 rows of 32 B at a 64 B pitch.
7685 let t = padded::<f32>(PixelFormat::PlanarRgb, 64);
7686 fill_pattern(&t);
7687 let mut dst = vec![0u8; 384];
7688 t.copy_to_flat(&mut dst).unwrap();
7689 assert_eq!(dst, expected_rows(12, 64, 32));
7690 }
7691
7692 #[test]
7693 fn copy_to_flat_of_a_tight_stride_is_the_mapping() {
7694 let t = padded::<u8>(PixelFormat::Rgb, 24);
7695 fill_pattern(&t);
7696 let mut dst = vec![0u8; 96];
7697 t.copy_to_flat(&mut dst).unwrap();
7698 assert_eq!(dst, expected_rows(1, 96, 96));
7699
7700 let t = Tensor::<u8>::image(
7701 8,
7702 4,
7703 PixelFormat::Rgb,
7704 Some(TensorMemory::Mem),
7705 CpuAccess::ReadWrite,
7706 )
7707 .unwrap();
7708 fill_pattern(&t);
7709 let mut dst = vec![0u8; 96];
7710 t.copy_to_flat(&mut dst).unwrap();
7711 assert_eq!(dst, expected_rows(1, 96, 96));
7712 }
7713
7714 #[test]
7715 fn copy_to_flat_needs_a_tight_sized_destination() {
7716 let t = padded::<u8>(PixelFormat::Rgb, 32);
7717 let mut dst = vec![0u8; 95];
7718 match t.copy_to_flat(&mut dst) {
7719 Err(Error::InvalidArgument(msg)) => {
7720 assert!(
7721 msg.contains("dst is 95 bytes") && msg.contains("96 bytes"),
7722 "{msg}"
7723 )
7724 }
7725 other => panic!("expected InvalidArgument, got {other:?}"),
7726 }
7727 }
7728
7729 #[test]
7730 fn copy_to_flat_compacts_a_padded_batch() {
7731 // Two 8x2 RGB tiles at a 32 B pitch: four rows in all.
7732 let mut t = padded::<u8>(PixelFormat::Rgb, 32);
7733 t.set_logical_shape(&[2, 2, 8, 3]).unwrap();
7734 fill_pattern(&t);
7735 let mut dst = vec![0u8; 96];
7736 t.copy_to_flat(&mut dst).unwrap();
7737 assert_eq!(dst, expected_rows(4, 32, 24));
7738 }
7739
7740 #[test]
7741 fn strided_planar_map_spans_every_plane_row() {
7742 let t = padded::<u8>(PixelFormat::PlanarRgb, 16);
7743 assert_eq!(t.map().unwrap().as_slice().len(), 16 * 12);
7744 }
7745
7746 #[test]
7747 fn batch_of_a_padded_batch_starts_at_its_padded_tile() {
7748 // Tile 1 starts after two 32 B rows, not after 48 tight bytes.
7749 let mut t = padded::<u8>(PixelFormat::Rgb, 32);
7750 t.set_logical_shape(&[2, 2, 8, 3]).unwrap();
7751 fill_pattern(&t);
7752 let e = t.batch(1).unwrap();
7753 assert_eq!(e.plane_offset(), Some(64));
7754 let mut dst = vec![0u8; 48];
7755 e.copy_to_flat(&mut dst).unwrap();
7756 let tile1: Vec<u8> = expected_rows(4, 32, 24).split_off(48);
7757 assert_eq!(dst, tile1);
7758 }
7759
7760 #[test]
7761 fn strided_packed_batch_map_spans_every_tile_row() {
7762 let mut t = padded::<u8>(PixelFormat::Rgb, 32);
7763 t.set_logical_shape(&[2, 2, 8, 3]).unwrap();
7764 assert_eq!(t.map().unwrap().as_slice().len(), 32 * 4);
7765 }
7766
7767 #[test]
7768 fn strided_access_to_a_flattened_image_is_refused() {
7769 let mut t = padded::<u8>(PixelFormat::PlanarRgb, 16);
7770 t.set_logical_shape(&[96]).unwrap();
7771 match t.map() {
7772 Err(Error::InvalidOperation(msg)) => assert!(msg.contains("[96]"), "{msg}"),
7773 other => panic!("expected InvalidOperation, got {other:?}"),
7774 }
7775 let mut dst = vec![0u8; 96];
7776 match t.copy_to_flat(&mut dst) {
7777 Err(Error::InvalidOperation(msg)) => assert!(msg.contains("[96]"), "{msg}"),
7778 other => panic!("expected InvalidOperation, got {other:?}"),
7779 }
7780 }
7781
7782 #[test]
7783 #[cfg(target_os = "linux")]
7784 fn strided_shm_map_is_bounded_by_the_segment() {
7785 let _lock = crate::tests::fd_lock_shared();
7786 let mut t = Tensor::<u8>::image(
7787 640,
7788 480,
7789 PixelFormat::Nv12,
7790 Some(TensorMemory::Shm),
7791 CpuAccess::ReadWrite,
7792 )
7793 .expect("Linux has POSIX shared memory");
7794 assert_eq!(t.capacity_bytes(), 640 * 720);
7795 assert_eq!(t.map().unwrap().as_slice().len(), 640 * 720);
7796
7797 t.configure_image(320, 240, PixelFormat::Nv12).unwrap();
7798 assert_eq!(t.row_stride(), Some(640), "the pool pitch fits and is kept");
7799 assert_eq!(t.map().unwrap().as_slice().len(), 640 * 360);
7800
7801 t.configure_image(640, 480, PixelFormat::Nv12).unwrap();
7802 t.set_row_stride(704).unwrap();
7803 match t.map() {
7804 Err(Error::InsufficientCapacity { needed, capacity }) => {
7805 assert_eq!((needed, capacity), (704 * 720, 640 * 720))
7806 }
7807 other => panic!("expected InsufficientCapacity, got {other:?}"),
7808 }
7809 }
7810
7811 /// Minimal host-memory `PboOps`: a fixed buffer, never resized.
7812 struct HostPbo(Vec<u8>);
7813
7814 // SAFETY: the pointer addresses `self.0`, allocated once and never
7815 // resized, and the tensor holds the `Arc` to `self` for every map's life.
7816 unsafe impl PboOps for HostPbo {
7817 fn map_buffer(&self, _buffer_id: u32, size: usize) -> Result<PboMapping> {
7818 assert!(size <= self.0.len());
7819 Ok(PboMapping {
7820 ptr: self.0.as_ptr() as *mut u8,
7821 size,
7822 })
7823 }
7824 fn unmap_buffer(&self, _buffer_id: u32) -> Result<()> {
7825 Ok(())
7826 }
7827 fn delete_buffer(&self, _buffer_id: u32) {}
7828 }
7829
7830 fn pbo_rgba(width: usize, height: usize) -> Tensor<u8> {
7831 let bytes = width * height * 4;
7832 let pbo = PboTensor::<u8>::from_pbo(
7833 7,
7834 bytes,
7835 &[height, width, 4],
7836 None,
7837 Arc::new(HostPbo(vec![0; bytes])),
7838 )
7839 .unwrap();
7840 let mut t = Tensor::from_pbo(pbo).unwrap();
7841 t.set_format(PixelFormat::Rgba).unwrap();
7842 t
7843 }
7844
7845 #[test]
7846 fn strided_pbo_map_is_bounded_by_the_buffer() {
7847 let mut t = pbo_rgba(8, 4);
7848 t.set_row_stride(32).unwrap();
7849 assert_eq!(t.map().unwrap().as_slice().len(), 128);
7850
7851 t.set_row_stride(64).unwrap();
7852 match t.map() {
7853 Err(Error::InsufficientCapacity { needed, capacity }) => {
7854 assert_eq!((needed, capacity), (256, 128))
7855 }
7856 other => panic!("expected InsufficientCapacity, got {other:?}"),
7857 }
7858 }
7859
7860 #[test]
7861 fn pbo_accessors_answer_only_for_a_pbo_tensor() {
7862 let mem = Tensor::<u8>::new(&[4], Some(TensorMemory::Mem), None).unwrap();
7863 assert_eq!(mem.pbo_is_mapped(), None);
7864 assert!(mem.pbo_keepalive().is_none());
7865
7866 let t = pbo_rgba(8, 4);
7867 assert_eq!(t.pbo_is_mapped(), Some(false));
7868 assert!(t.pbo_keepalive().is_some());
7869 let m = t.map().unwrap();
7870 assert_eq!(t.pbo_is_mapped(), Some(true));
7871 drop(m);
7872 assert_eq!(t.pbo_is_mapped(), Some(false));
7873 }
7874}
7875
7876#[cfg(all(test, feature = "static"))]
7877mod accessor_tests {
7878 use super::*;
7879 use std::sync::atomic::{AtomicUsize, Ordering};
7880
7881 #[test]
7882 #[cfg(unix)]
7883 fn plane_descriptor_carries_its_stride_and_offset() {
7884 use std::os::fd::AsFd;
7885 let _lock = crate::tests::fd_lock_shared();
7886 let file = std::fs::File::open("/dev/null").unwrap();
7887 let pd = PlaneDescriptor::new(file.as_fd()).unwrap();
7888 assert_eq!((pd.stride(), pd.offset()), (None, None));
7889 let pd = pd.with_stride(832).with_offset(4096);
7890 assert_eq!((pd.stride(), pd.offset()), (Some(832), Some(4096)));
7891 }
7892
7893 #[test]
7894 #[cfg(target_os = "linux")]
7895 fn as_dma_answers_only_for_a_dma_buffer() {
7896 use std::os::fd::AsRawFd;
7897 let _lock = crate::tests::fd_lock_shared();
7898 let mem = Tensor::<u8>::new(&[64], Some(TensorMemory::Mem), None).unwrap();
7899 assert!(mem.as_dma().is_none());
7900 if !crate::test_support::dma_or_skip("as_dma_answers_only_for_a_dma_buffer") {
7901 return;
7902 }
7903 let t = Tensor::<u8>::new(&[64], Some(TensorMemory::DmaBuf), None).unwrap();
7904 let dma = t.as_dma().expect("a DMA tensor downcasts");
7905 assert_eq!(
7906 dma.fd.as_raw_fd(),
7907 t.dmabuf().unwrap().as_raw_fd(),
7908 "the downcast is this tensor's own buffer"
7909 );
7910 }
7911
7912 #[test]
7913 fn derived_identity_mixes_the_kind_into_the_key() {
7914 let id = BufferIdentity::derived(IdentityKind::Shm, u64::MAX);
7915 assert_eq!(id.id(), ((IdentityKind::Shm as u64) << 56) ^ u64::MAX);
7916 assert_ne!(
7917 id.id(),
7918 BufferIdentity::derived(IdentityKind::DmaBuf, u64::MAX).id()
7919 );
7920 }
7921
7922 #[test]
7923 fn empty_is_a_zero_element_shape() {
7924 let mut t = Tensor::<f32>::new(&[2, 3], Some(TensorMemory::Mem), None).unwrap();
7925 assert!(!TensorTrait::is_empty(&t));
7926 {
7927 let m = t.map().unwrap();
7928 assert_eq!(TensorMapTrait::len(&m), 6);
7929 assert_eq!(TensorMapTrait::size(&m), 24);
7930 assert!(!TensorMapTrait::is_empty(&m));
7931 }
7932 t.set_logical_shape(&[0]).unwrap();
7933 assert!(TensorTrait::is_empty(&t));
7934 let m = t.map().unwrap();
7935 assert_eq!(TensorMapTrait::size(&m), 0);
7936 assert!(TensorMapTrait::is_empty(&m));
7937 }
7938
7939 struct CountingCudaOps {
7940 maps: AtomicUsize,
7941 }
7942
7943 impl CudaGlOps for CountingCudaOps {
7944 fn map(&self, _r: crate::cuda::GraphicsResource) -> Option<(*mut std::ffi::c_void, usize)> {
7945 self.maps.fetch_add(1, Ordering::SeqCst);
7946 Some((0x1000usize as *mut _, 64))
7947 }
7948 fn unmap(&self, _r: crate::cuda::GraphicsResource) {}
7949 fn unregister(&self, _r: crate::cuda::GraphicsResource) {}
7950 }
7951
7952 #[test]
7953 fn cuda_accessors_follow_the_attached_handle() {
7954 let mut t = Tensor::<f32>::new(&[16], Some(TensorMemory::Mem), None).unwrap();
7955 assert!(!t.is_cuda_attached());
7956 assert!(t.cuda().is_none());
7957 assert!(t.cuda_map().is_none());
7958 assert!(t.cuda_map_mut().is_none());
7959
7960 let ops = Arc::new(CountingCudaOps {
7961 maps: AtomicUsize::new(0),
7962 });
7963 t.set_cuda_handle(CudaHandle::new_gl(
7964 0x1usize as crate::cuda::GraphicsResource,
7965 64,
7966 ops.clone(),
7967 ));
7968 assert!(t.is_cuda_attached());
7969 assert!(t.cuda().is_some());
7970 assert_eq!(t.cuda_map().map(|m| m.device_ptr() as usize), Some(0x1000));
7971 assert_eq!(t.cuda_map_mut().map(|m| m.len()), Some(64));
7972 assert_eq!(ops.maps.load(Ordering::SeqCst), 2);
7973 }
7974
7975 #[test]
7976 fn from_slice_copies_values_of_a_matching_shape() {
7977 let t = Tensor::<u16>::from_slice(&[1, 2, 3, 4, 5, 6], &[2, 3]).unwrap();
7978 assert_eq!(t.shape(), &[2, 3]);
7979 assert_eq!(t.map().unwrap().as_slice(), &[1, 2, 3, 4, 5, 6]);
7980 match Tensor::<u16>::from_slice(&[1, 2, 3, 4, 5], &[2, 3]) {
7981 Err(Error::InvalidShape(msg)) => assert!(msg.contains("values.len()=5"), "{msg}"),
7982 other => panic!("expected InvalidShape, got {other:?}"),
7983 }
7984 }
7985
7986 /// An `Any` compression request with no pinned memory tries the
7987 /// platform's zero-copy image first and falls back when it cannot serve
7988 /// the format (packed RGB f32 has no IOSurface layout, for one).
7989 #[cfg(not(target_os = "android"))]
7990 #[test]
7991 fn any_compression_request_falls_back_when_zero_copy_cannot_serve() {
7992 let _lock = crate::tests::fd_lock_shared();
7993 let desc =
7994 ImageDesc::new(10, 10, PixelFormat::Rgb, DType::F32).with_compression(Compression::Any);
7995 let t = Tensor::<f32>::image_desc(&desc).expect("Any never fails for compression reasons");
7996 assert_eq!(t.shape(), &[10, 10, 3]);
7997 assert_eq!(t.compression(), None);
7998 }
7999}
8000
8001#[cfg(test)]
8002mod cpu_access_tests {
8003 use super::*;
8004
8005 #[test]
8006 fn covers_matrix() {
8007 use CpuAccess::*;
8008 // Every declaration covers a narrower or equal request…
8009 for a in [None, Read, Write, ReadWrite] {
8010 assert!(a.covers(None), "{a:?} must cover None");
8011 assert!(ReadWrite.covers(a), "ReadWrite must cover {a:?}");
8012 }
8013 assert!(Read.covers(Read));
8014 assert!(Write.covers(Write));
8015 // …and never a wider one.
8016 assert!(!None.covers(Read));
8017 assert!(!None.covers(Write));
8018 assert!(!Read.covers(Write));
8019 assert!(!Read.covers(ReadWrite));
8020 assert!(!Write.covers(Read));
8021 assert!(!Write.covers(ReadWrite));
8022 }
8023
8024 #[test]
8025 fn map_with_none_is_invalid() {
8026 let t = Tensor::<u8>::new(&[16], Some(TensorMemory::Mem), None).unwrap();
8027 match t.map_with(CpuAccess::None) {
8028 Err(Error::InvalidArgument(_)) => {}
8029 Err(other) => panic!("expected InvalidArgument, got {other:?}"),
8030 Ok(_) => panic!("map_with(CpuAccess::None) must not succeed"),
8031 }
8032 }
8033
8034 #[test]
8035 fn read_map_rejects_mutation_uniformly() {
8036 // Mem backend: map_read yields a working read view whose mutable
8037 // accessor panics (the uniform cross-backend contract).
8038 let t = Tensor::<u8>::new(&[8], Some(TensorMemory::Mem), None).unwrap();
8039 t.map_mut().unwrap().as_mut_slice().copy_from_slice(&[7; 8]);
8040 let ro = t.map_read().unwrap();
8041 assert_eq!(ro.as_slice(), &[7; 8]);
8042 drop(ro);
8043 let result = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
8044 let mut ro = t.map_read().unwrap();
8045 let _ = ro.as_mut_slice();
8046 }));
8047 assert!(result.is_err(), "as_mut_slice through map_read must panic");
8048 }
8049
8050 #[test]
8051 fn write_and_rw_maps_stay_mutable() {
8052 let t = Tensor::<u8>::new(&[4], Some(TensorMemory::Mem), None).unwrap();
8053 t.map_write()
8054 .unwrap()
8055 .as_mut_slice()
8056 .copy_from_slice(&[1; 4]);
8057 t.map().unwrap().as_mut_slice().copy_from_slice(&[2; 4]);
8058 assert_eq!(t.map_read().unwrap().as_slice(), &[2; 4]);
8059 }
8060
8061 /// The read-only IOSurface lock path: a `map_read` must observe data
8062 /// written through a prior read-write lock, and its unlock (which
8063 /// skips the cache flush) must not disturb subsequent reads.
8064 #[test]
8065 #[cfg(target_os = "macos")]
8066 fn iosurface_read_only_lock_roundtrip() {
8067 // Declares this test as an fd-opener; see FD_LOCK.
8068 let _lock = crate::tests::fd_lock_shared();
8069 let Ok(t) = Tensor::<u8>::new(&[64], Some(TensorMemory::DmaBuf), None) else {
8070 crate::test_support::report_skip("IOSurface unavailable");
8071 return;
8072 };
8073 {
8074 let mut m = t.map_mut().unwrap();
8075 for (i, b) in m.as_mut_slice().iter_mut().enumerate() {
8076 *b = (i * 3) as u8;
8077 }
8078 }
8079 for _ in 0..2 {
8080 let ro = t.map_read().unwrap();
8081 for (i, b) in ro.as_slice().iter().enumerate() {
8082 assert_eq!(*b, (i * 3) as u8, "byte {i} through read-only lock");
8083 }
8084 }
8085 let result = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
8086 let mut ro = t.map_read().unwrap();
8087 let _ = ro.as_mut_slice();
8088 }));
8089 assert!(result.is_err(), "IOSurface read map must reject mutation");
8090 }
8091}
8092
8093#[cfg(test)]
8094mod tests {
8095 use std::sync::{PoisonError, RwLock};
8096
8097 use super::*;
8098
8099 #[ctor::ctor(unsafe)]
8100 fn init() {
8101 env_logger::Builder::from_env(env_logger::Env::default().default_filter_or("info")).init();
8102 }
8103
8104 /// Macro to get the current function name for logging in tests.
8105 #[cfg(target_os = "linux")]
8106 macro_rules! function {
8107 () => {{
8108 fn f() {}
8109 fn type_name_of<T>(_: T) -> &'static str {
8110 std::any::type_name::<T>()
8111 }
8112 let name = type_name_of(f);
8113
8114 // Find and cut the rest of the path
8115 match &name[..name.len() - 3].rfind(':') {
8116 Some(pos) => &name[pos + 1..name.len() - 3],
8117 None => &name[..name.len() - 3],
8118 }
8119 }};
8120 }
8121
8122 #[test]
8123 #[cfg(target_os = "linux")]
8124 fn test_tensor() {
8125 let _lock = crate::tests::fd_lock_shared();
8126 let shape = vec![1];
8127 let tensor = DmaTensor::<f32>::new(&shape, Some("dma_tensor"));
8128 let dma_enabled = tensor.is_ok();
8129
8130 let tensor = Tensor::<f32>::new(&shape, None, None).expect("Failed to create tensor");
8131 // Auto-select priority is Dma > Mem; Shm is never auto-selected.
8132 match dma_enabled {
8133 true => assert_eq!(tensor.memory(), TensorMemory::DmaBuf),
8134 false => assert_eq!(tensor.memory(), TensorMemory::Mem),
8135 }
8136 }
8137
8138 #[test]
8139 #[cfg(any(target_os = "macos", target_os = "ios"))]
8140 fn test_tensor() {
8141 let shape = vec![1];
8142 let tensor = Tensor::<f32>::new(&shape, None, None).expect("Failed to create tensor");
8143 // macOS/iOS auto-fallback chain: IOSurface (Dma) → Mem. Healthy systems
8144 // return Dma; Mem only appears under memory pressure or sandboxed
8145 // contexts where IOSurfaceCreate fails. Shm is never auto-selected.
8146 let m = tensor.memory();
8147 assert!(
8148 matches!(m, TensorMemory::DmaBuf | TensorMemory::Mem),
8149 "Unexpected auto-fallback result on macOS/iOS: {m:?}"
8150 );
8151 }
8152
8153 #[test]
8154 #[cfg(all(
8155 unix,
8156 not(any(target_os = "linux", target_os = "macos", target_os = "ios"))
8157 ))]
8158 fn test_tensor() {
8159 let shape = vec![1];
8160 let tensor = Tensor::<f32>::new(&shape, None, None).expect("Failed to create tensor");
8161 // Other Unix (BSD): no DMA, so auto-selection is Mem (Shm is
8162 // explicit-only, never auto-selected).
8163 assert_eq!(tensor.memory(), TensorMemory::Mem);
8164 }
8165
8166 #[test]
8167 #[cfg(not(unix))]
8168 fn test_tensor() {
8169 let shape = vec![1];
8170 let tensor = Tensor::<f32>::new(&shape, None, None).expect("Failed to create tensor");
8171 assert_eq!(tensor.memory(), TensorMemory::Mem);
8172 }
8173
8174 #[test]
8175 #[cfg(target_os = "linux")]
8176 fn test_dma_tensor() {
8177 let _lock = crate::tests::fd_lock_shared();
8178 if !crate::test_support::dma_or_skip("test_dma_tensor") {
8179 return;
8180 }
8181
8182 let shape = vec![2, 3, 4];
8183 let tensor =
8184 DmaTensor::<f32>::new(&shape, Some("test_tensor")).expect("Failed to create tensor");
8185
8186 const DUMMY_VALUE: f32 = 12.34;
8187
8188 assert_eq!(tensor.memory(), TensorMemory::DmaBuf);
8189 assert_eq!(tensor.name(), "test_tensor");
8190 assert_eq!(tensor.shape(), &shape);
8191 assert_eq!(tensor.size(), 2 * 3 * 4 * std::mem::size_of::<f32>());
8192 assert_eq!(tensor.len(), 2 * 3 * 4);
8193
8194 {
8195 let mut tensor_map = tensor.map().expect("Failed to map DMA memory");
8196 tensor_map.fill(42.0);
8197 assert!(tensor_map.iter().all(|&x| x == 42.0));
8198 }
8199
8200 {
8201 let shared = Tensor::<f32>::from_fd(
8202 tensor
8203 .clone_fd()
8204 .expect("Failed to duplicate tensor file descriptor"),
8205 &shape,
8206 Some("test_tensor_shared"),
8207 )
8208 .expect("Failed to create tensor from fd");
8209
8210 assert_eq!(shared.memory(), TensorMemory::DmaBuf);
8211 assert_eq!(shared.name(), "test_tensor_shared");
8212 assert_eq!(shared.shape(), &shape);
8213
8214 let mut tensor_map = shared.map().expect("Failed to map DMA memory from fd");
8215 tensor_map.fill(DUMMY_VALUE);
8216 assert!(tensor_map.iter().all(|&x| x == DUMMY_VALUE));
8217 }
8218
8219 {
8220 let tensor_map = tensor.map().expect("Failed to map DMA memory");
8221 assert!(tensor_map.iter().all(|&x| x == DUMMY_VALUE));
8222 }
8223
8224 let mut tensor = DmaTensor::<u8>::new(&shape, None).expect("Failed to create tensor");
8225 assert_eq!(tensor.shape(), &shape);
8226 let new_shape = vec![3, 4, 4];
8227 assert!(
8228 tensor.reshape(&new_shape).is_err(),
8229 "Reshape should fail due to size mismatch"
8230 );
8231 assert_eq!(tensor.shape(), &shape, "Shape should remain unchanged");
8232
8233 let new_shape = vec![2, 3, 4];
8234 tensor.reshape(&new_shape).expect("Reshape should succeed");
8235 assert_eq!(
8236 tensor.shape(),
8237 &new_shape,
8238 "Shape should be updated after successful reshape"
8239 );
8240
8241 {
8242 let mut tensor_map = tensor.map().expect("Failed to map DMA memory");
8243 tensor_map.fill(1);
8244 assert!(tensor_map.iter().all(|&x| x == 1));
8245 }
8246
8247 {
8248 let mut tensor_map = tensor.map().expect("Failed to map DMA memory");
8249 tensor_map[2] = 42;
8250 assert_eq!(tensor_map[1], 1, "Value at index 1 should be 1");
8251 assert_eq!(tensor_map[2], 42, "Value at index 2 should be 42");
8252 }
8253 }
8254
8255 #[test]
8256 #[cfg(unix)]
8257 fn test_shm_tensor() {
8258 let _lock = crate::tests::fd_lock_shared();
8259 if !is_shm_available() {
8260 log::warn!("SKIPPED: test_shm_tensor - SHM allocation unavailable on this platform");
8261 return;
8262 }
8263 let shape = vec![2, 3, 4];
8264 let tensor =
8265 ShmTensor::<f32>::new(&shape, Some("test_tensor")).expect("Failed to create tensor");
8266 assert_eq!(tensor.shape(), &shape);
8267 assert_eq!(tensor.size(), 2 * 3 * 4 * std::mem::size_of::<f32>());
8268 assert_eq!(tensor.name(), "test_tensor");
8269
8270 const DUMMY_VALUE: f32 = 12.34;
8271 {
8272 let mut tensor_map = tensor.map().expect("Failed to map shared memory");
8273 tensor_map.fill(42.0);
8274 assert!(tensor_map.iter().all(|&x| x == 42.0));
8275 }
8276
8277 {
8278 let shared = Tensor::<f32>::from_fd(
8279 tensor
8280 .clone_fd()
8281 .expect("Failed to duplicate tensor file descriptor"),
8282 &shape,
8283 Some("test_tensor_shared"),
8284 )
8285 .expect("Failed to create tensor from fd");
8286
8287 assert_eq!(shared.memory(), TensorMemory::Shm);
8288 assert_eq!(shared.name(), "test_tensor_shared");
8289 assert_eq!(shared.shape(), &shape);
8290
8291 let mut tensor_map = shared.map().expect("Failed to map shared memory from fd");
8292 tensor_map.fill(DUMMY_VALUE);
8293 assert!(tensor_map.iter().all(|&x| x == DUMMY_VALUE));
8294 }
8295
8296 {
8297 let tensor_map = tensor.map().expect("Failed to map shared memory");
8298 assert!(tensor_map.iter().all(|&x| x == DUMMY_VALUE));
8299 }
8300
8301 let mut tensor = ShmTensor::<u8>::new(&shape, None).expect("Failed to create tensor");
8302 assert_eq!(tensor.shape(), &shape);
8303 let new_shape = vec![3, 4, 4];
8304 assert!(
8305 tensor.reshape(&new_shape).is_err(),
8306 "Reshape should fail due to size mismatch"
8307 );
8308 assert_eq!(tensor.shape(), &shape, "Shape should remain unchanged");
8309
8310 let new_shape = vec![2, 3, 4];
8311 tensor.reshape(&new_shape).expect("Reshape should succeed");
8312 assert_eq!(
8313 tensor.shape(),
8314 &new_shape,
8315 "Shape should be updated after successful reshape"
8316 );
8317
8318 {
8319 let mut tensor_map = tensor.map().expect("Failed to map shared memory");
8320 tensor_map.fill(1);
8321 assert!(tensor_map.iter().all(|&x| x == 1));
8322 }
8323
8324 {
8325 let mut tensor_map = tensor.map().expect("Failed to map shared memory");
8326 tensor_map[2] = 42;
8327 assert_eq!(tensor_map[1], 1, "Value at index 1 should be 1");
8328 assert_eq!(tensor_map[2], 42, "Value at index 2 should be 42");
8329 }
8330 }
8331
8332 #[test]
8333 fn mem_subview_partitions_parent_buffer() {
8334 // One heap [2,4] u8 parent (8 bytes). Two [1,4] sub-views at byte
8335 // offsets 0 and 4 must share the parent allocation (zero-copy) and be
8336 // independently writable: view 0 owns bytes [0,4), view 1 owns [4,8).
8337 // Today this is impossible — heap offset is rejected and there is no
8338 // shared sub-view constructor.
8339 let parent = Tensor::<u8>::new(&[2, 4], Some(TensorMemory::Mem), None).unwrap();
8340 let view0 = parent.subview(0, &[1, 4]).expect("subview at offset 0");
8341 let view1 = parent.subview(4, &[1, 4]).expect("subview at offset 4");
8342
8343 view1
8344 .map()
8345 .unwrap()
8346 .as_mut_slice()
8347 .copy_from_slice(&[10, 20, 30, 40]);
8348 view0
8349 .map()
8350 .unwrap()
8351 .as_mut_slice()
8352 .copy_from_slice(&[1, 2, 3, 4]);
8353
8354 // Each view sees only its own window.
8355 assert_eq!(view0.map().unwrap().as_slice(), &[1, 2, 3, 4]);
8356 assert_eq!(view1.map().unwrap().as_slice(), &[10, 20, 30, 40]);
8357 // The parent buffer is correctly partitioned (shared, zero-copy).
8358 assert_eq!(
8359 parent.map().unwrap().as_slice(),
8360 &[1, 2, 3, 4, 10, 20, 30, 40]
8361 );
8362 }
8363
8364 #[test]
8365 fn batch_partitions_leading_dim() {
8366 // Raw [4,2,2,3] u8 batched tensor: 4 elements of 12 bytes each. batch(n)
8367 // yields element n at offset n*12, sharing the parent buffer (zero-copy).
8368 let parent = Tensor::<u8>::new(&[4, 2, 2, 3], Some(TensorMemory::Mem), None).unwrap();
8369 for i in 0..4u8 {
8370 let e = parent.batch(i as usize).expect("batch element");
8371 assert_eq!(e.shape(), &[2, 2, 3]);
8372 // A batch element shares the parent's BufferIdentity.
8373 assert_eq!(e.buffer_identity().id(), parent.buffer_identity().id());
8374 for b in e.map().unwrap().as_mut_slice() {
8375 *b = i + 1;
8376 }
8377 }
8378 // Each element occupies its own 12-byte band of the parent.
8379 let whole = parent.map().unwrap();
8380 let s = whole.as_slice();
8381 for i in 0..4usize {
8382 assert!(
8383 s[i * 12..(i + 1) * 12].iter().all(|&b| b == (i as u8 + 1)),
8384 "band {i} not partitioned: {:?}",
8385 &s[i * 12..(i + 1) * 12]
8386 );
8387 }
8388 }
8389
8390 #[test]
8391 fn view_origin_snapshots_parent_and_composes() {
8392 // view() on a whole image snapshots the parent dims + the view's origin.
8393 let parent = Tensor::<u8>::image(
8394 100,
8395 80,
8396 PixelFormat::Rgba,
8397 Some(TensorMemory::Mem),
8398 crate::CpuAccess::ReadWrite,
8399 )
8400 .unwrap();
8401 assert_eq!(
8402 parent.view_origin(),
8403 None,
8404 "whole tensor has no view_origin"
8405 );
8406 let v = parent.view(Region::new(10, 20, 30, 40)).unwrap();
8407 assert_eq!(
8408 v.view_origin(),
8409 Some(ViewOrigin {
8410 parent_width: 100,
8411 parent_height: 80,
8412 parent_row_stride: 100 * 4, // tight RGBA pitch
8413 x: 10,
8414 y: 20
8415 })
8416 );
8417 // A view of a view keeps the ROOT parent and accumulates the origin.
8418 let v2 = v.view(Region::new(5, 5, 10, 10)).unwrap();
8419 assert_eq!(
8420 v2.view_origin(),
8421 Some(ViewOrigin {
8422 parent_width: 100,
8423 parent_height: 80,
8424 parent_row_stride: 100 * 4,
8425 x: 15,
8426 y: 25
8427 }),
8428 "nested view composes onto the root parent"
8429 );
8430 }
8431
8432 #[test]
8433 fn view_origin_none_for_raw_batch() {
8434 // A raw (unformatted) batched tensor has no pixel geometry, so batch()
8435 // leaves view_origin None (the per-slot path, not the one-import pivot).
8436 let parent = Tensor::<u8>::new(&[4, 2, 2, 3], Some(TensorMemory::Mem), None).unwrap();
8437 assert_eq!(parent.batch(2).unwrap().view_origin(), None);
8438 }
8439
8440 #[test]
8441 fn batch_rejects_out_of_bounds_index() {
8442 let parent = Tensor::<u8>::new(&[4, 2, 2, 3], Some(TensorMemory::Mem), None).unwrap();
8443 match parent.batch(4) {
8444 Err(Error::BatchIndexOutOfBounds { index, batch }) => {
8445 assert_eq!((index, batch), (4, 4));
8446 }
8447 other => panic!("expected BatchIndexOutOfBounds, got {other:?}"),
8448 }
8449 }
8450
8451 #[test]
8452 fn batch_zero_on_unit_n_is_whole() {
8453 // N == 1: batch(0) is the whole per-element block at offset 0 (no plane_offset).
8454 let parent = Tensor::<u8>::new(&[1, 2, 2, 3], Some(TensorMemory::Mem), None).unwrap();
8455 let e = parent.batch(0).unwrap();
8456 assert_eq!(e.shape(), &[2, 2, 3]);
8457 assert_eq!(e.plane_offset(), None);
8458 assert_eq!(e.buffer_identity().id(), parent.buffer_identity().id());
8459 }
8460
8461 #[test]
8462 fn mem_subview_rejects_unaligned_offset() {
8463 // f32 has align 4; a byte offset of 2 cannot back a valid `*const f32`.
8464 let parent = Tensor::<f32>::new(&[8], Some(TensorMemory::Mem), None).unwrap();
8465 assert!(parent.subview(2, &[1]).is_err());
8466 // A correctly aligned offset is accepted.
8467 assert!(parent.subview(4, &[1]).is_ok());
8468 }
8469
8470 #[test]
8471 fn mem_subview_rejects_out_of_bounds() {
8472 let parent = Tensor::<u8>::new(&[8], Some(TensorMemory::Mem), None).unwrap();
8473 // offset 6 + 4 bytes = 10 exceeds the 8-byte allocation.
8474 assert!(parent.subview(6, &[4]).is_err());
8475 }
8476
8477 /// Regression guard for the `TensorTrait::view` promotion (R2): a `subview`
8478 /// must share the parent's `BufferIdentity` on **every** backend, not mint a
8479 /// fresh one. Identity-keyed caches (the GL EGLImage import) rely on this to
8480 /// treat offset-distinct windows of one buffer as a single import; a fresh
8481 /// identity would silently break that and regress zero-copy import reuse.
8482 ///
8483 /// Runs each backend that can be allocated without a GPU/GL context on the
8484 /// test host: `Mem` (always); `Shm` (when POSIX shm is available); the
8485 /// platform-native zero-copy buffer `Dma` (DMA-BUF on Linux / IOSurface on
8486 /// macOS, when available). `Pbo` shares its identity the same way (see
8487 /// `pbo.rs` `view`) but needs a live GL context, so it is exercised by the
8488 /// image-crate GL tests rather than here.
8489 #[test]
8490 fn subview_shares_buffer_identity_all_backends() {
8491 // Declares this test as an fd-opener; see FD_LOCK.
8492 let _lock = crate::tests::fd_lock_shared();
8493 // u8 has align 1, so every byte offset is valid for the alignment check;
8494 // this isolates the identity-sharing contract from alignment concerns.
8495 // `window` is the view shape; a formatted parent validates it against
8496 // its own format, so an image parent needs an image-shaped window.
8497 let assert_parent_shares = |parent: Tensor<u8>, window: &[usize], label: &str| {
8498 let parent_id = parent.buffer_identity().id();
8499 // Two offset-distinct windows must both carry the parent's identity.
8500 let v0 = parent
8501 .subview(0, window)
8502 .unwrap_or_else(|e| panic!("{label}: subview(0) failed: {e:?}"));
8503 let v1 = parent
8504 .subview(16, window)
8505 .unwrap_or_else(|e| panic!("{label}: subview(16) failed: {e:?}"));
8506 assert_eq!(
8507 v0.buffer_identity().id(),
8508 parent_id,
8509 "{label}: subview(0) minted a fresh BufferIdentity"
8510 );
8511 assert_eq!(
8512 v1.buffer_identity().id(),
8513 parent_id,
8514 "{label}: subview(16) minted a fresh BufferIdentity"
8515 );
8516 };
8517 let assert_shares = |memory: TensorMemory, label: &str| {
8518 let parent = Tensor::<u8>::new(&[64], Some(memory), None)
8519 .unwrap_or_else(|e| panic!("{label}: parent alloc failed: {e:?}"));
8520 assert_parent_shares(parent, &[16], label);
8521 };
8522
8523 assert_shares(TensorMemory::Mem, "Mem");
8524
8525 #[cfg(unix)]
8526 if crate::is_shm_available() {
8527 assert_shares(TensorMemory::Shm, "Shm");
8528 }
8529
8530 // Dma == DMA-BUF on Linux, IOSurface on macOS, AHardwareBuffer on
8531 // Android; same public variant.
8532 #[cfg(target_os = "linux")]
8533 if crate::test_support::dma_or_skip("subview_shares_buffer_identity_all_backends (Dma)") {
8534 assert_shares(TensorMemory::DmaBuf, "Dma");
8535 }
8536 #[cfg(not(any(target_os = "linux", target_os = "windows")))]
8537 if crate::is_gpu_buffer_available() {
8538 assert_shares(TensorMemory::DmaBuf, "Dma");
8539 }
8540 // Windows' Dma is an image-formatted D3D11 texture, which
8541 // `Tensor::new` deliberately refuses -- allocate it the only way it
8542 // can be allocated and hold it to the same contract.
8543 #[cfg(target_os = "windows")]
8544 if crate::is_gpu_buffer_available() {
8545 let parent = Tensor::<u8>::image(
8546 64,
8547 1,
8548 PixelFormat::Grey,
8549 Some(TensorMemory::DmaBuf),
8550 CpuAccess::ReadWrite,
8551 )
8552 .expect("D3D11 texture alloc failed");
8553 assert_parent_shares(parent, &[1, 16, 1], "Dma");
8554 }
8555 }
8556
8557 /// A sub-view records a plane offset, and on Windows the D3D11 storage
8558 /// already applied that offset itself when the view was made -- so the
8559 /// map's plane-offset gate must accept it, exactly as it does for the
8560 /// other zero-copy backings. It used to refuse, with a message that named
8561 /// every backend but this one.
8562 #[cfg(target_os = "windows")]
8563 #[test]
8564 fn subview_of_a_texture_tensor_maps_at_its_offset() {
8565 if !crate::is_gpu_buffer_available() {
8566 return;
8567 }
8568 let t = Tensor::<u8>::image(
8569 32,
8570 8,
8571 PixelFormat::Rgba,
8572 Some(TensorMemory::DmaBuf),
8573 CpuAccess::ReadWrite,
8574 )
8575 .unwrap();
8576 // For a texture tensor this is the backing pitch either way: the
8577 // allocation records the staging pitch when it exceeds the tight row,
8578 // and the two are equal when it does not.
8579 let stride = t.effective_row_stride().unwrap();
8580 let v = t.subview(stride * 2, &[2, 32, 4]).unwrap();
8581 assert!(
8582 v.plane_offset().is_some_and(|o| o > 0),
8583 "a sub-view at a row offset must carry one"
8584 );
8585 let m = v
8586 .map_read()
8587 .expect("a sub-view of a texture tensor must map");
8588 assert_eq!(m.as_slice().len(), 2 * stride, "two rows of the window");
8589 }
8590
8591 #[test]
8592 fn mem_subview_four_views_no_aliasing() {
8593 // One [4,3] f32 parent; four [1,3] views at 12-byte strides, each
8594 // written independently. Exercises a multi-byte element type (offsets
8595 // must stay element-aligned) and N-way zero-copy sharing.
8596 let parent = Tensor::<f32>::new(&[4, 3], Some(TensorMemory::Mem), None).unwrap();
8597 let frame = 3 * std::mem::size_of::<f32>();
8598 for i in 0..4 {
8599 let v = parent.subview(i * frame, &[1, 3]).unwrap();
8600 let val = i as f32 + 1.0;
8601 v.map()
8602 .unwrap()
8603 .as_mut_slice()
8604 .copy_from_slice(&[val, val, val]);
8605 }
8606 assert_eq!(
8607 parent.map().unwrap().as_slice(),
8608 &[1.0, 1.0, 1.0, 2.0, 2.0, 2.0, 3.0, 3.0, 3.0, 4.0, 4.0, 4.0]
8609 );
8610 }
8611
8612 #[test]
8613 fn mem_subview_inherits_format_and_row_stride() {
8614 // A sub-view is a ready-to-use sub-image: it inherits the parent's
8615 // pixel format and (crucially) its padded row stride, so a strided
8616 // parent yields strided windows. Set a stride wider than the tight row
8617 // to exercise the row_stride inheritance path specifically.
8618 let mut parent = Tensor::<u8>::image(
8619 100,
8620 100,
8621 PixelFormat::Rgba,
8622 Some(TensorMemory::Mem),
8623 crate::CpuAccess::ReadWrite,
8624 )
8625 .unwrap();
8626 parent.set_row_stride_unchecked(512); // padded stride (> 100*4)
8627 let view = parent.subview(4096, &[10, 10, 4]).unwrap();
8628 assert_eq!(view.format(), Some(PixelFormat::Rgba), "format inherited");
8629 assert_eq!(view.row_stride(), Some(512), "row_stride inherited");
8630 }
8631
8632 #[test]
8633 fn mem_strided_subview_maps_offset_and_byte_size() {
8634 // Integration of the sub-region offset (PR #89) and the strided-map
8635 // `byte_size_override` (PR #90): a strided sub-view exposes its full
8636 // padded window (`row_stride × rows`) starting at the view's byte
8637 // offset, mapped zero-copy into the parent.
8638 let parent = Tensor::<u8>::new(&[2048], Some(TensorMemory::Mem), None).unwrap();
8639 let mut view = parent.subview(128, &[8, 16]).unwrap(); // 8 rows × 16 @ off 128
8640 assert_eq!(view.plane_offset(), Some(128));
8641 view.set_row_stride_unchecked(32); // padded stride (> 16)
8642
8643 {
8644 let mut m = view.map().unwrap();
8645 let s = m.as_mut_slice();
8646 // Strided map exposes the padded window: stride(32) × rows(8) = 256.
8647 assert_eq!(
8648 s.len(),
8649 256,
8650 "strided map exposes the full padded byte window"
8651 );
8652 s[0] = 0xAA; // row 0, col 0
8653 s[32] = 0xBB; // row 1, col 0 (one stride in)
8654 }
8655
8656 // Zero-copy: the writes land in the parent at the view's offset.
8657 let p = parent.map().unwrap();
8658 let pb = p.as_slice();
8659 assert_eq!(pb[128], 0xAA, "row 0 writes at parent offset 128");
8660 assert_eq!(
8661 pb[128 + 32],
8662 0xBB,
8663 "row 1 writes at parent offset 128 + stride"
8664 );
8665 }
8666
8667 #[test]
8668 #[cfg(unix)]
8669 fn shm_subview_partitions_parent_buffer() {
8670 // Declares this test as an fd-opener; see FD_LOCK.
8671 let _lock = crate::tests::fd_lock_shared();
8672 // Mirrors `mem_subview_partitions_parent_buffer` for Shm: one [2,4] u8
8673 // parent shared segment (8 bytes); two [1,4] sub-views at byte offsets 0
8674 // and 4 must share the segment (zero-copy, via cloned fd) and be
8675 // independently writable — view 0 owns [0,4), view 1 owns [4,8).
8676 if !crate::is_shm_available() {
8677 crate::test_support::report_skip("shm not available");
8678 return;
8679 }
8680 let parent = Tensor::<u8>::new(&[2, 4], Some(TensorMemory::Shm), None).unwrap();
8681 let view0 = parent.subview(0, &[1, 4]).expect("shm subview at offset 0");
8682 let view1 = parent.subview(4, &[1, 4]).expect("shm subview at offset 4");
8683
8684 view1
8685 .map()
8686 .unwrap()
8687 .as_mut_slice()
8688 .copy_from_slice(&[10, 20, 30, 40]);
8689 view0
8690 .map()
8691 .unwrap()
8692 .as_mut_slice()
8693 .copy_from_slice(&[1, 2, 3, 4]);
8694
8695 assert_eq!(view0.map().unwrap().as_slice(), &[1, 2, 3, 4]);
8696 assert_eq!(view1.map().unwrap().as_slice(), &[10, 20, 30, 40]);
8697 // The parent sees the full partitioned segment (shared, zero-copy).
8698 assert_eq!(
8699 parent.map().unwrap().as_slice(),
8700 &[1, 2, 3, 4, 10, 20, 30, 40]
8701 );
8702 // A sub-view of a sub-view composes the offset.
8703 let nested = view1.subview(2, &[1, 2]).expect("nested shm subview");
8704 assert_eq!(nested.map().unwrap().as_slice(), &[30, 40]);
8705 }
8706
8707 #[test]
8708 #[cfg(unix)]
8709 fn shm_subview_rejects_unaligned_and_oob() {
8710 // Declares this test as an fd-opener; see FD_LOCK.
8711 let _lock = crate::tests::fd_lock_shared();
8712 if !crate::is_shm_available() {
8713 crate::test_support::report_skip("shm not available");
8714 return;
8715 }
8716 // f32 align 4: a 2-byte offset cannot back a valid `*const f32`.
8717 let parent = Tensor::<f32>::new(&[8], Some(TensorMemory::Shm), None).unwrap();
8718 assert!(parent.subview(2, &[1]).is_err());
8719 assert!(parent.subview(4, &[1]).is_ok());
8720 // Out of bounds: offset 6 + 4 bytes = 10 > 8-byte (u8) segment.
8721 let p2 = Tensor::<u8>::new(&[8], Some(TensorMemory::Shm), None).unwrap();
8722 assert!(p2.subview(6, &[4]).is_err());
8723 }
8724
8725 #[test]
8726 #[cfg(target_os = "linux")]
8727 fn dma_subview_matches_mem_subview() {
8728 // Serialize against the fd-leak tests: this test opens DMA fds (alloc +
8729 // clone_fd), which would otherwise perturb their fd counts.
8730 let _lock = crate::tests::fd_lock_shared();
8731 // Identical sub-view semantics across Dma (shared fd) and Mem (shared
8732 // Arc): same offsets → same logical windows → same partition.
8733 if !crate::test_support::dma_or_skip("dma_subview_matches_mem_subview") {
8734 return;
8735 }
8736 let dma = Tensor::<u8>::new(&[8], Some(TensorMemory::DmaBuf), None).unwrap();
8737 let mem = Tensor::<u8>::new(&[8], Some(TensorMemory::Mem), None).unwrap();
8738 for parent in [&dma, &mem] {
8739 let v0 = parent.subview(0, &[4]).unwrap();
8740 let v1 = parent.subview(4, &[4]).unwrap();
8741 v0.map()
8742 .unwrap()
8743 .as_mut_slice()
8744 .copy_from_slice(&[1, 2, 3, 4]);
8745 v1.map()
8746 .unwrap()
8747 .as_mut_slice()
8748 .copy_from_slice(&[5, 6, 7, 8]);
8749 assert_eq!(parent.map().unwrap().as_slice(), &[1, 2, 3, 4, 5, 6, 7, 8]);
8750 }
8751 }
8752
8753 #[test]
8754 #[cfg(target_os = "linux")]
8755 fn dma_strided_subview_maps_padded_window() {
8756 // The strided-map path differs by backing: DMA maps through
8757 // `mmap_offset` + the `byte_size_override`, not the Mem `Arc` slice. A
8758 // padded sub-view of a DMA buffer must still expose its full
8759 // `row_stride × rows` window zero-copy at the view's offset (the GPU
8760 // batched-render-to-DMA case). Mirrors
8761 // `mem_strided_subview_maps_offset_and_byte_size` on a Dma parent.
8762 let _lock = crate::tests::fd_lock_shared();
8763 if !crate::test_support::dma_or_skip("dma_strided_subview_maps_padded_window") {
8764 return;
8765 }
8766 let parent = Tensor::<u8>::new(&[2048], Some(TensorMemory::DmaBuf), None).unwrap();
8767 let mut view = parent.subview(128, &[8, 16]).unwrap();
8768 assert_eq!(view.plane_offset(), Some(128));
8769 view.set_row_stride_unchecked(32); // padded stride (> 16)
8770
8771 {
8772 let mut m = view.map().unwrap();
8773 let s = m.as_mut_slice();
8774 assert_eq!(s.len(), 256, "strided DMA map exposes stride(32) × rows(8)");
8775 s[0] = 0xAA; // row 0, col 0
8776 s[32] = 0xBB; // row 1, col 0 (one stride in)
8777 }
8778
8779 let p = parent.map().unwrap();
8780 let pb = p.as_slice();
8781 assert_eq!(pb[128], 0xAA, "row 0 writes at parent offset 128");
8782 assert_eq!(
8783 pb[128 + 32],
8784 0xBB,
8785 "row 1 writes at parent offset 128 + stride"
8786 );
8787 }
8788
8789 #[test]
8790 #[cfg(target_os = "linux")]
8791 fn view_single_row_snapshots_parent_stride() {
8792 // A single-row `view()` keeps a TIGHT `row_stride` for map-span safety,
8793 // but its `view_origin` snapshots the PARENT row stride — the GL backend
8794 // keys its EGLImage import/pitch on that snapshot (not the view's tight
8795 // stride), so single-row and multi-row sibling views collapse onto the
8796 // same parent import.
8797 let _lock = crate::tests::fd_lock_shared();
8798 // 8x4 RGBA with a padded 64-byte row stride (tight row = 8*4 = 32).
8799 if !crate::test_support::dma_or_skip("view_single_row_snapshots_parent_stride") {
8800 return;
8801 }
8802 let parent = Tensor::<u8>::image_with_stride(
8803 8,
8804 4,
8805 PixelFormat::Rgba,
8806 64,
8807 Some(TensorMemory::DmaBuf),
8808 crate::CpuAccess::ReadWrite,
8809 )
8810 .unwrap();
8811 assert_eq!(parent.effective_row_stride(), Some(64));
8812 // Bottom row (y=3) at x>0 — the case the tight single-row stride guards.
8813 let row = parent.view(Region::new(2, 3, 4, 1)).unwrap();
8814 // The view's own stride is tight (4*4 = 16) so its strided map stays in
8815 // bounds; the GL-facing parent pitch (64) lives in `view_origin`.
8816 assert_eq!(row.effective_row_stride(), Some(16));
8817 let vo = row.view_origin().expect("a view carries a view_origin");
8818 assert_eq!(
8819 vo.parent_row_stride, 64,
8820 "GL keys/pitches a view on the parent stride, not its tight one"
8821 );
8822 // The tight stride keeps map() in-bounds for the bottom / x>0 single row.
8823 assert_eq!(row.map().unwrap().as_slice().len(), 16);
8824 }
8825
8826 #[test]
8827 fn test_mem_tensor() {
8828 let shape = vec![2, 3, 4];
8829 let tensor =
8830 MemTensor::<f32>::new(&shape, Some("test_tensor")).expect("Failed to create tensor");
8831 assert_eq!(tensor.shape(), &shape);
8832 assert_eq!(tensor.size(), 2 * 3 * 4 * std::mem::size_of::<f32>());
8833 assert_eq!(tensor.name(), "test_tensor");
8834
8835 {
8836 let mut tensor_map = tensor.map().expect("Failed to map memory");
8837 tensor_map.fill(42.0);
8838 assert!(tensor_map.iter().all(|&x| x == 42.0));
8839 }
8840
8841 let mut tensor = MemTensor::<u8>::new(&shape, None).expect("Failed to create tensor");
8842 assert_eq!(tensor.shape(), &shape);
8843 let new_shape = vec![3, 4, 4];
8844 assert!(
8845 tensor.reshape(&new_shape).is_err(),
8846 "Reshape should fail due to size mismatch"
8847 );
8848 assert_eq!(tensor.shape(), &shape, "Shape should remain unchanged");
8849
8850 let new_shape = vec![2, 3, 4];
8851 tensor.reshape(&new_shape).expect("Reshape should succeed");
8852 assert_eq!(
8853 tensor.shape(),
8854 &new_shape,
8855 "Shape should be updated after successful reshape"
8856 );
8857
8858 {
8859 let mut tensor_map = tensor.map().expect("Failed to map memory");
8860 tensor_map.fill(1);
8861 assert!(tensor_map.iter().all(|&x| x == 1));
8862 }
8863
8864 {
8865 let mut tensor_map = tensor.map().expect("Failed to map memory");
8866 tensor_map[2] = 42;
8867 assert_eq!(tensor_map[1], 1, "Value at index 1 should be 1");
8868 assert_eq!(tensor_map[2], 42, "Value at index 2 should be 42");
8869 }
8870 }
8871
8872 #[test]
8873 #[cfg(target_os = "linux")]
8874 fn test_dma_no_fd_leaks() {
8875 let _lock = fd_lock_exclusive();
8876 if !crate::test_support::dma_or_skip(function!()) {
8877 return;
8878 }
8879
8880 let baseline = open_fd_count(DMA_FD_TARGETS);
8881
8882 // Self-check — see `test_shm_no_fd_leaks`.
8883 {
8884 let probe = Tensor::<u8>::new(&[100, 100], Some(TensorMemory::DmaBuf), None)
8885 .expect("Failed to create tensor");
8886 assert!(
8887 open_fd_count(DMA_FD_TARGETS) > baseline,
8888 "a live DMA-BUF tensor must be visible to the dmabuf fd counter, \
8889 otherwise this test cannot detect a leak"
8890 );
8891 drop(probe);
8892 }
8893
8894 for _ in 0..100 {
8895 let tensor = Tensor::<u8>::new(&[100, 100], Some(TensorMemory::DmaBuf), None)
8896 .expect("Failed to create tensor");
8897 let mut map = tensor.map().unwrap();
8898 map.as_mut_slice().fill(233);
8899 }
8900
8901 let after = open_fd_count(DMA_FD_TARGETS);
8902 assert_eq!(
8903 baseline, after,
8904 "DMA-BUF file descriptor leak detected: {baseline} -> {after}"
8905 );
8906 }
8907
8908 #[test]
8909 #[cfg(target_os = "linux")]
8910 fn test_dma_from_fd_no_fd_leaks() {
8911 let _lock = fd_lock_exclusive();
8912 if !crate::test_support::dma_or_skip(function!()) {
8913 return;
8914 }
8915
8916 let baseline = open_fd_count(DMA_FD_TARGETS);
8917 let orig = Tensor::<u8>::new(&[100, 100], Some(TensorMemory::DmaBuf), None).unwrap();
8918
8919 // Self-check — see `test_shm_no_fd_leaks`.
8920 let held = open_fd_count(DMA_FD_TARGETS);
8921 assert!(
8922 held > baseline,
8923 "a live DMA-BUF tensor must be visible to the dmabuf fd counter, \
8924 otherwise this test cannot detect a leak"
8925 );
8926
8927 for _ in 0..100 {
8928 let tensor =
8929 Tensor::<u8>::from_fd(orig.clone_fd().unwrap(), orig.shape(), None).unwrap();
8930 assert_eq!(
8931 tensor.memory(),
8932 TensorMemory::DmaBuf,
8933 "DMA-BUF fd must import as Dma, not be silently downgraded"
8934 );
8935 let mut map = tensor.map().unwrap();
8936 map.as_mut_slice().fill(233);
8937 }
8938
8939 let after_loop = open_fd_count(DMA_FD_TARGETS);
8940 assert_eq!(
8941 held, after_loop,
8942 "DMA-BUF file descriptor leak detected: {held} -> {after_loop}"
8943 );
8944
8945 // And the source tensor's own fd goes back on drop.
8946 drop(orig);
8947 let after_drop = open_fd_count(DMA_FD_TARGETS);
8948 assert_eq!(
8949 baseline, after_drop,
8950 "dropping the source tensor must release its fd: {baseline} -> {after_drop}"
8951 );
8952 }
8953
8954 /// A filesystem magic must report its true 32-bit value regardless of
8955 /// how wide, and how signed, `fstatfs`'s `f_type` is on the target.
8956 ///
8957 /// `fs_type_t` is `__fsword_t` on Linux/gnu — `i64` on 64-bit but `i32`
8958 /// on 32-bit (armv7, i686, aarch64-ilp32) — `c_int` on uclibc, and
8959 /// unsigned on musl and s390x. Widening a *signed* 32-bit `f_type`
8960 /// sign-extends every magic with bit 31 set, so a naive widening cast
8961 /// reports e.g. `0xffffffff958458f6` for hugetlbfs. The docs tell
8962 /// callers to look the reported value up in `include/uapi/linux/magic.h`,
8963 /// so a sign-extended value is actively misleading.
8964 ///
8965 /// This cannot be reproduced on a 64-bit host — the defect only exists
8966 /// where `f_type` is a signed 32-bit type — so the test drives the
8967 /// conversion directly with the value such a target would produce.
8968 #[test]
8969 #[cfg(target_os = "linux")]
8970 fn test_fs_magic_normalizes_sign_extended_values() {
8971 // Declares this test as an fd-opener; see FD_LOCK.
8972 let _lock = crate::tests::fd_lock_shared();
8973 // Magics whose bit 31 is set. HUGETLBFS is the one that matters
8974 // most in practice: a MFD_HUGETLB memfd is the likeliest fd to land
8975 // in the UnknownBufferType arm.
8976 const HUGETLBFS_MAGIC: u32 = 0x9584_58f6;
8977 const F2FS_MAGIC: u32 = 0xf2f5_2010;
8978 const BTRFS_MAGIC: u32 = 0x9123_683e;
8979
8980 for magic in [HUGETLBFS_MAGIC, F2FS_MAGIC, BTRFS_MAGIC] {
8981 // 64-bit gnu: f_type is i64 and already holds the true value.
8982 assert_eq!(fs_magic(i64::from(magic)), magic);
8983
8984 // 32-bit gnu / uclibc: f_type is i32, so the value arrives
8985 // sign-extended once widened. It must still report its true
8986 // 32 bits.
8987 let sign_extended = i64::from(magic as i32);
8988 assert!(sign_extended < 0, "{magic:#x} should have bit 31 set");
8989 assert_eq!(
8990 fs_magic(sign_extended),
8991 magic,
8992 "{magic:#x} must survive a signed 32-bit f_type"
8993 );
8994 }
8995
8996 // The two magics we actually classify on are below 2^31, so they are
8997 // unaffected by signedness either way — this is why the DMA-vs-SHM
8998 // decision is correct on 32-bit even without this normalization.
8999 for magic in [DMA_BUF_MAGIC, TMPFS_MAGIC] {
9000 assert!(magic < 0x8000_0000);
9001 assert_eq!(fs_magic(i64::from(magic as i32)), magic);
9002 }
9003 }
9004
9005 /// A DMA-BUF fd must import as [`TensorMemory::DmaBuf`].
9006 ///
9007 /// Regression test for the `st_dev` minor-number classifier, which
9008 /// hardcoded `9 | 10` as "this is DMA". Those minors come from
9009 /// `get_anon_bdev()` and are assigned first-come-first-served at boot,
9010 /// so they vary by kernel build and boot order — a real DMA-BUF is
9011 /// minor 12 on x86 desktop and minor 8 on the ADIS Verdin. Both fell
9012 /// into the `_` arm and imported as SHM, which "works" (a DMA-BUF is
9013 /// mmap-able) but silently forfeits zero-copy.
9014 #[test]
9015 #[cfg(target_os = "linux")]
9016 fn test_from_fd_dma_imports_as_dma() {
9017 let _lock = crate::tests::fd_lock_shared();
9018 if !crate::test_support::dma_or_skip(function!()) {
9019 return;
9020 }
9021
9022 let orig = Tensor::<u8>::new(&[64, 64], Some(TensorMemory::DmaBuf), None).unwrap();
9023 assert_eq!(orig.memory(), TensorMemory::DmaBuf);
9024
9025 let imported = Tensor::<u8>::from_fd(orig.clone_fd().unwrap(), orig.shape(), None).unwrap();
9026
9027 assert_eq!(
9028 imported.memory(),
9029 TensorMemory::DmaBuf,
9030 "a DMA-BUF fd must import as Dma"
9031 );
9032 }
9033
9034 /// A tmpfs/SHM fd must import as [`TensorMemory::Shm`].
9035 ///
9036 /// The companion to `test_from_fd_dma_imports_as_dma`: confirms the
9037 /// magic-based classifier identifies SHM positively rather than by
9038 /// falling through.
9039 #[test]
9040 #[cfg(target_os = "linux")]
9041 fn test_from_fd_shm_imports_as_shm() {
9042 let _lock = crate::tests::fd_lock_shared();
9043 if !is_shm_available() {
9044 log::warn!("SKIPPED: {} - SHM memory not available", function!());
9045 return;
9046 }
9047
9048 let orig = Tensor::<u8>::new(&[64, 64], Some(TensorMemory::Shm), None).unwrap();
9049 assert_eq!(orig.memory(), TensorMemory::Shm);
9050
9051 let imported = Tensor::<u8>::from_fd(orig.clone_fd().unwrap(), orig.shape(), None).unwrap();
9052
9053 assert_eq!(
9054 imported.memory(),
9055 TensorMemory::Shm,
9056 "a tmpfs fd must import as Shm"
9057 );
9058 }
9059
9060 /// An fd that is neither a DMA-BUF nor tmpfs must be rejected.
9061 ///
9062 /// A pipe is the convenient probe: it lives on `pipefs`, another
9063 /// `get_anon_bdev()` pseudo-filesystem, so it shares major 0 with the
9064 /// buffer types we do support and is only distinguishable by magic.
9065 /// Importing one as SHM is meaningless — `mmap` on a pipe fails — so
9066 /// the classifier must say "unknown" rather than guess.
9067 #[test]
9068 #[cfg(target_os = "linux")]
9069 fn test_from_fd_rejects_unknown_filesystem() {
9070 let _lock = crate::tests::fd_lock_shared();
9071
9072 let (read_end, _write_end) = nix::unistd::pipe().unwrap();
9073
9074 let result = Tensor::<u8>::from_fd(read_end, &[64], None);
9075
9076 match result {
9077 Err(Error::UnknownBufferType(magic)) => {
9078 // PIPEFS_MAGIC, from include/uapi/linux/magic.h
9079 assert_eq!(magic, 0x5049_5045, "expected PIPEFS_MAGIC");
9080 }
9081 other => panic!("expected UnknownBufferType for a pipe fd, got {other:?}"),
9082 }
9083 }
9084
9085 #[test]
9086 #[cfg(target_os = "linux")]
9087 fn test_shm_no_fd_leaks() {
9088 let _lock = fd_lock_exclusive();
9089 if !is_shm_available() {
9090 log::warn!(
9091 "SKIPPED: {} - SHM memory allocation not available (permission denied or no SHM support)",
9092 function!()
9093 );
9094 return;
9095 }
9096
9097 let baseline = open_fd_count(SHM_FD_TARGETS);
9098
9099 // Self-check: the counter must actually see a live shm tensor. If
9100 // it does not — a changed backing object, a different `/proc`
9101 // rendering — the leak assertion below would hold no matter how
9102 // much leaked, so fail loudly here instead.
9103 {
9104 let probe = Tensor::<u8>::new(&[100, 100], Some(TensorMemory::Shm), None)
9105 .expect("Failed to create tensor");
9106 assert!(
9107 open_fd_count(SHM_FD_TARGETS) > baseline,
9108 "a live shm tensor must be visible to the shm fd counter, \
9109 otherwise this test cannot detect a leak"
9110 );
9111 drop(probe);
9112 }
9113
9114 for _ in 0..100 {
9115 let tensor = Tensor::<u8>::new(&[100, 100], Some(TensorMemory::Shm), None)
9116 .expect("Failed to create tensor");
9117 let mut map = tensor.map().unwrap();
9118 map.as_mut_slice().fill(233);
9119 }
9120
9121 let after = open_fd_count(SHM_FD_TARGETS);
9122 assert_eq!(
9123 baseline, after,
9124 "shm file descriptor leak detected: {baseline} -> {after}"
9125 );
9126 }
9127
9128 #[test]
9129 #[cfg(target_os = "linux")]
9130 fn test_shm_from_fd_no_fd_leaks() {
9131 let _lock = fd_lock_exclusive();
9132 if !is_shm_available() {
9133 log::warn!(
9134 "SKIPPED: {} - SHM memory allocation not available (permission denied or no SHM support)",
9135 function!()
9136 );
9137 return;
9138 }
9139
9140 let baseline = open_fd_count(SHM_FD_TARGETS);
9141 let orig = Tensor::<u8>::new(&[100, 100], Some(TensorMemory::Shm), None).unwrap();
9142
9143 // Self-check — see `test_shm_no_fd_leaks`.
9144 let held = open_fd_count(SHM_FD_TARGETS);
9145 assert!(
9146 held > baseline,
9147 "a live shm tensor must be visible to the shm fd counter, \
9148 otherwise this test cannot detect a leak"
9149 );
9150
9151 for _ in 0..100 {
9152 let tensor =
9153 Tensor::<u8>::from_fd(orig.clone_fd().unwrap(), orig.shape(), None).unwrap();
9154 let mut map = tensor.map().unwrap();
9155 map.as_mut_slice().fill(233);
9156 }
9157
9158 let after_loop = open_fd_count(SHM_FD_TARGETS);
9159 assert_eq!(
9160 held, after_loop,
9161 "shm file descriptor leak detected: {held} -> {after_loop}"
9162 );
9163
9164 // And the source tensor's own fd goes back on drop.
9165 drop(orig);
9166 let after_drop = open_fd_count(SHM_FD_TARGETS);
9167 assert_eq!(
9168 baseline, after_drop,
9169 "dropping the source tensor must release its fd: {baseline} -> {after_drop}"
9170 );
9171 }
9172
9173 #[cfg(feature = "ndarray")]
9174 #[test]
9175 fn test_ndarray() {
9176 let _lock = crate::tests::fd_lock_shared();
9177 let shape = vec![2, 3, 4];
9178 let tensor = Tensor::<f32>::new(&shape, None, None).expect("Failed to create tensor");
9179
9180 let mut tensor_map = tensor.map().expect("Failed to map tensor memory");
9181 tensor_map.fill(1.0);
9182
9183 let view = tensor_map.view().expect("Failed to get ndarray view");
9184 assert_eq!(view.shape(), &[2, 3, 4]);
9185 assert!(view.iter().all(|&x| x == 1.0));
9186
9187 let mut view_mut = tensor_map
9188 .view_mut()
9189 .expect("Failed to get mutable ndarray view");
9190 view_mut[[0, 0, 0]] = 42.0;
9191 assert_eq!(view_mut[[0, 0, 0]], 42.0);
9192 assert_eq!(tensor_map[0], 42.0, "Value at index 0 should be 42");
9193 }
9194
9195 #[test]
9196 fn test_buffer_identity_derived_equality() {
9197 // Two identities derived from the same kind and key must be equal —
9198 // that is the whole point of `derived`: two independently-linked
9199 // copies of this crate, each importing the same underlying buffer,
9200 // have to agree on its id without talking to each other.
9201 let id1 = BufferIdentity::derived(IdentityKind::DmaBuf, 4242);
9202 let id2 = BufferIdentity::derived(IdentityKind::DmaBuf, 4242);
9203 assert_eq!(
9204 id1.id(),
9205 id2.id(),
9206 "Identities derived from the same kind and key should match"
9207 );
9208
9209 // A different key still yields a different id.
9210 let id3 = BufferIdentity::derived(IdentityKind::DmaBuf, 4243);
9211 assert_ne!(
9212 id1.id(),
9213 id3.id(),
9214 "Identities derived from different keys should differ"
9215 );
9216 }
9217
9218 #[test]
9219 fn test_buffer_identity_clone_shares_guard() {
9220 let id1 = BufferIdentity::derived(IdentityKind::DmaBuf, 4242);
9221 let weak = id1.weak();
9222 assert!(
9223 weak.upgrade().is_some(),
9224 "Weak should be alive while original exists"
9225 );
9226
9227 let id2 = id1.clone();
9228 assert_eq!(id1.id(), id2.id(), "Cloned identity should have same id");
9229
9230 drop(id1);
9231 assert!(
9232 weak.upgrade().is_some(),
9233 "Weak should still be alive (clone holds Arc)"
9234 );
9235
9236 drop(id2);
9237 assert!(
9238 weak.upgrade().is_none(),
9239 "Weak should be dead after all clones dropped"
9240 );
9241 }
9242
9243 #[test]
9244 fn test_tensor_buffer_identity() {
9245 let t1 = Tensor::<u8>::new(&[100], Some(TensorMemory::Mem), Some("t1")).unwrap();
9246 let t2 = Tensor::<u8>::new(&[100], Some(TensorMemory::Mem), Some("t2")).unwrap();
9247 assert_ne!(
9248 t1.buffer_identity().id(),
9249 t2.buffer_identity().id(),
9250 "Different tensors should have different buffer ids"
9251 );
9252 }
9253
9254 // ------------------------------------------------------------------------
9255 // Quantization — constructor validation + accessor correctness.
9256 // ------------------------------------------------------------------------
9257
9258 #[test]
9259 fn test_quantization_per_tensor_constructors() {
9260 let q = Quantization::per_tensor(0.1, -5);
9261 assert!(q.is_per_tensor());
9262 assert!(!q.is_per_channel());
9263 assert!(!q.is_symmetric());
9264 assert_eq!(q.scale(), &[0.1]);
9265 assert_eq!(q.zero_point(), Some(&[-5][..]));
9266
9267 let qs = Quantization::per_tensor_symmetric(0.05);
9268 assert!(qs.is_per_tensor());
9269 assert!(qs.is_symmetric());
9270 assert_eq!(qs.zero_point(), None);
9271 }
9272
9273 #[test]
9274 fn test_quantization_per_channel_constructors() {
9275 let q = Quantization::per_channel(vec![0.1, 0.2, 0.3], vec![0, -1, 1], 2).unwrap();
9276 assert!(q.is_per_channel());
9277 assert!(!q.is_symmetric());
9278 assert_eq!(q.axis(), Some(2));
9279 assert_eq!(q.scale().len(), 3);
9280
9281 let qs = Quantization::per_channel_symmetric(vec![0.054, 0.089, 0.195], 0).unwrap();
9282 assert!(qs.is_per_channel());
9283 assert!(qs.is_symmetric());
9284 assert_eq!(qs.axis(), Some(0));
9285 }
9286
9287 #[test]
9288 fn test_quantization_per_channel_length_mismatch_rejected() {
9289 // len(scales) != len(zero_points) → rejected at construction.
9290 let err = Quantization::per_channel(vec![0.1, 0.2], vec![0, 0, 0], 0).unwrap_err();
9291 assert!(matches!(err, Error::QuantizationInvalid { .. }));
9292 }
9293
9294 #[test]
9295 fn test_quantization_per_channel_empty_rejected() {
9296 let err = Quantization::per_channel_symmetric(vec![], 0).unwrap_err();
9297 assert!(matches!(err, Error::QuantizationInvalid { .. }));
9298 }
9299
9300 /// Constructors guard scale/zero_point length invariants, but
9301 /// `Quantization` is `Deserialize`, so malformed JSON (e.g. an
9302 /// empty `scale` array, or `zero_point` length that disagrees with
9303 /// `scale`) bypasses the constructor checks. `set_quantization`
9304 /// must reject these via `validate()` so they don't poison
9305 /// downstream `mode()` selection or per-channel kernel indexing.
9306 #[test]
9307 fn test_quantization_validate_rejects_malformed_deserialize() {
9308 let mut t = Tensor::<i8>::new(&[1, 1, 4], Some(TensorMemory::Mem), None).unwrap();
9309
9310 // Empty scale array: must be rejected.
9311 let q: Quantization = serde_json::from_str(r#"{"scale": []}"#).unwrap();
9312 assert!(matches!(
9313 t.set_quantization(q).unwrap_err(),
9314 Error::QuantizationInvalid { .. }
9315 ));
9316
9317 // Per-tensor with multi-element zero_point: must be rejected.
9318 let q: Quantization =
9319 serde_json::from_str(r#"{"scale": 0.1, "zero_point": [0, 0, 0]}"#).unwrap();
9320 assert!(matches!(
9321 t.set_quantization(q).unwrap_err(),
9322 Error::QuantizationInvalid { .. }
9323 ));
9324
9325 // Per-channel zero_point length != scale length: must be rejected.
9326 let q: Quantization = serde_json::from_str(
9327 r#"{"scale": [0.1, 0.2, 0.3, 0.4], "zero_point": [0, 0], "axis": 2}"#,
9328 )
9329 .unwrap();
9330 assert!(matches!(
9331 t.set_quantization(q).unwrap_err(),
9332 Error::QuantizationInvalid { .. }
9333 ));
9334 }
9335
9336 #[test]
9337 fn test_quantization_mode_dispatch() {
9338 let pt = Quantization::per_tensor(0.1, -5);
9339 assert!(matches!(
9340 pt.mode(),
9341 QuantMode::PerTensor { scale, zero_point } if scale == 0.1 && zero_point == -5
9342 ));
9343
9344 let pts = Quantization::per_tensor_symmetric(0.05);
9345 assert!(matches!(
9346 pts.mode(),
9347 QuantMode::PerTensorSymmetric { scale } if scale == 0.05
9348 ));
9349
9350 let pc = Quantization::per_channel(vec![0.1, 0.2], vec![0, -1], 2).unwrap();
9351 assert!(matches!(pc.mode(), QuantMode::PerChannel { axis: 2, .. }));
9352
9353 let pcs = Quantization::per_channel_symmetric(vec![0.1, 0.2], 0).unwrap();
9354 assert!(matches!(
9355 pcs.mode(),
9356 QuantMode::PerChannelSymmetric { axis: 0, .. }
9357 ));
9358 }
9359
9360 #[test]
9361 fn test_tensor_quantization_roundtrip_integer() {
9362 let mut t = Tensor::<i8>::new(&[2, 3, 4], Some(TensorMemory::Mem), None).unwrap();
9363 assert!(t.quantization().is_none());
9364 t.set_quantization(Quantization::per_tensor(0.1, -5))
9365 .unwrap();
9366 let q = t.quantization().unwrap();
9367 assert_eq!(q.scale(), &[0.1]);
9368 t.clear_quantization();
9369 assert!(t.quantization().is_none());
9370 }
9371
9372 #[test]
9373 fn test_tensor_with_quantization_builder() {
9374 let t = Tensor::<i8>::new(&[4, 4], Some(TensorMemory::Mem), None)
9375 .unwrap()
9376 .with_quantization(Quantization::per_tensor_symmetric(0.05))
9377 .unwrap();
9378 assert!(t.quantization().is_some());
9379 }
9380
9381 #[test]
9382 fn test_tensor_dyn_quantization_float_arm_returns_none() {
9383 let t = Tensor::<f32>::new(&[2, 2], Some(TensorMemory::Mem), None).unwrap();
9384 let td = TensorDyn::F32(t);
9385 assert!(td.quantization().is_none());
9386 }
9387
9388 #[test]
9389 fn test_tensor_dyn_set_quantization_float_arm_errors() {
9390 let t = Tensor::<f32>::new(&[2, 2], Some(TensorMemory::Mem), None).unwrap();
9391 let mut td = TensorDyn::F32(t);
9392 let err = td
9393 .set_quantization(Quantization::per_tensor(0.1, 0))
9394 .unwrap_err();
9395 // float path returns a QuantizationInvalid error.
9396 assert!(matches!(err, Error::QuantizationInvalid { .. }));
9397 }
9398
9399 /// Compile-time type gate — calling `Tensor::<f32>::quantization()` must
9400 /// fail to compile (the `IntegerType` trait bound is not satisfied by
9401 /// `f32`). This doctest anchors the invariant.
9402 ///
9403 /// ```compile_fail
9404 /// use edgefirst_tensor::{Tensor, TensorMemory};
9405 /// let t = Tensor::<f32>::new(&[2, 2], Some(TensorMemory::Mem), None).unwrap();
9406 /// let _ = t.quantization(); // compile error: f32 not IntegerType
9407 /// ```
9408 fn _compile_fail_doctest_anchor() {}
9409
9410 // Any test that cares about the fd count must grab it exclusively.
9411 // Any tests which modifies the fd count by opening or closing fds must grab it
9412 // shared.
9413 pub static FD_LOCK: RwLock<()> = RwLock::new(());
9414
9415 /// Take [`FD_LOCK`] shared, tolerating poisoning.
9416 ///
9417 /// A test that panics while holding the lock poisons it, and every later
9418 /// `.unwrap()` on it panics too — one real failure turns into a handful
9419 /// of unrelated ones that bury it. The guarded data is `()`, so a panic
9420 /// cannot have left an invariant broken and there is nothing to protect
9421 /// against by refusing the lock.
9422 pub fn fd_lock_shared() -> std::sync::RwLockReadGuard<'static, ()> {
9423 FD_LOCK.read().unwrap_or_else(PoisonError::into_inner)
9424 }
9425
9426 /// Take [`FD_LOCK`] exclusively. Poison-tolerant, see [`fd_lock_shared`].
9427 ///
9428 /// Linux-only, like all four of its callers: the fd-leak tests read fd
9429 /// state out of `/proc`, which no other platform provides. Ungated, it
9430 /// is dead code on macOS and the lint denies that.
9431 /// [`fd_lock_shared`] stays available everywhere — macOS tests declare
9432 /// fd activity too.
9433 #[cfg(target_os = "linux")]
9434 pub fn fd_lock_exclusive() -> std::sync::RwLockWriteGuard<'static, ()> {
9435 FD_LOCK.write().unwrap_or_else(PoisonError::into_inner)
9436 }
9437
9438 /// `/proc` fd-target fragments identifying a POSIX shared memory
9439 /// segment. `shm_open` then `shm_unlink` leaves the target rendered as
9440 /// `/dev/shm/<name> (deleted)`.
9441 #[cfg(target_os = "linux")]
9442 const SHM_FD_TARGETS: &[&str] = &["/dev/shm/"];
9443
9444 /// `/proc` fd-target fragments identifying the fds a DMA tensor owns.
9445 ///
9446 /// Both the exported buffer and the heap it came from are listed, so a
9447 /// leak of either is still caught — narrowing this to the buffer alone
9448 /// would silently stop covering leaked heap handles.
9449 #[cfg(target_os = "linux")]
9450 const DMA_FD_TARGETS: &[&str] = &["dmabuf", "dma_heap"];
9451
9452 /// How many of this process's open fds point at an object whose `/proc`
9453 /// target contains `needle`.
9454 ///
9455 /// The fd-leak tests deliberately do NOT compare total fd counts. Under
9456 /// `cargo test` every test in this crate runs as a thread in one
9457 /// process, and roughly twenty of them create shm or DMA tensors —
9458 /// while anything else may open a file, a socket, or a pipe — so an
9459 /// unrelated test opening or closing an fd inside the measurement
9460 /// window moved the total and failed the assertion. Worse, the panic
9461 /// poisoned [`FD_LOCK`] and took four more tests down with it, so the
9462 /// visible failure had little to do with the cause.
9463 ///
9464 /// Counting only fds that name the object under test makes the
9465 /// measurement immune to every unrelated fd, while still catching the
9466 /// leak each test exists to catch. (CI runs nextest, one process per
9467 /// test, which is why this only ever flaked when run locally.)
9468 #[cfg(target_os = "linux")]
9469 fn open_fd_count(needles: &[&str]) -> usize {
9470 let proc = procfs::process::Process::myself()
9471 .expect("Failed to get current process using /proc/self");
9472 proc.fd()
9473 .expect("Failed to list open file descriptors")
9474 .filter_map(Result::ok)
9475 .filter(|entry| {
9476 let target = format!("{:?}", entry.target);
9477 needles.iter().any(|needle| target.contains(needle))
9478 })
9479 .count()
9480 }
9481
9482 /// Test that DMA is NOT available on non-Linux platforms.
9483 /// This verifies the cross-platform behavior of is_dma_available().
9484 #[test]
9485 #[cfg(not(target_os = "linux"))]
9486 fn test_dma_not_available_on_non_linux() {
9487 assert!(
9488 !is_dma_available(),
9489 "DMA memory allocation should NOT be available on non-Linux platforms"
9490 );
9491 }
9492
9493 #[test]
9494 fn colorimetry_defaults_none_and_roundtrips_without_auto_fill() {
9495 use crate::{ColorEncoding, ColorRange, Colorimetry, PixelFormat, TensorMemory};
9496 let mut t = Tensor::<u8>::image(
9497 1280,
9498 720,
9499 PixelFormat::Nv12,
9500 Some(TensorMemory::Mem),
9501 crate::CpuAccess::ReadWrite,
9502 )
9503 .unwrap();
9504 assert_eq!(t.colorimetry(), None); // default undefined
9505 let c = Colorimetry::default()
9506 .with_encoding(ColorEncoding::Bt709)
9507 .with_range(ColorRange::Limited);
9508 t.set_colorimetry(Some(c));
9509 assert_eq!(t.colorimetry(), Some(c));
9510 // configure_image must NOT touch colorimetry
9511 t.configure_image(640, 480, PixelFormat::Grey).unwrap();
9512 assert_eq!(t.colorimetry(), Some(c));
9513 }
9514
9515 #[test]
9516 fn configure_image_within_capacity() {
9517 let mut t = Tensor::<u8>::image_with_capacity(
9518 640,
9519 480,
9520 PixelFormat::Rgb,
9521 None,
9522 crate::CpuAccess::ReadWrite,
9523 )
9524 .unwrap();
9525 t.configure_image(320, 240, PixelFormat::Nv12).unwrap();
9526 assert_eq!(t.format(), Some(PixelFormat::Nv12));
9527 assert_eq!(t.width(), Some(320));
9528 assert_eq!(t.height(), Some(240));
9529 assert_eq!(t.shape(), &[360, 320]); // 240*3/2
9530 }
9531
9532 #[test]
9533 fn configure_image_too_large_errors() {
9534 let mut t = Tensor::<u8>::image_with_capacity(
9535 64,
9536 64,
9537 PixelFormat::Grey,
9538 None,
9539 crate::CpuAccess::ReadWrite,
9540 )
9541 .unwrap();
9542 let err = t
9543 .configure_image(1920, 1080, PixelFormat::Nv12)
9544 .unwrap_err();
9545 assert!(matches!(err, Error::InsufficientCapacity { .. }));
9546 }
9547
9548 /// A reused max-sized IOSurface pool keeps its physical `bytesPerRow` when
9549 /// reconfigured to a smaller logical image (physical-grid / logical-ROI
9550 /// decoupling), instead of collapsing to the frame's natural row stride.
9551 #[test]
9552 #[cfg(target_os = "macos")]
9553 fn configure_image_preserves_iosurface_physical_stride() {
9554 // Declares this test as an fd-opener; see FD_LOCK.
9555 let _lock = crate::tests::fd_lock_shared();
9556 // Pool: GREY/R8 IOSurface 100 wide → bytesPerRow padded to 128.
9557 let mut pool = Tensor::<u8>::image(
9558 100,
9559 64,
9560 PixelFormat::Grey,
9561 Some(TensorMemory::DmaBuf),
9562 crate::CpuAccess::ReadWrite,
9563 )
9564 .unwrap();
9565 let pitch = pool.effective_row_stride().unwrap();
9566 assert!(
9567 pitch >= 128 && pitch.is_multiple_of(64),
9568 "padded bytesPerRow, got {pitch}"
9569 );
9570
9571 // Reconfigure to a smaller NV12 frame; the physical pitch must survive
9572 // (natural would be 32, but the surface stride is the 128-padded pitch).
9573 pool.configure_image(32, 16, PixelFormat::Nv12).unwrap();
9574 assert_eq!(pool.format(), Some(PixelFormat::Nv12));
9575 assert_eq!(pool.width(), Some(32));
9576 assert_eq!(pool.height(), Some(16));
9577 assert_eq!(
9578 pool.effective_row_stride(),
9579 Some(pitch),
9580 "configure_image must preserve the IOSurface physical bytesPerRow"
9581 );
9582
9583 // Reconfigure again to NV24 — pitch still preserved.
9584 pool.configure_image(32, 16, PixelFormat::Nv24).unwrap();
9585 assert_eq!(pool.effective_row_stride(), Some(pitch));
9586 }
9587
9588 /// `configure_image` on a Mem backing reconfigures to the format's
9589 /// **64-byte-aligned** row stride (the odd-dim contract: every image tensor
9590 /// carries a 64-aligned `row_stride`). For NV12 32×16 the minimum is
9591 /// `even(32)=32`, rounded up to the 64-byte alignment → 64. The capacity
9592 /// (64×64×4 RGBA = 16 KiB) easily holds the 24×64 = 1.5 KiB NV12 layout.
9593 #[test]
9594 fn configure_image_mem_aligns_stride() {
9595 let mut t = Tensor::<u8>::image_with_capacity(
9596 64,
9597 64,
9598 PixelFormat::Rgba,
9599 Some(TensorMemory::Mem),
9600 crate::CpuAccess::ReadWrite,
9601 )
9602 .unwrap();
9603 t.configure_image(32, 16, PixelFormat::Nv12).unwrap();
9604 let s = t.effective_row_stride().unwrap();
9605 assert_eq!(s % 64, 0, "stride must be 64-aligned");
9606 assert!(s >= 32, "stride must cover the even-width minimum");
9607 assert_eq!(s, 64);
9608 }
9609
9610 #[test]
9611 fn strided_mem_tensor_cpu_maps_full_padded_buffer() {
9612 // A packed RGBA image with row padding (GPU-pitch style): logical width
9613 // 8 px (32 B/row) but a 48-byte row stride. Over-allocate capacity (for
9614 // 16 px), narrow the logical width, then record the padded stride.
9615 // Previously `map()` rejected this on non-Linux with
9616 // "DMA backing is Linux-only"; HAL-owned Mem is now mappable.
9617 let mut t = Tensor::<u8>::image_with_capacity(
9618 16,
9619 3,
9620 PixelFormat::Rgba,
9621 Some(TensorMemory::Mem),
9622 crate::CpuAccess::ReadWrite,
9623 )
9624 .unwrap(); // capacity 3 × 16 × 4 = 192 B
9625 t.configure_image(8, 3, PixelFormat::Rgba).unwrap(); // logical [3, 8, 4] = 96 B
9626 t.set_row_stride(48).unwrap(); // padded stride (>= 32 B min)
9627
9628 let map = t.map().expect("strided Mem tensor should CPU-map");
9629 // Full padded buffer (stride 48 × 3 rows = 144 B), not the 96 B logical
9630 // view — callers iterate rows via `effective_row_stride()`.
9631 assert_eq!(map.as_slice().len(), 144);
9632 // Logical shape is still reported for shape-aware consumers.
9633 assert_eq!(map.shape(), &[3, 8, 4]);
9634 }
9635
9636 #[test]
9637 fn strided_mem_tensor_over_capacity_errors() {
9638 // Stride larger than the allocation: 64 B × 3 rows = 192 B > 96 B cap.
9639 let mut t = Tensor::<u8>::new(&[3, 8, 4], Some(TensorMemory::Mem), None).unwrap();
9640 t.set_format(PixelFormat::Rgba).unwrap();
9641 t.set_row_stride(64).unwrap();
9642 assert!(matches!(t.map(), Err(Error::InsufficientCapacity { .. })));
9643 }
9644
9645 /// Test that SHM memory allocation is available and usable on Unix systems.
9646 /// This is a basic functional test; Linux has additional FD leak tests using procfs.
9647 #[test]
9648 #[cfg(unix)]
9649 fn test_shm_available_and_usable() {
9650 // Declares this test as an fd-opener; see FD_LOCK.
9651 let _lock = crate::tests::fd_lock_shared();
9652 // "Unix" is not one platform here: Android is unix but its bionic libc
9653 // has no POSIX shared memory, so allocation legitimately reports
9654 // NotImplemented (see shm.rs module docs). Asserting availability
9655 // outright fails there for a documented reason.
9656 if !is_shm_available() {
9657 log::warn!(
9658 "SKIPPED: test_shm_available_and_usable - SHM allocation \
9659 unavailable on this platform"
9660 );
9661 return;
9662 }
9663
9664 // Create a tensor with SHM backing
9665 let tensor = Tensor::<u8>::new(&[100, 100], Some(TensorMemory::Shm), None)
9666 .expect("Failed to create SHM tensor");
9667
9668 // Verify we can map and write to it
9669 let mut map = tensor.map().expect("Failed to map SHM tensor");
9670 map.as_mut_slice().fill(0xAB);
9671
9672 // Verify the data was written correctly
9673 assert!(
9674 map.as_slice().iter().all(|&b| b == 0xAB),
9675 "SHM tensor data should be writable and readable"
9676 );
9677 }
9678
9679 // =========================================================================
9680 // packed_rgba16f_layout — host-runnable geometry unit tests (TDD)
9681 // =========================================================================
9682
9683 #[test]
9684 fn packed_rgba16f_layout_planar_rgb_f16() {
9685 let layout =
9686 packed_rgba16f_layout(PixelFormat::PlanarRgb, DType::F16, 640, 640).expect("Some");
9687 assert_eq!(layout.surface_w, 160);
9688 assert_eq!(layout.surface_h, 1920);
9689 assert_eq!(layout.bytes_per_texel, 8);
9690 assert_eq!(layout.pitch, 1280);
9691 }
9692
9693 #[test]
9694 fn packed_rgba16f_layout_planar_rgba_f16() {
9695 let layout =
9696 packed_rgba16f_layout(PixelFormat::PlanarRgba, DType::F16, 640, 640).expect("Some");
9697 assert_eq!(layout.surface_w, 160);
9698 assert_eq!(layout.surface_h, 2560); // 4 planes
9699 assert_eq!(layout.bytes_per_texel, 8);
9700 assert_eq!(layout.pitch, 1280);
9701 }
9702
9703 #[test]
9704 fn packed_rgba16f_layout_rejects_misaligned() {
9705 assert!(packed_rgba16f_layout(PixelFormat::PlanarRgb, DType::F16, 642, 640).is_none());
9706 }
9707
9708 #[test]
9709 fn packed_rgba16f_layout_rejects_non_f16() {
9710 // Non-F16 dtype with planar RGB
9711 assert!(packed_rgba16f_layout(PixelFormat::PlanarRgb, DType::U8, 640, 640).is_none());
9712 // Non-planar format with F32
9713 assert!(packed_rgba16f_layout(PixelFormat::Rgb, DType::F32, 640, 640).is_none());
9714 // Packed Rgba with F16 is not a planar format → None
9715 assert!(packed_rgba16f_layout(PixelFormat::Rgba, DType::F16, 640, 640).is_none());
9716 }
9717
9718 #[test]
9719 fn cuda_map_fast_fails_to_none_without_handle() {
9720 let t = Tensor::<f32>::new(&[4], Some(TensorMemory::Mem), None).unwrap();
9721 assert!(t.cuda().is_none());
9722 assert!(t.cuda_map().is_none()); // pure local check, no GL routing
9723 }
9724
9725 #[test]
9726 fn cuda_returns_none_without_handle() {
9727 // A plain Mem-backed tensor has no CUDA handle attached.
9728 let t = Tensor::<f32>::new(&[2, 2], Some(TensorMemory::Mem), None).unwrap();
9729 assert!(t.cuda().is_none(), "no CUDA handle on a Mem tensor");
9730 assert!(t.cuda_map().is_none(), "fast-fail map → None");
9731 }
9732
9733 #[test]
9734 fn cuda_map_then_host_map_fallback() {
9735 // The documented client pattern: try cuda_map() first; when it is None
9736 // (no CUDA handle — the case for a plain Mem tensor), fall back to map().
9737 let t = Tensor::<f32>::new(&[2, 2], Some(TensorMemory::Mem), None).unwrap();
9738 // Bind to a named variable so the CudaMap guard (and its borrow of `t`)
9739 // is dropped at the end of this statement, before the else branch borrows `t` again.
9740 let cuda = t.cuda_map();
9741 if let Some(_c) = cuda {
9742 // On a CUDA-registered tensor we'd use the device ptr here.
9743 unreachable!("a Mem tensor has no CUDA handle");
9744 } else {
9745 let host = t.map().expect("host map fallback must succeed");
9746 // TensorMapTrait::len() returns the element count (not bytes).
9747 assert_eq!(host.len(), 4); // 2*2 f32 elements
9748 }
9749 }
9750
9751 // -------------------------------------------------------------------------
9752 // Tensor::from_foreign — public API tests at the Tensor<T> layer.
9753 //
9754 // The low-level MemTensor::from_foreign mechanics (owner-drop, view sharing)
9755 // are covered in mem.rs. These tests exercise the Tensor<T> guard paths
9756 // (null ptr, empty shape, size overflow) and the basic wrap+readback
9757 // contract, confirming the public unsafe API wires through correctly.
9758 // -------------------------------------------------------------------------
9759
9760 #[test]
9761 fn from_foreign_valid_wrap_and_readback() {
9762 // The canonical CUDA zero-copy shape: wrap a caller allocation as a
9763 // Mem tensor and verify the tensor reads the exact same bytes.
9764 let mut buf: Vec<f32> = vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0];
9765 let ptr = buf.as_mut_ptr();
9766 let t = unsafe { Tensor::<f32>::from_foreign(ptr, &[2, 3], None, Some("test_foreign")) }
9767 .expect("valid from_foreign must succeed");
9768 assert_eq!(t.shape(), &[2, 3]);
9769 assert_eq!(t.memory(), TensorMemory::Mem);
9770 assert_eq!(t.name(), "test_foreign");
9771 let m = t.map().unwrap();
9772 // The tensor is a zero-copy borrow — it sees the caller's data.
9773 assert_eq!(m.as_slice(), &[1.0_f32, 2.0, 3.0, 4.0, 5.0, 6.0]);
9774 }
9775
9776 #[test]
9777 fn from_foreign_write_visible_in_caller_allocation() {
9778 // Writes through the tensor's map land in the caller's buffer (zero-copy).
9779 let mut buf: Vec<u8> = vec![0u8; 6];
9780 let ptr = buf.as_mut_ptr();
9781 let t = unsafe { Tensor::<u8>::from_foreign(ptr, &[2, 3], None, None) }.unwrap();
9782 {
9783 let mut m = t.map().unwrap();
9784 m.as_mut_slice().copy_from_slice(&[10, 20, 30, 40, 50, 60]);
9785 }
9786 drop(t);
9787 // Mutations are visible in the original Vec — same physical buffer.
9788 assert_eq!(buf, vec![10, 20, 30, 40, 50, 60]);
9789 }
9790
9791 #[test]
9792 fn from_foreign_rejects_null_ptr() {
9793 let err = unsafe { Tensor::<u8>::from_foreign(std::ptr::null_mut(), &[4], None, None) }
9794 .unwrap_err();
9795 assert!(
9796 matches!(err, Error::InvalidArgument(ref m) if m.contains("non-null")),
9797 "expected InvalidArgument(non-null), got {err:?}"
9798 );
9799 }
9800
9801 #[test]
9802 fn from_foreign_rejects_empty_shape() {
9803 let mut dummy: u8 = 0;
9804 let err = unsafe { Tensor::<u8>::from_foreign(&mut dummy, &[], None, None) }.unwrap_err();
9805 assert!(
9806 matches!(err, Error::InvalidSize(0)),
9807 "expected InvalidSize(0) for empty shape, got {err:?}"
9808 );
9809 }
9810
9811 #[test]
9812 fn from_foreign_rejects_overflow_shape() {
9813 // Two dimensions whose product overflows usize — the overflow guard must
9814 // fire before any pointer arithmetic is attempted.
9815 let mut dummy: u8 = 0;
9816 let huge = [usize::MAX / 2 + 1, 2];
9817 let err = unsafe { Tensor::<u8>::from_foreign(&mut dummy, &huge, None, None) }.unwrap_err();
9818 assert!(
9819 matches!(err, Error::InvalidArgument(ref m) if m.contains("overflow")),
9820 "expected InvalidArgument(overflow), got {err:?}"
9821 );
9822 }
9823
9824 #[test]
9825 fn from_foreign_owner_keeps_allocation_alive() {
9826 // When `owner` is `Some`, dropping the Tensor must not free the backing
9827 // before the owner is also gone — the owner's Drop fires on last ref.
9828 use std::sync::atomic::{AtomicBool, Ordering};
9829 let flag = std::sync::Arc::new(AtomicBool::new(false));
9830 let flag2 = flag.clone();
9831 struct Guard(std::sync::Arc<AtomicBool>);
9832 impl Drop for Guard {
9833 fn drop(&mut self) {
9834 self.0.store(true, Ordering::SeqCst);
9835 }
9836 }
9837 let mut buf: Vec<u32> = vec![42u32; 4];
9838 let ptr = buf.as_mut_ptr();
9839 let owner: ForeignOwner = Box::new(Guard(flag2));
9840 let t = unsafe { Tensor::<u32>::from_foreign(ptr, &[4], Some(owner), None) }.unwrap();
9841 // Map co-owns the backing Arc; the owner must stay alive while the map lives.
9842 let m = t.map().unwrap();
9843 assert_eq!(m.as_slice()[0], 42);
9844 drop(t); // tensor dropped while map is still live
9845 assert!(
9846 !flag.load(Ordering::SeqCst),
9847 "owner must not drop while a map shares the backing"
9848 );
9849 drop(m);
9850 assert!(
9851 flag.load(Ordering::SeqCst),
9852 "owner Drop must fire when the last Arc reference is released"
9853 );
9854 }
9855}