logicaffeine_compile/spatial/passes/fuse_generic.rs
1//! `fuse_generic` — put back into one register window what legalization had to take apart.
2//!
3//! ★ THE POINT OF SPLITTING FIRST AND FUSING AFTER. A fused operator written by hand is a claim
4//! nobody can check: it is one opaque kernel, and whether it computes the composition it is named
5//! after is a matter of reading it carefully. Legalizing into stages and then fusing them makes the
6//! same kernel a REWRITE — and a rewrite has a before and an after that `equivalent_under`
7//! (`spatial/equiv.rs`) can evaluate against each other. The performance is identical; the
8//! difference is that one of them is checkable.
9//!
10//! What this fuses: two adjacent `ComputeOp::Generic` stages where the second consumes the first's
11//! only output, both walk the same fully-parallel iteration space, and the composition is one the
12//! destination registers can carry — the outer operation's second operand must be a plain tile,
13//! because the broadcast instruction family (`mul_tiles_bcast_cols` and friends) reads two circular
14//! buffers and has no dest-reuse spelling. Everything else is left alone rather than guessed at.
15//!
16//! It runs at `Phase::Template`, before `tile` and `place`, so removing a stage costs only a dense
17//! renumber of stages/buffers/edges — there is no placement, memory map or schedule to keep in step.
18
19use super::SpatialReject;
20use crate::codegen_sva::fabric_target::Precision;
21use crate::spatial::ir::*;
22use crate::spatial::target::SpatialTarget;
23
24/// Substitute `inner` for `Operand(slot)` inside `body`, shifting every other operand index by the
25/// map `shift` gives it.
26fn splice(
27 body: &TensorBody,
28 slot: usize,
29 inner: &TensorBody,
30 shift: &dyn Fn(usize) -> usize,
31) -> TensorBody {
32 match body {
33 TensorBody::Operand(i) if *i == slot => inner.clone(),
34 TensorBody::Operand(i) => TensorBody::Operand(shift(*i)),
35 TensorBody::Const(c) => TensorBody::Const(*c),
36 TensorBody::Add(a, b) => TensorBody::Add(
37 Box::new(splice(a, slot, inner, shift)),
38 Box::new(splice(b, slot, inner, shift)),
39 ),
40 TensorBody::Sub(a, b) => TensorBody::Sub(
41 Box::new(splice(a, slot, inner, shift)),
42 Box::new(splice(b, slot, inner, shift)),
43 ),
44 TensorBody::Mul(a, b) => TensorBody::Mul(
45 Box::new(splice(a, slot, inner, shift)),
46 Box::new(splice(b, slot, inner, shift)),
47 ),
48 TensorBody::Min(a, b) => TensorBody::Min(
49 Box::new(splice(a, slot, inner, shift)),
50 Box::new(splice(b, slot, inner, shift)),
51 ),
52 TensorBody::Max(a, b) => TensorBody::Max(
53 Box::new(splice(a, slot, inner, shift)),
54 Box::new(splice(b, slot, inner, shift)),
55 ),
56 TensorBody::FloorDiv(a, b) => TensorBody::FloorDiv(
57 Box::new(splice(a, slot, inner, shift)),
58 Box::new(splice(b, slot, inner, shift)),
59 ),
60 TensorBody::Shr { arg, by } => TensorBody::Shr {
61 arg: Box::new(splice(arg, slot, inner, shift)),
62 by: *by,
63 },
64 TensorBody::Unary { kind, arg } => TensorBody::Unary {
65 kind: *kind,
66 arg: Box::new(splice(arg, slot, inner, shift)),
67 },
68 // A predicate's operand-bound is re-indexed like any operand; it is never the spliced slot
69 // (`fuse_once` declines that case — a bound is an index into the operand list, not a body).
70 TensorBody::Select {
71 cond,
72 then,
73 otherwise,
74 } => TensorBody::Select {
75 cond: cond.map_operand(shift),
76 then: Box::new(splice(then, slot, inner, shift)),
77 otherwise: Box::new(splice(otherwise, slot, inner, shift)),
78 },
79 }
80}
81
82/// Re-index a body's operands through `shift` (used for the inner body, whose operands move to the
83/// front of the fused operand list).
84fn reindex(body: &TensorBody, shift: &dyn Fn(usize) -> usize) -> TensorBody {
85 splice(body, usize::MAX, &TensorBody::Const(0), shift)
86}
87
88/// A map as its points see it on `iters`: an axis of extent 1 is always 0, so a term along it reads nothing, and two
89/// maps that differ only there read the same element at every point.
90fn on_space(map: &ElemAccess, iters: &[IterSpec]) -> ElemAccess {
91 let mut terms: Vec<(LoopVar, i64)> = map
92 .terms
93 .iter()
94 .copied()
95 .filter(|(v, stride)| *stride != 0 && iters.iter().any(|ax| ax.var == *v && ax.extent > 1))
96 .collect();
97 terms.sort();
98 ElemAccess { terms, ..map.clone() }
99}
100
101/// All-parallel and the same extents — the precondition for two stages to share a window at all.
102fn same_parallel_space(a: &[IterSpec], b: &[IterSpec]) -> bool {
103 a.len() == b.len()
104 && a.iter().zip(b).all(|(x, y)| {
105 x.var == y.var
106 && x.extent == y.extent
107 && matches!(x.kind, IterKind::Parallel)
108 && matches!(y.kind, IterKind::Parallel)
109 })
110}
111
112/// One fusion, or `None` when nothing in the program qualifies.
113fn fuse_once(p: &SpatialProgram) -> Option<SpatialProgram> {
114 for (bi, b) in p.stages.iter().enumerate() {
115 let StageKind::Compute(ComputeOp::Generic {
116 iters: bit,
117 ins: bmaps,
118 out: bout,
119 body: bbody,
120 }) = &b.kind
121 else {
122 continue;
123 };
124 for (slot, feed) in b.ins.iter().enumerate() {
125 // The producer must be a Generic whose ONLY consumer is this stage.
126 let Some(ai) = p.stages.iter().position(|s| {
127 s.outs.first() == Some(feed)
128 && matches!(s.kind, StageKind::Compute(ComputeOp::Generic { .. }))
129 }) else {
130 continue;
131 };
132 if ai == bi {
133 continue;
134 }
135 if p.stages
136 .iter()
137 .enumerate()
138 .any(|(i, s)| i != bi && s.ins.contains(feed))
139 {
140 continue; // the intermediate is read elsewhere; fusing would delete a value someone needs
141 }
142 // ★ ONE READ OF THE INTERMEDIATE, OR NO FUSION. The splice replaces one slot's read; a consumer that reads the
143 // intermediate in another slot too (a rotation reads each element and its partner across the half of its
144 // head) would be left reading a buffer the fusion deletes.
145 if b.ins.iter().enumerate().any(|(i, x)| i != slot && x == feed) {
146 continue;
147 }
148 let a = &p.stages[ai];
149 let StageKind::Compute(ComputeOp::Generic {
150 iters: ait,
151 ins: amaps,
152 out: aout,
153 body: abody,
154 }) = &a.kind
155 else {
156 continue;
157 };
158 if !same_parallel_space(ait, bit) {
159 continue;
160 }
161 // ★ A FOLD KEEPS WHAT THE PAGE BETWEEN KEPT. The producer packs the intermediate into its page at the page's
162 // width and the consumer reads that back; the spliced body carries it in the destination registers instead.
163 // A bf16 or fp16 page rounds what a register holds, so across one the two computes stay two, the page between
164 // them rounding as the program apart rounds.
165 if matches!(p.buffer(*feed).dtype, Precision::Bf16 | Precision::Fp16 | Precision::Fp8) {
166 continue;
167 }
168 // ★ THE CONSUMER MUST READ THE INTERMEDIATE WHERE IT WAS WRITTEN. Splicing the producer's body
169 // into the slot replaces the READ, so the consumer's map for that slot is discarded — which is
170 // only the same computation when that map is the producer's own output map. A shifted read
171 // (`at(-1)`, a band's base) would pick up the value at the WRONG point, and an indirect term
172 // is not an address at all but a check: a stride-0 gather is how a row write proves its row
173 // is in range, and splicing it away deleted the proof (an out-of-range `Set row` wrote nothing
174 // and reported nothing).
175 if on_space(&bmaps[slot], bit) != on_space(aout, ait) || bmaps[slot].indirect.is_some() {
176 continue;
177 }
178 // A predicate whose BOUND is the intermediate cannot be spliced: a bound is an operand
179 // index, and the producer's body is an expression, not an operand.
180 if bbody.pred_names_operand(slot) {
181 continue;
182 }
183 // The consumer's OTHER operands must be plain tiles: a broadcast has no dest-reuse form.
184 if bmaps
185 .iter()
186 .enumerate()
187 .any(|(i, m)| i != slot && bit.iter().any(|ax| ax.extent > 1 && !m.varies_along(ax.var)))
188 {
189 continue;
190 }
191 // Operand list: the producer's operands first, then the consumer's others in order.
192 let mut ins: Vec<BufferId> = a.ins.clone();
193 let mut maps: Vec<ElemAccess> = amaps.clone();
194 let mut shift = vec![usize::MAX; b.ins.len()];
195 for (i, (buf, m)) in b.ins.iter().zip(bmaps).enumerate() {
196 if i == slot {
197 continue;
198 }
199 shift[i] = ins.len();
200 ins.push(*buf);
201 maps.push(m.clone());
202 }
203 // ★ THE COMPOSED SIDE MUST END UP ON THE LEFT. A fused body is emitted by leaving the
204 // inner result in the destination registers and having the outer operation consume DST
205 // via its dest-reuse form — and dest-reuse takes the OTHER operand from a circular
206 // buffer. So `f(x) ⊕ b` is emittable and `b ⊕ f(x)` is not. Substituting blindly produced
207 // the second shape (fusing `rsqrt` into `rescale` gave `x · rsqrt(...)`), which lowered to
208 // a refusal the printer was right to make and this pass was wrong to cause. Addition and
209 // multiplication commute, so the operands are swapped; subtraction does not, so it is
210 // left un-fused rather than reassociated into something that computes a different thing.
211 let bbody = match bbody {
212 TensorBody::Add(x, y) if matches!(**y, TensorBody::Operand(i) if i == slot) => {
213 TensorBody::Add(y.clone(), x.clone())
214 }
215 TensorBody::Mul(x, y) if matches!(**y, TensorBody::Operand(i) if i == slot) => {
216 TensorBody::Mul(y.clone(), x.clone())
217 }
218 TensorBody::Sub(_, y) if matches!(**y, TensorBody::Operand(i) if i == slot) => {
219 continue
220 }
221 other => other.clone(),
222 };
223 let inner = reindex(abody, &|i| i);
224 let body = splice(&bbody, slot, &inner, &|i| {
225 shift.get(i).copied().unwrap_or(i)
226 });
227
228 // ── Rebuild without stage `ai` and without the intermediate buffer ──────────────────
229 let dead_buf = *feed;
230 let mut q = p.clone();
231 q.stages[bi].kind = StageKind::Compute(ComputeOp::Generic {
232 iters: bit.clone(),
233 ins: maps,
234 out: bout.clone(),
235 body,
236 });
237 q.stages[bi].ins = ins;
238 let keep_stage: Vec<bool> = (0..q.stages.len()).map(|i| i != ai).collect();
239 let keep_buf: Vec<bool> = q.buffers.iter().map(|x| x.id != dead_buf).collect();
240 let smap: Vec<usize> = keep_stage
241 .iter()
242 .scan(0, |n, k| {
243 let v = *n;
244 if *k {
245 *n += 1
246 }
247 Some(v)
248 })
249 .collect();
250 let bmap: Vec<usize> = keep_buf
251 .iter()
252 .scan(0, |n, k| {
253 let v = *n;
254 if *k {
255 *n += 1
256 }
257 Some(v)
258 })
259 .collect();
260 let rs = |s: StageId| StageId::new(smap[s.idx()]);
261 let rb = |x: BufferId| BufferId::new(bmap[x.idx()]);
262
263 let mut stages: Vec<Stage> = Vec::new();
264 for (i, s) in q.stages.iter().enumerate() {
265 if !keep_stage[i] {
266 continue;
267 }
268 let mut s = s.clone();
269 s.id = StageId::new(stages.len());
270 s.ins = s.ins.iter().map(|x| rb(*x)).collect();
271 s.outs = s.outs.iter().map(|x| rb(*x)).collect();
272 s.kind.remap_buffers(&rb);
273 stages.push(s);
274 }
275 let mut buffers: Vec<Buffer> = Vec::new();
276 for (i, x) in q.buffers.iter().enumerate() {
277 if !keep_buf[i] {
278 continue;
279 }
280 let mut x = x.clone();
281 x.id = BufferId::new(buffers.len());
282 x.alias_of = x.alias_of.map(rb);
283 buffers.push(x);
284 }
285 let mut edges: Vec<Edge> = Vec::new();
286 for e in &q.edges {
287 if e.buffer == dead_buf {
288 continue; // the hand-off that fusion removed
289 }
290 let mut e = e.clone();
291 // The producer's operand edges now feed the fused stage.
292 let to: Vec<StageId> =
293 e.to.iter()
294 .map(|t| {
295 if *t == StageId::new(ai) {
296 StageId::new(bi)
297 } else {
298 *t
299 }
300 })
301 .collect();
302 e.id = EdgeId::new(edges.len());
303 e.from = rs(e.from);
304 e.to = to.iter().map(|t| rs(*t)).collect();
305 e.buffer = rb(e.buffer);
306 edges.push(e);
307 }
308 q.stages = stages;
309 q.buffers = buffers;
310 q.edges = edges;
311 if q.validate().is_err() {
312 continue; // never hand on a program this pass could not keep legal
313 }
314 return Some(q);
315 }
316 }
317 None
318}
319
320/// Fuse until nothing more qualifies.
321pub fn fuse_generic(
322 p: &SpatialProgram,
323 _target: &SpatialTarget,
324) -> Result<SpatialProgram, SpatialReject> {
325 let mut cur = p.clone();
326 // Bounded by the stage count: each fusion removes one stage, so this cannot spin.
327 for _ in 0..p.stages.len() {
328 match fuse_once(&cur) {
329 Some(next) => cur = next,
330 None => break,
331 }
332 }
333 Ok(cur)
334}