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_raw_host(
965 ptr: *mut u8,
966 shape: &[usize],
967 dtype: DType,
968 name: Option<&str>,
969 ) -> crate::Result<Self> {
970 unsafe { Self::from_foreign_ptr(ptr, shape, dtype, None, name) }
971 }
972
973 pub unsafe fn from_raw_host_with_capacity(
989 ptr: *mut u8,
990 shape: &[usize],
991 capacity_bytes: usize,
992 dtype: DType,
993 name: Option<&str>,
994 ) -> crate::Result<Self> {
995 unsafe {
996 Self::from_foreign_ptr_with_capacity(ptr, shape, dtype, capacity_bytes, None, name)
997 }
998 }
999
1000 pub(crate) fn import_storage(
1012 desc: &crate::TensorDesc,
1013 shape: &[usize],
1014 dtype: DType,
1015 ) -> crate::Result<Self> {
1016 match desc.kind {
1017 #[cfg(target_os = "linux")]
1018 crate::protocol::kind::DMABUF => {
1019 let owned = crate::protocol::dup_descriptor_fd(desc)?;
1024 Self::from_fd(owned, shape, dtype, None)
1025 }
1026 #[cfg(not(target_os = "linux"))]
1027 crate::protocol::kind::DMABUF => Err(crate::Error::NotImplemented(
1028 "dma-buf import off Linux".into(),
1029 )),
1030 #[cfg(any(target_os = "macos", target_os = "ios"))]
1031 crate::protocol::kind::IOSURFACE => {
1032 let id = crate::protocol::descriptor_surface_id(desc)?;
1033 let surface_ref = crate::iosurface::lookup_by_id(id).ok_or_else(|| {
1034 crate::Error::InvalidArgument(format!(
1035 "IOSurface id {id} is not live (the producer's capsule may \
1036 already have been released)"
1037 ))
1038 })?;
1039 let result = unsafe { Self::from_iosurface(surface_ref, shape, dtype, None) };
1044 crate::iosurface::release(surface_ref);
1045 result
1046 }
1047 #[cfg(not(any(target_os = "macos", target_os = "ios")))]
1048 crate::protocol::kind::IOSURFACE => Err(crate::Error::NotImplemented(
1049 "IOSurface import off Apple platforms".into(),
1050 )),
1051 crate::protocol::kind::HOST => {
1052 crate::protocol::check_descriptor_host_ptr(desc)?;
1053 unsafe {
1060 Self::from_raw_host_with_capacity(
1061 desc.ptr.0,
1062 shape,
1063 desc.capacity as usize,
1064 dtype,
1065 None,
1066 )
1067 }
1068 }
1069 crate::protocol::kind::PBO => {
1070 let buffer_id = crate::protocol::descriptor_pbo_buffer_id(desc)?;
1071 let ops = unsafe { crate::pbo::import_pbo_ops(desc.ptr.0 as *const _)? };
1081 Self::from_pbo_import(buffer_id, desc.capacity as usize, shape, dtype, ops)
1082 }
1083 #[cfg(target_os = "windows")]
1084 crate::protocol::kind::D3D11_TEXTURE => {
1085 let (tex, completion) = crate::protocol::descriptor_d3d11_handles(desc)?;
1086 unsafe {
1093 let (format, width, height) =
1094 crate::protocol::descriptor_d3d11_geometry(desc, tex, shape)?;
1095 Self::from_d3d11_shared_handle(
1096 tex,
1097 width,
1098 height,
1099 format,
1100 dtype,
1101 crate::CpuAccess::ReadWrite,
1102 completion,
1103 None,
1104 )
1105 }
1106 }
1107 #[cfg(not(target_os = "windows"))]
1108 crate::protocol::kind::D3D11_TEXTURE => Err(crate::Error::NotImplemented(
1109 "D3D11 texture import off Windows".into(),
1110 )),
1111 k => Err(crate::Error::NotImplemented(format!(
1112 "tensor interop kind {k} cannot be imported by this build"
1113 ))),
1114 }
1115 }
1116
1117 pub unsafe fn from_foreign_ptr(
1130 ptr: *mut u8,
1131 shape: &[usize],
1132 dtype: DType,
1133 owner: Option<crate::ForeignOwner>,
1134 name: Option<&str>,
1135 ) -> crate::Result<Self> {
1136 unsafe {
1141 match dtype {
1142 DType::U8 => {
1143 Tensor::<u8>::from_foreign(ptr.cast(), shape, owner, name).map(Self::U8)
1144 }
1145 DType::I8 => {
1146 Tensor::<i8>::from_foreign(ptr.cast(), shape, owner, name).map(Self::I8)
1147 }
1148 DType::U16 => {
1149 Tensor::<u16>::from_foreign(ptr.cast(), shape, owner, name).map(Self::U16)
1150 }
1151 DType::I16 => {
1152 Tensor::<i16>::from_foreign(ptr.cast(), shape, owner, name).map(Self::I16)
1153 }
1154 DType::U32 => {
1155 Tensor::<u32>::from_foreign(ptr.cast(), shape, owner, name).map(Self::U32)
1156 }
1157 DType::I32 => {
1158 Tensor::<i32>::from_foreign(ptr.cast(), shape, owner, name).map(Self::I32)
1159 }
1160 DType::U64 => {
1161 Tensor::<u64>::from_foreign(ptr.cast(), shape, owner, name).map(Self::U64)
1162 }
1163 DType::I64 => {
1164 Tensor::<i64>::from_foreign(ptr.cast(), shape, owner, name).map(Self::I64)
1165 }
1166 DType::F16 => {
1167 Tensor::<f16>::from_foreign(ptr.cast(), shape, owner, name).map(Self::F16)
1168 }
1169 DType::F32 => {
1170 Tensor::<f32>::from_foreign(ptr.cast(), shape, owner, name).map(Self::F32)
1171 }
1172 DType::F64 => {
1173 Tensor::<f64>::from_foreign(ptr.cast(), shape, owner, name).map(Self::F64)
1174 }
1175 }
1176 }
1177 }
1178
1179 pub unsafe fn from_foreign_ptr_with_capacity(
1190 ptr: *mut u8,
1191 shape: &[usize],
1192 dtype: DType,
1193 capacity_bytes: usize,
1194 owner: Option<crate::ForeignOwner>,
1195 name: Option<&str>,
1196 ) -> crate::Result<Self> {
1197 macro_rules! import {
1198 ($t:ty, $variant:ident) => {
1199 unsafe {
1204 Tensor::<$t>::from_foreign_with_capacity(
1205 ptr.cast(),
1206 shape,
1207 capacity_bytes,
1208 owner,
1209 name,
1210 )
1211 }
1212 .map(Self::$variant)
1213 };
1214 }
1215 match dtype {
1216 DType::U8 => import!(u8, U8),
1217 DType::I8 => import!(i8, I8),
1218 DType::U16 => import!(u16, U16),
1219 DType::I16 => import!(i16, I16),
1220 DType::U32 => import!(u32, U32),
1221 DType::I32 => import!(i32, I32),
1222 DType::U64 => import!(u64, U64),
1223 DType::I64 => import!(i64, I64),
1224 DType::F16 => import!(f16, F16),
1225 DType::F32 => import!(f32, F32),
1226 DType::F64 => import!(f64, F64),
1227 }
1228 }
1229
1230 #[cfg(any(target_os = "macos", target_os = "ios"))]
1241 pub fn from_iosurface_id(
1242 id: u32,
1243 shape: &[usize],
1244 dtype: DType,
1245 name: Option<&str>,
1246 ) -> crate::Result<Self> {
1247 let surface = crate::iosurface::lookup_by_id(id).ok_or_else(|| {
1248 crate::Error::InvalidArgument(format!("no live IOSurface with id {id}"))
1249 })?;
1250 let result = unsafe { Self::from_iosurface(surface, shape, dtype, name) };
1265 crate::iosurface::release(surface);
1266 result
1267 }
1268
1269 #[cfg(any(target_os = "macos", target_os = "ios"))]
1277 pub unsafe fn from_iosurface(
1278 surface_ref: *mut std::ffi::c_void,
1279 shape: &[usize],
1280 dtype: DType,
1281 name: Option<&str>,
1282 ) -> crate::Result<Self> {
1283 unsafe {
1284 match dtype {
1285 DType::U8 => Tensor::<u8>::from_iosurface(surface_ref, shape, name).map(Self::U8),
1286 DType::I8 => Tensor::<i8>::from_iosurface(surface_ref, shape, name).map(Self::I8),
1287 DType::U16 => {
1288 Tensor::<u16>::from_iosurface(surface_ref, shape, name).map(Self::U16)
1289 }
1290 DType::I16 => {
1291 Tensor::<i16>::from_iosurface(surface_ref, shape, name).map(Self::I16)
1292 }
1293 DType::U32 => {
1294 Tensor::<u32>::from_iosurface(surface_ref, shape, name).map(Self::U32)
1295 }
1296 DType::I32 => {
1297 Tensor::<i32>::from_iosurface(surface_ref, shape, name).map(Self::I32)
1298 }
1299 DType::U64 => {
1300 Tensor::<u64>::from_iosurface(surface_ref, shape, name).map(Self::U64)
1301 }
1302 DType::I64 => {
1303 Tensor::<i64>::from_iosurface(surface_ref, shape, name).map(Self::I64)
1304 }
1305 DType::F16 => {
1306 Tensor::<f16>::from_iosurface(surface_ref, shape, name).map(Self::F16)
1307 }
1308 DType::F32 => {
1309 Tensor::<f32>::from_iosurface(surface_ref, shape, name).map(Self::F32)
1310 }
1311 DType::F64 => {
1312 Tensor::<f64>::from_iosurface(surface_ref, shape, name).map(Self::F64)
1313 }
1314 }
1315 }
1316 }
1317
1318 #[cfg(any(target_os = "macos", target_os = "ios"))]
1321 pub fn iosurface_id(&self) -> Option<u32> {
1322 dispatch!(self, iosurface_id)
1323 }
1324
1325 pub fn pbo_id(&self) -> Option<u32> {
1328 dispatch!(self, pbo_id)
1329 }
1330
1331 pub fn pbo_vtable_ptr(&self) -> Option<*const std::ffi::c_void> {
1335 dispatch!(self, pbo_vtable_ptr)
1336 }
1337
1338 pub fn pbo_keepalive(&self) -> Option<std::sync::Arc<dyn Send + Sync>> {
1343 dispatch!(self, pbo_keepalive)
1344 }
1345
1346 #[cfg(any(target_os = "macos", target_os = "ios"))]
1350 pub fn iosurface_ref(&self) -> Option<*mut std::ffi::c_void> {
1351 dispatch!(self, iosurface_ref)
1352 }
1353
1354 #[cfg(any(target_os = "macos", target_os = "ios"))]
1359 pub fn iosurface_physical_dims(&self) -> Option<(usize, usize)> {
1360 dispatch!(self, iosurface_physical_dims)
1361 }
1362
1363 #[cfg(target_os = "windows")]
1366 pub fn d3d11_texture(&self) -> Option<*mut std::ffi::c_void> {
1367 dispatch!(self, d3d11_texture)
1368 }
1369
1370 #[cfg(target_os = "windows")]
1373 pub fn d3d11_layout(&self) -> Option<crate::d3d11_layout::D3d11ImageLayout> {
1374 dispatch!(self, d3d11_layout)
1375 }
1376
1377 #[cfg(target_os = "windows")]
1380 pub fn d3d11_shared_handle(&self) -> crate::Result<std::os::windows::io::OwnedHandle> {
1381 dispatch!(self, d3d11_shared_handle)
1382 }
1383
1384 #[cfg(target_os = "windows")]
1387 pub fn gpu_completion(&self) -> crate::Result<Option<crate::d3d11::GpuCompletion>> {
1388 dispatch!(self, gpu_completion)
1389 }
1390
1391 pub fn gpu_write_value(&self) -> u64 {
1395 dispatch!(self, gpu_write_value)
1396 }
1397
1398 #[cfg(target_os = "windows")]
1402 pub fn set_gpu_write(&self, value: u64) -> crate::Result<()> {
1403 dispatch!(self, set_gpu_write, value)
1404 }
1405
1406 #[cfg(target_os = "windows")]
1409 pub(crate) fn d3d11_shared_handle_value(&self) -> Option<usize> {
1410 dispatch!(self, d3d11_shared_handle_value)
1411 }
1412
1413 #[cfg(target_os = "windows")]
1421 #[allow(clippy::too_many_arguments)] pub unsafe fn from_d3d11_texture(
1423 texture: *mut std::ffi::c_void,
1424 width: usize,
1425 height: usize,
1426 format: PixelFormat,
1427 dtype: DType,
1428 access: crate::CpuAccess,
1429 name: Option<&str>,
1430 ) -> crate::Result<Self> {
1431 unsafe {
1434 match dtype {
1435 DType::U8 => {
1436 Tensor::<u8>::from_d3d11_texture(texture, width, height, format, access, name)
1437 .map(Self::U8)
1438 }
1439 DType::I8 => {
1440 Tensor::<i8>::from_d3d11_texture(texture, width, height, format, access, name)
1441 .map(Self::I8)
1442 }
1443 DType::U16 => {
1444 Tensor::<u16>::from_d3d11_texture(texture, width, height, format, access, name)
1445 .map(Self::U16)
1446 }
1447 DType::I16 => {
1448 Tensor::<i16>::from_d3d11_texture(texture, width, height, format, access, name)
1449 .map(Self::I16)
1450 }
1451 DType::U32 => {
1452 Tensor::<u32>::from_d3d11_texture(texture, width, height, format, access, name)
1453 .map(Self::U32)
1454 }
1455 DType::I32 => {
1456 Tensor::<i32>::from_d3d11_texture(texture, width, height, format, access, name)
1457 .map(Self::I32)
1458 }
1459 DType::U64 => {
1460 Tensor::<u64>::from_d3d11_texture(texture, width, height, format, access, name)
1461 .map(Self::U64)
1462 }
1463 DType::I64 => {
1464 Tensor::<i64>::from_d3d11_texture(texture, width, height, format, access, name)
1465 .map(Self::I64)
1466 }
1467 DType::F16 => {
1468 Tensor::<f16>::from_d3d11_texture(texture, width, height, format, access, name)
1469 .map(Self::F16)
1470 }
1471 DType::F32 => {
1472 Tensor::<f32>::from_d3d11_texture(texture, width, height, format, access, name)
1473 .map(Self::F32)
1474 }
1475 DType::F64 => {
1476 Tensor::<f64>::from_d3d11_texture(texture, width, height, format, access, name)
1477 .map(Self::F64)
1478 }
1479 }
1480 }
1481 }
1482
1483 #[cfg(target_os = "windows")]
1492 #[allow(clippy::too_many_arguments)] pub unsafe fn from_d3d11_shared_handle(
1494 handle: std::os::windows::io::RawHandle,
1495 width: usize,
1496 height: usize,
1497 format: PixelFormat,
1498 dtype: DType,
1499 access: crate::CpuAccess,
1500 completion: Option<(std::os::windows::io::RawHandle, u64)>,
1501 name: Option<&str>,
1502 ) -> crate::Result<Self> {
1503 unsafe {
1506 match dtype {
1507 DType::U8 => Tensor::<u8>::from_d3d11_shared_handle(
1508 handle, width, height, format, access, completion, name,
1509 )
1510 .map(Self::U8),
1511 DType::I8 => Tensor::<i8>::from_d3d11_shared_handle(
1512 handle, width, height, format, access, completion, name,
1513 )
1514 .map(Self::I8),
1515 DType::U16 => Tensor::<u16>::from_d3d11_shared_handle(
1516 handle, width, height, format, access, completion, name,
1517 )
1518 .map(Self::U16),
1519 DType::I16 => Tensor::<i16>::from_d3d11_shared_handle(
1520 handle, width, height, format, access, completion, name,
1521 )
1522 .map(Self::I16),
1523 DType::U32 => Tensor::<u32>::from_d3d11_shared_handle(
1524 handle, width, height, format, access, completion, name,
1525 )
1526 .map(Self::U32),
1527 DType::I32 => Tensor::<i32>::from_d3d11_shared_handle(
1528 handle, width, height, format, access, completion, name,
1529 )
1530 .map(Self::I32),
1531 DType::U64 => Tensor::<u64>::from_d3d11_shared_handle(
1532 handle, width, height, format, access, completion, name,
1533 )
1534 .map(Self::U64),
1535 DType::I64 => Tensor::<i64>::from_d3d11_shared_handle(
1536 handle, width, height, format, access, completion, name,
1537 )
1538 .map(Self::I64),
1539 DType::F16 => Tensor::<f16>::from_d3d11_shared_handle(
1540 handle, width, height, format, access, completion, name,
1541 )
1542 .map(Self::F16),
1543 DType::F32 => Tensor::<f32>::from_d3d11_shared_handle(
1544 handle, width, height, format, access, completion, name,
1545 )
1546 .map(Self::F32),
1547 DType::F64 => Tensor::<f64>::from_d3d11_shared_handle(
1548 handle, width, height, format, access, completion, name,
1549 )
1550 .map(Self::F64),
1551 }
1552 }
1553 }
1554
1555 #[cfg(target_os = "android")]
1565 pub unsafe fn from_hardware_buffer(
1566 buffer_ptr: *mut std::ffi::c_void,
1567 shape: &[usize],
1568 dtype: DType,
1569 name: Option<&str>,
1570 ) -> crate::Result<Self> {
1571 unsafe {
1572 match dtype {
1573 DType::U8 => {
1574 Tensor::<u8>::from_hardware_buffer(buffer_ptr, shape, name).map(Self::U8)
1575 }
1576 DType::I8 => {
1577 Tensor::<i8>::from_hardware_buffer(buffer_ptr, shape, name).map(Self::I8)
1578 }
1579 DType::U16 => {
1580 Tensor::<u16>::from_hardware_buffer(buffer_ptr, shape, name).map(Self::U16)
1581 }
1582 DType::I16 => {
1583 Tensor::<i16>::from_hardware_buffer(buffer_ptr, shape, name).map(Self::I16)
1584 }
1585 DType::U32 => {
1586 Tensor::<u32>::from_hardware_buffer(buffer_ptr, shape, name).map(Self::U32)
1587 }
1588 DType::I32 => {
1589 Tensor::<i32>::from_hardware_buffer(buffer_ptr, shape, name).map(Self::I32)
1590 }
1591 DType::U64 => {
1592 Tensor::<u64>::from_hardware_buffer(buffer_ptr, shape, name).map(Self::U64)
1593 }
1594 DType::I64 => {
1595 Tensor::<i64>::from_hardware_buffer(buffer_ptr, shape, name).map(Self::I64)
1596 }
1597 DType::F16 => {
1598 Tensor::<f16>::from_hardware_buffer(buffer_ptr, shape, name).map(Self::F16)
1599 }
1600 DType::F32 => {
1601 Tensor::<f32>::from_hardware_buffer(buffer_ptr, shape, name).map(Self::F32)
1602 }
1603 DType::F64 => {
1604 Tensor::<f64>::from_hardware_buffer(buffer_ptr, shape, name).map(Self::F64)
1605 }
1606 }
1607 }
1608 }
1609
1610 #[cfg(target_os = "android")]
1614 pub fn hardware_buffer_ptr(&self) -> Option<*mut std::ffi::c_void> {
1615 dispatch!(self, hardware_buffer_ptr)
1616 }
1617
1618 #[cfg(target_os = "android")]
1622 pub fn hardware_buffer_physical_dims(&self) -> Option<(usize, usize)> {
1623 dispatch!(self, hardware_buffer_physical_dims)
1624 }
1625
1626 pub fn copy_to_flat(&self, dst: &mut [u8]) -> crate::Result<()> {
1631 dispatch!(self, copy_to_flat, dst)
1632 }
1633
1634 pub fn image(
1652 width: usize,
1653 height: usize,
1654 format: PixelFormat,
1655 dtype: DType,
1656 memory: Option<TensorMemory>,
1657 access: crate::CpuAccess,
1658 ) -> crate::Result<Self> {
1659 match dtype {
1660 DType::U8 => Tensor::<u8>::image(width, height, format, memory, access).map(Self::U8),
1661 DType::I8 => Tensor::<i8>::image(width, height, format, memory, access).map(Self::I8),
1662 DType::U16 => {
1663 Tensor::<u16>::image(width, height, format, memory, access).map(Self::U16)
1664 }
1665 DType::I16 => {
1666 Tensor::<i16>::image(width, height, format, memory, access).map(Self::I16)
1667 }
1668 DType::U32 => {
1669 Tensor::<u32>::image(width, height, format, memory, access).map(Self::U32)
1670 }
1671 DType::I32 => {
1672 Tensor::<i32>::image(width, height, format, memory, access).map(Self::I32)
1673 }
1674 DType::U64 => {
1675 Tensor::<u64>::image(width, height, format, memory, access).map(Self::U64)
1676 }
1677 DType::I64 => {
1678 Tensor::<i64>::image(width, height, format, memory, access).map(Self::I64)
1679 }
1680 DType::F16 => {
1681 Tensor::<f16>::image(width, height, format, memory, access).map(Self::F16)
1682 }
1683 DType::F32 => {
1684 Tensor::<f32>::image(width, height, format, memory, access).map(Self::F32)
1685 }
1686 DType::F64 => {
1687 Tensor::<f64>::image(width, height, format, memory, access).map(Self::F64)
1688 }
1689 }
1690 }
1691
1692 pub fn image_desc(desc: &crate::ImageDesc) -> crate::Result<Self> {
1696 match desc.dtype() {
1697 DType::U8 => Tensor::<u8>::image_desc(desc).map(Self::U8),
1698 DType::I8 => Tensor::<i8>::image_desc(desc).map(Self::I8),
1699 DType::U16 => Tensor::<u16>::image_desc(desc).map(Self::U16),
1700 DType::I16 => Tensor::<i16>::image_desc(desc).map(Self::I16),
1701 DType::U32 => Tensor::<u32>::image_desc(desc).map(Self::U32),
1702 DType::I32 => Tensor::<i32>::image_desc(desc).map(Self::I32),
1703 DType::U64 => Tensor::<u64>::image_desc(desc).map(Self::U64),
1704 DType::I64 => Tensor::<i64>::image_desc(desc).map(Self::I64),
1705 DType::F16 => Tensor::<f16>::image_desc(desc).map(Self::F16),
1706 DType::F32 => Tensor::<f32>::image_desc(desc).map(Self::F32),
1707 DType::F64 => Tensor::<f64>::image_desc(desc).map(Self::F64),
1708 }
1709 }
1710
1711 pub fn compression(&self) -> Option<crate::CompressionScheme> {
1714 dispatch!(self, compression)
1715 }
1716
1717 pub fn image_with_stride(
1744 width: usize,
1745 height: usize,
1746 format: PixelFormat,
1747 dtype: DType,
1748 row_stride_bytes: usize,
1749 memory: Option<TensorMemory>,
1750 access: crate::CpuAccess,
1751 ) -> crate::Result<Self> {
1752 match dtype {
1753 DType::U8 => Tensor::<u8>::image_with_stride(
1754 width,
1755 height,
1756 format,
1757 row_stride_bytes,
1758 memory,
1759 access,
1760 )
1761 .map(Self::U8),
1762 DType::I8 => Tensor::<i8>::image_with_stride(
1763 width,
1764 height,
1765 format,
1766 row_stride_bytes,
1767 memory,
1768 access,
1769 )
1770 .map(Self::I8),
1771 DType::U16 => Tensor::<u16>::image_with_stride(
1772 width,
1773 height,
1774 format,
1775 row_stride_bytes,
1776 memory,
1777 access,
1778 )
1779 .map(Self::U16),
1780 DType::I16 => Tensor::<i16>::image_with_stride(
1781 width,
1782 height,
1783 format,
1784 row_stride_bytes,
1785 memory,
1786 access,
1787 )
1788 .map(Self::I16),
1789 DType::U32 => Tensor::<u32>::image_with_stride(
1790 width,
1791 height,
1792 format,
1793 row_stride_bytes,
1794 memory,
1795 access,
1796 )
1797 .map(Self::U32),
1798 DType::I32 => Tensor::<i32>::image_with_stride(
1799 width,
1800 height,
1801 format,
1802 row_stride_bytes,
1803 memory,
1804 access,
1805 )
1806 .map(Self::I32),
1807 DType::U64 => Tensor::<u64>::image_with_stride(
1808 width,
1809 height,
1810 format,
1811 row_stride_bytes,
1812 memory,
1813 access,
1814 )
1815 .map(Self::U64),
1816 DType::I64 => Tensor::<i64>::image_with_stride(
1817 width,
1818 height,
1819 format,
1820 row_stride_bytes,
1821 memory,
1822 access,
1823 )
1824 .map(Self::I64),
1825 DType::F16 => Tensor::<f16>::image_with_stride(
1826 width,
1827 height,
1828 format,
1829 row_stride_bytes,
1830 memory,
1831 access,
1832 )
1833 .map(Self::F16),
1834 DType::F32 => Tensor::<f32>::image_with_stride(
1835 width,
1836 height,
1837 format,
1838 row_stride_bytes,
1839 memory,
1840 access,
1841 )
1842 .map(Self::F32),
1843 DType::F64 => Tensor::<f64>::image_with_stride(
1844 width,
1845 height,
1846 format,
1847 row_stride_bytes,
1848 memory,
1849 access,
1850 )
1851 .map(Self::F64),
1852 }
1853 }
1854}
1855
1856impl From<Tensor<u8>> for TensorDyn {
1859 fn from(t: Tensor<u8>) -> Self {
1860 Self::U8(t)
1861 }
1862}
1863
1864impl From<Tensor<i8>> for TensorDyn {
1865 fn from(t: Tensor<i8>) -> Self {
1866 Self::I8(t)
1867 }
1868}
1869
1870impl From<Tensor<u16>> for TensorDyn {
1871 fn from(t: Tensor<u16>) -> Self {
1872 Self::U16(t)
1873 }
1874}
1875
1876impl From<Tensor<i16>> for TensorDyn {
1877 fn from(t: Tensor<i16>) -> Self {
1878 Self::I16(t)
1879 }
1880}
1881
1882impl From<Tensor<u32>> for TensorDyn {
1883 fn from(t: Tensor<u32>) -> Self {
1884 Self::U32(t)
1885 }
1886}
1887
1888impl From<Tensor<i32>> for TensorDyn {
1889 fn from(t: Tensor<i32>) -> Self {
1890 Self::I32(t)
1891 }
1892}
1893
1894impl From<Tensor<u64>> for TensorDyn {
1895 fn from(t: Tensor<u64>) -> Self {
1896 Self::U64(t)
1897 }
1898}
1899
1900impl From<Tensor<i64>> for TensorDyn {
1901 fn from(t: Tensor<i64>) -> Self {
1902 Self::I64(t)
1903 }
1904}
1905
1906impl From<Tensor<f16>> for TensorDyn {
1907 fn from(t: Tensor<f16>) -> Self {
1908 Self::F16(t)
1909 }
1910}
1911
1912impl From<Tensor<f32>> for TensorDyn {
1913 fn from(t: Tensor<f32>) -> Self {
1914 Self::F32(t)
1915 }
1916}
1917
1918impl From<Tensor<f64>> for TensorDyn {
1919 fn from(t: Tensor<f64>) -> Self {
1920 Self::F64(t)
1921 }
1922}
1923
1924impl fmt::Debug for TensorDyn {
1925 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
1926 dispatch!(self, fmt, f)
1927 }
1928}
1929
1930pub type Raw = *mut std::ffi::c_void;
1933
1934impl TensorDyn {
1935 pub fn into_raw(self) -> Raw {
1938 Box::into_raw(Box::new(self)) as Raw
1939 }
1940
1941 pub unsafe fn from_raw(p: Raw) -> TensorDyn {
1947 *unsafe { Box::from_raw(p as *mut TensorDyn) }
1948 }
1949
1950 pub unsafe fn with_raw<R>(p: Raw, f: impl FnOnce(&mut TensorDyn) -> R) -> R {
1956 let mut t = std::mem::ManuallyDrop::new(unsafe { std::ptr::read(p as *mut TensorDyn) });
1957 f(&mut t)
1958 }
1959}
1960
1961#[cfg(test)]
1962mod tests {
1963 use super::*;
1964
1965 #[test]
1966 fn from_typed_tensor() {
1967 let t = Tensor::<u8>::new(&[10], None, None).unwrap();
1968 let dyn_t: TensorDyn = t.into();
1969 assert_eq!(dyn_t.dtype(), DType::U8);
1970 assert_eq!(dyn_t.shape(), &[10]);
1971 }
1972
1973 #[test]
1974 fn from_foreign_ptr_wraps_borrowed_memory() {
1975 use crate::TensorMapTrait;
1976 let mut vec: Vec<f32> = vec![0.0; 4];
1979 let ptr = vec.as_mut_ptr() as *mut u8;
1980 let owner: crate::ForeignOwner = Box::new(vec);
1981 let t = unsafe {
1982 TensorDyn::from_foreign_ptr(ptr, &[2, 2], DType::F32, Some(owner), Some("trt_output"))
1983 }
1984 .unwrap();
1985 assert_eq!(t.dtype(), DType::F32);
1986 assert_eq!(t.memory(), TensorMemory::Mem);
1987 assert_eq!(t.shape(), &[2, 2]);
1988 {
1989 let mut m = t.as_f32().unwrap().map().unwrap();
1990 m.as_mut_slice().copy_from_slice(&[1.0, 2.0, 3.0, 4.0]);
1991 }
1992 let m = t.as_f32().unwrap().map().unwrap();
1993 assert_eq!(m.as_slice(), &[1.0, 2.0, 3.0, 4.0]);
1994 }
1995
1996 #[test]
2005 fn from_foreign_ptr_rejects_null_ptr() {
2006 let err = unsafe {
2007 TensorDyn::from_foreign_ptr(std::ptr::null_mut(), &[4], DType::U8, None, None)
2008 }
2009 .unwrap_err();
2010 assert!(
2012 matches!(err, crate::error::Error::InvalidArgument(ref m) if m.contains("non-null")),
2013 "expected InvalidArgument(non-null), got {err:?}"
2014 );
2015 }
2016
2017 #[test]
2018 fn from_foreign_ptr_rejects_empty_shape() {
2019 let mut dummy: u8 = 0;
2020 let err = unsafe {
2021 TensorDyn::from_foreign_ptr(&mut dummy as *mut u8, &[], DType::U8, None, None)
2022 }
2023 .unwrap_err();
2024 assert!(
2025 matches!(err, crate::error::Error::InvalidSize(0)),
2026 "expected InvalidSize(0) for empty shape, got {err:?}"
2027 );
2028 }
2029
2030 #[test]
2031 fn from_foreign_ptr_rejects_overflow_shape() {
2032 let mut dummy: u8 = 0;
2033 let huge = [usize::MAX / 2 + 1, 2];
2034 let err = unsafe { TensorDyn::from_foreign_ptr(&mut dummy, &huge, DType::U8, None, None) }
2035 .unwrap_err();
2036 assert!(
2037 matches!(err, crate::error::Error::InvalidArgument(ref m) if m.contains("overflow")),
2038 "expected InvalidArgument(overflow), got {err:?}"
2039 );
2040 }
2041
2042 #[test]
2043 fn from_foreign_ptr_u8_dtype_dispatch() {
2044 let mut buf: Vec<u8> = vec![1, 2, 3, 4];
2047 let ptr = buf.as_mut_ptr();
2048 let owner: crate::ForeignOwner = Box::new(buf);
2049 let t = unsafe {
2050 TensorDyn::from_foreign_ptr(ptr, &[4], DType::U8, Some(owner), Some("u8_foreign"))
2051 }
2052 .unwrap();
2053 assert_eq!(t.dtype(), DType::U8);
2054 assert_eq!(t.shape(), &[4]);
2055 let m = t.as_u8().unwrap().map().unwrap();
2056 use crate::TensorMapTrait;
2057 assert_eq!(m.as_slice(), &[1u8, 2, 3, 4]);
2058 }
2059
2060 #[test]
2061 fn map_bytes_roundtrips_and_is_typed_agnostic() {
2062 use crate::TensorMapTrait;
2063 let t = Tensor::<f32>::new(&[4], None, None).unwrap();
2066 let dyn_t: TensorDyn = t.into();
2067 {
2068 let mut view = dyn_t.map_bytes(crate::CpuAccess::ReadWrite).unwrap();
2069 assert_eq!(view.len(), 16, "4 x f32 is 16 bytes, not 4");
2070 view.as_mut_slice()
2071 .copy_from_slice(&[1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16]);
2072 }
2073 let view = dyn_t.map_bytes(crate::CpuAccess::Read).unwrap();
2074 assert_eq!(
2075 view.as_slice(),
2076 &[1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16]
2077 );
2078 }
2079
2080 #[test]
2081 fn map_bytes_none_is_invalid() {
2082 let t = Tensor::<u8>::new(&[4], None, None).unwrap();
2083 let dyn_t: TensorDyn = t.into();
2084 assert!(matches!(
2085 dyn_t.map_bytes(crate::CpuAccess::None),
2086 Err(crate::error::Error::InvalidArgument(_))
2087 ));
2088 }
2089
2090 #[test]
2091 fn downcast_ref() {
2092 let t = Tensor::<u8>::new(&[10], None, None).unwrap();
2093 let dyn_t: TensorDyn = t.into();
2094 assert!(dyn_t.as_u8().is_some());
2095 assert!(dyn_t.as_i8().is_none());
2096 }
2097
2098 #[test]
2099 fn downcast_into() {
2100 let t = Tensor::<u8>::new(&[10], None, None).unwrap();
2101 let dyn_t: TensorDyn = t.into();
2102 let back = dyn_t.into_u8().unwrap();
2103 assert_eq!(back.shape(), &[10]);
2104 }
2105
2106 #[test]
2107 fn image_accessors() {
2108 let t = Tensor::<u8>::image(
2109 640,
2110 480,
2111 PixelFormat::Rgba,
2112 None,
2113 crate::CpuAccess::ReadWrite,
2114 )
2115 .unwrap();
2116 let dyn_t: TensorDyn = t.into();
2117 assert_eq!(dyn_t.format(), Some(PixelFormat::Rgba));
2118 assert_eq!(dyn_t.width(), Some(640));
2119 assert_eq!(dyn_t.height(), Some(480));
2120 assert!(!dyn_t.is_multiplane());
2121 }
2122
2123 #[test]
2124 fn cpu_access_reports_the_declared_access() {
2125 let read_t =
2127 Tensor::<u8>::image(4, 4, PixelFormat::Rgba, None, crate::CpuAccess::Read).unwrap();
2128 let dyn_read: TensorDyn = read_t.into();
2129 assert_eq!(dyn_read.cpu_access(), crate::CpuAccess::Read);
2130 assert!(!dyn_read.cpu_access().writes());
2131
2132 let mem_t = Tensor::<u8>::new(&[4], None, None).unwrap();
2136 let dyn_mem: TensorDyn = mem_t.into();
2137 assert_eq!(dyn_mem.cpu_access(), crate::CpuAccess::ReadWrite);
2138 assert!(dyn_mem.cpu_access().writes());
2139 }
2140
2141 #[test]
2142 fn image_constructor() {
2143 let dyn_t = TensorDyn::image(
2144 640,
2145 480,
2146 PixelFormat::Rgb,
2147 DType::U8,
2148 None,
2149 crate::CpuAccess::ReadWrite,
2150 )
2151 .unwrap();
2152 assert_eq!(dyn_t.dtype(), DType::U8);
2153 assert_eq!(dyn_t.format(), Some(PixelFormat::Rgb));
2154 assert_eq!(dyn_t.width(), Some(640));
2155 }
2156
2157 #[test]
2158 fn image_constructor_i8() {
2159 let dyn_t = TensorDyn::image(
2160 640,
2161 480,
2162 PixelFormat::Rgb,
2163 DType::I8,
2164 None,
2165 crate::CpuAccess::ReadWrite,
2166 )
2167 .unwrap();
2168 assert_eq!(dyn_t.dtype(), DType::I8);
2169 assert_eq!(dyn_t.format(), Some(PixelFormat::Rgb));
2170 }
2171
2172 #[test]
2173 fn set_format_packed() {
2174 let mut t = TensorDyn::new(&[480, 640, 3], DType::U8, None, None).unwrap();
2175 assert_eq!(t.format(), None);
2176 t.set_format(PixelFormat::Rgb).unwrap();
2177 assert_eq!(t.format(), Some(PixelFormat::Rgb));
2178 assert_eq!(t.width(), Some(640));
2179 assert_eq!(t.height(), Some(480));
2180 }
2181
2182 #[test]
2183 fn set_format_planar() {
2184 let mut t = TensorDyn::new(&[3, 480, 640], DType::U8, None, None).unwrap();
2185 t.set_format(PixelFormat::PlanarRgb).unwrap();
2186 assert_eq!(t.format(), Some(PixelFormat::PlanarRgb));
2187 assert_eq!(t.width(), Some(640));
2188 assert_eq!(t.height(), Some(480));
2189 }
2190
2191 #[test]
2192 fn set_format_rejects_wrong_shape() {
2193 let mut t = TensorDyn::new(&[480, 640, 4], DType::U8, None, None).unwrap();
2194 assert!(t.set_format(PixelFormat::Rgb).is_err());
2195 }
2196
2197 #[test]
2198 fn with_format_builder() {
2199 let t = TensorDyn::new(&[480, 640, 4], DType::U8, None, None)
2200 .unwrap()
2201 .with_format(PixelFormat::Rgba)
2202 .unwrap();
2203 assert_eq!(t.format(), Some(PixelFormat::Rgba));
2204 assert_eq!(t.width(), Some(640));
2205 assert_eq!(t.height(), Some(480));
2206 }
2207
2208 #[cfg(target_os = "linux")]
2209 #[test]
2210 fn dmabuf_clone_mem_tensor_fails() {
2211 let t = TensorDyn::new(&[480, 640, 3], DType::U8, Some(TensorMemory::Mem), None).unwrap();
2212 assert_eq!(t.memory(), TensorMemory::Mem);
2213 assert!(t.dmabuf_clone().is_err());
2214 }
2215
2216 #[cfg(target_os = "linux")]
2217 #[test]
2218 fn dmabuf_mem_tensor_fails() {
2219 let t = TensorDyn::new(&[480, 640, 3], DType::U8, Some(TensorMemory::Mem), None).unwrap();
2220 assert!(t.dmabuf().is_err());
2221 }
2222
2223 #[test]
2224 fn set_format_semi_planar_nv12() {
2225 let mut t = TensorDyn::new(&[720, 640], DType::U8, Some(TensorMemory::Mem), None).unwrap();
2227 t.set_format(PixelFormat::Nv12).unwrap();
2228 assert_eq!(t.format(), Some(PixelFormat::Nv12));
2229 assert_eq!(t.width(), Some(640));
2230 assert_eq!(t.height(), Some(480));
2231 }
2232
2233 #[test]
2234 fn set_format_semi_planar_nv16() {
2235 let mut t = TensorDyn::new(&[960, 640], DType::U8, Some(TensorMemory::Mem), None).unwrap();
2237 t.set_format(PixelFormat::Nv16).unwrap();
2238 assert_eq!(t.format(), Some(PixelFormat::Nv16));
2239 assert_eq!(t.width(), Some(640));
2240 assert_eq!(t.height(), Some(480));
2241 }
2242
2243 #[test]
2244 fn with_format_rejects_wrong_shape() {
2245 let result = TensorDyn::new(&[480, 640, 4], DType::U8, None, None)
2246 .unwrap()
2247 .with_format(PixelFormat::Rgb);
2248 assert!(result.is_err());
2249 }
2250
2251 #[test]
2252 fn set_format_preserved_after_rejection() {
2253 let mut t = TensorDyn::new(&[480, 640, 3], DType::U8, None, None).unwrap();
2254 t.set_format(PixelFormat::Rgb).unwrap();
2255 assert_eq!(t.format(), Some(PixelFormat::Rgb));
2256
2257 assert!(t.set_format(PixelFormat::Rgba).is_err());
2259
2260 assert_eq!(t.format(), Some(PixelFormat::Rgb));
2262 }
2263
2264 #[test]
2265 fn set_format_idempotent() {
2266 let mut t = TensorDyn::new(&[480, 640, 3], DType::U8, None, None).unwrap();
2267 t.set_format(PixelFormat::Rgb).unwrap();
2268 t.set_format(PixelFormat::Rgb).unwrap();
2269 assert_eq!(t.format(), Some(PixelFormat::Rgb));
2270 assert_eq!(t.width(), Some(640));
2271 assert_eq!(t.height(), Some(480));
2272 }
2273
2274 #[test]
2277 fn set_row_stride_valid() {
2278 let mut t = TensorDyn::image(
2280 100,
2281 100,
2282 PixelFormat::Rgba,
2283 DType::U8,
2284 None,
2285 crate::CpuAccess::ReadWrite,
2286 )
2287 .unwrap();
2288 t.set_row_stride(512).unwrap();
2289 assert_eq!(t.row_stride(), Some(512));
2290 assert_eq!(t.effective_row_stride(), Some(512));
2291 }
2292
2293 #[test]
2294 fn set_row_stride_equals_min() {
2295 let mut t = TensorDyn::image(
2297 100,
2298 100,
2299 PixelFormat::Rgb,
2300 DType::U8,
2301 None,
2302 crate::CpuAccess::ReadWrite,
2303 )
2304 .unwrap();
2305 t.set_row_stride(300).unwrap();
2306 assert_eq!(t.row_stride(), Some(300));
2307 }
2308
2309 #[test]
2310 fn set_row_stride_too_small() {
2311 let mut t = TensorDyn::image(
2316 64,
2317 100,
2318 PixelFormat::Rgba,
2319 DType::U8,
2320 None,
2321 crate::CpuAccess::ReadWrite,
2322 )
2323 .unwrap();
2324 assert!(t.set_row_stride(200).is_err());
2325 assert_eq!(t.row_stride(), None);
2326 }
2327
2328 #[test]
2329 fn set_row_stride_zero() {
2330 let mut t = TensorDyn::image(
2331 100,
2332 100,
2333 PixelFormat::Rgb,
2334 DType::U8,
2335 None,
2336 crate::CpuAccess::ReadWrite,
2337 )
2338 .unwrap();
2339 assert!(t.set_row_stride(0).is_err());
2340 }
2341
2342 #[test]
2343 fn set_row_stride_requires_format() {
2344 let mut t = TensorDyn::new(&[480, 640, 3], DType::U8, None, None).unwrap();
2345 assert!(t.set_row_stride(2048).is_err());
2346 }
2347
2348 #[test]
2349 fn effective_row_stride_without_stride() {
2350 let t = TensorDyn::image(
2355 64,
2356 100,
2357 PixelFormat::Rgb,
2358 DType::U8,
2359 None,
2360 crate::CpuAccess::ReadWrite,
2361 )
2362 .unwrap();
2363 assert_eq!(t.row_stride(), None);
2364 assert_eq!(t.effective_row_stride(), Some(192)); }
2366
2367 #[test]
2368 fn effective_row_stride_padded_packed_dma() {
2369 let t = match TensorDyn::image(
2375 100,
2376 100,
2377 PixelFormat::Rgb,
2378 DType::U8,
2379 Some(TensorMemory::DmaBuf),
2380 crate::CpuAccess::ReadWrite,
2381 ) {
2382 Ok(t) if t.memory() == TensorMemory::DmaBuf => t,
2383 _ => return,
2384 };
2385 if cfg!(target_os = "android") {
2391 assert!(
2392 t.row_stride().is_none() || t.row_stride() >= Some(300),
2393 "AHardwareBuffer pitch must be absent or at least the tight row"
2394 );
2395 return;
2396 }
2397 if cfg!(target_os = "windows") {
2402 assert!(
2403 t.row_stride().is_none() || t.row_stride() >= Some(300),
2404 "D3D11 staging pitch must be absent or at least the tight row"
2405 );
2406 assert!(t.effective_row_stride() >= Some(300));
2407 return;
2408 }
2409 assert_eq!(t.row_stride(), Some(320));
2410 assert_eq!(t.effective_row_stride(), Some(320));
2411 }
2412
2413 #[test]
2414 fn effective_row_stride_no_format() {
2415 let t = TensorDyn::new(&[480, 640, 3], DType::U8, None, None).unwrap();
2416 assert_eq!(t.effective_row_stride(), None);
2417 }
2418
2419 #[test]
2420 fn with_row_stride_builder() {
2421 let t = TensorDyn::image(
2422 100,
2423 100,
2424 PixelFormat::Rgba,
2425 DType::U8,
2426 None,
2427 crate::CpuAccess::ReadWrite,
2428 )
2429 .unwrap()
2430 .with_row_stride(512)
2431 .unwrap();
2432 assert_eq!(t.row_stride(), Some(512));
2433 assert_eq!(t.effective_row_stride(), Some(512));
2434 }
2435
2436 #[test]
2437 fn with_row_stride_rejects_small() {
2438 let result = TensorDyn::image(
2439 100,
2440 100,
2441 PixelFormat::Rgba,
2442 DType::U8,
2443 None,
2444 crate::CpuAccess::ReadWrite,
2445 )
2446 .unwrap()
2447 .with_row_stride(200);
2448 assert!(result.is_err());
2449 }
2450
2451 #[test]
2452 fn set_format_clears_row_stride() {
2453 let mut t = TensorDyn::new(&[480, 640, 3], DType::U8, None, None).unwrap();
2454 t.set_format(PixelFormat::Rgb).unwrap();
2455 t.set_row_stride(2048).unwrap();
2456 assert_eq!(t.row_stride(), Some(2048));
2457
2458 let _ = t.set_format(PixelFormat::Bgra);
2460 assert_eq!(t.row_stride(), Some(2048));
2461
2462 t.set_format(PixelFormat::Rgb).unwrap();
2464 assert_eq!(t.row_stride(), Some(2048));
2465
2466 t.reshape(&[480 * 640 * 3]).unwrap();
2468 assert_eq!(t.row_stride(), None);
2469 assert_eq!(t.format(), None);
2470 }
2471
2472 #[test]
2473 fn set_format_different_compatible_clears_stride() {
2474 let mut t = TensorDyn::new(&[480, 640, 4], DType::U8, None, None).unwrap();
2477 t.set_format(PixelFormat::Rgba).unwrap();
2478 t.set_row_stride(4096).unwrap();
2479 assert_eq!(t.row_stride(), Some(4096));
2480
2481 t.set_format(PixelFormat::Bgra).unwrap();
2483 assert_eq!(t.format(), Some(PixelFormat::Bgra));
2484 assert_eq!(t.row_stride(), None);
2485 }
2486
2487 #[test]
2488 fn set_format_same_preserves_stride() {
2489 let mut t = TensorDyn::image(
2490 100,
2491 100,
2492 PixelFormat::Rgb,
2493 DType::U8,
2494 None,
2495 crate::CpuAccess::ReadWrite,
2496 )
2497 .unwrap();
2498 t.set_row_stride(512).unwrap();
2499 t.set_format(PixelFormat::Rgb).unwrap();
2501 assert_eq!(t.row_stride(), Some(512));
2502 }
2503
2504 #[test]
2505 fn effective_row_stride_planar() {
2506 let t = TensorDyn::image(
2507 640,
2508 480,
2509 PixelFormat::PlanarRgb,
2510 DType::U8,
2511 None,
2512 crate::CpuAccess::ReadWrite,
2513 )
2514 .unwrap();
2515 assert_eq!(t.effective_row_stride(), Some(640)); }
2517
2518 #[test]
2519 fn effective_row_stride_nv12() {
2520 let t = TensorDyn::image(
2521 640,
2522 480,
2523 PixelFormat::Nv12,
2524 DType::U8,
2525 None,
2526 crate::CpuAccess::ReadWrite,
2527 )
2528 .unwrap();
2529 assert_eq!(t.effective_row_stride(), Some(640)); }
2531
2532 #[test]
2533 fn map_rejects_strided_tensor() {
2534 let mut t = Tensor::<u8>::image(
2535 100,
2536 100,
2537 PixelFormat::Rgba,
2538 Some(TensorMemory::Mem),
2539 crate::CpuAccess::ReadWrite,
2540 )
2541 .unwrap();
2542 assert!(t.map().is_ok());
2544 t.set_row_stride(512).unwrap();
2546 let err = t.map();
2547 assert!(err.is_err());
2548 }
2549
2550 #[test]
2553 fn plane_offset_default_none() {
2554 let t = TensorDyn::image(
2555 100,
2556 100,
2557 PixelFormat::Rgba,
2558 DType::U8,
2559 None,
2560 crate::CpuAccess::ReadWrite,
2561 )
2562 .unwrap();
2563 assert_eq!(t.plane_offset(), None);
2564 }
2565
2566 #[test]
2567 fn set_plane_offset_basic() {
2568 let mut t = TensorDyn::image(
2569 100,
2570 100,
2571 PixelFormat::Rgba,
2572 DType::U8,
2573 None,
2574 crate::CpuAccess::ReadWrite,
2575 )
2576 .unwrap();
2577 t.set_plane_offset(4096);
2578 assert_eq!(t.plane_offset(), Some(4096));
2579 }
2580
2581 #[test]
2582 fn set_plane_offset_zero() {
2583 let mut t = TensorDyn::image(
2584 100,
2585 100,
2586 PixelFormat::Rgb,
2587 DType::U8,
2588 None,
2589 crate::CpuAccess::ReadWrite,
2590 )
2591 .unwrap();
2592 t.set_plane_offset(0);
2593 assert_eq!(t.plane_offset(), Some(0));
2594 }
2595
2596 #[test]
2597 fn set_plane_offset_no_format() {
2598 let mut t = TensorDyn::new(&[480, 640, 3], DType::U8, None, None).unwrap();
2600 t.set_plane_offset(4096);
2601 assert_eq!(t.plane_offset(), Some(4096));
2602 }
2603
2604 #[test]
2605 fn set_format_clears_plane_offset() {
2606 let mut t = TensorDyn::new(&[480, 640, 3], DType::U8, None, None).unwrap();
2607 t.set_format(PixelFormat::Rgb).unwrap();
2608 t.set_plane_offset(4096);
2609 assert_eq!(t.plane_offset(), Some(4096));
2610
2611 t.set_format(PixelFormat::Rgb).unwrap();
2613 assert_eq!(t.plane_offset(), Some(4096));
2614
2615 t.reshape(&[480 * 640 * 3]).unwrap();
2617 assert_eq!(t.plane_offset(), None);
2618 assert_eq!(t.format(), None);
2619 }
2620
2621 #[test]
2622 fn map_rejects_out_of_bounds_offset() {
2623 let mut t = Tensor::<u8>::image(
2624 100,
2625 100,
2626 PixelFormat::Rgba,
2627 Some(TensorMemory::Mem),
2628 crate::CpuAccess::ReadWrite,
2629 )
2630 .unwrap();
2631 assert!(t.map().is_ok());
2633 t.set_plane_offset(4096);
2636 assert!(t.map().is_err());
2637 }
2638
2639 #[test]
2640 fn mem_subview_in_bounds_maps_at_offset() {
2641 let parent = Tensor::<u8>::image(
2644 100,
2645 100,
2646 PixelFormat::Rgba,
2647 Some(TensorMemory::Mem),
2648 crate::CpuAccess::ReadWrite,
2649 )
2650 .unwrap();
2651 let view = parent.subview(4096, &[10, 10, 4]).unwrap();
2653 assert_eq!(view.plane_offset(), Some(4096));
2654 assert!(view.map().is_ok());
2655 }
2656
2657 #[test]
2658 fn dyn_batch_dispatches_every_dtype() {
2659 use DType::*;
2663 for dt in [U8, I8, U16, I16, U32, I32, U64, I64, F16, F32, F64] {
2664 let parent = TensorDyn::new(&[2, 4], dt, Some(TensorMemory::Mem), None).unwrap();
2665 let view = parent.batch(1).unwrap();
2666 assert_eq!(view.dtype(), dt, "batch must preserve dtype {dt:?}");
2667 assert_eq!(view.shape(), &[4], "{dt:?}");
2668 }
2669 }
2670
2671 #[test]
2672 fn map_accepts_zero_offset_tensor() {
2673 let mut t = Tensor::<u8>::image(
2674 100,
2675 100,
2676 PixelFormat::Rgba,
2677 Some(TensorMemory::Mem),
2678 crate::CpuAccess::ReadWrite,
2679 )
2680 .unwrap();
2681 t.set_plane_offset(0);
2682 assert!(t.map().is_ok());
2684 }
2685
2686 #[test]
2687 fn dyn_configure_image_nv12() {
2688 let mut t = TensorDyn::image(
2689 640,
2690 480,
2691 PixelFormat::Rgb,
2692 DType::U8,
2693 None,
2694 crate::CpuAccess::ReadWrite,
2695 )
2696 .unwrap();
2697 t.configure_image(320, 240, PixelFormat::Nv12).unwrap();
2698 assert_eq!(t.format(), Some(PixelFormat::Nv12));
2699 assert_eq!((t.width(), t.height()), (Some(320), Some(240)));
2700 }
2701
2702 #[test]
2703 fn tensordyn_colorimetry_roundtrip() {
2704 use crate::{ColorEncoding, Colorimetry, DType, PixelFormat};
2705 let mut t = TensorDyn::image(
2706 1280,
2707 720,
2708 PixelFormat::Nv12,
2709 DType::U8,
2710 None,
2711 crate::CpuAccess::ReadWrite,
2712 )
2713 .unwrap();
2714 assert_eq!(t.colorimetry(), None);
2715 let c = Colorimetry::default().with_encoding(ColorEncoding::Bt709);
2716 t.set_colorimetry(Some(c));
2717 assert_eq!(t.colorimetry(), Some(c));
2718 }
2719
2720 #[test]
2721 fn from_planes_propagates_plane_offset() {
2722 let mut luma =
2723 Tensor::<u8>::new(&[480, 640], Some(TensorMemory::Mem), Some("luma")).unwrap();
2724 luma.set_plane_offset(4096);
2725 let chroma =
2726 Tensor::<u8>::new(&[240, 640], Some(TensorMemory::Mem), Some("chroma")).unwrap();
2727 let combined = Tensor::<u8>::from_planes(luma, chroma, PixelFormat::Nv12).unwrap();
2728 assert_eq!(combined.plane_offset(), Some(4096));
2729 }
2730
2731 #[test]
2732 fn cuda_passthrough_none_for_mem_tensor() {
2733 let t: TensorDyn = Tensor::<f32>::new(&[10], Some(TensorMemory::Mem), None)
2736 .unwrap()
2737 .into();
2738 assert!(t.cuda().is_none());
2739 assert!(t.cuda_map().is_none());
2740 }
2741}