1use crate::{DType, PixelFormat, Tensor, TensorMemory, TensorTrait};
5use half::f16;
6use std::fmt;
7
8#[non_exhaustive]
10pub enum TensorDyn {
11 U8(Tensor<u8>),
13 I8(Tensor<i8>),
15 U16(Tensor<u16>),
17 I16(Tensor<i16>),
19 U32(Tensor<u32>),
21 I32(Tensor<i32>),
23 U64(Tensor<u64>),
25 I64(Tensor<i64>),
27 F16(Tensor<f16>),
29 F32(Tensor<f32>),
31 F64(Tensor<f64>),
33}
34
35#[allow(clippy::useless_transmute)]
60unsafe fn retag_same_width(t: TensorDyn, dtype: DType) -> TensorDyn {
61 macro_rules! group {
70 ($inner:expr, $src:ty, $same:ident, $( $d:ident => $ty:ty ),+ $(,)?) => {
71 match dtype {
72 $(
73 DType::$d => TensorDyn::$d(unsafe {
76 std::mem::transmute::<Tensor<$src>, Tensor<$ty>>($inner)
77 }),
78 )+
79 _ => TensorDyn::$same($inner),
80 }
81 };
82 }
83 match t {
84 TensorDyn::U8(i) => group!(i, u8, U8, U8 => u8, I8 => i8),
85 TensorDyn::I8(i) => group!(i, i8, I8, U8 => u8, I8 => i8),
86 TensorDyn::U16(i) => group!(i, u16, U16, U16 => u16, I16 => i16, F16 => f16),
87 TensorDyn::I16(i) => group!(i, i16, I16, U16 => u16, I16 => i16, F16 => f16),
88 TensorDyn::F16(i) => group!(i, f16, F16, U16 => u16, I16 => i16, F16 => f16),
89 TensorDyn::U32(i) => group!(i, u32, U32, U32 => u32, I32 => i32, F32 => f32),
90 TensorDyn::I32(i) => group!(i, i32, I32, U32 => u32, I32 => i32, F32 => f32),
91 TensorDyn::F32(i) => group!(i, f32, F32, U32 => u32, I32 => i32, F32 => f32),
92 TensorDyn::U64(i) => group!(i, u64, U64, U64 => u64, I64 => i64, F64 => f64),
93 TensorDyn::I64(i) => group!(i, i64, I64, U64 => u64, I64 => i64, F64 => f64),
94 TensorDyn::F64(i) => group!(i, f64, F64, U64 => u64, I64 => i64, F64 => f64),
95 }
96}
97
98macro_rules! dispatch {
99 ($self:expr, $method:ident $(, $arg:expr)*) => {
100 match $self {
101 TensorDyn::U8(t) => t.$method($($arg),*),
102 TensorDyn::I8(t) => t.$method($($arg),*),
103 TensorDyn::U16(t) => t.$method($($arg),*),
104 TensorDyn::I16(t) => t.$method($($arg),*),
105 TensorDyn::U32(t) => t.$method($($arg),*),
106 TensorDyn::I32(t) => t.$method($($arg),*),
107 TensorDyn::U64(t) => t.$method($($arg),*),
108 TensorDyn::I64(t) => t.$method($($arg),*),
109 TensorDyn::F16(t) => t.$method($($arg),*),
110 TensorDyn::F32(t) => t.$method($($arg),*),
111 TensorDyn::F64(t) => t.$method($($arg),*),
112 }
113 };
114}
115
116macro_rules! dyn_fanout {
120 ($self:expr, $method:ident $(, $arg:expr)*) => {
121 match $self {
122 TensorDyn::U8(t) => t.$method($($arg),*).map(TensorDyn::U8),
123 TensorDyn::I8(t) => t.$method($($arg),*).map(TensorDyn::I8),
124 TensorDyn::U16(t) => t.$method($($arg),*).map(TensorDyn::U16),
125 TensorDyn::I16(t) => t.$method($($arg),*).map(TensorDyn::I16),
126 TensorDyn::U32(t) => t.$method($($arg),*).map(TensorDyn::U32),
127 TensorDyn::I32(t) => t.$method($($arg),*).map(TensorDyn::I32),
128 TensorDyn::U64(t) => t.$method($($arg),*).map(TensorDyn::U64),
129 TensorDyn::I64(t) => t.$method($($arg),*).map(TensorDyn::I64),
130 TensorDyn::F16(t) => t.$method($($arg),*).map(TensorDyn::F16),
131 TensorDyn::F32(t) => t.$method($($arg),*).map(TensorDyn::F32),
132 TensorDyn::F64(t) => t.$method($($arg),*).map(TensorDyn::F64),
133 }
134 };
135}
136
137macro_rules! downcast_methods {
139 ($variant:ident, $ty:ty, $as_name:ident, $as_mut_name:ident, $into_name:ident) => {
140 pub fn $as_name(&self) -> Option<&Tensor<$ty>> {
142 match self {
143 Self::$variant(t) => Some(t),
144 _ => None,
145 }
146 }
147
148 pub fn $as_mut_name(&mut self) -> Option<&mut Tensor<$ty>> {
150 match self {
151 Self::$variant(t) => Some(t),
152 _ => None,
153 }
154 }
155
156 #[allow(clippy::result_large_err)]
159 pub fn $into_name(self) -> Result<Tensor<$ty>, Self> {
160 match self {
161 Self::$variant(t) => Ok(t),
162 other => Err(other),
163 }
164 }
165 };
166}
167
168impl TensorDyn {
169 pub fn dtype(&self) -> DType {
171 match self {
172 Self::U8(_) => DType::U8,
173 Self::I8(_) => DType::I8,
174 Self::U16(_) => DType::U16,
175 Self::I16(_) => DType::I16,
176 Self::U32(_) => DType::U32,
177 Self::I32(_) => DType::I32,
178 Self::U64(_) => DType::U64,
179 Self::I64(_) => DType::I64,
180 Self::F16(_) => DType::F16,
181 Self::F32(_) => DType::F32,
182 Self::F64(_) => DType::F64,
183 }
184 }
185
186 pub fn pin_host(&self, access: crate::CpuAccess) -> crate::Result<crate::HostPin<'static>> {
203 dispatch!(self, pin_host, access)
204 }
205
206 pub fn map_bytes(
239 &self,
240 access: crate::CpuAccess,
241 ) -> crate::Result<crate::view::HostView<'static, u8>> {
242 use crate::TensorTrait;
243 match self {
244 Self::U8(t) => t.map_with(access).map(crate::view::HostView::into_bytes),
245 Self::I8(t) => t.map_with(access).map(crate::view::HostView::into_bytes),
246 Self::U16(t) => t.map_with(access).map(crate::view::HostView::into_bytes),
247 Self::I16(t) => t.map_with(access).map(crate::view::HostView::into_bytes),
248 Self::U32(t) => t.map_with(access).map(crate::view::HostView::into_bytes),
249 Self::I32(t) => t.map_with(access).map(crate::view::HostView::into_bytes),
250 Self::U64(t) => t.map_with(access).map(crate::view::HostView::into_bytes),
251 Self::I64(t) => t.map_with(access).map(crate::view::HostView::into_bytes),
252 Self::F16(t) => t.map_with(access).map(crate::view::HostView::into_bytes),
253 Self::F32(t) => t.map_with(access).map(crate::view::HostView::into_bytes),
254 Self::F64(t) => t.map_with(access).map(crate::view::HostView::into_bytes),
255 }
256 }
257
258 pub fn sync_for_cpu(&self, access: crate::CpuAccess) -> crate::Result<()> {
260 dispatch!(self, sync_for_cpu, access)
261 }
262
263 pub fn sync_for_device(&self, access: crate::CpuAccess) -> crate::Result<()> {
265 dispatch!(self, sync_for_device, access)
266 }
267
268 pub fn descriptor(&self) -> crate::TensorDesc {
269 self.descriptor_pinned(None)
270 }
271
272 pub fn descriptor_pinned(&self, pin: Option<&crate::HostPin<'_>>) -> crate::TensorDesc {
280 let handle: i64 = match self.memory() {
285 #[cfg(target_os = "linux")]
286 TensorMemory::DmaBuf => {
287 use std::os::fd::AsRawFd;
288 self.dmabuf().map(|fd| fd.as_raw_fd() as i64).unwrap_or(-1)
289 }
290 #[cfg(any(target_os = "macos", target_os = "ios"))]
291 TensorMemory::DmaBuf => self.iosurface_id().map(|id| id as i64).unwrap_or(-1),
292 TensorMemory::Pbo => self.pbo_id().map(|id| id as i64).unwrap_or(-1),
293 _ => -1,
294 };
295 crate::protocol::from_parts(crate::protocol::DescParts {
296 dims: self.shape(),
297 memory: self.memory(),
298 dtype: self.dtype(),
299 fourcc: self.format().map(|f| f.to_fourcc()).unwrap_or(0),
300 format: self.format(),
301 row_stride: self.row_stride(),
302 handle,
303 colorimetry: self.colorimetry().map(|c| c.pack()).unwrap_or(0),
304 capacity: self.capacity_bytes() as u64,
305 pin,
306 pbo_vtable_ptr: self.pbo_vtable_ptr(),
307 })
308 }
309
310 pub fn shape(&self) -> &[usize] {
311 dispatch!(self, shape)
312 }
313
314 pub fn set_logical_shape(&mut self, shape: &[usize]) -> crate::Result<()> {
318 dispatch!(self, set_logical_shape, shape)
319 }
320
321 pub fn capacity_bytes(&self) -> usize {
323 dispatch!(self, capacity_bytes)
324 }
325
326 pub fn name(&self) -> String {
328 dispatch!(self, name)
329 }
330
331 pub fn format(&self) -> Option<PixelFormat> {
333 dispatch!(self, format)
334 }
335
336 pub fn width(&self) -> Option<usize> {
338 dispatch!(self, width)
339 }
340
341 pub fn height(&self) -> Option<usize> {
343 dispatch!(self, height)
344 }
345
346 pub fn size(&self) -> usize {
348 dispatch!(self, size)
349 }
350
351 pub fn set_dtype(&mut self, dtype: DType) -> crate::Result<()> {
370 let current = self.dtype();
371 if current == dtype {
372 return Ok(());
373 }
374 if current.size() != dtype.size() {
375 return Err(crate::Error::InvalidArgument(format!(
376 "set_dtype: {current:?} and {dtype:?} are different widths \
377 ({} vs {} bytes); retagging would change the element count \
378 over an allocation whose size did not change",
379 current.size(),
380 dtype.size()
381 )));
382 }
383 unsafe {
391 let taken = std::ptr::read(self);
392 std::ptr::write(self, retag_same_width(taken, dtype));
393 }
394 Ok(())
395 }
396
397 pub fn memory(&self) -> TensorMemory {
399 dispatch!(self, memory)
400 }
401
402 pub fn cpu_access(&self) -> crate::CpuAccess {
414 dispatch!(self, cpu_access)
415 }
416
417 pub fn reshape(&mut self, shape: &[usize]) -> crate::Result<()> {
419 dispatch!(self, reshape, shape)
420 }
421
422 pub fn set_format(&mut self, format: PixelFormat) -> crate::Result<()> {
440 dispatch!(self, set_format, format)
441 }
442
443 pub fn with_format(mut self, format: PixelFormat) -> crate::Result<Self> {
460 self.set_format(format)?;
461 Ok(self)
462 }
463
464 pub fn colorimetry(&self) -> Option<crate::Colorimetry> {
466 dispatch!(self, colorimetry)
467 }
468
469 pub fn set_colorimetry(&mut self, c: Option<crate::Colorimetry>) {
471 dispatch!(self, set_colorimetry, c)
472 }
473
474 pub fn with_colorimetry(mut self, c: crate::Colorimetry) -> Self {
476 self.set_colorimetry(Some(c));
477 self
478 }
479
480 pub fn row_stride(&self) -> Option<usize> {
482 dispatch!(self, row_stride)
483 }
484
485 pub fn effective_row_stride(&self) -> Option<usize> {
487 dispatch!(self, effective_row_stride)
488 }
489
490 pub fn configure_image(
493 &mut self,
494 width: usize,
495 height: usize,
496 format: PixelFormat,
497 ) -> crate::Result<()> {
498 dispatch!(self, configure_image, width, height, format)
499 }
500
501 pub fn set_row_stride(&mut self, stride: usize) -> crate::Result<()> {
507 dispatch!(self, set_row_stride, stride)
508 }
509
510 pub fn with_row_stride(mut self, stride: usize) -> crate::Result<Self> {
512 self.set_row_stride(stride)?;
513 Ok(self)
514 }
515
516 pub fn set_row_stride_unchecked(&mut self, stride: usize) {
522 dispatch!(self, set_row_stride_unchecked, stride)
523 }
524
525 pub fn plane_offset(&self) -> Option<usize> {
527 dispatch!(self, plane_offset)
528 }
529
530 pub fn view_origin(&self) -> Option<crate::ViewOrigin> {
534 dispatch!(self, view_origin)
535 }
536
537 pub fn set_plane_offset(&mut self, offset: usize) {
539 dispatch!(self, set_plane_offset, offset)
540 }
541
542 pub fn batch(&self, n: usize) -> crate::Result<TensorDyn> {
545 dyn_fanout!(self, batch, n)
546 }
547
548 pub fn view(&self, region: crate::Region) -> crate::Result<TensorDyn> {
551 dyn_fanout!(self, view, region)
552 }
553
554 pub fn cuda(&self) -> Option<&crate::cuda::CudaHandle> {
559 dispatch!(self, cuda)
560 }
561
562 pub fn cuda_map(&self) -> Option<crate::cuda::CudaMap<'_>> {
585 dispatch!(self, cuda_map)
586 }
587
588 pub fn quantization(&self) -> Option<&crate::Quantization> {
592 match self {
593 Self::U8(t) => t.quantization(),
594 Self::I8(t) => t.quantization(),
595 Self::U16(t) => t.quantization(),
596 Self::I16(t) => t.quantization(),
597 Self::U32(t) => t.quantization(),
598 Self::I32(t) => t.quantization(),
599 Self::U64(t) => t.quantization(),
600 Self::I64(t) => t.quantization(),
601 Self::F16(_) | Self::F32(_) | Self::F64(_) => None,
602 }
603 }
604
605 pub fn set_quantization(&mut self, q: crate::Quantization) -> crate::Result<()> {
609 match self {
610 Self::U8(t) => t.set_quantization(q),
611 Self::I8(t) => t.set_quantization(q),
612 Self::U16(t) => t.set_quantization(q),
613 Self::I16(t) => t.set_quantization(q),
614 Self::U32(t) => t.set_quantization(q),
615 Self::I32(t) => t.set_quantization(q),
616 Self::U64(t) => t.set_quantization(q),
617 Self::I64(t) => t.set_quantization(q),
618 Self::F16(_) | Self::F32(_) | Self::F64(_) => Err(crate::Error::QuantizationInvalid {
619 field: "dtype_is_integer",
620 expected: "integer tensor dtype (u8/i8/u16/i16/u32/i32/u64/i64)".to_string(),
621 got: format!("{:?}", self.dtype()),
622 }),
623 }
624 }
625
626 pub fn with_quantization(mut self, q: crate::Quantization) -> crate::Result<Self> {
629 self.set_quantization(q)?;
630 Ok(self)
631 }
632
633 pub fn clear_quantization(&mut self) {
635 match self {
636 Self::U8(t) => t.clear_quantization(),
637 Self::I8(t) => t.clear_quantization(),
638 Self::U16(t) => t.clear_quantization(),
639 Self::I16(t) => t.clear_quantization(),
640 Self::U32(t) => t.clear_quantization(),
641 Self::I32(t) => t.clear_quantization(),
642 Self::U64(t) => t.clear_quantization(),
643 Self::I64(t) => t.clear_quantization(),
644 Self::F16(_) | Self::F32(_) | Self::F64(_) => {}
645 }
646 }
647
648 #[cfg(unix)]
650 pub fn clone_fd(&self) -> crate::Result<std::os::fd::OwnedFd> {
651 dispatch!(self, clone_fd)
652 }
653
654 #[cfg(target_os = "linux")]
665 pub fn dmabuf_clone(&self) -> crate::Result<std::os::fd::OwnedFd> {
666 if self.memory() != TensorMemory::DmaBuf {
667 return Err(crate::Error::NotImplemented(format!(
668 "dmabuf_clone requires DMA-backed tensor, got {:?}",
669 self.memory()
670 )));
671 }
672 self.clone_fd()
673 }
674
675 #[cfg(target_os = "linux")]
686 pub fn dmabuf(&self) -> crate::Result<std::os::fd::BorrowedFd<'_>> {
687 dispatch!(self, dmabuf)
688 }
689
690 pub fn is_multiplane(&self) -> bool {
692 dispatch!(self, is_multiplane)
693 }
694
695 pub fn buffer_identity(&self) -> &crate::BufferIdentity {
705 dispatch!(self, buffer_identity)
706 }
707
708 pub fn aliases(&self, other: &Self) -> bool {
726 if self.buffer_identity().id() == other.buffer_identity().id() {
727 return true;
728 }
729 if self.memory() != other.memory() {
730 return false;
731 }
732 #[cfg(target_os = "linux")]
733 if self.memory() == TensorMemory::DmaBuf {
734 use std::os::fd::AsRawFd;
735 if let (Ok(a), Ok(b)) = (self.dmabuf(), other.dmabuf()) {
736 return a.as_raw_fd() == b.as_raw_fd();
737 }
738 }
739 false
740 }
741
742 downcast_methods!(U8, u8, as_u8, as_u8_mut, into_u8);
745 downcast_methods!(I8, i8, as_i8, as_i8_mut, into_i8);
746 downcast_methods!(U16, u16, as_u16, as_u16_mut, into_u16);
747 downcast_methods!(I16, i16, as_i16, as_i16_mut, into_i16);
748 downcast_methods!(U32, u32, as_u32, as_u32_mut, into_u32);
749 downcast_methods!(I32, i32, as_i32, as_i32_mut, into_i32);
750 downcast_methods!(U64, u64, as_u64, as_u64_mut, into_u64);
751 downcast_methods!(I64, i64, as_i64, as_i64_mut, into_i64);
752 downcast_methods!(F16, f16, as_f16, as_f16_mut, into_f16);
753 downcast_methods!(F32, f32, as_f32, as_f32_mut, into_f32);
754 downcast_methods!(F64, f64, as_f64, as_f64_mut, into_f64);
755
756 pub fn new(
758 shape: &[usize],
759 dtype: DType,
760 memory: Option<TensorMemory>,
761 name: Option<&str>,
762 ) -> crate::Result<Self> {
763 match dtype {
764 DType::U8 => Tensor::<u8>::new(shape, memory, name).map(Self::U8),
765 DType::I8 => Tensor::<i8>::new(shape, memory, name).map(Self::I8),
766 DType::U16 => Tensor::<u16>::new(shape, memory, name).map(Self::U16),
767 DType::I16 => Tensor::<i16>::new(shape, memory, name).map(Self::I16),
768 DType::U32 => Tensor::<u32>::new(shape, memory, name).map(Self::U32),
769 DType::I32 => Tensor::<i32>::new(shape, memory, name).map(Self::I32),
770 DType::U64 => Tensor::<u64>::new(shape, memory, name).map(Self::U64),
771 DType::I64 => Tensor::<i64>::new(shape, memory, name).map(Self::I64),
772 DType::F16 => Tensor::<f16>::new(shape, memory, name).map(Self::F16),
773 DType::F32 => Tensor::<f32>::new(shape, memory, name).map(Self::F32),
774 DType::F64 => Tensor::<f64>::new(shape, memory, name).map(Self::F64),
775 }
776 }
777
778 #[cfg(unix)]
805 pub fn from_fd(
806 fd: std::os::fd::OwnedFd,
807 shape: &[usize],
808 dtype: DType,
809 name: Option<&str>,
810 ) -> crate::Result<Self> {
811 match dtype {
812 DType::U8 => Tensor::<u8>::from_fd(fd, shape, name).map(Self::U8),
813 DType::I8 => Tensor::<i8>::from_fd(fd, shape, name).map(Self::I8),
814 DType::U16 => Tensor::<u16>::from_fd(fd, shape, name).map(Self::U16),
815 DType::I16 => Tensor::<i16>::from_fd(fd, shape, name).map(Self::I16),
816 DType::U32 => Tensor::<u32>::from_fd(fd, shape, name).map(Self::U32),
817 DType::I32 => Tensor::<i32>::from_fd(fd, shape, name).map(Self::I32),
818 DType::U64 => Tensor::<u64>::from_fd(fd, shape, name).map(Self::U64),
819 DType::I64 => Tensor::<i64>::from_fd(fd, shape, name).map(Self::I64),
820 DType::F16 => Tensor::<f16>::from_fd(fd, shape, name).map(Self::F16),
821 DType::F32 => Tensor::<f32>::from_fd(fd, shape, name).map(Self::F32),
822 DType::F64 => Tensor::<f64>::from_fd(fd, shape, name).map(Self::F64),
823 }
824 }
825
826 pub(crate) fn from_pbo_import(
832 buffer_id: u32,
833 size: usize,
834 shape: &[usize],
835 dtype: DType,
836 ops: std::sync::Arc<dyn crate::PboOps>,
837 ) -> crate::Result<Self> {
838 match dtype {
839 DType::U8 => crate::PboTensor::<u8>::from_pbo(buffer_id, size, shape, None, ops)
840 .and_then(|p| Tensor::from_pbo(p).map(Self::U8)),
841 DType::I8 => crate::PboTensor::<i8>::from_pbo(buffer_id, size, shape, None, ops)
842 .and_then(|p| Tensor::from_pbo(p).map(Self::I8)),
843 DType::U16 => crate::PboTensor::<u16>::from_pbo(buffer_id, size, shape, None, ops)
844 .and_then(|p| Tensor::from_pbo(p).map(Self::U16)),
845 DType::I16 => crate::PboTensor::<i16>::from_pbo(buffer_id, size, shape, None, ops)
846 .and_then(|p| Tensor::from_pbo(p).map(Self::I16)),
847 DType::U32 => crate::PboTensor::<u32>::from_pbo(buffer_id, size, shape, None, ops)
848 .and_then(|p| Tensor::from_pbo(p).map(Self::U32)),
849 DType::I32 => crate::PboTensor::<i32>::from_pbo(buffer_id, size, shape, None, ops)
850 .and_then(|p| Tensor::from_pbo(p).map(Self::I32)),
851 DType::U64 => crate::PboTensor::<u64>::from_pbo(buffer_id, size, shape, None, ops)
852 .and_then(|p| Tensor::from_pbo(p).map(Self::U64)),
853 DType::I64 => crate::PboTensor::<i64>::from_pbo(buffer_id, size, shape, None, ops)
854 .and_then(|p| Tensor::from_pbo(p).map(Self::I64)),
855 DType::F16 => crate::PboTensor::<f16>::from_pbo(buffer_id, size, shape, None, ops)
856 .and_then(|p| Tensor::from_pbo(p).map(Self::F16)),
857 DType::F32 => crate::PboTensor::<f32>::from_pbo(buffer_id, size, shape, None, ops)
858 .and_then(|p| Tensor::from_pbo(p).map(Self::F32)),
859 DType::F64 => crate::PboTensor::<f64>::from_pbo(buffer_id, size, shape, None, ops)
860 .and_then(|p| Tensor::from_pbo(p).map(Self::F64)),
861 }
862 }
863
864 pub unsafe fn from_raw_host(
882 ptr: *mut u8,
883 shape: &[usize],
884 dtype: DType,
885 name: Option<&str>,
886 ) -> crate::Result<Self> {
887 unsafe { Self::from_foreign_ptr(ptr, shape, dtype, None, name) }
888 }
889
890 pub unsafe fn from_raw_host_with_capacity(
906 ptr: *mut u8,
907 shape: &[usize],
908 capacity_bytes: usize,
909 dtype: DType,
910 name: Option<&str>,
911 ) -> crate::Result<Self> {
912 unsafe {
913 Self::from_foreign_ptr_with_capacity(ptr, shape, dtype, capacity_bytes, None, name)
914 }
915 }
916
917 pub(crate) fn import_storage(
929 desc: &crate::TensorDesc,
930 shape: &[usize],
931 dtype: DType,
932 ) -> crate::Result<Self> {
933 match desc.kind {
934 #[cfg(target_os = "linux")]
935 crate::protocol::kind::DMABUF => {
936 let owned = crate::protocol::dup_descriptor_fd(desc)?;
941 Self::from_fd(owned, shape, dtype, None)
942 }
943 #[cfg(not(target_os = "linux"))]
944 crate::protocol::kind::DMABUF => Err(crate::Error::NotImplemented(
945 "dma-buf import off Linux".into(),
946 )),
947 #[cfg(any(target_os = "macos", target_os = "ios"))]
948 crate::protocol::kind::IOSURFACE => {
949 let id = crate::protocol::descriptor_surface_id(desc)?;
950 let surface_ref = crate::iosurface::lookup_by_id(id).ok_or_else(|| {
951 crate::Error::InvalidArgument(format!(
952 "IOSurface id {id} is not live (the producer's capsule may \
953 already have been released)"
954 ))
955 })?;
956 let result = unsafe { Self::from_iosurface(surface_ref, shape, dtype, None) };
961 crate::iosurface::release(surface_ref);
962 result
963 }
964 #[cfg(not(any(target_os = "macos", target_os = "ios")))]
965 crate::protocol::kind::IOSURFACE => Err(crate::Error::NotImplemented(
966 "IOSurface import off Apple platforms".into(),
967 )),
968 crate::protocol::kind::HOST => {
969 crate::protocol::check_descriptor_host_ptr(desc)?;
970 unsafe {
977 Self::from_raw_host_with_capacity(
978 desc.ptr.0,
979 shape,
980 desc.capacity as usize,
981 dtype,
982 None,
983 )
984 }
985 }
986 crate::protocol::kind::PBO => {
987 let buffer_id = crate::protocol::descriptor_pbo_buffer_id(desc)?;
988 let ops = unsafe { crate::pbo::import_pbo_ops(desc.ptr.0 as *const _)? };
998 Self::from_pbo_import(buffer_id, desc.capacity as usize, shape, dtype, ops)
999 }
1000 k => Err(crate::Error::NotImplemented(format!(
1001 "tensor interop kind {k} cannot be imported by this build"
1002 ))),
1003 }
1004 }
1005
1006 pub unsafe fn from_foreign_ptr(
1019 ptr: *mut u8,
1020 shape: &[usize],
1021 dtype: DType,
1022 owner: Option<crate::ForeignOwner>,
1023 name: Option<&str>,
1024 ) -> crate::Result<Self> {
1025 unsafe {
1030 match dtype {
1031 DType::U8 => {
1032 Tensor::<u8>::from_foreign(ptr.cast(), shape, owner, name).map(Self::U8)
1033 }
1034 DType::I8 => {
1035 Tensor::<i8>::from_foreign(ptr.cast(), shape, owner, name).map(Self::I8)
1036 }
1037 DType::U16 => {
1038 Tensor::<u16>::from_foreign(ptr.cast(), shape, owner, name).map(Self::U16)
1039 }
1040 DType::I16 => {
1041 Tensor::<i16>::from_foreign(ptr.cast(), shape, owner, name).map(Self::I16)
1042 }
1043 DType::U32 => {
1044 Tensor::<u32>::from_foreign(ptr.cast(), shape, owner, name).map(Self::U32)
1045 }
1046 DType::I32 => {
1047 Tensor::<i32>::from_foreign(ptr.cast(), shape, owner, name).map(Self::I32)
1048 }
1049 DType::U64 => {
1050 Tensor::<u64>::from_foreign(ptr.cast(), shape, owner, name).map(Self::U64)
1051 }
1052 DType::I64 => {
1053 Tensor::<i64>::from_foreign(ptr.cast(), shape, owner, name).map(Self::I64)
1054 }
1055 DType::F16 => {
1056 Tensor::<f16>::from_foreign(ptr.cast(), shape, owner, name).map(Self::F16)
1057 }
1058 DType::F32 => {
1059 Tensor::<f32>::from_foreign(ptr.cast(), shape, owner, name).map(Self::F32)
1060 }
1061 DType::F64 => {
1062 Tensor::<f64>::from_foreign(ptr.cast(), shape, owner, name).map(Self::F64)
1063 }
1064 }
1065 }
1066 }
1067
1068 pub unsafe fn from_foreign_ptr_with_capacity(
1079 ptr: *mut u8,
1080 shape: &[usize],
1081 dtype: DType,
1082 capacity_bytes: usize,
1083 owner: Option<crate::ForeignOwner>,
1084 name: Option<&str>,
1085 ) -> crate::Result<Self> {
1086 macro_rules! import {
1087 ($t:ty, $variant:ident) => {
1088 unsafe {
1093 Tensor::<$t>::from_foreign_with_capacity(
1094 ptr.cast(),
1095 shape,
1096 capacity_bytes,
1097 owner,
1098 name,
1099 )
1100 }
1101 .map(Self::$variant)
1102 };
1103 }
1104 match dtype {
1105 DType::U8 => import!(u8, U8),
1106 DType::I8 => import!(i8, I8),
1107 DType::U16 => import!(u16, U16),
1108 DType::I16 => import!(i16, I16),
1109 DType::U32 => import!(u32, U32),
1110 DType::I32 => import!(i32, I32),
1111 DType::U64 => import!(u64, U64),
1112 DType::I64 => import!(i64, I64),
1113 DType::F16 => import!(f16, F16),
1114 DType::F32 => import!(f32, F32),
1115 DType::F64 => import!(f64, F64),
1116 }
1117 }
1118
1119 #[cfg(any(target_os = "macos", target_os = "ios"))]
1130 pub fn from_iosurface_id(
1131 id: u32,
1132 shape: &[usize],
1133 dtype: DType,
1134 name: Option<&str>,
1135 ) -> crate::Result<Self> {
1136 let surface = crate::iosurface::lookup_by_id(id).ok_or_else(|| {
1137 crate::Error::InvalidArgument(format!("no live IOSurface with id {id}"))
1138 })?;
1139 let result = unsafe { Self::from_iosurface(surface, shape, dtype, name) };
1154 crate::iosurface::release(surface);
1155 result
1156 }
1157
1158 #[cfg(any(target_os = "macos", target_os = "ios"))]
1166 pub unsafe fn from_iosurface(
1167 surface_ref: *mut std::ffi::c_void,
1168 shape: &[usize],
1169 dtype: DType,
1170 name: Option<&str>,
1171 ) -> crate::Result<Self> {
1172 unsafe {
1173 match dtype {
1174 DType::U8 => Tensor::<u8>::from_iosurface(surface_ref, shape, name).map(Self::U8),
1175 DType::I8 => Tensor::<i8>::from_iosurface(surface_ref, shape, name).map(Self::I8),
1176 DType::U16 => {
1177 Tensor::<u16>::from_iosurface(surface_ref, shape, name).map(Self::U16)
1178 }
1179 DType::I16 => {
1180 Tensor::<i16>::from_iosurface(surface_ref, shape, name).map(Self::I16)
1181 }
1182 DType::U32 => {
1183 Tensor::<u32>::from_iosurface(surface_ref, shape, name).map(Self::U32)
1184 }
1185 DType::I32 => {
1186 Tensor::<i32>::from_iosurface(surface_ref, shape, name).map(Self::I32)
1187 }
1188 DType::U64 => {
1189 Tensor::<u64>::from_iosurface(surface_ref, shape, name).map(Self::U64)
1190 }
1191 DType::I64 => {
1192 Tensor::<i64>::from_iosurface(surface_ref, shape, name).map(Self::I64)
1193 }
1194 DType::F16 => {
1195 Tensor::<f16>::from_iosurface(surface_ref, shape, name).map(Self::F16)
1196 }
1197 DType::F32 => {
1198 Tensor::<f32>::from_iosurface(surface_ref, shape, name).map(Self::F32)
1199 }
1200 DType::F64 => {
1201 Tensor::<f64>::from_iosurface(surface_ref, shape, name).map(Self::F64)
1202 }
1203 }
1204 }
1205 }
1206
1207 #[cfg(any(target_os = "macos", target_os = "ios"))]
1210 pub fn iosurface_id(&self) -> Option<u32> {
1211 dispatch!(self, iosurface_id)
1212 }
1213
1214 pub fn pbo_id(&self) -> Option<u32> {
1217 dispatch!(self, pbo_id)
1218 }
1219
1220 pub fn pbo_vtable_ptr(&self) -> Option<*const std::ffi::c_void> {
1224 dispatch!(self, pbo_vtable_ptr)
1225 }
1226
1227 pub fn pbo_keepalive(&self) -> Option<std::sync::Arc<dyn Send + Sync>> {
1232 dispatch!(self, pbo_keepalive)
1233 }
1234
1235 #[cfg(any(target_os = "macos", target_os = "ios"))]
1239 pub fn iosurface_ref(&self) -> Option<*mut std::ffi::c_void> {
1240 dispatch!(self, iosurface_ref)
1241 }
1242
1243 #[cfg(any(target_os = "macos", target_os = "ios"))]
1248 pub fn iosurface_physical_dims(&self) -> Option<(usize, usize)> {
1249 dispatch!(self, iosurface_physical_dims)
1250 }
1251
1252 #[cfg(target_os = "android")]
1262 pub unsafe fn from_hardware_buffer(
1263 buffer_ptr: *mut std::ffi::c_void,
1264 shape: &[usize],
1265 dtype: DType,
1266 name: Option<&str>,
1267 ) -> crate::Result<Self> {
1268 unsafe {
1269 match dtype {
1270 DType::U8 => {
1271 Tensor::<u8>::from_hardware_buffer(buffer_ptr, shape, name).map(Self::U8)
1272 }
1273 DType::I8 => {
1274 Tensor::<i8>::from_hardware_buffer(buffer_ptr, shape, name).map(Self::I8)
1275 }
1276 DType::U16 => {
1277 Tensor::<u16>::from_hardware_buffer(buffer_ptr, shape, name).map(Self::U16)
1278 }
1279 DType::I16 => {
1280 Tensor::<i16>::from_hardware_buffer(buffer_ptr, shape, name).map(Self::I16)
1281 }
1282 DType::U32 => {
1283 Tensor::<u32>::from_hardware_buffer(buffer_ptr, shape, name).map(Self::U32)
1284 }
1285 DType::I32 => {
1286 Tensor::<i32>::from_hardware_buffer(buffer_ptr, shape, name).map(Self::I32)
1287 }
1288 DType::U64 => {
1289 Tensor::<u64>::from_hardware_buffer(buffer_ptr, shape, name).map(Self::U64)
1290 }
1291 DType::I64 => {
1292 Tensor::<i64>::from_hardware_buffer(buffer_ptr, shape, name).map(Self::I64)
1293 }
1294 DType::F16 => {
1295 Tensor::<f16>::from_hardware_buffer(buffer_ptr, shape, name).map(Self::F16)
1296 }
1297 DType::F32 => {
1298 Tensor::<f32>::from_hardware_buffer(buffer_ptr, shape, name).map(Self::F32)
1299 }
1300 DType::F64 => {
1301 Tensor::<f64>::from_hardware_buffer(buffer_ptr, shape, name).map(Self::F64)
1302 }
1303 }
1304 }
1305 }
1306
1307 #[cfg(target_os = "android")]
1311 pub fn hardware_buffer_ptr(&self) -> Option<*mut std::ffi::c_void> {
1312 dispatch!(self, hardware_buffer_ptr)
1313 }
1314
1315 #[cfg(target_os = "android")]
1319 pub fn hardware_buffer_physical_dims(&self) -> Option<(usize, usize)> {
1320 dispatch!(self, hardware_buffer_physical_dims)
1321 }
1322
1323 pub fn copy_to_flat(&self, dst: &mut [u8]) -> crate::Result<()> {
1328 dispatch!(self, copy_to_flat, dst)
1329 }
1330
1331 pub fn image(
1349 width: usize,
1350 height: usize,
1351 format: PixelFormat,
1352 dtype: DType,
1353 memory: Option<TensorMemory>,
1354 access: crate::CpuAccess,
1355 ) -> crate::Result<Self> {
1356 match dtype {
1357 DType::U8 => Tensor::<u8>::image(width, height, format, memory, access).map(Self::U8),
1358 DType::I8 => Tensor::<i8>::image(width, height, format, memory, access).map(Self::I8),
1359 DType::U16 => {
1360 Tensor::<u16>::image(width, height, format, memory, access).map(Self::U16)
1361 }
1362 DType::I16 => {
1363 Tensor::<i16>::image(width, height, format, memory, access).map(Self::I16)
1364 }
1365 DType::U32 => {
1366 Tensor::<u32>::image(width, height, format, memory, access).map(Self::U32)
1367 }
1368 DType::I32 => {
1369 Tensor::<i32>::image(width, height, format, memory, access).map(Self::I32)
1370 }
1371 DType::U64 => {
1372 Tensor::<u64>::image(width, height, format, memory, access).map(Self::U64)
1373 }
1374 DType::I64 => {
1375 Tensor::<i64>::image(width, height, format, memory, access).map(Self::I64)
1376 }
1377 DType::F16 => {
1378 Tensor::<f16>::image(width, height, format, memory, access).map(Self::F16)
1379 }
1380 DType::F32 => {
1381 Tensor::<f32>::image(width, height, format, memory, access).map(Self::F32)
1382 }
1383 DType::F64 => {
1384 Tensor::<f64>::image(width, height, format, memory, access).map(Self::F64)
1385 }
1386 }
1387 }
1388
1389 pub fn image_desc(desc: &crate::ImageDesc) -> crate::Result<Self> {
1393 match desc.dtype() {
1394 DType::U8 => Tensor::<u8>::image_desc(desc).map(Self::U8),
1395 DType::I8 => Tensor::<i8>::image_desc(desc).map(Self::I8),
1396 DType::U16 => Tensor::<u16>::image_desc(desc).map(Self::U16),
1397 DType::I16 => Tensor::<i16>::image_desc(desc).map(Self::I16),
1398 DType::U32 => Tensor::<u32>::image_desc(desc).map(Self::U32),
1399 DType::I32 => Tensor::<i32>::image_desc(desc).map(Self::I32),
1400 DType::U64 => Tensor::<u64>::image_desc(desc).map(Self::U64),
1401 DType::I64 => Tensor::<i64>::image_desc(desc).map(Self::I64),
1402 DType::F16 => Tensor::<f16>::image_desc(desc).map(Self::F16),
1403 DType::F32 => Tensor::<f32>::image_desc(desc).map(Self::F32),
1404 DType::F64 => Tensor::<f64>::image_desc(desc).map(Self::F64),
1405 }
1406 }
1407
1408 pub fn compression(&self) -> Option<crate::CompressionScheme> {
1411 dispatch!(self, compression)
1412 }
1413
1414 pub fn image_with_stride(
1441 width: usize,
1442 height: usize,
1443 format: PixelFormat,
1444 dtype: DType,
1445 row_stride_bytes: usize,
1446 memory: Option<TensorMemory>,
1447 access: crate::CpuAccess,
1448 ) -> crate::Result<Self> {
1449 match dtype {
1450 DType::U8 => Tensor::<u8>::image_with_stride(
1451 width,
1452 height,
1453 format,
1454 row_stride_bytes,
1455 memory,
1456 access,
1457 )
1458 .map(Self::U8),
1459 DType::I8 => Tensor::<i8>::image_with_stride(
1460 width,
1461 height,
1462 format,
1463 row_stride_bytes,
1464 memory,
1465 access,
1466 )
1467 .map(Self::I8),
1468 DType::U16 => Tensor::<u16>::image_with_stride(
1469 width,
1470 height,
1471 format,
1472 row_stride_bytes,
1473 memory,
1474 access,
1475 )
1476 .map(Self::U16),
1477 DType::I16 => Tensor::<i16>::image_with_stride(
1478 width,
1479 height,
1480 format,
1481 row_stride_bytes,
1482 memory,
1483 access,
1484 )
1485 .map(Self::I16),
1486 DType::U32 => Tensor::<u32>::image_with_stride(
1487 width,
1488 height,
1489 format,
1490 row_stride_bytes,
1491 memory,
1492 access,
1493 )
1494 .map(Self::U32),
1495 DType::I32 => Tensor::<i32>::image_with_stride(
1496 width,
1497 height,
1498 format,
1499 row_stride_bytes,
1500 memory,
1501 access,
1502 )
1503 .map(Self::I32),
1504 DType::U64 => Tensor::<u64>::image_with_stride(
1505 width,
1506 height,
1507 format,
1508 row_stride_bytes,
1509 memory,
1510 access,
1511 )
1512 .map(Self::U64),
1513 DType::I64 => Tensor::<i64>::image_with_stride(
1514 width,
1515 height,
1516 format,
1517 row_stride_bytes,
1518 memory,
1519 access,
1520 )
1521 .map(Self::I64),
1522 DType::F16 => Tensor::<f16>::image_with_stride(
1523 width,
1524 height,
1525 format,
1526 row_stride_bytes,
1527 memory,
1528 access,
1529 )
1530 .map(Self::F16),
1531 DType::F32 => Tensor::<f32>::image_with_stride(
1532 width,
1533 height,
1534 format,
1535 row_stride_bytes,
1536 memory,
1537 access,
1538 )
1539 .map(Self::F32),
1540 DType::F64 => Tensor::<f64>::image_with_stride(
1541 width,
1542 height,
1543 format,
1544 row_stride_bytes,
1545 memory,
1546 access,
1547 )
1548 .map(Self::F64),
1549 }
1550 }
1551}
1552
1553impl From<Tensor<u8>> for TensorDyn {
1556 fn from(t: Tensor<u8>) -> Self {
1557 Self::U8(t)
1558 }
1559}
1560
1561impl From<Tensor<i8>> for TensorDyn {
1562 fn from(t: Tensor<i8>) -> Self {
1563 Self::I8(t)
1564 }
1565}
1566
1567impl From<Tensor<u16>> for TensorDyn {
1568 fn from(t: Tensor<u16>) -> Self {
1569 Self::U16(t)
1570 }
1571}
1572
1573impl From<Tensor<i16>> for TensorDyn {
1574 fn from(t: Tensor<i16>) -> Self {
1575 Self::I16(t)
1576 }
1577}
1578
1579impl From<Tensor<u32>> for TensorDyn {
1580 fn from(t: Tensor<u32>) -> Self {
1581 Self::U32(t)
1582 }
1583}
1584
1585impl From<Tensor<i32>> for TensorDyn {
1586 fn from(t: Tensor<i32>) -> Self {
1587 Self::I32(t)
1588 }
1589}
1590
1591impl From<Tensor<u64>> for TensorDyn {
1592 fn from(t: Tensor<u64>) -> Self {
1593 Self::U64(t)
1594 }
1595}
1596
1597impl From<Tensor<i64>> for TensorDyn {
1598 fn from(t: Tensor<i64>) -> Self {
1599 Self::I64(t)
1600 }
1601}
1602
1603impl From<Tensor<f16>> for TensorDyn {
1604 fn from(t: Tensor<f16>) -> Self {
1605 Self::F16(t)
1606 }
1607}
1608
1609impl From<Tensor<f32>> for TensorDyn {
1610 fn from(t: Tensor<f32>) -> Self {
1611 Self::F32(t)
1612 }
1613}
1614
1615impl From<Tensor<f64>> for TensorDyn {
1616 fn from(t: Tensor<f64>) -> Self {
1617 Self::F64(t)
1618 }
1619}
1620
1621impl fmt::Debug for TensorDyn {
1622 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
1623 dispatch!(self, fmt, f)
1624 }
1625}
1626
1627pub type Raw = *mut std::ffi::c_void;
1630
1631impl TensorDyn {
1632 pub fn into_raw(self) -> Raw {
1635 Box::into_raw(Box::new(self)) as Raw
1636 }
1637
1638 pub unsafe fn from_raw(p: Raw) -> TensorDyn {
1644 *unsafe { Box::from_raw(p as *mut TensorDyn) }
1645 }
1646
1647 pub unsafe fn with_raw<R>(p: Raw, f: impl FnOnce(&mut TensorDyn) -> R) -> R {
1653 let mut t = std::mem::ManuallyDrop::new(unsafe { std::ptr::read(p as *mut TensorDyn) });
1654 f(&mut t)
1655 }
1656}
1657
1658#[cfg(test)]
1659mod tests {
1660 use super::*;
1661
1662 #[test]
1663 fn from_typed_tensor() {
1664 let t = Tensor::<u8>::new(&[10], None, None).unwrap();
1665 let dyn_t: TensorDyn = t.into();
1666 assert_eq!(dyn_t.dtype(), DType::U8);
1667 assert_eq!(dyn_t.shape(), &[10]);
1668 }
1669
1670 #[test]
1671 fn from_foreign_ptr_wraps_borrowed_memory() {
1672 use crate::TensorMapTrait;
1673 let mut vec: Vec<f32> = vec![0.0; 4];
1676 let ptr = vec.as_mut_ptr() as *mut u8;
1677 let owner: crate::ForeignOwner = Box::new(vec);
1678 let t = unsafe {
1679 TensorDyn::from_foreign_ptr(ptr, &[2, 2], DType::F32, Some(owner), Some("trt_output"))
1680 }
1681 .unwrap();
1682 assert_eq!(t.dtype(), DType::F32);
1683 assert_eq!(t.memory(), TensorMemory::Mem);
1684 assert_eq!(t.shape(), &[2, 2]);
1685 {
1686 let mut m = t.as_f32().unwrap().map().unwrap();
1687 m.as_mut_slice().copy_from_slice(&[1.0, 2.0, 3.0, 4.0]);
1688 }
1689 let m = t.as_f32().unwrap().map().unwrap();
1690 assert_eq!(m.as_slice(), &[1.0, 2.0, 3.0, 4.0]);
1691 }
1692
1693 #[test]
1702 fn from_foreign_ptr_rejects_null_ptr() {
1703 let err = unsafe {
1704 TensorDyn::from_foreign_ptr(std::ptr::null_mut(), &[4], DType::U8, None, None)
1705 }
1706 .unwrap_err();
1707 assert!(
1709 matches!(err, crate::error::Error::InvalidArgument(ref m) if m.contains("non-null")),
1710 "expected InvalidArgument(non-null), got {err:?}"
1711 );
1712 }
1713
1714 #[test]
1715 fn from_foreign_ptr_rejects_empty_shape() {
1716 let mut dummy: u8 = 0;
1717 let err = unsafe {
1718 TensorDyn::from_foreign_ptr(&mut dummy as *mut u8, &[], DType::U8, None, None)
1719 }
1720 .unwrap_err();
1721 assert!(
1722 matches!(err, crate::error::Error::InvalidSize(0)),
1723 "expected InvalidSize(0) for empty shape, got {err:?}"
1724 );
1725 }
1726
1727 #[test]
1728 fn from_foreign_ptr_rejects_overflow_shape() {
1729 let mut dummy: u8 = 0;
1730 let huge = [usize::MAX / 2 + 1, 2];
1731 let err = unsafe { TensorDyn::from_foreign_ptr(&mut dummy, &huge, DType::U8, None, None) }
1732 .unwrap_err();
1733 assert!(
1734 matches!(err, crate::error::Error::InvalidArgument(ref m) if m.contains("overflow")),
1735 "expected InvalidArgument(overflow), got {err:?}"
1736 );
1737 }
1738
1739 #[test]
1740 fn from_foreign_ptr_u8_dtype_dispatch() {
1741 let mut buf: Vec<u8> = vec![1, 2, 3, 4];
1744 let ptr = buf.as_mut_ptr();
1745 let owner: crate::ForeignOwner = Box::new(buf);
1746 let t = unsafe {
1747 TensorDyn::from_foreign_ptr(ptr, &[4], DType::U8, Some(owner), Some("u8_foreign"))
1748 }
1749 .unwrap();
1750 assert_eq!(t.dtype(), DType::U8);
1751 assert_eq!(t.shape(), &[4]);
1752 let m = t.as_u8().unwrap().map().unwrap();
1753 use crate::TensorMapTrait;
1754 assert_eq!(m.as_slice(), &[1u8, 2, 3, 4]);
1755 }
1756
1757 #[test]
1758 fn map_bytes_roundtrips_and_is_typed_agnostic() {
1759 use crate::TensorMapTrait;
1760 let t = Tensor::<f32>::new(&[4], None, None).unwrap();
1763 let dyn_t: TensorDyn = t.into();
1764 {
1765 let mut view = dyn_t.map_bytes(crate::CpuAccess::ReadWrite).unwrap();
1766 assert_eq!(view.len(), 16, "4 x f32 is 16 bytes, not 4");
1767 view.as_mut_slice()
1768 .copy_from_slice(&[1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16]);
1769 }
1770 let view = dyn_t.map_bytes(crate::CpuAccess::Read).unwrap();
1771 assert_eq!(
1772 view.as_slice(),
1773 &[1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16]
1774 );
1775 }
1776
1777 #[test]
1778 fn map_bytes_none_is_invalid() {
1779 let t = Tensor::<u8>::new(&[4], None, None).unwrap();
1780 let dyn_t: TensorDyn = t.into();
1781 assert!(matches!(
1782 dyn_t.map_bytes(crate::CpuAccess::None),
1783 Err(crate::error::Error::InvalidArgument(_))
1784 ));
1785 }
1786
1787 #[test]
1788 fn downcast_ref() {
1789 let t = Tensor::<u8>::new(&[10], None, None).unwrap();
1790 let dyn_t: TensorDyn = t.into();
1791 assert!(dyn_t.as_u8().is_some());
1792 assert!(dyn_t.as_i8().is_none());
1793 }
1794
1795 #[test]
1796 fn downcast_into() {
1797 let t = Tensor::<u8>::new(&[10], None, None).unwrap();
1798 let dyn_t: TensorDyn = t.into();
1799 let back = dyn_t.into_u8().unwrap();
1800 assert_eq!(back.shape(), &[10]);
1801 }
1802
1803 #[test]
1804 fn image_accessors() {
1805 let t = Tensor::<u8>::image(
1806 640,
1807 480,
1808 PixelFormat::Rgba,
1809 None,
1810 crate::CpuAccess::ReadWrite,
1811 )
1812 .unwrap();
1813 let dyn_t: TensorDyn = t.into();
1814 assert_eq!(dyn_t.format(), Some(PixelFormat::Rgba));
1815 assert_eq!(dyn_t.width(), Some(640));
1816 assert_eq!(dyn_t.height(), Some(480));
1817 assert!(!dyn_t.is_multiplane());
1818 }
1819
1820 #[test]
1821 fn cpu_access_reports_the_declared_access() {
1822 let read_t =
1824 Tensor::<u8>::image(4, 4, PixelFormat::Rgba, None, crate::CpuAccess::Read).unwrap();
1825 let dyn_read: TensorDyn = read_t.into();
1826 assert_eq!(dyn_read.cpu_access(), crate::CpuAccess::Read);
1827 assert!(!dyn_read.cpu_access().writes());
1828
1829 let mem_t = Tensor::<u8>::new(&[4], None, None).unwrap();
1833 let dyn_mem: TensorDyn = mem_t.into();
1834 assert_eq!(dyn_mem.cpu_access(), crate::CpuAccess::ReadWrite);
1835 assert!(dyn_mem.cpu_access().writes());
1836 }
1837
1838 #[test]
1839 fn image_constructor() {
1840 let dyn_t = TensorDyn::image(
1841 640,
1842 480,
1843 PixelFormat::Rgb,
1844 DType::U8,
1845 None,
1846 crate::CpuAccess::ReadWrite,
1847 )
1848 .unwrap();
1849 assert_eq!(dyn_t.dtype(), DType::U8);
1850 assert_eq!(dyn_t.format(), Some(PixelFormat::Rgb));
1851 assert_eq!(dyn_t.width(), Some(640));
1852 }
1853
1854 #[test]
1855 fn image_constructor_i8() {
1856 let dyn_t = TensorDyn::image(
1857 640,
1858 480,
1859 PixelFormat::Rgb,
1860 DType::I8,
1861 None,
1862 crate::CpuAccess::ReadWrite,
1863 )
1864 .unwrap();
1865 assert_eq!(dyn_t.dtype(), DType::I8);
1866 assert_eq!(dyn_t.format(), Some(PixelFormat::Rgb));
1867 }
1868
1869 #[test]
1870 fn set_format_packed() {
1871 let mut t = TensorDyn::new(&[480, 640, 3], DType::U8, None, None).unwrap();
1872 assert_eq!(t.format(), None);
1873 t.set_format(PixelFormat::Rgb).unwrap();
1874 assert_eq!(t.format(), Some(PixelFormat::Rgb));
1875 assert_eq!(t.width(), Some(640));
1876 assert_eq!(t.height(), Some(480));
1877 }
1878
1879 #[test]
1880 fn set_format_planar() {
1881 let mut t = TensorDyn::new(&[3, 480, 640], DType::U8, None, None).unwrap();
1882 t.set_format(PixelFormat::PlanarRgb).unwrap();
1883 assert_eq!(t.format(), Some(PixelFormat::PlanarRgb));
1884 assert_eq!(t.width(), Some(640));
1885 assert_eq!(t.height(), Some(480));
1886 }
1887
1888 #[test]
1889 fn set_format_rejects_wrong_shape() {
1890 let mut t = TensorDyn::new(&[480, 640, 4], DType::U8, None, None).unwrap();
1891 assert!(t.set_format(PixelFormat::Rgb).is_err());
1892 }
1893
1894 #[test]
1895 fn with_format_builder() {
1896 let t = TensorDyn::new(&[480, 640, 4], DType::U8, None, None)
1897 .unwrap()
1898 .with_format(PixelFormat::Rgba)
1899 .unwrap();
1900 assert_eq!(t.format(), Some(PixelFormat::Rgba));
1901 assert_eq!(t.width(), Some(640));
1902 assert_eq!(t.height(), Some(480));
1903 }
1904
1905 #[cfg(target_os = "linux")]
1906 #[test]
1907 fn dmabuf_clone_mem_tensor_fails() {
1908 let t = TensorDyn::new(&[480, 640, 3], DType::U8, Some(TensorMemory::Mem), None).unwrap();
1909 assert_eq!(t.memory(), TensorMemory::Mem);
1910 assert!(t.dmabuf_clone().is_err());
1911 }
1912
1913 #[cfg(target_os = "linux")]
1914 #[test]
1915 fn dmabuf_mem_tensor_fails() {
1916 let t = TensorDyn::new(&[480, 640, 3], DType::U8, Some(TensorMemory::Mem), None).unwrap();
1917 assert!(t.dmabuf().is_err());
1918 }
1919
1920 #[test]
1921 fn set_format_semi_planar_nv12() {
1922 let mut t = TensorDyn::new(&[720, 640], DType::U8, Some(TensorMemory::Mem), None).unwrap();
1924 t.set_format(PixelFormat::Nv12).unwrap();
1925 assert_eq!(t.format(), Some(PixelFormat::Nv12));
1926 assert_eq!(t.width(), Some(640));
1927 assert_eq!(t.height(), Some(480));
1928 }
1929
1930 #[test]
1931 fn set_format_semi_planar_nv16() {
1932 let mut t = TensorDyn::new(&[960, 640], DType::U8, Some(TensorMemory::Mem), None).unwrap();
1934 t.set_format(PixelFormat::Nv16).unwrap();
1935 assert_eq!(t.format(), Some(PixelFormat::Nv16));
1936 assert_eq!(t.width(), Some(640));
1937 assert_eq!(t.height(), Some(480));
1938 }
1939
1940 #[test]
1941 fn with_format_rejects_wrong_shape() {
1942 let result = TensorDyn::new(&[480, 640, 4], DType::U8, None, None)
1943 .unwrap()
1944 .with_format(PixelFormat::Rgb);
1945 assert!(result.is_err());
1946 }
1947
1948 #[test]
1949 fn set_format_preserved_after_rejection() {
1950 let mut t = TensorDyn::new(&[480, 640, 3], DType::U8, None, None).unwrap();
1951 t.set_format(PixelFormat::Rgb).unwrap();
1952 assert_eq!(t.format(), Some(PixelFormat::Rgb));
1953
1954 assert!(t.set_format(PixelFormat::Rgba).is_err());
1956
1957 assert_eq!(t.format(), Some(PixelFormat::Rgb));
1959 }
1960
1961 #[test]
1962 fn set_format_idempotent() {
1963 let mut t = TensorDyn::new(&[480, 640, 3], DType::U8, None, None).unwrap();
1964 t.set_format(PixelFormat::Rgb).unwrap();
1965 t.set_format(PixelFormat::Rgb).unwrap();
1966 assert_eq!(t.format(), Some(PixelFormat::Rgb));
1967 assert_eq!(t.width(), Some(640));
1968 assert_eq!(t.height(), Some(480));
1969 }
1970
1971 #[test]
1974 fn set_row_stride_valid() {
1975 let mut t = TensorDyn::image(
1977 100,
1978 100,
1979 PixelFormat::Rgba,
1980 DType::U8,
1981 None,
1982 crate::CpuAccess::ReadWrite,
1983 )
1984 .unwrap();
1985 t.set_row_stride(512).unwrap();
1986 assert_eq!(t.row_stride(), Some(512));
1987 assert_eq!(t.effective_row_stride(), Some(512));
1988 }
1989
1990 #[test]
1991 fn set_row_stride_equals_min() {
1992 let mut t = TensorDyn::image(
1994 100,
1995 100,
1996 PixelFormat::Rgb,
1997 DType::U8,
1998 None,
1999 crate::CpuAccess::ReadWrite,
2000 )
2001 .unwrap();
2002 t.set_row_stride(300).unwrap();
2003 assert_eq!(t.row_stride(), Some(300));
2004 }
2005
2006 #[test]
2007 fn set_row_stride_too_small() {
2008 let mut t = TensorDyn::image(
2013 64,
2014 100,
2015 PixelFormat::Rgba,
2016 DType::U8,
2017 None,
2018 crate::CpuAccess::ReadWrite,
2019 )
2020 .unwrap();
2021 assert!(t.set_row_stride(200).is_err());
2022 assert_eq!(t.row_stride(), None);
2023 }
2024
2025 #[test]
2026 fn set_row_stride_zero() {
2027 let mut t = TensorDyn::image(
2028 100,
2029 100,
2030 PixelFormat::Rgb,
2031 DType::U8,
2032 None,
2033 crate::CpuAccess::ReadWrite,
2034 )
2035 .unwrap();
2036 assert!(t.set_row_stride(0).is_err());
2037 }
2038
2039 #[test]
2040 fn set_row_stride_requires_format() {
2041 let mut t = TensorDyn::new(&[480, 640, 3], DType::U8, None, None).unwrap();
2042 assert!(t.set_row_stride(2048).is_err());
2043 }
2044
2045 #[test]
2046 fn effective_row_stride_without_stride() {
2047 let t = TensorDyn::image(
2052 64,
2053 100,
2054 PixelFormat::Rgb,
2055 DType::U8,
2056 None,
2057 crate::CpuAccess::ReadWrite,
2058 )
2059 .unwrap();
2060 assert_eq!(t.row_stride(), None);
2061 assert_eq!(t.effective_row_stride(), Some(192)); }
2063
2064 #[test]
2065 fn effective_row_stride_padded_packed_dma() {
2066 let t = match TensorDyn::image(
2072 100,
2073 100,
2074 PixelFormat::Rgb,
2075 DType::U8,
2076 Some(TensorMemory::DmaBuf),
2077 crate::CpuAccess::ReadWrite,
2078 ) {
2079 Ok(t) if t.memory() == TensorMemory::DmaBuf => t,
2080 _ => return,
2081 };
2082 if cfg!(target_os = "android") {
2088 assert!(
2089 t.row_stride().is_none() || t.row_stride() >= Some(300),
2090 "AHardwareBuffer pitch must be absent or at least the tight row"
2091 );
2092 return;
2093 }
2094 assert_eq!(t.row_stride(), Some(320));
2095 assert_eq!(t.effective_row_stride(), Some(320));
2096 }
2097
2098 #[test]
2099 fn effective_row_stride_no_format() {
2100 let t = TensorDyn::new(&[480, 640, 3], DType::U8, None, None).unwrap();
2101 assert_eq!(t.effective_row_stride(), None);
2102 }
2103
2104 #[test]
2105 fn with_row_stride_builder() {
2106 let t = TensorDyn::image(
2107 100,
2108 100,
2109 PixelFormat::Rgba,
2110 DType::U8,
2111 None,
2112 crate::CpuAccess::ReadWrite,
2113 )
2114 .unwrap()
2115 .with_row_stride(512)
2116 .unwrap();
2117 assert_eq!(t.row_stride(), Some(512));
2118 assert_eq!(t.effective_row_stride(), Some(512));
2119 }
2120
2121 #[test]
2122 fn with_row_stride_rejects_small() {
2123 let result = TensorDyn::image(
2124 100,
2125 100,
2126 PixelFormat::Rgba,
2127 DType::U8,
2128 None,
2129 crate::CpuAccess::ReadWrite,
2130 )
2131 .unwrap()
2132 .with_row_stride(200);
2133 assert!(result.is_err());
2134 }
2135
2136 #[test]
2137 fn set_format_clears_row_stride() {
2138 let mut t = TensorDyn::new(&[480, 640, 3], DType::U8, None, None).unwrap();
2139 t.set_format(PixelFormat::Rgb).unwrap();
2140 t.set_row_stride(2048).unwrap();
2141 assert_eq!(t.row_stride(), Some(2048));
2142
2143 let _ = t.set_format(PixelFormat::Bgra);
2145 assert_eq!(t.row_stride(), Some(2048));
2146
2147 t.set_format(PixelFormat::Rgb).unwrap();
2149 assert_eq!(t.row_stride(), Some(2048));
2150
2151 t.reshape(&[480 * 640 * 3]).unwrap();
2153 assert_eq!(t.row_stride(), None);
2154 assert_eq!(t.format(), None);
2155 }
2156
2157 #[test]
2158 fn set_format_different_compatible_clears_stride() {
2159 let mut t = TensorDyn::new(&[480, 640, 4], DType::U8, None, None).unwrap();
2162 t.set_format(PixelFormat::Rgba).unwrap();
2163 t.set_row_stride(4096).unwrap();
2164 assert_eq!(t.row_stride(), Some(4096));
2165
2166 t.set_format(PixelFormat::Bgra).unwrap();
2168 assert_eq!(t.format(), Some(PixelFormat::Bgra));
2169 assert_eq!(t.row_stride(), None);
2170 }
2171
2172 #[test]
2173 fn set_format_same_preserves_stride() {
2174 let mut t = TensorDyn::image(
2175 100,
2176 100,
2177 PixelFormat::Rgb,
2178 DType::U8,
2179 None,
2180 crate::CpuAccess::ReadWrite,
2181 )
2182 .unwrap();
2183 t.set_row_stride(512).unwrap();
2184 t.set_format(PixelFormat::Rgb).unwrap();
2186 assert_eq!(t.row_stride(), Some(512));
2187 }
2188
2189 #[test]
2190 fn effective_row_stride_planar() {
2191 let t = TensorDyn::image(
2192 640,
2193 480,
2194 PixelFormat::PlanarRgb,
2195 DType::U8,
2196 None,
2197 crate::CpuAccess::ReadWrite,
2198 )
2199 .unwrap();
2200 assert_eq!(t.effective_row_stride(), Some(640)); }
2202
2203 #[test]
2204 fn effective_row_stride_nv12() {
2205 let t = TensorDyn::image(
2206 640,
2207 480,
2208 PixelFormat::Nv12,
2209 DType::U8,
2210 None,
2211 crate::CpuAccess::ReadWrite,
2212 )
2213 .unwrap();
2214 assert_eq!(t.effective_row_stride(), Some(640)); }
2216
2217 #[test]
2218 fn map_rejects_strided_tensor() {
2219 let mut t = Tensor::<u8>::image(
2220 100,
2221 100,
2222 PixelFormat::Rgba,
2223 Some(TensorMemory::Mem),
2224 crate::CpuAccess::ReadWrite,
2225 )
2226 .unwrap();
2227 assert!(t.map().is_ok());
2229 t.set_row_stride(512).unwrap();
2231 let err = t.map();
2232 assert!(err.is_err());
2233 }
2234
2235 #[test]
2238 fn plane_offset_default_none() {
2239 let t = TensorDyn::image(
2240 100,
2241 100,
2242 PixelFormat::Rgba,
2243 DType::U8,
2244 None,
2245 crate::CpuAccess::ReadWrite,
2246 )
2247 .unwrap();
2248 assert_eq!(t.plane_offset(), None);
2249 }
2250
2251 #[test]
2252 fn set_plane_offset_basic() {
2253 let mut t = TensorDyn::image(
2254 100,
2255 100,
2256 PixelFormat::Rgba,
2257 DType::U8,
2258 None,
2259 crate::CpuAccess::ReadWrite,
2260 )
2261 .unwrap();
2262 t.set_plane_offset(4096);
2263 assert_eq!(t.plane_offset(), Some(4096));
2264 }
2265
2266 #[test]
2267 fn set_plane_offset_zero() {
2268 let mut t = TensorDyn::image(
2269 100,
2270 100,
2271 PixelFormat::Rgb,
2272 DType::U8,
2273 None,
2274 crate::CpuAccess::ReadWrite,
2275 )
2276 .unwrap();
2277 t.set_plane_offset(0);
2278 assert_eq!(t.plane_offset(), Some(0));
2279 }
2280
2281 #[test]
2282 fn set_plane_offset_no_format() {
2283 let mut t = TensorDyn::new(&[480, 640, 3], DType::U8, None, None).unwrap();
2285 t.set_plane_offset(4096);
2286 assert_eq!(t.plane_offset(), Some(4096));
2287 }
2288
2289 #[test]
2290 fn set_format_clears_plane_offset() {
2291 let mut t = TensorDyn::new(&[480, 640, 3], DType::U8, None, None).unwrap();
2292 t.set_format(PixelFormat::Rgb).unwrap();
2293 t.set_plane_offset(4096);
2294 assert_eq!(t.plane_offset(), Some(4096));
2295
2296 t.set_format(PixelFormat::Rgb).unwrap();
2298 assert_eq!(t.plane_offset(), Some(4096));
2299
2300 t.reshape(&[480 * 640 * 3]).unwrap();
2302 assert_eq!(t.plane_offset(), None);
2303 assert_eq!(t.format(), None);
2304 }
2305
2306 #[test]
2307 fn map_rejects_out_of_bounds_offset() {
2308 let mut t = Tensor::<u8>::image(
2309 100,
2310 100,
2311 PixelFormat::Rgba,
2312 Some(TensorMemory::Mem),
2313 crate::CpuAccess::ReadWrite,
2314 )
2315 .unwrap();
2316 assert!(t.map().is_ok());
2318 t.set_plane_offset(4096);
2321 assert!(t.map().is_err());
2322 }
2323
2324 #[test]
2325 fn mem_subview_in_bounds_maps_at_offset() {
2326 let parent = Tensor::<u8>::image(
2329 100,
2330 100,
2331 PixelFormat::Rgba,
2332 Some(TensorMemory::Mem),
2333 crate::CpuAccess::ReadWrite,
2334 )
2335 .unwrap();
2336 let view = parent.subview(4096, &[10, 10, 4]).unwrap();
2338 assert_eq!(view.plane_offset(), Some(4096));
2339 assert!(view.map().is_ok());
2340 }
2341
2342 #[test]
2343 fn dyn_batch_dispatches_every_dtype() {
2344 use DType::*;
2348 for dt in [U8, I8, U16, I16, U32, I32, U64, I64, F16, F32, F64] {
2349 let parent = TensorDyn::new(&[2, 4], dt, Some(TensorMemory::Mem), None).unwrap();
2350 let view = parent.batch(1).unwrap();
2351 assert_eq!(view.dtype(), dt, "batch must preserve dtype {dt:?}");
2352 assert_eq!(view.shape(), &[4], "{dt:?}");
2353 }
2354 }
2355
2356 #[test]
2357 fn map_accepts_zero_offset_tensor() {
2358 let mut t = Tensor::<u8>::image(
2359 100,
2360 100,
2361 PixelFormat::Rgba,
2362 Some(TensorMemory::Mem),
2363 crate::CpuAccess::ReadWrite,
2364 )
2365 .unwrap();
2366 t.set_plane_offset(0);
2367 assert!(t.map().is_ok());
2369 }
2370
2371 #[test]
2372 fn dyn_configure_image_nv12() {
2373 let mut t = TensorDyn::image(
2374 640,
2375 480,
2376 PixelFormat::Rgb,
2377 DType::U8,
2378 None,
2379 crate::CpuAccess::ReadWrite,
2380 )
2381 .unwrap();
2382 t.configure_image(320, 240, PixelFormat::Nv12).unwrap();
2383 assert_eq!(t.format(), Some(PixelFormat::Nv12));
2384 assert_eq!((t.width(), t.height()), (Some(320), Some(240)));
2385 }
2386
2387 #[test]
2388 fn tensordyn_colorimetry_roundtrip() {
2389 use crate::{ColorEncoding, Colorimetry, DType, PixelFormat};
2390 let mut t = TensorDyn::image(
2391 1280,
2392 720,
2393 PixelFormat::Nv12,
2394 DType::U8,
2395 None,
2396 crate::CpuAccess::ReadWrite,
2397 )
2398 .unwrap();
2399 assert_eq!(t.colorimetry(), None);
2400 let c = Colorimetry::default().with_encoding(ColorEncoding::Bt709);
2401 t.set_colorimetry(Some(c));
2402 assert_eq!(t.colorimetry(), Some(c));
2403 }
2404
2405 #[test]
2406 fn from_planes_propagates_plane_offset() {
2407 let mut luma =
2408 Tensor::<u8>::new(&[480, 640], Some(TensorMemory::Mem), Some("luma")).unwrap();
2409 luma.set_plane_offset(4096);
2410 let chroma =
2411 Tensor::<u8>::new(&[240, 640], Some(TensorMemory::Mem), Some("chroma")).unwrap();
2412 let combined = Tensor::<u8>::from_planes(luma, chroma, PixelFormat::Nv12).unwrap();
2413 assert_eq!(combined.plane_offset(), Some(4096));
2414 }
2415
2416 #[test]
2417 fn cuda_passthrough_none_for_mem_tensor() {
2418 let t: TensorDyn = Tensor::<f32>::new(&[10], Some(TensorMemory::Mem), None)
2421 .unwrap()
2422 .into();
2423 assert!(t.cuda().is_none());
2424 assert!(t.cuda_map().is_none());
2425 }
2426}