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 cpu_access: self.cpu_access(),
384 view: self.is_view(),
385 })
386 }
387
388 pub fn shape(&self) -> &[usize] {
389 dispatch!(self, shape)
390 }
391
392 pub fn set_logical_shape(&mut self, shape: &[usize]) -> crate::Result<()> {
396 dispatch!(self, set_logical_shape, shape)
397 }
398
399 pub fn capacity_bytes(&self) -> usize {
401 dispatch!(self, capacity_bytes)
402 }
403
404 pub fn name(&self) -> String {
406 dispatch!(self, name)
407 }
408
409 pub fn format(&self) -> Option<PixelFormat> {
411 dispatch!(self, format)
412 }
413
414 pub fn width(&self) -> Option<usize> {
416 dispatch!(self, width)
417 }
418
419 pub fn height(&self) -> Option<usize> {
421 dispatch!(self, height)
422 }
423
424 pub fn size(&self) -> usize {
426 dispatch!(self, size)
427 }
428
429 pub fn set_dtype(&mut self, dtype: DType) -> crate::Result<()> {
448 let current = self.dtype();
449 if current == dtype {
450 return Ok(());
451 }
452 if current.size() != dtype.size() {
453 return Err(crate::Error::InvalidArgument(format!(
454 "set_dtype: {current:?} and {dtype:?} are different widths \
455 ({} vs {} bytes); retagging would change the element count \
456 over an allocation whose size did not change",
457 current.size(),
458 dtype.size()
459 )));
460 }
461 unsafe {
469 let taken = std::ptr::read(self);
470 std::ptr::write(self, retag_same_width(taken, dtype));
471 }
472 Ok(())
473 }
474
475 pub fn memory(&self) -> TensorMemory {
477 dispatch!(self, memory)
478 }
479
480 pub fn cpu_access(&self) -> crate::CpuAccess {
492 dispatch!(self, cpu_access)
493 }
494
495 pub fn restrict_cpu_access(&mut self, access: crate::CpuAccess) {
497 dispatch!(self, restrict_cpu_access, access)
498 }
499
500 pub fn reshape(&mut self, shape: &[usize]) -> crate::Result<()> {
502 dispatch!(self, reshape, shape)
503 }
504
505 pub fn set_format(&mut self, format: PixelFormat) -> crate::Result<()> {
523 dispatch!(self, set_format, format)
524 }
525
526 pub fn with_format(mut self, format: PixelFormat) -> crate::Result<Self> {
543 self.set_format(format)?;
544 Ok(self)
545 }
546
547 pub fn colorimetry(&self) -> Option<crate::Colorimetry> {
549 dispatch!(self, colorimetry)
550 }
551
552 pub fn set_colorimetry(&mut self, c: Option<crate::Colorimetry>) {
554 dispatch!(self, set_colorimetry, c)
555 }
556
557 pub fn with_colorimetry(mut self, c: crate::Colorimetry) -> Self {
559 self.set_colorimetry(Some(c));
560 self
561 }
562
563 pub fn row_stride(&self) -> Option<usize> {
565 dispatch!(self, row_stride)
566 }
567
568 pub fn effective_row_stride(&self) -> Option<usize> {
570 dispatch!(self, effective_row_stride)
571 }
572
573 pub fn configure_image(
576 &mut self,
577 width: usize,
578 height: usize,
579 format: PixelFormat,
580 ) -> crate::Result<()> {
581 dispatch!(self, configure_image, width, height, format)
582 }
583
584 pub fn set_row_stride(&mut self, stride: usize) -> crate::Result<()> {
590 dispatch!(self, set_row_stride, stride)
591 }
592
593 pub fn with_row_stride(mut self, stride: usize) -> crate::Result<Self> {
595 self.set_row_stride(stride)?;
596 Ok(self)
597 }
598
599 pub fn set_row_stride_unchecked(&mut self, stride: usize) {
605 dispatch!(self, set_row_stride_unchecked, stride)
606 }
607
608 pub fn plane_offset(&self) -> Option<usize> {
610 dispatch!(self, plane_offset)
611 }
612
613 pub fn view_origin(&self) -> Option<crate::ViewOrigin> {
617 dispatch!(self, view_origin)
618 }
619
620 pub fn is_view(&self) -> bool {
622 dispatch!(self, is_view)
623 }
624
625 pub fn mark_view(&mut self) {
627 dispatch!(self, mark_view)
628 }
629
630 pub fn set_plane_offset(&mut self, offset: usize) {
632 dispatch!(self, set_plane_offset, offset)
633 }
634
635 pub fn batch(&self, n: usize) -> crate::Result<TensorDyn> {
638 dyn_fanout!(self, batch, n)
639 }
640
641 pub fn view(&self, region: crate::Region) -> crate::Result<TensorDyn> {
644 dyn_fanout!(self, view, region)
645 }
646
647 pub fn cuda(&self) -> Option<&crate::cuda::CudaHandle> {
652 dispatch!(self, cuda)
653 }
654
655 pub fn cuda_map(&self) -> Option<crate::cuda::CudaMap<'_>> {
678 dispatch!(self, cuda_map)
679 }
680
681 pub fn cuda_map_mut(&self) -> Option<crate::cuda::CudaMap<'_>> {
685 dispatch!(self, cuda_map_mut)
686 }
687
688 pub fn set_cuda_handle(&mut self, h: crate::cuda::CudaHandle) {
690 dispatch!(self, set_cuda_handle, h)
691 }
692
693 pub fn is_cuda_attached(&self) -> bool {
695 dispatch!(self, is_cuda_attached)
696 }
697
698 pub fn quantization(&self) -> Option<&crate::Quantization> {
702 match self {
703 Self::U8(t) => t.quantization(),
704 Self::I8(t) => t.quantization(),
705 Self::U16(t) => t.quantization(),
706 Self::I16(t) => t.quantization(),
707 Self::U32(t) => t.quantization(),
708 Self::I32(t) => t.quantization(),
709 Self::U64(t) => t.quantization(),
710 Self::I64(t) => t.quantization(),
711 Self::F16(_) | Self::F32(_) | Self::F64(_) => None,
712 }
713 }
714
715 pub fn set_quantization(&mut self, q: crate::Quantization) -> crate::Result<()> {
719 match self {
720 Self::U8(t) => t.set_quantization(q),
721 Self::I8(t) => t.set_quantization(q),
722 Self::U16(t) => t.set_quantization(q),
723 Self::I16(t) => t.set_quantization(q),
724 Self::U32(t) => t.set_quantization(q),
725 Self::I32(t) => t.set_quantization(q),
726 Self::U64(t) => t.set_quantization(q),
727 Self::I64(t) => t.set_quantization(q),
728 Self::F16(_) | Self::F32(_) | Self::F64(_) => Err(crate::Error::QuantizationInvalid {
729 field: "dtype_is_integer",
730 expected: "integer tensor dtype (u8/i8/u16/i16/u32/i32/u64/i64)".to_string(),
731 got: format!("{:?}", self.dtype()),
732 }),
733 }
734 }
735
736 pub fn with_quantization(mut self, q: crate::Quantization) -> crate::Result<Self> {
739 self.set_quantization(q)?;
740 Ok(self)
741 }
742
743 pub fn clear_quantization(&mut self) {
745 match self {
746 Self::U8(t) => t.clear_quantization(),
747 Self::I8(t) => t.clear_quantization(),
748 Self::U16(t) => t.clear_quantization(),
749 Self::I16(t) => t.clear_quantization(),
750 Self::U32(t) => t.clear_quantization(),
751 Self::I32(t) => t.clear_quantization(),
752 Self::U64(t) => t.clear_quantization(),
753 Self::I64(t) => t.clear_quantization(),
754 Self::F16(_) | Self::F32(_) | Self::F64(_) => {}
755 }
756 }
757
758 #[cfg(unix)]
760 pub fn clone_fd(&self) -> crate::Result<std::os::fd::OwnedFd> {
761 dispatch!(self, clone_fd)
762 }
763
764 #[cfg(target_os = "linux")]
775 pub fn dmabuf_clone(&self) -> crate::Result<std::os::fd::OwnedFd> {
776 if self.memory() != TensorMemory::DmaBuf {
777 return Err(crate::Error::NotImplemented(format!(
778 "dmabuf_clone requires DMA-backed tensor, got {:?}",
779 self.memory()
780 )));
781 }
782 self.clone_fd()
783 }
784
785 #[cfg(target_os = "linux")]
796 pub fn dmabuf(&self) -> crate::Result<std::os::fd::BorrowedFd<'_>> {
797 dispatch!(self, dmabuf)
798 }
799
800 pub fn is_multiplane(&self) -> bool {
802 dispatch!(self, is_multiplane)
803 }
804
805 pub fn chroma_dyn(&self) -> Option<TensorDyn> {
808 match self {
809 TensorDyn::U8(t) => t
810 .chroma()
811 .and_then(|c| c.clone_chroma_plane().ok())
812 .map(TensorDyn::U8),
813 TensorDyn::I8(t) => t
814 .chroma()
815 .and_then(|c| c.clone_chroma_plane().ok())
816 .map(TensorDyn::I8),
817 TensorDyn::U16(t) => t
818 .chroma()
819 .and_then(|c| c.clone_chroma_plane().ok())
820 .map(TensorDyn::U16),
821 TensorDyn::I16(t) => t
822 .chroma()
823 .and_then(|c| c.clone_chroma_plane().ok())
824 .map(TensorDyn::I16),
825 TensorDyn::U32(t) => t
826 .chroma()
827 .and_then(|c| c.clone_chroma_plane().ok())
828 .map(TensorDyn::U32),
829 TensorDyn::I32(t) => t
830 .chroma()
831 .and_then(|c| c.clone_chroma_plane().ok())
832 .map(TensorDyn::I32),
833 TensorDyn::U64(t) => t
834 .chroma()
835 .and_then(|c| c.clone_chroma_plane().ok())
836 .map(TensorDyn::U64),
837 TensorDyn::I64(t) => t
838 .chroma()
839 .and_then(|c| c.clone_chroma_plane().ok())
840 .map(TensorDyn::I64),
841 TensorDyn::F16(t) => t
842 .chroma()
843 .and_then(|c| c.clone_chroma_plane().ok())
844 .map(TensorDyn::F16),
845 TensorDyn::F32(t) => t
846 .chroma()
847 .and_then(|c| c.clone_chroma_plane().ok())
848 .map(TensorDyn::F32),
849 TensorDyn::F64(t) => t
850 .chroma()
851 .and_then(|c| c.clone_chroma_plane().ok())
852 .map(TensorDyn::F64),
853 }
854 }
855
856 pub fn buffer_identity(&self) -> &crate::BufferIdentity {
866 dispatch!(self, buffer_identity)
867 }
868
869 pub fn aliases(&self, other: &Self) -> bool {
887 if self.buffer_identity().id() == other.buffer_identity().id() {
888 return true;
889 }
890 if self.memory() != other.memory() {
891 return false;
892 }
893 #[cfg(target_os = "linux")]
894 if self.memory() == TensorMemory::DmaBuf {
895 use std::os::fd::AsRawFd;
896 if let (Ok(a), Ok(b)) = (self.dmabuf(), other.dmabuf()) {
897 return a.as_raw_fd() == b.as_raw_fd();
898 }
899 }
900 false
901 }
902
903 downcast_methods!(U8, u8, as_u8, as_u8_mut, into_u8);
906 downcast_methods!(I8, i8, as_i8, as_i8_mut, into_i8);
907 downcast_methods!(U16, u16, as_u16, as_u16_mut, into_u16);
908 downcast_methods!(I16, i16, as_i16, as_i16_mut, into_i16);
909 downcast_methods!(U32, u32, as_u32, as_u32_mut, into_u32);
910 downcast_methods!(I32, i32, as_i32, as_i32_mut, into_i32);
911 downcast_methods!(U64, u64, as_u64, as_u64_mut, into_u64);
912 downcast_methods!(I64, i64, as_i64, as_i64_mut, into_i64);
913 downcast_methods!(F16, f16, as_f16, as_f16_mut, into_f16);
914 downcast_methods!(F32, f32, as_f32, as_f32_mut, into_f32);
915 downcast_methods!(F64, f64, as_f64, as_f64_mut, into_f64);
916
917 pub fn new(
919 shape: &[usize],
920 dtype: DType,
921 memory: Option<TensorMemory>,
922 name: Option<&str>,
923 ) -> crate::Result<Self> {
924 match dtype {
925 DType::U8 => Tensor::<u8>::new(shape, memory, name).map(Self::U8),
926 DType::I8 => Tensor::<i8>::new(shape, memory, name).map(Self::I8),
927 DType::U16 => Tensor::<u16>::new(shape, memory, name).map(Self::U16),
928 DType::I16 => Tensor::<i16>::new(shape, memory, name).map(Self::I16),
929 DType::U32 => Tensor::<u32>::new(shape, memory, name).map(Self::U32),
930 DType::I32 => Tensor::<i32>::new(shape, memory, name).map(Self::I32),
931 DType::U64 => Tensor::<u64>::new(shape, memory, name).map(Self::U64),
932 DType::I64 => Tensor::<i64>::new(shape, memory, name).map(Self::I64),
933 DType::F16 => Tensor::<f16>::new(shape, memory, name).map(Self::F16),
934 DType::F32 => Tensor::<f32>::new(shape, memory, name).map(Self::F32),
935 DType::F64 => Tensor::<f64>::new(shape, memory, name).map(Self::F64),
936 }
937 }
938
939 #[cfg(unix)]
966 pub fn from_fd(
967 fd: std::os::fd::OwnedFd,
968 shape: &[usize],
969 dtype: DType,
970 name: Option<&str>,
971 ) -> crate::Result<Self> {
972 match dtype {
973 DType::U8 => Tensor::<u8>::from_fd(fd, shape, name).map(Self::U8),
974 DType::I8 => Tensor::<i8>::from_fd(fd, shape, name).map(Self::I8),
975 DType::U16 => Tensor::<u16>::from_fd(fd, shape, name).map(Self::U16),
976 DType::I16 => Tensor::<i16>::from_fd(fd, shape, name).map(Self::I16),
977 DType::U32 => Tensor::<u32>::from_fd(fd, shape, name).map(Self::U32),
978 DType::I32 => Tensor::<i32>::from_fd(fd, shape, name).map(Self::I32),
979 DType::U64 => Tensor::<u64>::from_fd(fd, shape, name).map(Self::U64),
980 DType::I64 => Tensor::<i64>::from_fd(fd, shape, name).map(Self::I64),
981 DType::F16 => Tensor::<f16>::from_fd(fd, shape, name).map(Self::F16),
982 DType::F32 => Tensor::<f32>::from_fd(fd, shape, name).map(Self::F32),
983 DType::F64 => Tensor::<f64>::from_fd(fd, shape, name).map(Self::F64),
984 }
985 }
986
987 pub(crate) fn from_pbo_import(
993 buffer_id: u32,
994 size: usize,
995 shape: &[usize],
996 dtype: DType,
997 ops: std::sync::Arc<dyn crate::PboOps>,
998 ) -> crate::Result<Self> {
999 match dtype {
1000 DType::U8 => crate::PboTensor::<u8>::from_pbo(buffer_id, size, shape, None, ops)
1001 .and_then(|p| Tensor::from_pbo(p).map(Self::U8)),
1002 DType::I8 => crate::PboTensor::<i8>::from_pbo(buffer_id, size, shape, None, ops)
1003 .and_then(|p| Tensor::from_pbo(p).map(Self::I8)),
1004 DType::U16 => crate::PboTensor::<u16>::from_pbo(buffer_id, size, shape, None, ops)
1005 .and_then(|p| Tensor::from_pbo(p).map(Self::U16)),
1006 DType::I16 => crate::PboTensor::<i16>::from_pbo(buffer_id, size, shape, None, ops)
1007 .and_then(|p| Tensor::from_pbo(p).map(Self::I16)),
1008 DType::U32 => crate::PboTensor::<u32>::from_pbo(buffer_id, size, shape, None, ops)
1009 .and_then(|p| Tensor::from_pbo(p).map(Self::U32)),
1010 DType::I32 => crate::PboTensor::<i32>::from_pbo(buffer_id, size, shape, None, ops)
1011 .and_then(|p| Tensor::from_pbo(p).map(Self::I32)),
1012 DType::U64 => crate::PboTensor::<u64>::from_pbo(buffer_id, size, shape, None, ops)
1013 .and_then(|p| Tensor::from_pbo(p).map(Self::U64)),
1014 DType::I64 => crate::PboTensor::<i64>::from_pbo(buffer_id, size, shape, None, ops)
1015 .and_then(|p| Tensor::from_pbo(p).map(Self::I64)),
1016 DType::F16 => crate::PboTensor::<f16>::from_pbo(buffer_id, size, shape, None, ops)
1017 .and_then(|p| Tensor::from_pbo(p).map(Self::F16)),
1018 DType::F32 => crate::PboTensor::<f32>::from_pbo(buffer_id, size, shape, None, ops)
1019 .and_then(|p| Tensor::from_pbo(p).map(Self::F32)),
1020 DType::F64 => crate::PboTensor::<f64>::from_pbo(buffer_id, size, shape, None, ops)
1021 .and_then(|p| Tensor::from_pbo(p).map(Self::F64)),
1022 }
1023 }
1024
1025 pub unsafe fn from_client_pbo(
1038 state: crate::EfClientState,
1039 map_fn: crate::EfPboMapFn,
1040 unmap_fn: crate::EfPboUnmapFn,
1041 buffer_id: u32,
1042 size: usize,
1043 shape: &[usize],
1044 dtype: DType,
1045 ) -> crate::Result<TensorDyn> {
1046 let ops = unsafe { crate::client_state_pbo_ops(state, map_fn, unmap_fn) }?;
1048 Self::from_pbo_import(buffer_id, size, shape, dtype, ops)
1049 }
1050
1051 pub unsafe fn from_raw_host(
1069 ptr: *mut u8,
1070 shape: &[usize],
1071 dtype: DType,
1072 name: Option<&str>,
1073 ) -> crate::Result<Self> {
1074 unsafe { Self::from_foreign_ptr(ptr, shape, dtype, None, name) }
1075 }
1076
1077 pub unsafe fn from_raw_host_with_capacity(
1093 ptr: *mut u8,
1094 shape: &[usize],
1095 capacity_bytes: usize,
1096 dtype: DType,
1097 name: Option<&str>,
1098 ) -> crate::Result<Self> {
1099 unsafe {
1100 Self::from_foreign_ptr_with_capacity(ptr, shape, dtype, capacity_bytes, None, name)
1101 }
1102 }
1103
1104 pub(crate) fn import_storage(
1116 desc: &crate::TensorDesc,
1117 shape: &[usize],
1118 dtype: DType,
1119 ) -> crate::Result<Self> {
1120 match desc.kind {
1121 #[cfg(target_os = "linux")]
1122 crate::protocol::kind::DMABUF => {
1123 let owned = crate::protocol::dup_descriptor_fd(desc)?;
1128 Self::from_fd(owned, shape, dtype, None)
1129 }
1130 #[cfg(not(target_os = "linux"))]
1131 crate::protocol::kind::DMABUF => Err(crate::Error::NotImplemented(
1132 "dma-buf import off Linux".into(),
1133 )),
1134 #[cfg(any(target_os = "macos", target_os = "ios"))]
1135 crate::protocol::kind::IOSURFACE => {
1136 let id = crate::protocol::descriptor_surface_id(desc)?;
1137 let surface_ref = crate::iosurface::lookup_by_id(id).ok_or_else(|| {
1138 crate::Error::InvalidArgument(format!(
1139 "IOSurface id {id} is not live (the producer's capsule may \
1140 already have been released)"
1141 ))
1142 })?;
1143 let result = unsafe { Self::from_iosurface(surface_ref, shape, dtype, None) };
1148 crate::iosurface::release(surface_ref);
1149 result
1150 }
1151 #[cfg(not(any(target_os = "macos", target_os = "ios")))]
1152 crate::protocol::kind::IOSURFACE => Err(crate::Error::NotImplemented(
1153 "IOSurface import off Apple platforms".into(),
1154 )),
1155 crate::protocol::kind::MEM | crate::protocol::kind::SHM => {
1156 crate::protocol::check_descriptor_host_ptr(desc)?;
1157 unsafe {
1164 Self::from_raw_host_with_capacity(
1165 desc.ptr.0,
1166 shape,
1167 desc.capacity as usize,
1168 dtype,
1169 None,
1170 )
1171 }
1172 }
1173 crate::protocol::kind::PBO => {
1174 let buffer_id = crate::protocol::descriptor_pbo_buffer_id(desc)?;
1175 let parts = unsafe { crate::read_pbo_vtable_parts(desc.ptr.0 as *const _) }?;
1185 unsafe {
1188 Self::from_client_pbo(
1189 parts.state,
1190 parts.map_fn,
1191 parts.unmap_fn,
1192 buffer_id,
1193 desc.capacity as usize,
1194 shape,
1195 dtype,
1196 )
1197 }
1198 }
1199 #[cfg(target_os = "windows")]
1200 crate::protocol::kind::D3D11_TEXTURE => {
1201 let (tex, completion) = crate::protocol::descriptor_d3d11_handles(desc)?;
1202 unsafe {
1208 let (format, width, height) =
1209 crate::protocol::descriptor_d3d11_geometry(desc, tex, shape)?;
1210 Self::from_d3d11_shared_handle(
1211 tex,
1212 width,
1213 height,
1214 format,
1215 dtype,
1216 crate::protocol::flags::cpu_access(desc.flags)
1217 .unwrap_or(crate::CpuAccess::ReadWrite),
1218 completion,
1219 None,
1220 )
1221 }
1222 }
1223 #[cfg(not(target_os = "windows"))]
1224 crate::protocol::kind::D3D11_TEXTURE => Err(crate::Error::NotImplemented(
1225 "D3D11 texture import off Windows".into(),
1226 )),
1227 k => Err(crate::Error::NotImplemented(format!(
1228 "tensor interop kind {k} cannot be imported by this build"
1229 ))),
1230 }
1231 }
1232
1233 pub unsafe fn from_foreign_ptr(
1246 ptr: *mut u8,
1247 shape: &[usize],
1248 dtype: DType,
1249 owner: Option<crate::ForeignOwner>,
1250 name: Option<&str>,
1251 ) -> crate::Result<Self> {
1252 unsafe {
1257 match dtype {
1258 DType::U8 => {
1259 Tensor::<u8>::from_foreign(ptr.cast(), shape, owner, name).map(Self::U8)
1260 }
1261 DType::I8 => {
1262 Tensor::<i8>::from_foreign(ptr.cast(), shape, owner, name).map(Self::I8)
1263 }
1264 DType::U16 => {
1265 Tensor::<u16>::from_foreign(ptr.cast(), shape, owner, name).map(Self::U16)
1266 }
1267 DType::I16 => {
1268 Tensor::<i16>::from_foreign(ptr.cast(), shape, owner, name).map(Self::I16)
1269 }
1270 DType::U32 => {
1271 Tensor::<u32>::from_foreign(ptr.cast(), shape, owner, name).map(Self::U32)
1272 }
1273 DType::I32 => {
1274 Tensor::<i32>::from_foreign(ptr.cast(), shape, owner, name).map(Self::I32)
1275 }
1276 DType::U64 => {
1277 Tensor::<u64>::from_foreign(ptr.cast(), shape, owner, name).map(Self::U64)
1278 }
1279 DType::I64 => {
1280 Tensor::<i64>::from_foreign(ptr.cast(), shape, owner, name).map(Self::I64)
1281 }
1282 DType::F16 => {
1283 Tensor::<f16>::from_foreign(ptr.cast(), shape, owner, name).map(Self::F16)
1284 }
1285 DType::F32 => {
1286 Tensor::<f32>::from_foreign(ptr.cast(), shape, owner, name).map(Self::F32)
1287 }
1288 DType::F64 => {
1289 Tensor::<f64>::from_foreign(ptr.cast(), shape, owner, name).map(Self::F64)
1290 }
1291 }
1292 }
1293 }
1294
1295 pub unsafe fn from_foreign_ptr_with_capacity(
1306 ptr: *mut u8,
1307 shape: &[usize],
1308 dtype: DType,
1309 capacity_bytes: usize,
1310 owner: Option<crate::ForeignOwner>,
1311 name: Option<&str>,
1312 ) -> crate::Result<Self> {
1313 macro_rules! import {
1314 ($t:ty, $variant:ident) => {
1315 unsafe {
1320 Tensor::<$t>::from_foreign_with_capacity(
1321 ptr.cast(),
1322 shape,
1323 capacity_bytes,
1324 owner,
1325 name,
1326 )
1327 }
1328 .map(Self::$variant)
1329 };
1330 }
1331 match dtype {
1332 DType::U8 => import!(u8, U8),
1333 DType::I8 => import!(i8, I8),
1334 DType::U16 => import!(u16, U16),
1335 DType::I16 => import!(i16, I16),
1336 DType::U32 => import!(u32, U32),
1337 DType::I32 => import!(i32, I32),
1338 DType::U64 => import!(u64, U64),
1339 DType::I64 => import!(i64, I64),
1340 DType::F16 => import!(f16, F16),
1341 DType::F32 => import!(f32, F32),
1342 DType::F64 => import!(f64, F64),
1343 }
1344 }
1345
1346 #[cfg(any(target_os = "macos", target_os = "ios"))]
1357 pub fn from_iosurface_id(
1358 id: u32,
1359 shape: &[usize],
1360 dtype: DType,
1361 name: Option<&str>,
1362 ) -> crate::Result<Self> {
1363 let surface = crate::iosurface::lookup_by_id(id).ok_or_else(|| {
1364 crate::Error::InvalidArgument(format!("no live IOSurface with id {id}"))
1365 })?;
1366 let result = unsafe { Self::from_iosurface(surface, shape, dtype, name) };
1381 crate::iosurface::release(surface);
1382 result
1383 }
1384
1385 #[cfg(any(target_os = "macos", target_os = "ios"))]
1393 pub unsafe fn from_iosurface(
1394 surface_ref: *mut std::ffi::c_void,
1395 shape: &[usize],
1396 dtype: DType,
1397 name: Option<&str>,
1398 ) -> crate::Result<Self> {
1399 unsafe {
1400 match dtype {
1401 DType::U8 => Tensor::<u8>::from_iosurface(surface_ref, shape, name).map(Self::U8),
1402 DType::I8 => Tensor::<i8>::from_iosurface(surface_ref, shape, name).map(Self::I8),
1403 DType::U16 => {
1404 Tensor::<u16>::from_iosurface(surface_ref, shape, name).map(Self::U16)
1405 }
1406 DType::I16 => {
1407 Tensor::<i16>::from_iosurface(surface_ref, shape, name).map(Self::I16)
1408 }
1409 DType::U32 => {
1410 Tensor::<u32>::from_iosurface(surface_ref, shape, name).map(Self::U32)
1411 }
1412 DType::I32 => {
1413 Tensor::<i32>::from_iosurface(surface_ref, shape, name).map(Self::I32)
1414 }
1415 DType::U64 => {
1416 Tensor::<u64>::from_iosurface(surface_ref, shape, name).map(Self::U64)
1417 }
1418 DType::I64 => {
1419 Tensor::<i64>::from_iosurface(surface_ref, shape, name).map(Self::I64)
1420 }
1421 DType::F16 => {
1422 Tensor::<f16>::from_iosurface(surface_ref, shape, name).map(Self::F16)
1423 }
1424 DType::F32 => {
1425 Tensor::<f32>::from_iosurface(surface_ref, shape, name).map(Self::F32)
1426 }
1427 DType::F64 => {
1428 Tensor::<f64>::from_iosurface(surface_ref, shape, name).map(Self::F64)
1429 }
1430 }
1431 }
1432 }
1433
1434 #[cfg(any(target_os = "macos", target_os = "ios"))]
1437 pub fn iosurface_id(&self) -> Option<u32> {
1438 dispatch!(self, iosurface_id)
1439 }
1440
1441 pub fn pbo_id(&self) -> Option<u32> {
1444 dispatch!(self, pbo_id)
1445 }
1446
1447 pub fn pbo_is_mapped(&self) -> Option<bool> {
1451 dispatch!(self, pbo_is_mapped)
1452 }
1453
1454 pub fn pbo_vtable_ptr(&self) -> Option<*const std::ffi::c_void> {
1458 dispatch!(self, pbo_vtable_ptr)
1459 }
1460
1461 pub fn pbo_keepalive(&self) -> Option<std::sync::Arc<dyn Send + Sync>> {
1466 dispatch!(self, pbo_keepalive)
1467 }
1468
1469 #[cfg(any(target_os = "macos", target_os = "ios"))]
1473 pub fn iosurface_ref(&self) -> Option<*mut std::ffi::c_void> {
1474 dispatch!(self, iosurface_ref)
1475 }
1476
1477 #[cfg(any(target_os = "macos", target_os = "ios"))]
1482 pub fn iosurface_physical_dims(&self) -> Option<(usize, usize)> {
1483 dispatch!(self, iosurface_physical_dims)
1484 }
1485
1486 #[cfg(target_os = "windows")]
1489 pub fn d3d11_texture(&self) -> Option<*mut std::ffi::c_void> {
1490 dispatch!(self, d3d11_texture)
1491 }
1492
1493 #[cfg(target_os = "windows")]
1496 pub fn d3d11_layout(&self) -> Option<crate::d3d11_layout::D3d11ImageLayout> {
1497 dispatch!(self, d3d11_layout)
1498 }
1499
1500 #[cfg(target_os = "windows")]
1503 pub fn d3d11_shared_handle(&self) -> crate::Result<std::os::windows::io::OwnedHandle> {
1504 dispatch!(self, d3d11_shared_handle)
1505 }
1506
1507 #[cfg(target_os = "windows")]
1510 pub fn gpu_completion(&self) -> crate::Result<Option<crate::d3d11::GpuCompletion>> {
1511 dispatch!(self, gpu_completion)
1512 }
1513
1514 pub fn gpu_write_value(&self) -> u64 {
1518 dispatch!(self, gpu_write_value)
1519 }
1520
1521 #[cfg(target_os = "windows")]
1525 pub fn set_gpu_write(&self, value: u64) -> crate::Result<()> {
1526 dispatch!(self, set_gpu_write, value)
1527 }
1528
1529 #[cfg(target_os = "windows")]
1532 pub(crate) fn d3d11_shared_handle_value(&self) -> Option<usize> {
1533 dispatch!(self, d3d11_shared_handle_value)
1534 }
1535
1536 #[cfg(target_os = "windows")]
1544 #[allow(clippy::too_many_arguments)] pub unsafe fn from_d3d11_texture(
1546 texture: *mut std::ffi::c_void,
1547 width: usize,
1548 height: usize,
1549 format: PixelFormat,
1550 dtype: DType,
1551 access: crate::CpuAccess,
1552 name: Option<&str>,
1553 ) -> crate::Result<Self> {
1554 unsafe {
1557 match dtype {
1558 DType::U8 => {
1559 Tensor::<u8>::from_d3d11_texture(texture, width, height, format, access, name)
1560 .map(Self::U8)
1561 }
1562 DType::I8 => {
1563 Tensor::<i8>::from_d3d11_texture(texture, width, height, format, access, name)
1564 .map(Self::I8)
1565 }
1566 DType::U16 => {
1567 Tensor::<u16>::from_d3d11_texture(texture, width, height, format, access, name)
1568 .map(Self::U16)
1569 }
1570 DType::I16 => {
1571 Tensor::<i16>::from_d3d11_texture(texture, width, height, format, access, name)
1572 .map(Self::I16)
1573 }
1574 DType::U32 => {
1575 Tensor::<u32>::from_d3d11_texture(texture, width, height, format, access, name)
1576 .map(Self::U32)
1577 }
1578 DType::I32 => {
1579 Tensor::<i32>::from_d3d11_texture(texture, width, height, format, access, name)
1580 .map(Self::I32)
1581 }
1582 DType::U64 => {
1583 Tensor::<u64>::from_d3d11_texture(texture, width, height, format, access, name)
1584 .map(Self::U64)
1585 }
1586 DType::I64 => {
1587 Tensor::<i64>::from_d3d11_texture(texture, width, height, format, access, name)
1588 .map(Self::I64)
1589 }
1590 DType::F16 => {
1591 Tensor::<f16>::from_d3d11_texture(texture, width, height, format, access, name)
1592 .map(Self::F16)
1593 }
1594 DType::F32 => {
1595 Tensor::<f32>::from_d3d11_texture(texture, width, height, format, access, name)
1596 .map(Self::F32)
1597 }
1598 DType::F64 => {
1599 Tensor::<f64>::from_d3d11_texture(texture, width, height, format, access, name)
1600 .map(Self::F64)
1601 }
1602 }
1603 }
1604 }
1605
1606 #[cfg(target_os = "windows")]
1615 #[allow(clippy::too_many_arguments)] pub unsafe fn from_d3d11_shared_handle(
1617 handle: std::os::windows::io::RawHandle,
1618 width: usize,
1619 height: usize,
1620 format: PixelFormat,
1621 dtype: DType,
1622 access: crate::CpuAccess,
1623 completion: Option<(std::os::windows::io::RawHandle, u64)>,
1624 name: Option<&str>,
1625 ) -> crate::Result<Self> {
1626 unsafe {
1629 match dtype {
1630 DType::U8 => Tensor::<u8>::from_d3d11_shared_handle(
1631 handle, width, height, format, access, completion, name,
1632 )
1633 .map(Self::U8),
1634 DType::I8 => Tensor::<i8>::from_d3d11_shared_handle(
1635 handle, width, height, format, access, completion, name,
1636 )
1637 .map(Self::I8),
1638 DType::U16 => Tensor::<u16>::from_d3d11_shared_handle(
1639 handle, width, height, format, access, completion, name,
1640 )
1641 .map(Self::U16),
1642 DType::I16 => Tensor::<i16>::from_d3d11_shared_handle(
1643 handle, width, height, format, access, completion, name,
1644 )
1645 .map(Self::I16),
1646 DType::U32 => Tensor::<u32>::from_d3d11_shared_handle(
1647 handle, width, height, format, access, completion, name,
1648 )
1649 .map(Self::U32),
1650 DType::I32 => Tensor::<i32>::from_d3d11_shared_handle(
1651 handle, width, height, format, access, completion, name,
1652 )
1653 .map(Self::I32),
1654 DType::U64 => Tensor::<u64>::from_d3d11_shared_handle(
1655 handle, width, height, format, access, completion, name,
1656 )
1657 .map(Self::U64),
1658 DType::I64 => Tensor::<i64>::from_d3d11_shared_handle(
1659 handle, width, height, format, access, completion, name,
1660 )
1661 .map(Self::I64),
1662 DType::F16 => Tensor::<f16>::from_d3d11_shared_handle(
1663 handle, width, height, format, access, completion, name,
1664 )
1665 .map(Self::F16),
1666 DType::F32 => Tensor::<f32>::from_d3d11_shared_handle(
1667 handle, width, height, format, access, completion, name,
1668 )
1669 .map(Self::F32),
1670 DType::F64 => Tensor::<f64>::from_d3d11_shared_handle(
1671 handle, width, height, format, access, completion, name,
1672 )
1673 .map(Self::F64),
1674 }
1675 }
1676 }
1677
1678 #[cfg(target_os = "android")]
1688 pub unsafe fn from_hardware_buffer(
1689 buffer_ptr: *mut std::ffi::c_void,
1690 shape: &[usize],
1691 dtype: DType,
1692 name: Option<&str>,
1693 ) -> crate::Result<Self> {
1694 unsafe {
1695 match dtype {
1696 DType::U8 => {
1697 Tensor::<u8>::from_hardware_buffer(buffer_ptr, shape, name).map(Self::U8)
1698 }
1699 DType::I8 => {
1700 Tensor::<i8>::from_hardware_buffer(buffer_ptr, shape, name).map(Self::I8)
1701 }
1702 DType::U16 => {
1703 Tensor::<u16>::from_hardware_buffer(buffer_ptr, shape, name).map(Self::U16)
1704 }
1705 DType::I16 => {
1706 Tensor::<i16>::from_hardware_buffer(buffer_ptr, shape, name).map(Self::I16)
1707 }
1708 DType::U32 => {
1709 Tensor::<u32>::from_hardware_buffer(buffer_ptr, shape, name).map(Self::U32)
1710 }
1711 DType::I32 => {
1712 Tensor::<i32>::from_hardware_buffer(buffer_ptr, shape, name).map(Self::I32)
1713 }
1714 DType::U64 => {
1715 Tensor::<u64>::from_hardware_buffer(buffer_ptr, shape, name).map(Self::U64)
1716 }
1717 DType::I64 => {
1718 Tensor::<i64>::from_hardware_buffer(buffer_ptr, shape, name).map(Self::I64)
1719 }
1720 DType::F16 => {
1721 Tensor::<f16>::from_hardware_buffer(buffer_ptr, shape, name).map(Self::F16)
1722 }
1723 DType::F32 => {
1724 Tensor::<f32>::from_hardware_buffer(buffer_ptr, shape, name).map(Self::F32)
1725 }
1726 DType::F64 => {
1727 Tensor::<f64>::from_hardware_buffer(buffer_ptr, shape, name).map(Self::F64)
1728 }
1729 }
1730 }
1731 }
1732
1733 #[cfg(target_os = "android")]
1737 pub fn hardware_buffer_ptr(&self) -> Option<*mut std::ffi::c_void> {
1738 dispatch!(self, hardware_buffer_ptr)
1739 }
1740
1741 #[cfg(target_os = "android")]
1745 pub fn hardware_buffer_physical_dims(&self) -> Option<(usize, usize)> {
1746 dispatch!(self, hardware_buffer_physical_dims)
1747 }
1748
1749 pub fn copy_to_flat(&self, dst: &mut [u8]) -> crate::Result<()> {
1754 dispatch!(self, copy_to_flat, dst)
1755 }
1756
1757 pub fn image(
1775 width: usize,
1776 height: usize,
1777 format: PixelFormat,
1778 dtype: DType,
1779 memory: Option<TensorMemory>,
1780 access: crate::CpuAccess,
1781 ) -> crate::Result<Self> {
1782 match dtype {
1783 DType::U8 => Tensor::<u8>::image(width, height, format, memory, access).map(Self::U8),
1784 DType::I8 => Tensor::<i8>::image(width, height, format, memory, access).map(Self::I8),
1785 DType::U16 => {
1786 Tensor::<u16>::image(width, height, format, memory, access).map(Self::U16)
1787 }
1788 DType::I16 => {
1789 Tensor::<i16>::image(width, height, format, memory, access).map(Self::I16)
1790 }
1791 DType::U32 => {
1792 Tensor::<u32>::image(width, height, format, memory, access).map(Self::U32)
1793 }
1794 DType::I32 => {
1795 Tensor::<i32>::image(width, height, format, memory, access).map(Self::I32)
1796 }
1797 DType::U64 => {
1798 Tensor::<u64>::image(width, height, format, memory, access).map(Self::U64)
1799 }
1800 DType::I64 => {
1801 Tensor::<i64>::image(width, height, format, memory, access).map(Self::I64)
1802 }
1803 DType::F16 => {
1804 Tensor::<f16>::image(width, height, format, memory, access).map(Self::F16)
1805 }
1806 DType::F32 => {
1807 Tensor::<f32>::image(width, height, format, memory, access).map(Self::F32)
1808 }
1809 DType::F64 => {
1810 Tensor::<f64>::image(width, height, format, memory, access).map(Self::F64)
1811 }
1812 }
1813 }
1814
1815 pub fn image_desc(desc: &crate::ImageDesc) -> crate::Result<Self> {
1819 match desc.dtype() {
1820 DType::U8 => Tensor::<u8>::image_desc(desc).map(Self::U8),
1821 DType::I8 => Tensor::<i8>::image_desc(desc).map(Self::I8),
1822 DType::U16 => Tensor::<u16>::image_desc(desc).map(Self::U16),
1823 DType::I16 => Tensor::<i16>::image_desc(desc).map(Self::I16),
1824 DType::U32 => Tensor::<u32>::image_desc(desc).map(Self::U32),
1825 DType::I32 => Tensor::<i32>::image_desc(desc).map(Self::I32),
1826 DType::U64 => Tensor::<u64>::image_desc(desc).map(Self::U64),
1827 DType::I64 => Tensor::<i64>::image_desc(desc).map(Self::I64),
1828 DType::F16 => Tensor::<f16>::image_desc(desc).map(Self::F16),
1829 DType::F32 => Tensor::<f32>::image_desc(desc).map(Self::F32),
1830 DType::F64 => Tensor::<f64>::image_desc(desc).map(Self::F64),
1831 }
1832 }
1833
1834 pub fn compression(&self) -> Option<crate::CompressionScheme> {
1837 dispatch!(self, compression)
1838 }
1839
1840 pub fn contiguity(&self) -> crate::Contiguity {
1843 dispatch!(self, contiguity)
1844 }
1845
1846 pub fn image_with_stride(
1873 width: usize,
1874 height: usize,
1875 format: PixelFormat,
1876 dtype: DType,
1877 row_stride_bytes: usize,
1878 memory: Option<TensorMemory>,
1879 access: crate::CpuAccess,
1880 ) -> crate::Result<Self> {
1881 match dtype {
1882 DType::U8 => Tensor::<u8>::image_with_stride(
1883 width,
1884 height,
1885 format,
1886 row_stride_bytes,
1887 memory,
1888 access,
1889 )
1890 .map(Self::U8),
1891 DType::I8 => Tensor::<i8>::image_with_stride(
1892 width,
1893 height,
1894 format,
1895 row_stride_bytes,
1896 memory,
1897 access,
1898 )
1899 .map(Self::I8),
1900 DType::U16 => Tensor::<u16>::image_with_stride(
1901 width,
1902 height,
1903 format,
1904 row_stride_bytes,
1905 memory,
1906 access,
1907 )
1908 .map(Self::U16),
1909 DType::I16 => Tensor::<i16>::image_with_stride(
1910 width,
1911 height,
1912 format,
1913 row_stride_bytes,
1914 memory,
1915 access,
1916 )
1917 .map(Self::I16),
1918 DType::U32 => Tensor::<u32>::image_with_stride(
1919 width,
1920 height,
1921 format,
1922 row_stride_bytes,
1923 memory,
1924 access,
1925 )
1926 .map(Self::U32),
1927 DType::I32 => Tensor::<i32>::image_with_stride(
1928 width,
1929 height,
1930 format,
1931 row_stride_bytes,
1932 memory,
1933 access,
1934 )
1935 .map(Self::I32),
1936 DType::U64 => Tensor::<u64>::image_with_stride(
1937 width,
1938 height,
1939 format,
1940 row_stride_bytes,
1941 memory,
1942 access,
1943 )
1944 .map(Self::U64),
1945 DType::I64 => Tensor::<i64>::image_with_stride(
1946 width,
1947 height,
1948 format,
1949 row_stride_bytes,
1950 memory,
1951 access,
1952 )
1953 .map(Self::I64),
1954 DType::F16 => Tensor::<f16>::image_with_stride(
1955 width,
1956 height,
1957 format,
1958 row_stride_bytes,
1959 memory,
1960 access,
1961 )
1962 .map(Self::F16),
1963 DType::F32 => Tensor::<f32>::image_with_stride(
1964 width,
1965 height,
1966 format,
1967 row_stride_bytes,
1968 memory,
1969 access,
1970 )
1971 .map(Self::F32),
1972 DType::F64 => Tensor::<f64>::image_with_stride(
1973 width,
1974 height,
1975 format,
1976 row_stride_bytes,
1977 memory,
1978 access,
1979 )
1980 .map(Self::F64),
1981 }
1982 }
1983}
1984
1985impl From<Tensor<u8>> for TensorDyn {
1988 fn from(t: Tensor<u8>) -> Self {
1989 Self::U8(t)
1990 }
1991}
1992
1993impl From<Tensor<i8>> for TensorDyn {
1994 fn from(t: Tensor<i8>) -> Self {
1995 Self::I8(t)
1996 }
1997}
1998
1999impl From<Tensor<u16>> for TensorDyn {
2000 fn from(t: Tensor<u16>) -> Self {
2001 Self::U16(t)
2002 }
2003}
2004
2005impl From<Tensor<i16>> for TensorDyn {
2006 fn from(t: Tensor<i16>) -> Self {
2007 Self::I16(t)
2008 }
2009}
2010
2011impl From<Tensor<u32>> for TensorDyn {
2012 fn from(t: Tensor<u32>) -> Self {
2013 Self::U32(t)
2014 }
2015}
2016
2017impl From<Tensor<i32>> for TensorDyn {
2018 fn from(t: Tensor<i32>) -> Self {
2019 Self::I32(t)
2020 }
2021}
2022
2023impl From<Tensor<u64>> for TensorDyn {
2024 fn from(t: Tensor<u64>) -> Self {
2025 Self::U64(t)
2026 }
2027}
2028
2029impl From<Tensor<i64>> for TensorDyn {
2030 fn from(t: Tensor<i64>) -> Self {
2031 Self::I64(t)
2032 }
2033}
2034
2035impl From<Tensor<f16>> for TensorDyn {
2036 fn from(t: Tensor<f16>) -> Self {
2037 Self::F16(t)
2038 }
2039}
2040
2041impl From<Tensor<f32>> for TensorDyn {
2042 fn from(t: Tensor<f32>) -> Self {
2043 Self::F32(t)
2044 }
2045}
2046
2047impl From<Tensor<f64>> for TensorDyn {
2048 fn from(t: Tensor<f64>) -> Self {
2049 Self::F64(t)
2050 }
2051}
2052
2053impl fmt::Debug for TensorDyn {
2054 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
2055 dispatch!(self, fmt, f)
2056 }
2057}
2058
2059pub type Raw = *mut std::ffi::c_void;
2062
2063impl TensorDyn {
2064 pub fn into_raw(self) -> Raw {
2067 Box::into_raw(Box::new(self)) as Raw
2068 }
2069
2070 pub unsafe fn from_raw(p: Raw) -> TensorDyn {
2076 *unsafe { Box::from_raw(p as *mut TensorDyn) }
2077 }
2078
2079 pub unsafe fn borrow_raw<'a>(p: Raw) -> BorrowedTensor<'a> {
2089 BorrowedTensor {
2090 tensor: unsafe { std::ptr::NonNull::new_unchecked(p as *mut TensorDyn) },
2091 _handle: std::marker::PhantomData,
2092 }
2093 }
2094
2095 pub unsafe fn with_raw<R>(p: Raw, f: impl FnOnce(&mut TensorDyn) -> R) -> R {
2100 let mut t = unsafe { Self::borrow_raw(p) };
2101 f(&mut t)
2102 }
2103}
2104
2105pub struct BorrowedTensor<'a> {
2110 tensor: std::ptr::NonNull<TensorDyn>,
2111 _handle: std::marker::PhantomData<&'a mut TensorDyn>,
2112}
2113
2114unsafe impl Send for BorrowedTensor<'_> {}
2117unsafe impl Sync for BorrowedTensor<'_> {}
2118
2119impl std::ops::Deref for BorrowedTensor<'_> {
2120 type Target = TensorDyn;
2121 fn deref(&self) -> &TensorDyn {
2122 unsafe { self.tensor.as_ref() }
2124 }
2125}
2126
2127impl std::ops::DerefMut for BorrowedTensor<'_> {
2128 fn deref_mut(&mut self) -> &mut TensorDyn {
2129 unsafe { self.tensor.as_mut() }
2131 }
2132}
2133
2134impl fmt::Debug for BorrowedTensor<'_> {
2135 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
2136 fmt::Debug::fmt(&**self, f)
2137 }
2138}
2139
2140#[cfg(test)]
2141mod tests {
2142 use super::*;
2143
2144 #[test]
2145 fn from_typed_tensor() {
2146 let t = Tensor::<u8>::new(&[10], None, None).unwrap();
2147 let dyn_t: TensorDyn = t.into();
2148 assert_eq!(dyn_t.dtype(), DType::U8);
2149 assert_eq!(dyn_t.shape(), &[10]);
2150 }
2151
2152 #[test]
2153 fn from_foreign_ptr_wraps_borrowed_memory() {
2154 use crate::TensorMapTrait;
2155 let mut vec: Vec<f32> = vec![0.0; 4];
2158 let ptr = vec.as_mut_ptr() as *mut u8;
2159 let owner: crate::ForeignOwner = Box::new(vec);
2160 let t = unsafe {
2161 TensorDyn::from_foreign_ptr(ptr, &[2, 2], DType::F32, Some(owner), Some("trt_output"))
2162 }
2163 .unwrap();
2164 assert_eq!(t.dtype(), DType::F32);
2165 assert_eq!(t.memory(), TensorMemory::Mem);
2166 assert_eq!(t.shape(), &[2, 2]);
2167 {
2168 let mut m = t.as_f32().unwrap().map().unwrap();
2169 m.as_mut_slice().copy_from_slice(&[1.0, 2.0, 3.0, 4.0]);
2170 }
2171 let m = t.as_f32().unwrap().map().unwrap();
2172 assert_eq!(m.as_slice(), &[1.0, 2.0, 3.0, 4.0]);
2173 }
2174
2175 #[test]
2184 fn from_foreign_ptr_rejects_null_ptr() {
2185 let err = unsafe {
2186 TensorDyn::from_foreign_ptr(std::ptr::null_mut(), &[4], DType::U8, None, None)
2187 }
2188 .unwrap_err();
2189 assert!(
2191 matches!(err, crate::error::Error::InvalidArgument(ref m) if m.contains("non-null")),
2192 "expected InvalidArgument(non-null), got {err:?}"
2193 );
2194 }
2195
2196 #[test]
2197 fn from_foreign_ptr_rejects_empty_shape() {
2198 let mut dummy: u8 = 0;
2199 let err = unsafe {
2200 TensorDyn::from_foreign_ptr(&mut dummy as *mut u8, &[], DType::U8, None, None)
2201 }
2202 .unwrap_err();
2203 assert!(
2204 matches!(err, crate::error::Error::InvalidSize(0)),
2205 "expected InvalidSize(0) for empty shape, got {err:?}"
2206 );
2207 }
2208
2209 #[test]
2210 fn from_foreign_ptr_rejects_overflow_shape() {
2211 let mut dummy: u8 = 0;
2212 let huge = [usize::MAX / 2 + 1, 2];
2213 let err = unsafe { TensorDyn::from_foreign_ptr(&mut dummy, &huge, DType::U8, None, None) }
2214 .unwrap_err();
2215 assert!(
2216 matches!(err, crate::error::Error::InvalidArgument(ref m) if m.contains("overflow")),
2217 "expected InvalidArgument(overflow), got {err:?}"
2218 );
2219 }
2220
2221 #[test]
2222 fn from_foreign_ptr_u8_dtype_dispatch() {
2223 let mut buf: Vec<u8> = vec![1, 2, 3, 4];
2226 let ptr = buf.as_mut_ptr();
2227 let owner: crate::ForeignOwner = Box::new(buf);
2228 let t = unsafe {
2229 TensorDyn::from_foreign_ptr(ptr, &[4], DType::U8, Some(owner), Some("u8_foreign"))
2230 }
2231 .unwrap();
2232 assert_eq!(t.dtype(), DType::U8);
2233 assert_eq!(t.shape(), &[4]);
2234 let m = t.as_u8().unwrap().map().unwrap();
2235 use crate::TensorMapTrait;
2236 assert_eq!(m.as_slice(), &[1u8, 2, 3, 4]);
2237 }
2238
2239 #[test]
2240 fn map_bytes_roundtrips_and_is_typed_agnostic() {
2241 use crate::TensorMapTrait;
2242 let t = Tensor::<f32>::new(&[4], None, None).unwrap();
2245 let dyn_t: TensorDyn = t.into();
2246 {
2247 let mut view = dyn_t.map_bytes(crate::CpuAccess::ReadWrite).unwrap();
2248 assert_eq!(view.len(), 16, "4 x f32 is 16 bytes, not 4");
2249 view.as_mut_slice()
2250 .copy_from_slice(&[1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16]);
2251 }
2252 let view = dyn_t.map_bytes(crate::CpuAccess::Read).unwrap();
2253 assert_eq!(
2254 view.as_slice(),
2255 &[1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16]
2256 );
2257 }
2258
2259 #[test]
2260 fn map_bytes_none_is_invalid() {
2261 let t = Tensor::<u8>::new(&[4], None, None).unwrap();
2262 let dyn_t: TensorDyn = t.into();
2263 assert!(matches!(
2264 dyn_t.map_bytes(crate::CpuAccess::None),
2265 Err(crate::error::Error::InvalidArgument(_))
2266 ));
2267 }
2268
2269 #[test]
2270 fn downcast_ref() {
2271 let t = Tensor::<u8>::new(&[10], None, None).unwrap();
2272 let dyn_t: TensorDyn = t.into();
2273 assert!(dyn_t.as_u8().is_some());
2274 assert!(dyn_t.as_i8().is_none());
2275 }
2276
2277 #[test]
2278 fn downcast_into() {
2279 let t = Tensor::<u8>::new(&[10], None, None).unwrap();
2280 let dyn_t: TensorDyn = t.into();
2281 let back = dyn_t.into_u8().unwrap();
2282 assert_eq!(back.shape(), &[10]);
2283 }
2284
2285 #[test]
2286 fn image_accessors() {
2287 let t = Tensor::<u8>::image(
2288 640,
2289 480,
2290 PixelFormat::Rgba,
2291 None,
2292 crate::CpuAccess::ReadWrite,
2293 )
2294 .unwrap();
2295 let dyn_t: TensorDyn = t.into();
2296 assert_eq!(dyn_t.format(), Some(PixelFormat::Rgba));
2297 assert_eq!(dyn_t.width(), Some(640));
2298 assert_eq!(dyn_t.height(), Some(480));
2299 assert!(!dyn_t.is_multiplane());
2300 }
2301
2302 #[test]
2303 fn cpu_access_reports_the_declared_access() {
2304 let read_t =
2306 Tensor::<u8>::image(4, 4, PixelFormat::Rgba, None, crate::CpuAccess::Read).unwrap();
2307 let dyn_read: TensorDyn = read_t.into();
2308 assert_eq!(dyn_read.cpu_access(), crate::CpuAccess::Read);
2309 assert!(!dyn_read.cpu_access().writes());
2310
2311 let mem_t = Tensor::<u8>::new(&[4], None, None).unwrap();
2315 let dyn_mem: TensorDyn = mem_t.into();
2316 assert_eq!(dyn_mem.cpu_access(), crate::CpuAccess::ReadWrite);
2317 assert!(dyn_mem.cpu_access().writes());
2318 }
2319
2320 #[test]
2321 fn image_constructor() {
2322 let dyn_t = TensorDyn::image(
2323 640,
2324 480,
2325 PixelFormat::Rgb,
2326 DType::U8,
2327 None,
2328 crate::CpuAccess::ReadWrite,
2329 )
2330 .unwrap();
2331 assert_eq!(dyn_t.dtype(), DType::U8);
2332 assert_eq!(dyn_t.format(), Some(PixelFormat::Rgb));
2333 assert_eq!(dyn_t.width(), Some(640));
2334 }
2335
2336 #[test]
2337 fn image_constructor_i8() {
2338 let dyn_t = TensorDyn::image(
2339 640,
2340 480,
2341 PixelFormat::Rgb,
2342 DType::I8,
2343 None,
2344 crate::CpuAccess::ReadWrite,
2345 )
2346 .unwrap();
2347 assert_eq!(dyn_t.dtype(), DType::I8);
2348 assert_eq!(dyn_t.format(), Some(PixelFormat::Rgb));
2349 }
2350
2351 #[test]
2352 fn set_format_packed() {
2353 let mut t = TensorDyn::new(&[480, 640, 3], DType::U8, None, None).unwrap();
2354 assert_eq!(t.format(), None);
2355 t.set_format(PixelFormat::Rgb).unwrap();
2356 assert_eq!(t.format(), Some(PixelFormat::Rgb));
2357 assert_eq!(t.width(), Some(640));
2358 assert_eq!(t.height(), Some(480));
2359 }
2360
2361 #[test]
2362 fn set_format_planar() {
2363 let mut t = TensorDyn::new(&[3, 480, 640], DType::U8, None, None).unwrap();
2364 t.set_format(PixelFormat::PlanarRgb).unwrap();
2365 assert_eq!(t.format(), Some(PixelFormat::PlanarRgb));
2366 assert_eq!(t.width(), Some(640));
2367 assert_eq!(t.height(), Some(480));
2368 }
2369
2370 #[test]
2371 fn set_format_rejects_wrong_shape() {
2372 let mut t = TensorDyn::new(&[480, 640, 4], DType::U8, None, None).unwrap();
2373 assert!(t.set_format(PixelFormat::Rgb).is_err());
2374 }
2375
2376 #[test]
2377 fn with_format_builder() {
2378 let t = TensorDyn::new(&[480, 640, 4], DType::U8, None, None)
2379 .unwrap()
2380 .with_format(PixelFormat::Rgba)
2381 .unwrap();
2382 assert_eq!(t.format(), Some(PixelFormat::Rgba));
2383 assert_eq!(t.width(), Some(640));
2384 assert_eq!(t.height(), Some(480));
2385 }
2386
2387 #[cfg(target_os = "linux")]
2388 #[test]
2389 fn dmabuf_clone_mem_tensor_fails() {
2390 let t = TensorDyn::new(&[480, 640, 3], DType::U8, Some(TensorMemory::Mem), None).unwrap();
2391 assert_eq!(t.memory(), TensorMemory::Mem);
2392 assert!(t.dmabuf_clone().is_err());
2393 }
2394
2395 #[cfg(target_os = "linux")]
2396 #[test]
2397 fn dmabuf_mem_tensor_fails() {
2398 let t = TensorDyn::new(&[480, 640, 3], DType::U8, Some(TensorMemory::Mem), None).unwrap();
2399 assert!(t.dmabuf().is_err());
2400 }
2401
2402 #[test]
2403 fn set_format_semi_planar_nv12() {
2404 let mut t = TensorDyn::new(&[720, 640], DType::U8, Some(TensorMemory::Mem), None).unwrap();
2406 t.set_format(PixelFormat::Nv12).unwrap();
2407 assert_eq!(t.format(), Some(PixelFormat::Nv12));
2408 assert_eq!(t.width(), Some(640));
2409 assert_eq!(t.height(), Some(480));
2410 }
2411
2412 #[test]
2413 fn set_format_semi_planar_nv16() {
2414 let mut t = TensorDyn::new(&[960, 640], DType::U8, Some(TensorMemory::Mem), None).unwrap();
2416 t.set_format(PixelFormat::Nv16).unwrap();
2417 assert_eq!(t.format(), Some(PixelFormat::Nv16));
2418 assert_eq!(t.width(), Some(640));
2419 assert_eq!(t.height(), Some(480));
2420 }
2421
2422 #[test]
2423 fn with_format_rejects_wrong_shape() {
2424 let result = TensorDyn::new(&[480, 640, 4], DType::U8, None, None)
2425 .unwrap()
2426 .with_format(PixelFormat::Rgb);
2427 assert!(result.is_err());
2428 }
2429
2430 #[test]
2431 fn set_format_preserved_after_rejection() {
2432 let mut t = TensorDyn::new(&[480, 640, 3], DType::U8, None, None).unwrap();
2433 t.set_format(PixelFormat::Rgb).unwrap();
2434 assert_eq!(t.format(), Some(PixelFormat::Rgb));
2435
2436 assert!(t.set_format(PixelFormat::Rgba).is_err());
2438
2439 assert_eq!(t.format(), Some(PixelFormat::Rgb));
2441 }
2442
2443 #[test]
2444 fn set_format_idempotent() {
2445 let mut t = TensorDyn::new(&[480, 640, 3], DType::U8, None, None).unwrap();
2446 t.set_format(PixelFormat::Rgb).unwrap();
2447 t.set_format(PixelFormat::Rgb).unwrap();
2448 assert_eq!(t.format(), Some(PixelFormat::Rgb));
2449 assert_eq!(t.width(), Some(640));
2450 assert_eq!(t.height(), Some(480));
2451 }
2452
2453 #[test]
2456 fn set_row_stride_valid() {
2457 let mut t = TensorDyn::image(
2459 100,
2460 100,
2461 PixelFormat::Rgba,
2462 DType::U8,
2463 None,
2464 crate::CpuAccess::ReadWrite,
2465 )
2466 .unwrap();
2467 t.set_row_stride(512).unwrap();
2468 assert_eq!(t.row_stride(), Some(512));
2469 assert_eq!(t.effective_row_stride(), Some(512));
2470 }
2471
2472 #[test]
2473 fn set_row_stride_equals_min() {
2474 let mut t = TensorDyn::image(
2476 100,
2477 100,
2478 PixelFormat::Rgb,
2479 DType::U8,
2480 None,
2481 crate::CpuAccess::ReadWrite,
2482 )
2483 .unwrap();
2484 t.set_row_stride(300).unwrap();
2485 assert_eq!(t.row_stride(), Some(300));
2486 }
2487
2488 #[test]
2489 fn set_row_stride_too_small() {
2490 let mut t = TensorDyn::image(
2495 64,
2496 100,
2497 PixelFormat::Rgba,
2498 DType::U8,
2499 None,
2500 crate::CpuAccess::ReadWrite,
2501 )
2502 .unwrap();
2503 assert!(t.set_row_stride(200).is_err());
2504 assert_eq!(t.row_stride(), None);
2505 }
2506
2507 #[test]
2508 fn set_row_stride_zero() {
2509 let mut t = TensorDyn::image(
2510 100,
2511 100,
2512 PixelFormat::Rgb,
2513 DType::U8,
2514 None,
2515 crate::CpuAccess::ReadWrite,
2516 )
2517 .unwrap();
2518 assert!(t.set_row_stride(0).is_err());
2519 }
2520
2521 #[test]
2522 fn set_row_stride_requires_format() {
2523 let mut t = TensorDyn::new(&[480, 640, 3], DType::U8, None, None).unwrap();
2524 assert!(t.set_row_stride(2048).is_err());
2525 }
2526
2527 #[test]
2528 fn effective_row_stride_without_stride() {
2529 let t = TensorDyn::image(
2534 64,
2535 100,
2536 PixelFormat::Rgb,
2537 DType::U8,
2538 None,
2539 crate::CpuAccess::ReadWrite,
2540 )
2541 .unwrap();
2542 assert_eq!(t.row_stride(), None);
2543 assert_eq!(t.effective_row_stride(), Some(192)); }
2545
2546 #[test]
2547 fn effective_row_stride_padded_packed_dma() {
2548 let result = TensorDyn::image(
2554 100,
2555 100,
2556 PixelFormat::Rgb,
2557 DType::U8,
2558 Some(TensorMemory::DmaBuf),
2559 crate::CpuAccess::ReadWrite,
2560 );
2561 #[cfg(target_os = "linux")]
2562 let t = crate::test_support::dma_allocated(result, "DMA-BUF RGB image allocation");
2563 #[cfg(not(target_os = "linux"))]
2564 let t = crate::test_support::allocated(
2565 result,
2566 "zero-copy RGB image allocation",
2567 if crate::is_gpu_buffer_available() {
2568 crate::test_support::Expect::Succeeds
2569 } else {
2570 crate::test_support::Expect::MayFail(
2571 "no zero-copy GPU buffer allocator on this host",
2572 )
2573 },
2574 );
2575 let Some(t) = t else {
2576 return;
2577 };
2578 assert_eq!(t.memory(), TensorMemory::DmaBuf);
2579 if cfg!(target_os = "android") {
2585 assert!(
2586 t.row_stride().is_none() || t.row_stride() >= Some(300),
2587 "AHardwareBuffer pitch must be absent or at least the tight row"
2588 );
2589 return;
2590 }
2591 if cfg!(target_os = "windows") {
2596 assert!(
2597 t.row_stride().is_none() || t.row_stride() >= Some(300),
2598 "D3D11 staging pitch must be absent or at least the tight row"
2599 );
2600 assert!(t.effective_row_stride() >= Some(300));
2601 return;
2602 }
2603 assert_eq!(t.row_stride(), Some(320));
2604 assert_eq!(t.effective_row_stride(), Some(320));
2605 }
2606
2607 #[test]
2608 fn effective_row_stride_no_format() {
2609 let t = TensorDyn::new(&[480, 640, 3], DType::U8, None, None).unwrap();
2610 assert_eq!(t.effective_row_stride(), None);
2611 }
2612
2613 #[test]
2614 fn with_row_stride_builder() {
2615 let t = TensorDyn::image(
2616 100,
2617 100,
2618 PixelFormat::Rgba,
2619 DType::U8,
2620 None,
2621 crate::CpuAccess::ReadWrite,
2622 )
2623 .unwrap()
2624 .with_row_stride(512)
2625 .unwrap();
2626 assert_eq!(t.row_stride(), Some(512));
2627 assert_eq!(t.effective_row_stride(), Some(512));
2628 }
2629
2630 #[test]
2631 fn with_row_stride_rejects_small() {
2632 let result = TensorDyn::image(
2633 100,
2634 100,
2635 PixelFormat::Rgba,
2636 DType::U8,
2637 None,
2638 crate::CpuAccess::ReadWrite,
2639 )
2640 .unwrap()
2641 .with_row_stride(200);
2642 assert!(result.is_err());
2643 }
2644
2645 #[test]
2646 fn set_format_clears_row_stride() {
2647 let mut t = TensorDyn::new(&[480, 640, 3], DType::U8, None, None).unwrap();
2648 t.set_format(PixelFormat::Rgb).unwrap();
2649 t.set_row_stride(2048).unwrap();
2650 assert_eq!(t.row_stride(), Some(2048));
2651
2652 let _ = t.set_format(PixelFormat::Bgra);
2654 assert_eq!(t.row_stride(), Some(2048));
2655
2656 t.set_format(PixelFormat::Rgb).unwrap();
2658 assert_eq!(t.row_stride(), Some(2048));
2659
2660 t.reshape(&[480 * 640 * 3]).unwrap();
2662 assert_eq!(t.row_stride(), None);
2663 assert_eq!(t.format(), None);
2664 }
2665
2666 #[test]
2667 fn set_format_different_compatible_clears_stride() {
2668 let mut t = TensorDyn::new(&[480, 640, 4], DType::U8, None, None).unwrap();
2671 t.set_format(PixelFormat::Rgba).unwrap();
2672 t.set_row_stride(4096).unwrap();
2673 assert_eq!(t.row_stride(), Some(4096));
2674
2675 t.set_format(PixelFormat::Bgra).unwrap();
2677 assert_eq!(t.format(), Some(PixelFormat::Bgra));
2678 assert_eq!(t.row_stride(), None);
2679 }
2680
2681 #[test]
2682 fn set_format_same_preserves_stride() {
2683 let mut t = TensorDyn::image(
2684 100,
2685 100,
2686 PixelFormat::Rgb,
2687 DType::U8,
2688 None,
2689 crate::CpuAccess::ReadWrite,
2690 )
2691 .unwrap();
2692 t.set_row_stride(512).unwrap();
2693 t.set_format(PixelFormat::Rgb).unwrap();
2695 assert_eq!(t.row_stride(), Some(512));
2696 }
2697
2698 #[test]
2699 fn effective_row_stride_planar() {
2700 let t = TensorDyn::image(
2701 640,
2702 480,
2703 PixelFormat::PlanarRgb,
2704 DType::U8,
2705 None,
2706 crate::CpuAccess::ReadWrite,
2707 )
2708 .unwrap();
2709 assert_eq!(t.effective_row_stride(), Some(640)); }
2711
2712 #[test]
2713 fn effective_row_stride_nv12() {
2714 let t = TensorDyn::image(
2715 640,
2716 480,
2717 PixelFormat::Nv12,
2718 DType::U8,
2719 None,
2720 crate::CpuAccess::ReadWrite,
2721 )
2722 .unwrap();
2723 assert_eq!(t.effective_row_stride(), Some(640)); }
2725
2726 #[test]
2727 fn map_rejects_strided_tensor() {
2728 let mut t = Tensor::<u8>::image(
2729 100,
2730 100,
2731 PixelFormat::Rgba,
2732 Some(TensorMemory::Mem),
2733 crate::CpuAccess::ReadWrite,
2734 )
2735 .unwrap();
2736 assert!(t.map().is_ok());
2738 t.set_row_stride(512).unwrap();
2740 let err = t.map();
2741 assert!(err.is_err());
2742 }
2743
2744 #[test]
2747 fn plane_offset_default_none() {
2748 let t = TensorDyn::image(
2749 100,
2750 100,
2751 PixelFormat::Rgba,
2752 DType::U8,
2753 None,
2754 crate::CpuAccess::ReadWrite,
2755 )
2756 .unwrap();
2757 assert_eq!(t.plane_offset(), None);
2758 }
2759
2760 #[test]
2761 fn set_plane_offset_basic() {
2762 let mut t = TensorDyn::image(
2763 100,
2764 100,
2765 PixelFormat::Rgba,
2766 DType::U8,
2767 None,
2768 crate::CpuAccess::ReadWrite,
2769 )
2770 .unwrap();
2771 t.set_plane_offset(4096);
2772 assert_eq!(t.plane_offset(), Some(4096));
2773 }
2774
2775 #[test]
2776 fn set_plane_offset_zero() {
2777 let mut t = TensorDyn::image(
2778 100,
2779 100,
2780 PixelFormat::Rgb,
2781 DType::U8,
2782 None,
2783 crate::CpuAccess::ReadWrite,
2784 )
2785 .unwrap();
2786 t.set_plane_offset(0);
2787 assert_eq!(t.plane_offset(), Some(0));
2788 }
2789
2790 #[test]
2791 fn set_plane_offset_no_format() {
2792 let mut t = TensorDyn::new(&[480, 640, 3], DType::U8, None, None).unwrap();
2794 t.set_plane_offset(4096);
2795 assert_eq!(t.plane_offset(), Some(4096));
2796 }
2797
2798 #[test]
2799 fn set_format_clears_plane_offset() {
2800 let mut t = TensorDyn::new(&[480, 640, 3], DType::U8, None, None).unwrap();
2801 t.set_format(PixelFormat::Rgb).unwrap();
2802 t.set_plane_offset(4096);
2803 assert_eq!(t.plane_offset(), Some(4096));
2804
2805 t.set_format(PixelFormat::Rgb).unwrap();
2807 assert_eq!(t.plane_offset(), Some(4096));
2808
2809 t.reshape(&[480 * 640 * 3]).unwrap();
2811 assert_eq!(t.plane_offset(), None);
2812 assert_eq!(t.format(), None);
2813 }
2814
2815 #[test]
2816 fn map_rejects_out_of_bounds_offset() {
2817 let mut t = Tensor::<u8>::image(
2818 100,
2819 100,
2820 PixelFormat::Rgba,
2821 Some(TensorMemory::Mem),
2822 crate::CpuAccess::ReadWrite,
2823 )
2824 .unwrap();
2825 assert!(t.map().is_ok());
2827 t.set_plane_offset(4096);
2830 assert!(t.map().is_err());
2831 }
2832
2833 #[test]
2834 fn mem_subview_in_bounds_maps_at_offset() {
2835 let parent = Tensor::<u8>::image(
2838 100,
2839 100,
2840 PixelFormat::Rgba,
2841 Some(TensorMemory::Mem),
2842 crate::CpuAccess::ReadWrite,
2843 )
2844 .unwrap();
2845 let view = parent.subview(4096, &[10, 10, 4]).unwrap();
2847 assert_eq!(view.plane_offset(), Some(4096));
2848 assert!(view.map().is_ok());
2849 }
2850
2851 #[test]
2852 fn dyn_batch_dispatches_every_dtype() {
2853 use DType::*;
2857 for dt in [U8, I8, U16, I16, U32, I32, U64, I64, F16, F32, F64] {
2858 let parent = TensorDyn::new(&[2, 4], dt, Some(TensorMemory::Mem), None).unwrap();
2859 let view = parent.batch(1).unwrap();
2860 assert_eq!(view.dtype(), dt, "batch must preserve dtype {dt:?}");
2861 assert_eq!(view.shape(), &[4], "{dt:?}");
2862 }
2863 }
2864
2865 #[test]
2866 fn map_accepts_zero_offset_tensor() {
2867 let mut t = Tensor::<u8>::image(
2868 100,
2869 100,
2870 PixelFormat::Rgba,
2871 Some(TensorMemory::Mem),
2872 crate::CpuAccess::ReadWrite,
2873 )
2874 .unwrap();
2875 t.set_plane_offset(0);
2876 assert!(t.map().is_ok());
2878 }
2879
2880 #[test]
2881 fn dyn_configure_image_nv12() {
2882 let mut t = TensorDyn::image(
2883 640,
2884 480,
2885 PixelFormat::Rgb,
2886 DType::U8,
2887 None,
2888 crate::CpuAccess::ReadWrite,
2889 )
2890 .unwrap();
2891 t.configure_image(320, 240, PixelFormat::Nv12).unwrap();
2892 assert_eq!(t.format(), Some(PixelFormat::Nv12));
2893 assert_eq!((t.width(), t.height()), (Some(320), Some(240)));
2894 }
2895
2896 #[test]
2897 fn tensordyn_colorimetry_roundtrip() {
2898 use crate::{ColorEncoding, Colorimetry, DType, PixelFormat};
2899 let mut t = TensorDyn::image(
2900 1280,
2901 720,
2902 PixelFormat::Nv12,
2903 DType::U8,
2904 None,
2905 crate::CpuAccess::ReadWrite,
2906 )
2907 .unwrap();
2908 assert_eq!(t.colorimetry(), None);
2909 let c = Colorimetry::default().with_encoding(ColorEncoding::Bt709);
2910 t.set_colorimetry(Some(c));
2911 assert_eq!(t.colorimetry(), Some(c));
2912 }
2913
2914 #[test]
2915 fn from_planes_propagates_plane_offset() {
2916 let mut luma =
2917 Tensor::<u8>::new(&[480, 640], Some(TensorMemory::Mem), Some("luma")).unwrap();
2918 luma.set_plane_offset(4096);
2919 let chroma =
2920 Tensor::<u8>::new(&[240, 640], Some(TensorMemory::Mem), Some("chroma")).unwrap();
2921 let combined = Tensor::<u8>::from_planes(luma, chroma, PixelFormat::Nv12).unwrap();
2922 assert_eq!(combined.plane_offset(), Some(4096));
2923 }
2924
2925 #[test]
2926 fn cuda_passthrough_none_for_mem_tensor() {
2927 let t: TensorDyn = Tensor::<f32>::new(&[10], Some(TensorMemory::Mem), None)
2930 .unwrap()
2931 .into();
2932 assert!(t.cuda().is_none());
2933 assert!(t.cuda_map().is_none());
2934 }
2935}