Merge branch 'master' into matthias/simplify-allocator

scroll-tech · Dec 12, 2024 · f03b5d0 · f03b5d0
2 parents 6f6d1bb + 979ee30
commit f03b5d0
Show file tree

Hide file tree

Showing 56 changed files with 312 additions and 370 deletions.
diff --git a/.github/workflows/integration.yml b/.github/workflows/integration.yml
@@ -32,13 +32,6 @@ jobs:
     timeout-minutes: 30
     runs-on: [self-hosted, Linux, X64]
 
-    strategy:
-      matrix:
-        target: [x86_64-unknown-linux-gnu, riscv32im-unknown-none-elf]
-        # Exclude the riscv32im-unknown-none-elf target
-        exclude:
-          - target: riscv32im-unknown-none-elf
-
     steps:
       - uses: actions/checkout@v2
       - name: Cargo cache
@@ -57,11 +50,11 @@ jobs:
         env:
           RAYON_NUM_THREADS: 2
           RUSTFLAGS: "-C opt-level=3"
-        run: cargo run --package ceno_zkvm --example riscv_opcodes --target ${{ matrix.target }} -- --start 10 --end 11
+        run: cargo run --package ceno_zkvm --example riscv_opcodes -- --start 10 --end 11
 
       - name: Run fibonacci
         env:
           RAYON_NUM_THREADS: 8
           RUST_LOG: debug
           RUSTFLAGS: "-C opt-level=3"
-        run: cargo run --package ceno_zkvm --bin e2e --target ${{ matrix.target }} -- --platform=sp1 ceno_zkvm/examples/fibonacci.elf
+        run: cargo run --package ceno_zkvm --bin e2e -- --platform=sp1 ceno_zkvm/examples/fibonacci.elf
diff --git a/.github/workflows/lints.yml b/.github/workflows/lints.yml
@@ -32,13 +32,6 @@ jobs:
     timeout-minutes: 30
     runs-on: [self-hosted, Linux, X64]
 
-    strategy:
-      matrix:
-        target: [x86_64-unknown-linux-gnu, riscv32im-unknown-none-elf]
-        # Exclude the riscv32im-unknown-none-elf target
-        exclude:
-          - target: riscv32im-unknown-none-elf
-
     steps:
       - uses: actions/checkout@v2
       - uses: dtolnay/rust-toolchain@master
@@ -67,7 +60,6 @@ jobs:
 
       - name: Run clippy
         env:
-          TARGET: ${{ matrix.target }}
           RUSTFLAGS: "-Dwarnings"
         run: |
           cargo check --workspace --all-targets --exclude ceno_rt

diff --git a/.github/workflows/tests.yml b/.github/workflows/tests.yml
@@ -31,13 +31,6 @@ jobs:
     timeout-minutes: 30
     runs-on: [self-hosted, Linux, X64]
 
-    strategy:
-      matrix:
-        target: [x86_64-unknown-linux-gnu, riscv32im-unknown-none-elf]
-        # Exclude the riscv32im-unknown-none-elf target
-        exclude:
-          - target: riscv32im-unknown-none-elf
-
     steps:
       - uses: actions/checkout@v2
       - uses: dtolnay/rust-toolchain@master
@@ -61,6 +54,4 @@ jobs:
         run: |
           cargo make --version || cargo install cargo-make
       - name: run test
-        env:
-          TARGET: ${{ matrix.target }}
         run: cargo make tests
diff --git a/Makefile.toml b/Makefile.toml
@@ -1,14 +1,6 @@
 [env]
 CARGO_MAKE_EXTEND_WORKSPACE_MAKEFILE = true
 CORE = { script = ["nproc"] }
-CUR_TARGET = { script = ['''
-        if [ -z "${TARGET}" ]; then
-            TARGET=$(rustc -vV | grep "host" | awk '{print $2}')
-            echo "${TARGET}"
-        else
-            echo "${TARGET}"
-        fi
-        '''] }
 RAYON_NUM_THREADS = "${CORE}"
 
 [tasks.tests]
