logicaffeine_compile/spatial/passes/
fuse_readers.rs

1//! `fuse_readers` — reader fusion / multicast as a GENERAL pass, decided by the cost model.
2//!
3//! Sites whose Load stages read the identical tile stream (a mesh row shares its A blocks, a column its B
4//! blocks) form a GROUP: one site keeps its Load and becomes the sender (it multicasts each block after
5//! landing it), the others become receivers whose Load is replaced by a `Recv`. Every group gets a semaphore
6//! pair (ready counter at the sender, valid flag at the receivers) and its members' NOCs follow the target's
7//! direction rule. The decision is the cost model's comparison of the fused and unfused arrangements of the
8//! SAME plan — the comparison `tile` already made — never a flag.
9
10use super::SpatialReject;
11use crate::spatial::cost::is_streaming;
12use crate::spatial::ir::*;
13use crate::spatial::target::{HandshakeKind, Interconnect, SpatialTarget};
14
15/// Whether the cost model prefers the fused (multicast) arrangement for this placed plan.
16pub fn fusion_wins(p: &SpatialProgram, target: &SpatialTarget) -> bool {
17    let Some(plan) = p.work.plan else {
18        return false;
19    };
20    if is_streaming(&plan) {
21        return false;
22    }
23    // Fusion cannot win on a target with no multicast at all. Costing it first and discovering the missing
24    // interconnect afterwards turns an option the target simply lacks into a hard refusal of a program it
25    // could run perfectly well unfused.
26    //
27    // The SEMAPHORE capability is deliberately NOT checked here. A target that can multicast but cannot
28    // synchronise is an inconsistent description rather than a target without a feature, and `fuse_readers`
29    // refuses it typed so the contradiction surfaces instead of silently degrading to unfused.
30    if !target.provides_interconnect(Interconnect::MulticastRect) {
31        return false;
32    }
33    let (ny, nx) = plan.grid;
34    if nx <= 1 && ny <= 1 {
35        return false;
36    }
37    // A multicast group shares ONE operand panel along a mesh row or column. Once the block grid outgrows the
38    // mesh, `place` wraps items onto sites, so a mesh row holds blocks from several block-rows that share no
39    // panel and there is nothing left to broadcast. Past that point fusion is not a losing option, it is not
40    // an option — which is why the tile pass offers no fused candidate there either.
41    if ny > target.sites.y || nx > target.sites.x {
42        return false;
43    }
44    // ★ THE SAME QUESTION THE PLAN SEARCH ASKED, ASKED THE SAME WAY. Fusing shares an operand across a
45    // rectangle: each shared tile is read from DRAM once and multicast once, instead of being fetched by
46    // every site that needs it. Whether that wins is a movement question, and it is settled by the work
47    // the two choreographies will actually issue — not by a counter that does not track work.
48    let objective = crate::spatial::work::Objective::default();
49    let (Some(fused), Some(unfused)) = (
50        crate::spatial::work::work_of_plan(p, &plan, true, target),
51        crate::spatial::work::work_of_plan(p, &plan, false, target),
52    ) else {
53        return false;
54    };
55    objective.key(&fused) < objective.key(&unfused)
56}
57
58/// The pass: fuse iff the cost model says so.
59pub fn fuse_readers(
60    p: &SpatialProgram,
61    target: &SpatialTarget,
62) -> Result<SpatialProgram, SpatialReject> {
63    let has_matmul = p
64        .stages
65        .iter()
66        .any(|s| matches!(s.kind, StageKind::Compute(ComputeOp::MatmulTiles { .. })));
67    fuse_readers_forced(p, target, has_matmul && fusion_wins(p, target))
68}
69
70/// Fuse (or not) regardless of the cost model — a calibration facility, never a pass.
71pub fn fuse_readers_forced(
72    p: &SpatialProgram,
73    target: &SpatialTarget,
74    fused: bool,
75) -> Result<SpatialProgram, SpatialReject> {
76    let Some(pl) = &p.placement else {
77        return Err(SpatialReject::WrongPhase {
78            pass: "fuse_readers",
79            expected: "Placed",
80            found: p.phase(),
81        });
82    };
83    if !pl.groups.is_empty() || !fused {
84        return Ok(p.clone());
85    }
86    if !target.provides_interconnect(Interconnect::MulticastRect) {
87        return Err(SpatialReject::MissingInterconnect {
88            needed: Interconnect::MulticastRect,
89            target: target.name,
90        });
91    }
92    if !target.provides_handshake(HandshakeKind::Semaphore) {
93        return Err(SpatialReject::MissingHandshake {
94            needed: HandshakeKind::Semaphore,
95            target: target.name,
96        });
97    }
98    let plan = p.work.plan.expect("placed programs are tiled");
99    let (ny, nx) = plan.grid;
100    let rule = target.noc_rule;
101    let mut q = p.clone();
102    let pl = q.placement.as_mut().unwrap();
103    let sites: Vec<Site> = pl.items_of_site.iter().map(|(s, _)| *s).collect();
104    let mut groups = Vec::new();
105    // Row groups share A (operand 0); column groups share B (operand 1).
106    let mut specs: Vec<(Direction, Operand, Rect, Site)> = Vec::new();
107    if nx > 1 {
108        for y in 0..ny {
109            let rect = Rect {
110                x0: 0,
111                y0: y as u16,
112                x1: (nx - 1) as u16,
113                y1: y as u16,
114            };
115            specs.push((Direction::Row, Operand::A, rect, Site::new(0, y)));
116        }
117    }
118    if ny > 1 {
119        for x in 0..nx {
120            let rect = Rect {
121                x0: x as u16,
122                y0: 0,
123                x1: x as u16,
124                y1: (ny - 1) as u16,
125            };
126            specs.push((Direction::Col, Operand::B, rect, Site::new(x, 0)));
127        }
128    }
129    let mut sems_of: std::collections::BTreeMap<Operand, (SemId, SemId)> = Default::default();
130    for (direction, operand, rect, sender) in specs {
131        let members = rect.sites();
132        if !members.iter().all(|m| sites.contains(m)) {
133            continue;
134        }
135        let counts: std::collections::BTreeSet<usize> =
136            members.iter().map(|m| pl.items_at(*m).len()).collect();
137        if counts.len() != 1 {
138            return Err(SpatialReject::GroupCountMismatch {
139                group: GroupId::new(groups.len()),
140                counts: members
141                    .iter()
142                    .map(|m| (*m, pl.items_at(*m).len()))
143                    .collect(),
144            });
145        }
146        let gid = GroupId::new(groups.len());
147        let noc = NocId(rule.map(|r| r.sender_noc).unwrap_or(0));
148        groups.push(Group {
149            id: gid,
150            rect,
151            sender,
152            direction,
153            operand,
154            noc,
155        });
156        let tensor = match operand {
157            Operand::A => TensorId(0),
158            Operand::B => TensorId(1),
159        };
160        let (ready, valid) = match sems_of.get(&operand) {
161            Some(pair) => *pair,
162            None => {
163                let ready = SemId::new(q.semaphores.len());
164                q.semaphores.push(Semaphore {
165                    id: ready,
166                    init: 0,
167                    role: SemRole::ReadyCounter {
168                        expected: members.len() - 1,
169                    },
170                    operand: Some(operand),
171                    sites: Vec::new(),
172                });
173                let valid = SemId::new(q.semaphores.len());
174                q.semaphores.push(Semaphore {
175                    id: valid,
176                    init: 0,
177                    role: SemRole::ValidFlag,
178                    operand: Some(operand),
179                    sites: Vec::new(),
180                });
181                sems_of.insert(operand, (ready, valid));
182                (ready, valid)
183            }
184        };
185        for sem in [ready, valid] {
186            for m in &members {
187                if !q.semaphores[sem.idx()].sites.contains(m) {
188                    q.semaphores[sem.idx()].sites.push(*m);
189                }
190            }
191        }
192        let mut sender_stage = None;
193        let mut sender_buf = None;
194        let mut receivers = Vec::new();
195        for s in q.stages.iter_mut() {
196            let Some(site) = s.site else { continue };
197            if !rect.contains(site) {
198                continue;
199            }
200            let is_load = matches!(&s.kind, StageKind::Load { tensor: t, .. } if *t == tensor);
201            if !is_load {
202                continue;
203            }
204            if site == sender {
205                if let StageKind::Load { multicast, .. } = &mut s.kind {
206                    *multicast = Some(gid);
207                }
208                sender_stage = Some(s.id);
209                sender_buf = s.outs.first().copied();
210            } else {
211                s.kind = StageKind::Recv { group: gid };
212                s.name = s.name.replace("load_", "recv_");
213                s.ins.clear();
214                receivers.push(s.id);
215            }
216        }
217        let (Some(from), Some(buffer)) = (sender_stage, sender_buf) else {
218            return Err(SpatialReject::Invalid(IrViolation::GroupMalformed {
219                group: gid,
220                reason: "no sender load stage",
221            }));
222        };
223        q.edges.push(Edge {
224            id: EdgeId::new(q.edges.len()),
225            from,
226            buffer,
227            to: receivers,
228            handshake: Handshake::Semaphore {
229                ready,
230                valid,
231                ready_count: members.len() - 1,
232            },
233        });
234    }
235    let pl = q.placement.as_mut().unwrap();
236    pl.groups = groups;
237    // NOC assignment is keyed off the SITE's role in its multicast groups, which only exist once fusion has
238    // run — `place` cannot know it. A row sender drives the sender interface; a site that senders operand A
239    // keeps it for its own reads; everything else reads on the receiver interface; and the writer takes the
240    // complement, with the first column swapped because its multicasts travel the other way.
241    //
242    // This is the vendor's documented convention. We briefly keyed it off the ENGINE instead, on the strength
243    // of a simulator measurement that looked like a 30% win — and that counter was later proved not to track
244    // work at all, so the evidence for the change evaporated. Reverted rather than kept on a number we can no
245    // longer stand behind.
246    if let Some(r) = rule {
247        let groups = pl.groups.clone();
248        let a_senders: std::collections::BTreeSet<Site> = groups
249            .iter()
250            .filter(|g| g.operand == Operand::A)
251            .map(|g| g.sender)
252            .collect();
253        for (i, s) in q.stages.iter().enumerate() {
254            let Some(site) = s.site else { continue };
255            let noc = match &s.kind {
256                StageKind::Load {
257                    multicast: Some(g), ..
258                } if groups[g.idx()].direction == Direction::Row => r.sender_noc,
259                StageKind::Load { .. } | StageKind::Recv { .. } => {
260                    if a_senders.contains(&site) {
261                        r.sender_noc
262                    } else {
263                        r.receiver_noc
264                    }
265                }
266                StageKind::Store { .. } => {
267                    if site.x == 0 {
268                        r.writer_noc_first_column
269                    } else {
270                        r.writer_noc_default
271                    }
272                }
273                _ => continue,
274            };
275            pl.engine_of[i].noc = Some(NocId(noc));
276        }
277        // ★ TWO MOVEMENT ENGINES ON ONE SITE MAY NOT SHARE A NOC.
278        //
279        // The role rule above assigns by ROLE, and a site can hold two roles: an A-multicast sender is
280        // also its own B reader, so both of its engines asked for `sender_noc`; a receiver site asked for
281        // `receiver_noc` twice. Each RISC-V keeps its OWN software count of issued NOC transactions but
282        // they then share one hardware counter, so `noc_async_read_barrier()` compares its count against a
283        // counter carrying the other engine's traffic and NEVER CLEARS.
284        //
285        // Measured: at 128³ the multicast program wedged with `dm1_a_send` parked in
286        // `noc_async_read_barrier()` (tt-metal's watcher: core (0,0) `NRBW,NRBW`), while the blocked
287        // program at the same shape — whose engines land on NOC 0 and NOC 1 — ran in seconds at PCC 1.0.
288        //
289        // The role preference is kept where it is free, and only the COLLIDING engine is moved, so the
290        // vendor's convention still holds wherever it can.
291        let mut taken: std::collections::BTreeMap<(Site, NocId), Engine> = Default::default();
292        for slot in pl.engine_of.iter_mut() {
293            let Some(noc) = slot.noc else { continue };
294            match taken.get(&(slot.site, noc)) {
295                Some(e) if *e != slot.engine => {
296                    let other = NocId(1 - noc.0);
297                    slot.noc = Some(other);
298                    taken.insert((slot.site, other), slot.engine);
299                }
300                _ => {
301                    taken.insert((slot.site, noc), slot.engine);
302                }
303            }
304        }
305    }
306    q.requires.handshakes.insert(HandshakeKind::Semaphore);
307    q.requires.interconnect.insert(Interconnect::MulticastRect);
308    q.validate()?;
309    Ok(q)
310}
logicaffeine_compile/spatial/passes/fuse_readers.rs

logicaffeine_compile/spatial/passes/
fuse_readers.rs
