rlx_fusion/fusion/
shared_input_matmul.rs

1// RLX — versatile ML compiler + runtime.
2// Copyright (C) 2026 Eugene Hauptmann, Nataliya Kosmyna.
3//
4// This program is free software: you can redistribute it and/or modify
5// it under the terms of the GNU General Public License as published by
6// the Free Software Foundation, version 3.
7//
8// This program is distributed in the hope that it will be useful,
9// but WITHOUT ANY WARRANTY; without even the implied warranty of
10// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
11// GNU General Public License for more details.
12//
13// You should have received a copy of the GNU General Public License
14// along with this program. If not, see <https://www.gnu.org/licenses/>.
15
16//! `shared_input_matmul` — extracted from the `fusion` module for navigability (see `mod.rs`).
17
18#![allow(unused_imports)]
19
20use crate::pass::Pass;
21use rlx_ir::op::*;
22use rlx_ir::*;
23use std::collections::HashMap;
24
25// ── Helper: graph rewriter ──────────────────────────────────────────────
26
27use crate::graph_rewrite::Rewriter;
28
29// ── Pass 1: MatMul + Bias + Activation → FusedMatMulBiasAct ─────────────
30
31use super::*;
32
33/// Detects two MatMul nodes with the same input and concatenates their
34/// weight matrices into a single larger MatMul.
35///
36/// Pattern:
37///   %a = matmul(%x, %w1)
38///   %b = matmul(%x, %w2)
39/// Becomes:
40///   %ab = matmul(%x, concat(%w1, %w2))
41///   %a = narrow(%ab, ..., 0, n1)
42///   %b = narrow(%ab, ..., n1, n2)
43///
44/// This saves one full input read (the shared input is read once instead
45/// of twice). Critical for SwiGLU (fc11+fc12) and QKV fusion.
46pub struct FuseSharedInputMatMul;
47
48impl Pass for FuseSharedInputMatMul {
49    fn name(&self) -> &str {
50        "fuse_shared_input_matmul"
51    }
52
53    fn run(&self, graph: Graph) -> Graph {
54        struct FuseGroup {
55            input_id: NodeId,
56            matmul_ids: Vec<NodeId>,
57        }
58
59        let mut input_to_matmuls: HashMap<NodeId, Vec<NodeId>> = HashMap::new();
60        for node in graph.nodes() {
61            if matches!(node.op, Op::MatMul) {
62                input_to_matmuls
63                    .entry(node.inputs[0])
64                    .or_default()
65                    .push(node.id);
66            }
67        }
68
69        let mut groups: Vec<FuseGroup> = Vec::new();
70        for (input_id, matmul_ids) in input_to_matmuls {
71            if matmul_ids.len() < 2 {
72                continue;
73            }
74            let first = graph.node(matmul_ids[0]);
75            let w0 = graph.shape(first.inputs[1]);
76            if w0.rank() != 2 {
77                continue;
78            }
79            let compatible = matmul_ids.iter().all(|&id| {
80                let m = graph.node(id);
81                matches!(m.op, Op::MatMul)
82                    && graph.shape(m.inputs[1]).rank() == 2
83                    && graph.shape(m.inputs[1]).dim(0) == w0.dim(0)
84            });
85            if compatible {
86                groups.push(FuseGroup {
87                    input_id,
88                    matmul_ids,
89                });
90            }
91        }
92
93        if groups.is_empty() {
94            return graph;
95        }
96
97        let group_by_first: HashMap<NodeId, &FuseGroup> =
98            groups.iter().map(|g| (g.matmul_ids[0], g)).collect();
99
100        let mut fused_away: HashMap<NodeId, ()> = HashMap::new();
101        for g in &groups {
102            for &id in &g.matmul_ids[1..] {
103                fused_away.insert(id, ());
104            }
105        }
106
107        let mut rw = Rewriter::new(&graph.name);
108        for node in graph.nodes() {
109            if fused_away.contains_key(&node.id) {
110                continue;
111            }
112
113            if let Some(group) = group_by_first.get(&node.id) {
114                let matmuls: Vec<_> = group.matmul_ids.iter().map(|&id| graph.node(id)).collect();
115                let weight_ids: Vec<NodeId> = matmuls.iter().map(|m| m.inputs[1]).collect();
116                rw.ensure_mapped(&graph, std::slice::from_ref(&group.input_id));
117                rw.ensure_mapped(&graph, &weight_ids);
118
119                let w0_shape = graph.shape(weight_ids[0]);
120                let k = w0_shape.dim(0).unwrap_static();
121                let ns: Vec<usize> = weight_ids
122                    .iter()
123                    .map(|&w| graph.shape(w).dim(1).unwrap_static())
124                    .collect();
125                let combined_n: usize = ns.iter().sum();
126
127                let concat_shape = Shape::new(&[k, combined_n], w0_shape.dtype());
128                let concat_id = rw.add_fused(Op::Concat { axis: 1 }, &weight_ids, concat_shape);
129
130                let out_rank = matmuls[0].shape.rank();
131                let mut mm_dims: Vec<Dim> =
132                    (0..out_rank).map(|i| matmuls[0].shape.dim(i)).collect();
133                mm_dims[out_rank - 1] = Dim::Static(combined_n);
134                let mm_shape = Shape::from_dims(&mm_dims, matmuls[0].shape.dtype());
135                let mm_id = rw.new_graph.add_node(
136                    Op::MatMul,
137                    vec![rw.map(group.input_id), concat_id],
138                    mm_shape,
139                );
140
141                let mut start = 0usize;
142                for (mm, &n) in matmuls.iter().zip(&ns) {
143                    let narrow = rw.new_graph.add_node(
144                        Op::Narrow {
145                            axis: out_rank - 1,
146                            start,
147                            len: n,
148                        },
149                        vec![mm_id],
150                        mm.shape.clone(),
151                    );
152                    rw.replace(mm.id, narrow);
153                    start += n;
154                }
155                continue;
156            }
157
158            rw.copy_node(node);
159        }
160
161        rw.finish(&graph.outputs)
162    }
163}
164
165// ── Pass 4: Detect SwiGLU pattern → FusedSwiGLU ────────────────────────
rlx_fusion/fusion/shared_input_matmul.rs

rlx_fusion/fusion/
shared_input_matmul.rs
