nxpu_opt/
fma_fusion.rs

1//! FMA (fused multiply-add) fusion pass.
2//!
3//! Detects `a * b + c` patterns and replaces them with `fma(a, b, c)`.
4//! This is a common NPU/GPU optimization that reduces two operations to one.
5
6use std::collections::HashMap;
7
8use nxpu_ir::{Arena, BinaryOp, Expression, Function, Handle, MathFunction, Module};
9
10use crate::Pass;
11use crate::dce::expression_operands;
12
13/// Fuses multiply-add patterns into `fma` math operations.
14#[derive(Debug)]
15pub struct FmaFusion;
16
17impl Pass for FmaFusion {
18    fn name(&self) -> &str {
19        "fma-fusion"
20    }
21
22    fn run(&self, module: &mut Module) -> bool {
23        let mut changed = false;
24        for (_, func) in module.functions.iter_mut() {
25            changed |= run_on_function(func);
26        }
27        for ep in &mut module.entry_points {
28            changed |= run_on_function(&mut ep.function);
29        }
30        changed
31    }
32}
33
34/// Count how many times each expression handle is referenced as an operand.
35fn build_ref_counts(arena: &Arena<Expression>) -> HashMap<Handle<Expression>, usize> {
36    let mut counts: HashMap<Handle<Expression>, usize> = HashMap::new();
37    for (_, expr) in arena.iter() {
38        for operand in expression_operands(expr) {
39            *counts.entry(operand).or_insert(0) += 1;
40        }
41    }
42    counts
43}
44
45fn run_on_function(func: &mut Function) -> bool {
46    let mut changed = false;
47
48    let ref_counts = build_ref_counts(&func.expressions);
49    let handles: Vec<Handle<Expression>> = func.expressions.iter().map(|(h, _)| h).collect();
50
51    for handle in handles {
52        let replacement = match &func.expressions[handle] {
53            Expression::Binary {
54                op: BinaryOp::Add,
55                left,
56                right,
57            } => try_fuse_fma(&func.expressions, *left, *right, &ref_counts),
58            _ => None,
59        };
60
61        if let Some(new_expr) = replacement {
62            func.expressions[handle] = new_expr;
63            changed = true;
64        }
65    }
66
67    changed
68}
69
70/// Checks if `left + right` can be fused into `fma(a, b, c)`.
71///
72/// Matches:
73/// - `(a * b) + c`  →  `fma(a, b, c)`
74/// - `c + (a * b)`  →  `fma(a, b, c)`
75///
76/// Skips fusion if the multiply expression is referenced more than once,
77/// since replacing it would change semantics for other users.
78fn try_fuse_fma(
79    exprs: &Arena<Expression>,
80    left: Handle<Expression>,
81    right: Handle<Expression>,
82    ref_counts: &HashMap<Handle<Expression>, usize>,
83) -> Option<Expression> {
84    // Pattern 1: left = a * b, addend = right
85    if let Expression::Binary {
86        op: BinaryOp::Multiply,
87        left: a,
88        right: b,
89    } = &exprs[left]
90    {
91        if ref_counts.get(&left).copied().unwrap_or(0) > 1 {
92            return None;
93        }
94        return Some(Expression::Math {
95            fun: MathFunction::Fma,
96            arg: *a,
97            arg1: Some(*b),
98            arg2: Some(right),
99            arg3: None,
100        });
101    }
102
103    // Pattern 2: right = a * b, addend = left
104    if let Expression::Binary {
105        op: BinaryOp::Multiply,
106        left: a,
107        right: b,
108    } = &exprs[right]
109    {
110        if ref_counts.get(&right).copied().unwrap_or(0) > 1 {
111            return None;
112        }
113        return Some(Expression::Math {
114            fun: MathFunction::Fma,
115            arg: *a,
116            arg1: Some(*b),
117            arg2: Some(left),
118            arg3: None,
119        });
120    }
121
122    None
123}
124
125#[cfg(test)]
126mod tests {
127    use super::*;
128    use nxpu_ir::Literal;
129
130    #[test]
131    fn fuse_mul_add() {
132        let mut func = Function::new("test");
133        let a = func
134            .expressions
135            .append(Expression::Literal(Literal::F32(2.0)));
136        let b = func
137            .expressions
138            .append(Expression::Literal(Literal::F32(3.0)));
139        let c = func
140            .expressions
141            .append(Expression::Literal(Literal::F32(4.0)));
142        let mul = func.expressions.append(Expression::Binary {
143            op: BinaryOp::Multiply,
144            left: a,
145            right: b,
146        });
147        let add = func.expressions.append(Expression::Binary {
148            op: BinaryOp::Add,
149            left: mul,
150            right: c,
151        });
152
153        let changed = run_on_function(&mut func);
154        assert!(changed);
155        match &func.expressions[add] {
156            Expression::Math {
157                fun: MathFunction::Fma,
158                arg,
159                arg1,
160                arg2,
161                arg3,
162            } => {
163                assert_eq!(*arg, a);
164                assert_eq!(*arg1, Some(b));
165                assert_eq!(*arg2, Some(c));
166                assert!(arg3.is_none());
167            }
168            other => panic!("expected Fma, got {other:?}"),
169        }
170    }
171
172    #[test]
173    fn fuse_commuted_add_mul() {
174        // c + (a * b)
175        let mut func = Function::new("test");
176        let a = func
177            .expressions
178            .append(Expression::Literal(Literal::F32(2.0)));
179        let b = func
180            .expressions
181            .append(Expression::Literal(Literal::F32(3.0)));
182        let c = func
183            .expressions
184            .append(Expression::Literal(Literal::F32(4.0)));
185        let mul = func.expressions.append(Expression::Binary {
186            op: BinaryOp::Multiply,
187            left: a,
188            right: b,
189        });
190        let add = func.expressions.append(Expression::Binary {
191            op: BinaryOp::Add,
192            left: c,
193            right: mul,
194        });
195
196        let changed = run_on_function(&mut func);
197        assert!(changed);
198        match &func.expressions[add] {
199            Expression::Math {
200                fun: MathFunction::Fma,
201                arg,
202                arg1,
203                arg2,
204                arg3,
205            } => {
206                assert_eq!(*arg, a);
207                assert_eq!(*arg1, Some(b));
208                assert_eq!(*arg2, Some(c));
209                assert!(arg3.is_none());
210            }
211            other => panic!("expected Fma, got {other:?}"),
212        }
213    }
214
215    #[test]
216    fn no_fusion_without_multiply() {
217        let mut func = Function::new("test");
218        let a = func
219            .expressions
220            .append(Expression::Literal(Literal::F32(1.0)));
221        let b = func
222            .expressions
223            .append(Expression::Literal(Literal::F32(2.0)));
224        let _add = func.expressions.append(Expression::Binary {
225            op: BinaryOp::Add,
226            left: a,
227            right: b,
228        });
229
230        let changed = run_on_function(&mut func);
231        assert!(!changed);
232    }
233
234    #[test]
235    fn no_fusion_on_subtract() {
236        let mut func = Function::new("test");
237        let a = func
238            .expressions
239            .append(Expression::Literal(Literal::F32(2.0)));
240        let b = func
241            .expressions
242            .append(Expression::Literal(Literal::F32(3.0)));
243        let c = func
244            .expressions
245            .append(Expression::Literal(Literal::F32(4.0)));
246        let mul = func.expressions.append(Expression::Binary {
247            op: BinaryOp::Multiply,
248            left: a,
249            right: b,
250        });
251        let _sub = func.expressions.append(Expression::Binary {
252            op: BinaryOp::Subtract,
253            left: mul,
254            right: c,
255        });
256
257        let changed = run_on_function(&mut func);
258        assert!(!changed);
259    }
260
261    #[test]
262    fn no_fusion_multi_use_multiply() {
263        // mul is used by both add1 and add2, so fusion should be skipped.
264        let mut func = Function::new("test");
265        let a = func
266            .expressions
267            .append(Expression::Literal(Literal::F32(2.0)));
268        let b = func
269            .expressions
270            .append(Expression::Literal(Literal::F32(3.0)));
271        let c = func
272            .expressions
273            .append(Expression::Literal(Literal::F32(4.0)));
274        let d = func
275            .expressions
276            .append(Expression::Literal(Literal::F32(5.0)));
277        let mul = func.expressions.append(Expression::Binary {
278            op: BinaryOp::Multiply,
279            left: a,
280            right: b,
281        });
282        // Two additions using the same multiply.
283        let _add1 = func.expressions.append(Expression::Binary {
284            op: BinaryOp::Add,
285            left: mul,
286            right: c,
287        });
288        let _add2 = func.expressions.append(Expression::Binary {
289            op: BinaryOp::Add,
290            left: mul,
291            right: d,
292        });
293
294        let changed = run_on_function(&mut func);
295        assert!(!changed);
296    }
297
298    #[test]
299    fn fusion_single_use_multiply() {
300        // mul is only used once — fusion should proceed.
301        let mut func = Function::new("test");
302        let a = func
303            .expressions
304            .append(Expression::Literal(Literal::F32(2.0)));
305        let b = func
306            .expressions
307            .append(Expression::Literal(Literal::F32(3.0)));
308        let c = func
309            .expressions
310            .append(Expression::Literal(Literal::F32(4.0)));
311        let mul = func.expressions.append(Expression::Binary {
312            op: BinaryOp::Multiply,
313            left: a,
314            right: b,
315        });
316        let add = func.expressions.append(Expression::Binary {
317            op: BinaryOp::Add,
318            left: mul,
319            right: c,
320        });
321
322        let changed = run_on_function(&mut func);
323        assert!(changed);
324        assert!(matches!(
325            &func.expressions[add],
326            Expression::Math {
327                fun: MathFunction::Fma,
328                ..
329            }
330        ));
331    }
332}
nxpu_opt/fma_fusion.rs

nxpu_opt/
fma_fusion.rs
