mnn/tensor/
list.rs

1use crate::{AnyTensorRef, Device, TensorRef, prelude::*};
2use mnn_sys::HalideType;
3
4/// A list of tensors returned by interpreter output queries.
5///
6/// This list only hands out shared references to the tensors it describes. Use
7/// [`TensorListMut`] (returned by [`Interpreter::inputs`](crate::Interpreter::inputs))
8/// when you need to mutate the tensors.
9#[repr(transparent)]
10pub struct TensorList<'t> {
11    pub(crate) inner: *const mnn_sys::TensorInfoArray,
12    pub(crate) __marker: PhantomData<&'t ()>,
13}
14
15impl core::fmt::Debug for TensorList<'_> {
16    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> core::fmt::Result {
17        f.debug_list().entries(self.iter()).finish()
18    }
19}
20
21impl Drop for TensorList<'_> {
22    fn drop(&mut self) {
23        unsafe { mnn_sys::destroyTensorInfoArray(self.inner.cast_mut()) }
24    }
25}
26
27impl<'t> TensorList<'t> {
28    pub(crate) fn from_ptr(inner: *const mnn_sys::TensorInfoArray) -> Self {
29        Self {
30            inner,
31            __marker: PhantomData,
32        }
33    }
34
35    /// Returns the size of the tensor list
36    pub fn size(&self) -> usize {
37        unsafe { (*self.inner).size }
38    }
39
40    /// Get the tensor at the given index
41    pub fn get(&self, index: usize) -> Option<TensorInfo<'t, '_>> {
42        if index >= self.size() {
43            None
44        } else {
45            let gtinfo = unsafe { mnn_sys::getTensorInfoArray(self.inner, index) };
46            if !gtinfo.is_null() {
47                Some(TensorInfo {
48                    tensor_info: gtinfo,
49                    __marker: PhantomData,
50                })
51            } else {
52                None
53            }
54        }
55    }
56
57    /// Get an iterator over the tensor list
58    pub fn iter<'tl>(&'tl self) -> TensorListIter<'t, 'tl> {
59        TensorListIter {
60            tensor_list: self,
61            idx: 0,
62        }
63    }
64}
65
66impl<'t, 'tl> IntoIterator for &'tl TensorList<'t> {
67    type Item = TensorInfo<'t, 'tl>;
68    type IntoIter = TensorListIter<'t, 'tl>;
69
70    fn into_iter(self) -> Self::IntoIter {
71        self.iter()
72    }
73}
74
75/// A list of tensors returned by interpreter input queries.
76///
77/// This list is created from a `&mut Session` and therefore has exclusive access to the
78/// tensors it describes, which is what allows it to hand out `&mut TensorRef`s.
79/// It dereferences to [`TensorList`] for all the read-only operations.
80#[repr(transparent)]
81pub struct TensorListMut<'t> {
82    inner: TensorList<'t>,
83}
84
85impl core::fmt::Debug for TensorListMut<'_> {
86    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> core::fmt::Result {
87        self.inner.fmt(f)
88    }
89}
90
91impl<'t> core::ops::Deref for TensorListMut<'t> {
92    type Target = TensorList<'t>;
93
94    fn deref(&self) -> &Self::Target {
95        &self.inner
96    }
97}
98
99impl<'t> TensorListMut<'t> {
100    pub(crate) fn from_ptr(inner: *const mnn_sys::TensorInfoArray) -> Self {
101        Self {
102            inner: TensorList::from_ptr(inner),
103        }
104    }
105
106    /// Get a mutable reference to the tensor at the given index
107    pub fn get_mut(&mut self, index: usize) -> Option<TensorInfoMut<'t, '_>> {
108        if index >= self.inner.size() {
109            None
110        } else {
111            let gtinfo = unsafe { mnn_sys::getTensorInfoArray(self.inner.inner, index) };
112            if !gtinfo.is_null() {
113                Some(TensorInfoMut {
114                    tensor_info: gtinfo,
115                    __marker: PhantomData,
116                })
117            } else {
118                None
119            }
120        }
121    }
122
123    /// Get a mutable iterator over the tensor list
124    pub fn iter_mut<'tl>(&'tl mut self) -> TensorListIterMut<'t, 'tl> {
125        TensorListIterMut {
126            tensor_list: self,
127            idx: 0,
128        }
129    }
130}
131
132impl<'t, 'tl> IntoIterator for &'tl TensorListMut<'t> {
133    type Item = TensorInfo<'t, 'tl>;
134    type IntoIter = TensorListIter<'t, 'tl>;
135
136    fn into_iter(self) -> Self::IntoIter {
137        self.inner.iter()
138    }
139}
140
141impl<'t, 'tl> IntoIterator for &'tl mut TensorListMut<'t> {
142    type Item = TensorInfoMut<'t, 'tl>;
143    type IntoIter = TensorListIterMut<'t, 'tl>;
144
145    fn into_iter(self) -> Self::IntoIter {
146        self.iter_mut()
147    }
148}
149
150/// Iterator over tensor list
151pub struct TensorListIter<'t, 'tl> {
152    tensor_list: &'tl TensorList<'t>,
153    idx: usize,
154}
155
156impl<'t, 'tl> Iterator for TensorListIter<'t, 'tl> {
157    type Item = TensorInfo<'t, 'tl>;
158    fn next(&mut self) -> Option<Self::Item> {
159        let idx = self.idx;
160        self.idx += 1;
161        self.tensor_list.get(idx)
162    }
163}
164
165/// Mutable iterator over tensor list
166pub struct TensorListIterMut<'t, 'tl> {
167    tensor_list: &'tl mut TensorListMut<'t>,
168    idx: usize,
169}
170
171impl<'t, 'tl> Iterator for TensorListIterMut<'t, 'tl> {
172    type Item = TensorInfoMut<'t, 'tl>;
173    fn next(&mut self) -> Option<Self::Item> {
174        let idx = self.idx;
175        self.idx += 1;
176        let out = self.tensor_list.get_mut(idx);
177        // SAFETY: Every index is visited at most once, so each `TensorInfoMut` handed out by
178        // this iterator refers to a distinct entry of the list. The iterator holds the list
179        // exclusively for `'tl`, so extending the borrow from the `&mut self` of `get_mut`
180        // to `'tl` cannot create overlapping mutable access. This is the same reasoning the
181        // standard library uses for `slice::IterMut`.
182        unsafe {
183            core::mem::transmute::<Option<TensorInfoMut<'t, '_>>, Option<TensorInfoMut<'t, 'tl>>>(
184                out,
185            )
186        }
187    }
188}
189
190/// Information about a tensor in a tensor list
191#[repr(transparent)]
192pub struct TensorInfo<'t, 'tl> {
193    pub(crate) tensor_info: *mut mnn_sys::TensorInfo,
194    pub(crate) __marker: PhantomData<&'tl TensorList<'t>>,
195}
196
197impl core::fmt::Debug for TensorInfo<'_, '_> {
198    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
199        let tensor = self.raw_tensor();
200        let shape = tensor.shape();
201        f.debug_struct("TensorInfo")
202            .field("name", &self.name())
203            .field("tensor", &shape)
204            .finish()
205    }
206}
207
208/// Mutable information about a tensor in a tensor list
209///
210/// Each `TensorInfoMut` represents exclusive access to one entry of a [`TensorListMut`] for
211/// the lifetime `'tl`.
212#[repr(transparent)]
213pub struct TensorInfoMut<'t, 'tl> {
214    pub(crate) tensor_info: *mut mnn_sys::TensorInfo,
215    pub(crate) __marker: PhantomData<&'tl mut TensorList<'t>>,
216}
217
218impl core::fmt::Debug for TensorInfoMut<'_, '_> {
219    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
220        let tensor = self.raw_tensor();
221        let shape = tensor.shape();
222        f.debug_struct("TensorInfoMut")
223            .field("name", &self.name())
224            .field("tensor", &shape)
225            .finish()
226    }
227}
228
229/// Type and shape checks shared by all the typed accessors
230fn check_tensor<H: HalideType>(tensor: &AnyTensorRef) -> Result<()> {
231    ensure!(!tensor.is_dynamic_unsized(), ErrorKind::DynamicTensorError);
232    ensure!(
233        tensor.is_type_of::<H>(),
234        ErrorKind::HalideTypeMismatch {
235            got: std::any::type_name::<H>(),
236        }
237    );
238    Ok(())
239}
240
241impl<'t, 'tl> TensorInfoMut<'t, 'tl> {
242    fn tensor_ptr(&self) -> *mut mnn_sys::Tensor {
243        debug_assert!(!self.tensor_info.is_null());
244        let tensor = unsafe { (*self.tensor_info).tensor };
245        debug_assert!(!tensor.is_null());
246        tensor.cast()
247    }
248
249    /// Get the name of the tensor
250    pub fn name(&self) -> &'tl str {
251        debug_assert!(!self.tensor_info.is_null());
252        unsafe { (*self.tensor_info).name.to_cstr() }
253            .to_str()
254            .expect("Tensor name is not utf-8")
255    }
256
257    /// Get the tensor with the specified type
258    pub fn tensor<H: HalideType>(&self) -> Result<&TensorRef<H, Device>> {
259        check_tensor::<H>(self.raw_tensor())?;
260        Ok(unsafe { crate::tensor::from_raw_parts(self.tensor_ptr()) })
261    }
262
263    /// Get a mutable reference to the tensor with the specified type
264    ///
265    /// The returned reference borrows `self` mutably, use [`Self::into_tensor_mut`] if you
266    /// need a reference that lives as long as the list borrow.
267    pub fn tensor_mut<H: HalideType>(&mut self) -> Result<&mut TensorRef<H, Device>> {
268        check_tensor::<H>(self.raw_tensor())?;
269        Ok(unsafe { crate::tensor::from_raw_parts_mut(self.tensor_ptr()) })
270    }
271
272    /// Consume this entry and get a mutable reference to the tensor that lives as long as
273    /// the borrow of the [`TensorListMut`] this entry came from.
274    pub fn into_tensor_mut<H: HalideType>(self) -> Result<&'tl mut TensorRef<H, Device>> {
275        check_tensor::<H>(self.raw_tensor())?;
276        // SAFETY: `self` represents exclusive access to this entry for `'tl` and is consumed
277        // here, so no other reference to the tensor can be created through it.
278        Ok(unsafe { crate::tensor::from_raw_parts_mut(self.tensor_ptr()) })
279    }
280
281    /// # Safety
282    /// The shape is not checked so it's marked unsafe since futher calls to interpreter might be **unsafe** with this
283    pub unsafe fn tensor_unresized<H: HalideType>(&mut self) -> Result<&mut TensorRef<H, Device>> {
284        let tensor = self.raw_tensor();
285        ensure!(
286            tensor.is_type_of::<H>(),
287            ErrorKind::HalideTypeMismatch {
288                got: std::any::type_name::<H>(),
289            }
290        );
291        Ok(unsafe { crate::tensor::from_raw_parts_mut(self.tensor_ptr()) })
292    }
293
294    /// This function return's the raw tensor without any sort of type-checking or shape-checking
295    pub fn raw_tensor(&self) -> &AnyTensorRef {
296        unsafe { AnyTensorRef::from_ptr(self.tensor_ptr()) }
297    }
298
299    /// This function return's the raw tensor mutably without any sort of type-checking or shape-checking
300    pub fn raw_tensor_mut(&mut self) -> &mut AnyTensorRef {
301        unsafe { AnyTensorRef::from_mut_ptr(self.tensor_ptr()) }
302    }
303}
304
305impl<'t, 'tl> TensorInfo<'t, 'tl> {
306    fn tensor_ptr(&self) -> *mut mnn_sys::Tensor {
307        debug_assert!(!self.tensor_info.is_null());
308        let tensor = unsafe { (*self.tensor_info).tensor };
309        debug_assert!(!tensor.is_null());
310        tensor.cast()
311    }
312
313    /// Get the name of the tensor
314    pub fn name(&self) -> &'tl str {
315        debug_assert!(!self.tensor_info.is_null());
316        unsafe { (*self.tensor_info).name.to_cstr() }
317            .to_str()
318            .expect("Tensor name is not utf-8")
319    }
320
321    /// Get the tensor with the specified type
322    pub fn tensor<H: HalideType>(&self) -> Result<&'tl TensorRef<H, Device>> {
323        check_tensor::<H>(self.raw_tensor())?;
324        Ok(unsafe { crate::tensor::from_raw_parts(self.tensor_ptr()) })
325    }
326
327    /// This function return's the raw tensor without any sort of type-checking or shape-checking
328    pub fn raw_tensor(&self) -> &'tl AnyTensorRef {
329        unsafe { AnyTensorRef::from_ptr(self.tensor_ptr()) }
330    }
331}
mnn/tensor/list.rs

mnn/tensor/
list.rs
