Skip to main content

oxfuse_runtime/
plan.rs

1//! A verified graph flattened into steps over numbered buffer slots, so the
2//! executor does no name or hash lookups at run time.
3
4use oxfuse_core::{Graph, Op, Source, ValueId};
5use std::sync::Arc;
6
7#[derive(Debug, Clone)]
8pub struct Step {
9    pub op: Op,
10    pub inputs: Vec<usize>,
11    pub output: usize,
12}
13
14#[derive(Debug, Clone)]
15pub struct InputBinding {
16    pub name: String,
17    pub slot: usize,
18    pub shape: Vec<usize>,
19}
20
21#[derive(Debug, Clone)]
22pub struct Plan {
23    pub slot_shapes: Vec<Vec<usize>>,
24    pub constants: Vec<(usize, Arc<[f32]>)>,
25    pub inputs: Vec<InputBinding>,
26    pub steps: Vec<Step>,
27    pub outputs: Vec<usize>,
28}
29
30impl Plan {
31    /// Slot `i` holds value `i`. Values are in definition order, so `steps`
32    /// is already topologically ordered.
33    pub fn build(graph: &Graph) -> Plan {
34        let mut plan = Plan {
35            slot_shapes: Vec::with_capacity(graph.values().len()),
36            constants: Vec::new(),
37            inputs: Vec::new(),
38            steps: Vec::new(),
39            outputs: graph.outputs().iter().map(|id| id.index()).collect(),
40        };
41        for (slot, value) in graph.values().iter().enumerate() {
42            plan.slot_shapes.push(value.ty.shape.clone());
43            match &value.source {
44                Source::Input => plan.inputs.push(InputBinding {
45                    name: graph.label(ValueId(slot as u32)),
46                    slot,
47                    shape: value.ty.shape.clone(),
48                }),
49                Source::Constant(data) => plan.constants.push((slot, Arc::clone(data))),
50                Source::Node { op, inputs } => plan.steps.push(Step {
51                    op: op.clone(),
52                    inputs: inputs.iter().map(|id| id.index()).collect(),
53                    output: slot,
54                }),
55            }
56        }
57        plan
58    }
59}