rlx-fusion 0.2.6

Fusion passes and fused-op decomposition for RLX MIR
Documentation 
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

// RLX — versatile ML compiler + runtime.
// Copyright (C) 2026 Eugene Hauptmann, Nataliya Kosmyna.
//
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU General Public License as published by
// the Free Software Foundation, version 3.

//! MIR fusion passes and fused-op decomposition.
//!
//! Pattern-matching fusion (`FuseMatMulBiasAct`, `FuseSwiGLU`, …) and
//! the inverse [`unfuse_fused_for_autodiff`] rewrite used before autodiff.

pub mod control_flow;
pub mod fk_fusion;
pub mod fk_graphs;
pub mod fusion;
pub mod fusion_fragment;
pub mod fusion_report;
pub mod graph_rewrite;
pub mod limits;
pub mod lower_backward_ops;
pub mod lower_dot_general;
pub mod lower_logical_kernels;
pub mod lower_loss_ops;
pub mod lower_reduce_axes;
pub mod lower_vae_ops;
pub mod pass;
pub mod unfuse;

pub use control_flow::{
    LowerControlFlow, inline_if, inline_subgraph_into, inline_subgraph_into_outputs, unroll_while,
};
pub use fk_fusion::{
    DecomposeFusionRegions, FuseBatchPreprocess, FuseRegionPrologue, MarkBatchSliceRegions,
    MarkTransformRegions,
};
pub use fk_graphs::{
    batch_narrow_relu_primitive_graph, batch_narrow_relu_regions_graph, nchw, resize_relu_graph,
    resize_relu_region_graph,
};
pub use fusion::{
    FuseAttentionBlock, FuseMatMulBiasAct, FuseResidualLN, FuseResidualRmsNorm, FuseRmsNormReshape,
    FuseSharedInputMatMul, FuseSwiGLU, FuseSwiGLUDualMatmul, FuseTransformerLayer,
    MarkElementwiseRegions, UnfuseElementwiseRegions, clip_elementwise_regions,
};
pub use fusion_fragment::{
    FusionFragment, FusionRole, fusion_fragments, is_registered_transform_op,
    prologue_for_transform_op, register_fusion_fragment, transform_chain_eligible,
};
pub use fusion_report::{FusionReport, MissReason, MissedFusion};
pub use limits::{FusionLimits, active_fusion_limits, with_fusion_limits};
pub use lower_backward_ops::LowerBackwardOps;
pub use lower_dot_general::LowerDotGeneral;
pub use lower_logical_kernels::lower_logical_kernels;
pub use lower_loss_ops::LowerSoftmaxCrossEntropy;
pub use lower_reduce_axes::LowerNonLastAxisReduce;
pub use lower_vae_ops::{LowerBatchNormInference, LowerGroupNorm, LowerResizeNearest2x};
pub use pass::{Pass, run_passes};
pub use unfuse::unfuse_fused_for_autodiff;