feat/#132 divu opcode

ceno_zkvm/src/chip_handler/general.rs

@@ -146,6 +146,29 @@ impl<'a, E: ExtensionField> CircuitBuilder<'a, E> {
         )
     }
 
+    pub fn condition_require_equal<NR, N>(
+        &mut self,
+        name_fn: N,
+        cond: Expression<E>,
+        target: Expression<E>,
+        true_expr: Expression<E>,
+        false_expr: Expression<E>,
+    ) -> Result<(), ZKVMError>
+    where
+        NR: Into<String>,
+        N: FnOnce() -> NR,
+    {
+        // cond * (true_expr) + (1 - cond) * false_expr
+        // => false_expr + cond * true_expr - cond * false_expr
+        self.namespace(
+            || "cond_require_equal",
+            |cb| {
+                let cond_target = false_expr.clone() + cond.clone() * true_expr - cond * false_expr;
+                cb.cs.require_zero(name_fn, target - cond_target)
+            },
+        )
+    }
+
     pub(crate) fn assert_ux<NR, N, const C: usize>(
         &mut self,
         name_fn: N,

ceno_zkvm/src/instructions/riscv.rs

@@ -5,6 +5,7 @@ pub mod blt;
 pub mod branch;
 pub mod config;
 pub mod constants;
+pub mod divu;
 pub mod logic;
 
 mod b_insn;

ceno_zkvm/src/instructions/riscv/config.rs

@@ -189,6 +189,7 @@ impl UIntLtInput<'_> {
     }
 }
 
