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
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
// RLX — versatile ML compiler + runtime.
// Copyright (C) 2026 Eugene Hauptmann, Nataliya Kosmyna.
// SPDX-License-Identifier: MIT OR Apache-2.0
//! Activation-scale calibration for post-training INT8 quantization.
//!
//! The runtime-side counterpart to `rlx_opt::quant_insert`. Compiles a
//! forward graph with calibration "tap" nodes wired in as outputs;
//! the caller runs one batch at a time (filling the input slots
//! between calls) and the [`Calibrator`] accumulates max-abs per tap.
//! At the end, `scales()` returns `max_abs / 127.0` per tap (clamped
//! up to `1e-6` to avoid division-by-zero on a flat tensor) — the
//! per-tensor scale that maps the calibration range into i8.
//!
//! Why max-abs and not e.g. the 99th percentile? Max-abs matches what
//! the cortexm Python trainer used to do (and what the Rust trainer
//! that replaced it does). It's symmetric (zero zero-point), maps
//! `[-max, +max] → [-127, 127]`, and gives the worst-case-correct
//! quantization for activations whose distributions are roughly
//! zero-centered. Percentile-based / KL-divergence calibration are
//! follow-ups for later.
use crate::arena::Arena;
use crate::thunk::{ThunkSchedule, compile_thunks, execute_thunks};
use rlx_ir::{Graph, NodeId};
/// Compiled calibration harness. The graph is owned by the caller —
/// we hold a reference and the compiled artifacts (arena + schedule).
/// The caller writes inputs and parameters into `arena_mut()` between
/// batches.
pub struct Calibrator<'g> {
graph: &'g Graph,
arena: Arena,
sched: ThunkSchedule,
/// `(tap_node_id, num_elements)` pairs — cached so each `step`
/// doesn't re-walk the graph for shape info.
taps: Vec<(NodeId, usize)>,
/// Running max-abs per tap. Index aligns with the `taps` order
/// the caller passed to `new`.
max_abs: Vec<f32>,
}
impl<'g> Calibrator<'g> {
/// Build a calibrator over `graph` that records max-abs at each
/// `tap` after every `step()`. The graph must already have those
/// taps in its `outputs` list (so the memory planner keeps their
/// arena slots alive to end-of-execution); this constructor
/// asserts the precondition.
pub fn new(graph: &'g Graph, taps: Vec<NodeId>) -> Self {
for &t in &taps {
assert!(
graph.outputs.contains(&t),
"Calibrator: tap {t} must be in graph.outputs so its slot \
survives the run; add it via graph.set_outputs(…)"
);
}
let plan = rlx_opt::memory::plan_memory(graph);
let arena = Arena::from_plan(plan);
let sched = compile_thunks(graph, &arena);
let n = taps.len();
let taps_with_len: Vec<(NodeId, usize)> = taps
.into_iter()
.map(|t| {
let len = graph.node(t).shape.num_elements().unwrap_or(0);
(t, len)
})
.collect();
Self {
graph,
arena,
sched,
taps: taps_with_len,
max_abs: vec![0.0; n],
}
}
/// Mutable arena access — for writing inputs/params before each
/// `step()` and (typically once at startup) for filling
/// `Op::Constant` data via `rlx_runtime`'s loader.
pub fn arena_mut(&mut self) -> &mut Arena {
&mut self.arena
}
/// Read-only arena view — for reading the tap values manually if
/// the caller wants something fancier than max-abs.
pub fn arena(&self) -> &Arena {
&self.arena
}
/// Run one forward batch, then update each tap's running max-abs.
pub fn step(&mut self) {
execute_thunks(&self.sched, self.arena.raw_buf_mut());
for ((tap, len), max) in self.taps.iter().zip(self.max_abs.iter_mut()) {
let off = self.arena.byte_offset(*tap);
unsafe {
let p = self.arena.raw_buf().as_ptr().add(off) as *const f32;
for i in 0..*len {
let v = (*p.add(i)).abs();
if v > *max {
*max = v;
}
}
}
}
}
/// Per-tap max-abs accumulated so far (in input order).
pub fn max_abs(&self) -> &[f32] {
&self.max_abs
}
/// Per-tap scale = `max_abs / 127.0`, clamped up to `1e-6`.
/// Use directly as the `scale` for `Op::Quantize` / `Op::Dequantize`
/// or `rlx_opt::CalibrationEntry::per_tensor`.
pub fn scales(&self) -> Vec<f32> {
self.max_abs.iter().map(|m| (m / 127.0).max(1e-6)).collect()
}
/// Borrow the inner graph (for the caller to re-look-up NodeIds
/// after compilation).
pub fn graph(&self) -> &Graph {
self.graph
}
}
#[cfg(test)]
mod tests {
use super::*;
use rlx_ir::op::*;
use rlx_ir::*;
/// One-tap calibration over a trivial graph: tap = `x` itself.
/// Hand-pack a couple of batches with known max-abs values and
/// verify `scales()` reflects them.
#[test]
fn calibrator_tracks_max_abs_across_batches() {
let f = DType::F32;
let mut g = Graph::new("calib_demo");
let x = g.input("x", Shape::new(&[4], f));
// Identity-ish: the tap *is* the input. Adding a Relu so the
// graph is non-trivial.
let y = g.activation(Activation::Relu, x, Shape::new(&[4], f));
g.set_outputs(vec![x, y]); // tap on `x` and `y`
let mut cal = Calibrator::new(&g, vec![x, y]);
// Batch 1: max-abs of x = 3.0; max-abs of y (Relu) = 3.0.
write_into(cal.arena_mut(), x, &[-3.0, 1.0, -2.0, 0.5]);
cal.step();
// Batch 2: x's max-abs grows to 7.0; y's stays since negatives
// get zeroed by Relu.
write_into(cal.arena_mut(), x, &[-7.0, 0.0, -7.0, -2.0]);
cal.step();
// Batch 3: both grow.
write_into(cal.arena_mut(), x, &[10.0, 0.0, 0.0, 5.0]);
cal.step();
let mx = cal.max_abs();
assert!((mx[0] - 10.0).abs() < 1e-6, "x max_abs: {}", mx[0]);
assert!((mx[1] - 10.0).abs() < 1e-6, "y max_abs: {}", mx[1]);
let s = cal.scales();
assert!((s[0] - 10.0 / 127.0).abs() < 1e-6);
assert!((s[1] - 10.0 / 127.0).abs() < 1e-6);
}
fn write_into(arena: &mut Arena, id: NodeId, data: &[f32]) {
let off = arena.byte_offset(id);
let buf = arena.raw_buf_mut();
unsafe {
let p = buf.as_mut_ptr().add(off) as *mut f32;
for (i, &v) in data.iter().enumerate() {
*p.add(i) = v;
}
}
}
}