#[cfg(all(feature = "static", feature = "dynamic"))]
compile_error!(
"edgefirst-tensor: `static` and `dynamic` are mutually exclusive. \
A consumer building the C-API leaves wants `--no-default-features \
--features dynamic`; everything else wants the default."
);
#[cfg(not(any(feature = "static", feature = "dynamic")))]
compile_error!(
"edgefirst-tensor: exactly one backend must be enabled. This crate has \
`default-features = false` at the workspace level, so every consumer \
must name one explicitly: `features = [\"static\"]` for Rust consumers, \
or `[\"dynamic\"]` for the C-API leaves."
);
#[cfg(all(target_os = "android", feature = "static"))]
mod ahardwarebuffer;
#[allow(dead_code)]
mod ahardwarebuffer_layout;
pub mod colorimetry;
#[cfg(feature = "tracing")]
pub mod trace;
pub mod covguard;
mod cuda;
mod detect;
#[cfg(all(target_os = "linux", feature = "static"))]
mod dma;
#[cfg(all(target_os = "linux", feature = "static"))]
mod dmabuf;
mod error;
mod format;
#[cfg(target_os = "windows")]
pub mod d3d11;
pub mod d3d11_layout;
#[cfg(any(target_os = "macos", target_os = "ios"))]
mod iosurface_layout;
#[cfg(all(any(target_os = "macos", target_os = "ios"), feature = "static"))]
mod iosurface;
pub mod lens;
#[cfg(feature = "static")]
mod mem;
mod pbo;
pub mod pin;
pub mod protocol;
#[cfg(all(unix, feature = "static"))]
mod shm;
mod tensor_dyn;
#[cfg(test)]
mod test_support;
pub mod view;
mod vocabulary;
pub use colorimetry::{
ColorEncoding, ColorRange, ColorSpace, ColorTransfer, Colorimetry, MatrixWeights, RangeScaling,
};
#[cfg(all(coverage, target_os = "linux"))]
#[used]
#[link_section = ".init_array"]
static __EDGEFIRST_COV_INSTALL: extern "C" fn() = {
extern "C" fn ctor() {
crate::covguard::install();
}
ctor
};
#[cfg(all(target_os = "android", feature = "static"))]
pub use crate::ahardwarebuffer::image_ahardwarebuffer_layout;
#[cfg(all(target_os = "android", feature = "static"))]
pub(crate) use crate::ahardwarebuffer::AHardwareBufferTensor;
#[cfg(all(target_os = "linux", feature = "static"))]
pub(crate) use crate::dma::DmaTensor;
#[cfg(all(any(target_os = "macos", target_os = "ios"), feature = "static"))]
pub(crate) use crate::iosurface::IoSurfaceTensor;
#[cfg(target_os = "windows")]
pub use crate::d3d11::geometry::{d3d11_shared_handle_geometry, d3d11_texture_geometry};
#[cfg(any(target_os = "macos", target_os = "ios"))]
pub use crate::iosurface_layout::image_iosurface_layout;
#[cfg(feature = "static")]
pub(crate) use crate::mem::MemTensor;
pub use crate::pbo::{
client_state_pbo_ops, read_pbo_vtable_parts, EfClientState, EfPboMapFn, EfPboUnmapFn,
PboClientParts, PboMapping, PboOps, PboOpsVtableParts, PboTensor,
};
#[cfg(all(unix, feature = "static"))]
pub(crate) use crate::shm::ShmTensor;
pub use cuda::{
client_state_cuda_ops, gl_map_resource, gl_register_buffer, gl_unmap_resource,
gl_unregister_resource, is_cuda_available, memcpy_device_to_host, memcpy_host_to_device,
runtime_path as cuda_runtime_path, stream_create, stream_destroy, stream_synchronize,
CudaGlOps, CudaHandle, CudaMap, CudaStream,
};
pub use detect::{unletter_norm, BoundingBox, DetectBox, ProtoData, ProtoLayout, Segmentation};
pub use error::{Error, Result};
#[cfg(feature = "static")]
pub mod blob;
#[cfg(feature = "static")]
pub use blob::{
BlobError, BlobHeader, BlobPlane, BlobStrings, BlobView, Regions, TransportMode, HEADER_LEN,
PLANE_RECORD_LEN,
};
pub use format::{ChromaLayout, PixelFormat, PixelLayout, PlaneGeometry};
pub use lens::Element;
use num_traits::Num;
pub use pin::HostPin;
pub use protocol::{
dtype as tensor_dtype, dtype_to_dtype, format as tensor_format, format_from_code,
kind as tensor_kind, SendPtr, TensorDesc, ABI_VERSION,
};
use serde::{Deserialize, Serialize};
#[cfg(unix)]
use std::os::fd::OwnedFd;
use std::{
fmt,
sync::{
atomic::{AtomicU64, Ordering},
Arc, Weak,
},
};
#[cfg(feature = "dynamic")]
pub use tensor_dyn::Tensor;
pub use tensor_dyn::{Raw, TensorDyn};
pub use view::HostView;
pub use vocabulary::vocabulary_demo;
pub type ForeignOwner = Box<dyn std::any::Any + Send + Sync>;
pub use half::f16;
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct PackedRgba16fLayout {
pub surface_w: usize,
pub surface_h: usize,
pub bytes_per_texel: usize,
pub pitch: usize,
}
pub fn packed_rgba16f_layout(
format: PixelFormat,
dtype: DType,
width: usize,
height: usize,
) -> Option<PackedRgba16fLayout> {
if dtype != DType::F16 {
return None;
}
let planes: usize = match format {
PixelFormat::PlanarRgb => 3,
PixelFormat::PlanarRgba => 4,
_ => return None,
};
if !width.is_multiple_of(4) {
return None;
}
let surface_w = width / 4;
let surface_h = planes.checked_mul(height)?;
let bytes_per_texel = 8;
let pitch = surface_w.checked_mul(bytes_per_texel)?;
Some(PackedRgba16fLayout {
surface_w,
surface_h,
bytes_per_texel,
pitch,
})
}
pub fn packed_rgb888_layout(width: usize, height: usize) -> Option<PackedRgb888Layout> {
if !width.is_multiple_of(4) {
return None;
}
let row_bytes = width.checked_mul(3)?;
let surface_w = row_bytes / 4;
Some(PackedRgb888Layout {
surface_w,
surface_h: height,
bytes_per_texel: 4,
pitch: row_bytes,
})
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct PackedRgb888Layout {
pub surface_w: usize,
pub surface_h: usize,
pub bytes_per_texel: usize,
pub pitch: usize,
}
pub fn image_dims_from_shape(format: PixelFormat, shape: &[usize]) -> Result<(usize, usize)> {
let wrong_rank = |expected: usize| {
Error::InvalidShape(format!(
"image_dims_from_shape: {format:?} describes a rank-{expected} allocation \
shape, got rank {} ({shape:?})",
shape.len()
))
};
match format.layout() {
PixelLayout::Packed => {
if shape.len() != 3 {
return Err(wrong_rank(3));
}
Ok((shape[1], shape[0]))
}
PixelLayout::Planar => {
if shape.len() != 3 {
return Err(wrong_rank(3));
}
Ok((shape[2], shape[1]))
}
PixelLayout::SemiPlanar => {
if shape.len() != 2 {
return Err(wrong_rank(2));
}
let height = format.image_height_from_combined(shape[0]).ok_or_else(|| {
Error::InvalidShape(format!(
"image_dims_from_shape: no {format:?} image height has combined plane \
height {} (shape {shape:?})",
shape[0]
))
})?;
Ok((shape[1], height))
}
}
}
#[cfg(unix)]
pub struct PlaneDescriptor {
fd: OwnedFd,
stride: Option<usize>,
offset: Option<usize>,
}
#[cfg(unix)]
impl PlaneDescriptor {
pub fn new(fd: std::os::fd::BorrowedFd<'_>) -> Result<Self> {
let owned = fd.try_clone_to_owned()?;
Ok(Self {
fd: owned,
stride: None,
offset: None,
})
}
pub fn with_stride(mut self, stride: usize) -> Self {
self.stride = Some(stride);
self
}
pub fn with_offset(mut self, offset: usize) -> Self {
self.offset = Some(offset);
self
}
pub fn into_fd(self) -> OwnedFd {
self.fd
}
pub fn stride(&self) -> Option<usize> {
self.stride
}
pub fn offset(&self) -> Option<usize> {
self.offset
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct Region {
pub x: usize,
pub y: usize,
pub width: usize,
pub height: usize,
}
impl Region {
pub fn new(x: usize, y: usize, width: usize, height: usize) -> Self {
Self {
x,
y,
width,
height,
}
}
pub fn fits_within(&self, width: usize, height: usize) -> bool {
self.x.saturating_add(self.width) <= width && self.y.saturating_add(self.height) <= height
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct ViewOrigin {
pub parent_width: usize,
pub parent_height: usize,
pub parent_row_stride: usize,
pub x: usize,
pub y: usize,
}
crate::ef_vocabulary! {
#[derive(Serialize, Deserialize)]
#[non_exhaustive]
pub enum DType {
U8 = 0, "u8", U8,
I8 = 1, "i8", I8,
U16 = 2, "u16", U16,
I16 = 3, "i16", I16,
U32 = 4, "u32", U32,
I32 = 5, "i32", I32,
U64 = 6, "u64", U64,
I64 = 7, "i64", I64,
F16 = 8, "f16", F16,
F32 = 9, "f32", F32,
F64 = 10, "f64", F64,
}
#[doc(hidden)]
pub mod dtype_wire;
}
impl DType {
pub const fn size(&self) -> usize {
match self {
Self::U8 | Self::I8 => 1,
Self::U16 | Self::I16 | Self::F16 => 2,
Self::U32 | Self::I32 | Self::F32 => 4,
Self::U64 | Self::I64 | Self::F64 => 8,
}
}
pub fn name(&self) -> &'static str {
self.as_str()
}
}
impl fmt::Display for DType {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.write_str(self.name())
}
}
#[cfg(feature = "static")]
pub(crate) fn dtype_of<T: 'static>() -> Option<DType> {
use std::any::TypeId;
let id = TypeId::of::<T>();
if id == TypeId::of::<u8>() {
Some(DType::U8)
} else if id == TypeId::of::<i8>() {
Some(DType::I8)
} else if id == TypeId::of::<u16>() {
Some(DType::U16)
} else if id == TypeId::of::<i16>() {
Some(DType::I16)
} else if id == TypeId::of::<u32>() {
Some(DType::U32)
} else if id == TypeId::of::<i32>() {
Some(DType::I32)
} else if id == TypeId::of::<u64>() {
Some(DType::U64)
} else if id == TypeId::of::<i64>() {
Some(DType::I64)
} else if id == TypeId::of::<half::f16>() {
Some(DType::F16)
} else if id == TypeId::of::<f32>() {
Some(DType::F32)
} else if id == TypeId::of::<f64>() {
Some(DType::F64)
} else {
None
}
}
mod sealed {
pub trait Sealed {}
impl Sealed for u8 {}
impl Sealed for i8 {}
impl Sealed for u16 {}
impl Sealed for i16 {}
impl Sealed for u32 {}
impl Sealed for i32 {}
impl Sealed for u64 {}
impl Sealed for i64 {}
impl Sealed for half::f16 {}
impl Sealed for f32 {}
impl Sealed for f64 {}
}
pub trait IntegerType: sealed::Sealed {}
impl IntegerType for u8 {}
impl IntegerType for i8 {}
impl IntegerType for u16 {}
impl IntegerType for i16 {}
impl IntegerType for u32 {}
impl IntegerType for i32 {}
impl IntegerType for u64 {}
impl IntegerType for i64 {}
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct Quantization {
#[serde(deserialize_with = "deserialize_scalar_or_vec_f32")]
scale: Vec<f32>,
#[serde(
default,
deserialize_with = "deserialize_opt_scalar_or_vec_i32",
skip_serializing_if = "Option::is_none"
)]
zero_point: Option<Vec<i32>>,
#[serde(default, skip_serializing_if = "Option::is_none")]
axis: Option<usize>,
}
#[derive(Debug, Clone, Copy)]
pub enum QuantMode<'a> {
PerTensorSymmetric {
scale: f32,
},
PerTensor {
scale: f32,
zero_point: i32,
},
PerChannelSymmetric {
scales: &'a [f32],
axis: usize,
},
PerChannel {
scales: &'a [f32],
zero_points: &'a [i32],
axis: usize,
},
}
impl Quantization {
pub fn per_tensor_symmetric(scale: f32) -> Self {
Self {
scale: vec![scale],
zero_point: None,
axis: None,
}
}
pub fn per_tensor(scale: f32, zero_point: i32) -> Self {
Self {
scale: vec![scale],
zero_point: Some(vec![zero_point]),
axis: None,
}
}
pub fn per_channel_symmetric(scales: Vec<f32>, axis: usize) -> Result<Self> {
if scales.is_empty() {
return Err(Error::QuantizationInvalid {
field: "scale.len",
expected: "non-empty per-channel scales".to_string(),
got: "length 0".to_string(),
});
}
Ok(Self {
scale: scales,
zero_point: None,
axis: Some(axis),
})
}
pub fn per_channel(scales: Vec<f32>, zero_points: Vec<i32>, axis: usize) -> Result<Self> {
if scales.is_empty() {
return Err(Error::QuantizationInvalid {
field: "scale.len",
expected: "non-empty per-channel scales".to_string(),
got: "length 0".to_string(),
});
}
if scales.len() != zero_points.len() {
return Err(Error::QuantizationInvalid {
field: "zero_point.len",
expected: format!("length matches scale ({})", scales.len()),
got: format!("length {}", zero_points.len()),
});
}
Ok(Self {
scale: scales,
zero_point: Some(zero_points),
axis: Some(axis),
})
}
pub fn mode(&self) -> QuantMode<'_> {
match (self.scale.len(), self.zero_point.as_deref(), self.axis) {
(1, None, _) => QuantMode::PerTensorSymmetric {
scale: self.scale[0],
},
(1, Some(zps), _) => QuantMode::PerTensor {
scale: self.scale[0],
zero_point: zps.first().copied().unwrap_or(0),
},
(_, None, Some(axis)) => QuantMode::PerChannelSymmetric {
scales: &self.scale,
axis,
},
(_, Some(zps), Some(axis)) => QuantMode::PerChannel {
scales: &self.scale,
zero_points: zps,
axis,
},
_ => {
debug_assert!(
false,
"Quantization::mode: per-channel without axis is unreachable"
);
QuantMode::PerTensorSymmetric {
scale: self.scale.first().copied().unwrap_or(1.0),
}
}
}
}
pub fn is_per_tensor(&self) -> bool {
self.scale.len() == 1
}
pub fn is_per_channel(&self) -> bool {
self.scale.len() > 1
}
pub fn is_symmetric(&self) -> bool {
match &self.zero_point {
None => true,
Some(zps) => zps.iter().all(|&z| z == 0),
}
}
pub fn scale(&self) -> &[f32] {
&self.scale
}
pub fn zero_point(&self) -> Option<&[i32]> {
self.zero_point.as_deref()
}
pub fn axis(&self) -> Option<usize> {
self.axis
}
#[cfg(feature = "static")]
pub(crate) fn validate(&self, shape: &[usize]) -> Result<()> {
if self.scale.is_empty() {
return Err(Error::QuantizationInvalid {
field: "scale.len",
expected: ">= 1".to_string(),
got: "0".to_string(),
});
}
if let Some(zps) = self.zero_point.as_ref() {
let expected = if self.scale.len() == 1 {
1
} else {
self.scale.len()
};
if zps.len() != expected {
return Err(Error::QuantizationInvalid {
field: "zero_point.len",
expected: format!(
"{expected} (matching {})",
if self.scale.len() == 1 {
"per-tensor scale"
} else {
"per-channel scale.len"
}
),
got: format!("length {}", zps.len()),
});
}
}
match (self.scale.len(), self.axis) {
(1, None) => Ok(()),
(1, Some(_)) => Err(Error::QuantizationInvalid {
field: "per_tensor_redundant_axis",
expected: "axis=None for per-tensor quantization".to_string(),
got: format!("axis={:?}", self.axis),
}),
(_, None) => Err(Error::QuantizationInvalid {
field: "per_channel_requires_axis",
expected: format!(
"axis=Some(_) for per-channel quantization (scale.len={})",
self.scale.len()
),
got: "axis=None".to_string(),
}),
(n, Some(axis)) => {
if axis >= shape.len() {
return Err(Error::QuantizationInvalid {
field: "axis",
expected: format!("axis < tensor rank ({})", shape.len()),
got: format!("axis={axis}"),
});
}
if shape[axis] != n {
return Err(Error::QuantizationInvalid {
field: "scale.len",
expected: format!("length matches shape[{axis}] ({})", shape[axis]),
got: format!("length {n}"),
});
}
Ok(())
}
}
}
}
impl From<(f32, i32)> for Quantization {
fn from((scale, zero_point): (f32, i32)) -> Self {
Self::per_tensor(scale, zero_point)
}
}
fn deserialize_scalar_or_vec_f32<'de, D: serde::Deserializer<'de>>(
de: D,
) -> std::result::Result<Vec<f32>, D::Error> {
use serde::de::{self, Visitor};
struct V;
impl<'de> Visitor<'de> for V {
type Value = Vec<f32>;
#[cfg_attr(test, mutants::skip)]
fn expecting(&self, f: &mut fmt::Formatter) -> fmt::Result {
f.write_str("f32 or array of f32")
}
fn visit_f64<E: de::Error>(self, v: f64) -> std::result::Result<Self::Value, E> {
Ok(vec![v as f32])
}
#[allow(clippy::cast_possible_truncation)]
fn visit_i64<E: de::Error>(self, v: i64) -> std::result::Result<Self::Value, E> {
Ok(vec![v as f32])
}
#[allow(clippy::cast_possible_truncation, clippy::cast_precision_loss)]
fn visit_u64<E: de::Error>(self, v: u64) -> std::result::Result<Self::Value, E> {
Ok(vec![v as f32])
}
fn visit_seq<A: de::SeqAccess<'de>>(
self,
mut seq: A,
) -> std::result::Result<Self::Value, A::Error> {
let mut out = Vec::with_capacity(seq.size_hint().unwrap_or(1));
while let Some(x) = seq.next_element::<f32>()? {
out.push(x);
}
Ok(out)
}
}
de.deserialize_any(V)
}
fn deserialize_opt_scalar_or_vec_i32<'de, D: serde::Deserializer<'de>>(
de: D,
) -> std::result::Result<Option<Vec<i32>>, D::Error> {
use serde::de::{self, Visitor};
struct V;
impl<'de> Visitor<'de> for V {
type Value = Option<Vec<i32>>;
#[cfg_attr(test, mutants::skip)]
fn expecting(&self, f: &mut fmt::Formatter) -> fmt::Result {
f.write_str("null, i32, or array of i32")
}
#[cfg_attr(test, mutants::skip)]
fn visit_none<E: de::Error>(self) -> std::result::Result<Self::Value, E> {
Ok(None)
}
#[cfg_attr(test, mutants::skip)]
fn visit_unit<E: de::Error>(self) -> std::result::Result<Self::Value, E> {
Ok(None)
}
fn visit_some<D2: serde::Deserializer<'de>>(
self,
de: D2,
) -> std::result::Result<Self::Value, D2::Error> {
struct Inner;
impl<'de> Visitor<'de> for Inner {
type Value = Vec<i32>;
#[cfg_attr(test, mutants::skip)]
fn expecting(&self, f: &mut fmt::Formatter) -> fmt::Result {
f.write_str("i32 or array of i32")
}
#[allow(clippy::cast_possible_truncation)]
fn visit_i64<E: de::Error>(self, v: i64) -> std::result::Result<Self::Value, E> {
Ok(vec![v as i32])
}
#[allow(clippy::cast_possible_truncation, clippy::cast_possible_wrap)]
fn visit_u64<E: de::Error>(self, v: u64) -> std::result::Result<Self::Value, E> {
Ok(vec![v as i32])
}
fn visit_seq<A: de::SeqAccess<'de>>(
self,
mut seq: A,
) -> std::result::Result<Self::Value, A::Error> {
let mut out = Vec::with_capacity(seq.size_hint().unwrap_or(1));
while let Some(x) = seq.next_element::<i32>()? {
out.push(x);
}
Ok(out)
}
}
de.deserialize_any(Inner).map(Some)
}
#[allow(clippy::cast_possible_truncation)]
fn visit_i64<E: de::Error>(self, v: i64) -> std::result::Result<Self::Value, E> {
Ok(Some(vec![v as i32]))
}
#[allow(clippy::cast_possible_truncation, clippy::cast_possible_wrap)]
fn visit_u64<E: de::Error>(self, v: u64) -> std::result::Result<Self::Value, E> {
Ok(Some(vec![v as i32]))
}
fn visit_seq<A: de::SeqAccess<'de>>(
self,
mut seq: A,
) -> std::result::Result<Self::Value, A::Error> {
let mut out = Vec::with_capacity(seq.size_hint().unwrap_or(1));
while let Some(x) = seq.next_element::<i32>()? {
out.push(x);
}
Ok(Some(out))
}
}
de.deserialize_option(V)
}
static UNPLANNED_CPU_ACCESS: AtomicU64 = AtomicU64::new(0);
pub fn unplanned_cpu_access_count() -> u64 {
UNPLANNED_CPU_ACCESS.load(Ordering::Relaxed)
}
#[cfg(feature = "static")]
pub(crate) fn note_unplanned_cpu_access(identity_id: u64, backend: &str, detail: &str) {
UNPLANNED_CPU_ACCESS.fetch_add(1, Ordering::Relaxed);
static WARNED: std::sync::OnceLock<std::sync::Mutex<std::collections::HashSet<u64>>> =
std::sync::OnceLock::new();
let warned = WARNED.get_or_init(Default::default);
if warned.lock().is_ok_and(|mut s| s.insert(identity_id)) {
log::warn!(
"unplanned CPU access on {backend} buffer (identity {identity_id}): {detail} — \
declare the intent at allocation (CpuAccess::Read/Write/ReadWrite) to make \
this a planned, optimized mapping"
);
}
}
#[inline]
pub(crate) fn assert_map_writable(writable: bool, backend: &str) {
assert!(
writable,
"{backend} map is read-only (obtained via map_read()/CpuAccess::Read) — \
use map_mut() or map_write() for mutable access"
);
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
pub enum CpuAccess {
#[default]
None,
Read,
Write,
ReadWrite,
}
impl CpuAccess {
pub fn reads(self) -> bool {
matches!(self, CpuAccess::Read | CpuAccess::ReadWrite)
}
pub fn writes(self) -> bool {
matches!(self, CpuAccess::Write | CpuAccess::ReadWrite)
}
pub fn covers(self, requested: CpuAccess) -> bool {
(!requested.reads() || self.reads()) && (!requested.writes() || self.writes())
}
}
#[non_exhaustive]
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum Compression {
Any,
Scheme(CompressionScheme),
}
crate::ef_vocabulary! {
#[non_exhaustive]
pub enum CompressionScheme {
Ubwc = 1, "ubwc", UBWC,
Afbc = 2, "afbc", AFBC,
Pvric = 3, "pvric", PVRIC,
Dcc = 4, "dcc", DCC,
}
#[doc(hidden)]
pub mod compression_wire;
}
static COMPRESSION_FALLBACKS: AtomicU64 = AtomicU64::new(0);
pub fn compression_fallback_count() -> u64 {
COMPRESSION_FALLBACKS.load(Ordering::Relaxed)
}
#[cfg(feature = "static")]
pub(crate) fn note_compression_fallback(detail: &str) {
COMPRESSION_FALLBACKS.fetch_add(1, Ordering::Relaxed);
log::debug!("compression request fell back to linear: {detail}");
}
pub fn compression_support(format: PixelFormat, dtype: DType) -> bool {
#[cfg(target_os = "android")]
{
crate::ahardwarebuffer_layout::compression_eligible(format, dtype)
&& crate::ahardwarebuffer::device_compression_scheme().is_some()
}
#[cfg(not(target_os = "android"))]
{
let _ = (format, dtype);
false
}
}
#[derive(Debug, Clone)]
pub struct ImageDesc {
width: usize,
height: usize,
format: PixelFormat,
dtype: DType,
memory: Option<TensorMemory>,
access: CpuAccess,
compression: Option<Compression>,
}
impl ImageDesc {
pub fn new(width: usize, height: usize, format: PixelFormat, dtype: DType) -> Self {
Self {
width,
height,
format,
dtype,
memory: None,
access: CpuAccess::None,
compression: None,
}
}
pub fn with_memory(mut self, memory: Option<TensorMemory>) -> Self {
self.memory = memory;
self
}
pub fn with_access(mut self, access: CpuAccess) -> Self {
self.access = access;
self
}
pub fn with_compression(mut self, compression: Compression) -> Self {
self.compression = Some(compression);
self
}
pub fn width(&self) -> usize {
self.width
}
pub fn height(&self) -> usize {
self.height
}
pub fn format(&self) -> PixelFormat {
self.format
}
pub fn dtype(&self) -> DType {
self.dtype
}
pub fn memory(&self) -> Option<TensorMemory> {
self.memory
}
pub fn access(&self) -> CpuAccess {
self.access
}
pub fn compression(&self) -> Option<Compression> {
self.compression
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum IdentityKind {
DmaBuf,
Shm,
IoSurface,
AHardwareBufferId,
AHardwareBufferPtr,
D3d11Texture,
HostPtr,
Pbo,
}
#[derive(Debug, Clone)]
pub struct BufferIdentity {
id: u64,
kind: IdentityKind,
guard: Arc<()>,
}
impl BufferIdentity {
pub fn derived(kind: IdentityKind, key: u64) -> Self {
let id = ((kind as u64) << 56) ^ key;
Self {
id,
kind,
guard: Arc::new(()),
}
}
pub fn id(&self) -> u64 {
self.id
}
pub fn kind(&self) -> IdentityKind {
self.kind
}
pub fn weak(&self) -> Weak<()> {
Arc::downgrade(&self.guard)
}
}
#[cfg(all(target_os = "linux", feature = "static"))]
use nix::sys::stat::{major, minor};
#[cfg(all(target_os = "linux", feature = "static"))]
const DMA_BUF_MAGIC: u32 = 0x444d_4142;
#[cfg(all(target_os = "linux", feature = "static"))]
const TMPFS_MAGIC: u32 = 0x0102_1994;
#[cfg(all(target_os = "linux", feature = "static"))]
const fn fs_magic(raw: i64) -> u32 {
raw as u32
}
pub trait TensorTrait<T>: Send + Sync
where
T: Num + Clone + fmt::Debug,
{
fn new(shape: &[usize], name: Option<&str>) -> Result<Self>
where
Self: Sized;
#[cfg(unix)]
fn from_fd(fd: std::os::fd::OwnedFd, shape: &[usize], name: Option<&str>) -> Result<Self>
where
Self: Sized;
#[cfg(unix)]
fn clone_fd(&self) -> Result<std::os::fd::OwnedFd>;
fn memory(&self) -> TensorMemory;
fn name(&self) -> String;
fn len(&self) -> usize {
self.shape().iter().product()
}
fn is_empty(&self) -> bool {
self.len() == 0
}
fn size(&self) -> usize {
self.len() * std::mem::size_of::<T>()
}
fn shape(&self) -> &[usize];
fn reshape(&mut self, shape: &[usize]) -> Result<()>;
fn capacity_bytes(&self) -> usize;
fn set_logical_shape(&mut self, shape: &[usize]) -> Result<()>;
fn map_with<'a>(&self, access: CpuAccess) -> Result<crate::view::HostView<'a, T>>
where
T: 'a;
fn map<'a>(&self) -> Result<crate::view::HostView<'a, T>>
where
T: 'a,
{
self.map_with(CpuAccess::ReadWrite)
}
fn map_read<'a>(&self) -> Result<crate::view::HostView<'a, T>>
where
T: 'a,
{
self.map_with(CpuAccess::Read)
}
fn map_write<'a>(&self) -> Result<crate::view::HostView<'a, T>>
where
T: 'a,
{
self.map_with(CpuAccess::Write)
}
fn map_mut<'a>(&self) -> Result<crate::view::HostView<'a, T>>
where
T: 'a,
{
self.map_with(CpuAccess::ReadWrite)
}
fn try_map_with<'a>(&self, access: CpuAccess) -> Result<crate::view::HostView<'a, T>>
where
T: 'a,
{
self.map_with(access)
}
fn try_map<'a>(&self) -> Result<crate::view::HostView<'a, T>>
where
T: 'a,
{
self.try_map_with(CpuAccess::ReadWrite)
}
fn buffer_identity(&self) -> &BufferIdentity;
fn view(&self, offset_bytes: usize, shape: &[usize]) -> Result<Self>
where
Self: Sized,
{
let _ = (offset_bytes, shape);
Err(Error::NotImplemented(
"view (zero-copy sub-region) is not supported for this tensor backend".to_owned(),
))
}
}
pub trait TensorMapTrait<T>
where
T: Num + Clone + fmt::Debug,
{
fn shape(&self) -> &[usize];
fn unmap(&mut self);
fn len(&self) -> usize {
self.shape().iter().product()
}
fn is_empty(&self) -> bool {
self.len() == 0
}
fn size(&self) -> usize {
self.len() * std::mem::size_of::<T>()
}
fn as_slice(&self) -> &[T];
fn as_mut_slice(&mut self) -> &mut [T];
fn is_writable(&self) -> bool;
#[cfg(feature = "ndarray")]
fn view(&'_ self) -> Result<ndarray::ArrayView<'_, T, ndarray::Dim<ndarray::IxDynImpl>>> {
Ok(ndarray::ArrayView::from_shape(
self.shape(),
self.as_slice(),
)?)
}
#[cfg(feature = "ndarray")]
fn view_mut(
&'_ mut self,
) -> Result<ndarray::ArrayViewMut<'_, T, ndarray::Dim<ndarray::IxDynImpl>>> {
let shape = self.shape().to_vec();
Ok(ndarray::ArrayViewMut::from_shape(
shape,
self.as_mut_slice(),
)?)
}
}
crate::ef_vocabulary! {
#[non_exhaustive]
pub enum TensorMemory {
Mem = 0, "mem", MEM,
Shm = 1, "shm", SHM,
DmaBuf = 2, "dmabuf", DMABUF,
IoSurface = 3, "iosurface", IOSURFACE,
Pbo = 4, "pbo", PBO,
Cuda = 5, "cuda", CUDA,
}
#[doc(hidden)]
pub mod tensor_memory_wire;
}
impl TensorMemory {
pub fn is_available(self) -> bool {
match self {
TensorMemory::Mem => true,
TensorMemory::Shm => is_shm_available(),
TensorMemory::DmaBuf => is_gpu_buffer_available(),
TensorMemory::IoSurface | TensorMemory::Pbo | TensorMemory::Cuda => false,
}
}
#[cfg(feature = "static")]
pub(crate) fn unsupported_here(self) -> Error {
let reason = match self {
TensorMemory::Mem => "Mem is servable on every platform, so this should not happen",
TensorMemory::Shm => "POSIX shared memory is unix-only",
TensorMemory::DmaBuf => {
"this target has no platform GPU-buffer allocator (DMA-BUF is Linux, \
IOSurface macOS/iOS, AHardwareBuffer Android, D3D11 textures on \
Windows)"
}
TensorMemory::IoSurface => {
"no backend produces or accepts IOSurface under its own code yet -- \
macOS/iOS allocate and report TensorMemory::DmaBuf"
}
TensorMemory::Pbo => {
"PBO tensors are created by ImageProcessor::create_image with a GL \
context current, never by Tensor::new"
}
TensorMemory::Cuda => "there is no CUDA backing in this build",
};
Error::NotImplemented(format!(
"TensorMemory::{self:?} is not supported by this build: {reason}"
))
}
}
impl From<TensorMemory> for String {
fn from(memory: TensorMemory) -> Self {
memory.as_str().to_owned()
}
}
impl TryFrom<&str> for TensorMemory {
type Error = Error;
fn try_from(s: &str) -> Result<Self> {
TensorMemory::from_str_code(s)
.or(if s == "dma" {
Some(TensorMemory::DmaBuf)
} else {
None
})
.ok_or_else(|| Error::InvalidMemoryType(s.to_owned()))
}
}
#[cfg(feature = "static")]
#[derive(Debug)]
#[allow(dead_code)] pub(crate) enum TensorStorage<T>
where
T: Num + Clone + fmt::Debug + Send + Sync,
{
#[cfg(target_os = "linux")]
Dma(DmaTensor<T>),
#[cfg(any(target_os = "macos", target_os = "ios"))]
Dma(IoSurfaceTensor<T>),
#[cfg(target_os = "android")]
Dma(AHardwareBufferTensor<T>),
#[cfg(target_os = "windows")]
Dma(crate::d3d11::texture::D3d11TextureTensor<T>),
#[cfg(unix)]
Shm(ShmTensor<T>),
Mem(MemTensor<T>),
Pbo(PboTensor<T>),
}
#[cfg(feature = "static")]
impl<T> TensorStorage<T>
where
T: Num + Clone + fmt::Debug + Send + Sync,
{
#[cfg(any(
target_os = "macos",
target_os = "ios",
target_os = "android",
target_os = "windows"
))]
pub(crate) fn backing_row_stride(&self) -> Option<usize> {
match self {
#[cfg(any(target_os = "macos", target_os = "ios"))]
TensorStorage::Dma(t) => t.image_backing_row_stride(),
#[cfg(target_os = "android")]
TensorStorage::Dma(t) => t.image_backing_row_stride(),
#[cfg(target_os = "windows")]
TensorStorage::Dma(t) => t.image_backing_row_stride(),
_ => None,
}
}
fn new(shape: &[usize], memory: Option<TensorMemory>, name: Option<&str>) -> Result<Self> {
match memory {
#[cfg(target_os = "linux")]
Some(TensorMemory::DmaBuf) => {
DmaTensor::<T>::new(shape, name).map(TensorStorage::Dma)
}
#[cfg(any(target_os = "macos", target_os = "ios"))]
Some(TensorMemory::DmaBuf) => {
IoSurfaceTensor::<T>::new(shape, name).map(TensorStorage::Dma)
}
#[cfg(target_os = "android")]
Some(TensorMemory::DmaBuf) => {
AHardwareBufferTensor::<T>::new(shape, name).map(TensorStorage::Dma)
}
#[cfg(target_os = "windows")]
Some(TensorMemory::DmaBuf) => Err(crate::error::Error::NotImplemented(
"D3D11 texture tensors are image-formatted; use Tensor::image with \
TensorMemory::DmaBuf"
.to_owned(),
)),
#[cfg(not(any(
target_os = "linux",
target_os = "macos",
target_os = "ios",
target_os = "android",
target_os = "windows"
)))]
Some(TensorMemory::DmaBuf) => Err(crate::error::Error::NotImplemented(
"TensorMemory::DmaBuf is only available on Linux (DMA-BUF), macOS/iOS (IOSurface), \
Android (AHardwareBuffer), and Windows (D3D11 textures, via Tensor::image)"
.to_owned(),
)),
#[cfg(unix)]
Some(TensorMemory::Shm) => {
ShmTensor::<T>::new(shape, name).map(TensorStorage::Shm)
}
#[cfg(not(unix))]
Some(m @ TensorMemory::Shm) => Err(m.unsupported_here()),
Some(TensorMemory::Mem) => {
MemTensor::<T>::new(shape, name).map(TensorStorage::Mem)
}
Some(TensorMemory::Pbo) => Err(crate::error::Error::NotImplemented(
"PboTensor cannot be created via Tensor::new() — use ImageProcessor::create_image()".to_owned(),
)),
Some(m @ (TensorMemory::IoSurface | TensorMemory::Cuda)) => Err(m.unsupported_here()),
None => {
if std::env::var("EDGEFIRST_TENSOR_FORCE_MEM")
.is_ok_and(|x| x != "0" && x.to_lowercase() != "false")
{
MemTensor::<T>::new(shape, name).map(TensorStorage::Mem)
} else {
#[cfg(target_os = "linux")]
{
match DmaTensor::<T>::new(shape, name) {
Ok(tensor) => Ok(TensorStorage::Dma(tensor)),
Err(_) => MemTensor::<T>::new(shape, name).map(TensorStorage::Mem),
}
}
#[cfg(any(target_os = "macos", target_os = "ios"))]
{
match IoSurfaceTensor::<T>::new(shape, name) {
Ok(tensor) => Ok(TensorStorage::Dma(tensor)),
Err(_) => MemTensor::<T>::new(shape, name).map(TensorStorage::Mem),
}
}
#[cfg(target_os = "android")]
{
MemTensor::<T>::new(shape, name).map(TensorStorage::Mem)
}
#[cfg(all(
unix,
not(any(
target_os = "linux",
target_os = "macos",
target_os = "ios",
target_os = "android"
))
))]
{
MemTensor::<T>::new(shape, name).map(TensorStorage::Mem)
}
#[cfg(not(unix))]
{
MemTensor::<T>::new(shape, name).map(TensorStorage::Mem)
}
}
}
}
}
#[cfg(target_os = "linux")]
pub(crate) fn new_dma_with_byte_size(
shape: &[usize],
byte_size: usize,
name: Option<&str>,
) -> Result<Self> {
DmaTensor::<T>::new_with_byte_size(shape, byte_size, name).map(TensorStorage::Dma)
}
pub(crate) fn new_mem_with_byte_size(
shape: &[usize],
byte_size: usize,
name: Option<&str>,
) -> Result<Self>
where
T: 'static,
{
MemTensor::<T>::with_capacity_bytes(shape, byte_size, name).map(TensorStorage::Mem)
}
#[cfg(unix)]
pub(crate) fn new_shm_with_byte_size(
shape: &[usize],
byte_size: usize,
name: Option<&str>,
) -> Result<Self> {
ShmTensor::<T>::new_with_byte_size(shape, byte_size, name).map(TensorStorage::Shm)
}
#[cfg(any(target_os = "macos", target_os = "ios"))]
pub(crate) fn new_image_iosurface(
width: usize,
height: usize,
format: PixelFormat,
dtype: DType,
shape: &[usize],
name: Option<&str>,
) -> Result<Self> {
IoSurfaceTensor::<T>::new_image(width, height, format, dtype, shape, name)
.map(TensorStorage::Dma)
}
#[cfg(target_os = "android")]
pub(crate) fn new_image_ahardwarebuffer(
width: usize,
height: usize,
format: PixelFormat,
dtype: DType,
shape: &[usize],
name: Option<&str>,
access: CpuAccess,
) -> Result<Self> {
AHardwareBufferTensor::<T>::new_image(width, height, format, dtype, shape, name, access)
.map(TensorStorage::Dma)
}
#[cfg(target_os = "windows")]
pub(crate) fn new_image_d3d11(
width: usize,
height: usize,
format: PixelFormat,
dtype: DType,
shape: &[usize],
name: Option<&str>,
access: CpuAccess,
) -> Result<Self> {
crate::d3d11::texture::D3d11TextureTensor::<T>::new_image(
width, height, format, dtype, shape, name, access,
)
.map(TensorStorage::Dma)
}
#[cfg(unix)]
fn from_fd(fd: OwnedFd, shape: &[usize], name: Option<&str>) -> Result<Self> {
#[cfg(target_os = "linux")]
{
use nix::sys::stat::fstat;
use nix::sys::statfs::fstatfs;
let stat = fstat(&fd)?;
let major = major(stat.st_dev);
let minor = minor(stat.st_dev);
if major != 0 {
return Err(Error::UnknownDeviceType(major, minor));
}
let magic = fs_magic(fstatfs(&fd)?.filesystem_type().0 as i64);
log::debug!(
"Creating tensor from fd: major={major}, minor={minor}, magic={magic:#010x}"
);
match magic {
DMA_BUF_MAGIC => DmaTensor::<T>::from_fd(fd, shape, name).map(TensorStorage::Dma),
TMPFS_MAGIC => ShmTensor::<T>::from_fd(fd, shape, name).map(TensorStorage::Shm),
other => Err(Error::UnknownBufferType(other)),
}
}
#[cfg(all(unix, not(target_os = "linux")))]
{
ShmTensor::<T>::from_fd(fd, shape, name).map(TensorStorage::Shm)
}
}
}
#[cfg(feature = "static")]
impl<T> TensorTrait<T> for TensorStorage<T>
where
T: Num + Clone + fmt::Debug + Send + Sync,
{
fn new(shape: &[usize], name: Option<&str>) -> Result<Self> {
Self::new(shape, None, name)
}
#[cfg(unix)]
fn from_fd(fd: OwnedFd, shape: &[usize], name: Option<&str>) -> Result<Self> {
Self::from_fd(fd, shape, name)
}
#[cfg(unix)]
fn clone_fd(&self) -> Result<OwnedFd> {
match self {
#[cfg(any(
target_os = "linux",
target_os = "macos",
target_os = "ios",
target_os = "android"
))]
TensorStorage::Dma(t) => t.clone_fd(),
TensorStorage::Shm(t) => t.clone_fd(),
TensorStorage::Mem(t) => t.clone_fd(),
TensorStorage::Pbo(t) => t.clone_fd(),
}
}
fn memory(&self) -> TensorMemory {
match self {
#[cfg(any(
target_os = "linux",
target_os = "macos",
target_os = "ios",
target_os = "android",
target_os = "windows"
))]
TensorStorage::Dma(_) => TensorMemory::DmaBuf,
#[cfg(unix)]
TensorStorage::Shm(_) => TensorMemory::Shm,
TensorStorage::Mem(_) => TensorMemory::Mem,
TensorStorage::Pbo(_) => TensorMemory::Pbo,
}
}
fn name(&self) -> String {
match self {
#[cfg(any(
target_os = "linux",
target_os = "macos",
target_os = "ios",
target_os = "android",
target_os = "windows"
))]
TensorStorage::Dma(t) => t.name(),
#[cfg(unix)]
TensorStorage::Shm(t) => t.name(),
TensorStorage::Mem(t) => t.name(),
TensorStorage::Pbo(t) => t.name(),
}
}
fn shape(&self) -> &[usize] {
match self {
#[cfg(any(
target_os = "linux",
target_os = "macos",
target_os = "ios",
target_os = "android",
target_os = "windows"
))]
TensorStorage::Dma(t) => t.shape(),
#[cfg(unix)]
TensorStorage::Shm(t) => t.shape(),
TensorStorage::Mem(t) => t.shape(),
TensorStorage::Pbo(t) => t.shape(),
}
}
fn reshape(&mut self, shape: &[usize]) -> Result<()> {
match self {
#[cfg(any(
target_os = "linux",
target_os = "macos",
target_os = "ios",
target_os = "android",
target_os = "windows"
))]
TensorStorage::Dma(t) => t.reshape(shape),
#[cfg(unix)]
TensorStorage::Shm(t) => t.reshape(shape),
TensorStorage::Mem(t) => t.reshape(shape),
TensorStorage::Pbo(t) => t.reshape(shape),
}
}
fn capacity_bytes(&self) -> usize {
match self {
#[cfg(any(
target_os = "linux",
target_os = "macos",
target_os = "ios",
target_os = "android",
target_os = "windows"
))]
TensorStorage::Dma(t) => t.capacity_bytes(),
#[cfg(unix)]
TensorStorage::Shm(t) => t.capacity_bytes(),
TensorStorage::Mem(t) => t.capacity_bytes(),
TensorStorage::Pbo(t) => t.capacity_bytes(),
}
}
fn set_logical_shape(&mut self, shape: &[usize]) -> Result<()> {
match self {
#[cfg(any(
target_os = "linux",
target_os = "macos",
target_os = "ios",
target_os = "android",
target_os = "windows"
))]
TensorStorage::Dma(t) => t.set_logical_shape(shape),
#[cfg(unix)]
TensorStorage::Shm(t) => t.set_logical_shape(shape),
TensorStorage::Mem(t) => t.set_logical_shape(shape),
TensorStorage::Pbo(t) => t.set_logical_shape(shape),
}
}
fn map_with<'a>(&self, access: CpuAccess) -> Result<crate::view::HostView<'a, T>>
where
T: 'a,
{
match self {
#[cfg(any(
target_os = "linux",
target_os = "macos",
target_os = "ios",
target_os = "android",
target_os = "windows"
))]
TensorStorage::Dma(t) => t.map_with(access),
#[cfg(unix)]
TensorStorage::Shm(t) => t.map_with(access),
TensorStorage::Mem(t) => t.map_with(access),
TensorStorage::Pbo(t) => t.map_with(access),
}
}
fn try_map_with<'a>(&self, access: CpuAccess) -> Result<crate::view::HostView<'a, T>>
where
T: 'a,
{
match self {
#[cfg(any(
target_os = "linux",
target_os = "macos",
target_os = "ios",
target_os = "android",
target_os = "windows"
))]
TensorStorage::Dma(t) => t.try_map_with(access),
#[cfg(unix)]
TensorStorage::Shm(t) => t.try_map_with(access),
TensorStorage::Mem(t) => t.try_map_with(access),
TensorStorage::Pbo(t) => t.try_map_with(access),
}
}
fn buffer_identity(&self) -> &BufferIdentity {
match self {
#[cfg(any(
target_os = "linux",
target_os = "macos",
target_os = "ios",
target_os = "android",
target_os = "windows"
))]
TensorStorage::Dma(t) => t.buffer_identity(),
#[cfg(unix)]
TensorStorage::Shm(t) => t.buffer_identity(),
TensorStorage::Mem(t) => t.buffer_identity(),
TensorStorage::Pbo(t) => t.buffer_identity(),
}
}
fn view(&self, offset_bytes: usize, shape: &[usize]) -> Result<Self> {
match self {
#[cfg(any(
target_os = "linux",
target_os = "macos",
target_os = "ios",
target_os = "android",
target_os = "windows"
))]
TensorStorage::Dma(t) => t.view(offset_bytes, shape).map(TensorStorage::Dma),
#[cfg(unix)]
TensorStorage::Shm(t) => t.view(offset_bytes, shape).map(TensorStorage::Shm),
TensorStorage::Mem(t) => t.view(offset_bytes, shape).map(TensorStorage::Mem),
TensorStorage::Pbo(t) => t.view(offset_bytes, shape).map(TensorStorage::Pbo),
}
}
}
#[cfg(feature = "static")]
#[derive(Debug)]
pub struct Tensor<T>
where
T: Num + Clone + fmt::Debug + Send + Sync,
{
cuda: Option<crate::cuda::CudaHandle>,
pub(crate) storage: TensorStorage<T>,
format: Option<PixelFormat>,
chroma: Option<Box<Tensor<T>>>,
row_stride: Option<usize>,
plane_offset: Option<usize>,
pub(crate) quantization: Option<Quantization>,
colorimetry: Option<crate::Colorimetry>,
cpu_access: CpuAccess,
compression: Option<CompressionScheme>,
view_origin: Option<ViewOrigin>,
}
#[cfg(feature = "static")]
impl<T> Tensor<T>
where
T: Num + Clone + fmt::Debug + Send + Sync,
{
pub(crate) fn wrap(storage: TensorStorage<T>) -> Self {
Self {
storage,
format: None,
chroma: None,
row_stride: None,
plane_offset: None,
quantization: None,
cuda: None,
colorimetry: None,
cpu_access: CpuAccess::ReadWrite,
compression: None,
view_origin: None,
}
}
pub fn from_slice(values: &[T], shape: &[usize]) -> Result<Self>
where
T: Copy,
{
let expected: usize = shape.iter().product();
if values.len() != expected {
return Err(Error::InvalidShape(format!(
"from_slice: values.len()={} but shape product={expected} (shape={shape:?})",
values.len()
)));
}
let t = Self::new(shape, Some(TensorMemory::Mem), None)?;
{
let mut m = t.map()?;
m.as_mut_slice().copy_from_slice(values);
}
Ok(t)
}
pub unsafe fn from_foreign(
ptr: *mut T,
shape: &[usize],
owner: Option<crate::ForeignOwner>,
name: Option<&str>,
) -> Result<Self> {
unsafe { Self::from_foreign_with_capacity(ptr, shape, 0, owner, name) }
}
pub unsafe fn from_foreign_with_capacity(
ptr: *mut T,
shape: &[usize],
capacity_bytes: usize,
owner: Option<crate::ForeignOwner>,
name: Option<&str>,
) -> Result<Self> {
if shape.is_empty() {
return Err(Error::InvalidSize(0));
}
if ptr.is_null() {
return Err(Error::InvalidArgument(
"from_foreign: ptr must be non-null".to_owned(),
));
}
shape
.iter()
.copied()
.try_fold(1usize, |acc, dim| acc.checked_mul(dim))
.ok_or_else(|| {
Error::InvalidArgument(format!(
"from_foreign: shape.product() overflows usize (shape={shape:?})"
))
})?;
let capacity_elems = capacity_bytes / std::mem::size_of::<T>();
let mem = unsafe {
MemTensor::<T>::from_foreign_with_capacity(ptr, shape, capacity_elems, owner, name)
};
Ok(Self::wrap(TensorStorage::Mem(mem)))
}
pub fn pin_host<'a>(&self, access: CpuAccess) -> Result<HostPin<'a>>
where
T: 'a,
{
if access == CpuAccess::None {
return Err(Error::InvalidArgument(
"pin_host: CpuAccess::None is not a valid pin; a pinned mapping \
is CPU access by definition"
.to_owned(),
));
}
let logical = self.len() * std::mem::size_of::<T>();
let pin = match &self.storage {
TensorStorage::Mem(m) => Ok(m.host_pin()),
#[cfg(unix)]
TensorStorage::Shm(s) => s.host_pin(),
#[cfg(target_os = "linux")]
TensorStorage::Dma(d) => d.host_pin(),
#[cfg(any(target_os = "macos", target_os = "ios"))]
TensorStorage::Dma(s) => s.host_pin(),
#[cfg(target_os = "windows")]
TensorStorage::Dma(t) => t.host_pin(access),
#[cfg(target_os = "android")]
TensorStorage::Dma(_) => Err(Error::NotImplemented(
"pin_host: AHardwareBuffer has no host address outside \
AHardwareBuffer_lock, so it cannot hand out one that outlives \
a guard. Use map()/map_with() for CPU access; for zero-copy \
handoff pass the buffer itself."
.to_owned(),
)),
TensorStorage::Pbo(_) => Err(Error::NotImplemented(
crate::pbo::PIN_HOST_PBO_REFUSAL.to_owned(),
)),
}?;
Ok(pin.narrowed(logical))
}
pub fn sync_for_cpu(&self, access: CpuAccess) -> Result<()> {
Self::check_sync_access(access)?;
match &self.storage {
TensorStorage::Mem(_) => Ok(()),
#[cfg(unix)]
TensorStorage::Shm(_) => Ok(()),
#[cfg(target_os = "linux")]
TensorStorage::Dma(d) => {
crate::dmabuf::sync_access(&d.fd, true, access).map_err(Error::NixError)
}
#[cfg(any(target_os = "macos", target_os = "ios"))]
TensorStorage::Dma(s) => s.sync_for_cpu(access),
#[cfg(target_os = "windows")]
TensorStorage::Dma(t) => t.sync_for_cpu(access),
#[cfg(target_os = "android")]
TensorStorage::Dma(_) => Err(Error::NotImplemented(
"sync_for_cpu: AHardwareBuffer has no sync independent of its \
lock. AHardwareBuffer_lock both establishes the address and \
performs the cache invalidate, and the NDK leaves locking an \
already-locked buffer undefined, so a bracket on top of a map \
guard would nest. Use map()/map_with(), which owns the pair."
.to_owned(),
)),
TensorStorage::Pbo(_) => Err(Error::NotImplemented(
"sync_for_cpu: a PBO has no coherency window independent of its \
map -- glMapBufferRange establishes the address and the \
visibility together, and glUnmapBuffer publishes the writes. \
Use map()/map_with(), which owns the pair."
.to_owned(),
)),
}
}
pub fn sync_for_device(&self, access: CpuAccess) -> Result<()> {
Self::check_sync_access(access)?;
match &self.storage {
TensorStorage::Mem(_) => Ok(()),
#[cfg(unix)]
TensorStorage::Shm(_) => Ok(()),
#[cfg(target_os = "linux")]
TensorStorage::Dma(d) => {
crate::dmabuf::sync_access(&d.fd, false, access).map_err(Error::NixError)
}
#[cfg(any(target_os = "macos", target_os = "ios"))]
TensorStorage::Dma(s) => s.sync_for_device(access),
#[cfg(target_os = "windows")]
TensorStorage::Dma(t) => t.sync_for_device(access),
#[cfg(target_os = "android")]
TensorStorage::Dma(_) => Err(Error::NotImplemented(
"sync_for_device: AHardwareBuffer has no sync independent of \
its lock — see sync_for_cpu. Use map()/map_with()."
.to_owned(),
)),
TensorStorage::Pbo(_) => Err(Error::NotImplemented(
"sync_for_device: a PBO has no coherency window independent of \
its map -- see sync_for_cpu. Use map()/map_with()."
.to_owned(),
)),
}
}
fn check_sync_access(access: CpuAccess) -> Result<()> {
if access == CpuAccess::None {
return Err(Error::InvalidArgument(
"sync_for_cpu/sync_for_device: CpuAccess::None describes no \
access to synchronise"
.to_owned(),
));
}
Ok(())
}
#[cfg(feature = "ndarray")]
pub fn from_arrayview3(view: ndarray::ArrayView3<'_, T>) -> Result<Self>
where
T: Copy,
{
let (h, w, c) = view.dim();
let t = Self::new(&[h, w, c], Some(TensorMemory::Mem), None)?;
{
let mut m = t.map()?;
let dst = m.as_mut_slice();
if let Some(src) = view.as_slice() {
dst.copy_from_slice(src);
} else {
for (d, &s) in dst.iter_mut().zip(view.iter()) {
*d = s;
}
}
}
Ok(t)
}
pub fn new(shape: &[usize], memory: Option<TensorMemory>, name: Option<&str>) -> Result<Self> {
let _span = tracing::trace_span!(
"tensor.alloc",
?shape,
memory = ?memory,
dtype = std::any::type_name::<T>(),
)
.entered();
#[cfg_attr(not(target_os = "linux"), allow(unused_mut))]
let mut t = TensorStorage::new(shape, memory, name).map(Self::wrap)?;
#[cfg(target_os = "linux")]
t.try_init_dma_cuda();
Ok(t)
}
pub fn image_desc(desc: &ImageDesc) -> Result<Self>
where
T: 'static,
{
let Some(t_dtype) = dtype_of::<T>() else {
return Err(Error::InvalidArgument(
"image_desc: element type has no DType mapping".into(),
));
};
if t_dtype != desc.dtype {
return Err(Error::InvalidArgument(format!(
"image_desc: desc.dtype is {:?} but the tensor element type is {t_dtype:?}",
desc.dtype
)));
}
if desc.compression.is_some() && desc.access != CpuAccess::None {
return Err(Error::InvalidArgument(format!(
"image_desc: a compression request requires CpuAccess::None (declared {:?}) — CPU mapping pins the layout linear",
desc.access
)));
}
if let Some(Compression::Scheme(requested)) = desc.compression {
#[cfg(not(target_os = "android"))]
{
return Err(Error::NotImplemented(format!(
"image_desc: Compression::Scheme({requested:?}) — no vendor tile compression on this platform (request Compression::Any for a portable fallback)"
)));
}
#[cfg(target_os = "android")]
{
if matches!(
desc.memory,
Some(TensorMemory::Mem) | Some(TensorMemory::Shm) | Some(TensorMemory::Pbo)
) {
return Err(Error::InvalidArgument(format!(
"image_desc: Compression::Scheme({requested:?}) requires hardware memory (TensorMemory::DmaBuf or auto-select), got {:?}",
desc.memory
)));
}
if !crate::ahardwarebuffer_layout::compression_eligible(desc.format, desc.dtype) {
return Err(Error::InvalidArgument(format!(
"image_desc: ({:?}, {:?}) is not compression-eligible (RGBA8888 u8/i8 initially)",
desc.format, desc.dtype
)));
}
let device = crate::ahardwarebuffer::device_compression_scheme();
if device != Some(requested) {
return Err(Error::InvalidArgument(format!(
"image_desc: Compression::Scheme({requested:?}) but the device's native scheme is {device:?}"
)));
}
}
}
#[allow(unused_mut)]
let mut t = match (desc.memory, desc.compression) {
(None, Some(request)) => {
match Self::image(
desc.width,
desc.height,
desc.format,
Some(TensorMemory::DmaBuf),
desc.access,
) {
Ok(t) => t,
Err(e) if matches!(request, Compression::Scheme(_)) => return Err(e),
Err(_) => Self::image(desc.width, desc.height, desc.format, None, desc.access)?,
}
}
(memory, _) => Self::image(desc.width, desc.height, desc.format, memory, desc.access)?,
};
if let Some(request) = desc.compression {
#[cfg(target_os = "android")]
{
let eligible =
crate::ahardwarebuffer_layout::compression_eligible(desc.format, desc.dtype);
let scheme = crate::ahardwarebuffer::device_compression_scheme();
let is_ahb = t.memory() == TensorMemory::DmaBuf;
match (request, scheme) {
(_, Some(s)) if eligible && is_ahb => {
t.set_compression_unchecked(Some(s));
}
(Compression::Scheme(requested), _) => {
return Err(Error::InvalidOperation(format!(
"image_desc: Compression::Scheme({requested:?}) requested but the allocation resolved to {:?} (not an AHardwareBuffer)",
t.memory()
)));
}
(Compression::Any, _) => {
note_compression_fallback(&format!(
"({:?}, {:?}) {}x{}: eligible={eligible}, scheme={scheme:?}, memory={:?}",
desc.format,
desc.dtype,
desc.width,
desc.height,
t.memory()
));
}
}
}
#[cfg(not(target_os = "android"))]
{
debug_assert!(matches!(request, Compression::Any));
let _ = request;
note_compression_fallback(&format!(
"({:?}, {:?}) {}x{}: no vendor tile compression on this platform",
desc.format, desc.dtype, desc.width, desc.height
));
}
}
Ok(t)
}
pub fn image(
width: usize,
height: usize,
format: PixelFormat,
memory: Option<TensorMemory>,
access: CpuAccess,
) -> Result<Self>
where
T: 'static,
{
let shape = format.allocation_shape(width, height).ok_or_else(|| {
Error::InvalidArgument(format!(
"invalid dimensions {width}x{height} for format {format:?}"
))
})?;
#[cfg(any(target_os = "macos", target_os = "ios"))]
if matches!(memory, Some(TensorMemory::DmaBuf)) {
let natural_row_bytes = match format.layout() {
PixelLayout::Planar => width * std::mem::size_of::<T>(),
_ => width * format.channels() * std::mem::size_of::<T>(),
};
let has_image_fourcc = dtype_of::<T>()
.and_then(|dt| crate::iosurface_layout::image_iosurface_layout(format, dt))
.is_some();
let padded_ok = has_image_fourcc && format.layout() != PixelLayout::Planar;
if !natural_row_bytes.is_multiple_of(64) && !padded_ok {
let elem_size = std::mem::size_of::<T>();
let per_pixel_bytes = match format.layout() {
PixelLayout::Planar => elem_size.max(1),
_ => format.channels().max(1) * elem_size.max(1),
};
let aligned_row_bytes = natural_row_bytes.next_multiple_of(64);
let pad_hint =
if per_pixel_bytes > 0 && aligned_row_bytes.is_multiple_of(per_pixel_bytes) {
let w = aligned_row_bytes / per_pixel_bytes;
format!("Pad width to {w} (the next 64-byte-aligned stride), ")
} else {
String::new()
};
return Err(Error::InvalidArgument(format!(
"Tensor::image: {format:?} {width}x{height} with element \
size {elem_size} produces a {natural_row_bytes}-byte natural \
row pitch, which is not 64-byte aligned. \
IOSurface rounds bytes_per_row up to 64 bytes, so a \
contiguous CPU map of this tensor would read garbage. \
{pad_hint}pass memory=None to auto-fall-back to SHM, or \
pass memory=Some(TensorMemory::Shm) or \
Some(TensorMemory::Mem) explicitly."
)));
}
let dtype = dtype_of::<T>().ok_or_else(|| {
Error::InvalidArgument(format!(
"Tensor::image: element type {} has no DType, so no \
image-formatted IOSurface exists. Pass memory=None or \
Some(TensorMemory::Mem) for a CPU tensor.",
std::any::type_name::<T>()
))
})?;
if crate::iosurface_layout::image_iosurface_layout(format, dtype).is_none() {
return Err(Error::InvalidArgument(format!(
"Tensor::image: no zero-copy IOSurface mapping exists for \
{format:?}/{dtype:?} on macOS/iOS (supported: \
Rgba/Rgb @ U8/I8, Bgra/Yuyv/Grey/Nv12/Nv16/Nv24 @ U8, \
Rgba/PlanarRgb/PlanarRgba @ F16). Pass memory=None to \
auto-select, or Some(TensorMemory::Mem) explicitly, for \
a CPU tensor."
)));
}
if format == PixelFormat::Rgb
&& matches!(dtype, DType::U8 | DType::I8)
&& packed_rgb888_layout(width, height).is_none()
{
return Err(Error::InvalidArgument(format!(
"Tensor::image: Rgb {dtype:?} requires width%4==0 for the RGBA8888 \
IOSurface packing (got width={width}). Pad the width, or pass \
memory=Some(TensorMemory::Mem) for a CPU tensor."
)));
}
let storage = TensorStorage::<T>::new_image_iosurface(
width, height, format, dtype, &shape, None,
)?;
let mut t = Self::wrap(storage);
t.format = Some(format);
if let TensorStorage::Dma(ref io) = t.storage {
let bpr = io.bytes_per_row();
if let Some(natural) = t.effective_row_stride() {
if bpr > natural {
t.set_row_stride_unchecked(bpr);
}
}
}
t.cpu_access = access;
return Ok(t);
}
#[cfg(target_os = "android")]
if matches!(memory, Some(TensorMemory::DmaBuf)) {
let dtype = dtype_of::<T>().ok_or_else(|| {
Error::InvalidArgument(format!(
"Tensor::image: element type {} has no DType, so no \
image-formatted AHardwareBuffer exists. Pass memory=None \
or Some(TensorMemory::Mem) for a CPU tensor.",
std::any::type_name::<T>()
))
})?;
if format.layout() == PixelLayout::Planar
&& dtype == DType::F16
&& packed_rgba16f_layout(format, dtype, width, height).is_none()
{
return Err(Error::InvalidArgument(format!(
"Tensor::image: {format:?} F16 requires width%4==0 for the RGBA16F \
AHardwareBuffer packing (got width={width}). Pad the width, or pass \
memory=Some(TensorMemory::Mem) for a CPU tensor."
)));
}
if format == PixelFormat::Rgb
&& matches!(dtype, DType::U8 | DType::I8)
&& packed_rgb888_layout(width, height).is_none()
{
return Err(Error::InvalidArgument(format!(
"Tensor::image: Rgb {dtype:?} requires width%4==0 for the RGBA8888 \
AHardwareBuffer packing (got width={width}). Pad the width, or pass \
memory=Some(TensorMemory::Mem) for a CPU tensor."
)));
}
if crate::ahardwarebuffer::image_ahardwarebuffer_layout(format, dtype).is_none() {
return Err(Error::InvalidArgument(format!(
"Tensor::image: no zero-copy AHardwareBuffer mapping exists \
for {format:?}/{dtype:?} on Android (Grey/NV* need \
R8-format buffers, API 29+; the HAL floor is API 26 — see \
`image_ahardwarebuffer_layout`). Pass memory=None to \
auto-select, or Some(TensorMemory::Mem) explicitly, for a \
CPU tensor; camera NV12 stays zero-copy by wrapping the \
camera's own AHardwareBuffer instead of allocating one."
)));
}
let storage = TensorStorage::<T>::new_image_ahardwarebuffer(
width, height, format, dtype, &shape, None, access,
)?;
let mut t = Self::wrap(storage);
t.format = Some(format);
if let TensorStorage::Dma(ref ahb) = t.storage {
let bpr = ahb.bytes_per_row();
if let Some(natural) = t.effective_row_stride() {
if bpr > natural {
log::debug!(
"Tensor::image: gralloc padded the {format:?} \
AHardwareBuffer pitch to {bpr} bytes (natural \
{natural}); recording row stride"
);
t.set_row_stride_unchecked(bpr);
}
}
}
t.cpu_access = access;
return Ok(t);
}
#[cfg(target_os = "windows")]
if matches!(memory, Some(TensorMemory::DmaBuf)) {
let dtype = dtype_of::<T>().ok_or_else(|| {
Error::InvalidArgument(format!(
"Tensor::image: element type {} has no DType, so no D3D11 texture \
layout exists. Pass memory=None or Some(TensorMemory::Mem) for a \
CPU tensor.",
std::any::type_name::<T>()
))
})?;
let storage = TensorStorage::<T>::new_image_d3d11(
width, height, format, dtype, &shape, None, access,
)?;
let mut t = Self::wrap(storage);
t.format = Some(format);
record_d3d11_row_stride(&mut t, format, "Tensor::image");
t.cpu_access = access;
t.try_init_d3d11_cuda();
return Ok(t);
}
let elem = std::mem::size_of::<T>();
let channels = format.channels();
let (natural_stride, total_rows) = match format.layout() {
PixelLayout::SemiPlanar => (width.next_multiple_of(2) * elem, shape[0]),
PixelLayout::Packed => (width * channels * elem, height),
PixelLayout::Planar => (width * elem, channels * height),
};
let aligned_stride = natural_stride.next_multiple_of(64);
let semi = format.layout() == PixelLayout::SemiPlanar;
let host_stride = if semi { aligned_stride } else { natural_stride };
let host_byte_size = host_stride * total_rows;
#[cfg(target_os = "linux")]
let dma_byte_size = aligned_stride * total_rows;
let (storage, used_stride) = match memory {
#[cfg(target_os = "linux")]
Some(TensorMemory::DmaBuf) => (
TensorStorage::<T>::new_dma_with_byte_size(&shape, dma_byte_size, None)?,
aligned_stride,
),
#[cfg(unix)]
Some(TensorMemory::Shm) => (
TensorStorage::<T>::new_shm_with_byte_size(&shape, host_byte_size, None)?,
host_stride,
),
Some(TensorMemory::Mem) => (
TensorStorage::<T>::new_mem_with_byte_size(&shape, host_byte_size, None)?,
host_stride,
),
#[allow(unused_variables)]
Some(other) => {
return {
let mut t = Self::new(&shape, Some(other), None)?;
t.format = Some(format);
t.cpu_access = access;
Ok(t)
};
}
None => {
#[cfg(target_os = "linux")]
{
match TensorStorage::<T>::new_dma_with_byte_size(&shape, dma_byte_size, None) {
Ok(s) => (s, aligned_stride),
Err(_) => (
TensorStorage::<T>::new_mem_with_byte_size(
&shape,
host_byte_size,
None,
)?,
host_stride,
),
}
}
#[cfg(not(target_os = "linux"))]
{
(
TensorStorage::<T>::new_mem_with_byte_size(&shape, host_byte_size, None)?,
host_stride,
)
}
}
};
let mut t = Self::wrap(storage);
t.format = Some(format);
if semi || used_stride > natural_stride {
t.set_row_stride_unchecked(used_stride);
}
debug_assert!(
t.row_stride.is_some() || !semi,
"image() must always set row_stride for semi-planar tensors"
);
t.cpu_access = access;
#[cfg(target_os = "linux")]
t.try_init_dma_cuda();
Ok(t)
}
pub fn image_with_stride(
width: usize,
height: usize,
format: PixelFormat,
row_stride_bytes: usize,
memory: Option<TensorMemory>,
access: CpuAccess,
) -> Result<Self> {
#[cfg(not(target_os = "linux"))]
let _ = access;
#[cfg(not(target_os = "linux"))]
{
let _ = (width, height, format, row_stride_bytes, memory);
Err(Error::NotImplemented(
"image_with_stride requires DMA support (Linux only)".to_owned(),
))
}
#[cfg(target_os = "linux")]
{
if format.layout() != PixelLayout::Packed {
return Err(Error::NotImplemented(format!(
"Tensor::image_with_stride only supports packed pixel layouts, got {format:?}"
)));
}
let elem = std::mem::size_of::<T>();
let min_stride = width
.checked_mul(format.channels())
.and_then(|p| p.checked_mul(elem))
.ok_or_else(|| {
Error::InvalidArgument(format!(
"image_with_stride: width {width} × channels {} × sizeof::<T>={elem} \
overflows usize",
format.channels()
))
})?;
if row_stride_bytes < min_stride {
return Err(Error::InvalidArgument(format!(
"image_with_stride: row_stride {row_stride_bytes} < minimum {min_stride} \
({width} px × {} ch × {elem} B)",
format.channels()
)));
}
let total_byte_size = row_stride_bytes.checked_mul(height).ok_or_else(|| {
Error::InvalidArgument(format!(
"image_with_stride: row_stride {row_stride_bytes} × height {height} overflows usize"
))
})?;
let shape = vec![height, width, format.channels()];
let storage = match memory {
Some(TensorMemory::DmaBuf) | None => {
TensorStorage::<T>::new_dma_with_byte_size(&shape, total_byte_size, None)?
}
Some(other) => {
return Err(Error::NotImplemented(format!(
"image_with_stride: only TensorMemory::DmaBuf is supported, got {other:?}"
)));
}
};
let mut t = Self::wrap(storage);
t.format = Some(format);
t.row_stride = Some(row_stride_bytes);
t.cpu_access = access;
t.try_init_dma_cuda();
Ok(t)
}
}
pub fn set_format(&mut self, format: PixelFormat) -> Result<()> {
let shape = self.shape();
match format.layout() {
PixelLayout::Packed => {
if shape.len() != 3 || shape[2] != format.channels() {
return Err(Error::InvalidShape(format!(
"packed format {format:?} expects [H, W, {}], got {shape:?}",
format.channels()
)));
}
}
PixelLayout::Planar => {
if shape.len() != 3 || shape[0] != format.channels() {
return Err(Error::InvalidShape(format!(
"planar format {format:?} expects [{}, H, W], got {shape:?}",
format.channels()
)));
}
}
PixelLayout::SemiPlanar => {
if shape.len() != 2 {
return Err(Error::InvalidShape(format!(
"semi-planar format {format:?} expects [H*k, W], got {shape:?}"
)));
}
match format {
PixelFormat::Nv12 if shape[0] % 3 == 1 => {
return Err(Error::InvalidShape(format!(
"NV12 contiguous shape[0] must be H + ceil(H/2) for some height; \
{} is unreachable (≡ 1 mod 3)",
shape[0]
)));
}
PixelFormat::Nv16 if !shape[0].is_multiple_of(2) => {
return Err(Error::InvalidShape(format!(
"NV16 contiguous shape[0] must be even, got {}",
shape[0]
)));
}
PixelFormat::Nv24 if !shape[0].is_multiple_of(3) => {
return Err(Error::InvalidShape(format!(
"NV24 contiguous shape[0] must be a multiple of 3 (= 3H), got {}",
shape[0]
)));
}
_ => {}
}
}
}
if self.format != Some(format) {
self.row_stride = None;
self.plane_offset = None;
match self.storage {
TensorStorage::Mem(ref mut m) => m.set_offset(0),
#[cfg(target_os = "linux")]
TensorStorage::Dma(ref mut dma) => dma.mmap_offset = 0,
#[cfg(any(target_os = "macos", target_os = "ios"))]
TensorStorage::Dma(ref mut io) => io.view_offset = 0,
#[cfg(target_os = "windows")]
TensorStorage::Dma(ref mut tex) => tex.set_view_offset(0),
TensorStorage::Pbo(ref mut p) => p.view_offset = 0,
#[cfg(unix)]
TensorStorage::Shm(ref mut s) => s.set_offset(0),
#[cfg(target_os = "android")]
TensorStorage::Dma(_) => {}
}
}
self.format = Some(format);
Ok(())
}
pub fn configure_image(
&mut self,
width: usize,
height: usize,
format: PixelFormat,
) -> Result<()> {
let shape = format.allocation_shape(width, height).ok_or_else(|| {
Error::InvalidArgument(format!(
"invalid dimensions {width}x{height} for format {format:?}"
))
})?;
let prior_stride = self.row_stride;
self.storage.set_logical_shape(&shape)?;
self.set_format(format)?;
self.row_stride = None;
let elem = std::mem::size_of::<T>();
let channels = format.channels();
let (min_stride, total_rows) = match format.layout() {
PixelLayout::SemiPlanar => (width.next_multiple_of(2) * elem, shape[0]),
PixelLayout::Packed => (width * channels * elem, height),
PixelLayout::Planar => (width * elem, channels * height),
};
let needs_align = self.storage.memory() == TensorMemory::DmaBuf
|| format.layout() == PixelLayout::SemiPlanar;
#[cfg(any(
target_os = "macos",
target_os = "ios",
target_os = "android",
target_os = "windows"
))]
if let Some(pitch) = self.storage.backing_row_stride() {
let active_stride = if pitch > min_stride {
self.set_row_stride_unchecked(pitch);
pitch
} else {
min_stride
};
if !needs_align {
return Ok(());
}
return self.check_configured_stride(active_stride, total_rows);
}
if !needs_align {
return Ok(());
}
let aligned = min_stride.next_multiple_of(64);
let capacity = self.storage.capacity_bytes();
let candidate = if let Some(ps) = prior_stride {
if ps >= min_stride && ps % 64 == 0 && ps * total_rows <= capacity {
ps
} else {
aligned
}
} else {
aligned
};
let active_stride = if candidate * total_rows <= capacity {
self.set_row_stride_unchecked(candidate);
candidate
} else {
min_stride
};
self.check_configured_stride(active_stride, total_rows)
}
fn check_configured_stride(&self, active_stride: usize, total_rows: usize) -> Result<()> {
let needed = active_stride * total_rows;
let capacity = self.storage.capacity_bytes();
if needed > capacity {
return Err(Error::InsufficientCapacity { needed, capacity });
}
Ok(())
}
pub fn image_with_capacity(
width: usize,
height: usize,
format: PixelFormat,
memory: Option<TensorMemory>,
access: CpuAccess,
) -> Result<Self>
where
T: 'static,
{
Self::image(width, height, format, memory, access)
}
pub fn format(&self) -> Option<PixelFormat> {
self.format
}
fn image_dims(&self) -> Option<(PixelFormat, &[usize])> {
let fmt = self.format?;
let shape = self.shape();
let rank = image_rank(fmt.layout());
let batch_dims = shape.len().checked_sub(rank).filter(|&n| n <= 1)?;
Some((fmt, &shape[batch_dims..]))
}
pub fn width(&self) -> Option<usize> {
let (fmt, dims) = self.image_dims()?;
match fmt.layout() {
PixelLayout::Packed => Some(dims[1]),
PixelLayout::Planar => Some(dims[2]),
PixelLayout::SemiPlanar => Some(dims[1]),
}
}
pub fn height(&self) -> Option<usize> {
let (fmt, dims) = self.image_dims()?;
match fmt.layout() {
PixelLayout::Packed => Some(dims[0]),
PixelLayout::Planar => Some(dims[1]),
PixelLayout::SemiPlanar => {
if self.is_multiplane() {
Some(dims[0])
} else {
match fmt {
PixelFormat::Nv12 => Some(dims[0] * 2 / 3),
PixelFormat::Nv16 => Some(dims[0] / 2),
PixelFormat::Nv24 => Some(dims[0] / 3),
_ => None,
}
}
}
}
}
pub fn from_planes(luma: Tensor<T>, chroma: Tensor<T>, format: PixelFormat) -> Result<Self> {
if format.layout() != PixelLayout::SemiPlanar {
return Err(Error::InvalidArgument(format!(
"from_planes requires a semi-planar format, got {format:?}"
)));
}
if chroma.format.is_some() || chroma.chroma.is_some() {
return Err(Error::InvalidArgument(
"chroma tensor must be a raw tensor (no format or chroma metadata)".into(),
));
}
let luma_shape = luma.shape();
let chroma_shape = chroma.shape();
if luma_shape.len() != 2 || chroma_shape.len() != 2 {
return Err(Error::InvalidArgument(format!(
"from_planes expects 2D shapes, got luma={luma_shape:?} chroma={chroma_shape:?}"
)));
}
if luma_shape[1] != chroma_shape[1] {
return Err(Error::InvalidArgument(format!(
"luma width {} != chroma width {}",
luma_shape[1], chroma_shape[1]
)));
}
match format {
PixelFormat::Nv12 => {
if luma_shape[0] % 2 != 0 {
return Err(Error::InvalidArgument(format!(
"NV12 requires even luma height, got {}",
luma_shape[0]
)));
}
if chroma_shape[0] != luma_shape[0] / 2 {
return Err(Error::InvalidArgument(format!(
"NV12 chroma height {} != luma height / 2 ({})",
chroma_shape[0],
luma_shape[0] / 2
)));
}
}
PixelFormat::Nv16 => {
if chroma_shape[0] != luma_shape[0] {
return Err(Error::InvalidArgument(format!(
"NV16 chroma height {} != luma height {}",
chroma_shape[0], luma_shape[0]
)));
}
}
_ => {
return Err(Error::InvalidArgument(format!(
"from_planes only supports NV12 and NV16 (NV24 multiplane not yet \
supported — use a contiguous NV24 tensor), got {format:?}"
)));
}
}
Ok(Tensor {
storage: luma.storage,
format: Some(format),
chroma: Some(Box::new(chroma)),
row_stride: luma.row_stride,
plane_offset: luma.plane_offset,
quantization: luma.quantization,
cuda: None,
colorimetry: luma.colorimetry,
cpu_access: luma.cpu_access,
compression: luma.compression,
view_origin: None,
})
}
pub fn is_multiplane(&self) -> bool {
self.chroma.is_some()
}
pub fn chroma(&self) -> Option<&Tensor<T>> {
self.chroma.as_deref()
}
pub fn chroma_mut(&mut self) -> Option<&mut Tensor<T>> {
self.chroma.as_deref_mut()
}
pub(crate) fn clone_chroma_plane(&self) -> Result<Self> {
self.subview(0, self.shape())
}
pub fn row_stride(&self) -> Option<usize> {
self.row_stride
}
pub fn effective_row_stride(&self) -> Option<usize> {
if let Some(s) = self.row_stride {
return Some(s);
}
let fmt = self.format?;
let w = self.width()?;
let elem = std::mem::size_of::<T>();
Some(match fmt.layout() {
PixelLayout::Packed => w * fmt.channels() * elem,
PixelLayout::Planar => w * elem,
PixelLayout::SemiPlanar => w.next_multiple_of(2) * elem,
})
}
pub fn copy_to_flat(&self, dst: &mut [u8]) -> Result<()> {
let tight_bytes = crate::ahardwarebuffer_layout::checked_shape_bytes::<T>(self.shape())?;
if dst.len() != tight_bytes {
return Err(Error::InvalidArgument(format!(
"copy_to_flat: dst is {} bytes but the tensor's tight \
footprint is {tight_bytes} bytes (shape {:?})",
dst.len(),
self.shape()
)));
}
let map = self.map_read()?;
let src: &[u8] = unsafe {
std::slice::from_raw_parts(
map.as_slice().as_ptr() as *const u8,
std::mem::size_of_val(map.as_slice()),
)
};
let Some(stride) = self.row_stride else {
let Some(tight) = src.get(..tight_bytes) else {
return Err(Error::InvalidOperation(format!(
"copy_to_flat: mapped {} bytes < tight footprint {tight_bytes}",
src.len()
)));
};
dst.copy_from_slice(tight);
return Ok(());
};
let rows = strided_rows(self)?;
if rows == 0 {
return Err(Error::InvalidOperation(
"copy_to_flat: a strided tensor with no rows has nothing to copy".into(),
));
}
let row_bytes = tight_bytes / rows;
let need = (rows - 1)
.checked_mul(stride)
.and_then(|b| b.checked_add(row_bytes))
.ok_or_else(|| {
Error::InvalidOperation(format!(
"copy_to_flat: stride {stride} × rows {rows} overflows usize"
))
})?;
if src.len() < need {
return Err(Error::InvalidOperation(format!(
"copy_to_flat: mapped {} bytes but strided rows need {need}",
src.len()
)));
}
for r in 0..rows {
dst[r * row_bytes..(r + 1) * row_bytes]
.copy_from_slice(&src[r * stride..r * stride + row_bytes]);
}
Ok(())
}
pub fn set_row_stride(&mut self, stride: usize) -> Result<()> {
let fmt = self.format.ok_or_else(|| {
Error::InvalidArgument("cannot set row_stride without a pixel format".into())
})?;
let w = self.width().ok_or_else(|| {
Error::InvalidArgument("cannot determine width for row_stride validation".into())
})?;
let elem = std::mem::size_of::<T>();
let min_stride = match fmt.layout() {
PixelLayout::Packed => w * fmt.channels() * elem,
PixelLayout::Planar => w * elem,
PixelLayout::SemiPlanar => w.next_multiple_of(2) * elem,
};
if stride < min_stride {
return Err(Error::InvalidArgument(format!(
"row_stride {stride} < minimum {min_stride} for {fmt:?} at width {w}"
)));
}
self.row_stride = Some(stride);
Ok(())
}
pub fn set_row_stride_unchecked(&mut self, stride: usize) {
self.row_stride = Some(stride);
}
pub fn with_row_stride(mut self, stride: usize) -> Result<Self> {
self.set_row_stride(stride)?;
Ok(self)
}
pub fn plane_offset(&self) -> Option<usize> {
self.plane_offset
}
pub fn view_origin(&self) -> Option<ViewOrigin> {
self.view_origin
}
pub fn set_plane_offset(&mut self, offset: usize) {
self.plane_offset = Some(offset);
match self.storage {
TensorStorage::Mem(ref mut m) => m.set_offset(offset),
#[cfg(target_os = "linux")]
TensorStorage::Dma(ref mut dma) => dma.mmap_offset = offset,
#[cfg(any(target_os = "macos", target_os = "ios"))]
TensorStorage::Dma(ref mut io) => io.view_offset = offset,
#[cfg(target_os = "windows")]
TensorStorage::Dma(ref mut tex) => tex.set_view_offset(offset),
TensorStorage::Pbo(ref mut p) => p.view_offset = offset,
#[cfg(unix)]
TensorStorage::Shm(ref mut s) => s.set_offset(offset),
#[cfg(target_os = "android")]
TensorStorage::Dma(_) => {}
}
}
pub fn cpu_access(&self) -> CpuAccess {
self.cpu_access
}
#[cfg(target_os = "android")]
pub(crate) fn set_cpu_access_unchecked(&mut self, access: CpuAccess) {
self.cpu_access = access;
}
pub fn compression(&self) -> Option<CompressionScheme> {
self.compression
}
#[cfg(target_os = "android")]
pub(crate) fn set_compression_unchecked(&mut self, scheme: Option<CompressionScheme>) {
self.compression = scheme;
}
pub fn colorimetry(&self) -> Option<crate::Colorimetry> {
self.colorimetry
}
pub fn set_colorimetry(&mut self, c: Option<crate::Colorimetry>) {
self.colorimetry = c;
}
pub fn with_colorimetry(mut self, c: crate::Colorimetry) -> Self {
self.colorimetry = Some(c);
self
}
pub(crate) fn subview(&self, offset_bytes: usize, shape: &[usize]) -> Result<Tensor<T>> {
let abs_offset = self
.plane_offset
.unwrap_or(0)
.checked_add(offset_bytes)
.ok_or(Error::InvalidSize(offset_bytes))?;
let mut t = Tensor::wrap(self.storage.view(offset_bytes, shape)?);
if let Some(fmt) = self.format {
t.set_format(fmt)?;
}
if let Some(rs) = self.row_stride {
t.set_row_stride_unchecked(rs);
}
t.quantization = self.quantization.clone();
t.set_colorimetry(self.colorimetry);
t.cpu_access = self.cpu_access;
t.compression = self.compression;
if abs_offset > 0 {
t.set_plane_offset(abs_offset);
}
Ok(t)
}
pub fn batch(&self, n: usize) -> Result<Tensor<T>> {
let shape = self.shape();
if let Some(fmt) = self.format {
let elem_rank = image_rank(fmt.layout());
if shape.len() != elem_rank + 1 {
return Err(Error::InvalidShape(format!(
"batch(): tensor is not batched ({fmt:?} expects a leading N over a \
{elem_rank}-D element, got shape {shape:?})"
)));
}
}
let batch = *shape
.first()
.ok_or_else(|| Error::InvalidShape("batch(): empty shape".into()))?;
if n >= batch {
return Err(Error::BatchIndexOutOfBounds { index: n, batch });
}
let elem_shape: Vec<usize> = shape[1..].to_vec();
let elem_count: usize = elem_shape.iter().product();
elem_count
.checked_mul(std::mem::size_of::<T>())
.ok_or(Error::InvalidSize(elem_count))?;
let dims: Vec<u64> = shape.iter().map(|&d| d as u64).collect();
let elem_bytes = crate::protocol::c_byte_strides(
&dims,
std::mem::size_of::<T>() as i64,
self.format.map(|f| f.layout()),
self.row_stride,
)[0] as usize;
let offset = n.checked_mul(elem_bytes).ok_or(Error::InvalidSize(n))?;
let view_origin = match self.format.map(|f| f.layout()) {
Some(PixelLayout::Packed) => {
let tile_h = elem_shape[0];
let tile_w = elem_shape[1];
let parent_stride = self
.effective_row_stride()
.expect("a formatted tensor of its batched rank has a width");
Some(self.compose_view_origin(tile_w, batch * tile_h, parent_stride, 0, n * tile_h))
}
_ => None,
};
let mut t = self.subview(offset, &elem_shape)?;
t.view_origin = view_origin;
Ok(t)
}
pub fn view(&self, region: Region) -> Result<Tensor<T>> {
let fmt = self.format.ok_or_else(|| {
Error::InvalidOperation("view() requires a formatted image tensor".into())
})?;
if fmt.layout() != PixelLayout::Packed {
return Err(Error::InvalidOperation(format!(
"view() supports packed formats only (got {fmt:?}); use batch(n) for batched \
planar tensors"
)));
}
let w = self
.width()
.ok_or_else(|| Error::InvalidOperation("view(): tensor has no image width".into()))?;
let h = self
.height()
.ok_or_else(|| Error::InvalidOperation("view(): tensor has no image height".into()))?;
if !region.fits_within(w, h) {
return Err(Error::RegionOutOfBounds {
region,
bounds: (w, h),
});
}
let elem = std::mem::size_of::<T>();
let bpp = fmt.channels() * elem;
let stride = self
.effective_row_stride()
.expect("a tensor with an image width has a row stride");
let offset = region
.y
.checked_mul(stride)
.and_then(|yo| yo.checked_add(region.x.checked_mul(bpp)?))
.ok_or(Error::InvalidSize(region.y))?;
let sub_shape = fmt
.allocation_shape(region.width, region.height)
.ok_or_else(|| Error::InvalidShape(format!("view(): invalid shape for {fmt:?}")))?;
let mut t = self.subview(offset, &sub_shape)?;
let view_stride = if region.height > 1 {
stride
} else {
region.width * bpp
};
t.set_row_stride_unchecked(view_stride);
t.view_origin = Some(self.compose_view_origin(w, h, stride, region.x, region.y));
Ok(t)
}
fn compose_view_origin(
&self,
parent_width: usize,
parent_height: usize,
parent_row_stride: usize,
x: usize,
y: usize,
) -> ViewOrigin {
match self.view_origin {
Some(root) => ViewOrigin {
parent_width: root.parent_width,
parent_height: root.parent_height,
parent_row_stride: root.parent_row_stride,
x: root.x.saturating_add(x),
y: root.y.saturating_add(y),
},
None => ViewOrigin {
parent_width,
parent_height,
parent_row_stride,
x,
y,
},
}
}
pub(crate) fn as_pbo(&self) -> Option<&PboTensor<T>> {
match &self.storage {
TensorStorage::Pbo(p) => Some(p),
_ => None,
}
}
pub fn pbo_id(&self) -> Option<u32> {
self.as_pbo().map(PboTensor::buffer_id)
}
pub fn pbo_is_mapped(&self) -> Option<bool> {
self.as_pbo().map(PboTensor::is_mapped)
}
pub fn pbo_vtable_ptr(&self) -> Option<*const std::ffi::c_void> {
self.as_pbo()
.map(|p| p.pbo_vtable() as *const _ as *const std::ffi::c_void)
}
pub fn pbo_keepalive(&self) -> Option<Arc<dyn Send + Sync>> {
self.as_pbo().map(PboTensor::pbo_keepalive)
}
#[cfg(target_os = "linux")]
pub fn as_dma(&self) -> Option<&DmaTensor<T>> {
match &self.storage {
TensorStorage::Dma(d) => Some(d),
_ => None,
}
}
#[cfg(target_os = "linux")]
pub fn dmabuf(&self) -> Result<std::os::fd::BorrowedFd<'_>> {
use std::os::fd::AsFd;
match &self.storage {
TensorStorage::Dma(dma) => Ok(dma.fd.as_fd()),
_ => Err(Error::NotImplemented(format!(
"dmabuf requires DMA-backed tensor, got {:?}",
self.storage.memory()
))),
}
}
pub fn from_pbo(pbo: PboTensor<T>) -> Result<Self> {
Ok(Self {
storage: TensorStorage::Pbo(pbo),
format: None,
chroma: None,
row_stride: None,
plane_offset: None,
quantization: None,
cuda: None,
colorimetry: None,
cpu_access: CpuAccess::ReadWrite,
compression: None,
view_origin: None,
})
}
pub fn cuda(&self) -> Option<&crate::cuda::CudaHandle> {
self.cuda.as_ref()
}
pub fn is_cuda_attached(&self) -> bool {
self.cuda.is_some()
}
pub fn set_cuda_handle(&mut self, h: crate::cuda::CudaHandle) {
self.cuda = Some(h);
}
pub fn cuda_map(&self) -> Option<crate::cuda::CudaMap<'_>> {
self.cuda.as_ref()?.map()
}
pub fn cuda_map_mut(&self) -> Option<crate::cuda::CudaMap<'_>> {
self.cuda.as_ref()?.map_mut()
}
#[cfg(target_os = "linux")]
pub fn try_init_dma_cuda(&mut self) {
if self.cuda.is_some() || !crate::cuda::is_cuda_available() {
return;
}
let (raw_fd, buf_size) = match &self.storage {
TensorStorage::Dma(dma) => {
use std::os::fd::AsRawFd;
(dma.fd.as_raw_fd(), dma.buf_size)
}
_ => return,
};
if let Some((ext, dptr)) = crate::cuda::import_dma_fd(raw_fd, buf_size) {
self.cuda = Some(crate::cuda::CudaHandle::new_external(ext, dptr, buf_size));
}
}
#[cfg(target_os = "windows")]
pub fn try_init_d3d11_cuda(&mut self) {
if self.cuda.is_some() {
return;
}
let TensorStorage::Dma(tex) = &self.storage else {
return;
};
if let Some(h) = tex.try_init_cuda() {
self.cuda = Some(h);
}
}
}
#[cfg(all(feature = "static", target_os = "windows"))]
impl<T> Tensor<T>
where
T: Num + Clone + fmt::Debug + Send + Sync,
{
pub fn d3d11_texture(&self) -> Option<*mut std::ffi::c_void> {
match &self.storage {
TensorStorage::Dma(t) => Some(t.texture_ptr()),
_ => None,
}
}
pub fn d3d11_layout(&self) -> Option<crate::d3d11_layout::D3d11ImageLayout> {
match &self.storage {
TensorStorage::Dma(t) => Some(t.layout()),
_ => None,
}
}
pub fn d3d11_shared_handle(&self) -> Result<std::os::windows::io::OwnedHandle> {
match &self.storage {
TensorStorage::Dma(t) => t.shared_handle(),
_ => Err(Error::NotImplemented(
"not a D3D11 texture tensor".to_owned(),
)),
}
}
pub fn gpu_completion(&self) -> Result<Option<crate::d3d11::GpuCompletion>> {
match &self.storage {
TensorStorage::Dma(t) => t.gpu_completion(),
_ => Err(Error::NotImplemented(
"not a D3D11 texture tensor".to_owned(),
)),
}
}
pub fn set_gpu_write(&self, value: u64) -> Result<()> {
match &self.storage {
TensorStorage::Dma(t) => {
t.set_gpu_write(value);
Ok(())
}
_ => Err(Error::NotImplemented(
"not a D3D11 texture tensor".to_owned(),
)),
}
}
pub(crate) fn d3d11_shared_handle_value(&self) -> Option<usize> {
match &self.storage {
TensorStorage::Dma(t) => Some(t.shared_handle_value()),
_ => None,
}
}
}
#[cfg(feature = "static")]
impl<T> Tensor<T>
where
T: Num + Clone + fmt::Debug + Send + Sync,
{
pub fn gpu_write_value(&self) -> u64 {
#[cfg(target_os = "windows")]
{
if let TensorStorage::Dma(t) = &self.storage {
return t.last_gpu_write();
}
}
0
}
}
#[cfg(all(feature = "static", target_os = "windows"))]
fn record_d3d11_row_stride<T>(t: &mut Tensor<T>, format: PixelFormat, what: &str)
where
T: Num + Clone + fmt::Debug + Send + Sync,
{
let pitch = match &t.storage {
TensorStorage::Dma(tex) => tex.image_backing_row_stride(),
_ => None,
};
let (Some(pitch), Some(natural)) = (pitch, t.effective_row_stride()) else {
return;
};
if pitch > natural {
log::debug!(
"{what}: the D3D11 backing pitch for the {format:?} texture is {pitch} bytes \
(natural {natural}); recording row stride"
);
t.set_row_stride_unchecked(pitch);
}
}
#[cfg(all(feature = "static", target_os = "windows"))]
impl<T> Tensor<T>
where
T: Num + Clone + fmt::Debug + Send + Sync + 'static,
{
pub unsafe fn from_d3d11_texture(
texture: *mut std::ffi::c_void,
width: usize,
height: usize,
format: PixelFormat,
access: CpuAccess,
name: Option<&str>,
) -> Result<Self> {
let dtype = dtype_of::<T>().ok_or_else(|| {
Error::InvalidArgument(format!(
"from_d3d11_texture: element type {} has no DType",
std::any::type_name::<T>()
))
})?;
let shape = format.allocation_shape(width, height).ok_or_else(|| {
Error::InvalidArgument(format!(
"invalid dimensions {width}x{height} for format {format:?}"
))
})?;
let storage = unsafe {
crate::d3d11::texture::D3d11TextureTensor::<T>::from_d3d11_texture(
texture, width, height, format, dtype, &shape, access, name,
)
}?;
let mut t = Self::wrap(TensorStorage::Dma(storage));
t.format = Some(format);
record_d3d11_row_stride(&mut t, format, "from_d3d11_texture");
t.cpu_access = access;
t.try_init_d3d11_cuda();
Ok(t)
}
pub unsafe fn from_d3d11_shared_handle(
handle: std::os::windows::io::RawHandle,
width: usize,
height: usize,
format: PixelFormat,
access: CpuAccess,
completion: Option<(std::os::windows::io::RawHandle, u64)>,
name: Option<&str>,
) -> Result<Self> {
let dtype = dtype_of::<T>().ok_or_else(|| {
Error::InvalidArgument(format!(
"from_d3d11_shared_handle: element type {} has no DType",
std::any::type_name::<T>()
))
})?;
let shape = format.allocation_shape(width, height).ok_or_else(|| {
Error::InvalidArgument(format!(
"invalid dimensions {width}x{height} for format {format:?}"
))
})?;
let storage = unsafe {
crate::d3d11::texture::D3d11TextureTensor::<T>::from_d3d11_shared_handle(
handle, width, height, format, dtype, &shape, access, completion, name,
)
}?;
let mut t = Self::wrap(TensorStorage::Dma(storage));
t.format = Some(format);
record_d3d11_row_stride(&mut t, format, "from_d3d11_shared_handle");
t.cpu_access = access;
t.try_init_d3d11_cuda();
Ok(t)
}
}
#[cfg(feature = "static")]
impl<T> Tensor<T>
where
T: IntegerType + Num + Clone + fmt::Debug + Send + Sync,
{
pub fn quantization(&self) -> Option<&Quantization> {
self.quantization.as_ref()
}
pub fn set_quantization(&mut self, q: Quantization) -> Result<()> {
q.validate(self.shape())?;
self.quantization = Some(q);
Ok(())
}
pub fn with_quantization(mut self, q: Quantization) -> Result<Self> {
self.set_quantization(q)?;
Ok(self)
}
pub fn clear_quantization(&mut self) {
self.quantization = None;
}
}
#[cfg(feature = "static")]
fn image_rank(layout: PixelLayout) -> usize {
match layout {
PixelLayout::SemiPlanar => 2,
PixelLayout::Packed | PixelLayout::Planar => 3,
}
}
#[cfg(feature = "static")]
fn map_strided_host_view<'a, T>(
tensor: &Tensor<T>,
stride: usize,
access: CpuAccess,
non_blocking: bool,
) -> Result<crate::view::HostView<'a, T>>
where
T: Num + Clone + fmt::Debug + Send + Sync + 'a,
{
let rows = strided_rows(tensor)?;
let total_bytes = stride.checked_mul(rows).ok_or_else(|| {
Error::InvalidOperation(format!(
"Tensor::map: row_stride {stride} × rows {rows} overflows usize"
))
})?;
map_strided_storage(tensor, stride, rows, total_bytes, access, non_blocking)
}
#[cfg(feature = "static")]
fn strided_rows<T>(tensor: &Tensor<T>) -> Result<usize>
where
T: Num + Clone + fmt::Debug + Send + Sync,
{
let shape = tensor.shape();
let Some(fmt) = tensor.format else {
return shape.first().copied().ok_or_else(|| {
Error::InvalidOperation("strided access requires a non-empty shape".into())
});
};
let (_, dims) = tensor.image_dims().ok_or_else(|| {
Error::InvalidOperation(format!(
"strided access to a {fmt:?} tensor requires its image shape, optionally \
batched; got {shape:?}"
))
})?;
let per_image = match fmt.layout() {
PixelLayout::Planar => dims[0] * dims[1],
PixelLayout::Packed | PixelLayout::SemiPlanar => dims[0],
};
let batch: usize = shape[..shape.len() - dims.len()].iter().product();
Ok(batch * per_image)
}
#[cfg(feature = "static")]
fn map_strided_storage<'a, T>(
tensor: &Tensor<T>,
_stride: usize,
_rows: usize,
total_bytes: usize,
access: CpuAccess,
_non_blocking: bool,
) -> Result<crate::view::HostView<'a, T>>
where
T: Num + Clone + fmt::Debug + Send + Sync + 'a,
{
match &tensor.storage {
#[cfg(target_os = "linux")]
TensorStorage::Dma(dma) => {
let available_bytes = dma.buf_size.saturating_sub(dma.mmap_offset);
if total_bytes > available_bytes {
return Err(Error::InvalidOperation(format!(
"Tensor::map: strided mapping needs {total_bytes} bytes \
but DMA buffer only has {available_bytes} available \
(buf_size={}, mmap_offset={}, stride={_stride}, rows={_rows}); \
the row_stride was likely set larger than the original allocation",
dma.buf_size, dma.mmap_offset
)));
}
dma.map_with_byte_size(total_bytes, access)
}
TensorStorage::Mem(mem) => {
let capacity = tensor.storage.capacity_bytes();
if total_bytes > capacity {
return Err(Error::InsufficientCapacity {
needed: total_bytes,
capacity,
});
}
mem.map_with_byte_size(total_bytes, access)
}
#[cfg(unix)]
TensorStorage::Shm(shm) => {
let capacity = tensor.storage.capacity_bytes();
if total_bytes > capacity {
return Err(Error::InsufficientCapacity {
needed: total_bytes,
capacity,
});
}
shm.map_with_byte_size(total_bytes, access)
}
#[cfg(any(target_os = "macos", target_os = "ios"))]
TensorStorage::Dma(io) => {
let available = io.buf_size.saturating_sub(io.view_offset);
if total_bytes > available {
return Err(Error::InsufficientCapacity {
needed: total_bytes,
capacity: available,
});
}
io.map_with_byte_size(total_bytes, access)
}
#[cfg(target_os = "android")]
TensorStorage::Dma(ahb) => {
let available = ahb.buf_size.saturating_sub(ahb.view_offset);
if total_bytes > available {
return Err(Error::InsufficientCapacity {
needed: total_bytes,
capacity: available,
});
}
ahb.map_with_byte_size(total_bytes, access)
}
#[cfg(target_os = "windows")]
TensorStorage::Dma(tex) => {
if _non_blocking {
tex.try_map_with_byte_size(total_bytes, access)
} else {
tex.map_with_byte_size(total_bytes, access)
}
}
TensorStorage::Pbo(pbo) => {
let available = pbo.capacity_bytes().saturating_sub(pbo.view_offset);
if total_bytes > available {
return Err(Error::InsufficientCapacity {
needed: total_bytes,
capacity: available,
});
}
pbo.map_with_byte_size(total_bytes, access)
}
#[allow(unreachable_patterns)]
_ => Err(Error::InvalidOperation(
"CPU mapping of strided tensors is supported only for HAL-allocated \
Mem/Shm (any platform), DMA (Linux, self-allocated or imported), \
IOSurface (macOS), AHardwareBuffer (Android), D3D11 textures \
(Windows), and PBO"
.into(),
)),
}
}
#[cfg(feature = "static")]
impl<T> Tensor<T>
where
T: Num + Clone + fmt::Debug + Send + Sync,
{
fn map_impl<'a>(
&self,
access: CpuAccess,
non_blocking: bool,
) -> Result<crate::view::HostView<'a, T>>
where
T: 'a,
{
let _span = tracing::trace_span!(
"tensor.map",
memory = ?self.storage.memory(),
?access,
non_blocking,
)
.entered();
if access == CpuAccess::None {
return Err(Error::InvalidArgument(
"map_with(CpuAccess::None) is not a mappable direction — use \
map_read()/map_write()/map_mut()"
.into(),
));
}
if !self.cpu_access.covers(access) {
note_unplanned_cpu_access(
self.buffer_identity().id(),
&format!("{:?}", self.storage.memory()),
"map access exceeds the declared CpuAccess",
);
}
if let Some(stride) = self.row_stride {
return map_strided_host_view(self, stride, access, non_blocking);
}
if non_blocking {
self.storage.try_map_with(access)
} else {
self.storage.map_with(access)
}
}
}
#[cfg(all(
feature = "static",
any(
target_os = "macos",
target_os = "ios",
target_os = "android",
target_os = "windows"
)
))]
impl<T> Tensor<T>
where
T: Num + Clone + fmt::Debug + Send + Sync,
{
fn adopt_backing_pitch(&mut self, shape: &[usize]) {
let Some(pitch) = self.storage.backing_row_stride() else {
return;
};
let Some((rank, rows)) = self.format.map(|f| match f.layout() {
PixelLayout::Packed => (3, shape.first().copied()),
PixelLayout::Planar => (3, shape.get(1).copied()),
PixelLayout::SemiPlanar => (2, shape.first().copied()),
}) else {
return;
};
if rank != shape.len() {
return;
}
let prior = self.row_stride.take();
let Some(natural) = self.effective_row_stride() else {
self.row_stride = prior;
return;
};
if rows == Some(1) {
self.set_row_stride_unchecked(natural);
} else if pitch > natural {
self.set_row_stride_unchecked(pitch);
} else {
self.row_stride = prior;
}
}
}
#[cfg(feature = "static")]
impl<T> TensorTrait<T> for Tensor<T>
where
T: Num + Clone + fmt::Debug + Send + Sync,
{
fn new(shape: &[usize], name: Option<&str>) -> Result<Self>
where
Self: Sized,
{
Self::new(shape, None, name)
}
#[cfg(unix)]
fn from_fd(fd: std::os::fd::OwnedFd, shape: &[usize], name: Option<&str>) -> Result<Self>
where
Self: Sized,
{
#[cfg_attr(not(target_os = "linux"), allow(unused_mut))]
let mut t = Self::wrap(TensorStorage::from_fd(fd, shape, name)?);
#[cfg(target_os = "linux")]
t.try_init_dma_cuda();
Ok(t)
}
#[cfg(unix)]
fn clone_fd(&self) -> Result<std::os::fd::OwnedFd> {
self.storage.clone_fd()
}
fn memory(&self) -> TensorMemory {
self.storage.memory()
}
fn name(&self) -> String {
self.storage.name()
}
fn shape(&self) -> &[usize] {
self.storage.shape()
}
fn capacity_bytes(&self) -> usize {
self.storage.capacity_bytes()
}
fn reshape(&mut self, shape: &[usize]) -> Result<()> {
if self.chroma.is_some() {
return Err(Error::InvalidOperation(
"cannot reshape a multiplane tensor — decompose planes first".into(),
));
}
self.storage.reshape(shape)?;
self.format = None;
self.row_stride = None;
self.plane_offset = None;
match self.storage {
TensorStorage::Mem(ref mut m) => m.set_offset(0),
#[cfg(target_os = "linux")]
TensorStorage::Dma(ref mut dma) => dma.mmap_offset = 0,
#[cfg(any(target_os = "macos", target_os = "ios"))]
TensorStorage::Dma(_) => {}
#[cfg(target_os = "windows")]
TensorStorage::Dma(ref mut tex) => tex.set_view_offset(0),
#[cfg(target_os = "android")]
TensorStorage::Dma(_) => {}
TensorStorage::Pbo(_) => {}
#[cfg(unix)]
TensorStorage::Shm(ref mut s) => s.set_offset(0),
}
Ok(())
}
fn set_logical_shape(&mut self, shape: &[usize]) -> Result<()> {
self.storage.set_logical_shape(shape)?;
#[cfg(any(
target_os = "macos",
target_os = "ios",
target_os = "android",
target_os = "windows"
))]
self.adopt_backing_pitch(shape);
Ok(())
}
fn map_with<'a>(&self, access: CpuAccess) -> Result<crate::view::HostView<'a, T>>
where
T: 'a,
{
self.map_impl(access, false)
}
fn try_map_with<'a>(&self, access: CpuAccess) -> Result<crate::view::HostView<'a, T>>
where
T: 'a,
{
self.map_impl(access, true)
}
fn buffer_identity(&self) -> &BufferIdentity {
self.storage.buffer_identity()
}
}
#[cfg(target_os = "linux")]
static DMA_AVAILABLE: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
#[cfg(any(target_os = "macos", target_os = "ios"))]
static IOSURFACE_AVAILABLE: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
#[cfg(target_os = "linux")]
pub fn is_dma_available() -> bool {
*DMA_AVAILABLE
.get_or_init(|| Tensor::<u8>::new(&[64], Some(TensorMemory::DmaBuf), None).is_ok())
}
#[cfg(not(target_os = "linux"))]
pub fn is_dma_available() -> bool {
false
}
#[cfg(any(target_os = "macos", target_os = "ios"))]
pub fn is_iosurface_available() -> bool {
*IOSURFACE_AVAILABLE.get_or_init(|| {
Tensor::<u8>::new(&[64], Some(TensorMemory::DmaBuf), None).is_ok()
})
}
#[cfg(not(any(target_os = "macos", target_os = "ios")))]
pub fn is_iosurface_available() -> bool {
false
}
#[cfg(target_os = "android")]
static AHARDWAREBUFFER_AVAILABLE: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
#[cfg(target_os = "android")]
pub fn is_ahardwarebuffer_available() -> bool {
*AHARDWAREBUFFER_AVAILABLE.get_or_init(|| {
Tensor::<u8>::new(&[64], Some(TensorMemory::DmaBuf), None).is_ok()
})
}
#[cfg(not(target_os = "android"))]
pub fn is_ahardwarebuffer_available() -> bool {
false
}
pub fn is_gpu_buffer_available() -> bool {
#[cfg(target_os = "linux")]
{
is_dma_available()
}
#[cfg(any(target_os = "macos", target_os = "ios"))]
{
is_iosurface_available()
}
#[cfg(target_os = "android")]
{
is_ahardwarebuffer_available()
}
#[cfg(target_os = "windows")]
{
crate::d3d11::device().is_ok()
}
#[cfg(not(any(
target_os = "linux",
target_os = "macos",
target_os = "ios",
target_os = "android",
target_os = "windows"
)))]
{
false
}
}
#[cfg(unix)]
static SHM_AVAILABLE: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
#[cfg(unix)]
pub fn is_shm_available() -> bool {
*SHM_AVAILABLE.get_or_init(|| Tensor::<u8>::new(&[64], Some(TensorMemory::Shm), None).is_ok())
}
#[cfg(not(unix))]
pub fn is_shm_available() -> bool {
false
}
#[cfg(test)]
mod dtype_tests {
use super::*;
#[test]
fn dtype_size() {
assert_eq!(DType::U8.size(), 1);
assert_eq!(DType::I8.size(), 1);
assert_eq!(DType::U16.size(), 2);
assert_eq!(DType::I16.size(), 2);
assert_eq!(DType::U32.size(), 4);
assert_eq!(DType::I32.size(), 4);
assert_eq!(DType::U64.size(), 8);
assert_eq!(DType::I64.size(), 8);
assert_eq!(DType::F16.size(), 2);
assert_eq!(DType::F32.size(), 4);
assert_eq!(DType::F64.size(), 8);
}
#[test]
fn dtype_name() {
assert_eq!(DType::U8.name(), "u8");
assert_eq!(DType::F16.name(), "f16");
assert_eq!(DType::F32.name(), "f32");
}
#[test]
fn dtype_serde_roundtrip() {
use serde_json;
let dt = DType::F16;
let json = serde_json::to_string(&dt).unwrap();
let back: DType = serde_json::from_str(&json).unwrap();
assert_eq!(dt, back);
}
}
#[cfg(test)]
mod image_tests {
use super::*;
#[test]
fn image_dims_from_shape_inverts_allocation_shape() {
for (format, w, h) in [
(PixelFormat::Rgba, 640usize, 480usize),
(PixelFormat::Rgb, 641, 3),
(PixelFormat::Grey, 32, 32),
(PixelFormat::PlanarRgb, 640, 480),
(PixelFormat::PlanarRgba, 8, 5),
(PixelFormat::Nv12, 640, 481),
(PixelFormat::Nv16, 33, 17),
(PixelFormat::Nv24, 16, 9),
] {
let shape = format.allocation_shape(w, h).unwrap();
assert_eq!(
image_dims_from_shape(format, &shape).unwrap(),
(w, h),
"{format:?} {w}x{h} via {shape:?}"
);
}
}
#[test]
fn image_dims_from_shape_rejects_the_wrong_rank() {
for (format, shape) in [
(PixelFormat::Rgba, vec![480usize, 640usize]),
(PixelFormat::PlanarRgb, vec![3, 480]),
(PixelFormat::Nv12, vec![720, 640, 1]),
] {
let err = image_dims_from_shape(format, &shape).unwrap_err();
assert!(
matches!(err, Error::InvalidShape(ref m) if m.contains("rank")),
"{format:?} {shape:?}: {err}"
);
}
}
#[test]
fn image_dims_from_shape_rejects_an_unproduced_combined_height() {
let err = image_dims_from_shape(PixelFormat::Nv12, &[721, 640]).unwrap_err();
assert!(
matches!(err, Error::InvalidShape(ref m) if m.contains("combined plane height")),
"{err}"
);
}
#[test]
fn allocation_shape_per_layout() {
assert_eq!(
PixelFormat::Rgb.allocation_shape(640, 480),
Some(vec![480, 640, 3])
);
assert_eq!(
PixelFormat::Grey.allocation_shape(640, 480),
Some(vec![480, 640, 1])
);
assert_eq!(
PixelFormat::Nv12.allocation_shape(640, 480),
Some(vec![720, 640])
);
assert_eq!(
PixelFormat::Nv12.allocation_shape(640, 481),
Some(vec![722, 640])
);
assert_eq!(
PixelFormat::Nv12.allocation_shape(641, 480),
Some(vec![720, 641])
);
assert_eq!(
PixelFormat::Nv16.allocation_shape(641, 480),
Some(vec![960, 641])
);
assert_eq!(
PixelFormat::PlanarRgb.allocation_shape(640, 480),
Some(vec![3, 480, 640])
);
assert_eq!(
PixelFormat::Nv16.allocation_shape(640, 480),
Some(vec![960, 640])
);
}
#[test]
fn raw_tensor_has_no_format() {
let t = Tensor::<u8>::new(&[480, 640, 3], None, None).unwrap();
assert!(t.format().is_none());
assert!(t.width().is_none());
assert!(t.height().is_none());
assert!(!t.is_multiplane());
assert!(t.chroma().is_none());
}
#[test]
fn image_tensor_packed() {
let t = Tensor::<u8>::image(
640,
480,
PixelFormat::Rgba,
None,
crate::CpuAccess::ReadWrite,
)
.unwrap();
assert_eq!(t.format(), Some(PixelFormat::Rgba));
assert_eq!(t.width(), Some(640));
assert_eq!(t.height(), Some(480));
assert_eq!(t.shape(), &[480, 640, 4]);
assert!(!t.is_multiplane());
}
#[test]
fn image_tensor_planar() {
let t = Tensor::<u8>::image(
640,
480,
PixelFormat::PlanarRgb,
None,
crate::CpuAccess::ReadWrite,
)
.unwrap();
assert_eq!(t.format(), Some(PixelFormat::PlanarRgb));
assert_eq!(t.width(), Some(640));
assert_eq!(t.height(), Some(480));
assert_eq!(t.shape(), &[3, 480, 640]);
}
#[test]
#[cfg(target_os = "macos")]
fn image_tensor_dma_non_aligned_packed_width_pads_zero_copy() {
let _lock = crate::tests::fd_lock_shared();
let t = Tensor::<u8>::image(
4,
4,
PixelFormat::Rgba,
Some(TensorMemory::DmaBuf),
crate::CpuAccess::ReadWrite,
)
.expect("padded RGBA IOSurface should allocate");
assert_eq!(t.format(), Some(PixelFormat::Rgba));
assert_eq!(t.width(), Some(4));
assert_eq!(t.height(), Some(4));
let stride = t.effective_row_stride().expect("stride");
assert_eq!(stride % 64, 0, "padded to 64-byte row alignment");
assert!(stride >= 16);
let m = t.map().expect("strided IOSurface map");
assert_eq!(m.as_slice().len(), stride * 4);
}
#[test]
#[cfg(target_os = "macos")]
fn image_tensor_dma_rejects_indivisible_pixel_pitch_without_pad_hint() {
let _lock = crate::tests::fd_lock_shared();
let err = Tensor::<f32>::image(
10,
10,
PixelFormat::Rgb,
Some(TensorMemory::DmaBuf),
crate::CpuAccess::ReadWrite,
)
.expect_err("RGB f32 with 12 B/pixel and non-aligned width must be rejected");
match err {
Error::InvalidArgument(msg) => {
assert!(
msg.contains("64-byte aligned"),
"error must still name the alignment requirement: {msg}"
);
assert!(
!msg.contains("Pad width"),
"indivisible per-pixel pitch makes a width suggestion impossible; \
hint must be omitted, got: {msg}"
);
assert!(
msg.contains("memory=None") && msg.contains("TensorMemory::Mem"),
"error must still list the always-applicable alternatives: {msg}"
);
}
other => panic!("expected InvalidArgument, got {other:?}"),
}
}
#[test]
#[cfg(target_os = "macos")]
fn image_tensor_dma_packed_rgb_u8_contract() {
let _lock = crate::tests::fd_lock_shared();
let err = Tensor::<u8>::image(
10,
10,
PixelFormat::Rgb,
Some(TensorMemory::DmaBuf),
crate::CpuAccess::ReadWrite,
)
.expect_err("Rgb u8 width%4!=0 must be rejected");
assert!(
matches!(&err, Error::InvalidArgument(m) if m.contains("width%4==0")),
"got {err:?}"
);
let t = Tensor::<u8>::image(
12,
4,
PixelFormat::Rgb,
Some(TensorMemory::DmaBuf),
crate::CpuAccess::ReadWrite,
)
.expect("width 12 Rgb u8 must allocate padded");
assert_eq!(t.memory(), TensorMemory::DmaBuf);
assert!(
t.row_stride()
.is_some_and(|s| s >= 64 && s.is_multiple_of(64)),
"padded pitch must be recorded: {:?}",
t.row_stride()
);
let t = Tensor::<u8>::image(
640,
8,
PixelFormat::Rgb,
Some(TensorMemory::DmaBuf),
crate::CpuAccess::ReadWrite,
)
.expect("width 640 Rgb u8 must allocate flat");
assert_eq!(t.row_stride(), None);
let t = Tensor::<i8>::image(
640,
8,
PixelFormat::Rgb,
Some(TensorMemory::DmaBuf),
crate::CpuAccess::ReadWrite,
)
.expect("Rgb i8 shares the RGBA8888 mapping");
assert_eq!(t.memory(), TensorMemory::DmaBuf);
}
#[test]
#[cfg(target_os = "macos")]
fn image_tensor_dma_planar_f16_alignment() {
let _lock = crate::tests::fd_lock_shared();
let err = Tensor::<half::f16>::image(
16,
16,
PixelFormat::PlanarRgb,
Some(TensorMemory::DmaBuf),
crate::CpuAccess::ReadWrite,
)
.expect_err("width=16 PlanarRgb F16 is 32-byte row, must reject");
assert!(matches!(err, Error::InvalidArgument(_)), "got {err:?}");
let t = Tensor::<half::f16>::image(
32,
8,
PixelFormat::PlanarRgb,
Some(TensorMemory::DmaBuf),
crate::CpuAccess::ReadWrite,
)
.expect("width=32 PlanarRgb F16 is 64-byte row, must succeed");
assert_eq!(t.format(), Some(PixelFormat::PlanarRgb));
}
#[test]
fn image_tensor_semi_planar_contiguous() {
let t = Tensor::<u8>::image(
640,
480,
PixelFormat::Nv12,
None,
crate::CpuAccess::ReadWrite,
)
.unwrap();
assert_eq!(t.format(), Some(PixelFormat::Nv12));
assert_eq!(t.width(), Some(640));
assert_eq!(t.height(), Some(480));
assert_eq!(t.shape(), &[720, 640]);
assert!(!t.is_multiplane());
}
#[test]
#[cfg(target_os = "linux")]
fn image_tensor_with_stride_preserves_logical_width() {
let _lock = crate::tests::fd_lock_shared();
if !crate::test_support::dma_or_skip("image_tensor_with_stride_preserves_logical_width") {
return;
}
let stride = 12032;
let t = Tensor::<u8>::image_with_stride(
3004,
1688,
PixelFormat::Rgba,
stride,
Some(TensorMemory::DmaBuf),
crate::CpuAccess::ReadWrite,
)
.unwrap();
assert_eq!(t.width(), Some(3004));
assert_eq!(t.height(), Some(1688));
assert_eq!(t.shape(), &[1688, 3004, 4]);
assert_eq!(t.effective_row_stride(), Some(stride));
use crate::TensorMapTrait;
{
let map = t.map().unwrap();
assert!(
map.as_slice().len() >= stride * 1688,
"mapped buffer {} bytes < expected {}",
map.as_slice().len(),
stride * 1688
);
}
{
let mut map = t.map().unwrap();
let slice = map.as_mut_slice();
for y in 0..1688 {
let row_start = y * stride;
for x in 0..3004 {
let p = row_start + x * 4;
slice[p] = (y & 0xFF) as u8;
slice[p + 1] = (x & 0xFF) as u8;
slice[p + 2] = 0x42;
slice[p + 3] = 0xFF;
}
}
}
{
let map = t.map().unwrap();
let slice = map.as_slice();
assert_eq!(slice[0], 0x00);
assert_eq!(slice[1], 0x00);
assert_eq!(slice[2], 0x42);
assert_eq!(slice[3], 0xFF);
let mid = 100 * stride + 50 * 4;
assert_eq!(slice[mid], 100);
assert_eq!(slice[mid + 1], 50);
assert_eq!(slice[mid + 2], 0x42);
}
}
#[test]
#[cfg(target_os = "linux")]
fn image_tensor_foreign_strided_map_is_bounds_checked_not_refused() {
if !crate::test_support::dma_or_skip(
"image_tensor_foreign_strided_map_is_bounds_checked_not_refused",
) {
return;
}
let backing =
Tensor::<u8>::new(&[240 * 320 * 4], Some(TensorMemory::DmaBuf), None).unwrap();
{
let fd = backing.clone_fd().unwrap();
let shape = [240usize, 320, 4];
let storage = TensorStorage::<u8>::from_fd(fd, &shape, None).unwrap();
let mut t = Tensor::<u8>::wrap(storage);
t.set_format(PixelFormat::Bgra).unwrap();
t.set_row_stride(320 * 4).unwrap();
assert!(
t.map().is_ok(),
"an in-bounds strided map of an imported DMA-BUF must be \
accepted -- refusing it is what broke decode-into-GPU-tensor"
);
}
{
let fd = backing.clone_fd().unwrap();
let shape = [240usize, 320, 4];
let storage = TensorStorage::<u8>::from_fd(fd, &shape, None).unwrap();
let mut t = Tensor::<u8>::wrap(storage);
t.set_format(PixelFormat::Bgra).unwrap();
t.set_row_stride(320 * 4 * 2).unwrap();
let err = t.map();
assert!(
matches!(err, Err(Error::InvalidOperation(_))),
"a strided map larger than the DMA allocation must be refused"
);
}
}
#[test]
#[cfg(target_os = "linux")]
fn image_tensor_with_stride_map_rejects_tampered_stride() {
let _lock = crate::tests::fd_lock_shared();
if !crate::test_support::dma_or_skip("image_tensor_with_stride_map_rejects_tampered_stride")
{
return;
}
let mut t = Tensor::<u8>::image_with_stride(
640,
480,
PixelFormat::Rgba,
3072,
Some(TensorMemory::DmaBuf),
crate::CpuAccess::ReadWrite,
)
.unwrap();
t.set_row_stride(12288).unwrap();
let err = t.map();
assert!(
matches!(err, Err(Error::InvalidOperation(_))),
"map() with oversized stride must return InvalidOperation"
);
}
#[test]
fn dma_tensor_new_with_byte_size_rejects_shape_overflow() {
#[cfg(target_os = "linux")]
{
let err = crate::dma::DmaTensor::<u64>::new_with_byte_size(
&[usize::MAX, 2, 2],
usize::MAX,
None,
);
assert!(
matches!(err, Err(Error::InvalidArgument(_))),
"new_with_byte_size must detect shape.product() overflow"
);
}
}
#[test]
#[cfg(target_os = "linux")]
fn image_tensor_with_stride_rejects_too_small_stride() {
let _lock = crate::tests::fd_lock_shared();
let err = Tensor::<u8>::image_with_stride(
640,
480,
PixelFormat::Rgba,
2400,
Some(TensorMemory::DmaBuf),
crate::CpuAccess::ReadWrite,
);
assert!(matches!(err, Err(Error::InvalidArgument(_))));
}
#[test]
#[cfg(target_os = "linux")]
fn image_tensor_with_stride_rejects_non_packed() {
let _lock = crate::tests::fd_lock_shared();
let err = Tensor::<u8>::image_with_stride(
640,
480,
PixelFormat::Nv12,
640,
Some(TensorMemory::DmaBuf),
crate::CpuAccess::ReadWrite,
);
assert!(matches!(err, Err(Error::NotImplemented(_))));
}
#[test]
fn set_format_valid() {
let mut t = Tensor::<u8>::new(&[480, 640, 3], None, None).unwrap();
assert!(t.format().is_none());
t.set_format(PixelFormat::Rgb).unwrap();
assert_eq!(t.format(), Some(PixelFormat::Rgb));
assert_eq!(t.width(), Some(640));
assert_eq!(t.height(), Some(480));
}
#[test]
fn set_format_invalid_shape() {
let mut t = Tensor::<u8>::new(&[480, 640, 4], None, None).unwrap();
let err = t.set_format(PixelFormat::Rgb);
assert!(err.is_err());
assert!(t.format().is_none());
}
#[test]
fn reshape_clears_format() {
let mut t = Tensor::<u8>::image(
640,
480,
PixelFormat::Rgba,
None,
crate::CpuAccess::ReadWrite,
)
.unwrap();
assert_eq!(t.format(), Some(PixelFormat::Rgba));
t.reshape(&[480 * 640 * 4]).unwrap();
assert!(t.format().is_none());
}
#[test]
fn from_planes_nv12() {
let y = Tensor::<u8>::new(&[480, 640], None, None).unwrap();
let uv = Tensor::<u8>::new(&[240, 640], None, None).unwrap();
let img = Tensor::from_planes(y, uv, PixelFormat::Nv12).unwrap();
assert_eq!(img.format(), Some(PixelFormat::Nv12));
assert!(img.is_multiplane());
assert!(img.chroma().is_some());
assert_eq!(img.width(), Some(640));
assert_eq!(img.height(), Some(480));
}
#[test]
fn from_planes_rejects_non_semiplanar() {
let y = Tensor::<u8>::new(&[480, 640], None, None).unwrap();
let uv = Tensor::<u8>::new(&[240, 640], None, None).unwrap();
let err = Tensor::from_planes(y, uv, PixelFormat::Rgb);
assert!(err.is_err());
}
#[test]
fn reshape_multiplane_errors() {
let y = Tensor::<u8>::new(&[480, 640], None, None).unwrap();
let uv = Tensor::<u8>::new(&[240, 640], None, None).unwrap();
let mut img = Tensor::from_planes(y, uv, PixelFormat::Nv12).unwrap();
let err = img.reshape(&[480 * 640 + 240 * 640]);
assert!(err.is_err());
}
}
#[cfg(test)]
mod compression_tests {
use super::*;
#[test]
fn desc_builder_roundtrips() {
let desc = ImageDesc::new(640, 480, PixelFormat::Rgba, DType::U8)
.with_memory(Some(TensorMemory::Mem))
.with_access(CpuAccess::Read)
.with_compression(Compression::Any);
assert_eq!(desc.width(), 640);
assert_eq!(desc.height(), 480);
assert_eq!(desc.format(), PixelFormat::Rgba);
assert_eq!(desc.dtype(), DType::U8);
assert_eq!(desc.memory(), Some(TensorMemory::Mem));
assert_eq!(desc.access(), CpuAccess::Read);
assert_eq!(desc.compression(), Some(Compression::Any));
let plain = ImageDesc::new(2, 2, PixelFormat::Grey, DType::U8);
assert_eq!(plain.memory(), None);
assert_eq!(plain.access(), CpuAccess::None);
assert_eq!(plain.compression(), None);
}
#[test]
fn desc_dtype_must_match_element_type() {
let desc = ImageDesc::new(4, 4, PixelFormat::Rgba, DType::F32);
match Tensor::<u8>::image_desc(&desc) {
Err(Error::InvalidArgument(msg)) => assert!(msg.contains("dtype")),
other => panic!("expected InvalidArgument, got {other:?}"),
}
}
#[test]
fn compression_with_cpu_access_is_invalid() {
let desc = ImageDesc::new(4, 4, PixelFormat::Rgba, DType::U8)
.with_access(CpuAccess::ReadWrite)
.with_compression(Compression::Any);
match Tensor::<u8>::image_desc(&desc) {
Err(Error::InvalidArgument(msg)) => assert!(msg.contains("CpuAccess::None")),
other => panic!("expected InvalidArgument, got {other:?}"),
}
}
#[cfg(not(target_os = "android"))]
#[test]
fn scheme_request_off_android_is_not_implemented() {
let desc = ImageDesc::new(4, 4, PixelFormat::Rgba, DType::U8)
.with_compression(Compression::Scheme(CompressionScheme::Ubwc));
match Tensor::<u8>::image_desc(&desc) {
Err(Error::NotImplemented(msg)) => assert!(msg.contains("Ubwc")),
other => panic!("expected NotImplemented, got {other:?}"),
}
}
#[cfg(not(target_os = "android"))]
#[test]
fn any_request_off_android_resolves_linear_and_counts() {
let before = compression_fallback_count();
let desc = ImageDesc::new(64, 64, PixelFormat::Rgba, DType::U8)
.with_memory(Some(TensorMemory::Mem))
.with_compression(Compression::Any);
let t = Tensor::<u8>::image_desc(&desc).unwrap();
assert_eq!(t.compression(), None);
assert!(compression_fallback_count() > before);
}
#[test]
fn desc_without_request_matches_classic_constructor() {
let desc = ImageDesc::new(32, 32, PixelFormat::Rgba, DType::U8)
.with_memory(Some(TensorMemory::Mem))
.with_access(CpuAccess::ReadWrite);
let t = Tensor::<u8>::image_desc(&desc).unwrap();
assert_eq!(t.compression(), None);
assert_eq!(t.cpu_access(), CpuAccess::ReadWrite);
assert_eq!(t.width(), Some(32));
let m = t.map_read().unwrap();
assert_eq!(m.as_slice().len(), 32 * 32 * 4);
}
#[test]
fn configure_image_preserves_compression_and_views_inherit() {
let mut t = Tensor::<u8>::image(
64,
64,
PixelFormat::Rgba,
Some(TensorMemory::Mem),
CpuAccess::None,
)
.unwrap();
t.compression = Some(CompressionScheme::Ubwc);
t.configure_image(32, 32, PixelFormat::Rgba).unwrap();
assert_eq!(t.compression(), Some(CompressionScheme::Ubwc));
let view = t.subview(0, &[16, 32, 4]).unwrap();
assert_eq!(view.compression(), Some(CompressionScheme::Ubwc));
}
#[test]
fn tensor_dyn_dispatches_desc_and_compression() {
let desc = ImageDesc::new(16, 16, PixelFormat::Rgba, DType::U8)
.with_memory(Some(TensorMemory::Mem))
.with_access(CpuAccess::ReadWrite);
let t = TensorDyn::image_desc(&desc).unwrap();
assert_eq!(t.compression(), None);
assert!(matches!(t, TensorDyn::U8(_)));
}
}
#[cfg(all(test, feature = "static"))]
mod vocabulary_and_probe_tests {
use super::*;
#[test]
fn dtype_displays_its_wire_name() {
for (dt, name) in [
(DType::U8, "u8"),
(DType::I8, "i8"),
(DType::U16, "u16"),
(DType::I16, "i16"),
(DType::U32, "u32"),
(DType::I32, "i32"),
(DType::U64, "u64"),
(DType::I64, "i64"),
(DType::F16, "f16"),
(DType::F32, "f32"),
(DType::F64, "f64"),
] {
assert_eq!(dt.to_string(), name);
}
}
#[test]
fn dtype_of_maps_every_element_type_and_nothing_else() {
assert_eq!(dtype_of::<u8>(), Some(DType::U8));
assert_eq!(dtype_of::<i8>(), Some(DType::I8));
assert_eq!(dtype_of::<u16>(), Some(DType::U16));
assert_eq!(dtype_of::<i16>(), Some(DType::I16));
assert_eq!(dtype_of::<u32>(), Some(DType::U32));
assert_eq!(dtype_of::<i32>(), Some(DType::I32));
assert_eq!(dtype_of::<u64>(), Some(DType::U64));
assert_eq!(dtype_of::<i64>(), Some(DType::I64));
assert_eq!(dtype_of::<half::f16>(), Some(DType::F16));
assert_eq!(dtype_of::<f32>(), Some(DType::F32));
assert_eq!(dtype_of::<f64>(), Some(DType::F64));
assert_eq!(dtype_of::<bool>(), None);
}
#[test]
fn tensor_memory_converts_to_and_from_its_wire_string() {
assert_eq!(String::from(TensorMemory::DmaBuf), "dmabuf");
assert_eq!(String::from(TensorMemory::Shm), "shm");
assert_eq!(
TensorMemory::try_from("dmabuf").unwrap(),
TensorMemory::DmaBuf
);
assert_eq!(
TensorMemory::try_from("dma").unwrap(),
TensorMemory::DmaBuf,
"the legacy spelling still parses"
);
assert_eq!(TensorMemory::try_from("mem").unwrap(), TensorMemory::Mem);
match TensorMemory::try_from("bogus") {
Err(Error::InvalidMemoryType(s)) => assert_eq!(s, "bogus"),
other => panic!("expected InvalidMemoryType, got {other:?}"),
}
}
#[test]
fn tensor_memory_is_available_matches_a_real_allocation() {
let _lock = crate::tests::fd_lock_shared();
assert!(TensorMemory::Mem.is_available());
assert!(!TensorMemory::IoSurface.is_available());
assert!(!TensorMemory::Pbo.is_available());
assert!(!TensorMemory::Cuda.is_available());
#[cfg(unix)]
assert_eq!(
TensorMemory::Shm.is_available(),
Tensor::<u8>::new(&[64], Some(TensorMemory::Shm), None).is_ok()
);
#[cfg(not(target_os = "windows"))]
assert_eq!(
TensorMemory::DmaBuf.is_available(),
Tensor::<u8>::new(&[64], Some(TensorMemory::DmaBuf), None).is_ok()
);
}
#[test]
#[cfg(target_os = "linux")]
fn linux_probes_report_what_linux_can_allocate() {
let _lock = crate::tests::fd_lock_shared();
let dma = DmaTensor::<u8>::new(&[64], None).is_ok();
assert_eq!(is_dma_available(), dma);
assert_eq!(is_gpu_buffer_available(), dma);
assert!(!is_iosurface_available());
assert!(!is_ahardwarebuffer_available());
assert!(is_shm_available());
}
}
#[cfg(all(test, feature = "static"))]
mod quantization_serde_tests {
use super::*;
fn parse(json: &str) -> Quantization {
serde_json::from_str(json).unwrap_or_else(|e| panic!("{json}: {e}"))
}
#[test]
fn scale_accepts_every_json_number_form() {
assert_eq!(parse(r#"{"scale": 0.5}"#).scale(), &[0.5]);
assert_eq!(parse(r#"{"scale": 2}"#).scale(), &[2.0]);
assert_eq!(parse(r#"{"scale": -2}"#).scale(), &[-2.0]);
assert_eq!(parse(r#"{"scale": [0.25, 4]}"#).scale(), &[0.25, 4.0]);
}
#[test]
fn zero_point_accepts_null_scalar_and_array() {
assert_eq!(
parse(r#"{"scale": 0.1, "zero_point": -5}"#).zero_point(),
Some(&[-5][..])
);
assert_eq!(
parse(r#"{"scale": 0.1, "zero_point": 3}"#).zero_point(),
Some(&[3][..])
);
assert_eq!(
parse(r#"{"scale": [0.1, 0.2], "zero_point": [1, -2], "axis": 0}"#).zero_point(),
Some(&[1, -2][..])
);
assert_eq!(
parse(r#"{"scale": 0.1, "zero_point": null}"#).zero_point(),
None
);
assert_eq!(parse(r#"{"scale": 0.1}"#).zero_point(), None);
}
#[test]
fn zero_point_visitor_accepts_values_without_an_option_wrapper() {
use serde::de::value::{
Error as ValueError, I64Deserializer, SeqDeserializer, U64Deserializer,
};
assert_eq!(
deserialize_opt_scalar_or_vec_i32(I64Deserializer::<ValueError>::new(-7)).unwrap(),
Some(vec![-7])
);
assert_eq!(
deserialize_opt_scalar_or_vec_i32(U64Deserializer::<ValueError>::new(9)).unwrap(),
Some(vec![9])
);
let seq = SeqDeserializer::<_, ValueError>::new(vec![4i32, -4].into_iter());
assert_eq!(
deserialize_opt_scalar_or_vec_i32(seq).unwrap(),
Some(vec![4, -4])
);
}
#[test]
fn per_channel_is_not_per_tensor() {
let q = Quantization::per_channel(vec![0.1, 0.2], vec![0, 1], 0).unwrap();
assert!(!q.is_per_tensor());
assert!(Quantization::per_tensor(0.1, 0).is_per_tensor());
}
#[test]
fn validate_accepts_matching_per_channel_and_names_the_failing_field() {
let mut t = Tensor::<i8>::new(&[1, 1, 4], Some(TensorMemory::Mem), None).unwrap();
let q = Quantization::per_channel(vec![0.1, 0.2, 0.3, 0.4], vec![0, 1, 2, 3], 2).unwrap();
t.set_quantization(q)
.expect("four channels on axis 2 of [1, 1, 4]");
assert_eq!(
t.quantization().and_then(Quantization::zero_point),
Some(&[0, 1, 2, 3][..])
);
let q = Quantization::per_channel(vec![0.1, 0.2, 0.3, 0.4], vec![0, 1, 2, 3], 3).unwrap();
match t.set_quantization(q) {
Err(Error::QuantizationInvalid { field, got, .. }) => {
assert_eq!(field, "axis");
assert_eq!(got, "axis=3");
}
other => panic!("axis 3 of a rank-3 tensor must be refused, got {other:?}"),
}
let q = Quantization::per_channel_symmetric(vec![0.1, 0.2, 0.3, 0.4], 1).unwrap();
match t.set_quantization(q) {
Err(Error::QuantizationInvalid { field, got, .. }) => {
assert_eq!(field, "scale.len");
assert_eq!(got, "length 4");
}
other => panic!("four scales on a length-1 axis must be refused, got {other:?}"),
}
}
}
#[cfg(all(test, feature = "static"))]
mod image_geometry_tests {
use super::*;
fn mem_image<T>(width: usize, height: usize, format: PixelFormat) -> Tensor<T>
where
T: Num + Clone + fmt::Debug + Send + Sync + 'static,
{
Tensor::<T>::image(
width,
height,
format,
Some(TensorMemory::Mem),
CpuAccess::ReadWrite,
)
.unwrap_or_else(|e| panic!("Mem {format:?} {width}x{height}: {e}"))
}
fn raw_mem<T>(shape: &[usize]) -> Tensor<T>
where
T: Num + Clone + fmt::Debug + Send + Sync + 'static,
{
Tensor::<T>::new(shape, Some(TensorMemory::Mem), None).unwrap()
}
#[test]
fn semi_planar_image_stride_scales_with_the_element_size() {
let t = mem_image::<u16>(100, 64, PixelFormat::Nv12);
assert_eq!(t.row_stride(), Some(256));
assert_eq!(t.capacity_bytes(), 256 * 96);
}
#[test]
fn host_planar_image_is_tight_at_every_element_size() {
let t = mem_image::<u8>(320, 240, PixelFormat::PlanarRgb);
assert_eq!(t.row_stride(), None);
assert_eq!(t.effective_row_stride(), Some(320));
assert_eq!(t.capacity_bytes(), 320 * 240 * 3);
let t = mem_image::<f32>(100, 40, PixelFormat::PlanarRgb);
assert_eq!(t.row_stride(), None);
assert_eq!(t.effective_row_stride(), Some(400));
assert_eq!(t.capacity_bytes(), 400 * 120);
}
#[test]
#[cfg(target_os = "linux")]
fn dma_packed_image_records_the_aligned_pitch_only_when_padded() {
let _lock = crate::tests::fd_lock_shared();
if !crate::test_support::dma_or_skip(
"dma_packed_image_records_the_aligned_pitch_only_when_padded",
) {
return;
}
let t = Tensor::<u8>::image(
100,
10,
PixelFormat::Rgba,
Some(TensorMemory::DmaBuf),
CpuAccess::ReadWrite,
)
.unwrap();
assert_eq!(t.row_stride(), Some(448), "400 B rows pad to 448");
let t = Tensor::<u8>::image(
64,
10,
PixelFormat::Rgba,
Some(TensorMemory::DmaBuf),
CpuAccess::ReadWrite,
)
.unwrap();
assert_eq!(t.row_stride(), None, "256 B rows are already aligned");
}
#[test]
#[cfg(target_os = "linux")]
fn image_with_stride_validates_the_stride_before_the_memory() {
match Tensor::<u8>::image_with_stride(
640,
480,
PixelFormat::Rgba,
2559,
Some(TensorMemory::Mem),
CpuAccess::ReadWrite,
) {
Err(Error::InvalidArgument(msg)) => {
assert!(msg.contains("2559 < minimum 2560"), "{msg}")
}
other => panic!("expected InvalidArgument, got {other:?}"),
}
match Tensor::<u8>::image_with_stride(
640,
480,
PixelFormat::Rgba,
2560,
Some(TensorMemory::Mem),
CpuAccess::ReadWrite,
) {
Err(Error::NotImplemented(msg)) => assert!(msg.contains("only TensorMemory::DmaBuf")),
other => panic!("a valid stride on Mem must reach the memory check, got {other:?}"),
}
}
#[test]
fn set_format_checks_the_planar_and_semi_planar_shape() {
let mut t = raw_mem::<u8>(&[480, 640, 3]);
match t.set_format(PixelFormat::PlanarRgb) {
Err(Error::InvalidShape(msg)) => assert!(msg.contains("[3, H, W]"), "{msg}"),
other => panic!("[480, 640, 3] is not planar, got {other:?}"),
}
assert!(raw_mem::<u8>(&[3, 480, 640])
.set_format(PixelFormat::PlanarRgb)
.is_ok());
for (format, bad, good) in [
(PixelFormat::Nv12, 721, [720usize, 722]),
(PixelFormat::Nv16, 481, [480, 482]),
(PixelFormat::Nv24, 1441, [1440, 1443]),
] {
match raw_mem::<u8>(&[bad, 640]).set_format(format) {
Err(Error::InvalidShape(msg)) => {
assert!(msg.contains(&bad.to_string()), "{format:?}: {msg}")
}
other => panic!("{format:?} [{bad}, 640] must be refused, got {other:?}"),
}
for rows in good {
raw_mem::<u8>(&[rows, 640])
.set_format(format)
.unwrap_or_else(|e| panic!("{format:?} [{rows}, 640]: {e}"));
}
}
}
#[test]
fn contiguous_semi_planar_height_divides_the_combined_rows() {
assert_eq!(
mem_image::<u8>(320, 240, PixelFormat::Nv24).height(),
Some(240)
);
assert_eq!(
mem_image::<u8>(320, 240, PixelFormat::Nv16).height(),
Some(240)
);
assert_eq!(
mem_image::<u8>(320, 241, PixelFormat::Nv12).height(),
Some(241)
);
}
#[test]
fn effective_row_stride_scales_with_the_element_size() {
assert_eq!(
mem_image::<f32>(10, 4, PixelFormat::Rgb).effective_row_stride(),
Some(120)
);
assert_eq!(
mem_image::<f32>(10, 4, PixelFormat::PlanarRgb).effective_row_stride(),
Some(40)
);
let mut t = raw_mem::<u16>(&[360, 320]);
t.set_format(PixelFormat::Nv12).unwrap();
assert_eq!(t.effective_row_stride(), Some(640));
let mut t = raw_mem::<u16>(&[360, 321]);
t.set_format(PixelFormat::Nv12).unwrap();
assert_eq!(
t.effective_row_stride(),
Some(644),
"odd width rounds up to even"
);
}
#[test]
fn set_row_stride_minimum_scales_with_the_element_size() {
let refuse = |mut t: Tensor<f32>, stride: usize, min: usize| match t.set_row_stride(stride)
{
Err(Error::InvalidArgument(msg)) => {
assert!(msg.contains(&format!("{stride} < minimum {min}")), "{msg}")
}
other => panic!("stride {stride} under {min} must be refused, got {other:?}"),
};
refuse(mem_image::<f32>(10, 4, PixelFormat::Rgb), 119, 120);
refuse(mem_image::<f32>(10, 4, PixelFormat::PlanarRgb), 39, 40);
let mut t = mem_image::<f32>(10, 4, PixelFormat::Rgb);
t.set_row_stride(120).unwrap();
assert_eq!(t.row_stride(), Some(120));
let mut t = mem_image::<f32>(10, 4, PixelFormat::PlanarRgb);
t.set_row_stride(40).unwrap();
assert_eq!(t.row_stride(), Some(40));
let mut t = mem_image::<u16>(101, 64, PixelFormat::Nv12);
match t.set_row_stride(203) {
Err(Error::InvalidArgument(msg)) => assert!(msg.contains("203 < minimum 204"), "{msg}"),
other => panic!("expected InvalidArgument, got {other:?}"),
}
t.set_row_stride(204).unwrap();
}
#[test]
fn configure_image_semi_planar_stride_scales_with_the_element_size() {
let mut t = mem_image::<u16>(640, 480, PixelFormat::Rgba);
t.configure_image(100, 64, PixelFormat::Nv12).unwrap();
assert_eq!(t.row_stride(), Some(256));
}
#[test]
fn configure_image_keeps_host_packed_odd_widths_tight() {
let mut t = mem_image::<u8>(640, 480, PixelFormat::Rgb);
t.configure_image(321, 240, PixelFormat::Rgb).unwrap();
assert_eq!(t.row_stride(), None);
assert_eq!(t.effective_row_stride(), Some(963));
}
#[test]
fn configure_image_does_not_inherit_a_stride_under_the_same_format() {
let mut t = mem_image::<u8>(640, 480, PixelFormat::Rgb);
t.configure_image(100, 100, PixelFormat::Rgb).unwrap();
t.set_row_stride(400).unwrap();
t.configure_image(300, 100, PixelFormat::Rgb).unwrap();
assert_eq!(t.row_stride(), None);
assert_eq!(t.effective_row_stride(), Some(900));
let mut t = mem_image::<u8>(64, 1000, PixelFormat::Nv12);
assert_eq!(t.row_stride(), Some(64));
t.configure_image(200, 300, PixelFormat::Nv12).unwrap();
assert_eq!(t.row_stride(), None);
assert_eq!(t.effective_row_stride(), Some(200));
}
#[test]
fn configure_image_pool_reuse_picks_the_stride_that_fits() {
let mut t = mem_image::<u8>(64, 1024, PixelFormat::Nv12);
t.configure_image(200, 100, PixelFormat::Nv12).unwrap();
assert_eq!(t.row_stride(), Some(256));
let mut t = mem_image::<u8>(64, 64, PixelFormat::Nv12);
t.set_row_stride(100).unwrap();
t.configure_image(64, 64, PixelFormat::Nv12).unwrap();
assert_eq!(t.row_stride(), Some(64));
let mut t = mem_image::<u8>(640, 480, PixelFormat::Nv12);
t.configure_image(320, 900, PixelFormat::Nv12).unwrap();
assert_eq!(t.row_stride(), Some(320));
let mut t = raw_mem::<u8>(&[150 * 100]);
t.configure_image(100, 100, PixelFormat::Nv12).unwrap();
assert_eq!(t.row_stride(), None);
assert_eq!(t.effective_row_stride(), Some(100));
}
#[test]
fn configure_image_refuses_rows_that_overrun_the_allocation() {
let mut t = raw_mem::<u8>(&[150 * 99]);
match t.configure_image(99, 100, PixelFormat::Nv12) {
Err(Error::InsufficientCapacity { needed, capacity }) => {
assert_eq!((needed, capacity), (15000, 14850))
}
other => panic!("expected InsufficientCapacity, got {other:?}"),
}
}
#[test]
#[cfg(target_os = "linux")]
fn configure_image_on_dma_aligns_packed_and_planar_rows() {
let _lock = crate::tests::fd_lock_shared();
if !crate::test_support::dma_or_skip("configure_image_on_dma_aligns_packed_and_planar_rows")
{
return;
}
let mut t = Tensor::<u16>::image(
640,
480,
PixelFormat::Rgba,
Some(TensorMemory::DmaBuf),
CpuAccess::ReadWrite,
)
.unwrap();
t.configure_image(100, 50, PixelFormat::Rgb).unwrap();
assert_eq!(t.row_stride(), Some(640));
let mut t = Tensor::<u16>::new(&[16 * 1024], Some(TensorMemory::DmaBuf), None).unwrap();
assert_eq!(t.capacity_bytes(), 32 * 1024);
t.configure_image(100, 50, PixelFormat::PlanarRgb).unwrap();
assert_eq!(t.row_stride(), None);
assert_eq!(t.effective_row_stride(), Some(200));
t.configure_image(100, 40, PixelFormat::PlanarRgb).unwrap();
assert_eq!(t.row_stride(), Some(256), "120 rows of 256 B fit");
}
#[test]
fn from_planes_requires_a_raw_two_dimensional_chroma() {
let formatted_chroma = mem_image::<u8>(640, 240, PixelFormat::Grey);
match Tensor::from_planes(
raw_mem::<u8>(&[480, 640]),
formatted_chroma,
PixelFormat::Nv12,
) {
Err(Error::InvalidArgument(msg)) => assert!(msg.contains("raw tensor"), "{msg}"),
other => panic!("a formatted chroma plane must be refused, got {other:?}"),
}
match Tensor::from_planes(
raw_mem::<u8>(&[480, 640, 1]),
raw_mem::<u8>(&[240, 640]),
PixelFormat::Nv12,
) {
Err(Error::InvalidArgument(msg)) => assert!(msg.contains("2D shapes"), "{msg}"),
other => panic!("a 3-D luma plane must be refused, got {other:?}"),
}
}
#[test]
fn from_planes_nv16_needs_full_height_chroma() {
let img = Tensor::from_planes(
raw_mem::<u8>(&[480, 640]),
raw_mem::<u8>(&[480, 640]),
PixelFormat::Nv16,
)
.expect("NV16 chroma is full height");
assert_eq!(img.format(), Some(PixelFormat::Nv16));
assert_eq!(img.height(), Some(480));
match Tensor::from_planes(
raw_mem::<u8>(&[480, 640]),
raw_mem::<u8>(&[240, 640]),
PixelFormat::Nv16,
) {
Err(Error::InvalidArgument(msg)) => {
assert!(msg.contains("NV16 chroma height"), "{msg}")
}
other => panic!("half-height NV16 chroma must be refused, got {other:?}"),
}
}
#[test]
fn chroma_mut_writes_the_chroma_plane() {
let mut img = Tensor::from_planes(
raw_mem::<u8>(&[4, 8]),
raw_mem::<u8>(&[2, 8]),
PixelFormat::Nv12,
)
.unwrap();
let chroma = img.chroma_mut().expect("a multiplane image has chroma");
chroma.map().unwrap().as_mut_slice().fill(0x80);
let chroma = img.chroma().unwrap();
assert!(chroma.map().unwrap().as_slice().iter().all(|&b| b == 0x80));
}
#[test]
fn view_offset_and_stride_scale_with_the_element_size() {
let t = mem_image::<u8>(100, 80, PixelFormat::Rgba);
let v = t.view(Region::new(10, 20, 30, 40)).unwrap();
assert_eq!(v.plane_offset(), Some(20 * 400 + 10 * 4));
assert_eq!(v.row_stride(), Some(400));
let t = mem_image::<u16>(100, 80, PixelFormat::Rgba);
let v = t.view(Region::new(10, 20, 30, 40)).unwrap();
assert_eq!(v.plane_offset(), Some(20 * 800 + 10 * 8));
assert_eq!(v.row_stride(), Some(800));
assert_eq!(v.view_origin().map(|o| o.parent_row_stride), Some(800));
}
#[test]
fn single_row_view_keeps_a_tight_stride() {
let t = mem_image::<u8>(8, 4, PixelFormat::Rgba);
let row = t.view(Region::new(2, 3, 4, 1)).unwrap();
assert_eq!(row.effective_row_stride(), Some(16));
assert_eq!(row.map().unwrap().as_slice().len(), 16);
assert_eq!(row.view_origin().map(|o| o.parent_row_stride), Some(32));
let two_rows = t.view(Region::new(2, 2, 4, 2)).unwrap();
assert_eq!(two_rows.effective_row_stride(), Some(32));
}
#[test]
fn view_refuses_a_region_past_either_edge() {
let t = mem_image::<u8>(64, 40, PixelFormat::Rgba);
for region in [Region::new(0, 30, 10, 20), Region::new(60, 0, 10, 1)] {
match t.view(region) {
Err(Error::RegionOutOfBounds { region: r, bounds }) => {
assert_eq!((r, bounds), (region, (64, 40)))
}
other => panic!("{region:?} must be refused, got {other:?}"),
}
}
assert!(Region::new(0, 0, 64, 40).fits_within(64, 40));
assert!(!Region::new(1, 0, 64, 40).fits_within(64, 40));
assert!(!Region::new(0, 1, 64, 40).fits_within(64, 40));
}
}
#[cfg(all(test, feature = "static"))]
mod strided_access_tests {
use super::*;
fn fill_pattern<T>(t: &Tensor<T>)
where
T: Num + Clone + fmt::Debug + Send + Sync,
{
let mut m = t.map().unwrap();
let s = m.as_mut_slice();
let bytes = unsafe {
std::slice::from_raw_parts_mut(s.as_mut_ptr() as *mut u8, std::mem::size_of_val(s))
};
for (i, b) in bytes.iter_mut().enumerate() {
*b = (i % 251) as u8;
}
}
fn expected_rows(rows: usize, stride: usize, row_bytes: usize) -> Vec<u8> {
(0..rows)
.flat_map(|r| (r * stride..r * stride + row_bytes).map(|i| (i % 251) as u8))
.collect()
}
fn padded<T>(format: PixelFormat, stride: usize) -> Tensor<T>
where
T: Num + Clone + fmt::Debug + Send + Sync + 'static,
{
let mut t = Tensor::<T>::image_with_capacity(
16,
4,
format,
Some(TensorMemory::Mem),
CpuAccess::ReadWrite,
)
.unwrap();
t.configure_image(8, 4, format).unwrap();
t.set_row_stride(stride).unwrap();
t
}
#[test]
fn copy_to_flat_compacts_padded_packed_rows() {
let t = padded::<u8>(PixelFormat::Rgb, 32);
fill_pattern(&t);
let mut dst = vec![0u8; 96];
t.copy_to_flat(&mut dst).unwrap();
assert_eq!(dst, expected_rows(4, 32, 24));
}
#[test]
fn copy_to_flat_compacts_padded_planar_rows() {
let t = padded::<f32>(PixelFormat::PlanarRgb, 64);
fill_pattern(&t);
let mut dst = vec![0u8; 384];
t.copy_to_flat(&mut dst).unwrap();
assert_eq!(dst, expected_rows(12, 64, 32));
}
#[test]
fn copy_to_flat_of_a_tight_stride_is_the_mapping() {
let t = padded::<u8>(PixelFormat::Rgb, 24);
fill_pattern(&t);
let mut dst = vec![0u8; 96];
t.copy_to_flat(&mut dst).unwrap();
assert_eq!(dst, expected_rows(1, 96, 96));
let t = Tensor::<u8>::image(
8,
4,
PixelFormat::Rgb,
Some(TensorMemory::Mem),
CpuAccess::ReadWrite,
)
.unwrap();
fill_pattern(&t);
let mut dst = vec![0u8; 96];
t.copy_to_flat(&mut dst).unwrap();
assert_eq!(dst, expected_rows(1, 96, 96));
}
#[test]
fn copy_to_flat_needs_a_tight_sized_destination() {
let t = padded::<u8>(PixelFormat::Rgb, 32);
let mut dst = vec![0u8; 95];
match t.copy_to_flat(&mut dst) {
Err(Error::InvalidArgument(msg)) => {
assert!(
msg.contains("dst is 95 bytes") && msg.contains("96 bytes"),
"{msg}"
)
}
other => panic!("expected InvalidArgument, got {other:?}"),
}
}
#[test]
fn copy_to_flat_compacts_a_padded_batch() {
let mut t = padded::<u8>(PixelFormat::Rgb, 32);
t.set_logical_shape(&[2, 2, 8, 3]).unwrap();
fill_pattern(&t);
let mut dst = vec![0u8; 96];
t.copy_to_flat(&mut dst).unwrap();
assert_eq!(dst, expected_rows(4, 32, 24));
}
#[test]
fn strided_planar_map_spans_every_plane_row() {
let t = padded::<u8>(PixelFormat::PlanarRgb, 16);
assert_eq!(t.map().unwrap().as_slice().len(), 16 * 12);
}
#[test]
fn batch_of_a_padded_batch_starts_at_its_padded_tile() {
let mut t = padded::<u8>(PixelFormat::Rgb, 32);
t.set_logical_shape(&[2, 2, 8, 3]).unwrap();
fill_pattern(&t);
let e = t.batch(1).unwrap();
assert_eq!(e.plane_offset(), Some(64));
let mut dst = vec![0u8; 48];
e.copy_to_flat(&mut dst).unwrap();
let tile1: Vec<u8> = expected_rows(4, 32, 24).split_off(48);
assert_eq!(dst, tile1);
}
#[test]
fn strided_packed_batch_map_spans_every_tile_row() {
let mut t = padded::<u8>(PixelFormat::Rgb, 32);
t.set_logical_shape(&[2, 2, 8, 3]).unwrap();
assert_eq!(t.map().unwrap().as_slice().len(), 32 * 4);
}
#[test]
fn strided_access_to_a_flattened_image_is_refused() {
let mut t = padded::<u8>(PixelFormat::PlanarRgb, 16);
t.set_logical_shape(&[96]).unwrap();
match t.map() {
Err(Error::InvalidOperation(msg)) => assert!(msg.contains("[96]"), "{msg}"),
other => panic!("expected InvalidOperation, got {other:?}"),
}
let mut dst = vec![0u8; 96];
match t.copy_to_flat(&mut dst) {
Err(Error::InvalidOperation(msg)) => assert!(msg.contains("[96]"), "{msg}"),
other => panic!("expected InvalidOperation, got {other:?}"),
}
}
#[test]
#[cfg(target_os = "linux")]
fn strided_shm_map_is_bounded_by_the_segment() {
let _lock = crate::tests::fd_lock_shared();
let mut t = Tensor::<u8>::image(
640,
480,
PixelFormat::Nv12,
Some(TensorMemory::Shm),
CpuAccess::ReadWrite,
)
.expect("Linux has POSIX shared memory");
assert_eq!(t.capacity_bytes(), 640 * 720);
assert_eq!(t.map().unwrap().as_slice().len(), 640 * 720);
t.configure_image(320, 240, PixelFormat::Nv12).unwrap();
assert_eq!(t.row_stride(), Some(640), "the pool pitch fits and is kept");
assert_eq!(t.map().unwrap().as_slice().len(), 640 * 360);
t.configure_image(640, 480, PixelFormat::Nv12).unwrap();
t.set_row_stride(704).unwrap();
match t.map() {
Err(Error::InsufficientCapacity { needed, capacity }) => {
assert_eq!((needed, capacity), (704 * 720, 640 * 720))
}
other => panic!("expected InsufficientCapacity, got {other:?}"),
}
}
struct HostPbo(Vec<u8>);
unsafe impl PboOps for HostPbo {
fn map_buffer(&self, _buffer_id: u32, size: usize) -> Result<PboMapping> {
assert!(size <= self.0.len());
Ok(PboMapping {
ptr: self.0.as_ptr() as *mut u8,
size,
})
}
fn unmap_buffer(&self, _buffer_id: u32) -> Result<()> {
Ok(())
}
fn delete_buffer(&self, _buffer_id: u32) {}
}
fn pbo_rgba(width: usize, height: usize) -> Tensor<u8> {
let bytes = width * height * 4;
let pbo = PboTensor::<u8>::from_pbo(
7,
bytes,
&[height, width, 4],
None,
Arc::new(HostPbo(vec![0; bytes])),
)
.unwrap();
let mut t = Tensor::from_pbo(pbo).unwrap();
t.set_format(PixelFormat::Rgba).unwrap();
t
}
#[test]
fn strided_pbo_map_is_bounded_by_the_buffer() {
let mut t = pbo_rgba(8, 4);
t.set_row_stride(32).unwrap();
assert_eq!(t.map().unwrap().as_slice().len(), 128);
t.set_row_stride(64).unwrap();
match t.map() {
Err(Error::InsufficientCapacity { needed, capacity }) => {
assert_eq!((needed, capacity), (256, 128))
}
other => panic!("expected InsufficientCapacity, got {other:?}"),
}
}
#[test]
fn pbo_accessors_answer_only_for_a_pbo_tensor() {
let mem = Tensor::<u8>::new(&[4], Some(TensorMemory::Mem), None).unwrap();
assert_eq!(mem.pbo_is_mapped(), None);
assert!(mem.pbo_keepalive().is_none());
let t = pbo_rgba(8, 4);
assert_eq!(t.pbo_is_mapped(), Some(false));
assert!(t.pbo_keepalive().is_some());
let m = t.map().unwrap();
assert_eq!(t.pbo_is_mapped(), Some(true));
drop(m);
assert_eq!(t.pbo_is_mapped(), Some(false));
}
}
#[cfg(all(test, feature = "static"))]
mod accessor_tests {
use super::*;
use std::sync::atomic::{AtomicUsize, Ordering};
#[test]
#[cfg(unix)]
fn plane_descriptor_carries_its_stride_and_offset() {
use std::os::fd::AsFd;
let _lock = crate::tests::fd_lock_shared();
let file = std::fs::File::open("/dev/null").unwrap();
let pd = PlaneDescriptor::new(file.as_fd()).unwrap();
assert_eq!((pd.stride(), pd.offset()), (None, None));
let pd = pd.with_stride(832).with_offset(4096);
assert_eq!((pd.stride(), pd.offset()), (Some(832), Some(4096)));
}
#[test]
#[cfg(target_os = "linux")]
fn as_dma_answers_only_for_a_dma_buffer() {
use std::os::fd::AsRawFd;
let _lock = crate::tests::fd_lock_shared();
let mem = Tensor::<u8>::new(&[64], Some(TensorMemory::Mem), None).unwrap();
assert!(mem.as_dma().is_none());
if !crate::test_support::dma_or_skip("as_dma_answers_only_for_a_dma_buffer") {
return;
}
let t = Tensor::<u8>::new(&[64], Some(TensorMemory::DmaBuf), None).unwrap();
let dma = t.as_dma().expect("a DMA tensor downcasts");
assert_eq!(
dma.fd.as_raw_fd(),
t.dmabuf().unwrap().as_raw_fd(),
"the downcast is this tensor's own buffer"
);
}
#[test]
fn derived_identity_mixes_the_kind_into_the_key() {
let id = BufferIdentity::derived(IdentityKind::Shm, u64::MAX);
assert_eq!(id.id(), ((IdentityKind::Shm as u64) << 56) ^ u64::MAX);
assert_ne!(
id.id(),
BufferIdentity::derived(IdentityKind::DmaBuf, u64::MAX).id()
);
}
#[test]
fn empty_is_a_zero_element_shape() {
let mut t = Tensor::<f32>::new(&[2, 3], Some(TensorMemory::Mem), None).unwrap();
assert!(!TensorTrait::is_empty(&t));
{
let m = t.map().unwrap();
assert_eq!(TensorMapTrait::len(&m), 6);
assert_eq!(TensorMapTrait::size(&m), 24);
assert!(!TensorMapTrait::is_empty(&m));
}
t.set_logical_shape(&[0]).unwrap();
assert!(TensorTrait::is_empty(&t));
let m = t.map().unwrap();
assert_eq!(TensorMapTrait::size(&m), 0);
assert!(TensorMapTrait::is_empty(&m));
}
struct CountingCudaOps {
maps: AtomicUsize,
}
impl CudaGlOps for CountingCudaOps {
fn map(&self, _r: crate::cuda::GraphicsResource) -> Option<(*mut std::ffi::c_void, usize)> {
self.maps.fetch_add(1, Ordering::SeqCst);
Some((0x1000usize as *mut _, 64))
}
fn unmap(&self, _r: crate::cuda::GraphicsResource) {}
fn unregister(&self, _r: crate::cuda::GraphicsResource) {}
}
#[test]
fn cuda_accessors_follow_the_attached_handle() {
let mut t = Tensor::<f32>::new(&[16], Some(TensorMemory::Mem), None).unwrap();
assert!(!t.is_cuda_attached());
assert!(t.cuda().is_none());
assert!(t.cuda_map().is_none());
assert!(t.cuda_map_mut().is_none());
let ops = Arc::new(CountingCudaOps {
maps: AtomicUsize::new(0),
});
t.set_cuda_handle(CudaHandle::new_gl(
0x1usize as crate::cuda::GraphicsResource,
64,
ops.clone(),
));
assert!(t.is_cuda_attached());
assert!(t.cuda().is_some());
assert_eq!(t.cuda_map().map(|m| m.device_ptr() as usize), Some(0x1000));
assert_eq!(t.cuda_map_mut().map(|m| m.len()), Some(64));
assert_eq!(ops.maps.load(Ordering::SeqCst), 2);
}
#[test]
fn from_slice_copies_values_of_a_matching_shape() {
let t = Tensor::<u16>::from_slice(&[1, 2, 3, 4, 5, 6], &[2, 3]).unwrap();
assert_eq!(t.shape(), &[2, 3]);
assert_eq!(t.map().unwrap().as_slice(), &[1, 2, 3, 4, 5, 6]);
match Tensor::<u16>::from_slice(&[1, 2, 3, 4, 5], &[2, 3]) {
Err(Error::InvalidShape(msg)) => assert!(msg.contains("values.len()=5"), "{msg}"),
other => panic!("expected InvalidShape, got {other:?}"),
}
}
#[cfg(not(target_os = "android"))]
#[test]
fn any_compression_request_falls_back_when_zero_copy_cannot_serve() {
let _lock = crate::tests::fd_lock_shared();
let desc =
ImageDesc::new(10, 10, PixelFormat::Rgb, DType::F32).with_compression(Compression::Any);
let t = Tensor::<f32>::image_desc(&desc).expect("Any never fails for compression reasons");
assert_eq!(t.shape(), &[10, 10, 3]);
assert_eq!(t.compression(), None);
}
}
#[cfg(test)]
mod cpu_access_tests {
use super::*;
#[test]
fn covers_matrix() {
use CpuAccess::*;
for a in [None, Read, Write, ReadWrite] {
assert!(a.covers(None), "{a:?} must cover None");
assert!(ReadWrite.covers(a), "ReadWrite must cover {a:?}");
}
assert!(Read.covers(Read));
assert!(Write.covers(Write));
assert!(!None.covers(Read));
assert!(!None.covers(Write));
assert!(!Read.covers(Write));
assert!(!Read.covers(ReadWrite));
assert!(!Write.covers(Read));
assert!(!Write.covers(ReadWrite));
}
#[test]
fn map_with_none_is_invalid() {
let t = Tensor::<u8>::new(&[16], Some(TensorMemory::Mem), None).unwrap();
match t.map_with(CpuAccess::None) {
Err(Error::InvalidArgument(_)) => {}
Err(other) => panic!("expected InvalidArgument, got {other:?}"),
Ok(_) => panic!("map_with(CpuAccess::None) must not succeed"),
}
}
#[test]
fn read_map_rejects_mutation_uniformly() {
let t = Tensor::<u8>::new(&[8], Some(TensorMemory::Mem), None).unwrap();
t.map_mut().unwrap().as_mut_slice().copy_from_slice(&[7; 8]);
let ro = t.map_read().unwrap();
assert_eq!(ro.as_slice(), &[7; 8]);
drop(ro);
let result = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
let mut ro = t.map_read().unwrap();
let _ = ro.as_mut_slice();
}));
assert!(result.is_err(), "as_mut_slice through map_read must panic");
}
#[test]
fn write_and_rw_maps_stay_mutable() {
let t = Tensor::<u8>::new(&[4], Some(TensorMemory::Mem), None).unwrap();
t.map_write()
.unwrap()
.as_mut_slice()
.copy_from_slice(&[1; 4]);
t.map().unwrap().as_mut_slice().copy_from_slice(&[2; 4]);
assert_eq!(t.map_read().unwrap().as_slice(), &[2; 4]);
}
#[test]
#[cfg(target_os = "macos")]
fn iosurface_read_only_lock_roundtrip() {
let _lock = crate::tests::fd_lock_shared();
let Ok(t) = Tensor::<u8>::new(&[64], Some(TensorMemory::DmaBuf), None) else {
crate::test_support::report_skip("IOSurface unavailable");
return;
};
{
let mut m = t.map_mut().unwrap();
for (i, b) in m.as_mut_slice().iter_mut().enumerate() {
*b = (i * 3) as u8;
}
}
for _ in 0..2 {
let ro = t.map_read().unwrap();
for (i, b) in ro.as_slice().iter().enumerate() {
assert_eq!(*b, (i * 3) as u8, "byte {i} through read-only lock");
}
}
let result = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
let mut ro = t.map_read().unwrap();
let _ = ro.as_mut_slice();
}));
assert!(result.is_err(), "IOSurface read map must reject mutation");
}
}
#[cfg(test)]
mod tests {
use std::sync::{PoisonError, RwLock};
use super::*;
#[ctor::ctor(unsafe)]
fn init() {
env_logger::Builder::from_env(env_logger::Env::default().default_filter_or("info")).init();
}
#[cfg(target_os = "linux")]
macro_rules! function {
() => {{
fn f() {}
fn type_name_of<T>(_: T) -> &'static str {
std::any::type_name::<T>()
}
let name = type_name_of(f);
match &name[..name.len() - 3].rfind(':') {
Some(pos) => &name[pos + 1..name.len() - 3],
None => &name[..name.len() - 3],
}
}};
}
#[test]
#[cfg(target_os = "linux")]
fn test_tensor() {
let _lock = crate::tests::fd_lock_shared();
let shape = vec![1];
let tensor = DmaTensor::<f32>::new(&shape, Some("dma_tensor"));
let dma_enabled = tensor.is_ok();
let tensor = Tensor::<f32>::new(&shape, None, None).expect("Failed to create tensor");
match dma_enabled {
true => assert_eq!(tensor.memory(), TensorMemory::DmaBuf),
false => assert_eq!(tensor.memory(), TensorMemory::Mem),
}
}
#[test]
#[cfg(any(target_os = "macos", target_os = "ios"))]
fn test_tensor() {
let shape = vec![1];
let tensor = Tensor::<f32>::new(&shape, None, None).expect("Failed to create tensor");
let m = tensor.memory();
assert!(
matches!(m, TensorMemory::DmaBuf | TensorMemory::Mem),
"Unexpected auto-fallback result on macOS/iOS: {m:?}"
);
}
#[test]
#[cfg(all(
unix,
not(any(target_os = "linux", target_os = "macos", target_os = "ios"))
))]
fn test_tensor() {
let shape = vec![1];
let tensor = Tensor::<f32>::new(&shape, None, None).expect("Failed to create tensor");
assert_eq!(tensor.memory(), TensorMemory::Mem);
}
#[test]
#[cfg(not(unix))]
fn test_tensor() {
let shape = vec![1];
let tensor = Tensor::<f32>::new(&shape, None, None).expect("Failed to create tensor");
assert_eq!(tensor.memory(), TensorMemory::Mem);
}
#[test]
#[cfg(target_os = "linux")]
fn test_dma_tensor() {
let _lock = crate::tests::fd_lock_shared();
if !crate::test_support::dma_or_skip("test_dma_tensor") {
return;
}
let shape = vec![2, 3, 4];
let tensor =
DmaTensor::<f32>::new(&shape, Some("test_tensor")).expect("Failed to create tensor");
const DUMMY_VALUE: f32 = 12.34;
assert_eq!(tensor.memory(), TensorMemory::DmaBuf);
assert_eq!(tensor.name(), "test_tensor");
assert_eq!(tensor.shape(), &shape);
assert_eq!(tensor.size(), 2 * 3 * 4 * std::mem::size_of::<f32>());
assert_eq!(tensor.len(), 2 * 3 * 4);
{
let mut tensor_map = tensor.map().expect("Failed to map DMA memory");
tensor_map.fill(42.0);
assert!(tensor_map.iter().all(|&x| x == 42.0));
}
{
let shared = Tensor::<f32>::from_fd(
tensor
.clone_fd()
.expect("Failed to duplicate tensor file descriptor"),
&shape,
Some("test_tensor_shared"),
)
.expect("Failed to create tensor from fd");
assert_eq!(shared.memory(), TensorMemory::DmaBuf);
assert_eq!(shared.name(), "test_tensor_shared");
assert_eq!(shared.shape(), &shape);
let mut tensor_map = shared.map().expect("Failed to map DMA memory from fd");
tensor_map.fill(DUMMY_VALUE);
assert!(tensor_map.iter().all(|&x| x == DUMMY_VALUE));
}
{
let tensor_map = tensor.map().expect("Failed to map DMA memory");
assert!(tensor_map.iter().all(|&x| x == DUMMY_VALUE));
}
let mut tensor = DmaTensor::<u8>::new(&shape, None).expect("Failed to create tensor");
assert_eq!(tensor.shape(), &shape);
let new_shape = vec![3, 4, 4];
assert!(
tensor.reshape(&new_shape).is_err(),
"Reshape should fail due to size mismatch"
);
assert_eq!(tensor.shape(), &shape, "Shape should remain unchanged");
let new_shape = vec![2, 3, 4];
tensor.reshape(&new_shape).expect("Reshape should succeed");
assert_eq!(
tensor.shape(),
&new_shape,
"Shape should be updated after successful reshape"
);
{
let mut tensor_map = tensor.map().expect("Failed to map DMA memory");
tensor_map.fill(1);
assert!(tensor_map.iter().all(|&x| x == 1));
}
{
let mut tensor_map = tensor.map().expect("Failed to map DMA memory");
tensor_map[2] = 42;
assert_eq!(tensor_map[1], 1, "Value at index 1 should be 1");
assert_eq!(tensor_map[2], 42, "Value at index 2 should be 42");
}
}
#[test]
#[cfg(unix)]
fn test_shm_tensor() {
let _lock = crate::tests::fd_lock_shared();
if !is_shm_available() {
log::warn!("SKIPPED: test_shm_tensor - SHM allocation unavailable on this platform");
return;
}
let shape = vec![2, 3, 4];
let tensor =
ShmTensor::<f32>::new(&shape, Some("test_tensor")).expect("Failed to create tensor");
assert_eq!(tensor.shape(), &shape);
assert_eq!(tensor.size(), 2 * 3 * 4 * std::mem::size_of::<f32>());
assert_eq!(tensor.name(), "test_tensor");
const DUMMY_VALUE: f32 = 12.34;
{
let mut tensor_map = tensor.map().expect("Failed to map shared memory");
tensor_map.fill(42.0);
assert!(tensor_map.iter().all(|&x| x == 42.0));
}
{
let shared = Tensor::<f32>::from_fd(
tensor
.clone_fd()
.expect("Failed to duplicate tensor file descriptor"),
&shape,
Some("test_tensor_shared"),
)
.expect("Failed to create tensor from fd");
assert_eq!(shared.memory(), TensorMemory::Shm);
assert_eq!(shared.name(), "test_tensor_shared");
assert_eq!(shared.shape(), &shape);
let mut tensor_map = shared.map().expect("Failed to map shared memory from fd");
tensor_map.fill(DUMMY_VALUE);
assert!(tensor_map.iter().all(|&x| x == DUMMY_VALUE));
}
{
let tensor_map = tensor.map().expect("Failed to map shared memory");
assert!(tensor_map.iter().all(|&x| x == DUMMY_VALUE));
}
let mut tensor = ShmTensor::<u8>::new(&shape, None).expect("Failed to create tensor");
assert_eq!(tensor.shape(), &shape);
let new_shape = vec![3, 4, 4];
assert!(
tensor.reshape(&new_shape).is_err(),
"Reshape should fail due to size mismatch"
);
assert_eq!(tensor.shape(), &shape, "Shape should remain unchanged");
let new_shape = vec![2, 3, 4];
tensor.reshape(&new_shape).expect("Reshape should succeed");
assert_eq!(
tensor.shape(),
&new_shape,
"Shape should be updated after successful reshape"
);
{
let mut tensor_map = tensor.map().expect("Failed to map shared memory");
tensor_map.fill(1);
assert!(tensor_map.iter().all(|&x| x == 1));
}
{
let mut tensor_map = tensor.map().expect("Failed to map shared memory");
tensor_map[2] = 42;
assert_eq!(tensor_map[1], 1, "Value at index 1 should be 1");
assert_eq!(tensor_map[2], 42, "Value at index 2 should be 42");
}
}
#[test]
fn mem_subview_partitions_parent_buffer() {
let parent = Tensor::<u8>::new(&[2, 4], Some(TensorMemory::Mem), None).unwrap();
let view0 = parent.subview(0, &[1, 4]).expect("subview at offset 0");
let view1 = parent.subview(4, &[1, 4]).expect("subview at offset 4");
view1
.map()
.unwrap()
.as_mut_slice()
.copy_from_slice(&[10, 20, 30, 40]);
view0
.map()
.unwrap()
.as_mut_slice()
.copy_from_slice(&[1, 2, 3, 4]);
assert_eq!(view0.map().unwrap().as_slice(), &[1, 2, 3, 4]);
assert_eq!(view1.map().unwrap().as_slice(), &[10, 20, 30, 40]);
assert_eq!(
parent.map().unwrap().as_slice(),
&[1, 2, 3, 4, 10, 20, 30, 40]
);
}
#[test]
fn batch_partitions_leading_dim() {
let parent = Tensor::<u8>::new(&[4, 2, 2, 3], Some(TensorMemory::Mem), None).unwrap();
for i in 0..4u8 {
let e = parent.batch(i as usize).expect("batch element");
assert_eq!(e.shape(), &[2, 2, 3]);
assert_eq!(e.buffer_identity().id(), parent.buffer_identity().id());
for b in e.map().unwrap().as_mut_slice() {
*b = i + 1;
}
}
let whole = parent.map().unwrap();
let s = whole.as_slice();
for i in 0..4usize {
assert!(
s[i * 12..(i + 1) * 12].iter().all(|&b| b == (i as u8 + 1)),
"band {i} not partitioned: {:?}",
&s[i * 12..(i + 1) * 12]
);
}
}
#[test]
fn view_origin_snapshots_parent_and_composes() {
let parent = Tensor::<u8>::image(
100,
80,
PixelFormat::Rgba,
Some(TensorMemory::Mem),
crate::CpuAccess::ReadWrite,
)
.unwrap();
assert_eq!(
parent.view_origin(),
None,
"whole tensor has no view_origin"
);
let v = parent.view(Region::new(10, 20, 30, 40)).unwrap();
assert_eq!(
v.view_origin(),
Some(ViewOrigin {
parent_width: 100,
parent_height: 80,
parent_row_stride: 100 * 4, x: 10,
y: 20
})
);
let v2 = v.view(Region::new(5, 5, 10, 10)).unwrap();
assert_eq!(
v2.view_origin(),
Some(ViewOrigin {
parent_width: 100,
parent_height: 80,
parent_row_stride: 100 * 4,
x: 15,
y: 25
}),
"nested view composes onto the root parent"
);
}
#[test]
fn view_origin_none_for_raw_batch() {
let parent = Tensor::<u8>::new(&[4, 2, 2, 3], Some(TensorMemory::Mem), None).unwrap();
assert_eq!(parent.batch(2).unwrap().view_origin(), None);
}
#[test]
fn batch_rejects_out_of_bounds_index() {
let parent = Tensor::<u8>::new(&[4, 2, 2, 3], Some(TensorMemory::Mem), None).unwrap();
match parent.batch(4) {
Err(Error::BatchIndexOutOfBounds { index, batch }) => {
assert_eq!((index, batch), (4, 4));
}
other => panic!("expected BatchIndexOutOfBounds, got {other:?}"),
}
}
#[test]
fn batch_zero_on_unit_n_is_whole() {
let parent = Tensor::<u8>::new(&[1, 2, 2, 3], Some(TensorMemory::Mem), None).unwrap();
let e = parent.batch(0).unwrap();
assert_eq!(e.shape(), &[2, 2, 3]);
assert_eq!(e.plane_offset(), None);
assert_eq!(e.buffer_identity().id(), parent.buffer_identity().id());
}
#[test]
fn mem_subview_rejects_unaligned_offset() {
let parent = Tensor::<f32>::new(&[8], Some(TensorMemory::Mem), None).unwrap();
assert!(parent.subview(2, &[1]).is_err());
assert!(parent.subview(4, &[1]).is_ok());
}
#[test]
fn mem_subview_rejects_out_of_bounds() {
let parent = Tensor::<u8>::new(&[8], Some(TensorMemory::Mem), None).unwrap();
assert!(parent.subview(6, &[4]).is_err());
}
#[test]
fn subview_shares_buffer_identity_all_backends() {
let _lock = crate::tests::fd_lock_shared();
let assert_parent_shares = |parent: Tensor<u8>, window: &[usize], label: &str| {
let parent_id = parent.buffer_identity().id();
let v0 = parent
.subview(0, window)
.unwrap_or_else(|e| panic!("{label}: subview(0) failed: {e:?}"));
let v1 = parent
.subview(16, window)
.unwrap_or_else(|e| panic!("{label}: subview(16) failed: {e:?}"));
assert_eq!(
v0.buffer_identity().id(),
parent_id,
"{label}: subview(0) minted a fresh BufferIdentity"
);
assert_eq!(
v1.buffer_identity().id(),
parent_id,
"{label}: subview(16) minted a fresh BufferIdentity"
);
};
let assert_shares = |memory: TensorMemory, label: &str| {
let parent = Tensor::<u8>::new(&[64], Some(memory), None)
.unwrap_or_else(|e| panic!("{label}: parent alloc failed: {e:?}"));
assert_parent_shares(parent, &[16], label);
};
assert_shares(TensorMemory::Mem, "Mem");
#[cfg(unix)]
if crate::is_shm_available() {
assert_shares(TensorMemory::Shm, "Shm");
}
#[cfg(target_os = "linux")]
if crate::test_support::dma_or_skip("subview_shares_buffer_identity_all_backends (Dma)") {
assert_shares(TensorMemory::DmaBuf, "Dma");
}
#[cfg(not(any(target_os = "linux", target_os = "windows")))]
if crate::is_gpu_buffer_available() {
assert_shares(TensorMemory::DmaBuf, "Dma");
}
#[cfg(target_os = "windows")]
if crate::is_gpu_buffer_available() {
let parent = Tensor::<u8>::image(
64,
1,
PixelFormat::Grey,
Some(TensorMemory::DmaBuf),
CpuAccess::ReadWrite,
)
.expect("D3D11 texture alloc failed");
assert_parent_shares(parent, &[1, 16, 1], "Dma");
}
}
#[cfg(target_os = "windows")]
#[test]
fn subview_of_a_texture_tensor_maps_at_its_offset() {
if !crate::is_gpu_buffer_available() {
return;
}
let t = Tensor::<u8>::image(
32,
8,
PixelFormat::Rgba,
Some(TensorMemory::DmaBuf),
CpuAccess::ReadWrite,
)
.unwrap();
let stride = t.effective_row_stride().unwrap();
let v = t.subview(stride * 2, &[2, 32, 4]).unwrap();
assert!(
v.plane_offset().is_some_and(|o| o > 0),
"a sub-view at a row offset must carry one"
);
let m = v
.map_read()
.expect("a sub-view of a texture tensor must map");
assert_eq!(m.as_slice().len(), 2 * stride, "two rows of the window");
}
#[test]
fn mem_subview_four_views_no_aliasing() {
let parent = Tensor::<f32>::new(&[4, 3], Some(TensorMemory::Mem), None).unwrap();
let frame = 3 * std::mem::size_of::<f32>();
for i in 0..4 {
let v = parent.subview(i * frame, &[1, 3]).unwrap();
let val = i as f32 + 1.0;
v.map()
.unwrap()
.as_mut_slice()
.copy_from_slice(&[val, val, val]);
}
assert_eq!(
parent.map().unwrap().as_slice(),
&[1.0, 1.0, 1.0, 2.0, 2.0, 2.0, 3.0, 3.0, 3.0, 4.0, 4.0, 4.0]
);
}
#[test]
fn mem_subview_inherits_format_and_row_stride() {
let mut parent = Tensor::<u8>::image(
100,
100,
PixelFormat::Rgba,
Some(TensorMemory::Mem),
crate::CpuAccess::ReadWrite,
)
.unwrap();
parent.set_row_stride_unchecked(512); let view = parent.subview(4096, &[10, 10, 4]).unwrap();
assert_eq!(view.format(), Some(PixelFormat::Rgba), "format inherited");
assert_eq!(view.row_stride(), Some(512), "row_stride inherited");
}
#[test]
fn mem_strided_subview_maps_offset_and_byte_size() {
let parent = Tensor::<u8>::new(&[2048], Some(TensorMemory::Mem), None).unwrap();
let mut view = parent.subview(128, &[8, 16]).unwrap(); assert_eq!(view.plane_offset(), Some(128));
view.set_row_stride_unchecked(32);
{
let mut m = view.map().unwrap();
let s = m.as_mut_slice();
assert_eq!(
s.len(),
256,
"strided map exposes the full padded byte window"
);
s[0] = 0xAA; s[32] = 0xBB; }
let p = parent.map().unwrap();
let pb = p.as_slice();
assert_eq!(pb[128], 0xAA, "row 0 writes at parent offset 128");
assert_eq!(
pb[128 + 32],
0xBB,
"row 1 writes at parent offset 128 + stride"
);
}
#[test]
#[cfg(unix)]
fn shm_subview_partitions_parent_buffer() {
let _lock = crate::tests::fd_lock_shared();
if !crate::is_shm_available() {
crate::test_support::report_skip("shm not available");
return;
}
let parent = Tensor::<u8>::new(&[2, 4], Some(TensorMemory::Shm), None).unwrap();
let view0 = parent.subview(0, &[1, 4]).expect("shm subview at offset 0");
let view1 = parent.subview(4, &[1, 4]).expect("shm subview at offset 4");
view1
.map()
.unwrap()
.as_mut_slice()
.copy_from_slice(&[10, 20, 30, 40]);
view0
.map()
.unwrap()
.as_mut_slice()
.copy_from_slice(&[1, 2, 3, 4]);
assert_eq!(view0.map().unwrap().as_slice(), &[1, 2, 3, 4]);
assert_eq!(view1.map().unwrap().as_slice(), &[10, 20, 30, 40]);
assert_eq!(
parent.map().unwrap().as_slice(),
&[1, 2, 3, 4, 10, 20, 30, 40]
);
let nested = view1.subview(2, &[1, 2]).expect("nested shm subview");
assert_eq!(nested.map().unwrap().as_slice(), &[30, 40]);
}
#[test]
#[cfg(unix)]
fn shm_subview_rejects_unaligned_and_oob() {
let _lock = crate::tests::fd_lock_shared();
if !crate::is_shm_available() {
crate::test_support::report_skip("shm not available");
return;
}
let parent = Tensor::<f32>::new(&[8], Some(TensorMemory::Shm), None).unwrap();
assert!(parent.subview(2, &[1]).is_err());
assert!(parent.subview(4, &[1]).is_ok());
let p2 = Tensor::<u8>::new(&[8], Some(TensorMemory::Shm), None).unwrap();
assert!(p2.subview(6, &[4]).is_err());
}
#[test]
#[cfg(target_os = "linux")]
fn dma_subview_matches_mem_subview() {
let _lock = crate::tests::fd_lock_shared();
if !crate::test_support::dma_or_skip("dma_subview_matches_mem_subview") {
return;
}
let dma = Tensor::<u8>::new(&[8], Some(TensorMemory::DmaBuf), None).unwrap();
let mem = Tensor::<u8>::new(&[8], Some(TensorMemory::Mem), None).unwrap();
for parent in [&dma, &mem] {
let v0 = parent.subview(0, &[4]).unwrap();
let v1 = parent.subview(4, &[4]).unwrap();
v0.map()
.unwrap()
.as_mut_slice()
.copy_from_slice(&[1, 2, 3, 4]);
v1.map()
.unwrap()
.as_mut_slice()
.copy_from_slice(&[5, 6, 7, 8]);
assert_eq!(parent.map().unwrap().as_slice(), &[1, 2, 3, 4, 5, 6, 7, 8]);
}
}
#[test]
#[cfg(target_os = "linux")]
fn dma_strided_subview_maps_padded_window() {
let _lock = crate::tests::fd_lock_shared();
if !crate::test_support::dma_or_skip("dma_strided_subview_maps_padded_window") {
return;
}
let parent = Tensor::<u8>::new(&[2048], Some(TensorMemory::DmaBuf), None).unwrap();
let mut view = parent.subview(128, &[8, 16]).unwrap();
assert_eq!(view.plane_offset(), Some(128));
view.set_row_stride_unchecked(32);
{
let mut m = view.map().unwrap();
let s = m.as_mut_slice();
assert_eq!(s.len(), 256, "strided DMA map exposes stride(32) × rows(8)");
s[0] = 0xAA; s[32] = 0xBB; }
let p = parent.map().unwrap();
let pb = p.as_slice();
assert_eq!(pb[128], 0xAA, "row 0 writes at parent offset 128");
assert_eq!(
pb[128 + 32],
0xBB,
"row 1 writes at parent offset 128 + stride"
);
}
#[test]
#[cfg(target_os = "linux")]
fn view_single_row_snapshots_parent_stride() {
let _lock = crate::tests::fd_lock_shared();
if !crate::test_support::dma_or_skip("view_single_row_snapshots_parent_stride") {
return;
}
let parent = Tensor::<u8>::image_with_stride(
8,
4,
PixelFormat::Rgba,
64,
Some(TensorMemory::DmaBuf),
crate::CpuAccess::ReadWrite,
)
.unwrap();
assert_eq!(parent.effective_row_stride(), Some(64));
let row = parent.view(Region::new(2, 3, 4, 1)).unwrap();
assert_eq!(row.effective_row_stride(), Some(16));
let vo = row.view_origin().expect("a view carries a view_origin");
assert_eq!(
vo.parent_row_stride, 64,
"GL keys/pitches a view on the parent stride, not its tight one"
);
assert_eq!(row.map().unwrap().as_slice().len(), 16);
}
#[test]
fn test_mem_tensor() {
let shape = vec![2, 3, 4];
let tensor =
MemTensor::<f32>::new(&shape, Some("test_tensor")).expect("Failed to create tensor");
assert_eq!(tensor.shape(), &shape);
assert_eq!(tensor.size(), 2 * 3 * 4 * std::mem::size_of::<f32>());
assert_eq!(tensor.name(), "test_tensor");
{
let mut tensor_map = tensor.map().expect("Failed to map memory");
tensor_map.fill(42.0);
assert!(tensor_map.iter().all(|&x| x == 42.0));
}
let mut tensor = MemTensor::<u8>::new(&shape, None).expect("Failed to create tensor");
assert_eq!(tensor.shape(), &shape);
let new_shape = vec![3, 4, 4];
assert!(
tensor.reshape(&new_shape).is_err(),
"Reshape should fail due to size mismatch"
);
assert_eq!(tensor.shape(), &shape, "Shape should remain unchanged");
let new_shape = vec![2, 3, 4];
tensor.reshape(&new_shape).expect("Reshape should succeed");
assert_eq!(
tensor.shape(),
&new_shape,
"Shape should be updated after successful reshape"
);
{
let mut tensor_map = tensor.map().expect("Failed to map memory");
tensor_map.fill(1);
assert!(tensor_map.iter().all(|&x| x == 1));
}
{
let mut tensor_map = tensor.map().expect("Failed to map memory");
tensor_map[2] = 42;
assert_eq!(tensor_map[1], 1, "Value at index 1 should be 1");
assert_eq!(tensor_map[2], 42, "Value at index 2 should be 42");
}
}
#[test]
#[cfg(target_os = "linux")]
fn test_dma_no_fd_leaks() {
let _lock = fd_lock_exclusive();
if !crate::test_support::dma_or_skip(function!()) {
return;
}
let baseline = open_fd_count(DMA_FD_TARGETS);
{
let probe = Tensor::<u8>::new(&[100, 100], Some(TensorMemory::DmaBuf), None)
.expect("Failed to create tensor");
assert!(
open_fd_count(DMA_FD_TARGETS) > baseline,
"a live DMA-BUF tensor must be visible to the dmabuf fd counter, \
otherwise this test cannot detect a leak"
);
drop(probe);
}
for _ in 0..100 {
let tensor = Tensor::<u8>::new(&[100, 100], Some(TensorMemory::DmaBuf), None)
.expect("Failed to create tensor");
let mut map = tensor.map().unwrap();
map.as_mut_slice().fill(233);
}
let after = open_fd_count(DMA_FD_TARGETS);
assert_eq!(
baseline, after,
"DMA-BUF file descriptor leak detected: {baseline} -> {after}"
);
}
#[test]
#[cfg(target_os = "linux")]
fn test_dma_from_fd_no_fd_leaks() {
let _lock = fd_lock_exclusive();
if !crate::test_support::dma_or_skip(function!()) {
return;
}
let baseline = open_fd_count(DMA_FD_TARGETS);
let orig = Tensor::<u8>::new(&[100, 100], Some(TensorMemory::DmaBuf), None).unwrap();
let held = open_fd_count(DMA_FD_TARGETS);
assert!(
held > baseline,
"a live DMA-BUF tensor must be visible to the dmabuf fd counter, \
otherwise this test cannot detect a leak"
);
for _ in 0..100 {
let tensor =
Tensor::<u8>::from_fd(orig.clone_fd().unwrap(), orig.shape(), None).unwrap();
assert_eq!(
tensor.memory(),
TensorMemory::DmaBuf,
"DMA-BUF fd must import as Dma, not be silently downgraded"
);
let mut map = tensor.map().unwrap();
map.as_mut_slice().fill(233);
}
let after_loop = open_fd_count(DMA_FD_TARGETS);
assert_eq!(
held, after_loop,
"DMA-BUF file descriptor leak detected: {held} -> {after_loop}"
);
drop(orig);
let after_drop = open_fd_count(DMA_FD_TARGETS);
assert_eq!(
baseline, after_drop,
"dropping the source tensor must release its fd: {baseline} -> {after_drop}"
);
}
#[test]
#[cfg(target_os = "linux")]
fn test_fs_magic_normalizes_sign_extended_values() {
let _lock = crate::tests::fd_lock_shared();
const HUGETLBFS_MAGIC: u32 = 0x9584_58f6;
const F2FS_MAGIC: u32 = 0xf2f5_2010;
const BTRFS_MAGIC: u32 = 0x9123_683e;
for magic in [HUGETLBFS_MAGIC, F2FS_MAGIC, BTRFS_MAGIC] {
assert_eq!(fs_magic(i64::from(magic)), magic);
let sign_extended = i64::from(magic as i32);
assert!(sign_extended < 0, "{magic:#x} should have bit 31 set");
assert_eq!(
fs_magic(sign_extended),
magic,
"{magic:#x} must survive a signed 32-bit f_type"
);
}
for magic in [DMA_BUF_MAGIC, TMPFS_MAGIC] {
assert!(magic < 0x8000_0000);
assert_eq!(fs_magic(i64::from(magic as i32)), magic);
}
}
#[test]
#[cfg(target_os = "linux")]
fn test_from_fd_dma_imports_as_dma() {
let _lock = crate::tests::fd_lock_shared();
if !crate::test_support::dma_or_skip(function!()) {
return;
}
let orig = Tensor::<u8>::new(&[64, 64], Some(TensorMemory::DmaBuf), None).unwrap();
assert_eq!(orig.memory(), TensorMemory::DmaBuf);
let imported = Tensor::<u8>::from_fd(orig.clone_fd().unwrap(), orig.shape(), None).unwrap();
assert_eq!(
imported.memory(),
TensorMemory::DmaBuf,
"a DMA-BUF fd must import as Dma"
);
}
#[test]
#[cfg(target_os = "linux")]
fn test_from_fd_shm_imports_as_shm() {
let _lock = crate::tests::fd_lock_shared();
if !is_shm_available() {
log::warn!("SKIPPED: {} - SHM memory not available", function!());
return;
}
let orig = Tensor::<u8>::new(&[64, 64], Some(TensorMemory::Shm), None).unwrap();
assert_eq!(orig.memory(), TensorMemory::Shm);
let imported = Tensor::<u8>::from_fd(orig.clone_fd().unwrap(), orig.shape(), None).unwrap();
assert_eq!(
imported.memory(),
TensorMemory::Shm,
"a tmpfs fd must import as Shm"
);
}
#[test]
#[cfg(target_os = "linux")]
fn test_from_fd_rejects_unknown_filesystem() {
let _lock = crate::tests::fd_lock_shared();
let (read_end, _write_end) = nix::unistd::pipe().unwrap();
let result = Tensor::<u8>::from_fd(read_end, &[64], None);
match result {
Err(Error::UnknownBufferType(magic)) => {
assert_eq!(magic, 0x5049_5045, "expected PIPEFS_MAGIC");
}
other => panic!("expected UnknownBufferType for a pipe fd, got {other:?}"),
}
}
#[test]
#[cfg(target_os = "linux")]
fn test_shm_no_fd_leaks() {
let _lock = fd_lock_exclusive();
if !is_shm_available() {
log::warn!(
"SKIPPED: {} - SHM memory allocation not available (permission denied or no SHM support)",
function!()
);
return;
}
let baseline = open_fd_count(SHM_FD_TARGETS);
{
let probe = Tensor::<u8>::new(&[100, 100], Some(TensorMemory::Shm), None)
.expect("Failed to create tensor");
assert!(
open_fd_count(SHM_FD_TARGETS) > baseline,
"a live shm tensor must be visible to the shm fd counter, \
otherwise this test cannot detect a leak"
);
drop(probe);
}
for _ in 0..100 {
let tensor = Tensor::<u8>::new(&[100, 100], Some(TensorMemory::Shm), None)
.expect("Failed to create tensor");
let mut map = tensor.map().unwrap();
map.as_mut_slice().fill(233);
}
let after = open_fd_count(SHM_FD_TARGETS);
assert_eq!(
baseline, after,
"shm file descriptor leak detected: {baseline} -> {after}"
);
}
#[test]
#[cfg(target_os = "linux")]
fn test_shm_from_fd_no_fd_leaks() {
let _lock = fd_lock_exclusive();
if !is_shm_available() {
log::warn!(
"SKIPPED: {} - SHM memory allocation not available (permission denied or no SHM support)",
function!()
);
return;
}
let baseline = open_fd_count(SHM_FD_TARGETS);
let orig = Tensor::<u8>::new(&[100, 100], Some(TensorMemory::Shm), None).unwrap();
let held = open_fd_count(SHM_FD_TARGETS);
assert!(
held > baseline,
"a live shm tensor must be visible to the shm fd counter, \
otherwise this test cannot detect a leak"
);
for _ in 0..100 {
let tensor =
Tensor::<u8>::from_fd(orig.clone_fd().unwrap(), orig.shape(), None).unwrap();
let mut map = tensor.map().unwrap();
map.as_mut_slice().fill(233);
}
let after_loop = open_fd_count(SHM_FD_TARGETS);
assert_eq!(
held, after_loop,
"shm file descriptor leak detected: {held} -> {after_loop}"
);
drop(orig);
let after_drop = open_fd_count(SHM_FD_TARGETS);
assert_eq!(
baseline, after_drop,
"dropping the source tensor must release its fd: {baseline} -> {after_drop}"
);
}
#[cfg(feature = "ndarray")]
#[test]
fn test_ndarray() {
let _lock = crate::tests::fd_lock_shared();
let shape = vec![2, 3, 4];
let tensor = Tensor::<f32>::new(&shape, None, None).expect("Failed to create tensor");
let mut tensor_map = tensor.map().expect("Failed to map tensor memory");
tensor_map.fill(1.0);
let view = tensor_map.view().expect("Failed to get ndarray view");
assert_eq!(view.shape(), &[2, 3, 4]);
assert!(view.iter().all(|&x| x == 1.0));
let mut view_mut = tensor_map
.view_mut()
.expect("Failed to get mutable ndarray view");
view_mut[[0, 0, 0]] = 42.0;
assert_eq!(view_mut[[0, 0, 0]], 42.0);
assert_eq!(tensor_map[0], 42.0, "Value at index 0 should be 42");
}
#[test]
fn test_buffer_identity_derived_equality() {
let id1 = BufferIdentity::derived(IdentityKind::DmaBuf, 4242);
let id2 = BufferIdentity::derived(IdentityKind::DmaBuf, 4242);
assert_eq!(
id1.id(),
id2.id(),
"Identities derived from the same kind and key should match"
);
let id3 = BufferIdentity::derived(IdentityKind::DmaBuf, 4243);
assert_ne!(
id1.id(),
id3.id(),
"Identities derived from different keys should differ"
);
}
#[test]
fn test_buffer_identity_clone_shares_guard() {
let id1 = BufferIdentity::derived(IdentityKind::DmaBuf, 4242);
let weak = id1.weak();
assert!(
weak.upgrade().is_some(),
"Weak should be alive while original exists"
);
let id2 = id1.clone();
assert_eq!(id1.id(), id2.id(), "Cloned identity should have same id");
drop(id1);
assert!(
weak.upgrade().is_some(),
"Weak should still be alive (clone holds Arc)"
);
drop(id2);
assert!(
weak.upgrade().is_none(),
"Weak should be dead after all clones dropped"
);
}
#[test]
fn test_tensor_buffer_identity() {
let t1 = Tensor::<u8>::new(&[100], Some(TensorMemory::Mem), Some("t1")).unwrap();
let t2 = Tensor::<u8>::new(&[100], Some(TensorMemory::Mem), Some("t2")).unwrap();
assert_ne!(
t1.buffer_identity().id(),
t2.buffer_identity().id(),
"Different tensors should have different buffer ids"
);
}
#[test]
fn test_quantization_per_tensor_constructors() {
let q = Quantization::per_tensor(0.1, -5);
assert!(q.is_per_tensor());
assert!(!q.is_per_channel());
assert!(!q.is_symmetric());
assert_eq!(q.scale(), &[0.1]);
assert_eq!(q.zero_point(), Some(&[-5][..]));
let qs = Quantization::per_tensor_symmetric(0.05);
assert!(qs.is_per_tensor());
assert!(qs.is_symmetric());
assert_eq!(qs.zero_point(), None);
}
#[test]
fn test_quantization_per_channel_constructors() {
let q = Quantization::per_channel(vec![0.1, 0.2, 0.3], vec![0, -1, 1], 2).unwrap();
assert!(q.is_per_channel());
assert!(!q.is_symmetric());
assert_eq!(q.axis(), Some(2));
assert_eq!(q.scale().len(), 3);
let qs = Quantization::per_channel_symmetric(vec![0.054, 0.089, 0.195], 0).unwrap();
assert!(qs.is_per_channel());
assert!(qs.is_symmetric());
assert_eq!(qs.axis(), Some(0));
}
#[test]
fn test_quantization_per_channel_length_mismatch_rejected() {
let err = Quantization::per_channel(vec![0.1, 0.2], vec![0, 0, 0], 0).unwrap_err();
assert!(matches!(err, Error::QuantizationInvalid { .. }));
}
#[test]
fn test_quantization_per_channel_empty_rejected() {
let err = Quantization::per_channel_symmetric(vec![], 0).unwrap_err();
assert!(matches!(err, Error::QuantizationInvalid { .. }));
}
#[test]
fn test_quantization_validate_rejects_malformed_deserialize() {
let mut t = Tensor::<i8>::new(&[1, 1, 4], Some(TensorMemory::Mem), None).unwrap();
let q: Quantization = serde_json::from_str(r#"{"scale": []}"#).unwrap();
assert!(matches!(
t.set_quantization(q).unwrap_err(),
Error::QuantizationInvalid { .. }
));
let q: Quantization =
serde_json::from_str(r#"{"scale": 0.1, "zero_point": [0, 0, 0]}"#).unwrap();
assert!(matches!(
t.set_quantization(q).unwrap_err(),
Error::QuantizationInvalid { .. }
));
let q: Quantization = serde_json::from_str(
r#"{"scale": [0.1, 0.2, 0.3, 0.4], "zero_point": [0, 0], "axis": 2}"#,
)
.unwrap();
assert!(matches!(
t.set_quantization(q).unwrap_err(),
Error::QuantizationInvalid { .. }
));
}
#[test]
fn test_quantization_mode_dispatch() {
let pt = Quantization::per_tensor(0.1, -5);
assert!(matches!(
pt.mode(),
QuantMode::PerTensor { scale, zero_point } if scale == 0.1 && zero_point == -5
));
let pts = Quantization::per_tensor_symmetric(0.05);
assert!(matches!(
pts.mode(),
QuantMode::PerTensorSymmetric { scale } if scale == 0.05
));
let pc = Quantization::per_channel(vec![0.1, 0.2], vec![0, -1], 2).unwrap();
assert!(matches!(pc.mode(), QuantMode::PerChannel { axis: 2, .. }));
let pcs = Quantization::per_channel_symmetric(vec![0.1, 0.2], 0).unwrap();
assert!(matches!(
pcs.mode(),
QuantMode::PerChannelSymmetric { axis: 0, .. }
));
}
#[test]
fn test_tensor_quantization_roundtrip_integer() {
let mut t = Tensor::<i8>::new(&[2, 3, 4], Some(TensorMemory::Mem), None).unwrap();
assert!(t.quantization().is_none());
t.set_quantization(Quantization::per_tensor(0.1, -5))
.unwrap();
let q = t.quantization().unwrap();
assert_eq!(q.scale(), &[0.1]);
t.clear_quantization();
assert!(t.quantization().is_none());
}
#[test]
fn test_tensor_with_quantization_builder() {
let t = Tensor::<i8>::new(&[4, 4], Some(TensorMemory::Mem), None)
.unwrap()
.with_quantization(Quantization::per_tensor_symmetric(0.05))
.unwrap();
assert!(t.quantization().is_some());
}
#[test]
fn test_tensor_dyn_quantization_float_arm_returns_none() {
let t = Tensor::<f32>::new(&[2, 2], Some(TensorMemory::Mem), None).unwrap();
let td = TensorDyn::F32(t);
assert!(td.quantization().is_none());
}
#[test]
fn test_tensor_dyn_set_quantization_float_arm_errors() {
let t = Tensor::<f32>::new(&[2, 2], Some(TensorMemory::Mem), None).unwrap();
let mut td = TensorDyn::F32(t);
let err = td
.set_quantization(Quantization::per_tensor(0.1, 0))
.unwrap_err();
assert!(matches!(err, Error::QuantizationInvalid { .. }));
}
fn _compile_fail_doctest_anchor() {}
pub static FD_LOCK: RwLock<()> = RwLock::new(());
pub fn fd_lock_shared() -> std::sync::RwLockReadGuard<'static, ()> {
FD_LOCK.read().unwrap_or_else(PoisonError::into_inner)
}
#[cfg(target_os = "linux")]
pub fn fd_lock_exclusive() -> std::sync::RwLockWriteGuard<'static, ()> {
FD_LOCK.write().unwrap_or_else(PoisonError::into_inner)
}
#[cfg(target_os = "linux")]
const SHM_FD_TARGETS: &[&str] = &["/dev/shm/"];
#[cfg(target_os = "linux")]
const DMA_FD_TARGETS: &[&str] = &["dmabuf", "dma_heap"];
#[cfg(target_os = "linux")]
fn open_fd_count(needles: &[&str]) -> usize {
let proc = procfs::process::Process::myself()
.expect("Failed to get current process using /proc/self");
proc.fd()
.expect("Failed to list open file descriptors")
.filter_map(Result::ok)
.filter(|entry| {
let target = format!("{:?}", entry.target);
needles.iter().any(|needle| target.contains(needle))
})
.count()
}
#[test]
#[cfg(not(target_os = "linux"))]
fn test_dma_not_available_on_non_linux() {
assert!(
!is_dma_available(),
"DMA memory allocation should NOT be available on non-Linux platforms"
);
}
#[test]
fn colorimetry_defaults_none_and_roundtrips_without_auto_fill() {
use crate::{ColorEncoding, ColorRange, Colorimetry, PixelFormat, TensorMemory};
let mut t = Tensor::<u8>::image(
1280,
720,
PixelFormat::Nv12,
Some(TensorMemory::Mem),
crate::CpuAccess::ReadWrite,
)
.unwrap();
assert_eq!(t.colorimetry(), None); let c = Colorimetry::default()
.with_encoding(ColorEncoding::Bt709)
.with_range(ColorRange::Limited);
t.set_colorimetry(Some(c));
assert_eq!(t.colorimetry(), Some(c));
t.configure_image(640, 480, PixelFormat::Grey).unwrap();
assert_eq!(t.colorimetry(), Some(c));
}
#[test]
fn configure_image_within_capacity() {
let mut t = Tensor::<u8>::image_with_capacity(
640,
480,
PixelFormat::Rgb,
None,
crate::CpuAccess::ReadWrite,
)
.unwrap();
t.configure_image(320, 240, PixelFormat::Nv12).unwrap();
assert_eq!(t.format(), Some(PixelFormat::Nv12));
assert_eq!(t.width(), Some(320));
assert_eq!(t.height(), Some(240));
assert_eq!(t.shape(), &[360, 320]); }
#[test]
fn configure_image_too_large_errors() {
let mut t = Tensor::<u8>::image_with_capacity(
64,
64,
PixelFormat::Grey,
None,
crate::CpuAccess::ReadWrite,
)
.unwrap();
let err = t
.configure_image(1920, 1080, PixelFormat::Nv12)
.unwrap_err();
assert!(matches!(err, Error::InsufficientCapacity { .. }));
}
#[test]
#[cfg(target_os = "macos")]
fn configure_image_preserves_iosurface_physical_stride() {
let _lock = crate::tests::fd_lock_shared();
let mut pool = Tensor::<u8>::image(
100,
64,
PixelFormat::Grey,
Some(TensorMemory::DmaBuf),
crate::CpuAccess::ReadWrite,
)
.unwrap();
let pitch = pool.effective_row_stride().unwrap();
assert!(
pitch >= 128 && pitch.is_multiple_of(64),
"padded bytesPerRow, got {pitch}"
);
pool.configure_image(32, 16, PixelFormat::Nv12).unwrap();
assert_eq!(pool.format(), Some(PixelFormat::Nv12));
assert_eq!(pool.width(), Some(32));
assert_eq!(pool.height(), Some(16));
assert_eq!(
pool.effective_row_stride(),
Some(pitch),
"configure_image must preserve the IOSurface physical bytesPerRow"
);
pool.configure_image(32, 16, PixelFormat::Nv24).unwrap();
assert_eq!(pool.effective_row_stride(), Some(pitch));
}
#[test]
fn configure_image_mem_aligns_stride() {
let mut t = Tensor::<u8>::image_with_capacity(
64,
64,
PixelFormat::Rgba,
Some(TensorMemory::Mem),
crate::CpuAccess::ReadWrite,
)
.unwrap();
t.configure_image(32, 16, PixelFormat::Nv12).unwrap();
let s = t.effective_row_stride().unwrap();
assert_eq!(s % 64, 0, "stride must be 64-aligned");
assert!(s >= 32, "stride must cover the even-width minimum");
assert_eq!(s, 64);
}
#[test]
fn strided_mem_tensor_cpu_maps_full_padded_buffer() {
let mut t = Tensor::<u8>::image_with_capacity(
16,
3,
PixelFormat::Rgba,
Some(TensorMemory::Mem),
crate::CpuAccess::ReadWrite,
)
.unwrap(); t.configure_image(8, 3, PixelFormat::Rgba).unwrap(); t.set_row_stride(48).unwrap();
let map = t.map().expect("strided Mem tensor should CPU-map");
assert_eq!(map.as_slice().len(), 144);
assert_eq!(map.shape(), &[3, 8, 4]);
}
#[test]
fn strided_mem_tensor_over_capacity_errors() {
let mut t = Tensor::<u8>::new(&[3, 8, 4], Some(TensorMemory::Mem), None).unwrap();
t.set_format(PixelFormat::Rgba).unwrap();
t.set_row_stride(64).unwrap();
assert!(matches!(t.map(), Err(Error::InsufficientCapacity { .. })));
}
#[test]
#[cfg(unix)]
fn test_shm_available_and_usable() {
let _lock = crate::tests::fd_lock_shared();
if !is_shm_available() {
log::warn!(
"SKIPPED: test_shm_available_and_usable - SHM allocation \
unavailable on this platform"
);
return;
}
let tensor = Tensor::<u8>::new(&[100, 100], Some(TensorMemory::Shm), None)
.expect("Failed to create SHM tensor");
let mut map = tensor.map().expect("Failed to map SHM tensor");
map.as_mut_slice().fill(0xAB);
assert!(
map.as_slice().iter().all(|&b| b == 0xAB),
"SHM tensor data should be writable and readable"
);
}
#[test]
fn packed_rgba16f_layout_planar_rgb_f16() {
let layout =
packed_rgba16f_layout(PixelFormat::PlanarRgb, DType::F16, 640, 640).expect("Some");
assert_eq!(layout.surface_w, 160);
assert_eq!(layout.surface_h, 1920);
assert_eq!(layout.bytes_per_texel, 8);
assert_eq!(layout.pitch, 1280);
}
#[test]
fn packed_rgba16f_layout_planar_rgba_f16() {
let layout =
packed_rgba16f_layout(PixelFormat::PlanarRgba, DType::F16, 640, 640).expect("Some");
assert_eq!(layout.surface_w, 160);
assert_eq!(layout.surface_h, 2560); assert_eq!(layout.bytes_per_texel, 8);
assert_eq!(layout.pitch, 1280);
}
#[test]
fn packed_rgba16f_layout_rejects_misaligned() {
assert!(packed_rgba16f_layout(PixelFormat::PlanarRgb, DType::F16, 642, 640).is_none());
}
#[test]
fn packed_rgba16f_layout_rejects_non_f16() {
assert!(packed_rgba16f_layout(PixelFormat::PlanarRgb, DType::U8, 640, 640).is_none());
assert!(packed_rgba16f_layout(PixelFormat::Rgb, DType::F32, 640, 640).is_none());
assert!(packed_rgba16f_layout(PixelFormat::Rgba, DType::F16, 640, 640).is_none());
}
#[test]
fn cuda_map_fast_fails_to_none_without_handle() {
let t = Tensor::<f32>::new(&[4], Some(TensorMemory::Mem), None).unwrap();
assert!(t.cuda().is_none());
assert!(t.cuda_map().is_none()); }
#[test]
fn cuda_returns_none_without_handle() {
let t = Tensor::<f32>::new(&[2, 2], Some(TensorMemory::Mem), None).unwrap();
assert!(t.cuda().is_none(), "no CUDA handle on a Mem tensor");
assert!(t.cuda_map().is_none(), "fast-fail map → None");
}
#[test]
fn cuda_map_then_host_map_fallback() {
let t = Tensor::<f32>::new(&[2, 2], Some(TensorMemory::Mem), None).unwrap();
let cuda = t.cuda_map();
if let Some(_c) = cuda {
unreachable!("a Mem tensor has no CUDA handle");
} else {
let host = t.map().expect("host map fallback must succeed");
assert_eq!(host.len(), 4); }
}
#[test]
fn from_foreign_valid_wrap_and_readback() {
let mut buf: Vec<f32> = vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0];
let ptr = buf.as_mut_ptr();
let t = unsafe { Tensor::<f32>::from_foreign(ptr, &[2, 3], None, Some("test_foreign")) }
.expect("valid from_foreign must succeed");
assert_eq!(t.shape(), &[2, 3]);
assert_eq!(t.memory(), TensorMemory::Mem);
assert_eq!(t.name(), "test_foreign");
let m = t.map().unwrap();
assert_eq!(m.as_slice(), &[1.0_f32, 2.0, 3.0, 4.0, 5.0, 6.0]);
}
#[test]
fn from_foreign_write_visible_in_caller_allocation() {
let mut buf: Vec<u8> = vec![0u8; 6];
let ptr = buf.as_mut_ptr();
let t = unsafe { Tensor::<u8>::from_foreign(ptr, &[2, 3], None, None) }.unwrap();
{
let mut m = t.map().unwrap();
m.as_mut_slice().copy_from_slice(&[10, 20, 30, 40, 50, 60]);
}
drop(t);
assert_eq!(buf, vec![10, 20, 30, 40, 50, 60]);
}
#[test]
fn from_foreign_rejects_null_ptr() {
let err = unsafe { Tensor::<u8>::from_foreign(std::ptr::null_mut(), &[4], None, None) }
.unwrap_err();
assert!(
matches!(err, Error::InvalidArgument(ref m) if m.contains("non-null")),
"expected InvalidArgument(non-null), got {err:?}"
);
}
#[test]
fn from_foreign_rejects_empty_shape() {
let mut dummy: u8 = 0;
let err = unsafe { Tensor::<u8>::from_foreign(&mut dummy, &[], None, None) }.unwrap_err();
assert!(
matches!(err, Error::InvalidSize(0)),
"expected InvalidSize(0) for empty shape, got {err:?}"
);
}
#[test]
fn from_foreign_rejects_overflow_shape() {
let mut dummy: u8 = 0;
let huge = [usize::MAX / 2 + 1, 2];
let err = unsafe { Tensor::<u8>::from_foreign(&mut dummy, &huge, None, None) }.unwrap_err();
assert!(
matches!(err, Error::InvalidArgument(ref m) if m.contains("overflow")),
"expected InvalidArgument(overflow), got {err:?}"
);
}
#[test]
fn from_foreign_owner_keeps_allocation_alive() {
use std::sync::atomic::{AtomicBool, Ordering};
let flag = std::sync::Arc::new(AtomicBool::new(false));
let flag2 = flag.clone();
struct Guard(std::sync::Arc<AtomicBool>);
impl Drop for Guard {
fn drop(&mut self) {
self.0.store(true, Ordering::SeqCst);
}
}
let mut buf: Vec<u32> = vec![42u32; 4];
let ptr = buf.as_mut_ptr();
let owner: ForeignOwner = Box::new(Guard(flag2));
let t = unsafe { Tensor::<u32>::from_foreign(ptr, &[4], Some(owner), None) }.unwrap();
let m = t.map().unwrap();
assert_eq!(m.as_slice()[0], 42);
drop(t); assert!(
!flag.load(Ordering::SeqCst),
"owner must not drop while a map shares the backing"
);
drop(m);
assert!(
flag.load(Ordering::SeqCst),
"owner Drop must fire when the last Arc reference is released"
);
}
}