mnn/tensor/
tensor_ref.rs

1use super::{Tensor, *};
2
3/// This is modeled as a reference to a tensor, you should never be able to dereference this.
4/// At any point you should only have &TensorRef or &mut TensorRef, never an owned TensorRef.
5#[repr(transparent)]
6pub struct TensorRef<H, M>
7where
8    H: HalideType,
9    M: TensorMachine,
10{
11    __marker: PhantomData<(M, mnn_sys::Tensor, [H])>,
12}
13
14impl<H, M> core::fmt::Debug for TensorRef<H, M>
15where
16    H: HalideType,
17    M: TensorMachine,
18{
19    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
20        let dynamic = self.is_dynamic_unsized();
21        // The element count and byte size are meaningless (and computed from negative
22        // extents) until a dynamic tensor has been resized.
23        let elements = (!dynamic).then(|| self.element_size());
24        let size = (!dynamic).then(|| self.size());
25        f.debug_struct("TensorRef")
26            .field("shape", &self.shape())
27            .field("device_id", &self.device_id())
28            .field("dimensions", &self.get_dimension_type())
29            .field("elements", &elements)
30            .field("dynamic", &dynamic)
31            .field("size", &size)
32            .finish()
33    }
34}
35
36impl<H: HalideType, M: TensorMachine> TensorRef<H, M> {
37    /// Get a raw pointer to the underlying MNN tensor
38    pub(crate) fn as_ptr(&self) -> *mut mnn_sys::Tensor {
39        self as *const Self as *mut mnn_sys::Tensor
40    }
41
42    /// Copies the data from a host tensor to the self.as_ptr()
43    ///
44    /// The tensors must have the same logical shape; MNN converts between layouts
45    /// (e.g. NCHW <-> NC4HW4) as part of the copy.
46    pub fn copy_from_host_tensor(&mut self, tensor: &TensorRef<H, Host>) -> Result<()> {
47        AnyTensorRef::copy_from_host_tensor(self, tensor)
48    }
49
50    /// Copies the data from the self.as_ptr() to a host tensor
51    ///
52    /// The tensors must have the same logical shape; MNN converts between layouts
53    /// (e.g. NCHW <-> NC4HW4) as part of the copy.
54    pub fn copy_to_host_tensor(&self, tensor: &mut TensorRef<H, Host>) -> Result<()> {
55        AnyTensorRef::copy_to_host_tensor(self, tensor)
56    }
57}
58
59impl<H, M> core::ops::Deref for TensorRef<H, M>
60where
61    H: HalideType,
62    M: TensorMachine,
63{
64    type Target = AnyTensorRef;
65    fn deref(&self) -> &Self::Target {
66        unsafe { &*(self.as_ptr() as *const AnyTensorRef) }
67    }
68}
69
70impl<H, M> core::ops::DerefMut for TensorRef<H, M>
71where
72    H: HalideType,
73    M: TensorMachine,
74{
75    fn deref_mut(&mut self) -> &mut Self::Target {
76        unsafe { &mut *(self.as_ptr().cast()) }
77    }
78}
79
80impl<T> TensorRef<T, Device>
81where
82    T: HalideType,
83{
84    /// Try to wait for the device tensor to finish processing
85    pub fn wait(&self, map_type: MapType, finish: bool) {
86        unsafe {
87            mnn_sys::Tensor_wait(self.as_ptr(), map_type, finish as i32);
88        }
89    }
90
91    /// Create a host tensor from the device tensor with same dimensions and data type and
92    /// optionally copy the data from the device tensor
93    ///
94    /// The host tensor uses the unpacked layout (NCHW / NHWC) even when the device tensor is
95    /// stored as NC4HW4, so `host()` yields the elements in logical order.
96    pub fn create_host_tensor_from_device(&self, copy_data: bool) -> super::Tensor<Owned<T>, Host> {
97        let shape = self.shape();
98        assert!(
99            !self.is_dynamic_unsized(),
100            "Cannot create a host tensor from a device tensor with a dynamic shape {shape:?}; resize the session first"
101        );
102        let dm_type = self.get_dimension_type();
103        let mut out = super::Tensor::new(shape, dm_type);
104
105        if copy_data {
106            self.copy_to_host_tensor(&mut out)
107                .expect("Failed to copy data from device tensor");
108        }
109        out
110    }
111}
112
113/// Construct a tensor reference from a raw pointer to an MNN tensor
114///
115/// # Safety
116/// The caller must ensure that the provided pointer is valid and points to a properly initialized MNN
117#[inline]
118pub unsafe fn from_raw_parts<'a, H, M>(tensor: *const mnn_sys::Tensor) -> &'a TensorRef<H, M>
119where
120    H: HalideType,
121    M: TensorMachine,
122{
123    unsafe { &*tensor.cast::<TensorRef<H, M>>() }
124}
125
126/// Construct a mutable tensor reference from a raw pointer to an MNN tensor
127///
128/// # Safety
129/// The caller must ensure that the provided pointer is valid and points to a properly initialized
130#[inline]
131pub unsafe fn from_raw_parts_mut<'a, H, M>(tensor: *mut mnn_sys::Tensor) -> &'a mut TensorRef<H, M>
132where
133    H: HalideType,
134    M: TensorMachine,
135{
136    unsafe { &mut *tensor.cast::<TensorRef<H, M>>() }
137}
138
139impl<S, M, H> core::ops::Deref for Tensor<S, M, H>
140where
141    S: TensorType<H = H>,
142    H: HalideType,
143    M: TensorMachine,
144{
145    type Target = TensorRef<H, M>;
146
147    fn deref(&self) -> &Self::Target {
148        unsafe { from_raw_parts(self.tensor) }
149    }
150}
151
152/// Only tensors that own their data or hold a mutable borrow of it can be mutated through
153/// the [`TensorRef`]; a `Tensor<View<&H>, _>` built over a shared slice must not be able to
154/// write through that slice.
155impl<S, M, H> core::ops::DerefMut for Tensor<S, M, H>
156where
157    S: MutableTensorType<H = H>,
158    H: HalideType,
159    M: TensorMachine,
160{
161    fn deref_mut(&mut self) -> &mut Self::Target {
162        unsafe { from_raw_parts_mut(self.tensor) }
163    }
164}
165
166impl<T, M> TensorRef<T, M>
167where
168    T: HalideType,
169    M: TensorMachine,
170{
171    /// Fill the tensor with the specified value
172    ///
173    /// Fails if the tensor has a different element type, has not been resized yet (dynamic
174    /// shape) or, for device tensors, if MNN refuses the copy (e.g. no backend is attached).
175    pub fn fill(&mut self, value: T) -> Result<()>
176    where
177        T: Copy,
178    {
179        ensure!(
180            self.is_type_of::<T>(),
181            ErrorKind::HalideTypeMismatch {
182                got: std::any::type_name::<T>(),
183            }
184        );
185        ensure!(!self.is_dynamic_unsized(), ErrorKind::DynamicTensorError);
186        if M::host() {
187            let size = self.element_size();
188            let result: &mut [T] = unsafe {
189                let data = mnn_sys::Tensor_host_mut(self.as_ptr()).cast();
190                core::slice::from_raw_parts_mut(data, size)
191            };
192            result.fill(value);
193        } else {
194            let shape = self.shape();
195            let dm_type = self.get_dimension_type();
196            let mut host: TensorOwned<T, Host> = Tensor::new(shape, dm_type);
197            host.fill(value)?;
198            self.copy_from_host_tensor(&host)?;
199        }
200        Ok(())
201    }
202}
203
204impl<T> TensorRef<T, Host>
205where
206    T: HalideType,
207{
208    /// Try to map the device tensor to the host memory and get the slice
209    pub fn try_host(&self) -> Result<&[T]> {
210        let size = self.element_size();
211        ensure!(
212            self.is_type_of::<T>(),
213            ErrorKind::HalideTypeMismatch {
214                got: std::any::type_name::<T>(),
215            }
216        );
217        let result = unsafe {
218            let data = mnn_sys::Tensor_host(self.as_ptr()).cast();
219            core::slice::from_raw_parts(data, size)
220        };
221        Ok(result)
222    }
223
224    /// Try to map the device tensor to the host memory and get the mutable slice
225    pub fn try_host_mut(&mut self) -> Result<&mut [T]> {
226        let size = self.element_size();
227        ensure!(
228            self.is_type_of::<T>(),
229            ErrorKind::HalideTypeMismatch {
230                got: std::any::type_name::<T>(),
231            }
232        );
233
234        let result = unsafe {
235            let data: *mut T = mnn_sys::Tensor_host_mut(self.as_ptr()).cast();
236            debug_assert!(!data.is_null());
237            core::slice::from_raw_parts_mut(data, size)
238        };
239        Ok(result)
240    }
241
242    /// Get the host memory slice of the tensor
243    pub fn host(&self) -> &[T] {
244        self.try_host().expect("Failed to get tensor host")
245    }
246
247    /// Get the mutable host memory slice of the tensor
248    pub fn host_mut(&mut self) -> &mut [T] {
249        self.try_host_mut().expect("Failed to get tensor host_mut")
250    }
251}
252
253#[cfg(test)]
254mod tests {
255    use super::*;
256
257    #[test]
258    fn copy_from_host_tensor_reports_failure() {
259        // A plain host tensor has no backend, so MNN refuses the copy; the error must not be
260        // swallowed.
261        let src = TensorOwned::<f32, Host>::from_shape_data(
262            [1, 2, 3],
263            DimensionType::Caffe,
264            &[1., 2., 3., 4., 5., 6.],
265        );
266        let mut dst = TensorOwned::<f32, Host>::new([1, 2, 3], DimensionType::Caffe);
267        dst.copy_from_host_tensor(&src)
268            .expect_err("copy without a backend must fail");
269        src.copy_to_host_tensor(&mut dst)
270            .expect_err("copy without a backend must fail");
271    }
272
273    #[test]
274    fn fill_on_device_tensor_without_backend_reports_failure() {
275        let mut dev = TensorOwned::<f32, Device>::new([1, 2, 3], DimensionType::Caffe);
276        dev.fill(1.0).expect_err("fill without a backend must fail");
277    }
278
279    #[test]
280    fn fill_rejects_wrong_element_type() {
281        let t = TensorOwned::<f32, Host>::new([1, 2, 3], DimensionType::Caffe);
282        // Reinterpret the typed reference to request a fill with the wrong element type.
283        let wrong: &mut TensorRef<i32, Host> = unsafe { from_raw_parts_mut(t.as_ptr()) };
284        wrong
285            .fill(1)
286            .expect_err("filling an f32 tensor with i32 must fail");
287    }
288
289    #[test]
290    fn debug_of_dynamic_tensor_does_not_panic() {
291        // A device tensor is never allocated, so a dynamic (-1) extent is representable.
292        let t = TensorOwned::<f32, Device>::new([1, -1, 3], DimensionType::Caffe);
293        let s = format!("{:?}", &*t);
294        assert!(s.contains("dynamic: true"), "{s}");
295        assert!(s.contains("elements: None"), "{s}");
296    }
297}
mnn/tensor/tensor_ref.rs

mnn/tensor/
tensor_ref.rs
