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, ViewOrigin,
};
pub struct TensorDyn {
handle: NonNull<EfTensor>,
shape_cache: Vec<usize>,
identity: BufferIdentity,
quantization_cache: std::sync::OnceLock<Option<crate::Quantization>>,
chroma_cache: std::sync::OnceLock<Option<Box<TensorDyn>>>,
#[cfg(target_os = "windows")]
descriptor_texture_handle: std::sync::OnceLock<Option<std::os::windows::io::OwnedHandle>>,
}
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(),
chroma_cache: std::sync::OnceLock::new(),
#[cfg(target_os = "windows")]
descriptor_texture_handle: std::sync::OnceLock::new(),
}
}
fn derive_identity(handle: NonNull<EfTensor>) -> BufferIdentity {
#[cfg(target_os = "windows")]
{
let code = unsafe { edgefirst_tensor_ffi::ef_tensor_storage_kind(handle.as_ptr()) };
if TensorMemory::from_code(code) == Some(TensorMemory::DmaBuf) {
let tex = unsafe { edgefirst_tensor_ffi::ef_tensor_d3d11_texture(handle.as_ptr()) };
if !tex.is_null() {
return BufferIdentity::derived(
IdentityKind::D3d11Texture,
tex as usize as u64,
);
}
}
}
#[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)
}
pub fn is_multiplane(&self) -> bool {
unsafe { edgefirst_tensor_ffi::ef_tensor_is_multiplane(self.handle.as_ptr()) != 0 }
}
pub fn chroma(&self) -> Option<&TensorDyn> {
self.chroma_cache
.get_or_init(|| {
let p = unsafe { edgefirst_tensor_ffi::ef_tensor_chroma(self.handle.as_ptr()) };
std::ptr::NonNull::new(p).map(|h| Box::new(unsafe { Self::from_raw(h.as_ptr()) }))
})
.as_deref()
}
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);
Ok(())
}
pub fn name(&self) -> String {
String::new()
}
pub fn memory(&self) -> TensorMemory {
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>> {
self.map_pin_with(access, false)
}
pub(crate) fn map_pin_with(
&self,
access: CpuAccess,
non_blocking: bool,
) -> Result<crate::pin::HostPin<'static>> {
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 {
if non_blocking {
edgefirst_tensor_ffi::ef_tensor_try_map(self.handle.as_ptr(), code, &mut view)
} else {
edgefirst_tensor_ffi::ef_tensor_map(self.handle.as_ptr(), code, &mut view)
}
};
if rc == libc::EAGAIN && non_blocking {
return Err(Error::IoError(std::io::Error::new(
std::io::ErrorKind::WouldBlock,
ffi_last_error(),
)));
}
if rc != 0 {
let what = if non_blocking {
"ef_tensor_try_map"
} else {
"ef_tensor_map"
};
return Err(Error::InvalidOperation(format!(
"{what} 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))
}
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 try_map_bytes(&self, access: CpuAccess) -> Result<crate::view::HostView<'static, u8>> {
let pin = self.map_pin_with(access, true)?;
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>> {
if self.memory() == TensorMemory::Pbo {
return Err(Error::NotImplemented(
crate::pbo::PIN_HOST_PBO_REFUSAL.to_owned(),
));
}
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, dims) = self.image_dims()?;
match fmt.layout() {
PixelLayout::Packed => dims.get(1).copied(),
PixelLayout::Planar => dims.get(2).copied(),
PixelLayout::SemiPlanar => dims.get(1).copied(),
}
}
fn image_dims(&self) -> Option<(PixelFormat, &[usize])> {
let fmt = self.format()?;
let shape = self.shape();
let rank = match fmt.layout() {
PixelLayout::SemiPlanar => 2,
PixelLayout::Packed | PixelLayout::Planar => 3,
};
let batch_dims = shape.len().checked_sub(rank).filter(|&n| n <= 1)?;
Some((fmt, &shape[batch_dims..]))
}
pub fn height(&self) -> Option<usize> {
let (fmt, dims) = self.image_dims()?;
match fmt.layout() {
PixelLayout::Packed => dims.first().copied(),
PixelLayout::Planar => dims.get(1).copied(),
PixelLayout::SemiPlanar => {
if self.is_multiplane() {
dims.first().copied()
} else {
match fmt {
PixelFormat::Nv12 => dims.first().map(|h| h * 2 / 3),
PixelFormat::Nv16 => dims.first().map(|h| h / 2),
PixelFormat::Nv24 => dims.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);
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 {
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"
};
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))) }
}
#[cfg(target_os = "windows")]
pub fn d3d11_texture(&self) -> Option<*mut std::ffi::c_void> {
let p = unsafe { edgefirst_tensor_ffi::ef_tensor_d3d11_texture(self.handle.as_ptr()) };
(!p.is_null()).then_some(p)
}
#[cfg(target_os = "windows")]
pub fn d3d11_layout(&self) -> Option<crate::d3d11_layout::D3d11ImageLayout> {
let mut out = edgefirst_tensor_ffi::EfD3d11Layout::default();
let rc =
unsafe { edgefirst_tensor_ffi::ef_tensor_d3d11_layout(self.handle.as_ptr(), &mut out) };
(rc == 0).then_some(crate::d3d11_layout::D3d11ImageLayout {
dxgi_format: out.dxgi_format,
texture_width: out.texture_width as usize,
texture_height: out.texture_height as usize,
bytes_per_texel: out.bytes_per_texel as usize,
gl_internal_format: out.gl_internal_format,
})
}
#[cfg(target_os = "windows")]
pub fn d3d11_shared_handle(&self) -> Result<std::os::windows::io::OwnedHandle> {
use std::os::windows::io::FromRawHandle;
let h =
unsafe { edgefirst_tensor_ffi::ef_tensor_d3d11_shared_handle(self.handle.as_ptr()) };
if h.is_null() {
return Err(ffi_error(Error::NotImplemented));
}
Ok(unsafe { std::os::windows::io::OwnedHandle::from_raw_handle(h) })
}
#[cfg(target_os = "windows")]
pub fn gpu_completion(&self) -> Result<Option<crate::d3d11::GpuCompletion>> {
use std::os::windows::io::FromRawHandle;
let mut fence = std::ptr::null_mut();
let mut value = 0u64;
let rc = unsafe {
edgefirst_tensor_ffi::ef_tensor_gpu_completion(
self.handle.as_ptr(),
&mut fence,
&mut value,
)
};
if rc != 0 {
return Err(ffi_error(Error::NotImplemented));
}
if fence.is_null() {
return Ok(None);
}
let fence = unsafe { std::os::windows::io::OwnedHandle::from_raw_handle(fence) };
Ok(Some(crate::d3d11::GpuCompletion { fence, value }))
}
pub fn gpu_write_value(&self) -> u64 {
if self.memory() != TensorMemory::DmaBuf {
return 0;
}
unsafe { edgefirst_tensor_ffi::ef_tensor_gpu_write_value(self.handle.as_ptr()) }
}
#[cfg(target_os = "windows")]
pub fn set_gpu_write(&self, value: u64) -> Result<()> {
let rc =
unsafe { edgefirst_tensor_ffi::ef_tensor_set_gpu_write(self.handle.as_ptr(), value) };
if rc != 0 {
return Err(ffi_error(Error::NotImplemented));
}
Ok(())
}
#[deprecated(
since = "0.32.0",
note = "On the dynamic backend use is_cuda_attached() or cuda_map() — the CudaHandle is not borrowable through this lens."
)]
pub fn cuda(&self) -> Option<&crate::cuda::CudaHandle> {
None
}
pub fn is_cuda_attached(&self) -> bool {
unsafe { edgefirst_tensor_ffi::ef_tensor_cuda_attached(self.handle.as_ptr()) != 0 }
}
pub fn set_cuda_handle(&mut self, h: crate::cuda::CudaHandle) {
match h.into_client_parts() {
Ok(parts) => {
let rc = unsafe {
edgefirst_tensor_ffi::ef_tensor_cuda_attach(
self.handle.as_ptr(),
parts.state,
parts.resource,
parts.size,
Some(parts.map_fn),
Some(parts.unmap_fn),
Some(parts.unregister_fn),
)
};
if rc != 0 {
log::debug!("ef_tensor_cuda_attach failed: {rc}");
}
}
Err(e) => log::debug!("set_cuda_handle: {e}"),
}
}
fn ffi_cuda_map(
&self,
mk: unsafe extern "C" fn(*const EfTensor) -> *mut std::ffi::c_void,
) -> Option<crate::cuda::CudaMap<'_>> {
let map = unsafe { mk(self.handle.as_ptr()) };
if map.is_null() {
return None;
}
unsafe { crate::cuda::CudaMap::from_ffi(map) }
}
pub fn cuda_map(&self) -> Option<crate::cuda::CudaMap<'_>> {
self.ffi_cuda_map(edgefirst_tensor_ffi::ef_tensor_cuda_map)
}
pub fn cuda_map_mut(&self) -> Option<crate::cuda::CudaMap<'_>> {
self.ffi_cuda_map(edgefirst_tensor_ffi::ef_tensor_cuda_map_mut)
}
pub fn pbo_id(&self) -> Option<u32> {
let mut id: u32 = 0;
let rc = unsafe { edgefirst_tensor_ffi::ef_tensor_pbo_id(self.handle.as_ptr(), &mut id) };
(rc == 0).then_some(id)
}
pub fn pbo_is_mapped(&self) -> Option<bool> {
let rc = unsafe { edgefirst_tensor_ffi::ef_tensor_pbo_is_mapped(self.handle.as_ptr()) };
match rc {
0 => Some(false),
1 => Some(true),
_ => None,
}
}
pub fn pbo_vtable_ptr(&self) -> Option<*const std::ffi::c_void> {
let p = unsafe { edgefirst_tensor_ffi::ef_tensor_pbo_vtable(self.handle.as_ptr()) };
(!p.is_null()).then_some(p)
}
pub fn pbo_keepalive(&self) -> Option<std::sync::Arc<dyn Send + Sync>> {
self.pbo_id()?;
let rc = unsafe { edgefirst_tensor_ffi::ef_tensor_retain(self.handle.as_ptr()) };
if rc != 0 {
return None;
}
Some(std::sync::Arc::new(HandleKeepalive(self.handle.as_ptr())))
}
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()
))),
}
}
#[cfg(target_os = "windows")]
#[allow(clippy::too_many_arguments)] pub unsafe fn from_d3d11_texture(
texture: *mut std::ffi::c_void,
width: usize,
height: usize,
format: PixelFormat,
dtype: DType,
access: CpuAccess,
name: Option<&str>,
) -> Result<Self> {
let dims = allocation_dims(format, width, height, "from_d3d11_texture")?;
let format = format_cstring(format)?;
let name = name_cstring(name)?;
let handle = unsafe {
edgefirst_tensor_ffi::ef_tensor_from_d3d11_texture(
texture,
dtype.code(),
dims.as_ptr(),
dims.len() as u32,
format.as_ptr(),
access_code(access),
name.as_ref().map_or(std::ptr::null(), |n| n.as_ptr()),
)
};
match NonNull::new(handle) {
Some(h) => Ok(Self::from_handle(h)),
None => Err(ffi_error(Error::InvalidArgument)),
}
}
#[cfg(target_os = "windows")]
#[allow(clippy::too_many_arguments)] pub unsafe fn from_d3d11_shared_handle(
handle: std::os::windows::io::RawHandle,
width: usize,
height: usize,
format: PixelFormat,
dtype: DType,
access: CpuAccess,
completion: Option<(std::os::windows::io::RawHandle, u64)>,
name: Option<&str>,
) -> Result<Self> {
let dims = allocation_dims(format, width, height, "from_d3d11_shared_handle")?;
let format = format_cstring(format)?;
let name = name_cstring(name)?;
let (fence, fence_value) = completion.unwrap_or((std::ptr::null_mut(), 0));
let tensor = unsafe {
edgefirst_tensor_ffi::ef_tensor_from_d3d11_shared_handle(
handle,
dtype.code(),
dims.as_ptr(),
dims.len() as u32,
format.as_ptr(),
access_code(access),
fence,
fence_value,
name.as_ref().map_or(std::ptr::null(), |n| n.as_ptr()),
)
};
match NonNull::new(tensor) {
Some(h) => Ok(Self::from_handle(h)),
None => Err(ffi_error(Error::InvalidArgument)),
}
}
pub unsafe fn from_client_pbo(
state: crate::EfClientState,
map_fn: crate::EfPboMapFn,
unmap_fn: crate::EfPboUnmapFn,
buffer_id: u32,
size: usize,
shape: &[usize],
dtype: DType,
) -> Result<TensorDyn> {
let dims: Vec<u64> = shape.iter().map(|d| *d as u64).collect();
let handle = unsafe {
edgefirst_tensor_ffi::ef_tensor_wrap_pbo(
state,
buffer_id,
size,
dtype.code(),
dims.as_ptr(),
dims.len() as u32,
Some(map_fn),
Some(unmap_fn),
)
};
match NonNull::new(handle) {
Some(h) => Ok(Self::from_handle(h)),
None => Err(ffi_error(Error::InvalidArgument)),
}
}
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 parts = unsafe { crate::read_pbo_vtable_parts(desc.ptr.0 as *const _) }?;
unsafe {
Self::from_client_pbo(
parts.state,
parts.map_fn,
parts.unmap_fn,
buffer_id,
desc.capacity as usize,
shape,
dtype,
)
}
}
#[cfg(target_os = "windows")]
crate::protocol::kind::D3D11_TEXTURE => {
let (tex, completion) = crate::protocol::descriptor_d3d11_handles(desc)?;
unsafe {
let (format, width, height) =
crate::protocol::descriptor_d3d11_geometry(desc, tex, shape)?;
Self::from_d3d11_shared_handle(
tex,
width,
height,
format,
dtype,
CpuAccess::ReadWrite,
completion,
None,
)
}
}
#[cfg(not(target_os = "windows"))]
crate::protocol::kind::D3D11_TEXTURE => Err(Error::NotImplemented(
"D3D11 texture import off Windows".into(),
)),
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),
#[cfg(target_os = "windows")]
TensorMemory::DmaBuf => self.descriptor_texture_handle_value(),
_ => self.plane0().map(|p| p.handle).unwrap_or(-1),
};
#[cfg(target_os = "windows")]
let (fence_handle, sync) = match memory {
TensorMemory::DmaBuf => (
crate::d3d11::device()
.map(|d| d.exported_fence_handle_value())
.unwrap_or(0),
Some(self.gpu_write_value()).filter(|v| *v != 0),
),
_ => (0, None),
};
#[cfg(not(target_os = "windows"))]
let (fence_handle, sync) = (0, None);
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(),
fence_handle,
sync,
})
}
#[cfg(target_os = "windows")]
fn descriptor_texture_handle_value(&self) -> i64 {
use std::os::windows::io::AsRawHandle;
self.descriptor_texture_handle
.get_or_init(|| self.d3d11_shared_handle().ok())
.as_ref()
.map(|h| h.as_raw_handle() as usize as i64)
.unwrap_or(-1)
}
}
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}")))
}
#[cfg(target_os = "windows")]
fn name_cstring(name: Option<&str>) -> Result<Option<std::ffi::CString>> {
name.map(std::ffi::CString::new)
.transpose()
.map_err(|e| Error::InvalidArgument(format!("name contains a NUL: {e}")))
}
#[cfg(target_os = "windows")]
fn allocation_dims(
format: PixelFormat,
width: usize,
height: usize,
what: &str,
) -> Result<Vec<u64>> {
format
.allocation_shape(width, height)
.map(|s| s.iter().map(|&d| d as u64).collect())
.ok_or_else(|| {
Error::InvalidArgument(format!(
"{what}: invalid dimensions {width}x{height} for format {format:?}"
))
})
}
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 HandleKeepalive(*mut EfTensor);
unsafe impl Send for HandleKeepalive {}
unsafe impl Sync for HandleKeepalive {}
impl Drop for HandleKeepalive {
fn drop(&mut self) {
unsafe { edgefirst_tensor_ffi::ef_tensor_free(self.0) }
}
}
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;
}