1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
// Copyright (C) 2025 zk4x
// SPDX-License-Identifier: LGPL-3.0-only WITH Classpath-exception-2.0
//! Fuse multiply-add operations.
//!
//! This module provides optimization for fusing multiply-add (MAD) operations,
//! which combines `x * y + z` patterns into a single MAD instruction.
//! This reduces instruction count and can improve performance.
use super::autotune::Optimization;
use crate::{
Map,
kernel::{BOp, Kernel, Op, UOp},
};
/// Fuse multiply-add operations into MAD instructions.
#[derive(Debug)]
pub struct FuseMad;
impl Optimization for FuseMad {
fn nconfigs(&self) -> u64 {
1
}
fn apply(&self, kernel: &mut Kernel, _config: u64) {
kernel.fuse_mad();
}
}
impl Kernel {
/// Make the `FuseMad` optimization.
pub fn opt_fuse_mad(&self) -> Box<dyn Optimization> {
Box::new(FuseMad)
}
/// Fuse multiply-add operations into MAD instructions.
///
/// This method identifies patterns of the form `x * y + z` and
/// fuses them into a single MAD instruction, reducing instruction
/// count and potentially improving performance.
///
/// The optimization looks for:
///
/// - Binary add where one operand is a multiply
/// - The multiply has a reference count of 1 (used only once)
/// - The multiply and add can be fused into a MAD
pub fn fuse_mad(&mut self) {
let mut op_id = self.head;
let mut rcs = Map::default();
while !op_id.is_null() {
for param in self.ops[op_id].op.parameters() {
rcs.entry(param).and_modify(|rc| *rc += 1).or_insert(1);
}
if let Op::Binary { x: xo, y: yo, bop } = self.ops[op_id].op
&& bop == BOp::Add
{
if let Op::Binary { x, y, bop } = self.ops[xo].op {
if bop == BOp::Mul && rcs[&xo] == 1 {
self.ops[op_id].op = Op::Mad { x, y, z: yo };
}
} else if let Op::Binary { x, y, bop } = self.ops[yo].op
&& bop == BOp::Mul
&& rcs[&yo] == 1
{
self.ops[op_id].op = Op::Mad { x, y, z: xo };
}
}
op_id = self.next_op(op_id);
}
self.verify();
}
/// Find all multiply add operations and unfuse them
pub fn unfuse_mad(&mut self) {
#[cfg(feature = "time")]
let _timer = crate::Timer::new("unfuse_mad");
let mut op_id = self.head;
while !op_id.is_null() {
if let Op::Mad { x, y, z } = self.ops[op_id].op {
let x = self.insert_before(op_id, Op::Binary { x, y, bop: BOp::Mul });
self.ops[op_id].op = Op::Binary { x, y: z, bop: BOp::Add };
}
op_id = self.next_op(op_id);
}
}
/// Fuse reciprocal-of-square-root into a single rsqrt.
///
/// This method identifies `1/sqrt(x)` spelled as
/// `Reciprocal(Sqrt(x))` and fuses it into one `Rsqrt`, which
/// lowers to a single tile op (`rsqrt_tile`, `rsqrt.approx`,
/// `InverseSqrt`, ...) instead of two.
///
/// Like [`Kernel::fuse_mad`], the inner sqrt must be single-use
/// (reference count 1).
pub fn fuse_rsqrt(&mut self) {
let mut op_id = self.head;
let mut rcs = Map::default();
while !op_id.is_null() {
for param in self.ops[op_id].op.parameters() {
rcs.entry(param).and_modify(|rc| *rc += 1).or_insert(1);
}
if let Op::Unary { x: xo, uop } = self.ops[op_id].op
&& uop == UOp::Reciprocal
&& let Op::Unary { x, uop } = self.ops[xo].op
&& uop == UOp::Sqrt
&& rcs[&xo] == 1
{
self.ops[op_id].op = Op::Unary { x, uop: UOp::Rsqrt };
}
op_id = self.next_op(op_id);
}
self.verify();
}
}