pub const HEADER_LEN: usize = 72;
pub const PLANE_RECORD_LEN: usize = 56;
pub const SUPPORTED_REQUIRED_MASK: u64 = 0;
pub const HEADER_PID_OFFSET: usize = 28;
pub const D3D11_HANDLE_BYTES: usize = 24;
#[cfg(unix)]
type BlobFd = std::os::fd::RawFd;
#[cfg(not(unix))]
type BlobFd = i32;
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum BlobError {
Truncated {
need: usize,
have: usize,
},
UnsupportedRequirement(u64),
ForgedLength {
field: &'static str,
value: u64,
},
InvalidUtf8(&'static str),
Malformed(&'static str),
}
impl std::fmt::Display for BlobError {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
Self::Truncated { need, have } => {
write!(f, "blob truncated: needs {need} bytes, buffer has {have}")
}
Self::UnsupportedRequirement(m) => write!(
f,
"blob requires unsupported features (required_mask {m:#x}); \
refusing rather than misreading it"
),
Self::ForgedLength { field, value } => {
write!(f, "blob field `{field}` has implausible value {value}")
}
Self::InvalidUtf8(which) => write!(f, "blob string `{which}` is not valid UTF-8"),
Self::Malformed(why) => write!(f, "malformed blob: {why}"),
}
}
}
impl std::error::Error for BlobError {}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
pub struct BlobHeader {
pub size: u64,
pub required_mask: u64,
pub planes_bytes: u64,
pub storage_kind: u32,
pub pid: u32,
pub fence_fd: i32,
pub dtype: u32,
pub quant_axis: i32,
pub ndim: u32,
pub plane_count: u32,
pub quant_scales_len: u32,
pub quant_zero_points_len: u32,
pub strings_bytes: u32,
pub strides_len: u32,
}
impl BlobHeader {
#[cfg_attr(test, mutants::skip)]
pub fn empty() -> Self {
Self::default()
}
pub fn write_to(&self, buf: &mut [u8]) {
assert!(buf.len() >= HEADER_LEN, "header buffer too small");
buf[0..8].copy_from_slice(&self.size.to_le_bytes());
buf[8..16].copy_from_slice(&self.required_mask.to_le_bytes());
buf[16..24].copy_from_slice(&self.planes_bytes.to_le_bytes());
buf[24..28].copy_from_slice(&self.storage_kind.to_le_bytes());
buf[28..32].copy_from_slice(&self.pid.to_le_bytes());
buf[32..36].copy_from_slice(&self.fence_fd.to_le_bytes());
buf[36..40].copy_from_slice(&self.dtype.to_le_bytes());
buf[40..44].copy_from_slice(&self.quant_axis.to_le_bytes());
buf[44..48].copy_from_slice(&self.ndim.to_le_bytes());
buf[48..52].copy_from_slice(&self.plane_count.to_le_bytes());
buf[52..56].copy_from_slice(&self.quant_scales_len.to_le_bytes());
buf[56..60].copy_from_slice(&self.quant_zero_points_len.to_le_bytes());
buf[60..64].copy_from_slice(&self.strings_bytes.to_le_bytes());
buf[64..68].copy_from_slice(&self.strides_len.to_le_bytes());
buf[68..72].copy_from_slice(&0u32.to_le_bytes()); }
}
fn rd_u64(buf: &[u8], off: usize) -> u64 {
let mut b = [0u8; 8];
b.copy_from_slice(&buf[off..off + 8]);
u64::from_le_bytes(b)
}
fn rd_u32(buf: &[u8], off: usize) -> u32 {
let mut b = [0u8; 4];
b.copy_from_slice(&buf[off..off + 4]);
u32::from_le_bytes(b)
}
pub fn parse_header(buf: &[u8]) -> Result<BlobHeader, BlobError> {
if buf.len() < HEADER_LEN {
return Err(BlobError::Truncated {
need: HEADER_LEN,
have: buf.len(),
});
}
let required_mask = rd_u64(buf, 8);
if required_mask & !SUPPORTED_REQUIRED_MASK != 0 {
return Err(BlobError::UnsupportedRequirement(required_mask));
}
let size = rd_u64(buf, 0);
if size < HEADER_LEN as u64 {
return Err(BlobError::ForgedLength {
field: "size",
value: size,
});
}
if size > buf.len() as u64 {
return Err(BlobError::Truncated {
need: size as usize,
have: buf.len(),
});
}
if rd_u32(buf, 68) != 0 {
return Err(BlobError::Malformed("reserved header word is not zero"));
}
Ok(BlobHeader {
size,
required_mask,
planes_bytes: rd_u64(buf, 16),
storage_kind: rd_u32(buf, 24),
pid: rd_u32(buf, 28),
fence_fd: rd_u32(buf, 32) as i32,
dtype: rd_u32(buf, 36),
quant_axis: rd_u32(buf, 40) as i32,
ndim: rd_u32(buf, 44),
plane_count: rd_u32(buf, 48),
quant_scales_len: rd_u32(buf, 52),
quant_zero_points_len: rd_u32(buf, 56),
strings_bytes: rd_u32(buf, 60),
strides_len: rd_u32(buf, 64),
})
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct Regions {
pub shape: usize,
pub strides: usize,
pub quant_scales: usize,
pub quant_zero_points: usize,
pub strings: usize,
pub planes: usize,
pub end: usize,
}
fn advance(
off: usize,
count: u32,
elem: usize,
limit: usize,
field: &'static str,
) -> Result<usize, BlobError> {
let bytes = (count as usize)
.checked_mul(elem)
.ok_or(BlobError::ForgedLength {
field,
value: count as u64,
})?;
let next = off.checked_add(bytes).ok_or(BlobError::ForgedLength {
field,
value: count as u64,
})?;
if next > limit {
return Err(BlobError::ForgedLength {
field,
value: count as u64,
});
}
Ok(next)
}
pub fn region_offsets(h: &BlobHeader) -> Result<Regions, BlobError> {
let limit = h.size as usize;
if limit < HEADER_LEN {
return Err(BlobError::ForgedLength {
field: "size",
value: h.size,
});
}
if h.strides_len != 0 && h.strides_len != h.ndim {
return Err(BlobError::ForgedLength {
field: "strides_len",
value: h.strides_len as u64,
});
}
let shape = HEADER_LEN;
let strides = advance(shape, h.ndim, 8, limit, "ndim")?;
let quant_scales = advance(strides, h.strides_len, 8, limit, "strides_len")?;
let quant_zero_points = advance(
quant_scales,
h.quant_scales_len,
4,
limit,
"quant_scales_len",
)?;
let strings = advance(
quant_zero_points,
h.quant_zero_points_len,
4,
limit,
"quant_zero_points_len",
)?;
let planes = advance(strings, h.strings_bytes, 1, limit, "strings_bytes")?;
if !planes.is_multiple_of(8) {
return Err(BlobError::Malformed(
"planes region is not 8-aligned; strings_bytes must include padding",
));
}
let planes_bytes = usize::try_from(h.planes_bytes).map_err(|_| BlobError::ForgedLength {
field: "planes_bytes",
value: h.planes_bytes,
})?;
let end = planes
.checked_add(planes_bytes)
.ok_or(BlobError::ForgedLength {
field: "planes_bytes",
value: h.planes_bytes,
})?;
if end > limit {
return Err(BlobError::ForgedLength {
field: "planes_bytes",
value: h.planes_bytes,
});
}
let max_planes = planes_bytes / PLANE_RECORD_LEN;
if h.plane_count as usize > max_planes {
return Err(BlobError::ForgedLength {
field: "plane_count",
value: h.plane_count as u64,
});
}
Ok(Regions {
shape,
strides,
quant_scales,
quant_zero_points,
strings,
planes,
end,
})
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
pub struct BlobStrings<'a> {
pub format: &'a str,
pub color_space: &'a str,
pub color_transfer: &'a str,
pub color_encoding: &'a str,
pub color_range: &'a str,
}
const STRING_COUNT: usize = 5;
pub fn strings_encoded_len(s: &BlobStrings<'_>, start: usize) -> usize {
let raw: usize = STRING_COUNT * 4
+ s.format.len()
+ s.color_space.len()
+ s.color_transfer.len()
+ s.color_encoding.len()
+ s.color_range.len();
(start + raw).next_multiple_of(8) - start
}
pub fn write_strings(s: &BlobStrings<'_>, out: &mut Vec<u8>) {
for part in [
s.format,
s.color_space,
s.color_transfer,
s.color_encoding,
s.color_range,
] {
out.extend_from_slice(&(part.len() as u32).to_le_bytes());
out.extend_from_slice(part.as_bytes());
}
while !out.len().is_multiple_of(8) {
out.push(0);
}
}
pub fn parse_strings(buf: &[u8]) -> Result<BlobStrings<'_>, BlobError> {
const NAMES: [&str; STRING_COUNT] = [
"format",
"color_space",
"color_transfer",
"color_encoding",
"color_range",
];
let mut out: [&str; STRING_COUNT] = [""; STRING_COUNT];
let mut off = 0usize;
for (i, name) in NAMES.iter().enumerate() {
let name = *name;
let end = off
.checked_add(4)
.ok_or(BlobError::Malformed("string count overflows"))?;
if end > buf.len() {
return Err(BlobError::Truncated {
need: end,
have: buf.len(),
});
}
let len = rd_u32(buf, off) as usize;
off = end;
let stop = off.checked_add(len).ok_or(BlobError::ForgedLength {
field: name,
value: len as u64,
})?;
if stop > buf.len() {
return Err(BlobError::ForgedLength {
field: name,
value: len as u64,
});
}
out[i] = std::str::from_utf8(&buf[off..stop]).map_err(|_| BlobError::InvalidUtf8(name))?;
off = stop;
}
Ok(BlobStrings {
format: out[0],
color_space: out[1],
color_transfer: out[2],
color_encoding: out[3],
color_range: out[4],
})
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct BlobPlane<'a> {
pub handle: i64,
pub offset: u64,
pub stride: u64,
pub size: u64,
pub used: u64,
pub modifier: u64,
pub handle_bytes: &'a [u8],
pub data: &'a [u8],
}
impl BlobPlane<'_> {
pub fn is_inline(&self) -> bool {
self.handle < 0
}
pub fn encoded_len(&self) -> usize {
(PLANE_RECORD_LEN + self.handle_bytes.len() + self.data.len()).next_multiple_of(8)
}
}
pub fn write_plane(p: &BlobPlane<'_>, out: &mut Vec<u8>) {
let start = out.len();
out.extend_from_slice(&p.handle.to_le_bytes());
out.extend_from_slice(&p.offset.to_le_bytes());
out.extend_from_slice(&p.stride.to_le_bytes());
out.extend_from_slice(&p.size.to_le_bytes());
out.extend_from_slice(&p.used.to_le_bytes());
out.extend_from_slice(&p.modifier.to_le_bytes());
out.extend_from_slice(&(p.handle_bytes.len() as u32).to_le_bytes());
out.extend_from_slice(&(p.data.len() as u32).to_le_bytes());
out.extend_from_slice(p.handle_bytes);
out.extend_from_slice(p.data);
while !(out.len() - start).is_multiple_of(8) {
out.push(0);
}
}
fn rd_i64(buf: &[u8], off: usize) -> i64 {
let mut b = [0u8; 8];
b.copy_from_slice(&buf[off..off + 8]);
i64::from_le_bytes(b)
}
pub fn parse_planes(buf: &[u8], count: u32) -> Result<Vec<BlobPlane<'_>>, BlobError> {
let mut planes = Vec::with_capacity(count as usize);
let mut off = 0usize;
for _ in 0..count {
let head = off
.checked_add(PLANE_RECORD_LEN)
.ok_or(BlobError::Malformed("plane record offset overflows"))?;
if head > buf.len() {
return Err(BlobError::Truncated {
need: head,
have: buf.len(),
});
}
let hb_len = rd_u32(buf, off + 48) as usize;
let data_len = rd_u32(buf, off + 52) as usize;
let hb_end = head.checked_add(hb_len).ok_or(BlobError::ForgedLength {
field: "handle_bytes_len",
value: hb_len as u64,
})?;
let data_end = hb_end
.checked_add(data_len)
.ok_or(BlobError::ForgedLength {
field: "data_len",
value: data_len as u64,
})?;
if hb_end > buf.len() {
return Err(BlobError::ForgedLength {
field: "handle_bytes_len",
value: hb_len as u64,
});
}
if data_end > buf.len() {
return Err(BlobError::ForgedLength {
field: "data_len",
value: data_len as u64,
});
}
planes.push(BlobPlane {
handle: rd_i64(buf, off),
offset: rd_u64(buf, off + 8),
stride: rd_u64(buf, off + 16),
size: rd_u64(buf, off + 24),
used: rd_u64(buf, off + 32),
modifier: rd_u64(buf, off + 40),
handle_bytes: &buf[head..hb_end],
data: &buf[hb_end..data_end],
});
off = data_end.next_multiple_of(8);
}
Ok(planes)
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum TransportMode {
Reference,
Inline,
}
#[derive(Debug, Clone, Copy)]
pub struct BlobView<'a> {
buf: &'a [u8],
header: BlobHeader,
regions: Regions,
}
impl<'a> BlobView<'a> {
pub fn parse(buf: &'a [u8]) -> Result<Self, BlobError> {
let header = parse_header(buf)?;
let regions = region_offsets(&header)?;
Ok(Self {
buf,
header,
regions,
})
}
pub fn header(&self) -> BlobHeader {
self.header
}
pub fn shape(&self) -> Vec<u64> {
(0..self.header.ndim as usize)
.map(|i| rd_u64(self.buf, self.regions.shape + i * 8))
.collect()
}
pub fn strides(&self) -> Vec<i64> {
(0..self.header.strides_len as usize)
.map(|i| rd_i64(self.buf, self.regions.strides + i * 8))
.collect()
}
pub fn strings(&self) -> BlobStrings<'a> {
parse_strings(&self.buf[self.regions.strings..self.regions.planes]).unwrap_or_default()
}
pub fn planes(&self) -> Result<Vec<BlobPlane<'a>>, BlobError> {
parse_planes(
&self.buf[self.regions.planes..self.regions.end],
self.header.plane_count,
)
}
}
pub fn export(t: &crate::TensorDyn, mode: TransportMode) -> crate::Result<(Vec<u8>, Vec<BlobFd>)> {
let fmt = t.format();
let shape = export_addressing_shape(t, fmt)?;
let esz = t.dtype().size() as i64;
let strides =
crate::protocol::c_byte_strides(&shape, esz, fmt.map(|f| f.layout()), t.row_stride());
let quant = t.quantization();
let (quant_axis, scales, zeros): (i32, &[f32], &[i32]) = match quant {
None => (-2, &[], &[]),
Some(q) => (
q.axis().map(|a| a as i32).unwrap_or(-1),
q.scale(),
q.zero_point().unwrap_or(&[]),
),
};
let colorimetry = t.colorimetry().unwrap_or_default();
let strings = BlobStrings {
format: fmt.map(|f| f.as_str()).unwrap_or(""),
color_space: colorimetry.space.map(|v| v.as_str()).unwrap_or(""),
color_transfer: colorimetry.transfer.map(|v| v.as_str()).unwrap_or(""),
color_encoding: colorimetry.encoding.map(|v| v.as_str()).unwrap_or(""),
color_range: colorimetry.range.map(|v| v.as_str()).unwrap_or(""),
};
let geoms = export_plane_geoms(t, fmt)?;
let pin = match mode {
TransportMode::Inline => Some(t.pin_host(crate::CpuAccess::Read)?),
TransportMode::Reference => None,
};
let all_bytes: &[u8] = match &pin {
Some(p) => unsafe { std::slice::from_raw_parts(p.as_mut_ptr(), p.len()) },
None => &[],
};
let (handle, handle_bytes, fds): (i64, Vec<u8>, Vec<BlobFd>) = match mode {
TransportMode::Inline => (-1, Vec::new(), Vec::new()),
TransportMode::Reference => {
let raw = reference_handle(t)?;
#[cfg(target_os = "windows")]
{
(0, reference_handle_bytes(t, raw), Vec::new())
}
#[cfg(not(target_os = "windows"))]
{
let raw = i32::try_from(raw).map_err(|_| {
crate::Error::InvalidOperation("handle is not representable as an fd".into())
})?;
(0, Vec::new(), vec![raw])
}
}
};
let mut planes_buf = Vec::new();
write_export_planes(
&geoms,
mode,
handle,
&handle_bytes,
all_bytes,
&mut planes_buf,
)?;
let header = BlobHeader {
size: 0, required_mask: 0,
planes_bytes: planes_buf.len() as u64,
storage_kind: t.memory().code(),
pid: match mode {
TransportMode::Inline => 0,
TransportMode::Reference => std::process::id(),
},
fence_fd: -1,
dtype: t.dtype().code(),
quant_axis,
ndim: shape.len() as u32,
plane_count: geoms.len() as u32,
quant_scales_len: scales.len() as u32,
quant_zero_points_len: zeros.len() as u32,
strings_bytes: 0, strides_len: strides.len() as u32,
};
let mut out = vec![0u8; HEADER_LEN];
for d in &shape {
out.extend_from_slice(&d.to_le_bytes());
}
for st in &strides {
out.extend_from_slice(&st.to_le_bytes());
}
for sc in scales {
out.extend_from_slice(&sc.to_le_bytes());
}
for z in zeros {
out.extend_from_slice(&z.to_le_bytes());
}
let strings_start = out.len();
write_strings(&strings, &mut out);
let strings_bytes = out.len() - strings_start;
out.extend_from_slice(&planes_buf);
let mut header = header;
header.strings_bytes = strings_bytes as u32;
header.size = out.len() as u64;
header.write_to(&mut out);
Ok((out, fds))
}
fn export_addressing_shape(
t: &crate::TensorDyn,
fmt: Option<crate::PixelFormat>,
) -> crate::Result<Vec<u64>> {
match fmt {
Some(f) => {
let (w, h) = image_dims(t, f)?;
Ok(f.addressing_shape(w, h)
.ok_or_else(|| {
crate::Error::InvalidArgument(format!("no addressing shape for {f:?}"))
})?
.iter()
.map(|d| *d as u64)
.collect())
}
None => Ok(t.shape().iter().map(|d| *d as u64).collect()),
}
}
fn export_plane_geoms(
t: &crate::TensorDyn,
fmt: Option<crate::PixelFormat>,
) -> crate::Result<Vec<crate::PlaneGeometry>> {
match fmt {
Some(f) => {
let (w, h) = image_dims(t, f)?;
let rs = image_row_stride(t, f)?;
f.plane_table(w, h, rs)
.ok_or_else(|| crate::Error::InvalidArgument(format!("no plane table for {f:?}")))
}
None => Ok(vec![crate::PlaneGeometry {
offset: 0,
stride: t.capacity_bytes() as u64,
size: t.capacity_bytes() as u64,
}]),
}
}
fn write_export_planes(
geoms: &[crate::PlaneGeometry],
mode: TransportMode,
handle: i64,
handle_bytes: &[u8],
all_bytes: &[u8],
planes_buf: &mut Vec<u8>,
) -> crate::Result<()> {
for g in geoms {
let data: &[u8] = match mode {
TransportMode::Inline => {
let start = g.offset as usize;
let end = start
.checked_add(g.size as usize)
.ok_or(crate::Error::InvalidSize(g.size as usize))?;
all_bytes
.get(start..end)
.ok_or(crate::Error::InvalidSize(end))?
}
TransportMode::Reference => &[],
};
write_plane(
&BlobPlane {
handle,
offset: if mode == TransportMode::Inline {
0
} else {
g.offset
},
stride: g.stride,
size: if mode == TransportMode::Inline {
data.len() as u64
} else {
g.size
},
used: if mode == TransportMode::Inline {
data.len() as u64
} else {
g.size
},
modifier: 0,
handle_bytes,
data,
},
planes_buf,
);
}
Ok(())
}
fn image_dims(t: &crate::TensorDyn, f: crate::PixelFormat) -> crate::Result<(usize, usize)> {
let shape = t.shape();
let bad = || crate::Error::InvalidShape(format!("shape {shape:?} is not an image for {f:?}"));
match (f.layout(), shape) {
(crate::PixelLayout::Packed, [h, w, _]) => Ok((*w, *h)),
(crate::PixelLayout::Planar, [_, h, w]) => Ok((*w, *h)),
(crate::PixelLayout::SemiPlanar, [total, w]) => {
let h = (0..=*total)
.find(|&h| f.combined_plane_height(h) == Some(*total))
.ok_or_else(bad)?;
Ok((*w, h))
}
_ => Err(bad()),
}
}
fn image_row_stride(t: &crate::TensorDyn, f: crate::PixelFormat) -> crate::Result<usize> {
t.effective_row_stride().ok_or_else(|| {
crate::Error::InvalidShape(format!(
"{f:?} tensor of shape {:?} reports no row pitch",
t.shape()
))
})
}
fn reference_handle(t: &crate::TensorDyn) -> crate::Result<i64> {
match t.memory() {
#[cfg(target_os = "linux")]
crate::TensorMemory::DmaBuf => {
use std::os::fd::AsRawFd;
t.dmabuf().map(|fd| fd.as_raw_fd() as i64).map_err(|e| {
crate::Error::InvalidOperation(format!("dma-buf tensor has no fd: {e}"))
})
}
#[cfg(any(target_os = "macos", target_os = "ios"))]
crate::TensorMemory::DmaBuf => t
.iosurface_id()
.map(|id| id as i64)
.ok_or_else(|| crate::Error::InvalidOperation("IOSurface tensor has no id".into())),
#[cfg(target_os = "windows")]
crate::TensorMemory::DmaBuf => {
t.d3d11_shared_handle_value()
.map(|h| h as i64)
.ok_or_else(|| {
crate::Error::InvalidOperation(
"D3D11 texture tensor has no shared handle".into(),
)
})
}
other => Err(crate::Error::InvalidOperation(format!(
"{other:?} has no shareable handle; export it with TransportMode::Inline"
))),
}
}
#[cfg(target_os = "windows")]
fn reference_handle_bytes(t: &crate::TensorDyn, raw: i64) -> Vec<u8> {
let fence = crate::d3d11::device()
.map(|d| d.exported_fence_handle_value())
.unwrap_or(0);
let (fence, value) = reference_completion(fence, t.gpu_write_value());
let mut out = Vec::with_capacity(D3D11_HANDLE_BYTES);
out.extend_from_slice(&(raw as u64).to_le_bytes());
out.extend_from_slice(&fence.to_le_bytes());
out.extend_from_slice(&value.to_le_bytes());
out
}
#[cfg(target_os = "windows")]
fn reference_completion(fence: usize, recorded: u64) -> (u64, u64) {
if fence == 0 {
(0, 0)
} else {
(fence as u64, recorded)
}
}
pub fn import(blob: &[u8], fds: &[BlobFd]) -> crate::Result<crate::TensorDyn> {
#[cfg(not(unix))]
let _ = fds;
let v = BlobView::parse(blob).map_err(|e| crate::Error::InvalidArgument(e.to_string()))?;
let h = v.header();
let planes = v
.planes()
.map_err(|e| crate::Error::InvalidArgument(e.to_string()))?;
let inline = planes.first().map(BlobPlane::is_inline).unwrap_or(true);
validate_import_planes(&planes, inline)?;
let dtype = crate::protocol::dtype_to_dtype(h.dtype)
.ok_or_else(|| crate::Error::NotImplemented(format!("blob dtype code {}", h.dtype)))?;
let strings = v.strings();
let format = if strings.format.is_empty() {
None
} else {
Some(
crate::PixelFormat::from_str_code(strings.format).ok_or_else(|| {
crate::Error::NotImplemented(format!("blob format {:?}", strings.format))
})?,
)
};
if !inline {
#[cfg(target_os = "windows")]
{
return import_referenced_d3d11_blob(&v, &h, &planes, format, &strings, dtype);
}
#[cfg(not(any(unix, target_os = "windows")))]
{
return Err(crate::Error::NotImplemented(
"reference-mode blob import is not available on this platform".into(),
));
}
#[cfg(unix)]
{
return import_referenced_blob(&v, &h, &planes, fds, format, &strings, dtype);
}
}
let grid: Vec<usize> = v.shape().iter().map(|d| *d as usize).collect();
let alloc_shape = alloc_shape_from_grid(format, &grid)?;
let mut t = crate::TensorDyn::new(&alloc_shape, dtype, Some(crate::TensorMemory::Mem), None)?;
if let Some(f) = format {
t.set_format(f)?;
}
blit_inline_planes(&mut t, &planes, format, &grid)?;
t.set_colorimetry(colorimetry_from(&strings));
if let Some(q) = import_quantization(&v)? {
t.set_quantization(q)?;
}
Ok(t)
}
fn validate_import_planes(planes: &[BlobPlane<'_>], inline: bool) -> crate::Result<()> {
if planes.iter().any(|p| p.is_inline() != inline) {
return Err(crate::Error::InvalidArgument(
"blob mixes inline and referenced planes; a tensor has one transport mode".into(),
));
}
for p in planes {
if p.is_inline() {
if p.size as usize != p.data.len() {
return Err(crate::Error::InvalidArgument(format!(
"inline plane declares size {} but carries {} bytes",
p.size,
p.data.len()
)));
}
if p.modifier != 0 || !p.handle_bytes.is_empty() {
return Err(crate::Error::InvalidArgument(
"inline plane must have modifier 0 and no handle_bytes".into(),
));
}
} else if !p.data.is_empty() {
return Err(crate::Error::InvalidArgument(
"referenced plane must not also carry inline bytes".into(),
));
}
}
Ok(())
}
fn alloc_shape_from_grid(
format: Option<crate::PixelFormat>,
grid: &[usize],
) -> crate::Result<Vec<usize>> {
match format {
Some(f) => {
let (w, h_px) = dims_from_grid(f, grid)?;
f.allocation_shape(w, h_px)
.ok_or_else(|| crate::Error::InvalidShape(format!("no allocation shape for {f:?}")))
}
None => Ok(grid.to_vec()),
}
}
#[cfg(unix)]
fn import_referenced_blob(
v: &BlobView<'_>,
h: &BlobHeader,
planes: &[BlobPlane<'_>],
fds: &[BlobFd],
format: Option<crate::PixelFormat>,
strings: &BlobStrings<'_>,
dtype: crate::DType,
) -> crate::Result<crate::TensorDyn> {
let first = planes
.first()
.ok_or_else(|| crate::Error::InvalidArgument("referenced blob carries no planes".into()))?;
let idx = usize::try_from(first.handle).map_err(|_| {
crate::Error::InvalidArgument(format!("plane handle {} is not an index", first.handle))
})?;
let raw = *fds.get(idx).ok_or_else(|| {
crate::Error::InvalidArgument(format!(
"plane handle indexes fd {idx} but only {} were provided",
fds.len()
))
})?;
for p in planes {
let i = usize::try_from(p.handle).unwrap_or(usize::MAX);
if i >= fds.len() {
return Err(crate::Error::InvalidArgument(format!(
"plane handle indexes fd {i} but only {} were provided",
fds.len()
)));
}
}
let grid: Vec<usize> = v.shape().iter().map(|d| *d as usize).collect();
let alloc_shape = alloc_shape_from_grid(format, &grid)?;
let mut t = open_referenced_tensor(h.storage_kind, raw, &alloc_shape, dtype)?;
if let Some(f) = format {
t.set_format(f)?;
if first.stride > 0 {
t.set_row_stride(first.stride as usize)?;
}
}
t.set_colorimetry(colorimetry_from(strings));
if let Some(q) = import_quantization(v)? {
t.set_quantization(q)?;
}
Ok(t)
}
#[cfg(unix)]
fn open_referenced_tensor(
storage_kind: u32,
raw: BlobFd,
alloc_shape: &[usize],
dtype: crate::DType,
) -> crate::Result<crate::TensorDyn> {
match crate::TensorMemory::from_code(storage_kind) {
#[cfg(any(target_os = "macos", target_os = "ios"))]
Some(crate::TensorMemory::DmaBuf) | Some(crate::TensorMemory::IoSurface) => {
let id = u32::try_from(raw).map_err(|_| {
crate::Error::InvalidArgument(format!("{raw} is not an IOSurface id"))
})?;
crate::TensorDyn::from_iosurface_id(id, alloc_shape, dtype, None)
}
#[cfg(all(unix, not(any(target_os = "macos", target_os = "ios"))))]
Some(crate::TensorMemory::DmaBuf) | Some(crate::TensorMemory::Shm) => {
use std::os::fd::{BorrowedFd, OwnedFd};
let borrowed = unsafe { BorrowedFd::borrow_raw(raw) };
let owned: OwnedFd = borrowed
.try_clone_to_owned()
.map_err(crate::Error::IoError)?;
crate::TensorDyn::from_fd(owned, alloc_shape, dtype, None)
}
other => Err(crate::Error::NotImplemented(format!(
"reference-mode import for {other:?} on this platform"
))),
}
}
#[cfg(target_os = "windows")]
fn import_referenced_d3d11_blob(
v: &BlobView<'_>,
h: &BlobHeader,
planes: &[BlobPlane<'_>],
format: Option<crate::PixelFormat>,
strings: &BlobStrings<'_>,
dtype: crate::DType,
) -> crate::Result<crate::TensorDyn> {
let first = planes
.first()
.ok_or_else(|| crate::Error::InvalidArgument("referenced blob carries no planes".into()))?;
for p in planes {
if p.handle != 0 || p.handle_bytes != first.handle_bytes {
return Err(crate::Error::InvalidArgument(
"D3D11 reference blob planes must all carry handle 0 and identical \
handle bytes; they name one texture"
.into(),
));
}
}
let grid: Vec<usize> = v.shape().iter().map(|d| *d as usize).collect();
let mut t = open_referenced_texture(h.storage_kind, first, h.pid, &grid, dtype, format)?;
t.set_colorimetry(colorimetry_from(strings));
if let Some(q) = import_quantization(v)? {
t.set_quantization(q)?;
}
Ok(t)
}
#[cfg(target_os = "windows")]
fn open_referenced_texture(
storage_kind: u32,
first: &BlobPlane<'_>,
pid: u32,
grid: &[usize],
dtype: crate::DType,
format: Option<crate::PixelFormat>,
) -> crate::Result<crate::TensorDyn> {
use std::os::windows::io::{AsRawHandle, FromRawHandle, OwnedHandle};
use windows::Win32::Foundation::HANDLE;
use windows::Win32::System::Threading::{OpenProcess, PROCESS_DUP_HANDLE};
if storage_kind != crate::TensorMemory::DmaBuf.code() {
return Err(crate::Error::InvalidArgument(format!(
"storage kind {storage_kind} is not referenceable on Windows"
)));
}
if first.handle_bytes.len() != D3D11_HANDLE_BYTES {
return Err(crate::Error::InvalidArgument(format!(
"D3D11 reference plane needs {D3D11_HANDLE_BYTES} handle bytes, not {}",
first.handle_bytes.len()
)));
}
let format = format.ok_or_else(|| {
crate::Error::InvalidArgument("D3D11 reference blob without a pixel format".into())
})?;
let word = |i: usize| {
let mut b = [0u8; 8];
b.copy_from_slice(&first.handle_bytes[i * 8..(i + 1) * 8]);
u64::from_le_bytes(b)
};
let handle_at = |i: usize| HANDLE(word(i) as usize as *mut std::ffi::c_void);
let (mut tex, mut fence, value) = (handle_at(0), handle_at(1), word(2));
if tex.is_invalid() {
return Err(crate::Error::InvalidArgument(
"D3D11 reference plane carries no texture handle".into(),
));
}
if fence.is_invalid() && value != 0 {
return Err(crate::Error::InvalidArgument(format!(
"D3D11 reference plane names fence value {value} with no fence handle; \
the completion it names cannot be waited on"
)));
}
let mut duplicates: Vec<OwnedHandle> = Vec::new();
if std::process::id() != pid {
let src = unsafe { OpenProcess(PROCESS_DUP_HANDLE, false, pid) }
.map_err(|e| open_process_error(pid, &e))?;
let src = unsafe { OwnedHandle::from_raw_handle(src.0) };
for (h, which) in [(&mut tex, "texture"), (&mut fence, "fence")] {
if h.is_invalid() {
continue;
}
let dup = duplicate_out_of(&src, *h, pid, which)?;
*h = HANDLE(dup.as_raw_handle());
duplicates.push(dup);
}
}
let geometry = unsafe { crate::protocol::d3d11_geometry_checked(format, tex.0, grid) };
let completion = (value != 0).then_some((fence.0, value));
let opened = match geometry {
Ok((width, height)) => unsafe {
crate::TensorDyn::from_d3d11_shared_handle(
tex.0,
width,
height,
format,
dtype,
crate::CpuAccess::ReadWrite,
completion,
None,
)
},
Err(e) => Err(e),
};
drop(duplicates);
opened
}
#[cfg(target_os = "windows")]
fn win32_error_kind(e: &windows::core::Error) -> std::io::ErrorKind {
use std::io::ErrorKind;
use windows::core::HRESULT;
use windows::Win32::Foundation::{ERROR_ACCESS_DENIED, ERROR_INVALID_PARAMETER};
let code = e.code();
if code == HRESULT::from_win32(ERROR_ACCESS_DENIED.0) {
ErrorKind::PermissionDenied
} else if code == HRESULT::from_win32(ERROR_INVALID_PARAMETER.0) {
ErrorKind::NotFound
} else {
ErrorKind::Other
}
}
#[cfg(target_os = "windows")]
fn open_process_error(pid: u32, e: &windows::core::Error) -> crate::Error {
use std::io::ErrorKind;
let kind = win32_error_kind(e);
let what = match kind {
ErrorKind::PermissionDenied => {
format!("cannot open exporting process {pid} to duplicate its handles: access denied")
}
ErrorKind::NotFound => format!("exporting process {pid} is gone"),
_ => format!(
"cannot open exporting process {pid} to duplicate its handles: {} \
(HRESULT 0x{:08X})",
e.message(),
e.code().0 as u32
),
};
crate::Error::IoError(std::io::Error::new(kind, what))
}
#[cfg(target_os = "windows")]
fn duplicate_out_of(
src: &std::os::windows::io::OwnedHandle,
handle: windows::Win32::Foundation::HANDLE,
pid: u32,
which: &str,
) -> crate::Result<std::os::windows::io::OwnedHandle> {
use std::os::windows::io::{AsRawHandle, FromRawHandle, OwnedHandle};
use windows::Win32::Foundation::{DuplicateHandle, DUPLICATE_SAME_ACCESS, HANDLE};
use windows::Win32::System::Threading::GetCurrentProcess;
let mut out = HANDLE(std::ptr::null_mut());
unsafe {
DuplicateHandle(
HANDLE(src.as_raw_handle()),
handle,
GetCurrentProcess(),
&mut out,
0,
false,
DUPLICATE_SAME_ACCESS,
)
}
.map_err(|e| {
crate::Error::IoError(std::io::Error::new(
win32_error_kind(&e),
format!(
"cannot duplicate the {which} handle out of exporting process {pid}: {} \
(HRESULT 0x{:08X})",
e.message(),
e.code().0 as u32
),
))
})?;
Ok(unsafe { OwnedHandle::from_raw_handle(out.0) })
}
fn blit_inline_planes(
t: &mut crate::TensorDyn,
planes: &[BlobPlane<'_>],
format: Option<crate::PixelFormat>,
grid: &[usize],
) -> crate::Result<()> {
let pin = t.pin_host(crate::CpuAccess::Write)?;
let dst = unsafe { std::slice::from_raw_parts_mut(pin.as_mut_ptr(), pin.len()) };
let geoms: Vec<crate::PlaneGeometry> = match format {
Some(f) => {
let (w, h_px) = dims_from_grid(f, grid)?;
let rs = image_row_stride(t, f)?;
f.plane_table(w, h_px, rs)
.ok_or_else(|| crate::Error::InvalidShape(format!("no plane table for {f:?}")))?
}
None => vec![crate::PlaneGeometry {
offset: 0,
stride: dst.len() as u64,
size: dst.len() as u64,
}],
};
if geoms.len() != planes.len() {
return Err(crate::Error::InvalidArgument(format!(
"blob carries {} planes but {:?} has {}",
planes.len(),
format,
geoms.len()
)));
}
for (g, p) in geoms.iter().zip(planes) {
let start = g.offset as usize;
let end = start
.checked_add(p.data.len())
.ok_or(crate::Error::InvalidSize(p.data.len()))?;
let slot = dst
.get_mut(start..end)
.ok_or(crate::Error::InvalidSize(end))?;
slot.copy_from_slice(p.data);
}
Ok(())
}
fn colorimetry_from(s: &BlobStrings<'_>) -> Option<crate::Colorimetry> {
let c = crate::Colorimetry {
space: crate::ColorSpace::from_str_code(s.color_space),
transfer: crate::ColorTransfer::from_str_code(s.color_transfer),
encoding: crate::ColorEncoding::from_str_code(s.color_encoding),
range: crate::ColorRange::from_str_code(s.color_range),
};
if c == crate::Colorimetry::default() {
None
} else {
Some(c)
}
}
fn import_quantization(v: &BlobView<'_>) -> crate::Result<Option<crate::Quantization>> {
let h = v.header;
if h.quant_axis == -2 {
return Ok(None);
}
let scales: Vec<f32> = (0..h.quant_scales_len as usize)
.map(|i| f32::from_bits(rd_u32(v.buf, v.regions.quant_scales + i * 4)))
.collect();
let zeros: Vec<i32> = (0..h.quant_zero_points_len as usize)
.map(|i| rd_u32(v.buf, v.regions.quant_zero_points + i * 4) as i32)
.collect();
quantization_from(h.quant_axis, scales, zeros).map(Some)
}
fn quantization_from(
axis: i32,
scales: Vec<f32>,
zeros: Vec<i32>,
) -> crate::Result<crate::Quantization> {
let malformed = |why: String| {
crate::Error::InvalidArgument(format!("blob quantization is malformed: {why}"))
};
if axis == -1 {
return match (scales.as_slice(), zeros.as_slice()) {
([scale], []) => Ok(crate::Quantization::per_tensor_symmetric(*scale)),
([scale], [zero]) => Ok(crate::Quantization::per_tensor(*scale, *zero)),
(s, z) => Err(malformed(format!(
"per-tensor quantization needs 1 scale and 0 or 1 zero points, \
got {} and {}",
s.len(),
z.len()
))),
};
}
let axis = usize::try_from(axis).map_err(|_| {
malformed(format!(
"quant_axis {axis} is neither -2 (none), -1 (per-tensor) nor a channel index"
))
})?;
let q = if zeros.is_empty() {
crate::Quantization::per_channel_symmetric(scales, axis)
} else {
crate::Quantization::per_channel(scales, zeros, axis)
};
q.map_err(|e| malformed(e.to_string()))
}
fn dims_from_grid(f: crate::PixelFormat, grid: &[usize]) -> crate::Result<(usize, usize)> {
let bad = || crate::Error::InvalidShape(format!("grid {grid:?} is not an image for {f:?}"));
match (f.layout(), grid) {
(crate::PixelLayout::Planar, [_, h, w, ..]) => Ok((*w, *h)),
(crate::PixelLayout::Planar, _) => Err(bad()),
(_, [h, w, ..]) => Ok((*w, *h)),
_ => Err(bad()),
}
}
#[cfg(test)]
mod tests {
use super::*;
#[cfg(target_os = "windows")]
#[test]
fn a_reference_plane_drops_the_value_with_the_fence() {
assert_eq!(super::reference_completion(0, 42), (0, 0));
assert_eq!(super::reference_completion(0, 0), (0, 0));
assert_eq!(super::reference_completion(0x1234, 0), (0x1234, 0));
assert_eq!(super::reference_completion(0x1234, 42), (0x1234, 42));
}
#[test]
fn every_region_offset_is_constant_time_arithmetic() {
let h = BlobHeader {
size: 4096,
ndim: 3,
strides_len: 3,
quant_scales_len: 4,
quant_zero_points_len: 4,
strings_bytes: 24,
plane_count: 2,
planes_bytes: 128,
..BlobHeader::empty()
};
let r = region_offsets(&h).unwrap();
assert_eq!(r.shape, HEADER_LEN);
assert_eq!(r.strides, HEADER_LEN + 3 * 8);
assert_eq!(r.quant_scales, HEADER_LEN + 3 * 8 + 3 * 8);
assert_eq!(r.quant_zero_points, r.quant_scales + 4 * 4);
assert_eq!(r.strings, r.quant_zero_points + 4 * 4);
assert_eq!(r.planes, r.strings + 24);
assert_eq!(r.end, r.planes + 128);
}
#[test]
fn empty_strides_leave_the_region_zero_width() {
let h = BlobHeader {
size: 4096,
ndim: 2,
strides_len: 0,
..BlobHeader::empty()
};
let r = region_offsets(&h).unwrap();
assert_eq!(r.strides, HEADER_LEN + 2 * 8);
assert_eq!(
r.quant_scales, r.strides,
"densely-packed strides occupy no bytes"
);
}
#[test]
fn strides_must_be_ndim_entries_or_none() {
let h = BlobHeader {
size: 4096,
ndim: 3,
strides_len: 2,
..BlobHeader::empty()
};
assert!(matches!(
region_offsets(&h),
Err(BlobError::ForgedLength {
field: "strides_len",
..
})
));
}
#[test]
fn a_forged_ndim_cannot_derive_an_offset_past_the_buffer() {
let h = BlobHeader {
size: 72,
ndim: u32::MAX,
strides_len: u32::MAX,
..BlobHeader::empty()
};
assert!(
region_offsets(&h).is_err(),
"ndim * 8 must be checked, not wrapped"
);
}
#[test]
fn regions_that_overrun_size_are_refused() {
let h = BlobHeader {
size: 80,
ndim: 2,
strides_len: 2,
plane_count: 1,
planes_bytes: 4096,
..BlobHeader::empty()
};
assert!(region_offsets(&h).is_err());
}
#[test]
fn plane_count_is_bounded_by_the_planes_region_not_by_a_constant() {
let h = BlobHeader {
size: 4096,
plane_count: u32::MAX,
planes_bytes: 56,
..BlobHeader::empty()
};
assert!(matches!(
region_offsets(&h),
Err(BlobError::ForgedLength {
field: "plane_count",
..
})
));
}
fn fill_ramp(t: &TensorDyn) {
let pin = t.pin_host(CpuAccess::ReadWrite).expect("pin");
t.sync_for_cpu(CpuAccess::ReadWrite).expect("sync for cpu");
let s = unsafe { std::slice::from_raw_parts_mut(pin.as_mut_ptr(), pin.len()) };
for (i, b) in s.iter_mut().enumerate() {
*b = (i % 251) as u8;
}
t.sync_for_device(CpuAccess::ReadWrite)
.expect("sync for device");
}
fn bytes_of(t: &TensorDyn) -> Vec<u8> {
let pin = t.pin_host(CpuAccess::Read).expect("pin");
t.sync_for_cpu(CpuAccess::Read).expect("sync for cpu");
let bytes = unsafe { std::slice::from_raw_parts(pin.as_mut_ptr(), pin.len()) }.to_vec();
t.sync_for_device(CpuAccess::Read).expect("sync for device");
bytes
}
fn reference_capable(w: usize, h: usize, what: &str) -> Option<TensorDyn> {
if cfg!(target_os = "linux") && !crate::test_support::dma_or_skip(what) {
return None;
}
match Tensor::<u8>::image(
w,
h,
PixelFormat::Nv12,
Some(TensorMemory::DmaBuf),
CpuAccess::ReadWrite,
) {
Ok(t) => Some(TensorDyn::from(t)),
Err(e) if cfg!(target_os = "linux") => {
panic!("{what}: a usable DMA heap must allocate NV12 {w}x{h}: {e:?}")
}
Err(e) => {
crate::test_support::report_skip(&format!(
"no shareable backing on this host ({e:?}); reference \
transport not exercised"
));
None
}
}
}
#[test]
fn a_reference_export_puts_an_index_in_the_handle_not_a_raw_fd() {
let Some(src) = reference_capable(
64,
48,
"a_reference_export_puts_an_index_in_the_handle_not_a_raw_fd",
) else {
return;
};
let (blob, fds) = export(&src, TransportMode::Reference).unwrap();
let v = BlobView::parse(&blob).unwrap();
let planes = v.planes().unwrap();
for p in &planes {
assert!(!p.is_inline());
}
if cfg!(target_os = "windows") {
assert!(fds.is_empty(), "an NT handle does not travel out of band");
assert!(planes
.iter()
.all(|p| p.handle_bytes.len() == D3D11_HANDLE_BYTES));
} else {
assert!(!fds.is_empty(), "reference mode must carry an fd table");
for p in &planes {
assert!(
(p.handle as usize) < fds.len(),
"handle {} must index the {}-entry fd table, not be a raw fd",
p.handle,
fds.len()
);
}
assert_eq!(fds.len(), 1, "planes sharing a buffer share a table entry");
}
assert!(planes.iter().all(|p| p.handle == 0));
}
#[test]
fn a_reference_round_trip_preserves_content_and_the_source_may_die() {
let Some(src) = reference_capable(
64,
48,
"a_reference_round_trip_preserves_content_and_the_source_may_die",
) else {
return;
};
fill_ramp(&src);
let expected = bytes_of(&src);
let (blob, fds) = export(&src, TransportMode::Reference).unwrap();
let src_id = src.buffer_identity().id();
let got = import(&blob, &fds).unwrap();
if !cfg!(target_os = "windows") {
assert_eq!(
got.buffer_identity().id(),
src_id,
"reference mode must share the buffer, not copy it"
);
}
drop(src);
assert_eq!(
bytes_of(&got),
expected,
"the imported tensor sees the bytes"
);
assert_eq!(got.format(), Some(PixelFormat::Nv12));
}
#[test]
fn an_fd_index_past_the_table_is_refused() {
let Some(src) = reference_capable(64, 48, "an_fd_index_past_the_table_is_refused") else {
return;
};
let (blob, fds) = export(&src, TransportMode::Reference).unwrap();
if cfg!(target_os = "windows") {
let err = import(&patch_handle_word(&blob, 0, 0), &[])
.expect_err("a null texture handle must be refused, not opened");
assert!(
matches!(err, crate::Error::InvalidArgument(_)),
"expected InvalidArgument, got {err:?}"
);
let named = patch_handle_word(&patch_handle_word(&blob, 1, 0), 2, 7);
let err = import(&named, &[])
.expect_err("a fence value with no fence handle must be refused");
let msg = format!("{err}");
assert!(
matches!(err, crate::Error::InvalidArgument(_))
&& msg.contains("cannot be waited on"),
"expected the unwaitable-completion refusal, got {err:?}"
);
let mut tampered = blob.clone();
tampered[HEADER_PID_OFFSET..HEADER_PID_OFFSET + 4]
.copy_from_slice(&u32::MAX.to_le_bytes());
let err =
import(&tampered, &[]).expect_err("a pid that does not exist must be refused");
match &err {
crate::Error::IoError(e) => assert_eq!(
e.kind(),
std::io::ErrorKind::NotFound,
"a pid no process holds is reported as gone, got {e:?}"
),
other => panic!("expected IoError, got {other:?}"),
}
return;
}
let r = region_offsets(&BlobView::parse(&blob).unwrap().header()).unwrap();
assert!(
import(&blob, &[]).is_err(),
"an empty table cannot satisfy an index"
);
let mut tampered = blob.clone();
tampered[r.planes..r.planes + 8].copy_from_slice(&99i64.to_le_bytes());
assert!(
import(&tampered, &fds).is_err(),
"an out-of-range index must be refused, not read"
);
}
fn patch_handle_word(blob: &[u8], word: usize, value: u64) -> Vec<u8> {
let v = BlobView::parse(blob).unwrap();
let r = region_offsets(&v.header()).unwrap();
let stride = (PLANE_RECORD_LEN + D3D11_HANDLE_BYTES).next_multiple_of(8);
let mut out = blob.to_vec();
for i in 0..v.header().plane_count as usize {
let at = r.planes + i * stride + PLANE_RECORD_LEN + word * 8;
out[at..at + 8].copy_from_slice(&value.to_le_bytes());
}
out
}
#[test]
#[cfg(target_os = "windows")]
fn a_garbage_texture_handle_is_refused_by_the_constructor() {
let Some(src) = reference_capable(
64,
48,
"a_garbage_texture_handle_is_refused_by_the_constructor",
) else {
return;
};
let (blob, _) = export(&src, TransportMode::Reference).unwrap();
let err = import(&patch_handle_word(&blob, 0, 0x1234), &[])
.expect_err("a handle that is not a shared texture must be refused");
let msg = format!("{err}");
assert!(
msg.contains("OpenSharedResource"),
"the refusal must come from opening the handle, got: {msg}"
);
}
#[test]
#[cfg(target_os = "windows")]
fn planes_that_disagree_about_the_texture_are_refused() {
let Some(src) =
reference_capable(64, 48, "planes_that_disagree_about_the_texture_are_refused")
else {
return;
};
let (blob, _) = export(&src, TransportMode::Reference).unwrap();
let v = BlobView::parse(&blob).unwrap();
let r = region_offsets(&v.header()).unwrap();
assert_eq!(v.planes().unwrap().len(), 2, "NV12 exports two planes");
let second = r.planes + (PLANE_RECORD_LEN + D3D11_HANDLE_BYTES).next_multiple_of(8);
let mut tampered = blob.clone();
tampered[second + PLANE_RECORD_LEN..second + PLANE_RECORD_LEN + 8]
.copy_from_slice(&0xBEEFu64.to_le_bytes());
let err = import(&tampered, &[]).expect_err("planes naming two textures must be refused");
let msg = format!("{err}");
assert!(
matches!(err, crate::Error::InvalidArgument(_)) && msg.contains("one texture"),
"expected the one-texture refusal, got {err:?}"
);
}
#[test]
fn pid_offset_matches_the_writer() {
let h = BlobHeader {
size: HEADER_LEN as u64,
pid: 0x0BAD_F00D,
..BlobHeader::empty()
};
let mut buf = vec![0u8; HEADER_LEN];
h.write_to(&mut buf);
assert_eq!(rd_u32(&buf, HEADER_PID_OFFSET), h.pid);
}
#[test]
fn inline_round_trip_preserves_content_and_every_metadata_field() {
let src = nv12_mem(64, 48).with_colorimetry(crate::Colorimetry::jfif());
fill_ramp(&src);
let blob = export(&src, TransportMode::Inline).unwrap().0;
let got = import(&blob, &[]).unwrap();
assert_eq!(got.dtype(), src.dtype());
assert_eq!(got.format(), src.format());
assert_eq!(got.colorimetry(), src.colorimetry());
assert_eq!(got.shape(), src.shape(), "allocation geometry is rebuilt");
assert_eq!(bytes_of(&got), bytes_of(&src), "every byte survived");
}
#[test]
fn inline_import_is_a_new_allocation_and_the_source_may_die() {
let src = nv12_mem(64, 48);
fill_ramp(&src);
let expected = bytes_of(&src);
let blob = export(&src, TransportMode::Inline).unwrap().0;
drop(src);
let got = import(&blob, &[]).unwrap();
assert_eq!(bytes_of(&got), expected);
}
#[test]
fn a_batched_image_is_refused_rather_than_exported_as_one_image() {
let mut t = Tensor::<u8>::image(
8,
4,
crate::PixelFormat::Rgb,
Some(TensorMemory::Mem),
crate::CpuAccess::ReadWrite,
)
.unwrap();
crate::TensorTrait::set_logical_shape(&mut t, &[2, 2, 8, 3]).unwrap();
match export(&TensorDyn::from(t), TransportMode::Inline) {
Err(crate::Error::InvalidShape(msg)) => {
assert!(msg.contains("[2, 2, 8, 3]"), "{msg}")
}
other => panic!("expected InvalidShape, got {other:?}"),
}
}
#[test]
fn a_bare_tensor_round_trips_including_byte_strides() {
let src =
TensorDyn::from(Tensor::<f32>::new(&[3, 8], Some(TensorMemory::Mem), None).unwrap());
fill_ramp(&src);
let blob = export(&src, TransportMode::Inline).unwrap().0;
let got = import(&blob, &[]).unwrap();
assert_eq!(got.shape(), &[3, 8]);
assert_eq!(got.dtype(), crate::DType::F32);
assert_eq!(got.format(), None);
assert_eq!(bytes_of(&got), bytes_of(&src));
}
fn build_blob(planes: &[BlobPlane<'_>]) -> Vec<u8> {
build(&BlobSpec::new(planes))
}
struct BlobSpec<'a> {
planes: &'a [BlobPlane<'a>],
storage_kind: u32,
dtype: crate::DType,
grid: &'a [u64],
strings: BlobStrings<'a>,
quant_axis: i32,
scales: &'a [f32],
zeros: &'a [i32],
}
impl<'a> BlobSpec<'a> {
fn new(planes: &'a [BlobPlane<'a>]) -> Self {
Self {
planes,
storage_kind: 0,
dtype: crate::DType::U8,
grid: &[1],
strings: BlobStrings::default(),
quant_axis: -2,
scales: &[],
zeros: &[],
}
}
}
fn build(spec: &BlobSpec<'_>) -> Vec<u8> {
let mut planes_buf = Vec::new();
for p in spec.planes {
write_plane(p, &mut planes_buf);
}
let mut out = vec![0u8; HEADER_LEN];
for d in spec.grid {
out.extend_from_slice(&d.to_le_bytes());
}
for s in spec.scales {
out.extend_from_slice(&s.to_le_bytes());
}
for z in spec.zeros {
out.extend_from_slice(&z.to_le_bytes());
}
let strings_bytes = strings_encoded_len(&spec.strings, out.len());
write_strings(&spec.strings, &mut out);
out.extend_from_slice(&planes_buf);
let h = BlobHeader {
size: out.len() as u64,
planes_bytes: planes_buf.len() as u64,
storage_kind: spec.storage_kind,
dtype: spec.dtype.code(),
quant_axis: spec.quant_axis,
ndim: spec.grid.len() as u32,
plane_count: spec.planes.len() as u32,
quant_scales_len: spec.scales.len() as u32,
quant_zero_points_len: spec.zeros.len() as u32,
strings_bytes: strings_bytes as u32,
..BlobHeader::empty()
};
h.write_to(&mut out);
out
}
#[test]
fn no_corrupted_header_value_can_panic_the_parser() {
let good = export(&nv12_mem(64, 48), TransportMode::Inline).unwrap().0;
for word in 0..HEADER_LEN / 4 {
for pattern in [u32::MAX, u32::MAX / 2, 1, 0x8000_0000] {
let mut bad = good.clone();
bad[word * 4..word * 4 + 4].copy_from_slice(&pattern.to_le_bytes());
if let Ok(v) = BlobView::parse(&bad) {
let _ = v.shape();
let _ = v.strides();
let _ = v.strings();
let _ = v.planes();
}
let _ = import(&bad, &[]);
}
}
}
#[test]
fn no_corrupted_plane_record_can_panic_the_parser() {
let good = export(&nv12_mem(64, 48), TransportMode::Inline).unwrap().0;
let r = region_offsets(&BlobView::parse(&good).unwrap().header()).unwrap();
for word in 0..PLANE_RECORD_LEN / 4 {
for pattern in [u32::MAX, u32::MAX / 2, 0x8000_0000] {
let mut bad = good.clone();
let at = r.planes + word * 4;
bad[at..at + 4].copy_from_slice(&pattern.to_le_bytes());
if let Ok(v) = BlobView::parse(&bad) {
let _ = v.planes();
}
let _ = import(&bad, &[]);
}
}
}
#[test]
fn a_mixed_inline_and_referenced_plane_set_is_refused() {
let payload = [1u8, 2, 3, 4];
let referenced = BlobPlane {
handle: 7,
offset: 0,
stride: 4,
size: 4,
used: 4,
modifier: 0,
handle_bytes: &[],
data: &[],
};
let inlined = BlobPlane {
handle: -1,
offset: 0,
stride: 4,
size: payload.len() as u64,
used: payload.len() as u64,
modifier: 0,
handle_bytes: &[],
data: &payload,
};
let blob = build_blob(&[referenced, inlined]);
let err =
import(&blob, &[]).expect_err("a plane set mixing transport modes must be refused");
assert!(
err.to_string().contains("one transport mode"),
"must fail on the mode mix, not incidentally: {err}"
);
}
#[test]
fn a_uniformly_referenced_plane_set_is_not_rejected_as_mixed() {
let referenced = BlobPlane {
handle: 7,
offset: 0,
stride: 4,
size: 4,
used: 4,
modifier: 0,
handle_bytes: &[],
data: &[],
};
let blob = build_blob(&[referenced, referenced]);
let err = import(&blob, &[]).expect_err("an index into nothing cannot be opened");
let expected = if cfg!(target_os = "windows") {
"must all carry handle 0"
} else {
"indexes fd 7"
};
assert!(
err.to_string().contains(expected),
"a uniform plane set must clear the mode check: {err}"
);
}
#[test]
fn an_inline_plane_whose_size_disagrees_with_its_data_is_refused() {
let blob = export(&nv12_mem(64, 48), TransportMode::Inline).unwrap().0;
let mut tampered = blob.clone();
let r = region_offsets(&BlobView::parse(&blob).unwrap().header()).unwrap();
tampered[r.planes + 24..r.planes + 32].copy_from_slice(&9999u64.to_le_bytes());
assert!(import(&tampered, &[]).is_err());
}
#[test]
fn quantization_survives_the_round_trip() {
let mut src =
TensorDyn::from(Tensor::<i8>::new(&[3, 8], Some(TensorMemory::Mem), None).unwrap());
src.set_quantization(crate::Quantization::per_tensor(0.5, -7))
.unwrap();
let blob = export(&src, TransportMode::Inline).unwrap().0;
let got = import(&blob, &[]).unwrap();
let q = got.quantization().expect("quantization survived");
assert_eq!(q.scale(), &[0.5]);
assert_eq!(q.zero_point(), Some(&[-7i32][..]));
}
use crate::{CpuAccess, PixelFormat, Tensor, TensorDyn, TensorMemory};
fn nv12_mem(w: usize, h: usize) -> TensorDyn {
TensorDyn::from(
Tensor::<u8>::image(
w,
h,
PixelFormat::Nv12,
Some(TensorMemory::Mem),
CpuAccess::ReadWrite,
)
.expect("alloc"),
)
}
#[test]
fn export_carries_the_addressing_grid_not_the_allocation() {
let blob = export(&nv12_mem(640, 480), TransportMode::Inline)
.unwrap()
.0;
let v = BlobView::parse(&blob).unwrap();
assert_eq!(v.shape(), &[480, 640]);
assert_eq!(v.header().dtype, crate::DType::U8.code());
assert_eq!(v.strings().format, "NV12");
}
#[test]
fn export_emits_one_plane_per_plane_table_entry() {
let blob = export(&nv12_mem(640, 480), TransportMode::Inline)
.unwrap()
.0;
let v = BlobView::parse(&blob).unwrap();
let planes = v.planes().unwrap();
assert_eq!(planes.len(), 2, "NV12 is two planes");
assert_eq!(planes[0].size as usize, planes[0].data.len());
assert_eq!(planes[1].size as usize, planes[1].data.len());
assert_eq!(planes[0].size, 307_200);
assert_eq!(planes[1].size, 153_600);
assert_eq!(planes[1].offset, 0, "an inline plane carries no offset");
}
#[test]
fn inline_export_clears_pid_and_fence_fd() {
let blob = export(&nv12_mem(64, 48), TransportMode::Inline).unwrap().0;
let h = BlobView::parse(&blob).unwrap().header();
assert_eq!(h.pid, 0, "no handle to reopen");
assert_eq!(h.fence_fd, -1, "inlining means the producer already waited");
}
#[test]
fn inline_planes_carry_no_modifier_and_no_handle_bytes() {
let blob = export(&nv12_mem(64, 48), TransportMode::Inline).unwrap().0;
for p in BlobView::parse(&blob).unwrap().planes().unwrap() {
assert!(p.is_inline());
assert_eq!(p.modifier, 0);
assert!(p.handle_bytes.is_empty());
assert_eq!(p.size as usize, p.data.len());
}
}
#[test]
fn reference_export_of_a_mem_tensor_is_refused() {
assert!(export(&nv12_mem(64, 48), TransportMode::Reference).is_err());
}
#[test]
fn a_blob_parses_correctly_from_a_non_zero_offset() {
let blob = export(&nv12_mem(64, 48), TransportMode::Inline).unwrap().0;
let mut framed = vec![0xAAu8; 16];
framed.extend_from_slice(&blob);
let direct = BlobView::parse(&blob).unwrap();
let framed_view = BlobView::parse(&framed[16..]).unwrap();
assert_eq!(framed_view.header(), direct.header());
assert_eq!(framed_view.shape(), direct.shape());
assert_eq!(framed_view.strings(), direct.strings());
}
#[test]
fn colorimetry_axes_are_carried_as_strings() {
let t = nv12_mem(64, 48).with_colorimetry(crate::Colorimetry::jfif());
let blob = export(&t, TransportMode::Inline).unwrap().0;
let v = BlobView::parse(&blob).unwrap();
let s = v.strings();
assert_eq!(s.color_space, "srgb");
assert_eq!(s.color_transfer, "srgb");
assert_eq!(s.color_encoding, "bt601");
assert_eq!(s.color_range, "full");
}
#[test]
fn a_bare_tensor_carries_no_format_and_its_shape_is_its_grid() {
let t =
TensorDyn::from(Tensor::<f32>::new(&[3, 8], Some(TensorMemory::Mem), None).unwrap());
let blob = export(&t, TransportMode::Inline).unwrap().0;
let v = BlobView::parse(&blob).unwrap();
assert_eq!(v.strings().format, "", "not an image");
assert_eq!(v.shape(), &[3, 8]);
assert_eq!(v.strides(), &[32i64, 4], "byte strides for f32");
}
fn ref_plane() -> BlobPlane<'static> {
BlobPlane {
handle: 7,
offset: 0,
stride: 640,
size: 307_200,
used: 307_200,
modifier: 0,
handle_bytes: &[],
data: &[],
}
}
#[test]
fn a_plane_record_round_trips() {
let mut out = Vec::new();
write_plane(&ref_plane(), &mut out);
let got = parse_planes(&out, 1).unwrap();
assert_eq!(got.len(), 1);
assert_eq!(got[0], ref_plane());
}
#[test]
fn a_plane_scalar_block_is_56_bytes() {
let mut out = Vec::new();
write_plane(&ref_plane(), &mut out);
assert_eq!(out.len(), PLANE_RECORD_LEN);
assert_eq!(PLANE_RECORD_LEN, 56);
}
#[test]
fn every_plane_record_starts_8_aligned_whatever_the_payload() {
for n in 0..17usize {
let payload: Vec<u8> = (0..n as u8).collect();
let mut out = Vec::new();
write_plane(
&BlobPlane {
handle: -1,
size: n as u64,
data: &payload,
..ref_plane()
},
&mut out,
);
write_plane(&ref_plane(), &mut out);
assert_eq!(
out.len() % 8,
0,
"payload of {n} bytes left the region unaligned"
);
let got = parse_planes(&out, 2).unwrap();
assert_eq!(got[0].data, &payload[..], "payload {n} survived padding");
assert_eq!(
got[1],
ref_plane(),
"the second record parsed from an aligned offset"
);
}
}
#[test]
fn plane_payloads_round_trip_independently() {
let hb = [0xDEu8, 0xAD, 0xBE, 0xEF];
let data = [1u8, 2, 3];
let mut out = Vec::new();
write_plane(
&BlobPlane {
handle_bytes: &hb,
data: &data,
..ref_plane()
},
&mut out,
);
let got = parse_planes(&out, 1).unwrap();
assert_eq!(got[0].handle_bytes, &hb[..]);
assert_eq!(got[0].data, &data[..]);
}
#[test]
fn a_forged_plane_payload_length_is_refused() {
let mut out = vec![0u8; PLANE_RECORD_LEN];
out[48..52].copy_from_slice(&u32::MAX.to_le_bytes()); assert!(matches!(
parse_planes(&out, 1),
Err(BlobError::ForgedLength { .. })
));
}
#[test]
fn a_plane_count_larger_than_the_records_present_is_refused() {
let mut out = Vec::new();
write_plane(&ref_plane(), &mut out);
assert!(
parse_planes(&out, 2).is_err(),
"claims two records, holds one"
);
}
fn sample_strings() -> BlobStrings<'static> {
BlobStrings {
format: "NV12",
color_space: "bt709",
color_transfer: "bt709",
color_encoding: "bt709",
color_range: "limited",
}
}
#[test]
fn the_five_strings_round_trip_in_message_order() {
let mut out = Vec::new();
write_strings(&sample_strings(), &mut out);
let got = parse_strings(&out).unwrap();
assert_eq!(got, sample_strings());
}
#[test]
fn string_order_is_positional_and_matches_the_message() {
let s = BlobStrings {
format: "a",
color_space: "bb",
color_transfer: "ccc",
color_encoding: "dddd",
color_range: "eeeee",
};
let mut out = Vec::new();
write_strings(&s, &mut out);
assert_eq!(parse_strings(&out).unwrap(), s);
}
#[test]
fn empty_strings_are_legal_and_distinct_from_absent() {
let s = BlobStrings::default();
let mut out = Vec::new();
write_strings(&s, &mut out);
assert_eq!(parse_strings(&out).unwrap(), s);
assert!(
!out.is_empty(),
"five empty strings still occupy their counts"
);
}
#[test]
fn the_strings_region_is_padded_so_planes_start_8_aligned() {
for f in ["", "a", "ab", "abc", "NV12", "rgb8_planar"] {
let mut out = Vec::new();
write_strings(
&BlobStrings {
format: f,
..Default::default()
},
&mut out,
);
assert_eq!(out.len() % 8, 0, "region for format {f:?} is not 8-aligned");
}
}
#[test]
fn a_non_utf8_string_is_refused() {
let mut out = Vec::new();
out.extend_from_slice(&2u32.to_le_bytes());
out.extend_from_slice(&[0xff, 0xfe]);
for _ in 0..4 {
out.extend_from_slice(&0u32.to_le_bytes());
}
while out.len() % 8 != 0 {
out.push(0);
}
assert!(matches!(
parse_strings(&out),
Err(BlobError::InvalidUtf8(_))
));
}
#[test]
fn a_forged_string_length_cannot_read_past_the_region() {
let mut out = Vec::new();
out.extend_from_slice(&u32::MAX.to_le_bytes());
out.extend_from_slice(b"NV12");
while out.len() % 8 != 0 {
out.push(0);
}
assert!(matches!(
parse_strings(&out),
Err(BlobError::ForgedLength { .. })
));
}
#[test]
fn a_truncated_strings_region_is_refused() {
assert!(parse_strings(&[0u8; 4]).is_err(), "needs five counts");
}
#[test]
fn an_overrunning_region_is_attributed_to_its_own_field() {
let h = BlobHeader {
size: 4096,
ndim: 2,
strides_len: 2,
strings_bytes: u32::MAX,
..BlobHeader::empty()
};
assert!(
matches!(
region_offsets(&h),
Err(BlobError::ForgedLength {
field: "strings_bytes",
..
})
),
"the overrun must name strings_bytes, not a later field"
);
}
#[test]
fn header_is_72_bytes_and_every_scalar_is_naturally_aligned() {
assert_eq!(HEADER_LEN, 72);
assert_eq!(HEADER_LEN % 8, 0, "the variable tail starts 8-aligned");
for off in [0usize, 8, 16] {
assert_eq!(off % 8, 0, "u64 field at {off} is not 8-aligned");
}
for off in [24usize, 28, 32, 36, 40, 44, 48, 52, 56, 60, 64] {
assert_eq!(off % 4, 0, "u32 field at {off} is not 4-aligned");
}
}
#[test]
fn a_header_round_trips_through_bytes() {
let h = BlobHeader {
size: 72,
required_mask: 0,
planes_bytes: 0,
storage_kind: 2,
pid: 4242,
fence_fd: -1,
dtype: 1,
quant_axis: -2,
ndim: 2,
plane_count: 0,
quant_scales_len: 0,
quant_zero_points_len: 0,
strings_bytes: 0,
strides_len: 2,
};
let mut buf = vec![0u8; HEADER_LEN];
h.write_to(&mut buf);
assert_eq!(parse_header(&buf).unwrap(), h);
}
#[test]
fn the_header_is_little_endian_on_every_host() {
let mut buf = vec![0u8; HEADER_LEN];
BlobHeader {
size: 0x0102_0304_0506_0708,
..BlobHeader::empty()
}
.write_to(&mut buf);
assert_eq!(
&buf[0..8],
&[0x08, 0x07, 0x06, 0x05, 0x04, 0x03, 0x02, 0x01],
"size must be little-endian regardless of host byte order"
);
}
#[test]
fn a_truncated_buffer_is_refused_not_read() {
assert!(matches!(
parse_header(&[0u8; 71]),
Err(BlobError::Truncated { .. })
));
}
#[test]
fn an_unknown_required_bit_is_refused() {
let mut buf = vec![0u8; HEADER_LEN];
BlobHeader {
size: 72,
required_mask: 1 << 63,
..BlobHeader::empty()
}
.write_to(&mut buf);
assert!(matches!(
parse_header(&buf),
Err(BlobError::UnsupportedRequirement(_))
));
}
#[test]
fn a_size_smaller_than_the_header_is_refused() {
let mut buf = vec![0u8; HEADER_LEN];
BlobHeader {
size: 8,
..BlobHeader::empty()
}
.write_to(&mut buf);
assert!(
parse_header(&buf).is_err(),
"size must cover at least the header"
);
}
#[test]
fn blob_error_display_names_the_field_and_the_values() {
let cases = [
(
BlobError::Truncated { need: 80, have: 72 },
"blob truncated: needs 80 bytes, buffer has 72",
),
(
BlobError::UnsupportedRequirement(0x10),
"blob requires unsupported features (required_mask 0x10); refusing rather \
than misreading it",
),
(
BlobError::ForgedLength {
field: "ndim",
value: 9,
},
"blob field `ndim` has implausible value 9",
),
(
BlobError::InvalidUtf8("format"),
"blob string `format` is not valid UTF-8",
),
(
BlobError::Malformed("no planes"),
"malformed blob: no planes",
),
];
for (err, want) in cases {
assert_eq!(err.to_string(), want);
}
}
#[test]
fn a_size_past_the_end_of_the_buffer_is_truncation() {
let mut buf = vec![0u8; HEADER_LEN];
BlobHeader {
size: 80,
..BlobHeader::empty()
}
.write_to(&mut buf);
assert_eq!(
parse_header(&buf),
Err(BlobError::Truncated { need: 80, have: 72 })
);
}
#[test]
fn bytes_after_the_blob_are_not_part_of_it() {
let blob = export(&nv12_mem(64, 48), TransportMode::Inline).unwrap().0;
let mut framed = blob.clone();
framed.extend_from_slice(&[0xAA; 16]);
let v = BlobView::parse(&framed).expect("a buffer longer than `size` parses");
assert_eq!(v.header().size, blob.len() as u64);
assert_eq!(v.planes().unwrap().len(), 2);
}
#[test]
fn a_header_only_blob_has_every_region_empty_and_ending_at_the_header() {
let r = region_offsets(&BlobHeader {
size: HEADER_LEN as u64,
..BlobHeader::empty()
})
.expect("a blob of exactly one header is well-formed");
assert_eq!(
r,
Regions {
shape: HEADER_LEN,
strides: HEADER_LEN,
quant_scales: HEADER_LEN,
quant_zero_points: HEADER_LEN,
strings: HEADER_LEN,
planes: HEADER_LEN,
end: HEADER_LEN,
}
);
}
#[test]
fn a_size_below_the_header_is_attributed_to_size() {
assert_eq!(
region_offsets(&BlobHeader {
size: 8,
..BlobHeader::empty()
}),
Err(BlobError::ForgedLength {
field: "size",
value: 8
})
);
}
#[test]
fn plane_count_is_bounded_by_the_whole_records_that_fit() {
let header = |plane_count| BlobHeader {
size: 4096,
planes_bytes: 2 * PLANE_RECORD_LEN as u64,
plane_count,
..BlobHeader::empty()
};
assert!(
region_offsets(&header(2)).is_ok(),
"two minimum-size records fit in {} bytes",
2 * PLANE_RECORD_LEN
);
assert_eq!(
region_offsets(&header(3)),
Err(BlobError::ForgedLength {
field: "plane_count",
value: 3
})
);
}
#[test]
fn the_strings_region_ends_8_aligned_in_the_blob_wherever_it_starts() {
let formats = ["", "a", "abc", "NV12", "rgb8_planar"].map(|format| BlobStrings {
format,
..Default::default()
});
let all_five = BlobStrings {
format: "a",
color_space: "bb",
color_transfer: "cccc",
color_encoding: "dddddddd",
color_range: "eeeeeeeeeeeeeeee",
};
for start in [0usize, 3, 4] {
for strings in formats.iter().chain([&all_five]) {
let mut out = vec![0xAAu8; start];
write_strings(strings, &mut out);
assert_eq!(out.len() % 8, 0, "start {start}, {strings:?}");
assert_eq!(
out.len() - start,
strings_encoded_len(strings, start),
"start {start}, {strings:?}"
);
assert_eq!(parse_strings(&out[start..]).unwrap(), *strings);
}
}
}
#[test]
fn an_export_whose_quantization_ends_4_aligned_still_parses() {
let mut src =
TensorDyn::from(Tensor::<i8>::new(&[3, 8], Some(TensorMemory::Mem), None).unwrap());
src.set_quantization(crate::Quantization::per_tensor_symmetric(0.5))
.unwrap();
let blob = export(&src, TransportMode::Inline).unwrap().0;
let v = BlobView::parse(&blob).expect("export must write a blob its parser accepts");
assert_eq!(v.planes().unwrap().len(), 1);
}
#[test]
fn encoded_len_is_exactly_what_write_plane_appends() {
let handle_bytes = [0xDEu8, 0xAD, 0xBE, 0xEF];
for hb in [&[][..], &handle_bytes[..]] {
for n in 0..17usize {
let data: Vec<u8> = (0..n as u8).collect();
let plane = BlobPlane {
handle_bytes: hb,
data: &data,
..ref_plane()
};
let mut out = vec![0xAAu8; 3];
write_plane(&plane, &mut out);
assert_eq!(
out.len() - 3,
plane.encoded_len(),
"{} handle bytes, {n} data bytes",
hb.len()
);
}
}
}
#[test]
fn a_plane_is_inline_only_when_its_handle_is_negative() {
assert!(
!BlobPlane {
handle: 0,
..ref_plane()
}
.is_inline(),
"index 0 is the first fd-table entry, a reference"
);
assert!(BlobPlane {
handle: -1,
..ref_plane()
}
.is_inline());
}
#[test]
fn a_reference_plane_keeps_its_offset_and_extent() {
let geoms = [crate::PlaneGeometry {
offset: 4096,
stride: 64,
size: 3072,
}];
let mut buf = Vec::new();
write_export_planes(&geoms, TransportMode::Reference, 0, &[], &[], &mut buf).unwrap();
let planes = parse_planes(&buf, 1).unwrap();
let p = planes[0];
assert_eq!(
(p.handle, p.offset, p.stride, p.size, p.used),
(0, 4096, 64, 3072, 3072)
);
assert!(p.data.is_empty(), "a reference carries no bytes");
}
#[test]
fn every_layout_round_trips_inline_through_its_grid() {
for (format, grid) in [
(PixelFormat::Rgb, vec![48u64, 64, 3]),
(PixelFormat::Grey, vec![48, 64]),
(PixelFormat::PlanarRgb, vec![3, 48, 64]),
] {
let src = TensorDyn::from(
Tensor::<u8>::image(
64,
48,
format,
Some(TensorMemory::Mem),
CpuAccess::ReadWrite,
)
.expect("alloc"),
);
fill_ramp(&src);
let blob = export(&src, TransportMode::Inline).unwrap().0;
assert_eq!(BlobView::parse(&blob).unwrap().shape(), grid, "{format:?}");
let got = import(&blob, &[]).unwrap();
assert_eq!(got.format(), Some(format));
assert_eq!(got.shape(), src.shape(), "{format:?}");
assert_eq!(bytes_of(&got), bytes_of(&src), "{format:?}");
}
}
fn inline_plane(data: &[u8]) -> BlobPlane<'_> {
BlobPlane {
handle: -1,
offset: 0,
stride: data.len() as u64,
size: data.len() as u64,
used: data.len() as u64,
modifier: 0,
handle_bytes: &[],
data,
}
}
#[test]
fn a_well_formed_hand_built_inline_blob_imports() {
let got = import(&build_blob(&[inline_plane(&[7])]), &[]).expect("well-formed");
assert_eq!(got.shape(), &[1]);
assert_eq!(bytes_of(&got), vec![7]);
}
#[test]
fn an_inline_plane_with_a_modifier_or_handle_bytes_is_refused() {
let data = [7u8];
for plane in [
BlobPlane {
modifier: 1,
..inline_plane(&data)
},
BlobPlane {
handle_bytes: &[1],
..inline_plane(&data)
},
] {
let err = import(&build_blob(&[plane]), &[]).expect_err("must be refused");
assert!(
err.to_string()
.contains("inline plane must have modifier 0 and no handle_bytes"),
"got: {err}"
);
}
}
#[test]
fn a_referenced_plane_carrying_bytes_is_refused() {
let plane = BlobPlane {
handle: 0,
..inline_plane(&[7])
};
let err = import(&build_blob(&[plane]), &[]).expect_err("must be refused");
assert!(
err.to_string()
.contains("referenced plane must not also carry inline bytes"),
"got: {err}"
);
}
#[test]
fn a_grid_too_short_for_its_format_is_refused() {
for (format, grid) in [
(PixelFormat::Nv12, &[1u64][..]),
(PixelFormat::PlanarRgb, &[48, 64][..]),
] {
let data = [0u8];
let planes = [inline_plane(&data)];
let blob = build(&BlobSpec {
grid,
strings: BlobStrings {
format: format.as_str(),
..Default::default()
},
..BlobSpec::new(&planes)
});
let err = import(&blob, &[]).expect_err("must be refused");
let want = format!("grid {grid:?} is not an image for {format:?}");
assert!(
matches!(&err, crate::Error::InvalidShape(m) if m.contains(&want)),
"{format:?}: got {err:?}"
);
}
}
#[test]
fn an_unquantized_export_declares_no_quantization() {
let src =
TensorDyn::from(Tensor::<i8>::new(&[3, 8], Some(TensorMemory::Mem), None).unwrap());
let blob = export(&src, TransportMode::Inline).unwrap().0;
let h = BlobView::parse(&blob).unwrap().header();
assert_eq!(h.quant_axis, -2);
assert_eq!((h.quant_scales_len, h.quant_zero_points_len), (0, 0));
assert!(import(&blob, &[]).unwrap().quantization().is_none());
}
fn round_trip_quantization(q: crate::Quantization) -> crate::Quantization {
let mut src =
TensorDyn::from(Tensor::<i8>::new(&[2, 2, 3], Some(TensorMemory::Mem), None).unwrap());
src.set_quantization(q).unwrap();
let blob = export(&src, TransportMode::Inline).unwrap().0;
import(&blob, &[])
.unwrap()
.quantization()
.expect("quantization survived")
.clone()
}
#[test]
fn per_tensor_symmetric_quantization_survives_without_a_zero_point() {
let q = round_trip_quantization(crate::Quantization::per_tensor_symmetric(0.5));
assert_eq!(q.scale(), &[0.5]);
assert_eq!(q.zero_point(), None);
assert_eq!(q.axis(), None);
}
#[test]
fn per_channel_symmetric_quantization_survives_with_its_axis() {
let q = round_trip_quantization(
crate::Quantization::per_channel_symmetric(vec![0.1, 0.2, 0.3], 2).unwrap(),
);
assert_eq!(q.scale(), &[0.1, 0.2, 0.3]);
assert_eq!(q.zero_point(), None);
assert_eq!(q.axis(), Some(2));
}
#[test]
fn per_channel_asymmetric_quantization_keeps_its_zero_points() {
let q = round_trip_quantization(
crate::Quantization::per_channel(vec![0.1, 0.2, 0.3], vec![-3, 0, 7], 2).unwrap(),
);
assert_eq!(q.scale(), &[0.1, 0.2, 0.3]);
assert_eq!(q.zero_point(), Some(&[-3, 0, 7][..]));
assert_eq!(q.axis(), Some(2));
}
#[test]
fn declared_but_malformed_quantization_is_refused() {
let data = [0u8; 3];
let planes = [inline_plane(&data)];
let cases: [(i32, &[f32], &[i32], &str); 4] = [
(-3, &[0.5], &[], "quant_axis -3 is neither"),
(
-1,
&[],
&[],
"needs 1 scale and 0 or 1 zero points, got 0 and 0",
),
(-1, &[0.5, 0.25], &[], "got 2 and 0"),
(0, &[0.5, 0.25, 0.125], &[1, 2], "zero_point.len"),
];
for (quant_axis, scales, zeros, want) in cases {
let blob = build(&BlobSpec {
dtype: crate::DType::I8,
grid: &[3],
quant_axis,
scales,
zeros,
..BlobSpec::new(&planes)
});
let err = import(&blob, &[]).expect_err("malformed quantization must be refused");
let msg = err.to_string();
assert!(
msg.contains("blob quantization is malformed") && msg.contains(want),
"axis {quant_axis}, {} scales, {} zeros: got {msg}",
scales.len(),
zeros.len()
);
}
}
#[cfg(all(unix, not(any(target_os = "macos", target_os = "ios"))))]
mod shm_reference {
use super::*;
use std::os::fd::AsRawFd;
fn reference_planes(
format: PixelFormat,
w: usize,
h: usize,
pitch: usize,
) -> Vec<BlobPlane<'static>> {
format
.plane_table(w, h, pitch)
.unwrap()
.iter()
.map(|g| BlobPlane {
handle: 0,
offset: g.offset,
stride: g.stride,
size: g.size,
used: g.size,
modifier: 0,
handle_bytes: &[],
data: &[],
})
.collect()
}
fn nv12_spec<'a>(planes: &'a [BlobPlane<'a>]) -> BlobSpec<'a> {
BlobSpec {
storage_kind: TensorMemory::Shm.code(),
grid: &[48, 64],
strings: BlobStrings {
format: "NV12",
..Default::default()
},
..BlobSpec::new(planes)
}
}
fn wide_shm() -> TensorDyn {
TensorDyn::from(
Tensor::<u8>::image(
128,
48,
PixelFormat::Nv12,
Some(TensorMemory::Shm),
CpuAccess::ReadWrite,
)
.expect("shm allocation"),
)
}
#[test]
fn a_referenced_shm_blob_aliases_its_producer_at_the_producers_pitch() {
let src = wide_shm();
fill_ramp(&src);
let fd = src.clone_fd().expect("shm fd");
let planes = reference_planes(PixelFormat::Nv12, 64, 48, 128);
let got = import(&build(&nv12_spec(&planes)), &[fd.as_raw_fd()])
.expect("a Shm reference imports");
assert_eq!(got.memory(), TensorMemory::Shm);
assert_eq!(got.format(), Some(PixelFormat::Nv12));
assert_eq!(got.shape(), &[72, 64], "the allocation shape of the grid");
assert_eq!(
got.row_stride(),
Some(128),
"the producer's pitch, which the shape cannot express"
);
let expected = bytes_of(&src);
let seen = bytes_of(&got);
assert_eq!(
seen[..],
expected[..seen.len()],
"the same bytes, not a copy"
);
}
#[test]
fn a_zero_pitch_means_none_was_recorded() {
let src = wide_shm();
let fd = src.clone_fd().expect("shm fd");
let mut planes = reference_planes(PixelFormat::Nv12, 64, 48, 128);
planes[0].stride = 0;
let got = import(&build(&nv12_spec(&planes)), &[fd.as_raw_fd()])
.expect("a zero pitch is not an error");
assert_ne!(got.row_stride(), Some(0));
}
#[test]
fn a_pitch_the_format_refuses_is_refused_not_ignored() {
let src = wide_shm();
let fd = src.clone_fd().expect("shm fd");
let mut planes = reference_planes(PixelFormat::Nv12, 64, 48, 128);
planes[0].stride = 32;
let err = import(&build(&nv12_spec(&planes)), &[fd.as_raw_fd()])
.expect_err("a pitch narrower than a row is malformed");
assert!(
err.to_string().contains("row_stride 32 < minimum 64"),
"got: {err}"
);
}
#[test]
fn every_plane_index_is_bounded_not_only_the_first() {
let src = wide_shm();
let fd = src.clone_fd().expect("shm fd");
let mut planes = reference_planes(PixelFormat::Nv12, 64, 48, 128);
planes[1].handle = 5;
let err = import(&build(&nv12_spec(&planes)), &[fd.as_raw_fd()])
.expect_err("index 5 is past a one-entry table");
assert!(
err.to_string()
.contains("plane handle indexes fd 5 but only 1 were provided"),
"got: {err}"
);
}
#[test]
fn quantization_survives_a_reference_round_trip() {
let src = TensorDyn::from(
Tensor::<i8>::new(&[2, 2, 3], Some(TensorMemory::Shm), None).expect("shm"),
);
let fd = src.clone_fd().expect("shm fd");
let bytes = src.capacity_bytes() as u64;
let planes = [BlobPlane {
handle: 0,
offset: 0,
stride: bytes,
size: bytes,
used: bytes,
modifier: 0,
handle_bytes: &[],
data: &[],
}];
let blob = build(&BlobSpec {
storage_kind: TensorMemory::Shm.code(),
dtype: crate::DType::I8,
grid: &[2, 2, 3],
quant_axis: 2,
scales: &[0.1, 0.2, 0.3],
zeros: &[-3, 0, 7],
..BlobSpec::new(&planes)
});
let got = import(&blob, &[fd.as_raw_fd()]).expect("a Shm reference imports");
let q = got.quantization().expect("quantization survived");
assert_eq!(q.scale(), &[0.1, 0.2, 0.3]);
assert_eq!(q.zero_point(), Some(&[-3, 0, 7][..]));
assert_eq!(q.axis(), Some(2));
}
}
#[cfg(target_os = "linux")]
#[test]
fn a_padded_planar_reference_export_steps_planes_by_the_pitch() {
const WHAT: &str = "a_padded_planar_reference_export_steps_planes_by_the_pitch";
if !crate::test_support::dma_or_skip(WHAT) {
return;
}
let src = TensorDyn::from(
Tensor::<u8>::image(
50,
48,
PixelFormat::PlanarRgb,
Some(TensorMemory::DmaBuf),
CpuAccess::ReadWrite,
)
.expect("a usable DMA heap allocates a planar image"),
);
let pitch = src.effective_row_stride().expect("an image has a pitch");
assert!(pitch > 50, "precondition: the pitch pads, got {pitch}");
let (blob, _fds) = export(&src, TransportMode::Reference).unwrap();
let v = BlobView::parse(&blob).unwrap();
let planes = v.planes().unwrap();
assert_eq!(
v.strides(),
&[(pitch * 48) as i64, pitch as i64, 1],
"row stride is the pitch; plane stride is pitch * height"
);
assert_eq!(
planes[1].offset as i64,
v.strides()[0],
"the stride array and the plane table agree on where plane 1 starts"
);
}
}