1use std::collections::{HashMap, HashSet};
47use std::sync::Arc;
48
49use onnx_runtime_ep_api::{
50 DeviceBuffer, DevicePtr, DevicePtrMut, ExecutionProvider, KernelMatch, TensorMut, TensorView,
51};
52use onnx_runtime_ep_cpu::strided::view_in_bounds;
53use onnx_runtime_ep_cpu::CpuExecutionProvider;
54use onnx_runtime_ir::{
55 as_static_shape, compute_contiguous_strides, DataType, Dim, Graph, Node, NodeId, Shape,
56 SymbolId, TensorLayout, ValueId,
57};
58use onnx_runtime_loader::WeightStore;
59use onnx_runtime_shape_inference::{InferenceRegistry, MergePolicy};
60
61use crate::error::{Result, SessionError};
62use crate::sequence::{
63 concat_axis, split_axis, stack_new_axis, SeqTensor, SequenceValue,
64};
65use crate::tensor::{host_bytes, write_host, Tensor};
66
67#[derive(Debug)]
71pub(crate) struct NodePlan {
72 pub node_id: NodeId,
73 pub inputs: Vec<Option<ValueId>>,
79 pub outputs: Vec<ValueId>,
81 pub input_dtypes: Vec<DataType>,
84 pub output_dtypes: Vec<DataType>,
86}
87
88#[derive(Clone, PartialEq, Eq, Hash, Debug)]
95struct KernelKey {
96 node: u32,
97 shapes: Vec<Vec<usize>>,
98}
99
100#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
102pub struct CacheStats {
103 pub entries: usize,
105 pub hits: u64,
107 pub misses: u64,
109}
110
111#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
114pub struct ControlFlowStats {
115 pub subgraph_builds: u64,
117 pub subgraph_runs: u64,
119}
120
121#[derive(Default)]
123pub(crate) struct KernelCache {
124 entries: HashMap<KernelKey, Box<dyn onnx_runtime_ep_api::Kernel>>,
125 hits: u64,
126 misses: u64,
127}
128
129impl KernelCache {
130 fn stats(&self) -> CacheStats {
131 CacheStats {
132 entries: self.entries.len(),
133 hits: self.hits,
134 misses: self.misses,
135 }
136 }
137
138 fn get_or_create(
143 &mut self,
144 node_id: NodeId,
145 node: &Node,
146 input_shapes: &[Vec<usize>],
147 opset: u64,
148 ep: &CpuExecutionProvider,
149 ) -> Result<&dyn onnx_runtime_ep_api::Kernel> {
150 let key = KernelKey {
151 node: node_id.0,
152 shapes: input_shapes.to_vec(),
153 };
154 if self.entries.contains_key(&key) {
155 self.hits += 1;
156 } else {
157 let shape_dims: Vec<Shape> = input_shapes
160 .iter()
161 .map(|s| s.iter().map(|&d| Dim::Static(d)).collect())
162 .collect();
163 let layouts = vec![TensorLayout::contiguous(); input_shapes.len()];
164 if !matches!(
165 ep.supports_op(node, &shape_dims, &layouts),
166 KernelMatch::Supported { .. }
167 ) {
168 return Err(SessionError::unsupported_op(node, node_id, opset, ep.name()));
169 }
170 let kernel = ep.get_kernel(node, input_shapes, opset)?;
171 self.entries.insert(key.clone(), kernel);
172 self.misses += 1;
173 }
174 Ok(self.entries.get(&key).expect("just inserted").as_ref())
175 }
176}
177
178pub(crate) struct Executor {
181 graph: Graph,
182 weights: Arc<WeightStore>,
192 ep: Arc<CpuExecutionProvider>,
193 buffers: HashMap<ValueId, DeviceBuffer>,
197 buffer_shapes: HashMap<ValueId, Vec<usize>>,
200 value_shapes: HashMap<ValueId, Shape>,
202 value_dtypes: HashMap<ValueId, DataType>,
204 plan: Vec<NodePlan>,
206 input_index: HashMap<String, ValueId>,
208 required_inputs: Vec<ValueId>,
210 has_symbols: bool,
214 cache: KernelCache,
215 name_index: HashMap<String, ValueId>,
219 subgraph_execs: HashMap<(NodeId, String), CompiledSubgraph>,
228 control_flow_stats: ControlFlowStats,
229 views: HashMap<ValueId, ValueView>,
234 pinned: HashSet<ValueId>,
239 sequence_values: HashSet<ValueId>,
245 sequences: HashMap<ValueId, SequenceValue>,
251 seq_elem_values: HashMap<ValueId, Arc<SeqTensor>>,
258}
259
260#[derive(Clone, Debug)]
266struct ValueView {
267 source: ValueId,
268 shape: Vec<usize>,
269 strides: Vec<i64>,
270 byte_offset: usize,
271}
272
273struct InInfo {
280 present: bool,
281 dtype: DataType,
282 shape: Vec<usize>,
283 strides: Vec<i64>,
284 byte_offset: usize,
285 base_ptr: *const std::ffi::c_void,
286 root_len: usize,
288}
289
290struct CompiledSubgraph {
294 exec: Executor,
295 input_names: Vec<String>,
298 built_shapes: Vec<Vec<usize>>,
300}
301
302struct PreparedSubgraph {
306 key: (NodeId, String),
307 formal_names: Vec<String>,
308 capture_names: Vec<String>,
309 captures: HashMap<String, Tensor>,
311}
312
313fn view_bounds(
318 shape: &[usize],
319 strides: &[i64],
320 byte_offset: usize,
321 dtype: DataType,
322 buffer_len: usize,
323) -> Result<()> {
324 let esize = dtype.byte_size();
325 if esize == 0 {
326 let numel: usize = shape.iter().product();
328 let need = byte_offset + dtype.storage_bytes(numel);
329 if need > buffer_len {
330 return Err(SessionError::from(
331 onnx_runtime_ep_api::EpError::InvalidTensorView {
332 reason: format!(
333 "sub-byte view needs {need} bytes but backing allocation is {buffer_len}"
334 ),
335 },
336 ));
337 }
338 return Ok(());
339 }
340 view_in_bounds(shape, strides, byte_offset, esize, buffer_len)?;
341 Ok(())
342}
343
344fn gather_view(
351 src: &[u8],
352 shape: &[usize],
353 strides: &[i64],
354 byte_offset: usize,
355 esize: usize,
356) -> Vec<u8> {
357 let n: usize = shape.iter().product();
358 let mut out = vec![0u8; n * esize];
359 if n == 0 {
360 return out;
361 }
362 let rank = shape.len();
363 let mut idx = vec![0usize; rank];
364 let mut w = 0usize;
365 loop {
366 let mut off = byte_offset as i64;
367 for d in 0..rank {
368 off += strides[d] * idx[d] as i64 * esize as i64;
369 }
370 let s = off as usize;
371 out[w..w + esize].copy_from_slice(&src[s..s + esize]);
372 w += esize;
373 let mut carried = true;
375 for axis in (0..rank).rev() {
376 idx[axis] += 1;
377 if idx[axis] < shape[axis] {
378 carried = false;
379 break;
380 }
381 idx[axis] = 0;
382 }
383 if carried {
384 break;
385 }
386 }
387 out
388}
389
390fn checked_numel(dims: &[usize], value: impl FnOnce() -> String) -> Result<usize> {
395 let mut acc = 1usize;
396 for &d in dims {
397 acc = match acc.checked_mul(d) {
398 Some(n) => n,
399 None => {
400 return Err(SessionError::ShapeOverflow {
401 value: value(),
402 dims: dims.to_vec(),
403 })
404 }
405 };
406 }
407 Ok(acc)
408}
409
410fn checked_storage_bytes(
416 dtype: DataType,
417 numel: usize,
418 value: impl FnOnce() -> String,
419 dims: &[usize],
420) -> Result<usize> {
421 dtype
422 .checked_storage_bytes(numel)
423 .ok_or_else(|| SessionError::ShapeOverflow {
424 value: value(),
425 dims: dims.to_vec(),
426 })
427}
428
429fn effective_opset(graph: &Graph, node: &Node) -> u64 {
434 let domain = node.domain.as_str();
435 graph
436 .opset_imports
437 .get(domain)
438 .or_else(|| {
439 if domain.is_empty() {
440 graph.opset_imports.get("ai.onnx")
441 } else if domain == "ai.onnx" {
442 graph.opset_imports.get("")
443 } else {
444 None
445 }
446 })
447 .copied()
448 .unwrap_or_else(|| {
449 unreachable!(
450 "internal invariant violated: node #{} ({}::{}) has no opset import",
451 node.id.0,
452 if node.domain.is_empty() {
453 "ai.onnx"
454 } else {
455 &node.domain
456 },
457 node.op_type
458 )
459 })
460}
461
462fn substitute(shape: &Shape, bindings: &HashMap<SymbolId, usize>) -> Option<Vec<usize>> {
465 shape
466 .iter()
467 .map(|d| match d {
468 Dim::Static(n) => Some(*n),
469 Dim::Symbolic(s) => bindings.get(s).copied(),
470 })
471 .collect()
472}
473
474fn buffer_as_i64(buffer: &DeviceBuffer, dtype: DataType) -> Option<Vec<i64>> {
478 bytes_as_i64(crate::tensor::host_bytes(buffer), dtype)
479}
480
481fn bytes_as_i64(bytes: &[u8], dtype: DataType) -> Option<Vec<i64>> {
485 match dtype {
486 DataType::Int64 => Some(
487 bytes
488 .chunks_exact(8)
489 .map(|c| i64::from_le_bytes(c.try_into().unwrap()))
490 .collect(),
491 ),
492 DataType::Int32 => Some(
493 bytes
494 .chunks_exact(4)
495 .map(|c| i32::from_le_bytes(c.try_into().unwrap()) as i64)
496 .collect(),
497 ),
498 _ => None,
499 }
500}
501
502fn dynamic_output_shapes(
513 node: &Node,
514 input_shapes: &[Vec<usize>],
515 input_values: &[Option<Vec<i64>>],
516) -> Option<Vec<Vec<usize>>> {
517 match node.op_type.as_str() {
518 "Slice" => {
523 let data_shape = input_shapes.first()?;
524 let starts = input_values.get(1)?.as_ref()?;
525 let ends = input_values.get(2)?.as_ref()?;
526 let (axes, steps) = onnx_runtime_ep_cpu::slice_axes_steps(
527 starts.len(),
528 input_values.get(3).and_then(|v| v.as_deref()),
529 input_values.get(4).and_then(|v| v.as_deref()),
530 );
531 let plan =
535 onnx_runtime_ep_cpu::slice_plan(data_shape, starts, ends, &axes, &steps).ok()?;
536 let count: Vec<usize> = plan.iter().map(|p| p.count).collect();
537 Some(vec![count])
538 }
539 _ => None,
540 }
541}
542
543impl Executor {
544 pub(crate) fn build(
546 graph: Graph,
547 weights: Arc<WeightStore>,
548 ep: Arc<CpuExecutionProvider>,
549 ) -> Result<Self> {
550 let order = graph.topological_order()?;
552
553 let mut value_shapes: HashMap<ValueId, Shape> = HashMap::new();
554 let mut value_dtypes: HashMap<ValueId, DataType> = HashMap::new();
555 let mut buffers: HashMap<ValueId, DeviceBuffer> = HashMap::new();
556 let mut buffer_shapes: HashMap<ValueId, Vec<usize>> = HashMap::new();
557
558 let init_align = TensorLayout::contiguous().alignment;
565 for (&vid, weight) in &graph.initializers {
566 let dtype = weight.dtype();
567 let dims = weight.dims().to_vec();
568 let bytes = weights.bytes(weight).ok_or_else(|| {
569 SessionError::Internal(format!("weight bytes unavailable for value#{}", vid.0))
570 })?;
571 let producer_less = graph.value(vid).producer.is_none();
580 let buf = if producer_less
581 && !bytes.is_empty()
582 && (bytes.as_ptr() as usize).is_multiple_of(init_align)
583 {
584 unsafe {
593 DeviceBuffer::from_borrowed_parts(
594 bytes.as_ptr() as *mut std::ffi::c_void,
595 ep.device_id(),
596 bytes.len(),
597 init_align,
598 )
599 }
600 } else {
601 let mut owned = ep.allocate(bytes.len().max(1), init_align)?;
602 write_host(&mut owned, bytes)?;
603 owned
604 };
605 value_dtypes.insert(vid, dtype);
606 value_shapes.insert(vid, dims.iter().map(|&d| Dim::Static(d)).collect());
607 buffer_shapes.insert(vid, dims);
608 buffers.insert(vid, buf);
609 }
610
611 for &vid in &graph.inputs {
615 value_shapes
616 .entry(vid)
617 .or_insert_with(|| graph.value(vid).shape.clone());
618 value_dtypes.entry(vid).or_insert(graph.value(vid).dtype);
619 }
620 for &nid in &order {
621 for &out in &graph.node(nid).outputs {
622 value_shapes
623 .entry(out)
624 .or_insert_with(|| graph.value(out).shape.clone());
625 value_dtypes.entry(out).or_insert(graph.value(out).dtype);
626 }
627 }
628
629 let has_symbols = value_shapes.values().any(|s| as_static_shape(s).is_none());
630
631 let mut sequence_values: HashSet<ValueId> = HashSet::new();
636 for &nid in &order {
637 let node = graph.node(nid);
638 if produces_sequence_output(&node.op_type, &node.domain) {
639 for &out in &node.outputs {
640 sequence_values.insert(out);
641 }
642 }
643 }
644
645 let mut plan = Vec::with_capacity(order.len());
647 for &nid in &order {
648 let node = graph.node(nid);
649 if onnx_runtime_loader::is_ep_context_op(&node.op_type, &node.domain) {
655 continue;
656 }
657 let mut slots: Vec<Option<ValueId>> = node.inputs.clone();
662 while matches!(slots.last(), Some(None)) {
663 slots.pop();
664 }
665 let inputs = slots;
666 let outputs: Vec<ValueId> = node.outputs.clone();
667 let input_dtypes: Vec<DataType> = inputs
668 .iter()
669 .map(|v| v.map(|vid| value_dtypes[&vid]).unwrap_or(DataType::Float32))
670 .collect();
671 let output_dtypes: Vec<DataType> = outputs.iter().map(|v| value_dtypes[v]).collect();
672 plan.push(NodePlan {
673 node_id: nid,
674 inputs,
675 outputs,
676 input_dtypes,
677 output_dtypes,
678 });
679 }
680
681 let mut input_index = HashMap::new();
683 let mut required_inputs = Vec::new();
684 for &vid in &graph.inputs {
685 if graph.initializers.contains_key(&vid) {
686 continue; }
688 required_inputs.push(vid);
689 if let Some(name) = &graph.value(vid).name {
690 input_index.insert(name.clone(), vid);
691 }
692 }
693
694 let mut name_index = HashMap::new();
697 for (vid, value) in graph.values.iter() {
698 if let Some(name) = &value.name {
699 name_index.insert(name.clone(), vid);
700 }
701 }
702
703 let mut exec = Self {
704 graph,
705 weights,
706 ep,
707 buffers,
708 buffer_shapes,
709 value_shapes,
710 value_dtypes,
711 plan,
712 input_index,
713 required_inputs,
714 has_symbols,
715 cache: KernelCache::default(),
716 name_index,
717 subgraph_execs: HashMap::new(),
718 control_flow_stats: ControlFlowStats::default(),
719 views: HashMap::new(),
720 pinned: HashSet::new(),
721 sequence_values,
722 sequences: HashMap::new(),
723 seq_elem_values: HashMap::new(),
724 };
725
726 if !exec.has_symbols {
731 let empty = HashMap::new();
732 let resolved = exec.resolve_all(&empty)?;
733 exec.size_buffers(&resolved)?;
734 exec.compile_all(&resolved)?;
735 }
736 Ok(exec)
737 }
738
739 fn ensure_buffer(&mut self, vid: ValueId, dtype: DataType, dims: &[usize]) -> Result<()> {
742 if self.buffer_shapes.get(&vid).map(|s| s.as_slice()) == Some(dims) {
743 return Ok(()); }
745 if let Some(old) = self.buffers.remove(&vid) {
746 self.ep.deallocate(old)?;
747 }
748 let numel = checked_numel(dims, || format!("value#{}", vid.0))?;
749 let size = checked_storage_bytes(dtype, numel, || format!("value#{}", vid.0), dims)?;
750 let buf = self
751 .ep
752 .allocate(size.max(1), TensorLayout::contiguous().alignment)?;
753 self.buffers.insert(vid, buf);
754 self.buffer_shapes.insert(vid, dims.to_vec());
755 Ok(())
756 }
757
758 fn resolve_all(
762 &self,
763 bindings: &HashMap<SymbolId, usize>,
764 ) -> Result<HashMap<ValueId, Vec<usize>>> {
765 let mut resolved = HashMap::with_capacity(self.value_shapes.len());
766 for (&vid, shape) in &self.value_shapes {
767 if self.sequence_values.contains(&vid) {
771 continue;
772 }
773 match substitute(shape, bindings) {
774 Some(dims) => {
775 resolved.insert(vid, dims);
776 }
777 None => {
778 let value = self.graph.value(vid);
779 let name = value
780 .name
781 .clone()
782 .unwrap_or_else(|| format!("value#{}", vid.0));
783 let op = value
784 .producer
785 .map(|nid| self.graph.node(nid).op_type.clone())
786 .unwrap_or_else(|| "<graph input>".to_string());
787 return Err(SessionError::UnresolvedShape { value: name, op });
788 }
789 }
790 }
791 Ok(resolved)
792 }
793
794 fn resolve_soft(&self, bindings: &HashMap<SymbolId, usize>) -> HashMap<ValueId, Vec<usize>> {
799 let mut resolved = HashMap::with_capacity(self.value_shapes.len());
800 for (&vid, shape) in &self.value_shapes {
801 if let Some(dims) = substitute(shape, bindings) {
802 resolved.insert(vid, dims);
803 }
804 }
805 resolved
806 }
807
808 fn size_buffers(&mut self, resolved: &HashMap<ValueId, Vec<usize>>) -> Result<()> {
814 let vids: Vec<ValueId> = self.value_shapes.keys().copied().collect();
815 for vid in vids {
816 if self.graph.initializers.contains_key(&vid) {
817 continue;
818 }
819 if self.sequence_values.contains(&vid) {
822 continue;
823 }
824 let dtype = self.value_dtypes[&vid];
825 let Some(dims) = resolved.get(&vid).cloned() else {
826 continue;
827 };
828 self.ensure_buffer(vid, dtype, &dims)?;
829 }
830 Ok(())
831 }
832
833 fn node_input_shapes(
837 plan: &NodePlan,
838 resolved: &HashMap<ValueId, Vec<usize>>,
839 ) -> Vec<Vec<usize>> {
840 plan.inputs
841 .iter()
842 .map(|v| v.map(|vid| resolved[&vid].clone()).unwrap_or_default())
843 .collect()
844 }
845
846 fn compile_all(&mut self, resolved: &HashMap<ValueId, Vec<usize>>) -> Result<()> {
848 for i in 0..self.plan.len() {
849 let node_id = self.plan[i].node_id;
850 let node = self.graph.node(node_id);
851 if is_control_flow_op(&node.op_type, &node.domain) {
855 continue;
856 }
857 if is_sequence_op(&node.op_type, &node.domain) {
861 continue;
862 }
863 let input_shapes = Self::node_input_shapes(&self.plan[i], resolved);
864 let node = self.graph.node(node_id);
865 let opset = effective_opset(&self.graph, node);
866 self.cache
867 .get_or_create(node_id, node, &input_shapes, opset, &self.ep)?;
868 }
869 Ok(())
870 }
871
872 pub(crate) fn cache_stats(&self) -> CacheStats {
873 self.cache.stats()
874 }
875
876 pub(crate) fn control_flow_stats(&self) -> ControlFlowStats {
877 self.control_flow_stats
878 }
879
880 pub(crate) fn graph(&self) -> &Graph {
884 &self.graph
885 }
886
887 pub(crate) fn weights(&self) -> &Arc<WeightStore> {
890 &self.weights
891 }
892
893 pub(crate) fn warmup(&mut self) -> Result<()> {
898 if self.has_symbols {
899 return Ok(());
900 }
901 let empty = HashMap::new();
902 let resolved = self.resolve_all(&empty)?;
903 self.compile_all(&resolved)
904 }
905
906 fn bind_symbols(
909 &self,
910 inputs: &[(&str, &Tensor)],
911 ) -> Result<HashMap<SymbolId, usize>> {
912 let mut bindings: HashMap<SymbolId, usize> = HashMap::new();
913 for (name, tensor) in inputs {
914 let vid = *self
915 .input_index
916 .get(*name)
917 .ok_or_else(|| SessionError::InputNotFound {
918 name: (*name).to_string(),
919 })?;
920 let want_dtype = self.value_dtypes[&vid];
921 if tensor.dtype != want_dtype {
922 return Err(SessionError::DtypeMismatch {
923 name: (*name).to_string(),
924 expected: format!("{want_dtype:?}"),
925 got: format!("{:?}", tensor.dtype),
926 });
927 }
928 let decl = &self.value_shapes[&vid];
929 if decl.len() != tensor.shape.len() {
930 return Err(SessionError::RankMismatch {
931 name: (*name).to_string(),
932 expected: decl.len(),
933 got: tensor.shape.len(),
934 });
935 }
936 for (dim, &actual) in decl.iter().zip(&tensor.shape) {
937 match dim {
938 Dim::Static(n) => {
939 if *n != actual {
940 return Err(SessionError::ShapeMismatch {
941 name: (*name).to_string(),
942 expected: as_static_shape(decl).unwrap_or_default(),
943 got: tensor.shape.clone(),
944 });
945 }
946 }
947 Dim::Symbolic(s) => {
948 if let Some(&prev) = bindings.get(s) {
949 if prev != actual {
950 let sym = self
951 .symbol_name(*s)
952 .unwrap_or_else(|| format!("symbol#{}", s.0));
953 return Err(SessionError::SymbolConflict {
954 symbol: sym,
955 first: prev,
956 second: actual,
957 });
958 }
959 } else {
960 bindings.insert(*s, actual);
961 }
962 }
963 }
964 }
965 }
966 Ok(bindings)
967 }
968
969 fn symbol_name(&self, s: SymbolId) -> Option<String> {
971 self.graph
972 .symbol_constraints
973 .get(&s)
974 .and_then(|c| c.name.clone())
975 }
976
977 pub(crate) fn run(&mut self, inputs: &[(&str, &Tensor)]) -> Result<Vec<Tensor>> {
979 self.run_scoped(inputs, &HashMap::new())
980 }
981
982 fn run_scoped(
988 &mut self,
989 inputs: &[(&str, &Tensor)],
990 outer_scope: &HashMap<String, Tensor>,
991 ) -> Result<Vec<Tensor>> {
992 self.views.clear();
995 self.pinned.clear();
996 self.sequences.clear();
999 self.seq_elem_values.clear();
1000
1001 let bindings = self.bind_symbols(inputs)?;
1003
1004 let provided: Vec<ValueId> = inputs
1006 .iter()
1007 .filter_map(|(name, _)| self.input_index.get(*name).copied())
1008 .collect();
1009 for &vid in &self.required_inputs {
1010 if !provided.contains(&vid) {
1011 let name = self
1012 .graph
1013 .value(vid)
1014 .name
1015 .clone()
1016 .unwrap_or_else(|| format!("value#{}", vid.0));
1017 return Err(SessionError::InputNotFound { name });
1018 }
1019 }
1020
1021 let mut resolved = self.resolve_soft(&bindings);
1026 self.size_buffers(&resolved)?;
1027
1028 for (name, tensor) in inputs {
1030 let vid = self.input_index[*name];
1031 let buf = self
1032 .buffers
1033 .get_mut(&vid)
1034 .expect("input value has a buffer");
1035 write_host(buf, tensor.as_bytes())?;
1036 }
1037
1038 for pi in 0..self.plan.len() {
1043 let node_id = self.plan[pi].node_id;
1044 let node = self.graph.node(node_id);
1045 if is_control_flow_op(&node.op_type, &node.domain) {
1046 self.exec_control_flow(pi, &mut resolved, outer_scope)?;
1047 } else if is_sequence_op(&node.op_type, &node.domain) {
1048 self.exec_sequence_node(pi, &mut resolved)?;
1049 } else {
1050 self.exec_kernel_node(pi, &mut resolved)?;
1051 }
1052 }
1053
1054 let mut results = Vec::with_capacity(self.graph.outputs.len());
1059 for &vid in &self.graph.outputs {
1060 if self.sequence_values.contains(&vid) {
1065 let name = self
1066 .graph
1067 .try_value(vid)
1068 .and_then(|v| v.name.clone())
1069 .unwrap_or_else(|| format!("value#{}", vid.0));
1070 return Err(SessionError::SequenceOp {
1071 op: "<graph output>".to_string(),
1072 reason: format!(
1073 "graph output {name} is a Sequence value, which cannot be \
1074 returned through the tensor `run` API. To fix: end the graph \
1075 with ConcatFromSequence or SequenceAt to produce tensor \
1076 output(s)"
1077 ),
1078 });
1079 }
1080 let dtype = self.value_dtypes[&vid];
1081 let shape = resolved[&vid].clone();
1082 let bytes = self.contiguous_bytes(vid, &shape, dtype)?;
1083 results.push(Tensor::from_raw_in(self.ep.clone(), dtype, shape, &bytes)?);
1084 }
1085 Ok(results)
1086 }
1087
1088 fn exec_kernel_node(
1092 &mut self,
1093 pi: usize,
1094 resolved: &mut HashMap<ValueId, Vec<usize>>,
1095 ) -> Result<()> {
1096 let node_id = self.plan[pi].node_id;
1099 let inputs = self.plan[pi].inputs.clone();
1100 let outputs = self.plan[pi].outputs.clone();
1101 let input_dtypes = self.plan[pi].input_dtypes.clone();
1102 let output_dtypes = self.plan[pi].output_dtypes.clone();
1103
1104 let input_shapes: Vec<Vec<usize>> = inputs
1105 .iter()
1106 .map(|v| v.map(|vid| resolved[&vid].clone()).unwrap_or_default())
1107 .collect();
1108
1109 if outputs.iter().any(|v| !resolved.contains_key(v)) {
1114 let input_values: Vec<Option<Vec<i64>>> = inputs
1115 .iter()
1116 .enumerate()
1117 .map(|(i, v)| {
1118 v.and_then(|vid| self.input_i64(vid, &input_shapes[i], input_dtypes[i]))
1119 })
1120 .collect();
1121 let node = self.graph.node(node_id);
1122 let out_shapes = dynamic_output_shapes(node, &input_shapes, &input_values)
1123 .ok_or_else(|| {
1124 let vid = outputs
1125 .iter()
1126 .find(|v| !resolved.contains_key(v))
1127 .copied()
1128 .unwrap_or(outputs[0]);
1129 let value = self.graph.value(vid);
1130 SessionError::UnresolvedShape {
1131 value: value
1132 .name
1133 .clone()
1134 .unwrap_or_else(|| format!("value#{}", vid.0)),
1135 op: node.op_type.clone(),
1136 }
1137 })?;
1138 if out_shapes.len() != outputs.len() {
1139 return Err(SessionError::OutputShapeCountMismatch {
1140 op: self.graph.node(node_id).op_type.clone(),
1141 expected: outputs.len(),
1142 got: out_shapes.len(),
1143 });
1144 }
1145 for (oi, &ovid) in outputs.iter().enumerate() {
1146 resolved.insert(ovid, out_shapes[oi].clone());
1147 }
1148 }
1149
1150 let output_shapes: Vec<Vec<usize>> =
1151 outputs.iter().map(|v| resolved[v].clone()).collect();
1152
1153 let mut in_infos: Vec<InInfo> = Vec::with_capacity(inputs.len());
1156 for (i, slot) in inputs.iter().enumerate() {
1157 let Some(vid) = *slot else {
1158 in_infos.push(InInfo {
1159 present: false,
1160 dtype: input_dtypes[i],
1161 shape: Vec::new(),
1162 strides: Vec::new(),
1163 byte_offset: 0,
1164 base_ptr: std::ptr::null(),
1165 root_len: 0,
1166 });
1167 continue;
1168 };
1169 if let Some(elem) = self.seq_elem_values.get(&vid) {
1175 let shape = input_shapes[i].clone();
1176 let strides = compute_contiguous_strides(&shape);
1177 let root_len = elem.data.len();
1178 let base_ptr = elem.as_ptr() as *const std::ffi::c_void;
1179 view_bounds(&shape, &strides, 0, input_dtypes[i], root_len)?;
1180 in_infos.push(InInfo {
1181 present: true,
1182 dtype: input_dtypes[i],
1183 shape,
1184 strides,
1185 byte_offset: 0,
1186 base_ptr,
1187 root_len,
1188 });
1189 continue;
1190 }
1191 let root = self.root_of(vid);
1192 let buf = self.buffers.get(&root).ok_or_else(|| {
1193 SessionError::Internal(format!("missing buffer for input value#{}", vid.0))
1194 })?;
1195 let root_len = buf.len();
1196 let base_ptr = buf.as_ptr();
1197 let (shape, strides, byte_offset) = match self.views.get(&vid) {
1198 Some(view) => (view.shape.clone(), view.strides.clone(), view.byte_offset),
1199 None => {
1200 let shape = input_shapes[i].clone();
1201 let strides = compute_contiguous_strides(&shape);
1202 (shape, strides, 0)
1203 }
1204 };
1205 view_bounds(&shape, &strides, byte_offset, input_dtypes[i], root_len)?;
1206 in_infos.push(InInfo {
1207 present: true,
1208 dtype: input_dtypes[i],
1209 shape,
1210 strides,
1211 byte_offset,
1212 base_ptr,
1213 root_len,
1214 });
1215 }
1216
1217 let ep = self.ep.clone();
1218
1219 let graph = &self.graph;
1222 let cache = &mut self.cache;
1223 let buffers = &mut self.buffers;
1224 let buffer_shapes = &mut self.buffer_shapes;
1225 let views_meta = &mut self.views;
1226 let pinned = &mut self.pinned;
1227
1228 let mut views: Vec<TensorView> = Vec::with_capacity(in_infos.len());
1231 for info in &in_infos {
1232 if !info.present {
1233 views.push(TensorView::absent(info.dtype));
1234 continue;
1235 }
1236 views.push(
1237 TensorView::new(
1238 DevicePtr(info.base_ptr),
1239 info.dtype,
1240 &info.shape,
1241 &info.strides,
1242 onnx_runtime_ir::DeviceId::cpu(),
1243 )
1244 .with_byte_offset(info.byte_offset),
1245 );
1246 }
1247
1248 let node = graph.node(node_id);
1249 let opset = effective_opset(graph, node);
1250 let kernel = cache.get_or_create(node_id, node, &input_shapes, opset, &ep)?;
1251
1252 if let Some(specs) = kernel.view_outputs(&views, outputs.len()) {
1257 drop(views);
1258 if specs.len() != outputs.len() {
1259 return Err(SessionError::Internal(format!(
1260 "op '{}' returned {} view outputs for {} outputs",
1261 node.op_type,
1262 specs.len(),
1263 outputs.len()
1264 )));
1265 }
1266 for (oi, spec) in specs.into_iter().enumerate() {
1267 let ovid = outputs[oi];
1268 let Some(in_vid) = inputs.get(spec.input_index).copied().flatten() else {
1269 return Err(SessionError::Internal(format!(
1270 "op '{}' view output {} references invalid input index {}",
1271 node.op_type, oi, spec.input_index
1272 )));
1273 };
1274 let root = match views_meta.get(&in_vid) {
1275 Some(v) => v.source,
1276 None => in_vid,
1277 };
1278 let root_len = buffers.get(&root).map(|b| b.len()).ok_or_else(|| {
1279 SessionError::Internal(format!(
1280 "view source value#{} has no buffer",
1281 root.0
1282 ))
1283 })?;
1284 view_bounds(
1286 &spec.shape,
1287 &spec.strides,
1288 spec.byte_offset,
1289 output_dtypes[oi],
1290 root_len,
1291 )?;
1292 debug_assert!(
1299 !pinned.contains(&ovid),
1300 "value#{} is pinned as a live view source yet is being reproduced",
1301 ovid.0
1302 );
1303 if let Some(old) = buffers.remove(&ovid) {
1304 ep.deallocate(old)?;
1305 }
1306 buffer_shapes.remove(&ovid);
1307 views_meta.insert(
1308 ovid,
1309 ValueView {
1310 source: root,
1311 shape: spec.shape.clone(),
1312 strides: spec.strides,
1313 byte_offset: spec.byte_offset,
1314 },
1315 );
1316 pinned.insert(root);
1317 resolved.insert(ovid, spec.shape);
1318 }
1319 return Ok(());
1320 }
1321
1322 for (oi, &ovid) in outputs.iter().enumerate() {
1327 let dims = &output_shapes[oi];
1328 let numel = checked_numel(dims, || format!("value#{}", ovid.0))?;
1329 let need = checked_storage_bytes(
1330 output_dtypes[oi],
1331 numel,
1332 || format!("value#{}", ovid.0),
1333 dims,
1334 )?
1335 .max(1);
1336 let fits = buffers.get(&ovid).map(|b| b.len() == need).unwrap_or(false);
1337 if !fits {
1338 debug_assert!(
1341 !pinned.contains(&ovid),
1342 "value#{} is pinned as a live view source yet is being resized",
1343 ovid.0
1344 );
1345 if let Some(old) = buffers.remove(&ovid) {
1346 ep.deallocate(old)?;
1347 }
1348 let buf = ep.allocate(need, TensorLayout::contiguous().alignment)?;
1349 buffers.insert(ovid, buf);
1350 }
1351 }
1352
1353 let mut mat: Vec<Option<(Vec<u8>, Vec<i64>)>> = Vec::with_capacity(in_infos.len());
1358 for (i, info) in in_infos.iter().enumerate() {
1359 if !info.present {
1360 mat.push(None);
1361 continue;
1362 }
1363 let contiguous = onnx_runtime_ir::is_contiguous(&info.shape, &info.strides);
1364 if contiguous || kernel.supports_strided_input(i) {
1365 mat.push(None);
1366 continue;
1367 }
1368 let esize = info.dtype.byte_size();
1369 if esize == 0 {
1370 return Err(SessionError::from(
1371 onnx_runtime_ep_api::EpError::InvalidTensorView {
1372 reason: format!(
1373 "cannot materialize sub-byte strided input {i} of op '{}'",
1374 node.op_type
1375 ),
1376 },
1377 ));
1378 }
1379 let src = unsafe {
1380 std::slice::from_raw_parts(info.base_ptr as *const u8, info.root_len)
1381 };
1382 let gathered = gather_view(src, &info.shape, &info.strides, info.byte_offset, esize);
1383 let strides = compute_contiguous_strides(&info.shape);
1384 mat.push(Some((gathered, strides)));
1385 }
1386
1387 drop(views);
1390 let mut views: Vec<TensorView> = Vec::with_capacity(in_infos.len());
1391 for (i, info) in in_infos.iter().enumerate() {
1392 if !info.present {
1393 views.push(TensorView::absent(info.dtype));
1394 continue;
1395 }
1396 match &mat[i] {
1397 Some((buf, strides)) => views.push(TensorView::new(
1398 DevicePtr(buf.as_ptr() as *const std::ffi::c_void),
1399 info.dtype,
1400 &info.shape,
1401 strides,
1402 onnx_runtime_ir::DeviceId::cpu(),
1403 )),
1404 None => views.push(
1405 TensorView::new(
1406 DevicePtr(info.base_ptr),
1407 info.dtype,
1408 &info.shape,
1409 &info.strides,
1410 onnx_runtime_ir::DeviceId::cpu(),
1411 )
1412 .with_byte_offset(info.byte_offset),
1413 ),
1414 }
1415 }
1416
1417 let out_strides: Vec<Vec<i64>> = output_shapes
1420 .iter()
1421 .map(|s| compute_contiguous_strides(s))
1422 .collect();
1423 let mut out_bufs: Vec<(ValueId, DeviceBuffer)> = Vec::with_capacity(outputs.len());
1424 for &vid in &outputs {
1425 let buf = buffers.remove(&vid).ok_or_else(|| {
1426 SessionError::Internal(format!("missing buffer for output value#{}", vid.0))
1427 })?;
1428 out_bufs.push((vid, buf));
1429 }
1430 let mut outs: Vec<TensorMut> = Vec::with_capacity(out_bufs.len());
1431 for (i, (_, buf)) in out_bufs.iter_mut().enumerate() {
1432 view_bounds(
1433 &output_shapes[i],
1434 &out_strides[i],
1435 0,
1436 output_dtypes[i],
1437 buf.len(),
1438 )?;
1439 let ptr = buf.as_mut_ptr();
1440 outs.push(TensorMut::new(
1441 DevicePtrMut(ptr),
1442 output_dtypes[i],
1443 &output_shapes[i],
1444 &out_strides[i],
1445 onnx_runtime_ir::DeviceId::cpu(),
1446 ));
1447 }
1448
1449 kernel.execute(&views, &mut outs)?;
1450
1451 drop(views);
1452 drop(outs);
1453 for (vid, buf) in out_bufs {
1454 buffers.insert(vid, buf);
1455 }
1456 Ok(())
1457 }
1458
1459 fn input_i64(&self, vid: ValueId, shape: &[usize], dtype: DataType) -> Option<Vec<i64>> {
1464 if self.views.contains_key(&vid) || self.seq_elem_values.contains_key(&vid) {
1465 let bytes = self.contiguous_bytes(vid, shape, dtype).ok()?;
1466 bytes_as_i64(&bytes, dtype)
1467 } else {
1468 self.buffers.get(&vid).and_then(|b| buffer_as_i64(b, dtype))
1469 }
1470 }
1471}
1472
1473impl Executor {
1501 fn exec_sequence_node(
1503 &mut self,
1504 pi: usize,
1505 resolved: &mut HashMap<ValueId, Vec<usize>>,
1506 ) -> Result<()> {
1507 let node_id = self.plan[pi].node_id;
1508 let inputs = self.plan[pi].inputs.clone();
1509 let outputs = self.plan[pi].outputs.clone();
1510 let op = self.graph.node(node_id).op_type.clone();
1511
1512 match op.as_str() {
1513 "SequenceEmpty" => {
1514 let dtype_attr = self.graph.node(node_id).attr("dtype").and_then(|a| a.as_int());
1515 let dtype = match dtype_attr {
1516 None => DataType::Float32, Some(raw) => DataType::from_onnx(raw as i32).ok_or_else(|| {
1518 SessionError::SequenceOp {
1519 op: op.clone(),
1520 reason: format!(
1521 "attribute 'dtype' = {raw} is not a known ONNX \
1522 TensorProto.DataType. To fix: use a valid element \
1523 dtype id (e.g. 1=float32, 7=int64)"
1524 ),
1525 }
1526 })?,
1527 };
1528 self.sequences
1529 .insert(outputs[0], SequenceValue::empty(dtype));
1530 Ok(())
1531 }
1532 "SequenceConstruct" => {
1533 let mut items = Vec::with_capacity(inputs.len());
1534 for slot in &inputs {
1535 let vid = slot.ok_or_else(|| self.seq_missing_input(&op))?;
1536 items.push(self.read_seq_element(vid, resolved)?);
1537 }
1538 let seq = SequenceValue::construct(items).map_err(seq_err)?;
1539 self.sequences.insert(outputs[0], seq);
1540 Ok(())
1541 }
1542 "SequenceInsert" => {
1543 let seq = self.get_sequence(inputs.first().copied().flatten(), &op)?;
1544 let tvid = inputs.get(1).copied().flatten().ok_or_else(|| {
1545 self.seq_missing_input(&op)
1546 })?;
1547 let tensor = self.read_seq_element(tvid, resolved)?;
1548 let position = match inputs.get(2).copied().flatten() {
1549 Some(pvid) => Some(self.read_scalar_i64(pvid, resolved, &op)?),
1550 None => None,
1551 };
1552 let out = seq.insert(tensor, position).map_err(seq_err)?;
1553 self.sequences.insert(outputs[0], out);
1554 Ok(())
1555 }
1556 "SequenceErase" => {
1557 let seq = self.get_sequence(inputs.first().copied().flatten(), &op)?;
1558 let position = match inputs.get(1).copied().flatten() {
1559 Some(pvid) => Some(self.read_scalar_i64(pvid, resolved, &op)?),
1560 None => None,
1561 };
1562 let out = seq.erase(position).map_err(seq_err)?;
1563 self.sequences.insert(outputs[0], out);
1564 Ok(())
1565 }
1566 "SequenceAt" => {
1567 let seq = self.get_sequence(inputs.first().copied().flatten(), &op)?;
1568 let pvid = inputs.get(1).copied().flatten().ok_or_else(|| {
1569 SessionError::SequenceOp {
1570 op: op.clone(),
1571 reason: "requires a 'position' input. To fix: supply the \
1572 index tensor of the element to read"
1573 .to_string(),
1574 }
1575 })?;
1576 let pos = self.read_scalar_i64(pvid, resolved, &op)?;
1577 let elem = seq.at(pos).map_err(seq_err)?;
1578 self.store_seq_element_output(outputs[0], elem, resolved)
1579 }
1580 "SequenceLength" => {
1581 let seq = self.get_sequence(inputs.first().copied().flatten(), &op)?;
1582 let len = seq.len() as i64;
1583 self.store_raw_tensor_output(
1584 outputs[0],
1585 DataType::Int64,
1586 Vec::new(),
1587 &len.to_le_bytes(),
1588 resolved,
1589 )
1590 }
1591 "SplitToSequence" => self.exec_split_to_sequence(&op, &inputs, &outputs, resolved),
1592 "ConcatFromSequence" => {
1593 self.exec_concat_from_sequence(node_id, &op, &inputs, &outputs, resolved)
1594 }
1595 other => Err(SessionError::SequenceOp {
1596 op: other.to_string(),
1597 reason: "unrecognized Sequence op (executor routing bug)".to_string(),
1598 }),
1599 }
1600 }
1601
1602 fn exec_split_to_sequence(
1604 &mut self,
1605 op: &str,
1606 inputs: &[Option<ValueId>],
1607 outputs: &[ValueId],
1608 resolved: &mut HashMap<ValueId, Vec<usize>>,
1609 ) -> Result<()> {
1610 let node = self.graph.node(self.plan_node_of(outputs[0]));
1611 let axis_attr = node.attr("axis").and_then(|a| a.as_int()).unwrap_or(0);
1612 let keepdims = node.attr("keepdims").and_then(|a| a.as_int()).unwrap_or(1) != 0;
1613
1614 let ivid = inputs.first().copied().flatten().ok_or_else(|| self.seq_missing_input(op))?;
1615 let dtype = self.value_dtypes[&ivid];
1616 let esize = dtype.byte_size();
1617 if esize == 0 {
1618 return Err(SessionError::SequenceOp {
1619 op: op.to_string(),
1620 reason: format!(
1621 "sub-byte dtype {dtype:?} is not supported for SplitToSequence. \
1622 To fix: Cast to a byte-addressable dtype before splitting"
1623 ),
1624 });
1625 }
1626 let shape = resolved
1627 .get(&ivid)
1628 .cloned()
1629 .ok_or_else(|| self.seq_unresolved(op, ivid))?;
1630 let rank = shape.len();
1631 if rank == 0 {
1632 return Err(SessionError::SequenceOp {
1633 op: op.to_string(),
1634 reason: "cannot split a scalar (rank-0) tensor. To fix: split a \
1635 tensor with at least one dimension"
1636 .to_string(),
1637 });
1638 }
1639 let axis = normalize_axis(axis_attr, rank).ok_or_else(|| SessionError::SequenceOp {
1640 op: op.to_string(),
1641 reason: format!(
1642 "attribute 'axis' = {axis_attr} is out of range for a rank-{rank} \
1643 input (valid range is [{}, {}])",
1644 -(rank as i64),
1645 rank as i64 - 1
1646 ),
1647 })?;
1648 let axis_dim = shape[axis];
1649 let bytes = self.contiguous_bytes(ivid, &shape, dtype)?;
1650
1651 let mut squeeze = false;
1653 let sizes: Vec<usize> = match inputs.get(1).copied().flatten() {
1654 None => {
1655 squeeze = !keepdims;
1658 vec![1; axis_dim]
1659 }
1660 Some(svid) => {
1661 let sshape = resolved.get(&svid).cloned().unwrap_or_default();
1662 let svals = self.read_i64_vec(svid, &sshape, op)?;
1663 let is_scalar = sshape.is_empty();
1664 if is_scalar {
1665 let chunk = *svals.first().ok_or_else(|| SessionError::SequenceOp {
1666 op: op.to_string(),
1667 reason: "'split' scalar is empty".to_string(),
1668 })?;
1669 if chunk <= 0 {
1670 return Err(SessionError::SequenceOp {
1671 op: op.to_string(),
1672 reason: format!(
1673 "'split' chunk size {chunk} must be positive"
1674 ),
1675 });
1676 }
1677 let chunk = chunk as usize;
1678 let mut v = Vec::new();
1679 let mut rem = axis_dim;
1680 while rem > 0 {
1681 let k = rem.min(chunk);
1682 v.push(k);
1683 rem -= k;
1684 }
1685 v
1686 } else {
1687 let v: Vec<usize> = svals.iter().map(|&x| x.max(0) as usize).collect();
1688 let sum: usize = v.iter().sum();
1689 if sum != axis_dim {
1690 return Err(SessionError::SequenceOp {
1691 op: op.to_string(),
1692 reason: format!(
1693 "'split' sizes {v:?} sum to {sum} but axis {axis} has \
1694 extent {axis_dim}. To fix: make the split sizes sum to \
1695 the axis length"
1696 ),
1697 });
1698 }
1699 v
1700 }
1701 }
1702 };
1703
1704 let parts = split_axis(&bytes, &shape, axis, &sizes, esize);
1705 let items: Vec<std::sync::Arc<SeqTensor>> = parts
1706 .into_iter()
1707 .map(|(mut sh, data)| {
1708 if squeeze {
1709 sh.remove(axis);
1710 }
1711 SeqTensor::shared(dtype, sh, data)
1712 })
1713 .collect();
1714 self.sequences.insert(
1715 outputs[0],
1716 SequenceValue {
1717 elem_dtype: dtype,
1718 items,
1719 },
1720 );
1721 Ok(())
1722 }
1723
1724 fn exec_concat_from_sequence(
1727 &mut self,
1728 node_id: NodeId,
1729 op: &str,
1730 inputs: &[Option<ValueId>],
1731 outputs: &[ValueId],
1732 resolved: &mut HashMap<ValueId, Vec<usize>>,
1733 ) -> Result<()> {
1734 let node = self.graph.node(node_id);
1735 let axis_attr = node.attr("axis").and_then(|a| a.as_int()).ok_or_else(|| {
1736 SessionError::SequenceOp {
1737 op: op.to_string(),
1738 reason: "requires the mandatory 'axis' attribute. To fix: set 'axis'"
1739 .to_string(),
1740 }
1741 })?;
1742 let new_axis = node.attr("new_axis").and_then(|a| a.as_int()).unwrap_or(0) != 0;
1743
1744 let seq = self.get_sequence(inputs.first().copied().flatten(), op)?;
1745 if seq.is_empty() {
1746 return Err(SessionError::SequenceOp {
1747 op: op.to_string(),
1748 reason: "cannot concatenate an empty sequence (output shape is \
1749 undefined). To fix: guard with SequenceLength"
1750 .to_string(),
1751 });
1752 }
1753 let dtype = seq.elem_dtype;
1754 let esize = dtype.byte_size();
1755 if esize == 0 {
1756 return Err(SessionError::SequenceOp {
1757 op: op.to_string(),
1758 reason: format!(
1759 "sub-byte dtype {dtype:?} is not supported for ConcatFromSequence"
1760 ),
1761 });
1762 }
1763 let elem_shapes: Vec<Vec<usize>> = seq.items.iter().map(|t| t.shape.clone()).collect();
1764 let elem_datas: Vec<&[u8]> = seq.items.iter().map(|t| t.data.as_slice()).collect();
1765 let rank = elem_shapes[0].len();
1766
1767 let (oshape, out) = if new_axis {
1768 for (i, s) in elem_shapes.iter().enumerate() {
1770 if s != &elem_shapes[0] {
1771 return Err(SessionError::SequenceOp {
1772 op: op.to_string(),
1773 reason: format!(
1774 "ConcatFromSequence(new_axis=1) requires identical element \
1775 shapes, but element {i} has shape {s:?} vs {:?}",
1776 elem_shapes[0]
1777 ),
1778 });
1779 }
1780 }
1781 let axis = normalize_axis(axis_attr, rank + 1).ok_or_else(|| {
1782 SessionError::SequenceOp {
1783 op: op.to_string(),
1784 reason: format!(
1785 "'axis' = {axis_attr} is out of range for new_axis=1 stacking \
1786 of rank-{rank} elements (valid range is [{}, {}])",
1787 -(rank as i64) - 1,
1788 rank as i64
1789 ),
1790 }
1791 })?;
1792 stack_new_axis(&elem_datas, &elem_shapes[0], axis, esize)
1793 } else {
1794 let axis = normalize_axis(axis_attr, rank).ok_or_else(|| SessionError::SequenceOp {
1795 op: op.to_string(),
1796 reason: format!(
1797 "'axis' = {axis_attr} is out of range for rank-{rank} elements \
1798 (valid range is [{}, {}])",
1799 -(rank as i64),
1800 rank as i64 - 1
1801 ),
1802 })?;
1803 for (i, s) in elem_shapes.iter().enumerate() {
1805 let mismatch = s.len() != rank
1806 || s.iter().enumerate().any(|(d, &v)| d != axis && v != elem_shapes[0][d]);
1807 if mismatch {
1808 return Err(SessionError::SequenceOp {
1809 op: op.to_string(),
1810 reason: format!(
1811 "ConcatFromSequence requires elements to match on all axes \
1812 except {axis}, but element {i} has shape {s:?} vs {:?}",
1813 elem_shapes[0]
1814 ),
1815 });
1816 }
1817 }
1818 concat_axis(&elem_datas, &elem_shapes, axis, esize)
1819 };
1820 drop(seq);
1821 self.store_raw_tensor_output(outputs[0], dtype, oshape, &out, resolved)
1822 }
1823
1824 fn read_seq_element(
1830 &self,
1831 vid: ValueId,
1832 resolved: &HashMap<ValueId, Vec<usize>>,
1833 ) -> Result<std::sync::Arc<SeqTensor>> {
1834 if let Some(elem) = self.seq_elem_values.get(&vid) {
1835 return Ok(std::sync::Arc::clone(elem)); }
1837 let dtype = self.value_dtypes[&vid];
1838 let shape = resolved
1839 .get(&vid)
1840 .cloned()
1841 .ok_or_else(|| self.seq_unresolved("Sequence", vid))?;
1842 let bytes = self.contiguous_bytes(vid, &shape, dtype)?;
1843 Ok(SeqTensor::shared(dtype, shape, bytes))
1844 }
1845
1846 fn get_sequence(&self, vid: Option<ValueId>, op: &str) -> Result<SequenceValue> {
1849 let vid = vid.ok_or_else(|| self.seq_missing_input(op))?;
1850 self.sequences.get(&vid).cloned().ok_or_else(|| SessionError::SequenceOp {
1851 op: op.to_string(),
1852 reason: format!(
1853 "input value#{} is not a live sequence. To fix: ensure it is produced \
1854 by a Sequence-producing op (SequenceEmpty/Construct/Insert/Erase/\
1855 SplitToSequence)",
1856 vid.0
1857 ),
1858 })
1859 }
1860
1861 fn read_scalar_i64(
1863 &self,
1864 vid: ValueId,
1865 resolved: &HashMap<ValueId, Vec<usize>>,
1866 op: &str,
1867 ) -> Result<i64> {
1868 let shape = resolved.get(&vid).cloned().unwrap_or_default();
1869 let dtype = self.value_dtypes[&vid];
1870 let vals = self.input_i64(vid, &shape, dtype).ok_or_else(|| SessionError::SequenceOp {
1871 op: op.to_string(),
1872 reason: format!(
1873 "position input has dtype {dtype:?}, expected an integer (int32/int64). \
1874 To fix: provide an int64 scalar index"
1875 ),
1876 })?;
1877 vals.first().copied().ok_or_else(|| SessionError::SequenceOp {
1878 op: op.to_string(),
1879 reason: "position input is empty; expected a single scalar index".to_string(),
1880 })
1881 }
1882
1883 fn read_i64_vec(&self, vid: ValueId, shape: &[usize], op: &str) -> Result<Vec<i64>> {
1886 let dtype = self.value_dtypes[&vid];
1887 self.input_i64(vid, shape, dtype).ok_or_else(|| SessionError::SequenceOp {
1888 op: op.to_string(),
1889 reason: format!(
1890 "'split' input has dtype {dtype:?}, expected int32/int64. To fix: \
1891 provide integer split sizes"
1892 ),
1893 })
1894 }
1895
1896 fn store_seq_element_output(
1900 &mut self,
1901 vid: ValueId,
1902 elem: std::sync::Arc<SeqTensor>,
1903 resolved: &mut HashMap<ValueId, Vec<usize>>,
1904 ) -> Result<()> {
1905 if let Some(old) = self.buffers.remove(&vid) {
1906 self.ep.deallocate(old)?;
1907 }
1908 self.buffer_shapes.remove(&vid);
1909 self.views.remove(&vid);
1910 resolved.insert(vid, elem.shape.clone());
1911 self.value_dtypes.insert(vid, elem.dtype);
1912 self.seq_elem_values.insert(vid, elem);
1913 Ok(())
1914 }
1915
1916 fn store_raw_tensor_output(
1920 &mut self,
1921 vid: ValueId,
1922 dtype: DataType,
1923 dims: Vec<usize>,
1924 bytes: &[u8],
1925 resolved: &mut HashMap<ValueId, Vec<usize>>,
1926 ) -> Result<()> {
1927 self.seq_elem_values.remove(&vid);
1930 self.views.remove(&vid);
1931 let need = bytes.len().max(1);
1932 let fits = self.buffers.get(&vid).map(|b| b.len() == need).unwrap_or(false);
1933 if !fits {
1934 if let Some(old) = self.buffers.remove(&vid) {
1935 self.ep.deallocate(old)?;
1936 }
1937 let buf = self.ep.allocate(need, TensorLayout::contiguous().alignment)?;
1938 self.buffers.insert(vid, buf);
1939 }
1940 let buf = self.buffers.get_mut(&vid).expect("just ensured");
1941 write_host(buf, bytes)?;
1942 self.value_dtypes.insert(vid, dtype);
1943 self.buffer_shapes.insert(vid, dims.clone());
1944 resolved.insert(vid, dims);
1945 Ok(())
1946 }
1947
1948 fn plan_node_of(&self, vid: ValueId) -> NodeId {
1951 self.graph
1952 .value(vid)
1953 .producer
1954 .expect("sequence op output has a producer")
1955 }
1956
1957 fn seq_missing_input(&self, op: &str) -> SessionError {
1958 SessionError::SequenceOp {
1959 op: op.to_string(),
1960 reason: "a required input is missing (omitted None slot). To fix: connect \
1961 all required inputs of this Sequence op"
1962 .to_string(),
1963 }
1964 }
1965
1966 fn seq_unresolved(&self, op: &str, vid: ValueId) -> SessionError {
1967 let name = self
1968 .graph
1969 .try_value(vid)
1970 .and_then(|v| v.name.clone())
1971 .unwrap_or_else(|| format!("value#{}", vid.0));
1972 SessionError::SequenceOp {
1973 op: op.to_string(),
1974 reason: format!(
1975 "input {name} has no resolved shape yet. To fix: ensure its producer \
1976 runs before this Sequence op"
1977 ),
1978 }
1979 }
1980}
1981
1982fn seq_err(e: crate::sequence::SeqOpError) -> SessionError {
1984 SessionError::SequenceOp {
1985 op: e.op.to_string(),
1986 reason: e.reason,
1987 }
1988}
1989
1990fn normalize_axis(axis: i64, rank: usize) -> Option<usize> {
1993 let r = rank as i64;
1994 let a = if axis < 0 { axis + r } else { axis };
1995 if a < 0 || a >= r {
1996 None
1997 } else {
1998 Some(a as usize)
1999 }
2000}
2001
2002
2003fn is_control_flow_op(op_type: &str, domain: &str) -> bool {
2008 (domain.is_empty() || domain == "ai.onnx") && matches!(op_type, "If" | "Loop" | "Scan")
2009}
2010
2011fn is_sequence_op(op_type: &str, domain: &str) -> bool {
2018 (domain.is_empty() || domain == "ai.onnx")
2019 && matches!(
2020 op_type,
2021 "SequenceEmpty"
2022 | "SequenceConstruct"
2023 | "SequenceInsert"
2024 | "SequenceErase"
2025 | "SequenceAt"
2026 | "SequenceLength"
2027 | "SplitToSequence"
2028 | "ConcatFromSequence"
2029 )
2030}
2031
2032fn produces_sequence_output(op_type: &str, domain: &str) -> bool {
2035 (domain.is_empty() || domain == "ai.onnx")
2036 && matches!(
2037 op_type,
2038 "SequenceEmpty"
2039 | "SequenceConstruct"
2040 | "SequenceInsert"
2041 | "SequenceErase"
2042 | "SplitToSequence"
2043 )
2044}
2045
2046fn tensor_scalar_i64(t: &Tensor) -> Option<i64> {
2048 match t.dtype {
2049 DataType::Int64 => t
2050 .as_bytes()
2051 .get(..8)
2052 .map(|c| i64::from_le_bytes(c.try_into().unwrap())),
2053 DataType::Int32 => t
2054 .as_bytes()
2055 .get(..4)
2056 .map(|c| i32::from_le_bytes(c.try_into().unwrap()) as i64),
2057 _ => None,
2058 }
2059}
2060
2061fn tensor_scalar_bool(t: &Tensor) -> Option<bool> {
2064 if t.dtype != DataType::Bool {
2065 return None;
2066 }
2067 t.as_bytes().first().map(|&b| b != 0)
2068}
2069
2070fn scalar_i64_tensor(v: i64) -> Result<Tensor> {
2072 Tensor::from_raw(DataType::Int64, vec![], &v.to_le_bytes())
2073}
2074
2075fn scalar_bool_tensor(v: bool) -> Result<Tensor> {
2077 Tensor::from_raw(DataType::Bool, vec![], &[u8::from(v)])
2078}
2079
2080fn missing_capture_error(attr_key: &str, name: &str) -> SessionError {
2081 SessionError::Internal(format!(
2082 "control-flow body '{attr_key}' captures free variable '{name}', but it is not \
2083 available in the enclosing scope. RULES #1: a subgraph may only reference outer \
2084 values that are graph inputs, initializers, or produced by an upstream node in an \
2085 enclosing graph; '{name}' matches none of these"
2086 ))
2087}
2088
2089fn required_outer_names(graph: &Graph) -> HashSet<String> {
2093 let formal_set: HashSet<ValueId> = graph.inputs.iter().copied().collect();
2094 let local_names: HashSet<&str> = graph
2095 .values
2096 .iter()
2097 .filter_map(|(_, value)| value.name.as_deref())
2098 .collect();
2099 let mut required = HashSet::new();
2100 for (vid, value) in graph.values.iter() {
2101 if value.producer.is_none()
2102 && !formal_set.contains(&vid)
2103 && !graph.initializers.contains_key(&vid)
2104 && let Some(name) = &value.name
2105 {
2106 required.insert(name.clone());
2107 }
2108 }
2109 for nested in graph.subgraphs.values() {
2110 for name in required_outer_names(nested) {
2111 if !local_names.contains(name.as_str()) {
2112 required.insert(name);
2113 }
2114 }
2115 }
2116 required
2117}
2118
2119impl Executor {
2127 fn value_tensor(
2130 &self,
2131 vid: ValueId,
2132 resolved: &HashMap<ValueId, Vec<usize>>,
2133 ) -> Result<Tensor> {
2134 let dtype = self.value_dtypes[&vid];
2135 let shape = resolved.get(&vid).cloned().ok_or_else(|| {
2136 let name = self
2137 .graph
2138 .try_value(vid)
2139 .and_then(|v| v.name.clone())
2140 .unwrap_or_else(|| format!("value#{}", vid.0));
2141 SessionError::UnresolvedShape {
2142 value: name,
2143 op: "<control-flow input>".to_string(),
2144 }
2145 })?;
2146 let bytes = self.contiguous_bytes(vid, &shape, dtype)?;
2148 Tensor::from_raw_in(self.ep.clone(), dtype, shape, &bytes)
2149 }
2150
2151 fn root_of(&self, vid: ValueId) -> ValueId {
2155 match self.views.get(&vid) {
2156 Some(v) => v.source,
2157 None => vid,
2158 }
2159 }
2160
2161 fn contiguous_bytes(
2166 &self,
2167 vid: ValueId,
2168 shape: &[usize],
2169 dtype: DataType,
2170 ) -> Result<Vec<u8>> {
2171 let numel: usize = shape.iter().product();
2172 let n = dtype.storage_bytes(numel);
2173 if let Some(elem) = self.seq_elem_values.get(&vid) {
2178 return Ok(elem.data[..n.min(elem.data.len())].to_vec());
2179 }
2180 if let Some(view) = self.views.get(&vid) {
2181 let buf = self.buffers.get(&view.source).ok_or_else(|| {
2182 SessionError::Internal(format!(
2183 "view value#{} aliases missing source buffer value#{}",
2184 vid.0, view.source.0
2185 ))
2186 })?;
2187 let esize = dtype.byte_size();
2188 if esize == 0 {
2189 return Err(SessionError::Internal(format!(
2192 "cannot materialize sub-byte view value#{}",
2193 vid.0
2194 )));
2195 }
2196 Ok(gather_view(
2197 host_bytes(buf),
2198 &view.shape,
2199 &view.strides,
2200 view.byte_offset,
2201 esize,
2202 ))
2203 } else {
2204 let buf = self.buffers.get(&vid).ok_or_else(|| {
2205 SessionError::Internal(format!("value#{} not produced", vid.0))
2206 })?;
2207 Ok(host_bytes(buf)[..n].to_vec())
2208 }
2209 }
2210
2211 fn store_output_tensor(
2218 &mut self,
2219 vid: ValueId,
2220 tensor: &Tensor,
2221 resolved: &mut HashMap<ValueId, Vec<usize>>,
2222 ) -> Result<()> {
2223 self.store_output_bytes(
2224 vid,
2225 tensor.dtype,
2226 tensor.shape.clone(),
2227 tensor.as_bytes(),
2228 resolved,
2229 )
2230 }
2231
2232 fn store_output_bytes(
2233 &mut self,
2234 vid: ValueId,
2235 dtype: DataType,
2236 dims: Vec<usize>,
2237 bytes: &[u8],
2238 resolved: &mut HashMap<ValueId, Vec<usize>>,
2239 ) -> Result<()> {
2240 let numel = checked_numel(&dims, || format!("value#{}", vid.0))?;
2241 let need = checked_storage_bytes(dtype, numel, || format!("value#{}", vid.0), &dims)?
2242 .max(1);
2243 let fits = self.buffers.get(&vid).map(|b| b.len() == need).unwrap_or(false);
2244 if !fits {
2245 if let Some(old) = self.buffers.remove(&vid) {
2246 self.ep.deallocate(old)?;
2247 }
2248 let buf = self
2249 .ep
2250 .allocate(need, TensorLayout::contiguous().alignment)?;
2251 self.buffers.insert(vid, buf);
2252 }
2253 let buf = self.buffers.get_mut(&vid).expect("just ensured");
2254 write_host(buf, bytes)?;
2255 self.value_dtypes.insert(vid, dtype);
2256 self.buffer_shapes.insert(vid, dims.clone());
2257 resolved.insert(vid, dims);
2258 Ok(())
2259 }
2260
2261 fn prepare_subgraph(
2266 &self,
2267 node_id: NodeId,
2268 attr_key: &str,
2269 resolved: &HashMap<ValueId, Vec<usize>>,
2270 outer_scope: &HashMap<String, Tensor>,
2271 ) -> Result<PreparedSubgraph> {
2272 let key = (node_id, attr_key.to_string());
2273 let body = self.graph.subgraphs.get(&key).ok_or_else(|| {
2274 SessionError::Internal(format!(
2275 "control-flow node #{} references missing subgraph '{attr_key}'",
2276 node_id.0
2277 ))
2278 })?;
2279
2280 let formal_names: Vec<String> = body
2281 .inputs
2282 .iter()
2283 .map(|&vid| {
2284 body.value(vid)
2285 .name
2286 .clone()
2287 .unwrap_or_else(|| format!("value#{}", vid.0))
2288 })
2289 .collect();
2290 let formal_set: HashSet<ValueId> = body.inputs.iter().copied().collect();
2291 let mut capture_names = Vec::new();
2292 for (vid, value) in body.values.iter() {
2293 if value.producer.is_none()
2294 && !formal_set.contains(&vid)
2295 && !body.initializers.contains_key(&vid)
2296 && let Some(name) = &value.name
2297 {
2298 capture_names.push(name.clone());
2299 }
2300 }
2301 capture_names.sort();
2302
2303 let mut scope_names = required_outer_names(body);
2304 scope_names.extend(capture_names.iter().cloned());
2305 let mut captures = HashMap::with_capacity(scope_names.len());
2306 for name in scope_names {
2307 let tensor = if let Some(&vid) = self.name_index.get(&name) {
2308 let materialized = self.buffers.contains_key(&vid)
2309 || self.views.contains_key(&vid)
2310 || self.seq_elem_values.contains_key(&vid);
2311 if resolved.contains_key(&vid) && materialized {
2312 self.value_tensor(vid, resolved)?
2313 } else {
2314 outer_scope.get(&name).cloned().ok_or_else(|| {
2315 missing_capture_error(attr_key, &name)
2316 })?
2317 }
2318 } else {
2319 outer_scope.get(&name).cloned().ok_or_else(|| {
2320 missing_capture_error(attr_key, &name)
2321 })?
2322 };
2323 captures.insert(name, tensor);
2324 }
2325
2326 Ok(PreparedSubgraph {
2327 key,
2328 formal_names,
2329 capture_names,
2330 captures,
2331 })
2332 }
2333
2334 fn build_subgraph_exec(
2339 &self,
2340 prepared: &PreparedSubgraph,
2341 externals: &[&Tensor],
2342 ) -> Result<CompiledSubgraph> {
2343 let key = &prepared.key;
2344 let body = self.graph.subgraphs.get(key).ok_or_else(|| {
2345 SessionError::Internal(format!(
2346 "control-flow node #{} has no registered subgraph '{}'",
2347 key.0 .0, key.1
2348 ))
2349 })?;
2350 let mut g = body.clone();
2351
2352 let mut body_names: HashMap<String, ValueId> = HashMap::new();
2354 for (vid, value) in g.values.iter() {
2355 if let Some(n) = &value.name {
2356 body_names.insert(n.clone(), vid);
2357 }
2358 }
2359
2360 for cname in &prepared.capture_names {
2362 let vid = *body_names.get(cname).ok_or_else(|| {
2363 SessionError::Internal(format!(
2364 "control-flow body '{}' lost capture value '{cname}'",
2365 key.1
2366 ))
2367 })?;
2368 if !g.inputs.contains(&vid) {
2369 g.add_input(vid);
2370 }
2371 }
2372
2373 let all_names = prepared
2376 .formal_names
2377 .iter()
2378 .chain(prepared.capture_names.iter());
2379 for (name, tensor) in all_names.zip(externals.iter()) {
2380 let vid = *body_names.get(name).ok_or_else(|| {
2381 SessionError::Internal(format!(
2382 "control-flow body '{}' missing formal/captured input '{name}'",
2383 key.1
2384 ))
2385 })?;
2386 let v = g.value_mut(vid);
2387 v.dtype = tensor.dtype;
2388 v.shape = tensor.shape.iter().map(|&d| Dim::Static(d)).collect();
2389 }
2390
2391 let registry = InferenceRegistry::default_registry();
2395 let opset_imports = self.graph.opset_imports.clone();
2396 registry.infer_graph(&mut g, &opset_imports, MergePolicy::Permissive)?;
2397
2398 let exec = Executor::build(g, self.weights.clone(), self.ep.clone())?;
2399 Ok(CompiledSubgraph {
2400 exec,
2401 input_names: prepared
2402 .formal_names
2403 .iter()
2404 .chain(prepared.capture_names.iter())
2405 .cloned()
2406 .collect(),
2407 built_shapes: externals.iter().map(|t| t.shape.clone()).collect(),
2408 })
2409 }
2410
2411 fn run_subgraph(
2415 &mut self,
2416 prepared: &PreparedSubgraph,
2417 formal_inputs: &[&Tensor],
2418 ) -> Result<Vec<Tensor>> {
2419 if prepared.formal_names.len() != formal_inputs.len() {
2420 return Err(SessionError::Internal(format!(
2421 "control-flow body '{}' expects {} formal input(s) but {} were supplied",
2422 prepared.key.1,
2423 prepared.formal_names.len(),
2424 formal_inputs.len()
2425 )));
2426 }
2427
2428 let mut externals: Vec<&Tensor> =
2429 Vec::with_capacity(formal_inputs.len() + prepared.capture_names.len());
2430 externals.extend_from_slice(formal_inputs);
2431 for name in &prepared.capture_names {
2432 externals.push(prepared.captures.get(name).expect("prepared capture must be present"));
2433 }
2434 let rebuild = match self.subgraph_execs.get(&prepared.key) {
2435 Some(cs) => {
2436 cs.built_shapes.len() != externals.len()
2437 || cs
2438 .built_shapes
2439 .iter()
2440 .zip(externals.iter())
2441 .any(|(built, tensor)| built != &tensor.shape)
2442 }
2443 None => true,
2444 };
2445 if rebuild {
2446 let child = self.build_subgraph_exec(prepared, &externals)?;
2447 self.subgraph_execs.insert(prepared.key.clone(), child);
2448 self.control_flow_stats.subgraph_builds += 1;
2449 }
2450
2451 self.control_flow_stats.subgraph_runs += 1;
2452 let cs = self.subgraph_execs.get_mut(&prepared.key).expect("child present");
2453 let inputs: Vec<(&str, &Tensor)> = cs
2454 .input_names
2455 .iter()
2456 .map(String::as_str)
2457 .zip(externals)
2458 .collect();
2459 cs.exec.run_scoped(&inputs, &prepared.captures)
2460 }
2461
2462 fn exec_control_flow(
2464 &mut self,
2465 pi: usize,
2466 resolved: &mut HashMap<ValueId, Vec<usize>>,
2467 outer_scope: &HashMap<String, Tensor>,
2468 ) -> Result<()> {
2469 let node = self.graph.node(self.plan[pi].node_id).clone();
2470 match node.op_type.as_str() {
2471 "If" => self.exec_if(&node, resolved, outer_scope),
2472 "Loop" => self.exec_loop(&node, resolved, outer_scope),
2473 "Scan" => self.exec_scan(&node, resolved, outer_scope),
2474 other => Err(SessionError::Internal(format!(
2475 "exec_control_flow reached non-control-flow op {other:?}"
2476 ))),
2477 }
2478 }
2479
2480 fn exec_if(
2483 &mut self,
2484 node: &Node,
2485 resolved: &mut HashMap<ValueId, Vec<usize>>,
2486 outer_scope: &HashMap<String, Tensor>,
2487 ) -> Result<()> {
2488 let cond_vid = node.inputs.first().and_then(|s| *s).ok_or_else(|| {
2489 SessionError::Internal("If node is missing its required 'cond' input".to_string())
2490 })?;
2491 let cond_t = self.value_tensor(cond_vid, resolved)?;
2492 let cond = tensor_scalar_bool(&cond_t).ok_or_else(|| SessionError::Internal(format!(
2493 "If: 'cond' must be a BOOL scalar, got dtype {:?} shape {:?}",
2494 cond_t.dtype, cond_t.shape
2495 )))?;
2496
2497 let attr_key = if cond { "then_branch" } else { "else_branch" };
2498 let prepared = self.prepare_subgraph(node.id, attr_key, resolved, outer_scope)?;
2499 let outs = self.run_subgraph(&prepared, &[])?;
2500
2501 if outs.len() != node.outputs.len() {
2502 return Err(SessionError::OutputShapeCountMismatch {
2503 op: format!("If/{attr_key}"),
2504 expected: node.outputs.len(),
2505 got: outs.len(),
2506 });
2507 }
2508 for (vid, t) in node.outputs.iter().zip(outs.iter()) {
2509 self.store_output_tensor(*vid, t, resolved)?;
2510 }
2511 Ok(())
2512 }
2513
2514 fn exec_loop(
2520 &mut self,
2521 node: &Node,
2522 resolved: &mut HashMap<ValueId, Vec<usize>>,
2523 outer_scope: &HashMap<String, Tensor>,
2524 ) -> Result<()> {
2525 let m: Option<i64> = match node.inputs.first().and_then(|s| *s) {
2528 Some(vid) => {
2529 let t = self.value_tensor(vid, resolved)?;
2530 let m = tensor_scalar_i64(&t).ok_or_else(|| SessionError::Internal(format!(
2531 "Loop: trip-count 'M' must be an INT64/INT32 scalar, got dtype {:?}",
2532 t.dtype
2533 )))?;
2534 Some(m)
2535 }
2536 None => None,
2537 };
2538 let mut cond: Option<bool> = match node.inputs.get(1).and_then(|s| *s) {
2539 Some(vid) => {
2540 let t = self.value_tensor(vid, resolved)?;
2541 Some(tensor_scalar_bool(&t).ok_or_else(|| SessionError::Internal(format!(
2542 "Loop: 'cond' must be a BOOL scalar, got dtype {:?}",
2543 t.dtype
2544 )))?)
2545 }
2546 None => None,
2547 };
2548
2549 let mut carried: Vec<Tensor> = Vec::new();
2551 for slot in node.inputs.iter().skip(2) {
2552 let vid = slot.ok_or_else(|| SessionError::Internal(
2553 "Loop: an interior loop-carried input is omitted (empty), which ONNX does not \
2554 allow — every v_initial must be provided".to_string(),
2555 ))?;
2556 carried.push(self.value_tensor(vid, resolved)?);
2557 }
2558 let num_carried = carried.len();
2559 let num_outputs = node.outputs.len();
2562 if num_outputs < num_carried {
2563 return Err(SessionError::Internal(format!(
2564 "Loop: node declares {num_outputs} output(s) but has {num_carried} loop-carried \
2565 dependency(ies); outputs must be carried-finals followed by scan-outputs"
2566 )));
2567 }
2568 let num_scan = num_outputs - num_carried;
2569 let expected_iterations = m.and_then(|n| usize::try_from(n).ok());
2570 let mut scan_acc: Vec<TensorStackAccumulator> = (0..num_scan)
2571 .map(|_| TensorStackAccumulator::new(expected_iterations))
2572 .collect();
2573 let prepared = self.prepare_subgraph(node.id, "body", resolved, outer_scope)?;
2574 let mut iter_tensor = scalar_i64_tensor(0)?;
2575 let mut cond_tensor = scalar_bool_tensor(cond.unwrap_or(true))?;
2576
2577 let mut iter: i64 = 0;
2578 loop {
2579 if let Some(m) = m
2580 && iter >= m
2581 {
2582 break;
2583 }
2584 if cond == Some(false) {
2585 break;
2586 }
2587
2588 iter_tensor.overwrite_bytes(&iter.to_le_bytes())?;
2589 cond_tensor.overwrite_bytes(&[u8::from(cond.unwrap_or(true))])?;
2590 let mut formal: Vec<&Tensor> = Vec::with_capacity(2 + num_carried);
2591 formal.push(&iter_tensor);
2592 formal.push(&cond_tensor);
2593 formal.extend(carried.iter());
2594
2595 let outs = self.run_subgraph(&prepared, &formal)?;
2596 drop(formal);
2597 let expected = 1 + num_carried + num_scan;
2599 if outs.len() != expected {
2600 return Err(SessionError::OutputShapeCountMismatch {
2601 op: "Loop/body".to_string(),
2602 expected,
2603 got: outs.len(),
2604 });
2605 }
2606 let mut it = outs.into_iter();
2607 let cond_out = it.next().expect("cond_out present");
2608 cond = Some(tensor_scalar_bool(&cond_out).ok_or_else(|| SessionError::Internal(
2609 format!(
2610 "Loop: body's first output 'cond_out' must be a BOOL scalar, got dtype {:?}",
2611 cond_out.dtype
2612 ),
2613 ))?);
2614 carried.clear();
2615 carried.extend((&mut it).take(num_carried));
2616 for acc in scan_acc.iter_mut() {
2617 acc.push(it.next().expect("scan output present"))?;
2618 }
2619
2620 iter += 1;
2621 }
2622
2623 for (i, t) in carried.iter().enumerate() {
2625 self.store_output_tensor(node.outputs[i], t, resolved)?;
2626 }
2627 for (s, acc) in scan_acc.into_iter().enumerate() {
2628 let (dtype, shape, bytes) = acc.finish();
2629 self.store_output_bytes(
2630 node.outputs[num_carried + s],
2631 dtype,
2632 shape,
2633 &bytes,
2634 resolved,
2635 )?;
2636 }
2637 Ok(())
2638 }
2639
2640 fn exec_scan(
2647 &mut self,
2648 node: &Node,
2649 resolved: &mut HashMap<ValueId, Vec<usize>>,
2650 outer_scope: &HashMap<String, Tensor>,
2651 ) -> Result<()> {
2652 let num_scan_inputs = node
2653 .attr("num_scan_inputs")
2654 .and_then(|a| a.as_int())
2655 .ok_or_else(|| SessionError::Internal(
2656 "Scan: required attribute 'num_scan_inputs' is missing or not an INT".to_string(),
2657 ))? as usize;
2658
2659 for attr in ["scan_input_axes", "scan_output_axes"] {
2662 if let Some(a) = node.attr(attr)
2663 && let Some(axes) = a.as_ints()
2664 && axes.iter().any(|&ax| ax != 0)
2665 {
2666 return Err(SessionError::Internal(format!(
2667 "Scan: attribute '{attr}' = {axes:?} requests a non-zero scan axis, \
2668 which this runtime does not yet support. Expected axis 0 for every \
2669 scan input/output; re-export with axis 0 or wait for full Scan-axis \
2670 support"
2671 )));
2672 }
2673 }
2674 for attr in ["scan_input_directions", "scan_output_directions"] {
2675 if let Some(a) = node.attr(attr)
2676 && let Some(dirs) = a.as_ints()
2677 && dirs.iter().any(|&d| d != 0)
2678 {
2679 return Err(SessionError::Internal(format!(
2680 "Scan: attribute '{attr}' = {dirs:?} requests reverse iteration, which \
2681 this runtime does not yet support (forward only). Re-export forward or \
2682 wait for reverse-Scan support"
2683 )));
2684 }
2685 }
2686
2687 let total_inputs = node.inputs.len();
2688 if total_inputs < num_scan_inputs {
2689 return Err(SessionError::Internal(format!(
2690 "Scan: node has {total_inputs} input(s) but num_scan_inputs={num_scan_inputs}"
2691 )));
2692 }
2693 let num_state = total_inputs - num_scan_inputs;
2694
2695 let mut state: Vec<Tensor> = Vec::with_capacity(num_state);
2697 for slot in node.inputs.iter().take(num_state) {
2698 let vid = slot.ok_or_else(|| SessionError::Internal(
2699 "Scan: an initial-state input is omitted (empty), which ONNX does not allow"
2700 .to_string(),
2701 ))?;
2702 state.push(self.value_tensor(vid, resolved)?);
2703 }
2704 let mut scan_inputs: Vec<Tensor> = Vec::with_capacity(num_scan_inputs);
2705 for slot in node.inputs.iter().skip(num_state) {
2706 let vid = slot.ok_or_else(|| SessionError::Internal(
2707 "Scan: a scan input is omitted (empty), which ONNX does not allow".to_string(),
2708 ))?;
2709 scan_inputs.push(self.value_tensor(vid, resolved)?);
2710 }
2711
2712 let seq_len = scan_inputs
2714 .first()
2715 .and_then(|t| t.shape.first().copied())
2716 .ok_or_else(|| SessionError::Internal(
2717 "Scan: requires at least one scan input with rank >= 1".to_string(),
2718 ))?;
2719 for (i, t) in scan_inputs.iter().enumerate() {
2720 let this = t.shape.first().copied().unwrap_or(0);
2721 if this != seq_len {
2722 return Err(SessionError::Internal(format!(
2723 "Scan: scan input #{i} has scan-axis length {this} but the first scan input has \
2724 {seq_len}; all scan inputs must share the same scan-axis length"
2725 )));
2726 }
2727 }
2728
2729 let num_outputs = node.outputs.len();
2730 if num_outputs < num_state {
2731 return Err(SessionError::Internal(format!(
2732 "Scan: declares {num_outputs} output(s) but has {num_state} state variable(s); \
2733 outputs must be final-state followed by scan-outputs"
2734 )));
2735 }
2736 let num_scan_out = num_outputs - num_state;
2737 let mut scan_acc: Vec<TensorStackAccumulator> = (0..num_scan_out)
2738 .map(|_| TensorStackAccumulator::new(Some(seq_len)))
2739 .collect();
2740 let prepared = self.prepare_subgraph(node.id, "body", resolved, outer_scope)?;
2741 let mut scan_slices = Vec::with_capacity(num_scan_inputs);
2742 if seq_len != 0 {
2743 for t in &scan_inputs {
2744 let (shape, bytes) = leading_slice(t, 0)?;
2745 scan_slices.push(Tensor::from_raw_in(
2746 self.ep.clone(),
2747 t.dtype,
2748 shape,
2749 bytes,
2750 )?);
2751 }
2752 }
2753 for step in 0..seq_len {
2754 if step != 0 {
2755 for (source, slice) in scan_inputs.iter().zip(scan_slices.iter_mut()) {
2756 let (_, bytes) = leading_slice(source, step)?;
2757 slice.overwrite_bytes(bytes)?;
2758 }
2759 }
2760 let mut formal: Vec<&Tensor> = Vec::with_capacity(num_state + num_scan_inputs);
2761 formal.extend(state.iter());
2762 formal.extend(scan_slices.iter());
2763
2764 let outs = self.run_subgraph(&prepared, &formal)?;
2765 drop(formal);
2766 let expected = num_state + num_scan_out;
2767 if outs.len() != expected {
2768 return Err(SessionError::OutputShapeCountMismatch {
2769 op: "Scan/body".to_string(),
2770 expected,
2771 got: outs.len(),
2772 });
2773 }
2774 let mut it = outs.into_iter();
2775 state.clear();
2776 state.extend((&mut it).take(num_state));
2777 for acc in scan_acc.iter_mut() {
2778 acc.push(it.next().expect("scan output present"))?;
2779 }
2780 }
2781
2782 for (i, t) in state.iter().enumerate() {
2783 self.store_output_tensor(node.outputs[i], t, resolved)?;
2784 }
2785 for (s, acc) in scan_acc.into_iter().enumerate() {
2786 let (dtype, shape, bytes) = acc.finish();
2787 self.store_output_bytes(node.outputs[num_state + s], dtype, shape, &bytes, resolved)?;
2788 }
2789 Ok(())
2790 }
2791}
2792
2793fn leading_slice(t: &Tensor, index: usize) -> Result<(Vec<usize>, &[u8])> {
2796 if t.shape.is_empty() {
2797 return Err(SessionError::Internal(
2798 "Scan: cannot slice a scalar scan input along axis 0".to_string(),
2799 ));
2800 }
2801 let outer = t.shape[0];
2802 if index >= outer {
2803 return Err(SessionError::Internal(format!(
2804 "Scan: slice index {index} out of range for scan-axis length {outer}"
2805 )));
2806 }
2807 let inner_shape = t.shape[1..].to_vec();
2808 let inner_numel: usize = inner_shape.iter().product();
2809 let esize = t.dtype.byte_size();
2810 if esize == 0 {
2812 return Err(SessionError::Internal(format!(
2813 "Scan: sub-byte dtype {:?} scan inputs are not supported",
2814 t.dtype
2815 )));
2816 }
2817 let slice_bytes = inner_numel * esize;
2818 let start = index * slice_bytes;
2819 let bytes = &t.as_bytes()[start..start + slice_bytes];
2820 Ok((inner_shape, bytes))
2821}
2822
2823struct TensorStackAccumulator {
2828 expected_len: Option<usize>,
2829 dtype: Option<DataType>,
2830 elem_shape: Vec<usize>,
2831 len: usize,
2832 bytes: Vec<u8>,
2833}
2834
2835impl TensorStackAccumulator {
2836 fn new(expected_len: Option<usize>) -> Self {
2837 Self {
2838 expected_len,
2839 dtype: None,
2840 elem_shape: Vec::new(),
2841 len: 0,
2842 bytes: Vec::new(),
2843 }
2844 }
2845
2846 fn push(&mut self, tensor: Tensor) -> Result<()> {
2847 if let Some(dtype) = self.dtype {
2848 if tensor.shape != self.elem_shape || tensor.dtype != dtype {
2849 return Err(SessionError::Internal(format!(
2850 "Loop/Scan: scan output slice {} has shape {:?} dtype {:?} but the first slice \
2851 is shape {:?} dtype {:?}; every iteration's scan output must match",
2852 self.len, tensor.shape, tensor.dtype, self.elem_shape, dtype
2853 )));
2854 }
2855 } else {
2856 if tensor.dtype.byte_size() == 0 {
2857 return Err(SessionError::Internal(format!(
2858 "Loop/Scan: sub-byte dtype {:?} scan outputs are not supported",
2859 tensor.dtype
2860 )));
2861 }
2862 self.dtype = Some(tensor.dtype);
2863 self.elem_shape = tensor.shape.clone();
2864 if let Some(expected) = self.expected_len {
2865 self.bytes.reserve(expected.saturating_mul(tensor.as_bytes().len()));
2866 }
2867 }
2868 self.bytes.extend_from_slice(tensor.as_bytes());
2869 self.len += 1;
2870 Ok(())
2871 }
2872
2873 fn finish(self) -> (DataType, Vec<usize>, Vec<u8>) {
2874 if self.len == 0 {
2875 return (DataType::Float32, vec![0], Vec::new());
2876 }
2877 let dtype = self.dtype.expect("non-empty accumulator has dtype");
2878 let mut shape = Vec::with_capacity(1 + self.elem_shape.len());
2879 shape.push(self.len);
2880 shape.extend(self.elem_shape);
2881 (dtype, shape, self.bytes)
2882 }
2883}
2884
2885impl Drop for Executor {
2886 fn drop(&mut self) {
2887 for (_, buf) in self.buffers.drain() {
2889 let _ = self.ep.deallocate(buf);
2890 }
2891 }
2892}
2893
2894pub(crate) fn auto_detect_cpu_ep() -> Result<Arc<CpuExecutionProvider>> {
2898 let mut ep = CpuExecutionProvider::new();
2899 ep.initialize(&Default::default())?;
2900 Ok(Arc::new(ep))
2901}
2902
2903#[cfg(test)]
2904mod tests {
2905 use super::*;
2906
2907 #[test]
2911 fn view_bounds_rejects_out_of_bounds_view() {
2912 let shape = [2usize, 3];
2914 let strides = compute_contiguous_strides(&shape);
2915 let err = view_bounds(&shape, &strides, 0, DataType::Float32, 16);
2916 assert!(err.is_err(), "gate must reject an oversized view");
2917
2918 assert!(view_bounds(&shape, &strides, 0, DataType::Float32, 24).is_ok());
2920 }
2921
2922 #[test]
2925 fn view_bounds_rejects_offset_overrun() {
2926 let shape = [4usize];
2927 let strides = compute_contiguous_strides(&shape);
2928 assert!(view_bounds(&shape, &strides, 8, DataType::Float32, 16).is_err());
2930 assert!(view_bounds(&shape, &strides, 0, DataType::Float32, 16).is_ok());
2931 }
2932
2933 #[test]
2936 fn substitute_resolves_bound_symbols_only() {
2937 let mut bindings = HashMap::new();
2938 bindings.insert(SymbolId(0), 7usize);
2939 let shape = vec![Dim::Symbolic(SymbolId(0)), Dim::Static(4)];
2940 assert_eq!(substitute(&shape, &bindings), Some(vec![7, 4]));
2941
2942 let unbound = vec![Dim::Symbolic(SymbolId(1)), Dim::Static(4)];
2943 assert_eq!(substitute(&unbound, &bindings), None);
2944 }
2945
2946 #[test]
2950 fn checked_numel_detects_overflow() {
2951 assert_eq!(checked_numel(&[2, 3, 4], || "v".into()).unwrap(), 24);
2953 assert_eq!(checked_numel(&[], || "v".into()).unwrap(), 1);
2954
2955 let huge = [usize::MAX, 2];
2957 let err = checked_numel(&huge, || "value#9".into());
2958 assert!(matches!(
2959 err,
2960 Err(SessionError::ShapeOverflow { .. })
2961 ));
2962 }
2963
2964 #[test]
2968 fn checked_storage_bytes_detects_byte_overflow() {
2969 let numel = usize::MAX / 4;
2972 let err = checked_storage_bytes(DataType::Float64, numel, || "value#9".into(), &[numel]);
2973 assert!(matches!(err, Err(SessionError::ShapeOverflow { .. })));
2974
2975 assert_eq!(
2977 checked_storage_bytes(DataType::Float32, 4, || "v".into(), &[4]).unwrap(),
2978 16
2979 );
2980 }
2981
2982 #[test]
2987 fn dynamic_output_shapes_slice_is_single_output() {
2988 let node = Node::new(NodeId(0), "Slice", vec![], vec![]);
2989 let input_shapes = vec![vec![4usize, 2]];
2990 let input_values = vec![
2991 None, Some(vec![1]), Some(vec![3]), Some(vec![0]), Some(vec![1]), ];
2997 let out = dynamic_output_shapes(&node, &input_shapes, &input_values).unwrap();
2998 assert_eq!(out.len(), 1, "Slice must resolve exactly one output shape");
2999 assert_eq!(out[0], vec![2, 2]);
3000
3001 let other = Node::new(NodeId(1), "Conv", vec![], vec![]);
3003 assert!(dynamic_output_shapes(&other, &input_shapes, &input_values).is_none());
3004 }
3005
3006 #[test]
3009 fn effective_opset_reads_graph_import() {
3010 let mut graph = Graph::default();
3011 graph.opset_imports.insert(String::new(), 12);
3012 let node = Node::new(NodeId(0), "Softmax", vec![], vec![]);
3013 assert_eq!(effective_opset(&graph, &node), 12);
3014
3015 graph.opset_imports.insert(String::new(), 0);
3016 assert_eq!(effective_opset(&graph, &node), 0);
3017
3018 }
3019
3020 #[test]
3021 #[should_panic(expected = "internal invariant violated")]
3022 fn effective_opset_requires_validated_import() {
3023 effective_opset(
3024 &Graph::default(),
3025 &Node::new(NodeId(0), "Softmax", vec![], vec![]),
3026 );
3027 }
3028
3029 use onnx_runtime_ir::{static_shape, WeightRef};
3032 use std::path::PathBuf;
3033
3034 fn weightstream_tmp_dir() -> PathBuf {
3036 let dir = PathBuf::from(concat!(
3037 env!("CARGO_MANIFEST_DIR"),
3038 "/../../target/weightstream_test"
3039 ));
3040 std::fs::create_dir_all(&dir).expect("create weight-streaming test dir");
3041 dir
3042 }
3043
3044 fn f32_le(data: &[f32]) -> Vec<u8> {
3045 data.iter().flat_map(|v| v.to_le_bytes()).collect()
3046 }
3047
3048 #[test]
3052 fn aligned_external_initializer_is_borrowed_zero_copy() {
3053 let align = TensorLayout::contiguous().alignment;
3054 let path = weightstream_tmp_dir().join("aligned_init.bin");
3055 let w_data = [1.0f32, 2.0, 3.0, 4.0];
3056 std::fs::write(&path, f32_le(&w_data)).unwrap();
3057
3058 let mut store = WeightStore::new();
3059 store.map_external(&path).unwrap();
3060
3061 let mut g = Graph::new();
3062 g.opset_imports.insert(String::new(), 17);
3063 let w = g.create_named_value("W", DataType::Float32, static_shape([4]));
3064 g.set_initializer(
3065 w,
3066 WeightRef::External {
3067 path: path.clone(),
3068 offset: 0, length: 16,
3070 dtype: DataType::Float32,
3071 dims: vec![4],
3072 },
3073 );
3074 let y = g.create_value(DataType::Float32, static_shape([4]));
3075 g.insert_node(Node::new(NodeId(0), "Relu", vec![Some(w)], vec![y]));
3076 g.add_output(y);
3077
3078 let ep = auto_detect_cpu_ep().unwrap();
3079 let exec = Executor::build(g, Arc::new(store), ep).unwrap();
3080
3081 let weight = &exec.graph.initializers[&w];
3082 let src = exec.weights().bytes(weight).unwrap();
3083 assert!(
3084 (src.as_ptr() as usize).is_multiple_of(align),
3085 "mmap window must be aligned for this test to exercise the zero-copy path"
3086 );
3087 let buf = &exec.buffers[&w];
3088 assert!(buf.is_borrowed(), "aligned initializer must be borrowed, not copied");
3089 assert_eq!(
3090 buf.as_ptr() as *const u8,
3091 src.as_ptr(),
3092 "zero-copy: the buffer must alias the mmap bytes (no copy)"
3093 );
3094
3095 let _ = std::fs::remove_file(&path);
3096 }
3097
3098 #[test]
3101 fn unaligned_external_initializer_falls_back_to_owned_copy() {
3102 let align = TensorLayout::contiguous().alignment;
3103 let path = weightstream_tmp_dir().join("unaligned_init.bin");
3104 let offset = 8usize;
3107 let w_data = [5.0f32, 6.0, 7.0, 8.0];
3108 let mut file = vec![0u8; offset];
3109 file.extend_from_slice(&f32_le(&w_data));
3110 std::fs::write(&path, &file).unwrap();
3111
3112 let mut store = WeightStore::new();
3113 store.map_external(&path).unwrap();
3114
3115 let mut g = Graph::new();
3116 g.opset_imports.insert(String::new(), 17);
3117 let w = g.create_named_value("W", DataType::Float32, static_shape([4]));
3118 g.set_initializer(
3119 w,
3120 WeightRef::External {
3121 path: path.clone(),
3122 offset,
3123 length: 16,
3124 dtype: DataType::Float32,
3125 dims: vec![4],
3126 },
3127 );
3128 let x = g.create_named_value("X", DataType::Float32, static_shape([4]));
3129 g.add_input(x);
3130 let y = g.create_value(DataType::Float32, static_shape([4]));
3131 g.insert_node(Node::new(NodeId(0), "Add", vec![Some(x), Some(w)], vec![y]));
3132 g.add_output(y);
3133
3134 let ep = auto_detect_cpu_ep().unwrap();
3135 let mut exec = Executor::build(g, Arc::new(store), ep).unwrap();
3136
3137 let weight = &exec.graph.initializers[&w];
3138 let src = exec.weights().bytes(weight).unwrap();
3139 assert!(
3140 !(src.as_ptr() as usize).is_multiple_of(align),
3141 "window must be unaligned for this test to exercise the fallback"
3142 );
3143 let buf = &exec.buffers[&w];
3144 assert!(
3145 !buf.is_borrowed(),
3146 "unaligned initializer must fall back to an owned copy"
3147 );
3148 assert_ne!(
3149 buf.as_ptr() as *const u8,
3150 src.as_ptr(),
3151 "fallback: the buffer must be a fresh copy, not an alias"
3152 );
3153
3154 let x_tensor = Tensor::from_f32(&[4], &[10.0, 20.0, 30.0, 40.0]).unwrap();
3156 let out = exec.run(&[("X", &x_tensor)]).unwrap();
3157 assert_eq!(out.len(), 1);
3158 let got = out[0].to_vec_f32();
3159 let want = [15.0f32, 26.0, 37.0, 48.0];
3160 assert_eq!(got.len(), want.len());
3161 for (g, w) in got.iter().zip(want.iter()) {
3162 assert!((g - w).abs() < 1e-5, "got {g}, want {w}");
3163 }
3164
3165 let _ = std::fs::remove_file(&path);
3166 }
3167
3168 #[test]
3176 fn producer_backed_initializer_is_not_borrowed() {
3177 let align = TensorLayout::contiguous().alignment;
3178 let path = weightstream_tmp_dir().join("producer_backed_init.bin");
3179 let w_data = [1.0f32, 2.0, 3.0, 4.0];
3180 std::fs::write(&path, f32_le(&w_data)).unwrap();
3181
3182 let mut store = WeightStore::new();
3183 store.map_external(&path).unwrap();
3184
3185 let mut g = Graph::new();
3186 g.opset_imports.insert(String::new(), 17);
3187 let x = g.create_named_value("X", DataType::Float32, static_shape([4]));
3188 g.add_input(x);
3189 let w = g.create_named_value("W", DataType::Float32, static_shape([4]));
3190 g.set_initializer(
3191 w,
3192 WeightRef::External {
3193 path: path.clone(),
3194 offset: 0, length: 16,
3196 dtype: DataType::Float32,
3197 dims: vec![4],
3198 },
3199 );
3200 g.insert_node(Node::new(NodeId(0), "Identity", vec![Some(x)], vec![w]));
3203 let y = g.create_value(DataType::Float32, static_shape([4]));
3204 g.insert_node(Node::new(NodeId(1), "Add", vec![Some(x), Some(w)], vec![y]));
3205 g.add_output(y);
3206
3207 assert!(
3208 g.value(w).producer.is_some(),
3209 "test setup: initializer value must have a producer",
3210 );
3211
3212 let ep = auto_detect_cpu_ep().unwrap();
3213 let exec = Executor::build(g, Arc::new(store), ep).unwrap();
3214
3215 let weight = &exec.graph.initializers[&w];
3216 let src = exec.weights().bytes(weight).unwrap();
3217 assert!(
3218 (src.as_ptr() as usize).is_multiple_of(align),
3219 "mmap window must be aligned so only the producer guard prevents borrowing",
3220 );
3221 let buf = &exec.buffers[&w];
3222 assert!(
3223 !buf.is_borrowed(),
3224 "producer-backed initializer must fall back to an owned writable copy",
3225 );
3226 assert_ne!(
3227 buf.as_ptr() as *const u8,
3228 src.as_ptr(),
3229 "producer-backed initializer must not alias read-only mmap bytes",
3230 );
3231
3232 let _ = std::fs::remove_file(&path);
3233 }
3234}