openvm_stark_backend/prover/logup_zerocheck/
fractional_sumcheck_gkr.rs

1use std::ops::Add;
2
3use p3_field::{Field, PrimeCharacteristicRing};
4use p3_util::log2_strict_usize;
5use tracing::{debug, instrument};
6
7use crate::{
8    proof::GkrLayerClaims,
9    prover::{
10        error::LogupZerocheckError,
11        poly::evals_eq_hypercube,
12        sumcheck::{fold_mle_evals, sumcheck_round_poly_evals},
13        ColMajorMatrix,
14    },
15    FiatShamirTranscript, StarkProtocolConfig,
16};
17
18/// Proof for fractional sumcheck protocol
19pub struct FracSumcheckProof<SC: StarkProtocolConfig> {
20    /// The fractional sum p_0 / q_0
21    pub fractional_sum: (SC::EF, SC::EF),
22    /// The claims for p_j(0, rho), p_j(1, rho), q_j(0, rho), and q_j(1, rho) for each layer j > 0.
23    pub claims_per_layer: Vec<GkrLayerClaims<SC>>,
24    /// Sumcheck polynomials for each layer, for each sumcheck round, given by their evaluations on
25    /// {1, 2, 3}.
26    pub sumcheck_polys: Vec<Vec<[SC::EF; 3]>>,
27}
28
29#[derive(Clone, Copy, Debug, Default, derive_new::new)]
30#[repr(C)]
31pub struct Frac<EF> {
32    // PERF[jpw]: in the initial round, we can keep `p` in base field
33    pub p: EF,
34    pub q: EF,
35}
36
37impl<EF: Field> Add<Frac<EF>> for Frac<EF> {
38    type Output = Frac<EF>;
39
40    fn add(self, other: Frac<EF>) -> Self::Output {
41        Frac {
42            p: self.p * other.q + self.q * other.p,
43            q: self.q * other.q,
44        }
45    }
46}
47
48/// Runs the fractional sumcheck protocol using GKR layered circuit.
49///
50/// # Arguments
51/// * `transcript` - The Fiat-Shamir transcript
52/// * `evals` - list of `(p, q)` pairs of fractions in projective coordinates representing
53///   evaluations on the hypercube
54/// * `assert_zero` - Whether to assert that the final sum is zero. If `true`, then the transcript
55///   will not observe the numerator of the final sum.
56///
57/// # Returns
58/// The fractional sumcheck proof and the final random evaluation vector.
59#[instrument(level = "info", skip_all)]
60pub fn fractional_sumcheck<SC: StarkProtocolConfig, TS: FiatShamirTranscript<SC>>(
61    transcript: &mut TS,
62    evals: &[Frac<SC::EF>],
63    assert_zero: bool,
64) -> Result<(FracSumcheckProof<SC>, Vec<SC::EF>), LogupZerocheckError> {
65    if evals.is_empty() {
66        return Ok((
67            FracSumcheckProof {
68                fractional_sum: (SC::EF::ZERO, SC::EF::ONE),
69                claims_per_layer: vec![],
70                sumcheck_polys: vec![],
71            },
72            vec![],
73        ));
74    }
75    // total_rounds = l_skip + n_logup
76    let total_rounds = log2_strict_usize(evals.len());
77    // sumcheck polys for layers j=2,...,total_rounds
78    let mut sumcheck_polys = Vec::with_capacity(total_rounds);
79
80    // segment tree: layer i=0,...,total_rounds starts at 2^i (index 0 unused)
81    let mut tree_evals: Vec<Frac<SC::EF>> = vec![Frac::default(); 2 << total_rounds];
82    tree_evals[(1 << total_rounds)..].copy_from_slice(evals);
83
84    for node_idx in (1..(1 << total_rounds)).rev() {
85        tree_evals[node_idx] = tree_evals[2 * node_idx] + tree_evals[2 * node_idx + 1];
86    }
87    let frac_sum = tree_evals[1];
88    if assert_zero {
89        if frac_sum.p != SC::EF::ZERO {
90            return Err(LogupZerocheckError::NonZeroRootSum);
91        }
92    } else {
93        transcript.observe_ext(frac_sum.p);
94    }
95    transcript.observe_ext(frac_sum.q);
96
97    // Index i is for layer i+1
98    let mut claims_per_layer: Vec<GkrLayerClaims<SC>> = Vec::with_capacity(total_rounds);
99
100    // Process each GKR round
101    // `j = round + 1` goes from `1, ..., total_rounds`
102    //
103    // Round `j = 1` is special since "sumcheck" is directly checked by verifier
104    claims_per_layer.push(GkrLayerClaims::<SC> {
105        p_xi_0: tree_evals[2].p,
106        q_xi_0: tree_evals[2].q,
107        p_xi_1: tree_evals[3].p,
108        q_xi_1: tree_evals[3].q,
109    });
110    transcript.observe_ext(claims_per_layer[0].p_xi_0);
111    transcript.observe_ext(claims_per_layer[0].q_xi_0);
112    transcript.observe_ext(claims_per_layer[0].p_xi_1);
113    transcript.observe_ext(claims_per_layer[0].q_xi_1);
114    let mu_1 = transcript.sample_ext();
115    debug!(gkr_round = 0, mu = %mu_1);
116    // ξ^{(j-1)}
117    let mut xi_prev = vec![mu_1];
118
119    // GKR rounds
120    for round in 1..total_rounds {
121        // Number of hypercube points
122        let eval_size = 1 << round;
123        // We apply batch sumcheck to the polynomials
124        // \eq(ξ^{(j-1)}, Y) (\hat p_j(0, Y) \hat q_j(1, Y) + \hat p_j(1, Y) \hat q_j(0, Y))
125        // \eq(ξ^{(j-1)}, Y) (\hat q_j(0, Y) \hat q_j(1, Y))
126        // Note: these are polynomials of degree 3 in each Y_i coordinate.
127
128        // Sample λ_j for batching
129        let lambda = transcript.sample_ext();
130        debug!(gkr_round = round, %lambda);
131
132        // Columns are p_j0, q_j0, p_j1, q_j1
133        // PERF: use a view instead of re-allocating memory
134        let mut pq_j_evals = SC::EF::zero_vec(4 * eval_size);
135        let segment = &tree_evals[2 * eval_size..4 * eval_size];
136        for x in 0..eval_size {
137            pq_j_evals[x] = segment[2 * x].p;
138            pq_j_evals[eval_size + x] = segment[2 * x].q;
139            pq_j_evals[2 * eval_size + x] = segment[2 * x + 1].p;
140            pq_j_evals[3 * eval_size + x] = segment[2 * x + 1].q;
141        }
142        let mut pq_j_evals = ColMajorMatrix::new(pq_j_evals, 4);
143        let mut eq_xis = ColMajorMatrix::new(evals_eq_hypercube(&xi_prev), 1);
144
145        // Batch sumcheck where the round polynomials are evaluated at {1, 2, 3}
146        let (round_polys_eval, rho) = {
147            let n = round;
148            let mut round_polys_eval = Vec::with_capacity(n);
149            let mut r_vec = Vec::with_capacity(n);
150
151            // Sumcheck rounds: apply fraction addition in projective coordinates to MLEs
152            for sumcheck_round in 0..n {
153                // Evaluate the univariate polynomial at {1, 2, 3}
154                // :projective fraction addition is degree 2, and then another +1 for eq
155                let [s_evals] = sumcheck_round_poly_evals(
156                    n - sumcheck_round,
157                    3,
158                    &[eq_xis.as_view(), pq_j_evals.as_view()],
159                    |_x, _y, row| {
160                        let eq_xi = row[0][0];
161                        let &[p_j0, q_j0, p_j1, q_j1] = row[1].as_slice() else {
162                            unreachable!("pq_j_evals always has 4 columns")
163                        };
164                        let p_prev = p_j0 * q_j1 + p_j1 * q_j0;
165                        let q_prev = q_j0 * q_j1;
166                        // batch using lambda
167                        [eq_xi * (p_prev + lambda * q_prev)]
168                    },
169                );
170                let s_evals: [SC::EF; 3] = s_evals.try_into().unwrap();
171                for &eval in &s_evals {
172                    transcript.observe_ext(eval);
173                }
174                round_polys_eval.push(s_evals);
175
176                let r_round = transcript.sample_ext();
177                pq_j_evals = fold_mle_evals(pq_j_evals, r_round);
178                eq_xis = fold_mle_evals(eq_xis, r_round);
179                r_vec.push(r_round);
180                debug!(gkr_round = round, %sumcheck_round, %r_round);
181            }
182            (round_polys_eval, r_vec)
183        };
184        claims_per_layer.push(GkrLayerClaims::<SC> {
185            p_xi_0: pq_j_evals.column(0)[0],
186            q_xi_0: pq_j_evals.column(1)[0],
187            p_xi_1: pq_j_evals.column(2)[0],
188            q_xi_1: pq_j_evals.column(3)[0],
189        });
190        transcript.observe_ext(claims_per_layer[round].p_xi_0);
191        transcript.observe_ext(claims_per_layer[round].q_xi_0);
192        transcript.observe_ext(claims_per_layer[round].p_xi_1);
193        transcript.observe_ext(claims_per_layer[round].q_xi_1);
194
195        // Sample μ_j for reduction to single evaluation point
196        let mu = transcript.sample_ext();
197        debug!(gkr_round = round, %mu);
198
199        // Update ξ^{(j)} = (μ_j, ρ^{(j-1)})
200        xi_prev = [vec![mu], rho].concat();
201
202        sumcheck_polys.push(round_polys_eval);
203    }
204
205    Ok((
206        FracSumcheckProof {
207            fractional_sum: (frac_sum.p, frac_sum.q),
208            claims_per_layer,
209            sumcheck_polys,
210        },
211        xi_prev,
212    ))
213}
openvm_stark_backend/prover/logup_zerocheck/fractional_sumcheck_gkr.rs

openvm_stark_backend/prover/logup_zerocheck/
fractional_sumcheck_gkr.rs
