logicaffeine_compile/spatial/lower/tensix/
roles.rs

1//! Engine programs → kernel classes. Two programs whose site-independent signatures agree share ONE kernel
2//! source; a launch (`CreateKernel`) is a distinct (source, engine, NOC) over the maximal rectangles of the
3//! sites it covers. The 2-D GEMM's four reader variants fall out of the signature, unnamed.
4
5use std::collections::{BTreeMap, BTreeSet};
6
7use super::SpatialLowerError;
8use crate::spatial::ir::*;
9
10#[derive(Clone, Debug, PartialEq, Eq)]
11pub struct Launch {
12    pub noc: Option<NocId>,
13    pub rects: Vec<Rect>,
14    pub sites: Vec<Site>,
15}
16
17/// One kernel source: a signature class of engine programs, launched once per NOC it runs on.
18#[derive(Clone, Debug, PartialEq, Eq)]
19pub struct KernelClass {
20    pub file_name: String,
21    pub engine: Engine,
22    pub signature: String,
23    /// The representative program (the first site in order) the kernel text is printed from.
24    pub representative: usize,
25    /// Every site of the class, in site order.
26    pub sites: Vec<Site>,
27    /// One `CreateKernel` per NOC.
28    pub launches: Vec<Launch>,
29}
30
31/// Greedy maximal-rectangle decomposition of a site set, deterministic (row-major seeds).
32pub fn rect_decompose(sites: &BTreeSet<Site>) -> Vec<Rect> {
33    let mut left: BTreeSet<(u16, u16)> = sites.iter().map(|s| (s.y, s.x)).collect();
34    let mut out = Vec::new();
35    while let Some(&(y0, x0)) = left.iter().next() {
36        let mut x1 = x0;
37        while left.contains(&(y0, x1 + 1)) {
38            x1 += 1;
39        }
40        let mut y1 = y0;
41        while (x0..=x1).all(|x| left.contains(&(y1 + 1, x))) {
42            y1 += 1;
43        }
44        for y in y0..=y1 {
45            for x in x0..=x1 {
46                left.remove(&(y, x));
47            }
48        }
49        out.push(Rect { x0, y0, x1, y1 });
50    }
51    out
52}
53
54fn describe(p: &SpatialProgram, prog: &EngineProgram) -> String {
55    // A readable role name from the ops: per operand `send` / `recv` / `read`, plus `writer`.
56    let mut a = None;
57    let mut b = None;
58    let mut writer = false;
59    let mut compute = false;
60    let sched_ops: Vec<ChoreoOp> = prog.ops.iter().map(|t| t.op.clone()).collect();
61    walk_ops(&sched_ops, 1, &mut |op, _| match op {
62        ChoreoOp::ReadTiles { tensor, .. } => {
63            let slot = if tensor.idx() == 0 { &mut a } else { &mut b };
64            if slot.is_none() {
65                *slot = Some("read");
66            }
67        }
68        ChoreoOp::McastSend { group, .. } => {
69            let op = p.placement.as_ref().unwrap().groups[group.idx()].operand;
70            let slot = if op == Operand::A { &mut a } else { &mut b };
71            *slot = Some("send");
72        }
73        ChoreoOp::SemInc { sem, .. } => {
74            if let Some(op) = p.semaphores[sem.idx()].operand {
75                let slot = if op == Operand::A { &mut a } else { &mut b };
76                *slot = Some("recv");
77            }
78        }
79        ChoreoOp::WriteTiles { .. } => writer = true,
80        ChoreoOp::Compute { .. } => compute = true,
81        _ => {}
82    });
83    let mut parts = Vec::new();
84    if let Some(r) = a {
85        parts.push(format!("a_{r}"));
86    }
87    if let Some(r) = b {
88        parts.push(format!("b_{r}"));
89    }
90    if writer {
91        parts.push("writer".into());
92    }
93    if compute {
94        parts.push("compute".into());
95    }
96    if parts.is_empty() {
97        parts.push("idle".into());
98    }
99    parts.join("_")
100}
101
102pub fn classes(p: &SpatialProgram) -> Result<Vec<KernelClass>, SpatialLowerError> {
103    let pl = p
104        .placement
105        .as_ref()
106        .ok_or(SpatialLowerError::Unsupported("a placed program"))?;
107    let sched = p
108        .schedule
109        .as_ref()
110        .ok_or(SpatialLowerError::Unsupported("a scheduled program"))?;
111    // (engine, signature) → members with their NOC.
112    let mut groups: BTreeMap<(u8, String), Vec<(usize, Option<u8>)>> = Default::default();
113    for (i, prog) in sched.programs.iter().enumerate() {
114        let noc = pl
115            .engine_of
116            .iter()
117            .find(|e| e.site == prog.site && e.engine == prog.engine)
118            .and_then(|e| e.noc);
119        let ekey = match prog.engine {
120            Engine::Movement(e) => e,
121            Engine::Compute => 255,
122        };
123        groups
124            .entry((ekey, p.engine_signature(prog)))
125            .or_default()
126            .push((i, noc.map(|n| n.0)));
127    }
128    let mut classes = Vec::new();
129    let mut names: BTreeMap<String, usize> = Default::default();
130    for ((_, sig), members) in groups {
131        let rep = &sched.programs[members[0].0];
132        let prefix = match rep.engine {
133            Engine::Movement(e) => format!("dm{e}"),
134            Engine::Compute => "compute".into(),
135        };
136        let base = format!("{prefix}_{}", describe(p, rep));
137        let n = names.entry(base.clone()).or_insert(0);
138        *n += 1;
139        let file_name = if *n == 1 {
140            format!("{base}.cpp")
141        } else {
142            format!("{base}_{}.cpp", *n)
143        };
144        let sites: Vec<Site> = members
145            .iter()
146            .map(|(i, _)| sched.programs[*i].site)
147            .collect();
148        let mut by_noc: BTreeMap<Option<u8>, BTreeSet<Site>> = Default::default();
149        for (i, noc) in &members {
150            by_noc
151                .entry(*noc)
152                .or_default()
153                .insert(sched.programs[*i].site);
154        }
155        let launches = by_noc
156            .into_iter()
157            .map(|(noc, set)| Launch {
158                noc: noc.map(NocId),
159                sites: set.iter().copied().collect(),
160                rects: rect_decompose(&set),
161            })
162            .collect();
163        classes.push(KernelClass {
164            file_name,
165            engine: rep.engine,
166            signature: sig,
167            representative: members[0].0,
168            sites,
169            launches,
170        });
171    }
172    Ok(classes)
173}
logicaffeine_compile/spatial/lower/tensix/roles.rs

logicaffeine_compile/spatial/lower/tensix/
roles.rs