+// TODO move ExprLtConfig to gadgets
 #[derive(Debug)]
 pub struct ExprLtConfig {
     pub is_lt: Option<WitIn>,

ceno_zkvm/src/instructions/riscv/divu.rs

@@ -0,0 +1,220 @@
+use ceno_emul::{InsnKind, StepRecord};
+use ff_ext::ExtensionField;
+use itertools::Itertools;
+
+use super::{constants::UInt, r_insn::RInstructionConfig, RIVInstruction};
+use crate::{
+    circuit_builder::CircuitBuilder, error::ZKVMError, gadgets::IsZeroConfig,
+    instructions::Instruction, uint::Value, witness::LkMultiplicity,
+};
+use core::mem::MaybeUninit;
+use std::marker::PhantomData;
+
+pub struct ArithConfig<E: ExtensionField> {
+    r_insn: RInstructionConfig<E>,
+
+    dividend: UInt<E>,
+    divisor: UInt<E>,
+    outcome: UInt<E>,
+
+    remainder: UInt<E>,
+    inter_mul_value: UInt<E>,
+    is_zero: IsZeroConfig,
+}
+
+pub struct ArithInstruction<E, I>(PhantomData<(E, I)>);
+
+pub struct DivUOp;
+impl RIVInstruction for DivUOp {
+    const INST_KIND: InsnKind = InsnKind::DIVU;
+}
+pub type DivUInstruction<E> = ArithInstruction<E, DivUOp>;
+
+impl<E: ExtensionField, I: RIVInstruction> Instruction<E> for ArithInstruction<E, I> {
+    type InstructionConfig = ArithConfig<E>;
+
+    fn name() -> String {
+        format!("{:?}", I::INST_KIND)
+    }
+
+    fn construct_circuit(
+        circuit_builder: &mut CircuitBuilder<E>,
+    ) -> Result<Self::InstructionConfig, ZKVMError> {
+        // outcome = dividend / divisor + remainder => dividend = divisor * outcome + r
+        let mut divisor = UInt::new_unchecked(|| "divisor", circuit_builder)?;
+        let mut outcome = UInt::new(|| "outcome", circuit_builder)?;
+        let mut r = UInt::new(|| "remainder", circuit_builder)?;
+
+        let (inter_mul_value, dividend) =
+            divisor.mul_add(|| "dividend", circuit_builder, &mut outcome, &mut r, true)?;
+
+        // div by zero check
+        let is_zero = IsZeroConfig::construct_circuit(circuit_builder, divisor.value())?;
+
+        let outcome_value = outcome.value();
+        circuit_builder
+            .condition_require_equal(
+                || "outcome_is_zero",
+                is_zero.expr(),
+                outcome_value.clone(),
+                ((1 << UInt::<E>::M) - 1).into(),
+                outcome_value,
+            )
+            .unwrap();
+
+        let r_insn = RInstructionConfig::<E>::construct_circuit(
+            circuit_builder,
+            I::INST_KIND,
+            dividend.register_expr(),
+            divisor.register_expr(),
+            outcome.register_expr(),
+        )?;
+
+        Ok(ArithConfig {
+            r_insn,
+            dividend,
+            divisor,
+            outcome,
+            remainder: r,
+            inter_mul_value,
+            is_zero,
+        })
+    }
+
+    fn assign_instance(
+        config: &Self::InstructionConfig,
+        instance: &mut [MaybeUninit<E::BaseField>],
+        lkm: &mut LkMultiplicity,
+        step: &StepRecord,
+    ) -> Result<(), ZKVMError> {
+        let rs1 = step.rs1().unwrap().value;
+        let rs2 = step.rs2().unwrap().value;
+        let rd = step.rd().unwrap().value.after;
+
+        // dividend = divisor * outcome + r
+        let dividend = Value::new_unchecked(rs1);
+        let divisor = Value::new_unchecked(rs2);
+        let outcome = Value::new(rd, lkm);
+
+        // divisor * outcome
+        let inter_mul_value = Value::new(rs2 * rd, lkm);
+        let r = if rs2 == 0 {
+            Value::new_unchecked(0)
+        } else {
+            Value::new(rs1 % rs2, lkm)
+        };
+
+        // assignment
+        config.r_insn.assign_instance(instance, lkm, step)?;
+        config.divisor.assign_limbs(instance, divisor.u16_fields());
+        config.outcome.assign_limbs(instance, outcome.u16_fields());
+
+        let (_, mul_carries, add_carries) = divisor.mul_add(&outcome, &r, lkm, true);
+        config
+            .inter_mul_value
+            .assign_limbs(instance, inter_mul_value.u16_fields());
+        config.inter_mul_value.assign_carries(
+            instance,
+            mul_carries
+                .into_iter()
+                .map(|carry| E::BaseField::from(carry as u64))
+                .collect_vec(),
+        );
+        config.remainder.assign_limbs(instance, r.u16_fields());
+
+        config
+            .dividend
+            .assign_limbs(instance, dividend.u16_fields());
+        config.dividend.assign_carries(
+            instance,
+            add_carries
+                .into_iter()
+                .map(|carry| E::BaseField::from(carry as u64))
+                .collect_vec(),
+        );
+
+        config
+            .is_zero
+            .assign_instance(instance, (rs2 as u64).into())?;
+
+        Ok(())
+    }
+}
+
+#[cfg(test)]
+mod test {
+
+    mod divu {
+        use std::u32;
+
+        use ceno_emul::{Change, StepRecord, Word};
+        use goldilocks::GoldilocksExt2;
+        use itertools::Itertools;
+        use multilinear_extensions::mle::IntoMLEs;
+        use rand::Rng;
+
+        use crate::{
+            circuit_builder::{CircuitBuilder, ConstraintSystem},
+            instructions::{riscv::divu::DivUInstruction, Instruction},
+            scheme::mock_prover::{MockProver, MOCK_PC_DIVU, MOCK_PROGRAM},
+        };
+
+        fn verify(name: &'static str, dividend: Word, divisor: Word, outcome: Word) {
+            let mut cs = ConstraintSystem::<GoldilocksExt2>::new(|| "riscv");
+            let mut cb = CircuitBuilder::new(&mut cs);
+            let config = cb
+                .namespace(
+                    || format!("divu_{name}"),
+                    |cb| Ok(DivUInstruction::construct_circuit(cb)),
+                )
+                .unwrap()
+                .unwrap();
+
+            // values assignment
+            let (raw_witin, _) = DivUInstruction::assign_instances(
+                &config,
+                cb.cs.num_witin as usize,
+                vec![StepRecord::new_r_instruction(
+                    3,
+                    MOCK_PC_DIVU,
+                    MOCK_PROGRAM[9],
+                    dividend,
+                    divisor,
+                    Change::new(0, outcome),
+                    0,
+                )],
+            )
+            .unwrap();
+
+            MockProver::assert_satisfied(
+                &mut cb,
+                &raw_witin
+                    .de_interleaving()
+                    .into_mles()
+                    .into_iter()
+                    .map(|v| v.into())
+                    .collect_vec(),
+                None,
+            );
+        }
+        #[test]
+        fn test_opcode_divu() {
+            verify("basic", 10, 2, 5);
+            verify("dividend > divisor", 10, 11, 0);
+            verify("remainder", 11, 2, 5);
+            verify("u32::MAX", u32::MAX, u32::MAX, 1);
+            verify("u32::MAX div by 2", u32::MAX, 2, u32::MAX / 2);
+            verify("div by zero", 10, 0, u32::MAX);
+            verify("mul carry", 1202729773, 171818539, 7);
+        }
+
+        #[test]
+        fn test_opcode_divu_random() {
+            let mut rng = rand::thread_rng();
+            let a: u32 = rng.gen();
+            let b: u32 = rng.gen_range(1..u32::MAX);
+            println!("random: {} / {} = {}", a, b, a / b);
+            verify("random", a, b, a / b);
+        }
+    }
+}

ceno_zkvm/src/instructions/riscv/r_insn.rs

@@ -77,32 +77,31 @@ impl<E: ExtensionField> RInstructionConfig<E> {
         let prev_rd_value = UInt::new_unchecked(|| "prev_rd_value", circuit_builder)?;
 
         // Register read and write.
-        let (ts, lt_rs1_cfg) = circuit_builder.register_read(
+        let (next_ts, lt_rs1_cfg) = circuit_builder.register_read(
             || "read_rs1",
             &rs1_id,
             prev_rs1_ts.expr(),
             cur_ts.expr(),
             rs1_read,
         )?;
-        let (ts, lt_rs2_cfg) = circuit_builder.register_read(
+        let (next_ts, lt_rs2_cfg) = circuit_builder.register_read(
             || "read_rs2",
             &rs2_id,
             prev_rs2_ts.expr(),
-            ts,
+            next_ts,
             rs2_read,
         )?;
-        let (ts, lt_prev_ts_cfg) = circuit_builder.register_write(
+        let (next_ts, lt_prev_ts_cfg) = circuit_builder.register_write(
             || "write_rd",
             &rd_id,
             prev_rd_ts.expr(),
-            ts,
+            next_ts,
             prev_rd_value.register_expr(),
             rd_written,
         )?;
 
         // State out.
         let next_pc = pc.expr() + PC_STEP_SIZE.into();
-        let next_ts = ts + 1.into();
         circuit_builder.state_out(next_pc, next_ts)?;
 
         Ok(RInstructionConfig {

ceno_zkvm/src/scheme/mock_prover.rs

@@ -58,6 +58,8 @@ pub const MOCK_PROGRAM: &[u32] = &[
     0x00 << 25 | MOCK_RS2 << 20 | MOCK_RS1 << 15 | 0b001 << 12 | 0x08 << 7 | 0x63,
     // blt x2, x3, -8
     0b_1_111111 << 25 | MOCK_RS2 << 20 | MOCK_RS1 << 15 | 0b_100 << 12 | 0b_1100_1 << 7 | 0x63,
+    // divu (0x01, 0x05, 0x33)
+    0x01 << 25 | MOCK_RS2 << 20 | MOCK_RS1 << 15 | 0x05 << 12 | MOCK_RD << 7 | 0x33,
 ];
 // Addresses of particular instructions in the mock program.
 pub const MOCK_PC_ADD: ByteAddr = ByteAddr(CENO_PLATFORM.pc_start());
@@ -72,6 +74,7 @@ pub const MOCK_PC_BLT: Change<ByteAddr> = Change {
     before: ByteAddr(CENO_PLATFORM.pc_start() + 32),
     after: ByteAddr(CENO_PLATFORM.pc_start() + 24),
 };
+pub const MOCK_PC_DIVU: ByteAddr = ByteAddr(CENO_PLATFORM.pc_start() + 36);
 
 #[allow(clippy::enum_variant_names)]
 #[derive(Debug, PartialEq, Clone)]