logicaffeine_compile/spatial/
residency.rs

1//! Residency — which tensors the SRAM holds, and what leaves the chip.
2//!
3//! ★ THE SRAM IS THE PMUS' MEMORY, AND IT IS ALLOCATED. A spatial chip's SRAM tier is the sum of its sites' local
4//! memories: every buffer a program lays out takes a share of it, and a tensor held there between sections — an
5//! intermediate one section writes and a later one reads, a state one step leaves for the next — takes more. It is
6//! only a tier a tensor can LIVE in where any site reaches any other's memory (a crossbar, the RDU); a target whose
7//! sites share data through a global memory (Tensix, a GPU) hands intermediates through that memory.
8//!
9//! Two decisions, at the two scopes that can make them:
10//!   • a PROGRAM places its intermediates ([`place_activations`]): each lives from the phase that writes it to the
11//!     last phase that reads it, and the live set of every phase must fit the SRAM its buffers leave — past it, the
12//!     intermediate whose last use is furthest away spills first (linear-scan allocation over live ranges);
13//!   • a DEPLOYMENT places the carried state ([`plan_states`]): a step's state outlives the program, every layer
14//!     holds its own, and only the whole model knows what else is resident. The state that moves the most bytes per
15//!     byte held — a recurrent memory read and written every step before a cache read once — is kept first, within
16//!     each chip's SRAM less the largest program's own footprint.
17//! What leaves the SRAM goes to the next tier down ([`plan_off_chip`]): an RDU's HBM while it holds it, densest first,
18//! then the DDR behind it (`Dram`); a chip with no HBM keeps it all in its DRAM.
19
20use std::collections::{BTreeMap, BTreeSet};
21
22use super::ir::*;
23use super::target::{precision_bytes, Interconnect, SpatialTarget};
24
25/// Bytes `t` takes in memory: its tiles, whole.
26pub fn tensor_bytes(p: &SpatialProgram, t: &TensorDecl) -> u64 {
27    (t.tiles.rows_t * t.tiles.cols_t * p.tile.elems() * precision_bytes(t.dtype)) as u64
28}
29
30/// The SRAM each chip's buffers take: over every site of the chip, the highest byte its local memory map reaches.
31pub fn buffer_footprint(p: &SpatialProgram) -> BTreeMap<u16, u64> {
32    let mut high: BTreeMap<Site, u64> = BTreeMap::new();
33    if let Some(mm) = &p.memory {
34        for at in mm.buffer_at.iter().filter(|a| a.space != MemSpace::Acc) {
35            let top = high.entry(at.site).or_default();
36            *top = (*top).max(at.offset_bytes as u64 + at.bytes as u64);
37        }
38    }
39    let mut out = BTreeMap::new();
40    for (site, top) in high {
41        *out.entry(site.chip).or_default() += top;
42    }
43    out
44}
45
46/// The chips whose sites load or store each tensor, by tensor id.
47fn tensor_chips(p: &SpatialProgram) -> Vec<BTreeSet<u16>> {
48    let mut out = vec![BTreeSet::new(); p.tensors.len()];
49    let pl = p.placement.as_ref();
50    for s in &p.stages {
51        let (StageKind::Load { tensor, .. } | StageKind::Store { tensor, .. }) = &s.kind else {
52            continue;
53        };
54        let site = pl
55            .and_then(|pl| pl.site_of.get(s.id.idx()).copied())
56            .or(s.site);
57        out[tensor.idx()].insert(site.map(|x| x.chip).unwrap_or(0));
58    }
59    out
60}
61
62/// Each intermediate's live range over the program's phases: from the first phase that stores it to the last that
63/// loads it.
64fn live_ranges(p: &SpatialProgram) -> BTreeMap<TensorId, (usize, usize)> {
65    let phases = p.section_phases();
66    let phase_of = |s: StageId| {
67        p.section_of_stage(s)
68            .and_then(|k| phases.get(k).copied())
69            .unwrap_or(0)
70    };
71    let mut out: BTreeMap<TensorId, (usize, usize)> = BTreeMap::new();
72    for s in &p.stages {
73        let (tensor, stores) = match &s.kind {
74            StageKind::Load { tensor, .. } => (*tensor, false),
75            StageKind::Store { tensor, .. } => (*tensor, true),
76            _ => continue,
77        };
78        if p.tensor(tensor).role != TensorRole::Activation {
79            continue;
80        }
81        let ph = phase_of(s.id);
82        let r = out.entry(tensor).or_insert((ph, ph));
83        if stores {
84            r.0 = r.0.min(ph);
85        }
86        r.1 = r.1.max(ph);
87    }
88    out
89}
90
91/// Whether `target` can hold a tensor in its SRAM between sections: an SRAM tier any site reaches.
92pub fn holds_tensors_on_chip(target: &SpatialTarget) -> bool {
93    target.dram.sram_kib > 0 && target.provides_interconnect(Interconnect::Crossbar)
94}
95
96/// The chips each intermediate is held on: those whose sites load or store it, chip 0 for one no site touches.
97fn held_on(chips: &[BTreeSet<u16>], t: TensorId) -> Vec<u16> {
98    if chips[t.idx()].is_empty() {
99        vec![0]
100    } else {
101        chips[t.idx()].iter().copied().collect()
102    }
103}
104
105/// ★ PLACE THE INTERMEDIATES. Where `target` holds tensors on chip, each chip's intermediates are allocated over their
106/// live ranges in the SRAM its buffers leave; where it does not, or past what fits, they stay in the tier they were
107/// declared in. A live set past the SRAM spills, at the first phase it overflows, the intermediate live there whose
108/// last use is furthest away — then the largest, then the lowest id — until every phase fits.
109///
110/// One sweep over the phases per chip (linear scan): a range joins the live set at its first phase and leaves after its
111/// last, the live bytes are kept as a running sum, and an overflow spills from the live set ordered by last use. A
112/// spill only lowers the phases after it, and every phase before already fit, so the sweep never looks back — the
113/// same intermediates spill as re-searching from the first phase would, in `O(n log n)`.
114pub fn place_activations(p: &mut SpatialProgram, target: &SpatialTarget) {
115    let ranges = live_ranges(p);
116    let chips = tensor_chips(p);
117    let mut spilled: BTreeSet<TensorId> = BTreeSet::new();
118    if holds_tensors_on_chip(target) {
119        let sram = target.dram.sram_kib as u64 * 1024;
120        let buffers = buffer_footprint(p);
121        let mut by_chip: BTreeMap<u16, Vec<(TensorId, (usize, usize), u64)>> = BTreeMap::new();
122        for (t, r) in &ranges {
123            let b = tensor_bytes(p, p.tensor(*t));
124            for c in held_on(&chips, *t) {
125                by_chip.entry(c).or_default().push((*t, *r, b));
126            }
127        }
128        for (chip, mut live) in by_chip {
129            let room = sram.saturating_sub(buffers.get(&chip).copied().unwrap_or(0));
130            live.sort_by_key(|(t, r, _)| (r.0, t.idx()));
131            // The live set, ordered so its last element is the one to spill: furthest last use, then largest, then
132            // lowest id. Ranges leave it in order of their last phase.
133            let mut active: BTreeSet<(usize, u64, std::cmp::Reverse<usize>)> = BTreeSet::new();
134            let mut ending: BTreeSet<(usize, usize, u64)> = BTreeSet::new();
135            let mut sum = 0u64;
136            let mut next = 0usize;
137            let mut phases: Vec<usize> = live.iter().map(|(_, r, _)| r.0).collect();
138            phases.dedup();
139            for ph in phases {
140                while let Some(&(last, id, b)) = ending.first() {
141                    if last >= ph {
142                        break;
143                    }
144                    ending.pop_first();
145                    if active.remove(&(last, b, std::cmp::Reverse(id))) {
146                        sum -= b;
147                    }
148                }
149                while next < live.len() && live[next].1 .0 == ph {
150                    let (t, r, b) = live[next];
151                    next += 1;
152                    if spilled.contains(&t) {
153                        continue;
154                    }
155                    active.insert((r.1, b, std::cmp::Reverse(t.idx())));
156                    ending.insert((r.1, t.idx(), b));
157                    sum += b;
158                }
159                while sum > room {
160                    let (last, b, std::cmp::Reverse(id)) = active
161                        .pop_last()
162                        .expect("an overflowing phase holds an intermediate");
163                    ending.remove(&(last, id, b));
164                    sum -= b;
165                    spilled.insert(TensorId::new(id));
166                }
167            }
168        }
169    }
170    let on_chip = holds_tensors_on_chip(target);
171    for t in &mut p.tensors {
172        if t.role == TensorRole::Activation {
173            t.space = if on_chip && !spilled.contains(&t.id) {
174                TensorSpace::Sram
175            } else {
176                TensorSpace::Dram
177            };
178        }
179    }
180}
181
182/// The SRAM each chip of `p` holds while it runs: its buffers and the peak of its on-chip intermediates — the running
183/// sum of the live bytes over the phases, a range added at its first phase and taken away after its last.
184pub fn sram_footprint(p: &SpatialProgram) -> BTreeMap<u16, u64> {
185    let ranges = live_ranges(p);
186    let chips = tensor_chips(p);
187    let mut out = buffer_footprint(p);
188    let mut deltas: BTreeMap<u16, BTreeMap<usize, i128>> = BTreeMap::new();
189    for (t, r) in &ranges {
190        let d = p.tensor(*t);
191        if d.space != TensorSpace::Sram {
192            continue;
193        }
194        let b = tensor_bytes(p, d) as i128;
195        for c in held_on(&chips, *t) {
196            let e = deltas.entry(c).or_default();
197            *e.entry(r.0).or_default() += b;
198            *e.entry(r.1 + 1).or_default() -= b;
199        }
200    }
201    for (c, d) in deltas {
202        let (mut sum, mut peak) = (0i128, 0i128);
203        for (_, delta) in d {
204            sum += delta;
205            peak = peak.max(sum);
206        }
207        *out.entry(c).or_default() += peak as u64;
208    }
209    out
210}
211
212/// A carried state of a deployment: `tensor` of the program `program`, held on chip `chip` once for each of the
213/// `instances` layers that run the program, each `bytes` in memory and moving `moved` bytes a step.
214#[derive(Clone, Debug, PartialEq, Eq)]
215pub struct StateClass {
216    pub program: String,
217    pub tensor: String,
218    pub chip: u16,
219    pub instances: usize,
220    pub bytes: u64,
221    pub moved: u64,
222}
223
224impl StateClass {
225    /// Every instance's bytes.
226    pub fn total(&self) -> u64 {
227        self.bytes * self.instances as u64
228    }
229}
230
231/// Where a deployment keeps its carried state: each chip's SRAM, the reserve its largest program takes while it runs,
232/// and every state class with the tier it was given.
233#[derive(Clone, Debug, PartialEq, Eq, Default)]
234pub struct ResidencyPlan {
235    pub sram_bytes: u64,
236    /// Each chip's HBM; `0` where the chip has none.
237    pub hbm_bytes: u64,
238    pub reserve: BTreeMap<u16, u64>,
239    /// Every carried state, with its tier.
240    pub placed: Vec<(StateClass, TensorSpace)>,
241    /// Every other tensor off the SRAM — weights, runtime inputs, results, spilled intermediates — with its tier.
242    pub off_chip: Vec<(StateClass, TensorSpace)>,
243}
244
245impl ResidencyPlan {
246    /// The state `chip` keeps in its SRAM.
247    pub fn on_chip(&self, chip: u16) -> u64 {
248        self.placed
249            .iter()
250            .filter(|(c, s)| c.chip == chip && *s == TensorSpace::Sram)
251            .map(|(c, _)| c.total())
252            .sum()
253    }
254
255    /// The tier the plan gives `tensor` of `program`.
256    pub fn tier_of(&self, program: &str, tensor: &str) -> Option<TensorSpace> {
257        self.placed
258            .iter()
259            .find(|(c, _)| c.program == program && c.tensor == tensor)
260            .map(|(_, s)| *s)
261    }
262}
263
264/// ★ KEEP WHAT MOVES MOST PER BYTE HELD. In order of bytes moved a step per byte held — then the smaller class, then
265/// by name — each class goes on chip whole if its chip has room left beyond `reserve`, and stays in HBM/DDR
266/// (`Dram`) otherwise; a class that does not fit leaves the room to the classes after it. Every byte a class keeps
267/// on chip is a byte a step does not move off it, `moved / bytes` times over, so filling the SRAM densest-first is
268/// the most traffic whole classes can keep.
269pub fn plan_states(
270    classes: Vec<StateClass>,
271    sram_bytes: u64,
272    reserve: &BTreeMap<u16, u64>,
273) -> ResidencyPlan {
274    let mut order: Vec<usize> = (0..classes.len()).collect();
275    order.sort_by(|&a, &b| {
276        let (x, y) = (&classes[a], &classes[b]);
277        (y.moved as u128 * x.bytes as u128)
278            .cmp(&(x.moved as u128 * y.bytes as u128))
279            .then(x.total().cmp(&y.total()))
280            .then((&x.program, &x.tensor).cmp(&(&y.program, &y.tensor)))
281    });
282    let mut room: BTreeMap<u16, u64> = BTreeMap::new();
283    let mut tier = vec![TensorSpace::Dram; classes.len()];
284    for i in order {
285        let c = &classes[i];
286        let left = room.entry(c.chip).or_insert_with(|| {
287            sram_bytes.saturating_sub(reserve.get(&c.chip).copied().unwrap_or(0))
288        });
289        if c.total() <= *left {
290            *left -= c.total();
291            tier[i] = TensorSpace::Sram;
292        }
293    }
294    ResidencyPlan {
295        sram_bytes,
296        hbm_bytes: 0,
297        reserve: reserve.clone(),
298        placed: classes.into_iter().zip(tier).collect(),
299        off_chip: Vec::new(),
300    }
301}
302
303/// ★ THE NEXT TIER DOWN. Each class the SRAM does not hold goes, in the order [`plan_states`] ranks by, to its chip's
304/// HBM while `hbm_bytes` holds it whole, and to the DDR (`Dram`) past it; with no HBM, all of it to the DRAM.
305pub fn plan_off_chip(classes: &[StateClass], hbm_bytes: u64) -> Vec<TensorSpace> {
306    let mut order: Vec<usize> = (0..classes.len()).collect();
307    order.sort_by(|&a, &b| {
308        let (x, y) = (&classes[a], &classes[b]);
309        (y.moved as u128 * x.bytes as u128)
310            .cmp(&(x.moved as u128 * y.bytes as u128))
311            .then(x.total().cmp(&y.total()))
312            .then((&x.program, &x.tensor).cmp(&(&y.program, &y.tensor)))
313    });
314    let mut room: BTreeMap<u16, u64> = BTreeMap::new();
315    let mut tier = vec![TensorSpace::Dram; classes.len()];
316    for i in order {
317        let c = &classes[i];
318        let left = room.entry(c.chip).or_insert(hbm_bytes);
319        if hbm_bytes > 0 && c.total() <= *left {
320            *left -= c.total();
321            tier[i] = TensorSpace::Hbm;
322        }
323    }
324    tier
325}
326
327/// Bytes each tensor of `p` moves a step, every loop at its bound.
328pub fn moved_bytes(p: &SpatialProgram) -> Vec<u64> {
329    let mut out = vec![0u64; p.tensors.len()];
330    let tile = p.tile.elems() as u64;
331    if let Some(sch) = &p.schedule {
332        for prog in &sch.programs {
333            walk_tagged(&prog.ops, 1, &mut |t, m| {
334                if let ChoreoOp::ReadTiles {
335                    tensor,
336                    block_h,
337                    block_w,
338                    ..
339                }
340                | ChoreoOp::WriteTiles {
341                    tensor,
342                    block_h,
343                    block_w,
344                    ..
345                } = &t.op
346                {
347                    out[tensor.idx()] += m
348                        * (*block_h * *block_w) as u64
349                        * tile
350                        * precision_bytes(p.tensor(*tensor).dtype) as u64;
351                }
352            });
353        }
354    }
355    out
356}
357
358/// One program of a deployment: its name, the layers that run it, the program, and — when the lowering knows it —
359/// which of its `State` tensors a step carries to the next. A state the step assembles whole each time is the step's
360/// own table: it lives while the program runs, like an intermediate, and is not the deployment's to keep.
361pub type DeployedProgram<'a> = (
362    &'a str,
363    usize,
364    &'a mut SpatialProgram,
365    Option<&'a BTreeSet<String>>,
366);
367
368/// A deployment's carried state planned and placed: every CARRIED `State` tensor of every program becomes a
369/// [`StateClass`], once per layer that runs the program; the reserve is each chip's largest footprint while a program
370/// runs — [`sram_footprint`] and the program's own tables, kept on chip when the program has room for them — and
371/// every carried state takes the tier [`plan_states`] gives it.
372pub fn plan_deployment(programs: &mut [DeployedProgram], target: &SpatialTarget) -> ResidencyPlan {
373    let sram = if holds_tensors_on_chip(target) {
374        target.dram.sram_kib as u64 * 1024
375    } else {
376        0
377    };
378    let mut reserve: BTreeMap<u16, u64> = BTreeMap::new();
379    let mut classes = Vec::new();
380    for (name, instances, p, carried) in programs.iter_mut() {
381        let chips = tensor_chips(p);
382        let chip_of = |t: &TensorDecl| chips[t.id.idx()].iter().next().copied().unwrap_or(0);
383        let own = |t: &TensorDecl| {
384            t.role == TensorRole::State && carried.is_some_and(|c| !c.contains(&t.name))
385        };
386        let mut footprint = sram_footprint(p);
387        let mut tables: BTreeMap<u16, u64> = BTreeMap::new();
388        for t in p.tensors.iter().filter(|t| own(t)) {
389            *tables.entry(chip_of(t)).or_default() += tensor_bytes(p, t);
390        }
391        let fits = tables
392            .iter()
393            .all(|(c, b)| footprint.get(c).copied().unwrap_or(0) + b <= sram);
394        let table_ids: Vec<TensorId> = p.tensors.iter().filter(|t| own(t)).map(|t| t.id).collect();
395        for id in table_ids {
396            p.tensors[id.idx()].space = if fits {
397                TensorSpace::Sram
398            } else {
399                TensorSpace::Dram
400            };
401        }
402        if fits {
403            for (c, b) in tables {
404                *footprint.entry(c).or_default() += b;
405            }
406        }
407        for (c, b) in footprint {
408            let r = reserve.entry(c).or_default();
409            *r = (*r).max(b);
410        }
411        let moved = moved_bytes(p);
412        for t in p
413            .tensors
414            .iter()
415            .filter(|t| t.role == TensorRole::State && !own(t))
416        {
417            classes.push(StateClass {
418                program: name.to_string(),
419                tensor: t.name.clone(),
420                chip: chip_of(t),
421                instances: *instances,
422                bytes: tensor_bytes(p, t),
423                moved: moved[t.id.idx()],
424            });
425        }
426    }
427    let mut plan = plan_states(classes, sram, &reserve);
428    // Everything the SRAM does not hold — a carried state it could not keep, and every tensor that was never a
429    // candidate — takes the next tier down, one class per tensor of each program, once per layer that runs it.
430    let hbm = target.dram.hbm_gib as u64 * (1 << 30);
431    let placed_at: std::collections::HashMap<(&str, &str), usize> = plan
432        .placed
433        .iter()
434        .enumerate()
435        .map(|(i, (c, _))| ((c.program.as_str(), c.tensor.as_str()), i))
436        .collect();
437    let mut rest: Vec<(StateClass, Option<usize>)> = Vec::new();
438    for (name, instances, p, _) in programs.iter() {
439        let chips = tensor_chips(p);
440        let moved = moved_bytes(p);
441        for t in &p.tensors {
442            let carried = placed_at.get(&(*name, t.name.as_str())).copied();
443            let on_chip = match carried {
444                Some(i) => plan.placed[i].1 == TensorSpace::Sram,
445                None => t.space == TensorSpace::Sram,
446            };
447            if !on_chip {
448                let chip = chips[t.id.idx()].iter().next().copied().unwrap_or(0);
449                rest.push((
450                    StateClass {
451                        program: name.to_string(),
452                        tensor: t.name.clone(),
453                        chip,
454                        instances: *instances,
455                        bytes: tensor_bytes(p, t),
456                        moved: moved[t.id.idx()],
457                    },
458                    carried,
459                ));
460            }
461        }
462    }
463    drop(placed_at);
464    let classes: Vec<StateClass> = rest.iter().map(|(c, _)| c.clone()).collect();
465    let tiers = plan_off_chip(&classes, hbm);
466    plan.hbm_bytes = hbm;
467    for ((c, carried), tier) in rest.into_iter().zip(tiers) {
468        match carried {
469            Some(i) => plan.placed[i].1 = tier,
470            None => plan.off_chip.push((c, tier)),
471        }
472    }
473    let tier_at: std::collections::HashMap<(&str, &str), TensorSpace> = plan
474        .placed
475        .iter()
476        .chain(plan.off_chip.iter())
477        .map(|(c, s)| ((c.program.as_str(), c.tensor.as_str()), *s))
478        .collect();
479    for (name, _, p, _) in programs.iter_mut() {
480        for t in p.tensors.iter_mut() {
481            if let Some(tier) = tier_at.get(&(*name, t.name.as_str())) {
482                t.space = *tier;
483            }
484        }
485    }
486    drop(tier_at);
487    plan
488}
logicaffeine_compile/spatial/residency.rs

logicaffeine_compile/spatial/
residency.rs
