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
// RLX — versatile ML compiler + runtime.
// Copyright (C) 2026 Eugene Hauptmann, Nataliya Kosmyna.
// SPDX-License-Identifier: MIT OR Apache-2.0
//! 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,
)
}
}