stwo_gpu_constraint_framework/
info.rs

1use core::array;
2use core::cell::{RefCell, RefMut};
3use core::ops::{Add, AddAssign, Mul, MulAssign, Neg, Sub};
4
5use hashbrown::HashMap;
6use num_traits::{One, Zero};
7use std_shims::{vec, Rc, String, Vec};
8use stwo::core::fields::m31::BaseField;
9use stwo::core::fields::qm31::SecureField;
10use stwo::core::fields::FieldExpOps;
11use stwo::core::pcs::TreeVec;
12use stwo::core::verifier::PREPROCESSED_TRACE_IDX;
13use stwo::core::Fraction;
14
15use super::logup::LogupAtRow;
16use super::preprocessed_columns::PreProcessedColumnId;
17use super::{EvalAtRow, Relation, RelationEntry, INTERACTION_TRACE_IDX};
18
19/// Collects information about the constraints.
20/// This includes mask offsets and columns at each interaction, the number of constraints and number
21/// of arithmetic operations.
22#[derive(Default)]
23pub struct InfoEvaluator {
24    pub mask_offsets: TreeVec<Vec<Vec<isize>>>,
25    pub n_constraints: usize,
26    pub preprocessed_columns: Vec<PreProcessedColumnId>,
27    pub logup: LogupAtRow<Self>,
28    pub arithmetic_counts: ArithmeticCounts,
29    pub logup_counts: LogupCountPerRow,
30}
31impl InfoEvaluator {
32    pub fn new(
33        log_size: u32,
34        preprocessed_columns: Vec<PreProcessedColumnId>,
35        claimed_sum: SecureField,
36    ) -> Self {
37        Self {
38            mask_offsets: Default::default(),
39            n_constraints: Default::default(),
40            preprocessed_columns,
41            logup: LogupAtRow::new(INTERACTION_TRACE_IDX, claimed_sum, log_size),
42            arithmetic_counts: Default::default(),
43            logup_counts: LogupCountPerRow::new(),
44        }
45    }
46
47    /// Create an empty `InfoEvaluator`, to measure components before their size and logup sums are
48    /// available.
49    pub fn empty() -> Self {
50        Self::new(16, vec![], SecureField::default())
51    }
52}
53impl EvalAtRow for InfoEvaluator {
54    type F = FieldCounter;
55    type EF = ExtensionFieldCounter;
56
57    fn next_interaction_mask<const N: usize>(
58        &mut self,
59        interaction: usize,
60        offsets: [isize; N],
61    ) -> [Self::F; N] {
62        assert!(
63            interaction != PREPROCESSED_TRACE_IDX,
64            "Preprocessed should be accesses with `get_preprocessed_column`",
65        );
66
67        // Check if requested a mask from a new interaction
68        if self.mask_offsets.len() <= interaction {
69            // Extend `mask_offsets` so that `interaction` is the last index.
70            self.mask_offsets.resize(interaction + 1, vec![]);
71        }
72        self.mask_offsets[interaction].push(offsets.into_iter().collect());
73        array::from_fn(|_| FieldCounter::one())
74    }
75
76    fn get_preprocessed_column(&mut self, column: PreProcessedColumnId) -> Self::F {
77        self.preprocessed_columns.push(column);
78        FieldCounter::one()
79    }
80
81    fn add_constraint<G>(&mut self, constraint: G)
82    where
83        Self::EF: Mul<G, Output = Self::EF>,
84    {
85        let lin_combination =
86            ExtensionFieldCounter::one() + ExtensionFieldCounter::one() * constraint;
87        self.arithmetic_counts.merge(lin_combination.drain());
88        self.n_constraints += 1;
89    }
90
91    fn combine_ef(values: [Self::F; 4]) -> Self::EF {
92        let mut res = ExtensionFieldCounter::zero();
93        let _ = values.map(|v| res.merge(v));
94        res
95    }
96
97    fn add_to_relation<R: Relation<Self::F, Self::EF>>(
98        &mut self,
99        entry: RelationEntry<'_, Self::F, Self::EF, R>,
100    ) {
101        self.logup_counts.inc(entry.relation.get_name());
102        let frac = Fraction::new(
103            entry.multiplicity.clone(),
104            entry.relation.combine(entry.values),
105        );
106        self.write_logup_frac(frac);
107    }
108
109    super::logup_proxy!();
110}
111
112/// Stores a count of field operations.
113#[derive(Debug, Default, Clone, Copy, PartialEq, Eq)]
114pub struct ArithmeticCounts {
115    /// Number of [`EvalAtRow::EF`] * [`EvalAtRow::EF`] operations.
116    pub n_ef_mul_ef: usize,
117    /// Number of [`EvalAtRow::EF`] * [`EvalAtRow::F`] operations.
118    pub n_ef_mul_f: usize,
119    /// Number of [`EvalAtRow::F`] * [`BaseField`] operations.
120    pub n_ef_mul_base_field: usize,
121    /// Number of [`EvalAtRow::EF`] + [`EvalAtRow::EF`] operations.
122    pub n_ef_add_ef: usize,
123    /// Number of [`EvalAtRow::EF`] + [`EvalAtRow::F`] operations.
124    pub n_ef_add_f: usize,
125    /// Number of [`EvalAtRow::EF`] * [`BaseField`] operations.
126    pub n_ef_add_base_field: usize,
127    /// Number of [`EvalAtRow::F`] * [`EvalAtRow::F`] operations.
128    pub n_f_mul_f: usize,
129    /// Number of [`EvalAtRow::F`] * [`BaseField`] operations.
130    pub n_f_mul_base_field: usize,
131    /// Number of [`EvalAtRow::F`] + [`EvalAtRow::F`] operations.
132    pub n_f_add_f: usize,
133    /// Number of [`EvalAtRow::F`] + [`BaseField`] operations.
134    pub n_f_add_base_field: usize,
135}
136
137impl ArithmeticCounts {
138    const fn merge(&mut self, other: ArithmeticCounts) {
139        let Self {
140            n_ef_mul_ef,
141            n_ef_mul_f,
142            n_ef_mul_base_field,
143            n_ef_add_ef,
144            n_ef_add_f,
145            n_ef_add_base_field,
146            n_f_mul_f,
147            n_f_mul_base_field,
148            n_f_add_f,
149            n_f_add_base_field,
150        } = self;
151
152        *n_ef_mul_ef += other.n_ef_mul_ef;
153        *n_ef_mul_f += other.n_ef_mul_f;
154        *n_ef_mul_base_field += other.n_ef_mul_base_field;
155        *n_ef_add_f += other.n_ef_add_f;
156        *n_ef_add_base_field += other.n_ef_add_base_field;
157        *n_ef_add_ef += other.n_ef_add_ef;
158        *n_f_mul_f += other.n_f_mul_f;
159        *n_f_mul_base_field += other.n_f_mul_base_field;
160        *n_f_add_f += other.n_f_add_f;
161        *n_f_add_base_field += other.n_f_add_base_field;
162    }
163}
164
165#[derive(Debug, Default, Clone)]
166pub struct ArithmeticCounter<const IS_EXT_FIELD: bool>(Rc<RefCell<ArithmeticCounts>>);
167
168/// Counts operations on [`EvalAtRow::F`].
169pub type FieldCounter = ArithmeticCounter<false>;
170
171/// Counts operations on [`EvalAtRow::EF`].
172pub type ExtensionFieldCounter = ArithmeticCounter<true>;
173
174impl<const IS_EXT_FIELD: bool> ArithmeticCounter<IS_EXT_FIELD> {
175    fn merge<const OTHER_IS_EXT_FIELD: bool>(
176        &mut self,
177        other: ArithmeticCounter<OTHER_IS_EXT_FIELD>,
178    ) {
179        // Skip if they come from the same source.
180        if Rc::ptr_eq(&self.0, &other.0) {
181            return;
182        }
183
184        self.counts().merge(other.drain());
185    }
186
187    fn drain(self) -> ArithmeticCounts {
188        self.0.take()
189    }
190
191    fn counts(&mut self) -> RefMut<'_, ArithmeticCounts> {
192        self.0.borrow_mut()
193    }
194}
195
196impl<const IS_EXT_FIELD: bool> Zero for ArithmeticCounter<IS_EXT_FIELD> {
197    fn zero() -> Self {
198        Self::default()
199    }
200
201    fn is_zero(&self) -> bool {
202        // TODO(andrew): Consider removing Zero from EvalAtRow::F, EvalAtRow::EF since is_zero
203        // doesn't make sense. Creating zero elements does though.
204        panic!()
205    }
206}
207
208impl<const IS_EXT_FIELD: bool> One for ArithmeticCounter<IS_EXT_FIELD> {
209    fn one() -> Self {
210        Self::default()
211    }
212}
213
214impl<const IS_EXT_FIELD: bool> Add for ArithmeticCounter<IS_EXT_FIELD> {
215    type Output = Self;
216
217    fn add(mut self, rhs: Self) -> Self {
218        self.merge(rhs);
219        {
220            let mut counts = self.counts();
221            match IS_EXT_FIELD {
222                true => counts.n_ef_add_ef += 1,
223                false => counts.n_f_add_f += 1,
224            }
225        }
226        self
227    }
228}
229
230impl<const IS_EXT_FIELD: bool> Sub for ArithmeticCounter<IS_EXT_FIELD> {
231    type Output = Self;
232
233    #[allow(clippy::suspicious_arithmetic_impl)]
234    fn sub(self, rhs: Self) -> Self {
235        // Treat as addition.
236        self + rhs
237    }
238}
239
240impl Add<FieldCounter> for ExtensionFieldCounter {
241    type Output = Self;
242
243    fn add(mut self, rhs: FieldCounter) -> Self {
244        self.merge(rhs);
245        self.counts().n_ef_add_f += 1;
246        self
247    }
248}
249
250impl<const IS_EXT_FIELD: bool> Mul for ArithmeticCounter<IS_EXT_FIELD> {
251    type Output = Self;
252
253    fn mul(mut self, rhs: Self) -> Self {
254        self.merge(rhs);
255        {
256            let mut counts = self.counts();
257            match IS_EXT_FIELD {
258                true => counts.n_ef_mul_ef += 1,
259                false => counts.n_f_mul_f += 1,
260            }
261        }
262        self
263    }
264}
265
266impl Mul<FieldCounter> for ExtensionFieldCounter {
267    type Output = ExtensionFieldCounter;
268
269    #[allow(clippy::suspicious_arithmetic_impl)]
270    fn mul(mut self, rhs: FieldCounter) -> Self {
271        self.merge(rhs);
272        self.counts().n_ef_mul_f += 1;
273        self
274    }
275}
276
277impl<const IS_EXT_FIELD: bool> MulAssign for ArithmeticCounter<IS_EXT_FIELD> {
278    fn mul_assign(&mut self, rhs: Self) {
279        *self = self.clone() * rhs
280    }
281}
282
283impl<const IS_EXT_FIELD: bool> AddAssign for ArithmeticCounter<IS_EXT_FIELD> {
284    fn add_assign(&mut self, rhs: Self) {
285        *self = self.clone() + rhs
286    }
287}
288
289impl AddAssign<BaseField> for FieldCounter {
290    fn add_assign(&mut self, _rhs: BaseField) {
291        self.counts().n_f_add_base_field += 1;
292    }
293}
294
295impl Mul<BaseField> for FieldCounter {
296    type Output = Self;
297
298    #[allow(clippy::suspicious_arithmetic_impl)]
299    fn mul(mut self, _rhs: BaseField) -> Self {
300        self.counts().n_f_mul_base_field += 1;
301        self
302    }
303}
304
305impl Mul<BaseField> for ExtensionFieldCounter {
306    type Output = Self;
307
308    #[allow(clippy::suspicious_arithmetic_impl)]
309    fn mul(mut self, _rhs: BaseField) -> Self {
310        self.counts().n_ef_mul_base_field += 1;
311        self
312    }
313}
314
315impl Mul<SecureField> for ExtensionFieldCounter {
316    type Output = Self;
317
318    fn mul(self, _rhs: SecureField) -> Self {
319        self * ExtensionFieldCounter::zero()
320    }
321}
322
323impl Add<SecureField> for FieldCounter {
324    type Output = ExtensionFieldCounter;
325
326    fn add(self, _rhs: SecureField) -> ExtensionFieldCounter {
327        ExtensionFieldCounter::zero() + self
328    }
329}
330
331impl Add<BaseField> for ExtensionFieldCounter {
332    type Output = Self;
333
334    fn add(mut self, _rhs: BaseField) -> Self {
335        self.counts().n_ef_add_base_field += 1;
336        self
337    }
338}
339
340impl Add<SecureField> for ExtensionFieldCounter {
341    type Output = Self;
342
343    fn add(self, _rhs: SecureField) -> Self {
344        self + ExtensionFieldCounter::zero()
345    }
346}
347
348impl Sub<SecureField> for ExtensionFieldCounter {
349    type Output = Self;
350
351    #[allow(clippy::suspicious_arithmetic_impl)]
352    fn sub(self, rhs: SecureField) -> Self {
353        // Tread subtraction as addition
354        self + rhs
355    }
356}
357
358impl Mul<SecureField> for FieldCounter {
359    type Output = ExtensionFieldCounter;
360
361    fn mul(self, _rhs: SecureField) -> ExtensionFieldCounter {
362        ExtensionFieldCounter::zero() * self
363    }
364}
365
366impl From<BaseField> for FieldCounter {
367    fn from(_value: BaseField) -> Self {
368        Self::one()
369    }
370}
371
372impl From<SecureField> for ExtensionFieldCounter {
373    fn from(_value: SecureField) -> Self {
374        Self::one()
375    }
376}
377
378impl From<FieldCounter> for ExtensionFieldCounter {
379    fn from(value: FieldCounter) -> Self {
380        Self(value.0)
381    }
382}
383
384impl<const IS_EXT_FIELD: bool> Neg for ArithmeticCounter<IS_EXT_FIELD> {
385    type Output = Self;
386
387    fn neg(self) -> Self {
388        // Treat as addition.
389        self + ArithmeticCounter::<IS_EXT_FIELD>::zero()
390    }
391}
392
393impl<const IS_EXT_FIELD: bool> FieldExpOps for ArithmeticCounter<IS_EXT_FIELD> {
394    fn inverse(&self) -> Self {
395        todo!()
396    }
397}
398
399#[derive(Debug, Default, Clone)]
400pub struct LogupCountPerRow {
401    data: HashMap<String, usize>,
402}
403impl LogupCountPerRow {
404    pub fn new() -> Self {
405        Self {
406            data: HashMap::new(),
407        }
408    }
409
410    pub fn inc(&mut self, relation: &str) {
411        *self
412            .data
413            .entry(std_shims::ToString::to_string(relation))
414            .or_insert(0) += 1;
415    }
416
417    pub fn iter(&self) -> impl Iterator<Item = (&String, &usize)> {
418        self.data.iter()
419    }
420}
421
422#[cfg(test)]
423mod tests {
424    use num_traits::{One, Zero};
425    use stwo::core::fields::m31::BaseField;
426    use stwo::core::fields::qm31::SecureField;
427
428    use super::{ExtensionFieldCounter, InfoEvaluator};
429    use crate::info::{ArithmeticCounts, FieldCounter};
430    use crate::{relation, EvalAtRow, FrameworkEval, RelationEntry};
431
432    #[test]
433    fn test_arithmetic_counter() {
434        const N_EF_MUL_EF: usize = 1;
435        const N_EF_MUL_F: usize = 2;
436        const N_EF_MUL_BASE_FIELD: usize = 3;
437        const N_EF_MUL_ASSIGN_EF: usize = 4;
438        const N_EF_MUL_SECURE_FIELD: usize = 5;
439        const N_EF_ADD_EF: usize = 6;
440        const N_EF_ADD_ASSIGN_EF: usize = 7;
441        const N_EF_ADD_F: usize = 8;
442        const N_EF_NEG: usize = 9;
443        const N_EF_SUB_EF: usize = 10;
444        const N_EF_ADD_BASE_FIELD: usize = 11;
445        const N_F_MUL_F: usize = 12;
446        const N_F_MUL_ASSIGN_F: usize = 13;
447        const N_F_MUL_BASE_FIELD: usize = 14;
448        const N_F_ADD_F: usize = 15;
449        const N_F_ADD_ASSIGN_F: usize = 16;
450        const N_F_ADD_ASSIGN_BASE_FIELD: usize = 17;
451        const N_F_NEG: usize = 18;
452        const N_F_SUB_F: usize = 19;
453        let mut ef = ExtensionFieldCounter::zero();
454        let mut f = FieldCounter::zero();
455
456        (0..N_EF_MUL_EF).for_each(|_| ef = ef.clone() * ef.clone());
457        (0..N_EF_MUL_F).for_each(|_| ef = ef.clone() * f.clone());
458        (0..N_EF_MUL_BASE_FIELD).for_each(|_| ef = ef.clone() * BaseField::one());
459        (0..N_EF_MUL_SECURE_FIELD).for_each(|_| ef = ef.clone() * SecureField::one());
460        (0..N_EF_MUL_ASSIGN_EF).for_each(|_| ef *= ef.clone());
461        (0..N_EF_ADD_EF).for_each(|_| ef = ef.clone() + ef.clone());
462        (0..N_EF_ADD_ASSIGN_EF).for_each(|_| ef += ef.clone());
463        (0..N_EF_ADD_F).for_each(|_| ef = ef.clone() + f.clone());
464        (0..N_EF_ADD_BASE_FIELD).for_each(|_| ef = ef.clone() + BaseField::one());
465        (0..N_EF_NEG).for_each(|_| ef = -ef.clone());
466        (0..N_EF_SUB_EF).for_each(|_| ef = ef.clone() - ef.clone());
467        (0..N_F_MUL_F).for_each(|_| f = f.clone() * f.clone());
468        (0..N_F_MUL_ASSIGN_F).for_each(|_| f *= f.clone());
469        (0..N_F_MUL_BASE_FIELD).for_each(|_| f = f.clone() * BaseField::one());
470        (0..N_F_ADD_F).for_each(|_| f = f.clone() + f.clone());
471        (0..N_F_ADD_ASSIGN_F).for_each(|_| f += f.clone());
472        (0..N_F_ADD_ASSIGN_BASE_FIELD).for_each(|_| f += BaseField::one());
473        (0..N_F_NEG).for_each(|_| f = -f.clone());
474        (0..N_F_SUB_F).for_each(|_| f = f.clone() - f.clone());
475        let mut res = f.drain();
476        res.merge(ef.drain());
477
478        assert_eq!(
479            res,
480            ArithmeticCounts {
481                n_ef_mul_ef: N_EF_MUL_EF + N_EF_MUL_SECURE_FIELD + N_EF_MUL_ASSIGN_EF,
482                n_ef_mul_base_field: N_EF_MUL_BASE_FIELD,
483                n_ef_mul_f: N_EF_MUL_F,
484                n_ef_add_ef: N_EF_ADD_EF + N_EF_NEG + N_EF_SUB_EF + N_EF_ADD_ASSIGN_EF,
485                n_ef_add_f: N_EF_ADD_F,
486                n_ef_add_base_field: N_EF_ADD_BASE_FIELD,
487                n_f_mul_f: N_F_MUL_F + N_F_MUL_ASSIGN_F,
488                n_f_mul_base_field: N_F_MUL_BASE_FIELD,
489                n_f_add_f: N_F_ADD_F + N_F_NEG + N_F_SUB_F + N_F_ADD_ASSIGN_F,
490                n_f_add_base_field: N_F_ADD_ASSIGN_BASE_FIELD,
491            }
492        );
493    }
494
495    relation!(TestRelation, 0);
496    relation!(TestRelation2, 1);
497
498    struct TestEval {
499        relation: TestRelation,
500        relation2: TestRelation2,
501    }
502    impl FrameworkEval for TestEval {
503        fn log_size(&self) -> u32 {
504            unimplemented!()
505        }
506
507        fn max_constraint_log_degree_bound(&self) -> u32 {
508            unimplemented!()
509        }
510
511        fn evaluate<E: EvalAtRow>(&self, mut eval: E) -> E {
512            eval.add_to_relation(RelationEntry::new(
513                &self.relation,
514                SecureField::one().into(),
515                &[],
516            ));
517            eval.add_to_relation(RelationEntry::new(
518                &self.relation2,
519                SecureField::one().into(),
520                &[],
521            ));
522            eval.add_to_relation(RelationEntry::new(
523                &self.relation2,
524                SecureField::one().into(),
525                &[],
526            ));
527
528            eval.finalize_logup();
529            eval
530        }
531    }
532
533    #[test]
534    fn test_logup_count_per_row() {
535        let eval = TestEval {
536            relation: TestRelation::dummy(),
537            relation2: TestRelation2::dummy(),
538        };
539        let info = eval.evaluate(InfoEvaluator::empty());
540
541        assert_eq!(info.logup_counts.data["TestRelation"], 1);
542        assert_eq!(info.logup_counts.data["TestRelation2"], 2);
543    }
544}
stwo_gpu_constraint_framework/info.rs

stwo_gpu_constraint_framework/
info.rs
