logicaffeine_compile/spatial/passes/
section_shape.rs

1//! The canonical SHAPE of a section — its choreography with every name forgotten.
2//!
3//! ★ WHY A SHAPE IS NOT A SECTION. Sixty-four decoder layers are the same handful of dataflows run over
4//! and over against different weights. The compiler proves that today by unrolling: four real-width
5//! layers compose into 153 sections, and sixty-four would be hundreds of gigabytes of IR describing a
6//! program no SN50 could hold anyway — 432 MB of SRAM against 27 GB of weights means the chip runs ONE
7//! layer's dataflow and streams the layers through it.
8//!
9//! Turning that unrolled form into a loop needs the regularity to be a CHECKED FACT, not an assumption
10//! about how the model was authored. A shape is that fact: a section reduced to its op sequence with
11//! every buffer, semaphore, group, site and engine renumbered by first appearance, and every TENSOR
12//! forgotten. Two sections share a shape exactly when their event graphs are isomorphic — which is the
13//! precise sense in which layer 7 and layer 8 "are the same".
14//!
15//! What a shape deliberately keeps: loop counts, buffer depths and page sizes, semaphore roles, access
16//! patterns, and the compute bodies. Two sections that differ in any of those run different dataflows,
17//! however alike they look in the source.
18
19use std::collections::BTreeMap;
20
21use super::super::ir::*;
22
23/// FNV-1a over a canonical byte stream. The digest is a function of the program, never of an address.
24struct Shape(u64);
25
26impl Shape {
27    fn new() -> Shape {
28        Shape(0xcbf29ce484222325)
29    }
30    fn bytes(&mut self, b: &[u8]) {
31        for x in b {
32            self.0 ^= u64::from(*x);
33            self.0 = self.0.wrapping_mul(0x100000001b3);
34        }
35    }
36    fn tag(&mut self, s: &str) {
37        self.bytes(s.as_bytes());
38        self.bytes(&[0xff]);
39    }
40    fn num(&mut self, n: usize) {
41        self.bytes(&(n as u64).to_le_bytes());
42    }
43    fn json<T: serde::Serialize>(&mut self, v: &T) {
44        self.bytes(serde_json::to_string(v).unwrap_or_default().as_bytes());
45        self.bytes(&[0xfe]);
46    }
47}
48
49/// Ids renumbered by first appearance, so two isomorphic sections canonicalize identically.
50#[derive(Default)]
51struct Canon {
52    site: BTreeMap<(u16, u16, u16), usize>,
53    engine: BTreeMap<String, usize>,
54    buf: BTreeMap<usize, usize>,
55    sem: BTreeMap<usize, usize>,
56    grp: BTreeMap<usize, usize>,
57}
58
59impl Canon {
60    fn of<K: Ord>(m: &mut BTreeMap<K, usize>, k: K) -> usize {
61        let n = m.len();
62        *m.entry(k).or_insert(n)
63    }
64    fn site(&mut self, s: Site) -> usize {
65        Canon::of(&mut self.site, (s.chip, s.y, s.x))
66    }
67    fn engine(&mut self, e: Engine) -> usize {
68        Canon::of(&mut self.engine, format!("{e:?}"))
69    }
70    fn buf(&mut self, b: BufferId) -> usize {
71        Canon::of(&mut self.buf, b.idx())
72    }
73    fn sem(&mut self, s: SemId) -> usize {
74        Canon::of(&mut self.sem, s.idx())
75    }
76    fn grp(&mut self, g: GroupId) -> usize {
77        Canon::of(&mut self.grp, g.idx())
78    }
79}
80
81/// Fold one op into the shape, renumbering every id it names. A TENSOR is never folded in: that is the
82/// one thing a shape forgets, and the reason two layers of the same kind share one.
83fn fold_op(p: &SpatialProgram, c: &mut Canon, h: &mut Shape, op: &ChoreoOp) {
84    // A buffer's own geometry is part of the dataflow, so it travels with the renumbered id.
85    let buf = |c: &mut Canon, h: &mut Shape, b: BufferId| {
86        h.num(c.buf(b));
87        if let Some(bf) = p.buffers.get(b.idx()) {
88            h.json(&bf.kind);
89            h.json(&bf.dtype);
90            h.json(&bf.space);
91        }
92    };
93    let sem = |c: &mut Canon, h: &mut Shape, s: SemId| {
94        h.num(c.sem(s));
95        if let Some(sm) = p.semaphores.get(s.idx()) {
96            h.json(&sm.role);
97            h.num(sm.init as usize);
98        }
99    };
100    match op {
101        ChoreoOp::Reserve { buf: b, pages } => {
102            h.tag("reserve");
103            buf(c, h, *b);
104            h.num(*pages);
105        }
106        ChoreoOp::Push { buf: b, pages } => {
107            h.tag("push");
108            buf(c, h, *b);
109            h.num(*pages);
110        }
111        ChoreoOp::Wait { buf: b, pages } => {
112            h.tag("wait");
113            buf(c, h, *b);
114            h.num(*pages);
115        }
116        ChoreoOp::Pop { buf: b, pages } => {
117            h.tag("pop");
118            buf(c, h, *b);
119            h.num(*pages);
120        }
121        ChoreoOp::ReadTiles {
122            at,
123            block_h,
124            block_w,
125            stride_w,
126            stride_h,
127            into,
128            ..
129        } => {
130            h.tag("read");
131            h.json(at);
132            h.num(*block_h);
133            h.num(*block_w);
134            h.num(*stride_w);
135            h.num(*stride_h);
136            buf(c, h, *into);
137        }
138        ChoreoOp::WriteTiles {
139            from,
140            at,
141            block_h,
142            block_w,
143            stride_w,
144            stride_h,
145            ..
146        } => {
147            h.tag("write");
148            buf(c, h, *from);
149            h.json(at);
150            h.num(*block_h);
151            h.num(*block_w);
152            h.num(*stride_w);
153            h.num(*stride_h);
154        }
155        ChoreoOp::ReadBarrier => h.tag("rbar"),
156        ChoreoOp::WriteBarrier => h.tag("wbar"),
157        ChoreoOp::McastSend {
158            buf: b,
159            pages,
160            group,
161            hold_path,
162        } => {
163            h.tag("mcast");
164            buf(c, h, *b);
165            h.num(*pages);
166            h.num(c.grp(*group));
167            h.num(usize::from(*hold_path));
168        }
169        ChoreoOp::McastSetValid { sem: s, group } => {
170            h.tag("mvalid");
171            sem(c, h, *s);
172            h.num(c.grp(*group));
173        }
174        ChoreoOp::Flush { cond } => {
175            h.tag("flush");
176            h.json(cond);
177        }
178        ChoreoOp::SemSet { sem: s, value } => {
179            h.tag("semset");
180            sem(c, h, *s);
181            h.json(value);
182        }
183        ChoreoOp::SemInc { sem: s, at, by } => {
184            h.tag("seminc");
185            sem(c, h, *s);
186            h.num(c.site(*at));
187            h.num(*by);
188        }
189        ChoreoOp::SemWait { sem: s, until } => {
190            h.tag("semwait");
191            sem(c, h, *s);
192            h.json(until);
193        }
194        ChoreoOp::AcquireAcc { tiles } => {
195            h.tag("acqacc");
196            h.num(*tiles);
197        }
198        ChoreoOp::ReleaseAcc => h.tag("relacc"),
199        ChoreoOp::Compute {
200            op,
201            ins,
202            acc_reload,
203        } => {
204            h.tag("compute");
205            h.json(op);
206            h.num(ins.len());
207            for b in ins {
208                buf(c, h, *b);
209            }
210            match acc_reload {
211                Some(b) => buf(c, h, *b),
212                None => h.tag("noreload"),
213            }
214        }
215        ChoreoOp::PackAcc {
216            to,
217            tiles,
218            convert,
219            accumulate,
220        } => {
221            h.tag("packacc");
222            buf(c, h, *to);
223            h.num(*tiles);
224            h.json(convert);
225            h.num(usize::from(*accumulate));
226        }
227        ChoreoOp::Barrier { group } => {
228            h.tag("barrier");
229            h.num(c.grp(*group));
230        }
231        ChoreoOp::Loop { var, count, body } => {
232            h.tag("loop");
233            h.json(var);
234            // A trip keeps its arithmetic and forgets the tensor it reads, like every other tensor here.
235            match count {
236                LoopCount::Const(n) => h.num(*n),
237                LoopCount::Trip { trip, max } => {
238                    let t = p.trip(*trip);
239                    h.tag("trip");
240                    h.num(*max);
241                    h.num(t.elem);
242                    h.bytes(&t.admit.0.to_le_bytes());
243                    h.bytes(&t.admit.1.to_le_bytes());
244                    h.bytes(&t.bias.to_le_bytes());
245                    h.num(t.per);
246                }
247            }
248            h.num(body.len());
249            for t in body {
250                fold_op(p, c, h, &t.op);
251            }
252            h.tag("endloop");
253        }
254    }
255}
256
257/// The canonical shape of every section, in section order.
258///
259/// A program with no sections has no shapes — the function is total and reports nothing rather than
260/// treating a single-section program as a degenerate repetition of itself.
261pub fn section_shapes(p: &SpatialProgram) -> Vec<u64> {
262    let Some(sched) = p.schedule.as_ref() else {
263        return Vec::new();
264    };
265    // ★ ONE WALK OVER THE SCHEDULE, NOT ONE PER SECTION.
266    //
267    // Filtering every engine program's whole op list once per section is quadratic in the pair, and the
268    // pair is not small: a real-width eight-layer forward is 301 sections over 178,436,487 top-level ops
269    // across 1,088 engine programs, so the naive form is ~5e10 comparisons and never returns. Stage ids
270    // already say which section an op belongs to, so bucketing the ops ONCE costs one pass and makes the
271    // whole analysis linear. (The same shape of defect as the `validate_schedule` walk that cost 16
272    // minutes and 13 GB before it was folded into a single pass.)
273    let mut sec_of_stage = vec![usize::MAX; p.stages.len()];
274    for (si, sp) in p.sections.iter().enumerate() {
275        for slot in sec_of_stage
276            .iter_mut()
277            .take(sp.stages.1.min(p.stages.len()))
278            .skip(sp.stages.0)
279        {
280            *slot = si;
281        }
282    }
283    // `buckets[section][program]` = that program's ops belonging to that section, in program order.
284    let np = sched.programs.len();
285    let mut buckets: Vec<Vec<Vec<&TaggedOp>>> = vec![Vec::new(); p.sections.len()];
286    for b in buckets.iter_mut() {
287        b.resize_with(np, Vec::new);
288    }
289    for (pi, prog) in sched.programs.iter().enumerate() {
290        for t in &prog.ops {
291            let si = sec_of_stage
292                .get(t.stage.idx())
293                .copied()
294                .unwrap_or(usize::MAX);
295            if let Some(rows) = buckets.get_mut(si) {
296                rows[pi].push(t);
297            }
298        }
299    }
300    p.sections
301        .iter()
302        .enumerate()
303        .map(|(si, sp)| {
304            let mut c = Canon::default();
305            let mut h = Shape::new();
306            h.num(sp.k_tiles);
307            h.json(&sp.plan);
308            h.num(sp.items.1.saturating_sub(sp.items.0));
309            for (pi, prog) in sched.programs.iter().enumerate() {
310                let ops = &buckets[si][pi];
311                if ops.is_empty() {
312                    continue;
313                }
314                h.tag("prog");
315                h.num(c.site(prog.site));
316                h.num(c.engine(prog.engine));
317                h.num(ops.len());
318                for t in ops {
319                    fold_op(p, &mut c, &mut h, &t.op);
320                }
321            }
322            h.0
323        })
324        .collect()
325}
326
327/// Sections grouped by shape: `(shape, the sections that have it)`, ordered by each group's first
328/// section so the answer is stable. The groups partition the sections.
329pub fn section_repeat_groups(p: &SpatialProgram) -> Vec<(u64, Vec<usize>)> {
330    let mut by: BTreeMap<u64, Vec<usize>> = BTreeMap::new();
331    for (i, s) in section_shapes(p).into_iter().enumerate() {
332        by.entry(s).or_default().push(i);
333    }
334    let mut v: Vec<(u64, Vec<usize>)> = by.into_iter().collect();
335    v.sort_by_key(|(_, m)| m.first().copied().unwrap_or(usize::MAX));
336    v
337}
logicaffeine_compile/spatial/passes/section_shape.rs

logicaffeine_compile/spatial/passes/
section_shape.rs
