spatialrust_runtime/
graph.rs

1//! Bounded spatial operator graph with explicit fusion and transfer receipts.
2
3use std::collections::{HashMap, HashSet, VecDeque};
4
5use spatialrust_distribute::{
6    BackpressurePolicy, BackpressureSignal, NamedTransfer, TransferLedger,
7};
8
9use crate::{RuntimeError, RuntimeResult};
10
11/// Static placement and fusion properties for one operator.
12#[derive(Clone, Debug, PartialEq, Eq)]
13pub struct GraphNodeSpec {
14    /// Stable node identifier.
15    pub id: String,
16    /// Host/device placement label.
17    pub device: String,
18    /// Whether this operator permits queue-eliding fusion with neighbors.
19    pub fusable: bool,
20}
21
22impl GraphNodeSpec {
23    /// Creates a checked node specification.
24    pub fn try_new(
25        id: impl Into<String>,
26        device: impl Into<String>,
27        fusable: bool,
28    ) -> RuntimeResult<Self> {
29        let (id, device) = (id.into(), device.into());
30        if id.is_empty() || device.is_empty() {
31            return Err(RuntimeError::InvalidConfiguration(
32                "graph node id and device must be non-empty".into(),
33            ));
34        }
35        Ok(Self { id, device, fusable })
36    }
37}
38
39/// Stateful transformation executed at one graph node.
40pub trait GraphOperator<T>: Send {
41    /// Processes one owned value without an implicit device transfer.
42    fn process(&mut self, value: T) -> RuntimeResult<T>;
43}
44
45/// Closure-backed graph operator.
46pub struct FnOperator<F>(F);
47
48impl<F> FnOperator<F> {
49    /// Wraps a stateful transformation closure.
50    pub fn new(function: F) -> Self {
51        Self(function)
52    }
53}
54
55impl<T, F> GraphOperator<T> for FnOperator<F>
56where
57    F: FnMut(T) -> RuntimeResult<T> + Send,
58{
59    fn process(&mut self, value: T) -> RuntimeResult<T> {
60        (self.0)(value)
61    }
62}
63
64struct Node<T> {
65    spec: GraphNodeSpec,
66    operator: Box<dyn GraphOperator<T>>,
67}
68
69/// Mutable graph builder. Compilation freezes topology and validates transfers.
70pub struct SpatialExecutionGraph<T> {
71    nodes: HashMap<String, Node<T>>,
72    edges: Vec<(String, String)>,
73    transfers: Vec<NamedTransfer>,
74    pressure: BackpressurePolicy,
75}
76
77impl<T> SpatialExecutionGraph<T> {
78    /// Creates an empty graph with explicit input watermarks.
79    pub fn new(pressure: BackpressurePolicy) -> Self {
80        Self { nodes: HashMap::new(), edges: Vec::new(), transfers: Vec::new(), pressure }
81    }
82
83    /// Adds a uniquely named operator.
84    pub fn add_node(
85        &mut self,
86        spec: GraphNodeSpec,
87        operator: impl GraphOperator<T> + 'static,
88    ) -> RuntimeResult<()> {
89        if self.nodes.contains_key(&spec.id) {
90            return Err(RuntimeError::InvalidConfiguration(format!(
91                "duplicate graph node `{}`",
92                spec.id
93            )));
94        }
95        self.nodes.insert(spec.id.clone(), Node { spec, operator: Box::new(operator) });
96        Ok(())
97    }
98
99    /// Adds a directed data edge.
100    pub fn connect(&mut self, from: impl Into<String>, to: impl Into<String>) -> RuntimeResult<()> {
101        let (from, to) = (from.into(), to.into());
102        if from == to || !self.nodes.contains_key(&from) || !self.nodes.contains_key(&to) {
103            return Err(RuntimeError::InvalidConfiguration(
104                "graph edge requires distinct existing endpoints".into(),
105            ));
106        }
107        if !self.edges.iter().any(|edge| edge == &(from.clone(), to.clone())) {
108            self.edges.push((from, to));
109        }
110        Ok(())
111    }
112
113    /// Declares a named cross-placement transfer for one edge.
114    pub fn declare_transfer(&mut self, transfer: NamedTransfer) {
115        self.transfers.push(transfer);
116    }
117
118    /// Validates topology/placements and builds deterministic fusion groups.
119    pub fn compile(self) -> RuntimeResult<CompiledSpatialGraph<T>> {
120        if self.nodes.is_empty() {
121            return Err(RuntimeError::InvalidConfiguration("execution graph is empty".into()));
122        }
123        let order = topological_order(&self.nodes, &self.edges)?;
124        let sources = order.iter().filter(|id| indegree(id, &self.edges) == 0).count();
125        if sources != 1 {
126            return Err(RuntimeError::InvalidConfiguration(
127                "execution graph requires exactly one source".into(),
128            ));
129        }
130        let transfer_edges = self
131            .transfers
132            .iter()
133            .map(|transfer| (transfer.from.as_str(), transfer.to.as_str()))
134            .collect::<HashSet<_>>();
135        for (from, to) in &self.edges {
136            let source = &self.nodes[from].spec;
137            let target = &self.nodes[to].spec;
138            if source.device != target.device && !transfer_edges.contains(&(from, to)) {
139                return Err(RuntimeError::InvalidConfiguration(format!(
140                    "cross-device edge `{from}` -> `{to}` requires a named transfer"
141                )));
142            }
143        }
144        for transfer in &self.transfers {
145            if !self.edges.iter().any(|edge| edge == &(transfer.from.clone(), transfer.to.clone()))
146            {
147                return Err(RuntimeError::InvalidConfiguration(format!(
148                    "transfer `{}` does not name a graph edge",
149                    transfer.name
150                )));
151            }
152        }
153        let groups = fusion_groups(&order, &self.nodes, &self.edges);
154        Ok(CompiledSpatialGraph {
155            nodes: self.nodes,
156            edges: self.edges,
157            transfers: self.transfers,
158            order,
159            groups,
160            pressure: self.pressure,
161            inputs: VecDeque::new(),
162            soft_trips: 0,
163            hard_rejects: 0,
164        })
165    }
166}
167
168/// Per-run schedule, transfer, and backpressure evidence.
169#[derive(Clone, Debug, Default, PartialEq, Eq)]
170pub struct ExecutionReceipt {
171    /// Nodes executed in deterministic topological order.
172    pub stages: Vec<String>,
173    /// Queue-eliding same-device linear groups (singletons included).
174    pub fusion_groups: Vec<Vec<String>>,
175    /// Named transfers completed during the run.
176    pub transfers: TransferLedger,
177    /// Highest observed input queue depth.
178    pub max_input_depth: usize,
179    /// Cumulative soft-watermark admissions.
180    pub soft_trips: u64,
181    /// Cumulative hard-watermark rejections.
182    pub hard_rejects: u64,
183}
184
185/// Validated graph with a bounded input queue.
186pub struct CompiledSpatialGraph<T> {
187    nodes: HashMap<String, Node<T>>,
188    edges: Vec<(String, String)>,
189    transfers: Vec<NamedTransfer>,
190    order: Vec<String>,
191    groups: Vec<Vec<String>>,
192    pressure: BackpressurePolicy,
193    inputs: VecDeque<T>,
194    soft_trips: u64,
195    hard_rejects: u64,
196}
197
198impl<T: Clone> CompiledSpatialGraph<T> {
199    /// Admits one source value or rejects at the hard watermark.
200    pub fn try_submit(&mut self, value: T) -> RuntimeResult<BackpressureSignal> {
201        let signal = self.pressure.evaluate(self.inputs.len());
202        if signal == BackpressureSignal::HardLimit {
203            self.hard_rejects += 1;
204            return Err(RuntimeError::CapacityExceeded("spatial-graph-input".into()));
205        }
206        if signal == BackpressureSignal::SoftLimit {
207            self.soft_trips += 1;
208        }
209        self.inputs.push_back(value);
210        Ok(self.pressure.evaluate(self.inputs.len()))
211    }
212
213    /// Executes the oldest input and returns one value per sink plus a receipt.
214    pub fn run_next(&mut self) -> RuntimeResult<Option<(Vec<T>, ExecutionReceipt)>> {
215        let Some(value) = self.inputs.pop_front() else {
216            return Ok(None);
217        };
218        let max_input_depth = self.inputs.len() + 1;
219        let source = self
220            .order
221            .iter()
222            .find(|id| indegree(id, &self.edges) == 0)
223            .expect("compiled graph has one source")
224            .clone();
225        let mut pending = HashMap::<String, Vec<T>>::new();
226        pending.insert(source, vec![value]);
227        let mut sinks = Vec::new();
228        let mut ledger = TransferLedger::new();
229        for id in &self.order {
230            let values = pending.remove(id).unwrap_or_default();
231            for value in values {
232                let output =
233                    self.nodes.get_mut(id).expect("compiled node").operator.process(value)?;
234                let successors = successors(id, &self.edges);
235                if successors.is_empty() {
236                    sinks.push(output);
237                    continue;
238                }
239                for successor in &successors {
240                    if let Some(transfer) = self
241                        .transfers
242                        .iter()
243                        .find(|item| item.from == *id && item.to == **successor)
244                    {
245                        ledger.record(transfer.clone());
246                    }
247                    pending.entry((*successor).to_string()).or_default().push(output.clone());
248                }
249            }
250        }
251        Ok(Some((
252            sinks,
253            ExecutionReceipt {
254                stages: self.order.clone(),
255                fusion_groups: self.groups.clone(),
256                transfers: ledger,
257                max_input_depth,
258                soft_trips: self.soft_trips,
259                hard_rejects: self.hard_rejects,
260            },
261        )))
262    }
263
264    /// Returns the compiled queue-elision schedule.
265    pub fn fusion_groups(&self) -> &[Vec<String>] {
266        &self.groups
267    }
268}
269
270fn indegree(id: &str, edges: &[(String, String)]) -> usize {
271    edges.iter().filter(|(_, to)| to == id).count()
272}
273
274fn successors<'a>(id: &str, edges: &'a [(String, String)]) -> Vec<&'a str> {
275    edges.iter().filter_map(|(from, to)| (from == id).then_some(to.as_str())).collect()
276}
277
278fn topological_order<T>(
279    nodes: &HashMap<String, Node<T>>,
280    edges: &[(String, String)],
281) -> RuntimeResult<Vec<String>> {
282    let mut degree =
283        nodes.keys().map(|id| (id.clone(), indegree(id, edges))).collect::<HashMap<_, _>>();
284    let mut order: Vec<String> = Vec::with_capacity(nodes.len());
285    while order.len() < nodes.len() {
286        let next = degree
287            .iter()
288            .filter(|(_, &value)| value == 0)
289            .map(|(id, _)| id)
290            .filter(|id| !order.contains(id))
291            .min()
292            .cloned();
293        let Some(next) = next else {
294            return Err(RuntimeError::InvalidConfiguration(
295                "execution graph contains a cycle".into(),
296            ));
297        };
298        order.push(next.clone());
299        for successor in successors(&next, edges) {
300            *degree.get_mut(successor).expect("validated endpoint") -= 1;
301        }
302    }
303    Ok(order)
304}
305
306fn fusion_groups<T>(
307    order: &[String],
308    nodes: &HashMap<String, Node<T>>,
309    edges: &[(String, String)],
310) -> Vec<Vec<String>> {
311    let mut groups: Vec<Vec<String>> = Vec::new();
312    for id in order {
313        let append = groups.last().and_then(|group| group.last()).is_some_and(|previous| {
314            edges.iter().any(|edge| edge == &(previous.clone(), id.clone()))
315                && successors(previous, edges).len() == 1
316                && indegree(id, edges) == 1
317                && nodes[previous].spec.fusable
318                && nodes[id].spec.fusable
319                && nodes[previous].spec.device == nodes[id].spec.device
320        });
321        if append {
322            groups.last_mut().unwrap().push(id.clone());
323        } else {
324            groups.push(vec![id.clone()]);
325        }
326    }
327    groups
328}
329
330#[cfg(test)]
331mod tests {
332    use super::{FnOperator, GraphNodeSpec, SpatialExecutionGraph};
333    use spatialrust_distribute::{
334        BackpressurePolicy, NamedTransfer, TransferDirection, TransferKind,
335    };
336
337    #[test]
338    fn fuses_local_chain_and_receipts_named_device_copy() {
339        let mut graph = SpatialExecutionGraph::new(BackpressurePolicy::try_new(2, 3).unwrap());
340        for (id, device, add) in [("decode", "cpu", 1), ("gray", "cpu", 2), ("infer", "gpu", 4)] {
341            graph
342                .add_node(
343                    GraphNodeSpec::try_new(id, device, true).unwrap(),
344                    FnOperator::new(move |value: i32| Ok(value + add)),
345                )
346                .unwrap();
347        }
348        graph.connect("decode", "gray").unwrap();
349        graph.connect("gray", "infer").unwrap();
350        graph.declare_transfer(
351            NamedTransfer::try_new(
352                "gray-upload",
353                TransferDirection::HostToDevice,
354                TransferKind::ExplicitCopy,
355                "gray",
356                "infer",
357                1024,
358            )
359            .unwrap(),
360        );
361        let mut compiled = graph.compile().unwrap();
362        assert_eq!(
363            compiled.fusion_groups(),
364            &[vec!["decode".to_string(), "gray".to_string()], vec!["infer".to_string()]]
365        );
366        compiled.try_submit(10).unwrap();
367        let (outputs, receipt) = compiled.run_next().unwrap().unwrap();
368        assert_eq!(outputs, vec![17]);
369        assert_eq!(receipt.transfers.counted_copy_bytes(), 1024);
370        assert_eq!(receipt.transfers.completed()[0].name, "gray-upload");
371    }
372
373    #[test]
374    fn rejects_missing_transfer_cycle_and_full_input() {
375        let mut graph =
376            SpatialExecutionGraph::<i32>::new(BackpressurePolicy::try_new(1, 1).unwrap());
377        graph
378            .add_node(
379                GraphNodeSpec::try_new("a", "cpu", true).unwrap(),
380                FnOperator::new(Ok::<_, crate::RuntimeError>),
381            )
382            .unwrap();
383        graph
384            .add_node(
385                GraphNodeSpec::try_new("b", "gpu", true).unwrap(),
386                FnOperator::new(Ok::<_, crate::RuntimeError>),
387            )
388            .unwrap();
389        graph.connect("a", "b").unwrap();
390        assert!(graph.compile().is_err());
391
392        let mut graph =
393            SpatialExecutionGraph::<i32>::new(BackpressurePolicy::try_new(1, 1).unwrap());
394        graph
395            .add_node(
396                GraphNodeSpec::try_new("a", "cpu", true).unwrap(),
397                FnOperator::new(Ok::<_, crate::RuntimeError>),
398            )
399            .unwrap();
400        let mut compiled = graph.compile().unwrap();
401        compiled.try_submit(1).unwrap();
402        assert!(compiled.try_submit(2).is_err());
403    }
404}
spatialrust_runtime/graph.rs

spatialrust_runtime/
graph.rs
