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 quantization(&self) -> Option<&crate::Quantization> {
675 match self {
676 Self::U8(t) => t.quantization(),
677 Self::I8(t) => t.quantization(),
678 Self::U16(t) => t.quantization(),
679 Self::I16(t) => t.quantization(),
680 Self::U32(t) => t.quantization(),
681 Self::I32(t) => t.quantization(),
682 Self::U64(t) => t.quantization(),
683 Self::I64(t) => t.quantization(),
684 Self::F16(_) | Self::F32(_) | Self::F64(_) => None,
685 }
686 }
687
688 pub fn set_quantization(&mut self, q: crate::Quantization) -> crate::Result<()> {
692 match self {
693 Self::U8(t) => t.set_quantization(q),
694 Self::I8(t) => t.set_quantization(q),
695 Self::U16(t) => t.set_quantization(q),
696 Self::I16(t) => t.set_quantization(q),
697 Self::U32(t) => t.set_quantization(q),
698 Self::I32(t) => t.set_quantization(q),
699 Self::U64(t) => t.set_quantization(q),
700 Self::I64(t) => t.set_quantization(q),
701 Self::F16(_) | Self::F32(_) | Self::F64(_) => Err(crate::Error::QuantizationInvalid {
702 field: "dtype_is_integer",
703 expected: "integer tensor dtype (u8/i8/u16/i16/u32/i32/u64/i64)".to_string(),
704 got: format!("{:?}", self.dtype()),
705 }),
706 }
707 }
708
709 pub fn with_quantization(mut self, q: crate::Quantization) -> crate::Result<Self> {
712 self.set_quantization(q)?;
713 Ok(self)
714 }
715
716 pub fn clear_quantization(&mut self) {
718 match self {
719 Self::U8(t) => t.clear_quantization(),
720 Self::I8(t) => t.clear_quantization(),
721 Self::U16(t) => t.clear_quantization(),
722 Self::I16(t) => t.clear_quantization(),
723 Self::U32(t) => t.clear_quantization(),
724 Self::I32(t) => t.clear_quantization(),
725 Self::U64(t) => t.clear_quantization(),
726 Self::I64(t) => t.clear_quantization(),
727 Self::F16(_) | Self::F32(_) | Self::F64(_) => {}
728 }
729 }
730
731 #[cfg(unix)]
733 pub fn clone_fd(&self) -> crate::Result<std::os::fd::OwnedFd> {
734 dispatch!(self, clone_fd)
735 }
736
737 #[cfg(target_os = "linux")]
748 pub fn dmabuf_clone(&self) -> crate::Result<std::os::fd::OwnedFd> {
749 if self.memory() != TensorMemory::DmaBuf {
750 return Err(crate::Error::NotImplemented(format!(
751 "dmabuf_clone requires DMA-backed tensor, got {:?}",
752 self.memory()
753 )));
754 }
755 self.clone_fd()
756 }
757
758 #[cfg(target_os = "linux")]
769 pub fn dmabuf(&self) -> crate::Result<std::os::fd::BorrowedFd<'_>> {
770 dispatch!(self, dmabuf)
771 }
772
773 pub fn is_multiplane(&self) -> bool {
775 dispatch!(self, is_multiplane)
776 }
777
778 pub fn buffer_identity(&self) -> &crate::BufferIdentity {
788 dispatch!(self, buffer_identity)
789 }
790
791 pub fn aliases(&self, other: &Self) -> bool {
809 if self.buffer_identity().id() == other.buffer_identity().id() {
810 return true;
811 }
812 if self.memory() != other.memory() {
813 return false;
814 }
815 #[cfg(target_os = "linux")]
816 if self.memory() == TensorMemory::DmaBuf {
817 use std::os::fd::AsRawFd;
818 if let (Ok(a), Ok(b)) = (self.dmabuf(), other.dmabuf()) {
819 return a.as_raw_fd() == b.as_raw_fd();
820 }
821 }
822 false
823 }
824
825 downcast_methods!(U8, u8, as_u8, as_u8_mut, into_u8);
828 downcast_methods!(I8, i8, as_i8, as_i8_mut, into_i8);
829 downcast_methods!(U16, u16, as_u16, as_u16_mut, into_u16);
830 downcast_methods!(I16, i16, as_i16, as_i16_mut, into_i16);
831 downcast_methods!(U32, u32, as_u32, as_u32_mut, into_u32);
832 downcast_methods!(I32, i32, as_i32, as_i32_mut, into_i32);
833 downcast_methods!(U64, u64, as_u64, as_u64_mut, into_u64);
834 downcast_methods!(I64, i64, as_i64, as_i64_mut, into_i64);
835 downcast_methods!(F16, f16, as_f16, as_f16_mut, into_f16);
836 downcast_methods!(F32, f32, as_f32, as_f32_mut, into_f32);
837 downcast_methods!(F64, f64, as_f64, as_f64_mut, into_f64);
838
839 pub fn new(
841 shape: &[usize],
842 dtype: DType,
843 memory: Option<TensorMemory>,
844 name: Option<&str>,
845 ) -> crate::Result<Self> {
846 match dtype {
847 DType::U8 => Tensor::<u8>::new(shape, memory, name).map(Self::U8),
848 DType::I8 => Tensor::<i8>::new(shape, memory, name).map(Self::I8),
849 DType::U16 => Tensor::<u16>::new(shape, memory, name).map(Self::U16),
850 DType::I16 => Tensor::<i16>::new(shape, memory, name).map(Self::I16),
851 DType::U32 => Tensor::<u32>::new(shape, memory, name).map(Self::U32),
852 DType::I32 => Tensor::<i32>::new(shape, memory, name).map(Self::I32),
853 DType::U64 => Tensor::<u64>::new(shape, memory, name).map(Self::U64),
854 DType::I64 => Tensor::<i64>::new(shape, memory, name).map(Self::I64),
855 DType::F16 => Tensor::<f16>::new(shape, memory, name).map(Self::F16),
856 DType::F32 => Tensor::<f32>::new(shape, memory, name).map(Self::F32),
857 DType::F64 => Tensor::<f64>::new(shape, memory, name).map(Self::F64),
858 }
859 }
860
861 #[cfg(unix)]
888 pub fn from_fd(
889 fd: std::os::fd::OwnedFd,
890 shape: &[usize],
891 dtype: DType,
892 name: Option<&str>,
893 ) -> crate::Result<Self> {
894 match dtype {
895 DType::U8 => Tensor::<u8>::from_fd(fd, shape, name).map(Self::U8),
896 DType::I8 => Tensor::<i8>::from_fd(fd, shape, name).map(Self::I8),
897 DType::U16 => Tensor::<u16>::from_fd(fd, shape, name).map(Self::U16),
898 DType::I16 => Tensor::<i16>::from_fd(fd, shape, name).map(Self::I16),
899 DType::U32 => Tensor::<u32>::from_fd(fd, shape, name).map(Self::U32),
900 DType::I32 => Tensor::<i32>::from_fd(fd, shape, name).map(Self::I32),
901 DType::U64 => Tensor::<u64>::from_fd(fd, shape, name).map(Self::U64),
902 DType::I64 => Tensor::<i64>::from_fd(fd, shape, name).map(Self::I64),
903 DType::F16 => Tensor::<f16>::from_fd(fd, shape, name).map(Self::F16),
904 DType::F32 => Tensor::<f32>::from_fd(fd, shape, name).map(Self::F32),
905 DType::F64 => Tensor::<f64>::from_fd(fd, shape, name).map(Self::F64),
906 }
907 }
908
909 pub(crate) fn from_pbo_import(
915 buffer_id: u32,
916 size: usize,
917 shape: &[usize],
918 dtype: DType,
919 ops: std::sync::Arc<dyn crate::PboOps>,
920 ) -> crate::Result<Self> {
921 match dtype {
922 DType::U8 => crate::PboTensor::<u8>::from_pbo(buffer_id, size, shape, None, ops)
923 .and_then(|p| Tensor::from_pbo(p).map(Self::U8)),
924 DType::I8 => crate::PboTensor::<i8>::from_pbo(buffer_id, size, shape, None, ops)
925 .and_then(|p| Tensor::from_pbo(p).map(Self::I8)),
926 DType::U16 => crate::PboTensor::<u16>::from_pbo(buffer_id, size, shape, None, ops)
927 .and_then(|p| Tensor::from_pbo(p).map(Self::U16)),
928 DType::I16 => crate::PboTensor::<i16>::from_pbo(buffer_id, size, shape, None, ops)
929 .and_then(|p| Tensor::from_pbo(p).map(Self::I16)),
930 DType::U32 => crate::PboTensor::<u32>::from_pbo(buffer_id, size, shape, None, ops)
931 .and_then(|p| Tensor::from_pbo(p).map(Self::U32)),
932 DType::I32 => crate::PboTensor::<i32>::from_pbo(buffer_id, size, shape, None, ops)
933 .and_then(|p| Tensor::from_pbo(p).map(Self::I32)),
934 DType::U64 => crate::PboTensor::<u64>::from_pbo(buffer_id, size, shape, None, ops)
935 .and_then(|p| Tensor::from_pbo(p).map(Self::U64)),
936 DType::I64 => crate::PboTensor::<i64>::from_pbo(buffer_id, size, shape, None, ops)
937 .and_then(|p| Tensor::from_pbo(p).map(Self::I64)),
938 DType::F16 => crate::PboTensor::<f16>::from_pbo(buffer_id, size, shape, None, ops)
939 .and_then(|p| Tensor::from_pbo(p).map(Self::F16)),
940 DType::F32 => crate::PboTensor::<f32>::from_pbo(buffer_id, size, shape, None, ops)
941 .and_then(|p| Tensor::from_pbo(p).map(Self::F32)),
942 DType::F64 => crate::PboTensor::<f64>::from_pbo(buffer_id, size, shape, None, ops)
943 .and_then(|p| Tensor::from_pbo(p).map(Self::F64)),
944 }
945 }
946
947 pub unsafe fn from_client_pbo(
960 state: crate::EfClientState,
961 map_fn: crate::EfPboMapFn,
962 unmap_fn: crate::EfPboUnmapFn,
963 buffer_id: u32,
964 size: usize,
965 shape: &[usize],
966 dtype: DType,
967 ) -> crate::Result<TensorDyn> {
968 let ops = unsafe { crate::client_state_pbo_ops(state, map_fn, unmap_fn) }?;
970 Self::from_pbo_import(buffer_id, size, shape, dtype, ops)
971 }
972
973 pub unsafe fn from_raw_host(
991 ptr: *mut u8,
992 shape: &[usize],
993 dtype: DType,
994 name: Option<&str>,
995 ) -> crate::Result<Self> {
996 unsafe { Self::from_foreign_ptr(ptr, shape, dtype, None, name) }
997 }
998
999 pub unsafe fn from_raw_host_with_capacity(
1015 ptr: *mut u8,
1016 shape: &[usize],
1017 capacity_bytes: usize,
1018 dtype: DType,
1019 name: Option<&str>,
1020 ) -> crate::Result<Self> {
1021 unsafe {
1022 Self::from_foreign_ptr_with_capacity(ptr, shape, dtype, capacity_bytes, None, name)
1023 }
1024 }
1025
1026 pub(crate) fn import_storage(
1038 desc: &crate::TensorDesc,
1039 shape: &[usize],
1040 dtype: DType,
1041 ) -> crate::Result<Self> {
1042 match desc.kind {
1043 #[cfg(target_os = "linux")]
1044 crate::protocol::kind::DMABUF => {
1045 let owned = crate::protocol::dup_descriptor_fd(desc)?;
1050 Self::from_fd(owned, shape, dtype, None)
1051 }
1052 #[cfg(not(target_os = "linux"))]
1053 crate::protocol::kind::DMABUF => Err(crate::Error::NotImplemented(
1054 "dma-buf import off Linux".into(),
1055 )),
1056 #[cfg(any(target_os = "macos", target_os = "ios"))]
1057 crate::protocol::kind::IOSURFACE => {
1058 let id = crate::protocol::descriptor_surface_id(desc)?;
1059 let surface_ref = crate::iosurface::lookup_by_id(id).ok_or_else(|| {
1060 crate::Error::InvalidArgument(format!(
1061 "IOSurface id {id} is not live (the producer's capsule may \
1062 already have been released)"
1063 ))
1064 })?;
1065 let result = unsafe { Self::from_iosurface(surface_ref, shape, dtype, None) };
1070 crate::iosurface::release(surface_ref);
1071 result
1072 }
1073 #[cfg(not(any(target_os = "macos", target_os = "ios")))]
1074 crate::protocol::kind::IOSURFACE => Err(crate::Error::NotImplemented(
1075 "IOSurface import off Apple platforms".into(),
1076 )),
1077 crate::protocol::kind::HOST => {
1078 crate::protocol::check_descriptor_host_ptr(desc)?;
1079 unsafe {
1086 Self::from_raw_host_with_capacity(
1087 desc.ptr.0,
1088 shape,
1089 desc.capacity as usize,
1090 dtype,
1091 None,
1092 )
1093 }
1094 }
1095 crate::protocol::kind::PBO => {
1096 let buffer_id = crate::protocol::descriptor_pbo_buffer_id(desc)?;
1097 let parts = unsafe { crate::read_pbo_vtable_parts(desc.ptr.0 as *const _) }?;
1107 unsafe {
1110 Self::from_client_pbo(
1111 parts.state,
1112 parts.map_fn,
1113 parts.unmap_fn,
1114 buffer_id,
1115 desc.capacity as usize,
1116 shape,
1117 dtype,
1118 )
1119 }
1120 }
1121 #[cfg(target_os = "windows")]
1122 crate::protocol::kind::D3D11_TEXTURE => {
1123 let (tex, completion) = crate::protocol::descriptor_d3d11_handles(desc)?;
1124 unsafe {
1131 let (format, width, height) =
1132 crate::protocol::descriptor_d3d11_geometry(desc, tex, shape)?;
1133 Self::from_d3d11_shared_handle(
1134 tex,
1135 width,
1136 height,
1137 format,
1138 dtype,
1139 crate::CpuAccess::ReadWrite,
1140 completion,
1141 None,
1142 )
1143 }
1144 }
1145 #[cfg(not(target_os = "windows"))]
1146 crate::protocol::kind::D3D11_TEXTURE => Err(crate::Error::NotImplemented(
1147 "D3D11 texture import off Windows".into(),
1148 )),
1149 k => Err(crate::Error::NotImplemented(format!(
1150 "tensor interop kind {k} cannot be imported by this build"
1151 ))),
1152 }
1153 }
1154
1155 pub unsafe fn from_foreign_ptr(
1168 ptr: *mut u8,
1169 shape: &[usize],
1170 dtype: DType,
1171 owner: Option<crate::ForeignOwner>,
1172 name: Option<&str>,
1173 ) -> crate::Result<Self> {
1174 unsafe {
1179 match dtype {
1180 DType::U8 => {
1181 Tensor::<u8>::from_foreign(ptr.cast(), shape, owner, name).map(Self::U8)
1182 }
1183 DType::I8 => {
1184 Tensor::<i8>::from_foreign(ptr.cast(), shape, owner, name).map(Self::I8)
1185 }
1186 DType::U16 => {
1187 Tensor::<u16>::from_foreign(ptr.cast(), shape, owner, name).map(Self::U16)
1188 }
1189 DType::I16 => {
1190 Tensor::<i16>::from_foreign(ptr.cast(), shape, owner, name).map(Self::I16)
1191 }
1192 DType::U32 => {
1193 Tensor::<u32>::from_foreign(ptr.cast(), shape, owner, name).map(Self::U32)
1194 }
1195 DType::I32 => {
1196 Tensor::<i32>::from_foreign(ptr.cast(), shape, owner, name).map(Self::I32)
1197 }
1198 DType::U64 => {
1199 Tensor::<u64>::from_foreign(ptr.cast(), shape, owner, name).map(Self::U64)
1200 }
1201 DType::I64 => {
1202 Tensor::<i64>::from_foreign(ptr.cast(), shape, owner, name).map(Self::I64)
1203 }
1204 DType::F16 => {
1205 Tensor::<f16>::from_foreign(ptr.cast(), shape, owner, name).map(Self::F16)
1206 }
1207 DType::F32 => {
1208 Tensor::<f32>::from_foreign(ptr.cast(), shape, owner, name).map(Self::F32)
1209 }
1210 DType::F64 => {
1211 Tensor::<f64>::from_foreign(ptr.cast(), shape, owner, name).map(Self::F64)
1212 }
1213 }
1214 }
1215 }
1216
1217 pub unsafe fn from_foreign_ptr_with_capacity(
1228 ptr: *mut u8,
1229 shape: &[usize],
1230 dtype: DType,
1231 capacity_bytes: usize,
1232 owner: Option<crate::ForeignOwner>,
1233 name: Option<&str>,
1234 ) -> crate::Result<Self> {
1235 macro_rules! import {
1236 ($t:ty, $variant:ident) => {
1237 unsafe {
1242 Tensor::<$t>::from_foreign_with_capacity(
1243 ptr.cast(),
1244 shape,
1245 capacity_bytes,
1246 owner,
1247 name,
1248 )
1249 }
1250 .map(Self::$variant)
1251 };
1252 }
1253 match dtype {
1254 DType::U8 => import!(u8, U8),
1255 DType::I8 => import!(i8, I8),
1256 DType::U16 => import!(u16, U16),
1257 DType::I16 => import!(i16, I16),
1258 DType::U32 => import!(u32, U32),
1259 DType::I32 => import!(i32, I32),
1260 DType::U64 => import!(u64, U64),
1261 DType::I64 => import!(i64, I64),
1262 DType::F16 => import!(f16, F16),
1263 DType::F32 => import!(f32, F32),
1264 DType::F64 => import!(f64, F64),
1265 }
1266 }
1267
1268 #[cfg(any(target_os = "macos", target_os = "ios"))]
1279 pub fn from_iosurface_id(
1280 id: u32,
1281 shape: &[usize],
1282 dtype: DType,
1283 name: Option<&str>,
1284 ) -> crate::Result<Self> {
1285 let surface = crate::iosurface::lookup_by_id(id).ok_or_else(|| {
1286 crate::Error::InvalidArgument(format!("no live IOSurface with id {id}"))
1287 })?;
1288 let result = unsafe { Self::from_iosurface(surface, shape, dtype, name) };
1303 crate::iosurface::release(surface);
1304 result
1305 }
1306
1307 #[cfg(any(target_os = "macos", target_os = "ios"))]
1315 pub unsafe fn from_iosurface(
1316 surface_ref: *mut std::ffi::c_void,
1317 shape: &[usize],
1318 dtype: DType,
1319 name: Option<&str>,
1320 ) -> crate::Result<Self> {
1321 unsafe {
1322 match dtype {
1323 DType::U8 => Tensor::<u8>::from_iosurface(surface_ref, shape, name).map(Self::U8),
1324 DType::I8 => Tensor::<i8>::from_iosurface(surface_ref, shape, name).map(Self::I8),
1325 DType::U16 => {
1326 Tensor::<u16>::from_iosurface(surface_ref, shape, name).map(Self::U16)
1327 }
1328 DType::I16 => {
1329 Tensor::<i16>::from_iosurface(surface_ref, shape, name).map(Self::I16)
1330 }
1331 DType::U32 => {
1332 Tensor::<u32>::from_iosurface(surface_ref, shape, name).map(Self::U32)
1333 }
1334 DType::I32 => {
1335 Tensor::<i32>::from_iosurface(surface_ref, shape, name).map(Self::I32)
1336 }
1337 DType::U64 => {
1338 Tensor::<u64>::from_iosurface(surface_ref, shape, name).map(Self::U64)
1339 }
1340 DType::I64 => {
1341 Tensor::<i64>::from_iosurface(surface_ref, shape, name).map(Self::I64)
1342 }
1343 DType::F16 => {
1344 Tensor::<f16>::from_iosurface(surface_ref, shape, name).map(Self::F16)
1345 }
1346 DType::F32 => {
1347 Tensor::<f32>::from_iosurface(surface_ref, shape, name).map(Self::F32)
1348 }
1349 DType::F64 => {
1350 Tensor::<f64>::from_iosurface(surface_ref, shape, name).map(Self::F64)
1351 }
1352 }
1353 }
1354 }
1355
1356 #[cfg(any(target_os = "macos", target_os = "ios"))]
1359 pub fn iosurface_id(&self) -> Option<u32> {
1360 dispatch!(self, iosurface_id)
1361 }
1362
1363 pub fn pbo_id(&self) -> Option<u32> {
1366 dispatch!(self, pbo_id)
1367 }
1368
1369 pub fn pbo_is_mapped(&self) -> Option<bool> {
1373 dispatch!(self, pbo_is_mapped)
1374 }
1375
1376 pub fn pbo_vtable_ptr(&self) -> Option<*const std::ffi::c_void> {
1380 dispatch!(self, pbo_vtable_ptr)
1381 }
1382
1383 pub fn pbo_keepalive(&self) -> Option<std::sync::Arc<dyn Send + Sync>> {
1388 dispatch!(self, pbo_keepalive)
1389 }
1390
1391 #[cfg(any(target_os = "macos", target_os = "ios"))]
1395 pub fn iosurface_ref(&self) -> Option<*mut std::ffi::c_void> {
1396 dispatch!(self, iosurface_ref)
1397 }
1398
1399 #[cfg(any(target_os = "macos", target_os = "ios"))]
1404 pub fn iosurface_physical_dims(&self) -> Option<(usize, usize)> {
1405 dispatch!(self, iosurface_physical_dims)
1406 }
1407
1408 #[cfg(target_os = "windows")]
1411 pub fn d3d11_texture(&self) -> Option<*mut std::ffi::c_void> {
1412 dispatch!(self, d3d11_texture)
1413 }
1414
1415 #[cfg(target_os = "windows")]
1418 pub fn d3d11_layout(&self) -> Option<crate::d3d11_layout::D3d11ImageLayout> {
1419 dispatch!(self, d3d11_layout)
1420 }
1421
1422 #[cfg(target_os = "windows")]
1425 pub fn d3d11_shared_handle(&self) -> crate::Result<std::os::windows::io::OwnedHandle> {
1426 dispatch!(self, d3d11_shared_handle)
1427 }
1428
1429 #[cfg(target_os = "windows")]
1432 pub fn gpu_completion(&self) -> crate::Result<Option<crate::d3d11::GpuCompletion>> {
1433 dispatch!(self, gpu_completion)
1434 }
1435
1436 pub fn gpu_write_value(&self) -> u64 {
1440 dispatch!(self, gpu_write_value)
1441 }
1442
1443 #[cfg(target_os = "windows")]
1447 pub fn set_gpu_write(&self, value: u64) -> crate::Result<()> {
1448 dispatch!(self, set_gpu_write, value)
1449 }
1450
1451 #[cfg(target_os = "windows")]
1454 pub(crate) fn d3d11_shared_handle_value(&self) -> Option<usize> {
1455 dispatch!(self, d3d11_shared_handle_value)
1456 }
1457
1458 #[cfg(target_os = "windows")]
1466 #[allow(clippy::too_many_arguments)] pub unsafe fn from_d3d11_texture(
1468 texture: *mut std::ffi::c_void,
1469 width: usize,
1470 height: usize,
1471 format: PixelFormat,
1472 dtype: DType,
1473 access: crate::CpuAccess,
1474 name: Option<&str>,
1475 ) -> crate::Result<Self> {
1476 unsafe {
1479 match dtype {
1480 DType::U8 => {
1481 Tensor::<u8>::from_d3d11_texture(texture, width, height, format, access, name)
1482 .map(Self::U8)
1483 }
1484 DType::I8 => {
1485 Tensor::<i8>::from_d3d11_texture(texture, width, height, format, access, name)
1486 .map(Self::I8)
1487 }
1488 DType::U16 => {
1489 Tensor::<u16>::from_d3d11_texture(texture, width, height, format, access, name)
1490 .map(Self::U16)
1491 }
1492 DType::I16 => {
1493 Tensor::<i16>::from_d3d11_texture(texture, width, height, format, access, name)
1494 .map(Self::I16)
1495 }
1496 DType::U32 => {
1497 Tensor::<u32>::from_d3d11_texture(texture, width, height, format, access, name)
1498 .map(Self::U32)
1499 }
1500 DType::I32 => {
1501 Tensor::<i32>::from_d3d11_texture(texture, width, height, format, access, name)
1502 .map(Self::I32)
1503 }
1504 DType::U64 => {
1505 Tensor::<u64>::from_d3d11_texture(texture, width, height, format, access, name)
1506 .map(Self::U64)
1507 }
1508 DType::I64 => {
1509 Tensor::<i64>::from_d3d11_texture(texture, width, height, format, access, name)
1510 .map(Self::I64)
1511 }
1512 DType::F16 => {
1513 Tensor::<f16>::from_d3d11_texture(texture, width, height, format, access, name)
1514 .map(Self::F16)
1515 }
1516 DType::F32 => {
1517 Tensor::<f32>::from_d3d11_texture(texture, width, height, format, access, name)
1518 .map(Self::F32)
1519 }
1520 DType::F64 => {
1521 Tensor::<f64>::from_d3d11_texture(texture, width, height, format, access, name)
1522 .map(Self::F64)
1523 }
1524 }
1525 }
1526 }
1527
1528 #[cfg(target_os = "windows")]
1537 #[allow(clippy::too_many_arguments)] pub unsafe fn from_d3d11_shared_handle(
1539 handle: std::os::windows::io::RawHandle,
1540 width: usize,
1541 height: usize,
1542 format: PixelFormat,
1543 dtype: DType,
1544 access: crate::CpuAccess,
1545 completion: Option<(std::os::windows::io::RawHandle, u64)>,
1546 name: Option<&str>,
1547 ) -> crate::Result<Self> {
1548 unsafe {
1551 match dtype {
1552 DType::U8 => Tensor::<u8>::from_d3d11_shared_handle(
1553 handle, width, height, format, access, completion, name,
1554 )
1555 .map(Self::U8),
1556 DType::I8 => Tensor::<i8>::from_d3d11_shared_handle(
1557 handle, width, height, format, access, completion, name,
1558 )
1559 .map(Self::I8),
1560 DType::U16 => Tensor::<u16>::from_d3d11_shared_handle(
1561 handle, width, height, format, access, completion, name,
1562 )
1563 .map(Self::U16),
1564 DType::I16 => Tensor::<i16>::from_d3d11_shared_handle(
1565 handle, width, height, format, access, completion, name,
1566 )
1567 .map(Self::I16),
1568 DType::U32 => Tensor::<u32>::from_d3d11_shared_handle(
1569 handle, width, height, format, access, completion, name,
1570 )
1571 .map(Self::U32),
1572 DType::I32 => Tensor::<i32>::from_d3d11_shared_handle(
1573 handle, width, height, format, access, completion, name,
1574 )
1575 .map(Self::I32),
1576 DType::U64 => Tensor::<u64>::from_d3d11_shared_handle(
1577 handle, width, height, format, access, completion, name,
1578 )
1579 .map(Self::U64),
1580 DType::I64 => Tensor::<i64>::from_d3d11_shared_handle(
1581 handle, width, height, format, access, completion, name,
1582 )
1583 .map(Self::I64),
1584 DType::F16 => Tensor::<f16>::from_d3d11_shared_handle(
1585 handle, width, height, format, access, completion, name,
1586 )
1587 .map(Self::F16),
1588 DType::F32 => Tensor::<f32>::from_d3d11_shared_handle(
1589 handle, width, height, format, access, completion, name,
1590 )
1591 .map(Self::F32),
1592 DType::F64 => Tensor::<f64>::from_d3d11_shared_handle(
1593 handle, width, height, format, access, completion, name,
1594 )
1595 .map(Self::F64),
1596 }
1597 }
1598 }
1599
1600 #[cfg(target_os = "android")]
1610 pub unsafe fn from_hardware_buffer(
1611 buffer_ptr: *mut std::ffi::c_void,
1612 shape: &[usize],
1613 dtype: DType,
1614 name: Option<&str>,
1615 ) -> crate::Result<Self> {
1616 unsafe {
1617 match dtype {
1618 DType::U8 => {
1619 Tensor::<u8>::from_hardware_buffer(buffer_ptr, shape, name).map(Self::U8)
1620 }
1621 DType::I8 => {
1622 Tensor::<i8>::from_hardware_buffer(buffer_ptr, shape, name).map(Self::I8)
1623 }
1624 DType::U16 => {
1625 Tensor::<u16>::from_hardware_buffer(buffer_ptr, shape, name).map(Self::U16)
1626 }
1627 DType::I16 => {
1628 Tensor::<i16>::from_hardware_buffer(buffer_ptr, shape, name).map(Self::I16)
1629 }
1630 DType::U32 => {
1631 Tensor::<u32>::from_hardware_buffer(buffer_ptr, shape, name).map(Self::U32)
1632 }
1633 DType::I32 => {
1634 Tensor::<i32>::from_hardware_buffer(buffer_ptr, shape, name).map(Self::I32)
1635 }
1636 DType::U64 => {
1637 Tensor::<u64>::from_hardware_buffer(buffer_ptr, shape, name).map(Self::U64)
1638 }
1639 DType::I64 => {
1640 Tensor::<i64>::from_hardware_buffer(buffer_ptr, shape, name).map(Self::I64)
1641 }
1642 DType::F16 => {
1643 Tensor::<f16>::from_hardware_buffer(buffer_ptr, shape, name).map(Self::F16)
1644 }
1645 DType::F32 => {
1646 Tensor::<f32>::from_hardware_buffer(buffer_ptr, shape, name).map(Self::F32)
1647 }
1648 DType::F64 => {
1649 Tensor::<f64>::from_hardware_buffer(buffer_ptr, shape, name).map(Self::F64)
1650 }
1651 }
1652 }
1653 }
1654
1655 #[cfg(target_os = "android")]
1659 pub fn hardware_buffer_ptr(&self) -> Option<*mut std::ffi::c_void> {
1660 dispatch!(self, hardware_buffer_ptr)
1661 }
1662
1663 #[cfg(target_os = "android")]
1667 pub fn hardware_buffer_physical_dims(&self) -> Option<(usize, usize)> {
1668 dispatch!(self, hardware_buffer_physical_dims)
1669 }
1670
1671 pub fn copy_to_flat(&self, dst: &mut [u8]) -> crate::Result<()> {
1676 dispatch!(self, copy_to_flat, dst)
1677 }
1678
1679 pub fn image(
1697 width: usize,
1698 height: usize,
1699 format: PixelFormat,
1700 dtype: DType,
1701 memory: Option<TensorMemory>,
1702 access: crate::CpuAccess,
1703 ) -> crate::Result<Self> {
1704 match dtype {
1705 DType::U8 => Tensor::<u8>::image(width, height, format, memory, access).map(Self::U8),
1706 DType::I8 => Tensor::<i8>::image(width, height, format, memory, access).map(Self::I8),
1707 DType::U16 => {
1708 Tensor::<u16>::image(width, height, format, memory, access).map(Self::U16)
1709 }
1710 DType::I16 => {
1711 Tensor::<i16>::image(width, height, format, memory, access).map(Self::I16)
1712 }
1713 DType::U32 => {
1714 Tensor::<u32>::image(width, height, format, memory, access).map(Self::U32)
1715 }
1716 DType::I32 => {
1717 Tensor::<i32>::image(width, height, format, memory, access).map(Self::I32)
1718 }
1719 DType::U64 => {
1720 Tensor::<u64>::image(width, height, format, memory, access).map(Self::U64)
1721 }
1722 DType::I64 => {
1723 Tensor::<i64>::image(width, height, format, memory, access).map(Self::I64)
1724 }
1725 DType::F16 => {
1726 Tensor::<f16>::image(width, height, format, memory, access).map(Self::F16)
1727 }
1728 DType::F32 => {
1729 Tensor::<f32>::image(width, height, format, memory, access).map(Self::F32)
1730 }
1731 DType::F64 => {
1732 Tensor::<f64>::image(width, height, format, memory, access).map(Self::F64)
1733 }
1734 }
1735 }
1736
1737 pub fn image_desc(desc: &crate::ImageDesc) -> crate::Result<Self> {
1741 match desc.dtype() {
1742 DType::U8 => Tensor::<u8>::image_desc(desc).map(Self::U8),
1743 DType::I8 => Tensor::<i8>::image_desc(desc).map(Self::I8),
1744 DType::U16 => Tensor::<u16>::image_desc(desc).map(Self::U16),
1745 DType::I16 => Tensor::<i16>::image_desc(desc).map(Self::I16),
1746 DType::U32 => Tensor::<u32>::image_desc(desc).map(Self::U32),
1747 DType::I32 => Tensor::<i32>::image_desc(desc).map(Self::I32),
1748 DType::U64 => Tensor::<u64>::image_desc(desc).map(Self::U64),
1749 DType::I64 => Tensor::<i64>::image_desc(desc).map(Self::I64),
1750 DType::F16 => Tensor::<f16>::image_desc(desc).map(Self::F16),
1751 DType::F32 => Tensor::<f32>::image_desc(desc).map(Self::F32),
1752 DType::F64 => Tensor::<f64>::image_desc(desc).map(Self::F64),
1753 }
1754 }
1755
1756 pub fn compression(&self) -> Option<crate::CompressionScheme> {
1759 dispatch!(self, compression)
1760 }
1761
1762 pub fn image_with_stride(
1789 width: usize,
1790 height: usize,
1791 format: PixelFormat,
1792 dtype: DType,
1793 row_stride_bytes: usize,
1794 memory: Option<TensorMemory>,
1795 access: crate::CpuAccess,
1796 ) -> crate::Result<Self> {
1797 match dtype {
1798 DType::U8 => Tensor::<u8>::image_with_stride(
1799 width,
1800 height,
1801 format,
1802 row_stride_bytes,
1803 memory,
1804 access,
1805 )
1806 .map(Self::U8),
1807 DType::I8 => Tensor::<i8>::image_with_stride(
1808 width,
1809 height,
1810 format,
1811 row_stride_bytes,
1812 memory,
1813 access,
1814 )
1815 .map(Self::I8),
1816 DType::U16 => Tensor::<u16>::image_with_stride(
1817 width,
1818 height,
1819 format,
1820 row_stride_bytes,
1821 memory,
1822 access,
1823 )
1824 .map(Self::U16),
1825 DType::I16 => Tensor::<i16>::image_with_stride(
1826 width,
1827 height,
1828 format,
1829 row_stride_bytes,
1830 memory,
1831 access,
1832 )
1833 .map(Self::I16),
1834 DType::U32 => Tensor::<u32>::image_with_stride(
1835 width,
1836 height,
1837 format,
1838 row_stride_bytes,
1839 memory,
1840 access,
1841 )
1842 .map(Self::U32),
1843 DType::I32 => Tensor::<i32>::image_with_stride(
1844 width,
1845 height,
1846 format,
1847 row_stride_bytes,
1848 memory,
1849 access,
1850 )
1851 .map(Self::I32),
1852 DType::U64 => Tensor::<u64>::image_with_stride(
1853 width,
1854 height,
1855 format,
1856 row_stride_bytes,
1857 memory,
1858 access,
1859 )
1860 .map(Self::U64),
1861 DType::I64 => Tensor::<i64>::image_with_stride(
1862 width,
1863 height,
1864 format,
1865 row_stride_bytes,
1866 memory,
1867 access,
1868 )
1869 .map(Self::I64),
1870 DType::F16 => Tensor::<f16>::image_with_stride(
1871 width,
1872 height,
1873 format,
1874 row_stride_bytes,
1875 memory,
1876 access,
1877 )
1878 .map(Self::F16),
1879 DType::F32 => Tensor::<f32>::image_with_stride(
1880 width,
1881 height,
1882 format,
1883 row_stride_bytes,
1884 memory,
1885 access,
1886 )
1887 .map(Self::F32),
1888 DType::F64 => Tensor::<f64>::image_with_stride(
1889 width,
1890 height,
1891 format,
1892 row_stride_bytes,
1893 memory,
1894 access,
1895 )
1896 .map(Self::F64),
1897 }
1898 }
1899}
1900
1901impl From<Tensor<u8>> for TensorDyn {
1904 fn from(t: Tensor<u8>) -> Self {
1905 Self::U8(t)
1906 }
1907}
1908
1909impl From<Tensor<i8>> for TensorDyn {
1910 fn from(t: Tensor<i8>) -> Self {
1911 Self::I8(t)
1912 }
1913}
1914
1915impl From<Tensor<u16>> for TensorDyn {
1916 fn from(t: Tensor<u16>) -> Self {
1917 Self::U16(t)
1918 }
1919}
1920
1921impl From<Tensor<i16>> for TensorDyn {
1922 fn from(t: Tensor<i16>) -> Self {
1923 Self::I16(t)
1924 }
1925}
1926
1927impl From<Tensor<u32>> for TensorDyn {
1928 fn from(t: Tensor<u32>) -> Self {
1929 Self::U32(t)
1930 }
1931}
1932
1933impl From<Tensor<i32>> for TensorDyn {
1934 fn from(t: Tensor<i32>) -> Self {
1935 Self::I32(t)
1936 }
1937}
1938
1939impl From<Tensor<u64>> for TensorDyn {
1940 fn from(t: Tensor<u64>) -> Self {
1941 Self::U64(t)
1942 }
1943}
1944
1945impl From<Tensor<i64>> for TensorDyn {
1946 fn from(t: Tensor<i64>) -> Self {
1947 Self::I64(t)
1948 }
1949}
1950
1951impl From<Tensor<f16>> for TensorDyn {
1952 fn from(t: Tensor<f16>) -> Self {
1953 Self::F16(t)
1954 }
1955}
1956
1957impl From<Tensor<f32>> for TensorDyn {
1958 fn from(t: Tensor<f32>) -> Self {
1959 Self::F32(t)
1960 }
1961}
1962
1963impl From<Tensor<f64>> for TensorDyn {
1964 fn from(t: Tensor<f64>) -> Self {
1965 Self::F64(t)
1966 }
1967}
1968
1969impl fmt::Debug for TensorDyn {
1970 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
1971 dispatch!(self, fmt, f)
1972 }
1973}
1974
1975pub type Raw = *mut std::ffi::c_void;
1978
1979impl TensorDyn {
1980 pub fn into_raw(self) -> Raw {
1983 Box::into_raw(Box::new(self)) as Raw
1984 }
1985
1986 pub unsafe fn from_raw(p: Raw) -> TensorDyn {
1992 *unsafe { Box::from_raw(p as *mut TensorDyn) }
1993 }
1994
1995 pub unsafe fn with_raw<R>(p: Raw, f: impl FnOnce(&mut TensorDyn) -> R) -> R {
2001 let mut t = std::mem::ManuallyDrop::new(unsafe { std::ptr::read(p as *mut TensorDyn) });
2002 f(&mut t)
2003 }
2004}
2005
2006#[cfg(test)]
2007mod tests {
2008 use super::*;
2009
2010 #[test]
2011 fn from_typed_tensor() {
2012 let t = Tensor::<u8>::new(&[10], None, None).unwrap();
2013 let dyn_t: TensorDyn = t.into();
2014 assert_eq!(dyn_t.dtype(), DType::U8);
2015 assert_eq!(dyn_t.shape(), &[10]);
2016 }
2017
2018 #[test]
2019 fn from_foreign_ptr_wraps_borrowed_memory() {
2020 use crate::TensorMapTrait;
2021 let mut vec: Vec<f32> = vec![0.0; 4];
2024 let ptr = vec.as_mut_ptr() as *mut u8;
2025 let owner: crate::ForeignOwner = Box::new(vec);
2026 let t = unsafe {
2027 TensorDyn::from_foreign_ptr(ptr, &[2, 2], DType::F32, Some(owner), Some("trt_output"))
2028 }
2029 .unwrap();
2030 assert_eq!(t.dtype(), DType::F32);
2031 assert_eq!(t.memory(), TensorMemory::Mem);
2032 assert_eq!(t.shape(), &[2, 2]);
2033 {
2034 let mut m = t.as_f32().unwrap().map().unwrap();
2035 m.as_mut_slice().copy_from_slice(&[1.0, 2.0, 3.0, 4.0]);
2036 }
2037 let m = t.as_f32().unwrap().map().unwrap();
2038 assert_eq!(m.as_slice(), &[1.0, 2.0, 3.0, 4.0]);
2039 }
2040
2041 #[test]
2050 fn from_foreign_ptr_rejects_null_ptr() {
2051 let err = unsafe {
2052 TensorDyn::from_foreign_ptr(std::ptr::null_mut(), &[4], DType::U8, None, None)
2053 }
2054 .unwrap_err();
2055 assert!(
2057 matches!(err, crate::error::Error::InvalidArgument(ref m) if m.contains("non-null")),
2058 "expected InvalidArgument(non-null), got {err:?}"
2059 );
2060 }
2061
2062 #[test]
2063 fn from_foreign_ptr_rejects_empty_shape() {
2064 let mut dummy: u8 = 0;
2065 let err = unsafe {
2066 TensorDyn::from_foreign_ptr(&mut dummy as *mut u8, &[], DType::U8, None, None)
2067 }
2068 .unwrap_err();
2069 assert!(
2070 matches!(err, crate::error::Error::InvalidSize(0)),
2071 "expected InvalidSize(0) for empty shape, got {err:?}"
2072 );
2073 }
2074
2075 #[test]
2076 fn from_foreign_ptr_rejects_overflow_shape() {
2077 let mut dummy: u8 = 0;
2078 let huge = [usize::MAX / 2 + 1, 2];
2079 let err = unsafe { TensorDyn::from_foreign_ptr(&mut dummy, &huge, DType::U8, None, None) }
2080 .unwrap_err();
2081 assert!(
2082 matches!(err, crate::error::Error::InvalidArgument(ref m) if m.contains("overflow")),
2083 "expected InvalidArgument(overflow), got {err:?}"
2084 );
2085 }
2086
2087 #[test]
2088 fn from_foreign_ptr_u8_dtype_dispatch() {
2089 let mut buf: Vec<u8> = vec![1, 2, 3, 4];
2092 let ptr = buf.as_mut_ptr();
2093 let owner: crate::ForeignOwner = Box::new(buf);
2094 let t = unsafe {
2095 TensorDyn::from_foreign_ptr(ptr, &[4], DType::U8, Some(owner), Some("u8_foreign"))
2096 }
2097 .unwrap();
2098 assert_eq!(t.dtype(), DType::U8);
2099 assert_eq!(t.shape(), &[4]);
2100 let m = t.as_u8().unwrap().map().unwrap();
2101 use crate::TensorMapTrait;
2102 assert_eq!(m.as_slice(), &[1u8, 2, 3, 4]);
2103 }
2104
2105 #[test]
2106 fn map_bytes_roundtrips_and_is_typed_agnostic() {
2107 use crate::TensorMapTrait;
2108 let t = Tensor::<f32>::new(&[4], None, None).unwrap();
2111 let dyn_t: TensorDyn = t.into();
2112 {
2113 let mut view = dyn_t.map_bytes(crate::CpuAccess::ReadWrite).unwrap();
2114 assert_eq!(view.len(), 16, "4 x f32 is 16 bytes, not 4");
2115 view.as_mut_slice()
2116 .copy_from_slice(&[1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16]);
2117 }
2118 let view = dyn_t.map_bytes(crate::CpuAccess::Read).unwrap();
2119 assert_eq!(
2120 view.as_slice(),
2121 &[1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16]
2122 );
2123 }
2124
2125 #[test]
2126 fn map_bytes_none_is_invalid() {
2127 let t = Tensor::<u8>::new(&[4], None, None).unwrap();
2128 let dyn_t: TensorDyn = t.into();
2129 assert!(matches!(
2130 dyn_t.map_bytes(crate::CpuAccess::None),
2131 Err(crate::error::Error::InvalidArgument(_))
2132 ));
2133 }
2134
2135 #[test]
2136 fn downcast_ref() {
2137 let t = Tensor::<u8>::new(&[10], None, None).unwrap();
2138 let dyn_t: TensorDyn = t.into();
2139 assert!(dyn_t.as_u8().is_some());
2140 assert!(dyn_t.as_i8().is_none());
2141 }
2142
2143 #[test]
2144 fn downcast_into() {
2145 let t = Tensor::<u8>::new(&[10], None, None).unwrap();
2146 let dyn_t: TensorDyn = t.into();
2147 let back = dyn_t.into_u8().unwrap();
2148 assert_eq!(back.shape(), &[10]);
2149 }
2150
2151 #[test]
2152 fn image_accessors() {
2153 let t = Tensor::<u8>::image(
2154 640,
2155 480,
2156 PixelFormat::Rgba,
2157 None,
2158 crate::CpuAccess::ReadWrite,
2159 )
2160 .unwrap();
2161 let dyn_t: TensorDyn = t.into();
2162 assert_eq!(dyn_t.format(), Some(PixelFormat::Rgba));
2163 assert_eq!(dyn_t.width(), Some(640));
2164 assert_eq!(dyn_t.height(), Some(480));
2165 assert!(!dyn_t.is_multiplane());
2166 }
2167
2168 #[test]
2169 fn cpu_access_reports_the_declared_access() {
2170 let read_t =
2172 Tensor::<u8>::image(4, 4, PixelFormat::Rgba, None, crate::CpuAccess::Read).unwrap();
2173 let dyn_read: TensorDyn = read_t.into();
2174 assert_eq!(dyn_read.cpu_access(), crate::CpuAccess::Read);
2175 assert!(!dyn_read.cpu_access().writes());
2176
2177 let mem_t = Tensor::<u8>::new(&[4], None, None).unwrap();
2181 let dyn_mem: TensorDyn = mem_t.into();
2182 assert_eq!(dyn_mem.cpu_access(), crate::CpuAccess::ReadWrite);
2183 assert!(dyn_mem.cpu_access().writes());
2184 }
2185
2186 #[test]
2187 fn image_constructor() {
2188 let dyn_t = TensorDyn::image(
2189 640,
2190 480,
2191 PixelFormat::Rgb,
2192 DType::U8,
2193 None,
2194 crate::CpuAccess::ReadWrite,
2195 )
2196 .unwrap();
2197 assert_eq!(dyn_t.dtype(), DType::U8);
2198 assert_eq!(dyn_t.format(), Some(PixelFormat::Rgb));
2199 assert_eq!(dyn_t.width(), Some(640));
2200 }
2201
2202 #[test]
2203 fn image_constructor_i8() {
2204 let dyn_t = TensorDyn::image(
2205 640,
2206 480,
2207 PixelFormat::Rgb,
2208 DType::I8,
2209 None,
2210 crate::CpuAccess::ReadWrite,
2211 )
2212 .unwrap();
2213 assert_eq!(dyn_t.dtype(), DType::I8);
2214 assert_eq!(dyn_t.format(), Some(PixelFormat::Rgb));
2215 }
2216
2217 #[test]
2218 fn set_format_packed() {
2219 let mut t = TensorDyn::new(&[480, 640, 3], DType::U8, None, None).unwrap();
2220 assert_eq!(t.format(), None);
2221 t.set_format(PixelFormat::Rgb).unwrap();
2222 assert_eq!(t.format(), Some(PixelFormat::Rgb));
2223 assert_eq!(t.width(), Some(640));
2224 assert_eq!(t.height(), Some(480));
2225 }
2226
2227 #[test]
2228 fn set_format_planar() {
2229 let mut t = TensorDyn::new(&[3, 480, 640], DType::U8, None, None).unwrap();
2230 t.set_format(PixelFormat::PlanarRgb).unwrap();
2231 assert_eq!(t.format(), Some(PixelFormat::PlanarRgb));
2232 assert_eq!(t.width(), Some(640));
2233 assert_eq!(t.height(), Some(480));
2234 }
2235
2236 #[test]
2237 fn set_format_rejects_wrong_shape() {
2238 let mut t = TensorDyn::new(&[480, 640, 4], DType::U8, None, None).unwrap();
2239 assert!(t.set_format(PixelFormat::Rgb).is_err());
2240 }
2241
2242 #[test]
2243 fn with_format_builder() {
2244 let t = TensorDyn::new(&[480, 640, 4], DType::U8, None, None)
2245 .unwrap()
2246 .with_format(PixelFormat::Rgba)
2247 .unwrap();
2248 assert_eq!(t.format(), Some(PixelFormat::Rgba));
2249 assert_eq!(t.width(), Some(640));
2250 assert_eq!(t.height(), Some(480));
2251 }
2252
2253 #[cfg(target_os = "linux")]
2254 #[test]
2255 fn dmabuf_clone_mem_tensor_fails() {
2256 let t = TensorDyn::new(&[480, 640, 3], DType::U8, Some(TensorMemory::Mem), None).unwrap();
2257 assert_eq!(t.memory(), TensorMemory::Mem);
2258 assert!(t.dmabuf_clone().is_err());
2259 }
2260
2261 #[cfg(target_os = "linux")]
2262 #[test]
2263 fn dmabuf_mem_tensor_fails() {
2264 let t = TensorDyn::new(&[480, 640, 3], DType::U8, Some(TensorMemory::Mem), None).unwrap();
2265 assert!(t.dmabuf().is_err());
2266 }
2267
2268 #[test]
2269 fn set_format_semi_planar_nv12() {
2270 let mut t = TensorDyn::new(&[720, 640], DType::U8, Some(TensorMemory::Mem), None).unwrap();
2272 t.set_format(PixelFormat::Nv12).unwrap();
2273 assert_eq!(t.format(), Some(PixelFormat::Nv12));
2274 assert_eq!(t.width(), Some(640));
2275 assert_eq!(t.height(), Some(480));
2276 }
2277
2278 #[test]
2279 fn set_format_semi_planar_nv16() {
2280 let mut t = TensorDyn::new(&[960, 640], DType::U8, Some(TensorMemory::Mem), None).unwrap();
2282 t.set_format(PixelFormat::Nv16).unwrap();
2283 assert_eq!(t.format(), Some(PixelFormat::Nv16));
2284 assert_eq!(t.width(), Some(640));
2285 assert_eq!(t.height(), Some(480));
2286 }
2287
2288 #[test]
2289 fn with_format_rejects_wrong_shape() {
2290 let result = TensorDyn::new(&[480, 640, 4], DType::U8, None, None)
2291 .unwrap()
2292 .with_format(PixelFormat::Rgb);
2293 assert!(result.is_err());
2294 }
2295
2296 #[test]
2297 fn set_format_preserved_after_rejection() {
2298 let mut t = TensorDyn::new(&[480, 640, 3], DType::U8, None, None).unwrap();
2299 t.set_format(PixelFormat::Rgb).unwrap();
2300 assert_eq!(t.format(), Some(PixelFormat::Rgb));
2301
2302 assert!(t.set_format(PixelFormat::Rgba).is_err());
2304
2305 assert_eq!(t.format(), Some(PixelFormat::Rgb));
2307 }
2308
2309 #[test]
2310 fn set_format_idempotent() {
2311 let mut t = TensorDyn::new(&[480, 640, 3], DType::U8, None, None).unwrap();
2312 t.set_format(PixelFormat::Rgb).unwrap();
2313 t.set_format(PixelFormat::Rgb).unwrap();
2314 assert_eq!(t.format(), Some(PixelFormat::Rgb));
2315 assert_eq!(t.width(), Some(640));
2316 assert_eq!(t.height(), Some(480));
2317 }
2318
2319 #[test]
2322 fn set_row_stride_valid() {
2323 let mut t = TensorDyn::image(
2325 100,
2326 100,
2327 PixelFormat::Rgba,
2328 DType::U8,
2329 None,
2330 crate::CpuAccess::ReadWrite,
2331 )
2332 .unwrap();
2333 t.set_row_stride(512).unwrap();
2334 assert_eq!(t.row_stride(), Some(512));
2335 assert_eq!(t.effective_row_stride(), Some(512));
2336 }
2337
2338 #[test]
2339 fn set_row_stride_equals_min() {
2340 let mut t = TensorDyn::image(
2342 100,
2343 100,
2344 PixelFormat::Rgb,
2345 DType::U8,
2346 None,
2347 crate::CpuAccess::ReadWrite,
2348 )
2349 .unwrap();
2350 t.set_row_stride(300).unwrap();
2351 assert_eq!(t.row_stride(), Some(300));
2352 }
2353
2354 #[test]
2355 fn set_row_stride_too_small() {
2356 let mut t = TensorDyn::image(
2361 64,
2362 100,
2363 PixelFormat::Rgba,
2364 DType::U8,
2365 None,
2366 crate::CpuAccess::ReadWrite,
2367 )
2368 .unwrap();
2369 assert!(t.set_row_stride(200).is_err());
2370 assert_eq!(t.row_stride(), None);
2371 }
2372
2373 #[test]
2374 fn set_row_stride_zero() {
2375 let mut t = TensorDyn::image(
2376 100,
2377 100,
2378 PixelFormat::Rgb,
2379 DType::U8,
2380 None,
2381 crate::CpuAccess::ReadWrite,
2382 )
2383 .unwrap();
2384 assert!(t.set_row_stride(0).is_err());
2385 }
2386
2387 #[test]
2388 fn set_row_stride_requires_format() {
2389 let mut t = TensorDyn::new(&[480, 640, 3], DType::U8, None, None).unwrap();
2390 assert!(t.set_row_stride(2048).is_err());
2391 }
2392
2393 #[test]
2394 fn effective_row_stride_without_stride() {
2395 let t = TensorDyn::image(
2400 64,
2401 100,
2402 PixelFormat::Rgb,
2403 DType::U8,
2404 None,
2405 crate::CpuAccess::ReadWrite,
2406 )
2407 .unwrap();
2408 assert_eq!(t.row_stride(), None);
2409 assert_eq!(t.effective_row_stride(), Some(192)); }
2411
2412 #[test]
2413 fn effective_row_stride_padded_packed_dma() {
2414 let t = match TensorDyn::image(
2420 100,
2421 100,
2422 PixelFormat::Rgb,
2423 DType::U8,
2424 Some(TensorMemory::DmaBuf),
2425 crate::CpuAccess::ReadWrite,
2426 ) {
2427 Ok(t) if t.memory() == TensorMemory::DmaBuf => t,
2428 _ => return,
2429 };
2430 if cfg!(target_os = "android") {
2436 assert!(
2437 t.row_stride().is_none() || t.row_stride() >= Some(300),
2438 "AHardwareBuffer pitch must be absent or at least the tight row"
2439 );
2440 return;
2441 }
2442 if cfg!(target_os = "windows") {
2447 assert!(
2448 t.row_stride().is_none() || t.row_stride() >= Some(300),
2449 "D3D11 staging pitch must be absent or at least the tight row"
2450 );
2451 assert!(t.effective_row_stride() >= Some(300));
2452 return;
2453 }
2454 assert_eq!(t.row_stride(), Some(320));
2455 assert_eq!(t.effective_row_stride(), Some(320));
2456 }
2457
2458 #[test]
2459 fn effective_row_stride_no_format() {
2460 let t = TensorDyn::new(&[480, 640, 3], DType::U8, None, None).unwrap();
2461 assert_eq!(t.effective_row_stride(), None);
2462 }
2463
2464 #[test]
2465 fn with_row_stride_builder() {
2466 let t = TensorDyn::image(
2467 100,
2468 100,
2469 PixelFormat::Rgba,
2470 DType::U8,
2471 None,
2472 crate::CpuAccess::ReadWrite,
2473 )
2474 .unwrap()
2475 .with_row_stride(512)
2476 .unwrap();
2477 assert_eq!(t.row_stride(), Some(512));
2478 assert_eq!(t.effective_row_stride(), Some(512));
2479 }
2480
2481 #[test]
2482 fn with_row_stride_rejects_small() {
2483 let result = TensorDyn::image(
2484 100,
2485 100,
2486 PixelFormat::Rgba,
2487 DType::U8,
2488 None,
2489 crate::CpuAccess::ReadWrite,
2490 )
2491 .unwrap()
2492 .with_row_stride(200);
2493 assert!(result.is_err());
2494 }
2495
2496 #[test]
2497 fn set_format_clears_row_stride() {
2498 let mut t = TensorDyn::new(&[480, 640, 3], DType::U8, None, None).unwrap();
2499 t.set_format(PixelFormat::Rgb).unwrap();
2500 t.set_row_stride(2048).unwrap();
2501 assert_eq!(t.row_stride(), Some(2048));
2502
2503 let _ = t.set_format(PixelFormat::Bgra);
2505 assert_eq!(t.row_stride(), Some(2048));
2506
2507 t.set_format(PixelFormat::Rgb).unwrap();
2509 assert_eq!(t.row_stride(), Some(2048));
2510
2511 t.reshape(&[480 * 640 * 3]).unwrap();
2513 assert_eq!(t.row_stride(), None);
2514 assert_eq!(t.format(), None);
2515 }
2516
2517 #[test]
2518 fn set_format_different_compatible_clears_stride() {
2519 let mut t = TensorDyn::new(&[480, 640, 4], DType::U8, None, None).unwrap();
2522 t.set_format(PixelFormat::Rgba).unwrap();
2523 t.set_row_stride(4096).unwrap();
2524 assert_eq!(t.row_stride(), Some(4096));
2525
2526 t.set_format(PixelFormat::Bgra).unwrap();
2528 assert_eq!(t.format(), Some(PixelFormat::Bgra));
2529 assert_eq!(t.row_stride(), None);
2530 }
2531
2532 #[test]
2533 fn set_format_same_preserves_stride() {
2534 let mut t = TensorDyn::image(
2535 100,
2536 100,
2537 PixelFormat::Rgb,
2538 DType::U8,
2539 None,
2540 crate::CpuAccess::ReadWrite,
2541 )
2542 .unwrap();
2543 t.set_row_stride(512).unwrap();
2544 t.set_format(PixelFormat::Rgb).unwrap();
2546 assert_eq!(t.row_stride(), Some(512));
2547 }
2548
2549 #[test]
2550 fn effective_row_stride_planar() {
2551 let t = TensorDyn::image(
2552 640,
2553 480,
2554 PixelFormat::PlanarRgb,
2555 DType::U8,
2556 None,
2557 crate::CpuAccess::ReadWrite,
2558 )
2559 .unwrap();
2560 assert_eq!(t.effective_row_stride(), Some(640)); }
2562
2563 #[test]
2564 fn effective_row_stride_nv12() {
2565 let t = TensorDyn::image(
2566 640,
2567 480,
2568 PixelFormat::Nv12,
2569 DType::U8,
2570 None,
2571 crate::CpuAccess::ReadWrite,
2572 )
2573 .unwrap();
2574 assert_eq!(t.effective_row_stride(), Some(640)); }
2576
2577 #[test]
2578 fn map_rejects_strided_tensor() {
2579 let mut t = Tensor::<u8>::image(
2580 100,
2581 100,
2582 PixelFormat::Rgba,
2583 Some(TensorMemory::Mem),
2584 crate::CpuAccess::ReadWrite,
2585 )
2586 .unwrap();
2587 assert!(t.map().is_ok());
2589 t.set_row_stride(512).unwrap();
2591 let err = t.map();
2592 assert!(err.is_err());
2593 }
2594
2595 #[test]
2598 fn plane_offset_default_none() {
2599 let t = TensorDyn::image(
2600 100,
2601 100,
2602 PixelFormat::Rgba,
2603 DType::U8,
2604 None,
2605 crate::CpuAccess::ReadWrite,
2606 )
2607 .unwrap();
2608 assert_eq!(t.plane_offset(), None);
2609 }
2610
2611 #[test]
2612 fn set_plane_offset_basic() {
2613 let mut t = TensorDyn::image(
2614 100,
2615 100,
2616 PixelFormat::Rgba,
2617 DType::U8,
2618 None,
2619 crate::CpuAccess::ReadWrite,
2620 )
2621 .unwrap();
2622 t.set_plane_offset(4096);
2623 assert_eq!(t.plane_offset(), Some(4096));
2624 }
2625
2626 #[test]
2627 fn set_plane_offset_zero() {
2628 let mut t = TensorDyn::image(
2629 100,
2630 100,
2631 PixelFormat::Rgb,
2632 DType::U8,
2633 None,
2634 crate::CpuAccess::ReadWrite,
2635 )
2636 .unwrap();
2637 t.set_plane_offset(0);
2638 assert_eq!(t.plane_offset(), Some(0));
2639 }
2640
2641 #[test]
2642 fn set_plane_offset_no_format() {
2643 let mut t = TensorDyn::new(&[480, 640, 3], DType::U8, None, None).unwrap();
2645 t.set_plane_offset(4096);
2646 assert_eq!(t.plane_offset(), Some(4096));
2647 }
2648
2649 #[test]
2650 fn set_format_clears_plane_offset() {
2651 let mut t = TensorDyn::new(&[480, 640, 3], DType::U8, None, None).unwrap();
2652 t.set_format(PixelFormat::Rgb).unwrap();
2653 t.set_plane_offset(4096);
2654 assert_eq!(t.plane_offset(), Some(4096));
2655
2656 t.set_format(PixelFormat::Rgb).unwrap();
2658 assert_eq!(t.plane_offset(), Some(4096));
2659
2660 t.reshape(&[480 * 640 * 3]).unwrap();
2662 assert_eq!(t.plane_offset(), None);
2663 assert_eq!(t.format(), None);
2664 }
2665
2666 #[test]
2667 fn map_rejects_out_of_bounds_offset() {
2668 let mut t = Tensor::<u8>::image(
2669 100,
2670 100,
2671 PixelFormat::Rgba,
2672 Some(TensorMemory::Mem),
2673 crate::CpuAccess::ReadWrite,
2674 )
2675 .unwrap();
2676 assert!(t.map().is_ok());
2678 t.set_plane_offset(4096);
2681 assert!(t.map().is_err());
2682 }
2683
2684 #[test]
2685 fn mem_subview_in_bounds_maps_at_offset() {
2686 let parent = Tensor::<u8>::image(
2689 100,
2690 100,
2691 PixelFormat::Rgba,
2692 Some(TensorMemory::Mem),
2693 crate::CpuAccess::ReadWrite,
2694 )
2695 .unwrap();
2696 let view = parent.subview(4096, &[10, 10, 4]).unwrap();
2698 assert_eq!(view.plane_offset(), Some(4096));
2699 assert!(view.map().is_ok());
2700 }
2701
2702 #[test]
2703 fn dyn_batch_dispatches_every_dtype() {
2704 use DType::*;
2708 for dt in [U8, I8, U16, I16, U32, I32, U64, I64, F16, F32, F64] {
2709 let parent = TensorDyn::new(&[2, 4], dt, Some(TensorMemory::Mem), None).unwrap();
2710 let view = parent.batch(1).unwrap();
2711 assert_eq!(view.dtype(), dt, "batch must preserve dtype {dt:?}");
2712 assert_eq!(view.shape(), &[4], "{dt:?}");
2713 }
2714 }
2715
2716 #[test]
2717 fn map_accepts_zero_offset_tensor() {
2718 let mut t = Tensor::<u8>::image(
2719 100,
2720 100,
2721 PixelFormat::Rgba,
2722 Some(TensorMemory::Mem),
2723 crate::CpuAccess::ReadWrite,
2724 )
2725 .unwrap();
2726 t.set_plane_offset(0);
2727 assert!(t.map().is_ok());
2729 }
2730
2731 #[test]
2732 fn dyn_configure_image_nv12() {
2733 let mut t = TensorDyn::image(
2734 640,
2735 480,
2736 PixelFormat::Rgb,
2737 DType::U8,
2738 None,
2739 crate::CpuAccess::ReadWrite,
2740 )
2741 .unwrap();
2742 t.configure_image(320, 240, PixelFormat::Nv12).unwrap();
2743 assert_eq!(t.format(), Some(PixelFormat::Nv12));
2744 assert_eq!((t.width(), t.height()), (Some(320), Some(240)));
2745 }
2746
2747 #[test]
2748 fn tensordyn_colorimetry_roundtrip() {
2749 use crate::{ColorEncoding, Colorimetry, DType, PixelFormat};
2750 let mut t = TensorDyn::image(
2751 1280,
2752 720,
2753 PixelFormat::Nv12,
2754 DType::U8,
2755 None,
2756 crate::CpuAccess::ReadWrite,
2757 )
2758 .unwrap();
2759 assert_eq!(t.colorimetry(), None);
2760 let c = Colorimetry::default().with_encoding(ColorEncoding::Bt709);
2761 t.set_colorimetry(Some(c));
2762 assert_eq!(t.colorimetry(), Some(c));
2763 }
2764
2765 #[test]
2766 fn from_planes_propagates_plane_offset() {
2767 let mut luma =
2768 Tensor::<u8>::new(&[480, 640], Some(TensorMemory::Mem), Some("luma")).unwrap();
2769 luma.set_plane_offset(4096);
2770 let chroma =
2771 Tensor::<u8>::new(&[240, 640], Some(TensorMemory::Mem), Some("chroma")).unwrap();
2772 let combined = Tensor::<u8>::from_planes(luma, chroma, PixelFormat::Nv12).unwrap();
2773 assert_eq!(combined.plane_offset(), Some(4096));
2774 }
2775
2776 #[test]
2777 fn cuda_passthrough_none_for_mem_tensor() {
2778 let t: TensorDyn = Tensor::<f32>::new(&[10], Some(TensorMemory::Mem), None)
2781 .unwrap()
2782 .into();
2783 assert!(t.cuda().is_none());
2784 assert!(t.cuda_map().is_none());
2785 }
2786}