onnx_runtime_eager/dispatch.rs
1//! The core single-op dispatch flow (`docs/EAGER.md` §10.1), reconciled to the
2//! real runtime APIs.
3//!
4//! ## Design-vs-real-API reconciliation
5//!
6//! The design pseudocode (`docs/EAGER.md` §10.1) calls
7//! `registry.lookup(op_type, domain, opset)` then `factory.create(attrs)`. The
8//! real APIs are:
9//!
10//! * [`OpRegistry::lookup`](onnx_runtime_ep_api::OpRegistry::lookup) /
11//! [`KernelFactory::create`](onnx_runtime_ep_api::KernelFactory) take a
12//! [`Node`] and `input_shapes`, and the CPU EP already wraps both behind
13//! [`ExecutionProvider::get_kernel`](onnx_runtime_ep_api::ExecutionProvider::get_kernel)
14//! `(node, shapes, opset)`. Eager dispatch therefore builds an **ephemeral
15//! single [`Node`]** (op_type, domain, attributes, and placeholder
16//! input/output value slots) and calls `get_kernel`, mapping the EP's
17//! `NoEpForOp` into [`EagerError::NoKernel`]. This mirrors how
18//! `onnx-runtime-session/src/executor.rs` drives kernels.
19//! * Output shapes come from
20//! [`InferenceRegistry::infer_node`](onnx_runtime_shape_inference::InferenceRegistry)
21//! (§9), fed the same ephemeral node plus per-input [`NodeIo`] built from the
22//! concrete input shapes/dtypes. Because eager inputs have fully static
23//! shapes, the inferred [`DimExpr`]s resolve to constants; a symbolic or
24//! missing output is a [`EagerError::ShapeInference`] error (the kernel-
25//! provided fallback of §9.2 is DEFERRED).
26
27use std::collections::HashMap;
28use std::ffi::c_void;
29
30use onnx_runtime_ep_api::{DevicePtr, DevicePtrMut, EpError, TensorMut, TensorView};
31use onnx_runtime_ir::{Attribute, DataType, Node, NodeId, ValueId, compute_contiguous_strides};
32use onnx_runtime_shape_inference::{DimExpr, MergePolicy, NodeIo, SymbolInterner, TypeInfo};
33
34use crate::EagerContext;
35use crate::cache::KernelCacheKey;
36use crate::error::{EagerError, Result};
37use crate::tensor::Tensor;
38
39/// Build the ephemeral [`Node`] that feeds both `get_kernel` and shape
40/// inference. Input/output value ids are positional placeholders — no real
41/// graph exists in eager mode.
42///
43/// DEFERRED (EAGER.md §9): multi-output arity. Phase-1 dispatch materialises a
44/// single output slot, which covers every Phase-1 op. Multi-output ops (e.g.
45/// `TopK`, `Split`) need the true output count here (from an op schema or an
46/// explicit argument) before their extra outputs are inferred/allocated.
47fn ephemeral_node(
48 op_type: &str,
49 domain: &str,
50 attrs: &HashMap<String, Attribute>,
51 num_inputs: usize,
52 num_outputs: usize,
53) -> Node {
54 let inputs: Vec<Option<ValueId>> = (0..num_inputs).map(|i| Some(ValueId(i as u32))).collect();
55 let outputs: Vec<ValueId> = (0..num_outputs)
56 .map(|i| ValueId((num_inputs + i) as u32))
57 .collect();
58 let mut node = Node::new(NodeId(0), op_type, inputs, outputs);
59 node.domain = domain.to_string();
60 node.attributes = attrs.clone();
61 node
62}
63
64impl EagerContext {
65 /// Dispatch a single ONNX op to a kernel and return its outputs
66 /// (`docs/EAGER.md` §10.1). The 7-step flow, reconciled to the real APIs:
67 ///
68 /// 1. resolve the effective opset (explicit per-call value > domain default),
69 /// 2. resolve the target device from the inputs (mixed devices = error),
70 /// 3. build the compiled-kernel cache key,
71 /// 4. get-or-compile the kernel via the device's EP (missing kernel =
72 /// [`EagerError::NoKernel`]),
73 /// 5. infer output shapes/dtypes for the single op,
74 /// 6. allocate the output tensors on the device,
75 /// 7. build zero-copy views and execute.
76 pub fn dispatch(
77 &self,
78 op_type: &str,
79 domain: &str,
80 inputs: &[&Tensor],
81 attrs: &HashMap<String, Attribute>,
82 explicit_opset: Option<u64>,
83 ) -> Result<Vec<Tensor>> {
84 // 1. Effective opset for this domain.
85 let opset = self
86 .domains
87 .read()
88 .expect("domain registry lock poisoned")
89 .resolve_opset(domain, explicit_opset);
90
91 // 2. Target device (mixed-device inputs are rejected, §1.6).
92 let device = self.resolve_device(inputs)?;
93 let ep = self.ep_for_device(device)?;
94
95 // Concrete input shapes, reused for the kernel, cache key, and views.
96 let input_shapes: Vec<Vec<usize>> = inputs.iter().map(|t| t.shape().to_vec()).collect();
97
98 // DEFERRED (EAGER.md §9): single output slot only (see `ephemeral_node`).
99 let node = ephemeral_node(op_type, domain, attrs, inputs.len(), 1);
100
101 // 3 + 4. Cache key, then get-or-compile the kernel through the EP.
102 let cache_key = KernelCacheKey {
103 op_type: op_type.to_string(),
104 domain: domain.to_string(),
105 opset,
106 input_shapes: input_shapes.clone(),
107 input_dtypes: inputs.iter().map(|t| t.dtype()).collect(),
108 device,
109 };
110 let kernel = self
111 .cache
112 .lock()
113 .expect("kernel cache lock poisoned")
114 .get_or_create(
115 cache_key,
116 || -> Result<Box<dyn onnx_runtime_ep_api::Kernel>> {
117 ep.get_kernel(&node, &input_shapes, opset)
118 .map_err(|e| match e {
119 EpError::NoEpForOp { .. } => EagerError::NoKernel {
120 op_type: op_type.to_string(),
121 domain: domain.to_string(),
122 device,
123 },
124 other => EagerError::Kernel(other),
125 })
126 },
127 )?;
128
129 // 5. Infer output shapes/dtypes for the single op.
130 let output_meta = self.infer_output_meta(&node, op_type, domain, opset, inputs)?;
131
132 // 6. Allocate output tensors on the target device.
133 let mut outputs: Vec<Tensor> = Vec::with_capacity(output_meta.len());
134 for (dtype, shape) in &output_meta {
135 outputs.push(Tensor::zeros_in(ep.clone(), *dtype, shape.clone())?);
136 }
137
138 // 7. Build zero-copy views over the raw device buffers and execute.
139 // Stride/shape holders must outlive the views that borrow them.
140 let in_strides: Vec<Vec<i64>> = input_shapes
141 .iter()
142 .map(|s| compute_contiguous_strides(s))
143 .collect();
144 let out_shapes: Vec<Vec<usize>> = output_meta.iter().map(|(_, s)| s.clone()).collect();
145 let out_strides: Vec<Vec<i64>> = out_shapes
146 .iter()
147 .map(|s| compute_contiguous_strides(s))
148 .collect();
149
150 let in_ptrs: Vec<*const c_void> = inputs.iter().map(|t| t.device_ptr()).collect();
151 let out_ptrs: Vec<*mut c_void> = outputs.iter_mut().map(|t| t.device_ptr_mut()).collect();
152
153 let input_views: Vec<TensorView> = (0..inputs.len())
154 .map(|i| {
155 TensorView::new(
156 DevicePtr(in_ptrs[i]),
157 inputs[i].dtype(),
158 &input_shapes[i],
159 &in_strides[i],
160 device,
161 )
162 })
163 .collect();
164 let mut output_views: Vec<TensorMut> = (0..outputs.len())
165 .map(|i| {
166 TensorMut::new(
167 DevicePtrMut(out_ptrs[i]),
168 output_meta[i].0,
169 &out_shapes[i],
170 &out_strides[i],
171 device,
172 )
173 })
174 .collect();
175
176 {
177 let guard = kernel.lock().expect("cached kernel mutex poisoned");
178 guard.execute(&input_views, &mut output_views)?;
179 }
180 // Drop the views (they borrow the ptr/stride holders) before the owned
181 // output tensors leave the function.
182 drop(output_views);
183 drop(input_views);
184
185 Ok(outputs)
186 }
187
188 /// Per-op output shape/dtype inference (`docs/EAGER.md` §9), driven by the
189 /// shared [`InferenceRegistry`](onnx_runtime_shape_inference::InferenceRegistry).
190 /// Returns one `(dtype, shape)` per output slot.
191 fn infer_output_meta(
192 &self,
193 node: &Node,
194 op_type: &str,
195 domain: &str,
196 opset: u64,
197 inputs: &[&Tensor],
198 ) -> Result<Vec<(DataType, Vec<usize>)>> {
199 let input_ios: Vec<NodeIo> = inputs
200 .iter()
201 .map(|t| {
202 let shape: Vec<DimExpr> = t
203 .shape()
204 .iter()
205 .map(|&d| DimExpr::constant(d as i64))
206 .collect();
207 NodeIo::typed(TypeInfo::new(t.dtype(), shape))
208 })
209 .collect();
210
211 // The registry keys inference on `(domain, op)` and reads the opset for
212 // the op's domain from `opset_imports`; supply exactly that entry.
213 let mut opset_imports: HashMap<String, u64> = HashMap::new();
214 opset_imports.insert(domain.to_string(), opset);
215
216 let mut interner = SymbolInterner::new(0);
217 let out_ios = self.inference.infer_node(
218 node,
219 &opset_imports,
220 input_ios,
221 MergePolicy::Permissive,
222 &mut interner,
223 )?;
224
225 let mut meta = Vec::with_capacity(out_ios.len());
226 for (i, io) in out_ios.iter().enumerate() {
227 let type_info = io
228 .type_info
229 .as_ref()
230 .ok_or_else(|| EagerError::ShapeInference {
231 op_type: op_type.to_string(),
232 domain: domain.to_string(),
233 reason: format!(
234 "no shape-inference rule resolved output {i} \
235 (kernel-provided fallback is DEFERRED, EAGER.md §9.2)"
236 ),
237 })?;
238 let mut dims = Vec::with_capacity(type_info.shape.len());
239 for (axis, d) in type_info.shape.iter().enumerate() {
240 let c = d.as_const().ok_or_else(|| EagerError::ShapeInference {
241 op_type: op_type.to_string(),
242 domain: domain.to_string(),
243 reason: format!("output {i} axis {axis} is symbolic, not allocatable"),
244 })?;
245 if c < 0 {
246 return Err(EagerError::ShapeInference {
247 op_type: op_type.to_string(),
248 domain: domain.to_string(),
249 reason: format!("output {i} axis {axis} has a negative extent {c}"),
250 });
251 }
252 dims.push(c as usize);
253 }
254 meta.push((type_info.dtype, dims));
255 }
256 Ok(meta)
257 }
258}