feat: saturation for add,sub,sum

Author: Forpee

atlas-onnx-tracer/src/model/mod.rs

@@ -261,6 +261,11 @@ impl Model {
             .nodes
             .values()
             .map(|node| match &node.operator {
+                // Saturating Add/Sub commit a one-hot clamp read-address
+                // decomposition. Its chunks are `k_chunk`-wide just like the
+                // activation lookups, so the largest committed poly is
+                // `k_chunk * T` (the 64-bit address only changes the chunk
+                // *count*, not the per-chunk size).
                 Operator::Tanh(_)
                 | Operator::Cos(_)
                 | Operator::Div(_)
@@ -269,6 +274,9 @@ impl Model {
                 | Operator::Rsqrt(_)
                 | Operator::Sigmoid(_)
                 | Operator::Sin(_)
+                | Operator::Add(_)
+                | Operator::Sub(_)
+                | Operator::Sum(_)
                 | Operator::SoftmaxLastAxis(_) => {
                     LOG_K_CHUNK + log_2(node.pow2_padded_num_output_elements())
                 }

atlas-onnx-tracer/src/ops/add.rs

@@ -1,10 +1,10 @@
 use super::{Add, Op};
-use crate::tensor::{self, Tensor};
+use crate::tensor::Tensor;
 
 impl Op for Add {
     #[tracing::instrument(name = "Add::f", skip_all)]
     fn f(&self, inputs: Vec<&Tensor<i32>>) -> Tensor<i32> {
-        tensor::ops::add(&inputs).unwrap()
+        super::sat_binop(inputs, "Add", |a, b| a + b)
     }
 
     fn requires_shape_equality(&self) -> bool {

atlas-onnx-tracer/src/ops/mod.rs

@@ -180,6 +180,90 @@ pub trait Op {
     }
 }
 
+/// Broadcast-expand `lhs` and `rhs`, apply `combine` element-wise in `i64`,
+/// and return the unclamped intermediate as `Tensor<i64>`.
+///
+/// This is the re-executable kernel the proof system calls to recover the
+/// pre-saturation intermediate without storing it in the trace.
+pub(super) fn sat_accumulate_pair(
+    lhs: &crate::tensor::Tensor<i32>,
+    rhs: &crate::tensor::Tensor<i32>,
+    op_name: &str,
+    combine: impl Fn(i64, i64) -> i64 + Sync,
+) -> crate::tensor::Tensor<i64> {
+    use crate::tensor::get_broadcasted_shape;
+    use common::parallel::par_enabled;
+    use rayon::prelude::*;
+
+    let shape = get_broadcasted_shape(lhs.dims(), rhs.dims())
+        .unwrap_or_else(|_| panic!("{op_name}: incompatible broadcast shapes"));
+    let lhs_exp = lhs
+        .expand(&shape)
+        .unwrap_or_else(|_| panic!("{op_name}: expand lhs"));
+    let rhs_exp = rhs
+        .expand(&shape)
+        .unwrap_or_else(|_| panic!("{op_name}: expand rhs"));
+    let data: Vec<i64> = lhs_exp
+        .data()
+        .par_iter()
+        .zip(rhs_exp.data().par_iter())
+        .with_min_len(par_enabled())
+        .map(|(&a, &b)| combine(a as i64, b as i64))
+        .collect();
+    crate::tensor::Tensor::new(Some(&data), &shape)
+        .unwrap_or_else(|_| panic!("{op_name}: Tensor::new"))
+}
+
+/// Clamp each element of a `Tensor<i64>` into `[i32::MIN, i32::MAX]`.
+pub(super) fn clamp_to_i32(t: &crate::tensor::Tensor<i64>) -> crate::tensor::Tensor<i32> {
+    use common::parallel::par_enabled;
+    use rayon::prelude::*;
+
+    let data: Vec<i32> = t
+        .data()
+        .par_iter()
+        .with_min_len(par_enabled())
+        .map(|&v| v.clamp(i32::MIN as i64, i32::MAX as i64) as i32)
+        .collect();
+    crate::tensor::Tensor::new(Some(&data), t.dims())
+        .unwrap_or_else(|_| panic!("clamp_to_i32: Tensor::new"))
+}
+
+/// Saturating element-wise binary operation for [`Add`] and [`Sub`].
+///
+/// For each consecutive input pair: accumulates via [`sat_accumulate_pair`]
+/// (in `i64`), then clamps to `i32` via [`clamp_to_i32`].  Per-step clamping
+/// preserves correct saturation semantics for variadic inputs.
+pub(super) fn sat_binop(
+    inputs: Vec<&crate::tensor::Tensor<i32>>,
+    op_name: &str,
+    combine: impl Fn(i64, i64) -> i64 + Sync + Copy,
+) -> crate::tensor::Tensor<i32> {
+    let mut output = inputs[0].clone();
+    for &rhs in &inputs[1..] {
+        output = clamp_to_i32(&sat_accumulate_pair(&output, rhs, op_name, combine));
+    }
+    output
+}
+
+/// Re-execute a binary [`Add`]/[`Sub`] node's accumulation, returning the
+/// pre-clamp `i64` intermediate.
+///
+/// The proof system uses this to recover the accumulation (lookup-index)
+/// polynomial for the saturating-clamp lookup without storing it in the trace.
+/// Panics on non-`Add`/`Sub` operators or a non-binary operand list.
+pub fn sat_binop_intermediate(
+    operator: &Operator,
+    lhs: &crate::tensor::Tensor<i32>,
+    rhs: &crate::tensor::Tensor<i32>,
+) -> crate::tensor::Tensor<i64> {
+    match operator {
+        Operator::Add(_) => sat_accumulate_pair(lhs, rhs, "Add", |a, b| a + b),
+        Operator::Sub(_) => sat_accumulate_pair(lhs, rhs, "Sub", |a, b| a - b),
+        _ => panic!("sat_binop_intermediate: expected Add or Sub, got {operator:?}"),
+    }
+}
+
 /// Shared evaluation for the periodic trig operators ([`Sin`], [`Cos`]).
 ///
 /// Both reduce the input modulo a `4π` approximation before applying their

atlas-onnx-tracer/src/ops/sub.rs

@@ -1,12 +1,12 @@
 use crate::{
     ops::{Op, Sub},
-    tensor::{self, Tensor},
+    tensor::Tensor,
 };
 
 impl Op for Sub {
     #[tracing::instrument(name = "Sub::f", skip_all)]
     fn f(&self, inputs: Vec<&Tensor<i32>>) -> Tensor<i32> {
-        tensor::ops::sub(&inputs).unwrap()
+        super::sat_binop(inputs, "Sub", |a, b| a - b)
     }
 
     fn requires_shape_equality(&self) -> bool {

atlas-onnx-tracer/src/ops/sum.rs

@@ -1,3 +1,4 @@
+use rayon::prelude::*;
 use tract_onnx::prelude::tract_itertools::Itertools;
 
 use crate::{
@@ -8,22 +9,21 @@ use crate::{
 impl Op for Sum {
     #[tracing::instrument(name = "Sum::f", skip_all)]
     fn f(&self, inputs: Vec<&Tensor<i32>>) -> Tensor<i32> {
-        #[cfg(feature = "fused-ops")]
-        {
-            sum_axes_i32(inputs[0], &self.axes).unwrap()
-        }
-        #[cfg(not(feature = "fused-ops"))]
-        {
-            crate::tensor::ops::sum_axes(inputs[0], &self.axes).unwrap()
-        }
+        sum_axes_i32(inputs[0], &self.axes).unwrap()
     }
 }
 
-/// Sum over specified axes of an i32 tensor, accumulating in i64 to prevent overflow,
-#[tracing::instrument(name = "tensor::ops::sum_axes_i32", skip_all)]
-pub fn sum_axes_i32(a: &Tensor<i32>, axes: &[usize]) -> Result<Tensor<i32>, TensorError> {
+/// Sum over specified axes, accumulating in `i64` to avoid overflow (the
+/// pre-clamp intermediate). Each output element is the exact integer sum over
+/// the reduced axes; summed axes collapse to length 1.
+///
+/// This is the polynomial the saturating-clamp proof reasons about: the proof
+/// system recovers it to bridge `output = SatClamp(acc)`.
+#[tracing::instrument(name = "tensor::ops::sum_axes_i64", skip_all)]
+pub fn sum_axes_i64(a: &Tensor<i32>, axes: &[usize]) -> Result<Tensor<i64>, TensorError> {
     if axes.is_empty() {
-        return Ok(a.clone());
+        let data: Vec<i64> = a.data().par_iter().map(|&v| v as i64).collect();
+        return Tensor::new(Some(&data), a.dims());
     }
 
     let mut new_dims = vec![];
@@ -35,15 +35,15 @@ pub fn sum_axes_i32(a: &Tensor<i32>, axes: &[usize]) -> Result<Tensor<i32>, Tens
         }
     }
 
-    let res = Tensor::new(None, &new_dims)?;
+    let res = Tensor::<i64>::new(None, &new_dims)?;
 
     let cartesian_coord = new_dims
         .iter()
         .map(|x| 0..*x)
         .multi_cartesian_product()
         .collect::<Vec<_>>();
 
-    let res = res.par_enum_map(|i, _: i32| {
+    res.par_enum_map(|i, _: i64| {
         let coord = cartesian_coord[i].clone();
         let mut prod_dims = vec![];
         for (j, c) in coord.iter().enumerate() {
@@ -55,13 +55,22 @@ pub fn sum_axes_i32(a: &Tensor<i32>, axes: &[usize]) -> Result<Tensor<i32>, Tens
         }
 
         let slice = a.get_slice(&prod_dims)?;
-        // Accumulate in i64
         let mut acc: i64 = 0;
         let _ = slice.map(|v| acc += v as i64);
+        Ok(acc)
+    })
+}
 
-        // Saturate instead of wrapping — prevents sign-flip corruption
-        Ok(acc.clamp(i32::MIN as i64, i32::MAX as i64) as i32)
-    })?;
-
-    Ok(res)
+/// Saturating sum over specified axes: [`sum_axes_i64`] clamped back to `i32`.
+///
+/// Saturating (rather than wrapping) prevents sign-flip corruption on overflow.
+#[tracing::instrument(name = "tensor::ops::sum_axes_i32", skip_all)]
+pub fn sum_axes_i32(a: &Tensor<i32>, axes: &[usize]) -> Result<Tensor<i32>, TensorError> {
+    let acc = sum_axes_i64(a, axes)?;
+    let data: Vec<i32> = acc
+        .data()
+        .par_iter()
+        .map(|&v| v.clamp(i32::MIN as i64, i32::MAX as i64) as i32)
+        .collect();
+    Tensor::new(Some(&data), acc.dims())
 }

common/src/lib.rs

@@ -160,6 +160,14 @@ canonical_serde_enum! {
         /// * `0` – node index
         /// * `1` – decomposition index `d`
         SoftmaxSatDiffRaD(usize, usize),
+
+        /// One-hot read-address decomposition for the saturating-clamp lookup of
+        /// a `Add`/`Sub` node. The clamp is indexed by the 64-bit accumulation, so
+        /// there are `64 / log_k_chunk` chunks (e.g. 16 for `log_k_chunk = 4`).
+        ///
+        /// * `0` – node index
+        /// * `1` – decomposition index `d`
+        ClampRaD(usize, usize),
     }
 }
 
@@ -367,5 +375,23 @@ canonical_serde_enum! {
         ///
         /// * `0` – producer node index
         NTEvalShiftOutput(usize),
+
+        // ----- Saturating clamp (Add/Sub) -----
+        /// MLE of the pre-clamp **i64 accumulation** of a saturating `Add`/`Sub`.
+        ///
+        /// This is the `raf` (read-address-field) polynomial of the saturating-clamp
+        /// lookup: its values are the indices into [`SatClampTable`](joltworks). It is
+        /// tied to the operands by the linear identity `acc(r) = left(r) ± right(r)`,
+        /// and the lookup proves `output(r) = clamp(acc(r))`.
+        ///
+        /// * `0` – node index
+        ClampAcc(usize),
+
+        /// One-hot read-address polynomial for the saturating-clamp lookup
+        /// (`Add`/`Sub`). Virtualized — [`CommittedPoly::ClampRaD`] commits its
+        /// decomposition.
+        ///
+        /// * `0` – node index
+        ClampRa(usize),
     }
 }