Fix: ShardRamCircuit runs too slow (#1117)

# Issue

There is recurring pattern in Poseidon2 gadget which incurs a lot of
monomial terms. I.e.
1. `state'[i] = \sum_j M[i][j] * state[j]`;
2. `state''[i] = state'[i]^3` => this will create a lot of terms due to
the nature of distributive law.

## Walk around

The temporary solution to this issue is to allocate polynomials to store
the value after each linear layer operation (1).

| before/after| num_polys | num monomial terms of zerocheck |
|------------|------------|------------------------------------|
| before        |  329            | 7139                     |
| after | 374 (16*2 + 13 ⬆️) | 957 |

Author: kunxian-xia

ceno_zkvm/src/gadgets/poseidon2.rs

@@ -78,7 +78,8 @@ pub struct Poseidon2Config<
     const HALF_FULL_ROUNDS: usize,
     const PARTIAL_ROUNDS: usize,
 > {
-    cols: Vec<WitIn>,
+    p3_cols: Vec<WitIn>,                // columns in the plonky3-air
+    post_linear_layer_cols: Vec<WitIn>, /* additional columns to hold the state after linear layers */
     constants: RoundConstants<E::BaseField, STATE_WIDTH, HALF_FULL_ROUNDS, PARTIAL_ROUNDS>,
 }
 
@@ -141,10 +142,6 @@ impl<
             }
             (7, 1) => {
                 let committed_x3: Expression<E> = sbox.0[0].clone();
-                // TODO: avoid x^3 as x may have ~STATE_WIDTH terms after the linear layer
-                //       we can allocate one more column to store x^2 (which has ~STATE_WIDTH^2 terms)
-                //       then x^3 = x * x^2
-                //       but this will increase the number of columns (by FULL_ROUNDS * STATE_WIDTH + PARTIAL_ROUNDS)
                 cb.require_zero(|| "x3 = x.cube()", committed_x3.clone() - x.cube())?;
                 committed_x3.square() * x.clone()
             }
@@ -169,7 +166,7 @@ impl<
         }
         Self::external_linear_layer(state);
         for (state_i, post_i) in state.iter_mut().zip_eq(full_round.post.iter()) {
-            cb.require_zero(|| "post_i = state_i", state_i.clone() - post_i)?;
+            cb.require_equal(|| "post_i = state_i", state_i.clone(), post_i.clone())?;
             *state_i = post_i.clone();
         }
 
