logicaffeine_compile/codegen_cuda/lane/
graph.rs

1//! The step graph: an authored step (a decode step, a prefill chunk) as the ops it runs and the tensors they
2//! touch, each tensor given the role a compiled lane needs to place it.
3//!
4//! The graph is the step lifted into the one IR ([`TensorGraph::lift`]): every top-level tensor statement is one op,
5//! and a statement the front end cannot lift — or a write it cannot account for — is refused by name. A tensor's role
6//! is the IR's, from how the program uses it:
7//!
8//!   • a **parameter** is a tensor the program declares and reads as a position or as row indices (the token, a
9//!     chunk's tokens, the position): a compiled step takes it at run time;
10//!   • **state** is a declared tensor an op writes in place (a KV cache by `Set row`, a convolution ring, a
11//!     delta-rule memory), carried from one step to the next;
12//!   • a **constant** is any other declared tensor (weights, norm vectors, rotary tables), which the lane's binding
13//!     fills;
14//!   • an **activation** is what an op computes.
15
16use std::collections::{BTreeMap, BTreeSet};
17
18use crate::spatial::passes::recognize_stages::{RecognizedProgram, TensorOpSpec};
19use crate::tensor_graph::graph::{Storage, TensorGraph};
20
21pub use crate::tensor_graph::graph::{GraphError, Role};
22
23/// One tensor of the step: its shape as the program gives it (`rows × cols`), and its role.
24#[derive(Clone, Debug, PartialEq, Eq)]
25pub struct Tensor {
26    pub name: String,
27    pub rows: usize,
28    pub cols: usize,
29    pub role: Role,
30}
31
32impl Tensor {
33    pub fn len(&self) -> usize {
34        self.rows * self.cols
35    }
36}
37
38/// One op of the step, its operands and result by name.
39#[derive(Clone, Debug)]
40pub struct Op {
41    pub spec: TensorOpSpec,
42    pub ins: Vec<String>,
43    pub out: String,
44}
45
46/// An authored step as a graph: its tensors in their roles, its ops, and the step as the one IR holds it, which the
47/// shared optimizer reads.
48#[derive(Clone, Debug)]
49pub struct StepGraph {
50    pub tensors: BTreeMap<String, Tensor>,
51    pub ops: Vec<Op>,
52    pub lifted: TensorGraph,
53}
54
55impl StepGraph {
56    /// Build the graph of a walked program.
57    pub fn build(rp: &RecognizedProgram) -> Result<StepGraph, GraphError> {
58        let lifted = TensorGraph::lift(rp, None, None, &|_| Storage::Values)?;
59        let tensors = lifted
60            .tensors
61            .values()
62            .map(|t| (t.name.clone(), Tensor { name: t.name.clone(), rows: t.rows, cols: t.cols, role: t.role }))
63            .collect();
64        let ops = lifted
65            .ops
66            .iter()
67            .map(|o| Op { spec: o.prim().expect("a lifted op is a primitive").clone(), ins: o.ins.clone(), out: o.out.clone() })
68            .collect();
69        Ok(StepGraph { tensors, ops, lifted })
70    }
71
72    /// The tensors of one role, by name.
73    pub fn of_role(&self, role: Role) -> impl Iterator<Item = &Tensor> {
74        self.tensors.values().filter(move |t| t.role == role)
75    }
76
77    /// The ops that read `tensor`, by index.
78    pub fn readers(&self, tensor: &str) -> Vec<usize> {
79        self.ops.iter().enumerate().filter(|(_, o)| o.ins.iter().any(|x| x == tensor)).map(|(i, _)| i).collect()
80    }
81
82    /// The attention caches: every state tensor an attention reads its keys or values from, itself or through the
83    /// head slices cut from it (a latent cache's values are its rows' first part).
84    pub fn caches(&self) -> BTreeSet<String> {
85        let producer: BTreeMap<&str, &Op> = self.ops.iter().map(|o| (o.out.as_str(), o)).collect();
86        let mut out = BTreeSet::new();
87        for op in self.ops.iter().filter(|o| matches!(o.spec, TensorOpSpec::Attention(_))) {
88            for read in &op.ins[1..3] {
89                let mut t = read.as_str();
90                loop {
91                    if self.tensors.get(t).is_some_and(|x| x.role == Role::State) {
92                        out.insert(t.to_string());
93                        break;
94                    }
95                    match producer.get(t) {
96                        Some(p) if matches!(p.spec, TensorOpSpec::HeadSlice(_)) => t = &p.ins[0],
97                        _ => break,
98                    }
99                }
100            }
101        }
102        out
103    }
104}
logicaffeine_compile/codegen_cuda/lane/graph.rs

logicaffeine_compile/codegen_cuda/lane/
graph.rs
