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 caller's
1182 /// maps, pins and sync brackets, and the HAL's own CPU work (the CPU
1183 /// image backend's converts and draws, codec decodes, decoder reads and
1184 /// writes). Hardware access (ISP, codec, GPU, NPU, G2D, a compositor)
1185 /// needs no declaration and is always allowed.
1186 ///
1187 /// The HAL **enforces** it. A CPU map, pin or sync whose direction the
1188 /// declaration does not cover fails with [`Error::CpuAccessDenied`] on
1189 /// every backend, before the memory is touched. The HAL's own CPU passes
1190 /// go through the same check, so, for example, a CPU draw that would
1191 /// blend over a `Write` image, or a CPU convert that would read a `None`
1192 /// source, is refused instead of running.
1193 ///
1194 /// | Level | CPU may | Memory selected |
1195 /// |---|---|---|
1196 /// | `None` | nothing | Linux DMA-BUF: uncached CMA where available. Android: no CPU usage bits, so tile compression stays possible. |
1197 /// | `Read` | read | cached; read-direction syncs only |
1198 /// | `Write` | write whole passes | Linux DMA-BUF: uncached CMA where available, written write-combined |
1199 /// | `ReadWrite` (default) | read and write | cached |
1200 ///
1201 /// **Why uncached.** A device that maps a cached buffer has to clean and
1202 /// invalidate the CPU cache over all of it on every access: a compositor
1203 /// blitting a client's buffer every frame pays that for the whole buffer
1204 /// each time. An uncached buffer needs no cache maintenance. CPU reads
1205 /// of uncached memory are slow, which is why only declarations that
1206 /// exclude CPU reads select it; with the declaration enforced, nothing
1207 /// reads such a buffer on the CPU. On Linux a `None` or `Write` DMA-BUF
1208 /// comes from `/dev/dma_heap/linux,cma-uncached` when the system has it,
1209 /// and from the cached heaps otherwise. A write-only pass must write the
1210 /// whole buffer (see [`TensorTrait::map_for_write`]).
1211 ///
1212 /// On Android the `ReadWrite` default asks gralloc for
1213 /// `CPU_READ_OFTEN | CPU_WRITE_OFTEN`, which may give linear, CPU-cached
1214 /// memory. Pass `None` for a buffer that is never mapped on the CPU, so
1215 /// gralloc can keep a tiled or compressed allocation.
1216 ///
1217 /// # Examples
1218 ///
1219 /// A CPU-drawn overlay: written whole every frame, never read back.
1220 ///
1221 /// ```rust,no_run
1222 /// use edgefirst_tensor::{
1223 /// CpuAccess, Error, PixelFormat, Tensor, TensorMapTrait, TensorMemory, TensorTrait,
1224 /// };
1225 ///
1226 /// let overlay = Tensor::<u8>::image(
1227 /// 640,
1228 /// 480,
1229 /// PixelFormat::Rgba,
1230 /// Some(TensorMemory::DmaBuf),
1231 /// CpuAccess::Write,
1232 /// )?;
1233 /// overlay.map_for_write()?.as_mut_slice().fill(0);
1234 /// assert!(matches!(overlay.map_read(), Err(Error::CpuAccessDenied { .. })));
1235 /// # Ok::<(), Error>(())
1236 /// ```
1237 ///
1238 /// [`CpuAccess::code`] is the wire code the C header's `ef_cpu_access`
1239 /// enumerators are asserted against.
1240 #[derive(Default)]
1241 pub enum CpuAccess {
1242 /// Hardware-only buffer: no CPU access. Required for vendor tile
1243 /// compression on platforms with compressed layouts; uncached CMA
1244 /// for a Linux DMA-BUF where available.
1245 None = 0, "none", NONE,
1246 /// CPU reads only (verification, CPU consumers) — cached mapping.
1247 Read = 1, "read", READ,
1248 /// CPU writes only, whole-buffer passes (decode targets, CPU-drawn
1249 /// overlays) — write-combined mapping; uncached CMA for a Linux
1250 /// DMA-BUF where available.
1251 Write = 2, "write", WRITE,
1252 /// CPU reads and writes (the default).
1253 #[default]
1254 ReadWrite = 3, "readwrite", READ_WRITE,
1255 }
1256 // Hidden like `dtype_wire`: emission plumbing, not a second documented API.
1257 #[doc(hidden)]
1258 pub mod cpu_access_wire;
1259}
1260
1261impl CpuAccess {
1262 /// The declaration with exactly the given directions.
1263 pub(crate) const fn from_directions(read: bool, write: bool) -> CpuAccess {
1264 match (read, write) {
1265 (true, true) => CpuAccess::ReadWrite,
1266 (true, false) => CpuAccess::Read,
1267 (false, true) => CpuAccess::Write,
1268 (false, false) => CpuAccess::None,
1269 }
1270 }
1271
1272 /// Whether this declaration includes CPU reads.
1273 pub fn reads(self) -> bool {
1274 matches!(self, CpuAccess::Read | CpuAccess::ReadWrite)
1275 }
1276
1277 /// Whether this declaration includes CPU writes.
1278 pub fn writes(self) -> bool {
1279 matches!(self, CpuAccess::Write | CpuAccess::ReadWrite)
1280 }
1281
1282 /// Whether this declaration covers `requested` (every direction the
1283 /// request needs is declared).
1284 pub fn covers(self, requested: CpuAccess) -> bool {
1285 (!requested.reads() || self.reads()) && (!requested.writes() || self.writes())
1286 }
1287
1288 /// The declaration covering every direction of both `self` and `other`.
1289 ///
1290 /// ```
1291 /// use edgefirst_tensor::CpuAccess;
1292 /// assert_eq!(CpuAccess::Read.union(CpuAccess::Write), CpuAccess::ReadWrite);
1293 /// assert_eq!(CpuAccess::None.union(CpuAccess::Write), CpuAccess::Write);
1294 /// ```
1295 pub fn union(self, other: CpuAccess) -> CpuAccess {
1296 Self::from_directions(
1297 self.reads() || other.reads(),
1298 self.writes() || other.writes(),
1299 )
1300 }
1301
1302 /// The declaration covering only the directions both `self` and `other`
1303 /// cover.
1304 ///
1305 /// ```
1306 /// use edgefirst_tensor::CpuAccess;
1307 /// assert_eq!(CpuAccess::ReadWrite.intersection(CpuAccess::Write), CpuAccess::Write);
1308 /// assert_eq!(CpuAccess::Read.intersection(CpuAccess::Write), CpuAccess::None);
1309 /// ```
1310 pub fn intersection(self, other: CpuAccess) -> CpuAccess {
1311 Self::from_directions(
1312 self.reads() && other.reads(),
1313 self.writes() && other.writes(),
1314 )
1315 }
1316
1317 /// `Ok` when this declaration covers `requested`, otherwise
1318 /// [`Error::CpuAccessDenied`] naming both.
1319 ///
1320 /// ```
1321 /// use edgefirst_tensor::{CpuAccess, Error};
1322 /// assert!(CpuAccess::ReadWrite.require(CpuAccess::Write).is_ok());
1323 /// assert!(matches!(
1324 /// CpuAccess::Read.require(CpuAccess::Write),
1325 /// Err(Error::CpuAccessDenied { .. })
1326 /// ));
1327 /// ```
1328 pub fn require(self, requested: CpuAccess) -> Result<()> {
1329 if self.covers(requested) {
1330 Ok(())
1331 } else {
1332 Err(Error::CpuAccessDenied {
1333 declared: self,
1334 requested,
1335 })
1336 }
1337 }
1338
1339 /// [`require`](Self::require) for a map, pin or sync, where
1340 /// `CpuAccess::None` is no direction at all: [`Error::InvalidArgument`]
1341 /// naming `what`.
1342 pub(crate) fn require_direction(self, requested: CpuAccess, what: &str) -> Result<()> {
1343 if requested == CpuAccess::None {
1344 return Err(Error::InvalidArgument(format!(
1345 "{what}: CpuAccess::None is not a CPU access direction"
1346 )));
1347 }
1348 self.require(requested)
1349 }
1350}
1351
1352/// Requested tile-compression behavior for an image allocation (set via
1353/// [`ImageDesc::with_compression`]).
1354///
1355/// Vendor GPUs store textures in proprietary compressed tile layouts
1356/// (UBWC, AFBC, PVRIC, DCC) that cut memory bandwidth; eligibility
1357/// requires a hardware-only buffer ([`CpuAccess::None`]) because CPU
1358/// mapping pins the layout linear. The request records best knowledge on
1359/// the tensor ([`Tensor::compression`]); it never changes what bytes a
1360/// consumer sees through the GPU/NPU.
1361#[non_exhaustive]
1362#[derive(Debug, Clone, Copy, PartialEq, Eq)]
1363pub enum Compression {
1364 /// Let the platform use its native scheme when the format is
1365 /// eligible; otherwise allocate linear and count the fallback
1366 /// ([`compression_fallback_count`]). The right default for pipelines
1367 /// that want the bandwidth win without portability failures.
1368 Any,
1369 /// Require one specific scheme: allocation fails with
1370 /// [`Error::InvalidArgument`] when the device's scheme differs and
1371 /// [`Error::NotImplemented`] on platforms without vendor tile
1372 /// compression. For consumers whose ABI names a layout (e.g. a
1373 /// QNN context binary declaring UBWC inputs).
1374 Scheme(CompressionScheme),
1375}
1376
1377// Declared through `ef_vocabulary!` (see `vocabulary.rs`). This vocabulary
1378// gained a wire numbering when `ef_tensor_compression` was added: the
1379// dynamic backend reads a scheme back across the C boundary, so the enum
1380// needs codes that cannot drift from the C enumerators. Hand-writing that
1381// mapping in `edgefirst-tensor-capi` was the alternative and it is strictly
1382// worse -- `#[non_exhaustive]` makes a cross-crate `match` require a
1383// wildcard arm, so a scheme added here would silently map to "linear"
1384// there. Generated from one declaration, it cannot.
1385//
1386// Code 0 is deliberately unassigned and reserved for "linear" -- the state
1387// this crate spells `Option::None`, which has no variant to carry a code of
1388// its own. The C side names it `EF_COMPRESSION_NONE`.
1389crate::ef_vocabulary! {
1390 /// Vendor tile-compression schemes the HAL recognizes and records.
1391 ///
1392 /// Unrecognized platforms record `None` (linear) — there is deliberately
1393 /// no `Unknown` variant; a scheme is only recorded when the vendor is
1394 /// positively identified.
1395 #[non_exhaustive]
1396 pub enum CompressionScheme {
1397 /// Qualcomm Adreno Universal Bandwidth Compression.
1398 Ubwc = 1, "ubwc", UBWC,
1399 /// Arm Mali/Immortalis Framebuffer Compression.
1400 Afbc = 2, "afbc", AFBC,
1401 /// Imagination PowerVR Image Compression (Google Tensor G5+).
1402 Pvric = 3, "pvric", PVRIC,
1403 /// Samsung Xclipse (AMD RDNA) Delta Color Compression.
1404 Dcc = 4, "dcc", DCC,
1405 }
1406 // `#[doc(hidden)]` for the same reason `dtype_wire` is: emission
1407 // plumbing for `const`-only consumers, not a second documented API.
1408 #[doc(hidden)]
1409 pub mod compression_wire;
1410}
1411
1412crate::ef_vocabulary! {
1413 /// Whether a tensor's memory is known to be physically contiguous
1414 /// ([`Tensor::contiguity`]).
1415 ///
1416 /// Capture and other DMA engines without an IOMMU (the i.MX 8M Plus ISI,
1417 /// the vvcam ISP) can only write into physically contiguous memory.
1418 /// Only memory this crate allocated from a Linux DMA heap is known
1419 /// either way: the CMA heap is contiguous and the system heap is not.
1420 /// An imported fd, and every other kind of memory, is `Unknown`.
1421 #[non_exhaustive]
1422 pub enum Contiguity {
1423 /// Not known: an imported fd, or memory not allocated from a DMA heap.
1424 Unknown = 0, "unknown", UNKNOWN,
1425 /// Allocated from the CMA heap.
1426 Contiguous = 1, "contiguous", CONTIGUOUS,
1427 /// Allocated from the system heap.
1428 NonContiguous = 2, "non_contiguous", NON_CONTIGUOUS,
1429 }
1430 #[doc(hidden)]
1431 pub mod contiguity_wire;
1432}
1433
1434/// Count of image allocations that requested [`Compression::Any`] but
1435/// resolved to a linear layout (ineligible format/dtype, unrecognized
1436/// vendor, or a platform without vendor tile compression). Read via
1437/// [`compression_fallback_count`].
1438static COMPRESSION_FALLBACKS: AtomicU64 = AtomicU64::new(0);
1439
1440/// Number of [`Compression::Any`] requests since process start that
1441/// resolved to a linear layout instead of a vendor tile scheme. A
1442/// steady-state pipeline holds this flat after warmup; growth means an
1443/// allocation path keeps requesting compression it never gets.
1444pub fn compression_fallback_count() -> u64 {
1445 COMPRESSION_FALLBACKS.load(Ordering::Relaxed)
1446}
1447
1448/// Record a `Compression::Any` request resolving linear.
1449///
1450/// `static`-only: the sole caller is `Tensor::image_desc` (`impl<T>
1451/// Tensor<T>`), itself `#[cfg(not(edgefirst_tensor_dynamic))]`.
1452/// [`compression_fallback_count`] stays available under `dynamic` too --
1453/// it just always reads 0 there.
1454#[cfg(not(edgefirst_tensor_dynamic))]
1455pub(crate) fn note_compression_fallback(detail: &str) {
1456 COMPRESSION_FALLBACKS.fetch_add(1, Ordering::Relaxed);
1457 log::debug!("compression request fell back to linear: {detail}");
1458}
1459
1460/// Whether this platform can allocate `(format, dtype)` images in a
1461/// vendor tile-compressed layout. `true` requires an Android build, an
1462/// eligible format (RGBA8888 `u8`/`i8` initially), and a positively
1463/// identified GPU vendor; everywhere else the answer is `false` and
1464/// [`Compression::Any`] requests fall back to linear.
1465pub fn compression_support(format: PixelFormat, dtype: DType) -> bool {
1466 #[cfg(target_os = "android")]
1467 {
1468 crate::ahardwarebuffer_layout::compression_eligible(format, dtype)
1469 && crate::ahardwarebuffer::device_compression_scheme().is_some()
1470 }
1471 #[cfg(not(target_os = "android"))]
1472 {
1473 let _ = (format, dtype);
1474 false
1475 }
1476}
1477
1478/// Declarative image-allocation request — the full-featured front door
1479/// for image tensors ([`TensorDyn::image_desc`] and the image crate's
1480/// `ImageProcessor::create_image_desc`).
1481///
1482/// The classic constructors (`image`, `create_image`) cover the common
1483/// cases; the desc carries the optional requests — today the
1484/// [`Compression`] request — without another constructor-parameter
1485/// sweep. Fields are private and the builders consume/return by value,
1486/// so future options are non-breaking.
1487///
1488/// ```
1489/// use edgefirst_tensor::{Compression, CpuAccess, DType, ImageDesc, PixelFormat};
1490/// let desc = ImageDesc::new(640, 640, PixelFormat::Rgba, DType::U8)
1491/// .with_access(CpuAccess::None)
1492/// .with_compression(Compression::Any);
1493/// ```
1494#[derive(Debug, Clone)]
1495pub struct ImageDesc {
1496 width: usize,
1497 height: usize,
1498 format: PixelFormat,
1499 dtype: DType,
1500 memory: Option<TensorMemory>,
1501 access: CpuAccess,
1502 compression: Option<Compression>,
1503 contiguous: bool,
1504}
1505
1506impl ImageDesc {
1507 /// A new image request: auto-selected memory, [`CpuAccess::ReadWrite`],
1508 /// no compression request.
1509 pub fn new(width: usize, height: usize, format: PixelFormat, dtype: DType) -> Self {
1510 Self {
1511 width,
1512 height,
1513 format,
1514 dtype,
1515 memory: None,
1516 access: CpuAccess::default(),
1517 compression: None,
1518 contiguous: false,
1519 }
1520 }
1521
1522 /// Request a specific memory backing (`None` = auto-select).
1523 pub fn with_memory(mut self, memory: Option<TensorMemory>) -> Self {
1524 self.memory = memory;
1525 self
1526 }
1527
1528 /// Declare the CPU access (see [`CpuAccess`]). Any declaration other
1529 /// than `None` makes a compression request invalid, so a compression
1530 /// request must also call `with_access(CpuAccess::None)`.
1531 pub fn with_access(mut self, access: CpuAccess) -> Self {
1532 self.access = access;
1533 self
1534 }
1535
1536 /// Request a tile-compressed layout (see [`Compression`]).
1537 pub fn with_compression(mut self, compression: Compression) -> Self {
1538 self.compression = Some(compression);
1539 self
1540 }
1541
1542 /// Require physically contiguous memory (see [`Contiguity`]).
1543 ///
1544 /// The image is allocated from the Linux CMA DMA heap only. Allocation
1545 /// fails, naming the heap, when there is no CMA heap or it cannot fit
1546 /// the image, instead of falling back to the system heap or to another
1547 /// kind of memory. Memory must be auto-selected or
1548 /// [`TensorMemory::DmaBuf`]. Other platforms return
1549 /// [`Error::NotImplemented`].
1550 pub fn with_contiguous(mut self, contiguous: bool) -> Self {
1551 self.contiguous = contiguous;
1552 self
1553 }
1554
1555 /// Requested width in pixels.
1556 pub fn width(&self) -> usize {
1557 self.width
1558 }
1559
1560 /// Requested height in pixels.
1561 pub fn height(&self) -> usize {
1562 self.height
1563 }
1564
1565 /// Requested pixel format.
1566 pub fn format(&self) -> PixelFormat {
1567 self.format
1568 }
1569
1570 /// Requested element type.
1571 pub fn dtype(&self) -> DType {
1572 self.dtype
1573 }
1574
1575 /// Requested memory backing (`None` = auto-select).
1576 pub fn memory(&self) -> Option<TensorMemory> {
1577 self.memory
1578 }
1579
1580 /// Declared CPU access.
1581 pub fn access(&self) -> CpuAccess {
1582 self.access
1583 }
1584
1585 /// The compression request, if any.
1586 pub fn compression(&self) -> Option<Compression> {
1587 self.compression
1588 }
1589
1590 /// Whether physically contiguous memory is required.
1591 pub fn contiguous(&self) -> bool {
1592 self.contiguous
1593 }
1594}
1595
1596/// What kind of system key an identity was derived from. Participates in the
1597/// id so that, say, inode 7 and IOSurfaceID 7 cannot collide.
1598///
1599/// Invariant every variant must uphold: the key passed to
1600/// [`BufferIdentity::derived`] is unique among live buffers of that kind
1601/// **that own an allocation**. There is no "no identity" case — where no
1602/// system-wide key exists a process-local one does (a pointer, a GL name),
1603/// and a process-local key is exactly as sound as a system-wide one:
1604/// pointer/name reuse after free is the same ABA hazard as inode reuse, and
1605/// is mitigated the same way, by the retaining cache (see the design spec's
1606/// Task 6).
1607///
1608/// The allocation qualifier is load-bearing, not a hedge. A zero-element
1609/// tensor (`[0, num_protos]`, the tracker's "no detections this frame"
1610/// sentinel — see [`crate::mem`]) owns no allocation, so its `Vec` yields the
1611/// dangling well-aligned pointer `align_of::<T>()`. Every zero-element
1612/// `Tensor<f32>` in a process therefore shares one identity, and can even
1613/// collide across kinds with a small integer key from another variant.
1614///
1615/// That is safe rather than tolerated: an identity exists to answer "are
1616/// these the same bytes?", and buffers with no bytes have nothing to alias.
1617/// Nor can they reach an import cache. There are four GPU import paths —
1618/// a dma-buf fd on Linux, an IOSurface ref on macOS, an `AHardwareBuffer` on
1619/// Android, and an `ID3D11Texture2D` on Windows — and each takes only its own
1620/// object, so a tensor with no allocation reaches none of them. The Windows
1621/// path additionally keys on the texture the import itself holds a reference
1622/// to, and `edgefirst-image` **refuses** any texture whose identity is not
1623/// `D3d11Texture` over that same pointer — in every build, and on every
1624/// import including one served from its cache
1625/// (`GlPlatform::validate_import_identity`). That refusal exists because the
1626/// rule was once only written
1627/// down: the dynamic backend handed the image crate `HostPtr` identities over
1628/// recyclable handle addresses, and converts silently rendered into dropped
1629/// textures. Do not "fix" the zero-element case by minting a distinct id per
1630/// empty tensor: that reintroduces a counter, which is the state this whole
1631/// design exists to remove.
1632#[derive(Debug, Clone, Copy, PartialEq, Eq)]
1633pub enum IdentityKind {
1634 /// `(st_dev, st_ino)` of a dma-buf or memfd. Survives `dup`, and is the
1635 /// same value in every process that holds the buffer.
1636 DmaBuf,
1637 Shm,
1638 /// `IOSurfaceGetID`.
1639 IoSurface,
1640 /// `AHardwareBuffer_getId`, API 31+.
1641 AHardwareBufferId,
1642 /// The `AHardwareBuffer*` pointer, pre-API-31 where no id exists.
1643 /// Process-local: unique among live buffers here, meaningless elsewhere.
1644 AHardwareBufferPtr,
1645 /// An `ID3D11Texture2D*`. Process-local, like `AHardwareBufferPtr`, and
1646 /// the same pointer for every kind of texture tensor: one the HAL
1647 /// allocated, one wrapped from a caller's pointer, and one opened from a
1648 /// shared NT handle all key on the texture object they hold.
1649 ///
1650 /// **Both backends derive it from the same pointer, so their ids are
1651 /// equal byte for byte.** The static backend reads the `ComPtr` it owns
1652 /// (`d3d11::texture::tex_key`); the dynamic backend reads the same
1653 /// pointer back across the ABI with `ef_tensor_d3d11_texture`
1654 /// (`TensorDyn::derive_identity`). That equality is the point, not a
1655 /// coincidence: [`BufferIdentity::derived`] exists so two independently
1656 /// linked copies of this crate name one buffer identically, and
1657 /// `edgefirst-image` links the dynamic backend inside every `-capi` leaf
1658 /// while its own tests link the static one.
1659 ///
1660 /// **An `ef_tensor` handle's address must never be the key.** It is not a
1661 /// buffer at all, nothing the cache holds keeps it alive, and this
1662 /// library's allocator hands a freed handle's address to the very next
1663 /// one it mints — measured at 400 texture handles over 4 distinct
1664 /// addresses. The dynamic backend did exactly that on Windows until
1665 /// 2026-09-04, and a convert whose destination key collided rendered into
1666 /// the dropped texture while returning success.
1667 ///
1668 /// A handle import therefore does **not** compare equal to the tensor it
1669 /// was exported from, nor to a second import of the same handle:
1670 /// `OpenSharedResource1` mints a fresh `ID3D11Texture2D` per open. Keying
1671 /// on a handle value would make those compare equal, and would be unsound:
1672 /// Windows recycles a closed handle value straight into the next
1673 /// `DuplicateHandle`, so a consumer that caches on this id past the
1674 /// producing tensor's lifetime -- as the image crate's EGLImage cache
1675 /// does -- could serve one texture's entry for a later, unrelated import.
1676 /// A COM object's address cannot be recycled that way: a cache that holds
1677 /// a reference to the texture keeps the address taken for as long as its
1678 /// entry can be found. The cost is a cache miss per re-import instead of
1679 /// a wrong image.
1680 D3d11Texture,
1681 /// A host mapping address. Process-local.
1682 HostPtr,
1683 /// A GL buffer name. Valid only in the context that created it.
1684 Pbo,
1685}
1686
1687/// Unique identity for a tensor's underlying buffer.
1688///
1689/// Created fresh on every buffer allocation or import. The `id` is a cache
1690/// key, `kind` records what system key (if any) it was derived from, and the
1691/// `guard` is an `Arc<()>` whose weak references allow downstream caches to
1692/// detect when the buffer has been dropped.
1693#[derive(Debug, Clone)]
1694pub struct BufferIdentity {
1695 id: u64,
1696 kind: IdentityKind,
1697 guard: Arc<()>,
1698}
1699
1700impl BufferIdentity {
1701 /// Derive an identity from a key the operating system already guarantees.
1702 ///
1703 /// Deliberately NOT a counter: two independently-linked copies of this
1704 /// crate must derive the same id for the same buffer, or a consumer's
1705 /// import cache misses every frame. See the design spec, "The core finding".
1706 pub fn derived(kind: IdentityKind, key: u64) -> Self {
1707 // Mix the discriminant into the high bits so different kinds with the
1708 // same numeric key never collide.
1709 let id = ((kind as u64) << 56) ^ key;
1710 Self {
1711 id,
1712 kind,
1713 guard: Arc::new(()),
1714 }
1715 }
1716
1717 /// Unique identifier for this buffer. Changes when the buffer changes.
1718 pub fn id(&self) -> u64 {
1719 self.id
1720 }
1721
1722 /// What kind of system key this identity was derived from.
1723 pub fn kind(&self) -> IdentityKind {
1724 self.kind
1725 }
1726
1727 /// Returns a weak reference to the buffer guard. Goes dead when the
1728 /// owning Tensor is dropped (and no clones remain).
1729 pub fn weak(&self) -> Weak<()> {
1730 Arc::downgrade(&self.guard)
1731 }
1732}
1733
1734// `static`-only from here through `fs_magic` below: every use site lives in
1735// `impl<T> TensorStorage<T>` (`from_fd`'s DMA-BUF-vs-tmpfs classification),
1736// itself `#[cfg(not(edgefirst_tensor_dynamic))]`. `dynamic`'s `from_fd`
1737// (`dynamic_backend.rs`) drives `ef_tensor_builder_wrap` instead, which
1738// does this same classification inside `libedgefirst_tensor.so`.
1739#[cfg(all(target_os = "linux", not(edgefirst_tensor_dynamic)))]
1740use nix::sys::stat::{major, minor};
1741
1742/// Filesystem magic of the internal `dma_buf` mount, from
1743/// `include/uapi/linux/magic.h` (`DMA_BUF_MAGIC`, "DMAB").
1744///
1745/// This is stable UAPI and is the only reliable way to recognize a DMA-BUF
1746/// fd — see [`TensorStorage::from_fd`] for why `st_dev` cannot be used.
1747#[cfg(all(target_os = "linux", not(edgefirst_tensor_dynamic)))]
1748const DMA_BUF_MAGIC: u32 = 0x444d_4142;
1749
1750/// Filesystem magic of tmpfs, from `include/uapi/linux/magic.h`
1751/// (`TMPFS_MAGIC`). Covers both POSIX `shm_open` segments under `/dev/shm`
1752/// and anonymous `memfd_create` files.
1753#[cfg(all(target_os = "linux", not(edgefirst_tensor_dynamic)))]
1754const TMPFS_MAGIC: u32 = 0x0102_1994;
1755
1756/// Normalize a raw `fstatfs` `f_type` to the 32-bit unsigned magic that
1757/// `include/uapi/linux/magic.h` documents.
1758///
1759/// The width and signedness of `f_type` are target-dependent: `__fsword_t`
1760/// on Linux/gnu (`i64` on 64-bit, `i32` on 32-bit targets such as armv7,
1761/// i686 and aarch64-ilp32), `c_int` on uclibc, `c_ulong` on musl, and
1762/// `c_uint` on s390x. Where it is signed and 32 bits wide, widening it
1763/// sign-extends every magic with bit 31 set — hugetlbfs (`0x958458f6`),
1764/// btrfs (`0x9123683e`), f2fs (`0xf2f52010`) — which would make the value
1765/// reported by [`Error::UnknownBufferType`] impossible to find in the
1766/// header the docs point callers at.
1767///
1768/// Every magic is a 32-bit quantity, so truncating back to `u32` recovers
1769/// the true value from all four representations and loses nothing.
1770#[cfg(all(target_os = "linux", not(edgefirst_tensor_dynamic)))]
1771const fn fs_magic(raw: i64) -> u32 {
1772 raw as u32
1773}
1774
1775/// The operations every memory backend implements.
1776///
1777/// This is the seam that lets [`Tensor<T>`] hide which backend is in play.
1778/// The implementors are `DmaTensor` (Linux DMA-BUF), `IoSurfaceTensor`
1779/// (macOS/iOS), `AHardwareBufferTensor` (Android), `ShmTensor`, `MemTensor`,
1780/// and [`PboTensor`]; `TensorStorage<T>` dispatches to whichever is active.
1781/// Backend types are not public — allocate a [`Tensor<T>`] or [`TensorDyn`]
1782/// and call these methods through it.
1783///
1784/// Import the trait to get `shape`, `size`, `map`, `clone_fd`,
1785/// `buffer_identity`, and the zero-copy sub-region `view` on a tensor value.
1786/// `Tensor::view` and `Tensor::batch` route through [`TensorTrait::view`], so
1787/// each backend's identity-sharing rule lives in exactly one place.
1788///
1789/// [`TensorDyn`]: crate::TensorDyn
1790pub trait TensorTrait<T>: Send + Sync
1791where
1792 T: Num + Clone + fmt::Debug,
1793{
1794 /// Create a new tensor with the given shape and optional name. If no name
1795 /// is given, a random name will be generated.
1796 fn new(shape: &[usize], name: Option<&str>) -> Result<Self>
1797 where
1798 Self: Sized;
1799
1800 #[cfg(unix)]
1801 /// Import an existing buffer as a tensor, taking ownership of its file
1802 /// descriptor. The buffer is adopted in place — no bytes are copied.
1803 ///
1804 /// The backend is **detected**, not chosen: the fd already belongs to a
1805 /// buffer type, and this call's job is to recognize which. On Linux that
1806 /// is decided by the fd's filesystem magic, which is stable UAPI
1807 /// (`include/uapi/linux/magic.h`):
1808 ///
1809 /// | Filesystem | Magic | Resulting [`TensorMemory`] |
1810 /// |------------|-------|----------------------------|
1811 /// | `dma_buf` | `DMA_BUF_MAGIC` (`0x444d4142`) | [`TensorMemory::DmaBuf`] |
1812 /// | `tmpfs` (`/dev/shm` **and** `memfd`) | `TMPFS_MAGIC` (`0x01021994`) | [`TensorMemory::Shm`] |
1813 /// | anything else | — | *rejected* — see Errors |
1814 ///
1815 /// Both supported types are identified **positively**. An unrecognized
1816 /// filesystem is an error, never a fallback to shared memory: the wrong
1817 /// branch does not fail loudly (a DMA-BUF is mmap-able, so it would
1818 /// import as a perfectly functional tensor that merely isn't DMA, and a
1819 /// pipe would import as a zero-length one), so guessing would trade a
1820 /// clear error here for silent loss of zero-copy far downstream.
1821 ///
1822 /// The device number is deliberately **not** consulted. `dma_buf` files
1823 /// live on an internal kernel mount whose `st_dev` comes from
1824 /// `get_anon_bdev()` — an IDA shared with every other anonymous
1825 /// pseudo-filesystem and allocated in boot order — so the minor a
1826 /// DMA-BUF lands on varies by kernel build and is not part of any ABI.
1827 ///
1828 /// On non-Linux Unix (macOS/iOS/Android) there is no fd-based DMA import;
1829 /// the fd is always adopted as [`TensorMemory::Shm`].
1830 ///
1831 /// # Arguments
1832 ///
1833 /// * `fd` - Owned descriptor for the buffer to import. Ownership
1834 /// transfers to the returned tensor and the fd is closed on drop; pass
1835 /// [`clone_fd`](TensorTrait::clone_fd) output to keep your own handle.
1836 /// * `shape` - Logical shape to interpret the buffer with. Must describe
1837 /// no more elements than the buffer holds.
1838 /// * `name` - Optional name; a random one is generated when `None`.
1839 ///
1840 /// # Returns
1841 ///
1842 /// A tensor sharing the imported buffer's memory, whose
1843 /// [`memory()`](TensorTrait::memory) reports the detected backend.
1844 ///
1845 /// # Errors
1846 ///
1847 /// * [`Error::UnknownBufferType`] - the fd is on a filesystem that is
1848 /// neither `dma_buf` nor `tmpfs`, so its buffer type cannot be
1849 /// determined. Carries the observed `fstatfs` magic as a `u32`,
1850 /// normalized so it can be looked up in `magic.h` directly on both
1851 /// 32- and 64-bit targets. Typical causes: a regular file, a pipe or
1852 /// socket, or a `MFD_HUGETLB` memfd (hugetlbfs, not tmpfs). Linux only.
1853 /// * [`Error::UnknownDeviceType`] - the fd's `st_dev` major is non-zero,
1854 /// i.e. it lives on a real block device rather than an anonymous or
1855 /// in-memory filesystem. Linux only.
1856 /// * [`Error::InvalidSize`] - `shape` is empty or describes zero
1857 /// elements.
1858 /// * [`Error::NixError`] - `fstat`, `fstatfs`, or `mmap` failed on the
1859 /// descriptor.
1860 ///
1861 /// # Examples
1862 ///
1863 /// ```no_run
1864 /// use edgefirst_tensor::{Tensor, TensorMemory, TensorTrait};
1865 ///
1866 /// # fn main() -> edgefirst_tensor::Result<()> {
1867 /// let src = Tensor::<u8>::new(&[480, 640, 3], Some(TensorMemory::DmaBuf), None)?;
1868 ///
1869 /// // Round-tripping a DMA-BUF fd preserves the backend — the import is
1870 /// // still zero-copy, and still eligible for GPU/NPU paths.
1871 /// let imported = Tensor::<u8>::from_fd(src.clone_fd()?, src.shape(), None)?;
1872 /// assert_eq!(imported.memory(), TensorMemory::DmaBuf);
1873 /// # Ok(())
1874 /// # }
1875 /// ```
1876 fn from_fd(fd: std::os::fd::OwnedFd, shape: &[usize], name: Option<&str>) -> Result<Self>
1877 where
1878 Self: Sized;
1879
1880 #[cfg(unix)]
1881 /// Clone the file descriptor associated with this tensor.
1882 fn clone_fd(&self) -> Result<std::os::fd::OwnedFd>;
1883
1884 /// Get the memory type of this tensor.
1885 fn memory(&self) -> TensorMemory;
1886
1887 /// Get the name of this tensor.
1888 fn name(&self) -> String;
1889
1890 /// Get the number of elements in this tensor.
1891 fn len(&self) -> usize {
1892 self.shape().iter().product()
1893 }
1894
1895 /// Check if the tensor is empty.
1896 fn is_empty(&self) -> bool {
1897 self.len() == 0
1898 }
1899
1900 /// Get the size in bytes of this tensor.
1901 fn size(&self) -> usize {
1902 self.len() * std::mem::size_of::<T>()
1903 }
1904
1905 /// Get the shape of this tensor.
1906 fn shape(&self) -> &[usize];
1907
1908 /// Reshape this tensor to the given shape. The total number of elements
1909 /// must remain the same.
1910 fn reshape(&mut self, shape: &[usize]) -> Result<()>;
1911
1912 /// Bytes of the underlying allocation (>= the current logical `size()`).
1913 fn capacity_bytes(&self) -> usize;
1914
1915 /// Set the logical shape to any shape whose byte size fits the allocation
1916 /// capacity, without the equal-size constraint of `reshape`.
1917 fn set_logical_shape(&mut self, shape: &[usize]) -> Result<()>;
1918
1919 /// Map the tensor into memory with the given access direction and
1920 /// return a TensorMap for accessing the data.
1921 ///
1922 /// `access` selects the platform mapping mode (read-only IOSurface
1923 /// lock, dma-buf sync direction, AHardwareBuffer lock usage) and the
1924 /// map's mutability: a map obtained with [`CpuAccess::Read`] rejects
1925 /// `as_mut_slice`. [`CpuAccess::None`] is not a mappable direction
1926 /// and returns [`Error::InvalidArgument`]; a direction the tensor's
1927 /// declared [`CpuAccess`] does not cover returns
1928 /// [`Error::CpuAccessDenied`].
1929 ///
1930 /// Prefer the typed wrappers [`map_read`](Self::map_read) /
1931 /// [`map_write`](Self::map_write) / [`map_mut`](Self::map_mut).
1932 ///
1933 /// Like [`pin_host`](Tensor::pin_host), `'a` bounds how long the backing
1934 /// allocation lives rather than borrowing the tensor: the returned view
1935 /// shares ownership through the pin's keepalive, so a map can coexist with
1936 /// a `&mut` on the same tensor.
1937 ///
1938 /// # Windows D3D11 textures: a write-only map of a partial window
1939 ///
1940 /// A [`CpuAccess::Write`] map is not refreshed from the texture first, so
1941 /// the bytes the caller leaves untouched are undefined -- the documented
1942 /// contract a whole-tensor map can meet and a sub-view cannot, since the
1943 /// window spans its own rows only.
1944 ///
1945 /// A tensor allocated [`CpuAccess::Write`] (no staging texture) publishes
1946 /// only the window's own rows on unmap, so a sub-view -- a
1947 /// [`view`](Tensor::view), or a shape narrowed with `set_logical_shape`
1948 /// -- is refused only when its window is not a whole, pitch-aligned row
1949 /// range (a plain-offset [`view`](Tensor::view) always is; an arbitrary
1950 /// byte offset from `set_plane_offset` might not be), which fails with
1951 /// [`Error::InvalidArgument`].
1952 ///
1953 /// A tensor allocated [`CpuAccess::Read`] or [`CpuAccess::ReadWrite`]
1954 /// (staging texture) has no such per-window publish: unmapping always
1955 /// copies the whole staging texture back, so a write-only sub-view is
1956 /// refused with [`Error::PartialWriteRequiresReadWrite`], naming
1957 /// [`CpuAccess::ReadWrite`] as the same window with the refresh that
1958 /// makes the untouched rows well-defined -- and, unlike the
1959 /// `CpuAccess::Write` case above, that retry always succeeds.
1960 fn map_with<'a>(&self, access: CpuAccess) -> Result<crate::view::HostView<'a, T>>
1961 where
1962 T: 'a;
1963
1964 /// Map the tensor read-write (equivalent to
1965 /// `map_with(CpuAccess::ReadWrite)`).
1966 ///
1967 /// Fails with [`Error::CpuAccessDenied`] on a `Read`- or `Write`-declared
1968 /// tensor: use [`map_read`](Self::map_read) to read one and
1969 /// [`map_for_write`](Self::map_for_write) to fill one.
1970 fn map<'a>(&self) -> Result<crate::view::HostView<'a, T>>
1971 where
1972 T: 'a,
1973 {
1974 self.map_with(CpuAccess::ReadWrite)
1975 }
1976
1977 /// Map the tensor for CPU reading only. The returned map rejects
1978 /// `as_mut_slice`; on macOS this takes the read-only IOSurface lock
1979 /// (skips the unlock flush), on Linux the dma-buf read-direction
1980 /// sync. Fails with [`Error::CpuAccessDenied`] on a `Write`- or
1981 /// `None`-declared tensor.
1982 fn map_read<'a>(&self) -> Result<crate::view::HostView<'a, T>>
1983 where
1984 T: 'a,
1985 {
1986 self.map_with(CpuAccess::Read)
1987 }
1988
1989 /// Map the tensor for CPU writing (fill-only: reading through a
1990 /// write map may see write-combined memory — do not read the slice).
1991 /// Valid only for a pass that overwrites the whole mapping; see
1992 /// [`write_access`](Self::write_access). Fails with
1993 /// [`Error::CpuAccessDenied`] on a `Read`- or `None`-declared tensor.
1994 fn map_write<'a>(&self) -> Result<crate::view::HostView<'a, T>>
1995 where
1996 T: 'a,
1997 {
1998 self.map_with(CpuAccess::Write)
1999 }
2000
2001 /// Map the tensor read-write (alias of [`map`](Self::map) with the
2002 /// intent spelled out); refused on the same declarations.
2003 fn map_mut<'a>(&self) -> Result<crate::view::HostView<'a, T>>
2004 where
2005 T: 'a,
2006 {
2007 self.map_with(CpuAccess::ReadWrite)
2008 }
2009
2010 /// The CPU access declared for this tensor (see [`CpuAccess`]). A bare
2011 /// backing carries no declaration of its own and reports `ReadWrite`.
2012 fn cpu_access(&self) -> CpuAccess {
2013 CpuAccess::ReadWrite
2014 }
2015
2016 /// The parent-image snapshot when this tensor is a sub-region of a larger
2017 /// buffer (see [`ViewOrigin`]); a bare backing reports `None`.
2018 fn view_origin(&self) -> Option<ViewOrigin> {
2019 None
2020 }
2021
2022 /// Whether this tensor addresses part of a larger buffer (see
2023 /// [`Tensor::is_view`]); a bare backing reports `false`.
2024 fn is_view(&self) -> bool {
2025 self.view_origin().is_some()
2026 }
2027
2028 /// The direction a CPU pass that overwrites every element of this tensor
2029 /// maps with.
2030 ///
2031 /// A write-only map does not invalidate the CPU cache (on arm64,
2032 /// `DMA_BUF_SYNC_START|WRITE` only orders), so a stale cache line the pass
2033 /// writes only partly is written back over whatever a device left in the
2034 /// rest of it. The map may therefore be `Write` only when the pass covers
2035 /// the whole buffer: a whole tensor takes its declaration plus `Write`
2036 /// (`Write` on a write-only declaration), a [`view`](Tensor::view), whose
2037 /// rows share cache lines with bytes outside it, takes `ReadWrite`.
2038 fn write_access(&self) -> CpuAccess {
2039 if self.is_view() {
2040 CpuAccess::ReadWrite
2041 } else {
2042 self.cpu_access().union(CpuAccess::Write)
2043 }
2044 }
2045
2046 /// Map for a CPU pass that overwrites every element, in
2047 /// [`write_access`](Self::write_access)'s direction.
2048 ///
2049 /// # Errors
2050 ///
2051 /// [`Error::CpuAccessDenied`] when the declaration does not cover that
2052 /// direction, otherwise whatever [`map_with`](Self::map_with) returns.
2053 fn map_for_write<'a>(&self) -> Result<crate::view::HostView<'a, T>>
2054 where
2055 T: 'a,
2056 {
2057 self.map_with(self.write_access())
2058 }
2059
2060 /// Map for a CPU pass that writes only part of the buffer and keeps the
2061 /// rest (a letterbox border around a device-rendered interior): always
2062 /// `ReadWrite`, see [`write_access`](Self::write_access).
2063 ///
2064 /// # Errors
2065 ///
2066 /// [`Error::CpuAccessDenied`] unless the declaration is `ReadWrite`,
2067 /// otherwise whatever [`map_with`](Self::map_with) returns.
2068 fn map_for_partial_write<'a>(&self) -> Result<crate::view::HostView<'a, T>>
2069 where
2070 T: 'a,
2071 {
2072 self.map_with(CpuAccess::ReadWrite)
2073 }
2074
2075 /// Non-blocking [`map_with`](Self::map_with): returns
2076 /// `Err(Error::IoError(WouldBlock))` while a GPU copy the map depends on
2077 /// is still in flight, and makes progress on a retry.
2078 ///
2079 /// Backings with no such copy (host memory, shared memory, dma-buf,
2080 /// IOSurface, AHardwareBuffer, PBO) alias `map_with`, so a caller can use
2081 /// this unconditionally and only ever see `WouldBlock` where the wait is
2082 /// real -- today the Windows D3D11 texture's staging refresh.
2083 ///
2084 /// A retry loop must yield or sleep between attempts. On the WARP
2085 /// software adapter the CPU threads are the GPU, so a tight loop starves
2086 /// the very copy it is waiting for.
2087 ///
2088 /// [`CpuAccess::Write`] is refused on a Windows D3D11 texture tensor whose
2089 /// window is shorter than its backing, exactly as it is for
2090 /// [`map_with`](Self::map_with); see that method.
2091 fn try_map_with<'a>(&self, access: CpuAccess) -> Result<crate::view::HostView<'a, T>>
2092 where
2093 T: 'a,
2094 {
2095 self.map_with(access)
2096 }
2097
2098 /// Non-blocking [`map`](Self::map): read-write, and `WouldBlock` rather
2099 /// than a stall. See [`try_map_with`](Self::try_map_with).
2100 fn try_map<'a>(&self) -> Result<crate::view::HostView<'a, T>>
2101 where
2102 T: 'a,
2103 {
2104 self.try_map_with(CpuAccess::ReadWrite)
2105 }
2106
2107 /// Get the buffer identity for cache keying and liveness tracking.
2108 fn buffer_identity(&self) -> &BufferIdentity;
2109
2110 /// Create a zero-copy sub-region view of this backing that shares the
2111 /// underlying allocation **and** [`BufferIdentity`].
2112 ///
2113 /// The window is `[offset_bytes, offset_bytes + shape.product() *
2114 /// size_of::<T>())` measured from this tensor's own logical start, so a
2115 /// sub-view of a sub-view composes by adding offsets. Sharing the parent's
2116 /// identity is the contract that lets identity-keyed caches (e.g. the GL
2117 /// EGLImage import cache) treat offset-distinct windows as one buffer rather
2118 /// than unrelated allocations — `view` must never mint a fresh identity.
2119 ///
2120 /// Defaults to [`Error::NotImplemented`]; every backend that supports
2121 /// sub-views overrides it (`Mem`, `Shm`, Linux DMA, macOS IOSurface, `Pbo`).
2122 fn view(&self, offset_bytes: usize, shape: &[usize]) -> Result<Self>
2123 where
2124 Self: Sized,
2125 {
2126 let _ = (offset_bytes, shape);
2127 Err(Error::NotImplemented(
2128 "view (zero-copy sub-region) is not supported for this tensor backend".to_owned(),
2129 ))
2130 }
2131}
2132
2133pub trait TensorMapTrait<T>
2134where
2135 T: Num + Clone + fmt::Debug,
2136{
2137 /// Get the shape of this tensor map.
2138 fn shape(&self) -> &[usize];
2139
2140 /// Unmap the tensor from memory.
2141 fn unmap(&mut self);
2142
2143 /// Get the number of elements in this tensor map.
2144 fn len(&self) -> usize {
2145 self.shape().iter().product()
2146 }
2147
2148 /// Check if the tensor map is empty.
2149 fn is_empty(&self) -> bool {
2150 self.len() == 0
2151 }
2152
2153 /// Get the size in bytes of this tensor map.
2154 fn size(&self) -> usize {
2155 self.len() * std::mem::size_of::<T>()
2156 }
2157
2158 /// Get a slice to the data in this tensor map.
2159 fn as_slice(&self) -> &[T];
2160
2161 /// Get a mutable slice to the data in this tensor map.
2162 fn as_mut_slice(&mut self) -> &mut [T];
2163
2164 /// Whether this mapping permits writes -- i.e. whether the coherency
2165 /// bracket it holds will flush on release.
2166 ///
2167 /// The single source of truth for a binding that has to advertise
2168 /// mutability to a foreign consumer (the Python buffer protocol's
2169 /// `readonly` flag, say). Deriving that from the access a caller
2170 /// *requested* instead would leave the advertisement agreeing with the
2171 /// request while both disagreed with the mapping actually taken.
2172 fn is_writable(&self) -> bool;
2173
2174 #[cfg(feature = "ndarray")]
2175 /// Get an ndarray ArrayView of the tensor data.
2176 fn view(&'_ self) -> Result<ndarray::ArrayView<'_, T, ndarray::Dim<ndarray::IxDynImpl>>> {
2177 Ok(ndarray::ArrayView::from_shape(
2178 self.shape(),
2179 self.as_slice(),
2180 )?)
2181 }
2182
2183 #[cfg(feature = "ndarray")]
2184 /// Get an ndarray ArrayViewMut of the tensor data.
2185 fn view_mut(
2186 &'_ mut self,
2187 ) -> Result<ndarray::ArrayViewMut<'_, T, ndarray::Dim<ndarray::IxDynImpl>>> {
2188 let shape = self.shape().to_vec();
2189 Ok(ndarray::ArrayViewMut::from_shape(
2190 shape,
2191 self.as_mut_slice(),
2192 )?)
2193 }
2194}
2195
2196// Declared through `ef_vocabulary!` (see `vocabulary.rs`). These codes are
2197// the single numbering of backing stores across the stack: `code()`, the
2198// `tensor_memory_wire` constants, the C header's `ef_storage_kind`, the
2199// Python `TensorMemory` enum and the cross-package descriptor's
2200// `TensorDesc::kind` (`protocol::kind` re-exports `tensor_memory_wire`) all
2201// read the same declared literal per variant.
2202crate::ef_vocabulary! {
2203 /// Which memory backend a tensor is (or should be) allocated from.
2204 ///
2205 /// Pass `Some(..)` to a constructor to pin the backend, or `None` to
2206 /// auto-select. Auto-selection differs by constructor: [`Tensor::new`]
2207 /// tries `DmaBuf` → `Mem`, while the image constructors try `DmaBuf` →
2208 /// `Mem` and skip `Shm`. `EDGEFIRST_TENSOR_FORCE_MEM=1` short-circuits
2209 /// either chain to `Mem`, which is how tests run on hosts without
2210 /// DMA-heap permissions.
2211 ///
2212 /// A pinned request has no fallback: if the backend cannot serve it, the
2213 /// constructor fails rather than quietly giving you something slower.
2214 /// [`TensorTrait::memory`] reports what a tensor actually ended up with.
2215 ///
2216 /// # Every variant exists on every platform
2217 ///
2218 /// This is a namespace of codes, not a list of what this build can
2219 /// allocate -- [`TensorMemory::is_available`] answers the latter, at
2220 /// runtime. `Shm` used to be `#[cfg(unix)]`, which shifted every
2221 /// variant declared after it on a non-unix build and made the
2222 /// numbering a compile-time detail of the host. It also made a
2223 /// cross-platform recording unreadable: a blob captured on Linux
2224 /// carries `"dmabuf"`, and a macOS consumer has to be able to NAME
2225 /// that backing in order to report "cannot materialise here" instead
2226 /// of failing with an unrecognised code.
2227 ///
2228 /// # `#[non_exhaustive]`
2229 ///
2230 /// Match this with a wildcard arm. The variant list is expected to
2231 /// keep growing -- `IoSurface` and `Cuda` were added here, Windows
2232 /// backings and a real CUDA device backing are named as coming -- and
2233 /// this attribute is what keeps a closed-source consumer's
2234 /// `match memory { .. }` compiling across those additions. It protects
2235 /// DOWNSTREAM, not this crate: inside `edgefirst-tensor` the matches
2236 /// stay exhaustive and compile-enforced. Same rule as [`DType`],
2237 /// [`Compression`] and [`PixelFormat`], which `TensorMemory` was the
2238 /// lone exception to.
2239 #[non_exhaustive]
2240 pub enum TensorMemory {
2241 /// Regular system memory allocation. Available everywhere; the
2242 /// floor every auto-selection chain ends on.
2243 Mem = 0, "mem", MEM,
2244 /// POSIX Shared Memory allocation. Suitable for inter-process
2245 /// communication, but not for hardware acceleration. Only
2246 /// allocatable on unix; nameable everywhere.
2247 Shm = 1, "shm", SHM,
2248 /// Platform-native zero-copy GPU buffer.
2249 ///
2250 /// On Linux this is a DMA-BUF (`DmaTensor` in
2251 /// `crates/tensor/src/dma.rs`) allocated via the DRM/dma-heap
2252 /// subsystem. On macOS/iOS it is currently also this variant that
2253 /// yields an IOSurface, and on Android an AHardwareBuffer: they
2254 /// share the `TensorStorage::Dma` slot at the trait level, so the
2255 /// one code covers all of them. `IoSurface` and `D3d11Texture`
2256 /// below are the codes for naming those backings specifically; no
2257 /// backend reports either as its `memory()` yet.
2258 ///
2259 /// Allows hardware-accelerated paths (OpenGL backend on Linux via
2260 /// `EGL_EXT_image_dma_buf_import`; macOS via
2261 /// `EGL_ANGLE_iosurface_client_buffer`). CPU access via `map()`
2262 /// incurs cache-coherency overhead on Linux DMA-BUF and is similar
2263 /// in cost on IOSurface; SHM/Mem are cheaper for CPU-only workloads.
2264 ///
2265 /// On Windows it is an `ID3D11Texture2D` on the process D3D11 device
2266 /// (`crates/tensor/src/d3d11/`), allocated only through
2267 /// [`Tensor::image`](crate::Tensor::image); CPU access goes through a
2268 /// staging copy. The first such allocation creates the process
2269 /// device, so a host supplying its own device must call
2270 /// `d3d11::use_external_device` before it -- and before
2271 /// [`crate::is_gpu_buffer_available`], which creates the device too.
2272 DmaBuf = 2, "dmabuf", DMABUF,
2273 /// Apple IOSurface, named specifically rather than through the
2274 /// portable `DmaBuf` spelling.
2275 ///
2276 /// Defined as a code so a descriptor or recording that names an
2277 /// IOSurface is parseable on every platform. **No backend produces
2278 /// or accepts it yet** -- macOS/iOS allocate and report `DmaBuf`.
2279 /// Pinning a request to it fails with `NotImplemented`.
2280 IoSurface = 3, "iosurface", IOSURFACE,
2281 /// OpenGL Pixel Buffer Object memory. Created by `ImageProcessor`
2282 /// when DMA-buf is unavailable but OpenGL is present -- never by
2283 /// [`Tensor::new`], which has no GL context to create it in.
2284 Pbo = 4, "pbo", PBO,
2285 /// CUDA device memory: a device pointer, not host-addressable.
2286 ///
2287 /// Defined as a code for the same reason as `IoSurface`. **No
2288 /// backend produces or accepts it yet**; pinning a request to it
2289 /// fails with `NotImplemented`.
2290 Cuda = 5, "cuda", CUDA,
2291 /// Windows `ID3D11Texture2D`, named specifically rather than through
2292 /// the portable `DmaBuf` spelling.
2293 ///
2294 /// This is the code a Windows producer writes into
2295 /// [`TensorDesc::kind`](crate::TensorDesc::kind), where the handle
2296 /// field means an NT shared handle rather than an fd. As a
2297 /// [`TensorMemory`] it is defined for the same reason as
2298 /// `IoSurface`: **no backend allocates or reports it** -- Windows
2299 /// allocates and reports `DmaBuf` -- and pinning a request to it
2300 /// fails with `NotImplemented`.
2301 D3d11Texture = 6, "d3d11", D3D11_TEXTURE,
2302 }
2303 // `#[doc(hidden)]`: `pub` so the const-only form is reachable by an FFI
2304 // or cbindgen consumer, but it is emission plumbing rather than a second
2305 // documented API -- `TensorMemory::code()` is the documented path, and
2306 // the two are generated from the same literal.
2307 #[doc(hidden)]
2308 pub mod tensor_memory_wire;
2309}
2310
2311impl TensorMemory {
2312 /// Whether this backing can actually be allocated in this process, now.
2313 ///
2314 /// A runtime question, not a compile-time one, which is why it is a
2315 /// method on the vocabulary rather than a `cfg` on the variant:
2316 /// `DmaBuf` needs `/dev/dma_heap` to exist *and* this process to have
2317 /// permission on it (a container, a user outside the `video`/`render`
2318 /// groups, and a stock desktop give three different answers on the same
2319 /// kernel), and `Shm` needs a writable `/dev/shm`. Each probe allocates
2320 /// once and caches.
2321 ///
2322 /// Reports what can be *allocated*, which is not the same as what can be
2323 /// *named*: every variant is nameable on every platform (that is the
2324 /// point of the vocabulary), and the ones this build cannot serve say so
2325 /// here rather than by being absent from the enum.
2326 ///
2327 /// Four variants answer `false` unconditionally today, for reasons
2328 /// worth stating rather than hiding behind a probe that cannot fail:
2329 ///
2330 /// * `IoSurface` -- no backend produces or accepts it yet; macOS/iOS
2331 /// serve IOSurface under the `DmaBuf` spelling, so `DmaBuf` is the
2332 /// variant that answers `true` there.
2333 /// * `Pbo` -- PBO tensors are created by `ImageProcessor::create_image`
2334 /// with a GL context current, never by [`Tensor::new`]. This crate
2335 /// holds no GL context and cannot see one; ask the image crate
2336 /// whether a GL backend is live.
2337 /// * `Cuda` -- there is no CUDA backing in this crate at all. Probing
2338 /// for `libcuda` would answer `true` on a Jetson while every
2339 /// allocation still failed, which is worse than answering `false`.
2340 /// * `D3d11Texture` -- same as `IoSurface`: Windows serves D3D11
2341 /// textures under the `DmaBuf` spelling.
2342 ///
2343 /// One variant answers `true` for a narrower constructor than the rest:
2344 /// on Windows `DmaBuf` is an `ID3D11Texture2D`, which is image-formatted,
2345 /// so `true` there means [`Tensor::image`] can serve it and
2346 /// [`Tensor::new`] with `Some(DmaBuf)` still answers `NotImplemented`.
2347 /// [`is_gpu_buffer_available`](crate::is_gpu_buffer_available), which this
2348 /// forwards to, states the same split.
2349 pub fn is_available(self) -> bool {
2350 match self {
2351 // Always: the fallback every constructor chain ends on. If this
2352 // could fail, allocation would have no floor left.
2353 TensorMemory::Mem => true,
2354 TensorMemory::Shm => is_shm_available(),
2355 // Today `DmaBuf` is the portable spelling of "platform-native
2356 // zero-copy GPU buffer", so it is the portable probe that
2357 // answers for it -- true on macOS/iOS (IOSurface) and Android
2358 // (AHardwareBuffer) as well as Linux.
2359 TensorMemory::DmaBuf => is_gpu_buffer_available(),
2360 TensorMemory::IoSurface
2361 | TensorMemory::Pbo
2362 | TensorMemory::Cuda
2363 | TensorMemory::D3d11Texture => false,
2364 }
2365 }
2366
2367 /// The error to return when a caller pins a request to a backing this
2368 /// build knows the name of but cannot serve.
2369 ///
2370 /// Deliberately not `unreachable!()`. `IoSurface` and `Cuda` are
2371 /// unreachable *today* only because no backend reports them yet; the
2372 /// day one does, an `unreachable!()` written now becomes a panic in
2373 /// library code reached from an ordinary caller request. A "this build
2374 /// cannot do that" error is the same information without the crash, and
2375 /// it is what every other unservable request here already returns.
2376 ///
2377 /// `static`-only: every call site is inside `impl<T> TensorStorage<T>`
2378 /// (`TensorStorage::new`), itself `#[cfg(not(edgefirst_tensor_dynamic))]`.
2379 #[cfg(not(edgefirst_tensor_dynamic))]
2380 pub(crate) fn unsupported_here(self) -> Error {
2381 let reason = match self {
2382 // Not reachable: `Mem` is servable on every target. Stated
2383 // rather than panicked on, for the reason above.
2384 TensorMemory::Mem => "Mem is servable on every platform, so this should not happen",
2385 TensorMemory::Shm => "POSIX shared memory is unix-only",
2386 TensorMemory::DmaBuf => {
2387 "this target has no platform GPU-buffer allocator (DMA-BUF is Linux, \
2388 IOSurface macOS/iOS, AHardwareBuffer Android, D3D11 textures on \
2389 Windows)"
2390 }
2391 TensorMemory::IoSurface => {
2392 "no backend produces or accepts IOSurface under its own code yet -- \
2393 macOS/iOS allocate and report TensorMemory::DmaBuf"
2394 }
2395 TensorMemory::Pbo => {
2396 "PBO tensors are created by ImageProcessor::create_image with a GL \
2397 context current, never by Tensor::new"
2398 }
2399 TensorMemory::Cuda => "there is no CUDA backing in this build",
2400 TensorMemory::D3d11Texture => {
2401 "no backend produces or accepts a D3D11 texture under its own code yet -- \
2402 Windows allocates and reports TensorMemory::DmaBuf"
2403 }
2404 };
2405 Error::NotImplemented(format!(
2406 "TensorMemory::{self:?} is not supported by this build: {reason}"
2407 ))
2408 }
2409}
2410
2411impl From<TensorMemory> for String {
2412 /// The wire string, via [`TensorMemory::as_str`] -- one spelling, not a
2413 /// second one maintained alongside it.
2414 ///
2415 /// `DmaBuf` renders as `"dmabuf"`, not the `"dma"` this impl used to
2416 /// emit. `TryFrom<&str>` still accepts `"dma"` so anything that
2417 /// persisted the old spelling keeps parsing; the benchmark result keys
2418 /// in `.github/scripts` spell it `"dma"` but build that string
2419 /// themselves and never went through this impl.
2420 fn from(memory: TensorMemory) -> Self {
2421 memory.as_str().to_owned()
2422 }
2423}
2424
2425impl TryFrom<&str> for TensorMemory {
2426 type Error = Error;
2427
2428 /// Parses [`TensorMemory::as_str`]'s spelling, plus the legacy `"dma"`
2429 /// alias for `DmaBuf`.
2430 fn try_from(s: &str) -> Result<Self> {
2431 TensorMemory::from_str_code(s)
2432 .or(if s == "dma" {
2433 Some(TensorMemory::DmaBuf)
2434 } else {
2435 None
2436 })
2437 .ok_or_else(|| Error::InvalidMemoryType(s.to_owned()))
2438 }
2439}
2440
2441// Everything from here through the end of `impl<T> TensorTrait<T> for
2442// Tensor<T>` below is the `static` backend's storage-owning `Tensor<T>`:
2443// `TensorStorage<T>` (the per-memory-backend enum) and the `Tensor<T>`
2444// struct that wraps it. `dynamic`'s `Tensor<T>` is a different, much
2445// smaller type -- a typed lens over `TensorDyn` with no storage of its
2446// own -- defined in `tensor_dyn/dynamic_tensor.rs` and re-exported below.
2447#[cfg(not(edgefirst_tensor_dynamic))]
2448#[derive(Debug)]
2449#[allow(dead_code)] // Variants are constructed by downstream crates via pub(crate) helpers
2450pub(crate) enum TensorStorage<T>
2451where
2452 T: Num + Clone + fmt::Debug + Send + Sync,
2453{
2454 /// Platform-native zero-copy GPU buffer. Inner type differs per
2455 /// target: `DmaTensor` on Linux (DMA-BUF fd), `IoSurfaceTensor` on
2456 /// macOS (CFRetained IOSurface). The shared variant name keeps the
2457 /// public `TensorMemory::DmaBuf` discriminant stable across platforms.
2458 #[cfg(target_os = "linux")]
2459 Dma(DmaTensor<T>),
2460 #[cfg(any(target_os = "macos", target_os = "ios"))]
2461 Dma(IoSurfaceTensor<T>),
2462 #[cfg(target_os = "android")]
2463 Dma(AHardwareBufferTensor<T>),
2464 #[cfg(target_os = "windows")]
2465 Dma(crate::d3d11::texture::D3d11TextureTensor<T>),
2466 #[cfg(unix)]
2467 Shm(ShmTensor<T>),
2468 Mem(MemTensor<T>),
2469 Pbo(PboTensor<T>),
2470}
2471
2472#[cfg(not(edgefirst_tensor_dynamic))]
2473impl<T> TensorStorage<T>
2474where
2475 T: Num + Clone + fmt::Debug + Send + Sync,
2476{
2477 /// The backing allocation's intrinsic physical row pitch in bytes, if it
2478 /// has one that is fixed independent of the logical shape. macOS IOSurface
2479 /// reports its 64-aligned `bytesPerRow`; other backings (Linux DMA, SHM,
2480 /// Mem, PBO) have no fixed pitch beyond the logical shape and return `None`.
2481 ///
2482 /// Used by `configure_image` to preserve the physical pitch when a reused
2483 /// pool tensor is reconfigured to a smaller logical image — so the decode
2484 /// writes rows at the surface's real stride and the GPU samples them with
2485 /// the same stride (the physical-grid / logical-ROI decoupling).
2486 ///
2487 /// Compiled only where such a backing exists; on Linux and other unix
2488 /// targets it would always answer `None`.
2489 #[cfg(any(
2490 target_os = "macos",
2491 target_os = "ios",
2492 target_os = "android",
2493 target_os = "windows"
2494 ))]
2495 pub(crate) fn backing_row_stride(&self) -> Option<usize> {
2496 match self {
2497 // Only genuine image-formatted IOSurfaces (height > 1) carry a real
2498 // per-row pitch; a generic byte-bag (height == 1) returns `None` so
2499 // `configure_image` does not adopt its whole-buffer "row" as a stride.
2500 #[cfg(any(target_os = "macos", target_os = "ios"))]
2501 TensorStorage::Dma(t) => t.image_backing_row_stride(),
2502 // Android AHardwareBuffer: same rule — only genuine 2D
2503 // image-formatted buffers (height > 1) carry a real row pitch.
2504 #[cfg(target_os = "android")]
2505 TensorStorage::Dma(t) => t.image_backing_row_stride(),
2506 // Windows D3D11: the driver's staging row pitch. `None` when the
2507 // tensor has no staging (write-only or hardware-only), whose host
2508 // buffer has tight rows -- except for a semi-planar tensor, whose
2509 // texture is allocated as wide as that pitch for every access, so
2510 // it always has one to report.
2511 #[cfg(target_os = "windows")]
2512 TensorStorage::Dma(t) => t.image_backing_row_stride(),
2513 _ => None,
2514 }
2515 }
2516
2517 /// Create a new tensor storage with the given shape, memory type, and
2518 /// optional name. If no name is given, a random name will be generated.
2519 /// If no memory type is given, the best available memory type will be
2520 /// chosen based on the platform and environment variables.
2521 fn new(shape: &[usize], memory: Option<TensorMemory>, name: Option<&str>) -> Result<Self> {
2522 match memory {
2523 #[cfg(target_os = "linux")]
2524 Some(TensorMemory::DmaBuf) => {
2525 DmaTensor::<T>::new(shape, name).map(TensorStorage::Dma)
2526 }
2527 #[cfg(any(target_os = "macos", target_os = "ios"))]
2528 Some(TensorMemory::DmaBuf) => {
2529 IoSurfaceTensor::<T>::new(shape, name).map(TensorStorage::Dma)
2530 }
2531 #[cfg(target_os = "android")]
2532 Some(TensorMemory::DmaBuf) => {
2533 AHardwareBufferTensor::<T>::new(shape, name).map(TensorStorage::Dma)
2534 }
2535 // Windows has a DmaBuf backing, but it is a D3D11 texture: an image
2536 // layout, not a byte bag, so a shape alone cannot describe one.
2537 // Name the constructor that can, rather than quietly handing back
2538 // host memory the caller did not ask for.
2539 #[cfg(target_os = "windows")]
2540 Some(TensorMemory::DmaBuf) => Err(crate::error::Error::NotImplemented(
2541 "D3D11 texture tensors are image-formatted; use Tensor::image with \
2542 TensorMemory::DmaBuf"
2543 .to_owned(),
2544 )),
2545 #[cfg(not(any(
2546 target_os = "linux",
2547 target_os = "macos",
2548 target_os = "ios",
2549 target_os = "android",
2550 target_os = "windows"
2551 )))]
2552 Some(TensorMemory::DmaBuf) => Err(crate::error::Error::NotImplemented(
2553 "TensorMemory::DmaBuf is only available on Linux (DMA-BUF), macOS/iOS (IOSurface), \
2554 Android (AHardwareBuffer), and Windows (D3D11 textures, via Tensor::image)"
2555 .to_owned(),
2556 )),
2557 #[cfg(unix)]
2558 Some(TensorMemory::Shm) => {
2559 ShmTensor::<T>::new(shape, name).map(TensorStorage::Shm)
2560 }
2561 #[cfg(not(unix))]
2562 Some(m @ TensorMemory::Shm) => Err(m.unsupported_here()),
2563 Some(TensorMemory::Mem) => {
2564 MemTensor::<T>::new(shape, name).map(TensorStorage::Mem)
2565 }
2566 Some(TensorMemory::Pbo) => Err(crate::error::Error::NotImplemented(
2567 "PboTensor cannot be created via Tensor::new() — use ImageProcessor::create_image()".to_owned(),
2568 )),
2569 // Defined codes with no backing behind them yet. An error, not
2570 // an `unreachable!()` -- see `unsupported_here`.
2571 Some(
2572 m @ (TensorMemory::IoSurface | TensorMemory::Cuda | TensorMemory::D3d11Texture),
2573 ) => Err(m.unsupported_here()),
2574 None => {
2575 if std::env::var("EDGEFIRST_TENSOR_FORCE_MEM")
2576 .is_ok_and(|x| x != "0" && x.to_lowercase() != "false")
2577 {
2578 MemTensor::<T>::new(shape, name).map(TensorStorage::Mem)
2579 } else {
2580 // Auto-select priority: Dma > Mem. Shm is intentionally NOT
2581 // auto-selected — it offers no advantage over Mem for an
2582 // in-process tensor and Mem always succeeds, so Shm below Mem
2583 // is effectively never reached. Request Shm explicitly via
2584 // `TensorMemory::Shm` when cross-process sharing is needed.
2585 // (PBO sits between Dma and Mem but is GL-backed and created
2586 // only via `ImageProcessor::create_image`.)
2587 #[cfg(target_os = "linux")]
2588 {
2589 // Linux: Try DMA -> Mem
2590 match DmaTensor::<T>::new(shape, name) {
2591 Ok(tensor) => Ok(TensorStorage::Dma(tensor)),
2592 Err(_) => MemTensor::<T>::new(shape, name).map(TensorStorage::Mem),
2593 }
2594 }
2595 #[cfg(any(target_os = "macos", target_os = "ios"))]
2596 {
2597 // macOS/iOS: Try IOSurface -> Mem. IOSurface is the
2598 // GPU-shareable backend (zero-copy via ANGLE), filling the
2599 // same role as DMA-BUF on Linux.
2600 match IoSurfaceTensor::<T>::new(shape, name) {
2601 Ok(tensor) => Ok(TensorStorage::Dma(tensor)),
2602 Err(_) => MemTensor::<T>::new(shape, name).map(TensorStorage::Mem),
2603 }
2604 }
2605 #[cfg(target_os = "android")]
2606 {
2607 // Android: Mem. Unlike macOS (where a byte-bag
2608 // IOSurface is still GL-importable as R8) a generic
2609 // BLOB AHardwareBuffer cannot back a GPU texture, so
2610 // auto-selecting it buys nothing and costs plenty: a
2611 // gralloc allocator-HAL ioctl per allocation (orders
2612 // slower than malloc, ≥1 page + a dmabuf fd even for
2613 // tiny tensors) and a lock/unlock cache-maintenance
2614 // round trip on every map(). Zero-copy image tensors
2615 // come from `Tensor::image(..)`; callers that want a
2616 // BLOB (NNAPI handoff) request `TensorMemory::DmaBuf`
2617 // explicitly.
2618 MemTensor::<T>::new(shape, name).map(TensorStorage::Mem)
2619 }
2620 #[cfg(all(
2621 unix,
2622 not(any(
2623 target_os = "linux",
2624 target_os = "macos",
2625 target_os = "ios",
2626 target_os = "android"
2627 ))
2628 ))]
2629 {
2630 // Other Unix (BSD): Mem only (no DMA; Shm is explicit-only)
2631 MemTensor::<T>::new(shape, name).map(TensorStorage::Mem)
2632 }
2633 #[cfg(not(unix))]
2634 {
2635 // Windows/other: Mem only
2636 MemTensor::<T>::new(shape, name).map(TensorStorage::Mem)
2637 }
2638 }
2639 }
2640 }
2641 }
2642
2643 /// Create a DMA-backed tensor storage with an explicit byte size that
2644 /// may exceed `shape.product() * sizeof(T)`. Used for image tensors
2645 /// with row-padded layouts (see `DmaTensor::new_with_byte_size`).
2646 ///
2647 /// This is intentionally DMA-only: padding is only meaningful for
2648 /// buffers that will be imported as GPU textures via EGLImage. PBO,
2649 /// Shm, and Mem storage doesn't benefit from pitch alignment and
2650 /// shouldn't pay the memory cost.
2651 /// `policy` picks the heaps (see [`crate::dma::HeapPolicy`]): with
2652 /// `contiguous` the CMA heaps only (see [`ImageDesc::with_contiguous`]),
2653 /// and with `uncached` the uncached CMA heap first.
2654 #[cfg(target_os = "linux")]
2655 pub(crate) fn new_dma_with_byte_size(
2656 shape: &[usize],
2657 byte_size: usize,
2658 name: Option<&str>,
2659 policy: crate::dma::HeapPolicy,
2660 ) -> Result<Self> {
2661 DmaTensor::<T>::new_with_byte_size(shape, byte_size, name, policy).map(TensorStorage::Dma)
2662 }
2663
2664 // No non-Linux stub: the only caller (`Tensor::image_with_stride`)
2665 // returns `NotImplemented` directly on non-Linux without ever
2666 // reaching the storage layer, so defining a stub here would be
2667 // dead code and fail the `-D warnings` clippy gate on macOS CI.
2668
2669 /// Create a Mem-backed tensor storage with an explicit byte size that may
2670 /// exceed `shape.product() * sizeof(T)`. Used for image tensors with
2671 /// 64-byte-aligned row strides (see `MemTensor::with_capacity_bytes`).
2672 pub(crate) fn new_mem_with_byte_size(
2673 shape: &[usize],
2674 byte_size: usize,
2675 name: Option<&str>,
2676 ) -> Result<Self>
2677 where
2678 T: 'static,
2679 {
2680 MemTensor::<T>::with_capacity_bytes(shape, byte_size, name).map(TensorStorage::Mem)
2681 }
2682
2683 /// Create a Shm-backed tensor storage with an explicit byte size that may
2684 /// exceed `shape.product() * sizeof(T)`. Used for image tensors with
2685 /// 64-byte-aligned row strides (see `ShmTensor::new_with_byte_size`).
2686 #[cfg(unix)]
2687 pub(crate) fn new_shm_with_byte_size(
2688 shape: &[usize],
2689 byte_size: usize,
2690 name: Option<&str>,
2691 ) -> Result<Self> {
2692 ShmTensor::<T>::new_with_byte_size(shape, byte_size, name).map(TensorStorage::Shm)
2693 }
2694
2695 /// Allocate an image-formatted IOSurface-backed storage (macOS).
2696 ///
2697 /// Used by `Tensor::image()` when the caller requests
2698 /// `TensorMemory::DmaBuf` and the format has an IOSurface FourCC
2699 /// mapping (YUYV, RGBA, BGRA today). Falls back to `new_with_byte_size`
2700 /// otherwise.
2701 #[cfg(any(target_os = "macos", target_os = "ios"))]
2702 pub(crate) fn new_image_iosurface(
2703 width: usize,
2704 height: usize,
2705 format: PixelFormat,
2706 dtype: DType,
2707 shape: &[usize],
2708 name: Option<&str>,
2709 ) -> Result<Self> {
2710 IoSurfaceTensor::<T>::new_image(width, height, format, dtype, shape, name)
2711 .map(TensorStorage::Dma)
2712 }
2713
2714 /// Allocate an image-formatted AHardwareBuffer-backed storage (Android).
2715 ///
2716 /// Used by `Tensor::image()` when the caller requests
2717 /// `TensorMemory::DmaBuf` and the format has an AHardwareBuffer format
2718 /// mapping (RGBA8 and the RGBA16F float paths today). Falls back to
2719 /// `new_with_byte_size` otherwise.
2720 #[cfg(target_os = "android")]
2721 pub(crate) fn new_image_ahardwarebuffer(
2722 width: usize,
2723 height: usize,
2724 format: PixelFormat,
2725 dtype: DType,
2726 shape: &[usize],
2727 name: Option<&str>,
2728 access: CpuAccess,
2729 ) -> Result<Self> {
2730 AHardwareBufferTensor::<T>::new_image(width, height, format, dtype, shape, name, access)
2731 .map(TensorStorage::Dma)
2732 }
2733
2734 /// Allocate an image-formatted D3D11 texture storage (Windows).
2735 ///
2736 /// Used by `Tensor::image()` when the caller requests
2737 /// `TensorMemory::DmaBuf`. Unlike the IOSurface and AHardwareBuffer
2738 /// constructors there is no byte-bag fallback: a format with no texture
2739 /// layout is an error, because a texture is the only thing this backing
2740 /// can be.
2741 #[cfg(target_os = "windows")]
2742 pub(crate) fn new_image_d3d11(
2743 width: usize,
2744 height: usize,
2745 format: PixelFormat,
2746 dtype: DType,
2747 shape: &[usize],
2748 name: Option<&str>,
2749 access: CpuAccess,
2750 ) -> Result<Self> {
2751 crate::d3d11::texture::D3d11TextureTensor::<T>::new_image(
2752 width, height, format, dtype, shape, name, access,
2753 )
2754 .map(TensorStorage::Dma)
2755 }
2756
2757 /// Create a new tensor storage using the given file descriptor, shape,
2758 /// and optional name.
2759 #[cfg(unix)]
2760 fn from_fd(fd: OwnedFd, shape: &[usize], name: Option<&str>) -> Result<Self> {
2761 #[cfg(target_os = "linux")]
2762 {
2763 use nix::sys::stat::fstat;
2764 use nix::sys::statfs::fstatfs;
2765
2766 let stat = fstat(&fd)?;
2767 let major = major(stat.st_dev);
2768 let minor = minor(stat.st_dev);
2769
2770 if major != 0 {
2771 // Dma and Shm tensors are expected to have major number 0
2772 return Err(Error::UnknownDeviceType(major, minor));
2773 }
2774
2775 // Classify by filesystem magic, never by the st_dev minor.
2776 //
2777 // dma_buf files live on an internal kernel mount (`dma_buf_mnt`,
2778 // created with `kern_mount`) whose superblock draws its device
2779 // number from `get_anon_bdev()` — an IDA allocated first-come,
2780 // first-served during boot, shared with every other anonymous
2781 // pseudo-filesystem (pipefs, sockfs, anon_inodefs, nsfs, bdev,
2782 // tracefs, ...). The minor a DMA-BUF lands on is therefore a
2783 // function of which pseudo-filesystems registered first on that
2784 // kernel build: it moves with kernel config, driver load order,
2785 // initramfs use, and kernel version. It is not part of any ABI.
2786 //
2787 // Observed values for a genuine DMA-BUF: 12 on x86 desktop, 8 on
2788 // the ADIS Verdin. An earlier revision hardcoded `9 | 10` here,
2789 // which silently imported real DMA-BUFs as shared memory
2790 // everywhere else.
2791 //
2792 // The filesystem magic, by contrast, is stable UAPI
2793 // (include/uapi/linux/magic.h). Both branches are identified
2794 // positively; anything else is genuinely unknown and is rejected
2795 // rather than guessed at. Falling back to SHM would "work" — a
2796 // DMA-BUF is mmap-able, and even a pipe imports as a zero-length
2797 // tensor — which is exactly why it must not be the default.
2798 let magic = fs_magic(fstatfs(&fd)?.filesystem_type().0 as i64);
2799
2800 log::debug!(
2801 "Creating tensor from fd: major={major}, minor={minor}, magic={magic:#010x}"
2802 );
2803
2804 match magic {
2805 DMA_BUF_MAGIC => DmaTensor::<T>::from_fd(fd, shape, name).map(TensorStorage::Dma),
2806 TMPFS_MAGIC => ShmTensor::<T>::from_fd(fd, shape, name).map(TensorStorage::Shm),
2807 other => Err(Error::UnknownBufferType(other)),
2808 }
2809 }
2810 #[cfg(all(unix, not(target_os = "linux")))]
2811 {
2812 // On macOS/iOS/BSD, always use SHM (no DMA-BUF fd import)
2813 ShmTensor::<T>::from_fd(fd, shape, name).map(TensorStorage::Shm)
2814 }
2815 }
2816}
2817
2818#[cfg(not(edgefirst_tensor_dynamic))]
2819impl<T> TensorTrait<T> for TensorStorage<T>
2820where
2821 T: Num + Clone + fmt::Debug + Send + Sync,
2822{
2823 fn new(shape: &[usize], name: Option<&str>) -> Result<Self> {
2824 Self::new(shape, None, name)
2825 }
2826
2827 #[cfg(unix)]
2828 fn from_fd(fd: OwnedFd, shape: &[usize], name: Option<&str>) -> Result<Self> {
2829 Self::from_fd(fd, shape, name)
2830 }
2831
2832 #[cfg(unix)]
2833 fn clone_fd(&self) -> Result<OwnedFd> {
2834 match self {
2835 #[cfg(any(
2836 target_os = "linux",
2837 target_os = "macos",
2838 target_os = "ios",
2839 target_os = "android"
2840 ))]
2841 TensorStorage::Dma(t) => t.clone_fd(),
2842 TensorStorage::Shm(t) => t.clone_fd(),
2843 TensorStorage::Mem(t) => t.clone_fd(),
2844 TensorStorage::Pbo(t) => t.clone_fd(),
2845 }
2846 }
2847
2848 fn memory(&self) -> TensorMemory {
2849 match self {
2850 #[cfg(any(
2851 target_os = "linux",
2852 target_os = "macos",
2853 target_os = "ios",
2854 target_os = "android",
2855 target_os = "windows"
2856 ))]
2857 TensorStorage::Dma(_) => TensorMemory::DmaBuf,
2858 #[cfg(unix)]
2859 TensorStorage::Shm(_) => TensorMemory::Shm,
2860 TensorStorage::Mem(_) => TensorMemory::Mem,
2861 TensorStorage::Pbo(_) => TensorMemory::Pbo,
2862 }
2863 }
2864
2865 fn name(&self) -> String {
2866 match self {
2867 #[cfg(any(
2868 target_os = "linux",
2869 target_os = "macos",
2870 target_os = "ios",
2871 target_os = "android",
2872 target_os = "windows"
2873 ))]
2874 TensorStorage::Dma(t) => t.name(),
2875 #[cfg(unix)]
2876 TensorStorage::Shm(t) => t.name(),
2877 TensorStorage::Mem(t) => t.name(),
2878 TensorStorage::Pbo(t) => t.name(),
2879 }
2880 }
2881
2882 fn shape(&self) -> &[usize] {
2883 match self {
2884 #[cfg(any(
2885 target_os = "linux",
2886 target_os = "macos",
2887 target_os = "ios",
2888 target_os = "android",
2889 target_os = "windows"
2890 ))]
2891 TensorStorage::Dma(t) => t.shape(),
2892 #[cfg(unix)]
2893 TensorStorage::Shm(t) => t.shape(),
2894 TensorStorage::Mem(t) => t.shape(),
2895 TensorStorage::Pbo(t) => t.shape(),
2896 }
2897 }
2898
2899 fn reshape(&mut self, shape: &[usize]) -> Result<()> {
2900 match self {
2901 #[cfg(any(
2902 target_os = "linux",
2903 target_os = "macos",
2904 target_os = "ios",
2905 target_os = "android",
2906 target_os = "windows"
2907 ))]
2908 TensorStorage::Dma(t) => t.reshape(shape),
2909 #[cfg(unix)]
2910 TensorStorage::Shm(t) => t.reshape(shape),
2911 TensorStorage::Mem(t) => t.reshape(shape),
2912 TensorStorage::Pbo(t) => t.reshape(shape),
2913 }
2914 }
2915
2916 fn capacity_bytes(&self) -> usize {
2917 match self {
2918 #[cfg(any(
2919 target_os = "linux",
2920 target_os = "macos",
2921 target_os = "ios",
2922 target_os = "android",
2923 target_os = "windows"
2924 ))]
2925 TensorStorage::Dma(t) => t.capacity_bytes(),
2926 #[cfg(unix)]
2927 TensorStorage::Shm(t) => t.capacity_bytes(),
2928 TensorStorage::Mem(t) => t.capacity_bytes(),
2929 TensorStorage::Pbo(t) => t.capacity_bytes(),
2930 }
2931 }
2932
2933 fn set_logical_shape(&mut self, shape: &[usize]) -> Result<()> {
2934 match self {
2935 #[cfg(any(
2936 target_os = "linux",
2937 target_os = "macos",
2938 target_os = "ios",
2939 target_os = "android",
2940 target_os = "windows"
2941 ))]
2942 TensorStorage::Dma(t) => t.set_logical_shape(shape),
2943 #[cfg(unix)]
2944 TensorStorage::Shm(t) => t.set_logical_shape(shape),
2945 TensorStorage::Mem(t) => t.set_logical_shape(shape),
2946 TensorStorage::Pbo(t) => t.set_logical_shape(shape),
2947 }
2948 }
2949
2950 fn map_with<'a>(&self, access: CpuAccess) -> Result<crate::view::HostView<'a, T>>
2951 where
2952 T: 'a,
2953 {
2954 match self {
2955 #[cfg(any(
2956 target_os = "linux",
2957 target_os = "macos",
2958 target_os = "ios",
2959 target_os = "android",
2960 target_os = "windows"
2961 ))]
2962 TensorStorage::Dma(t) => t.map_with(access),
2963 #[cfg(unix)]
2964 TensorStorage::Shm(t) => t.map_with(access),
2965 TensorStorage::Mem(t) => t.map_with(access),
2966 TensorStorage::Pbo(t) => t.map_with(access),
2967 }
2968 }
2969
2970 fn try_map_with<'a>(&self, access: CpuAccess) -> Result<crate::view::HostView<'a, T>>
2971 where
2972 T: 'a,
2973 {
2974 match self {
2975 #[cfg(any(
2976 target_os = "linux",
2977 target_os = "macos",
2978 target_os = "ios",
2979 target_os = "android",
2980 target_os = "windows"
2981 ))]
2982 TensorStorage::Dma(t) => t.try_map_with(access),
2983 #[cfg(unix)]
2984 TensorStorage::Shm(t) => t.try_map_with(access),
2985 TensorStorage::Mem(t) => t.try_map_with(access),
2986 TensorStorage::Pbo(t) => t.try_map_with(access),
2987 }
2988 }
2989
2990 fn buffer_identity(&self) -> &BufferIdentity {
2991 match self {
2992 #[cfg(any(
2993 target_os = "linux",
2994 target_os = "macos",
2995 target_os = "ios",
2996 target_os = "android",
2997 target_os = "windows"
2998 ))]
2999 TensorStorage::Dma(t) => t.buffer_identity(),
3000 #[cfg(unix)]
3001 TensorStorage::Shm(t) => t.buffer_identity(),
3002 TensorStorage::Mem(t) => t.buffer_identity(),
3003 TensorStorage::Pbo(t) => t.buffer_identity(),
3004 }
3005 }
3006
3007 /// Forward a sub-region view to the active backend, re-wrapping the
3008 /// backend's view (which shares the parent's allocation and
3009 /// [`BufferIdentity`]) back into the matching `TensorStorage` variant. Each
3010 /// backend's `view` is its own `TensorTrait::view` override; this single
3011 /// match is the only per-variant dispatch (`Tensor::subview` calls through
3012 /// here rather than matching the storage itself).
3013 fn view(&self, offset_bytes: usize, shape: &[usize]) -> Result<Self> {
3014 match self {
3015 #[cfg(any(
3016 target_os = "linux",
3017 target_os = "macos",
3018 target_os = "ios",
3019 target_os = "android",
3020 target_os = "windows"
3021 ))]
3022 TensorStorage::Dma(t) => t.view(offset_bytes, shape).map(TensorStorage::Dma),
3023 #[cfg(unix)]
3024 TensorStorage::Shm(t) => t.view(offset_bytes, shape).map(TensorStorage::Shm),
3025 TensorStorage::Mem(t) => t.view(offset_bytes, shape).map(TensorStorage::Mem),
3026 TensorStorage::Pbo(t) => t.view(offset_bytes, shape).map(TensorStorage::Pbo),
3027 }
3028 }
3029}
3030
3031/// Multi-backend tensor with optional image format metadata.
3032///
3033/// When `format` is `Some`, this tensor represents an image. Width, height,
3034/// and channels are derived from `shape` + `format`. When `format` is `None`,
3035/// this is a raw tensor (identical to the pre-refactoring behavior).
3036///
3037/// This is the `static` backend's `Tensor<T>` -- it owns its storage. See
3038/// `tensor_dyn/dynamic_tensor.rs` for `dynamic`'s `Tensor<T>`, a typed lens
3039/// with no storage of its own.
3040#[cfg(not(edgefirst_tensor_dynamic))]
3041#[derive(Debug)]
3042pub struct Tensor<T>
3043where
3044 T: Num + Clone + fmt::Debug + Send + Sync,
3045{
3046 /// CUDA registration for this tensor, if any. Set after creation by
3047 /// the image crate once a PBO is registered with CUDA interop.
3048 ///
3049 /// MUST be declared before `storage`: CUDA must unregister the GL buffer
3050 /// before storage's Drop deletes it (cudaGraphicsUnregisterResource before
3051 /// glDeleteBuffers). Rust drops fields in declaration order.
3052 cuda: Option<crate::cuda::CudaHandle>,
3053 pub(crate) storage: TensorStorage<T>,
3054 format: Option<PixelFormat>,
3055 chroma: Option<Box<Tensor<T>>>,
3056 /// Row stride in bytes for externally allocated buffers with row padding.
3057 /// `None` means tightly packed (stride == width * bytes_per_pixel).
3058 row_stride: Option<usize>,
3059 /// Byte offset within the DMA-BUF where image data starts.
3060 /// `None` means offset 0 (data starts at the beginning of the buffer).
3061 plane_offset: Option<usize>,
3062 /// Quantization metadata for integer-typed tensors. Public access is
3063 /// gated by the `IntegerType` trait — `Tensor<f32>` etc. carry the
3064 /// field for layout uniformity but have no way to read or write it.
3065 pub(crate) quantization: Option<Quantization>,
3066 /// Optional colorimetry metadata. `None` = undefined; never auto-filled.
3067 colorimetry: Option<crate::Colorimetry>,
3068 /// Declared CPU access (see [`CpuAccess`]). Image constructors set it
3069 /// from their `access` parameter; non-image tensors (`new`, imports,
3070 /// numpy) default to `ReadWrite` — they are CPU-centric by nature.
3071 /// Maps, pins and sync brackets beyond this declaration are refused.
3072 cpu_access: CpuAccess,
3073 /// Recorded vendor tile-compression scheme — best knowledge from
3074 /// allocation time (see [`Compression`]). `Some` only for Android
3075 /// hardware-only AHardwareBuffers whose allocation requested
3076 /// compression on an eligible format with a recognized vendor. When
3077 /// set, the row-stride accessors describe no meaningful linear
3078 /// layout. `configure_image` preserves it (physical layout, unlike
3079 /// colorimetry); views inherit it.
3080 compression: Option<CompressionScheme>,
3081 /// Parent-image snapshot when this tensor is a [`view`](Self::view)/
3082 /// [`batch`](Self::batch) sub-region; `None` for a whole tensor. Lets the GL
3083 /// backend key its import on the parent and render the view as a
3084 /// `glViewport`/`glScissor` ROI. See [`ViewOrigin`].
3085 view_origin: Option<ViewOrigin>,
3086 /// Set by [`mark_view`](Self::mark_view) on a tensor rebuilt from a view's
3087 /// descriptor, which addresses part of a larger buffer but has no
3088 /// [`ViewOrigin`], and inherited by its sub-views. See
3089 /// [`is_view`](Self::is_view).
3090 view_marked: bool,
3091}
3092
3093#[cfg(not(edgefirst_tensor_dynamic))]
3094impl<T> Tensor<T>
3095where
3096 T: Num + Clone + fmt::Debug + Send + Sync,
3097{
3098 /// Wrap a TensorStorage in a Tensor with no image metadata.
3099 pub(crate) fn wrap(storage: TensorStorage<T>) -> Self {
3100 Self {
3101 storage,
3102 format: None,
3103 chroma: None,
3104 row_stride: None,
3105 plane_offset: None,
3106 quantization: None,
3107 cuda: None,
3108 colorimetry: None,
3109 cpu_access: CpuAccess::ReadWrite,
3110 compression: None,
3111 view_origin: None,
3112 view_marked: false,
3113 }
3114 }
3115
3116 /// Construct a tensor from a row-major element slice + shape. Allocates a
3117 /// new buffer (`TensorMemory::Mem`) and memcpys the contents; caller
3118 /// retains ownership of the input slice.
3119 ///
3120 /// # Errors
3121 ///
3122 /// - [`Error::InvalidShape`] if `values.len() != shape.iter().product()`.
3123 /// - Propagates any allocation error from [`Self::new`].
3124 pub fn from_slice(values: &[T], shape: &[usize]) -> Result<Self>
3125 where
3126 T: Copy,
3127 {
3128 let expected: usize = shape.iter().product();
3129 if values.len() != expected {
3130 return Err(Error::InvalidShape(format!(
3131 "from_slice: values.len()={} but shape product={expected} (shape={shape:?})",
3132 values.len()
3133 )));
3134 }
3135 let t = Self::new(shape, Some(TensorMemory::Mem), None)?;
3136 {
3137 let mut m = t.map()?;
3138 m.as_mut_slice().copy_from_slice(values);
3139 }
3140 Ok(t)
3141 }
3142
3143 /// Wrap externally-owned memory as a tensor without copying. The tensor
3144 /// borrows `[ptr, ptr + shape.product() * size_of::<T>())` as
3145 /// [`TensorMemory::Mem`]; `owner`, when `Some`, co-owns the source so it
3146 /// outlives the tensor (and all derived views/maps). See [`ForeignOwner`].
3147 ///
3148 /// The canonical use is CUDA zero-copy: allocate host-coherent memory
3149 /// (`cudaHostAlloc`), wrap the host pointer here, and bind the matching
3150 /// device pointer to the inference engine — reads and writes hit the same
3151 /// physical buffer with no host copy. The identical primitive backs the
3152 /// Python `Tensor.from_numpy` zero-copy borrow (owner = the NumPy object).
3153 ///
3154 /// # Safety
3155 ///
3156 /// `ptr` must be non-null, aligned to `align_of::<T>()`, and valid for
3157 /// `shape.product()` elements of `T` for as long as the returned tensor —
3158 /// and every view/map sharing its backing — is alive. Pass an `owner` that
3159 /// co-owns the source to uphold that contract.
3160 ///
3161 /// # Errors
3162 ///
3163 /// [`Error::InvalidSize`] if `shape` is empty.
3164 pub unsafe fn from_foreign(
3165 ptr: *mut T,
3166 shape: &[usize],
3167 owner: Option<crate::ForeignOwner>,
3168 name: Option<&str>,
3169 ) -> Result<Self> {
3170 // SAFETY: forwarded to `from_foreign_with_capacity`'s contract; `0`
3171 // clamps to the tight shape-derived footprint below, i.e. no extra
3172 // headroom -- identical to this function's own prior behaviour.
3173 unsafe { Self::from_foreign_with_capacity(ptr, shape, 0, owner, name) }
3174 }
3175
3176 /// [`Self::from_foreign`], additionally recording that the allocation
3177 /// behind `ptr` holds at least `capacity_bytes` bytes -- headroom beyond
3178 /// the tight `shape.product() * size_of::<T>()` footprint that a later
3179 /// [`configure_image`](Self::configure_image) / `set_logical_shape` can
3180 /// grow into without erroring.
3181 ///
3182 /// This is the consumer half of the cross-package capsule protocol's
3183 /// [`crate::protocol::TensorDesc::capacity`]: a producer's pool tensor
3184 /// (allocated for its largest expected image, holding a smaller one
3185 /// today) reports its real allocation size so the imported alias
3186 /// inherits the same headroom instead of being clamped to today's
3187 /// logical shape. `capacity_bytes` below the tight footprint is clamped
3188 /// up to it, never down -- the shape-implied floor always holds.
3189 ///
3190 /// # Safety
3191 ///
3192 /// Same as [`Self::from_foreign`], except `ptr` must be valid for
3193 /// `capacity_bytes` bytes (or the tight shape footprint, whichever is
3194 /// larger), not just the shape-implied footprint.
3195 ///
3196 /// # Errors
3197 ///
3198 /// [`Error::InvalidSize`] if `shape` is empty.
3199 pub unsafe fn from_foreign_with_capacity(
3200 ptr: *mut T,
3201 shape: &[usize],
3202 capacity_bytes: usize,
3203 owner: Option<crate::ForeignOwner>,
3204 name: Option<&str>,
3205 ) -> Result<Self> {
3206 if shape.is_empty() {
3207 return Err(Error::InvalidSize(0));
3208 }
3209 if ptr.is_null() {
3210 return Err(Error::InvalidArgument(
3211 "from_foreign: ptr must be non-null".to_owned(),
3212 ));
3213 }
3214 shape
3215 .iter()
3216 .copied()
3217 .try_fold(1usize, usize::checked_mul)
3218 .ok_or_else(|| {
3219 Error::InvalidArgument(format!(
3220 "from_foreign: shape.product() overflows usize (shape={shape:?})"
3221 ))
3222 })?;
3223 let capacity_elems = capacity_bytes / std::mem::size_of::<T>();
3224 // SAFETY: this function's own contract (see its doc comment) is
3225 // exactly `MemTensor::from_foreign_with_capacity`'s contract.
3226 let mem = unsafe {
3227 MemTensor::<T>::from_foreign_with_capacity(ptr, shape, capacity_elems, owner, name)
3228 };
3229 Ok(Self::wrap(TensorStorage::Mem(mem)))
3230 }
3231
3232 /// Pin a stable host address for this tensor's data.
3233 ///
3234 /// The returned [`HostPin`] carries **no borrow of the tensor**, so a
3235 /// caller can pin once at init and still pass `&mut self` to
3236 /// `ImageProcessor::convert` every frame. That borrow conflict is the
3237 /// blocker reported in [#134](https://github.com/EdgeFirstAI/hal/issues/134).
3238 ///
3239 /// The address stays valid until the pin is dropped, and survives
3240 /// re-layout: `configure_image` and `set_logical_shape` only rewrite the
3241 /// logical shape after a capacity check, and every backend's pin owns its
3242 /// own mapping, so nothing can move underneath it.
3243 ///
3244 /// A pin is therefore a **snapshot**, not a tracker: [`HostPin::len`]
3245 /// reports the extent at pin time and does not follow a later re-layout.
3246 /// That is the safe direction — the window always stays inside the
3247 /// allocation — but it means a caller who re-layouts should re-pin if it
3248 /// wants the new extent.
3249 ///
3250 /// The address is **not** a coherency guarantee: on non-coherent backends
3251 /// bracket CPU access with
3252 /// [`sync_for_cpu`](Self::sync_for_cpu) and
3253 /// [`sync_for_device`](Self::sync_for_device), which are no-ops where the
3254 /// backend does not require them.
3255 ///
3256 /// # Cost
3257 ///
3258 /// On `Mem` this is an `Arc` refcount bump — the backing allocation is
3259 /// never resized, so its base pointer is already stable. Other backends
3260 /// establish a mapping once and reuse it.
3261 ///
3262 /// # Errors
3263 ///
3264 /// - [`Error::InvalidArgument`] when `access` is [`CpuAccess::None`]: a
3265 /// pinned mapping is CPU access by definition, so this is an error
3266 /// rather than a warn-and-stage.
3267 /// - [`Error::CpuAccessDenied`] when the tensor's declared
3268 /// [`CpuAccess`] does not cover `access`.
3269 /// - [`Error::NotImplemented`] for the two backends that **cannot** pin,
3270 /// rather than merely not having been wired yet: a PBO
3271 /// (`glMapBufferRange` establishes the address and the visibility
3272 /// together, and there is no `GL_MAP_PERSISTENT_BIT` in GLES) and an
3273 /// AHardwareBuffer (`AHardwareBuffer_lock` is what yields the address at
3274 /// all, so there is no address outside the lock). Both name
3275 /// `map()`/`map_with()` as the alternative. `Mem`, `Shm`, DMA-BUF,
3276 /// IOSurface and the Windows D3D11 texture all pin — the texture through
3277 /// a persistent host buffer the storage owns, which
3278 /// [`sync_for_cpu`](Self::sync_for_cpu) and
3279 /// [`sync_for_device`](Self::sync_for_device) move bytes between and the
3280 /// texture, rather than through the staging mapping (a pin outlives every
3281 /// map guard, so holding the staging texture mapped would freeze it).
3282 pub fn pin_host<'a>(&self, access: CpuAccess) -> Result<HostPin<'a>>
3283 where
3284 T: 'a,
3285 {
3286 self.cpu_access.require_direction(access, "pin_host")?;
3287 // Backends yield everything addressable from the tensor's offset;
3288 // narrow to the logical extent so a consumer cannot mistake pitch
3289 // padding for data. A map guard keeps the wider window.
3290 let logical = self.len() * std::mem::size_of::<T>();
3291 let pin = match &self.storage {
3292 TensorStorage::Mem(m) => Ok(m.host_pin()),
3293 #[cfg(unix)]
3294 TensorStorage::Shm(s) => s.host_pin(),
3295 #[cfg(target_os = "linux")]
3296 TensorStorage::Dma(d) => d.host_pin(),
3297 #[cfg(any(target_os = "macos", target_os = "ios"))]
3298 TensorStorage::Dma(s) => s.host_pin(),
3299 // Windows D3D11: the pin is the storage's own persistent host
3300 // buffer, not the staging mapping -- see `D3d11TextureTensor::host_pin`.
3301 // The direction goes with it: a writable pin is what tells
3302 // `sync_for_device` the host buffer, not the staging copy, carries
3303 // the caller's writes.
3304 #[cfg(target_os = "windows")]
3305 TensorStorage::Dma(t) => t.host_pin(access),
3306 // AHardwareBuffer cannot support a pin at all, and that is a
3307 // platform contract rather than missing work: the NDK exposes no
3308 // way to obtain a CPU address outside AHardwareBuffer_lock, so a
3309 // "persistent" address would mean holding the lock for the pin's
3310 // lifetime -- which blocks the GPU from the buffer and makes any
3311 // further lock undefined. Both are the opposite of what a pin is
3312 // for. Say so, rather than implying a TODO.
3313 #[cfg(target_os = "android")]
3314 TensorStorage::Dma(_) => Err(Error::NotImplemented(
3315 "pin_host: AHardwareBuffer has no host address outside \
3316 AHardwareBuffer_lock, so it cannot hand out one that outlives \
3317 a guard. Use map()/map_with() for CPU access; for zero-copy \
3318 handoff pass the buffer itself."
3319 .to_owned(),
3320 )),
3321 // PBO is the same shape of constraint as AHardwareBuffer. The
3322 // message is a shared constant because `TensorDyn::pin_host`
3323 // (`tensor_dyn/dynamic_backend.rs`) must repeat this refusal --
3324 // the ABI has `ef_tensor_map` and no `ef_tensor_pin_host`, so the
3325 // library cannot tell a pin from a map. See
3326 // `crate::pbo::PIN_HOST_PBO_REFUSAL` for the reasoning at length.
3327 TensorStorage::Pbo(_) => Err(Error::NotImplemented(
3328 crate::pbo::PIN_HOST_PBO_REFUSAL.to_owned(),
3329 )),
3330 }?;
3331 Ok(pin.narrowed(logical))
3332 }
3333
3334 /// Make CPU reads/writes of this tensor coherent — the CPU is about to
3335 /// access the buffer.
3336 ///
3337 /// Corresponds to `DMA_BUF_SYNC_START`, and the name matches the Linux
3338 /// kernel DMA API for the same reason: `for_cpu` says who is about to
3339 /// touch the memory, which is what the caller actually knows.
3340 ///
3341 /// Pair with [`sync_for_device`](Self::sync_for_device).
3342 ///
3343 /// # Which backends owe maintenance
3344 ///
3345 /// | backend | behaviour |
3346 /// |---|---|
3347 /// | `Mem`, `Shm` | **no-op** — plain CPU memory, coherent by construction |
3348 /// | `Dma` (Linux DMA-BUF) | one `DMA_BUF_IOCTL_SYNC`, direction-tagged |
3349 /// | `Dma` (macOS/iOS IOSurface) | `IOSurfaceLock` / `IOSurfaceUnlock`, options matched to `access` |
3350 /// | `Dma` (Android AHardwareBuffer) | [`Error::NotImplemented`] — no sync independent of the lock |
3351 /// | `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 |
3352 /// | `Pbo` | [`Error::NotImplemented`] — no coherency window independent of the map |
3353 ///
3354 /// Only the first row is a no-op, so "bracket unconditionally and pay
3355 /// nothing" holds for `Mem`/`Shm` alone. The two that report
3356 /// `NotImplemented` do so because their platform fuses addressing and
3357 /// coherency into one call — `AHardwareBuffer_lock` and
3358 /// `glMapBufferRange` each perform the maintenance *as* they hand out the
3359 /// address, so a separate bracket would have to nest inside a map guard
3360 /// that already owns the pair. Both name `map()`/`map_with()` as the
3361 /// alternative. They previously fell through a wildcard and answered
3362 /// `Ok(())`, which is worse than an error because the symptom is
3363 /// stale-but-plausible pixels rather than a failure. The match is
3364 /// exhaustive so a new backend cannot inherit that silence — it must state
3365 /// its coherency story to compile.
3366 ///
3367 /// # Cost
3368 ///
3369 /// `Mem`/`Shm`: nothing. `Dma` on Linux: one ioctl. IOSurface: one
3370 /// lock/unlock pair. Windows D3D11 texture: one `CopyResource` plus a
3371 /// staging map per direction, and nothing at all on a tensor no one has
3372 /// pinned.
3373 ///
3374 /// # Errors
3375 ///
3376 /// [`Error::InvalidArgument`] when `access` is [`CpuAccess::None`], which
3377 /// describes no access to make coherent; [`Error::CpuAccessDenied`] when
3378 /// the tensor's declared [`CpuAccess`] does not cover `access`.
3379 pub fn sync_for_cpu(&self, access: CpuAccess) -> Result<()> {
3380 self.check_sync_access(access)?;
3381 match &self.storage {
3382 // Plain CPU memory: coherent by construction.
3383 TensorStorage::Mem(_) => Ok(()),
3384 #[cfg(unix)]
3385 TensorStorage::Shm(_) => Ok(()),
3386 #[cfg(target_os = "linux")]
3387 TensorStorage::Dma(d) => {
3388 // The direction tells the kernel which maintenance this access
3389 // needs: a read-only bracket skips the writeback at END, a
3390 // write-only one skips the invalidate at START. It MUST match
3391 // the direction passed to sync_for_device.
3392 crate::dmabuf::sync_access(&d.fd, true, access).map_err(Error::NixError)
3393 }
3394 #[cfg(any(target_os = "macos", target_os = "ios"))]
3395 TensorStorage::Dma(s) => s.sync_for_cpu(access),
3396 // Windows D3D11: refreshes the pinned host buffer from the texture
3397 // through the staging copy.
3398 #[cfg(target_os = "windows")]
3399 TensorStorage::Dma(t) => t.sync_for_cpu(access),
3400 #[cfg(target_os = "android")]
3401 TensorStorage::Dma(_) => Err(Error::NotImplemented(
3402 "sync_for_cpu: AHardwareBuffer has no sync independent of its \
3403 lock. AHardwareBuffer_lock both establishes the address and \
3404 performs the cache invalidate, and the NDK leaves locking an \
3405 already-locked buffer undefined, so a bracket on top of a map \
3406 guard would nest. Use map()/map_with(), which owns the pair."
3407 .to_owned(),
3408 )),
3409 TensorStorage::Pbo(_) => Err(Error::NotImplemented(
3410 "sync_for_cpu: a PBO has no coherency window independent of its \
3411 map -- glMapBufferRange establishes the address and the \
3412 visibility together, and glUnmapBuffer publishes the writes. \
3413 Use map()/map_with(), which owns the pair."
3414 .to_owned(),
3415 )),
3416 }
3417 }
3418
3419 /// Release the buffer back to the device — the CPU is done accessing it.
3420 ///
3421 /// Corresponds to `DMA_BUF_SYNC_END`. See
3422 /// [`sync_for_cpu`](Self::sync_for_cpu) for the pairing and cost.
3423 pub fn sync_for_device(&self, access: CpuAccess) -> Result<()> {
3424 self.check_sync_access(access)?;
3425 match &self.storage {
3426 TensorStorage::Mem(_) => Ok(()),
3427 #[cfg(unix)]
3428 TensorStorage::Shm(_) => Ok(()),
3429 #[cfg(target_os = "linux")]
3430 TensorStorage::Dma(d) => {
3431 crate::dmabuf::sync_access(&d.fd, false, access).map_err(Error::NixError)
3432 }
3433 #[cfg(any(target_os = "macos", target_os = "ios"))]
3434 TensorStorage::Dma(s) => s.sync_for_device(access),
3435 // Windows D3D11: uploads the pinned host buffer to the texture.
3436 #[cfg(target_os = "windows")]
3437 TensorStorage::Dma(t) => t.sync_for_device(access),
3438 #[cfg(target_os = "android")]
3439 TensorStorage::Dma(_) => Err(Error::NotImplemented(
3440 "sync_for_device: AHardwareBuffer has no sync independent of \
3441 its lock — see sync_for_cpu. Use map()/map_with()."
3442 .to_owned(),
3443 )),
3444 TensorStorage::Pbo(_) => Err(Error::NotImplemented(
3445 "sync_for_device: a PBO has no coherency window independent of \
3446 its map -- see sync_for_cpu. Use map()/map_with()."
3447 .to_owned(),
3448 )),
3449 }
3450 }
3451
3452 fn check_sync_access(&self, access: CpuAccess) -> Result<()> {
3453 self.cpu_access
3454 .require_direction(access, "sync_for_cpu/sync_for_device")
3455 }
3456
3457 /// Construct a tensor from a 3-D ndarray view. Respects strides — one
3458 /// copy in all cases; contiguous views take a memcpy fast path.
3459 ///
3460 /// Only available when the `ndarray` feature is enabled.
3461 #[cfg(feature = "ndarray")]
3462 pub fn from_arrayview3(view: ndarray::ArrayView3<'_, T>) -> Result<Self>
3463 where
3464 T: Copy,
3465 {
3466 let (h, w, c) = view.dim();
3467 let t = Self::new(&[h, w, c], Some(TensorMemory::Mem), None)?;
3468 {
3469 let mut m = t.map()?;
3470 let dst = m.as_mut_slice();
3471 if let Some(src) = view.as_slice() {
3472 dst.copy_from_slice(src);
3473 } else {
3474 for (d, &s) in dst.iter_mut().zip(view.iter()) {
3475 *d = s;
3476 }
3477 }
3478 }
3479 Ok(t)
3480 }
3481
3482 /// Create a new tensor with the given shape, memory type, and optional
3483 /// name. If no name is given, a random name will be generated. If no
3484 /// memory type is given, the best available memory type will be chosen
3485 /// based on the platform and environment variables.
3486 ///
3487 /// On Linux platforms, the order of preference is: Dma -> Shm -> Mem.
3488 /// On other Unix platforms (macOS), the order is: Shm -> Mem.
3489 /// On non-Unix platforms, only Mem is available.
3490 ///
3491 /// # Environment Variables
3492 /// - `EDGEFIRST_TENSOR_FORCE_MEM`: If set to a non-zero and non-false
3493 /// value, forces the use of regular system memory allocation
3494 /// (`TensorMemory::Mem`) regardless of platform capabilities.
3495 ///
3496 /// # Example
3497 /// ```rust
3498 /// use edgefirst_tensor::{Error, Tensor, TensorMemory, TensorTrait};
3499 /// # fn main() -> Result<(), Error> {
3500 /// let tensor = Tensor::<f32>::new(&[2, 3, 4], Some(TensorMemory::Mem), Some("test_tensor"))?;
3501 /// assert_eq!(tensor.memory(), TensorMemory::Mem);
3502 /// assert_eq!(tensor.name(), "test_tensor");
3503 /// # Ok(())
3504 /// # }
3505 /// ```
3506 pub fn new(shape: &[usize], memory: Option<TensorMemory>, name: Option<&str>) -> Result<Self> {
3507 let _span = tracing::trace_span!(
3508 "tensor.alloc",
3509 ?shape,
3510 memory = ?memory,
3511 dtype = std::any::type_name::<T>(),
3512 )
3513 .entered();
3514 #[cfg_attr(not(target_os = "linux"), allow(unused_mut))]
3515 let mut t = TensorStorage::new(shape, memory, name).map(Self::wrap)?;
3516 // Best-effort: attach a CUDA ExternalMemory handle for DMA tensors on
3517 // CUDA-capable hosts. Never blocks tensor creation on failure.
3518 // RUNTIME-UNVALIDATED: no CUDA+dma_heap test platform available; ABI
3519 // layout-asserted vs. CUDA 12.6 driver_types.h; mechanism proven by
3520 // gpu-probe O5 on Orin.
3521 #[cfg(target_os = "linux")]
3522 t.try_init_dma_cuda();
3523 Ok(t)
3524 }
3525
3526 /// Import an existing buffer as a tensor, taking ownership of its file
3527 /// descriptor. No bytes are copied.
3528 #[cfg(unix)]
3529 pub fn from_fd(fd: std::os::fd::OwnedFd, shape: &[usize], name: Option<&str>) -> Result<Self> {
3530 #[cfg_attr(not(target_os = "linux"), allow(unused_mut))]
3531 let mut t = Self::wrap(TensorStorage::from_fd(fd, shape, name)?);
3532 // Best-effort CUDA external memory import for DMA-backed tensors.
3533 // RUNTIME-UNVALIDATED: see try_init_dma_cuda().
3534 #[cfg(target_os = "linux")]
3535 t.try_init_dma_cuda();
3536 Ok(t)
3537 }
3538
3539 /// Create an image tensor with the given format.
3540 /// Allocate an image tensor from a declarative request — the
3541 /// full-featured constructor behind [`Self::image`] and friends.
3542 ///
3543 /// Adds the [`Compression`] request to the classic parameters:
3544 ///
3545 /// - a request with any CPU access other than [`CpuAccess::None`] is
3546 /// [`Error::InvalidArgument`] (CPU mapping pins the layout linear);
3547 /// - [`Compression::Scheme`] fails with [`Error::NotImplemented`] on
3548 /// platforms without vendor tile compression, and with
3549 /// [`Error::InvalidArgument`] when the device's native scheme or
3550 /// the format eligibility doesn't match;
3551 /// - [`Compression::Any`] never fails for compression reasons: it
3552 /// records the scheme when the allocation is eligible and
3553 /// otherwise resolves linear, incrementing
3554 /// [`compression_fallback_count`].
3555 ///
3556 /// The recorded outcome is readable via [`Tensor::compression`].
3557 ///
3558 /// A [`ImageDesc::with_contiguous`] request allocates from the CMA heap
3559 /// only, and fails instead of falling back; [`Tensor::contiguity`]
3560 /// reports the heap.
3561 pub fn image_desc(desc: &ImageDesc) -> Result<Self>
3562 where
3563 T: 'static,
3564 {
3565 let Some(t_dtype) = dtype_of::<T>() else {
3566 return Err(Error::InvalidArgument(
3567 "image_desc: element type has no DType mapping".into(),
3568 ));
3569 };
3570 if t_dtype != desc.dtype {
3571 return Err(Error::InvalidArgument(format!(
3572 "image_desc: desc.dtype is {:?} but the tensor element type is {t_dtype:?}",
3573 desc.dtype
3574 )));
3575 }
3576
3577 // A contiguous request is a CMA DMA-BUF or an error: no fallback to
3578 // another heap or another kind of memory.
3579 if desc.contiguous {
3580 #[cfg(not(target_os = "linux"))]
3581 {
3582 return Err(Error::NotImplemented(
3583 "image_desc: physically contiguous memory needs the Linux CMA DMA heap".into(),
3584 ));
3585 }
3586 #[cfg(target_os = "linux")]
3587 if !matches!(desc.memory, None | Some(TensorMemory::DmaBuf)) {
3588 return Err(Error::InvalidArgument(format!(
3589 "image_desc: physically contiguous memory is DMA-BUF from the CMA heap, \
3590 but {:?} was requested",
3591 desc.memory
3592 )));
3593 }
3594 }
3595
3596 // Compression-request guards. CPU access pins the layout linear,
3597 // so a request combined with any declared access is a
3598 // contradiction the caller should hear about immediately.
3599 if desc.compression.is_some() && desc.access != CpuAccess::None {
3600 return Err(Error::InvalidArgument(format!(
3601 "image_desc: a compression request requires CpuAccess::None (declared {:?}) — CPU mapping pins the layout linear",
3602 desc.access
3603 )));
3604 }
3605 if let Some(Compression::Scheme(requested)) = desc.compression {
3606 #[cfg(not(target_os = "android"))]
3607 {
3608 return Err(Error::NotImplemented(format!(
3609 "image_desc: Compression::Scheme({requested:?}) — no vendor tile compression on this platform (request Compression::Any for a portable fallback)"
3610 )));
3611 }
3612 #[cfg(target_os = "android")]
3613 {
3614 if matches!(
3615 desc.memory,
3616 Some(TensorMemory::Mem) | Some(TensorMemory::Shm) | Some(TensorMemory::Pbo)
3617 ) {
3618 return Err(Error::InvalidArgument(format!(
3619 "image_desc: Compression::Scheme({requested:?}) requires hardware memory (TensorMemory::DmaBuf or auto-select), got {:?}",
3620 desc.memory
3621 )));
3622 }
3623 if !crate::ahardwarebuffer_layout::compression_eligible(desc.format, desc.dtype) {
3624 return Err(Error::InvalidArgument(format!(
3625 "image_desc: ({:?}, {:?}) is not compression-eligible (RGBA8888 u8/i8 initially)",
3626 desc.format, desc.dtype
3627 )));
3628 }
3629 let device = crate::ahardwarebuffer::device_compression_scheme();
3630 if device != Some(requested) {
3631 return Err(Error::InvalidArgument(format!(
3632 "image_desc: Compression::Scheme({requested:?}) but the device's native scheme is {device:?}"
3633 )));
3634 }
3635 }
3636 }
3637
3638 // A compression request implies a hardware pipeline, so auto
3639 // memory promotes to the platform's zero-copy allocation first
3640 // (`Tensor::image` only takes the AHardwareBuffer/IOSurface path
3641 // under an explicit Dma request). `Scheme` propagates the Dma
3642 // failure — the caller demanded a layout only that allocator can
3643 // produce; `Any` falls back to plain auto-select.
3644 #[allow(unused_mut)]
3645 let mut t = if desc.contiguous {
3646 Self::image_impl(
3647 desc.width,
3648 desc.height,
3649 desc.format,
3650 Some(TensorMemory::DmaBuf),
3651 desc.access,
3652 true,
3653 )?
3654 } else {
3655 match (desc.memory, desc.compression) {
3656 (None, Some(request)) => {
3657 match Self::image(
3658 desc.width,
3659 desc.height,
3660 desc.format,
3661 Some(TensorMemory::DmaBuf),
3662 desc.access,
3663 ) {
3664 Ok(t) => t,
3665 Err(e) if matches!(request, Compression::Scheme(_)) => return Err(e),
3666 Err(_) => {
3667 Self::image(desc.width, desc.height, desc.format, None, desc.access)?
3668 }
3669 }
3670 }
3671 (memory, _) => {
3672 Self::image(desc.width, desc.height, desc.format, memory, desc.access)?
3673 }
3674 }
3675 };
3676
3677 // Record best knowledge / count fallbacks. Only an Android
3678 // hardware-only AHardwareBuffer allocation can actually hold a
3679 // vendor tile layout; everywhere else an Any request resolves
3680 // linear and is counted.
3681 if let Some(request) = desc.compression {
3682 #[cfg(target_os = "android")]
3683 {
3684 let eligible =
3685 crate::ahardwarebuffer_layout::compression_eligible(desc.format, desc.dtype);
3686 let scheme = crate::ahardwarebuffer::device_compression_scheme();
3687 let is_ahb = t.memory() == TensorMemory::DmaBuf;
3688 match (request, scheme) {
3689 (_, Some(s)) if eligible && is_ahb => {
3690 t.set_compression_unchecked(Some(s));
3691 }
3692 (Compression::Scheme(requested), _) => {
3693 // Pre-validated eligible + scheme match, so the only
3694 // way here is the allocation resolving off-AHB.
3695 return Err(Error::InvalidOperation(format!(
3696 "image_desc: Compression::Scheme({requested:?}) requested but the allocation resolved to {:?} (not an AHardwareBuffer)",
3697 t.memory()
3698 )));
3699 }
3700 (Compression::Any, _) => {
3701 note_compression_fallback(&format!(
3702 "({:?}, {:?}) {}x{}: eligible={eligible}, scheme={scheme:?}, memory={:?}",
3703 desc.format,
3704 desc.dtype,
3705 desc.width,
3706 desc.height,
3707 t.memory()
3708 ));
3709 }
3710 }
3711 }
3712 #[cfg(not(target_os = "android"))]
3713 {
3714 debug_assert!(matches!(request, Compression::Any));
3715 let _ = request;
3716 note_compression_fallback(&format!(
3717 "({:?}, {:?}) {}x{}: no vendor tile compression on this platform",
3718 desc.format, desc.dtype, desc.width, desc.height
3719 ));
3720 }
3721 }
3722 Ok(t)
3723 }
3724
3725 /// Allocate an image tensor of `width` × `height` in `format`.
3726 ///
3727 /// This is the common-case image constructor. The shape follows from the
3728 /// format's layout — packed `[H, W, C]`, planar `[C, H, W]`, semi-planar
3729 /// `[H + chroma_rows, W]` — and the format is recorded on the tensor, so
3730 /// the image crate and the codec can work from it without out-of-band
3731 /// parameters.
3732 ///
3733 /// `memory` requests a backend, or `None` to auto-select. Auto-select for
3734 /// images is **DMA → Mem**: unlike [`Tensor::new`] it does not try SHM,
3735 /// which offers an image nothing that Mem doesn't (it is not
3736 /// GPU-importable, and Mem always succeeds). Ask for
3737 /// [`TensorMemory::Shm`] explicitly if you need it.
3738 ///
3739 /// `access` declares what the **CPU** intends to do with the buffer and is
3740 /// required; see [`CpuAccess`]. Hardware access — GPU, NPU, ISP, codec —
3741 /// is always implied and never declared. Prefer the narrowest declaration
3742 /// that fits: `None` keeps the allocation eligible for Android vendor tile
3743 /// compression, and anything else pins the layout linear.
3744 ///
3745 /// DMA-backed images always get a **64-byte-aligned row stride**, in every
3746 /// layout, because Mali and Vivante reject an EGLImage import at an
3747 /// unaligned pitch. Read the real pitch back with
3748 /// [`Tensor::effective_row_stride`] rather than assuming `width × bpp`.
3749 ///
3750 /// For a compression request or any other less common knob, build an
3751 /// [`ImageDesc`] and call [`Tensor::image_desc`].
3752 ///
3753 /// # Errors
3754 ///
3755 /// - [`Error::InvalidArgument`] if `width` × `height` is not a valid size
3756 /// for `format` (for example odd dimensions where the format's chroma
3757 /// subsampling forbids them), or, on macOS with an explicit
3758 /// [`TensorMemory::DmaBuf`], if the format cannot be expressed as an
3759 /// image-formatted IOSurface at this width. The message names the
3760 /// aligned width to use.
3761 /// - Whatever the chosen backend returns if the allocation itself fails —
3762 /// for an explicit `memory` request there is no fallback.
3763 ///
3764 /// # Example
3765 ///
3766 /// ```rust
3767 /// use edgefirst_tensor::{CpuAccess, PixelFormat, Tensor, TensorMemory};
3768 ///
3769 /// # fn main() -> Result<(), edgefirst_tensor::Error> {
3770 /// // A camera frame the GPU converts and the CPU never touches.
3771 /// let frame = Tensor::<u8>::image(1920, 1080, PixelFormat::Nv12,
3772 /// Some(TensorMemory::Mem), CpuAccess::None)?;
3773 /// assert_eq!(frame.format(), Some(PixelFormat::Nv12));
3774 /// # Ok(())
3775 /// # }
3776 /// ```
3777 pub fn image(
3778 width: usize,
3779 height: usize,
3780 format: PixelFormat,
3781 memory: Option<TensorMemory>,
3782 access: CpuAccess,
3783 ) -> Result<Self>
3784 where
3785 T: 'static,
3786 {
3787 Self::image_impl(width, height, format, memory, access, false)
3788 }
3789
3790 /// [`Tensor::image`], plus the [`ImageDesc::with_contiguous`] request.
3791 /// `contiguous` is only honoured with `Some(TensorMemory::DmaBuf)` on
3792 /// Linux; [`Tensor::image_desc`] checks that before calling.
3793 fn image_impl(
3794 width: usize,
3795 height: usize,
3796 format: PixelFormat,
3797 memory: Option<TensorMemory>,
3798 access: CpuAccess,
3799 contiguous: bool,
3800 ) -> Result<Self>
3801 where
3802 T: 'static,
3803 {
3804 debug_assert!(!contiguous || memory == Some(TensorMemory::DmaBuf));
3805 #[cfg(not(target_os = "linux"))]
3806 let _ = contiguous;
3807 // Shape comes from the shared `PixelFormat::allocation_shape` helper (packed /
3808 // planar / semi-planar NV12·NV16). NV12 supports odd dimensions via the
3809 // `H + ceil(H/2)` combined-plane height.
3810 // The `T: 'static` bound is required by the macOS IOSurface path below.
3811 let shape = format.allocation_shape(width, height).ok_or_else(|| {
3812 Error::InvalidArgument(format!(
3813 "invalid dimensions {width}x{height} for format {format:?}"
3814 ))
3815 })?;
3816
3817 // macOS Dma path: allocate a format-aware IOSurface (FourCC +
3818 // 2D dimensions) so the GL backend can bind it via
3819 // `EGL_ANGLE_iosurface_client_buffer`. Without this, the IOSurface
3820 // would default to a generic byte buffer (FourCC 'L008') and
3821 // ANGLE would reject the import with `EGL_BAD_ATTRIBUTE`.
3822 //
3823 // Guard: IOSurface rounds `bytes_per_row` up to 64-byte alignment.
3824 // If the natural row pitch (`width * channels * sizeof(T)`) is not
3825 // already 64-byte aligned, the padded allocation cannot be mapped
3826 // as a contiguous packed tensor — CPU reads/writes would use the
3827 // wrong stride.
3828 //
3829 // Explicit-Dma contract: when the caller passes
3830 // `Some(TensorMemory::DmaBuf)` they have asked for an
3831 // **image-formatted IOSurface**. Silently downgrading to the
3832 // generic 'L008' byte-bag when alignment fails buries the
3833 // mismatch — the caller only finds out hours later when ANGLE
3834 // (or any GL importer) rejects the bind with
3835 // `EGL_BAD_ATTRIBUTE`. Same anti-pattern bit us previously on
3836 // Mali GPUs with DMA-BUF padding. The right behaviour is to
3837 // fail loudly here with the alignment requirement spelled out
3838 // so the caller can either pick aligned dimensions, request
3839 // SHM/Mem explicitly, or pass `memory=None` for auto-select.
3840 #[cfg(any(target_os = "macos", target_os = "ios"))]
3841 if matches!(memory, Some(TensorMemory::DmaBuf)) {
3842 // For planar formats the IOSurface stacks channels
3843 // vertically (channels * height rows), so the row stride is
3844 // single-channel width * sizeof(T). Packed formats keep the
3845 // natural width * channels * sizeof(T) stride.
3846 let natural_row_bytes = match format.layout() {
3847 PixelLayout::Planar => width * std::mem::size_of::<T>(),
3848 _ => width * format.channels() * std::mem::size_of::<T>(),
3849 };
3850 // A format with a real IOSurface FourCC (RGBA/BGRA/YUYV packed,
3851 // GREY/NV12/NV16/NV24 as R8) tolerates a non-64-aligned natural
3852 // pitch: the surface is allocated with its own 64-aligned
3853 // `bytes_per_row`, the tensor records that stride below, and a CPU
3854 // map iterates rows correctly via the strided-map path while the GL
3855 // import uses the surface's pitch directly — fully zero-copy.
3856 // Planar (the F16 RGBA16F packing) is consumed flat as
3857 // `[1, C, H, W]` with no stride, so it still requires an aligned
3858 // pitch; and formats without a FourCC would fall through to a
3859 // generic byte-bag GL can't bind. Both fail loudly rather than
3860 // silently downgrade.
3861 let has_image_fourcc = dtype_of::<T>()
3862 .and_then(|dt| crate::iosurface_layout::image_iosurface_layout(format, dt))
3863 .is_some();
3864 let padded_ok = has_image_fourcc && format.layout() != PixelLayout::Planar;
3865 if !natural_row_bytes.is_multiple_of(64) && !padded_ok {
3866 let elem_size = std::mem::size_of::<T>();
3867 let per_pixel_bytes = match format.layout() {
3868 PixelLayout::Planar => elem_size.max(1),
3869 _ => format.channels().max(1) * elem_size.max(1),
3870 };
3871 // Compute the next 64-byte-aligned width by rounding the
3872 // natural row pitch up to the next multiple of 64 and
3873 // dividing back by per-pixel bytes. This handles every
3874 // `per_pixel_bytes` value correctly:
3875 //
3876 // * Divisors of 64 (1/2/4/8/16/32/64) → the suggestion
3877 // is always 64-byte aligned.
3878 // * Non-divisors of 64 (e.g. RGB u8 with 3 B/pixel) →
3879 // the next aligned row pitch may not be an integer
3880 // multiple of per_pixel_bytes (3 doesn't divide 64
3881 // in any way), so a "pad width to N" suggestion is
3882 // structurally impossible — omit the suggestion
3883 // instead of printing a wrong number.
3884 // * per_pixel_bytes > 64 → same situation, also
3885 // omitted; the previous formula divided by zero.
3886 //
3887 // The error always names the alignment requirement
3888 // verbatim and lists the two non-DMA alternatives so
3889 // the caller has at least one always-applicable fix.
3890 let aligned_row_bytes = natural_row_bytes.next_multiple_of(64);
3891 let pad_hint =
3892 if per_pixel_bytes > 0 && aligned_row_bytes.is_multiple_of(per_pixel_bytes) {
3893 let w = aligned_row_bytes / per_pixel_bytes;
3894 format!("Pad width to {w} (the next 64-byte-aligned stride), ")
3895 } else {
3896 String::new()
3897 };
3898 return Err(Error::InvalidArgument(format!(
3899 "Tensor::image: {format:?} {width}x{height} with element \
3900 size {elem_size} produces a {natural_row_bytes}-byte natural \
3901 row pitch, which is not 64-byte aligned. \
3902 IOSurface rounds bytes_per_row up to 64 bytes, so a \
3903 contiguous CPU map of this tensor would read garbage. \
3904 {pad_hint}pass memory=None to auto-fall-back to SHM, or \
3905 pass memory=Some(TensorMemory::Shm) or \
3906 Some(TensorMemory::Mem) explicitly."
3907 )));
3908 }
3909 // Alignment OK. Explicit-Dma contract: the caller asked for an
3910 // **image-formatted, GL-importable** IOSurface, so every failure
3911 // from here on is loud. The old behaviour fell through to the
3912 // generic 'L008' byte-bag `Some(other)` arm below, and the caller
3913 // only found out when the GL import rejected the bind with
3914 // `EGL_BAD_ATTRIBUTE` — the same silent-downgrade anti-pattern as
3915 // the alignment case above. (Semi-planar/Grey u8 map to 'L008'
3916 // *by design* in `image_iosurface_layout` — the R8-plane
3917 // representation the YUV shaders sample — so they return through
3918 // the mapped path and never reach these errors.)
3919 let dtype = dtype_of::<T>().ok_or_else(|| {
3920 Error::InvalidArgument(format!(
3921 "Tensor::image: element type {} has no DType, so no \
3922 image-formatted IOSurface exists. Pass memory=None or \
3923 Some(TensorMemory::Mem) for a CPU tensor.",
3924 std::any::type_name::<T>()
3925 ))
3926 })?;
3927 if crate::iosurface_layout::image_iosurface_layout(format, dtype).is_none() {
3928 return Err(Error::InvalidArgument(format!(
3929 "Tensor::image: no zero-copy IOSurface mapping exists for \
3930 {format:?}/{dtype:?} on macOS/iOS (supported: \
3931 Rgba/Rgb @ U8/I8, Bgra/Yuyv/Grey/Nv12/Nv16/Nv24 @ U8, \
3932 Rgba/PlanarRgb/PlanarRgba @ F16). Pass memory=None to \
3933 auto-select, or Some(TensorMemory::Mem) explicitly, for \
3934 a CPU tensor."
3935 )));
3936 }
3937 // Packed RGB u8/i8 rides an RGBA8888 surface at (W*3/4, H) —
3938 // reject a width the texel packing cannot express up front
3939 // (mirrors the Android pre-guard).
3940 if format == PixelFormat::Rgb
3941 && matches!(dtype, DType::U8 | DType::I8)
3942 && packed_rgb888_layout(width, height).is_none()
3943 {
3944 return Err(Error::InvalidArgument(format!(
3945 "Tensor::image: Rgb {dtype:?} requires width%4==0 for the RGBA8888 \
3946 IOSurface packing (got width={width}). Pad the width, or pass \
3947 memory=Some(TensorMemory::Mem) for a CPU tensor."
3948 )));
3949 }
3950 let storage = TensorStorage::<T>::new_image_iosurface(
3951 width, height, format, dtype, &shape, None,
3952 )?;
3953 let mut t = Self::wrap(storage);
3954 t.format = Some(format);
3955 // IOSurface rounds `bytes_per_row` up to 64 bytes. When that
3956 // pitch exceeds the natural packed/planar row stride, record
3957 // it so CPU consumers iterate rows correctly (the GL import
3958 // already uses the surface's own pitch). For 64-aligned rows
3959 // — the common model-input case — the two match and no stride
3960 // is stored, leaving the flat mapping unchanged.
3961 if let TensorStorage::Dma(ref io) = t.storage {
3962 let bpr = io.bytes_per_row();
3963 if let Some(natural) = t.effective_row_stride() {
3964 if bpr > natural {
3965 t.set_row_stride_unchecked(bpr);
3966 }
3967 }
3968 }
3969 t.cpu_access = access;
3970 return Ok(t);
3971 }
3972
3973 // Android Dma path: allocate a format-aware AHardwareBuffer so the
3974 // GL backend can import it as an EGLImage
3975 // (`eglGetNativeClientBufferANDROID` → `eglCreateImageKHR`).
3976 //
3977 // Geometry: gralloc chooses the row pitch (`desc.stride`) at
3978 // allocation, and pads freely (validated on the Galaxy S26 Ultra,
3979 // where SnapAlloc pads the planar-F16 RGBA16F surface). A padded
3980 // pitch is recorded on the tensor so CPU maps iterate rows via the
3981 // strided-map path; the GL render uses the buffer's own pitch
3982 // through the EGLImage either way, so the GPU path stays fully
3983 // zero-copy. Consumers needing a FLAT layout (the future NPU
3984 // handoff's `[1, C, H, W]` contract) must check `row_stride()` and
3985 // repack when set — flatness is a per-device property here, unlike
3986 // macOS where IOSurface's 64-BYTE alignment keeps model-sized F16
3987 // surfaces naturally flat.
3988 #[cfg(target_os = "android")]
3989 if matches!(memory, Some(TensorMemory::DmaBuf)) {
3990 // Explicit-Dma contract (mirrors the macOS block above): the
3991 // caller asked for an image-formatted, GL-importable
3992 // AHardwareBuffer, so an unmapped combination or a gralloc
3993 // refusal errors here instead of falling through to a BLOB
3994 // byte-bag the GL backend can never import — that downgrade
3995 // only surfaced as a silent per-frame CPU upload.
3996 let dtype = dtype_of::<T>().ok_or_else(|| {
3997 Error::InvalidArgument(format!(
3998 "Tensor::image: element type {} has no DType, so no \
3999 image-formatted AHardwareBuffer exists. Pass memory=None \
4000 or Some(TensorMemory::Mem) for a CPU tensor.",
4001 std::any::type_name::<T>()
4002 ))
4003 })?;
4004 // Planar F16 requires width % 4 == 0 for the RGBA16F
4005 // packing — reject up front with the requirement spelled
4006 // out (mirrors the macOS pre-allocation guard) instead of
4007 // falling through to a byte-bag GL cannot bind.
4008 if format.layout() == PixelLayout::Planar
4009 && dtype == DType::F16
4010 && packed_rgba16f_layout(format, dtype, width, height).is_none()
4011 {
4012 return Err(Error::InvalidArgument(format!(
4013 "Tensor::image: {format:?} F16 requires width%4==0 for the RGBA16F \
4014 AHardwareBuffer packing (got width={width}). Pad the width, or pass \
4015 memory=Some(TensorMemory::Mem) for a CPU tensor."
4016 )));
4017 }
4018 // Packed RGB u8/i8 rides an RGBA8888 surface at (W*3/4, H) —
4019 // the same whole-texel constraint as the F16 packing above.
4020 if format == PixelFormat::Rgb
4021 && matches!(dtype, DType::U8 | DType::I8)
4022 && packed_rgb888_layout(width, height).is_none()
4023 {
4024 return Err(Error::InvalidArgument(format!(
4025 "Tensor::image: Rgb {dtype:?} requires width%4==0 for the RGBA8888 \
4026 AHardwareBuffer packing (got width={width}). Pad the width, or pass \
4027 memory=Some(TensorMemory::Mem) for a CPU tensor."
4028 )));
4029 }
4030 if crate::ahardwarebuffer::image_ahardwarebuffer_layout(format, dtype).is_none() {
4031 return Err(Error::InvalidArgument(format!(
4032 "Tensor::image: no zero-copy AHardwareBuffer mapping exists \
4033 for {format:?}/{dtype:?} on Android (Grey/NV* need \
4034 R8-format buffers, API 29+; the HAL floor is API 26 — see \
4035 `image_ahardwarebuffer_layout`). Pass memory=None to \
4036 auto-select, or Some(TensorMemory::Mem) explicitly, for a \
4037 CPU tensor; camera NV12 stays zero-copy by wrapping the \
4038 camera's own AHardwareBuffer instead of allocating one."
4039 )));
4040 }
4041 let storage = TensorStorage::<T>::new_image_ahardwarebuffer(
4042 width, height, format, dtype, &shape, None, access,
4043 )?;
4044 let mut t = Self::wrap(storage);
4045 t.format = Some(format);
4046 if let TensorStorage::Dma(ref ahb) = t.storage {
4047 // gralloc chooses the row pitch (Qualcomm's
4048 // SnapAlloc pads e.g. the 160-px-wide RGBA16F
4049 // surface of a 640-wide planar F16 target).
4050 // When it exceeds the natural stride, record
4051 // it so CPU consumers iterate rows correctly
4052 // via `effective_row_stride()` — the GPU
4053 // renders through the EGLImage at the
4054 // buffer's real pitch regardless, so the
4055 // render stays fully zero-copy. Consumers
4056 // that need the FLAT `[1, C, H, W]` layout
4057 // (the future NPU handoff) must check
4058 // `row_stride()` is unset and repack — or
4059 // pick an aligned width — rather than assume
4060 // flatness (see the module docs).
4061 let bpr = ahb.bytes_per_row();
4062 if let Some(natural) = t.effective_row_stride() {
4063 if bpr > natural {
4064 log::debug!(
4065 "Tensor::image: gralloc padded the {format:?} \
4066 AHardwareBuffer pitch to {bpr} bytes (natural \
4067 {natural}); recording row stride"
4068 );
4069 t.set_row_stride_unchecked(bpr);
4070 }
4071 }
4072 }
4073 t.cpu_access = access;
4074 return Ok(t);
4075 }
4076
4077 // Windows Dma path: allocate an `ID3D11Texture2D` on the process
4078 // device so the GL backend can bind it through ANGLE's D3D11 texture
4079 // import and a D3D12/CUDA consumer can open its shared NT handle.
4080 //
4081 // Explicit-Dma contract, as on macOS and Android: the caller asked for
4082 // a texture, so a format with no texture layout is an error here
4083 // rather than a silent fall-through to the host-memory arms below,
4084 // which would only surface much later as a failed GPU import.
4085 #[cfg(target_os = "windows")]
4086 if matches!(memory, Some(TensorMemory::DmaBuf)) {
4087 let dtype = dtype_of::<T>().ok_or_else(|| {
4088 Error::InvalidArgument(format!(
4089 "Tensor::image: element type {} has no DType, so no D3D11 texture \
4090 layout exists. Pass memory=None or Some(TensorMemory::Mem) for a \
4091 CPU tensor.",
4092 std::any::type_name::<T>()
4093 ))
4094 })?;
4095 let storage = TensorStorage::<T>::new_image_d3d11(
4096 width, height, format, dtype, &shape, None, access,
4097 )?;
4098 let mut t = Self::wrap(storage);
4099 t.format = Some(format);
4100 record_d3d11_row_stride(&mut t, format, "Tensor::image");
4101 t.cpu_access = access;
4102 // Best-effort CUDA import, as the Linux DMA arms do: a texture a
4103 // CUDA consumer can map is worth the one import at allocation.
4104 t.try_init_d3d11_cuda();
4105 return Ok(t);
4106 }
4107
4108 // Compute the **64-byte-aligned** row stride for every image layout.
4109 //
4110 // Embedded GPUs reject `eglCreateImage` DMA-BUF imports whose row pitch
4111 // is not 64-byte aligned: Mali returns `EGL_BAD_ALLOC`, Vivante
4112 // `EGL_BAD_ACCESS`. This bit packed RGBA/RGB destinations at odd widths
4113 // AND at even non-multiple-of-16 widths (e.g. 321→1284, 322→1288 bytes —
4114 // neither divisible by 64), so an odd-source → RGBA convert failed on
4115 // imx95/imx8mp while succeeding on V3D/Tegra. Semi-planar already aligned
4116 // here; we now align packed and planar identically so every image()
4117 // allocation is GPU-importable regardless of width.
4118 //
4119 // The per-layout natural pitch and total row count:
4120 // * SemiPlanar `[total_h, width]` — pitch = even(width)·elem, rows = total_h
4121 // * Packed `[height, width, ch]` — pitch = width·ch·elem, rows = height
4122 // * Planar `[ch, height, width]` — pitch = width·elem, rows = ch·height
4123 // Allocation byte size = `aligned_stride · total_rows` (NOT the shape
4124 // product, which reflects only the logical width and under-allocates the
4125 // padding on odd / unaligned widths).
4126 let elem = std::mem::size_of::<T>();
4127 let channels = format.channels();
4128 let (natural_stride, total_rows) = match format.layout() {
4129 PixelLayout::SemiPlanar => (width.next_multiple_of(2) * elem, shape[0]),
4130 PixelLayout::Packed => (width * channels * elem, height),
4131 PixelLayout::Planar => (width * elem, channels * height),
4132 };
4133 let aligned_stride = natural_stride.next_multiple_of(64);
4134 let semi = format.layout() == PixelLayout::SemiPlanar;
4135
4136 // DMA buffers MUST carry a 64-aligned row pitch — Mali/Vivante reject a
4137 // DMA-BUF EGLImage whose pitch is not 64-aligned. Semi-planar also needs
4138 // the aligned pitch on every backend (its chroma-plane offset math
4139 // assumes it). Packed/planar on host-only memory (Mem/Shm) keep the
4140 // natural tight pitch so the many flat CPU consumers are unaffected.
4141 let host_stride = if semi { aligned_stride } else { natural_stride };
4142 let host_byte_size = host_stride * total_rows;
4143 #[cfg(target_os = "linux")]
4144 let dma_byte_size = aligned_stride * total_rows;
4145
4146 // `used_stride` is the actual row pitch of the storage created below.
4147 let (storage, used_stride) = match memory {
4148 #[cfg(target_os = "linux")]
4149 Some(TensorMemory::DmaBuf) => (
4150 TensorStorage::<T>::new_dma_with_byte_size(
4151 &shape,
4152 dma_byte_size,
4153 None,
4154 crate::dma::HeapPolicy::for_access(access, contiguous),
4155 )?,
4156 aligned_stride,
4157 ),
4158 #[cfg(unix)]
4159 Some(TensorMemory::Shm) => (
4160 TensorStorage::<T>::new_shm_with_byte_size(&shape, host_byte_size, None)?,
4161 host_stride,
4162 ),
4163 Some(TensorMemory::Mem) => (
4164 TensorStorage::<T>::new_mem_with_byte_size(&shape, host_byte_size, None)?,
4165 host_stride,
4166 ),
4167 #[allow(unused_variables)]
4168 Some(other) => {
4169 // PBO and any future variants: fall through to standard new().
4170 return {
4171 let mut t = Self::new(&shape, Some(other), None)?;
4172 t.format = Some(format);
4173 t.cpu_access = access;
4174 Ok(t)
4175 };
4176 }
4177 None => {
4178 // Auto-select priority: DMA → Mem (DMA gets the 64-aligned
4179 // pitch; the Mem fallback keeps the tight host pitch, so the
4180 // recorded stride matches the storage used). Shm is NOT
4181 // auto-selected — it offers no advantage over Mem for an
4182 // in-process image and Mem always succeeds, so it sits below Mem
4183 // and is reached only via an explicit `TensorMemory::Shm`.
4184 #[cfg(target_os = "linux")]
4185 {
4186 match TensorStorage::<T>::new_dma_with_byte_size(
4187 &shape,
4188 dma_byte_size,
4189 None,
4190 crate::dma::HeapPolicy::for_access(access, false),
4191 ) {
4192 Ok(s) => (s, aligned_stride),
4193 Err(_) => (
4194 TensorStorage::<T>::new_mem_with_byte_size(
4195 &shape,
4196 host_byte_size,
4197 None,
4198 )?,
4199 host_stride,
4200 ),
4201 }
4202 }
4203 #[cfg(not(target_os = "linux"))]
4204 {
4205 (
4206 TensorStorage::<T>::new_mem_with_byte_size(&shape, host_byte_size, None)?,
4207 host_stride,
4208 )
4209 }
4210 }
4211 };
4212
4213 let mut t = Self::wrap(storage);
4214 t.format = Some(format);
4215 // Record the row stride when it exceeds the natural tight pitch (padding
4216 // is present — DMA packed/planar at an unaligned width, or always for
4217 // semi-planar), mirroring the IOSurface path above. Aligned-width and
4218 // host-only packed/planar images keep their flat layout with no explicit
4219 // stride; `effective_row_stride()` then falls back to the identical
4220 // computed pitch. When padding IS present, consumers must iterate rows by
4221 // `effective_row_stride()` to skip it.
4222 if semi || used_stride > natural_stride {
4223 t.set_row_stride_unchecked(used_stride);
4224 }
4225 debug_assert!(
4226 t.row_stride.is_some() || !semi,
4227 "image() must always set row_stride for semi-planar tensors"
4228 );
4229 t.cpu_access = access;
4230 #[cfg(target_os = "linux")]
4231 t.try_init_dma_cuda();
4232 Ok(t)
4233 }
4234
4235 /// Create a DMA-backed image tensor with an explicit row stride that
4236 /// may exceed the natural `width * channels * sizeof(T)` pitch.
4237 ///
4238 /// Used for image tensors that need GPU pitch alignment padding: the
4239 /// underlying DMA-BUF is sized to `row_stride * height` bytes, but
4240 /// the tensor's logical shape stays at `[height, width, channels]`.
4241 /// `width()` / `height()` / `shape()` continue to report the
4242 /// user-requested values; the padding is visible only via
4243 /// `row_stride()` / `effective_row_stride()` and is automatically
4244 /// propagated to the GL backend's EGLImage import so Mali Valhall
4245 /// accepts the buffer.
4246 ///
4247 /// # Supported formats
4248 ///
4249 /// Currently only **packed** pixel layouts (RGBA8, BGRA8, RGB888,
4250 /// Grey, etc.) are supported — the formats the GL backend uses as
4251 /// render destinations. Semi-planar formats (NV12, NV16) come from
4252 /// external allocators (camera capture, video decoders) and are
4253 /// imported via `TensorDyn::from_fd` + `set_row_stride`, which
4254 /// already supports padded strides.
4255 ///
4256 /// # Supported memory
4257 ///
4258 /// Currently only `TensorMemory::DmaBuf` is supported. PBO and Mem
4259 /// storage don't go through EGLImage import so they don't need
4260 /// pitch alignment; if you pass any other memory type this returns
4261 /// `NotImplemented`. `None` (auto-select) is treated as `Dma`.
4262 ///
4263 /// # Errors
4264 ///
4265 /// - `InvalidArgument` if `row_stride_bytes < width * channels * sizeof(T)`
4266 /// (the requested stride would not fit a single row)
4267 /// - `NotImplemented` for non-packed formats or non-DMA memory
4268 /// - `IoError` if the DMA-heap allocation fails (propagated from
4269 /// `DmaTensor::new_with_byte_size`)
4270 pub fn image_with_stride(
4271 width: usize,
4272 height: usize,
4273 format: PixelFormat,
4274 row_stride_bytes: usize,
4275 memory: Option<TensorMemory>,
4276 access: CpuAccess,
4277 ) -> Result<Self> {
4278 #[cfg(not(target_os = "linux"))]
4279 let _ = access;
4280 // DMA backing (the only thing this constructor produces) is
4281 // Linux-only. On macOS/BSD/Windows the non-Linux block below is
4282 // the only compiled body and returns `NotImplemented` directly;
4283 // on Linux the non-Linux block is cfg-removed and the function
4284 // falls through to the real validation + allocation path. Each
4285 // target compiles exactly one of the two blocks, and the block
4286 // serves as the function's tail expression in both cases — so
4287 // neither needs an explicit `return` (avoids
4288 // `clippy::needless_return` on the macOS CI gate).
4289 #[cfg(not(target_os = "linux"))]
4290 {
4291 let _ = (width, height, format, row_stride_bytes, memory);
4292 Err(Error::NotImplemented(
4293 "image_with_stride requires DMA support (Linux only)".to_owned(),
4294 ))
4295 }
4296
4297 #[cfg(target_os = "linux")]
4298 {
4299 if format.layout() != PixelLayout::Packed {
4300 return Err(Error::NotImplemented(format!(
4301 "Tensor::image_with_stride only supports packed pixel layouts, got {format:?}"
4302 )));
4303 }
4304 let elem = std::mem::size_of::<T>();
4305 let min_stride = width
4306 .checked_mul(format.channels())
4307 .and_then(|p| p.checked_mul(elem))
4308 .ok_or_else(|| {
4309 Error::InvalidArgument(format!(
4310 "image_with_stride: width {width} × channels {} × sizeof::<T>={elem} \
4311 overflows usize",
4312 format.channels()
4313 ))
4314 })?;
4315 if row_stride_bytes < min_stride {
4316 return Err(Error::InvalidArgument(format!(
4317 "image_with_stride: row_stride {row_stride_bytes} < minimum {min_stride} \
4318 ({width} px × {} ch × {elem} B)",
4319 format.channels()
4320 )));
4321 }
4322 let total_byte_size = row_stride_bytes.checked_mul(height).ok_or_else(|| {
4323 Error::InvalidArgument(format!(
4324 "image_with_stride: row_stride {row_stride_bytes} × height {height} overflows usize"
4325 ))
4326 })?;
4327
4328 let shape = vec![height, width, format.channels()];
4329
4330 let storage = match memory {
4331 Some(TensorMemory::DmaBuf) | None => TensorStorage::<T>::new_dma_with_byte_size(
4332 &shape,
4333 total_byte_size,
4334 None,
4335 crate::dma::HeapPolicy::for_access(access, false),
4336 )?,
4337 Some(other) => {
4338 return Err(Error::NotImplemented(format!(
4339 "image_with_stride: only TensorMemory::DmaBuf is supported, got {other:?}"
4340 )));
4341 }
4342 };
4343
4344 let mut t = Self::wrap(storage);
4345 t.format = Some(format);
4346 t.row_stride = Some(row_stride_bytes);
4347 t.cpu_access = access;
4348 // Match new()/from_fd(): a DMA tensor must attempt CUDA external-
4349 // memory import so a strided DMA buffer is also zero-copy
4350 // CUDA-mappable (no-op when libcudart is absent).
4351 t.try_init_dma_cuda();
4352 Ok(t)
4353 }
4354 }
4355
4356 /// Attach format metadata to an existing tensor.
4357 ///
4358 /// # Arguments
4359 ///
4360 /// * `format` - The pixel format to attach
4361 ///
4362 /// # Returns
4363 ///
4364 /// `Ok(())` on success, with the format stored as metadata on the tensor.
4365 ///
4366 /// # Errors
4367 ///
4368 /// Returns `Error::InvalidShape` if the tensor shape is incompatible with
4369 /// the format's layout (packed expects `[H, W, C]`, planar expects
4370 /// `[C, H, W]`, semi-planar expects `[H*k, W]` with format-specific
4371 /// height constraints).
4372 pub fn set_format(&mut self, format: PixelFormat) -> Result<()> {
4373 let shape = self.shape();
4374 match format.layout() {
4375 PixelLayout::Packed => {
4376 if shape.len() != 3 || shape[2] != format.channels() {
4377 return Err(Error::InvalidShape(format!(
4378 "packed format {format:?} expects [H, W, {}], got {shape:?}",
4379 format.channels()
4380 )));
4381 }
4382 }
4383 PixelLayout::Planar => {
4384 if shape.len() != 3 || shape[0] != format.channels() {
4385 return Err(Error::InvalidShape(format!(
4386 "planar format {format:?} expects [{}, H, W], got {shape:?}",
4387 format.channels()
4388 )));
4389 }
4390 }
4391 PixelLayout::SemiPlanar => {
4392 if shape.len() != 2 {
4393 return Err(Error::InvalidShape(format!(
4394 "semi-planar format {format:?} expects [H*k, W], got {shape:?}"
4395 )));
4396 }
4397 match format {
4398 // Combined-plane height is `H + ceil(H/2)` (luma + chroma
4399 // rows). For even H that is `3H/2` (≡ 0 mod 3); for odd H it
4400 // is `(3H+1)/2` (≡ 2 mod 3). Only totals ≡ 1 mod 3 are
4401 // unreachable, so reject just those — odd-height NV12 is
4402 // valid (e.g. 725 rows for a 483-tall image).
4403 PixelFormat::Nv12 if shape[0] % 3 == 1 => {
4404 return Err(Error::InvalidShape(format!(
4405 "NV12 contiguous shape[0] must be H + ceil(H/2) for some height; \
4406 {} is unreachable (≡ 1 mod 3)",
4407 shape[0]
4408 )));
4409 }
4410 PixelFormat::Nv16 if !shape[0].is_multiple_of(2) => {
4411 return Err(Error::InvalidShape(format!(
4412 "NV16 contiguous shape[0] must be even, got {}",
4413 shape[0]
4414 )));
4415 }
4416 // NV24 (4:4:4): combined-plane height is 3H (Y + 2H chroma).
4417 PixelFormat::Nv24 if !shape[0].is_multiple_of(3) => {
4418 return Err(Error::InvalidShape(format!(
4419 "NV24 contiguous shape[0] must be a multiple of 3 (= 3H), got {}",
4420 shape[0]
4421 )));
4422 }
4423 _ => {}
4424 }
4425 }
4426 }
4427 // Clear stored stride/offset when format changes — they may be invalid
4428 // for the new format. Caller must re-set after changing format.
4429 if self.format != Some(format) {
4430 self.row_stride = None;
4431 self.plane_offset = None;
4432 match self.storage {
4433 TensorStorage::Mem(ref mut m) => m.set_offset(0),
4434 #[cfg(target_os = "linux")]
4435 TensorStorage::Dma(ref mut dma) => dma.mmap_offset = 0,
4436 // Paired with the arm in `set_plane_offset`: a setter that
4437 // takes effect on this backing needs a clear that does too,
4438 // or the wrapper reports `None` while `map()` still starts at
4439 // the old offset -- a stale window rather than a lost one.
4440 // `reshape`'s clear site needs no such arm:
4441 // `IoSurfaceTensor::reshape` zeroes `view_offset` itself
4442 // before the match there runs.
4443 #[cfg(any(target_os = "macos", target_os = "ios"))]
4444 TensorStorage::Dma(ref mut io) => io.view_offset = 0,
4445 // The same pair for D3D11. Unlike `IoSurfaceTensor`,
4446 // `D3d11TextureTensor::reshape` leaves `view_offset` alone, so
4447 // `reshape`'s clear site carries this arm as well.
4448 #[cfg(target_os = "windows")]
4449 TensorStorage::Dma(ref mut tex) => tex.set_view_offset(0),
4450 // Paired with the arm in `set_plane_offset`: a setter that
4451 // takes effect on this backing needs a clear that does too.
4452 // `reshape`'s clear site needs no `Pbo` arm --
4453 // `PboTensor::reshape` zeroes `view_offset` itself before the
4454 // match there runs, exactly as `IoSurfaceTensor::reshape`
4455 // does. `set_logical_shape` deliberately does not, matching
4456 // `Mem`.
4457 TensorStorage::Pbo(ref mut p) => p.view_offset = 0,
4458 // Same pair, `Shm` side. `cfg(unix)` mirrors the `Shm`
4459 // variant's own gate on `TensorStorage`, not a per-OS
4460 // multiplex -- see the comment on the `set_plane_offset`
4461 // arm.
4462 #[cfg(unix)]
4463 TensorStorage::Shm(ref mut s) => s.set_offset(0),
4464 // No-op, pairing with `set_plane_offset`'s Android arm:
4465 // nothing writes `AHardwareBufferTensor::view_offset` there
4466 // either yet, so there is nothing here for a format change
4467 // to clear. See that arm's comment for the full accounting
4468 // of the still-incomplete gap.
4469 #[cfg(target_os = "android")]
4470 TensorStorage::Dma(_) => {}
4471 }
4472 }
4473 self.format = Some(format);
4474 Ok(())
4475 }
4476
4477 /// Set this tensor's logical dimensions and pixel format to a decoded
4478 /// image, reusing the existing allocation. The shape is derived from the
4479 /// format layout; fails with `Error::InsufficientCapacity` if the
4480 /// allocation cannot hold `width`×`height` in `format`, or
4481 /// `Error::InvalidArgument` if the dimensions are invalid for the format.
4482 ///
4483 /// For NV12/NV16/NV24 the buffer width is rounded up to even (a chroma-plane
4484 /// interleaving requirement); the true odd width is reported by the decoder
4485 /// in `ImageInfo` and trimmed by a `convert()` crop. See
4486 /// [`PixelFormat::allocation_shape`].
4487 ///
4488 /// When the backing has a fixed physical row pitch (an IOSurface's
4489 /// 64-aligned `bytesPerRow`) that exceeds the new format's natural row
4490 /// stride — i.e. a reused max-sized pool tensor reconfigured to a smaller
4491 /// image — the physical pitch is preserved as the tensor's `row_stride`.
4492 /// This keeps the **physical grid** (allocation stride/surface) fixed while
4493 /// the **logical ROI** (this image's W×H) changes, so the decode writes rows
4494 /// at the surface's real stride and the GPU samples them at the same stride.
4495 /// Exact-sized buffers (pitch == natural) stay tightly packed unchanged.
4496 pub fn configure_image(
4497 &mut self,
4498 width: usize,
4499 height: usize,
4500 format: PixelFormat,
4501 ) -> Result<()> {
4502 let shape = format.allocation_shape(width, height).ok_or_else(|| {
4503 Error::InvalidArgument(format!(
4504 "invalid dimensions {width}x{height} for format {format:?}"
4505 ))
4506 })?;
4507 // Capture the pre-existing row stride before `set_format` clears it.
4508 // For pool tensors that were allocated at a larger width (e.g. 1920-wide
4509 // pool decoding a 789-wide image), this preserves the backing pitch so
4510 // rows are still written at the correct physical stride.
4511 let prior_stride = self.row_stride;
4512
4513 self.storage.set_logical_shape(&shape)?;
4514 self.set_format(format)?;
4515 // `set_format` keeps the stride when the format is unchanged; the new
4516 // geometry's stride is chosen below, never inherited by accident.
4517 self.row_stride = None;
4518
4519 // Restore the correct row pitch for the new geometry. A DMA buffer of
4520 // ANY layout, and a semi-planar buffer on any backing, MUST carry a
4521 // 64-byte-aligned pitch: Mali/Vivante reject a DMA-BUF EGLImage whose
4522 // row pitch is not 64-aligned (`EGL_BAD_ALLOC` / `EGL_BAD_ACCESS`), and
4523 // semi-planar chroma-offset math assumes it. This mirrors the pitch
4524 // `Tensor::image()` allocates, so a recycled pool buffer imports
4525 // identically to a fresh one — without it a `configure_image()`'d packed
4526 // buffer (e.g. Y800 96-wide → tight pitch 96) failed convert() on
4527 // imx95/imx8mp via the texture-upload fallback while the fresh oracle
4528 // (aligned 128) succeeded. Packed/planar on host-only memory (Mem/Shm)
4529 // keep the tight pitch so flat CPU consumers stay unaffected — matching
4530 // `image()`'s `host_stride` rule.
4531 let elem = std::mem::size_of::<T>();
4532 let channels = format.channels();
4533 let (min_stride, total_rows) = match format.layout() {
4534 PixelLayout::SemiPlanar => (width.next_multiple_of(2) * elem, shape[0]),
4535 PixelLayout::Packed => (width * channels * elem, height),
4536 PixelLayout::Planar => (width * elem, channels * height),
4537 };
4538 let needs_align = self.storage.memory() == TensorMemory::DmaBuf
4539 || format.layout() == PixelLayout::SemiPlanar;
4540
4541 // A backing with its own fixed pitch (IOSurface, AHardwareBuffer, a
4542 // staged D3D11 texture) keeps it whenever it exceeds the new row.
4543 #[cfg(any(
4544 target_os = "macos",
4545 target_os = "ios",
4546 target_os = "android",
4547 target_os = "windows"
4548 ))]
4549 if let Some(pitch) = self.storage.backing_row_stride() {
4550 let active_stride = if pitch > min_stride {
4551 self.set_row_stride_unchecked(pitch);
4552 pitch
4553 } else {
4554 min_stride
4555 };
4556 if !needs_align {
4557 return Ok(());
4558 }
4559 return self.check_configured_stride(active_stride, total_rows);
4560 }
4561
4562 if !needs_align {
4563 return Ok(());
4564 }
4565 // Priority:
4566 // 1. Prior stride (pool reuse): if the pre-existing stride is
4567 // 64-aligned, >= this layout's minimum row, and fits the
4568 // allocation, keep it. This is the hot-loop reuse case (large
4569 // pool, small image).
4570 // 2. Compute a fresh 64-aligned stride for the current width.
4571 let aligned = min_stride.next_multiple_of(64);
4572 let capacity = self.storage.capacity_bytes();
4573
4574 let candidate = if let Some(ps) = prior_stride {
4575 if ps >= min_stride && ps % 64 == 0 && ps * total_rows <= capacity {
4576 ps
4577 } else {
4578 aligned
4579 }
4580 } else {
4581 aligned
4582 };
4583
4584 let active_stride = if candidate * total_rows <= capacity {
4585 self.set_row_stride_unchecked(candidate);
4586 candidate
4587 } else {
4588 // The tight row may still fit where the aligned one does not.
4589 min_stride
4590 };
4591 self.check_configured_stride(active_stride, total_rows)
4592 }
4593
4594 /// The last step of [`configure_image`](Self::configure_image) for an
4595 /// aligned layout: rows that do not fit the allocation are refused here,
4596 /// rather than left to fault on a later map or write.
4597 fn check_configured_stride(&self, active_stride: usize, total_rows: usize) -> Result<()> {
4598 let needed = active_stride * total_rows;
4599 let capacity = self.storage.capacity_bytes();
4600 if needed > capacity {
4601 return Err(Error::InsufficientCapacity { needed, capacity });
4602 }
4603 Ok(())
4604 }
4605
4606 /// Allocate an image tensor sized to hold up to `width`×`height` in
4607 /// `format`, reusable for any smaller image via `configure_image`.
4608 pub fn image_with_capacity(
4609 width: usize,
4610 height: usize,
4611 format: PixelFormat,
4612 memory: Option<TensorMemory>,
4613 access: CpuAccess,
4614 ) -> Result<Self>
4615 where
4616 T: 'static,
4617 {
4618 Self::image(width, height, format, memory, access)
4619 }
4620
4621 /// Pixel format (None if not an image).
4622 pub fn format(&self) -> Option<PixelFormat> {
4623 self.format
4624 }
4625
4626 /// The format and the trailing shape dims that describe one image: packed
4627 /// `[H, W, C]`, planar `[C, H, W]` or semi-planar `[H*k, W]`, optionally
4628 /// behind a leading batch `N` (see [`batch`](Self::batch)). `None` when
4629 /// unformatted or when the shape has any other rank --
4630 /// `set_logical_shape` keeps the format on a shape it no longer describes.
4631 fn image_dims(&self) -> Option<(PixelFormat, &[usize])> {
4632 let fmt = self.format?;
4633 let shape = self.shape();
4634 let rank = image_rank(fmt.layout());
4635 let batch_dims = shape.len().checked_sub(rank).filter(|&n| n <= 1)?;
4636 Some((fmt, &shape[batch_dims..]))
4637 }
4638
4639 /// Image width (None if not an image). For a batched `[N, ...]` tensor
4640 /// this is the width of one element.
4641 pub fn width(&self) -> Option<usize> {
4642 let (fmt, dims) = self.image_dims()?;
4643 match fmt.layout() {
4644 PixelLayout::Packed => Some(dims[1]),
4645 PixelLayout::Planar => Some(dims[2]),
4646 PixelLayout::SemiPlanar => Some(dims[1]),
4647 }
4648 }
4649
4650 /// Image height (None if not an image). For a batched `[N, ...]` tensor
4651 /// this is the height of one element.
4652 ///
4653 /// For semi-planar formats the combined-plane shape row count is divided
4654 /// by the format's luma-to-total ratio to recover logical height. This
4655 /// returns the exact logical height (including odd heights) only because
4656 /// the logical dimensions are tracked separately from the physical shape —
4657 /// `configure_image` stores the actual `(width, height)` in the format's
4658 /// `allocation_shape`, which round-trips losslessly via these accessors.
4659 pub fn height(&self) -> Option<usize> {
4660 let (fmt, dims) = self.image_dims()?;
4661 match fmt.layout() {
4662 PixelLayout::Packed => Some(dims[0]),
4663 PixelLayout::Planar => Some(dims[1]),
4664 PixelLayout::SemiPlanar => {
4665 if self.is_multiplane() {
4666 Some(dims[0])
4667 } else {
4668 match fmt {
4669 PixelFormat::Nv12 => Some(dims[0] * 2 / 3),
4670 PixelFormat::Nv16 => Some(dims[0] / 2),
4671 PixelFormat::Nv24 => Some(dims[0] / 3),
4672 _ => None,
4673 }
4674 }
4675 }
4676 }
4677 }
4678
4679 /// Create from separate Y and UV planes (multiplane NV12/NV16).
4680 pub fn from_planes(luma: Tensor<T>, chroma: Tensor<T>, format: PixelFormat) -> Result<Self> {
4681 if format.layout() != PixelLayout::SemiPlanar {
4682 return Err(Error::InvalidArgument(format!(
4683 "from_planes requires a semi-planar format, got {format:?}"
4684 )));
4685 }
4686 if chroma.format.is_some() || chroma.chroma.is_some() {
4687 return Err(Error::InvalidArgument(
4688 "chroma tensor must be a raw tensor (no format or chroma metadata)".into(),
4689 ));
4690 }
4691 let luma_shape = luma.shape();
4692 let chroma_shape = chroma.shape();
4693 if luma_shape.len() != 2 || chroma_shape.len() != 2 {
4694 return Err(Error::InvalidArgument(format!(
4695 "from_planes expects 2D shapes, got luma={luma_shape:?} chroma={chroma_shape:?}"
4696 )));
4697 }
4698 if luma_shape[1] != chroma_shape[1] {
4699 return Err(Error::InvalidArgument(format!(
4700 "luma width {} != chroma width {}",
4701 luma_shape[1], chroma_shape[1]
4702 )));
4703 }
4704 match format {
4705 PixelFormat::Nv12 => {
4706 if luma_shape[0] % 2 != 0 {
4707 return Err(Error::InvalidArgument(format!(
4708 "NV12 requires even luma height, got {}",
4709 luma_shape[0]
4710 )));
4711 }
4712 if chroma_shape[0] != luma_shape[0] / 2 {
4713 return Err(Error::InvalidArgument(format!(
4714 "NV12 chroma height {} != luma height / 2 ({})",
4715 chroma_shape[0],
4716 luma_shape[0] / 2
4717 )));
4718 }
4719 }
4720 PixelFormat::Nv16 => {
4721 if chroma_shape[0] != luma_shape[0] {
4722 return Err(Error::InvalidArgument(format!(
4723 "NV16 chroma height {} != luma height {}",
4724 chroma_shape[0], luma_shape[0]
4725 )));
4726 }
4727 }
4728 // NV24's chroma plane is full-resolution (2×-wide interleaved UV),
4729 // which the equal-width plane check above doesn't model. Multiplane
4730 // NV24 is unused (the JPEG decoder emits a contiguous NV24 buffer),
4731 // so it's not supported here yet.
4732 _ => {
4733 return Err(Error::InvalidArgument(format!(
4734 "from_planes only supports NV12 and NV16 (NV24 multiplane not yet \
4735 supported — use a contiguous NV24 tensor), got {format:?}"
4736 )));
4737 }
4738 }
4739
4740 Ok(Tensor {
4741 storage: luma.storage,
4742 format: Some(format),
4743 chroma: Some(Box::new(chroma)),
4744 row_stride: luma.row_stride,
4745 plane_offset: luma.plane_offset,
4746 quantization: luma.quantization,
4747 // A multiplane tensor spans two DMA-BUFs (luma + chroma); CUDA
4748 // external-memory import is per-fd, so there is no single device
4749 // pointer for the composite. Any CUDA handle the luma plane carried
4750 // is intentionally dropped — consumers needing CUDA access to
4751 // multiplane data must import each plane independently.
4752 cuda: None,
4753 colorimetry: luma.colorimetry,
4754 cpu_access: luma.cpu_access,
4755 compression: luma.compression,
4756 // A composed multiplane tensor is a whole image, not a sub-view.
4757 view_origin: None,
4758 view_marked: false,
4759 })
4760 }
4761
4762 /// Whether this tensor uses separate plane allocations.
4763 pub fn is_multiplane(&self) -> bool {
4764 self.chroma.is_some()
4765 }
4766
4767 /// Access the chroma plane for multiplane semi-planar images.
4768 pub fn chroma(&self) -> Option<&Tensor<T>> {
4769 self.chroma.as_deref()
4770 }
4771
4772 /// Mutable access to the chroma plane for multiplane semi-planar images.
4773 pub fn chroma_mut(&mut self) -> Option<&mut Tensor<T>> {
4774 self.chroma.as_deref_mut()
4775 }
4776
4777 /// Independent tensor over this plane's allocation, copying geometry
4778 /// metadata. Used to answer `ef_tensor_chroma` from the handle without a
4779 /// client-side shadow field.
4780 pub(crate) fn clone_chroma_plane(&self) -> Result<Self> {
4781 // Every storage accepted by `from_planes` implements `view` as a
4782 // shared-backing lens. Using that common path keeps
4783 // `is_multiplane() => chroma().is_some()` true for SHM, PBO and
4784 // D3D11 too, instead of silently losing chroma when a hand-written
4785 // per-storage clone omits a variant.
4786 self.subview(0, self.shape())
4787 }
4788
4789 /// Row stride in bytes (`None` = tightly packed).
4790 pub fn row_stride(&self) -> Option<usize> {
4791 self.row_stride
4792 }
4793
4794 /// Effective row stride in bytes: the stored stride if set, otherwise the
4795 /// minimum stride computed from the format, width, and element size.
4796 /// Returns `None` only when no format is set and no explicit stride was
4797 /// stored via [`set_row_stride`](Self::set_row_stride).
4798 ///
4799 /// **GREY note:** `effective_row_stride()` for a GREY tensor returns the
4800 /// tight `width` bytes (no padding), which is what `normalize_to_numpy` and
4801 /// the CPU convert path expect. The codec's internal `native_row_stride`
4802 /// (64-byte-aligned) is used only during decoding and is not propagated to
4803 /// the tensor's stored stride, so callers reading via
4804 /// `effective_row_stride()` always see the tight value for GREY.
4805 pub fn effective_row_stride(&self) -> Option<usize> {
4806 if let Some(s) = self.row_stride {
4807 return Some(s);
4808 }
4809 let fmt = self.format?;
4810 let w = self.width()?;
4811 let elem = std::mem::size_of::<T>();
4812 Some(match fmt.layout() {
4813 PixelLayout::Packed => w * fmt.channels() * elem,
4814 PixelLayout::Planar => w * elem,
4815 // Semi-planar: minimum stride must cover the even width so the
4816 // interleaved chroma columns are byte-aligned on odd-width images.
4817 PixelLayout::SemiPlanar => w.next_multiple_of(2) * elem,
4818 })
4819 }
4820
4821 /// Copy the tensor's logical bytes into `dst`, compacting away any
4822 /// recorded row-stride padding.
4823 ///
4824 /// The flatness helper for NPU handoff: consumers that need a FLAT
4825 /// layout (e.g. `[1, C, H, W]` for NNAPI/LiteRT when the runtime cannot
4826 /// take a padded pitch) call this when [`row_stride`](Self::row_stride)
4827 /// is `Some` — on a tight tensor it degenerates to one memcpy, so it is
4828 /// safe to call unconditionally. `dst.len()` must equal the tight byte
4829 /// footprint (`shape` product × element size). Zero-copy consumers
4830 /// should prefer the buffer handle + [`Tensor::effective_row_stride`]
4831 /// (Self::effective_row_stride) and skip this copy entirely.
4832 pub fn copy_to_flat(&self, dst: &mut [u8]) -> Result<()> {
4833 let tight_bytes = crate::ahardwarebuffer_layout::checked_shape_bytes::<T>(self.shape())?;
4834 if dst.len() != tight_bytes {
4835 return Err(Error::InvalidArgument(format!(
4836 "copy_to_flat: dst is {} bytes but the tensor's tight \
4837 footprint is {tight_bytes} bytes (shape {:?})",
4838 dst.len(),
4839 self.shape()
4840 )));
4841 }
4842 // Read-only: this helper never writes back into the tensor, so a
4843 // read map suffices — and is required on backends (D3D11) that
4844 // reject a writable map on a tensor allocated CpuAccess::Read.
4845 let map = self.map_read()?;
4846 // SAFETY: T is a plain numeric type (crate-wide bound); viewing the
4847 // mapped elements as bytes is sound.
4848 let src: &[u8] = unsafe {
4849 std::slice::from_raw_parts(
4850 map.as_slice().as_ptr() as *const u8,
4851 std::mem::size_of_val(map.as_slice()),
4852 )
4853 };
4854 let Some(stride) = self.row_stride else {
4855 // Tight layout: the mapped window is exactly the logical bytes.
4856 let Some(tight) = src.get(..tight_bytes) else {
4857 return Err(Error::InvalidOperation(format!(
4858 "copy_to_flat: mapped {} bytes < tight footprint {tight_bytes}",
4859 src.len()
4860 )));
4861 };
4862 dst.copy_from_slice(tight);
4863 return Ok(());
4864 };
4865 // Strided: the rows the strided map spans, each holding one tight
4866 // logical row. `rows` is a factor of the element count, so it divides
4867 // the tight footprint exactly.
4868 let rows = strided_rows(self)?;
4869 if rows == 0 {
4870 return Err(Error::InvalidOperation(
4871 "copy_to_flat: a strided tensor with no rows has nothing to copy".into(),
4872 ));
4873 }
4874 let row_bytes = tight_bytes / rows;
4875 let need = (rows - 1)
4876 .checked_mul(stride)
4877 .and_then(|b| b.checked_add(row_bytes))
4878 .ok_or_else(|| {
4879 Error::InvalidOperation(format!(
4880 "copy_to_flat: stride {stride} × rows {rows} overflows usize"
4881 ))
4882 })?;
4883 if src.len() < need {
4884 return Err(Error::InvalidOperation(format!(
4885 "copy_to_flat: mapped {} bytes but strided rows need {need}",
4886 src.len()
4887 )));
4888 }
4889 for r in 0..rows {
4890 dst[r * row_bytes..(r + 1) * row_bytes]
4891 .copy_from_slice(&src[r * stride..r * stride + row_bytes]);
4892 }
4893 Ok(())
4894 }
4895
4896 /// Set the row stride in bytes for externally allocated buffers with
4897 /// row padding (e.g. V4L2 or GStreamer allocators).
4898 ///
4899 /// The stride is propagated to the EGL DMA-BUF import attributes so
4900 /// the GPU interprets the padded buffer layout correctly. Must be
4901 /// called after [`set_format`](Self::set_format) and before the tensor
4902 /// is first passed to `ImageProcessor::convert`. The stored stride
4903 /// is cleared automatically if the pixel format is later changed.
4904 ///
4905 /// No stride-vs-buffer-size validation is performed because the
4906 /// backing allocation size is not reliably known: external DMA-BUFs
4907 /// may be over-allocated by the allocator, and internal tensors store
4908 /// a logical (unpadded) shape. An incorrect stride will be caught by
4909 /// the EGL driver at import time.
4910 ///
4911 /// # Arguments
4912 ///
4913 /// * `stride` - Row stride in bytes. Must be >= the minimum stride for
4914 /// the format (width * channels * sizeof(T) for packed,
4915 /// width * sizeof(T) for planar/semi-planar).
4916 ///
4917 /// # Errors
4918 ///
4919 /// * `InvalidArgument` if no pixel format is set on this tensor
4920 /// * `InvalidArgument` if `stride` is less than the minimum for the
4921 /// format and width
4922 pub fn set_row_stride(&mut self, stride: usize) -> Result<()> {
4923 let fmt = self.format.ok_or_else(|| {
4924 Error::InvalidArgument("cannot set row_stride without a pixel format".into())
4925 })?;
4926 let w = self.width().ok_or_else(|| {
4927 Error::InvalidArgument("cannot determine width for row_stride validation".into())
4928 })?;
4929 let elem = std::mem::size_of::<T>();
4930 let min_stride = match fmt.layout() {
4931 PixelLayout::Packed => w * fmt.channels() * elem,
4932 PixelLayout::Planar => w * elem,
4933 // Semi-planar: minimum must cover even width for chroma alignment.
4934 PixelLayout::SemiPlanar => w.next_multiple_of(2) * elem,
4935 };
4936 if stride < min_stride {
4937 return Err(Error::InvalidArgument(format!(
4938 "row_stride {stride} < minimum {min_stride} for {fmt:?} at width {w}"
4939 )));
4940 }
4941 self.row_stride = Some(stride);
4942 Ok(())
4943 }
4944
4945 /// Set the row stride without format validation.
4946 ///
4947 /// Use this for raw sub-tensors (e.g. chroma planes) that don't carry
4948 /// format metadata. The caller is responsible for ensuring the stride
4949 /// is valid.
4950 pub fn set_row_stride_unchecked(&mut self, stride: usize) {
4951 self.row_stride = Some(stride);
4952 }
4953
4954 /// Builder-style variant of [`set_row_stride`](Self::set_row_stride),
4955 /// consuming and returning `self`.
4956 ///
4957 /// # Errors
4958 ///
4959 /// Same conditions as [`set_row_stride`](Self::set_row_stride).
4960 pub fn with_row_stride(mut self, stride: usize) -> Result<Self> {
4961 self.set_row_stride(stride)?;
4962 Ok(self)
4963 }
4964
4965 /// Byte offset within the DMA-BUF where image data starts (`None` = 0).
4966 pub fn plane_offset(&self) -> Option<usize> {
4967 self.plane_offset
4968 }
4969
4970 /// The parent-image snapshot if this tensor is a [`view`](Self::view)/
4971 /// [`batch`](Self::batch) sub-region; `None` for a whole tensor. The GL
4972 /// backend keys its import on the parent geometry and renders this view as a
4973 /// `glViewport`/`glScissor` ROI at `(x, y, width, height)`. See [`ViewOrigin`].
4974 pub fn view_origin(&self) -> Option<ViewOrigin> {
4975 self.view_origin
4976 }
4977
4978 /// Whether this tensor addresses part of a larger buffer: a
4979 /// [`view`](Self::view)/[`batch`](Self::batch) with a [`ViewOrigin`], or a
4980 /// tensor rebuilt from such a view's descriptor, which carries none.
4981 ///
4982 /// A view's rows share cache lines with bytes outside it, so a CPU pass
4983 /// that writes it maps `ReadWrite` (see [`TensorTrait::write_access`]).
4984 pub fn is_view(&self) -> bool {
4985 self.view_origin.is_some() || self.view_marked
4986 }
4987
4988 /// Record that this tensor addresses part of a larger buffer, as
4989 /// [`is_view`](Self::is_view) reports. The descriptor import calls this
4990 /// for a producer that was a view; it is never cleared.
4991 pub fn mark_view(&mut self) {
4992 self.view_marked = true;
4993 }
4994
4995 /// Set the byte offset within the DMA-BUF where image data starts.
4996 ///
4997 /// Propagated to `EGL_DMA_BUF_PLANE0_OFFSET_EXT` on GPU import.
4998 /// Unlike [`set_row_stride`](Self::set_row_stride), no format is required
4999 /// since the offset is format-independent.
5000 pub fn set_plane_offset(&mut self, offset: usize) {
5001 self.plane_offset = Some(offset);
5002 // The offset consulted by `map()` lives inside the storage variant.
5003 // Keep it in sync with the wrapper field for every backing that
5004 // honors it; see also the clear sites in `set_format` and `reshape`.
5005 //
5006 // Note that `TensorStorage::Dma` is a cfg-multiplexed *name* rather
5007 // than one type -- `DmaTensor` on Linux, `IoSurfaceTensor` on
5008 // macOS/iOS, `AHardwareBufferTensor` on Android,
5009 // `D3d11TextureTensor` on Windows -- so each platform needs its own
5010 // arm below; a missing one would silently become a fall-through
5011 // rather than a compile error, which is how macOS went without one
5012 // for a while. The match is kept exhaustive (no `_ => {}`)
5013 // specifically so that hazard is a compile error, not a silent gap.
5014 match self.storage {
5015 TensorStorage::Mem(ref mut m) => m.set_offset(offset),
5016 #[cfg(target_os = "linux")]
5017 TensorStorage::Dma(ref mut dma) => dma.mmap_offset = offset,
5018 // The offset `IoSurfaceTensor::scoped_pin`/`map_inner` add to the
5019 // locked base address. Writing it here is what makes a
5020 // *reconstructed* view address its own sub-region instead of the
5021 // parent surface's origin; `IoSurfaceTensor::view` already set it
5022 // on the fresh-view path, which is why only the descriptor round
5023 // trip was ever wrong. Idempotent against `subview`, which sets
5024 // the same absolute value by both routes -- pinned by
5025 // `nested_iosurface_subviews_do_not_compound_their_plane_offset`.
5026 #[cfg(any(target_os = "macos", target_os = "ios"))]
5027 TensorStorage::Dma(ref mut io) => io.view_offset = offset,
5028 // The offset every `D3d11TextureTensor` pin adds to its host or
5029 // staging base, in backing (pitched) space like `row_stride()`.
5030 // `view()` bounds-checks the value it computes; this route does
5031 // not, so the pins refuse an offset past the backing instead
5032 // (`check_view_offset` in `d3d11/texture.rs`). Pinned by the
5033 // `d3d11` plane-offset tests in `crates/tensor/tests`.
5034 #[cfg(target_os = "windows")]
5035 TensorStorage::Dma(ref mut tex) => tex.set_view_offset(offset),
5036 // The offset `PboTensor::scoped_pin` adds to the mapped base.
5037 // Writing it here is what makes a *reconstructed* view address
5038 // its own sub-region instead of the parent buffer's origin;
5039 // `PboTensor::view` already set it on the fresh-view path, which
5040 // is why only the round trip was ever wrong -- the same shape
5041 // IOSurface had. Idempotent against `subview`, which sets the
5042 // same absolute value by both routes -- pinned by
5043 // `nested_pbo_subviews_do_not_compound_their_plane_offset`.
5044 //
5045 // Not `cfg`-gated, and deliberately so: unlike
5046 // `TensorStorage::Dma` -- a cfg-multiplexed *name* whose gated
5047 // arm silently becomes `_ => {}` on the other targets, the
5048 // mechanism behind this whole issue class -- `TensorStorage::Pbo`
5049 // is one type on every platform.
5050 TensorStorage::Pbo(ref mut p) => p.view_offset = offset,
5051 // The offset `ShmTensor::map_inner`/`host_pin` add to the mmap
5052 // base, exactly like `MemTensor::offset`. `cfg(unix)` here
5053 // mirrors `TensorStorage::Shm`'s own gate on the enum
5054 // definition -- it is not a per-OS multiplex the way
5055 // `TensorStorage::Dma` is, so this arm cannot silently become
5056 // `_ => {}` on a target where the `Shm` variant still exists.
5057 #[cfg(unix)]
5058 TensorStorage::Shm(ref mut s) => s.set_offset(offset),
5059 // STILL INCOMPLETE: `AHardwareBufferTensor` (Android) has a
5060 // `view_offset` that its own `map()` adds and its own `view()`
5061 // sets, but nothing writes it here yet, so restoring an offset
5062 // onto an already-built tensor is silently a no-op on this
5063 // backing -- a caller gets a map at the parent's origin. Not
5064 // reachable through `TensorDyn::import_descriptor` today
5065 // (`import_storage`'s DMABUF arm is `cfg(target_os = "linux")`),
5066 // so it is latent rather than silent. See
5067 // `edgefirst-python-common`'s `interop::apply_plane_offset` for
5068 // the full per-backing accounting.
5069 #[cfg(target_os = "android")]
5070 TensorStorage::Dma(_) => {}
5071 }
5072 }
5073
5074 /// The CPU access declared for this tensor (see [`CpuAccess`]): set at
5075 /// allocation, possibly narrowed by
5076 /// [`restrict_cpu_access`](Self::restrict_cpu_access). A view copies its
5077 /// parent's declaration when it is taken.
5078 pub fn cpu_access(&self) -> CpuAccess {
5079 self.cpu_access
5080 }
5081
5082 /// Narrow the declared CPU access to the directions both it and `access`
5083 /// cover, on the chroma plane too. Never widens: a declaration must not
5084 /// exceed what the allocation was made for, so a wider `access` leaves
5085 /// it unchanged. An importer uses this to carry a producer's declaration
5086 /// onto an alias of the same buffer.
5087 ///
5088 /// ```
5089 /// # use edgefirst_tensor::{CpuAccess, PixelFormat, Tensor, TensorMemory};
5090 /// # fn main() -> Result<(), edgefirst_tensor::Error> {
5091 /// let mut t = Tensor::<u8>::image(4, 4, PixelFormat::Rgba, Some(TensorMemory::Mem), CpuAccess::ReadWrite)?;
5092 /// t.restrict_cpu_access(CpuAccess::Read);
5093 /// assert_eq!(t.cpu_access(), CpuAccess::Read);
5094 /// t.restrict_cpu_access(CpuAccess::ReadWrite);
5095 /// assert_eq!(t.cpu_access(), CpuAccess::Read);
5096 /// # Ok(())
5097 /// # }
5098 /// ```
5099 pub fn restrict_cpu_access(&mut self, access: CpuAccess) {
5100 self.cpu_access = self.cpu_access.intersection(access);
5101 if let Some(chroma) = self.chroma.as_mut() {
5102 chroma.restrict_cpu_access(access);
5103 }
5104 }
5105
5106 /// Set the declared CPU access without re-allocating — crate-private:
5107 /// the declaration must reflect the underlying allocation's real
5108 /// capabilities (constructors and importers set it; arbitrary widening
5109 /// would defeat the contract).
5110 // Only the Android AHardwareBuffer importer derives a declaration from usage bits.
5111 #[cfg(target_os = "android")]
5112 pub(crate) fn set_cpu_access_unchecked(&mut self, access: CpuAccess) {
5113 self.cpu_access = access;
5114 }
5115
5116 /// The vendor tile-compression scheme recorded at allocation, or
5117 /// `None` for a linear layout. `Some` means the pixels live in a
5118 /// proprietary tile order: the row-stride accessors describe no
5119 /// meaningful linear layout and CPU maps are best-effort (see
5120 /// [`Compression`]). Only Android hardware-only allocations that
5121 /// requested compression record a scheme.
5122 pub fn compression(&self) -> Option<CompressionScheme> {
5123 self.compression
5124 }
5125
5126 /// Whether this tensor's memory is known to be physically contiguous:
5127 /// [`Contiguity::Contiguous`] for a DMA-BUF this crate allocated from
5128 /// the CMA heap, [`Contiguity::NonContiguous`] for one from the system
5129 /// heap, and [`Contiguity::Unknown`] for an imported fd and every other
5130 /// kind of memory.
5131 pub fn contiguity(&self) -> Contiguity {
5132 match &self.storage {
5133 #[cfg(target_os = "linux")]
5134 TensorStorage::Dma(d) => d.contiguity,
5135 _ => Contiguity::Unknown,
5136 }
5137 }
5138
5139 /// Record the compression scheme — crate-private: recording is an
5140 /// allocation-time fact ([`Tensor::image_desc`] sets it; arbitrary
5141 /// mutation would misdescribe the physical layout).
5142 // Only the Android allocation path records a scheme.
5143 #[cfg(target_os = "android")]
5144 pub(crate) fn set_compression_unchecked(&mut self, scheme: Option<CompressionScheme>) {
5145 self.compression = scheme;
5146 }
5147
5148 pub fn colorimetry(&self) -> Option<crate::Colorimetry> {
5149 self.colorimetry
5150 }
5151
5152 /// Attach/clear colorimetry metadata.
5153 pub fn set_colorimetry(&mut self, c: Option<crate::Colorimetry>) {
5154 self.colorimetry = c;
5155 }
5156
5157 /// Builder-style colorimetry attach.
5158 pub fn with_colorimetry(mut self, c: crate::Colorimetry) -> Self {
5159 self.colorimetry = Some(c);
5160 self
5161 }
5162
5163 /// Create a zero-copy sub-region view of this tensor's backing buffer.
5164 ///
5165 /// The returned tensor shares this tensor's allocation (no copy) and maps
5166 /// the window `[offset_bytes, offset_bytes + shape.product()*size_of::<T>())`
5167 /// measured from this tensor's own logical start. N sub-views into one
5168 /// parent can be written independently, enabling batched assembly into a
5169 /// single buffer. Identical semantics across `Mem` (shared `Arc`) and
5170 /// `Dma` (shared fd) backings.
5171 ///
5172 /// # Disjointness
5173 ///
5174 /// Independent writes are sound *only* when the windows do not overlap. The
5175 /// shared backing uses interior mutability (`UnsafeCell` cells), so two
5176 /// sub-views whose byte ranges intersect alias the same cells: writing one
5177 /// while reading or writing the other is a data race and therefore
5178 /// **undefined behaviour**. The caller is responsible for keeping the
5179 /// windows disjoint; this method does not check for overlap.
5180 ///
5181 /// # Errors
5182 ///
5183 /// - [`Error::InvalidOperation`] if the backing is not `Mem` or `Dma`, or
5184 /// if `offset_bytes` is not a multiple of `align_of::<T>()`.
5185 /// - [`Error::InsufficientCapacity`] / [`Error::InvalidSize`] if the window
5186 /// exceeds the parent allocation.
5187 pub(crate) fn subview(&self, offset_bytes: usize, shape: &[usize]) -> Result<Tensor<T>> {
5188 // Offset is absolute into the backing allocation: a sub-view of a
5189 // sub-view composes by adding this tensor's own offset.
5190 let abs_offset = self
5191 .plane_offset
5192 .unwrap_or(0)
5193 .checked_add(offset_bytes)
5194 .ok_or(Error::InvalidSize(offset_bytes))?;
5195 // Every backend exposes `view(offset, shape)` via `TensorTrait`, sharing
5196 // the resource AND `BufferIdentity` (unlike `from_fd`/`from_surface`,
5197 // which mint a fresh identity). The GL backend keys the import on the
5198 // shared identity so offset-distinct sub-views of one buffer reuse a
5199 // single import and address their window via `glViewport`. `Mem`/`Shm`
5200 // share via the allocation `Arc` / a cloned fd; `Pbo` via the GL-buffer
5201 // `Arc`; Linux DMA-BUF / macOS IOSurface via the shared fd / CFRetain.
5202 // `TensorStorage::view` performs the one remaining per-variant dispatch.
5203 let mut t = Tensor::wrap(self.storage.view(offset_bytes, shape)?);
5204 // Inherit the parent's image metadata so the view is a ready-to-use
5205 // sub-image (e.g. a `convert()` destination). The offset is applied
5206 // LAST because `set_format` deliberately clears it — the offset is a
5207 // structural property of the sub-region, not format-dependent metadata.
5208 if let Some(fmt) = self.format {
5209 t.set_format(fmt)?;
5210 }
5211 if let Some(rs) = self.row_stride {
5212 t.set_row_stride_unchecked(rs);
5213 }
5214 t.quantization = self.quantization.clone();
5215 // A sub-region of an image carries the parent's colorimetry — it is the
5216 // same pixels, same color encoding. Inherit it like the other image
5217 // metadata above so a sub-view is a faithful convert() source/target.
5218 t.set_colorimetry(self.colorimetry);
5219 // The declared CPU access is a property of the underlying
5220 // allocation, so every view shares the parent's declaration.
5221 t.cpu_access = self.cpu_access;
5222 // Likewise the recorded compression scheme: the view shares the
5223 // parent's physical layout.
5224 t.compression = self.compression;
5225 t.view_marked = self.is_view();
5226 if abs_offset > 0 {
5227 t.set_plane_offset(abs_offset);
5228 }
5229 Ok(t)
5230 }
5231
5232 /// Borrow batch element `n` of a batched tensor as a zero-copy view.
5233 ///
5234 /// A batched tensor prepends `N` as the leading dimension over the
5235 /// per-element image layout (`[N, H, W, C]` packed or `[N, C, H, W]`
5236 /// planar) — `N` is over the whole per-element block regardless of
5237 /// `HWC`/`CHW`. `batch(n)` returns element `n`: the contiguous per-element
5238 /// region at byte offset `n * element_size`, sharing the parent's
5239 /// `BufferIdentity` and inheriting its format / row stride / colorimetry.
5240 /// `batch(0)` on a tensor with `N == 1` is equivalent to the whole tensor.
5241 ///
5242 /// # Errors
5243 ///
5244 /// - [`Error::BatchIndexOutOfBounds`] if `n >= N`.
5245 /// - [`Error::InvalidShape`] if the tensor is not batched (a formatted
5246 /// tensor whose rank lacks the leading `N`, or an empty shape).
5247 pub fn batch(&self, n: usize) -> Result<Tensor<T>> {
5248 let shape = self.shape();
5249 // With a format we know the exact per-element rank, so a missing leading
5250 // `N` is a misuse we reject rather than silently treating a spatial dim
5251 // as the batch. Raw tensors take shape[0] as `N` by contract.
5252 if let Some(fmt) = self.format {
5253 let elem_rank = image_rank(fmt.layout());
5254 if shape.len() != elem_rank + 1 {
5255 return Err(Error::InvalidShape(format!(
5256 "batch(): tensor is not batched ({fmt:?} expects a leading N over a \
5257 {elem_rank}-D element, got shape {shape:?})"
5258 )));
5259 }
5260 }
5261 let batch = *shape
5262 .first()
5263 .ok_or_else(|| Error::InvalidShape("batch(): empty shape".into()))?;
5264 if n >= batch {
5265 return Err(Error::BatchIndexOutOfBounds { index: n, batch });
5266 }
5267 let elem_shape: Vec<usize> = shape[1..].to_vec();
5268 let elem_count: usize = elem_shape.iter().product();
5269 elem_count
5270 .checked_mul(std::mem::size_of::<T>())
5271 .ok_or(Error::InvalidSize(elem_count))?;
5272 // One element's footprint including row padding: the descriptor's
5273 // stride for `N`, so a padded batch is addressed the way it is mapped.
5274 let dims: Vec<u64> = shape.iter().map(|&d| d as u64).collect();
5275 let elem_bytes = crate::protocol::c_byte_strides(
5276 &dims,
5277 std::mem::size_of::<T>() as i64,
5278 self.format.map(|f| f.layout()),
5279 self.row_stride,
5280 )[0] as usize;
5281 let offset = n.checked_mul(elem_bytes).ok_or(Error::InvalidSize(n))?;
5282 // For a packed `[N, H, W, C]` tensor the N tiles stack vertically in the
5283 // shared buffer, so the GL import sees one `(W, N*H)` parent and each
5284 // tile is the row-band at `y = n*H`. Snapshot that parent so the backend
5285 // imports once and renders the tile via `glViewport`. Non-packed
5286 // (planar/semi-planar) batching keeps the per-slot path for now (planar
5287 // NCHW tiling is a separate step); raw tensors have no pixel geometry.
5288 let view_origin = match self.format.map(|f| f.layout()) {
5289 Some(PixelLayout::Packed) => {
5290 let tile_h = elem_shape[0];
5291 let tile_w = elem_shape[1];
5292 // Per-row pitch of the tall `(W, N*H)` parent — padded stride if
5293 // set, else the tight row width. The GL import keys on this.
5294 let parent_stride = self
5295 .effective_row_stride()
5296 .expect("a formatted tensor of its batched rank has a width");
5297 Some(self.compose_view_origin(tile_w, batch * tile_h, parent_stride, 0, n * tile_h))
5298 }
5299 _ => None,
5300 };
5301 let mut t = self.subview(offset, &elem_shape)?;
5302 t.view_origin = view_origin;
5303 Ok(t)
5304 }
5305
5306 /// Borrow a rectangular spatial sub-region of an image tensor as a
5307 /// zero-copy view — the **destination/source crop** primitive.
5308 ///
5309 /// `region` is in pixels of the image's leading frame. The returned view
5310 /// shares the parent's `BufferIdentity` and addresses the sub-rectangle by
5311 /// offset + the **parent's** row pitch (so each row lands at the correct
5312 /// columns). `convert(src, &mut dst.view(rect), …)` renders into that
5313 /// sub-rectangle of `dst`; a letterbox fit then clears the view and renders
5314 /// the aspect-preserved content into its inner region. `view`/`batch`/the
5315 /// whole tensor are the one coherent destination model — there is no
5316 /// separate `dst_rect`.
5317 ///
5318 /// # Known limitation: bottom-right multi-row views
5319 ///
5320 /// A multi-row view is mapped as `parent_pitch × rows` bytes **from its own
5321 /// offset**, so the window it asks for overruns the allocation by
5322 /// `region.x * bpp` whenever the view is both offset horizontally *and*
5323 /// reaches the parent's last row. Such a view constructs fine but fails at
5324 /// [`map`](TensorTrait::map) time with [`Error::InsufficientCapacity`] —
5325 /// loudly, never as an out-of-bounds access. For a tiled destination this is
5326 /// the bottom row of tiles at every column but the first. The single-row
5327 /// case is already special-cased below (it keeps its own tight stride); the
5328 /// general fix is to clamp the last row's exposure to the view's own row
5329 /// bytes, which is not yet implemented.
5330 ///
5331 /// # Errors
5332 ///
5333 /// - [`Error::RegionOutOfBounds`] if `region` exceeds the image bounds.
5334 /// - [`Error::InvalidOperation`] if the tensor is not a packed-format image
5335 /// (planar/semi-planar spatial sub-rects are not a single strided window;
5336 /// use [`batch`](Self::batch) for batched planar tensors).
5337 pub fn view(&self, region: Region) -> Result<Tensor<T>> {
5338 let fmt = self.format.ok_or_else(|| {
5339 Error::InvalidOperation("view() requires a formatted image tensor".into())
5340 })?;
5341 if fmt.layout() != PixelLayout::Packed {
5342 return Err(Error::InvalidOperation(format!(
5343 "view() supports packed formats only (got {fmt:?}); use batch(n) for batched \
5344 planar tensors"
5345 )));
5346 }
5347 let w = self
5348 .width()
5349 .ok_or_else(|| Error::InvalidOperation("view(): tensor has no image width".into()))?;
5350 let h = self
5351 .height()
5352 .ok_or_else(|| Error::InvalidOperation("view(): tensor has no image height".into()))?;
5353 if !region.fits_within(w, h) {
5354 return Err(Error::RegionOutOfBounds {
5355 region,
5356 bounds: (w, h),
5357 });
5358 }
5359 let elem = std::mem::size_of::<T>();
5360 let bpp = fmt.channels() * elem;
5361 let stride = self
5362 .effective_row_stride()
5363 .expect("a tensor with an image width has a row stride");
5364 let offset = region
5365 .y
5366 .checked_mul(stride)
5367 .and_then(|yo| yo.checked_add(region.x.checked_mul(bpp)?))
5368 .ok_or(Error::InvalidSize(region.y))?;
5369 let sub_shape = fmt
5370 .allocation_shape(region.width, region.height)
5371 .ok_or_else(|| Error::InvalidShape(format!("view(): invalid shape for {fmt:?}")))?;
5372 let mut t = self.subview(offset, &sub_shape)?;
5373 // A multi-row sub-rect must advance rows by the PARENT pitch so each row
5374 // addresses the correct columns. A single-row view uses its own tight
5375 // stride — the parent pitch would make the strided `map()` (which exposes
5376 // `stride × rows`) expose a trailing row that runs past the buffer tail
5377 // for an offset (x>0 / bottom) view. The GL backend does NOT rely on this
5378 // (single-row-tight) `row_stride`: it reads the parent pitch from
5379 // `view_origin.parent_row_stride` so its import/cache pitch stays
5380 // parent-consistent for views of any height — see `ViewOrigin`.
5381 let view_stride = if region.height > 1 {
5382 stride
5383 } else {
5384 region.width * bpp
5385 };
5386 t.set_row_stride_unchecked(view_stride);
5387 // Snapshot the parent `(w, h, row_stride)` so the GL backend imports the
5388 // parent once (keyed on the parent pitch, not this view's possibly-tight
5389 // single-row stride) and renders this sub-rect as a `glViewport`/
5390 // `glScissor` ROI at `(region.x, region.y)`. Composes when viewing an
5391 // existing view.
5392 t.view_origin = Some(self.compose_view_origin(w, h, stride, region.x, region.y));
5393 Ok(t)
5394 }
5395
5396 /// Build the [`ViewOrigin`] for a new sub-region of `self`. When `self` is a
5397 /// whole tensor the snapshot names `self` as the parent; when `self` is
5398 /// already a view, the snapshot keeps the **root** parent and accumulates the
5399 /// local origin so nested views still resolve to one import.
5400 fn compose_view_origin(
5401 &self,
5402 parent_width: usize,
5403 parent_height: usize,
5404 parent_row_stride: usize,
5405 x: usize,
5406 y: usize,
5407 ) -> ViewOrigin {
5408 match self.view_origin {
5409 Some(root) => ViewOrigin {
5410 parent_width: root.parent_width,
5411 parent_height: root.parent_height,
5412 parent_row_stride: root.parent_row_stride,
5413 x: root.x.saturating_add(x),
5414 y: root.y.saturating_add(y),
5415 },
5416 None => ViewOrigin {
5417 parent_width,
5418 parent_height,
5419 parent_row_stride,
5420 x,
5421 y,
5422 },
5423 }
5424 }
5425
5426 /// The wrapped `PboTensor`, for this module's own PBO accessors.
5427 ///
5428 /// Crate-private: `PboTensor<T>` is a `static`-side value with no
5429 /// `dynamic` counterpart to lend (there the buffer lives inside
5430 /// `libedgefirst_tensor.so`), so a public borrow of it could not have
5431 /// the same signature on both backends. Callers use
5432 /// [`Self::pbo_id`]/[`Self::pbo_is_mapped`], which are the whole surface
5433 /// `edgefirst-image` ever used.
5434 pub(crate) fn as_pbo(&self) -> Option<&PboTensor<T>> {
5435 match &self.storage {
5436 TensorStorage::Pbo(p) => Some(p),
5437 _ => None,
5438 }
5439 }
5440
5441 /// The GL buffer ID backing this tensor (for GL backends / the
5442 /// cross-package protocol's `handle`). `None` when the tensor is not
5443 /// PBO-backed.
5444 pub fn pbo_id(&self) -> Option<u32> {
5445 self.as_pbo().map(PboTensor::buffer_id)
5446 }
5447
5448 /// Whether this PBO currently holds -- or is establishing -- a CPU
5449 /// mapping. `None` when the tensor is not PBO-backed.
5450 ///
5451 /// The question a GL caller is really asking is "is this buffer free
5452 /// for GL operations?", so a map or unmap still in flight answers
5453 /// `Some(true)`: the GL side must not touch the buffer until the
5454 /// transition finishes. See [`PboTensor::is_mapped`].
5455 pub fn pbo_is_mapped(&self) -> Option<bool> {
5456 self.as_pbo().map(PboTensor::is_mapped)
5457 }
5458
5459 /// The C-ABI `PboOpsVtable` address backing this tensor (for the
5460 /// cross-package protocol's `ptr` under `kind::PBO`; see
5461 /// [`crate::protocol::TensorDesc::ptr`]'s doc comment). `None` when the
5462 /// tensor is not PBO-backed.
5463 pub fn pbo_vtable_ptr(&self) -> Option<*const std::ffi::c_void> {
5464 self.as_pbo()
5465 .map(|p| p.pbo_vtable() as *const _ as *const std::ffi::c_void)
5466 }
5467
5468 /// A type-erased keepalive that must stay alive for at least as long as
5469 /// [`Self::pbo_vtable_ptr`]'s address is used -- see
5470 /// [`PboTensor::pbo_keepalive`]'s own doc comment. `None` when the
5471 /// tensor is not PBO-backed.
5472 pub fn pbo_keepalive(&self) -> Option<Arc<dyn Send + Sync>> {
5473 self.as_pbo().map(PboTensor::pbo_keepalive)
5474 }
5475
5476 /// Downcast to DMA tensor reference (for EGL import, G2D).
5477 #[cfg(target_os = "linux")]
5478 pub fn as_dma(&self) -> Option<&DmaTensor<T>> {
5479 match &self.storage {
5480 TensorStorage::Dma(d) => Some(d),
5481 _ => None,
5482 }
5483 }
5484
5485 /// Borrow the DMA-BUF file descriptor backing this tensor.
5486 ///
5487 /// # Returns
5488 ///
5489 /// A borrowed reference to the DMA-BUF file descriptor, tied to `self`'s
5490 /// lifetime.
5491 ///
5492 /// # Errors
5493 ///
5494 /// Returns `Error::NotImplemented` if the tensor is not DMA-backed.
5495 #[cfg(target_os = "linux")]
5496 pub fn dmabuf(&self) -> Result<std::os::fd::BorrowedFd<'_>> {
5497 use std::os::fd::AsFd;
5498 match &self.storage {
5499 TensorStorage::Dma(dma) => Ok(dma.fd.as_fd()),
5500 _ => Err(Error::NotImplemented(format!(
5501 "dmabuf requires DMA-backed tensor, got {:?}",
5502 self.storage.memory()
5503 ))),
5504 }
5505 }
5506
5507 /// Construct a Tensor from a PBO tensor (for GL backends that allocate PBOs).
5508 ///
5509 /// Infallible in this backend (`Ok` always) -- kept as `Result` to match
5510 /// `dynamic`'s own `Tensor::from_pbo` (`dynamic_tensor.rs`), which
5511 /// genuinely can fail: it must allocate a real `ef_tensor_*` handle
5512 /// alongside `pbo` to carry shape/dtype/format metadata (this backend
5513 /// has nowhere else to put it), an allocation this backend never
5514 /// needed since it just wraps `pbo` into `TensorStorage::Pbo` directly.
5515 /// One shared, fallible signature lets `edgefirst-image`'s own call
5516 /// sites (`gl/threaded.rs`) stay backend-agnostic.
5517 pub fn from_pbo(pbo: PboTensor<T>) -> Result<Self> {
5518 Ok(Self {
5519 storage: TensorStorage::Pbo(pbo),
5520 format: None,
5521 chroma: None,
5522 row_stride: None,
5523 plane_offset: None,
5524 quantization: None,
5525 cuda: None,
5526 colorimetry: None,
5527 cpu_access: CpuAccess::ReadWrite,
5528 compression: None,
5529 view_origin: None,
5530 view_marked: false,
5531 })
5532 }
5533
5534 /// The CUDA registration for this tensor, if any (set at creation on CUDA devices).
5535 pub fn cuda(&self) -> Option<&crate::cuda::CudaHandle> {
5536 self.cuda.as_ref()
5537 }
5538
5539 /// Whether a CUDA registration is attached to this tensor.
5540 pub fn is_cuda_attached(&self) -> bool {
5541 self.cuda.is_some()
5542 }
5543
5544 /// Attach a CUDA handle (called by ImageProcessor::create_image after registering a PBO).
5545 pub fn set_cuda_handle(&mut self, h: crate::cuda::CudaHandle) {
5546 self.cuda = Some(h);
5547 }
5548
5549 /// Fast-fail CUDA map: None (no GL routing) when no handle; else map (PBO routes to the GL worker).
5550 ///
5551 /// Returns a scoped [`CudaMap`] guard holding the raw CUDA device pointer
5552 /// for the duration of the mapping. For GL-buffer-backed tensors the unmap is deferred until the
5553 /// guard drops, freeing the PBO for the next `convert()` call. When no CUDA handle is attached
5554 /// (the common case for plain `Mem`/`DMA` tensors without CUDA registration), returns `None`
5555 /// immediately — no GL routing, no allocation.
5556 ///
5557 /// # Example — zero-copy CUDA input with host fallback
5558 ///
5559 /// ```no_run
5560 /// use edgefirst_tensor::{Tensor, TensorMemory, TensorTrait};
5561 /// # fn feed_tensorrt(_dptr: *mut std::ffi::c_void, _bytes: usize) {}
5562 /// # fn demo(t: &Tensor<f32>) {
5563 /// // Try the zero-copy CUDA device pointer first.
5564 /// if let Some(cuda) = t.cuda_map() {
5565 /// feed_tensorrt(cuda.device_ptr(), cuda.len());
5566 /// // `cuda` (a CudaMap guard) unmaps when it goes out of scope, freeing
5567 /// // the GPU buffer for the next convert().
5568 /// } else {
5569 /// // Fall back to the host mapping when no CUDA handle is attached.
5570 /// let _host = t.map().expect("host map fallback must succeed");
5571 /// // `_host` is a TensorMap<f32> that derefs to &[f32].
5572 /// }
5573 /// # }
5574 /// ```
5575 ///
5576 /// # D3D11 texture tensors: the mapping is tight rows
5577 ///
5578 /// On Windows the mapping is a linear copy of the texture laid out at
5579 /// `width * bytes_per_texel`, and `len()` is that times the row count --
5580 /// *not* [`row_stride`](Self::row_stride), which reports the D3D11 staging
5581 /// pitch a CPU map sees and is the larger number on a padded backing. A
5582 /// consumer that strides the device pointer by `row_stride()` reads past
5583 /// the end of the mapping.
5584 ///
5585 /// Semi-planar formats (NV12, NV16, NV24) are the exception: their texture
5586 /// is allocated as wide as its own row pitch, so the tight row *is* the
5587 /// pitch and the mapping equals `row_stride() x combined_height` bytes.
5588 /// Striding the device pointer by `row_stride()` is exactly right there.
5589 pub fn cuda_map(&self) -> Option<crate::cuda::CudaMap<'_>> {
5590 self.cuda.as_ref()?.map()
5591 }
5592
5593 /// Writable counterpart of [`cuda_map`](Self::cuda_map). See
5594 /// [`CudaHandle::map_mut`](crate::cuda::CudaHandle::map_mut) for which
5595 /// backings distinguish the two.
5596 ///
5597 /// The D3D11 arm is the one that does: the guard's release copies the
5598 /// buffer back into the texture. Its layout is [`cuda_map`](Self::cuda_map)'s
5599 /// -- tight rows of `width * bytes_per_texel`, not `row_stride()` -- and
5600 /// writing it at the wrong pitch scrambles the image rather than failing.
5601 pub fn cuda_map_mut(&self) -> Option<crate::cuda::CudaMap<'_>> {
5602 self.cuda.as_ref()?.map_mut()
5603 }
5604
5605 /// Attempt to attach a CUDA `ExternalMemory` handle for DMA-backed tensors.
5606 ///
5607 /// On a CUDA-capable host, imports the DMA-BUF fd via
5608 /// `cudaImportExternalMemory(OpaqueFd)` and maps it to a device pointer.
5609 /// Sets `self.cuda` to a persistent `ExternalMem` handle on success. No-op
5610 /// if CUDA is unavailable, the tensor is not DMA-backed, or a handle is
5611 /// already set. Import failure is silently ignored — the tensor remains
5612 /// usable without a CUDA handle.
5613 ///
5614 /// # RUNTIME-UNVALIDATED
5615 ///
5616 /// No test platform has both `/dev/dma_heap` and a CUDA device. ABI is
5617 /// layout-asserted vs. CUDA 12.6 `driver_types.h`; the mechanism is proven
5618 /// by gpu-probe O5 on Orin. Best-effort: tensor creation never fails here.
5619 #[cfg(target_os = "linux")]
5620 pub fn try_init_dma_cuda(&mut self) {
5621 // Fast-path: already imported, CUDA not available, or not a DMA tensor.
5622 if self.cuda.is_some() || !crate::cuda::is_cuda_available() {
5623 return;
5624 }
5625 let (raw_fd, buf_size) = match &self.storage {
5626 TensorStorage::Dma(dma) => {
5627 use std::os::fd::AsRawFd;
5628 (dma.fd.as_raw_fd(), dma.buf_size)
5629 }
5630 _ => return,
5631 };
5632 if let Some((ext, dptr)) = crate::cuda::import_dma_fd(raw_fd, buf_size) {
5633 self.cuda = Some(crate::cuda::CudaHandle::new_external(ext, dptr, buf_size));
5634 }
5635 }
5636
5637 /// Attempt to attach a CUDA external-memory handle for D3D11 texture
5638 /// tensors — the Windows counterpart of
5639 /// [`try_init_dma_cuda`](Self::try_init_dma_cuda).
5640 ///
5641 /// Imports the texture's shared NT handle via
5642 /// `cudaImportExternalMemory(D3D11Resource)` and the process device's
5643 /// fence as an external semaphore, so [`cuda_map`](Self::cuda_map) hands
5644 /// out a linear device buffer whose refresh is ordered after the D3D11
5645 /// work a producer recorded with [`set_gpu_write`](Self::set_gpu_write).
5646 ///
5647 /// No-op when a handle is already attached or the tensor is not
5648 /// texture-backed. Import failure is silently ignored: the WARP adapter
5649 /// has no CUDA device at all, and a texture whose import fails is still a
5650 /// valid tensor whose `cuda_map()` is `None`.
5651 #[cfg(target_os = "windows")]
5652 pub fn try_init_d3d11_cuda(&mut self) {
5653 if self.cuda.is_some() {
5654 return;
5655 }
5656 let TensorStorage::Dma(tex) = &self.storage else {
5657 return;
5658 };
5659 if let Some(h) = tex.try_init_cuda() {
5660 self.cuda = Some(h);
5661 }
5662 }
5663}
5664
5665/// Windows D3D11 texture accessors. No `'static` bound: none of these needs
5666/// `dtype_of::<T>()`.
5667#[cfg(all(not(edgefirst_tensor_dynamic), target_os = "windows"))]
5668impl<T> Tensor<T>
5669where
5670 T: Num + Clone + fmt::Debug + Send + Sync,
5671{
5672 /// The borrowed `ID3D11Texture2D*` backing this tensor, valid while the
5673 /// tensor lives. `None` when the tensor is not texture-backed.
5674 pub fn d3d11_texture(&self) -> Option<*mut std::ffi::c_void> {
5675 match &self.storage {
5676 TensorStorage::Dma(t) => Some(t.texture_ptr()),
5677 _ => None,
5678 }
5679 }
5680
5681 /// The texture geometry the HAL chose for this image: DXGI format, texel
5682 /// dimensions and the matching GL internal format. `None` when the tensor
5683 /// is not texture-backed.
5684 pub fn d3d11_layout(&self) -> Option<crate::d3d11_layout::D3d11ImageLayout> {
5685 match &self.storage {
5686 TensorStorage::Dma(t) => Some(t.layout()),
5687 _ => None,
5688 }
5689 }
5690
5691 /// A duplicated NT handle the caller owns: open it on a D3D12 device,
5692 /// another D3D11 device, or CUDA, or duplicate it into another process.
5693 /// Closing it does not affect this tensor.
5694 pub fn d3d11_shared_handle(&self) -> Result<std::os::windows::io::OwnedHandle> {
5695 match &self.storage {
5696 TensorStorage::Dma(t) => t.shared_handle(),
5697 _ => Err(Error::NotImplemented(
5698 "not a D3D11 texture tensor".to_owned(),
5699 )),
5700 }
5701 }
5702
5703 /// The fence handle plus value a GPU consumer waits on before reading this
5704 /// texture, or `None` when no GPU write has been recorded on it.
5705 pub fn gpu_completion(&self) -> Result<Option<crate::d3d11::GpuCompletion>> {
5706 match &self.storage {
5707 TensorStorage::Dma(t) => t.gpu_completion(),
5708 _ => Err(Error::NotImplemented(
5709 "not a D3D11 texture tensor".to_owned(),
5710 )),
5711 }
5712 }
5713
5714 /// Record that GPU work writing this texture completes at `value` of the
5715 /// process device's fence. Monotonic: an older value never displaces a
5716 /// newer one, so recording out of order is safe rather than merely
5717 /// tolerated.
5718 ///
5719 /// Takes `&self`, not `&mut self`: the fence slot is an `AtomicU64` and
5720 /// this is a `fetch_max` on it. A producer records a write on the very
5721 /// tensor its consumers already hold and are reading, so demanding
5722 /// exclusive access here would forbid the one call this exists for.
5723 pub fn set_gpu_write(&self, value: u64) -> Result<()> {
5724 match &self.storage {
5725 TensorStorage::Dma(t) => {
5726 t.set_gpu_write(value);
5727 Ok(())
5728 }
5729 _ => Err(Error::NotImplemented(
5730 "not a D3D11 texture tensor".to_owned(),
5731 )),
5732 }
5733 }
5734
5735 /// The tensor's own NT handle value -- not a duplicate -- for descriptors
5736 /// whose consumer keeps the producing tensor alive.
5737 ///
5738 /// Borrowed rather than duplicated because the C exports (task A6) hand
5739 /// out owned duplicates, which a `Copy` descriptor cannot close; the
5740 /// cross-package protocol's `kind::D3D11_TEXTURE` reads this instead.
5741 pub(crate) fn d3d11_shared_handle_value(&self) -> Option<usize> {
5742 match &self.storage {
5743 TensorStorage::Dma(t) => Some(t.shared_handle_value()),
5744 _ => None,
5745 }
5746 }
5747}
5748
5749/// The recorded GPU write value, on every platform: the descriptor path and
5750/// the Python bindings forward it without a `cfg` of their own, and only
5751/// Windows can record one.
5752#[cfg(not(edgefirst_tensor_dynamic))]
5753impl<T> Tensor<T>
5754where
5755 T: Num + Clone + fmt::Debug + Send + Sync,
5756{
5757 /// The fence value of the newest GPU write recorded on this tensor, or 0
5758 /// when there is none: the `value` of `gpu_completion` without the
5759 /// duplicated fence handle, for a consumer that already holds the process
5760 /// fence. 0 for a tensor that is not a D3D11 texture, and on every
5761 /// platform other than Windows.
5762 pub fn gpu_write_value(&self) -> u64 {
5763 #[cfg(target_os = "windows")]
5764 {
5765 if let TensorStorage::Dma(t) = &self.storage {
5766 return t.last_gpu_write();
5767 }
5768 }
5769 0
5770 }
5771}
5772
5773/// Records the D3D11 backing's row pitch on `t` when it exceeds the natural
5774/// stride for its format and width.
5775///
5776/// The staging pitch is the row stride CPU consumers see, and for a
5777/// semi-planar tensor it is also the texture's own width, which the GL
5778/// backend's Path B shader takes as its `tex_width`. Three constructors need
5779/// it identically -- `Tensor::image`'s Windows arm and the two wrap
5780/// constructors below -- and a wrapped texture is the case that makes it
5781/// load-bearing: without it such a tensor reports the natural `even(width)`
5782/// stride while its texture is the driver's pitch wide, and the shader samples
5783/// a padded texture at the image width.
5784///
5785/// `t.format` must already be set: the natural stride comes from it.
5786#[cfg(all(not(edgefirst_tensor_dynamic), target_os = "windows"))]
5787fn record_d3d11_row_stride<T>(t: &mut Tensor<T>, format: PixelFormat, what: &str)
5788where
5789 T: Num + Clone + fmt::Debug + Send + Sync,
5790{
5791 let pitch = match &t.storage {
5792 TensorStorage::Dma(tex) => tex.image_backing_row_stride(),
5793 _ => None,
5794 };
5795 let (Some(pitch), Some(natural)) = (pitch, t.effective_row_stride()) else {
5796 return;
5797 };
5798 if pitch > natural {
5799 log::debug!(
5800 "{what}: the D3D11 backing pitch for the {format:?} texture is {pitch} bytes \
5801 (natural {natural}); recording row stride"
5802 );
5803 t.set_row_stride_unchecked(pitch);
5804 }
5805}
5806
5807/// The two texture-importing constructors. Separate from the accessors above
5808/// because `dtype_of::<T>()` needs `T: 'static`.
5809#[cfg(all(not(edgefirst_tensor_dynamic), target_os = "windows"))]
5810impl<T> Tensor<T>
5811where
5812 T: Num + Clone + fmt::Debug + Send + Sync + 'static,
5813{
5814 /// Wrap an existing `ID3D11Texture2D` as a tensor. The texture is checked
5815 /// against the layout the HAL would have chosen for `format`/`width`/
5816 /// `height` and must live on the HAL's process device; ownership stays
5817 /// with the caller, since this takes its own reference.
5818 ///
5819 /// The backing's row pitch is recorded as the tensor's
5820 /// [`row_stride`](Self::row_stride) when it exceeds the natural stride, so
5821 /// a wrapped semi-planar texture -- whose width *is* the pitch -- reports
5822 /// the geometry its consumers address it by.
5823 ///
5824 /// # Safety
5825 ///
5826 /// `texture` must be null or a live `ID3D11Texture2D` created on the HAL
5827 /// device ([`crate::d3d11::device()`]).
5828 pub unsafe fn from_d3d11_texture(
5829 texture: *mut std::ffi::c_void,
5830 width: usize,
5831 height: usize,
5832 format: PixelFormat,
5833 access: CpuAccess,
5834 name: Option<&str>,
5835 ) -> Result<Self> {
5836 let dtype = dtype_of::<T>().ok_or_else(|| {
5837 Error::InvalidArgument(format!(
5838 "from_d3d11_texture: element type {} has no DType",
5839 std::any::type_name::<T>()
5840 ))
5841 })?;
5842 let shape = format.allocation_shape(width, height).ok_or_else(|| {
5843 Error::InvalidArgument(format!(
5844 "invalid dimensions {width}x{height} for format {format:?}"
5845 ))
5846 })?;
5847 // SAFETY: the caller guarantees `texture` is null or a live
5848 // `ID3D11Texture2D`; the storage takes its own reference.
5849 let storage = unsafe {
5850 crate::d3d11::texture::D3d11TextureTensor::<T>::from_d3d11_texture(
5851 texture, width, height, format, dtype, &shape, access, name,
5852 )
5853 }?;
5854 let mut t = Self::wrap(TensorStorage::Dma(storage));
5855 t.format = Some(format);
5856 record_d3d11_row_stride(&mut t, format, "from_d3d11_texture");
5857 t.cpu_access = access;
5858 t.try_init_d3d11_cuda();
5859 Ok(t)
5860 }
5861
5862 /// Open a shared texture by its NT handle on the HAL's device. When
5863 /// `completion` is given, its fence value is waited for on the immediate
5864 /// context, so a same-device reader needs no further ordering. Both
5865 /// handles stay owned by the caller: this duplicates what it keeps.
5866 ///
5867 /// Records the backing's row pitch as
5868 /// [`row_stride`](Self::row_stride) on the same rule as
5869 /// [`from_d3d11_texture`](Self::from_d3d11_texture).
5870 ///
5871 /// # Safety
5872 ///
5873 /// `handle` must be an NT shared handle of a D3D11 texture, valid in this
5874 /// process, and `completion`'s handle a shared fence handle.
5875 pub unsafe fn from_d3d11_shared_handle(
5876 handle: std::os::windows::io::RawHandle,
5877 width: usize,
5878 height: usize,
5879 format: PixelFormat,
5880 access: CpuAccess,
5881 completion: Option<(std::os::windows::io::RawHandle, u64)>,
5882 name: Option<&str>,
5883 ) -> Result<Self> {
5884 let dtype = dtype_of::<T>().ok_or_else(|| {
5885 Error::InvalidArgument(format!(
5886 "from_d3d11_shared_handle: element type {} has no DType",
5887 std::any::type_name::<T>()
5888 ))
5889 })?;
5890 let shape = format.allocation_shape(width, height).ok_or_else(|| {
5891 Error::InvalidArgument(format!(
5892 "invalid dimensions {width}x{height} for format {format:?}"
5893 ))
5894 })?;
5895 // SAFETY: the caller guarantees the handles are valid in this process.
5896 let storage = unsafe {
5897 crate::d3d11::texture::D3d11TextureTensor::<T>::from_d3d11_shared_handle(
5898 handle, width, height, format, dtype, &shape, access, completion, name,
5899 )
5900 }?;
5901 let mut t = Self::wrap(TensorStorage::Dma(storage));
5902 t.format = Some(format);
5903 record_d3d11_row_stride(&mut t, format, "from_d3d11_shared_handle");
5904 t.cpu_access = access;
5905 t.try_init_d3d11_cuda();
5906 Ok(t)
5907 }
5908}
5909
5910// Quantization accessors — type-gated to integer element types via the
5911// sealed `IntegerType` trait. Calling `.quantization()` on a `Tensor<f32>`
5912// produces a compile error, not a runtime one.
5913#[cfg(not(edgefirst_tensor_dynamic))]
5914impl<T> Tensor<T>
5915where
5916 T: IntegerType + Num + Clone + fmt::Debug + Send + Sync,
5917{
5918 /// Quantization metadata for this tensor, if set.
5919 pub fn quantization(&self) -> Option<&Quantization> {
5920 self.quantization.as_ref()
5921 }
5922
5923 /// Attach quantization metadata to this tensor. Validates against the
5924 /// tensor's shape — returns [`Error::QuantizationInvalid`] on any
5925 /// inconsistency (mismatched scale/zp lengths, out-of-range axis, etc.).
5926 pub fn set_quantization(&mut self, q: Quantization) -> Result<()> {
5927 q.validate(self.shape())?;
5928 self.quantization = Some(q);
5929 Ok(())
5930 }
5931
5932 /// Builder-style variant of [`Self::set_quantization`]. Consumes `self`
5933 /// and returns `Result<Self>` — on success yields the tensor with the
5934 /// attached quantization; on validation failure returns
5935 /// [`Error::QuantizationInvalid`] and drops `self` (the tensor is not
5936 /// returned in the error arm).
5937 pub fn with_quantization(mut self, q: Quantization) -> Result<Self> {
5938 self.set_quantization(q)?;
5939 Ok(self)
5940 }
5941
5942 /// Clear any quantization metadata on this tensor.
5943 pub fn clear_quantization(&mut self) {
5944 self.quantization = None;
5945 }
5946}
5947
5948/// Rank of one image's shape in `layout`: `[H, W, C]` / `[C, H, W]` are 3-D,
5949/// semi-planar `[H*k, W]` is 2-D. A batched tensor has one more.
5950#[cfg(not(edgefirst_tensor_dynamic))]
5951fn image_rank(layout: PixelLayout) -> usize {
5952 match layout {
5953 PixelLayout::SemiPlanar => 2,
5954 PixelLayout::Packed | PixelLayout::Planar => 3,
5955 }
5956}
5957
5958#[cfg(not(edgefirst_tensor_dynamic))]
5959fn map_strided_host_view<'a, T>(
5960 tensor: &Tensor<T>,
5961 stride: usize,
5962 access: CpuAccess,
5963 non_blocking: bool,
5964) -> Result<crate::view::HostView<'a, T>>
5965where
5966 T: Num + Clone + fmt::Debug + Send + Sync + 'a,
5967{
5968 // Rows sit at `stride`-byte spacing; see `strided_rows` for the count.
5969 let rows = strided_rows(tensor)?;
5970 let total_bytes = stride.checked_mul(rows).ok_or_else(|| {
5971 Error::InvalidOperation(format!(
5972 "Tensor::map: row_stride {stride} × rows {rows} overflows usize"
5973 ))
5974 })?;
5975 map_strided_storage(tensor, stride, rows, total_bytes, access, non_blocking)
5976}
5977
5978/// Rows a strided map or flat copy of `tensor` spans at its row stride: a
5979/// raw tensor's leading dim, else every image row of every batch element --
5980/// `H` per packed image, `C * H` per planar one (C planes of H rows), and the
5981/// combined `H*k` per semi-planar one. No product here can overflow: each is
5982/// a factor of the shape's element count, which the allocation already holds.
5983#[cfg(not(edgefirst_tensor_dynamic))]
5984fn strided_rows<T>(tensor: &Tensor<T>) -> Result<usize>
5985where
5986 T: Num + Clone + fmt::Debug + Send + Sync,
5987{
5988 let shape = tensor.shape();
5989 let Some(fmt) = tensor.format else {
5990 return shape.first().copied().ok_or_else(|| {
5991 Error::InvalidOperation("strided access requires a non-empty shape".into())
5992 });
5993 };
5994 let (_, dims) = tensor.image_dims().ok_or_else(|| {
5995 Error::InvalidOperation(format!(
5996 "strided access to a {fmt:?} tensor requires its image shape, optionally \
5997 batched; got {shape:?}"
5998 ))
5999 })?;
6000 let per_image = match fmt.layout() {
6001 PixelLayout::Planar => dims[0] * dims[1],
6002 PixelLayout::Packed | PixelLayout::SemiPlanar => dims[0],
6003 };
6004 let batch: usize = shape[..shape.len() - dims.len()].iter().product();
6005 Ok(batch * per_image)
6006}
6007
6008#[cfg(not(edgefirst_tensor_dynamic))]
6009fn map_strided_storage<'a, T>(
6010 tensor: &Tensor<T>,
6011 _stride: usize,
6012 _rows: usize,
6013 total_bytes: usize,
6014 access: CpuAccess,
6015 _non_blocking: bool,
6016) -> Result<crate::view::HostView<'a, T>>
6017where
6018 T: Num + Clone + fmt::Debug + Send + Sync + 'a,
6019{
6020 match &tensor.storage {
6021 #[cfg(target_os = "linux")]
6022 TensorStorage::Dma(dma) => {
6023 // `set_row_stride()` only validates `stride >= min_stride`,
6024 // not that `stride × rows` fits the DMA-BUF, so re-check
6025 // here — mapping past `buf_size` would SIGBUS on access.
6026 // `buf_size` is real kernel truth either way, not the
6027 // importer's say-so: for a self-allocated tensor it is
6028 // the allocation size; for an imported one it is
6029 // `fstat`-derived when the kernel reports a usable size
6030 // (`DmaTensor::from_fd`), or the declared logical size in
6031 // the rare fallback -- either way an out-of-range mmap
6032 // against the real dma-buf is still rejected by the
6033 // kernel, so this bound stays sound regardless.
6034 let available_bytes = dma.buf_size.saturating_sub(dma.mmap_offset);
6035 if total_bytes > available_bytes {
6036 return Err(Error::InvalidOperation(format!(
6037 "Tensor::map: strided mapping needs {total_bytes} bytes \
6038 but DMA buffer only has {available_bytes} available \
6039 (buf_size={}, mmap_offset={}, stride={_stride}, rows={_rows}); \
6040 the row_stride was likely set larger than the original allocation",
6041 dma.buf_size, dma.mmap_offset
6042 )));
6043 }
6044 dma.map_with_byte_size(total_bytes, access)
6045 }
6046 TensorStorage::Mem(mem) => {
6047 let capacity = tensor.storage.capacity_bytes();
6048 if total_bytes > capacity {
6049 return Err(Error::InsufficientCapacity {
6050 needed: total_bytes,
6051 capacity,
6052 });
6053 }
6054 mem.map_with_byte_size(total_bytes, access)
6055 }
6056 #[cfg(unix)]
6057 TensorStorage::Shm(shm) => {
6058 let capacity = tensor.storage.capacity_bytes();
6059 if total_bytes > capacity {
6060 return Err(Error::InsufficientCapacity {
6061 needed: total_bytes,
6062 capacity,
6063 });
6064 }
6065 shm.map_with_byte_size(total_bytes, access)
6066 }
6067 // macOS/iOS: `TensorStorage::Dma` is the IOSurface. The lock yields
6068 // the full surface base address, and the row pitch
6069 // (`IOSurfaceGetBytesPerRow`) is known from the API for both
6070 // self-allocated and imported surfaces — unlike a foreign
6071 // DMA-BUF — so a strided CPU view is sound and zero-copy.
6072 #[cfg(any(target_os = "macos", target_os = "ios"))]
6073 TensorStorage::Dma(io) => {
6074 // A sub-view's window is `buf_size − view_offset`; the strided
6075 // span must fit the window, not the whole surface.
6076 let available = io.buf_size.saturating_sub(io.view_offset);
6077 if total_bytes > available {
6078 return Err(Error::InsufficientCapacity {
6079 needed: total_bytes,
6080 capacity: available,
6081 });
6082 }
6083 io.map_with_byte_size(total_bytes, access)
6084 }
6085 // Android: `TensorStorage::Dma` is the AHardwareBuffer. The lock
6086 // yields the full buffer base address, and the row pitch is
6087 // known from the allocator-filled descriptor — so a strided CPU
6088 // view is sound and zero-copy, same as IOSurface.
6089 #[cfg(target_os = "android")]
6090 TensorStorage::Dma(ahb) => {
6091 // A sub-view's window is `buf_size − view_offset`; the strided
6092 // span must fit the window, not the whole buffer.
6093 let available = ahb.buf_size.saturating_sub(ahb.view_offset);
6094 if total_bytes > available {
6095 return Err(Error::InsufficientCapacity {
6096 needed: total_bytes,
6097 capacity: available,
6098 });
6099 }
6100 ahb.map_with_byte_size(total_bytes, access)
6101 }
6102 // Windows: `TensorStorage::Dma` is the D3D11 texture, whose CPU side
6103 // is the staging copy at the driver's pitch -- the very pitch
6104 // `Tensor::image` recorded as this tensor's row stride. Its
6105 // `map_with_byte_size` bounds-checks the span against this tensor's
6106 // own window, as the IOSurface arm above does. This is also the only
6107 // backing with a real non-blocking map, so it is the only arm that
6108 // reads `_non_blocking`.
6109 #[cfg(target_os = "windows")]
6110 TensorStorage::Dma(tex) => {
6111 if _non_blocking {
6112 tex.try_map_with_byte_size(total_bytes, access)
6113 } else {
6114 tex.map_with_byte_size(total_bytes, access)
6115 }
6116 }
6117 TensorStorage::Pbo(pbo) => {
6118 // PBO: the GPU-side allocation may have a padded row stride
6119 // (e.g. 64-byte aligned). Expose the full padded buffer so a
6120 // CPU producer (JPEG decoder) and a strided convert source
6121 // can iterate rows via `effective_row_stride()` without
6122 // running past the slice — the logical `pbo.map()` view would
6123 // stop after `shape.product()` and lose bytes past row 0.
6124 // A sub-view's window is `capacity − view_offset`.
6125 let available = pbo.capacity_bytes().saturating_sub(pbo.view_offset);
6126 if total_bytes > available {
6127 return Err(Error::InsufficientCapacity {
6128 needed: total_bytes,
6129 capacity: available,
6130 });
6131 }
6132 pbo.map_with_byte_size(total_bytes, access)
6133 }
6134 // Every `TensorStorage` variant this build can construct is
6135 // matched explicitly above (Mem/Shm/Dma/Pbo, each cfg-gated
6136 // to the platforms it exists on) -- this catch-all exists
6137 // only so the match stays exhaustive across every
6138 // OS/cfg combination without a matching arm on every one of
6139 // them, not because any variant is meant to fall through to
6140 // it. `#[allow(unreachable_patterns)]` because on any given
6141 // platform it genuinely is unreachable.
6142 #[allow(unreachable_patterns)]
6143 _ => Err(Error::InvalidOperation(
6144 "CPU mapping of strided tensors is supported only for HAL-allocated \
6145 Mem/Shm (any platform), DMA (Linux, self-allocated or imported), \
6146 IOSurface (macOS), AHardwareBuffer (Android), D3D11 textures \
6147 (Windows), and PBO"
6148 .into(),
6149 )),
6150 }
6151}
6152
6153#[cfg(not(edgefirst_tensor_dynamic))]
6154impl<T> Tensor<T>
6155where
6156 T: Num + Clone + fmt::Debug + Send + Sync,
6157{
6158 /// The shared body of [`TensorTrait::map_with`] and
6159 /// [`TensorTrait::try_map_with`]: the direction check, the
6160 /// declared-vs-requested telemetry, the strided and plane-offset routing,
6161 /// and finally the backing's own map. `non_blocking` chooses which of the
6162 /// backing's two maps that last step takes; everything before it is the
6163 /// same either way.
6164 fn map_impl<'a>(
6165 &self,
6166 access: CpuAccess,
6167 non_blocking: bool,
6168 ) -> Result<crate::view::HostView<'a, T>>
6169 where
6170 T: 'a,
6171 {
6172 let _span = tracing::trace_span!(
6173 "tensor.map",
6174 memory = ?self.storage.memory(),
6175 ?access,
6176 non_blocking,
6177 )
6178 .entered();
6179 self.cpu_access.require_direction(access, "map_with")?;
6180 // CPU mapping of a strided tensor exposes the full padded buffer
6181 // (`row_stride × rows`) so callers can iterate rows via
6182 // `effective_row_stride()` without running past the slice. This is sound
6183 // only when the HAL owns and can size-check the allocation:
6184 //
6185 // * Self-allocated Mem / Shm tensors (any platform) — the backing
6186 // `Vec` / shm segment is sized by `capacity_bytes()`, checked here.
6187 // * Self-allocated DMA tensors (Linux) — pitch padding from
6188 // `image_with_stride()`; checked against the DMA-BUF `buf_size`.
6189 //
6190 // * Self-allocated PBO tensors (any platform with GL) — the GL buffer
6191 // is sized by `capacity_bytes()` and may carry 64-byte row padding;
6192 // the JPEG decoder mmaps it and convert() reads it, both iterating
6193 // by `row_stride`. Checked against the PBO capacity below.
6194 //
6195 // A strided **imported** DMA-BUF is accepted here too (Task 10b): the
6196 // `stride × rows <= buf_size` check just below is exactly what a
6197 // self-allocated Dma tensor relies on for its own soundness, and
6198 // `buf_size` for an imported fd is `fstat`-derived when the kernel
6199 // reports a usable size (`stat.st_size > 0` and at least the
6200 // declared shape's logical size; `DmaTensor::from_fd`) -- kernel
6201 // truth, not a caller's claim, in that case. The rare fallback (a
6202 // kernel that reports `st_size <= 0` for a DMA-BUF fd) uses the
6203 // declared logical size instead; that is not a hole here either, since
6204 // an out-of-range mmap against the real dma-buf is still rejected by
6205 // the kernel regardless of what this crate believed `buf_size` was.
6206 // So the check is equally load-bearing regardless of who allocated
6207 // the buffer. What
6208 // changed is only *which* callers can now reach this path: the
6209 // cross-package capsule protocol reconstructs an imported alias of a
6210 // HAL-owned Dma tensor and then calls `configure_image`, which
6211 // computes its stride with the exact same alignment rule a
6212 // self-allocated `Tensor::image()` used — so the value being
6213 // bounds-checked here is HAL-computed and format/width-derived, not
6214 // read from an external driver. Rejecting it unconditionally used to
6215 // make every foreign-imported GPU-backed destination fail any CPU
6216 // write path (the JPEG decoder's MCU writer among them) even when the
6217 // write was provably in-bounds. `set_row_stride()` was already
6218 // "trust the caller, bounds-check the map" for every *other* backing
6219 // (Mem/Shm/self-allocated Dma/PBO) — singling out imported Dma out
6220 // for a stricter standard was the inconsistency, not the trust model
6221 // itself. IOSurface keeps its own always-allowed strided path
6222 // unconditionally (see its arm below): a live `IOSurfaceRef` reports
6223 // its real `bytesPerRow` from the OS regardless of who allocated it,
6224 // so there was never anything to distrust there.
6225 if let Some(stride) = self.row_stride {
6226 return map_strided_host_view(self, stride, access, non_blocking);
6227 }
6228 // A non-zero plane offset needs no check here: every storage applies
6229 // its own offset inside its map -- Mem/Shm adjust the slice base, the
6230 // Linux DMA-BUF and Android AHardwareBuffer the mapped base, IOSurface
6231 // and D3D11 textures add their `view_offset`, and PBO starts the
6232 // staged copy at it.
6233 if non_blocking {
6234 self.storage.try_map_with(access)
6235 } else {
6236 self.storage.map_with(access)
6237 }
6238 }
6239}
6240
6241#[cfg(all(
6242 not(edgefirst_tensor_dynamic),
6243 any(
6244 target_os = "macos",
6245 target_os = "ios",
6246 target_os = "android",
6247 target_os = "windows"
6248 )
6249))]
6250impl<T> Tensor<T>
6251where
6252 T: Num + Clone + fmt::Debug + Send + Sync,
6253{
6254 /// The pitched-backing half of [`TensorTrait::set_logical_shape`]: record
6255 /// the backing's own pitch as the row stride of the new `shape`.
6256 fn adopt_backing_pitch(&mut self, shape: &[usize]) {
6257 // Only a backing with a pitch of its own, and only while the shape
6258 // still has the rank the format's row is measured on: a caller may
6259 // flatten an image to one dimension and leave the format behind.
6260 let Some(pitch) = self.storage.backing_row_stride() else {
6261 return;
6262 };
6263 let Some((rank, rows)) = self.format.map(|f| match f.layout() {
6264 PixelLayout::Packed => (3, shape.first().copied()),
6265 PixelLayout::Planar => (3, shape.get(1).copied()),
6266 PixelLayout::SemiPlanar => (2, shape.first().copied()),
6267 }) else {
6268 return;
6269 };
6270 if rank != shape.len() {
6271 return;
6272 }
6273 // The natural stride of the new shape, not of the stride recorded for
6274 // the old one.
6275 let prior = self.row_stride.take();
6276 let Some(natural) = self.effective_row_stride() else {
6277 self.row_stride = prior;
6278 return;
6279 };
6280 if rows == Some(1) {
6281 self.set_row_stride_unchecked(natural);
6282 } else if pitch > natural {
6283 self.set_row_stride_unchecked(pitch);
6284 } else {
6285 self.row_stride = prior;
6286 }
6287 }
6288}
6289
6290#[cfg(not(edgefirst_tensor_dynamic))]
6291impl<T> TensorTrait<T> for Tensor<T>
6292where
6293 T: Num + Clone + fmt::Debug + Send + Sync,
6294{
6295 fn cpu_access(&self) -> CpuAccess {
6296 self.cpu_access
6297 }
6298
6299 fn view_origin(&self) -> Option<ViewOrigin> {
6300 self.view_origin
6301 }
6302
6303 fn is_view(&self) -> bool {
6304 Tensor::is_view(self)
6305 }
6306
6307 fn new(shape: &[usize], name: Option<&str>) -> Result<Self>
6308 where
6309 Self: Sized,
6310 {
6311 Self::new(shape, None, name)
6312 }
6313
6314 #[cfg(unix)]
6315 fn from_fd(fd: std::os::fd::OwnedFd, shape: &[usize], name: Option<&str>) -> Result<Self>
6316 where
6317 Self: Sized,
6318 {
6319 Self::from_fd(fd, shape, name)
6320 }
6321
6322 #[cfg(unix)]
6323 fn clone_fd(&self) -> Result<std::os::fd::OwnedFd> {
6324 self.storage.clone_fd()
6325 }
6326
6327 fn memory(&self) -> TensorMemory {
6328 self.storage.memory()
6329 }
6330
6331 fn name(&self) -> String {
6332 self.storage.name()
6333 }
6334
6335 fn shape(&self) -> &[usize] {
6336 self.storage.shape()
6337 }
6338
6339 /// Bytes available in the underlying allocation, from this tensor's
6340 /// window start -- may exceed [`TensorTrait::size`] when the backing
6341 /// carries spare room beyond the current logical shape (a pitch-aligned
6342 /// image, or a pool tensor sized for a larger image than it holds
6343 /// today). See [`crate::protocol::TensorDesc::capacity`], the
6344 /// cross-package descriptor field this backs.
6345 fn capacity_bytes(&self) -> usize {
6346 self.storage.capacity_bytes()
6347 }
6348
6349 fn reshape(&mut self, shape: &[usize]) -> Result<()> {
6350 if self.chroma.is_some() {
6351 return Err(Error::InvalidOperation(
6352 "cannot reshape a multiplane tensor — decompose planes first".into(),
6353 ));
6354 }
6355 self.storage.reshape(shape)?;
6356 self.format = None;
6357 self.row_stride = None;
6358 self.plane_offset = None;
6359 match self.storage {
6360 TensorStorage::Mem(ref mut m) => m.set_offset(0),
6361 #[cfg(target_os = "linux")]
6362 TensorStorage::Dma(ref mut dma) => dma.mmap_offset = 0,
6363 // No-op: `IoSurfaceTensor::reshape` (iosurface.rs) zeroes its
6364 // own `view_offset` before this match runs.
6365 #[cfg(any(target_os = "macos", target_os = "ios"))]
6366 TensorStorage::Dma(_) => {}
6367 // `D3d11TextureTensor::reshape` does not self-zero, so the clear
6368 // that pairs with `set_plane_offset`'s D3D11 arm lives here.
6369 #[cfg(target_os = "windows")]
6370 TensorStorage::Dma(ref mut tex) => tex.set_view_offset(0),
6371 // No-op: `AHardwareBufferTensor::reshape` (ahardwarebuffer.rs)
6372 // zeroes its own `view_offset` before this match runs, the same
6373 // as `IoSurfaceTensor`. Unlike `set_plane_offset`'s Android arm,
6374 // this one is a genuine self-zeroing no-op, not a still-
6375 // incomplete gap: `view_offset` never survives a reshape here
6376 // regardless of whether `set_plane_offset` can restore it.
6377 #[cfg(target_os = "android")]
6378 TensorStorage::Dma(_) => {}
6379 // No-op: `PboTensor::reshape` (pbo.rs) zeroes its own
6380 // `view_offset` before this match runs, the same as
6381 // `IoSurfaceTensor`/`AHardwareBufferTensor` above -- which is
6382 // why, unlike `set_format`'s clear block, this match never
6383 // needed a `Pbo` arm that writes anything.
6384 TensorStorage::Pbo(_) => {}
6385 // Unlike `PboTensor::reshape`, `ShmTensor::reshape` (shm.rs)
6386 // only revalidates the element count and swaps `self.shape` --
6387 // it never touches `self.offset` -- so this clear site needs
6388 // the arm `set_format`'s does. `cfg(unix)` mirrors the `Shm`
6389 // variant's own gate, not a per-OS multiplex.
6390 #[cfg(unix)]
6391 TensorStorage::Shm(ref mut s) => s.set_offset(0),
6392 }
6393 Ok(())
6394 }
6395
6396 /// Set the logical shape to any shape whose bytes fit the allocation,
6397 /// without [`reshape`](TensorTrait::reshape)'s equal-count constraint.
6398 ///
6399 /// Forwards to the storage's capacity-aware implementation -- the same
6400 /// one [`Tensor::configure_image`] uses for pool reuse, a decode into an
6401 /// oversized reusable buffer reconfiguring it to a smaller image without
6402 /// reallocating -- so the two agree. The pixel format is kept.
6403 ///
6404 /// A narrower image inside a pitched backing (an IOSurface, an
6405 /// AHardwareBuffer, a D3D11 texture) still advances rows by the backing's
6406 /// pitch, so that pitch is recorded as the row stride whenever it exceeds
6407 /// the new shape's natural stride -- the same adoption `configure_image`
6408 /// makes. Without it a texture narrowed from 64 to 16 texels wide reports
6409 /// a 64-byte stride over rows that are 256 bytes apart. A single-row
6410 /// shape records its own tight stride instead, by [`view`](Self::view)'s
6411 /// rule: a strided `map()` spans `stride * rows`, and one pitched row
6412 /// runs past the backing when the window sits in the last row.
6413 fn set_logical_shape(&mut self, shape: &[usize]) -> Result<()> {
6414 self.storage.set_logical_shape(shape)?;
6415 #[cfg(any(
6416 target_os = "macos",
6417 target_os = "ios",
6418 target_os = "android",
6419 target_os = "windows"
6420 ))]
6421 self.adopt_backing_pitch(shape);
6422 Ok(())
6423 }
6424
6425 fn map_with<'a>(&self, access: CpuAccess) -> Result<crate::view::HostView<'a, T>>
6426 where
6427 T: 'a,
6428 {
6429 self.map_impl(access, false)
6430 }
6431
6432 fn try_map_with<'a>(&self, access: CpuAccess) -> Result<crate::view::HostView<'a, T>>
6433 where
6434 T: 'a,
6435 {
6436 self.map_impl(access, true)
6437 }
6438
6439 fn buffer_identity(&self) -> &BufferIdentity {
6440 self.storage.buffer_identity()
6441 }
6442}
6443
6444// ============================================================================
6445// Platform availability helpers
6446// ============================================================================
6447
6448/// Cached result of the Linux DMA-BUF availability probe.
6449#[cfg(target_os = "linux")]
6450static DMA_AVAILABLE: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
6451/// Cached result of the macOS/iOS IOSurface availability probe.
6452#[cfg(any(target_os = "macos", target_os = "ios"))]
6453static IOSURFACE_AVAILABLE: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
6454
6455/// Check if Linux DMA-BUF allocation is available on this system.
6456///
6457/// Linux-specific availability check (typically requires `/dev/dma_heap`
6458/// access — running as root or membership in a video/render group). For
6459/// portable code that wants "any zero-copy GPU buffer", use
6460/// [`is_gpu_buffer_available`] which also covers IOSurface on macOS.
6461///
6462/// This function caches its result after the first call.
6463#[cfg(target_os = "linux")]
6464pub fn is_dma_available() -> bool {
6465 *DMA_AVAILABLE
6466 .get_or_init(|| Tensor::<u8>::new(&[64], Some(TensorMemory::DmaBuf), None).is_ok())
6467}
6468
6469/// Always returns `false` on non-Linux platforms.
6470#[cfg(not(target_os = "linux"))]
6471pub fn is_dma_available() -> bool {
6472 false
6473}
6474
6475/// Check if macOS/iOS IOSurface allocation is available on this system.
6476///
6477/// IOSurface is part of the macOS/iOS OS and is essentially always present;
6478/// this probe catches degraded scenarios such as memory pressure or
6479/// sandboxed contexts where `IOSurfaceCreate` fails. The result is
6480/// cached after the first call.
6481///
6482/// Always returns `false` on non-Apple platforms.
6483#[cfg(any(target_os = "macos", target_os = "ios"))]
6484pub fn is_iosurface_available() -> bool {
6485 *IOSURFACE_AVAILABLE.get_or_init(|| {
6486 // Probe via the same Dma path — on macOS/iOS this routes through
6487 // IoSurfaceTensor::new.
6488 Tensor::<u8>::new(&[64], Some(TensorMemory::DmaBuf), None).is_ok()
6489 })
6490}
6491
6492#[cfg(not(any(target_os = "macos", target_os = "ios")))]
6493pub fn is_iosurface_available() -> bool {
6494 false
6495}
6496
6497/// Cached result of the Android AHardwareBuffer availability probe.
6498#[cfg(target_os = "android")]
6499static AHARDWAREBUFFER_AVAILABLE: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
6500
6501/// Check if Android AHardwareBuffer allocation is available on this system.
6502///
6503/// AHardwareBuffer is part of the Android OS (public NDK ABI since API
6504/// 26) and is essentially always present; this probe catches degraded
6505/// scenarios such as memory pressure or gralloc failures. The result is
6506/// cached after the first call.
6507#[cfg(target_os = "android")]
6508pub fn is_ahardwarebuffer_available() -> bool {
6509 *AHARDWAREBUFFER_AVAILABLE.get_or_init(|| {
6510 // Probe via the same Dma path — on Android this routes through
6511 // AHardwareBufferTensor::new.
6512 Tensor::<u8>::new(&[64], Some(TensorMemory::DmaBuf), None).is_ok()
6513 })
6514}
6515
6516/// Always returns `false` on non-Android platforms.
6517#[cfg(not(target_os = "android"))]
6518pub fn is_ahardwarebuffer_available() -> bool {
6519 false
6520}
6521
6522/// Portable probe for the platform's native zero-copy GPU buffer allocator
6523/// (DMA-BUF on Linux, IOSurface on macOS/iOS, AHardwareBuffer on Android,
6524/// D3D11 textures on Windows). Returns `false` on platforms with no
6525/// equivalent. Use this when writing cross-platform code that cares whether
6526/// the `Dma` tensor variant will work, not which underlying mechanism is used.
6527///
6528/// On Windows the `Dma` variant is image-formatted: `true` here means
6529/// [`Tensor::image`] can serve `TensorMemory::DmaBuf`, not [`Tensor::new`].
6530///
6531/// # Ordering on Windows
6532///
6533/// This probe **creates the process D3D11 device**, because the only honest
6534/// answer to "can this host allocate one" is to try. A host that must supply
6535/// its own device calls `d3d11::use_external_device` **before** this probe or
6536/// any allocation: the device is cached on first creation and installing one
6537/// afterwards is refused.
6538pub fn is_gpu_buffer_available() -> bool {
6539 #[cfg(target_os = "linux")]
6540 {
6541 is_dma_available()
6542 }
6543 #[cfg(any(target_os = "macos", target_os = "ios"))]
6544 {
6545 is_iosurface_available()
6546 }
6547 #[cfg(target_os = "android")]
6548 {
6549 is_ahardwarebuffer_available()
6550 }
6551 // No cached probe: `d3d11::device()` is itself a `OnceLock` over the
6552 // process device, so asking twice costs one atomic load.
6553 #[cfg(target_os = "windows")]
6554 {
6555 crate::d3d11::device().is_ok()
6556 }
6557 #[cfg(not(any(
6558 target_os = "linux",
6559 target_os = "macos",
6560 target_os = "ios",
6561 target_os = "android",
6562 target_os = "windows"
6563 )))]
6564 {
6565 false
6566 }
6567}
6568
6569/// Check if POSIX shared memory allocation is available on this system.
6570///
6571/// Returns `true` on Unix systems (Linux, macOS, BSD) where POSIX shared memory
6572/// is supported. Always returns `false` on non-Unix platforms (Windows).
6573///
6574/// This function caches its result after the first call for efficiency.
6575#[cfg(unix)]
6576static SHM_AVAILABLE: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
6577
6578/// Check if POSIX shared memory allocation is available on this system.
6579#[cfg(unix)]
6580pub fn is_shm_available() -> bool {
6581 *SHM_AVAILABLE.get_or_init(|| Tensor::<u8>::new(&[64], Some(TensorMemory::Shm), None).is_ok())
6582}
6583
6584/// Check if POSIX shared memory allocation is available on this system.
6585///
6586/// Always returns `false` on non-Unix platforms since POSIX SHM is Unix-specific.
6587#[cfg(not(unix))]
6588pub fn is_shm_available() -> bool {
6589 false
6590}
6591
6592#[cfg(test)]
6593mod dtype_tests {
6594 use super::*;
6595
6596 #[test]
6597 fn dtype_size() {
6598 assert_eq!(DType::U8.size(), 1);
6599 assert_eq!(DType::I8.size(), 1);
6600 assert_eq!(DType::U16.size(), 2);
6601 assert_eq!(DType::I16.size(), 2);
6602 assert_eq!(DType::U32.size(), 4);
6603 assert_eq!(DType::I32.size(), 4);
6604 assert_eq!(DType::U64.size(), 8);
6605 assert_eq!(DType::I64.size(), 8);
6606 assert_eq!(DType::F16.size(), 2);
6607 assert_eq!(DType::F32.size(), 4);
6608 assert_eq!(DType::F64.size(), 8);
6609 }
6610
6611 #[test]
6612 fn dtype_name() {
6613 assert_eq!(DType::U8.name(), "u8");
6614 assert_eq!(DType::F16.name(), "f16");
6615 assert_eq!(DType::F32.name(), "f32");
6616 }
6617
6618 #[test]
6619 fn dtype_serde_roundtrip() {
6620 use serde_json;
6621 let dt = DType::F16;
6622 let json = serde_json::to_string(&dt).unwrap();
6623 let back: DType = serde_json::from_str(&json).unwrap();
6624 assert_eq!(dt, back);
6625 }
6626}
6627
6628#[cfg(test)]
6629mod image_tests {
6630 use super::*;
6631
6632 #[test]
6633 fn image_dims_from_shape_inverts_allocation_shape() {
6634 for (format, w, h) in [
6635 (PixelFormat::Rgba, 640usize, 480usize),
6636 (PixelFormat::Rgb, 641, 3),
6637 (PixelFormat::Grey, 32, 32),
6638 (PixelFormat::PlanarRgb, 640, 480),
6639 (PixelFormat::PlanarRgba, 8, 5),
6640 (PixelFormat::Nv12, 640, 481),
6641 (PixelFormat::Nv16, 33, 17),
6642 (PixelFormat::Nv24, 16, 9),
6643 ] {
6644 let shape = format.allocation_shape(w, h).unwrap();
6645 assert_eq!(
6646 image_dims_from_shape(format, &shape).unwrap(),
6647 (w, h),
6648 "{format:?} {w}x{h} via {shape:?}"
6649 );
6650 }
6651 }
6652
6653 #[test]
6654 fn image_dims_from_shape_rejects_the_wrong_rank() {
6655 // Packed and planar want rank 3, semi-planar rank 2.
6656 for (format, shape) in [
6657 (PixelFormat::Rgba, vec![480usize, 640usize]),
6658 (PixelFormat::PlanarRgb, vec![3, 480]),
6659 (PixelFormat::Nv12, vec![720, 640, 1]),
6660 ] {
6661 let err = image_dims_from_shape(format, &shape).unwrap_err();
6662 assert!(
6663 matches!(err, Error::InvalidShape(ref m) if m.contains("rank")),
6664 "{format:?} {shape:?}: {err}"
6665 );
6666 }
6667 }
6668
6669 #[test]
6670 fn image_dims_from_shape_rejects_an_unproduced_combined_height() {
6671 // NV12 maps 480 -> 720 and 481 -> 722; nothing maps to 721, so the
6672 // shape is rejected rather than rounded to 480.
6673 let err = image_dims_from_shape(PixelFormat::Nv12, &[721, 640]).unwrap_err();
6674 assert!(
6675 matches!(err, Error::InvalidShape(ref m) if m.contains("combined plane height")),
6676 "{err}"
6677 );
6678 }
6679
6680 #[test]
6681 fn allocation_shape_per_layout() {
6682 assert_eq!(
6683 PixelFormat::Rgb.allocation_shape(640, 480),
6684 Some(vec![480, 640, 3])
6685 );
6686 assert_eq!(
6687 PixelFormat::Grey.allocation_shape(640, 480),
6688 Some(vec![480, 640, 1])
6689 );
6690 assert_eq!(
6691 PixelFormat::Nv12.allocation_shape(640, 480),
6692 Some(vec![720, 640])
6693 );
6694 // Odd height: combined-plane height is `481 + ceil(481/2)` = 481 + 241
6695 // = 722 rows. Logical height is recovered as `722 * 2 / 3` = 481.
6696 assert_eq!(
6697 PixelFormat::Nv12.allocation_shape(640, 481),
6698 Some(vec![722, 640])
6699 );
6700 // Odd width: shape carries the LOGICAL width (641).
6701 // The 64-aligned stride (>= 642) is stored separately on the Tensor.
6702 assert_eq!(
6703 PixelFormat::Nv12.allocation_shape(641, 480),
6704 Some(vec![720, 641])
6705 );
6706 // NV16 odd width: same — logical width in shape, stride separate.
6707 assert_eq!(
6708 PixelFormat::Nv16.allocation_shape(641, 480),
6709 Some(vec![960, 641])
6710 );
6711 assert_eq!(
6712 PixelFormat::PlanarRgb.allocation_shape(640, 480),
6713 Some(vec![3, 480, 640])
6714 );
6715 assert_eq!(
6716 PixelFormat::Nv16.allocation_shape(640, 480),
6717 Some(vec![960, 640])
6718 );
6719 }
6720
6721 #[test]
6722 fn raw_tensor_has_no_format() {
6723 let t = Tensor::<u8>::new(&[480, 640, 3], None, None).unwrap();
6724 assert!(t.format().is_none());
6725 assert!(t.width().is_none());
6726 assert!(t.height().is_none());
6727 assert!(!t.is_multiplane());
6728 assert!(t.chroma().is_none());
6729 }
6730
6731 #[test]
6732 fn image_tensor_packed() {
6733 let t = Tensor::<u8>::image(
6734 640,
6735 480,
6736 PixelFormat::Rgba,
6737 None,
6738 crate::CpuAccess::ReadWrite,
6739 )
6740 .unwrap();
6741 assert_eq!(t.format(), Some(PixelFormat::Rgba));
6742 assert_eq!(t.width(), Some(640));
6743 assert_eq!(t.height(), Some(480));
6744 assert_eq!(t.shape(), &[480, 640, 4]);
6745 assert!(!t.is_multiplane());
6746 }
6747
6748 #[test]
6749 fn image_tensor_planar() {
6750 let t = Tensor::<u8>::image(
6751 640,
6752 480,
6753 PixelFormat::PlanarRgb,
6754 None,
6755 crate::CpuAccess::ReadWrite,
6756 )
6757 .unwrap();
6758 assert_eq!(t.format(), Some(PixelFormat::PlanarRgb));
6759 assert_eq!(t.width(), Some(640));
6760 assert_eq!(t.height(), Some(480));
6761 assert_eq!(t.shape(), &[3, 480, 640]);
6762 }
6763
6764 #[test]
6765 #[cfg(target_os = "macos")]
6766 fn image_tensor_dma_non_aligned_packed_width_pads_zero_copy() {
6767 // Declares this test as an fd-opener; see FD_LOCK.
6768 let _lock = crate::tests::fd_lock_shared();
6769 // RGBA u8 at width=4 → 4*4 = 16 bytes/row, not 64-byte aligned. RGBA has
6770 // a real IOSurface FourCC, so an explicit `Some(TensorMemory::DmaBuf)`
6771 // request now allocates a padded image IOSurface (64-aligned
6772 // `bytes_per_row`) and records the stride — a fully zero-copy buffer GL
6773 // can bind and the CPU can map via the strided path. (Previously this
6774 // failed loudly to avoid an 'L008' byte-bag downgrade; with a real
6775 // FourCC surface that concern no longer applies.)
6776 let t = Tensor::<u8>::image(
6777 4,
6778 4,
6779 PixelFormat::Rgba,
6780 Some(TensorMemory::DmaBuf),
6781 crate::CpuAccess::ReadWrite,
6782 )
6783 .expect("padded RGBA IOSurface should allocate");
6784 assert_eq!(t.format(), Some(PixelFormat::Rgba));
6785 assert_eq!(t.width(), Some(4));
6786 assert_eq!(t.height(), Some(4));
6787 let stride = t.effective_row_stride().expect("stride");
6788 assert_eq!(stride % 64, 0, "padded to 64-byte row alignment");
6789 assert!(stride >= 16);
6790 // A CPU map exposes the full padded surface for strided iteration.
6791 let m = t.map().expect("strided IOSurface map");
6792 assert_eq!(m.as_slice().len(), stride * 4);
6793 }
6794
6795 /// `per_pixel_bytes` that doesn't divide 64 evenly (e.g. RGB u8 with
6796 /// 3 B/pixel) makes a "Pad width to N" suggestion structurally
6797 /// impossible — there is no integer width whose `width * 3` is a
6798 /// multiple of 64. The error must still fire (no silent SHM
6799 /// fallback for explicit-DMA requests) and must spell out the
6800 /// alignment requirement; it just omits the misleading "pad to N"
6801 /// hint instead of printing a number whose row pitch still won't
6802 /// align.
6803 #[test]
6804 #[cfg(target_os = "macos")]
6805 fn image_tensor_dma_rejects_indivisible_pixel_pitch_without_pad_hint() {
6806 // Declares this test as an fd-opener; see FD_LOCK.
6807 let _lock = crate::tests::fd_lock_shared();
6808 // Width=10 RGB f32 → 120 B/row, not 64-byte aligned, and (Rgb,
6809 // F32) has no IOSurface mapping so the padded-stride tolerance
6810 // does not apply. The next 64-multiple (128 B) isn't an integer
6811 // multiple of 12 B/pixel, so the "pad width to N" hint can't
6812 // produce a valid number and must be omitted. (Rgb u8 used to be
6813 // this test's subject but now has a real RGBA8888 mapping with
6814 // padded-stride tolerance — see the test below.)
6815 let err = Tensor::<f32>::image(
6816 10,
6817 10,
6818 PixelFormat::Rgb,
6819 Some(TensorMemory::DmaBuf),
6820 crate::CpuAccess::ReadWrite,
6821 )
6822 .expect_err("RGB f32 with 12 B/pixel and non-aligned width must be rejected");
6823 match err {
6824 Error::InvalidArgument(msg) => {
6825 assert!(
6826 msg.contains("64-byte aligned"),
6827 "error must still name the alignment requirement: {msg}"
6828 );
6829 assert!(
6830 !msg.contains("Pad width"),
6831 "indivisible per-pixel pitch makes a width suggestion impossible; \
6832 hint must be omitted, got: {msg}"
6833 );
6834 assert!(
6835 msg.contains("memory=None") && msg.contains("TensorMemory::Mem"),
6836 "error must still list the always-applicable alternatives: {msg}"
6837 );
6838 }
6839 other => panic!("expected InvalidArgument, got {other:?}"),
6840 }
6841 }
6842
6843 #[test]
6844 #[cfg(target_os = "macos")]
6845 fn image_tensor_dma_packed_rgb_u8_contract() {
6846 // Declares this test as an fd-opener; see FD_LOCK.
6847 let _lock = crate::tests::fd_lock_shared();
6848 // Packed RGB u8 @Dma is a designed RGBA8888 mapping at
6849 // (W*3/4, H) — the INT8 NPU input layout, shared with Android.
6850 // width%4 != 0 cannot form whole texels → loud InvalidArgument…
6851 let err = Tensor::<u8>::image(
6852 10,
6853 10,
6854 PixelFormat::Rgb,
6855 Some(TensorMemory::DmaBuf),
6856 crate::CpuAccess::ReadWrite,
6857 )
6858 .expect_err("Rgb u8 width%4!=0 must be rejected");
6859 assert!(
6860 matches!(&err, Error::InvalidArgument(m) if m.contains("width%4==0")),
6861 "got {err:?}"
6862 );
6863 // …width%4 == 0 with a non-64-aligned pitch allocates PADDED
6864 // (36 B rows → 64 B surface pitch, recorded on the tensor)…
6865 let t = Tensor::<u8>::image(
6866 12,
6867 4,
6868 PixelFormat::Rgb,
6869 Some(TensorMemory::DmaBuf),
6870 crate::CpuAccess::ReadWrite,
6871 )
6872 .expect("width 12 Rgb u8 must allocate padded");
6873 assert_eq!(t.memory(), TensorMemory::DmaBuf);
6874 assert!(
6875 t.row_stride()
6876 .is_some_and(|s| s >= 64 && s.is_multiple_of(64)),
6877 "padded pitch must be recorded: {:?}",
6878 t.row_stride()
6879 );
6880 // …and the aligned model-input width stays flat (640*3 = 1920 is
6881 // 64-aligned → no recorded stride, the buffer IS [H, W, 3]).
6882 let t = Tensor::<u8>::image(
6883 640,
6884 8,
6885 PixelFormat::Rgb,
6886 Some(TensorMemory::DmaBuf),
6887 crate::CpuAccess::ReadWrite,
6888 )
6889 .expect("width 640 Rgb u8 must allocate flat");
6890 assert_eq!(t.row_stride(), None);
6891 // I8 shares the layout (INT8 shader bias, not a format change).
6892 let t = Tensor::<i8>::image(
6893 640,
6894 8,
6895 PixelFormat::Rgb,
6896 Some(TensorMemory::DmaBuf),
6897 crate::CpuAccess::ReadWrite,
6898 )
6899 .expect("Rgb i8 shares the RGBA8888 mapping");
6900 assert_eq!(t.memory(), TensorMemory::DmaBuf);
6901 }
6902
6903 #[test]
6904 #[cfg(target_os = "macos")]
6905 fn image_tensor_dma_planar_f16_alignment() {
6906 // Declares this test as an fd-opener; see FD_LOCK.
6907 let _lock = crate::tests::fd_lock_shared();
6908 // PlanarRgb F16 uses single-channel row pitch (width * 2 bytes).
6909 // Width=16 → 32 bytes/row (not aligned); width=32 → 64 bytes/row (aligned).
6910 let err = Tensor::<half::f16>::image(
6911 16,
6912 16,
6913 PixelFormat::PlanarRgb,
6914 Some(TensorMemory::DmaBuf),
6915 crate::CpuAccess::ReadWrite,
6916 )
6917 .expect_err("width=16 PlanarRgb F16 is 32-byte row, must reject");
6918 assert!(matches!(err, Error::InvalidArgument(_)), "got {err:?}");
6919 // 32 wide should work.
6920 let t = Tensor::<half::f16>::image(
6921 32,
6922 8,
6923 PixelFormat::PlanarRgb,
6924 Some(TensorMemory::DmaBuf),
6925 crate::CpuAccess::ReadWrite,
6926 )
6927 .expect("width=32 PlanarRgb F16 is 64-byte row, must succeed");
6928 assert_eq!(t.format(), Some(PixelFormat::PlanarRgb));
6929 }
6930
6931 #[test]
6932 fn image_tensor_semi_planar_contiguous() {
6933 let t = Tensor::<u8>::image(
6934 640,
6935 480,
6936 PixelFormat::Nv12,
6937 None,
6938 crate::CpuAccess::ReadWrite,
6939 )
6940 .unwrap();
6941 assert_eq!(t.format(), Some(PixelFormat::Nv12));
6942 assert_eq!(t.width(), Some(640));
6943 assert_eq!(t.height(), Some(480));
6944 // NV12: H*3/2 = 720
6945 assert_eq!(t.shape(), &[720, 640]);
6946 assert!(!t.is_multiplane());
6947 }
6948
6949 #[test]
6950 #[cfg(target_os = "linux")]
6951 fn image_tensor_with_stride_preserves_logical_width() {
6952 // Declares this test as an fd-opener; see FD_LOCK.
6953 let _lock = crate::tests::fd_lock_shared();
6954 if !crate::test_support::dma_or_skip("image_tensor_with_stride_preserves_logical_width") {
6955 return;
6956 }
6957 // 3004×1688 RGBA8: natural pitch 12016, padded to 12032 (64-aligned).
6958 let stride = 12032;
6959 let t = Tensor::<u8>::image_with_stride(
6960 3004,
6961 1688,
6962 PixelFormat::Rgba,
6963 stride,
6964 Some(TensorMemory::DmaBuf),
6965 crate::CpuAccess::ReadWrite,
6966 )
6967 .unwrap();
6968 // Logical dimensions unchanged by padding — this is the contract.
6969 assert_eq!(t.width(), Some(3004));
6970 assert_eq!(t.height(), Some(1688));
6971 assert_eq!(t.shape(), &[1688, 3004, 4]);
6972 // Stride is carried separately and reports the padded pitch.
6973 assert_eq!(t.effective_row_stride(), Some(stride));
6974 // Buffer is sized to stride × height so the full padded layout fits,
6975 // and CPU map() works for self-allocated strided DMA tensors.
6976 use crate::TensorMapTrait;
6977 {
6978 let map = t.map().unwrap();
6979 assert!(
6980 map.as_slice().len() >= stride * 1688,
6981 "mapped buffer {} bytes < expected {}",
6982 map.as_slice().len(),
6983 stride * 1688
6984 );
6985 }
6986 // CPU write access works too — iterate rows using the padded stride,
6987 // touch only the active `width × bpp` region, verify it round-trips.
6988 {
6989 let mut map = t.map().unwrap();
6990 let slice = map.as_mut_slice();
6991 for y in 0..1688 {
6992 let row_start = y * stride;
6993 for x in 0..3004 {
6994 let p = row_start + x * 4;
6995 slice[p] = (y & 0xFF) as u8;
6996 slice[p + 1] = (x & 0xFF) as u8;
6997 slice[p + 2] = 0x42;
6998 slice[p + 3] = 0xFF;
6999 }
7000 }
7001 }
7002 {
7003 let map = t.map().unwrap();
7004 let slice = map.as_slice();
7005 // Sample a few pixels to confirm the round-trip.
7006 assert_eq!(slice[0], 0x00);
7007 assert_eq!(slice[1], 0x00);
7008 assert_eq!(slice[2], 0x42);
7009 assert_eq!(slice[3], 0xFF);
7010 let mid = 100 * stride + 50 * 4;
7011 assert_eq!(slice[mid], 100);
7012 assert_eq!(slice[mid + 1], 50);
7013 assert_eq!(slice[mid + 2], 0x42);
7014 }
7015 }
7016
7017 #[test]
7018 #[cfg(target_os = "linux")]
7019 fn image_tensor_foreign_strided_map_is_bounds_checked_not_refused() {
7020 // An imported (foreign) DMA-BUF with a row stride used to be refused
7021 // outright. That changed deliberately: the cross-package capsule
7022 // protocol reconstructs an imported alias of a HAL-owned Dma tensor,
7023 // and refusing every strided CPU map broke the flagship
7024 // create_image() -> decode_file_into() -> convert() pipeline on Linux.
7025 //
7026 // What replaced the refusal is a bounds check against `buf_size`,
7027 // which for an imported fd is `fstat`-derived kernel truth rather
7028 // than the caller's claim. So the property worth testing is no longer
7029 // "foreign strided maps fail" but "they succeed exactly when they
7030 // fit" -- which is what actually keeps the mapping sound.
7031 //
7032 // This test previously asserted the old refusal and went red on every
7033 // board with a dma-heap when the behaviour changed, unnoticed because
7034 // it needs both Linux and a dma-heap: it skips on macOS and on Orin
7035 // (nvmap, no CONFIG_DMABUF_HEAPS).
7036 if !crate::test_support::dma_or_skip(
7037 "image_tensor_foreign_strided_map_is_bounds_checked_not_refused",
7038 ) {
7039 return;
7040 }
7041 let backing =
7042 Tensor::<u8>::new(&[240 * 320 * 4], Some(TensorMemory::DmaBuf), None).unwrap();
7043
7044 // In bounds: the natural stride for the declared image fits the
7045 // allocation, so the map must succeed.
7046 {
7047 let fd = backing.clone_fd().unwrap();
7048 let shape = [240usize, 320, 4];
7049 let storage = TensorStorage::<u8>::from_fd(fd, &shape, None).unwrap();
7050 let mut t = Tensor::<u8>::wrap(storage);
7051 t.set_format(PixelFormat::Bgra).unwrap();
7052 t.set_row_stride(320 * 4).unwrap();
7053 assert!(
7054 t.map().is_ok(),
7055 "an in-bounds strided map of an imported DMA-BUF must be \
7056 accepted -- refusing it is what broke decode-into-GPU-tensor"
7057 );
7058 }
7059
7060 // Out of bounds: a stride the allocation cannot back must still be
7061 // refused, and the error must say so. This is the half that carries
7062 // the soundness.
7063 {
7064 let fd = backing.clone_fd().unwrap();
7065 let shape = [240usize, 320, 4];
7066 let storage = TensorStorage::<u8>::from_fd(fd, &shape, None).unwrap();
7067 let mut t = Tensor::<u8>::wrap(storage);
7068 t.set_format(PixelFormat::Bgra).unwrap();
7069 // Double the pitch: 240 rows x 2560 bytes is twice the allocation.
7070 t.set_row_stride(320 * 4 * 2).unwrap();
7071 let err = t.map();
7072 assert!(
7073 matches!(err, Err(Error::InvalidOperation(_))),
7074 "a strided map larger than the DMA allocation must be refused"
7075 );
7076 }
7077 }
7078
7079 #[test]
7080 #[cfg(target_os = "linux")]
7081 fn image_tensor_with_stride_map_rejects_tampered_stride() {
7082 // Declares this test as an fd-opener; see FD_LOCK.
7083 let _lock = crate::tests::fd_lock_shared();
7084 // Round-3 PR feedback (C1): `set_row_stride` is public and only
7085 // validates `stride >= min_stride`, not that the new stride × height
7086 // fits the underlying buffer. A caller that tampers with the stride
7087 // after allocation must not be able to coerce `Tensor::map()` into
7088 // returning a slice larger than the backing mmap (that would be UB
7089 // in `DmaMap::as_slice`).
7090 if !crate::test_support::dma_or_skip("image_tensor_with_stride_map_rejects_tampered_stride")
7091 {
7092 return;
7093 }
7094 // Allocate a 640×480 RGBA8 padded canvas (stride = 3072 = 768 px).
7095 // Backing buffer is 3072 × 480 = 1,474,560 bytes.
7096 let mut t = Tensor::<u8>::image_with_stride(
7097 640,
7098 480,
7099 PixelFormat::Rgba,
7100 3072,
7101 Some(TensorMemory::DmaBuf),
7102 crate::CpuAccess::ReadWrite,
7103 )
7104 .unwrap();
7105 // Tamper: push the stride up to 4 × the original. This is >=
7106 // min_stride (2560), so `set_row_stride` accepts it.
7107 t.set_row_stride(12288).unwrap();
7108 // Map must now refuse — 12288 × 480 = 5,898,240 > 1,474,560.
7109 let err = t.map();
7110 assert!(
7111 matches!(err, Err(Error::InvalidOperation(_))),
7112 "map() with oversized stride must return InvalidOperation"
7113 );
7114 }
7115
7116 #[test]
7117 fn dma_tensor_new_with_byte_size_rejects_shape_overflow() {
7118 // Round-3 PR feedback (C3): shape.product() * sizeof(T) must use
7119 // checked arithmetic so a pathological shape can't wrap usize and
7120 // make the byte_size-vs-logical-size comparison incorrect.
7121 //
7122 // This test only exercises the overflow rejection path, which is
7123 // pure-Rust and doesn't touch dma_heap — safe to run on any target.
7124 #[cfg(target_os = "linux")]
7125 {
7126 let err = crate::dma::DmaTensor::<u64>::new_with_byte_size(
7127 &[usize::MAX, 2, 2],
7128 usize::MAX,
7129 None,
7130 crate::dma::HeapPolicy::default(),
7131 );
7132 assert!(
7133 matches!(err, Err(Error::InvalidArgument(_))),
7134 "new_with_byte_size must detect shape.product() overflow"
7135 );
7136 }
7137 }
7138
7139 #[test]
7140 #[cfg(target_os = "linux")]
7141 fn image_tensor_with_stride_rejects_too_small_stride() {
7142 // Declares this test as an fd-opener; see FD_LOCK.
7143 let _lock = crate::tests::fd_lock_shared();
7144 // 640×480 RGBA8 natural pitch = 2560, request 2400 → should error.
7145 let err = Tensor::<u8>::image_with_stride(
7146 640,
7147 480,
7148 PixelFormat::Rgba,
7149 2400,
7150 Some(TensorMemory::DmaBuf),
7151 crate::CpuAccess::ReadWrite,
7152 );
7153 assert!(matches!(err, Err(Error::InvalidArgument(_))));
7154 }
7155
7156 #[test]
7157 #[cfg(target_os = "linux")]
7158 fn image_tensor_with_stride_rejects_non_packed() {
7159 // Declares this test as an fd-opener; see FD_LOCK.
7160 let _lock = crate::tests::fd_lock_shared();
7161 // NV12 is SemiPlanar → not supported. (Linux-only because
7162 // `TensorMemory::DmaBuf` itself is a Linux-only enum variant.)
7163 let err = Tensor::<u8>::image_with_stride(
7164 640,
7165 480,
7166 PixelFormat::Nv12,
7167 640,
7168 Some(TensorMemory::DmaBuf),
7169 crate::CpuAccess::ReadWrite,
7170 );
7171 assert!(matches!(err, Err(Error::NotImplemented(_))));
7172 }
7173
7174 #[test]
7175 fn set_format_valid() {
7176 let mut t = Tensor::<u8>::new(&[480, 640, 3], None, None).unwrap();
7177 assert!(t.format().is_none());
7178 t.set_format(PixelFormat::Rgb).unwrap();
7179 assert_eq!(t.format(), Some(PixelFormat::Rgb));
7180 assert_eq!(t.width(), Some(640));
7181 assert_eq!(t.height(), Some(480));
7182 }
7183
7184 #[test]
7185 fn set_format_invalid_shape() {
7186 let mut t = Tensor::<u8>::new(&[480, 640, 4], None, None).unwrap();
7187 // RGB expects 3 channels, not 4
7188 let err = t.set_format(PixelFormat::Rgb);
7189 assert!(err.is_err());
7190 // Original tensor is unmodified
7191 assert!(t.format().is_none());
7192 }
7193
7194 #[test]
7195 fn reshape_clears_format() {
7196 let mut t = Tensor::<u8>::image(
7197 640,
7198 480,
7199 PixelFormat::Rgba,
7200 None,
7201 crate::CpuAccess::ReadWrite,
7202 )
7203 .unwrap();
7204 assert_eq!(t.format(), Some(PixelFormat::Rgba));
7205 // Reshape to flat — format cleared
7206 t.reshape(&[480 * 640 * 4]).unwrap();
7207 assert!(t.format().is_none());
7208 }
7209
7210 #[test]
7211 fn from_planes_nv12() {
7212 let y = Tensor::<u8>::new(&[480, 640], None, None).unwrap();
7213 let uv = Tensor::<u8>::new(&[240, 640], None, None).unwrap();
7214 let img = Tensor::from_planes(y, uv, PixelFormat::Nv12).unwrap();
7215 assert_eq!(img.format(), Some(PixelFormat::Nv12));
7216 assert!(img.is_multiplane());
7217 assert!(img.chroma().is_some());
7218 assert_eq!(img.width(), Some(640));
7219 assert_eq!(img.height(), Some(480));
7220 }
7221
7222 #[test]
7223 fn from_planes_rejects_non_semiplanar() {
7224 let y = Tensor::<u8>::new(&[480, 640], None, None).unwrap();
7225 let uv = Tensor::<u8>::new(&[240, 640], None, None).unwrap();
7226 let err = Tensor::from_planes(y, uv, PixelFormat::Rgb);
7227 assert!(err.is_err());
7228 }
7229
7230 #[test]
7231 fn reshape_multiplane_errors() {
7232 let y = Tensor::<u8>::new(&[480, 640], None, None).unwrap();
7233 let uv = Tensor::<u8>::new(&[240, 640], None, None).unwrap();
7234 let mut img = Tensor::from_planes(y, uv, PixelFormat::Nv12).unwrap();
7235 let err = img.reshape(&[480 * 640 + 240 * 640]);
7236 assert!(err.is_err());
7237 }
7238}
7239
7240#[cfg(test)]
7241mod compression_tests {
7242 use super::*;
7243
7244 #[test]
7245 fn desc_builder_roundtrips() {
7246 let desc = ImageDesc::new(640, 480, PixelFormat::Rgba, DType::U8)
7247 .with_memory(Some(TensorMemory::Mem))
7248 .with_access(CpuAccess::Read)
7249 .with_compression(Compression::Any);
7250 assert_eq!(desc.width(), 640);
7251 assert_eq!(desc.height(), 480);
7252 assert_eq!(desc.format(), PixelFormat::Rgba);
7253 assert_eq!(desc.dtype(), DType::U8);
7254 assert_eq!(desc.memory(), Some(TensorMemory::Mem));
7255 assert_eq!(desc.access(), CpuAccess::Read);
7256 assert_eq!(desc.compression(), Some(Compression::Any));
7257
7258 // Defaults: auto memory, CPU read/write, no request.
7259 let plain = ImageDesc::new(2, 2, PixelFormat::Grey, DType::U8);
7260 assert_eq!(plain.memory(), None);
7261 assert_eq!(plain.access(), CpuAccess::ReadWrite);
7262 assert_eq!(plain.compression(), None);
7263 }
7264
7265 #[test]
7266 fn desc_dtype_must_match_element_type() {
7267 let desc = ImageDesc::new(4, 4, PixelFormat::Rgba, DType::F32);
7268 match Tensor::<u8>::image_desc(&desc) {
7269 Err(Error::InvalidArgument(msg)) => assert!(msg.contains("dtype")),
7270 other => panic!("expected InvalidArgument, got {other:?}"),
7271 }
7272 }
7273
7274 #[test]
7275 fn compression_with_cpu_access_is_invalid() {
7276 let desc = ImageDesc::new(4, 4, PixelFormat::Rgba, DType::U8)
7277 .with_access(CpuAccess::ReadWrite)
7278 .with_compression(Compression::Any);
7279 match Tensor::<u8>::image_desc(&desc) {
7280 Err(Error::InvalidArgument(msg)) => assert!(msg.contains("CpuAccess::None")),
7281 other => panic!("expected InvalidArgument, got {other:?}"),
7282 }
7283 }
7284
7285 #[cfg(not(target_os = "android"))]
7286 #[test]
7287 fn scheme_request_off_android_is_not_implemented() {
7288 let desc = ImageDesc::new(4, 4, PixelFormat::Rgba, DType::U8)
7289 .with_access(CpuAccess::None)
7290 .with_compression(Compression::Scheme(CompressionScheme::Ubwc));
7291 match Tensor::<u8>::image_desc(&desc) {
7292 Err(Error::NotImplemented(msg)) => assert!(msg.contains("Ubwc")),
7293 other => panic!("expected NotImplemented, got {other:?}"),
7294 }
7295 }
7296
7297 #[cfg(not(target_os = "android"))]
7298 #[test]
7299 fn any_request_off_android_resolves_linear_and_counts() {
7300 let before = compression_fallback_count();
7301 let desc = ImageDesc::new(64, 64, PixelFormat::Rgba, DType::U8)
7302 .with_memory(Some(TensorMemory::Mem))
7303 .with_access(CpuAccess::None)
7304 .with_compression(Compression::Any);
7305 let t = Tensor::<u8>::image_desc(&desc).unwrap();
7306 assert_eq!(t.compression(), None);
7307 assert!(compression_fallback_count() > before);
7308 }
7309
7310 #[test]
7311 fn desc_without_request_matches_classic_constructor() {
7312 let desc = ImageDesc::new(32, 32, PixelFormat::Rgba, DType::U8)
7313 .with_memory(Some(TensorMemory::Mem))
7314 .with_access(CpuAccess::ReadWrite);
7315 let t = Tensor::<u8>::image_desc(&desc).unwrap();
7316 assert_eq!(t.compression(), None);
7317 assert_eq!(t.cpu_access(), CpuAccess::ReadWrite);
7318 assert_eq!(t.width(), Some(32));
7319 // Mappable exactly like the classic constructor's result.
7320 let m = t.map_read().unwrap();
7321 assert_eq!(m.as_slice().len(), 32 * 32 * 4);
7322 }
7323
7324 #[test]
7325 fn configure_image_preserves_compression_and_views_inherit() {
7326 // No host platform records a scheme, so emulate the recording to
7327 // pin the preserve/inherit semantics (the physical layout does
7328 // not change when the logical image is reconfigured).
7329 let mut t = Tensor::<u8>::image(
7330 64,
7331 64,
7332 PixelFormat::Rgba,
7333 Some(TensorMemory::Mem),
7334 CpuAccess::None,
7335 )
7336 .unwrap();
7337 t.compression = Some(CompressionScheme::Ubwc);
7338 t.configure_image(32, 32, PixelFormat::Rgba).unwrap();
7339 assert_eq!(t.compression(), Some(CompressionScheme::Ubwc));
7340 let view = t.subview(0, &[16, 32, 4]).unwrap();
7341 assert_eq!(view.compression(), Some(CompressionScheme::Ubwc));
7342 }
7343
7344 #[test]
7345 fn tensor_dyn_dispatches_desc_and_compression() {
7346 let desc = ImageDesc::new(16, 16, PixelFormat::Rgba, DType::U8)
7347 .with_memory(Some(TensorMemory::Mem))
7348 .with_access(CpuAccess::ReadWrite);
7349 let t = TensorDyn::image_desc(&desc).unwrap();
7350 assert_eq!(t.compression(), None);
7351 assert!(matches!(t, TensorDyn::U8(_)));
7352 }
7353}
7354
7355#[cfg(all(test, not(edgefirst_tensor_dynamic)))]
7356mod vocabulary_and_probe_tests {
7357 use super::*;
7358
7359 #[test]
7360 fn dtype_displays_its_wire_name() {
7361 for (dt, name) in [
7362 (DType::U8, "u8"),
7363 (DType::I8, "i8"),
7364 (DType::U16, "u16"),
7365 (DType::I16, "i16"),
7366 (DType::U32, "u32"),
7367 (DType::I32, "i32"),
7368 (DType::U64, "u64"),
7369 (DType::I64, "i64"),
7370 (DType::F16, "f16"),
7371 (DType::F32, "f32"),
7372 (DType::F64, "f64"),
7373 ] {
7374 assert_eq!(dt.to_string(), name);
7375 }
7376 }
7377
7378 #[test]
7379 fn dtype_of_maps_every_element_type_and_nothing_else() {
7380 assert_eq!(dtype_of::<u8>(), Some(DType::U8));
7381 assert_eq!(dtype_of::<i8>(), Some(DType::I8));
7382 assert_eq!(dtype_of::<u16>(), Some(DType::U16));
7383 assert_eq!(dtype_of::<i16>(), Some(DType::I16));
7384 assert_eq!(dtype_of::<u32>(), Some(DType::U32));
7385 assert_eq!(dtype_of::<i32>(), Some(DType::I32));
7386 assert_eq!(dtype_of::<u64>(), Some(DType::U64));
7387 assert_eq!(dtype_of::<i64>(), Some(DType::I64));
7388 assert_eq!(dtype_of::<half::f16>(), Some(DType::F16));
7389 assert_eq!(dtype_of::<f32>(), Some(DType::F32));
7390 assert_eq!(dtype_of::<f64>(), Some(DType::F64));
7391 assert_eq!(dtype_of::<bool>(), None);
7392 }
7393
7394 #[test]
7395 fn tensor_memory_converts_to_and_from_its_wire_string() {
7396 assert_eq!(String::from(TensorMemory::DmaBuf), "dmabuf");
7397 assert_eq!(String::from(TensorMemory::Shm), "shm");
7398 assert_eq!(
7399 TensorMemory::try_from("dmabuf").unwrap(),
7400 TensorMemory::DmaBuf
7401 );
7402 assert_eq!(
7403 TensorMemory::try_from("dma").unwrap(),
7404 TensorMemory::DmaBuf,
7405 "the legacy spelling still parses"
7406 );
7407 assert_eq!(TensorMemory::try_from("mem").unwrap(), TensorMemory::Mem);
7408 match TensorMemory::try_from("bogus") {
7409 Err(Error::InvalidMemoryType(s)) => assert_eq!(s, "bogus"),
7410 other => panic!("expected InvalidMemoryType, got {other:?}"),
7411 }
7412 }
7413
7414 #[test]
7415 fn tensor_memory_is_available_matches_a_real_allocation() {
7416 let _lock = crate::tests::fd_lock_shared();
7417 assert!(TensorMemory::Mem.is_available());
7418 assert!(!TensorMemory::IoSurface.is_available());
7419 assert!(!TensorMemory::Pbo.is_available());
7420 assert!(!TensorMemory::Cuda.is_available());
7421 assert!(!TensorMemory::D3d11Texture.is_available());
7422 #[cfg(unix)]
7423 assert_eq!(
7424 TensorMemory::Shm.is_available(),
7425 Tensor::<u8>::new(&[64], Some(TensorMemory::Shm), None).is_ok()
7426 );
7427 #[cfg(not(target_os = "windows"))]
7428 assert_eq!(
7429 TensorMemory::DmaBuf.is_available(),
7430 Tensor::<u8>::new(&[64], Some(TensorMemory::DmaBuf), None).is_ok()
7431 );
7432 }
7433
7434 #[test]
7435 #[cfg(target_os = "linux")]
7436 fn linux_probes_report_what_linux_can_allocate() {
7437 let _lock = crate::tests::fd_lock_shared();
7438 let dma = DmaTensor::<u8>::new(&[64], None).is_ok();
7439 assert_eq!(is_dma_available(), dma);
7440 assert_eq!(is_gpu_buffer_available(), dma);
7441 assert!(!is_iosurface_available());
7442 assert!(!is_ahardwarebuffer_available());
7443 // Linux (unlike Android's bionic) always has POSIX shared memory.
7444 assert!(is_shm_available());
7445 }
7446}
7447
7448#[cfg(all(test, not(edgefirst_tensor_dynamic)))]
7449mod quantization_serde_tests {
7450 use super::*;
7451
7452 fn parse(json: &str) -> Quantization {
7453 serde_json::from_str(json).unwrap_or_else(|e| panic!("{json}: {e}"))
7454 }
7455
7456 #[test]
7457 fn scale_accepts_every_json_number_form() {
7458 assert_eq!(parse(r#"{"scale": 0.5}"#).scale(), &[0.5]);
7459 assert_eq!(parse(r#"{"scale": 2}"#).scale(), &[2.0]);
7460 assert_eq!(parse(r#"{"scale": -2}"#).scale(), &[-2.0]);
7461 assert_eq!(parse(r#"{"scale": [0.25, 4]}"#).scale(), &[0.25, 4.0]);
7462 }
7463
7464 #[test]
7465 fn zero_point_accepts_null_scalar_and_array() {
7466 assert_eq!(
7467 parse(r#"{"scale": 0.1, "zero_point": -5}"#).zero_point(),
7468 Some(&[-5][..])
7469 );
7470 assert_eq!(
7471 parse(r#"{"scale": 0.1, "zero_point": 3}"#).zero_point(),
7472 Some(&[3][..])
7473 );
7474 assert_eq!(
7475 parse(r#"{"scale": [0.1, 0.2], "zero_point": [1, -2], "axis": 0}"#).zero_point(),
7476 Some(&[1, -2][..])
7477 );
7478 assert_eq!(
7479 parse(r#"{"scale": 0.1, "zero_point": null}"#).zero_point(),
7480 None
7481 );
7482 assert_eq!(parse(r#"{"scale": 0.1}"#).zero_point(), None);
7483 }
7484
7485 /// Deserializers whose `deserialize_option` forwards to
7486 /// `deserialize_any` (serde's value deserializers among them) reach the
7487 /// zero-point visitor's scalar and sequence arms directly, not through
7488 /// `visit_some`.
7489 #[test]
7490 fn zero_point_visitor_accepts_values_without_an_option_wrapper() {
7491 use serde::de::value::{
7492 Error as ValueError, I64Deserializer, SeqDeserializer, U64Deserializer,
7493 };
7494 assert_eq!(
7495 deserialize_opt_scalar_or_vec_i32(I64Deserializer::<ValueError>::new(-7)).unwrap(),
7496 Some(vec![-7])
7497 );
7498 assert_eq!(
7499 deserialize_opt_scalar_or_vec_i32(U64Deserializer::<ValueError>::new(9)).unwrap(),
7500 Some(vec![9])
7501 );
7502 let seq = SeqDeserializer::<_, ValueError>::new(vec![4i32, -4].into_iter());
7503 assert_eq!(
7504 deserialize_opt_scalar_or_vec_i32(seq).unwrap(),
7505 Some(vec![4, -4])
7506 );
7507 }
7508
7509 #[test]
7510 fn per_channel_is_not_per_tensor() {
7511 let q = Quantization::per_channel(vec![0.1, 0.2], vec![0, 1], 0).unwrap();
7512 assert!(!q.is_per_tensor());
7513 assert!(Quantization::per_tensor(0.1, 0).is_per_tensor());
7514 }
7515
7516 #[test]
7517 fn validate_accepts_matching_per_channel_and_names_the_failing_field() {
7518 let mut t = Tensor::<i8>::new(&[1, 1, 4], Some(TensorMemory::Mem), None).unwrap();
7519 let q = Quantization::per_channel(vec![0.1, 0.2, 0.3, 0.4], vec![0, 1, 2, 3], 2).unwrap();
7520 t.set_quantization(q)
7521 .expect("four channels on axis 2 of [1, 1, 4]");
7522 assert_eq!(
7523 t.quantization().and_then(Quantization::zero_point),
7524 Some(&[0, 1, 2, 3][..])
7525 );
7526
7527 let q = Quantization::per_channel(vec![0.1, 0.2, 0.3, 0.4], vec![0, 1, 2, 3], 3).unwrap();
7528 match t.set_quantization(q) {
7529 Err(Error::QuantizationInvalid { field, got, .. }) => {
7530 assert_eq!(field, "axis");
7531 assert_eq!(got, "axis=3");
7532 }
7533 other => panic!("axis 3 of a rank-3 tensor must be refused, got {other:?}"),
7534 }
7535
7536 let q = Quantization::per_channel_symmetric(vec![0.1, 0.2, 0.3, 0.4], 1).unwrap();
7537 match t.set_quantization(q) {
7538 Err(Error::QuantizationInvalid { field, got, .. }) => {
7539 assert_eq!(field, "scale.len");
7540 assert_eq!(got, "length 4");
7541 }
7542 other => panic!("four scales on a length-1 axis must be refused, got {other:?}"),
7543 }
7544 }
7545}
7546
7547#[cfg(all(test, not(edgefirst_tensor_dynamic)))]
7548mod image_geometry_tests {
7549 use super::*;
7550
7551 fn mem_image<T>(width: usize, height: usize, format: PixelFormat) -> Tensor<T>
7552 where
7553 T: Num + Clone + fmt::Debug + Send + Sync + 'static,
7554 {
7555 Tensor::<T>::image(
7556 width,
7557 height,
7558 format,
7559 Some(TensorMemory::Mem),
7560 CpuAccess::ReadWrite,
7561 )
7562 .unwrap_or_else(|e| panic!("Mem {format:?} {width}x{height}: {e}"))
7563 }
7564
7565 fn raw_mem<T>(shape: &[usize]) -> Tensor<T>
7566 where
7567 T: Num + Clone + fmt::Debug + Send + Sync + 'static,
7568 {
7569 Tensor::<T>::new(shape, Some(TensorMemory::Mem), None).unwrap()
7570 }
7571
7572 #[test]
7573 fn semi_planar_image_stride_scales_with_the_element_size() {
7574 // u16 NV12 at width 100: 200 B rows, aligned to 256.
7575 let t = mem_image::<u16>(100, 64, PixelFormat::Nv12);
7576 assert_eq!(t.row_stride(), Some(256));
7577 assert_eq!(t.capacity_bytes(), 256 * 96);
7578 }
7579
7580 #[test]
7581 fn host_planar_image_is_tight_at_every_element_size() {
7582 let t = mem_image::<u8>(320, 240, PixelFormat::PlanarRgb);
7583 assert_eq!(t.row_stride(), None);
7584 assert_eq!(t.effective_row_stride(), Some(320));
7585 assert_eq!(t.capacity_bytes(), 320 * 240 * 3);
7586
7587 let t = mem_image::<f32>(100, 40, PixelFormat::PlanarRgb);
7588 assert_eq!(t.row_stride(), None);
7589 assert_eq!(t.effective_row_stride(), Some(400));
7590 assert_eq!(t.capacity_bytes(), 400 * 120);
7591 }
7592
7593 #[test]
7594 #[cfg(target_os = "linux")]
7595 fn dma_packed_image_records_the_aligned_pitch_only_when_padded() {
7596 let _lock = crate::tests::fd_lock_shared();
7597 if !crate::test_support::dma_or_skip(
7598 "dma_packed_image_records_the_aligned_pitch_only_when_padded",
7599 ) {
7600 return;
7601 }
7602 let t = Tensor::<u8>::image(
7603 100,
7604 10,
7605 PixelFormat::Rgba,
7606 Some(TensorMemory::DmaBuf),
7607 CpuAccess::ReadWrite,
7608 )
7609 .unwrap();
7610 assert_eq!(t.row_stride(), Some(448), "400 B rows pad to 448");
7611 let t = Tensor::<u8>::image(
7612 64,
7613 10,
7614 PixelFormat::Rgba,
7615 Some(TensorMemory::DmaBuf),
7616 CpuAccess::ReadWrite,
7617 )
7618 .unwrap();
7619 assert_eq!(t.row_stride(), None, "256 B rows are already aligned");
7620 }
7621
7622 #[test]
7623 #[cfg(target_os = "linux")]
7624 fn image_with_stride_validates_the_stride_before_the_memory() {
7625 match Tensor::<u8>::image_with_stride(
7626 640,
7627 480,
7628 PixelFormat::Rgba,
7629 2559,
7630 Some(TensorMemory::Mem),
7631 CpuAccess::ReadWrite,
7632 ) {
7633 Err(Error::InvalidArgument(msg)) => {
7634 assert!(msg.contains("2559 < minimum 2560"), "{msg}")
7635 }
7636 other => panic!("expected InvalidArgument, got {other:?}"),
7637 }
7638 match Tensor::<u8>::image_with_stride(
7639 640,
7640 480,
7641 PixelFormat::Rgba,
7642 2560,
7643 Some(TensorMemory::Mem),
7644 CpuAccess::ReadWrite,
7645 ) {
7646 Err(Error::NotImplemented(msg)) => assert!(msg.contains("only TensorMemory::DmaBuf")),
7647 other => panic!("a valid stride on Mem must reach the memory check, got {other:?}"),
7648 }
7649 }
7650
7651 #[test]
7652 fn set_format_checks_the_planar_and_semi_planar_shape() {
7653 let mut t = raw_mem::<u8>(&[480, 640, 3]);
7654 match t.set_format(PixelFormat::PlanarRgb) {
7655 Err(Error::InvalidShape(msg)) => assert!(msg.contains("[3, H, W]"), "{msg}"),
7656 other => panic!("[480, 640, 3] is not planar, got {other:?}"),
7657 }
7658 assert!(raw_mem::<u8>(&[3, 480, 640])
7659 .set_format(PixelFormat::PlanarRgb)
7660 .is_ok());
7661
7662 for (format, bad, good) in [
7663 (PixelFormat::Nv12, 721, [720usize, 722]),
7664 (PixelFormat::Nv16, 481, [480, 482]),
7665 (PixelFormat::Nv24, 1441, [1440, 1443]),
7666 ] {
7667 match raw_mem::<u8>(&[bad, 640]).set_format(format) {
7668 Err(Error::InvalidShape(msg)) => {
7669 assert!(msg.contains(&bad.to_string()), "{format:?}: {msg}")
7670 }
7671 other => panic!("{format:?} [{bad}, 640] must be refused, got {other:?}"),
7672 }
7673 for rows in good {
7674 raw_mem::<u8>(&[rows, 640])
7675 .set_format(format)
7676 .unwrap_or_else(|e| panic!("{format:?} [{rows}, 640]: {e}"));
7677 }
7678 }
7679 }
7680
7681 #[test]
7682 fn contiguous_semi_planar_height_divides_the_combined_rows() {
7683 assert_eq!(
7684 mem_image::<u8>(320, 240, PixelFormat::Nv24).height(),
7685 Some(240)
7686 );
7687 assert_eq!(
7688 mem_image::<u8>(320, 240, PixelFormat::Nv16).height(),
7689 Some(240)
7690 );
7691 assert_eq!(
7692 mem_image::<u8>(320, 241, PixelFormat::Nv12).height(),
7693 Some(241)
7694 );
7695 }
7696
7697 #[test]
7698 fn effective_row_stride_scales_with_the_element_size() {
7699 assert_eq!(
7700 mem_image::<f32>(10, 4, PixelFormat::Rgb).effective_row_stride(),
7701 Some(120)
7702 );
7703 assert_eq!(
7704 mem_image::<f32>(10, 4, PixelFormat::PlanarRgb).effective_row_stride(),
7705 Some(40)
7706 );
7707 let mut t = raw_mem::<u16>(&[360, 320]);
7708 t.set_format(PixelFormat::Nv12).unwrap();
7709 assert_eq!(t.effective_row_stride(), Some(640));
7710 let mut t = raw_mem::<u16>(&[360, 321]);
7711 t.set_format(PixelFormat::Nv12).unwrap();
7712 assert_eq!(
7713 t.effective_row_stride(),
7714 Some(644),
7715 "odd width rounds up to even"
7716 );
7717 }
7718
7719 #[test]
7720 fn set_row_stride_minimum_scales_with_the_element_size() {
7721 let refuse = |mut t: Tensor<f32>, stride: usize, min: usize| match t.set_row_stride(stride)
7722 {
7723 Err(Error::InvalidArgument(msg)) => {
7724 assert!(msg.contains(&format!("{stride} < minimum {min}")), "{msg}")
7725 }
7726 other => panic!("stride {stride} under {min} must be refused, got {other:?}"),
7727 };
7728 refuse(mem_image::<f32>(10, 4, PixelFormat::Rgb), 119, 120);
7729 refuse(mem_image::<f32>(10, 4, PixelFormat::PlanarRgb), 39, 40);
7730 let mut t = mem_image::<f32>(10, 4, PixelFormat::Rgb);
7731 t.set_row_stride(120).unwrap();
7732 assert_eq!(t.row_stride(), Some(120));
7733 let mut t = mem_image::<f32>(10, 4, PixelFormat::PlanarRgb);
7734 t.set_row_stride(40).unwrap();
7735 assert_eq!(t.row_stride(), Some(40));
7736
7737 let mut t = mem_image::<u16>(101, 64, PixelFormat::Nv12);
7738 match t.set_row_stride(203) {
7739 Err(Error::InvalidArgument(msg)) => assert!(msg.contains("203 < minimum 204"), "{msg}"),
7740 other => panic!("expected InvalidArgument, got {other:?}"),
7741 }
7742 t.set_row_stride(204).unwrap();
7743 }
7744
7745 #[test]
7746 fn configure_image_semi_planar_stride_scales_with_the_element_size() {
7747 let mut t = mem_image::<u16>(640, 480, PixelFormat::Rgba);
7748 t.configure_image(100, 64, PixelFormat::Nv12).unwrap();
7749 assert_eq!(t.row_stride(), Some(256));
7750 }
7751
7752 #[test]
7753 fn configure_image_keeps_host_packed_odd_widths_tight() {
7754 let mut t = mem_image::<u8>(640, 480, PixelFormat::Rgb);
7755 t.configure_image(321, 240, PixelFormat::Rgb).unwrap();
7756 assert_eq!(t.row_stride(), None);
7757 assert_eq!(t.effective_row_stride(), Some(963));
7758 }
7759
7760 #[test]
7761 fn configure_image_does_not_inherit_a_stride_under_the_same_format() {
7762 // Host packed: a stride set for one geometry must not survive into
7763 // the next, where it is below the new minimum row.
7764 let mut t = mem_image::<u8>(640, 480, PixelFormat::Rgb);
7765 t.configure_image(100, 100, PixelFormat::Rgb).unwrap();
7766 t.set_row_stride(400).unwrap();
7767 t.configure_image(300, 100, PixelFormat::Rgb).unwrap();
7768 assert_eq!(t.row_stride(), None);
7769 assert_eq!(t.effective_row_stride(), Some(900));
7770
7771 // Semi-planar whose aligned stride does not fit: the tight row does,
7772 // and must be what the tensor reports -- not the old 64.
7773 let mut t = mem_image::<u8>(64, 1000, PixelFormat::Nv12);
7774 assert_eq!(t.row_stride(), Some(64));
7775 t.configure_image(200, 300, PixelFormat::Nv12).unwrap();
7776 assert_eq!(t.row_stride(), None);
7777 assert_eq!(t.effective_row_stride(), Some(200));
7778 }
7779
7780 #[test]
7781 fn configure_image_pool_reuse_picks_the_stride_that_fits() {
7782 // A prior stride narrower than the new row is replaced by the
7783 // aligned one.
7784 let mut t = mem_image::<u8>(64, 1024, PixelFormat::Nv12);
7785 t.configure_image(200, 100, PixelFormat::Nv12).unwrap();
7786 assert_eq!(t.row_stride(), Some(256));
7787
7788 // An unaligned prior stride is replaced too.
7789 let mut t = mem_image::<u8>(64, 64, PixelFormat::Nv12);
7790 t.set_row_stride(100).unwrap();
7791 t.configure_image(64, 64, PixelFormat::Nv12).unwrap();
7792 assert_eq!(t.row_stride(), Some(64));
7793
7794 // A wide prior stride whose rows would overrun the allocation gives
7795 // way to the aligned stride of the new width.
7796 let mut t = mem_image::<u8>(640, 480, PixelFormat::Nv12);
7797 t.configure_image(320, 900, PixelFormat::Nv12).unwrap();
7798 assert_eq!(t.row_stride(), Some(320));
7799
7800 // Neither the aligned stride fits nor is it needed: the tight row
7801 // exactly fills a raw allocation.
7802 let mut t = raw_mem::<u8>(&[150 * 100]);
7803 t.configure_image(100, 100, PixelFormat::Nv12).unwrap();
7804 assert_eq!(t.row_stride(), None);
7805 assert_eq!(t.effective_row_stride(), Some(100));
7806 }
7807
7808 #[test]
7809 fn configure_image_refuses_rows_that_overrun_the_allocation() {
7810 // Odd width 99: the shape [150, 99] fits 14850 bytes, but NV12 rows
7811 // are at least the even width, 100 * 150 = 15000 bytes.
7812 let mut t = raw_mem::<u8>(&[150 * 99]);
7813 match t.configure_image(99, 100, PixelFormat::Nv12) {
7814 Err(Error::InsufficientCapacity { needed, capacity }) => {
7815 assert_eq!((needed, capacity), (15000, 14850))
7816 }
7817 other => panic!("expected InsufficientCapacity, got {other:?}"),
7818 }
7819 }
7820
7821 #[test]
7822 #[cfg(target_os = "linux")]
7823 fn configure_image_on_dma_aligns_packed_and_planar_rows() {
7824 let _lock = crate::tests::fd_lock_shared();
7825 if !crate::test_support::dma_or_skip("configure_image_on_dma_aligns_packed_and_planar_rows")
7826 {
7827 return;
7828 }
7829 // Packed u16 Rgb at width 100: 600 B rows align to 640.
7830 let mut t = Tensor::<u16>::image(
7831 640,
7832 480,
7833 PixelFormat::Rgba,
7834 Some(TensorMemory::DmaBuf),
7835 CpuAccess::ReadWrite,
7836 )
7837 .unwrap();
7838 t.configure_image(100, 50, PixelFormat::Rgb).unwrap();
7839 assert_eq!(t.row_stride(), Some(640));
7840
7841 // Planar u16 at width 100: 200 B rows would align to 256, but 150
7842 // rows of 256 B overrun the 32 KiB buffer while 150 rows of 200 B fit.
7843 let mut t = Tensor::<u16>::new(&[16 * 1024], Some(TensorMemory::DmaBuf), None).unwrap();
7844 assert_eq!(t.capacity_bytes(), 32 * 1024);
7845 t.configure_image(100, 50, PixelFormat::PlanarRgb).unwrap();
7846 assert_eq!(t.row_stride(), None);
7847 assert_eq!(t.effective_row_stride(), Some(200));
7848 t.configure_image(100, 40, PixelFormat::PlanarRgb).unwrap();
7849 assert_eq!(t.row_stride(), Some(256), "120 rows of 256 B fit");
7850 }
7851
7852 #[test]
7853 fn from_planes_requires_a_raw_two_dimensional_chroma() {
7854 let formatted_chroma = mem_image::<u8>(640, 240, PixelFormat::Grey);
7855 match Tensor::from_planes(
7856 raw_mem::<u8>(&[480, 640]),
7857 formatted_chroma,
7858 PixelFormat::Nv12,
7859 ) {
7860 Err(Error::InvalidArgument(msg)) => assert!(msg.contains("raw tensor"), "{msg}"),
7861 other => panic!("a formatted chroma plane must be refused, got {other:?}"),
7862 }
7863 match Tensor::from_planes(
7864 raw_mem::<u8>(&[480, 640, 1]),
7865 raw_mem::<u8>(&[240, 640]),
7866 PixelFormat::Nv12,
7867 ) {
7868 Err(Error::InvalidArgument(msg)) => assert!(msg.contains("2D shapes"), "{msg}"),
7869 other => panic!("a 3-D luma plane must be refused, got {other:?}"),
7870 }
7871 }
7872
7873 #[test]
7874 fn from_planes_nv16_needs_full_height_chroma() {
7875 let img = Tensor::from_planes(
7876 raw_mem::<u8>(&[480, 640]),
7877 raw_mem::<u8>(&[480, 640]),
7878 PixelFormat::Nv16,
7879 )
7880 .expect("NV16 chroma is full height");
7881 assert_eq!(img.format(), Some(PixelFormat::Nv16));
7882 assert_eq!(img.height(), Some(480));
7883
7884 match Tensor::from_planes(
7885 raw_mem::<u8>(&[480, 640]),
7886 raw_mem::<u8>(&[240, 640]),
7887 PixelFormat::Nv16,
7888 ) {
7889 Err(Error::InvalidArgument(msg)) => {
7890 assert!(msg.contains("NV16 chroma height"), "{msg}")
7891 }
7892 other => panic!("half-height NV16 chroma must be refused, got {other:?}"),
7893 }
7894 }
7895
7896 #[test]
7897 fn chroma_mut_writes_the_chroma_plane() {
7898 let mut img = Tensor::from_planes(
7899 raw_mem::<u8>(&[4, 8]),
7900 raw_mem::<u8>(&[2, 8]),
7901 PixelFormat::Nv12,
7902 )
7903 .unwrap();
7904 let chroma = img.chroma_mut().expect("a multiplane image has chroma");
7905 chroma.map().unwrap().as_mut_slice().fill(0x80);
7906 let chroma = img.chroma().unwrap();
7907 assert!(chroma.map().unwrap().as_slice().iter().all(|&b| b == 0x80));
7908 }
7909
7910 #[test]
7911 fn view_offset_and_stride_scale_with_the_element_size() {
7912 let t = mem_image::<u8>(100, 80, PixelFormat::Rgba);
7913 let v = t.view(Region::new(10, 20, 30, 40)).unwrap();
7914 assert_eq!(v.plane_offset(), Some(20 * 400 + 10 * 4));
7915 assert_eq!(v.row_stride(), Some(400));
7916
7917 let t = mem_image::<u16>(100, 80, PixelFormat::Rgba);
7918 let v = t.view(Region::new(10, 20, 30, 40)).unwrap();
7919 assert_eq!(v.plane_offset(), Some(20 * 800 + 10 * 8));
7920 assert_eq!(v.row_stride(), Some(800));
7921 assert_eq!(v.view_origin().map(|o| o.parent_row_stride), Some(800));
7922 }
7923
7924 #[test]
7925 fn single_row_view_keeps_a_tight_stride() {
7926 let t = mem_image::<u8>(8, 4, PixelFormat::Rgba);
7927 let row = t.view(Region::new(2, 3, 4, 1)).unwrap();
7928 assert_eq!(row.effective_row_stride(), Some(16));
7929 assert_eq!(row.map().unwrap().as_slice().len(), 16);
7930 assert_eq!(row.view_origin().map(|o| o.parent_row_stride), Some(32));
7931
7932 let two_rows = t.view(Region::new(2, 2, 4, 2)).unwrap();
7933 assert_eq!(two_rows.effective_row_stride(), Some(32));
7934 }
7935
7936 #[test]
7937 fn view_refuses_a_region_past_either_edge() {
7938 let t = mem_image::<u8>(64, 40, PixelFormat::Rgba);
7939 for region in [Region::new(0, 30, 10, 20), Region::new(60, 0, 10, 1)] {
7940 match t.view(region) {
7941 Err(Error::RegionOutOfBounds { region: r, bounds }) => {
7942 assert_eq!((r, bounds), (region, (64, 40)))
7943 }
7944 other => panic!("{region:?} must be refused, got {other:?}"),
7945 }
7946 }
7947 assert!(Region::new(0, 0, 64, 40).fits_within(64, 40));
7948 assert!(!Region::new(1, 0, 64, 40).fits_within(64, 40));
7949 assert!(!Region::new(0, 1, 64, 40).fits_within(64, 40));
7950 }
7951}
7952
7953#[cfg(all(test, not(edgefirst_tensor_dynamic)))]
7954mod strided_access_tests {
7955 use super::*;
7956
7957 /// A patterned byte view of `t`'s whole mapping: byte `i` is `i % 251`.
7958 fn fill_pattern<T>(t: &Tensor<T>)
7959 where
7960 T: Num + Clone + fmt::Debug + Send + Sync,
7961 {
7962 let mut m = t.map().unwrap();
7963 let s = m.as_mut_slice();
7964 // SAFETY: plain numeric elements; every byte pattern is a valid value
7965 // for the integer and float types these tests use.
7966 let bytes = unsafe {
7967 std::slice::from_raw_parts_mut(s.as_mut_ptr() as *mut u8, std::mem::size_of_val(s))
7968 };
7969 for (i, b) in bytes.iter_mut().enumerate() {
7970 *b = (i % 251) as u8;
7971 }
7972 }
7973
7974 fn expected_rows(rows: usize, stride: usize, row_bytes: usize) -> Vec<u8> {
7975 (0..rows)
7976 .flat_map(|r| (r * stride..r * stride + row_bytes).map(|i| (i % 251) as u8))
7977 .collect()
7978 }
7979
7980 fn padded<T>(format: PixelFormat, stride: usize) -> Tensor<T>
7981 where
7982 T: Num + Clone + fmt::Debug + Send + Sync + 'static,
7983 {
7984 let mut t = Tensor::<T>::image_with_capacity(
7985 16,
7986 4,
7987 format,
7988 Some(TensorMemory::Mem),
7989 CpuAccess::ReadWrite,
7990 )
7991 .unwrap();
7992 t.configure_image(8, 4, format).unwrap();
7993 t.set_row_stride(stride).unwrap();
7994 t
7995 }
7996
7997 #[test]
7998 fn copy_to_flat_compacts_padded_packed_rows() {
7999 // 8 px * 3 B = 24 B rows at a 32 B pitch.
8000 let t = padded::<u8>(PixelFormat::Rgb, 32);
8001 fill_pattern(&t);
8002 let mut dst = vec![0u8; 96];
8003 t.copy_to_flat(&mut dst).unwrap();
8004 assert_eq!(dst, expected_rows(4, 32, 24));
8005 }
8006
8007 #[test]
8008 fn copy_to_flat_compacts_padded_planar_rows() {
8009 // f32 planar: 3 planes x 4 rows of 32 B at a 64 B pitch.
8010 let t = padded::<f32>(PixelFormat::PlanarRgb, 64);
8011 fill_pattern(&t);
8012 let mut dst = vec![0u8; 384];
8013 t.copy_to_flat(&mut dst).unwrap();
8014 assert_eq!(dst, expected_rows(12, 64, 32));
8015 }
8016
8017 #[test]
8018 fn copy_to_flat_of_a_tight_stride_is_the_mapping() {
8019 let t = padded::<u8>(PixelFormat::Rgb, 24);
8020 fill_pattern(&t);
8021 let mut dst = vec![0u8; 96];
8022 t.copy_to_flat(&mut dst).unwrap();
8023 assert_eq!(dst, expected_rows(1, 96, 96));
8024
8025 let t = Tensor::<u8>::image(
8026 8,
8027 4,
8028 PixelFormat::Rgb,
8029 Some(TensorMemory::Mem),
8030 CpuAccess::ReadWrite,
8031 )
8032 .unwrap();
8033 fill_pattern(&t);
8034 let mut dst = vec![0u8; 96];
8035 t.copy_to_flat(&mut dst).unwrap();
8036 assert_eq!(dst, expected_rows(1, 96, 96));
8037 }
8038
8039 #[test]
8040 fn copy_to_flat_needs_a_tight_sized_destination() {
8041 let t = padded::<u8>(PixelFormat::Rgb, 32);
8042 let mut dst = vec![0u8; 95];
8043 match t.copy_to_flat(&mut dst) {
8044 Err(Error::InvalidArgument(msg)) => {
8045 assert!(
8046 msg.contains("dst is 95 bytes") && msg.contains("96 bytes"),
8047 "{msg}"
8048 )
8049 }
8050 other => panic!("expected InvalidArgument, got {other:?}"),
8051 }
8052 }
8053
8054 #[test]
8055 fn copy_to_flat_compacts_a_padded_batch() {
8056 // Two 8x2 RGB tiles at a 32 B pitch: four rows in all.
8057 let mut t = padded::<u8>(PixelFormat::Rgb, 32);
8058 t.set_logical_shape(&[2, 2, 8, 3]).unwrap();
8059 fill_pattern(&t);
8060 let mut dst = vec![0u8; 96];
8061 t.copy_to_flat(&mut dst).unwrap();
8062 assert_eq!(dst, expected_rows(4, 32, 24));
8063 }
8064
8065 #[test]
8066 fn strided_planar_map_spans_every_plane_row() {
8067 let t = padded::<u8>(PixelFormat::PlanarRgb, 16);
8068 assert_eq!(t.map().unwrap().as_slice().len(), 16 * 12);
8069 }
8070
8071 #[test]
8072 fn batch_of_a_padded_batch_starts_at_its_padded_tile() {
8073 // Tile 1 starts after two 32 B rows, not after 48 tight bytes.
8074 let mut t = padded::<u8>(PixelFormat::Rgb, 32);
8075 t.set_logical_shape(&[2, 2, 8, 3]).unwrap();
8076 fill_pattern(&t);
8077 let e = t.batch(1).unwrap();
8078 assert_eq!(e.plane_offset(), Some(64));
8079 let mut dst = vec![0u8; 48];
8080 e.copy_to_flat(&mut dst).unwrap();
8081 let tile1: Vec<u8> = expected_rows(4, 32, 24).split_off(48);
8082 assert_eq!(dst, tile1);
8083 }
8084
8085 #[test]
8086 fn strided_packed_batch_map_spans_every_tile_row() {
8087 let mut t = padded::<u8>(PixelFormat::Rgb, 32);
8088 t.set_logical_shape(&[2, 2, 8, 3]).unwrap();
8089 assert_eq!(t.map().unwrap().as_slice().len(), 32 * 4);
8090 }
8091
8092 #[test]
8093 fn strided_access_to_a_flattened_image_is_refused() {
8094 let mut t = padded::<u8>(PixelFormat::PlanarRgb, 16);
8095 t.set_logical_shape(&[96]).unwrap();
8096 match t.map() {
8097 Err(Error::InvalidOperation(msg)) => assert!(msg.contains("[96]"), "{msg}"),
8098 other => panic!("expected InvalidOperation, got {other:?}"),
8099 }
8100 let mut dst = vec![0u8; 96];
8101 match t.copy_to_flat(&mut dst) {
8102 Err(Error::InvalidOperation(msg)) => assert!(msg.contains("[96]"), "{msg}"),
8103 other => panic!("expected InvalidOperation, got {other:?}"),
8104 }
8105 }
8106
8107 #[test]
8108 #[cfg(target_os = "linux")]
8109 fn strided_shm_map_is_bounded_by_the_segment() {
8110 let _lock = crate::tests::fd_lock_shared();
8111 let mut t = Tensor::<u8>::image(
8112 640,
8113 480,
8114 PixelFormat::Nv12,
8115 Some(TensorMemory::Shm),
8116 CpuAccess::ReadWrite,
8117 )
8118 .expect("Linux has POSIX shared memory");
8119 assert_eq!(t.capacity_bytes(), 640 * 720);
8120 assert_eq!(t.map().unwrap().as_slice().len(), 640 * 720);
8121
8122 t.configure_image(320, 240, PixelFormat::Nv12).unwrap();
8123 assert_eq!(t.row_stride(), Some(640), "the pool pitch fits and is kept");
8124 assert_eq!(t.map().unwrap().as_slice().len(), 640 * 360);
8125
8126 t.configure_image(640, 480, PixelFormat::Nv12).unwrap();
8127 t.set_row_stride(704).unwrap();
8128 match t.map() {
8129 Err(Error::InsufficientCapacity { needed, capacity }) => {
8130 assert_eq!((needed, capacity), (704 * 720, 640 * 720))
8131 }
8132 other => panic!("expected InsufficientCapacity, got {other:?}"),
8133 }
8134 }
8135
8136 /// Minimal host-memory `PboOps`: a fixed buffer, never resized.
8137 struct HostPbo(Vec<u8>);
8138
8139 // SAFETY: the pointer addresses `self.0`, allocated once and never
8140 // resized, and the tensor holds the `Arc` to `self` for every map's life.
8141 unsafe impl PboOps for HostPbo {
8142 fn map_buffer(&self, _buffer_id: u32, size: usize) -> Result<PboMapping> {
8143 assert!(size <= self.0.len());
8144 Ok(PboMapping {
8145 ptr: self.0.as_ptr() as *mut u8,
8146 size,
8147 })
8148 }
8149 fn unmap_buffer(&self, _buffer_id: u32) -> Result<()> {
8150 Ok(())
8151 }
8152 fn delete_buffer(&self, _buffer_id: u32) {}
8153 }
8154
8155 fn pbo_rgba(width: usize, height: usize) -> Tensor<u8> {
8156 let bytes = width * height * 4;
8157 let pbo = PboTensor::<u8>::from_pbo(
8158 7,
8159 bytes,
8160 &[height, width, 4],
8161 None,
8162 Arc::new(HostPbo(vec![0; bytes])),
8163 )
8164 .unwrap();
8165 let mut t = Tensor::from_pbo(pbo).unwrap();
8166 t.set_format(PixelFormat::Rgba).unwrap();
8167 t
8168 }
8169
8170 #[test]
8171 fn strided_pbo_map_is_bounded_by_the_buffer() {
8172 let mut t = pbo_rgba(8, 4);
8173 t.set_row_stride(32).unwrap();
8174 assert_eq!(t.map().unwrap().as_slice().len(), 128);
8175
8176 t.set_row_stride(64).unwrap();
8177 match t.map() {
8178 Err(Error::InsufficientCapacity { needed, capacity }) => {
8179 assert_eq!((needed, capacity), (256, 128))
8180 }
8181 other => panic!("expected InsufficientCapacity, got {other:?}"),
8182 }
8183 }
8184
8185 #[test]
8186 fn pbo_accessors_answer_only_for_a_pbo_tensor() {
8187 let mem = Tensor::<u8>::new(&[4], Some(TensorMemory::Mem), None).unwrap();
8188 assert_eq!(mem.pbo_is_mapped(), None);
8189 assert!(mem.pbo_keepalive().is_none());
8190
8191 let t = pbo_rgba(8, 4);
8192 assert_eq!(t.pbo_is_mapped(), Some(false));
8193 assert!(t.pbo_keepalive().is_some());
8194 let m = t.map().unwrap();
8195 assert_eq!(t.pbo_is_mapped(), Some(true));
8196 drop(m);
8197 assert_eq!(t.pbo_is_mapped(), Some(false));
8198 }
8199}
8200
8201#[cfg(all(test, not(edgefirst_tensor_dynamic)))]
8202mod accessor_tests {
8203 use super::*;
8204 use std::sync::atomic::{AtomicUsize, Ordering};
8205
8206 #[test]
8207 #[cfg(unix)]
8208 fn plane_descriptor_carries_its_stride_and_offset() {
8209 use std::os::fd::AsFd;
8210 let _lock = crate::tests::fd_lock_shared();
8211 let file = std::fs::File::open("/dev/null").unwrap();
8212 let pd = PlaneDescriptor::new(file.as_fd()).unwrap();
8213 assert_eq!((pd.stride(), pd.offset()), (None, None));
8214 let pd = pd.with_stride(832).with_offset(4096);
8215 assert_eq!((pd.stride(), pd.offset()), (Some(832), Some(4096)));
8216 }
8217
8218 #[test]
8219 #[cfg(target_os = "linux")]
8220 fn as_dma_answers_only_for_a_dma_buffer() {
8221 use std::os::fd::AsRawFd;
8222 let _lock = crate::tests::fd_lock_shared();
8223 let mem = Tensor::<u8>::new(&[64], Some(TensorMemory::Mem), None).unwrap();
8224 assert!(mem.as_dma().is_none());
8225 if !crate::test_support::dma_or_skip("as_dma_answers_only_for_a_dma_buffer") {
8226 return;
8227 }
8228 let t = Tensor::<u8>::new(&[64], Some(TensorMemory::DmaBuf), None).unwrap();
8229 let dma = t.as_dma().expect("a DMA tensor downcasts");
8230 assert_eq!(
8231 dma.fd.as_raw_fd(),
8232 t.dmabuf().unwrap().as_raw_fd(),
8233 "the downcast is this tensor's own buffer"
8234 );
8235 }
8236
8237 #[test]
8238 fn derived_identity_mixes_the_kind_into_the_key() {
8239 let id = BufferIdentity::derived(IdentityKind::Shm, u64::MAX);
8240 assert_eq!(id.id(), ((IdentityKind::Shm as u64) << 56) ^ u64::MAX);
8241 assert_ne!(
8242 id.id(),
8243 BufferIdentity::derived(IdentityKind::DmaBuf, u64::MAX).id()
8244 );
8245 }
8246
8247 #[test]
8248 fn empty_is_a_zero_element_shape() {
8249 let mut t = Tensor::<f32>::new(&[2, 3], Some(TensorMemory::Mem), None).unwrap();
8250 assert!(!TensorTrait::is_empty(&t));
8251 {
8252 let m = t.map().unwrap();
8253 assert_eq!(TensorMapTrait::len(&m), 6);
8254 assert_eq!(TensorMapTrait::size(&m), 24);
8255 assert!(!TensorMapTrait::is_empty(&m));
8256 }
8257 t.set_logical_shape(&[0]).unwrap();
8258 assert!(TensorTrait::is_empty(&t));
8259 let m = t.map().unwrap();
8260 assert_eq!(TensorMapTrait::size(&m), 0);
8261 assert!(TensorMapTrait::is_empty(&m));
8262 }
8263
8264 struct CountingCudaOps {
8265 maps: AtomicUsize,
8266 }
8267
8268 impl CudaGlOps for CountingCudaOps {
8269 fn map(&self, _r: crate::cuda::GraphicsResource) -> Option<(*mut std::ffi::c_void, usize)> {
8270 self.maps.fetch_add(1, Ordering::SeqCst);
8271 Some((0x1000usize as *mut _, 64))
8272 }
8273 fn unmap(&self, _r: crate::cuda::GraphicsResource) {}
8274 fn unregister(&self, _r: crate::cuda::GraphicsResource) {}
8275 }
8276
8277 #[test]
8278 fn cuda_accessors_follow_the_attached_handle() {
8279 let mut t = Tensor::<f32>::new(&[16], Some(TensorMemory::Mem), None).unwrap();
8280 assert!(!t.is_cuda_attached());
8281 assert!(t.cuda().is_none());
8282 assert!(t.cuda_map().is_none());
8283 assert!(t.cuda_map_mut().is_none());
8284
8285 let ops = Arc::new(CountingCudaOps {
8286 maps: AtomicUsize::new(0),
8287 });
8288 t.set_cuda_handle(CudaHandle::new_gl(
8289 0x1usize as crate::cuda::GraphicsResource,
8290 64,
8291 ops.clone(),
8292 ));
8293 assert!(t.is_cuda_attached());
8294 assert!(t.cuda().is_some());
8295 assert_eq!(t.cuda_map().map(|m| m.device_ptr() as usize), Some(0x1000));
8296 assert_eq!(t.cuda_map_mut().map(|m| m.len()), Some(64));
8297 assert_eq!(ops.maps.load(Ordering::SeqCst), 2);
8298 }
8299
8300 #[test]
8301 fn from_slice_copies_values_of_a_matching_shape() {
8302 let t = Tensor::<u16>::from_slice(&[1, 2, 3, 4, 5, 6], &[2, 3]).unwrap();
8303 assert_eq!(t.shape(), &[2, 3]);
8304 assert_eq!(t.map().unwrap().as_slice(), &[1, 2, 3, 4, 5, 6]);
8305 match Tensor::<u16>::from_slice(&[1, 2, 3, 4, 5], &[2, 3]) {
8306 Err(Error::InvalidShape(msg)) => assert!(msg.contains("values.len()=5"), "{msg}"),
8307 other => panic!("expected InvalidShape, got {other:?}"),
8308 }
8309 }
8310
8311 /// An `Any` compression request with no pinned memory tries the
8312 /// platform's zero-copy image first and falls back when it cannot serve
8313 /// the format (packed RGB f32 has no IOSurface layout, for one).
8314 #[cfg(not(target_os = "android"))]
8315 #[test]
8316 fn any_compression_request_falls_back_when_zero_copy_cannot_serve() {
8317 let _lock = crate::tests::fd_lock_shared();
8318 let desc = ImageDesc::new(10, 10, PixelFormat::Rgb, DType::F32)
8319 .with_access(CpuAccess::None)
8320 .with_compression(Compression::Any);
8321 let t = Tensor::<f32>::image_desc(&desc).expect("Any never fails for compression reasons");
8322 assert_eq!(t.shape(), &[10, 10, 3]);
8323 assert_eq!(t.compression(), None);
8324 }
8325}
8326
8327#[cfg(test)]
8328mod cpu_access_tests {
8329 use super::*;
8330
8331 #[test]
8332 fn covers_matrix() {
8333 use CpuAccess::*;
8334 // Every declaration covers a narrower or equal request…
8335 for a in [None, Read, Write, ReadWrite] {
8336 assert!(a.covers(None), "{a:?} must cover None");
8337 assert!(ReadWrite.covers(a), "ReadWrite must cover {a:?}");
8338 }
8339 assert!(Read.covers(Read));
8340 assert!(Write.covers(Write));
8341 // …and never a wider one.
8342 assert!(!None.covers(Read));
8343 assert!(!None.covers(Write));
8344 assert!(!Read.covers(Write));
8345 assert!(!Read.covers(ReadWrite));
8346 assert!(!Write.covers(Read));
8347 assert!(!Write.covers(ReadWrite));
8348 }
8349
8350 #[test]
8351 fn union_covers_both_and_require_names_both() {
8352 use CpuAccess::*;
8353 for a in [None, Read, Write, ReadWrite] {
8354 for b in [None, Read, Write, ReadWrite] {
8355 let u = a.union(b);
8356 assert!(u.covers(a) && u.covers(b), "{a:?} | {b:?} = {u:?}");
8357 assert_eq!(a.require(b).is_ok(), a.covers(b));
8358 }
8359 }
8360 match Read.require(Write) {
8361 Err(Error::CpuAccessDenied {
8362 declared: Read,
8363 requested: Write,
8364 }) => {}
8365 other => panic!("expected CpuAccessDenied, got {other:?}"),
8366 }
8367 }
8368
8369 #[test]
8370 fn map_with_none_is_invalid() {
8371 let t = Tensor::<u8>::new(&[16], Some(TensorMemory::Mem), None).unwrap();
8372 match t.map_with(CpuAccess::None) {
8373 Err(Error::InvalidArgument(_)) => {}
8374 Err(other) => panic!("expected InvalidArgument, got {other:?}"),
8375 Ok(_) => panic!("map_with(CpuAccess::None) must not succeed"),
8376 }
8377 }
8378
8379 #[test]
8380 fn read_map_rejects_mutation_uniformly() {
8381 // Mem backend: map_read yields a working read view whose mutable
8382 // accessor panics (the uniform cross-backend contract).
8383 let t = Tensor::<u8>::new(&[8], Some(TensorMemory::Mem), None).unwrap();
8384 t.map_mut().unwrap().as_mut_slice().copy_from_slice(&[7; 8]);
8385 let ro = t.map_read().unwrap();
8386 assert_eq!(ro.as_slice(), &[7; 8]);
8387 drop(ro);
8388 let result = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
8389 let mut ro = t.map_read().unwrap();
8390 let _ = ro.as_mut_slice();
8391 }));
8392 assert!(result.is_err(), "as_mut_slice through map_read must panic");
8393 }
8394
8395 #[test]
8396 fn write_and_rw_maps_stay_mutable() {
8397 let t = Tensor::<u8>::new(&[4], Some(TensorMemory::Mem), None).unwrap();
8398 t.map_write()
8399 .unwrap()
8400 .as_mut_slice()
8401 .copy_from_slice(&[1; 4]);
8402 t.map().unwrap().as_mut_slice().copy_from_slice(&[2; 4]);
8403 assert_eq!(t.map_read().unwrap().as_slice(), &[2; 4]);
8404 }
8405
8406 /// The read-only IOSurface lock path: a `map_read` must observe data
8407 /// written through a prior read-write lock, and its unlock (which
8408 /// skips the cache flush) must not disturb subsequent reads.
8409 #[test]
8410 #[cfg(target_os = "macos")]
8411 fn iosurface_read_only_lock_roundtrip() {
8412 // Declares this test as an fd-opener; see FD_LOCK.
8413 let _lock = crate::tests::fd_lock_shared();
8414 let expect = if crate::is_iosurface_available() {
8415 crate::test_support::Expect::Succeeds
8416 } else {
8417 crate::test_support::Expect::MayFail("the IOSurface probe failed on this host")
8418 };
8419 let Some(t) = crate::test_support::allocated(
8420 Tensor::<u8>::new(&[64], Some(TensorMemory::DmaBuf), None),
8421 "IOSurface allocation",
8422 expect,
8423 ) else {
8424 return;
8425 };
8426 {
8427 let mut m = t.map_mut().unwrap();
8428 for (i, b) in m.as_mut_slice().iter_mut().enumerate() {
8429 *b = (i * 3) as u8;
8430 }
8431 }
8432 for _ in 0..2 {
8433 let ro = t.map_read().unwrap();
8434 for (i, b) in ro.as_slice().iter().enumerate() {
8435 assert_eq!(*b, (i * 3) as u8, "byte {i} through read-only lock");
8436 }
8437 }
8438 let result = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
8439 let mut ro = t.map_read().unwrap();
8440 let _ = ro.as_mut_slice();
8441 }));
8442 assert!(result.is_err(), "IOSurface read map must reject mutation");
8443 }
8444}
8445
8446#[cfg(test)]
8447mod tests {
8448 use std::sync::{PoisonError, RwLock};
8449
8450 use super::*;
8451
8452 #[ctor::ctor(unsafe)]
8453 fn init() {
8454 env_logger::Builder::from_env(env_logger::Env::default().default_filter_or("info")).init();
8455 }
8456
8457 /// Macro to get the current function name for logging in tests.
8458 #[cfg(target_os = "linux")]
8459 macro_rules! function {
8460 () => {{
8461 fn f() {}
8462 fn type_name_of<T>(_: T) -> &'static str {
8463 std::any::type_name::<T>()
8464 }
8465 let name = type_name_of(f);
8466
8467 // Find and cut the rest of the path
8468 match &name[..name.len() - 3].rfind(':') {
8469 Some(pos) => &name[pos + 1..name.len() - 3],
8470 None => &name[..name.len() - 3],
8471 }
8472 }};
8473 }
8474
8475 #[test]
8476 #[cfg(target_os = "linux")]
8477 fn test_tensor() {
8478 let _lock = crate::tests::fd_lock_shared();
8479 let shape = vec![1];
8480 let tensor = DmaTensor::<f32>::new(&shape, Some("dma_tensor"));
8481 let dma_enabled = tensor.is_ok();
8482
8483 let tensor = Tensor::<f32>::new(&shape, None, None).expect("Failed to create tensor");
8484 // Auto-select priority is Dma > Mem; Shm is never auto-selected.
8485 match dma_enabled {
8486 true => assert_eq!(tensor.memory(), TensorMemory::DmaBuf),
8487 false => assert_eq!(tensor.memory(), TensorMemory::Mem),
8488 }
8489 }
8490
8491 #[test]
8492 #[cfg(any(target_os = "macos", target_os = "ios"))]
8493 fn test_tensor() {
8494 let shape = vec![1];
8495 let tensor = Tensor::<f32>::new(&shape, None, None).expect("Failed to create tensor");
8496 // macOS/iOS auto-fallback chain: IOSurface (Dma) → Mem. Healthy systems
8497 // return Dma; Mem only appears under memory pressure or sandboxed
8498 // contexts where IOSurfaceCreate fails. Shm is never auto-selected.
8499 let m = tensor.memory();
8500 assert!(
8501 matches!(m, TensorMemory::DmaBuf | TensorMemory::Mem),
8502 "Unexpected auto-fallback result on macOS/iOS: {m:?}"
8503 );
8504 }
8505
8506 #[test]
8507 #[cfg(all(
8508 unix,
8509 not(any(target_os = "linux", target_os = "macos", target_os = "ios"))
8510 ))]
8511 fn test_tensor() {
8512 let shape = vec![1];
8513 let tensor = Tensor::<f32>::new(&shape, None, None).expect("Failed to create tensor");
8514 // Other Unix (BSD): no DMA, so auto-selection is Mem (Shm is
8515 // explicit-only, never auto-selected).
8516 assert_eq!(tensor.memory(), TensorMemory::Mem);
8517 }
8518
8519 #[test]
8520 #[cfg(not(unix))]
8521 fn test_tensor() {
8522 let shape = vec![1];
8523 let tensor = Tensor::<f32>::new(&shape, None, None).expect("Failed to create tensor");
8524 assert_eq!(tensor.memory(), TensorMemory::Mem);
8525 }
8526
8527 #[test]
8528 #[cfg(target_os = "linux")]
8529 fn test_dma_tensor() {
8530 let _lock = crate::tests::fd_lock_shared();
8531 if !crate::test_support::dma_or_skip("test_dma_tensor") {
8532 return;
8533 }
8534
8535 let shape = vec![2, 3, 4];
8536 let tensor =
8537 DmaTensor::<f32>::new(&shape, Some("test_tensor")).expect("Failed to create tensor");
8538
8539 const DUMMY_VALUE: f32 = 12.34;
8540
8541 assert_eq!(tensor.memory(), TensorMemory::DmaBuf);
8542 assert_eq!(tensor.name(), "test_tensor");
8543 assert_eq!(tensor.shape(), &shape);
8544 assert_eq!(tensor.size(), 2 * 3 * 4 * std::mem::size_of::<f32>());
8545 assert_eq!(tensor.len(), 2 * 3 * 4);
8546
8547 {
8548 let mut tensor_map = tensor.map().expect("Failed to map DMA memory");
8549 tensor_map.fill(42.0);
8550 assert!(tensor_map.iter().all(|&x| x == 42.0));
8551 }
8552
8553 {
8554 let shared = Tensor::<f32>::from_fd(
8555 tensor
8556 .clone_fd()
8557 .expect("Failed to duplicate tensor file descriptor"),
8558 &shape,
8559 Some("test_tensor_shared"),
8560 )
8561 .expect("Failed to create tensor from fd");
8562
8563 assert_eq!(shared.memory(), TensorMemory::DmaBuf);
8564 assert_eq!(shared.name(), "test_tensor_shared");
8565 assert_eq!(shared.shape(), &shape);
8566
8567 let mut tensor_map = shared.map().expect("Failed to map DMA memory from fd");
8568 tensor_map.fill(DUMMY_VALUE);
8569 assert!(tensor_map.iter().all(|&x| x == DUMMY_VALUE));
8570 }
8571
8572 {
8573 let tensor_map = tensor.map().expect("Failed to map DMA memory");
8574 assert!(tensor_map.iter().all(|&x| x == DUMMY_VALUE));
8575 }
8576
8577 let mut tensor = DmaTensor::<u8>::new(&shape, None).expect("Failed to create tensor");
8578 assert_eq!(tensor.shape(), &shape);
8579 let new_shape = vec![3, 4, 4];
8580 assert!(
8581 tensor.reshape(&new_shape).is_err(),
8582 "Reshape should fail due to size mismatch"
8583 );
8584 assert_eq!(tensor.shape(), &shape, "Shape should remain unchanged");
8585
8586 let new_shape = vec![2, 3, 4];
8587 tensor.reshape(&new_shape).expect("Reshape should succeed");
8588 assert_eq!(
8589 tensor.shape(),
8590 &new_shape,
8591 "Shape should be updated after successful reshape"
8592 );
8593
8594 {
8595 let mut tensor_map = tensor.map().expect("Failed to map DMA memory");
8596 tensor_map.fill(1);
8597 assert!(tensor_map.iter().all(|&x| x == 1));
8598 }
8599
8600 {
8601 let mut tensor_map = tensor.map().expect("Failed to map DMA memory");
8602 tensor_map[2] = 42;
8603 assert_eq!(tensor_map[1], 1, "Value at index 1 should be 1");
8604 assert_eq!(tensor_map[2], 42, "Value at index 2 should be 42");
8605 }
8606 }
8607
8608 #[test]
8609 #[cfg(unix)]
8610 fn test_shm_tensor() {
8611 let _lock = crate::tests::fd_lock_shared();
8612 if !is_shm_available() {
8613 crate::test_support::report_skip(
8614 "test_shm_tensor - SHM allocation unavailable on this platform",
8615 );
8616 return;
8617 }
8618 let shape = vec![2, 3, 4];
8619 let tensor =
8620 ShmTensor::<f32>::new(&shape, Some("test_tensor")).expect("Failed to create tensor");
8621 assert_eq!(tensor.shape(), &shape);
8622 assert_eq!(tensor.size(), 2 * 3 * 4 * std::mem::size_of::<f32>());
8623 assert_eq!(tensor.name(), "test_tensor");
8624
8625 const DUMMY_VALUE: f32 = 12.34;
8626 {
8627 let mut tensor_map = tensor.map().expect("Failed to map shared memory");
8628 tensor_map.fill(42.0);
8629 assert!(tensor_map.iter().all(|&x| x == 42.0));
8630 }
8631
8632 {
8633 let shared = Tensor::<f32>::from_fd(
8634 tensor
8635 .clone_fd()
8636 .expect("Failed to duplicate tensor file descriptor"),
8637 &shape,
8638 Some("test_tensor_shared"),
8639 )
8640 .expect("Failed to create tensor from fd");
8641
8642 assert_eq!(shared.memory(), TensorMemory::Shm);
8643 assert_eq!(shared.name(), "test_tensor_shared");
8644 assert_eq!(shared.shape(), &shape);
8645
8646 let mut tensor_map = shared.map().expect("Failed to map shared memory from fd");
8647 tensor_map.fill(DUMMY_VALUE);
8648 assert!(tensor_map.iter().all(|&x| x == DUMMY_VALUE));
8649 }
8650
8651 {
8652 let tensor_map = tensor.map().expect("Failed to map shared memory");
8653 assert!(tensor_map.iter().all(|&x| x == DUMMY_VALUE));
8654 }
8655
8656 let mut tensor = ShmTensor::<u8>::new(&shape, None).expect("Failed to create tensor");
8657 assert_eq!(tensor.shape(), &shape);
8658 let new_shape = vec![3, 4, 4];
8659 assert!(
8660 tensor.reshape(&new_shape).is_err(),
8661 "Reshape should fail due to size mismatch"
8662 );
8663 assert_eq!(tensor.shape(), &shape, "Shape should remain unchanged");
8664
8665 let new_shape = vec![2, 3, 4];
8666 tensor.reshape(&new_shape).expect("Reshape should succeed");
8667 assert_eq!(
8668 tensor.shape(),
8669 &new_shape,
8670 "Shape should be updated after successful reshape"
8671 );
8672
8673 {
8674 let mut tensor_map = tensor.map().expect("Failed to map shared memory");
8675 tensor_map.fill(1);
8676 assert!(tensor_map.iter().all(|&x| x == 1));
8677 }
8678
8679 {
8680 let mut tensor_map = tensor.map().expect("Failed to map shared memory");
8681 tensor_map[2] = 42;
8682 assert_eq!(tensor_map[1], 1, "Value at index 1 should be 1");
8683 assert_eq!(tensor_map[2], 42, "Value at index 2 should be 42");
8684 }
8685 }
8686
8687 #[test]
8688 fn mem_subview_partitions_parent_buffer() {
8689 // One heap [2,4] u8 parent (8 bytes). Two [1,4] sub-views at byte
8690 // offsets 0 and 4 must share the parent allocation (zero-copy) and be
8691 // independently writable: view 0 owns bytes [0,4), view 1 owns [4,8).
8692 // Today this is impossible — heap offset is rejected and there is no
8693 // shared sub-view constructor.
8694 let parent = Tensor::<u8>::new(&[2, 4], Some(TensorMemory::Mem), None).unwrap();
8695 let view0 = parent.subview(0, &[1, 4]).expect("subview at offset 0");
8696 let view1 = parent.subview(4, &[1, 4]).expect("subview at offset 4");
8697
8698 view1
8699 .map()
8700 .unwrap()
8701 .as_mut_slice()
8702 .copy_from_slice(&[10, 20, 30, 40]);
8703 view0
8704 .map()
8705 .unwrap()
8706 .as_mut_slice()
8707 .copy_from_slice(&[1, 2, 3, 4]);
8708
8709 // Each view sees only its own window.
8710 assert_eq!(view0.map().unwrap().as_slice(), &[1, 2, 3, 4]);
8711 assert_eq!(view1.map().unwrap().as_slice(), &[10, 20, 30, 40]);
8712 // The parent buffer is correctly partitioned (shared, zero-copy).
8713 assert_eq!(
8714 parent.map().unwrap().as_slice(),
8715 &[1, 2, 3, 4, 10, 20, 30, 40]
8716 );
8717 }
8718
8719 #[test]
8720 fn batch_partitions_leading_dim() {
8721 // Raw [4,2,2,3] u8 batched tensor: 4 elements of 12 bytes each. batch(n)
8722 // yields element n at offset n*12, sharing the parent buffer (zero-copy).
8723 let parent = Tensor::<u8>::new(&[4, 2, 2, 3], Some(TensorMemory::Mem), None).unwrap();
8724 for i in 0..4u8 {
8725 let e = parent.batch(i as usize).expect("batch element");
8726 assert_eq!(e.shape(), &[2, 2, 3]);
8727 // A batch element shares the parent's BufferIdentity.
8728 assert_eq!(e.buffer_identity().id(), parent.buffer_identity().id());
8729 for b in e.map().unwrap().as_mut_slice() {
8730 *b = i + 1;
8731 }
8732 }
8733 // Each element occupies its own 12-byte band of the parent.
8734 let whole = parent.map().unwrap();
8735 let s = whole.as_slice();
8736 for i in 0..4usize {
8737 assert!(
8738 s[i * 12..(i + 1) * 12].iter().all(|&b| b == (i as u8 + 1)),
8739 "band {i} not partitioned: {:?}",
8740 &s[i * 12..(i + 1) * 12]
8741 );
8742 }
8743 }
8744
8745 #[test]
8746 fn view_origin_snapshots_parent_and_composes() {
8747 // view() on a whole image snapshots the parent dims + the view's origin.
8748 let parent = Tensor::<u8>::image(
8749 100,
8750 80,
8751 PixelFormat::Rgba,
8752 Some(TensorMemory::Mem),
8753 crate::CpuAccess::ReadWrite,
8754 )
8755 .unwrap();
8756 assert_eq!(
8757 parent.view_origin(),
8758 None,
8759 "whole tensor has no view_origin"
8760 );
8761 let v = parent.view(Region::new(10, 20, 30, 40)).unwrap();
8762 assert_eq!(
8763 v.view_origin(),
8764 Some(ViewOrigin {
8765 parent_width: 100,
8766 parent_height: 80,
8767 parent_row_stride: 100 * 4, // tight RGBA pitch
8768 x: 10,
8769 y: 20
8770 })
8771 );
8772 // A view of a view keeps the ROOT parent and accumulates the origin.
8773 let v2 = v.view(Region::new(5, 5, 10, 10)).unwrap();
8774 assert_eq!(
8775 v2.view_origin(),
8776 Some(ViewOrigin {
8777 parent_width: 100,
8778 parent_height: 80,
8779 parent_row_stride: 100 * 4,
8780 x: 15,
8781 y: 25
8782 }),
8783 "nested view composes onto the root parent"
8784 );
8785 }
8786
8787 #[test]
8788 fn view_origin_none_for_raw_batch() {
8789 // A raw (unformatted) batched tensor has no pixel geometry, so batch()
8790 // leaves view_origin None (the per-slot path, not the one-import pivot).
8791 let parent = Tensor::<u8>::new(&[4, 2, 2, 3], Some(TensorMemory::Mem), None).unwrap();
8792 assert_eq!(parent.batch(2).unwrap().view_origin(), None);
8793 }
8794
8795 #[test]
8796 fn batch_rejects_out_of_bounds_index() {
8797 let parent = Tensor::<u8>::new(&[4, 2, 2, 3], Some(TensorMemory::Mem), None).unwrap();
8798 match parent.batch(4) {
8799 Err(Error::BatchIndexOutOfBounds { index, batch }) => {
8800 assert_eq!((index, batch), (4, 4));
8801 }
8802 other => panic!("expected BatchIndexOutOfBounds, got {other:?}"),
8803 }
8804 }
8805
8806 #[test]
8807 fn batch_zero_on_unit_n_is_whole() {
8808 // N == 1: batch(0) is the whole per-element block at offset 0 (no plane_offset).
8809 let parent = Tensor::<u8>::new(&[1, 2, 2, 3], Some(TensorMemory::Mem), None).unwrap();
8810 let e = parent.batch(0).unwrap();
8811 assert_eq!(e.shape(), &[2, 2, 3]);
8812 assert_eq!(e.plane_offset(), None);
8813 assert_eq!(e.buffer_identity().id(), parent.buffer_identity().id());
8814 }
8815
8816 #[test]
8817 fn mem_subview_rejects_unaligned_offset() {
8818 // f32 has align 4; a byte offset of 2 cannot back a valid `*const f32`.
8819 let parent = Tensor::<f32>::new(&[8], Some(TensorMemory::Mem), None).unwrap();
8820 assert!(parent.subview(2, &[1]).is_err());
8821 // A correctly aligned offset is accepted.
8822 assert!(parent.subview(4, &[1]).is_ok());
8823 }
8824
8825 #[test]
8826 fn mem_subview_rejects_out_of_bounds() {
8827 let parent = Tensor::<u8>::new(&[8], Some(TensorMemory::Mem), None).unwrap();
8828 // offset 6 + 4 bytes = 10 exceeds the 8-byte allocation.
8829 assert!(parent.subview(6, &[4]).is_err());
8830 }
8831
8832 /// Regression guard for the `TensorTrait::view` promotion (R2): a `subview`
8833 /// must share the parent's `BufferIdentity` on **every** backend, not mint a
8834 /// fresh one. Identity-keyed caches (the GL EGLImage import) rely on this to
8835 /// treat offset-distinct windows of one buffer as a single import; a fresh
8836 /// identity would silently break that and regress zero-copy import reuse.
8837 ///
8838 /// Runs each backend that can be allocated without a GPU/GL context on the
8839 /// test host: `Mem` (always); `Shm` (when POSIX shm is available); the
8840 /// platform-native zero-copy buffer `Dma` (DMA-BUF on Linux / IOSurface on
8841 /// macOS, when available). `Pbo` shares its identity the same way (see
8842 /// `pbo.rs` `view`) but needs a live GL context, so it is exercised by the
8843 /// image-crate GL tests rather than here.
8844 #[test]
8845 fn subview_shares_buffer_identity_all_backends() {
8846 // Declares this test as an fd-opener; see FD_LOCK.
8847 let _lock = crate::tests::fd_lock_shared();
8848 // u8 has align 1, so every byte offset is valid for the alignment check;
8849 // this isolates the identity-sharing contract from alignment concerns.
8850 // `window` is the view shape; a formatted parent validates it against
8851 // its own format, so an image parent needs an image-shaped window.
8852 let assert_parent_shares = |parent: Tensor<u8>, window: &[usize], label: &str| {
8853 let parent_id = parent.buffer_identity().id();
8854 // Two offset-distinct windows must both carry the parent's identity.
8855 let v0 = parent
8856 .subview(0, window)
8857 .unwrap_or_else(|e| panic!("{label}: subview(0) failed: {e:?}"));
8858 let v1 = parent
8859 .subview(16, window)
8860 .unwrap_or_else(|e| panic!("{label}: subview(16) failed: {e:?}"));
8861 assert_eq!(
8862 v0.buffer_identity().id(),
8863 parent_id,
8864 "{label}: subview(0) minted a fresh BufferIdentity"
8865 );
8866 assert_eq!(
8867 v1.buffer_identity().id(),
8868 parent_id,
8869 "{label}: subview(16) minted a fresh BufferIdentity"
8870 );
8871 };
8872 let assert_shares = |memory: TensorMemory, label: &str| {
8873 let parent = Tensor::<u8>::new(&[64], Some(memory), None)
8874 .unwrap_or_else(|e| panic!("{label}: parent alloc failed: {e:?}"));
8875 assert_parent_shares(parent, &[16], label);
8876 };
8877
8878 assert_shares(TensorMemory::Mem, "Mem");
8879
8880 #[cfg(unix)]
8881 if crate::is_shm_available() {
8882 assert_shares(TensorMemory::Shm, "Shm");
8883 }
8884
8885 // Dma == DMA-BUF on Linux, IOSurface on macOS, AHardwareBuffer on
8886 // Android; same public variant.
8887 #[cfg(target_os = "linux")]
8888 if crate::test_support::dma_or_skip("subview_shares_buffer_identity_all_backends (Dma)") {
8889 assert_shares(TensorMemory::DmaBuf, "Dma");
8890 }
8891 #[cfg(not(any(target_os = "linux", target_os = "windows")))]
8892 if crate::is_gpu_buffer_available() {
8893 assert_shares(TensorMemory::DmaBuf, "Dma");
8894 }
8895 // Windows' Dma is an image-formatted D3D11 texture, which
8896 // `Tensor::new` deliberately refuses -- allocate it the only way it
8897 // can be allocated and hold it to the same contract.
8898 #[cfg(target_os = "windows")]
8899 if crate::is_gpu_buffer_available() {
8900 let parent = Tensor::<u8>::image(
8901 64,
8902 1,
8903 PixelFormat::Grey,
8904 Some(TensorMemory::DmaBuf),
8905 CpuAccess::ReadWrite,
8906 )
8907 .expect("D3D11 texture alloc failed");
8908 assert_parent_shares(parent, &[1, 16, 1], "Dma");
8909 }
8910 }
8911
8912 /// A sub-view records a plane offset, and on Windows the D3D11 storage
8913 /// already applied that offset itself when the view was made -- so the
8914 /// map's plane-offset gate must accept it, exactly as it does for the
8915 /// other zero-copy backings. It used to refuse, with a message that named
8916 /// every backend but this one.
8917 #[cfg(target_os = "windows")]
8918 #[test]
8919 fn subview_of_a_texture_tensor_maps_at_its_offset() {
8920 if !crate::is_gpu_buffer_available() {
8921 crate::test_support::report_skip(
8922 "subview_of_a_texture_tensor_maps_at_its_offset - no D3D11 device on this host",
8923 );
8924 return;
8925 }
8926 let t = Tensor::<u8>::image(
8927 32,
8928 8,
8929 PixelFormat::Rgba,
8930 Some(TensorMemory::DmaBuf),
8931 CpuAccess::ReadWrite,
8932 )
8933 .unwrap();
8934 // For a texture tensor this is the backing pitch either way: the
8935 // allocation records the staging pitch when it exceeds the tight row,
8936 // and the two are equal when it does not.
8937 let stride = t.effective_row_stride().unwrap();
8938 let v = t.subview(stride * 2, &[2, 32, 4]).unwrap();
8939 assert!(
8940 v.plane_offset().is_some_and(|o| o > 0),
8941 "a sub-view at a row offset must carry one"
8942 );
8943 let m = v
8944 .map_read()
8945 .expect("a sub-view of a texture tensor must map");
8946 assert_eq!(m.as_slice().len(), 2 * stride, "two rows of the window");
8947 }
8948
8949 #[test]
8950 fn mem_subview_four_views_no_aliasing() {
8951 // One [4,3] f32 parent; four [1,3] views at 12-byte strides, each
8952 // written independently. Exercises a multi-byte element type (offsets
8953 // must stay element-aligned) and N-way zero-copy sharing.
8954 let parent = Tensor::<f32>::new(&[4, 3], Some(TensorMemory::Mem), None).unwrap();
8955 let frame = 3 * std::mem::size_of::<f32>();
8956 for i in 0..4 {
8957 let v = parent.subview(i * frame, &[1, 3]).unwrap();
8958 let val = i as f32 + 1.0;
8959 v.map()
8960 .unwrap()
8961 .as_mut_slice()
8962 .copy_from_slice(&[val, val, val]);
8963 }
8964 assert_eq!(
8965 parent.map().unwrap().as_slice(),
8966 &[1.0, 1.0, 1.0, 2.0, 2.0, 2.0, 3.0, 3.0, 3.0, 4.0, 4.0, 4.0]
8967 );
8968 }
8969
8970 #[test]
8971 fn mem_subview_inherits_format_and_row_stride() {
8972 // A sub-view is a ready-to-use sub-image: it inherits the parent's
8973 // pixel format and (crucially) its padded row stride, so a strided
8974 // parent yields strided windows. Set a stride wider than the tight row
8975 // to exercise the row_stride inheritance path specifically.
8976 let mut parent = Tensor::<u8>::image(
8977 100,
8978 100,
8979 PixelFormat::Rgba,
8980 Some(TensorMemory::Mem),
8981 crate::CpuAccess::ReadWrite,
8982 )
8983 .unwrap();
8984 parent.set_row_stride_unchecked(512); // padded stride (> 100*4)
8985 let view = parent.subview(4096, &[10, 10, 4]).unwrap();
8986 assert_eq!(view.format(), Some(PixelFormat::Rgba), "format inherited");
8987 assert_eq!(view.row_stride(), Some(512), "row_stride inherited");
8988 }
8989
8990 #[test]
8991 fn mem_strided_subview_maps_offset_and_byte_size() {
8992 // Integration of the sub-region offset (PR #89) and the strided-map
8993 // `byte_size_override` (PR #90): a strided sub-view exposes its full
8994 // padded window (`row_stride × rows`) starting at the view's byte
8995 // offset, mapped zero-copy into the parent.
8996 let parent = Tensor::<u8>::new(&[2048], Some(TensorMemory::Mem), None).unwrap();
8997 let mut view = parent.subview(128, &[8, 16]).unwrap(); // 8 rows × 16 @ off 128
8998 assert_eq!(view.plane_offset(), Some(128));
8999 view.set_row_stride_unchecked(32); // padded stride (> 16)
9000
9001 {
9002 let mut m = view.map().unwrap();
9003 let s = m.as_mut_slice();
9004 // Strided map exposes the padded window: stride(32) × rows(8) = 256.
9005 assert_eq!(
9006 s.len(),
9007 256,
9008 "strided map exposes the full padded byte window"
9009 );
9010 s[0] = 0xAA; // row 0, col 0
9011 s[32] = 0xBB; // row 1, col 0 (one stride in)
9012 }
9013
9014 // Zero-copy: the writes land in the parent at the view's offset.
9015 let p = parent.map().unwrap();
9016 let pb = p.as_slice();
9017 assert_eq!(pb[128], 0xAA, "row 0 writes at parent offset 128");
9018 assert_eq!(
9019 pb[128 + 32],
9020 0xBB,
9021 "row 1 writes at parent offset 128 + stride"
9022 );
9023 }
9024
9025 #[test]
9026 #[cfg(unix)]
9027 fn shm_subview_partitions_parent_buffer() {
9028 // Declares this test as an fd-opener; see FD_LOCK.
9029 let _lock = crate::tests::fd_lock_shared();
9030 // Mirrors `mem_subview_partitions_parent_buffer` for Shm: one [2,4] u8
9031 // parent shared segment (8 bytes); two [1,4] sub-views at byte offsets 0
9032 // and 4 must share the segment (zero-copy, via cloned fd) and be
9033 // independently writable — view 0 owns [0,4), view 1 owns [4,8).
9034 if !crate::is_shm_available() {
9035 crate::test_support::report_skip("shm not available");
9036 return;
9037 }
9038 let parent = Tensor::<u8>::new(&[2, 4], Some(TensorMemory::Shm), None).unwrap();
9039 let view0 = parent.subview(0, &[1, 4]).expect("shm subview at offset 0");
9040 let view1 = parent.subview(4, &[1, 4]).expect("shm subview at offset 4");
9041
9042 view1
9043 .map()
9044 .unwrap()
9045 .as_mut_slice()
9046 .copy_from_slice(&[10, 20, 30, 40]);
9047 view0
9048 .map()
9049 .unwrap()
9050 .as_mut_slice()
9051 .copy_from_slice(&[1, 2, 3, 4]);
9052
9053 assert_eq!(view0.map().unwrap().as_slice(), &[1, 2, 3, 4]);
9054 assert_eq!(view1.map().unwrap().as_slice(), &[10, 20, 30, 40]);
9055 // The parent sees the full partitioned segment (shared, zero-copy).
9056 assert_eq!(
9057 parent.map().unwrap().as_slice(),
9058 &[1, 2, 3, 4, 10, 20, 30, 40]
9059 );
9060 // A sub-view of a sub-view composes the offset.
9061 let nested = view1.subview(2, &[1, 2]).expect("nested shm subview");
9062 assert_eq!(nested.map().unwrap().as_slice(), &[30, 40]);
9063 }
9064
9065 #[test]
9066 #[cfg(unix)]
9067 fn shm_subview_rejects_unaligned_and_oob() {
9068 // Declares this test as an fd-opener; see FD_LOCK.
9069 let _lock = crate::tests::fd_lock_shared();
9070 if !crate::is_shm_available() {
9071 crate::test_support::report_skip("shm not available");
9072 return;
9073 }
9074 // f32 align 4: a 2-byte offset cannot back a valid `*const f32`.
9075 let parent = Tensor::<f32>::new(&[8], Some(TensorMemory::Shm), None).unwrap();
9076 assert!(parent.subview(2, &[1]).is_err());
9077 assert!(parent.subview(4, &[1]).is_ok());
9078 // Out of bounds: offset 6 + 4 bytes = 10 > 8-byte (u8) segment.
9079 let p2 = Tensor::<u8>::new(&[8], Some(TensorMemory::Shm), None).unwrap();
9080 assert!(p2.subview(6, &[4]).is_err());
9081 }
9082
9083 #[test]
9084 #[cfg(target_os = "linux")]
9085 fn dma_subview_matches_mem_subview() {
9086 // Serialize against the fd-leak tests: this test opens DMA fds (alloc +
9087 // clone_fd), which would otherwise perturb their fd counts.
9088 let _lock = crate::tests::fd_lock_shared();
9089 // Identical sub-view semantics across Dma (shared fd) and Mem (shared
9090 // Arc): same offsets → same logical windows → same partition.
9091 if !crate::test_support::dma_or_skip("dma_subview_matches_mem_subview") {
9092 return;
9093 }
9094 let dma = Tensor::<u8>::new(&[8], Some(TensorMemory::DmaBuf), None).unwrap();
9095 let mem = Tensor::<u8>::new(&[8], Some(TensorMemory::Mem), None).unwrap();
9096 for parent in [&dma, &mem] {
9097 let v0 = parent.subview(0, &[4]).unwrap();
9098 let v1 = parent.subview(4, &[4]).unwrap();
9099 v0.map()
9100 .unwrap()
9101 .as_mut_slice()
9102 .copy_from_slice(&[1, 2, 3, 4]);
9103 v1.map()
9104 .unwrap()
9105 .as_mut_slice()
9106 .copy_from_slice(&[5, 6, 7, 8]);
9107 assert_eq!(parent.map().unwrap().as_slice(), &[1, 2, 3, 4, 5, 6, 7, 8]);
9108 }
9109 }
9110
9111 #[test]
9112 #[cfg(target_os = "linux")]
9113 fn dma_strided_subview_maps_padded_window() {
9114 // The strided-map path differs by backing: DMA maps through
9115 // `mmap_offset` + the `byte_size_override`, not the Mem `Arc` slice. A
9116 // padded sub-view of a DMA buffer must still expose its full
9117 // `row_stride × rows` window zero-copy at the view's offset (the GPU
9118 // batched-render-to-DMA case). Mirrors
9119 // `mem_strided_subview_maps_offset_and_byte_size` on a Dma parent.
9120 let _lock = crate::tests::fd_lock_shared();
9121 if !crate::test_support::dma_or_skip("dma_strided_subview_maps_padded_window") {
9122 return;
9123 }
9124 let parent = Tensor::<u8>::new(&[2048], Some(TensorMemory::DmaBuf), None).unwrap();
9125 let mut view = parent.subview(128, &[8, 16]).unwrap();
9126 assert_eq!(view.plane_offset(), Some(128));
9127 view.set_row_stride_unchecked(32); // padded stride (> 16)
9128
9129 {
9130 let mut m = view.map().unwrap();
9131 let s = m.as_mut_slice();
9132 assert_eq!(s.len(), 256, "strided DMA map exposes stride(32) × rows(8)");
9133 s[0] = 0xAA; // row 0, col 0
9134 s[32] = 0xBB; // row 1, col 0 (one stride in)
9135 }
9136
9137 let p = parent.map().unwrap();
9138 let pb = p.as_slice();
9139 assert_eq!(pb[128], 0xAA, "row 0 writes at parent offset 128");
9140 assert_eq!(
9141 pb[128 + 32],
9142 0xBB,
9143 "row 1 writes at parent offset 128 + stride"
9144 );
9145 }
9146
9147 #[test]
9148 #[cfg(target_os = "linux")]
9149 fn view_single_row_snapshots_parent_stride() {
9150 // A single-row `view()` keeps a TIGHT `row_stride` for map-span safety,
9151 // but its `view_origin` snapshots the PARENT row stride — the GL backend
9152 // keys its EGLImage import/pitch on that snapshot (not the view's tight
9153 // stride), so single-row and multi-row sibling views collapse onto the
9154 // same parent import.
9155 let _lock = crate::tests::fd_lock_shared();
9156 // 8x4 RGBA with a padded 64-byte row stride (tight row = 8*4 = 32).
9157 if !crate::test_support::dma_or_skip("view_single_row_snapshots_parent_stride") {
9158 return;
9159 }
9160 let parent = Tensor::<u8>::image_with_stride(
9161 8,
9162 4,
9163 PixelFormat::Rgba,
9164 64,
9165 Some(TensorMemory::DmaBuf),
9166 crate::CpuAccess::ReadWrite,
9167 )
9168 .unwrap();
9169 assert_eq!(parent.effective_row_stride(), Some(64));
9170 // Bottom row (y=3) at x>0 — the case the tight single-row stride guards.
9171 let row = parent.view(Region::new(2, 3, 4, 1)).unwrap();
9172 // The view's own stride is tight (4*4 = 16) so its strided map stays in
9173 // bounds; the GL-facing parent pitch (64) lives in `view_origin`.
9174 assert_eq!(row.effective_row_stride(), Some(16));
9175 let vo = row.view_origin().expect("a view carries a view_origin");
9176 assert_eq!(
9177 vo.parent_row_stride, 64,
9178 "GL keys/pitches a view on the parent stride, not its tight one"
9179 );
9180 // The tight stride keeps map() in-bounds for the bottom / x>0 single row.
9181 assert_eq!(row.map().unwrap().as_slice().len(), 16);
9182 }
9183
9184 #[test]
9185 fn test_mem_tensor() {
9186 let shape = vec![2, 3, 4];
9187 let tensor =
9188 MemTensor::<f32>::new(&shape, Some("test_tensor")).expect("Failed to create tensor");
9189 assert_eq!(tensor.shape(), &shape);
9190 assert_eq!(tensor.size(), 2 * 3 * 4 * std::mem::size_of::<f32>());
9191 assert_eq!(tensor.name(), "test_tensor");
9192
9193 {
9194 let mut tensor_map = tensor.map().expect("Failed to map memory");
9195 tensor_map.fill(42.0);
9196 assert!(tensor_map.iter().all(|&x| x == 42.0));
9197 }
9198
9199 let mut tensor = MemTensor::<u8>::new(&shape, None).expect("Failed to create tensor");
9200 assert_eq!(tensor.shape(), &shape);
9201 let new_shape = vec![3, 4, 4];
9202 assert!(
9203 tensor.reshape(&new_shape).is_err(),
9204 "Reshape should fail due to size mismatch"
9205 );
9206 assert_eq!(tensor.shape(), &shape, "Shape should remain unchanged");
9207
9208 let new_shape = vec![2, 3, 4];
9209 tensor.reshape(&new_shape).expect("Reshape should succeed");
9210 assert_eq!(
9211 tensor.shape(),
9212 &new_shape,
9213 "Shape should be updated after successful reshape"
9214 );
9215
9216 {
9217 let mut tensor_map = tensor.map().expect("Failed to map memory");
9218 tensor_map.fill(1);
9219 assert!(tensor_map.iter().all(|&x| x == 1));
9220 }
9221
9222 {
9223 let mut tensor_map = tensor.map().expect("Failed to map memory");
9224 tensor_map[2] = 42;
9225 assert_eq!(tensor_map[1], 1, "Value at index 1 should be 1");
9226 assert_eq!(tensor_map[2], 42, "Value at index 2 should be 42");
9227 }
9228 }
9229
9230 #[test]
9231 #[cfg(target_os = "linux")]
9232 fn test_dma_no_fd_leaks() {
9233 let _lock = fd_lock_exclusive();
9234 if !crate::test_support::dma_or_skip(function!()) {
9235 return;
9236 }
9237
9238 let baseline = open_fd_count(DMA_FD_TARGETS);
9239
9240 // Self-check — see `test_shm_no_fd_leaks`.
9241 {
9242 let probe = Tensor::<u8>::new(&[100, 100], Some(TensorMemory::DmaBuf), None)
9243 .expect("Failed to create tensor");
9244 assert!(
9245 open_fd_count(DMA_FD_TARGETS) > baseline,
9246 "a live DMA-BUF tensor must be visible to the dmabuf fd counter, \
9247 otherwise this test cannot detect a leak"
9248 );
9249 drop(probe);
9250 }
9251
9252 for _ in 0..100 {
9253 let tensor = Tensor::<u8>::new(&[100, 100], Some(TensorMemory::DmaBuf), None)
9254 .expect("Failed to create tensor");
9255 let mut map = tensor.map().unwrap();
9256 map.as_mut_slice().fill(233);
9257 }
9258
9259 let after = open_fd_count(DMA_FD_TARGETS);
9260 assert_eq!(
9261 baseline, after,
9262 "DMA-BUF file descriptor leak detected: {baseline} -> {after}"
9263 );
9264 }
9265
9266 #[test]
9267 #[cfg(target_os = "linux")]
9268 fn test_dma_from_fd_no_fd_leaks() {
9269 let _lock = fd_lock_exclusive();
9270 if !crate::test_support::dma_or_skip(function!()) {
9271 return;
9272 }
9273
9274 let baseline = open_fd_count(DMA_FD_TARGETS);
9275 let orig = Tensor::<u8>::new(&[100, 100], Some(TensorMemory::DmaBuf), None).unwrap();
9276
9277 // Self-check — see `test_shm_no_fd_leaks`.
9278 let held = open_fd_count(DMA_FD_TARGETS);
9279 assert!(
9280 held > baseline,
9281 "a live DMA-BUF tensor must be visible to the dmabuf fd counter, \
9282 otherwise this test cannot detect a leak"
9283 );
9284
9285 for _ in 0..100 {
9286 let tensor =
9287 Tensor::<u8>::from_fd(orig.clone_fd().unwrap(), orig.shape(), None).unwrap();
9288 assert_eq!(
9289 tensor.memory(),
9290 TensorMemory::DmaBuf,
9291 "DMA-BUF fd must import as Dma, not be silently downgraded"
9292 );
9293 let mut map = tensor.map().unwrap();
9294 map.as_mut_slice().fill(233);
9295 }
9296
9297 let after_loop = open_fd_count(DMA_FD_TARGETS);
9298 assert_eq!(
9299 held, after_loop,
9300 "DMA-BUF file descriptor leak detected: {held} -> {after_loop}"
9301 );
9302
9303 // And the source tensor's own fd goes back on drop.
9304 drop(orig);
9305 let after_drop = open_fd_count(DMA_FD_TARGETS);
9306 assert_eq!(
9307 baseline, after_drop,
9308 "dropping the source tensor must release its fd: {baseline} -> {after_drop}"
9309 );
9310 }
9311
9312 /// A filesystem magic must report its true 32-bit value regardless of
9313 /// how wide, and how signed, `fstatfs`'s `f_type` is on the target.
9314 ///
9315 /// `fs_type_t` is `__fsword_t` on Linux/gnu — `i64` on 64-bit but `i32`
9316 /// on 32-bit (armv7, i686, aarch64-ilp32) — `c_int` on uclibc, and
9317 /// unsigned on musl and s390x. Widening a *signed* 32-bit `f_type`
9318 /// sign-extends every magic with bit 31 set, so a naive widening cast
9319 /// reports e.g. `0xffffffff958458f6` for hugetlbfs. The docs tell
9320 /// callers to look the reported value up in `include/uapi/linux/magic.h`,
9321 /// so a sign-extended value is actively misleading.
9322 ///
9323 /// This cannot be reproduced on a 64-bit host — the defect only exists
9324 /// where `f_type` is a signed 32-bit type — so the test drives the
9325 /// conversion directly with the value such a target would produce.
9326 #[test]
9327 #[cfg(target_os = "linux")]
9328 fn test_fs_magic_normalizes_sign_extended_values() {
9329 // Declares this test as an fd-opener; see FD_LOCK.
9330 let _lock = crate::tests::fd_lock_shared();
9331 // Magics whose bit 31 is set. HUGETLBFS is the one that matters
9332 // most in practice: a MFD_HUGETLB memfd is the likeliest fd to land
9333 // in the UnknownBufferType arm.
9334 const HUGETLBFS_MAGIC: u32 = 0x9584_58f6;
9335 const F2FS_MAGIC: u32 = 0xf2f5_2010;
9336 const BTRFS_MAGIC: u32 = 0x9123_683e;
9337
9338 for magic in [HUGETLBFS_MAGIC, F2FS_MAGIC, BTRFS_MAGIC] {
9339 // 64-bit gnu: f_type is i64 and already holds the true value.
9340 assert_eq!(fs_magic(i64::from(magic)), magic);
9341
9342 // 32-bit gnu / uclibc: f_type is i32, so the value arrives
9343 // sign-extended once widened. It must still report its true
9344 // 32 bits.
9345 let sign_extended = i64::from(magic as i32);
9346 assert!(sign_extended < 0, "{magic:#x} should have bit 31 set");
9347 assert_eq!(
9348 fs_magic(sign_extended),
9349 magic,
9350 "{magic:#x} must survive a signed 32-bit f_type"
9351 );
9352 }
9353
9354 // The two magics we actually classify on are below 2^31, so they are
9355 // unaffected by signedness either way — this is why the DMA-vs-SHM
9356 // decision is correct on 32-bit even without this normalization.
9357 for magic in [DMA_BUF_MAGIC, TMPFS_MAGIC] {
9358 assert!(magic < 0x8000_0000);
9359 assert_eq!(fs_magic(i64::from(magic as i32)), magic);
9360 }
9361 }
9362
9363 /// A DMA-BUF fd must import as [`TensorMemory::DmaBuf`].
9364 ///
9365 /// Regression test for the `st_dev` minor-number classifier, which
9366 /// hardcoded `9 | 10` as "this is DMA". Those minors come from
9367 /// `get_anon_bdev()` and are assigned first-come-first-served at boot,
9368 /// so they vary by kernel build and boot order — a real DMA-BUF is
9369 /// minor 12 on x86 desktop and minor 8 on the ADIS Verdin. Both fell
9370 /// into the `_` arm and imported as SHM, which "works" (a DMA-BUF is
9371 /// mmap-able) but silently forfeits zero-copy.
9372 #[test]
9373 #[cfg(target_os = "linux")]
9374 fn test_from_fd_dma_imports_as_dma() {
9375 let _lock = crate::tests::fd_lock_shared();
9376 if !crate::test_support::dma_or_skip(function!()) {
9377 return;
9378 }
9379
9380 let orig = Tensor::<u8>::new(&[64, 64], Some(TensorMemory::DmaBuf), None).unwrap();
9381 assert_eq!(orig.memory(), TensorMemory::DmaBuf);
9382
9383 let imported = Tensor::<u8>::from_fd(orig.clone_fd().unwrap(), orig.shape(), None).unwrap();
9384
9385 assert_eq!(
9386 imported.memory(),
9387 TensorMemory::DmaBuf,
9388 "a DMA-BUF fd must import as Dma"
9389 );
9390 }
9391
9392 /// A tmpfs/SHM fd must import as [`TensorMemory::Shm`].
9393 ///
9394 /// The companion to `test_from_fd_dma_imports_as_dma`: confirms the
9395 /// magic-based classifier identifies SHM positively rather than by
9396 /// falling through.
9397 #[test]
9398 #[cfg(target_os = "linux")]
9399 fn test_from_fd_shm_imports_as_shm() {
9400 let _lock = crate::tests::fd_lock_shared();
9401 if !is_shm_available() {
9402 crate::test_support::report_skip(&format!(
9403 "{} - SHM memory not available",
9404 function!()
9405 ));
9406 return;
9407 }
9408
9409 let orig = Tensor::<u8>::new(&[64, 64], Some(TensorMemory::Shm), None).unwrap();
9410 assert_eq!(orig.memory(), TensorMemory::Shm);
9411
9412 let imported = Tensor::<u8>::from_fd(orig.clone_fd().unwrap(), orig.shape(), None).unwrap();
9413
9414 assert_eq!(
9415 imported.memory(),
9416 TensorMemory::Shm,
9417 "a tmpfs fd must import as Shm"
9418 );
9419 }
9420
9421 /// An fd that is neither a DMA-BUF nor tmpfs must be rejected.
9422 ///
9423 /// A pipe is the convenient probe: it lives on `pipefs`, another
9424 /// `get_anon_bdev()` pseudo-filesystem, so it shares major 0 with the
9425 /// buffer types we do support and is only distinguishable by magic.
9426 /// Importing one as SHM is meaningless — `mmap` on a pipe fails — so
9427 /// the classifier must say "unknown" rather than guess.
9428 #[test]
9429 #[cfg(target_os = "linux")]
9430 fn test_from_fd_rejects_unknown_filesystem() {
9431 let _lock = crate::tests::fd_lock_shared();
9432
9433 let (read_end, _write_end) = nix::unistd::pipe().unwrap();
9434
9435 let result = Tensor::<u8>::from_fd(read_end, &[64], None);
9436
9437 match result {
9438 Err(Error::UnknownBufferType(magic)) => {
9439 // PIPEFS_MAGIC, from include/uapi/linux/magic.h
9440 assert_eq!(magic, 0x5049_5045, "expected PIPEFS_MAGIC");
9441 }
9442 other => panic!("expected UnknownBufferType for a pipe fd, got {other:?}"),
9443 }
9444 }
9445
9446 #[test]
9447 #[cfg(target_os = "linux")]
9448 fn test_shm_no_fd_leaks() {
9449 let _lock = fd_lock_exclusive();
9450 if !is_shm_available() {
9451 crate::test_support::report_skip(&format!(
9452 "{} - SHM memory allocation not available (permission denied or no SHM support)",
9453 function!()
9454 ));
9455 return;
9456 }
9457
9458 let baseline = open_fd_count(SHM_FD_TARGETS);
9459
9460 // Self-check: the counter must actually see a live shm tensor. If
9461 // it does not — a changed backing object, a different `/proc`
9462 // rendering — the leak assertion below would hold no matter how
9463 // much leaked, so fail loudly here instead.
9464 {
9465 let probe = Tensor::<u8>::new(&[100, 100], Some(TensorMemory::Shm), None)
9466 .expect("Failed to create tensor");
9467 assert!(
9468 open_fd_count(SHM_FD_TARGETS) > baseline,
9469 "a live shm tensor must be visible to the shm fd counter, \
9470 otherwise this test cannot detect a leak"
9471 );
9472 drop(probe);
9473 }
9474
9475 for _ in 0..100 {
9476 let tensor = Tensor::<u8>::new(&[100, 100], Some(TensorMemory::Shm), None)
9477 .expect("Failed to create tensor");
9478 let mut map = tensor.map().unwrap();
9479 map.as_mut_slice().fill(233);
9480 }
9481
9482 let after = open_fd_count(SHM_FD_TARGETS);
9483 assert_eq!(
9484 baseline, after,
9485 "shm file descriptor leak detected: {baseline} -> {after}"
9486 );
9487 }
9488
9489 #[test]
9490 #[cfg(target_os = "linux")]
9491 fn test_shm_from_fd_no_fd_leaks() {
9492 let _lock = fd_lock_exclusive();
9493 if !is_shm_available() {
9494 crate::test_support::report_skip(&format!(
9495 "{} - SHM memory allocation not available (permission denied or no SHM support)",
9496 function!()
9497 ));
9498 return;
9499 }
9500
9501 let baseline = open_fd_count(SHM_FD_TARGETS);
9502 let orig = Tensor::<u8>::new(&[100, 100], Some(TensorMemory::Shm), None).unwrap();
9503
9504 // Self-check — see `test_shm_no_fd_leaks`.
9505 let held = open_fd_count(SHM_FD_TARGETS);
9506 assert!(
9507 held > baseline,
9508 "a live shm tensor must be visible to the shm fd counter, \
9509 otherwise this test cannot detect a leak"
9510 );
9511
9512 for _ in 0..100 {
9513 let tensor =
9514 Tensor::<u8>::from_fd(orig.clone_fd().unwrap(), orig.shape(), None).unwrap();
9515 let mut map = tensor.map().unwrap();
9516 map.as_mut_slice().fill(233);
9517 }
9518
9519 let after_loop = open_fd_count(SHM_FD_TARGETS);
9520 assert_eq!(
9521 held, after_loop,
9522 "shm file descriptor leak detected: {held} -> {after_loop}"
9523 );
9524
9525 // And the source tensor's own fd goes back on drop.
9526 drop(orig);
9527 let after_drop = open_fd_count(SHM_FD_TARGETS);
9528 assert_eq!(
9529 baseline, after_drop,
9530 "dropping the source tensor must release its fd: {baseline} -> {after_drop}"
9531 );
9532 }
9533
9534 #[cfg(feature = "ndarray")]
9535 #[test]
9536 fn test_ndarray() {
9537 let _lock = crate::tests::fd_lock_shared();
9538 let shape = vec![2, 3, 4];
9539 let tensor = Tensor::<f32>::new(&shape, None, None).expect("Failed to create tensor");
9540
9541 let mut tensor_map = tensor.map().expect("Failed to map tensor memory");
9542 tensor_map.fill(1.0);
9543
9544 let view = tensor_map.view().expect("Failed to get ndarray view");
9545 assert_eq!(view.shape(), &[2, 3, 4]);
9546 assert!(view.iter().all(|&x| x == 1.0));
9547
9548 let mut view_mut = tensor_map
9549 .view_mut()
9550 .expect("Failed to get mutable ndarray view");
9551 view_mut[[0, 0, 0]] = 42.0;
9552 assert_eq!(view_mut[[0, 0, 0]], 42.0);
9553 assert_eq!(tensor_map[0], 42.0, "Value at index 0 should be 42");
9554 }
9555
9556 #[test]
9557 fn test_buffer_identity_derived_equality() {
9558 // Two identities derived from the same kind and key must be equal —
9559 // that is the whole point of `derived`: two independently-linked
9560 // copies of this crate, each importing the same underlying buffer,
9561 // have to agree on its id without talking to each other.
9562 let id1 = BufferIdentity::derived(IdentityKind::DmaBuf, 4242);
9563 let id2 = BufferIdentity::derived(IdentityKind::DmaBuf, 4242);
9564 assert_eq!(
9565 id1.id(),
9566 id2.id(),
9567 "Identities derived from the same kind and key should match"
9568 );
9569
9570 // A different key still yields a different id.
9571 let id3 = BufferIdentity::derived(IdentityKind::DmaBuf, 4243);
9572 assert_ne!(
9573 id1.id(),
9574 id3.id(),
9575 "Identities derived from different keys should differ"
9576 );
9577 }
9578
9579 #[test]
9580 fn test_buffer_identity_clone_shares_guard() {
9581 let id1 = BufferIdentity::derived(IdentityKind::DmaBuf, 4242);
9582 let weak = id1.weak();
9583 assert!(
9584 weak.upgrade().is_some(),
9585 "Weak should be alive while original exists"
9586 );
9587
9588 let id2 = id1.clone();
9589 assert_eq!(id1.id(), id2.id(), "Cloned identity should have same id");
9590
9591 drop(id1);
9592 assert!(
9593 weak.upgrade().is_some(),
9594 "Weak should still be alive (clone holds Arc)"
9595 );
9596
9597 drop(id2);
9598 assert!(
9599 weak.upgrade().is_none(),
9600 "Weak should be dead after all clones dropped"
9601 );
9602 }
9603
9604 #[test]
9605 fn test_tensor_buffer_identity() {
9606 let t1 = Tensor::<u8>::new(&[100], Some(TensorMemory::Mem), Some("t1")).unwrap();
9607 let t2 = Tensor::<u8>::new(&[100], Some(TensorMemory::Mem), Some("t2")).unwrap();
9608 assert_ne!(
9609 t1.buffer_identity().id(),
9610 t2.buffer_identity().id(),
9611 "Different tensors should have different buffer ids"
9612 );
9613 }
9614
9615 // ------------------------------------------------------------------------
9616 // Quantization — constructor validation + accessor correctness.
9617 // ------------------------------------------------------------------------
9618
9619 #[test]
9620 fn test_quantization_per_tensor_constructors() {
9621 let q = Quantization::per_tensor(0.1, -5);
9622 assert!(q.is_per_tensor());
9623 assert!(!q.is_per_channel());
9624 assert!(!q.is_symmetric());
9625 assert_eq!(q.scale(), &[0.1]);
9626 assert_eq!(q.zero_point(), Some(&[-5][..]));
9627
9628 let qs = Quantization::per_tensor_symmetric(0.05);
9629 assert!(qs.is_per_tensor());
9630 assert!(qs.is_symmetric());
9631 assert_eq!(qs.zero_point(), None);
9632 }
9633
9634 #[test]
9635 fn test_quantization_per_channel_constructors() {
9636 let q = Quantization::per_channel(vec![0.1, 0.2, 0.3], vec![0, -1, 1], 2).unwrap();
9637 assert!(q.is_per_channel());
9638 assert!(!q.is_symmetric());
9639 assert_eq!(q.axis(), Some(2));
9640 assert_eq!(q.scale().len(), 3);
9641
9642 let qs = Quantization::per_channel_symmetric(vec![0.054, 0.089, 0.195], 0).unwrap();
9643 assert!(qs.is_per_channel());
9644 assert!(qs.is_symmetric());
9645 assert_eq!(qs.axis(), Some(0));
9646 }
9647
9648 #[test]
9649 fn test_quantization_per_channel_length_mismatch_rejected() {
9650 // len(scales) != len(zero_points) → rejected at construction.
9651 let err = Quantization::per_channel(vec![0.1, 0.2], vec![0, 0, 0], 0).unwrap_err();
9652 assert!(matches!(err, Error::QuantizationInvalid { .. }));
9653 }
9654
9655 #[test]
9656 fn test_quantization_per_channel_empty_rejected() {
9657 let err = Quantization::per_channel_symmetric(vec![], 0).unwrap_err();
9658 assert!(matches!(err, Error::QuantizationInvalid { .. }));
9659 }
9660
9661 /// Constructors guard scale/zero_point length invariants, but
9662 /// `Quantization` is `Deserialize`, so malformed JSON (e.g. an
9663 /// empty `scale` array, or `zero_point` length that disagrees with
9664 /// `scale`) bypasses the constructor checks. `set_quantization`
9665 /// must reject these via `validate()` so they don't poison
9666 /// downstream `mode()` selection or per-channel kernel indexing.
9667 #[test]
9668 fn test_quantization_validate_rejects_malformed_deserialize() {
9669 let mut t = Tensor::<i8>::new(&[1, 1, 4], Some(TensorMemory::Mem), None).unwrap();
9670
9671 // Empty scale array: must be rejected.
9672 let q: Quantization = serde_json::from_str(r#"{"scale": []}"#).unwrap();
9673 assert!(matches!(
9674 t.set_quantization(q).unwrap_err(),
9675 Error::QuantizationInvalid { .. }
9676 ));
9677
9678 // Per-tensor with multi-element zero_point: must be rejected.
9679 let q: Quantization =
9680 serde_json::from_str(r#"{"scale": 0.1, "zero_point": [0, 0, 0]}"#).unwrap();
9681 assert!(matches!(
9682 t.set_quantization(q).unwrap_err(),
9683 Error::QuantizationInvalid { .. }
9684 ));
9685
9686 // Per-channel zero_point length != scale length: must be rejected.
9687 let q: Quantization = serde_json::from_str(
9688 r#"{"scale": [0.1, 0.2, 0.3, 0.4], "zero_point": [0, 0], "axis": 2}"#,
9689 )
9690 .unwrap();
9691 assert!(matches!(
9692 t.set_quantization(q).unwrap_err(),
9693 Error::QuantizationInvalid { .. }
9694 ));
9695 }
9696
9697 #[test]
9698 fn test_quantization_mode_dispatch() {
9699 let pt = Quantization::per_tensor(0.1, -5);
9700 assert!(matches!(
9701 pt.mode(),
9702 QuantMode::PerTensor { scale, zero_point } if scale == 0.1 && zero_point == -5
9703 ));
9704
9705 let pts = Quantization::per_tensor_symmetric(0.05);
9706 assert!(matches!(
9707 pts.mode(),
9708 QuantMode::PerTensorSymmetric { scale } if scale == 0.05
9709 ));
9710
9711 let pc = Quantization::per_channel(vec![0.1, 0.2], vec![0, -1], 2).unwrap();
9712 assert!(matches!(pc.mode(), QuantMode::PerChannel { axis: 2, .. }));
9713
9714 let pcs = Quantization::per_channel_symmetric(vec![0.1, 0.2], 0).unwrap();
9715 assert!(matches!(
9716 pcs.mode(),
9717 QuantMode::PerChannelSymmetric { axis: 0, .. }
9718 ));
9719 }
9720
9721 #[test]
9722 fn test_tensor_quantization_roundtrip_integer() {
9723 let mut t = Tensor::<i8>::new(&[2, 3, 4], Some(TensorMemory::Mem), None).unwrap();
9724 assert!(t.quantization().is_none());
9725 t.set_quantization(Quantization::per_tensor(0.1, -5))
9726 .unwrap();
9727 let q = t.quantization().unwrap();
9728 assert_eq!(q.scale(), &[0.1]);
9729 t.clear_quantization();
9730 assert!(t.quantization().is_none());
9731 }
9732
9733 #[test]
9734 fn test_tensor_with_quantization_builder() {
9735 let t = Tensor::<i8>::new(&[4, 4], Some(TensorMemory::Mem), None)
9736 .unwrap()
9737 .with_quantization(Quantization::per_tensor_symmetric(0.05))
9738 .unwrap();
9739 assert!(t.quantization().is_some());
9740 }
9741
9742 #[test]
9743 fn test_tensor_dyn_quantization_float_arm_returns_none() {
9744 let t = Tensor::<f32>::new(&[2, 2], Some(TensorMemory::Mem), None).unwrap();
9745 let td = TensorDyn::F32(t);
9746 assert!(td.quantization().is_none());
9747 }
9748
9749 #[test]
9750 fn test_tensor_dyn_set_quantization_float_arm_errors() {
9751 let t = Tensor::<f32>::new(&[2, 2], Some(TensorMemory::Mem), None).unwrap();
9752 let mut td = TensorDyn::F32(t);
9753 let err = td
9754 .set_quantization(Quantization::per_tensor(0.1, 0))
9755 .unwrap_err();
9756 // float path returns a QuantizationInvalid error.
9757 assert!(matches!(err, Error::QuantizationInvalid { .. }));
9758 }
9759
9760 /// Compile-time type gate — calling `Tensor::<f32>::quantization()` must
9761 /// fail to compile (the `IntegerType` trait bound is not satisfied by
9762 /// `f32`). This doctest anchors the invariant.
9763 ///
9764 /// ```compile_fail
9765 /// use edgefirst_tensor::{Tensor, TensorMemory};
9766 /// let t = Tensor::<f32>::new(&[2, 2], Some(TensorMemory::Mem), None).unwrap();
9767 /// let _ = t.quantization(); // compile error: f32 not IntegerType
9768 /// ```
9769 fn _compile_fail_doctest_anchor() {}
9770
9771 // Any test that cares about the fd count must grab it exclusively.
9772 // Any tests which modifies the fd count by opening or closing fds must grab it
9773 // shared.
9774 pub static FD_LOCK: RwLock<()> = RwLock::new(());
9775
9776 /// Take [`FD_LOCK`] shared, tolerating poisoning.
9777 ///
9778 /// A test that panics while holding the lock poisons it, and every later
9779 /// `.unwrap()` on it panics too — one real failure turns into a handful
9780 /// of unrelated ones that bury it. The guarded data is `()`, so a panic
9781 /// cannot have left an invariant broken and there is nothing to protect
9782 /// against by refusing the lock.
9783 pub fn fd_lock_shared() -> std::sync::RwLockReadGuard<'static, ()> {
9784 FD_LOCK.read().unwrap_or_else(PoisonError::into_inner)
9785 }
9786
9787 /// Take [`FD_LOCK`] exclusively. Poison-tolerant, see [`fd_lock_shared`].
9788 ///
9789 /// Linux-only, like all four of its callers: the fd-leak tests read fd
9790 /// state out of `/proc`, which no other platform provides. Ungated, it
9791 /// is dead code on macOS and the lint denies that.
9792 /// [`fd_lock_shared`] stays available everywhere — macOS tests declare
9793 /// fd activity too.
9794 #[cfg(target_os = "linux")]
9795 pub fn fd_lock_exclusive() -> std::sync::RwLockWriteGuard<'static, ()> {
9796 FD_LOCK.write().unwrap_or_else(PoisonError::into_inner)
9797 }
9798
9799 /// `/proc` fd-target fragments identifying a POSIX shared memory
9800 /// segment. `shm_open` then `shm_unlink` leaves the target rendered as
9801 /// `/dev/shm/<name> (deleted)`.
9802 #[cfg(target_os = "linux")]
9803 const SHM_FD_TARGETS: &[&str] = &["/dev/shm/"];
9804
9805 /// `/proc` fd-target fragments identifying the fds a DMA tensor owns.
9806 ///
9807 /// Both the exported buffer and the heap it came from are listed, so a
9808 /// leak of either is still caught — narrowing this to the buffer alone
9809 /// would silently stop covering leaked heap handles.
9810 #[cfg(target_os = "linux")]
9811 const DMA_FD_TARGETS: &[&str] = &["dmabuf", "dma_heap"];
9812
9813 /// How many of this process's open fds point at an object whose `/proc`
9814 /// target contains `needle`.
9815 ///
9816 /// The fd-leak tests deliberately do NOT compare total fd counts. Under
9817 /// `cargo test` every test in this crate runs as a thread in one
9818 /// process, and roughly twenty of them create shm or DMA tensors —
9819 /// while anything else may open a file, a socket, or a pipe — so an
9820 /// unrelated test opening or closing an fd inside the measurement
9821 /// window moved the total and failed the assertion. Worse, the panic
9822 /// poisoned [`FD_LOCK`] and took four more tests down with it, so the
9823 /// visible failure had little to do with the cause.
9824 ///
9825 /// Counting only fds that name the object under test makes the
9826 /// measurement immune to every unrelated fd, while still catching the
9827 /// leak each test exists to catch. (CI runs nextest, one process per
9828 /// test, which is why this only ever flaked when run locally.)
9829 #[cfg(target_os = "linux")]
9830 fn open_fd_count(needles: &[&str]) -> usize {
9831 let proc = procfs::process::Process::myself()
9832 .expect("Failed to get current process using /proc/self");
9833 proc.fd()
9834 .expect("Failed to list open file descriptors")
9835 .filter_map(Result::ok)
9836 .filter(|entry| {
9837 let target = format!("{:?}", entry.target);
9838 needles.iter().any(|needle| target.contains(needle))
9839 })
9840 .count()
9841 }
9842
9843 /// Test that DMA is NOT available on non-Linux platforms.
9844 /// This verifies the cross-platform behavior of is_dma_available().
9845 #[test]
9846 #[cfg(not(target_os = "linux"))]
9847 fn test_dma_not_available_on_non_linux() {
9848 assert!(
9849 !is_dma_available(),
9850 "DMA memory allocation should NOT be available on non-Linux platforms"
9851 );
9852 }
9853
9854 #[test]
9855 fn colorimetry_defaults_none_and_roundtrips_without_auto_fill() {
9856 use crate::{ColorEncoding, ColorRange, Colorimetry, PixelFormat, TensorMemory};
9857 let mut t = Tensor::<u8>::image(
9858 1280,
9859 720,
9860 PixelFormat::Nv12,
9861 Some(TensorMemory::Mem),
9862 crate::CpuAccess::ReadWrite,
9863 )
9864 .unwrap();
9865 assert_eq!(t.colorimetry(), None); // default undefined
9866 let c = Colorimetry::default()
9867 .with_encoding(ColorEncoding::Bt709)
9868 .with_range(ColorRange::Limited);
9869 t.set_colorimetry(Some(c));
9870 assert_eq!(t.colorimetry(), Some(c));
9871 // configure_image must NOT touch colorimetry
9872 t.configure_image(640, 480, PixelFormat::Grey).unwrap();
9873 assert_eq!(t.colorimetry(), Some(c));
9874 }
9875
9876 #[test]
9877 fn configure_image_within_capacity() {
9878 let mut t = Tensor::<u8>::image_with_capacity(
9879 640,
9880 480,
9881 PixelFormat::Rgb,
9882 None,
9883 crate::CpuAccess::ReadWrite,
9884 )
9885 .unwrap();
9886 t.configure_image(320, 240, PixelFormat::Nv12).unwrap();
9887 assert_eq!(t.format(), Some(PixelFormat::Nv12));
9888 assert_eq!(t.width(), Some(320));
9889 assert_eq!(t.height(), Some(240));
9890 assert_eq!(t.shape(), &[360, 320]); // 240*3/2
9891 }
9892
9893 #[test]
9894 fn configure_image_too_large_errors() {
9895 let mut t = Tensor::<u8>::image_with_capacity(
9896 64,
9897 64,
9898 PixelFormat::Grey,
9899 None,
9900 crate::CpuAccess::ReadWrite,
9901 )
9902 .unwrap();
9903 let err = t
9904 .configure_image(1920, 1080, PixelFormat::Nv12)
9905 .unwrap_err();
9906 assert!(matches!(err, Error::InsufficientCapacity { .. }));
9907 }
9908
9909 /// A reused max-sized IOSurface pool keeps its physical `bytesPerRow` when
9910 /// reconfigured to a smaller logical image (physical-grid / logical-ROI
9911 /// decoupling), instead of collapsing to the frame's natural row stride.
9912 #[test]
9913 #[cfg(target_os = "macos")]
9914 fn configure_image_preserves_iosurface_physical_stride() {
9915 // Declares this test as an fd-opener; see FD_LOCK.
9916 let _lock = crate::tests::fd_lock_shared();
9917 // Pool: GREY/R8 IOSurface 100 wide → bytesPerRow padded to 128.
9918 let mut pool = Tensor::<u8>::image(
9919 100,
9920 64,
9921 PixelFormat::Grey,
9922 Some(TensorMemory::DmaBuf),
9923 crate::CpuAccess::ReadWrite,
9924 )
9925 .unwrap();
9926 let pitch = pool.effective_row_stride().unwrap();
9927 assert!(
9928 pitch >= 128 && pitch.is_multiple_of(64),
9929 "padded bytesPerRow, got {pitch}"
9930 );
9931
9932 // Reconfigure to a smaller NV12 frame; the physical pitch must survive
9933 // (natural would be 32, but the surface stride is the 128-padded pitch).
9934 pool.configure_image(32, 16, PixelFormat::Nv12).unwrap();
9935 assert_eq!(pool.format(), Some(PixelFormat::Nv12));
9936 assert_eq!(pool.width(), Some(32));
9937 assert_eq!(pool.height(), Some(16));
9938 assert_eq!(
9939 pool.effective_row_stride(),
9940 Some(pitch),
9941 "configure_image must preserve the IOSurface physical bytesPerRow"
9942 );
9943
9944 // Reconfigure again to NV24 — pitch still preserved.
9945 pool.configure_image(32, 16, PixelFormat::Nv24).unwrap();
9946 assert_eq!(pool.effective_row_stride(), Some(pitch));
9947 }
9948
9949 /// `configure_image` on a Mem backing reconfigures to the format's
9950 /// **64-byte-aligned** row stride (the odd-dim contract: every image tensor
9951 /// carries a 64-aligned `row_stride`). For NV12 32×16 the minimum is
9952 /// `even(32)=32`, rounded up to the 64-byte alignment → 64. The capacity
9953 /// (64×64×4 RGBA = 16 KiB) easily holds the 24×64 = 1.5 KiB NV12 layout.
9954 #[test]
9955 fn configure_image_mem_aligns_stride() {
9956 let mut t = Tensor::<u8>::image_with_capacity(
9957 64,
9958 64,
9959 PixelFormat::Rgba,
9960 Some(TensorMemory::Mem),
9961 crate::CpuAccess::ReadWrite,
9962 )
9963 .unwrap();
9964 t.configure_image(32, 16, PixelFormat::Nv12).unwrap();
9965 let s = t.effective_row_stride().unwrap();
9966 assert_eq!(s % 64, 0, "stride must be 64-aligned");
9967 assert!(s >= 32, "stride must cover the even-width minimum");
9968 assert_eq!(s, 64);
9969 }
9970
9971 #[test]
9972 fn strided_mem_tensor_cpu_maps_full_padded_buffer() {
9973 // A packed RGBA image with row padding (GPU-pitch style): logical width
9974 // 8 px (32 B/row) but a 48-byte row stride. Over-allocate capacity (for
9975 // 16 px), narrow the logical width, then record the padded stride.
9976 // Previously `map()` rejected this on non-Linux with
9977 // "DMA backing is Linux-only"; HAL-owned Mem is now mappable.
9978 let mut t = Tensor::<u8>::image_with_capacity(
9979 16,
9980 3,
9981 PixelFormat::Rgba,
9982 Some(TensorMemory::Mem),
9983 crate::CpuAccess::ReadWrite,
9984 )
9985 .unwrap(); // capacity 3 × 16 × 4 = 192 B
9986 t.configure_image(8, 3, PixelFormat::Rgba).unwrap(); // logical [3, 8, 4] = 96 B
9987 t.set_row_stride(48).unwrap(); // padded stride (>= 32 B min)
9988
9989 let map = t.map().expect("strided Mem tensor should CPU-map");
9990 // Full padded buffer (stride 48 × 3 rows = 144 B), not the 96 B logical
9991 // view — callers iterate rows via `effective_row_stride()`.
9992 assert_eq!(map.as_slice().len(), 144);
9993 // Logical shape is still reported for shape-aware consumers.
9994 assert_eq!(map.shape(), &[3, 8, 4]);
9995 }
9996
9997 #[test]
9998 fn strided_mem_tensor_over_capacity_errors() {
9999 // Stride larger than the allocation: 64 B × 3 rows = 192 B > 96 B cap.
10000 let mut t = Tensor::<u8>::new(&[3, 8, 4], Some(TensorMemory::Mem), None).unwrap();
10001 t.set_format(PixelFormat::Rgba).unwrap();
10002 t.set_row_stride(64).unwrap();
10003 assert!(matches!(t.map(), Err(Error::InsufficientCapacity { .. })));
10004 }
10005
10006 /// Test that SHM memory allocation is available and usable on Unix systems.
10007 /// This is a basic functional test; Linux has additional FD leak tests using procfs.
10008 #[test]
10009 #[cfg(unix)]
10010 fn test_shm_available_and_usable() {
10011 // Declares this test as an fd-opener; see FD_LOCK.
10012 let _lock = crate::tests::fd_lock_shared();
10013 // "Unix" is not one platform here: Android is unix but its bionic libc
10014 // has no POSIX shared memory, so allocation legitimately reports
10015 // NotImplemented (see shm.rs module docs). Asserting availability
10016 // outright fails there for a documented reason.
10017 if !is_shm_available() {
10018 crate::test_support::report_skip(
10019 "test_shm_available_and_usable - SHM allocation \
10020 unavailable on this platform",
10021 );
10022 return;
10023 }
10024
10025 // Create a tensor with SHM backing
10026 let tensor = Tensor::<u8>::new(&[100, 100], Some(TensorMemory::Shm), None)
10027 .expect("Failed to create SHM tensor");
10028
10029 // Verify we can map and write to it
10030 let mut map = tensor.map().expect("Failed to map SHM tensor");
10031 map.as_mut_slice().fill(0xAB);
10032
10033 // Verify the data was written correctly
10034 assert!(
10035 map.as_slice().iter().all(|&b| b == 0xAB),
10036 "SHM tensor data should be writable and readable"
10037 );
10038 }
10039
10040 // =========================================================================
10041 // packed_rgba16f_layout — host-runnable geometry unit tests (TDD)
10042 // =========================================================================
10043
10044 #[test]
10045 fn packed_rgba16f_layout_planar_rgb_f16() {
10046 let layout =
10047 packed_rgba16f_layout(PixelFormat::PlanarRgb, DType::F16, 640, 640).expect("Some");
10048 assert_eq!(layout.surface_w, 160);
10049 assert_eq!(layout.surface_h, 1920);
10050 assert_eq!(layout.bytes_per_texel, 8);
10051 assert_eq!(layout.pitch, 1280);
10052 }
10053
10054 #[test]
10055 fn packed_rgba16f_layout_planar_rgba_f16() {
10056 let layout =
10057 packed_rgba16f_layout(PixelFormat::PlanarRgba, DType::F16, 640, 640).expect("Some");
10058 assert_eq!(layout.surface_w, 160);
10059 assert_eq!(layout.surface_h, 2560); // 4 planes
10060 assert_eq!(layout.bytes_per_texel, 8);
10061 assert_eq!(layout.pitch, 1280);
10062 }
10063
10064 #[test]
10065 fn packed_rgba16f_layout_rejects_misaligned() {
10066 assert!(packed_rgba16f_layout(PixelFormat::PlanarRgb, DType::F16, 642, 640).is_none());
10067 }
10068
10069 #[test]
10070 fn packed_rgba16f_layout_rejects_non_f16() {
10071 // Non-F16 dtype with planar RGB
10072 assert!(packed_rgba16f_layout(PixelFormat::PlanarRgb, DType::U8, 640, 640).is_none());
10073 // Non-planar format with F32
10074 assert!(packed_rgba16f_layout(PixelFormat::Rgb, DType::F32, 640, 640).is_none());
10075 // Packed Rgba with F16 is not a planar format → None
10076 assert!(packed_rgba16f_layout(PixelFormat::Rgba, DType::F16, 640, 640).is_none());
10077 }
10078
10079 #[test]
10080 fn cuda_map_fast_fails_to_none_without_handle() {
10081 let t = Tensor::<f32>::new(&[4], Some(TensorMemory::Mem), None).unwrap();
10082 assert!(t.cuda().is_none());
10083 assert!(t.cuda_map().is_none()); // pure local check, no GL routing
10084 }
10085
10086 #[test]
10087 fn cuda_returns_none_without_handle() {
10088 // A plain Mem-backed tensor has no CUDA handle attached.
10089 let t = Tensor::<f32>::new(&[2, 2], Some(TensorMemory::Mem), None).unwrap();
10090 assert!(t.cuda().is_none(), "no CUDA handle on a Mem tensor");
10091 assert!(t.cuda_map().is_none(), "fast-fail map → None");
10092 }
10093
10094 #[test]
10095 fn cuda_map_then_host_map_fallback() {
10096 // The documented client pattern: try cuda_map() first; when it is None
10097 // (no CUDA handle — the case for a plain Mem tensor), fall back to map().
10098 let t = Tensor::<f32>::new(&[2, 2], Some(TensorMemory::Mem), None).unwrap();
10099 // Bind to a named variable so the CudaMap guard (and its borrow of `t`)
10100 // is dropped at the end of this statement, before the else branch borrows `t` again.
10101 let cuda = t.cuda_map();
10102 if let Some(_c) = cuda {
10103 // On a CUDA-registered tensor we'd use the device ptr here.
10104 unreachable!("a Mem tensor has no CUDA handle");
10105 } else {
10106 let host = t.map().expect("host map fallback must succeed");
10107 // TensorMapTrait::len() returns the element count (not bytes).
10108 assert_eq!(host.len(), 4); // 2*2 f32 elements
10109 }
10110 }
10111
10112 // -------------------------------------------------------------------------
10113 // Tensor::from_foreign — public API tests at the Tensor<T> layer.
10114 //
10115 // The low-level MemTensor::from_foreign mechanics (owner-drop, view sharing)
10116 // are covered in mem.rs. These tests exercise the Tensor<T> guard paths
10117 // (null ptr, empty shape, size overflow) and the basic wrap+readback
10118 // contract, confirming the public unsafe API wires through correctly.
10119 // -------------------------------------------------------------------------
10120
10121 #[test]
10122 fn from_foreign_valid_wrap_and_readback() {
10123 // The canonical CUDA zero-copy shape: wrap a caller allocation as a
10124 // Mem tensor and verify the tensor reads the exact same bytes.
10125 let mut buf: Vec<f32> = vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0];
10126 let ptr = buf.as_mut_ptr();
10127 let t = unsafe { Tensor::<f32>::from_foreign(ptr, &[2, 3], None, Some("test_foreign")) }
10128 .expect("valid from_foreign must succeed");
10129 assert_eq!(t.shape(), &[2, 3]);
10130 assert_eq!(t.memory(), TensorMemory::Mem);
10131 assert_eq!(t.name(), "test_foreign");
10132 let m = t.map().unwrap();
10133 // The tensor is a zero-copy borrow — it sees the caller's data.
10134 assert_eq!(m.as_slice(), &[1.0_f32, 2.0, 3.0, 4.0, 5.0, 6.0]);
10135 }
10136
10137 #[test]
10138 fn from_foreign_write_visible_in_caller_allocation() {
10139 // Writes through the tensor's map land in the caller's buffer (zero-copy).
10140 let mut buf: Vec<u8> = vec![0u8; 6];
10141 let ptr = buf.as_mut_ptr();
10142 let t = unsafe { Tensor::<u8>::from_foreign(ptr, &[2, 3], None, None) }.unwrap();
10143 {
10144 let mut m = t.map().unwrap();
10145 m.as_mut_slice().copy_from_slice(&[10, 20, 30, 40, 50, 60]);
10146 }
10147 drop(t);
10148 // Mutations are visible in the original Vec — same physical buffer.
10149 assert_eq!(buf, vec![10, 20, 30, 40, 50, 60]);
10150 }
10151
10152 #[test]
10153 fn from_foreign_rejects_null_ptr() {
10154 let err = unsafe { Tensor::<u8>::from_foreign(std::ptr::null_mut(), &[4], None, None) }
10155 .unwrap_err();
10156 assert!(
10157 matches!(err, Error::InvalidArgument(ref m) if m.contains("non-null")),
10158 "expected InvalidArgument(non-null), got {err:?}"
10159 );
10160 }
10161
10162 #[test]
10163 fn from_foreign_rejects_empty_shape() {
10164 let mut dummy: u8 = 0;
10165 let err = unsafe { Tensor::<u8>::from_foreign(&mut dummy, &[], None, None) }.unwrap_err();
10166 assert!(
10167 matches!(err, Error::InvalidSize(0)),
10168 "expected InvalidSize(0) for empty shape, got {err:?}"
10169 );
10170 }
10171
10172 #[test]
10173 fn from_foreign_rejects_overflow_shape() {
10174 // Two dimensions whose product overflows usize — the overflow guard must
10175 // fire before any pointer arithmetic is attempted.
10176 let mut dummy: u8 = 0;
10177 let huge = [usize::MAX / 2 + 1, 2];
10178 let err = unsafe { Tensor::<u8>::from_foreign(&mut dummy, &huge, None, None) }.unwrap_err();
10179 assert!(
10180 matches!(err, Error::InvalidArgument(ref m) if m.contains("overflow")),
10181 "expected InvalidArgument(overflow), got {err:?}"
10182 );
10183 }
10184
10185 #[test]
10186 fn from_foreign_owner_keeps_allocation_alive() {
10187 // When `owner` is `Some`, dropping the Tensor must not free the backing
10188 // before the owner is also gone — the owner's Drop fires on last ref.
10189 use std::sync::atomic::{AtomicBool, Ordering};
10190 let flag = std::sync::Arc::new(AtomicBool::new(false));
10191 let flag2 = flag.clone();
10192 struct Guard(std::sync::Arc<AtomicBool>);
10193 impl Drop for Guard {
10194 fn drop(&mut self) {
10195 self.0.store(true, Ordering::SeqCst);
10196 }
10197 }
10198 let mut buf: Vec<u32> = vec![42u32; 4];
10199 let ptr = buf.as_mut_ptr();
10200 let owner: ForeignOwner = Box::new(Guard(flag2));
10201 let t = unsafe { Tensor::<u32>::from_foreign(ptr, &[4], Some(owner), None) }.unwrap();
10202 // Map co-owns the backing Arc; the owner must stay alive while the map lives.
10203 let m = t.map().unwrap();
10204 assert_eq!(m.as_slice()[0], 42);
10205 drop(t); // tensor dropped while map is still live
10206 assert!(
10207 !flag.load(Ordering::SeqCst),
10208 "owner must not drop while a map shares the backing"
10209 );
10210 drop(m);
10211 assert!(
10212 flag.load(Ordering::SeqCst),
10213 "owner Drop must fire when the last Arc reference is released"
10214 );
10215 }
10216}