rlx_fusion/fusion/transformer_layer.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//! `transformer_layer` — 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/// Fuses an entire BERT-style transformer layer (attention block + residual+LN +
34/// FFN + residual+LN) into one [`Op::FusedTransformerLayer`] node.
35///
36/// Pattern (after [`FuseMatMulBiasAct`], [`FuseResidualLN`], and
37/// [`FuseAttentionBlock`] have run — order matters):
38///
39/// ```text
40/// skip ──┬─→ FusedAttentionBlock(qkv_w, out_w, mask, qkv_b, out_b) ─→ attn_out
41/// └─→ FusedResidualLN(attn_out, skip, ln1_g, ln1_b) ─→ h1
42/// ├─→ FusedMatMulBiasAct(fc1_w, fc1_b, GeLU) ─→ ffn_int
43/// │ ↓
44/// │ FusedMatMulBiasAct(fc2_w, fc2_b, None) ─→ ffn_out
45/// └────────────────────→ FusedResidualLN(ffn_out, h1, ln2_g, ln2_b) ─→ out
46/// ```
47///
48/// All five nodes collapse into a single `FusedTransformerLayer { num_heads,
49/// head_dim, intermediate_size, eps1, eps2, activation, has_bias: true }`
50/// with the 14-input layout consumed by `rlx-mlx`'s lowering at
51/// `rlx-mlx/src/lower.rs:1528`:
52/// `[hidden, qkv_w, qkv_b, out_w, out_b, ln1_g, ln1_b, fc1_w, fc1_b, fc2_w, fc2_b, ln2_g, ln2_b, mask]`.
53///
54/// Threshold is the same as [`FuseAttentionBlock`] (`RLX_FUSE_ATTN_THRESHOLD`,
55/// default 64). Backends that don't natively support `FusedTransformerLayer`
56/// un-fuse it back to primitives at compile time; backends that do (MLX) can
57/// emit one monolithic kernel per layer.
58pub struct FuseTransformerLayer;
59
60impl FuseTransformerLayer {
61 fn should_fuse(graph: &Graph) -> bool {
62 // Same gate as FuseAttentionBlock — single-source of truth for
63 // "this graph is small enough for L1-resident block fusion".
64 FuseAttentionBlock::should_fuse(graph)
65 }
66}
67
68impl Pass for FuseTransformerLayer {
69 fn name(&self) -> &str {
70 "fuse_transformer_layer"
71 }
72
73 fn run(&self, graph: Graph) -> Graph {
74 if !Self::should_fuse(&graph) {
75 return graph;
76 }
77
78 // Graph-output guard: any intermediate we'd absorb must not be an
79 // explicit output, otherwise a downstream caller would see the
80 // collapsed result instead of the per-stage tensor it expects.
81 let mut is_output: HashMap<NodeId, ()> = HashMap::new();
82 for &oid in &graph.outputs {
83 is_output.insert(oid, ());
84 }
85
86 struct LayerMatch {
87 attn_id: NodeId,
88 ln1_id: NodeId,
89 fc1_id: NodeId,
90 fc2_id: NodeId,
91 ln2_id: NodeId,
92 inputs: [NodeId; 14],
93 num_heads: usize,
94 head_dim: usize,
95 intermediate_size: usize,
96 eps1: f32,
97 eps2: f32,
98 activation: Activation,
99 out_shape: Shape,
100 }
101
102 let mut matches: Vec<LayerMatch> = Vec::new();
103 let mut fused_away: HashMap<NodeId, ()> = HashMap::new();
104
105 for node in graph.nodes() {
106 // Anchor on each FusedAttentionBlock — every BERT layer starts here.
107 let Some((num_heads, head_dim, hidden_id, qkv_w, out_w, mask, qkv_b, out_b)) =
108 fused_attn_block_bert(node)
109 else {
110 continue;
111 };
112 let attn_id = node.id;
113 // Attention's only consumer must be the post-attn FusedResidualLN.
114 if graph.use_count(attn_id) != 1 || is_output.contains_key(&attn_id) {
115 continue;
116 }
117 let ln1_id = match graph
118 .nodes()
119 .iter()
120 .find(|n| n.inputs.contains(&attn_id))
121 .map(|n| n.id)
122 {
123 Some(id) => id,
124 None => continue,
125 };
126 let ln1_node = graph.node(ln1_id);
127 let Some((ln1_x, ln1_res, ln1_g, ln1_b, eps1)) = fused_residual_ln_no_bias(ln1_node)
128 else {
129 continue;
130 };
131 // Order in the residual+LN: x = attn_out, residual = skip (= hidden).
132 if ln1_x != attn_id || ln1_res != hidden_id {
133 continue;
134 }
135 // h1 must have exactly 2 consumers (FFN.1 input AND ln2 residual).
136 if graph.use_count(ln1_id) != 2 || is_output.contains_key(&ln1_id) {
137 continue;
138 }
139
140 // Find FFN.1: FusedMatMulBiasAct(h1, fc1_w, fc1_b) with GeLU.
141 let mut fc1_candidate: Option<NodeId> = None;
142 let mut ln2_candidate: Option<NodeId> = None;
143 for cn in graph.nodes() {
144 if !cn.inputs.contains(&ln1_id) {
145 continue;
146 }
147 if fused_mm_bias_act(cn).is_some() && cn.inputs[0] == ln1_id {
148 fc1_candidate = Some(cn.id);
149 } else if fused_residual_ln_no_bias(cn).is_some() && cn.inputs[1] == ln1_id {
150 ln2_candidate = Some(cn.id);
151 }
152 }
153 let (Some(fc1_id), Some(ln2_id)) = (fc1_candidate, ln2_candidate) else {
154 continue;
155 };
156 let fc1_node = graph.node(fc1_id);
157 let Some((_, fc1_w, fc1_b, activation)) = fused_mm_bias_act(fc1_node) else {
158 continue;
159 };
160 // FFN.1 output → FFN.2 (single consumer).
161 if graph.use_count(fc1_id) != 1 || is_output.contains_key(&fc1_id) {
162 continue;
163 }
164 let fc2_id = match graph
165 .nodes()
166 .iter()
167 .find(|n| n.inputs.contains(&fc1_id))
168 .map(|n| n.id)
169 {
170 Some(id) => id,
171 None => continue,
172 };
173 let fc2_node = graph.node(fc2_id);
174 // FFN.2 must be FusedMatMulBiasAct with activation=None.
175 let Some((fc2_in, fc2_w, fc2_b)) = fused_mm_bias_none(fc2_node) else {
176 continue;
177 };
178 if fc2_in != fc1_id {
179 continue;
180 }
181 if graph.use_count(fc2_id) != 1 || is_output.contains_key(&fc2_id) {
182 continue;
183 }
184 // Final residual+LN: x = ffn_out, residual = h1, gamma/beta + eps2.
185 let ln2_node = graph.node(ln2_id);
186 let Some((ln2_x, ln2_res, ln2_g, ln2_b, eps2)) = fused_residual_ln_no_bias(ln2_node)
187 else {
188 continue;
189 };
190 if ln2_x != fc2_id || ln2_res != ln1_id {
191 continue;
192 }
193 // intermediate_size from fc1_w (`[H, intermediate_size]`).
194 let intermediate_size = {
195 let s = &graph.node(fc1_w).shape;
196 if s.rank() != 2 {
197 continue;
198 }
199 let d = s.dim(s.rank() - 1);
200 if !d.is_static() {
201 continue;
202 }
203 d.unwrap_static()
204 };
205
206 matches.push(LayerMatch {
207 attn_id,
208 ln1_id,
209 fc1_id,
210 fc2_id,
211 ln2_id,
212 inputs: [
213 hidden_id, qkv_w, qkv_b, out_w, out_b, ln1_g, ln1_b, fc1_w, fc1_b, fc2_w,
214 fc2_b, ln2_g, ln2_b, mask,
215 ],
216 num_heads,
217 head_dim,
218 intermediate_size,
219 eps1,
220 eps2,
221 activation,
222 out_shape: ln2_node.shape.clone(),
223 });
224 fused_away.insert(attn_id, ());
225 fused_away.insert(ln1_id, ());
226 fused_away.insert(fc1_id, ());
227 fused_away.insert(fc2_id, ());
228 fused_away.insert(ln2_id, ());
229 }
230
231 if matches.is_empty() {
232 return graph;
233 }
234
235 // Index by ln2 (the layer's terminal node) so we know when to emit.
236 let mut by_terminal: HashMap<NodeId, &LayerMatch> = HashMap::new();
237 for m in &matches {
238 by_terminal.insert(m.ln2_id, m);
239 }
240
241 let mut rw = Rewriter::new(&graph.name);
242 for node in graph.nodes() {
243 if fused_away.contains_key(&node.id) {
244 if let Some(m) = by_terminal.get(&node.id) {
245 rw.ensure_mapped(&graph, &m.inputs);
246 let fused_id = rw.add_fused(
247 Op::FusedTransformerLayer {
248 num_heads: m.num_heads,
249 head_dim: m.head_dim,
250 intermediate_size: m.intermediate_size,
251 eps1: m.eps1,
252 eps2: m.eps2,
253 activation: m.activation,
254 has_bias: true,
255 },
256 &m.inputs,
257 m.out_shape.clone(),
258 );
259 rw.replace(m.attn_id, fused_id);
260 rw.replace(m.ln1_id, fused_id);
261 rw.replace(m.fc1_id, fused_id);
262 rw.replace(m.fc2_id, fused_id);
263 rw.replace(node.id, fused_id);
264 }
265 continue;
266 }
267 rw.copy_node(node);
268 }
269 rw.finish(&graph.outputs)
270 }
271}
272
273// ── PLAN L2: MarkElementwiseRegions ─────────────────────────────────────
274//
275// Walk the graph and collapse maximal chains of element-wise ops
276// (Activation / Cast / Binary / Compare) into a single
277// `Op::ElementwiseRegion`. Conditions for inclusion in a chain:
278// - Op is element-wise per `is_elementwise()` (excluding Where which
279// has a 3-input mask semantic that doesn't compose into a single
280// scalar register chain cleanly — keep as separate op for now).
281// - Output shape exactly equals every input shape (no broadcast —
282// broadcast scalar/vector adds register-pattern complexity, defer).
283// - Every intermediate (chain-internal) value has exactly one
284// consumer in the *whole* graph. Multi-consumer values must
285// materialize.
286// The chain start can read graph-level inputs / params / earlier-fused
287// nodes; the chain end is the last single-consumer or terminal node.
288// This is the simplest correct cut — N-ary chain fusion replaces the
289// pairwise `fuse_elementwise_chains` pattern in each backend with one
290// IR-level pass + a single backend kernel. See PLAN L2.
291//
292// Fusion boundaries: chains do not extend across inputs whose producer
293// satisfies [`rlx_ir::Op::is_fusion_boundary`] (BLAS, Gaussian splat, …).