cu29_runtime/planner/
pipeline.rs1use 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
25pub 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}