rlx-compile 0.2.5

HIR → MIR → LIR compile pipeline for RLX
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
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120

// 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.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
// GNU General Public License for more details.
//
// You should have received a copy of the GNU General Public License
// along with this program. If not, see <https://www.gnu.org/licenses/>.

//! Replace selected `Op::Param` nodes with `Op::Constant` before compile-time opts.
//!
//! Deploy graphs (e.g. pruned ternary FFT) often fix gate masks and zero twiddles
//! at specialization time while still building the graph with `Graph::param`.
//! Baking those values here lets constant folding and DCE remove dead paths.

use rlx_fusion::pass::Pass;
use rlx_ir::{Graph, NodeId, Op};
use std::collections::HashMap;

fn encode_f32(data: &[f32]) -> Vec<u8> {
    let mut bytes = Vec::with_capacity(data.len() * 4);
    for &v in data {
        bytes.extend_from_slice(&v.to_le_bytes());
    }
    bytes
}

/// Substitute listed params with constants. Unlisted params are unchanged.
pub fn specialize_params(graph: &Graph, bindings: &HashMap<String, Vec<f32>>) -> Graph {
    if bindings.is_empty() {
        return graph.clone();
    }
    let mut out = Graph::new(graph.name.clone());
    let mut id_map: HashMap<NodeId, NodeId> = HashMap::new();

    for node in graph.nodes() {
        let new_id = match &node.op {
            Op::Param { name } => {
                if let Some(values) = bindings.get(name) {
                    let expected = node.shape.num_elements().unwrap_or(values.len());
                    assert_eq!(
                        values.len(),
                        expected,
                        "param '{name}' binding len {} != shape elements {expected}",
                        values.len()
                    );
                    out.add_node(
                        Op::Constant {
                            data: encode_f32(values),
                        },
                        vec![],
                        node.shape.clone(),
                    )
                } else {
                    out.add_node(node.op.clone(), vec![], node.shape.clone())
                }
            }
            _ => {
                let new_inputs: Vec<NodeId> = node.inputs.iter().map(|i| id_map[i]).collect();
                out.add_node(node.op.clone(), new_inputs, node.shape.clone())
            }
        };
        id_map.insert(node.id, new_id);
    }

    let new_outputs: Vec<NodeId> = graph.outputs.iter().map(|o| id_map[o]).collect();
    out.set_outputs(new_outputs);
    out
}

/// Pass wrapper for the fusion pipeline / runtime preprocess hook.
pub struct SpecializeParams {
    pub bindings: HashMap<String, Vec<f32>>,
}

impl SpecializeParams {
    pub fn new(bindings: HashMap<String, Vec<f32>>) -> Self {
        Self { bindings }
    }
}

impl Pass for SpecializeParams {
    fn name(&self) -> &str {
        "specialize_params"
    }

    fn run(&self, graph: Graph) -> Graph {
        specialize_params(&graph, &self.bindings)
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use rlx_ir::Shape;
    use rlx_ir::op::BinaryOp;
    use rlx_ir::*;

    #[test]
    fn replaces_bound_param_with_constant() {
        let s = Shape::new(&[2], DType::F32);
        let mut g = Graph::new("t");
        let x = g.input("x", s.clone());
        let w = g.param("w", s.clone());
        let y = g.binary(BinaryOp::Mul, x, w, s.clone());
        g.set_outputs(vec![y]);

        let mut bindings = HashMap::new();
        bindings.insert("w".into(), vec![0.0, 1.0]);
        let out = specialize_params(&g, &bindings);
        let w_node = out.node(out.nodes()[1].id);
        assert!(matches!(w_node.op, Op::Constant { .. }));
    }
}