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::Arc;
42use std::sync::{Mutex, OnceLock, RwLock};
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::{LATEST_ONNX_OPSET, resolve_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
134            .lock()
135            .expect("kernel cache lock poisoned")
136            .stats()
137    }
138
139    /// Resolve the single target device from `inputs`; mixed devices are an
140    /// error (`docs/EAGER.md` §1.6 / §4.3, Design Decision).
141    pub(crate) fn resolve_device(&self, inputs: &[&Tensor]) -> Result<DeviceId> {
142        let devices: Vec<DeviceId> = inputs.iter().map(|t| t.device()).collect();
143        resolve_device_from_ids(&devices, self.default_device)
144    }
145
146    /// Get the EP (as an owned `Arc` clone) that owns `device`.
147    pub(crate) fn ep_for_device(&self, device: DeviceId) -> Result<Arc<dyn ExecutionProvider>> {
148        self.eps
149            .iter()
150            .find(|ep| ep.device_id() == device)
151            .cloned()
152            .ok_or(EagerError::NoEpForDevice(device))
153    }
154}
155
156/// Resolve a single device from a list of input devices (`docs/EAGER.md` §4.3):
157/// no inputs → `default`; all on one device → that device; otherwise a
158/// [`EagerError::MixedDeviceInputs`] error.
159///
160/// Split out from [`EagerContext::resolve_device`] so the mixed-device policy is
161/// unit-testable with fabricated [`DeviceId`]s even on a CPU-only build.
162pub(crate) fn resolve_device_from_ids(devices: &[DeviceId], default: DeviceId) -> Result<DeviceId> {
163    let mut unique: Vec<DeviceId> = Vec::new();
164    for &d in devices {
165        if !unique.contains(&d) {
166            unique.push(d);
167        }
168    }
169    match unique.len() {
170        0 => Ok(default),
171        1 => Ok(unique[0]),
172        _ => Err(EagerError::MixedDeviceInputs {
173            devices: unique,
174            hint: "use .to(device) to move all inputs to the same device".to_string(),
175        }),
176    }
177}
178
179static GLOBAL_CONTEXT: OnceLock<EagerContext> = OnceLock::new();
180
181/// Get or lazily initialize the process-global eager context (`docs/EAGER.md`
182/// §10.2).
183pub fn global_context() -> &'static EagerContext {
184    GLOBAL_CONTEXT.get_or_init(|| EagerContext::new().expect("failed to initialize eager context"))
185}
186
187#[cfg(test)]
188mod tests {
189    use super::*;
190
191    fn _assert_send_sync<T: Send + Sync>() {}
192
193    #[test]
194    fn context_is_send_sync() {
195        // Required for the `OnceLock<EagerContext>` global singleton.
196        _assert_send_sync::<EagerContext>();
197    }
198
199    #[test]
200    fn new_context_detects_cpu() {
201        let ctx = EagerContext::new().unwrap();
202        assert_eq!(ctx.default_device(), DeviceId::cpu());
203        assert!(ctx.domains().iter().any(|(d, _)| d.is_empty()));
204    }
205
206    #[test]
207    fn resolve_device_same_device_ok() {
208        let ctx = EagerContext::new().unwrap();
209        let a = Tensor::from_raw_in(
210            ctx.ep_for_device(DeviceId::cpu()).unwrap(),
211            onnx_runtime_ir::DataType::Float32,
212            vec![2],
213            &1.0f32.to_le_bytes().repeat(2),
214        )
215        .unwrap();
216        let b = a.clone();
217        assert_eq!(ctx.resolve_device(&[&a, &b]).unwrap(), DeviceId::cpu());
218    }
219
220    #[test]
221    fn resolve_device_no_inputs_uses_default() {
222        let ctx = EagerContext::new().unwrap();
223        assert_eq!(ctx.resolve_device(&[]).unwrap(), ctx.default_device());
224    }
225
226    #[test]
227    fn resolve_device_mixed_is_error() {
228        // Fabricated CPU + CUDA device ids exercise the mixed-device guard even
229        // though only a CPU EP exists (a real cross-device tensor is infeasible
230        // CPU-only), per docs/EAGER.md §1.6.
231        use onnx_runtime_ir::DeviceType;
232        let cpu = DeviceId::cpu();
233        let cuda = DeviceId::new(DeviceType::Cuda, 0);
234        assert!(matches!(
235            resolve_device_from_ids(&[cpu, cuda], cpu),
236            Err(EagerError::MixedDeviceInputs { .. })
237        ));
238        // Same-device and empty cases resolve fine.
239        assert_eq!(resolve_device_from_ids(&[cpu, cpu], cpu).unwrap(), cpu);
240        assert_eq!(resolve_device_from_ids(&[], cpu).unwrap(), cpu);
241    }
242}
onnx_runtime_eager/lib.rs

onnx_runtime_eager/
lib.rs
