Skip to main content

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}