@@ -178,11 +175,18 @@ impl<
 
     fn eval_partial_round(
         state: &mut [Expression<E>; STATE_WIDTH],
+        post_linear_layer: &WitIn,
         partial_round: &PartialRound<Expression<E>, STATE_WIDTH, SBOX_DEGREE, SBOX_REGISTERS>,
         round_constant: &E::BaseField,
         cb: &mut CircuitBuilder<E>,
     ) -> Result<(), CircuitBuilderError> {
-        state[0] = state[0].clone() + round_constant.expr();
+        cb.require_equal(
+            || "post_linear_layer[0] = state[0]",
+            post_linear_layer.expr(),
+            state[0].clone() + round_constant.expr(),
+        )?;
+        state[0] = post_linear_layer.expr();
+
         Self::eval_sbox(&partial_round.sbox, &mut state[0], cb)?;
 
         cb.require_zero(
@@ -231,17 +235,40 @@ impl<
             PARTIAL_ROUNDS,
         >,
     ) -> Self {
-        let num_cols =
+        let num_p3_cols =
             num_cols::<STATE_WIDTH, SBOX_DEGREE, SBOX_REGISTERS, HALF_FULL_ROUNDS, PARTIAL_ROUNDS>(
             );
-        let cols = from_fn(|| Some(cb.create_witin(|| "poseidon2 col")))
-            .take(num_cols)
+        let p3_cols = from_fn(|| Some(cb.create_witin(|| "poseidon2 col")))
+            .take(num_p3_cols)
             .collect::<Vec<_>>();
-        let mut col_exprs = cols
+        let mut col_exprs = p3_cols
             .iter()
             .map(|c| c.expr())
             .collect::<Vec<Expression<E>>>();
 
+        // allocate columns to cache the state after each linear layer
+        // 1. before 0th full round
+        let mut post_linear_layer_cols = (0..STATE_WIDTH)
+            .map(|j| {
+                cb.create_witin(|| format!("[before 0th full round] post linear layer col[{j}]"))
+            })
+            .collect::<Vec<WitIn>>();
+        // 2. before each partial round
+        for i in 0..PARTIAL_ROUNDS {
+            post_linear_layer_cols.push(cb.create_witin(|| {
+                format!("[round {}] post linear layer col", i + HALF_FULL_ROUNDS)
+            }));
+        }
+        // 3. before HALF_FULL_ROUNDS-th full round
+        post_linear_layer_cols.extend((0..STATE_WIDTH).map(|j| {
+            cb.create_witin(|| {
+                format!(
+                    "[before {}th full round] post linear layer col[{j}]",
+                    HALF_FULL_ROUNDS
+                )
+            })
+        }));
+
         let poseidon2_cols: &mut Poseidon2Cols<
             Expression<E>,
             STATE_WIDTH,
@@ -254,6 +281,23 @@ impl<
         // external linear layer
         Self::external_linear_layer(&mut poseidon2_cols.inputs);
 
+        // after linear layer, each state_i has ~STATE_WIDTH terms
+        // therefore, we want to reduce that to one as the number of terms
+        // after sbox(state_i + rc_i) = (state_i + rc_i)^d will explode
+        poseidon2_cols
+            .inputs
+            .iter_mut()
+            .zip_eq(post_linear_layer_cols[0..STATE_WIDTH].iter())
+            .for_each(|(input, post_linear)| {
+                cb.require_equal(
+                    || "post_linear_layer = input",
+                    post_linear.expr(),
+                    input.clone(),
+                )
+                .unwrap();
+                *input = post_linear.expr();
+            });
+
         // eval full round
         for round in 0..HALF_FULL_ROUNDS {
             Self::eval_full_round(
@@ -269,15 +313,27 @@ impl<
         for round in 0..PARTIAL_ROUNDS {
             Self::eval_partial_round(
                 &mut poseidon2_cols.inputs,
+                &post_linear_layer_cols[STATE_WIDTH + round],
                 &poseidon2_cols.partial_rounds[round],
                 &round_constants.partial_round_constants[round],
                 cb,
             )
             .unwrap();
         }
 
-        // TODO: after the last partial round, each state_i has ~STATE_WIDTH terms
-        //       which will make the next full round to have many terms
+        poseidon2_cols
+            .inputs
+            .iter_mut()
+            .zip_eq(post_linear_layer_cols[STATE_WIDTH + PARTIAL_ROUNDS..].iter())
+            .for_each(|(input, post_linear)| {
+                cb.require_equal(
+                    || "post_linear_layer = input",
+                    post_linear.expr(),
+                    input.clone(),
+                )
+                .unwrap();
+                *input = post_linear.expr();
+            });
 
         // eval full round
         for round in 0..HALF_FULL_ROUNDS {
@@ -291,13 +347,14 @@ impl<
         }
 
         Poseidon2Config {
-            cols,
+            p3_cols,
+            post_linear_layer_cols,
             constants: round_constants,
         }
     }
 
     pub fn inputs(&self) -> Vec<Expression<E>> {
-        let col_exprs = self.cols.iter().map(|c| c.expr()).collect::<Vec<_>>();
+        let col_exprs = self.p3_cols.iter().map(|c| c.expr()).collect::<Vec<_>>();
 
         let poseidon2_cols: &Poseidon2Cols<
             Expression<E>,
@@ -312,7 +369,7 @@ impl<
     }
 
     pub fn output(&self) -> Vec<Expression<E>> {
-        let col_exprs = self.cols.iter().map(|c| c.expr()).collect::<Vec<_>>();
+        let col_exprs = self.p3_cols.iter().map(|c| c.expr()).collect::<Vec<_>>();
 
         let poseidon2_cols: &Poseidon2Cols<
             Expression<E>,
@@ -330,19 +387,29 @@ impl<
             .unwrap()
     }
 
+    fn num_p3_cols(&self) -> usize {
+        self.p3_cols.len()
+    }
+
+    pub fn num_cols(&self) -> usize {
+        self.p3_cols.len() + self.post_linear_layer_cols.len()
+    }
+
     pub fn assign_instance(
         &self,
         instance: &mut [E::BaseField],
         state: [E::BaseField; STATE_WIDTH],
     ) {
+        let (p3_cols, post_linear_layer_cols) = instance.split_at_mut(self.num_p3_cols());
+
         let poseidon2_cols: &mut Poseidon2Cols<
             E::BaseField,
             STATE_WIDTH,
             SBOX_DEGREE,
             SBOX_REGISTERS,
             HALF_FULL_ROUNDS,
             PARTIAL_ROUNDS,
-        > = instance.borrow_mut();
+        > = p3_cols.borrow_mut();
 
         generate_trace_rows_for_perm::<
             E::BaseField,
@@ -352,7 +419,12 @@ impl<
             SBOX_REGISTERS,
             HALF_FULL_ROUNDS,
             PARTIAL_ROUNDS,
-        >(poseidon2_cols, state, &self.constants);
+        >(
+            poseidon2_cols,
+            post_linear_layer_cols,
+            state,
+            &self.constants,
+        );
     }
 }
 
@@ -376,6 +448,7 @@ fn generate_trace_rows_for_perm<
         HALF_FULL_ROUNDS,
         PARTIAL_ROUNDS,
     >,
+    post_linear_layers: &mut [F],
     mut state: [F; WIDTH],
     constants: &RoundConstants<F, WIDTH, HALF_FULL_ROUNDS, PARTIAL_ROUNDS>,
 ) {
@@ -389,6 +462,15 @@ fn generate_trace_rows_for_perm<
 
     LinearLayers::external_linear_layer(&mut state);
 
+    // 1. before 0th full round
+    // post_linear_layer[i] = state[i]
+    post_linear_layers[0..WIDTH]
+        .iter_mut()
+        .zip(state.iter())
+        .for_each(|(post, &x)| {
+            *post = x;
+        });
+
     for (full_round, constants) in perm
         .beginning_full_rounds
         .iter_mut()
@@ -399,18 +481,29 @@ fn generate_trace_rows_for_perm<
         );
     }
 
