mnn/
backend.rs

1//! The backend module contains the data types for the backend configuration
2
3use crate::prelude::*;
4use std::str::FromStr;
5
6use mnn_sys::*;
7
8/// BackendConfig is a struct that holds the configuration for the backend
9/// memory: [MemoryMode]
10/// power: [PowerMode]
11/// precision: [PrecisionMode]
12#[repr(transparent)]
13pub struct BackendConfig {
14    pub(crate) inner: *mut MNNBackendConfig,
15    __marker: core::marker::PhantomData<()>,
16}
17
18impl core::fmt::Debug for BackendConfig {
19    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
20        f.debug_struct("BackendConfig")
21            .field("memory", &self.get_memory_mode())
22            .field("power", &self.get_power_mode())
23            .field("precision", &self.get_precision_mode())
24            .finish()
25    }
26}
27
28#[cfg(feature = "serde")]
29impl serde::Serialize for BackendConfig {
30    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
31    where
32        S: serde::ser::Serializer,
33    {
34        use serde::ser::SerializeStruct;
35        let mut state = serializer.serialize_struct("BackendConfig", 3)?;
36        state.serialize_field("memory", &self.get_memory_mode())?;
37        state.serialize_field("power", &self.get_power_mode())?;
38        state.serialize_field("precision", &self.get_precision_mode())?;
39        state.end()
40    }
41}
42
43impl Clone for BackendConfig {
44    fn clone(&self) -> Self {
45        unsafe {
46            let inner = mnn_sys::mnnbc_clone(self.inner);
47            Self {
48                inner,
49                __marker: core::marker::PhantomData,
50            }
51        }
52    }
53}
54
55impl Drop for BackendConfig {
56    fn drop(&mut self) {
57        unsafe {
58            mnn_sys::mnnbc_destroy(self.inner);
59        }
60    }
61}
62
63impl Default for BackendConfig {
64    fn default() -> Self {
65        Self::new()
66    }
67}
68
69/// PowerModes depend on if the specific backend has support for it
70#[derive(Debug, Clone, Copy, Eq, PartialEq, Hash)]
71#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
72pub enum PowerMode {
73    /// Low power mode
74    Low,
75    /// Normal power mode
76    Normal,
77    /// High power mode
78    High,
79}
80
81impl PowerMode {
82    fn to_mnn_sys(self) -> mnn_sys::PowerMode {
83        match self {
84            Self::Low => mnn_sys::PowerMode::Power_Low,
85            Self::Normal => mnn_sys::PowerMode::Power_Normal,
86            Self::High => mnn_sys::PowerMode::Power_High,
87        }
88    }
89
90    /// Returns a string representation of the power mode
91    pub fn to_str(self) -> &'static str {
92        match self {
93            Self::Low => "low",
94            Self::Normal => "normal",
95            Self::High => "high",
96        }
97    }
98
99    fn from_mnn_sys(mode: mnn_sys::PowerMode) -> Self {
100        match mode {
101            mnn_sys::PowerMode::Power_Low => Self::Low,
102            mnn_sys::PowerMode::Power_Normal => Self::Normal,
103            mnn_sys::PowerMode::Power_High => Self::High,
104            _ => Self::Normal,
105        }
106    }
107}
108
109impl FromStr for PowerMode {
110    type Err = MNNError;
111    fn from_str(s: &str) -> Result<Self, Self::Err> {
112        match s {
113            "low" => Ok(Self::Low),
114            "normal" => Ok(Self::Normal),
115            "high" => Ok(Self::High),
116            _ => Err(error!(ErrorKind::ParseError).attach(format!("invalid power mode: {s}"))),
117        }
118    }
119}
120
121impl FromStr for MemoryMode {
122    type Err = MNNError;
123    fn from_str(s: &str) -> Result<Self, Self::Err> {
124        match s {
125            "low" => Ok(Self::Low),
126            "normal" => Ok(Self::Normal),
127            "high" => Ok(Self::High),
128            _ => Err(error!(ErrorKind::ParseError).attach(format!("invalid memory mode: {s}"))),
129        }
130    }
131}
132
133impl FromStr for PrecisionMode {
134    type Err = MNNError;
135    fn from_str(s: &str) -> Result<Self, Self::Err> {
136        match s {
137            "low" => Ok(Self::Low),
138            "normal" => Ok(Self::Normal),
139            "high" => Ok(Self::High),
140            "low_bf16" => Ok(Self::LowBf16),
141            _ => Err(error!(ErrorKind::ParseError).attach(format!("invalid precision mode: {s}"))),
142        }
143    }
144}
145
146/// MemoryModes depend on if the specific backend has support for it
147#[derive(Debug, Clone, Copy, Eq, PartialEq, Hash)]
148#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
149pub enum MemoryMode {
150    /// Low memory mode
151    Low,
152    /// Normal memory mode
153    Normal,
154    /// High memory mode
155    High,
156}
157
158impl MemoryMode {
159    fn to_mnn_sys(self) -> mnn_sys::MemoryMode {
160        match self {
161            Self::Low => mnn_sys::MemoryMode::Memory_Low,
162            Self::Normal => mnn_sys::MemoryMode::Memory_Normal,
163            Self::High => mnn_sys::MemoryMode::Memory_High,
164        }
165    }
166
167    /// Returns a string representation of the memory mode
168    pub fn to_str(self) -> &'static str {
169        match self {
170            Self::Low => "low",
171            Self::Normal => "normal",
172            Self::High => "high",
173        }
174    }
175
176    fn from_mnn_sys(mode: mnn_sys::MemoryMode) -> Self {
177        match mode {
178            mnn_sys::MemoryMode::Memory_Low => Self::Low,
179            mnn_sys::MemoryMode::Memory_Normal => Self::Normal,
180            mnn_sys::MemoryMode::Memory_High => Self::High,
181            _ => Self::Normal,
182        }
183    }
184}
185
186/// PrecisionModes depend on if the specific backend has support for it
187#[derive(Debug, Clone, Copy, Eq, PartialEq, Hash)]
188#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
189pub enum PrecisionMode {
190    /// Normal precision mode
191    Normal = 0,
192    /// High precision mode
193    High,
194    /// Low precision mode
195    Low,
196    /// Low precision mode with BF16
197    LowBf16,
198}
199impl PrecisionMode {
200    pub(crate) fn to_mnn_sys(self) -> mnn_sys::PrecisionMode {
201        match self {
202            Self::LowBf16 => mnn_sys::PrecisionMode::Precision_Low_BF16,
203            Self::Low => mnn_sys::PrecisionMode::Precision_Low,
204            Self::Normal => mnn_sys::PrecisionMode::Precision_Normal,
205            Self::High => mnn_sys::PrecisionMode::Precision_High,
206        }
207    }
208
209    /// Returns a string representation of the precision mode
210    pub fn to_str(self) -> &'static str {
211        match self {
212            Self::LowBf16 => "low_bf16",
213            Self::Low => "low",
214            Self::Normal => "normal",
215            Self::High => "high",
216        }
217    }
218
219    fn from_mnn_sys(mode: mnn_sys::PrecisionMode) -> Self {
220        match mode {
221            mnn_sys::PrecisionMode::Precision_Low_BF16 => Self::LowBf16,
222            mnn_sys::PrecisionMode::Precision_Low => Self::Low,
223            mnn_sys::PrecisionMode::Precision_Normal => Self::Normal,
224            mnn_sys::PrecisionMode::Precision_High => Self::High,
225            _ => Self::Normal,
226        }
227    }
228}
229
230impl BackendConfig {
231    /// Create a new backend config
232    pub fn new() -> Self {
233        unsafe {
234            let inner = mnnbc_create();
235            Self {
236                inner,
237                __marker: core::marker::PhantomData,
238            }
239        }
240    }
241
242    /// Sets the [MemoryMode] for the backend
243    pub fn set_memory_mode(&mut self, mode: MemoryMode) {
244        unsafe {
245            mnn_sys::mnnbc_set_memory_mode(self.inner, mode.to_mnn_sys());
246        }
247    }
248
249    /// Sets the [MemoryMode] for the backend
250    pub fn with_memory_mode(mut self, mode: MemoryMode) -> Self {
251        self.set_memory_mode(mode);
252        self
253    }
254
255    /// Gets the [MemoryMode] for the backend
256    pub fn get_memory_mode(&self) -> MemoryMode {
257        unsafe { MemoryMode::from_mnn_sys(mnn_sys::mnnbc_get_memory_mode(self.inner)) }
258    }
259
260    /// Sets the [PowerMode] for the backend
261    pub fn set_power_mode(&mut self, mode: PowerMode) {
262        unsafe {
263            mnn_sys::mnnbc_set_power_mode(self.inner, mode.to_mnn_sys());
264        }
265    }
266
267    /// Sets the [PowerMode] for the backend
268    pub fn with_power_mode(mut self, mode: PowerMode) -> Self {
269        self.set_power_mode(mode);
270        self
271    }
272
273    /// Gets the [PowerMode] for the backend
274    pub fn get_power_mode(&self) -> PowerMode {
275        unsafe { PowerMode::from_mnn_sys(mnn_sys::mnnbc_get_power_mode(self.inner)) }
276    }
277
278    /// Sets the [PrecisionMode] for the backend
279    pub fn set_precision_mode(&mut self, mode: PrecisionMode) {
280        unsafe {
281            mnn_sys::mnnbc_set_precision_mode(self.inner, mode.to_mnn_sys());
282        }
283    }
284
285    /// Sets the [PrecisionMode] for the backend
286    pub fn with_precision_mode(mut self, mode: PrecisionMode) -> Self {
287        self.set_precision_mode(mode);
288        self
289    }
290
291    /// Gets the [PrecisionMode] for the backend
292    pub fn get_precision_mode(&self) -> PrecisionMode {
293        unsafe { PrecisionMode::from_mnn_sys(mnn_sys::mnnbc_get_precision_mode(self.inner)) }
294    }
295
296    /// Sets the flags for the backend
297    /// What the flag represents is depends on each backend or isn't documented
298    pub fn set_flags(&mut self, flags: usize) {
299        unsafe {
300            mnn_sys::mnnbc_set_flags(self.inner, flags);
301        }
302    }
303
304    /// Sets the flags for the backend
305    pub fn with_flags(mut self, flags: usize) -> Self {
306        self.set_flags(flags);
307        self
308    }
309
310    /// # Safety
311    /// This just binds to the underlying unsafe api and should be used only if you know what you
312    /// are doing
313    pub unsafe fn set_shared_context(&mut self, shared_context: *mut libc::c_void) {
314        unsafe {
315            mnn_sys::mnnbc_set_shared_context(self.inner, shared_context);
316        }
317    }
318
319    /// # Safety
320    /// This just binds to the underlying unsafe api and should be used only if you know what you
321    /// are doing
322    pub unsafe fn with_shared_context(mut self, shared_context: *mut libc::c_void) -> Self {
323        unsafe {
324            self.set_shared_context(shared_context);
325        }
326        self
327    }
328}
329
330#[test]
331fn test_backend_config() {
332    let mut config = BackendConfig::new();
333    config.set_memory_mode(MemoryMode::Low);
334    config.set_power_mode(PowerMode::Low);
335    config.set_precision_mode(PrecisionMode::Low);
336    let config = std::hint::black_box(config.clone());
337    assert_eq!(config.get_memory_mode(), MemoryMode::Low);
338    assert_eq!(config.get_power_mode(), PowerMode::Low);
339    assert_eq!(config.get_precision_mode(), PrecisionMode::Low);
340    let config = config
341        .with_memory_mode(MemoryMode::Normal)
342        .with_power_mode(PowerMode::Normal)
343        .with_precision_mode(PrecisionMode::Normal);
344    assert_eq!(config.get_memory_mode(), MemoryMode::Normal);
345    assert_eq!(config.get_power_mode(), PowerMode::Normal);
346    assert_eq!(config.get_precision_mode(), PrecisionMode::Normal);
347}
mnn/backend.rs

mnn/
backend.rs
