logicaffeine_compile/spatial/lower/tensix/
cb.rs

1//! Buffers → circular buffers. One `CircularBufferConfig` per buffer TEMPLATE (the pre-placement buffer
2//! every site instantiates), with a deterministic index class: operand inputs `c_0..`, outputs `c_16..`,
3//! spill buffers `c_24..`.
4
5use std::collections::BTreeMap;
6
7use super::{template_name, template_of, SpatialLowerError};
8use crate::codegen_sva::fabric_target::Precision;
9use crate::spatial::ir::*;
10use crate::spatial::target::SpatialTarget;
11
12#[derive(Clone, Debug, PartialEq, Eq)]
13pub struct CbSlot {
14    pub template: BufferId,
15    pub name: String,
16    pub index: u32,
17    pub dtype: Precision,
18    /// The template's page: the LARGEST any site's instance holds. Every site of the template reserves,
19    /// pushes, waits and pops in it, so each reservation divides the buffer; a boundary site whose own page
20    /// is narrower fills only the front of each slot (and multicasts only that — `kernel::sent_tiles`).
21    pub page_tiles: usize,
22    pub depth_pages: usize,
23    pub tile_bytes: usize,
24    /// Every site that instantiates the template.
25    pub sites: Vec<Site>,
26    /// The template whose memory this one's pages live in (`Buffer::alias_of`): both are indices of one
27    /// `CircularBufferConfig`, and this slot has none of its own.
28    pub alias_of: Option<BufferId>,
29}
30
31impl CbSlot {
32    pub fn total_bytes(&self) -> usize {
33        self.page_tiles * self.depth_pages * self.tile_bytes
34    }
35    pub fn c_name(&self) -> String {
36        format!("CB_{}", self.name.to_uppercase())
37    }
38}
39
40#[derive(Clone, Debug, PartialEq, Eq, Default)]
41pub struct CbPlan {
42    pub slots: Vec<CbSlot>,
43    by_template: BTreeMap<BufferId, usize>,
44    by_instance: BTreeMap<BufferId, usize>,
45}
46
47/// The tt-metal data format of a precision the Tensix engines carry.
48pub fn data_format(p: Precision) -> Result<&'static str, SpatialLowerError> {
49    Ok(match p {
50        Precision::Bf16 => "tt::DataFormat::Float16_b",
51        Precision::Fp16 => "tt::DataFormat::Float16",
52        Precision::Fp32 => "tt::DataFormat::Float32",
53        Precision::Int8 => "tt::DataFormat::Int8",
54        Precision::Int16 | Precision::Int32 | Precision::Int64 | Precision::Fp8 => {
55            return Err(SpatialLowerError::Precision { needed: p })
56        }
57    })
58}
59
60impl CbPlan {
61    pub fn of(p: &SpatialProgram, target: &SpatialTarget) -> Result<CbPlan, SpatialLowerError> {
62        let mut inputs: BTreeMap<BufferId, ()> = Default::default();
63        let mut outputs: BTreeMap<BufferId, ()> = Default::default();
64        for s in &p.stages {
65            match &s.kind {
66                StageKind::Load { .. } | StageKind::Recv { .. } => {
67                    for b in &s.outs {
68                        inputs.insert(template_of(p.buffer(*b)), ());
69                    }
70                }
71                StageKind::Store { .. } => {
72                    for b in &s.ins {
73                        outputs.insert(template_of(p.buffer(*b)), ());
74                    }
75                }
76                _ => {}
77            }
78        }
79        let mut templates: BTreeMap<BufferId, (&Buffer, Vec<Site>, usize)> = Default::default();
80        for b in &p.buffers {
81            let e = templates
82                .entry(template_of(b))
83                .or_insert((b, Vec::new(), 0));
84            if let Some(site) = b.site {
85                e.1.push(site);
86            }
87            e.2 = e.2.max(b.kind.page_tiles());
88        }
89        let (mut ni, mut no, mut nm) = (0u32, 0u32, 0u32);
90        let mut slots = Vec::new();
91        let mut by_template = BTreeMap::new();
92        for (tid, (b, sites, page_tiles)) in &templates {
93            data_format(b.dtype)?;
94            let index = if matches!(b.kind, BufferKind::Interm { .. }) {
95                nm += 1;
96                24 + nm - 1
97            } else if inputs.contains_key(tid) {
98                ni += 1;
99                ni - 1
100            } else if outputs.contains_key(tid) {
101                no += 1;
102                16 + no - 1
103            } else {
104                return Err(SpatialLowerError::Unsupported(
105                    "a buffer neither read by a mover nor written by one",
106                ));
107            };
108            let max = if index < 16 { 8 } else { 8 };
109            if (index < 16 && ni > max)
110                || (16..24).contains(&index) && no > max
111                || index >= 24 && nm > max
112            {
113                return Err(SpatialLowerError::TooBig {
114                    what: "circular buffers",
115                    need: (ni + no + nm) as usize,
116                    have: 32,
117                });
118            }
119            let depth_pages = match b.kind {
120                BufferKind::Fifo { depth_pages, .. } | BufferKind::Interm { depth_pages, .. } => {
121                    depth_pages
122                }
123                BufferKind::Accumulator { .. } => {
124                    return Err(SpatialLowerError::Unsupported(
125                        "an accumulator buffer in local memory",
126                    ))
127                }
128            };
129            by_template.insert(*tid, slots.len());
130            let alias_of = b.alias_of.map(|a| template_of(p.buffer(a)));
131            slots.push(CbSlot {
132                template: *tid,
133                name: template_name(p, b),
134                index,
135                dtype: b.dtype,
136                page_tiles: *page_tiles,
137                depth_pages,
138                tile_bytes: target.tile_bytes(b.dtype),
139                sites: sites.clone(),
140                alias_of,
141            });
142        }
143        let mut by_instance = BTreeMap::new();
144        for b in &p.buffers {
145            by_instance.insert(b.id, by_template[&template_of(b)]);
146        }
147        Ok(CbPlan {
148            slots,
149            by_template,
150            by_instance,
151        })
152    }
153
154    pub fn slot(&self, instance: BufferId) -> &CbSlot {
155        &self.slots[self.by_instance[&instance]]
156    }
157}
logicaffeine_compile/spatial/lower/tensix/cb.rs

logicaffeine_compile/spatial/lower/tensix/
cb.rs
