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 gl_map_resource, gl_register_buffer, gl_unmap_resource, gl_unregister_resource,
189 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 fn expecting(&self, f: &mut fmt::Formatter) -> fmt::Result {
1082 f.write_str("f32 or array of f32")
1083 }
1084 fn visit_f64<E: de::Error>(self, v: f64) -> std::result::Result<Self::Value, E> {
1085 Ok(vec![v as f32])
1086 }
1087 #[allow(clippy::cast_possible_truncation)]
1088 fn visit_i64<E: de::Error>(self, v: i64) -> std::result::Result<Self::Value, E> {
1089 Ok(vec![v as f32])
1090 }
1091 #[allow(clippy::cast_possible_truncation, clippy::cast_precision_loss)]
1092 fn visit_u64<E: de::Error>(self, v: u64) -> std::result::Result<Self::Value, E> {
1093 Ok(vec![v as f32])
1094 }
1095 fn visit_seq<A: de::SeqAccess<'de>>(
1096 self,
1097 mut seq: A,
1098 ) -> std::result::Result<Self::Value, A::Error> {
1099 let mut out = Vec::with_capacity(seq.size_hint().unwrap_or(1));
1100 while let Some(x) = seq.next_element::<f32>()? {
1101 out.push(x);
1102 }
1103 Ok(out)
1104 }
1105 }
1106 de.deserialize_any(V)
1107}
1108
1109fn deserialize_opt_scalar_or_vec_i32<'de, D: serde::Deserializer<'de>>(
1110 de: D,
1111) -> std::result::Result<Option<Vec<i32>>, D::Error> {
1112 use serde::de::{self, Visitor};
1113 struct V;
1114 impl<'de> Visitor<'de> for V {
1115 type Value = Option<Vec<i32>>;
1116 fn expecting(&self, f: &mut fmt::Formatter) -> fmt::Result {
1117 f.write_str("null, i32, or array of i32")
1118 }
1119 fn visit_none<E: de::Error>(self) -> std::result::Result<Self::Value, E> {
1120 Ok(None)
1121 }
1122 fn visit_unit<E: de::Error>(self) -> std::result::Result<Self::Value, E> {
1123 Ok(None)
1124 }
1125 fn visit_some<D2: serde::Deserializer<'de>>(
1126 self,
1127 de: D2,
1128 ) -> std::result::Result<Self::Value, D2::Error> {
1129 struct Inner;
1130 impl<'de> Visitor<'de> for Inner {
1131 type Value = Vec<i32>;
1132 fn expecting(&self, f: &mut fmt::Formatter) -> fmt::Result {
1133 f.write_str("i32 or array of i32")
1134 }
1135 #[allow(clippy::cast_possible_truncation)]
1136 fn visit_i64<E: de::Error>(self, v: i64) -> std::result::Result<Self::Value, E> {
1137 Ok(vec![v as i32])
1138 }
1139 #[allow(clippy::cast_possible_truncation, clippy::cast_possible_wrap)]
1140 fn visit_u64<E: de::Error>(self, v: u64) -> std::result::Result<Self::Value, E> {
1141 Ok(vec![v as i32])
1142 }
1143 fn visit_seq<A: de::SeqAccess<'de>>(
1144 self,
1145 mut seq: A,
1146 ) -> std::result::Result<Self::Value, A::Error> {
1147 let mut out = Vec::with_capacity(seq.size_hint().unwrap_or(1));
1148 while let Some(x) = seq.next_element::<i32>()? {
1149 out.push(x);
1150 }
1151 Ok(out)
1152 }
1153 }
1154 de.deserialize_any(Inner).map(Some)
1155 }
1156 #[allow(clippy::cast_possible_truncation)]
1157 fn visit_i64<E: de::Error>(self, v: i64) -> std::result::Result<Self::Value, E> {
1158 Ok(Some(vec![v as i32]))
1159 }
1160 #[allow(clippy::cast_possible_truncation, clippy::cast_possible_wrap)]
1161 fn visit_u64<E: de::Error>(self, v: u64) -> std::result::Result<Self::Value, E> {
1162 Ok(Some(vec![v as i32]))
1163 }
1164 fn visit_seq<A: de::SeqAccess<'de>>(
1165 self,
1166 mut seq: A,
1167 ) -> std::result::Result<Self::Value, A::Error> {
1168 let mut out = Vec::with_capacity(seq.size_hint().unwrap_or(1));
1169 while let Some(x) = seq.next_element::<i32>()? {
1170 out.push(x);
1171 }
1172 Ok(Some(out))
1173 }
1174 }
1175 de.deserialize_option(V)
1176}
1177
1178/// Count of tensor maps requested beyond the buffer's declared
1179/// [`CpuAccess`] (including any CPU map of a `CpuAccess::None` buffer).
1180/// Undeclared CPU access is a pipeline smell: it forfeits layout
1181/// optimizations (write-combined mappings, tile compression) and may be
1182/// slow or refused on Android. Read via [`unplanned_cpu_access_count`];
1183/// each offending buffer also logs one warning.
1184static UNPLANNED_CPU_ACCESS: AtomicU64 = AtomicU64::new(0);
1185
1186/// Number of tensor maps that exceeded the buffer's declared
1187/// [`CpuAccess`] since process start. A pipeline that declares its CPU
1188/// access correctly holds this flat; see [`CpuAccess`] for the contract.
1189pub fn unplanned_cpu_access_count() -> u64 {
1190 UNPLANNED_CPU_ACCESS.load(Ordering::Relaxed)
1191}
1192
1193/// Record an unplanned CPU access on `identity_id`, warning once per
1194/// buffer (a steady-state pipeline maps the same buffer every frame — a
1195/// per-map warn would flood the log; repeats count silently).
1196///
1197/// `static`-only: the sole caller is `impl<T> TensorMapTrait<T> for
1198/// Tensor<T>` in `lib.rs`, itself `#[cfg(feature = "static")]`.
1199/// [`unplanned_cpu_access_count`] stays available under `dynamic` too --
1200/// it just always reads 0 there, honestly reflecting that this backend
1201/// produces no such telemetry.
1202#[cfg(feature = "static")]
1203pub(crate) fn note_unplanned_cpu_access(identity_id: u64, backend: &str, detail: &str) {
1204 UNPLANNED_CPU_ACCESS.fetch_add(1, Ordering::Relaxed);
1205 static WARNED: std::sync::OnceLock<std::sync::Mutex<std::collections::HashSet<u64>>> =
1206 std::sync::OnceLock::new();
1207 let warned = WARNED.get_or_init(Default::default);
1208 if warned.lock().is_ok_and(|mut s| s.insert(identity_id)) {
1209 log::warn!(
1210 "unplanned CPU access on {backend} buffer (identity {identity_id}): {detail} — \
1211 declare the intent at allocation (CpuAccess::Read/Write/ReadWrite) to make \
1212 this a planned, optimized mapping"
1213 );
1214 }
1215}
1216
1217/// Uniform guard for mutable access through a read-only map. Every
1218/// `TensorMap` backend calls this from `as_mut_slice`/`deref_mut` so a
1219/// `map_read()` misuse fails identically on all platforms instead of
1220/// passing on tolerant ones and exploding only on Android.
1221#[inline]
1222pub(crate) fn assert_map_writable(writable: bool, backend: &str) {
1223 assert!(
1224 writable,
1225 "{backend} map is read-only (obtained via map_read()/CpuAccess::Read) — \
1226 use map_mut() or map_write() for mutable access"
1227 );
1228}
1229
1230/// Declared CPU involvement for an image tensor, chosen at allocation.
1231///
1232/// The HAL assumes buffers are produced and consumed by hardware (ISP,
1233/// codec, GPU, NPU) — hardware access needs no declaration. CPU access is
1234/// the opt-in: it selects the CPU usage/mapping mode at allocation
1235/// (write-combined for `Write`, cached for `Read`) and, on Android, pins
1236/// the layout linear (vendor tile compression requires `None`).
1237///
1238/// Mapping beyond the declared access is best-effort, never silent: it
1239/// may be refused ([`Error::NotImplemented`]) or take a slow path, and it
1240/// always increments [`unplanned_cpu_access_count`] with a once-per-buffer
1241/// warning.
1242#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
1243pub enum CpuAccess {
1244 /// Hardware-only buffer (the default): no CPU mapping declared.
1245 /// Compression-eligible on platforms with vendor tile layouts.
1246 #[default]
1247 None,
1248 /// CPU reads (verification, CPU consumers) — cached mapping.
1249 Read,
1250 /// CPU writes (decode targets) — write-combined mapping where the
1251 /// platform supports it; reading through a `Write` map is undeclared.
1252 Write,
1253 /// CPU reads and writes — the pre-CpuAccess implicit behavior.
1254 ReadWrite,
1255}
1256
1257impl CpuAccess {
1258 /// Whether this declaration includes CPU reads.
1259 pub fn reads(self) -> bool {
1260 matches!(self, CpuAccess::Read | CpuAccess::ReadWrite)
1261 }
1262
1263 /// Whether this declaration includes CPU writes.
1264 pub fn writes(self) -> bool {
1265 matches!(self, CpuAccess::Write | CpuAccess::ReadWrite)
1266 }
1267
1268 /// Whether this declaration covers `requested` (every direction the
1269 /// request needs is declared).
1270 pub fn covers(self, requested: CpuAccess) -> bool {
1271 (!requested.reads() || self.reads()) && (!requested.writes() || self.writes())
1272 }
1273}
1274
1275/// Requested tile-compression behavior for an image allocation (set via
1276/// [`ImageDesc::with_compression`]).
1277///
1278/// Vendor GPUs store textures in proprietary compressed tile layouts
1279/// (UBWC, AFBC, PVRIC, DCC) that cut memory bandwidth; eligibility
1280/// requires a hardware-only buffer ([`CpuAccess::None`]) because CPU
1281/// mapping pins the layout linear. The request records best knowledge on
1282/// the tensor ([`Tensor::compression`]); it never changes what bytes a
1283/// consumer sees through the GPU/NPU.
1284#[non_exhaustive]
1285#[derive(Debug, Clone, Copy, PartialEq, Eq)]
1286pub enum Compression {
1287 /// Let the platform use its native scheme when the format is
1288 /// eligible; otherwise allocate linear and count the fallback
1289 /// ([`compression_fallback_count`]). The right default for pipelines
1290 /// that want the bandwidth win without portability failures.
1291 Any,
1292 /// Require one specific scheme: allocation fails with
1293 /// [`Error::InvalidArgument`] when the device's scheme differs and
1294 /// [`Error::NotImplemented`] on platforms without vendor tile
1295 /// compression. For consumers whose ABI names a layout (e.g. a
1296 /// QNN context binary declaring UBWC inputs).
1297 Scheme(CompressionScheme),
1298}
1299
1300// Declared through `ef_vocabulary!` (see `vocabulary.rs`). This vocabulary
1301// gained a wire numbering when `ef_tensor_compression` was added: the
1302// dynamic backend reads a scheme back across the C boundary, so the enum
1303// needs codes that cannot drift from the C enumerators. Hand-writing that
1304// mapping in `edgefirst-tensor-capi` was the alternative and it is strictly
1305// worse -- `#[non_exhaustive]` makes a cross-crate `match` require a
1306// wildcard arm, so a scheme added here would silently map to "linear"
1307// there. Generated from one declaration, it cannot.
1308//
1309// Code 0 is deliberately unassigned and reserved for "linear" -- the state
1310// this crate spells `Option::None`, which has no variant to carry a code of
1311// its own. The C side names it `EF_COMPRESSION_NONE`.
1312crate::ef_vocabulary! {
1313 /// Vendor tile-compression schemes the HAL recognizes and records.
1314 ///
1315 /// Unrecognized platforms record `None` (linear) — there is deliberately
1316 /// no `Unknown` variant; a scheme is only recorded when the vendor is
1317 /// positively identified.
1318 #[non_exhaustive]
1319 pub enum CompressionScheme {
1320 /// Qualcomm Adreno Universal Bandwidth Compression.
1321 Ubwc = 1, "ubwc", UBWC,
1322 /// Arm Mali/Immortalis Framebuffer Compression.
1323 Afbc = 2, "afbc", AFBC,
1324 /// Imagination PowerVR Image Compression (Google Tensor G5+).
1325 Pvric = 3, "pvric", PVRIC,
1326 /// Samsung Xclipse (AMD RDNA) Delta Color Compression.
1327 Dcc = 4, "dcc", DCC,
1328 }
1329 // `#[doc(hidden)]` for the same reason `dtype_wire` is: emission
1330 // plumbing for `const`-only consumers, not a second documented API.
1331 #[doc(hidden)]
1332 pub mod compression_wire;
1333}
1334
1335/// Count of image allocations that requested [`Compression::Any`] but
1336/// resolved to a linear layout (ineligible format/dtype, unrecognized
1337/// vendor, or a platform without vendor tile compression). Read via
1338/// [`compression_fallback_count`].
1339static COMPRESSION_FALLBACKS: AtomicU64 = AtomicU64::new(0);
1340
1341/// Number of [`Compression::Any`] requests since process start that
1342/// resolved to a linear layout instead of a vendor tile scheme. A
1343/// steady-state pipeline holds this flat after warmup; growth means an
1344/// allocation path keeps requesting compression it never gets.
1345pub fn compression_fallback_count() -> u64 {
1346 COMPRESSION_FALLBACKS.load(Ordering::Relaxed)
1347}
1348
1349/// Record a `Compression::Any` request resolving linear.
1350///
1351/// `static`-only: the sole caller is `Tensor::image_desc` (`impl<T>
1352/// Tensor<T>`), itself `#[cfg(feature = "static")]`.
1353/// [`compression_fallback_count`] stays available under `dynamic` too --
1354/// it just always reads 0 there.
1355#[cfg(feature = "static")]
1356pub(crate) fn note_compression_fallback(detail: &str) {
1357 COMPRESSION_FALLBACKS.fetch_add(1, Ordering::Relaxed);
1358 log::debug!("compression request fell back to linear: {detail}");
1359}
1360
1361/// Whether this platform can allocate `(format, dtype)` images in a
1362/// vendor tile-compressed layout. `true` requires an Android build, an
1363/// eligible format (RGBA8888 `u8`/`i8` initially), and a positively
1364/// identified GPU vendor; everywhere else the answer is `false` and
1365/// [`Compression::Any`] requests fall back to linear.
1366pub fn compression_support(format: PixelFormat, dtype: DType) -> bool {
1367 #[cfg(target_os = "android")]
1368 {
1369 crate::ahardwarebuffer_layout::compression_eligible(format, dtype)
1370 && crate::ahardwarebuffer::device_compression_scheme().is_some()
1371 }
1372 #[cfg(not(target_os = "android"))]
1373 {
1374 let _ = (format, dtype);
1375 false
1376 }
1377}
1378
1379/// Declarative image-allocation request — the full-featured front door
1380/// for image tensors ([`TensorDyn::image_desc`] and the image crate's
1381/// `ImageProcessor::create_image_desc`).
1382///
1383/// The classic constructors (`image`, `create_image`) cover the common
1384/// cases; the desc carries the optional requests — today the
1385/// [`Compression`] request — without another constructor-parameter
1386/// sweep. Fields are private and the builders consume/return by value,
1387/// so future options are non-breaking.
1388///
1389/// ```
1390/// use edgefirst_tensor::{Compression, CpuAccess, DType, ImageDesc, PixelFormat};
1391/// let desc = ImageDesc::new(640, 640, PixelFormat::Rgba, DType::U8)
1392/// .with_access(CpuAccess::None)
1393/// .with_compression(Compression::Any);
1394/// ```
1395#[derive(Debug, Clone)]
1396pub struct ImageDesc {
1397 width: usize,
1398 height: usize,
1399 format: PixelFormat,
1400 dtype: DType,
1401 memory: Option<TensorMemory>,
1402 access: CpuAccess,
1403 compression: Option<Compression>,
1404}
1405
1406impl ImageDesc {
1407 /// A new image request: auto-selected memory, [`CpuAccess::None`]
1408 /// (hardware-only), no compression request.
1409 pub fn new(width: usize, height: usize, format: PixelFormat, dtype: DType) -> Self {
1410 Self {
1411 width,
1412 height,
1413 format,
1414 dtype,
1415 memory: None,
1416 access: CpuAccess::None,
1417 compression: None,
1418 }
1419 }
1420
1421 /// Request a specific memory backing (`None` = auto-select).
1422 pub fn with_memory(mut self, memory: Option<TensorMemory>) -> Self {
1423 self.memory = memory;
1424 self
1425 }
1426
1427 /// Declare the CPU access (see [`CpuAccess`]). Any declaration other
1428 /// than `None` makes a compression request invalid.
1429 pub fn with_access(mut self, access: CpuAccess) -> Self {
1430 self.access = access;
1431 self
1432 }
1433
1434 /// Request a tile-compressed layout (see [`Compression`]).
1435 pub fn with_compression(mut self, compression: Compression) -> Self {
1436 self.compression = Some(compression);
1437 self
1438 }
1439
1440 /// Requested width in pixels.
1441 pub fn width(&self) -> usize {
1442 self.width
1443 }
1444
1445 /// Requested height in pixels.
1446 pub fn height(&self) -> usize {
1447 self.height
1448 }
1449
1450 /// Requested pixel format.
1451 pub fn format(&self) -> PixelFormat {
1452 self.format
1453 }
1454
1455 /// Requested element type.
1456 pub fn dtype(&self) -> DType {
1457 self.dtype
1458 }
1459
1460 /// Requested memory backing (`None` = auto-select).
1461 pub fn memory(&self) -> Option<TensorMemory> {
1462 self.memory
1463 }
1464
1465 /// Declared CPU access.
1466 pub fn access(&self) -> CpuAccess {
1467 self.access
1468 }
1469
1470 /// The compression request, if any.
1471 pub fn compression(&self) -> Option<Compression> {
1472 self.compression
1473 }
1474}
1475
1476/// What kind of system key an identity was derived from. Participates in the
1477/// id so that, say, inode 7 and IOSurfaceID 7 cannot collide.
1478///
1479/// Invariant every variant must uphold: the key passed to
1480/// [`BufferIdentity::derived`] is unique among live buffers of that kind
1481/// **that own an allocation**. There is no "no identity" case — where no
1482/// system-wide key exists a process-local one does (a pointer, a GL name),
1483/// and a process-local key is exactly as sound as a system-wide one:
1484/// pointer/name reuse after free is the same ABA hazard as inode reuse, and
1485/// is mitigated the same way, by the retaining cache (see the design spec's
1486/// Task 6).
1487///
1488/// The allocation qualifier is load-bearing, not a hedge. A zero-element
1489/// tensor (`[0, num_protos]`, the tracker's "no detections this frame"
1490/// sentinel — see [`crate::mem`]) owns no allocation, so its `Vec` yields the
1491/// dangling well-aligned pointer `align_of::<T>()`. Every zero-element
1492/// `Tensor<f32>` in a process therefore shares one identity, and can even
1493/// collide across kinds with a small integer key from another variant.
1494///
1495/// That is safe rather than tolerated: an identity exists to answer "are
1496/// these the same bytes?", and buffers with no bytes have nothing to alias.
1497/// Nor can they reach an import cache. There are four GPU import paths —
1498/// a dma-buf fd on Linux, an IOSurface ref on macOS, an `AHardwareBuffer` on
1499/// Android, and an `ID3D11Texture2D` on Windows — and each takes only its own
1500/// object, so a tensor with no allocation reaches none of them. The Windows
1501/// path additionally keys on the texture the import itself holds a reference
1502/// to, and `edgefirst-image` **refuses** any texture whose identity is not
1503/// `D3d11Texture` over that same pointer — in every build, and on every
1504/// import including one served from its cache
1505/// (`GlPlatform::validate_import_identity`). That refusal exists because the
1506/// rule was once only written
1507/// down: the dynamic backend handed the image crate `HostPtr` identities over
1508/// recyclable handle addresses, and converts silently rendered into dropped
1509/// textures. Do not "fix" the zero-element case by minting a distinct id per
1510/// empty tensor: that reintroduces a counter, which is the state this whole
1511/// design exists to remove.
1512#[derive(Debug, Clone, Copy, PartialEq, Eq)]
1513pub enum IdentityKind {
1514 /// `(st_dev, st_ino)` of a dma-buf or memfd. Survives `dup`, and is the
1515 /// same value in every process that holds the buffer.
1516 DmaBuf,
1517 Shm,
1518 /// `IOSurfaceGetID`.
1519 IoSurface,
1520 /// `AHardwareBuffer_getId`, API 31+.
1521 AHardwareBufferId,
1522 /// The `AHardwareBuffer*` pointer, pre-API-31 where no id exists.
1523 /// Process-local: unique among live buffers here, meaningless elsewhere.
1524 AHardwareBufferPtr,
1525 /// An `ID3D11Texture2D*`. Process-local, like `AHardwareBufferPtr`, and
1526 /// the same pointer for every kind of texture tensor: one the HAL
1527 /// allocated, one wrapped from a caller's pointer, and one opened from a
1528 /// shared NT handle all key on the texture object they hold.
1529 ///
1530 /// **Both backends derive it from the same pointer, so their ids are
1531 /// equal byte for byte.** The static backend reads the `ComPtr` it owns
1532 /// (`d3d11::texture::tex_key`); the dynamic backend reads the same
1533 /// pointer back across the ABI with `ef_tensor_d3d11_texture`
1534 /// (`TensorDyn::derive_identity`). That equality is the point, not a
1535 /// coincidence: [`BufferIdentity::derived`] exists so two independently
1536 /// linked copies of this crate name one buffer identically, and
1537 /// `edgefirst-image` links the dynamic backend inside every `-capi` leaf
1538 /// while its own tests link the static one.
1539 ///
1540 /// **An `ef_tensor` handle's address must never be the key.** It is not a
1541 /// buffer at all, nothing the cache holds keeps it alive, and this
1542 /// library's allocator hands a freed handle's address to the very next
1543 /// one it mints — measured at 400 texture handles over 4 distinct
1544 /// addresses. The dynamic backend did exactly that on Windows until
1545 /// 2026-09-04, and a convert whose destination key collided rendered into
1546 /// the dropped texture while returning success.
1547 ///
1548 /// A handle import therefore does **not** compare equal to the tensor it
1549 /// was exported from, nor to a second import of the same handle:
1550 /// `OpenSharedResource1` mints a fresh `ID3D11Texture2D` per open. Keying
1551 /// on a handle value would make those compare equal, and would be unsound:
1552 /// Windows recycles a closed handle value straight into the next
1553 /// `DuplicateHandle`, so a consumer that caches on this id past the
1554 /// producing tensor's lifetime -- as the image crate's EGLImage cache
1555 /// does -- could serve one texture's entry for a later, unrelated import.
1556 /// A COM object's address cannot be recycled that way: a cache that holds
1557 /// a reference to the texture keeps the address taken for as long as its
1558 /// entry can be found. The cost is a cache miss per re-import instead of
1559 /// a wrong image.
1560 D3d11Texture,
1561 /// A host mapping address. Process-local.
1562 HostPtr,
1563 /// A GL buffer name. Valid only in the context that created it.
1564 Pbo,
1565}
1566
1567/// Unique identity for a tensor's underlying buffer.
1568///
1569/// Created fresh on every buffer allocation or import. The `id` is a cache
1570/// key, `kind` records what system key (if any) it was derived from, and the
1571/// `guard` is an `Arc<()>` whose weak references allow downstream caches to
1572/// detect when the buffer has been dropped.
1573#[derive(Debug, Clone)]
1574pub struct BufferIdentity {
1575 id: u64,
1576 kind: IdentityKind,
1577 guard: Arc<()>,
1578}
1579
1580impl BufferIdentity {
1581 /// Derive an identity from a key the operating system already guarantees.
1582 ///
1583 /// Deliberately NOT a counter: two independently-linked copies of this
1584 /// crate must derive the same id for the same buffer, or a consumer's
1585 /// import cache misses every frame. See the design spec, "The core finding".
1586 pub fn derived(kind: IdentityKind, key: u64) -> Self {
1587 // Mix the discriminant into the high bits so different kinds with the
1588 // same numeric key never collide.
1589 let id = ((kind as u64) << 56) ^ key;
1590 Self {
1591 id,
1592 kind,
1593 guard: Arc::new(()),
1594 }
1595 }
1596
1597 /// Unique identifier for this buffer. Changes when the buffer changes.
1598 pub fn id(&self) -> u64 {
1599 self.id
1600 }
1601
1602 /// What kind of system key this identity was derived from.
1603 pub fn kind(&self) -> IdentityKind {
1604 self.kind
1605 }
1606
1607 /// Returns a weak reference to the buffer guard. Goes dead when the
1608 /// owning Tensor is dropped (and no clones remain).
1609 pub fn weak(&self) -> Weak<()> {
1610 Arc::downgrade(&self.guard)
1611 }
1612}
1613
1614// `static`-only from here through `fs_magic` below: every use site lives in
1615// `impl<T> TensorStorage<T>` (`from_fd`'s DMA-BUF-vs-tmpfs classification),
1616// itself `#[cfg(feature = "static")]`. `dynamic`'s `from_fd`
1617// (`dynamic_backend.rs`) drives `ef_tensor_builder_wrap` instead, which
1618// does this same classification inside `libedgefirst_tensor.so`.
1619#[cfg(all(target_os = "linux", feature = "static"))]
1620use nix::sys::stat::{major, minor};
1621
1622/// Filesystem magic of the internal `dma_buf` mount, from
1623/// `include/uapi/linux/magic.h` (`DMA_BUF_MAGIC`, "DMAB").
1624///
1625/// This is stable UAPI and is the only reliable way to recognize a DMA-BUF
1626/// fd — see [`TensorStorage::from_fd`] for why `st_dev` cannot be used.
1627#[cfg(all(target_os = "linux", feature = "static"))]
1628const DMA_BUF_MAGIC: u32 = 0x444d_4142;
1629
1630/// Filesystem magic of tmpfs, from `include/uapi/linux/magic.h`
1631/// (`TMPFS_MAGIC`). Covers both POSIX `shm_open` segments under `/dev/shm`
1632/// and anonymous `memfd_create` files.
1633#[cfg(all(target_os = "linux", feature = "static"))]
1634const TMPFS_MAGIC: u32 = 0x0102_1994;
1635
1636/// Normalize a raw `fstatfs` `f_type` to the 32-bit unsigned magic that
1637/// `include/uapi/linux/magic.h` documents.
1638///
1639/// The width and signedness of `f_type` are target-dependent: `__fsword_t`
1640/// on Linux/gnu (`i64` on 64-bit, `i32` on 32-bit targets such as armv7,
1641/// i686 and aarch64-ilp32), `c_int` on uclibc, `c_ulong` on musl, and
1642/// `c_uint` on s390x. Where it is signed and 32 bits wide, widening it
1643/// sign-extends every magic with bit 31 set — hugetlbfs (`0x958458f6`),
1644/// btrfs (`0x9123683e`), f2fs (`0xf2f52010`) — which would make the value
1645/// reported by [`Error::UnknownBufferType`] impossible to find in the
1646/// header the docs point callers at.
1647///
1648/// Every magic is a 32-bit quantity, so truncating back to `u32` recovers
1649/// the true value from all four representations and loses nothing.
1650#[cfg(all(target_os = "linux", feature = "static"))]
1651const fn fs_magic(raw: i64) -> u32 {
1652 raw as u32
1653}
1654
1655/// The operations every memory backend implements.
1656///
1657/// This is the seam that lets [`Tensor<T>`] hide which backend is in play.
1658/// The implementors are `DmaTensor` (Linux DMA-BUF), `IoSurfaceTensor`
1659/// (macOS/iOS), `AHardwareBufferTensor` (Android), `ShmTensor`, `MemTensor`,
1660/// and [`PboTensor`]; `TensorStorage<T>` dispatches to whichever is active.
1661/// Backend types are not public — allocate a [`Tensor<T>`] or [`TensorDyn`]
1662/// and call these methods through it.
1663///
1664/// Import the trait to get `shape`, `size`, `map`, `clone_fd`,
1665/// `buffer_identity`, and the zero-copy sub-region `view` on a tensor value.
1666/// `Tensor::view` and `Tensor::batch` route through [`TensorTrait::view`], so
1667/// each backend's identity-sharing rule lives in exactly one place.
1668///
1669/// [`TensorDyn`]: crate::TensorDyn
1670pub trait TensorTrait<T>: Send + Sync
1671where
1672 T: Num + Clone + fmt::Debug,
1673{
1674 /// Create a new tensor with the given shape and optional name. If no name
1675 /// is given, a random name will be generated.
1676 fn new(shape: &[usize], name: Option<&str>) -> Result<Self>
1677 where
1678 Self: Sized;
1679
1680 #[cfg(unix)]
1681 /// Import an existing buffer as a tensor, taking ownership of its file
1682 /// descriptor. The buffer is adopted in place — no bytes are copied.
1683 ///
1684 /// The backend is **detected**, not chosen: the fd already belongs to a
1685 /// buffer type, and this call's job is to recognize which. On Linux that
1686 /// is decided by the fd's filesystem magic, which is stable UAPI
1687 /// (`include/uapi/linux/magic.h`):
1688 ///
1689 /// | Filesystem | Magic | Resulting [`TensorMemory`] |
1690 /// |------------|-------|----------------------------|
1691 /// | `dma_buf` | `DMA_BUF_MAGIC` (`0x444d4142`) | [`TensorMemory::DmaBuf`] |
1692 /// | `tmpfs` (`/dev/shm` **and** `memfd`) | `TMPFS_MAGIC` (`0x01021994`) | [`TensorMemory::Shm`] |
1693 /// | anything else | — | *rejected* — see Errors |
1694 ///
1695 /// Both supported types are identified **positively**. An unrecognized
1696 /// filesystem is an error, never a fallback to shared memory: the wrong
1697 /// branch does not fail loudly (a DMA-BUF is mmap-able, so it would
1698 /// import as a perfectly functional tensor that merely isn't DMA, and a
1699 /// pipe would import as a zero-length one), so guessing would trade a
1700 /// clear error here for silent loss of zero-copy far downstream.
1701 ///
1702 /// The device number is deliberately **not** consulted. `dma_buf` files
1703 /// live on an internal kernel mount whose `st_dev` comes from
1704 /// `get_anon_bdev()` — an IDA shared with every other anonymous
1705 /// pseudo-filesystem and allocated in boot order — so the minor a
1706 /// DMA-BUF lands on varies by kernel build and is not part of any ABI.
1707 ///
1708 /// On non-Linux Unix (macOS/iOS/Android) there is no fd-based DMA import;
1709 /// the fd is always adopted as [`TensorMemory::Shm`].
1710 ///
1711 /// # Arguments
1712 ///
1713 /// * `fd` - Owned descriptor for the buffer to import. Ownership
1714 /// transfers to the returned tensor and the fd is closed on drop; pass
1715 /// [`clone_fd`](TensorTrait::clone_fd) output to keep your own handle.
1716 /// * `shape` - Logical shape to interpret the buffer with. Must describe
1717 /// no more elements than the buffer holds.
1718 /// * `name` - Optional name; a random one is generated when `None`.
1719 ///
1720 /// # Returns
1721 ///
1722 /// A tensor sharing the imported buffer's memory, whose
1723 /// [`memory()`](TensorTrait::memory) reports the detected backend.
1724 ///
1725 /// # Errors
1726 ///
1727 /// * [`Error::UnknownBufferType`] - the fd is on a filesystem that is
1728 /// neither `dma_buf` nor `tmpfs`, so its buffer type cannot be
1729 /// determined. Carries the observed `fstatfs` magic as a `u32`,
1730 /// normalized so it can be looked up in `magic.h` directly on both
1731 /// 32- and 64-bit targets. Typical causes: a regular file, a pipe or
1732 /// socket, or a `MFD_HUGETLB` memfd (hugetlbfs, not tmpfs). Linux only.
1733 /// * [`Error::UnknownDeviceType`] - the fd's `st_dev` major is non-zero,
1734 /// i.e. it lives on a real block device rather than an anonymous or
1735 /// in-memory filesystem. Linux only.
1736 /// * [`Error::InvalidSize`] - `shape` is empty or describes zero
1737 /// elements.
1738 /// * [`Error::NixError`] - `fstat`, `fstatfs`, or `mmap` failed on the
1739 /// descriptor.
1740 ///
1741 /// # Examples
1742 ///
1743 /// ```no_run
1744 /// use edgefirst_tensor::{Tensor, TensorMemory, TensorTrait};
1745 ///
1746 /// # fn main() -> edgefirst_tensor::Result<()> {
1747 /// let src = Tensor::<u8>::new(&[480, 640, 3], Some(TensorMemory::DmaBuf), None)?;
1748 ///
1749 /// // Round-tripping a DMA-BUF fd preserves the backend — the import is
1750 /// // still zero-copy, and still eligible for GPU/NPU paths.
1751 /// let imported = Tensor::<u8>::from_fd(src.clone_fd()?, src.shape(), None)?;
1752 /// assert_eq!(imported.memory(), TensorMemory::DmaBuf);
1753 /// # Ok(())
1754 /// # }
1755 /// ```
1756 fn from_fd(fd: std::os::fd::OwnedFd, shape: &[usize], name: Option<&str>) -> Result<Self>
1757 where
1758 Self: Sized;
1759
1760 #[cfg(unix)]
1761 /// Clone the file descriptor associated with this tensor.
1762 fn clone_fd(&self) -> Result<std::os::fd::OwnedFd>;
1763
1764 /// Get the memory type of this tensor.
1765 fn memory(&self) -> TensorMemory;
1766
1767 /// Get the name of this tensor.
1768 fn name(&self) -> String;
1769
1770 /// Get the number of elements in this tensor.
1771 fn len(&self) -> usize {
1772 self.shape().iter().product()
1773 }
1774
1775 /// Check if the tensor is empty.
1776 fn is_empty(&self) -> bool {
1777 self.len() == 0
1778 }
1779
1780 /// Get the size in bytes of this tensor.
1781 fn size(&self) -> usize {
1782 self.len() * std::mem::size_of::<T>()
1783 }
1784
1785 /// Get the shape of this tensor.
1786 fn shape(&self) -> &[usize];
1787
1788 /// Reshape this tensor to the given shape. The total number of elements
1789 /// must remain the same.
1790 fn reshape(&mut self, shape: &[usize]) -> Result<()>;
1791
1792 /// Bytes of the underlying allocation (>= the current logical `size()`).
1793 /// Defaults to the logical size for storages without spare capacity.
1794 fn capacity_bytes(&self) -> usize {
1795 self.size()
1796 }
1797
1798 /// Set the logical shape to any shape whose byte size fits the allocation
1799 /// capacity, without the equal-size constraint of `reshape`.
1800 fn set_logical_shape(&mut self, shape: &[usize]) -> Result<()> {
1801 self.reshape(shape)
1802 }
1803
1804 /// Map the tensor into memory with the given access direction and
1805 /// return a TensorMap for accessing the data.
1806 ///
1807 /// `access` selects the platform mapping mode (read-only IOSurface
1808 /// lock, dma-buf sync direction, AHardwareBuffer lock usage) and the
1809 /// map's mutability: a map obtained with [`CpuAccess::Read`] rejects
1810 /// `as_mut_slice`. [`CpuAccess::None`] is not a mappable direction
1811 /// and returns [`Error::InvalidArgument`].
1812 ///
1813 /// Prefer the typed wrappers [`map_read`](Self::map_read) /
1814 /// [`map_write`](Self::map_write) / [`map_mut`](Self::map_mut).
1815 /// Map for CPU access.
1816 ///
1817 /// Like [`pin_host`](Tensor::pin_host), `'a` bounds how long the backing
1818 /// allocation lives rather than borrowing the tensor: the returned view
1819 /// shares ownership through the pin's keepalive, so a map can coexist with
1820 /// a `&mut` on the same tensor.
1821 ///
1822 /// # Windows D3D11 textures: no write-only map of a partial window
1823 ///
1824 /// A [`CpuAccess::Write`] map is not refreshed from the texture and
1825 /// publishes the whole texture when it unmaps, so the bytes the caller
1826 /// leaves untouched are undefined. That is a contract a whole-tensor map
1827 /// can meet and a sub-view cannot: the window spans its own rows only.
1828 /// A write-only map of a tensor whose window is shorter than its backing
1829 /// -- a [`view`](Tensor::view), or a shape narrowed with
1830 /// `set_logical_shape` -- is therefore refused with
1831 /// [`Error::InvalidArgument`] naming [`CpuAccess::ReadWrite`], which is
1832 /// the same window with the refresh that makes the untouched rows
1833 /// well-defined.
1834 fn map_with<'a>(&self, access: CpuAccess) -> Result<crate::view::HostView<'a, T>>
1835 where
1836 T: 'a;
1837
1838 /// Map the tensor read-write (equivalent to
1839 /// `map_with(CpuAccess::ReadWrite)` — the historical `map()`
1840 /// behavior).
1841 fn map<'a>(&self) -> Result<crate::view::HostView<'a, T>>
1842 where
1843 T: 'a,
1844 {
1845 self.map_with(CpuAccess::ReadWrite)
1846 }
1847
1848 /// Map the tensor for CPU reading only. The returned map rejects
1849 /// `as_mut_slice`; on macOS this takes the read-only IOSurface lock
1850 /// (skips the unlock flush), on Linux the dma-buf read-direction
1851 /// sync.
1852 fn map_read<'a>(&self) -> Result<crate::view::HostView<'a, T>>
1853 where
1854 T: 'a,
1855 {
1856 self.map_with(CpuAccess::Read)
1857 }
1858
1859 /// Map the tensor for CPU writing (fill-only: reading through a
1860 /// write map may see write-combined memory — do not read the slice).
1861 fn map_write<'a>(&self) -> Result<crate::view::HostView<'a, T>>
1862 where
1863 T: 'a,
1864 {
1865 self.map_with(CpuAccess::Write)
1866 }
1867
1868 /// Map the tensor read-write (alias of [`map`](Self::map) with the
1869 /// intent spelled out).
1870 fn map_mut<'a>(&self) -> Result<crate::view::HostView<'a, T>>
1871 where
1872 T: 'a,
1873 {
1874 self.map_with(CpuAccess::ReadWrite)
1875 }
1876
1877 /// Non-blocking [`map_with`](Self::map_with): returns
1878 /// `Err(Error::IoError(WouldBlock))` while a GPU copy the map depends on
1879 /// is still in flight, and makes progress on a retry.
1880 ///
1881 /// Backings with no such copy (host memory, shared memory, dma-buf,
1882 /// IOSurface, AHardwareBuffer, PBO) alias `map_with`, so a caller can use
1883 /// this unconditionally and only ever see `WouldBlock` where the wait is
1884 /// real -- today the Windows D3D11 texture's staging refresh.
1885 ///
1886 /// A retry loop must yield or sleep between attempts. On the WARP
1887 /// software adapter the CPU threads are the GPU, so a tight loop starves
1888 /// the very copy it is waiting for.
1889 ///
1890 /// [`CpuAccess::Write`] is refused on a Windows D3D11 texture tensor whose
1891 /// window is shorter than its backing, exactly as it is for
1892 /// [`map_with`](Self::map_with); see that method.
1893 fn try_map_with<'a>(&self, access: CpuAccess) -> Result<crate::view::HostView<'a, T>>
1894 where
1895 T: 'a,
1896 {
1897 self.map_with(access)
1898 }
1899
1900 /// Non-blocking [`map`](Self::map): read-write, and `WouldBlock` rather
1901 /// than a stall. See [`try_map_with`](Self::try_map_with).
1902 fn try_map<'a>(&self) -> Result<crate::view::HostView<'a, T>>
1903 where
1904 T: 'a,
1905 {
1906 self.try_map_with(CpuAccess::ReadWrite)
1907 }
1908
1909 /// Get the buffer identity for cache keying and liveness tracking.
1910 fn buffer_identity(&self) -> &BufferIdentity;
1911
1912 /// Create a zero-copy sub-region view of this backing that shares the
1913 /// underlying allocation **and** [`BufferIdentity`].
1914 ///
1915 /// The window is `[offset_bytes, offset_bytes + shape.product() *
1916 /// size_of::<T>())` measured from this tensor's own logical start, so a
1917 /// sub-view of a sub-view composes by adding offsets. Sharing the parent's
1918 /// identity is the contract that lets identity-keyed caches (e.g. the GL
1919 /// EGLImage import cache) treat offset-distinct windows as one buffer rather
1920 /// than unrelated allocations — `view` must never mint a fresh identity.
1921 ///
1922 /// Defaults to [`Error::NotImplemented`]; every backend that supports
1923 /// sub-views overrides it (`Mem`, `Shm`, Linux DMA, macOS IOSurface, `Pbo`).
1924 fn view(&self, offset_bytes: usize, shape: &[usize]) -> Result<Self>
1925 where
1926 Self: Sized,
1927 {
1928 let _ = (offset_bytes, shape);
1929 Err(Error::NotImplemented(
1930 "view (zero-copy sub-region) is not supported for this tensor backend".to_owned(),
1931 ))
1932 }
1933}
1934
1935pub trait TensorMapTrait<T>
1936where
1937 T: Num + Clone + fmt::Debug,
1938{
1939 /// Get the shape of this tensor map.
1940 fn shape(&self) -> &[usize];
1941
1942 /// Unmap the tensor from memory.
1943 fn unmap(&mut self);
1944
1945 /// Get the number of elements in this tensor map.
1946 fn len(&self) -> usize {
1947 self.shape().iter().product()
1948 }
1949
1950 /// Check if the tensor map is empty.
1951 fn is_empty(&self) -> bool {
1952 self.len() == 0
1953 }
1954
1955 /// Get the size in bytes of this tensor map.
1956 fn size(&self) -> usize {
1957 self.len() * std::mem::size_of::<T>()
1958 }
1959
1960 /// Get a slice to the data in this tensor map.
1961 fn as_slice(&self) -> &[T];
1962
1963 /// Get a mutable slice to the data in this tensor map.
1964 fn as_mut_slice(&mut self) -> &mut [T];
1965
1966 /// Whether this mapping permits writes -- i.e. whether the coherency
1967 /// bracket it holds will flush on release.
1968 ///
1969 /// The single source of truth for a binding that has to advertise
1970 /// mutability to a foreign consumer (the Python buffer protocol's
1971 /// `readonly` flag, say). Deriving that from the access a caller
1972 /// *requested* instead would leave the advertisement agreeing with the
1973 /// request while both disagreed with the mapping actually taken.
1974 fn is_writable(&self) -> bool;
1975
1976 #[cfg(feature = "ndarray")]
1977 /// Get an ndarray ArrayView of the tensor data.
1978 fn view(&'_ self) -> Result<ndarray::ArrayView<'_, T, ndarray::Dim<ndarray::IxDynImpl>>> {
1979 Ok(ndarray::ArrayView::from_shape(
1980 self.shape(),
1981 self.as_slice(),
1982 )?)
1983 }
1984
1985 #[cfg(feature = "ndarray")]
1986 /// Get an ndarray ArrayViewMut of the tensor data.
1987 fn view_mut(
1988 &'_ mut self,
1989 ) -> Result<ndarray::ArrayViewMut<'_, T, ndarray::Dim<ndarray::IxDynImpl>>> {
1990 let shape = self.shape().to_vec();
1991 Ok(ndarray::ArrayViewMut::from_shape(
1992 shape,
1993 self.as_mut_slice(),
1994 )?)
1995 }
1996}
1997
1998// Declared through `ef_vocabulary!` (see `vocabulary.rs`). Before this, the
1999// same four backings were numbered three different ways -- Rust's implicit
2000// discriminants (`Dma=0, Shm=1, Mem=2, Pbo=3`), the C ABI's
2001// `hal_tensor_memory` (`MEM=0, DMA=1, SHM=2, PBO=3`), and Python's
2002// `PyTensorMemory`, which had no explicit discriminants at all and a
2003// `#[cfg(unix)]` variant in the middle of the list, so its numbering
2004// depended on the target OS. The canonical assignment below is the one all
2005// of them converge on; `hal_tensor_memory` is the single explicitly-mapped
2006// outlier that keeps its own released values (see `to_hal_tensor_memory` in
2007// the capi crate).
2008//
2009// Note this is NOT `protocol::kind`, which is a genuinely different
2010// vocabulary and stays so: `kind` describes what a `TensorDesc`'s `handle`
2011// and `ptr` fields MEAN to an importer, and both `Mem` and `Shm` mean the
2012// same thing there (`kind::HOST`, host-addressable pointer, no handle). A
2013// vocabulary that collapses two variants into one cannot also be the
2014// vocabulary that distinguishes them. `kind_of()` in `protocol.rs` is the
2015// mapping between them.
2016crate::ef_vocabulary! {
2017 /// Which memory backend a tensor is (or should be) allocated from.
2018 ///
2019 /// Pass `Some(..)` to a constructor to pin the backend, or `None` to
2020 /// auto-select. Auto-selection differs by constructor: [`Tensor::new`]
2021 /// tries `DmaBuf` → `Mem`, while the image constructors try `DmaBuf` →
2022 /// `Mem` and skip `Shm`. `EDGEFIRST_TENSOR_FORCE_MEM=1` short-circuits
2023 /// either chain to `Mem`, which is how tests run on hosts without
2024 /// DMA-heap permissions.
2025 ///
2026 /// A pinned request has no fallback: if the backend cannot serve it, the
2027 /// constructor fails rather than quietly giving you something slower.
2028 /// [`TensorTrait::memory`] reports what a tensor actually ended up with.
2029 ///
2030 /// # Every variant exists on every platform
2031 ///
2032 /// This is a namespace of codes, not a list of what this build can
2033 /// allocate -- [`TensorMemory::is_available`] answers the latter, at
2034 /// runtime. `Shm` used to be `#[cfg(unix)]`, which shifted every
2035 /// variant declared after it on a non-unix build and made the
2036 /// numbering a compile-time detail of the host. It also made a
2037 /// cross-platform recording unreadable: a blob captured on Linux
2038 /// carries `"dmabuf"`, and a macOS consumer has to be able to NAME
2039 /// that backing in order to report "cannot materialise here" instead
2040 /// of failing with an unrecognised code.
2041 ///
2042 /// # `#[non_exhaustive]`
2043 ///
2044 /// Match this with a wildcard arm. The variant list is expected to
2045 /// keep growing -- `IoSurface` and `Cuda` were added here, Windows
2046 /// backings and a real CUDA device backing are named as coming -- and
2047 /// this attribute is what keeps a closed-source consumer's
2048 /// `match memory { .. }` compiling across those additions. It protects
2049 /// DOWNSTREAM, not this crate: inside `edgefirst-tensor` the matches
2050 /// stay exhaustive and compile-enforced. Same rule as [`DType`],
2051 /// [`Compression`] and [`PixelFormat`], which `TensorMemory` was the
2052 /// lone exception to.
2053 #[non_exhaustive]
2054 pub enum TensorMemory {
2055 /// Regular system memory allocation. Available everywhere; the
2056 /// floor every auto-selection chain ends on.
2057 Mem = 0, "mem", MEM,
2058 /// POSIX Shared Memory allocation. Suitable for inter-process
2059 /// communication, but not for hardware acceleration. Only
2060 /// allocatable on unix; nameable everywhere.
2061 Shm = 1, "shm", SHM,
2062 /// Platform-native zero-copy GPU buffer.
2063 ///
2064 /// On Linux this is a DMA-BUF (`DmaTensor` in
2065 /// `crates/tensor/src/dma.rs`) allocated via the DRM/dma-heap
2066 /// subsystem. On macOS/iOS it is currently also this variant that
2067 /// yields an IOSurface, and on Android an AHardwareBuffer: they
2068 /// share the `TensorStorage::Dma` slot at the trait level, and the
2069 /// public C API discriminant (`HAL_TENSOR_MEMORY_DMA = 1`) covers
2070 /// all three with no ABI break. `IoSurface` below is the code for
2071 /// naming that backing specifically; no backend reports it yet.
2072 ///
2073 /// Allows hardware-accelerated paths (OpenGL backend on Linux via
2074 /// `EGL_EXT_image_dma_buf_import`; macOS via
2075 /// `EGL_ANGLE_iosurface_client_buffer`). CPU access via `map()`
2076 /// incurs cache-coherency overhead on Linux DMA-BUF and is similar
2077 /// in cost on IOSurface; SHM/Mem are cheaper for CPU-only workloads.
2078 ///
2079 /// On Windows it is an `ID3D11Texture2D` on the process D3D11 device
2080 /// (`crates/tensor/src/d3d11/`), allocated only through
2081 /// [`Tensor::image`](crate::Tensor::image); CPU access goes through a
2082 /// staging copy. The first such allocation creates the process
2083 /// device, so a host supplying its own device must call
2084 /// `d3d11::use_external_device` before it -- and before
2085 /// [`crate::is_gpu_buffer_available`], which creates the device too.
2086 DmaBuf = 2, "dmabuf", DMABUF,
2087 /// Apple IOSurface, named specifically rather than through the
2088 /// portable `DmaBuf` spelling.
2089 ///
2090 /// Defined as a code so a descriptor or recording that names an
2091 /// IOSurface is parseable on every platform. **No backend produces
2092 /// or accepts it yet** -- macOS/iOS allocate and report `DmaBuf`.
2093 /// Pinning a request to it fails with `NotImplemented`.
2094 IoSurface = 3, "iosurface", IOSURFACE,
2095 /// OpenGL Pixel Buffer Object memory. Created by `ImageProcessor`
2096 /// when DMA-buf is unavailable but OpenGL is present -- never by
2097 /// [`Tensor::new`], which has no GL context to create it in.
2098 Pbo = 4, "pbo", PBO,
2099 /// CUDA device memory: a device pointer, not host-addressable.
2100 ///
2101 /// Defined as a code for the same reason as `IoSurface`. **No
2102 /// backend produces or accepts it yet**; pinning a request to it
2103 /// fails with `NotImplemented`.
2104 Cuda = 5, "cuda", CUDA,
2105 }
2106 // `#[doc(hidden)]`: `pub` so the const-only form is reachable by an FFI
2107 // or cbindgen consumer, but it is emission plumbing rather than a second
2108 // documented API -- `TensorMemory::code()` is the documented path, and
2109 // the two are generated from the same literal.
2110 #[doc(hidden)]
2111 pub mod tensor_memory_wire;
2112}
2113
2114impl TensorMemory {
2115 /// Whether this backing can actually be allocated in this process, now.
2116 ///
2117 /// A runtime question, not a compile-time one, which is why it is a
2118 /// method on the vocabulary rather than a `cfg` on the variant:
2119 /// `DmaBuf` needs `/dev/dma_heap` to exist *and* this process to have
2120 /// permission on it (a container, a user outside the `video`/`render`
2121 /// groups, and a stock desktop give three different answers on the same
2122 /// kernel), and `Shm` needs a writable `/dev/shm`. Each probe allocates
2123 /// once and caches.
2124 ///
2125 /// Reports what can be *allocated*, which is not the same as what can be
2126 /// *named*: every variant is nameable on every platform (that is the
2127 /// point of the vocabulary), and the ones this build cannot serve say so
2128 /// here rather than by being absent from the enum.
2129 ///
2130 /// Three variants answer `false` unconditionally today, for reasons
2131 /// worth stating rather than hiding behind a probe that cannot fail:
2132 ///
2133 /// * `IoSurface` -- no backend produces or accepts it yet; macOS/iOS
2134 /// serve IOSurface under the `DmaBuf` spelling, so `DmaBuf` is the
2135 /// variant that answers `true` there.
2136 /// * `Pbo` -- PBO tensors are created by `ImageProcessor::create_image`
2137 /// with a GL context current, never by [`Tensor::new`]. This crate
2138 /// holds no GL context and cannot see one; ask the image crate
2139 /// whether a GL backend is live.
2140 /// * `Cuda` -- there is no CUDA backing in this crate at all. Probing
2141 /// for `libcuda` would answer `true` on a Jetson while every
2142 /// allocation still failed, which is worse than answering `false`.
2143 ///
2144 /// One variant answers `true` for a narrower constructor than the rest:
2145 /// on Windows `DmaBuf` is an `ID3D11Texture2D`, which is image-formatted,
2146 /// so `true` there means [`Tensor::image`] can serve it and
2147 /// [`Tensor::new`] with `Some(DmaBuf)` still answers `NotImplemented`.
2148 /// [`is_gpu_buffer_available`](crate::is_gpu_buffer_available), which this
2149 /// forwards to, states the same split.
2150 pub fn is_available(self) -> bool {
2151 match self {
2152 // Always: the fallback every constructor chain ends on. If this
2153 // could fail, allocation would have no floor left.
2154 TensorMemory::Mem => true,
2155 TensorMemory::Shm => is_shm_available(),
2156 // Today `DmaBuf` is the portable spelling of "platform-native
2157 // zero-copy GPU buffer", so it is the portable probe that
2158 // answers for it -- true on macOS/iOS (IOSurface) and Android
2159 // (AHardwareBuffer) as well as Linux.
2160 TensorMemory::DmaBuf => is_gpu_buffer_available(),
2161 TensorMemory::IoSurface | TensorMemory::Pbo | TensorMemory::Cuda => false,
2162 }
2163 }
2164
2165 /// The error to return when a caller pins a request to a backing this
2166 /// build knows the name of but cannot serve.
2167 ///
2168 /// Deliberately not `unreachable!()`. `IoSurface` and `Cuda` are
2169 /// unreachable *today* only because no backend reports them yet; the
2170 /// day one does, an `unreachable!()` written now becomes a panic in
2171 /// library code reached from an ordinary caller request. A "this build
2172 /// cannot do that" error is the same information without the crash, and
2173 /// it is what every other unservable request here already returns.
2174 ///
2175 /// `static`-only: every call site is inside `impl<T> TensorStorage<T>`
2176 /// (`TensorStorage::new`), itself `#[cfg(feature = "static")]`.
2177 #[cfg(feature = "static")]
2178 pub(crate) fn unsupported_here(self) -> Error {
2179 let reason = match self {
2180 // Not reachable: `Mem` is servable on every target. Stated
2181 // rather than panicked on, for the reason above.
2182 TensorMemory::Mem => "Mem is servable on every platform, so this should not happen",
2183 TensorMemory::Shm => "POSIX shared memory is unix-only",
2184 TensorMemory::DmaBuf => {
2185 "this target has no platform GPU-buffer allocator (DMA-BUF is Linux, \
2186 IOSurface macOS/iOS, AHardwareBuffer Android, D3D11 textures on \
2187 Windows)"
2188 }
2189 TensorMemory::IoSurface => {
2190 "no backend produces or accepts IOSurface under its own code yet -- \
2191 macOS/iOS allocate and report TensorMemory::DmaBuf"
2192 }
2193 TensorMemory::Pbo => {
2194 "PBO tensors are created by ImageProcessor::create_image with a GL \
2195 context current, never by Tensor::new"
2196 }
2197 TensorMemory::Cuda => "there is no CUDA backing in this build",
2198 };
2199 Error::NotImplemented(format!(
2200 "TensorMemory::{self:?} is not supported by this build: {reason}"
2201 ))
2202 }
2203}
2204
2205impl From<TensorMemory> for String {
2206 /// The wire string, via [`TensorMemory::as_str`] -- one spelling, not a
2207 /// second one maintained alongside it.
2208 ///
2209 /// `DmaBuf` renders as `"dmabuf"`, not the `"dma"` this impl used to
2210 /// emit. `TryFrom<&str>` still accepts `"dma"` so anything that
2211 /// persisted the old spelling keeps parsing; the benchmark result keys
2212 /// in `.github/scripts` spell it `"dma"` but build that string
2213 /// themselves and never went through this impl.
2214 fn from(memory: TensorMemory) -> Self {
2215 memory.as_str().to_owned()
2216 }
2217}
2218
2219impl TryFrom<&str> for TensorMemory {
2220 type Error = Error;
2221
2222 /// Parses [`TensorMemory::as_str`]'s spelling, plus the legacy `"dma"`
2223 /// alias for `DmaBuf`.
2224 fn try_from(s: &str) -> Result<Self> {
2225 TensorMemory::from_str_code(s)
2226 .or(if s == "dma" {
2227 Some(TensorMemory::DmaBuf)
2228 } else {
2229 None
2230 })
2231 .ok_or_else(|| Error::InvalidMemoryType(s.to_owned()))
2232 }
2233}
2234
2235// Everything from here through the end of `impl<T> TensorTrait<T> for
2236// Tensor<T>` below is the `static` backend's storage-owning `Tensor<T>`:
2237// `TensorStorage<T>` (the per-memory-backend enum) and the `Tensor<T>`
2238// struct that wraps it. `dynamic`'s `Tensor<T>` is a different, much
2239// smaller type -- a typed lens over `TensorDyn` with no storage of its
2240// own -- defined in `tensor_dyn/dynamic_tensor.rs` and re-exported below.
2241#[cfg(feature = "static")]
2242#[derive(Debug)]
2243#[allow(dead_code)] // Variants are constructed by downstream crates via pub(crate) helpers
2244pub(crate) enum TensorStorage<T>
2245where
2246 T: Num + Clone + fmt::Debug + Send + Sync,
2247{
2248 /// Platform-native zero-copy GPU buffer. Inner type differs per
2249 /// target: `DmaTensor` on Linux (DMA-BUF fd), `IoSurfaceTensor` on
2250 /// macOS (CFRetained IOSurface). The shared variant name keeps the
2251 /// public `TensorMemory::DmaBuf` discriminant stable across platforms.
2252 #[cfg(target_os = "linux")]
2253 Dma(DmaTensor<T>),
2254 #[cfg(any(target_os = "macos", target_os = "ios"))]
2255 Dma(IoSurfaceTensor<T>),
2256 #[cfg(target_os = "android")]
2257 Dma(AHardwareBufferTensor<T>),
2258 #[cfg(target_os = "windows")]
2259 Dma(crate::d3d11::texture::D3d11TextureTensor<T>),
2260 #[cfg(unix)]
2261 Shm(ShmTensor<T>),
2262 Mem(MemTensor<T>),
2263 Pbo(PboTensor<T>),
2264}
2265
2266#[cfg(feature = "static")]
2267impl<T> TensorStorage<T>
2268where
2269 T: Num + Clone + fmt::Debug + Send + Sync,
2270{
2271 /// The backing allocation's intrinsic physical row pitch in bytes, if it
2272 /// has one that is fixed independent of the logical shape. macOS IOSurface
2273 /// reports its 64-aligned `bytesPerRow`; other backings (Linux DMA, SHM,
2274 /// Mem, PBO) have no fixed pitch beyond the logical shape and return `None`.
2275 ///
2276 /// Used by `configure_image` to preserve the physical pitch when a reused
2277 /// pool tensor is reconfigured to a smaller logical image — so the decode
2278 /// writes rows at the surface's real stride and the GPU samples them with
2279 /// the same stride (the physical-grid / logical-ROI decoupling).
2280 pub(crate) fn backing_row_stride(&self) -> Option<usize> {
2281 match self {
2282 // Only genuine image-formatted IOSurfaces (height > 1) carry a real
2283 // per-row pitch; a generic byte-bag (height == 1) returns `None` so
2284 // `configure_image` does not adopt its whole-buffer "row" as a stride.
2285 #[cfg(any(target_os = "macos", target_os = "ios"))]
2286 TensorStorage::Dma(t) => t.image_backing_row_stride(),
2287 // Android AHardwareBuffer: same rule — only genuine 2D
2288 // image-formatted buffers (height > 1) carry a real row pitch.
2289 #[cfg(target_os = "android")]
2290 TensorStorage::Dma(t) => t.image_backing_row_stride(),
2291 // Windows D3D11: the driver's staging row pitch. `None` when the
2292 // tensor has no staging (write-only or hardware-only), whose host
2293 // buffer has tight rows -- except for a semi-planar tensor, whose
2294 // texture is allocated as wide as that pitch for every access, so
2295 // it always has one to report.
2296 #[cfg(target_os = "windows")]
2297 TensorStorage::Dma(t) => t.image_backing_row_stride(),
2298 _ => None,
2299 }
2300 }
2301
2302 /// Create a new tensor storage with the given shape, memory type, and
2303 /// optional name. If no name is given, a random name will be generated.
2304 /// If no memory type is given, the best available memory type will be
2305 /// chosen based on the platform and environment variables.
2306 fn new(shape: &[usize], memory: Option<TensorMemory>, name: Option<&str>) -> Result<Self> {
2307 match memory {
2308 #[cfg(target_os = "linux")]
2309 Some(TensorMemory::DmaBuf) => {
2310 DmaTensor::<T>::new(shape, name).map(TensorStorage::Dma)
2311 }
2312 #[cfg(any(target_os = "macos", target_os = "ios"))]
2313 Some(TensorMemory::DmaBuf) => {
2314 IoSurfaceTensor::<T>::new(shape, name).map(TensorStorage::Dma)
2315 }
2316 #[cfg(target_os = "android")]
2317 Some(TensorMemory::DmaBuf) => {
2318 AHardwareBufferTensor::<T>::new(shape, name).map(TensorStorage::Dma)
2319 }
2320 // Windows has a DmaBuf backing, but it is a D3D11 texture: an image
2321 // layout, not a byte bag, so a shape alone cannot describe one.
2322 // Name the constructor that can, rather than quietly handing back
2323 // host memory the caller did not ask for.
2324 #[cfg(target_os = "windows")]
2325 Some(TensorMemory::DmaBuf) => Err(crate::error::Error::NotImplemented(
2326 "D3D11 texture tensors are image-formatted; use Tensor::image with \
2327 TensorMemory::DmaBuf"
2328 .to_owned(),
2329 )),
2330 #[cfg(not(any(
2331 target_os = "linux",
2332 target_os = "macos",
2333 target_os = "ios",
2334 target_os = "android",
2335 target_os = "windows"
2336 )))]
2337 Some(TensorMemory::DmaBuf) => Err(crate::error::Error::NotImplemented(
2338 "TensorMemory::DmaBuf is only available on Linux (DMA-BUF), macOS/iOS (IOSurface), \
2339 Android (AHardwareBuffer), and Windows (D3D11 textures, via Tensor::image)"
2340 .to_owned(),
2341 )),
2342 #[cfg(unix)]
2343 Some(TensorMemory::Shm) => {
2344 ShmTensor::<T>::new(shape, name).map(TensorStorage::Shm)
2345 }
2346 #[cfg(not(unix))]
2347 Some(m @ TensorMemory::Shm) => Err(m.unsupported_here()),
2348 Some(TensorMemory::Mem) => {
2349 MemTensor::<T>::new(shape, name).map(TensorStorage::Mem)
2350 }
2351 Some(TensorMemory::Pbo) => Err(crate::error::Error::NotImplemented(
2352 "PboTensor cannot be created via Tensor::new() — use ImageProcessor::create_image()".to_owned(),
2353 )),
2354 // Defined codes with no backing behind them yet. An error, not
2355 // an `unreachable!()` -- see `unsupported_here`.
2356 Some(m @ (TensorMemory::IoSurface | TensorMemory::Cuda)) => Err(m.unsupported_here()),
2357 None => {
2358 if std::env::var("EDGEFIRST_TENSOR_FORCE_MEM")
2359 .is_ok_and(|x| x != "0" && x.to_lowercase() != "false")
2360 {
2361 MemTensor::<T>::new(shape, name).map(TensorStorage::Mem)
2362 } else {
2363 // Auto-select priority: Dma > Mem. Shm is intentionally NOT
2364 // auto-selected — it offers no advantage over Mem for an
2365 // in-process tensor and Mem always succeeds, so Shm below Mem
2366 // is effectively never reached. Request Shm explicitly via
2367 // `TensorMemory::Shm` when cross-process sharing is needed.
2368 // (PBO sits between Dma and Mem but is GL-backed and created
2369 // only via `ImageProcessor::create_image`.)
2370 #[cfg(target_os = "linux")]
2371 {
2372 // Linux: Try DMA -> Mem
2373 match DmaTensor::<T>::new(shape, name) {
2374 Ok(tensor) => Ok(TensorStorage::Dma(tensor)),
2375 Err(_) => MemTensor::<T>::new(shape, name).map(TensorStorage::Mem),
2376 }
2377 }
2378 #[cfg(any(target_os = "macos", target_os = "ios"))]
2379 {
2380 // macOS/iOS: Try IOSurface -> Mem. IOSurface is the
2381 // GPU-shareable backend (zero-copy via ANGLE), filling the
2382 // same role as DMA-BUF on Linux.
2383 match IoSurfaceTensor::<T>::new(shape, name) {
2384 Ok(tensor) => Ok(TensorStorage::Dma(tensor)),
2385 Err(_) => MemTensor::<T>::new(shape, name).map(TensorStorage::Mem),
2386 }
2387 }
2388 #[cfg(target_os = "android")]
2389 {
2390 // Android: Mem. Unlike macOS (where a byte-bag
2391 // IOSurface is still GL-importable as R8) a generic
2392 // BLOB AHardwareBuffer cannot back a GPU texture, so
2393 // auto-selecting it buys nothing and costs plenty: a
2394 // gralloc allocator-HAL ioctl per allocation (orders
2395 // slower than malloc, ≥1 page + a dmabuf fd even for
2396 // tiny tensors) and a lock/unlock cache-maintenance
2397 // round trip on every map(). Zero-copy image tensors
2398 // come from `Tensor::image(..)`; callers that want a
2399 // BLOB (NNAPI handoff) request `TensorMemory::DmaBuf`
2400 // explicitly.
2401 MemTensor::<T>::new(shape, name).map(TensorStorage::Mem)
2402 }
2403 #[cfg(all(
2404 unix,
2405 not(any(
2406 target_os = "linux",
2407 target_os = "macos",
2408 target_os = "ios",
2409 target_os = "android"
2410 ))
2411 ))]
2412 {
2413 // Other Unix (BSD): Mem only (no DMA; Shm is explicit-only)
2414 MemTensor::<T>::new(shape, name).map(TensorStorage::Mem)
2415 }
2416 #[cfg(not(unix))]
2417 {
2418 // Windows/other: Mem only
2419 MemTensor::<T>::new(shape, name).map(TensorStorage::Mem)
2420 }
2421 }
2422 }
2423 }
2424 }
2425
2426 /// Create a DMA-backed tensor storage with an explicit byte size that
2427 /// may exceed `shape.product() * sizeof(T)`. Used for image tensors
2428 /// with row-padded layouts (see `DmaTensor::new_with_byte_size`).
2429 ///
2430 /// This is intentionally DMA-only: padding is only meaningful for
2431 /// buffers that will be imported as GPU textures via EGLImage. PBO,
2432 /// Shm, and Mem storage doesn't benefit from pitch alignment and
2433 /// shouldn't pay the memory cost.
2434 #[cfg(target_os = "linux")]
2435 pub(crate) fn new_dma_with_byte_size(
2436 shape: &[usize],
2437 byte_size: usize,
2438 name: Option<&str>,
2439 ) -> Result<Self> {
2440 DmaTensor::<T>::new_with_byte_size(shape, byte_size, name).map(TensorStorage::Dma)
2441 }
2442
2443 // No non-Linux stub: the only caller (`Tensor::image_with_stride`)
2444 // returns `NotImplemented` directly on non-Linux without ever
2445 // reaching the storage layer, so defining a stub here would be
2446 // dead code and fail the `-D warnings` clippy gate on macOS CI.
2447
2448 /// Create a Mem-backed tensor storage with an explicit byte size that may
2449 /// exceed `shape.product() * sizeof(T)`. Used for image tensors with
2450 /// 64-byte-aligned row strides (see `MemTensor::with_capacity_bytes`).
2451 pub(crate) fn new_mem_with_byte_size(
2452 shape: &[usize],
2453 byte_size: usize,
2454 name: Option<&str>,
2455 ) -> Result<Self>
2456 where
2457 T: 'static,
2458 {
2459 MemTensor::<T>::with_capacity_bytes(shape, byte_size, name).map(TensorStorage::Mem)
2460 }
2461
2462 /// Create a Shm-backed tensor storage with an explicit byte size that may
2463 /// exceed `shape.product() * sizeof(T)`. Used for image tensors with
2464 /// 64-byte-aligned row strides (see `ShmTensor::new_with_byte_size`).
2465 #[cfg(unix)]
2466 pub(crate) fn new_shm_with_byte_size(
2467 shape: &[usize],
2468 byte_size: usize,
2469 name: Option<&str>,
2470 ) -> Result<Self> {
2471 ShmTensor::<T>::new_with_byte_size(shape, byte_size, name).map(TensorStorage::Shm)
2472 }
2473
2474 /// Allocate an image-formatted IOSurface-backed storage (macOS).
2475 ///
2476 /// Used by `Tensor::image()` when the caller requests
2477 /// `TensorMemory::DmaBuf` and the format has an IOSurface FourCC
2478 /// mapping (YUYV, RGBA, BGRA today). Falls back to `new_with_byte_size`
2479 /// otherwise.
2480 #[cfg(any(target_os = "macos", target_os = "ios"))]
2481 pub(crate) fn new_image_iosurface(
2482 width: usize,
2483 height: usize,
2484 format: PixelFormat,
2485 dtype: DType,
2486 shape: &[usize],
2487 name: Option<&str>,
2488 ) -> Result<Self> {
2489 IoSurfaceTensor::<T>::new_image(width, height, format, dtype, shape, name)
2490 .map(TensorStorage::Dma)
2491 }
2492
2493 /// Allocate an image-formatted AHardwareBuffer-backed storage (Android).
2494 ///
2495 /// Used by `Tensor::image()` when the caller requests
2496 /// `TensorMemory::DmaBuf` and the format has an AHardwareBuffer format
2497 /// mapping (RGBA8 and the RGBA16F float paths today). Falls back to
2498 /// `new_with_byte_size` otherwise.
2499 #[cfg(target_os = "android")]
2500 pub(crate) fn new_image_ahardwarebuffer(
2501 width: usize,
2502 height: usize,
2503 format: PixelFormat,
2504 dtype: DType,
2505 shape: &[usize],
2506 name: Option<&str>,
2507 access: CpuAccess,
2508 ) -> Result<Self> {
2509 AHardwareBufferTensor::<T>::new_image(width, height, format, dtype, shape, name, access)
2510 .map(TensorStorage::Dma)
2511 }
2512
2513 /// Allocate an image-formatted D3D11 texture storage (Windows).
2514 ///
2515 /// Used by `Tensor::image()` when the caller requests
2516 /// `TensorMemory::DmaBuf`. Unlike the IOSurface and AHardwareBuffer
2517 /// constructors there is no byte-bag fallback: a format with no texture
2518 /// layout is an error, because a texture is the only thing this backing
2519 /// can be.
2520 #[cfg(target_os = "windows")]
2521 pub(crate) fn new_image_d3d11(
2522 width: usize,
2523 height: usize,
2524 format: PixelFormat,
2525 dtype: DType,
2526 shape: &[usize],
2527 name: Option<&str>,
2528 access: CpuAccess,
2529 ) -> Result<Self> {
2530 crate::d3d11::texture::D3d11TextureTensor::<T>::new_image(
2531 width, height, format, dtype, shape, name, access,
2532 )
2533 .map(TensorStorage::Dma)
2534 }
2535
2536 /// Create a new tensor storage using the given file descriptor, shape,
2537 /// and optional name.
2538 #[cfg(unix)]
2539 fn from_fd(fd: OwnedFd, shape: &[usize], name: Option<&str>) -> Result<Self> {
2540 #[cfg(target_os = "linux")]
2541 {
2542 use nix::sys::stat::fstat;
2543 use nix::sys::statfs::fstatfs;
2544
2545 let stat = fstat(&fd)?;
2546 let major = major(stat.st_dev);
2547 let minor = minor(stat.st_dev);
2548
2549 if major != 0 {
2550 // Dma and Shm tensors are expected to have major number 0
2551 return Err(Error::UnknownDeviceType(major, minor));
2552 }
2553
2554 // Classify by filesystem magic, never by the st_dev minor.
2555 //
2556 // dma_buf files live on an internal kernel mount (`dma_buf_mnt`,
2557 // created with `kern_mount`) whose superblock draws its device
2558 // number from `get_anon_bdev()` — an IDA allocated first-come,
2559 // first-served during boot, shared with every other anonymous
2560 // pseudo-filesystem (pipefs, sockfs, anon_inodefs, nsfs, bdev,
2561 // tracefs, ...). The minor a DMA-BUF lands on is therefore a
2562 // function of which pseudo-filesystems registered first on that
2563 // kernel build: it moves with kernel config, driver load order,
2564 // initramfs use, and kernel version. It is not part of any ABI.
2565 //
2566 // Observed values for a genuine DMA-BUF: 12 on x86 desktop, 8 on
2567 // the ADIS Verdin. An earlier revision hardcoded `9 | 10` here,
2568 // which silently imported real DMA-BUFs as shared memory
2569 // everywhere else.
2570 //
2571 // The filesystem magic, by contrast, is stable UAPI
2572 // (include/uapi/linux/magic.h). Both branches are identified
2573 // positively; anything else is genuinely unknown and is rejected
2574 // rather than guessed at. Falling back to SHM would "work" — a
2575 // DMA-BUF is mmap-able, and even a pipe imports as a zero-length
2576 // tensor — which is exactly why it must not be the default.
2577 let magic = fs_magic(fstatfs(&fd)?.filesystem_type().0 as i64);
2578
2579 log::debug!(
2580 "Creating tensor from fd: major={major}, minor={minor}, magic={magic:#010x}"
2581 );
2582
2583 match magic {
2584 DMA_BUF_MAGIC => DmaTensor::<T>::from_fd(fd, shape, name).map(TensorStorage::Dma),
2585 TMPFS_MAGIC => ShmTensor::<T>::from_fd(fd, shape, name).map(TensorStorage::Shm),
2586 other => Err(Error::UnknownBufferType(other)),
2587 }
2588 }
2589 #[cfg(all(unix, not(target_os = "linux")))]
2590 {
2591 // On macOS/iOS/BSD, always use SHM (no DMA-BUF fd import)
2592 ShmTensor::<T>::from_fd(fd, shape, name).map(TensorStorage::Shm)
2593 }
2594 }
2595}
2596
2597#[cfg(feature = "static")]
2598impl<T> TensorTrait<T> for TensorStorage<T>
2599where
2600 T: Num + Clone + fmt::Debug + Send + Sync,
2601{
2602 fn new(shape: &[usize], name: Option<&str>) -> Result<Self> {
2603 Self::new(shape, None, name)
2604 }
2605
2606 #[cfg(unix)]
2607 fn from_fd(fd: OwnedFd, shape: &[usize], name: Option<&str>) -> Result<Self> {
2608 Self::from_fd(fd, shape, name)
2609 }
2610
2611 #[cfg(unix)]
2612 fn clone_fd(&self) -> Result<OwnedFd> {
2613 match self {
2614 #[cfg(any(
2615 target_os = "linux",
2616 target_os = "macos",
2617 target_os = "ios",
2618 target_os = "android"
2619 ))]
2620 TensorStorage::Dma(t) => t.clone_fd(),
2621 TensorStorage::Shm(t) => t.clone_fd(),
2622 TensorStorage::Mem(t) => t.clone_fd(),
2623 TensorStorage::Pbo(t) => t.clone_fd(),
2624 }
2625 }
2626
2627 fn memory(&self) -> TensorMemory {
2628 match self {
2629 #[cfg(any(
2630 target_os = "linux",
2631 target_os = "macos",
2632 target_os = "ios",
2633 target_os = "android",
2634 target_os = "windows"
2635 ))]
2636 TensorStorage::Dma(_) => TensorMemory::DmaBuf,
2637 #[cfg(unix)]
2638 TensorStorage::Shm(_) => TensorMemory::Shm,
2639 TensorStorage::Mem(_) => TensorMemory::Mem,
2640 TensorStorage::Pbo(_) => TensorMemory::Pbo,
2641 }
2642 }
2643
2644 fn name(&self) -> String {
2645 match self {
2646 #[cfg(any(
2647 target_os = "linux",
2648 target_os = "macos",
2649 target_os = "ios",
2650 target_os = "android",
2651 target_os = "windows"
2652 ))]
2653 TensorStorage::Dma(t) => t.name(),
2654 #[cfg(unix)]
2655 TensorStorage::Shm(t) => t.name(),
2656 TensorStorage::Mem(t) => t.name(),
2657 TensorStorage::Pbo(t) => t.name(),
2658 }
2659 }
2660
2661 fn shape(&self) -> &[usize] {
2662 match self {
2663 #[cfg(any(
2664 target_os = "linux",
2665 target_os = "macos",
2666 target_os = "ios",
2667 target_os = "android",
2668 target_os = "windows"
2669 ))]
2670 TensorStorage::Dma(t) => t.shape(),
2671 #[cfg(unix)]
2672 TensorStorage::Shm(t) => t.shape(),
2673 TensorStorage::Mem(t) => t.shape(),
2674 TensorStorage::Pbo(t) => t.shape(),
2675 }
2676 }
2677
2678 fn reshape(&mut self, shape: &[usize]) -> Result<()> {
2679 match self {
2680 #[cfg(any(
2681 target_os = "linux",
2682 target_os = "macos",
2683 target_os = "ios",
2684 target_os = "android",
2685 target_os = "windows"
2686 ))]
2687 TensorStorage::Dma(t) => t.reshape(shape),
2688 #[cfg(unix)]
2689 TensorStorage::Shm(t) => t.reshape(shape),
2690 TensorStorage::Mem(t) => t.reshape(shape),
2691 TensorStorage::Pbo(t) => t.reshape(shape),
2692 }
2693 }
2694
2695 fn capacity_bytes(&self) -> usize {
2696 match self {
2697 #[cfg(any(
2698 target_os = "linux",
2699 target_os = "macos",
2700 target_os = "ios",
2701 target_os = "android",
2702 target_os = "windows"
2703 ))]
2704 TensorStorage::Dma(t) => t.capacity_bytes(),
2705 #[cfg(unix)]
2706 TensorStorage::Shm(t) => t.capacity_bytes(),
2707 TensorStorage::Mem(t) => t.capacity_bytes(),
2708 TensorStorage::Pbo(t) => t.capacity_bytes(),
2709 }
2710 }
2711
2712 fn set_logical_shape(&mut self, shape: &[usize]) -> Result<()> {
2713 match self {
2714 #[cfg(any(
2715 target_os = "linux",
2716 target_os = "macos",
2717 target_os = "ios",
2718 target_os = "android",
2719 target_os = "windows"
2720 ))]
2721 TensorStorage::Dma(t) => t.set_logical_shape(shape),
2722 #[cfg(unix)]
2723 TensorStorage::Shm(t) => t.set_logical_shape(shape),
2724 TensorStorage::Mem(t) => t.set_logical_shape(shape),
2725 TensorStorage::Pbo(t) => t.set_logical_shape(shape),
2726 }
2727 }
2728
2729 fn map_with<'a>(&self, access: CpuAccess) -> Result<crate::view::HostView<'a, T>>
2730 where
2731 T: 'a,
2732 {
2733 match self {
2734 #[cfg(any(
2735 target_os = "linux",
2736 target_os = "macos",
2737 target_os = "ios",
2738 target_os = "android",
2739 target_os = "windows"
2740 ))]
2741 TensorStorage::Dma(t) => t.map_with(access),
2742 #[cfg(unix)]
2743 TensorStorage::Shm(t) => t.map_with(access),
2744 TensorStorage::Mem(t) => t.map_with(access),
2745 TensorStorage::Pbo(t) => t.map_with(access),
2746 }
2747 }
2748
2749 fn try_map_with<'a>(&self, access: CpuAccess) -> Result<crate::view::HostView<'a, T>>
2750 where
2751 T: 'a,
2752 {
2753 match self {
2754 #[cfg(any(
2755 target_os = "linux",
2756 target_os = "macos",
2757 target_os = "ios",
2758 target_os = "android",
2759 target_os = "windows"
2760 ))]
2761 TensorStorage::Dma(t) => t.try_map_with(access),
2762 #[cfg(unix)]
2763 TensorStorage::Shm(t) => t.try_map_with(access),
2764 TensorStorage::Mem(t) => t.try_map_with(access),
2765 TensorStorage::Pbo(t) => t.try_map_with(access),
2766 }
2767 }
2768
2769 fn buffer_identity(&self) -> &BufferIdentity {
2770 match self {
2771 #[cfg(any(
2772 target_os = "linux",
2773 target_os = "macos",
2774 target_os = "ios",
2775 target_os = "android",
2776 target_os = "windows"
2777 ))]
2778 TensorStorage::Dma(t) => t.buffer_identity(),
2779 #[cfg(unix)]
2780 TensorStorage::Shm(t) => t.buffer_identity(),
2781 TensorStorage::Mem(t) => t.buffer_identity(),
2782 TensorStorage::Pbo(t) => t.buffer_identity(),
2783 }
2784 }
2785
2786 /// Forward a sub-region view to the active backend, re-wrapping the
2787 /// backend's view (which shares the parent's allocation and
2788 /// [`BufferIdentity`]) back into the matching `TensorStorage` variant. Each
2789 /// backend's `view` is its own `TensorTrait::view` override; this single
2790 /// match is the only per-variant dispatch (`Tensor::subview` calls through
2791 /// here rather than matching the storage itself).
2792 fn view(&self, offset_bytes: usize, shape: &[usize]) -> Result<Self> {
2793 match self {
2794 #[cfg(any(
2795 target_os = "linux",
2796 target_os = "macos",
2797 target_os = "ios",
2798 target_os = "android",
2799 target_os = "windows"
2800 ))]
2801 TensorStorage::Dma(t) => t.view(offset_bytes, shape).map(TensorStorage::Dma),
2802 #[cfg(unix)]
2803 TensorStorage::Shm(t) => t.view(offset_bytes, shape).map(TensorStorage::Shm),
2804 TensorStorage::Mem(t) => t.view(offset_bytes, shape).map(TensorStorage::Mem),
2805 TensorStorage::Pbo(t) => t.view(offset_bytes, shape).map(TensorStorage::Pbo),
2806 }
2807 }
2808}
2809
2810/// Multi-backend tensor with optional image format metadata.
2811///
2812/// When `format` is `Some`, this tensor represents an image. Width, height,
2813/// and channels are derived from `shape` + `format`. When `format` is `None`,
2814/// this is a raw tensor (identical to the pre-refactoring behavior).
2815///
2816/// This is the `static` backend's `Tensor<T>` -- it owns its storage. See
2817/// `tensor_dyn/dynamic_tensor.rs` for `dynamic`'s `Tensor<T>`, a typed lens
2818/// with no storage of its own.
2819#[cfg(feature = "static")]
2820#[derive(Debug)]
2821pub struct Tensor<T>
2822where
2823 T: Num + Clone + fmt::Debug + Send + Sync,
2824{
2825 /// CUDA registration for this tensor, if any. Set after creation by
2826 /// the image crate once a PBO is registered with CUDA interop.
2827 ///
2828 /// MUST be declared before `storage`: CUDA must unregister the GL buffer
2829 /// before storage's Drop deletes it (cudaGraphicsUnregisterResource before
2830 /// glDeleteBuffers). Rust drops fields in declaration order.
2831 cuda: Option<crate::cuda::CudaHandle>,
2832 pub(crate) storage: TensorStorage<T>,
2833 format: Option<PixelFormat>,
2834 chroma: Option<Box<Tensor<T>>>,
2835 /// Row stride in bytes for externally allocated buffers with row padding.
2836 /// `None` means tightly packed (stride == width * bytes_per_pixel).
2837 row_stride: Option<usize>,
2838 /// Byte offset within the DMA-BUF where image data starts.
2839 /// `None` means offset 0 (data starts at the beginning of the buffer).
2840 plane_offset: Option<usize>,
2841 /// Quantization metadata for integer-typed tensors. Public access is
2842 /// gated by the `IntegerType` trait — `Tensor<f32>` etc. carry the
2843 /// field for layout uniformity but have no way to read or write it.
2844 pub(crate) quantization: Option<Quantization>,
2845 /// Optional colorimetry metadata. `None` = undefined; never auto-filled.
2846 colorimetry: Option<crate::Colorimetry>,
2847 /// Declared CPU access (see [`CpuAccess`]). Image constructors set it
2848 /// from their `access` parameter; non-image tensors (`new`, imports,
2849 /// numpy) default to `ReadWrite` — they are CPU-centric by nature.
2850 /// `map_with` counts requests beyond this declaration as unplanned.
2851 cpu_access: CpuAccess,
2852 /// Recorded vendor tile-compression scheme — best knowledge from
2853 /// allocation time (see [`Compression`]). `Some` only for Android
2854 /// hardware-only AHardwareBuffers whose allocation requested
2855 /// compression on an eligible format with a recognized vendor. When
2856 /// set, the row-stride accessors describe no meaningful linear
2857 /// layout. `configure_image` preserves it (physical layout, unlike
2858 /// colorimetry); views inherit it.
2859 compression: Option<CompressionScheme>,
2860 /// Parent-image snapshot when this tensor is a [`view`](Self::view)/
2861 /// [`batch`](Self::batch) sub-region; `None` for a whole tensor. Lets the GL
2862 /// backend key its import on the parent and render the view as a
2863 /// `glViewport`/`glScissor` ROI. See [`ViewOrigin`].
2864 view_origin: Option<ViewOrigin>,
2865}
2866
2867#[cfg(feature = "static")]
2868impl<T> Tensor<T>
2869where
2870 T: Num + Clone + fmt::Debug + Send + Sync,
2871{
2872 /// Wrap a TensorStorage in a Tensor with no image metadata.
2873 pub(crate) fn wrap(storage: TensorStorage<T>) -> Self {
2874 Self {
2875 storage,
2876 format: None,
2877 chroma: None,
2878 row_stride: None,
2879 plane_offset: None,
2880 quantization: None,
2881 cuda: None,
2882 colorimetry: None,
2883 cpu_access: CpuAccess::ReadWrite,
2884 compression: None,
2885 view_origin: None,
2886 }
2887 }
2888
2889 /// Construct a tensor from a row-major element slice + shape. Allocates a
2890 /// new buffer (`TensorMemory::Mem`) and memcpys the contents; caller
2891 /// retains ownership of the input slice.
2892 ///
2893 /// # Errors
2894 ///
2895 /// - [`Error::InvalidShape`] if `values.len() != shape.iter().product()`.
2896 /// - Propagates any allocation error from [`Self::new`].
2897 pub fn from_slice(values: &[T], shape: &[usize]) -> Result<Self>
2898 where
2899 T: Copy,
2900 {
2901 let expected: usize = shape.iter().product();
2902 if values.len() != expected {
2903 return Err(Error::InvalidShape(format!(
2904 "from_slice: values.len()={} but shape product={expected} (shape={shape:?})",
2905 values.len()
2906 )));
2907 }
2908 let t = Self::new(shape, Some(TensorMemory::Mem), None)?;
2909 {
2910 let mut m = t.map()?;
2911 m.as_mut_slice().copy_from_slice(values);
2912 }
2913 Ok(t)
2914 }
2915
2916 /// Wrap externally-owned memory as a tensor without copying. The tensor
2917 /// borrows `[ptr, ptr + shape.product() * size_of::<T>())` as
2918 /// [`TensorMemory::Mem`]; `owner`, when `Some`, co-owns the source so it
2919 /// outlives the tensor (and all derived views/maps). See [`ForeignOwner`].
2920 ///
2921 /// The canonical use is CUDA zero-copy: allocate host-coherent memory
2922 /// (`cudaHostAlloc`), wrap the host pointer here, and bind the matching
2923 /// device pointer to the inference engine — reads and writes hit the same
2924 /// physical buffer with no host copy. The identical primitive backs the
2925 /// Python `Tensor.from_numpy` zero-copy borrow (owner = the NumPy object).
2926 ///
2927 /// # Safety
2928 ///
2929 /// `ptr` must be non-null, aligned to `align_of::<T>()`, and valid for
2930 /// `shape.product()` elements of `T` for as long as the returned tensor —
2931 /// and every view/map sharing its backing — is alive. Pass an `owner` that
2932 /// co-owns the source to uphold that contract.
2933 ///
2934 /// # Errors
2935 ///
2936 /// [`Error::InvalidSize`] if `shape` is empty.
2937 pub unsafe fn from_foreign(
2938 ptr: *mut T,
2939 shape: &[usize],
2940 owner: Option<crate::ForeignOwner>,
2941 name: Option<&str>,
2942 ) -> Result<Self> {
2943 // SAFETY: forwarded to `from_foreign_with_capacity`'s contract; `0`
2944 // clamps to the tight shape-derived footprint below, i.e. no extra
2945 // headroom -- identical to this function's own prior behaviour.
2946 unsafe { Self::from_foreign_with_capacity(ptr, shape, 0, owner, name) }
2947 }
2948
2949 /// [`Self::from_foreign`], additionally recording that the allocation
2950 /// behind `ptr` holds at least `capacity_bytes` bytes -- headroom beyond
2951 /// the tight `shape.product() * size_of::<T>()` footprint that a later
2952 /// [`configure_image`](Self::configure_image) / `set_logical_shape` can
2953 /// grow into without erroring.
2954 ///
2955 /// This is the consumer half of the cross-package capsule protocol's
2956 /// [`crate::protocol::TensorDesc::capacity`]: a producer's pool tensor
2957 /// (allocated for its largest expected image, holding a smaller one
2958 /// today) reports its real allocation size so the imported alias
2959 /// inherits the same headroom instead of being clamped to today's
2960 /// logical shape. `capacity_bytes` below the tight footprint is clamped
2961 /// up to it, never down -- the shape-implied floor always holds.
2962 ///
2963 /// # Safety
2964 ///
2965 /// Same as [`Self::from_foreign`], except `ptr` must be valid for
2966 /// `capacity_bytes` bytes (or the tight shape footprint, whichever is
2967 /// larger), not just the shape-implied footprint.
2968 ///
2969 /// # Errors
2970 ///
2971 /// [`Error::InvalidSize`] if `shape` is empty.
2972 pub unsafe fn from_foreign_with_capacity(
2973 ptr: *mut T,
2974 shape: &[usize],
2975 capacity_bytes: usize,
2976 owner: Option<crate::ForeignOwner>,
2977 name: Option<&str>,
2978 ) -> Result<Self> {
2979 if shape.is_empty() {
2980 return Err(Error::InvalidSize(0));
2981 }
2982 if ptr.is_null() {
2983 return Err(Error::InvalidArgument(
2984 "from_foreign: ptr must be non-null".to_owned(),
2985 ));
2986 }
2987 shape
2988 .iter()
2989 .copied()
2990 .try_fold(1usize, |acc, dim| acc.checked_mul(dim))
2991 .ok_or_else(|| {
2992 Error::InvalidArgument(format!(
2993 "from_foreign: shape.product() overflows usize (shape={shape:?})"
2994 ))
2995 })?;
2996 let capacity_elems = capacity_bytes / std::mem::size_of::<T>();
2997 // SAFETY: this function's own contract (see its doc comment) is
2998 // exactly `MemTensor::from_foreign_with_capacity`'s contract.
2999 let mem = unsafe {
3000 MemTensor::<T>::from_foreign_with_capacity(ptr, shape, capacity_elems, owner, name)
3001 };
3002 Ok(Self::wrap(TensorStorage::Mem(mem)))
3003 }
3004
3005 /// Pin a stable host address for this tensor's data.
3006 ///
3007 /// The returned [`HostPin`] carries **no borrow of the tensor**, so a
3008 /// caller can pin once at init and still pass `&mut self` to
3009 /// `ImageProcessor::convert` every frame. That borrow conflict is the
3010 /// blocker reported in [#134](https://github.com/EdgeFirstAI/hal/issues/134).
3011 ///
3012 /// The address stays valid until the pin is dropped, and survives
3013 /// re-layout: `configure_image` and `set_logical_shape` only rewrite the
3014 /// logical shape after a capacity check, and every backend's pin owns its
3015 /// own mapping, so nothing can move underneath it.
3016 ///
3017 /// A pin is therefore a **snapshot**, not a tracker: [`HostPin::len`]
3018 /// reports the extent at pin time and does not follow a later re-layout.
3019 /// That is the safe direction — the window always stays inside the
3020 /// allocation — but it means a caller who re-layouts should re-pin if it
3021 /// wants the new extent.
3022 ///
3023 /// The address is **not** a coherency guarantee: on non-coherent backends
3024 /// bracket CPU access with
3025 /// [`sync_for_cpu`](Self::sync_for_cpu) and
3026 /// [`sync_for_device`](Self::sync_for_device), which are no-ops where the
3027 /// backend does not require them.
3028 ///
3029 /// # Cost
3030 ///
3031 /// On `Mem` this is an `Arc` refcount bump — the backing allocation is
3032 /// never resized, so its base pointer is already stable. Other backends
3033 /// establish a mapping once and reuse it.
3034 ///
3035 /// # Errors
3036 ///
3037 /// - [`Error::InvalidArgument`] when `access` is [`CpuAccess::None`]: a
3038 /// pinned mapping is CPU access by definition, so this is an error
3039 /// rather than a warn-and-stage.
3040 /// - [`Error::NotImplemented`] for the two backends that **cannot** pin,
3041 /// rather than merely not having been wired yet: a PBO
3042 /// (`glMapBufferRange` establishes the address and the visibility
3043 /// together, and there is no `GL_MAP_PERSISTENT_BIT` in GLES) and an
3044 /// AHardwareBuffer (`AHardwareBuffer_lock` is what yields the address at
3045 /// all, so there is no address outside the lock). Both name
3046 /// `map()`/`map_with()` as the alternative. `Mem`, `Shm`, DMA-BUF,
3047 /// IOSurface and the Windows D3D11 texture all pin — the texture through
3048 /// a persistent host buffer the storage owns, which
3049 /// [`sync_for_cpu`](Self::sync_for_cpu) and
3050 /// [`sync_for_device`](Self::sync_for_device) move bytes between and the
3051 /// texture, rather than through the staging mapping (a pin outlives every
3052 /// map guard, so holding the staging texture mapped would freeze it).
3053 pub fn pin_host<'a>(&self, access: CpuAccess) -> Result<HostPin<'a>>
3054 where
3055 T: 'a,
3056 {
3057 if access == CpuAccess::None {
3058 return Err(Error::InvalidArgument(
3059 "pin_host: CpuAccess::None is not a valid pin; a pinned mapping \
3060 is CPU access by definition"
3061 .to_owned(),
3062 ));
3063 }
3064 // Backends yield everything addressable from the tensor's offset;
3065 // narrow to the logical extent so a consumer cannot mistake pitch
3066 // padding for data. A map guard keeps the wider window.
3067 let logical = self.len() * std::mem::size_of::<T>();
3068 let pin = match &self.storage {
3069 TensorStorage::Mem(m) => Ok(m.host_pin()),
3070 #[cfg(unix)]
3071 TensorStorage::Shm(s) => s.host_pin(),
3072 #[cfg(target_os = "linux")]
3073 TensorStorage::Dma(d) => d.host_pin(),
3074 #[cfg(any(target_os = "macos", target_os = "ios"))]
3075 TensorStorage::Dma(s) => s.host_pin(),
3076 // Windows D3D11: the pin is the storage's own persistent host
3077 // buffer, not the staging mapping -- see `D3d11TextureTensor::host_pin`.
3078 // The direction goes with it: a writable pin is what tells
3079 // `sync_for_device` the host buffer, not the staging copy, carries
3080 // the caller's writes.
3081 #[cfg(target_os = "windows")]
3082 TensorStorage::Dma(t) => t.host_pin(access),
3083 // AHardwareBuffer cannot support a pin at all, and that is a
3084 // platform contract rather than missing work: the NDK exposes no
3085 // way to obtain a CPU address outside AHardwareBuffer_lock, so a
3086 // "persistent" address would mean holding the lock for the pin's
3087 // lifetime -- which blocks the GPU from the buffer and makes any
3088 // further lock undefined. Both are the opposite of what a pin is
3089 // for. Say so, rather than implying a TODO.
3090 #[cfg(target_os = "android")]
3091 TensorStorage::Dma(_) => Err(Error::NotImplemented(
3092 "pin_host: AHardwareBuffer has no host address outside \
3093 AHardwareBuffer_lock, so it cannot hand out one that outlives \
3094 a guard. Use map()/map_with() for CPU access; for zero-copy \
3095 handoff pass the buffer itself."
3096 .to_owned(),
3097 )),
3098 // PBO is the same shape of constraint as AHardwareBuffer. The
3099 // message is a shared constant because `TensorDyn::pin_host`
3100 // (`tensor_dyn/dynamic_backend.rs`) must repeat this refusal --
3101 // the ABI has `ef_tensor_map` and no `ef_tensor_pin_host`, so the
3102 // library cannot tell a pin from a map. See
3103 // `crate::pbo::PIN_HOST_PBO_REFUSAL` for the reasoning at length.
3104 TensorStorage::Pbo(_) => Err(Error::NotImplemented(
3105 crate::pbo::PIN_HOST_PBO_REFUSAL.to_owned(),
3106 )),
3107 }?;
3108 Ok(pin.narrowed(logical))
3109 }
3110
3111 /// Make CPU reads/writes of this tensor coherent — the CPU is about to
3112 /// access the buffer.
3113 ///
3114 /// Corresponds to `DMA_BUF_SYNC_START`, and the name matches the Linux
3115 /// kernel DMA API for the same reason: `for_cpu` says who is about to
3116 /// touch the memory, which is what the caller actually knows.
3117 ///
3118 /// Pair with [`sync_for_device`](Self::sync_for_device).
3119 ///
3120 /// # Which backends owe maintenance
3121 ///
3122 /// | backend | behaviour |
3123 /// |---|---|
3124 /// | `Mem`, `Shm` | **no-op** — plain CPU memory, coherent by construction |
3125 /// | `Dma` (Linux DMA-BUF) | one `DMA_BUF_IOCTL_SYNC`, direction-tagged |
3126 /// | `Dma` (macOS/iOS IOSurface) | `IOSurfaceLock` / `IOSurfaceUnlock`, options matched to `access` |
3127 /// | `Dma` (Android AHardwareBuffer) | [`Error::NotImplemented`] — no sync independent of the lock |
3128 /// | `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 |
3129 /// | `Pbo` | [`Error::NotImplemented`] — no coherency window independent of the map |
3130 ///
3131 /// Only the first row is a no-op, so "bracket unconditionally and pay
3132 /// nothing" holds for `Mem`/`Shm` alone. The two that report
3133 /// `NotImplemented` do so because their platform fuses addressing and
3134 /// coherency into one call — `AHardwareBuffer_lock` and
3135 /// `glMapBufferRange` each perform the maintenance *as* they hand out the
3136 /// address, so a separate bracket would have to nest inside a map guard
3137 /// that already owns the pair. Both name `map()`/`map_with()` as the
3138 /// alternative. They previously fell through a wildcard and answered
3139 /// `Ok(())`, which is worse than an error because the symptom is
3140 /// stale-but-plausible pixels rather than a failure. The match is
3141 /// exhaustive so a new backend cannot inherit that silence — it must state
3142 /// its coherency story to compile.
3143 ///
3144 /// # Cost
3145 ///
3146 /// `Mem`/`Shm`: nothing. `Dma` on Linux: one ioctl. IOSurface: one
3147 /// lock/unlock pair. Windows D3D11 texture: one `CopyResource` plus a
3148 /// staging map per direction, and nothing at all on a tensor no one has
3149 /// pinned.
3150 ///
3151 /// # Errors
3152 ///
3153 /// [`Error::InvalidArgument`] when `access` is [`CpuAccess::None`], which
3154 /// describes no access to make coherent.
3155 pub fn sync_for_cpu(&self, access: CpuAccess) -> Result<()> {
3156 Self::check_sync_access(access)?;
3157 match &self.storage {
3158 // Plain CPU memory: coherent by construction.
3159 TensorStorage::Mem(_) => Ok(()),
3160 #[cfg(unix)]
3161 TensorStorage::Shm(_) => Ok(()),
3162 #[cfg(target_os = "linux")]
3163 TensorStorage::Dma(d) => {
3164 // The direction tells the kernel which maintenance this access
3165 // needs: a read-only bracket skips the writeback at END, a
3166 // write-only one skips the invalidate at START. It MUST match
3167 // the direction passed to sync_for_device.
3168 crate::dmabuf::sync_access(&d.fd, true, access).map_err(Error::NixError)
3169 }
3170 #[cfg(any(target_os = "macos", target_os = "ios"))]
3171 TensorStorage::Dma(s) => s.sync_for_cpu(access),
3172 // Windows D3D11: refreshes the pinned host buffer from the texture
3173 // through the staging copy.
3174 #[cfg(target_os = "windows")]
3175 TensorStorage::Dma(t) => t.sync_for_cpu(access),
3176 #[cfg(target_os = "android")]
3177 TensorStorage::Dma(_) => Err(Error::NotImplemented(
3178 "sync_for_cpu: AHardwareBuffer has no sync independent of its \
3179 lock. AHardwareBuffer_lock both establishes the address and \
3180 performs the cache invalidate, and the NDK leaves locking an \
3181 already-locked buffer undefined, so a bracket on top of a map \
3182 guard would nest. Use map()/map_with(), which owns the pair."
3183 .to_owned(),
3184 )),
3185 TensorStorage::Pbo(_) => Err(Error::NotImplemented(
3186 "sync_for_cpu: a PBO has no coherency window independent of its \
3187 map -- glMapBufferRange establishes the address and the \
3188 visibility together, and glUnmapBuffer publishes the writes. \
3189 Use map()/map_with(), which owns the pair."
3190 .to_owned(),
3191 )),
3192 }
3193 }
3194
3195 /// Release the buffer back to the device — the CPU is done accessing it.
3196 ///
3197 /// Corresponds to `DMA_BUF_SYNC_END`. See
3198 /// [`sync_for_cpu`](Self::sync_for_cpu) for the pairing and cost.
3199 pub fn sync_for_device(&self, access: CpuAccess) -> Result<()> {
3200 Self::check_sync_access(access)?;
3201 match &self.storage {
3202 TensorStorage::Mem(_) => Ok(()),
3203 #[cfg(unix)]
3204 TensorStorage::Shm(_) => Ok(()),
3205 #[cfg(target_os = "linux")]
3206 TensorStorage::Dma(d) => {
3207 crate::dmabuf::sync_access(&d.fd, false, access).map_err(Error::NixError)
3208 }
3209 #[cfg(any(target_os = "macos", target_os = "ios"))]
3210 TensorStorage::Dma(s) => s.sync_for_device(access),
3211 // Windows D3D11: uploads the pinned host buffer to the texture.
3212 #[cfg(target_os = "windows")]
3213 TensorStorage::Dma(t) => t.sync_for_device(access),
3214 #[cfg(target_os = "android")]
3215 TensorStorage::Dma(_) => Err(Error::NotImplemented(
3216 "sync_for_device: AHardwareBuffer has no sync independent of \
3217 its lock — see sync_for_cpu. Use map()/map_with()."
3218 .to_owned(),
3219 )),
3220 TensorStorage::Pbo(_) => Err(Error::NotImplemented(
3221 "sync_for_device: a PBO has no coherency window independent of \
3222 its map -- see sync_for_cpu. Use map()/map_with()."
3223 .to_owned(),
3224 )),
3225 }
3226 }
3227
3228 fn check_sync_access(access: CpuAccess) -> Result<()> {
3229 if access == CpuAccess::None {
3230 return Err(Error::InvalidArgument(
3231 "sync_for_cpu/sync_for_device: CpuAccess::None describes no \
3232 access to synchronise"
3233 .to_owned(),
3234 ));
3235 }
3236 Ok(())
3237 }
3238
3239 /// Construct a tensor from a 3-D ndarray view. Respects strides — one
3240 /// copy in all cases; contiguous views take a memcpy fast path.
3241 ///
3242 /// Only available when the `ndarray` feature is enabled.
3243 #[cfg(feature = "ndarray")]
3244 pub fn from_arrayview3(view: ndarray::ArrayView3<'_, T>) -> Result<Self>
3245 where
3246 T: Copy,
3247 {
3248 let (h, w, c) = view.dim();
3249 let t = Self::new(&[h, w, c], Some(TensorMemory::Mem), None)?;
3250 {
3251 let mut m = t.map()?;
3252 let dst = m.as_mut_slice();
3253 if let Some(src) = view.as_slice() {
3254 dst.copy_from_slice(src);
3255 } else {
3256 for (d, &s) in dst.iter_mut().zip(view.iter()) {
3257 *d = s;
3258 }
3259 }
3260 }
3261 Ok(t)
3262 }
3263
3264 /// Create a new tensor with the given shape, memory type, and optional
3265 /// name. If no name is given, a random name will be generated. If no
3266 /// memory type is given, the best available memory type will be chosen
3267 /// based on the platform and environment variables.
3268 ///
3269 /// On Linux platforms, the order of preference is: Dma -> Shm -> Mem.
3270 /// On other Unix platforms (macOS), the order is: Shm -> Mem.
3271 /// On non-Unix platforms, only Mem is available.
3272 ///
3273 /// # Environment Variables
3274 /// - `EDGEFIRST_TENSOR_FORCE_MEM`: If set to a non-zero and non-false
3275 /// value, forces the use of regular system memory allocation
3276 /// (`TensorMemory::Mem`) regardless of platform capabilities.
3277 ///
3278 /// # Example
3279 /// ```rust
3280 /// use edgefirst_tensor::{Error, Tensor, TensorMemory, TensorTrait};
3281 /// # fn main() -> Result<(), Error> {
3282 /// let tensor = Tensor::<f32>::new(&[2, 3, 4], Some(TensorMemory::Mem), Some("test_tensor"))?;
3283 /// assert_eq!(tensor.memory(), TensorMemory::Mem);
3284 /// assert_eq!(tensor.name(), "test_tensor");
3285 /// # Ok(())
3286 /// # }
3287 /// ```
3288 pub fn new(shape: &[usize], memory: Option<TensorMemory>, name: Option<&str>) -> Result<Self> {
3289 let _span = tracing::trace_span!(
3290 "tensor.alloc",
3291 ?shape,
3292 memory = ?memory,
3293 dtype = std::any::type_name::<T>(),
3294 )
3295 .entered();
3296 #[cfg_attr(not(target_os = "linux"), allow(unused_mut))]
3297 let mut t = TensorStorage::new(shape, memory, name).map(Self::wrap)?;
3298 // Best-effort: attach a CUDA ExternalMemory handle for DMA tensors on
3299 // CUDA-capable hosts. Never blocks tensor creation on failure.
3300 // RUNTIME-UNVALIDATED: no CUDA+dma_heap test platform available; ABI
3301 // layout-asserted vs. CUDA 12.6 driver_types.h; mechanism proven by
3302 // gpu-probe O5 on Orin.
3303 #[cfg(target_os = "linux")]
3304 t.try_init_dma_cuda();
3305 Ok(t)
3306 }
3307
3308 /// Create an image tensor with the given format.
3309 /// Allocate an image tensor from a declarative request — the
3310 /// full-featured constructor behind [`Self::image`] and friends.
3311 ///
3312 /// Adds the [`Compression`] request to the classic parameters:
3313 ///
3314 /// - a request with any CPU access other than [`CpuAccess::None`] is
3315 /// [`Error::InvalidArgument`] (CPU mapping pins the layout linear);
3316 /// - [`Compression::Scheme`] fails with [`Error::NotImplemented`] on
3317 /// platforms without vendor tile compression, and with
3318 /// [`Error::InvalidArgument`] when the device's native scheme or
3319 /// the format eligibility doesn't match;
3320 /// - [`Compression::Any`] never fails for compression reasons: it
3321 /// records the scheme when the allocation is eligible and
3322 /// otherwise resolves linear, incrementing
3323 /// [`compression_fallback_count`].
3324 ///
3325 /// The recorded outcome is readable via [`Tensor::compression`].
3326 pub fn image_desc(desc: &ImageDesc) -> Result<Self>
3327 where
3328 T: 'static,
3329 {
3330 let Some(t_dtype) = dtype_of::<T>() else {
3331 return Err(Error::InvalidArgument(
3332 "image_desc: element type has no DType mapping".into(),
3333 ));
3334 };
3335 if t_dtype != desc.dtype {
3336 return Err(Error::InvalidArgument(format!(
3337 "image_desc: desc.dtype is {:?} but the tensor element type is {t_dtype:?}",
3338 desc.dtype
3339 )));
3340 }
3341
3342 // Compression-request guards. CPU access pins the layout linear,
3343 // so a request combined with any declared access is a
3344 // contradiction the caller should hear about immediately.
3345 if desc.compression.is_some() && desc.access != CpuAccess::None {
3346 return Err(Error::InvalidArgument(format!(
3347 "image_desc: a compression request requires CpuAccess::None (declared {:?}) — CPU mapping pins the layout linear",
3348 desc.access
3349 )));
3350 }
3351 if let Some(Compression::Scheme(requested)) = desc.compression {
3352 #[cfg(not(target_os = "android"))]
3353 {
3354 return Err(Error::NotImplemented(format!(
3355 "image_desc: Compression::Scheme({requested:?}) — no vendor tile compression on this platform (request Compression::Any for a portable fallback)"
3356 )));
3357 }
3358 #[cfg(target_os = "android")]
3359 {
3360 if matches!(
3361 desc.memory,
3362 Some(TensorMemory::Mem) | Some(TensorMemory::Shm) | Some(TensorMemory::Pbo)
3363 ) {
3364 return Err(Error::InvalidArgument(format!(
3365 "image_desc: Compression::Scheme({requested:?}) requires hardware memory (TensorMemory::DmaBuf or auto-select), got {:?}",
3366 desc.memory
3367 )));
3368 }
3369 if !crate::ahardwarebuffer_layout::compression_eligible(desc.format, desc.dtype) {
3370 return Err(Error::InvalidArgument(format!(
3371 "image_desc: ({:?}, {:?}) is not compression-eligible (RGBA8888 u8/i8 initially)",
3372 desc.format, desc.dtype
3373 )));
3374 }
3375 let device = crate::ahardwarebuffer::device_compression_scheme();
3376 if device != Some(requested) {
3377 return Err(Error::InvalidArgument(format!(
3378 "image_desc: Compression::Scheme({requested:?}) but the device's native scheme is {device:?}"
3379 )));
3380 }
3381 }
3382 }
3383
3384 // A compression request implies a hardware pipeline, so auto
3385 // memory promotes to the platform's zero-copy allocation first
3386 // (`Tensor::image` only takes the AHardwareBuffer/IOSurface path
3387 // under an explicit Dma request). `Scheme` propagates the Dma
3388 // failure — the caller demanded a layout only that allocator can
3389 // produce; `Any` falls back to plain auto-select.
3390 #[allow(unused_mut)]
3391 let mut t = match (desc.memory, desc.compression) {
3392 (None, Some(request)) => {
3393 match Self::image(
3394 desc.width,
3395 desc.height,
3396 desc.format,
3397 Some(TensorMemory::DmaBuf),
3398 desc.access,
3399 ) {
3400 Ok(t) => t,
3401 Err(e) if matches!(request, Compression::Scheme(_)) => return Err(e),
3402 Err(_) => Self::image(desc.width, desc.height, desc.format, None, desc.access)?,
3403 }
3404 }
3405 (memory, _) => Self::image(desc.width, desc.height, desc.format, memory, desc.access)?,
3406 };
3407
3408 // Record best knowledge / count fallbacks. Only an Android
3409 // hardware-only AHardwareBuffer allocation can actually hold a
3410 // vendor tile layout; everywhere else an Any request resolves
3411 // linear and is counted.
3412 if let Some(request) = desc.compression {
3413 #[cfg(target_os = "android")]
3414 {
3415 let eligible =
3416 crate::ahardwarebuffer_layout::compression_eligible(desc.format, desc.dtype);
3417 let scheme = crate::ahardwarebuffer::device_compression_scheme();
3418 let is_ahb = t.memory() == TensorMemory::DmaBuf;
3419 match (request, scheme) {
3420 (_, Some(s)) if eligible && is_ahb => {
3421 t.set_compression_unchecked(Some(s));
3422 }
3423 (Compression::Scheme(requested), _) => {
3424 // Pre-validated eligible + scheme match, so the only
3425 // way here is the allocation resolving off-AHB.
3426 return Err(Error::InvalidOperation(format!(
3427 "image_desc: Compression::Scheme({requested:?}) requested but the allocation resolved to {:?} (not an AHardwareBuffer)",
3428 t.memory()
3429 )));
3430 }
3431 (Compression::Any, _) => {
3432 note_compression_fallback(&format!(
3433 "({:?}, {:?}) {}x{}: eligible={eligible}, scheme={scheme:?}, memory={:?}",
3434 desc.format,
3435 desc.dtype,
3436 desc.width,
3437 desc.height,
3438 t.memory()
3439 ));
3440 }
3441 }
3442 }
3443 #[cfg(not(target_os = "android"))]
3444 {
3445 debug_assert!(matches!(request, Compression::Any));
3446 let _ = request;
3447 note_compression_fallback(&format!(
3448 "({:?}, {:?}) {}x{}: no vendor tile compression on this platform",
3449 desc.format, desc.dtype, desc.width, desc.height
3450 ));
3451 }
3452 }
3453 Ok(t)
3454 }
3455
3456 /// Allocate an image tensor of `width` × `height` in `format`.
3457 ///
3458 /// This is the common-case image constructor. The shape follows from the
3459 /// format's layout — packed `[H, W, C]`, planar `[C, H, W]`, semi-planar
3460 /// `[H + chroma_rows, W]` — and the format is recorded on the tensor, so
3461 /// the image crate and the codec can work from it without out-of-band
3462 /// parameters.
3463 ///
3464 /// `memory` requests a backend, or `None` to auto-select. Auto-select for
3465 /// images is **DMA → Mem**: unlike [`Tensor::new`] it does not try SHM,
3466 /// which offers an image nothing that Mem doesn't (it is not
3467 /// GPU-importable, and Mem always succeeds). Ask for
3468 /// [`TensorMemory::Shm`] explicitly if you need it.
3469 ///
3470 /// `access` declares what the **CPU** intends to do with the buffer and is
3471 /// required; see [`CpuAccess`]. Hardware access — GPU, NPU, ISP, codec —
3472 /// is always implied and never declared. Prefer the narrowest declaration
3473 /// that fits: `None` keeps the allocation eligible for Android vendor tile
3474 /// compression, and anything else pins the layout linear.
3475 ///
3476 /// DMA-backed images always get a **64-byte-aligned row stride**, in every
3477 /// layout, because Mali and Vivante reject an EGLImage import at an
3478 /// unaligned pitch. Read the real pitch back with
3479 /// [`Tensor::effective_row_stride`] rather than assuming `width × bpp`.
3480 ///
3481 /// For a compression request or any other less common knob, build an
3482 /// [`ImageDesc`] and call [`Tensor::image_desc`].
3483 ///
3484 /// # Errors
3485 ///
3486 /// - [`Error::InvalidArgument`] if `width` × `height` is not a valid size
3487 /// for `format` (for example odd dimensions where the format's chroma
3488 /// subsampling forbids them), or, on macOS with an explicit
3489 /// [`TensorMemory::DmaBuf`], if the format cannot be expressed as an
3490 /// image-formatted IOSurface at this width. The message names the
3491 /// aligned width to use.
3492 /// - Whatever the chosen backend returns if the allocation itself fails —
3493 /// for an explicit `memory` request there is no fallback.
3494 ///
3495 /// # Example
3496 ///
3497 /// ```rust
3498 /// use edgefirst_tensor::{CpuAccess, PixelFormat, Tensor, TensorMemory};
3499 ///
3500 /// # fn main() -> Result<(), edgefirst_tensor::Error> {
3501 /// // A camera frame the GPU converts and the CPU never touches.
3502 /// let frame = Tensor::<u8>::image(1920, 1080, PixelFormat::Nv12,
3503 /// Some(TensorMemory::Mem), CpuAccess::None)?;
3504 /// assert_eq!(frame.format(), Some(PixelFormat::Nv12));
3505 /// # Ok(())
3506 /// # }
3507 /// ```
3508 pub fn image(
3509 width: usize,
3510 height: usize,
3511 format: PixelFormat,
3512 memory: Option<TensorMemory>,
3513 access: CpuAccess,
3514 ) -> Result<Self>
3515 where
3516 T: 'static,
3517 {
3518 // Shape comes from the shared `PixelFormat::allocation_shape` helper (packed /
3519 // planar / semi-planar NV12·NV16). NV12 supports odd dimensions via the
3520 // `H + ceil(H/2)` combined-plane height.
3521 // The `T: 'static` bound is required by the macOS IOSurface path below.
3522 let shape = format.allocation_shape(width, height).ok_or_else(|| {
3523 Error::InvalidArgument(format!(
3524 "invalid dimensions {width}x{height} for format {format:?}"
3525 ))
3526 })?;
3527
3528 // macOS Dma path: allocate a format-aware IOSurface (FourCC +
3529 // 2D dimensions) so the GL backend can bind it via
3530 // `EGL_ANGLE_iosurface_client_buffer`. Without this, the IOSurface
3531 // would default to a generic byte buffer (FourCC 'L008') and
3532 // ANGLE would reject the import with `EGL_BAD_ATTRIBUTE`.
3533 //
3534 // Guard: IOSurface rounds `bytes_per_row` up to 64-byte alignment.
3535 // If the natural row pitch (`width * channels * sizeof(T)`) is not
3536 // already 64-byte aligned, the padded allocation cannot be mapped
3537 // as a contiguous packed tensor — CPU reads/writes would use the
3538 // wrong stride.
3539 //
3540 // Explicit-Dma contract: when the caller passes
3541 // `Some(TensorMemory::DmaBuf)` they have asked for an
3542 // **image-formatted IOSurface**. Silently downgrading to the
3543 // generic 'L008' byte-bag when alignment fails buries the
3544 // mismatch — the caller only finds out hours later when ANGLE
3545 // (or any GL importer) rejects the bind with
3546 // `EGL_BAD_ATTRIBUTE`. Same anti-pattern bit us previously on
3547 // Mali GPUs with DMA-BUF padding. The right behaviour is to
3548 // fail loudly here with the alignment requirement spelled out
3549 // so the caller can either pick aligned dimensions, request
3550 // SHM/Mem explicitly, or pass `memory=None` for auto-select.
3551 #[cfg(any(target_os = "macos", target_os = "ios"))]
3552 if matches!(memory, Some(TensorMemory::DmaBuf)) {
3553 // For planar formats the IOSurface stacks channels
3554 // vertically (channels * height rows), so the row stride is
3555 // single-channel width * sizeof(T). Packed formats keep the
3556 // natural width * channels * sizeof(T) stride.
3557 let natural_row_bytes = match format.layout() {
3558 PixelLayout::Planar => width * std::mem::size_of::<T>(),
3559 _ => width * format.channels() * std::mem::size_of::<T>(),
3560 };
3561 // A format with a real IOSurface FourCC (RGBA/BGRA/YUYV packed,
3562 // GREY/NV12/NV16/NV24 as R8) tolerates a non-64-aligned natural
3563 // pitch: the surface is allocated with its own 64-aligned
3564 // `bytes_per_row`, the tensor records that stride below, and a CPU
3565 // map iterates rows correctly via the strided-map path while the GL
3566 // import uses the surface's pitch directly — fully zero-copy.
3567 // Planar (the F16 RGBA16F packing) is consumed flat as
3568 // `[1, C, H, W]` with no stride, so it still requires an aligned
3569 // pitch; and formats without a FourCC would fall through to a
3570 // generic byte-bag GL can't bind. Both fail loudly rather than
3571 // silently downgrade.
3572 let has_image_fourcc = dtype_of::<T>()
3573 .and_then(|dt| crate::iosurface_layout::image_iosurface_layout(format, dt))
3574 .is_some();
3575 let padded_ok = has_image_fourcc && format.layout() != PixelLayout::Planar;
3576 if !natural_row_bytes.is_multiple_of(64) && !padded_ok {
3577 let elem_size = std::mem::size_of::<T>();
3578 let per_pixel_bytes = match format.layout() {
3579 PixelLayout::Planar => elem_size.max(1),
3580 _ => format.channels().max(1) * elem_size.max(1),
3581 };
3582 // Compute the next 64-byte-aligned width by rounding the
3583 // natural row pitch up to the next multiple of 64 and
3584 // dividing back by per-pixel bytes. This handles every
3585 // `per_pixel_bytes` value correctly:
3586 //
3587 // * Divisors of 64 (1/2/4/8/16/32/64) → the suggestion
3588 // is always 64-byte aligned.
3589 // * Non-divisors of 64 (e.g. RGB u8 with 3 B/pixel) →
3590 // the next aligned row pitch may not be an integer
3591 // multiple of per_pixel_bytes (3 doesn't divide 64
3592 // in any way), so a "pad width to N" suggestion is
3593 // structurally impossible — omit the suggestion
3594 // instead of printing a wrong number.
3595 // * per_pixel_bytes > 64 → same situation, also
3596 // omitted; the previous formula divided by zero.
3597 //
3598 // The error always names the alignment requirement
3599 // verbatim and lists the two non-DMA alternatives so
3600 // the caller has at least one always-applicable fix.
3601 let aligned_row_bytes = natural_row_bytes.next_multiple_of(64);
3602 let pad_hint =
3603 if per_pixel_bytes > 0 && aligned_row_bytes.is_multiple_of(per_pixel_bytes) {
3604 let w = aligned_row_bytes / per_pixel_bytes;
3605 format!("Pad width to {w} (the next 64-byte-aligned stride), ")
3606 } else {
3607 String::new()
3608 };
3609 return Err(Error::InvalidArgument(format!(
3610 "Tensor::image: {format:?} {width}x{height} with element \
3611 size {elem_size} produces a {natural_row_bytes}-byte natural \
3612 row pitch, which is not 64-byte aligned. \
3613 IOSurface rounds bytes_per_row up to 64 bytes, so a \
3614 contiguous CPU map of this tensor would read garbage. \
3615 {pad_hint}pass memory=None to auto-fall-back to SHM, or \
3616 pass memory=Some(TensorMemory::Shm) or \
3617 Some(TensorMemory::Mem) explicitly."
3618 )));
3619 }
3620 // Alignment OK. Explicit-Dma contract: the caller asked for an
3621 // **image-formatted, GL-importable** IOSurface, so every failure
3622 // from here on is loud. The old behaviour fell through to the
3623 // generic 'L008' byte-bag `Some(other)` arm below, and the caller
3624 // only found out when the GL import rejected the bind with
3625 // `EGL_BAD_ATTRIBUTE` — the same silent-downgrade anti-pattern as
3626 // the alignment case above. (Semi-planar/Grey u8 map to 'L008'
3627 // *by design* in `image_iosurface_layout` — the R8-plane
3628 // representation the YUV shaders sample — so they return through
3629 // the mapped path and never reach these errors.)
3630 let dtype = dtype_of::<T>().ok_or_else(|| {
3631 Error::InvalidArgument(format!(
3632 "Tensor::image: element type {} has no DType, so no \
3633 image-formatted IOSurface exists. Pass memory=None or \
3634 Some(TensorMemory::Mem) for a CPU tensor.",
3635 std::any::type_name::<T>()
3636 ))
3637 })?;
3638 if crate::iosurface_layout::image_iosurface_layout(format, dtype).is_none() {
3639 return Err(Error::InvalidArgument(format!(
3640 "Tensor::image: no zero-copy IOSurface mapping exists for \
3641 {format:?}/{dtype:?} on macOS/iOS (supported: \
3642 Rgba/Rgb @ U8/I8, Bgra/Yuyv/Grey/Nv12/Nv16/Nv24 @ U8, \
3643 Rgba/PlanarRgb/PlanarRgba @ F16). Pass memory=None to \
3644 auto-select, or Some(TensorMemory::Mem) explicitly, for \
3645 a CPU tensor."
3646 )));
3647 }
3648 // Packed RGB u8/i8 rides an RGBA8888 surface at (W*3/4, H) —
3649 // reject a width the texel packing cannot express up front
3650 // (mirrors the Android pre-guard).
3651 if format == PixelFormat::Rgb
3652 && matches!(dtype, DType::U8 | DType::I8)
3653 && packed_rgb888_layout(width, height).is_none()
3654 {
3655 return Err(Error::InvalidArgument(format!(
3656 "Tensor::image: Rgb {dtype:?} requires width%4==0 for the RGBA8888 \
3657 IOSurface packing (got width={width}). Pad the width, or pass \
3658 memory=Some(TensorMemory::Mem) for a CPU tensor."
3659 )));
3660 }
3661 let storage = TensorStorage::<T>::new_image_iosurface(
3662 width, height, format, dtype, &shape, None,
3663 )?;
3664 let mut t = Self::wrap(storage);
3665 t.format = Some(format);
3666 // IOSurface rounds `bytes_per_row` up to 64 bytes. When that
3667 // pitch exceeds the natural packed/planar row stride, record
3668 // it so CPU consumers iterate rows correctly (the GL import
3669 // already uses the surface's own pitch). For 64-aligned rows
3670 // — the common model-input case — the two match and no stride
3671 // is stored, leaving the flat mapping unchanged.
3672 if let TensorStorage::Dma(ref io) = t.storage {
3673 let bpr = io.bytes_per_row();
3674 if let Some(natural) = t.effective_row_stride() {
3675 if bpr > natural {
3676 t.set_row_stride_unchecked(bpr);
3677 }
3678 }
3679 }
3680 t.cpu_access = access;
3681 return Ok(t);
3682 }
3683
3684 // Android Dma path: allocate a format-aware AHardwareBuffer so the
3685 // GL backend can import it as an EGLImage
3686 // (`eglGetNativeClientBufferANDROID` → `eglCreateImageKHR`).
3687 //
3688 // Geometry: gralloc chooses the row pitch (`desc.stride`) at
3689 // allocation, and pads freely (validated on the Galaxy S26 Ultra,
3690 // where SnapAlloc pads the planar-F16 RGBA16F surface). A padded
3691 // pitch is recorded on the tensor so CPU maps iterate rows via the
3692 // strided-map path; the GL render uses the buffer's own pitch
3693 // through the EGLImage either way, so the GPU path stays fully
3694 // zero-copy. Consumers needing a FLAT layout (the future NPU
3695 // handoff's `[1, C, H, W]` contract) must check `row_stride()` and
3696 // repack when set — flatness is a per-device property here, unlike
3697 // macOS where IOSurface's 64-BYTE alignment keeps model-sized F16
3698 // surfaces naturally flat.
3699 #[cfg(target_os = "android")]
3700 if matches!(memory, Some(TensorMemory::DmaBuf)) {
3701 // Explicit-Dma contract (mirrors the macOS block above): the
3702 // caller asked for an image-formatted, GL-importable
3703 // AHardwareBuffer, so an unmapped combination or a gralloc
3704 // refusal errors here instead of falling through to a BLOB
3705 // byte-bag the GL backend can never import — that downgrade
3706 // only surfaced as a silent per-frame CPU upload.
3707 let dtype = dtype_of::<T>().ok_or_else(|| {
3708 Error::InvalidArgument(format!(
3709 "Tensor::image: element type {} has no DType, so no \
3710 image-formatted AHardwareBuffer exists. Pass memory=None \
3711 or Some(TensorMemory::Mem) for a CPU tensor.",
3712 std::any::type_name::<T>()
3713 ))
3714 })?;
3715 // Planar F16 requires width % 4 == 0 for the RGBA16F
3716 // packing — reject up front with the requirement spelled
3717 // out (mirrors the macOS pre-allocation guard) instead of
3718 // falling through to a byte-bag GL cannot bind.
3719 if format.layout() == PixelLayout::Planar
3720 && dtype == DType::F16
3721 && packed_rgba16f_layout(format, dtype, width, height).is_none()
3722 {
3723 return Err(Error::InvalidArgument(format!(
3724 "Tensor::image: {format:?} F16 requires width%4==0 for the RGBA16F \
3725 AHardwareBuffer packing (got width={width}). Pad the width, or pass \
3726 memory=Some(TensorMemory::Mem) for a CPU tensor."
3727 )));
3728 }
3729 // Packed RGB u8/i8 rides an RGBA8888 surface at (W*3/4, H) —
3730 // the same whole-texel constraint as the F16 packing above.
3731 if format == PixelFormat::Rgb
3732 && matches!(dtype, DType::U8 | DType::I8)
3733 && packed_rgb888_layout(width, height).is_none()
3734 {
3735 return Err(Error::InvalidArgument(format!(
3736 "Tensor::image: Rgb {dtype:?} requires width%4==0 for the RGBA8888 \
3737 AHardwareBuffer packing (got width={width}). Pad the width, or pass \
3738 memory=Some(TensorMemory::Mem) for a CPU tensor."
3739 )));
3740 }
3741 if crate::ahardwarebuffer::image_ahardwarebuffer_layout(format, dtype).is_none() {
3742 return Err(Error::InvalidArgument(format!(
3743 "Tensor::image: no zero-copy AHardwareBuffer mapping exists \
3744 for {format:?}/{dtype:?} on Android (Grey/NV* need \
3745 R8-format buffers, API 29+; the HAL floor is API 26 — see \
3746 `image_ahardwarebuffer_layout`). Pass memory=None to \
3747 auto-select, or Some(TensorMemory::Mem) explicitly, for a \
3748 CPU tensor; camera NV12 stays zero-copy by wrapping the \
3749 camera's own AHardwareBuffer instead of allocating one."
3750 )));
3751 }
3752 let storage = TensorStorage::<T>::new_image_ahardwarebuffer(
3753 width, height, format, dtype, &shape, None, access,
3754 )?;
3755 let mut t = Self::wrap(storage);
3756 t.format = Some(format);
3757 if let TensorStorage::Dma(ref ahb) = t.storage {
3758 // gralloc chooses the row pitch (Qualcomm's
3759 // SnapAlloc pads e.g. the 160-px-wide RGBA16F
3760 // surface of a 640-wide planar F16 target).
3761 // When it exceeds the natural stride, record
3762 // it so CPU consumers iterate rows correctly
3763 // via `effective_row_stride()` — the GPU
3764 // renders through the EGLImage at the
3765 // buffer's real pitch regardless, so the
3766 // render stays fully zero-copy. Consumers
3767 // that need the FLAT `[1, C, H, W]` layout
3768 // (the future NPU handoff) must check
3769 // `row_stride()` is unset and repack — or
3770 // pick an aligned width — rather than assume
3771 // flatness (see the module docs).
3772 let bpr = ahb.bytes_per_row();
3773 if let Some(natural) = t.effective_row_stride() {
3774 if bpr > natural {
3775 log::debug!(
3776 "Tensor::image: gralloc padded the {format:?} \
3777 AHardwareBuffer pitch to {bpr} bytes (natural \
3778 {natural}); recording row stride"
3779 );
3780 t.set_row_stride_unchecked(bpr);
3781 }
3782 }
3783 }
3784 t.cpu_access = access;
3785 return Ok(t);
3786 }
3787
3788 // Windows Dma path: allocate an `ID3D11Texture2D` on the process
3789 // device so the GL backend can bind it through ANGLE's D3D11 texture
3790 // import and a D3D12/CUDA consumer can open its shared NT handle.
3791 //
3792 // Explicit-Dma contract, as on macOS and Android: the caller asked for
3793 // a texture, so a format with no texture layout is an error here
3794 // rather than a silent fall-through to the host-memory arms below,
3795 // which would only surface much later as a failed GPU import.
3796 #[cfg(target_os = "windows")]
3797 if matches!(memory, Some(TensorMemory::DmaBuf)) {
3798 let dtype = dtype_of::<T>().ok_or_else(|| {
3799 Error::InvalidArgument(format!(
3800 "Tensor::image: element type {} has no DType, so no D3D11 texture \
3801 layout exists. Pass memory=None or Some(TensorMemory::Mem) for a \
3802 CPU tensor.",
3803 std::any::type_name::<T>()
3804 ))
3805 })?;
3806 let storage = TensorStorage::<T>::new_image_d3d11(
3807 width, height, format, dtype, &shape, None, access,
3808 )?;
3809 let mut t = Self::wrap(storage);
3810 t.format = Some(format);
3811 record_d3d11_row_stride(&mut t, format, "Tensor::image");
3812 t.cpu_access = access;
3813 // Best-effort CUDA import, as the Linux DMA arms do: a texture a
3814 // CUDA consumer can map is worth the one import at allocation.
3815 t.try_init_d3d11_cuda();
3816 return Ok(t);
3817 }
3818
3819 // Compute the **64-byte-aligned** row stride for every image layout.
3820 //
3821 // Embedded GPUs reject `eglCreateImage` DMA-BUF imports whose row pitch
3822 // is not 64-byte aligned: Mali returns `EGL_BAD_ALLOC`, Vivante
3823 // `EGL_BAD_ACCESS`. This bit packed RGBA/RGB destinations at odd widths
3824 // AND at even non-multiple-of-16 widths (e.g. 321→1284, 322→1288 bytes —
3825 // neither divisible by 64), so an odd-source → RGBA convert failed on
3826 // imx95/imx8mp while succeeding on V3D/Tegra. Semi-planar already aligned
3827 // here; we now align packed and planar identically so every image()
3828 // allocation is GPU-importable regardless of width.
3829 //
3830 // The per-layout natural pitch and total row count:
3831 // * SemiPlanar `[total_h, width]` — pitch = even(width)·elem, rows = total_h
3832 // * Packed `[height, width, ch]` — pitch = width·ch·elem, rows = height
3833 // * Planar `[ch, height, width]` — pitch = width·elem, rows = ch·height
3834 // Allocation byte size = `aligned_stride · total_rows` (NOT the shape
3835 // product, which reflects only the logical width and under-allocates the
3836 // padding on odd / unaligned widths).
3837 let elem = std::mem::size_of::<T>();
3838 let channels = format.channels();
3839 let (natural_stride, total_rows) = match format.layout() {
3840 PixelLayout::SemiPlanar => (width.next_multiple_of(2) * elem, shape[0]),
3841 PixelLayout::Packed => (width * channels * elem, height),
3842 PixelLayout::Planar => (width * elem, channels * height),
3843 };
3844 let aligned_stride = natural_stride.next_multiple_of(64);
3845 let semi = format.layout() == PixelLayout::SemiPlanar;
3846
3847 // DMA buffers MUST carry a 64-aligned row pitch — Mali/Vivante reject a
3848 // DMA-BUF EGLImage whose pitch is not 64-aligned. Semi-planar also needs
3849 // the aligned pitch on every backend (its chroma-plane offset math
3850 // assumes it). Packed/planar on host-only memory (Mem/Shm) keep the
3851 // natural tight pitch so the many flat CPU consumers are unaffected.
3852 let host_stride = if semi { aligned_stride } else { natural_stride };
3853 let host_byte_size = host_stride * total_rows;
3854 #[cfg(target_os = "linux")]
3855 let dma_byte_size = aligned_stride * total_rows;
3856
3857 // `used_stride` is the actual row pitch of the storage created below.
3858 let (storage, used_stride) = match memory {
3859 #[cfg(target_os = "linux")]
3860 Some(TensorMemory::DmaBuf) => (
3861 TensorStorage::<T>::new_dma_with_byte_size(&shape, dma_byte_size, None)?,
3862 aligned_stride,
3863 ),
3864 #[cfg(unix)]
3865 Some(TensorMemory::Shm) => (
3866 TensorStorage::<T>::new_shm_with_byte_size(&shape, host_byte_size, None)?,
3867 host_stride,
3868 ),
3869 Some(TensorMemory::Mem) => (
3870 TensorStorage::<T>::new_mem_with_byte_size(&shape, host_byte_size, None)?,
3871 host_stride,
3872 ),
3873 #[allow(unused_variables)]
3874 Some(other) => {
3875 // PBO and any future variants: fall through to standard new().
3876 return {
3877 let mut t = Self::new(&shape, Some(other), None)?;
3878 t.format = Some(format);
3879 t.cpu_access = access;
3880 Ok(t)
3881 };
3882 }
3883 None => {
3884 // Auto-select priority: DMA → Mem (DMA gets the 64-aligned
3885 // pitch; the Mem fallback keeps the tight host pitch, so the
3886 // recorded stride matches the storage used). Shm is NOT
3887 // auto-selected — it offers no advantage over Mem for an
3888 // in-process image and Mem always succeeds, so it sits below Mem
3889 // and is reached only via an explicit `TensorMemory::Shm`.
3890 #[cfg(target_os = "linux")]
3891 {
3892 match TensorStorage::<T>::new_dma_with_byte_size(&shape, dma_byte_size, None) {
3893 Ok(s) => (s, aligned_stride),
3894 Err(_) => (
3895 TensorStorage::<T>::new_mem_with_byte_size(
3896 &shape,
3897 host_byte_size,
3898 None,
3899 )?,
3900 host_stride,
3901 ),
3902 }
3903 }
3904 #[cfg(not(target_os = "linux"))]
3905 {
3906 (
3907 TensorStorage::<T>::new_mem_with_byte_size(&shape, host_byte_size, None)?,
3908 host_stride,
3909 )
3910 }
3911 }
3912 };
3913
3914 let mut t = Self::wrap(storage);
3915 t.format = Some(format);
3916 // Record the row stride when it exceeds the natural tight pitch (padding
3917 // is present — DMA packed/planar at an unaligned width, or always for
3918 // semi-planar), mirroring the IOSurface path above. Aligned-width and
3919 // host-only packed/planar images keep their flat layout with no explicit
3920 // stride; `effective_row_stride()` then falls back to the identical
3921 // computed pitch. When padding IS present, consumers must iterate rows by
3922 // `effective_row_stride()` to skip it.
3923 if semi || used_stride > natural_stride {
3924 t.set_row_stride_unchecked(used_stride);
3925 }
3926 debug_assert!(
3927 t.row_stride.is_some() || !semi,
3928 "image() must always set row_stride for semi-planar tensors"
3929 );
3930 t.cpu_access = access;
3931 #[cfg(target_os = "linux")]
3932 t.try_init_dma_cuda();
3933 Ok(t)
3934 }
3935
3936 /// Create a DMA-backed image tensor with an explicit row stride that
3937 /// may exceed the natural `width * channels * sizeof(T)` pitch.
3938 ///
3939 /// Used for image tensors that need GPU pitch alignment padding: the
3940 /// underlying DMA-BUF is sized to `row_stride * height` bytes, but
3941 /// the tensor's logical shape stays at `[height, width, channels]`.
3942 /// `width()` / `height()` / `shape()` continue to report the
3943 /// user-requested values; the padding is visible only via
3944 /// `row_stride()` / `effective_row_stride()` and is automatically
3945 /// propagated to the GL backend's EGLImage import so Mali Valhall
3946 /// accepts the buffer.
3947 ///
3948 /// # Supported formats
3949 ///
3950 /// Currently only **packed** pixel layouts (RGBA8, BGRA8, RGB888,
3951 /// Grey, etc.) are supported — the formats the GL backend uses as
3952 /// render destinations. Semi-planar formats (NV12, NV16) come from
3953 /// external allocators (camera capture, video decoders) and are
3954 /// imported via `TensorDyn::from_fd` + `set_row_stride`, which
3955 /// already supports padded strides.
3956 ///
3957 /// # Supported memory
3958 ///
3959 /// Currently only `TensorMemory::DmaBuf` is supported. PBO and Mem
3960 /// storage don't go through EGLImage import so they don't need
3961 /// pitch alignment; if you pass any other memory type this returns
3962 /// `NotImplemented`. `None` (auto-select) is treated as `Dma`.
3963 ///
3964 /// # Errors
3965 ///
3966 /// - `InvalidArgument` if `row_stride_bytes < width * channels * sizeof(T)`
3967 /// (the requested stride would not fit a single row)
3968 /// - `NotImplemented` for non-packed formats or non-DMA memory
3969 /// - `IoError` if the DMA-heap allocation fails (propagated from
3970 /// `DmaTensor::new_with_byte_size`)
3971 pub fn image_with_stride(
3972 width: usize,
3973 height: usize,
3974 format: PixelFormat,
3975 row_stride_bytes: usize,
3976 memory: Option<TensorMemory>,
3977 access: CpuAccess,
3978 ) -> Result<Self> {
3979 #[cfg(not(target_os = "linux"))]
3980 let _ = access;
3981 // DMA backing (the only thing this constructor produces) is
3982 // Linux-only. On macOS/BSD/Windows the non-Linux block below is
3983 // the only compiled body and returns `NotImplemented` directly;
3984 // on Linux the non-Linux block is cfg-removed and the function
3985 // falls through to the real validation + allocation path. Each
3986 // target compiles exactly one of the two blocks, and the block
3987 // serves as the function's tail expression in both cases — so
3988 // neither needs an explicit `return` (avoids
3989 // `clippy::needless_return` on the macOS CI gate).
3990 #[cfg(not(target_os = "linux"))]
3991 {
3992 let _ = (width, height, format, row_stride_bytes, memory);
3993 Err(Error::NotImplemented(
3994 "image_with_stride requires DMA support (Linux only)".to_owned(),
3995 ))
3996 }
3997
3998 #[cfg(target_os = "linux")]
3999 {
4000 if format.layout() != PixelLayout::Packed {
4001 return Err(Error::NotImplemented(format!(
4002 "Tensor::image_with_stride only supports packed pixel layouts, got {format:?}"
4003 )));
4004 }
4005 let elem = std::mem::size_of::<T>();
4006 let min_stride = width
4007 .checked_mul(format.channels())
4008 .and_then(|p| p.checked_mul(elem))
4009 .ok_or_else(|| {
4010 Error::InvalidArgument(format!(
4011 "image_with_stride: width {width} × channels {} × sizeof::<T>={elem} \
4012 overflows usize",
4013 format.channels()
4014 ))
4015 })?;
4016 if row_stride_bytes < min_stride {
4017 return Err(Error::InvalidArgument(format!(
4018 "image_with_stride: row_stride {row_stride_bytes} < minimum {min_stride} \
4019 ({width} px × {} ch × {elem} B)",
4020 format.channels()
4021 )));
4022 }
4023 let total_byte_size = row_stride_bytes.checked_mul(height).ok_or_else(|| {
4024 Error::InvalidArgument(format!(
4025 "image_with_stride: row_stride {row_stride_bytes} × height {height} overflows usize"
4026 ))
4027 })?;
4028
4029 let shape = vec![height, width, format.channels()];
4030
4031 let storage = match memory {
4032 Some(TensorMemory::DmaBuf) | None => {
4033 TensorStorage::<T>::new_dma_with_byte_size(&shape, total_byte_size, None)?
4034 }
4035 Some(other) => {
4036 return Err(Error::NotImplemented(format!(
4037 "image_with_stride: only TensorMemory::DmaBuf is supported, got {other:?}"
4038 )));
4039 }
4040 };
4041
4042 let mut t = Self::wrap(storage);
4043 t.format = Some(format);
4044 t.row_stride = Some(row_stride_bytes);
4045 t.cpu_access = access;
4046 // Match new()/from_fd(): a DMA tensor must attempt CUDA external-
4047 // memory import so a strided DMA buffer is also zero-copy
4048 // CUDA-mappable (no-op when libcudart is absent).
4049 t.try_init_dma_cuda();
4050 Ok(t)
4051 }
4052 }
4053
4054 /// Attach format metadata to an existing tensor.
4055 ///
4056 /// # Arguments
4057 ///
4058 /// * `format` - The pixel format to attach
4059 ///
4060 /// # Returns
4061 ///
4062 /// `Ok(())` on success, with the format stored as metadata on the tensor.
4063 ///
4064 /// # Errors
4065 ///
4066 /// Returns `Error::InvalidShape` if the tensor shape is incompatible with
4067 /// the format's layout (packed expects `[H, W, C]`, planar expects
4068 /// `[C, H, W]`, semi-planar expects `[H*k, W]` with format-specific
4069 /// height constraints).
4070 pub fn set_format(&mut self, format: PixelFormat) -> Result<()> {
4071 let shape = self.shape();
4072 match format.layout() {
4073 PixelLayout::Packed => {
4074 if shape.len() != 3 || shape[2] != format.channels() {
4075 return Err(Error::InvalidShape(format!(
4076 "packed format {format:?} expects [H, W, {}], got {shape:?}",
4077 format.channels()
4078 )));
4079 }
4080 }
4081 PixelLayout::Planar => {
4082 if shape.len() != 3 || shape[0] != format.channels() {
4083 return Err(Error::InvalidShape(format!(
4084 "planar format {format:?} expects [{}, H, W], got {shape:?}",
4085 format.channels()
4086 )));
4087 }
4088 }
4089 PixelLayout::SemiPlanar => {
4090 if shape.len() != 2 {
4091 return Err(Error::InvalidShape(format!(
4092 "semi-planar format {format:?} expects [H*k, W], got {shape:?}"
4093 )));
4094 }
4095 match format {
4096 // Combined-plane height is `H + ceil(H/2)` (luma + chroma
4097 // rows). For even H that is `3H/2` (≡ 0 mod 3); for odd H it
4098 // is `(3H+1)/2` (≡ 2 mod 3). Only totals ≡ 1 mod 3 are
4099 // unreachable, so reject just those — odd-height NV12 is
4100 // valid (e.g. 725 rows for a 483-tall image).
4101 PixelFormat::Nv12 if shape[0] % 3 == 1 => {
4102 return Err(Error::InvalidShape(format!(
4103 "NV12 contiguous shape[0] must be H + ceil(H/2) for some height; \
4104 {} is unreachable (≡ 1 mod 3)",
4105 shape[0]
4106 )));
4107 }
4108 PixelFormat::Nv16 if !shape[0].is_multiple_of(2) => {
4109 return Err(Error::InvalidShape(format!(
4110 "NV16 contiguous shape[0] must be even, got {}",
4111 shape[0]
4112 )));
4113 }
4114 // NV24 (4:4:4): combined-plane height is 3H (Y + 2H chroma).
4115 PixelFormat::Nv24 if !shape[0].is_multiple_of(3) => {
4116 return Err(Error::InvalidShape(format!(
4117 "NV24 contiguous shape[0] must be a multiple of 3 (= 3H), got {}",
4118 shape[0]
4119 )));
4120 }
4121 _ => {}
4122 }
4123 }
4124 }
4125 // Clear stored stride/offset when format changes — they may be invalid
4126 // for the new format. Caller must re-set after changing format.
4127 if self.format != Some(format) {
4128 self.row_stride = None;
4129 self.plane_offset = None;
4130 match self.storage {
4131 TensorStorage::Mem(ref mut m) => m.set_offset(0),
4132 #[cfg(target_os = "linux")]
4133 TensorStorage::Dma(ref mut dma) => dma.mmap_offset = 0,
4134 // Paired with the arm in `set_plane_offset`: a setter that
4135 // takes effect on this backing needs a clear that does too,
4136 // or the wrapper reports `None` while `map()` still starts at
4137 // the old offset -- a stale window rather than a lost one.
4138 // `reshape`'s clear site needs no such arm:
4139 // `IoSurfaceTensor::reshape` zeroes `view_offset` itself
4140 // before the match there runs.
4141 #[cfg(any(target_os = "macos", target_os = "ios"))]
4142 TensorStorage::Dma(ref mut io) => io.view_offset = 0,
4143 // The same pair for D3D11. Unlike `IoSurfaceTensor`,
4144 // `D3d11TextureTensor::reshape` leaves `view_offset` alone, so
4145 // `reshape`'s clear site carries this arm as well.
4146 #[cfg(target_os = "windows")]
4147 TensorStorage::Dma(ref mut tex) => tex.set_view_offset(0),
4148 // Paired with the arm in `set_plane_offset`: a setter that
4149 // takes effect on this backing needs a clear that does too.
4150 // `reshape`'s clear site needs no `Pbo` arm --
4151 // `PboTensor::reshape` zeroes `view_offset` itself before the
4152 // match there runs, exactly as `IoSurfaceTensor::reshape`
4153 // does. `set_logical_shape` deliberately does not, matching
4154 // `Mem`.
4155 TensorStorage::Pbo(ref mut p) => p.view_offset = 0,
4156 // Same pair, `Shm` side. `cfg(unix)` mirrors the `Shm`
4157 // variant's own gate on `TensorStorage`, not a per-OS
4158 // multiplex -- see the comment on the `set_plane_offset`
4159 // arm.
4160 #[cfg(unix)]
4161 TensorStorage::Shm(ref mut s) => s.set_offset(0),
4162 // No-op, pairing with `set_plane_offset`'s Android arm:
4163 // nothing writes `AHardwareBufferTensor::view_offset` there
4164 // either yet, so there is nothing here for a format change
4165 // to clear. See that arm's comment for the full accounting
4166 // of the still-incomplete gap.
4167 #[cfg(target_os = "android")]
4168 TensorStorage::Dma(_) => {}
4169 }
4170 }
4171 self.format = Some(format);
4172 Ok(())
4173 }
4174
4175 /// Set this tensor's logical dimensions and pixel format to a decoded
4176 /// image, reusing the existing allocation. The shape is derived from the
4177 /// format layout; fails with `Error::InsufficientCapacity` if the
4178 /// allocation cannot hold `width`×`height` in `format`, or
4179 /// `Error::InvalidArgument` if the dimensions are invalid for the format.
4180 ///
4181 /// For NV12/NV16/NV24 the buffer width is rounded up to even (a chroma-plane
4182 /// interleaving requirement); the true odd width is reported by the decoder
4183 /// in `ImageInfo` and trimmed by a `convert()` crop. See
4184 /// [`PixelFormat::allocation_shape`].
4185 ///
4186 /// When the backing has a fixed physical row pitch (an IOSurface's
4187 /// 64-aligned `bytesPerRow`) that exceeds the new format's natural row
4188 /// stride — i.e. a reused max-sized pool tensor reconfigured to a smaller
4189 /// image — the physical pitch is preserved as the tensor's `row_stride`.
4190 /// This keeps the **physical grid** (allocation stride/surface) fixed while
4191 /// the **logical ROI** (this image's W×H) changes, so the decode writes rows
4192 /// at the surface's real stride and the GPU samples them at the same stride.
4193 /// Exact-sized buffers (pitch == natural) stay tightly packed unchanged.
4194 pub fn configure_image(
4195 &mut self,
4196 width: usize,
4197 height: usize,
4198 format: PixelFormat,
4199 ) -> Result<()> {
4200 let shape = format.allocation_shape(width, height).ok_or_else(|| {
4201 Error::InvalidArgument(format!(
4202 "invalid dimensions {width}x{height} for format {format:?}"
4203 ))
4204 })?;
4205 // Capture the pre-existing row stride before `set_format` clears it.
4206 // For pool tensors that were allocated at a larger width (e.g. 1920-wide
4207 // pool decoding a 789-wide image), this preserves the backing pitch so
4208 // rows are still written at the correct physical stride.
4209 let prior_stride = self.row_stride;
4210
4211 self.storage.set_logical_shape(&shape)?;
4212 self.set_format(format)?; // clears any stale row_stride
4213
4214 // Restore the correct row pitch for the new geometry. A DMA buffer of
4215 // ANY layout, and a semi-planar buffer on any backing, MUST carry a
4216 // 64-byte-aligned pitch: Mali/Vivante reject a DMA-BUF EGLImage whose
4217 // row pitch is not 64-aligned (`EGL_BAD_ALLOC` / `EGL_BAD_ACCESS`), and
4218 // semi-planar chroma-offset math assumes it. This mirrors the pitch
4219 // `Tensor::image()` allocates, so a recycled pool buffer imports
4220 // identically to a fresh one — without it a `configure_image()`'d packed
4221 // buffer (e.g. Y800 96-wide → tight pitch 96) failed convert() on
4222 // imx95/imx8mp via the texture-upload fallback while the fresh oracle
4223 // (aligned 128) succeeded. Packed/planar on host-only memory (Mem/Shm)
4224 // keep the tight pitch so flat CPU consumers stay unaffected — matching
4225 // `image()`'s `host_stride` rule.
4226 let elem = std::mem::size_of::<T>();
4227 let channels = format.channels();
4228 let (min_stride, total_rows) = match format.layout() {
4229 PixelLayout::SemiPlanar => (width.next_multiple_of(2) * elem, shape[0]),
4230 PixelLayout::Packed => (width * channels * elem, height),
4231 PixelLayout::Planar => (width * elem, channels * height),
4232 };
4233 let needs_align = self.storage.memory() == TensorMemory::DmaBuf
4234 || format.layout() == PixelLayout::SemiPlanar;
4235
4236 let active_stride = if let Some(pitch) = self.storage.backing_row_stride() {
4237 // macOS IOSurface: use the surface's native pitch.
4238 let natural = self.effective_row_stride().unwrap_or(0);
4239 if pitch > natural {
4240 self.set_row_stride_unchecked(pitch);
4241 pitch
4242 } else {
4243 natural
4244 }
4245 } else if needs_align {
4246 // Priority:
4247 // 1. Prior stride (pool reuse): if the pre-existing stride is
4248 // 64-aligned, >= this layout's minimum row, and fits the
4249 // allocation, keep it. This is the hot-loop reuse case (large
4250 // pool, small image).
4251 // 2. Compute a fresh 64-aligned stride for the current width.
4252 let aligned = min_stride.next_multiple_of(64);
4253 let capacity = self.storage.capacity_bytes();
4254
4255 let candidate = if let Some(ps) = prior_stride {
4256 if ps >= min_stride && ps % 64 == 0 && ps * total_rows <= capacity {
4257 ps
4258 } else {
4259 aligned
4260 }
4261 } else {
4262 aligned
4263 };
4264
4265 if candidate * total_rows <= capacity {
4266 self.set_row_stride_unchecked(candidate);
4267 candidate
4268 } else {
4269 // Shouldn't happen for legitimate pools, but don't crash.
4270 self.effective_row_stride().unwrap_or(0)
4271 }
4272 } else {
4273 self.effective_row_stride().unwrap_or(0)
4274 };
4275
4276 // Ensure the active stride fits the allocation. A pool reconfigured to a
4277 // wider image than its backing would silently SIGBUS on any subsequent
4278 // map/write — catch it here instead.
4279 if needs_align && active_stride > 0 {
4280 let needed = active_stride * total_rows;
4281 let capacity = self.storage.capacity_bytes();
4282 if needed > capacity {
4283 return Err(Error::InsufficientCapacity { needed, capacity });
4284 }
4285 }
4286 Ok(())
4287 }
4288
4289 /// Allocate an image tensor sized to hold up to `width`×`height` in
4290 /// `format`, reusable for any smaller image via `configure_image`.
4291 pub fn image_with_capacity(
4292 width: usize,
4293 height: usize,
4294 format: PixelFormat,
4295 memory: Option<TensorMemory>,
4296 access: CpuAccess,
4297 ) -> Result<Self>
4298 where
4299 T: 'static,
4300 {
4301 Self::image(width, height, format, memory, access)
4302 }
4303
4304 /// Pixel format (None if not an image).
4305 pub fn format(&self) -> Option<PixelFormat> {
4306 self.format
4307 }
4308
4309 /// Image width (None if not an image).
4310 pub fn width(&self) -> Option<usize> {
4311 let fmt = self.format?;
4312 let shape = self.shape();
4313 match fmt.layout() {
4314 PixelLayout::Packed => Some(shape[1]),
4315 PixelLayout::Planar => Some(shape[2]),
4316 PixelLayout::SemiPlanar => Some(shape[1]),
4317 }
4318 }
4319
4320 /// Image height (None if not an image).
4321 ///
4322 /// For semi-planar formats the combined-plane shape row count is divided
4323 /// by the format's luma-to-total ratio to recover logical height. This
4324 /// returns the exact logical height (including odd heights) only because
4325 /// the logical dimensions are tracked separately from the physical shape —
4326 /// `configure_image` stores the actual `(width, height)` in the format's
4327 /// `allocation_shape`, which round-trips losslessly via these accessors.
4328 pub fn height(&self) -> Option<usize> {
4329 let fmt = self.format?;
4330 let shape = self.shape();
4331 match fmt.layout() {
4332 PixelLayout::Packed => Some(shape[0]),
4333 PixelLayout::Planar => Some(shape[1]),
4334 PixelLayout::SemiPlanar => {
4335 if self.is_multiplane() {
4336 Some(shape[0])
4337 } else {
4338 match fmt {
4339 PixelFormat::Nv12 => Some(shape[0] * 2 / 3),
4340 PixelFormat::Nv16 => Some(shape[0] / 2),
4341 PixelFormat::Nv24 => Some(shape[0] / 3),
4342 _ => None,
4343 }
4344 }
4345 }
4346 }
4347 }
4348
4349 /// Create from separate Y and UV planes (multiplane NV12/NV16).
4350 pub fn from_planes(luma: Tensor<T>, chroma: Tensor<T>, format: PixelFormat) -> Result<Self> {
4351 if format.layout() != PixelLayout::SemiPlanar {
4352 return Err(Error::InvalidArgument(format!(
4353 "from_planes requires a semi-planar format, got {format:?}"
4354 )));
4355 }
4356 if chroma.format.is_some() || chroma.chroma.is_some() {
4357 return Err(Error::InvalidArgument(
4358 "chroma tensor must be a raw tensor (no format or chroma metadata)".into(),
4359 ));
4360 }
4361 let luma_shape = luma.shape();
4362 let chroma_shape = chroma.shape();
4363 if luma_shape.len() != 2 || chroma_shape.len() != 2 {
4364 return Err(Error::InvalidArgument(format!(
4365 "from_planes expects 2D shapes, got luma={luma_shape:?} chroma={chroma_shape:?}"
4366 )));
4367 }
4368 if luma_shape[1] != chroma_shape[1] {
4369 return Err(Error::InvalidArgument(format!(
4370 "luma width {} != chroma width {}",
4371 luma_shape[1], chroma_shape[1]
4372 )));
4373 }
4374 match format {
4375 PixelFormat::Nv12 => {
4376 if luma_shape[0] % 2 != 0 {
4377 return Err(Error::InvalidArgument(format!(
4378 "NV12 requires even luma height, got {}",
4379 luma_shape[0]
4380 )));
4381 }
4382 if chroma_shape[0] != luma_shape[0] / 2 {
4383 return Err(Error::InvalidArgument(format!(
4384 "NV12 chroma height {} != luma height / 2 ({})",
4385 chroma_shape[0],
4386 luma_shape[0] / 2
4387 )));
4388 }
4389 }
4390 PixelFormat::Nv16 => {
4391 if chroma_shape[0] != luma_shape[0] {
4392 return Err(Error::InvalidArgument(format!(
4393 "NV16 chroma height {} != luma height {}",
4394 chroma_shape[0], luma_shape[0]
4395 )));
4396 }
4397 }
4398 // NV24's chroma plane is full-resolution (2×-wide interleaved UV),
4399 // which the equal-width plane check above doesn't model. Multiplane
4400 // NV24 is unused (the JPEG decoder emits a contiguous NV24 buffer),
4401 // so it's not supported here yet.
4402 _ => {
4403 return Err(Error::InvalidArgument(format!(
4404 "from_planes only supports NV12 and NV16 (NV24 multiplane not yet \
4405 supported — use a contiguous NV24 tensor), got {format:?}"
4406 )));
4407 }
4408 }
4409
4410 Ok(Tensor {
4411 storage: luma.storage,
4412 format: Some(format),
4413 chroma: Some(Box::new(chroma)),
4414 row_stride: luma.row_stride,
4415 plane_offset: luma.plane_offset,
4416 quantization: luma.quantization,
4417 // A multiplane tensor spans two DMA-BUFs (luma + chroma); CUDA
4418 // external-memory import is per-fd, so there is no single device
4419 // pointer for the composite. Any CUDA handle the luma plane carried
4420 // is intentionally dropped — consumers needing CUDA access to
4421 // multiplane data must import each plane independently.
4422 cuda: None,
4423 colorimetry: luma.colorimetry,
4424 cpu_access: luma.cpu_access,
4425 compression: luma.compression,
4426 // A composed multiplane tensor is a whole image, not a sub-view.
4427 view_origin: None,
4428 })
4429 }
4430
4431 /// Whether this tensor uses separate plane allocations.
4432 pub fn is_multiplane(&self) -> bool {
4433 self.chroma.is_some()
4434 }
4435
4436 /// Access the chroma plane for multiplane semi-planar images.
4437 pub fn chroma(&self) -> Option<&Tensor<T>> {
4438 self.chroma.as_deref()
4439 }
4440
4441 /// Mutable access to the chroma plane for multiplane semi-planar images.
4442 pub fn chroma_mut(&mut self) -> Option<&mut Tensor<T>> {
4443 self.chroma.as_deref_mut()
4444 }
4445
4446 /// Row stride in bytes (`None` = tightly packed).
4447 pub fn row_stride(&self) -> Option<usize> {
4448 self.row_stride
4449 }
4450
4451 /// Effective row stride in bytes: the stored stride if set, otherwise the
4452 /// minimum stride computed from the format, width, and element size.
4453 /// Returns `None` only when no format is set and no explicit stride was
4454 /// stored via [`set_row_stride`](Self::set_row_stride).
4455 ///
4456 /// **GREY note:** `effective_row_stride()` for a GREY tensor returns the
4457 /// tight `width` bytes (no padding), which is what `normalize_to_numpy` and
4458 /// the CPU convert path expect. The codec's internal `native_row_stride`
4459 /// (64-byte-aligned) is used only during decoding and is not propagated to
4460 /// the tensor's stored stride, so callers reading via
4461 /// `effective_row_stride()` always see the tight value for GREY.
4462 pub fn effective_row_stride(&self) -> Option<usize> {
4463 if let Some(s) = self.row_stride {
4464 return Some(s);
4465 }
4466 let fmt = self.format?;
4467 let w = self.width()?;
4468 let elem = std::mem::size_of::<T>();
4469 Some(match fmt.layout() {
4470 PixelLayout::Packed => w * fmt.channels() * elem,
4471 PixelLayout::Planar => w * elem,
4472 // Semi-planar: minimum stride must cover the even width so the
4473 // interleaved chroma columns are byte-aligned on odd-width images.
4474 PixelLayout::SemiPlanar => w.next_multiple_of(2) * elem,
4475 })
4476 }
4477
4478 /// Copy the tensor's logical bytes into `dst`, compacting away any
4479 /// recorded row-stride padding.
4480 ///
4481 /// The flatness helper for NPU handoff: consumers that need a FLAT
4482 /// layout (e.g. `[1, C, H, W]` for NNAPI/LiteRT when the runtime cannot
4483 /// take a padded pitch) call this when [`row_stride`](Self::row_stride)
4484 /// is `Some` — on a tight tensor it degenerates to one memcpy, so it is
4485 /// safe to call unconditionally. `dst.len()` must equal the tight byte
4486 /// footprint (`shape` product × element size). Zero-copy consumers
4487 /// should prefer the buffer handle + [`Tensor::effective_row_stride`]
4488 /// (Self::effective_row_stride) and skip this copy entirely.
4489 pub fn copy_to_flat(&self, dst: &mut [u8]) -> Result<()> {
4490 let tight_bytes = crate::ahardwarebuffer_layout::checked_shape_bytes::<T>(self.shape())?;
4491 if dst.len() != tight_bytes {
4492 return Err(Error::InvalidArgument(format!(
4493 "copy_to_flat: dst is {} bytes but the tensor's tight \
4494 footprint is {tight_bytes} bytes (shape {:?})",
4495 dst.len(),
4496 self.shape()
4497 )));
4498 }
4499 // Read-only: this helper never writes back into the tensor, so a
4500 // read map suffices — and is required on backends (D3D11) that
4501 // reject a writable map on a tensor allocated CpuAccess::Read.
4502 let map = self.map_read()?;
4503 // SAFETY: T is a plain numeric type (crate-wide bound); viewing the
4504 // mapped elements as bytes is sound.
4505 let src: &[u8] = unsafe {
4506 std::slice::from_raw_parts(
4507 map.as_slice().as_ptr() as *const u8,
4508 std::mem::size_of_val(map.as_slice()),
4509 )
4510 };
4511 let Some(stride) = self.row_stride else {
4512 // Tight layout: the mapped window is exactly the logical bytes.
4513 let got = src.len().min(tight_bytes);
4514 if got < tight_bytes {
4515 return Err(Error::InvalidOperation(format!(
4516 "copy_to_flat: mapped {got} bytes < tight footprint {tight_bytes}"
4517 )));
4518 }
4519 dst.copy_from_slice(&src[..tight_bytes]);
4520 return Ok(());
4521 };
4522 // Strided: row count follows the strided-map convention (planar
4523 // stacks C planes of H rows; packed/semi-planar use shape[0]), and
4524 // the logical row is tight_bytes / rows for every layout.
4525 let rows = match self.format.map(|f| f.layout()) {
4526 Some(PixelLayout::Planar) => {
4527 let s = self.shape();
4528 if s.len() < 2 {
4529 return Err(Error::InvalidOperation(
4530 "copy_to_flat: strided planar tensor requires [C, H, W] shape".into(),
4531 ));
4532 }
4533 s[0].checked_mul(s[1]).ok_or_else(|| {
4534 Error::InvalidOperation(format!(
4535 "copy_to_flat: planar rows {} × {} overflows usize",
4536 s[0], s[1]
4537 ))
4538 })?
4539 }
4540 _ => *self.shape().first().ok_or_else(|| {
4541 Error::InvalidOperation("copy_to_flat: tensor has an empty shape".into())
4542 })?,
4543 };
4544 if rows == 0 || !tight_bytes.is_multiple_of(rows) {
4545 return Err(Error::InvalidOperation(format!(
4546 "copy_to_flat: tight footprint {tight_bytes} does not divide \
4547 into {rows} rows"
4548 )));
4549 }
4550 let row_bytes = tight_bytes / rows;
4551 let need = (rows - 1)
4552 .checked_mul(stride)
4553 .and_then(|b| b.checked_add(row_bytes))
4554 .ok_or_else(|| {
4555 Error::InvalidOperation(format!(
4556 "copy_to_flat: stride {stride} × rows {rows} overflows usize"
4557 ))
4558 })?;
4559 if src.len() < need {
4560 return Err(Error::InvalidOperation(format!(
4561 "copy_to_flat: mapped {} bytes but strided rows need {need}",
4562 src.len()
4563 )));
4564 }
4565 for r in 0..rows {
4566 dst[r * row_bytes..(r + 1) * row_bytes]
4567 .copy_from_slice(&src[r * stride..r * stride + row_bytes]);
4568 }
4569 Ok(())
4570 }
4571
4572 /// Set the row stride in bytes for externally allocated buffers with
4573 /// row padding (e.g. V4L2 or GStreamer allocators).
4574 ///
4575 /// The stride is propagated to the EGL DMA-BUF import attributes so
4576 /// the GPU interprets the padded buffer layout correctly. Must be
4577 /// called after [`set_format`](Self::set_format) and before the tensor
4578 /// is first passed to `ImageProcessor::convert`. The stored stride
4579 /// is cleared automatically if the pixel format is later changed.
4580 ///
4581 /// No stride-vs-buffer-size validation is performed because the
4582 /// backing allocation size is not reliably known: external DMA-BUFs
4583 /// may be over-allocated by the allocator, and internal tensors store
4584 /// a logical (unpadded) shape. An incorrect stride will be caught by
4585 /// the EGL driver at import time.
4586 ///
4587 /// # Arguments
4588 ///
4589 /// * `stride` - Row stride in bytes. Must be >= the minimum stride for
4590 /// the format (width * channels * sizeof(T) for packed,
4591 /// width * sizeof(T) for planar/semi-planar).
4592 ///
4593 /// # Errors
4594 ///
4595 /// * `InvalidArgument` if no pixel format is set on this tensor
4596 /// * `InvalidArgument` if `stride` is less than the minimum for the
4597 /// format and width
4598 pub fn set_row_stride(&mut self, stride: usize) -> Result<()> {
4599 let fmt = self.format.ok_or_else(|| {
4600 Error::InvalidArgument("cannot set row_stride without a pixel format".into())
4601 })?;
4602 let w = self.width().ok_or_else(|| {
4603 Error::InvalidArgument("cannot determine width for row_stride validation".into())
4604 })?;
4605 let elem = std::mem::size_of::<T>();
4606 let min_stride = match fmt.layout() {
4607 PixelLayout::Packed => w * fmt.channels() * elem,
4608 PixelLayout::Planar => w * elem,
4609 // Semi-planar: minimum must cover even width for chroma alignment.
4610 PixelLayout::SemiPlanar => w.next_multiple_of(2) * elem,
4611 };
4612 if stride < min_stride {
4613 return Err(Error::InvalidArgument(format!(
4614 "row_stride {stride} < minimum {min_stride} for {fmt:?} at width {w}"
4615 )));
4616 }
4617 self.row_stride = Some(stride);
4618 Ok(())
4619 }
4620
4621 /// Set the row stride without format validation.
4622 ///
4623 /// Use this for raw sub-tensors (e.g. chroma planes) that don't carry
4624 /// format metadata. The caller is responsible for ensuring the stride
4625 /// is valid.
4626 pub fn set_row_stride_unchecked(&mut self, stride: usize) {
4627 self.row_stride = Some(stride);
4628 }
4629
4630 /// Builder-style variant of [`set_row_stride`](Self::set_row_stride),
4631 /// consuming and returning `self`.
4632 ///
4633 /// # Errors
4634 ///
4635 /// Same conditions as [`set_row_stride`](Self::set_row_stride).
4636 pub fn with_row_stride(mut self, stride: usize) -> Result<Self> {
4637 self.set_row_stride(stride)?;
4638 Ok(self)
4639 }
4640
4641 /// Byte offset within the DMA-BUF where image data starts (`None` = 0).
4642 pub fn plane_offset(&self) -> Option<usize> {
4643 self.plane_offset
4644 }
4645
4646 /// The parent-image snapshot if this tensor is a [`view`](Self::view)/
4647 /// [`batch`](Self::batch) sub-region; `None` for a whole tensor. The GL
4648 /// backend keys its import on the parent geometry and renders this view as a
4649 /// `glViewport`/`glScissor` ROI at `(x, y, width, height)`. See [`ViewOrigin`].
4650 pub fn view_origin(&self) -> Option<ViewOrigin> {
4651 self.view_origin
4652 }
4653
4654 /// Set the byte offset within the DMA-BUF where image data starts.
4655 ///
4656 /// Propagated to `EGL_DMA_BUF_PLANE0_OFFSET_EXT` on GPU import.
4657 /// Unlike [`set_row_stride`](Self::set_row_stride), no format is required
4658 /// since the offset is format-independent.
4659 pub fn set_plane_offset(&mut self, offset: usize) {
4660 self.plane_offset = Some(offset);
4661 // The offset consulted by `map()` lives inside the storage variant.
4662 // Keep it in sync with the wrapper field for every backing that
4663 // honors it; see also the clear sites in `set_format` and `reshape`.
4664 //
4665 // Note that `TensorStorage::Dma` is a cfg-multiplexed *name* rather
4666 // than one type -- `DmaTensor` on Linux, `IoSurfaceTensor` on
4667 // macOS/iOS, `AHardwareBufferTensor` on Android,
4668 // `D3d11TextureTensor` on Windows -- so each platform needs its own
4669 // arm below; a missing one would silently become a fall-through
4670 // rather than a compile error, which is how macOS went without one
4671 // for a while. The match is kept exhaustive (no `_ => {}`)
4672 // specifically so that hazard is a compile error, not a silent gap.
4673 match self.storage {
4674 TensorStorage::Mem(ref mut m) => m.set_offset(offset),
4675 #[cfg(target_os = "linux")]
4676 TensorStorage::Dma(ref mut dma) => dma.mmap_offset = offset,
4677 // The offset `IoSurfaceTensor::scoped_pin`/`map_inner` add to the
4678 // locked base address. Writing it here is what makes a
4679 // *reconstructed* view address its own sub-region instead of the
4680 // parent surface's origin; `IoSurfaceTensor::view` already set it
4681 // on the fresh-view path, which is why only the descriptor round
4682 // trip was ever wrong. Idempotent against `subview`, which sets
4683 // the same absolute value by both routes -- pinned by
4684 // `nested_iosurface_subviews_do_not_compound_their_plane_offset`.
4685 #[cfg(any(target_os = "macos", target_os = "ios"))]
4686 TensorStorage::Dma(ref mut io) => io.view_offset = offset,
4687 // The offset every `D3d11TextureTensor` pin adds to its host or
4688 // staging base, in backing (pitched) space like `row_stride()`.
4689 // `view()` bounds-checks the value it computes; this route does
4690 // not, so the pins refuse an offset past the backing instead
4691 // (`check_view_offset` in `d3d11/texture.rs`). Pinned by the
4692 // `d3d11` plane-offset tests in `crates/tensor/tests`.
4693 #[cfg(target_os = "windows")]
4694 TensorStorage::Dma(ref mut tex) => tex.set_view_offset(offset),
4695 // The offset `PboTensor::scoped_pin` adds to the mapped base.
4696 // Writing it here is what makes a *reconstructed* view address
4697 // its own sub-region instead of the parent buffer's origin;
4698 // `PboTensor::view` already set it on the fresh-view path, which
4699 // is why only the round trip was ever wrong -- the same shape
4700 // IOSurface had. Idempotent against `subview`, which sets the
4701 // same absolute value by both routes -- pinned by
4702 // `nested_pbo_subviews_do_not_compound_their_plane_offset`.
4703 //
4704 // Not `cfg`-gated, and deliberately so: unlike
4705 // `TensorStorage::Dma` -- a cfg-multiplexed *name* whose gated
4706 // arm silently becomes `_ => {}` on the other targets, the
4707 // mechanism behind this whole issue class -- `TensorStorage::Pbo`
4708 // is one type on every platform.
4709 TensorStorage::Pbo(ref mut p) => p.view_offset = offset,
4710 // The offset `ShmTensor::map_inner`/`host_pin` add to the mmap
4711 // base, exactly like `MemTensor::offset`. `cfg(unix)` here
4712 // mirrors `TensorStorage::Shm`'s own gate on the enum
4713 // definition -- it is not a per-OS multiplex the way
4714 // `TensorStorage::Dma` is, so this arm cannot silently become
4715 // `_ => {}` on a target where the `Shm` variant still exists.
4716 #[cfg(unix)]
4717 TensorStorage::Shm(ref mut s) => s.set_offset(offset),
4718 // STILL INCOMPLETE: `AHardwareBufferTensor` (Android) has a
4719 // `view_offset` that its own `map()` adds and its own `view()`
4720 // sets, but nothing writes it here yet, so restoring an offset
4721 // onto an already-built tensor is silently a no-op on this
4722 // backing -- a caller gets a map at the parent's origin. Not
4723 // reachable through `TensorDyn::import_descriptor` today
4724 // (`import_storage`'s DMABUF arm is `cfg(target_os = "linux")`),
4725 // so it is latent rather than silent. See
4726 // `edgefirst-python-common`'s `interop::apply_plane_offset` for
4727 // the full per-backing accounting. PBO and Shm were on this list
4728 // until Stage B of the client-side-state work gave them the arms
4729 // above.
4730 #[cfg(target_os = "android")]
4731 TensorStorage::Dma(_) => {}
4732 }
4733 }
4734
4735 /// Colorimetry metadata (`None` = undefined; never auto-filled).
4736 /// The CPU access declared for this tensor at allocation (see
4737 /// [`CpuAccess`]). Views share their parent's declaration.
4738 pub fn cpu_access(&self) -> CpuAccess {
4739 self.cpu_access
4740 }
4741
4742 /// Set the declared CPU access without re-allocating — crate-private:
4743 /// the declaration must reflect the underlying allocation's real
4744 /// capabilities (constructors and importers set it; arbitrary widening
4745 /// would defeat the contract).
4746 // Only the Android AHardwareBuffer importer derives a declaration from
4747 // an existing allocation's usage bits today.
4748 #[cfg_attr(not(target_os = "android"), allow(dead_code))]
4749 pub(crate) fn set_cpu_access_unchecked(&mut self, access: CpuAccess) {
4750 self.cpu_access = access;
4751 }
4752
4753 /// The vendor tile-compression scheme recorded at allocation, or
4754 /// `None` for a linear layout. `Some` means the pixels live in a
4755 /// proprietary tile order: the row-stride accessors describe no
4756 /// meaningful linear layout and CPU maps are best-effort (see
4757 /// [`Compression`]). Only Android hardware-only allocations that
4758 /// requested compression record a scheme.
4759 pub fn compression(&self) -> Option<CompressionScheme> {
4760 self.compression
4761 }
4762
4763 /// Record the compression scheme — crate-private: recording is an
4764 /// allocation-time fact ([`Tensor::image_desc`] sets it; arbitrary
4765 /// mutation would misdescribe the physical layout).
4766 // Only the Android allocation path records a scheme today.
4767 #[cfg_attr(not(target_os = "android"), allow(dead_code))]
4768 pub(crate) fn set_compression_unchecked(&mut self, scheme: Option<CompressionScheme>) {
4769 self.compression = scheme;
4770 }
4771
4772 pub fn colorimetry(&self) -> Option<crate::Colorimetry> {
4773 self.colorimetry
4774 }
4775
4776 /// Attach/clear colorimetry metadata.
4777 pub fn set_colorimetry(&mut self, c: Option<crate::Colorimetry>) {
4778 self.colorimetry = c;
4779 }
4780
4781 /// Builder-style colorimetry attach.
4782 pub fn with_colorimetry(mut self, c: crate::Colorimetry) -> Self {
4783 self.colorimetry = Some(c);
4784 self
4785 }
4786
4787 /// Create a zero-copy sub-region view of this tensor's backing buffer.
4788 ///
4789 /// The returned tensor shares this tensor's allocation (no copy) and maps
4790 /// the window `[offset_bytes, offset_bytes + shape.product()*size_of::<T>())`
4791 /// measured from this tensor's own logical start. N sub-views into one
4792 /// parent can be written independently, enabling batched assembly into a
4793 /// single buffer. Identical semantics across `Mem` (shared `Arc`) and
4794 /// `Dma` (shared fd) backings.
4795 ///
4796 /// # Disjointness
4797 ///
4798 /// Independent writes are sound *only* when the windows do not overlap. The
4799 /// shared backing uses interior mutability (`UnsafeCell` cells), so two
4800 /// sub-views whose byte ranges intersect alias the same cells: writing one
4801 /// while reading or writing the other is a data race and therefore
4802 /// **undefined behaviour**. The caller is responsible for keeping the
4803 /// windows disjoint; this method does not check for overlap.
4804 ///
4805 /// # Errors
4806 ///
4807 /// - [`Error::InvalidOperation`] if the backing is not `Mem` or `Dma`, or
4808 /// if `offset_bytes` is not a multiple of `align_of::<T>()`.
4809 /// - [`Error::InsufficientCapacity`] / [`Error::InvalidSize`] if the window
4810 /// exceeds the parent allocation.
4811 pub(crate) fn subview(&self, offset_bytes: usize, shape: &[usize]) -> Result<Tensor<T>> {
4812 // Offset is absolute into the backing allocation: a sub-view of a
4813 // sub-view composes by adding this tensor's own offset.
4814 let abs_offset = self
4815 .plane_offset
4816 .unwrap_or(0)
4817 .checked_add(offset_bytes)
4818 .ok_or(Error::InvalidSize(offset_bytes))?;
4819 // Every backend exposes `view(offset, shape)` via `TensorTrait`, sharing
4820 // the resource AND `BufferIdentity` (unlike `from_fd`/`from_surface`,
4821 // which mint a fresh identity). The GL backend keys the import on the
4822 // shared identity so offset-distinct sub-views of one buffer reuse a
4823 // single import and address their window via `glViewport`. `Mem`/`Shm`
4824 // share via the allocation `Arc` / a cloned fd; `Pbo` via the GL-buffer
4825 // `Arc`; Linux DMA-BUF / macOS IOSurface via the shared fd / CFRetain.
4826 // `TensorStorage::view` performs the one remaining per-variant dispatch.
4827 let mut t = Tensor::wrap(self.storage.view(offset_bytes, shape)?);
4828 // Inherit the parent's image metadata so the view is a ready-to-use
4829 // sub-image (e.g. a `convert()` destination). The offset is applied
4830 // LAST because `set_format` deliberately clears it — the offset is a
4831 // structural property of the sub-region, not format-dependent metadata.
4832 if let Some(fmt) = self.format {
4833 t.set_format(fmt)?;
4834 }
4835 if let Some(rs) = self.row_stride {
4836 t.set_row_stride_unchecked(rs);
4837 }
4838 t.quantization = self.quantization.clone();
4839 // A sub-region of an image carries the parent's colorimetry — it is the
4840 // same pixels, same color encoding. Inherit it like the other image
4841 // metadata above so a sub-view is a faithful convert() source/target.
4842 t.set_colorimetry(self.colorimetry);
4843 // The declared CPU access is a property of the underlying
4844 // allocation, so every view shares the parent's declaration.
4845 t.cpu_access = self.cpu_access;
4846 // Likewise the recorded compression scheme: the view shares the
4847 // parent's physical layout.
4848 t.compression = self.compression;
4849 if abs_offset > 0 {
4850 t.set_plane_offset(abs_offset);
4851 }
4852 Ok(t)
4853 }
4854
4855 /// Borrow batch element `n` of a batched tensor as a zero-copy view.
4856 ///
4857 /// A batched tensor prepends `N` as the leading dimension over the
4858 /// per-element image layout (`[N, H, W, C]` packed or `[N, C, H, W]`
4859 /// planar) — `N` is over the whole per-element block regardless of
4860 /// `HWC`/`CHW`. `batch(n)` returns element `n`: the contiguous per-element
4861 /// region at byte offset `n * element_size`, sharing the parent's
4862 /// `BufferIdentity` and inheriting its format / row stride / colorimetry.
4863 /// `batch(0)` on a tensor with `N == 1` is equivalent to the whole tensor.
4864 ///
4865 /// # Errors
4866 ///
4867 /// - [`Error::BatchIndexOutOfBounds`] if `n >= N`.
4868 /// - [`Error::InvalidShape`] if the tensor is not batched (a formatted
4869 /// tensor whose rank lacks the leading `N`, or an empty shape).
4870 pub fn batch(&self, n: usize) -> Result<Tensor<T>> {
4871 let shape = self.shape();
4872 // With a format we know the exact per-element rank, so a missing leading
4873 // `N` is a misuse we reject rather than silently treating a spatial dim
4874 // as the batch. Raw tensors take shape[0] as `N` by contract.
4875 if let Some(fmt) = self.format {
4876 let elem_rank = match fmt.layout() {
4877 PixelLayout::SemiPlanar => 2,
4878 _ => 3,
4879 };
4880 if shape.len() != elem_rank + 1 {
4881 return Err(Error::InvalidShape(format!(
4882 "batch(): tensor is not batched ({fmt:?} expects a leading N over a \
4883 {elem_rank}-D element, got shape {shape:?})"
4884 )));
4885 }
4886 }
4887 let batch = *shape
4888 .first()
4889 .ok_or_else(|| Error::InvalidShape("batch(): empty shape".into()))?;
4890 if n >= batch {
4891 return Err(Error::BatchIndexOutOfBounds { index: n, batch });
4892 }
4893 let elem_shape: Vec<usize> = shape[1..].to_vec();
4894 let elem_count: usize = elem_shape.iter().product();
4895 let elem_bytes = elem_count
4896 .checked_mul(std::mem::size_of::<T>())
4897 .ok_or(Error::InvalidSize(elem_count))?;
4898 let offset = n.checked_mul(elem_bytes).ok_or(Error::InvalidSize(n))?;
4899 // For a packed `[N, H, W, C]` tensor the N tiles stack vertically in the
4900 // shared buffer, so the GL import sees one `(W, N*H)` parent and each
4901 // tile is the row-band at `y = n*H`. Snapshot that parent so the backend
4902 // imports once and renders the tile via `glViewport`. Non-packed
4903 // (planar/semi-planar) batching keeps the per-slot path for now (planar
4904 // NCHW tiling is a separate step); raw tensors have no pixel geometry.
4905 let view_origin = match self.format.map(|f| f.layout()) {
4906 Some(PixelLayout::Packed) => {
4907 let tile_h = elem_shape[0];
4908 let tile_w = elem_shape[1];
4909 // Per-row pitch of the tall `(W, N*H)` parent — padded stride if
4910 // set, else the tight row width. The GL import keys on this.
4911 let bpp = elem_shape[2] * std::mem::size_of::<T>();
4912 let parent_stride = self.effective_row_stride().unwrap_or(tile_w * bpp);
4913 Some(self.compose_view_origin(tile_w, batch * tile_h, parent_stride, 0, n * tile_h))
4914 }
4915 _ => None,
4916 };
4917 let mut t = self.subview(offset, &elem_shape)?;
4918 t.view_origin = view_origin;
4919 Ok(t)
4920 }
4921
4922 /// Borrow a rectangular spatial sub-region of an image tensor as a
4923 /// zero-copy view — the **destination/source crop** primitive.
4924 ///
4925 /// `region` is in pixels of the image's leading frame. The returned view
4926 /// shares the parent's `BufferIdentity` and addresses the sub-rectangle by
4927 /// offset + the **parent's** row pitch (so each row lands at the correct
4928 /// columns). `convert(src, &mut dst.view(rect), …)` renders into that
4929 /// sub-rectangle of `dst`; a letterbox fit then clears the view and renders
4930 /// the aspect-preserved content into its inner region. `view`/`batch`/the
4931 /// whole tensor are the one coherent destination model — there is no
4932 /// separate `dst_rect`.
4933 ///
4934 /// # Known limitation: bottom-right multi-row views
4935 ///
4936 /// A multi-row view is mapped as `parent_pitch × rows` bytes **from its own
4937 /// offset**, so the window it asks for overruns the allocation by
4938 /// `region.x * bpp` whenever the view is both offset horizontally *and*
4939 /// reaches the parent's last row. Such a view constructs fine but fails at
4940 /// [`map`](TensorTrait::map) time with [`Error::InsufficientCapacity`] —
4941 /// loudly, never as an out-of-bounds access. For a tiled destination this is
4942 /// the bottom row of tiles at every column but the first. The single-row
4943 /// case is already special-cased below (it keeps its own tight stride); the
4944 /// general fix is to clamp the last row's exposure to the view's own row
4945 /// bytes, which is not yet implemented.
4946 ///
4947 /// # Errors
4948 ///
4949 /// - [`Error::RegionOutOfBounds`] if `region` exceeds the image bounds.
4950 /// - [`Error::InvalidOperation`] if the tensor is not a packed-format image
4951 /// (planar/semi-planar spatial sub-rects are not a single strided window;
4952 /// use [`batch`](Self::batch) for batched planar tensors).
4953 pub fn view(&self, region: Region) -> Result<Tensor<T>> {
4954 let fmt = self.format.ok_or_else(|| {
4955 Error::InvalidOperation("view() requires a formatted image tensor".into())
4956 })?;
4957 if fmt.layout() != PixelLayout::Packed {
4958 return Err(Error::InvalidOperation(format!(
4959 "view() supports packed formats only (got {fmt:?}); use batch(n) for batched \
4960 planar tensors"
4961 )));
4962 }
4963 let w = self
4964 .width()
4965 .ok_or_else(|| Error::InvalidOperation("view(): tensor has no image width".into()))?;
4966 let h = self
4967 .height()
4968 .ok_or_else(|| Error::InvalidOperation("view(): tensor has no image height".into()))?;
4969 if !region.fits_within(w, h) {
4970 return Err(Error::RegionOutOfBounds {
4971 region,
4972 bounds: (w, h),
4973 });
4974 }
4975 let elem = std::mem::size_of::<T>();
4976 let bpp = fmt.channels() * elem;
4977 let stride = self.effective_row_stride().unwrap_or(w * bpp);
4978 let offset = region
4979 .y
4980 .checked_mul(stride)
4981 .and_then(|yo| yo.checked_add(region.x.checked_mul(bpp)?))
4982 .ok_or(Error::InvalidSize(region.y))?;
4983 let sub_shape = fmt
4984 .allocation_shape(region.width, region.height)
4985 .ok_or_else(|| Error::InvalidShape(format!("view(): invalid shape for {fmt:?}")))?;
4986 let mut t = self.subview(offset, &sub_shape)?;
4987 // A multi-row sub-rect must advance rows by the PARENT pitch so each row
4988 // addresses the correct columns. A single-row view uses its own tight
4989 // stride — the parent pitch would make the strided `map()` (which exposes
4990 // `stride × rows`) expose a trailing row that runs past the buffer tail
4991 // for an offset (x>0 / bottom) view. The GL backend does NOT rely on this
4992 // (single-row-tight) `row_stride`: it reads the parent pitch from
4993 // `view_origin.parent_row_stride` so its import/cache pitch stays
4994 // parent-consistent for views of any height — see `ViewOrigin`.
4995 let view_stride = if region.height > 1 {
4996 stride
4997 } else {
4998 region.width * bpp
4999 };
5000 t.set_row_stride_unchecked(view_stride);
5001 // Snapshot the parent `(w, h, row_stride)` so the GL backend imports the
5002 // parent once (keyed on the parent pitch, not this view's possibly-tight
5003 // single-row stride) and renders this sub-rect as a `glViewport`/
5004 // `glScissor` ROI at `(region.x, region.y)`. Composes when viewing an
5005 // existing view.
5006 t.view_origin = Some(self.compose_view_origin(w, h, stride, region.x, region.y));
5007 Ok(t)
5008 }
5009
5010 /// Build the [`ViewOrigin`] for a new sub-region of `self`. When `self` is a
5011 /// whole tensor the snapshot names `self` as the parent; when `self` is
5012 /// already a view, the snapshot keeps the **root** parent and accumulates the
5013 /// local origin so nested views still resolve to one import.
5014 fn compose_view_origin(
5015 &self,
5016 parent_width: usize,
5017 parent_height: usize,
5018 parent_row_stride: usize,
5019 x: usize,
5020 y: usize,
5021 ) -> ViewOrigin {
5022 match self.view_origin {
5023 Some(root) => ViewOrigin {
5024 parent_width: root.parent_width,
5025 parent_height: root.parent_height,
5026 parent_row_stride: root.parent_row_stride,
5027 x: root.x.saturating_add(x),
5028 y: root.y.saturating_add(y),
5029 },
5030 None => ViewOrigin {
5031 parent_width,
5032 parent_height,
5033 parent_row_stride,
5034 x,
5035 y,
5036 },
5037 }
5038 }
5039
5040 /// The wrapped `PboTensor`, for this module's own PBO accessors.
5041 ///
5042 /// Crate-private: `PboTensor<T>` is a `static`-side value with no
5043 /// `dynamic` counterpart to lend (there the buffer lives inside
5044 /// `libedgefirst_tensor.so`), so a public borrow of it could not have
5045 /// the same signature on both backends. Callers use
5046 /// [`Self::pbo_id`]/[`Self::pbo_is_mapped`], which are the whole surface
5047 /// `edgefirst-image` ever used.
5048 pub(crate) fn as_pbo(&self) -> Option<&PboTensor<T>> {
5049 match &self.storage {
5050 TensorStorage::Pbo(p) => Some(p),
5051 _ => None,
5052 }
5053 }
5054
5055 /// The GL buffer ID backing this tensor (for GL backends / the
5056 /// cross-package protocol's `handle`). `None` when the tensor is not
5057 /// PBO-backed.
5058 pub fn pbo_id(&self) -> Option<u32> {
5059 self.as_pbo().map(PboTensor::buffer_id)
5060 }
5061
5062 /// Whether this PBO currently holds -- or is establishing -- a CPU
5063 /// mapping. `None` when the tensor is not PBO-backed.
5064 ///
5065 /// The question a GL caller is really asking is "is this buffer free
5066 /// for GL operations?", so a map or unmap still in flight answers
5067 /// `Some(true)`: the GL side must not touch the buffer until the
5068 /// transition finishes. See [`PboTensor::is_mapped`].
5069 pub fn pbo_is_mapped(&self) -> Option<bool> {
5070 self.as_pbo().map(PboTensor::is_mapped)
5071 }
5072
5073 /// The C-ABI `PboOpsVtable` address backing this tensor (for the
5074 /// cross-package protocol's `ptr` under `kind::PBO`; see
5075 /// [`crate::protocol::TensorDesc::ptr`]'s doc comment). `None` when the
5076 /// tensor is not PBO-backed.
5077 pub fn pbo_vtable_ptr(&self) -> Option<*const std::ffi::c_void> {
5078 self.as_pbo()
5079 .map(|p| p.pbo_vtable() as *const _ as *const std::ffi::c_void)
5080 }
5081
5082 /// A type-erased keepalive that must stay alive for at least as long as
5083 /// [`Self::pbo_vtable_ptr`]'s address is used -- see
5084 /// [`PboTensor::pbo_keepalive`]'s own doc comment. `None` when the
5085 /// tensor is not PBO-backed.
5086 pub fn pbo_keepalive(&self) -> Option<Arc<dyn Send + Sync>> {
5087 self.as_pbo().map(PboTensor::pbo_keepalive)
5088 }
5089
5090 /// Downcast to DMA tensor reference (for EGL import, G2D).
5091 #[cfg(target_os = "linux")]
5092 pub fn as_dma(&self) -> Option<&DmaTensor<T>> {
5093 match &self.storage {
5094 TensorStorage::Dma(d) => Some(d),
5095 _ => None,
5096 }
5097 }
5098
5099 /// Borrow the DMA-BUF file descriptor backing this tensor.
5100 ///
5101 /// # Returns
5102 ///
5103 /// A borrowed reference to the DMA-BUF file descriptor, tied to `self`'s
5104 /// lifetime.
5105 ///
5106 /// # Errors
5107 ///
5108 /// Returns `Error::NotImplemented` if the tensor is not DMA-backed.
5109 #[cfg(target_os = "linux")]
5110 pub fn dmabuf(&self) -> Result<std::os::fd::BorrowedFd<'_>> {
5111 use std::os::fd::AsFd;
5112 match &self.storage {
5113 TensorStorage::Dma(dma) => Ok(dma.fd.as_fd()),
5114 _ => Err(Error::NotImplemented(format!(
5115 "dmabuf requires DMA-backed tensor, got {:?}",
5116 self.storage.memory()
5117 ))),
5118 }
5119 }
5120
5121 /// Construct a Tensor from a PBO tensor (for GL backends that allocate PBOs).
5122 ///
5123 /// Infallible in this backend (`Ok` always) -- kept as `Result` to match
5124 /// `dynamic`'s own `Tensor::from_pbo` (`dynamic_tensor.rs`), which
5125 /// genuinely can fail: it must allocate a real `ef_tensor_*` handle
5126 /// alongside `pbo` to carry shape/dtype/format metadata (this backend
5127 /// has nowhere else to put it), an allocation this backend never
5128 /// needed since it just wraps `pbo` into `TensorStorage::Pbo` directly.
5129 /// One shared, fallible signature lets `edgefirst-image`'s own call
5130 /// sites (`gl/threaded.rs`) stay backend-agnostic.
5131 pub fn from_pbo(pbo: PboTensor<T>) -> Result<Self> {
5132 Ok(Self {
5133 storage: TensorStorage::Pbo(pbo),
5134 format: None,
5135 chroma: None,
5136 row_stride: None,
5137 plane_offset: None,
5138 quantization: None,
5139 cuda: None,
5140 colorimetry: None,
5141 cpu_access: CpuAccess::ReadWrite,
5142 compression: None,
5143 view_origin: None,
5144 })
5145 }
5146
5147 /// The CUDA registration for this tensor, if any (set at creation on CUDA devices).
5148 pub fn cuda(&self) -> Option<&crate::cuda::CudaHandle> {
5149 self.cuda.as_ref()
5150 }
5151
5152 /// Attach a CUDA handle (called by ImageProcessor::create_image after registering a PBO).
5153 pub fn set_cuda_handle(&mut self, h: crate::cuda::CudaHandle) {
5154 self.cuda = Some(h);
5155 }
5156
5157 /// Fast-fail CUDA map: None (no GL routing) when no handle; else map (PBO routes to the GL worker).
5158 ///
5159 /// Returns a scoped [`CudaMap`] guard holding the raw CUDA device pointer
5160 /// for the duration of the mapping. For GL-buffer-backed tensors the unmap is deferred until the
5161 /// guard drops, freeing the PBO for the next `convert()` call. When no CUDA handle is attached
5162 /// (the common case for plain `Mem`/`DMA` tensors without CUDA registration), returns `None`
5163 /// immediately — no GL routing, no allocation.
5164 ///
5165 /// # Example — zero-copy CUDA input with host fallback
5166 ///
5167 /// ```no_run
5168 /// use edgefirst_tensor::{Tensor, TensorMemory, TensorTrait};
5169 /// # fn feed_tensorrt(_dptr: *mut std::ffi::c_void, _bytes: usize) {}
5170 /// # fn demo(t: &Tensor<f32>) {
5171 /// // Try the zero-copy CUDA device pointer first.
5172 /// if let Some(cuda) = t.cuda_map() {
5173 /// feed_tensorrt(cuda.device_ptr(), cuda.len());
5174 /// // `cuda` (a CudaMap guard) unmaps when it goes out of scope, freeing
5175 /// // the GPU buffer for the next convert().
5176 /// } else {
5177 /// // Fall back to the host mapping when no CUDA handle is attached.
5178 /// let _host = t.map().expect("host map fallback must succeed");
5179 /// // `_host` is a TensorMap<f32> that derefs to &[f32].
5180 /// }
5181 /// # }
5182 /// ```
5183 ///
5184 /// # D3D11 texture tensors: the mapping is tight rows
5185 ///
5186 /// On Windows the mapping is a linear copy of the texture laid out at
5187 /// `width * bytes_per_texel`, and `len()` is that times the row count --
5188 /// *not* [`row_stride`](Self::row_stride), which reports the D3D11 staging
5189 /// pitch a CPU map sees and is the larger number on a padded backing. A
5190 /// consumer that strides the device pointer by `row_stride()` reads past
5191 /// the end of the mapping.
5192 ///
5193 /// Semi-planar formats (NV12, NV16, NV24) are the exception: their texture
5194 /// is allocated as wide as its own row pitch, so the tight row *is* the
5195 /// pitch and the mapping equals `row_stride() x combined_height` bytes.
5196 /// Striding the device pointer by `row_stride()` is exactly right there.
5197 pub fn cuda_map(&self) -> Option<crate::cuda::CudaMap<'_>> {
5198 self.cuda.as_ref()?.map()
5199 }
5200
5201 /// Writable counterpart of [`cuda_map`](Self::cuda_map). See
5202 /// [`CudaHandle::map_mut`](crate::cuda::CudaHandle::map_mut) for which
5203 /// backings distinguish the two.
5204 ///
5205 /// The D3D11 arm is the one that does: the guard's release copies the
5206 /// buffer back into the texture. Its layout is [`cuda_map`](Self::cuda_map)'s
5207 /// -- tight rows of `width * bytes_per_texel`, not `row_stride()` -- and
5208 /// writing it at the wrong pitch scrambles the image rather than failing.
5209 pub fn cuda_map_mut(&self) -> Option<crate::cuda::CudaMap<'_>> {
5210 self.cuda.as_ref()?.map_mut()
5211 }
5212
5213 /// Attempt to attach a CUDA `ExternalMemory` handle for DMA-backed tensors.
5214 ///
5215 /// On a CUDA-capable host, imports the DMA-BUF fd via
5216 /// `cudaImportExternalMemory(OpaqueFd)` and maps it to a device pointer.
5217 /// Sets `self.cuda` to a persistent `ExternalMem` handle on success. No-op
5218 /// if CUDA is unavailable, the tensor is not DMA-backed, or a handle is
5219 /// already set. Import failure is silently ignored — the tensor remains
5220 /// usable without a CUDA handle.
5221 ///
5222 /// # RUNTIME-UNVALIDATED
5223 ///
5224 /// No test platform has both `/dev/dma_heap` and a CUDA device. ABI is
5225 /// layout-asserted vs. CUDA 12.6 `driver_types.h`; the mechanism is proven
5226 /// by gpu-probe O5 on Orin. Best-effort: tensor creation never fails here.
5227 #[cfg(target_os = "linux")]
5228 pub fn try_init_dma_cuda(&mut self) {
5229 // Fast-path: already imported, CUDA not available, or not a DMA tensor.
5230 if self.cuda.is_some() || !crate::cuda::is_cuda_available() {
5231 return;
5232 }
5233 let (raw_fd, buf_size) = match &self.storage {
5234 TensorStorage::Dma(dma) => {
5235 use std::os::fd::AsRawFd;
5236 (dma.fd.as_raw_fd(), dma.buf_size)
5237 }
5238 _ => return,
5239 };
5240 if let Some((ext, dptr)) = crate::cuda::import_dma_fd(raw_fd, buf_size) {
5241 self.cuda = Some(crate::cuda::CudaHandle::new_external(ext, dptr, buf_size));
5242 }
5243 }
5244
5245 /// Attempt to attach a CUDA external-memory handle for D3D11 texture
5246 /// tensors — the Windows counterpart of
5247 /// [`try_init_dma_cuda`](Self::try_init_dma_cuda).
5248 ///
5249 /// Imports the texture's shared NT handle via
5250 /// `cudaImportExternalMemory(D3D11Resource)` and the process device's
5251 /// fence as an external semaphore, so [`cuda_map`](Self::cuda_map) hands
5252 /// out a linear device buffer whose refresh is ordered after the D3D11
5253 /// work a producer recorded with [`set_gpu_write`](Self::set_gpu_write).
5254 ///
5255 /// No-op when a handle is already attached or the tensor is not
5256 /// texture-backed. Import failure is silently ignored: the WARP adapter
5257 /// has no CUDA device at all, and a texture whose import fails is still a
5258 /// valid tensor whose `cuda_map()` is `None`.
5259 #[cfg(target_os = "windows")]
5260 pub fn try_init_d3d11_cuda(&mut self) {
5261 if self.cuda.is_some() {
5262 return;
5263 }
5264 let TensorStorage::Dma(tex) = &self.storage else {
5265 return;
5266 };
5267 if let Some(h) = tex.try_init_cuda() {
5268 self.cuda = Some(h);
5269 }
5270 }
5271}
5272
5273/// Windows D3D11 texture accessors. No `'static` bound: none of these needs
5274/// `dtype_of::<T>()`.
5275#[cfg(all(feature = "static", target_os = "windows"))]
5276impl<T> Tensor<T>
5277where
5278 T: Num + Clone + fmt::Debug + Send + Sync,
5279{
5280 /// The borrowed `ID3D11Texture2D*` backing this tensor, valid while the
5281 /// tensor lives. `None` when the tensor is not texture-backed.
5282 pub fn d3d11_texture(&self) -> Option<*mut std::ffi::c_void> {
5283 match &self.storage {
5284 TensorStorage::Dma(t) => Some(t.texture_ptr()),
5285 _ => None,
5286 }
5287 }
5288
5289 /// The texture geometry the HAL chose for this image: DXGI format, texel
5290 /// dimensions and the matching GL internal format. `None` when the tensor
5291 /// is not texture-backed.
5292 pub fn d3d11_layout(&self) -> Option<crate::d3d11_layout::D3d11ImageLayout> {
5293 match &self.storage {
5294 TensorStorage::Dma(t) => Some(t.layout()),
5295 _ => None,
5296 }
5297 }
5298
5299 /// A duplicated NT handle the caller owns: open it on a D3D12 device,
5300 /// another D3D11 device, or CUDA, or duplicate it into another process.
5301 /// Closing it does not affect this tensor.
5302 pub fn d3d11_shared_handle(&self) -> Result<std::os::windows::io::OwnedHandle> {
5303 match &self.storage {
5304 TensorStorage::Dma(t) => t.shared_handle(),
5305 _ => Err(Error::NotImplemented(
5306 "not a D3D11 texture tensor".to_owned(),
5307 )),
5308 }
5309 }
5310
5311 /// The fence handle plus value a GPU consumer waits on before reading this
5312 /// texture, or `None` when no GPU write has been recorded on it.
5313 pub fn gpu_completion(&self) -> Result<Option<crate::d3d11::GpuCompletion>> {
5314 match &self.storage {
5315 TensorStorage::Dma(t) => t.gpu_completion(),
5316 _ => Err(Error::NotImplemented(
5317 "not a D3D11 texture tensor".to_owned(),
5318 )),
5319 }
5320 }
5321
5322 /// Record that GPU work writing this texture completes at `value` of the
5323 /// process device's fence. Monotonic: an older value never displaces a
5324 /// newer one, so recording out of order is safe rather than merely
5325 /// tolerated.
5326 ///
5327 /// Takes `&self`, not `&mut self`: the fence slot is an `AtomicU64` and
5328 /// this is a `fetch_max` on it. A producer records a write on the very
5329 /// tensor its consumers already hold and are reading, so demanding
5330 /// exclusive access here would forbid the one call this exists for.
5331 pub fn set_gpu_write(&self, value: u64) -> Result<()> {
5332 match &self.storage {
5333 TensorStorage::Dma(t) => {
5334 t.set_gpu_write(value);
5335 Ok(())
5336 }
5337 _ => Err(Error::NotImplemented(
5338 "not a D3D11 texture tensor".to_owned(),
5339 )),
5340 }
5341 }
5342
5343 /// The tensor's own NT handle value -- not a duplicate -- for descriptors
5344 /// whose consumer keeps the producing tensor alive.
5345 ///
5346 /// Borrowed rather than duplicated because the C exports (task A6) hand
5347 /// out owned duplicates, which a `Copy` descriptor cannot close; the
5348 /// cross-package protocol's `kind::D3D11_TEXTURE` reads this instead.
5349 pub(crate) fn d3d11_shared_handle_value(&self) -> Option<usize> {
5350 match &self.storage {
5351 TensorStorage::Dma(t) => Some(t.shared_handle_value()),
5352 _ => None,
5353 }
5354 }
5355}
5356
5357/// The recorded GPU write value, on every platform: the descriptor path and
5358/// the Python bindings forward it without a `cfg` of their own, and only
5359/// Windows can record one.
5360#[cfg(feature = "static")]
5361impl<T> Tensor<T>
5362where
5363 T: Num + Clone + fmt::Debug + Send + Sync,
5364{
5365 /// The fence value of the newest GPU write recorded on this tensor, or 0
5366 /// when there is none: the `value` of `gpu_completion` without the
5367 /// duplicated fence handle, for a consumer that already holds the process
5368 /// fence. 0 for a tensor that is not a D3D11 texture, and on every
5369 /// platform other than Windows.
5370 pub fn gpu_write_value(&self) -> u64 {
5371 #[cfg(target_os = "windows")]
5372 {
5373 if let TensorStorage::Dma(t) = &self.storage {
5374 return t.last_gpu_write();
5375 }
5376 }
5377 0
5378 }
5379}
5380
5381/// Records the D3D11 backing's row pitch on `t` when it exceeds the natural
5382/// stride for its format and width.
5383///
5384/// The staging pitch is the row stride CPU consumers see, and for a
5385/// semi-planar tensor it is also the texture's own width, which the GL
5386/// backend's Path B shader takes as its `tex_width`. Three constructors need
5387/// it identically -- `Tensor::image`'s Windows arm and the two wrap
5388/// constructors below -- and a wrapped texture is the case that makes it
5389/// load-bearing: without it such a tensor reports the natural `even(width)`
5390/// stride while its texture is the driver's pitch wide, and the shader samples
5391/// a padded texture at the image width.
5392///
5393/// `t.format` must already be set: the natural stride comes from it.
5394#[cfg(all(feature = "static", target_os = "windows"))]
5395fn record_d3d11_row_stride<T>(t: &mut Tensor<T>, format: PixelFormat, what: &str)
5396where
5397 T: Num + Clone + fmt::Debug + Send + Sync,
5398{
5399 let pitch = match &t.storage {
5400 TensorStorage::Dma(tex) => tex.image_backing_row_stride(),
5401 _ => None,
5402 };
5403 let (Some(pitch), Some(natural)) = (pitch, t.effective_row_stride()) else {
5404 return;
5405 };
5406 if pitch > natural {
5407 log::debug!(
5408 "{what}: the D3D11 backing pitch for the {format:?} texture is {pitch} bytes \
5409 (natural {natural}); recording row stride"
5410 );
5411 t.set_row_stride_unchecked(pitch);
5412 }
5413}
5414
5415/// The two texture-importing constructors. Separate from the accessors above
5416/// because `dtype_of::<T>()` needs `T: 'static`.
5417#[cfg(all(feature = "static", target_os = "windows"))]
5418impl<T> Tensor<T>
5419where
5420 T: Num + Clone + fmt::Debug + Send + Sync + 'static,
5421{
5422 /// Wrap an existing `ID3D11Texture2D` as a tensor. The texture is checked
5423 /// against the layout the HAL would have chosen for `format`/`width`/
5424 /// `height` and must live on the HAL's process device; ownership stays
5425 /// with the caller, since this takes its own reference.
5426 ///
5427 /// The backing's row pitch is recorded as the tensor's
5428 /// [`row_stride`](Self::row_stride) when it exceeds the natural stride, so
5429 /// a wrapped semi-planar texture -- whose width *is* the pitch -- reports
5430 /// the geometry its consumers address it by.
5431 ///
5432 /// # Safety
5433 ///
5434 /// `texture` must be null or a live `ID3D11Texture2D` created on the HAL
5435 /// device ([`crate::d3d11::device()`]).
5436 pub unsafe fn from_d3d11_texture(
5437 texture: *mut std::ffi::c_void,
5438 width: usize,
5439 height: usize,
5440 format: PixelFormat,
5441 access: CpuAccess,
5442 name: Option<&str>,
5443 ) -> Result<Self> {
5444 let dtype = dtype_of::<T>().ok_or_else(|| {
5445 Error::InvalidArgument(format!(
5446 "from_d3d11_texture: element type {} has no DType",
5447 std::any::type_name::<T>()
5448 ))
5449 })?;
5450 let shape = format.allocation_shape(width, height).ok_or_else(|| {
5451 Error::InvalidArgument(format!(
5452 "invalid dimensions {width}x{height} for format {format:?}"
5453 ))
5454 })?;
5455 // SAFETY: the caller guarantees `texture` is null or a live
5456 // `ID3D11Texture2D`; the storage takes its own reference.
5457 let storage = unsafe {
5458 crate::d3d11::texture::D3d11TextureTensor::<T>::from_d3d11_texture(
5459 texture, width, height, format, dtype, &shape, access, name,
5460 )
5461 }?;
5462 let mut t = Self::wrap(TensorStorage::Dma(storage));
5463 t.format = Some(format);
5464 record_d3d11_row_stride(&mut t, format, "from_d3d11_texture");
5465 t.cpu_access = access;
5466 t.try_init_d3d11_cuda();
5467 Ok(t)
5468 }
5469
5470 /// Open a shared texture by its NT handle on the HAL's device. When
5471 /// `completion` is given, its fence value is waited for on the immediate
5472 /// context, so a same-device reader needs no further ordering. Both
5473 /// handles stay owned by the caller: this duplicates what it keeps.
5474 ///
5475 /// Records the backing's row pitch as
5476 /// [`row_stride`](Self::row_stride) on the same rule as
5477 /// [`from_d3d11_texture`](Self::from_d3d11_texture).
5478 ///
5479 /// # Safety
5480 ///
5481 /// `handle` must be an NT shared handle of a D3D11 texture, valid in this
5482 /// process, and `completion`'s handle a shared fence handle.
5483 pub unsafe fn from_d3d11_shared_handle(
5484 handle: std::os::windows::io::RawHandle,
5485 width: usize,
5486 height: usize,
5487 format: PixelFormat,
5488 access: CpuAccess,
5489 completion: Option<(std::os::windows::io::RawHandle, u64)>,
5490 name: Option<&str>,
5491 ) -> Result<Self> {
5492 let dtype = dtype_of::<T>().ok_or_else(|| {
5493 Error::InvalidArgument(format!(
5494 "from_d3d11_shared_handle: element type {} has no DType",
5495 std::any::type_name::<T>()
5496 ))
5497 })?;
5498 let shape = format.allocation_shape(width, height).ok_or_else(|| {
5499 Error::InvalidArgument(format!(
5500 "invalid dimensions {width}x{height} for format {format:?}"
5501 ))
5502 })?;
5503 // SAFETY: the caller guarantees the handles are valid in this process.
5504 let storage = unsafe {
5505 crate::d3d11::texture::D3d11TextureTensor::<T>::from_d3d11_shared_handle(
5506 handle, width, height, format, dtype, &shape, access, completion, name,
5507 )
5508 }?;
5509 let mut t = Self::wrap(TensorStorage::Dma(storage));
5510 t.format = Some(format);
5511 record_d3d11_row_stride(&mut t, format, "from_d3d11_shared_handle");
5512 t.cpu_access = access;
5513 t.try_init_d3d11_cuda();
5514 Ok(t)
5515 }
5516}
5517
5518// Quantization accessors — type-gated to integer element types via the
5519// sealed `IntegerType` trait. Calling `.quantization()` on a `Tensor<f32>`
5520// produces a compile error, not a runtime one.
5521#[cfg(feature = "static")]
5522impl<T> Tensor<T>
5523where
5524 T: IntegerType + Num + Clone + fmt::Debug + Send + Sync,
5525{
5526 /// Quantization metadata for this tensor, if set.
5527 pub fn quantization(&self) -> Option<&Quantization> {
5528 self.quantization.as_ref()
5529 }
5530
5531 /// Attach quantization metadata to this tensor. Validates against the
5532 /// tensor's shape — returns [`Error::QuantizationInvalid`] on any
5533 /// inconsistency (mismatched scale/zp lengths, out-of-range axis, etc.).
5534 pub fn set_quantization(&mut self, q: Quantization) -> Result<()> {
5535 q.validate(self.shape())?;
5536 self.quantization = Some(q);
5537 Ok(())
5538 }
5539
5540 /// Builder-style variant of [`Self::set_quantization`]. Consumes `self`
5541 /// and returns `Result<Self>` — on success yields the tensor with the
5542 /// attached quantization; on validation failure returns
5543 /// [`Error::QuantizationInvalid`] and drops `self` (the tensor is not
5544 /// returned in the error arm).
5545 pub fn with_quantization(mut self, q: Quantization) -> Result<Self> {
5546 self.set_quantization(q)?;
5547 Ok(self)
5548 }
5549
5550 /// Clear any quantization metadata on this tensor.
5551 pub fn clear_quantization(&mut self) {
5552 self.quantization = None;
5553 }
5554}
5555
5556#[cfg(feature = "static")]
5557fn map_strided_host_view<'a, T>(
5558 tensor: &Tensor<T>,
5559 stride: usize,
5560 access: CpuAccess,
5561 non_blocking: bool,
5562) -> Result<crate::view::HostView<'a, T>>
5563where
5564 T: Num + Clone + fmt::Debug + Send + Sync + 'a,
5565{
5566 // Rows sit at `stride`-byte spacing. The row count is the first
5567 // shape dim for packed `[H, W, C]` and semi-planar `[H*k, W]`,
5568 // but planar `[C, H, W]` stacks C planes of H rows — its surface
5569 // row count is `C × H` (`shape[0]` alone would expose a 3-row
5570 // window and truncate the map; first hit by Android planar-F16
5571 // AHardwareBuffers, whose gralloc pads the pitch — macOS/Linux
5572 // planar pitches happen to be naturally aligned so no stride was
5573 // ever recorded there).
5574 let rows = strided_map_rows(tensor)?;
5575 let total_bytes = stride.checked_mul(rows).ok_or_else(|| {
5576 Error::InvalidOperation(format!(
5577 "Tensor::map: row_stride {stride} × rows {rows} overflows usize"
5578 ))
5579 })?;
5580 map_strided_storage(tensor, stride, rows, total_bytes, access, non_blocking)
5581}
5582
5583#[cfg(feature = "static")]
5584fn strided_map_rows<T>(tensor: &Tensor<T>) -> Result<usize>
5585where
5586 T: Num + Clone + fmt::Debug + Send + Sync,
5587{
5588 match tensor.format.map(|f| f.layout()) {
5589 Some(PixelLayout::Planar) => {
5590 let s = tensor.shape();
5591 if s.len() < 2 {
5592 return Err(Error::InvalidOperation(
5593 "Tensor::map: strided planar mapping requires [C, H, W] shape".into(),
5594 ));
5595 }
5596 s[0].checked_mul(s[1]).ok_or_else(|| {
5597 Error::InvalidOperation(format!(
5598 "Tensor::map: planar rows {} × {} overflows usize",
5599 s[0], s[1]
5600 ))
5601 })
5602 }
5603 _ => tensor.shape().first().copied().ok_or_else(|| {
5604 Error::InvalidOperation(
5605 "Tensor::map: strided mapping requires a non-empty shape".into(),
5606 )
5607 }),
5608 }
5609}
5610
5611#[cfg(feature = "static")]
5612fn map_strided_storage<'a, T>(
5613 tensor: &Tensor<T>,
5614 _stride: usize,
5615 _rows: usize,
5616 total_bytes: usize,
5617 access: CpuAccess,
5618 _non_blocking: bool,
5619) -> Result<crate::view::HostView<'a, T>>
5620where
5621 T: Num + Clone + fmt::Debug + Send + Sync + 'a,
5622{
5623 match &tensor.storage {
5624 #[cfg(target_os = "linux")]
5625 TensorStorage::Dma(dma) => {
5626 // `set_row_stride()` only validates `stride >= min_stride`,
5627 // not that `stride × rows` fits the DMA-BUF, so re-check
5628 // here — mapping past `buf_size` would SIGBUS on access.
5629 // `buf_size` is real kernel truth either way, not the
5630 // importer's say-so: for a self-allocated tensor it is
5631 // the allocation size; for an imported one it is
5632 // `fstat`-derived when the kernel reports a usable size
5633 // (`DmaTensor::from_fd`), or the declared logical size in
5634 // the rare fallback -- either way an out-of-range mmap
5635 // against the real dma-buf is still rejected by the
5636 // kernel, so this bound stays sound regardless.
5637 let available_bytes = dma.buf_size.saturating_sub(dma.mmap_offset);
5638 if total_bytes > available_bytes {
5639 return Err(Error::InvalidOperation(format!(
5640 "Tensor::map: strided mapping needs {total_bytes} bytes \
5641 but DMA buffer only has {available_bytes} available \
5642 (buf_size={}, mmap_offset={}, stride={_stride}, rows={_rows}); \
5643 the row_stride was likely set larger than the original allocation",
5644 dma.buf_size, dma.mmap_offset
5645 )));
5646 }
5647 dma.map_with_byte_size(total_bytes, access)
5648 }
5649 TensorStorage::Mem(mem) => {
5650 let capacity = tensor.storage.capacity_bytes();
5651 if total_bytes > capacity {
5652 return Err(Error::InsufficientCapacity {
5653 needed: total_bytes,
5654 capacity,
5655 });
5656 }
5657 mem.map_with_byte_size(total_bytes, access)
5658 }
5659 #[cfg(unix)]
5660 TensorStorage::Shm(shm) => {
5661 let capacity = tensor.storage.capacity_bytes();
5662 if total_bytes > capacity {
5663 return Err(Error::InsufficientCapacity {
5664 needed: total_bytes,
5665 capacity,
5666 });
5667 }
5668 shm.map_with_byte_size(total_bytes, access)
5669 }
5670 // macOS/iOS: `TensorStorage::Dma` is the IOSurface. The lock yields
5671 // the full surface base address, and the row pitch
5672 // (`IOSurfaceGetBytesPerRow`) is known from the API for both
5673 // self-allocated and imported surfaces — unlike a foreign
5674 // DMA-BUF — so a strided CPU view is sound and zero-copy.
5675 #[cfg(any(target_os = "macos", target_os = "ios"))]
5676 TensorStorage::Dma(io) => {
5677 // A sub-view's window is `buf_size − view_offset`; the strided
5678 // span must fit the window, not the whole surface.
5679 let available = io.buf_size.saturating_sub(io.view_offset);
5680 if total_bytes > available {
5681 return Err(Error::InsufficientCapacity {
5682 needed: total_bytes,
5683 capacity: available,
5684 });
5685 }
5686 io.map_with_byte_size(total_bytes, access)
5687 }
5688 // Android: `TensorStorage::Dma` is the AHardwareBuffer. The lock
5689 // yields the full buffer base address, and the row pitch is
5690 // known from the allocator-filled descriptor — so a strided CPU
5691 // view is sound and zero-copy, same as IOSurface.
5692 #[cfg(target_os = "android")]
5693 TensorStorage::Dma(ahb) => {
5694 // A sub-view's window is `buf_size − view_offset`; the strided
5695 // span must fit the window, not the whole buffer.
5696 let available = ahb.buf_size.saturating_sub(ahb.view_offset);
5697 if total_bytes > available {
5698 return Err(Error::InsufficientCapacity {
5699 needed: total_bytes,
5700 capacity: available,
5701 });
5702 }
5703 ahb.map_with_byte_size(total_bytes, access)
5704 }
5705 // Windows: `TensorStorage::Dma` is the D3D11 texture, whose CPU side
5706 // is the staging copy at the driver's pitch -- the very pitch
5707 // `Tensor::image` recorded as this tensor's row stride. Its
5708 // `map_with_byte_size` bounds-checks the span against this tensor's
5709 // own window, as the IOSurface arm above does. This is also the only
5710 // backing with a real non-blocking map, so it is the only arm that
5711 // reads `_non_blocking`.
5712 #[cfg(target_os = "windows")]
5713 TensorStorage::Dma(tex) => {
5714 if _non_blocking {
5715 tex.try_map_with_byte_size(total_bytes, access)
5716 } else {
5717 tex.map_with_byte_size(total_bytes, access)
5718 }
5719 }
5720 TensorStorage::Pbo(pbo) => {
5721 // PBO: the GPU-side allocation may have a padded row stride
5722 // (e.g. 64-byte aligned). Expose the full padded buffer so a
5723 // CPU producer (JPEG decoder) and a strided convert source
5724 // can iterate rows via `effective_row_stride()` without
5725 // running past the slice — the logical `pbo.map()` view would
5726 // stop after `shape.product()` and lose bytes past row 0.
5727 // A sub-view's window is `capacity − view_offset`.
5728 let available = pbo.capacity_bytes().saturating_sub(pbo.view_offset);
5729 if total_bytes > available {
5730 return Err(Error::InsufficientCapacity {
5731 needed: total_bytes,
5732 capacity: available,
5733 });
5734 }
5735 pbo.map_with_byte_size(total_bytes, access)
5736 }
5737 // Every `TensorStorage` variant this build can construct is
5738 // matched explicitly above (Mem/Shm/Dma/Pbo, each cfg-gated
5739 // to the platforms it exists on) -- this catch-all exists
5740 // only so the match stays exhaustive across every
5741 // OS/cfg combination without a matching arm on every one of
5742 // them, not because any variant is meant to fall through to
5743 // it. `#[allow(unreachable_patterns)]` because on any given
5744 // platform it genuinely is unreachable.
5745 #[allow(unreachable_patterns)]
5746 _ => Err(Error::InvalidOperation(
5747 "CPU mapping of strided tensors is supported only for HAL-allocated \
5748 Mem/Shm (any platform), DMA (Linux, self-allocated or imported), \
5749 IOSurface (macOS), AHardwareBuffer (Android), D3D11 textures \
5750 (Windows), and PBO"
5751 .into(),
5752 )),
5753 }
5754}
5755
5756#[cfg(feature = "static")]
5757impl<T> Tensor<T>
5758where
5759 T: Num + Clone + fmt::Debug + Send + Sync,
5760{
5761 /// The shared body of [`TensorTrait::map_with`] and
5762 /// [`TensorTrait::try_map_with`]: the direction check, the
5763 /// declared-vs-requested telemetry, the strided and plane-offset routing,
5764 /// and finally the backing's own map. `non_blocking` chooses which of the
5765 /// backing's two maps that last step takes; everything before it is the
5766 /// same either way.
5767 fn map_impl<'a>(
5768 &self,
5769 access: CpuAccess,
5770 non_blocking: bool,
5771 ) -> Result<crate::view::HostView<'a, T>>
5772 where
5773 T: 'a,
5774 {
5775 let _span = tracing::trace_span!(
5776 "tensor.map",
5777 memory = ?self.storage.memory(),
5778 ?access,
5779 non_blocking,
5780 )
5781 .entered();
5782 if access == CpuAccess::None {
5783 return Err(Error::InvalidArgument(
5784 "map_with(CpuAccess::None) is not a mappable direction — use \
5785 map_read()/map_write()/map_mut()"
5786 .into(),
5787 ));
5788 }
5789 // Declared-vs-requested telemetry (all platforms): mapping beyond
5790 // the allocation-time declaration is best-effort — tolerated where
5791 // the backing is CPU-mappable regardless (Mem/Shm/dma-buf/
5792 // IOSurface), refused by the Android backend for CpuAccess::None
5793 // buffers — but always loud and counted, never silent.
5794 if !self.cpu_access.covers(access) {
5795 note_unplanned_cpu_access(
5796 self.buffer_identity().id(),
5797 &format!("{:?}", self.storage.memory()),
5798 "map access exceeds the declared CpuAccess",
5799 );
5800 }
5801 // CPU mapping of a strided tensor exposes the full padded buffer
5802 // (`row_stride × rows`) so callers can iterate rows via
5803 // `effective_row_stride()` without running past the slice. This is sound
5804 // only when the HAL owns and can size-check the allocation:
5805 //
5806 // * Self-allocated Mem / Shm tensors (any platform) — the backing
5807 // `Vec` / shm segment is sized by `capacity_bytes()`, checked here.
5808 // * Self-allocated DMA tensors (Linux) — pitch padding from
5809 // `image_with_stride()`; checked against the DMA-BUF `buf_size`.
5810 //
5811 // * Self-allocated PBO tensors (any platform with GL) — the GL buffer
5812 // is sized by `capacity_bytes()` and may carry 64-byte row padding;
5813 // the JPEG decoder mmaps it and convert() reads it, both iterating
5814 // by `row_stride`. Checked against the PBO capacity below.
5815 //
5816 // A strided **imported** DMA-BUF is accepted here too (Task 10b): the
5817 // `stride × rows <= buf_size` check just below is exactly what a
5818 // self-allocated Dma tensor relies on for its own soundness, and
5819 // `buf_size` for an imported fd is `fstat`-derived when the kernel
5820 // reports a usable size (`stat.st_size > 0` and at least the
5821 // declared shape's logical size; `DmaTensor::from_fd`) -- kernel
5822 // truth, not a caller's claim, in that case. The rare fallback (a
5823 // kernel that reports `st_size <= 0` for a DMA-BUF fd) uses the
5824 // declared logical size instead; that is not a hole here either, since
5825 // an out-of-range mmap against the real dma-buf is still rejected by
5826 // the kernel regardless of what this crate believed `buf_size` was.
5827 // So the check is equally load-bearing regardless of who allocated
5828 // the buffer. What
5829 // changed is only *which* callers can now reach this path: the
5830 // cross-package capsule protocol reconstructs an imported alias of a
5831 // HAL-owned Dma tensor and then calls `configure_image`, which
5832 // computes its stride with the exact same alignment rule a
5833 // self-allocated `Tensor::image()` used — so the value being
5834 // bounds-checked here is HAL-computed and format/width-derived, not
5835 // read from an external driver. Rejecting it unconditionally used to
5836 // make every foreign-imported GPU-backed destination fail any CPU
5837 // write path (the JPEG decoder's MCU writer among them) even when the
5838 // write was provably in-bounds. `set_row_stride()` was already
5839 // "trust the caller, bounds-check the map" for every *other* backing
5840 // (Mem/Shm/self-allocated Dma/PBO) — singling out imported Dma out
5841 // for a stricter standard was the inconsistency, not the trust model
5842 // itself. IOSurface keeps its own always-allowed strided path
5843 // unconditionally (see its arm below): a live `IOSurfaceRef` reports
5844 // its real `bytesPerRow` from the OS regardless of who allocated it,
5845 // so there was never anything to distrust there.
5846 if let Some(stride) = self.row_stride {
5847 return map_strided_host_view(self, stride, access, non_blocking);
5848 }
5849 // Offset tensors are supported for storages that apply the offset
5850 // inside their own `map()`: DMA (`DmaMap`/IOSurface adjust the mapped
5851 // base), Mem (`MemMap` adjusts the slice base), Shm (`ShmMap` adjusts
5852 // the slice base), and PBO (the staged copy starts at the offset). Every
5853 // self-allocated backing now carries a sub-region concept via `view`, so
5854 // a non-zero offset is honoured rather than rejected.
5855 if self.plane_offset.is_some_and(|o| o > 0) {
5856 let supported = matches!(self.storage, TensorStorage::Mem(_) | TensorStorage::Pbo(_));
5857 // macOS `Dma` is the IOSurface; Linux `Dma` is the DMA-BUF; Android
5858 // `Dma` is the AHardwareBuffer, Windows `Dma` is the D3D11 texture
5859 // (whose pin and map both add the storage's own `view_offset`, the
5860 // same way IOSurface does) — all apply the offset in their map.
5861 // (`Dma` is the same variant name on each, hence one `cfg(any(...))`
5862 // arm rather than four.)
5863 #[cfg(any(
5864 target_os = "linux",
5865 target_os = "macos",
5866 target_os = "ios",
5867 target_os = "android",
5868 target_os = "windows"
5869 ))]
5870 let supported = supported || matches!(self.storage, TensorStorage::Dma(_));
5871 #[cfg(unix)]
5872 let supported = supported || matches!(self.storage, TensorStorage::Shm(_));
5873 if !supported {
5874 return Err(Error::InvalidOperation(
5875 "plane offset only supported for DMA, Mem, Shm, and PBO tensors".into(),
5876 ));
5877 }
5878 }
5879 if non_blocking {
5880 self.storage.try_map_with(access)
5881 } else {
5882 self.storage.map_with(access)
5883 }
5884 }
5885}
5886
5887#[cfg(feature = "static")]
5888impl<T> TensorTrait<T> for Tensor<T>
5889where
5890 T: Num + Clone + fmt::Debug + Send + Sync,
5891{
5892 fn new(shape: &[usize], name: Option<&str>) -> Result<Self>
5893 where
5894 Self: Sized,
5895 {
5896 Self::new(shape, None, name)
5897 }
5898
5899 #[cfg(unix)]
5900 fn from_fd(fd: std::os::fd::OwnedFd, shape: &[usize], name: Option<&str>) -> Result<Self>
5901 where
5902 Self: Sized,
5903 {
5904 #[cfg_attr(not(target_os = "linux"), allow(unused_mut))]
5905 let mut t = Self::wrap(TensorStorage::from_fd(fd, shape, name)?);
5906 // Best-effort CUDA external memory import for DMA-backed tensors.
5907 // RUNTIME-UNVALIDATED: see try_init_dma_cuda().
5908 #[cfg(target_os = "linux")]
5909 t.try_init_dma_cuda();
5910 Ok(t)
5911 }
5912
5913 #[cfg(unix)]
5914 fn clone_fd(&self) -> Result<std::os::fd::OwnedFd> {
5915 self.storage.clone_fd()
5916 }
5917
5918 fn memory(&self) -> TensorMemory {
5919 self.storage.memory()
5920 }
5921
5922 fn name(&self) -> String {
5923 self.storage.name()
5924 }
5925
5926 fn shape(&self) -> &[usize] {
5927 self.storage.shape()
5928 }
5929
5930 /// Bytes available in the underlying allocation, from this tensor's
5931 /// window start -- may exceed [`TensorTrait::size`] when the backing
5932 /// carries spare room beyond the current logical shape (a pitch-aligned
5933 /// image, or a pool tensor sized for a larger image than it holds
5934 /// today). See [`crate::protocol::TensorDesc::capacity`], the
5935 /// cross-package descriptor field this backs.
5936 fn capacity_bytes(&self) -> usize {
5937 self.storage.capacity_bytes()
5938 }
5939
5940 fn reshape(&mut self, shape: &[usize]) -> Result<()> {
5941 if self.chroma.is_some() {
5942 return Err(Error::InvalidOperation(
5943 "cannot reshape a multiplane tensor — decompose planes first".into(),
5944 ));
5945 }
5946 self.storage.reshape(shape)?;
5947 self.format = None;
5948 self.row_stride = None;
5949 self.plane_offset = None;
5950 match self.storage {
5951 TensorStorage::Mem(ref mut m) => m.set_offset(0),
5952 #[cfg(target_os = "linux")]
5953 TensorStorage::Dma(ref mut dma) => dma.mmap_offset = 0,
5954 // No-op: `IoSurfaceTensor::reshape` (iosurface.rs) zeroes its
5955 // own `view_offset` before this match runs.
5956 #[cfg(any(target_os = "macos", target_os = "ios"))]
5957 TensorStorage::Dma(_) => {}
5958 // `D3d11TextureTensor::reshape` does not self-zero, so the clear
5959 // that pairs with `set_plane_offset`'s D3D11 arm lives here.
5960 #[cfg(target_os = "windows")]
5961 TensorStorage::Dma(ref mut tex) => tex.set_view_offset(0),
5962 // No-op: `AHardwareBufferTensor::reshape` (ahardwarebuffer.rs)
5963 // zeroes its own `view_offset` before this match runs, the same
5964 // as `IoSurfaceTensor`. Unlike `set_plane_offset`'s Android arm,
5965 // this one is a genuine self-zeroing no-op, not a still-
5966 // incomplete gap: `view_offset` never survives a reshape here
5967 // regardless of whether `set_plane_offset` can restore it.
5968 #[cfg(target_os = "android")]
5969 TensorStorage::Dma(_) => {}
5970 // No-op: `PboTensor::reshape` (pbo.rs) zeroes its own
5971 // `view_offset` before this match runs, the same as
5972 // `IoSurfaceTensor`/`AHardwareBufferTensor` above -- which is
5973 // why, unlike `set_format`'s clear block, this match never
5974 // needed a `Pbo` arm that writes anything.
5975 TensorStorage::Pbo(_) => {}
5976 // Unlike `PboTensor::reshape`, `ShmTensor::reshape` (shm.rs)
5977 // only revalidates the element count and swaps `self.shape` --
5978 // it never touches `self.offset` -- so this clear site needs
5979 // the arm `set_format`'s does. `cfg(unix)` mirrors the `Shm`
5980 // variant's own gate, not a per-OS multiplex.
5981 #[cfg(unix)]
5982 TensorStorage::Shm(ref mut s) => s.set_offset(0),
5983 }
5984 Ok(())
5985 }
5986
5987 /// Set the logical shape to any shape whose bytes fit the allocation,
5988 /// without [`reshape`](TensorTrait::reshape)'s equal-count constraint.
5989 ///
5990 /// Overridden rather than inherited, which it was not for most of this
5991 /// crate's life. The trait's default body is `self.reshape(shape)`, so
5992 /// a caller reaching this through `Tensor<T>` silently got the strict
5993 /// rule under a name that promises the opposite -- on both backends,
5994 /// which is why G13 cannot see it.
5995 ///
5996 /// The capacity-aware implementations were never dead: every storage
5997 /// type has one, and [`Tensor::configure_image`] reaches them by
5998 /// calling `self.storage.set_logical_shape(..)` directly. That is the
5999 /// pool-reuse path -- a decode into an oversized reusable buffer
6000 /// reconfiguring it to a smaller image without reallocating -- and it
6001 /// worked precisely because it bypassed this method. This forwards the
6002 /// same way, so the two agree.
6003 ///
6004 /// A narrower image inside a pitched backing (an IOSurface, an
6005 /// AHardwareBuffer, a D3D11 texture) still advances rows by the backing's
6006 /// pitch, so that pitch is recorded as the row stride whenever it exceeds
6007 /// the new shape's natural stride -- the same adoption `configure_image`
6008 /// makes. Without it a texture narrowed from 64 to 16 texels wide reports
6009 /// a 64-byte stride over rows that are 256 bytes apart. A single-row
6010 /// shape records its own tight stride instead, by [`view`](Self::view)'s
6011 /// rule: a strided `map()` spans `stride * rows`, and one pitched row
6012 /// runs past the backing when the window sits in the last row.
6013 fn set_logical_shape(&mut self, shape: &[usize]) -> Result<()> {
6014 self.storage.set_logical_shape(shape)?;
6015 // Only a backing with a pitch of its own, and only while the shape
6016 // still has the rank the format's row is measured on: a caller may
6017 // flatten an image to one dimension and leave the format behind.
6018 let Some(pitch) = self.storage.backing_row_stride() else {
6019 return Ok(());
6020 };
6021 let Some((rank, rows)) = self.format.map(|f| match f.layout() {
6022 PixelLayout::Packed => (3, shape.first().copied()),
6023 PixelLayout::Planar => (3, shape.get(1).copied()),
6024 PixelLayout::SemiPlanar => (2, shape.first().copied()),
6025 }) else {
6026 return Ok(());
6027 };
6028 if rank != shape.len() {
6029 return Ok(());
6030 }
6031 // The natural stride of the new shape, not of the stride recorded for
6032 // the old one.
6033 let prior = self.row_stride.take();
6034 let Some(natural) = self.effective_row_stride() else {
6035 self.row_stride = prior;
6036 return Ok(());
6037 };
6038 if rows == Some(1) {
6039 self.set_row_stride_unchecked(natural);
6040 } else if pitch > natural {
6041 self.set_row_stride_unchecked(pitch);
6042 } else {
6043 self.row_stride = prior;
6044 }
6045 Ok(())
6046 }
6047
6048 fn map_with<'a>(&self, access: CpuAccess) -> Result<crate::view::HostView<'a, T>>
6049 where
6050 T: 'a,
6051 {
6052 self.map_impl(access, false)
6053 }
6054
6055 fn try_map_with<'a>(&self, access: CpuAccess) -> Result<crate::view::HostView<'a, T>>
6056 where
6057 T: 'a,
6058 {
6059 self.map_impl(access, true)
6060 }
6061
6062 fn buffer_identity(&self) -> &BufferIdentity {
6063 self.storage.buffer_identity()
6064 }
6065}
6066
6067// ============================================================================
6068// Platform availability helpers
6069// ============================================================================
6070
6071/// Cached result of the Linux DMA-BUF availability probe.
6072#[cfg(target_os = "linux")]
6073static DMA_AVAILABLE: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
6074/// Cached result of the macOS/iOS IOSurface availability probe.
6075#[cfg(any(target_os = "macos", target_os = "ios"))]
6076static IOSURFACE_AVAILABLE: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
6077
6078/// Check if Linux DMA-BUF allocation is available on this system.
6079///
6080/// Linux-specific availability check (typically requires `/dev/dma_heap`
6081/// access — running as root or membership in a video/render group). For
6082/// portable code that wants "any zero-copy GPU buffer", use
6083/// [`is_gpu_buffer_available`] which also covers IOSurface on macOS.
6084///
6085/// This function caches its result after the first call.
6086#[cfg(target_os = "linux")]
6087pub fn is_dma_available() -> bool {
6088 *DMA_AVAILABLE
6089 .get_or_init(|| Tensor::<u8>::new(&[64], Some(TensorMemory::DmaBuf), None).is_ok())
6090}
6091
6092/// Always returns `false` on non-Linux platforms.
6093#[cfg(not(target_os = "linux"))]
6094pub fn is_dma_available() -> bool {
6095 false
6096}
6097
6098/// Check if macOS/iOS IOSurface allocation is available on this system.
6099///
6100/// IOSurface is part of the macOS/iOS OS and is essentially always present;
6101/// this probe catches degraded scenarios such as memory pressure or
6102/// sandboxed contexts where `IOSurfaceCreate` fails. The result is
6103/// cached after the first call.
6104///
6105/// Always returns `false` on non-Apple platforms.
6106#[cfg(any(target_os = "macos", target_os = "ios"))]
6107pub fn is_iosurface_available() -> bool {
6108 *IOSURFACE_AVAILABLE.get_or_init(|| {
6109 // Probe via the same Dma path — on macOS/iOS this routes through
6110 // IoSurfaceTensor::new.
6111 Tensor::<u8>::new(&[64], Some(TensorMemory::DmaBuf), None).is_ok()
6112 })
6113}
6114
6115#[cfg(not(any(target_os = "macos", target_os = "ios")))]
6116pub fn is_iosurface_available() -> bool {
6117 false
6118}
6119
6120/// Cached result of the Android AHardwareBuffer availability probe.
6121#[cfg(target_os = "android")]
6122static AHARDWAREBUFFER_AVAILABLE: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
6123
6124/// Check if Android AHardwareBuffer allocation is available on this system.
6125///
6126/// AHardwareBuffer is part of the Android OS (public NDK ABI since API
6127/// 26) and is essentially always present; this probe catches degraded
6128/// scenarios such as memory pressure or gralloc failures. The result is
6129/// cached after the first call.
6130#[cfg(target_os = "android")]
6131pub fn is_ahardwarebuffer_available() -> bool {
6132 *AHARDWAREBUFFER_AVAILABLE.get_or_init(|| {
6133 // Probe via the same Dma path — on Android this routes through
6134 // AHardwareBufferTensor::new.
6135 Tensor::<u8>::new(&[64], Some(TensorMemory::DmaBuf), None).is_ok()
6136 })
6137}
6138
6139/// Always returns `false` on non-Android platforms.
6140#[cfg(not(target_os = "android"))]
6141pub fn is_ahardwarebuffer_available() -> bool {
6142 false
6143}
6144
6145/// Portable probe for the platform's native zero-copy GPU buffer allocator
6146/// (DMA-BUF on Linux, IOSurface on macOS/iOS, AHardwareBuffer on Android,
6147/// D3D11 textures on Windows). Returns `false` on platforms with no
6148/// equivalent. Use this when writing cross-platform code that cares whether
6149/// the `Dma` tensor variant will work, not which underlying mechanism is used.
6150///
6151/// On Windows the `Dma` variant is image-formatted: `true` here means
6152/// [`Tensor::image`] can serve `TensorMemory::DmaBuf`, not [`Tensor::new`].
6153///
6154/// # Ordering on Windows
6155///
6156/// This probe **creates the process D3D11 device**, because the only honest
6157/// answer to "can this host allocate one" is to try. A host that must supply
6158/// its own device calls `d3d11::use_external_device` **before** this probe or
6159/// any allocation: the device is cached on first creation and installing one
6160/// afterwards is refused.
6161pub fn is_gpu_buffer_available() -> bool {
6162 #[cfg(target_os = "linux")]
6163 {
6164 is_dma_available()
6165 }
6166 #[cfg(any(target_os = "macos", target_os = "ios"))]
6167 {
6168 is_iosurface_available()
6169 }
6170 #[cfg(target_os = "android")]
6171 {
6172 is_ahardwarebuffer_available()
6173 }
6174 // No cached probe: `d3d11::device()` is itself a `OnceLock` over the
6175 // process device, so asking twice costs one atomic load.
6176 #[cfg(target_os = "windows")]
6177 {
6178 crate::d3d11::device().is_ok()
6179 }
6180 #[cfg(not(any(
6181 target_os = "linux",
6182 target_os = "macos",
6183 target_os = "ios",
6184 target_os = "android",
6185 target_os = "windows"
6186 )))]
6187 {
6188 false
6189 }
6190}
6191
6192/// Check if POSIX shared memory allocation is available on this system.
6193///
6194/// Returns `true` on Unix systems (Linux, macOS, BSD) where POSIX shared memory
6195/// is supported. Always returns `false` on non-Unix platforms (Windows).
6196///
6197/// This function caches its result after the first call for efficiency.
6198#[cfg(unix)]
6199static SHM_AVAILABLE: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
6200
6201/// Check if POSIX shared memory allocation is available on this system.
6202#[cfg(unix)]
6203pub fn is_shm_available() -> bool {
6204 *SHM_AVAILABLE.get_or_init(|| Tensor::<u8>::new(&[64], Some(TensorMemory::Shm), None).is_ok())
6205}
6206
6207/// Check if POSIX shared memory allocation is available on this system.
6208///
6209/// Always returns `false` on non-Unix platforms since POSIX SHM is Unix-specific.
6210#[cfg(not(unix))]
6211pub fn is_shm_available() -> bool {
6212 false
6213}
6214
6215#[cfg(test)]
6216mod dtype_tests {
6217 use super::*;
6218
6219 #[test]
6220 fn dtype_size() {
6221 assert_eq!(DType::U8.size(), 1);
6222 assert_eq!(DType::I8.size(), 1);
6223 assert_eq!(DType::U16.size(), 2);
6224 assert_eq!(DType::I16.size(), 2);
6225 assert_eq!(DType::U32.size(), 4);
6226 assert_eq!(DType::I32.size(), 4);
6227 assert_eq!(DType::U64.size(), 8);
6228 assert_eq!(DType::I64.size(), 8);
6229 assert_eq!(DType::F16.size(), 2);
6230 assert_eq!(DType::F32.size(), 4);
6231 assert_eq!(DType::F64.size(), 8);
6232 }
6233
6234 #[test]
6235 fn dtype_name() {
6236 assert_eq!(DType::U8.name(), "u8");
6237 assert_eq!(DType::F16.name(), "f16");
6238 assert_eq!(DType::F32.name(), "f32");
6239 }
6240
6241 #[test]
6242 fn dtype_serde_roundtrip() {
6243 use serde_json;
6244 let dt = DType::F16;
6245 let json = serde_json::to_string(&dt).unwrap();
6246 let back: DType = serde_json::from_str(&json).unwrap();
6247 assert_eq!(dt, back);
6248 }
6249}
6250
6251#[cfg(test)]
6252mod image_tests {
6253 use super::*;
6254
6255 #[test]
6256 fn image_dims_from_shape_inverts_allocation_shape() {
6257 for (format, w, h) in [
6258 (PixelFormat::Rgba, 640usize, 480usize),
6259 (PixelFormat::Rgb, 641, 3),
6260 (PixelFormat::Grey, 32, 32),
6261 (PixelFormat::PlanarRgb, 640, 480),
6262 (PixelFormat::PlanarRgba, 8, 5),
6263 (PixelFormat::Nv12, 640, 481),
6264 (PixelFormat::Nv16, 33, 17),
6265 (PixelFormat::Nv24, 16, 9),
6266 ] {
6267 let shape = format.allocation_shape(w, h).unwrap();
6268 assert_eq!(
6269 image_dims_from_shape(format, &shape).unwrap(),
6270 (w, h),
6271 "{format:?} {w}x{h} via {shape:?}"
6272 );
6273 }
6274 }
6275
6276 #[test]
6277 fn image_dims_from_shape_rejects_the_wrong_rank() {
6278 // Packed and planar want rank 3, semi-planar rank 2.
6279 for (format, shape) in [
6280 (PixelFormat::Rgba, vec![480usize, 640usize]),
6281 (PixelFormat::PlanarRgb, vec![3, 480]),
6282 (PixelFormat::Nv12, vec![720, 640, 1]),
6283 ] {
6284 let err = image_dims_from_shape(format, &shape).unwrap_err();
6285 assert!(
6286 matches!(err, Error::InvalidShape(ref m) if m.contains("rank")),
6287 "{format:?} {shape:?}: {err}"
6288 );
6289 }
6290 }
6291
6292 #[test]
6293 fn image_dims_from_shape_rejects_an_unproduced_combined_height() {
6294 // NV12 maps 480 -> 720 and 481 -> 722; nothing maps to 721, so the
6295 // shape is rejected rather than rounded to 480.
6296 let err = image_dims_from_shape(PixelFormat::Nv12, &[721, 640]).unwrap_err();
6297 assert!(
6298 matches!(err, Error::InvalidShape(ref m) if m.contains("combined plane height")),
6299 "{err}"
6300 );
6301 }
6302
6303 #[test]
6304 fn allocation_shape_per_layout() {
6305 assert_eq!(
6306 PixelFormat::Rgb.allocation_shape(640, 480),
6307 Some(vec![480, 640, 3])
6308 );
6309 assert_eq!(
6310 PixelFormat::Grey.allocation_shape(640, 480),
6311 Some(vec![480, 640, 1])
6312 );
6313 assert_eq!(
6314 PixelFormat::Nv12.allocation_shape(640, 480),
6315 Some(vec![720, 640])
6316 );
6317 // Odd height: combined-plane height is `481 + ceil(481/2)` = 481 + 241
6318 // = 722 rows. Logical height is recovered as `722 * 2 / 3` = 481.
6319 assert_eq!(
6320 PixelFormat::Nv12.allocation_shape(640, 481),
6321 Some(vec![722, 640])
6322 );
6323 // Odd width: shape carries the LOGICAL width (641).
6324 // The 64-aligned stride (>= 642) is stored separately on the Tensor.
6325 assert_eq!(
6326 PixelFormat::Nv12.allocation_shape(641, 480),
6327 Some(vec![720, 641])
6328 );
6329 // NV16 odd width: same — logical width in shape, stride separate.
6330 assert_eq!(
6331 PixelFormat::Nv16.allocation_shape(641, 480),
6332 Some(vec![960, 641])
6333 );
6334 assert_eq!(
6335 PixelFormat::PlanarRgb.allocation_shape(640, 480),
6336 Some(vec![3, 480, 640])
6337 );
6338 assert_eq!(
6339 PixelFormat::Nv16.allocation_shape(640, 480),
6340 Some(vec![960, 640])
6341 );
6342 }
6343
6344 #[test]
6345 fn raw_tensor_has_no_format() {
6346 let t = Tensor::<u8>::new(&[480, 640, 3], None, None).unwrap();
6347 assert!(t.format().is_none());
6348 assert!(t.width().is_none());
6349 assert!(t.height().is_none());
6350 assert!(!t.is_multiplane());
6351 assert!(t.chroma().is_none());
6352 }
6353
6354 #[test]
6355 fn image_tensor_packed() {
6356 let t = Tensor::<u8>::image(
6357 640,
6358 480,
6359 PixelFormat::Rgba,
6360 None,
6361 crate::CpuAccess::ReadWrite,
6362 )
6363 .unwrap();
6364 assert_eq!(t.format(), Some(PixelFormat::Rgba));
6365 assert_eq!(t.width(), Some(640));
6366 assert_eq!(t.height(), Some(480));
6367 assert_eq!(t.shape(), &[480, 640, 4]);
6368 assert!(!t.is_multiplane());
6369 }
6370
6371 #[test]
6372 fn image_tensor_planar() {
6373 let t = Tensor::<u8>::image(
6374 640,
6375 480,
6376 PixelFormat::PlanarRgb,
6377 None,
6378 crate::CpuAccess::ReadWrite,
6379 )
6380 .unwrap();
6381 assert_eq!(t.format(), Some(PixelFormat::PlanarRgb));
6382 assert_eq!(t.width(), Some(640));
6383 assert_eq!(t.height(), Some(480));
6384 assert_eq!(t.shape(), &[3, 480, 640]);
6385 }
6386
6387 #[test]
6388 #[cfg(target_os = "macos")]
6389 fn image_tensor_dma_non_aligned_packed_width_pads_zero_copy() {
6390 // Declares this test as an fd-opener; see FD_LOCK.
6391 let _lock = crate::tests::fd_lock_shared();
6392 // RGBA u8 at width=4 → 4*4 = 16 bytes/row, not 64-byte aligned. RGBA has
6393 // a real IOSurface FourCC, so an explicit `Some(TensorMemory::DmaBuf)`
6394 // request now allocates a padded image IOSurface (64-aligned
6395 // `bytes_per_row`) and records the stride — a fully zero-copy buffer GL
6396 // can bind and the CPU can map via the strided path. (Previously this
6397 // failed loudly to avoid an 'L008' byte-bag downgrade; with a real
6398 // FourCC surface that concern no longer applies.)
6399 let t = Tensor::<u8>::image(
6400 4,
6401 4,
6402 PixelFormat::Rgba,
6403 Some(TensorMemory::DmaBuf),
6404 crate::CpuAccess::ReadWrite,
6405 )
6406 .expect("padded RGBA IOSurface should allocate");
6407 assert_eq!(t.format(), Some(PixelFormat::Rgba));
6408 assert_eq!(t.width(), Some(4));
6409 assert_eq!(t.height(), Some(4));
6410 let stride = t.effective_row_stride().expect("stride");
6411 assert_eq!(stride % 64, 0, "padded to 64-byte row alignment");
6412 assert!(stride >= 16);
6413 // A CPU map exposes the full padded surface for strided iteration.
6414 let m = t.map().expect("strided IOSurface map");
6415 assert_eq!(m.as_slice().len(), stride * 4);
6416 }
6417
6418 /// `per_pixel_bytes` that doesn't divide 64 evenly (e.g. RGB u8 with
6419 /// 3 B/pixel) makes a "Pad width to N" suggestion structurally
6420 /// impossible — there is no integer width whose `width * 3` is a
6421 /// multiple of 64. The error must still fire (no silent SHM
6422 /// fallback for explicit-DMA requests) and must spell out the
6423 /// alignment requirement; it just omits the misleading "pad to N"
6424 /// hint instead of printing a number whose row pitch still won't
6425 /// align.
6426 #[test]
6427 #[cfg(target_os = "macos")]
6428 fn image_tensor_dma_rejects_indivisible_pixel_pitch_without_pad_hint() {
6429 // Declares this test as an fd-opener; see FD_LOCK.
6430 let _lock = crate::tests::fd_lock_shared();
6431 // Width=10 RGB f32 → 120 B/row, not 64-byte aligned, and (Rgb,
6432 // F32) has no IOSurface mapping so the padded-stride tolerance
6433 // does not apply. The next 64-multiple (128 B) isn't an integer
6434 // multiple of 12 B/pixel, so the "pad width to N" hint can't
6435 // produce a valid number and must be omitted. (Rgb u8 used to be
6436 // this test's subject but now has a real RGBA8888 mapping with
6437 // padded-stride tolerance — see the test below.)
6438 let err = Tensor::<f32>::image(
6439 10,
6440 10,
6441 PixelFormat::Rgb,
6442 Some(TensorMemory::DmaBuf),
6443 crate::CpuAccess::ReadWrite,
6444 )
6445 .expect_err("RGB f32 with 12 B/pixel and non-aligned width must be rejected");
6446 match err {
6447 Error::InvalidArgument(msg) => {
6448 assert!(
6449 msg.contains("64-byte aligned"),
6450 "error must still name the alignment requirement: {msg}"
6451 );
6452 assert!(
6453 !msg.contains("Pad width"),
6454 "indivisible per-pixel pitch makes a width suggestion impossible; \
6455 hint must be omitted, got: {msg}"
6456 );
6457 assert!(
6458 msg.contains("memory=None") && msg.contains("TensorMemory::Mem"),
6459 "error must still list the always-applicable alternatives: {msg}"
6460 );
6461 }
6462 other => panic!("expected InvalidArgument, got {other:?}"),
6463 }
6464 }
6465
6466 #[test]
6467 #[cfg(target_os = "macos")]
6468 fn image_tensor_dma_packed_rgb_u8_contract() {
6469 // Declares this test as an fd-opener; see FD_LOCK.
6470 let _lock = crate::tests::fd_lock_shared();
6471 // Packed RGB u8 @Dma is a designed RGBA8888 mapping at
6472 // (W*3/4, H) — the INT8 NPU input layout, shared with Android.
6473 // width%4 != 0 cannot form whole texels → loud InvalidArgument…
6474 let err = Tensor::<u8>::image(
6475 10,
6476 10,
6477 PixelFormat::Rgb,
6478 Some(TensorMemory::DmaBuf),
6479 crate::CpuAccess::ReadWrite,
6480 )
6481 .expect_err("Rgb u8 width%4!=0 must be rejected");
6482 assert!(
6483 matches!(&err, Error::InvalidArgument(m) if m.contains("width%4==0")),
6484 "got {err:?}"
6485 );
6486 // …width%4 == 0 with a non-64-aligned pitch allocates PADDED
6487 // (36 B rows → 64 B surface pitch, recorded on the tensor)…
6488 let t = Tensor::<u8>::image(
6489 12,
6490 4,
6491 PixelFormat::Rgb,
6492 Some(TensorMemory::DmaBuf),
6493 crate::CpuAccess::ReadWrite,
6494 )
6495 .expect("width 12 Rgb u8 must allocate padded");
6496 assert_eq!(t.memory(), TensorMemory::DmaBuf);
6497 assert!(
6498 t.row_stride()
6499 .is_some_and(|s| s >= 64 && s.is_multiple_of(64)),
6500 "padded pitch must be recorded: {:?}",
6501 t.row_stride()
6502 );
6503 // …and the aligned model-input width stays flat (640*3 = 1920 is
6504 // 64-aligned → no recorded stride, the buffer IS [H, W, 3]).
6505 let t = Tensor::<u8>::image(
6506 640,
6507 8,
6508 PixelFormat::Rgb,
6509 Some(TensorMemory::DmaBuf),
6510 crate::CpuAccess::ReadWrite,
6511 )
6512 .expect("width 640 Rgb u8 must allocate flat");
6513 assert_eq!(t.row_stride(), None);
6514 // I8 shares the layout (INT8 shader bias, not a format change).
6515 let t = Tensor::<i8>::image(
6516 640,
6517 8,
6518 PixelFormat::Rgb,
6519 Some(TensorMemory::DmaBuf),
6520 crate::CpuAccess::ReadWrite,
6521 )
6522 .expect("Rgb i8 shares the RGBA8888 mapping");
6523 assert_eq!(t.memory(), TensorMemory::DmaBuf);
6524 }
6525
6526 #[test]
6527 #[cfg(target_os = "macos")]
6528 fn image_tensor_dma_planar_f16_alignment() {
6529 // Declares this test as an fd-opener; see FD_LOCK.
6530 let _lock = crate::tests::fd_lock_shared();
6531 // PlanarRgb F16 uses single-channel row pitch (width * 2 bytes).
6532 // Width=16 → 32 bytes/row (not aligned); width=32 → 64 bytes/row (aligned).
6533 let err = Tensor::<half::f16>::image(
6534 16,
6535 16,
6536 PixelFormat::PlanarRgb,
6537 Some(TensorMemory::DmaBuf),
6538 crate::CpuAccess::ReadWrite,
6539 )
6540 .expect_err("width=16 PlanarRgb F16 is 32-byte row, must reject");
6541 assert!(matches!(err, Error::InvalidArgument(_)), "got {err:?}");
6542 // 32 wide should work.
6543 let t = Tensor::<half::f16>::image(
6544 32,
6545 8,
6546 PixelFormat::PlanarRgb,
6547 Some(TensorMemory::DmaBuf),
6548 crate::CpuAccess::ReadWrite,
6549 )
6550 .expect("width=32 PlanarRgb F16 is 64-byte row, must succeed");
6551 assert_eq!(t.format(), Some(PixelFormat::PlanarRgb));
6552 }
6553
6554 #[test]
6555 fn image_tensor_semi_planar_contiguous() {
6556 let t = Tensor::<u8>::image(
6557 640,
6558 480,
6559 PixelFormat::Nv12,
6560 None,
6561 crate::CpuAccess::ReadWrite,
6562 )
6563 .unwrap();
6564 assert_eq!(t.format(), Some(PixelFormat::Nv12));
6565 assert_eq!(t.width(), Some(640));
6566 assert_eq!(t.height(), Some(480));
6567 // NV12: H*3/2 = 720
6568 assert_eq!(t.shape(), &[720, 640]);
6569 assert!(!t.is_multiplane());
6570 }
6571
6572 #[test]
6573 #[cfg(target_os = "linux")]
6574 fn image_tensor_with_stride_preserves_logical_width() {
6575 // Declares this test as an fd-opener; see FD_LOCK.
6576 let _lock = crate::tests::fd_lock_shared();
6577 // Skip if DMA not available (e.g. sandboxed CI lacking dma_heap access).
6578 if !is_dma_available() {
6579 crate::test_support::report_skip("DMA heap not available");
6580 return;
6581 }
6582 // 3004×1688 RGBA8: natural pitch 12016, padded to 12032 (64-aligned).
6583 let stride = 12032;
6584 let t = Tensor::<u8>::image_with_stride(
6585 3004,
6586 1688,
6587 PixelFormat::Rgba,
6588 stride,
6589 Some(TensorMemory::DmaBuf),
6590 crate::CpuAccess::ReadWrite,
6591 )
6592 .unwrap();
6593 // Logical dimensions unchanged by padding — this is the contract.
6594 assert_eq!(t.width(), Some(3004));
6595 assert_eq!(t.height(), Some(1688));
6596 assert_eq!(t.shape(), &[1688, 3004, 4]);
6597 // Stride is carried separately and reports the padded pitch.
6598 assert_eq!(t.effective_row_stride(), Some(stride));
6599 // Buffer is sized to stride × height so the full padded layout fits,
6600 // and CPU map() works for self-allocated strided DMA tensors.
6601 use crate::TensorMapTrait;
6602 {
6603 let map = t.map().unwrap();
6604 assert!(
6605 map.as_slice().len() >= stride * 1688,
6606 "mapped buffer {} bytes < expected {}",
6607 map.as_slice().len(),
6608 stride * 1688
6609 );
6610 }
6611 // CPU write access works too — iterate rows using the padded stride,
6612 // touch only the active `width × bpp` region, verify it round-trips.
6613 {
6614 let mut map = t.map().unwrap();
6615 let slice = map.as_mut_slice();
6616 for y in 0..1688 {
6617 let row_start = y * stride;
6618 for x in 0..3004 {
6619 let p = row_start + x * 4;
6620 slice[p] = (y & 0xFF) as u8;
6621 slice[p + 1] = (x & 0xFF) as u8;
6622 slice[p + 2] = 0x42;
6623 slice[p + 3] = 0xFF;
6624 }
6625 }
6626 }
6627 {
6628 let map = t.map().unwrap();
6629 let slice = map.as_slice();
6630 // Sample a few pixels to confirm the round-trip.
6631 assert_eq!(slice[0], 0x00);
6632 assert_eq!(slice[1], 0x00);
6633 assert_eq!(slice[2], 0x42);
6634 assert_eq!(slice[3], 0xFF);
6635 let mid = 100 * stride + 50 * 4;
6636 assert_eq!(slice[mid], 100);
6637 assert_eq!(slice[mid + 1], 50);
6638 assert_eq!(slice[mid + 2], 0x42);
6639 }
6640 }
6641
6642 #[test]
6643 #[cfg(target_os = "linux")]
6644 fn image_tensor_foreign_strided_map_is_bounds_checked_not_refused() {
6645 // An imported (foreign) DMA-BUF with a row stride used to be refused
6646 // outright. That changed deliberately: the cross-package capsule
6647 // protocol reconstructs an imported alias of a HAL-owned Dma tensor,
6648 // and refusing every strided CPU map broke the flagship
6649 // create_image() -> decode_file_into() -> convert() pipeline on Linux.
6650 //
6651 // What replaced the refusal is a bounds check against `buf_size`,
6652 // which for an imported fd is `fstat`-derived kernel truth rather
6653 // than the caller's claim. So the property worth testing is no longer
6654 // "foreign strided maps fail" but "they succeed exactly when they
6655 // fit" -- which is what actually keeps the mapping sound.
6656 //
6657 // This test previously asserted the old refusal and went red on every
6658 // board with a dma-heap when the behaviour changed, unnoticed because
6659 // it needs both Linux and a dma-heap: it skips on macOS and on Orin
6660 // (nvmap, no CONFIG_DMABUF_HEAPS).
6661 if !is_dma_available() {
6662 crate::test_support::report_skip("DMA heap not available");
6663 return;
6664 }
6665 let backing =
6666 Tensor::<u8>::new(&[240 * 320 * 4], Some(TensorMemory::DmaBuf), None).unwrap();
6667
6668 // In bounds: the natural stride for the declared image fits the
6669 // allocation, so the map must succeed.
6670 {
6671 let fd = backing.clone_fd().unwrap();
6672 let shape = [240usize, 320, 4];
6673 let storage = TensorStorage::<u8>::from_fd(fd, &shape, None).unwrap();
6674 let mut t = Tensor::<u8>::wrap(storage);
6675 t.set_format(PixelFormat::Bgra).unwrap();
6676 t.set_row_stride(320 * 4).unwrap();
6677 assert!(
6678 t.map().is_ok(),
6679 "an in-bounds strided map of an imported DMA-BUF must be \
6680 accepted -- refusing it is what broke decode-into-GPU-tensor"
6681 );
6682 }
6683
6684 // Out of bounds: a stride the allocation cannot back must still be
6685 // refused, and the error must say so. This is the half that carries
6686 // the soundness.
6687 {
6688 let fd = backing.clone_fd().unwrap();
6689 let shape = [240usize, 320, 4];
6690 let storage = TensorStorage::<u8>::from_fd(fd, &shape, None).unwrap();
6691 let mut t = Tensor::<u8>::wrap(storage);
6692 t.set_format(PixelFormat::Bgra).unwrap();
6693 // Double the pitch: 240 rows x 2560 bytes is twice the allocation.
6694 t.set_row_stride(320 * 4 * 2).unwrap();
6695 let err = t.map();
6696 assert!(
6697 matches!(err, Err(Error::InvalidOperation(_))),
6698 "a strided map larger than the DMA allocation must be refused"
6699 );
6700 }
6701 }
6702
6703 #[test]
6704 #[cfg(target_os = "linux")]
6705 fn image_tensor_with_stride_map_rejects_tampered_stride() {
6706 // Declares this test as an fd-opener; see FD_LOCK.
6707 let _lock = crate::tests::fd_lock_shared();
6708 // Round-3 PR feedback (C1): `set_row_stride` is public and only
6709 // validates `stride >= min_stride`, not that the new stride × height
6710 // fits the underlying buffer. A caller that tampers with the stride
6711 // after allocation must not be able to coerce `Tensor::map()` into
6712 // returning a slice larger than the backing mmap (that would be UB
6713 // in `DmaMap::as_slice`).
6714 if !is_dma_available() {
6715 crate::test_support::report_skip("DMA heap not available");
6716 return;
6717 }
6718 // Allocate a 640×480 RGBA8 padded canvas (stride = 3072 = 768 px).
6719 // Backing buffer is 3072 × 480 = 1,474,560 bytes.
6720 let mut t = Tensor::<u8>::image_with_stride(
6721 640,
6722 480,
6723 PixelFormat::Rgba,
6724 3072,
6725 Some(TensorMemory::DmaBuf),
6726 crate::CpuAccess::ReadWrite,
6727 )
6728 .unwrap();
6729 // Tamper: push the stride up to 4 × the original. This is >=
6730 // min_stride (2560), so `set_row_stride` accepts it.
6731 t.set_row_stride(12288).unwrap();
6732 // Map must now refuse — 12288 × 480 = 5,898,240 > 1,474,560.
6733 let err = t.map();
6734 assert!(
6735 matches!(err, Err(Error::InvalidOperation(_))),
6736 "map() with oversized stride must return InvalidOperation"
6737 );
6738 }
6739
6740 #[test]
6741 fn dma_tensor_new_with_byte_size_rejects_shape_overflow() {
6742 // Round-3 PR feedback (C3): shape.product() * sizeof(T) must use
6743 // checked arithmetic so a pathological shape can't wrap usize and
6744 // make the byte_size-vs-logical-size comparison incorrect.
6745 //
6746 // This test only exercises the overflow rejection path, which is
6747 // pure-Rust and doesn't touch dma_heap — safe to run on any target.
6748 #[cfg(target_os = "linux")]
6749 {
6750 let err = crate::dma::DmaTensor::<u64>::new_with_byte_size(
6751 &[usize::MAX, 2, 2],
6752 usize::MAX,
6753 None,
6754 );
6755 assert!(
6756 matches!(err, Err(Error::InvalidArgument(_))),
6757 "new_with_byte_size must detect shape.product() overflow"
6758 );
6759 }
6760 }
6761
6762 #[test]
6763 #[cfg(target_os = "linux")]
6764 fn image_tensor_with_stride_rejects_too_small_stride() {
6765 // Declares this test as an fd-opener; see FD_LOCK.
6766 let _lock = crate::tests::fd_lock_shared();
6767 // 640×480 RGBA8 natural pitch = 2560, request 2400 → should error.
6768 let err = Tensor::<u8>::image_with_stride(
6769 640,
6770 480,
6771 PixelFormat::Rgba,
6772 2400,
6773 Some(TensorMemory::DmaBuf),
6774 crate::CpuAccess::ReadWrite,
6775 );
6776 assert!(matches!(err, Err(Error::InvalidArgument(_))));
6777 }
6778
6779 #[test]
6780 #[cfg(target_os = "linux")]
6781 fn image_tensor_with_stride_rejects_non_packed() {
6782 // Declares this test as an fd-opener; see FD_LOCK.
6783 let _lock = crate::tests::fd_lock_shared();
6784 // NV12 is SemiPlanar → not supported. (Linux-only because
6785 // `TensorMemory::DmaBuf` itself is a Linux-only enum variant.)
6786 let err = Tensor::<u8>::image_with_stride(
6787 640,
6788 480,
6789 PixelFormat::Nv12,
6790 640,
6791 Some(TensorMemory::DmaBuf),
6792 crate::CpuAccess::ReadWrite,
6793 );
6794 assert!(matches!(err, Err(Error::NotImplemented(_))));
6795 }
6796
6797 #[test]
6798 fn set_format_valid() {
6799 let mut t = Tensor::<u8>::new(&[480, 640, 3], None, None).unwrap();
6800 assert!(t.format().is_none());
6801 t.set_format(PixelFormat::Rgb).unwrap();
6802 assert_eq!(t.format(), Some(PixelFormat::Rgb));
6803 assert_eq!(t.width(), Some(640));
6804 assert_eq!(t.height(), Some(480));
6805 }
6806
6807 #[test]
6808 fn set_format_invalid_shape() {
6809 let mut t = Tensor::<u8>::new(&[480, 640, 4], None, None).unwrap();
6810 // RGB expects 3 channels, not 4
6811 let err = t.set_format(PixelFormat::Rgb);
6812 assert!(err.is_err());
6813 // Original tensor is unmodified
6814 assert!(t.format().is_none());
6815 }
6816
6817 #[test]
6818 fn reshape_clears_format() {
6819 let mut t = Tensor::<u8>::image(
6820 640,
6821 480,
6822 PixelFormat::Rgba,
6823 None,
6824 crate::CpuAccess::ReadWrite,
6825 )
6826 .unwrap();
6827 assert_eq!(t.format(), Some(PixelFormat::Rgba));
6828 // Reshape to flat — format cleared
6829 t.reshape(&[480 * 640 * 4]).unwrap();
6830 assert!(t.format().is_none());
6831 }
6832
6833 #[test]
6834 fn from_planes_nv12() {
6835 let y = Tensor::<u8>::new(&[480, 640], None, None).unwrap();
6836 let uv = Tensor::<u8>::new(&[240, 640], None, None).unwrap();
6837 let img = Tensor::from_planes(y, uv, PixelFormat::Nv12).unwrap();
6838 assert_eq!(img.format(), Some(PixelFormat::Nv12));
6839 assert!(img.is_multiplane());
6840 assert!(img.chroma().is_some());
6841 assert_eq!(img.width(), Some(640));
6842 assert_eq!(img.height(), Some(480));
6843 }
6844
6845 #[test]
6846 fn from_planes_rejects_non_semiplanar() {
6847 let y = Tensor::<u8>::new(&[480, 640], None, None).unwrap();
6848 let uv = Tensor::<u8>::new(&[240, 640], None, None).unwrap();
6849 let err = Tensor::from_planes(y, uv, PixelFormat::Rgb);
6850 assert!(err.is_err());
6851 }
6852
6853 #[test]
6854 fn reshape_multiplane_errors() {
6855 let y = Tensor::<u8>::new(&[480, 640], None, None).unwrap();
6856 let uv = Tensor::<u8>::new(&[240, 640], None, None).unwrap();
6857 let mut img = Tensor::from_planes(y, uv, PixelFormat::Nv12).unwrap();
6858 let err = img.reshape(&[480 * 640 + 240 * 640]);
6859 assert!(err.is_err());
6860 }
6861}
6862
6863#[cfg(test)]
6864mod compression_tests {
6865 use super::*;
6866
6867 #[test]
6868 fn desc_builder_roundtrips() {
6869 let desc = ImageDesc::new(640, 480, PixelFormat::Rgba, DType::U8)
6870 .with_memory(Some(TensorMemory::Mem))
6871 .with_access(CpuAccess::Read)
6872 .with_compression(Compression::Any);
6873 assert_eq!(desc.width(), 640);
6874 assert_eq!(desc.height(), 480);
6875 assert_eq!(desc.format(), PixelFormat::Rgba);
6876 assert_eq!(desc.dtype(), DType::U8);
6877 assert_eq!(desc.memory(), Some(TensorMemory::Mem));
6878 assert_eq!(desc.access(), CpuAccess::Read);
6879 assert_eq!(desc.compression(), Some(Compression::Any));
6880
6881 // Defaults: auto memory, hardware-only, no request.
6882 let plain = ImageDesc::new(2, 2, PixelFormat::Grey, DType::U8);
6883 assert_eq!(plain.memory(), None);
6884 assert_eq!(plain.access(), CpuAccess::None);
6885 assert_eq!(plain.compression(), None);
6886 }
6887
6888 #[test]
6889 fn desc_dtype_must_match_element_type() {
6890 let desc = ImageDesc::new(4, 4, PixelFormat::Rgba, DType::F32);
6891 match Tensor::<u8>::image_desc(&desc) {
6892 Err(Error::InvalidArgument(msg)) => assert!(msg.contains("dtype")),
6893 other => panic!("expected InvalidArgument, got {other:?}"),
6894 }
6895 }
6896
6897 #[test]
6898 fn compression_with_cpu_access_is_invalid() {
6899 let desc = ImageDesc::new(4, 4, PixelFormat::Rgba, DType::U8)
6900 .with_access(CpuAccess::ReadWrite)
6901 .with_compression(Compression::Any);
6902 match Tensor::<u8>::image_desc(&desc) {
6903 Err(Error::InvalidArgument(msg)) => assert!(msg.contains("CpuAccess::None")),
6904 other => panic!("expected InvalidArgument, got {other:?}"),
6905 }
6906 }
6907
6908 #[cfg(not(target_os = "android"))]
6909 #[test]
6910 fn scheme_request_off_android_is_not_implemented() {
6911 let desc = ImageDesc::new(4, 4, PixelFormat::Rgba, DType::U8)
6912 .with_compression(Compression::Scheme(CompressionScheme::Ubwc));
6913 match Tensor::<u8>::image_desc(&desc) {
6914 Err(Error::NotImplemented(msg)) => assert!(msg.contains("Ubwc")),
6915 other => panic!("expected NotImplemented, got {other:?}"),
6916 }
6917 }
6918
6919 #[cfg(not(target_os = "android"))]
6920 #[test]
6921 fn any_request_off_android_resolves_linear_and_counts() {
6922 let before = compression_fallback_count();
6923 let desc = ImageDesc::new(64, 64, PixelFormat::Rgba, DType::U8)
6924 .with_memory(Some(TensorMemory::Mem))
6925 .with_compression(Compression::Any);
6926 let t = Tensor::<u8>::image_desc(&desc).unwrap();
6927 assert_eq!(t.compression(), None);
6928 assert!(compression_fallback_count() > before);
6929 }
6930
6931 #[test]
6932 fn desc_without_request_matches_classic_constructor() {
6933 let desc = ImageDesc::new(32, 32, PixelFormat::Rgba, DType::U8)
6934 .with_memory(Some(TensorMemory::Mem))
6935 .with_access(CpuAccess::ReadWrite);
6936 let t = Tensor::<u8>::image_desc(&desc).unwrap();
6937 assert_eq!(t.compression(), None);
6938 assert_eq!(t.cpu_access(), CpuAccess::ReadWrite);
6939 assert_eq!(t.width(), Some(32));
6940 // Mappable exactly like the classic constructor's result.
6941 let m = t.map_read().unwrap();
6942 assert_eq!(m.as_slice().len(), 32 * 32 * 4);
6943 }
6944
6945 #[test]
6946 fn configure_image_preserves_compression_and_views_inherit() {
6947 // No host platform records a scheme, so emulate the recording to
6948 // pin the preserve/inherit semantics (the physical layout does
6949 // not change when the logical image is reconfigured).
6950 let mut t = Tensor::<u8>::image(
6951 64,
6952 64,
6953 PixelFormat::Rgba,
6954 Some(TensorMemory::Mem),
6955 CpuAccess::None,
6956 )
6957 .unwrap();
6958 t.compression = Some(CompressionScheme::Ubwc);
6959 t.configure_image(32, 32, PixelFormat::Rgba).unwrap();
6960 assert_eq!(t.compression(), Some(CompressionScheme::Ubwc));
6961 let view = t.subview(0, &[16, 32, 4]).unwrap();
6962 assert_eq!(view.compression(), Some(CompressionScheme::Ubwc));
6963 }
6964
6965 #[test]
6966 fn tensor_dyn_dispatches_desc_and_compression() {
6967 let desc = ImageDesc::new(16, 16, PixelFormat::Rgba, DType::U8)
6968 .with_memory(Some(TensorMemory::Mem))
6969 .with_access(CpuAccess::ReadWrite);
6970 let t = TensorDyn::image_desc(&desc).unwrap();
6971 assert_eq!(t.compression(), None);
6972 assert!(matches!(t, TensorDyn::U8(_)));
6973 }
6974}
6975
6976#[cfg(test)]
6977mod cpu_access_tests {
6978 use super::*;
6979
6980 #[test]
6981 fn covers_matrix() {
6982 use CpuAccess::*;
6983 // Every declaration covers a narrower or equal request…
6984 for a in [None, Read, Write, ReadWrite] {
6985 assert!(a.covers(None), "{a:?} must cover None");
6986 assert!(ReadWrite.covers(a), "ReadWrite must cover {a:?}");
6987 }
6988 assert!(Read.covers(Read));
6989 assert!(Write.covers(Write));
6990 // …and never a wider one.
6991 assert!(!None.covers(Read));
6992 assert!(!None.covers(Write));
6993 assert!(!Read.covers(Write));
6994 assert!(!Read.covers(ReadWrite));
6995 assert!(!Write.covers(Read));
6996 assert!(!Write.covers(ReadWrite));
6997 }
6998
6999 #[test]
7000 fn map_with_none_is_invalid() {
7001 let t = Tensor::<u8>::new(&[16], Some(TensorMemory::Mem), None).unwrap();
7002 match t.map_with(CpuAccess::None) {
7003 Err(Error::InvalidArgument(_)) => {}
7004 Err(other) => panic!("expected InvalidArgument, got {other:?}"),
7005 Ok(_) => panic!("map_with(CpuAccess::None) must not succeed"),
7006 }
7007 }
7008
7009 #[test]
7010 fn read_map_rejects_mutation_uniformly() {
7011 // Mem backend: map_read yields a working read view whose mutable
7012 // accessor panics (the uniform cross-backend contract).
7013 let t = Tensor::<u8>::new(&[8], Some(TensorMemory::Mem), None).unwrap();
7014 t.map_mut().unwrap().as_mut_slice().copy_from_slice(&[7; 8]);
7015 let ro = t.map_read().unwrap();
7016 assert_eq!(ro.as_slice(), &[7; 8]);
7017 drop(ro);
7018 let result = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
7019 let mut ro = t.map_read().unwrap();
7020 let _ = ro.as_mut_slice();
7021 }));
7022 assert!(result.is_err(), "as_mut_slice through map_read must panic");
7023 }
7024
7025 #[test]
7026 fn write_and_rw_maps_stay_mutable() {
7027 let t = Tensor::<u8>::new(&[4], Some(TensorMemory::Mem), None).unwrap();
7028 t.map_write()
7029 .unwrap()
7030 .as_mut_slice()
7031 .copy_from_slice(&[1; 4]);
7032 t.map().unwrap().as_mut_slice().copy_from_slice(&[2; 4]);
7033 assert_eq!(t.map_read().unwrap().as_slice(), &[2; 4]);
7034 }
7035
7036 /// The read-only IOSurface lock path: a `map_read` must observe data
7037 /// written through a prior read-write lock, and its unlock (which
7038 /// skips the cache flush) must not disturb subsequent reads.
7039 #[test]
7040 #[cfg(target_os = "macos")]
7041 fn iosurface_read_only_lock_roundtrip() {
7042 // Declares this test as an fd-opener; see FD_LOCK.
7043 let _lock = crate::tests::fd_lock_shared();
7044 let Ok(t) = Tensor::<u8>::new(&[64], Some(TensorMemory::DmaBuf), None) else {
7045 crate::test_support::report_skip("IOSurface unavailable");
7046 return;
7047 };
7048 {
7049 let mut m = t.map_mut().unwrap();
7050 for (i, b) in m.as_mut_slice().iter_mut().enumerate() {
7051 *b = (i * 3) as u8;
7052 }
7053 }
7054 for _ in 0..2 {
7055 let ro = t.map_read().unwrap();
7056 for (i, b) in ro.as_slice().iter().enumerate() {
7057 assert_eq!(*b, (i * 3) as u8, "byte {i} through read-only lock");
7058 }
7059 }
7060 let result = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
7061 let mut ro = t.map_read().unwrap();
7062 let _ = ro.as_mut_slice();
7063 }));
7064 assert!(result.is_err(), "IOSurface read map must reject mutation");
7065 }
7066}
7067
7068#[cfg(test)]
7069mod tests {
7070 #[cfg(target_os = "linux")]
7071 use nix::unistd::{access, AccessFlags};
7072 #[cfg(target_os = "linux")]
7073 use std::io::Write as _;
7074 use std::sync::{PoisonError, RwLock};
7075
7076 use super::*;
7077
7078 #[ctor::ctor(unsafe)]
7079 fn init() {
7080 env_logger::Builder::from_env(env_logger::Env::default().default_filter_or("info")).init();
7081 }
7082
7083 /// Macro to get the current function name for logging in tests.
7084 #[cfg(target_os = "linux")]
7085 macro_rules! function {
7086 () => {{
7087 fn f() {}
7088 fn type_name_of<T>(_: T) -> &'static str {
7089 std::any::type_name::<T>()
7090 }
7091 let name = type_name_of(f);
7092
7093 // Find and cut the rest of the path
7094 match &name[..name.len() - 3].rfind(':') {
7095 Some(pos) => &name[pos + 1..name.len() - 3],
7096 None => &name[..name.len() - 3],
7097 }
7098 }};
7099 }
7100
7101 #[test]
7102 #[cfg(target_os = "linux")]
7103 fn test_tensor() {
7104 let _lock = crate::tests::fd_lock_shared();
7105 let shape = vec![1];
7106 let tensor = DmaTensor::<f32>::new(&shape, Some("dma_tensor"));
7107 let dma_enabled = tensor.is_ok();
7108
7109 let tensor = Tensor::<f32>::new(&shape, None, None).expect("Failed to create tensor");
7110 // Auto-select priority is Dma > Mem; Shm is never auto-selected.
7111 match dma_enabled {
7112 true => assert_eq!(tensor.memory(), TensorMemory::DmaBuf),
7113 false => assert_eq!(tensor.memory(), TensorMemory::Mem),
7114 }
7115 }
7116
7117 #[test]
7118 #[cfg(any(target_os = "macos", target_os = "ios"))]
7119 fn test_tensor() {
7120 let shape = vec![1];
7121 let tensor = Tensor::<f32>::new(&shape, None, None).expect("Failed to create tensor");
7122 // macOS/iOS auto-fallback chain: IOSurface (Dma) → Mem. Healthy systems
7123 // return Dma; Mem only appears under memory pressure or sandboxed
7124 // contexts where IOSurfaceCreate fails. Shm is never auto-selected.
7125 let m = tensor.memory();
7126 assert!(
7127 matches!(m, TensorMemory::DmaBuf | TensorMemory::Mem),
7128 "Unexpected auto-fallback result on macOS/iOS: {m:?}"
7129 );
7130 }
7131
7132 #[test]
7133 #[cfg(all(
7134 unix,
7135 not(any(target_os = "linux", target_os = "macos", target_os = "ios"))
7136 ))]
7137 fn test_tensor() {
7138 let shape = vec![1];
7139 let tensor = Tensor::<f32>::new(&shape, None, None).expect("Failed to create tensor");
7140 // Other Unix (BSD): no DMA, so auto-selection is Mem (Shm is
7141 // explicit-only, never auto-selected).
7142 assert_eq!(tensor.memory(), TensorMemory::Mem);
7143 }
7144
7145 #[test]
7146 #[cfg(not(unix))]
7147 fn test_tensor() {
7148 let shape = vec![1];
7149 let tensor = Tensor::<f32>::new(&shape, None, None).expect("Failed to create tensor");
7150 assert_eq!(tensor.memory(), TensorMemory::Mem);
7151 }
7152
7153 #[test]
7154 #[cfg(target_os = "linux")]
7155 fn test_dma_tensor() {
7156 let _lock = crate::tests::fd_lock_shared();
7157 match access(
7158 "/dev/dma_heap/linux,cma",
7159 AccessFlags::R_OK | AccessFlags::W_OK,
7160 ) {
7161 Ok(_) => println!("/dev/dma_heap/linux,cma is available"),
7162 Err(_) => match access(
7163 "/dev/dma_heap/system",
7164 AccessFlags::R_OK | AccessFlags::W_OK,
7165 ) {
7166 Ok(_) => println!("/dev/dma_heap/system is available"),
7167 Err(e) => {
7168 writeln!(
7169 &mut std::io::stdout(),
7170 "[WARNING] DMA Heap is unavailable: {e}"
7171 )
7172 .unwrap();
7173 return;
7174 }
7175 },
7176 }
7177
7178 let shape = vec![2, 3, 4];
7179 let tensor =
7180 DmaTensor::<f32>::new(&shape, Some("test_tensor")).expect("Failed to create tensor");
7181
7182 const DUMMY_VALUE: f32 = 12.34;
7183
7184 assert_eq!(tensor.memory(), TensorMemory::DmaBuf);
7185 assert_eq!(tensor.name(), "test_tensor");
7186 assert_eq!(tensor.shape(), &shape);
7187 assert_eq!(tensor.size(), 2 * 3 * 4 * std::mem::size_of::<f32>());
7188 assert_eq!(tensor.len(), 2 * 3 * 4);
7189
7190 {
7191 let mut tensor_map = tensor.map().expect("Failed to map DMA memory");
7192 tensor_map.fill(42.0);
7193 assert!(tensor_map.iter().all(|&x| x == 42.0));
7194 }
7195
7196 {
7197 let shared = Tensor::<f32>::from_fd(
7198 tensor
7199 .clone_fd()
7200 .expect("Failed to duplicate tensor file descriptor"),
7201 &shape,
7202 Some("test_tensor_shared"),
7203 )
7204 .expect("Failed to create tensor from fd");
7205
7206 assert_eq!(shared.memory(), TensorMemory::DmaBuf);
7207 assert_eq!(shared.name(), "test_tensor_shared");
7208 assert_eq!(shared.shape(), &shape);
7209
7210 let mut tensor_map = shared.map().expect("Failed to map DMA memory from fd");
7211 tensor_map.fill(DUMMY_VALUE);
7212 assert!(tensor_map.iter().all(|&x| x == DUMMY_VALUE));
7213 }
7214
7215 {
7216 let tensor_map = tensor.map().expect("Failed to map DMA memory");
7217 assert!(tensor_map.iter().all(|&x| x == DUMMY_VALUE));
7218 }
7219
7220 let mut tensor = DmaTensor::<u8>::new(&shape, None).expect("Failed to create tensor");
7221 assert_eq!(tensor.shape(), &shape);
7222 let new_shape = vec![3, 4, 4];
7223 assert!(
7224 tensor.reshape(&new_shape).is_err(),
7225 "Reshape should fail due to size mismatch"
7226 );
7227 assert_eq!(tensor.shape(), &shape, "Shape should remain unchanged");
7228
7229 let new_shape = vec![2, 3, 4];
7230 tensor.reshape(&new_shape).expect("Reshape should succeed");
7231 assert_eq!(
7232 tensor.shape(),
7233 &new_shape,
7234 "Shape should be updated after successful reshape"
7235 );
7236
7237 {
7238 let mut tensor_map = tensor.map().expect("Failed to map DMA memory");
7239 tensor_map.fill(1);
7240 assert!(tensor_map.iter().all(|&x| x == 1));
7241 }
7242
7243 {
7244 let mut tensor_map = tensor.map().expect("Failed to map DMA memory");
7245 tensor_map[2] = 42;
7246 assert_eq!(tensor_map[1], 1, "Value at index 1 should be 1");
7247 assert_eq!(tensor_map[2], 42, "Value at index 2 should be 42");
7248 }
7249 }
7250
7251 #[test]
7252 #[cfg(unix)]
7253 fn test_shm_tensor() {
7254 let _lock = crate::tests::fd_lock_shared();
7255 if !is_shm_available() {
7256 log::warn!("SKIPPED: test_shm_tensor - SHM allocation unavailable on this platform");
7257 return;
7258 }
7259 let shape = vec![2, 3, 4];
7260 let tensor =
7261 ShmTensor::<f32>::new(&shape, Some("test_tensor")).expect("Failed to create tensor");
7262 assert_eq!(tensor.shape(), &shape);
7263 assert_eq!(tensor.size(), 2 * 3 * 4 * std::mem::size_of::<f32>());
7264 assert_eq!(tensor.name(), "test_tensor");
7265
7266 const DUMMY_VALUE: f32 = 12.34;
7267 {
7268 let mut tensor_map = tensor.map().expect("Failed to map shared memory");
7269 tensor_map.fill(42.0);
7270 assert!(tensor_map.iter().all(|&x| x == 42.0));
7271 }
7272
7273 {
7274 let shared = Tensor::<f32>::from_fd(
7275 tensor
7276 .clone_fd()
7277 .expect("Failed to duplicate tensor file descriptor"),
7278 &shape,
7279 Some("test_tensor_shared"),
7280 )
7281 .expect("Failed to create tensor from fd");
7282
7283 assert_eq!(shared.memory(), TensorMemory::Shm);
7284 assert_eq!(shared.name(), "test_tensor_shared");
7285 assert_eq!(shared.shape(), &shape);
7286
7287 let mut tensor_map = shared.map().expect("Failed to map shared memory from fd");
7288 tensor_map.fill(DUMMY_VALUE);
7289 assert!(tensor_map.iter().all(|&x| x == DUMMY_VALUE));
7290 }
7291
7292 {
7293 let tensor_map = tensor.map().expect("Failed to map shared memory");
7294 assert!(tensor_map.iter().all(|&x| x == DUMMY_VALUE));
7295 }
7296
7297 let mut tensor = ShmTensor::<u8>::new(&shape, None).expect("Failed to create tensor");
7298 assert_eq!(tensor.shape(), &shape);
7299 let new_shape = vec![3, 4, 4];
7300 assert!(
7301 tensor.reshape(&new_shape).is_err(),
7302 "Reshape should fail due to size mismatch"
7303 );
7304 assert_eq!(tensor.shape(), &shape, "Shape should remain unchanged");
7305
7306 let new_shape = vec![2, 3, 4];
7307 tensor.reshape(&new_shape).expect("Reshape should succeed");
7308 assert_eq!(
7309 tensor.shape(),
7310 &new_shape,
7311 "Shape should be updated after successful reshape"
7312 );
7313
7314 {
7315 let mut tensor_map = tensor.map().expect("Failed to map shared memory");
7316 tensor_map.fill(1);
7317 assert!(tensor_map.iter().all(|&x| x == 1));
7318 }
7319
7320 {
7321 let mut tensor_map = tensor.map().expect("Failed to map shared memory");
7322 tensor_map[2] = 42;
7323 assert_eq!(tensor_map[1], 1, "Value at index 1 should be 1");
7324 assert_eq!(tensor_map[2], 42, "Value at index 2 should be 42");
7325 }
7326 }
7327
7328 #[test]
7329 fn mem_subview_partitions_parent_buffer() {
7330 // One heap [2,4] u8 parent (8 bytes). Two [1,4] sub-views at byte
7331 // offsets 0 and 4 must share the parent allocation (zero-copy) and be
7332 // independently writable: view 0 owns bytes [0,4), view 1 owns [4,8).
7333 // Today this is impossible — heap offset is rejected and there is no
7334 // shared sub-view constructor.
7335 let parent = Tensor::<u8>::new(&[2, 4], Some(TensorMemory::Mem), None).unwrap();
7336 let view0 = parent.subview(0, &[1, 4]).expect("subview at offset 0");
7337 let view1 = parent.subview(4, &[1, 4]).expect("subview at offset 4");
7338
7339 view1
7340 .map()
7341 .unwrap()
7342 .as_mut_slice()
7343 .copy_from_slice(&[10, 20, 30, 40]);
7344 view0
7345 .map()
7346 .unwrap()
7347 .as_mut_slice()
7348 .copy_from_slice(&[1, 2, 3, 4]);
7349
7350 // Each view sees only its own window.
7351 assert_eq!(view0.map().unwrap().as_slice(), &[1, 2, 3, 4]);
7352 assert_eq!(view1.map().unwrap().as_slice(), &[10, 20, 30, 40]);
7353 // The parent buffer is correctly partitioned (shared, zero-copy).
7354 assert_eq!(
7355 parent.map().unwrap().as_slice(),
7356 &[1, 2, 3, 4, 10, 20, 30, 40]
7357 );
7358 }
7359
7360 #[test]
7361 fn batch_partitions_leading_dim() {
7362 // Raw [4,2,2,3] u8 batched tensor: 4 elements of 12 bytes each. batch(n)
7363 // yields element n at offset n*12, sharing the parent buffer (zero-copy).
7364 let parent = Tensor::<u8>::new(&[4, 2, 2, 3], Some(TensorMemory::Mem), None).unwrap();
7365 for i in 0..4u8 {
7366 let e = parent.batch(i as usize).expect("batch element");
7367 assert_eq!(e.shape(), &[2, 2, 3]);
7368 // A batch element shares the parent's BufferIdentity.
7369 assert_eq!(e.buffer_identity().id(), parent.buffer_identity().id());
7370 for b in e.map().unwrap().as_mut_slice() {
7371 *b = i + 1;
7372 }
7373 }
7374 // Each element occupies its own 12-byte band of the parent.
7375 let whole = parent.map().unwrap();
7376 let s = whole.as_slice();
7377 for i in 0..4usize {
7378 assert!(
7379 s[i * 12..(i + 1) * 12].iter().all(|&b| b == (i as u8 + 1)),
7380 "band {i} not partitioned: {:?}",
7381 &s[i * 12..(i + 1) * 12]
7382 );
7383 }
7384 }
7385
7386 #[test]
7387 fn view_origin_snapshots_parent_and_composes() {
7388 // view() on a whole image snapshots the parent dims + the view's origin.
7389 let parent = Tensor::<u8>::image(
7390 100,
7391 80,
7392 PixelFormat::Rgba,
7393 Some(TensorMemory::Mem),
7394 crate::CpuAccess::ReadWrite,
7395 )
7396 .unwrap();
7397 assert_eq!(
7398 parent.view_origin(),
7399 None,
7400 "whole tensor has no view_origin"
7401 );
7402 let v = parent.view(Region::new(10, 20, 30, 40)).unwrap();
7403 assert_eq!(
7404 v.view_origin(),
7405 Some(ViewOrigin {
7406 parent_width: 100,
7407 parent_height: 80,
7408 parent_row_stride: 100 * 4, // tight RGBA pitch
7409 x: 10,
7410 y: 20
7411 })
7412 );
7413 // A view of a view keeps the ROOT parent and accumulates the origin.
7414 let v2 = v.view(Region::new(5, 5, 10, 10)).unwrap();
7415 assert_eq!(
7416 v2.view_origin(),
7417 Some(ViewOrigin {
7418 parent_width: 100,
7419 parent_height: 80,
7420 parent_row_stride: 100 * 4,
7421 x: 15,
7422 y: 25
7423 }),
7424 "nested view composes onto the root parent"
7425 );
7426 }
7427
7428 #[test]
7429 fn view_origin_none_for_raw_batch() {
7430 // A raw (unformatted) batched tensor has no pixel geometry, so batch()
7431 // leaves view_origin None (the per-slot path, not the one-import pivot).
7432 let parent = Tensor::<u8>::new(&[4, 2, 2, 3], Some(TensorMemory::Mem), None).unwrap();
7433 assert_eq!(parent.batch(2).unwrap().view_origin(), None);
7434 }
7435
7436 #[test]
7437 fn batch_rejects_out_of_bounds_index() {
7438 let parent = Tensor::<u8>::new(&[4, 2, 2, 3], Some(TensorMemory::Mem), None).unwrap();
7439 match parent.batch(4) {
7440 Err(Error::BatchIndexOutOfBounds { index, batch }) => {
7441 assert_eq!((index, batch), (4, 4));
7442 }
7443 other => panic!("expected BatchIndexOutOfBounds, got {other:?}"),
7444 }
7445 }
7446
7447 #[test]
7448 fn batch_zero_on_unit_n_is_whole() {
7449 // N == 1: batch(0) is the whole per-element block at offset 0 (no plane_offset).
7450 let parent = Tensor::<u8>::new(&[1, 2, 2, 3], Some(TensorMemory::Mem), None).unwrap();
7451 let e = parent.batch(0).unwrap();
7452 assert_eq!(e.shape(), &[2, 2, 3]);
7453 assert_eq!(e.plane_offset(), None);
7454 assert_eq!(e.buffer_identity().id(), parent.buffer_identity().id());
7455 }
7456
7457 #[test]
7458 fn mem_subview_rejects_unaligned_offset() {
7459 // f32 has align 4; a byte offset of 2 cannot back a valid `*const f32`.
7460 let parent = Tensor::<f32>::new(&[8], Some(TensorMemory::Mem), None).unwrap();
7461 assert!(parent.subview(2, &[1]).is_err());
7462 // A correctly aligned offset is accepted.
7463 assert!(parent.subview(4, &[1]).is_ok());
7464 }
7465
7466 #[test]
7467 fn mem_subview_rejects_out_of_bounds() {
7468 let parent = Tensor::<u8>::new(&[8], Some(TensorMemory::Mem), None).unwrap();
7469 // offset 6 + 4 bytes = 10 exceeds the 8-byte allocation.
7470 assert!(parent.subview(6, &[4]).is_err());
7471 }
7472
7473 /// Regression guard for the `TensorTrait::view` promotion (R2): a `subview`
7474 /// must share the parent's `BufferIdentity` on **every** backend, not mint a
7475 /// fresh one. Identity-keyed caches (the GL EGLImage import) rely on this to
7476 /// treat offset-distinct windows of one buffer as a single import; a fresh
7477 /// identity would silently break that and regress zero-copy import reuse.
7478 ///
7479 /// Runs each backend that can be allocated without a GPU/GL context on the
7480 /// test host: `Mem` (always); `Shm` (when POSIX shm is available); the
7481 /// platform-native zero-copy buffer `Dma` (DMA-BUF on Linux / IOSurface on
7482 /// macOS, when available). `Pbo` shares its identity the same way (see
7483 /// `pbo.rs` `view`) but needs a live GL context, so it is exercised by the
7484 /// image-crate GL tests rather than here.
7485 #[test]
7486 fn subview_shares_buffer_identity_all_backends() {
7487 // Declares this test as an fd-opener; see FD_LOCK.
7488 let _lock = crate::tests::fd_lock_shared();
7489 // u8 has align 1, so every byte offset is valid for the alignment check;
7490 // this isolates the identity-sharing contract from alignment concerns.
7491 // `window` is the view shape; a formatted parent validates it against
7492 // its own format, so an image parent needs an image-shaped window.
7493 let assert_parent_shares = |parent: Tensor<u8>, window: &[usize], label: &str| {
7494 let parent_id = parent.buffer_identity().id();
7495 // Two offset-distinct windows must both carry the parent's identity.
7496 let v0 = parent
7497 .subview(0, window)
7498 .unwrap_or_else(|e| panic!("{label}: subview(0) failed: {e:?}"));
7499 let v1 = parent
7500 .subview(16, window)
7501 .unwrap_or_else(|e| panic!("{label}: subview(16) failed: {e:?}"));
7502 assert_eq!(
7503 v0.buffer_identity().id(),
7504 parent_id,
7505 "{label}: subview(0) minted a fresh BufferIdentity"
7506 );
7507 assert_eq!(
7508 v1.buffer_identity().id(),
7509 parent_id,
7510 "{label}: subview(16) minted a fresh BufferIdentity"
7511 );
7512 };
7513 let assert_shares = |memory: TensorMemory, label: &str| {
7514 let parent = Tensor::<u8>::new(&[64], Some(memory), None)
7515 .unwrap_or_else(|e| panic!("{label}: parent alloc failed: {e:?}"));
7516 assert_parent_shares(parent, &[16], label);
7517 };
7518
7519 assert_shares(TensorMemory::Mem, "Mem");
7520
7521 #[cfg(unix)]
7522 if crate::is_shm_available() {
7523 assert_shares(TensorMemory::Shm, "Shm");
7524 }
7525
7526 // Dma == DMA-BUF on Linux, IOSurface on macOS, AHardwareBuffer on
7527 // Android; same public variant.
7528 #[cfg(not(target_os = "windows"))]
7529 if crate::is_gpu_buffer_available() {
7530 assert_shares(TensorMemory::DmaBuf, "Dma");
7531 }
7532 // Windows' Dma is an image-formatted D3D11 texture, which
7533 // `Tensor::new` deliberately refuses -- allocate it the only way it
7534 // can be allocated and hold it to the same contract.
7535 #[cfg(target_os = "windows")]
7536 if crate::is_gpu_buffer_available() {
7537 let parent = Tensor::<u8>::image(
7538 64,
7539 1,
7540 PixelFormat::Grey,
7541 Some(TensorMemory::DmaBuf),
7542 CpuAccess::ReadWrite,
7543 )
7544 .expect("D3D11 texture alloc failed");
7545 assert_parent_shares(parent, &[1, 16, 1], "Dma");
7546 }
7547 }
7548
7549 /// A sub-view records a plane offset, and on Windows the D3D11 storage
7550 /// already applied that offset itself when the view was made -- so the
7551 /// map's plane-offset gate must accept it, exactly as it does for the
7552 /// other zero-copy backings. It used to refuse, with a message that named
7553 /// every backend but this one.
7554 #[cfg(target_os = "windows")]
7555 #[test]
7556 fn subview_of_a_texture_tensor_maps_at_its_offset() {
7557 if !crate::is_gpu_buffer_available() {
7558 return;
7559 }
7560 let t = Tensor::<u8>::image(
7561 32,
7562 8,
7563 PixelFormat::Rgba,
7564 Some(TensorMemory::DmaBuf),
7565 CpuAccess::ReadWrite,
7566 )
7567 .unwrap();
7568 // For a texture tensor this is the backing pitch either way: the
7569 // allocation records the staging pitch when it exceeds the tight row,
7570 // and the two are equal when it does not.
7571 let stride = t.effective_row_stride().unwrap();
7572 let v = t.subview(stride * 2, &[2, 32, 4]).unwrap();
7573 assert!(
7574 v.plane_offset().is_some_and(|o| o > 0),
7575 "a sub-view at a row offset must carry one"
7576 );
7577 let m = v
7578 .map_read()
7579 .expect("a sub-view of a texture tensor must map");
7580 assert_eq!(m.as_slice().len(), 2 * stride, "two rows of the window");
7581 }
7582
7583 #[test]
7584 fn mem_subview_four_views_no_aliasing() {
7585 // One [4,3] f32 parent; four [1,3] views at 12-byte strides, each
7586 // written independently. Exercises a multi-byte element type (offsets
7587 // must stay element-aligned) and N-way zero-copy sharing.
7588 let parent = Tensor::<f32>::new(&[4, 3], Some(TensorMemory::Mem), None).unwrap();
7589 let frame = 3 * std::mem::size_of::<f32>();
7590 for i in 0..4 {
7591 let v = parent.subview(i * frame, &[1, 3]).unwrap();
7592 let val = i as f32 + 1.0;
7593 v.map()
7594 .unwrap()
7595 .as_mut_slice()
7596 .copy_from_slice(&[val, val, val]);
7597 }
7598 assert_eq!(
7599 parent.map().unwrap().as_slice(),
7600 &[1.0, 1.0, 1.0, 2.0, 2.0, 2.0, 3.0, 3.0, 3.0, 4.0, 4.0, 4.0]
7601 );
7602 }
7603
7604 #[test]
7605 fn mem_subview_inherits_format_and_row_stride() {
7606 // A sub-view is a ready-to-use sub-image: it inherits the parent's
7607 // pixel format and (crucially) its padded row stride, so a strided
7608 // parent yields strided windows. Set a stride wider than the tight row
7609 // to exercise the row_stride inheritance path specifically.
7610 let mut parent = Tensor::<u8>::image(
7611 100,
7612 100,
7613 PixelFormat::Rgba,
7614 Some(TensorMemory::Mem),
7615 crate::CpuAccess::ReadWrite,
7616 )
7617 .unwrap();
7618 parent.set_row_stride_unchecked(512); // padded stride (> 100*4)
7619 let view = parent.subview(4096, &[10, 10, 4]).unwrap();
7620 assert_eq!(view.format(), Some(PixelFormat::Rgba), "format inherited");
7621 assert_eq!(view.row_stride(), Some(512), "row_stride inherited");
7622 }
7623
7624 #[test]
7625 fn mem_strided_subview_maps_offset_and_byte_size() {
7626 // Integration of the sub-region offset (PR #89) and the strided-map
7627 // `byte_size_override` (PR #90): a strided sub-view exposes its full
7628 // padded window (`row_stride × rows`) starting at the view's byte
7629 // offset, mapped zero-copy into the parent.
7630 let parent = Tensor::<u8>::new(&[2048], Some(TensorMemory::Mem), None).unwrap();
7631 let mut view = parent.subview(128, &[8, 16]).unwrap(); // 8 rows × 16 @ off 128
7632 assert_eq!(view.plane_offset(), Some(128));
7633 view.set_row_stride_unchecked(32); // padded stride (> 16)
7634
7635 {
7636 let mut m = view.map().unwrap();
7637 let s = m.as_mut_slice();
7638 // Strided map exposes the padded window: stride(32) × rows(8) = 256.
7639 assert_eq!(
7640 s.len(),
7641 256,
7642 "strided map exposes the full padded byte window"
7643 );
7644 s[0] = 0xAA; // row 0, col 0
7645 s[32] = 0xBB; // row 1, col 0 (one stride in)
7646 }
7647
7648 // Zero-copy: the writes land in the parent at the view's offset.
7649 let p = parent.map().unwrap();
7650 let pb = p.as_slice();
7651 assert_eq!(pb[128], 0xAA, "row 0 writes at parent offset 128");
7652 assert_eq!(
7653 pb[128 + 32],
7654 0xBB,
7655 "row 1 writes at parent offset 128 + stride"
7656 );
7657 }
7658
7659 #[test]
7660 #[cfg(unix)]
7661 fn shm_subview_partitions_parent_buffer() {
7662 // Declares this test as an fd-opener; see FD_LOCK.
7663 let _lock = crate::tests::fd_lock_shared();
7664 // Mirrors `mem_subview_partitions_parent_buffer` for Shm: one [2,4] u8
7665 // parent shared segment (8 bytes); two [1,4] sub-views at byte offsets 0
7666 // and 4 must share the segment (zero-copy, via cloned fd) and be
7667 // independently writable — view 0 owns [0,4), view 1 owns [4,8).
7668 if !crate::is_shm_available() {
7669 crate::test_support::report_skip("shm not available");
7670 return;
7671 }
7672 let parent = Tensor::<u8>::new(&[2, 4], Some(TensorMemory::Shm), None).unwrap();
7673 let view0 = parent.subview(0, &[1, 4]).expect("shm subview at offset 0");
7674 let view1 = parent.subview(4, &[1, 4]).expect("shm subview at offset 4");
7675
7676 view1
7677 .map()
7678 .unwrap()
7679 .as_mut_slice()
7680 .copy_from_slice(&[10, 20, 30, 40]);
7681 view0
7682 .map()
7683 .unwrap()
7684 .as_mut_slice()
7685 .copy_from_slice(&[1, 2, 3, 4]);
7686
7687 assert_eq!(view0.map().unwrap().as_slice(), &[1, 2, 3, 4]);
7688 assert_eq!(view1.map().unwrap().as_slice(), &[10, 20, 30, 40]);
7689 // The parent sees the full partitioned segment (shared, zero-copy).
7690 assert_eq!(
7691 parent.map().unwrap().as_slice(),
7692 &[1, 2, 3, 4, 10, 20, 30, 40]
7693 );
7694 // A sub-view of a sub-view composes the offset.
7695 let nested = view1.subview(2, &[1, 2]).expect("nested shm subview");
7696 assert_eq!(nested.map().unwrap().as_slice(), &[30, 40]);
7697 }
7698
7699 #[test]
7700 #[cfg(unix)]
7701 fn shm_subview_rejects_unaligned_and_oob() {
7702 // Declares this test as an fd-opener; see FD_LOCK.
7703 let _lock = crate::tests::fd_lock_shared();
7704 if !crate::is_shm_available() {
7705 crate::test_support::report_skip("shm not available");
7706 return;
7707 }
7708 // f32 align 4: a 2-byte offset cannot back a valid `*const f32`.
7709 let parent = Tensor::<f32>::new(&[8], Some(TensorMemory::Shm), None).unwrap();
7710 assert!(parent.subview(2, &[1]).is_err());
7711 assert!(parent.subview(4, &[1]).is_ok());
7712 // Out of bounds: offset 6 + 4 bytes = 10 > 8-byte (u8) segment.
7713 let p2 = Tensor::<u8>::new(&[8], Some(TensorMemory::Shm), None).unwrap();
7714 assert!(p2.subview(6, &[4]).is_err());
7715 }
7716
7717 #[test]
7718 #[cfg(target_os = "linux")]
7719 fn dma_subview_matches_mem_subview() {
7720 // Serialize against the fd-leak tests: this test opens DMA fds (alloc +
7721 // clone_fd), which would otherwise perturb their fd counts.
7722 let _lock = crate::tests::fd_lock_shared();
7723 // Identical sub-view semantics across Dma (shared fd) and Mem (shared
7724 // Arc): same offsets → same logical windows → same partition.
7725 let dma = match Tensor::<u8>::new(&[8], Some(TensorMemory::DmaBuf), None) {
7726 Ok(t) => t,
7727 Err(_) => {
7728 crate::test_support::report_skip("DMA not available");
7729 return;
7730 }
7731 };
7732 let mem = Tensor::<u8>::new(&[8], Some(TensorMemory::Mem), None).unwrap();
7733 for parent in [&dma, &mem] {
7734 let v0 = parent.subview(0, &[4]).unwrap();
7735 let v1 = parent.subview(4, &[4]).unwrap();
7736 v0.map()
7737 .unwrap()
7738 .as_mut_slice()
7739 .copy_from_slice(&[1, 2, 3, 4]);
7740 v1.map()
7741 .unwrap()
7742 .as_mut_slice()
7743 .copy_from_slice(&[5, 6, 7, 8]);
7744 assert_eq!(parent.map().unwrap().as_slice(), &[1, 2, 3, 4, 5, 6, 7, 8]);
7745 }
7746 }
7747
7748 #[test]
7749 #[cfg(target_os = "linux")]
7750 fn dma_strided_subview_maps_padded_window() {
7751 // The strided-map path differs by backing: DMA maps through
7752 // `mmap_offset` + the `byte_size_override`, not the Mem `Arc` slice. A
7753 // padded sub-view of a DMA buffer must still expose its full
7754 // `row_stride × rows` window zero-copy at the view's offset (the GPU
7755 // batched-render-to-DMA case). Mirrors
7756 // `mem_strided_subview_maps_offset_and_byte_size` on a Dma parent.
7757 let _lock = crate::tests::fd_lock_shared();
7758 let parent = match Tensor::<u8>::new(&[2048], Some(TensorMemory::DmaBuf), None) {
7759 Ok(t) => t,
7760 Err(_) => {
7761 crate::test_support::report_skip("DMA not available");
7762 return;
7763 }
7764 };
7765 let mut view = parent.subview(128, &[8, 16]).unwrap();
7766 assert_eq!(view.plane_offset(), Some(128));
7767 view.set_row_stride_unchecked(32); // padded stride (> 16)
7768
7769 {
7770 let mut m = view.map().unwrap();
7771 let s = m.as_mut_slice();
7772 assert_eq!(s.len(), 256, "strided DMA map exposes stride(32) × rows(8)");
7773 s[0] = 0xAA; // row 0, col 0
7774 s[32] = 0xBB; // row 1, col 0 (one stride in)
7775 }
7776
7777 let p = parent.map().unwrap();
7778 let pb = p.as_slice();
7779 assert_eq!(pb[128], 0xAA, "row 0 writes at parent offset 128");
7780 assert_eq!(
7781 pb[128 + 32],
7782 0xBB,
7783 "row 1 writes at parent offset 128 + stride"
7784 );
7785 }
7786
7787 #[test]
7788 #[cfg(target_os = "linux")]
7789 fn view_single_row_snapshots_parent_stride() {
7790 // A single-row `view()` keeps a TIGHT `row_stride` for map-span safety,
7791 // but its `view_origin` snapshots the PARENT row stride — the GL backend
7792 // keys its EGLImage import/pitch on that snapshot (not the view's tight
7793 // stride), so single-row and multi-row sibling views collapse onto the
7794 // same parent import.
7795 let _lock = crate::tests::fd_lock_shared();
7796 // 8x4 RGBA with a padded 64-byte row stride (tight row = 8*4 = 32).
7797 let parent = match Tensor::<u8>::image_with_stride(
7798 8,
7799 4,
7800 PixelFormat::Rgba,
7801 64,
7802 Some(TensorMemory::DmaBuf),
7803 crate::CpuAccess::ReadWrite,
7804 ) {
7805 Ok(t) => t,
7806 Err(_) => {
7807 crate::test_support::report_skip("DMA not available");
7808 return;
7809 }
7810 };
7811 assert_eq!(parent.effective_row_stride(), Some(64));
7812 // Bottom row (y=3) at x>0 — the case the tight single-row stride guards.
7813 let row = parent.view(Region::new(2, 3, 4, 1)).unwrap();
7814 // The view's own stride is tight (4*4 = 16) so its strided map stays in
7815 // bounds; the GL-facing parent pitch (64) lives in `view_origin`.
7816 assert_eq!(row.effective_row_stride(), Some(16));
7817 let vo = row.view_origin().expect("a view carries a view_origin");
7818 assert_eq!(
7819 vo.parent_row_stride, 64,
7820 "GL keys/pitches a view on the parent stride, not its tight one"
7821 );
7822 // The tight stride keeps map() in-bounds for the bottom / x>0 single row.
7823 assert_eq!(row.map().unwrap().as_slice().len(), 16);
7824 }
7825
7826 #[test]
7827 fn test_mem_tensor() {
7828 let shape = vec![2, 3, 4];
7829 let tensor =
7830 MemTensor::<f32>::new(&shape, Some("test_tensor")).expect("Failed to create tensor");
7831 assert_eq!(tensor.shape(), &shape);
7832 assert_eq!(tensor.size(), 2 * 3 * 4 * std::mem::size_of::<f32>());
7833 assert_eq!(tensor.name(), "test_tensor");
7834
7835 {
7836 let mut tensor_map = tensor.map().expect("Failed to map memory");
7837 tensor_map.fill(42.0);
7838 assert!(tensor_map.iter().all(|&x| x == 42.0));
7839 }
7840
7841 let mut tensor = MemTensor::<u8>::new(&shape, None).expect("Failed to create tensor");
7842 assert_eq!(tensor.shape(), &shape);
7843 let new_shape = vec![3, 4, 4];
7844 assert!(
7845 tensor.reshape(&new_shape).is_err(),
7846 "Reshape should fail due to size mismatch"
7847 );
7848 assert_eq!(tensor.shape(), &shape, "Shape should remain unchanged");
7849
7850 let new_shape = vec![2, 3, 4];
7851 tensor.reshape(&new_shape).expect("Reshape should succeed");
7852 assert_eq!(
7853 tensor.shape(),
7854 &new_shape,
7855 "Shape should be updated after successful reshape"
7856 );
7857
7858 {
7859 let mut tensor_map = tensor.map().expect("Failed to map memory");
7860 tensor_map.fill(1);
7861 assert!(tensor_map.iter().all(|&x| x == 1));
7862 }
7863
7864 {
7865 let mut tensor_map = tensor.map().expect("Failed to map memory");
7866 tensor_map[2] = 42;
7867 assert_eq!(tensor_map[1], 1, "Value at index 1 should be 1");
7868 assert_eq!(tensor_map[2], 42, "Value at index 2 should be 42");
7869 }
7870 }
7871
7872 #[test]
7873 #[cfg(target_os = "linux")]
7874 fn test_dma_no_fd_leaks() {
7875 let _lock = fd_lock_exclusive();
7876 if !is_dma_available() {
7877 log::warn!(
7878 "SKIPPED: {} - DMA memory allocation not available (permission denied or no DMA-BUF support)",
7879 function!()
7880 );
7881 return;
7882 }
7883
7884 let baseline = open_fd_count(DMA_FD_TARGETS);
7885
7886 // Self-check — see `test_shm_no_fd_leaks`.
7887 {
7888 let probe = Tensor::<u8>::new(&[100, 100], Some(TensorMemory::DmaBuf), None)
7889 .expect("Failed to create tensor");
7890 assert!(
7891 open_fd_count(DMA_FD_TARGETS) > baseline,
7892 "a live DMA-BUF tensor must be visible to the dmabuf fd counter, \
7893 otherwise this test cannot detect a leak"
7894 );
7895 drop(probe);
7896 }
7897
7898 for _ in 0..100 {
7899 let tensor = Tensor::<u8>::new(&[100, 100], Some(TensorMemory::DmaBuf), None)
7900 .expect("Failed to create tensor");
7901 let mut map = tensor.map().unwrap();
7902 map.as_mut_slice().fill(233);
7903 }
7904
7905 let after = open_fd_count(DMA_FD_TARGETS);
7906 assert_eq!(
7907 baseline, after,
7908 "DMA-BUF file descriptor leak detected: {baseline} -> {after}"
7909 );
7910 }
7911
7912 #[test]
7913 #[cfg(target_os = "linux")]
7914 fn test_dma_from_fd_no_fd_leaks() {
7915 let _lock = fd_lock_exclusive();
7916 if !is_dma_available() {
7917 log::warn!(
7918 "SKIPPED: {} - DMA memory allocation not available (permission denied or no DMA-BUF support)",
7919 function!()
7920 );
7921 return;
7922 }
7923
7924 let baseline = open_fd_count(DMA_FD_TARGETS);
7925 let orig = Tensor::<u8>::new(&[100, 100], Some(TensorMemory::DmaBuf), None).unwrap();
7926
7927 // Self-check — see `test_shm_no_fd_leaks`.
7928 let held = open_fd_count(DMA_FD_TARGETS);
7929 assert!(
7930 held > baseline,
7931 "a live DMA-BUF tensor must be visible to the dmabuf fd counter, \
7932 otherwise this test cannot detect a leak"
7933 );
7934
7935 for _ in 0..100 {
7936 let tensor =
7937 Tensor::<u8>::from_fd(orig.clone_fd().unwrap(), orig.shape(), None).unwrap();
7938 assert_eq!(
7939 tensor.memory(),
7940 TensorMemory::DmaBuf,
7941 "DMA-BUF fd must import as Dma, not be silently downgraded"
7942 );
7943 let mut map = tensor.map().unwrap();
7944 map.as_mut_slice().fill(233);
7945 }
7946
7947 let after_loop = open_fd_count(DMA_FD_TARGETS);
7948 assert_eq!(
7949 held, after_loop,
7950 "DMA-BUF file descriptor leak detected: {held} -> {after_loop}"
7951 );
7952
7953 // And the source tensor's own fd goes back on drop.
7954 drop(orig);
7955 let after_drop = open_fd_count(DMA_FD_TARGETS);
7956 assert_eq!(
7957 baseline, after_drop,
7958 "dropping the source tensor must release its fd: {baseline} -> {after_drop}"
7959 );
7960 }
7961
7962 /// A filesystem magic must report its true 32-bit value regardless of
7963 /// how wide, and how signed, `fstatfs`'s `f_type` is on the target.
7964 ///
7965 /// `fs_type_t` is `__fsword_t` on Linux/gnu — `i64` on 64-bit but `i32`
7966 /// on 32-bit (armv7, i686, aarch64-ilp32) — `c_int` on uclibc, and
7967 /// unsigned on musl and s390x. Widening a *signed* 32-bit `f_type`
7968 /// sign-extends every magic with bit 31 set, so a naive widening cast
7969 /// reports e.g. `0xffffffff958458f6` for hugetlbfs. The docs tell
7970 /// callers to look the reported value up in `include/uapi/linux/magic.h`,
7971 /// so a sign-extended value is actively misleading.
7972 ///
7973 /// This cannot be reproduced on a 64-bit host — the defect only exists
7974 /// where `f_type` is a signed 32-bit type — so the test drives the
7975 /// conversion directly with the value such a target would produce.
7976 #[test]
7977 #[cfg(target_os = "linux")]
7978 fn test_fs_magic_normalizes_sign_extended_values() {
7979 // Declares this test as an fd-opener; see FD_LOCK.
7980 let _lock = crate::tests::fd_lock_shared();
7981 // Magics whose bit 31 is set. HUGETLBFS is the one that matters
7982 // most in practice: a MFD_HUGETLB memfd is the likeliest fd to land
7983 // in the UnknownBufferType arm.
7984 const HUGETLBFS_MAGIC: u32 = 0x9584_58f6;
7985 const F2FS_MAGIC: u32 = 0xf2f5_2010;
7986 const BTRFS_MAGIC: u32 = 0x9123_683e;
7987
7988 for magic in [HUGETLBFS_MAGIC, F2FS_MAGIC, BTRFS_MAGIC] {
7989 // 64-bit gnu: f_type is i64 and already holds the true value.
7990 assert_eq!(fs_magic(i64::from(magic)), magic);
7991
7992 // 32-bit gnu / uclibc: f_type is i32, so the value arrives
7993 // sign-extended once widened. It must still report its true
7994 // 32 bits.
7995 let sign_extended = i64::from(magic as i32);
7996 assert!(sign_extended < 0, "{magic:#x} should have bit 31 set");
7997 assert_eq!(
7998 fs_magic(sign_extended),
7999 magic,
8000 "{magic:#x} must survive a signed 32-bit f_type"
8001 );
8002 }
8003
8004 // The two magics we actually classify on are below 2^31, so they are
8005 // unaffected by signedness either way — this is why the DMA-vs-SHM
8006 // decision is correct on 32-bit even without this normalization.
8007 for magic in [DMA_BUF_MAGIC, TMPFS_MAGIC] {
8008 assert!(magic < 0x8000_0000);
8009 assert_eq!(fs_magic(i64::from(magic as i32)), magic);
8010 }
8011 }
8012
8013 /// A DMA-BUF fd must import as [`TensorMemory::DmaBuf`].
8014 ///
8015 /// Regression test for the `st_dev` minor-number classifier, which
8016 /// hardcoded `9 | 10` as "this is DMA". Those minors come from
8017 /// `get_anon_bdev()` and are assigned first-come-first-served at boot,
8018 /// so they vary by kernel build and boot order — a real DMA-BUF is
8019 /// minor 12 on x86 desktop and minor 8 on the ADIS Verdin. Both fell
8020 /// into the `_` arm and imported as SHM, which "works" (a DMA-BUF is
8021 /// mmap-able) but silently forfeits zero-copy.
8022 #[test]
8023 #[cfg(target_os = "linux")]
8024 fn test_from_fd_dma_imports_as_dma() {
8025 let _lock = crate::tests::fd_lock_shared();
8026 if !is_dma_available() {
8027 log::warn!("SKIPPED: {} - DMA memory not available", function!());
8028 return;
8029 }
8030
8031 let orig = Tensor::<u8>::new(&[64, 64], Some(TensorMemory::DmaBuf), None).unwrap();
8032 assert_eq!(orig.memory(), TensorMemory::DmaBuf);
8033
8034 let imported = Tensor::<u8>::from_fd(orig.clone_fd().unwrap(), orig.shape(), None).unwrap();
8035
8036 assert_eq!(
8037 imported.memory(),
8038 TensorMemory::DmaBuf,
8039 "a DMA-BUF fd must import as Dma"
8040 );
8041 }
8042
8043 /// A tmpfs/SHM fd must import as [`TensorMemory::Shm`].
8044 ///
8045 /// The companion to `test_from_fd_dma_imports_as_dma`: confirms the
8046 /// magic-based classifier identifies SHM positively rather than by
8047 /// falling through.
8048 #[test]
8049 #[cfg(target_os = "linux")]
8050 fn test_from_fd_shm_imports_as_shm() {
8051 let _lock = crate::tests::fd_lock_shared();
8052 if !is_shm_available() {
8053 log::warn!("SKIPPED: {} - SHM memory not available", function!());
8054 return;
8055 }
8056
8057 let orig = Tensor::<u8>::new(&[64, 64], Some(TensorMemory::Shm), None).unwrap();
8058 assert_eq!(orig.memory(), TensorMemory::Shm);
8059
8060 let imported = Tensor::<u8>::from_fd(orig.clone_fd().unwrap(), orig.shape(), None).unwrap();
8061
8062 assert_eq!(
8063 imported.memory(),
8064 TensorMemory::Shm,
8065 "a tmpfs fd must import as Shm"
8066 );
8067 }
8068
8069 /// An fd that is neither a DMA-BUF nor tmpfs must be rejected.
8070 ///
8071 /// A pipe is the convenient probe: it lives on `pipefs`, another
8072 /// `get_anon_bdev()` pseudo-filesystem, so it shares major 0 with the
8073 /// buffer types we do support and is only distinguishable by magic.
8074 /// Importing one as SHM is meaningless — `mmap` on a pipe fails — so
8075 /// the classifier must say "unknown" rather than guess.
8076 #[test]
8077 #[cfg(target_os = "linux")]
8078 fn test_from_fd_rejects_unknown_filesystem() {
8079 let _lock = crate::tests::fd_lock_shared();
8080
8081 let (read_end, _write_end) = nix::unistd::pipe().unwrap();
8082
8083 let result = Tensor::<u8>::from_fd(read_end, &[64], None);
8084
8085 match result {
8086 Err(Error::UnknownBufferType(magic)) => {
8087 // PIPEFS_MAGIC, from include/uapi/linux/magic.h
8088 assert_eq!(magic, 0x5049_5045, "expected PIPEFS_MAGIC");
8089 }
8090 other => panic!("expected UnknownBufferType for a pipe fd, got {other:?}"),
8091 }
8092 }
8093
8094 #[test]
8095 #[cfg(target_os = "linux")]
8096 fn test_shm_no_fd_leaks() {
8097 let _lock = fd_lock_exclusive();
8098 if !is_shm_available() {
8099 log::warn!(
8100 "SKIPPED: {} - SHM memory allocation not available (permission denied or no SHM support)",
8101 function!()
8102 );
8103 return;
8104 }
8105
8106 let baseline = open_fd_count(SHM_FD_TARGETS);
8107
8108 // Self-check: the counter must actually see a live shm tensor. If
8109 // it does not — a changed backing object, a different `/proc`
8110 // rendering — the leak assertion below would hold no matter how
8111 // much leaked, so fail loudly here instead.
8112 {
8113 let probe = Tensor::<u8>::new(&[100, 100], Some(TensorMemory::Shm), None)
8114 .expect("Failed to create tensor");
8115 assert!(
8116 open_fd_count(SHM_FD_TARGETS) > baseline,
8117 "a live shm tensor must be visible to the shm fd counter, \
8118 otherwise this test cannot detect a leak"
8119 );
8120 drop(probe);
8121 }
8122
8123 for _ in 0..100 {
8124 let tensor = Tensor::<u8>::new(&[100, 100], Some(TensorMemory::Shm), None)
8125 .expect("Failed to create tensor");
8126 let mut map = tensor.map().unwrap();
8127 map.as_mut_slice().fill(233);
8128 }
8129
8130 let after = open_fd_count(SHM_FD_TARGETS);
8131 assert_eq!(
8132 baseline, after,
8133 "shm file descriptor leak detected: {baseline} -> {after}"
8134 );
8135 }
8136
8137 #[test]
8138 #[cfg(target_os = "linux")]
8139 fn test_shm_from_fd_no_fd_leaks() {
8140 let _lock = fd_lock_exclusive();
8141 if !is_shm_available() {
8142 log::warn!(
8143 "SKIPPED: {} - SHM memory allocation not available (permission denied or no SHM support)",
8144 function!()
8145 );
8146 return;
8147 }
8148
8149 let baseline = open_fd_count(SHM_FD_TARGETS);
8150 let orig = Tensor::<u8>::new(&[100, 100], Some(TensorMemory::Shm), None).unwrap();
8151
8152 // Self-check — see `test_shm_no_fd_leaks`.
8153 let held = open_fd_count(SHM_FD_TARGETS);
8154 assert!(
8155 held > baseline,
8156 "a live shm tensor must be visible to the shm fd counter, \
8157 otherwise this test cannot detect a leak"
8158 );
8159
8160 for _ in 0..100 {
8161 let tensor =
8162 Tensor::<u8>::from_fd(orig.clone_fd().unwrap(), orig.shape(), None).unwrap();
8163 let mut map = tensor.map().unwrap();
8164 map.as_mut_slice().fill(233);
8165 }
8166
8167 let after_loop = open_fd_count(SHM_FD_TARGETS);
8168 assert_eq!(
8169 held, after_loop,
8170 "shm file descriptor leak detected: {held} -> {after_loop}"
8171 );
8172
8173 // And the source tensor's own fd goes back on drop.
8174 drop(orig);
8175 let after_drop = open_fd_count(SHM_FD_TARGETS);
8176 assert_eq!(
8177 baseline, after_drop,
8178 "dropping the source tensor must release its fd: {baseline} -> {after_drop}"
8179 );
8180 }
8181
8182 #[cfg(feature = "ndarray")]
8183 #[test]
8184 fn test_ndarray() {
8185 let _lock = crate::tests::fd_lock_shared();
8186 let shape = vec![2, 3, 4];
8187 let tensor = Tensor::<f32>::new(&shape, None, None).expect("Failed to create tensor");
8188
8189 let mut tensor_map = tensor.map().expect("Failed to map tensor memory");
8190 tensor_map.fill(1.0);
8191
8192 let view = tensor_map.view().expect("Failed to get ndarray view");
8193 assert_eq!(view.shape(), &[2, 3, 4]);
8194 assert!(view.iter().all(|&x| x == 1.0));
8195
8196 let mut view_mut = tensor_map
8197 .view_mut()
8198 .expect("Failed to get mutable ndarray view");
8199 view_mut[[0, 0, 0]] = 42.0;
8200 assert_eq!(view_mut[[0, 0, 0]], 42.0);
8201 assert_eq!(tensor_map[0], 42.0, "Value at index 0 should be 42");
8202 }
8203
8204 #[test]
8205 fn test_buffer_identity_derived_equality() {
8206 // Two identities derived from the same kind and key must be equal —
8207 // that is the whole point of `derived`: two independently-linked
8208 // copies of this crate, each importing the same underlying buffer,
8209 // have to agree on its id without talking to each other.
8210 let id1 = BufferIdentity::derived(IdentityKind::DmaBuf, 4242);
8211 let id2 = BufferIdentity::derived(IdentityKind::DmaBuf, 4242);
8212 assert_eq!(
8213 id1.id(),
8214 id2.id(),
8215 "Identities derived from the same kind and key should match"
8216 );
8217
8218 // A different key still yields a different id.
8219 let id3 = BufferIdentity::derived(IdentityKind::DmaBuf, 4243);
8220 assert_ne!(
8221 id1.id(),
8222 id3.id(),
8223 "Identities derived from different keys should differ"
8224 );
8225 }
8226
8227 #[test]
8228 fn test_buffer_identity_clone_shares_guard() {
8229 let id1 = BufferIdentity::derived(IdentityKind::DmaBuf, 4242);
8230 let weak = id1.weak();
8231 assert!(
8232 weak.upgrade().is_some(),
8233 "Weak should be alive while original exists"
8234 );
8235
8236 let id2 = id1.clone();
8237 assert_eq!(id1.id(), id2.id(), "Cloned identity should have same id");
8238
8239 drop(id1);
8240 assert!(
8241 weak.upgrade().is_some(),
8242 "Weak should still be alive (clone holds Arc)"
8243 );
8244
8245 drop(id2);
8246 assert!(
8247 weak.upgrade().is_none(),
8248 "Weak should be dead after all clones dropped"
8249 );
8250 }
8251
8252 #[test]
8253 fn test_tensor_buffer_identity() {
8254 let t1 = Tensor::<u8>::new(&[100], Some(TensorMemory::Mem), Some("t1")).unwrap();
8255 let t2 = Tensor::<u8>::new(&[100], Some(TensorMemory::Mem), Some("t2")).unwrap();
8256 assert_ne!(
8257 t1.buffer_identity().id(),
8258 t2.buffer_identity().id(),
8259 "Different tensors should have different buffer ids"
8260 );
8261 }
8262
8263 // ------------------------------------------------------------------------
8264 // Quantization — constructor validation + accessor correctness.
8265 // ------------------------------------------------------------------------
8266
8267 #[test]
8268 fn test_quantization_per_tensor_constructors() {
8269 let q = Quantization::per_tensor(0.1, -5);
8270 assert!(q.is_per_tensor());
8271 assert!(!q.is_per_channel());
8272 assert!(!q.is_symmetric());
8273 assert_eq!(q.scale(), &[0.1]);
8274 assert_eq!(q.zero_point(), Some(&[-5][..]));
8275
8276 let qs = Quantization::per_tensor_symmetric(0.05);
8277 assert!(qs.is_per_tensor());
8278 assert!(qs.is_symmetric());
8279 assert_eq!(qs.zero_point(), None);
8280 }
8281
8282 #[test]
8283 fn test_quantization_per_channel_constructors() {
8284 let q = Quantization::per_channel(vec![0.1, 0.2, 0.3], vec![0, -1, 1], 2).unwrap();
8285 assert!(q.is_per_channel());
8286 assert!(!q.is_symmetric());
8287 assert_eq!(q.axis(), Some(2));
8288 assert_eq!(q.scale().len(), 3);
8289
8290 let qs = Quantization::per_channel_symmetric(vec![0.054, 0.089, 0.195], 0).unwrap();
8291 assert!(qs.is_per_channel());
8292 assert!(qs.is_symmetric());
8293 assert_eq!(qs.axis(), Some(0));
8294 }
8295
8296 #[test]
8297 fn test_quantization_per_channel_length_mismatch_rejected() {
8298 // len(scales) != len(zero_points) → rejected at construction.
8299 let err = Quantization::per_channel(vec![0.1, 0.2], vec![0, 0, 0], 0).unwrap_err();
8300 assert!(matches!(err, Error::QuantizationInvalid { .. }));
8301 }
8302
8303 #[test]
8304 fn test_quantization_per_channel_empty_rejected() {
8305 let err = Quantization::per_channel_symmetric(vec![], 0).unwrap_err();
8306 assert!(matches!(err, Error::QuantizationInvalid { .. }));
8307 }
8308
8309 /// Constructors guard scale/zero_point length invariants, but
8310 /// `Quantization` is `Deserialize`, so malformed JSON (e.g. an
8311 /// empty `scale` array, or `zero_point` length that disagrees with
8312 /// `scale`) bypasses the constructor checks. `set_quantization`
8313 /// must reject these via `validate()` so they don't poison
8314 /// downstream `mode()` selection or per-channel kernel indexing.
8315 #[test]
8316 fn test_quantization_validate_rejects_malformed_deserialize() {
8317 let mut t = Tensor::<i8>::new(&[1, 1, 4], Some(TensorMemory::Mem), None).unwrap();
8318
8319 // Empty scale array: must be rejected.
8320 let q: Quantization = serde_json::from_str(r#"{"scale": []}"#).unwrap();
8321 assert!(matches!(
8322 t.set_quantization(q).unwrap_err(),
8323 Error::QuantizationInvalid { .. }
8324 ));
8325
8326 // Per-tensor with multi-element zero_point: must be rejected.
8327 let q: Quantization =
8328 serde_json::from_str(r#"{"scale": 0.1, "zero_point": [0, 0, 0]}"#).unwrap();
8329 assert!(matches!(
8330 t.set_quantization(q).unwrap_err(),
8331 Error::QuantizationInvalid { .. }
8332 ));
8333
8334 // Per-channel zero_point length != scale length: must be rejected.
8335 let q: Quantization = serde_json::from_str(
8336 r#"{"scale": [0.1, 0.2, 0.3, 0.4], "zero_point": [0, 0], "axis": 2}"#,
8337 )
8338 .unwrap();
8339 assert!(matches!(
8340 t.set_quantization(q).unwrap_err(),
8341 Error::QuantizationInvalid { .. }
8342 ));
8343 }
8344
8345 #[test]
8346 fn test_quantization_mode_dispatch() {
8347 let pt = Quantization::per_tensor(0.1, -5);
8348 assert!(matches!(
8349 pt.mode(),
8350 QuantMode::PerTensor { scale, zero_point } if scale == 0.1 && zero_point == -5
8351 ));
8352
8353 let pts = Quantization::per_tensor_symmetric(0.05);
8354 assert!(matches!(
8355 pts.mode(),
8356 QuantMode::PerTensorSymmetric { scale } if scale == 0.05
8357 ));
8358
8359 let pc = Quantization::per_channel(vec![0.1, 0.2], vec![0, -1], 2).unwrap();
8360 assert!(matches!(pc.mode(), QuantMode::PerChannel { axis: 2, .. }));
8361
8362 let pcs = Quantization::per_channel_symmetric(vec![0.1, 0.2], 0).unwrap();
8363 assert!(matches!(
8364 pcs.mode(),
8365 QuantMode::PerChannelSymmetric { axis: 0, .. }
8366 ));
8367 }
8368
8369 #[test]
8370 fn test_tensor_quantization_roundtrip_integer() {
8371 let mut t = Tensor::<i8>::new(&[2, 3, 4], Some(TensorMemory::Mem), None).unwrap();
8372 assert!(t.quantization().is_none());
8373 t.set_quantization(Quantization::per_tensor(0.1, -5))
8374 .unwrap();
8375 let q = t.quantization().unwrap();
8376 assert_eq!(q.scale(), &[0.1]);
8377 t.clear_quantization();
8378 assert!(t.quantization().is_none());
8379 }
8380
8381 #[test]
8382 fn test_tensor_with_quantization_builder() {
8383 let t = Tensor::<i8>::new(&[4, 4], Some(TensorMemory::Mem), None)
8384 .unwrap()
8385 .with_quantization(Quantization::per_tensor_symmetric(0.05))
8386 .unwrap();
8387 assert!(t.quantization().is_some());
8388 }
8389
8390 #[test]
8391 fn test_tensor_dyn_quantization_float_arm_returns_none() {
8392 let t = Tensor::<f32>::new(&[2, 2], Some(TensorMemory::Mem), None).unwrap();
8393 let td = TensorDyn::F32(t);
8394 assert!(td.quantization().is_none());
8395 }
8396
8397 #[test]
8398 fn test_tensor_dyn_set_quantization_float_arm_errors() {
8399 let t = Tensor::<f32>::new(&[2, 2], Some(TensorMemory::Mem), None).unwrap();
8400 let mut td = TensorDyn::F32(t);
8401 let err = td
8402 .set_quantization(Quantization::per_tensor(0.1, 0))
8403 .unwrap_err();
8404 // float path returns a QuantizationInvalid error.
8405 assert!(matches!(err, Error::QuantizationInvalid { .. }));
8406 }
8407
8408 /// Compile-time type gate — calling `Tensor::<f32>::quantization()` must
8409 /// fail to compile (the `IntegerType` trait bound is not satisfied by
8410 /// `f32`). This doctest anchors the invariant.
8411 ///
8412 /// ```compile_fail
8413 /// use edgefirst_tensor::{Tensor, TensorMemory};
8414 /// let t = Tensor::<f32>::new(&[2, 2], Some(TensorMemory::Mem), None).unwrap();
8415 /// let _ = t.quantization(); // compile error: f32 not IntegerType
8416 /// ```
8417 fn _compile_fail_doctest_anchor() {}
8418
8419 // Any test that cares about the fd count must grab it exclusively.
8420 // Any tests which modifies the fd count by opening or closing fds must grab it
8421 // shared.
8422 pub static FD_LOCK: RwLock<()> = RwLock::new(());
8423
8424 /// Take [`FD_LOCK`] shared, tolerating poisoning.
8425 ///
8426 /// A test that panics while holding the lock poisons it, and every later
8427 /// `.unwrap()` on it panics too — one real failure turns into a handful
8428 /// of unrelated ones that bury it. The guarded data is `()`, so a panic
8429 /// cannot have left an invariant broken and there is nothing to protect
8430 /// against by refusing the lock.
8431 pub fn fd_lock_shared() -> std::sync::RwLockReadGuard<'static, ()> {
8432 FD_LOCK.read().unwrap_or_else(PoisonError::into_inner)
8433 }
8434
8435 /// Take [`FD_LOCK`] exclusively. Poison-tolerant, see [`fd_lock_shared`].
8436 ///
8437 /// Linux-only, like all four of its callers: the fd-leak tests read fd
8438 /// state out of `/proc`, which no other platform provides. Ungated, it
8439 /// is dead code on macOS and the lint denies that.
8440 /// [`fd_lock_shared`] stays available everywhere — macOS tests declare
8441 /// fd activity too.
8442 #[cfg(target_os = "linux")]
8443 pub fn fd_lock_exclusive() -> std::sync::RwLockWriteGuard<'static, ()> {
8444 FD_LOCK.write().unwrap_or_else(PoisonError::into_inner)
8445 }
8446
8447 /// `/proc` fd-target fragments identifying a POSIX shared memory
8448 /// segment. `shm_open` then `shm_unlink` leaves the target rendered as
8449 /// `/dev/shm/<name> (deleted)`.
8450 #[cfg(target_os = "linux")]
8451 const SHM_FD_TARGETS: &[&str] = &["/dev/shm/"];
8452
8453 /// `/proc` fd-target fragments identifying the fds a DMA tensor owns.
8454 ///
8455 /// Both the exported buffer and the heap it came from are listed, so a
8456 /// leak of either is still caught — narrowing this to the buffer alone
8457 /// would silently stop covering leaked heap handles.
8458 #[cfg(target_os = "linux")]
8459 const DMA_FD_TARGETS: &[&str] = &["dmabuf", "dma_heap"];
8460
8461 /// How many of this process's open fds point at an object whose `/proc`
8462 /// target contains `needle`.
8463 ///
8464 /// The fd-leak tests deliberately do NOT compare total fd counts. Under
8465 /// `cargo test` every test in this crate runs as a thread in one
8466 /// process, and roughly twenty of them create shm or DMA tensors —
8467 /// while anything else may open a file, a socket, or a pipe — so an
8468 /// unrelated test opening or closing an fd inside the measurement
8469 /// window moved the total and failed the assertion. Worse, the panic
8470 /// poisoned [`FD_LOCK`] and took four more tests down with it, so the
8471 /// visible failure had little to do with the cause.
8472 ///
8473 /// Counting only fds that name the object under test makes the
8474 /// measurement immune to every unrelated fd, while still catching the
8475 /// leak each test exists to catch. (CI runs nextest, one process per
8476 /// test, which is why this only ever flaked when run locally.)
8477 #[cfg(target_os = "linux")]
8478 fn open_fd_count(needles: &[&str]) -> usize {
8479 let proc = procfs::process::Process::myself()
8480 .expect("Failed to get current process using /proc/self");
8481 proc.fd()
8482 .expect("Failed to list open file descriptors")
8483 .filter_map(Result::ok)
8484 .filter(|entry| {
8485 let target = format!("{:?}", entry.target);
8486 needles.iter().any(|needle| target.contains(needle))
8487 })
8488 .count()
8489 }
8490
8491 /// Test that DMA is NOT available on non-Linux platforms.
8492 /// This verifies the cross-platform behavior of is_dma_available().
8493 #[test]
8494 #[cfg(not(target_os = "linux"))]
8495 fn test_dma_not_available_on_non_linux() {
8496 assert!(
8497 !is_dma_available(),
8498 "DMA memory allocation should NOT be available on non-Linux platforms"
8499 );
8500 }
8501
8502 #[test]
8503 fn colorimetry_defaults_none_and_roundtrips_without_auto_fill() {
8504 use crate::{ColorEncoding, ColorRange, Colorimetry, PixelFormat, TensorMemory};
8505 let mut t = Tensor::<u8>::image(
8506 1280,
8507 720,
8508 PixelFormat::Nv12,
8509 Some(TensorMemory::Mem),
8510 crate::CpuAccess::ReadWrite,
8511 )
8512 .unwrap();
8513 assert_eq!(t.colorimetry(), None); // default undefined
8514 let c = Colorimetry::default()
8515 .with_encoding(ColorEncoding::Bt709)
8516 .with_range(ColorRange::Limited);
8517 t.set_colorimetry(Some(c));
8518 assert_eq!(t.colorimetry(), Some(c));
8519 // configure_image must NOT touch colorimetry
8520 t.configure_image(640, 480, PixelFormat::Grey).unwrap();
8521 assert_eq!(t.colorimetry(), Some(c));
8522 }
8523
8524 #[test]
8525 fn configure_image_within_capacity() {
8526 let mut t = Tensor::<u8>::image_with_capacity(
8527 640,
8528 480,
8529 PixelFormat::Rgb,
8530 None,
8531 crate::CpuAccess::ReadWrite,
8532 )
8533 .unwrap();
8534 t.configure_image(320, 240, PixelFormat::Nv12).unwrap();
8535 assert_eq!(t.format(), Some(PixelFormat::Nv12));
8536 assert_eq!(t.width(), Some(320));
8537 assert_eq!(t.height(), Some(240));
8538 assert_eq!(t.shape(), &[360, 320]); // 240*3/2
8539 }
8540
8541 #[test]
8542 fn configure_image_too_large_errors() {
8543 let mut t = Tensor::<u8>::image_with_capacity(
8544 64,
8545 64,
8546 PixelFormat::Grey,
8547 None,
8548 crate::CpuAccess::ReadWrite,
8549 )
8550 .unwrap();
8551 let err = t
8552 .configure_image(1920, 1080, PixelFormat::Nv12)
8553 .unwrap_err();
8554 assert!(matches!(err, Error::InsufficientCapacity { .. }));
8555 }
8556
8557 /// A reused max-sized IOSurface pool keeps its physical `bytesPerRow` when
8558 /// reconfigured to a smaller logical image (physical-grid / logical-ROI
8559 /// decoupling), instead of collapsing to the frame's natural row stride.
8560 #[test]
8561 #[cfg(target_os = "macos")]
8562 fn configure_image_preserves_iosurface_physical_stride() {
8563 // Declares this test as an fd-opener; see FD_LOCK.
8564 let _lock = crate::tests::fd_lock_shared();
8565 // Pool: GREY/R8 IOSurface 100 wide → bytesPerRow padded to 128.
8566 let mut pool = Tensor::<u8>::image(
8567 100,
8568 64,
8569 PixelFormat::Grey,
8570 Some(TensorMemory::DmaBuf),
8571 crate::CpuAccess::ReadWrite,
8572 )
8573 .unwrap();
8574 let pitch = pool.effective_row_stride().unwrap();
8575 assert!(
8576 pitch >= 128 && pitch.is_multiple_of(64),
8577 "padded bytesPerRow, got {pitch}"
8578 );
8579
8580 // Reconfigure to a smaller NV12 frame; the physical pitch must survive
8581 // (natural would be 32, but the surface stride is the 128-padded pitch).
8582 pool.configure_image(32, 16, PixelFormat::Nv12).unwrap();
8583 assert_eq!(pool.format(), Some(PixelFormat::Nv12));
8584 assert_eq!(pool.width(), Some(32));
8585 assert_eq!(pool.height(), Some(16));
8586 assert_eq!(
8587 pool.effective_row_stride(),
8588 Some(pitch),
8589 "configure_image must preserve the IOSurface physical bytesPerRow"
8590 );
8591
8592 // Reconfigure again to NV24 — pitch still preserved.
8593 pool.configure_image(32, 16, PixelFormat::Nv24).unwrap();
8594 assert_eq!(pool.effective_row_stride(), Some(pitch));
8595 }
8596
8597 /// `configure_image` on a Mem backing reconfigures to the format's
8598 /// **64-byte-aligned** row stride (the odd-dim contract: every image tensor
8599 /// carries a 64-aligned `row_stride`). For NV12 32×16 the minimum is
8600 /// `even(32)=32`, rounded up to the 64-byte alignment → 64. The capacity
8601 /// (64×64×4 RGBA = 16 KiB) easily holds the 24×64 = 1.5 KiB NV12 layout.
8602 #[test]
8603 fn configure_image_mem_aligns_stride() {
8604 let mut t = Tensor::<u8>::image_with_capacity(
8605 64,
8606 64,
8607 PixelFormat::Rgba,
8608 Some(TensorMemory::Mem),
8609 crate::CpuAccess::ReadWrite,
8610 )
8611 .unwrap();
8612 t.configure_image(32, 16, PixelFormat::Nv12).unwrap();
8613 let s = t.effective_row_stride().unwrap();
8614 assert_eq!(s % 64, 0, "stride must be 64-aligned");
8615 assert!(s >= 32, "stride must cover the even-width minimum");
8616 assert_eq!(s, 64);
8617 }
8618
8619 #[test]
8620 fn strided_mem_tensor_cpu_maps_full_padded_buffer() {
8621 // A packed RGBA image with row padding (GPU-pitch style): logical width
8622 // 8 px (32 B/row) but a 48-byte row stride. Over-allocate capacity (for
8623 // 16 px), narrow the logical width, then record the padded stride.
8624 // Previously `map()` rejected this on non-Linux with
8625 // "DMA backing is Linux-only"; HAL-owned Mem is now mappable.
8626 let mut t = Tensor::<u8>::image_with_capacity(
8627 16,
8628 3,
8629 PixelFormat::Rgba,
8630 Some(TensorMemory::Mem),
8631 crate::CpuAccess::ReadWrite,
8632 )
8633 .unwrap(); // capacity 3 × 16 × 4 = 192 B
8634 t.configure_image(8, 3, PixelFormat::Rgba).unwrap(); // logical [3, 8, 4] = 96 B
8635 t.set_row_stride(48).unwrap(); // padded stride (>= 32 B min)
8636
8637 let map = t.map().expect("strided Mem tensor should CPU-map");
8638 // Full padded buffer (stride 48 × 3 rows = 144 B), not the 96 B logical
8639 // view — callers iterate rows via `effective_row_stride()`.
8640 assert_eq!(map.as_slice().len(), 144);
8641 // Logical shape is still reported for shape-aware consumers.
8642 assert_eq!(map.shape(), &[3, 8, 4]);
8643 }
8644
8645 #[test]
8646 fn strided_mem_tensor_over_capacity_errors() {
8647 // Stride larger than the allocation: 64 B × 3 rows = 192 B > 96 B cap.
8648 let mut t = Tensor::<u8>::new(&[3, 8, 4], Some(TensorMemory::Mem), None).unwrap();
8649 t.set_format(PixelFormat::Rgba).unwrap();
8650 t.set_row_stride(64).unwrap();
8651 assert!(matches!(t.map(), Err(Error::InsufficientCapacity { .. })));
8652 }
8653
8654 /// Test that SHM memory allocation is available and usable on Unix systems.
8655 /// This is a basic functional test; Linux has additional FD leak tests using procfs.
8656 #[test]
8657 #[cfg(unix)]
8658 fn test_shm_available_and_usable() {
8659 // Declares this test as an fd-opener; see FD_LOCK.
8660 let _lock = crate::tests::fd_lock_shared();
8661 // "Unix" is not one platform here: Android is unix but its bionic libc
8662 // has no POSIX shared memory, so allocation legitimately reports
8663 // NotImplemented (see shm.rs module docs). Asserting availability
8664 // outright fails there for a documented reason.
8665 if !is_shm_available() {
8666 log::warn!(
8667 "SKIPPED: test_shm_available_and_usable - SHM allocation \
8668 unavailable on this platform"
8669 );
8670 return;
8671 }
8672
8673 // Create a tensor with SHM backing
8674 let tensor = Tensor::<u8>::new(&[100, 100], Some(TensorMemory::Shm), None)
8675 .expect("Failed to create SHM tensor");
8676
8677 // Verify we can map and write to it
8678 let mut map = tensor.map().expect("Failed to map SHM tensor");
8679 map.as_mut_slice().fill(0xAB);
8680
8681 // Verify the data was written correctly
8682 assert!(
8683 map.as_slice().iter().all(|&b| b == 0xAB),
8684 "SHM tensor data should be writable and readable"
8685 );
8686 }
8687
8688 // =========================================================================
8689 // packed_rgba16f_layout — host-runnable geometry unit tests (TDD)
8690 // =========================================================================
8691
8692 #[test]
8693 fn packed_rgba16f_layout_planar_rgb_f16() {
8694 let layout =
8695 packed_rgba16f_layout(PixelFormat::PlanarRgb, DType::F16, 640, 640).expect("Some");
8696 assert_eq!(layout.surface_w, 160);
8697 assert_eq!(layout.surface_h, 1920);
8698 assert_eq!(layout.bytes_per_texel, 8);
8699 assert_eq!(layout.pitch, 1280);
8700 }
8701
8702 #[test]
8703 fn packed_rgba16f_layout_planar_rgba_f16() {
8704 let layout =
8705 packed_rgba16f_layout(PixelFormat::PlanarRgba, DType::F16, 640, 640).expect("Some");
8706 assert_eq!(layout.surface_w, 160);
8707 assert_eq!(layout.surface_h, 2560); // 4 planes
8708 assert_eq!(layout.bytes_per_texel, 8);
8709 assert_eq!(layout.pitch, 1280);
8710 }
8711
8712 #[test]
8713 fn packed_rgba16f_layout_rejects_misaligned() {
8714 assert!(packed_rgba16f_layout(PixelFormat::PlanarRgb, DType::F16, 642, 640).is_none());
8715 }
8716
8717 #[test]
8718 fn packed_rgba16f_layout_rejects_non_f16() {
8719 // Non-F16 dtype with planar RGB
8720 assert!(packed_rgba16f_layout(PixelFormat::PlanarRgb, DType::U8, 640, 640).is_none());
8721 // Non-planar format with F32
8722 assert!(packed_rgba16f_layout(PixelFormat::Rgb, DType::F32, 640, 640).is_none());
8723 // Packed Rgba with F16 is not a planar format → None
8724 assert!(packed_rgba16f_layout(PixelFormat::Rgba, DType::F16, 640, 640).is_none());
8725 }
8726
8727 #[test]
8728 fn cuda_map_fast_fails_to_none_without_handle() {
8729 let t = Tensor::<f32>::new(&[4], Some(TensorMemory::Mem), None).unwrap();
8730 assert!(t.cuda().is_none());
8731 assert!(t.cuda_map().is_none()); // pure local check, no GL routing
8732 }
8733
8734 #[test]
8735 fn cuda_returns_none_without_handle() {
8736 // A plain Mem-backed tensor has no CUDA handle attached.
8737 let t = Tensor::<f32>::new(&[2, 2], Some(TensorMemory::Mem), None).unwrap();
8738 assert!(t.cuda().is_none(), "no CUDA handle on a Mem tensor");
8739 assert!(t.cuda_map().is_none(), "fast-fail map → None");
8740 }
8741
8742 #[test]
8743 fn cuda_map_then_host_map_fallback() {
8744 // The documented client pattern: try cuda_map() first; when it is None
8745 // (no CUDA handle — the case for a plain Mem tensor), fall back to map().
8746 let t = Tensor::<f32>::new(&[2, 2], Some(TensorMemory::Mem), None).unwrap();
8747 // Bind to a named variable so the CudaMap guard (and its borrow of `t`)
8748 // is dropped at the end of this statement, before the else branch borrows `t` again.
8749 let cuda = t.cuda_map();
8750 if let Some(_c) = cuda {
8751 // On a CUDA-registered tensor we'd use the device ptr here.
8752 unreachable!("a Mem tensor has no CUDA handle");
8753 } else {
8754 let host = t.map().expect("host map fallback must succeed");
8755 // TensorMapTrait::len() returns the element count (not bytes).
8756 assert_eq!(host.len(), 4); // 2*2 f32 elements
8757 }
8758 }
8759
8760 // -------------------------------------------------------------------------
8761 // Tensor::from_foreign — public API tests at the Tensor<T> layer.
8762 //
8763 // The low-level MemTensor::from_foreign mechanics (owner-drop, view sharing)
8764 // are covered in mem.rs. These tests exercise the Tensor<T> guard paths
8765 // (null ptr, empty shape, size overflow) and the basic wrap+readback
8766 // contract, confirming the public unsafe API wires through correctly.
8767 // -------------------------------------------------------------------------
8768
8769 #[test]
8770 fn from_foreign_valid_wrap_and_readback() {
8771 // The canonical CUDA zero-copy shape: wrap a caller allocation as a
8772 // Mem tensor and verify the tensor reads the exact same bytes.
8773 let mut buf: Vec<f32> = vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0];
8774 let ptr = buf.as_mut_ptr();
8775 let t = unsafe { Tensor::<f32>::from_foreign(ptr, &[2, 3], None, Some("test_foreign")) }
8776 .expect("valid from_foreign must succeed");
8777 assert_eq!(t.shape(), &[2, 3]);
8778 assert_eq!(t.memory(), TensorMemory::Mem);
8779 assert_eq!(t.name(), "test_foreign");
8780 let m = t.map().unwrap();
8781 // The tensor is a zero-copy borrow — it sees the caller's data.
8782 assert_eq!(m.as_slice(), &[1.0_f32, 2.0, 3.0, 4.0, 5.0, 6.0]);
8783 }
8784
8785 #[test]
8786 fn from_foreign_write_visible_in_caller_allocation() {
8787 // Writes through the tensor's map land in the caller's buffer (zero-copy).
8788 let mut buf: Vec<u8> = vec![0u8; 6];
8789 let ptr = buf.as_mut_ptr();
8790 let t = unsafe { Tensor::<u8>::from_foreign(ptr, &[2, 3], None, None) }.unwrap();
8791 {
8792 let mut m = t.map().unwrap();
8793 m.as_mut_slice().copy_from_slice(&[10, 20, 30, 40, 50, 60]);
8794 }
8795 drop(t);
8796 // Mutations are visible in the original Vec — same physical buffer.
8797 assert_eq!(buf, vec![10, 20, 30, 40, 50, 60]);
8798 }
8799
8800 #[test]
8801 fn from_foreign_rejects_null_ptr() {
8802 let err = unsafe { Tensor::<u8>::from_foreign(std::ptr::null_mut(), &[4], None, None) }
8803 .unwrap_err();
8804 assert!(
8805 matches!(err, Error::InvalidArgument(ref m) if m.contains("non-null")),
8806 "expected InvalidArgument(non-null), got {err:?}"
8807 );
8808 }
8809
8810 #[test]
8811 fn from_foreign_rejects_empty_shape() {
8812 let mut dummy: u8 = 0;
8813 let err = unsafe { Tensor::<u8>::from_foreign(&mut dummy, &[], None, None) }.unwrap_err();
8814 assert!(
8815 matches!(err, Error::InvalidSize(0)),
8816 "expected InvalidSize(0) for empty shape, got {err:?}"
8817 );
8818 }
8819
8820 #[test]
8821 fn from_foreign_rejects_overflow_shape() {
8822 // Two dimensions whose product overflows usize — the overflow guard must
8823 // fire before any pointer arithmetic is attempted.
8824 let mut dummy: u8 = 0;
8825 let huge = [usize::MAX / 2 + 1, 2];
8826 let err = unsafe { Tensor::<u8>::from_foreign(&mut dummy, &huge, None, None) }.unwrap_err();
8827 assert!(
8828 matches!(err, Error::InvalidArgument(ref m) if m.contains("overflow")),
8829 "expected InvalidArgument(overflow), got {err:?}"
8830 );
8831 }
8832
8833 #[test]
8834 fn from_foreign_owner_keeps_allocation_alive() {
8835 // When `owner` is `Some`, dropping the Tensor must not free the backing
8836 // before the owner is also gone — the owner's Drop fires on last ref.
8837 use std::sync::atomic::{AtomicBool, Ordering};
8838 let flag = std::sync::Arc::new(AtomicBool::new(false));
8839 let flag2 = flag.clone();
8840 struct Guard(std::sync::Arc<AtomicBool>);
8841 impl Drop for Guard {
8842 fn drop(&mut self) {
8843 self.0.store(true, Ordering::SeqCst);
8844 }
8845 }
8846 let mut buf: Vec<u32> = vec![42u32; 4];
8847 let ptr = buf.as_mut_ptr();
8848 let owner: ForeignOwner = Box::new(Guard(flag2));
8849 let t = unsafe { Tensor::<u32>::from_foreign(ptr, &[4], Some(owner), None) }.unwrap();
8850 // Map co-owns the backing Arc; the owner must stay alive while the map lives.
8851 let m = t.map().unwrap();
8852 assert_eq!(m.as_slice()[0], 42);
8853 drop(t); // tensor dropped while map is still live
8854 assert!(
8855 !flag.load(Ordering::SeqCst),
8856 "owner must not drop while a map shares the backing"
8857 );
8858 drop(m);
8859 assert!(
8860 flag.load(Ordering::SeqCst),
8861 "owner Drop must fire when the last Arc reference is released"
8862 );
8863 }
8864}