torsh_jit/graph/
builder.rs

1//! Graph builder for constructing computation graphs
2
3use crate::graph::core::{ComputationGraph, Edge, Node, NodeId};
4use crate::graph::metadata::GraphMetadata;
5use crate::graph::operations::{ConstantInfo, ConstantValue, Operation, ParameterInfo};
6use crate::{JitError, JitResult};
7use std::collections::HashMap;
8use torsh_core::{DType, DeviceType, Shape};
9
10/// Builder for constructing computation graphs
11#[derive(Debug)]
12pub struct GraphBuilder {
13    graph: ComputationGraph,
14}
15
16impl GraphBuilder {
17    /// Create a new graph builder
18    pub fn new() -> Self {
19        Self {
20            graph: ComputationGraph::new(),
21        }
22    }
23
24    /// Create a new graph builder with metadata
25    pub fn with_metadata(metadata: GraphMetadata) -> Self {
26        let mut graph = ComputationGraph::new();
27        graph.metadata = metadata;
28        Self { graph }
29    }
30
31    /// Add an input node
32    pub fn add_input(&mut self, name: String, shape: Shape, dtype: DType) -> NodeId {
33        let node = Node::new(Operation::Input, name)
34            .with_output_shapes(vec![Some(shape)])
35            .with_dtypes(vec![dtype]);
36        let node_id = self.graph.add_node(node);
37        self.graph.add_input(node_id);
38        node_id
39    }
40
41    /// Add a parameter node
42    pub fn add_parameter(
43        &mut self,
44        name: String,
45        shape: Shape,
46        dtype: DType,
47        trainable: bool,
48    ) -> NodeId {
49        let param_info = ParameterInfo {
50            name: name.clone(),
51            trainable,
52        };
53        let node = Node::new(Operation::Parameter(param_info), name)
54            .with_output_shapes(vec![Some(shape)])
55            .with_dtypes(vec![dtype]);
56        let node_id = self.graph.add_node(node);
57        self.graph.add_input(node_id);
58        node_id
59    }
60
61    /// Add a constant node
62    pub fn add_constant(
63        &mut self,
64        name: String,
65        value: ConstantValue,
66        shape: Shape,
67        dtype: DType,
68    ) -> NodeId {
69        let const_info = ConstantInfo { value };
70        let node = Node::new(Operation::Constant(const_info), name)
71            .with_output_shapes(vec![Some(shape)])
72            .with_dtypes(vec![dtype]);
73        self.graph.add_node(node)
74    }
75
76    /// Add a generic operation node
77    pub fn add_operation(&mut self, name: String, operation: Operation) -> NodeId {
78        let node = Node::new(operation, name);
79        self.graph.add_node(node)
80    }
81
82    /// Add an operation node with shape information
83    pub fn add_operation_with_shapes(
84        &mut self,
85        name: String,
86        operation: Operation,
87        input_shapes: Vec<Option<Shape>>,
88        output_shapes: Vec<Option<Shape>>,
89        dtypes: Vec<DType>,
90    ) -> NodeId {
91        let node = Node::new(operation, name)
92            .with_input_shapes(input_shapes)
93            .with_output_shapes(output_shapes)
94            .with_dtypes(dtypes);
95        self.graph.add_node(node)
96    }
97
98    /// Connect two nodes with an edge
99    pub fn connect(&mut self, from: NodeId, to: NodeId) -> JitResult<()> {
100        self.connect_with_ports(from, 0, to, 0)
101    }
102
103    /// Connect two nodes with specific ports
104    pub fn connect_with_ports(
105        &mut self,
106        from: NodeId,
107        from_output: usize,
108        to: NodeId,
109        to_input: usize,
110    ) -> JitResult<()> {
111        // Validate nodes exist
112        if self.graph.get_node(from).is_none() {
113            return Err(JitError::GraphError(format!(
114                "Source node {:?} does not exist",
115                from
116            )));
117        }
118        if self.graph.get_node(to).is_none() {
119            return Err(JitError::GraphError(format!(
120                "Destination node {:?} does not exist",
121                to
122            )));
123        }
124
125        let edge = Edge {
126            src_output: from_output,
127            dst_input: to_input,
128        };
129        self.graph.add_edge(from, to, edge);
130        Ok(())
131    }
132
133    /// Mark a node as output
134    pub fn mark_output(&mut self, node_id: NodeId) -> JitResult<()> {
135        if self.graph.get_node(node_id).is_none() {
136            return Err(JitError::GraphError(format!(
137                "Node {:?} does not exist",
138                node_id
139            )));
140        }
141        self.graph.add_output(node_id);
142        Ok(())
143    }
144
145    /// Set node shape information
146    pub fn set_node_shapes(
147        &mut self,
148        node_id: NodeId,
149        input_shapes: Vec<Option<Shape>>,
150        output_shapes: Vec<Option<Shape>>,
151    ) -> JitResult<()> {
152        if let Some(node) = self.graph.get_node_mut(node_id) {
153            node.input_shapes = input_shapes;
154            node.output_shapes = output_shapes;
155            Ok(())
156        } else {
157            Err(JitError::GraphError(format!(
158                "Node {:?} does not exist",
159                node_id
160            )))
161        }
162    }
163
164    /// Set node data types
165    pub fn set_node_dtypes(&mut self, node_id: NodeId, dtypes: Vec<DType>) -> JitResult<()> {
166        if let Some(node) = self.graph.get_node_mut(node_id) {
167            node.dtypes = dtypes;
168            Ok(())
169        } else {
170            Err(JitError::GraphError(format!(
171                "Node {:?} does not exist",
172                node_id
173            )))
174        }
175    }
176
177    /// Set node device
178    pub fn set_node_device(&mut self, node_id: NodeId, device: DeviceType) -> JitResult<()> {
179        if let Some(node) = self.graph.get_node_mut(node_id) {
180            node.device = device;
181            Ok(())
182        } else {
183            Err(JitError::GraphError(format!(
184                "Node {:?} does not exist",
185                node_id
186            )))
187        }
188    }
189
190    /// Build and return the computation graph
191    pub fn build(self) -> JitResult<ComputationGraph> {
192        self.graph.validate()?;
193        Ok(self.graph)
194    }
195
196    /// Build without validation (for testing purposes)
197    pub fn build_unchecked(self) -> ComputationGraph {
198        self.graph
199    }
200
201    /// Get a reference to the current graph being built
202    pub fn graph(&self) -> &ComputationGraph {
203        &self.graph
204    }
205
206    /// Get a mutable reference to the current graph being built
207    pub fn graph_mut(&mut self) -> &mut ComputationGraph {
208        &mut self.graph
209    }
210
211    /// Clone the current state of the graph
212    pub fn clone_graph(&self) -> ComputationGraph {
213        self.graph.clone()
214    }
215
216    // Convenience methods for common operations
217
218    /// Add an element-wise binary operation
219    pub fn add_binary_op(
220        &mut self,
221        name: String,
222        operation: Operation,
223        left: NodeId,
224        right: NodeId,
225    ) -> JitResult<NodeId> {
226        let node_id = self.add_operation(name, operation);
227        self.connect(left, node_id)?;
228        self.connect(right, node_id)?;
229        Ok(node_id)
230    }
231
232    /// Add an element-wise unary operation
233    pub fn add_unary_op(
234        &mut self,
235        name: String,
236        operation: Operation,
237        input: NodeId,
238    ) -> JitResult<NodeId> {
239        let node_id = self.add_operation(name, operation);
240        self.connect(input, node_id)?;
241        Ok(node_id)
242    }
243
244    /// Add a reduction operation
245    pub fn add_reduction_op(
246        &mut self,
247        name: String,
248        operation: Operation,
249        input: NodeId,
250    ) -> JitResult<NodeId> {
251        let node_id = self.add_operation(name, operation);
252        self.connect(input, node_id)?;
253        Ok(node_id)
254    }
255
256    /// Create a linear chain of operations
257    pub fn create_linear_chain(
258        &mut self,
259        operations: Vec<(String, Operation)>,
260        input: NodeId,
261    ) -> JitResult<NodeId> {
262        let mut current = input;
263        for (name, op) in operations {
264            current = self.add_unary_op(name, op, current)?;
265        }
266        Ok(current)
267    }
268
269    /// Create a residual connection (input + f(input))
270    pub fn create_residual_connection(
271        &mut self,
272        input: NodeId,
273        transform_ops: Vec<(String, Operation)>,
274    ) -> JitResult<NodeId> {
275        let transformed = self.create_linear_chain(transform_ops, input)?;
276        self.add_binary_op(
277            "residual_add".to_string(),
278            Operation::Add,
279            input,
280            transformed,
281        )
282    }
283
284    /// Validate the current graph state
285    pub fn validate(&self) -> JitResult<()> {
286        self.graph.validate()
287    }
288
289    /// Get statistics about the current graph
290    pub fn statistics(&self) -> GraphStatistics {
291        GraphStatistics::from_graph(&self.graph)
292    }
293}
294
295impl Default for GraphBuilder {
296    fn default() -> Self {
297        Self::new()
298    }
299}
300
301/// Statistics about a computation graph
302#[derive(Debug, Clone)]
303pub struct GraphStatistics {
304    pub node_count: usize,
305    pub edge_count: usize,
306    pub input_count: usize,
307    pub output_count: usize,
308    pub operation_counts: HashMap<String, usize>,
309    pub memory_estimate: usize,
310    pub complexity_estimate: usize,
311}
312
313impl GraphStatistics {
314    /// Generate statistics from a computation graph
315    pub fn from_graph(graph: &ComputationGraph) -> Self {
316        let mut operation_counts = HashMap::new();
317
318        for (_, node) in graph.nodes() {
319            let op_name = node.operation.as_str().to_string();
320            *operation_counts.entry(op_name).or_insert(0) += 1;
321        }
322
323        Self {
324            node_count: graph.node_count(),
325            edge_count: graph.edge_count(),
326            input_count: graph.inputs.len(),
327            output_count: graph.outputs.len(),
328            operation_counts,
329            memory_estimate: graph.memory_estimate(),
330            complexity_estimate: graph.complexity_estimate(),
331        }
332    }
333}
torsh_jit/graph/builder.rs

torsh_jit/graph/
builder.rs
