Skip to main content

cu29_runtime/planner/
pipeline.rs

1//! Mechanical stage placement using the shared schedule representation.
2
3use super::CuPlan;
4use super::CuPlanDependency;
5use super::CuPlanPlacement;
6use super::CuPlanWorker;
7use super::CuPlanner;
8use super::PlanningInput;
9use super::Serial;
10use super::assemble_runtime_plan_with_planner;
11use super::mission_graphs;
12use super::schedule::PlanShape;
13use crate::config::ComponentConfig;
14use crate::config::RT_POOL;
15use crate::config::SchedulingPolicy;
16use crate::curuntime::CuExecutionUnit;
17use crate::curuntime::CuStepPhase;
18use alloc::collections::BTreeMap;
19use alloc::format;
20use alloc::vec;
21use alloc::vec::Vec;
22use cu29_traits::CuError;
23use cu29_traits::CuResult;
24
25/// Place each process stage on a dedicated worker in graph order.
26pub struct Pipeline {
27    max_in_flight: Option<u32>,
28}
29
30impl Pipeline {
31    pub(super) fn new(params: Option<&ComponentConfig>) -> CuResult<Self> {
32        let max_in_flight = params
33            .map(|params| params.get_value::<u32>("max_in_flight"))
34            .transpose()
35            .map_err(|error| CuError::from(format!("Pipeline: {error}")))?
36            .flatten();
37        Ok(Self { max_in_flight })
38    }
39}
40
41impl CuPlanner for Pipeline {
42    fn plan(&self, input: &PlanningInput<'_>) -> CuResult<CuPlan> {
43        let config = input.config;
44        let capacity = config
45            .logging
46            .as_ref()
47            .and_then(|logging| logging.copperlist_count)
48            .unwrap_or(super::DEFAULT_COPPERLIST_COUNT);
49        let max_in_flight = self.max_in_flight.unwrap_or(
50            u32::try_from(capacity).map_err(|_| CuError::from("Pipeline capacity exceeds u32"))?,
51        );
52        let pool = config
53            .runtime
54            .as_ref()
55            .and_then(|runtime| runtime.thread_pools.iter().find(|pool| pool.id == RT_POOL));
56        let mut missions = BTreeMap::new();
57        for (mission, graph) in mission_graphs(config) {
58            let mut assembled = assemble_runtime_plan_with_planner(config, graph, &Serial)?;
59            let mut groups = BTreeMap::new();
60            let mut order = Vec::new();
61            for unit in core::mem::take(&mut assembled.execution.steps) {
62                let CuExecutionUnit::Step(step) = &unit else {
63                    return Err(CuError::from("Pipeline requires process steps"));
64                };
65                if step.phase != CuStepPhase::AnytimeRefine {
66                    order.push(step.node_id);
67                }
68                groups
69                    .entry(step.node_id)
70                    .or_insert_with(Vec::new)
71                    .push(unit);
72            }
73            assembled.execution.steps = order
74                .into_iter()
75                .flat_map(|node| groups.remove(&node).expect("stage group exists"))
76                .collect();
77            let mut plan =
78                PlanShape::new(&assembled, config, graph, &mission, &[])?.serial_plan()?;
79            plan.max_in_flight = max_in_flight;
80            let mut workers: Vec<CuPlanWorker> = Vec::new();
81            for (index, unit) in assembled.execution.steps.iter().enumerate() {
82                let CuExecutionUnit::Step(step) = unit else {
83                    return Err(CuError::from("Pipeline requires process steps"));
84                };
85                let occurrence = u32::try_from(index)
86                    .map_err(|_| CuError::from("Pipeline has too many steps"))?;
87                if step.phase == CuStepPhase::AnytimeRefine {
88                    workers
89                        .last_mut()
90                        .ok_or_else(|| CuError::from("Refinement has no base stage"))?
91                        .steps
92                        .push(occurrence);
93                } else {
94                    let worker = workers.len();
95                    let cpu = pool
96                        .and_then(|pool| pool.affinity.as_ref())
97                        .filter(|cpus| !cpus.is_empty())
98                        .map(|cpus| cpus[worker % cpus.len()]);
99                    workers.push(CuPlanWorker {
100                        id: format!("stage_{worker}"),
101                        placement: CuPlanPlacement::Thread {
102                            cpu,
103                            policy: pool.map_or(SchedulingPolicy::Fair, |pool| pool.policy),
104                        },
105                        steps: vec![occurrence],
106                    });
107                }
108                if occurrence > 0 {
109                    let edge = CuPlanDependency {
110                        from: occurrence - 1,
111                        to: occurrence,
112                        cycle_lag: 0,
113                    };
114                    if !plan.dependencies.contains(&edge) {
115                        plan.dependencies.push(edge);
116                    }
117                }
118            }
119            plan.workers = workers;
120            missions.insert(mission, plan);
121        }
122        let plan = CuPlan {
123            concurrent_resources: Vec::new(),
124            missions,
125        };
126        plan.validate(config)?;
127        Ok(plan)
128    }
129}