onnx_runtime_eager/
lib.rs

1//! # `onnx-runtime-eager`
2//!
3//! Eager (single-op) execution for the ORT 2.0 runtime (`docs/EAGER.md`).
4//! Dispatch individual ONNX ops to execution-provider kernels **without building
5//! a graph** — a PyTorch-style experience with ONNX op semantics.
6//!
7//! This crate is the pure-Rust core of the design. It reuses the existing
8//! runtime abstractions rather than inventing parallel ones:
9//!
10//! * [`ExecutionProvider`](onnx_runtime_ep_api::ExecutionProvider) /
11//!   [`Kernel`](onnx_runtime_ep_api::Kernel) /
12//!   [`OpRegistry`](onnx_runtime_ep_api::OpRegistry) from `onnx-runtime-ep-api`,
13//! * the populated CPU registry from `onnx-runtime-ep-cpu`,
14//! * per-op shape/dtype inference from `onnx-runtime-shape-inference`,
15//! * the IR vocabulary ([`Node`](onnx_runtime_ir::Node),
16//!   [`DataType`](onnx_runtime_ir::DataType),
17//!   [`Attribute`](onnx_runtime_ir::Attribute), …) from `onnx-runtime-ir`.
18//!
19//! ## Phase-1 scope (`docs/EAGER.md`)
20//!
21//! Implemented: [`EagerContext`], [`EagerContext::dispatch`] (single-op CPU
22//! dispatch), the [`DomainRegistry`](domain::DomainRegistry), opset resolution
23//! ([`opset`]), and the compiled-[`KernelCache`](cache::KernelCache).
24//!
25//! DEFERRED (each is flagged at its hook point):
26//! * PyO3 / Python bindings (§11), including the `Tensor` single→tensor /
27//!   multi→tuple sugar and `ops.*` typed wrappers (§4.1).
28//! * Subgraph ops `If`/`Loop`/`Scan` (§7).
29//! * The opset context-manager / `nxrt.device()` context (§5.3, §2.2).
30//! * GPU/CUDA EP dispatch and implicit cross-device transfer.
31//! * Kernel-provided shape-inference fallback (§9.2).
32//! * DLPack / numpy interop (§3).
33
34mod cache;
35mod dispatch;
36mod domain;
37mod error;
38mod opset;
39mod tensor;
40
41use std::sync::{Mutex, OnceLock, RwLock};
42use std::sync::Arc;
43
44use onnx_runtime_ep_api::{EpConfig, ExecutionProvider};
45use onnx_runtime_ep_cpu::CpuExecutionProvider;
46use onnx_runtime_ir::DeviceId;
47use onnx_runtime_shape_inference::InferenceRegistry;
48
49pub use cache::{CacheStats, KernelCache, KernelCacheKey};
50pub use domain::{DomainInfo, DomainRegistry};
51pub use error::{EagerError, Result};
52pub use opset::{resolve_opset, LATEST_ONNX_OPSET};
53pub use tensor::Tensor;
54
55/// Default compiled-kernel cache capacity (`docs/EAGER.md` §13: per-process LRU,
56/// 4096).
57const DEFAULT_CACHE_CAPACITY: usize = 4096;
58
59/// The process-wide eager execution context (`docs/EAGER.md` §10.1).
60///
61/// Manages the available EPs, the domain registry, per-op shape inference, the
62/// compiled-kernel cache, and the default device. It is thread-safe via internal
63/// locking so it can back the [`global_context`] singleton.
64pub struct EagerContext {
65    /// Available EPs, auto-detected at initialization. CPU only in Phase-1;
66    /// GPU EPs are added here as a *registration*, not a rewrite.
67    pub(crate) eps: Vec<Arc<dyn ExecutionProvider>>,
68    /// Compiled-kernel cache (keyed by op + shapes + dtypes + device).
69    pub(crate) cache: Mutex<KernelCache>,
70    /// Domain registry (standard + custom).
71    pub(crate) domains: RwLock<DomainRegistry>,
72    /// Per-op output shape/dtype inference (§9).
73    pub(crate) inference: InferenceRegistry,
74    /// Default device for ops with no inputs (`docs/EAGER.md` §4.3 case 3).
75    pub(crate) default_device: DeviceId,
76}
77
78/// Auto-detect the available execution providers (`docs/EAGER.md` §10.1
79/// `detect_available_eps`).
80///
81/// Phase-1: the always-available CPU EP only. Adding a GPU EP later is a matter
82/// of pushing another initialized provider here — the dispatch path selects an
83/// EP generically by device, with no hardcoded backend/vendor name.
84fn detect_available_eps() -> Result<Vec<Arc<dyn ExecutionProvider>>> {
85    // DEFERRED (EAGER.md §4.3 / §2.2): GPU/CUDA EP discovery + auto-transfer.
86    let mut cpu = CpuExecutionProvider::new();
87    cpu.initialize(&EpConfig::default())?;
88    Ok(vec![Arc::new(cpu) as Arc<dyn ExecutionProvider>])
89}
90
91impl EagerContext {
92    /// Initialize a context with auto-detected devices (`docs/EAGER.md` §10.1).
93    pub fn new() -> Result<Self> {
94        let eps = detect_available_eps()?;
95        let default_device = eps
96            .first()
97            .map(|ep| ep.device_id())
98            .unwrap_or_else(DeviceId::cpu);
99        Ok(Self {
100            eps,
101            cache: Mutex::new(KernelCache::new(DEFAULT_CACHE_CAPACITY)),
102            domains: RwLock::new(DomainRegistry::new()),
103            inference: InferenceRegistry::default_registry(),
104            default_device,
105        })
106    }
107
108    /// The default device for input-less ops (`docs/EAGER.md` §4.3).
109    pub fn default_device(&self) -> DeviceId {
110        self.default_device
111    }
112
113    /// Register (or update) a custom domain's default opset (`docs/EAGER.md`
114    /// §6.1 `register_domain`).
115    pub fn register_domain(&self, domain: &str, default_opset: u64) {
116        self.domains
117            .write()
118            .expect("domain registry lock poisoned")
119            .register(domain, default_opset);
120    }
121
122    /// All registered domains and their default opsets (`docs/EAGER.md` §6.1
123    /// `domains()`).
124    pub fn domains(&self) -> Vec<(String, u64)> {
125        self.domains
126            .read()
127            .expect("domain registry lock poisoned")
128            .domains()
129    }
130
131    /// Current kernel-cache statistics (entries / hits / misses).
132    pub fn cache_stats(&self) -> CacheStats {
133        self.cache.lock().expect("kernel cache lock poisoned").stats()
134    }
135
136    /// Resolve the single target device from `inputs`; mixed devices are an
137    /// error (`docs/EAGER.md` §1.6 / §4.3, Design Decision).
138    pub(crate) fn resolve_device(&self, inputs: &[&Tensor]) -> Result<DeviceId> {
139        let devices: Vec<DeviceId> = inputs.iter().map(|t| t.device()).collect();
140        resolve_device_from_ids(&devices, self.default_device)
141    }
142
143    /// Get the EP (as an owned `Arc` clone) that owns `device`.
144    pub(crate) fn ep_for_device(&self, device: DeviceId) -> Result<Arc<dyn ExecutionProvider>> {
145        self.eps
146            .iter()
147            .find(|ep| ep.device_id() == device)
148            .cloned()
149            .ok_or(EagerError::NoEpForDevice(device))
150    }
151}
152
153/// Resolve a single device from a list of input devices (`docs/EAGER.md` §4.3):
154/// no inputs → `default`; all on one device → that device; otherwise a
155/// [`EagerError::MixedDeviceInputs`] error.
156///
157/// Split out from [`EagerContext::resolve_device`] so the mixed-device policy is
158/// unit-testable with fabricated [`DeviceId`]s even on a CPU-only build.
159pub(crate) fn resolve_device_from_ids(devices: &[DeviceId], default: DeviceId) -> Result<DeviceId> {
160    let mut unique: Vec<DeviceId> = Vec::new();
161    for &d in devices {
162        if !unique.contains(&d) {
163            unique.push(d);
164        }
165    }
166    match unique.len() {
167        0 => Ok(default),
168        1 => Ok(unique[0]),
169        _ => Err(EagerError::MixedDeviceInputs {
170            devices: unique,
171            hint: "use .to(device) to move all inputs to the same device".to_string(),
172        }),
173    }
174}
175
176static GLOBAL_CONTEXT: OnceLock<EagerContext> = OnceLock::new();
177
178/// Get or lazily initialize the process-global eager context (`docs/EAGER.md`
179/// §10.2).
180pub fn global_context() -> &'static EagerContext {
181    GLOBAL_CONTEXT.get_or_init(|| EagerContext::new().expect("failed to initialize eager context"))
182}
183
184#[cfg(test)]
185mod tests {
186    use super::*;
187
188    fn _assert_send_sync<T: Send + Sync>() {}
189
190    #[test]
191    fn context_is_send_sync() {
192        // Required for the `OnceLock<EagerContext>` global singleton.
193        _assert_send_sync::<EagerContext>();
194    }
195
196    #[test]
197    fn new_context_detects_cpu() {
198        let ctx = EagerContext::new().unwrap();
199        assert_eq!(ctx.default_device(), DeviceId::cpu());
200        assert!(ctx.domains().iter().any(|(d, _)| d.is_empty()));
201    }
202
203    #[test]
204    fn resolve_device_same_device_ok() {
205        let ctx = EagerContext::new().unwrap();
206        let a = Tensor::from_raw_in(
207            ctx.ep_for_device(DeviceId::cpu()).unwrap(),
208            onnx_runtime_ir::DataType::Float32,
209            vec![2],
210            &1.0f32.to_le_bytes().repeat(2),
211        )
212        .unwrap();
213        let b = a.clone();
214        assert_eq!(ctx.resolve_device(&[&a, &b]).unwrap(), DeviceId::cpu());
215    }
216
217    #[test]
218    fn resolve_device_no_inputs_uses_default() {
219        let ctx = EagerContext::new().unwrap();
220        assert_eq!(ctx.resolve_device(&[]).unwrap(), ctx.default_device());
221    }
222
223    #[test]
224    fn resolve_device_mixed_is_error() {
225        // Fabricated CPU + CUDA device ids exercise the mixed-device guard even
226        // though only a CPU EP exists (a real cross-device tensor is infeasible
227        // CPU-only), per docs/EAGER.md §1.6.
228        use onnx_runtime_ir::DeviceType;
229        let cpu = DeviceId::cpu();
230        let cuda = DeviceId::new(DeviceType::Cuda, 0);
231        assert!(matches!(
232            resolve_device_from_ids(&[cpu, cuda], cpu),
233            Err(EagerError::MixedDeviceInputs { .. })
234        ));
235        // Same-device and empty cases resolve fine.
236        assert_eq!(resolve_device_from_ids(&[cpu, cpu], cpu).unwrap(), cpu);
237        assert_eq!(resolve_device_from_ids(&[], cpu).unwrap(), cpu);
238    }
239}
onnx_runtime_eager/lib.rs

onnx_runtime_eager/
lib.rs
