rlx_runtime/precompile.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 graph cleanup before the fusion / backend pipeline.
17
18use crate::CompileOptions;
19use rlx_ir::Graph;
20use rlx_opt::pass::Pass as _;
21
22/// Param specialization, algebraic simplify, DCE, and constant folding.
23pub fn precompile_cleanup(graph: Graph, options: &CompileOptions) -> Graph {
24 let mut graph = graph;
25 if options
26 .param_bindings
27 .as_ref()
28 .is_some_and(|b| !b.is_empty())
29 {
30 let bindings = options.param_bindings.as_ref().unwrap();
31 graph = rlx_opt::specialize_params(&graph, bindings);
32 }
33 post_specialize_cleanup(graph, options)
34}
35
36/// DCE / fold after fusion — skips param specialization (already applied pre-fusion).
37pub fn post_fusion_cleanup(graph: Graph, options: &CompileOptions) -> Graph {
38 post_specialize_cleanup(graph, options)
39}
40
41fn post_specialize_cleanup(graph: Graph, options: &CompileOptions) -> Graph {
42 let mut graph = rlx_opt::AlgebraicSimplify.run(graph);
43 if options.dce {
44 graph = rlx_opt::DeadCodeElimination.run(graph);
45 }
46 if options.constant_folding {
47 graph = rlx_opt::ConstantFolding.run(graph);
48 }
49 graph
50}