onnx_runtime_ep_api/
registry.rs

1//! Op → kernel-factory registry and the EP registry (§4.3, §4.6).
2
3use std::collections::HashMap;
4use std::path::Path;
5
6use onnx_runtime_ir::{Node, Shape, TensorLayout};
7
8use crate::error::Result;
9use crate::kernel::{Kernel, KernelMatch};
10use crate::provider::{EpConfig, EpId, ExecutionProvider};
11
12/// Registry key: an operator identity plus the opset version it was introduced.
13#[derive(Clone, PartialEq, Eq, Hash, Debug)]
14pub struct OpKey {
15    pub op_type: String,
16    pub domain: String,
17    pub since_version: u64,
18}
19
20impl OpKey {
21    pub fn new(op_type: impl Into<String>, domain: impl Into<String>, since_version: u64) -> Self {
22        Self {
23            op_type: op_type.into(),
24            domain: domain.into(),
25            since_version,
26        }
27    }
28}
29
30/// Normalise the default ONNX domain: the empty string and `"ai.onnx"` name the
31/// same (standard) domain. Contrib domains (e.g. `"com.microsoft"`) are left
32/// untouched. Keeps dispatch keyed on `(op_type, domain)` model-agnostically.
33fn norm_domain(domain: &str) -> &str {
34    if domain == "ai.onnx" { "" } else { domain }
35}
36
37/// Creates kernels for a specific op.
38pub trait KernelFactory: Send + Sync {
39    fn create(&self, node: &Node, input_shapes: &[Vec<usize>]) -> Result<Box<dyn Kernel>>;
40}
41
42/// Maps `(op_type, domain, opset)` → kernel factory (§4.3).
43#[derive(Default)]
44pub struct OpRegistry {
45    entries: HashMap<OpKey, Box<dyn KernelFactory>>,
46}
47
48impl OpRegistry {
49    pub fn new() -> Self {
50        Self::default()
51    }
52
53    /// Register a factory under `key`.
54    pub fn register(&mut self, key: OpKey, factory: Box<dyn KernelFactory>) {
55        self.entries.insert(key, factory);
56    }
57
58    /// Look up the best matching factory: the highest `since_version` that is
59    /// `<= opset` for the given `(op_type, domain)`.
60    pub fn lookup(&self, op_type: &str, domain: &str, opset: u64) -> Option<&dyn KernelFactory> {
61        let domain = norm_domain(domain);
62        self.entries
63            .iter()
64            .filter(|(k, _)| {
65                k.op_type == op_type && k.domain == domain && k.since_version <= opset
66            })
67            .max_by_key(|(k, _)| k.since_version)
68            .map(|(_, f)| f.as_ref())
69    }
70
71    /// Whether any factory is registered for `(op_type, domain)`, ignoring the
72    /// opset version. Used by providers to answer "is this an op/domain we
73    /// support?" without a concrete opset in hand (`supports_op`). Keeps the
74    /// support decision keyed on `(op_type, domain)` via the registry rather
75    /// than a hardcoded op/domain whitelist.
76    pub fn supports(&self, op_type: &str, domain: &str) -> bool {
77        let domain = norm_domain(domain);
78        self.entries
79            .keys()
80            .any(|k| k.op_type == op_type && k.domain == domain)
81    }
82
83    /// Number of registered entries.
84    pub fn len(&self) -> usize {
85        self.entries.len()
86    }
87
88    /// Whether the registry is empty.
89    pub fn is_empty(&self) -> bool {
90        self.entries.is_empty()
91    }
92}
93
94/// Ordered set of execution providers with a priority list (§4.6).
95#[derive(Default)]
96pub struct EpRegistry {
97    eps: Vec<Box<dyn ExecutionProvider>>,
98    /// Priority order as indices into `eps` (front = highest priority).
99    priority: Vec<EpId>,
100}
101
102impl EpRegistry {
103    pub fn new() -> Self {
104        Self::default()
105    }
106
107    /// Register an EP, returning its [`EpId`]. Appended to the priority list.
108    pub fn register(&mut self, ep: Box<dyn ExecutionProvider>) -> EpId {
109        let id = EpId(self.eps.len() as u32);
110        self.eps.push(ep);
111        self.priority.push(id);
112        id
113    }
114
115    /// Load a legacy ORT plugin EP from a shared library (Phase 2).
116    pub fn load_legacy(&mut self, path: &Path, config: &EpConfig) -> Result<EpId> {
117        let _ = (path, config);
118        todo!("ort2-ep-api Phase 2: dlopen legacy ORT plugin EP and adapt its vtable")
119    }
120
121    /// Override the priority order.
122    pub fn set_priority(&mut self, order: Vec<EpId>) {
123        self.priority = order;
124    }
125
126    /// Borrow an EP by id.
127    pub fn get(&self, id: EpId) -> Option<&dyn ExecutionProvider> {
128        self.eps.get(id.0 as usize).map(|b| b.as_ref())
129    }
130
131    /// The priority order.
132    pub fn priority(&self) -> &[EpId] {
133        &self.priority
134    }
135
136    /// All EPs (in priority order) that can handle `op`, with their match info.
137    pub fn candidates_for_op(
138        &self,
139        op: &Node,
140        shapes: &[Shape],
141        layouts: &[TensorLayout],
142    ) -> Vec<(EpId, KernelMatch)> {
143        let mut out = Vec::new();
144        for &id in &self.priority {
145            if let Some(ep) = self.get(id) {
146                let m = ep.supports_op(op, shapes, layouts);
147                if m.is_supported() {
148                    out.push((id, m));
149                }
150            }
151        }
152        out
153    }
154}
onnx_runtime_ep_api/registry.rs

onnx_runtime_ep_api/
registry.rs
