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