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    // Decompose registered custom ops that opt into a `lower` rule into
25    // primitives BEFORE fusion / legalize / kernel dispatch. Every backend that
26    // whitelists `OpKind::Custom` (all of them) would otherwise pass legalize and
27    // then hard-fail at kernel-dispatch time for a custom op it has no kernel
28    // for; lowering here lets such an op run on any backend with no kernel. A
29    // no-op unless the graph carries a custom op, and idempotent with the same
30    // pass inside `rewrite_for_backend`.
31    let mut graph = rlx_opt::lower_custom_ops(graph);
32    if options
33        .param_bindings
34        .as_ref()
35        .is_some_and(|b| !b.is_empty())
36    {
37        let bindings = options.param_bindings.as_ref().unwrap();
38        graph = rlx_opt::specialize_params(&graph, bindings);
39    }
40    post_specialize_cleanup(graph, options)
41}
42
43/// DCE / fold after fusion — skips param specialization (already applied pre-fusion).
44pub fn post_fusion_cleanup(graph: Graph, options: &CompileOptions) -> Graph {
45    post_specialize_cleanup(graph, options)
46}
47
48fn post_specialize_cleanup(graph: Graph, options: &CompileOptions) -> Graph {
49    let mut graph = rlx_opt::AlgebraicSimplify.run(graph);
50    if options.dce {
51        graph = rlx_opt::DeadCodeElimination.run(graph);
52    }
53    if options.constant_folding {
54        graph = rlx_opt::ConstantFolding.run(graph);
55    }
56    graph
57}
rlx_runtime/precompile.rs

rlx_runtime/
precompile.rs
