rlx_ir/ops/
elementwise.rs

1// RLX — versatile ML compiler + runtime.
2// Copyright (C) 2026 Eugene Hauptmann, Nataliya Kosmyna.
3// SPDX-License-Identifier: MIT OR Apache-2.0
4
5//! Element-wise builders: binary ops, activations (plan #53).
6
7use crate::op::{Activation, BinaryOp};
8use crate::{Graph, NodeId, Op, Shape};
9
10impl Graph {
11    /// Binary element-wise operation.
12    pub fn binary(&mut self, op: BinaryOp, lhs: NodeId, rhs: NodeId, out_shape: Shape) -> NodeId {
13        self.push(Op::Binary(op), vec![lhs, rhs], out_shape, None)
14    }
15
16    /// Unary activation.
17    pub fn activation(&mut self, act: Activation, input: NodeId, shape: Shape) -> NodeId {
18        self.push(Op::Activation(act), vec![input], shape, None)
19    }
20
21    /// Per-tensor INT8 quantization. Output dtype = `I8`, same shape
22    /// otherwise. `scale` and `zero_point` apply uniformly to every
23    /// element. Use `quantize_per_channel` when weights deserve
24    /// per-channel scales (the standard PTQ improvement).
25    pub fn quantize(&mut self, x: NodeId, scale: f32, zero_point: i32) -> NodeId {
26        let shape = self.shape(x).clone().with_dtype(crate::DType::I8);
27        self.push(
28            Op::Quantize {
29                axis: None,
30                scales: vec![scale],
31                zero_points: vec![zero_point],
32            },
33            vec![x],
34            shape,
35            None,
36        )
37    }
38
39    /// Per-channel INT8 quantization. `scales` and `zero_points` must
40    /// each have length `input.dim(axis)`; the kernel picks the i-th
41    /// pair when quantizing the i-th slice along `axis`. The most
42    /// common usage is `axis = 0` for a `[C_out, C_in, kH, kW]`
43    /// conv weight (one scale per output channel).
44    pub fn quantize_per_channel(
45        &mut self,
46        x: NodeId,
47        axis: usize,
48        scales: Vec<f32>,
49        zero_points: Vec<i32>,
50    ) -> NodeId {
51        debug_assert_eq!(scales.len(), zero_points.len());
52        let shape = self.shape(x).clone().with_dtype(crate::DType::I8);
53        debug_assert_eq!(
54            shape.dim(axis),
55            crate::shape::Dim::Static(scales.len()),
56            "quantize_per_channel: scales.len() must match input.dim(axis)"
57        );
58        self.push(
59            Op::Quantize {
60                axis: Some(axis),
61                scales,
62                zero_points,
63            },
64            vec![x],
65            shape,
66            None,
67        )
68    }
69
70    /// Per-tensor INT8 dequantization (inverse of `quantize`). Output
71    /// dtype is f32.
72    pub fn dequantize(&mut self, x: NodeId, scale: f32, zero_point: i32) -> NodeId {
73        let shape = self.shape(x).clone().with_dtype(crate::DType::F32);
74        self.push(
75            Op::Dequantize {
76                axis: None,
77                scales: vec![scale],
78                zero_points: vec![zero_point],
79            },
80            vec![x],
81            shape,
82            None,
83        )
84    }
85
86    /// Per-channel INT8 dequantization (inverse of
87    /// `quantize_per_channel`).
88    pub fn dequantize_per_channel(
89        &mut self,
90        x: NodeId,
91        axis: usize,
92        scales: Vec<f32>,
93        zero_points: Vec<i32>,
94    ) -> NodeId {
95        debug_assert_eq!(scales.len(), zero_points.len());
96        let shape = self.shape(x).clone().with_dtype(crate::DType::F32);
97        debug_assert_eq!(
98            shape.dim(axis),
99            crate::shape::Dim::Static(scales.len()),
100            "dequantize_per_channel: scales.len() must match input.dim(axis)"
101        );
102        self.push(
103            Op::Dequantize {
104                axis: Some(axis),
105                scales,
106                zero_points,
107            },
108            vec![x],
109            shape,
110            None,
111        )
112    }
113}
rlx_ir/ops/elementwise.rs

rlx_ir/ops/
elementwise.rs
