oar_ocr_core/core/config/
onnx.rs

1//! ONNX Runtime configuration types and utilities.
2
3use serde::{Deserialize, Serialize};
4
5/// Graph optimization levels for ONNX Runtime.
6///
7/// This enum represents the different levels of graph optimization that can be applied
8/// during ONNX Runtime session creation.
9#[derive(Debug, Clone, Copy, Serialize, Deserialize, Default)]
10pub enum OrtGraphOptimizationLevel {
11    /// Disable all optimizations.
12    DisableAll,
13    /// Enable basic optimizations.
14    #[default]
15    Level1,
16    /// Enable extended optimizations.
17    Level2,
18    /// Enable all optimizations.
19    Level3,
20    /// Enable all optimizations (alias for Level3).
21    All,
22}
23
24/// Execution providers for ONNX Runtime.
25///
26/// This enum represents the different execution providers that can be used
27/// with ONNX Runtime for model inference.
28#[derive(Debug, Clone, PartialEq, Serialize, Deserialize, Default)]
29pub enum OrtExecutionProvider {
30    /// CPU execution provider (always available)
31    #[default]
32    CPU,
33    /// NVIDIA CUDA execution provider
34    CUDA {
35        /// CUDA device ID (default: 0)
36        device_id: Option<i32>,
37        /// Memory limit in bytes (optional)
38        gpu_mem_limit: Option<usize>,
39        /// Arena extend strategy: "NextPowerOfTwo" or "SameAsRequested"
40        arena_extend_strategy: Option<String>,
41        /// CUDNN convolution algorithm search: "Exhaustive", "Heuristic", or "Default"
42        cudnn_conv_algo_search: Option<String>,
43        /// CUDNN convolution use max workspace (default: true)
44        cudnn_conv_use_max_workspace: Option<bool>,
45    },
46    /// DirectML execution provider (Windows only)
47    DirectML {
48        /// DirectML device ID (default: 0)
49        device_id: Option<i32>,
50    },
51    /// OpenVINO execution provider
52    OpenVINO {
53        /// Device type (e.g., "CPU", "GPU", "MYRIAD")
54        device_type: Option<String>,
55        /// Number of threads (optional)
56        num_threads: Option<usize>,
57    },
58    /// TensorRT execution provider
59    TensorRT {
60        /// TensorRT device ID (default: 0)
61        device_id: Option<i32>,
62        /// Maximum workspace size in bytes
63        max_workspace_size: Option<usize>,
64        /// Minimum subgraph size for TensorRT acceleration
65        min_subgraph_size: Option<usize>,
66        /// FP16 enable flag
67        fp16_enable: Option<bool>,
68        /// Enable use of timing cache to speed up builds
69        timing_cache: Option<bool>,
70        /// Set path for storing timing cache
71        timing_cache_path: Option<String>,
72        /// Force use of timing cache regardless of GPU match
73        force_timing_cache: Option<bool>,
74        /// Enable caching of TensorRT engines
75        engine_cache: Option<bool>,
76        /// Set path to store cached TensorRT engines
77        engine_cache_path: Option<String>,
78        /// Dump ep context model
79        dump_ep_context_model: Option<bool>,
80        /// The path of an embedded engine model
81        ep_context_file_path: Option<String>,
82    },
83    /// CoreML execution provider (macOS/iOS only)
84    CoreML {
85        /// Use Apple Neural Engine only
86        ane_only: Option<bool>,
87        /// Enable subgraphs
88        subgraphs: Option<bool>,
89    },
90    /// WebGPU execution provider
91    WebGPU,
92}
93
94/// Configuration for ONNX Runtime sessions.
95///
96/// This struct contains various configuration options for ONNX Runtime sessions,
97/// including threading, memory management, and optimization settings.
98#[derive(Debug, Clone, Default, Serialize, Deserialize)]
99pub struct OrtSessionConfig {
100    /// Number of threads used to parallelize execution within nodes
101    pub intra_threads: Option<usize>,
102    /// Number of threads used to parallelize execution across nodes
103    pub inter_threads: Option<usize>,
104    /// Enable parallel execution mode
105    pub parallel_execution: Option<bool>,
106    /// Graph optimization level
107    pub optimization_level: Option<OrtGraphOptimizationLevel>,
108    /// Execution providers in order of preference
109    pub execution_providers: Option<Vec<OrtExecutionProvider>>,
110    /// Enable memory pattern optimization
111    pub enable_mem_pattern: Option<bool>,
112    /// Log severity level (0=Verbose, 1=Info, 2=Warning, 3=Error, 4=Fatal)
113    pub log_severity_level: Option<i32>,
114    /// Log verbosity level
115    pub log_verbosity_level: Option<i32>,
116    /// Session configuration entries (key-value pairs)
117    pub session_config_entries: Option<std::collections::HashMap<String, String>>,
118}
119
120impl OrtSessionConfig {
121    /// Creates a new OrtSessionConfig with default values.
122    pub fn new() -> Self {
123        Self::default()
124    }
125
126    /// Sets the number of intra-op threads.
127    pub fn with_intra_threads(mut self, threads: usize) -> Self {
128        self.intra_threads = Some(threads);
129        self
130    }
131
132    /// Sets the number of inter-op threads.
133    pub fn with_inter_threads(mut self, threads: usize) -> Self {
134        self.inter_threads = Some(threads);
135        self
136    }
137
138    /// Enables or disables parallel execution.
139    pub fn with_parallel_execution(mut self, enabled: bool) -> Self {
140        self.parallel_execution = Some(enabled);
141        self
142    }
143
144    /// Sets the graph optimization level.
145    pub fn with_optimization_level(mut self, level: OrtGraphOptimizationLevel) -> Self {
146        self.optimization_level = Some(level);
147        self
148    }
149
150    /// Sets the execution providers, in order of preference.
151    pub fn with_execution_providers(mut self, providers: Vec<OrtExecutionProvider>) -> Self {
152        self.execution_providers = Some(providers);
153        self
154    }
155
156    /// Appends a single execution provider.
157    pub fn add_execution_provider(mut self, provider: OrtExecutionProvider) -> Self {
158        if let Some(ref mut providers) = self.execution_providers {
159            providers.push(provider);
160        } else {
161            self.execution_providers = Some(vec![provider]);
162        }
163        self
164    }
165
166    /// Enables or disables memory pattern optimization.
167    pub fn with_memory_pattern(mut self, enable: bool) -> Self {
168        self.enable_mem_pattern = Some(enable);
169        self
170    }
171
172    /// Sets the log severity level (0=Verbose, 1=Info, 2=Warning, 3=Error, 4=Fatal).
173    pub fn with_log_severity_level(mut self, level: i32) -> Self {
174        self.log_severity_level = Some(level);
175        self
176    }
177
178    /// Sets the log verbosity level.
179    pub fn with_log_verbosity_level(mut self, level: i32) -> Self {
180        self.log_verbosity_level = Some(level);
181        self
182    }
183
184    /// Adds a session configuration entry.
185    pub fn add_config_entry<K: Into<String>, V: Into<String>>(mut self, key: K, value: V) -> Self {
186        if let Some(ref mut entries) = self.session_config_entries {
187            entries.insert(key.into(), value.into());
188        } else {
189            let mut entries = std::collections::HashMap::new();
190            entries.insert(key.into(), value.into());
191            self.session_config_entries = Some(entries);
192        }
193        self
194    }
195
196    /// Effective intra-op thread count, defaulting to available parallelism.
197    pub fn get_intra_threads(&self) -> usize {
198        self.intra_threads.unwrap_or_else(|| {
199            std::thread::available_parallelism()
200                .map(|n| n.get())
201                .unwrap_or(1)
202        })
203    }
204
205    /// Effective inter-op thread count, defaulting to 1.
206    pub fn get_inter_threads(&self) -> usize {
207        self.inter_threads.unwrap_or(1)
208    }
209
210    /// Effective graph optimization level, defaulting to `OrtGraphOptimizationLevel::default()`.
211    pub fn get_optimization_level(&self) -> OrtGraphOptimizationLevel {
212        self.optimization_level.unwrap_or_default()
213    }
214
215    /// Configured execution providers, defaulting to CPU.
216    pub fn get_execution_providers(&self) -> Vec<OrtExecutionProvider> {
217        self.execution_providers
218            .clone()
219            .unwrap_or_else(|| vec![OrtExecutionProvider::CPU])
220    }
221}
222
223#[cfg(test)]
224mod tests {
225    use super::*;
226
227    #[test]
228    fn test_ort_session_config_builder() {
229        let config = OrtSessionConfig::new()
230            .with_intra_threads(4)
231            .with_inter_threads(2)
232            .with_optimization_level(OrtGraphOptimizationLevel::Level2)
233            .with_memory_pattern(true)
234            .add_execution_provider(OrtExecutionProvider::CPU);
235
236        assert_eq!(config.intra_threads, Some(4));
237        assert_eq!(config.inter_threads, Some(2));
238        assert!(matches!(
239            config.optimization_level,
240            Some(OrtGraphOptimizationLevel::Level2)
241        ));
242        assert_eq!(config.enable_mem_pattern, Some(true));
243        assert!(config.execution_providers.is_some());
244    }
245
246    #[test]
247    fn test_ort_session_config_getters() {
248        let config = OrtSessionConfig::new()
249            .with_intra_threads(8)
250            .with_inter_threads(4)
251            .with_optimization_level(OrtGraphOptimizationLevel::All);
252
253        assert_eq!(config.get_intra_threads(), 8);
254        assert_eq!(config.get_inter_threads(), 4);
255        assert!(matches!(
256            config.get_optimization_level(),
257            OrtGraphOptimizationLevel::All
258        ));
259    }
260}
oar_ocr_core/core/config/onnx.rs

oar_ocr_core/core/config/
onnx.rs
