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::{compute_contiguous_strides, Attribute, DataType, Node, NodeId, ValueId};
32use onnx_runtime_shape_inference::{DimExpr, MergePolicy, NodeIo, SymbolInterner, TypeInfo};
33
34use crate::cache::KernelCacheKey;
35use crate::error::{EagerError, Result};
36use crate::tensor::Tensor;
37use crate::EagerContext;
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(cache_key, || -> Result<Box<dyn onnx_runtime_ep_api::Kernel>> {
115 ep.get_kernel(&node, &input_shapes, opset).map_err(|e| match e {
116 EpError::NoEpForOp { .. } => EagerError::NoKernel {
117 op_type: op_type.to_string(),
118 domain: domain.to_string(),
119 device,
120 },
121 other => EagerError::Kernel(other),
122 })
123 })?;
124
125 // 5. Infer output shapes/dtypes for the single op.
126 let output_meta = self.infer_output_meta(&node, op_type, domain, opset, inputs)?;
127
128 // 6. Allocate output tensors on the target device.
129 let mut outputs: Vec<Tensor> = Vec::with_capacity(output_meta.len());
130 for (dtype, shape) in &output_meta {
131 outputs.push(Tensor::zeros_in(ep.clone(), *dtype, shape.clone())?);
132 }
133
134 // 7. Build zero-copy views over the raw device buffers and execute.
135 // Stride/shape holders must outlive the views that borrow them.
136 let in_strides: Vec<Vec<i64>> = input_shapes
137 .iter()
138 .map(|s| compute_contiguous_strides(s))
139 .collect();
140 let out_shapes: Vec<Vec<usize>> = output_meta.iter().map(|(_, s)| s.clone()).collect();
141 let out_strides: Vec<Vec<i64>> = out_shapes
142 .iter()
143 .map(|s| compute_contiguous_strides(s))
144 .collect();
145
146 let in_ptrs: Vec<*const c_void> = inputs.iter().map(|t| t.device_ptr()).collect();
147 let out_ptrs: Vec<*mut c_void> =
148 outputs.iter_mut().map(|t| t.device_ptr_mut()).collect();
149
150 let input_views: Vec<TensorView> = (0..inputs.len())
151 .map(|i| {
152 TensorView::new(
153 DevicePtr(in_ptrs[i]),
154 inputs[i].dtype(),
155 &input_shapes[i],
156 &in_strides[i],
157 device,
158 )
159 })
160 .collect();
161 let mut output_views: Vec<TensorMut> = (0..outputs.len())
162 .map(|i| {
163 TensorMut::new(
164 DevicePtrMut(out_ptrs[i]),
165 output_meta[i].0,
166 &out_shapes[i],
167 &out_strides[i],
168 device,
169 )
170 })
171 .collect();
172
173 {
174 let guard = kernel.lock().expect("cached kernel mutex poisoned");
175 guard.execute(&input_views, &mut output_views)?;
176 }
177 // Drop the views (they borrow the ptr/stride holders) before the owned
178 // output tensors leave the function.
179 drop(output_views);
180 drop(input_views);
181
182 Ok(outputs)
183 }
184
185 /// Per-op output shape/dtype inference (`docs/EAGER.md` §9), driven by the
186 /// shared [`InferenceRegistry`](onnx_runtime_shape_inference::InferenceRegistry).
187 /// Returns one `(dtype, shape)` per output slot.
188 fn infer_output_meta(
189 &self,
190 node: &Node,
191 op_type: &str,
192 domain: &str,
193 opset: u64,
194 inputs: &[&Tensor],
195 ) -> Result<Vec<(DataType, Vec<usize>)>> {
196 let input_ios: Vec<NodeIo> = inputs
197 .iter()
198 .map(|t| {
199 let shape: Vec<DimExpr> =
200 t.shape().iter().map(|&d| DimExpr::constant(d as i64)).collect();
201 NodeIo::typed(TypeInfo::new(t.dtype(), shape))
202 })
203 .collect();
204
205 // The registry keys inference on `(domain, op)` and reads the opset for
206 // the op's domain from `opset_imports`; supply exactly that entry.
207 let mut opset_imports: HashMap<String, u64> = HashMap::new();
208 opset_imports.insert(domain.to_string(), opset);
209
210 let mut interner = SymbolInterner::new(0);
211 let out_ios = self.inference.infer_node(
212 node,
213 &opset_imports,
214 input_ios,
215 MergePolicy::Permissive,
216 &mut interner,
217 )?;
218
219 let mut meta = Vec::with_capacity(out_ios.len());
220 for (i, io) in out_ios.iter().enumerate() {
221 let type_info = io.type_info.as_ref().ok_or_else(|| EagerError::ShapeInference {
222 op_type: op_type.to_string(),
223 domain: domain.to_string(),
224 reason: format!(
225 "no shape-inference rule resolved output {i} \
226 (kernel-provided fallback is DEFERRED, EAGER.md §9.2)"
227 ),
228 })?;
229 let mut dims = Vec::with_capacity(type_info.shape.len());
230 for (axis, d) in type_info.shape.iter().enumerate() {
231 let c = d.as_const().ok_or_else(|| EagerError::ShapeInference {
232 op_type: op_type.to_string(),
233 domain: domain.to_string(),
234 reason: format!("output {i} axis {axis} is symbolic, not allocatable"),
235 })?;
236 if c < 0 {
237 return Err(EagerError::ShapeInference {
238 op_type: op_type.to_string(),
239 domain: domain.to_string(),
240 reason: format!("output {i} axis {axis} has a negative extent {c}"),
241 });
242 }
243 dims.push(c as usize);
244 }
245 meta.push((type_info.dtype, dims));
246 }
247 Ok(meta)
248 }
249}