pub const ABI_VERSION: u32 = 1;
pub use crate::dtype_wire as dtype;
pub(crate) fn dtype_of(d: crate::DType) -> u32 {
d.code()
}
pub fn dtype_to_dtype(code: u32) -> Option<crate::DType> {
crate::DType::from_code(code)
}
pub mod format {
pub const NONE: u32 = 0;
pub use crate::format::pixel_format_wire::*;
}
pub(crate) fn format_of(f: crate::PixelFormat) -> u32 {
f.code()
}
pub fn format_from_code(code: u32) -> Option<crate::PixelFormat> {
crate::PixelFormat::from_code(code)
}
pub mod flags {
pub const SYNC_PRESENT: u32 = 1 << 0;
}
pub mod kind {
pub const HOST: u32 = 0;
pub const DMABUF: u32 = 1;
pub const IOSURFACE: u32 = 2;
pub const PBO: u32 = 3;
pub const CUDA_DEVICE: u32 = 4;
pub const D3D11_TEXTURE: u32 = 5;
}
#[repr(transparent)]
#[derive(Clone, Copy, Debug)]
pub struct SendPtr(pub *mut u8);
unsafe impl Send for SendPtr {}
unsafe impl Sync for SendPtr {}
impl SendPtr {
pub const fn null() -> Self {
SendPtr(std::ptr::null_mut())
}
pub fn is_null(&self) -> bool {
self.0.is_null()
}
}
#[repr(C)]
#[derive(Clone, Copy, Debug)]
pub struct TensorDesc {
pub version: u32,
pub kind: u32,
pub handle: i64,
pub ptr: SendPtr,
pub ndim: u32,
pub dtype: u32,
pub shape: [u64; 8],
pub strides: [i64; 8],
pub fourcc: u32,
pub format: u32,
pub colorimetry: u32,
pub flags: u32,
pub capacity: u64,
pub sync: u64,
}
impl TensorDesc {
pub const MAX_NDIM: usize = 8;
pub fn shape(&self) -> &[u64] {
&self.shape[..(self.ndim as usize).min(Self::MAX_NDIM)]
}
pub fn strides(&self) -> &[i64] {
&self.strides[..(self.ndim as usize).min(Self::MAX_NDIM)]
}
pub fn len(&self) -> u64 {
self.shape().iter().product()
}
pub fn is_empty(&self) -> bool {
self.len() == 0
}
}
use crate::TensorMemory;
pub(crate) fn kind_of(memory: TensorMemory) -> u32 {
match memory {
TensorMemory::Mem | TensorMemory::Shm => kind::HOST,
TensorMemory::Pbo => kind::PBO,
TensorMemory::Cuda => kind::CUDA_DEVICE,
TensorMemory::IoSurface => kind::IOSURFACE,
TensorMemory::DmaBuf => {
#[cfg(any(target_os = "macos", target_os = "ios"))]
{
kind::IOSURFACE
}
#[cfg(target_os = "windows")]
{
kind::D3D11_TEXTURE
}
#[cfg(not(any(target_os = "macos", target_os = "ios", target_os = "windows")))]
{
kind::DMABUF
}
}
}
}
#[cfg(target_os = "linux")]
pub(crate) fn dup_descriptor_fd(desc: &TensorDesc) -> crate::Result<std::os::fd::OwnedFd> {
use std::os::fd::BorrowedFd;
if desc.handle < 0 {
return Err(crate::Error::InvalidArgument(
"dma-buf descriptor carries no fd".into(),
));
}
let raw_fd = i32::try_from(desc.handle).map_err(|_| {
crate::Error::InvalidArgument(format!(
"dma-buf descriptor fd {} is not representable as a file descriptor",
desc.handle
))
})?;
let borrowed = unsafe { BorrowedFd::borrow_raw(raw_fd) };
Ok(nix::unistd::dup(borrowed)?)
}
#[cfg(any(target_os = "macos", target_os = "ios"))]
pub(crate) fn descriptor_surface_id(desc: &TensorDesc) -> crate::Result<u32> {
if desc.handle < 0 {
return Err(crate::Error::InvalidArgument(
"IOSurface descriptor carries no surface id".into(),
));
}
u32::try_from(desc.handle).map_err(|_| {
crate::Error::InvalidArgument(format!(
"IOSurface descriptor id {} is out of range",
desc.handle
))
})
}
pub(crate) fn check_descriptor_host_ptr(desc: &TensorDesc) -> crate::Result<()> {
if desc.ptr.is_null() {
return Err(crate::Error::InvalidArgument(
"host descriptor has no address: the producer was asked \
for a descriptive-only capsule (request access=\"read\" \
or \"readwrite\")"
.into(),
));
}
Ok(())
}
pub(crate) fn descriptor_pbo_buffer_id(desc: &TensorDesc) -> crate::Result<u32> {
if desc.handle < 0 {
return Err(crate::Error::InvalidArgument(
"PBO descriptor carries no buffer id".into(),
));
}
let buffer_id = u32::try_from(desc.handle).map_err(|_| {
crate::Error::InvalidArgument(format!(
"PBO descriptor buffer id {} is out of range",
desc.handle
))
})?;
if desc.ptr.is_null() {
return Err(crate::Error::InvalidArgument(
"PBO descriptor carries no ops vtable -- the producer either \
predates this build's cross-cdylib PBO support or its own \
PboTensor could not build one"
.into(),
));
}
Ok(buffer_id)
}
#[cfg(target_os = "windows")]
pub(crate) fn descriptor_d3d11_handles(
desc: &TensorDesc,
) -> crate::Result<(
std::os::windows::io::RawHandle,
Option<(std::os::windows::io::RawHandle, u64)>,
)> {
if desc.handle <= 0 {
return Err(crate::Error::InvalidArgument(
"D3D11_TEXTURE descriptor without a texture handle".into(),
));
}
let tex = desc.handle as usize as std::os::windows::io::RawHandle;
let completion = (desc.flags & flags::SYNC_PRESENT != 0 && !desc.ptr.is_null())
.then_some((desc.ptr.0 as std::os::windows::io::RawHandle, desc.sync));
Ok((tex, completion))
}
#[cfg(target_os = "windows")]
pub(crate) unsafe fn descriptor_d3d11_geometry(
desc: &TensorDesc,
texture: std::os::windows::io::RawHandle,
shape: &[usize],
) -> crate::Result<(crate::PixelFormat, usize, usize)> {
let format = crate::PixelFormat::from_code(desc.format).ok_or_else(|| {
crate::Error::InvalidArgument(
"D3D11_TEXTURE descriptor without a pixel format: the texture's own \
description cannot say which format it was allocated for"
.into(),
)
})?;
let (width, height) =
unsafe { d3d11_geometry_with_rule(format, texture, shape, ShapeRule::OrWindow) }?;
Ok((format, width, height))
}
#[cfg(target_os = "windows")]
#[derive(Clone, Copy, PartialEq, Eq)]
enum ShapeRule {
#[cfg_attr(feature = "dynamic", allow(dead_code))]
Exact,
OrWindow,
}
#[cfg(target_os = "windows")]
fn is_packed_window_of(
format: crate::PixelFormat,
shape: &[usize],
width: usize,
height: usize,
) -> bool {
if format.layout() != crate::PixelLayout::Packed {
return false;
}
let [h, w, c] = shape else {
return false;
};
*c == format.channels() && (1..=height).contains(h) && (1..=width).contains(w)
}
#[cfg(target_os = "windows")]
#[cfg_attr(feature = "dynamic", allow(dead_code))]
pub(crate) unsafe fn d3d11_geometry_checked(
format: crate::PixelFormat,
texture: std::os::windows::io::RawHandle,
shape: &[usize],
) -> crate::Result<(usize, usize)> {
unsafe { d3d11_geometry_with_rule(format, texture, shape, ShapeRule::Exact) }
}
#[cfg(target_os = "windows")]
unsafe fn d3d11_geometry_with_rule(
format: crate::PixelFormat,
texture: std::os::windows::io::RawHandle,
shape: &[usize],
rule: ShapeRule,
) -> crate::Result<(usize, usize)> {
let (width, height) =
unsafe { crate::d3d11_shared_handle_geometry(texture, format, Some(shape)) }?;
let allocation = format.allocation_shape(width, height);
let addressing = format.addressing_shape(width, height);
if allocation.as_deref() == Some(shape) || addressing.as_deref() == Some(shape) {
return Ok((width, height));
}
if rule == ShapeRule::OrWindow && is_packed_window_of(format, shape, width, height) {
return Ok((width, height));
}
let window = match rule {
ShapeRule::OrWindow => " nor a packed window",
ShapeRule::Exact => "",
};
Err(crate::Error::InvalidArgument(format!(
"D3D11 texture shape {shape:?} is neither the {format:?} allocation \
shape {allocation:?} nor the addressing shape {addressing:?}{window} of the \
{width}x{height} texture it names"
)))
}
pub(crate) struct DescParts<'a, 'p> {
pub dims: &'a [usize],
pub memory: TensorMemory,
pub dtype: crate::DType,
pub fourcc: u32,
pub format: Option<crate::PixelFormat>,
pub row_stride: Option<usize>,
pub handle: i64,
pub colorimetry: u32,
pub capacity: u64,
pub pin: Option<&'a crate::HostPin<'p>>,
pub pbo_vtable_ptr: Option<*const std::ffi::c_void>,
pub fence_handle: usize,
pub sync: Option<u64>,
}
pub(crate) fn from_parts(parts: DescParts) -> TensorDesc {
let DescParts {
dims,
memory,
dtype,
fourcc,
format,
row_stride,
handle,
colorimetry,
capacity,
pin,
pbo_vtable_ptr,
fence_handle,
sync,
} = parts;
let desc_kind = kind_of(memory);
let completion = (desc_kind == kind::D3D11_TEXTURE && fence_handle != 0)
.then_some(sync)
.flatten()
.map(|value| (fence_handle, value));
let ndim = dims.len();
let filled = ndim.min(TensorDesc::MAX_NDIM);
if ndim > TensorDesc::MAX_NDIM {
log::warn!(
"tensor of rank {ndim} exceeds the descriptor's {} shape slots; \
the descriptor will be rejected on import",
TensorDesc::MAX_NDIM
);
}
let mut shape = [0u64; TensorDesc::MAX_NDIM];
let mut strides = [0i64; TensorDesc::MAX_NDIM];
for (i, d) in dims.iter().take(filled).enumerate() {
shape[i] = *d as u64;
}
let mut acc: i64 = dtype.size() as i64;
for i in (0..filled).rev() {
strides[i] = acc;
acc *= shape[i] as i64;
}
if ndim >= 2 {
if let Some(rs) = row_stride {
let row_dim = match format.map(|f| f.layout()) {
Some(crate::PixelLayout::Planar) if ndim >= 3 => 1,
_ => 0,
};
strides[row_dim] = rs as i64;
}
}
TensorDesc {
version: ABI_VERSION,
kind: desc_kind,
handle,
ptr: if desc_kind == kind::D3D11_TEXTURE {
completion.map_or_else(SendPtr::null, |(h, _)| SendPtr(h as *mut u8))
} else {
pin.map(|p| SendPtr(p.as_mut_ptr()))
.or_else(|| pbo_vtable_ptr.map(|p| SendPtr(p as *mut u8)))
.unwrap_or_else(SendPtr::null)
},
ndim: ndim as u32,
dtype: dtype_of(dtype),
shape,
strides,
fourcc,
format: format.map(format_of).unwrap_or(format::NONE),
colorimetry,
flags: if completion.is_some() {
flags::SYNC_PRESENT
} else {
0
},
capacity,
sync: completion.map_or(0, |(_, value)| value),
}
}
#[repr(transparent)]
#[derive(Clone, Copy, Debug)]
pub struct SendConstPtr<T>(pub *const T);
unsafe impl<T> Send for SendConstPtr<T> {}
unsafe impl<T> Sync for SendConstPtr<T> {}
impl<T> SendConstPtr<T> {
pub const fn null() -> Self {
SendConstPtr(std::ptr::null())
}
}
#[repr(C)]
#[derive(Clone, Copy, Debug)]
pub struct QuantDesc {
pub len: u64,
pub axis_plus_one: u64,
pub scales: SendConstPtr<f32>,
pub zero_points: SendConstPtr<i32>,
}
impl QuantDesc {
pub const fn absent() -> Self {
Self {
len: 0,
axis_plus_one: 0,
scales: SendConstPtr::null(),
zero_points: SendConstPtr::null(),
}
}
pub fn borrowing(quant: &crate::Quantization) -> Self {
Self {
len: quant.scale().len() as u64,
axis_plus_one: quant.axis().map_or(0, |a| a as u64 + 1),
scales: SendConstPtr(quant.scale().as_ptr()),
zero_points: quant
.zero_point()
.map_or(SendConstPtr::null(), |zp| SendConstPtr(zp.as_ptr())),
}
}
pub unsafe fn to_quantization(&self) -> crate::Result<Option<crate::Quantization>> {
if self.len == 0 {
return Ok(None);
}
if self.scales.0.is_null() {
return Err(crate::Error::InvalidArgument(format!(
"malformed tensor capsule: quantization declares {} scale(s) but \
carries a null scale array",
self.len
)));
}
let len = self.len as usize;
let scales = unsafe { std::slice::from_raw_parts(self.scales.0, len) }.to_vec();
let zero_points = if self.zero_points.0.is_null() {
None
} else {
Some(unsafe { std::slice::from_raw_parts(self.zero_points.0, len) }.to_vec())
};
let axis = self.axis_plus_one.checked_sub(1).map(|a| a as usize);
let quant = match (axis, zero_points) {
(None, None) if len == 1 => Ok(crate::Quantization::per_tensor_symmetric(scales[0])),
(None, Some(zp)) if len == 1 => Ok(crate::Quantization::per_tensor(scales[0], zp[0])),
(Some(a), None) => crate::Quantization::per_channel_symmetric(scales, a),
(Some(a), Some(zp)) => crate::Quantization::per_channel(scales, zp, a),
(None, _) => {
return Err(crate::Error::InvalidArgument(format!(
"malformed tensor capsule: {len} quantization scales but no \
channel axis; per-channel quantization must name one"
)));
}
};
quant.map(Some).map_err(|e| {
crate::Error::InvalidArgument(format!(
"malformed tensor capsule: quantization metadata is not valid: {e}"
))
})
}
}
#[cfg(test)]
mod quant_desc_tests {
use super::QuantDesc;
use crate::Quantization;
fn round_trip(q: &Quantization) -> Option<Quantization> {
let desc = QuantDesc::borrowing(q);
unsafe { desc.to_quantization() }.expect("a well-formed descriptor must import")
}
#[test]
fn absent_descriptor_is_none_not_an_error() {
let imported = unsafe { QuantDesc::absent().to_quantization() };
assert!(imported
.expect("absent is well-formed, not an error")
.is_none());
}
#[test]
fn per_tensor_round_trips() {
for q in [
Quantization::per_tensor_symmetric(0.5),
Quantization::per_tensor(0.125, -7),
] {
assert_eq!(round_trip(&q).as_ref(), Some(&q));
}
}
#[test]
fn per_channel_symmetric_round_trips_with_its_length_and_axis() {
let q = Quantization::per_channel_symmetric(vec![0.5, 0.25, 0.125], 0)
.expect("per-channel symmetric");
let imported = round_trip(&q).expect("per-channel must survive the descriptor");
assert_eq!(imported, q);
assert_eq!(imported.scale(), &[0.5, 0.25, 0.125]);
assert_eq!(imported.axis(), Some(0));
assert_eq!(imported.zero_point(), None);
}
#[test]
fn per_channel_asymmetric_round_trips_with_its_zero_points() {
let q = Quantization::per_channel(vec![0.5, 0.25, 0.125], vec![1, 2, 3], 2)
.expect("per-channel asymmetric");
let imported = round_trip(&q).expect("per-channel must survive the descriptor");
assert_eq!(imported, q);
assert_eq!(imported.scale().len(), 3);
assert_eq!(imported.axis(), Some(2));
assert_eq!(imported.zero_point(), Some(&[1, 2, 3][..]));
}
#[test]
fn the_axis_plus_one_encoding_is_reversible() {
for axis in 0..4usize {
let q =
Quantization::per_channel_symmetric(vec![0.5, 0.25], axis).expect("per-channel");
assert_eq!(QuantDesc::borrowing(&q).axis_plus_one, axis as u64 + 1);
assert_eq!(round_trip(&q).and_then(|i| i.axis()), Some(axis));
}
}
#[test]
fn per_channel_without_an_axis_is_an_error_not_a_silent_none() {
let q = Quantization::per_channel_symmetric(vec![0.5, 0.25], 0).expect("per-channel");
let mut desc = QuantDesc::borrowing(&q);
desc.axis_plus_one = 0;
let err = unsafe { desc.to_quantization() }
.expect_err("per-channel scales naming no axis is malformed");
assert!(
err.to_string().contains("no channel axis"),
"the error must say what is wrong with the descriptor, got: {err}"
);
}
#[test]
fn a_null_scale_array_with_a_non_zero_length_is_an_error() {
let mut desc = QuantDesc::absent();
desc.len = 3;
let err =
unsafe { desc.to_quantization() }.expect_err("a length with a null array is malformed");
assert!(err.to_string().contains("null scale array"), "got: {err}");
}
}