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