fix/program-size2: refactor padding_zero (#615)

Follow-up to #611. --------- Co-authored-by: Aurélien Nicolas <[email protected]>
scroll-tech · Nov 21, 2024 · c548dc8 · c548dc8
1 parent a1e7462
commit c548dc8
Show file tree

Hide file tree

Showing 4 changed files with 39 additions and 43 deletions.
diff --git a/ceno_zkvm/src/instructions.rs b/ceno_zkvm/src/instructions.rs
@@ -94,7 +94,7 @@ pub trait Instruction<E: ExtensionField> {
                 num_padding_instances
             };
             raw_witin
-                .par_batch_iter_padding_mut(num_padding_instance_per_batch)
+                .par_batch_iter_padding_mut(None, num_padding_instance_per_batch)
                 .with_min_len(MIN_PAR_SIZE)
                 .for_each(|row| {
                     row.chunks_mut(num_witin)

diff --git a/ceno_zkvm/src/tables/mod.rs b/ceno_zkvm/src/tables/mod.rs
@@ -2,8 +2,8 @@ use crate::{
     circuit_builder::CircuitBuilder, error::ZKVMError, scheme::constants::MIN_PAR_SIZE,
     witness::RowMajorMatrix,
 };
-use ff::Field;
 use ff_ext::ExtensionField;
+use goldilocks::SmallField;
 use rayon::iter::{IndexedParallelIterator, ParallelIterator};
 use std::{collections::HashMap, mem::MaybeUninit};
 mod range;
@@ -46,25 +46,34 @@ pub trait TableCircuit<E: ExtensionField> {
         table: &mut RowMajorMatrix<E::BaseField>,
         num_witin: usize,
     ) -> Result<(), ZKVMError> {
-        // Fill the padding with zeros, if any.
-        let num_padding_instances = table.num_padding_instances();
-        if num_padding_instances > 0 {
-            let nthreads =
-                std::env::var("RAYON_NUM_THREADS").map_or(8, |s| s.parse::<usize>().unwrap_or(8));
-            let padding_instance = vec![MaybeUninit::new(E::BaseField::ZERO); num_witin];
-            let num_padding_instance_per_batch = if num_padding_instances > 256 {
-                num_padding_instances.div_ceil(nthreads)
-            } else {
-                num_padding_instances
-            };
-            table
-                .par_batch_iter_padding_mut(num_padding_instance_per_batch)
-                .with_min_len(MIN_PAR_SIZE)
-                .for_each(|row| {
-                    row.chunks_mut(num_witin)
-                        .for_each(|instance| instance.copy_from_slice(padding_instance.as_slice()));
-                });
-        }
+        padding_zero(table, num_witin, None);
         Ok(())
     }
 }
+
+/// Fill the padding with zeros. Start after the given `num_instances`, or detect it from the table.
+pub fn padding_zero<F: SmallField>(
+    table: &mut RowMajorMatrix<F>,
+    num_cols: usize,
+    num_instances: Option<usize>,
+) {
+    // Fill the padding with zeros, if any.
+    let num_padding_instances = table.num_padding_instances();
+    if num_padding_instances > 0 {
+        let nthreads =
+            std::env::var("RAYON_NUM_THREADS").map_or(8, |s| s.parse::<usize>().unwrap_or(8));
+        let padding_instance = vec![MaybeUninit::new(F::ZERO); num_cols];
+        let num_padding_instance_per_batch = if num_padding_instances > 256 {
+            num_padding_instances.div_ceil(nthreads)
+        } else {
+            num_padding_instances
+        };
+        table
+            .par_batch_iter_padding_mut(num_instances, num_padding_instance_per_batch)
+            .with_min_len(MIN_PAR_SIZE)
+            .for_each(|row| {
+                row.chunks_mut(num_cols)
+                    .for_each(|instance| instance.copy_from_slice(padding_instance.as_slice()));
+            });
+    }
+}
diff --git a/ceno_zkvm/src/tables/program.rs b/ceno_zkvm/src/tables/program.rs
@@ -7,7 +7,7 @@ use crate::{
     scheme::constants::MIN_PAR_SIZE,
     set_fixed_val, set_val,
     structs::ROMType,
-    tables::TableCircuit,
+    tables::{TableCircuit, padding_zero},
     utils::i64_to_base,
     witness::RowMajorMatrix,
 };
@@ -175,16 +175,8 @@ impl<E: ExtensionField> TableCircuit<E> for ProgramTableCircuit<E> {
                 }
             });
 
-        assert_eq!(INVALID as u64, 0, "cannot use 0 as program padding");
-        fixed
-            .par_iter_mut()
-            .with_min_len(MIN_PAR_SIZE)
-            .skip(num_instructions)
-            .for_each(|row| {
-                for col in config.record.as_slice() {
-                    set_fixed_val!(row, *col, 0_u64.into());
-                }
-            });
+        assert_eq!(INVALID as u64, 0, "0 padding must be invalid instructions");
+        padding_zero(&mut fixed, num_fixed, Some(num_instructions));
 
         fixed
     }
@@ -212,13 +204,7 @@ impl<E: ExtensionField> TableCircuit<E> for ProgramTableCircuit<E> {
                 set_val!(row, config.mlt, E::BaseField::from(mlt as u64));
             });
 
-        witness
-            .par_iter_mut()
-            .with_min_len(MIN_PAR_SIZE)
-            .skip(program.instructions.len())
-            .for_each(|row| {
-                set_val!(row, config.mlt, 0_u64);
-            });
+        padding_zero(&mut witness, num_witin, Some(program.instructions.len()));
 
         Ok(witness)
     }

diff --git a/ceno_zkvm/src/witness.rs b/ceno_zkvm/src/witness.rs
@@ -87,12 +87,13 @@ impl<T: Sized + Sync + Clone + Send + Copy> RowMajorMatrix<T> {
 
     pub fn par_batch_iter_padding_mut(
         &mut self,
-        num_rows: usize,
+        num_instances: Option<usize>,
+        batch_size: usize,
     ) -> rayon::slice::ChunksMut<'_, MaybeUninit<T>> {
-        let valid_instance = self.num_instances();
-        self.values[valid_instance * self.num_col..]
+        let num_instances = num_instances.unwrap_or(self.num_instances());
+        self.values[num_instances * self.num_col..]
             .as_mut()
-            .par_chunks_mut(num_rows * self.num_col)
+            .par_chunks_mut(batch_size * self.num_col)
     }
 
     pub fn de_interleaving(mut self) -> Vec<Vec<T>> {