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/// CoreML hardware selection used by the ONNX Runtime CoreML execution provider.
25#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, Default)]
26pub enum OrtCoreMLComputeUnits {
27    /// Let CoreML select from CPU, GPU, and Neural Engine.
28    #[default]
29    All,
30    /// Restrict CoreML to CPU and GPU.
31    CPUAndGPU,
32    /// Restrict CoreML to CPU and Neural Engine.
33    CPUAndNeuralEngine,
34    /// Restrict CoreML to CPU.
35    CPUOnly,
36}
37
38/// CoreML model representation created by ONNX Runtime.
39#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, Default)]
40pub enum OrtCoreMLModelFormat {
41    /// The modern CoreML representation (macOS 12+), with broader operator support.
42    #[default]
43    MLProgram,
44    /// The legacy CoreML neural-network representation.
45    NeuralNetwork,
46}
47
48/// CoreML graph-specialization policy.
49#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, Default)]
50pub enum OrtCoreMLSpecializationStrategy {
51    /// CoreML's balanced default.
52    #[default]
53    Default,
54    /// Prefer steady-state prediction latency over specialization time and size.
55    FastPrediction,
56}
57
58/// Advanced CoreML execution-provider options.
59///
60/// These options live on [`OrtSessionConfig`] instead of adding fields to
61/// [`OrtExecutionProvider::CoreML`], preserving source compatibility for code
62/// that constructs or exhaustively matches the provider variant.
63#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, Default)]
64pub struct OrtCoreMLConfig {
65    /// Hardware units available to CoreML.
66    pub compute_units: Option<OrtCoreMLComputeUnits>,
67    /// CoreML model representation.
68    pub model_format: Option<OrtCoreMLModelFormat>,
69    /// Only claim nodes whose model inputs have static shapes.
70    pub static_input_shapes: Option<bool>,
71    /// CoreML graph-specialization policy.
72    pub specialization_strategy: Option<OrtCoreMLSpecializationStrategy>,
73    /// Permit FP16 accumulation on the GPU.
74    pub allow_low_precision_accumulation_on_gpu: Option<bool>,
75    /// Log CoreML's hardware assignment and estimated cost.
76    pub profile_compute_plan: Option<bool>,
77    /// Directory used to cache compiled CoreML models.
78    pub model_cache_dir: Option<String>,
79}
80
81pub(crate) const COREML_CONFIG_ENTRY: &str = "oar.internal.coreml_config";
82
83/// Execution providers for ONNX Runtime.
84///
85/// This enum represents the different execution providers that can be used
86/// with ONNX Runtime for model inference.
87#[derive(Debug, Clone, PartialEq, Serialize, Deserialize, Default)]
88pub enum OrtExecutionProvider {
89    /// CPU execution provider (always available)
90    #[default]
91    CPU,
92    /// NVIDIA CUDA execution provider
93    CUDA {
94        /// CUDA device ID (default: 0)
95        device_id: Option<i32>,
96        /// Memory limit in bytes (optional)
97        gpu_mem_limit: Option<usize>,
98        /// Arena extend strategy: "NextPowerOfTwo" or "SameAsRequested"
99        arena_extend_strategy: Option<String>,
100        /// CUDNN convolution algorithm search: "Exhaustive", "Heuristic", or "Default"
101        cudnn_conv_algo_search: Option<String>,
102        /// CUDNN convolution use max workspace (default: true)
103        cudnn_conv_use_max_workspace: Option<bool>,
104    },
105    /// DirectML execution provider (Windows only)
106    DirectML {
107        /// DirectML device ID (default: 0)
108        device_id: Option<i32>,
109    },
110    /// OpenVINO execution provider
111    OpenVINO {
112        /// Device type (e.g., "CPU", "GPU", "MYRIAD")
113        device_type: Option<String>,
114        /// Number of threads (optional)
115        num_threads: Option<usize>,
116    },
117    /// TensorRT execution provider
118    TensorRT {
119        /// TensorRT device ID (default: 0)
120        device_id: Option<i32>,
121        /// Maximum workspace size in bytes
122        max_workspace_size: Option<usize>,
123        /// Minimum subgraph size for TensorRT acceleration
124        min_subgraph_size: Option<usize>,
125        /// FP16 enable flag
126        fp16_enable: Option<bool>,
127        /// Enable use of timing cache to speed up builds
128        timing_cache: Option<bool>,
129        /// Set path for storing timing cache
130        timing_cache_path: Option<String>,
131        /// Force use of timing cache regardless of GPU match
132        force_timing_cache: Option<bool>,
133        /// Enable caching of TensorRT engines
134        engine_cache: Option<bool>,
135        /// Set path to store cached TensorRT engines
136        engine_cache_path: Option<String>,
137        /// Dump ep context model
138        dump_ep_context_model: Option<bool>,
139        /// The path of an embedded engine model
140        ep_context_file_path: Option<String>,
141    },
142    /// CoreML execution provider (macOS/iOS only)
143    CoreML {
144        /// Use CPU and Apple Neural Engine compute units. Despite the
145        /// historical name, unsupported nodes may still execute on CPU.
146        ane_only: Option<bool>,
147        /// Enable subgraphs
148        subgraphs: Option<bool>,
149    },
150    /// WebGPU execution provider
151    WebGPU,
152}
153
154/// Configuration for ONNX Runtime sessions.
155///
156/// This struct contains various configuration options for ONNX Runtime sessions,
157/// including threading, memory management, and optimization settings.
158#[derive(Debug, Clone, Default, Serialize, Deserialize)]
159pub struct OrtSessionConfig {
160    /// Number of threads used to parallelize execution within nodes
161    pub intra_threads: Option<usize>,
162    /// Number of threads used to parallelize execution across nodes
163    pub inter_threads: Option<usize>,
164    /// Enable parallel execution mode
165    pub parallel_execution: Option<bool>,
166    /// Graph optimization level
167    pub optimization_level: Option<OrtGraphOptimizationLevel>,
168    /// Execution providers in order of preference
169    pub execution_providers: Option<Vec<OrtExecutionProvider>>,
170    /// Enable memory pattern optimization
171    pub enable_mem_pattern: Option<bool>,
172    /// Log severity level (0=Verbose, 1=Info, 2=Warning, 3=Error, 4=Fatal)
173    pub log_severity_level: Option<i32>,
174    /// Log verbosity level
175    pub log_verbosity_level: Option<i32>,
176    /// Session configuration entries (key-value pairs)
177    pub session_config_entries: Option<std::collections::HashMap<String, String>>,
178}
179
180impl OrtSessionConfig {
181    /// Creates a new OrtSessionConfig with default values.
182    pub fn new() -> Self {
183        Self::default()
184    }
185
186    /// Sets the number of intra-op threads.
187    pub fn with_intra_threads(mut self, threads: usize) -> Self {
188        self.intra_threads = Some(threads);
189        self
190    }
191
192    /// Sets the number of inter-op threads.
193    pub fn with_inter_threads(mut self, threads: usize) -> Self {
194        self.inter_threads = Some(threads);
195        self
196    }
197
198    /// Enables or disables parallel execution.
199    pub fn with_parallel_execution(mut self, enabled: bool) -> Self {
200        self.parallel_execution = Some(enabled);
201        self
202    }
203
204    /// Sets the graph optimization level.
205    pub fn with_optimization_level(mut self, level: OrtGraphOptimizationLevel) -> Self {
206        self.optimization_level = Some(level);
207        self
208    }
209
210    /// Sets the execution providers, in order of preference.
211    pub fn with_execution_providers(mut self, providers: Vec<OrtExecutionProvider>) -> Self {
212        self.execution_providers = Some(providers);
213        self
214    }
215
216    /// Appends a single execution provider.
217    pub fn add_execution_provider(mut self, provider: OrtExecutionProvider) -> Self {
218        if let Some(ref mut providers) = self.execution_providers {
219            providers.push(provider);
220        } else {
221            self.execution_providers = Some(vec![provider]);
222        }
223        self
224    }
225
226    /// Enables or disables memory pattern optimization.
227    pub fn with_memory_pattern(mut self, enable: bool) -> Self {
228        self.enable_mem_pattern = Some(enable);
229        self
230    }
231
232    /// Sets the log severity level (0=Verbose, 1=Info, 2=Warning, 3=Error, 4=Fatal).
233    pub fn with_log_severity_level(mut self, level: i32) -> Self {
234        self.log_severity_level = Some(level);
235        self
236    }
237
238    /// Sets the log verbosity level.
239    pub fn with_log_verbosity_level(mut self, level: i32) -> Self {
240        self.log_verbosity_level = Some(level);
241        self
242    }
243
244    /// Adds a session configuration entry.
245    pub fn add_config_entry<K: Into<String>, V: Into<String>>(mut self, key: K, value: V) -> Self {
246        if let Some(ref mut entries) = self.session_config_entries {
247            entries.insert(key.into(), value.into());
248        } else {
249            let mut entries = std::collections::HashMap::new();
250            entries.insert(key.into(), value.into());
251            self.session_config_entries = Some(entries);
252        }
253        self
254    }
255
256    /// Sets advanced options for any CoreML execution provider in this session.
257    pub fn with_coreml_config(mut self, config: OrtCoreMLConfig) -> Self {
258        let value =
259            serde_json::to_string(&config).expect("serializing OrtCoreMLConfig cannot fail");
260        self.session_config_entries
261            .get_or_insert_with(Default::default)
262            .insert(COREML_CONFIG_ENTRY.to_owned(), value);
263        self
264    }
265
266    pub(crate) fn coreml_config(&self) -> Result<Option<OrtCoreMLConfig>, serde_json::Error> {
267        self.session_config_entries
268            .as_ref()
269            .and_then(|entries| entries.get(COREML_CONFIG_ENTRY))
270            .map(|value| serde_json::from_str(value))
271            .transpose()
272    }
273
274    /// Effective intra-op thread count, defaulting to available parallelism.
275    pub fn get_intra_threads(&self) -> usize {
276        self.intra_threads.unwrap_or_else(|| {
277            std::thread::available_parallelism()
278                .map(|n| n.get())
279                .unwrap_or(1)
280        })
281    }
282
283    /// Effective inter-op thread count, defaulting to 1.
284    pub fn get_inter_threads(&self) -> usize {
285        self.inter_threads.unwrap_or(1)
286    }
287
288    /// Effective graph optimization level, defaulting to `OrtGraphOptimizationLevel::default()`.
289    pub fn get_optimization_level(&self) -> OrtGraphOptimizationLevel {
290        self.optimization_level.unwrap_or_default()
291    }
292
293    /// Configured execution providers, defaulting to CPU.
294    pub fn get_execution_providers(&self) -> Vec<OrtExecutionProvider> {
295        self.execution_providers
296            .clone()
297            .unwrap_or_else(|| vec![OrtExecutionProvider::CPU])
298    }
299
300    /// Returns whether an explicitly configured hardware accelerator is present.
301    ///
302    /// `execution_providers` is a preference-ordered list: ONNX Runtime lets
303    /// each provider claim graph nodes in list order, and CPU can claim
304    /// almost any node, so a CPU entry listed first effectively runs the
305    /// session on CPU regardless of what accelerators follow it. Only the
306    /// first provider therefore determines whether this is an accelerated
307    /// configuration. No provider configuration, an empty provider list, and
308    /// a CPU-first list (including CPU alone) all use CPU-oriented pipeline
309    /// defaults; an accelerator listed first with a CPU fallback after it
310    /// (CUDA, TensorRT, DirectML, OpenVINO, CoreML, or WebGPU) counts as
311    /// accelerated.
312    pub fn has_accelerator_provider(&self) -> bool {
313        self.execution_providers
314            .as_ref()
315            .and_then(|providers| providers.first())
316            .is_some_and(|provider| !matches!(provider, OrtExecutionProvider::CPU))
317    }
318}
319
320#[cfg(test)]
321mod tests {
322    use super::*;
323
324    #[test]
325    fn test_ort_session_config_builder() {
326        let config = OrtSessionConfig::new()
327            .with_intra_threads(4)
328            .with_inter_threads(2)
329            .with_optimization_level(OrtGraphOptimizationLevel::Level2)
330            .with_memory_pattern(true)
331            .add_execution_provider(OrtExecutionProvider::CPU);
332
333        assert_eq!(config.intra_threads, Some(4));
334        assert_eq!(config.inter_threads, Some(2));
335        assert!(matches!(
336            config.optimization_level,
337            Some(OrtGraphOptimizationLevel::Level2)
338        ));
339        assert_eq!(config.enable_mem_pattern, Some(true));
340        assert!(config.execution_providers.is_some());
341    }
342
343    #[test]
344    fn test_ort_session_config_getters() {
345        let config = OrtSessionConfig::new()
346            .with_intra_threads(8)
347            .with_inter_threads(4)
348            .with_optimization_level(OrtGraphOptimizationLevel::All);
349
350        assert_eq!(config.get_intra_threads(), 8);
351        assert_eq!(config.get_inter_threads(), 4);
352        assert!(matches!(
353            config.get_optimization_level(),
354            OrtGraphOptimizationLevel::All
355        ));
356    }
357
358    #[test]
359    fn test_accelerator_provider_detection() {
360        assert!(!OrtSessionConfig::new().has_accelerator_provider());
361        assert!(
362            !OrtSessionConfig::new()
363                .with_execution_providers(vec![OrtExecutionProvider::CPU])
364                .has_accelerator_provider()
365        );
366        assert!(
367            OrtSessionConfig::new()
368                .with_execution_providers(vec![
369                    OrtExecutionProvider::DirectML { device_id: Some(0) },
370                    OrtExecutionProvider::CPU,
371                ])
372                .has_accelerator_provider()
373        );
374        // CPU listed first claims nearly every node before the accelerator
375        // gets a chance to, so this is a CPU-preferred configuration despite
376        // the accelerator appearing later in the list.
377        assert!(
378            !OrtSessionConfig::new()
379                .with_execution_providers(vec![
380                    OrtExecutionProvider::CPU,
381                    OrtExecutionProvider::DirectML { device_id: Some(0) },
382                ])
383                .has_accelerator_provider()
384        );
385    }
386
387    #[test]
388    fn coreml_provider_keeps_legacy_variant_shape() {
389        let provider = OrtExecutionProvider::CoreML {
390            ane_only: Some(true),
391            subgraphs: Some(false),
392        };
393        let OrtExecutionProvider::CoreML {
394            ane_only,
395            subgraphs,
396        } = provider
397        else {
398            unreachable!()
399        };
400        assert_eq!(ane_only, Some(true));
401        assert_eq!(subgraphs, Some(false));
402    }
403
404    #[test]
405    fn coreml_advanced_config_round_trips_through_session_config() {
406        let expected = OrtCoreMLConfig {
407            compute_units: Some(OrtCoreMLComputeUnits::CPUAndGPU),
408            model_format: Some(OrtCoreMLModelFormat::MLProgram),
409            static_input_shapes: Some(true),
410            specialization_strategy: Some(OrtCoreMLSpecializationStrategy::FastPrediction),
411            allow_low_precision_accumulation_on_gpu: Some(true),
412            profile_compute_plan: None,
413            model_cache_dir: Some("cache".to_owned()),
414        };
415        let config = OrtSessionConfig::new().with_coreml_config(expected.clone());
416        assert_eq!(config.coreml_config().unwrap(), Some(expected));
417    }
418}
oar_ocr_core/core/config/onnx.rs

oar_ocr_core/core/config/
onnx.rs
