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