onnx_runtime_eager/
domain.rs

1//! The domain registry: standard + custom ONNX operator domains and their
2//! default opset versions (`docs/EAGER.md` §6).
3//!
4//! Each domain carries a `default_opset`. Unlike the design's `DomainInfo`
5//! (`docs/EAGER.md` §6.4), which also nests a per-domain `KernelRegistry`, the
6//! kernels here are **not** owned by the domain registry: they live in the
7//! execution provider's [`OpRegistry`](onnx_runtime_ep_api::OpRegistry), keyed
8//! on `(op_type, domain, opset)`. Reusing the EP registry keeps a single source
9//! of truth for op support and avoids duplicating the kernel abstraction. The
10//! per-domain kernel registry and the domain-namespace handle
11//! (`ms.dispatch(...)`, §6.2) are DEFERRED.
12
13use std::collections::HashMap;
14
15use crate::opset::{LATEST_ONNX_OPSET, resolve_opset};
16
17/// Metadata for one registered operator domain.
18#[derive(Clone, Debug)]
19pub struct DomainInfo {
20    /// The domain name (`""` == the default ONNX domain).
21    pub name: String,
22    /// The default opset version for ops in this domain.
23    pub default_opset: u64,
24}
25
26/// Registry of known operator domains and their default opsets (`docs/EAGER.md`
27/// §6.4).
28#[derive(Debug)]
29pub struct DomainRegistry {
30    domains: HashMap<String, DomainInfo>,
31}
32
33impl DomainRegistry {
34    /// A registry pre-populated with the standard and common contrib domains
35    /// (`docs/EAGER.md` §6.1): the default ONNX domain at [`LATEST_ONNX_OPSET`],
36    /// `ai.onnx.ml` at 3, and `com.microsoft` at 1.
37    pub fn new() -> Self {
38        let mut reg = Self {
39            domains: HashMap::new(),
40        };
41        reg.register("", LATEST_ONNX_OPSET);
42        reg.register("ai.onnx.ml", 3);
43        reg.register("com.microsoft", 1);
44        reg
45    }
46
47    /// Register (or update) a domain with a default opset (`docs/EAGER.md` §6.1
48    /// `register_domain`).
49    pub fn register(&mut self, domain: &str, default_opset: u64) {
50        self.domains.insert(
51            domain.to_string(),
52            DomainInfo {
53                name: domain.to_string(),
54                default_opset,
55            },
56        );
57    }
58
59    /// The registered default opset for `domain`, or [`LATEST_ONNX_OPSET`] for
60    /// an unregistered domain (`docs/EAGER.md` §6.4 `resolve_opset`).
61    pub fn default_opset(&self, domain: &str) -> u64 {
62        self.domains
63            .get(domain)
64            .map(|d| d.default_opset)
65            .unwrap_or(LATEST_ONNX_OPSET)
66    }
67
68    /// Resolve the effective opset for a dispatch: an explicit per-call value
69    /// wins, otherwise the domain's registered default (`docs/EAGER.md` §5.2).
70    pub fn resolve_opset(&self, domain: &str, explicit: Option<u64>) -> u64 {
71        resolve_opset(self.default_opset(domain), explicit)
72    }
73
74    /// Whether `domain` is registered.
75    pub fn contains(&self, domain: &str) -> bool {
76        self.domains.contains_key(domain)
77    }
78
79    /// All registered domains and their default opsets (`docs/EAGER.md` §6.1
80    /// `domains()`).
81    pub fn domains(&self) -> Vec<(String, u64)> {
82        let mut out: Vec<(String, u64)> = self
83            .domains
84            .values()
85            .map(|d| (d.name.clone(), d.default_opset))
86            .collect();
87        out.sort_by(|a, b| a.0.cmp(&b.0));
88        out
89    }
90}
91
92impl Default for DomainRegistry {
93    fn default() -> Self {
94        Self::new()
95    }
96}
97
98#[cfg(test)]
99mod tests {
100    use super::*;
101
102    #[test]
103    fn defaults_are_registered() {
104        let reg = DomainRegistry::new();
105        assert_eq!(reg.default_opset(""), LATEST_ONNX_OPSET);
106        assert_eq!(reg.default_opset("ai.onnx.ml"), 3);
107        assert_eq!(reg.default_opset("com.microsoft"), 1);
108    }
109
110    #[test]
111    fn unregistered_domain_defaults_to_latest() {
112        let reg = DomainRegistry::new();
113        assert!(!reg.contains("com.acme"));
114        assert_eq!(reg.default_opset("com.acme"), LATEST_ONNX_OPSET);
115    }
116
117    #[test]
118    fn explicit_opset_overrides_domain_default() {
119        let reg = DomainRegistry::new();
120        assert_eq!(reg.resolve_opset("com.microsoft", Some(2)), 2);
121        assert_eq!(reg.resolve_opset("com.microsoft", None), 1);
122    }
123
124    #[test]
125    fn custom_domain_registration() {
126        let mut reg = DomainRegistry::new();
127        reg.register("com.acme", 5);
128        assert_eq!(reg.default_opset("com.acme"), 5);
129        assert!(reg.domains().contains(&("com.acme".to_string(), 5)));
130    }
131}
onnx_runtime_eager/domain.rs

onnx_runtime_eager/
domain.rs
