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 image_with_stride(
1850 width: usize,
1851 height: usize,
1852 format: PixelFormat,
1853 dtype: DType,
1854 row_stride_bytes: usize,
1855 memory: Option<TensorMemory>,
1856 access: crate::CpuAccess,
1857 ) -> crate::Result<Self> {
1858 match dtype {
1859 DType::U8 => Tensor::<u8>::image_with_stride(
1860 width,
1861 height,
1862 format,
1863 row_stride_bytes,
1864 memory,
1865 access,
1866 )
1867 .map(Self::U8),
1868 DType::I8 => Tensor::<i8>::image_with_stride(
1869 width,
1870 height,
1871 format,
1872 row_stride_bytes,
1873 memory,
1874 access,
1875 )
1876 .map(Self::I8),
1877 DType::U16 => Tensor::<u16>::image_with_stride(
1878 width,
1879 height,
1880 format,
1881 row_stride_bytes,
1882 memory,
1883 access,
1884 )
1885 .map(Self::U16),
1886 DType::I16 => Tensor::<i16>::image_with_stride(
1887 width,
1888 height,
1889 format,
1890 row_stride_bytes,
1891 memory,
1892 access,
1893 )
1894 .map(Self::I16),
1895 DType::U32 => Tensor::<u32>::image_with_stride(
1896 width,
1897 height,
1898 format,
1899 row_stride_bytes,
1900 memory,
1901 access,
1902 )
1903 .map(Self::U32),
1904 DType::I32 => Tensor::<i32>::image_with_stride(
1905 width,
1906 height,
1907 format,
1908 row_stride_bytes,
1909 memory,
1910 access,
1911 )
1912 .map(Self::I32),
1913 DType::U64 => Tensor::<u64>::image_with_stride(
1914 width,
1915 height,
1916 format,
1917 row_stride_bytes,
1918 memory,
1919 access,
1920 )
1921 .map(Self::U64),
1922 DType::I64 => Tensor::<i64>::image_with_stride(
1923 width,
1924 height,
1925 format,
1926 row_stride_bytes,
1927 memory,
1928 access,
1929 )
1930 .map(Self::I64),
1931 DType::F16 => Tensor::<f16>::image_with_stride(
1932 width,
1933 height,
1934 format,
1935 row_stride_bytes,
1936 memory,
1937 access,
1938 )
1939 .map(Self::F16),
1940 DType::F32 => Tensor::<f32>::image_with_stride(
1941 width,
1942 height,
1943 format,
1944 row_stride_bytes,
1945 memory,
1946 access,
1947 )
1948 .map(Self::F32),
1949 DType::F64 => Tensor::<f64>::image_with_stride(
1950 width,
1951 height,
1952 format,
1953 row_stride_bytes,
1954 memory,
1955 access,
1956 )
1957 .map(Self::F64),
1958 }
1959 }
1960}
1961
1962impl From<Tensor<u8>> for TensorDyn {
1965 fn from(t: Tensor<u8>) -> Self {
1966 Self::U8(t)
1967 }
1968}
1969
1970impl From<Tensor<i8>> for TensorDyn {
1971 fn from(t: Tensor<i8>) -> Self {
1972 Self::I8(t)
1973 }
1974}
1975
1976impl From<Tensor<u16>> for TensorDyn {
1977 fn from(t: Tensor<u16>) -> Self {
1978 Self::U16(t)
1979 }
1980}
1981
1982impl From<Tensor<i16>> for TensorDyn {
1983 fn from(t: Tensor<i16>) -> Self {
1984 Self::I16(t)
1985 }
1986}
1987
1988impl From<Tensor<u32>> for TensorDyn {
1989 fn from(t: Tensor<u32>) -> Self {
1990 Self::U32(t)
1991 }
1992}
1993
1994impl From<Tensor<i32>> for TensorDyn {
1995 fn from(t: Tensor<i32>) -> Self {
1996 Self::I32(t)
1997 }
1998}
1999
2000impl From<Tensor<u64>> for TensorDyn {
2001 fn from(t: Tensor<u64>) -> Self {
2002 Self::U64(t)
2003 }
2004}
2005
2006impl From<Tensor<i64>> for TensorDyn {
2007 fn from(t: Tensor<i64>) -> Self {
2008 Self::I64(t)
2009 }
2010}
2011
2012impl From<Tensor<f16>> for TensorDyn {
2013 fn from(t: Tensor<f16>) -> Self {
2014 Self::F16(t)
2015 }
2016}
2017
2018impl From<Tensor<f32>> for TensorDyn {
2019 fn from(t: Tensor<f32>) -> Self {
2020 Self::F32(t)
2021 }
2022}
2023
2024impl From<Tensor<f64>> for TensorDyn {
2025 fn from(t: Tensor<f64>) -> Self {
2026 Self::F64(t)
2027 }
2028}
2029
2030impl fmt::Debug for TensorDyn {
2031 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
2032 dispatch!(self, fmt, f)
2033 }
2034}
2035
2036pub type Raw = *mut std::ffi::c_void;
2039
2040impl TensorDyn {
2041 pub fn into_raw(self) -> Raw {
2044 Box::into_raw(Box::new(self)) as Raw
2045 }
2046
2047 pub unsafe fn from_raw(p: Raw) -> TensorDyn {
2053 *unsafe { Box::from_raw(p as *mut TensorDyn) }
2054 }
2055
2056 pub unsafe fn with_raw<R>(p: Raw, f: impl FnOnce(&mut TensorDyn) -> R) -> R {
2062 let mut t = std::mem::ManuallyDrop::new(unsafe { std::ptr::read(p as *mut TensorDyn) });
2063 f(&mut t)
2064 }
2065}
2066
2067#[cfg(test)]
2068mod tests {
2069 use super::*;
2070
2071 #[test]
2072 fn from_typed_tensor() {
2073 let t = Tensor::<u8>::new(&[10], None, None).unwrap();
2074 let dyn_t: TensorDyn = t.into();
2075 assert_eq!(dyn_t.dtype(), DType::U8);
2076 assert_eq!(dyn_t.shape(), &[10]);
2077 }
2078
2079 #[test]
2080 fn from_foreign_ptr_wraps_borrowed_memory() {
2081 use crate::TensorMapTrait;
2082 let mut vec: Vec<f32> = vec![0.0; 4];
2085 let ptr = vec.as_mut_ptr() as *mut u8;
2086 let owner: crate::ForeignOwner = Box::new(vec);
2087 let t = unsafe {
2088 TensorDyn::from_foreign_ptr(ptr, &[2, 2], DType::F32, Some(owner), Some("trt_output"))
2089 }
2090 .unwrap();
2091 assert_eq!(t.dtype(), DType::F32);
2092 assert_eq!(t.memory(), TensorMemory::Mem);
2093 assert_eq!(t.shape(), &[2, 2]);
2094 {
2095 let mut m = t.as_f32().unwrap().map().unwrap();
2096 m.as_mut_slice().copy_from_slice(&[1.0, 2.0, 3.0, 4.0]);
2097 }
2098 let m = t.as_f32().unwrap().map().unwrap();
2099 assert_eq!(m.as_slice(), &[1.0, 2.0, 3.0, 4.0]);
2100 }
2101
2102 #[test]
2111 fn from_foreign_ptr_rejects_null_ptr() {
2112 let err = unsafe {
2113 TensorDyn::from_foreign_ptr(std::ptr::null_mut(), &[4], DType::U8, None, None)
2114 }
2115 .unwrap_err();
2116 assert!(
2118 matches!(err, crate::error::Error::InvalidArgument(ref m) if m.contains("non-null")),
2119 "expected InvalidArgument(non-null), got {err:?}"
2120 );
2121 }
2122
2123 #[test]
2124 fn from_foreign_ptr_rejects_empty_shape() {
2125 let mut dummy: u8 = 0;
2126 let err = unsafe {
2127 TensorDyn::from_foreign_ptr(&mut dummy as *mut u8, &[], DType::U8, None, None)
2128 }
2129 .unwrap_err();
2130 assert!(
2131 matches!(err, crate::error::Error::InvalidSize(0)),
2132 "expected InvalidSize(0) for empty shape, got {err:?}"
2133 );
2134 }
2135
2136 #[test]
2137 fn from_foreign_ptr_rejects_overflow_shape() {
2138 let mut dummy: u8 = 0;
2139 let huge = [usize::MAX / 2 + 1, 2];
2140 let err = unsafe { TensorDyn::from_foreign_ptr(&mut dummy, &huge, DType::U8, None, None) }
2141 .unwrap_err();
2142 assert!(
2143 matches!(err, crate::error::Error::InvalidArgument(ref m) if m.contains("overflow")),
2144 "expected InvalidArgument(overflow), got {err:?}"
2145 );
2146 }
2147
2148 #[test]
2149 fn from_foreign_ptr_u8_dtype_dispatch() {
2150 let mut buf: Vec<u8> = vec![1, 2, 3, 4];
2153 let ptr = buf.as_mut_ptr();
2154 let owner: crate::ForeignOwner = Box::new(buf);
2155 let t = unsafe {
2156 TensorDyn::from_foreign_ptr(ptr, &[4], DType::U8, Some(owner), Some("u8_foreign"))
2157 }
2158 .unwrap();
2159 assert_eq!(t.dtype(), DType::U8);
2160 assert_eq!(t.shape(), &[4]);
2161 let m = t.as_u8().unwrap().map().unwrap();
2162 use crate::TensorMapTrait;
2163 assert_eq!(m.as_slice(), &[1u8, 2, 3, 4]);
2164 }
2165
2166 #[test]
2167 fn map_bytes_roundtrips_and_is_typed_agnostic() {
2168 use crate::TensorMapTrait;
2169 let t = Tensor::<f32>::new(&[4], None, None).unwrap();
2172 let dyn_t: TensorDyn = t.into();
2173 {
2174 let mut view = dyn_t.map_bytes(crate::CpuAccess::ReadWrite).unwrap();
2175 assert_eq!(view.len(), 16, "4 x f32 is 16 bytes, not 4");
2176 view.as_mut_slice()
2177 .copy_from_slice(&[1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16]);
2178 }
2179 let view = dyn_t.map_bytes(crate::CpuAccess::Read).unwrap();
2180 assert_eq!(
2181 view.as_slice(),
2182 &[1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16]
2183 );
2184 }
2185
2186 #[test]
2187 fn map_bytes_none_is_invalid() {
2188 let t = Tensor::<u8>::new(&[4], None, None).unwrap();
2189 let dyn_t: TensorDyn = t.into();
2190 assert!(matches!(
2191 dyn_t.map_bytes(crate::CpuAccess::None),
2192 Err(crate::error::Error::InvalidArgument(_))
2193 ));
2194 }
2195
2196 #[test]
2197 fn downcast_ref() {
2198 let t = Tensor::<u8>::new(&[10], None, None).unwrap();
2199 let dyn_t: TensorDyn = t.into();
2200 assert!(dyn_t.as_u8().is_some());
2201 assert!(dyn_t.as_i8().is_none());
2202 }
2203
2204 #[test]
2205 fn downcast_into() {
2206 let t = Tensor::<u8>::new(&[10], None, None).unwrap();
2207 let dyn_t: TensorDyn = t.into();
2208 let back = dyn_t.into_u8().unwrap();
2209 assert_eq!(back.shape(), &[10]);
2210 }
2211
2212 #[test]
2213 fn image_accessors() {
2214 let t = Tensor::<u8>::image(
2215 640,
2216 480,
2217 PixelFormat::Rgba,
2218 None,
2219 crate::CpuAccess::ReadWrite,
2220 )
2221 .unwrap();
2222 let dyn_t: TensorDyn = t.into();
2223 assert_eq!(dyn_t.format(), Some(PixelFormat::Rgba));
2224 assert_eq!(dyn_t.width(), Some(640));
2225 assert_eq!(dyn_t.height(), Some(480));
2226 assert!(!dyn_t.is_multiplane());
2227 }
2228
2229 #[test]
2230 fn cpu_access_reports_the_declared_access() {
2231 let read_t =
2233 Tensor::<u8>::image(4, 4, PixelFormat::Rgba, None, crate::CpuAccess::Read).unwrap();
2234 let dyn_read: TensorDyn = read_t.into();
2235 assert_eq!(dyn_read.cpu_access(), crate::CpuAccess::Read);
2236 assert!(!dyn_read.cpu_access().writes());
2237
2238 let mem_t = Tensor::<u8>::new(&[4], None, None).unwrap();
2242 let dyn_mem: TensorDyn = mem_t.into();
2243 assert_eq!(dyn_mem.cpu_access(), crate::CpuAccess::ReadWrite);
2244 assert!(dyn_mem.cpu_access().writes());
2245 }
2246
2247 #[test]
2248 fn image_constructor() {
2249 let dyn_t = TensorDyn::image(
2250 640,
2251 480,
2252 PixelFormat::Rgb,
2253 DType::U8,
2254 None,
2255 crate::CpuAccess::ReadWrite,
2256 )
2257 .unwrap();
2258 assert_eq!(dyn_t.dtype(), DType::U8);
2259 assert_eq!(dyn_t.format(), Some(PixelFormat::Rgb));
2260 assert_eq!(dyn_t.width(), Some(640));
2261 }
2262
2263 #[test]
2264 fn image_constructor_i8() {
2265 let dyn_t = TensorDyn::image(
2266 640,
2267 480,
2268 PixelFormat::Rgb,
2269 DType::I8,
2270 None,
2271 crate::CpuAccess::ReadWrite,
2272 )
2273 .unwrap();
2274 assert_eq!(dyn_t.dtype(), DType::I8);
2275 assert_eq!(dyn_t.format(), Some(PixelFormat::Rgb));
2276 }
2277
2278 #[test]
2279 fn set_format_packed() {
2280 let mut t = TensorDyn::new(&[480, 640, 3], DType::U8, None, None).unwrap();
2281 assert_eq!(t.format(), None);
2282 t.set_format(PixelFormat::Rgb).unwrap();
2283 assert_eq!(t.format(), Some(PixelFormat::Rgb));
2284 assert_eq!(t.width(), Some(640));
2285 assert_eq!(t.height(), Some(480));
2286 }
2287
2288 #[test]
2289 fn set_format_planar() {
2290 let mut t = TensorDyn::new(&[3, 480, 640], DType::U8, None, None).unwrap();
2291 t.set_format(PixelFormat::PlanarRgb).unwrap();
2292 assert_eq!(t.format(), Some(PixelFormat::PlanarRgb));
2293 assert_eq!(t.width(), Some(640));
2294 assert_eq!(t.height(), Some(480));
2295 }
2296
2297 #[test]
2298 fn set_format_rejects_wrong_shape() {
2299 let mut t = TensorDyn::new(&[480, 640, 4], DType::U8, None, None).unwrap();
2300 assert!(t.set_format(PixelFormat::Rgb).is_err());
2301 }
2302
2303 #[test]
2304 fn with_format_builder() {
2305 let t = TensorDyn::new(&[480, 640, 4], DType::U8, None, None)
2306 .unwrap()
2307 .with_format(PixelFormat::Rgba)
2308 .unwrap();
2309 assert_eq!(t.format(), Some(PixelFormat::Rgba));
2310 assert_eq!(t.width(), Some(640));
2311 assert_eq!(t.height(), Some(480));
2312 }
2313
2314 #[cfg(target_os = "linux")]
2315 #[test]
2316 fn dmabuf_clone_mem_tensor_fails() {
2317 let t = TensorDyn::new(&[480, 640, 3], DType::U8, Some(TensorMemory::Mem), None).unwrap();
2318 assert_eq!(t.memory(), TensorMemory::Mem);
2319 assert!(t.dmabuf_clone().is_err());
2320 }
2321
2322 #[cfg(target_os = "linux")]
2323 #[test]
2324 fn dmabuf_mem_tensor_fails() {
2325 let t = TensorDyn::new(&[480, 640, 3], DType::U8, Some(TensorMemory::Mem), None).unwrap();
2326 assert!(t.dmabuf().is_err());
2327 }
2328
2329 #[test]
2330 fn set_format_semi_planar_nv12() {
2331 let mut t = TensorDyn::new(&[720, 640], DType::U8, Some(TensorMemory::Mem), None).unwrap();
2333 t.set_format(PixelFormat::Nv12).unwrap();
2334 assert_eq!(t.format(), Some(PixelFormat::Nv12));
2335 assert_eq!(t.width(), Some(640));
2336 assert_eq!(t.height(), Some(480));
2337 }
2338
2339 #[test]
2340 fn set_format_semi_planar_nv16() {
2341 let mut t = TensorDyn::new(&[960, 640], DType::U8, Some(TensorMemory::Mem), None).unwrap();
2343 t.set_format(PixelFormat::Nv16).unwrap();
2344 assert_eq!(t.format(), Some(PixelFormat::Nv16));
2345 assert_eq!(t.width(), Some(640));
2346 assert_eq!(t.height(), Some(480));
2347 }
2348
2349 #[test]
2350 fn with_format_rejects_wrong_shape() {
2351 let result = TensorDyn::new(&[480, 640, 4], DType::U8, None, None)
2352 .unwrap()
2353 .with_format(PixelFormat::Rgb);
2354 assert!(result.is_err());
2355 }
2356
2357 #[test]
2358 fn set_format_preserved_after_rejection() {
2359 let mut t = TensorDyn::new(&[480, 640, 3], DType::U8, None, None).unwrap();
2360 t.set_format(PixelFormat::Rgb).unwrap();
2361 assert_eq!(t.format(), Some(PixelFormat::Rgb));
2362
2363 assert!(t.set_format(PixelFormat::Rgba).is_err());
2365
2366 assert_eq!(t.format(), Some(PixelFormat::Rgb));
2368 }
2369
2370 #[test]
2371 fn set_format_idempotent() {
2372 let mut t = TensorDyn::new(&[480, 640, 3], DType::U8, None, None).unwrap();
2373 t.set_format(PixelFormat::Rgb).unwrap();
2374 t.set_format(PixelFormat::Rgb).unwrap();
2375 assert_eq!(t.format(), Some(PixelFormat::Rgb));
2376 assert_eq!(t.width(), Some(640));
2377 assert_eq!(t.height(), Some(480));
2378 }
2379
2380 #[test]
2383 fn set_row_stride_valid() {
2384 let mut t = TensorDyn::image(
2386 100,
2387 100,
2388 PixelFormat::Rgba,
2389 DType::U8,
2390 None,
2391 crate::CpuAccess::ReadWrite,
2392 )
2393 .unwrap();
2394 t.set_row_stride(512).unwrap();
2395 assert_eq!(t.row_stride(), Some(512));
2396 assert_eq!(t.effective_row_stride(), Some(512));
2397 }
2398
2399 #[test]
2400 fn set_row_stride_equals_min() {
2401 let mut t = TensorDyn::image(
2403 100,
2404 100,
2405 PixelFormat::Rgb,
2406 DType::U8,
2407 None,
2408 crate::CpuAccess::ReadWrite,
2409 )
2410 .unwrap();
2411 t.set_row_stride(300).unwrap();
2412 assert_eq!(t.row_stride(), Some(300));
2413 }
2414
2415 #[test]
2416 fn set_row_stride_too_small() {
2417 let mut t = TensorDyn::image(
2422 64,
2423 100,
2424 PixelFormat::Rgba,
2425 DType::U8,
2426 None,
2427 crate::CpuAccess::ReadWrite,
2428 )
2429 .unwrap();
2430 assert!(t.set_row_stride(200).is_err());
2431 assert_eq!(t.row_stride(), None);
2432 }
2433
2434 #[test]
2435 fn set_row_stride_zero() {
2436 let mut t = TensorDyn::image(
2437 100,
2438 100,
2439 PixelFormat::Rgb,
2440 DType::U8,
2441 None,
2442 crate::CpuAccess::ReadWrite,
2443 )
2444 .unwrap();
2445 assert!(t.set_row_stride(0).is_err());
2446 }
2447
2448 #[test]
2449 fn set_row_stride_requires_format() {
2450 let mut t = TensorDyn::new(&[480, 640, 3], DType::U8, None, None).unwrap();
2451 assert!(t.set_row_stride(2048).is_err());
2452 }
2453
2454 #[test]
2455 fn effective_row_stride_without_stride() {
2456 let t = TensorDyn::image(
2461 64,
2462 100,
2463 PixelFormat::Rgb,
2464 DType::U8,
2465 None,
2466 crate::CpuAccess::ReadWrite,
2467 )
2468 .unwrap();
2469 assert_eq!(t.row_stride(), None);
2470 assert_eq!(t.effective_row_stride(), Some(192)); }
2472
2473 #[test]
2474 fn effective_row_stride_padded_packed_dma() {
2475 let t = match TensorDyn::image(
2481 100,
2482 100,
2483 PixelFormat::Rgb,
2484 DType::U8,
2485 Some(TensorMemory::DmaBuf),
2486 crate::CpuAccess::ReadWrite,
2487 ) {
2488 Ok(t) if t.memory() == TensorMemory::DmaBuf => t,
2489 _ => return,
2490 };
2491 if cfg!(target_os = "android") {
2497 assert!(
2498 t.row_stride().is_none() || t.row_stride() >= Some(300),
2499 "AHardwareBuffer pitch must be absent or at least the tight row"
2500 );
2501 return;
2502 }
2503 if cfg!(target_os = "windows") {
2508 assert!(
2509 t.row_stride().is_none() || t.row_stride() >= Some(300),
2510 "D3D11 staging pitch must be absent or at least the tight row"
2511 );
2512 assert!(t.effective_row_stride() >= Some(300));
2513 return;
2514 }
2515 assert_eq!(t.row_stride(), Some(320));
2516 assert_eq!(t.effective_row_stride(), Some(320));
2517 }
2518
2519 #[test]
2520 fn effective_row_stride_no_format() {
2521 let t = TensorDyn::new(&[480, 640, 3], DType::U8, None, None).unwrap();
2522 assert_eq!(t.effective_row_stride(), None);
2523 }
2524
2525 #[test]
2526 fn with_row_stride_builder() {
2527 let t = TensorDyn::image(
2528 100,
2529 100,
2530 PixelFormat::Rgba,
2531 DType::U8,
2532 None,
2533 crate::CpuAccess::ReadWrite,
2534 )
2535 .unwrap()
2536 .with_row_stride(512)
2537 .unwrap();
2538 assert_eq!(t.row_stride(), Some(512));
2539 assert_eq!(t.effective_row_stride(), Some(512));
2540 }
2541
2542 #[test]
2543 fn with_row_stride_rejects_small() {
2544 let result = TensorDyn::image(
2545 100,
2546 100,
2547 PixelFormat::Rgba,
2548 DType::U8,
2549 None,
2550 crate::CpuAccess::ReadWrite,
2551 )
2552 .unwrap()
2553 .with_row_stride(200);
2554 assert!(result.is_err());
2555 }
2556
2557 #[test]
2558 fn set_format_clears_row_stride() {
2559 let mut t = TensorDyn::new(&[480, 640, 3], DType::U8, None, None).unwrap();
2560 t.set_format(PixelFormat::Rgb).unwrap();
2561 t.set_row_stride(2048).unwrap();
2562 assert_eq!(t.row_stride(), Some(2048));
2563
2564 let _ = t.set_format(PixelFormat::Bgra);
2566 assert_eq!(t.row_stride(), Some(2048));
2567
2568 t.set_format(PixelFormat::Rgb).unwrap();
2570 assert_eq!(t.row_stride(), Some(2048));
2571
2572 t.reshape(&[480 * 640 * 3]).unwrap();
2574 assert_eq!(t.row_stride(), None);
2575 assert_eq!(t.format(), None);
2576 }
2577
2578 #[test]
2579 fn set_format_different_compatible_clears_stride() {
2580 let mut t = TensorDyn::new(&[480, 640, 4], DType::U8, None, None).unwrap();
2583 t.set_format(PixelFormat::Rgba).unwrap();
2584 t.set_row_stride(4096).unwrap();
2585 assert_eq!(t.row_stride(), Some(4096));
2586
2587 t.set_format(PixelFormat::Bgra).unwrap();
2589 assert_eq!(t.format(), Some(PixelFormat::Bgra));
2590 assert_eq!(t.row_stride(), None);
2591 }
2592
2593 #[test]
2594 fn set_format_same_preserves_stride() {
2595 let mut t = TensorDyn::image(
2596 100,
2597 100,
2598 PixelFormat::Rgb,
2599 DType::U8,
2600 None,
2601 crate::CpuAccess::ReadWrite,
2602 )
2603 .unwrap();
2604 t.set_row_stride(512).unwrap();
2605 t.set_format(PixelFormat::Rgb).unwrap();
2607 assert_eq!(t.row_stride(), Some(512));
2608 }
2609
2610 #[test]
2611 fn effective_row_stride_planar() {
2612 let t = TensorDyn::image(
2613 640,
2614 480,
2615 PixelFormat::PlanarRgb,
2616 DType::U8,
2617 None,
2618 crate::CpuAccess::ReadWrite,
2619 )
2620 .unwrap();
2621 assert_eq!(t.effective_row_stride(), Some(640)); }
2623
2624 #[test]
2625 fn effective_row_stride_nv12() {
2626 let t = TensorDyn::image(
2627 640,
2628 480,
2629 PixelFormat::Nv12,
2630 DType::U8,
2631 None,
2632 crate::CpuAccess::ReadWrite,
2633 )
2634 .unwrap();
2635 assert_eq!(t.effective_row_stride(), Some(640)); }
2637
2638 #[test]
2639 fn map_rejects_strided_tensor() {
2640 let mut t = Tensor::<u8>::image(
2641 100,
2642 100,
2643 PixelFormat::Rgba,
2644 Some(TensorMemory::Mem),
2645 crate::CpuAccess::ReadWrite,
2646 )
2647 .unwrap();
2648 assert!(t.map().is_ok());
2650 t.set_row_stride(512).unwrap();
2652 let err = t.map();
2653 assert!(err.is_err());
2654 }
2655
2656 #[test]
2659 fn plane_offset_default_none() {
2660 let t = TensorDyn::image(
2661 100,
2662 100,
2663 PixelFormat::Rgba,
2664 DType::U8,
2665 None,
2666 crate::CpuAccess::ReadWrite,
2667 )
2668 .unwrap();
2669 assert_eq!(t.plane_offset(), None);
2670 }
2671
2672 #[test]
2673 fn set_plane_offset_basic() {
2674 let mut t = TensorDyn::image(
2675 100,
2676 100,
2677 PixelFormat::Rgba,
2678 DType::U8,
2679 None,
2680 crate::CpuAccess::ReadWrite,
2681 )
2682 .unwrap();
2683 t.set_plane_offset(4096);
2684 assert_eq!(t.plane_offset(), Some(4096));
2685 }
2686
2687 #[test]
2688 fn set_plane_offset_zero() {
2689 let mut t = TensorDyn::image(
2690 100,
2691 100,
2692 PixelFormat::Rgb,
2693 DType::U8,
2694 None,
2695 crate::CpuAccess::ReadWrite,
2696 )
2697 .unwrap();
2698 t.set_plane_offset(0);
2699 assert_eq!(t.plane_offset(), Some(0));
2700 }
2701
2702 #[test]
2703 fn set_plane_offset_no_format() {
2704 let mut t = TensorDyn::new(&[480, 640, 3], DType::U8, None, None).unwrap();
2706 t.set_plane_offset(4096);
2707 assert_eq!(t.plane_offset(), Some(4096));
2708 }
2709
2710 #[test]
2711 fn set_format_clears_plane_offset() {
2712 let mut t = TensorDyn::new(&[480, 640, 3], DType::U8, None, None).unwrap();
2713 t.set_format(PixelFormat::Rgb).unwrap();
2714 t.set_plane_offset(4096);
2715 assert_eq!(t.plane_offset(), Some(4096));
2716
2717 t.set_format(PixelFormat::Rgb).unwrap();
2719 assert_eq!(t.plane_offset(), Some(4096));
2720
2721 t.reshape(&[480 * 640 * 3]).unwrap();
2723 assert_eq!(t.plane_offset(), None);
2724 assert_eq!(t.format(), None);
2725 }
2726
2727 #[test]
2728 fn map_rejects_out_of_bounds_offset() {
2729 let mut t = Tensor::<u8>::image(
2730 100,
2731 100,
2732 PixelFormat::Rgba,
2733 Some(TensorMemory::Mem),
2734 crate::CpuAccess::ReadWrite,
2735 )
2736 .unwrap();
2737 assert!(t.map().is_ok());
2739 t.set_plane_offset(4096);
2742 assert!(t.map().is_err());
2743 }
2744
2745 #[test]
2746 fn mem_subview_in_bounds_maps_at_offset() {
2747 let parent = Tensor::<u8>::image(
2750 100,
2751 100,
2752 PixelFormat::Rgba,
2753 Some(TensorMemory::Mem),
2754 crate::CpuAccess::ReadWrite,
2755 )
2756 .unwrap();
2757 let view = parent.subview(4096, &[10, 10, 4]).unwrap();
2759 assert_eq!(view.plane_offset(), Some(4096));
2760 assert!(view.map().is_ok());
2761 }
2762
2763 #[test]
2764 fn dyn_batch_dispatches_every_dtype() {
2765 use DType::*;
2769 for dt in [U8, I8, U16, I16, U32, I32, U64, I64, F16, F32, F64] {
2770 let parent = TensorDyn::new(&[2, 4], dt, Some(TensorMemory::Mem), None).unwrap();
2771 let view = parent.batch(1).unwrap();
2772 assert_eq!(view.dtype(), dt, "batch must preserve dtype {dt:?}");
2773 assert_eq!(view.shape(), &[4], "{dt:?}");
2774 }
2775 }
2776
2777 #[test]
2778 fn map_accepts_zero_offset_tensor() {
2779 let mut t = Tensor::<u8>::image(
2780 100,
2781 100,
2782 PixelFormat::Rgba,
2783 Some(TensorMemory::Mem),
2784 crate::CpuAccess::ReadWrite,
2785 )
2786 .unwrap();
2787 t.set_plane_offset(0);
2788 assert!(t.map().is_ok());
2790 }
2791
2792 #[test]
2793 fn dyn_configure_image_nv12() {
2794 let mut t = TensorDyn::image(
2795 640,
2796 480,
2797 PixelFormat::Rgb,
2798 DType::U8,
2799 None,
2800 crate::CpuAccess::ReadWrite,
2801 )
2802 .unwrap();
2803 t.configure_image(320, 240, PixelFormat::Nv12).unwrap();
2804 assert_eq!(t.format(), Some(PixelFormat::Nv12));
2805 assert_eq!((t.width(), t.height()), (Some(320), Some(240)));
2806 }
2807
2808 #[test]
2809 fn tensordyn_colorimetry_roundtrip() {
2810 use crate::{ColorEncoding, Colorimetry, DType, PixelFormat};
2811 let mut t = TensorDyn::image(
2812 1280,
2813 720,
2814 PixelFormat::Nv12,
2815 DType::U8,
2816 None,
2817 crate::CpuAccess::ReadWrite,
2818 )
2819 .unwrap();
2820 assert_eq!(t.colorimetry(), None);
2821 let c = Colorimetry::default().with_encoding(ColorEncoding::Bt709);
2822 t.set_colorimetry(Some(c));
2823 assert_eq!(t.colorimetry(), Some(c));
2824 }
2825
2826 #[test]
2827 fn from_planes_propagates_plane_offset() {
2828 let mut luma =
2829 Tensor::<u8>::new(&[480, 640], Some(TensorMemory::Mem), Some("luma")).unwrap();
2830 luma.set_plane_offset(4096);
2831 let chroma =
2832 Tensor::<u8>::new(&[240, 640], Some(TensorMemory::Mem), Some("chroma")).unwrap();
2833 let combined = Tensor::<u8>::from_planes(luma, chroma, PixelFormat::Nv12).unwrap();
2834 assert_eq!(combined.plane_offset(), Some(4096));
2835 }
2836
2837 #[test]
2838 fn cuda_passthrough_none_for_mem_tensor() {
2839 let t: TensorDyn = Tensor::<f32>::new(&[10], Some(TensorMemory::Mem), None)
2842 .unwrap()
2843 .into();
2844 assert!(t.cuda().is_none());
2845 assert!(t.cuda_map().is_none());
2846 }
2847}