rlx_fusion/
lib.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//! MIR fusion passes and fused-op decomposition.
9//!
10//! Pattern-matching fusion (`FuseMatMulBiasAct`, `FuseSwiGLU`, …) and
11//! the inverse [`unfuse_fused_for_autodiff`] rewrite used before autodiff.
12
13pub mod control_flow;
14pub mod fk_fusion;
15pub mod fk_graphs;
16pub mod fusion;
17pub mod fusion_fragment;
18pub mod fusion_report;
19pub mod graph_rewrite;
20pub mod limits;
21pub mod lower_backward_ops;
22pub mod lower_dot_general;
23pub mod lower_logical_kernels;
24pub mod lower_loss_ops;
25pub mod lower_reduce_axes;
26pub mod lower_vae_ops;
27pub mod pass;
28pub mod unfuse;
29
30pub use control_flow::{
31    LowerControlFlow, inline_if, inline_subgraph_into, inline_subgraph_into_outputs, unroll_while,
32};
33pub use fk_fusion::{
34    DecomposeFusionRegions, FuseBatchPreprocess, FuseRegionPrologue, MarkBatchSliceRegions,
35    MarkTransformRegions,
36};
37pub use fk_graphs::{
38    batch_narrow_relu_primitive_graph, batch_narrow_relu_regions_graph, nchw, resize_relu_graph,
39    resize_relu_region_graph,
40};
41pub use fusion::{
42    FuseAttentionBlock, FuseMatMulBiasAct, FuseResidualLN, FuseResidualRmsNorm, FuseRmsNormReshape,
43    FuseSharedInputMatMul, FuseSwiGLU, FuseSwiGLUDualMatmul, FuseTransformerLayer,
44    MarkElementwiseRegions, UnfuseElementwiseRegions, clip_elementwise_regions,
45};
46pub use fusion_fragment::{
47    FusionFragment, FusionRole, fusion_fragments, is_registered_transform_op,
48    prologue_for_transform_op, register_fusion_fragment, transform_chain_eligible,
49};
50pub use fusion_report::{FusionReport, MissReason, MissedFusion};
51pub use limits::{FusionLimits, active_fusion_limits, with_fusion_limits};
52pub use lower_backward_ops::LowerBackwardOps;
53pub use lower_dot_general::LowerDotGeneral;
54pub use lower_logical_kernels::lower_logical_kernels;
55pub use lower_loss_ops::LowerSoftmaxCrossEntropy;
56pub use lower_reduce_axes::LowerNonLastAxisReduce;
57pub use lower_vae_ops::{LowerBatchNormInference, LowerGroupNorm, LowerResizeNearest2x};
58pub use pass::{Pass, run_passes};
59pub use unfuse::unfuse_fused_for_autodiff;
rlx_fusion/lib.rs

rlx_fusion/
lib.rs
