native_onnx/options.rs
1//! Public backend selection, compilation, and inference configuration.
2use ort::ep::ExecutionProviderDispatch;
3use std::{collections::BTreeMap, path::PathBuf};
4
5/// An installed ONNX Runtime execution provider.
6#[derive(Clone, Debug)]
7pub enum Backend {
8 /// NVIDIA TensorRT execution provider.
9 TensorRt,
10 /// NVIDIA CUDA execution provider.
11 Cuda,
12 /// Apple CoreML execution provider (macOS only).
13 #[cfg(target_os = "macos")]
14 CoreMl,
15 /// ONNX Runtime's CPU execution provider.
16 Cpu,
17 /// Configure any existing ORT provider without adding its native dependencies.
18 Custom {
19 /// Display name used in backend reports and fallback events.
20 name: String,
21 /// Configured provider to register with ONNX Runtime.
22 provider: ExecutionProviderDispatch,
23 },
24}
25impl Backend {
26 /// Display name used in backend reports and fallback events.
27 pub fn name(&self) -> &str {
28 match self {
29 Self::TensorRt => "TensorRT",
30 Self::Cuda => "CUDA",
31 #[cfg(target_os = "macos")]
32 Self::CoreMl => "CoreML",
33 Self::Cpu => "CPU",
34 Self::Custom { name, .. } => name,
35 }
36 }
37}
38
39/// Build and execution options used only by TensorRT.
40#[derive(Clone, Debug)]
41pub struct TensorRtOptions {
42 /// GPU ordinal, default 0.
43 pub device_id: i32,
44 /// Allow FP16 computations internally; does not change model I/O types. Default false.
45 pub fp16: bool,
46 /// Allow TF32 for TensorRT FP32 operations (default true).
47 /// The process-wide NVIDIA_TF32_OVERRIDE must be unset for true, or 0 for false.
48 pub tf32: bool,
49 /// Enable CUDA graphs for fixed-shape, supported tensor I/O.
50 pub cuda_graph: bool,
51 /// Maximum builder workspace in bytes, default 4 GiB on 64-bit targets.
52 pub workspace_bytes: usize,
53 /// Builder optimization level from 0 through 5, default 3.
54 pub builder_optimization_level: u8,
55 /// Maximum auxiliary streams; -1 lets TensorRT choose (the default).
56 pub auxiliary_streams: i8,
57 /// Allow sparse weights where supported, default true.
58 pub sparsity: bool,
59 /// ORT syntax, e.g. `images:1x3x640x640`. Set all three for dynamic models.
60 pub min_shapes: Option<String>,
61 /// Preferred dynamic shapes, in the same ORT syntax as [`Self::min_shapes`].
62 pub opt_shapes: Option<String>,
63 /// Maximum dynamic shapes, in the same ORT syntax as [`Self::min_shapes`].
64 pub max_shapes: Option<String>,
65}
66impl Default for TensorRtOptions {
67 fn default() -> Self {
68 Self {
69 device_id: 0,
70 fp16: false,
71 tf32: true,
72 cuda_graph: true,
73 workspace_bytes: usize::try_from(4_u64 * 1024 * 1024 * 1024).unwrap_or(usize::MAX),
74 builder_optimization_level: 3,
75 auxiliary_streams: -1,
76 sparsity: true,
77 min_shapes: None,
78 opt_shapes: None,
79 max_shapes: None,
80 }
81 }
82}
83
84/// Options used only by the built-in CUDA execution provider.
85#[derive(Clone, Copy, Debug)]
86pub struct CudaOptions {
87 /// GPU ordinal, default 0.
88 pub device_id: i32,
89 /// Allow TF32 for CUDA FP32 operations, default true.
90 pub tf32: bool,
91 /// Capture and replay fixed-shape models through persistent I/O bindings.
92 /// Disabled by default; ignored when graph I/O is not eligible for prepared buffers.
93 pub cuda_graph: bool,
94}
95impl Default for CudaOptions {
96 fn default() -> Self {
97 Self {
98 device_id: 0,
99 tf32: true,
100 cuda_graph: false,
101 }
102 }
103}
104impl CudaOptions {
105 pub(crate) fn validate(&self) -> crate::Result<()> {
106 anyhow::ensure!(self.device_id >= 0, "Negative CUDA device id");
107 Ok(())
108 }
109}
110
111/// One compilation target together with its provider-specific configuration.
112/// Compilation never selects a target automatically or falls back to another provider.
113#[derive(Clone, Debug)]
114pub enum CompileTarget {
115 /// Compile an embedded TensorRT engine with these builder/execution settings.
116 TensorRt(TensorRtOptions),
117 /// Optimize an ONNX graph for CUDA with these provider settings.
118 Cuda(CudaOptions),
119 /// Compile for CoreML using MLProgram (macOS only).
120 #[cfg(target_os = "macos")]
121 CoreMl,
122 /// Optimize for CPU; threading is configured in [`CompileOptions`].
123 Cpu,
124 /// Compile with an explicitly configured installed ORT provider.
125 Custom {
126 /// Display name used in compilation reports and errors.
127 name: String,
128 /// Provider configured with its compilation settings.
129 provider: ExecutionProviderDispatch,
130 },
131}
132
133/// Compilation configuration with a mandatory target and no automatic fallback.
134/// Construct with [`Self::new`], then adjust shared settings as needed.
135#[derive(Clone, Debug)]
136pub struct CompileOptions {
137 /// Required target and its provider-specific compilation settings.
138 pub target: CompileTarget,
139 /// Existing ORT shared library; see [`OnnxOptions::runtime_path`].
140 pub runtime_path: Option<PathBuf>,
141 /// Threads within an operator during compilation/validation, default 1.
142 /// Set to 0 to let ONNX Runtime choose.
143 pub intra_threads: usize,
144 /// Threads across operators, used when parallel execution is enabled; default 1.
145 pub inter_threads: usize,
146 /// Execute independent operators in parallel, default false.
147 /// CUDA graph capture is disabled in this mode.
148 pub parallel_execution: bool,
149 /// Resolve symbolic ONNX dimensions to fixed sizes before compilation.
150 pub dimension_overrides: BTreeMap<String, usize>,
151}
152
153impl CompileOptions {
154 /// Select a target explicitly and use defaults for the shared settings.
155 /// GPU targets must carry their configuration, including when using defaults.
156 pub fn new(target: CompileTarget) -> Self {
157 Self {
158 target,
159 runtime_path: None,
160 intra_threads: 1,
161 inter_threads: 1,
162 parallel_execution: false,
163 dimension_overrides: BTreeMap::new(),
164 }
165 }
166
167 pub(crate) fn into_runtime_options(self) -> crate::Result<OnnxOptions> {
168 let (backend, tensorrt, cuda) = match self.target {
169 CompileTarget::TensorRt(config) => {
170 config.validate()?;
171 (Backend::TensorRt, Some(config), None)
172 }
173 CompileTarget::Cuda(config) => {
174 config.validate()?;
175 (Backend::Cuda, None, Some(config))
176 }
177 #[cfg(target_os = "macos")]
178 CompileTarget::CoreMl => (Backend::CoreMl, None, None),
179 CompileTarget::Cpu => (Backend::Cpu, None, None),
180 CompileTarget::Custom { name, provider } => {
181 (Backend::Custom { name, provider }, None, None)
182 }
183 };
184 Ok(OnnxOptions {
185 runtime_path: self.runtime_path,
186 backend: BackendSelection::Require(backend),
187 tensorrt,
188 cuda,
189 intra_threads: self.intra_threads,
190 inter_threads: self.inter_threads,
191 parallel_execution: self.parallel_execution,
192 dimension_overrides: self.dimension_overrides,
193 profiling: None,
194 })
195 }
196}
197
198/// Select one required provider, or try candidates in order with fallback.
199#[derive(Clone, Debug)]
200pub enum BackendSelection {
201 /// Fail if registration, node placement, or inference cannot use this provider.
202 Require(Backend),
203 /// Allow ORT node partitioning and whole-session fallback, in this order.
204 Auto(Vec<Backend>),
205}
206impl BackendSelection {
207 pub(crate) fn candidates(&self) -> &[Backend] {
208 match self {
209 Self::Require(backend) => std::slice::from_ref(backend),
210 Self::Auto(backends) => backends,
211 }
212 }
213}
214impl Default for BackendSelection {
215 fn default() -> Self {
216 #[cfg(target_os = "macos")]
217 {
218 Self::Auto(vec![Backend::CoreMl, Backend::Cpu])
219 }
220 #[cfg(cuda_platform)]
221 {
222 Self::Auto(vec![Backend::TensorRt, Backend::Cuda, Backend::Cpu])
223 }
224 #[cfg(not(any(target_os = "macos", cuda_platform)))]
225 {
226 Self::Auto(vec![Backend::Cpu])
227 }
228 }
229}
230
231/// Runtime library, backend selection, threading, and model-shape configuration.
232#[derive(Clone, Debug)]
233pub struct OnnxOptions {
234 /// Existing ORT shared library; otherwise use ORT_DYLIB_PATH or the platform loader:
235 /// onnxruntime.dll on Windows, libonnxruntime.so on Linux, libonnxruntime.dylib on macOS.
236 /// The process shares one ORT runtime; select it before the first ORT operation.
237 pub runtime_path: Option<PathBuf>,
238 /// Required provider or ordered fallback candidates; defaults depend on the platform.
239 pub backend: BackendSelection,
240 /// Optional TensorRT overrides. `None` uses [`TensorRtOptions::default()`] when selected.
241 /// Ignored by other providers; does not enable/disable TensorRT.
242 pub tensorrt: Option<TensorRtOptions>,
243 /// Optional CUDA overrides. `None` uses [`CudaOptions::default()`] when selected.
244 /// Ignored by other providers; does not enable/disable CUDA.
245 pub cuda: Option<CudaOptions>,
246 /// Threads within an operator, default 1; 0 lets ONNX Runtime choose.
247 pub intra_threads: usize,
248 /// Threads across operators, default 1; used with parallel ORT execution.
249 pub inter_threads: usize,
250 /// Execute independent operators in parallel, default false.
251 /// Enables `inter_threads`; CUDA graph capture is disabled in this mode.
252 pub parallel_execution: bool,
253 /// Resolve symbolic ONNX dimensions to fixed sizes, enabling the prepared path.
254 pub dimension_overrides: BTreeMap<String, usize>,
255 /// Optional ORT profiling prefix. Leave unset for timing measurements.
256 pub profiling: Option<PathBuf>,
257}
258impl Default for OnnxOptions {
259 fn default() -> Self {
260 Self {
261 runtime_path: None,
262 backend: BackendSelection::default(),
263 tensorrt: None,
264 cuda: None,
265 intra_threads: 1,
266 inter_threads: 1,
267 parallel_execution: false,
268 dimension_overrides: BTreeMap::new(),
269 profiling: None,
270 }
271 }
272}
273impl OnnxOptions {
274 /// Require CPU inference without attempting GPU providers.
275 pub fn cpu() -> Self {
276 Self {
277 backend: BackendSelection::Require(Backend::Cpu),
278 ..Self::default()
279 }
280 }
281 pub(crate) fn validate(&self) -> crate::Result<()> {
282 anyhow::ensure!(
283 !self.backend.candidates().is_empty(),
284 "At least one backend is required"
285 );
286 anyhow::ensure!(
287 self.intra_threads <= i32::MAX as usize && self.inter_threads <= i32::MAX as usize,
288 "Thread counts must fit in a signed 32-bit integer"
289 );
290 anyhow::ensure!(
291 self.dimension_overrides
292 .values()
293 .all(|&d| d > 0 && i64::try_from(d).is_ok()),
294 "Invalid dimension override"
295 );
296 Ok(())
297 }
298}
299
300impl TensorRtOptions {
301 /// Validate only when TensorRT is selected, so unrelated settings cannot block CPU/CUDA.
302 pub(crate) fn validate(&self) -> crate::Result<()> {
303 anyhow::ensure!(self.device_id >= 0, "Negative TensorRT device id");
304 anyhow::ensure!(
305 self.builder_optimization_level <= 5,
306 "TensorRT build level must be 0..=5"
307 );
308 anyhow::ensure!(self.workspace_bytes > 0, "Workspace must be positive");
309 anyhow::ensure!(
310 self.auxiliary_streams >= -1,
311 "TensorRT auxiliary streams must be -1 or nonnegative"
312 );
313 let profiles = [&self.min_shapes, &self.opt_shapes, &self.max_shapes];
314 anyhow::ensure!(
315 profiles.iter().all(|x| x.is_some()) || profiles.iter().all(|x| x.is_none()),
316 "Set all three TensorRT profiles or none"
317 );
318 Ok(())
319 }
320}