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 [`CpuAccess`] declaration (`ReadWrite` is the default). Hardware access — GPU, NPU,
40 ISP, codec — is always implied; what the CPU intends is the one thing the allocator cannot guess.
41 A CPU map, pin or sync beyond the declaration fails with [`Error::CpuAccessDenied`].
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// Static backend only: reached solely through `TensorStorage`'s AHardwareBuffer variant.
47#[cfg(all(target_os = "android", not(edgefirst_tensor_dynamic)))]
48mod ahardwarebuffer;
49// Pure AHardwareBuffer layout logic (format table, descriptor geometry,
50// overflow-checked shape math) — cfg-free so it compiles and unit-tests
51// on every host; the android module above consumes it.
52#[allow(dead_code)]
53mod ahardwarebuffer_layout;
54pub mod colorimetry;
55#[cfg(feature = "tracing")]
56pub mod trace;
57
58pub mod covguard;
59mod cuda;
60mod detect;
61// Static backend only: the dynamic backend reaches DMA-BUF fds through `ef_tensor_plane_at`.
62#[cfg(all(target_os = "linux", not(edgefirst_tensor_dynamic)))]
63mod dma;
64#[cfg(all(target_os = "linux", not(edgefirst_tensor_dynamic)))]
65mod dmabuf;
66mod error;
67mod format;
68// Layout table is backend-agnostic (pure DXGI-format math), same reasoning
69// as `iosurface_layout` below: compiled on every platform so it is
70// unit-tested on every CI lane. The allocator that consumes it
71// (`crate::d3d11`) is Windows-only.
72#[cfg(target_os = "windows")]
73pub mod d3d11;
74pub mod d3d11_layout;
75// Layout table is backend-agnostic (pure FourCC math). Storage
76// (`IoSurfaceTensor`) is `static`-only; `dynamic` reads IOSurface handles
77// through `ef_tensor_*`.
78#[cfg(any(target_os = "macos", target_os = "ios"))]
79mod iosurface_layout;
80// `static`-only, same reasoning as `ahardwarebuffer` above: reached only
81// through `TensorStorage`'s `IoSurface` variant.
82#[cfg(all(
83 any(target_os = "macos", target_os = "ios"),
84 not(edgefirst_tensor_dynamic)
85))]
86mod iosurface;
87pub mod lens;
88// `static`-only: `MemTensor` backs `TensorStorage::Mem`, `impl<T> Tensor<T>`
89// only. `dynamic`'s host-memory allocation goes through `ef_tensor_new`/
90// `ef_tensor_builder_alloc` in `libedgefirst_tensor.so` instead -- there is
91// no local `MemTensor` under `dynamic` to be dead code, so the whole module
92// gates rather than an internal `#[allow]`.
93#[cfg(not(edgefirst_tensor_dynamic))]
94mod mem;
95mod pbo;
96pub mod pin;
97mod preprocess;
98pub mod protocol;
99// `static`-only, same reasoning as `mem` above: `ShmTensor` backs
100// `TensorStorage::Shm`, `impl<T> Tensor<T>` only.
101#[cfg(all(unix, not(edgefirst_tensor_dynamic)))]
102mod shm;
103mod tensor_dyn;
104#[cfg(test)]
105mod test_support;
106pub mod view;
107mod vocabulary;
108pub use colorimetry::{
109 ColorEncoding, ColorRange, ColorSpace, ColorTransfer, Colorimetry, MatrixWeights, RangeScaling,
110};
111
112/// Retained constructor: installs the coverage flush-on-abort handler for this
113/// crate's instrumented test binary. See `covguard`. Only present under
114/// coverage on Linux (`.init_array` is ELF-only; the i.MX flush is Linux-only).
115#[cfg(all(coverage, target_os = "linux"))]
116#[used]
117#[link_section = ".init_array"]
118static __EDGEFIRST_COV_INSTALL: extern "C" fn() = {
119 extern "C" fn ctor() {
120 crate::covguard::install();
121 }
122 ctor
123};
124
125// Backing tensor/map types are internal implementation details: callers
126// allocate `Tensor<T>` / `TensorDyn` and map them, never naming the per-memory
127// backing types directly. They are `pub(crate)` so they stay nameable for the
128// `TensorStorage` / `TensorMap` enums without leaking into the public API.
129// Exceptions kept public: `Pbo*` is a GL extension point implemented by the
130// image crate, and `image_iosurface_layout` is a public helper.
131#[cfg(all(target_os = "android", not(edgefirst_tensor_dynamic)))]
132pub use crate::ahardwarebuffer::image_ahardwarebuffer_layout;
133#[cfg(all(target_os = "android", not(edgefirst_tensor_dynamic)))]
134pub(crate) use crate::ahardwarebuffer::AHardwareBufferTensor;
135#[cfg(all(target_os = "linux", not(edgefirst_tensor_dynamic)))]
136pub(crate) use crate::dma::DmaTensor;
137#[cfg(all(
138 any(target_os = "macos", target_os = "ios"),
139 not(edgefirst_tensor_dynamic)
140))]
141pub(crate) use crate::iosurface::IoSurfaceTensor;
142// The two geometry helpers are compiled for both backends: three of the four
143// surfaces that read a raw texture's geometry (the Python bindings, and the
144// protocol and blob imports) link `dynamic`, and only the C ABI links
145// `static`. See `crate::d3d11::geometry`.
146#[cfg(target_os = "windows")]
147pub use crate::d3d11::geometry::{d3d11_shared_handle_geometry, d3d11_texture_geometry};
148#[cfg(any(target_os = "macos", target_os = "ios"))]
149pub use crate::iosurface_layout::image_iosurface_layout;
150#[cfg(not(edgefirst_tensor_dynamic))]
151pub(crate) use crate::mem::MemTensor;
152pub use crate::pbo::{
153 client_state_pbo_ops, read_pbo_vtable_parts, EfClientState, EfPboMapFn, EfPboUnmapFn,
154 PboClientParts, PboMapping, PboOps, PboOpsVtableParts, PboTensor,
155};
156#[cfg(all(unix, not(edgefirst_tensor_dynamic)))]
157pub(crate) use crate::shm::ShmTensor;
158pub use cuda::{
159 client_state_cuda_ops, gl_map_resource, gl_register_buffer, gl_unmap_resource,
160 gl_unregister_resource, is_cuda_available, memcpy_device_to_host, memcpy_host_to_device,
161 runtime_path as cuda_runtime_path, stream_create, stream_destroy, stream_synchronize,
162 CudaGlOps, CudaHandle, CudaMap, CudaStream,
163};
164pub use detect::{unletter_norm, BoundingBox, DetectBox, ProtoData, ProtoLayout, Segmentation};
165pub use error::{Error, Result};
166pub use preprocess::{
167 broadcast_rgb, fit_code, interpolation_code, letterbox_placement_code, letterbox_rect,
168 normalization_code, pad_space_code, ContentRegion, ContentRegionSource, Fit, FitMode,
169 Interpolation, LetterboxPlacement, Normalization, NormalizationPreset, PadSpace,
170 PreprocessSpec,
171};
172// Static-backend-implementation code: backs `tensor-capi`'s own
173// `ef_tensor_export`/`ef_tensor_import` (the static cdylib's `TensorDyn`
174// serialization), the one caller in the whole workspace. It calls
175// `TensorDyn::new`/`from_fd` and reads fields (`.quantization()`,
176// `.colorimetry()`, `.pin_host()`, ...) that only the `static` backend's
177// richer `TensorDyn` surface has, so it never compiled under `dynamic` in
178// fact even though nothing gated it -- see `PRIMITIVE-INVENTORY.md`'s
179// "structural findings".
180#[cfg(not(edgefirst_tensor_dynamic))]
181pub mod blob;
182#[cfg(not(edgefirst_tensor_dynamic))]
183pub use blob::{
184 BlobError, BlobHeader, BlobPlane, BlobStrings, BlobView, Regions, TransportMode, HEADER_LEN,
185 PLANE_RECORD_LEN,
186};
187pub use format::{ChromaLayout, PixelFormat, PixelLayout, PlaneGeometry};
188pub use lens::Element;
189use num_traits::Num;
190pub use pin::HostPin;
191pub use protocol::{
192 dtype as tensor_dtype, dtype_to_dtype, format as tensor_format, format_from_code,
193 kind as tensor_kind, SendPtr, TensorDesc, ABI_VERSION,
194};
195use serde::{Deserialize, Serialize};
196#[cfg(unix)]
197use std::os::fd::OwnedFd;
198use std::{
199 fmt,
200 sync::{
201 atomic::{AtomicU64, Ordering},
202 Arc, Weak,
203 },
204};
205// `static`'s `Tensor<T>` is declared directly above, in this file; there is
206// nothing to re-export for it. `dynamic`'s lives in `tensor_dyn` (see that
207// module's docs) and is re-exported here under the same crate-root name so
208// both backends present the same `edgefirst_tensor::Tensor` to callers.
209#[cfg(edgefirst_tensor_dynamic)]
210pub use tensor_dyn::Tensor;
211pub use tensor_dyn::{BorrowedTensor, Raw, TensorDyn};
212pub use view::HostView;
213pub use vocabulary::vocabulary_demo;
214
215/// Opaque keep-alive handle for a foreign-memory tensor (see
216/// [`Tensor::from_foreign`] / [`TensorDyn::from_foreign_ptr`]).
217///
218/// The HAL borrows the foreign buffer without owning it; this handle co-owns
219/// the *source* so the borrowed memory stays valid for the tensor's life. Its
220/// `Drop` releases the source — e.g. a small struct that calls `cudaFreeHost`,
221/// or a `Py<PyAny>` that decrements a NumPy array's refcount. Wrapping it in an
222/// `Arc` (then boxing each clone) makes the release fire exactly once, after
223/// the last sharing tensor/view/map drops, regardless of drop order.
224pub type ForeignOwner = Box<dyn std::any::Any + Send + Sync>;
225
226/// Re-export of `half::f16` so downstream crates can write
227/// `Tensor::<edgefirst_tensor::f16>::from_iosurface(…)` without
228/// adding `half` to their own dependency list. The version stays in
229/// lockstep with the `half` workspace dep.
230pub use half::f16;
231
232// =============================================================================
233// RGBA16F packed-layout geometry — single source of truth
234//
235// A `PlanarRgb` [3,H,W] or `PlanarRgba` [4,H,W] f16 tensor is represented
236// on the GPU as an RGBA16F surface (the only float format accepted by the
237// ANGLE IOSurface extension). Four contiguous f16 elements are packed into
238// each 8-byte RGBA16F texel, yielding a `(W/4, C*H)` surface.
239//
240// All call sites that need these dimensions must use `packed_rgba16f_layout`
241// so the rule lives in exactly one place. Currently consumed by:
242// - `crates/tensor/src/iosurface.rs` `new_image` (macOS IOSurface alloc)
243// - `crates/image/src/gl/iosurface_import.rs` (macOS GL IOSurface import)
244// - `crates/image/src/gl/processor/float.rs` (Linux GL float render — PBO
245// readback and DMA-BUF, also via the `dma_f16_packed_layout` wrapper)
246// =============================================================================
247
248/// Geometry of the RGBA16F-packed surface backing a planar F16 image tensor.
249///
250/// ANGLE only supports one float `(type, internal_format)` pair for IOSurface
251/// import: `(GL_HALF_FLOAT, GL_RGBA)` = RGBA16F (8 bytes/texel). To map a
252/// `[C, H, W]` f16 planar tensor onto such a surface, 4 contiguous f16
253/// elements are packed into each RGBA16F texel, yielding a surface of
254/// `(W/4, C*H)` texels at 8 bytes/texel. The byte stream is identical to a
255/// (nonexistent) R16F `(W, C*H)` surface and can be consumed as `&[f16]`
256/// with shape `[1, C, H, W]` without rearrangement.
257///
258/// Obtain via [`packed_rgba16f_layout`] — never construct directly.
259#[derive(Debug, Clone, Copy, PartialEq, Eq)]
260pub struct PackedRgba16fLayout {
261 /// Surface width in texels (`width / 4`).
262 pub surface_w: usize,
263 /// Surface height in texels (`planes * height`).
264 pub surface_h: usize,
265 /// Bytes per RGBA16F texel (always 8).
266 pub bytes_per_texel: usize,
267 /// Row pitch in bytes (`surface_w * 8`).
268 pub pitch: usize,
269}
270
271/// Canonical geometry for the RGBA16F-packed surface backing a planar F16
272/// image tensor.
273///
274/// Returns `Some(layout)` only when **all** of the following hold:
275///
276/// - `dtype == DType::F16`
277/// - `format` is `PixelFormat::PlanarRgb` (3 planes) or
278/// `PixelFormat::PlanarRgba` (4 planes)
279/// - `width % 4 == 0`
280///
281/// Returns `None` for any other `(format, dtype)` combination, misaligned
282/// width, or when the surface geometry would overflow `usize` — callers
283/// must fall back to a non-packed path or return a context-appropriate
284/// error.
285///
286/// # Examples
287///
288/// ```rust
289/// use edgefirst_tensor::{packed_rgba16f_layout, PixelFormat, DType};
290///
291/// let layout = packed_rgba16f_layout(PixelFormat::PlanarRgb, DType::F16, 640, 480).unwrap();
292/// assert_eq!(layout.surface_w, 160);
293/// assert_eq!(layout.surface_h, 1440);
294/// assert_eq!(layout.bytes_per_texel, 8);
295/// assert_eq!(layout.pitch, 1280);
296/// ```
297pub fn packed_rgba16f_layout(
298 format: PixelFormat,
299 dtype: DType,
300 width: usize,
301 height: usize,
302) -> Option<PackedRgba16fLayout> {
303 if dtype != DType::F16 {
304 return None;
305 }
306 let planes: usize = match format {
307 PixelFormat::PlanarRgb => 3,
308 PixelFormat::PlanarRgba => 4,
309 _ => return None,
310 };
311 if !width.is_multiple_of(4) {
312 return None;
313 }
314 let surface_w = width / 4;
315 // Checked arithmetic: a degenerate (height, width) could otherwise wrap
316 // and yield an under-sized layout, which downstream allocators trust for
317 // GPU/CPU buffer sizing. Overflow → None (handled like any other
318 // unsupported geometry).
319 let surface_h = planes.checked_mul(height)?;
320 let bytes_per_texel = 8;
321 let pitch = surface_w.checked_mul(bytes_per_texel)?;
322 Some(PackedRgba16fLayout {
323 surface_w,
324 surface_h,
325 bytes_per_texel,
326 pitch,
327 })
328}
329
330/// Geometry of the RGBA8888-packed surface backing a packed RGB u8/i8 image
331/// tensor.
332///
333/// GPUs have no 3-channel renderable format, so the GL engine's two-pass
334/// packed-RGB shader writes the tight `[H, W, 3]` byte stream into an
335/// RGBA8888 surface: each texel carries 4 consecutive RGB bytes, giving a
336/// `(W*3/4, H)` surface at 4 bytes/texel whose rows are byte-identical to
337/// tight RGB — consumable flat as `[H, W, 3]` with no rearrangement (the
338/// u8/i8 analog of [`packed_rgba16f_layout`]; i8 shares the layout since
339/// INT8 quantization is a per-byte `^0x80` bias, not a format change).
340///
341/// Returns `Some(layout)` only when `width % 4 == 0` (so `W*3` bytes divide
342/// into whole texels) and the geometry does not overflow `usize`.
343///
344/// # Examples
345///
346/// ```rust
347/// use edgefirst_tensor::packed_rgb888_layout;
348///
349/// let layout = packed_rgb888_layout(640, 480).unwrap();
350/// assert_eq!(layout.surface_w, 480); // 640*3/4
351/// assert_eq!(layout.surface_h, 480);
352/// assert_eq!(layout.bytes_per_texel, 4);
353/// assert_eq!(layout.pitch, 1920); // 640*3
354/// assert!(packed_rgb888_layout(641, 480).is_none());
355/// ```
356pub fn packed_rgb888_layout(width: usize, height: usize) -> Option<PackedRgb888Layout> {
357 if !width.is_multiple_of(4) {
358 return None;
359 }
360 let row_bytes = width.checked_mul(3)?;
361 let surface_w = row_bytes / 4;
362 Some(PackedRgb888Layout {
363 surface_w,
364 surface_h: height,
365 bytes_per_texel: 4,
366 pitch: row_bytes,
367 })
368}
369
370/// Geometry of the RGBA8888 surface backing a packed RGB u8/i8 image tensor.
371///
372/// Obtain via [`packed_rgb888_layout`] — never construct directly.
373#[derive(Debug, Clone, Copy, PartialEq, Eq)]
374pub struct PackedRgb888Layout {
375 /// Surface width in texels (`width * 3 / 4`).
376 pub surface_w: usize,
377 /// Surface height in texels (`height`).
378 pub surface_h: usize,
379 /// Bytes per RGBA8888 texel (always 4).
380 pub bytes_per_texel: usize,
381 /// Row pitch in bytes (`surface_w * 4` = `width * 3`).
382 pub pitch: usize,
383}
384
385/// The logical `(width, height)` an **allocation** shape describes for
386/// `format` -- the inverse of [`PixelFormat::allocation_shape`].
387///
388/// Used wherever a tensor arrives carrying a shape and a format but not its
389/// dimensions: the C constructors, the Python bindings and the cross-package
390/// protocol import.
391///
392/// # The input is the allocation shape, not the addressing shape
393///
394/// [`PixelFormat::allocation_shape`] is what [`Tensor::image`] allocates from
395/// and what `descriptor_pinned` emits from `shape()`, so that is what this
396/// inverts. [`PixelFormat::addressing_shape`] is a different, smaller grid,
397/// and for a semi-planar format the two collide: NV12's addressing shape is
398/// `[h, w]` and its allocation shape is `[combined_h, w]`, both rank 2. Handed
399/// an addressing shape this cannot tell -- NV12 `[480, 640]` reads as a
400/// combined height of 480 and answers `(640, 320)`, silently. Nothing in the
401/// shape distinguishes them, so callers must pass the allocation shape.
402///
403/// A D3D11 import is the one place this cannot bite: the constructors check
404/// the derived size against the texture description, so a mismatch fails with
405/// [`Error::InvalidArgument`] naming both. On Windows, prefer
406/// `d3d11_texture_geometry` and `d3d11_shared_handle_geometry`, which read the
407/// dimensions off the texture itself and leave the shape as a consistency
408/// check -- except for a semi-planar texture, whose width is its row stride
409/// rather than the image width, so those helpers take the width from the shape
410/// they are given and refuse to answer without one.
411///
412/// The rank is part of the format's contract, so a shape of the wrong rank is
413/// an error rather than an interpretation of whatever dimensions are present:
414///
415/// | Layout | Shape | Result |
416/// |--------|-------|--------|
417/// | Packed | `[h, w, c]` | `(shape[1], shape[0])` |
418/// | Planar | `[c, h, w]` | `(shape[2], shape[1])` |
419/// | SemiPlanar | `[combined_h, w]` | `(shape[1], image height)` |
420///
421/// A semi-planar shape carries the *combined* plane height, so the logical
422/// height is recovered through
423/// [`PixelFormat::image_height_from_combined`], which rejects a combined
424/// height no image height produces (721 for NV12, say) instead of rounding to
425/// a nearby one.
426///
427/// # Errors
428///
429/// [`Error::InvalidShape`] for a shape of the wrong rank, or a semi-planar
430/// combined height with no image height behind it.
431///
432/// # Example
433///
434/// ```
435/// use edgefirst_tensor::{image_dims_from_shape, PixelFormat};
436/// assert_eq!(image_dims_from_shape(PixelFormat::Rgba, &[480, 640, 4]).unwrap(), (640, 480));
437/// assert_eq!(image_dims_from_shape(PixelFormat::PlanarRgb, &[3, 480, 640]).unwrap(), (640, 480));
438/// assert_eq!(image_dims_from_shape(PixelFormat::Nv12, &[722, 640]).unwrap(), (640, 481));
439/// assert!(image_dims_from_shape(PixelFormat::Nv12, &[721, 640]).is_err());
440/// ```
441pub fn image_dims_from_shape(format: PixelFormat, shape: &[usize]) -> Result<(usize, usize)> {
442 let wrong_rank = |expected: usize| {
443 Error::InvalidShape(format!(
444 "image_dims_from_shape: {format:?} describes a rank-{expected} allocation \
445 shape, got rank {} ({shape:?})",
446 shape.len()
447 ))
448 };
449 match format.layout() {
450 PixelLayout::Packed => {
451 if shape.len() != 3 {
452 return Err(wrong_rank(3));
453 }
454 Ok((shape[1], shape[0]))
455 }
456 PixelLayout::Planar => {
457 if shape.len() != 3 {
458 return Err(wrong_rank(3));
459 }
460 Ok((shape[2], shape[1]))
461 }
462 PixelLayout::SemiPlanar => {
463 if shape.len() != 2 {
464 return Err(wrong_rank(2));
465 }
466 let height = format.image_height_from_combined(shape[0]).ok_or_else(|| {
467 Error::InvalidShape(format!(
468 "image_dims_from_shape: no {format:?} image height has combined plane \
469 height {} (shape {shape:?})",
470 shape[0]
471 ))
472 })?;
473 Ok((shape[1], height))
474 }
475 }
476}
477
478/// Per-plane DMA-BUF descriptor for external buffer import.
479///
480/// Owns a duplicated file descriptor plus optional stride and offset metadata.
481/// The fd is duplicated eagerly in [`new()`](Self::new) so that a bad fd is
482/// caught immediately. `import_image` consumes the descriptor and takes
483/// ownership of the duped fd — no further cleanup is needed by the caller.
484///
485/// # Examples
486///
487/// ```rust,no_run
488/// use edgefirst_tensor::PlaneDescriptor;
489/// use std::os::fd::BorrowedFd;
490///
491/// // SAFETY: fd 42 is hypothetical; real code must pass a valid fd.
492/// let pd = unsafe { PlaneDescriptor::new(BorrowedFd::borrow_raw(42)) }
493/// .unwrap()
494/// .with_stride(2048)
495/// .with_offset(0);
496/// ```
497#[cfg(unix)]
498pub struct PlaneDescriptor {
499 fd: OwnedFd,
500 stride: Option<usize>,
501 offset: Option<usize>,
502}
503
504#[cfg(unix)]
505impl PlaneDescriptor {
506 /// Create a new plane descriptor by duplicating the given file descriptor.
507 ///
508 /// The fd is duped immediately — a bad fd fails here rather than inside
509 /// `import_image`. The caller retains ownership of the original fd.
510 ///
511 /// # Errors
512 ///
513 /// Returns an error if the `dup()` syscall fails (e.g. invalid fd or
514 /// fd limit reached).
515 pub fn new(fd: std::os::fd::BorrowedFd<'_>) -> Result<Self> {
516 let owned = fd.try_clone_to_owned()?;
517 Ok(Self {
518 fd: owned,
519 stride: None,
520 offset: None,
521 })
522 }
523
524 /// Set the row stride in bytes (consuming builder).
525 pub fn with_stride(mut self, stride: usize) -> Self {
526 self.stride = Some(stride);
527 self
528 }
529
530 /// Set the plane offset in bytes (consuming builder).
531 pub fn with_offset(mut self, offset: usize) -> Self {
532 self.offset = Some(offset);
533 self
534 }
535
536 /// Consume the descriptor and return the owned file descriptor.
537 pub fn into_fd(self) -> OwnedFd {
538 self.fd
539 }
540
541 /// Row stride in bytes, if set.
542 pub fn stride(&self) -> Option<usize> {
543 self.stride
544 }
545
546 /// Plane offset in bytes, if set.
547 pub fn offset(&self) -> Option<usize> {
548 self.offset
549 }
550}
551
552/// A rectangular sub-region of a tensor's leading spatial frame, in pixels.
553///
554/// `Region` is the single rectangle type in the workspace: the argument to
555/// [`Tensor::view`], the source sampling window in the image crate's `Crop`,
556/// and the geometry the image backend lowers to a `glViewport` (a destination
557/// tile) or a sampling rectangle (a source). Coordinates are pixel/element
558/// units of the leading spatial axes; byte addressing is derived from the
559/// parent's row stride, not stored here.
560#[derive(Debug, Clone, Copy, PartialEq, Eq)]
561pub struct Region {
562 pub x: usize,
563 pub y: usize,
564 pub width: usize,
565 pub height: usize,
566}
567
568impl Region {
569 /// Create a region at `(x, y)` spanning `width` × `height` pixels.
570 pub fn new(x: usize, y: usize, width: usize, height: usize) -> Self {
571 Self {
572 x,
573 y,
574 width,
575 height,
576 }
577 }
578
579 /// True when the region lies fully within a `width` × `height` frame.
580 pub fn fits_within(&self, width: usize, height: usize) -> bool {
581 self.x.saturating_add(self.width) <= width && self.y.saturating_add(self.height) <= height
582 }
583}
584
585/// The parent image a [`view`](Tensor::view)/[`batch`](Tensor::batch) sub-region
586/// was carved from, snapshotted at the time the view was created.
587///
588/// A view shares the parent's `BufferIdentity` and addresses a sub-rectangle of
589/// it. The GL backend keys its EGLImage import on the **parent** geometry (so all
590/// sibling views of one buffer collapse to a single import) and renders each
591/// view as a `glViewport`+`glScissor` ROI at `(x, y, width, height)` within that
592/// parent — the view is render state, never a distinct import. `parent_width`/
593/// `parent_height` are the parent's logical pixel dimensions; `x`/`y` are this
594/// view's top-left origin within the parent (pixels). Nested views compose:
595/// the snapshot always names the **root** parent, with offsets accumulated.
596#[derive(Debug, Clone, Copy, PartialEq, Eq)]
597pub struct ViewOrigin {
598 /// Logical width of the root parent image, in pixels.
599 pub parent_width: usize,
600 /// Logical height of the root parent image, in pixels. For a `batch(n)` view
601 /// of an `[N, H, W, C]` tensor this is `N * H` (the tiles stack vertically in
602 /// the shared buffer).
603 pub parent_height: usize,
604 /// The parent's row stride in **bytes**. The GL backend keys its EGLImage
605 /// import/cache and pitch on this — NOT on the view's own `row_stride`, which
606 /// a single-row view sets tight (for map-span safety). Using the parent
607 /// stride keeps the import pitch parent-consistent so single-row and
608 /// multi-row sibling views collapse onto the same parent import.
609 pub parent_row_stride: usize,
610 /// This view's top-left x origin within the root parent, in pixels.
611 pub x: usize,
612 /// This view's top-left y origin within the root parent, in pixels.
613 pub y: usize,
614}
615
616// Declared through `ef_vocabulary!` (see `vocabulary.rs`): `.code()` is this
617// crate's wire numbering, mirrored by `protocol::dtype` and the C ABI's
618// `hal_dtype`. Both surfaces are generated from the same per-variant literal
619// below rather than re-typed by hand, and `dtype_codes_are_unchanged_by_the_macro`
620// in `tests/vocabulary.rs` pins them against the values already shipping.
621crate::ef_vocabulary! {
622 /// Element type discriminant for runtime type identification.
623 #[derive(Serialize, Deserialize)]
624 #[non_exhaustive]
625 pub enum DType {
626 U8 = 0, "u8", U8,
627 I8 = 1, "i8", I8,
628 U16 = 2, "u16", U16,
629 I16 = 3, "i16", I16,
630 U32 = 4, "u32", U32,
631 I32 = 5, "i32", I32,
632 U64 = 6, "u64", U64,
633 I64 = 7, "i64", I64,
634 F16 = 8, "f16", F16,
635 F32 = 9, "f32", F32,
636 F64 = 10, "f64", F64,
637 }
638 // `#[doc(hidden)]`: this module has to be `pub` for `protocol::dtype`'s
639 // `pub use` re-export below to compile, but it is emission plumbing, not
640 // a second documented API for the same eleven codes -- `protocol::dtype`
641 // is the canonical, documented path.
642 #[doc(hidden)]
643 pub mod dtype_wire;
644}
645
646impl DType {
647 /// Size of one element in bytes.
648 pub const fn size(&self) -> usize {
649 match self {
650 Self::U8 | Self::I8 => 1,
651 Self::U16 | Self::I16 | Self::F16 => 2,
652 Self::U32 | Self::I32 | Self::F32 => 4,
653 Self::U64 | Self::I64 | Self::F64 => 8,
654 }
655 }
656
657 /// Short type name (e.g., "u8", "f32", "f16"). Identical to
658 /// [`DType::as_str`]; kept as its own method for existing callers.
659 pub fn name(&self) -> &'static str {
660 self.as_str()
661 }
662}
663
664impl fmt::Display for DType {
665 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
666 f.write_str(self.name())
667 }
668}
669
670/// Map a static numeric type `T` to its `DType` discriminant, returning
671/// `None` for types that do not have a `DType` representation (e.g.
672/// user-defined wrappers in tests).
673///
674/// Runtime dtype of a static `Tensor<T>` element type — a safe `TypeId`
675/// allowlist of HAL's primitive numeric types (`Some` for `u8`..`i64` /
676/// `f16` / `f32` / `f64`, `None` otherwise). Used by the macOS IOSurface
677/// image constructors (FourCC / pixel-format lookup) and by the `Mem`
678/// backing's `alloc_zeroed` fast path (these types' `T::zero()` is the
679/// all-zeros bit pattern), so it is compiled on every target -- but only
680/// under `static`: every call site lives in `impl<T> Tensor<T>`, itself
681/// `#[cfg(not(edgefirst_tensor_dynamic))]`.
682#[cfg(not(edgefirst_tensor_dynamic))]
683pub(crate) fn dtype_of<T: 'static>() -> Option<DType> {
684 use std::any::TypeId;
685 let id = TypeId::of::<T>();
686 if id == TypeId::of::<u8>() {
687 Some(DType::U8)
688 } else if id == TypeId::of::<i8>() {
689 Some(DType::I8)
690 } else if id == TypeId::of::<u16>() {
691 Some(DType::U16)
692 } else if id == TypeId::of::<i16>() {
693 Some(DType::I16)
694 } else if id == TypeId::of::<u32>() {
695 Some(DType::U32)
696 } else if id == TypeId::of::<i32>() {
697 Some(DType::I32)
698 } else if id == TypeId::of::<u64>() {
699 Some(DType::U64)
700 } else if id == TypeId::of::<i64>() {
701 Some(DType::I64)
702 } else if id == TypeId::of::<half::f16>() {
703 Some(DType::F16)
704 } else if id == TypeId::of::<f32>() {
705 Some(DType::F32)
706 } else if id == TypeId::of::<f64>() {
707 Some(DType::F64)
708 } else {
709 None
710 }
711}
712
713// =============================================================================
714// Quantization metadata — type-gated to integer element types via sealed
715// `IntegerType` trait. Accessors on `Tensor<T>` only compile when `T` is
716// an integer type; calling them on `Tensor<f32>` / `Tensor<f16>` etc. is a
717// compile error, not a runtime one.
718// =============================================================================
719
720mod sealed {
721 pub trait Sealed {}
722 impl Sealed for u8 {}
723 impl Sealed for i8 {}
724 impl Sealed for u16 {}
725 impl Sealed for i16 {}
726 impl Sealed for u32 {}
727 impl Sealed for i32 {}
728 impl Sealed for u64 {}
729 impl Sealed for i64 {}
730 // f16 / f32 / f64 are sealed too (crate::lens::Element covers all
731 // eleven element types), but deliberately NOT given `IntegerType` below.
732 impl Sealed for half::f16 {}
733 impl Sealed for f32 {}
734 impl Sealed for f64 {}
735}
736
737/// Integer element types that may carry quantization metadata.
738///
739/// Sealed trait: implemented for `u8`, `i8`, `u16`, `i16`, `u32`, `i32`,
740/// `u64`, `i64`. Cannot be implemented downstream. Float element types
741/// (`half::f16`, `f32`, `f64`) are explicitly excluded — quantization
742/// metadata does not apply to float tensors per the edgefirst.json spec.
743pub trait IntegerType: sealed::Sealed {}
744impl IntegerType for u8 {}
745impl IntegerType for i8 {}
746impl IntegerType for u16 {}
747impl IntegerType for i16 {}
748impl IntegerType for u32 {}
749impl IntegerType for i32 {}
750impl IntegerType for u64 {}
751impl IntegerType for i64 {}
752
753/// Quantization parameters for an integer tensor.
754///
755/// Covers all four modes the edgefirst.json spec defines:
756///
757/// | Mode | `scale.len()` | `zero_point` | `axis` |
758/// |---|---|---|---|
759/// | Per-tensor symmetric | 1 | `None` | `None` |
760/// | Per-tensor asymmetric | 1 | `Some(len == 1)` | `None` |
761/// | Per-channel symmetric | >1 | `None` | `Some(c)` |
762/// | Per-channel asymmetric | >1 | `Some(len == scale.len())` | `Some(c)` |
763///
764/// The quantized storage type is carried on the parent [`Tensor<T>`]; this
765/// struct does not duplicate it. Construct via the four named constructors
766/// (the only public entry points); direct field mutation is not allowed so
767/// invalid combinations cannot be represented.
768///
769/// Dequantization formula:
770///
771/// ```text
772/// real_value = scale[c] × (quantized_value[c] - zero_point[c])
773/// ```
774///
775/// where `c` is the channel index (always `0` for per-tensor).
776#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
777pub struct Quantization {
778 /// Per-tensor: `vec![scale]`. Per-channel: `vec![scale_0, scale_1, ...]`.
779 #[serde(deserialize_with = "deserialize_scalar_or_vec_f32")]
780 scale: Vec<f32>,
781
782 /// `None` means symmetric (zero-point is 0). `Some(vec)` must have the
783 /// same length as `scale`.
784 #[serde(
785 default,
786 deserialize_with = "deserialize_opt_scalar_or_vec_i32",
787 skip_serializing_if = "Option::is_none"
788 )]
789 zero_point: Option<Vec<i32>>,
790
791 /// Channel axis for per-channel quantization. `Some(_)` iff
792 /// `scale.len() > 1`. Validated against the parent tensor's shape at
793 /// `set_quantization()` time.
794 #[serde(default, skip_serializing_if = "Option::is_none")]
795 axis: Option<usize>,
796}
797
798/// Semantic mode discriminant for hot-path kernel dispatch.
799///
800/// Obtain via [`Quantization::mode`] once at kernel entry; never inside a
801/// pixel-level loop. The enum is borrow-based so the hot kernel receives
802/// the scales / zero-points as slices without reallocation.
803#[derive(Debug, Clone, Copy)]
804pub enum QuantMode<'a> {
805 PerTensorSymmetric {
806 scale: f32,
807 },
808 PerTensor {
809 scale: f32,
810 zero_point: i32,
811 },
812 PerChannelSymmetric {
813 scales: &'a [f32],
814 axis: usize,
815 },
816 PerChannel {
817 scales: &'a [f32],
818 zero_points: &'a [i32],
819 axis: usize,
820 },
821}
822
823impl Quantization {
824 /// Per-tensor symmetric (zero_point = 0).
825 pub fn per_tensor_symmetric(scale: f32) -> Self {
826 Self {
827 scale: vec![scale],
828 zero_point: None,
829 axis: None,
830 }
831 }
832
833 /// Per-tensor asymmetric — the most common runtime shape.
834 pub fn per_tensor(scale: f32, zero_point: i32) -> Self {
835 Self {
836 scale: vec![scale],
837 zero_point: Some(vec![zero_point]),
838 axis: None,
839 }
840 }
841
842 /// Per-channel symmetric. Errors on empty `scales`.
843 pub fn per_channel_symmetric(scales: Vec<f32>, axis: usize) -> Result<Self> {
844 if scales.is_empty() {
845 return Err(Error::QuantizationInvalid {
846 field: "scale.len",
847 expected: "non-empty per-channel scales".to_string(),
848 got: "length 0".to_string(),
849 });
850 }
851 Ok(Self {
852 scale: scales,
853 zero_point: None,
854 axis: Some(axis),
855 })
856 }
857
858 /// Per-channel asymmetric. Errors on length mismatch between `scales`
859 /// and `zero_points`, or empty arrays.
860 pub fn per_channel(scales: Vec<f32>, zero_points: Vec<i32>, axis: usize) -> Result<Self> {
861 if scales.is_empty() {
862 return Err(Error::QuantizationInvalid {
863 field: "scale.len",
864 expected: "non-empty per-channel scales".to_string(),
865 got: "length 0".to_string(),
866 });
867 }
868 if scales.len() != zero_points.len() {
869 return Err(Error::QuantizationInvalid {
870 field: "zero_point.len",
871 expected: format!("length matches scale ({})", scales.len()),
872 got: format!("length {}", zero_points.len()),
873 });
874 }
875 Ok(Self {
876 scale: scales,
877 zero_point: Some(zero_points),
878 axis: Some(axis),
879 })
880 }
881
882 /// Borrow-based dispatch view. Match once at kernel entry.
883 pub fn mode(&self) -> QuantMode<'_> {
884 match (self.scale.len(), self.zero_point.as_deref(), self.axis) {
885 (1, None, _) => QuantMode::PerTensorSymmetric {
886 scale: self.scale[0],
887 },
888 (1, Some(zps), _) => QuantMode::PerTensor {
889 scale: self.scale[0],
890 zero_point: zps.first().copied().unwrap_or(0),
891 },
892 (_, None, Some(axis)) => QuantMode::PerChannelSymmetric {
893 scales: &self.scale,
894 axis,
895 },
896 (_, Some(zps), Some(axis)) => QuantMode::PerChannel {
897 scales: &self.scale,
898 zero_points: zps,
899 axis,
900 },
901 // The `validate()` path prevents constructing a
902 // per-channel Quantization without an axis, so the remaining
903 // pattern is unreachable in practice. Fall back to
904 // per-tensor symmetric using scale[0] to avoid panicking in
905 // release; debug builds assert.
906 _ => {
907 debug_assert!(
908 false,
909 "Quantization::mode: per-channel without axis is unreachable"
910 );
911 QuantMode::PerTensorSymmetric {
912 scale: self.scale.first().copied().unwrap_or(1.0),
913 }
914 }
915 }
916 }
917
918 /// Returns `true` for per-tensor quantization (`scale.len() == 1`).
919 pub fn is_per_tensor(&self) -> bool {
920 self.scale.len() == 1
921 }
922
923 /// Returns `true` for per-channel quantization (`scale.len() > 1`).
924 pub fn is_per_channel(&self) -> bool {
925 self.scale.len() > 1
926 }
927
928 /// Returns `true` for symmetric quantization (no zero-point, or
929 /// zero-point vector of all zeros).
930 pub fn is_symmetric(&self) -> bool {
931 match &self.zero_point {
932 None => true,
933 Some(zps) => zps.iter().all(|&z| z == 0),
934 }
935 }
936
937 /// Borrow the scale array. Length 1 for per-tensor; `num_channels` for
938 /// per-channel.
939 pub fn scale(&self) -> &[f32] {
940 &self.scale
941 }
942
943 /// Borrow the zero-point array. `None` for symmetric.
944 pub fn zero_point(&self) -> Option<&[i32]> {
945 self.zero_point.as_deref()
946 }
947
948 /// Channel axis for per-channel quantization. `None` for per-tensor.
949 pub fn axis(&self) -> Option<usize> {
950 self.axis
951 }
952
953 /// Validate against a target tensor shape. Runs in
954 /// `Tensor::set_quantization()`. Catches:
955 /// - empty `scale` (reject — must declare at least one factor)
956 /// - `zero_point` length inconsistent with `scale` (reject —
957 /// per-tensor must have len 1, per-channel must match `scale.len`)
958 /// - `axis >= shape.len()` (axis out of range)
959 /// - `scale.len() != shape[axis]` for per-channel
960 /// - per-channel without axis (reject)
961 /// - per-tensor with redundant axis (reject)
962 ///
963 /// `static`-only: `Tensor::set_quantization` is `impl<T> Tensor<T>`,
964 /// itself `#[cfg(not(edgefirst_tensor_dynamic))]`. `dynamic`'s
965 /// `TensorDyn::set_quantization` (`dynamic_backend.rs`) sends the raw
966 /// scale/zero-point/axis over FFI and lets `libedgefirst_tensor.so`
967 /// validate them on its own (`static`) side instead of calling this
968 /// method locally.
969 #[cfg(not(edgefirst_tensor_dynamic))]
970 pub(crate) fn validate(&self, shape: &[usize]) -> Result<()> {
971 // `Quantization` is `Deserialize`, so malformed JSON like
972 // `{"scale": [], "zero_point": []}` could otherwise produce an
973 // ill-defined value that confuses `mode()` selection and the
974 // per-channel kernels' indexing.
975 if self.scale.is_empty() {
976 return Err(Error::QuantizationInvalid {
977 field: "scale.len",
978 expected: ">= 1".to_string(),
979 got: "0".to_string(),
980 });
981 }
982 if let Some(zps) = self.zero_point.as_ref() {
983 // Per-tensor: scale.len() == 1 and zero_point.len() must == 1.
984 // Per-channel: zero_point.len() must == scale.len().
985 let expected = if self.scale.len() == 1 {
986 1
987 } else {
988 self.scale.len()
989 };
990 if zps.len() != expected {
991 return Err(Error::QuantizationInvalid {
992 field: "zero_point.len",
993 expected: format!(
994 "{expected} (matching {})",
995 if self.scale.len() == 1 {
996 "per-tensor scale"
997 } else {
998 "per-channel scale.len"
999 }
1000 ),
1001 got: format!("length {}", zps.len()),
1002 });
1003 }
1004 }
1005
1006 match (self.scale.len(), self.axis) {
1007 (1, None) => Ok(()),
1008 (1, Some(_)) => Err(Error::QuantizationInvalid {
1009 field: "per_tensor_redundant_axis",
1010 expected: "axis=None for per-tensor quantization".to_string(),
1011 got: format!("axis={:?}", self.axis),
1012 }),
1013 (_, None) => Err(Error::QuantizationInvalid {
1014 field: "per_channel_requires_axis",
1015 expected: format!(
1016 "axis=Some(_) for per-channel quantization (scale.len={})",
1017 self.scale.len()
1018 ),
1019 got: "axis=None".to_string(),
1020 }),
1021 (n, Some(axis)) => {
1022 if axis >= shape.len() {
1023 return Err(Error::QuantizationInvalid {
1024 field: "axis",
1025 expected: format!("axis < tensor rank ({})", shape.len()),
1026 got: format!("axis={axis}"),
1027 });
1028 }
1029 if shape[axis] != n {
1030 return Err(Error::QuantizationInvalid {
1031 field: "scale.len",
1032 expected: format!("length matches shape[{axis}] ({})", shape[axis]),
1033 got: format!("length {n}"),
1034 });
1035 }
1036 Ok(())
1037 }
1038 }
1039 }
1040}
1041
1042impl From<(f32, i32)> for Quantization {
1043 /// Convenience construction from a `(scale, zero_point)` tuple. Matches
1044 /// the legacy `QuantTuple` / `Quantization::new` calling convention so
1045 /// existing `(0.1, -128).into()` sites keep working.
1046 fn from((scale, zero_point): (f32, i32)) -> Self {
1047 Self::per_tensor(scale, zero_point)
1048 }
1049}
1050
1051fn deserialize_scalar_or_vec_f32<'de, D: serde::Deserializer<'de>>(
1052 de: D,
1053) -> std::result::Result<Vec<f32>, D::Error> {
1054 use serde::de::{self, Visitor};
1055 struct V;
1056 impl<'de> Visitor<'de> for V {
1057 type Value = Vec<f32>;
1058 // Only serde's error text reads this; no value flows from it.
1059 #[cfg_attr(test, mutants::skip)]
1060 fn expecting(&self, f: &mut fmt::Formatter) -> fmt::Result {
1061 f.write_str("f32 or array of f32")
1062 }
1063 fn visit_f64<E: de::Error>(self, v: f64) -> std::result::Result<Self::Value, E> {
1064 Ok(vec![v as f32])
1065 }
1066 #[allow(clippy::cast_possible_truncation)]
1067 fn visit_i64<E: de::Error>(self, v: i64) -> std::result::Result<Self::Value, E> {
1068 Ok(vec![v as f32])
1069 }
1070 #[allow(clippy::cast_possible_truncation, clippy::cast_precision_loss)]
1071 fn visit_u64<E: de::Error>(self, v: u64) -> std::result::Result<Self::Value, E> {
1072 Ok(vec![v as f32])
1073 }
1074 fn visit_seq<A: de::SeqAccess<'de>>(
1075 self,
1076 mut seq: A,
1077 ) -> std::result::Result<Self::Value, A::Error> {
1078 let mut out = Vec::with_capacity(seq.size_hint().unwrap_or(1));
1079 while let Some(x) = seq.next_element::<f32>()? {
1080 out.push(x);
1081 }
1082 Ok(out)
1083 }
1084 }
1085 de.deserialize_any(V)
1086}
1087
1088fn deserialize_opt_scalar_or_vec_i32<'de, D: serde::Deserializer<'de>>(
1089 de: D,
1090) -> std::result::Result<Option<Vec<i32>>, D::Error> {
1091 use serde::de::{self, Visitor};
1092 struct V;
1093 impl<'de> Visitor<'de> for V {
1094 type Value = Option<Vec<i32>>;
1095 // Only serde's error text reads this; no value flows from it.
1096 #[cfg_attr(test, mutants::skip)]
1097 fn expecting(&self, f: &mut fmt::Formatter) -> fmt::Result {
1098 f.write_str("null, i32, or array of i32")
1099 }
1100 // `Ok(None)` is `Ok(Default::default())`, the only replacement.
1101 #[cfg_attr(test, mutants::skip)]
1102 fn visit_none<E: de::Error>(self) -> std::result::Result<Self::Value, E> {
1103 Ok(None)
1104 }
1105 // As `visit_none`.
1106 #[cfg_attr(test, mutants::skip)]
1107 fn visit_unit<E: de::Error>(self) -> std::result::Result<Self::Value, E> {
1108 Ok(None)
1109 }
1110 fn visit_some<D2: serde::Deserializer<'de>>(
1111 self,
1112 de: D2,
1113 ) -> std::result::Result<Self::Value, D2::Error> {
1114 struct Inner;
1115 impl<'de> Visitor<'de> for Inner {
1116 type Value = Vec<i32>;
1117 // Only serde's error text reads this; no value flows from it.
1118 #[cfg_attr(test, mutants::skip)]
1119 fn expecting(&self, f: &mut fmt::Formatter) -> fmt::Result {
1120 f.write_str("i32 or array of i32")
1121 }
1122 #[allow(clippy::cast_possible_truncation)]
1123 fn visit_i64<E: de::Error>(self, v: i64) -> std::result::Result<Self::Value, E> {
1124 Ok(vec![v as i32])
1125 }
1126 #[allow(clippy::cast_possible_truncation, clippy::cast_possible_wrap)]
1127 fn visit_u64<E: de::Error>(self, v: u64) -> std::result::Result<Self::Value, E> {
1128 Ok(vec![v as i32])
1129 }
1130 fn visit_seq<A: de::SeqAccess<'de>>(
1131 self,
1132 mut seq: A,
1133 ) -> std::result::Result<Self::Value, A::Error> {
1134 let mut out = Vec::with_capacity(seq.size_hint().unwrap_or(1));
1135 while let Some(x) = seq.next_element::<i32>()? {
1136 out.push(x);
1137 }
1138 Ok(out)
1139 }
1140 }
1141 de.deserialize_any(Inner).map(Some)
1142 }
1143 #[allow(clippy::cast_possible_truncation)]
1144 fn visit_i64<E: de::Error>(self, v: i64) -> std::result::Result<Self::Value, E> {
1145 Ok(Some(vec![v as i32]))
1146 }
1147 #[allow(clippy::cast_possible_truncation, clippy::cast_possible_wrap)]
1148 fn visit_u64<E: de::Error>(self, v: u64) -> std::result::Result<Self::Value, E> {
1149 Ok(Some(vec![v as i32]))
1150 }
1151 fn visit_seq<A: de::SeqAccess<'de>>(
1152 self,
1153 mut seq: A,
1154 ) -> std::result::Result<Self::Value, A::Error> {
1155 let mut out = Vec::with_capacity(seq.size_hint().unwrap_or(1));
1156 while let Some(x) = seq.next_element::<i32>()? {
1157 out.push(x);
1158 }
1159 Ok(Some(out))
1160 }
1161 }
1162 de.deserialize_option(V)
1163}
1164
1165/// Uniform guard for mutable access through a read-only map. Every
1166/// `TensorMap` backend calls this from `as_mut_slice`/`deref_mut` so a
1167/// `map_read()` misuse fails identically on all platforms instead of
1168/// passing on tolerant ones and exploding only on Android.
1169#[inline]
1170pub(crate) fn assert_map_writable(writable: bool, backend: &str) {
1171 assert!(
1172 writable,
1173 "{backend} map is read-only (obtained via map_read()/CpuAccess::Read) — \
1174 use map_mut() or map_write() for mutable access"
1175 );
1176}
1177
1178crate::ef_vocabulary! {
1179 /// Declared CPU involvement for a tensor, chosen at allocation.
1180 ///
1181 /// The declaration covers **all** CPU access to the buffer — the
1182 /// caller's maps, pins and sync brackets, and the HAL's own CPU work
1183 /// (the CPU image backend, codec decodes, decoder outputs). Hardware
1184 /// access (ISP, codec, GPU, NPU) needs no declaration. It selects the CPU
1185 /// usage/mapping mode at allocation (write-combined for `Write`, cached
1186 /// for `Read`) and, on Android, pins the layout linear (vendor tile
1187 /// compression requires `None`).
1188 ///
1189 /// A CPU map, pin or sync whose direction the declaration does not cover
1190 /// fails with [`Error::CpuAccessDenied`] on every backend.
1191 ///
1192 /// On Android the `ReadWrite` default asks gralloc for
1193 /// `CPU_READ_OFTEN | CPU_WRITE_OFTEN`, which may give linear, CPU-cached
1194 /// memory. Pass `None` for a buffer that is never mapped on the CPU, so
1195 /// gralloc can keep a tiled or compressed allocation. Linux DMA-BUF
1196 /// allocation does not depend on the declaration.
1197 ///
1198 /// [`CpuAccess::code`] is the wire code the C header's `ef_cpu_access`
1199 /// enumerators are asserted against.
1200 #[derive(Default)]
1201 pub enum CpuAccess {
1202 /// Hardware-only buffer: no CPU access. Required for vendor tile
1203 /// compression on platforms with compressed layouts.
1204 None = 0, "none", NONE,
1205 /// CPU reads only (verification, CPU consumers) — cached mapping.
1206 Read = 1, "read", READ,
1207 /// CPU writes only (decode targets) — write-combined mapping where
1208 /// the platform supports it.
1209 Write = 2, "write", WRITE,
1210 /// CPU reads and writes (the default).
1211 #[default]
1212 ReadWrite = 3, "readwrite", READ_WRITE,
1213 }
1214 // Hidden like `dtype_wire`: emission plumbing, not a second documented API.
1215 #[doc(hidden)]
1216 pub mod cpu_access_wire;
1217}
1218
1219impl CpuAccess {
1220 /// The declaration with exactly the given directions.
1221 pub(crate) const fn from_directions(read: bool, write: bool) -> CpuAccess {
1222 match (read, write) {
1223 (true, true) => CpuAccess::ReadWrite,
1224 (true, false) => CpuAccess::Read,
1225 (false, true) => CpuAccess::Write,
1226 (false, false) => CpuAccess::None,
1227 }
1228 }
1229
1230 /// Whether this declaration includes CPU reads.
1231 pub fn reads(self) -> bool {
1232 matches!(self, CpuAccess::Read | CpuAccess::ReadWrite)
1233 }
1234
1235 /// Whether this declaration includes CPU writes.
1236 pub fn writes(self) -> bool {
1237 matches!(self, CpuAccess::Write | CpuAccess::ReadWrite)
1238 }
1239
1240 /// Whether this declaration covers `requested` (every direction the
1241 /// request needs is declared).
1242 pub fn covers(self, requested: CpuAccess) -> bool {
1243 (!requested.reads() || self.reads()) && (!requested.writes() || self.writes())
1244 }
1245
1246 /// The declaration covering every direction of both `self` and `other`.
1247 ///
1248 /// ```
1249 /// use edgefirst_tensor::CpuAccess;
1250 /// assert_eq!(CpuAccess::Read.union(CpuAccess::Write), CpuAccess::ReadWrite);
1251 /// assert_eq!(CpuAccess::None.union(CpuAccess::Write), CpuAccess::Write);
1252 /// ```
1253 pub fn union(self, other: CpuAccess) -> CpuAccess {
1254 Self::from_directions(
1255 self.reads() || other.reads(),
1256 self.writes() || other.writes(),
1257 )
1258 }
1259
1260 /// The declaration covering only the directions both `self` and `other`
1261 /// cover.
1262 ///
1263 /// ```
1264 /// use edgefirst_tensor::CpuAccess;
1265 /// assert_eq!(CpuAccess::ReadWrite.intersection(CpuAccess::Write), CpuAccess::Write);
1266 /// assert_eq!(CpuAccess::Read.intersection(CpuAccess::Write), CpuAccess::None);
1267 /// ```
1268 pub fn intersection(self, other: CpuAccess) -> CpuAccess {
1269 Self::from_directions(
1270 self.reads() && other.reads(),
1271 self.writes() && other.writes(),
1272 )
1273 }
1274
1275 /// `Ok` when this declaration covers `requested`, otherwise
1276 /// [`Error::CpuAccessDenied`] naming both.
1277 ///
1278 /// ```
1279 /// use edgefirst_tensor::{CpuAccess, Error};
1280 /// assert!(CpuAccess::ReadWrite.require(CpuAccess::Write).is_ok());
1281 /// assert!(matches!(
1282 /// CpuAccess::Read.require(CpuAccess::Write),
1283 /// Err(Error::CpuAccessDenied { .. })
1284 /// ));
1285 /// ```
1286 pub fn require(self, requested: CpuAccess) -> Result<()> {
1287 if self.covers(requested) {
1288 Ok(())
1289 } else {
1290 Err(Error::CpuAccessDenied {
1291 declared: self,
1292 requested,
1293 })
1294 }
1295 }
1296
1297 /// [`require`](Self::require) for a map, pin or sync, where
1298 /// `CpuAccess::None` is no direction at all: [`Error::InvalidArgument`]
1299 /// naming `what`.
1300 pub(crate) fn require_direction(self, requested: CpuAccess, what: &str) -> Result<()> {
1301 if requested == CpuAccess::None {
1302 return Err(Error::InvalidArgument(format!(
1303 "{what}: CpuAccess::None is not a CPU access direction"
1304 )));
1305 }
1306 self.require(requested)
1307 }
1308}
1309
1310/// Requested tile-compression behavior for an image allocation (set via
1311/// [`ImageDesc::with_compression`]).
1312///
1313/// Vendor GPUs store textures in proprietary compressed tile layouts
1314/// (UBWC, AFBC, PVRIC, DCC) that cut memory bandwidth; eligibility
1315/// requires a hardware-only buffer ([`CpuAccess::None`]) because CPU
1316/// mapping pins the layout linear. The request records best knowledge on
1317/// the tensor ([`Tensor::compression`]); it never changes what bytes a
1318/// consumer sees through the GPU/NPU.
1319#[non_exhaustive]
1320#[derive(Debug, Clone, Copy, PartialEq, Eq)]
1321pub enum Compression {
1322 /// Let the platform use its native scheme when the format is
1323 /// eligible; otherwise allocate linear and count the fallback
1324 /// ([`compression_fallback_count`]). The right default for pipelines
1325 /// that want the bandwidth win without portability failures.
1326 Any,
1327 /// Require one specific scheme: allocation fails with
1328 /// [`Error::InvalidArgument`] when the device's scheme differs and
1329 /// [`Error::NotImplemented`] on platforms without vendor tile
1330 /// compression. For consumers whose ABI names a layout (e.g. a
1331 /// QNN context binary declaring UBWC inputs).
1332 Scheme(CompressionScheme),
1333}
1334
1335// Declared through `ef_vocabulary!` (see `vocabulary.rs`). This vocabulary
1336// gained a wire numbering when `ef_tensor_compression` was added: the
1337// dynamic backend reads a scheme back across the C boundary, so the enum
1338// needs codes that cannot drift from the C enumerators. Hand-writing that
1339// mapping in `edgefirst-tensor-capi` was the alternative and it is strictly
1340// worse -- `#[non_exhaustive]` makes a cross-crate `match` require a
1341// wildcard arm, so a scheme added here would silently map to "linear"
1342// there. Generated from one declaration, it cannot.
1343//
1344// Code 0 is deliberately unassigned and reserved for "linear" -- the state
1345// this crate spells `Option::None`, which has no variant to carry a code of
1346// its own. The C side names it `EF_COMPRESSION_NONE`.
1347crate::ef_vocabulary! {
1348 /// Vendor tile-compression schemes the HAL recognizes and records.
1349 ///
1350 /// Unrecognized platforms record `None` (linear) — there is deliberately
1351 /// no `Unknown` variant; a scheme is only recorded when the vendor is
1352 /// positively identified.
1353 #[non_exhaustive]
1354 pub enum CompressionScheme {
1355 /// Qualcomm Adreno Universal Bandwidth Compression.
1356 Ubwc = 1, "ubwc", UBWC,
1357 /// Arm Mali/Immortalis Framebuffer Compression.
1358 Afbc = 2, "afbc", AFBC,
1359 /// Imagination PowerVR Image Compression (Google Tensor G5+).
1360 Pvric = 3, "pvric", PVRIC,
1361 /// Samsung Xclipse (AMD RDNA) Delta Color Compression.
1362 Dcc = 4, "dcc", DCC,
1363 }
1364 // `#[doc(hidden)]` for the same reason `dtype_wire` is: emission
1365 // plumbing for `const`-only consumers, not a second documented API.
1366 #[doc(hidden)]
1367 pub mod compression_wire;
1368}
1369
1370crate::ef_vocabulary! {
1371 /// Whether a tensor's memory is known to be physically contiguous
1372 /// ([`Tensor::contiguity`]).
1373 ///
1374 /// Capture and other DMA engines without an IOMMU (the i.MX 8M Plus ISI,
1375 /// the vvcam ISP) can only write into physically contiguous memory.
1376 /// Only memory this crate allocated from a Linux DMA heap is known
1377 /// either way: the CMA heap is contiguous and the system heap is not.
1378 /// An imported fd, and every other kind of memory, is `Unknown`.
1379 #[non_exhaustive]
1380 pub enum Contiguity {
1381 /// Not known: an imported fd, or memory not allocated from a DMA heap.
1382 Unknown = 0, "unknown", UNKNOWN,
1383 /// Allocated from the CMA heap.
1384 Contiguous = 1, "contiguous", CONTIGUOUS,
1385 /// Allocated from the system heap.
1386 NonContiguous = 2, "non_contiguous", NON_CONTIGUOUS,
1387 }
1388 #[doc(hidden)]
1389 pub mod contiguity_wire;
1390}
1391
1392/// Count of image allocations that requested [`Compression::Any`] but
1393/// resolved to a linear layout (ineligible format/dtype, unrecognized
1394/// vendor, or a platform without vendor tile compression). Read via
1395/// [`compression_fallback_count`].
1396static COMPRESSION_FALLBACKS: AtomicU64 = AtomicU64::new(0);
1397
1398/// Number of [`Compression::Any`] requests since process start that
1399/// resolved to a linear layout instead of a vendor tile scheme. A
1400/// steady-state pipeline holds this flat after warmup; growth means an
1401/// allocation path keeps requesting compression it never gets.
1402pub fn compression_fallback_count() -> u64 {
1403 COMPRESSION_FALLBACKS.load(Ordering::Relaxed)
1404}
1405
1406/// Record a `Compression::Any` request resolving linear.
1407///
1408/// `static`-only: the sole caller is `Tensor::image_desc` (`impl<T>
1409/// Tensor<T>`), itself `#[cfg(not(edgefirst_tensor_dynamic))]`.
1410/// [`compression_fallback_count`] stays available under `dynamic` too --
1411/// it just always reads 0 there.
1412#[cfg(not(edgefirst_tensor_dynamic))]
1413pub(crate) fn note_compression_fallback(detail: &str) {
1414 COMPRESSION_FALLBACKS.fetch_add(1, Ordering::Relaxed);
1415 log::debug!("compression request fell back to linear: {detail}");
1416}
1417
1418/// Whether this platform can allocate `(format, dtype)` images in a
1419/// vendor tile-compressed layout. `true` requires an Android build, an
1420/// eligible format (RGBA8888 `u8`/`i8` initially), and a positively
1421/// identified GPU vendor; everywhere else the answer is `false` and
1422/// [`Compression::Any`] requests fall back to linear.
1423pub fn compression_support(format: PixelFormat, dtype: DType) -> bool {
1424 #[cfg(target_os = "android")]
1425 {
1426 crate::ahardwarebuffer_layout::compression_eligible(format, dtype)
1427 && crate::ahardwarebuffer::device_compression_scheme().is_some()
1428 }
1429 #[cfg(not(target_os = "android"))]
1430 {
1431 let _ = (format, dtype);
1432 false
1433 }
1434}
1435
1436/// Declarative image-allocation request — the full-featured front door
1437/// for image tensors ([`TensorDyn::image_desc`] and the image crate's
1438/// `ImageProcessor::create_image_desc`).
1439///
1440/// The classic constructors (`image`, `create_image`) cover the common
1441/// cases; the desc carries the optional requests — today the
1442/// [`Compression`] request — without another constructor-parameter
1443/// sweep. Fields are private and the builders consume/return by value,
1444/// so future options are non-breaking.
1445///
1446/// ```
1447/// use edgefirst_tensor::{Compression, CpuAccess, DType, ImageDesc, PixelFormat};
1448/// let desc = ImageDesc::new(640, 640, PixelFormat::Rgba, DType::U8)
1449/// .with_access(CpuAccess::None)
1450/// .with_compression(Compression::Any);
1451/// ```
1452#[derive(Debug, Clone)]
1453pub struct ImageDesc {
1454 width: usize,
1455 height: usize,
1456 format: PixelFormat,
1457 dtype: DType,
1458 memory: Option<TensorMemory>,
1459 access: CpuAccess,
1460 compression: Option<Compression>,
1461 contiguous: bool,
1462}
1463
1464impl ImageDesc {
1465 /// A new image request: auto-selected memory, [`CpuAccess::ReadWrite`],
1466 /// no compression request.
1467 pub fn new(width: usize, height: usize, format: PixelFormat, dtype: DType) -> Self {
1468 Self {
1469 width,
1470 height,
1471 format,
1472 dtype,
1473 memory: None,
1474 access: CpuAccess::default(),
1475 compression: None,
1476 contiguous: false,
1477 }
1478 }
1479
1480 /// Request a specific memory backing (`None` = auto-select).
1481 pub fn with_memory(mut self, memory: Option<TensorMemory>) -> Self {
1482 self.memory = memory;
1483 self
1484 }
1485
1486 /// Declare the CPU access (see [`CpuAccess`]). Any declaration other
1487 /// than `None` makes a compression request invalid, so a compression
1488 /// request must also call `with_access(CpuAccess::None)`.
1489 pub fn with_access(mut self, access: CpuAccess) -> Self {
1490 self.access = access;
1491 self
1492 }
1493
1494 /// Request a tile-compressed layout (see [`Compression`]).
1495 pub fn with_compression(mut self, compression: Compression) -> Self {
1496 self.compression = Some(compression);
1497 self
1498 }
1499
1500 /// Require physically contiguous memory (see [`Contiguity`]).
1501 ///
1502 /// The image is allocated from the Linux CMA DMA heap only. Allocation
1503 /// fails, naming the heap, when there is no CMA heap or it cannot fit
1504 /// the image, instead of falling back to the system heap or to another
1505 /// kind of memory. Memory must be auto-selected or
1506 /// [`TensorMemory::DmaBuf`]. Other platforms return
1507 /// [`Error::NotImplemented`].
1508 pub fn with_contiguous(mut self, contiguous: bool) -> Self {
1509 self.contiguous = contiguous;
1510 self
1511 }
1512
1513 /// Requested width in pixels.
1514 pub fn width(&self) -> usize {
1515 self.width
1516 }
1517
1518 /// Requested height in pixels.
1519 pub fn height(&self) -> usize {
1520 self.height
1521 }
1522
1523 /// Requested pixel format.
1524 pub fn format(&self) -> PixelFormat {
1525 self.format
1526 }
1527
1528 /// Requested element type.
1529 pub fn dtype(&self) -> DType {
1530 self.dtype
1531 }
1532
1533 /// Requested memory backing (`None` = auto-select).
1534 pub fn memory(&self) -> Option<TensorMemory> {
1535 self.memory
1536 }
1537
1538 /// Declared CPU access.
1539 pub fn access(&self) -> CpuAccess {
1540 self.access
1541 }
1542
1543 /// The compression request, if any.
1544 pub fn compression(&self) -> Option<Compression> {
1545 self.compression
1546 }
1547
1548 /// Whether physically contiguous memory is required.
1549 pub fn contiguous(&self) -> bool {
1550 self.contiguous
1551 }
1552}
1553
1554/// What kind of system key an identity was derived from. Participates in the
1555/// id so that, say, inode 7 and IOSurfaceID 7 cannot collide.
1556///
1557/// Invariant every variant must uphold: the key passed to
1558/// [`BufferIdentity::derived`] is unique among live buffers of that kind
1559/// **that own an allocation**. There is no "no identity" case — where no
1560/// system-wide key exists a process-local one does (a pointer, a GL name),
1561/// and a process-local key is exactly as sound as a system-wide one:
1562/// pointer/name reuse after free is the same ABA hazard as inode reuse, and
1563/// is mitigated the same way, by the retaining cache (see the design spec's
1564/// Task 6).
1565///
1566/// The allocation qualifier is load-bearing, not a hedge. A zero-element
1567/// tensor (`[0, num_protos]`, the tracker's "no detections this frame"
1568/// sentinel — see [`crate::mem`]) owns no allocation, so its `Vec` yields the
1569/// dangling well-aligned pointer `align_of::<T>()`. Every zero-element
1570/// `Tensor<f32>` in a process therefore shares one identity, and can even
1571/// collide across kinds with a small integer key from another variant.
1572///
1573/// That is safe rather than tolerated: an identity exists to answer "are
1574/// these the same bytes?", and buffers with no bytes have nothing to alias.
1575/// Nor can they reach an import cache. There are four GPU import paths —
1576/// a dma-buf fd on Linux, an IOSurface ref on macOS, an `AHardwareBuffer` on
1577/// Android, and an `ID3D11Texture2D` on Windows — and each takes only its own
1578/// object, so a tensor with no allocation reaches none of them. The Windows
1579/// path additionally keys on the texture the import itself holds a reference
1580/// to, and `edgefirst-image` **refuses** any texture whose identity is not
1581/// `D3d11Texture` over that same pointer — in every build, and on every
1582/// import including one served from its cache
1583/// (`GlPlatform::validate_import_identity`). That refusal exists because the
1584/// rule was once only written
1585/// down: the dynamic backend handed the image crate `HostPtr` identities over
1586/// recyclable handle addresses, and converts silently rendered into dropped
1587/// textures. Do not "fix" the zero-element case by minting a distinct id per
1588/// empty tensor: that reintroduces a counter, which is the state this whole
1589/// design exists to remove.
1590#[derive(Debug, Clone, Copy, PartialEq, Eq)]
1591pub enum IdentityKind {
1592 /// `(st_dev, st_ino)` of a dma-buf or memfd. Survives `dup`, and is the
1593 /// same value in every process that holds the buffer.
1594 DmaBuf,
1595 Shm,
1596 /// `IOSurfaceGetID`.
1597 IoSurface,
1598 /// `AHardwareBuffer_getId`, API 31+.
1599 AHardwareBufferId,
1600 /// The `AHardwareBuffer*` pointer, pre-API-31 where no id exists.
1601 /// Process-local: unique among live buffers here, meaningless elsewhere.
1602 AHardwareBufferPtr,
1603 /// An `ID3D11Texture2D*`. Process-local, like `AHardwareBufferPtr`, and
1604 /// the same pointer for every kind of texture tensor: one the HAL
1605 /// allocated, one wrapped from a caller's pointer, and one opened from a
1606 /// shared NT handle all key on the texture object they hold.
1607 ///
1608 /// **Both backends derive it from the same pointer, so their ids are
1609 /// equal byte for byte.** The static backend reads the `ComPtr` it owns
1610 /// (`d3d11::texture::tex_key`); the dynamic backend reads the same
1611 /// pointer back across the ABI with `ef_tensor_d3d11_texture`
1612 /// (`TensorDyn::derive_identity`). That equality is the point, not a
1613 /// coincidence: [`BufferIdentity::derived`] exists so two independently
1614 /// linked copies of this crate name one buffer identically, and
1615 /// `edgefirst-image` links the dynamic backend inside every `-capi` leaf
1616 /// while its own tests link the static one.
1617 ///
1618 /// **An `ef_tensor` handle's address must never be the key.** It is not a
1619 /// buffer at all, nothing the cache holds keeps it alive, and this
1620 /// library's allocator hands a freed handle's address to the very next
1621 /// one it mints — measured at 400 texture handles over 4 distinct
1622 /// addresses. The dynamic backend did exactly that on Windows until
1623 /// 2026-09-04, and a convert whose destination key collided rendered into
1624 /// the dropped texture while returning success.
1625 ///
1626 /// A handle import therefore does **not** compare equal to the tensor it
1627 /// was exported from, nor to a second import of the same handle:
1628 /// `OpenSharedResource1` mints a fresh `ID3D11Texture2D` per open. Keying
1629 /// on a handle value would make those compare equal, and would be unsound:
1630 /// Windows recycles a closed handle value straight into the next
1631 /// `DuplicateHandle`, so a consumer that caches on this id past the
1632 /// producing tensor's lifetime -- as the image crate's EGLImage cache
1633 /// does -- could serve one texture's entry for a later, unrelated import.
1634 /// A COM object's address cannot be recycled that way: a cache that holds
1635 /// a reference to the texture keeps the address taken for as long as its
1636 /// entry can be found. The cost is a cache miss per re-import instead of
1637 /// a wrong image.
1638 D3d11Texture,
1639 /// A host mapping address. Process-local.
1640 HostPtr,
1641 /// A GL buffer name. Valid only in the context that created it.
1642 Pbo,
1643}
1644
1645/// Unique identity for a tensor's underlying buffer.
1646///
1647/// Created fresh on every buffer allocation or import. The `id` is a cache
1648/// key, `kind` records what system key (if any) it was derived from, and the
1649/// `guard` is an `Arc<()>` whose weak references allow downstream caches to
1650/// detect when the buffer has been dropped.
1651#[derive(Debug, Clone)]
1652pub struct BufferIdentity {
1653 id: u64,
1654 kind: IdentityKind,
1655 guard: Arc<()>,
1656}
1657
1658impl BufferIdentity {
1659 /// Derive an identity from a key the operating system already guarantees.
1660 ///
1661 /// Deliberately NOT a counter: two independently-linked copies of this
1662 /// crate must derive the same id for the same buffer, or a consumer's
1663 /// import cache misses every frame. See the design spec, "The core finding".
1664 pub fn derived(kind: IdentityKind, key: u64) -> Self {
1665 // Mix the discriminant into the high bits so different kinds with the
1666 // same numeric key never collide.
1667 let id = ((kind as u64) << 56) ^ key;
1668 Self {
1669 id,
1670 kind,
1671 guard: Arc::new(()),
1672 }
1673 }
1674
1675 /// Unique identifier for this buffer. Changes when the buffer changes.
1676 pub fn id(&self) -> u64 {
1677 self.id
1678 }
1679
1680 /// What kind of system key this identity was derived from.
1681 pub fn kind(&self) -> IdentityKind {
1682 self.kind
1683 }
1684
1685 /// Returns a weak reference to the buffer guard. Goes dead when the
1686 /// owning Tensor is dropped (and no clones remain).
1687 pub fn weak(&self) -> Weak<()> {
1688 Arc::downgrade(&self.guard)
1689 }
1690}
1691
1692// `static`-only from here through `fs_magic` below: every use site lives in
1693// `impl<T> TensorStorage<T>` (`from_fd`'s DMA-BUF-vs-tmpfs classification),
1694// itself `#[cfg(not(edgefirst_tensor_dynamic))]`. `dynamic`'s `from_fd`
1695// (`dynamic_backend.rs`) drives `ef_tensor_builder_wrap` instead, which
1696// does this same classification inside `libedgefirst_tensor.so`.
1697#[cfg(all(target_os = "linux", not(edgefirst_tensor_dynamic)))]
1698use nix::sys::stat::{major, minor};
1699
1700/// Filesystem magic of the internal `dma_buf` mount, from
1701/// `include/uapi/linux/magic.h` (`DMA_BUF_MAGIC`, "DMAB").
1702///
1703/// This is stable UAPI and is the only reliable way to recognize a DMA-BUF
1704/// fd — see [`TensorStorage::from_fd`] for why `st_dev` cannot be used.
1705#[cfg(all(target_os = "linux", not(edgefirst_tensor_dynamic)))]
1706const DMA_BUF_MAGIC: u32 = 0x444d_4142;
1707
1708/// Filesystem magic of tmpfs, from `include/uapi/linux/magic.h`
1709/// (`TMPFS_MAGIC`). Covers both POSIX `shm_open` segments under `/dev/shm`
1710/// and anonymous `memfd_create` files.
1711#[cfg(all(target_os = "linux", not(edgefirst_tensor_dynamic)))]
1712const TMPFS_MAGIC: u32 = 0x0102_1994;
1713
1714/// Normalize a raw `fstatfs` `f_type` to the 32-bit unsigned magic that
1715/// `include/uapi/linux/magic.h` documents.
1716///
1717/// The width and signedness of `f_type` are target-dependent: `__fsword_t`
1718/// on Linux/gnu (`i64` on 64-bit, `i32` on 32-bit targets such as armv7,
1719/// i686 and aarch64-ilp32), `c_int` on uclibc, `c_ulong` on musl, and
1720/// `c_uint` on s390x. Where it is signed and 32 bits wide, widening it
1721/// sign-extends every magic with bit 31 set — hugetlbfs (`0x958458f6`),
1722/// btrfs (`0x9123683e`), f2fs (`0xf2f52010`) — which would make the value
1723/// reported by [`Error::UnknownBufferType`] impossible to find in the
1724/// header the docs point callers at.
1725///
1726/// Every magic is a 32-bit quantity, so truncating back to `u32` recovers
1727/// the true value from all four representations and loses nothing.
1728#[cfg(all(target_os = "linux", not(edgefirst_tensor_dynamic)))]
1729const fn fs_magic(raw: i64) -> u32 {
1730 raw as u32
1731}
1732
1733/// The operations every memory backend implements.
1734///
1735/// This is the seam that lets [`Tensor<T>`] hide which backend is in play.
1736/// The implementors are `DmaTensor` (Linux DMA-BUF), `IoSurfaceTensor`
1737/// (macOS/iOS), `AHardwareBufferTensor` (Android), `ShmTensor`, `MemTensor`,
1738/// and [`PboTensor`]; `TensorStorage<T>` dispatches to whichever is active.
1739/// Backend types are not public — allocate a [`Tensor<T>`] or [`TensorDyn`]
1740/// and call these methods through it.
1741///
1742/// Import the trait to get `shape`, `size`, `map`, `clone_fd`,
1743/// `buffer_identity`, and the zero-copy sub-region `view` on a tensor value.
1744/// `Tensor::view` and `Tensor::batch` route through [`TensorTrait::view`], so
1745/// each backend's identity-sharing rule lives in exactly one place.
1746///
1747/// [`TensorDyn`]: crate::TensorDyn
1748pub trait TensorTrait<T>: Send + Sync
1749where
1750 T: Num + Clone + fmt::Debug,
1751{
1752 /// Create a new tensor with the given shape and optional name. If no name
1753 /// is given, a random name will be generated.
1754 fn new(shape: &[usize], name: Option<&str>) -> Result<Self>
1755 where
1756 Self: Sized;
1757
1758 #[cfg(unix)]
1759 /// Import an existing buffer as a tensor, taking ownership of its file
1760 /// descriptor. The buffer is adopted in place — no bytes are copied.
1761 ///
1762 /// The backend is **detected**, not chosen: the fd already belongs to a
1763 /// buffer type, and this call's job is to recognize which. On Linux that
1764 /// is decided by the fd's filesystem magic, which is stable UAPI
1765 /// (`include/uapi/linux/magic.h`):
1766 ///
1767 /// | Filesystem | Magic | Resulting [`TensorMemory`] |
1768 /// |------------|-------|----------------------------|
1769 /// | `dma_buf` | `DMA_BUF_MAGIC` (`0x444d4142`) | [`TensorMemory::DmaBuf`] |
1770 /// | `tmpfs` (`/dev/shm` **and** `memfd`) | `TMPFS_MAGIC` (`0x01021994`) | [`TensorMemory::Shm`] |
1771 /// | anything else | — | *rejected* — see Errors |
1772 ///
1773 /// Both supported types are identified **positively**. An unrecognized
1774 /// filesystem is an error, never a fallback to shared memory: the wrong
1775 /// branch does not fail loudly (a DMA-BUF is mmap-able, so it would
1776 /// import as a perfectly functional tensor that merely isn't DMA, and a
1777 /// pipe would import as a zero-length one), so guessing would trade a
1778 /// clear error here for silent loss of zero-copy far downstream.
1779 ///
1780 /// The device number is deliberately **not** consulted. `dma_buf` files
1781 /// live on an internal kernel mount whose `st_dev` comes from
1782 /// `get_anon_bdev()` — an IDA shared with every other anonymous
1783 /// pseudo-filesystem and allocated in boot order — so the minor a
1784 /// DMA-BUF lands on varies by kernel build and is not part of any ABI.
1785 ///
1786 /// On non-Linux Unix (macOS/iOS/Android) there is no fd-based DMA import;
1787 /// the fd is always adopted as [`TensorMemory::Shm`].
1788 ///
1789 /// # Arguments
1790 ///
1791 /// * `fd` - Owned descriptor for the buffer to import. Ownership
1792 /// transfers to the returned tensor and the fd is closed on drop; pass
1793 /// [`clone_fd`](TensorTrait::clone_fd) output to keep your own handle.
1794 /// * `shape` - Logical shape to interpret the buffer with. Must describe
1795 /// no more elements than the buffer holds.
1796 /// * `name` - Optional name; a random one is generated when `None`.
1797 ///
1798 /// # Returns
1799 ///
1800 /// A tensor sharing the imported buffer's memory, whose
1801 /// [`memory()`](TensorTrait::memory) reports the detected backend.
1802 ///
1803 /// # Errors
1804 ///
1805 /// * [`Error::UnknownBufferType`] - the fd is on a filesystem that is
1806 /// neither `dma_buf` nor `tmpfs`, so its buffer type cannot be
1807 /// determined. Carries the observed `fstatfs` magic as a `u32`,
1808 /// normalized so it can be looked up in `magic.h` directly on both
1809 /// 32- and 64-bit targets. Typical causes: a regular file, a pipe or
1810 /// socket, or a `MFD_HUGETLB` memfd (hugetlbfs, not tmpfs). Linux only.
1811 /// * [`Error::UnknownDeviceType`] - the fd's `st_dev` major is non-zero,
1812 /// i.e. it lives on a real block device rather than an anonymous or
1813 /// in-memory filesystem. Linux only.
1814 /// * [`Error::InvalidSize`] - `shape` is empty or describes zero
1815 /// elements.
1816 /// * [`Error::NixError`] - `fstat`, `fstatfs`, or `mmap` failed on the
1817 /// descriptor.
1818 ///
1819 /// # Examples
1820 ///
1821 /// ```no_run
1822 /// use edgefirst_tensor::{Tensor, TensorMemory, TensorTrait};
1823 ///
1824 /// # fn main() -> edgefirst_tensor::Result<()> {
1825 /// let src = Tensor::<u8>::new(&[480, 640, 3], Some(TensorMemory::DmaBuf), None)?;
1826 ///
1827 /// // Round-tripping a DMA-BUF fd preserves the backend — the import is
1828 /// // still zero-copy, and still eligible for GPU/NPU paths.
1829 /// let imported = Tensor::<u8>::from_fd(src.clone_fd()?, src.shape(), None)?;
1830 /// assert_eq!(imported.memory(), TensorMemory::DmaBuf);
1831 /// # Ok(())
1832 /// # }
1833 /// ```
1834 fn from_fd(fd: std::os::fd::OwnedFd, shape: &[usize], name: Option<&str>) -> Result<Self>
1835 where
1836 Self: Sized;
1837
1838 #[cfg(unix)]
1839 /// Clone the file descriptor associated with this tensor.
1840 fn clone_fd(&self) -> Result<std::os::fd::OwnedFd>;
1841
1842 /// Get the memory type of this tensor.
1843 fn memory(&self) -> TensorMemory;
1844
1845 /// Get the name of this tensor.
1846 fn name(&self) -> String;
1847
1848 /// Get the number of elements in this tensor.
1849 fn len(&self) -> usize {
1850 self.shape().iter().product()
1851 }
1852
1853 /// Check if the tensor is empty.
1854 fn is_empty(&self) -> bool {
1855 self.len() == 0
1856 }
1857
1858 /// Get the size in bytes of this tensor.
1859 fn size(&self) -> usize {
1860 self.len() * std::mem::size_of::<T>()
1861 }
1862
1863 /// Get the shape of this tensor.
1864 fn shape(&self) -> &[usize];
1865
1866 /// Reshape this tensor to the given shape. The total number of elements
1867 /// must remain the same.
1868 fn reshape(&mut self, shape: &[usize]) -> Result<()>;
1869
1870 /// Bytes of the underlying allocation (>= the current logical `size()`).
1871 fn capacity_bytes(&self) -> usize;
1872
1873 /// Set the logical shape to any shape whose byte size fits the allocation
1874 /// capacity, without the equal-size constraint of `reshape`.
1875 fn set_logical_shape(&mut self, shape: &[usize]) -> Result<()>;
1876
1877 /// Map the tensor into memory with the given access direction and
1878 /// return a TensorMap for accessing the data.
1879 ///
1880 /// `access` selects the platform mapping mode (read-only IOSurface
1881 /// lock, dma-buf sync direction, AHardwareBuffer lock usage) and the
1882 /// map's mutability: a map obtained with [`CpuAccess::Read`] rejects
1883 /// `as_mut_slice`. [`CpuAccess::None`] is not a mappable direction
1884 /// and returns [`Error::InvalidArgument`]; a direction the tensor's
1885 /// declared [`CpuAccess`] does not cover returns
1886 /// [`Error::CpuAccessDenied`].
1887 ///
1888 /// Prefer the typed wrappers [`map_read`](Self::map_read) /
1889 /// [`map_write`](Self::map_write) / [`map_mut`](Self::map_mut).
1890 ///
1891 /// Like [`pin_host`](Tensor::pin_host), `'a` bounds how long the backing
1892 /// allocation lives rather than borrowing the tensor: the returned view
1893 /// shares ownership through the pin's keepalive, so a map can coexist with
1894 /// a `&mut` on the same tensor.
1895 ///
1896 /// # Windows D3D11 textures: a write-only map of a partial window
1897 ///
1898 /// A [`CpuAccess::Write`] map is not refreshed from the texture first, so
1899 /// the bytes the caller leaves untouched are undefined -- the documented
1900 /// contract a whole-tensor map can meet and a sub-view cannot, since the
1901 /// window spans its own rows only.
1902 ///
1903 /// A tensor allocated [`CpuAccess::Write`] (no staging texture) publishes
1904 /// only the window's own rows on unmap, so a sub-view -- a
1905 /// [`view`](Tensor::view), or a shape narrowed with `set_logical_shape`
1906 /// -- is refused only when its window is not a whole, pitch-aligned row
1907 /// range (a plain-offset [`view`](Tensor::view) always is; an arbitrary
1908 /// byte offset from `set_plane_offset` might not be), which fails with
1909 /// [`Error::InvalidArgument`].
1910 ///
1911 /// A tensor allocated [`CpuAccess::Read`] or [`CpuAccess::ReadWrite`]
1912 /// (staging texture) has no such per-window publish: unmapping always
1913 /// copies the whole staging texture back, so a write-only sub-view is
1914 /// refused with [`Error::PartialWriteRequiresReadWrite`], naming
1915 /// [`CpuAccess::ReadWrite`] as the same window with the refresh that
1916 /// makes the untouched rows well-defined -- and, unlike the
1917 /// `CpuAccess::Write` case above, that retry always succeeds.
1918 fn map_with<'a>(&self, access: CpuAccess) -> Result<crate::view::HostView<'a, T>>
1919 where
1920 T: 'a;
1921
1922 /// Map the tensor read-write (equivalent to
1923 /// `map_with(CpuAccess::ReadWrite)`).
1924 ///
1925 /// Fails with [`Error::CpuAccessDenied`] on a `Read`- or `Write`-declared
1926 /// tensor: use [`map_read`](Self::map_read) to read one and
1927 /// [`map_for_write`](Self::map_for_write) to fill one.
1928 fn map<'a>(&self) -> Result<crate::view::HostView<'a, T>>
1929 where
1930 T: 'a,
1931 {
1932 self.map_with(CpuAccess::ReadWrite)
1933 }
1934
1935 /// Map the tensor for CPU reading only. The returned map rejects
1936 /// `as_mut_slice`; on macOS this takes the read-only IOSurface lock
1937 /// (skips the unlock flush), on Linux the dma-buf read-direction
1938 /// sync. Fails with [`Error::CpuAccessDenied`] on a `Write`- or
1939 /// `None`-declared tensor.
1940 fn map_read<'a>(&self) -> Result<crate::view::HostView<'a, T>>
1941 where
1942 T: 'a,
1943 {
1944 self.map_with(CpuAccess::Read)
1945 }
1946
1947 /// Map the tensor for CPU writing (fill-only: reading through a
1948 /// write map may see write-combined memory — do not read the slice).
1949 /// Valid only for a pass that overwrites the whole mapping; see
1950 /// [`write_access`](Self::write_access). Fails with
1951 /// [`Error::CpuAccessDenied`] on a `Read`- or `None`-declared tensor.
1952 fn map_write<'a>(&self) -> Result<crate::view::HostView<'a, T>>
1953 where
1954 T: 'a,
1955 {
1956 self.map_with(CpuAccess::Write)
1957 }
1958
1959 /// Map the tensor read-write (alias of [`map`](Self::map) with the
1960 /// intent spelled out); refused on the same declarations.
1961 fn map_mut<'a>(&self) -> Result<crate::view::HostView<'a, T>>
1962 where
1963 T: 'a,
1964 {
1965 self.map_with(CpuAccess::ReadWrite)
1966 }
1967
1968 /// The CPU access declared for this tensor (see [`CpuAccess`]). A bare
1969 /// backing carries no declaration of its own and reports `ReadWrite`.
1970 fn cpu_access(&self) -> CpuAccess {
1971 CpuAccess::ReadWrite
1972 }
1973
1974 /// The parent-image snapshot when this tensor is a sub-region of a larger
1975 /// buffer (see [`ViewOrigin`]); a bare backing reports `None`.
1976 fn view_origin(&self) -> Option<ViewOrigin> {
1977 None
1978 }
1979
1980 /// Whether this tensor addresses part of a larger buffer (see
1981 /// [`Tensor::is_view`]); a bare backing reports `false`.
1982 fn is_view(&self) -> bool {
1983 self.view_origin().is_some()
1984 }
1985
1986 /// The direction a CPU pass that overwrites every element of this tensor
1987 /// maps with.
1988 ///
1989 /// A write-only map does not invalidate the CPU cache (on arm64,
1990 /// `DMA_BUF_SYNC_START|WRITE` only orders), so a stale cache line the pass
1991 /// writes only partly is written back over whatever a device left in the
1992 /// rest of it. The map may therefore be `Write` only when the pass covers
1993 /// the whole buffer: a whole tensor takes its declaration plus `Write`
1994 /// (`Write` on a write-only declaration), a [`view`](Tensor::view), whose
1995 /// rows share cache lines with bytes outside it, takes `ReadWrite`.
1996 fn write_access(&self) -> CpuAccess {
1997 if self.is_view() {
1998 CpuAccess::ReadWrite
1999 } else {
2000 self.cpu_access().union(CpuAccess::Write)
2001 }
2002 }
2003
2004 /// Map for a CPU pass that overwrites every element, in
2005 /// [`write_access`](Self::write_access)'s direction.
2006 ///
2007 /// # Errors
2008 ///
2009 /// [`Error::CpuAccessDenied`] when the declaration does not cover that
2010 /// direction, otherwise whatever [`map_with`](Self::map_with) returns.
2011 fn map_for_write<'a>(&self) -> Result<crate::view::HostView<'a, T>>
2012 where
2013 T: 'a,
2014 {
2015 self.map_with(self.write_access())
2016 }
2017
2018 /// Map for a CPU pass that writes only part of the buffer and keeps the
2019 /// rest (a letterbox border around a device-rendered interior): always
2020 /// `ReadWrite`, see [`write_access`](Self::write_access).
2021 ///
2022 /// # Errors
2023 ///
2024 /// [`Error::CpuAccessDenied`] unless the declaration is `ReadWrite`,
2025 /// otherwise whatever [`map_with`](Self::map_with) returns.
2026 fn map_for_partial_write<'a>(&self) -> Result<crate::view::HostView<'a, T>>
2027 where
2028 T: 'a,
2029 {
2030 self.map_with(CpuAccess::ReadWrite)
2031 }
2032
2033 /// Non-blocking [`map_with`](Self::map_with): returns
2034 /// `Err(Error::IoError(WouldBlock))` while a GPU copy the map depends on
2035 /// is still in flight, and makes progress on a retry.
2036 ///
2037 /// Backings with no such copy (host memory, shared memory, dma-buf,
2038 /// IOSurface, AHardwareBuffer, PBO) alias `map_with`, so a caller can use
2039 /// this unconditionally and only ever see `WouldBlock` where the wait is
2040 /// real -- today the Windows D3D11 texture's staging refresh.
2041 ///
2042 /// A retry loop must yield or sleep between attempts. On the WARP
2043 /// software adapter the CPU threads are the GPU, so a tight loop starves
2044 /// the very copy it is waiting for.
2045 ///
2046 /// [`CpuAccess::Write`] is refused on a Windows D3D11 texture tensor whose
2047 /// window is shorter than its backing, exactly as it is for
2048 /// [`map_with`](Self::map_with); see that method.
2049 fn try_map_with<'a>(&self, access: CpuAccess) -> Result<crate::view::HostView<'a, T>>
2050 where
2051 T: 'a,
2052 {
2053 self.map_with(access)
2054 }
2055
2056 /// Non-blocking [`map`](Self::map): read-write, and `WouldBlock` rather
2057 /// than a stall. See [`try_map_with`](Self::try_map_with).
2058 fn try_map<'a>(&self) -> Result<crate::view::HostView<'a, T>>
2059 where
2060 T: 'a,
2061 {
2062 self.try_map_with(CpuAccess::ReadWrite)
2063 }
2064
2065 /// Get the buffer identity for cache keying and liveness tracking.
2066 fn buffer_identity(&self) -> &BufferIdentity;
2067
2068 /// Create a zero-copy sub-region view of this backing that shares the
2069 /// underlying allocation **and** [`BufferIdentity`].
2070 ///
2071 /// The window is `[offset_bytes, offset_bytes + shape.product() *
2072 /// size_of::<T>())` measured from this tensor's own logical start, so a
2073 /// sub-view of a sub-view composes by adding offsets. Sharing the parent's
2074 /// identity is the contract that lets identity-keyed caches (e.g. the GL
2075 /// EGLImage import cache) treat offset-distinct windows as one buffer rather
2076 /// than unrelated allocations — `view` must never mint a fresh identity.
2077 ///
2078 /// Defaults to [`Error::NotImplemented`]; every backend that supports
2079 /// sub-views overrides it (`Mem`, `Shm`, Linux DMA, macOS IOSurface, `Pbo`).
2080 fn view(&self, offset_bytes: usize, shape: &[usize]) -> Result<Self>
2081 where
2082 Self: Sized,
2083 {
2084 let _ = (offset_bytes, shape);
2085 Err(Error::NotImplemented(
2086 "view (zero-copy sub-region) is not supported for this tensor backend".to_owned(),
2087 ))
2088 }
2089}
2090
2091pub trait TensorMapTrait<T>
2092where
2093 T: Num + Clone + fmt::Debug,
2094{
2095 /// Get the shape of this tensor map.
2096 fn shape(&self) -> &[usize];
2097
2098 /// Unmap the tensor from memory.
2099 fn unmap(&mut self);
2100
2101 /// Get the number of elements in this tensor map.
2102 fn len(&self) -> usize {
2103 self.shape().iter().product()
2104 }
2105
2106 /// Check if the tensor map is empty.
2107 fn is_empty(&self) -> bool {
2108 self.len() == 0
2109 }
2110
2111 /// Get the size in bytes of this tensor map.
2112 fn size(&self) -> usize {
2113 self.len() * std::mem::size_of::<T>()
2114 }
2115
2116 /// Get a slice to the data in this tensor map.
2117 fn as_slice(&self) -> &[T];
2118
2119 /// Get a mutable slice to the data in this tensor map.
2120 fn as_mut_slice(&mut self) -> &mut [T];
2121
2122 /// Whether this mapping permits writes -- i.e. whether the coherency
2123 /// bracket it holds will flush on release.
2124 ///
2125 /// The single source of truth for a binding that has to advertise
2126 /// mutability to a foreign consumer (the Python buffer protocol's
2127 /// `readonly` flag, say). Deriving that from the access a caller
2128 /// *requested* instead would leave the advertisement agreeing with the
2129 /// request while both disagreed with the mapping actually taken.
2130 fn is_writable(&self) -> bool;
2131
2132 #[cfg(feature = "ndarray")]
2133 /// Get an ndarray ArrayView of the tensor data.
2134 fn view(&'_ self) -> Result<ndarray::ArrayView<'_, T, ndarray::Dim<ndarray::IxDynImpl>>> {
2135 Ok(ndarray::ArrayView::from_shape(
2136 self.shape(),
2137 self.as_slice(),
2138 )?)
2139 }
2140
2141 #[cfg(feature = "ndarray")]
2142 /// Get an ndarray ArrayViewMut of the tensor data.
2143 fn view_mut(
2144 &'_ mut self,
2145 ) -> Result<ndarray::ArrayViewMut<'_, T, ndarray::Dim<ndarray::IxDynImpl>>> {
2146 let shape = self.shape().to_vec();
2147 Ok(ndarray::ArrayViewMut::from_shape(
2148 shape,
2149 self.as_mut_slice(),
2150 )?)
2151 }
2152}
2153
2154// Declared through `ef_vocabulary!` (see `vocabulary.rs`). These codes are
2155// the single numbering of backing stores across the stack: `code()`, the
2156// `tensor_memory_wire` constants, the C header's `ef_storage_kind`, the
2157// Python `TensorMemory` enum and the cross-package descriptor's
2158// `TensorDesc::kind` (`protocol::kind` re-exports `tensor_memory_wire`) all
2159// read the same declared literal per variant.
2160crate::ef_vocabulary! {
2161 /// Which memory backend a tensor is (or should be) allocated from.
2162 ///
2163 /// Pass `Some(..)` to a constructor to pin the backend, or `None` to
2164 /// auto-select. Auto-selection differs by constructor: [`Tensor::new`]
2165 /// tries `DmaBuf` → `Mem`, while the image constructors try `DmaBuf` →
2166 /// `Mem` and skip `Shm`. `EDGEFIRST_TENSOR_FORCE_MEM=1` short-circuits
2167 /// either chain to `Mem`, which is how tests run on hosts without
2168 /// DMA-heap permissions.
2169 ///
2170 /// A pinned request has no fallback: if the backend cannot serve it, the
2171 /// constructor fails rather than quietly giving you something slower.
2172 /// [`TensorTrait::memory`] reports what a tensor actually ended up with.
2173 ///
2174 /// # Every variant exists on every platform
2175 ///
2176 /// This is a namespace of codes, not a list of what this build can
2177 /// allocate -- [`TensorMemory::is_available`] answers the latter, at
2178 /// runtime. `Shm` used to be `#[cfg(unix)]`, which shifted every
2179 /// variant declared after it on a non-unix build and made the
2180 /// numbering a compile-time detail of the host. It also made a
2181 /// cross-platform recording unreadable: a blob captured on Linux
2182 /// carries `"dmabuf"`, and a macOS consumer has to be able to NAME
2183 /// that backing in order to report "cannot materialise here" instead
2184 /// of failing with an unrecognised code.
2185 ///
2186 /// # `#[non_exhaustive]`
2187 ///
2188 /// Match this with a wildcard arm. The variant list is expected to
2189 /// keep growing -- `IoSurface` and `Cuda` were added here, Windows
2190 /// backings and a real CUDA device backing are named as coming -- and
2191 /// this attribute is what keeps a closed-source consumer's
2192 /// `match memory { .. }` compiling across those additions. It protects
2193 /// DOWNSTREAM, not this crate: inside `edgefirst-tensor` the matches
2194 /// stay exhaustive and compile-enforced. Same rule as [`DType`],
2195 /// [`Compression`] and [`PixelFormat`], which `TensorMemory` was the
2196 /// lone exception to.
2197 #[non_exhaustive]
2198 pub enum TensorMemory {
2199 /// Regular system memory allocation. Available everywhere; the
2200 /// floor every auto-selection chain ends on.
2201 Mem = 0, "mem", MEM,
2202 /// POSIX Shared Memory allocation. Suitable for inter-process
2203 /// communication, but not for hardware acceleration. Only
2204 /// allocatable on unix; nameable everywhere.
2205 Shm = 1, "shm", SHM,
2206 /// Platform-native zero-copy GPU buffer.
2207 ///
2208 /// On Linux this is a DMA-BUF (`DmaTensor` in
2209 /// `crates/tensor/src/dma.rs`) allocated via the DRM/dma-heap
2210 /// subsystem. On macOS/iOS it is currently also this variant that
2211 /// yields an IOSurface, and on Android an AHardwareBuffer: they
2212 /// share the `TensorStorage::Dma` slot at the trait level, so the
2213 /// one code covers all of them. `IoSurface` and `D3d11Texture`
2214 /// below are the codes for naming those backings specifically; no
2215 /// backend reports either as its `memory()` yet.
2216 ///
2217 /// Allows hardware-accelerated paths (OpenGL backend on Linux via
2218 /// `EGL_EXT_image_dma_buf_import`; macOS via
2219 /// `EGL_ANGLE_iosurface_client_buffer`). CPU access via `map()`
2220 /// incurs cache-coherency overhead on Linux DMA-BUF and is similar
2221 /// in cost on IOSurface; SHM/Mem are cheaper for CPU-only workloads.
2222 ///
2223 /// On Windows it is an `ID3D11Texture2D` on the process D3D11 device
2224 /// (`crates/tensor/src/d3d11/`), allocated only through
2225 /// [`Tensor::image`](crate::Tensor::image); CPU access goes through a
2226 /// staging copy. The first such allocation creates the process
2227 /// device, so a host supplying its own device must call
2228 /// `d3d11::use_external_device` before it -- and before
2229 /// [`crate::is_gpu_buffer_available`], which creates the device too.
2230 DmaBuf = 2, "dmabuf", DMABUF,
2231 /// Apple IOSurface, named specifically rather than through the
2232 /// portable `DmaBuf` spelling.
2233 ///
2234 /// Defined as a code so a descriptor or recording that names an
2235 /// IOSurface is parseable on every platform. **No backend produces
2236 /// or accepts it yet** -- macOS/iOS allocate and report `DmaBuf`.
2237 /// Pinning a request to it fails with `NotImplemented`.
2238 IoSurface = 3, "iosurface", IOSURFACE,
2239 /// OpenGL Pixel Buffer Object memory. Created by `ImageProcessor`
2240 /// when DMA-buf is unavailable but OpenGL is present -- never by
2241 /// [`Tensor::new`], which has no GL context to create it in.
2242 Pbo = 4, "pbo", PBO,
2243 /// CUDA device memory: a device pointer, not host-addressable.
2244 ///
2245 /// Defined as a code for the same reason as `IoSurface`. **No
2246 /// backend produces or accepts it yet**; pinning a request to it
2247 /// fails with `NotImplemented`.
2248 Cuda = 5, "cuda", CUDA,
2249 /// Windows `ID3D11Texture2D`, named specifically rather than through
2250 /// the portable `DmaBuf` spelling.
2251 ///
2252 /// This is the code a Windows producer writes into
2253 /// [`TensorDesc::kind`](crate::TensorDesc::kind), where the handle
2254 /// field means an NT shared handle rather than an fd. As a
2255 /// [`TensorMemory`] it is defined for the same reason as
2256 /// `IoSurface`: **no backend allocates or reports it** -- Windows
2257 /// allocates and reports `DmaBuf` -- and pinning a request to it
2258 /// fails with `NotImplemented`.
2259 D3d11Texture = 6, "d3d11", D3D11_TEXTURE,
2260 }
2261 // `#[doc(hidden)]`: `pub` so the const-only form is reachable by an FFI
2262 // or cbindgen consumer, but it is emission plumbing rather than a second
2263 // documented API -- `TensorMemory::code()` is the documented path, and
2264 // the two are generated from the same literal.
2265 #[doc(hidden)]
2266 pub mod tensor_memory_wire;
2267}
2268
2269impl TensorMemory {
2270 /// Whether this backing can actually be allocated in this process, now.
2271 ///
2272 /// A runtime question, not a compile-time one, which is why it is a
2273 /// method on the vocabulary rather than a `cfg` on the variant:
2274 /// `DmaBuf` needs `/dev/dma_heap` to exist *and* this process to have
2275 /// permission on it (a container, a user outside the `video`/`render`
2276 /// groups, and a stock desktop give three different answers on the same
2277 /// kernel), and `Shm` needs a writable `/dev/shm`. Each probe allocates
2278 /// once and caches.
2279 ///
2280 /// Reports what can be *allocated*, which is not the same as what can be
2281 /// *named*: every variant is nameable on every platform (that is the
2282 /// point of the vocabulary), and the ones this build cannot serve say so
2283 /// here rather than by being absent from the enum.
2284 ///
2285 /// Four variants answer `false` unconditionally today, for reasons
2286 /// worth stating rather than hiding behind a probe that cannot fail:
2287 ///
2288 /// * `IoSurface` -- no backend produces or accepts it yet; macOS/iOS
2289 /// serve IOSurface under the `DmaBuf` spelling, so `DmaBuf` is the
2290 /// variant that answers `true` there.
2291 /// * `Pbo` -- PBO tensors are created by `ImageProcessor::create_image`
2292 /// with a GL context current, never by [`Tensor::new`]. This crate
2293 /// holds no GL context and cannot see one; ask the image crate
2294 /// whether a GL backend is live.
2295 /// * `Cuda` -- there is no CUDA backing in this crate at all. Probing
2296 /// for `libcuda` would answer `true` on a Jetson while every
2297 /// allocation still failed, which is worse than answering `false`.
2298 /// * `D3d11Texture` -- same as `IoSurface`: Windows serves D3D11
2299 /// textures under the `DmaBuf` spelling.
2300 ///
2301 /// One variant answers `true` for a narrower constructor than the rest:
2302 /// on Windows `DmaBuf` is an `ID3D11Texture2D`, which is image-formatted,
2303 /// so `true` there means [`Tensor::image`] can serve it and
2304 /// [`Tensor::new`] with `Some(DmaBuf)` still answers `NotImplemented`.
2305 /// [`is_gpu_buffer_available`](crate::is_gpu_buffer_available), which this
2306 /// forwards to, states the same split.
2307 pub fn is_available(self) -> bool {
2308 match self {
2309 // Always: the fallback every constructor chain ends on. If this
2310 // could fail, allocation would have no floor left.
2311 TensorMemory::Mem => true,
2312 TensorMemory::Shm => is_shm_available(),
2313 // Today `DmaBuf` is the portable spelling of "platform-native
2314 // zero-copy GPU buffer", so it is the portable probe that
2315 // answers for it -- true on macOS/iOS (IOSurface) and Android
2316 // (AHardwareBuffer) as well as Linux.
2317 TensorMemory::DmaBuf => is_gpu_buffer_available(),
2318 TensorMemory::IoSurface
2319 | TensorMemory::Pbo
2320 | TensorMemory::Cuda
2321 | TensorMemory::D3d11Texture => false,
2322 }
2323 }
2324
2325 /// The error to return when a caller pins a request to a backing this
2326 /// build knows the name of but cannot serve.
2327 ///
2328 /// Deliberately not `unreachable!()`. `IoSurface` and `Cuda` are
2329 /// unreachable *today* only because no backend reports them yet; the
2330 /// day one does, an `unreachable!()` written now becomes a panic in
2331 /// library code reached from an ordinary caller request. A "this build
2332 /// cannot do that" error is the same information without the crash, and
2333 /// it is what every other unservable request here already returns.
2334 ///
2335 /// `static`-only: every call site is inside `impl<T> TensorStorage<T>`
2336 /// (`TensorStorage::new`), itself `#[cfg(not(edgefirst_tensor_dynamic))]`.
2337 #[cfg(not(edgefirst_tensor_dynamic))]
2338 pub(crate) fn unsupported_here(self) -> Error {
2339 let reason = match self {
2340 // Not reachable: `Mem` is servable on every target. Stated
2341 // rather than panicked on, for the reason above.
2342 TensorMemory::Mem => "Mem is servable on every platform, so this should not happen",
2343 TensorMemory::Shm => "POSIX shared memory is unix-only",
2344 TensorMemory::DmaBuf => {
2345 "this target has no platform GPU-buffer allocator (DMA-BUF is Linux, \
2346 IOSurface macOS/iOS, AHardwareBuffer Android, D3D11 textures on \
2347 Windows)"
2348 }
2349 TensorMemory::IoSurface => {
2350 "no backend produces or accepts IOSurface under its own code yet -- \
2351 macOS/iOS allocate and report TensorMemory::DmaBuf"
2352 }
2353 TensorMemory::Pbo => {
2354 "PBO tensors are created by ImageProcessor::create_image with a GL \
2355 context current, never by Tensor::new"
2356 }
2357 TensorMemory::Cuda => "there is no CUDA backing in this build",
2358 TensorMemory::D3d11Texture => {
2359 "no backend produces or accepts a D3D11 texture under its own code yet -- \
2360 Windows allocates and reports TensorMemory::DmaBuf"
2361 }
2362 };
2363 Error::NotImplemented(format!(
2364 "TensorMemory::{self:?} is not supported by this build: {reason}"
2365 ))
2366 }
2367}
2368
2369impl From<TensorMemory> for String {
2370 /// The wire string, via [`TensorMemory::as_str`] -- one spelling, not a
2371 /// second one maintained alongside it.
2372 ///
2373 /// `DmaBuf` renders as `"dmabuf"`, not the `"dma"` this impl used to
2374 /// emit. `TryFrom<&str>` still accepts `"dma"` so anything that
2375 /// persisted the old spelling keeps parsing; the benchmark result keys
2376 /// in `.github/scripts` spell it `"dma"` but build that string
2377 /// themselves and never went through this impl.
2378 fn from(memory: TensorMemory) -> Self {
2379 memory.as_str().to_owned()
2380 }
2381}
2382
2383impl TryFrom<&str> for TensorMemory {
2384 type Error = Error;
2385
2386 /// Parses [`TensorMemory::as_str`]'s spelling, plus the legacy `"dma"`
2387 /// alias for `DmaBuf`.
2388 fn try_from(s: &str) -> Result<Self> {
2389 TensorMemory::from_str_code(s)
2390 .or(if s == "dma" {
2391 Some(TensorMemory::DmaBuf)
2392 } else {
2393 None
2394 })
2395 .ok_or_else(|| Error::InvalidMemoryType(s.to_owned()))
2396 }
2397}
2398
2399// Everything from here through the end of `impl<T> TensorTrait<T> for
2400// Tensor<T>` below is the `static` backend's storage-owning `Tensor<T>`:
2401// `TensorStorage<T>` (the per-memory-backend enum) and the `Tensor<T>`
2402// struct that wraps it. `dynamic`'s `Tensor<T>` is a different, much
2403// smaller type -- a typed lens over `TensorDyn` with no storage of its
2404// own -- defined in `tensor_dyn/dynamic_tensor.rs` and re-exported below.
2405#[cfg(not(edgefirst_tensor_dynamic))]
2406#[derive(Debug)]
2407#[allow(dead_code)] // Variants are constructed by downstream crates via pub(crate) helpers
2408pub(crate) enum TensorStorage<T>
2409where
2410 T: Num + Clone + fmt::Debug + Send + Sync,
2411{
2412 /// Platform-native zero-copy GPU buffer. Inner type differs per
2413 /// target: `DmaTensor` on Linux (DMA-BUF fd), `IoSurfaceTensor` on
2414 /// macOS (CFRetained IOSurface). The shared variant name keeps the
2415 /// public `TensorMemory::DmaBuf` discriminant stable across platforms.
2416 #[cfg(target_os = "linux")]
2417 Dma(DmaTensor<T>),
2418 #[cfg(any(target_os = "macos", target_os = "ios"))]
2419 Dma(IoSurfaceTensor<T>),
2420 #[cfg(target_os = "android")]
2421 Dma(AHardwareBufferTensor<T>),
2422 #[cfg(target_os = "windows")]
2423 Dma(crate::d3d11::texture::D3d11TextureTensor<T>),
2424 #[cfg(unix)]
2425 Shm(ShmTensor<T>),
2426 Mem(MemTensor<T>),
2427 Pbo(PboTensor<T>),
2428}
2429
2430#[cfg(not(edgefirst_tensor_dynamic))]
2431impl<T> TensorStorage<T>
2432where
2433 T: Num + Clone + fmt::Debug + Send + Sync,
2434{
2435 /// The backing allocation's intrinsic physical row pitch in bytes, if it
2436 /// has one that is fixed independent of the logical shape. macOS IOSurface
2437 /// reports its 64-aligned `bytesPerRow`; other backings (Linux DMA, SHM,
2438 /// Mem, PBO) have no fixed pitch beyond the logical shape and return `None`.
2439 ///
2440 /// Used by `configure_image` to preserve the physical pitch when a reused
2441 /// pool tensor is reconfigured to a smaller logical image — so the decode
2442 /// writes rows at the surface's real stride and the GPU samples them with
2443 /// the same stride (the physical-grid / logical-ROI decoupling).
2444 ///
2445 /// Compiled only where such a backing exists; on Linux and other unix
2446 /// targets it would always answer `None`.
2447 #[cfg(any(
2448 target_os = "macos",
2449 target_os = "ios",
2450 target_os = "android",
2451 target_os = "windows"
2452 ))]
2453 pub(crate) fn backing_row_stride(&self) -> Option<usize> {
2454 match self {
2455 // Only genuine image-formatted IOSurfaces (height > 1) carry a real
2456 // per-row pitch; a generic byte-bag (height == 1) returns `None` so
2457 // `configure_image` does not adopt its whole-buffer "row" as a stride.
2458 #[cfg(any(target_os = "macos", target_os = "ios"))]
2459 TensorStorage::Dma(t) => t.image_backing_row_stride(),
2460 // Android AHardwareBuffer: same rule — only genuine 2D
2461 // image-formatted buffers (height > 1) carry a real row pitch.
2462 #[cfg(target_os = "android")]
2463 TensorStorage::Dma(t) => t.image_backing_row_stride(),
2464 // Windows D3D11: the driver's staging row pitch. `None` when the
2465 // tensor has no staging (write-only or hardware-only), whose host
2466 // buffer has tight rows -- except for a semi-planar tensor, whose
2467 // texture is allocated as wide as that pitch for every access, so
2468 // it always has one to report.
2469 #[cfg(target_os = "windows")]
2470 TensorStorage::Dma(t) => t.image_backing_row_stride(),
2471 _ => None,
2472 }
2473 }
2474
2475 /// Create a new tensor storage with the given shape, memory type, and
2476 /// optional name. If no name is given, a random name will be generated.
2477 /// If no memory type is given, the best available memory type will be
2478 /// chosen based on the platform and environment variables.
2479 fn new(shape: &[usize], memory: Option<TensorMemory>, name: Option<&str>) -> Result<Self> {
2480 match memory {
2481 #[cfg(target_os = "linux")]
2482 Some(TensorMemory::DmaBuf) => {
2483 DmaTensor::<T>::new(shape, name).map(TensorStorage::Dma)
2484 }
2485 #[cfg(any(target_os = "macos", target_os = "ios"))]
2486 Some(TensorMemory::DmaBuf) => {
2487 IoSurfaceTensor::<T>::new(shape, name).map(TensorStorage::Dma)
2488 }
2489 #[cfg(target_os = "android")]
2490 Some(TensorMemory::DmaBuf) => {
2491 AHardwareBufferTensor::<T>::new(shape, name).map(TensorStorage::Dma)
2492 }
2493 // Windows has a DmaBuf backing, but it is a D3D11 texture: an image
2494 // layout, not a byte bag, so a shape alone cannot describe one.
2495 // Name the constructor that can, rather than quietly handing back
2496 // host memory the caller did not ask for.
2497 #[cfg(target_os = "windows")]
2498 Some(TensorMemory::DmaBuf) => Err(crate::error::Error::NotImplemented(
2499 "D3D11 texture tensors are image-formatted; use Tensor::image with \
2500 TensorMemory::DmaBuf"
2501 .to_owned(),
2502 )),
2503 #[cfg(not(any(
2504 target_os = "linux",
2505 target_os = "macos",
2506 target_os = "ios",
2507 target_os = "android",
2508 target_os = "windows"
2509 )))]
2510 Some(TensorMemory::DmaBuf) => Err(crate::error::Error::NotImplemented(
2511 "TensorMemory::DmaBuf is only available on Linux (DMA-BUF), macOS/iOS (IOSurface), \
2512 Android (AHardwareBuffer), and Windows (D3D11 textures, via Tensor::image)"
2513 .to_owned(),
2514 )),
2515 #[cfg(unix)]
2516 Some(TensorMemory::Shm) => {
2517 ShmTensor::<T>::new(shape, name).map(TensorStorage::Shm)
2518 }
2519 #[cfg(not(unix))]
2520 Some(m @ TensorMemory::Shm) => Err(m.unsupported_here()),
2521 Some(TensorMemory::Mem) => {
2522 MemTensor::<T>::new(shape, name).map(TensorStorage::Mem)
2523 }
2524 Some(TensorMemory::Pbo) => Err(crate::error::Error::NotImplemented(
2525 "PboTensor cannot be created via Tensor::new() — use ImageProcessor::create_image()".to_owned(),
2526 )),
2527 // Defined codes with no backing behind them yet. An error, not
2528 // an `unreachable!()` -- see `unsupported_here`.
2529 Some(
2530 m @ (TensorMemory::IoSurface | TensorMemory::Cuda | TensorMemory::D3d11Texture),
2531 ) => Err(m.unsupported_here()),
2532 None => {
2533 if std::env::var("EDGEFIRST_TENSOR_FORCE_MEM")
2534 .is_ok_and(|x| x != "0" && x.to_lowercase() != "false")
2535 {
2536 MemTensor::<T>::new(shape, name).map(TensorStorage::Mem)
2537 } else {
2538 // Auto-select priority: Dma > Mem. Shm is intentionally NOT
2539 // auto-selected — it offers no advantage over Mem for an
2540 // in-process tensor and Mem always succeeds, so Shm below Mem
2541 // is effectively never reached. Request Shm explicitly via
2542 // `TensorMemory::Shm` when cross-process sharing is needed.
2543 // (PBO sits between Dma and Mem but is GL-backed and created
2544 // only via `ImageProcessor::create_image`.)
2545 #[cfg(target_os = "linux")]
2546 {
2547 // Linux: Try DMA -> Mem
2548 match DmaTensor::<T>::new(shape, name) {
2549 Ok(tensor) => Ok(TensorStorage::Dma(tensor)),
2550 Err(_) => MemTensor::<T>::new(shape, name).map(TensorStorage::Mem),
2551 }
2552 }
2553 #[cfg(any(target_os = "macos", target_os = "ios"))]
2554 {
2555 // macOS/iOS: Try IOSurface -> Mem. IOSurface is the
2556 // GPU-shareable backend (zero-copy via ANGLE), filling the
2557 // same role as DMA-BUF on Linux.
2558 match IoSurfaceTensor::<T>::new(shape, name) {
2559 Ok(tensor) => Ok(TensorStorage::Dma(tensor)),
2560 Err(_) => MemTensor::<T>::new(shape, name).map(TensorStorage::Mem),
2561 }
2562 }
2563 #[cfg(target_os = "android")]
2564 {
2565 // Android: Mem. Unlike macOS (where a byte-bag
2566 // IOSurface is still GL-importable as R8) a generic
2567 // BLOB AHardwareBuffer cannot back a GPU texture, so
2568 // auto-selecting it buys nothing and costs plenty: a
2569 // gralloc allocator-HAL ioctl per allocation (orders
2570 // slower than malloc, ≥1 page + a dmabuf fd even for
2571 // tiny tensors) and a lock/unlock cache-maintenance
2572 // round trip on every map(). Zero-copy image tensors
2573 // come from `Tensor::image(..)`; callers that want a
2574 // BLOB (NNAPI handoff) request `TensorMemory::DmaBuf`
2575 // explicitly.
2576 MemTensor::<T>::new(shape, name).map(TensorStorage::Mem)
2577 }
2578 #[cfg(all(
2579 unix,
2580 not(any(
2581 target_os = "linux",
2582 target_os = "macos",
2583 target_os = "ios",
2584 target_os = "android"
2585 ))
2586 ))]
2587 {
2588 // Other Unix (BSD): Mem only (no DMA; Shm is explicit-only)
2589 MemTensor::<T>::new(shape, name).map(TensorStorage::Mem)
2590 }
2591 #[cfg(not(unix))]
2592 {
2593 // Windows/other: Mem only
2594 MemTensor::<T>::new(shape, name).map(TensorStorage::Mem)
2595 }
2596 }
2597 }
2598 }
2599 }
2600
2601 /// Create a DMA-backed tensor storage with an explicit byte size that
2602 /// may exceed `shape.product() * sizeof(T)`. Used for image tensors
2603 /// with row-padded layouts (see `DmaTensor::new_with_byte_size`).
2604 ///
2605 /// This is intentionally DMA-only: padding is only meaningful for
2606 /// buffers that will be imported as GPU textures via EGLImage. PBO,
2607 /// Shm, and Mem storage doesn't benefit from pitch alignment and
2608 /// shouldn't pay the memory cost.
2609 /// With `contiguous`, the memory comes from the CMA heap only (see
2610 /// [`ImageDesc::with_contiguous`]).
2611 #[cfg(target_os = "linux")]
2612 pub(crate) fn new_dma_with_byte_size(
2613 shape: &[usize],
2614 byte_size: usize,
2615 name: Option<&str>,
2616 contiguous: bool,
2617 ) -> Result<Self> {
2618 DmaTensor::<T>::new_with_byte_size(shape, byte_size, name, contiguous)
2619 .map(TensorStorage::Dma)
2620 }
2621
2622 // No non-Linux stub: the only caller (`Tensor::image_with_stride`)
2623 // returns `NotImplemented` directly on non-Linux without ever
2624 // reaching the storage layer, so defining a stub here would be
2625 // dead code and fail the `-D warnings` clippy gate on macOS CI.
2626
2627 /// Create a Mem-backed tensor storage with an explicit byte size that may
2628 /// exceed `shape.product() * sizeof(T)`. Used for image tensors with
2629 /// 64-byte-aligned row strides (see `MemTensor::with_capacity_bytes`).
2630 pub(crate) fn new_mem_with_byte_size(
2631 shape: &[usize],
2632 byte_size: usize,
2633 name: Option<&str>,
2634 ) -> Result<Self>
2635 where
2636 T: 'static,
2637 {
2638 MemTensor::<T>::with_capacity_bytes(shape, byte_size, name).map(TensorStorage::Mem)
2639 }
2640
2641 /// Create a Shm-backed tensor storage with an explicit byte size that may
2642 /// exceed `shape.product() * sizeof(T)`. Used for image tensors with
2643 /// 64-byte-aligned row strides (see `ShmTensor::new_with_byte_size`).
2644 #[cfg(unix)]
2645 pub(crate) fn new_shm_with_byte_size(
2646 shape: &[usize],
2647 byte_size: usize,
2648 name: Option<&str>,
2649 ) -> Result<Self> {
2650 ShmTensor::<T>::new_with_byte_size(shape, byte_size, name).map(TensorStorage::Shm)
2651 }
2652
2653 /// Allocate an image-formatted IOSurface-backed storage (macOS).
2654 ///
2655 /// Used by `Tensor::image()` when the caller requests
2656 /// `TensorMemory::DmaBuf` and the format has an IOSurface FourCC
2657 /// mapping (YUYV, RGBA, BGRA today). Falls back to `new_with_byte_size`
2658 /// otherwise.
2659 #[cfg(any(target_os = "macos", target_os = "ios"))]
2660 pub(crate) fn new_image_iosurface(
2661 width: usize,
2662 height: usize,
2663 format: PixelFormat,
2664 dtype: DType,
2665 shape: &[usize],
2666 name: Option<&str>,
2667 ) -> Result<Self> {
2668 IoSurfaceTensor::<T>::new_image(width, height, format, dtype, shape, name)
2669 .map(TensorStorage::Dma)
2670 }
2671
2672 /// Allocate an image-formatted AHardwareBuffer-backed storage (Android).
2673 ///
2674 /// Used by `Tensor::image()` when the caller requests
2675 /// `TensorMemory::DmaBuf` and the format has an AHardwareBuffer format
2676 /// mapping (RGBA8 and the RGBA16F float paths today). Falls back to
2677 /// `new_with_byte_size` otherwise.
2678 #[cfg(target_os = "android")]
2679 pub(crate) fn new_image_ahardwarebuffer(
2680 width: usize,
2681 height: usize,
2682 format: PixelFormat,
2683 dtype: DType,
2684 shape: &[usize],
2685 name: Option<&str>,
2686 access: CpuAccess,
2687 ) -> Result<Self> {
2688 AHardwareBufferTensor::<T>::new_image(width, height, format, dtype, shape, name, access)
2689 .map(TensorStorage::Dma)
2690 }
2691
2692 /// Allocate an image-formatted D3D11 texture storage (Windows).
2693 ///
2694 /// Used by `Tensor::image()` when the caller requests
2695 /// `TensorMemory::DmaBuf`. Unlike the IOSurface and AHardwareBuffer
2696 /// constructors there is no byte-bag fallback: a format with no texture
2697 /// layout is an error, because a texture is the only thing this backing
2698 /// can be.
2699 #[cfg(target_os = "windows")]
2700 pub(crate) fn new_image_d3d11(
2701 width: usize,
2702 height: usize,
2703 format: PixelFormat,
2704 dtype: DType,
2705 shape: &[usize],
2706 name: Option<&str>,
2707 access: CpuAccess,
2708 ) -> Result<Self> {
2709 crate::d3d11::texture::D3d11TextureTensor::<T>::new_image(
2710 width, height, format, dtype, shape, name, access,
2711 )
2712 .map(TensorStorage::Dma)
2713 }
2714
2715 /// Create a new tensor storage using the given file descriptor, shape,
2716 /// and optional name.
2717 #[cfg(unix)]
2718 fn from_fd(fd: OwnedFd, shape: &[usize], name: Option<&str>) -> Result<Self> {
2719 #[cfg(target_os = "linux")]
2720 {
2721 use nix::sys::stat::fstat;
2722 use nix::sys::statfs::fstatfs;
2723
2724 let stat = fstat(&fd)?;
2725 let major = major(stat.st_dev);
2726 let minor = minor(stat.st_dev);
2727
2728 if major != 0 {
2729 // Dma and Shm tensors are expected to have major number 0
2730 return Err(Error::UnknownDeviceType(major, minor));
2731 }
2732
2733 // Classify by filesystem magic, never by the st_dev minor.
2734 //
2735 // dma_buf files live on an internal kernel mount (`dma_buf_mnt`,
2736 // created with `kern_mount`) whose superblock draws its device
2737 // number from `get_anon_bdev()` — an IDA allocated first-come,
2738 // first-served during boot, shared with every other anonymous
2739 // pseudo-filesystem (pipefs, sockfs, anon_inodefs, nsfs, bdev,
2740 // tracefs, ...). The minor a DMA-BUF lands on is therefore a
2741 // function of which pseudo-filesystems registered first on that
2742 // kernel build: it moves with kernel config, driver load order,
2743 // initramfs use, and kernel version. It is not part of any ABI.
2744 //
2745 // Observed values for a genuine DMA-BUF: 12 on x86 desktop, 8 on
2746 // the ADIS Verdin. An earlier revision hardcoded `9 | 10` here,
2747 // which silently imported real DMA-BUFs as shared memory
2748 // everywhere else.
2749 //
2750 // The filesystem magic, by contrast, is stable UAPI
2751 // (include/uapi/linux/magic.h). Both branches are identified
2752 // positively; anything else is genuinely unknown and is rejected
2753 // rather than guessed at. Falling back to SHM would "work" — a
2754 // DMA-BUF is mmap-able, and even a pipe imports as a zero-length
2755 // tensor — which is exactly why it must not be the default.
2756 let magic = fs_magic(fstatfs(&fd)?.filesystem_type().0 as i64);
2757
2758 log::debug!(
2759 "Creating tensor from fd: major={major}, minor={minor}, magic={magic:#010x}"
2760 );
2761
2762 match magic {
2763 DMA_BUF_MAGIC => DmaTensor::<T>::from_fd(fd, shape, name).map(TensorStorage::Dma),
2764 TMPFS_MAGIC => ShmTensor::<T>::from_fd(fd, shape, name).map(TensorStorage::Shm),
2765 other => Err(Error::UnknownBufferType(other)),
2766 }
2767 }
2768 #[cfg(all(unix, not(target_os = "linux")))]
2769 {
2770 // On macOS/iOS/BSD, always use SHM (no DMA-BUF fd import)
2771 ShmTensor::<T>::from_fd(fd, shape, name).map(TensorStorage::Shm)
2772 }
2773 }
2774}
2775
2776#[cfg(not(edgefirst_tensor_dynamic))]
2777impl<T> TensorTrait<T> for TensorStorage<T>
2778where
2779 T: Num + Clone + fmt::Debug + Send + Sync,
2780{
2781 fn new(shape: &[usize], name: Option<&str>) -> Result<Self> {
2782 Self::new(shape, None, name)
2783 }
2784
2785 #[cfg(unix)]
2786 fn from_fd(fd: OwnedFd, shape: &[usize], name: Option<&str>) -> Result<Self> {
2787 Self::from_fd(fd, shape, name)
2788 }
2789
2790 #[cfg(unix)]
2791 fn clone_fd(&self) -> Result<OwnedFd> {
2792 match self {
2793 #[cfg(any(
2794 target_os = "linux",
2795 target_os = "macos",
2796 target_os = "ios",
2797 target_os = "android"
2798 ))]
2799 TensorStorage::Dma(t) => t.clone_fd(),
2800 TensorStorage::Shm(t) => t.clone_fd(),
2801 TensorStorage::Mem(t) => t.clone_fd(),
2802 TensorStorage::Pbo(t) => t.clone_fd(),
2803 }
2804 }
2805
2806 fn memory(&self) -> TensorMemory {
2807 match self {
2808 #[cfg(any(
2809 target_os = "linux",
2810 target_os = "macos",
2811 target_os = "ios",
2812 target_os = "android",
2813 target_os = "windows"
2814 ))]
2815 TensorStorage::Dma(_) => TensorMemory::DmaBuf,
2816 #[cfg(unix)]
2817 TensorStorage::Shm(_) => TensorMemory::Shm,
2818 TensorStorage::Mem(_) => TensorMemory::Mem,
2819 TensorStorage::Pbo(_) => TensorMemory::Pbo,
2820 }
2821 }
2822
2823 fn name(&self) -> String {
2824 match self {
2825 #[cfg(any(
2826 target_os = "linux",
2827 target_os = "macos",
2828 target_os = "ios",
2829 target_os = "android",
2830 target_os = "windows"
2831 ))]
2832 TensorStorage::Dma(t) => t.name(),
2833 #[cfg(unix)]
2834 TensorStorage::Shm(t) => t.name(),
2835 TensorStorage::Mem(t) => t.name(),
2836 TensorStorage::Pbo(t) => t.name(),
2837 }
2838 }
2839
2840 fn shape(&self) -> &[usize] {
2841 match self {
2842 #[cfg(any(
2843 target_os = "linux",
2844 target_os = "macos",
2845 target_os = "ios",
2846 target_os = "android",
2847 target_os = "windows"
2848 ))]
2849 TensorStorage::Dma(t) => t.shape(),
2850 #[cfg(unix)]
2851 TensorStorage::Shm(t) => t.shape(),
2852 TensorStorage::Mem(t) => t.shape(),
2853 TensorStorage::Pbo(t) => t.shape(),
2854 }
2855 }
2856
2857 fn reshape(&mut self, shape: &[usize]) -> Result<()> {
2858 match self {
2859 #[cfg(any(
2860 target_os = "linux",
2861 target_os = "macos",
2862 target_os = "ios",
2863 target_os = "android",
2864 target_os = "windows"
2865 ))]
2866 TensorStorage::Dma(t) => t.reshape(shape),
2867 #[cfg(unix)]
2868 TensorStorage::Shm(t) => t.reshape(shape),
2869 TensorStorage::Mem(t) => t.reshape(shape),
2870 TensorStorage::Pbo(t) => t.reshape(shape),
2871 }
2872 }
2873
2874 fn capacity_bytes(&self) -> usize {
2875 match self {
2876 #[cfg(any(
2877 target_os = "linux",
2878 target_os = "macos",
2879 target_os = "ios",
2880 target_os = "android",
2881 target_os = "windows"
2882 ))]
2883 TensorStorage::Dma(t) => t.capacity_bytes(),
2884 #[cfg(unix)]
2885 TensorStorage::Shm(t) => t.capacity_bytes(),
2886 TensorStorage::Mem(t) => t.capacity_bytes(),
2887 TensorStorage::Pbo(t) => t.capacity_bytes(),
2888 }
2889 }
2890
2891 fn set_logical_shape(&mut self, shape: &[usize]) -> Result<()> {
2892 match self {
2893 #[cfg(any(
2894 target_os = "linux",
2895 target_os = "macos",
2896 target_os = "ios",
2897 target_os = "android",
2898 target_os = "windows"
2899 ))]
2900 TensorStorage::Dma(t) => t.set_logical_shape(shape),
2901 #[cfg(unix)]
2902 TensorStorage::Shm(t) => t.set_logical_shape(shape),
2903 TensorStorage::Mem(t) => t.set_logical_shape(shape),
2904 TensorStorage::Pbo(t) => t.set_logical_shape(shape),
2905 }
2906 }
2907
2908 fn map_with<'a>(&self, access: CpuAccess) -> Result<crate::view::HostView<'a, T>>
2909 where
2910 T: 'a,
2911 {
2912 match self {
2913 #[cfg(any(
2914 target_os = "linux",
2915 target_os = "macos",
2916 target_os = "ios",
2917 target_os = "android",
2918 target_os = "windows"
2919 ))]
2920 TensorStorage::Dma(t) => t.map_with(access),
2921 #[cfg(unix)]
2922 TensorStorage::Shm(t) => t.map_with(access),
2923 TensorStorage::Mem(t) => t.map_with(access),
2924 TensorStorage::Pbo(t) => t.map_with(access),
2925 }
2926 }
2927
2928 fn try_map_with<'a>(&self, access: CpuAccess) -> Result<crate::view::HostView<'a, T>>
2929 where
2930 T: 'a,
2931 {
2932 match self {
2933 #[cfg(any(
2934 target_os = "linux",
2935 target_os = "macos",
2936 target_os = "ios",
2937 target_os = "android",
2938 target_os = "windows"
2939 ))]
2940 TensorStorage::Dma(t) => t.try_map_with(access),
2941 #[cfg(unix)]
2942 TensorStorage::Shm(t) => t.try_map_with(access),
2943 TensorStorage::Mem(t) => t.try_map_with(access),
2944 TensorStorage::Pbo(t) => t.try_map_with(access),
2945 }
2946 }
2947
2948 fn buffer_identity(&self) -> &BufferIdentity {
2949 match self {
2950 #[cfg(any(
2951 target_os = "linux",
2952 target_os = "macos",
2953 target_os = "ios",
2954 target_os = "android",
2955 target_os = "windows"
2956 ))]
2957 TensorStorage::Dma(t) => t.buffer_identity(),
2958 #[cfg(unix)]
2959 TensorStorage::Shm(t) => t.buffer_identity(),
2960 TensorStorage::Mem(t) => t.buffer_identity(),
2961 TensorStorage::Pbo(t) => t.buffer_identity(),
2962 }
2963 }
2964
2965 /// Forward a sub-region view to the active backend, re-wrapping the
2966 /// backend's view (which shares the parent's allocation and
2967 /// [`BufferIdentity`]) back into the matching `TensorStorage` variant. Each
2968 /// backend's `view` is its own `TensorTrait::view` override; this single
2969 /// match is the only per-variant dispatch (`Tensor::subview` calls through
2970 /// here rather than matching the storage itself).
2971 fn view(&self, offset_bytes: usize, shape: &[usize]) -> Result<Self> {
2972 match self {
2973 #[cfg(any(
2974 target_os = "linux",
2975 target_os = "macos",
2976 target_os = "ios",
2977 target_os = "android",
2978 target_os = "windows"
2979 ))]
2980 TensorStorage::Dma(t) => t.view(offset_bytes, shape).map(TensorStorage::Dma),
2981 #[cfg(unix)]
2982 TensorStorage::Shm(t) => t.view(offset_bytes, shape).map(TensorStorage::Shm),
2983 TensorStorage::Mem(t) => t.view(offset_bytes, shape).map(TensorStorage::Mem),
2984 TensorStorage::Pbo(t) => t.view(offset_bytes, shape).map(TensorStorage::Pbo),
2985 }
2986 }
2987}
2988
2989/// Multi-backend tensor with optional image format metadata.
2990///
2991/// When `format` is `Some`, this tensor represents an image. Width, height,
2992/// and channels are derived from `shape` + `format`. When `format` is `None`,
2993/// this is a raw tensor (identical to the pre-refactoring behavior).
2994///
2995/// This is the `static` backend's `Tensor<T>` -- it owns its storage. See
2996/// `tensor_dyn/dynamic_tensor.rs` for `dynamic`'s `Tensor<T>`, a typed lens
2997/// with no storage of its own.
2998#[cfg(not(edgefirst_tensor_dynamic))]
2999#[derive(Debug)]
3000pub struct Tensor<T>
3001where
3002 T: Num + Clone + fmt::Debug + Send + Sync,
3003{
3004 /// CUDA registration for this tensor, if any. Set after creation by
3005 /// the image crate once a PBO is registered with CUDA interop.
3006 ///
3007 /// MUST be declared before `storage`: CUDA must unregister the GL buffer
3008 /// before storage's Drop deletes it (cudaGraphicsUnregisterResource before
3009 /// glDeleteBuffers). Rust drops fields in declaration order.
3010 cuda: Option<crate::cuda::CudaHandle>,
3011 pub(crate) storage: TensorStorage<T>,
3012 format: Option<PixelFormat>,
3013 chroma: Option<Box<Tensor<T>>>,
3014 /// Row stride in bytes for externally allocated buffers with row padding.
3015 /// `None` means tightly packed (stride == width * bytes_per_pixel).
3016 row_stride: Option<usize>,
3017 /// Byte offset within the DMA-BUF where image data starts.
3018 /// `None` means offset 0 (data starts at the beginning of the buffer).
3019 plane_offset: Option<usize>,
3020 /// Quantization metadata for integer-typed tensors. Public access is
3021 /// gated by the `IntegerType` trait — `Tensor<f32>` etc. carry the
3022 /// field for layout uniformity but have no way to read or write it.
3023 pub(crate) quantization: Option<Quantization>,
3024 /// Optional colorimetry metadata. `None` = undefined; never auto-filled.
3025 colorimetry: Option<crate::Colorimetry>,
3026 /// Declared CPU access (see [`CpuAccess`]). Image constructors set it
3027 /// from their `access` parameter; non-image tensors (`new`, imports,
3028 /// numpy) default to `ReadWrite` — they are CPU-centric by nature.
3029 /// Maps, pins and sync brackets beyond this declaration are refused.
3030 cpu_access: CpuAccess,
3031 /// Recorded vendor tile-compression scheme — best knowledge from
3032 /// allocation time (see [`Compression`]). `Some` only for Android
3033 /// hardware-only AHardwareBuffers whose allocation requested
3034 /// compression on an eligible format with a recognized vendor. When
3035 /// set, the row-stride accessors describe no meaningful linear
3036 /// layout. `configure_image` preserves it (physical layout, unlike
3037 /// colorimetry); views inherit it.
3038 compression: Option<CompressionScheme>,
3039 /// Parent-image snapshot when this tensor is a [`view`](Self::view)/
3040 /// [`batch`](Self::batch) sub-region; `None` for a whole tensor. Lets the GL
3041 /// backend key its import on the parent and render the view as a
3042 /// `glViewport`/`glScissor` ROI. See [`ViewOrigin`].
3043 view_origin: Option<ViewOrigin>,
3044 /// Set by [`mark_view`](Self::mark_view) on a tensor rebuilt from a view's
3045 /// descriptor, which addresses part of a larger buffer but has no
3046 /// [`ViewOrigin`], and inherited by its sub-views. See
3047 /// [`is_view`](Self::is_view).
3048 view_marked: bool,
3049}
3050
3051#[cfg(not(edgefirst_tensor_dynamic))]
3052impl<T> Tensor<T>
3053where
3054 T: Num + Clone + fmt::Debug + Send + Sync,
3055{
3056 /// Wrap a TensorStorage in a Tensor with no image metadata.
3057 pub(crate) fn wrap(storage: TensorStorage<T>) -> Self {
3058 Self {
3059 storage,
3060 format: None,
3061 chroma: None,
3062 row_stride: None,
3063 plane_offset: None,
3064 quantization: None,
3065 cuda: None,
3066 colorimetry: None,
3067 cpu_access: CpuAccess::ReadWrite,
3068 compression: None,
3069 view_origin: None,
3070 view_marked: false,
3071 }
3072 }
3073
3074 /// Construct a tensor from a row-major element slice + shape. Allocates a
3075 /// new buffer (`TensorMemory::Mem`) and memcpys the contents; caller
3076 /// retains ownership of the input slice.
3077 ///
3078 /// # Errors
3079 ///
3080 /// - [`Error::InvalidShape`] if `values.len() != shape.iter().product()`.
3081 /// - Propagates any allocation error from [`Self::new`].
3082 pub fn from_slice(values: &[T], shape: &[usize]) -> Result<Self>
3083 where
3084 T: Copy,
3085 {
3086 let expected: usize = shape.iter().product();
3087 if values.len() != expected {
3088 return Err(Error::InvalidShape(format!(
3089 "from_slice: values.len()={} but shape product={expected} (shape={shape:?})",
3090 values.len()
3091 )));
3092 }
3093 let t = Self::new(shape, Some(TensorMemory::Mem), None)?;
3094 {
3095 let mut m = t.map()?;
3096 m.as_mut_slice().copy_from_slice(values);
3097 }
3098 Ok(t)
3099 }
3100
3101 /// Wrap externally-owned memory as a tensor without copying. The tensor
3102 /// borrows `[ptr, ptr + shape.product() * size_of::<T>())` as
3103 /// [`TensorMemory::Mem`]; `owner`, when `Some`, co-owns the source so it
3104 /// outlives the tensor (and all derived views/maps). See [`ForeignOwner`].
3105 ///
3106 /// The canonical use is CUDA zero-copy: allocate host-coherent memory
3107 /// (`cudaHostAlloc`), wrap the host pointer here, and bind the matching
3108 /// device pointer to the inference engine — reads and writes hit the same
3109 /// physical buffer with no host copy. The identical primitive backs the
3110 /// Python `Tensor.from_numpy` zero-copy borrow (owner = the NumPy object).
3111 ///
3112 /// # Safety
3113 ///
3114 /// `ptr` must be non-null, aligned to `align_of::<T>()`, and valid for
3115 /// `shape.product()` elements of `T` for as long as the returned tensor —
3116 /// and every view/map sharing its backing — is alive. Pass an `owner` that
3117 /// co-owns the source to uphold that contract.
3118 ///
3119 /// # Errors
3120 ///
3121 /// [`Error::InvalidSize`] if `shape` is empty.
3122 pub unsafe fn from_foreign(
3123 ptr: *mut T,
3124 shape: &[usize],
3125 owner: Option<crate::ForeignOwner>,
3126 name: Option<&str>,
3127 ) -> Result<Self> {
3128 // SAFETY: forwarded to `from_foreign_with_capacity`'s contract; `0`
3129 // clamps to the tight shape-derived footprint below, i.e. no extra
3130 // headroom -- identical to this function's own prior behaviour.
3131 unsafe { Self::from_foreign_with_capacity(ptr, shape, 0, owner, name) }
3132 }
3133
3134 /// [`Self::from_foreign`], additionally recording that the allocation
3135 /// behind `ptr` holds at least `capacity_bytes` bytes -- headroom beyond
3136 /// the tight `shape.product() * size_of::<T>()` footprint that a later
3137 /// [`configure_image`](Self::configure_image) / `set_logical_shape` can
3138 /// grow into without erroring.
3139 ///
3140 /// This is the consumer half of the cross-package capsule protocol's
3141 /// [`crate::protocol::TensorDesc::capacity`]: a producer's pool tensor
3142 /// (allocated for its largest expected image, holding a smaller one
3143 /// today) reports its real allocation size so the imported alias
3144 /// inherits the same headroom instead of being clamped to today's
3145 /// logical shape. `capacity_bytes` below the tight footprint is clamped
3146 /// up to it, never down -- the shape-implied floor always holds.
3147 ///
3148 /// # Safety
3149 ///
3150 /// Same as [`Self::from_foreign`], except `ptr` must be valid for
3151 /// `capacity_bytes` bytes (or the tight shape footprint, whichever is
3152 /// larger), not just the shape-implied footprint.
3153 ///
3154 /// # Errors
3155 ///
3156 /// [`Error::InvalidSize`] if `shape` is empty.
3157 pub unsafe fn from_foreign_with_capacity(
3158 ptr: *mut T,
3159 shape: &[usize],
3160 capacity_bytes: usize,
3161 owner: Option<crate::ForeignOwner>,
3162 name: Option<&str>,
3163 ) -> Result<Self> {
3164 if shape.is_empty() {
3165 return Err(Error::InvalidSize(0));
3166 }
3167 if ptr.is_null() {
3168 return Err(Error::InvalidArgument(
3169 "from_foreign: ptr must be non-null".to_owned(),
3170 ));
3171 }
3172 shape
3173 .iter()
3174 .copied()
3175 .try_fold(1usize, usize::checked_mul)
3176 .ok_or_else(|| {
3177 Error::InvalidArgument(format!(
3178 "from_foreign: shape.product() overflows usize (shape={shape:?})"
3179 ))
3180 })?;
3181 let capacity_elems = capacity_bytes / std::mem::size_of::<T>();
3182 // SAFETY: this function's own contract (see its doc comment) is
3183 // exactly `MemTensor::from_foreign_with_capacity`'s contract.
3184 let mem = unsafe {
3185 MemTensor::<T>::from_foreign_with_capacity(ptr, shape, capacity_elems, owner, name)
3186 };
3187 Ok(Self::wrap(TensorStorage::Mem(mem)))
3188 }
3189
3190 /// Pin a stable host address for this tensor's data.
3191 ///
3192 /// The returned [`HostPin`] carries **no borrow of the tensor**, so a
3193 /// caller can pin once at init and still pass `&mut self` to
3194 /// `ImageProcessor::convert` every frame. That borrow conflict is the
3195 /// blocker reported in [#134](https://github.com/EdgeFirstAI/hal/issues/134).
3196 ///
3197 /// The address stays valid until the pin is dropped, and survives
3198 /// re-layout: `configure_image` and `set_logical_shape` only rewrite the
3199 /// logical shape after a capacity check, and every backend's pin owns its
3200 /// own mapping, so nothing can move underneath it.
3201 ///
3202 /// A pin is therefore a **snapshot**, not a tracker: [`HostPin::len`]
3203 /// reports the extent at pin time and does not follow a later re-layout.
3204 /// That is the safe direction — the window always stays inside the
3205 /// allocation — but it means a caller who re-layouts should re-pin if it
3206 /// wants the new extent.
3207 ///
3208 /// The address is **not** a coherency guarantee: on non-coherent backends
3209 /// bracket CPU access with
3210 /// [`sync_for_cpu`](Self::sync_for_cpu) and
3211 /// [`sync_for_device`](Self::sync_for_device), which are no-ops where the
3212 /// backend does not require them.
3213 ///
3214 /// # Cost
3215 ///
3216 /// On `Mem` this is an `Arc` refcount bump — the backing allocation is
3217 /// never resized, so its base pointer is already stable. Other backends
3218 /// establish a mapping once and reuse it.
3219 ///
3220 /// # Errors
3221 ///
3222 /// - [`Error::InvalidArgument`] when `access` is [`CpuAccess::None`]: a
3223 /// pinned mapping is CPU access by definition, so this is an error
3224 /// rather than a warn-and-stage.
3225 /// - [`Error::CpuAccessDenied`] when the tensor's declared
3226 /// [`CpuAccess`] does not cover `access`.
3227 /// - [`Error::NotImplemented`] for the two backends that **cannot** pin,
3228 /// rather than merely not having been wired yet: a PBO
3229 /// (`glMapBufferRange` establishes the address and the visibility
3230 /// together, and there is no `GL_MAP_PERSISTENT_BIT` in GLES) and an
3231 /// AHardwareBuffer (`AHardwareBuffer_lock` is what yields the address at
3232 /// all, so there is no address outside the lock). Both name
3233 /// `map()`/`map_with()` as the alternative. `Mem`, `Shm`, DMA-BUF,
3234 /// IOSurface and the Windows D3D11 texture all pin — the texture through
3235 /// a persistent host buffer the storage owns, which
3236 /// [`sync_for_cpu`](Self::sync_for_cpu) and
3237 /// [`sync_for_device`](Self::sync_for_device) move bytes between and the
3238 /// texture, rather than through the staging mapping (a pin outlives every
3239 /// map guard, so holding the staging texture mapped would freeze it).
3240 pub fn pin_host<'a>(&self, access: CpuAccess) -> Result<HostPin<'a>>
3241 where
3242 T: 'a,
3243 {
3244 self.cpu_access.require_direction(access, "pin_host")?;
3245 // Backends yield everything addressable from the tensor's offset;
3246 // narrow to the logical extent so a consumer cannot mistake pitch
3247 // padding for data. A map guard keeps the wider window.
3248 let logical = self.len() * std::mem::size_of::<T>();
3249 let pin = match &self.storage {
3250 TensorStorage::Mem(m) => Ok(m.host_pin()),
3251 #[cfg(unix)]
3252 TensorStorage::Shm(s) => s.host_pin(),
3253 #[cfg(target_os = "linux")]
3254 TensorStorage::Dma(d) => d.host_pin(),
3255 #[cfg(any(target_os = "macos", target_os = "ios"))]
3256 TensorStorage::Dma(s) => s.host_pin(),
3257 // Windows D3D11: the pin is the storage's own persistent host
3258 // buffer, not the staging mapping -- see `D3d11TextureTensor::host_pin`.
3259 // The direction goes with it: a writable pin is what tells
3260 // `sync_for_device` the host buffer, not the staging copy, carries
3261 // the caller's writes.
3262 #[cfg(target_os = "windows")]
3263 TensorStorage::Dma(t) => t.host_pin(access),
3264 // AHardwareBuffer cannot support a pin at all, and that is a
3265 // platform contract rather than missing work: the NDK exposes no
3266 // way to obtain a CPU address outside AHardwareBuffer_lock, so a
3267 // "persistent" address would mean holding the lock for the pin's
3268 // lifetime -- which blocks the GPU from the buffer and makes any
3269 // further lock undefined. Both are the opposite of what a pin is
3270 // for. Say so, rather than implying a TODO.
3271 #[cfg(target_os = "android")]
3272 TensorStorage::Dma(_) => Err(Error::NotImplemented(
3273 "pin_host: AHardwareBuffer has no host address outside \
3274 AHardwareBuffer_lock, so it cannot hand out one that outlives \
3275 a guard. Use map()/map_with() for CPU access; for zero-copy \
3276 handoff pass the buffer itself."
3277 .to_owned(),
3278 )),
3279 // PBO is the same shape of constraint as AHardwareBuffer. The
3280 // message is a shared constant because `TensorDyn::pin_host`
3281 // (`tensor_dyn/dynamic_backend.rs`) must repeat this refusal --
3282 // the ABI has `ef_tensor_map` and no `ef_tensor_pin_host`, so the
3283 // library cannot tell a pin from a map. See
3284 // `crate::pbo::PIN_HOST_PBO_REFUSAL` for the reasoning at length.
3285 TensorStorage::Pbo(_) => Err(Error::NotImplemented(
3286 crate::pbo::PIN_HOST_PBO_REFUSAL.to_owned(),
3287 )),
3288 }?;
3289 Ok(pin.narrowed(logical))
3290 }
3291
3292 /// Make CPU reads/writes of this tensor coherent — the CPU is about to
3293 /// access the buffer.
3294 ///
3295 /// Corresponds to `DMA_BUF_SYNC_START`, and the name matches the Linux
3296 /// kernel DMA API for the same reason: `for_cpu` says who is about to
3297 /// touch the memory, which is what the caller actually knows.
3298 ///
3299 /// Pair with [`sync_for_device`](Self::sync_for_device).
3300 ///
3301 /// # Which backends owe maintenance
3302 ///
3303 /// | backend | behaviour |
3304 /// |---|---|
3305 /// | `Mem`, `Shm` | **no-op** — plain CPU memory, coherent by construction |
3306 /// | `Dma` (Linux DMA-BUF) | one `DMA_BUF_IOCTL_SYNC`, direction-tagged |
3307 /// | `Dma` (macOS/iOS IOSurface) | `IOSurfaceLock` / `IOSurfaceUnlock`, options matched to `access` |
3308 /// | `Dma` (Android AHardwareBuffer) | [`Error::NotImplemented`] — no sync independent of the lock |
3309 /// | `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 |
3310 /// | `Pbo` | [`Error::NotImplemented`] — no coherency window independent of the map |
3311 ///
3312 /// Only the first row is a no-op, so "bracket unconditionally and pay
3313 /// nothing" holds for `Mem`/`Shm` alone. The two that report
3314 /// `NotImplemented` do so because their platform fuses addressing and
3315 /// coherency into one call — `AHardwareBuffer_lock` and
3316 /// `glMapBufferRange` each perform the maintenance *as* they hand out the
3317 /// address, so a separate bracket would have to nest inside a map guard
3318 /// that already owns the pair. Both name `map()`/`map_with()` as the
3319 /// alternative. They previously fell through a wildcard and answered
3320 /// `Ok(())`, which is worse than an error because the symptom is
3321 /// stale-but-plausible pixels rather than a failure. The match is
3322 /// exhaustive so a new backend cannot inherit that silence — it must state
3323 /// its coherency story to compile.
3324 ///
3325 /// # Cost
3326 ///
3327 /// `Mem`/`Shm`: nothing. `Dma` on Linux: one ioctl. IOSurface: one
3328 /// lock/unlock pair. Windows D3D11 texture: one `CopyResource` plus a
3329 /// staging map per direction, and nothing at all on a tensor no one has
3330 /// pinned.
3331 ///
3332 /// # Errors
3333 ///
3334 /// [`Error::InvalidArgument`] when `access` is [`CpuAccess::None`], which
3335 /// describes no access to make coherent; [`Error::CpuAccessDenied`] when
3336 /// the tensor's declared [`CpuAccess`] does not cover `access`.
3337 pub fn sync_for_cpu(&self, access: CpuAccess) -> Result<()> {
3338 self.check_sync_access(access)?;
3339 match &self.storage {
3340 // Plain CPU memory: coherent by construction.
3341 TensorStorage::Mem(_) => Ok(()),
3342 #[cfg(unix)]
3343 TensorStorage::Shm(_) => Ok(()),
3344 #[cfg(target_os = "linux")]
3345 TensorStorage::Dma(d) => {
3346 // The direction tells the kernel which maintenance this access
3347 // needs: a read-only bracket skips the writeback at END, a
3348 // write-only one skips the invalidate at START. It MUST match
3349 // the direction passed to sync_for_device.
3350 crate::dmabuf::sync_access(&d.fd, true, access).map_err(Error::NixError)
3351 }
3352 #[cfg(any(target_os = "macos", target_os = "ios"))]
3353 TensorStorage::Dma(s) => s.sync_for_cpu(access),
3354 // Windows D3D11: refreshes the pinned host buffer from the texture
3355 // through the staging copy.
3356 #[cfg(target_os = "windows")]
3357 TensorStorage::Dma(t) => t.sync_for_cpu(access),
3358 #[cfg(target_os = "android")]
3359 TensorStorage::Dma(_) => Err(Error::NotImplemented(
3360 "sync_for_cpu: AHardwareBuffer has no sync independent of its \
3361 lock. AHardwareBuffer_lock both establishes the address and \
3362 performs the cache invalidate, and the NDK leaves locking an \
3363 already-locked buffer undefined, so a bracket on top of a map \
3364 guard would nest. Use map()/map_with(), which owns the pair."
3365 .to_owned(),
3366 )),
3367 TensorStorage::Pbo(_) => Err(Error::NotImplemented(
3368 "sync_for_cpu: a PBO has no coherency window independent of its \
3369 map -- glMapBufferRange establishes the address and the \
3370 visibility together, and glUnmapBuffer publishes the writes. \
3371 Use map()/map_with(), which owns the pair."
3372 .to_owned(),
3373 )),
3374 }
3375 }
3376
3377 /// Release the buffer back to the device — the CPU is done accessing it.
3378 ///
3379 /// Corresponds to `DMA_BUF_SYNC_END`. See
3380 /// [`sync_for_cpu`](Self::sync_for_cpu) for the pairing and cost.
3381 pub fn sync_for_device(&self, access: CpuAccess) -> Result<()> {
3382 self.check_sync_access(access)?;
3383 match &self.storage {
3384 TensorStorage::Mem(_) => Ok(()),
3385 #[cfg(unix)]
3386 TensorStorage::Shm(_) => Ok(()),
3387 #[cfg(target_os = "linux")]
3388 TensorStorage::Dma(d) => {
3389 crate::dmabuf::sync_access(&d.fd, false, access).map_err(Error::NixError)
3390 }
3391 #[cfg(any(target_os = "macos", target_os = "ios"))]
3392 TensorStorage::Dma(s) => s.sync_for_device(access),
3393 // Windows D3D11: uploads the pinned host buffer to the texture.
3394 #[cfg(target_os = "windows")]
3395 TensorStorage::Dma(t) => t.sync_for_device(access),
3396 #[cfg(target_os = "android")]
3397 TensorStorage::Dma(_) => Err(Error::NotImplemented(
3398 "sync_for_device: AHardwareBuffer has no sync independent of \
3399 its lock — see sync_for_cpu. Use map()/map_with()."
3400 .to_owned(),
3401 )),
3402 TensorStorage::Pbo(_) => Err(Error::NotImplemented(
3403 "sync_for_device: a PBO has no coherency window independent of \
3404 its map -- see sync_for_cpu. Use map()/map_with()."
3405 .to_owned(),
3406 )),
3407 }
3408 }
3409
3410 fn check_sync_access(&self, access: CpuAccess) -> Result<()> {
3411 self.cpu_access
3412 .require_direction(access, "sync_for_cpu/sync_for_device")
3413 }
3414
3415 /// Construct a tensor from a 3-D ndarray view. Respects strides — one
3416 /// copy in all cases; contiguous views take a memcpy fast path.
3417 ///
3418 /// Only available when the `ndarray` feature is enabled.
3419 #[cfg(feature = "ndarray")]
3420 pub fn from_arrayview3(view: ndarray::ArrayView3<'_, T>) -> Result<Self>
3421 where
3422 T: Copy,
3423 {
3424 let (h, w, c) = view.dim();
3425 let t = Self::new(&[h, w, c], Some(TensorMemory::Mem), None)?;
3426 {
3427 let mut m = t.map()?;
3428 let dst = m.as_mut_slice();
3429 if let Some(src) = view.as_slice() {
3430 dst.copy_from_slice(src);
3431 } else {
3432 for (d, &s) in dst.iter_mut().zip(view.iter()) {
3433 *d = s;
3434 }
3435 }
3436 }
3437 Ok(t)
3438 }
3439
3440 /// Create a new tensor with the given shape, memory type, and optional
3441 /// name. If no name is given, a random name will be generated. If no
3442 /// memory type is given, the best available memory type will be chosen
3443 /// based on the platform and environment variables.
3444 ///
3445 /// On Linux platforms, the order of preference is: Dma -> Shm -> Mem.
3446 /// On other Unix platforms (macOS), the order is: Shm -> Mem.
3447 /// On non-Unix platforms, only Mem is available.
3448 ///
3449 /// # Environment Variables
3450 /// - `EDGEFIRST_TENSOR_FORCE_MEM`: If set to a non-zero and non-false
3451 /// value, forces the use of regular system memory allocation
3452 /// (`TensorMemory::Mem`) regardless of platform capabilities.
3453 ///
3454 /// # Example
3455 /// ```rust
3456 /// use edgefirst_tensor::{Error, Tensor, TensorMemory, TensorTrait};
3457 /// # fn main() -> Result<(), Error> {
3458 /// let tensor = Tensor::<f32>::new(&[2, 3, 4], Some(TensorMemory::Mem), Some("test_tensor"))?;
3459 /// assert_eq!(tensor.memory(), TensorMemory::Mem);
3460 /// assert_eq!(tensor.name(), "test_tensor");
3461 /// # Ok(())
3462 /// # }
3463 /// ```
3464 pub fn new(shape: &[usize], memory: Option<TensorMemory>, name: Option<&str>) -> Result<Self> {
3465 let _span = tracing::trace_span!(
3466 "tensor.alloc",
3467 ?shape,
3468 memory = ?memory,
3469 dtype = std::any::type_name::<T>(),
3470 )
3471 .entered();
3472 #[cfg_attr(not(target_os = "linux"), allow(unused_mut))]
3473 let mut t = TensorStorage::new(shape, memory, name).map(Self::wrap)?;
3474 // Best-effort: attach a CUDA ExternalMemory handle for DMA tensors on
3475 // CUDA-capable hosts. Never blocks tensor creation on failure.
3476 // RUNTIME-UNVALIDATED: no CUDA+dma_heap test platform available; ABI
3477 // layout-asserted vs. CUDA 12.6 driver_types.h; mechanism proven by
3478 // gpu-probe O5 on Orin.
3479 #[cfg(target_os = "linux")]
3480 t.try_init_dma_cuda();
3481 Ok(t)
3482 }
3483
3484 /// Import an existing buffer as a tensor, taking ownership of its file
3485 /// descriptor. No bytes are copied.
3486 #[cfg(unix)]
3487 pub fn from_fd(fd: std::os::fd::OwnedFd, shape: &[usize], name: Option<&str>) -> Result<Self> {
3488 #[cfg_attr(not(target_os = "linux"), allow(unused_mut))]
3489 let mut t = Self::wrap(TensorStorage::from_fd(fd, shape, name)?);
3490 // Best-effort CUDA external memory import for DMA-backed tensors.
3491 // RUNTIME-UNVALIDATED: see try_init_dma_cuda().
3492 #[cfg(target_os = "linux")]
3493 t.try_init_dma_cuda();
3494 Ok(t)
3495 }
3496
3497 /// Create an image tensor with the given format.
3498 /// Allocate an image tensor from a declarative request — the
3499 /// full-featured constructor behind [`Self::image`] and friends.
3500 ///
3501 /// Adds the [`Compression`] request to the classic parameters:
3502 ///
3503 /// - a request with any CPU access other than [`CpuAccess::None`] is
3504 /// [`Error::InvalidArgument`] (CPU mapping pins the layout linear);
3505 /// - [`Compression::Scheme`] fails with [`Error::NotImplemented`] on
3506 /// platforms without vendor tile compression, and with
3507 /// [`Error::InvalidArgument`] when the device's native scheme or
3508 /// the format eligibility doesn't match;
3509 /// - [`Compression::Any`] never fails for compression reasons: it
3510 /// records the scheme when the allocation is eligible and
3511 /// otherwise resolves linear, incrementing
3512 /// [`compression_fallback_count`].
3513 ///
3514 /// The recorded outcome is readable via [`Tensor::compression`].
3515 ///
3516 /// A [`ImageDesc::with_contiguous`] request allocates from the CMA heap
3517 /// only, and fails instead of falling back; [`Tensor::contiguity`]
3518 /// reports the heap.
3519 pub fn image_desc(desc: &ImageDesc) -> Result<Self>
3520 where
3521 T: 'static,
3522 {
3523 let Some(t_dtype) = dtype_of::<T>() else {
3524 return Err(Error::InvalidArgument(
3525 "image_desc: element type has no DType mapping".into(),
3526 ));
3527 };
3528 if t_dtype != desc.dtype {
3529 return Err(Error::InvalidArgument(format!(
3530 "image_desc: desc.dtype is {:?} but the tensor element type is {t_dtype:?}",
3531 desc.dtype
3532 )));
3533 }
3534
3535 // A contiguous request is a CMA DMA-BUF or an error: no fallback to
3536 // another heap or another kind of memory.
3537 if desc.contiguous {
3538 #[cfg(not(target_os = "linux"))]
3539 {
3540 return Err(Error::NotImplemented(
3541 "image_desc: physically contiguous memory needs the Linux CMA DMA heap".into(),
3542 ));
3543 }
3544 #[cfg(target_os = "linux")]
3545 if !matches!(desc.memory, None | Some(TensorMemory::DmaBuf)) {
3546 return Err(Error::InvalidArgument(format!(
3547 "image_desc: physically contiguous memory is DMA-BUF from the CMA heap, \
3548 but {:?} was requested",
3549 desc.memory
3550 )));
3551 }
3552 }
3553
3554 // Compression-request guards. CPU access pins the layout linear,
3555 // so a request combined with any declared access is a
3556 // contradiction the caller should hear about immediately.
3557 if desc.compression.is_some() && desc.access != CpuAccess::None {
3558 return Err(Error::InvalidArgument(format!(
3559 "image_desc: a compression request requires CpuAccess::None (declared {:?}) — CPU mapping pins the layout linear",
3560 desc.access
3561 )));
3562 }
3563 if let Some(Compression::Scheme(requested)) = desc.compression {
3564 #[cfg(not(target_os = "android"))]
3565 {
3566 return Err(Error::NotImplemented(format!(
3567 "image_desc: Compression::Scheme({requested:?}) — no vendor tile compression on this platform (request Compression::Any for a portable fallback)"
3568 )));
3569 }
3570 #[cfg(target_os = "android")]
3571 {
3572 if matches!(
3573 desc.memory,
3574 Some(TensorMemory::Mem) | Some(TensorMemory::Shm) | Some(TensorMemory::Pbo)
3575 ) {
3576 return Err(Error::InvalidArgument(format!(
3577 "image_desc: Compression::Scheme({requested:?}) requires hardware memory (TensorMemory::DmaBuf or auto-select), got {:?}",
3578 desc.memory
3579 )));
3580 }
3581 if !crate::ahardwarebuffer_layout::compression_eligible(desc.format, desc.dtype) {
3582 return Err(Error::InvalidArgument(format!(
3583 "image_desc: ({:?}, {:?}) is not compression-eligible (RGBA8888 u8/i8 initially)",
3584 desc.format, desc.dtype
3585 )));
3586 }
3587 let device = crate::ahardwarebuffer::device_compression_scheme();
3588 if device != Some(requested) {
3589 return Err(Error::InvalidArgument(format!(
3590 "image_desc: Compression::Scheme({requested:?}) but the device's native scheme is {device:?}"
3591 )));
3592 }
3593 }
3594 }
3595
3596 // A compression request implies a hardware pipeline, so auto
3597 // memory promotes to the platform's zero-copy allocation first
3598 // (`Tensor::image` only takes the AHardwareBuffer/IOSurface path
3599 // under an explicit Dma request). `Scheme` propagates the Dma
3600 // failure — the caller demanded a layout only that allocator can
3601 // produce; `Any` falls back to plain auto-select.
3602 #[allow(unused_mut)]
3603 let mut t = if desc.contiguous {
3604 Self::image_impl(
3605 desc.width,
3606 desc.height,
3607 desc.format,
3608 Some(TensorMemory::DmaBuf),
3609 desc.access,
3610 true,
3611 )?
3612 } else {
3613 match (desc.memory, desc.compression) {
3614 (None, Some(request)) => {
3615 match Self::image(
3616 desc.width,
3617 desc.height,
3618 desc.format,
3619 Some(TensorMemory::DmaBuf),
3620 desc.access,
3621 ) {
3622 Ok(t) => t,
3623 Err(e) if matches!(request, Compression::Scheme(_)) => return Err(e),
3624 Err(_) => {
3625 Self::image(desc.width, desc.height, desc.format, None, desc.access)?
3626 }
3627 }
3628 }
3629 (memory, _) => {
3630 Self::image(desc.width, desc.height, desc.format, memory, desc.access)?
3631 }
3632 }
3633 };
3634
3635 // Record best knowledge / count fallbacks. Only an Android
3636 // hardware-only AHardwareBuffer allocation can actually hold a
3637 // vendor tile layout; everywhere else an Any request resolves
3638 // linear and is counted.
3639 if let Some(request) = desc.compression {
3640 #[cfg(target_os = "android")]
3641 {
3642 let eligible =
3643 crate::ahardwarebuffer_layout::compression_eligible(desc.format, desc.dtype);
3644 let scheme = crate::ahardwarebuffer::device_compression_scheme();
3645 let is_ahb = t.memory() == TensorMemory::DmaBuf;
3646 match (request, scheme) {
3647 (_, Some(s)) if eligible && is_ahb => {
3648 t.set_compression_unchecked(Some(s));
3649 }
3650 (Compression::Scheme(requested), _) => {
3651 // Pre-validated eligible + scheme match, so the only
3652 // way here is the allocation resolving off-AHB.
3653 return Err(Error::InvalidOperation(format!(
3654 "image_desc: Compression::Scheme({requested:?}) requested but the allocation resolved to {:?} (not an AHardwareBuffer)",
3655 t.memory()
3656 )));
3657 }
3658 (Compression::Any, _) => {
3659 note_compression_fallback(&format!(
3660 "({:?}, {:?}) {}x{}: eligible={eligible}, scheme={scheme:?}, memory={:?}",
3661 desc.format,
3662 desc.dtype,
3663 desc.width,
3664 desc.height,
3665 t.memory()
3666 ));
3667 }
3668 }
3669 }
3670 #[cfg(not(target_os = "android"))]
3671 {
3672 debug_assert!(matches!(request, Compression::Any));
3673 let _ = request;
3674 note_compression_fallback(&format!(
3675 "({:?}, {:?}) {}x{}: no vendor tile compression on this platform",
3676 desc.format, desc.dtype, desc.width, desc.height
3677 ));
3678 }
3679 }
3680 Ok(t)
3681 }
3682
3683 /// Allocate an image tensor of `width` × `height` in `format`.
3684 ///
3685 /// This is the common-case image constructor. The shape follows from the
3686 /// format's layout — packed `[H, W, C]`, planar `[C, H, W]`, semi-planar
3687 /// `[H + chroma_rows, W]` — and the format is recorded on the tensor, so
3688 /// the image crate and the codec can work from it without out-of-band
3689 /// parameters.
3690 ///
3691 /// `memory` requests a backend, or `None` to auto-select. Auto-select for
3692 /// images is **DMA → Mem**: unlike [`Tensor::new`] it does not try SHM,
3693 /// which offers an image nothing that Mem doesn't (it is not
3694 /// GPU-importable, and Mem always succeeds). Ask for
3695 /// [`TensorMemory::Shm`] explicitly if you need it.
3696 ///
3697 /// `access` declares what the **CPU** intends to do with the buffer and is
3698 /// required; see [`CpuAccess`]. Hardware access — GPU, NPU, ISP, codec —
3699 /// is always implied and never declared. Prefer the narrowest declaration
3700 /// that fits: `None` keeps the allocation eligible for Android vendor tile
3701 /// compression, and anything else pins the layout linear.
3702 ///
3703 /// DMA-backed images always get a **64-byte-aligned row stride**, in every
3704 /// layout, because Mali and Vivante reject an EGLImage import at an
3705 /// unaligned pitch. Read the real pitch back with
3706 /// [`Tensor::effective_row_stride`] rather than assuming `width × bpp`.
3707 ///
3708 /// For a compression request or any other less common knob, build an
3709 /// [`ImageDesc`] and call [`Tensor::image_desc`].
3710 ///
3711 /// # Errors
3712 ///
3713 /// - [`Error::InvalidArgument`] if `width` × `height` is not a valid size
3714 /// for `format` (for example odd dimensions where the format's chroma
3715 /// subsampling forbids them), or, on macOS with an explicit
3716 /// [`TensorMemory::DmaBuf`], if the format cannot be expressed as an
3717 /// image-formatted IOSurface at this width. The message names the
3718 /// aligned width to use.
3719 /// - Whatever the chosen backend returns if the allocation itself fails —
3720 /// for an explicit `memory` request there is no fallback.
3721 ///
3722 /// # Example
3723 ///
3724 /// ```rust
3725 /// use edgefirst_tensor::{CpuAccess, PixelFormat, Tensor, TensorMemory};
3726 ///
3727 /// # fn main() -> Result<(), edgefirst_tensor::Error> {
3728 /// // A camera frame the GPU converts and the CPU never touches.
3729 /// let frame = Tensor::<u8>::image(1920, 1080, PixelFormat::Nv12,
3730 /// Some(TensorMemory::Mem), CpuAccess::None)?;
3731 /// assert_eq!(frame.format(), Some(PixelFormat::Nv12));
3732 /// # Ok(())
3733 /// # }
3734 /// ```
3735 pub fn image(
3736 width: usize,
3737 height: usize,
3738 format: PixelFormat,
3739 memory: Option<TensorMemory>,
3740 access: CpuAccess,
3741 ) -> Result<Self>
3742 where
3743 T: 'static,
3744 {
3745 Self::image_impl(width, height, format, memory, access, false)
3746 }
3747
3748 /// [`Tensor::image`], plus the [`ImageDesc::with_contiguous`] request.
3749 /// `contiguous` is only honoured with `Some(TensorMemory::DmaBuf)` on
3750 /// Linux; [`Tensor::image_desc`] checks that before calling.
3751 fn image_impl(
3752 width: usize,
3753 height: usize,
3754 format: PixelFormat,
3755 memory: Option<TensorMemory>,
3756 access: CpuAccess,
3757 contiguous: bool,
3758 ) -> Result<Self>
3759 where
3760 T: 'static,
3761 {
3762 debug_assert!(!contiguous || memory == Some(TensorMemory::DmaBuf));
3763 #[cfg(not(target_os = "linux"))]
3764 let _ = contiguous;
3765 // Shape comes from the shared `PixelFormat::allocation_shape` helper (packed /
3766 // planar / semi-planar NV12·NV16). NV12 supports odd dimensions via the
3767 // `H + ceil(H/2)` combined-plane height.
3768 // The `T: 'static` bound is required by the macOS IOSurface path below.
3769 let shape = format.allocation_shape(width, height).ok_or_else(|| {
3770 Error::InvalidArgument(format!(
3771 "invalid dimensions {width}x{height} for format {format:?}"
3772 ))
3773 })?;
3774
3775 // macOS Dma path: allocate a format-aware IOSurface (FourCC +
3776 // 2D dimensions) so the GL backend can bind it via
3777 // `EGL_ANGLE_iosurface_client_buffer`. Without this, the IOSurface
3778 // would default to a generic byte buffer (FourCC 'L008') and
3779 // ANGLE would reject the import with `EGL_BAD_ATTRIBUTE`.
3780 //
3781 // Guard: IOSurface rounds `bytes_per_row` up to 64-byte alignment.
3782 // If the natural row pitch (`width * channels * sizeof(T)`) is not
3783 // already 64-byte aligned, the padded allocation cannot be mapped
3784 // as a contiguous packed tensor — CPU reads/writes would use the
3785 // wrong stride.
3786 //
3787 // Explicit-Dma contract: when the caller passes
3788 // `Some(TensorMemory::DmaBuf)` they have asked for an
3789 // **image-formatted IOSurface**. Silently downgrading to the
3790 // generic 'L008' byte-bag when alignment fails buries the
3791 // mismatch — the caller only finds out hours later when ANGLE
3792 // (or any GL importer) rejects the bind with
3793 // `EGL_BAD_ATTRIBUTE`. Same anti-pattern bit us previously on
3794 // Mali GPUs with DMA-BUF padding. The right behaviour is to
3795 // fail loudly here with the alignment requirement spelled out
3796 // so the caller can either pick aligned dimensions, request
3797 // SHM/Mem explicitly, or pass `memory=None` for auto-select.
3798 #[cfg(any(target_os = "macos", target_os = "ios"))]
3799 if matches!(memory, Some(TensorMemory::DmaBuf)) {
3800 // For planar formats the IOSurface stacks channels
3801 // vertically (channels * height rows), so the row stride is
3802 // single-channel width * sizeof(T). Packed formats keep the
3803 // natural width * channels * sizeof(T) stride.
3804 let natural_row_bytes = match format.layout() {
3805 PixelLayout::Planar => width * std::mem::size_of::<T>(),
3806 _ => width * format.channels() * std::mem::size_of::<T>(),
3807 };
3808 // A format with a real IOSurface FourCC (RGBA/BGRA/YUYV packed,
3809 // GREY/NV12/NV16/NV24 as R8) tolerates a non-64-aligned natural
3810 // pitch: the surface is allocated with its own 64-aligned
3811 // `bytes_per_row`, the tensor records that stride below, and a CPU
3812 // map iterates rows correctly via the strided-map path while the GL
3813 // import uses the surface's pitch directly — fully zero-copy.
3814 // Planar (the F16 RGBA16F packing) is consumed flat as
3815 // `[1, C, H, W]` with no stride, so it still requires an aligned
3816 // pitch; and formats without a FourCC would fall through to a
3817 // generic byte-bag GL can't bind. Both fail loudly rather than
3818 // silently downgrade.
3819 let has_image_fourcc = dtype_of::<T>()
3820 .and_then(|dt| crate::iosurface_layout::image_iosurface_layout(format, dt))
3821 .is_some();
3822 let padded_ok = has_image_fourcc && format.layout() != PixelLayout::Planar;
3823 if !natural_row_bytes.is_multiple_of(64) && !padded_ok {
3824 let elem_size = std::mem::size_of::<T>();
3825 let per_pixel_bytes = match format.layout() {
3826 PixelLayout::Planar => elem_size.max(1),
3827 _ => format.channels().max(1) * elem_size.max(1),
3828 };
3829 // Compute the next 64-byte-aligned width by rounding the
3830 // natural row pitch up to the next multiple of 64 and
3831 // dividing back by per-pixel bytes. This handles every
3832 // `per_pixel_bytes` value correctly:
3833 //
3834 // * Divisors of 64 (1/2/4/8/16/32/64) → the suggestion
3835 // is always 64-byte aligned.
3836 // * Non-divisors of 64 (e.g. RGB u8 with 3 B/pixel) →
3837 // the next aligned row pitch may not be an integer
3838 // multiple of per_pixel_bytes (3 doesn't divide 64
3839 // in any way), so a "pad width to N" suggestion is
3840 // structurally impossible — omit the suggestion
3841 // instead of printing a wrong number.
3842 // * per_pixel_bytes > 64 → same situation, also
3843 // omitted; the previous formula divided by zero.
3844 //
3845 // The error always names the alignment requirement
3846 // verbatim and lists the two non-DMA alternatives so
3847 // the caller has at least one always-applicable fix.
3848 let aligned_row_bytes = natural_row_bytes.next_multiple_of(64);
3849 let pad_hint =
3850 if per_pixel_bytes > 0 && aligned_row_bytes.is_multiple_of(per_pixel_bytes) {
3851 let w = aligned_row_bytes / per_pixel_bytes;
3852 format!("Pad width to {w} (the next 64-byte-aligned stride), ")
3853 } else {
3854 String::new()
3855 };
3856 return Err(Error::InvalidArgument(format!(
3857 "Tensor::image: {format:?} {width}x{height} with element \
3858 size {elem_size} produces a {natural_row_bytes}-byte natural \
3859 row pitch, which is not 64-byte aligned. \
3860 IOSurface rounds bytes_per_row up to 64 bytes, so a \
3861 contiguous CPU map of this tensor would read garbage. \
3862 {pad_hint}pass memory=None to auto-fall-back to SHM, or \
3863 pass memory=Some(TensorMemory::Shm) or \
3864 Some(TensorMemory::Mem) explicitly."
3865 )));
3866 }
3867 // Alignment OK. Explicit-Dma contract: the caller asked for an
3868 // **image-formatted, GL-importable** IOSurface, so every failure
3869 // from here on is loud. The old behaviour fell through to the
3870 // generic 'L008' byte-bag `Some(other)` arm below, and the caller
3871 // only found out when the GL import rejected the bind with
3872 // `EGL_BAD_ATTRIBUTE` — the same silent-downgrade anti-pattern as
3873 // the alignment case above. (Semi-planar/Grey u8 map to 'L008'
3874 // *by design* in `image_iosurface_layout` — the R8-plane
3875 // representation the YUV shaders sample — so they return through
3876 // the mapped path and never reach these errors.)
3877 let dtype = dtype_of::<T>().ok_or_else(|| {
3878 Error::InvalidArgument(format!(
3879 "Tensor::image: element type {} has no DType, so no \
3880 image-formatted IOSurface exists. Pass memory=None or \
3881 Some(TensorMemory::Mem) for a CPU tensor.",
3882 std::any::type_name::<T>()
3883 ))
3884 })?;
3885 if crate::iosurface_layout::image_iosurface_layout(format, dtype).is_none() {
3886 return Err(Error::InvalidArgument(format!(
3887 "Tensor::image: no zero-copy IOSurface mapping exists for \
3888 {format:?}/{dtype:?} on macOS/iOS (supported: \
3889 Rgba/Rgb @ U8/I8, Bgra/Yuyv/Grey/Nv12/Nv16/Nv24 @ U8, \
3890 Rgba/PlanarRgb/PlanarRgba @ F16). Pass memory=None to \
3891 auto-select, or Some(TensorMemory::Mem) explicitly, for \
3892 a CPU tensor."
3893 )));
3894 }
3895 // Packed RGB u8/i8 rides an RGBA8888 surface at (W*3/4, H) —
3896 // reject a width the texel packing cannot express up front
3897 // (mirrors the Android pre-guard).
3898 if format == PixelFormat::Rgb
3899 && matches!(dtype, DType::U8 | DType::I8)
3900 && packed_rgb888_layout(width, height).is_none()
3901 {
3902 return Err(Error::InvalidArgument(format!(
3903 "Tensor::image: Rgb {dtype:?} requires width%4==0 for the RGBA8888 \
3904 IOSurface packing (got width={width}). Pad the width, or pass \
3905 memory=Some(TensorMemory::Mem) for a CPU tensor."
3906 )));
3907 }
3908 let storage = TensorStorage::<T>::new_image_iosurface(
3909 width, height, format, dtype, &shape, None,
3910 )?;
3911 let mut t = Self::wrap(storage);
3912 t.format = Some(format);
3913 // IOSurface rounds `bytes_per_row` up to 64 bytes. When that
3914 // pitch exceeds the natural packed/planar row stride, record
3915 // it so CPU consumers iterate rows correctly (the GL import
3916 // already uses the surface's own pitch). For 64-aligned rows
3917 // — the common model-input case — the two match and no stride
3918 // is stored, leaving the flat mapping unchanged.
3919 if let TensorStorage::Dma(ref io) = t.storage {
3920 let bpr = io.bytes_per_row();
3921 if let Some(natural) = t.effective_row_stride() {
3922 if bpr > natural {
3923 t.set_row_stride_unchecked(bpr);
3924 }
3925 }
3926 }
3927 t.cpu_access = access;
3928 return Ok(t);
3929 }
3930
3931 // Android Dma path: allocate a format-aware AHardwareBuffer so the
3932 // GL backend can import it as an EGLImage
3933 // (`eglGetNativeClientBufferANDROID` → `eglCreateImageKHR`).
3934 //
3935 // Geometry: gralloc chooses the row pitch (`desc.stride`) at
3936 // allocation, and pads freely (validated on the Galaxy S26 Ultra,
3937 // where SnapAlloc pads the planar-F16 RGBA16F surface). A padded
3938 // pitch is recorded on the tensor so CPU maps iterate rows via the
3939 // strided-map path; the GL render uses the buffer's own pitch
3940 // through the EGLImage either way, so the GPU path stays fully
3941 // zero-copy. Consumers needing a FLAT layout (the future NPU
3942 // handoff's `[1, C, H, W]` contract) must check `row_stride()` and
3943 // repack when set — flatness is a per-device property here, unlike
3944 // macOS where IOSurface's 64-BYTE alignment keeps model-sized F16
3945 // surfaces naturally flat.
3946 #[cfg(target_os = "android")]
3947 if matches!(memory, Some(TensorMemory::DmaBuf)) {
3948 // Explicit-Dma contract (mirrors the macOS block above): the
3949 // caller asked for an image-formatted, GL-importable
3950 // AHardwareBuffer, so an unmapped combination or a gralloc
3951 // refusal errors here instead of falling through to a BLOB
3952 // byte-bag the GL backend can never import — that downgrade
3953 // only surfaced as a silent per-frame CPU upload.
3954 let dtype = dtype_of::<T>().ok_or_else(|| {
3955 Error::InvalidArgument(format!(
3956 "Tensor::image: element type {} has no DType, so no \
3957 image-formatted AHardwareBuffer exists. Pass memory=None \
3958 or Some(TensorMemory::Mem) for a CPU tensor.",
3959 std::any::type_name::<T>()
3960 ))
3961 })?;
3962 // Planar F16 requires width % 4 == 0 for the RGBA16F
3963 // packing — reject up front with the requirement spelled
3964 // out (mirrors the macOS pre-allocation guard) instead of
3965 // falling through to a byte-bag GL cannot bind.
3966 if format.layout() == PixelLayout::Planar
3967 && dtype == DType::F16
3968 && packed_rgba16f_layout(format, dtype, width, height).is_none()
3969 {
3970 return Err(Error::InvalidArgument(format!(
3971 "Tensor::image: {format:?} F16 requires width%4==0 for the RGBA16F \
3972 AHardwareBuffer packing (got width={width}). Pad the width, or pass \
3973 memory=Some(TensorMemory::Mem) for a CPU tensor."
3974 )));
3975 }
3976 // Packed RGB u8/i8 rides an RGBA8888 surface at (W*3/4, H) —
3977 // the same whole-texel constraint as the F16 packing above.
3978 if format == PixelFormat::Rgb
3979 && matches!(dtype, DType::U8 | DType::I8)
3980 && packed_rgb888_layout(width, height).is_none()
3981 {
3982 return Err(Error::InvalidArgument(format!(
3983 "Tensor::image: Rgb {dtype:?} requires width%4==0 for the RGBA8888 \
3984 AHardwareBuffer packing (got width={width}). Pad the width, or pass \
3985 memory=Some(TensorMemory::Mem) for a CPU tensor."
3986 )));
3987 }
3988 if crate::ahardwarebuffer::image_ahardwarebuffer_layout(format, dtype).is_none() {
3989 return Err(Error::InvalidArgument(format!(
3990 "Tensor::image: no zero-copy AHardwareBuffer mapping exists \
3991 for {format:?}/{dtype:?} on Android (Grey/NV* need \
3992 R8-format buffers, API 29+; the HAL floor is API 26 — see \
3993 `image_ahardwarebuffer_layout`). Pass memory=None to \
3994 auto-select, or Some(TensorMemory::Mem) explicitly, for a \
3995 CPU tensor; camera NV12 stays zero-copy by wrapping the \
3996 camera's own AHardwareBuffer instead of allocating one."
3997 )));
3998 }
3999 let storage = TensorStorage::<T>::new_image_ahardwarebuffer(
4000 width, height, format, dtype, &shape, None, access,
4001 )?;
4002 let mut t = Self::wrap(storage);
4003 t.format = Some(format);
4004 if let TensorStorage::Dma(ref ahb) = t.storage {
4005 // gralloc chooses the row pitch (Qualcomm's
4006 // SnapAlloc pads e.g. the 160-px-wide RGBA16F
4007 // surface of a 640-wide planar F16 target).
4008 // When it exceeds the natural stride, record
4009 // it so CPU consumers iterate rows correctly
4010 // via `effective_row_stride()` — the GPU
4011 // renders through the EGLImage at the
4012 // buffer's real pitch regardless, so the
4013 // render stays fully zero-copy. Consumers
4014 // that need the FLAT `[1, C, H, W]` layout
4015 // (the future NPU handoff) must check
4016 // `row_stride()` is unset and repack — or
4017 // pick an aligned width — rather than assume
4018 // flatness (see the module docs).
4019 let bpr = ahb.bytes_per_row();
4020 if let Some(natural) = t.effective_row_stride() {
4021 if bpr > natural {
4022 log::debug!(
4023 "Tensor::image: gralloc padded the {format:?} \
4024 AHardwareBuffer pitch to {bpr} bytes (natural \
4025 {natural}); recording row stride"
4026 );
4027 t.set_row_stride_unchecked(bpr);
4028 }
4029 }
4030 }
4031 t.cpu_access = access;
4032 return Ok(t);
4033 }
4034
4035 // Windows Dma path: allocate an `ID3D11Texture2D` on the process
4036 // device so the GL backend can bind it through ANGLE's D3D11 texture
4037 // import and a D3D12/CUDA consumer can open its shared NT handle.
4038 //
4039 // Explicit-Dma contract, as on macOS and Android: the caller asked for
4040 // a texture, so a format with no texture layout is an error here
4041 // rather than a silent fall-through to the host-memory arms below,
4042 // which would only surface much later as a failed GPU import.
4043 #[cfg(target_os = "windows")]
4044 if matches!(memory, Some(TensorMemory::DmaBuf)) {
4045 let dtype = dtype_of::<T>().ok_or_else(|| {
4046 Error::InvalidArgument(format!(
4047 "Tensor::image: element type {} has no DType, so no D3D11 texture \
4048 layout exists. Pass memory=None or Some(TensorMemory::Mem) for a \
4049 CPU tensor.",
4050 std::any::type_name::<T>()
4051 ))
4052 })?;
4053 let storage = TensorStorage::<T>::new_image_d3d11(
4054 width, height, format, dtype, &shape, None, access,
4055 )?;
4056 let mut t = Self::wrap(storage);
4057 t.format = Some(format);
4058 record_d3d11_row_stride(&mut t, format, "Tensor::image");
4059 t.cpu_access = access;
4060 // Best-effort CUDA import, as the Linux DMA arms do: a texture a
4061 // CUDA consumer can map is worth the one import at allocation.
4062 t.try_init_d3d11_cuda();
4063 return Ok(t);
4064 }
4065
4066 // Compute the **64-byte-aligned** row stride for every image layout.
4067 //
4068 // Embedded GPUs reject `eglCreateImage` DMA-BUF imports whose row pitch
4069 // is not 64-byte aligned: Mali returns `EGL_BAD_ALLOC`, Vivante
4070 // `EGL_BAD_ACCESS`. This bit packed RGBA/RGB destinations at odd widths
4071 // AND at even non-multiple-of-16 widths (e.g. 321→1284, 322→1288 bytes —
4072 // neither divisible by 64), so an odd-source → RGBA convert failed on
4073 // imx95/imx8mp while succeeding on V3D/Tegra. Semi-planar already aligned
4074 // here; we now align packed and planar identically so every image()
4075 // allocation is GPU-importable regardless of width.
4076 //
4077 // The per-layout natural pitch and total row count:
4078 // * SemiPlanar `[total_h, width]` — pitch = even(width)·elem, rows = total_h
4079 // * Packed `[height, width, ch]` — pitch = width·ch·elem, rows = height
4080 // * Planar `[ch, height, width]` — pitch = width·elem, rows = ch·height
4081 // Allocation byte size = `aligned_stride · total_rows` (NOT the shape
4082 // product, which reflects only the logical width and under-allocates the
4083 // padding on odd / unaligned widths).
4084 let elem = std::mem::size_of::<T>();
4085 let channels = format.channels();
4086 let (natural_stride, total_rows) = match format.layout() {
4087 PixelLayout::SemiPlanar => (width.next_multiple_of(2) * elem, shape[0]),
4088 PixelLayout::Packed => (width * channels * elem, height),
4089 PixelLayout::Planar => (width * elem, channels * height),
4090 };
4091 let aligned_stride = natural_stride.next_multiple_of(64);
4092 let semi = format.layout() == PixelLayout::SemiPlanar;
4093
4094 // DMA buffers MUST carry a 64-aligned row pitch — Mali/Vivante reject a
4095 // DMA-BUF EGLImage whose pitch is not 64-aligned. Semi-planar also needs
4096 // the aligned pitch on every backend (its chroma-plane offset math
4097 // assumes it). Packed/planar on host-only memory (Mem/Shm) keep the
4098 // natural tight pitch so the many flat CPU consumers are unaffected.
4099 let host_stride = if semi { aligned_stride } else { natural_stride };
4100 let host_byte_size = host_stride * total_rows;
4101 #[cfg(target_os = "linux")]
4102 let dma_byte_size = aligned_stride * total_rows;
4103
4104 // `used_stride` is the actual row pitch of the storage created below.
4105 let (storage, used_stride) = match memory {
4106 #[cfg(target_os = "linux")]
4107 Some(TensorMemory::DmaBuf) => (
4108 TensorStorage::<T>::new_dma_with_byte_size(
4109 &shape,
4110 dma_byte_size,
4111 None,
4112 contiguous,
4113 )?,
4114 aligned_stride,
4115 ),
4116 #[cfg(unix)]
4117 Some(TensorMemory::Shm) => (
4118 TensorStorage::<T>::new_shm_with_byte_size(&shape, host_byte_size, None)?,
4119 host_stride,
4120 ),
4121 Some(TensorMemory::Mem) => (
4122 TensorStorage::<T>::new_mem_with_byte_size(&shape, host_byte_size, None)?,
4123 host_stride,
4124 ),
4125 #[allow(unused_variables)]
4126 Some(other) => {
4127 // PBO and any future variants: fall through to standard new().
4128 return {
4129 let mut t = Self::new(&shape, Some(other), None)?;
4130 t.format = Some(format);
4131 t.cpu_access = access;
4132 Ok(t)
4133 };
4134 }
4135 None => {
4136 // Auto-select priority: DMA → Mem (DMA gets the 64-aligned
4137 // pitch; the Mem fallback keeps the tight host pitch, so the
4138 // recorded stride matches the storage used). Shm is NOT
4139 // auto-selected — it offers no advantage over Mem for an
4140 // in-process image and Mem always succeeds, so it sits below Mem
4141 // and is reached only via an explicit `TensorMemory::Shm`.
4142 #[cfg(target_os = "linux")]
4143 {
4144 match TensorStorage::<T>::new_dma_with_byte_size(
4145 &shape,
4146 dma_byte_size,
4147 None,
4148 false,
4149 ) {
4150 Ok(s) => (s, aligned_stride),
4151 Err(_) => (
4152 TensorStorage::<T>::new_mem_with_byte_size(
4153 &shape,
4154 host_byte_size,
4155 None,
4156 )?,
4157 host_stride,
4158 ),
4159 }
4160 }
4161 #[cfg(not(target_os = "linux"))]
4162 {
4163 (
4164 TensorStorage::<T>::new_mem_with_byte_size(&shape, host_byte_size, None)?,
4165 host_stride,
4166 )
4167 }
4168 }
4169 };
4170
4171 let mut t = Self::wrap(storage);
4172 t.format = Some(format);
4173 // Record the row stride when it exceeds the natural tight pitch (padding
4174 // is present — DMA packed/planar at an unaligned width, or always for
4175 // semi-planar), mirroring the IOSurface path above. Aligned-width and
4176 // host-only packed/planar images keep their flat layout with no explicit
4177 // stride; `effective_row_stride()` then falls back to the identical
4178 // computed pitch. When padding IS present, consumers must iterate rows by
4179 // `effective_row_stride()` to skip it.
4180 if semi || used_stride > natural_stride {
4181 t.set_row_stride_unchecked(used_stride);
4182 }
4183 debug_assert!(
4184 t.row_stride.is_some() || !semi,
4185 "image() must always set row_stride for semi-planar tensors"
4186 );
4187 t.cpu_access = access;
4188 #[cfg(target_os = "linux")]
4189 t.try_init_dma_cuda();
4190 Ok(t)
4191 }
4192
4193 /// Create a DMA-backed image tensor with an explicit row stride that
4194 /// may exceed the natural `width * channels * sizeof(T)` pitch.
4195 ///
4196 /// Used for image tensors that need GPU pitch alignment padding: the
4197 /// underlying DMA-BUF is sized to `row_stride * height` bytes, but
4198 /// the tensor's logical shape stays at `[height, width, channels]`.
4199 /// `width()` / `height()` / `shape()` continue to report the
4200 /// user-requested values; the padding is visible only via
4201 /// `row_stride()` / `effective_row_stride()` and is automatically
4202 /// propagated to the GL backend's EGLImage import so Mali Valhall
4203 /// accepts the buffer.
4204 ///
4205 /// # Supported formats
4206 ///
4207 /// Currently only **packed** pixel layouts (RGBA8, BGRA8, RGB888,
4208 /// Grey, etc.) are supported — the formats the GL backend uses as
4209 /// render destinations. Semi-planar formats (NV12, NV16) come from
4210 /// external allocators (camera capture, video decoders) and are
4211 /// imported via `TensorDyn::from_fd` + `set_row_stride`, which
4212 /// already supports padded strides.
4213 ///
4214 /// # Supported memory
4215 ///
4216 /// Currently only `TensorMemory::DmaBuf` is supported. PBO and Mem
4217 /// storage don't go through EGLImage import so they don't need
4218 /// pitch alignment; if you pass any other memory type this returns
4219 /// `NotImplemented`. `None` (auto-select) is treated as `Dma`.
4220 ///
4221 /// # Errors
4222 ///
4223 /// - `InvalidArgument` if `row_stride_bytes < width * channels * sizeof(T)`
4224 /// (the requested stride would not fit a single row)
4225 /// - `NotImplemented` for non-packed formats or non-DMA memory
4226 /// - `IoError` if the DMA-heap allocation fails (propagated from
4227 /// `DmaTensor::new_with_byte_size`)
4228 pub fn image_with_stride(
4229 width: usize,
4230 height: usize,
4231 format: PixelFormat,
4232 row_stride_bytes: usize,
4233 memory: Option<TensorMemory>,
4234 access: CpuAccess,
4235 ) -> Result<Self> {
4236 #[cfg(not(target_os = "linux"))]
4237 let _ = access;
4238 // DMA backing (the only thing this constructor produces) is
4239 // Linux-only. On macOS/BSD/Windows the non-Linux block below is
4240 // the only compiled body and returns `NotImplemented` directly;
4241 // on Linux the non-Linux block is cfg-removed and the function
4242 // falls through to the real validation + allocation path. Each
4243 // target compiles exactly one of the two blocks, and the block
4244 // serves as the function's tail expression in both cases — so
4245 // neither needs an explicit `return` (avoids
4246 // `clippy::needless_return` on the macOS CI gate).
4247 #[cfg(not(target_os = "linux"))]
4248 {
4249 let _ = (width, height, format, row_stride_bytes, memory);
4250 Err(Error::NotImplemented(
4251 "image_with_stride requires DMA support (Linux only)".to_owned(),
4252 ))
4253 }
4254
4255 #[cfg(target_os = "linux")]
4256 {
4257 if format.layout() != PixelLayout::Packed {
4258 return Err(Error::NotImplemented(format!(
4259 "Tensor::image_with_stride only supports packed pixel layouts, got {format:?}"
4260 )));
4261 }
4262 let elem = std::mem::size_of::<T>();
4263 let min_stride = width
4264 .checked_mul(format.channels())
4265 .and_then(|p| p.checked_mul(elem))
4266 .ok_or_else(|| {
4267 Error::InvalidArgument(format!(
4268 "image_with_stride: width {width} × channels {} × sizeof::<T>={elem} \
4269 overflows usize",
4270 format.channels()
4271 ))
4272 })?;
4273 if row_stride_bytes < min_stride {
4274 return Err(Error::InvalidArgument(format!(
4275 "image_with_stride: row_stride {row_stride_bytes} < minimum {min_stride} \
4276 ({width} px × {} ch × {elem} B)",
4277 format.channels()
4278 )));
4279 }
4280 let total_byte_size = row_stride_bytes.checked_mul(height).ok_or_else(|| {
4281 Error::InvalidArgument(format!(
4282 "image_with_stride: row_stride {row_stride_bytes} × height {height} overflows usize"
4283 ))
4284 })?;
4285
4286 let shape = vec![height, width, format.channels()];
4287
4288 let storage = match memory {
4289 Some(TensorMemory::DmaBuf) | None => TensorStorage::<T>::new_dma_with_byte_size(
4290 &shape,
4291 total_byte_size,
4292 None,
4293 false,
4294 )?,
4295 Some(other) => {
4296 return Err(Error::NotImplemented(format!(
4297 "image_with_stride: only TensorMemory::DmaBuf is supported, got {other:?}"
4298 )));
4299 }
4300 };
4301
4302 let mut t = Self::wrap(storage);
4303 t.format = Some(format);
4304 t.row_stride = Some(row_stride_bytes);
4305 t.cpu_access = access;
4306 // Match new()/from_fd(): a DMA tensor must attempt CUDA external-
4307 // memory import so a strided DMA buffer is also zero-copy
4308 // CUDA-mappable (no-op when libcudart is absent).
4309 t.try_init_dma_cuda();
4310 Ok(t)
4311 }
4312 }
4313
4314 /// Attach format metadata to an existing tensor.
4315 ///
4316 /// # Arguments
4317 ///
4318 /// * `format` - The pixel format to attach
4319 ///
4320 /// # Returns
4321 ///
4322 /// `Ok(())` on success, with the format stored as metadata on the tensor.
4323 ///
4324 /// # Errors
4325 ///
4326 /// Returns `Error::InvalidShape` if the tensor shape is incompatible with
4327 /// the format's layout (packed expects `[H, W, C]`, planar expects
4328 /// `[C, H, W]`, semi-planar expects `[H*k, W]` with format-specific
4329 /// height constraints).
4330 pub fn set_format(&mut self, format: PixelFormat) -> Result<()> {
4331 let shape = self.shape();
4332 match format.layout() {
4333 PixelLayout::Packed => {
4334 if shape.len() != 3 || shape[2] != format.channels() {
4335 return Err(Error::InvalidShape(format!(
4336 "packed format {format:?} expects [H, W, {}], got {shape:?}",
4337 format.channels()
4338 )));
4339 }
4340 }
4341 PixelLayout::Planar => {
4342 if shape.len() != 3 || shape[0] != format.channels() {
4343 return Err(Error::InvalidShape(format!(
4344 "planar format {format:?} expects [{}, H, W], got {shape:?}",
4345 format.channels()
4346 )));
4347 }
4348 }
4349 PixelLayout::SemiPlanar => {
4350 if shape.len() != 2 {
4351 return Err(Error::InvalidShape(format!(
4352 "semi-planar format {format:?} expects [H*k, W], got {shape:?}"
4353 )));
4354 }
4355 match format {
4356 // Combined-plane height is `H + ceil(H/2)` (luma + chroma
4357 // rows). For even H that is `3H/2` (≡ 0 mod 3); for odd H it
4358 // is `(3H+1)/2` (≡ 2 mod 3). Only totals ≡ 1 mod 3 are
4359 // unreachable, so reject just those — odd-height NV12 is
4360 // valid (e.g. 725 rows for a 483-tall image).
4361 PixelFormat::Nv12 if shape[0] % 3 == 1 => {
4362 return Err(Error::InvalidShape(format!(
4363 "NV12 contiguous shape[0] must be H + ceil(H/2) for some height; \
4364 {} is unreachable (≡ 1 mod 3)",
4365 shape[0]
4366 )));
4367 }
4368 PixelFormat::Nv16 if !shape[0].is_multiple_of(2) => {
4369 return Err(Error::InvalidShape(format!(
4370 "NV16 contiguous shape[0] must be even, got {}",
4371 shape[0]
4372 )));
4373 }
4374 // NV24 (4:4:4): combined-plane height is 3H (Y + 2H chroma).
4375 PixelFormat::Nv24 if !shape[0].is_multiple_of(3) => {
4376 return Err(Error::InvalidShape(format!(
4377 "NV24 contiguous shape[0] must be a multiple of 3 (= 3H), got {}",
4378 shape[0]
4379 )));
4380 }
4381 _ => {}
4382 }
4383 }
4384 }
4385 // Clear stored stride/offset when format changes — they may be invalid
4386 // for the new format. Caller must re-set after changing format.
4387 if self.format != Some(format) {
4388 self.row_stride = None;
4389 self.plane_offset = None;
4390 match self.storage {
4391 TensorStorage::Mem(ref mut m) => m.set_offset(0),
4392 #[cfg(target_os = "linux")]
4393 TensorStorage::Dma(ref mut dma) => dma.mmap_offset = 0,
4394 // Paired with the arm in `set_plane_offset`: a setter that
4395 // takes effect on this backing needs a clear that does too,
4396 // or the wrapper reports `None` while `map()` still starts at
4397 // the old offset -- a stale window rather than a lost one.
4398 // `reshape`'s clear site needs no such arm:
4399 // `IoSurfaceTensor::reshape` zeroes `view_offset` itself
4400 // before the match there runs.
4401 #[cfg(any(target_os = "macos", target_os = "ios"))]
4402 TensorStorage::Dma(ref mut io) => io.view_offset = 0,
4403 // The same pair for D3D11. Unlike `IoSurfaceTensor`,
4404 // `D3d11TextureTensor::reshape` leaves `view_offset` alone, so
4405 // `reshape`'s clear site carries this arm as well.
4406 #[cfg(target_os = "windows")]
4407 TensorStorage::Dma(ref mut tex) => tex.set_view_offset(0),
4408 // Paired with the arm in `set_plane_offset`: a setter that
4409 // takes effect on this backing needs a clear that does too.
4410 // `reshape`'s clear site needs no `Pbo` arm --
4411 // `PboTensor::reshape` zeroes `view_offset` itself before the
4412 // match there runs, exactly as `IoSurfaceTensor::reshape`
4413 // does. `set_logical_shape` deliberately does not, matching
4414 // `Mem`.
4415 TensorStorage::Pbo(ref mut p) => p.view_offset = 0,
4416 // Same pair, `Shm` side. `cfg(unix)` mirrors the `Shm`
4417 // variant's own gate on `TensorStorage`, not a per-OS
4418 // multiplex -- see the comment on the `set_plane_offset`
4419 // arm.
4420 #[cfg(unix)]
4421 TensorStorage::Shm(ref mut s) => s.set_offset(0),
4422 // No-op, pairing with `set_plane_offset`'s Android arm:
4423 // nothing writes `AHardwareBufferTensor::view_offset` there
4424 // either yet, so there is nothing here for a format change
4425 // to clear. See that arm's comment for the full accounting
4426 // of the still-incomplete gap.
4427 #[cfg(target_os = "android")]
4428 TensorStorage::Dma(_) => {}
4429 }
4430 }
4431 self.format = Some(format);
4432 Ok(())
4433 }
4434
4435 /// Set this tensor's logical dimensions and pixel format to a decoded
4436 /// image, reusing the existing allocation. The shape is derived from the
4437 /// format layout; fails with `Error::InsufficientCapacity` if the
4438 /// allocation cannot hold `width`×`height` in `format`, or
4439 /// `Error::InvalidArgument` if the dimensions are invalid for the format.
4440 ///
4441 /// For NV12/NV16/NV24 the buffer width is rounded up to even (a chroma-plane
4442 /// interleaving requirement); the true odd width is reported by the decoder
4443 /// in `ImageInfo` and trimmed by a `convert()` crop. See
4444 /// [`PixelFormat::allocation_shape`].
4445 ///
4446 /// When the backing has a fixed physical row pitch (an IOSurface's
4447 /// 64-aligned `bytesPerRow`) that exceeds the new format's natural row
4448 /// stride — i.e. a reused max-sized pool tensor reconfigured to a smaller
4449 /// image — the physical pitch is preserved as the tensor's `row_stride`.
4450 /// This keeps the **physical grid** (allocation stride/surface) fixed while
4451 /// the **logical ROI** (this image's W×H) changes, so the decode writes rows
4452 /// at the surface's real stride and the GPU samples them at the same stride.
4453 /// Exact-sized buffers (pitch == natural) stay tightly packed unchanged.
4454 pub fn configure_image(
4455 &mut self,
4456 width: usize,
4457 height: usize,
4458 format: PixelFormat,
4459 ) -> Result<()> {
4460 let shape = format.allocation_shape(width, height).ok_or_else(|| {
4461 Error::InvalidArgument(format!(
4462 "invalid dimensions {width}x{height} for format {format:?}"
4463 ))
4464 })?;
4465 // Capture the pre-existing row stride before `set_format` clears it.
4466 // For pool tensors that were allocated at a larger width (e.g. 1920-wide
4467 // pool decoding a 789-wide image), this preserves the backing pitch so
4468 // rows are still written at the correct physical stride.
4469 let prior_stride = self.row_stride;
4470
4471 self.storage.set_logical_shape(&shape)?;
4472 self.set_format(format)?;
4473 // `set_format` keeps the stride when the format is unchanged; the new
4474 // geometry's stride is chosen below, never inherited by accident.
4475 self.row_stride = None;
4476
4477 // Restore the correct row pitch for the new geometry. A DMA buffer of
4478 // ANY layout, and a semi-planar buffer on any backing, MUST carry a
4479 // 64-byte-aligned pitch: Mali/Vivante reject a DMA-BUF EGLImage whose
4480 // row pitch is not 64-aligned (`EGL_BAD_ALLOC` / `EGL_BAD_ACCESS`), and
4481 // semi-planar chroma-offset math assumes it. This mirrors the pitch
4482 // `Tensor::image()` allocates, so a recycled pool buffer imports
4483 // identically to a fresh one — without it a `configure_image()`'d packed
4484 // buffer (e.g. Y800 96-wide → tight pitch 96) failed convert() on
4485 // imx95/imx8mp via the texture-upload fallback while the fresh oracle
4486 // (aligned 128) succeeded. Packed/planar on host-only memory (Mem/Shm)
4487 // keep the tight pitch so flat CPU consumers stay unaffected — matching
4488 // `image()`'s `host_stride` rule.
4489 let elem = std::mem::size_of::<T>();
4490 let channels = format.channels();
4491 let (min_stride, total_rows) = match format.layout() {
4492 PixelLayout::SemiPlanar => (width.next_multiple_of(2) * elem, shape[0]),
4493 PixelLayout::Packed => (width * channels * elem, height),
4494 PixelLayout::Planar => (width * elem, channels * height),
4495 };
4496 let needs_align = self.storage.memory() == TensorMemory::DmaBuf
4497 || format.layout() == PixelLayout::SemiPlanar;
4498
4499 // A backing with its own fixed pitch (IOSurface, AHardwareBuffer, a
4500 // staged D3D11 texture) keeps it whenever it exceeds the new row.
4501 #[cfg(any(
4502 target_os = "macos",
4503 target_os = "ios",
4504 target_os = "android",
4505 target_os = "windows"
4506 ))]
4507 if let Some(pitch) = self.storage.backing_row_stride() {
4508 let active_stride = if pitch > min_stride {
4509 self.set_row_stride_unchecked(pitch);
4510 pitch
4511 } else {
4512 min_stride
4513 };
4514 if !needs_align {
4515 return Ok(());
4516 }
4517 return self.check_configured_stride(active_stride, total_rows);
4518 }
4519
4520 if !needs_align {
4521 return Ok(());
4522 }
4523 // Priority:
4524 // 1. Prior stride (pool reuse): if the pre-existing stride is
4525 // 64-aligned, >= this layout's minimum row, and fits the
4526 // allocation, keep it. This is the hot-loop reuse case (large
4527 // pool, small image).
4528 // 2. Compute a fresh 64-aligned stride for the current width.
4529 let aligned = min_stride.next_multiple_of(64);
4530 let capacity = self.storage.capacity_bytes();
4531
4532 let candidate = if let Some(ps) = prior_stride {
4533 if ps >= min_stride && ps % 64 == 0 && ps * total_rows <= capacity {
4534 ps
4535 } else {
4536 aligned
4537 }
4538 } else {
4539 aligned
4540 };
4541
4542 let active_stride = if candidate * total_rows <= capacity {
4543 self.set_row_stride_unchecked(candidate);
4544 candidate
4545 } else {
4546 // The tight row may still fit where the aligned one does not.
4547 min_stride
4548 };
4549 self.check_configured_stride(active_stride, total_rows)
4550 }
4551
4552 /// The last step of [`configure_image`](Self::configure_image) for an
4553 /// aligned layout: rows that do not fit the allocation are refused here,
4554 /// rather than left to fault on a later map or write.
4555 fn check_configured_stride(&self, active_stride: usize, total_rows: usize) -> Result<()> {
4556 let needed = active_stride * total_rows;
4557 let capacity = self.storage.capacity_bytes();
4558 if needed > capacity {
4559 return Err(Error::InsufficientCapacity { needed, capacity });
4560 }
4561 Ok(())
4562 }
4563
4564 /// Allocate an image tensor sized to hold up to `width`×`height` in
4565 /// `format`, reusable for any smaller image via `configure_image`.
4566 pub fn image_with_capacity(
4567 width: usize,
4568 height: usize,
4569 format: PixelFormat,
4570 memory: Option<TensorMemory>,
4571 access: CpuAccess,
4572 ) -> Result<Self>
4573 where
4574 T: 'static,
4575 {
4576 Self::image(width, height, format, memory, access)
4577 }
4578
4579 /// Pixel format (None if not an image).
4580 pub fn format(&self) -> Option<PixelFormat> {
4581 self.format
4582 }
4583
4584 /// The format and the trailing shape dims that describe one image: packed
4585 /// `[H, W, C]`, planar `[C, H, W]` or semi-planar `[H*k, W]`, optionally
4586 /// behind a leading batch `N` (see [`batch`](Self::batch)). `None` when
4587 /// unformatted or when the shape has any other rank --
4588 /// `set_logical_shape` keeps the format on a shape it no longer describes.
4589 fn image_dims(&self) -> Option<(PixelFormat, &[usize])> {
4590 let fmt = self.format?;
4591 let shape = self.shape();
4592 let rank = image_rank(fmt.layout());
4593 let batch_dims = shape.len().checked_sub(rank).filter(|&n| n <= 1)?;
4594 Some((fmt, &shape[batch_dims..]))
4595 }
4596
4597 /// Image width (None if not an image). For a batched `[N, ...]` tensor
4598 /// this is the width of one element.
4599 pub fn width(&self) -> Option<usize> {
4600 let (fmt, dims) = self.image_dims()?;
4601 match fmt.layout() {
4602 PixelLayout::Packed => Some(dims[1]),
4603 PixelLayout::Planar => Some(dims[2]),
4604 PixelLayout::SemiPlanar => Some(dims[1]),
4605 }
4606 }
4607
4608 /// Image height (None if not an image). For a batched `[N, ...]` tensor
4609 /// this is the height of one element.
4610 ///
4611 /// For semi-planar formats the combined-plane shape row count is divided
4612 /// by the format's luma-to-total ratio to recover logical height. This
4613 /// returns the exact logical height (including odd heights) only because
4614 /// the logical dimensions are tracked separately from the physical shape —
4615 /// `configure_image` stores the actual `(width, height)` in the format's
4616 /// `allocation_shape`, which round-trips losslessly via these accessors.
4617 pub fn height(&self) -> Option<usize> {
4618 let (fmt, dims) = self.image_dims()?;
4619 match fmt.layout() {
4620 PixelLayout::Packed => Some(dims[0]),
4621 PixelLayout::Planar => Some(dims[1]),
4622 PixelLayout::SemiPlanar => {
4623 if self.is_multiplane() {
4624 Some(dims[0])
4625 } else {
4626 match fmt {
4627 PixelFormat::Nv12 => Some(dims[0] * 2 / 3),
4628 PixelFormat::Nv16 => Some(dims[0] / 2),
4629 PixelFormat::Nv24 => Some(dims[0] / 3),
4630 _ => None,
4631 }
4632 }
4633 }
4634 }
4635 }
4636
4637 /// Create from separate Y and UV planes (multiplane NV12/NV16).
4638 pub fn from_planes(luma: Tensor<T>, chroma: Tensor<T>, format: PixelFormat) -> Result<Self> {
4639 if format.layout() != PixelLayout::SemiPlanar {
4640 return Err(Error::InvalidArgument(format!(
4641 "from_planes requires a semi-planar format, got {format:?}"
4642 )));
4643 }
4644 if chroma.format.is_some() || chroma.chroma.is_some() {
4645 return Err(Error::InvalidArgument(
4646 "chroma tensor must be a raw tensor (no format or chroma metadata)".into(),
4647 ));
4648 }
4649 let luma_shape = luma.shape();
4650 let chroma_shape = chroma.shape();
4651 if luma_shape.len() != 2 || chroma_shape.len() != 2 {
4652 return Err(Error::InvalidArgument(format!(
4653 "from_planes expects 2D shapes, got luma={luma_shape:?} chroma={chroma_shape:?}"
4654 )));
4655 }
4656 if luma_shape[1] != chroma_shape[1] {
4657 return Err(Error::InvalidArgument(format!(
4658 "luma width {} != chroma width {}",
4659 luma_shape[1], chroma_shape[1]
4660 )));
4661 }
4662 match format {
4663 PixelFormat::Nv12 => {
4664 if luma_shape[0] % 2 != 0 {
4665 return Err(Error::InvalidArgument(format!(
4666 "NV12 requires even luma height, got {}",
4667 luma_shape[0]
4668 )));
4669 }
4670 if chroma_shape[0] != luma_shape[0] / 2 {
4671 return Err(Error::InvalidArgument(format!(
4672 "NV12 chroma height {} != luma height / 2 ({})",
4673 chroma_shape[0],
4674 luma_shape[0] / 2
4675 )));
4676 }
4677 }
4678 PixelFormat::Nv16 => {
4679 if chroma_shape[0] != luma_shape[0] {
4680 return Err(Error::InvalidArgument(format!(
4681 "NV16 chroma height {} != luma height {}",
4682 chroma_shape[0], luma_shape[0]
4683 )));
4684 }
4685 }
4686 // NV24's chroma plane is full-resolution (2×-wide interleaved UV),
4687 // which the equal-width plane check above doesn't model. Multiplane
4688 // NV24 is unused (the JPEG decoder emits a contiguous NV24 buffer),
4689 // so it's not supported here yet.
4690 _ => {
4691 return Err(Error::InvalidArgument(format!(
4692 "from_planes only supports NV12 and NV16 (NV24 multiplane not yet \
4693 supported — use a contiguous NV24 tensor), got {format:?}"
4694 )));
4695 }
4696 }
4697
4698 Ok(Tensor {
4699 storage: luma.storage,
4700 format: Some(format),
4701 chroma: Some(Box::new(chroma)),
4702 row_stride: luma.row_stride,
4703 plane_offset: luma.plane_offset,
4704 quantization: luma.quantization,
4705 // A multiplane tensor spans two DMA-BUFs (luma + chroma); CUDA
4706 // external-memory import is per-fd, so there is no single device
4707 // pointer for the composite. Any CUDA handle the luma plane carried
4708 // is intentionally dropped — consumers needing CUDA access to
4709 // multiplane data must import each plane independently.
4710 cuda: None,
4711 colorimetry: luma.colorimetry,
4712 cpu_access: luma.cpu_access,
4713 compression: luma.compression,
4714 // A composed multiplane tensor is a whole image, not a sub-view.
4715 view_origin: None,
4716 view_marked: false,
4717 })
4718 }
4719
4720 /// Whether this tensor uses separate plane allocations.
4721 pub fn is_multiplane(&self) -> bool {
4722 self.chroma.is_some()
4723 }
4724
4725 /// Access the chroma plane for multiplane semi-planar images.
4726 pub fn chroma(&self) -> Option<&Tensor<T>> {
4727 self.chroma.as_deref()
4728 }
4729
4730 /// Mutable access to the chroma plane for multiplane semi-planar images.
4731 pub fn chroma_mut(&mut self) -> Option<&mut Tensor<T>> {
4732 self.chroma.as_deref_mut()
4733 }
4734
4735 /// Independent tensor over this plane's allocation, copying geometry
4736 /// metadata. Used to answer `ef_tensor_chroma` from the handle without a
4737 /// client-side shadow field.
4738 pub(crate) fn clone_chroma_plane(&self) -> Result<Self> {
4739 // Every storage accepted by `from_planes` implements `view` as a
4740 // shared-backing lens. Using that common path keeps
4741 // `is_multiplane() => chroma().is_some()` true for SHM, PBO and
4742 // D3D11 too, instead of silently losing chroma when a hand-written
4743 // per-storage clone omits a variant.
4744 self.subview(0, self.shape())
4745 }
4746
4747 /// Row stride in bytes (`None` = tightly packed).
4748 pub fn row_stride(&self) -> Option<usize> {
4749 self.row_stride
4750 }
4751
4752 /// Effective row stride in bytes: the stored stride if set, otherwise the
4753 /// minimum stride computed from the format, width, and element size.
4754 /// Returns `None` only when no format is set and no explicit stride was
4755 /// stored via [`set_row_stride`](Self::set_row_stride).
4756 ///
4757 /// **GREY note:** `effective_row_stride()` for a GREY tensor returns the
4758 /// tight `width` bytes (no padding), which is what `normalize_to_numpy` and
4759 /// the CPU convert path expect. The codec's internal `native_row_stride`
4760 /// (64-byte-aligned) is used only during decoding and is not propagated to
4761 /// the tensor's stored stride, so callers reading via
4762 /// `effective_row_stride()` always see the tight value for GREY.
4763 pub fn effective_row_stride(&self) -> Option<usize> {
4764 if let Some(s) = self.row_stride {
4765 return Some(s);
4766 }
4767 let fmt = self.format?;
4768 let w = self.width()?;
4769 let elem = std::mem::size_of::<T>();
4770 Some(match fmt.layout() {
4771 PixelLayout::Packed => w * fmt.channels() * elem,
4772 PixelLayout::Planar => w * elem,
4773 // Semi-planar: minimum stride must cover the even width so the
4774 // interleaved chroma columns are byte-aligned on odd-width images.
4775 PixelLayout::SemiPlanar => w.next_multiple_of(2) * elem,
4776 })
4777 }
4778
4779 /// Copy the tensor's logical bytes into `dst`, compacting away any
4780 /// recorded row-stride padding.
4781 ///
4782 /// The flatness helper for NPU handoff: consumers that need a FLAT
4783 /// layout (e.g. `[1, C, H, W]` for NNAPI/LiteRT when the runtime cannot
4784 /// take a padded pitch) call this when [`row_stride`](Self::row_stride)
4785 /// is `Some` — on a tight tensor it degenerates to one memcpy, so it is
4786 /// safe to call unconditionally. `dst.len()` must equal the tight byte
4787 /// footprint (`shape` product × element size). Zero-copy consumers
4788 /// should prefer the buffer handle + [`Tensor::effective_row_stride`]
4789 /// (Self::effective_row_stride) and skip this copy entirely.
4790 pub fn copy_to_flat(&self, dst: &mut [u8]) -> Result<()> {
4791 let tight_bytes = crate::ahardwarebuffer_layout::checked_shape_bytes::<T>(self.shape())?;
4792 if dst.len() != tight_bytes {
4793 return Err(Error::InvalidArgument(format!(
4794 "copy_to_flat: dst is {} bytes but the tensor's tight \
4795 footprint is {tight_bytes} bytes (shape {:?})",
4796 dst.len(),
4797 self.shape()
4798 )));
4799 }
4800 // Read-only: this helper never writes back into the tensor, so a
4801 // read map suffices — and is required on backends (D3D11) that
4802 // reject a writable map on a tensor allocated CpuAccess::Read.
4803 let map = self.map_read()?;
4804 // SAFETY: T is a plain numeric type (crate-wide bound); viewing the
4805 // mapped elements as bytes is sound.
4806 let src: &[u8] = unsafe {
4807 std::slice::from_raw_parts(
4808 map.as_slice().as_ptr() as *const u8,
4809 std::mem::size_of_val(map.as_slice()),
4810 )
4811 };
4812 let Some(stride) = self.row_stride else {
4813 // Tight layout: the mapped window is exactly the logical bytes.
4814 let Some(tight) = src.get(..tight_bytes) else {
4815 return Err(Error::InvalidOperation(format!(
4816 "copy_to_flat: mapped {} bytes < tight footprint {tight_bytes}",
4817 src.len()
4818 )));
4819 };
4820 dst.copy_from_slice(tight);
4821 return Ok(());
4822 };
4823 // Strided: the rows the strided map spans, each holding one tight
4824 // logical row. `rows` is a factor of the element count, so it divides
4825 // the tight footprint exactly.
4826 let rows = strided_rows(self)?;
4827 if rows == 0 {
4828 return Err(Error::InvalidOperation(
4829 "copy_to_flat: a strided tensor with no rows has nothing to copy".into(),
4830 ));
4831 }
4832 let row_bytes = tight_bytes / rows;
4833 let need = (rows - 1)
4834 .checked_mul(stride)
4835 .and_then(|b| b.checked_add(row_bytes))
4836 .ok_or_else(|| {
4837 Error::InvalidOperation(format!(
4838 "copy_to_flat: stride {stride} × rows {rows} overflows usize"
4839 ))
4840 })?;
4841 if src.len() < need {
4842 return Err(Error::InvalidOperation(format!(
4843 "copy_to_flat: mapped {} bytes but strided rows need {need}",
4844 src.len()
4845 )));
4846 }
4847 for r in 0..rows {
4848 dst[r * row_bytes..(r + 1) * row_bytes]
4849 .copy_from_slice(&src[r * stride..r * stride + row_bytes]);
4850 }
4851 Ok(())
4852 }
4853
4854 /// Set the row stride in bytes for externally allocated buffers with
4855 /// row padding (e.g. V4L2 or GStreamer allocators).
4856 ///
4857 /// The stride is propagated to the EGL DMA-BUF import attributes so
4858 /// the GPU interprets the padded buffer layout correctly. Must be
4859 /// called after [`set_format`](Self::set_format) and before the tensor
4860 /// is first passed to `ImageProcessor::convert`. The stored stride
4861 /// is cleared automatically if the pixel format is later changed.
4862 ///
4863 /// No stride-vs-buffer-size validation is performed because the
4864 /// backing allocation size is not reliably known: external DMA-BUFs
4865 /// may be over-allocated by the allocator, and internal tensors store
4866 /// a logical (unpadded) shape. An incorrect stride will be caught by
4867 /// the EGL driver at import time.
4868 ///
4869 /// # Arguments
4870 ///
4871 /// * `stride` - Row stride in bytes. Must be >= the minimum stride for
4872 /// the format (width * channels * sizeof(T) for packed,
4873 /// width * sizeof(T) for planar/semi-planar).
4874 ///
4875 /// # Errors
4876 ///
4877 /// * `InvalidArgument` if no pixel format is set on this tensor
4878 /// * `InvalidArgument` if `stride` is less than the minimum for the
4879 /// format and width
4880 pub fn set_row_stride(&mut self, stride: usize) -> Result<()> {
4881 let fmt = self.format.ok_or_else(|| {
4882 Error::InvalidArgument("cannot set row_stride without a pixel format".into())
4883 })?;
4884 let w = self.width().ok_or_else(|| {
4885 Error::InvalidArgument("cannot determine width for row_stride validation".into())
4886 })?;
4887 let elem = std::mem::size_of::<T>();
4888 let min_stride = match fmt.layout() {
4889 PixelLayout::Packed => w * fmt.channels() * elem,
4890 PixelLayout::Planar => w * elem,
4891 // Semi-planar: minimum must cover even width for chroma alignment.
4892 PixelLayout::SemiPlanar => w.next_multiple_of(2) * elem,
4893 };
4894 if stride < min_stride {
4895 return Err(Error::InvalidArgument(format!(
4896 "row_stride {stride} < minimum {min_stride} for {fmt:?} at width {w}"
4897 )));
4898 }
4899 self.row_stride = Some(stride);
4900 Ok(())
4901 }
4902
4903 /// Set the row stride without format validation.
4904 ///
4905 /// Use this for raw sub-tensors (e.g. chroma planes) that don't carry
4906 /// format metadata. The caller is responsible for ensuring the stride
4907 /// is valid.
4908 pub fn set_row_stride_unchecked(&mut self, stride: usize) {
4909 self.row_stride = Some(stride);
4910 }
4911
4912 /// Builder-style variant of [`set_row_stride`](Self::set_row_stride),
4913 /// consuming and returning `self`.
4914 ///
4915 /// # Errors
4916 ///
4917 /// Same conditions as [`set_row_stride`](Self::set_row_stride).
4918 pub fn with_row_stride(mut self, stride: usize) -> Result<Self> {
4919 self.set_row_stride(stride)?;
4920 Ok(self)
4921 }
4922
4923 /// Byte offset within the DMA-BUF where image data starts (`None` = 0).
4924 pub fn plane_offset(&self) -> Option<usize> {
4925 self.plane_offset
4926 }
4927
4928 /// The parent-image snapshot if this tensor is a [`view`](Self::view)/
4929 /// [`batch`](Self::batch) sub-region; `None` for a whole tensor. The GL
4930 /// backend keys its import on the parent geometry and renders this view as a
4931 /// `glViewport`/`glScissor` ROI at `(x, y, width, height)`. See [`ViewOrigin`].
4932 pub fn view_origin(&self) -> Option<ViewOrigin> {
4933 self.view_origin
4934 }
4935
4936 /// Whether this tensor addresses part of a larger buffer: a
4937 /// [`view`](Self::view)/[`batch`](Self::batch) with a [`ViewOrigin`], or a
4938 /// tensor rebuilt from such a view's descriptor, which carries none.
4939 ///
4940 /// A view's rows share cache lines with bytes outside it, so a CPU pass
4941 /// that writes it maps `ReadWrite` (see [`TensorTrait::write_access`]).
4942 pub fn is_view(&self) -> bool {
4943 self.view_origin.is_some() || self.view_marked
4944 }
4945
4946 /// Record that this tensor addresses part of a larger buffer, as
4947 /// [`is_view`](Self::is_view) reports. The descriptor import calls this
4948 /// for a producer that was a view; it is never cleared.
4949 pub fn mark_view(&mut self) {
4950 self.view_marked = true;
4951 }
4952
4953 /// Set the byte offset within the DMA-BUF where image data starts.
4954 ///
4955 /// Propagated to `EGL_DMA_BUF_PLANE0_OFFSET_EXT` on GPU import.
4956 /// Unlike [`set_row_stride`](Self::set_row_stride), no format is required
4957 /// since the offset is format-independent.
4958 pub fn set_plane_offset(&mut self, offset: usize) {
4959 self.plane_offset = Some(offset);
4960 // The offset consulted by `map()` lives inside the storage variant.
4961 // Keep it in sync with the wrapper field for every backing that
4962 // honors it; see also the clear sites in `set_format` and `reshape`.
4963 //
4964 // Note that `TensorStorage::Dma` is a cfg-multiplexed *name* rather
4965 // than one type -- `DmaTensor` on Linux, `IoSurfaceTensor` on
4966 // macOS/iOS, `AHardwareBufferTensor` on Android,
4967 // `D3d11TextureTensor` on Windows -- so each platform needs its own
4968 // arm below; a missing one would silently become a fall-through
4969 // rather than a compile error, which is how macOS went without one
4970 // for a while. The match is kept exhaustive (no `_ => {}`)
4971 // specifically so that hazard is a compile error, not a silent gap.
4972 match self.storage {
4973 TensorStorage::Mem(ref mut m) => m.set_offset(offset),
4974 #[cfg(target_os = "linux")]
4975 TensorStorage::Dma(ref mut dma) => dma.mmap_offset = offset,
4976 // The offset `IoSurfaceTensor::scoped_pin`/`map_inner` add to the
4977 // locked base address. Writing it here is what makes a
4978 // *reconstructed* view address its own sub-region instead of the
4979 // parent surface's origin; `IoSurfaceTensor::view` already set it
4980 // on the fresh-view path, which is why only the descriptor round
4981 // trip was ever wrong. Idempotent against `subview`, which sets
4982 // the same absolute value by both routes -- pinned by
4983 // `nested_iosurface_subviews_do_not_compound_their_plane_offset`.
4984 #[cfg(any(target_os = "macos", target_os = "ios"))]
4985 TensorStorage::Dma(ref mut io) => io.view_offset = offset,
4986 // The offset every `D3d11TextureTensor` pin adds to its host or
4987 // staging base, in backing (pitched) space like `row_stride()`.
4988 // `view()` bounds-checks the value it computes; this route does
4989 // not, so the pins refuse an offset past the backing instead
4990 // (`check_view_offset` in `d3d11/texture.rs`). Pinned by the
4991 // `d3d11` plane-offset tests in `crates/tensor/tests`.
4992 #[cfg(target_os = "windows")]
4993 TensorStorage::Dma(ref mut tex) => tex.set_view_offset(offset),
4994 // The offset `PboTensor::scoped_pin` adds to the mapped base.
4995 // Writing it here is what makes a *reconstructed* view address
4996 // its own sub-region instead of the parent buffer's origin;
4997 // `PboTensor::view` already set it on the fresh-view path, which
4998 // is why only the round trip was ever wrong -- the same shape
4999 // IOSurface had. Idempotent against `subview`, which sets the
5000 // same absolute value by both routes -- pinned by
5001 // `nested_pbo_subviews_do_not_compound_their_plane_offset`.
5002 //
5003 // Not `cfg`-gated, and deliberately so: unlike
5004 // `TensorStorage::Dma` -- a cfg-multiplexed *name* whose gated
5005 // arm silently becomes `_ => {}` on the other targets, the
5006 // mechanism behind this whole issue class -- `TensorStorage::Pbo`
5007 // is one type on every platform.
5008 TensorStorage::Pbo(ref mut p) => p.view_offset = offset,
5009 // The offset `ShmTensor::map_inner`/`host_pin` add to the mmap
5010 // base, exactly like `MemTensor::offset`. `cfg(unix)` here
5011 // mirrors `TensorStorage::Shm`'s own gate on the enum
5012 // definition -- it is not a per-OS multiplex the way
5013 // `TensorStorage::Dma` is, so this arm cannot silently become
5014 // `_ => {}` on a target where the `Shm` variant still exists.
5015 #[cfg(unix)]
5016 TensorStorage::Shm(ref mut s) => s.set_offset(offset),
5017 // STILL INCOMPLETE: `AHardwareBufferTensor` (Android) has a
5018 // `view_offset` that its own `map()` adds and its own `view()`
5019 // sets, but nothing writes it here yet, so restoring an offset
5020 // onto an already-built tensor is silently a no-op on this
5021 // backing -- a caller gets a map at the parent's origin. Not
5022 // reachable through `TensorDyn::import_descriptor` today
5023 // (`import_storage`'s DMABUF arm is `cfg(target_os = "linux")`),
5024 // so it is latent rather than silent. See
5025 // `edgefirst-python-common`'s `interop::apply_plane_offset` for
5026 // the full per-backing accounting.
5027 #[cfg(target_os = "android")]
5028 TensorStorage::Dma(_) => {}
5029 }
5030 }
5031
5032 /// The CPU access declared for this tensor (see [`CpuAccess`]): set at
5033 /// allocation, possibly narrowed by
5034 /// [`restrict_cpu_access`](Self::restrict_cpu_access). A view copies its
5035 /// parent's declaration when it is taken.
5036 pub fn cpu_access(&self) -> CpuAccess {
5037 self.cpu_access
5038 }
5039
5040 /// Narrow the declared CPU access to the directions both it and `access`
5041 /// cover, on the chroma plane too. Never widens: a declaration must not
5042 /// exceed what the allocation was made for, so a wider `access` leaves
5043 /// it unchanged. An importer uses this to carry a producer's declaration
5044 /// onto an alias of the same buffer.
5045 ///
5046 /// ```
5047 /// # use edgefirst_tensor::{CpuAccess, PixelFormat, Tensor, TensorMemory};
5048 /// # fn main() -> Result<(), edgefirst_tensor::Error> {
5049 /// let mut t = Tensor::<u8>::image(4, 4, PixelFormat::Rgba, Some(TensorMemory::Mem), CpuAccess::ReadWrite)?;
5050 /// t.restrict_cpu_access(CpuAccess::Read);
5051 /// assert_eq!(t.cpu_access(), CpuAccess::Read);
5052 /// t.restrict_cpu_access(CpuAccess::ReadWrite);
5053 /// assert_eq!(t.cpu_access(), CpuAccess::Read);
5054 /// # Ok(())
5055 /// # }
5056 /// ```
5057 pub fn restrict_cpu_access(&mut self, access: CpuAccess) {
5058 self.cpu_access = self.cpu_access.intersection(access);
5059 if let Some(chroma) = self.chroma.as_mut() {
5060 chroma.restrict_cpu_access(access);
5061 }
5062 }
5063
5064 /// Set the declared CPU access without re-allocating — crate-private:
5065 /// the declaration must reflect the underlying allocation's real
5066 /// capabilities (constructors and importers set it; arbitrary widening
5067 /// would defeat the contract).
5068 // Only the Android AHardwareBuffer importer derives a declaration from usage bits.
5069 #[cfg(target_os = "android")]
5070 pub(crate) fn set_cpu_access_unchecked(&mut self, access: CpuAccess) {
5071 self.cpu_access = access;
5072 }
5073
5074 /// The vendor tile-compression scheme recorded at allocation, or
5075 /// `None` for a linear layout. `Some` means the pixels live in a
5076 /// proprietary tile order: the row-stride accessors describe no
5077 /// meaningful linear layout and CPU maps are best-effort (see
5078 /// [`Compression`]). Only Android hardware-only allocations that
5079 /// requested compression record a scheme.
5080 pub fn compression(&self) -> Option<CompressionScheme> {
5081 self.compression
5082 }
5083
5084 /// Whether this tensor's memory is known to be physically contiguous:
5085 /// [`Contiguity::Contiguous`] for a DMA-BUF this crate allocated from
5086 /// the CMA heap, [`Contiguity::NonContiguous`] for one from the system
5087 /// heap, and [`Contiguity::Unknown`] for an imported fd and every other
5088 /// kind of memory.
5089 pub fn contiguity(&self) -> Contiguity {
5090 match &self.storage {
5091 #[cfg(target_os = "linux")]
5092 TensorStorage::Dma(d) => d.contiguity,
5093 _ => Contiguity::Unknown,
5094 }
5095 }
5096
5097 /// Record the compression scheme — crate-private: recording is an
5098 /// allocation-time fact ([`Tensor::image_desc`] sets it; arbitrary
5099 /// mutation would misdescribe the physical layout).
5100 // Only the Android allocation path records a scheme.
5101 #[cfg(target_os = "android")]
5102 pub(crate) fn set_compression_unchecked(&mut self, scheme: Option<CompressionScheme>) {
5103 self.compression = scheme;
5104 }
5105
5106 pub fn colorimetry(&self) -> Option<crate::Colorimetry> {
5107 self.colorimetry
5108 }
5109
5110 /// Attach/clear colorimetry metadata.
5111 pub fn set_colorimetry(&mut self, c: Option<crate::Colorimetry>) {
5112 self.colorimetry = c;
5113 }
5114
5115 /// Builder-style colorimetry attach.
5116 pub fn with_colorimetry(mut self, c: crate::Colorimetry) -> Self {
5117 self.colorimetry = Some(c);
5118 self
5119 }
5120
5121 /// Create a zero-copy sub-region view of this tensor's backing buffer.
5122 ///
5123 /// The returned tensor shares this tensor's allocation (no copy) and maps
5124 /// the window `[offset_bytes, offset_bytes + shape.product()*size_of::<T>())`
5125 /// measured from this tensor's own logical start. N sub-views into one
5126 /// parent can be written independently, enabling batched assembly into a
5127 /// single buffer. Identical semantics across `Mem` (shared `Arc`) and
5128 /// `Dma` (shared fd) backings.
5129 ///
5130 /// # Disjointness
5131 ///
5132 /// Independent writes are sound *only* when the windows do not overlap. The
5133 /// shared backing uses interior mutability (`UnsafeCell` cells), so two
5134 /// sub-views whose byte ranges intersect alias the same cells: writing one
5135 /// while reading or writing the other is a data race and therefore
5136 /// **undefined behaviour**. The caller is responsible for keeping the
5137 /// windows disjoint; this method does not check for overlap.
5138 ///
5139 /// # Errors
5140 ///
5141 /// - [`Error::InvalidOperation`] if the backing is not `Mem` or `Dma`, or
5142 /// if `offset_bytes` is not a multiple of `align_of::<T>()`.
5143 /// - [`Error::InsufficientCapacity`] / [`Error::InvalidSize`] if the window
5144 /// exceeds the parent allocation.
5145 pub(crate) fn subview(&self, offset_bytes: usize, shape: &[usize]) -> Result<Tensor<T>> {
5146 // Offset is absolute into the backing allocation: a sub-view of a
5147 // sub-view composes by adding this tensor's own offset.
5148 let abs_offset = self
5149 .plane_offset
5150 .unwrap_or(0)
5151 .checked_add(offset_bytes)
5152 .ok_or(Error::InvalidSize(offset_bytes))?;
5153 // Every backend exposes `view(offset, shape)` via `TensorTrait`, sharing
5154 // the resource AND `BufferIdentity` (unlike `from_fd`/`from_surface`,
5155 // which mint a fresh identity). The GL backend keys the import on the
5156 // shared identity so offset-distinct sub-views of one buffer reuse a
5157 // single import and address their window via `glViewport`. `Mem`/`Shm`
5158 // share via the allocation `Arc` / a cloned fd; `Pbo` via the GL-buffer
5159 // `Arc`; Linux DMA-BUF / macOS IOSurface via the shared fd / CFRetain.
5160 // `TensorStorage::view` performs the one remaining per-variant dispatch.
5161 let mut t = Tensor::wrap(self.storage.view(offset_bytes, shape)?);
5162 // Inherit the parent's image metadata so the view is a ready-to-use
5163 // sub-image (e.g. a `convert()` destination). The offset is applied
5164 // LAST because `set_format` deliberately clears it — the offset is a
5165 // structural property of the sub-region, not format-dependent metadata.
5166 if let Some(fmt) = self.format {
5167 t.set_format(fmt)?;
5168 }
5169 if let Some(rs) = self.row_stride {
5170 t.set_row_stride_unchecked(rs);
5171 }
5172 t.quantization = self.quantization.clone();
5173 // A sub-region of an image carries the parent's colorimetry — it is the
5174 // same pixels, same color encoding. Inherit it like the other image
5175 // metadata above so a sub-view is a faithful convert() source/target.
5176 t.set_colorimetry(self.colorimetry);
5177 // The declared CPU access is a property of the underlying
5178 // allocation, so every view shares the parent's declaration.
5179 t.cpu_access = self.cpu_access;
5180 // Likewise the recorded compression scheme: the view shares the
5181 // parent's physical layout.
5182 t.compression = self.compression;
5183 t.view_marked = self.is_view();
5184 if abs_offset > 0 {
5185 t.set_plane_offset(abs_offset);
5186 }
5187 Ok(t)
5188 }
5189
5190 /// Borrow batch element `n` of a batched tensor as a zero-copy view.
5191 ///
5192 /// A batched tensor prepends `N` as the leading dimension over the
5193 /// per-element image layout (`[N, H, W, C]` packed or `[N, C, H, W]`
5194 /// planar) — `N` is over the whole per-element block regardless of
5195 /// `HWC`/`CHW`. `batch(n)` returns element `n`: the contiguous per-element
5196 /// region at byte offset `n * element_size`, sharing the parent's
5197 /// `BufferIdentity` and inheriting its format / row stride / colorimetry.
5198 /// `batch(0)` on a tensor with `N == 1` is equivalent to the whole tensor.
5199 ///
5200 /// # Errors
5201 ///
5202 /// - [`Error::BatchIndexOutOfBounds`] if `n >= N`.
5203 /// - [`Error::InvalidShape`] if the tensor is not batched (a formatted
5204 /// tensor whose rank lacks the leading `N`, or an empty shape).
5205 pub fn batch(&self, n: usize) -> Result<Tensor<T>> {
5206 let shape = self.shape();
5207 // With a format we know the exact per-element rank, so a missing leading
5208 // `N` is a misuse we reject rather than silently treating a spatial dim
5209 // as the batch. Raw tensors take shape[0] as `N` by contract.
5210 if let Some(fmt) = self.format {
5211 let elem_rank = image_rank(fmt.layout());
5212 if shape.len() != elem_rank + 1 {
5213 return Err(Error::InvalidShape(format!(
5214 "batch(): tensor is not batched ({fmt:?} expects a leading N over a \
5215 {elem_rank}-D element, got shape {shape:?})"
5216 )));
5217 }
5218 }
5219 let batch = *shape
5220 .first()
5221 .ok_or_else(|| Error::InvalidShape("batch(): empty shape".into()))?;
5222 if n >= batch {
5223 return Err(Error::BatchIndexOutOfBounds { index: n, batch });
5224 }
5225 let elem_shape: Vec<usize> = shape[1..].to_vec();
5226 let elem_count: usize = elem_shape.iter().product();
5227 elem_count
5228 .checked_mul(std::mem::size_of::<T>())
5229 .ok_or(Error::InvalidSize(elem_count))?;
5230 // One element's footprint including row padding: the descriptor's
5231 // stride for `N`, so a padded batch is addressed the way it is mapped.
5232 let dims: Vec<u64> = shape.iter().map(|&d| d as u64).collect();
5233 let elem_bytes = crate::protocol::c_byte_strides(
5234 &dims,
5235 std::mem::size_of::<T>() as i64,
5236 self.format.map(|f| f.layout()),
5237 self.row_stride,
5238 )[0] as usize;
5239 let offset = n.checked_mul(elem_bytes).ok_or(Error::InvalidSize(n))?;
5240 // For a packed `[N, H, W, C]` tensor the N tiles stack vertically in the
5241 // shared buffer, so the GL import sees one `(W, N*H)` parent and each
5242 // tile is the row-band at `y = n*H`. Snapshot that parent so the backend
5243 // imports once and renders the tile via `glViewport`. Non-packed
5244 // (planar/semi-planar) batching keeps the per-slot path for now (planar
5245 // NCHW tiling is a separate step); raw tensors have no pixel geometry.
5246 let view_origin = match self.format.map(|f| f.layout()) {
5247 Some(PixelLayout::Packed) => {
5248 let tile_h = elem_shape[0];
5249 let tile_w = elem_shape[1];
5250 // Per-row pitch of the tall `(W, N*H)` parent — padded stride if
5251 // set, else the tight row width. The GL import keys on this.
5252 let parent_stride = self
5253 .effective_row_stride()
5254 .expect("a formatted tensor of its batched rank has a width");
5255 Some(self.compose_view_origin(tile_w, batch * tile_h, parent_stride, 0, n * tile_h))
5256 }
5257 _ => None,
5258 };
5259 let mut t = self.subview(offset, &elem_shape)?;
5260 t.view_origin = view_origin;
5261 Ok(t)
5262 }
5263
5264 /// Borrow a rectangular spatial sub-region of an image tensor as a
5265 /// zero-copy view — the **destination/source crop** primitive.
5266 ///
5267 /// `region` is in pixels of the image's leading frame. The returned view
5268 /// shares the parent's `BufferIdentity` and addresses the sub-rectangle by
5269 /// offset + the **parent's** row pitch (so each row lands at the correct
5270 /// columns). `convert(src, &mut dst.view(rect), …)` renders into that
5271 /// sub-rectangle of `dst`; a letterbox fit then clears the view and renders
5272 /// the aspect-preserved content into its inner region. `view`/`batch`/the
5273 /// whole tensor are the one coherent destination model — there is no
5274 /// separate `dst_rect`.
5275 ///
5276 /// # Known limitation: bottom-right multi-row views
5277 ///
5278 /// A multi-row view is mapped as `parent_pitch × rows` bytes **from its own
5279 /// offset**, so the window it asks for overruns the allocation by
5280 /// `region.x * bpp` whenever the view is both offset horizontally *and*
5281 /// reaches the parent's last row. Such a view constructs fine but fails at
5282 /// [`map`](TensorTrait::map) time with [`Error::InsufficientCapacity`] —
5283 /// loudly, never as an out-of-bounds access. For a tiled destination this is
5284 /// the bottom row of tiles at every column but the first. The single-row
5285 /// case is already special-cased below (it keeps its own tight stride); the
5286 /// general fix is to clamp the last row's exposure to the view's own row
5287 /// bytes, which is not yet implemented.
5288 ///
5289 /// # Errors
5290 ///
5291 /// - [`Error::RegionOutOfBounds`] if `region` exceeds the image bounds.
5292 /// - [`Error::InvalidOperation`] if the tensor is not a packed-format image
5293 /// (planar/semi-planar spatial sub-rects are not a single strided window;
5294 /// use [`batch`](Self::batch) for batched planar tensors).
5295 pub fn view(&self, region: Region) -> Result<Tensor<T>> {
5296 let fmt = self.format.ok_or_else(|| {
5297 Error::InvalidOperation("view() requires a formatted image tensor".into())
5298 })?;
5299 if fmt.layout() != PixelLayout::Packed {
5300 return Err(Error::InvalidOperation(format!(
5301 "view() supports packed formats only (got {fmt:?}); use batch(n) for batched \
5302 planar tensors"
5303 )));
5304 }
5305 let w = self
5306 .width()
5307 .ok_or_else(|| Error::InvalidOperation("view(): tensor has no image width".into()))?;
5308 let h = self
5309 .height()
5310 .ok_or_else(|| Error::InvalidOperation("view(): tensor has no image height".into()))?;
5311 if !region.fits_within(w, h) {
5312 return Err(Error::RegionOutOfBounds {
5313 region,
5314 bounds: (w, h),
5315 });
5316 }
5317 let elem = std::mem::size_of::<T>();
5318 let bpp = fmt.channels() * elem;
5319 let stride = self
5320 .effective_row_stride()
5321 .expect("a tensor with an image width has a row stride");
5322 let offset = region
5323 .y
5324 .checked_mul(stride)
5325 .and_then(|yo| yo.checked_add(region.x.checked_mul(bpp)?))
5326 .ok_or(Error::InvalidSize(region.y))?;
5327 let sub_shape = fmt
5328 .allocation_shape(region.width, region.height)
5329 .ok_or_else(|| Error::InvalidShape(format!("view(): invalid shape for {fmt:?}")))?;
5330 let mut t = self.subview(offset, &sub_shape)?;
5331 // A multi-row sub-rect must advance rows by the PARENT pitch so each row
5332 // addresses the correct columns. A single-row view uses its own tight
5333 // stride — the parent pitch would make the strided `map()` (which exposes
5334 // `stride × rows`) expose a trailing row that runs past the buffer tail
5335 // for an offset (x>0 / bottom) view. The GL backend does NOT rely on this
5336 // (single-row-tight) `row_stride`: it reads the parent pitch from
5337 // `view_origin.parent_row_stride` so its import/cache pitch stays
5338 // parent-consistent for views of any height — see `ViewOrigin`.
5339 let view_stride = if region.height > 1 {
5340 stride
5341 } else {
5342 region.width * bpp
5343 };
5344 t.set_row_stride_unchecked(view_stride);
5345 // Snapshot the parent `(w, h, row_stride)` so the GL backend imports the
5346 // parent once (keyed on the parent pitch, not this view's possibly-tight
5347 // single-row stride) and renders this sub-rect as a `glViewport`/
5348 // `glScissor` ROI at `(region.x, region.y)`. Composes when viewing an
5349 // existing view.
5350 t.view_origin = Some(self.compose_view_origin(w, h, stride, region.x, region.y));
5351 Ok(t)
5352 }
5353
5354 /// Build the [`ViewOrigin`] for a new sub-region of `self`. When `self` is a
5355 /// whole tensor the snapshot names `self` as the parent; when `self` is
5356 /// already a view, the snapshot keeps the **root** parent and accumulates the
5357 /// local origin so nested views still resolve to one import.
5358 fn compose_view_origin(
5359 &self,
5360 parent_width: usize,
5361 parent_height: usize,
5362 parent_row_stride: usize,
5363 x: usize,
5364 y: usize,
5365 ) -> ViewOrigin {
5366 match self.view_origin {
5367 Some(root) => ViewOrigin {
5368 parent_width: root.parent_width,
5369 parent_height: root.parent_height,
5370 parent_row_stride: root.parent_row_stride,
5371 x: root.x.saturating_add(x),
5372 y: root.y.saturating_add(y),
5373 },
5374 None => ViewOrigin {
5375 parent_width,
5376 parent_height,
5377 parent_row_stride,
5378 x,
5379 y,
5380 },
5381 }
5382 }
5383
5384 /// The wrapped `PboTensor`, for this module's own PBO accessors.
5385 ///
5386 /// Crate-private: `PboTensor<T>` is a `static`-side value with no
5387 /// `dynamic` counterpart to lend (there the buffer lives inside
5388 /// `libedgefirst_tensor.so`), so a public borrow of it could not have
5389 /// the same signature on both backends. Callers use
5390 /// [`Self::pbo_id`]/[`Self::pbo_is_mapped`], which are the whole surface
5391 /// `edgefirst-image` ever used.
5392 pub(crate) fn as_pbo(&self) -> Option<&PboTensor<T>> {
5393 match &self.storage {
5394 TensorStorage::Pbo(p) => Some(p),
5395 _ => None,
5396 }
5397 }
5398
5399 /// The GL buffer ID backing this tensor (for GL backends / the
5400 /// cross-package protocol's `handle`). `None` when the tensor is not
5401 /// PBO-backed.
5402 pub fn pbo_id(&self) -> Option<u32> {
5403 self.as_pbo().map(PboTensor::buffer_id)
5404 }
5405
5406 /// Whether this PBO currently holds -- or is establishing -- a CPU
5407 /// mapping. `None` when the tensor is not PBO-backed.
5408 ///
5409 /// The question a GL caller is really asking is "is this buffer free
5410 /// for GL operations?", so a map or unmap still in flight answers
5411 /// `Some(true)`: the GL side must not touch the buffer until the
5412 /// transition finishes. See [`PboTensor::is_mapped`].
5413 pub fn pbo_is_mapped(&self) -> Option<bool> {
5414 self.as_pbo().map(PboTensor::is_mapped)
5415 }
5416
5417 /// The C-ABI `PboOpsVtable` address backing this tensor (for the
5418 /// cross-package protocol's `ptr` under `kind::PBO`; see
5419 /// [`crate::protocol::TensorDesc::ptr`]'s doc comment). `None` when the
5420 /// tensor is not PBO-backed.
5421 pub fn pbo_vtable_ptr(&self) -> Option<*const std::ffi::c_void> {
5422 self.as_pbo()
5423 .map(|p| p.pbo_vtable() as *const _ as *const std::ffi::c_void)
5424 }
5425
5426 /// A type-erased keepalive that must stay alive for at least as long as
5427 /// [`Self::pbo_vtable_ptr`]'s address is used -- see
5428 /// [`PboTensor::pbo_keepalive`]'s own doc comment. `None` when the
5429 /// tensor is not PBO-backed.
5430 pub fn pbo_keepalive(&self) -> Option<Arc<dyn Send + Sync>> {
5431 self.as_pbo().map(PboTensor::pbo_keepalive)
5432 }
5433
5434 /// Downcast to DMA tensor reference (for EGL import, G2D).
5435 #[cfg(target_os = "linux")]
5436 pub fn as_dma(&self) -> Option<&DmaTensor<T>> {
5437 match &self.storage {
5438 TensorStorage::Dma(d) => Some(d),
5439 _ => None,
5440 }
5441 }
5442
5443 /// Borrow the DMA-BUF file descriptor backing this tensor.
5444 ///
5445 /// # Returns
5446 ///
5447 /// A borrowed reference to the DMA-BUF file descriptor, tied to `self`'s
5448 /// lifetime.
5449 ///
5450 /// # Errors
5451 ///
5452 /// Returns `Error::NotImplemented` if the tensor is not DMA-backed.
5453 #[cfg(target_os = "linux")]
5454 pub fn dmabuf(&self) -> Result<std::os::fd::BorrowedFd<'_>> {
5455 use std::os::fd::AsFd;
5456 match &self.storage {
5457 TensorStorage::Dma(dma) => Ok(dma.fd.as_fd()),
5458 _ => Err(Error::NotImplemented(format!(
5459 "dmabuf requires DMA-backed tensor, got {:?}",
5460 self.storage.memory()
5461 ))),
5462 }
5463 }
5464
5465 /// Construct a Tensor from a PBO tensor (for GL backends that allocate PBOs).
5466 ///
5467 /// Infallible in this backend (`Ok` always) -- kept as `Result` to match
5468 /// `dynamic`'s own `Tensor::from_pbo` (`dynamic_tensor.rs`), which
5469 /// genuinely can fail: it must allocate a real `ef_tensor_*` handle
5470 /// alongside `pbo` to carry shape/dtype/format metadata (this backend
5471 /// has nowhere else to put it), an allocation this backend never
5472 /// needed since it just wraps `pbo` into `TensorStorage::Pbo` directly.
5473 /// One shared, fallible signature lets `edgefirst-image`'s own call
5474 /// sites (`gl/threaded.rs`) stay backend-agnostic.
5475 pub fn from_pbo(pbo: PboTensor<T>) -> Result<Self> {
5476 Ok(Self {
5477 storage: TensorStorage::Pbo(pbo),
5478 format: None,
5479 chroma: None,
5480 row_stride: None,
5481 plane_offset: None,
5482 quantization: None,
5483 cuda: None,
5484 colorimetry: None,
5485 cpu_access: CpuAccess::ReadWrite,
5486 compression: None,
5487 view_origin: None,
5488 view_marked: false,
5489 })
5490 }
5491
5492 /// The CUDA registration for this tensor, if any (set at creation on CUDA devices).
5493 pub fn cuda(&self) -> Option<&crate::cuda::CudaHandle> {
5494 self.cuda.as_ref()
5495 }
5496
5497 /// Whether a CUDA registration is attached to this tensor.
5498 pub fn is_cuda_attached(&self) -> bool {
5499 self.cuda.is_some()
5500 }
5501
5502 /// Attach a CUDA handle (called by ImageProcessor::create_image after registering a PBO).
5503 pub fn set_cuda_handle(&mut self, h: crate::cuda::CudaHandle) {
5504 self.cuda = Some(h);
5505 }
5506
5507 /// Fast-fail CUDA map: None (no GL routing) when no handle; else map (PBO routes to the GL worker).
5508 ///
5509 /// Returns a scoped [`CudaMap`] guard holding the raw CUDA device pointer
5510 /// for the duration of the mapping. For GL-buffer-backed tensors the unmap is deferred until the
5511 /// guard drops, freeing the PBO for the next `convert()` call. When no CUDA handle is attached
5512 /// (the common case for plain `Mem`/`DMA` tensors without CUDA registration), returns `None`
5513 /// immediately — no GL routing, no allocation.
5514 ///
5515 /// # Example — zero-copy CUDA input with host fallback
5516 ///
5517 /// ```no_run
5518 /// use edgefirst_tensor::{Tensor, TensorMemory, TensorTrait};
5519 /// # fn feed_tensorrt(_dptr: *mut std::ffi::c_void, _bytes: usize) {}
5520 /// # fn demo(t: &Tensor<f32>) {
5521 /// // Try the zero-copy CUDA device pointer first.
5522 /// if let Some(cuda) = t.cuda_map() {
5523 /// feed_tensorrt(cuda.device_ptr(), cuda.len());
5524 /// // `cuda` (a CudaMap guard) unmaps when it goes out of scope, freeing
5525 /// // the GPU buffer for the next convert().
5526 /// } else {
5527 /// // Fall back to the host mapping when no CUDA handle is attached.
5528 /// let _host = t.map().expect("host map fallback must succeed");
5529 /// // `_host` is a TensorMap<f32> that derefs to &[f32].
5530 /// }
5531 /// # }
5532 /// ```
5533 ///
5534 /// # D3D11 texture tensors: the mapping is tight rows
5535 ///
5536 /// On Windows the mapping is a linear copy of the texture laid out at
5537 /// `width * bytes_per_texel`, and `len()` is that times the row count --
5538 /// *not* [`row_stride`](Self::row_stride), which reports the D3D11 staging
5539 /// pitch a CPU map sees and is the larger number on a padded backing. A
5540 /// consumer that strides the device pointer by `row_stride()` reads past
5541 /// the end of the mapping.
5542 ///
5543 /// Semi-planar formats (NV12, NV16, NV24) are the exception: their texture
5544 /// is allocated as wide as its own row pitch, so the tight row *is* the
5545 /// pitch and the mapping equals `row_stride() x combined_height` bytes.
5546 /// Striding the device pointer by `row_stride()` is exactly right there.
5547 pub fn cuda_map(&self) -> Option<crate::cuda::CudaMap<'_>> {
5548 self.cuda.as_ref()?.map()
5549 }
5550
5551 /// Writable counterpart of [`cuda_map`](Self::cuda_map). See
5552 /// [`CudaHandle::map_mut`](crate::cuda::CudaHandle::map_mut) for which
5553 /// backings distinguish the two.
5554 ///
5555 /// The D3D11 arm is the one that does: the guard's release copies the
5556 /// buffer back into the texture. Its layout is [`cuda_map`](Self::cuda_map)'s
5557 /// -- tight rows of `width * bytes_per_texel`, not `row_stride()` -- and
5558 /// writing it at the wrong pitch scrambles the image rather than failing.
5559 pub fn cuda_map_mut(&self) -> Option<crate::cuda::CudaMap<'_>> {
5560 self.cuda.as_ref()?.map_mut()
5561 }
5562
5563 /// Attempt to attach a CUDA `ExternalMemory` handle for DMA-backed tensors.
5564 ///
5565 /// On a CUDA-capable host, imports the DMA-BUF fd via
5566 /// `cudaImportExternalMemory(OpaqueFd)` and maps it to a device pointer.
5567 /// Sets `self.cuda` to a persistent `ExternalMem` handle on success. No-op
5568 /// if CUDA is unavailable, the tensor is not DMA-backed, or a handle is
5569 /// already set. Import failure is silently ignored — the tensor remains
5570 /// usable without a CUDA handle.
5571 ///
5572 /// # RUNTIME-UNVALIDATED
5573 ///
5574 /// No test platform has both `/dev/dma_heap` and a CUDA device. ABI is
5575 /// layout-asserted vs. CUDA 12.6 `driver_types.h`; the mechanism is proven
5576 /// by gpu-probe O5 on Orin. Best-effort: tensor creation never fails here.
5577 #[cfg(target_os = "linux")]
5578 pub fn try_init_dma_cuda(&mut self) {
5579 // Fast-path: already imported, CUDA not available, or not a DMA tensor.
5580 if self.cuda.is_some() || !crate::cuda::is_cuda_available() {
5581 return;
5582 }
5583 let (raw_fd, buf_size) = match &self.storage {
5584 TensorStorage::Dma(dma) => {
5585 use std::os::fd::AsRawFd;
5586 (dma.fd.as_raw_fd(), dma.buf_size)
5587 }
5588 _ => return,
5589 };
5590 if let Some((ext, dptr)) = crate::cuda::import_dma_fd(raw_fd, buf_size) {
5591 self.cuda = Some(crate::cuda::CudaHandle::new_external(ext, dptr, buf_size));
5592 }
5593 }
5594
5595 /// Attempt to attach a CUDA external-memory handle for D3D11 texture
5596 /// tensors — the Windows counterpart of
5597 /// [`try_init_dma_cuda`](Self::try_init_dma_cuda).
5598 ///
5599 /// Imports the texture's shared NT handle via
5600 /// `cudaImportExternalMemory(D3D11Resource)` and the process device's
5601 /// fence as an external semaphore, so [`cuda_map`](Self::cuda_map) hands
5602 /// out a linear device buffer whose refresh is ordered after the D3D11
5603 /// work a producer recorded with [`set_gpu_write`](Self::set_gpu_write).
5604 ///
5605 /// No-op when a handle is already attached or the tensor is not
5606 /// texture-backed. Import failure is silently ignored: the WARP adapter
5607 /// has no CUDA device at all, and a texture whose import fails is still a
5608 /// valid tensor whose `cuda_map()` is `None`.
5609 #[cfg(target_os = "windows")]
5610 pub fn try_init_d3d11_cuda(&mut self) {
5611 if self.cuda.is_some() {
5612 return;
5613 }
5614 let TensorStorage::Dma(tex) = &self.storage else {
5615 return;
5616 };
5617 if let Some(h) = tex.try_init_cuda() {
5618 self.cuda = Some(h);
5619 }
5620 }
5621}
5622
5623/// Windows D3D11 texture accessors. No `'static` bound: none of these needs
5624/// `dtype_of::<T>()`.
5625#[cfg(all(not(edgefirst_tensor_dynamic), target_os = "windows"))]
5626impl<T> Tensor<T>
5627where
5628 T: Num + Clone + fmt::Debug + Send + Sync,
5629{
5630 /// The borrowed `ID3D11Texture2D*` backing this tensor, valid while the
5631 /// tensor lives. `None` when the tensor is not texture-backed.
5632 pub fn d3d11_texture(&self) -> Option<*mut std::ffi::c_void> {
5633 match &self.storage {
5634 TensorStorage::Dma(t) => Some(t.texture_ptr()),
5635 _ => None,
5636 }
5637 }
5638
5639 /// The texture geometry the HAL chose for this image: DXGI format, texel
5640 /// dimensions and the matching GL internal format. `None` when the tensor
5641 /// is not texture-backed.
5642 pub fn d3d11_layout(&self) -> Option<crate::d3d11_layout::D3d11ImageLayout> {
5643 match &self.storage {
5644 TensorStorage::Dma(t) => Some(t.layout()),
5645 _ => None,
5646 }
5647 }
5648
5649 /// A duplicated NT handle the caller owns: open it on a D3D12 device,
5650 /// another D3D11 device, or CUDA, or duplicate it into another process.
5651 /// Closing it does not affect this tensor.
5652 pub fn d3d11_shared_handle(&self) -> Result<std::os::windows::io::OwnedHandle> {
5653 match &self.storage {
5654 TensorStorage::Dma(t) => t.shared_handle(),
5655 _ => Err(Error::NotImplemented(
5656 "not a D3D11 texture tensor".to_owned(),
5657 )),
5658 }
5659 }
5660
5661 /// The fence handle plus value a GPU consumer waits on before reading this
5662 /// texture, or `None` when no GPU write has been recorded on it.
5663 pub fn gpu_completion(&self) -> Result<Option<crate::d3d11::GpuCompletion>> {
5664 match &self.storage {
5665 TensorStorage::Dma(t) => t.gpu_completion(),
5666 _ => Err(Error::NotImplemented(
5667 "not a D3D11 texture tensor".to_owned(),
5668 )),
5669 }
5670 }
5671
5672 /// Record that GPU work writing this texture completes at `value` of the
5673 /// process device's fence. Monotonic: an older value never displaces a
5674 /// newer one, so recording out of order is safe rather than merely
5675 /// tolerated.
5676 ///
5677 /// Takes `&self`, not `&mut self`: the fence slot is an `AtomicU64` and
5678 /// this is a `fetch_max` on it. A producer records a write on the very
5679 /// tensor its consumers already hold and are reading, so demanding
5680 /// exclusive access here would forbid the one call this exists for.
5681 pub fn set_gpu_write(&self, value: u64) -> Result<()> {
5682 match &self.storage {
5683 TensorStorage::Dma(t) => {
5684 t.set_gpu_write(value);
5685 Ok(())
5686 }
5687 _ => Err(Error::NotImplemented(
5688 "not a D3D11 texture tensor".to_owned(),
5689 )),
5690 }
5691 }
5692
5693 /// The tensor's own NT handle value -- not a duplicate -- for descriptors
5694 /// whose consumer keeps the producing tensor alive.
5695 ///
5696 /// Borrowed rather than duplicated because the C exports (task A6) hand
5697 /// out owned duplicates, which a `Copy` descriptor cannot close; the
5698 /// cross-package protocol's `kind::D3D11_TEXTURE` reads this instead.
5699 pub(crate) fn d3d11_shared_handle_value(&self) -> Option<usize> {
5700 match &self.storage {
5701 TensorStorage::Dma(t) => Some(t.shared_handle_value()),
5702 _ => None,
5703 }
5704 }
5705}
5706
5707/// The recorded GPU write value, on every platform: the descriptor path and
5708/// the Python bindings forward it without a `cfg` of their own, and only
5709/// Windows can record one.
5710#[cfg(not(edgefirst_tensor_dynamic))]
5711impl<T> Tensor<T>
5712where
5713 T: Num + Clone + fmt::Debug + Send + Sync,
5714{
5715 /// The fence value of the newest GPU write recorded on this tensor, or 0
5716 /// when there is none: the `value` of `gpu_completion` without the
5717 /// duplicated fence handle, for a consumer that already holds the process
5718 /// fence. 0 for a tensor that is not a D3D11 texture, and on every
5719 /// platform other than Windows.
5720 pub fn gpu_write_value(&self) -> u64 {
5721 #[cfg(target_os = "windows")]
5722 {
5723 if let TensorStorage::Dma(t) = &self.storage {
5724 return t.last_gpu_write();
5725 }
5726 }
5727 0
5728 }
5729}
5730
5731/// Records the D3D11 backing's row pitch on `t` when it exceeds the natural
5732/// stride for its format and width.
5733///
5734/// The staging pitch is the row stride CPU consumers see, and for a
5735/// semi-planar tensor it is also the texture's own width, which the GL
5736/// backend's Path B shader takes as its `tex_width`. Three constructors need
5737/// it identically -- `Tensor::image`'s Windows arm and the two wrap
5738/// constructors below -- and a wrapped texture is the case that makes it
5739/// load-bearing: without it such a tensor reports the natural `even(width)`
5740/// stride while its texture is the driver's pitch wide, and the shader samples
5741/// a padded texture at the image width.
5742///
5743/// `t.format` must already be set: the natural stride comes from it.
5744#[cfg(all(not(edgefirst_tensor_dynamic), target_os = "windows"))]
5745fn record_d3d11_row_stride<T>(t: &mut Tensor<T>, format: PixelFormat, what: &str)
5746where
5747 T: Num + Clone + fmt::Debug + Send + Sync,
5748{
5749 let pitch = match &t.storage {
5750 TensorStorage::Dma(tex) => tex.image_backing_row_stride(),
5751 _ => None,
5752 };
5753 let (Some(pitch), Some(natural)) = (pitch, t.effective_row_stride()) else {
5754 return;
5755 };
5756 if pitch > natural {
5757 log::debug!(
5758 "{what}: the D3D11 backing pitch for the {format:?} texture is {pitch} bytes \
5759 (natural {natural}); recording row stride"
5760 );
5761 t.set_row_stride_unchecked(pitch);
5762 }
5763}
5764
5765/// The two texture-importing constructors. Separate from the accessors above
5766/// because `dtype_of::<T>()` needs `T: 'static`.
5767#[cfg(all(not(edgefirst_tensor_dynamic), target_os = "windows"))]
5768impl<T> Tensor<T>
5769where
5770 T: Num + Clone + fmt::Debug + Send + Sync + 'static,
5771{
5772 /// Wrap an existing `ID3D11Texture2D` as a tensor. The texture is checked
5773 /// against the layout the HAL would have chosen for `format`/`width`/
5774 /// `height` and must live on the HAL's process device; ownership stays
5775 /// with the caller, since this takes its own reference.
5776 ///
5777 /// The backing's row pitch is recorded as the tensor's
5778 /// [`row_stride`](Self::row_stride) when it exceeds the natural stride, so
5779 /// a wrapped semi-planar texture -- whose width *is* the pitch -- reports
5780 /// the geometry its consumers address it by.
5781 ///
5782 /// # Safety
5783 ///
5784 /// `texture` must be null or a live `ID3D11Texture2D` created on the HAL
5785 /// device ([`crate::d3d11::device()`]).
5786 pub unsafe fn from_d3d11_texture(
5787 texture: *mut std::ffi::c_void,
5788 width: usize,
5789 height: usize,
5790 format: PixelFormat,
5791 access: CpuAccess,
5792 name: Option<&str>,
5793 ) -> Result<Self> {
5794 let dtype = dtype_of::<T>().ok_or_else(|| {
5795 Error::InvalidArgument(format!(
5796 "from_d3d11_texture: element type {} has no DType",
5797 std::any::type_name::<T>()
5798 ))
5799 })?;
5800 let shape = format.allocation_shape(width, height).ok_or_else(|| {
5801 Error::InvalidArgument(format!(
5802 "invalid dimensions {width}x{height} for format {format:?}"
5803 ))
5804 })?;
5805 // SAFETY: the caller guarantees `texture` is null or a live
5806 // `ID3D11Texture2D`; the storage takes its own reference.
5807 let storage = unsafe {
5808 crate::d3d11::texture::D3d11TextureTensor::<T>::from_d3d11_texture(
5809 texture, width, height, format, dtype, &shape, access, name,
5810 )
5811 }?;
5812 let mut t = Self::wrap(TensorStorage::Dma(storage));
5813 t.format = Some(format);
5814 record_d3d11_row_stride(&mut t, format, "from_d3d11_texture");
5815 t.cpu_access = access;
5816 t.try_init_d3d11_cuda();
5817 Ok(t)
5818 }
5819
5820 /// Open a shared texture by its NT handle on the HAL's device. When
5821 /// `completion` is given, its fence value is waited for on the immediate
5822 /// context, so a same-device reader needs no further ordering. Both
5823 /// handles stay owned by the caller: this duplicates what it keeps.
5824 ///
5825 /// Records the backing's row pitch as
5826 /// [`row_stride`](Self::row_stride) on the same rule as
5827 /// [`from_d3d11_texture`](Self::from_d3d11_texture).
5828 ///
5829 /// # Safety
5830 ///
5831 /// `handle` must be an NT shared handle of a D3D11 texture, valid in this
5832 /// process, and `completion`'s handle a shared fence handle.
5833 pub unsafe fn from_d3d11_shared_handle(
5834 handle: std::os::windows::io::RawHandle,
5835 width: usize,
5836 height: usize,
5837 format: PixelFormat,
5838 access: CpuAccess,
5839 completion: Option<(std::os::windows::io::RawHandle, u64)>,
5840 name: Option<&str>,
5841 ) -> Result<Self> {
5842 let dtype = dtype_of::<T>().ok_or_else(|| {
5843 Error::InvalidArgument(format!(
5844 "from_d3d11_shared_handle: element type {} has no DType",
5845 std::any::type_name::<T>()
5846 ))
5847 })?;
5848 let shape = format.allocation_shape(width, height).ok_or_else(|| {
5849 Error::InvalidArgument(format!(
5850 "invalid dimensions {width}x{height} for format {format:?}"
5851 ))
5852 })?;
5853 // SAFETY: the caller guarantees the handles are valid in this process.
5854 let storage = unsafe {
5855 crate::d3d11::texture::D3d11TextureTensor::<T>::from_d3d11_shared_handle(
5856 handle, width, height, format, dtype, &shape, access, completion, name,
5857 )
5858 }?;
5859 let mut t = Self::wrap(TensorStorage::Dma(storage));
5860 t.format = Some(format);
5861 record_d3d11_row_stride(&mut t, format, "from_d3d11_shared_handle");
5862 t.cpu_access = access;
5863 t.try_init_d3d11_cuda();
5864 Ok(t)
5865 }
5866}
5867
5868// Quantization accessors — type-gated to integer element types via the
5869// sealed `IntegerType` trait. Calling `.quantization()` on a `Tensor<f32>`
5870// produces a compile error, not a runtime one.
5871#[cfg(not(edgefirst_tensor_dynamic))]
5872impl<T> Tensor<T>
5873where
5874 T: IntegerType + Num + Clone + fmt::Debug + Send + Sync,
5875{
5876 /// Quantization metadata for this tensor, if set.
5877 pub fn quantization(&self) -> Option<&Quantization> {
5878 self.quantization.as_ref()
5879 }
5880
5881 /// Attach quantization metadata to this tensor. Validates against the
5882 /// tensor's shape — returns [`Error::QuantizationInvalid`] on any
5883 /// inconsistency (mismatched scale/zp lengths, out-of-range axis, etc.).
5884 pub fn set_quantization(&mut self, q: Quantization) -> Result<()> {
5885 q.validate(self.shape())?;
5886 self.quantization = Some(q);
5887 Ok(())
5888 }
5889
5890 /// Builder-style variant of [`Self::set_quantization`]. Consumes `self`
5891 /// and returns `Result<Self>` — on success yields the tensor with the
5892 /// attached quantization; on validation failure returns
5893 /// [`Error::QuantizationInvalid`] and drops `self` (the tensor is not
5894 /// returned in the error arm).
5895 pub fn with_quantization(mut self, q: Quantization) -> Result<Self> {
5896 self.set_quantization(q)?;
5897 Ok(self)
5898 }
5899
5900 /// Clear any quantization metadata on this tensor.
5901 pub fn clear_quantization(&mut self) {
5902 self.quantization = None;
5903 }
5904}
5905
5906/// Rank of one image's shape in `layout`: `[H, W, C]` / `[C, H, W]` are 3-D,
5907/// semi-planar `[H*k, W]` is 2-D. A batched tensor has one more.
5908#[cfg(not(edgefirst_tensor_dynamic))]
5909fn image_rank(layout: PixelLayout) -> usize {
5910 match layout {
5911 PixelLayout::SemiPlanar => 2,
5912 PixelLayout::Packed | PixelLayout::Planar => 3,
5913 }
5914}
5915
5916#[cfg(not(edgefirst_tensor_dynamic))]
5917fn map_strided_host_view<'a, T>(
5918 tensor: &Tensor<T>,
5919 stride: usize,
5920 access: CpuAccess,
5921 non_blocking: bool,
5922) -> Result<crate::view::HostView<'a, T>>
5923where
5924 T: Num + Clone + fmt::Debug + Send + Sync + 'a,
5925{
5926 // Rows sit at `stride`-byte spacing; see `strided_rows` for the count.
5927 let rows = strided_rows(tensor)?;
5928 let total_bytes = stride.checked_mul(rows).ok_or_else(|| {
5929 Error::InvalidOperation(format!(
5930 "Tensor::map: row_stride {stride} × rows {rows} overflows usize"
5931 ))
5932 })?;
5933 map_strided_storage(tensor, stride, rows, total_bytes, access, non_blocking)
5934}
5935
5936/// Rows a strided map or flat copy of `tensor` spans at its row stride: a
5937/// raw tensor's leading dim, else every image row of every batch element --
5938/// `H` per packed image, `C * H` per planar one (C planes of H rows), and the
5939/// combined `H*k` per semi-planar one. No product here can overflow: each is
5940/// a factor of the shape's element count, which the allocation already holds.
5941#[cfg(not(edgefirst_tensor_dynamic))]
5942fn strided_rows<T>(tensor: &Tensor<T>) -> Result<usize>
5943where
5944 T: Num + Clone + fmt::Debug + Send + Sync,
5945{
5946 let shape = tensor.shape();
5947 let Some(fmt) = tensor.format else {
5948 return shape.first().copied().ok_or_else(|| {
5949 Error::InvalidOperation("strided access requires a non-empty shape".into())
5950 });
5951 };
5952 let (_, dims) = tensor.image_dims().ok_or_else(|| {
5953 Error::InvalidOperation(format!(
5954 "strided access to a {fmt:?} tensor requires its image shape, optionally \
5955 batched; got {shape:?}"
5956 ))
5957 })?;
5958 let per_image = match fmt.layout() {
5959 PixelLayout::Planar => dims[0] * dims[1],
5960 PixelLayout::Packed | PixelLayout::SemiPlanar => dims[0],
5961 };
5962 let batch: usize = shape[..shape.len() - dims.len()].iter().product();
5963 Ok(batch * per_image)
5964}
5965
5966#[cfg(not(edgefirst_tensor_dynamic))]
5967fn map_strided_storage<'a, T>(
5968 tensor: &Tensor<T>,
5969 _stride: usize,
5970 _rows: usize,
5971 total_bytes: usize,
5972 access: CpuAccess,
5973 _non_blocking: bool,
5974) -> Result<crate::view::HostView<'a, T>>
5975where
5976 T: Num + Clone + fmt::Debug + Send + Sync + 'a,
5977{
5978 match &tensor.storage {
5979 #[cfg(target_os = "linux")]
5980 TensorStorage::Dma(dma) => {
5981 // `set_row_stride()` only validates `stride >= min_stride`,
5982 // not that `stride × rows` fits the DMA-BUF, so re-check
5983 // here — mapping past `buf_size` would SIGBUS on access.
5984 // `buf_size` is real kernel truth either way, not the
5985 // importer's say-so: for a self-allocated tensor it is
5986 // the allocation size; for an imported one it is
5987 // `fstat`-derived when the kernel reports a usable size
5988 // (`DmaTensor::from_fd`), or the declared logical size in
5989 // the rare fallback -- either way an out-of-range mmap
5990 // against the real dma-buf is still rejected by the
5991 // kernel, so this bound stays sound regardless.
5992 let available_bytes = dma.buf_size.saturating_sub(dma.mmap_offset);
5993 if total_bytes > available_bytes {
5994 return Err(Error::InvalidOperation(format!(
5995 "Tensor::map: strided mapping needs {total_bytes} bytes \
5996 but DMA buffer only has {available_bytes} available \
5997 (buf_size={}, mmap_offset={}, stride={_stride}, rows={_rows}); \
5998 the row_stride was likely set larger than the original allocation",
5999 dma.buf_size, dma.mmap_offset
6000 )));
6001 }
6002 dma.map_with_byte_size(total_bytes, access)
6003 }
6004 TensorStorage::Mem(mem) => {
6005 let capacity = tensor.storage.capacity_bytes();
6006 if total_bytes > capacity {
6007 return Err(Error::InsufficientCapacity {
6008 needed: total_bytes,
6009 capacity,
6010 });
6011 }
6012 mem.map_with_byte_size(total_bytes, access)
6013 }
6014 #[cfg(unix)]
6015 TensorStorage::Shm(shm) => {
6016 let capacity = tensor.storage.capacity_bytes();
6017 if total_bytes > capacity {
6018 return Err(Error::InsufficientCapacity {
6019 needed: total_bytes,
6020 capacity,
6021 });
6022 }
6023 shm.map_with_byte_size(total_bytes, access)
6024 }
6025 // macOS/iOS: `TensorStorage::Dma` is the IOSurface. The lock yields
6026 // the full surface base address, and the row pitch
6027 // (`IOSurfaceGetBytesPerRow`) is known from the API for both
6028 // self-allocated and imported surfaces — unlike a foreign
6029 // DMA-BUF — so a strided CPU view is sound and zero-copy.
6030 #[cfg(any(target_os = "macos", target_os = "ios"))]
6031 TensorStorage::Dma(io) => {
6032 // A sub-view's window is `buf_size − view_offset`; the strided
6033 // span must fit the window, not the whole surface.
6034 let available = io.buf_size.saturating_sub(io.view_offset);
6035 if total_bytes > available {
6036 return Err(Error::InsufficientCapacity {
6037 needed: total_bytes,
6038 capacity: available,
6039 });
6040 }
6041 io.map_with_byte_size(total_bytes, access)
6042 }
6043 // Android: `TensorStorage::Dma` is the AHardwareBuffer. The lock
6044 // yields the full buffer base address, and the row pitch is
6045 // known from the allocator-filled descriptor — so a strided CPU
6046 // view is sound and zero-copy, same as IOSurface.
6047 #[cfg(target_os = "android")]
6048 TensorStorage::Dma(ahb) => {
6049 // A sub-view's window is `buf_size − view_offset`; the strided
6050 // span must fit the window, not the whole buffer.
6051 let available = ahb.buf_size.saturating_sub(ahb.view_offset);
6052 if total_bytes > available {
6053 return Err(Error::InsufficientCapacity {
6054 needed: total_bytes,
6055 capacity: available,
6056 });
6057 }
6058 ahb.map_with_byte_size(total_bytes, access)
6059 }
6060 // Windows: `TensorStorage::Dma` is the D3D11 texture, whose CPU side
6061 // is the staging copy at the driver's pitch -- the very pitch
6062 // `Tensor::image` recorded as this tensor's row stride. Its
6063 // `map_with_byte_size` bounds-checks the span against this tensor's
6064 // own window, as the IOSurface arm above does. This is also the only
6065 // backing with a real non-blocking map, so it is the only arm that
6066 // reads `_non_blocking`.
6067 #[cfg(target_os = "windows")]
6068 TensorStorage::Dma(tex) => {
6069 if _non_blocking {
6070 tex.try_map_with_byte_size(total_bytes, access)
6071 } else {
6072 tex.map_with_byte_size(total_bytes, access)
6073 }
6074 }
6075 TensorStorage::Pbo(pbo) => {
6076 // PBO: the GPU-side allocation may have a padded row stride
6077 // (e.g. 64-byte aligned). Expose the full padded buffer so a
6078 // CPU producer (JPEG decoder) and a strided convert source
6079 // can iterate rows via `effective_row_stride()` without
6080 // running past the slice — the logical `pbo.map()` view would
6081 // stop after `shape.product()` and lose bytes past row 0.
6082 // A sub-view's window is `capacity − view_offset`.
6083 let available = pbo.capacity_bytes().saturating_sub(pbo.view_offset);
6084 if total_bytes > available {
6085 return Err(Error::InsufficientCapacity {
6086 needed: total_bytes,
6087 capacity: available,
6088 });
6089 }
6090 pbo.map_with_byte_size(total_bytes, access)
6091 }
6092 // Every `TensorStorage` variant this build can construct is
6093 // matched explicitly above (Mem/Shm/Dma/Pbo, each cfg-gated
6094 // to the platforms it exists on) -- this catch-all exists
6095 // only so the match stays exhaustive across every
6096 // OS/cfg combination without a matching arm on every one of
6097 // them, not because any variant is meant to fall through to
6098 // it. `#[allow(unreachable_patterns)]` because on any given
6099 // platform it genuinely is unreachable.
6100 #[allow(unreachable_patterns)]
6101 _ => Err(Error::InvalidOperation(
6102 "CPU mapping of strided tensors is supported only for HAL-allocated \
6103 Mem/Shm (any platform), DMA (Linux, self-allocated or imported), \
6104 IOSurface (macOS), AHardwareBuffer (Android), D3D11 textures \
6105 (Windows), and PBO"
6106 .into(),
6107 )),
6108 }
6109}
6110
6111#[cfg(not(edgefirst_tensor_dynamic))]
6112impl<T> Tensor<T>
6113where
6114 T: Num + Clone + fmt::Debug + Send + Sync,
6115{
6116 /// The shared body of [`TensorTrait::map_with`] and
6117 /// [`TensorTrait::try_map_with`]: the direction check, the
6118 /// declared-vs-requested telemetry, the strided and plane-offset routing,
6119 /// and finally the backing's own map. `non_blocking` chooses which of the
6120 /// backing's two maps that last step takes; everything before it is the
6121 /// same either way.
6122 fn map_impl<'a>(
6123 &self,
6124 access: CpuAccess,
6125 non_blocking: bool,
6126 ) -> Result<crate::view::HostView<'a, T>>
6127 where
6128 T: 'a,
6129 {
6130 let _span = tracing::trace_span!(
6131 "tensor.map",
6132 memory = ?self.storage.memory(),
6133 ?access,
6134 non_blocking,
6135 )
6136 .entered();
6137 self.cpu_access.require_direction(access, "map_with")?;
6138 // CPU mapping of a strided tensor exposes the full padded buffer
6139 // (`row_stride × rows`) so callers can iterate rows via
6140 // `effective_row_stride()` without running past the slice. This is sound
6141 // only when the HAL owns and can size-check the allocation:
6142 //
6143 // * Self-allocated Mem / Shm tensors (any platform) — the backing
6144 // `Vec` / shm segment is sized by `capacity_bytes()`, checked here.
6145 // * Self-allocated DMA tensors (Linux) — pitch padding from
6146 // `image_with_stride()`; checked against the DMA-BUF `buf_size`.
6147 //
6148 // * Self-allocated PBO tensors (any platform with GL) — the GL buffer
6149 // is sized by `capacity_bytes()` and may carry 64-byte row padding;
6150 // the JPEG decoder mmaps it and convert() reads it, both iterating
6151 // by `row_stride`. Checked against the PBO capacity below.
6152 //
6153 // A strided **imported** DMA-BUF is accepted here too (Task 10b): the
6154 // `stride × rows <= buf_size` check just below is exactly what a
6155 // self-allocated Dma tensor relies on for its own soundness, and
6156 // `buf_size` for an imported fd is `fstat`-derived when the kernel
6157 // reports a usable size (`stat.st_size > 0` and at least the
6158 // declared shape's logical size; `DmaTensor::from_fd`) -- kernel
6159 // truth, not a caller's claim, in that case. The rare fallback (a
6160 // kernel that reports `st_size <= 0` for a DMA-BUF fd) uses the
6161 // declared logical size instead; that is not a hole here either, since
6162 // an out-of-range mmap against the real dma-buf is still rejected by
6163 // the kernel regardless of what this crate believed `buf_size` was.
6164 // So the check is equally load-bearing regardless of who allocated
6165 // the buffer. What
6166 // changed is only *which* callers can now reach this path: the
6167 // cross-package capsule protocol reconstructs an imported alias of a
6168 // HAL-owned Dma tensor and then calls `configure_image`, which
6169 // computes its stride with the exact same alignment rule a
6170 // self-allocated `Tensor::image()` used — so the value being
6171 // bounds-checked here is HAL-computed and format/width-derived, not
6172 // read from an external driver. Rejecting it unconditionally used to
6173 // make every foreign-imported GPU-backed destination fail any CPU
6174 // write path (the JPEG decoder's MCU writer among them) even when the
6175 // write was provably in-bounds. `set_row_stride()` was already
6176 // "trust the caller, bounds-check the map" for every *other* backing
6177 // (Mem/Shm/self-allocated Dma/PBO) — singling out imported Dma out
6178 // for a stricter standard was the inconsistency, not the trust model
6179 // itself. IOSurface keeps its own always-allowed strided path
6180 // unconditionally (see its arm below): a live `IOSurfaceRef` reports
6181 // its real `bytesPerRow` from the OS regardless of who allocated it,
6182 // so there was never anything to distrust there.
6183 if let Some(stride) = self.row_stride {
6184 return map_strided_host_view(self, stride, access, non_blocking);
6185 }
6186 // A non-zero plane offset needs no check here: every storage applies
6187 // its own offset inside its map -- Mem/Shm adjust the slice base, the
6188 // Linux DMA-BUF and Android AHardwareBuffer the mapped base, IOSurface
6189 // and D3D11 textures add their `view_offset`, and PBO starts the
6190 // staged copy at it.
6191 if non_blocking {
6192 self.storage.try_map_with(access)
6193 } else {
6194 self.storage.map_with(access)
6195 }
6196 }
6197}
6198
6199#[cfg(all(
6200 not(edgefirst_tensor_dynamic),
6201 any(
6202 target_os = "macos",
6203 target_os = "ios",
6204 target_os = "android",
6205 target_os = "windows"
6206 )
6207))]
6208impl<T> Tensor<T>
6209where
6210 T: Num + Clone + fmt::Debug + Send + Sync,
6211{
6212 /// The pitched-backing half of [`TensorTrait::set_logical_shape`]: record
6213 /// the backing's own pitch as the row stride of the new `shape`.
6214 fn adopt_backing_pitch(&mut self, shape: &[usize]) {
6215 // Only a backing with a pitch of its own, and only while the shape
6216 // still has the rank the format's row is measured on: a caller may
6217 // flatten an image to one dimension and leave the format behind.
6218 let Some(pitch) = self.storage.backing_row_stride() else {
6219 return;
6220 };
6221 let Some((rank, rows)) = self.format.map(|f| match f.layout() {
6222 PixelLayout::Packed => (3, shape.first().copied()),
6223 PixelLayout::Planar => (3, shape.get(1).copied()),
6224 PixelLayout::SemiPlanar => (2, shape.first().copied()),
6225 }) else {
6226 return;
6227 };
6228 if rank != shape.len() {
6229 return;
6230 }
6231 // The natural stride of the new shape, not of the stride recorded for
6232 // the old one.
6233 let prior = self.row_stride.take();
6234 let Some(natural) = self.effective_row_stride() else {
6235 self.row_stride = prior;
6236 return;
6237 };
6238 if rows == Some(1) {
6239 self.set_row_stride_unchecked(natural);
6240 } else if pitch > natural {
6241 self.set_row_stride_unchecked(pitch);
6242 } else {
6243 self.row_stride = prior;
6244 }
6245 }
6246}
6247
6248#[cfg(not(edgefirst_tensor_dynamic))]
6249impl<T> TensorTrait<T> for Tensor<T>
6250where
6251 T: Num + Clone + fmt::Debug + Send + Sync,
6252{
6253 fn cpu_access(&self) -> CpuAccess {
6254 self.cpu_access
6255 }
6256
6257 fn view_origin(&self) -> Option<ViewOrigin> {
6258 self.view_origin
6259 }
6260
6261 fn is_view(&self) -> bool {
6262 Tensor::is_view(self)
6263 }
6264
6265 fn new(shape: &[usize], name: Option<&str>) -> Result<Self>
6266 where
6267 Self: Sized,
6268 {
6269 Self::new(shape, None, name)
6270 }
6271
6272 #[cfg(unix)]
6273 fn from_fd(fd: std::os::fd::OwnedFd, shape: &[usize], name: Option<&str>) -> Result<Self>
6274 where
6275 Self: Sized,
6276 {
6277 Self::from_fd(fd, shape, name)
6278 }
6279
6280 #[cfg(unix)]
6281 fn clone_fd(&self) -> Result<std::os::fd::OwnedFd> {
6282 self.storage.clone_fd()
6283 }
6284
6285 fn memory(&self) -> TensorMemory {
6286 self.storage.memory()
6287 }
6288
6289 fn name(&self) -> String {
6290 self.storage.name()
6291 }
6292
6293 fn shape(&self) -> &[usize] {
6294 self.storage.shape()
6295 }
6296
6297 /// Bytes available in the underlying allocation, from this tensor's
6298 /// window start -- may exceed [`TensorTrait::size`] when the backing
6299 /// carries spare room beyond the current logical shape (a pitch-aligned
6300 /// image, or a pool tensor sized for a larger image than it holds
6301 /// today). See [`crate::protocol::TensorDesc::capacity`], the
6302 /// cross-package descriptor field this backs.
6303 fn capacity_bytes(&self) -> usize {
6304 self.storage.capacity_bytes()
6305 }
6306
6307 fn reshape(&mut self, shape: &[usize]) -> Result<()> {
6308 if self.chroma.is_some() {
6309 return Err(Error::InvalidOperation(
6310 "cannot reshape a multiplane tensor — decompose planes first".into(),
6311 ));
6312 }
6313 self.storage.reshape(shape)?;
6314 self.format = None;
6315 self.row_stride = None;
6316 self.plane_offset = None;
6317 match self.storage {
6318 TensorStorage::Mem(ref mut m) => m.set_offset(0),
6319 #[cfg(target_os = "linux")]
6320 TensorStorage::Dma(ref mut dma) => dma.mmap_offset = 0,
6321 // No-op: `IoSurfaceTensor::reshape` (iosurface.rs) zeroes its
6322 // own `view_offset` before this match runs.
6323 #[cfg(any(target_os = "macos", target_os = "ios"))]
6324 TensorStorage::Dma(_) => {}
6325 // `D3d11TextureTensor::reshape` does not self-zero, so the clear
6326 // that pairs with `set_plane_offset`'s D3D11 arm lives here.
6327 #[cfg(target_os = "windows")]
6328 TensorStorage::Dma(ref mut tex) => tex.set_view_offset(0),
6329 // No-op: `AHardwareBufferTensor::reshape` (ahardwarebuffer.rs)
6330 // zeroes its own `view_offset` before this match runs, the same
6331 // as `IoSurfaceTensor`. Unlike `set_plane_offset`'s Android arm,
6332 // this one is a genuine self-zeroing no-op, not a still-
6333 // incomplete gap: `view_offset` never survives a reshape here
6334 // regardless of whether `set_plane_offset` can restore it.
6335 #[cfg(target_os = "android")]
6336 TensorStorage::Dma(_) => {}
6337 // No-op: `PboTensor::reshape` (pbo.rs) zeroes its own
6338 // `view_offset` before this match runs, the same as
6339 // `IoSurfaceTensor`/`AHardwareBufferTensor` above -- which is
6340 // why, unlike `set_format`'s clear block, this match never
6341 // needed a `Pbo` arm that writes anything.
6342 TensorStorage::Pbo(_) => {}
6343 // Unlike `PboTensor::reshape`, `ShmTensor::reshape` (shm.rs)
6344 // only revalidates the element count and swaps `self.shape` --
6345 // it never touches `self.offset` -- so this clear site needs
6346 // the arm `set_format`'s does. `cfg(unix)` mirrors the `Shm`
6347 // variant's own gate, not a per-OS multiplex.
6348 #[cfg(unix)]
6349 TensorStorage::Shm(ref mut s) => s.set_offset(0),
6350 }
6351 Ok(())
6352 }
6353
6354 /// Set the logical shape to any shape whose bytes fit the allocation,
6355 /// without [`reshape`](TensorTrait::reshape)'s equal-count constraint.
6356 ///
6357 /// Forwards to the storage's capacity-aware implementation -- the same
6358 /// one [`Tensor::configure_image`] uses for pool reuse, a decode into an
6359 /// oversized reusable buffer reconfiguring it to a smaller image without
6360 /// reallocating -- so the two agree. The pixel format is kept.
6361 ///
6362 /// A narrower image inside a pitched backing (an IOSurface, an
6363 /// AHardwareBuffer, a D3D11 texture) still advances rows by the backing's
6364 /// pitch, so that pitch is recorded as the row stride whenever it exceeds
6365 /// the new shape's natural stride -- the same adoption `configure_image`
6366 /// makes. Without it a texture narrowed from 64 to 16 texels wide reports
6367 /// a 64-byte stride over rows that are 256 bytes apart. A single-row
6368 /// shape records its own tight stride instead, by [`view`](Self::view)'s
6369 /// rule: a strided `map()` spans `stride * rows`, and one pitched row
6370 /// runs past the backing when the window sits in the last row.
6371 fn set_logical_shape(&mut self, shape: &[usize]) -> Result<()> {
6372 self.storage.set_logical_shape(shape)?;
6373 #[cfg(any(
6374 target_os = "macos",
6375 target_os = "ios",
6376 target_os = "android",
6377 target_os = "windows"
6378 ))]
6379 self.adopt_backing_pitch(shape);
6380 Ok(())
6381 }
6382
6383 fn map_with<'a>(&self, access: CpuAccess) -> Result<crate::view::HostView<'a, T>>
6384 where
6385 T: 'a,
6386 {
6387 self.map_impl(access, false)
6388 }
6389
6390 fn try_map_with<'a>(&self, access: CpuAccess) -> Result<crate::view::HostView<'a, T>>
6391 where
6392 T: 'a,
6393 {
6394 self.map_impl(access, true)
6395 }
6396
6397 fn buffer_identity(&self) -> &BufferIdentity {
6398 self.storage.buffer_identity()
6399 }
6400}
6401
6402// ============================================================================
6403// Platform availability helpers
6404// ============================================================================
6405
6406/// Cached result of the Linux DMA-BUF availability probe.
6407#[cfg(target_os = "linux")]
6408static DMA_AVAILABLE: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
6409/// Cached result of the macOS/iOS IOSurface availability probe.
6410#[cfg(any(target_os = "macos", target_os = "ios"))]
6411static IOSURFACE_AVAILABLE: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
6412
6413/// Check if Linux DMA-BUF allocation is available on this system.
6414///
6415/// Linux-specific availability check (typically requires `/dev/dma_heap`
6416/// access — running as root or membership in a video/render group). For
6417/// portable code that wants "any zero-copy GPU buffer", use
6418/// [`is_gpu_buffer_available`] which also covers IOSurface on macOS.
6419///
6420/// This function caches its result after the first call.
6421#[cfg(target_os = "linux")]
6422pub fn is_dma_available() -> bool {
6423 *DMA_AVAILABLE
6424 .get_or_init(|| Tensor::<u8>::new(&[64], Some(TensorMemory::DmaBuf), None).is_ok())
6425}
6426
6427/// Always returns `false` on non-Linux platforms.
6428#[cfg(not(target_os = "linux"))]
6429pub fn is_dma_available() -> bool {
6430 false
6431}
6432
6433/// Check if macOS/iOS IOSurface allocation is available on this system.
6434///
6435/// IOSurface is part of the macOS/iOS OS and is essentially always present;
6436/// this probe catches degraded scenarios such as memory pressure or
6437/// sandboxed contexts where `IOSurfaceCreate` fails. The result is
6438/// cached after the first call.
6439///
6440/// Always returns `false` on non-Apple platforms.
6441#[cfg(any(target_os = "macos", target_os = "ios"))]
6442pub fn is_iosurface_available() -> bool {
6443 *IOSURFACE_AVAILABLE.get_or_init(|| {
6444 // Probe via the same Dma path — on macOS/iOS this routes through
6445 // IoSurfaceTensor::new.
6446 Tensor::<u8>::new(&[64], Some(TensorMemory::DmaBuf), None).is_ok()
6447 })
6448}
6449
6450#[cfg(not(any(target_os = "macos", target_os = "ios")))]
6451pub fn is_iosurface_available() -> bool {
6452 false
6453}
6454
6455/// Cached result of the Android AHardwareBuffer availability probe.
6456#[cfg(target_os = "android")]
6457static AHARDWAREBUFFER_AVAILABLE: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
6458
6459/// Check if Android AHardwareBuffer allocation is available on this system.
6460///
6461/// AHardwareBuffer is part of the Android OS (public NDK ABI since API
6462/// 26) and is essentially always present; this probe catches degraded
6463/// scenarios such as memory pressure or gralloc failures. The result is
6464/// cached after the first call.
6465#[cfg(target_os = "android")]
6466pub fn is_ahardwarebuffer_available() -> bool {
6467 *AHARDWAREBUFFER_AVAILABLE.get_or_init(|| {
6468 // Probe via the same Dma path — on Android this routes through
6469 // AHardwareBufferTensor::new.
6470 Tensor::<u8>::new(&[64], Some(TensorMemory::DmaBuf), None).is_ok()
6471 })
6472}
6473
6474/// Always returns `false` on non-Android platforms.
6475#[cfg(not(target_os = "android"))]
6476pub fn is_ahardwarebuffer_available() -> bool {
6477 false
6478}
6479
6480/// Portable probe for the platform's native zero-copy GPU buffer allocator
6481/// (DMA-BUF on Linux, IOSurface on macOS/iOS, AHardwareBuffer on Android,
6482/// D3D11 textures on Windows). Returns `false` on platforms with no
6483/// equivalent. Use this when writing cross-platform code that cares whether
6484/// the `Dma` tensor variant will work, not which underlying mechanism is used.
6485///
6486/// On Windows the `Dma` variant is image-formatted: `true` here means
6487/// [`Tensor::image`] can serve `TensorMemory::DmaBuf`, not [`Tensor::new`].
6488///
6489/// # Ordering on Windows
6490///
6491/// This probe **creates the process D3D11 device**, because the only honest
6492/// answer to "can this host allocate one" is to try. A host that must supply
6493/// its own device calls `d3d11::use_external_device` **before** this probe or
6494/// any allocation: the device is cached on first creation and installing one
6495/// afterwards is refused.
6496pub fn is_gpu_buffer_available() -> bool {
6497 #[cfg(target_os = "linux")]
6498 {
6499 is_dma_available()
6500 }
6501 #[cfg(any(target_os = "macos", target_os = "ios"))]
6502 {
6503 is_iosurface_available()
6504 }
6505 #[cfg(target_os = "android")]
6506 {
6507 is_ahardwarebuffer_available()
6508 }
6509 // No cached probe: `d3d11::device()` is itself a `OnceLock` over the
6510 // process device, so asking twice costs one atomic load.
6511 #[cfg(target_os = "windows")]
6512 {
6513 crate::d3d11::device().is_ok()
6514 }
6515 #[cfg(not(any(
6516 target_os = "linux",
6517 target_os = "macos",
6518 target_os = "ios",
6519 target_os = "android",
6520 target_os = "windows"
6521 )))]
6522 {
6523 false
6524 }
6525}
6526
6527/// Check if POSIX shared memory allocation is available on this system.
6528///
6529/// Returns `true` on Unix systems (Linux, macOS, BSD) where POSIX shared memory
6530/// is supported. Always returns `false` on non-Unix platforms (Windows).
6531///
6532/// This function caches its result after the first call for efficiency.
6533#[cfg(unix)]
6534static SHM_AVAILABLE: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
6535
6536/// Check if POSIX shared memory allocation is available on this system.
6537#[cfg(unix)]
6538pub fn is_shm_available() -> bool {
6539 *SHM_AVAILABLE.get_or_init(|| Tensor::<u8>::new(&[64], Some(TensorMemory::Shm), None).is_ok())
6540}
6541
6542/// Check if POSIX shared memory allocation is available on this system.
6543///
6544/// Always returns `false` on non-Unix platforms since POSIX SHM is Unix-specific.
6545#[cfg(not(unix))]
6546pub fn is_shm_available() -> bool {
6547 false
6548}
6549
6550#[cfg(test)]
6551mod dtype_tests {
6552 use super::*;
6553
6554 #[test]
6555 fn dtype_size() {
6556 assert_eq!(DType::U8.size(), 1);
6557 assert_eq!(DType::I8.size(), 1);
6558 assert_eq!(DType::U16.size(), 2);
6559 assert_eq!(DType::I16.size(), 2);
6560 assert_eq!(DType::U32.size(), 4);
6561 assert_eq!(DType::I32.size(), 4);
6562 assert_eq!(DType::U64.size(), 8);
6563 assert_eq!(DType::I64.size(), 8);
6564 assert_eq!(DType::F16.size(), 2);
6565 assert_eq!(DType::F32.size(), 4);
6566 assert_eq!(DType::F64.size(), 8);
6567 }
6568
6569 #[test]
6570 fn dtype_name() {
6571 assert_eq!(DType::U8.name(), "u8");
6572 assert_eq!(DType::F16.name(), "f16");
6573 assert_eq!(DType::F32.name(), "f32");
6574 }
6575
6576 #[test]
6577 fn dtype_serde_roundtrip() {
6578 use serde_json;
6579 let dt = DType::F16;
6580 let json = serde_json::to_string(&dt).unwrap();
6581 let back: DType = serde_json::from_str(&json).unwrap();
6582 assert_eq!(dt, back);
6583 }
6584}
6585
6586#[cfg(test)]
6587mod image_tests {
6588 use super::*;
6589
6590 #[test]
6591 fn image_dims_from_shape_inverts_allocation_shape() {
6592 for (format, w, h) in [
6593 (PixelFormat::Rgba, 640usize, 480usize),
6594 (PixelFormat::Rgb, 641, 3),
6595 (PixelFormat::Grey, 32, 32),
6596 (PixelFormat::PlanarRgb, 640, 480),
6597 (PixelFormat::PlanarRgba, 8, 5),
6598 (PixelFormat::Nv12, 640, 481),
6599 (PixelFormat::Nv16, 33, 17),
6600 (PixelFormat::Nv24, 16, 9),
6601 ] {
6602 let shape = format.allocation_shape(w, h).unwrap();
6603 assert_eq!(
6604 image_dims_from_shape(format, &shape).unwrap(),
6605 (w, h),
6606 "{format:?} {w}x{h} via {shape:?}"
6607 );
6608 }
6609 }
6610
6611 #[test]
6612 fn image_dims_from_shape_rejects_the_wrong_rank() {
6613 // Packed and planar want rank 3, semi-planar rank 2.
6614 for (format, shape) in [
6615 (PixelFormat::Rgba, vec![480usize, 640usize]),
6616 (PixelFormat::PlanarRgb, vec![3, 480]),
6617 (PixelFormat::Nv12, vec![720, 640, 1]),
6618 ] {
6619 let err = image_dims_from_shape(format, &shape).unwrap_err();
6620 assert!(
6621 matches!(err, Error::InvalidShape(ref m) if m.contains("rank")),
6622 "{format:?} {shape:?}: {err}"
6623 );
6624 }
6625 }
6626
6627 #[test]
6628 fn image_dims_from_shape_rejects_an_unproduced_combined_height() {
6629 // NV12 maps 480 -> 720 and 481 -> 722; nothing maps to 721, so the
6630 // shape is rejected rather than rounded to 480.
6631 let err = image_dims_from_shape(PixelFormat::Nv12, &[721, 640]).unwrap_err();
6632 assert!(
6633 matches!(err, Error::InvalidShape(ref m) if m.contains("combined plane height")),
6634 "{err}"
6635 );
6636 }
6637
6638 #[test]
6639 fn allocation_shape_per_layout() {
6640 assert_eq!(
6641 PixelFormat::Rgb.allocation_shape(640, 480),
6642 Some(vec![480, 640, 3])
6643 );
6644 assert_eq!(
6645 PixelFormat::Grey.allocation_shape(640, 480),
6646 Some(vec![480, 640, 1])
6647 );
6648 assert_eq!(
6649 PixelFormat::Nv12.allocation_shape(640, 480),
6650 Some(vec![720, 640])
6651 );
6652 // Odd height: combined-plane height is `481 + ceil(481/2)` = 481 + 241
6653 // = 722 rows. Logical height is recovered as `722 * 2 / 3` = 481.
6654 assert_eq!(
6655 PixelFormat::Nv12.allocation_shape(640, 481),
6656 Some(vec![722, 640])
6657 );
6658 // Odd width: shape carries the LOGICAL width (641).
6659 // The 64-aligned stride (>= 642) is stored separately on the Tensor.
6660 assert_eq!(
6661 PixelFormat::Nv12.allocation_shape(641, 480),
6662 Some(vec![720, 641])
6663 );
6664 // NV16 odd width: same — logical width in shape, stride separate.
6665 assert_eq!(
6666 PixelFormat::Nv16.allocation_shape(641, 480),
6667 Some(vec![960, 641])
6668 );
6669 assert_eq!(
6670 PixelFormat::PlanarRgb.allocation_shape(640, 480),
6671 Some(vec![3, 480, 640])
6672 );
6673 assert_eq!(
6674 PixelFormat::Nv16.allocation_shape(640, 480),
6675 Some(vec![960, 640])
6676 );
6677 }
6678
6679 #[test]
6680 fn raw_tensor_has_no_format() {
6681 let t = Tensor::<u8>::new(&[480, 640, 3], None, None).unwrap();
6682 assert!(t.format().is_none());
6683 assert!(t.width().is_none());
6684 assert!(t.height().is_none());
6685 assert!(!t.is_multiplane());
6686 assert!(t.chroma().is_none());
6687 }
6688
6689 #[test]
6690 fn image_tensor_packed() {
6691 let t = Tensor::<u8>::image(
6692 640,
6693 480,
6694 PixelFormat::Rgba,
6695 None,
6696 crate::CpuAccess::ReadWrite,
6697 )
6698 .unwrap();
6699 assert_eq!(t.format(), Some(PixelFormat::Rgba));
6700 assert_eq!(t.width(), Some(640));
6701 assert_eq!(t.height(), Some(480));
6702 assert_eq!(t.shape(), &[480, 640, 4]);
6703 assert!(!t.is_multiplane());
6704 }
6705
6706 #[test]
6707 fn image_tensor_planar() {
6708 let t = Tensor::<u8>::image(
6709 640,
6710 480,
6711 PixelFormat::PlanarRgb,
6712 None,
6713 crate::CpuAccess::ReadWrite,
6714 )
6715 .unwrap();
6716 assert_eq!(t.format(), Some(PixelFormat::PlanarRgb));
6717 assert_eq!(t.width(), Some(640));
6718 assert_eq!(t.height(), Some(480));
6719 assert_eq!(t.shape(), &[3, 480, 640]);
6720 }
6721
6722 #[test]
6723 #[cfg(target_os = "macos")]
6724 fn image_tensor_dma_non_aligned_packed_width_pads_zero_copy() {
6725 // Declares this test as an fd-opener; see FD_LOCK.
6726 let _lock = crate::tests::fd_lock_shared();
6727 // RGBA u8 at width=4 → 4*4 = 16 bytes/row, not 64-byte aligned. RGBA has
6728 // a real IOSurface FourCC, so an explicit `Some(TensorMemory::DmaBuf)`
6729 // request now allocates a padded image IOSurface (64-aligned
6730 // `bytes_per_row`) and records the stride — a fully zero-copy buffer GL
6731 // can bind and the CPU can map via the strided path. (Previously this
6732 // failed loudly to avoid an 'L008' byte-bag downgrade; with a real
6733 // FourCC surface that concern no longer applies.)
6734 let t = Tensor::<u8>::image(
6735 4,
6736 4,
6737 PixelFormat::Rgba,
6738 Some(TensorMemory::DmaBuf),
6739 crate::CpuAccess::ReadWrite,
6740 )
6741 .expect("padded RGBA IOSurface should allocate");
6742 assert_eq!(t.format(), Some(PixelFormat::Rgba));
6743 assert_eq!(t.width(), Some(4));
6744 assert_eq!(t.height(), Some(4));
6745 let stride = t.effective_row_stride().expect("stride");
6746 assert_eq!(stride % 64, 0, "padded to 64-byte row alignment");
6747 assert!(stride >= 16);
6748 // A CPU map exposes the full padded surface for strided iteration.
6749 let m = t.map().expect("strided IOSurface map");
6750 assert_eq!(m.as_slice().len(), stride * 4);
6751 }
6752
6753 /// `per_pixel_bytes` that doesn't divide 64 evenly (e.g. RGB u8 with
6754 /// 3 B/pixel) makes a "Pad width to N" suggestion structurally
6755 /// impossible — there is no integer width whose `width * 3` is a
6756 /// multiple of 64. The error must still fire (no silent SHM
6757 /// fallback for explicit-DMA requests) and must spell out the
6758 /// alignment requirement; it just omits the misleading "pad to N"
6759 /// hint instead of printing a number whose row pitch still won't
6760 /// align.
6761 #[test]
6762 #[cfg(target_os = "macos")]
6763 fn image_tensor_dma_rejects_indivisible_pixel_pitch_without_pad_hint() {
6764 // Declares this test as an fd-opener; see FD_LOCK.
6765 let _lock = crate::tests::fd_lock_shared();
6766 // Width=10 RGB f32 → 120 B/row, not 64-byte aligned, and (Rgb,
6767 // F32) has no IOSurface mapping so the padded-stride tolerance
6768 // does not apply. The next 64-multiple (128 B) isn't an integer
6769 // multiple of 12 B/pixel, so the "pad width to N" hint can't
6770 // produce a valid number and must be omitted. (Rgb u8 used to be
6771 // this test's subject but now has a real RGBA8888 mapping with
6772 // padded-stride tolerance — see the test below.)
6773 let err = Tensor::<f32>::image(
6774 10,
6775 10,
6776 PixelFormat::Rgb,
6777 Some(TensorMemory::DmaBuf),
6778 crate::CpuAccess::ReadWrite,
6779 )
6780 .expect_err("RGB f32 with 12 B/pixel and non-aligned width must be rejected");
6781 match err {
6782 Error::InvalidArgument(msg) => {
6783 assert!(
6784 msg.contains("64-byte aligned"),
6785 "error must still name the alignment requirement: {msg}"
6786 );
6787 assert!(
6788 !msg.contains("Pad width"),
6789 "indivisible per-pixel pitch makes a width suggestion impossible; \
6790 hint must be omitted, got: {msg}"
6791 );
6792 assert!(
6793 msg.contains("memory=None") && msg.contains("TensorMemory::Mem"),
6794 "error must still list the always-applicable alternatives: {msg}"
6795 );
6796 }
6797 other => panic!("expected InvalidArgument, got {other:?}"),
6798 }
6799 }
6800
6801 #[test]
6802 #[cfg(target_os = "macos")]
6803 fn image_tensor_dma_packed_rgb_u8_contract() {
6804 // Declares this test as an fd-opener; see FD_LOCK.
6805 let _lock = crate::tests::fd_lock_shared();
6806 // Packed RGB u8 @Dma is a designed RGBA8888 mapping at
6807 // (W*3/4, H) — the INT8 NPU input layout, shared with Android.
6808 // width%4 != 0 cannot form whole texels → loud InvalidArgument…
6809 let err = Tensor::<u8>::image(
6810 10,
6811 10,
6812 PixelFormat::Rgb,
6813 Some(TensorMemory::DmaBuf),
6814 crate::CpuAccess::ReadWrite,
6815 )
6816 .expect_err("Rgb u8 width%4!=0 must be rejected");
6817 assert!(
6818 matches!(&err, Error::InvalidArgument(m) if m.contains("width%4==0")),
6819 "got {err:?}"
6820 );
6821 // …width%4 == 0 with a non-64-aligned pitch allocates PADDED
6822 // (36 B rows → 64 B surface pitch, recorded on the tensor)…
6823 let t = Tensor::<u8>::image(
6824 12,
6825 4,
6826 PixelFormat::Rgb,
6827 Some(TensorMemory::DmaBuf),
6828 crate::CpuAccess::ReadWrite,
6829 )
6830 .expect("width 12 Rgb u8 must allocate padded");
6831 assert_eq!(t.memory(), TensorMemory::DmaBuf);
6832 assert!(
6833 t.row_stride()
6834 .is_some_and(|s| s >= 64 && s.is_multiple_of(64)),
6835 "padded pitch must be recorded: {:?}",
6836 t.row_stride()
6837 );
6838 // …and the aligned model-input width stays flat (640*3 = 1920 is
6839 // 64-aligned → no recorded stride, the buffer IS [H, W, 3]).
6840 let t = Tensor::<u8>::image(
6841 640,
6842 8,
6843 PixelFormat::Rgb,
6844 Some(TensorMemory::DmaBuf),
6845 crate::CpuAccess::ReadWrite,
6846 )
6847 .expect("width 640 Rgb u8 must allocate flat");
6848 assert_eq!(t.row_stride(), None);
6849 // I8 shares the layout (INT8 shader bias, not a format change).
6850 let t = Tensor::<i8>::image(
6851 640,
6852 8,
6853 PixelFormat::Rgb,
6854 Some(TensorMemory::DmaBuf),
6855 crate::CpuAccess::ReadWrite,
6856 )
6857 .expect("Rgb i8 shares the RGBA8888 mapping");
6858 assert_eq!(t.memory(), TensorMemory::DmaBuf);
6859 }
6860
6861 #[test]
6862 #[cfg(target_os = "macos")]
6863 fn image_tensor_dma_planar_f16_alignment() {
6864 // Declares this test as an fd-opener; see FD_LOCK.
6865 let _lock = crate::tests::fd_lock_shared();
6866 // PlanarRgb F16 uses single-channel row pitch (width * 2 bytes).
6867 // Width=16 → 32 bytes/row (not aligned); width=32 → 64 bytes/row (aligned).
6868 let err = Tensor::<half::f16>::image(
6869 16,
6870 16,
6871 PixelFormat::PlanarRgb,
6872 Some(TensorMemory::DmaBuf),
6873 crate::CpuAccess::ReadWrite,
6874 )
6875 .expect_err("width=16 PlanarRgb F16 is 32-byte row, must reject");
6876 assert!(matches!(err, Error::InvalidArgument(_)), "got {err:?}");
6877 // 32 wide should work.
6878 let t = Tensor::<half::f16>::image(
6879 32,
6880 8,
6881 PixelFormat::PlanarRgb,
6882 Some(TensorMemory::DmaBuf),
6883 crate::CpuAccess::ReadWrite,
6884 )
6885 .expect("width=32 PlanarRgb F16 is 64-byte row, must succeed");
6886 assert_eq!(t.format(), Some(PixelFormat::PlanarRgb));
6887 }
6888
6889 #[test]
6890 fn image_tensor_semi_planar_contiguous() {
6891 let t = Tensor::<u8>::image(
6892 640,
6893 480,
6894 PixelFormat::Nv12,
6895 None,
6896 crate::CpuAccess::ReadWrite,
6897 )
6898 .unwrap();
6899 assert_eq!(t.format(), Some(PixelFormat::Nv12));
6900 assert_eq!(t.width(), Some(640));
6901 assert_eq!(t.height(), Some(480));
6902 // NV12: H*3/2 = 720
6903 assert_eq!(t.shape(), &[720, 640]);
6904 assert!(!t.is_multiplane());
6905 }
6906
6907 #[test]
6908 #[cfg(target_os = "linux")]
6909 fn image_tensor_with_stride_preserves_logical_width() {
6910 // Declares this test as an fd-opener; see FD_LOCK.
6911 let _lock = crate::tests::fd_lock_shared();
6912 if !crate::test_support::dma_or_skip("image_tensor_with_stride_preserves_logical_width") {
6913 return;
6914 }
6915 // 3004×1688 RGBA8: natural pitch 12016, padded to 12032 (64-aligned).
6916 let stride = 12032;
6917 let t = Tensor::<u8>::image_with_stride(
6918 3004,
6919 1688,
6920 PixelFormat::Rgba,
6921 stride,
6922 Some(TensorMemory::DmaBuf),
6923 crate::CpuAccess::ReadWrite,
6924 )
6925 .unwrap();
6926 // Logical dimensions unchanged by padding — this is the contract.
6927 assert_eq!(t.width(), Some(3004));
6928 assert_eq!(t.height(), Some(1688));
6929 assert_eq!(t.shape(), &[1688, 3004, 4]);
6930 // Stride is carried separately and reports the padded pitch.
6931 assert_eq!(t.effective_row_stride(), Some(stride));
6932 // Buffer is sized to stride × height so the full padded layout fits,
6933 // and CPU map() works for self-allocated strided DMA tensors.
6934 use crate::TensorMapTrait;
6935 {
6936 let map = t.map().unwrap();
6937 assert!(
6938 map.as_slice().len() >= stride * 1688,
6939 "mapped buffer {} bytes < expected {}",
6940 map.as_slice().len(),
6941 stride * 1688
6942 );
6943 }
6944 // CPU write access works too — iterate rows using the padded stride,
6945 // touch only the active `width × bpp` region, verify it round-trips.
6946 {
6947 let mut map = t.map().unwrap();
6948 let slice = map.as_mut_slice();
6949 for y in 0..1688 {
6950 let row_start = y * stride;
6951 for x in 0..3004 {
6952 let p = row_start + x * 4;
6953 slice[p] = (y & 0xFF) as u8;
6954 slice[p + 1] = (x & 0xFF) as u8;
6955 slice[p + 2] = 0x42;
6956 slice[p + 3] = 0xFF;
6957 }
6958 }
6959 }
6960 {
6961 let map = t.map().unwrap();
6962 let slice = map.as_slice();
6963 // Sample a few pixels to confirm the round-trip.
6964 assert_eq!(slice[0], 0x00);
6965 assert_eq!(slice[1], 0x00);
6966 assert_eq!(slice[2], 0x42);
6967 assert_eq!(slice[3], 0xFF);
6968 let mid = 100 * stride + 50 * 4;
6969 assert_eq!(slice[mid], 100);
6970 assert_eq!(slice[mid + 1], 50);
6971 assert_eq!(slice[mid + 2], 0x42);
6972 }
6973 }
6974
6975 #[test]
6976 #[cfg(target_os = "linux")]
6977 fn image_tensor_foreign_strided_map_is_bounds_checked_not_refused() {
6978 // An imported (foreign) DMA-BUF with a row stride used to be refused
6979 // outright. That changed deliberately: the cross-package capsule
6980 // protocol reconstructs an imported alias of a HAL-owned Dma tensor,
6981 // and refusing every strided CPU map broke the flagship
6982 // create_image() -> decode_file_into() -> convert() pipeline on Linux.
6983 //
6984 // What replaced the refusal is a bounds check against `buf_size`,
6985 // which for an imported fd is `fstat`-derived kernel truth rather
6986 // than the caller's claim. So the property worth testing is no longer
6987 // "foreign strided maps fail" but "they succeed exactly when they
6988 // fit" -- which is what actually keeps the mapping sound.
6989 //
6990 // This test previously asserted the old refusal and went red on every
6991 // board with a dma-heap when the behaviour changed, unnoticed because
6992 // it needs both Linux and a dma-heap: it skips on macOS and on Orin
6993 // (nvmap, no CONFIG_DMABUF_HEAPS).
6994 if !crate::test_support::dma_or_skip(
6995 "image_tensor_foreign_strided_map_is_bounds_checked_not_refused",
6996 ) {
6997 return;
6998 }
6999 let backing =
7000 Tensor::<u8>::new(&[240 * 320 * 4], Some(TensorMemory::DmaBuf), None).unwrap();
7001
7002 // In bounds: the natural stride for the declared image fits the
7003 // allocation, so the map must succeed.
7004 {
7005 let fd = backing.clone_fd().unwrap();
7006 let shape = [240usize, 320, 4];
7007 let storage = TensorStorage::<u8>::from_fd(fd, &shape, None).unwrap();
7008 let mut t = Tensor::<u8>::wrap(storage);
7009 t.set_format(PixelFormat::Bgra).unwrap();
7010 t.set_row_stride(320 * 4).unwrap();
7011 assert!(
7012 t.map().is_ok(),
7013 "an in-bounds strided map of an imported DMA-BUF must be \
7014 accepted -- refusing it is what broke decode-into-GPU-tensor"
7015 );
7016 }
7017
7018 // Out of bounds: a stride the allocation cannot back must still be
7019 // refused, and the error must say so. This is the half that carries
7020 // the soundness.
7021 {
7022 let fd = backing.clone_fd().unwrap();
7023 let shape = [240usize, 320, 4];
7024 let storage = TensorStorage::<u8>::from_fd(fd, &shape, None).unwrap();
7025 let mut t = Tensor::<u8>::wrap(storage);
7026 t.set_format(PixelFormat::Bgra).unwrap();
7027 // Double the pitch: 240 rows x 2560 bytes is twice the allocation.
7028 t.set_row_stride(320 * 4 * 2).unwrap();
7029 let err = t.map();
7030 assert!(
7031 matches!(err, Err(Error::InvalidOperation(_))),
7032 "a strided map larger than the DMA allocation must be refused"
7033 );
7034 }
7035 }
7036
7037 #[test]
7038 #[cfg(target_os = "linux")]
7039 fn image_tensor_with_stride_map_rejects_tampered_stride() {
7040 // Declares this test as an fd-opener; see FD_LOCK.
7041 let _lock = crate::tests::fd_lock_shared();
7042 // Round-3 PR feedback (C1): `set_row_stride` is public and only
7043 // validates `stride >= min_stride`, not that the new stride × height
7044 // fits the underlying buffer. A caller that tampers with the stride
7045 // after allocation must not be able to coerce `Tensor::map()` into
7046 // returning a slice larger than the backing mmap (that would be UB
7047 // in `DmaMap::as_slice`).
7048 if !crate::test_support::dma_or_skip("image_tensor_with_stride_map_rejects_tampered_stride")
7049 {
7050 return;
7051 }
7052 // Allocate a 640×480 RGBA8 padded canvas (stride = 3072 = 768 px).
7053 // Backing buffer is 3072 × 480 = 1,474,560 bytes.
7054 let mut t = Tensor::<u8>::image_with_stride(
7055 640,
7056 480,
7057 PixelFormat::Rgba,
7058 3072,
7059 Some(TensorMemory::DmaBuf),
7060 crate::CpuAccess::ReadWrite,
7061 )
7062 .unwrap();
7063 // Tamper: push the stride up to 4 × the original. This is >=
7064 // min_stride (2560), so `set_row_stride` accepts it.
7065 t.set_row_stride(12288).unwrap();
7066 // Map must now refuse — 12288 × 480 = 5,898,240 > 1,474,560.
7067 let err = t.map();
7068 assert!(
7069 matches!(err, Err(Error::InvalidOperation(_))),
7070 "map() with oversized stride must return InvalidOperation"
7071 );
7072 }
7073
7074 #[test]
7075 fn dma_tensor_new_with_byte_size_rejects_shape_overflow() {
7076 // Round-3 PR feedback (C3): shape.product() * sizeof(T) must use
7077 // checked arithmetic so a pathological shape can't wrap usize and
7078 // make the byte_size-vs-logical-size comparison incorrect.
7079 //
7080 // This test only exercises the overflow rejection path, which is
7081 // pure-Rust and doesn't touch dma_heap — safe to run on any target.
7082 #[cfg(target_os = "linux")]
7083 {
7084 let err = crate::dma::DmaTensor::<u64>::new_with_byte_size(
7085 &[usize::MAX, 2, 2],
7086 usize::MAX,
7087 None,
7088 false,
7089 );
7090 assert!(
7091 matches!(err, Err(Error::InvalidArgument(_))),
7092 "new_with_byte_size must detect shape.product() overflow"
7093 );
7094 }
7095 }
7096
7097 #[test]
7098 #[cfg(target_os = "linux")]
7099 fn image_tensor_with_stride_rejects_too_small_stride() {
7100 // Declares this test as an fd-opener; see FD_LOCK.
7101 let _lock = crate::tests::fd_lock_shared();
7102 // 640×480 RGBA8 natural pitch = 2560, request 2400 → should error.
7103 let err = Tensor::<u8>::image_with_stride(
7104 640,
7105 480,
7106 PixelFormat::Rgba,
7107 2400,
7108 Some(TensorMemory::DmaBuf),
7109 crate::CpuAccess::ReadWrite,
7110 );
7111 assert!(matches!(err, Err(Error::InvalidArgument(_))));
7112 }
7113
7114 #[test]
7115 #[cfg(target_os = "linux")]
7116 fn image_tensor_with_stride_rejects_non_packed() {
7117 // Declares this test as an fd-opener; see FD_LOCK.
7118 let _lock = crate::tests::fd_lock_shared();
7119 // NV12 is SemiPlanar → not supported. (Linux-only because
7120 // `TensorMemory::DmaBuf` itself is a Linux-only enum variant.)
7121 let err = Tensor::<u8>::image_with_stride(
7122 640,
7123 480,
7124 PixelFormat::Nv12,
7125 640,
7126 Some(TensorMemory::DmaBuf),
7127 crate::CpuAccess::ReadWrite,
7128 );
7129 assert!(matches!(err, Err(Error::NotImplemented(_))));
7130 }
7131
7132 #[test]
7133 fn set_format_valid() {
7134 let mut t = Tensor::<u8>::new(&[480, 640, 3], None, None).unwrap();
7135 assert!(t.format().is_none());
7136 t.set_format(PixelFormat::Rgb).unwrap();
7137 assert_eq!(t.format(), Some(PixelFormat::Rgb));
7138 assert_eq!(t.width(), Some(640));
7139 assert_eq!(t.height(), Some(480));
7140 }
7141
7142 #[test]
7143 fn set_format_invalid_shape() {
7144 let mut t = Tensor::<u8>::new(&[480, 640, 4], None, None).unwrap();
7145 // RGB expects 3 channels, not 4
7146 let err = t.set_format(PixelFormat::Rgb);
7147 assert!(err.is_err());
7148 // Original tensor is unmodified
7149 assert!(t.format().is_none());
7150 }
7151
7152 #[test]
7153 fn reshape_clears_format() {
7154 let mut t = Tensor::<u8>::image(
7155 640,
7156 480,
7157 PixelFormat::Rgba,
7158 None,
7159 crate::CpuAccess::ReadWrite,
7160 )
7161 .unwrap();
7162 assert_eq!(t.format(), Some(PixelFormat::Rgba));
7163 // Reshape to flat — format cleared
7164 t.reshape(&[480 * 640 * 4]).unwrap();
7165 assert!(t.format().is_none());
7166 }
7167
7168 #[test]
7169 fn from_planes_nv12() {
7170 let y = Tensor::<u8>::new(&[480, 640], None, None).unwrap();
7171 let uv = Tensor::<u8>::new(&[240, 640], None, None).unwrap();
7172 let img = Tensor::from_planes(y, uv, PixelFormat::Nv12).unwrap();
7173 assert_eq!(img.format(), Some(PixelFormat::Nv12));
7174 assert!(img.is_multiplane());
7175 assert!(img.chroma().is_some());
7176 assert_eq!(img.width(), Some(640));
7177 assert_eq!(img.height(), Some(480));
7178 }
7179
7180 #[test]
7181 fn from_planes_rejects_non_semiplanar() {
7182 let y = Tensor::<u8>::new(&[480, 640], None, None).unwrap();
7183 let uv = Tensor::<u8>::new(&[240, 640], None, None).unwrap();
7184 let err = Tensor::from_planes(y, uv, PixelFormat::Rgb);
7185 assert!(err.is_err());
7186 }
7187
7188 #[test]
7189 fn reshape_multiplane_errors() {
7190 let y = Tensor::<u8>::new(&[480, 640], None, None).unwrap();
7191 let uv = Tensor::<u8>::new(&[240, 640], None, None).unwrap();
7192 let mut img = Tensor::from_planes(y, uv, PixelFormat::Nv12).unwrap();
7193 let err = img.reshape(&[480 * 640 + 240 * 640]);
7194 assert!(err.is_err());
7195 }
7196}
7197
7198#[cfg(test)]
7199mod compression_tests {
7200 use super::*;
7201
7202 #[test]
7203 fn desc_builder_roundtrips() {
7204 let desc = ImageDesc::new(640, 480, PixelFormat::Rgba, DType::U8)
7205 .with_memory(Some(TensorMemory::Mem))
7206 .with_access(CpuAccess::Read)
7207 .with_compression(Compression::Any);
7208 assert_eq!(desc.width(), 640);
7209 assert_eq!(desc.height(), 480);
7210 assert_eq!(desc.format(), PixelFormat::Rgba);
7211 assert_eq!(desc.dtype(), DType::U8);
7212 assert_eq!(desc.memory(), Some(TensorMemory::Mem));
7213 assert_eq!(desc.access(), CpuAccess::Read);
7214 assert_eq!(desc.compression(), Some(Compression::Any));
7215
7216 // Defaults: auto memory, CPU read/write, no request.
7217 let plain = ImageDesc::new(2, 2, PixelFormat::Grey, DType::U8);
7218 assert_eq!(plain.memory(), None);
7219 assert_eq!(plain.access(), CpuAccess::ReadWrite);
7220 assert_eq!(plain.compression(), None);
7221 }
7222
7223 #[test]
7224 fn desc_dtype_must_match_element_type() {
7225 let desc = ImageDesc::new(4, 4, PixelFormat::Rgba, DType::F32);
7226 match Tensor::<u8>::image_desc(&desc) {
7227 Err(Error::InvalidArgument(msg)) => assert!(msg.contains("dtype")),
7228 other => panic!("expected InvalidArgument, got {other:?}"),
7229 }
7230 }
7231
7232 #[test]
7233 fn compression_with_cpu_access_is_invalid() {
7234 let desc = ImageDesc::new(4, 4, PixelFormat::Rgba, DType::U8)
7235 .with_access(CpuAccess::ReadWrite)
7236 .with_compression(Compression::Any);
7237 match Tensor::<u8>::image_desc(&desc) {
7238 Err(Error::InvalidArgument(msg)) => assert!(msg.contains("CpuAccess::None")),
7239 other => panic!("expected InvalidArgument, got {other:?}"),
7240 }
7241 }
7242
7243 #[cfg(not(target_os = "android"))]
7244 #[test]
7245 fn scheme_request_off_android_is_not_implemented() {
7246 let desc = ImageDesc::new(4, 4, PixelFormat::Rgba, DType::U8)
7247 .with_access(CpuAccess::None)
7248 .with_compression(Compression::Scheme(CompressionScheme::Ubwc));
7249 match Tensor::<u8>::image_desc(&desc) {
7250 Err(Error::NotImplemented(msg)) => assert!(msg.contains("Ubwc")),
7251 other => panic!("expected NotImplemented, got {other:?}"),
7252 }
7253 }
7254
7255 #[cfg(not(target_os = "android"))]
7256 #[test]
7257 fn any_request_off_android_resolves_linear_and_counts() {
7258 let before = compression_fallback_count();
7259 let desc = ImageDesc::new(64, 64, PixelFormat::Rgba, DType::U8)
7260 .with_memory(Some(TensorMemory::Mem))
7261 .with_access(CpuAccess::None)
7262 .with_compression(Compression::Any);
7263 let t = Tensor::<u8>::image_desc(&desc).unwrap();
7264 assert_eq!(t.compression(), None);
7265 assert!(compression_fallback_count() > before);
7266 }
7267
7268 #[test]
7269 fn desc_without_request_matches_classic_constructor() {
7270 let desc = ImageDesc::new(32, 32, PixelFormat::Rgba, DType::U8)
7271 .with_memory(Some(TensorMemory::Mem))
7272 .with_access(CpuAccess::ReadWrite);
7273 let t = Tensor::<u8>::image_desc(&desc).unwrap();
7274 assert_eq!(t.compression(), None);
7275 assert_eq!(t.cpu_access(), CpuAccess::ReadWrite);
7276 assert_eq!(t.width(), Some(32));
7277 // Mappable exactly like the classic constructor's result.
7278 let m = t.map_read().unwrap();
7279 assert_eq!(m.as_slice().len(), 32 * 32 * 4);
7280 }
7281
7282 #[test]
7283 fn configure_image_preserves_compression_and_views_inherit() {
7284 // No host platform records a scheme, so emulate the recording to
7285 // pin the preserve/inherit semantics (the physical layout does
7286 // not change when the logical image is reconfigured).
7287 let mut t = Tensor::<u8>::image(
7288 64,
7289 64,
7290 PixelFormat::Rgba,
7291 Some(TensorMemory::Mem),
7292 CpuAccess::None,
7293 )
7294 .unwrap();
7295 t.compression = Some(CompressionScheme::Ubwc);
7296 t.configure_image(32, 32, PixelFormat::Rgba).unwrap();
7297 assert_eq!(t.compression(), Some(CompressionScheme::Ubwc));
7298 let view = t.subview(0, &[16, 32, 4]).unwrap();
7299 assert_eq!(view.compression(), Some(CompressionScheme::Ubwc));
7300 }
7301
7302 #[test]
7303 fn tensor_dyn_dispatches_desc_and_compression() {
7304 let desc = ImageDesc::new(16, 16, PixelFormat::Rgba, DType::U8)
7305 .with_memory(Some(TensorMemory::Mem))
7306 .with_access(CpuAccess::ReadWrite);
7307 let t = TensorDyn::image_desc(&desc).unwrap();
7308 assert_eq!(t.compression(), None);
7309 assert!(matches!(t, TensorDyn::U8(_)));
7310 }
7311}
7312
7313#[cfg(all(test, not(edgefirst_tensor_dynamic)))]
7314mod vocabulary_and_probe_tests {
7315 use super::*;
7316
7317 #[test]
7318 fn dtype_displays_its_wire_name() {
7319 for (dt, name) in [
7320 (DType::U8, "u8"),
7321 (DType::I8, "i8"),
7322 (DType::U16, "u16"),
7323 (DType::I16, "i16"),
7324 (DType::U32, "u32"),
7325 (DType::I32, "i32"),
7326 (DType::U64, "u64"),
7327 (DType::I64, "i64"),
7328 (DType::F16, "f16"),
7329 (DType::F32, "f32"),
7330 (DType::F64, "f64"),
7331 ] {
7332 assert_eq!(dt.to_string(), name);
7333 }
7334 }
7335
7336 #[test]
7337 fn dtype_of_maps_every_element_type_and_nothing_else() {
7338 assert_eq!(dtype_of::<u8>(), Some(DType::U8));
7339 assert_eq!(dtype_of::<i8>(), Some(DType::I8));
7340 assert_eq!(dtype_of::<u16>(), Some(DType::U16));
7341 assert_eq!(dtype_of::<i16>(), Some(DType::I16));
7342 assert_eq!(dtype_of::<u32>(), Some(DType::U32));
7343 assert_eq!(dtype_of::<i32>(), Some(DType::I32));
7344 assert_eq!(dtype_of::<u64>(), Some(DType::U64));
7345 assert_eq!(dtype_of::<i64>(), Some(DType::I64));
7346 assert_eq!(dtype_of::<half::f16>(), Some(DType::F16));
7347 assert_eq!(dtype_of::<f32>(), Some(DType::F32));
7348 assert_eq!(dtype_of::<f64>(), Some(DType::F64));
7349 assert_eq!(dtype_of::<bool>(), None);
7350 }
7351
7352 #[test]
7353 fn tensor_memory_converts_to_and_from_its_wire_string() {
7354 assert_eq!(String::from(TensorMemory::DmaBuf), "dmabuf");
7355 assert_eq!(String::from(TensorMemory::Shm), "shm");
7356 assert_eq!(
7357 TensorMemory::try_from("dmabuf").unwrap(),
7358 TensorMemory::DmaBuf
7359 );
7360 assert_eq!(
7361 TensorMemory::try_from("dma").unwrap(),
7362 TensorMemory::DmaBuf,
7363 "the legacy spelling still parses"
7364 );
7365 assert_eq!(TensorMemory::try_from("mem").unwrap(), TensorMemory::Mem);
7366 match TensorMemory::try_from("bogus") {
7367 Err(Error::InvalidMemoryType(s)) => assert_eq!(s, "bogus"),
7368 other => panic!("expected InvalidMemoryType, got {other:?}"),
7369 }
7370 }
7371
7372 #[test]
7373 fn tensor_memory_is_available_matches_a_real_allocation() {
7374 let _lock = crate::tests::fd_lock_shared();
7375 assert!(TensorMemory::Mem.is_available());
7376 assert!(!TensorMemory::IoSurface.is_available());
7377 assert!(!TensorMemory::Pbo.is_available());
7378 assert!(!TensorMemory::Cuda.is_available());
7379 assert!(!TensorMemory::D3d11Texture.is_available());
7380 #[cfg(unix)]
7381 assert_eq!(
7382 TensorMemory::Shm.is_available(),
7383 Tensor::<u8>::new(&[64], Some(TensorMemory::Shm), None).is_ok()
7384 );
7385 #[cfg(not(target_os = "windows"))]
7386 assert_eq!(
7387 TensorMemory::DmaBuf.is_available(),
7388 Tensor::<u8>::new(&[64], Some(TensorMemory::DmaBuf), None).is_ok()
7389 );
7390 }
7391
7392 #[test]
7393 #[cfg(target_os = "linux")]
7394 fn linux_probes_report_what_linux_can_allocate() {
7395 let _lock = crate::tests::fd_lock_shared();
7396 let dma = DmaTensor::<u8>::new(&[64], None).is_ok();
7397 assert_eq!(is_dma_available(), dma);
7398 assert_eq!(is_gpu_buffer_available(), dma);
7399 assert!(!is_iosurface_available());
7400 assert!(!is_ahardwarebuffer_available());
7401 // Linux (unlike Android's bionic) always has POSIX shared memory.
7402 assert!(is_shm_available());
7403 }
7404}
7405
7406#[cfg(all(test, not(edgefirst_tensor_dynamic)))]
7407mod quantization_serde_tests {
7408 use super::*;
7409
7410 fn parse(json: &str) -> Quantization {
7411 serde_json::from_str(json).unwrap_or_else(|e| panic!("{json}: {e}"))
7412 }
7413
7414 #[test]
7415 fn scale_accepts_every_json_number_form() {
7416 assert_eq!(parse(r#"{"scale": 0.5}"#).scale(), &[0.5]);
7417 assert_eq!(parse(r#"{"scale": 2}"#).scale(), &[2.0]);
7418 assert_eq!(parse(r#"{"scale": -2}"#).scale(), &[-2.0]);
7419 assert_eq!(parse(r#"{"scale": [0.25, 4]}"#).scale(), &[0.25, 4.0]);
7420 }
7421
7422 #[test]
7423 fn zero_point_accepts_null_scalar_and_array() {
7424 assert_eq!(
7425 parse(r#"{"scale": 0.1, "zero_point": -5}"#).zero_point(),
7426 Some(&[-5][..])
7427 );
7428 assert_eq!(
7429 parse(r#"{"scale": 0.1, "zero_point": 3}"#).zero_point(),
7430 Some(&[3][..])
7431 );
7432 assert_eq!(
7433 parse(r#"{"scale": [0.1, 0.2], "zero_point": [1, -2], "axis": 0}"#).zero_point(),
7434 Some(&[1, -2][..])
7435 );
7436 assert_eq!(
7437 parse(r#"{"scale": 0.1, "zero_point": null}"#).zero_point(),
7438 None
7439 );
7440 assert_eq!(parse(r#"{"scale": 0.1}"#).zero_point(), None);
7441 }
7442
7443 /// Deserializers whose `deserialize_option` forwards to
7444 /// `deserialize_any` (serde's value deserializers among them) reach the
7445 /// zero-point visitor's scalar and sequence arms directly, not through
7446 /// `visit_some`.
7447 #[test]
7448 fn zero_point_visitor_accepts_values_without_an_option_wrapper() {
7449 use serde::de::value::{
7450 Error as ValueError, I64Deserializer, SeqDeserializer, U64Deserializer,
7451 };
7452 assert_eq!(
7453 deserialize_opt_scalar_or_vec_i32(I64Deserializer::<ValueError>::new(-7)).unwrap(),
7454 Some(vec![-7])
7455 );
7456 assert_eq!(
7457 deserialize_opt_scalar_or_vec_i32(U64Deserializer::<ValueError>::new(9)).unwrap(),
7458 Some(vec![9])
7459 );
7460 let seq = SeqDeserializer::<_, ValueError>::new(vec![4i32, -4].into_iter());
7461 assert_eq!(
7462 deserialize_opt_scalar_or_vec_i32(seq).unwrap(),
7463 Some(vec![4, -4])
7464 );
7465 }
7466
7467 #[test]
7468 fn per_channel_is_not_per_tensor() {
7469 let q = Quantization::per_channel(vec![0.1, 0.2], vec![0, 1], 0).unwrap();
7470 assert!(!q.is_per_tensor());
7471 assert!(Quantization::per_tensor(0.1, 0).is_per_tensor());
7472 }
7473
7474 #[test]
7475 fn validate_accepts_matching_per_channel_and_names_the_failing_field() {
7476 let mut t = Tensor::<i8>::new(&[1, 1, 4], Some(TensorMemory::Mem), None).unwrap();
7477 let q = Quantization::per_channel(vec![0.1, 0.2, 0.3, 0.4], vec![0, 1, 2, 3], 2).unwrap();
7478 t.set_quantization(q)
7479 .expect("four channels on axis 2 of [1, 1, 4]");
7480 assert_eq!(
7481 t.quantization().and_then(Quantization::zero_point),
7482 Some(&[0, 1, 2, 3][..])
7483 );
7484
7485 let q = Quantization::per_channel(vec![0.1, 0.2, 0.3, 0.4], vec![0, 1, 2, 3], 3).unwrap();
7486 match t.set_quantization(q) {
7487 Err(Error::QuantizationInvalid { field, got, .. }) => {
7488 assert_eq!(field, "axis");
7489 assert_eq!(got, "axis=3");
7490 }
7491 other => panic!("axis 3 of a rank-3 tensor must be refused, got {other:?}"),
7492 }
7493
7494 let q = Quantization::per_channel_symmetric(vec![0.1, 0.2, 0.3, 0.4], 1).unwrap();
7495 match t.set_quantization(q) {
7496 Err(Error::QuantizationInvalid { field, got, .. }) => {
7497 assert_eq!(field, "scale.len");
7498 assert_eq!(got, "length 4");
7499 }
7500 other => panic!("four scales on a length-1 axis must be refused, got {other:?}"),
7501 }
7502 }
7503}
7504
7505#[cfg(all(test, not(edgefirst_tensor_dynamic)))]
7506mod image_geometry_tests {
7507 use super::*;
7508
7509 fn mem_image<T>(width: usize, height: usize, format: PixelFormat) -> Tensor<T>
7510 where
7511 T: Num + Clone + fmt::Debug + Send + Sync + 'static,
7512 {
7513 Tensor::<T>::image(
7514 width,
7515 height,
7516 format,
7517 Some(TensorMemory::Mem),
7518 CpuAccess::ReadWrite,
7519 )
7520 .unwrap_or_else(|e| panic!("Mem {format:?} {width}x{height}: {e}"))
7521 }
7522
7523 fn raw_mem<T>(shape: &[usize]) -> Tensor<T>
7524 where
7525 T: Num + Clone + fmt::Debug + Send + Sync + 'static,
7526 {
7527 Tensor::<T>::new(shape, Some(TensorMemory::Mem), None).unwrap()
7528 }
7529
7530 #[test]
7531 fn semi_planar_image_stride_scales_with_the_element_size() {
7532 // u16 NV12 at width 100: 200 B rows, aligned to 256.
7533 let t = mem_image::<u16>(100, 64, PixelFormat::Nv12);
7534 assert_eq!(t.row_stride(), Some(256));
7535 assert_eq!(t.capacity_bytes(), 256 * 96);
7536 }
7537
7538 #[test]
7539 fn host_planar_image_is_tight_at_every_element_size() {
7540 let t = mem_image::<u8>(320, 240, PixelFormat::PlanarRgb);
7541 assert_eq!(t.row_stride(), None);
7542 assert_eq!(t.effective_row_stride(), Some(320));
7543 assert_eq!(t.capacity_bytes(), 320 * 240 * 3);
7544
7545 let t = mem_image::<f32>(100, 40, PixelFormat::PlanarRgb);
7546 assert_eq!(t.row_stride(), None);
7547 assert_eq!(t.effective_row_stride(), Some(400));
7548 assert_eq!(t.capacity_bytes(), 400 * 120);
7549 }
7550
7551 #[test]
7552 #[cfg(target_os = "linux")]
7553 fn dma_packed_image_records_the_aligned_pitch_only_when_padded() {
7554 let _lock = crate::tests::fd_lock_shared();
7555 if !crate::test_support::dma_or_skip(
7556 "dma_packed_image_records_the_aligned_pitch_only_when_padded",
7557 ) {
7558 return;
7559 }
7560 let t = Tensor::<u8>::image(
7561 100,
7562 10,
7563 PixelFormat::Rgba,
7564 Some(TensorMemory::DmaBuf),
7565 CpuAccess::ReadWrite,
7566 )
7567 .unwrap();
7568 assert_eq!(t.row_stride(), Some(448), "400 B rows pad to 448");
7569 let t = Tensor::<u8>::image(
7570 64,
7571 10,
7572 PixelFormat::Rgba,
7573 Some(TensorMemory::DmaBuf),
7574 CpuAccess::ReadWrite,
7575 )
7576 .unwrap();
7577 assert_eq!(t.row_stride(), None, "256 B rows are already aligned");
7578 }
7579
7580 #[test]
7581 #[cfg(target_os = "linux")]
7582 fn image_with_stride_validates_the_stride_before_the_memory() {
7583 match Tensor::<u8>::image_with_stride(
7584 640,
7585 480,
7586 PixelFormat::Rgba,
7587 2559,
7588 Some(TensorMemory::Mem),
7589 CpuAccess::ReadWrite,
7590 ) {
7591 Err(Error::InvalidArgument(msg)) => {
7592 assert!(msg.contains("2559 < minimum 2560"), "{msg}")
7593 }
7594 other => panic!("expected InvalidArgument, got {other:?}"),
7595 }
7596 match Tensor::<u8>::image_with_stride(
7597 640,
7598 480,
7599 PixelFormat::Rgba,
7600 2560,
7601 Some(TensorMemory::Mem),
7602 CpuAccess::ReadWrite,
7603 ) {
7604 Err(Error::NotImplemented(msg)) => assert!(msg.contains("only TensorMemory::DmaBuf")),
7605 other => panic!("a valid stride on Mem must reach the memory check, got {other:?}"),
7606 }
7607 }
7608
7609 #[test]
7610 fn set_format_checks_the_planar_and_semi_planar_shape() {
7611 let mut t = raw_mem::<u8>(&[480, 640, 3]);
7612 match t.set_format(PixelFormat::PlanarRgb) {
7613 Err(Error::InvalidShape(msg)) => assert!(msg.contains("[3, H, W]"), "{msg}"),
7614 other => panic!("[480, 640, 3] is not planar, got {other:?}"),
7615 }
7616 assert!(raw_mem::<u8>(&[3, 480, 640])
7617 .set_format(PixelFormat::PlanarRgb)
7618 .is_ok());
7619
7620 for (format, bad, good) in [
7621 (PixelFormat::Nv12, 721, [720usize, 722]),
7622 (PixelFormat::Nv16, 481, [480, 482]),
7623 (PixelFormat::Nv24, 1441, [1440, 1443]),
7624 ] {
7625 match raw_mem::<u8>(&[bad, 640]).set_format(format) {
7626 Err(Error::InvalidShape(msg)) => {
7627 assert!(msg.contains(&bad.to_string()), "{format:?}: {msg}")
7628 }
7629 other => panic!("{format:?} [{bad}, 640] must be refused, got {other:?}"),
7630 }
7631 for rows in good {
7632 raw_mem::<u8>(&[rows, 640])
7633 .set_format(format)
7634 .unwrap_or_else(|e| panic!("{format:?} [{rows}, 640]: {e}"));
7635 }
7636 }
7637 }
7638
7639 #[test]
7640 fn contiguous_semi_planar_height_divides_the_combined_rows() {
7641 assert_eq!(
7642 mem_image::<u8>(320, 240, PixelFormat::Nv24).height(),
7643 Some(240)
7644 );
7645 assert_eq!(
7646 mem_image::<u8>(320, 240, PixelFormat::Nv16).height(),
7647 Some(240)
7648 );
7649 assert_eq!(
7650 mem_image::<u8>(320, 241, PixelFormat::Nv12).height(),
7651 Some(241)
7652 );
7653 }
7654
7655 #[test]
7656 fn effective_row_stride_scales_with_the_element_size() {
7657 assert_eq!(
7658 mem_image::<f32>(10, 4, PixelFormat::Rgb).effective_row_stride(),
7659 Some(120)
7660 );
7661 assert_eq!(
7662 mem_image::<f32>(10, 4, PixelFormat::PlanarRgb).effective_row_stride(),
7663 Some(40)
7664 );
7665 let mut t = raw_mem::<u16>(&[360, 320]);
7666 t.set_format(PixelFormat::Nv12).unwrap();
7667 assert_eq!(t.effective_row_stride(), Some(640));
7668 let mut t = raw_mem::<u16>(&[360, 321]);
7669 t.set_format(PixelFormat::Nv12).unwrap();
7670 assert_eq!(
7671 t.effective_row_stride(),
7672 Some(644),
7673 "odd width rounds up to even"
7674 );
7675 }
7676
7677 #[test]
7678 fn set_row_stride_minimum_scales_with_the_element_size() {
7679 let refuse = |mut t: Tensor<f32>, stride: usize, min: usize| match t.set_row_stride(stride)
7680 {
7681 Err(Error::InvalidArgument(msg)) => {
7682 assert!(msg.contains(&format!("{stride} < minimum {min}")), "{msg}")
7683 }
7684 other => panic!("stride {stride} under {min} must be refused, got {other:?}"),
7685 };
7686 refuse(mem_image::<f32>(10, 4, PixelFormat::Rgb), 119, 120);
7687 refuse(mem_image::<f32>(10, 4, PixelFormat::PlanarRgb), 39, 40);
7688 let mut t = mem_image::<f32>(10, 4, PixelFormat::Rgb);
7689 t.set_row_stride(120).unwrap();
7690 assert_eq!(t.row_stride(), Some(120));
7691 let mut t = mem_image::<f32>(10, 4, PixelFormat::PlanarRgb);
7692 t.set_row_stride(40).unwrap();
7693 assert_eq!(t.row_stride(), Some(40));
7694
7695 let mut t = mem_image::<u16>(101, 64, PixelFormat::Nv12);
7696 match t.set_row_stride(203) {
7697 Err(Error::InvalidArgument(msg)) => assert!(msg.contains("203 < minimum 204"), "{msg}"),
7698 other => panic!("expected InvalidArgument, got {other:?}"),
7699 }
7700 t.set_row_stride(204).unwrap();
7701 }
7702
7703 #[test]
7704 fn configure_image_semi_planar_stride_scales_with_the_element_size() {
7705 let mut t = mem_image::<u16>(640, 480, PixelFormat::Rgba);
7706 t.configure_image(100, 64, PixelFormat::Nv12).unwrap();
7707 assert_eq!(t.row_stride(), Some(256));
7708 }
7709
7710 #[test]
7711 fn configure_image_keeps_host_packed_odd_widths_tight() {
7712 let mut t = mem_image::<u8>(640, 480, PixelFormat::Rgb);
7713 t.configure_image(321, 240, PixelFormat::Rgb).unwrap();
7714 assert_eq!(t.row_stride(), None);
7715 assert_eq!(t.effective_row_stride(), Some(963));
7716 }
7717
7718 #[test]
7719 fn configure_image_does_not_inherit_a_stride_under_the_same_format() {
7720 // Host packed: a stride set for one geometry must not survive into
7721 // the next, where it is below the new minimum row.
7722 let mut t = mem_image::<u8>(640, 480, PixelFormat::Rgb);
7723 t.configure_image(100, 100, PixelFormat::Rgb).unwrap();
7724 t.set_row_stride(400).unwrap();
7725 t.configure_image(300, 100, PixelFormat::Rgb).unwrap();
7726 assert_eq!(t.row_stride(), None);
7727 assert_eq!(t.effective_row_stride(), Some(900));
7728
7729 // Semi-planar whose aligned stride does not fit: the tight row does,
7730 // and must be what the tensor reports -- not the old 64.
7731 let mut t = mem_image::<u8>(64, 1000, PixelFormat::Nv12);
7732 assert_eq!(t.row_stride(), Some(64));
7733 t.configure_image(200, 300, PixelFormat::Nv12).unwrap();
7734 assert_eq!(t.row_stride(), None);
7735 assert_eq!(t.effective_row_stride(), Some(200));
7736 }
7737
7738 #[test]
7739 fn configure_image_pool_reuse_picks_the_stride_that_fits() {
7740 // A prior stride narrower than the new row is replaced by the
7741 // aligned one.
7742 let mut t = mem_image::<u8>(64, 1024, PixelFormat::Nv12);
7743 t.configure_image(200, 100, PixelFormat::Nv12).unwrap();
7744 assert_eq!(t.row_stride(), Some(256));
7745
7746 // An unaligned prior stride is replaced too.
7747 let mut t = mem_image::<u8>(64, 64, PixelFormat::Nv12);
7748 t.set_row_stride(100).unwrap();
7749 t.configure_image(64, 64, PixelFormat::Nv12).unwrap();
7750 assert_eq!(t.row_stride(), Some(64));
7751
7752 // A wide prior stride whose rows would overrun the allocation gives
7753 // way to the aligned stride of the new width.
7754 let mut t = mem_image::<u8>(640, 480, PixelFormat::Nv12);
7755 t.configure_image(320, 900, PixelFormat::Nv12).unwrap();
7756 assert_eq!(t.row_stride(), Some(320));
7757
7758 // Neither the aligned stride fits nor is it needed: the tight row
7759 // exactly fills a raw allocation.
7760 let mut t = raw_mem::<u8>(&[150 * 100]);
7761 t.configure_image(100, 100, PixelFormat::Nv12).unwrap();
7762 assert_eq!(t.row_stride(), None);
7763 assert_eq!(t.effective_row_stride(), Some(100));
7764 }
7765
7766 #[test]
7767 fn configure_image_refuses_rows_that_overrun_the_allocation() {
7768 // Odd width 99: the shape [150, 99] fits 14850 bytes, but NV12 rows
7769 // are at least the even width, 100 * 150 = 15000 bytes.
7770 let mut t = raw_mem::<u8>(&[150 * 99]);
7771 match t.configure_image(99, 100, PixelFormat::Nv12) {
7772 Err(Error::InsufficientCapacity { needed, capacity }) => {
7773 assert_eq!((needed, capacity), (15000, 14850))
7774 }
7775 other => panic!("expected InsufficientCapacity, got {other:?}"),
7776 }
7777 }
7778
7779 #[test]
7780 #[cfg(target_os = "linux")]
7781 fn configure_image_on_dma_aligns_packed_and_planar_rows() {
7782 let _lock = crate::tests::fd_lock_shared();
7783 if !crate::test_support::dma_or_skip("configure_image_on_dma_aligns_packed_and_planar_rows")
7784 {
7785 return;
7786 }
7787 // Packed u16 Rgb at width 100: 600 B rows align to 640.
7788 let mut t = Tensor::<u16>::image(
7789 640,
7790 480,
7791 PixelFormat::Rgba,
7792 Some(TensorMemory::DmaBuf),
7793 CpuAccess::ReadWrite,
7794 )
7795 .unwrap();
7796 t.configure_image(100, 50, PixelFormat::Rgb).unwrap();
7797 assert_eq!(t.row_stride(), Some(640));
7798
7799 // Planar u16 at width 100: 200 B rows would align to 256, but 150
7800 // rows of 256 B overrun the 32 KiB buffer while 150 rows of 200 B fit.
7801 let mut t = Tensor::<u16>::new(&[16 * 1024], Some(TensorMemory::DmaBuf), None).unwrap();
7802 assert_eq!(t.capacity_bytes(), 32 * 1024);
7803 t.configure_image(100, 50, PixelFormat::PlanarRgb).unwrap();
7804 assert_eq!(t.row_stride(), None);
7805 assert_eq!(t.effective_row_stride(), Some(200));
7806 t.configure_image(100, 40, PixelFormat::PlanarRgb).unwrap();
7807 assert_eq!(t.row_stride(), Some(256), "120 rows of 256 B fit");
7808 }
7809
7810 #[test]
7811 fn from_planes_requires_a_raw_two_dimensional_chroma() {
7812 let formatted_chroma = mem_image::<u8>(640, 240, PixelFormat::Grey);
7813 match Tensor::from_planes(
7814 raw_mem::<u8>(&[480, 640]),
7815 formatted_chroma,
7816 PixelFormat::Nv12,
7817 ) {
7818 Err(Error::InvalidArgument(msg)) => assert!(msg.contains("raw tensor"), "{msg}"),
7819 other => panic!("a formatted chroma plane must be refused, got {other:?}"),
7820 }
7821 match Tensor::from_planes(
7822 raw_mem::<u8>(&[480, 640, 1]),
7823 raw_mem::<u8>(&[240, 640]),
7824 PixelFormat::Nv12,
7825 ) {
7826 Err(Error::InvalidArgument(msg)) => assert!(msg.contains("2D shapes"), "{msg}"),
7827 other => panic!("a 3-D luma plane must be refused, got {other:?}"),
7828 }
7829 }
7830
7831 #[test]
7832 fn from_planes_nv16_needs_full_height_chroma() {
7833 let img = Tensor::from_planes(
7834 raw_mem::<u8>(&[480, 640]),
7835 raw_mem::<u8>(&[480, 640]),
7836 PixelFormat::Nv16,
7837 )
7838 .expect("NV16 chroma is full height");
7839 assert_eq!(img.format(), Some(PixelFormat::Nv16));
7840 assert_eq!(img.height(), Some(480));
7841
7842 match Tensor::from_planes(
7843 raw_mem::<u8>(&[480, 640]),
7844 raw_mem::<u8>(&[240, 640]),
7845 PixelFormat::Nv16,
7846 ) {
7847 Err(Error::InvalidArgument(msg)) => {
7848 assert!(msg.contains("NV16 chroma height"), "{msg}")
7849 }
7850 other => panic!("half-height NV16 chroma must be refused, got {other:?}"),
7851 }
7852 }
7853
7854 #[test]
7855 fn chroma_mut_writes_the_chroma_plane() {
7856 let mut img = Tensor::from_planes(
7857 raw_mem::<u8>(&[4, 8]),
7858 raw_mem::<u8>(&[2, 8]),
7859 PixelFormat::Nv12,
7860 )
7861 .unwrap();
7862 let chroma = img.chroma_mut().expect("a multiplane image has chroma");
7863 chroma.map().unwrap().as_mut_slice().fill(0x80);
7864 let chroma = img.chroma().unwrap();
7865 assert!(chroma.map().unwrap().as_slice().iter().all(|&b| b == 0x80));
7866 }
7867
7868 #[test]
7869 fn view_offset_and_stride_scale_with_the_element_size() {
7870 let t = mem_image::<u8>(100, 80, PixelFormat::Rgba);
7871 let v = t.view(Region::new(10, 20, 30, 40)).unwrap();
7872 assert_eq!(v.plane_offset(), Some(20 * 400 + 10 * 4));
7873 assert_eq!(v.row_stride(), Some(400));
7874
7875 let t = mem_image::<u16>(100, 80, PixelFormat::Rgba);
7876 let v = t.view(Region::new(10, 20, 30, 40)).unwrap();
7877 assert_eq!(v.plane_offset(), Some(20 * 800 + 10 * 8));
7878 assert_eq!(v.row_stride(), Some(800));
7879 assert_eq!(v.view_origin().map(|o| o.parent_row_stride), Some(800));
7880 }
7881
7882 #[test]
7883 fn single_row_view_keeps_a_tight_stride() {
7884 let t = mem_image::<u8>(8, 4, PixelFormat::Rgba);
7885 let row = t.view(Region::new(2, 3, 4, 1)).unwrap();
7886 assert_eq!(row.effective_row_stride(), Some(16));
7887 assert_eq!(row.map().unwrap().as_slice().len(), 16);
7888 assert_eq!(row.view_origin().map(|o| o.parent_row_stride), Some(32));
7889
7890 let two_rows = t.view(Region::new(2, 2, 4, 2)).unwrap();
7891 assert_eq!(two_rows.effective_row_stride(), Some(32));
7892 }
7893
7894 #[test]
7895 fn view_refuses_a_region_past_either_edge() {
7896 let t = mem_image::<u8>(64, 40, PixelFormat::Rgba);
7897 for region in [Region::new(0, 30, 10, 20), Region::new(60, 0, 10, 1)] {
7898 match t.view(region) {
7899 Err(Error::RegionOutOfBounds { region: r, bounds }) => {
7900 assert_eq!((r, bounds), (region, (64, 40)))
7901 }
7902 other => panic!("{region:?} must be refused, got {other:?}"),
7903 }
7904 }
7905 assert!(Region::new(0, 0, 64, 40).fits_within(64, 40));
7906 assert!(!Region::new(1, 0, 64, 40).fits_within(64, 40));
7907 assert!(!Region::new(0, 1, 64, 40).fits_within(64, 40));
7908 }
7909}
7910
7911#[cfg(all(test, not(edgefirst_tensor_dynamic)))]
7912mod strided_access_tests {
7913 use super::*;
7914
7915 /// A patterned byte view of `t`'s whole mapping: byte `i` is `i % 251`.
7916 fn fill_pattern<T>(t: &Tensor<T>)
7917 where
7918 T: Num + Clone + fmt::Debug + Send + Sync,
7919 {
7920 let mut m = t.map().unwrap();
7921 let s = m.as_mut_slice();
7922 // SAFETY: plain numeric elements; every byte pattern is a valid value
7923 // for the integer and float types these tests use.
7924 let bytes = unsafe {
7925 std::slice::from_raw_parts_mut(s.as_mut_ptr() as *mut u8, std::mem::size_of_val(s))
7926 };
7927 for (i, b) in bytes.iter_mut().enumerate() {
7928 *b = (i % 251) as u8;
7929 }
7930 }
7931
7932 fn expected_rows(rows: usize, stride: usize, row_bytes: usize) -> Vec<u8> {
7933 (0..rows)
7934 .flat_map(|r| (r * stride..r * stride + row_bytes).map(|i| (i % 251) as u8))
7935 .collect()
7936 }
7937
7938 fn padded<T>(format: PixelFormat, stride: usize) -> Tensor<T>
7939 where
7940 T: Num + Clone + fmt::Debug + Send + Sync + 'static,
7941 {
7942 let mut t = Tensor::<T>::image_with_capacity(
7943 16,
7944 4,
7945 format,
7946 Some(TensorMemory::Mem),
7947 CpuAccess::ReadWrite,
7948 )
7949 .unwrap();
7950 t.configure_image(8, 4, format).unwrap();
7951 t.set_row_stride(stride).unwrap();
7952 t
7953 }
7954
7955 #[test]
7956 fn copy_to_flat_compacts_padded_packed_rows() {
7957 // 8 px * 3 B = 24 B rows at a 32 B pitch.
7958 let t = padded::<u8>(PixelFormat::Rgb, 32);
7959 fill_pattern(&t);
7960 let mut dst = vec![0u8; 96];
7961 t.copy_to_flat(&mut dst).unwrap();
7962 assert_eq!(dst, expected_rows(4, 32, 24));
7963 }
7964
7965 #[test]
7966 fn copy_to_flat_compacts_padded_planar_rows() {
7967 // f32 planar: 3 planes x 4 rows of 32 B at a 64 B pitch.
7968 let t = padded::<f32>(PixelFormat::PlanarRgb, 64);
7969 fill_pattern(&t);
7970 let mut dst = vec![0u8; 384];
7971 t.copy_to_flat(&mut dst).unwrap();
7972 assert_eq!(dst, expected_rows(12, 64, 32));
7973 }
7974
7975 #[test]
7976 fn copy_to_flat_of_a_tight_stride_is_the_mapping() {
7977 let t = padded::<u8>(PixelFormat::Rgb, 24);
7978 fill_pattern(&t);
7979 let mut dst = vec![0u8; 96];
7980 t.copy_to_flat(&mut dst).unwrap();
7981 assert_eq!(dst, expected_rows(1, 96, 96));
7982
7983 let t = Tensor::<u8>::image(
7984 8,
7985 4,
7986 PixelFormat::Rgb,
7987 Some(TensorMemory::Mem),
7988 CpuAccess::ReadWrite,
7989 )
7990 .unwrap();
7991 fill_pattern(&t);
7992 let mut dst = vec![0u8; 96];
7993 t.copy_to_flat(&mut dst).unwrap();
7994 assert_eq!(dst, expected_rows(1, 96, 96));
7995 }
7996
7997 #[test]
7998 fn copy_to_flat_needs_a_tight_sized_destination() {
7999 let t = padded::<u8>(PixelFormat::Rgb, 32);
8000 let mut dst = vec![0u8; 95];
8001 match t.copy_to_flat(&mut dst) {
8002 Err(Error::InvalidArgument(msg)) => {
8003 assert!(
8004 msg.contains("dst is 95 bytes") && msg.contains("96 bytes"),
8005 "{msg}"
8006 )
8007 }
8008 other => panic!("expected InvalidArgument, got {other:?}"),
8009 }
8010 }
8011
8012 #[test]
8013 fn copy_to_flat_compacts_a_padded_batch() {
8014 // Two 8x2 RGB tiles at a 32 B pitch: four rows in all.
8015 let mut t = padded::<u8>(PixelFormat::Rgb, 32);
8016 t.set_logical_shape(&[2, 2, 8, 3]).unwrap();
8017 fill_pattern(&t);
8018 let mut dst = vec![0u8; 96];
8019 t.copy_to_flat(&mut dst).unwrap();
8020 assert_eq!(dst, expected_rows(4, 32, 24));
8021 }
8022
8023 #[test]
8024 fn strided_planar_map_spans_every_plane_row() {
8025 let t = padded::<u8>(PixelFormat::PlanarRgb, 16);
8026 assert_eq!(t.map().unwrap().as_slice().len(), 16 * 12);
8027 }
8028
8029 #[test]
8030 fn batch_of_a_padded_batch_starts_at_its_padded_tile() {
8031 // Tile 1 starts after two 32 B rows, not after 48 tight bytes.
8032 let mut t = padded::<u8>(PixelFormat::Rgb, 32);
8033 t.set_logical_shape(&[2, 2, 8, 3]).unwrap();
8034 fill_pattern(&t);
8035 let e = t.batch(1).unwrap();
8036 assert_eq!(e.plane_offset(), Some(64));
8037 let mut dst = vec![0u8; 48];
8038 e.copy_to_flat(&mut dst).unwrap();
8039 let tile1: Vec<u8> = expected_rows(4, 32, 24).split_off(48);
8040 assert_eq!(dst, tile1);
8041 }
8042
8043 #[test]
8044 fn strided_packed_batch_map_spans_every_tile_row() {
8045 let mut t = padded::<u8>(PixelFormat::Rgb, 32);
8046 t.set_logical_shape(&[2, 2, 8, 3]).unwrap();
8047 assert_eq!(t.map().unwrap().as_slice().len(), 32 * 4);
8048 }
8049
8050 #[test]
8051 fn strided_access_to_a_flattened_image_is_refused() {
8052 let mut t = padded::<u8>(PixelFormat::PlanarRgb, 16);
8053 t.set_logical_shape(&[96]).unwrap();
8054 match t.map() {
8055 Err(Error::InvalidOperation(msg)) => assert!(msg.contains("[96]"), "{msg}"),
8056 other => panic!("expected InvalidOperation, got {other:?}"),
8057 }
8058 let mut dst = vec![0u8; 96];
8059 match t.copy_to_flat(&mut dst) {
8060 Err(Error::InvalidOperation(msg)) => assert!(msg.contains("[96]"), "{msg}"),
8061 other => panic!("expected InvalidOperation, got {other:?}"),
8062 }
8063 }
8064
8065 #[test]
8066 #[cfg(target_os = "linux")]
8067 fn strided_shm_map_is_bounded_by_the_segment() {
8068 let _lock = crate::tests::fd_lock_shared();
8069 let mut t = Tensor::<u8>::image(
8070 640,
8071 480,
8072 PixelFormat::Nv12,
8073 Some(TensorMemory::Shm),
8074 CpuAccess::ReadWrite,
8075 )
8076 .expect("Linux has POSIX shared memory");
8077 assert_eq!(t.capacity_bytes(), 640 * 720);
8078 assert_eq!(t.map().unwrap().as_slice().len(), 640 * 720);
8079
8080 t.configure_image(320, 240, PixelFormat::Nv12).unwrap();
8081 assert_eq!(t.row_stride(), Some(640), "the pool pitch fits and is kept");
8082 assert_eq!(t.map().unwrap().as_slice().len(), 640 * 360);
8083
8084 t.configure_image(640, 480, PixelFormat::Nv12).unwrap();
8085 t.set_row_stride(704).unwrap();
8086 match t.map() {
8087 Err(Error::InsufficientCapacity { needed, capacity }) => {
8088 assert_eq!((needed, capacity), (704 * 720, 640 * 720))
8089 }
8090 other => panic!("expected InsufficientCapacity, got {other:?}"),
8091 }
8092 }
8093
8094 /// Minimal host-memory `PboOps`: a fixed buffer, never resized.
8095 struct HostPbo(Vec<u8>);
8096
8097 // SAFETY: the pointer addresses `self.0`, allocated once and never
8098 // resized, and the tensor holds the `Arc` to `self` for every map's life.
8099 unsafe impl PboOps for HostPbo {
8100 fn map_buffer(&self, _buffer_id: u32, size: usize) -> Result<PboMapping> {
8101 assert!(size <= self.0.len());
8102 Ok(PboMapping {
8103 ptr: self.0.as_ptr() as *mut u8,
8104 size,
8105 })
8106 }
8107 fn unmap_buffer(&self, _buffer_id: u32) -> Result<()> {
8108 Ok(())
8109 }
8110 fn delete_buffer(&self, _buffer_id: u32) {}
8111 }
8112
8113 fn pbo_rgba(width: usize, height: usize) -> Tensor<u8> {
8114 let bytes = width * height * 4;
8115 let pbo = PboTensor::<u8>::from_pbo(
8116 7,
8117 bytes,
8118 &[height, width, 4],
8119 None,
8120 Arc::new(HostPbo(vec![0; bytes])),
8121 )
8122 .unwrap();
8123 let mut t = Tensor::from_pbo(pbo).unwrap();
8124 t.set_format(PixelFormat::Rgba).unwrap();
8125 t
8126 }
8127
8128 #[test]
8129 fn strided_pbo_map_is_bounded_by_the_buffer() {
8130 let mut t = pbo_rgba(8, 4);
8131 t.set_row_stride(32).unwrap();
8132 assert_eq!(t.map().unwrap().as_slice().len(), 128);
8133
8134 t.set_row_stride(64).unwrap();
8135 match t.map() {
8136 Err(Error::InsufficientCapacity { needed, capacity }) => {
8137 assert_eq!((needed, capacity), (256, 128))
8138 }
8139 other => panic!("expected InsufficientCapacity, got {other:?}"),
8140 }
8141 }
8142
8143 #[test]
8144 fn pbo_accessors_answer_only_for_a_pbo_tensor() {
8145 let mem = Tensor::<u8>::new(&[4], Some(TensorMemory::Mem), None).unwrap();
8146 assert_eq!(mem.pbo_is_mapped(), None);
8147 assert!(mem.pbo_keepalive().is_none());
8148
8149 let t = pbo_rgba(8, 4);
8150 assert_eq!(t.pbo_is_mapped(), Some(false));
8151 assert!(t.pbo_keepalive().is_some());
8152 let m = t.map().unwrap();
8153 assert_eq!(t.pbo_is_mapped(), Some(true));
8154 drop(m);
8155 assert_eq!(t.pbo_is_mapped(), Some(false));
8156 }
8157}
8158
8159#[cfg(all(test, not(edgefirst_tensor_dynamic)))]
8160mod accessor_tests {
8161 use super::*;
8162 use std::sync::atomic::{AtomicUsize, Ordering};
8163
8164 #[test]
8165 #[cfg(unix)]
8166 fn plane_descriptor_carries_its_stride_and_offset() {
8167 use std::os::fd::AsFd;
8168 let _lock = crate::tests::fd_lock_shared();
8169 let file = std::fs::File::open("/dev/null").unwrap();
8170 let pd = PlaneDescriptor::new(file.as_fd()).unwrap();
8171 assert_eq!((pd.stride(), pd.offset()), (None, None));
8172 let pd = pd.with_stride(832).with_offset(4096);
8173 assert_eq!((pd.stride(), pd.offset()), (Some(832), Some(4096)));
8174 }
8175
8176 #[test]
8177 #[cfg(target_os = "linux")]
8178 fn as_dma_answers_only_for_a_dma_buffer() {
8179 use std::os::fd::AsRawFd;
8180 let _lock = crate::tests::fd_lock_shared();
8181 let mem = Tensor::<u8>::new(&[64], Some(TensorMemory::Mem), None).unwrap();
8182 assert!(mem.as_dma().is_none());
8183 if !crate::test_support::dma_or_skip("as_dma_answers_only_for_a_dma_buffer") {
8184 return;
8185 }
8186 let t = Tensor::<u8>::new(&[64], Some(TensorMemory::DmaBuf), None).unwrap();
8187 let dma = t.as_dma().expect("a DMA tensor downcasts");
8188 assert_eq!(
8189 dma.fd.as_raw_fd(),
8190 t.dmabuf().unwrap().as_raw_fd(),
8191 "the downcast is this tensor's own buffer"
8192 );
8193 }
8194
8195 #[test]
8196 fn derived_identity_mixes_the_kind_into_the_key() {
8197 let id = BufferIdentity::derived(IdentityKind::Shm, u64::MAX);
8198 assert_eq!(id.id(), ((IdentityKind::Shm as u64) << 56) ^ u64::MAX);
8199 assert_ne!(
8200 id.id(),
8201 BufferIdentity::derived(IdentityKind::DmaBuf, u64::MAX).id()
8202 );
8203 }
8204
8205 #[test]
8206 fn empty_is_a_zero_element_shape() {
8207 let mut t = Tensor::<f32>::new(&[2, 3], Some(TensorMemory::Mem), None).unwrap();
8208 assert!(!TensorTrait::is_empty(&t));
8209 {
8210 let m = t.map().unwrap();
8211 assert_eq!(TensorMapTrait::len(&m), 6);
8212 assert_eq!(TensorMapTrait::size(&m), 24);
8213 assert!(!TensorMapTrait::is_empty(&m));
8214 }
8215 t.set_logical_shape(&[0]).unwrap();
8216 assert!(TensorTrait::is_empty(&t));
8217 let m = t.map().unwrap();
8218 assert_eq!(TensorMapTrait::size(&m), 0);
8219 assert!(TensorMapTrait::is_empty(&m));
8220 }
8221
8222 struct CountingCudaOps {
8223 maps: AtomicUsize,
8224 }
8225
8226 impl CudaGlOps for CountingCudaOps {
8227 fn map(&self, _r: crate::cuda::GraphicsResource) -> Option<(*mut std::ffi::c_void, usize)> {
8228 self.maps.fetch_add(1, Ordering::SeqCst);
8229 Some((0x1000usize as *mut _, 64))
8230 }
8231 fn unmap(&self, _r: crate::cuda::GraphicsResource) {}
8232 fn unregister(&self, _r: crate::cuda::GraphicsResource) {}
8233 }
8234
8235 #[test]
8236 fn cuda_accessors_follow_the_attached_handle() {
8237 let mut t = Tensor::<f32>::new(&[16], Some(TensorMemory::Mem), None).unwrap();
8238 assert!(!t.is_cuda_attached());
8239 assert!(t.cuda().is_none());
8240 assert!(t.cuda_map().is_none());
8241 assert!(t.cuda_map_mut().is_none());
8242
8243 let ops = Arc::new(CountingCudaOps {
8244 maps: AtomicUsize::new(0),
8245 });
8246 t.set_cuda_handle(CudaHandle::new_gl(
8247 0x1usize as crate::cuda::GraphicsResource,
8248 64,
8249 ops.clone(),
8250 ));
8251 assert!(t.is_cuda_attached());
8252 assert!(t.cuda().is_some());
8253 assert_eq!(t.cuda_map().map(|m| m.device_ptr() as usize), Some(0x1000));
8254 assert_eq!(t.cuda_map_mut().map(|m| m.len()), Some(64));
8255 assert_eq!(ops.maps.load(Ordering::SeqCst), 2);
8256 }
8257
8258 #[test]
8259 fn from_slice_copies_values_of_a_matching_shape() {
8260 let t = Tensor::<u16>::from_slice(&[1, 2, 3, 4, 5, 6], &[2, 3]).unwrap();
8261 assert_eq!(t.shape(), &[2, 3]);
8262 assert_eq!(t.map().unwrap().as_slice(), &[1, 2, 3, 4, 5, 6]);
8263 match Tensor::<u16>::from_slice(&[1, 2, 3, 4, 5], &[2, 3]) {
8264 Err(Error::InvalidShape(msg)) => assert!(msg.contains("values.len()=5"), "{msg}"),
8265 other => panic!("expected InvalidShape, got {other:?}"),
8266 }
8267 }
8268
8269 /// An `Any` compression request with no pinned memory tries the
8270 /// platform's zero-copy image first and falls back when it cannot serve
8271 /// the format (packed RGB f32 has no IOSurface layout, for one).
8272 #[cfg(not(target_os = "android"))]
8273 #[test]
8274 fn any_compression_request_falls_back_when_zero_copy_cannot_serve() {
8275 let _lock = crate::tests::fd_lock_shared();
8276 let desc = ImageDesc::new(10, 10, PixelFormat::Rgb, DType::F32)
8277 .with_access(CpuAccess::None)
8278 .with_compression(Compression::Any);
8279 let t = Tensor::<f32>::image_desc(&desc).expect("Any never fails for compression reasons");
8280 assert_eq!(t.shape(), &[10, 10, 3]);
8281 assert_eq!(t.compression(), None);
8282 }
8283}
8284
8285#[cfg(test)]
8286mod cpu_access_tests {
8287 use super::*;
8288
8289 #[test]
8290 fn covers_matrix() {
8291 use CpuAccess::*;
8292 // Every declaration covers a narrower or equal request…
8293 for a in [None, Read, Write, ReadWrite] {
8294 assert!(a.covers(None), "{a:?} must cover None");
8295 assert!(ReadWrite.covers(a), "ReadWrite must cover {a:?}");
8296 }
8297 assert!(Read.covers(Read));
8298 assert!(Write.covers(Write));
8299 // …and never a wider one.
8300 assert!(!None.covers(Read));
8301 assert!(!None.covers(Write));
8302 assert!(!Read.covers(Write));
8303 assert!(!Read.covers(ReadWrite));
8304 assert!(!Write.covers(Read));
8305 assert!(!Write.covers(ReadWrite));
8306 }
8307
8308 #[test]
8309 fn union_covers_both_and_require_names_both() {
8310 use CpuAccess::*;
8311 for a in [None, Read, Write, ReadWrite] {
8312 for b in [None, Read, Write, ReadWrite] {
8313 let u = a.union(b);
8314 assert!(u.covers(a) && u.covers(b), "{a:?} | {b:?} = {u:?}");
8315 assert_eq!(a.require(b).is_ok(), a.covers(b));
8316 }
8317 }
8318 match Read.require(Write) {
8319 Err(Error::CpuAccessDenied {
8320 declared: Read,
8321 requested: Write,
8322 }) => {}
8323 other => panic!("expected CpuAccessDenied, got {other:?}"),
8324 }
8325 }
8326
8327 #[test]
8328 fn map_with_none_is_invalid() {
8329 let t = Tensor::<u8>::new(&[16], Some(TensorMemory::Mem), None).unwrap();
8330 match t.map_with(CpuAccess::None) {
8331 Err(Error::InvalidArgument(_)) => {}
8332 Err(other) => panic!("expected InvalidArgument, got {other:?}"),
8333 Ok(_) => panic!("map_with(CpuAccess::None) must not succeed"),
8334 }
8335 }
8336
8337 #[test]
8338 fn read_map_rejects_mutation_uniformly() {
8339 // Mem backend: map_read yields a working read view whose mutable
8340 // accessor panics (the uniform cross-backend contract).
8341 let t = Tensor::<u8>::new(&[8], Some(TensorMemory::Mem), None).unwrap();
8342 t.map_mut().unwrap().as_mut_slice().copy_from_slice(&[7; 8]);
8343 let ro = t.map_read().unwrap();
8344 assert_eq!(ro.as_slice(), &[7; 8]);
8345 drop(ro);
8346 let result = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
8347 let mut ro = t.map_read().unwrap();
8348 let _ = ro.as_mut_slice();
8349 }));
8350 assert!(result.is_err(), "as_mut_slice through map_read must panic");
8351 }
8352
8353 #[test]
8354 fn write_and_rw_maps_stay_mutable() {
8355 let t = Tensor::<u8>::new(&[4], Some(TensorMemory::Mem), None).unwrap();
8356 t.map_write()
8357 .unwrap()
8358 .as_mut_slice()
8359 .copy_from_slice(&[1; 4]);
8360 t.map().unwrap().as_mut_slice().copy_from_slice(&[2; 4]);
8361 assert_eq!(t.map_read().unwrap().as_slice(), &[2; 4]);
8362 }
8363
8364 /// The read-only IOSurface lock path: a `map_read` must observe data
8365 /// written through a prior read-write lock, and its unlock (which
8366 /// skips the cache flush) must not disturb subsequent reads.
8367 #[test]
8368 #[cfg(target_os = "macos")]
8369 fn iosurface_read_only_lock_roundtrip() {
8370 // Declares this test as an fd-opener; see FD_LOCK.
8371 let _lock = crate::tests::fd_lock_shared();
8372 let expect = if crate::is_iosurface_available() {
8373 crate::test_support::Expect::Succeeds
8374 } else {
8375 crate::test_support::Expect::MayFail("the IOSurface probe failed on this host")
8376 };
8377 let Some(t) = crate::test_support::allocated(
8378 Tensor::<u8>::new(&[64], Some(TensorMemory::DmaBuf), None),
8379 "IOSurface allocation",
8380 expect,
8381 ) else {
8382 return;
8383 };
8384 {
8385 let mut m = t.map_mut().unwrap();
8386 for (i, b) in m.as_mut_slice().iter_mut().enumerate() {
8387 *b = (i * 3) as u8;
8388 }
8389 }
8390 for _ in 0..2 {
8391 let ro = t.map_read().unwrap();
8392 for (i, b) in ro.as_slice().iter().enumerate() {
8393 assert_eq!(*b, (i * 3) as u8, "byte {i} through read-only lock");
8394 }
8395 }
8396 let result = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
8397 let mut ro = t.map_read().unwrap();
8398 let _ = ro.as_mut_slice();
8399 }));
8400 assert!(result.is_err(), "IOSurface read map must reject mutation");
8401 }
8402}
8403
8404#[cfg(test)]
8405mod tests {
8406 use std::sync::{PoisonError, RwLock};
8407
8408 use super::*;
8409
8410 #[ctor::ctor(unsafe)]
8411 fn init() {
8412 env_logger::Builder::from_env(env_logger::Env::default().default_filter_or("info")).init();
8413 }
8414
8415 /// Macro to get the current function name for logging in tests.
8416 #[cfg(target_os = "linux")]
8417 macro_rules! function {
8418 () => {{
8419 fn f() {}
8420 fn type_name_of<T>(_: T) -> &'static str {
8421 std::any::type_name::<T>()
8422 }
8423 let name = type_name_of(f);
8424
8425 // Find and cut the rest of the path
8426 match &name[..name.len() - 3].rfind(':') {
8427 Some(pos) => &name[pos + 1..name.len() - 3],
8428 None => &name[..name.len() - 3],
8429 }
8430 }};
8431 }
8432
8433 #[test]
8434 #[cfg(target_os = "linux")]
8435 fn test_tensor() {
8436 let _lock = crate::tests::fd_lock_shared();
8437 let shape = vec![1];
8438 let tensor = DmaTensor::<f32>::new(&shape, Some("dma_tensor"));
8439 let dma_enabled = tensor.is_ok();
8440
8441 let tensor = Tensor::<f32>::new(&shape, None, None).expect("Failed to create tensor");
8442 // Auto-select priority is Dma > Mem; Shm is never auto-selected.
8443 match dma_enabled {
8444 true => assert_eq!(tensor.memory(), TensorMemory::DmaBuf),
8445 false => assert_eq!(tensor.memory(), TensorMemory::Mem),
8446 }
8447 }
8448
8449 #[test]
8450 #[cfg(any(target_os = "macos", target_os = "ios"))]
8451 fn test_tensor() {
8452 let shape = vec![1];
8453 let tensor = Tensor::<f32>::new(&shape, None, None).expect("Failed to create tensor");
8454 // macOS/iOS auto-fallback chain: IOSurface (Dma) → Mem. Healthy systems
8455 // return Dma; Mem only appears under memory pressure or sandboxed
8456 // contexts where IOSurfaceCreate fails. Shm is never auto-selected.
8457 let m = tensor.memory();
8458 assert!(
8459 matches!(m, TensorMemory::DmaBuf | TensorMemory::Mem),
8460 "Unexpected auto-fallback result on macOS/iOS: {m:?}"
8461 );
8462 }
8463
8464 #[test]
8465 #[cfg(all(
8466 unix,
8467 not(any(target_os = "linux", target_os = "macos", target_os = "ios"))
8468 ))]
8469 fn test_tensor() {
8470 let shape = vec![1];
8471 let tensor = Tensor::<f32>::new(&shape, None, None).expect("Failed to create tensor");
8472 // Other Unix (BSD): no DMA, so auto-selection is Mem (Shm is
8473 // explicit-only, never auto-selected).
8474 assert_eq!(tensor.memory(), TensorMemory::Mem);
8475 }
8476
8477 #[test]
8478 #[cfg(not(unix))]
8479 fn test_tensor() {
8480 let shape = vec![1];
8481 let tensor = Tensor::<f32>::new(&shape, None, None).expect("Failed to create tensor");
8482 assert_eq!(tensor.memory(), TensorMemory::Mem);
8483 }
8484
8485 #[test]
8486 #[cfg(target_os = "linux")]
8487 fn test_dma_tensor() {
8488 let _lock = crate::tests::fd_lock_shared();
8489 if !crate::test_support::dma_or_skip("test_dma_tensor") {
8490 return;
8491 }
8492
8493 let shape = vec![2, 3, 4];
8494 let tensor =
8495 DmaTensor::<f32>::new(&shape, Some("test_tensor")).expect("Failed to create tensor");
8496
8497 const DUMMY_VALUE: f32 = 12.34;
8498
8499 assert_eq!(tensor.memory(), TensorMemory::DmaBuf);
8500 assert_eq!(tensor.name(), "test_tensor");
8501 assert_eq!(tensor.shape(), &shape);
8502 assert_eq!(tensor.size(), 2 * 3 * 4 * std::mem::size_of::<f32>());
8503 assert_eq!(tensor.len(), 2 * 3 * 4);
8504
8505 {
8506 let mut tensor_map = tensor.map().expect("Failed to map DMA memory");
8507 tensor_map.fill(42.0);
8508 assert!(tensor_map.iter().all(|&x| x == 42.0));
8509 }
8510
8511 {
8512 let shared = Tensor::<f32>::from_fd(
8513 tensor
8514 .clone_fd()
8515 .expect("Failed to duplicate tensor file descriptor"),
8516 &shape,
8517 Some("test_tensor_shared"),
8518 )
8519 .expect("Failed to create tensor from fd");
8520
8521 assert_eq!(shared.memory(), TensorMemory::DmaBuf);
8522 assert_eq!(shared.name(), "test_tensor_shared");
8523 assert_eq!(shared.shape(), &shape);
8524
8525 let mut tensor_map = shared.map().expect("Failed to map DMA memory from fd");
8526 tensor_map.fill(DUMMY_VALUE);
8527 assert!(tensor_map.iter().all(|&x| x == DUMMY_VALUE));
8528 }
8529
8530 {
8531 let tensor_map = tensor.map().expect("Failed to map DMA memory");
8532 assert!(tensor_map.iter().all(|&x| x == DUMMY_VALUE));
8533 }
8534
8535 let mut tensor = DmaTensor::<u8>::new(&shape, None).expect("Failed to create tensor");
8536 assert_eq!(tensor.shape(), &shape);
8537 let new_shape = vec![3, 4, 4];
8538 assert!(
8539 tensor.reshape(&new_shape).is_err(),
8540 "Reshape should fail due to size mismatch"
8541 );
8542 assert_eq!(tensor.shape(), &shape, "Shape should remain unchanged");
8543
8544 let new_shape = vec![2, 3, 4];
8545 tensor.reshape(&new_shape).expect("Reshape should succeed");
8546 assert_eq!(
8547 tensor.shape(),
8548 &new_shape,
8549 "Shape should be updated after successful reshape"
8550 );
8551
8552 {
8553 let mut tensor_map = tensor.map().expect("Failed to map DMA memory");
8554 tensor_map.fill(1);
8555 assert!(tensor_map.iter().all(|&x| x == 1));
8556 }
8557
8558 {
8559 let mut tensor_map = tensor.map().expect("Failed to map DMA memory");
8560 tensor_map[2] = 42;
8561 assert_eq!(tensor_map[1], 1, "Value at index 1 should be 1");
8562 assert_eq!(tensor_map[2], 42, "Value at index 2 should be 42");
8563 }
8564 }
8565
8566 #[test]
8567 #[cfg(unix)]
8568 fn test_shm_tensor() {
8569 let _lock = crate::tests::fd_lock_shared();
8570 if !is_shm_available() {
8571 crate::test_support::report_skip(
8572 "test_shm_tensor - SHM allocation unavailable on this platform",
8573 );
8574 return;
8575 }
8576 let shape = vec![2, 3, 4];
8577 let tensor =
8578 ShmTensor::<f32>::new(&shape, Some("test_tensor")).expect("Failed to create tensor");
8579 assert_eq!(tensor.shape(), &shape);
8580 assert_eq!(tensor.size(), 2 * 3 * 4 * std::mem::size_of::<f32>());
8581 assert_eq!(tensor.name(), "test_tensor");
8582
8583 const DUMMY_VALUE: f32 = 12.34;
8584 {
8585 let mut tensor_map = tensor.map().expect("Failed to map shared memory");
8586 tensor_map.fill(42.0);
8587 assert!(tensor_map.iter().all(|&x| x == 42.0));
8588 }
8589
8590 {
8591 let shared = Tensor::<f32>::from_fd(
8592 tensor
8593 .clone_fd()
8594 .expect("Failed to duplicate tensor file descriptor"),
8595 &shape,
8596 Some("test_tensor_shared"),
8597 )
8598 .expect("Failed to create tensor from fd");
8599
8600 assert_eq!(shared.memory(), TensorMemory::Shm);
8601 assert_eq!(shared.name(), "test_tensor_shared");
8602 assert_eq!(shared.shape(), &shape);
8603
8604 let mut tensor_map = shared.map().expect("Failed to map shared memory from fd");
8605 tensor_map.fill(DUMMY_VALUE);
8606 assert!(tensor_map.iter().all(|&x| x == DUMMY_VALUE));
8607 }
8608
8609 {
8610 let tensor_map = tensor.map().expect("Failed to map shared memory");
8611 assert!(tensor_map.iter().all(|&x| x == DUMMY_VALUE));
8612 }
8613
8614 let mut tensor = ShmTensor::<u8>::new(&shape, None).expect("Failed to create tensor");
8615 assert_eq!(tensor.shape(), &shape);
8616 let new_shape = vec![3, 4, 4];
8617 assert!(
8618 tensor.reshape(&new_shape).is_err(),
8619 "Reshape should fail due to size mismatch"
8620 );
8621 assert_eq!(tensor.shape(), &shape, "Shape should remain unchanged");
8622
8623 let new_shape = vec![2, 3, 4];
8624 tensor.reshape(&new_shape).expect("Reshape should succeed");
8625 assert_eq!(
8626 tensor.shape(),
8627 &new_shape,
8628 "Shape should be updated after successful reshape"
8629 );
8630
8631 {
8632 let mut tensor_map = tensor.map().expect("Failed to map shared memory");
8633 tensor_map.fill(1);
8634 assert!(tensor_map.iter().all(|&x| x == 1));
8635 }
8636
8637 {
8638 let mut tensor_map = tensor.map().expect("Failed to map shared memory");
8639 tensor_map[2] = 42;
8640 assert_eq!(tensor_map[1], 1, "Value at index 1 should be 1");
8641 assert_eq!(tensor_map[2], 42, "Value at index 2 should be 42");
8642 }
8643 }
8644
8645 #[test]
8646 fn mem_subview_partitions_parent_buffer() {
8647 // One heap [2,4] u8 parent (8 bytes). Two [1,4] sub-views at byte
8648 // offsets 0 and 4 must share the parent allocation (zero-copy) and be
8649 // independently writable: view 0 owns bytes [0,4), view 1 owns [4,8).
8650 // Today this is impossible — heap offset is rejected and there is no
8651 // shared sub-view constructor.
8652 let parent = Tensor::<u8>::new(&[2, 4], Some(TensorMemory::Mem), None).unwrap();
8653 let view0 = parent.subview(0, &[1, 4]).expect("subview at offset 0");
8654 let view1 = parent.subview(4, &[1, 4]).expect("subview at offset 4");
8655
8656 view1
8657 .map()
8658 .unwrap()
8659 .as_mut_slice()
8660 .copy_from_slice(&[10, 20, 30, 40]);
8661 view0
8662 .map()
8663 .unwrap()
8664 .as_mut_slice()
8665 .copy_from_slice(&[1, 2, 3, 4]);
8666
8667 // Each view sees only its own window.
8668 assert_eq!(view0.map().unwrap().as_slice(), &[1, 2, 3, 4]);
8669 assert_eq!(view1.map().unwrap().as_slice(), &[10, 20, 30, 40]);
8670 // The parent buffer is correctly partitioned (shared, zero-copy).
8671 assert_eq!(
8672 parent.map().unwrap().as_slice(),
8673 &[1, 2, 3, 4, 10, 20, 30, 40]
8674 );
8675 }
8676
8677 #[test]
8678 fn batch_partitions_leading_dim() {
8679 // Raw [4,2,2,3] u8 batched tensor: 4 elements of 12 bytes each. batch(n)
8680 // yields element n at offset n*12, sharing the parent buffer (zero-copy).
8681 let parent = Tensor::<u8>::new(&[4, 2, 2, 3], Some(TensorMemory::Mem), None).unwrap();
8682 for i in 0..4u8 {
8683 let e = parent.batch(i as usize).expect("batch element");
8684 assert_eq!(e.shape(), &[2, 2, 3]);
8685 // A batch element shares the parent's BufferIdentity.
8686 assert_eq!(e.buffer_identity().id(), parent.buffer_identity().id());
8687 for b in e.map().unwrap().as_mut_slice() {
8688 *b = i + 1;
8689 }
8690 }
8691 // Each element occupies its own 12-byte band of the parent.
8692 let whole = parent.map().unwrap();
8693 let s = whole.as_slice();
8694 for i in 0..4usize {
8695 assert!(
8696 s[i * 12..(i + 1) * 12].iter().all(|&b| b == (i as u8 + 1)),
8697 "band {i} not partitioned: {:?}",
8698 &s[i * 12..(i + 1) * 12]
8699 );
8700 }
8701 }
8702
8703 #[test]
8704 fn view_origin_snapshots_parent_and_composes() {
8705 // view() on a whole image snapshots the parent dims + the view's origin.
8706 let parent = Tensor::<u8>::image(
8707 100,
8708 80,
8709 PixelFormat::Rgba,
8710 Some(TensorMemory::Mem),
8711 crate::CpuAccess::ReadWrite,
8712 )
8713 .unwrap();
8714 assert_eq!(
8715 parent.view_origin(),
8716 None,
8717 "whole tensor has no view_origin"
8718 );
8719 let v = parent.view(Region::new(10, 20, 30, 40)).unwrap();
8720 assert_eq!(
8721 v.view_origin(),
8722 Some(ViewOrigin {
8723 parent_width: 100,
8724 parent_height: 80,
8725 parent_row_stride: 100 * 4, // tight RGBA pitch
8726 x: 10,
8727 y: 20
8728 })
8729 );
8730 // A view of a view keeps the ROOT parent and accumulates the origin.
8731 let v2 = v.view(Region::new(5, 5, 10, 10)).unwrap();
8732 assert_eq!(
8733 v2.view_origin(),
8734 Some(ViewOrigin {
8735 parent_width: 100,
8736 parent_height: 80,
8737 parent_row_stride: 100 * 4,
8738 x: 15,
8739 y: 25
8740 }),
8741 "nested view composes onto the root parent"
8742 );
8743 }
8744
8745 #[test]
8746 fn view_origin_none_for_raw_batch() {
8747 // A raw (unformatted) batched tensor has no pixel geometry, so batch()
8748 // leaves view_origin None (the per-slot path, not the one-import pivot).
8749 let parent = Tensor::<u8>::new(&[4, 2, 2, 3], Some(TensorMemory::Mem), None).unwrap();
8750 assert_eq!(parent.batch(2).unwrap().view_origin(), None);
8751 }
8752
8753 #[test]
8754 fn batch_rejects_out_of_bounds_index() {
8755 let parent = Tensor::<u8>::new(&[4, 2, 2, 3], Some(TensorMemory::Mem), None).unwrap();
8756 match parent.batch(4) {
8757 Err(Error::BatchIndexOutOfBounds { index, batch }) => {
8758 assert_eq!((index, batch), (4, 4));
8759 }
8760 other => panic!("expected BatchIndexOutOfBounds, got {other:?}"),
8761 }
8762 }
8763
8764 #[test]
8765 fn batch_zero_on_unit_n_is_whole() {
8766 // N == 1: batch(0) is the whole per-element block at offset 0 (no plane_offset).
8767 let parent = Tensor::<u8>::new(&[1, 2, 2, 3], Some(TensorMemory::Mem), None).unwrap();
8768 let e = parent.batch(0).unwrap();
8769 assert_eq!(e.shape(), &[2, 2, 3]);
8770 assert_eq!(e.plane_offset(), None);
8771 assert_eq!(e.buffer_identity().id(), parent.buffer_identity().id());
8772 }
8773
8774 #[test]
8775 fn mem_subview_rejects_unaligned_offset() {
8776 // f32 has align 4; a byte offset of 2 cannot back a valid `*const f32`.
8777 let parent = Tensor::<f32>::new(&[8], Some(TensorMemory::Mem), None).unwrap();
8778 assert!(parent.subview(2, &[1]).is_err());
8779 // A correctly aligned offset is accepted.
8780 assert!(parent.subview(4, &[1]).is_ok());
8781 }
8782
8783 #[test]
8784 fn mem_subview_rejects_out_of_bounds() {
8785 let parent = Tensor::<u8>::new(&[8], Some(TensorMemory::Mem), None).unwrap();
8786 // offset 6 + 4 bytes = 10 exceeds the 8-byte allocation.
8787 assert!(parent.subview(6, &[4]).is_err());
8788 }
8789
8790 /// Regression guard for the `TensorTrait::view` promotion (R2): a `subview`
8791 /// must share the parent's `BufferIdentity` on **every** backend, not mint a
8792 /// fresh one. Identity-keyed caches (the GL EGLImage import) rely on this to
8793 /// treat offset-distinct windows of one buffer as a single import; a fresh
8794 /// identity would silently break that and regress zero-copy import reuse.
8795 ///
8796 /// Runs each backend that can be allocated without a GPU/GL context on the
8797 /// test host: `Mem` (always); `Shm` (when POSIX shm is available); the
8798 /// platform-native zero-copy buffer `Dma` (DMA-BUF on Linux / IOSurface on
8799 /// macOS, when available). `Pbo` shares its identity the same way (see
8800 /// `pbo.rs` `view`) but needs a live GL context, so it is exercised by the
8801 /// image-crate GL tests rather than here.
8802 #[test]
8803 fn subview_shares_buffer_identity_all_backends() {
8804 // Declares this test as an fd-opener; see FD_LOCK.
8805 let _lock = crate::tests::fd_lock_shared();
8806 // u8 has align 1, so every byte offset is valid for the alignment check;
8807 // this isolates the identity-sharing contract from alignment concerns.
8808 // `window` is the view shape; a formatted parent validates it against
8809 // its own format, so an image parent needs an image-shaped window.
8810 let assert_parent_shares = |parent: Tensor<u8>, window: &[usize], label: &str| {
8811 let parent_id = parent.buffer_identity().id();
8812 // Two offset-distinct windows must both carry the parent's identity.
8813 let v0 = parent
8814 .subview(0, window)
8815 .unwrap_or_else(|e| panic!("{label}: subview(0) failed: {e:?}"));
8816 let v1 = parent
8817 .subview(16, window)
8818 .unwrap_or_else(|e| panic!("{label}: subview(16) failed: {e:?}"));
8819 assert_eq!(
8820 v0.buffer_identity().id(),
8821 parent_id,
8822 "{label}: subview(0) minted a fresh BufferIdentity"
8823 );
8824 assert_eq!(
8825 v1.buffer_identity().id(),
8826 parent_id,
8827 "{label}: subview(16) minted a fresh BufferIdentity"
8828 );
8829 };
8830 let assert_shares = |memory: TensorMemory, label: &str| {
8831 let parent = Tensor::<u8>::new(&[64], Some(memory), None)
8832 .unwrap_or_else(|e| panic!("{label}: parent alloc failed: {e:?}"));
8833 assert_parent_shares(parent, &[16], label);
8834 };
8835
8836 assert_shares(TensorMemory::Mem, "Mem");
8837
8838 #[cfg(unix)]
8839 if crate::is_shm_available() {
8840 assert_shares(TensorMemory::Shm, "Shm");
8841 }
8842
8843 // Dma == DMA-BUF on Linux, IOSurface on macOS, AHardwareBuffer on
8844 // Android; same public variant.
8845 #[cfg(target_os = "linux")]
8846 if crate::test_support::dma_or_skip("subview_shares_buffer_identity_all_backends (Dma)") {
8847 assert_shares(TensorMemory::DmaBuf, "Dma");
8848 }
8849 #[cfg(not(any(target_os = "linux", target_os = "windows")))]
8850 if crate::is_gpu_buffer_available() {
8851 assert_shares(TensorMemory::DmaBuf, "Dma");
8852 }
8853 // Windows' Dma is an image-formatted D3D11 texture, which
8854 // `Tensor::new` deliberately refuses -- allocate it the only way it
8855 // can be allocated and hold it to the same contract.
8856 #[cfg(target_os = "windows")]
8857 if crate::is_gpu_buffer_available() {
8858 let parent = Tensor::<u8>::image(
8859 64,
8860 1,
8861 PixelFormat::Grey,
8862 Some(TensorMemory::DmaBuf),
8863 CpuAccess::ReadWrite,
8864 )
8865 .expect("D3D11 texture alloc failed");
8866 assert_parent_shares(parent, &[1, 16, 1], "Dma");
8867 }
8868 }
8869
8870 /// A sub-view records a plane offset, and on Windows the D3D11 storage
8871 /// already applied that offset itself when the view was made -- so the
8872 /// map's plane-offset gate must accept it, exactly as it does for the
8873 /// other zero-copy backings. It used to refuse, with a message that named
8874 /// every backend but this one.
8875 #[cfg(target_os = "windows")]
8876 #[test]
8877 fn subview_of_a_texture_tensor_maps_at_its_offset() {
8878 if !crate::is_gpu_buffer_available() {
8879 crate::test_support::report_skip(
8880 "subview_of_a_texture_tensor_maps_at_its_offset - no D3D11 device on this host",
8881 );
8882 return;
8883 }
8884 let t = Tensor::<u8>::image(
8885 32,
8886 8,
8887 PixelFormat::Rgba,
8888 Some(TensorMemory::DmaBuf),
8889 CpuAccess::ReadWrite,
8890 )
8891 .unwrap();
8892 // For a texture tensor this is the backing pitch either way: the
8893 // allocation records the staging pitch when it exceeds the tight row,
8894 // and the two are equal when it does not.
8895 let stride = t.effective_row_stride().unwrap();
8896 let v = t.subview(stride * 2, &[2, 32, 4]).unwrap();
8897 assert!(
8898 v.plane_offset().is_some_and(|o| o > 0),
8899 "a sub-view at a row offset must carry one"
8900 );
8901 let m = v
8902 .map_read()
8903 .expect("a sub-view of a texture tensor must map");
8904 assert_eq!(m.as_slice().len(), 2 * stride, "two rows of the window");
8905 }
8906
8907 #[test]
8908 fn mem_subview_four_views_no_aliasing() {
8909 // One [4,3] f32 parent; four [1,3] views at 12-byte strides, each
8910 // written independently. Exercises a multi-byte element type (offsets
8911 // must stay element-aligned) and N-way zero-copy sharing.
8912 let parent = Tensor::<f32>::new(&[4, 3], Some(TensorMemory::Mem), None).unwrap();
8913 let frame = 3 * std::mem::size_of::<f32>();
8914 for i in 0..4 {
8915 let v = parent.subview(i * frame, &[1, 3]).unwrap();
8916 let val = i as f32 + 1.0;
8917 v.map()
8918 .unwrap()
8919 .as_mut_slice()
8920 .copy_from_slice(&[val, val, val]);
8921 }
8922 assert_eq!(
8923 parent.map().unwrap().as_slice(),
8924 &[1.0, 1.0, 1.0, 2.0, 2.0, 2.0, 3.0, 3.0, 3.0, 4.0, 4.0, 4.0]
8925 );
8926 }
8927
8928 #[test]
8929 fn mem_subview_inherits_format_and_row_stride() {
8930 // A sub-view is a ready-to-use sub-image: it inherits the parent's
8931 // pixel format and (crucially) its padded row stride, so a strided
8932 // parent yields strided windows. Set a stride wider than the tight row
8933 // to exercise the row_stride inheritance path specifically.
8934 let mut parent = Tensor::<u8>::image(
8935 100,
8936 100,
8937 PixelFormat::Rgba,
8938 Some(TensorMemory::Mem),
8939 crate::CpuAccess::ReadWrite,
8940 )
8941 .unwrap();
8942 parent.set_row_stride_unchecked(512); // padded stride (> 100*4)
8943 let view = parent.subview(4096, &[10, 10, 4]).unwrap();
8944 assert_eq!(view.format(), Some(PixelFormat::Rgba), "format inherited");
8945 assert_eq!(view.row_stride(), Some(512), "row_stride inherited");
8946 }
8947
8948 #[test]
8949 fn mem_strided_subview_maps_offset_and_byte_size() {
8950 // Integration of the sub-region offset (PR #89) and the strided-map
8951 // `byte_size_override` (PR #90): a strided sub-view exposes its full
8952 // padded window (`row_stride × rows`) starting at the view's byte
8953 // offset, mapped zero-copy into the parent.
8954 let parent = Tensor::<u8>::new(&[2048], Some(TensorMemory::Mem), None).unwrap();
8955 let mut view = parent.subview(128, &[8, 16]).unwrap(); // 8 rows × 16 @ off 128
8956 assert_eq!(view.plane_offset(), Some(128));
8957 view.set_row_stride_unchecked(32); // padded stride (> 16)
8958
8959 {
8960 let mut m = view.map().unwrap();
8961 let s = m.as_mut_slice();
8962 // Strided map exposes the padded window: stride(32) × rows(8) = 256.
8963 assert_eq!(
8964 s.len(),
8965 256,
8966 "strided map exposes the full padded byte window"
8967 );
8968 s[0] = 0xAA; // row 0, col 0
8969 s[32] = 0xBB; // row 1, col 0 (one stride in)
8970 }
8971
8972 // Zero-copy: the writes land in the parent at the view's offset.
8973 let p = parent.map().unwrap();
8974 let pb = p.as_slice();
8975 assert_eq!(pb[128], 0xAA, "row 0 writes at parent offset 128");
8976 assert_eq!(
8977 pb[128 + 32],
8978 0xBB,
8979 "row 1 writes at parent offset 128 + stride"
8980 );
8981 }
8982
8983 #[test]
8984 #[cfg(unix)]
8985 fn shm_subview_partitions_parent_buffer() {
8986 // Declares this test as an fd-opener; see FD_LOCK.
8987 let _lock = crate::tests::fd_lock_shared();
8988 // Mirrors `mem_subview_partitions_parent_buffer` for Shm: one [2,4] u8
8989 // parent shared segment (8 bytes); two [1,4] sub-views at byte offsets 0
8990 // and 4 must share the segment (zero-copy, via cloned fd) and be
8991 // independently writable — view 0 owns [0,4), view 1 owns [4,8).
8992 if !crate::is_shm_available() {
8993 crate::test_support::report_skip("shm not available");
8994 return;
8995 }
8996 let parent = Tensor::<u8>::new(&[2, 4], Some(TensorMemory::Shm), None).unwrap();
8997 let view0 = parent.subview(0, &[1, 4]).expect("shm subview at offset 0");
8998 let view1 = parent.subview(4, &[1, 4]).expect("shm subview at offset 4");
8999
9000 view1
9001 .map()
9002 .unwrap()
9003 .as_mut_slice()
9004 .copy_from_slice(&[10, 20, 30, 40]);
9005 view0
9006 .map()
9007 .unwrap()
9008 .as_mut_slice()
9009 .copy_from_slice(&[1, 2, 3, 4]);
9010
9011 assert_eq!(view0.map().unwrap().as_slice(), &[1, 2, 3, 4]);
9012 assert_eq!(view1.map().unwrap().as_slice(), &[10, 20, 30, 40]);
9013 // The parent sees the full partitioned segment (shared, zero-copy).
9014 assert_eq!(
9015 parent.map().unwrap().as_slice(),
9016 &[1, 2, 3, 4, 10, 20, 30, 40]
9017 );
9018 // A sub-view of a sub-view composes the offset.
9019 let nested = view1.subview(2, &[1, 2]).expect("nested shm subview");
9020 assert_eq!(nested.map().unwrap().as_slice(), &[30, 40]);
9021 }
9022
9023 #[test]
9024 #[cfg(unix)]
9025 fn shm_subview_rejects_unaligned_and_oob() {
9026 // Declares this test as an fd-opener; see FD_LOCK.
9027 let _lock = crate::tests::fd_lock_shared();
9028 if !crate::is_shm_available() {
9029 crate::test_support::report_skip("shm not available");
9030 return;
9031 }
9032 // f32 align 4: a 2-byte offset cannot back a valid `*const f32`.
9033 let parent = Tensor::<f32>::new(&[8], Some(TensorMemory::Shm), None).unwrap();
9034 assert!(parent.subview(2, &[1]).is_err());
9035 assert!(parent.subview(4, &[1]).is_ok());
9036 // Out of bounds: offset 6 + 4 bytes = 10 > 8-byte (u8) segment.
9037 let p2 = Tensor::<u8>::new(&[8], Some(TensorMemory::Shm), None).unwrap();
9038 assert!(p2.subview(6, &[4]).is_err());
9039 }
9040
9041 #[test]
9042 #[cfg(target_os = "linux")]
9043 fn dma_subview_matches_mem_subview() {
9044 // Serialize against the fd-leak tests: this test opens DMA fds (alloc +
9045 // clone_fd), which would otherwise perturb their fd counts.
9046 let _lock = crate::tests::fd_lock_shared();
9047 // Identical sub-view semantics across Dma (shared fd) and Mem (shared
9048 // Arc): same offsets → same logical windows → same partition.
9049 if !crate::test_support::dma_or_skip("dma_subview_matches_mem_subview") {
9050 return;
9051 }
9052 let dma = Tensor::<u8>::new(&[8], Some(TensorMemory::DmaBuf), None).unwrap();
9053 let mem = Tensor::<u8>::new(&[8], Some(TensorMemory::Mem), None).unwrap();
9054 for parent in [&dma, &mem] {
9055 let v0 = parent.subview(0, &[4]).unwrap();
9056 let v1 = parent.subview(4, &[4]).unwrap();
9057 v0.map()
9058 .unwrap()
9059 .as_mut_slice()
9060 .copy_from_slice(&[1, 2, 3, 4]);
9061 v1.map()
9062 .unwrap()
9063 .as_mut_slice()
9064 .copy_from_slice(&[5, 6, 7, 8]);
9065 assert_eq!(parent.map().unwrap().as_slice(), &[1, 2, 3, 4, 5, 6, 7, 8]);
9066 }
9067 }
9068
9069 #[test]
9070 #[cfg(target_os = "linux")]
9071 fn dma_strided_subview_maps_padded_window() {
9072 // The strided-map path differs by backing: DMA maps through
9073 // `mmap_offset` + the `byte_size_override`, not the Mem `Arc` slice. A
9074 // padded sub-view of a DMA buffer must still expose its full
9075 // `row_stride × rows` window zero-copy at the view's offset (the GPU
9076 // batched-render-to-DMA case). Mirrors
9077 // `mem_strided_subview_maps_offset_and_byte_size` on a Dma parent.
9078 let _lock = crate::tests::fd_lock_shared();
9079 if !crate::test_support::dma_or_skip("dma_strided_subview_maps_padded_window") {
9080 return;
9081 }
9082 let parent = Tensor::<u8>::new(&[2048], Some(TensorMemory::DmaBuf), None).unwrap();
9083 let mut view = parent.subview(128, &[8, 16]).unwrap();
9084 assert_eq!(view.plane_offset(), Some(128));
9085 view.set_row_stride_unchecked(32); // padded stride (> 16)
9086
9087 {
9088 let mut m = view.map().unwrap();
9089 let s = m.as_mut_slice();
9090 assert_eq!(s.len(), 256, "strided DMA map exposes stride(32) × rows(8)");
9091 s[0] = 0xAA; // row 0, col 0
9092 s[32] = 0xBB; // row 1, col 0 (one stride in)
9093 }
9094
9095 let p = parent.map().unwrap();
9096 let pb = p.as_slice();
9097 assert_eq!(pb[128], 0xAA, "row 0 writes at parent offset 128");
9098 assert_eq!(
9099 pb[128 + 32],
9100 0xBB,
9101 "row 1 writes at parent offset 128 + stride"
9102 );
9103 }
9104
9105 #[test]
9106 #[cfg(target_os = "linux")]
9107 fn view_single_row_snapshots_parent_stride() {
9108 // A single-row `view()` keeps a TIGHT `row_stride` for map-span safety,
9109 // but its `view_origin` snapshots the PARENT row stride — the GL backend
9110 // keys its EGLImage import/pitch on that snapshot (not the view's tight
9111 // stride), so single-row and multi-row sibling views collapse onto the
9112 // same parent import.
9113 let _lock = crate::tests::fd_lock_shared();
9114 // 8x4 RGBA with a padded 64-byte row stride (tight row = 8*4 = 32).
9115 if !crate::test_support::dma_or_skip("view_single_row_snapshots_parent_stride") {
9116 return;
9117 }
9118 let parent = Tensor::<u8>::image_with_stride(
9119 8,
9120 4,
9121 PixelFormat::Rgba,
9122 64,
9123 Some(TensorMemory::DmaBuf),
9124 crate::CpuAccess::ReadWrite,
9125 )
9126 .unwrap();
9127 assert_eq!(parent.effective_row_stride(), Some(64));
9128 // Bottom row (y=3) at x>0 — the case the tight single-row stride guards.
9129 let row = parent.view(Region::new(2, 3, 4, 1)).unwrap();
9130 // The view's own stride is tight (4*4 = 16) so its strided map stays in
9131 // bounds; the GL-facing parent pitch (64) lives in `view_origin`.
9132 assert_eq!(row.effective_row_stride(), Some(16));
9133 let vo = row.view_origin().expect("a view carries a view_origin");
9134 assert_eq!(
9135 vo.parent_row_stride, 64,
9136 "GL keys/pitches a view on the parent stride, not its tight one"
9137 );
9138 // The tight stride keeps map() in-bounds for the bottom / x>0 single row.
9139 assert_eq!(row.map().unwrap().as_slice().len(), 16);
9140 }
9141
9142 #[test]
9143 fn test_mem_tensor() {
9144 let shape = vec![2, 3, 4];
9145 let tensor =
9146 MemTensor::<f32>::new(&shape, Some("test_tensor")).expect("Failed to create tensor");
9147 assert_eq!(tensor.shape(), &shape);
9148 assert_eq!(tensor.size(), 2 * 3 * 4 * std::mem::size_of::<f32>());
9149 assert_eq!(tensor.name(), "test_tensor");
9150
9151 {
9152 let mut tensor_map = tensor.map().expect("Failed to map memory");
9153 tensor_map.fill(42.0);
9154 assert!(tensor_map.iter().all(|&x| x == 42.0));
9155 }
9156
9157 let mut tensor = MemTensor::<u8>::new(&shape, None).expect("Failed to create tensor");
9158 assert_eq!(tensor.shape(), &shape);
9159 let new_shape = vec![3, 4, 4];
9160 assert!(
9161 tensor.reshape(&new_shape).is_err(),
9162 "Reshape should fail due to size mismatch"
9163 );
9164 assert_eq!(tensor.shape(), &shape, "Shape should remain unchanged");
9165
9166 let new_shape = vec![2, 3, 4];
9167 tensor.reshape(&new_shape).expect("Reshape should succeed");
9168 assert_eq!(
9169 tensor.shape(),
9170 &new_shape,
9171 "Shape should be updated after successful reshape"
9172 );
9173
9174 {
9175 let mut tensor_map = tensor.map().expect("Failed to map memory");
9176 tensor_map.fill(1);
9177 assert!(tensor_map.iter().all(|&x| x == 1));
9178 }
9179
9180 {
9181 let mut tensor_map = tensor.map().expect("Failed to map memory");
9182 tensor_map[2] = 42;
9183 assert_eq!(tensor_map[1], 1, "Value at index 1 should be 1");
9184 assert_eq!(tensor_map[2], 42, "Value at index 2 should be 42");
9185 }
9186 }
9187
9188 #[test]
9189 #[cfg(target_os = "linux")]
9190 fn test_dma_no_fd_leaks() {
9191 let _lock = fd_lock_exclusive();
9192 if !crate::test_support::dma_or_skip(function!()) {
9193 return;
9194 }
9195
9196 let baseline = open_fd_count(DMA_FD_TARGETS);
9197
9198 // Self-check — see `test_shm_no_fd_leaks`.
9199 {
9200 let probe = Tensor::<u8>::new(&[100, 100], Some(TensorMemory::DmaBuf), None)
9201 .expect("Failed to create tensor");
9202 assert!(
9203 open_fd_count(DMA_FD_TARGETS) > baseline,
9204 "a live DMA-BUF tensor must be visible to the dmabuf fd counter, \
9205 otherwise this test cannot detect a leak"
9206 );
9207 drop(probe);
9208 }
9209
9210 for _ in 0..100 {
9211 let tensor = Tensor::<u8>::new(&[100, 100], Some(TensorMemory::DmaBuf), None)
9212 .expect("Failed to create tensor");
9213 let mut map = tensor.map().unwrap();
9214 map.as_mut_slice().fill(233);
9215 }
9216
9217 let after = open_fd_count(DMA_FD_TARGETS);
9218 assert_eq!(
9219 baseline, after,
9220 "DMA-BUF file descriptor leak detected: {baseline} -> {after}"
9221 );
9222 }
9223
9224 #[test]
9225 #[cfg(target_os = "linux")]
9226 fn test_dma_from_fd_no_fd_leaks() {
9227 let _lock = fd_lock_exclusive();
9228 if !crate::test_support::dma_or_skip(function!()) {
9229 return;
9230 }
9231
9232 let baseline = open_fd_count(DMA_FD_TARGETS);
9233 let orig = Tensor::<u8>::new(&[100, 100], Some(TensorMemory::DmaBuf), None).unwrap();
9234
9235 // Self-check — see `test_shm_no_fd_leaks`.
9236 let held = open_fd_count(DMA_FD_TARGETS);
9237 assert!(
9238 held > baseline,
9239 "a live DMA-BUF tensor must be visible to the dmabuf fd counter, \
9240 otherwise this test cannot detect a leak"
9241 );
9242
9243 for _ in 0..100 {
9244 let tensor =
9245 Tensor::<u8>::from_fd(orig.clone_fd().unwrap(), orig.shape(), None).unwrap();
9246 assert_eq!(
9247 tensor.memory(),
9248 TensorMemory::DmaBuf,
9249 "DMA-BUF fd must import as Dma, not be silently downgraded"
9250 );
9251 let mut map = tensor.map().unwrap();
9252 map.as_mut_slice().fill(233);
9253 }
9254
9255 let after_loop = open_fd_count(DMA_FD_TARGETS);
9256 assert_eq!(
9257 held, after_loop,
9258 "DMA-BUF file descriptor leak detected: {held} -> {after_loop}"
9259 );
9260
9261 // And the source tensor's own fd goes back on drop.
9262 drop(orig);
9263 let after_drop = open_fd_count(DMA_FD_TARGETS);
9264 assert_eq!(
9265 baseline, after_drop,
9266 "dropping the source tensor must release its fd: {baseline} -> {after_drop}"
9267 );
9268 }
9269
9270 /// A filesystem magic must report its true 32-bit value regardless of
9271 /// how wide, and how signed, `fstatfs`'s `f_type` is on the target.
9272 ///
9273 /// `fs_type_t` is `__fsword_t` on Linux/gnu — `i64` on 64-bit but `i32`
9274 /// on 32-bit (armv7, i686, aarch64-ilp32) — `c_int` on uclibc, and
9275 /// unsigned on musl and s390x. Widening a *signed* 32-bit `f_type`
9276 /// sign-extends every magic with bit 31 set, so a naive widening cast
9277 /// reports e.g. `0xffffffff958458f6` for hugetlbfs. The docs tell
9278 /// callers to look the reported value up in `include/uapi/linux/magic.h`,
9279 /// so a sign-extended value is actively misleading.
9280 ///
9281 /// This cannot be reproduced on a 64-bit host — the defect only exists
9282 /// where `f_type` is a signed 32-bit type — so the test drives the
9283 /// conversion directly with the value such a target would produce.
9284 #[test]
9285 #[cfg(target_os = "linux")]
9286 fn test_fs_magic_normalizes_sign_extended_values() {
9287 // Declares this test as an fd-opener; see FD_LOCK.
9288 let _lock = crate::tests::fd_lock_shared();
9289 // Magics whose bit 31 is set. HUGETLBFS is the one that matters
9290 // most in practice: a MFD_HUGETLB memfd is the likeliest fd to land
9291 // in the UnknownBufferType arm.
9292 const HUGETLBFS_MAGIC: u32 = 0x9584_58f6;
9293 const F2FS_MAGIC: u32 = 0xf2f5_2010;
9294 const BTRFS_MAGIC: u32 = 0x9123_683e;
9295
9296 for magic in [HUGETLBFS_MAGIC, F2FS_MAGIC, BTRFS_MAGIC] {
9297 // 64-bit gnu: f_type is i64 and already holds the true value.
9298 assert_eq!(fs_magic(i64::from(magic)), magic);
9299
9300 // 32-bit gnu / uclibc: f_type is i32, so the value arrives
9301 // sign-extended once widened. It must still report its true
9302 // 32 bits.
9303 let sign_extended = i64::from(magic as i32);
9304 assert!(sign_extended < 0, "{magic:#x} should have bit 31 set");
9305 assert_eq!(
9306 fs_magic(sign_extended),
9307 magic,
9308 "{magic:#x} must survive a signed 32-bit f_type"
9309 );
9310 }
9311
9312 // The two magics we actually classify on are below 2^31, so they are
9313 // unaffected by signedness either way — this is why the DMA-vs-SHM
9314 // decision is correct on 32-bit even without this normalization.
9315 for magic in [DMA_BUF_MAGIC, TMPFS_MAGIC] {
9316 assert!(magic < 0x8000_0000);
9317 assert_eq!(fs_magic(i64::from(magic as i32)), magic);
9318 }
9319 }
9320
9321 /// A DMA-BUF fd must import as [`TensorMemory::DmaBuf`].
9322 ///
9323 /// Regression test for the `st_dev` minor-number classifier, which
9324 /// hardcoded `9 | 10` as "this is DMA". Those minors come from
9325 /// `get_anon_bdev()` and are assigned first-come-first-served at boot,
9326 /// so they vary by kernel build and boot order — a real DMA-BUF is
9327 /// minor 12 on x86 desktop and minor 8 on the ADIS Verdin. Both fell
9328 /// into the `_` arm and imported as SHM, which "works" (a DMA-BUF is
9329 /// mmap-able) but silently forfeits zero-copy.
9330 #[test]
9331 #[cfg(target_os = "linux")]
9332 fn test_from_fd_dma_imports_as_dma() {
9333 let _lock = crate::tests::fd_lock_shared();
9334 if !crate::test_support::dma_or_skip(function!()) {
9335 return;
9336 }
9337
9338 let orig = Tensor::<u8>::new(&[64, 64], Some(TensorMemory::DmaBuf), None).unwrap();
9339 assert_eq!(orig.memory(), TensorMemory::DmaBuf);
9340
9341 let imported = Tensor::<u8>::from_fd(orig.clone_fd().unwrap(), orig.shape(), None).unwrap();
9342
9343 assert_eq!(
9344 imported.memory(),
9345 TensorMemory::DmaBuf,
9346 "a DMA-BUF fd must import as Dma"
9347 );
9348 }
9349
9350 /// A tmpfs/SHM fd must import as [`TensorMemory::Shm`].
9351 ///
9352 /// The companion to `test_from_fd_dma_imports_as_dma`: confirms the
9353 /// magic-based classifier identifies SHM positively rather than by
9354 /// falling through.
9355 #[test]
9356 #[cfg(target_os = "linux")]
9357 fn test_from_fd_shm_imports_as_shm() {
9358 let _lock = crate::tests::fd_lock_shared();
9359 if !is_shm_available() {
9360 crate::test_support::report_skip(&format!(
9361 "{} - SHM memory not available",
9362 function!()
9363 ));
9364 return;
9365 }
9366
9367 let orig = Tensor::<u8>::new(&[64, 64], Some(TensorMemory::Shm), None).unwrap();
9368 assert_eq!(orig.memory(), TensorMemory::Shm);
9369
9370 let imported = Tensor::<u8>::from_fd(orig.clone_fd().unwrap(), orig.shape(), None).unwrap();
9371
9372 assert_eq!(
9373 imported.memory(),
9374 TensorMemory::Shm,
9375 "a tmpfs fd must import as Shm"
9376 );
9377 }
9378
9379 /// An fd that is neither a DMA-BUF nor tmpfs must be rejected.
9380 ///
9381 /// A pipe is the convenient probe: it lives on `pipefs`, another
9382 /// `get_anon_bdev()` pseudo-filesystem, so it shares major 0 with the
9383 /// buffer types we do support and is only distinguishable by magic.
9384 /// Importing one as SHM is meaningless — `mmap` on a pipe fails — so
9385 /// the classifier must say "unknown" rather than guess.
9386 #[test]
9387 #[cfg(target_os = "linux")]
9388 fn test_from_fd_rejects_unknown_filesystem() {
9389 let _lock = crate::tests::fd_lock_shared();
9390
9391 let (read_end, _write_end) = nix::unistd::pipe().unwrap();
9392
9393 let result = Tensor::<u8>::from_fd(read_end, &[64], None);
9394
9395 match result {
9396 Err(Error::UnknownBufferType(magic)) => {
9397 // PIPEFS_MAGIC, from include/uapi/linux/magic.h
9398 assert_eq!(magic, 0x5049_5045, "expected PIPEFS_MAGIC");
9399 }
9400 other => panic!("expected UnknownBufferType for a pipe fd, got {other:?}"),
9401 }
9402 }
9403
9404 #[test]
9405 #[cfg(target_os = "linux")]
9406 fn test_shm_no_fd_leaks() {
9407 let _lock = fd_lock_exclusive();
9408 if !is_shm_available() {
9409 crate::test_support::report_skip(&format!(
9410 "{} - SHM memory allocation not available (permission denied or no SHM support)",
9411 function!()
9412 ));
9413 return;
9414 }
9415
9416 let baseline = open_fd_count(SHM_FD_TARGETS);
9417
9418 // Self-check: the counter must actually see a live shm tensor. If
9419 // it does not — a changed backing object, a different `/proc`
9420 // rendering — the leak assertion below would hold no matter how
9421 // much leaked, so fail loudly here instead.
9422 {
9423 let probe = Tensor::<u8>::new(&[100, 100], Some(TensorMemory::Shm), None)
9424 .expect("Failed to create tensor");
9425 assert!(
9426 open_fd_count(SHM_FD_TARGETS) > baseline,
9427 "a live shm tensor must be visible to the shm fd counter, \
9428 otherwise this test cannot detect a leak"
9429 );
9430 drop(probe);
9431 }
9432
9433 for _ in 0..100 {
9434 let tensor = Tensor::<u8>::new(&[100, 100], Some(TensorMemory::Shm), None)
9435 .expect("Failed to create tensor");
9436 let mut map = tensor.map().unwrap();
9437 map.as_mut_slice().fill(233);
9438 }
9439
9440 let after = open_fd_count(SHM_FD_TARGETS);
9441 assert_eq!(
9442 baseline, after,
9443 "shm file descriptor leak detected: {baseline} -> {after}"
9444 );
9445 }
9446
9447 #[test]
9448 #[cfg(target_os = "linux")]
9449 fn test_shm_from_fd_no_fd_leaks() {
9450 let _lock = fd_lock_exclusive();
9451 if !is_shm_available() {
9452 crate::test_support::report_skip(&format!(
9453 "{} - SHM memory allocation not available (permission denied or no SHM support)",
9454 function!()
9455 ));
9456 return;
9457 }
9458
9459 let baseline = open_fd_count(SHM_FD_TARGETS);
9460 let orig = Tensor::<u8>::new(&[100, 100], Some(TensorMemory::Shm), None).unwrap();
9461
9462 // Self-check — see `test_shm_no_fd_leaks`.
9463 let held = open_fd_count(SHM_FD_TARGETS);
9464 assert!(
9465 held > baseline,
9466 "a live shm tensor must be visible to the shm fd counter, \
9467 otherwise this test cannot detect a leak"
9468 );
9469
9470 for _ in 0..100 {
9471 let tensor =
9472 Tensor::<u8>::from_fd(orig.clone_fd().unwrap(), orig.shape(), None).unwrap();
9473 let mut map = tensor.map().unwrap();
9474 map.as_mut_slice().fill(233);
9475 }
9476
9477 let after_loop = open_fd_count(SHM_FD_TARGETS);
9478 assert_eq!(
9479 held, after_loop,
9480 "shm file descriptor leak detected: {held} -> {after_loop}"
9481 );
9482
9483 // And the source tensor's own fd goes back on drop.
9484 drop(orig);
9485 let after_drop = open_fd_count(SHM_FD_TARGETS);
9486 assert_eq!(
9487 baseline, after_drop,
9488 "dropping the source tensor must release its fd: {baseline} -> {after_drop}"
9489 );
9490 }
9491
9492 #[cfg(feature = "ndarray")]
9493 #[test]
9494 fn test_ndarray() {
9495 let _lock = crate::tests::fd_lock_shared();
9496 let shape = vec![2, 3, 4];
9497 let tensor = Tensor::<f32>::new(&shape, None, None).expect("Failed to create tensor");
9498
9499 let mut tensor_map = tensor.map().expect("Failed to map tensor memory");
9500 tensor_map.fill(1.0);
9501
9502 let view = tensor_map.view().expect("Failed to get ndarray view");
9503 assert_eq!(view.shape(), &[2, 3, 4]);
9504 assert!(view.iter().all(|&x| x == 1.0));
9505
9506 let mut view_mut = tensor_map
9507 .view_mut()
9508 .expect("Failed to get mutable ndarray view");
9509 view_mut[[0, 0, 0]] = 42.0;
9510 assert_eq!(view_mut[[0, 0, 0]], 42.0);
9511 assert_eq!(tensor_map[0], 42.0, "Value at index 0 should be 42");
9512 }
9513
9514 #[test]
9515 fn test_buffer_identity_derived_equality() {
9516 // Two identities derived from the same kind and key must be equal —
9517 // that is the whole point of `derived`: two independently-linked
9518 // copies of this crate, each importing the same underlying buffer,
9519 // have to agree on its id without talking to each other.
9520 let id1 = BufferIdentity::derived(IdentityKind::DmaBuf, 4242);
9521 let id2 = BufferIdentity::derived(IdentityKind::DmaBuf, 4242);
9522 assert_eq!(
9523 id1.id(),
9524 id2.id(),
9525 "Identities derived from the same kind and key should match"
9526 );
9527
9528 // A different key still yields a different id.
9529 let id3 = BufferIdentity::derived(IdentityKind::DmaBuf, 4243);
9530 assert_ne!(
9531 id1.id(),
9532 id3.id(),
9533 "Identities derived from different keys should differ"
9534 );
9535 }
9536
9537 #[test]
9538 fn test_buffer_identity_clone_shares_guard() {
9539 let id1 = BufferIdentity::derived(IdentityKind::DmaBuf, 4242);
9540 let weak = id1.weak();
9541 assert!(
9542 weak.upgrade().is_some(),
9543 "Weak should be alive while original exists"
9544 );
9545
9546 let id2 = id1.clone();
9547 assert_eq!(id1.id(), id2.id(), "Cloned identity should have same id");
9548
9549 drop(id1);
9550 assert!(
9551 weak.upgrade().is_some(),
9552 "Weak should still be alive (clone holds Arc)"
9553 );
9554
9555 drop(id2);
9556 assert!(
9557 weak.upgrade().is_none(),
9558 "Weak should be dead after all clones dropped"
9559 );
9560 }
9561
9562 #[test]
9563 fn test_tensor_buffer_identity() {
9564 let t1 = Tensor::<u8>::new(&[100], Some(TensorMemory::Mem), Some("t1")).unwrap();
9565 let t2 = Tensor::<u8>::new(&[100], Some(TensorMemory::Mem), Some("t2")).unwrap();
9566 assert_ne!(
9567 t1.buffer_identity().id(),
9568 t2.buffer_identity().id(),
9569 "Different tensors should have different buffer ids"
9570 );
9571 }
9572
9573 // ------------------------------------------------------------------------
9574 // Quantization — constructor validation + accessor correctness.
9575 // ------------------------------------------------------------------------
9576
9577 #[test]
9578 fn test_quantization_per_tensor_constructors() {
9579 let q = Quantization::per_tensor(0.1, -5);
9580 assert!(q.is_per_tensor());
9581 assert!(!q.is_per_channel());
9582 assert!(!q.is_symmetric());
9583 assert_eq!(q.scale(), &[0.1]);
9584 assert_eq!(q.zero_point(), Some(&[-5][..]));
9585
9586 let qs = Quantization::per_tensor_symmetric(0.05);
9587 assert!(qs.is_per_tensor());
9588 assert!(qs.is_symmetric());
9589 assert_eq!(qs.zero_point(), None);
9590 }
9591
9592 #[test]
9593 fn test_quantization_per_channel_constructors() {
9594 let q = Quantization::per_channel(vec![0.1, 0.2, 0.3], vec![0, -1, 1], 2).unwrap();
9595 assert!(q.is_per_channel());
9596 assert!(!q.is_symmetric());
9597 assert_eq!(q.axis(), Some(2));
9598 assert_eq!(q.scale().len(), 3);
9599
9600 let qs = Quantization::per_channel_symmetric(vec![0.054, 0.089, 0.195], 0).unwrap();
9601 assert!(qs.is_per_channel());
9602 assert!(qs.is_symmetric());
9603 assert_eq!(qs.axis(), Some(0));
9604 }
9605
9606 #[test]
9607 fn test_quantization_per_channel_length_mismatch_rejected() {
9608 // len(scales) != len(zero_points) → rejected at construction.
9609 let err = Quantization::per_channel(vec![0.1, 0.2], vec![0, 0, 0], 0).unwrap_err();
9610 assert!(matches!(err, Error::QuantizationInvalid { .. }));
9611 }
9612
9613 #[test]
9614 fn test_quantization_per_channel_empty_rejected() {
9615 let err = Quantization::per_channel_symmetric(vec![], 0).unwrap_err();
9616 assert!(matches!(err, Error::QuantizationInvalid { .. }));
9617 }
9618
9619 /// Constructors guard scale/zero_point length invariants, but
9620 /// `Quantization` is `Deserialize`, so malformed JSON (e.g. an
9621 /// empty `scale` array, or `zero_point` length that disagrees with
9622 /// `scale`) bypasses the constructor checks. `set_quantization`
9623 /// must reject these via `validate()` so they don't poison
9624 /// downstream `mode()` selection or per-channel kernel indexing.
9625 #[test]
9626 fn test_quantization_validate_rejects_malformed_deserialize() {
9627 let mut t = Tensor::<i8>::new(&[1, 1, 4], Some(TensorMemory::Mem), None).unwrap();
9628
9629 // Empty scale array: must be rejected.
9630 let q: Quantization = serde_json::from_str(r#"{"scale": []}"#).unwrap();
9631 assert!(matches!(
9632 t.set_quantization(q).unwrap_err(),
9633 Error::QuantizationInvalid { .. }
9634 ));
9635
9636 // Per-tensor with multi-element zero_point: must be rejected.
9637 let q: Quantization =
9638 serde_json::from_str(r#"{"scale": 0.1, "zero_point": [0, 0, 0]}"#).unwrap();
9639 assert!(matches!(
9640 t.set_quantization(q).unwrap_err(),
9641 Error::QuantizationInvalid { .. }
9642 ));
9643
9644 // Per-channel zero_point length != scale length: must be rejected.
9645 let q: Quantization = serde_json::from_str(
9646 r#"{"scale": [0.1, 0.2, 0.3, 0.4], "zero_point": [0, 0], "axis": 2}"#,
9647 )
9648 .unwrap();
9649 assert!(matches!(
9650 t.set_quantization(q).unwrap_err(),
9651 Error::QuantizationInvalid { .. }
9652 ));
9653 }
9654
9655 #[test]
9656 fn test_quantization_mode_dispatch() {
9657 let pt = Quantization::per_tensor(0.1, -5);
9658 assert!(matches!(
9659 pt.mode(),
9660 QuantMode::PerTensor { scale, zero_point } if scale == 0.1 && zero_point == -5
9661 ));
9662
9663 let pts = Quantization::per_tensor_symmetric(0.05);
9664 assert!(matches!(
9665 pts.mode(),
9666 QuantMode::PerTensorSymmetric { scale } if scale == 0.05
9667 ));
9668
9669 let pc = Quantization::per_channel(vec![0.1, 0.2], vec![0, -1], 2).unwrap();
9670 assert!(matches!(pc.mode(), QuantMode::PerChannel { axis: 2, .. }));
9671
9672 let pcs = Quantization::per_channel_symmetric(vec![0.1, 0.2], 0).unwrap();
9673 assert!(matches!(
9674 pcs.mode(),
9675 QuantMode::PerChannelSymmetric { axis: 0, .. }
9676 ));
9677 }
9678
9679 #[test]
9680 fn test_tensor_quantization_roundtrip_integer() {
9681 let mut t = Tensor::<i8>::new(&[2, 3, 4], Some(TensorMemory::Mem), None).unwrap();
9682 assert!(t.quantization().is_none());
9683 t.set_quantization(Quantization::per_tensor(0.1, -5))
9684 .unwrap();
9685 let q = t.quantization().unwrap();
9686 assert_eq!(q.scale(), &[0.1]);
9687 t.clear_quantization();
9688 assert!(t.quantization().is_none());
9689 }
9690
9691 #[test]
9692 fn test_tensor_with_quantization_builder() {
9693 let t = Tensor::<i8>::new(&[4, 4], Some(TensorMemory::Mem), None)
9694 .unwrap()
9695 .with_quantization(Quantization::per_tensor_symmetric(0.05))
9696 .unwrap();
9697 assert!(t.quantization().is_some());
9698 }
9699
9700 #[test]
9701 fn test_tensor_dyn_quantization_float_arm_returns_none() {
9702 let t = Tensor::<f32>::new(&[2, 2], Some(TensorMemory::Mem), None).unwrap();
9703 let td = TensorDyn::F32(t);
9704 assert!(td.quantization().is_none());
9705 }
9706
9707 #[test]
9708 fn test_tensor_dyn_set_quantization_float_arm_errors() {
9709 let t = Tensor::<f32>::new(&[2, 2], Some(TensorMemory::Mem), None).unwrap();
9710 let mut td = TensorDyn::F32(t);
9711 let err = td
9712 .set_quantization(Quantization::per_tensor(0.1, 0))
9713 .unwrap_err();
9714 // float path returns a QuantizationInvalid error.
9715 assert!(matches!(err, Error::QuantizationInvalid { .. }));
9716 }
9717
9718 /// Compile-time type gate — calling `Tensor::<f32>::quantization()` must
9719 /// fail to compile (the `IntegerType` trait bound is not satisfied by
9720 /// `f32`). This doctest anchors the invariant.
9721 ///
9722 /// ```compile_fail
9723 /// use edgefirst_tensor::{Tensor, TensorMemory};
9724 /// let t = Tensor::<f32>::new(&[2, 2], Some(TensorMemory::Mem), None).unwrap();
9725 /// let _ = t.quantization(); // compile error: f32 not IntegerType
9726 /// ```
9727 fn _compile_fail_doctest_anchor() {}
9728
9729 // Any test that cares about the fd count must grab it exclusively.
9730 // Any tests which modifies the fd count by opening or closing fds must grab it
9731 // shared.
9732 pub static FD_LOCK: RwLock<()> = RwLock::new(());
9733
9734 /// Take [`FD_LOCK`] shared, tolerating poisoning.
9735 ///
9736 /// A test that panics while holding the lock poisons it, and every later
9737 /// `.unwrap()` on it panics too — one real failure turns into a handful
9738 /// of unrelated ones that bury it. The guarded data is `()`, so a panic
9739 /// cannot have left an invariant broken and there is nothing to protect
9740 /// against by refusing the lock.
9741 pub fn fd_lock_shared() -> std::sync::RwLockReadGuard<'static, ()> {
9742 FD_LOCK.read().unwrap_or_else(PoisonError::into_inner)
9743 }
9744
9745 /// Take [`FD_LOCK`] exclusively. Poison-tolerant, see [`fd_lock_shared`].
9746 ///
9747 /// Linux-only, like all four of its callers: the fd-leak tests read fd
9748 /// state out of `/proc`, which no other platform provides. Ungated, it
9749 /// is dead code on macOS and the lint denies that.
9750 /// [`fd_lock_shared`] stays available everywhere — macOS tests declare
9751 /// fd activity too.
9752 #[cfg(target_os = "linux")]
9753 pub fn fd_lock_exclusive() -> std::sync::RwLockWriteGuard<'static, ()> {
9754 FD_LOCK.write().unwrap_or_else(PoisonError::into_inner)
9755 }
9756
9757 /// `/proc` fd-target fragments identifying a POSIX shared memory
9758 /// segment. `shm_open` then `shm_unlink` leaves the target rendered as
9759 /// `/dev/shm/<name> (deleted)`.
9760 #[cfg(target_os = "linux")]
9761 const SHM_FD_TARGETS: &[&str] = &["/dev/shm/"];
9762
9763 /// `/proc` fd-target fragments identifying the fds a DMA tensor owns.
9764 ///
9765 /// Both the exported buffer and the heap it came from are listed, so a
9766 /// leak of either is still caught — narrowing this to the buffer alone
9767 /// would silently stop covering leaked heap handles.
9768 #[cfg(target_os = "linux")]
9769 const DMA_FD_TARGETS: &[&str] = &["dmabuf", "dma_heap"];
9770
9771 /// How many of this process's open fds point at an object whose `/proc`
9772 /// target contains `needle`.
9773 ///
9774 /// The fd-leak tests deliberately do NOT compare total fd counts. Under
9775 /// `cargo test` every test in this crate runs as a thread in one
9776 /// process, and roughly twenty of them create shm or DMA tensors —
9777 /// while anything else may open a file, a socket, or a pipe — so an
9778 /// unrelated test opening or closing an fd inside the measurement
9779 /// window moved the total and failed the assertion. Worse, the panic
9780 /// poisoned [`FD_LOCK`] and took four more tests down with it, so the
9781 /// visible failure had little to do with the cause.
9782 ///
9783 /// Counting only fds that name the object under test makes the
9784 /// measurement immune to every unrelated fd, while still catching the
9785 /// leak each test exists to catch. (CI runs nextest, one process per
9786 /// test, which is why this only ever flaked when run locally.)
9787 #[cfg(target_os = "linux")]
9788 fn open_fd_count(needles: &[&str]) -> usize {
9789 let proc = procfs::process::Process::myself()
9790 .expect("Failed to get current process using /proc/self");
9791 proc.fd()
9792 .expect("Failed to list open file descriptors")
9793 .filter_map(Result::ok)
9794 .filter(|entry| {
9795 let target = format!("{:?}", entry.target);
9796 needles.iter().any(|needle| target.contains(needle))
9797 })
9798 .count()
9799 }
9800
9801 /// Test that DMA is NOT available on non-Linux platforms.
9802 /// This verifies the cross-platform behavior of is_dma_available().
9803 #[test]
9804 #[cfg(not(target_os = "linux"))]
9805 fn test_dma_not_available_on_non_linux() {
9806 assert!(
9807 !is_dma_available(),
9808 "DMA memory allocation should NOT be available on non-Linux platforms"
9809 );
9810 }
9811
9812 #[test]
9813 fn colorimetry_defaults_none_and_roundtrips_without_auto_fill() {
9814 use crate::{ColorEncoding, ColorRange, Colorimetry, PixelFormat, TensorMemory};
9815 let mut t = Tensor::<u8>::image(
9816 1280,
9817 720,
9818 PixelFormat::Nv12,
9819 Some(TensorMemory::Mem),
9820 crate::CpuAccess::ReadWrite,
9821 )
9822 .unwrap();
9823 assert_eq!(t.colorimetry(), None); // default undefined
9824 let c = Colorimetry::default()
9825 .with_encoding(ColorEncoding::Bt709)
9826 .with_range(ColorRange::Limited);
9827 t.set_colorimetry(Some(c));
9828 assert_eq!(t.colorimetry(), Some(c));
9829 // configure_image must NOT touch colorimetry
9830 t.configure_image(640, 480, PixelFormat::Grey).unwrap();
9831 assert_eq!(t.colorimetry(), Some(c));
9832 }
9833
9834 #[test]
9835 fn configure_image_within_capacity() {
9836 let mut t = Tensor::<u8>::image_with_capacity(
9837 640,
9838 480,
9839 PixelFormat::Rgb,
9840 None,
9841 crate::CpuAccess::ReadWrite,
9842 )
9843 .unwrap();
9844 t.configure_image(320, 240, PixelFormat::Nv12).unwrap();
9845 assert_eq!(t.format(), Some(PixelFormat::Nv12));
9846 assert_eq!(t.width(), Some(320));
9847 assert_eq!(t.height(), Some(240));
9848 assert_eq!(t.shape(), &[360, 320]); // 240*3/2
9849 }
9850
9851 #[test]
9852 fn configure_image_too_large_errors() {
9853 let mut t = Tensor::<u8>::image_with_capacity(
9854 64,
9855 64,
9856 PixelFormat::Grey,
9857 None,
9858 crate::CpuAccess::ReadWrite,
9859 )
9860 .unwrap();
9861 let err = t
9862 .configure_image(1920, 1080, PixelFormat::Nv12)
9863 .unwrap_err();
9864 assert!(matches!(err, Error::InsufficientCapacity { .. }));
9865 }
9866
9867 /// A reused max-sized IOSurface pool keeps its physical `bytesPerRow` when
9868 /// reconfigured to a smaller logical image (physical-grid / logical-ROI
9869 /// decoupling), instead of collapsing to the frame's natural row stride.
9870 #[test]
9871 #[cfg(target_os = "macos")]
9872 fn configure_image_preserves_iosurface_physical_stride() {
9873 // Declares this test as an fd-opener; see FD_LOCK.
9874 let _lock = crate::tests::fd_lock_shared();
9875 // Pool: GREY/R8 IOSurface 100 wide → bytesPerRow padded to 128.
9876 let mut pool = Tensor::<u8>::image(
9877 100,
9878 64,
9879 PixelFormat::Grey,
9880 Some(TensorMemory::DmaBuf),
9881 crate::CpuAccess::ReadWrite,
9882 )
9883 .unwrap();
9884 let pitch = pool.effective_row_stride().unwrap();
9885 assert!(
9886 pitch >= 128 && pitch.is_multiple_of(64),
9887 "padded bytesPerRow, got {pitch}"
9888 );
9889
9890 // Reconfigure to a smaller NV12 frame; the physical pitch must survive
9891 // (natural would be 32, but the surface stride is the 128-padded pitch).
9892 pool.configure_image(32, 16, PixelFormat::Nv12).unwrap();
9893 assert_eq!(pool.format(), Some(PixelFormat::Nv12));
9894 assert_eq!(pool.width(), Some(32));
9895 assert_eq!(pool.height(), Some(16));
9896 assert_eq!(
9897 pool.effective_row_stride(),
9898 Some(pitch),
9899 "configure_image must preserve the IOSurface physical bytesPerRow"
9900 );
9901
9902 // Reconfigure again to NV24 — pitch still preserved.
9903 pool.configure_image(32, 16, PixelFormat::Nv24).unwrap();
9904 assert_eq!(pool.effective_row_stride(), Some(pitch));
9905 }
9906
9907 /// `configure_image` on a Mem backing reconfigures to the format's
9908 /// **64-byte-aligned** row stride (the odd-dim contract: every image tensor
9909 /// carries a 64-aligned `row_stride`). For NV12 32×16 the minimum is
9910 /// `even(32)=32`, rounded up to the 64-byte alignment → 64. The capacity
9911 /// (64×64×4 RGBA = 16 KiB) easily holds the 24×64 = 1.5 KiB NV12 layout.
9912 #[test]
9913 fn configure_image_mem_aligns_stride() {
9914 let mut t = Tensor::<u8>::image_with_capacity(
9915 64,
9916 64,
9917 PixelFormat::Rgba,
9918 Some(TensorMemory::Mem),
9919 crate::CpuAccess::ReadWrite,
9920 )
9921 .unwrap();
9922 t.configure_image(32, 16, PixelFormat::Nv12).unwrap();
9923 let s = t.effective_row_stride().unwrap();
9924 assert_eq!(s % 64, 0, "stride must be 64-aligned");
9925 assert!(s >= 32, "stride must cover the even-width minimum");
9926 assert_eq!(s, 64);
9927 }
9928
9929 #[test]
9930 fn strided_mem_tensor_cpu_maps_full_padded_buffer() {
9931 // A packed RGBA image with row padding (GPU-pitch style): logical width
9932 // 8 px (32 B/row) but a 48-byte row stride. Over-allocate capacity (for
9933 // 16 px), narrow the logical width, then record the padded stride.
9934 // Previously `map()` rejected this on non-Linux with
9935 // "DMA backing is Linux-only"; HAL-owned Mem is now mappable.
9936 let mut t = Tensor::<u8>::image_with_capacity(
9937 16,
9938 3,
9939 PixelFormat::Rgba,
9940 Some(TensorMemory::Mem),
9941 crate::CpuAccess::ReadWrite,
9942 )
9943 .unwrap(); // capacity 3 × 16 × 4 = 192 B
9944 t.configure_image(8, 3, PixelFormat::Rgba).unwrap(); // logical [3, 8, 4] = 96 B
9945 t.set_row_stride(48).unwrap(); // padded stride (>= 32 B min)
9946
9947 let map = t.map().expect("strided Mem tensor should CPU-map");
9948 // Full padded buffer (stride 48 × 3 rows = 144 B), not the 96 B logical
9949 // view — callers iterate rows via `effective_row_stride()`.
9950 assert_eq!(map.as_slice().len(), 144);
9951 // Logical shape is still reported for shape-aware consumers.
9952 assert_eq!(map.shape(), &[3, 8, 4]);
9953 }
9954
9955 #[test]
9956 fn strided_mem_tensor_over_capacity_errors() {
9957 // Stride larger than the allocation: 64 B × 3 rows = 192 B > 96 B cap.
9958 let mut t = Tensor::<u8>::new(&[3, 8, 4], Some(TensorMemory::Mem), None).unwrap();
9959 t.set_format(PixelFormat::Rgba).unwrap();
9960 t.set_row_stride(64).unwrap();
9961 assert!(matches!(t.map(), Err(Error::InsufficientCapacity { .. })));
9962 }
9963
9964 /// Test that SHM memory allocation is available and usable on Unix systems.
9965 /// This is a basic functional test; Linux has additional FD leak tests using procfs.
9966 #[test]
9967 #[cfg(unix)]
9968 fn test_shm_available_and_usable() {
9969 // Declares this test as an fd-opener; see FD_LOCK.
9970 let _lock = crate::tests::fd_lock_shared();
9971 // "Unix" is not one platform here: Android is unix but its bionic libc
9972 // has no POSIX shared memory, so allocation legitimately reports
9973 // NotImplemented (see shm.rs module docs). Asserting availability
9974 // outright fails there for a documented reason.
9975 if !is_shm_available() {
9976 crate::test_support::report_skip(
9977 "test_shm_available_and_usable - SHM allocation \
9978 unavailable on this platform",
9979 );
9980 return;
9981 }
9982
9983 // Create a tensor with SHM backing
9984 let tensor = Tensor::<u8>::new(&[100, 100], Some(TensorMemory::Shm), None)
9985 .expect("Failed to create SHM tensor");
9986
9987 // Verify we can map and write to it
9988 let mut map = tensor.map().expect("Failed to map SHM tensor");
9989 map.as_mut_slice().fill(0xAB);
9990
9991 // Verify the data was written correctly
9992 assert!(
9993 map.as_slice().iter().all(|&b| b == 0xAB),
9994 "SHM tensor data should be writable and readable"
9995 );
9996 }
9997
9998 // =========================================================================
9999 // packed_rgba16f_layout — host-runnable geometry unit tests (TDD)
10000 // =========================================================================
10001
10002 #[test]
10003 fn packed_rgba16f_layout_planar_rgb_f16() {
10004 let layout =
10005 packed_rgba16f_layout(PixelFormat::PlanarRgb, DType::F16, 640, 640).expect("Some");
10006 assert_eq!(layout.surface_w, 160);
10007 assert_eq!(layout.surface_h, 1920);
10008 assert_eq!(layout.bytes_per_texel, 8);
10009 assert_eq!(layout.pitch, 1280);
10010 }
10011
10012 #[test]
10013 fn packed_rgba16f_layout_planar_rgba_f16() {
10014 let layout =
10015 packed_rgba16f_layout(PixelFormat::PlanarRgba, DType::F16, 640, 640).expect("Some");
10016 assert_eq!(layout.surface_w, 160);
10017 assert_eq!(layout.surface_h, 2560); // 4 planes
10018 assert_eq!(layout.bytes_per_texel, 8);
10019 assert_eq!(layout.pitch, 1280);
10020 }
10021
10022 #[test]
10023 fn packed_rgba16f_layout_rejects_misaligned() {
10024 assert!(packed_rgba16f_layout(PixelFormat::PlanarRgb, DType::F16, 642, 640).is_none());
10025 }
10026
10027 #[test]
10028 fn packed_rgba16f_layout_rejects_non_f16() {
10029 // Non-F16 dtype with planar RGB
10030 assert!(packed_rgba16f_layout(PixelFormat::PlanarRgb, DType::U8, 640, 640).is_none());
10031 // Non-planar format with F32
10032 assert!(packed_rgba16f_layout(PixelFormat::Rgb, DType::F32, 640, 640).is_none());
10033 // Packed Rgba with F16 is not a planar format → None
10034 assert!(packed_rgba16f_layout(PixelFormat::Rgba, DType::F16, 640, 640).is_none());
10035 }
10036
10037 #[test]
10038 fn cuda_map_fast_fails_to_none_without_handle() {
10039 let t = Tensor::<f32>::new(&[4], Some(TensorMemory::Mem), None).unwrap();
10040 assert!(t.cuda().is_none());
10041 assert!(t.cuda_map().is_none()); // pure local check, no GL routing
10042 }
10043
10044 #[test]
10045 fn cuda_returns_none_without_handle() {
10046 // A plain Mem-backed tensor has no CUDA handle attached.
10047 let t = Tensor::<f32>::new(&[2, 2], Some(TensorMemory::Mem), None).unwrap();
10048 assert!(t.cuda().is_none(), "no CUDA handle on a Mem tensor");
10049 assert!(t.cuda_map().is_none(), "fast-fail map → None");
10050 }
10051
10052 #[test]
10053 fn cuda_map_then_host_map_fallback() {
10054 // The documented client pattern: try cuda_map() first; when it is None
10055 // (no CUDA handle — the case for a plain Mem tensor), fall back to map().
10056 let t = Tensor::<f32>::new(&[2, 2], Some(TensorMemory::Mem), None).unwrap();
10057 // Bind to a named variable so the CudaMap guard (and its borrow of `t`)
10058 // is dropped at the end of this statement, before the else branch borrows `t` again.
10059 let cuda = t.cuda_map();
10060 if let Some(_c) = cuda {
10061 // On a CUDA-registered tensor we'd use the device ptr here.
10062 unreachable!("a Mem tensor has no CUDA handle");
10063 } else {
10064 let host = t.map().expect("host map fallback must succeed");
10065 // TensorMapTrait::len() returns the element count (not bytes).
10066 assert_eq!(host.len(), 4); // 2*2 f32 elements
10067 }
10068 }
10069
10070 // -------------------------------------------------------------------------
10071 // Tensor::from_foreign — public API tests at the Tensor<T> layer.
10072 //
10073 // The low-level MemTensor::from_foreign mechanics (owner-drop, view sharing)
10074 // are covered in mem.rs. These tests exercise the Tensor<T> guard paths
10075 // (null ptr, empty shape, size overflow) and the basic wrap+readback
10076 // contract, confirming the public unsafe API wires through correctly.
10077 // -------------------------------------------------------------------------
10078
10079 #[test]
10080 fn from_foreign_valid_wrap_and_readback() {
10081 // The canonical CUDA zero-copy shape: wrap a caller allocation as a
10082 // Mem tensor and verify the tensor reads the exact same bytes.
10083 let mut buf: Vec<f32> = vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0];
10084 let ptr = buf.as_mut_ptr();
10085 let t = unsafe { Tensor::<f32>::from_foreign(ptr, &[2, 3], None, Some("test_foreign")) }
10086 .expect("valid from_foreign must succeed");
10087 assert_eq!(t.shape(), &[2, 3]);
10088 assert_eq!(t.memory(), TensorMemory::Mem);
10089 assert_eq!(t.name(), "test_foreign");
10090 let m = t.map().unwrap();
10091 // The tensor is a zero-copy borrow — it sees the caller's data.
10092 assert_eq!(m.as_slice(), &[1.0_f32, 2.0, 3.0, 4.0, 5.0, 6.0]);
10093 }
10094
10095 #[test]
10096 fn from_foreign_write_visible_in_caller_allocation() {
10097 // Writes through the tensor's map land in the caller's buffer (zero-copy).
10098 let mut buf: Vec<u8> = vec![0u8; 6];
10099 let ptr = buf.as_mut_ptr();
10100 let t = unsafe { Tensor::<u8>::from_foreign(ptr, &[2, 3], None, None) }.unwrap();
10101 {
10102 let mut m = t.map().unwrap();
10103 m.as_mut_slice().copy_from_slice(&[10, 20, 30, 40, 50, 60]);
10104 }
10105 drop(t);
10106 // Mutations are visible in the original Vec — same physical buffer.
10107 assert_eq!(buf, vec![10, 20, 30, 40, 50, 60]);
10108 }
10109
10110 #[test]
10111 fn from_foreign_rejects_null_ptr() {
10112 let err = unsafe { Tensor::<u8>::from_foreign(std::ptr::null_mut(), &[4], None, None) }
10113 .unwrap_err();
10114 assert!(
10115 matches!(err, Error::InvalidArgument(ref m) if m.contains("non-null")),
10116 "expected InvalidArgument(non-null), got {err:?}"
10117 );
10118 }
10119
10120 #[test]
10121 fn from_foreign_rejects_empty_shape() {
10122 let mut dummy: u8 = 0;
10123 let err = unsafe { Tensor::<u8>::from_foreign(&mut dummy, &[], None, None) }.unwrap_err();
10124 assert!(
10125 matches!(err, Error::InvalidSize(0)),
10126 "expected InvalidSize(0) for empty shape, got {err:?}"
10127 );
10128 }
10129
10130 #[test]
10131 fn from_foreign_rejects_overflow_shape() {
10132 // Two dimensions whose product overflows usize — the overflow guard must
10133 // fire before any pointer arithmetic is attempted.
10134 let mut dummy: u8 = 0;
10135 let huge = [usize::MAX / 2 + 1, 2];
10136 let err = unsafe { Tensor::<u8>::from_foreign(&mut dummy, &huge, None, None) }.unwrap_err();
10137 assert!(
10138 matches!(err, Error::InvalidArgument(ref m) if m.contains("overflow")),
10139 "expected InvalidArgument(overflow), got {err:?}"
10140 );
10141 }
10142
10143 #[test]
10144 fn from_foreign_owner_keeps_allocation_alive() {
10145 // When `owner` is `Some`, dropping the Tensor must not free the backing
10146 // before the owner is also gone — the owner's Drop fires on last ref.
10147 use std::sync::atomic::{AtomicBool, Ordering};
10148 let flag = std::sync::Arc::new(AtomicBool::new(false));
10149 let flag2 = flag.clone();
10150 struct Guard(std::sync::Arc<AtomicBool>);
10151 impl Drop for Guard {
10152 fn drop(&mut self) {
10153 self.0.store(true, Ordering::SeqCst);
10154 }
10155 }
10156 let mut buf: Vec<u32> = vec![42u32; 4];
10157 let ptr = buf.as_mut_ptr();
10158 let owner: ForeignOwner = Box::new(Guard(flag2));
10159 let t = unsafe { Tensor::<u32>::from_foreign(ptr, &[4], Some(owner), None) }.unwrap();
10160 // Map co-owns the backing Arc; the owner must stay alive while the map lives.
10161 let m = t.map().unwrap();
10162 assert_eq!(m.as_slice()[0], 42);
10163 drop(t); // tensor dropped while map is still live
10164 assert!(
10165 !flag.load(Ordering::SeqCst),
10166 "owner must not drop while a map shares the backing"
10167 );
10168 drop(m);
10169 assert!(
10170 flag.load(Ordering::SeqCst),
10171 "owner Drop must fire when the last Arc reference is released"
10172 );
10173 }
10174}