rlx_compile/
fusion_pipeline.rs

1// RLX — versatile ML compiler + runtime.
2// Copyright (C) 2026 Eugene Hauptmann, Nataliya Kosmyna.
3// SPDX-License-Identifier: MIT OR Apache-2.0
4
5//! Centralized fusion pass pipelines per backend target.
6//!
7//! [`fusion_passes_for_supported`] selects passes from a backend's
8//! [`rlx_ir::OpKind`] claim set so fusion never emits fused ops the
9//! target cannot lower. [`fusion_passes`] keeps the legacy
10//! [`FusionTarget`] entry point and delegates to the same selector.
11
12use rlx_ir::OpKind;
13
14use crate::DeadCodeElimination;
15use crate::io_output_gate::SelectPeaksOnlyOutputs;
16use rlx_fusion::control_flow::LowerControlFlow;
17use rlx_fusion::fk_fusion::{
18    DecomposeFusionRegions, FuseBatchPreprocess, FuseRegionPrologue, MarkBatchSliceRegions,
19    MarkTransformRegions,
20};
21use rlx_fusion::fusion::{
22    FuseAdaLayerNorm, FuseAttentionBackwardAll, FuseAttentionBlock, FuseConvAffineAct,
23    FuseConvBiasAct, FuseGatedResidual, FuseMatMulBiasAct, FuseMatMulResidual, FuseResidualLN,
24    FuseResidualRmsNorm, FuseRmsNormReshape, FuseSharedInputMatMul, FuseSwiGLU,
25    FuseSwiGLUDualMatmul, FuseTransformerLayer, MarkElementwiseRegions, UnfuseElementwiseRegions,
26};
27use rlx_fusion::limits::{FusionLimits, with_fusion_limits};
28use rlx_fusion::lower_dot_general::LowerDotGeneral;
29use rlx_fusion::pass::{Pass, run_passes};
30use rlx_ir::Graph;
31
32use crate::fusion_target::with_fusion_target;
33
34/// Compile target that selects a fusion pipeline.
35#[derive(Debug, Clone, Copy, PartialEq, Eq)]
36pub enum FusionTarget {
37    Cpu,
38    Metal,
39    Mlx,
40    Wgpu,
41    Cuda,
42    Rocm,
43    Tpu,
44}
45
46/// Per-target fusion toggles (env-driven on Metal today).
47#[derive(Debug, Clone, Copy, PartialEq, Eq)]
48pub struct FusionOptions {
49    /// Skip all pattern fusions (Metal: `RLX_METAL_NO_FUSION`).
50    pub skip_fusion: bool,
51    /// Break `ElementwiseRegion` back into primitives after marking.
52    pub unfuse_elementwise_regions: bool,
53    /// Keep fused `ElementwiseRegion` through lowering (env: `RLX_KEEP_ELEMENTWISE_REGIONS`).
54    pub keep_elementwise_regions: bool,
55    /// Decompose FKL-style transform / batch regions before backend lowering.
56    pub decompose_fusion_regions: bool,
57    /// Run FKL passes (`MarkTransformRegions`, prologue, batch). Env: `RLX_NO_FK_FUSION=1` disables.
58    pub fk_fusion: bool,
59    /// Fold `ResizeNearest2x` into `ElementwiseRegion` prologue. Env: `RLX_FUSE_REGION_PROLOGUE=0` disables.
60    pub fuse_region_prologue: bool,
61    /// Merge parallel region slices into `BatchElementwiseRegion`. Env: `RLX_FUSE_BATCH_PREPROCESS=0` disables.
62    pub fuse_batch_preprocess: bool,
63    /// Keep `TransformRegion` / `BatchElementwiseRegion` in MIR for native lowering. Env: `RLX_NATIVE_FK_REGIONS=1`.
64    pub native_fk_regions: bool,
65    /// Caps for fused elementwise chains (encoder / scratch limits).
66    pub fusion_limits: FusionLimits,
67    /// Skip Conv+Bias+Act fusion (`RLX_DISABLE_CONV_BIAS_ACT_FUSION`).
68    pub disable_conv_bias_act_fusion: bool,
69    /// Disable IO-gated peaks-only output selection (`RLX_NO_IO_PEAKS_OUTPUT`).
70    pub no_io_peaks_output: bool,
71}
72
73impl Default for FusionOptions {
74    fn default() -> Self {
75        Self {
76            skip_fusion: false,
77            unfuse_elementwise_regions: false,
78            keep_elementwise_regions: false,
79            decompose_fusion_regions: false,
80            fk_fusion: true,
81            fuse_region_prologue: true,
82            fuse_batch_preprocess: true,
83            native_fk_regions: false,
84            fusion_limits: FusionLimits::default(),
85            disable_conv_bias_act_fusion: false,
86            no_io_peaks_output: false,
87        }
88    }
89}
90
91impl FusionOptions {
92    /// Read Metal-specific env overrides.
93    pub fn from_metal_env() -> Self {
94        Self {
95            skip_fusion: rlx_ir::env::flag("RLX_METAL_NO_FUSION"),
96            unfuse_elementwise_regions: rlx_ir::env::flag("RLX_METAL_UNFUSE_REGIONS"),
97            keep_elementwise_regions: rlx_ir::env::flag("RLX_KEEP_ELEMENTWISE_REGIONS"),
98            decompose_fusion_regions: rlx_ir::env::flag("RLX_DECOMPOSE_FUSION_REGIONS"),
99            fk_fusion: !rlx_ir::env::flag("RLX_NO_FK_FUSION"),
100            fuse_region_prologue: if rlx_ir::env::is_unset("RLX_FUSE_REGION_PROLOGUE") {
101                true
102            } else {
103                rlx_ir::env::flag("RLX_FUSE_REGION_PROLOGUE")
104            },
105            fuse_batch_preprocess: if rlx_ir::env::is_unset("RLX_FUSE_BATCH_PREPROCESS") {
106                true
107            } else {
108                rlx_ir::env::flag("RLX_FUSE_BATCH_PREPROCESS")
109            },
110            native_fk_regions: rlx_ir::env::flag("RLX_NATIVE_FK_REGIONS"),
111            ..Self::default()
112        }
113    }
114
115    /// Merge session options with compile-time env overrides.
116    pub fn merge_env(mut self) -> Self {
117        if rlx_ir::env::flag("RLX_METAL_NO_FUSION") {
118            self.skip_fusion = true;
119        }
120        if rlx_ir::env::flag("RLX_METAL_UNFUSE_REGIONS") {
121            self.unfuse_elementwise_regions = true;
122        }
123        if rlx_ir::env::flag("RLX_KEEP_ELEMENTWISE_REGIONS") {
124            self.keep_elementwise_regions = true;
125        }
126        if rlx_ir::env::flag("RLX_DECOMPOSE_FUSION_REGIONS") {
127            self.decompose_fusion_regions = true;
128        }
129        if rlx_ir::env::flag("RLX_NO_FK_FUSION") {
130            self.fk_fusion = false;
131        }
132        if !rlx_ir::env::is_unset("RLX_FUSE_REGION_PROLOGUE") {
133            self.fuse_region_prologue = rlx_ir::env::flag("RLX_FUSE_REGION_PROLOGUE");
134        }
135        if !rlx_ir::env::is_unset("RLX_FUSE_BATCH_PREPROCESS") {
136            self.fuse_batch_preprocess = rlx_ir::env::flag("RLX_FUSE_BATCH_PREPROCESS");
137        }
138        if rlx_ir::env::flag("RLX_NATIVE_FK_REGIONS") {
139            self.native_fk_regions = true;
140        }
141        if rlx_ir::env_registry::flag("RLX_NO_NATIVE_FK_REGIONS") {
142            self.native_fk_regions = false;
143        }
144        if rlx_ir::env_registry::flag("RLX_DISABLE_CONV_BIAS_ACT_FUSION") {
145            self.disable_conv_bias_act_fusion = true;
146        }
147        if rlx_ir::env_registry::flag("RLX_NO_IO_PEAKS_OUTPUT") {
148            self.no_io_peaks_output = true;
149        }
150        self
151    }
152
153    /// GPU-class targets keep native FKL regions unless opted out.
154    pub fn apply_native_fk_defaults(mut self, target: FusionTarget) -> Self {
155        if rlx_ir::env::flag("RLX_NO_NATIVE_FK_REGIONS") {
156            self.native_fk_regions = false;
157            return self;
158        }
159        if self.native_fk_regions || rlx_ir::env::flag("RLX_NATIVE_FK_REGIONS") {
160            self.native_fk_regions = true;
161            return self;
162        }
163        if matches!(
164            target,
165            FusionTarget::Metal
166                | FusionTarget::Cuda
167                | FusionTarget::Rocm
168                | FusionTarget::Wgpu
169                | FusionTarget::Mlx
170                | FusionTarget::Tpu
171        ) {
172            self.native_fk_regions = true;
173        }
174        self
175    }
176
177    /// CPU executes element-wise chains as per-op thunks — mark then unfuse.
178    pub fn for_cpu() -> Self {
179        Self {
180            unfuse_elementwise_regions: true,
181            fusion_limits: FusionLimits::UNBOUNDED,
182            ..Self::default()
183        }
184    }
185
186    /// Metal keeps RMSNorm / matmul fusions but unfuses `ElementwiseRegion` by
187    /// default (fused MSL mis-lowers *long* chains on deep transformer graphs) —
188    /// UNLESS `RLX_KEEP_ELEMENTWISE_REGIONS` asks to keep them (short chains lower
189    /// fine and stay fused → fewer dispatches on the dispatch-bound backward).
190    pub fn for_metal() -> Self {
191        let mut opts = Self::from_metal_env();
192        if !opts.keep_elementwise_regions {
193            opts.unfuse_elementwise_regions = true;
194        }
195        opts
196    }
197
198    /// wgpu region kernel only supports trailing/scalar broadcast via
199    /// modulus — unfuse so LegalizeBroadcast Expand + Binary run separately.
200    pub fn for_wgpu() -> Self {
201        let keep = rlx_ir::env::flag("RLX_KEEP_ELEMENTWISE_REGIONS");
202        Self {
203            unfuse_elementwise_regions: !keep,
204            keep_elementwise_regions: keep,
205            ..Self::default()
206        }
207    }
208}
209
210/// Elementwise-region caps for `target` (matches GPU kernel encoders).
211pub fn fusion_limits_for_target(target: FusionTarget) -> FusionLimits {
212    match target {
213        FusionTarget::Cpu => FusionLimits::UNBOUNDED,
214        FusionTarget::Tpu => FusionLimits {
215            max_elementwise_steps: 32,
216            max_elementwise_inputs: 16,
217        },
218        _ => FusionLimits::GPU_NATIVE,
219    }
220}
221
222/// True when `supported` is empty (no claim) or contains `kind`.
223#[inline]
224pub fn supports_op(supported: &[OpKind], kind: OpKind) -> bool {
225    supported.is_empty() || supported.contains(&kind)
226}
227
228/// Return the ordered fusion passes allowed for `supported`.
229///
230/// When `supported` is empty every fusion pass runs (legacy "accept
231/// all" backends). When non-empty, each pattern fusion pass is
232/// included only if the backend claims the fused [`OpKind`] it
233/// emits. Lowering passes (`LowerControlFlow`, `LowerDotGeneral`) and
234/// `FuseRmsNormReshape` (topology-only) always run unless
235/// `skip_fusion` is set.
236pub fn fusion_passes_for_supported(
237    supported: &[OpKind],
238    opts: FusionOptions,
239    target: FusionTarget,
240) -> Vec<&'static dyn Pass> {
241    let opts = opts.apply_native_fk_defaults(target);
242    if opts.skip_fusion {
243        return vec![&LowerControlFlow, &LowerDotGeneral];
244    }
245
246    let mut passes: Vec<&'static dyn Pass> = vec![&LowerControlFlow, &LowerDotGeneral];
247
248    // ORDER: FuseMatMulBiasAct first, then FuseAttentionBlock. The block-level
249    // pass matches the post-fusion shape
250    //   FusedMatMulBiasAct(qkv) → narrow×3 → Attention → FusedMatMulBiasAct(out)
251    // which is the pattern BERT-family encoders actually present after the
252    // per-layer matmul+bias fusion has collapsed Q, K, V, and out projections.
253    // `RLX_DISABLE_MATMUL_BIAS_FUSION=1` keeps `matmul → add(bias)` decomposed.
254    // Ablation knob: two ports (`rlx-vieeg` on wgpu, `rlx-tinymyo` on every
255    // accelerator) diverge from CPU at the op immediately *after* such a pair,
256    // and there was no way to take the fusion out of the picture to confirm it.
257    if supports_op(supported, OpKind::FusedMatMulBiasAct)
258        && rlx_ir::env::var("RLX_DISABLE_MATMUL_BIAS_FUSION").as_deref() != Some("1")
259    {
260        passes.push(&FuseMatMulBiasAct);
261    }
262    // Conv + bias + activation → cuDNN's fused conv-bias-activation (CUDA only
263    // claims `FusedConvBiasAct`; every other backend keeps it decomposed).
264    // `RLX_DISABLE_CONV_BIAS_ACT_FUSION=1` skips it (ablation / A-B benchmarking
265    // vs the unfused conv+bias+act path). `FuseConvAffineAct` folds a
266    // host-pre-folded BatchNorm affine (`conv→Mul→Add→relu`) into the same op.
267    if supports_op(supported, OpKind::FusedConvBiasAct) && !opts.disable_conv_bias_act_fusion {
268        passes.push(&FuseConvBiasAct);
269        passes.push(&FuseConvAffineAct);
270    }
271    // Block-level fusion: `Op::FusedAttentionBlock`. All backends that claim
272    // this op now produce parity-correct output (the MLX
273    // `Op::FusedAttentionBlock` lowering at `rlx-mlx/src/lower.rs:1689`
274    // historically diverged on `MaskKind::Custom` BERT masks because it
275    // bypassed the binary→additive conversion and the contiguous
276    // materialization the unfused `Op::Attention` path applies — fixed
277    // alongside this pass landing).
278    if supports_op(supported, OpKind::FusedAttentionBlock) {
279        passes.push(&FuseAttentionBlock);
280    }
281    // Collapse the 3 sibling AttentionBackward ops (dQ/dK/dV) into one
282    // AttentionBackwardAll (single score+softmax recompute). Backend-gated:
283    // only fires where the op is lowered (CPU), leaving other backends' own
284    // attention-backward handling untouched.
285    if supports_op(supported, OpKind::AttentionBackwardAll) {
286        passes.push(&FuseAttentionBackwardAll);
287    }
288    // FuseResidualLN must run BEFORE FuseTransformerLayer: the layer-level
289    // pass matches `FAB → FusedResidualLN → FMBA(GeLU) → FMBA → FusedResidualLN`
290    // and needs the residual+LN ops already collapsed.
291    // `RLX_DISABLE_RESIDUAL_LN_FUSION=1` keeps `residual + LayerNorm` apart.
292    // The fused form *tees*: it writes both the normed output and the residual
293    // sum, so it is the only op in these graphs that writes a slot it does not
294    // own — worth being able to take out of the picture.
295    if rlx_ir::env::var("RLX_DISABLE_RESIDUAL_LN_FUSION").as_deref() != Some("1")
296        && supports_op(supported, OpKind::FusedResidualLN)
297    {
298        passes.push(&FuseResidualLN);
299    }
300    if supports_op(supported, OpKind::FusedResidualRmsNorm) {
301        passes.push(&FuseResidualRmsNorm);
302    }
303    // DiT adaLN-Zero / gated residual — after residual-LN so we don't
304    // compete for Add→LN patterns; gated residual is independent.
305    if supports_op(supported, OpKind::AdaLayerNorm) {
306        passes.push(&FuseAdaLayerNorm);
307    }
308    // `RLX_DISABLE_GATED_RESIDUAL_FUSION=1` keeps the gated residual decomposed.
309    // Ablation knob: it is one of the few fused ops with a distinct kernel per
310    // backend, so it is a candidate whenever CPU disagrees with every
311    // accelerator on an otherwise identical graph.
312    if supports_op(supported, OpKind::GatedResidual)
313        && rlx_ir::env::var("RLX_DISABLE_GATED_RESIDUAL_FUSION").as_deref() != Some("1")
314    {
315        passes.push(&FuseGatedResidual);
316    }
317    passes.push(&FuseRmsNormReshape);
318
319    // Layer-level fusion runs AFTER FuseResidualLN so it can match the
320    // post-fusion shape `FAB → FusedResidualLN → FMBA(GeLU) → FMBA →
321    // FusedResidualLN`. Opt-in via `RLX_ENABLE_FUSE_TRANSFORMER_LAYER`
322    // because backend perf wins are uneven: WGPU un-fuses with no
323    // dispatch reduction; MLX's lowering is correct (per the FAB fix
324    // above) but the MLX `compile()` already collapses sub-ops, so the
325    // extra IR-level fusion doesn't beat the natural pipeline. The pass
326    // exists for backends planning a monolithic transformer-layer kernel.
327    if rlx_ir::env::flag("RLX_ENABLE_FUSE_TRANSFORMER_LAYER")
328        && supports_op(supported, OpKind::FusedTransformerLayer)
329        && supports_op(supported, OpKind::FusedAttentionBlock)
330    {
331        passes.push(&FuseTransformerLayer);
332    }
333
334    // Option B (`RLX_NO_WEIGHT_CONCAT_FUSION`): skip the two fusions that build a
335    // combined weight via a runtime `concat` of the individual weight params
336    // (gate/up here, QKV below). For batch-1 decode that concat re-copies ~1.18GB
337    // of CONSTANT weights every token (≈2× the weight bandwidth) — reading the
338    // weights in place (separate matmuls) is faster. The activation-only
339    // `FuseSwiGLU` (silu·mul) still runs.
340    let no_weight_concat = rlx_ir::env::flag("RLX_NO_WEIGHT_CONCAT_FUSION");
341    if supports_op(supported, OpKind::FusedSwiGLU) && !no_weight_concat {
342        passes.push(&FuseSwiGLUDualMatmul);
343    }
344    // Fold `add(skip, matmul)` residuals into the matmul store. Metal-only
345    // today (one fewer dispatch on the launch-latency-bound decode); other
346    // backends never claim the op, so the pass never runs for them and the
347    // plain MatMul + Add stays.
348    if supports_op(supported, OpKind::FusedMatMulResidual) {
349        passes.push(&FuseMatMulResidual);
350    }
351    // Opt out: `RLX_NO_SHARED_INPUT_MATMUL=1` (debug / parity vs unfused AdaLN)
352    // or `RLX_NO_WEIGHT_CONCAT_FUSION=1` (Option B — read QKV weights in place).
353    if supports_op(supported, OpKind::MatMul)
354        && !rlx_ir::env::flag("RLX_NO_SHARED_INPUT_MATMUL")
355        && !no_weight_concat
356    {
357        passes.push(&FuseSharedInputMatMul);
358    }
359    if supports_op(supported, OpKind::FusedSwiGLU) {
360        passes.push(&FuseSwiGLU);
361    }
362
363    // Mark eligible element-wise chains only when the backend keeps regions.
364    // CPU/Metal unfuse immediately afterward — marking first duplicates the
365    // full graph for no benefit.
366    let keep_regions =
367        supports_op(supported, OpKind::ElementwiseRegion) && !opts.unfuse_elementwise_regions;
368    if keep_regions {
369        passes.push(&MarkElementwiseRegions);
370    }
371    if opts.fk_fusion {
372        passes.push(&MarkBatchSliceRegions);
373        passes.push(&MarkTransformRegions);
374        if opts.fuse_region_prologue {
375            passes.push(&FuseRegionPrologue);
376        }
377        if opts.fuse_batch_preprocess {
378            passes.push(&FuseBatchPreprocess);
379        }
380    }
381    let backend_native_fk = supports_op(supported, OpKind::TransformRegion)
382        && supports_op(supported, OpKind::BatchElementwiseRegion);
383    let keep_native_fk = opts.native_fk_regions && backend_native_fk;
384    if opts.decompose_fusion_regions || !keep_native_fk {
385        passes.push(&DecomposeFusionRegions);
386    }
387    if !keep_regions {
388        let unfuse = if matches!(target, FusionTarget::Cpu) {
389            &UnfuseElementwiseRegions::FOR_CPU
390        } else {
391            &UnfuseElementwiseRegions::FOR_GPU
392        };
393        passes.push(unfuse);
394    }
395
396    if supports_op(supported, OpKind::Fft) && supports_op(supported, OpKind::WelchPeaks) {
397        passes.push(&SelectPeaksOnlyOutputs);
398    }
399
400    finish_pipeline(passes)
401}
402
403/// FKL passes to run after [`MarkElementwiseRegions`] (e.g. `TpuExecutable::compile`).
404pub fn fk_passes_after_elementwise_regions(
405    supported: &[OpKind],
406    opts: FusionOptions,
407) -> Vec<&'static dyn Pass> {
408    let mut passes: Vec<&'static dyn Pass> = Vec::new();
409    if !opts.fk_fusion {
410        let backend_native_fk = supports_op(supported, OpKind::TransformRegion)
411            && supports_op(supported, OpKind::BatchElementwiseRegion);
412        let keep_native_fk = opts.native_fk_regions && backend_native_fk;
413        if opts.decompose_fusion_regions || !keep_native_fk {
414            passes.push(&DecomposeFusionRegions);
415        }
416        return finish_pipeline(passes);
417    }
418    passes.push(&MarkBatchSliceRegions);
419    passes.push(&MarkTransformRegions);
420    if opts.fuse_region_prologue {
421        passes.push(&FuseRegionPrologue);
422    }
423    if opts.fuse_batch_preprocess {
424        passes.push(&FuseBatchPreprocess);
425    }
426    let backend_native_fk = supports_op(supported, OpKind::TransformRegion)
427        && supports_op(supported, OpKind::BatchElementwiseRegion);
428    let keep_native_fk = opts.native_fk_regions && backend_native_fk;
429    if opts.decompose_fusion_regions || !keep_native_fk {
430        passes.push(&DecomposeFusionRegions);
431    }
432    finish_pipeline(passes)
433}
434
435/// IO gate decision for a rewrite on `target` (convenience for compile passes / model crates).
436pub fn should_fuse_with_target(
437    target: FusionTarget,
438    before: &crate::fusion_benefit::GraphIoProfile,
439    after: &crate::fusion_benefit::GraphIoProfile,
440) -> bool {
441    io_fusion_gate_for_target(target).should_fuse(before, after)
442}
443
444/// Phase 3 — IO-aware gate defaults for fusion rewrites on `target`.
445pub fn io_fusion_gate_for_target(target: FusionTarget) -> crate::fusion_benefit::IoFusionGate {
446    use crate::fusion_benefit::IoFusionGate;
447    match target {
448        FusionTarget::Metal | FusionTarget::Mlx => IoFusionGate {
449            dispatch_ns: 500.0,
450            roundtrip_ns: 5_000.0,
451            memory_bw: 200.0,
452            host_readback_bw: 200.0,
453            unified_memory: true,
454            host_thunk_penalty_ns: 2_000_000.0,
455            min_gain_ns: 1_000.0,
456        },
457        FusionTarget::Cuda | FusionTarget::Rocm => IoFusionGate {
458            dispatch_ns: 2_000.0,
459            roundtrip_ns: 20_000.0,
460            memory_bw: 800.0,
461            host_readback_bw: 50.0,
462            unified_memory: false,
463            host_thunk_penalty_ns: 15_000_000.0,
464            min_gain_ns: 5_000.0,
465        },
466        FusionTarget::Wgpu | FusionTarget::Tpu => IoFusionGate {
467            dispatch_ns: 3_000.0,
468            roundtrip_ns: 30_000.0,
469            memory_bw: 100.0,
470            host_readback_bw: 40.0,
471            unified_memory: false,
472            host_thunk_penalty_ns: 25_000_000.0,
473            min_gain_ns: 10_000.0,
474        },
475        FusionTarget::Cpu => IoFusionGate {
476            dispatch_ns: 50.0,
477            roundtrip_ns: 0.0,
478            memory_bw: 50.0,
479            host_readback_bw: 50.0,
480            unified_memory: true,
481            host_thunk_penalty_ns: 0.0,
482            min_gain_ns: 0.0,
483        },
484    }
485}
486
487/// Return the ordered fusion passes for `target`.
488pub fn fusion_passes(target: FusionTarget, opts: FusionOptions) -> Vec<&'static dyn Pass> {
489    let mut opts = opts;
490    // CPU thunks execute element-wise chains per-op. Metal's fused
491    // `ElementwiseRegion` MSL kernel mis-lowers long chains on deep
492    // transformer graphs (NaNs past ~14 blocks); keep FAB/RMSNorm fusions.
493    if !opts.keep_elementwise_regions
494        && matches!(target, FusionTarget::Cpu | FusionTarget::Metal)
495        && !opts.unfuse_elementwise_regions
496    {
497        opts.unfuse_elementwise_regions = true;
498    }
499    if opts.fusion_limits == FusionLimits::default() {
500        opts.fusion_limits = fusion_limits_for_target(target);
501    }
502    opts = opts.apply_native_fk_defaults(target);
503    fusion_passes_for_supported(supported_for_target(target), opts, target)
504}
505
506/// Per-target op claims used when a backend doesn't supply an explicit
507/// `supported_ops` slice. Must stay aligned with each backend's
508/// `*_SUPPORTED_OPS` in `rlx-runtime/src/backend.rs`.
509pub fn supported_for_target(target: FusionTarget) -> &'static [OpKind] {
510    use OpKind::*;
511    match target {
512        FusionTarget::Cpu => &[
513            MatMul,
514            DotGeneral,
515            ElementwiseRegion,
516            FusedSwiGLU,
517            FusedMatMulBiasAct,
518            FusedResidualLN,
519            FusedResidualRmsNorm,
520            FusedAttentionBlock,
521            AdaLayerNorm,
522            GatedResidual,
523        ],
524        FusionTarget::Metal => &[
525            MatMul,
526            DotGeneral,
527            ElementwiseRegion,
528            TransformRegion,
529            BatchElementwiseRegion,
530            FusedSwiGLU,
531            FusedMatMulBiasAct,
532            FusedResidualLN,
533            FusedResidualRmsNorm,
534            AdaLayerNorm,
535            GatedResidual,
536        ],
537        FusionTarget::Mlx => &[
538            MatMul,
539            DotGeneral,
540            ElementwiseRegion,
541            TransformRegion,
542            BatchElementwiseRegion,
543            FusedSwiGLU,
544            FusedMatMulBiasAct,
545            FusedResidualLN,
546            FusedResidualRmsNorm,
547            AdaLayerNorm,
548            GatedResidual,
549        ],
550        FusionTarget::Wgpu => &[
551            MatMul,
552            ElementwiseRegion,
553            TransformRegion,
554            BatchElementwiseRegion,
555            FusedSwiGLU,
556            FusedMatMulBiasAct,
557            FusedResidualLN,
558            FusedResidualRmsNorm,
559            AdaLayerNorm,
560            GatedResidual,
561            FusedAttentionBlock,
562            FusedTransformerLayer,
563        ],
564        // CUDA lowers `FusedConvBiasAct` to cuDNN's fused conv-bias-activation.
565        // ROCm/MIOpen has no fused-conv path yet, so it is NOT claimed there —
566        // its conv-bias-activation stays decomposed until a MIOpen lowering lands.
567        FusionTarget::Cuda => &[
568            MatMul,
569            DotGeneral,
570            ElementwiseRegion,
571            TransformRegion,
572            BatchElementwiseRegion,
573            FusedMatMulBiasAct,
574            FusedConvBiasAct,
575            FusedResidualLN,
576            FusedResidualRmsNorm,
577            AdaLayerNorm,
578            GatedResidual,
579        ],
580        FusionTarget::Rocm => &[
581            MatMul,
582            DotGeneral,
583            ElementwiseRegion,
584            TransformRegion,
585            BatchElementwiseRegion,
586            FusedMatMulBiasAct,
587            FusedResidualLN,
588            FusedResidualRmsNorm,
589            AdaLayerNorm,
590            GatedResidual,
591        ],
592        FusionTarget::Tpu => &[
593            MatMul,
594            ElementwiseRegion,
595            TransformRegion,
596            BatchElementwiseRegion,
597            FusedMatMulBiasAct,
598            FusedResidualLN,
599            AdaLayerNorm,
600            GatedResidual,
601        ],
602    }
603}
604
605fn finish_pipeline(mut passes: Vec<&'static dyn Pass>) -> Vec<&'static dyn Pass> {
606    passes.push(&DeadCodeElimination);
607    passes
608}
609
610/// Run the fusion pipeline for `target` on a MIR graph (IO-gated passes included).
611pub fn run_fusion_pipeline(
612    graph: Graph,
613    target: FusionTarget,
614    supported: &[OpKind],
615    opts: FusionOptions,
616) -> Graph {
617    let mut opts = opts.apply_native_fk_defaults(target);
618    if opts.fusion_limits == FusionLimits::default() {
619        opts.fusion_limits = fusion_limits_for_target(target);
620    }
621    let limits = opts.fusion_limits;
622    let passes = fusion_passes_for_supported(supported, opts, target);
623    with_fusion_target(target, || {
624        with_fusion_limits(limits, || run_passes(graph, &passes, false))
625    })
626}
627
628impl FusionOptions {
629    /// Canonical option set for a target, matching what each backend's
630    /// compile path uses: CPU unfuses element-wise regions (it executes them
631    /// as per-op thunks), Metal/MLX keep matmul/norm fusions but unfuse long
632    /// element-wise chains, wgpu unfuses for its modulus-broadcast kernel, and
633    /// the remaining GPU-class targets take the defaults (native FKL regions
634    /// are then enabled by [`apply_native_fk_defaults`](Self::apply_native_fk_defaults)).
635    /// This is the option set the one-call [`fuse`] / [`Fuse`] entries pick.
636    pub fn for_target(target: FusionTarget) -> Self {
637        match target {
638            FusionTarget::Cpu => Self::for_cpu(),
639            FusionTarget::Metal | FusionTarget::Mlx => Self::for_metal(),
640            FusionTarget::Wgpu => Self::for_wgpu(),
641            FusionTarget::Cuda | FusionTarget::Rocm | FusionTarget::Tpu => Self::default(),
642        }
643    }
644}
645
646/// One-call fusion: optimize `graph` for `target` with canonical defaults.
647///
648/// The ergonomic front door over [`run_fusion_pipeline`]. Callers no longer
649/// assemble the supported-op table or pick a per-target [`FusionOptions`] —
650/// both easy to get wrong. Exactly equivalent to:
651///
652/// ```ignore
653/// run_fusion_pipeline(
654///     graph, target,
655///     supported_for_target(target),
656///     FusionOptions::for_target(target),
657/// )
658/// ```
659///
660/// `Op::Input` / `Op::Param` leaves survive fusion, so recover handles into
661/// the rewritten graph by name with [`rlx_ir::Graph::node_id_by_name`];
662/// outputs stay positionally stable in `graph.outputs`. For custom limits,
663/// skip-fusion, or a change report, use the [`Fuse`] builder.
664pub fn fuse(graph: Graph, target: FusionTarget) -> Graph {
665    Fuse::new(target).run(graph)
666}
667
668/// Fluent fusion builder. `Fuse::new(target)` starts from the canonical
669/// [`FusionOptions::for_target`] defaults; setters override individual toggles
670/// before [`run`](Self::run) / [`run_with_report`](Self::run_with_report).
671///
672/// ```ignore
673/// let (fused, report) = Fuse::new(FusionTarget::Cpu)
674///     .limits(FusionLimits::UNBOUNDED)
675///     .run_with_report(graph);
676/// eprintln!("{}", report.summary_line());
677/// ```
678#[derive(Debug, Clone)]
679pub struct Fuse {
680    target: FusionTarget,
681    options: FusionOptions,
682}
683
684impl Fuse {
685    /// Start from the canonical defaults for `target`.
686    pub fn new(target: FusionTarget) -> Self {
687        Self {
688            target,
689            options: FusionOptions::for_target(target),
690        }
691    }
692
693    /// Replace the whole option set.
694    pub fn options(mut self, options: FusionOptions) -> Self {
695        self.options = options;
696        self
697    }
698
699    /// Override the element-wise fusion limits.
700    pub fn limits(mut self, limits: FusionLimits) -> Self {
701        self.options.fusion_limits = limits;
702        self
703    }
704
705    /// Disable all pattern fusion (lowering legalization only).
706    pub fn skip_fusion(mut self, skip: bool) -> Self {
707        self.options.skip_fusion = skip;
708        self
709    }
710
711    /// Layer compile-time env overrides (`RLX_*`) on top of the builder state.
712    pub fn merge_env(mut self) -> Self {
713        self.options = self.options.merge_env();
714        self
715    }
716
717    /// Run the pipeline, returning the optimized graph.
718    pub fn run(self, graph: Graph) -> Graph {
719        run_fusion_pipeline(
720            graph,
721            self.target,
722            supported_for_target(self.target),
723            self.options,
724        )
725    }
726
727    /// Run the pipeline and also return a [`FusionReport`] describing the
728    /// before→after change (op/region deltas, fusions left on the table) for
729    /// logging or assertions.
730    ///
731    /// [`FusionReport`]: rlx_fusion::fusion_report::FusionReport
732    pub fn run_with_report(self, graph: Graph) -> (Graph, rlx_fusion::fusion_report::FusionReport) {
733        let before = graph.clone();
734        let after = self.run(graph);
735        let report = rlx_fusion::fusion_report::FusionReport::analyze(&before, &after);
736        (after, report)
737    }
738}
739
740#[cfg(test)]
741mod tests {
742    use super::*;
743    use std::sync::Mutex;
744
745    /// Serializes every test that builds a pass list.
746    ///
747    /// `RLX_NO_NATIVE_FK_REGIONS` is a *process-global* env var that
748    /// `apply_native_fk_defaults` reads for **every** target, and
749    /// `tpu_native_fk_region_pass_policy` sets it. Any test calling
750    /// `fusion_passes` or `apply_native_fk_defaults` concurrently with that
751    /// one sees `native_fk_regions = false` and fails on an assertion that
752    /// has nothing to do with what it is testing. Hold this lock in any new
753    /// test here that builds a pass list.
754    static ENV_FK_TEST_LOCK: Mutex<()> = Mutex::new(());
755
756    /// Take [`ENV_FK_TEST_LOCK`], ignoring poisoning.
757    ///
758    /// The mutex orders env-var access; it guards no invariant that a panicking
759    /// test could leave broken. Unwrapping instead would turn one real failure
760    /// into a cascade of `PoisonError` panics in every sibling test.
761    fn env_fk_lock() -> std::sync::MutexGuard<'static, ()> {
762        ENV_FK_TEST_LOCK.lock().unwrap_or_else(|e| e.into_inner())
763    }
764
765    /// The one-call `fuse` / `Fuse` API folds matmul+bias+act on CPU and keeps
766    /// Input/Param leaves recoverable by name in the rewritten graph — the
767    /// ergonomic contract callers rely on instead of threading a remap table.
768    #[test]
769    fn fuse_one_call_collapses_matmul_bias_act_and_keeps_handles() {
770        use rlx_ir::op::{Activation, BinaryOp};
771        use rlx_ir::{DType, Graph, Shape};
772
773        let f = DType::F32;
774        let mut g = Graph::new("mm_bias_act");
775        let x = g.input("x", Shape::new(&[4, 8], f));
776        let w = g.param("w", Shape::new(&[8, 16], f));
777        let bias = g.param("bias", Shape::new(&[16], f));
778        let mm = g.matmul(x, w, Shape::new(&[4, 16], f));
779        let ba = g.binary(BinaryOp::Add, mm, bias, Shape::new(&[4, 16], f));
780        let act = g.activation(Activation::Relu, ba, Shape::new(&[4, 16], f));
781        g.set_outputs(vec![act]);
782        let before = g.len();
783
784        let (fused, report) = Fuse::new(FusionTarget::Cpu).run_with_report(g);
785
786        assert!(
787            fused.len() < before,
788            "matmul+bias+relu should fuse on CPU: {before} -> {} ({})",
789            fused.len(),
790            report.summary_line()
791        );
792        // Leaves survive a fusion rewrite; handles recoverable by name.
793        assert!(fused.input_id("x").is_some(), "input x lost after fusion");
794        assert!(fused.param_id("w").is_some(), "param w lost after fusion");
795        assert!(fused.param_id("bias").is_some(), "param bias lost");
796        assert!(fused.node_id_by_name("w").is_some());
797        // Outputs stay positionally stable.
798        assert_eq!(fused.outputs.len(), 1);
799    }
800
801    #[test]
802    fn cpu_pipeline_includes_attention_block() {
803        let _lock = env_fk_lock();
804        let passes = fusion_passes(FusionTarget::Cpu, FusionOptions::default());
805        assert_eq!(
806            passes.len(),
807            19,
808            "CPU default supported_ops omit Fft/WelchPeaks and mark_elementwise (unfuse-only backends skip mark); includes FuseAdaLayerNorm + FuseGatedResidual"
809        );
810        assert_eq!(passes[2].name(), "fuse_matmul_bias_act");
811        assert_eq!(passes[3].name(), "fuse_attention_block");
812        assert!(
813            passes.iter().any(|p| p.name() == "fuse_region_prologue"),
814            "default CPU pipeline should run FKL prologue fusion"
815        );
816        assert!(
817            !passes
818                .iter()
819                .any(|p| p.name() == "mark_elementwise_regions"),
820            "CPU unfuse backends should not mark elementwise regions before unfusing"
821        );
822        assert_eq!(passes.last().unwrap().name(), "dead_code_elimination");
823    }
824
825    #[test]
826    fn metal_skip_fusion_only_lowers_dot() {
827        let _lock = env_fk_lock();
828        let passes = fusion_passes(
829            FusionTarget::Metal,
830            FusionOptions {
831                skip_fusion: true,
832                ..FusionOptions::default()
833            },
834        );
835        assert_eq!(passes.len(), 2);
836        assert_eq!(passes[0].name(), "LowerControlFlow");
837        assert_eq!(passes[1].name(), "lower_dot_general");
838    }
839
840    #[test]
841    fn metal_supported_ops_omit_attention_block_fusion() {
842        let passes = fusion_passes_for_supported(
843            supported_for_target(FusionTarget::Metal),
844            FusionOptions::default(),
845            FusionTarget::Metal,
846        );
847        assert!(
848            !passes.iter().any(|p| p.name() == "fuse_attention_block"),
849            "Metal should not run FuseAttentionBlock"
850        );
851        assert!(
852            passes.iter().any(|p| p.name() == "fuse_matmul_bias_act"),
853            "Metal should fuse matmul+bias+act"
854        );
855    }
856
857    #[test]
858    fn cuda_supported_ops_fuse_matmul_bias_act() {
859        let passes = fusion_passes_for_supported(
860            supported_for_target(FusionTarget::Cuda),
861            FusionOptions::default(),
862            FusionTarget::Cuda,
863        );
864        assert!(
865            passes.iter().any(|p| p.name() == "fuse_matmul_bias_act"),
866            "CUDA should fuse matmul+bias+act when claimed"
867        );
868        assert!(
869            !passes.iter().any(|p| p.name() == "fuse_swiglu"),
870            "CUDA should not fuse SwiGLU"
871        );
872    }
873
874    #[test]
875    fn cpu_unfuses_elementwise_regions() {
876        let passes = fusion_passes_for_supported(
877            supported_for_target(FusionTarget::Cpu),
878            FusionOptions::for_cpu(),
879            FusionTarget::Cpu,
880        );
881        assert!(
882            passes
883                .iter()
884                .any(|p| p.name() == "unfuse_elementwise_regions")
885        );
886    }
887
888    #[test]
889    fn metal_unfuses_elementwise_regions_by_default() {
890        let _lock = env_fk_lock();
891        let passes = fusion_passes(FusionTarget::Metal, FusionOptions::default());
892        assert!(
893            passes
894                .iter()
895                .any(|p| p.name() == "unfuse_elementwise_regions")
896        );
897    }
898
899    #[test]
900    fn metal_default_unfuse_preserves_prologue_regions() {
901        let _lock = env_fk_lock();
902        let mut g = rlx_ir::Graph::new("t");
903        let shape_in = rlx_ir::Shape::new(&[1, 3, 8, 8], rlx_ir::DType::F32);
904        let shape_out = rlx_ir::Shape::new(&[1, 3, 16, 16], rlx_ir::DType::F32);
905        let x = g.input("x", shape_in);
906        let up = g.add_node(rlx_ir::Op::ResizeNearest2x, vec![x], shape_out.clone());
907        let r = g.add_node(
908            rlx_ir::Op::Activation(rlx_ir::op::Activation::Relu),
909            vec![up],
910            shape_out,
911        );
912        g.set_outputs(vec![r]);
913
914        let passes = fusion_passes(FusionTarget::Metal, FusionOptions::default());
915        let out = rlx_fusion::pass::run_passes(g, &passes, false);
916        assert!(out.nodes().iter().any(|n| {
917            matches!(
918                n.op,
919                rlx_ir::Op::ElementwiseRegion {
920                    prologue: rlx_ir::RegionPrologue::ResizeNearest2x,
921                    ..
922                }
923            )
924        }));
925    }
926
927    #[test]
928    fn fk_passes_after_elementwise_includes_batch_fusion() {
929        let _lock = env_fk_lock();
930        let opts = FusionOptions::default().apply_native_fk_defaults(FusionTarget::Tpu);
931        let passes =
932            fk_passes_after_elementwise_regions(supported_for_target(FusionTarget::Tpu), opts);
933        let names: Vec<_> = passes.iter().map(|p| p.name()).collect();
934        assert!(names.contains(&"mark_batch_slice_regions"));
935        assert!(names.contains(&"fuse_batch_preprocess"));
936        assert!(
937            !names.contains(&"decompose_fusion_regions"),
938            "TPU native FK defaults should keep batch/transform regions"
939        );
940    }
941
942    #[test]
943    fn tpu_native_fk_region_pass_policy() {
944        let _lock = env_fk_lock();
945        let default_passes = fusion_passes(FusionTarget::Tpu, FusionOptions::default());
946        assert!(
947            !default_passes
948                .iter()
949                .any(|p| p.name() == "decompose_fusion_regions"),
950            "default TPU pipeline keeps batch/transform regions via native_fk_defaults"
951        );
952
953        rlx_ir::env::set("RLX_NO_NATIVE_FK_REGIONS", "1");
954        let opt_out = fusion_passes(FusionTarget::Tpu, FusionOptions::default());
955        rlx_ir::env::unset("RLX_NO_NATIVE_FK_REGIONS");
956        assert!(
957            opt_out
958                .iter()
959                .any(|p| p.name() == "decompose_fusion_regions"),
960            "RLX_NO_NATIVE_FK_REGIONS should force decompose on TPU"
961        );
962    }
963
964    #[test]
965    fn native_fk_regions_skips_decompose_on_tpu() {
966        let _lock = env_fk_lock();
967        let passes = fusion_passes(
968            FusionTarget::Tpu,
969            FusionOptions {
970                native_fk_regions: true,
971                decompose_fusion_regions: false,
972                unfuse_elementwise_regions: false,
973                ..FusionOptions::default()
974            },
975        );
976        assert!(
977            !passes
978                .iter()
979                .any(|p| p.name() == "decompose_fusion_regions"),
980            "native_fk_regions should skip decompose on TPU when batch/transform are supported"
981        );
982    }
983
984    #[test]
985    fn native_fk_regions_skips_decompose_on_metal() {
986        let _lock = env_fk_lock();
987        let passes = fusion_passes(
988            FusionTarget::Metal,
989            FusionOptions {
990                native_fk_regions: true,
991                decompose_fusion_regions: false,
992                unfuse_elementwise_regions: false,
993                ..FusionOptions::default()
994            },
995        );
996        assert!(
997            !passes
998                .iter()
999                .any(|p| p.name() == "decompose_fusion_regions"),
1000            "native_fk_regions should skip decompose when backend claims batch/transform ops"
1001        );
1002    }
1003
1004    #[test]
1005    fn metal_keeps_elementwise_regions_when_requested() {
1006        let _lock = env_fk_lock();
1007        let passes = fusion_passes(
1008            FusionTarget::Metal,
1009            FusionOptions {
1010                keep_elementwise_regions: true,
1011                unfuse_elementwise_regions: false,
1012                ..FusionOptions::default()
1013            },
1014        );
1015        assert!(
1016            !passes
1017                .iter()
1018                .any(|p| p.name() == "unfuse_elementwise_regions"),
1019            "keep_elementwise_regions should skip unfuse pass"
1020        );
1021        assert!(
1022            passes.iter().any(|p| p.name() == "fuse_region_prologue"),
1023            "FKL prologue fusion should still run"
1024        );
1025    }
1026
1027    #[test]
1028    fn metal_audio_ops_pipeline_includes_peaks_output_gate() {
1029        let mut supported = supported_for_target(FusionTarget::Metal).to_vec();
1030        supported.push(OpKind::Fft);
1031        supported.push(OpKind::WelchPeaks);
1032        let passes =
1033            fusion_passes_for_supported(&supported, FusionOptions::default(), FusionTarget::Metal);
1034        assert!(
1035            passes
1036                .iter()
1037                .any(|p| p.name() == "select_peaks_only_outputs"),
1038            "Metal + Fft/WelchPeaks should run IO peaks-only output gate"
1039        );
1040    }
1041
1042    #[test]
1043    fn should_fuse_with_target_matches_gate() {
1044        use crate::fusion_benefit::GraphIoProfile;
1045        let dense = GraphIoProfile {
1046            kernel_launches: 3,
1047            sync_points: 0,
1048            host_output_bytes: 33_554_432,
1049            device_traffic_bytes: 184_549_376,
1050        };
1051        let fused = GraphIoProfile {
1052            kernel_launches: 4,
1053            sync_points: 1,
1054            host_output_bytes: 1_048_576,
1055            device_traffic_bytes: 219_152_384,
1056        };
1057        assert!(should_fuse_with_target(FusionTarget::Metal, &dense, &fused));
1058        assert!(!should_fuse_with_target(FusionTarget::Wgpu, &dense, &fused));
1059    }
1060}
rlx_compile/fusion_pipeline.rs

rlx_compile/
fusion_pipeline.rs
