use std::fmt;
use std::ptr::NonNull;
use edgefirst_tensor_ffi::EfTensor;
use crate::{
BufferIdentity, Colorimetry, CpuAccess, DType, Error, IdentityKind, PixelFormat, PixelLayout,
Region, Result, TensorMemory, TensorTrait, ViewOrigin,
};
pub struct TensorDyn {
handle: NonNull<EfTensor>,
shape_cache: Vec<usize>,
identity: BufferIdentity,
quantization_cache: std::sync::OnceLock<Option<crate::Quantization>>,
pub(crate) multiplane_chroma: Option<Box<TensorDyn>>,
pub(crate) pbo: Option<Box<dyn std::any::Any + Send + Sync>>,
pub(crate) cuda: Option<Box<crate::cuda::CudaHandle>>,
}
unsafe impl Send for TensorDyn {}
unsafe impl Sync for TensorDyn {}
impl Drop for TensorDyn {
fn drop(&mut self) {
unsafe { edgefirst_tensor_ffi::ef_tensor_free(self.handle.as_ptr()) }
}
}
pub type Raw = *mut EfTensor;
impl TensorDyn {
fn from_handle(handle: NonNull<EfTensor>) -> Self {
let shape_cache = Self::query_shape(handle);
let identity = Self::derive_identity(handle);
TensorDyn {
handle,
shape_cache,
identity,
quantization_cache: std::sync::OnceLock::new(),
multiplane_chroma: None,
pbo: None,
cuda: None,
}
}
fn derive_identity(handle: NonNull<EfTensor>) -> BufferIdentity {
#[cfg(unix)]
{
let code = unsafe { edgefirst_tensor_ffi::ef_tensor_storage_kind(handle.as_ptr()) };
if TensorMemory::from_code(code) == Some(TensorMemory::DmaBuf) {
let mut plane = edgefirst_tensor_ffi::EfTensorPlane::default();
let rc = unsafe {
edgefirst_tensor_ffi::ef_tensor_plane_at(handle.as_ptr(), 0, &mut plane)
};
if rc == 0 && plane.handle >= 0 {
let fd = unsafe { std::os::fd::BorrowedFd::borrow_raw(plane.handle as i32) };
if let Ok(stat) = nix::sys::stat::fstat(fd) {
return identity_from_stat(&stat);
}
}
}
}
BufferIdentity::derived(IdentityKind::HostPtr, handle.as_ptr() as u64)
}
#[cfg(unix)]
pub(crate) fn shadow_multiplane_chroma(&self) -> Result<TensorDyn> {
let fd = self.clone_fd()?;
Self::from_fd(fd, self.shape(), self.dtype(), None)
}
fn query_shape(handle: NonNull<EfTensor>) -> Vec<usize> {
unsafe {
let ndim = edgefirst_tensor_ffi::ef_tensor_ndim(handle.as_ptr());
let ptr = edgefirst_tensor_ffi::ef_tensor_shape(handle.as_ptr());
if ptr.is_null() || ndim == 0 {
return Vec::new();
}
std::slice::from_raw_parts(ptr, ndim as usize)
.iter()
.map(|&d| d as usize)
.collect()
}
}
pub fn into_raw(self) -> Raw {
let p = self.handle.as_ptr();
std::mem::forget(self); p
}
pub unsafe fn from_raw(p: Raw) -> TensorDyn {
Self::from_handle(unsafe { NonNull::new_unchecked(p) })
}
pub unsafe fn with_raw<R>(p: Raw, f: impl FnOnce(&mut TensorDyn) -> R) -> R {
let mut t =
std::mem::ManuallyDrop::new(Self::from_handle(unsafe { NonNull::new_unchecked(p) }));
f(&mut t)
}
pub fn new(
shape: &[usize],
dtype: DType,
memory: Option<TensorMemory>,
_name: Option<&str>,
) -> Result<Self> {
unsafe {
let b = edgefirst_tensor_ffi::ef_tensor_builder_new();
if b.is_null() {
return Err(Error::NotImplemented(
"ef_tensor_builder_new: allocation failed".into(),
));
}
let dims: Vec<u64> = shape.iter().map(|&d| d as u64).collect();
edgefirst_tensor_ffi::ef_tensor_builder_dtype(b, dtype.code());
edgefirst_tensor_ffi::ef_tensor_builder_shape(b, dims.as_ptr(), dims.len() as u32);
if let Some(mem) = memory {
edgefirst_tensor_ffi::ef_tensor_builder_storage(b, mem.code());
}
let handle = edgefirst_tensor_ffi::ef_tensor_builder_alloc(b);
let result = match NonNull::new(handle) {
Some(h) => Ok(Self::from_handle(h)),
None => Err(builder_error(
edgefirst_tensor_ffi::ef_tensor_builder_error(b),
)),
};
edgefirst_tensor_ffi::ef_tensor_builder_free(b);
result
}
}
#[cfg(unix)]
pub fn from_fd(
fd: std::os::fd::OwnedFd,
shape: &[usize],
dtype: DType,
_name: Option<&str>,
) -> Result<Self> {
use std::os::fd::{FromRawFd, IntoRawFd};
unsafe {
let b = edgefirst_tensor_ffi::ef_tensor_builder_new();
if b.is_null() {
return Err(Error::NotImplemented(
"ef_tensor_builder_new: allocation failed".into(),
));
}
let dims: Vec<u64> = shape.iter().map(|&d| d as u64).collect();
edgefirst_tensor_ffi::ef_tensor_builder_dtype(b, dtype.code());
edgefirst_tensor_ffi::ef_tensor_builder_shape(b, dims.as_ptr(), dims.len() as u32);
let raw_fd = fd.into_raw_fd();
edgefirst_tensor_ffi::ef_tensor_builder_add_plane(b, raw_fd as i64, 0, 0, 0, 0, 0);
let handle = edgefirst_tensor_ffi::ef_tensor_builder_wrap(b);
let result = match NonNull::new(handle) {
Some(h) => Ok(Self::from_handle(h)),
None => {
drop(std::os::fd::OwnedFd::from_raw_fd(raw_fd));
let errno = edgefirst_tensor_ffi::ef_tensor_builder_error(b);
Err(Error::NotImplemented(format!(
"ef_tensor_builder_wrap failed: errno {errno}"
)))
}
};
edgefirst_tensor_ffi::ef_tensor_builder_free(b);
result
}
}
pub fn shape(&self) -> &[usize] {
&self.shape_cache
}
pub fn reshape(&mut self, shape: &[usize]) -> Result<()> {
self.reshape_impl(shape, false)
}
pub fn set_logical_shape(&mut self, shape: &[usize]) -> Result<()> {
self.reshape_impl(shape, true)
}
fn reshape_impl(&mut self, shape: &[usize], by_capacity: bool) -> Result<()> {
let dims: Vec<u64> = shape.iter().map(|&d| d as u64).collect();
let rc = unsafe {
if by_capacity {
edgefirst_tensor_ffi::ef_tensor_set_logical_shape(
self.handle.as_ptr(),
dims.as_ptr(),
dims.len() as u32,
)
} else {
edgefirst_tensor_ffi::ef_tensor_reshape(
self.handle.as_ptr(),
dims.as_ptr(),
dims.len() as u32,
)
}
};
if rc != 0 {
let msg = ffi_last_error();
return Err(if rc == libc::ERANGE {
Error::ShapeMismatch(msg)
} else {
Error::InvalidShape(msg)
});
}
self.shape_cache = Self::query_shape(self.handle);
self.sync_pbo_shape(shape, by_capacity)?;
Ok(())
}
fn sync_pbo_shape(&mut self, shape: &[usize], by_capacity: bool) -> Result<()> {
let r = if by_capacity {
with_pbo_mut!(self, set_logical_shape, shape)
} else {
with_pbo_mut!(self, reshape, shape)
};
match r {
None => Ok(()), Some(res) => res,
}
}
pub fn name(&self) -> String {
String::new()
}
pub fn memory(&self) -> TensorMemory {
if self.pbo.is_some() {
return TensorMemory::Pbo;
}
let code = unsafe { edgefirst_tensor_ffi::ef_tensor_storage_kind(self.handle.as_ptr()) };
TensorMemory::from_code(code).unwrap_or(TensorMemory::Mem)
}
#[cfg(unix)]
pub fn clone_fd(&self) -> Result<std::os::fd::OwnedFd> {
use std::os::fd::FromRawFd;
let rc = unsafe { edgefirst_tensor_ffi::ef_tensor_clone_fd(self.handle.as_ptr()) };
if rc < 0 {
let msg = ffi_last_error();
return Err(if -rc == libc::ENOTSUP {
Error::NotImplemented(msg)
} else {
Error::IoError(std::io::Error::other(format!(
"clone_fd: {msg} (errno {})",
-rc
)))
});
}
Ok(unsafe { std::os::fd::OwnedFd::from_raw_fd(rc) })
}
pub fn buffer_identity(&self) -> &BufferIdentity {
&self.identity
}
pub(crate) fn map_pin(&self, access: CpuAccess) -> Result<crate::pin::HostPin<'static>> {
if self.pbo.is_some() {
return self.map_pin_pbo(access);
}
let code = match access {
CpuAccess::None => {
return Err(Error::InvalidArgument(
"map: CpuAccess::None is not a mappable direction".into(),
))
}
CpuAccess::Read => 1,
CpuAccess::Write => 2,
CpuAccess::ReadWrite => 3,
};
let mut view = edgefirst_tensor_ffi::EfTensorView {
ptr: std::ptr::null_mut(),
len: 0,
};
let rc =
unsafe { edgefirst_tensor_ffi::ef_tensor_map(self.handle.as_ptr(), code, &mut view) };
if rc != 0 {
return Err(Error::InvalidOperation(format!(
"ef_tensor_map failed: errno {rc}"
)));
}
let retain_rc = unsafe { edgefirst_tensor_ffi::ef_tensor_retain(self.handle.as_ptr()) };
if retain_rc != 0 {
unsafe { edgefirst_tensor_ffi::ef_tensor_unmap(self.handle.as_ptr()) };
return Err(Error::NotImplemented(format!(
"ef_tensor_retain failed: errno {retain_rc}"
)));
}
let keepalive: std::sync::Arc<dyn Send + Sync> =
std::sync::Arc::new(MapKeepalive(self.handle.as_ptr()));
Ok(crate::pin::HostPin::new(keepalive, view.ptr, view.len))
}
fn map_pin_pbo(&self, access: CpuAccess) -> Result<crate::pin::HostPin<'static>> {
macro_rules! map_arm {
($any:expr, $t:ty) => {
$any.downcast_ref::<crate::PboTensor<$t>>().map(|p| {
<crate::PboTensor<$t> as crate::TensorTrait<$t>>::map_with(p, access)
.map(|v| v.into_bytes())
})
};
}
let any = self
.pbo
.as_ref()
.ok_or_else(|| Error::NotImplemented("map: not a PBO-backed tensor".into()))?;
let mapped = None
.or_else(|| map_arm!(any, u8))
.or_else(|| map_arm!(any, i8))
.or_else(|| map_arm!(any, u16))
.or_else(|| map_arm!(any, i16))
.or_else(|| map_arm!(any, half::f16))
.or_else(|| map_arm!(any, u32))
.or_else(|| map_arm!(any, i32))
.or_else(|| map_arm!(any, f32))
.or_else(|| map_arm!(any, u64))
.or_else(|| map_arm!(any, i64))
.or_else(|| map_arm!(any, f64));
let mut view = mapped
.ok_or_else(|| {
Error::NotImplemented(
"map: the wrapped PboTensor's element type does not match this handle's dtype"
.into(),
)
})?
.map_err(|e| Error::InvalidOperation(format!("map: PBO map failed: {e}")))?;
let (ptr, len) = if access.writes() {
let s = crate::TensorMapTrait::as_mut_slice(&mut view);
(s.as_mut_ptr(), s.len())
} else {
let s = crate::TensorMapTrait::as_slice(&view);
(s.as_ptr() as *mut u8, s.len())
};
let keepalive: std::sync::Arc<dyn Send + Sync> = std::sync::Arc::new(view);
Ok(crate::pin::HostPin::new(keepalive, ptr, len))
}
pub fn map_bytes(&self, access: CpuAccess) -> Result<crate::view::HostView<'static, u8>> {
let pin = self.map_pin(access)?;
let len = pin.len();
Ok(crate::view::HostView::new(pin, vec![len], None, access))
}
pub fn pin_host(&self, access: CpuAccess) -> Result<crate::pin::HostPin<'static>> {
self.map_pin(access)
}
pub fn plane0(&self) -> Option<edgefirst_tensor_ffi::EfTensorPlane> {
let mut plane = edgefirst_tensor_ffi::EfTensorPlane::default();
let rc = unsafe {
edgefirst_tensor_ffi::ef_tensor_plane_at(self.handle.as_ptr(), 0, &mut plane)
};
(rc == 0).then_some(plane)
}
pub fn effective_row_stride(&self) -> Option<usize> {
self.plane0().map(|p| p.stride as usize)
}
pub fn width(&self) -> Option<usize> {
let fmt = self.format()?;
let shape = self.shape();
match fmt.layout() {
PixelLayout::Packed => shape.get(1).copied(),
PixelLayout::Planar => shape.get(2).copied(),
PixelLayout::SemiPlanar => shape.get(1).copied(),
}
}
pub fn height(&self) -> Option<usize> {
let fmt = self.format()?;
let shape = self.shape();
match fmt.layout() {
PixelLayout::Packed => shape.first().copied(),
PixelLayout::Planar => shape.get(1).copied(),
PixelLayout::SemiPlanar => match fmt {
PixelFormat::Nv12 => shape.first().map(|h| h * 2 / 3),
PixelFormat::Nv16 => shape.first().map(|h| h / 2),
PixelFormat::Nv24 => shape.first().map(|h| h / 3),
_ => None,
},
}
}
pub fn aliases(&self, other: &Self) -> bool {
if self.buffer_identity().id() == other.buffer_identity().id() {
return true;
}
if self.memory() != other.memory() {
return false;
}
if self.memory() == TensorMemory::DmaBuf {
if let (Some(a), Some(b)) = (self.plane0(), other.plane0()) {
if a.handle >= 0 && a.handle == b.handle {
return true;
}
}
}
false
}
pub fn colorimetry(&self) -> Option<Colorimetry> {
let packed = unsafe { edgefirst_tensor_ffi::ef_tensor_colorimetry(self.handle.as_ptr()) };
if packed == 0 {
None
} else {
Some(Colorimetry::unpack(packed))
}
}
pub fn set_colorimetry(&mut self, c: Option<Colorimetry>) {
let packed = c.map(|c| c.pack()).unwrap_or(0);
unsafe { edgefirst_tensor_ffi::ef_tensor_set_colorimetry(self.handle.as_ptr(), packed) };
}
pub fn with_colorimetry(mut self, c: Colorimetry) -> Self {
self.set_colorimetry(Some(c));
self
}
pub fn view_origin(&self) -> Option<ViewOrigin> {
let mut out = edgefirst_tensor_ffi::EfViewOrigin::default();
let rc =
unsafe { edgefirst_tensor_ffi::ef_tensor_view_origin(self.handle.as_ptr(), &mut out) };
if rc != 0 || out.has_origin == 0 {
return None;
}
Some(ViewOrigin {
parent_width: out.parent_width as usize,
parent_height: out.parent_height as usize,
parent_row_stride: out.parent_row_stride as usize,
x: out.x as usize,
y: out.y as usize,
})
}
pub fn dtype(&self) -> DType {
let code = unsafe { edgefirst_tensor_ffi::ef_tensor_dtype(self.handle.as_ptr()) };
DType::from_code(code).unwrap_or(DType::U8)
}
pub fn format(&self) -> Option<PixelFormat> {
let ptr = unsafe { edgefirst_tensor_ffi::ef_tensor_format(self.handle.as_ptr()) };
if ptr.is_null() {
return None;
}
let s = unsafe { std::ffi::CStr::from_ptr(ptr) }.to_str().ok()?;
if s.is_empty() {
None
} else {
PixelFormat::from_str_code(s)
}
}
pub fn image(
width: usize,
height: usize,
format: PixelFormat,
dtype: DType,
memory: Option<TensorMemory>,
access: CpuAccess,
) -> Result<Self> {
let c_format = format_cstring(format)?;
let (has_memory, memory_code) = memory_code(memory);
let handle = unsafe {
edgefirst_tensor_ffi::ef_tensor_image_alloc(
width,
height,
c_format.as_ptr(),
dtype.code(),
has_memory,
memory_code,
access_code(access),
)
};
match NonNull::new(handle) {
Some(h) => Ok(Self::from_handle(h)),
None => Err(ffi_error(Error::NotImplemented)),
}
}
pub fn image_desc(desc: &crate::ImageDesc) -> Result<Self> {
let c_format = format_cstring(desc.format())?;
unsafe {
let d = edgefirst_tensor_ffi::ef_tensor_image_desc_new(
desc.width(),
desc.height(),
c_format.as_ptr(),
desc.dtype().code(),
);
if d.is_null() {
return Err(Error::NotImplemented(
"ef_tensor_image_desc_new: allocation failed".into(),
));
}
if let Some(m) = desc.memory() {
edgefirst_tensor_ffi::ef_tensor_image_desc_set_memory(d, m.code());
}
edgefirst_tensor_ffi::ef_tensor_image_desc_set_access(d, access_code(desc.access()));
if let Some(c) = desc.compression() {
let code = match c {
crate::Compression::Any => 1,
crate::Compression::Scheme(_) => 1,
};
edgefirst_tensor_ffi::ef_tensor_image_desc_set_compression(d, code);
}
let handle = edgefirst_tensor_ffi::ef_tensor_image_desc_alloc(d);
let result = match NonNull::new(handle) {
Some(h) => Ok(Self::from_handle(h)),
None => Err(ffi_error(Error::NotImplemented)),
};
edgefirst_tensor_ffi::ef_tensor_image_desc_free(d);
result
}
}
pub fn image_with_stride(
width: usize,
height: usize,
format: PixelFormat,
dtype: DType,
row_stride_bytes: usize,
memory: Option<TensorMemory>,
access: CpuAccess,
) -> Result<Self> {
let c_format = format_cstring(format)?;
let (has_memory, memory_code) = memory_code(memory);
let handle = unsafe {
edgefirst_tensor_ffi::ef_tensor_image_with_stride_alloc(
width,
height,
c_format.as_ptr(),
dtype.code(),
row_stride_bytes,
has_memory,
memory_code,
access_code(access),
)
};
match NonNull::new(handle) {
Some(h) => Ok(Self::from_handle(h)),
None => Err(ffi_error(Error::NotImplemented)),
}
}
pub fn view(&self, region: Region) -> Result<TensorDyn> {
let handle = unsafe {
edgefirst_tensor_ffi::ef_tensor_view_region(
self.handle.as_ptr(),
region.x as u64,
region.y as u64,
region.width as u64,
region.height as u64,
)
};
match NonNull::new(handle) {
Some(h) => Ok(Self::from_handle(h)),
None => {
if ffi_last_error_class() == edgefirst_tensor_ffi::EfErrorClass::RegionOutOfBounds {
Err(Error::RegionOutOfBounds {
region,
bounds: (self.width().unwrap_or(0), self.height().unwrap_or(0)),
})
} else {
Err(ffi_error(Error::InvalidShape))
}
}
}
}
pub fn set_format(&mut self, format: PixelFormat) -> Result<()> {
let c_format = format_cstring(format)?;
let rc = unsafe {
edgefirst_tensor_ffi::ef_tensor_set_format(self.handle.as_ptr(), c_format.as_ptr())
};
if rc != 0 {
return Err(Error::InvalidArgument(format!(
"ef_tensor_set_format failed: {}",
ffi_last_error()
)));
}
Ok(())
}
pub fn with_format(mut self, format: PixelFormat) -> Result<Self> {
self.set_format(format)?;
Ok(self)
}
pub fn set_row_stride(&mut self, stride: usize) -> Result<()> {
let rc =
unsafe { edgefirst_tensor_ffi::ef_tensor_set_row_stride(self.handle.as_ptr(), stride) };
if rc != 0 {
return Err(Error::InvalidArgument(format!(
"ef_tensor_set_row_stride failed: {}",
ffi_last_error()
)));
}
Ok(())
}
pub fn set_row_stride_unchecked(&mut self, stride: usize) {
unsafe {
edgefirst_tensor_ffi::ef_tensor_set_row_stride_unchecked(self.handle.as_ptr(), stride)
};
}
pub fn set_dtype(&mut self, dtype: DType) -> Result<()> {
if self.dtype() == dtype {
return Ok(());
}
let rc = unsafe {
edgefirst_tensor_ffi::ef_tensor_set_dtype(self.handle.as_ptr(), dtype.code())
};
if rc == 0 {
return Ok(());
}
Err(Error::InvalidArgument(format!(
"ef_tensor_set_dtype failed: {} (errno {rc})",
ffi_last_error()
)))
}
pub fn set_plane_offset(&mut self, offset: usize) {
unsafe { edgefirst_tensor_ffi::ef_tensor_set_plane_offset(self.handle.as_ptr(), offset) };
}
pub fn plane_offset(&self) -> Option<usize> {
let raw = unsafe { edgefirst_tensor_ffi::ef_tensor_plane_offset(self.handle.as_ptr()) };
(raw >= 0).then_some(raw as usize)
}
pub fn configure_image(
&mut self,
width: usize,
height: usize,
format: PixelFormat,
) -> Result<()> {
let c_format = format_cstring(format)?;
let rc = unsafe {
edgefirst_tensor_ffi::ef_tensor_configure_image(
self.handle.as_ptr(),
width,
height,
c_format.as_ptr(),
)
};
if rc != 0 {
return Err(Error::InvalidArgument(format!(
"ef_tensor_configure_image failed: {}",
ffi_last_error()
)));
}
self.shape_cache = Self::query_shape(self.handle);
let new_shape = self.shape_cache.clone();
self.sync_pbo_shape(&new_shape, true)?;
Ok(())
}
pub fn quantization(&self) -> Option<&crate::Quantization> {
self.quantization_cache
.get_or_init(|| self.fetch_quantization())
.as_ref()
}
fn fetch_quantization(&self) -> Option<crate::Quantization> {
let mut info = edgefirst_tensor_ffi::EfQuantizationInfo::default();
let rc = unsafe {
edgefirst_tensor_ffi::ef_tensor_quantization_info(self.handle.as_ptr(), &mut info)
};
if rc != 0 || info.has_quantization == 0 {
return None;
}
let n = info.count as usize;
let mut scales = vec![0f32; n];
let mut zps = vec![0i32; n];
let rc = unsafe {
edgefirst_tensor_ffi::ef_tensor_quantization_get(
self.handle.as_ptr(),
scales.as_mut_ptr(),
zps.as_mut_ptr(),
n as u32,
)
};
if rc != 0 {
return None;
}
let symmetric = zps.iter().all(|&z| z == 0);
if info.axis == -1 {
Some(if symmetric {
crate::Quantization::per_tensor_symmetric(scales[0])
} else {
crate::Quantization::per_tensor(scales[0], zps[0])
})
} else {
let axis = info.axis as usize;
let q = if symmetric {
crate::Quantization::per_channel_symmetric(scales, axis)
} else {
crate::Quantization::per_channel(scales, zps, axis)
};
q.ok()
}
}
pub fn set_quantization(&mut self, q: crate::Quantization) -> Result<()> {
let axis = q.axis().map(|a| a as i32).unwrap_or(-1);
let scales = q.scale();
let n = scales.len() as u32;
let zps_owned = q.zero_point().map(|z| z.to_vec());
let zps_ptr = zps_owned
.as_ref()
.map(|z| z.as_ptr())
.unwrap_or(std::ptr::null());
let rc = unsafe {
edgefirst_tensor_ffi::ef_tensor_quantization_set(
self.handle.as_ptr(),
axis,
scales.as_ptr(),
zps_ptr,
n,
)
};
if rc != 0 {
return Err(Error::InvalidArgument(format!(
"ef_tensor_quantization_set failed: {}",
ffi_last_error()
)));
}
self.quantization_cache = std::sync::OnceLock::new();
let _ = self.quantization_cache.set(Some(q));
Ok(())
}
pub fn with_quantization(mut self, q: crate::Quantization) -> Result<Self> {
self.set_quantization(q)?;
Ok(self)
}
pub fn clear_quantization(&mut self) {
unsafe { edgefirst_tensor_ffi::ef_tensor_quantization_clear(self.handle.as_ptr()) };
self.quantization_cache = std::sync::OnceLock::new();
let _ = self.quantization_cache.set(None);
}
pub fn size(&self) -> usize {
self.shape().iter().product::<usize>() * self.dtype().size()
}
pub fn capacity_bytes(&self) -> usize {
if let Some(bytes) = with_pbo!(self, capacity_bytes) {
return bytes;
}
let n = unsafe { edgefirst_tensor_ffi::ef_tensor_capacity_bytes(self.handle.as_ptr()) };
usize::try_from(n).unwrap_or_else(|_| self.size())
}
pub fn row_stride(&self) -> Option<usize> {
let n = unsafe { edgefirst_tensor_ffi::ef_tensor_row_stride(self.handle.as_ptr()) };
usize::try_from(n).ok()
}
pub fn compression(&self) -> Option<crate::CompressionScheme> {
let code = unsafe { edgefirst_tensor_ffi::ef_tensor_compression(self.handle.as_ptr()) };
let scheme = crate::CompressionScheme::from_code(code);
if scheme.is_none() && code != 0 {
log::warn!(
"ef_tensor_compression returned scheme code {code}, which this build of \
edgefirst-tensor does not recognise; reporting a linear layout. The \
loaded libedgefirst_tensor.so is newer than this consumer."
);
}
scheme
}
pub fn sync_for_cpu(&self, access: CpuAccess) -> Result<()> {
self.sync(access, true)
}
pub fn sync_for_device(&self, access: CpuAccess) -> Result<()> {
self.sync(access, false)
}
fn sync(&self, access: CpuAccess, to_cpu: bool) -> Result<()> {
let what = if to_cpu {
"sync_for_cpu"
} else {
"sync_for_device"
};
if self.pbo.is_some() {
return Err(Error::NotImplemented(format!(
"{what}: a PBO has no coherency window independent of its map -- \
glMapBufferRange establishes the address and the visibility \
together, and glUnmapBuffer publishes the writes. Use \
Tensor::as_pbo().map() instead, which owns the pair."
)));
}
let code = match access {
CpuAccess::None => {
return Err(Error::InvalidArgument(format!(
"{what}: CpuAccess::None is not a sync direction"
)))
}
CpuAccess::Read => 1,
CpuAccess::Write => 2,
CpuAccess::ReadWrite => 3,
};
let rc = unsafe {
if to_cpu {
edgefirst_tensor_ffi::ef_tensor_sync_for_cpu(self.handle.as_ptr(), code)
} else {
edgefirst_tensor_ffi::ef_tensor_sync_for_device(self.handle.as_ptr(), code)
}
};
if rc == 0 {
return Ok(());
}
let msg = ffi_last_error();
if rc == libc::ENOTSUP {
return Err(Error::NotImplemented(msg));
}
Err(Error::InvalidOperation(format!(
"{what}: {msg} (errno {rc})"
)))
}
pub fn batch(&self, n: usize) -> Result<TensorDyn> {
let handle =
unsafe { edgefirst_tensor_ffi::ef_tensor_batch(self.handle.as_ptr(), n as u64) };
match NonNull::new(handle) {
Some(h) => Ok(Self::from_handle(h)),
None => {
let is_index = ffi_last_error_class()
== edgefirst_tensor_ffi::EfErrorClass::BatchIndexOutOfBounds;
Err(batch_error(n, self.shape().first().copied(), is_index))
}
}
}
#[cfg(target_os = "linux")]
pub fn dmabuf_clone(&self) -> Result<std::os::fd::OwnedFd> {
if self.memory() != TensorMemory::DmaBuf {
return Err(Error::NotImplemented(format!(
"dmabuf_clone requires DMA-backed tensor, got {:?}",
self.memory()
)));
}
self.clone_fd()
}
#[cfg(any(target_os = "macos", target_os = "ios"))]
pub fn iosurface_ref(&self) -> Option<*mut std::ffi::c_void> {
let p = unsafe { edgefirst_tensor_ffi::ef_tensor_iosurface_ref(self.handle.as_ptr()) };
if p.is_null() {
None
} else {
Some(p)
}
}
#[cfg(any(target_os = "macos", target_os = "ios"))]
pub fn iosurface_physical_dims(&self) -> Option<(usize, usize)> {
let surface = self.iosurface_ref()?;
unsafe { Some((IOSurfaceGetWidth(surface), IOSurfaceGetHeight(surface))) }
}
pub fn cuda(&self) -> Option<&crate::cuda::CudaHandle> {
self.cuda.as_deref()
}
pub fn cuda_map(&self) -> Option<crate::cuda::CudaMap<'_>> {
self.cuda()?.map()
}
pub fn pbo_id(&self) -> Option<u32> {
with_pbo!(self, buffer_id)
}
pub fn pbo_vtable_ptr(&self) -> Option<*const std::ffi::c_void> {
with_pbo!(self, pbo_vtable).map(|v| v as *const _ as *const std::ffi::c_void)
}
pub fn pbo_keepalive(&self) -> Option<std::sync::Arc<dyn Send + Sync>> {
with_pbo!(self, pbo_keepalive)
}
pub unsafe fn from_raw_host_with_capacity(
ptr: *mut u8,
shape: &[usize],
capacity_bytes: usize,
dtype: DType,
_name: Option<&str>,
) -> Result<Self> {
let dims: Vec<u64> = shape.iter().map(|&d| d as u64).collect();
let handle = unsafe {
edgefirst_tensor_ffi::ef_tensor_wrap_host(
ptr,
capacity_bytes,
dtype.code(),
dims.as_ptr(),
dims.len() as u32,
)
};
match NonNull::new(handle) {
Some(h) => Ok(Self::from_handle(h)),
None => Err(Error::InvalidArgument(format!(
"ef_tensor_wrap_host failed: {}",
ffi_last_error()
))),
}
}
#[cfg(any(target_os = "macos", target_os = "ios"))]
pub unsafe fn from_iosurface(
surface_ref: *mut std::ffi::c_void,
shape: &[usize],
dtype: DType,
name: Option<&str>,
) -> Result<Self> {
if surface_ref.is_null() {
return Err(Error::InvalidArgument(
"from_iosurface: surface_ref is null".into(),
));
}
let id = unsafe { IOSurfaceGetID(surface_ref) };
Self::from_iosurface_id(id, shape, dtype, name)
}
#[cfg(any(target_os = "macos", target_os = "ios"))]
pub fn iosurface_id(&self) -> Option<u32> {
let surface = self.iosurface_ref()?;
let id = unsafe { IOSurfaceGetID(surface) };
if id == 0 {
None
} else {
Some(id)
}
}
#[cfg(any(target_os = "macos", target_os = "ios"))]
pub fn from_iosurface_id(
id: u32,
shape: &[usize],
dtype: DType,
_name: Option<&str>,
) -> Result<Self> {
let dims: Vec<u64> = shape.iter().map(|&d| d as u64).collect();
let handle = unsafe {
edgefirst_tensor_ffi::ef_tensor_from_iosurface_id(
id,
dtype.code(),
dims.as_ptr(),
dims.len() as u32,
)
};
match NonNull::new(handle) {
Some(h) => Ok(Self::from_handle(h)),
None => Err(Error::InvalidArgument(format!(
"ef_tensor_from_iosurface_id failed: {}",
ffi_last_error()
))),
}
}
pub(crate) fn from_pbo_import(
buffer_id: u32,
size: usize,
shape: &[usize],
dtype: DType,
ops: std::sync::Arc<dyn crate::PboOps>,
) -> Result<Self> {
macro_rules! arm {
($t:ty) => {
crate::PboTensor::<$t>::from_pbo(buffer_id, size, shape, None, ops)
.and_then(crate::Tensor::<$t>::from_pbo)
.map(crate::Tensor::<$t>::into_inner)
};
}
match dtype {
DType::U8 => arm!(u8),
DType::I8 => arm!(i8),
DType::U16 => arm!(u16),
DType::I16 => arm!(i16),
DType::U32 => arm!(u32),
DType::I32 => arm!(i32),
DType::U64 => arm!(u64),
DType::I64 => arm!(i64),
DType::F16 => arm!(half::f16),
DType::F32 => arm!(f32),
DType::F64 => arm!(f64),
}
}
pub(crate) fn import_storage(
desc: &crate::TensorDesc,
shape: &[usize],
dtype: DType,
) -> Result<Self> {
match desc.kind {
#[cfg(target_os = "linux")]
crate::protocol::kind::DMABUF => {
let fd = crate::protocol::dup_descriptor_fd(desc)?;
Self::from_fd(fd, shape, dtype, None)
}
#[cfg(not(target_os = "linux"))]
crate::protocol::kind::DMABUF => {
Err(Error::NotImplemented("dma-buf import off Linux".into()))
}
#[cfg(any(target_os = "macos", target_os = "ios"))]
crate::protocol::kind::IOSURFACE => {
let id = crate::protocol::descriptor_surface_id(desc)?;
Self::from_iosurface_id(id, shape, dtype, None)
}
#[cfg(not(any(target_os = "macos", target_os = "ios")))]
crate::protocol::kind::IOSURFACE => Err(Error::NotImplemented(
"IOSurface import off Apple platforms".into(),
)),
crate::protocol::kind::HOST => {
crate::protocol::check_descriptor_host_ptr(desc)?;
unsafe {
Self::from_raw_host_with_capacity(
desc.ptr.0,
shape,
desc.capacity as usize,
dtype,
None,
)
}
}
crate::protocol::kind::PBO => {
let buffer_id = crate::protocol::descriptor_pbo_buffer_id(desc)?;
let ops = unsafe { crate::pbo::import_pbo_ops(desc.ptr.0 as *const _)? };
Self::from_pbo_import(buffer_id, desc.capacity as usize, shape, dtype, ops)
}
k => Err(Error::NotImplemented(format!(
"tensor interop kind {k} cannot be imported by this build"
))),
}
}
pub fn descriptor(&self) -> crate::TensorDesc {
self.descriptor_pinned(None)
}
pub fn descriptor_pinned(&self, pin: Option<&crate::HostPin<'_>>) -> crate::TensorDesc {
let memory = self.memory();
let handle: i64 = match memory {
TensorMemory::Pbo => self.pbo_id().map(|id| id as i64).unwrap_or(-1),
_ => self.plane0().map(|p| p.handle).unwrap_or(-1),
};
crate::protocol::from_parts(crate::protocol::DescParts {
dims: self.shape(),
memory,
dtype: self.dtype(),
fourcc: self.format().map(|f| f.to_fourcc()).unwrap_or(0),
format: self.format(),
row_stride: self.row_stride(),
handle,
colorimetry: self.colorimetry().map(|c| c.pack()).unwrap_or(0),
capacity: self.capacity_bytes() as u64,
pin,
pbo_vtable_ptr: self.pbo_vtable_ptr(),
})
}
}
macro_rules! with_pbo_mut {
($self:expr, $method:ident, $arg:expr) => {{
match $self.pbo.as_mut() {
None => None,
Some(any) => {
macro_rules! arm {
($t:ty) => {
any.downcast_mut::<crate::PboTensor<$t>>().map(|p| {
<crate::PboTensor<$t> as crate::TensorTrait<$t>>::$method(p, $arg)
})
};
}
None.or_else(|| arm!(u8))
.or_else(|| arm!(i8))
.or_else(|| arm!(u16))
.or_else(|| arm!(i16))
.or_else(|| arm!(half::f16))
.or_else(|| arm!(u32))
.or_else(|| arm!(i32))
.or_else(|| arm!(f32))
.or_else(|| arm!(u64))
.or_else(|| arm!(i64))
.or_else(|| arm!(f64))
}
}
}};
}
use with_pbo_mut;
macro_rules! with_pbo {
($self:expr, $method:ident) => {{
match $self.pbo.as_ref() {
None => None,
Some(any) => {
let p = None
.or_else(|| {
any.downcast_ref::<crate::PboTensor<u8>>()
.map(|p| p.$method())
})
.or_else(|| {
any.downcast_ref::<crate::PboTensor<i8>>()
.map(|p| p.$method())
})
.or_else(|| {
any.downcast_ref::<crate::PboTensor<u16>>()
.map(|p| p.$method())
})
.or_else(|| {
any.downcast_ref::<crate::PboTensor<i16>>()
.map(|p| p.$method())
})
.or_else(|| {
any.downcast_ref::<crate::PboTensor<half::f16>>()
.map(|p| p.$method())
})
.or_else(|| {
any.downcast_ref::<crate::PboTensor<u32>>()
.map(|p| p.$method())
})
.or_else(|| {
any.downcast_ref::<crate::PboTensor<i32>>()
.map(|p| p.$method())
})
.or_else(|| {
any.downcast_ref::<crate::PboTensor<f32>>()
.map(|p| p.$method())
})
.or_else(|| {
any.downcast_ref::<crate::PboTensor<u64>>()
.map(|p| p.$method())
})
.or_else(|| {
any.downcast_ref::<crate::PboTensor<i64>>()
.map(|p| p.$method())
})
.or_else(|| {
any.downcast_ref::<crate::PboTensor<f64>>()
.map(|p| p.$method())
});
p
}
}
}};
}
use with_pbo;
fn batch_error(n: usize, leading: Option<usize>, class_is_index: bool) -> Error {
match leading {
Some(batch) if class_is_index => Error::BatchIndexOutOfBounds { index: n, batch },
_ => ffi_error(Error::InvalidShape),
}
}
fn builder_error(errno: std::ffi::c_int) -> Error {
let msg = ffi_last_error();
match errno {
libc::ENOTSUP => Error::NotImplemented(msg),
libc::ENOMEM => Error::IoError(std::io::Error::from_raw_os_error(libc::ENOMEM)),
libc::ERANGE => Error::ShapeMismatch(msg),
libc::EACCES => Error::InvalidOperation(msg),
libc::EIO => Error::IoError(std::io::Error::other(msg)),
libc::EINVAL => Error::InvalidArgument(msg),
other => Error::InvalidArgument(format!("{msg} (errno {other})")),
}
}
#[cfg(unix)]
#[allow(clippy::unnecessary_cast)]
fn identity_from_stat(stat: &nix::sys::stat::FileStat) -> BufferIdentity {
let key = ((stat.st_dev as u64) << 32) ^ (stat.st_ino as u64);
BufferIdentity::derived(IdentityKind::DmaBuf, key)
}
fn format_cstring(format: PixelFormat) -> Result<std::ffi::CString> {
std::ffi::CString::new(format.as_str())
.map_err(|e| Error::InvalidArgument(format!("pixel format string contains a NUL: {e}")))
}
fn memory_code(memory: Option<TensorMemory>) -> (std::ffi::c_int, u32) {
match memory {
Some(m) => (1, m.code()),
None => (0, 0),
}
}
fn access_code(access: CpuAccess) -> u32 {
match access {
CpuAccess::None => 0,
CpuAccess::Read => 1,
CpuAccess::Write => 2,
CpuAccess::ReadWrite => 3,
}
}
pub(crate) fn ffi_last_error() -> String {
unsafe {
let ptr = edgefirst_tensor_ffi::ef_tensor_last_error_message();
if ptr.is_null() {
return String::new();
}
std::ffi::CStr::from_ptr(ptr).to_string_lossy().into_owned()
}
}
fn ffi_last_error_class() -> edgefirst_tensor_ffi::EfErrorClass {
use edgefirst_tensor_ffi::EfErrorClass as C;
let code = unsafe { edgefirst_tensor_ffi::ef_tensor_last_error_class() };
match code {
1 => C::InvalidArgument,
2 => C::InvalidShape,
3 => C::InsufficientCapacity,
4 => C::BatchIndexOutOfBounds,
5 => C::RegionOutOfBounds,
6 => C::NotSupported,
7 => C::InvalidOperation,
8 => C::AllocationFailed,
9 => C::QuantizationInvalid,
_ => C::Unspecified,
}
}
pub(crate) fn ffi_error(fallback: fn(String) -> Error) -> Error {
use edgefirst_tensor_ffi::EfErrorClass as C;
let msg = ffi_last_error();
match ffi_last_error_class() {
C::InvalidArgument => Error::InvalidArgument(msg),
C::InvalidShape => Error::InvalidShape(msg),
C::InsufficientCapacity | C::BatchIndexOutOfBounds | C::RegionOutOfBounds => {
Error::ShapeMismatch(msg)
}
C::NotSupported => Error::NotImplemented(msg),
C::InvalidOperation => Error::InvalidOperation(msg),
C::AllocationFailed => Error::IoError(std::io::Error::other(msg)),
C::QuantizationInvalid => Error::InvalidArgument(msg),
C::Unspecified => fallback(msg),
}
}
struct MapKeepalive(*mut EfTensor);
unsafe impl Send for MapKeepalive {}
unsafe impl Sync for MapKeepalive {}
impl Drop for MapKeepalive {
fn drop(&mut self) {
unsafe {
edgefirst_tensor_ffi::ef_tensor_unmap(self.0);
edgefirst_tensor_ffi::ef_tensor_free(self.0);
}
}
}
impl fmt::Debug for TensorDyn {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.debug_struct("TensorDyn")
.field("dtype", &self.dtype())
.field("shape", &self.shape())
.finish()
}
}
#[cfg(any(target_os = "macos", target_os = "ios"))]
#[link(name = "IOSurface", kind = "framework")]
extern "C" {
fn IOSurfaceGetWidth(surface: *mut std::ffi::c_void) -> usize;
fn IOSurfaceGetHeight(surface: *mut std::ffi::c_void) -> usize;
fn IOSurfaceGetID(surface: *mut std::ffi::c_void) -> u32;
}