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 try_map_bytes(
265 &self,
266 access: crate::CpuAccess,
267 ) -> crate::Result<crate::view::HostView<'static, u8>> {
268 use crate::TensorTrait;
269 match self {
270 Self::U8(t) => t
271 .try_map_with(access)
272 .map(crate::view::HostView::into_bytes),
273 Self::I8(t) => t
274 .try_map_with(access)
275 .map(crate::view::HostView::into_bytes),
276 Self::U16(t) => t
277 .try_map_with(access)
278 .map(crate::view::HostView::into_bytes),
279 Self::I16(t) => t
280 .try_map_with(access)
281 .map(crate::view::HostView::into_bytes),
282 Self::U32(t) => t
283 .try_map_with(access)
284 .map(crate::view::HostView::into_bytes),
285 Self::I32(t) => t
286 .try_map_with(access)
287 .map(crate::view::HostView::into_bytes),
288 Self::U64(t) => t
289 .try_map_with(access)
290 .map(crate::view::HostView::into_bytes),
291 Self::I64(t) => t
292 .try_map_with(access)
293 .map(crate::view::HostView::into_bytes),
294 Self::F16(t) => t
295 .try_map_with(access)
296 .map(crate::view::HostView::into_bytes),
297 Self::F32(t) => t
298 .try_map_with(access)
299 .map(crate::view::HostView::into_bytes),
300 Self::F64(t) => t
301 .try_map_with(access)
302 .map(crate::view::HostView::into_bytes),
303 }
304 }
305
306 pub fn sync_for_cpu(&self, access: crate::CpuAccess) -> crate::Result<()> {
308 dispatch!(self, sync_for_cpu, access)
309 }
310
311 pub fn sync_for_device(&self, access: crate::CpuAccess) -> crate::Result<()> {
313 dispatch!(self, sync_for_device, access)
314 }
315
316 pub fn descriptor(&self) -> crate::TensorDesc {
317 self.descriptor_pinned(None)
318 }
319
320 pub fn descriptor_pinned(&self, pin: Option<&crate::HostPin<'_>>) -> crate::TensorDesc {
328 let handle: i64 = match self.memory() {
333 #[cfg(target_os = "linux")]
334 TensorMemory::DmaBuf => {
335 use std::os::fd::AsRawFd;
336 self.dmabuf().map(|fd| fd.as_raw_fd() as i64).unwrap_or(-1)
337 }
338 #[cfg(any(target_os = "macos", target_os = "ios"))]
339 TensorMemory::DmaBuf => self.iosurface_id().map(|id| id as i64).unwrap_or(-1),
340 #[cfg(target_os = "windows")]
345 TensorMemory::DmaBuf => self
346 .d3d11_shared_handle_value()
347 .map(|h| h as i64)
348 .unwrap_or(-1),
349 TensorMemory::Pbo => self.pbo_id().map(|id| id as i64).unwrap_or(-1),
350 _ => -1,
351 };
352 #[cfg(target_os = "windows")]
358 let (fence_handle, sync) = match self.memory() {
359 TensorMemory::DmaBuf => (
360 crate::d3d11::device()
361 .map(|d| d.exported_fence_handle_value())
362 .unwrap_or(0),
363 Some(self.gpu_write_value()).filter(|v| *v != 0),
364 ),
365 _ => (0, None),
366 };
367 #[cfg(not(target_os = "windows"))]
368 let (fence_handle, sync) = (0, None);
369 crate::protocol::from_parts(crate::protocol::DescParts {
370 dims: self.shape(),
371 memory: self.memory(),
372 dtype: self.dtype(),
373 fourcc: self.format().map(|f| f.to_fourcc()).unwrap_or(0),
374 format: self.format(),
375 row_stride: self.row_stride(),
376 handle,
377 colorimetry: self.colorimetry().map(|c| c.pack()).unwrap_or(0),
378 capacity: self.capacity_bytes() as u64,
379 pin,
380 pbo_vtable_ptr: self.pbo_vtable_ptr(),
381 fence_handle,
382 sync,
383 })
384 }
385
386 pub fn shape(&self) -> &[usize] {
387 dispatch!(self, shape)
388 }
389
390 pub fn set_logical_shape(&mut self, shape: &[usize]) -> crate::Result<()> {
394 dispatch!(self, set_logical_shape, shape)
395 }
396
397 pub fn capacity_bytes(&self) -> usize {
399 dispatch!(self, capacity_bytes)
400 }
401
402 pub fn name(&self) -> String {
404 dispatch!(self, name)
405 }
406
407 pub fn format(&self) -> Option<PixelFormat> {
409 dispatch!(self, format)
410 }
411
412 pub fn width(&self) -> Option<usize> {
414 dispatch!(self, width)
415 }
416
417 pub fn height(&self) -> Option<usize> {
419 dispatch!(self, height)
420 }
421
422 pub fn size(&self) -> usize {
424 dispatch!(self, size)
425 }
426
427 pub fn set_dtype(&mut self, dtype: DType) -> crate::Result<()> {
446 let current = self.dtype();
447 if current == dtype {
448 return Ok(());
449 }
450 if current.size() != dtype.size() {
451 return Err(crate::Error::InvalidArgument(format!(
452 "set_dtype: {current:?} and {dtype:?} are different widths \
453 ({} vs {} bytes); retagging would change the element count \
454 over an allocation whose size did not change",
455 current.size(),
456 dtype.size()
457 )));
458 }
459 unsafe {
467 let taken = std::ptr::read(self);
468 std::ptr::write(self, retag_same_width(taken, dtype));
469 }
470 Ok(())
471 }
472
473 pub fn memory(&self) -> TensorMemory {
475 dispatch!(self, memory)
476 }
477
478 pub fn cpu_access(&self) -> crate::CpuAccess {
490 dispatch!(self, cpu_access)
491 }
492
493 pub fn reshape(&mut self, shape: &[usize]) -> crate::Result<()> {
495 dispatch!(self, reshape, shape)
496 }
497
498 pub fn set_format(&mut self, format: PixelFormat) -> crate::Result<()> {
516 dispatch!(self, set_format, format)
517 }
518
519 pub fn with_format(mut self, format: PixelFormat) -> crate::Result<Self> {
536 self.set_format(format)?;
537 Ok(self)
538 }
539
540 pub fn colorimetry(&self) -> Option<crate::Colorimetry> {
542 dispatch!(self, colorimetry)
543 }
544
545 pub fn set_colorimetry(&mut self, c: Option<crate::Colorimetry>) {
547 dispatch!(self, set_colorimetry, c)
548 }
549
550 pub fn with_colorimetry(mut self, c: crate::Colorimetry) -> Self {
552 self.set_colorimetry(Some(c));
553 self
554 }
555
556 pub fn row_stride(&self) -> Option<usize> {
558 dispatch!(self, row_stride)
559 }
560
561 pub fn effective_row_stride(&self) -> Option<usize> {
563 dispatch!(self, effective_row_stride)
564 }
565
566 pub fn configure_image(
569 &mut self,
570 width: usize,
571 height: usize,
572 format: PixelFormat,
573 ) -> crate::Result<()> {
574 dispatch!(self, configure_image, width, height, format)
575 }
576
577 pub fn set_row_stride(&mut self, stride: usize) -> crate::Result<()> {
583 dispatch!(self, set_row_stride, stride)
584 }
585
586 pub fn with_row_stride(mut self, stride: usize) -> crate::Result<Self> {
588 self.set_row_stride(stride)?;
589 Ok(self)
590 }
591
592 pub fn set_row_stride_unchecked(&mut self, stride: usize) {
598 dispatch!(self, set_row_stride_unchecked, stride)
599 }
600
601 pub fn plane_offset(&self) -> Option<usize> {
603 dispatch!(self, plane_offset)
604 }
605
606 pub fn view_origin(&self) -> Option<crate::ViewOrigin> {
610 dispatch!(self, view_origin)
611 }
612
613 pub fn set_plane_offset(&mut self, offset: usize) {
615 dispatch!(self, set_plane_offset, offset)
616 }
617
618 pub fn batch(&self, n: usize) -> crate::Result<TensorDyn> {
621 dyn_fanout!(self, batch, n)
622 }
623
624 pub fn view(&self, region: crate::Region) -> crate::Result<TensorDyn> {
627 dyn_fanout!(self, view, region)
628 }
629
630 pub fn cuda(&self) -> Option<&crate::cuda::CudaHandle> {
635 dispatch!(self, cuda)
636 }
637
638 pub fn cuda_map(&self) -> Option<crate::cuda::CudaMap<'_>> {
661 dispatch!(self, cuda_map)
662 }
663
664 pub fn cuda_map_mut(&self) -> Option<crate::cuda::CudaMap<'_>> {
668 dispatch!(self, cuda_map_mut)
669 }
670
671 pub fn set_cuda_handle(&mut self, h: crate::cuda::CudaHandle) {
673 dispatch!(self, set_cuda_handle, h)
674 }
675
676 pub fn is_cuda_attached(&self) -> bool {
678 dispatch!(self, is_cuda_attached)
679 }
680
681 pub fn quantization(&self) -> Option<&crate::Quantization> {
685 match self {
686 Self::U8(t) => t.quantization(),
687 Self::I8(t) => t.quantization(),
688 Self::U16(t) => t.quantization(),
689 Self::I16(t) => t.quantization(),
690 Self::U32(t) => t.quantization(),
691 Self::I32(t) => t.quantization(),
692 Self::U64(t) => t.quantization(),
693 Self::I64(t) => t.quantization(),
694 Self::F16(_) | Self::F32(_) | Self::F64(_) => None,
695 }
696 }
697
698 pub fn set_quantization(&mut self, q: crate::Quantization) -> crate::Result<()> {
702 match self {
703 Self::U8(t) => t.set_quantization(q),
704 Self::I8(t) => t.set_quantization(q),
705 Self::U16(t) => t.set_quantization(q),
706 Self::I16(t) => t.set_quantization(q),
707 Self::U32(t) => t.set_quantization(q),
708 Self::I32(t) => t.set_quantization(q),
709 Self::U64(t) => t.set_quantization(q),
710 Self::I64(t) => t.set_quantization(q),
711 Self::F16(_) | Self::F32(_) | Self::F64(_) => Err(crate::Error::QuantizationInvalid {
712 field: "dtype_is_integer",
713 expected: "integer tensor dtype (u8/i8/u16/i16/u32/i32/u64/i64)".to_string(),
714 got: format!("{:?}", self.dtype()),
715 }),
716 }
717 }
718
719 pub fn with_quantization(mut self, q: crate::Quantization) -> crate::Result<Self> {
722 self.set_quantization(q)?;
723 Ok(self)
724 }
725
726 pub fn clear_quantization(&mut self) {
728 match self {
729 Self::U8(t) => t.clear_quantization(),
730 Self::I8(t) => t.clear_quantization(),
731 Self::U16(t) => t.clear_quantization(),
732 Self::I16(t) => t.clear_quantization(),
733 Self::U32(t) => t.clear_quantization(),
734 Self::I32(t) => t.clear_quantization(),
735 Self::U64(t) => t.clear_quantization(),
736 Self::I64(t) => t.clear_quantization(),
737 Self::F16(_) | Self::F32(_) | Self::F64(_) => {}
738 }
739 }
740
741 #[cfg(unix)]
743 pub fn clone_fd(&self) -> crate::Result<std::os::fd::OwnedFd> {
744 dispatch!(self, clone_fd)
745 }
746
747 #[cfg(target_os = "linux")]
758 pub fn dmabuf_clone(&self) -> crate::Result<std::os::fd::OwnedFd> {
759 if self.memory() != TensorMemory::DmaBuf {
760 return Err(crate::Error::NotImplemented(format!(
761 "dmabuf_clone requires DMA-backed tensor, got {:?}",
762 self.memory()
763 )));
764 }
765 self.clone_fd()
766 }
767
768 #[cfg(target_os = "linux")]
779 pub fn dmabuf(&self) -> crate::Result<std::os::fd::BorrowedFd<'_>> {
780 dispatch!(self, dmabuf)
781 }
782
783 pub fn is_multiplane(&self) -> bool {
785 dispatch!(self, is_multiplane)
786 }
787
788 pub fn chroma_dyn(&self) -> Option<TensorDyn> {
791 match self {
792 TensorDyn::U8(t) => t
793 .chroma()
794 .and_then(|c| c.clone_chroma_plane().ok())
795 .map(TensorDyn::U8),
796 TensorDyn::I8(t) => t
797 .chroma()
798 .and_then(|c| c.clone_chroma_plane().ok())
799 .map(TensorDyn::I8),
800 TensorDyn::U16(t) => t
801 .chroma()
802 .and_then(|c| c.clone_chroma_plane().ok())
803 .map(TensorDyn::U16),
804 TensorDyn::I16(t) => t
805 .chroma()
806 .and_then(|c| c.clone_chroma_plane().ok())
807 .map(TensorDyn::I16),
808 TensorDyn::U32(t) => t
809 .chroma()
810 .and_then(|c| c.clone_chroma_plane().ok())
811 .map(TensorDyn::U32),
812 TensorDyn::I32(t) => t
813 .chroma()
814 .and_then(|c| c.clone_chroma_plane().ok())
815 .map(TensorDyn::I32),
816 TensorDyn::U64(t) => t
817 .chroma()
818 .and_then(|c| c.clone_chroma_plane().ok())
819 .map(TensorDyn::U64),
820 TensorDyn::I64(t) => t
821 .chroma()
822 .and_then(|c| c.clone_chroma_plane().ok())
823 .map(TensorDyn::I64),
824 TensorDyn::F16(t) => t
825 .chroma()
826 .and_then(|c| c.clone_chroma_plane().ok())
827 .map(TensorDyn::F16),
828 TensorDyn::F32(t) => t
829 .chroma()
830 .and_then(|c| c.clone_chroma_plane().ok())
831 .map(TensorDyn::F32),
832 TensorDyn::F64(t) => t
833 .chroma()
834 .and_then(|c| c.clone_chroma_plane().ok())
835 .map(TensorDyn::F64),
836 }
837 }
838
839 pub fn buffer_identity(&self) -> &crate::BufferIdentity {
849 dispatch!(self, buffer_identity)
850 }
851
852 pub fn aliases(&self, other: &Self) -> bool {
870 if self.buffer_identity().id() == other.buffer_identity().id() {
871 return true;
872 }
873 if self.memory() != other.memory() {
874 return false;
875 }
876 #[cfg(target_os = "linux")]
877 if self.memory() == TensorMemory::DmaBuf {
878 use std::os::fd::AsRawFd;
879 if let (Ok(a), Ok(b)) = (self.dmabuf(), other.dmabuf()) {
880 return a.as_raw_fd() == b.as_raw_fd();
881 }
882 }
883 false
884 }
885
886 downcast_methods!(U8, u8, as_u8, as_u8_mut, into_u8);
889 downcast_methods!(I8, i8, as_i8, as_i8_mut, into_i8);
890 downcast_methods!(U16, u16, as_u16, as_u16_mut, into_u16);
891 downcast_methods!(I16, i16, as_i16, as_i16_mut, into_i16);
892 downcast_methods!(U32, u32, as_u32, as_u32_mut, into_u32);
893 downcast_methods!(I32, i32, as_i32, as_i32_mut, into_i32);
894 downcast_methods!(U64, u64, as_u64, as_u64_mut, into_u64);
895 downcast_methods!(I64, i64, as_i64, as_i64_mut, into_i64);
896 downcast_methods!(F16, f16, as_f16, as_f16_mut, into_f16);
897 downcast_methods!(F32, f32, as_f32, as_f32_mut, into_f32);
898 downcast_methods!(F64, f64, as_f64, as_f64_mut, into_f64);
899
900 pub fn new(
902 shape: &[usize],
903 dtype: DType,
904 memory: Option<TensorMemory>,
905 name: Option<&str>,
906 ) -> crate::Result<Self> {
907 match dtype {
908 DType::U8 => Tensor::<u8>::new(shape, memory, name).map(Self::U8),
909 DType::I8 => Tensor::<i8>::new(shape, memory, name).map(Self::I8),
910 DType::U16 => Tensor::<u16>::new(shape, memory, name).map(Self::U16),
911 DType::I16 => Tensor::<i16>::new(shape, memory, name).map(Self::I16),
912 DType::U32 => Tensor::<u32>::new(shape, memory, name).map(Self::U32),
913 DType::I32 => Tensor::<i32>::new(shape, memory, name).map(Self::I32),
914 DType::U64 => Tensor::<u64>::new(shape, memory, name).map(Self::U64),
915 DType::I64 => Tensor::<i64>::new(shape, memory, name).map(Self::I64),
916 DType::F16 => Tensor::<f16>::new(shape, memory, name).map(Self::F16),
917 DType::F32 => Tensor::<f32>::new(shape, memory, name).map(Self::F32),
918 DType::F64 => Tensor::<f64>::new(shape, memory, name).map(Self::F64),
919 }
920 }
921
922 #[cfg(unix)]
949 pub fn from_fd(
950 fd: std::os::fd::OwnedFd,
951 shape: &[usize],
952 dtype: DType,
953 name: Option<&str>,
954 ) -> crate::Result<Self> {
955 match dtype {
956 DType::U8 => Tensor::<u8>::from_fd(fd, shape, name).map(Self::U8),
957 DType::I8 => Tensor::<i8>::from_fd(fd, shape, name).map(Self::I8),
958 DType::U16 => Tensor::<u16>::from_fd(fd, shape, name).map(Self::U16),
959 DType::I16 => Tensor::<i16>::from_fd(fd, shape, name).map(Self::I16),
960 DType::U32 => Tensor::<u32>::from_fd(fd, shape, name).map(Self::U32),
961 DType::I32 => Tensor::<i32>::from_fd(fd, shape, name).map(Self::I32),
962 DType::U64 => Tensor::<u64>::from_fd(fd, shape, name).map(Self::U64),
963 DType::I64 => Tensor::<i64>::from_fd(fd, shape, name).map(Self::I64),
964 DType::F16 => Tensor::<f16>::from_fd(fd, shape, name).map(Self::F16),
965 DType::F32 => Tensor::<f32>::from_fd(fd, shape, name).map(Self::F32),
966 DType::F64 => Tensor::<f64>::from_fd(fd, shape, name).map(Self::F64),
967 }
968 }
969
970 pub(crate) fn from_pbo_import(
976 buffer_id: u32,
977 size: usize,
978 shape: &[usize],
979 dtype: DType,
980 ops: std::sync::Arc<dyn crate::PboOps>,
981 ) -> crate::Result<Self> {
982 match dtype {
983 DType::U8 => crate::PboTensor::<u8>::from_pbo(buffer_id, size, shape, None, ops)
984 .and_then(|p| Tensor::from_pbo(p).map(Self::U8)),
985 DType::I8 => crate::PboTensor::<i8>::from_pbo(buffer_id, size, shape, None, ops)
986 .and_then(|p| Tensor::from_pbo(p).map(Self::I8)),
987 DType::U16 => crate::PboTensor::<u16>::from_pbo(buffer_id, size, shape, None, ops)
988 .and_then(|p| Tensor::from_pbo(p).map(Self::U16)),
989 DType::I16 => crate::PboTensor::<i16>::from_pbo(buffer_id, size, shape, None, ops)
990 .and_then(|p| Tensor::from_pbo(p).map(Self::I16)),
991 DType::U32 => crate::PboTensor::<u32>::from_pbo(buffer_id, size, shape, None, ops)
992 .and_then(|p| Tensor::from_pbo(p).map(Self::U32)),
993 DType::I32 => crate::PboTensor::<i32>::from_pbo(buffer_id, size, shape, None, ops)
994 .and_then(|p| Tensor::from_pbo(p).map(Self::I32)),
995 DType::U64 => crate::PboTensor::<u64>::from_pbo(buffer_id, size, shape, None, ops)
996 .and_then(|p| Tensor::from_pbo(p).map(Self::U64)),
997 DType::I64 => crate::PboTensor::<i64>::from_pbo(buffer_id, size, shape, None, ops)
998 .and_then(|p| Tensor::from_pbo(p).map(Self::I64)),
999 DType::F16 => crate::PboTensor::<f16>::from_pbo(buffer_id, size, shape, None, ops)
1000 .and_then(|p| Tensor::from_pbo(p).map(Self::F16)),
1001 DType::F32 => crate::PboTensor::<f32>::from_pbo(buffer_id, size, shape, None, ops)
1002 .and_then(|p| Tensor::from_pbo(p).map(Self::F32)),
1003 DType::F64 => crate::PboTensor::<f64>::from_pbo(buffer_id, size, shape, None, ops)
1004 .and_then(|p| Tensor::from_pbo(p).map(Self::F64)),
1005 }
1006 }
1007
1008 pub unsafe fn from_client_pbo(
1021 state: crate::EfClientState,
1022 map_fn: crate::EfPboMapFn,
1023 unmap_fn: crate::EfPboUnmapFn,
1024 buffer_id: u32,
1025 size: usize,
1026 shape: &[usize],
1027 dtype: DType,
1028 ) -> crate::Result<TensorDyn> {
1029 let ops = unsafe { crate::client_state_pbo_ops(state, map_fn, unmap_fn) }?;
1031 Self::from_pbo_import(buffer_id, size, shape, dtype, ops)
1032 }
1033
1034 pub unsafe fn from_raw_host(
1052 ptr: *mut u8,
1053 shape: &[usize],
1054 dtype: DType,
1055 name: Option<&str>,
1056 ) -> crate::Result<Self> {
1057 unsafe { Self::from_foreign_ptr(ptr, shape, dtype, None, name) }
1058 }
1059
1060 pub unsafe fn from_raw_host_with_capacity(
1076 ptr: *mut u8,
1077 shape: &[usize],
1078 capacity_bytes: usize,
1079 dtype: DType,
1080 name: Option<&str>,
1081 ) -> crate::Result<Self> {
1082 unsafe {
1083 Self::from_foreign_ptr_with_capacity(ptr, shape, dtype, capacity_bytes, None, name)
1084 }
1085 }
1086
1087 pub(crate) fn import_storage(
1099 desc: &crate::TensorDesc,
1100 shape: &[usize],
1101 dtype: DType,
1102 ) -> crate::Result<Self> {
1103 match desc.kind {
1104 #[cfg(target_os = "linux")]
1105 crate::protocol::kind::DMABUF => {
1106 let owned = crate::protocol::dup_descriptor_fd(desc)?;
1111 Self::from_fd(owned, shape, dtype, None)
1112 }
1113 #[cfg(not(target_os = "linux"))]
1114 crate::protocol::kind::DMABUF => Err(crate::Error::NotImplemented(
1115 "dma-buf import off Linux".into(),
1116 )),
1117 #[cfg(any(target_os = "macos", target_os = "ios"))]
1118 crate::protocol::kind::IOSURFACE => {
1119 let id = crate::protocol::descriptor_surface_id(desc)?;
1120 let surface_ref = crate::iosurface::lookup_by_id(id).ok_or_else(|| {
1121 crate::Error::InvalidArgument(format!(
1122 "IOSurface id {id} is not live (the producer's capsule may \
1123 already have been released)"
1124 ))
1125 })?;
1126 let result = unsafe { Self::from_iosurface(surface_ref, shape, dtype, None) };
1131 crate::iosurface::release(surface_ref);
1132 result
1133 }
1134 #[cfg(not(any(target_os = "macos", target_os = "ios")))]
1135 crate::protocol::kind::IOSURFACE => Err(crate::Error::NotImplemented(
1136 "IOSurface import off Apple platforms".into(),
1137 )),
1138 crate::protocol::kind::HOST => {
1139 crate::protocol::check_descriptor_host_ptr(desc)?;
1140 unsafe {
1147 Self::from_raw_host_with_capacity(
1148 desc.ptr.0,
1149 shape,
1150 desc.capacity as usize,
1151 dtype,
1152 None,
1153 )
1154 }
1155 }
1156 crate::protocol::kind::PBO => {
1157 let buffer_id = crate::protocol::descriptor_pbo_buffer_id(desc)?;
1158 let parts = unsafe { crate::read_pbo_vtable_parts(desc.ptr.0 as *const _) }?;
1168 unsafe {
1171 Self::from_client_pbo(
1172 parts.state,
1173 parts.map_fn,
1174 parts.unmap_fn,
1175 buffer_id,
1176 desc.capacity as usize,
1177 shape,
1178 dtype,
1179 )
1180 }
1181 }
1182 #[cfg(target_os = "windows")]
1183 crate::protocol::kind::D3D11_TEXTURE => {
1184 let (tex, completion) = crate::protocol::descriptor_d3d11_handles(desc)?;
1185 unsafe {
1192 let (format, width, height) =
1193 crate::protocol::descriptor_d3d11_geometry(desc, tex, shape)?;
1194 Self::from_d3d11_shared_handle(
1195 tex,
1196 width,
1197 height,
1198 format,
1199 dtype,
1200 crate::CpuAccess::ReadWrite,
1201 completion,
1202 None,
1203 )
1204 }
1205 }
1206 #[cfg(not(target_os = "windows"))]
1207 crate::protocol::kind::D3D11_TEXTURE => Err(crate::Error::NotImplemented(
1208 "D3D11 texture import off Windows".into(),
1209 )),
1210 k => Err(crate::Error::NotImplemented(format!(
1211 "tensor interop kind {k} cannot be imported by this build"
1212 ))),
1213 }
1214 }
1215
1216 pub unsafe fn from_foreign_ptr(
1229 ptr: *mut u8,
1230 shape: &[usize],
1231 dtype: DType,
1232 owner: Option<crate::ForeignOwner>,
1233 name: Option<&str>,
1234 ) -> crate::Result<Self> {
1235 unsafe {
1240 match dtype {
1241 DType::U8 => {
1242 Tensor::<u8>::from_foreign(ptr.cast(), shape, owner, name).map(Self::U8)
1243 }
1244 DType::I8 => {
1245 Tensor::<i8>::from_foreign(ptr.cast(), shape, owner, name).map(Self::I8)
1246 }
1247 DType::U16 => {
1248 Tensor::<u16>::from_foreign(ptr.cast(), shape, owner, name).map(Self::U16)
1249 }
1250 DType::I16 => {
1251 Tensor::<i16>::from_foreign(ptr.cast(), shape, owner, name).map(Self::I16)
1252 }
1253 DType::U32 => {
1254 Tensor::<u32>::from_foreign(ptr.cast(), shape, owner, name).map(Self::U32)
1255 }
1256 DType::I32 => {
1257 Tensor::<i32>::from_foreign(ptr.cast(), shape, owner, name).map(Self::I32)
1258 }
1259 DType::U64 => {
1260 Tensor::<u64>::from_foreign(ptr.cast(), shape, owner, name).map(Self::U64)
1261 }
1262 DType::I64 => {
1263 Tensor::<i64>::from_foreign(ptr.cast(), shape, owner, name).map(Self::I64)
1264 }
1265 DType::F16 => {
1266 Tensor::<f16>::from_foreign(ptr.cast(), shape, owner, name).map(Self::F16)
1267 }
1268 DType::F32 => {
1269 Tensor::<f32>::from_foreign(ptr.cast(), shape, owner, name).map(Self::F32)
1270 }
1271 DType::F64 => {
1272 Tensor::<f64>::from_foreign(ptr.cast(), shape, owner, name).map(Self::F64)
1273 }
1274 }
1275 }
1276 }
1277
1278 pub unsafe fn from_foreign_ptr_with_capacity(
1289 ptr: *mut u8,
1290 shape: &[usize],
1291 dtype: DType,
1292 capacity_bytes: usize,
1293 owner: Option<crate::ForeignOwner>,
1294 name: Option<&str>,
1295 ) -> crate::Result<Self> {
1296 macro_rules! import {
1297 ($t:ty, $variant:ident) => {
1298 unsafe {
1303 Tensor::<$t>::from_foreign_with_capacity(
1304 ptr.cast(),
1305 shape,
1306 capacity_bytes,
1307 owner,
1308 name,
1309 )
1310 }
1311 .map(Self::$variant)
1312 };
1313 }
1314 match dtype {
1315 DType::U8 => import!(u8, U8),
1316 DType::I8 => import!(i8, I8),
1317 DType::U16 => import!(u16, U16),
1318 DType::I16 => import!(i16, I16),
1319 DType::U32 => import!(u32, U32),
1320 DType::I32 => import!(i32, I32),
1321 DType::U64 => import!(u64, U64),
1322 DType::I64 => import!(i64, I64),
1323 DType::F16 => import!(f16, F16),
1324 DType::F32 => import!(f32, F32),
1325 DType::F64 => import!(f64, F64),
1326 }
1327 }
1328
1329 #[cfg(any(target_os = "macos", target_os = "ios"))]
1340 pub fn from_iosurface_id(
1341 id: u32,
1342 shape: &[usize],
1343 dtype: DType,
1344 name: Option<&str>,
1345 ) -> crate::Result<Self> {
1346 let surface = crate::iosurface::lookup_by_id(id).ok_or_else(|| {
1347 crate::Error::InvalidArgument(format!("no live IOSurface with id {id}"))
1348 })?;
1349 let result = unsafe { Self::from_iosurface(surface, shape, dtype, name) };
1364 crate::iosurface::release(surface);
1365 result
1366 }
1367
1368 #[cfg(any(target_os = "macos", target_os = "ios"))]
1376 pub unsafe fn from_iosurface(
1377 surface_ref: *mut std::ffi::c_void,
1378 shape: &[usize],
1379 dtype: DType,
1380 name: Option<&str>,
1381 ) -> crate::Result<Self> {
1382 unsafe {
1383 match dtype {
1384 DType::U8 => Tensor::<u8>::from_iosurface(surface_ref, shape, name).map(Self::U8),
1385 DType::I8 => Tensor::<i8>::from_iosurface(surface_ref, shape, name).map(Self::I8),
1386 DType::U16 => {
1387 Tensor::<u16>::from_iosurface(surface_ref, shape, name).map(Self::U16)
1388 }
1389 DType::I16 => {
1390 Tensor::<i16>::from_iosurface(surface_ref, shape, name).map(Self::I16)
1391 }
1392 DType::U32 => {
1393 Tensor::<u32>::from_iosurface(surface_ref, shape, name).map(Self::U32)
1394 }
1395 DType::I32 => {
1396 Tensor::<i32>::from_iosurface(surface_ref, shape, name).map(Self::I32)
1397 }
1398 DType::U64 => {
1399 Tensor::<u64>::from_iosurface(surface_ref, shape, name).map(Self::U64)
1400 }
1401 DType::I64 => {
1402 Tensor::<i64>::from_iosurface(surface_ref, shape, name).map(Self::I64)
1403 }
1404 DType::F16 => {
1405 Tensor::<f16>::from_iosurface(surface_ref, shape, name).map(Self::F16)
1406 }
1407 DType::F32 => {
1408 Tensor::<f32>::from_iosurface(surface_ref, shape, name).map(Self::F32)
1409 }
1410 DType::F64 => {
1411 Tensor::<f64>::from_iosurface(surface_ref, shape, name).map(Self::F64)
1412 }
1413 }
1414 }
1415 }
1416
1417 #[cfg(any(target_os = "macos", target_os = "ios"))]
1420 pub fn iosurface_id(&self) -> Option<u32> {
1421 dispatch!(self, iosurface_id)
1422 }
1423
1424 pub fn pbo_id(&self) -> Option<u32> {
1427 dispatch!(self, pbo_id)
1428 }
1429
1430 pub fn pbo_is_mapped(&self) -> Option<bool> {
1434 dispatch!(self, pbo_is_mapped)
1435 }
1436
1437 pub fn pbo_vtable_ptr(&self) -> Option<*const std::ffi::c_void> {
1441 dispatch!(self, pbo_vtable_ptr)
1442 }
1443
1444 pub fn pbo_keepalive(&self) -> Option<std::sync::Arc<dyn Send + Sync>> {
1449 dispatch!(self, pbo_keepalive)
1450 }
1451
1452 #[cfg(any(target_os = "macos", target_os = "ios"))]
1456 pub fn iosurface_ref(&self) -> Option<*mut std::ffi::c_void> {
1457 dispatch!(self, iosurface_ref)
1458 }
1459
1460 #[cfg(any(target_os = "macos", target_os = "ios"))]
1465 pub fn iosurface_physical_dims(&self) -> Option<(usize, usize)> {
1466 dispatch!(self, iosurface_physical_dims)
1467 }
1468
1469 #[cfg(target_os = "windows")]
1472 pub fn d3d11_texture(&self) -> Option<*mut std::ffi::c_void> {
1473 dispatch!(self, d3d11_texture)
1474 }
1475
1476 #[cfg(target_os = "windows")]
1479 pub fn d3d11_layout(&self) -> Option<crate::d3d11_layout::D3d11ImageLayout> {
1480 dispatch!(self, d3d11_layout)
1481 }
1482
1483 #[cfg(target_os = "windows")]
1486 pub fn d3d11_shared_handle(&self) -> crate::Result<std::os::windows::io::OwnedHandle> {
1487 dispatch!(self, d3d11_shared_handle)
1488 }
1489
1490 #[cfg(target_os = "windows")]
1493 pub fn gpu_completion(&self) -> crate::Result<Option<crate::d3d11::GpuCompletion>> {
1494 dispatch!(self, gpu_completion)
1495 }
1496
1497 pub fn gpu_write_value(&self) -> u64 {
1501 dispatch!(self, gpu_write_value)
1502 }
1503
1504 #[cfg(target_os = "windows")]
1508 pub fn set_gpu_write(&self, value: u64) -> crate::Result<()> {
1509 dispatch!(self, set_gpu_write, value)
1510 }
1511
1512 #[cfg(target_os = "windows")]
1515 pub(crate) fn d3d11_shared_handle_value(&self) -> Option<usize> {
1516 dispatch!(self, d3d11_shared_handle_value)
1517 }
1518
1519 #[cfg(target_os = "windows")]
1527 #[allow(clippy::too_many_arguments)] pub unsafe fn from_d3d11_texture(
1529 texture: *mut std::ffi::c_void,
1530 width: usize,
1531 height: usize,
1532 format: PixelFormat,
1533 dtype: DType,
1534 access: crate::CpuAccess,
1535 name: Option<&str>,
1536 ) -> crate::Result<Self> {
1537 unsafe {
1540 match dtype {
1541 DType::U8 => {
1542 Tensor::<u8>::from_d3d11_texture(texture, width, height, format, access, name)
1543 .map(Self::U8)
1544 }
1545 DType::I8 => {
1546 Tensor::<i8>::from_d3d11_texture(texture, width, height, format, access, name)
1547 .map(Self::I8)
1548 }
1549 DType::U16 => {
1550 Tensor::<u16>::from_d3d11_texture(texture, width, height, format, access, name)
1551 .map(Self::U16)
1552 }
1553 DType::I16 => {
1554 Tensor::<i16>::from_d3d11_texture(texture, width, height, format, access, name)
1555 .map(Self::I16)
1556 }
1557 DType::U32 => {
1558 Tensor::<u32>::from_d3d11_texture(texture, width, height, format, access, name)
1559 .map(Self::U32)
1560 }
1561 DType::I32 => {
1562 Tensor::<i32>::from_d3d11_texture(texture, width, height, format, access, name)
1563 .map(Self::I32)
1564 }
1565 DType::U64 => {
1566 Tensor::<u64>::from_d3d11_texture(texture, width, height, format, access, name)
1567 .map(Self::U64)
1568 }
1569 DType::I64 => {
1570 Tensor::<i64>::from_d3d11_texture(texture, width, height, format, access, name)
1571 .map(Self::I64)
1572 }
1573 DType::F16 => {
1574 Tensor::<f16>::from_d3d11_texture(texture, width, height, format, access, name)
1575 .map(Self::F16)
1576 }
1577 DType::F32 => {
1578 Tensor::<f32>::from_d3d11_texture(texture, width, height, format, access, name)
1579 .map(Self::F32)
1580 }
1581 DType::F64 => {
1582 Tensor::<f64>::from_d3d11_texture(texture, width, height, format, access, name)
1583 .map(Self::F64)
1584 }
1585 }
1586 }
1587 }
1588
1589 #[cfg(target_os = "windows")]
1598 #[allow(clippy::too_many_arguments)] pub unsafe fn from_d3d11_shared_handle(
1600 handle: std::os::windows::io::RawHandle,
1601 width: usize,
1602 height: usize,
1603 format: PixelFormat,
1604 dtype: DType,
1605 access: crate::CpuAccess,
1606 completion: Option<(std::os::windows::io::RawHandle, u64)>,
1607 name: Option<&str>,
1608 ) -> crate::Result<Self> {
1609 unsafe {
1612 match dtype {
1613 DType::U8 => Tensor::<u8>::from_d3d11_shared_handle(
1614 handle, width, height, format, access, completion, name,
1615 )
1616 .map(Self::U8),
1617 DType::I8 => Tensor::<i8>::from_d3d11_shared_handle(
1618 handle, width, height, format, access, completion, name,
1619 )
1620 .map(Self::I8),
1621 DType::U16 => Tensor::<u16>::from_d3d11_shared_handle(
1622 handle, width, height, format, access, completion, name,
1623 )
1624 .map(Self::U16),
1625 DType::I16 => Tensor::<i16>::from_d3d11_shared_handle(
1626 handle, width, height, format, access, completion, name,
1627 )
1628 .map(Self::I16),
1629 DType::U32 => Tensor::<u32>::from_d3d11_shared_handle(
1630 handle, width, height, format, access, completion, name,
1631 )
1632 .map(Self::U32),
1633 DType::I32 => Tensor::<i32>::from_d3d11_shared_handle(
1634 handle, width, height, format, access, completion, name,
1635 )
1636 .map(Self::I32),
1637 DType::U64 => Tensor::<u64>::from_d3d11_shared_handle(
1638 handle, width, height, format, access, completion, name,
1639 )
1640 .map(Self::U64),
1641 DType::I64 => Tensor::<i64>::from_d3d11_shared_handle(
1642 handle, width, height, format, access, completion, name,
1643 )
1644 .map(Self::I64),
1645 DType::F16 => Tensor::<f16>::from_d3d11_shared_handle(
1646 handle, width, height, format, access, completion, name,
1647 )
1648 .map(Self::F16),
1649 DType::F32 => Tensor::<f32>::from_d3d11_shared_handle(
1650 handle, width, height, format, access, completion, name,
1651 )
1652 .map(Self::F32),
1653 DType::F64 => Tensor::<f64>::from_d3d11_shared_handle(
1654 handle, width, height, format, access, completion, name,
1655 )
1656 .map(Self::F64),
1657 }
1658 }
1659 }
1660
1661 #[cfg(target_os = "android")]
1671 pub unsafe fn from_hardware_buffer(
1672 buffer_ptr: *mut std::ffi::c_void,
1673 shape: &[usize],
1674 dtype: DType,
1675 name: Option<&str>,
1676 ) -> crate::Result<Self> {
1677 unsafe {
1678 match dtype {
1679 DType::U8 => {
1680 Tensor::<u8>::from_hardware_buffer(buffer_ptr, shape, name).map(Self::U8)
1681 }
1682 DType::I8 => {
1683 Tensor::<i8>::from_hardware_buffer(buffer_ptr, shape, name).map(Self::I8)
1684 }
1685 DType::U16 => {
1686 Tensor::<u16>::from_hardware_buffer(buffer_ptr, shape, name).map(Self::U16)
1687 }
1688 DType::I16 => {
1689 Tensor::<i16>::from_hardware_buffer(buffer_ptr, shape, name).map(Self::I16)
1690 }
1691 DType::U32 => {
1692 Tensor::<u32>::from_hardware_buffer(buffer_ptr, shape, name).map(Self::U32)
1693 }
1694 DType::I32 => {
1695 Tensor::<i32>::from_hardware_buffer(buffer_ptr, shape, name).map(Self::I32)
1696 }
1697 DType::U64 => {
1698 Tensor::<u64>::from_hardware_buffer(buffer_ptr, shape, name).map(Self::U64)
1699 }
1700 DType::I64 => {
1701 Tensor::<i64>::from_hardware_buffer(buffer_ptr, shape, name).map(Self::I64)
1702 }
1703 DType::F16 => {
1704 Tensor::<f16>::from_hardware_buffer(buffer_ptr, shape, name).map(Self::F16)
1705 }
1706 DType::F32 => {
1707 Tensor::<f32>::from_hardware_buffer(buffer_ptr, shape, name).map(Self::F32)
1708 }
1709 DType::F64 => {
1710 Tensor::<f64>::from_hardware_buffer(buffer_ptr, shape, name).map(Self::F64)
1711 }
1712 }
1713 }
1714 }
1715
1716 #[cfg(target_os = "android")]
1720 pub fn hardware_buffer_ptr(&self) -> Option<*mut std::ffi::c_void> {
1721 dispatch!(self, hardware_buffer_ptr)
1722 }
1723
1724 #[cfg(target_os = "android")]
1728 pub fn hardware_buffer_physical_dims(&self) -> Option<(usize, usize)> {
1729 dispatch!(self, hardware_buffer_physical_dims)
1730 }
1731
1732 pub fn copy_to_flat(&self, dst: &mut [u8]) -> crate::Result<()> {
1737 dispatch!(self, copy_to_flat, dst)
1738 }
1739
1740 pub fn image(
1758 width: usize,
1759 height: usize,
1760 format: PixelFormat,
1761 dtype: DType,
1762 memory: Option<TensorMemory>,
1763 access: crate::CpuAccess,
1764 ) -> crate::Result<Self> {
1765 match dtype {
1766 DType::U8 => Tensor::<u8>::image(width, height, format, memory, access).map(Self::U8),
1767 DType::I8 => Tensor::<i8>::image(width, height, format, memory, access).map(Self::I8),
1768 DType::U16 => {
1769 Tensor::<u16>::image(width, height, format, memory, access).map(Self::U16)
1770 }
1771 DType::I16 => {
1772 Tensor::<i16>::image(width, height, format, memory, access).map(Self::I16)
1773 }
1774 DType::U32 => {
1775 Tensor::<u32>::image(width, height, format, memory, access).map(Self::U32)
1776 }
1777 DType::I32 => {
1778 Tensor::<i32>::image(width, height, format, memory, access).map(Self::I32)
1779 }
1780 DType::U64 => {
1781 Tensor::<u64>::image(width, height, format, memory, access).map(Self::U64)
1782 }
1783 DType::I64 => {
1784 Tensor::<i64>::image(width, height, format, memory, access).map(Self::I64)
1785 }
1786 DType::F16 => {
1787 Tensor::<f16>::image(width, height, format, memory, access).map(Self::F16)
1788 }
1789 DType::F32 => {
1790 Tensor::<f32>::image(width, height, format, memory, access).map(Self::F32)
1791 }
1792 DType::F64 => {
1793 Tensor::<f64>::image(width, height, format, memory, access).map(Self::F64)
1794 }
1795 }
1796 }
1797
1798 pub fn image_desc(desc: &crate::ImageDesc) -> crate::Result<Self> {
1802 match desc.dtype() {
1803 DType::U8 => Tensor::<u8>::image_desc(desc).map(Self::U8),
1804 DType::I8 => Tensor::<i8>::image_desc(desc).map(Self::I8),
1805 DType::U16 => Tensor::<u16>::image_desc(desc).map(Self::U16),
1806 DType::I16 => Tensor::<i16>::image_desc(desc).map(Self::I16),
1807 DType::U32 => Tensor::<u32>::image_desc(desc).map(Self::U32),
1808 DType::I32 => Tensor::<i32>::image_desc(desc).map(Self::I32),
1809 DType::U64 => Tensor::<u64>::image_desc(desc).map(Self::U64),
1810 DType::I64 => Tensor::<i64>::image_desc(desc).map(Self::I64),
1811 DType::F16 => Tensor::<f16>::image_desc(desc).map(Self::F16),
1812 DType::F32 => Tensor::<f32>::image_desc(desc).map(Self::F32),
1813 DType::F64 => Tensor::<f64>::image_desc(desc).map(Self::F64),
1814 }
1815 }
1816
1817 pub fn compression(&self) -> Option<crate::CompressionScheme> {
1820 dispatch!(self, compression)
1821 }
1822
1823 pub fn contiguity(&self) -> crate::Contiguity {
1826 dispatch!(self, contiguity)
1827 }
1828
1829 pub fn image_with_stride(
1856 width: usize,
1857 height: usize,
1858 format: PixelFormat,
1859 dtype: DType,
1860 row_stride_bytes: usize,
1861 memory: Option<TensorMemory>,
1862 access: crate::CpuAccess,
1863 ) -> crate::Result<Self> {
1864 match dtype {
1865 DType::U8 => Tensor::<u8>::image_with_stride(
1866 width,
1867 height,
1868 format,
1869 row_stride_bytes,
1870 memory,
1871 access,
1872 )
1873 .map(Self::U8),
1874 DType::I8 => Tensor::<i8>::image_with_stride(
1875 width,
1876 height,
1877 format,
1878 row_stride_bytes,
1879 memory,
1880 access,
1881 )
1882 .map(Self::I8),
1883 DType::U16 => Tensor::<u16>::image_with_stride(
1884 width,
1885 height,
1886 format,
1887 row_stride_bytes,
1888 memory,
1889 access,
1890 )
1891 .map(Self::U16),
1892 DType::I16 => Tensor::<i16>::image_with_stride(
1893 width,
1894 height,
1895 format,
1896 row_stride_bytes,
1897 memory,
1898 access,
1899 )
1900 .map(Self::I16),
1901 DType::U32 => Tensor::<u32>::image_with_stride(
1902 width,
1903 height,
1904 format,
1905 row_stride_bytes,
1906 memory,
1907 access,
1908 )
1909 .map(Self::U32),
1910 DType::I32 => Tensor::<i32>::image_with_stride(
1911 width,
1912 height,
1913 format,
1914 row_stride_bytes,
1915 memory,
1916 access,
1917 )
1918 .map(Self::I32),
1919 DType::U64 => Tensor::<u64>::image_with_stride(
1920 width,
1921 height,
1922 format,
1923 row_stride_bytes,
1924 memory,
1925 access,
1926 )
1927 .map(Self::U64),
1928 DType::I64 => Tensor::<i64>::image_with_stride(
1929 width,
1930 height,
1931 format,
1932 row_stride_bytes,
1933 memory,
1934 access,
1935 )
1936 .map(Self::I64),
1937 DType::F16 => Tensor::<f16>::image_with_stride(
1938 width,
1939 height,
1940 format,
1941 row_stride_bytes,
1942 memory,
1943 access,
1944 )
1945 .map(Self::F16),
1946 DType::F32 => Tensor::<f32>::image_with_stride(
1947 width,
1948 height,
1949 format,
1950 row_stride_bytes,
1951 memory,
1952 access,
1953 )
1954 .map(Self::F32),
1955 DType::F64 => Tensor::<f64>::image_with_stride(
1956 width,
1957 height,
1958 format,
1959 row_stride_bytes,
1960 memory,
1961 access,
1962 )
1963 .map(Self::F64),
1964 }
1965 }
1966}
1967
1968impl From<Tensor<u8>> for TensorDyn {
1971 fn from(t: Tensor<u8>) -> Self {
1972 Self::U8(t)
1973 }
1974}
1975
1976impl From<Tensor<i8>> for TensorDyn {
1977 fn from(t: Tensor<i8>) -> Self {
1978 Self::I8(t)
1979 }
1980}
1981
1982impl From<Tensor<u16>> for TensorDyn {
1983 fn from(t: Tensor<u16>) -> Self {
1984 Self::U16(t)
1985 }
1986}
1987
1988impl From<Tensor<i16>> for TensorDyn {
1989 fn from(t: Tensor<i16>) -> Self {
1990 Self::I16(t)
1991 }
1992}
1993
1994impl From<Tensor<u32>> for TensorDyn {
1995 fn from(t: Tensor<u32>) -> Self {
1996 Self::U32(t)
1997 }
1998}
1999
2000impl From<Tensor<i32>> for TensorDyn {
2001 fn from(t: Tensor<i32>) -> Self {
2002 Self::I32(t)
2003 }
2004}
2005
2006impl From<Tensor<u64>> for TensorDyn {
2007 fn from(t: Tensor<u64>) -> Self {
2008 Self::U64(t)
2009 }
2010}
2011
2012impl From<Tensor<i64>> for TensorDyn {
2013 fn from(t: Tensor<i64>) -> Self {
2014 Self::I64(t)
2015 }
2016}
2017
2018impl From<Tensor<f16>> for TensorDyn {
2019 fn from(t: Tensor<f16>) -> Self {
2020 Self::F16(t)
2021 }
2022}
2023
2024impl From<Tensor<f32>> for TensorDyn {
2025 fn from(t: Tensor<f32>) -> Self {
2026 Self::F32(t)
2027 }
2028}
2029
2030impl From<Tensor<f64>> for TensorDyn {
2031 fn from(t: Tensor<f64>) -> Self {
2032 Self::F64(t)
2033 }
2034}
2035
2036impl fmt::Debug for TensorDyn {
2037 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
2038 dispatch!(self, fmt, f)
2039 }
2040}
2041
2042pub type Raw = *mut std::ffi::c_void;
2045
2046impl TensorDyn {
2047 pub fn into_raw(self) -> Raw {
2050 Box::into_raw(Box::new(self)) as Raw
2051 }
2052
2053 pub unsafe fn from_raw(p: Raw) -> TensorDyn {
2059 *unsafe { Box::from_raw(p as *mut TensorDyn) }
2060 }
2061
2062 pub unsafe fn with_raw<R>(p: Raw, f: impl FnOnce(&mut TensorDyn) -> R) -> R {
2068 let mut t = std::mem::ManuallyDrop::new(unsafe { std::ptr::read(p as *mut TensorDyn) });
2069 f(&mut t)
2070 }
2071}
2072
2073#[cfg(test)]
2074mod tests {
2075 use super::*;
2076
2077 #[test]
2078 fn from_typed_tensor() {
2079 let t = Tensor::<u8>::new(&[10], None, None).unwrap();
2080 let dyn_t: TensorDyn = t.into();
2081 assert_eq!(dyn_t.dtype(), DType::U8);
2082 assert_eq!(dyn_t.shape(), &[10]);
2083 }
2084
2085 #[test]
2086 fn from_foreign_ptr_wraps_borrowed_memory() {
2087 use crate::TensorMapTrait;
2088 let mut vec: Vec<f32> = vec![0.0; 4];
2091 let ptr = vec.as_mut_ptr() as *mut u8;
2092 let owner: crate::ForeignOwner = Box::new(vec);
2093 let t = unsafe {
2094 TensorDyn::from_foreign_ptr(ptr, &[2, 2], DType::F32, Some(owner), Some("trt_output"))
2095 }
2096 .unwrap();
2097 assert_eq!(t.dtype(), DType::F32);
2098 assert_eq!(t.memory(), TensorMemory::Mem);
2099 assert_eq!(t.shape(), &[2, 2]);
2100 {
2101 let mut m = t.as_f32().unwrap().map().unwrap();
2102 m.as_mut_slice().copy_from_slice(&[1.0, 2.0, 3.0, 4.0]);
2103 }
2104 let m = t.as_f32().unwrap().map().unwrap();
2105 assert_eq!(m.as_slice(), &[1.0, 2.0, 3.0, 4.0]);
2106 }
2107
2108 #[test]
2117 fn from_foreign_ptr_rejects_null_ptr() {
2118 let err = unsafe {
2119 TensorDyn::from_foreign_ptr(std::ptr::null_mut(), &[4], DType::U8, None, None)
2120 }
2121 .unwrap_err();
2122 assert!(
2124 matches!(err, crate::error::Error::InvalidArgument(ref m) if m.contains("non-null")),
2125 "expected InvalidArgument(non-null), got {err:?}"
2126 );
2127 }
2128
2129 #[test]
2130 fn from_foreign_ptr_rejects_empty_shape() {
2131 let mut dummy: u8 = 0;
2132 let err = unsafe {
2133 TensorDyn::from_foreign_ptr(&mut dummy as *mut u8, &[], DType::U8, None, None)
2134 }
2135 .unwrap_err();
2136 assert!(
2137 matches!(err, crate::error::Error::InvalidSize(0)),
2138 "expected InvalidSize(0) for empty shape, got {err:?}"
2139 );
2140 }
2141
2142 #[test]
2143 fn from_foreign_ptr_rejects_overflow_shape() {
2144 let mut dummy: u8 = 0;
2145 let huge = [usize::MAX / 2 + 1, 2];
2146 let err = unsafe { TensorDyn::from_foreign_ptr(&mut dummy, &huge, DType::U8, None, None) }
2147 .unwrap_err();
2148 assert!(
2149 matches!(err, crate::error::Error::InvalidArgument(ref m) if m.contains("overflow")),
2150 "expected InvalidArgument(overflow), got {err:?}"
2151 );
2152 }
2153
2154 #[test]
2155 fn from_foreign_ptr_u8_dtype_dispatch() {
2156 let mut buf: Vec<u8> = vec![1, 2, 3, 4];
2159 let ptr = buf.as_mut_ptr();
2160 let owner: crate::ForeignOwner = Box::new(buf);
2161 let t = unsafe {
2162 TensorDyn::from_foreign_ptr(ptr, &[4], DType::U8, Some(owner), Some("u8_foreign"))
2163 }
2164 .unwrap();
2165 assert_eq!(t.dtype(), DType::U8);
2166 assert_eq!(t.shape(), &[4]);
2167 let m = t.as_u8().unwrap().map().unwrap();
2168 use crate::TensorMapTrait;
2169 assert_eq!(m.as_slice(), &[1u8, 2, 3, 4]);
2170 }
2171
2172 #[test]
2173 fn map_bytes_roundtrips_and_is_typed_agnostic() {
2174 use crate::TensorMapTrait;
2175 let t = Tensor::<f32>::new(&[4], None, None).unwrap();
2178 let dyn_t: TensorDyn = t.into();
2179 {
2180 let mut view = dyn_t.map_bytes(crate::CpuAccess::ReadWrite).unwrap();
2181 assert_eq!(view.len(), 16, "4 x f32 is 16 bytes, not 4");
2182 view.as_mut_slice()
2183 .copy_from_slice(&[1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16]);
2184 }
2185 let view = dyn_t.map_bytes(crate::CpuAccess::Read).unwrap();
2186 assert_eq!(
2187 view.as_slice(),
2188 &[1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16]
2189 );
2190 }
2191
2192 #[test]
2193 fn map_bytes_none_is_invalid() {
2194 let t = Tensor::<u8>::new(&[4], None, None).unwrap();
2195 let dyn_t: TensorDyn = t.into();
2196 assert!(matches!(
2197 dyn_t.map_bytes(crate::CpuAccess::None),
2198 Err(crate::error::Error::InvalidArgument(_))
2199 ));
2200 }
2201
2202 #[test]
2203 fn downcast_ref() {
2204 let t = Tensor::<u8>::new(&[10], None, None).unwrap();
2205 let dyn_t: TensorDyn = t.into();
2206 assert!(dyn_t.as_u8().is_some());
2207 assert!(dyn_t.as_i8().is_none());
2208 }
2209
2210 #[test]
2211 fn downcast_into() {
2212 let t = Tensor::<u8>::new(&[10], None, None).unwrap();
2213 let dyn_t: TensorDyn = t.into();
2214 let back = dyn_t.into_u8().unwrap();
2215 assert_eq!(back.shape(), &[10]);
2216 }
2217
2218 #[test]
2219 fn image_accessors() {
2220 let t = Tensor::<u8>::image(
2221 640,
2222 480,
2223 PixelFormat::Rgba,
2224 None,
2225 crate::CpuAccess::ReadWrite,
2226 )
2227 .unwrap();
2228 let dyn_t: TensorDyn = t.into();
2229 assert_eq!(dyn_t.format(), Some(PixelFormat::Rgba));
2230 assert_eq!(dyn_t.width(), Some(640));
2231 assert_eq!(dyn_t.height(), Some(480));
2232 assert!(!dyn_t.is_multiplane());
2233 }
2234
2235 #[test]
2236 fn cpu_access_reports_the_declared_access() {
2237 let read_t =
2239 Tensor::<u8>::image(4, 4, PixelFormat::Rgba, None, crate::CpuAccess::Read).unwrap();
2240 let dyn_read: TensorDyn = read_t.into();
2241 assert_eq!(dyn_read.cpu_access(), crate::CpuAccess::Read);
2242 assert!(!dyn_read.cpu_access().writes());
2243
2244 let mem_t = Tensor::<u8>::new(&[4], None, None).unwrap();
2248 let dyn_mem: TensorDyn = mem_t.into();
2249 assert_eq!(dyn_mem.cpu_access(), crate::CpuAccess::ReadWrite);
2250 assert!(dyn_mem.cpu_access().writes());
2251 }
2252
2253 #[test]
2254 fn image_constructor() {
2255 let dyn_t = TensorDyn::image(
2256 640,
2257 480,
2258 PixelFormat::Rgb,
2259 DType::U8,
2260 None,
2261 crate::CpuAccess::ReadWrite,
2262 )
2263 .unwrap();
2264 assert_eq!(dyn_t.dtype(), DType::U8);
2265 assert_eq!(dyn_t.format(), Some(PixelFormat::Rgb));
2266 assert_eq!(dyn_t.width(), Some(640));
2267 }
2268
2269 #[test]
2270 fn image_constructor_i8() {
2271 let dyn_t = TensorDyn::image(
2272 640,
2273 480,
2274 PixelFormat::Rgb,
2275 DType::I8,
2276 None,
2277 crate::CpuAccess::ReadWrite,
2278 )
2279 .unwrap();
2280 assert_eq!(dyn_t.dtype(), DType::I8);
2281 assert_eq!(dyn_t.format(), Some(PixelFormat::Rgb));
2282 }
2283
2284 #[test]
2285 fn set_format_packed() {
2286 let mut t = TensorDyn::new(&[480, 640, 3], DType::U8, None, None).unwrap();
2287 assert_eq!(t.format(), None);
2288 t.set_format(PixelFormat::Rgb).unwrap();
2289 assert_eq!(t.format(), Some(PixelFormat::Rgb));
2290 assert_eq!(t.width(), Some(640));
2291 assert_eq!(t.height(), Some(480));
2292 }
2293
2294 #[test]
2295 fn set_format_planar() {
2296 let mut t = TensorDyn::new(&[3, 480, 640], DType::U8, None, None).unwrap();
2297 t.set_format(PixelFormat::PlanarRgb).unwrap();
2298 assert_eq!(t.format(), Some(PixelFormat::PlanarRgb));
2299 assert_eq!(t.width(), Some(640));
2300 assert_eq!(t.height(), Some(480));
2301 }
2302
2303 #[test]
2304 fn set_format_rejects_wrong_shape() {
2305 let mut t = TensorDyn::new(&[480, 640, 4], DType::U8, None, None).unwrap();
2306 assert!(t.set_format(PixelFormat::Rgb).is_err());
2307 }
2308
2309 #[test]
2310 fn with_format_builder() {
2311 let t = TensorDyn::new(&[480, 640, 4], DType::U8, None, None)
2312 .unwrap()
2313 .with_format(PixelFormat::Rgba)
2314 .unwrap();
2315 assert_eq!(t.format(), Some(PixelFormat::Rgba));
2316 assert_eq!(t.width(), Some(640));
2317 assert_eq!(t.height(), Some(480));
2318 }
2319
2320 #[cfg(target_os = "linux")]
2321 #[test]
2322 fn dmabuf_clone_mem_tensor_fails() {
2323 let t = TensorDyn::new(&[480, 640, 3], DType::U8, Some(TensorMemory::Mem), None).unwrap();
2324 assert_eq!(t.memory(), TensorMemory::Mem);
2325 assert!(t.dmabuf_clone().is_err());
2326 }
2327
2328 #[cfg(target_os = "linux")]
2329 #[test]
2330 fn dmabuf_mem_tensor_fails() {
2331 let t = TensorDyn::new(&[480, 640, 3], DType::U8, Some(TensorMemory::Mem), None).unwrap();
2332 assert!(t.dmabuf().is_err());
2333 }
2334
2335 #[test]
2336 fn set_format_semi_planar_nv12() {
2337 let mut t = TensorDyn::new(&[720, 640], DType::U8, Some(TensorMemory::Mem), None).unwrap();
2339 t.set_format(PixelFormat::Nv12).unwrap();
2340 assert_eq!(t.format(), Some(PixelFormat::Nv12));
2341 assert_eq!(t.width(), Some(640));
2342 assert_eq!(t.height(), Some(480));
2343 }
2344
2345 #[test]
2346 fn set_format_semi_planar_nv16() {
2347 let mut t = TensorDyn::new(&[960, 640], DType::U8, Some(TensorMemory::Mem), None).unwrap();
2349 t.set_format(PixelFormat::Nv16).unwrap();
2350 assert_eq!(t.format(), Some(PixelFormat::Nv16));
2351 assert_eq!(t.width(), Some(640));
2352 assert_eq!(t.height(), Some(480));
2353 }
2354
2355 #[test]
2356 fn with_format_rejects_wrong_shape() {
2357 let result = TensorDyn::new(&[480, 640, 4], DType::U8, None, None)
2358 .unwrap()
2359 .with_format(PixelFormat::Rgb);
2360 assert!(result.is_err());
2361 }
2362
2363 #[test]
2364 fn set_format_preserved_after_rejection() {
2365 let mut t = TensorDyn::new(&[480, 640, 3], DType::U8, None, None).unwrap();
2366 t.set_format(PixelFormat::Rgb).unwrap();
2367 assert_eq!(t.format(), Some(PixelFormat::Rgb));
2368
2369 assert!(t.set_format(PixelFormat::Rgba).is_err());
2371
2372 assert_eq!(t.format(), Some(PixelFormat::Rgb));
2374 }
2375
2376 #[test]
2377 fn set_format_idempotent() {
2378 let mut t = TensorDyn::new(&[480, 640, 3], DType::U8, None, None).unwrap();
2379 t.set_format(PixelFormat::Rgb).unwrap();
2380 t.set_format(PixelFormat::Rgb).unwrap();
2381 assert_eq!(t.format(), Some(PixelFormat::Rgb));
2382 assert_eq!(t.width(), Some(640));
2383 assert_eq!(t.height(), Some(480));
2384 }
2385
2386 #[test]
2389 fn set_row_stride_valid() {
2390 let mut t = TensorDyn::image(
2392 100,
2393 100,
2394 PixelFormat::Rgba,
2395 DType::U8,
2396 None,
2397 crate::CpuAccess::ReadWrite,
2398 )
2399 .unwrap();
2400 t.set_row_stride(512).unwrap();
2401 assert_eq!(t.row_stride(), Some(512));
2402 assert_eq!(t.effective_row_stride(), Some(512));
2403 }
2404
2405 #[test]
2406 fn set_row_stride_equals_min() {
2407 let mut t = TensorDyn::image(
2409 100,
2410 100,
2411 PixelFormat::Rgb,
2412 DType::U8,
2413 None,
2414 crate::CpuAccess::ReadWrite,
2415 )
2416 .unwrap();
2417 t.set_row_stride(300).unwrap();
2418 assert_eq!(t.row_stride(), Some(300));
2419 }
2420
2421 #[test]
2422 fn set_row_stride_too_small() {
2423 let mut t = TensorDyn::image(
2428 64,
2429 100,
2430 PixelFormat::Rgba,
2431 DType::U8,
2432 None,
2433 crate::CpuAccess::ReadWrite,
2434 )
2435 .unwrap();
2436 assert!(t.set_row_stride(200).is_err());
2437 assert_eq!(t.row_stride(), None);
2438 }
2439
2440 #[test]
2441 fn set_row_stride_zero() {
2442 let mut t = TensorDyn::image(
2443 100,
2444 100,
2445 PixelFormat::Rgb,
2446 DType::U8,
2447 None,
2448 crate::CpuAccess::ReadWrite,
2449 )
2450 .unwrap();
2451 assert!(t.set_row_stride(0).is_err());
2452 }
2453
2454 #[test]
2455 fn set_row_stride_requires_format() {
2456 let mut t = TensorDyn::new(&[480, 640, 3], DType::U8, None, None).unwrap();
2457 assert!(t.set_row_stride(2048).is_err());
2458 }
2459
2460 #[test]
2461 fn effective_row_stride_without_stride() {
2462 let t = TensorDyn::image(
2467 64,
2468 100,
2469 PixelFormat::Rgb,
2470 DType::U8,
2471 None,
2472 crate::CpuAccess::ReadWrite,
2473 )
2474 .unwrap();
2475 assert_eq!(t.row_stride(), None);
2476 assert_eq!(t.effective_row_stride(), Some(192)); }
2478
2479 #[test]
2480 fn effective_row_stride_padded_packed_dma() {
2481 let t = match TensorDyn::image(
2487 100,
2488 100,
2489 PixelFormat::Rgb,
2490 DType::U8,
2491 Some(TensorMemory::DmaBuf),
2492 crate::CpuAccess::ReadWrite,
2493 ) {
2494 Ok(t) if t.memory() == TensorMemory::DmaBuf => t,
2495 _ => return,
2496 };
2497 if cfg!(target_os = "android") {
2503 assert!(
2504 t.row_stride().is_none() || t.row_stride() >= Some(300),
2505 "AHardwareBuffer pitch must be absent or at least the tight row"
2506 );
2507 return;
2508 }
2509 if cfg!(target_os = "windows") {
2514 assert!(
2515 t.row_stride().is_none() || t.row_stride() >= Some(300),
2516 "D3D11 staging pitch must be absent or at least the tight row"
2517 );
2518 assert!(t.effective_row_stride() >= Some(300));
2519 return;
2520 }
2521 assert_eq!(t.row_stride(), Some(320));
2522 assert_eq!(t.effective_row_stride(), Some(320));
2523 }
2524
2525 #[test]
2526 fn effective_row_stride_no_format() {
2527 let t = TensorDyn::new(&[480, 640, 3], DType::U8, None, None).unwrap();
2528 assert_eq!(t.effective_row_stride(), None);
2529 }
2530
2531 #[test]
2532 fn with_row_stride_builder() {
2533 let t = TensorDyn::image(
2534 100,
2535 100,
2536 PixelFormat::Rgba,
2537 DType::U8,
2538 None,
2539 crate::CpuAccess::ReadWrite,
2540 )
2541 .unwrap()
2542 .with_row_stride(512)
2543 .unwrap();
2544 assert_eq!(t.row_stride(), Some(512));
2545 assert_eq!(t.effective_row_stride(), Some(512));
2546 }
2547
2548 #[test]
2549 fn with_row_stride_rejects_small() {
2550 let result = TensorDyn::image(
2551 100,
2552 100,
2553 PixelFormat::Rgba,
2554 DType::U8,
2555 None,
2556 crate::CpuAccess::ReadWrite,
2557 )
2558 .unwrap()
2559 .with_row_stride(200);
2560 assert!(result.is_err());
2561 }
2562
2563 #[test]
2564 fn set_format_clears_row_stride() {
2565 let mut t = TensorDyn::new(&[480, 640, 3], DType::U8, None, None).unwrap();
2566 t.set_format(PixelFormat::Rgb).unwrap();
2567 t.set_row_stride(2048).unwrap();
2568 assert_eq!(t.row_stride(), Some(2048));
2569
2570 let _ = t.set_format(PixelFormat::Bgra);
2572 assert_eq!(t.row_stride(), Some(2048));
2573
2574 t.set_format(PixelFormat::Rgb).unwrap();
2576 assert_eq!(t.row_stride(), Some(2048));
2577
2578 t.reshape(&[480 * 640 * 3]).unwrap();
2580 assert_eq!(t.row_stride(), None);
2581 assert_eq!(t.format(), None);
2582 }
2583
2584 #[test]
2585 fn set_format_different_compatible_clears_stride() {
2586 let mut t = TensorDyn::new(&[480, 640, 4], DType::U8, None, None).unwrap();
2589 t.set_format(PixelFormat::Rgba).unwrap();
2590 t.set_row_stride(4096).unwrap();
2591 assert_eq!(t.row_stride(), Some(4096));
2592
2593 t.set_format(PixelFormat::Bgra).unwrap();
2595 assert_eq!(t.format(), Some(PixelFormat::Bgra));
2596 assert_eq!(t.row_stride(), None);
2597 }
2598
2599 #[test]
2600 fn set_format_same_preserves_stride() {
2601 let mut t = TensorDyn::image(
2602 100,
2603 100,
2604 PixelFormat::Rgb,
2605 DType::U8,
2606 None,
2607 crate::CpuAccess::ReadWrite,
2608 )
2609 .unwrap();
2610 t.set_row_stride(512).unwrap();
2611 t.set_format(PixelFormat::Rgb).unwrap();
2613 assert_eq!(t.row_stride(), Some(512));
2614 }
2615
2616 #[test]
2617 fn effective_row_stride_planar() {
2618 let t = TensorDyn::image(
2619 640,
2620 480,
2621 PixelFormat::PlanarRgb,
2622 DType::U8,
2623 None,
2624 crate::CpuAccess::ReadWrite,
2625 )
2626 .unwrap();
2627 assert_eq!(t.effective_row_stride(), Some(640)); }
2629
2630 #[test]
2631 fn effective_row_stride_nv12() {
2632 let t = TensorDyn::image(
2633 640,
2634 480,
2635 PixelFormat::Nv12,
2636 DType::U8,
2637 None,
2638 crate::CpuAccess::ReadWrite,
2639 )
2640 .unwrap();
2641 assert_eq!(t.effective_row_stride(), Some(640)); }
2643
2644 #[test]
2645 fn map_rejects_strided_tensor() {
2646 let mut t = Tensor::<u8>::image(
2647 100,
2648 100,
2649 PixelFormat::Rgba,
2650 Some(TensorMemory::Mem),
2651 crate::CpuAccess::ReadWrite,
2652 )
2653 .unwrap();
2654 assert!(t.map().is_ok());
2656 t.set_row_stride(512).unwrap();
2658 let err = t.map();
2659 assert!(err.is_err());
2660 }
2661
2662 #[test]
2665 fn plane_offset_default_none() {
2666 let t = TensorDyn::image(
2667 100,
2668 100,
2669 PixelFormat::Rgba,
2670 DType::U8,
2671 None,
2672 crate::CpuAccess::ReadWrite,
2673 )
2674 .unwrap();
2675 assert_eq!(t.plane_offset(), None);
2676 }
2677
2678 #[test]
2679 fn set_plane_offset_basic() {
2680 let mut t = TensorDyn::image(
2681 100,
2682 100,
2683 PixelFormat::Rgba,
2684 DType::U8,
2685 None,
2686 crate::CpuAccess::ReadWrite,
2687 )
2688 .unwrap();
2689 t.set_plane_offset(4096);
2690 assert_eq!(t.plane_offset(), Some(4096));
2691 }
2692
2693 #[test]
2694 fn set_plane_offset_zero() {
2695 let mut t = TensorDyn::image(
2696 100,
2697 100,
2698 PixelFormat::Rgb,
2699 DType::U8,
2700 None,
2701 crate::CpuAccess::ReadWrite,
2702 )
2703 .unwrap();
2704 t.set_plane_offset(0);
2705 assert_eq!(t.plane_offset(), Some(0));
2706 }
2707
2708 #[test]
2709 fn set_plane_offset_no_format() {
2710 let mut t = TensorDyn::new(&[480, 640, 3], DType::U8, None, None).unwrap();
2712 t.set_plane_offset(4096);
2713 assert_eq!(t.plane_offset(), Some(4096));
2714 }
2715
2716 #[test]
2717 fn set_format_clears_plane_offset() {
2718 let mut t = TensorDyn::new(&[480, 640, 3], DType::U8, None, None).unwrap();
2719 t.set_format(PixelFormat::Rgb).unwrap();
2720 t.set_plane_offset(4096);
2721 assert_eq!(t.plane_offset(), Some(4096));
2722
2723 t.set_format(PixelFormat::Rgb).unwrap();
2725 assert_eq!(t.plane_offset(), Some(4096));
2726
2727 t.reshape(&[480 * 640 * 3]).unwrap();
2729 assert_eq!(t.plane_offset(), None);
2730 assert_eq!(t.format(), None);
2731 }
2732
2733 #[test]
2734 fn map_rejects_out_of_bounds_offset() {
2735 let mut t = Tensor::<u8>::image(
2736 100,
2737 100,
2738 PixelFormat::Rgba,
2739 Some(TensorMemory::Mem),
2740 crate::CpuAccess::ReadWrite,
2741 )
2742 .unwrap();
2743 assert!(t.map().is_ok());
2745 t.set_plane_offset(4096);
2748 assert!(t.map().is_err());
2749 }
2750
2751 #[test]
2752 fn mem_subview_in_bounds_maps_at_offset() {
2753 let parent = Tensor::<u8>::image(
2756 100,
2757 100,
2758 PixelFormat::Rgba,
2759 Some(TensorMemory::Mem),
2760 crate::CpuAccess::ReadWrite,
2761 )
2762 .unwrap();
2763 let view = parent.subview(4096, &[10, 10, 4]).unwrap();
2765 assert_eq!(view.plane_offset(), Some(4096));
2766 assert!(view.map().is_ok());
2767 }
2768
2769 #[test]
2770 fn dyn_batch_dispatches_every_dtype() {
2771 use DType::*;
2775 for dt in [U8, I8, U16, I16, U32, I32, U64, I64, F16, F32, F64] {
2776 let parent = TensorDyn::new(&[2, 4], dt, Some(TensorMemory::Mem), None).unwrap();
2777 let view = parent.batch(1).unwrap();
2778 assert_eq!(view.dtype(), dt, "batch must preserve dtype {dt:?}");
2779 assert_eq!(view.shape(), &[4], "{dt:?}");
2780 }
2781 }
2782
2783 #[test]
2784 fn map_accepts_zero_offset_tensor() {
2785 let mut t = Tensor::<u8>::image(
2786 100,
2787 100,
2788 PixelFormat::Rgba,
2789 Some(TensorMemory::Mem),
2790 crate::CpuAccess::ReadWrite,
2791 )
2792 .unwrap();
2793 t.set_plane_offset(0);
2794 assert!(t.map().is_ok());
2796 }
2797
2798 #[test]
2799 fn dyn_configure_image_nv12() {
2800 let mut t = TensorDyn::image(
2801 640,
2802 480,
2803 PixelFormat::Rgb,
2804 DType::U8,
2805 None,
2806 crate::CpuAccess::ReadWrite,
2807 )
2808 .unwrap();
2809 t.configure_image(320, 240, PixelFormat::Nv12).unwrap();
2810 assert_eq!(t.format(), Some(PixelFormat::Nv12));
2811 assert_eq!((t.width(), t.height()), (Some(320), Some(240)));
2812 }
2813
2814 #[test]
2815 fn tensordyn_colorimetry_roundtrip() {
2816 use crate::{ColorEncoding, Colorimetry, DType, PixelFormat};
2817 let mut t = TensorDyn::image(
2818 1280,
2819 720,
2820 PixelFormat::Nv12,
2821 DType::U8,
2822 None,
2823 crate::CpuAccess::ReadWrite,
2824 )
2825 .unwrap();
2826 assert_eq!(t.colorimetry(), None);
2827 let c = Colorimetry::default().with_encoding(ColorEncoding::Bt709);
2828 t.set_colorimetry(Some(c));
2829 assert_eq!(t.colorimetry(), Some(c));
2830 }
2831
2832 #[test]
2833 fn from_planes_propagates_plane_offset() {
2834 let mut luma =
2835 Tensor::<u8>::new(&[480, 640], Some(TensorMemory::Mem), Some("luma")).unwrap();
2836 luma.set_plane_offset(4096);
2837 let chroma =
2838 Tensor::<u8>::new(&[240, 640], Some(TensorMemory::Mem), Some("chroma")).unwrap();
2839 let combined = Tensor::<u8>::from_planes(luma, chroma, PixelFormat::Nv12).unwrap();
2840 assert_eq!(combined.plane_offset(), Some(4096));
2841 }
2842
2843 #[test]
2844 fn cuda_passthrough_none_for_mem_tensor() {
2845 let t: TensorDyn = Tensor::<f32>::new(&[10], Some(TensorMemory::Mem), None)
2848 .unwrap()
2849 .into();
2850 assert!(t.cuda().is_none());
2851 assert!(t.cuda_map().is_none());
2852 }
2853}