-    for (partial_round, constant) in perm
+    for (i, (partial_round, constant)) in perm
         .partial_rounds
         .iter_mut()
         .zip(&constants.partial_round_constants)
+        .enumerate()
     {
         generate_partial_round::<F, LinearLayers, WIDTH, SBOX_DEGREE, SBOX_REGISTERS>(
             &mut state,
+            &mut post_linear_layers[WIDTH + i],
             partial_round,
             *constant,
         );
     }
 
+    // 3. before HALF_FULL_ROUNDS-th full round
+    // post_linear_layer[i] = state[i]
+    post_linear_layers[WIDTH + PARTIAL_ROUNDS..]
+        .iter_mut()
+        .zip(state.iter())
+        .for_each(|(post, &x)| {
+            *post = x;
+        });
+
     for (full_round, constants) in perm
         .ending_full_rounds
         .iter_mut()
@@ -459,10 +552,12 @@ fn generate_partial_round<
     const SBOX_REGISTERS: usize,
 >(
     state: &mut [F; WIDTH],
+    post_linear_layer: &mut F,
     partial_round: &mut PartialRound<F, WIDTH, SBOX_DEGREE, SBOX_REGISTERS>,
     round_constant: F,
 ) {
     state[0] += round_constant;
+    *post_linear_layer = state[0];
     generate_sbox(&mut partial_round.sbox, &mut state[0]);
     partial_round.post_sbox = state[0];
     LinearLayers::internal_linear_layer(state);

ceno_zkvm/src/tables/shard_ram.rs

@@ -690,7 +690,7 @@ mod tests {
             ShardRamCircuit::build_gkr_iop_circuit(&mut cb, &ProgramParams::default()).unwrap();
 
         // create a bunch of random memory read/write records
-        let n_global_reads = 1700;
+        let n_global_reads = 170000;
         let n_global_writes = 1420;
         let global_reads = (0..n_global_reads)
             .map(|i| {

gkr_iop/src/gkr/layer/zerocheck_layer.rs

@@ -165,11 +165,18 @@ impl<E: ExtensionField> ZerocheckLayer<E> for Layer<E> {
             self.n_fixed as WitnessId,
             self.n_instance,
         );
+        tracing::debug!("main sumcheck degree: {}", zero_expr.degree());
         self.main_sumcheck_expression = Some(zero_expr);
         self.main_sumcheck_expression_monomial_terms = self
             .main_sumcheck_expression
             .as_ref()
             .map(|expr| expr.get_monomial_terms());
+        tracing::debug!(
+            "main sumcheck monomial terms count: {}",
+            self.main_sumcheck_expression_monomial_terms
+                .as_ref()
+                .map_or(0, |terms| terms.len()),
+        );
         exit_span!(span);
     }