1use std::collections::{BinaryHeap, HashMap, HashSet};
5
6use crate::arena::Arena;
7use crate::dtype::DataType;
8use crate::error::GraphError;
9use crate::node::{Node, NodeId};
10use crate::shape::{Shape, SymbolConstraints, SymbolId};
11use crate::tensor::WeightRef;
12use crate::value::{Value, ValueId};
13
14#[derive(Clone, Debug, Default)]
20pub struct Graph {
21 pub nodes: Arena<NodeId, Node>,
22 pub values: Arena<ValueId, Value>,
23 pub inputs: Vec<ValueId>,
25 pub outputs: Vec<ValueId>,
27 pub initializers: HashMap<ValueId, WeightRef>,
29 pub symbol_constraints: HashMap<SymbolId, SymbolConstraints>,
31 pub opset_imports: HashMap<String, u64>,
33 pub subgraphs: HashMap<(NodeId, String), Graph>,
35
36 next_symbol: u32,
37 symbol_names: HashMap<String, SymbolId>,
38 unknown_value_types: HashSet<ValueId>,
39 unknown_value_shapes: HashSet<ValueId>,
40}
41
42impl Graph {
43 pub fn new() -> Self {
45 Self::default()
46 }
47
48 pub fn node(&self, id: NodeId) -> &Node {
53 self.nodes.get(id).expect("node id not live in graph")
54 }
55
56 pub fn node_mut(&mut self, id: NodeId) -> &mut Node {
58 self.nodes.get_mut(id).expect("node id not live in graph")
59 }
60
61 pub fn try_node(&self, id: NodeId) -> Option<&Node> {
63 self.nodes.get(id)
64 }
65
66 pub fn value(&self, id: ValueId) -> &Value {
69 self.values.get(id).expect("value id not live in graph")
70 }
71
72 pub fn value_mut(&mut self, id: ValueId) -> &mut Value {
74 self.values.get_mut(id).expect("value id not live in graph")
75 }
76
77 pub fn try_value(&self, id: ValueId) -> Option<&Value> {
79 self.values.get(id)
80 }
81
82 pub fn uses(&self, value: ValueId) -> Vec<(NodeId, u32)> {
84 self.value(value).consumers.uses()
85 }
86
87 pub fn consumers(&self, value: ValueId) -> Vec<NodeId> {
89 self.value(value).consumers.nodes()
90 }
91
92 pub fn num_uses(&self, value: ValueId) -> usize {
94 self.value(value).consumers.len()
95 }
96
97 pub fn has_uses(&self, value: ValueId) -> bool {
99 !self.value(value).consumers.is_empty()
100 }
101
102 pub fn num_nodes(&self) -> usize {
104 self.nodes.len()
105 }
106
107 pub fn num_values(&self) -> usize {
109 self.values.len()
110 }
111
112 pub fn value_type_is_known(&self, id: ValueId) -> bool {
114 !self.unknown_value_types.contains(&id)
115 }
116
117 pub fn value_shape_is_known(&self, id: ValueId) -> bool {
119 !self.unknown_value_shapes.contains(&id)
120 }
121
122 pub fn mark_value_type_unknown(&mut self, id: ValueId) {
124 self.unknown_value_types.insert(id);
125 }
126
127 pub fn mark_value_shape_unknown(&mut self, id: ValueId) {
129 self.unknown_value_shapes.insert(id);
130 }
131
132 pub fn create_symbol(&mut self, name: Option<String>) -> SymbolId {
136 let id = SymbolId(self.next_symbol);
137 self.next_symbol += 1;
138 self.symbol_constraints
139 .insert(id, SymbolConstraints::new(id, name.clone()));
140 if let Some(n) = name {
141 self.symbol_names.insert(n, id);
142 }
143 id
144 }
145
146 pub fn intern_symbol(&mut self, name: &str) -> SymbolId {
149 if let Some(id) = self.symbol_names.get(name) {
150 return *id;
151 }
152 self.create_symbol(Some(name.to_string()))
153 }
154
155 pub fn create_value(&mut self, dtype: DataType, shape: Shape) -> ValueId {
159 self.values
160 .insert_with(|vid| Value::new(vid, dtype, shape.clone()))
161 }
162
163 pub fn create_named_value(
165 &mut self,
166 name: impl Into<String>,
167 dtype: DataType,
168 shape: Shape,
169 ) -> ValueId {
170 let id = self.create_value(dtype, shape);
171 self.value_mut(id).name = Some(name.into());
172 id
173 }
174
175 pub fn add_input(&mut self, value: ValueId) {
177 self.value_mut(value).is_graph_input = true;
178 self.inputs.push(value);
179 debug_assert!(self.value(value).is_graph_input);
180 }
181
182 pub fn add_output(&mut self, value: ValueId) {
184 self.value_mut(value).is_graph_output = true;
185 self.outputs.push(value);
186 debug_assert!(self.value(value).is_graph_output);
187 }
188
189 pub fn insert_output(&mut self, index: usize, value: ValueId) {
191 self.value_mut(value).is_graph_output = true;
192 self.outputs.insert(index, value);
193 debug_assert!(self.value(value).is_graph_output);
194 }
195
196 pub fn remove_input(&mut self, index: usize) -> ValueId {
198 let value = self.inputs.remove(index);
199 if !self.inputs.contains(&value) {
200 self.value_mut(value).is_graph_input = false;
201 }
202 debug_assert_eq!(
203 self.value(value).is_graph_input,
204 self.inputs.contains(&value)
205 );
206 value
207 }
208
209 pub fn remove_output(&mut self, index: usize) -> ValueId {
211 let value = self.outputs.remove(index);
212 if !self.outputs.contains(&value) {
213 self.value_mut(value).is_graph_output = false;
214 }
215 debug_assert_eq!(
216 self.value(value).is_graph_output,
217 self.outputs.contains(&value)
218 );
219 value
220 }
221
222 pub fn set_inputs(&mut self, inputs: Vec<ValueId>) {
224 for value in self.inputs.drain(..) {
225 if let Some(metadata) = self.values.get_mut(value) {
226 metadata.is_graph_input = false;
227 }
228 }
229 for &value in &inputs {
230 self.value_mut(value).is_graph_input = true;
231 }
232 self.inputs = inputs;
233 debug_assert!(
234 self.inputs
235 .iter()
236 .all(|&value| self.value(value).is_graph_input)
237 );
238 }
239
240 pub fn set_outputs(&mut self, outputs: Vec<ValueId>) {
242 for value in self.outputs.drain(..) {
243 if let Some(metadata) = self.values.get_mut(value) {
244 metadata.is_graph_output = false;
245 }
246 }
247 for &value in &outputs {
248 self.value_mut(value).is_graph_output = true;
249 }
250 self.outputs = outputs;
251 debug_assert!(
252 self.outputs
253 .iter()
254 .all(|&value| self.value(value).is_graph_output)
255 );
256 }
257
258 pub fn set_initializer(&mut self, value: ValueId, weight: WeightRef) {
260 self.initializers.insert(value, weight);
261 }
262
263 pub fn predecessors(&self, node: NodeId) -> Vec<NodeId> {
267 let mut out = Vec::new();
268 let mut seen = HashSet::new();
269 for v in self.node(node).input_values() {
270 if let Some(val) = self.values.get(v)
271 && let Some(prod) = val.producer
272 && seen.insert(prod)
273 {
274 out.push(prod);
275 }
276 }
277 out.sort_unstable_by_key(|node| node.0);
278 out
279 }
280
281 pub fn successors(&self, node: NodeId) -> Vec<NodeId> {
283 let mut out = Vec::new();
284 let mut seen = HashSet::new();
285 for &v in &self.node(node).outputs {
286 if let Some(val) = self.values.get(v) {
287 for c in val.consumers.nodes() {
288 if seen.insert(c) {
289 out.push(c);
290 }
291 }
292 }
293 }
294 out.sort_unstable_by_key(|node| node.0);
295 out
296 }
297
298 pub fn nodes_between(&self, inputs: &[ValueId], outputs: &[ValueId]) -> Vec<NodeId> {
303 let boundary: HashSet<ValueId> = inputs.iter().copied().collect();
304 let mut nodes = Vec::new();
305 let mut seen_nodes = HashSet::new();
306 let mut seen_values = HashSet::new();
307 let mut stack: Vec<ValueId> = outputs.to_vec();
308 while let Some(v) = stack.pop() {
309 if boundary.contains(&v) || !seen_values.insert(v) {
310 continue;
311 }
312 let Some(val) = self.values.get(v) else {
313 continue;
314 };
315 if let Some(prod) = val.producer {
316 if seen_nodes.insert(prod) {
317 nodes.push(prod);
318 }
319 for iv in self.node(prod).input_values() {
320 stack.push(iv);
321 }
322 }
323 }
324 nodes
325 }
326
327 pub fn topological_order(&self) -> Result<Vec<NodeId>, GraphError> {
332 const VACANT: usize = usize::MAX;
333 let mut in_degree = vec![VACANT; self.nodes.capacity()];
334 let mut adj = vec![Vec::<NodeId>::new(); self.nodes.capacity()];
335 for node in self.nodes.keys() {
336 in_degree[node.0 as usize] = 0;
337 }
338
339 for (nid, node) in self.nodes.iter() {
340 for v in node.input_values() {
341 if let Some(val) = self.values.get(v)
342 && let Some(prod) = val.producer
343 && self.nodes.contains(prod)
344 {
345 adj[prod.0 as usize].push(nid);
346 in_degree[nid.0 as usize] += 1;
347 }
348 }
349 }
350
351 let mut ready: BinaryHeap<std::cmp::Reverse<u32>> = in_degree
353 .iter()
354 .enumerate()
355 .filter(|(_, degree)| **degree == 0)
356 .map(|(raw, _)| std::cmp::Reverse(raw as u32))
357 .collect();
358
359 let mut order = Vec::with_capacity(self.nodes.len());
360 while let Some(std::cmp::Reverse(raw)) = ready.pop() {
361 let nid = NodeId(raw);
362 order.push(nid);
363 for &successor in &adj[raw as usize] {
364 let degree = &mut in_degree[successor.0 as usize];
365 *degree -= 1;
366 if *degree == 0 {
367 ready.push(std::cmp::Reverse(successor.0));
368 }
369 }
370 }
371
372 if order.len() != self.nodes.len() {
373 return Err(GraphError::CycleDetected);
374 }
375 Ok(order)
376 }
377
378 pub fn insert_node(&mut self, node: Node) -> NodeId {
383 let id = self.nodes.insert_with(|nid| {
384 let mut node = node;
385 node.id = nid;
386 node
387 });
388 self.connect_edges(id);
389 id
390 }
391
392 pub fn remove_node(&mut self, id: NodeId) {
395 if !self.nodes.contains(id) {
396 return;
397 }
398 self.disconnect_edges(id);
399 let outputs = self.node(id).outputs.clone();
400 self.nodes.remove(id);
401 for v in outputs {
402 self.gc_value_if_orphan(v);
403 }
404 }
405
406 pub fn remove_nodes(&mut self, ids: &[NodeId]) {
412 for &id in ids {
413 self.remove_node(id);
414 }
415 }
416
417 pub fn replace_node_groups(
427 &mut self,
428 groups: Vec<(Vec<NodeId>, Node)>,
429 graph_outputs: &HashSet<ValueId>,
430 ) -> Vec<NodeId> {
431 debug_assert_eq!(
432 graph_outputs,
433 &self.outputs.iter().copied().collect::<HashSet<_>>()
434 );
435 let mut removed_nodes = HashSet::new();
436 for (node_ids, _) in &groups {
437 assert!(
438 !node_ids.is_empty(),
439 "replace_node_groups: group must not be empty"
440 );
441 for &id in node_ids {
442 assert!(
443 self.nodes.contains(id),
444 "replace_node_groups: node id not live"
445 );
446 assert!(
447 removed_nodes.insert(id),
448 "replace_node_groups: groups must be disjoint"
449 );
450 }
451 }
452
453 let mut inserted = Vec::with_capacity(groups.len());
454 for (node_ids, replacement) in groups {
455 for id in node_ids {
456 self.remove_node(id);
457 }
458 inserted.push(self.insert_node(replacement));
459 }
460
461 inserted
462 }
463
464 pub fn replace_node(&mut self, old: NodeId, new: Node) -> NodeId {
470 assert!(self.nodes.contains(old), "replace_node: old id not live");
471 self.disconnect_edges(old);
472 {
473 let slot = self.nodes.get_mut(old).expect("old live");
474 let mut new = new;
475 new.id = old;
476 *slot = new;
477 }
478 self.connect_edges(old);
479 old
480 }
481
482 pub fn insert_on_edge(&mut self, value: ValueId, new_node: Node) -> NodeId {
488 let (dtype, shape) = {
489 let v = self.value(value);
490 (v.dtype, v.shape.clone())
491 };
492 let new_value = self.create_value(dtype, shape);
493 self.replace_all_uses(value, new_value);
496
497 let mut new_node = new_node;
498 new_node.inputs = vec![Some(value)];
499 new_node.outputs = vec![new_value];
500 self.insert_node(new_node)
501 }
502
503 pub fn replace_input(
508 &mut self,
509 node: NodeId,
510 input_index: usize,
511 new_value: Option<ValueId>,
512 ) -> Option<ValueId> {
513 assert!(self.nodes.contains(node), "replace_input: node id not live");
514 assert!(
515 input_index < self.node(node).inputs.len(),
516 "replace_input: input index out of bounds"
517 );
518 if let Some(value) = new_value {
519 assert!(
520 self.values.contains(value),
521 "replace_input: value id not live"
522 );
523 }
524
525 let old_value = self.node(node).inputs[input_index];
526 if old_value == new_value {
527 return old_value;
528 }
529 if let Some(value) = old_value {
530 let removed = self
531 .value_mut(value)
532 .consumers
533 .remove(node, input_index as u32);
534 debug_assert!(removed, "old input edge must be present");
535 }
536 self.node_mut(node).inputs[input_index] = new_value;
537 if let Some(value) = new_value {
538 self.value_mut(value)
539 .consumers
540 .insert(node, input_index as u32);
541 }
542 old_value
543 }
544
545 pub fn replace_all_uses(&mut self, old_value: ValueId, new_value: ValueId) {
548 if old_value == new_value {
549 return;
550 }
551 let uses = match self.values.get(old_value) {
552 Some(value) => value.consumers.uses(),
553 None => return,
554 };
555 for (node, input_index) in uses {
556 if self.nodes.contains(node) {
557 self.replace_input(node, input_index as usize, Some(new_value));
558 }
559 }
560 if self.value(old_value).is_graph_output {
561 let mut outputs = self.outputs.clone();
562 for output in &mut outputs {
563 if *output == old_value {
564 *output = new_value;
565 }
566 }
567 self.set_outputs(outputs);
568 }
569 }
570
571 pub fn validate(&self) -> Result<(), Vec<GraphError>> {
575 let mut errors = Vec::new();
576 let graph_inputs: HashSet<_> = self.inputs.iter().copied().collect();
577 let graph_outputs: HashSet<_> = self.outputs.iter().copied().collect();
578 for (value, metadata) in self.values.iter() {
579 debug_assert_eq!(
580 metadata.is_graph_input,
581 graph_inputs.contains(&value),
582 "graph-input membership flag drifted for {value:?}"
583 );
584 debug_assert_eq!(
585 metadata.is_graph_output,
586 graph_outputs.contains(&value),
587 "graph-output membership flag drifted for {value:?}"
588 );
589 }
590
591 let mut produced: HashMap<ValueId, NodeId> = HashMap::new();
593 for (nid, node) in self.nodes.iter() {
594 for (input_index, input) in node.inputs.iter().enumerate() {
595 if let Some(value) = input {
596 if !self.values.contains(*value) {
597 errors.push(GraphError::DanglingValue(*value));
598 } else if !self
599 .value(*value)
600 .consumers
601 .contains(nid, input_index as u32)
602 {
603 errors.push(GraphError::ConsumerLinkMismatch(*value));
604 }
605 }
606 }
607 for &v in &node.outputs {
608 if !self.values.contains(v) {
609 errors.push(GraphError::DanglingValue(v));
610 continue;
611 }
612 if produced.insert(v, nid).is_some() {
613 errors.push(GraphError::DuplicateOutput(v));
614 }
615 }
616 }
617
618 for (vid, val) in self.values.iter() {
620 if let Some(p) = val.producer {
621 if !self.nodes.contains(p) {
622 errors.push(GraphError::DanglingNode(p));
623 } else if !self.node(p).outputs.contains(&vid) {
624 errors.push(GraphError::ProducerLinkMismatch(vid));
625 }
626 }
627 for (consumer, input_index) in val.consumers.uses() {
628 if !self.nodes.contains(consumer) {
629 errors.push(GraphError::DanglingNode(consumer));
630 } else if self.node(consumer).inputs.get(input_index as usize) != Some(&Some(vid)) {
631 errors.push(GraphError::ConsumerLinkMismatch(vid));
632 }
633 }
634 }
635
636 for &inp in &self.inputs {
638 if let Some(val) = self.values.get(inp) {
639 if val.producer.is_some() {
640 errors.push(GraphError::InputHasProducer(inp));
641 }
642 debug_assert!(val.is_graph_input);
643 } else {
644 errors.push(GraphError::DanglingValue(inp));
645 }
646 }
647
648 for &out in &self.outputs {
651 match self.values.get(out) {
652 Some(val) => {
653 debug_assert!(val.is_graph_output);
654 let is_source = val.is_graph_input || self.initializers.contains_key(&out);
655 if val.producer.is_none() && !is_source {
656 errors.push(GraphError::MissingProducer(out));
657 }
658 }
659 None => errors.push(GraphError::DanglingValue(out)),
660 }
661 }
662
663 if let Err(e) = self.topological_order() {
665 errors.push(e);
666 }
667
668 for (domain, &version) in &self.opset_imports {
670 if version == 0 {
671 errors.push(GraphError::InvalidOpsetImport {
672 domain: domain.clone(),
673 version,
674 });
675 }
676 }
677
678 for sub in self.subgraphs.values() {
680 if let Err(mut sub_errors) = sub.validate() {
681 errors.append(&mut sub_errors);
682 }
683 }
684
685 if errors.is_empty() {
686 Ok(())
687 } else {
688 Err(errors)
689 }
690 }
691
692 fn connect_edges(&mut self, id: NodeId) {
696 let inputs = self.node(id).inputs.clone();
697 let outputs = self.node(id).outputs.clone();
698 for (input_index, input) in inputs.into_iter().enumerate() {
699 if let Some(value) = input
700 && let Some(metadata) = self.values.get_mut(value)
701 {
702 metadata.consumers.insert(id, input_index as u32);
703 }
704 }
705 for v in outputs {
706 if let Some(val) = self.values.get_mut(v) {
707 val.producer = Some(id);
708 }
709 }
710 }
711
712 fn disconnect_edges(&mut self, id: NodeId) {
715 let input_count = self.node(id).inputs.len();
716 let outputs = self.node(id).outputs.clone();
717 for input_index in 0..input_count {
718 self.replace_input(id, input_index, None);
719 }
720 for v in outputs {
721 if let Some(val) = self.values.get_mut(v)
722 && val.producer == Some(id)
723 {
724 val.producer = None;
725 }
726 }
727 }
728
729 fn gc_value_if_orphan(&mut self, value: ValueId) {
732 let orphan = match self.values.get(value) {
733 Some(v) => {
734 v.producer.is_none()
735 && v.consumers.is_empty()
736 && !v.is_graph_input
737 && !v.is_graph_output
738 && !self.initializers.contains_key(&value)
739 }
740 None => false,
741 };
742 if orphan {
743 self.values.remove(value);
744 self.unknown_value_types.remove(&value);
745 self.unknown_value_shapes.remove(&value);
746 }
747 }
748}
749
750#[cfg(test)]
751mod tests {
752 use super::*;
753 use crate::shape::static_shape;
754 use crate::tensor::TensorData;
755
756 fn assert_graphs_identical(
757 mut batched: Graph,
758 mut sequential: Graph,
759 node_probes: usize,
760 value_probes: usize,
761 trial: usize,
762 ) {
763 assert_eq!(
764 format!("{batched:#?}"),
765 format!("{sequential:#?}"),
766 "graph mismatch on randomized trial {trial}"
767 );
768
769 for _ in 0..node_probes {
772 let batched_id =
773 batched.insert_node(Node::new(NodeId(0), "Probe", Vec::new(), Vec::new()));
774 let sequential_id =
775 sequential.insert_node(Node::new(NodeId(0), "Probe", Vec::new(), Vec::new()));
776 assert_eq!(
777 batched_id, sequential_id,
778 "node arena mismatch on randomized trial {trial}"
779 );
780 }
781 for _ in 0..value_probes {
782 let batched_id = batched.create_value(DataType::Float32, static_shape([1]));
783 let sequential_id = sequential.create_value(DataType::Float32, static_shape([1]));
784 assert_eq!(
785 batched_id, sequential_id,
786 "value arena mismatch on randomized trial {trial}"
787 );
788 }
789 }
790
791 struct TestRng(u64);
792
793 impl TestRng {
794 fn next(&mut self) -> u64 {
795 self.0 ^= self.0 << 13;
796 self.0 ^= self.0 >> 7;
797 self.0 ^= self.0 << 17;
798 self.0
799 }
800
801 fn usize(&mut self, upper: usize) -> usize {
802 (self.next() as usize) % upper
803 }
804 }
805
806 fn reference_remove_node(graph: &mut Graph, id: NodeId) {
807 if !graph.nodes.contains(id) {
808 return;
809 }
810 let (inputs, outputs) = {
811 let node = graph.node(id);
812 (
813 node.input_values().collect::<Vec<_>>(),
814 node.outputs.clone(),
815 )
816 };
817 let unique_inputs: HashSet<_> = inputs.into_iter().collect();
818 for value in unique_inputs {
819 if let Some(metadata) = graph.values.get_mut(value) {
820 for (consumer, input_index) in metadata.consumers.uses() {
821 if consumer == id {
822 metadata.consumers.remove(consumer, input_index);
823 }
824 }
825 }
826 }
827 for &value in &outputs {
828 if let Some(metadata) = graph.values.get_mut(value)
829 && metadata.producer == Some(id)
830 {
831 metadata.producer = None;
832 }
833 }
834 graph.nodes.remove(id);
835 for value in outputs {
836 let orphan = graph.values.get(value).is_some_and(|metadata| {
837 metadata.producer.is_none()
838 && metadata.consumers.is_empty()
839 && !graph.inputs.contains(&value)
840 && !graph.outputs.contains(&value)
841 && !graph.initializers.contains_key(&value)
842 });
843 if orphan {
844 graph.values.remove(value);
845 graph.unknown_value_types.remove(&value);
846 graph.unknown_value_shapes.remove(&value);
847 }
848 }
849 }
850
851 fn reference_topological_order(graph: &Graph) -> Result<Vec<NodeId>, GraphError> {
852 let mut in_degree: HashMap<NodeId, usize> =
853 graph.nodes.keys().map(|node| (node, 0)).collect();
854 let mut adjacency: HashMap<NodeId, Vec<NodeId>> = HashMap::new();
855 for (node_id, node) in graph.nodes.iter() {
856 for value in node.input_values() {
857 if let Some(producer) = graph.value(value).producer
858 && graph.nodes.contains(producer)
859 {
860 adjacency.entry(producer).or_default().push(node_id);
861 *in_degree.get_mut(&node_id).unwrap() += 1;
862 }
863 }
864 }
865 let mut ready: BinaryHeap<std::cmp::Reverse<u32>> = in_degree
866 .iter()
867 .filter(|(_, degree)| **degree == 0)
868 .map(|(node, _)| std::cmp::Reverse(node.0))
869 .collect();
870 let mut order = Vec::with_capacity(graph.num_nodes());
871 while let Some(std::cmp::Reverse(raw)) = ready.pop() {
872 let node = NodeId(raw);
873 order.push(node);
874 if let Some(successors) = adjacency.get(&node) {
875 for successor in successors {
876 let degree = in_degree.get_mut(successor).unwrap();
877 *degree -= 1;
878 if *degree == 0 {
879 ready.push(std::cmp::Reverse(successor.0));
880 }
881 }
882 }
883 }
884 if order.len() == graph.num_nodes() {
885 Ok(order)
886 } else {
887 Err(GraphError::CycleDetected)
888 }
889 }
890
891 fn sample_graph() -> Graph {
893 let mut g = Graph::new();
894 g.opset_imports.insert(String::new(), 17);
895 let a = g.create_named_value("a", DataType::Float32, static_shape([4]));
896 let c = g.create_named_value("c", DataType::Float32, static_shape([4]));
897 g.add_input(a);
898 g.add_input(c);
899
900 let b = g.create_value(DataType::Float32, static_shape([4]));
901 let relu = Node::new(NodeId(0), "Relu", vec![Some(a)], vec![b]);
902 g.insert_node(relu);
903
904 let d = g.create_named_value("d", DataType::Float32, static_shape([4]));
905 let add = Node::new(NodeId(0), "Add", vec![Some(b), Some(c)], vec![d]);
906 g.insert_node(add);
907 g.add_output(d);
908 g
909 }
910
911 #[test]
912 fn edges_are_wired_on_insert() {
913 let g = sample_graph();
914 assert_eq!(g.num_nodes(), 2);
915 let b = g.value(ValueId(2));
917 assert!(b.producer.is_some());
918 assert_eq!(b.consumers.len(), 1);
919 }
920
921 #[test]
922 fn topo_order_is_valid_and_deterministic() {
923 let g = sample_graph();
924 let order = g.topological_order().unwrap();
925 assert_eq!(order.len(), 2);
926 assert_eq!(order, vec![NodeId(0), NodeId(1)]);
928 }
929
930 #[test]
931 fn uses_preserve_input_multiplicity_and_replace_one_slot() {
932 let mut graph = Graph::new();
933 let old = graph.create_value(DataType::Float32, static_shape([1]));
934 let new = graph.create_value(DataType::Float32, static_shape([1]));
935 graph.add_input(old);
936 graph.add_input(new);
937 let output = graph.create_value(DataType::Float32, static_shape([1]));
938 let node = graph.insert_node(Node::new(
939 NodeId(0),
940 "Add",
941 vec![Some(old), Some(old)],
942 vec![output],
943 ));
944
945 assert_eq!(graph.uses(old), vec![(node, 0), (node, 1)]);
946 assert_eq!(graph.consumers(old), vec![node]);
947 assert_eq!(graph.replace_input(node, 1, Some(new)), Some(old));
948 assert_eq!(graph.uses(old), vec![(node, 0)]);
949 assert_eq!(graph.uses(new), vec![(node, 1)]);
950 assert!(graph.validate().is_ok());
951
952 graph.remove_node(node);
953 assert!(graph.uses(old).is_empty());
954 assert!(graph.uses(new).is_empty());
955 }
956
957 #[test]
958 fn io_membership_flags_follow_ordered_lists() {
959 let mut graph = Graph::new();
960 let a = graph.create_value(DataType::Float32, static_shape([1]));
961 let b = graph.create_value(DataType::Float32, static_shape([1]));
962 graph.add_input(a);
963 graph.add_output(a);
964 graph.insert_output(0, b);
965 assert!(graph.value(a).is_graph_input);
966 assert!(graph.value(a).is_graph_output);
967 assert!(graph.value(b).is_graph_output);
968
969 assert_eq!(graph.remove_output(1), a);
970 assert!(!graph.value(a).is_graph_output);
971 graph.set_inputs(vec![b]);
972 assert!(!graph.value(a).is_graph_input);
973 assert!(graph.value(b).is_graph_input);
974 graph.set_outputs(vec![a]);
975 assert!(graph.value(a).is_graph_output);
976 assert!(!graph.value(b).is_graph_output);
977 }
978
979 #[test]
980 fn consumer_hash_insertion_order_is_not_observable() {
981 let mut first = Graph::new();
982 let hub = first.create_value(DataType::Float32, static_shape([1]));
983 first.add_input(hub);
984 let mut nodes = Vec::new();
985 for _ in 0..32 {
986 let output = first.create_value(DataType::Float32, static_shape([1]));
987 nodes.push(first.insert_node(Node::new(
988 NodeId(0),
989 "Add",
990 vec![Some(hub), Some(hub)],
991 vec![output],
992 )));
993 }
994 let mut shuffled = first.clone();
995 let mut uses = shuffled.uses(hub);
996 for &(node, input_index) in &uses {
997 shuffled.replace_input(node, input_index as usize, None);
998 }
999 let mut rng = TestRng(0x6a09_e667_f3bc_c909);
1000 for index in (1..uses.len()).rev() {
1001 uses.swap(index, rng.usize(index + 1));
1002 }
1003 for (node, input_index) in uses {
1004 shuffled.replace_input(node, input_index as usize, Some(hub));
1005 }
1006
1007 assert_eq!(first.uses(hub), shuffled.uses(hub));
1008 assert_eq!(first.consumers(hub), shuffled.consumers(hub));
1009 assert_eq!(first.topological_order(), shuffled.topological_order());
1010 assert_eq!(format!("{first:#?}"), format!("{shuffled:#?}"));
1011 assert_eq!(
1012 format!("{:#?}", sample_graph()),
1013 format!("{:#?}", sample_graph())
1014 );
1015 assert_eq!(nodes, first.consumers(hub));
1016 }
1017
1018 #[test]
1019 fn validate_accepts_wellformed_graph() {
1020 let g = sample_graph();
1021 assert!(g.validate().is_ok());
1022 }
1023
1024 #[test]
1025 fn predecessors_and_successors() {
1026 let g = sample_graph();
1027 assert_eq!(g.successors(NodeId(0)), vec![NodeId(1)]);
1028 assert_eq!(g.predecessors(NodeId(1)), vec![NodeId(0)]);
1029 }
1030
1031 #[test]
1032 fn nodes_between_walks_back() {
1033 let g = sample_graph();
1034 let between = g.nodes_between(&[ValueId(0), ValueId(1)], &[ValueId(3)]);
1035 assert!(between.contains(&NodeId(0)));
1036 assert!(between.contains(&NodeId(1)));
1037 assert_eq!(between.len(), 2);
1038 }
1039
1040 #[test]
1041 fn replace_all_uses_redirects_consumers() {
1042 let mut g = sample_graph();
1043 let e = g.create_value(DataType::Float32, static_shape([4]));
1045 g.replace_all_uses(ValueId(2), e);
1046 let add = g.node(NodeId(1));
1048 assert!(add.input_values().any(|v| v == e));
1049 assert!(!add.input_values().any(|v| v == ValueId(2)));
1050 assert_eq!(g.consumers(e), vec![NodeId(1)]);
1051 assert!(g.value(ValueId(2)).consumers.is_empty());
1052 }
1053
1054 #[test]
1055 fn insert_on_edge_splices_node() {
1056 let mut g = sample_graph();
1057 let ident = Node::new(NodeId(0), "Identity", vec![], vec![]);
1059 let nid = g.insert_on_edge(ValueId(2), ident);
1060 let new_out = g.node(nid).outputs[0];
1062 assert_eq!(g.value(new_out).producer, Some(nid));
1063 assert!(g.node(NodeId(1)).input_values().any(|v| v == new_out));
1064 assert!(g.validate().is_ok());
1065 }
1066
1067 #[test]
1068 fn remove_node_disconnects_and_gcs() {
1069 let mut g = sample_graph();
1070 g.remove_node(NodeId(1)); assert_eq!(g.num_nodes(), 1);
1072 assert!(g.value(ValueId(3)).producer.is_none());
1074 assert!(g.value(ValueId(2)).consumers.is_empty());
1076 }
1077
1078 #[test]
1079 fn remove_nodes_filters_wide_shared_input_once() {
1080 let mut graph = Graph::new();
1081 let input = graph.create_value(DataType::Float32, static_shape([1]));
1082 graph.add_input(input);
1083 let mut dead = Vec::new();
1084 let mut outputs = Vec::new();
1085 for _ in 0..1_000 {
1086 let output = graph.create_value(DataType::Float32, static_shape([1]));
1087 dead.push(graph.insert_node(Node::new(
1088 NodeId(0),
1089 "Relu",
1090 vec![Some(input)],
1091 vec![output],
1092 )));
1093 outputs.push(output);
1094 }
1095
1096 graph.remove_nodes(&dead);
1097
1098 assert_eq!(graph.num_nodes(), 0);
1099 assert!(graph.value(input).consumers.is_empty());
1100 assert!(
1101 outputs
1102 .into_iter()
1103 .all(|output| graph.try_value(output).is_none())
1104 );
1105 }
1106
1107 #[test]
1108 fn remove_nodes_keeps_surviving_shared_consumer() {
1109 let mut graph = Graph::new();
1110 let input = graph.create_value(DataType::Float32, static_shape([1]));
1111 graph.add_input(input);
1112 let mut nodes = Vec::new();
1113 for op_type in ["Relu", "Neg", "Abs"] {
1114 let output = graph.create_value(DataType::Float32, static_shape([1]));
1115 nodes.push(graph.insert_node(Node::new(
1116 NodeId(0),
1117 op_type,
1118 vec![Some(input)],
1119 vec![output],
1120 )));
1121 }
1122
1123 graph.remove_nodes(&nodes[..2]);
1124
1125 assert_eq!(graph.consumers(input), vec![nodes[2]]);
1126 assert!(graph.try_node(nodes[2]).is_some());
1127 assert!(graph.validate().is_ok());
1128 }
1129
1130 #[test]
1131 fn remove_nodes_collects_value_after_all_consumers_are_removed() {
1132 let mut graph = Graph::new();
1133 let input = graph.create_value(DataType::Float32, static_shape([1]));
1134 graph.add_input(input);
1135 let shared = graph.create_value(DataType::Float32, static_shape([1]));
1136 let producer = graph.insert_node(Node::new(
1137 NodeId(0),
1138 "Relu",
1139 vec![Some(input)],
1140 vec![shared],
1141 ));
1142 let mut consumers = Vec::new();
1143 for op_type in ["Neg", "Abs"] {
1144 let output = graph.create_value(DataType::Float32, static_shape([1]));
1145 consumers.push(graph.insert_node(Node::new(
1146 NodeId(0),
1147 op_type,
1148 vec![Some(shared)],
1149 vec![output],
1150 )));
1151 }
1152
1153 graph.remove_nodes(&[consumers[0], consumers[1], producer]);
1154
1155 assert!(graph.try_value(shared).is_none());
1156 }
1157
1158 #[test]
1159 fn remove_nodes_keeps_graph_outputs_and_initializers() {
1160 let mut graph = Graph::new();
1161 let input = graph.create_value(DataType::Float32, static_shape([1]));
1162 graph.add_input(input);
1163
1164 let graph_output = graph.create_value(DataType::Float32, static_shape([1]));
1165 let output_node = graph.insert_node(Node::new(
1166 NodeId(0),
1167 "Relu",
1168 vec![Some(input)],
1169 vec![graph_output],
1170 ));
1171 graph.add_output(graph_output);
1172
1173 let initializer = graph.create_value(DataType::Float32, static_shape([1]));
1174 let initializer_node = graph.insert_node(Node::new(
1175 NodeId(0),
1176 "Neg",
1177 vec![Some(input)],
1178 vec![initializer],
1179 ));
1180 graph.set_initializer(
1181 initializer,
1182 WeightRef::Inline(TensorData::from_raw(
1183 DataType::Float32,
1184 vec![1],
1185 0.0f32.to_le_bytes().to_vec(),
1186 )),
1187 );
1188
1189 graph.remove_nodes(&[output_node, initializer_node]);
1190
1191 assert!(graph.try_value(graph_output).is_some());
1192 assert!(graph.value(graph_output).producer.is_none());
1193 assert!(graph.try_value(initializer).is_some());
1194 assert!(graph.value(initializer).producer.is_none());
1195 }
1196
1197 #[test]
1198 fn remove_nodes_ignores_duplicate_and_nonlive_ids() {
1199 let graph = sample_graph();
1200 let ids = [NodeId(u32::MAX), NodeId(1), NodeId(1)];
1201 let mut sequential = graph.clone();
1202 for id in ids {
1203 sequential.remove_node(id);
1204 }
1205 let mut batched = graph;
1206 batched.remove_nodes(&ids);
1207
1208 assert_graphs_identical(batched, sequential, 3, 5, 0);
1209 }
1210
1211 #[test]
1212 fn remove_nodes_matches_sequential_removal_on_random_dags() {
1213 let mut rng = TestRng(0x4d59_5df4_d0f3_3173);
1214
1215 for trial in 0..10_000 {
1216 let input_count = 1 + rng.usize(3);
1217 let node_count = 1 + rng.usize(12);
1218 let mut graph = Graph::new();
1219 let mut values = Vec::new();
1220 for _ in 0..input_count {
1221 let input = graph.create_value(DataType::Float32, static_shape([1]));
1222 graph.add_input(input);
1223 values.push(input);
1224 }
1225
1226 let mut nodes = Vec::with_capacity(node_count);
1227 for _ in 0..node_count {
1228 let input_arity = 1 + rng.usize(3);
1229 let inputs = (0..input_arity)
1230 .map(|_| Some(values[rng.usize(values.len())]))
1231 .collect();
1232 let output = graph.create_value(DataType::Float32, static_shape([1]));
1233 nodes.push(graph.insert_node(Node::new(NodeId(0), "Random", inputs, vec![output])));
1234 if rng.usize(5) == 0 {
1235 graph.mark_value_type_unknown(output);
1236 }
1237 if rng.usize(5) == 0 {
1238 graph.mark_value_shape_unknown(output);
1239 }
1240 values.push(output);
1241 }
1242
1243 for _ in 0..rng.usize(4) {
1244 let output = values[rng.usize(values.len())];
1245 graph.add_output(output);
1246 }
1247
1248 for i in (1..nodes.len()).rev() {
1249 let j = rng.usize(i + 1);
1250 nodes.swap(i, j);
1251 }
1252 nodes.truncate(rng.usize(nodes.len() + 1));
1253
1254 let mut sequential = graph.clone();
1255 for &id in &nodes {
1256 sequential.remove_node(id);
1257 }
1258 let mut batched = graph;
1259 batched.remove_nodes(&nodes);
1260
1261 assert_graphs_identical(
1262 batched,
1263 sequential,
1264 node_count + 1,
1265 input_count + node_count + 1,
1266 trial,
1267 );
1268 }
1269 }
1270
1271 #[test]
1272 fn single_node_removal_matches_vector_reference_on_random_dags() {
1273 let mut rng = TestRng(0xbb67_ae85_84ca_a73b);
1274 for trial in 0..2_000 {
1275 let mut graph = Graph::new();
1276 let input_count = 1 + rng.usize(3);
1277 let node_count = 1 + rng.usize(24);
1278 let mut values = Vec::new();
1279 for _ in 0..input_count {
1280 let value = graph.create_value(DataType::Float32, static_shape([1]));
1281 graph.add_input(value);
1282 values.push(value);
1283 }
1284 let mut nodes = Vec::new();
1285 for _ in 0..node_count {
1286 let input_count = 1 + rng.usize(4);
1287 let inputs = (0..input_count)
1288 .map(|_| Some(values[rng.usize(values.len())]))
1289 .collect();
1290 let output = graph.create_value(DataType::Float32, static_shape([1]));
1291 nodes.push(graph.insert_node(Node::new(NodeId(0), "Random", inputs, vec![output])));
1292 values.push(output);
1293 }
1294 for _ in 0..rng.usize(4) {
1295 graph.add_output(values[rng.usize(values.len())]);
1296 }
1297
1298 let mut removals = nodes.clone();
1299 for index in (1..removals.len()).rev() {
1300 removals.swap(index, rng.usize(index + 1));
1301 }
1302 removals.truncate(rng.usize(removals.len() + 1));
1303 if let Some(&duplicate) = removals.first() {
1304 removals.push(duplicate);
1305 }
1306 removals.push(NodeId(u32::MAX));
1307
1308 let mut actual = graph.clone();
1309 let mut reference = graph;
1310 for &node in &removals {
1311 actual.remove_node(node);
1312 reference_remove_node(&mut reference, node);
1313 }
1314
1315 assert_eq!(
1316 format!("{actual:#?}"),
1317 format!("{reference:#?}"),
1318 "debug mismatch on trial {trial}"
1319 );
1320 assert_eq!(
1321 actual.topological_order(),
1322 reference.topological_order(),
1323 "topology mismatch on trial {trial}"
1324 );
1325 for value in actual.values.keys() {
1326 assert_eq!(
1327 actual.uses(value),
1328 reference.uses(value),
1329 "uses mismatch for {value:?} on trial {trial}"
1330 );
1331 assert_eq!(
1332 actual.consumers(value),
1333 reference.consumers(value),
1334 "consumers mismatch for {value:?} on trial {trial}"
1335 );
1336 }
1337 }
1338 }
1339
1340 #[test]
1341 fn vec_indexed_topology_matches_hashmap_reference_on_random_dags() {
1342 let mut rng = TestRng(0x3c6e_f372_fe94_f82b);
1343 for trial in 0..2_000 {
1344 let mut graph = Graph::new();
1345 let input = graph.create_value(DataType::Float32, static_shape([1]));
1346 graph.add_input(input);
1347 let mut values = vec![input];
1348 let mut nodes = Vec::new();
1349 for _ in 0..(1 + rng.usize(48)) {
1350 let inputs = (0..(1 + rng.usize(4)))
1351 .map(|_| Some(values[rng.usize(values.len())]))
1352 .collect();
1353 let output = graph.create_value(DataType::Float32, static_shape([1]));
1354 nodes.push(graph.insert_node(Node::new(NodeId(0), "Random", inputs, vec![output])));
1355 values.push(output);
1356 }
1357 for &node in nodes.iter().filter(|_| rng.usize(5) == 0) {
1358 graph.remove_node(node);
1359 }
1360 assert_eq!(
1361 graph.topological_order(),
1362 reference_topological_order(&graph),
1363 "topology mismatch on trial {trial}"
1364 );
1365 }
1366 }
1367
1368 #[test]
1369 fn replace_node_groups_matches_sequential_mutation() {
1370 let mut sequential = Graph::new();
1371 let input = sequential.create_value(DataType::Float32, static_shape([1]));
1372 sequential.add_input(input);
1373 let interior = sequential.create_value(DataType::Float32, static_shape([1]));
1374 let first = sequential.insert_node(Node::new(
1375 NodeId(0),
1376 "Relu",
1377 vec![Some(input)],
1378 vec![interior],
1379 ));
1380 let output = sequential.create_value(DataType::Float32, static_shape([1]));
1381 let second = sequential.insert_node(Node::new(
1382 NodeId(0),
1383 "Relu",
1384 vec![Some(interior)],
1385 vec![output],
1386 ));
1387 sequential.add_output(output);
1388 let sibling_output = sequential.create_value(DataType::Float32, static_shape([1]));
1389 let sibling = sequential.insert_node(Node::new(
1390 NodeId(0),
1391 "Neg",
1392 vec![Some(input)],
1393 vec![sibling_output],
1394 ));
1395 sequential.add_output(sibling_output);
1396
1397 let mut batched = sequential.clone();
1398 let graph_outputs: HashSet<_> = batched.outputs.iter().copied().collect();
1399 let replacement0 = Node::new(NodeId(0), "EPContext", vec![Some(input)], vec![output]);
1400 let replacement1 = Node::new(
1401 NodeId(0),
1402 "EPContext",
1403 vec![Some(input)],
1404 vec![sibling_output],
1405 );
1406
1407 sequential.remove_node(first);
1408 sequential.remove_node(second);
1409 let replacement0_id = sequential.insert_node(replacement0.clone());
1410 sequential.remove_node(sibling);
1411 let replacement1_id = sequential.insert_node(replacement1.clone());
1412
1413 let inserted = batched.replace_node_groups(
1414 vec![
1415 (vec![first, second], replacement0),
1416 (vec![sibling], replacement1),
1417 ],
1418 &graph_outputs,
1419 );
1420 assert_eq!(inserted, vec![replacement0_id, replacement1_id]);
1421
1422 let sequential_nodes: Vec<_> = sequential
1423 .nodes
1424 .iter()
1425 .map(|(id, node)| {
1426 (
1427 id,
1428 node.op_type.clone(),
1429 node.inputs.clone(),
1430 node.outputs.clone(),
1431 )
1432 })
1433 .collect();
1434 let batched_nodes: Vec<_> = batched
1435 .nodes
1436 .iter()
1437 .map(|(id, node)| {
1438 (
1439 id,
1440 node.op_type.clone(),
1441 node.inputs.clone(),
1442 node.outputs.clone(),
1443 )
1444 })
1445 .collect();
1446 assert_eq!(batched_nodes, sequential_nodes);
1447
1448 let sequential_values: Vec<_> = sequential
1449 .values
1450 .iter()
1451 .map(|(id, value)| (id, value.producer, value.consumers.clone()))
1452 .collect();
1453 let batched_values: Vec<_> = batched
1454 .values
1455 .iter()
1456 .map(|(id, value)| (id, value.producer, value.consumers.clone()))
1457 .collect();
1458 assert_eq!(batched_values, sequential_values);
1459
1460 let next_sequential =
1461 sequential.insert_node(Node::new(NodeId(0), "Identity", Vec::new(), Vec::new()));
1462 let next_batched =
1463 batched.insert_node(Node::new(NodeId(0), "Identity", Vec::new(), Vec::new()));
1464 assert_eq!(next_batched, next_sequential);
1465 }
1466
1467 #[test]
1468 fn replace_node_groups_matches_sequential_orphan_collection() {
1469 let mut sequential = Graph::new();
1470 let input = sequential.create_value(DataType::Float32, static_shape([1]));
1471 sequential.add_input(input);
1472 let interior = sequential.create_value(DataType::Float32, static_shape([1]));
1473 let producer = sequential.insert_node(Node::new(
1474 NodeId(0),
1475 "Relu",
1476 vec![Some(input)],
1477 vec![interior],
1478 ));
1479 let output = sequential.create_value(DataType::Float32, static_shape([1]));
1480 let consumer = sequential.insert_node(Node::new(
1481 NodeId(0),
1482 "Relu",
1483 vec![Some(interior)],
1484 vec![output],
1485 ));
1486 sequential.add_output(output);
1487
1488 let mut batched = sequential.clone();
1489 let graph_outputs: HashSet<_> = batched.outputs.iter().copied().collect();
1490 let replacement = Node::new(NodeId(0), "EPContext", vec![Some(input)], vec![output]);
1491
1492 sequential.remove_node(consumer);
1493 sequential.remove_node(producer);
1494 sequential.insert_node(replacement.clone());
1495 batched.replace_node_groups(
1496 vec![(vec![consumer, producer], replacement)],
1497 &graph_outputs,
1498 );
1499
1500 assert!(sequential.try_value(interior).is_none());
1501 assert!(batched.try_value(interior).is_none());
1502 let next_sequential = sequential.create_value(DataType::Float32, static_shape([1]));
1503 let next_batched = batched.create_value(DataType::Float32, static_shape([1]));
1504 assert_eq!(next_batched, next_sequential);
1505 }
1506
1507 #[test]
1508 fn replace_node_preserves_id() {
1509 let mut g = sample_graph();
1510 let d = g.node(NodeId(1)).outputs[0];
1511 let b = g.node(NodeId(1)).inputs[0];
1512 let c = g.node(NodeId(1)).inputs[1];
1513 let sub = Node::new(NodeId(0), "Sub", vec![b, c], vec![d]);
1514 let id = g.replace_node(NodeId(1), sub);
1515 assert_eq!(id, NodeId(1));
1516 assert_eq!(g.node(NodeId(1)).op_type, "Sub");
1517 assert!(g.validate().is_ok());
1518 }
1519
1520 #[test]
1521 fn cycle_is_detected() {
1522 let mut g = Graph::new();
1523 let v0 = g.create_value(DataType::Float32, static_shape([1]));
1524 let v1 = g.create_value(DataType::Float32, static_shape([1]));
1525 g.insert_node(Node::new(NodeId(0), "A", vec![Some(v1)], vec![v0]));
1527 g.insert_node(Node::new(NodeId(0), "B", vec![Some(v0)], vec![v1]));
1528 assert_eq!(g.topological_order(), Err(GraphError::CycleDetected));
1529 assert!(g.validate().is_err());
1530 }
1531
1532 #[test]
1533 fn intern_symbol_dedups_by_name() {
1534 let mut g = Graph::new();
1535 let s1 = g.intern_symbol("batch");
1536 let s2 = g.intern_symbol("batch");
1537 let s3 = g.intern_symbol("seq");
1538 assert_eq!(s1, s2);
1539 assert_ne!(s1, s3);
1540 }
1541}