1use 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 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}