rlx-ir 0.2.1

Tensor IR for the RLX ML compiler — standalone, serializable, optimizable
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
121
122
123
124

// 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/>.

//! Element-wise builders: binary ops, activations (plan #53).

use crate::op::{Activation, BinaryOp};
use crate::{Graph, NodeId, Op, Shape};

impl Graph {
    /// Binary element-wise operation.
    pub fn binary(&mut self, op: BinaryOp, lhs: NodeId, rhs: NodeId, out_shape: Shape) -> NodeId {
        self.push(Op::Binary(op), vec![lhs, rhs], out_shape, None)
    }

    /// Unary activation.
    pub fn activation(&mut self, act: Activation, input: NodeId, shape: Shape) -> NodeId {
        self.push(Op::Activation(act), vec![input], shape, None)
    }

    /// Per-tensor INT8 quantization. Output dtype = `I8`, same shape
    /// otherwise. `scale` and `zero_point` apply uniformly to every
    /// element. Use `quantize_per_channel` when weights deserve
    /// per-channel scales (the standard PTQ improvement).
    pub fn quantize(&mut self, x: NodeId, scale: f32, zero_point: i32) -> NodeId {
        let shape = self.shape(x).clone().with_dtype(crate::DType::I8);
        self.push(
            Op::Quantize {
                axis: None,
                scales: vec![scale],
                zero_points: vec![zero_point],
            },
            vec![x],
            shape,
            None,
        )
    }

    /// Per-channel INT8 quantization. `scales` and `zero_points` must
    /// each have length `input.dim(axis)`; the kernel picks the i-th
    /// pair when quantizing the i-th slice along `axis`. The most
    /// common usage is `axis = 0` for a `[C_out, C_in, kH, kW]`
    /// conv weight (one scale per output channel).
    pub fn quantize_per_channel(
        &mut self,
        x: NodeId,
        axis: usize,
        scales: Vec<f32>,
        zero_points: Vec<i32>,
    ) -> NodeId {
        debug_assert_eq!(scales.len(), zero_points.len());
        let shape = self.shape(x).clone().with_dtype(crate::DType::I8);
        debug_assert_eq!(
            shape.dim(axis),
            crate::shape::Dim::Static(scales.len()),
            "quantize_per_channel: scales.len() must match input.dim(axis)"
        );
        self.push(
            Op::Quantize {
                axis: Some(axis),
                scales,
                zero_points,
            },
            vec![x],
            shape,
            None,
        )
    }

    /// Per-tensor INT8 dequantization (inverse of `quantize`). Output
    /// dtype is f32.
    pub fn dequantize(&mut self, x: NodeId, scale: f32, zero_point: i32) -> NodeId {
        let shape = self.shape(x).clone().with_dtype(crate::DType::F32);
        self.push(
            Op::Dequantize {
                axis: None,
                scales: vec![scale],
                zero_points: vec![zero_point],
            },
            vec![x],
            shape,
            None,
        )
    }

    /// Per-channel INT8 dequantization (inverse of
    /// `quantize_per_channel`).
    pub fn dequantize_per_channel(
        &mut self,
        x: NodeId,
        axis: usize,
        scales: Vec<f32>,
        zero_points: Vec<i32>,
    ) -> NodeId {
        debug_assert_eq!(scales.len(), zero_points.len());
        let shape = self.shape(x).clone().with_dtype(crate::DType::F32);
        debug_assert_eq!(
            shape.dim(axis),
            crate::shape::Dim::Static(scales.len()),
            "dequantize_per_channel: scales.len() must match input.dim(axis)"
        );
        self.push(
            Op::Dequantize {
                axis: Some(axis),
                scales,
                zero_points,
            },
            vec![x],
            shape,
            None,
        )
    }
}