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;
}
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 fn c_byte_strides(
shape: &[u64],
esz: i64,
layout: Option<crate::PixelLayout>,
row_stride: Option<usize>,
) -> Vec<i64> {
let mut strides = vec![0i64; shape.len()];
let mut acc = esz;
for (stride, dim) in strides.iter_mut().zip(shape).rev() {
*stride = acc;
acc *= *dim as i64;
}
let Some(rs) = row_stride.map(|rs| rs as i64) else {
return strides;
};
let image_rank = match layout {
Some(crate::PixelLayout::SemiPlanar) => Some(2),
Some(crate::PixelLayout::Packed | crate::PixelLayout::Planar) => Some(3),
None => None,
};
let batch_dims = match image_rank {
Some(rank) if shape.len() == rank + 1 => 1,
_ => 0,
};
let (lead, image) = strides.split_at_mut(batch_dims);
let image_bytes = match (layout, image, &shape[batch_dims..]) {
(Some(crate::PixelLayout::Planar), [plane, row, _, ..], [planes, rows, ..]) => {
*row = rs;
*plane = rs * *rows as i64;
*plane * *planes as i64
}
(_, [row, _, ..], [rows, ..]) => {
*row = rs;
rs * *rows as i64
}
_ => return strides,
};
if let [n] = lead {
*n = image_bytes;
}
strides
}
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 byte_strides = c_byte_strides(
&shape[..filled],
dtype.size() as i64,
format.map(|f| f.layout()),
row_stride,
);
strides[..filled].copy_from_slice(&byte_strides);
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}");
}
}
#[cfg(test)]
mod desc_tests {
use super::{c_byte_strides, from_parts, DescParts};
use crate::{DType, PixelFormat, PixelLayout, TensorMemory};
#[test]
fn a_padded_batch_steps_n_over_one_padded_image() {
assert_eq!(
c_byte_strides(&[2, 2, 8, 3], 1, Some(PixelLayout::Packed), Some(32)),
[64, 32, 3, 1]
);
assert_eq!(
c_byte_strides(&[2, 3, 4, 8], 1, Some(PixelLayout::Planar), Some(16)),
[192, 64, 16, 1]
);
assert_eq!(
c_byte_strides(&[2, 6, 8], 1, Some(PixelLayout::SemiPlanar), Some(32)),
[192, 32, 1]
);
assert_eq!(
c_byte_strides(&[2, 2, 8, 3], 1, None, Some(64)),
[64, 24, 3, 1]
);
}
#[test]
fn a_tight_tensor_gets_row_major_byte_strides() {
assert_eq!(c_byte_strides(&[3, 8], 4, None, None), [32, 4]);
assert_eq!(
c_byte_strides(&[48, 64, 3], 1, Some(PixelLayout::Packed), None),
[192, 3, 1]
);
}
#[test]
fn a_recorded_pitch_replaces_the_row_stride_of_a_packed_or_semi_planar_image() {
assert_eq!(
c_byte_strides(&[48, 64, 3], 1, Some(PixelLayout::Packed), Some(256)),
[256, 3, 1]
);
assert_eq!(
c_byte_strides(&[48, 64], 1, Some(PixelLayout::SemiPlanar), Some(128)),
[128, 1]
);
}
#[test]
fn a_recorded_pitch_moves_both_the_row_and_the_plane_stride_of_a_planar_image() {
assert_eq!(
c_byte_strides(&[3, 48, 64], 1, Some(PixelLayout::Planar), Some(128)),
[128 * 48, 128, 1]
);
assert_eq!(
c_byte_strides(&[3, 48, 64], 4, Some(PixelLayout::Planar), Some(320)),
[320 * 48, 320, 4]
);
}
#[test]
fn a_rank_one_tensor_has_no_row_to_pad() {
assert_eq!(c_byte_strides(&[16], 1, None, Some(64)), [1]);
}
fn parts(
memory: TensorMemory,
fence_handle: usize,
sync: Option<u64>,
) -> DescParts<'static, 'static> {
DescParts {
dims: &[48, 64, 3],
memory,
dtype: DType::U8,
fourcc: 0,
format: Some(PixelFormat::Rgb),
row_stride: None,
handle: -1,
colorimetry: 0,
capacity: 48 * 64 * 3,
pin: None,
pbo_vtable_ptr: None,
fence_handle,
sync,
}
}
#[test]
fn a_completion_on_a_kind_with_no_fence_is_not_advertised() {
let d = from_parts(parts(TensorMemory::Mem, 5, Some(9)));
assert_eq!(d.flags, 0);
assert_eq!(d.sync, 0);
assert!(d.ptr.is_null());
}
#[cfg(target_os = "windows")]
#[test]
fn a_texture_completion_needs_a_fence_handle() {
use super::flags;
let d = from_parts(parts(TensorMemory::DmaBuf, 5, Some(9)));
assert_eq!(d.flags, flags::SYNC_PRESENT);
assert_eq!(d.sync, 9);
assert_eq!(d.ptr.0 as usize, 5);
let d = from_parts(parts(TensorMemory::DmaBuf, 0, Some(9)));
assert_eq!(
d.flags, 0,
"a value with no fence handle cannot be waited on"
);
assert_eq!(d.sync, 0);
assert!(d.ptr.is_null());
}
}