Skip to main content

logicaffeine_compile/spatial/passes/
persist.rs

1//! `persist` — one resident worker per site loops over ITS work items (the "persistent kernel").
2//!
3//! Where a site owns more than one work item, every engine program of the site is wrapped in an outer
4//! item loop; the tile addresses are already affine in the item coordinates (`LV_ITEM_Y` / `LV_ITEM_X`),
5//! which the choreography derives from the item index. Multicast groups whose members own different item
6//! counts would pair blocks with nothing on the far side — refused typed.
7
8use 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}