use basefield in witness instead of extfield

ceno_emul/src/tracer.rs

@@ -19,16 +19,16 @@ use crate::{
 /// - Any pair of `rs1 / rs2 / rd` **may be the same**. Then, one op will point to the other op in the same instruction but a different subcycle. The circuits may follow the operations **without special handling** of repeated registers.
 #[derive(Clone, Debug, Default)]
 pub struct StepRecord {
-    cycle: Cycle,
-    pc: Change<ByteAddr>,
-    insn_code: Word,
+    pub cycle: Cycle,
+    pub pc: Change<ByteAddr>,
+    pub insn_code: Word,
 
-    rs1: Option<ReadOp>,
-    rs2: Option<ReadOp>,
+    pub rs1: Option<ReadOp>,
+    pub rs2: Option<ReadOp>,
 
-    rd: Option<WriteOp>,
+    pub rd: Option<WriteOp>,
 
-    memory_op: Option<WriteOp>,
+    pub memory_op: Option<WriteOp>,
 }
 
 #[derive(Clone, Debug, PartialEq, Eq, Hash)]

ceno_zkvm/src/instructions.rs

@@ -17,16 +17,16 @@ pub trait Instruction<E: ExtensionField> {
     // assign single instance giving step from trace
     fn assign_instance(
         config: &Self::InstructionConfig,
-        instance: &mut [MaybeUninit<E>],
+        instance: &mut [MaybeUninit<E::BaseField>],
         step: StepRecord,
     ) -> Result<(), ZKVMError>;
 
     fn assign_instances(
         config: &Self::InstructionConfig,
         num_witin: usize,
         steps: Vec<StepRecord>,
-    ) -> Result<RowMajorMatrix<E>, ZKVMError> {
-        let mut raw_witin = RowMajorMatrix::<E>::new(steps.len(), num_witin);
+    ) -> Result<RowMajorMatrix<E::BaseField>, ZKVMError> {
+        let mut raw_witin = RowMajorMatrix::<E::BaseField>::new(steps.len(), num_witin);
         let raw_witin_iter = raw_witin.par_iter_mut();
 
         raw_witin_iter

ceno_zkvm/src/instructions/riscv/addsub.rs

@@ -1,7 +1,9 @@
 use std::marker::PhantomData;
 
+use ark_std::iterable::Iterable;
 use ceno_emul::StepRecord;
 use ff_ext::ExtensionField;
+use itertools::Itertools;
 
 use super::{
     constants::{OPType, OpcodeType, RegUInt, PC_STEP_SIZE},
@@ -14,6 +16,7 @@ use crate::{
     expression::{ToExpr, WitIn},
     instructions::Instruction,
     set_val,
+    uint::UIntValue,
 };
 use core::mem::MaybeUninit;
 
@@ -147,28 +150,31 @@ impl<E: ExtensionField> Instruction<E> for AddInstruction {
     #[allow(clippy::option_map_unit_fn)]
     fn assign_instance(
         config: &Self::InstructionConfig,
-        instance: &mut [MaybeUninit<E>],
-        _step: StepRecord,
+        instance: &mut [MaybeUninit<E::BaseField>],
+        step: StepRecord,
     ) -> Result<(), ZKVMError> {
-        // TODO use field from step
+        // TODO use fields from step
         set_val!(instance, config.pc, 1);
         set_val!(instance, config.ts, 2);
-        config.prev_rd_value.wits_in().map(|prev_rd_value| {
-            set_val!(instance, prev_rd_value[0], 4);
-            set_val!(instance, prev_rd_value[1], 4);
-        });
-        config.addend_0.wits_in().map(|addend_0| {
-            set_val!(instance, addend_0[0], 4);
-            set_val!(instance, addend_0[1], 4);
-        });
-        config.addend_1.wits_in().map(|addend_1| {
-            set_val!(instance, addend_1[0], 4);
-            set_val!(instance, addend_1[1], 4);
-        });
-        config.outcome.carries.as_ref().map(|carry| {
-            set_val!(instance, carry[0], 0);
-            set_val!(instance, carry[1], 0);
-        });
+        let addend_0 = UIntValue(step.rs1().unwrap().value);
+        let addend_1 = UIntValue(step.rs2().unwrap().value);
+        config
+            .prev_rd_value
+            .assign_limbs(instance, [0, 0].iter().map(E::BaseField::from).collect());
+        config
+            .addend_0
+            .assign_limbs(instance, addend_0.as_u16_fields());
+        config
+            .addend_1
+            .assign_limbs(instance, addend_1.as_u16_fields());
+        let carries = addend_0.add_u16_carries(&addend_1);
+        config.outcome.assign_carries(
+            instance,
+            carries
+                .into_iter()
+                .map(|carry| E::BaseField::from(carry as u64))
+                .collect_vec(),
+        );
         // TODO #167
         set_val!(instance, config.rs1_id, 2);
         set_val!(instance, config.rs2_id, 2);
@@ -192,7 +198,7 @@ impl<E: ExtensionField> Instruction<E> for SubInstruction {
     #[allow(clippy::option_map_unit_fn)]
     fn assign_instance(
         config: &Self::InstructionConfig,
-        instance: &mut [MaybeUninit<E>],
+        instance: &mut [MaybeUninit<E::BaseField>],
         _step: StepRecord,
     ) -> Result<(), ZKVMError> {
         // TODO use field from step
@@ -229,14 +235,16 @@ impl<E: ExtensionField> Instruction<E> for SubInstruction {
 #[cfg(test)]
 mod test {
 
-    use ceno_emul::StepRecord;
+    use std::u32;
+
+    use ceno_emul::{ReadOp, StepRecord};
     use goldilocks::GoldilocksExt2;
     use itertools::Itertools;
     use multilinear_extensions::mle::IntoMLEs;
 
     use crate::{
         circuit_builder::{CircuitBuilder, ConstraintSystem},
-        instructions::{riscv::constants::RegUInt, Instruction},
+        instructions::Instruction,
         scheme::mock_prover::MockProver,
     };
 
@@ -261,7 +269,19 @@ mod test {
         let raw_witin = AddInstruction::assign_instances(
             &config,
             cb.cs.num_witin as usize,
-            vec![StepRecord::default()],
+            vec![StepRecord {
+                rs1: Some(ReadOp {
+                    addr: 0.into(),
+                    value: 11u32,
+                    previous_cycle: 0,
+                }),
+                rs2: Some(ReadOp {
+                    addr: 0.into(),
+                    value: 5u32,
+                    previous_cycle: 0,
+                }),
+                ..Default::default()
+            }],
         )
         .unwrap();
 
@@ -296,7 +316,19 @@ mod test {
         let raw_witin = AddInstruction::assign_instances(
             &config,
             cb.cs.num_witin as usize,
-            vec![StepRecord::default()],
+            vec![StepRecord {
+                rs1: Some(ReadOp {
+                    addr: 0.into(),
+                    value: u32::MAX - 1,
+                    previous_cycle: 0,
+                }),
+                rs2: Some(ReadOp {
+                    addr: 0.into(),
+                    value: u32::MAX - 1,
+                    previous_cycle: 0,
+                }),
+                ..Default::default()
+            }],
         )
         .unwrap();
 

ceno_zkvm/src/instructions/riscv/blt.rs

@@ -12,8 +12,8 @@ use crate::{
         riscv::config::{LtConfig, LtInput},
         Instruction,
     },
-    set_val,
-    utils::{i64_to_ext, limb_u8_to_u16},
+    set_field_val,
+    utils::{i64_to_base_field, limb_u8_to_u16},
 };
 
 use super::{
@@ -56,58 +56,68 @@ impl BltInput {
     pub fn assign<E: ExtensionField>(
         &self,
         config: &InstructionConfig<E>,
-        instance: &mut [MaybeUninit<E>],
+        instance: &mut [MaybeUninit<E::BaseField>],
     ) {
         assert!(!self.lhs_limb8.is_empty() && (self.lhs_limb8.len() == self.rhs_limb8.len()));
         // TODO: add boundary check for witin
         let lt_input = LtInput {
             lhs_limbs: &self.lhs_limb8,
             rhs_limbs: &self.rhs_limb8,
         };
-        let is_lt = lt_input.assign(instance, &config.is_lt);
+        let is_lt = lt_input.assign::<E>(instance, &config.is_lt);
 
-        set_val!(instance, config.pc, { i64_to_ext::<E>(self.pc as i64) });
-        set_val!(instance, config.next_pc, {
+        set_field_val!(instance, config.pc, {
+            i64_to_base_field::<E>(self.pc as i64)
+        });
+        set_field_val!(instance, config.next_pc, {
             if is_lt {
-                i64_to_ext::<E>(self.pc as i64 + self.imm as i64)
+                i64_to_base_field::<E>(self.pc as i64 + self.imm as i64)
             } else {
-                i64_to_ext::<E>(self.pc as i64 + PC_STEP_SIZE as i64)
+                i64_to_base_field::<E>(self.pc as i64 + PC_STEP_SIZE as i64)
             }
         });
-        set_val!(instance, config.ts, { i64_to_ext::<E>(self.ts as i64) });
-        set_val!(instance, config.imm, { i64_to_ext::<E>(self.imm as i64) });
-        set_val!(instance, config.rs1_id, {
-            i64_to_ext::<E>(self.rs1_id as i64)
+        set_field_val!(instance, config.ts, {
+            i64_to_base_field::<E>(self.ts as i64)
+        });
+        set_field_val!(instance, config.imm, {
+            i64_to_base_field::<E>(self.imm as i64)
+        });
+        set_field_val!(instance, config.rs1_id, {
+            i64_to_base_field::<E>(self.rs1_id as i64)
         });
-        set_val!(instance, config.rs2_id, {
-            i64_to_ext::<E>(self.rs2_id as i64)
+        set_field_val!(instance, config.rs2_id, {
+            i64_to_base_field::<E>(self.rs2_id as i64)
         });
-        set_val!(instance, config.prev_rs1_ts, {
-            i64_to_ext::<E>(self.prev_rs1_ts as i64)
+        set_field_val!(instance, config.prev_rs1_ts, {
+            i64_to_base_field::<E>(self.prev_rs1_ts as i64)
         });
-        set_val!(instance, config.prev_rs2_ts, {
-            i64_to_ext::<E>(self.prev_rs2_ts as i64)
+        set_field_val!(instance, config.prev_rs2_ts, {
+            i64_to_base_field::<E>(self.prev_rs2_ts as i64)
         });
 
-        config.lhs_limb8.assign(instance, {
+        config.lhs_limb8.assign_limbs(instance, {
             self.lhs_limb8
                 .iter()
-                .map(|&limb| i64_to_ext(limb as i64))
+                .map(|&limb| i64_to_base_field::<E>(limb as i64))
                 .collect()
         });
-        config.rhs_limb8.assign(instance, {
+        config.rhs_limb8.assign_limbs(instance, {
             self.rhs_limb8
                 .iter()
-                .map(|&limb| i64_to_ext(limb as i64))
+                .map(|&limb| i64_to_base_field::<E>(limb as i64))
                 .collect()
         });
         let lhs = limb_u8_to_u16(&self.lhs_limb8);
         let rhs = limb_u8_to_u16(&self.rhs_limb8);
-        config.lhs.assign(instance, {
-            lhs.iter().map(|&limb| i64_to_ext(limb as i64)).collect()
+        config.lhs.assign_limbs(instance, {
+            lhs.iter()
+                .map(|&limb| i64_to_base_field::<E>(limb as i64))
+                .collect()
         });
-        config.rhs.assign(instance, {
-            rhs.iter().map(|&limb| i64_to_ext(limb as i64)).collect()
+        config.rhs.assign_limbs(instance, {
+            rhs.iter()
+                .map(|&limb| i64_to_base_field::<E>(limb as i64))
+                .collect()
         });
     }
 
@@ -216,7 +226,7 @@ impl<E: ExtensionField> Instruction<E> for BltInstruction {
 
     fn assign_instance(
         config: &Self::InstructionConfig,
-        instance: &mut [std::mem::MaybeUninit<E>],
+        instance: &mut [std::mem::MaybeUninit<E::BaseField>],
         _step: ceno_emul::StepRecord,
     ) -> Result<(), ZKVMError> {
         // take input from _step