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
35macro_rules! dispatch {
37 ($self:expr, $method:ident $(, $arg:expr)*) => {
38 match $self {
39 TensorDyn::U8(t) => t.$method($($arg),*),
40 TensorDyn::I8(t) => t.$method($($arg),*),
41 TensorDyn::U16(t) => t.$method($($arg),*),
42 TensorDyn::I16(t) => t.$method($($arg),*),
43 TensorDyn::U32(t) => t.$method($($arg),*),
44 TensorDyn::I32(t) => t.$method($($arg),*),
45 TensorDyn::U64(t) => t.$method($($arg),*),
46 TensorDyn::I64(t) => t.$method($($arg),*),
47 TensorDyn::F16(t) => t.$method($($arg),*),
48 TensorDyn::F32(t) => t.$method($($arg),*),
49 TensorDyn::F64(t) => t.$method($($arg),*),
50 }
51 };
52}
53
54macro_rules! dyn_fanout {
58 ($self:expr, $method:ident $(, $arg:expr)*) => {
59 match $self {
60 TensorDyn::U8(t) => t.$method($($arg),*).map(TensorDyn::U8),
61 TensorDyn::I8(t) => t.$method($($arg),*).map(TensorDyn::I8),
62 TensorDyn::U16(t) => t.$method($($arg),*).map(TensorDyn::U16),
63 TensorDyn::I16(t) => t.$method($($arg),*).map(TensorDyn::I16),
64 TensorDyn::U32(t) => t.$method($($arg),*).map(TensorDyn::U32),
65 TensorDyn::I32(t) => t.$method($($arg),*).map(TensorDyn::I32),
66 TensorDyn::U64(t) => t.$method($($arg),*).map(TensorDyn::U64),
67 TensorDyn::I64(t) => t.$method($($arg),*).map(TensorDyn::I64),
68 TensorDyn::F16(t) => t.$method($($arg),*).map(TensorDyn::F16),
69 TensorDyn::F32(t) => t.$method($($arg),*).map(TensorDyn::F32),
70 TensorDyn::F64(t) => t.$method($($arg),*).map(TensorDyn::F64),
71 }
72 };
73}
74
75macro_rules! downcast_methods {
77 ($variant:ident, $ty:ty, $as_name:ident, $as_mut_name:ident, $into_name:ident) => {
78 pub fn $as_name(&self) -> Option<&Tensor<$ty>> {
80 match self {
81 Self::$variant(t) => Some(t),
82 _ => None,
83 }
84 }
85
86 pub fn $as_mut_name(&mut self) -> Option<&mut Tensor<$ty>> {
88 match self {
89 Self::$variant(t) => Some(t),
90 _ => None,
91 }
92 }
93
94 #[allow(clippy::result_large_err)]
97 pub fn $into_name(self) -> Result<Tensor<$ty>, Self> {
98 match self {
99 Self::$variant(t) => Ok(t),
100 other => Err(other),
101 }
102 }
103 };
104}
105
106impl TensorDyn {
107 pub fn dtype(&self) -> DType {
109 match self {
110 Self::U8(_) => DType::U8,
111 Self::I8(_) => DType::I8,
112 Self::U16(_) => DType::U16,
113 Self::I16(_) => DType::I16,
114 Self::U32(_) => DType::U32,
115 Self::I32(_) => DType::I32,
116 Self::U64(_) => DType::U64,
117 Self::I64(_) => DType::I64,
118 Self::F16(_) => DType::F16,
119 Self::F32(_) => DType::F32,
120 Self::F64(_) => DType::F64,
121 }
122 }
123
124 pub fn shape(&self) -> &[usize] {
126 dispatch!(self, shape)
127 }
128
129 pub fn name(&self) -> String {
131 dispatch!(self, name)
132 }
133
134 pub fn format(&self) -> Option<PixelFormat> {
136 dispatch!(self, format)
137 }
138
139 pub fn width(&self) -> Option<usize> {
141 dispatch!(self, width)
142 }
143
144 pub fn height(&self) -> Option<usize> {
146 dispatch!(self, height)
147 }
148
149 pub fn size(&self) -> usize {
151 dispatch!(self, size)
152 }
153
154 pub fn memory(&self) -> TensorMemory {
156 dispatch!(self, memory)
157 }
158
159 pub fn reshape(&mut self, shape: &[usize]) -> crate::Result<()> {
161 dispatch!(self, reshape, shape)
162 }
163
164 pub fn set_format(&mut self, format: PixelFormat) -> crate::Result<()> {
182 dispatch!(self, set_format, format)
183 }
184
185 pub fn with_format(mut self, format: PixelFormat) -> crate::Result<Self> {
202 self.set_format(format)?;
203 Ok(self)
204 }
205
206 pub fn colorimetry(&self) -> Option<crate::Colorimetry> {
208 dispatch!(self, colorimetry)
209 }
210
211 pub fn set_colorimetry(&mut self, c: Option<crate::Colorimetry>) {
213 dispatch!(self, set_colorimetry, c)
214 }
215
216 pub fn with_colorimetry(mut self, c: crate::Colorimetry) -> Self {
218 self.set_colorimetry(Some(c));
219 self
220 }
221
222 pub fn row_stride(&self) -> Option<usize> {
224 dispatch!(self, row_stride)
225 }
226
227 pub fn effective_row_stride(&self) -> Option<usize> {
229 dispatch!(self, effective_row_stride)
230 }
231
232 pub fn configure_image(
235 &mut self,
236 width: usize,
237 height: usize,
238 format: PixelFormat,
239 ) -> crate::Result<()> {
240 dispatch!(self, configure_image, width, height, format)
241 }
242
243 pub fn set_row_stride(&mut self, stride: usize) -> crate::Result<()> {
249 dispatch!(self, set_row_stride, stride)
250 }
251
252 pub fn with_row_stride(mut self, stride: usize) -> crate::Result<Self> {
254 self.set_row_stride(stride)?;
255 Ok(self)
256 }
257
258 pub fn plane_offset(&self) -> Option<usize> {
260 dispatch!(self, plane_offset)
261 }
262
263 pub fn view_origin(&self) -> Option<crate::ViewOrigin> {
267 dispatch!(self, view_origin)
268 }
269
270 pub fn set_plane_offset(&mut self, offset: usize) {
272 dispatch!(self, set_plane_offset, offset)
273 }
274
275 pub fn batch(&self, n: usize) -> crate::Result<TensorDyn> {
278 dyn_fanout!(self, batch, n)
279 }
280
281 pub fn view(&self, region: crate::Region) -> crate::Result<TensorDyn> {
284 dyn_fanout!(self, view, region)
285 }
286
287 pub fn cuda(&self) -> Option<&crate::cuda::CudaHandle> {
292 dispatch!(self, cuda)
293 }
294
295 pub fn cuda_map(&self) -> Option<crate::cuda::CudaMap<'_>> {
318 dispatch!(self, cuda_map)
319 }
320
321 pub fn quantization(&self) -> Option<&crate::Quantization> {
325 match self {
326 Self::U8(t) => t.quantization(),
327 Self::I8(t) => t.quantization(),
328 Self::U16(t) => t.quantization(),
329 Self::I16(t) => t.quantization(),
330 Self::U32(t) => t.quantization(),
331 Self::I32(t) => t.quantization(),
332 Self::U64(t) => t.quantization(),
333 Self::I64(t) => t.quantization(),
334 Self::F16(_) | Self::F32(_) | Self::F64(_) => None,
335 }
336 }
337
338 pub fn set_quantization(&mut self, q: crate::Quantization) -> crate::Result<()> {
342 match self {
343 Self::U8(t) => t.set_quantization(q),
344 Self::I8(t) => t.set_quantization(q),
345 Self::U16(t) => t.set_quantization(q),
346 Self::I16(t) => t.set_quantization(q),
347 Self::U32(t) => t.set_quantization(q),
348 Self::I32(t) => t.set_quantization(q),
349 Self::U64(t) => t.set_quantization(q),
350 Self::I64(t) => t.set_quantization(q),
351 Self::F16(_) | Self::F32(_) | Self::F64(_) => Err(crate::Error::QuantizationInvalid {
352 field: "dtype_is_integer",
353 expected: "integer tensor dtype (u8/i8/u16/i16/u32/i32/u64/i64)".to_string(),
354 got: format!("{:?}", self.dtype()),
355 }),
356 }
357 }
358
359 pub fn with_quantization(mut self, q: crate::Quantization) -> crate::Result<Self> {
362 self.set_quantization(q)?;
363 Ok(self)
364 }
365
366 pub fn clear_quantization(&mut self) {
368 match self {
369 Self::U8(t) => t.clear_quantization(),
370 Self::I8(t) => t.clear_quantization(),
371 Self::U16(t) => t.clear_quantization(),
372 Self::I16(t) => t.clear_quantization(),
373 Self::U32(t) => t.clear_quantization(),
374 Self::I32(t) => t.clear_quantization(),
375 Self::U64(t) => t.clear_quantization(),
376 Self::I64(t) => t.clear_quantization(),
377 Self::F16(_) | Self::F32(_) | Self::F64(_) => {}
378 }
379 }
380
381 #[cfg(unix)]
383 pub fn clone_fd(&self) -> crate::Result<std::os::fd::OwnedFd> {
384 dispatch!(self, clone_fd)
385 }
386
387 #[cfg(target_os = "linux")]
398 pub fn dmabuf_clone(&self) -> crate::Result<std::os::fd::OwnedFd> {
399 if self.memory() != TensorMemory::Dma {
400 return Err(crate::Error::NotImplemented(format!(
401 "dmabuf_clone requires DMA-backed tensor, got {:?}",
402 self.memory()
403 )));
404 }
405 self.clone_fd()
406 }
407
408 #[cfg(target_os = "linux")]
419 pub fn dmabuf(&self) -> crate::Result<std::os::fd::BorrowedFd<'_>> {
420 dispatch!(self, dmabuf)
421 }
422
423 pub fn is_multiplane(&self) -> bool {
425 dispatch!(self, is_multiplane)
426 }
427
428 pub fn buffer_identity(&self) -> &crate::BufferIdentity {
438 dispatch!(self, buffer_identity)
439 }
440
441 pub fn aliases(&self, other: &Self) -> bool {
459 if self.buffer_identity().id() == other.buffer_identity().id() {
460 return true;
461 }
462 if self.memory() != other.memory() {
463 return false;
464 }
465 #[cfg(target_os = "linux")]
466 if self.memory() == TensorMemory::Dma {
467 use std::os::fd::AsRawFd;
468 if let (Ok(a), Ok(b)) = (self.dmabuf(), other.dmabuf()) {
469 return a.as_raw_fd() == b.as_raw_fd();
470 }
471 }
472 false
473 }
474
475 downcast_methods!(U8, u8, as_u8, as_u8_mut, into_u8);
478 downcast_methods!(I8, i8, as_i8, as_i8_mut, into_i8);
479 downcast_methods!(U16, u16, as_u16, as_u16_mut, into_u16);
480 downcast_methods!(I16, i16, as_i16, as_i16_mut, into_i16);
481 downcast_methods!(U32, u32, as_u32, as_u32_mut, into_u32);
482 downcast_methods!(I32, i32, as_i32, as_i32_mut, into_i32);
483 downcast_methods!(U64, u64, as_u64, as_u64_mut, into_u64);
484 downcast_methods!(I64, i64, as_i64, as_i64_mut, into_i64);
485 downcast_methods!(F16, f16, as_f16, as_f16_mut, into_f16);
486 downcast_methods!(F32, f32, as_f32, as_f32_mut, into_f32);
487 downcast_methods!(F64, f64, as_f64, as_f64_mut, into_f64);
488
489 pub fn new(
491 shape: &[usize],
492 dtype: DType,
493 memory: Option<TensorMemory>,
494 name: Option<&str>,
495 ) -> crate::Result<Self> {
496 match dtype {
497 DType::U8 => Tensor::<u8>::new(shape, memory, name).map(Self::U8),
498 DType::I8 => Tensor::<i8>::new(shape, memory, name).map(Self::I8),
499 DType::U16 => Tensor::<u16>::new(shape, memory, name).map(Self::U16),
500 DType::I16 => Tensor::<i16>::new(shape, memory, name).map(Self::I16),
501 DType::U32 => Tensor::<u32>::new(shape, memory, name).map(Self::U32),
502 DType::I32 => Tensor::<i32>::new(shape, memory, name).map(Self::I32),
503 DType::U64 => Tensor::<u64>::new(shape, memory, name).map(Self::U64),
504 DType::I64 => Tensor::<i64>::new(shape, memory, name).map(Self::I64),
505 DType::F16 => Tensor::<f16>::new(shape, memory, name).map(Self::F16),
506 DType::F32 => Tensor::<f32>::new(shape, memory, name).map(Self::F32),
507 DType::F64 => Tensor::<f64>::new(shape, memory, name).map(Self::F64),
508 }
509 }
510
511 #[cfg(unix)]
513 pub fn from_fd(
514 fd: std::os::fd::OwnedFd,
515 shape: &[usize],
516 dtype: DType,
517 name: Option<&str>,
518 ) -> crate::Result<Self> {
519 match dtype {
520 DType::U8 => Tensor::<u8>::from_fd(fd, shape, name).map(Self::U8),
521 DType::I8 => Tensor::<i8>::from_fd(fd, shape, name).map(Self::I8),
522 DType::U16 => Tensor::<u16>::from_fd(fd, shape, name).map(Self::U16),
523 DType::I16 => Tensor::<i16>::from_fd(fd, shape, name).map(Self::I16),
524 DType::U32 => Tensor::<u32>::from_fd(fd, shape, name).map(Self::U32),
525 DType::I32 => Tensor::<i32>::from_fd(fd, shape, name).map(Self::I32),
526 DType::U64 => Tensor::<u64>::from_fd(fd, shape, name).map(Self::U64),
527 DType::I64 => Tensor::<i64>::from_fd(fd, shape, name).map(Self::I64),
528 DType::F16 => Tensor::<f16>::from_fd(fd, shape, name).map(Self::F16),
529 DType::F32 => Tensor::<f32>::from_fd(fd, shape, name).map(Self::F32),
530 DType::F64 => Tensor::<f64>::from_fd(fd, shape, name).map(Self::F64),
531 }
532 }
533
534 pub unsafe fn from_foreign_ptr(
547 ptr: *mut u8,
548 shape: &[usize],
549 dtype: DType,
550 owner: Option<crate::ForeignOwner>,
551 name: Option<&str>,
552 ) -> crate::Result<Self> {
553 match dtype {
554 DType::U8 => Tensor::<u8>::from_foreign(ptr.cast(), shape, owner, name).map(Self::U8),
555 DType::I8 => Tensor::<i8>::from_foreign(ptr.cast(), shape, owner, name).map(Self::I8),
556 DType::U16 => Tensor::<u16>::from_foreign(ptr.cast(), shape, owner, name).map(Self::U16),
557 DType::I16 => Tensor::<i16>::from_foreign(ptr.cast(), shape, owner, name).map(Self::I16),
558 DType::U32 => Tensor::<u32>::from_foreign(ptr.cast(), shape, owner, name).map(Self::U32),
559 DType::I32 => Tensor::<i32>::from_foreign(ptr.cast(), shape, owner, name).map(Self::I32),
560 DType::U64 => Tensor::<u64>::from_foreign(ptr.cast(), shape, owner, name).map(Self::U64),
561 DType::I64 => Tensor::<i64>::from_foreign(ptr.cast(), shape, owner, name).map(Self::I64),
562 DType::F16 => Tensor::<f16>::from_foreign(ptr.cast(), shape, owner, name).map(Self::F16),
563 DType::F32 => Tensor::<f32>::from_foreign(ptr.cast(), shape, owner, name).map(Self::F32),
564 DType::F64 => Tensor::<f64>::from_foreign(ptr.cast(), shape, owner, name).map(Self::F64),
565 }
566 }
567
568 #[cfg(target_os = "macos")]
576 pub unsafe fn from_iosurface(
577 surface_ref: *mut std::ffi::c_void,
578 shape: &[usize],
579 dtype: DType,
580 name: Option<&str>,
581 ) -> crate::Result<Self> {
582 unsafe {
583 match dtype {
584 DType::U8 => Tensor::<u8>::from_iosurface(surface_ref, shape, name).map(Self::U8),
585 DType::I8 => Tensor::<i8>::from_iosurface(surface_ref, shape, name).map(Self::I8),
586 DType::U16 => {
587 Tensor::<u16>::from_iosurface(surface_ref, shape, name).map(Self::U16)
588 }
589 DType::I16 => {
590 Tensor::<i16>::from_iosurface(surface_ref, shape, name).map(Self::I16)
591 }
592 DType::U32 => {
593 Tensor::<u32>::from_iosurface(surface_ref, shape, name).map(Self::U32)
594 }
595 DType::I32 => {
596 Tensor::<i32>::from_iosurface(surface_ref, shape, name).map(Self::I32)
597 }
598 DType::U64 => {
599 Tensor::<u64>::from_iosurface(surface_ref, shape, name).map(Self::U64)
600 }
601 DType::I64 => {
602 Tensor::<i64>::from_iosurface(surface_ref, shape, name).map(Self::I64)
603 }
604 DType::F16 => {
605 Tensor::<f16>::from_iosurface(surface_ref, shape, name).map(Self::F16)
606 }
607 DType::F32 => {
608 Tensor::<f32>::from_iosurface(surface_ref, shape, name).map(Self::F32)
609 }
610 DType::F64 => {
611 Tensor::<f64>::from_iosurface(surface_ref, shape, name).map(Self::F64)
612 }
613 }
614 }
615 }
616
617 #[cfg(target_os = "macos")]
620 pub fn iosurface_id(&self) -> Option<u32> {
621 dispatch!(self, iosurface_id)
622 }
623
624 #[cfg(target_os = "macos")]
628 pub fn iosurface_ref(&self) -> Option<*mut std::ffi::c_void> {
629 dispatch!(self, iosurface_ref)
630 }
631
632 #[cfg(target_os = "macos")]
637 pub fn iosurface_physical_dims(&self) -> Option<(usize, usize)> {
638 dispatch!(self, iosurface_physical_dims)
639 }
640
641 pub fn image(
659 width: usize,
660 height: usize,
661 format: PixelFormat,
662 dtype: DType,
663 memory: Option<TensorMemory>,
664 ) -> crate::Result<Self> {
665 match dtype {
666 DType::U8 => Tensor::<u8>::image(width, height, format, memory).map(Self::U8),
667 DType::I8 => Tensor::<i8>::image(width, height, format, memory).map(Self::I8),
668 DType::U16 => Tensor::<u16>::image(width, height, format, memory).map(Self::U16),
669 DType::I16 => Tensor::<i16>::image(width, height, format, memory).map(Self::I16),
670 DType::U32 => Tensor::<u32>::image(width, height, format, memory).map(Self::U32),
671 DType::I32 => Tensor::<i32>::image(width, height, format, memory).map(Self::I32),
672 DType::U64 => Tensor::<u64>::image(width, height, format, memory).map(Self::U64),
673 DType::I64 => Tensor::<i64>::image(width, height, format, memory).map(Self::I64),
674 DType::F16 => Tensor::<f16>::image(width, height, format, memory).map(Self::F16),
675 DType::F32 => Tensor::<f32>::image(width, height, format, memory).map(Self::F32),
676 DType::F64 => Tensor::<f64>::image(width, height, format, memory).map(Self::F64),
677 }
678 }
679
680 pub fn image_with_stride(
706 width: usize,
707 height: usize,
708 format: PixelFormat,
709 dtype: DType,
710 row_stride_bytes: usize,
711 memory: Option<TensorMemory>,
712 ) -> crate::Result<Self> {
713 match dtype {
714 DType::U8 => {
715 Tensor::<u8>::image_with_stride(width, height, format, row_stride_bytes, memory)
716 .map(Self::U8)
717 }
718 DType::I8 => {
719 Tensor::<i8>::image_with_stride(width, height, format, row_stride_bytes, memory)
720 .map(Self::I8)
721 }
722 DType::U16 => {
723 Tensor::<u16>::image_with_stride(width, height, format, row_stride_bytes, memory)
724 .map(Self::U16)
725 }
726 DType::I16 => {
727 Tensor::<i16>::image_with_stride(width, height, format, row_stride_bytes, memory)
728 .map(Self::I16)
729 }
730 DType::U32 => {
731 Tensor::<u32>::image_with_stride(width, height, format, row_stride_bytes, memory)
732 .map(Self::U32)
733 }
734 DType::I32 => {
735 Tensor::<i32>::image_with_stride(width, height, format, row_stride_bytes, memory)
736 .map(Self::I32)
737 }
738 DType::U64 => {
739 Tensor::<u64>::image_with_stride(width, height, format, row_stride_bytes, memory)
740 .map(Self::U64)
741 }
742 DType::I64 => {
743 Tensor::<i64>::image_with_stride(width, height, format, row_stride_bytes, memory)
744 .map(Self::I64)
745 }
746 DType::F16 => {
747 Tensor::<f16>::image_with_stride(width, height, format, row_stride_bytes, memory)
748 .map(Self::F16)
749 }
750 DType::F32 => {
751 Tensor::<f32>::image_with_stride(width, height, format, row_stride_bytes, memory)
752 .map(Self::F32)
753 }
754 DType::F64 => {
755 Tensor::<f64>::image_with_stride(width, height, format, row_stride_bytes, memory)
756 .map(Self::F64)
757 }
758 }
759 }
760}
761
762impl From<Tensor<u8>> for TensorDyn {
765 fn from(t: Tensor<u8>) -> Self {
766 Self::U8(t)
767 }
768}
769
770impl From<Tensor<i8>> for TensorDyn {
771 fn from(t: Tensor<i8>) -> Self {
772 Self::I8(t)
773 }
774}
775
776impl From<Tensor<u16>> for TensorDyn {
777 fn from(t: Tensor<u16>) -> Self {
778 Self::U16(t)
779 }
780}
781
782impl From<Tensor<i16>> for TensorDyn {
783 fn from(t: Tensor<i16>) -> Self {
784 Self::I16(t)
785 }
786}
787
788impl From<Tensor<u32>> for TensorDyn {
789 fn from(t: Tensor<u32>) -> Self {
790 Self::U32(t)
791 }
792}
793
794impl From<Tensor<i32>> for TensorDyn {
795 fn from(t: Tensor<i32>) -> Self {
796 Self::I32(t)
797 }
798}
799
800impl From<Tensor<u64>> for TensorDyn {
801 fn from(t: Tensor<u64>) -> Self {
802 Self::U64(t)
803 }
804}
805
806impl From<Tensor<i64>> for TensorDyn {
807 fn from(t: Tensor<i64>) -> Self {
808 Self::I64(t)
809 }
810}
811
812impl From<Tensor<f16>> for TensorDyn {
813 fn from(t: Tensor<f16>) -> Self {
814 Self::F16(t)
815 }
816}
817
818impl From<Tensor<f32>> for TensorDyn {
819 fn from(t: Tensor<f32>) -> Self {
820 Self::F32(t)
821 }
822}
823
824impl From<Tensor<f64>> for TensorDyn {
825 fn from(t: Tensor<f64>) -> Self {
826 Self::F64(t)
827 }
828}
829
830impl fmt::Debug for TensorDyn {
831 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
832 dispatch!(self, fmt, f)
833 }
834}
835
836#[cfg(test)]
837mod tests {
838 use super::*;
839
840 #[test]
841 fn from_typed_tensor() {
842 let t = Tensor::<u8>::new(&[10], None, None).unwrap();
843 let dyn_t: TensorDyn = t.into();
844 assert_eq!(dyn_t.dtype(), DType::U8);
845 assert_eq!(dyn_t.shape(), &[10]);
846 }
847
848 #[test]
849 fn from_foreign_ptr_wraps_borrowed_memory() {
850 use crate::TensorMapTrait;
851 let mut vec: Vec<f32> = vec![0.0; 4];
854 let ptr = vec.as_mut_ptr() as *mut u8;
855 let owner: crate::ForeignOwner = Box::new(vec);
856 let t = unsafe {
857 TensorDyn::from_foreign_ptr(ptr, &[2, 2], DType::F32, Some(owner), Some("trt_output"))
858 }
859 .unwrap();
860 assert_eq!(t.dtype(), DType::F32);
861 assert_eq!(t.memory(), TensorMemory::Mem);
862 assert_eq!(t.shape(), &[2, 2]);
863 {
864 let mut m = t.as_f32().unwrap().map().unwrap();
865 m.as_mut_slice().copy_from_slice(&[1.0, 2.0, 3.0, 4.0]);
866 }
867 let m = t.as_f32().unwrap().map().unwrap();
868 assert_eq!(m.as_slice(), &[1.0, 2.0, 3.0, 4.0]);
869 }
870
871 #[test]
880 fn from_foreign_ptr_rejects_null_ptr() {
881 let err = unsafe {
882 TensorDyn::from_foreign_ptr(std::ptr::null_mut(), &[4], DType::U8, None, None)
883 }
884 .unwrap_err();
885 assert!(
887 matches!(err, crate::error::Error::InvalidArgument(ref m) if m.contains("non-null")),
888 "expected InvalidArgument(non-null), got {err:?}"
889 );
890 }
891
892 #[test]
893 fn from_foreign_ptr_rejects_empty_shape() {
894 let mut dummy: u8 = 0;
895 let err = unsafe {
896 TensorDyn::from_foreign_ptr(&mut dummy as *mut u8, &[], DType::U8, None, None)
897 }
898 .unwrap_err();
899 assert!(
900 matches!(err, crate::error::Error::InvalidSize(0)),
901 "expected InvalidSize(0) for empty shape, got {err:?}"
902 );
903 }
904
905 #[test]
906 fn from_foreign_ptr_rejects_overflow_shape() {
907 let mut dummy: u8 = 0;
908 let huge = [usize::MAX / 2 + 1, 2];
909 let err = unsafe {
910 TensorDyn::from_foreign_ptr(&mut dummy, &huge, DType::U8, None, None)
911 }
912 .unwrap_err();
913 assert!(
914 matches!(err, crate::error::Error::InvalidArgument(ref m) if m.contains("overflow")),
915 "expected InvalidArgument(overflow), got {err:?}"
916 );
917 }
918
919 #[test]
920 fn from_foreign_ptr_u8_dtype_dispatch() {
921 let mut buf: Vec<u8> = vec![1, 2, 3, 4];
924 let ptr = buf.as_mut_ptr();
925 let owner: crate::ForeignOwner = Box::new(buf);
926 let t = unsafe {
927 TensorDyn::from_foreign_ptr(ptr, &[4], DType::U8, Some(owner), Some("u8_foreign"))
928 }
929 .unwrap();
930 assert_eq!(t.dtype(), DType::U8);
931 assert_eq!(t.shape(), &[4]);
932 let m = t.as_u8().unwrap().map().unwrap();
933 use crate::TensorMapTrait;
934 assert_eq!(m.as_slice(), &[1u8, 2, 3, 4]);
935 }
936
937 #[test]
938 fn downcast_ref() {
939 let t = Tensor::<u8>::new(&[10], None, None).unwrap();
940 let dyn_t: TensorDyn = t.into();
941 assert!(dyn_t.as_u8().is_some());
942 assert!(dyn_t.as_i8().is_none());
943 }
944
945 #[test]
946 fn downcast_into() {
947 let t = Tensor::<u8>::new(&[10], None, None).unwrap();
948 let dyn_t: TensorDyn = t.into();
949 let back = dyn_t.into_u8().unwrap();
950 assert_eq!(back.shape(), &[10]);
951 }
952
953 #[test]
954 fn image_accessors() {
955 let t = Tensor::<u8>::image(640, 480, PixelFormat::Rgba, None).unwrap();
956 let dyn_t: TensorDyn = t.into();
957 assert_eq!(dyn_t.format(), Some(PixelFormat::Rgba));
958 assert_eq!(dyn_t.width(), Some(640));
959 assert_eq!(dyn_t.height(), Some(480));
960 assert!(!dyn_t.is_multiplane());
961 }
962
963 #[test]
964 fn image_constructor() {
965 let dyn_t = TensorDyn::image(640, 480, PixelFormat::Rgb, DType::U8, None).unwrap();
966 assert_eq!(dyn_t.dtype(), DType::U8);
967 assert_eq!(dyn_t.format(), Some(PixelFormat::Rgb));
968 assert_eq!(dyn_t.width(), Some(640));
969 }
970
971 #[test]
972 fn image_constructor_i8() {
973 let dyn_t = TensorDyn::image(640, 480, PixelFormat::Rgb, DType::I8, None).unwrap();
974 assert_eq!(dyn_t.dtype(), DType::I8);
975 assert_eq!(dyn_t.format(), Some(PixelFormat::Rgb));
976 }
977
978 #[test]
979 fn set_format_packed() {
980 let mut t = TensorDyn::new(&[480, 640, 3], DType::U8, None, None).unwrap();
981 assert_eq!(t.format(), None);
982 t.set_format(PixelFormat::Rgb).unwrap();
983 assert_eq!(t.format(), Some(PixelFormat::Rgb));
984 assert_eq!(t.width(), Some(640));
985 assert_eq!(t.height(), Some(480));
986 }
987
988 #[test]
989 fn set_format_planar() {
990 let mut t = TensorDyn::new(&[3, 480, 640], DType::U8, None, None).unwrap();
991 t.set_format(PixelFormat::PlanarRgb).unwrap();
992 assert_eq!(t.format(), Some(PixelFormat::PlanarRgb));
993 assert_eq!(t.width(), Some(640));
994 assert_eq!(t.height(), Some(480));
995 }
996
997 #[test]
998 fn set_format_rejects_wrong_shape() {
999 let mut t = TensorDyn::new(&[480, 640, 4], DType::U8, None, None).unwrap();
1000 assert!(t.set_format(PixelFormat::Rgb).is_err());
1001 }
1002
1003 #[test]
1004 fn with_format_builder() {
1005 let t = TensorDyn::new(&[480, 640, 4], DType::U8, None, None)
1006 .unwrap()
1007 .with_format(PixelFormat::Rgba)
1008 .unwrap();
1009 assert_eq!(t.format(), Some(PixelFormat::Rgba));
1010 assert_eq!(t.width(), Some(640));
1011 assert_eq!(t.height(), Some(480));
1012 }
1013
1014 #[cfg(target_os = "linux")]
1015 #[test]
1016 fn dmabuf_clone_mem_tensor_fails() {
1017 let t = TensorDyn::new(&[480, 640, 3], DType::U8, Some(TensorMemory::Mem), None).unwrap();
1018 assert_eq!(t.memory(), TensorMemory::Mem);
1019 assert!(t.dmabuf_clone().is_err());
1020 }
1021
1022 #[cfg(target_os = "linux")]
1023 #[test]
1024 fn dmabuf_mem_tensor_fails() {
1025 let t = TensorDyn::new(&[480, 640, 3], DType::U8, Some(TensorMemory::Mem), None).unwrap();
1026 assert!(t.dmabuf().is_err());
1027 }
1028
1029 #[test]
1030 fn set_format_semi_planar_nv12() {
1031 let mut t = TensorDyn::new(&[720, 640], DType::U8, Some(TensorMemory::Mem), None).unwrap();
1033 t.set_format(PixelFormat::Nv12).unwrap();
1034 assert_eq!(t.format(), Some(PixelFormat::Nv12));
1035 assert_eq!(t.width(), Some(640));
1036 assert_eq!(t.height(), Some(480));
1037 }
1038
1039 #[test]
1040 fn set_format_semi_planar_nv16() {
1041 let mut t = TensorDyn::new(&[960, 640], DType::U8, Some(TensorMemory::Mem), None).unwrap();
1043 t.set_format(PixelFormat::Nv16).unwrap();
1044 assert_eq!(t.format(), Some(PixelFormat::Nv16));
1045 assert_eq!(t.width(), Some(640));
1046 assert_eq!(t.height(), Some(480));
1047 }
1048
1049 #[test]
1050 fn with_format_rejects_wrong_shape() {
1051 let result = TensorDyn::new(&[480, 640, 4], DType::U8, None, None)
1052 .unwrap()
1053 .with_format(PixelFormat::Rgb);
1054 assert!(result.is_err());
1055 }
1056
1057 #[test]
1058 fn set_format_preserved_after_rejection() {
1059 let mut t = TensorDyn::new(&[480, 640, 3], DType::U8, None, None).unwrap();
1060 t.set_format(PixelFormat::Rgb).unwrap();
1061 assert_eq!(t.format(), Some(PixelFormat::Rgb));
1062
1063 assert!(t.set_format(PixelFormat::Rgba).is_err());
1065
1066 assert_eq!(t.format(), Some(PixelFormat::Rgb));
1068 }
1069
1070 #[test]
1071 fn set_format_idempotent() {
1072 let mut t = TensorDyn::new(&[480, 640, 3], DType::U8, None, None).unwrap();
1073 t.set_format(PixelFormat::Rgb).unwrap();
1074 t.set_format(PixelFormat::Rgb).unwrap();
1075 assert_eq!(t.format(), Some(PixelFormat::Rgb));
1076 assert_eq!(t.width(), Some(640));
1077 assert_eq!(t.height(), Some(480));
1078 }
1079
1080 #[test]
1083 fn set_row_stride_valid() {
1084 let mut t = TensorDyn::image(100, 100, PixelFormat::Rgba, DType::U8, None).unwrap();
1086 t.set_row_stride(512).unwrap();
1087 assert_eq!(t.row_stride(), Some(512));
1088 assert_eq!(t.effective_row_stride(), Some(512));
1089 }
1090
1091 #[test]
1092 fn set_row_stride_equals_min() {
1093 let mut t = TensorDyn::image(100, 100, PixelFormat::Rgb, DType::U8, None).unwrap();
1095 t.set_row_stride(300).unwrap();
1096 assert_eq!(t.row_stride(), Some(300));
1097 }
1098
1099 #[test]
1100 fn set_row_stride_too_small() {
1101 let mut t = TensorDyn::image(64, 100, PixelFormat::Rgba, DType::U8, None).unwrap();
1106 assert!(t.set_row_stride(200).is_err());
1107 assert_eq!(t.row_stride(), None);
1108 }
1109
1110 #[test]
1111 fn set_row_stride_zero() {
1112 let mut t = TensorDyn::image(100, 100, PixelFormat::Rgb, DType::U8, None).unwrap();
1113 assert!(t.set_row_stride(0).is_err());
1114 }
1115
1116 #[test]
1117 fn set_row_stride_requires_format() {
1118 let mut t = TensorDyn::new(&[480, 640, 3], DType::U8, None, None).unwrap();
1119 assert!(t.set_row_stride(2048).is_err());
1120 }
1121
1122 #[test]
1123 fn effective_row_stride_without_stride() {
1124 let t = TensorDyn::image(64, 100, PixelFormat::Rgb, DType::U8, None).unwrap();
1129 assert_eq!(t.row_stride(), None);
1130 assert_eq!(t.effective_row_stride(), Some(192)); }
1132
1133 #[test]
1134 fn effective_row_stride_padded_packed_dma() {
1135 let t = match TensorDyn::image(
1141 100,
1142 100,
1143 PixelFormat::Rgb,
1144 DType::U8,
1145 Some(TensorMemory::Dma),
1146 ) {
1147 Ok(t) if t.memory() == TensorMemory::Dma => t,
1148 _ => return,
1149 };
1150 assert_eq!(t.row_stride(), Some(320));
1151 assert_eq!(t.effective_row_stride(), Some(320));
1152 }
1153
1154 #[test]
1155 fn effective_row_stride_no_format() {
1156 let t = TensorDyn::new(&[480, 640, 3], DType::U8, None, None).unwrap();
1157 assert_eq!(t.effective_row_stride(), None);
1158 }
1159
1160 #[test]
1161 fn with_row_stride_builder() {
1162 let t = TensorDyn::image(100, 100, PixelFormat::Rgba, DType::U8, None)
1163 .unwrap()
1164 .with_row_stride(512)
1165 .unwrap();
1166 assert_eq!(t.row_stride(), Some(512));
1167 assert_eq!(t.effective_row_stride(), Some(512));
1168 }
1169
1170 #[test]
1171 fn with_row_stride_rejects_small() {
1172 let result = TensorDyn::image(100, 100, PixelFormat::Rgba, DType::U8, None)
1173 .unwrap()
1174 .with_row_stride(200);
1175 assert!(result.is_err());
1176 }
1177
1178 #[test]
1179 fn set_format_clears_row_stride() {
1180 let mut t = TensorDyn::new(&[480, 640, 3], DType::U8, None, None).unwrap();
1181 t.set_format(PixelFormat::Rgb).unwrap();
1182 t.set_row_stride(2048).unwrap();
1183 assert_eq!(t.row_stride(), Some(2048));
1184
1185 let _ = t.set_format(PixelFormat::Bgra);
1187 assert_eq!(t.row_stride(), Some(2048));
1188
1189 t.set_format(PixelFormat::Rgb).unwrap();
1191 assert_eq!(t.row_stride(), Some(2048));
1192
1193 t.reshape(&[480 * 640 * 3]).unwrap();
1195 assert_eq!(t.row_stride(), None);
1196 assert_eq!(t.format(), None);
1197 }
1198
1199 #[test]
1200 fn set_format_different_compatible_clears_stride() {
1201 let mut t = TensorDyn::new(&[480, 640, 4], DType::U8, None, None).unwrap();
1204 t.set_format(PixelFormat::Rgba).unwrap();
1205 t.set_row_stride(4096).unwrap();
1206 assert_eq!(t.row_stride(), Some(4096));
1207
1208 t.set_format(PixelFormat::Bgra).unwrap();
1210 assert_eq!(t.format(), Some(PixelFormat::Bgra));
1211 assert_eq!(t.row_stride(), None);
1212 }
1213
1214 #[test]
1215 fn set_format_same_preserves_stride() {
1216 let mut t = TensorDyn::image(100, 100, PixelFormat::Rgb, DType::U8, None).unwrap();
1217 t.set_row_stride(512).unwrap();
1218 t.set_format(PixelFormat::Rgb).unwrap();
1220 assert_eq!(t.row_stride(), Some(512));
1221 }
1222
1223 #[test]
1224 fn effective_row_stride_planar() {
1225 let t = TensorDyn::image(640, 480, PixelFormat::PlanarRgb, DType::U8, None).unwrap();
1226 assert_eq!(t.effective_row_stride(), Some(640)); }
1228
1229 #[test]
1230 fn effective_row_stride_nv12() {
1231 let t = TensorDyn::image(640, 480, PixelFormat::Nv12, DType::U8, None).unwrap();
1232 assert_eq!(t.effective_row_stride(), Some(640)); }
1234
1235 #[test]
1236 fn map_rejects_strided_tensor() {
1237 let mut t =
1238 Tensor::<u8>::image(100, 100, PixelFormat::Rgba, Some(TensorMemory::Mem)).unwrap();
1239 assert!(t.map().is_ok());
1241 t.set_row_stride(512).unwrap();
1243 let err = t.map();
1244 assert!(err.is_err());
1245 }
1246
1247 #[test]
1250 fn plane_offset_default_none() {
1251 let t = TensorDyn::image(100, 100, PixelFormat::Rgba, DType::U8, None).unwrap();
1252 assert_eq!(t.plane_offset(), None);
1253 }
1254
1255 #[test]
1256 fn set_plane_offset_basic() {
1257 let mut t = TensorDyn::image(100, 100, PixelFormat::Rgba, DType::U8, None).unwrap();
1258 t.set_plane_offset(4096);
1259 assert_eq!(t.plane_offset(), Some(4096));
1260 }
1261
1262 #[test]
1263 fn set_plane_offset_zero() {
1264 let mut t = TensorDyn::image(100, 100, PixelFormat::Rgb, DType::U8, None).unwrap();
1265 t.set_plane_offset(0);
1266 assert_eq!(t.plane_offset(), Some(0));
1267 }
1268
1269 #[test]
1270 fn set_plane_offset_no_format() {
1271 let mut t = TensorDyn::new(&[480, 640, 3], DType::U8, None, None).unwrap();
1273 t.set_plane_offset(4096);
1274 assert_eq!(t.plane_offset(), Some(4096));
1275 }
1276
1277 #[test]
1278 fn set_format_clears_plane_offset() {
1279 let mut t = TensorDyn::new(&[480, 640, 3], DType::U8, None, None).unwrap();
1280 t.set_format(PixelFormat::Rgb).unwrap();
1281 t.set_plane_offset(4096);
1282 assert_eq!(t.plane_offset(), Some(4096));
1283
1284 t.set_format(PixelFormat::Rgb).unwrap();
1286 assert_eq!(t.plane_offset(), Some(4096));
1287
1288 t.reshape(&[480 * 640 * 3]).unwrap();
1290 assert_eq!(t.plane_offset(), None);
1291 assert_eq!(t.format(), None);
1292 }
1293
1294 #[test]
1295 fn map_rejects_out_of_bounds_offset() {
1296 let mut t =
1297 Tensor::<u8>::image(100, 100, PixelFormat::Rgba, Some(TensorMemory::Mem)).unwrap();
1298 assert!(t.map().is_ok());
1300 t.set_plane_offset(4096);
1303 assert!(t.map().is_err());
1304 }
1305
1306 #[test]
1307 fn mem_subview_in_bounds_maps_at_offset() {
1308 let parent =
1311 Tensor::<u8>::image(100, 100, PixelFormat::Rgba, Some(TensorMemory::Mem)).unwrap();
1312 let view = parent.subview(4096, &[10, 10, 4]).unwrap();
1314 assert_eq!(view.plane_offset(), Some(4096));
1315 assert!(view.map().is_ok());
1316 }
1317
1318 #[test]
1319 fn dyn_batch_dispatches_every_dtype() {
1320 use DType::*;
1324 for dt in [U8, I8, U16, I16, U32, I32, U64, I64, F16, F32, F64] {
1325 let parent = TensorDyn::new(&[2, 4], dt, Some(TensorMemory::Mem), None).unwrap();
1326 let view = parent.batch(1).unwrap();
1327 assert_eq!(view.dtype(), dt, "batch must preserve dtype {dt:?}");
1328 assert_eq!(view.shape(), &[4], "{dt:?}");
1329 }
1330 }
1331
1332 #[test]
1333 fn map_accepts_zero_offset_tensor() {
1334 let mut t =
1335 Tensor::<u8>::image(100, 100, PixelFormat::Rgba, Some(TensorMemory::Mem)).unwrap();
1336 t.set_plane_offset(0);
1337 assert!(t.map().is_ok());
1339 }
1340
1341 #[test]
1342 fn dyn_configure_image_nv12() {
1343 let mut t = TensorDyn::image(640, 480, PixelFormat::Rgb, DType::U8, None).unwrap();
1344 t.configure_image(320, 240, PixelFormat::Nv12).unwrap();
1345 assert_eq!(t.format(), Some(PixelFormat::Nv12));
1346 assert_eq!((t.width(), t.height()), (Some(320), Some(240)));
1347 }
1348
1349 #[test]
1350 fn tensordyn_colorimetry_roundtrip() {
1351 use crate::{ColorEncoding, Colorimetry, DType, PixelFormat};
1352 let mut t = TensorDyn::image(1280, 720, PixelFormat::Nv12, DType::U8, None).unwrap();
1353 assert_eq!(t.colorimetry(), None);
1354 let c = Colorimetry::default().with_encoding(ColorEncoding::Bt709);
1355 t.set_colorimetry(Some(c));
1356 assert_eq!(t.colorimetry(), Some(c));
1357 }
1358
1359 #[test]
1360 fn from_planes_propagates_plane_offset() {
1361 let mut luma =
1362 Tensor::<u8>::new(&[480, 640], Some(TensorMemory::Mem), Some("luma")).unwrap();
1363 luma.set_plane_offset(4096);
1364 let chroma =
1365 Tensor::<u8>::new(&[240, 640], Some(TensorMemory::Mem), Some("chroma")).unwrap();
1366 let combined = Tensor::<u8>::from_planes(luma, chroma, PixelFormat::Nv12).unwrap();
1367 assert_eq!(combined.plane_offset(), Some(4096));
1368 }
1369
1370 #[test]
1371 fn cuda_passthrough_none_for_mem_tensor() {
1372 let t: TensorDyn = Tensor::<f32>::new(&[10], Some(TensorMemory::Mem), None)
1375 .unwrap()
1376 .into();
1377 assert!(t.cuda().is_none());
1378 assert!(t.cuda_map().is_none());
1379 }
1380}