addition_graph/addition_graph.rs
1//! Minimal end-to-end example: build a graph, bind inputs, execute it, and read outputs.
2//!
3//! This example is the best starting point for understanding the crate's execution model.
4//!
5//! Key ideas:
6//! - A `Graph` is a DAG of tensor operations.
7//! - Input nodes declare required runtime inputs by `NodeId` and shape.
8//! - Operation nodes declare computation to perform later.
9//! - `Executor` evaluates the graph using kernels from a `KernelRegistry`.
10//! - The executor returns a map from output `NodeId` to computed `Tensor`.
11//!
12//! For kernel authors, this example shows the contract at the executor boundary:
13//! - graph construction performs structural validation,
14//! - execution provides concrete `Tensor` inputs,
15//! - kernels are responsible for producing the tensor value for a single op node.
16//!
17//! The Add kernel used here conceptually receives:
18//! - the Add node metadata,
19//! - the already-computed tensors for its input nodes,
20//! - enough context to produce a new output tensor.
21//!
22//! A custom kernel for Add would usually:
23//! - confirm input count and shape assumptions,
24//! - read element values from both input tensors,
25//! - allocate or construct the output tensor,
26//! - return the computed tensor or a `KernelError`.
27use tensor_forge::{Executor, Graph, KernelRegistry, Tensor};
28
29fn main() {
30 // Build the graph:
31 //
32 // a ----\
33 // add ---> out
34 // b ----/
35 //
36 // Here `a` and `b` are graph input nodes. They do not yet have runtime values;
37 // they only declare that the graph expects tensors of shape [2, 2].
38 let mut graph = Graph::new();
39
40 let a = graph.input_node(vec![2, 2]);
41 let b = graph.input_node(vec![2, 2]);
42
43 // Add an operation node.
44 //
45 // `graph.add(a, b)` does not perform arithmetic immediately. It adds a new node to
46 // the graph describing a future Add operation whose inputs are `a` and `b`.
47 //
48 // Shape validation happens here at graph-construction time. Since both inputs are
49 // [2, 2], the resulting Add node is also [2, 2].
50 let out = graph
51 .add(a, b)
52 .expect("Adding valid input nodes should succeed");
53
54 // Mark the node as an output. The executor will return a tensor for every node
55 // designated as a graph output.
56 graph
57 .set_output_node(out)
58 .expect("Setting output node should succeed");
59
60 // Create runtime tensors for the graph input nodes.
61 //
62 // These must match the shapes declared by the corresponding input nodes.
63 let a_tensor = Tensor::from_vec(vec![2, 2], vec![1.0, 2.0, 3.0, 4.0])
64 .expect("Tensor construction should succeed");
65 let b_tensor = Tensor::from_vec(vec![2, 2], vec![10.0, 20.0, 30.0, 40.0])
66 .expect("Tensor construction should succeed");
67
68 // Construct an executor with the default kernel registry.
69 //
70 // The registry determines which kernel implementation is used for each `OpKind`.
71 // In this example, the default registry is expected to contain an Add kernel.
72 let exec = Executor::new(KernelRegistry::default());
73
74 // Execute the graph.
75 //
76 // The bindings are `(NodeId, Tensor)` pairs. Each input node in the graph must be
77 // bound exactly once at runtime.
78 //
79 // Internally, the executor:
80 // 1. validates the bindings,
81 // 2. topologically orders the graph,
82 // 3. executes non-input nodes using registered kernels,
83 // 4. returns tensors for all declared output nodes.
84 let outputs = exec
85 .execute(&graph, vec![(a, a_tensor), (b, b_tensor)])
86 .expect("Execution should succeed");
87
88 let result = outputs
89 .get(&out)
90 .expect("Declared output should be present in executor results");
91
92 println!("Computed output for node {:?}: {:?}", out, result);
93}