1pub const HEADER_LEN: usize = 72;
59
60pub const PLANE_RECORD_LEN: usize = 56;
65
66pub const SUPPORTED_REQUIRED_MASK: u64 = 0;
70
71pub const HEADER_PID_OFFSET: usize = 28;
76
77pub const D3D11_HANDLE_BYTES: usize = 24;
86
87#[cfg(unix)]
92type BlobFd = std::os::fd::RawFd;
93#[cfg(not(unix))]
94type BlobFd = i32;
95
96#[derive(Debug, Clone, PartialEq, Eq)]
103pub enum BlobError {
104 Truncated {
106 need: usize,
108 have: usize,
110 },
111 UnsupportedRequirement(u64),
115 ForgedLength {
118 field: &'static str,
120 value: u64,
122 },
123 InvalidUtf8(&'static str),
125 Malformed(&'static str),
127}
128
129impl std::fmt::Display for BlobError {
130 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
131 match self {
132 Self::Truncated { need, have } => {
133 write!(f, "blob truncated: needs {need} bytes, buffer has {have}")
134 }
135 Self::UnsupportedRequirement(m) => write!(
136 f,
137 "blob requires unsupported features (required_mask {m:#x}); \
138 refusing rather than misreading it"
139 ),
140 Self::ForgedLength { field, value } => {
141 write!(f, "blob field `{field}` has implausible value {value}")
142 }
143 Self::InvalidUtf8(which) => write!(f, "blob string `{which}` is not valid UTF-8"),
144 Self::Malformed(why) => write!(f, "malformed blob: {why}"),
145 }
146 }
147}
148
149impl std::error::Error for BlobError {}
150
151#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
157pub struct BlobHeader {
158 pub size: u64,
162 pub required_mask: u64,
165 pub planes_bytes: u64,
167 pub storage_kind: u32,
169 pub pid: u32,
182 pub fence_fd: i32,
184 pub dtype: u32,
186 pub quant_axis: i32,
188 pub ndim: u32,
190 pub plane_count: u32,
192 pub quant_scales_len: u32,
194 pub quant_zero_points_len: u32,
196 pub strings_bytes: u32,
198 pub strides_len: u32,
204}
205
206impl BlobHeader {
207 #[cfg_attr(test, mutants::skip)]
210 pub fn empty() -> Self {
211 Self::default()
212 }
213
214 pub fn write_to(&self, buf: &mut [u8]) {
219 assert!(buf.len() >= HEADER_LEN, "header buffer too small");
220 buf[0..8].copy_from_slice(&self.size.to_le_bytes());
221 buf[8..16].copy_from_slice(&self.required_mask.to_le_bytes());
222 buf[16..24].copy_from_slice(&self.planes_bytes.to_le_bytes());
223 buf[24..28].copy_from_slice(&self.storage_kind.to_le_bytes());
224 buf[28..32].copy_from_slice(&self.pid.to_le_bytes());
225 buf[32..36].copy_from_slice(&self.fence_fd.to_le_bytes());
226 buf[36..40].copy_from_slice(&self.dtype.to_le_bytes());
227 buf[40..44].copy_from_slice(&self.quant_axis.to_le_bytes());
228 buf[44..48].copy_from_slice(&self.ndim.to_le_bytes());
229 buf[48..52].copy_from_slice(&self.plane_count.to_le_bytes());
230 buf[52..56].copy_from_slice(&self.quant_scales_len.to_le_bytes());
231 buf[56..60].copy_from_slice(&self.quant_zero_points_len.to_le_bytes());
232 buf[60..64].copy_from_slice(&self.strings_bytes.to_le_bytes());
233 buf[64..68].copy_from_slice(&self.strides_len.to_le_bytes());
234 buf[68..72].copy_from_slice(&0u32.to_le_bytes()); }
236}
237
238fn rd_u64(buf: &[u8], off: usize) -> u64 {
240 let mut b = [0u8; 8];
241 b.copy_from_slice(&buf[off..off + 8]);
242 u64::from_le_bytes(b)
243}
244
245fn rd_u32(buf: &[u8], off: usize) -> u32 {
247 let mut b = [0u8; 4];
248 b.copy_from_slice(&buf[off..off + 4]);
249 u32::from_le_bytes(b)
250}
251
252pub fn parse_header(buf: &[u8]) -> Result<BlobHeader, BlobError> {
263 if buf.len() < HEADER_LEN {
264 return Err(BlobError::Truncated {
265 need: HEADER_LEN,
266 have: buf.len(),
267 });
268 }
269 let required_mask = rd_u64(buf, 8);
270 if required_mask & !SUPPORTED_REQUIRED_MASK != 0 {
274 return Err(BlobError::UnsupportedRequirement(required_mask));
275 }
276 let size = rd_u64(buf, 0);
277 if size < HEADER_LEN as u64 {
278 return Err(BlobError::ForgedLength {
279 field: "size",
280 value: size,
281 });
282 }
283 if size > buf.len() as u64 {
286 return Err(BlobError::Truncated {
287 need: size as usize,
288 have: buf.len(),
289 });
290 }
291 if rd_u32(buf, 68) != 0 {
294 return Err(BlobError::Malformed("reserved header word is not zero"));
295 }
296 Ok(BlobHeader {
297 size,
298 required_mask,
299 planes_bytes: rd_u64(buf, 16),
300 storage_kind: rd_u32(buf, 24),
301 pid: rd_u32(buf, 28),
302 fence_fd: rd_u32(buf, 32) as i32,
303 dtype: rd_u32(buf, 36),
304 quant_axis: rd_u32(buf, 40) as i32,
305 ndim: rd_u32(buf, 44),
306 plane_count: rd_u32(buf, 48),
307 quant_scales_len: rd_u32(buf, 52),
308 quant_zero_points_len: rd_u32(buf, 56),
309 strings_bytes: rd_u32(buf, 60),
310 strides_len: rd_u32(buf, 64),
311 })
312}
313
314#[derive(Debug, Clone, Copy, PartialEq, Eq)]
319pub struct Regions {
320 pub shape: usize,
322 pub strides: usize,
324 pub quant_scales: usize,
326 pub quant_zero_points: usize,
328 pub strings: usize,
330 pub planes: usize,
332 pub end: usize,
334}
335
336fn advance(
343 off: usize,
344 count: u32,
345 elem: usize,
346 limit: usize,
347 field: &'static str,
348) -> Result<usize, BlobError> {
349 let bytes = (count as usize)
350 .checked_mul(elem)
351 .ok_or(BlobError::ForgedLength {
352 field,
353 value: count as u64,
354 })?;
355 let next = off.checked_add(bytes).ok_or(BlobError::ForgedLength {
356 field,
357 value: count as u64,
358 })?;
359 if next > limit {
360 return Err(BlobError::ForgedLength {
361 field,
362 value: count as u64,
363 });
364 }
365 Ok(next)
366}
367
368pub fn region_offsets(h: &BlobHeader) -> Result<Regions, BlobError> {
376 let limit = h.size as usize;
377 if limit < HEADER_LEN {
378 return Err(BlobError::ForgedLength {
379 field: "size",
380 value: h.size,
381 });
382 }
383 if h.strides_len != 0 && h.strides_len != h.ndim {
386 return Err(BlobError::ForgedLength {
387 field: "strides_len",
388 value: h.strides_len as u64,
389 });
390 }
391 let shape = HEADER_LEN;
392 let strides = advance(shape, h.ndim, 8, limit, "ndim")?;
393 let quant_scales = advance(strides, h.strides_len, 8, limit, "strides_len")?;
394 let quant_zero_points = advance(
395 quant_scales,
396 h.quant_scales_len,
397 4,
398 limit,
399 "quant_scales_len",
400 )?;
401 let strings = advance(
402 quant_zero_points,
403 h.quant_zero_points_len,
404 4,
405 limit,
406 "quant_zero_points_len",
407 )?;
408 let planes = advance(strings, h.strings_bytes, 1, limit, "strings_bytes")?;
409 if !planes.is_multiple_of(8) {
413 return Err(BlobError::Malformed(
414 "planes region is not 8-aligned; strings_bytes must include padding",
415 ));
416 }
417 let planes_bytes = usize::try_from(h.planes_bytes).map_err(|_| BlobError::ForgedLength {
418 field: "planes_bytes",
419 value: h.planes_bytes,
420 })?;
421 let end = planes
422 .checked_add(planes_bytes)
423 .ok_or(BlobError::ForgedLength {
424 field: "planes_bytes",
425 value: h.planes_bytes,
426 })?;
427 if end > limit {
428 return Err(BlobError::ForgedLength {
429 field: "planes_bytes",
430 value: h.planes_bytes,
431 });
432 }
433 let max_planes = planes_bytes / PLANE_RECORD_LEN;
437 if h.plane_count as usize > max_planes {
438 return Err(BlobError::ForgedLength {
439 field: "plane_count",
440 value: h.plane_count as u64,
441 });
442 }
443 Ok(Regions {
444 shape,
445 strides,
446 quant_scales,
447 quant_zero_points,
448 strings,
449 planes,
450 end,
451 })
452}
453
454#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
465pub struct BlobStrings<'a> {
466 pub format: &'a str,
468 pub color_space: &'a str,
470 pub color_transfer: &'a str,
472 pub color_encoding: &'a str,
474 pub color_range: &'a str,
476}
477
478const STRING_COUNT: usize = 5;
480
481pub fn strings_encoded_len(s: &BlobStrings<'_>, start: usize) -> usize {
487 let raw: usize = STRING_COUNT * 4
488 + s.format.len()
489 + s.color_space.len()
490 + s.color_transfer.len()
491 + s.color_encoding.len()
492 + s.color_range.len();
493 (start + raw).next_multiple_of(8) - start
494}
495
496pub fn write_strings(s: &BlobStrings<'_>, out: &mut Vec<u8>) {
504 for part in [
505 s.format,
506 s.color_space,
507 s.color_transfer,
508 s.color_encoding,
509 s.color_range,
510 ] {
511 out.extend_from_slice(&(part.len() as u32).to_le_bytes());
512 out.extend_from_slice(part.as_bytes());
513 }
514 while !out.len().is_multiple_of(8) {
515 out.push(0);
516 }
517}
518
519pub fn parse_strings(buf: &[u8]) -> Result<BlobStrings<'_>, BlobError> {
523 const NAMES: [&str; STRING_COUNT] = [
524 "format",
525 "color_space",
526 "color_transfer",
527 "color_encoding",
528 "color_range",
529 ];
530 let mut out: [&str; STRING_COUNT] = [""; STRING_COUNT];
531 let mut off = 0usize;
532 for (i, name) in NAMES.iter().enumerate() {
533 let name = *name;
534 let end = off
535 .checked_add(4)
536 .ok_or(BlobError::Malformed("string count overflows"))?;
537 if end > buf.len() {
538 return Err(BlobError::Truncated {
539 need: end,
540 have: buf.len(),
541 });
542 }
543 let len = rd_u32(buf, off) as usize;
544 off = end;
545 let stop = off.checked_add(len).ok_or(BlobError::ForgedLength {
546 field: name,
547 value: len as u64,
548 })?;
549 if stop > buf.len() {
550 return Err(BlobError::ForgedLength {
551 field: name,
552 value: len as u64,
553 });
554 }
555 out[i] = std::str::from_utf8(&buf[off..stop]).map_err(|_| BlobError::InvalidUtf8(name))?;
558 off = stop;
559 }
560 Ok(BlobStrings {
561 format: out[0],
562 color_space: out[1],
563 color_transfer: out[2],
564 color_encoding: out[3],
565 color_range: out[4],
566 })
567}
568
569#[derive(Debug, Clone, Copy, PartialEq, Eq)]
584pub struct BlobPlane<'a> {
585 pub handle: i64,
605 pub offset: u64,
607 pub stride: u64,
609 pub size: u64,
611 pub used: u64,
613 pub modifier: u64,
616 pub handle_bytes: &'a [u8],
628 pub data: &'a [u8],
630}
631
632impl BlobPlane<'_> {
633 pub fn is_inline(&self) -> bool {
635 self.handle < 0
636 }
637
638 pub fn encoded_len(&self) -> usize {
641 (PLANE_RECORD_LEN + self.handle_bytes.len() + self.data.len()).next_multiple_of(8)
642 }
643}
644
645pub fn write_plane(p: &BlobPlane<'_>, out: &mut Vec<u8>) {
647 let start = out.len();
648 out.extend_from_slice(&p.handle.to_le_bytes());
649 out.extend_from_slice(&p.offset.to_le_bytes());
650 out.extend_from_slice(&p.stride.to_le_bytes());
651 out.extend_from_slice(&p.size.to_le_bytes());
652 out.extend_from_slice(&p.used.to_le_bytes());
653 out.extend_from_slice(&p.modifier.to_le_bytes());
654 out.extend_from_slice(&(p.handle_bytes.len() as u32).to_le_bytes());
655 out.extend_from_slice(&(p.data.len() as u32).to_le_bytes());
656 out.extend_from_slice(p.handle_bytes);
657 out.extend_from_slice(p.data);
658 while !(out.len() - start).is_multiple_of(8) {
659 out.push(0);
660 }
661}
662
663fn rd_i64(buf: &[u8], off: usize) -> i64 {
665 let mut b = [0u8; 8];
666 b.copy_from_slice(&buf[off..off + 8]);
667 i64::from_le_bytes(b)
668}
669
670pub fn parse_planes(buf: &[u8], count: u32) -> Result<Vec<BlobPlane<'_>>, BlobError> {
677 let mut planes = Vec::with_capacity(count as usize);
678 let mut off = 0usize;
679 for _ in 0..count {
680 let head = off
681 .checked_add(PLANE_RECORD_LEN)
682 .ok_or(BlobError::Malformed("plane record offset overflows"))?;
683 if head > buf.len() {
684 return Err(BlobError::Truncated {
685 need: head,
686 have: buf.len(),
687 });
688 }
689 let hb_len = rd_u32(buf, off + 48) as usize;
690 let data_len = rd_u32(buf, off + 52) as usize;
691 let hb_end = head.checked_add(hb_len).ok_or(BlobError::ForgedLength {
692 field: "handle_bytes_len",
693 value: hb_len as u64,
694 })?;
695 let data_end = hb_end
696 .checked_add(data_len)
697 .ok_or(BlobError::ForgedLength {
698 field: "data_len",
699 value: data_len as u64,
700 })?;
701 if hb_end > buf.len() {
702 return Err(BlobError::ForgedLength {
703 field: "handle_bytes_len",
704 value: hb_len as u64,
705 });
706 }
707 if data_end > buf.len() {
708 return Err(BlobError::ForgedLength {
709 field: "data_len",
710 value: data_len as u64,
711 });
712 }
713 planes.push(BlobPlane {
714 handle: rd_i64(buf, off),
715 offset: rd_u64(buf, off + 8),
716 stride: rd_u64(buf, off + 16),
717 size: rd_u64(buf, off + 24),
718 used: rd_u64(buf, off + 32),
719 modifier: rd_u64(buf, off + 40),
720 handle_bytes: &buf[head..hb_end],
721 data: &buf[hb_end..data_end],
722 });
723 off = data_end.next_multiple_of(8);
724 }
725 Ok(planes)
726}
727
728#[derive(Debug, Clone, Copy, PartialEq, Eq)]
734pub enum TransportMode {
735 Reference,
738 Inline,
742}
743
744#[derive(Debug, Clone, Copy)]
748pub struct BlobView<'a> {
749 buf: &'a [u8],
750 header: BlobHeader,
751 regions: Regions,
752}
753
754impl<'a> BlobView<'a> {
755 pub fn parse(buf: &'a [u8]) -> Result<Self, BlobError> {
761 let header = parse_header(buf)?;
762 let regions = region_offsets(&header)?;
763 Ok(Self {
764 buf,
765 header,
766 regions,
767 })
768 }
769
770 pub fn header(&self) -> BlobHeader {
772 self.header
773 }
774
775 pub fn shape(&self) -> Vec<u64> {
777 (0..self.header.ndim as usize)
778 .map(|i| rd_u64(self.buf, self.regions.shape + i * 8))
779 .collect()
780 }
781
782 pub fn strides(&self) -> Vec<i64> {
784 (0..self.header.strides_len as usize)
785 .map(|i| rd_i64(self.buf, self.regions.strides + i * 8))
786 .collect()
787 }
788
789 pub fn strings(&self) -> BlobStrings<'a> {
791 parse_strings(&self.buf[self.regions.strings..self.regions.planes]).unwrap_or_default()
792 }
793
794 pub fn planes(&self) -> Result<Vec<BlobPlane<'a>>, BlobError> {
796 parse_planes(
797 &self.buf[self.regions.planes..self.regions.end],
798 self.header.plane_count,
799 )
800 }
801}
802
803pub fn export(t: &crate::TensorDyn, mode: TransportMode) -> crate::Result<(Vec<u8>, Vec<BlobFd>)> {
820 let fmt = t.format();
821 let shape = export_addressing_shape(t, fmt)?;
824 let esz = t.dtype().size() as i64;
825 let strides =
826 crate::protocol::c_byte_strides(&shape, esz, fmt.map(|f| f.layout()), t.row_stride());
827
828 let quant = t.quantization();
829 let (quant_axis, scales, zeros): (i32, &[f32], &[i32]) = match quant {
830 None => (-2, &[], &[]),
831 Some(q) => (
832 q.axis().map(|a| a as i32).unwrap_or(-1),
833 q.scale(),
834 q.zero_point().unwrap_or(&[]),
835 ),
836 };
837
838 let colorimetry = t.colorimetry().unwrap_or_default();
839 let strings = BlobStrings {
840 format: fmt.map(|f| f.as_str()).unwrap_or(""),
841 color_space: colorimetry.space.map(|v| v.as_str()).unwrap_or(""),
842 color_transfer: colorimetry.transfer.map(|v| v.as_str()).unwrap_or(""),
843 color_encoding: colorimetry.encoding.map(|v| v.as_str()).unwrap_or(""),
844 color_range: colorimetry.range.map(|v| v.as_str()).unwrap_or(""),
845 };
846
847 let geoms = export_plane_geoms(t, fmt)?;
850
851 let pin = match mode {
854 TransportMode::Inline => Some(t.pin_host(crate::CpuAccess::Read)?),
855 TransportMode::Reference => None,
856 };
857 let all_bytes: &[u8] = match &pin {
858 Some(p) => unsafe { std::slice::from_raw_parts(p.as_mut_ptr(), p.len()) },
859 None => &[],
860 };
861
862 let (handle, handle_bytes, fds): (i64, Vec<u8>, Vec<BlobFd>) = match mode {
866 TransportMode::Inline => (-1, Vec::new(), Vec::new()),
867 TransportMode::Reference => {
868 let raw = reference_handle(t)?;
869 #[cfg(target_os = "windows")]
870 {
871 (0, reference_handle_bytes(t, raw), Vec::new())
876 }
877 #[cfg(not(target_os = "windows"))]
878 {
879 let raw = i32::try_from(raw).map_err(|_| {
880 crate::Error::InvalidOperation("handle is not representable as an fd".into())
881 })?;
882 (0, Vec::new(), vec![raw])
884 }
885 }
886 };
887
888 let mut planes_buf = Vec::new();
889 write_export_planes(
890 &geoms,
891 mode,
892 handle,
893 &handle_bytes,
894 all_bytes,
895 &mut planes_buf,
896 )?;
897
898 let header = BlobHeader {
899 size: 0, required_mask: 0,
901 planes_bytes: planes_buf.len() as u64,
902 storage_kind: t.memory().code(),
903 pid: match mode {
905 TransportMode::Inline => 0,
906 TransportMode::Reference => std::process::id(),
907 },
908 fence_fd: -1,
911 dtype: t.dtype().code(),
912 quant_axis,
913 ndim: shape.len() as u32,
914 plane_count: geoms.len() as u32,
915 quant_scales_len: scales.len() as u32,
916 quant_zero_points_len: zeros.len() as u32,
917 strings_bytes: 0, strides_len: strides.len() as u32,
919 };
920
921 let mut out = vec![0u8; HEADER_LEN];
922 for d in &shape {
923 out.extend_from_slice(&d.to_le_bytes());
924 }
925 for st in &strides {
926 out.extend_from_slice(&st.to_le_bytes());
927 }
928 for sc in scales {
929 out.extend_from_slice(&sc.to_le_bytes());
930 }
931 for z in zeros {
932 out.extend_from_slice(&z.to_le_bytes());
933 }
934 let strings_start = out.len();
935 write_strings(&strings, &mut out);
936 let strings_bytes = out.len() - strings_start;
937 out.extend_from_slice(&planes_buf);
938
939 let mut header = header;
940 header.strings_bytes = strings_bytes as u32;
941 header.size = out.len() as u64;
942 header.write_to(&mut out);
943 Ok((out, fds))
944}
945
946fn export_addressing_shape(
947 t: &crate::TensorDyn,
948 fmt: Option<crate::PixelFormat>,
949) -> crate::Result<Vec<u64>> {
950 match fmt {
951 Some(f) => {
952 let (w, h) = image_dims(t, f)?;
953 Ok(f.addressing_shape(w, h)
954 .ok_or_else(|| {
955 crate::Error::InvalidArgument(format!("no addressing shape for {f:?}"))
956 })?
957 .iter()
958 .map(|d| *d as u64)
959 .collect())
960 }
961 None => Ok(t.shape().iter().map(|d| *d as u64).collect()),
962 }
963}
964
965fn export_plane_geoms(
966 t: &crate::TensorDyn,
967 fmt: Option<crate::PixelFormat>,
968) -> crate::Result<Vec<crate::PlaneGeometry>> {
969 match fmt {
970 Some(f) => {
971 let (w, h) = image_dims(t, f)?;
972 let rs = image_row_stride(t, f)?;
973 f.plane_table(w, h, rs)
974 .ok_or_else(|| crate::Error::InvalidArgument(format!("no plane table for {f:?}")))
975 }
976 None => Ok(vec![crate::PlaneGeometry {
977 offset: 0,
978 stride: t.capacity_bytes() as u64,
979 size: t.capacity_bytes() as u64,
980 }]),
981 }
982}
983
984fn write_export_planes(
985 geoms: &[crate::PlaneGeometry],
986 mode: TransportMode,
987 handle: i64,
988 handle_bytes: &[u8],
989 all_bytes: &[u8],
990 planes_buf: &mut Vec<u8>,
991) -> crate::Result<()> {
992 for g in geoms {
993 let data: &[u8] = match mode {
994 TransportMode::Inline => {
995 let start = g.offset as usize;
996 let end = start
997 .checked_add(g.size as usize)
998 .ok_or(crate::Error::InvalidSize(g.size as usize))?;
999 all_bytes
1000 .get(start..end)
1001 .ok_or(crate::Error::InvalidSize(end))?
1002 }
1003 TransportMode::Reference => &[],
1004 };
1005 write_plane(
1006 &BlobPlane {
1007 handle,
1008 offset: if mode == TransportMode::Inline {
1011 0
1012 } else {
1013 g.offset
1014 },
1015 stride: g.stride,
1016 size: if mode == TransportMode::Inline {
1018 data.len() as u64
1019 } else {
1020 g.size
1021 },
1022 used: if mode == TransportMode::Inline {
1023 data.len() as u64
1024 } else {
1025 g.size
1026 },
1027 modifier: 0,
1028 handle_bytes,
1033 data,
1034 },
1035 planes_buf,
1036 );
1037 }
1038 Ok(())
1039}
1040
1041fn image_dims(t: &crate::TensorDyn, f: crate::PixelFormat) -> crate::Result<(usize, usize)> {
1046 let shape = t.shape();
1047 let bad = || crate::Error::InvalidShape(format!("shape {shape:?} is not an image for {f:?}"));
1048 match (f.layout(), shape) {
1051 (crate::PixelLayout::Packed, [h, w, _]) => Ok((*w, *h)),
1052 (crate::PixelLayout::Planar, [_, h, w]) => Ok((*w, *h)),
1053 (crate::PixelLayout::SemiPlanar, [total, w]) => {
1054 let h = (0..=*total)
1058 .find(|&h| f.combined_plane_height(h) == Some(*total))
1059 .ok_or_else(bad)?;
1060 Ok((*w, h))
1061 }
1062 _ => Err(bad()),
1063 }
1064}
1065
1066fn image_row_stride(t: &crate::TensorDyn, f: crate::PixelFormat) -> crate::Result<usize> {
1072 t.effective_row_stride().ok_or_else(|| {
1073 crate::Error::InvalidShape(format!(
1074 "{f:?} tensor of shape {:?} reports no row pitch",
1075 t.shape()
1076 ))
1077 })
1078}
1079
1080fn reference_handle(t: &crate::TensorDyn) -> crate::Result<i64> {
1083 match t.memory() {
1084 #[cfg(target_os = "linux")]
1085 crate::TensorMemory::DmaBuf => {
1086 use std::os::fd::AsRawFd;
1087 t.dmabuf().map(|fd| fd.as_raw_fd() as i64).map_err(|e| {
1090 crate::Error::InvalidOperation(format!("dma-buf tensor has no fd: {e}"))
1091 })
1092 }
1093 #[cfg(any(target_os = "macos", target_os = "ios"))]
1094 crate::TensorMemory::DmaBuf => t
1095 .iosurface_id()
1096 .map(|id| id as i64)
1097 .ok_or_else(|| crate::Error::InvalidOperation("IOSurface tensor has no id".into())),
1098 #[cfg(target_os = "windows")]
1103 crate::TensorMemory::DmaBuf => {
1104 t.d3d11_shared_handle_value()
1105 .map(|h| h as i64)
1106 .ok_or_else(|| {
1107 crate::Error::InvalidOperation(
1108 "D3D11 texture tensor has no shared handle".into(),
1109 )
1110 })
1111 }
1112 other => Err(crate::Error::InvalidOperation(format!(
1116 "{other:?} has no shareable handle; export it with TransportMode::Inline"
1117 ))),
1118 }
1119}
1120
1121#[cfg(target_os = "windows")]
1136fn reference_handle_bytes(t: &crate::TensorDyn, raw: i64) -> Vec<u8> {
1137 let fence = crate::d3d11::device()
1138 .map(|d| d.exported_fence_handle_value())
1139 .unwrap_or(0);
1140 let (fence, value) = reference_completion(fence, t.gpu_write_value());
1141 let mut out = Vec::with_capacity(D3D11_HANDLE_BYTES);
1142 out.extend_from_slice(&(raw as u64).to_le_bytes());
1143 out.extend_from_slice(&fence.to_le_bytes());
1144 out.extend_from_slice(&value.to_le_bytes());
1145 out
1146}
1147
1148#[cfg(target_os = "windows")]
1157fn reference_completion(fence: usize, recorded: u64) -> (u64, u64) {
1158 if fence == 0 {
1159 (0, 0)
1160 } else {
1161 (fence as u64, recorded)
1162 }
1163}
1164
1165pub fn import(blob: &[u8], fds: &[BlobFd]) -> crate::Result<crate::TensorDyn> {
1187 #[cfg(not(unix))]
1188 let _ = fds;
1189 let v = BlobView::parse(blob).map_err(|e| crate::Error::InvalidArgument(e.to_string()))?;
1190 let h = v.header();
1191 let planes = v
1192 .planes()
1193 .map_err(|e| crate::Error::InvalidArgument(e.to_string()))?;
1194
1195 let inline = planes.first().map(BlobPlane::is_inline).unwrap_or(true);
1199 validate_import_planes(&planes, inline)?;
1200
1201 let dtype = crate::protocol::dtype_to_dtype(h.dtype)
1202 .ok_or_else(|| crate::Error::NotImplemented(format!("blob dtype code {}", h.dtype)))?;
1203 let strings = v.strings();
1204 let format = if strings.format.is_empty() {
1205 None
1206 } else {
1207 Some(
1208 crate::PixelFormat::from_str_code(strings.format).ok_or_else(|| {
1209 crate::Error::NotImplemented(format!("blob format {:?}", strings.format))
1210 })?,
1211 )
1212 };
1213
1214 if !inline {
1215 #[cfg(target_os = "windows")]
1216 {
1217 return import_referenced_d3d11_blob(&v, &h, &planes, format, &strings, dtype);
1218 }
1219 #[cfg(not(any(unix, target_os = "windows")))]
1220 {
1221 return Err(crate::Error::NotImplemented(
1222 "reference-mode blob import is not available on this platform".into(),
1223 ));
1224 }
1225 #[cfg(unix)]
1226 {
1227 return import_referenced_blob(&v, &h, &planes, fds, format, &strings, dtype);
1228 }
1229 }
1230
1231 let grid: Vec<usize> = v.shape().iter().map(|d| *d as usize).collect();
1235 let alloc_shape = alloc_shape_from_grid(format, &grid)?;
1236
1237 let mut t = crate::TensorDyn::new(&alloc_shape, dtype, Some(crate::TensorMemory::Mem), None)?;
1238 if let Some(f) = format {
1239 t.set_format(f)?;
1240 }
1241
1242 blit_inline_planes(&mut t, &planes, format, &grid)?;
1247
1248 t.set_colorimetry(colorimetry_from(&strings));
1249 if let Some(q) = import_quantization(&v)? {
1250 t.set_quantization(q)?;
1251 }
1252 Ok(t)
1253}
1254
1255fn validate_import_planes(planes: &[BlobPlane<'_>], inline: bool) -> crate::Result<()> {
1256 if planes.iter().any(|p| p.is_inline() != inline) {
1257 return Err(crate::Error::InvalidArgument(
1258 "blob mixes inline and referenced planes; a tensor has one transport mode".into(),
1259 ));
1260 }
1261 for p in planes {
1262 if p.is_inline() {
1263 if p.size as usize != p.data.len() {
1265 return Err(crate::Error::InvalidArgument(format!(
1266 "inline plane declares size {} but carries {} bytes",
1267 p.size,
1268 p.data.len()
1269 )));
1270 }
1271 if p.modifier != 0 || !p.handle_bytes.is_empty() {
1272 return Err(crate::Error::InvalidArgument(
1273 "inline plane must have modifier 0 and no handle_bytes".into(),
1274 ));
1275 }
1276 } else if !p.data.is_empty() {
1277 return Err(crate::Error::InvalidArgument(
1278 "referenced plane must not also carry inline bytes".into(),
1279 ));
1280 }
1281 }
1282 Ok(())
1283}
1284
1285fn alloc_shape_from_grid(
1286 format: Option<crate::PixelFormat>,
1287 grid: &[usize],
1288) -> crate::Result<Vec<usize>> {
1289 match format {
1290 Some(f) => {
1291 let (w, h_px) = dims_from_grid(f, grid)?;
1292 f.allocation_shape(w, h_px)
1293 .ok_or_else(|| crate::Error::InvalidShape(format!("no allocation shape for {f:?}")))
1294 }
1295 None => Ok(grid.to_vec()),
1296 }
1297}
1298
1299#[cfg(unix)]
1300fn import_referenced_blob(
1301 v: &BlobView<'_>,
1302 h: &BlobHeader,
1303 planes: &[BlobPlane<'_>],
1304 fds: &[BlobFd],
1305 format: Option<crate::PixelFormat>,
1306 strings: &BlobStrings<'_>,
1307 dtype: crate::DType,
1308) -> crate::Result<crate::TensorDyn> {
1309 let first = planes
1314 .first()
1315 .ok_or_else(|| crate::Error::InvalidArgument("referenced blob carries no planes".into()))?;
1316 let idx = usize::try_from(first.handle).map_err(|_| {
1317 crate::Error::InvalidArgument(format!("plane handle {} is not an index", first.handle))
1318 })?;
1319 let raw = *fds.get(idx).ok_or_else(|| {
1320 crate::Error::InvalidArgument(format!(
1321 "plane handle indexes fd {idx} but only {} were provided",
1322 fds.len()
1323 ))
1324 })?;
1325 for p in planes {
1326 let i = usize::try_from(p.handle).unwrap_or(usize::MAX);
1327 if i >= fds.len() {
1328 return Err(crate::Error::InvalidArgument(format!(
1329 "plane handle indexes fd {i} but only {} were provided",
1330 fds.len()
1331 )));
1332 }
1333 }
1334
1335 let grid: Vec<usize> = v.shape().iter().map(|d| *d as usize).collect();
1336 let alloc_shape = alloc_shape_from_grid(format, &grid)?;
1337 let mut t = open_referenced_tensor(h.storage_kind, raw, &alloc_shape, dtype)?;
1338 if let Some(f) = format {
1339 t.set_format(f)?;
1340 if first.stride > 0 {
1345 t.set_row_stride(first.stride as usize)?;
1346 }
1347 }
1348 t.set_colorimetry(colorimetry_from(strings));
1349 if let Some(q) = import_quantization(v)? {
1350 t.set_quantization(q)?;
1351 }
1352 Ok(t)
1353}
1354
1355#[cfg(unix)]
1356fn open_referenced_tensor(
1357 storage_kind: u32,
1358 raw: BlobFd,
1359 alloc_shape: &[usize],
1360 dtype: crate::DType,
1361) -> crate::Result<crate::TensorDyn> {
1362 match crate::TensorMemory::from_code(storage_kind) {
1368 #[cfg(any(target_os = "macos", target_os = "ios"))]
1369 Some(crate::TensorMemory::DmaBuf) | Some(crate::TensorMemory::IoSurface) => {
1370 let id = u32::try_from(raw).map_err(|_| {
1371 crate::Error::InvalidArgument(format!("{raw} is not an IOSurface id"))
1372 })?;
1373 crate::TensorDyn::from_iosurface_id(id, alloc_shape, dtype, None)
1374 }
1375 #[cfg(all(unix, not(any(target_os = "macos", target_os = "ios"))))]
1376 Some(crate::TensorMemory::DmaBuf) | Some(crate::TensorMemory::Shm) => {
1377 use std::os::fd::{BorrowedFd, OwnedFd};
1378 let borrowed = unsafe { BorrowedFd::borrow_raw(raw) };
1383 let owned: OwnedFd = borrowed
1384 .try_clone_to_owned()
1385 .map_err(crate::Error::IoError)?;
1386 crate::TensorDyn::from_fd(owned, alloc_shape, dtype, None)
1387 }
1388 other => Err(crate::Error::NotImplemented(format!(
1389 "reference-mode import for {other:?} on this platform"
1390 ))),
1391 }
1392}
1393
1394#[cfg(target_os = "windows")]
1403fn import_referenced_d3d11_blob(
1404 v: &BlobView<'_>,
1405 h: &BlobHeader,
1406 planes: &[BlobPlane<'_>],
1407 format: Option<crate::PixelFormat>,
1408 strings: &BlobStrings<'_>,
1409 dtype: crate::DType,
1410) -> crate::Result<crate::TensorDyn> {
1411 let first = planes
1412 .first()
1413 .ok_or_else(|| crate::Error::InvalidArgument("referenced blob carries no planes".into()))?;
1414 for p in planes {
1421 if p.handle != 0 || p.handle_bytes != first.handle_bytes {
1422 return Err(crate::Error::InvalidArgument(
1423 "D3D11 reference blob planes must all carry handle 0 and identical \
1424 handle bytes; they name one texture"
1425 .into(),
1426 ));
1427 }
1428 }
1429 let grid: Vec<usize> = v.shape().iter().map(|d| *d as usize).collect();
1434 let mut t = open_referenced_texture(h.storage_kind, first, h.pid, &grid, dtype, format)?;
1435 t.set_colorimetry(colorimetry_from(strings));
1443 if let Some(q) = import_quantization(v)? {
1444 t.set_quantization(q)?;
1445 }
1446 Ok(t)
1447}
1448
1449#[cfg(target_os = "windows")]
1464fn open_referenced_texture(
1465 storage_kind: u32,
1466 first: &BlobPlane<'_>,
1467 pid: u32,
1468 grid: &[usize],
1469 dtype: crate::DType,
1470 format: Option<crate::PixelFormat>,
1471) -> crate::Result<crate::TensorDyn> {
1472 use std::os::windows::io::{AsRawHandle, FromRawHandle, OwnedHandle};
1473 use windows::Win32::Foundation::HANDLE;
1474 use windows::Win32::System::Threading::{OpenProcess, PROCESS_DUP_HANDLE};
1475
1476 if storage_kind != crate::TensorMemory::DmaBuf.code() {
1477 return Err(crate::Error::InvalidArgument(format!(
1478 "storage kind {storage_kind} is not referenceable on Windows"
1479 )));
1480 }
1481 if first.handle_bytes.len() != D3D11_HANDLE_BYTES {
1482 return Err(crate::Error::InvalidArgument(format!(
1483 "D3D11 reference plane needs {D3D11_HANDLE_BYTES} handle bytes, not {}",
1484 first.handle_bytes.len()
1485 )));
1486 }
1487 let format = format.ok_or_else(|| {
1488 crate::Error::InvalidArgument("D3D11 reference blob without a pixel format".into())
1489 })?;
1490 let word = |i: usize| {
1491 let mut b = [0u8; 8];
1492 b.copy_from_slice(&first.handle_bytes[i * 8..(i + 1) * 8]);
1493 u64::from_le_bytes(b)
1494 };
1495 let handle_at = |i: usize| HANDLE(word(i) as usize as *mut std::ffi::c_void);
1496 let (mut tex, mut fence, value) = (handle_at(0), handle_at(1), word(2));
1497 if tex.is_invalid() {
1498 return Err(crate::Error::InvalidArgument(
1499 "D3D11 reference plane carries no texture handle".into(),
1500 ));
1501 }
1502 if fence.is_invalid() && value != 0 {
1508 return Err(crate::Error::InvalidArgument(format!(
1509 "D3D11 reference plane names fence value {value} with no fence handle; \
1510 the completion it names cannot be waited on"
1511 )));
1512 }
1513
1514 let mut duplicates: Vec<OwnedHandle> = Vec::new();
1518 if std::process::id() != pid {
1519 let src = unsafe { OpenProcess(PROCESS_DUP_HANDLE, false, pid) }
1523 .map_err(|e| open_process_error(pid, &e))?;
1524 let src = unsafe { OwnedHandle::from_raw_handle(src.0) };
1527 for (h, which) in [(&mut tex, "texture"), (&mut fence, "fence")] {
1528 if h.is_invalid() {
1529 continue;
1530 }
1531 let dup = duplicate_out_of(&src, *h, pid, which)?;
1532 *h = HANDLE(dup.as_raw_handle());
1533 duplicates.push(dup);
1534 }
1535 }
1536
1537 let geometry = unsafe { crate::protocol::d3d11_geometry_checked(format, tex.0, grid) };
1541 let completion = (value != 0).then_some((fence.0, value));
1544 let opened = match geometry {
1545 Ok((width, height)) => unsafe {
1548 crate::TensorDyn::from_d3d11_shared_handle(
1549 tex.0,
1550 width,
1551 height,
1552 format,
1553 dtype,
1554 crate::CpuAccess::ReadWrite,
1555 completion,
1556 None,
1557 )
1558 },
1559 Err(e) => Err(e),
1560 };
1561 drop(duplicates);
1564 opened
1565}
1566
1567#[cfg(target_os = "windows")]
1576fn win32_error_kind(e: &windows::core::Error) -> std::io::ErrorKind {
1577 use std::io::ErrorKind;
1578 use windows::core::HRESULT;
1579 use windows::Win32::Foundation::{ERROR_ACCESS_DENIED, ERROR_INVALID_PARAMETER};
1580
1581 let code = e.code();
1582 if code == HRESULT::from_win32(ERROR_ACCESS_DENIED.0) {
1583 ErrorKind::PermissionDenied
1584 } else if code == HRESULT::from_win32(ERROR_INVALID_PARAMETER.0) {
1585 ErrorKind::NotFound
1590 } else {
1591 ErrorKind::Other
1592 }
1593}
1594
1595#[cfg(target_os = "windows")]
1597fn open_process_error(pid: u32, e: &windows::core::Error) -> crate::Error {
1598 use std::io::ErrorKind;
1599
1600 let kind = win32_error_kind(e);
1601 let what = match kind {
1602 ErrorKind::PermissionDenied => {
1603 format!("cannot open exporting process {pid} to duplicate its handles: access denied")
1604 }
1605 ErrorKind::NotFound => format!("exporting process {pid} is gone"),
1606 _ => format!(
1607 "cannot open exporting process {pid} to duplicate its handles: {} \
1608 (HRESULT 0x{:08X})",
1609 e.message(),
1610 e.code().0 as u32
1611 ),
1612 };
1613 crate::Error::IoError(std::io::Error::new(kind, what))
1614}
1615
1616#[cfg(target_os = "windows")]
1620fn duplicate_out_of(
1621 src: &std::os::windows::io::OwnedHandle,
1622 handle: windows::Win32::Foundation::HANDLE,
1623 pid: u32,
1624 which: &str,
1625) -> crate::Result<std::os::windows::io::OwnedHandle> {
1626 use std::os::windows::io::{AsRawHandle, FromRawHandle, OwnedHandle};
1627 use windows::Win32::Foundation::{DuplicateHandle, DUPLICATE_SAME_ACCESS, HANDLE};
1628 use windows::Win32::System::Threading::GetCurrentProcess;
1629
1630 let mut out = HANDLE(std::ptr::null_mut());
1631 unsafe {
1634 DuplicateHandle(
1635 HANDLE(src.as_raw_handle()),
1636 handle,
1637 GetCurrentProcess(),
1638 &mut out,
1639 0,
1640 false,
1641 DUPLICATE_SAME_ACCESS,
1642 )
1643 }
1644 .map_err(|e| {
1645 crate::Error::IoError(std::io::Error::new(
1649 win32_error_kind(&e),
1650 format!(
1651 "cannot duplicate the {which} handle out of exporting process {pid}: {} \
1652 (HRESULT 0x{:08X})",
1653 e.message(),
1654 e.code().0 as u32
1655 ),
1656 ))
1657 })?;
1658 Ok(unsafe { OwnedHandle::from_raw_handle(out.0) })
1661}
1662
1663fn blit_inline_planes(
1664 t: &mut crate::TensorDyn,
1665 planes: &[BlobPlane<'_>],
1666 format: Option<crate::PixelFormat>,
1667 grid: &[usize],
1668) -> crate::Result<()> {
1669 let pin = t.pin_host(crate::CpuAccess::Write)?;
1670 let dst = unsafe { std::slice::from_raw_parts_mut(pin.as_mut_ptr(), pin.len()) };
1671 let geoms: Vec<crate::PlaneGeometry> = match format {
1672 Some(f) => {
1673 let (w, h_px) = dims_from_grid(f, grid)?;
1674 let rs = image_row_stride(t, f)?;
1675 f.plane_table(w, h_px, rs)
1676 .ok_or_else(|| crate::Error::InvalidShape(format!("no plane table for {f:?}")))?
1677 }
1678 None => vec![crate::PlaneGeometry {
1679 offset: 0,
1680 stride: dst.len() as u64,
1681 size: dst.len() as u64,
1682 }],
1683 };
1684 if geoms.len() != planes.len() {
1685 return Err(crate::Error::InvalidArgument(format!(
1686 "blob carries {} planes but {:?} has {}",
1687 planes.len(),
1688 format,
1689 geoms.len()
1690 )));
1691 }
1692 for (g, p) in geoms.iter().zip(planes) {
1693 let start = g.offset as usize;
1694 let end = start
1695 .checked_add(p.data.len())
1696 .ok_or(crate::Error::InvalidSize(p.data.len()))?;
1697 let slot = dst
1698 .get_mut(start..end)
1699 .ok_or(crate::Error::InvalidSize(end))?;
1700 slot.copy_from_slice(p.data);
1701 }
1702 Ok(())
1703}
1704
1705fn colorimetry_from(s: &BlobStrings<'_>) -> Option<crate::Colorimetry> {
1709 let c = crate::Colorimetry {
1710 space: crate::ColorSpace::from_str_code(s.color_space),
1711 transfer: crate::ColorTransfer::from_str_code(s.color_transfer),
1712 encoding: crate::ColorEncoding::from_str_code(s.color_encoding),
1713 range: crate::ColorRange::from_str_code(s.color_range),
1714 };
1715 if c == crate::Colorimetry::default() {
1716 None
1717 } else {
1718 Some(c)
1719 }
1720}
1721
1722fn import_quantization(v: &BlobView<'_>) -> crate::Result<Option<crate::Quantization>> {
1730 let h = v.header;
1731 if h.quant_axis == -2 {
1732 return Ok(None);
1733 }
1734 let scales: Vec<f32> = (0..h.quant_scales_len as usize)
1736 .map(|i| f32::from_bits(rd_u32(v.buf, v.regions.quant_scales + i * 4)))
1737 .collect();
1738 let zeros: Vec<i32> = (0..h.quant_zero_points_len as usize)
1739 .map(|i| rd_u32(v.buf, v.regions.quant_zero_points + i * 4) as i32)
1740 .collect();
1741 quantization_from(h.quant_axis, scales, zeros).map(Some)
1742}
1743
1744fn quantization_from(
1747 axis: i32,
1748 scales: Vec<f32>,
1749 zeros: Vec<i32>,
1750) -> crate::Result<crate::Quantization> {
1751 let malformed = |why: String| {
1752 crate::Error::InvalidArgument(format!("blob quantization is malformed: {why}"))
1753 };
1754 if axis == -1 {
1755 return match (scales.as_slice(), zeros.as_slice()) {
1756 ([scale], []) => Ok(crate::Quantization::per_tensor_symmetric(*scale)),
1757 ([scale], [zero]) => Ok(crate::Quantization::per_tensor(*scale, *zero)),
1758 (s, z) => Err(malformed(format!(
1759 "per-tensor quantization needs 1 scale and 0 or 1 zero points, \
1760 got {} and {}",
1761 s.len(),
1762 z.len()
1763 ))),
1764 };
1765 }
1766 let axis = usize::try_from(axis).map_err(|_| {
1767 malformed(format!(
1768 "quant_axis {axis} is neither -2 (none), -1 (per-tensor) nor a channel index"
1769 ))
1770 })?;
1771 let q = if zeros.is_empty() {
1772 crate::Quantization::per_channel_symmetric(scales, axis)
1773 } else {
1774 crate::Quantization::per_channel(scales, zeros, axis)
1775 };
1776 q.map_err(|e| malformed(e.to_string()))
1777}
1778
1779fn dims_from_grid(f: crate::PixelFormat, grid: &[usize]) -> crate::Result<(usize, usize)> {
1783 let bad = || crate::Error::InvalidShape(format!("grid {grid:?} is not an image for {f:?}"));
1784 match (f.layout(), grid) {
1785 (crate::PixelLayout::Planar, [_, h, w, ..]) => Ok((*w, *h)),
1786 (crate::PixelLayout::Planar, _) => Err(bad()),
1787 (_, [h, w, ..]) => Ok((*w, *h)),
1790 _ => Err(bad()),
1791 }
1792}
1793
1794#[cfg(test)]
1795mod tests {
1796 use super::*;
1797
1798 #[cfg(target_os = "windows")]
1803 #[test]
1804 fn a_reference_plane_drops_the_value_with_the_fence() {
1805 assert_eq!(super::reference_completion(0, 42), (0, 0));
1806 assert_eq!(super::reference_completion(0, 0), (0, 0));
1807 assert_eq!(super::reference_completion(0x1234, 0), (0x1234, 0));
1810 assert_eq!(super::reference_completion(0x1234, 42), (0x1234, 42));
1811 }
1812
1813 #[test]
1814 fn every_region_offset_is_constant_time_arithmetic() {
1815 let h = BlobHeader {
1817 size: 4096,
1818 ndim: 3,
1819 strides_len: 3,
1820 quant_scales_len: 4,
1821 quant_zero_points_len: 4,
1822 strings_bytes: 24,
1823 plane_count: 2,
1824 planes_bytes: 128,
1825 ..BlobHeader::empty()
1826 };
1827 let r = region_offsets(&h).unwrap();
1828 assert_eq!(r.shape, HEADER_LEN);
1829 assert_eq!(r.strides, HEADER_LEN + 3 * 8);
1830 assert_eq!(r.quant_scales, HEADER_LEN + 3 * 8 + 3 * 8);
1831 assert_eq!(r.quant_zero_points, r.quant_scales + 4 * 4);
1832 assert_eq!(r.strings, r.quant_zero_points + 4 * 4);
1833 assert_eq!(r.planes, r.strings + 24);
1834 assert_eq!(r.end, r.planes + 128);
1835 }
1836
1837 #[test]
1838 fn empty_strides_leave_the_region_zero_width() {
1839 let h = BlobHeader {
1840 size: 4096,
1841 ndim: 2,
1842 strides_len: 0,
1843 ..BlobHeader::empty()
1844 };
1845 let r = region_offsets(&h).unwrap();
1846 assert_eq!(r.strides, HEADER_LEN + 2 * 8);
1847 assert_eq!(
1848 r.quant_scales, r.strides,
1849 "densely-packed strides occupy no bytes"
1850 );
1851 }
1852
1853 #[test]
1854 fn strides_must_be_ndim_entries_or_none() {
1855 let h = BlobHeader {
1857 size: 4096,
1858 ndim: 3,
1859 strides_len: 2,
1860 ..BlobHeader::empty()
1861 };
1862 assert!(matches!(
1863 region_offsets(&h),
1864 Err(BlobError::ForgedLength {
1865 field: "strides_len",
1866 ..
1867 })
1868 ));
1869 }
1870
1871 #[test]
1872 fn a_forged_ndim_cannot_derive_an_offset_past_the_buffer() {
1873 let h = BlobHeader {
1874 size: 72,
1875 ndim: u32::MAX,
1876 strides_len: u32::MAX,
1877 ..BlobHeader::empty()
1878 };
1879 assert!(
1880 region_offsets(&h).is_err(),
1881 "ndim * 8 must be checked, not wrapped"
1882 );
1883 }
1884
1885 #[test]
1886 fn regions_that_overrun_size_are_refused() {
1887 let h = BlobHeader {
1890 size: 80,
1891 ndim: 2,
1892 strides_len: 2,
1893 plane_count: 1,
1894 planes_bytes: 4096,
1895 ..BlobHeader::empty()
1896 };
1897 assert!(region_offsets(&h).is_err());
1898 }
1899
1900 #[test]
1901 fn plane_count_is_bounded_by_the_planes_region_not_by_a_constant() {
1902 let h = BlobHeader {
1905 size: 4096,
1906 plane_count: u32::MAX,
1907 planes_bytes: 56,
1908 ..BlobHeader::empty()
1909 };
1910 assert!(matches!(
1911 region_offsets(&h),
1912 Err(BlobError::ForgedLength {
1913 field: "plane_count",
1914 ..
1915 })
1916 ));
1917 }
1918
1919 fn fill_ramp(t: &TensorDyn) {
1927 let pin = t.pin_host(CpuAccess::ReadWrite).expect("pin");
1928 t.sync_for_cpu(CpuAccess::ReadWrite).expect("sync for cpu");
1929 let s = unsafe { std::slice::from_raw_parts_mut(pin.as_mut_ptr(), pin.len()) };
1933 for (i, b) in s.iter_mut().enumerate() {
1934 *b = (i % 251) as u8;
1935 }
1936 t.sync_for_device(CpuAccess::ReadWrite)
1937 .expect("sync for device");
1938 }
1939
1940 fn bytes_of(t: &TensorDyn) -> Vec<u8> {
1941 let pin = t.pin_host(CpuAccess::Read).expect("pin");
1942 t.sync_for_cpu(CpuAccess::Read).expect("sync for cpu");
1943 let bytes = unsafe { std::slice::from_raw_parts(pin.as_mut_ptr(), pin.len()) }.to_vec();
1946 t.sync_for_device(CpuAccess::Read).expect("sync for device");
1947 bytes
1948 }
1949
1950 fn reference_capable(w: usize, h: usize, what: &str) -> Option<TensorDyn> {
1958 if cfg!(target_os = "linux") && !crate::test_support::dma_or_skip(what) {
1959 return None;
1960 }
1961 match Tensor::<u8>::image(
1962 w,
1963 h,
1964 PixelFormat::Nv12,
1965 Some(TensorMemory::DmaBuf),
1966 CpuAccess::ReadWrite,
1967 ) {
1968 Ok(t) => Some(TensorDyn::from(t)),
1969 Err(e) if cfg!(target_os = "linux") => {
1970 panic!("{what}: a usable DMA heap must allocate NV12 {w}x{h}: {e:?}")
1971 }
1972 Err(e) => {
1973 crate::test_support::report_skip(&format!(
1974 "no shareable backing on this host ({e:?}); reference \
1975 transport not exercised"
1976 ));
1977 None
1978 }
1979 }
1980 }
1981
1982 #[test]
1983 fn a_reference_export_puts_an_index_in_the_handle_not_a_raw_fd() {
1984 let Some(src) = reference_capable(
1985 64,
1986 48,
1987 "a_reference_export_puts_an_index_in_the_handle_not_a_raw_fd",
1988 ) else {
1989 return;
1990 };
1991 let (blob, fds) = export(&src, TransportMode::Reference).unwrap();
1992 let v = BlobView::parse(&blob).unwrap();
1993 let planes = v.planes().unwrap();
1994 for p in &planes {
1995 assert!(!p.is_inline());
1996 }
1997 if cfg!(target_os = "windows") {
1998 assert!(fds.is_empty(), "an NT handle does not travel out of band");
2003 assert!(planes
2004 .iter()
2005 .all(|p| p.handle_bytes.len() == D3D11_HANDLE_BYTES));
2006 } else {
2007 assert!(!fds.is_empty(), "reference mode must carry an fd table");
2008 for p in &planes {
2009 assert!(
2010 (p.handle as usize) < fds.len(),
2011 "handle {} must index the {}-entry fd table, not be a raw fd",
2012 p.handle,
2013 fds.len()
2014 );
2015 }
2016 assert_eq!(fds.len(), 1, "planes sharing a buffer share a table entry");
2020 }
2021 assert!(planes.iter().all(|p| p.handle == 0));
2022 }
2023
2024 #[test]
2025 fn a_reference_round_trip_preserves_content_and_the_source_may_die() {
2026 let Some(src) = reference_capable(
2027 64,
2028 48,
2029 "a_reference_round_trip_preserves_content_and_the_source_may_die",
2030 ) else {
2031 return;
2032 };
2033 fill_ramp(&src);
2034 let expected = bytes_of(&src);
2035 let (blob, fds) = export(&src, TransportMode::Reference).unwrap();
2036 let src_id = src.buffer_identity().id();
2037 let got = import(&blob, &fds).unwrap();
2038
2039 if !cfg!(target_os = "windows") {
2053 assert_eq!(
2054 got.buffer_identity().id(),
2055 src_id,
2056 "reference mode must share the buffer, not copy it"
2057 );
2058 }
2059
2060 drop(src);
2062 assert_eq!(
2063 bytes_of(&got),
2064 expected,
2065 "the imported tensor sees the bytes"
2066 );
2067 assert_eq!(got.format(), Some(PixelFormat::Nv12));
2068 }
2069
2070 #[test]
2071 fn an_fd_index_past_the_table_is_refused() {
2072 let Some(src) = reference_capable(64, 48, "an_fd_index_past_the_table_is_refused") else {
2075 return;
2076 };
2077 let (blob, fds) = export(&src, TransportMode::Reference).unwrap();
2078 if cfg!(target_os = "windows") {
2079 let err = import(&patch_handle_word(&blob, 0, 0), &[])
2084 .expect_err("a null texture handle must be refused, not opened");
2085 assert!(
2086 matches!(err, crate::Error::InvalidArgument(_)),
2087 "expected InvalidArgument, got {err:?}"
2088 );
2089
2090 let named = patch_handle_word(&patch_handle_word(&blob, 1, 0), 2, 7);
2096 let err = import(&named, &[])
2097 .expect_err("a fence value with no fence handle must be refused");
2098 let msg = format!("{err}");
2099 assert!(
2100 matches!(err, crate::Error::InvalidArgument(_))
2101 && msg.contains("cannot be waited on"),
2102 "expected the unwaitable-completion refusal, got {err:?}"
2103 );
2104
2105 let mut tampered = blob.clone();
2108 tampered[HEADER_PID_OFFSET..HEADER_PID_OFFSET + 4]
2109 .copy_from_slice(&u32::MAX.to_le_bytes());
2110 let err =
2111 import(&tampered, &[]).expect_err("a pid that does not exist must be refused");
2112 match &err {
2113 crate::Error::IoError(e) => assert_eq!(
2114 e.kind(),
2115 std::io::ErrorKind::NotFound,
2116 "a pid no process holds is reported as gone, got {e:?}"
2117 ),
2118 other => panic!("expected IoError, got {other:?}"),
2119 }
2120 return;
2121 }
2122 let r = region_offsets(&BlobView::parse(&blob).unwrap().header()).unwrap();
2123 assert!(
2124 import(&blob, &[]).is_err(),
2125 "an empty table cannot satisfy an index"
2126 );
2127 let mut tampered = blob.clone();
2128 tampered[r.planes..r.planes + 8].copy_from_slice(&99i64.to_le_bytes());
2129 assert!(
2130 import(&tampered, &fds).is_err(),
2131 "an out-of-range index must be refused, not read"
2132 );
2133 }
2134
2135 fn patch_handle_word(blob: &[u8], word: usize, value: u64) -> Vec<u8> {
2146 let v = BlobView::parse(blob).unwrap();
2147 let r = region_offsets(&v.header()).unwrap();
2148 let stride = (PLANE_RECORD_LEN + D3D11_HANDLE_BYTES).next_multiple_of(8);
2149 let mut out = blob.to_vec();
2150 for i in 0..v.header().plane_count as usize {
2151 let at = r.planes + i * stride + PLANE_RECORD_LEN + word * 8;
2152 out[at..at + 8].copy_from_slice(&value.to_le_bytes());
2153 }
2154 out
2155 }
2156
2157 #[test]
2162 #[cfg(target_os = "windows")]
2163 fn a_garbage_texture_handle_is_refused_by_the_constructor() {
2164 let Some(src) = reference_capable(
2165 64,
2166 48,
2167 "a_garbage_texture_handle_is_refused_by_the_constructor",
2168 ) else {
2169 return;
2170 };
2171 let (blob, _) = export(&src, TransportMode::Reference).unwrap();
2172 let err = import(&patch_handle_word(&blob, 0, 0x1234), &[])
2173 .expect_err("a handle that is not a shared texture must be refused");
2174 let msg = format!("{err}");
2175 assert!(
2176 msg.contains("OpenSharedResource"),
2177 "the refusal must come from opening the handle, got: {msg}"
2178 );
2179 }
2180
2181 #[test]
2184 #[cfg(target_os = "windows")]
2185 fn planes_that_disagree_about_the_texture_are_refused() {
2186 let Some(src) =
2187 reference_capable(64, 48, "planes_that_disagree_about_the_texture_are_refused")
2188 else {
2189 return;
2190 };
2191 let (blob, _) = export(&src, TransportMode::Reference).unwrap();
2192 let v = BlobView::parse(&blob).unwrap();
2193 let r = region_offsets(&v.header()).unwrap();
2194 assert_eq!(v.planes().unwrap().len(), 2, "NV12 exports two planes");
2195 let second = r.planes + (PLANE_RECORD_LEN + D3D11_HANDLE_BYTES).next_multiple_of(8);
2197 let mut tampered = blob.clone();
2198 tampered[second + PLANE_RECORD_LEN..second + PLANE_RECORD_LEN + 8]
2199 .copy_from_slice(&0xBEEFu64.to_le_bytes());
2200 let err = import(&tampered, &[]).expect_err("planes naming two textures must be refused");
2201 let msg = format!("{err}");
2202 assert!(
2203 matches!(err, crate::Error::InvalidArgument(_)) && msg.contains("one texture"),
2204 "expected the one-texture refusal, got {err:?}"
2205 );
2206 }
2207
2208 #[test]
2210 fn pid_offset_matches_the_writer() {
2211 let h = BlobHeader {
2212 size: HEADER_LEN as u64,
2213 pid: 0x0BAD_F00D,
2214 ..BlobHeader::empty()
2215 };
2216 let mut buf = vec![0u8; HEADER_LEN];
2217 h.write_to(&mut buf);
2218 assert_eq!(rd_u32(&buf, HEADER_PID_OFFSET), h.pid);
2219 }
2220
2221 #[test]
2222 fn inline_round_trip_preserves_content_and_every_metadata_field() {
2223 let src = nv12_mem(64, 48).with_colorimetry(crate::Colorimetry::jfif());
2224 fill_ramp(&src);
2225 let blob = export(&src, TransportMode::Inline).unwrap().0;
2226 let got = import(&blob, &[]).unwrap();
2227
2228 assert_eq!(got.dtype(), src.dtype());
2229 assert_eq!(got.format(), src.format());
2230 assert_eq!(got.colorimetry(), src.colorimetry());
2231 assert_eq!(got.shape(), src.shape(), "allocation geometry is rebuilt");
2232 assert_eq!(bytes_of(&got), bytes_of(&src), "every byte survived");
2233 }
2234
2235 #[test]
2236 fn inline_import_is_a_new_allocation_and_the_source_may_die() {
2237 let src = nv12_mem(64, 48);
2238 fill_ramp(&src);
2239 let expected = bytes_of(&src);
2240 let blob = export(&src, TransportMode::Inline).unwrap().0;
2241 drop(src);
2242 let got = import(&blob, &[]).unwrap();
2243 assert_eq!(bytes_of(&got), expected);
2244 }
2245
2246 #[test]
2247 fn a_batched_image_is_refused_rather_than_exported_as_one_image() {
2248 let mut t = Tensor::<u8>::image(
2249 8,
2250 4,
2251 crate::PixelFormat::Rgb,
2252 Some(TensorMemory::Mem),
2253 crate::CpuAccess::ReadWrite,
2254 )
2255 .unwrap();
2256 crate::TensorTrait::set_logical_shape(&mut t, &[2, 2, 8, 3]).unwrap();
2257 match export(&TensorDyn::from(t), TransportMode::Inline) {
2258 Err(crate::Error::InvalidShape(msg)) => {
2259 assert!(msg.contains("[2, 2, 8, 3]"), "{msg}")
2260 }
2261 other => panic!("expected InvalidShape, got {other:?}"),
2262 }
2263 }
2264
2265 #[test]
2266 fn a_bare_tensor_round_trips_including_byte_strides() {
2267 let src =
2268 TensorDyn::from(Tensor::<f32>::new(&[3, 8], Some(TensorMemory::Mem), None).unwrap());
2269 fill_ramp(&src);
2270 let blob = export(&src, TransportMode::Inline).unwrap().0;
2271 let got = import(&blob, &[]).unwrap();
2272 assert_eq!(got.shape(), &[3, 8]);
2273 assert_eq!(got.dtype(), crate::DType::F32);
2274 assert_eq!(got.format(), None);
2275 assert_eq!(bytes_of(&got), bytes_of(&src));
2276 }
2277
2278 fn build_blob(planes: &[BlobPlane<'_>]) -> Vec<u8> {
2286 build(&BlobSpec::new(planes))
2287 }
2288
2289 struct BlobSpec<'a> {
2291 planes: &'a [BlobPlane<'a>],
2292 storage_kind: u32,
2293 dtype: crate::DType,
2294 grid: &'a [u64],
2295 strings: BlobStrings<'a>,
2296 quant_axis: i32,
2297 scales: &'a [f32],
2298 zeros: &'a [i32],
2299 }
2300
2301 impl<'a> BlobSpec<'a> {
2302 fn new(planes: &'a [BlobPlane<'a>]) -> Self {
2304 Self {
2305 planes,
2306 storage_kind: 0,
2307 dtype: crate::DType::U8,
2308 grid: &[1],
2309 strings: BlobStrings::default(),
2310 quant_axis: -2,
2311 scales: &[],
2312 zeros: &[],
2313 }
2314 }
2315 }
2316
2317 fn build(spec: &BlobSpec<'_>) -> Vec<u8> {
2318 let mut planes_buf = Vec::new();
2319 for p in spec.planes {
2320 write_plane(p, &mut planes_buf);
2321 }
2322 let mut out = vec![0u8; HEADER_LEN];
2323 for d in spec.grid {
2324 out.extend_from_slice(&d.to_le_bytes());
2325 }
2326 for s in spec.scales {
2327 out.extend_from_slice(&s.to_le_bytes());
2328 }
2329 for z in spec.zeros {
2330 out.extend_from_slice(&z.to_le_bytes());
2331 }
2332 let strings_bytes = strings_encoded_len(&spec.strings, out.len());
2333 write_strings(&spec.strings, &mut out);
2334 out.extend_from_slice(&planes_buf);
2335 let h = BlobHeader {
2336 size: out.len() as u64,
2337 planes_bytes: planes_buf.len() as u64,
2338 storage_kind: spec.storage_kind,
2339 dtype: spec.dtype.code(),
2340 quant_axis: spec.quant_axis,
2341 ndim: spec.grid.len() as u32,
2342 plane_count: spec.planes.len() as u32,
2343 quant_scales_len: spec.scales.len() as u32,
2344 quant_zero_points_len: spec.zeros.len() as u32,
2345 strings_bytes: strings_bytes as u32,
2346 ..BlobHeader::empty()
2347 };
2348 h.write_to(&mut out);
2349 out
2350 }
2351
2352 #[test]
2353 fn no_corrupted_header_value_can_panic_the_parser() {
2354 let good = export(&nv12_mem(64, 48), TransportMode::Inline).unwrap().0;
2360 for word in 0..HEADER_LEN / 4 {
2361 for pattern in [u32::MAX, u32::MAX / 2, 1, 0x8000_0000] {
2362 let mut bad = good.clone();
2363 bad[word * 4..word * 4 + 4].copy_from_slice(&pattern.to_le_bytes());
2364 if let Ok(v) = BlobView::parse(&bad) {
2367 let _ = v.shape();
2368 let _ = v.strides();
2369 let _ = v.strings();
2370 let _ = v.planes();
2371 }
2372 let _ = import(&bad, &[]);
2373 }
2374 }
2375 }
2376
2377 #[test]
2378 fn no_corrupted_plane_record_can_panic_the_parser() {
2379 let good = export(&nv12_mem(64, 48), TransportMode::Inline).unwrap().0;
2380 let r = region_offsets(&BlobView::parse(&good).unwrap().header()).unwrap();
2381 for word in 0..PLANE_RECORD_LEN / 4 {
2382 for pattern in [u32::MAX, u32::MAX / 2, 0x8000_0000] {
2383 let mut bad = good.clone();
2384 let at = r.planes + word * 4;
2385 bad[at..at + 4].copy_from_slice(&pattern.to_le_bytes());
2386 if let Ok(v) = BlobView::parse(&bad) {
2387 let _ = v.planes();
2388 }
2389 let _ = import(&bad, &[]);
2390 }
2391 }
2392 }
2393
2394 #[test]
2395 fn a_mixed_inline_and_referenced_plane_set_is_refused() {
2396 let payload = [1u8, 2, 3, 4];
2404 let referenced = BlobPlane {
2405 handle: 7,
2406 offset: 0,
2407 stride: 4,
2408 size: 4,
2409 used: 4,
2410 modifier: 0,
2411 handle_bytes: &[],
2412 data: &[],
2413 };
2414 let inlined = BlobPlane {
2415 handle: -1,
2416 offset: 0,
2417 stride: 4,
2418 size: payload.len() as u64,
2419 used: payload.len() as u64,
2420 modifier: 0,
2421 handle_bytes: &[],
2422 data: &payload,
2423 };
2424 let blob = build_blob(&[referenced, inlined]);
2425 let err =
2426 import(&blob, &[]).expect_err("a plane set mixing transport modes must be refused");
2427 assert!(
2428 err.to_string().contains("one transport mode"),
2429 "must fail on the mode mix, not incidentally: {err}"
2430 );
2431 }
2432
2433 #[test]
2434 fn a_uniformly_referenced_plane_set_is_not_rejected_as_mixed() {
2435 let referenced = BlobPlane {
2440 handle: 7,
2441 offset: 0,
2442 stride: 4,
2443 size: 4,
2444 used: 4,
2445 modifier: 0,
2446 handle_bytes: &[],
2447 data: &[],
2448 };
2449 let blob = build_blob(&[referenced, referenced]);
2450 let err = import(&blob, &[]).expect_err("an index into nothing cannot be opened");
2451 let expected = if cfg!(target_os = "windows") {
2456 "must all carry handle 0"
2457 } else {
2458 "indexes fd 7"
2459 };
2460 assert!(
2461 err.to_string().contains(expected),
2462 "a uniform plane set must clear the mode check: {err}"
2463 );
2464 }
2465
2466 #[test]
2467 fn an_inline_plane_whose_size_disagrees_with_its_data_is_refused() {
2468 let blob = export(&nv12_mem(64, 48), TransportMode::Inline).unwrap().0;
2469 let mut tampered = blob.clone();
2470 let r = region_offsets(&BlobView::parse(&blob).unwrap().header()).unwrap();
2471 tampered[r.planes + 24..r.planes + 32].copy_from_slice(&9999u64.to_le_bytes());
2472 assert!(import(&tampered, &[]).is_err());
2473 }
2474
2475 #[test]
2476 fn quantization_survives_the_round_trip() {
2477 let mut src =
2478 TensorDyn::from(Tensor::<i8>::new(&[3, 8], Some(TensorMemory::Mem), None).unwrap());
2479 src.set_quantization(crate::Quantization::per_tensor(0.5, -7))
2480 .unwrap();
2481 let blob = export(&src, TransportMode::Inline).unwrap().0;
2482 let got = import(&blob, &[]).unwrap();
2483 let q = got.quantization().expect("quantization survived");
2484 assert_eq!(q.scale(), &[0.5]);
2485 assert_eq!(q.zero_point(), Some(&[-7i32][..]));
2486 }
2487
2488 use crate::{CpuAccess, PixelFormat, Tensor, TensorDyn, TensorMemory};
2489
2490 fn nv12_mem(w: usize, h: usize) -> TensorDyn {
2491 TensorDyn::from(
2492 Tensor::<u8>::image(
2493 w,
2494 h,
2495 PixelFormat::Nv12,
2496 Some(TensorMemory::Mem),
2497 CpuAccess::ReadWrite,
2498 )
2499 .expect("alloc"),
2500 )
2501 }
2502
2503 #[test]
2504 fn export_carries_the_addressing_grid_not_the_allocation() {
2505 let blob = export(&nv12_mem(640, 480), TransportMode::Inline)
2506 .unwrap()
2507 .0;
2508 let v = BlobView::parse(&blob).unwrap();
2509 assert_eq!(v.shape(), &[480, 640]);
2511 assert_eq!(v.header().dtype, crate::DType::U8.code());
2512 assert_eq!(v.strings().format, "NV12");
2513 }
2514
2515 #[test]
2516 fn export_emits_one_plane_per_plane_table_entry() {
2517 let blob = export(&nv12_mem(640, 480), TransportMode::Inline)
2518 .unwrap()
2519 .0;
2520 let v = BlobView::parse(&blob).unwrap();
2521 let planes = v.planes().unwrap();
2522 assert_eq!(planes.len(), 2, "NV12 is two planes");
2523 assert_eq!(planes[0].size as usize, planes[0].data.len());
2525 assert_eq!(planes[1].size as usize, planes[1].data.len());
2526 assert_eq!(planes[0].size, 307_200);
2527 assert_eq!(planes[1].size, 153_600);
2528 assert_eq!(planes[1].offset, 0, "an inline plane carries no offset");
2531 }
2532
2533 #[test]
2534 fn inline_export_clears_pid_and_fence_fd() {
2535 let blob = export(&nv12_mem(64, 48), TransportMode::Inline).unwrap().0;
2538 let h = BlobView::parse(&blob).unwrap().header();
2539 assert_eq!(h.pid, 0, "no handle to reopen");
2540 assert_eq!(h.fence_fd, -1, "inlining means the producer already waited");
2541 }
2542
2543 #[test]
2544 fn inline_planes_carry_no_modifier_and_no_handle_bytes() {
2545 let blob = export(&nv12_mem(64, 48), TransportMode::Inline).unwrap().0;
2547 for p in BlobView::parse(&blob).unwrap().planes().unwrap() {
2548 assert!(p.is_inline());
2549 assert_eq!(p.modifier, 0);
2550 assert!(p.handle_bytes.is_empty());
2551 assert_eq!(p.size as usize, p.data.len());
2552 }
2553 }
2554
2555 #[test]
2556 fn reference_export_of_a_mem_tensor_is_refused() {
2557 assert!(export(&nv12_mem(64, 48), TransportMode::Reference).is_err());
2561 }
2562
2563 #[test]
2564 fn a_blob_parses_correctly_from_a_non_zero_offset() {
2565 let blob = export(&nv12_mem(64, 48), TransportMode::Inline).unwrap().0;
2568 let mut framed = vec![0xAAu8; 16];
2569 framed.extend_from_slice(&blob);
2570 let direct = BlobView::parse(&blob).unwrap();
2571 let framed_view = BlobView::parse(&framed[16..]).unwrap();
2572 assert_eq!(framed_view.header(), direct.header());
2573 assert_eq!(framed_view.shape(), direct.shape());
2574 assert_eq!(framed_view.strings(), direct.strings());
2575 }
2576
2577 #[test]
2578 fn colorimetry_axes_are_carried_as_strings() {
2579 let t = nv12_mem(64, 48).with_colorimetry(crate::Colorimetry::jfif());
2580 let blob = export(&t, TransportMode::Inline).unwrap().0;
2581 let v = BlobView::parse(&blob).unwrap();
2582 let s = v.strings();
2583 assert_eq!(s.color_space, "srgb");
2584 assert_eq!(s.color_transfer, "srgb");
2585 assert_eq!(s.color_encoding, "bt601");
2586 assert_eq!(s.color_range, "full");
2587 }
2588
2589 #[test]
2590 fn a_bare_tensor_carries_no_format_and_its_shape_is_its_grid() {
2591 let t =
2592 TensorDyn::from(Tensor::<f32>::new(&[3, 8], Some(TensorMemory::Mem), None).unwrap());
2593 let blob = export(&t, TransportMode::Inline).unwrap().0;
2594 let v = BlobView::parse(&blob).unwrap();
2595 assert_eq!(v.strings().format, "", "not an image");
2596 assert_eq!(v.shape(), &[3, 8]);
2597 assert_eq!(v.strides(), &[32i64, 4], "byte strides for f32");
2598 }
2599
2600 fn ref_plane() -> BlobPlane<'static> {
2601 BlobPlane {
2602 handle: 7,
2603 offset: 0,
2604 stride: 640,
2605 size: 307_200,
2606 used: 307_200,
2607 modifier: 0,
2608 handle_bytes: &[],
2609 data: &[],
2610 }
2611 }
2612
2613 #[test]
2614 fn a_plane_record_round_trips() {
2615 let mut out = Vec::new();
2616 write_plane(&ref_plane(), &mut out);
2617 let got = parse_planes(&out, 1).unwrap();
2618 assert_eq!(got.len(), 1);
2619 assert_eq!(got[0], ref_plane());
2620 }
2621
2622 #[test]
2623 fn a_plane_scalar_block_is_56_bytes() {
2624 let mut out = Vec::new();
2627 write_plane(&ref_plane(), &mut out);
2628 assert_eq!(out.len(), PLANE_RECORD_LEN);
2629 assert_eq!(PLANE_RECORD_LEN, 56);
2630 }
2631
2632 #[test]
2633 fn every_plane_record_starts_8_aligned_whatever_the_payload() {
2634 for n in 0..17usize {
2637 let payload: Vec<u8> = (0..n as u8).collect();
2638 let mut out = Vec::new();
2639 write_plane(
2640 &BlobPlane {
2641 handle: -1,
2642 size: n as u64,
2643 data: &payload,
2644 ..ref_plane()
2645 },
2646 &mut out,
2647 );
2648 write_plane(&ref_plane(), &mut out);
2649 assert_eq!(
2650 out.len() % 8,
2651 0,
2652 "payload of {n} bytes left the region unaligned"
2653 );
2654 let got = parse_planes(&out, 2).unwrap();
2655 assert_eq!(got[0].data, &payload[..], "payload {n} survived padding");
2656 assert_eq!(
2657 got[1],
2658 ref_plane(),
2659 "the second record parsed from an aligned offset"
2660 );
2661 }
2662 }
2663
2664 #[test]
2665 fn plane_payloads_round_trip_independently() {
2666 let hb = [0xDEu8, 0xAD, 0xBE, 0xEF];
2669 let data = [1u8, 2, 3];
2670 let mut out = Vec::new();
2671 write_plane(
2672 &BlobPlane {
2673 handle_bytes: &hb,
2674 data: &data,
2675 ..ref_plane()
2676 },
2677 &mut out,
2678 );
2679 let got = parse_planes(&out, 1).unwrap();
2680 assert_eq!(got[0].handle_bytes, &hb[..]);
2681 assert_eq!(got[0].data, &data[..]);
2682 }
2683
2684 #[test]
2685 fn a_forged_plane_payload_length_is_refused() {
2686 let mut out = vec![0u8; PLANE_RECORD_LEN];
2687 out[48..52].copy_from_slice(&u32::MAX.to_le_bytes()); assert!(matches!(
2689 parse_planes(&out, 1),
2690 Err(BlobError::ForgedLength { .. })
2691 ));
2692 }
2693
2694 #[test]
2695 fn a_plane_count_larger_than_the_records_present_is_refused() {
2696 let mut out = Vec::new();
2697 write_plane(&ref_plane(), &mut out);
2698 assert!(
2699 parse_planes(&out, 2).is_err(),
2700 "claims two records, holds one"
2701 );
2702 }
2703
2704 fn sample_strings() -> BlobStrings<'static> {
2705 BlobStrings {
2706 format: "NV12",
2707 color_space: "bt709",
2708 color_transfer: "bt709",
2709 color_encoding: "bt709",
2710 color_range: "limited",
2711 }
2712 }
2713
2714 #[test]
2715 fn the_five_strings_round_trip_in_message_order() {
2716 let mut out = Vec::new();
2717 write_strings(&sample_strings(), &mut out);
2718 let got = parse_strings(&out).unwrap();
2719 assert_eq!(got, sample_strings());
2720 }
2721
2722 #[test]
2723 fn string_order_is_positional_and_matches_the_message() {
2724 let s = BlobStrings {
2728 format: "a",
2729 color_space: "bb",
2730 color_transfer: "ccc",
2731 color_encoding: "dddd",
2732 color_range: "eeeee",
2733 };
2734 let mut out = Vec::new();
2735 write_strings(&s, &mut out);
2736 assert_eq!(parse_strings(&out).unwrap(), s);
2737 }
2738
2739 #[test]
2740 fn empty_strings_are_legal_and_distinct_from_absent() {
2741 let s = BlobStrings::default();
2743 let mut out = Vec::new();
2744 write_strings(&s, &mut out);
2745 assert_eq!(parse_strings(&out).unwrap(), s);
2746 assert!(
2747 !out.is_empty(),
2748 "five empty strings still occupy their counts"
2749 );
2750 }
2751
2752 #[test]
2753 fn the_strings_region_is_padded_so_planes_start_8_aligned() {
2754 for f in ["", "a", "ab", "abc", "NV12", "rgb8_planar"] {
2757 let mut out = Vec::new();
2758 write_strings(
2759 &BlobStrings {
2760 format: f,
2761 ..Default::default()
2762 },
2763 &mut out,
2764 );
2765 assert_eq!(out.len() % 8, 0, "region for format {f:?} is not 8-aligned");
2766 }
2767 }
2768
2769 #[test]
2770 fn a_non_utf8_string_is_refused() {
2771 let mut out = Vec::new();
2774 out.extend_from_slice(&2u32.to_le_bytes());
2775 out.extend_from_slice(&[0xff, 0xfe]);
2776 for _ in 0..4 {
2777 out.extend_from_slice(&0u32.to_le_bytes());
2778 }
2779 while out.len() % 8 != 0 {
2780 out.push(0);
2781 }
2782 assert!(matches!(
2783 parse_strings(&out),
2784 Err(BlobError::InvalidUtf8(_))
2785 ));
2786 }
2787
2788 #[test]
2789 fn a_forged_string_length_cannot_read_past_the_region() {
2790 let mut out = Vec::new();
2791 out.extend_from_slice(&u32::MAX.to_le_bytes());
2792 out.extend_from_slice(b"NV12");
2793 while out.len() % 8 != 0 {
2794 out.push(0);
2795 }
2796 assert!(matches!(
2797 parse_strings(&out),
2798 Err(BlobError::ForgedLength { .. })
2799 ));
2800 }
2801
2802 #[test]
2803 fn a_truncated_strings_region_is_refused() {
2804 assert!(parse_strings(&[0u8; 4]).is_err(), "needs five counts");
2805 }
2806
2807 #[test]
2808 fn an_overrunning_region_is_attributed_to_its_own_field() {
2809 let h = BlobHeader {
2815 size: 4096,
2816 ndim: 2,
2817 strides_len: 2,
2818 strings_bytes: u32::MAX,
2819 ..BlobHeader::empty()
2820 };
2821 assert!(
2822 matches!(
2823 region_offsets(&h),
2824 Err(BlobError::ForgedLength {
2825 field: "strings_bytes",
2826 ..
2827 })
2828 ),
2829 "the overrun must name strings_bytes, not a later field"
2830 );
2831 }
2832
2833 #[test]
2834 fn header_is_72_bytes_and_every_scalar_is_naturally_aligned() {
2835 assert_eq!(HEADER_LEN, 72);
2836 assert_eq!(HEADER_LEN % 8, 0, "the variable tail starts 8-aligned");
2837 for off in [0usize, 8, 16] {
2838 assert_eq!(off % 8, 0, "u64 field at {off} is not 8-aligned");
2839 }
2840 for off in [24usize, 28, 32, 36, 40, 44, 48, 52, 56, 60, 64] {
2841 assert_eq!(off % 4, 0, "u32 field at {off} is not 4-aligned");
2842 }
2843 }
2844
2845 #[test]
2846 fn a_header_round_trips_through_bytes() {
2847 let h = BlobHeader {
2848 size: 72,
2849 required_mask: 0,
2850 planes_bytes: 0,
2851 storage_kind: 2,
2852 pid: 4242,
2853 fence_fd: -1,
2854 dtype: 1,
2855 quant_axis: -2,
2856 ndim: 2,
2857 plane_count: 0,
2858 quant_scales_len: 0,
2859 quant_zero_points_len: 0,
2860 strings_bytes: 0,
2861 strides_len: 2,
2862 };
2863 let mut buf = vec![0u8; HEADER_LEN];
2864 h.write_to(&mut buf);
2865 assert_eq!(parse_header(&buf).unwrap(), h);
2866 }
2867
2868 #[test]
2869 fn the_header_is_little_endian_on_every_host() {
2870 let mut buf = vec![0u8; HEADER_LEN];
2871 BlobHeader {
2872 size: 0x0102_0304_0506_0708,
2873 ..BlobHeader::empty()
2874 }
2875 .write_to(&mut buf);
2876 assert_eq!(
2877 &buf[0..8],
2878 &[0x08, 0x07, 0x06, 0x05, 0x04, 0x03, 0x02, 0x01],
2879 "size must be little-endian regardless of host byte order"
2880 );
2881 }
2882
2883 #[test]
2884 fn a_truncated_buffer_is_refused_not_read() {
2885 assert!(matches!(
2886 parse_header(&[0u8; 71]),
2887 Err(BlobError::Truncated { .. })
2888 ));
2889 }
2890
2891 #[test]
2892 fn an_unknown_required_bit_is_refused() {
2893 let mut buf = vec![0u8; HEADER_LEN];
2894 BlobHeader {
2895 size: 72,
2896 required_mask: 1 << 63,
2897 ..BlobHeader::empty()
2898 }
2899 .write_to(&mut buf);
2900 assert!(matches!(
2901 parse_header(&buf),
2902 Err(BlobError::UnsupportedRequirement(_))
2903 ));
2904 }
2905
2906 #[test]
2907 fn a_size_smaller_than_the_header_is_refused() {
2908 let mut buf = vec![0u8; HEADER_LEN];
2909 BlobHeader {
2910 size: 8,
2911 ..BlobHeader::empty()
2912 }
2913 .write_to(&mut buf);
2914 assert!(
2915 parse_header(&buf).is_err(),
2916 "size must cover at least the header"
2917 );
2918 }
2919
2920 #[test]
2921 fn blob_error_display_names_the_field_and_the_values() {
2922 let cases = [
2923 (
2924 BlobError::Truncated { need: 80, have: 72 },
2925 "blob truncated: needs 80 bytes, buffer has 72",
2926 ),
2927 (
2928 BlobError::UnsupportedRequirement(0x10),
2929 "blob requires unsupported features (required_mask 0x10); refusing rather \
2930 than misreading it",
2931 ),
2932 (
2933 BlobError::ForgedLength {
2934 field: "ndim",
2935 value: 9,
2936 },
2937 "blob field `ndim` has implausible value 9",
2938 ),
2939 (
2940 BlobError::InvalidUtf8("format"),
2941 "blob string `format` is not valid UTF-8",
2942 ),
2943 (
2944 BlobError::Malformed("no planes"),
2945 "malformed blob: no planes",
2946 ),
2947 ];
2948 for (err, want) in cases {
2949 assert_eq!(err.to_string(), want);
2950 }
2951 }
2952
2953 #[test]
2954 fn a_size_past_the_end_of_the_buffer_is_truncation() {
2955 let mut buf = vec![0u8; HEADER_LEN];
2956 BlobHeader {
2957 size: 80,
2958 ..BlobHeader::empty()
2959 }
2960 .write_to(&mut buf);
2961 assert_eq!(
2962 parse_header(&buf),
2963 Err(BlobError::Truncated { need: 80, have: 72 })
2964 );
2965 }
2966
2967 #[test]
2968 fn bytes_after_the_blob_are_not_part_of_it() {
2969 let blob = export(&nv12_mem(64, 48), TransportMode::Inline).unwrap().0;
2972 let mut framed = blob.clone();
2973 framed.extend_from_slice(&[0xAA; 16]);
2974 let v = BlobView::parse(&framed).expect("a buffer longer than `size` parses");
2975 assert_eq!(v.header().size, blob.len() as u64);
2976 assert_eq!(v.planes().unwrap().len(), 2);
2977 }
2978
2979 #[test]
2980 fn a_header_only_blob_has_every_region_empty_and_ending_at_the_header() {
2981 let r = region_offsets(&BlobHeader {
2982 size: HEADER_LEN as u64,
2983 ..BlobHeader::empty()
2984 })
2985 .expect("a blob of exactly one header is well-formed");
2986 assert_eq!(
2987 r,
2988 Regions {
2989 shape: HEADER_LEN,
2990 strides: HEADER_LEN,
2991 quant_scales: HEADER_LEN,
2992 quant_zero_points: HEADER_LEN,
2993 strings: HEADER_LEN,
2994 planes: HEADER_LEN,
2995 end: HEADER_LEN,
2996 }
2997 );
2998 }
2999
3000 #[test]
3001 fn a_size_below_the_header_is_attributed_to_size() {
3002 assert_eq!(
3003 region_offsets(&BlobHeader {
3004 size: 8,
3005 ..BlobHeader::empty()
3006 }),
3007 Err(BlobError::ForgedLength {
3008 field: "size",
3009 value: 8
3010 })
3011 );
3012 }
3013
3014 #[test]
3015 fn plane_count_is_bounded_by_the_whole_records_that_fit() {
3016 let header = |plane_count| BlobHeader {
3017 size: 4096,
3018 planes_bytes: 2 * PLANE_RECORD_LEN as u64,
3019 plane_count,
3020 ..BlobHeader::empty()
3021 };
3022 assert!(
3023 region_offsets(&header(2)).is_ok(),
3024 "two minimum-size records fit in {} bytes",
3025 2 * PLANE_RECORD_LEN
3026 );
3027 assert_eq!(
3028 region_offsets(&header(3)),
3029 Err(BlobError::ForgedLength {
3030 field: "plane_count",
3031 value: 3
3032 })
3033 );
3034 }
3035
3036 #[test]
3037 fn the_strings_region_ends_8_aligned_in_the_blob_wherever_it_starts() {
3038 let formats = ["", "a", "abc", "NV12", "rgb8_planar"].map(|format| BlobStrings {
3041 format,
3042 ..Default::default()
3043 });
3044 let all_five = BlobStrings {
3047 format: "a",
3048 color_space: "bb",
3049 color_transfer: "cccc",
3050 color_encoding: "dddddddd",
3051 color_range: "eeeeeeeeeeeeeeee",
3052 };
3053 for start in [0usize, 3, 4] {
3054 for strings in formats.iter().chain([&all_five]) {
3055 let mut out = vec![0xAAu8; start];
3056 write_strings(strings, &mut out);
3057 assert_eq!(out.len() % 8, 0, "start {start}, {strings:?}");
3058 assert_eq!(
3059 out.len() - start,
3060 strings_encoded_len(strings, start),
3061 "start {start}, {strings:?}"
3062 );
3063 assert_eq!(parse_strings(&out[start..]).unwrap(), *strings);
3064 }
3065 }
3066 }
3067
3068 #[test]
3071 fn an_export_whose_quantization_ends_4_aligned_still_parses() {
3072 let mut src =
3073 TensorDyn::from(Tensor::<i8>::new(&[3, 8], Some(TensorMemory::Mem), None).unwrap());
3074 src.set_quantization(crate::Quantization::per_tensor_symmetric(0.5))
3075 .unwrap();
3076 let blob = export(&src, TransportMode::Inline).unwrap().0;
3077 let v = BlobView::parse(&blob).expect("export must write a blob its parser accepts");
3078 assert_eq!(v.planes().unwrap().len(), 1);
3079 }
3080
3081 #[test]
3082 fn encoded_len_is_exactly_what_write_plane_appends() {
3083 let handle_bytes = [0xDEu8, 0xAD, 0xBE, 0xEF];
3084 for hb in [&[][..], &handle_bytes[..]] {
3085 for n in 0..17usize {
3086 let data: Vec<u8> = (0..n as u8).collect();
3087 let plane = BlobPlane {
3088 handle_bytes: hb,
3089 data: &data,
3090 ..ref_plane()
3091 };
3092 let mut out = vec![0xAAu8; 3];
3094 write_plane(&plane, &mut out);
3095 assert_eq!(
3096 out.len() - 3,
3097 plane.encoded_len(),
3098 "{} handle bytes, {n} data bytes",
3099 hb.len()
3100 );
3101 }
3102 }
3103 }
3104
3105 #[test]
3106 fn a_plane_is_inline_only_when_its_handle_is_negative() {
3107 assert!(
3108 !BlobPlane {
3109 handle: 0,
3110 ..ref_plane()
3111 }
3112 .is_inline(),
3113 "index 0 is the first fd-table entry, a reference"
3114 );
3115 assert!(BlobPlane {
3116 handle: -1,
3117 ..ref_plane()
3118 }
3119 .is_inline());
3120 }
3121
3122 #[test]
3126 fn a_reference_plane_keeps_its_offset_and_extent() {
3127 let geoms = [crate::PlaneGeometry {
3128 offset: 4096,
3129 stride: 64,
3130 size: 3072,
3131 }];
3132 let mut buf = Vec::new();
3133 write_export_planes(&geoms, TransportMode::Reference, 0, &[], &[], &mut buf).unwrap();
3134 let planes = parse_planes(&buf, 1).unwrap();
3135 let p = planes[0];
3136 assert_eq!(
3137 (p.handle, p.offset, p.stride, p.size, p.used),
3138 (0, 4096, 64, 3072, 3072)
3139 );
3140 assert!(p.data.is_empty(), "a reference carries no bytes");
3141 }
3142
3143 #[test]
3146 fn every_layout_round_trips_inline_through_its_grid() {
3147 for (format, grid) in [
3148 (PixelFormat::Rgb, vec![48u64, 64, 3]),
3149 (PixelFormat::Grey, vec![48, 64]),
3150 (PixelFormat::PlanarRgb, vec![3, 48, 64]),
3151 ] {
3152 let src = TensorDyn::from(
3153 Tensor::<u8>::image(
3154 64,
3155 48,
3156 format,
3157 Some(TensorMemory::Mem),
3158 CpuAccess::ReadWrite,
3159 )
3160 .expect("alloc"),
3161 );
3162 fill_ramp(&src);
3163 let blob = export(&src, TransportMode::Inline).unwrap().0;
3164 assert_eq!(BlobView::parse(&blob).unwrap().shape(), grid, "{format:?}");
3165 let got = import(&blob, &[]).unwrap();
3166 assert_eq!(got.format(), Some(format));
3167 assert_eq!(got.shape(), src.shape(), "{format:?}");
3168 assert_eq!(bytes_of(&got), bytes_of(&src), "{format:?}");
3169 }
3170 }
3171
3172 fn inline_plane(data: &[u8]) -> BlobPlane<'_> {
3173 BlobPlane {
3174 handle: -1,
3175 offset: 0,
3176 stride: data.len() as u64,
3177 size: data.len() as u64,
3178 used: data.len() as u64,
3179 modifier: 0,
3180 handle_bytes: &[],
3181 data,
3182 }
3183 }
3184
3185 #[test]
3186 fn a_well_formed_hand_built_inline_blob_imports() {
3187 let got = import(&build_blob(&[inline_plane(&[7])]), &[]).expect("well-formed");
3189 assert_eq!(got.shape(), &[1]);
3190 assert_eq!(bytes_of(&got), vec![7]);
3191 }
3192
3193 #[test]
3194 fn an_inline_plane_with_a_modifier_or_handle_bytes_is_refused() {
3195 let data = [7u8];
3196 for plane in [
3197 BlobPlane {
3198 modifier: 1,
3199 ..inline_plane(&data)
3200 },
3201 BlobPlane {
3202 handle_bytes: &[1],
3203 ..inline_plane(&data)
3204 },
3205 ] {
3206 let err = import(&build_blob(&[plane]), &[]).expect_err("must be refused");
3207 assert!(
3208 err.to_string()
3209 .contains("inline plane must have modifier 0 and no handle_bytes"),
3210 "got: {err}"
3211 );
3212 }
3213 }
3214
3215 #[test]
3216 fn a_referenced_plane_carrying_bytes_is_refused() {
3217 let plane = BlobPlane {
3218 handle: 0,
3219 ..inline_plane(&[7])
3220 };
3221 let err = import(&build_blob(&[plane]), &[]).expect_err("must be refused");
3222 assert!(
3223 err.to_string()
3224 .contains("referenced plane must not also carry inline bytes"),
3225 "got: {err}"
3226 );
3227 }
3228
3229 #[test]
3230 fn a_grid_too_short_for_its_format_is_refused() {
3231 for (format, grid) in [
3232 (PixelFormat::Nv12, &[1u64][..]),
3233 (PixelFormat::PlanarRgb, &[48, 64][..]),
3234 ] {
3235 let data = [0u8];
3236 let planes = [inline_plane(&data)];
3237 let blob = build(&BlobSpec {
3238 grid,
3239 strings: BlobStrings {
3240 format: format.as_str(),
3241 ..Default::default()
3242 },
3243 ..BlobSpec::new(&planes)
3244 });
3245 let err = import(&blob, &[]).expect_err("must be refused");
3246 let want = format!("grid {grid:?} is not an image for {format:?}");
3247 assert!(
3248 matches!(&err, crate::Error::InvalidShape(m) if m.contains(&want)),
3249 "{format:?}: got {err:?}"
3250 );
3251 }
3252 }
3253
3254 #[test]
3255 fn an_unquantized_export_declares_no_quantization() {
3256 let src =
3257 TensorDyn::from(Tensor::<i8>::new(&[3, 8], Some(TensorMemory::Mem), None).unwrap());
3258 let blob = export(&src, TransportMode::Inline).unwrap().0;
3259 let h = BlobView::parse(&blob).unwrap().header();
3260 assert_eq!(h.quant_axis, -2);
3261 assert_eq!((h.quant_scales_len, h.quant_zero_points_len), (0, 0));
3262 assert!(import(&blob, &[]).unwrap().quantization().is_none());
3263 }
3264
3265 fn round_trip_quantization(q: crate::Quantization) -> crate::Quantization {
3266 let mut src =
3267 TensorDyn::from(Tensor::<i8>::new(&[2, 2, 3], Some(TensorMemory::Mem), None).unwrap());
3268 src.set_quantization(q).unwrap();
3269 let blob = export(&src, TransportMode::Inline).unwrap().0;
3270 import(&blob, &[])
3271 .unwrap()
3272 .quantization()
3273 .expect("quantization survived")
3274 .clone()
3275 }
3276
3277 #[test]
3278 fn per_tensor_symmetric_quantization_survives_without_a_zero_point() {
3279 let q = round_trip_quantization(crate::Quantization::per_tensor_symmetric(0.5));
3280 assert_eq!(q.scale(), &[0.5]);
3281 assert_eq!(q.zero_point(), None);
3282 assert_eq!(q.axis(), None);
3283 }
3284
3285 #[test]
3286 fn per_channel_symmetric_quantization_survives_with_its_axis() {
3287 let q = round_trip_quantization(
3288 crate::Quantization::per_channel_symmetric(vec![0.1, 0.2, 0.3], 2).unwrap(),
3289 );
3290 assert_eq!(q.scale(), &[0.1, 0.2, 0.3]);
3291 assert_eq!(q.zero_point(), None);
3292 assert_eq!(q.axis(), Some(2));
3293 }
3294
3295 #[test]
3296 fn per_channel_asymmetric_quantization_keeps_its_zero_points() {
3297 let q = round_trip_quantization(
3298 crate::Quantization::per_channel(vec![0.1, 0.2, 0.3], vec![-3, 0, 7], 2).unwrap(),
3299 );
3300 assert_eq!(q.scale(), &[0.1, 0.2, 0.3]);
3301 assert_eq!(q.zero_point(), Some(&[-3, 0, 7][..]));
3302 assert_eq!(q.axis(), Some(2));
3303 }
3304
3305 #[test]
3308 fn declared_but_malformed_quantization_is_refused() {
3309 let data = [0u8; 3];
3310 let planes = [inline_plane(&data)];
3311 let cases: [(i32, &[f32], &[i32], &str); 4] = [
3312 (-3, &[0.5], &[], "quant_axis -3 is neither"),
3313 (
3314 -1,
3315 &[],
3316 &[],
3317 "needs 1 scale and 0 or 1 zero points, got 0 and 0",
3318 ),
3319 (-1, &[0.5, 0.25], &[], "got 2 and 0"),
3320 (0, &[0.5, 0.25, 0.125], &[1, 2], "zero_point.len"),
3321 ];
3322 for (quant_axis, scales, zeros, want) in cases {
3323 let blob = build(&BlobSpec {
3324 dtype: crate::DType::I8,
3325 grid: &[3],
3326 quant_axis,
3327 scales,
3328 zeros,
3329 ..BlobSpec::new(&planes)
3330 });
3331 let err = import(&blob, &[]).expect_err("malformed quantization must be refused");
3332 let msg = err.to_string();
3333 assert!(
3334 msg.contains("blob quantization is malformed") && msg.contains(want),
3335 "axis {quant_axis}, {} scales, {} zeros: got {msg}",
3336 scales.len(),
3337 zeros.len()
3338 );
3339 }
3340 }
3341
3342 #[cfg(all(unix, not(any(target_os = "macos", target_os = "ios"))))]
3345 mod shm_reference {
3346 use super::*;
3347 use std::os::fd::AsRawFd;
3348
3349 fn reference_planes(
3352 format: PixelFormat,
3353 w: usize,
3354 h: usize,
3355 pitch: usize,
3356 ) -> Vec<BlobPlane<'static>> {
3357 format
3358 .plane_table(w, h, pitch)
3359 .unwrap()
3360 .iter()
3361 .map(|g| BlobPlane {
3362 handle: 0,
3363 offset: g.offset,
3364 stride: g.stride,
3365 size: g.size,
3366 used: g.size,
3367 modifier: 0,
3368 handle_bytes: &[],
3369 data: &[],
3370 })
3371 .collect()
3372 }
3373
3374 fn nv12_spec<'a>(planes: &'a [BlobPlane<'a>]) -> BlobSpec<'a> {
3375 BlobSpec {
3376 storage_kind: TensorMemory::Shm.code(),
3377 grid: &[48, 64],
3378 strings: BlobStrings {
3379 format: "NV12",
3380 ..Default::default()
3381 },
3382 ..BlobSpec::new(planes)
3383 }
3384 }
3385
3386 fn wide_shm() -> TensorDyn {
3389 TensorDyn::from(
3390 Tensor::<u8>::image(
3391 128,
3392 48,
3393 PixelFormat::Nv12,
3394 Some(TensorMemory::Shm),
3395 CpuAccess::ReadWrite,
3396 )
3397 .expect("shm allocation"),
3398 )
3399 }
3400
3401 #[test]
3402 fn a_referenced_shm_blob_aliases_its_producer_at_the_producers_pitch() {
3403 let src = wide_shm();
3404 fill_ramp(&src);
3405 let fd = src.clone_fd().expect("shm fd");
3406 let planes = reference_planes(PixelFormat::Nv12, 64, 48, 128);
3407 let got = import(&build(&nv12_spec(&planes)), &[fd.as_raw_fd()])
3408 .expect("a Shm reference imports");
3409 assert_eq!(got.memory(), TensorMemory::Shm);
3410 assert_eq!(got.format(), Some(PixelFormat::Nv12));
3411 assert_eq!(got.shape(), &[72, 64], "the allocation shape of the grid");
3412 assert_eq!(
3413 got.row_stride(),
3414 Some(128),
3415 "the producer's pitch, which the shape cannot express"
3416 );
3417 let expected = bytes_of(&src);
3418 let seen = bytes_of(&got);
3419 assert_eq!(
3420 seen[..],
3421 expected[..seen.len()],
3422 "the same bytes, not a copy"
3423 );
3424 }
3425
3426 #[test]
3427 fn a_zero_pitch_means_none_was_recorded() {
3428 let src = wide_shm();
3429 let fd = src.clone_fd().expect("shm fd");
3430 let mut planes = reference_planes(PixelFormat::Nv12, 64, 48, 128);
3431 planes[0].stride = 0;
3432 let got = import(&build(&nv12_spec(&planes)), &[fd.as_raw_fd()])
3433 .expect("a zero pitch is not an error");
3434 assert_ne!(got.row_stride(), Some(0));
3435 }
3436
3437 #[test]
3438 fn a_pitch_the_format_refuses_is_refused_not_ignored() {
3439 let src = wide_shm();
3440 let fd = src.clone_fd().expect("shm fd");
3441 let mut planes = reference_planes(PixelFormat::Nv12, 64, 48, 128);
3442 planes[0].stride = 32;
3443 let err = import(&build(&nv12_spec(&planes)), &[fd.as_raw_fd()])
3444 .expect_err("a pitch narrower than a row is malformed");
3445 assert!(
3446 err.to_string().contains("row_stride 32 < minimum 64"),
3447 "got: {err}"
3448 );
3449 }
3450
3451 #[test]
3452 fn every_plane_index_is_bounded_not_only_the_first() {
3453 let src = wide_shm();
3454 let fd = src.clone_fd().expect("shm fd");
3455 let mut planes = reference_planes(PixelFormat::Nv12, 64, 48, 128);
3456 planes[1].handle = 5;
3457 let err = import(&build(&nv12_spec(&planes)), &[fd.as_raw_fd()])
3458 .expect_err("index 5 is past a one-entry table");
3459 assert!(
3460 err.to_string()
3461 .contains("plane handle indexes fd 5 but only 1 were provided"),
3462 "got: {err}"
3463 );
3464 }
3465
3466 #[test]
3467 fn quantization_survives_a_reference_round_trip() {
3468 let src = TensorDyn::from(
3469 Tensor::<i8>::new(&[2, 2, 3], Some(TensorMemory::Shm), None).expect("shm"),
3470 );
3471 let fd = src.clone_fd().expect("shm fd");
3472 let bytes = src.capacity_bytes() as u64;
3473 let planes = [BlobPlane {
3474 handle: 0,
3475 offset: 0,
3476 stride: bytes,
3477 size: bytes,
3478 used: bytes,
3479 modifier: 0,
3480 handle_bytes: &[],
3481 data: &[],
3482 }];
3483 let blob = build(&BlobSpec {
3484 storage_kind: TensorMemory::Shm.code(),
3485 dtype: crate::DType::I8,
3486 grid: &[2, 2, 3],
3487 quant_axis: 2,
3488 scales: &[0.1, 0.2, 0.3],
3489 zeros: &[-3, 0, 7],
3490 ..BlobSpec::new(&planes)
3491 });
3492 let got = import(&blob, &[fd.as_raw_fd()]).expect("a Shm reference imports");
3493 let q = got.quantization().expect("quantization survived");
3494 assert_eq!(q.scale(), &[0.1, 0.2, 0.3]);
3495 assert_eq!(q.zero_point(), Some(&[-3, 0, 7][..]));
3496 assert_eq!(q.axis(), Some(2));
3497 }
3498 }
3499
3500 #[cfg(target_os = "linux")]
3504 #[test]
3505 fn a_padded_planar_reference_export_steps_planes_by_the_pitch() {
3506 const WHAT: &str = "a_padded_planar_reference_export_steps_planes_by_the_pitch";
3507 if !crate::test_support::dma_or_skip(WHAT) {
3508 return;
3509 }
3510 let src = TensorDyn::from(
3512 Tensor::<u8>::image(
3513 50,
3514 48,
3515 PixelFormat::PlanarRgb,
3516 Some(TensorMemory::DmaBuf),
3517 CpuAccess::ReadWrite,
3518 )
3519 .expect("a usable DMA heap allocates a planar image"),
3520 );
3521 let pitch = src.effective_row_stride().expect("an image has a pitch");
3522 assert!(pitch > 50, "precondition: the pitch pads, got {pitch}");
3523 let (blob, _fds) = export(&src, TransportMode::Reference).unwrap();
3524 let v = BlobView::parse(&blob).unwrap();
3525 let planes = v.planes().unwrap();
3526 assert_eq!(
3527 v.strides(),
3528 &[(pitch * 48) as i64, pitch as i64, 1],
3529 "row stride is the pitch; plane stride is pitch * height"
3530 );
3531 assert_eq!(
3532 planes[1].offset as i64,
3533 v.strides()[0],
3534 "the stride array and the plane table agree on where plane 1 starts"
3535 );
3536 }
3537}