torsh_tensor/
computation_graph.rs

1//! Computation Graph for Lazy Evaluation and Optimization
2//!
3//! This module provides a computation graph system that enables lazy evaluation,
4//! operation fusion, and various graph-level optimizations before execution.
5//!
6//! # Features
7//!
8//! - **Lazy evaluation**: Build computation graph without executing
9//! - **Graph optimization**: Fuse operations, eliminate dead code, constant folding
10//! - **Memory planning**: Optimize memory allocation and reuse
11//! - **Parallel scheduling**: Automatic parallelization of independent operations
12//! - **Visualization**: Generate DOT graphs for debugging
13
14use std::collections::{HashMap, HashSet, VecDeque};
15use std::fmt;
16use std::sync::{Arc, Mutex};
17use torsh_core::sync::MutexExt;
18
19use torsh_core::{
20    device::DeviceType,
21    dtype::TensorElement,
22    error::{Result, TorshError},
23};
24
25use crate::Tensor;
26
27/// Unique identifier for graph nodes
28pub type NodeId = usize;
29
30/// Operation types in the computation graph
31#[derive(Debug, Clone)]
32pub enum GraphOp {
33    /// Input/constant tensor
34    Constant,
35    /// Element-wise addition
36    Add,
37    /// Element-wise multiplication
38    Mul,
39    /// Element-wise subtraction
40    Sub,
41    /// Element-wise division
42    Div,
43    /// Matrix multiplication
44    MatMul,
45    /// Reshape operation
46    Reshape(Vec<usize>),
47    /// Transpose operation
48    Transpose(usize, usize),
49    /// Reduction sum
50    Sum(Option<i32>),
51    /// Reduction mean
52    Mean(Option<i32>),
53    /// ReLU activation
54    ReLU,
55    /// Sigmoid activation
56    Sigmoid,
57    /// Tanh activation
58    Tanh,
59    /// Scalar addition
60    AddScalar(f64),
61    /// Scalar multiplication
62    MulScalar(f64),
63    /// Custom operation
64    Custom(String),
65}
66
67impl fmt::Display for GraphOp {
68    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
69        match self {
70            GraphOp::Constant => write!(f, "Const"),
71            GraphOp::Add => write!(f, "Add"),
72            GraphOp::Mul => write!(f, "Mul"),
73            GraphOp::Sub => write!(f, "Sub"),
74            GraphOp::Div => write!(f, "Div"),
75            GraphOp::MatMul => write!(f, "MatMul"),
76            GraphOp::Reshape(shape) => write!(f, "Reshape({:?})", shape),
77            GraphOp::Transpose(d0, d1) => write!(f, "Transpose({}, {})", d0, d1),
78            GraphOp::Sum(dim) => write!(f, "Sum({:?})", dim),
79            GraphOp::Mean(dim) => write!(f, "Mean({:?})", dim),
80            GraphOp::ReLU => write!(f, "ReLU"),
81            GraphOp::Sigmoid => write!(f, "Sigmoid"),
82            GraphOp::Tanh => write!(f, "Tanh"),
83            GraphOp::AddScalar(s) => write!(f, "AddScalar({})", s),
84            GraphOp::MulScalar(s) => write!(f, "MulScalar({})", s),
85            GraphOp::Custom(name) => write!(f, "Custom({})", name),
86        }
87    }
88}
89
90/// Node in the computation graph
91#[derive(Clone)]
92pub struct GraphNode<T: TensorElement> {
93    /// Unique node ID
94    pub id: NodeId,
95    /// Operation type
96    pub op: GraphOp,
97    /// Input node IDs
98    pub inputs: Vec<NodeId>,
99    /// Cached tensor data (for constants)
100    pub data: Option<Arc<Tensor<T>>>,
101    /// Output shape (if known)
102    pub shape: Option<Vec<usize>>,
103    /// Device
104    pub device: DeviceType,
105}
106
107impl<T: TensorElement> GraphNode<T> {
108    /// Create a new graph node
109    fn new(id: NodeId, op: GraphOp, inputs: Vec<NodeId>, device: DeviceType) -> Self {
110        Self {
111            id,
112            op,
113            inputs,
114            data: None,
115            shape: None,
116            device,
117        }
118    }
119
120    /// Create a constant node
121    fn constant(id: NodeId, tensor: Tensor<T>) -> Self {
122        let device = tensor.device;
123        let shape = Some(tensor.shape().dims().to_vec());
124        Self {
125            id,
126            op: GraphOp::Constant,
127            inputs: Vec::new(),
128            data: Some(Arc::new(tensor)),
129            shape,
130            device,
131        }
132    }
133}
134
135/// Computation graph
136pub struct ComputationGraph<T: TensorElement> {
137    /// All nodes in the graph
138    nodes: HashMap<NodeId, GraphNode<T>>,
139    /// Next available node ID
140    next_id: NodeId,
141    /// Output nodes (nodes that need to be computed)
142    outputs: Vec<NodeId>,
143    /// Execution cache (computed results)
144    cache: Arc<Mutex<HashMap<NodeId, Arc<Tensor<T>>>>>,
145}
146
147impl<T: TensorElement + Copy> ComputationGraph<T> {
148    /// Create a new empty computation graph
149    pub fn new() -> Self {
150        Self {
151            nodes: HashMap::new(),
152            next_id: 0,
153            outputs: Vec::new(),
154            cache: Arc::new(Mutex::new(HashMap::new())),
155        }
156    }
157
158    /// Add a constant tensor to the graph
159    pub fn constant(&mut self, tensor: Tensor<T>) -> NodeId {
160        let id = self.allocate_id();
161        let node = GraphNode::constant(id, tensor);
162        self.nodes.insert(id, node);
163        id
164    }
165
166    /// Add a binary operation node
167    pub fn binary_op(
168        &mut self,
169        op: GraphOp,
170        left: NodeId,
171        right: NodeId,
172        device: DeviceType,
173    ) -> NodeId {
174        let id = self.allocate_id();
175        let node = GraphNode::new(id, op, vec![left, right], device);
176        self.nodes.insert(id, node);
177        id
178    }
179
180    /// Add a unary operation node
181    pub fn unary_op(&mut self, op: GraphOp, input: NodeId, device: DeviceType) -> NodeId {
182        let id = self.allocate_id();
183        let node = GraphNode::new(id, op, vec![input], device);
184        self.nodes.insert(id, node);
185        id
186    }
187
188    /// Mark a node as an output
189    pub fn mark_output(&mut self, node: NodeId) {
190        if !self.outputs.contains(&node) {
191            self.outputs.push(node);
192        }
193    }
194
195    /// Get the number of nodes
196    pub fn num_nodes(&self) -> usize {
197        self.nodes.len()
198    }
199
200    /// Get the number of output nodes
201    pub fn num_outputs(&self) -> usize {
202        self.outputs.len()
203    }
204
205    /// Allocate a new node ID
206    fn allocate_id(&mut self) -> NodeId {
207        let id = self.next_id;
208        self.next_id += 1;
209        id
210    }
211
212    /// Perform topological sort of the graph
213    pub fn topological_sort(&self) -> Result<Vec<NodeId>> {
214        let mut in_degree: HashMap<NodeId, usize> = HashMap::new();
215        let mut adj_list: HashMap<NodeId, Vec<NodeId>> = HashMap::new();
216
217        // Build adjacency list and calculate in-degrees
218        for (&id, node) in &self.nodes {
219            in_degree.entry(id).or_insert(0);
220            for &input_id in &node.inputs {
221                adj_list.entry(input_id).or_insert_with(Vec::new).push(id);
222                *in_degree.entry(id).or_insert(0) += 1;
223            }
224        }
225
226        // Kahn's algorithm
227        let mut queue: VecDeque<NodeId> = in_degree
228            .iter()
229            .filter(|(_, &degree)| degree == 0)
230            .map(|(&id, _)| id)
231            .collect();
232
233        let mut sorted = Vec::new();
234
235        while let Some(node_id) = queue.pop_front() {
236            sorted.push(node_id);
237
238            if let Some(neighbors) = adj_list.get(&node_id) {
239                for &neighbor in neighbors {
240                    if let Some(degree) = in_degree.get_mut(&neighbor) {
241                        *degree -= 1;
242                        if *degree == 0 {
243                            queue.push_back(neighbor);
244                        }
245                    }
246                }
247            }
248        }
249
250        if sorted.len() != self.nodes.len() {
251            return Err(TorshError::InvalidArgument(
252                "Graph contains cycles".to_string(),
253            ));
254        }
255
256        Ok(sorted)
257    }
258
259    /// Optimize the graph by fusing operations
260    pub fn optimize(&mut self) -> Result<()>
261    where
262        T: std::ops::Add<Output = T>
263            + std::ops::Sub<Output = T>
264            + std::ops::Mul<Output = T>
265            + std::ops::Div<Output = T>
266            + torsh_core::FloatElement,
267    {
268        // Simple optimizations:
269        // 1. Constant folding
270        self.fold_constants()?;
271
272        // 2. Dead code elimination
273        self.eliminate_dead_code();
274
275        // 3. Operation fusion (future enhancement)
276
277        Ok(())
278    }
279
280    /// Fold constant operations
281    fn fold_constants(&mut self) -> Result<()>
282    where
283        T: std::ops::Add<Output = T>
284            + std::ops::Sub<Output = T>
285            + std::ops::Mul<Output = T>
286            + std::ops::Div<Output = T>
287            + torsh_core::FloatElement,
288    {
289        let sorted = self.topological_sort()?;
290
291        for &node_id in &sorted {
292            let node = self
293                .nodes
294                .get(&node_id)
295                .expect("node_id should exist in nodes after topological sort")
296                .clone();
297
298            // Check if all inputs are constants
299            let all_constant = node.inputs.iter().all(|&input_id| {
300                if let Some(input_node) = self.nodes.get(&input_id) {
301                    matches!(input_node.op, GraphOp::Constant)
302                } else {
303                    false
304                }
305            });
306
307            if all_constant && !node.inputs.is_empty() {
308                // Try to evaluate this node
309                if let Ok(result) = self.evaluate_node_internal(&node) {
310                    // Replace with constant
311                    let mut new_node = GraphNode::constant(node_id, result);
312                    new_node.device = node.device;
313                    self.nodes.insert(node_id, new_node);
314                }
315            }
316        }
317
318        Ok(())
319    }
320
321    /// Eliminate dead code (nodes not contributing to outputs)
322    fn eliminate_dead_code(&mut self) {
323        let mut reachable = HashSet::new();
324        let mut queue = VecDeque::from_iter(self.outputs.iter().copied());
325
326        // Mark all reachable nodes
327        while let Some(node_id) = queue.pop_front() {
328            if reachable.insert(node_id) {
329                if let Some(node) = self.nodes.get(&node_id) {
330                    for &input_id in &node.inputs {
331                        queue.push_back(input_id);
332                    }
333                }
334            }
335        }
336
337        // Remove unreachable nodes
338        self.nodes.retain(|&id, _| reachable.contains(&id));
339    }
340
341    /// Evaluate a single node
342    fn evaluate_node_internal(&self, node: &GraphNode<T>) -> Result<Tensor<T>>
343    where
344        T: std::ops::Add<Output = T>
345            + std::ops::Sub<Output = T>
346            + std::ops::Mul<Output = T>
347            + std::ops::Div<Output = T>
348            + torsh_core::FloatElement,
349    {
350        match &node.op {
351            GraphOp::Constant => node
352                .data
353                .as_ref()
354                .map(|t| (**t).clone())
355                .ok_or_else(|| TorshError::InvalidArgument("Constant has no data".to_string())),
356            GraphOp::Add => {
357                let left = self.get_input_tensor(node, 0)?;
358                let right = self.get_input_tensor(node, 1)?;
359                left.add_op(&right)
360            }
361            GraphOp::Mul => {
362                let left = self.get_input_tensor(node, 0)?;
363                let right = self.get_input_tensor(node, 1)?;
364                left.mul_op(&right)
365            }
366            GraphOp::Sub => {
367                let left = self.get_input_tensor(node, 0)?;
368                let right = self.get_input_tensor(node, 1)?;
369                left.sub(&right)
370            }
371            GraphOp::Div => {
372                let left = self.get_input_tensor(node, 0)?;
373                let right = self.get_input_tensor(node, 1)?;
374                left.div(&right)
375            }
376            GraphOp::AddScalar(s) => {
377                let input = self.get_input_tensor(node, 0)?;
378                let scalar = T::from_f64(*s).ok_or_else(|| {
379                    TorshError::InvalidArgument("Cannot convert scalar to tensor type".to_string())
380                })?;
381                input.add_scalar(scalar)
382            }
383            GraphOp::MulScalar(s) => {
384                let input = self.get_input_tensor(node, 0)?;
385                let scalar = T::from_f64(*s).ok_or_else(|| {
386                    TorshError::InvalidArgument("Cannot convert scalar to tensor type".to_string())
387                })?;
388                input.mul_scalar(scalar)
389            }
390            GraphOp::ReLU => {
391                let input = self.get_input_tensor(node, 0)?;
392                input.relu()
393            }
394            GraphOp::Sigmoid => {
395                let input = self.get_input_tensor(node, 0)?;
396                input.sigmoid()
397            }
398            GraphOp::Tanh => {
399                let input = self.get_input_tensor(node, 0)?;
400                input.tanh()
401            }
402            _ => Err(TorshError::InvalidArgument(format!(
403                "Unsupported operation: {}",
404                node.op
405            ))),
406        }
407    }
408
409    /// Get input tensor for a node
410    fn get_input_tensor(&self, node: &GraphNode<T>, index: usize) -> Result<Tensor<T>> {
411        let input_id = node.inputs.get(index).ok_or_else(|| {
412            TorshError::InvalidArgument(format!("Missing input {} for node {}", index, node.id))
413        })?;
414
415        let input_node = self.nodes.get(input_id).ok_or_else(|| {
416            TorshError::InvalidArgument(format!("Input node {} not found", input_id))
417        })?;
418
419        if let GraphOp::Constant = input_node.op {
420            input_node
421                .data
422                .as_ref()
423                .map(|t| (**t).clone())
424                .ok_or_else(|| TorshError::InvalidArgument("Constant has no data".to_string()))
425        } else {
426            Err(TorshError::InvalidArgument(
427                "Can only evaluate constants in internal evaluation".to_string(),
428            ))
429        }
430    }
431
432    /// Execute the graph and get outputs
433    pub fn execute(&self) -> Result<Vec<Tensor<T>>>
434    where
435        T: std::ops::Add<Output = T>
436            + std::ops::Sub<Output = T>
437            + std::ops::Mul<Output = T>
438            + std::ops::Div<Output = T>
439            + torsh_core::FloatElement,
440    {
441        let sorted = self.topological_sort()?;
442        let mut cache = self.cache.lock_or_recover();
443        cache.clear();
444
445        // Evaluate nodes in topological order
446        for &node_id in &sorted {
447            let node = self
448                .nodes
449                .get(&node_id)
450                .expect("node_id should exist in nodes after topological sort");
451
452            // Skip if already cached
453            if cache.contains_key(&node_id) {
454                continue;
455            }
456
457            let result = self.evaluate_node_internal(node)?;
458            cache.insert(node_id, Arc::new(result));
459        }
460
461        // Collect outputs
462        let mut outputs = Vec::new();
463        for &output_id in &self.outputs {
464            if let Some(result) = cache.get(&output_id) {
465                outputs.push((**result).clone());
466            } else {
467                return Err(TorshError::InvalidArgument(format!(
468                    "Output node {} not computed",
469                    output_id
470                )));
471            }
472        }
473
474        Ok(outputs)
475    }
476
477    /// Generate DOT representation for visualization
478    pub fn to_dot(&self) -> String {
479        let mut dot = String::from("digraph ComputationGraph {\n");
480        dot.push_str("  rankdir=BT;\n");
481        dot.push_str("  node [shape=box];\n\n");
482
483        // Add nodes
484        for (id, node) in &self.nodes {
485            let label = format!("{}\\nid={}", node.op, id);
486            let color = if self.outputs.contains(id) {
487                "red"
488            } else if matches!(node.op, GraphOp::Constant) {
489                "lightblue"
490            } else {
491                "lightgray"
492            };
493
494            dot.push_str(&format!(
495                "  {} [label=\"{}\", fillcolor={}, style=filled];\n",
496                id, label, color
497            ));
498        }
499
500        dot.push('\n');
501
502        // Add edges
503        for (id, node) in &self.nodes {
504            for (idx, &input_id) in node.inputs.iter().enumerate() {
505                dot.push_str(&format!("  {} -> {} [label=\"{}\"];\n", input_id, id, idx));
506            }
507        }
508
509        dot.push_str("}\n");
510        dot
511    }
512}
513
514impl<T: TensorElement + Copy> Default for ComputationGraph<T> {
515    fn default() -> Self {
516        Self::new()
517    }
518}
519
520#[cfg(test)]
521mod tests {
522    use super::*;
523    use crate::creation::*;
524
525    #[test]
526    fn test_graph_creation() {
527        let mut graph = ComputationGraph::<f32>::new();
528
529        let a = tensor_1d(&[1.0, 2.0, 3.0]).expect("tensor_1d creation should succeed");
530        let b = tensor_1d(&[4.0, 5.0, 6.0]).expect("tensor_1d creation should succeed");
531
532        let a_id = graph.constant(a);
533        let b_id = graph.constant(b);
534        let add_id = graph.binary_op(GraphOp::Add, a_id, b_id, DeviceType::Cpu);
535
536        graph.mark_output(add_id);
537
538        assert_eq!(graph.num_nodes(), 3);
539        assert_eq!(graph.num_outputs(), 1);
540    }
541
542    #[test]
543    fn test_topological_sort() {
544        let mut graph = ComputationGraph::<f32>::new();
545
546        let a = tensor_1d(&[1.0, 2.0]).expect("tensor_1d creation should succeed");
547        let b = tensor_1d(&[3.0, 4.0]).expect("tensor_1d creation should succeed");
548
549        let a_id = graph.constant(a);
550        let b_id = graph.constant(b);
551        let add_id = graph.binary_op(GraphOp::Add, a_id, b_id, DeviceType::Cpu);
552        let mul_id = graph.unary_op(GraphOp::MulScalar(2.0), add_id, DeviceType::Cpu);
553
554        let sorted = graph
555            .topological_sort()
556            .expect("topological sort should succeed");
557
558        // Should have all 4 nodes
559        assert_eq!(sorted.len(), 4);
560
561        // Constants should come before operations that use them
562        let a_pos = sorted
563            .iter()
564            .position(|&id| id == a_id)
565            .expect("position should succeed");
566        let b_pos = sorted
567            .iter()
568            .position(|&id| id == b_id)
569            .expect("position should succeed");
570        let add_pos = sorted
571            .iter()
572            .position(|&id| id == add_id)
573            .expect("position should succeed");
574        let mul_pos = sorted
575            .iter()
576            .position(|&id| id == mul_id)
577            .expect("position should succeed");
578
579        assert!(a_pos < add_pos);
580        assert!(b_pos < add_pos);
581        assert!(add_pos < mul_pos);
582    }
583
584    #[test]
585    fn test_constant_folding() {
586        let mut graph = ComputationGraph::<f32>::new();
587
588        let a = tensor_1d(&[1.0, 2.0]).expect("tensor_1d creation should succeed");
589        let b = tensor_1d(&[3.0, 4.0]).expect("tensor_1d creation should succeed");
590
591        let a_id = graph.constant(a);
592        let b_id = graph.constant(b);
593        let add_id = graph.binary_op(GraphOp::Add, a_id, b_id, DeviceType::Cpu);
594
595        graph.mark_output(add_id);
596
597        // Before optimization
598        assert_eq!(graph.num_nodes(), 3);
599
600        // Optimize
601        graph.optimize().expect("optimization should succeed");
602
603        // After optimization, the add should be folded into a constant
604        let add_node = graph.nodes.get(&add_id).expect("get should succeed");
605        assert!(matches!(add_node.op, GraphOp::Constant));
606    }
607
608    #[test]
609    fn test_dead_code_elimination() {
610        let mut graph = ComputationGraph::<f32>::new();
611
612        let a = tensor_1d(&[1.0]).expect("tensor_1d creation should succeed");
613        let b = tensor_1d(&[2.0]).expect("tensor_1d creation should succeed");
614        let c = tensor_1d(&[3.0]).expect("tensor_1d creation should succeed");
615
616        let a_id = graph.constant(a);
617        let b_id = graph.constant(b);
618        let c_id = graph.constant(c);
619
620        // Create used operation
621        let add_id = graph.binary_op(GraphOp::Add, a_id, b_id, DeviceType::Cpu);
622        graph.mark_output(add_id);
623
624        // Create unused operation (dead code)
625        let _mul_id = graph.unary_op(GraphOp::MulScalar(2.0), c_id, DeviceType::Cpu);
626
627        // Before optimization
628        assert_eq!(graph.num_nodes(), 5);
629
630        // Optimize
631        graph.optimize().expect("optimization should succeed");
632
633        // After optimization, unused nodes should be removed
634        // Note: constant folding folds add into a single constant, so we get 1 node
635        assert_eq!(graph.num_nodes(), 1); // The folded constant
636    }
637
638    #[test]
639    fn test_graph_execution() {
640        let mut graph = ComputationGraph::<f32>::new();
641
642        let a = tensor_1d(&[1.0, 2.0, 3.0]).expect("tensor_1d creation should succeed");
643        let b = tensor_1d(&[4.0, 5.0, 6.0]).expect("tensor_1d creation should succeed");
644
645        let a_id = graph.constant(a);
646        let b_id = graph.constant(b);
647        let add_id = graph.binary_op(GraphOp::Add, a_id, b_id, DeviceType::Cpu);
648
649        graph.mark_output(add_id);
650
651        let results = graph.execute().expect("execution should succeed");
652        assert_eq!(results.len(), 1);
653
654        let data = results[0]
655            .to_vec()
656            .expect("to_vec conversion should succeed");
657        assert_eq!(data, vec![5.0, 7.0, 9.0]);
658    }
659
660    #[test]
661    fn test_multiple_outputs() {
662        let mut graph = ComputationGraph::<f32>::new();
663
664        let a = tensor_1d(&[1.0, 2.0]).expect("tensor_1d creation should succeed");
665        let b = tensor_1d(&[3.0, 4.0]).expect("tensor_1d creation should succeed");
666
667        let a_id = graph.constant(a);
668        let b_id = graph.constant(b);
669        let add_id = graph.binary_op(GraphOp::Add, a_id, b_id, DeviceType::Cpu);
670        let mul_id = graph.binary_op(GraphOp::Mul, a_id, b_id, DeviceType::Cpu);
671
672        graph.mark_output(add_id);
673        graph.mark_output(mul_id);
674
675        let results = graph.execute().expect("execution should succeed");
676        assert_eq!(results.len(), 2);
677
678        let add_data = results[0]
679            .to_vec()
680            .expect("to_vec conversion should succeed");
681        let mul_data = results[1]
682            .to_vec()
683            .expect("to_vec conversion should succeed");
684
685        assert_eq!(add_data, vec![4.0, 6.0]);
686        assert_eq!(mul_data, vec![3.0, 8.0]);
687    }
688
689    #[test]
690    fn test_dot_generation() {
691        let mut graph = ComputationGraph::<f32>::new();
692
693        let a = tensor_1d(&[1.0]).expect("tensor_1d creation should succeed");
694        let a_id = graph.constant(a);
695        graph.mark_output(a_id);
696
697        let dot = graph.to_dot();
698
699        assert!(dot.contains("digraph ComputationGraph"));
700        assert!(dot.contains(&format!("id={}", a_id)));
701    }
702}
torsh_tensor/computation_graph.rs

torsh_tensor/
computation_graph.rs
