stwo_gpu_constraint_framework/
logup.rs

1use core::array;
2use core::ops::{Mul, Sub};
3
4use num_traits::{One, Zero};
5use std_shims::{vec, Vec};
6use stwo::core::channel::Channel;
7use stwo::core::fields::m31::BaseField;
8use stwo::core::fields::qm31::SecureField;
9use stwo::core::Fraction;
10
11use super::EvalAtRow;
12
13/// Evaluates constraints for batched logups.
14/// These constraint enforce the sum of multiplicity_i / (z + sum_j alpha^j * x_j) = claimed_sum.
15pub struct LogupAtRow<E: EvalAtRow> {
16    /// The index of the interaction used for the cumulative sum columns.
17    pub interaction: usize,
18    /// The total sum of all the fractions divided by n_rows.
19    pub cumsum_shift: SecureField,
20    /// The evaluation of the last cumulative sum column.
21    pub fracs: Vec<Fraction<E::EF, E::EF>>,
22    pub is_finalized: bool,
23    pub log_size: u32,
24}
25
26impl<E: EvalAtRow> Default for LogupAtRow<E> {
27    fn default() -> Self {
28        Self::dummy()
29    }
30}
31impl<E: EvalAtRow> LogupAtRow<E> {
32    pub fn new(interaction: usize, claimed_sum: SecureField, log_size: u32) -> Self {
33        Self {
34            interaction,
35            cumsum_shift: claimed_sum / BaseField::from_u32_unchecked(1 << log_size),
36            fracs: vec![],
37            is_finalized: true,
38            log_size,
39        }
40    }
41
42    // TODO(alont): Remove this once unnecessary LogupAtRows are gone.
43    pub fn dummy() -> Self {
44        Self {
45            interaction: 100,
46            cumsum_shift: SecureField::one(),
47            fracs: vec![],
48            is_finalized: true,
49            log_size: 10,
50        }
51    }
52}
53
54/// Ensures that the LogupAtRow is finalized.
55/// LogupAtRow should be finalized exactly once.
56impl<E: EvalAtRow> Drop for LogupAtRow<E> {
57    fn drop(&mut self) {
58        assert!(self.is_finalized, "LogupAtRow was not finalized");
59    }
60}
61
62/// Interaction elements for the logup protocol.
63#[derive(Clone, Debug, PartialEq, Eq)]
64pub struct LookupElements<const N: usize> {
65    pub z: SecureField,
66    pub alpha: SecureField,
67    pub alpha_powers: [SecureField; N],
68}
69impl<const N: usize> LookupElements<N> {
70    pub fn draw(channel: &mut impl Channel) -> Self {
71        let [z, alpha] = channel.draw_secure_felts(2).try_into().unwrap();
72        let mut cur = SecureField::one();
73        let alpha_powers = array::from_fn(|_| {
74            let res = cur;
75            cur *= alpha;
76            res
77        });
78        Self {
79            z,
80            alpha,
81            alpha_powers,
82        }
83    }
84    pub fn combine<F: Clone, EF>(&self, values: &[F]) -> EF
85    where
86        EF: Clone + Zero + From<F> + From<SecureField> + Mul<F, Output = EF> + Sub<EF, Output = EF>,
87    {
88        assert!(
89            self.alpha_powers.len() >= values.len(),
90            "Not enough alpha powers to combine values"
91        );
92        values
93            .iter()
94            .zip(self.alpha_powers)
95            .fold(EF::zero(), |acc, (value, power)| {
96                acc + EF::from(power) * value.clone()
97            })
98            - EF::from(self.z)
99    }
100
101    pub fn dummy() -> Self {
102        let z = SecureField::from_u32_unchecked(1, 2, 3, 4);
103        let alpha = SecureField::from_u32_unchecked(4, 3, 2, 1);
104        let mut cur = SecureField::one();
105        let alpha_powers = array::from_fn(|_| {
106            let res = cur;
107            cur *= alpha;
108            res
109        });
110        Self {
111            z,
112            alpha,
113            alpha_powers,
114        }
115    }
116}
117
118#[cfg(test)]
119mod tests {
120    use stwo::core::channel::Blake2sChannel;
121    use stwo::core::fields::m31::BaseField;
122    use stwo::core::fields::qm31::SecureField;
123    use stwo::core::fields::FieldExpOps;
124
125    use super::LookupElements;
126
127    #[test]
128    fn test_lookup_elements_combine() {
129        let mut channel = Blake2sChannel::default();
130        let lookup_elements = LookupElements::<3>::draw(&mut channel);
131        let values = [
132            BaseField::from_u32_unchecked(123),
133            BaseField::from_u32_unchecked(456),
134            BaseField::from_u32_unchecked(789),
135        ];
136
137        assert_eq!(
138            lookup_elements.combine::<BaseField, SecureField>(&values),
139            BaseField::from_u32_unchecked(123)
140                + BaseField::from_u32_unchecked(456) * lookup_elements.alpha
141                + BaseField::from_u32_unchecked(789) * lookup_elements.alpha.pow(2)
142                - lookup_elements.z
143        );
144    }
145}
stwo_gpu_constraint_framework/logup.rs

stwo_gpu_constraint_framework/
logup.rs
