Refactor row major matrix

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},
 };
 
@@ -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;
     }
 }
 

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>,

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,

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> {

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> {

ceno_zkvm/src/instructions.rs

@@ -1,26 +1,28 @@
-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,
     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 +40,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 +58,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 +76,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))
     }
 }

ceno_zkvm/src/instructions/riscv/arith.rs

@@ -8,7 +8,6 @@ use crate::{
     circuit_builder::CircuitBuilder, error::ZKVMError, instructions::Instruction, uint::Value,
     witness::LkMultiplicity,
 };
-use core::mem::MaybeUninit;
 
 /// This config handles R-Instructions that represent registers values as 2 * u16.
 #[derive(Debug)]
@@ -88,7 +87,7 @@ impl<E: ExtensionField, I: RIVInstruction> Instruction<E> for ArithInstruction<E
 
     fn assign_instance(
         config: &Self::InstructionConfig,
-        instance: &mut [MaybeUninit<E::BaseField>],
+        instance: &mut [<E as ExtensionField>::BaseField],
         lk_multiplicity: &mut LkMultiplicity,
         step: &StepRecord,
     ) -> Result<(), ZKVMError> {

ceno_zkvm/src/instructions/riscv/arith_imm.rs

@@ -1,4 +1,4 @@
-use std::{marker::PhantomData, mem::MaybeUninit};
+use std::marker::PhantomData;
 
 use ceno_emul::StepRecord;
 use ff_ext::ExtensionField;
@@ -57,7 +57,7 @@ impl<E: ExtensionField> Instruction<E> for AddiInstruction<E> {
 
     fn assign_instance(
         config: &Self::InstructionConfig,
-        instance: &mut [MaybeUninit<E::BaseField>],
+        instance: &mut [<E as ExtensionField>::BaseField],
         lk_multiplicity: &mut LkMultiplicity,
         step: &StepRecord,
     ) -> Result<(), ZKVMError> {

ceno_zkvm/src/instructions/riscv/b_insn.rs

@@ -13,7 +13,6 @@ use crate::{
     utils::i64_to_base,
     witness::LkMultiplicity,
 };
-use core::mem::MaybeUninit;
 
 // Opcode: 1100011
 // Funct3:
@@ -88,7 +87,7 @@ impl<E: ExtensionField> BInstructionConfig<E> {
 
     pub fn assign_instance(
         &self,
-        instance: &mut [MaybeUninit<E::BaseField>],
+        instance: &mut [<E as ExtensionField>::BaseField],
         lk_multiplicity: &mut LkMultiplicity,
         step: &StepRecord,
     ) -> Result<(), ZKVMError> {

ceno_zkvm/src/instructions/riscv/branch/beq_circuit.rs

@@ -1,4 +1,4 @@
-use std::{marker::PhantomData, mem::MaybeUninit};
+use std::marker::PhantomData;
 
 use ceno_emul::{InsnKind, StepRecord};
 use ff_ext::ExtensionField;
@@ -72,7 +72,7 @@ impl<E: ExtensionField, I: RIVInstruction> Instruction<E> for BeqCircuit<E, I> {
 
     fn assign_instance(
         config: &Self::InstructionConfig,
-        instance: &mut [MaybeUninit<E::BaseField>],
+        instance: &mut [<E as ExtensionField>::BaseField],
         lk_multiplicity: &mut LkMultiplicity,
         step: &StepRecord,
     ) -> Result<(), ZKVMError> {

ceno_zkvm/src/instructions/riscv/branch/blt.rs

@@ -1,4 +1,4 @@
-use std::{marker::PhantomData, mem::MaybeUninit};
+use std::marker::PhantomData;
 
 use ff_ext::ExtensionField;
 
@@ -66,7 +66,7 @@ impl<E: ExtensionField, I: RIVInstruction> Instruction<E> for BltCircuit<E, I> {
 
     fn assign_instance(
         config: &Self::InstructionConfig,
-        instance: &mut [MaybeUninit<E::BaseField>],
+        instance: &mut [E::BaseField],
         lk_multiplicity: &mut LkMultiplicity,
         step: &ceno_emul::StepRecord,
     ) -> Result<(), ZKVMError> {

ceno_zkvm/src/instructions/riscv/branch/bltu.rs

@@ -75,7 +75,7 @@ impl<E: ExtensionField, I: RIVInstruction> Instruction<E> for BltuCircuit<E, I>
 
     fn assign_instance(
         config: &Self::InstructionConfig,
-        instance: &mut [std::mem::MaybeUninit<E::BaseField>],
+        instance: &mut [E::BaseField],
         lk_multiplicity: &mut LkMultiplicity,
         step: &ceno_emul::StepRecord,
     ) -> Result<(), ZKVMError> {

ceno_zkvm/src/instructions/riscv/config.rs

@@ -1,5 +1,3 @@
-use std::mem::MaybeUninit;
-
 use crate::{expression::WitIn, set_val, utils::i64_to_base, witness::LkMultiplicity};
 use goldilocks::SmallField;
 use itertools::Itertools;
@@ -25,7 +23,7 @@ pub struct MsbInput<'a> {
 impl MsbInput<'_> {
     pub fn assign<F: SmallField>(
         &self,
-        instance: &mut [MaybeUninit<F>],
+        instance: &mut [F],
         config: &MsbConfig,
         lk_multiplicity: &mut LkMultiplicity,
     ) -> (u8, u8) {
@@ -61,7 +59,7 @@ pub struct UIntLtuInput<'a> {
 impl UIntLtuInput<'_> {
     pub fn assign<F: SmallField>(
         &self,
-        instance: &mut [MaybeUninit<F>],
+        instance: &mut [F],
         config: &UIntLtuConfig,
         lk_multiplicity: &mut LkMultiplicity,
     ) -> bool {
@@ -138,7 +136,7 @@ pub struct UIntLtInput<'a> {
 impl UIntLtInput<'_> {
     pub fn assign<F: SmallField>(
         &self,
-        instance: &mut [MaybeUninit<F>],
+        instance: &mut [F],
         config: &UIntLtConfig,
         lk_multiplicity: &mut LkMultiplicity,
     ) -> bool {

ceno_zkvm/src/instructions/riscv/div.rs

@@ -16,7 +16,6 @@ use crate::{
     uint::Value,
     witness::LkMultiplicity,
 };
-use core::mem::MaybeUninit;
 use std::marker::PhantomData;
 
 pub struct ArithConfig<E: ExtensionField> {
@@ -124,7 +123,7 @@ impl<E: ExtensionField, I: RIVInstruction> Instruction<E> for ArithInstruction<E
 
     fn assign_instance(
         config: &Self::InstructionConfig,
-        instance: &mut [MaybeUninit<E::BaseField>],
+        instance: &mut [E::BaseField],
         lkm: &mut LkMultiplicity,
         step: &StepRecord,
     ) -> Result<(), ZKVMError> {
@@ -173,7 +172,6 @@ mod test {
         use ceno_emul::{Change, InsnKind, StepRecord, Word, encode_rv32};
         use goldilocks::GoldilocksExt2;
         use itertools::Itertools;
-        use multilinear_extensions::mle::IntoMLEs;
         use rand::Rng;
 
         use crate::{
@@ -238,7 +236,6 @@ mod test {
             MockProver::assert_with_expected_errors(
                 &cb,
                 &raw_witin
-                    .de_interleaving()
                     .into_mles()
                     .into_iter()
                     .map(|v| v.into())