mnn/
interpreter.rs

1//! The interpreter module provides the `Interpreter` struct which is used to load and run models.
2use crate::{
3    AnyTensorRef, AsTensorShape, Device, ScheduleConfig, TensorList, TensorListMut, TensorRef,
4    prelude::*,
5};
6use mnn_sys::HalideType;
7use std::{ffi::CStr, path::Path, sync::Arc};
8
9pub(crate) type TensorCallbackT = Box<dyn Fn(&[&AnyTensorRef], OperatorInfo) -> bool>;
10
11#[repr(transparent)]
12pub(crate) struct TensorCallback {
13    inner: Arc<TensorCallbackT>,
14}
15
16impl Default for TensorCallback {
17    fn default() -> Self {
18        Self {
19            inner: Arc::new(Box::new(|_, _| true)),
20        }
21    }
22}
23
24impl TensorCallback {
25    pub(crate) fn from_ptr(f: *mut libc::c_void) -> Self {
26        debug_assert!(!f.is_null());
27        unsafe {
28            Self {
29                inner: Arc::from_raw(f.cast()),
30            }
31        }
32    }
33
34    pub(crate) fn into_ptr(self) -> *mut libc::c_void {
35        Arc::into_raw(self.inner) as *mut libc::c_void
36    }
37
38    #[cfg(test)]
39    pub(crate) fn identity() -> impl Fn(&[&AnyTensorRef], OperatorInfo) -> bool {
40        |_, _| true
41    }
42}
43
44impl<F> From<F> for TensorCallback
45where
46    F: Fn(&[&AnyTensorRef], OperatorInfo) -> bool + 'static,
47{
48    fn from(f: F) -> Self {
49        Self {
50            inner: Arc::new(Box::new(f)),
51        }
52    }
53}
54
55impl<T> From<Option<T>> for TensorCallback
56where
57    T: Fn(&[&AnyTensorRef], OperatorInfo) -> bool + 'static,
58{
59    fn from(f: Option<T>) -> Self {
60        match f {
61            Some(f) => Self {
62                inner: Arc::new(Box::new(f)),
63            },
64            None => Self::default(),
65        }
66    }
67}
68
69impl core::ops::Deref for TensorCallback {
70    type Target = TensorCallbackT;
71
72    fn deref(&self) -> &Self::Target {
73        &self.inner
74    }
75}
76
77/// The session mode to be used
78/// The items are mostly untested and are only documented 1:1 to the C++ codebase
79/// The only two items tested are
80/// - `Debug`
81/// - `Release`
82#[derive(Debug, Copy, Clone)]
83#[cfg_attr(windows, repr(i32))]
84#[cfg_attr(unix, repr(u32))]
85pub enum SessionMode {
86    #[doc = "About CallBack, Default Session_Debug*/\n/** runSessionWithCallBack is allowed and can get internal op info"]
87    Debug = mnn_sys::SessionMode::Session_Debug,
88    #[doc = "runSessionWithCallBack is not valid and can't get any info of op in\nsession"]
89    Release = mnn_sys::SessionMode::Session_Release,
90    #[doc = "About input tensor, Default Session_Input_Inside*/\n/** The input tensor is alloced by session, input data after session resized"]
91    InputInside = mnn_sys::SessionMode::Session_Input_Inside,
92    #[doc = "The input tensor is alloced by user, set input data before session\nresize"]
93    InputUser = mnn_sys::SessionMode::Session_Input_User,
94    #[doc = "The output tensor depends on session, and can't be separate used"]
95    OutputInside = mnn_sys::SessionMode::Session_Output_Inside,
96    #[doc = "The output tensor can be separated from session"]
97    OutputUser = mnn_sys::SessionMode::Session_Output_User,
98    #[doc = "Try Resize Session when create Session or not, default direct:"]
99    ResizeDirect = mnn_sys::SessionMode::Session_Resize_Direct,
100    #[doc = "Try Resize Session when create Session or not, default direct:"]
101    ResizeDefer = mnn_sys::SessionMode::Session_Resize_Defer,
102    #[doc = "Determine the Execution's forward type is determine by user or auto\ndetermine"]
103    BackendFix = mnn_sys::SessionMode::Session_Backend_Fix,
104    #[doc = "Determine the Execution's forward type is determine by user or auto\ndetermine"]
105    BackendAuto = mnn_sys::SessionMode::Session_Backend_Auto,
106    #[doc = "Determine static memory whether recyle in resizeSession or just cache the\nmemory"]
107    MemoryCollect = mnn_sys::SessionMode::Session_Memory_Collect,
108    #[doc = "Determine static memory whether recyle in resizeSession or just cache the\nmemory"]
109    MemoryCache = mnn_sys::SessionMode::Session_Memory_Cache,
110    #[doc = "Determine whether use codegen function"]
111    CodegenDisable = mnn_sys::SessionMode::Session_Codegen_Disable,
112    #[doc = "Determine whether use codegen function"]
113    CodegenEnable = mnn_sys::SessionMode::Session_Codegen_Enable,
114    #[doc = "Dynamic Reisze Optimization"]
115    ResizeCheck = mnn_sys::SessionMode::Session_Resize_Check,
116    #[doc = "Dynamic Reisze Optimization"]
117    ResizeFix = mnn_sys::SessionMode::Session_Resize_Fix,
118}
119
120impl SessionMode {
121    fn to_mnn_sys(self) -> mnn_sys::SessionMode::Type {
122        self as mnn_sys::SessionMode::Type
123    }
124}
125
126/// The Interpreter holds the model and manages sessions.
127///
128/// It's internals use a shared mutex so it's thread-safe.
129/// net data holder. multiple sessions could share same net.
130#[repr(transparent)]
131#[derive(Debug)]
132pub struct Interpreter {
133    pub(crate) inner: *mut mnn_sys::Interpreter,
134    pub(crate) __marker: PhantomData<()>,
135}
136
137unsafe impl Send for Interpreter {}
138
139impl Drop for Interpreter {
140    fn drop(&mut self) {
141        unsafe { mnn_sys::Interpreter_destroy(self.inner) }
142    }
143}
144
145impl Interpreter {
146    /// Create an net/interpreter from a file.
147    ///
148    /// `path`: the file path of the model
149    ///
150    /// return: the created net/interpreter
151    pub fn from_file(path: impl AsRef<Path>) -> Result<Self> {
152        let path = path.as_ref();
153        ensure!(path.exists(), ErrorKind::IOError; path.to_string_lossy().to_string(), "File not found");
154        let path = path.to_str().ok_or_else(|| error!(ErrorKind::AsciiError))?;
155        let c_path = std::ffi::CString::new(path).change_context(ErrorKind::AsciiError)?;
156        let interpreter = unsafe { mnn_sys::Interpreter_createFromFile(c_path.as_ptr()) };
157        ensure!(!interpreter.is_null(), ErrorKind::InterpreterError; "Failed to create interpreter", "Interpreter_createFromFile returned null");
158        Ok(Self {
159            inner: interpreter,
160            __marker: PhantomData,
161        })
162    }
163
164    /// Create an net/interpreter from a buffer.
165    ///
166    /// `bytes`: the buffer of the model
167    ///
168    /// return: the created net/interpreter
169    pub fn from_bytes(bytes: impl AsRef<[u8]>) -> Result<Self> {
170        let bytes = bytes.as_ref();
171        let size = bytes.len();
172        let interpreter =
173            unsafe { mnn_sys::Interpreter_createFromBuffer(bytes.as_ptr().cast(), size) };
174        ensure!(!interpreter.is_null(), ErrorKind::InterpreterError; "Failed to create interpreter", "Interpreter_createFromBuffer returned null");
175        Ok(Self {
176            inner: interpreter,
177            __marker: PhantomData,
178        })
179    }
180
181    /// Set session mode
182    ///
183    /// `mode`: the session mode
184    ///
185    /// **Warning:**
186    /// It should be called before create session!
187    pub fn set_session_mode(&self, mode: SessionMode) {
188        unsafe { mnn_sys::Interpreter_setSessionMode(self.inner, mode.to_mnn_sys()) }
189    }
190
191    ///call this function to get tensors ready.
192    ///
193    ///output tensor buffer (host or deviceId) should be retrieved after resize of any input tensor.
194    ///
195    ///`session`: the session to be prepared
196    pub fn resize_session(&self, session: &mut crate::Session) {
197        unsafe { mnn_sys::Interpreter_resizeSession(self.inner, session.inner) }
198    }
199
200    /// Resize session and reallocate the buffer.
201    ///
202    /// `session`: the session to be prepared.
203    ///
204    /// # Note
205    /// NeedRelloc is default to 1, 1 means need realloc!
206    pub fn resize_session_reallocate(&self, session: &mut crate::Session) {
207        unsafe { mnn_sys::Interpreter_resizeSessionWithFlag(self.inner, session.inner, 1i32) }
208    }
209
210    /// Resize the tensor using the given shape
211    pub fn resize_tensor<H: HalideType, M: TensorMachine>(
212        &self,
213        tensor: &mut TensorRef<H, M>,
214        dims: impl AsTensorShape,
215    ) {
216        let dims = dims.as_tensor_shape();
217        let dims_len = dims.size;
218        unsafe {
219            mnn_sys::Interpreter_resizeTensor(
220                self.inner,
221                tensor.as_ptr(),
222                dims.shape.as_ptr(),
223                dims_len,
224            )
225        }
226    }
227
228    /// Resize the tenror by name using the given shape
229    pub fn resize_tensor_by_name(
230        &self,
231        session: &mut crate::Session,
232        name: impl AsRef<str>,
233        dims: impl AsTensorShape,
234    ) -> Result<()> {
235        let tensor = self.raw_input(session, name)?;
236        let dims = dims.as_tensor_shape();
237        let dims_len = dims.size;
238        unsafe {
239            mnn_sys::Interpreter_resizeTensor(
240                self.inner,
241                tensor.as_ptr(),
242                dims.shape.as_ptr(),
243                dims_len,
244            )
245        }
246        Ok(())
247    }
248
249    /// Resize tensor by
250    /// - N -> batch
251    /// - C -> channel
252    /// - H -> height
253    /// - W -> width
254    pub fn resize_tensor_by_nchw<H: HalideType, M: TensorMachine>(
255        &self,
256        tensor: &mut TensorRef<H, M>,
257        batch: u16,
258        channel: u16,
259        height: u16,
260        width: u16,
261    ) {
262        unsafe {
263            mnn_sys::Interpreter_resizeTensorByNCHW(
264                self.inner,
265                tensor.as_ptr(),
266                batch.into(),
267                channel.into(),
268                height.into(),
269                width.into(),
270            )
271        }
272    }
273
274    /// Create a session with session config. Session will be managed in net/interpreter.
275    ///
276    /// `schedule` : the config of the session
277    ///
278    /// return: the created session
279    pub fn create_session(
280        &self,
281        schedule: crate::ScheduleConfig,
282    ) -> Result<crate::session::Session<'_>> {
283        profile!("Creating session"; {
284            let session = unsafe { mnn_sys::Interpreter_createSession(self.inner, schedule.inner) };
285            assert!(!session.is_null());
286            Ok(crate::session::Session {
287                inner: session,
288                net: self.inner,
289                __session_internals: crate::SessionInternals::Single(schedule),
290                __marker: PhantomData,
291            })
292        })
293    }
294
295    /// Release the model file buffer
296    /// # Safety
297    /// This function is marked unsafe since it's not clear what the safety guarantees are right
298    /// now. With a simple test it caused a segfault so it's marked unsafe
299    pub unsafe fn release_model(&mut self) {
300        unsafe { mnn_sys::Interpreter_releaseModel(self.inner) }
301    }
302
303    /// Create multi-path session with schedule configs and user-specified runtime. created session will be managed in net/interpreter.
304    ///
305    /// `schedule` : the config of the session
306    ///
307    /// return: the created session
308    pub fn create_multipath_session(
309        &self,
310        schedule: impl IntoIterator<Item = ScheduleConfig>,
311    ) -> Result<crate::session::Session<'_>> {
312        profile!("Creating multipath session"; {
313            let schedules: crate::ScheduleConfigs = schedule.into_iter().collect();
314            let sc: &[_] = schedules.inner.as_ref();
315            let session = unsafe { mnn_sys::Interpreter_createMultiPathSession(self.inner, sc.as_ptr(), sc.len()) };
316            assert!(!session.is_null());
317            Ok(crate::session::Session {
318                inner: session,
319                net: self.inner,
320                __session_internals: crate::SessionInternals::MultiSession(schedules),
321                __marker: PhantomData,
322            })
323        })
324    }
325
326    /// Print all input and output tensors info.
327    pub fn model_print_io(path: impl AsRef<Path>) -> Result<()> {
328        let path = path.as_ref();
329        crate::ensure!(path.exists(), ErrorKind::IOError);
330        let path = path.to_str().ok_or_else(|| error!(ErrorKind::AsciiError))?;
331        let c_path = std::ffi::CString::new(path).change_context(ErrorKind::AsciiError)?;
332        unsafe { mnn_sys::modelPrintIO(c_path.as_ptr()) }
333        Ok(())
334    }
335
336    /// Get the input tensor of the session.
337    ///
338    /// `session`: the session to get input tensor
339    ///
340    /// return: List of input tensors
341    pub fn inputs<'i>(&self, session: &'i mut crate::Session) -> TensorListMut<'i> {
342        let inputs = unsafe { mnn_sys::Interpreter_getSessionInputAll(self.inner, session.inner) };
343        TensorListMut::from_ptr(inputs)
344    }
345
346    /// Get the input tensor of the session by name.
347    ///
348    /// `session`: the session to get input tensor from
349    ///
350    /// `name`: the name of the input tensor
351    ///
352    /// return: the input tensor
353    pub fn input<'s, H: HalideType>(
354        &self,
355        session: &'s mut crate::Session,
356        name: impl AsRef<str>,
357    ) -> Result<&'s mut TensorRef<H, Device>> {
358        let name = name.as_ref();
359        let c_name = std::ffi::CString::new(name).change_context(ErrorKind::AsciiError)?;
360        let input = unsafe {
361            mnn_sys::Interpreter_getSessionInput(self.inner, session.inner, c_name.as_ptr())
362        };
363        ensure!(!input.is_null(), ErrorKind::TensorError; format!("Input tensor \"{name}\" not found"));
364        let tensor = unsafe { crate::tensor::from_raw_parts_mut(input) };
365        let shape = tensor.shape();
366        ensure!(!shape.as_ref().contains(&-1), ErrorKind::DynamicTensorError);
367        ensure!(
368            tensor.is_type_of::<H>(),
369            ErrorKind::HalideTypeMismatch {
370                got: std::any::type_name::<H>(),
371            };
372            format!("Input tensor \"{name}\" is not of type {}", std::any::type_name::<H>())
373        );
374        Ok(tensor)
375    }
376
377    /// Get the raw input tensor of a session by name
378    pub fn raw_input<'s>(
379        &self,
380        session: &'s mut crate::Session,
381        name: impl AsRef<str>,
382    ) -> Result<&'s AnyTensorRef> {
383        let name = name.as_ref();
384        let c_name = std::ffi::CString::new(name).change_context(ErrorKind::AsciiError)?;
385        let input = unsafe {
386            mnn_sys::Interpreter_getSessionInput(self.inner, session.inner, c_name.as_ptr())
387        };
388        ensure!(!input.is_null(), ErrorKind::TensorError; format!("Input tensor \"{name}\" not found"));
389        // let out =
390        Ok(unsafe { AnyTensorRef::from_ptr::<'s>(input) })
391    }
392
393    /// # Safety
394    /// **Warning**  We Still don't know the safety guarantees of this function so it's marked unsafe
395    pub unsafe fn input_unresized<'s, H: HalideType>(
396        &self,
397        session: &'s mut crate::Session,
398        name: impl AsRef<str>,
399    ) -> Result<&'s mut TensorRef<H, Device>> {
400        let name = name.as_ref();
401        let c_name = std::ffi::CString::new(name).change_context(ErrorKind::AsciiError)?;
402        let input = unsafe {
403            mnn_sys::Interpreter_getSessionInput(self.inner, session.inner, c_name.as_ptr())
404        };
405        ensure!(!input.is_null(), ErrorKind::TensorError; format!("Input tensor \"{name}\" not found"));
406        let tensor = unsafe { crate::tensor::from_raw_parts_mut(input) };
407        ensure!(
408            tensor.is_type_of::<H>(),
409            ErrorKind::HalideTypeMismatch {
410                got: std::any::type_name::<H>(),
411            }
412        );
413        Ok(tensor)
414    }
415
416    /// # Safety
417    /// Very **unsafe** since it doesn't check the type of the tensor
418    /// as well as the shape of the tensor
419    ///
420    /// **Panics** if the name is not ascii
421    /// **Undefined Behavior** if the tensor is not of type `H`
422    pub unsafe fn input_unchecked<'s, H: HalideType>(
423        &self,
424        session: &'s mut crate::Session,
425        name: impl AsRef<str>,
426    ) -> &'s mut TensorRef<H, Device> {
427        let name = name.as_ref();
428        let c_name = std::ffi::CString::new(name).expect("Input tensor name is not ascii");
429        unsafe {
430            let input =
431                mnn_sys::Interpreter_getSessionInput(self.inner, session.inner, c_name.as_ptr());
432            crate::tensor::from_raw_parts_mut(input)
433        }
434    }
435
436    /// Get the output tensor of a session by name
437    ///
438    /// `session` : the session to get output tensor from
439    ///
440    /// `name` : the name of the output tensor
441    pub fn output<'s, H: HalideType>(
442        &self,
443        session: &'s crate::Session,
444        name: impl AsRef<str>,
445    ) -> Result<&'s TensorRef<H, Device>> {
446        let name = name.as_ref();
447        let c_name = std::ffi::CString::new(name).change_context(ErrorKind::AsciiError)?;
448        let output = unsafe {
449            mnn_sys::Interpreter_getSessionOutput(self.inner, session.inner, c_name.as_ptr())
450        };
451        ensure!(!output.is_null(), ErrorKind::IOError;format!("Output tensor \"{name}\" not found"));
452        let tensor = unsafe { crate::tensor::from_raw_parts(output) };
453        let shape = tensor.shape();
454        ensure!(!shape.as_ref().contains(&-1), ErrorKind::DynamicTensorError);
455        ensure!(
456            tensor.is_type_of::<H>(),
457            ErrorKind::HalideTypeMismatch {
458                got: std::any::type_name::<H>(),
459            }
460        );
461        Ok(tensor)
462    }
463
464    /// Get the raw output tensor of a session by name
465    pub fn raw_output<'s>(
466        &self,
467        session: &'s crate::Session,
468        name: impl AsRef<str>,
469    ) -> Result<&'s AnyTensorRef> {
470        let name = name.as_ref();
471        let c_name = std::ffi::CString::new(name).change_context(ErrorKind::AsciiError)?;
472        let output = unsafe {
473            mnn_sys::Interpreter_getSessionOutput(self.inner, session.inner, c_name.as_ptr())
474        };
475        ensure!(!output.is_null(), ErrorKind::IOError;format!("Output tensor \"{name}\" not found"));
476        Ok(unsafe { AnyTensorRef::from_ptr::<'s>(output) })
477    }
478
479    /// Run a session
480    pub fn run_session(&self, session: &crate::session::Session) -> Result<()> {
481        profile!("Running session"; {
482            let ret = unsafe { mnn_sys::Interpreter_runSession(self.inner, session.inner) };
483            ensure!(
484                ret == mnn_sys::ErrorCode::ERROR_CODE_NO_ERROR,
485                ErrorKind::InternalError(ret)
486            );
487            Ok(())
488        })
489    }
490
491    /// Run a session with a callback
492    ///
493    /// `session` : the session to run
494    ///
495    /// `before` : a callback before each op. return true to run the op; return false to skip the op.
496    ///
497    /// `after` : a callback after each op. return true to continue running; return false to interrupt the session.
498    ///
499    /// `sync` : synchronously wait for finish of execution or not.
500    pub fn run_session_with_callback(
501        &self,
502        session: &crate::session::Session,
503        before: impl Fn(&[&AnyTensorRef], OperatorInfo) -> bool + 'static,
504        end: impl Fn(&[&AnyTensorRef], OperatorInfo) -> bool + 'static,
505        sync: bool,
506    ) -> Result<()> {
507        let sync = sync as libc::c_int;
508        let before = TensorCallback::from(before).into_ptr();
509        let end = TensorCallback::from(end).into_ptr();
510        let ret = unsafe {
511            mnn_sys::Interpreter_runSessionWithCallBackInfo(
512                self.inner,
513                session.inner,
514                before,
515                end,
516                sync,
517            )
518        };
519        ensure!(
520            ret == mnn_sys::ErrorCode::ERROR_CODE_NO_ERROR,
521            ErrorKind::InternalError(ret)
522        );
523        Ok(())
524    }
525
526    /// Get all output tensors of a session
527    pub fn outputs<'o>(&self, session: &'o crate::session::Session) -> TensorList<'o> {
528        let outputs =
529            unsafe { mnn_sys::Interpreter_getSessionOutputAll(self.inner, session.inner) };
530        TensorList::from_ptr(outputs)
531    }
532
533    /// If the cache exist, try to load cache from file.
534    /// After createSession, try to save cache to file.
535    ///
536    /// `cache_file` : the file path to save or load cache.
537    ///
538    /// `key_size` : the size of key
539    ///
540    /// # Note
541    /// The API should be called before create session.
542    ///
543    /// Key Depercerate, keeping for future use!
544    pub fn set_cache_file(&self, path: impl AsRef<Path>, key_size: usize) -> Result<()> {
545        let path = path.as_ref();
546        let path = dunce::simplified(path);
547        let path = path.to_str().ok_or_else(|| error!(ErrorKind::AsciiError))?;
548        let c_path = std::ffi::CString::new(path).change_context(ErrorKind::AsciiError)?;
549        unsafe { mnn_sys::Interpreter_setCacheFile(self.inner, c_path.as_ptr(), key_size) }
550        Ok(())
551    }
552
553    /// Update cache file
554    pub fn update_cache_file(&self, session: &mut crate::session::Session) -> Result<()> {
555        MNNError::from_error_code(unsafe {
556            mnn_sys::Interpreter_updateCacheFile(self.inner, session.inner)
557        });
558        Ok(())
559    }
560
561    /// Wait for all output tensors to be ready after computation
562    pub fn wait(&self, session: &crate::session::Session) {
563        self.outputs(session).iter().for_each(|tinfo| {
564            AnyTensorRef::wait(tinfo.raw_tensor(), mnn_sys::MapType::MAP_TENSOR_READ, true);
565        });
566    }
567
568    /// Get memory usage of a session in MB
569    pub fn memory(&self, session: &crate::session::Session) -> Result<f32> {
570        let mut memory = 0f32;
571        let memory_ptr = &mut memory as *mut f32;
572        let ret = unsafe {
573            mnn_sys::Interpreter_getSessionInfo(
574                self.inner,
575                session.inner,
576                mnn_sys::cpp::MNN_Interpreter_SessionInfoCode_MEMORY as _,
577                memory_ptr.cast(),
578            )
579        };
580        ensure!(
581            ret == 1,
582            ErrorKind::InterpreterError;
583            "Failed to get memory usage"
584        );
585        Ok(memory)
586    }
587
588    /// Get float operation needed in session in M
589    pub fn flops(&self, session: &crate::Session) -> Result<f32> {
590        let mut flop = 0.0f32;
591        let flop_ptr = &mut flop as *mut f32;
592        let ret = unsafe {
593            mnn_sys::Interpreter_getSessionInfo(
594                self.inner,
595                session.inner,
596                mnn_sys::cpp::MNN_Interpreter_SessionInfoCode_FLOPS as _,
597                flop_ptr.cast::<libc::c_void>(),
598            )
599        };
600        ensure!(
601            ret == 1,
602            ErrorKind::InterpreterError;
603            "Failed to get flops"
604        );
605        Ok(flop)
606    }
607
608    /// Get the resize status
609    pub fn resize_status(&self, session: &crate::Session) -> Result<ResizeStatus> {
610        let mut resize_status = 0i32;
611        let ptr = &mut resize_status as *mut i32;
612        let ret = unsafe {
613            mnn_sys::Interpreter_getSessionInfo(
614                self.inner,
615                session.inner,
616                mnn_sys::cpp::MNN_Interpreter_SessionInfoCode_RESIZE_STATUS as _,
617                ptr.cast(),
618            )
619        };
620        ensure!(
621        ret == 1,
622            ErrorKind::InterpreterError;
623            "Failed to get resize status"
624        );
625        match resize_status {
626            0 => Ok(ResizeStatus::None),
627            1 => Ok(ResizeStatus::NeedMalloc),
628            2 => Ok(ResizeStatus::NeedResize),
629            _ => Err(error!(ErrorKind::InterpreterError)),
630        }
631    }
632}
633
634/// The status of the resize operation
635#[derive(Debug, Copy, Clone, PartialEq, Eq)]
636#[repr(C)]
637pub enum ResizeStatus {
638    /// No resize needed
639    None = 0,
640    /// Need to malloc memory
641    NeedMalloc = 1,
642    /// Need to resize memory
643    NeedResize = 2,
644}
645
646#[unsafe(no_mangle)]
647extern "C" fn rust_closure_callback_runner_op(
648    f: *mut libc::c_void,
649    tensors: *const *mut mnn_sys::Tensor,
650    tensor_count: usize,
651    op: *mut libc::c_void,
652) -> libc::c_int {
653    let tensors = unsafe { std::slice::from_raw_parts(tensors.cast(), tensor_count) };
654    let f: TensorCallback = TensorCallback::from_ptr(f);
655    let op = OperatorInfo {
656        inner: op.cast(),
657        __marker: PhantomData,
658    };
659    let ret = f(tensors, op) as libc::c_int;
660
661    core::mem::forget(f);
662    ret
663}
664
665/// A struct that holds information about an operator
666#[repr(transparent)]
667pub struct OperatorInfo<'op> {
668    pub(crate) inner: *mut libc::c_void,
669    pub(crate) __marker: PhantomData<&'op ()>,
670}
671
672impl core::fmt::Debug for OperatorInfo<'_> {
673    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
674        f.debug_struct("OperatorInfo")
675            .field("name", &self.name())
676            .field("type", &self.type_name())
677            .field("flops", &self.flops())
678            .finish()
679    }
680}
681
682impl OperatorInfo<'_> {
683    /// Get the name of the operator
684    pub fn name(&self) -> &CStr {
685        unsafe { CStr::from_ptr(mnn_sys::OperatorInfo_name(self.inner)) }
686    }
687
688    /// Get the type of the operator
689    pub fn type_name(&self) -> &CStr {
690        unsafe { CStr::from_ptr(mnn_sys::OperatorInfo_type(self.inner)) }
691    }
692
693    /// Get the number of flops of the operator
694    pub fn flops(&self) -> f32 {
695        unsafe { mnn_sys::OperatorInfo_flops(self.inner) }
696    }
697}
698
699#[test]
700#[ignore = "This test doesn't work in CI"]
701fn test_run_session_with_callback_info_api() {
702    let file = Path::new("tests/assets/realesr.mnn")
703        .canonicalize()
704        .unwrap();
705    let interpreter = Interpreter::from_file(&file).unwrap();
706    let session = interpreter.create_session(ScheduleConfig::new()).unwrap();
707    interpreter
708        .run_session_with_callback(
709            &session,
710            TensorCallback::identity(),
711            TensorCallback::identity(),
712            true,
713        )
714        .unwrap();
715}
716
717#[test]
718#[ignore = "This test doesn't work in CI"]
719fn check_whether_sync_actually_works() {
720    let file = Path::new("tests/assets/realesr.mnn")
721        .canonicalize()
722        .unwrap();
723    let interpreter = Interpreter::from_file(&file).unwrap();
724    let session = interpreter.create_session(ScheduleConfig::new()).unwrap();
725    let time = std::time::Instant::now();
726    interpreter
727        .run_session_with_callback(
728            &session,
729            TensorCallback::identity(),
730            TensorCallback::identity(),
731            false,
732        )
733        .unwrap();
734    let time = time.elapsed();
735    let time2 = std::time::Instant::now();
736    interpreter
737        .run_session_with_callback(
738            &session,
739            TensorCallback::identity(),
740            TensorCallback::identity(),
741            true,
742        )
743        .unwrap();
744    let time2 = time2.elapsed();
745    assert!((time - time2) > std::time::Duration::from_millis(50));
746}
mnn/interpreter.rs

mnn/
interpreter.rs