@@ -20,8 +12,6 @@ args = [
   "--tests",
   "--examples",
   "--release",
-  "--target",
-  "${CUR_TARGET}",
   "--workspace",
   "--exclude",
   "ceno_rt",
@@ -46,8 +36,6 @@ args = [
   "--all-targets",
   "--exclude",
   "ceno_rt",
-  "--target",
-  "${CUR_TARGET}",
   "--",
   "-D",
   "warnings",

diff --git a/ceno_zkvm/src/expression.rs b/ceno_zkvm/src/expression.rs
@@ -4,7 +4,6 @@ use std::{
     cmp::max,
     fmt::Display,
     iter::{Product, Sum},
-    mem::MaybeUninit,
     ops::{Add, AddAssign, Deref, Mul, MulAssign, Neg, Shl, ShlAssign, Sub, SubAssign},
 };
 
@@ -591,17 +590,17 @@ macro_rules! mixed_binop_instances {
 mixed_binop_instances!(
     Add,
     add,
-    (u8, u16, u32, u64, usize, i8, i16, i32, i64, i128, isize)
+    (u8, u16, u32, u64, usize, i8, i16, i32, i64, isize)
 );
 mixed_binop_instances!(
     Sub,
     sub,
-    (u8, u16, u32, u64, usize, i8, i16, i32, i64, i128, isize)
+    (u8, u16, u32, u64, usize, i8, i16, i32, i64, isize)
 );
 mixed_binop_instances!(
     Mul,
     mul,
-    (u8, u16, u32, u64, usize, i8, i16, i32, i64, i128, isize)
+    (u8, u16, u32, u64, usize, i8, i16, i32, i64, isize)
 );
 
 impl<E: ExtensionField> Mul for Expression<E> {
@@ -756,12 +755,8 @@ impl WitIn {
         )
     }
 
-    pub fn assign<E: ExtensionField>(
-        &self,
-        instance: &mut [MaybeUninit<E::BaseField>],
-        value: E::BaseField,
-    ) {
-        instance[self.id as usize] = MaybeUninit::new(value);
+    pub fn assign<E: ExtensionField>(&self, instance: &mut [E::BaseField], value: E::BaseField) {
+        instance[self.id as usize] = value;
     }
 }
 
@@ -840,14 +835,6 @@ macro_rules! impl_from_unsigned {
 }
 impl_from_unsigned!(u8, u16, u32, u64, usize, RAMType, InsnKind);
 
-// Implement From trait for u128 separately since it requires explicit reduction
-impl<F: SmallField, E: ExtensionField<BaseField = F>> From<u128> for Expression<E> {
-    fn from(value: u128) -> Self {
-        let reduced = value.rem_euclid(F::MODULUS_U64 as u128) as u64;
-        Expression::Constant(F::from(reduced))
-    }
-}
-
 // Implement From trait for signed types
 macro_rules! impl_from_signed {
     ($($t:ty),*) => {
@@ -861,7 +848,7 @@ macro_rules! impl_from_signed {
         )*
     };
 }
-impl_from_signed!(i8, i16, i32, i64, i128, isize);
+impl_from_signed!(i8, i16, i32, i64, isize);
 
 impl<E: ExtensionField> Display for Expression<E> {
     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {

diff --git a/ceno_zkvm/src/gadgets/div.rs b/ceno_zkvm/src/gadgets/div.rs
@@ -1,4 +1,4 @@
-use std::{fmt::Display, mem::MaybeUninit};
+use std::fmt::Display;
 
 use ff_ext::ExtensionField;
 
@@ -53,7 +53,7 @@ impl<E: ExtensionField> DivConfig<E> {
 
     pub fn assign_instance<'a>(
         &self,
-        instance: &mut [MaybeUninit<E::BaseField>],
+        instance: &mut [E::BaseField],
         lkm: &mut LkMultiplicity,
         divisor: &Value<'a, u32>,
         quotient: &Value<'a, u32>,

diff --git a/ceno_zkvm/src/gadgets/is_lt.rs b/ceno_zkvm/src/gadgets/is_lt.rs
@@ -1,4 +1,4 @@
-use std::{fmt::Display, mem::MaybeUninit};
+use std::fmt::Display;
 
 use ceno_emul::{SWord, Word};
 use ff_ext::ExtensionField;
@@ -52,7 +52,7 @@ impl AssertLtConfig {
 
     pub fn assign_instance<F: SmallField>(
         &self,
-        instance: &mut [MaybeUninit<F>],
+        instance: &mut [F],
         lkm: &mut LkMultiplicity,
         lhs: u64,
         rhs: u64,
@@ -106,7 +106,7 @@ impl IsLtConfig {
 
     pub fn assign_instance<F: SmallField>(
         &self,
-        instance: &mut [MaybeUninit<F>],
+        instance: &mut [F],
         lkm: &mut LkMultiplicity,
         lhs: u64,
         rhs: u64,
@@ -118,7 +118,7 @@ impl IsLtConfig {
 
     pub fn assign_instance_signed<F: SmallField>(
         &self,
-        instance: &mut [MaybeUninit<F>],
+        instance: &mut [F],
         lkm: &mut LkMultiplicity,
         lhs: SWord,
         rhs: SWord,
@@ -184,7 +184,7 @@ impl InnerLtConfig {
 
     pub fn assign_instance<F: SmallField>(
         &self,
-        instance: &mut [MaybeUninit<F>],
+        instance: &mut [F],
         lkm: &mut LkMultiplicity,
         lhs: u64,
         rhs: u64,
@@ -202,7 +202,7 @@ impl InnerLtConfig {
     // TODO: refactor with the above function
     pub fn assign_instance_signed<F: SmallField>(
         &self,
-        instance: &mut [MaybeUninit<F>],
+        instance: &mut [F],
         lkm: &mut LkMultiplicity,
         lhs: SWord,
         rhs: SWord,
@@ -256,7 +256,7 @@ impl<E: ExtensionField> AssertSignedLtConfig<E> {
 
     pub fn assign_instance(
         &self,
-        instance: &mut [MaybeUninit<E::BaseField>],
+        instance: &mut [E::BaseField],
         lkm: &mut LkMultiplicity,
         lhs: SWord,
         rhs: SWord,
@@ -299,7 +299,7 @@ impl<E: ExtensionField> SignedLtConfig<E> {
 
     pub fn assign_instance(
         &self,
-        instance: &mut [MaybeUninit<E::BaseField>],
+        instance: &mut [E::BaseField],
         lkm: &mut LkMultiplicity,
         lhs: SWord,
         rhs: SWord,
@@ -351,7 +351,7 @@ impl<E: ExtensionField> InnerSignedLtConfig<E> {
 
     pub fn assign_instance(
         &self,
-        instance: &mut [MaybeUninit<E::BaseField>],
+        instance: &mut [E::BaseField],
         lkm: &mut LkMultiplicity,
         lhs: SWord,
         rhs: SWord,

diff --git a/ceno_zkvm/src/gadgets/is_zero.rs b/ceno_zkvm/src/gadgets/is_zero.rs
@@ -1,5 +1,3 @@
-use std::mem::MaybeUninit;
-
 use ff_ext::ExtensionField;
 use goldilocks::SmallField;
 
@@ -64,7 +62,7 @@ impl IsZeroConfig {
 
     pub fn assign_instance<F: SmallField>(
         &self,
-        instance: &mut [MaybeUninit<F>],
+        instance: &mut [F],
         x: F,
     ) -> Result<(), ZKVMError> {
         let (is_zero, inverse) = if x.is_zero_vartime() {
@@ -117,7 +115,7 @@ impl IsEqualConfig {
 
     pub fn assign_instance<F: SmallField>(
         &self,
-        instance: &mut [MaybeUninit<F>],
+        instance: &mut [F],
         a: F,
         b: F,
     ) -> Result<(), ZKVMError> {

diff --git a/ceno_zkvm/src/gadgets/signed_ext.rs b/ceno_zkvm/src/gadgets/signed_ext.rs
@@ -7,7 +7,7 @@ use crate::{
     witness::LkMultiplicity,
 };
 use ff_ext::ExtensionField;
-use std::{marker::PhantomData, mem::MaybeUninit};
+use std::marker::PhantomData;
 
 #[derive(Debug)]
 pub struct SignedExtendConfig<E> {
@@ -84,7 +84,7 @@ impl<E: ExtensionField> SignedExtendConfig<E> {
 
     pub fn assign_instance(
         &self,
-        instance: &mut [MaybeUninit<E::BaseField>],
+        instance: &mut [E::BaseField],
         lk_multiplicity: &mut LkMultiplicity,
         val: u64,
     ) -> Result<(), ZKVMError> {

diff --git a/ceno_zkvm/src/instructions.rs b/ceno_zkvm/src/instructions.rs
@@ -1,26 +1,31 @@
-use std::mem::MaybeUninit;
-
 use ceno_emul::StepRecord;
 use ff_ext::ExtensionField;
 use multilinear_extensions::util::max_usable_threads;
 use rayon::{
     iter::{IndexedParallelIterator, ParallelIterator},
     slice::ParallelSlice,
 };
+use std::sync::Arc;
 
 use crate::{
     circuit_builder::CircuitBuilder,
     error::ZKVMError,
-    scheme::constants::MIN_PAR_SIZE,
     witness::{LkMultiplicity, RowMajorMatrix},
 };
-use ff::Field;
 
 pub mod riscv;
 
+#[derive(Clone)]
 pub enum InstancePaddingStrategy {
-    Zero,
+    // Pads with default values of underlying type
+    // Usually zero, but check carefully
+    Default,
+    // Pads by repeating last row
     RepeatLast,
+    // Custom strategy consists of a closure
+    // `pad(i, j) = padding value for cell at row i, column j`
+    // pad should be able to cross thread boundaries
+    Custom(Arc<dyn Fn(u64, u64) -> u64 + Send + Sync>),
 }
 
 pub trait Instruction<E: ExtensionField> {
@@ -38,7 +43,7 @@ pub trait Instruction<E: ExtensionField> {
     // assign single instance giving step from trace
     fn assign_instance(
         config: &Self::InstructionConfig,
-        instance: &mut [MaybeUninit<E::BaseField>],
+        instance: &mut [E::BaseField],
         lk_multiplicity: &mut LkMultiplicity,
         step: &StepRecord,
     ) -> Result<(), ZKVMError>;
@@ -56,7 +61,8 @@ pub trait Instruction<E: ExtensionField> {
         }
         .max(1);
         let lk_multiplicity = LkMultiplicity::default();
-        let mut raw_witin = RowMajorMatrix::<E::BaseField>::new(steps.len(), num_witin);
+        let mut raw_witin =
+            RowMajorMatrix::<E::BaseField>::new(steps.len(), num_witin, Self::padding_strategy());
         let raw_witin_iter = raw_witin.par_batch_iter_mut(num_instance_per_batch);
 
         raw_witin_iter
@@ -73,35 +79,6 @@ pub trait Instruction<E: ExtensionField> {
             })
             .collect::<Result<(), ZKVMError>>()?;
 
-        let num_padding_instances = raw_witin.num_padding_instances();
-        if num_padding_instances > 0 {
-            // Fill the padding based on strategy
-
-            let padding_instance = match Self::padding_strategy() {
-                InstancePaddingStrategy::Zero => {
-                    vec![MaybeUninit::new(E::BaseField::ZERO); num_witin]
-                }
-                InstancePaddingStrategy::RepeatLast if steps.is_empty() => {
-                    tracing::debug!("No {} steps to repeat, using zero padding", Self::name());
-                    vec![MaybeUninit::new(E::BaseField::ZERO); num_witin]
-                }
-                InstancePaddingStrategy::RepeatLast => raw_witin[steps.len() - 1].to_vec(),
-            };
-
-            let num_padding_instance_per_batch = if num_padding_instances > 256 {
-                num_padding_instances.div_ceil(nthreads)
-            } else {
-                num_padding_instances
-            };
-            raw_witin
-                .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()));
-                });
-        }
-
         Ok((raw_witin, lk_multiplicity))
     }
 }