logicaffeine_compile/spatial/passes/
persist.rs1use super::{SpatialReject, LV_ITEM};
9use crate::spatial::ir::*;
10use crate::spatial::target::SpatialTarget;
11
12pub fn persist(
13 p: &SpatialProgram,
14 _target: &SpatialTarget,
15) -> Result<SpatialProgram, SpatialReject> {
16 let (Some(pl), Some(sched)) = (&p.placement, &p.schedule) else {
17 return Err(SpatialReject::WrongPhase {
18 pass: "persist",
19 expected: "Scheduled",
20 found: p.phase(),
21 });
22 };
23 for g in &pl.groups {
24 let counts: Vec<(Site, usize)> = g
25 .rect
26 .sites()
27 .iter()
28 .map(|s| (*s, pl.items_at(*s).len()))
29 .collect();
30 if counts
31 .iter()
32 .map(|(_, c)| *c)
33 .collect::<std::collections::BTreeSet<_>>()
34 .len()
35 > 1
36 {
37 return Err(SpatialReject::GroupCountMismatch {
38 group: g.id,
39 counts,
40 });
41 }
42 }
43 let mut q = p.clone();
44 let programs = q.schedule.as_mut().unwrap();
45 for prog in programs.programs.iter_mut() {
46 let items = pl.items_at(prog.site).len();
47 let already = matches!(prog.ops.first(), Some(TaggedOp { op: ChoreoOp::Loop { var, .. }, .. }) if *var == LV_ITEM);
48 if items > 1 && !already {
49 let stage = prog.ops[0].stage;
50 let body: Vec<TaggedOp> = prog.ops.drain(..).collect();
51 prog.ops = vec![TaggedOp {
52 stage,
53 op: ChoreoOp::Loop {
54 var: LV_ITEM,
55 count: LoopCount::Const(items),
56 body,
57 },
58 }];
59 }
60 }
61 let _ = sched;
62 q.validate()?;
63 Ok(q)
64}