logicaffeine_compile/spatial/lower/tensix/
args.rs

1//! Runtime arguments: the site-dependent values of a kernel class, declared ONCE as a layout that both the
2//! kernel text (`get_arg_val<uint32_t>(i)`) and the host's per-site `SetRuntimeArgs` are printed from.
3
4use super::roles::KernelClass;
5use super::SpatialLowerError;
6use crate::spatial::ir::*;
7use crate::spatial::passes::{LV_ITEM, LV_ITEM_X, LV_ITEM_Y};
8
9#[derive(Clone, Copy, Debug, PartialEq, Eq)]
10pub enum Axis {
11    X,
12    Y,
13}
14
15#[derive(Clone, Debug, PartialEq, Eq)]
16pub enum ArgSource {
17    /// DRAM base address of a tensor.
18    TensorAddr(TensorId),
19    /// First work item (linear id) of the site — items are consecutive in row-major block order.
20    ItemStart,
21    /// Number of work items of the site.
22    ItemCount,
23    /// Strided item enumeration (a wrapped blocked plan): first block row / column and the count per row.
24    ItemY0,
25    ItemX0,
26    ItemCountX,
27    /// Multicast rectangle of the RECEIVERS of the site's group for an operand (near corner / far corner).
28    McastStart {
29        operand: Operand,
30        axis: Axis,
31    },
32    McastEnd {
33        operand: Operand,
34        axis: Axis,
35    },
36    McastNumDests {
37        operand: Operand,
38    },
39    /// Physical coordinate of the sender of the site's group for an operand.
40    McastSender {
41        operand: Operand,
42        axis: Axis,
43    },
44    /// The program-wide semaphore handle (`get_semaphore(id)` in the kernel).
45    SemHandle(SemId),
46    /// Base of the device-cycle scratch, and this site's word in it. Stamped by the kernel that writes the
47    /// output, whose lifetime spans the whole program on its core.
48    CyclesAddr,
49    CyclesSlot(u8),
50    /// A tensor's tile-grid extent (`rows_t` / `cols_t`), handed to a data-movement kernel at run time: the
51    /// ONLY way shape enters a kernel. Every tile stride and loop count is printed as a block-extent literal
52    /// times or over one of these, so the address arithmetic has exactly the vendor's form — constant-folding
53    /// the shape restructures it and slows the NOC issue rate (+844 cycles at 640³ on ttsim), while a runtime
54    /// value the compiler cannot see is 1 costs a multiply per tile (+300 cycles). A movement kernel is thereby
55    /// shape-generic: the same binary serves every shape of the same blocking.
56    TensorDim {
57        tensor: TensorId,
58        axis: Axis,
59    },
60}
61
62#[derive(Clone, Debug, PartialEq, Eq)]
63pub struct ArgSlot {
64    pub name: String,
65    pub source: ArgSource,
66}
67
68#[derive(Clone, Debug, PartialEq, Default)]
69pub struct RuntimeArgLayout {
70    pub per_kernel: Vec<(String, Vec<ArgSlot>)>,
71}
72
73/// How a class enumerates its items: consecutive linear ids, or a (y, x) stride lattice.
74#[derive(Clone, Copy, Debug, PartialEq, Eq)]
75pub enum ItemMode {
76    Linear,
77    Strided,
78}
79
80pub fn operand_tag(o: Operand) -> &'static str {
81    match o {
82        Operand::A => "a",
83        Operand::B => "b",
84    }
85}
86
87fn operand_of_sem(p: &SpatialProgram, sem: SemId) -> Option<Operand> {
88    p.semaphores[sem.idx()].operand
89}
90
91/// Whether every member site's items are consecutive linear ids.
92pub fn item_mode(p: &SpatialProgram, class: &KernelClass) -> Result<ItemMode, SpatialLowerError> {
93    let pl = p.placement.as_ref().unwrap();
94    let mut linear = true;
95    for site in &class.sites {
96        let items = pl.items_at(*site);
97        if items.windows(2).any(|w| w[1].idx() != w[0].idx() + 1) {
98            linear = false;
99        }
100    }
101    if linear {
102        return Ok(ItemMode::Linear);
103    }
104    // Strided: items = {(y0 + iy·sy, x0 + ix·sx)} in row-major (iy, ix) order.
105    for site in &class.sites {
106        let coords: Vec<(usize, usize)> = pl
107            .items_at(*site)
108            .iter()
109            .map(|i| p.work.items[i.idx()].coord)
110            .collect();
111        let xs: std::collections::BTreeSet<usize> = coords.iter().map(|c| c.1).collect();
112        let ys: std::collections::BTreeSet<usize> = coords.iter().map(|c| c.0).collect();
113        let mut expect = Vec::new();
114        for y in &ys {
115            for x in &xs {
116                expect.push((*y, *x));
117            }
118        }
119        if expect != coords {
120            return Err(SpatialLowerError::Unsupported(
121                "a site's work items are neither a linear range nor a stride lattice",
122            ));
123        }
124    }
125    Ok(ItemMode::Strided)
126}
127
128/// The argument slots of a kernel class, in first-use order.
129pub fn slots_for(
130    p: &SpatialProgram,
131    class: &KernelClass,
132) -> Result<Vec<ArgSlot>, SpatialLowerError> {
133    let sched = p.schedule.as_ref().unwrap();
134    let rep = &sched.programs[class.representative];
135    let ops: Vec<ChoreoOp> = rep.ops.iter().map(|t| t.op.clone()).collect();
136    let mut slots: Vec<ArgSlot> = Vec::new();
137    let mut push = |name: String, source: ArgSource| {
138        if !slots.iter().any(|s| s.source == source) {
139            slots.push(ArgSlot { name, source });
140        }
141    };
142    let mut uses_coords = false;
143    let mut has_item_loop = false;
144    let mut err = None;
145    walk_ops(&ops, 1, &mut |op, _| match op {
146        ChoreoOp::ReadTiles { tensor, at, .. } | ChoreoOp::WriteTiles { tensor, at, .. } => {
147            push(
148                format!("t{}_addr", tensor.idx()),
149                ArgSource::TensorAddr(*tensor),
150            );
151            if at
152                .steps
153                .iter()
154                .any(|(v, _)| *v == LV_ITEM_Y || *v == LV_ITEM_X)
155            {
156                uses_coords = true;
157            }
158        }
159        ChoreoOp::Loop { var, .. } if *var == LV_ITEM => has_item_loop = true,
160        _ => {}
161    });
162    let mode = item_mode(p, class)?;
163    if uses_coords {
164        match mode {
165            ItemMode::Linear => {
166                push("item_start".into(), ArgSource::ItemStart);
167            }
168            ItemMode::Strided => {
169                push("item_y0".into(), ArgSource::ItemY0);
170                push("item_x0".into(), ArgSource::ItemX0);
171                push("item_count_x".into(), ArgSource::ItemCountX);
172            }
173        }
174    }
175    if has_item_loop {
176        push("item_count".into(), ArgSource::ItemCount);
177    }
178    walk_ops(&ops, 1, &mut |op, _| match op {
179        ChoreoOp::McastSend { group, .. } | ChoreoOp::McastSetValid { group, .. } => {
180            let operand = p.placement.as_ref().unwrap().groups[group.idx()].operand;
181            let t = operand_tag(operand);
182            push(
183                format!("mc_{t}_start_x"),
184                ArgSource::McastStart {
185                    operand,
186                    axis: Axis::X,
187                },
188            );
189            push(
190                format!("mc_{t}_start_y"),
191                ArgSource::McastStart {
192                    operand,
193                    axis: Axis::Y,
194                },
195            );
196            push(
197                format!("mc_{t}_end_x"),
198                ArgSource::McastEnd {
199                    operand,
200                    axis: Axis::X,
201                },
202            );
203            push(
204                format!("mc_{t}_end_y"),
205                ArgSource::McastEnd {
206                    operand,
207                    axis: Axis::Y,
208                },
209            );
210            push(
211                format!("mc_{t}_num_dests"),
212                ArgSource::McastNumDests { operand },
213            );
214            if let ChoreoOp::McastSetValid { sem, .. } = op {
215                push(format!("sem{}_id", sem.idx()), ArgSource::SemHandle(*sem));
216            }
217        }
218        ChoreoOp::SemInc { sem, .. } => match operand_of_sem(p, *sem) {
219            Some(operand) => {
220                let t = operand_tag(operand);
221                push(
222                    format!("mc_{t}_sender_x"),
223                    ArgSource::McastSender {
224                        operand,
225                        axis: Axis::X,
226                    },
227                );
228                push(
229                    format!("mc_{t}_sender_y"),
230                    ArgSource::McastSender {
231                        operand,
232                        axis: Axis::Y,
233                    },
234                );
235                push(format!("sem{}_id", sem.idx()), ArgSource::SemHandle(*sem));
236            }
237            None => {
238                err = Some(SpatialLowerError::Unsupported(
239                    "a semaphore increment outside a multicast group",
240                ))
241            }
242        },
243        ChoreoOp::SemSet { sem, .. } | ChoreoOp::SemWait { sem, .. } => {
244            push(format!("sem{}_id", sem.idx()), ArgSource::SemHandle(*sem));
245        }
246        ChoreoOp::Barrier { .. } => {
247            err = Some(SpatialLowerError::Unsupported("a group barrier on Tensix"))
248        }
249        _ => {}
250    });
251    if let Some(e) = err {
252        return Err(e);
253    }
254    // EVERY movement engine stamps its own span, and the host takes the max. Timing only the writer measures
255    // the writer, whose work is the output tiles and therefore independent of the reduction depth — MEASURED,
256    // 810 cycles at K=640 against 809 at K=1280. The operand reader is the engine whose work scales with K,
257    // and which of the two is the slower one is exactly what the plan search is trying to change, so the
258    // honest program length is the max over engines rather than any single one.
259    if stamps_cycles(&ops) {
260        let engine = match rep.engine {
261            Engine::Movement(e) => e,
262            Engine::Compute => 0,
263        };
264        push("cycles_addr".into(), ArgSource::CyclesAddr);
265        push("cycles_slot".into(), ArgSource::CyclesSlot(engine));
266    }
267    Ok(slots)
268}
269
270/// Does this engine program touch DRAM? Every such engine is timed; the compute engine cannot be, since a
271/// TRISC has no NOC of its own to write the result back with.
272pub fn stamps_cycles(ops: &[ChoreoOp]) -> bool {
273    let mut found = false;
274    walk_ops(ops, 1, &mut |op, _| {
275        if matches!(op, ChoreoOp::WriteTiles { .. } | ChoreoOp::ReadTiles { .. }) {
276            found = true;
277        }
278    });
279    found
280}
logicaffeine_compile/spatial/lower/tensix/args.rs

logicaffeine_compile/spatial/lower/tensix/
args.rs
