Merge pull request #125 from PolyhedraZK/fix-scratchpad

Fix scratchpad and Optimize Expander Eq
PolyhedraZK · Oct 14, 2024 · d03a433 · d03a433
2 parents 274cd76 + 70a15f0
commit d03a433
Show file tree

Hide file tree

Showing 2 changed files with 24 additions and 21 deletions.
diff --git a/arith/polynomials/src/eq.rs b/arith/polynomials/src/eq.rs
@@ -30,25 +30,16 @@ impl<F: Field> EqPolynomial<F> {
     #[inline]
     /// Expander's method of computing Eq(x, r)
     pub fn build_eq_x_r_with_buf(r: &[F], mul_factor: &F, eq_evals: &mut [F]) {
-        // well, Jolt's method to build eq(x,r) is slightly faster
-        //
-        let tmp = Self::scaled_evals_jolt(r, mul_factor);
-        eq_evals[..tmp.len()].copy_from_slice(&tmp);
-
-        // Expander's original method to build scaled eq(x,r)
-        //
-        // eq_evals[0] = *mul_factor;
-        // let mut cur_eval_num = 1;
-
-        // for r_i in r.iter() {
-        //     let eq_z_i_zero = F::one() - r_i;
-        //     let eq_z_i_one = r_i;
-        //     for j in 0..cur_eval_num {
-        //         eq_evals[j + cur_eval_num] = eq_evals[j] * eq_z_i_one;
-        //         eq_evals[j] *= eq_z_i_zero;
-        //     }
-        //     cur_eval_num <<= 1;
-        // }
+        eq_evals[0] = *mul_factor;
+        let mut cur_eval_num = 1;
+
+        for r_i in r.iter() {
+            for j in 0..cur_eval_num {
+                eq_evals[j + cur_eval_num] = eq_evals[j] * r_i;
+                eq_evals[j] -= eq_evals[j + cur_eval_num];
+            }
+            cur_eval_num <<= 1;
+        }
     }
 
     #[inline(always)]

diff --git a/sumcheck/src/scratch_pad.rs b/sumcheck/src/scratch_pad.rs
@@ -48,8 +48,20 @@ impl<C: GKRConfig> ProverScratchPad<C> {
             eq_evals_at_rz1: vec![C::ChallengeField::default(); max_output_num],
             eq_evals_at_r_simd0: vec![C::ChallengeField::default(); C::get_field_pack_size()],
             eq_evals_at_r_mpi0: vec![C::ChallengeField::default(); mpi_world_size],
-            eq_evals_first_half: vec![C::ChallengeField::default(); max_output_num],
-            eq_evals_second_half: vec![C::ChallengeField::default(); max_output_num],
+            eq_evals_first_half: vec![
+                C::ChallengeField::default();
+                max(
+                    max(max_output_num, C::get_field_pack_size()),
+                    mpi_world_size
+                )
+            ],
+            eq_evals_second_half: vec![
+                C::ChallengeField::default();
+                max(
+                    max(max_output_num, C::get_field_pack_size()),
+                    mpi_world_size
+                )
+            ],
 
             gate_exists_5: vec![false; max_input_num],
             gate_exists_1: vec![false; max_input_num],