oxicuda_graph/
graph.rs

1//! The `ComputeGraph` — a directed acyclic graph of GPU operations.
2//!
3//! `ComputeGraph` is the central data structure of `oxicuda-graph`. It
4//! stores nodes (GPU operations), directed dependency edges, and buffer
5//! descriptors. Analysis and optimisation passes operate on this structure
6//! before it is lowered to an `ExecutionPlan` by the executor.
7//!
8//! # Design
9//!
10//! * Nodes are stored in a `Vec` indexed by a dense `u32` slot.
11//! * Adjacency is maintained as two maps: `successors` and `predecessors`,
12//!   each mapping `NodeId → Vec<NodeId>`. This makes both forward and
13//!   backward traversals O(degree).
14//! * Cycle detection happens eagerly on `add_edge` using DFS, so the
15//!   invariant "this graph is a DAG" always holds after successful mutation.
16
17use std::collections::{HashMap, HashSet, VecDeque};
18
19use crate::error::{GraphError, GraphResult};
20use crate::node::{BufferDescriptor, BufferId, GraphNode, NodeId};
21
22// ---------------------------------------------------------------------------
23// ComputeGraph
24// ---------------------------------------------------------------------------
25
26/// A directed acyclic graph (DAG) of GPU operations.
27///
28/// Nodes represent individual GPU operations ([`GraphNode`]). Directed edges
29/// express execution-order dependencies: an edge `a → b` means `b` cannot
30/// begin until `a` has completed. Buffers referenced by nodes are registered
31/// separately via [`add_buffer`](ComputeGraph::add_buffer).
32///
33/// # Guarantees
34///
35/// * The graph is always a DAG — [`add_edge`](ComputeGraph::add_edge) returns
36///   [`GraphError::CycleDetected`] if adding the edge would create a cycle.
37/// * Node IDs are unique and stable: nodes are never moved or re-indexed.
38/// * Buffer IDs are unique and stable.
39#[derive(Debug, Clone)]
40pub struct ComputeGraph {
41    /// Nodes stored in insertion order (indexed by NodeId.0).
42    nodes: Vec<GraphNode>,
43    /// Successor adjacency: node → nodes that depend on it.
44    successors: Vec<Vec<NodeId>>,
45    /// Predecessor adjacency: node → nodes it depends on.
46    predecessors: Vec<Vec<NodeId>>,
47    /// Buffer metadata (indexed by BufferId.0).
48    buffers: Vec<BufferDescriptor>,
49    /// Next node id to allocate.
50    next_node: u32,
51    /// Next buffer id to allocate.
52    next_buf: u32,
53}
54
55impl Default for ComputeGraph {
56    fn default() -> Self {
57        Self::new()
58    }
59}
60
61impl ComputeGraph {
62    /// Creates an empty computation graph.
63    #[must_use]
64    pub fn new() -> Self {
65        Self {
66            nodes: Vec::new(),
67            successors: Vec::new(),
68            predecessors: Vec::new(),
69            buffers: Vec::new(),
70            next_node: 0,
71            next_buf: 0,
72        }
73    }
74
75    // -----------------------------------------------------------------------
76    // Node management
77    // -----------------------------------------------------------------------
78
79    /// Adds a node to the graph and returns its assigned `NodeId`.
80    ///
81    /// The node's `id` field will be overwritten with the allocated ID.
82    pub fn add_node(&mut self, mut node: GraphNode) -> NodeId {
83        let id = NodeId(self.next_node);
84        node.id = id;
85        self.next_node += 1;
86        self.nodes.push(node);
87        self.successors.push(Vec::new());
88        self.predecessors.push(Vec::new());
89        id
90    }
91
92    /// Returns a reference to the node with the given ID, or an error.
93    pub fn node(&self, id: NodeId) -> GraphResult<&GraphNode> {
94        self.nodes
95            .get(id.0 as usize)
96            .ok_or(GraphError::NodeNotFound(id))
97    }
98
99    /// Returns a mutable reference to the node with the given ID, or an error.
100    pub fn node_mut(&mut self, id: NodeId) -> GraphResult<&mut GraphNode> {
101        self.nodes
102            .get_mut(id.0 as usize)
103            .ok_or(GraphError::NodeNotFound(id))
104    }
105
106    /// Returns a slice of all nodes in insertion order.
107    #[inline]
108    pub fn nodes(&self) -> &[GraphNode] {
109        &self.nodes
110    }
111
112    /// Returns the total number of nodes.
113    #[inline]
114    pub fn node_count(&self) -> usize {
115        self.nodes.len()
116    }
117
118    // -----------------------------------------------------------------------
119    // Buffer management
120    // -----------------------------------------------------------------------
121
122    /// Registers a buffer and returns its assigned `BufferId`.
123    ///
124    /// The buffer's `id` field will be overwritten with the allocated ID.
125    pub fn add_buffer(&mut self, mut buf: BufferDescriptor) -> BufferId {
126        let id = BufferId(self.next_buf);
127        buf.id = id;
128        self.next_buf += 1;
129        self.buffers.push(buf);
130        id
131    }
132
133    /// Returns a reference to the buffer descriptor, or an error.
134    pub fn buffer(&self, id: BufferId) -> GraphResult<&BufferDescriptor> {
135        self.buffers
136            .get(id.0 as usize)
137            .ok_or(GraphError::NodeNotFound(NodeId(id.0)))
138    }
139
140    /// Returns a slice of all buffer descriptors.
141    #[inline]
142    pub fn buffers(&self) -> &[BufferDescriptor] {
143        &self.buffers
144    }
145
146    /// Returns the number of registered buffers.
147    #[inline]
148    pub fn buffer_count(&self) -> usize {
149        self.buffers.len()
150    }
151
152    // -----------------------------------------------------------------------
153    // Edge management
154    // -----------------------------------------------------------------------
155
156    /// Adds a directed dependency edge `from → to`.
157    ///
158    /// This means `to` will not begin execution until `from` has completed.
159    ///
160    /// # Errors
161    ///
162    /// * [`GraphError::NodeNotFound`] if either ID is invalid.
163    /// * [`GraphError::CycleDetected`] if the edge would create a cycle.
164    pub fn add_edge(&mut self, from: NodeId, to: NodeId) -> GraphResult<()> {
165        let n = self.nodes.len();
166        if from.0 as usize >= n {
167            return Err(GraphError::NodeNotFound(from));
168        }
169        if to.0 as usize >= n {
170            return Err(GraphError::NodeNotFound(to));
171        }
172        if from == to {
173            return Err(GraphError::CycleDetected { from, to });
174        }
175        // Check whether adding from→to would create a cycle: it does iff
176        // `from` is reachable from `to` in the current graph.
177        if self.is_reachable(to, from) {
178            return Err(GraphError::CycleDetected { from, to });
179        }
180        // Avoid duplicate edges.
181        if !self.successors[from.0 as usize].contains(&to) {
182            self.successors[from.0 as usize].push(to);
183            self.predecessors[to.0 as usize].push(from);
184        }
185        Ok(())
186    }
187
188    /// Returns whether node `src` can reach node `dst` via directed edges.
189    ///
190    /// Uses iterative BFS to avoid stack overflow on deep graphs.
191    pub fn is_reachable(&self, src: NodeId, dst: NodeId) -> bool {
192        if src == dst {
193            return true;
194        }
195        let n = self.nodes.len();
196        let mut visited = vec![false; n];
197        let mut queue = VecDeque::new();
198        queue.push_back(src);
199        visited[src.0 as usize] = true;
200        while let Some(curr) = queue.pop_front() {
201            for &next in &self.successors[curr.0 as usize] {
202                if next == dst {
203                    return true;
204                }
205                if !visited[next.0 as usize] {
206                    visited[next.0 as usize] = true;
207                    queue.push_back(next);
208                }
209            }
210        }
211        false
212    }
213
214    /// Returns successors of `id` (nodes that depend on it).
215    pub fn successors(&self, id: NodeId) -> GraphResult<&[NodeId]> {
216        if id.0 as usize >= self.nodes.len() {
217            return Err(GraphError::NodeNotFound(id));
218        }
219        Ok(&self.successors[id.0 as usize])
220    }
221
222    /// Returns predecessors of `id` (nodes it depends on).
223    pub fn predecessors(&self, id: NodeId) -> GraphResult<&[NodeId]> {
224        if id.0 as usize >= self.nodes.len() {
225            return Err(GraphError::NodeNotFound(id));
226        }
227        Ok(&self.predecessors[id.0 as usize])
228    }
229
230    /// Returns the total number of directed edges in the graph.
231    pub fn edge_count(&self) -> usize {
232        self.successors.iter().map(|v| v.len()).sum()
233    }
234
235    /// Returns all edges as `(from, to)` pairs.
236    pub fn edges(&self) -> Vec<(NodeId, NodeId)> {
237        let mut edges = Vec::new();
238        for (i, succs) in self.successors.iter().enumerate() {
239            for &to in succs {
240                edges.push((NodeId(i as u32), to));
241            }
242        }
243        edges
244    }
245
246    // -----------------------------------------------------------------------
247    // Source / sink queries
248    // -----------------------------------------------------------------------
249
250    /// Returns all nodes with no predecessors (entry points of the graph).
251    pub fn sources(&self) -> Vec<NodeId> {
252        self.predecessors
253            .iter()
254            .enumerate()
255            .filter(|(_, preds)| preds.is_empty())
256            .map(|(i, _)| NodeId(i as u32))
257            .collect()
258    }
259
260    /// Returns all nodes with no successors (terminal nodes of the graph).
261    pub fn sinks(&self) -> Vec<NodeId> {
262        self.successors
263            .iter()
264            .enumerate()
265            .filter(|(_, succs)| succs.is_empty())
266            .map(|(i, _)| NodeId(i as u32))
267            .collect()
268    }
269
270    // -----------------------------------------------------------------------
271    // Topological sort (Kahn's algorithm)
272    // -----------------------------------------------------------------------
273
274    /// Returns nodes in topological order (all predecessors before successors).
275    ///
276    /// Uses Kahn's BFS algorithm. Because the graph is guaranteed to be a DAG
277    /// after successful `add_edge` calls, this will always succeed unless the
278    /// graph is empty.
279    ///
280    /// # Errors
281    ///
282    /// Returns [`GraphError::EmptyGraph`] if there are no nodes.
283    pub fn topological_order(&self) -> GraphResult<Vec<NodeId>> {
284        if self.nodes.is_empty() {
285            return Err(GraphError::EmptyGraph);
286        }
287        let n = self.nodes.len();
288        let mut in_degree: Vec<u32> = self.predecessors.iter().map(|p| p.len() as u32).collect();
289        let mut queue: VecDeque<NodeId> = (0..n)
290            .filter(|&i| in_degree[i] == 0)
291            .map(|i| NodeId(i as u32))
292            .collect();
293        let mut order = Vec::with_capacity(n);
294        while let Some(id) = queue.pop_front() {
295            order.push(id);
296            for &succ in &self.successors[id.0 as usize] {
297                let d = &mut in_degree[succ.0 as usize];
298                *d -= 1;
299                if *d == 0 {
300                    queue.push_back(succ);
301                }
302            }
303        }
304        // Safety: the DAG invariant guarantees all nodes are reachable.
305        debug_assert_eq!(
306            order.len(),
307            n,
308            "topological sort incomplete — internal invariant broken"
309        );
310        Ok(order)
311    }
312
313    // -----------------------------------------------------------------------
314    // Buffer-derived data-flow edges
315    // -----------------------------------------------------------------------
316
317    /// Infers and adds control edges from buffer data-flow.
318    ///
319    /// For every buffer `b`, if node `a` writes `b` (has `b` in outputs) and
320    /// node `c` reads `b` (has `b` in inputs), adds the dependency edge `a → c`.
321    ///
322    /// # Errors
323    ///
324    /// Returns [`GraphError::CycleDetected`] if any inferred edge creates a cycle.
325    pub fn infer_data_edges(&mut self) -> GraphResult<()> {
326        // Build writer map: buffer → list of nodes that write it.
327        let mut writers: HashMap<BufferId, Vec<NodeId>> = HashMap::new();
328        for node in &self.nodes {
329            for &buf in &node.outputs {
330                writers.entry(buf).or_default().push(node.id);
331            }
332        }
333        // For each node that reads a buffer, add edge from all writers → this node.
334        let reader_data: Vec<(NodeId, Vec<BufferId>)> = self
335            .nodes
336            .iter()
337            .map(|n| (n.id, n.inputs.clone()))
338            .collect();
339        for (reader_id, inputs) in reader_data {
340            for buf in inputs {
341                if let Some(node_writers) = writers.get(&buf) {
342                    for &writer_id in node_writers {
343                        if writer_id != reader_id {
344                            // Ignore duplicate-edge or cycle errors here: cycles
345                            // introduced by explicit buffer write-after-write are
346                            // a user error caught by add_edge.
347                            self.add_edge(writer_id, reader_id)?;
348                        }
349                    }
350                }
351            }
352        }
353        Ok(())
354    }
355
356    // -----------------------------------------------------------------------
357    // Subgraph extraction
358    // -----------------------------------------------------------------------
359
360    /// Returns the set of all nodes reachable (forward) from the given roots.
361    pub fn reachable_from(&self, roots: &[NodeId]) -> HashSet<NodeId> {
362        let mut visited = HashSet::new();
363        let mut stack: Vec<NodeId> = roots.to_vec();
364        while let Some(id) = stack.pop() {
365            if visited.insert(id) {
366                for &s in &self.successors[id.0 as usize] {
367                    if !visited.contains(&s) {
368                        stack.push(s);
369                    }
370                }
371            }
372        }
373        visited
374    }
375
376    /// Returns the set of all nodes that can reach any of the given targets (reverse reachability).
377    pub fn reaching(&self, targets: &[NodeId]) -> HashSet<NodeId> {
378        let mut visited = HashSet::new();
379        let mut stack: Vec<NodeId> = targets.to_vec();
380        while let Some(id) = stack.pop() {
381            if visited.insert(id) {
382                for &p in &self.predecessors[id.0 as usize] {
383                    if !visited.contains(&p) {
384                        stack.push(p);
385                    }
386                }
387            }
388        }
389        visited
390    }
391
392    // -----------------------------------------------------------------------
393    // Graph properties
394    // -----------------------------------------------------------------------
395
396    /// Returns `true` if the graph has no nodes.
397    #[inline]
398    pub fn is_empty(&self) -> bool {
399        self.nodes.is_empty()
400    }
401
402    /// Returns the longest path length (in edges) from any source to any sink.
403    ///
404    /// This is the critical path length — a lower bound on the sequential
405    /// execution depth of the graph.
406    ///
407    /// # Errors
408    ///
409    /// Returns [`GraphError::EmptyGraph`] if the graph is empty.
410    pub fn critical_path_length(&self) -> GraphResult<usize> {
411        let order = self.topological_order()?;
412        let n = self.nodes.len();
413        let mut dist = vec![0usize; n];
414        let mut max_len = 0usize;
415        for id in &order {
416            let d = dist[id.0 as usize];
417            max_len = max_len.max(d);
418            for &succ in &self.successors[id.0 as usize] {
419                let nd = d + 1;
420                if nd > dist[succ.0 as usize] {
421                    dist[succ.0 as usize] = nd;
422                }
423            }
424        }
425        Ok(max_len)
426    }
427
428    /// Returns the maximum fan-in (in-degree) among all nodes.
429    pub fn max_in_degree(&self) -> usize {
430        self.predecessors.iter().map(|p| p.len()).max().unwrap_or(0)
431    }
432
433    /// Returns the maximum fan-out (out-degree) among all nodes.
434    pub fn max_out_degree(&self) -> usize {
435        self.successors.iter().map(|s| s.len()).max().unwrap_or(0)
436    }
437
438    /// Returns the number of parallel chains (independent execution paths).
439    ///
440    /// This is an upper bound on the number of streams that can be usefully
441    /// exploited for concurrent execution.
442    pub fn parallelism_width(&self) -> GraphResult<usize> {
443        // Width = maximum number of nodes at the same topological "level"
444        // (BFS layer from sources).
445        if self.nodes.is_empty() {
446            return Ok(0);
447        }
448        let n = self.nodes.len();
449        let mut level = vec![0usize; n];
450        let mut in_degree: Vec<u32> = self.predecessors.iter().map(|p| p.len() as u32).collect();
451        let mut queue: VecDeque<NodeId> = (0..n)
452            .filter(|&i| in_degree[i] == 0)
453            .map(|i| NodeId(i as u32))
454            .collect();
455        let mut max_width = queue.len();
456        while let Some(id) = queue.pop_front() {
457            for &succ in &self.successors[id.0 as usize] {
458                let nl = level[id.0 as usize] + 1;
459                if nl > level[succ.0 as usize] {
460                    level[succ.0 as usize] = nl;
461                }
462                let d = &mut in_degree[succ.0 as usize];
463                *d -= 1;
464                if *d == 0 {
465                    queue.push_back(succ);
466                }
467            }
468            // Compute width at each level by scanning all level values.
469        }
470        // Count nodes per level.
471        let max_level = *level.iter().max().unwrap_or(&0);
472        let mut width_at_level = vec![0usize; max_level + 1];
473        for &lv in &level {
474            width_at_level[lv] += 1;
475        }
476        max_width = max_width.max(*width_at_level.iter().max().unwrap_or(&0));
477        Ok(max_width)
478    }
479
480    // -----------------------------------------------------------------------
481    // Kernel-type queries
482    // -----------------------------------------------------------------------
483
484    /// Returns all nodes that are kernel launches (compute nodes).
485    pub fn kernel_nodes(&self) -> Vec<NodeId> {
486        self.nodes
487            .iter()
488            .filter(|n| n.kind.is_compute())
489            .map(|n| n.id)
490            .collect()
491    }
492
493    /// Returns all nodes that are fusible kernel launches.
494    pub fn fusible_nodes(&self) -> Vec<NodeId> {
495        self.nodes
496            .iter()
497            .filter(|n| n.kind.is_fusible())
498            .map(|n| n.id)
499            .collect()
500    }
501
502    // -----------------------------------------------------------------------
503    // DOT format serialisation
504    // -----------------------------------------------------------------------
505
506    /// Renders the graph in Graphviz DOT format for visualisation.
507    pub fn to_dot(&self) -> String {
508        let mut s = String::from("digraph ComputeGraph {\n  rankdir=TB;\n");
509        for node in &self.nodes {
510            let label = node.display_name();
511            let shape = if node.kind.is_compute() {
512                "box"
513            } else if node.kind.is_memory_op() {
514                "parallelogram"
515            } else {
516                "ellipse"
517            };
518            s.push_str(&format!(
519                "  {} [label=\"{} ({})\", shape={shape}];\n",
520                node.id.0,
521                label,
522                node.kind.tag()
523            ));
524        }
525        for (i, succs) in self.successors.iter().enumerate() {
526            for &to in succs {
527                s.push_str(&format!("  {} -> {};\n", i, to.0));
528            }
529        }
530        s.push('}');
531        s
532    }
533}
534
535impl std::fmt::Display for ComputeGraph {
536    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
537        write!(
538            f,
539            "ComputeGraph({} nodes, {} edges, {} buffers)",
540            self.node_count(),
541            self.edge_count(),
542            self.buffer_count()
543        )
544    }
545}
546
547// ---------------------------------------------------------------------------
548// Tests
549// ---------------------------------------------------------------------------
550
551#[cfg(test)]
552mod tests {
553    use super::*;
554    use crate::node::{BufferDescriptor, KernelConfig, MemcpyDir, NodeKind};
555
556    fn kernel_node(name: &str) -> GraphNode {
557        GraphNode::new(
558            NodeId(0),
559            NodeKind::KernelLaunch {
560                function_name: name.into(),
561                config: KernelConfig::linear(1, 32, 0),
562                fusible: true,
563            },
564        )
565    }
566
567    fn barrier_node() -> GraphNode {
568        GraphNode::new(NodeId(0), NodeKind::Barrier)
569    }
570
571    fn memcpy_node(dir: MemcpyDir, size: usize) -> GraphNode {
572        GraphNode::new(
573            NodeId(0),
574            NodeKind::Memcpy {
575                dir,
576                size_bytes: size,
577            },
578        )
579    }
580
581    // --- Basic construction ---
582
583    #[test]
584    fn new_graph_is_empty() {
585        let g = ComputeGraph::new();
586        assert!(g.is_empty());
587        assert_eq!(g.node_count(), 0);
588        assert_eq!(g.edge_count(), 0);
589        assert_eq!(g.buffer_count(), 0);
590    }
591
592    #[test]
593    fn default_is_empty() {
594        let g = ComputeGraph::default();
595        assert!(g.is_empty());
596    }
597
598    #[test]
599    fn add_node_assigns_sequential_ids() {
600        let mut g = ComputeGraph::new();
601        let a = g.add_node(barrier_node());
602        let b = g.add_node(barrier_node());
603        let c = g.add_node(barrier_node());
604        assert_eq!(a, NodeId(0));
605        assert_eq!(b, NodeId(1));
606        assert_eq!(c, NodeId(2));
607        assert_eq!(g.node_count(), 3);
608    }
609
610    #[test]
611    fn node_lookup_valid() {
612        let mut g = ComputeGraph::new();
613        let id = g.add_node(kernel_node("add"));
614        assert!(g.node(id).is_ok());
615        assert_eq!(g.node(id).unwrap().kind.function_name(), Some("add"));
616    }
617
618    #[test]
619    fn node_lookup_invalid() {
620        let g = ComputeGraph::new();
621        assert!(matches!(
622            g.node(NodeId(0)),
623            Err(GraphError::NodeNotFound(_))
624        ));
625    }
626
627    #[test]
628    fn node_mut_allows_modification() {
629        let mut g = ComputeGraph::new();
630        let id = g.add_node(barrier_node());
631        g.node_mut(id).unwrap().cost_hint = 42;
632        assert_eq!(g.node(id).unwrap().cost_hint, 42);
633    }
634
635    // --- Buffer management ---
636
637    #[test]
638    fn add_buffer_assigns_ids() {
639        let mut g = ComputeGraph::new();
640        let b0 = g.add_buffer(BufferDescriptor::new(BufferId(0), 1024));
641        let b1 = g.add_buffer(BufferDescriptor::new(BufferId(0), 2048));
642        assert_eq!(b0, BufferId(0));
643        assert_eq!(b1, BufferId(1));
644        assert_eq!(g.buffer_count(), 2);
645    }
646
647    #[test]
648    fn buffer_lookup_invalid() {
649        let g = ComputeGraph::new();
650        assert!(g.buffer(BufferId(0)).is_err());
651    }
652
653    // --- Edge management ---
654
655    #[test]
656    fn add_edge_valid() {
657        let mut g = ComputeGraph::new();
658        let a = g.add_node(kernel_node("a"));
659        let b = g.add_node(kernel_node("b"));
660        assert!(g.add_edge(a, b).is_ok());
661        assert_eq!(g.edge_count(), 1);
662        assert_eq!(g.successors(a).unwrap(), &[b]);
663        assert_eq!(g.predecessors(b).unwrap(), &[a]);
664    }
665
666    #[test]
667    fn add_edge_self_loop_rejected() {
668        let mut g = ComputeGraph::new();
669        let a = g.add_node(barrier_node());
670        assert!(matches!(
671            g.add_edge(a, a),
672            Err(GraphError::CycleDetected { .. })
673        ));
674    }
675
676    #[test]
677    fn add_edge_cycle_rejected() {
678        let mut g = ComputeGraph::new();
679        let a = g.add_node(barrier_node());
680        let b = g.add_node(barrier_node());
681        g.add_edge(a, b).unwrap();
682        assert!(matches!(
683            g.add_edge(b, a),
684            Err(GraphError::CycleDetected { .. })
685        ));
686    }
687
688    #[test]
689    fn add_edge_invalid_node_rejected() {
690        let mut g = ComputeGraph::new();
691        let a = g.add_node(barrier_node());
692        assert!(matches!(
693            g.add_edge(a, NodeId(99)),
694            Err(GraphError::NodeNotFound(_))
695        ));
696        assert!(matches!(
697            g.add_edge(NodeId(99), a),
698            Err(GraphError::NodeNotFound(_))
699        ));
700    }
701
702    #[test]
703    fn add_edge_duplicate_is_idempotent() {
704        let mut g = ComputeGraph::new();
705        let a = g.add_node(barrier_node());
706        let b = g.add_node(barrier_node());
707        g.add_edge(a, b).unwrap();
708        g.add_edge(a, b).unwrap(); // second call must succeed and NOT duplicate
709        assert_eq!(g.edge_count(), 1);
710    }
711
712    #[test]
713    fn edges_returns_all() {
714        let mut g = ComputeGraph::new();
715        let a = g.add_node(barrier_node());
716        let b = g.add_node(barrier_node());
717        let c = g.add_node(barrier_node());
718        g.add_edge(a, b).unwrap();
719        g.add_edge(a, c).unwrap();
720        let mut edges = g.edges();
721        edges.sort();
722        assert_eq!(edges, vec![(a, b), (a, c)]);
723    }
724
725    // --- Reachability ---
726
727    #[test]
728    fn is_reachable_direct() {
729        let mut g = ComputeGraph::new();
730        let a = g.add_node(barrier_node());
731        let b = g.add_node(barrier_node());
732        g.add_edge(a, b).unwrap();
733        assert!(g.is_reachable(a, b));
734        assert!(!g.is_reachable(b, a));
735    }
736
737    #[test]
738    fn is_reachable_transitive() {
739        let mut g = ComputeGraph::new();
740        let a = g.add_node(barrier_node());
741        let b = g.add_node(barrier_node());
742        let c = g.add_node(barrier_node());
743        g.add_edge(a, b).unwrap();
744        g.add_edge(b, c).unwrap();
745        assert!(g.is_reachable(a, c));
746        assert!(!g.is_reachable(c, a));
747    }
748
749    #[test]
750    fn is_reachable_disconnected() {
751        let mut g = ComputeGraph::new();
752        let a = g.add_node(barrier_node());
753        let b = g.add_node(barrier_node());
754        assert!(!g.is_reachable(a, b));
755        assert!(!g.is_reachable(b, a));
756    }
757
758    // --- Sources / sinks ---
759
760    #[test]
761    fn sources_and_sinks_linear_chain() {
762        let mut g = ComputeGraph::new();
763        let a = g.add_node(barrier_node());
764        let b = g.add_node(barrier_node());
765        let c = g.add_node(barrier_node());
766        g.add_edge(a, b).unwrap();
767        g.add_edge(b, c).unwrap();
768        let sources = g.sources();
769        let sinks = g.sinks();
770        assert_eq!(sources, vec![a]);
771        assert_eq!(sinks, vec![c]);
772    }
773
774    #[test]
775    fn sources_and_sinks_diamond() {
776        let mut g = ComputeGraph::new();
777        let a = g.add_node(barrier_node()); // source
778        let b = g.add_node(barrier_node());
779        let c = g.add_node(barrier_node());
780        let d = g.add_node(barrier_node()); // sink
781        g.add_edge(a, b).unwrap();
782        g.add_edge(a, c).unwrap();
783        g.add_edge(b, d).unwrap();
784        g.add_edge(c, d).unwrap();
785        assert_eq!(g.sources(), vec![a]);
786        assert_eq!(g.sinks(), vec![d]);
787    }
788
789    // --- Topological sort ---
790
791    #[test]
792    fn topological_order_linear() {
793        let mut g = ComputeGraph::new();
794        let a = g.add_node(barrier_node());
795        let b = g.add_node(barrier_node());
796        let c = g.add_node(barrier_node());
797        g.add_edge(a, b).unwrap();
798        g.add_edge(b, c).unwrap();
799        let order = g.topological_order().unwrap();
800        let pos_a = order.iter().position(|&x| x == a).unwrap();
801        let pos_b = order.iter().position(|&x| x == b).unwrap();
802        let pos_c = order.iter().position(|&x| x == c).unwrap();
803        assert!(pos_a < pos_b && pos_b < pos_c);
804    }
805
806    #[test]
807    fn topological_order_diamond() {
808        let mut g = ComputeGraph::new();
809        let a = g.add_node(barrier_node());
810        let b = g.add_node(barrier_node());
811        let c = g.add_node(barrier_node());
812        let d = g.add_node(barrier_node());
813        g.add_edge(a, b).unwrap();
814        g.add_edge(a, c).unwrap();
815        g.add_edge(b, d).unwrap();
816        g.add_edge(c, d).unwrap();
817        let order = g.topological_order().unwrap();
818        assert_eq!(order.len(), 4);
819        let pos = |n: NodeId| order.iter().position(|&x| x == n).unwrap();
820        assert!(pos(a) < pos(b));
821        assert!(pos(a) < pos(c));
822        assert!(pos(b) < pos(d));
823        assert!(pos(c) < pos(d));
824    }
825
826    #[test]
827    fn topological_order_empty_graph() {
828        let g = ComputeGraph::new();
829        assert!(matches!(g.topological_order(), Err(GraphError::EmptyGraph)));
830    }
831
832    #[test]
833    fn topological_order_isolated_nodes() {
834        let mut g = ComputeGraph::new();
835        g.add_node(barrier_node());
836        g.add_node(barrier_node());
837        g.add_node(barrier_node());
838        let order = g.topological_order().unwrap();
839        assert_eq!(order.len(), 3);
840    }
841
842    // --- Data-flow edge inference ---
843
844    #[test]
845    fn infer_data_edges_connects_writer_to_reader() {
846        let mut g = ComputeGraph::new();
847        let buf = g.add_buffer(BufferDescriptor::new(BufferId(0), 1024));
848        let writer = g.add_node(GraphNode::new(NodeId(0), NodeKind::Barrier).with_outputs([buf]));
849        let reader = g.add_node(GraphNode::new(NodeId(0), NodeKind::Barrier).with_inputs([buf]));
850        g.infer_data_edges().unwrap();
851        assert!(g.is_reachable(writer, reader));
852    }
853
854    #[test]
855    fn infer_data_edges_multiple_readers() {
856        let mut g = ComputeGraph::new();
857        let buf = g.add_buffer(BufferDescriptor::new(BufferId(0), 1024));
858        let writer = g.add_node(GraphNode::new(NodeId(0), NodeKind::Barrier).with_outputs([buf]));
859        let r1 = g.add_node(GraphNode::new(NodeId(0), NodeKind::Barrier).with_inputs([buf]));
860        let r2 = g.add_node(GraphNode::new(NodeId(0), NodeKind::Barrier).with_inputs([buf]));
861        g.infer_data_edges().unwrap();
862        assert!(g.is_reachable(writer, r1));
863        assert!(g.is_reachable(writer, r2));
864    }
865
866    // --- Critical path, degrees ---
867
868    #[test]
869    fn critical_path_linear_chain() {
870        let mut g = ComputeGraph::new();
871        let a = g.add_node(barrier_node());
872        let b = g.add_node(barrier_node());
873        let c = g.add_node(barrier_node());
874        let d = g.add_node(barrier_node());
875        g.add_edge(a, b).unwrap();
876        g.add_edge(b, c).unwrap();
877        g.add_edge(c, d).unwrap();
878        assert_eq!(g.critical_path_length().unwrap(), 3);
879    }
880
881    #[test]
882    fn critical_path_diamond() {
883        let mut g = ComputeGraph::new();
884        let a = g.add_node(barrier_node());
885        let b = g.add_node(barrier_node());
886        let c = g.add_node(barrier_node());
887        let d = g.add_node(barrier_node());
888        g.add_edge(a, b).unwrap();
889        g.add_edge(a, c).unwrap();
890        g.add_edge(b, d).unwrap();
891        g.add_edge(c, d).unwrap();
892        assert_eq!(g.critical_path_length().unwrap(), 2);
893    }
894
895    #[test]
896    fn max_degrees_computed_correctly() {
897        let mut g = ComputeGraph::new();
898        let a = g.add_node(barrier_node());
899        let b = g.add_node(barrier_node());
900        let c = g.add_node(barrier_node());
901        let d = g.add_node(barrier_node());
902        g.add_edge(a, b).unwrap();
903        g.add_edge(a, c).unwrap();
904        g.add_edge(a, d).unwrap();
905        assert_eq!(g.max_out_degree(), 3);
906        assert_eq!(g.max_in_degree(), 1);
907    }
908
909    // --- Subgraph reachability ---
910
911    #[test]
912    fn reachable_from_set() {
913        let mut g = ComputeGraph::new();
914        let a = g.add_node(barrier_node());
915        let b = g.add_node(barrier_node());
916        let c = g.add_node(barrier_node());
917        let d = g.add_node(barrier_node());
918        g.add_edge(a, b).unwrap();
919        g.add_edge(b, c).unwrap();
920        // d is isolated
921        let reach = g.reachable_from(&[a]);
922        assert!(reach.contains(&a));
923        assert!(reach.contains(&b));
924        assert!(reach.contains(&c));
925        assert!(!reach.contains(&d));
926    }
927
928    #[test]
929    fn reaching_set() {
930        let mut g = ComputeGraph::new();
931        let a = g.add_node(barrier_node());
932        let b = g.add_node(barrier_node());
933        let c = g.add_node(barrier_node());
934        g.add_edge(a, b).unwrap();
935        g.add_edge(b, c).unwrap();
936        let reaching = g.reaching(&[c]);
937        assert!(reaching.contains(&a));
938        assert!(reaching.contains(&b));
939        assert!(reaching.contains(&c));
940    }
941
942    // --- Kernel / fusible queries ---
943
944    #[test]
945    fn kernel_nodes_returns_compute_only() {
946        let mut g = ComputeGraph::new();
947        g.add_node(kernel_node("k0"));
948        g.add_node(barrier_node());
949        g.add_node(memcpy_node(MemcpyDir::HostToDevice, 1024));
950        g.add_node(kernel_node("k1"));
951        let kernels = g.kernel_nodes();
952        assert_eq!(kernels.len(), 2);
953    }
954
955    #[test]
956    fn fusible_nodes_returns_fusible_only() {
957        let mut g = ComputeGraph::new();
958        g.add_node(kernel_node("fusible")); // fusible=true in helper
959        g.add_node(GraphNode::new(
960            NodeId(0),
961            NodeKind::KernelLaunch {
962                function_name: "custom".into(),
963                config: KernelConfig::linear(1, 32, 0),
964                fusible: false,
965            },
966        ));
967        assert_eq!(g.fusible_nodes().len(), 1);
968    }
969
970    // --- Parallelism width ---
971
972    #[test]
973    fn parallelism_width_linear_is_one() {
974        let mut g = ComputeGraph::new();
975        let a = g.add_node(barrier_node());
976        let b = g.add_node(barrier_node());
977        let c = g.add_node(barrier_node());
978        g.add_edge(a, b).unwrap();
979        g.add_edge(b, c).unwrap();
980        assert_eq!(g.parallelism_width().unwrap(), 1);
981    }
982
983    #[test]
984    fn parallelism_width_fork_join() {
985        let mut g = ComputeGraph::new();
986        let src = g.add_node(barrier_node());
987        let b = g.add_node(barrier_node());
988        let c = g.add_node(barrier_node());
989        let d = g.add_node(barrier_node());
990        let sink = g.add_node(barrier_node());
991        g.add_edge(src, b).unwrap();
992        g.add_edge(src, c).unwrap();
993        g.add_edge(src, d).unwrap();
994        g.add_edge(b, sink).unwrap();
995        g.add_edge(c, sink).unwrap();
996        g.add_edge(d, sink).unwrap();
997        assert_eq!(g.parallelism_width().unwrap(), 3);
998    }
999
1000    // --- DOT output ---
1001
1002    #[test]
1003    fn to_dot_contains_node_labels() {
1004        let mut g = ComputeGraph::new();
1005        let a = g.add_node(kernel_node("my_kernel").with_name("k0"));
1006        let b = g.add_node(barrier_node());
1007        g.add_edge(a, b).unwrap();
1008        let dot = g.to_dot();
1009        assert!(dot.contains("digraph"));
1010        assert!(dot.contains("k0"));
1011        assert!(dot.contains("->"));
1012    }
1013
1014    // --- Display ---
1015
1016    #[test]
1017    fn display_shows_counts() {
1018        let mut g = ComputeGraph::new();
1019        g.add_node(barrier_node());
1020        g.add_node(barrier_node());
1021        g.add_buffer(BufferDescriptor::new(BufferId(0), 1));
1022        let s = g.to_string();
1023        assert!(s.contains("2 nodes"));
1024        assert!(s.contains("1 buffers"));
1025    }
1026}
oxicuda_graph/graph.rs

oxicuda_graph/
graph.rs
