i_insn config and SRLI opcode (#229)

Closes #122

---------

Co-authored-by: sm.wu <[email protected]>

ceno_emul/src/tracer.rs

@@ -114,6 +114,34 @@ impl StepRecord {
         }
     }
 
+    pub fn new_i_instruction(
+        cycle: Cycle,
+        pc: ByteAddr,
+        insn_code: Word,
+        rs1_read: Word,
+        rd: Change<Word>,
+        previous_cycle: Cycle,
+    ) -> StepRecord {
+        let insn = DecodedInstruction::new(insn_code);
+        StepRecord {
+            cycle,
+            pc: Change::new(pc, pc + PC_STEP_SIZE),
+            insn_code,
+            rs1: Some(ReadOp {
+                addr: CENO_PLATFORM.register_vma(insn.rs1() as RegIdx).into(),
+                value: rs1_read,
+                previous_cycle,
+            }),
+            rs2: None,
+            rd: Some(WriteOp {
+                addr: CENO_PLATFORM.register_vma(insn.rd() as RegIdx).into(),
+                value: rd,
+                previous_cycle,
+            }),
+            memory_op: None,
+        }
+    }
+
     pub fn cycle(&self) -> Cycle {
         self.cycle
     }

ceno_zkvm/src/instructions/riscv.rs

@@ -1,15 +1,16 @@
 use ceno_emul::InsnKind;
 
 pub mod arith;
+mod b_insn;
 pub mod blt;
 pub mod branch;
 pub mod config;
 pub mod constants;
 pub mod divu;
+mod i_insn;
 pub mod logic;
-
-mod b_insn;
 mod r_insn;
+pub mod shift_imm;
 
 #[cfg(test)]
 mod test;

ceno_zkvm/src/instructions/riscv/divu.rs

@@ -43,10 +43,10 @@ impl<E: ExtensionField, I: RIVInstruction> Instruction<E> for ArithInstruction<E
         // 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 r = UInt::new(|| "remainder", circuit_builder)?;
 
         let (inter_mul_value, dividend) =
-            divisor.mul_add(|| "dividend", circuit_builder, &mut outcome, &mut r, true)?;
+            divisor.mul_add(|| "dividend", circuit_builder, &mut outcome, &r, true)?;
 
         // div by zero check
         let is_zero = IsZeroConfig::construct_circuit(circuit_builder, divisor.value())?;

ceno_zkvm/src/instructions/riscv/i_insn.rs

@@ -0,0 +1,151 @@
+use ceno_emul::{InsnKind, StepRecord};
+use ff_ext::ExtensionField;
+
+use super::{
+    config::ExprLtConfig,
+    constants::{UInt, PC_STEP_SIZE},
+};
+use crate::{
+    chip_handler::{
+        GlobalStateRegisterMachineChipOperations, RegisterChipOperations, RegisterExpr,
+    },
+    circuit_builder::CircuitBuilder,
+    error::ZKVMError,
+    expression::{Expression, ToExpr, WitIn},
+    instructions::riscv::config::ExprLtInput,
+    set_val,
+    tables::InsnRecord,
+    uint::Value,
+    witness::LkMultiplicity,
+};
+use core::mem::MaybeUninit;
+
+/// This config handles the common part of I-type instructions:
+/// - PC, cycle, fetch.
+/// - Registers read and write.
+///
+/// It does not witness of the register values, nor the actual function (e.g. srli, addi, etc).
+#[derive(Debug)]
+pub struct IInstructionConfig<E: ExtensionField> {
+    pub pc: WitIn,
+    pub ts: WitIn,
+    pub rs1_id: WitIn,
+    pub rd_id: WitIn,
+    pub prev_rd_value: UInt<E>,
+    pub prev_rs1_ts: WitIn,
+    pub prev_rd_ts: WitIn,
+    pub lt_rs1_cfg: ExprLtConfig,
+    pub lt_rd_cfg: ExprLtConfig,
+}
+
+impl<E: ExtensionField> IInstructionConfig<E> {
+    pub fn construct_circuit(
+        circuit_builder: &mut CircuitBuilder<E>,
+        insn_kind: InsnKind,
+        imm: &Expression<E>,
+        rs1_read: RegisterExpr<E>,
+        rd_written: RegisterExpr<E>,
+    ) -> Result<Self, ZKVMError> {
+        // State in.
+        let pc = circuit_builder.create_witin(|| "pc")?;
+        let cur_ts = circuit_builder.create_witin(|| "cur_ts")?;
+        circuit_builder.state_in(pc.expr(), cur_ts.expr())?;
+
+        // Register indexes.
+        let rs1_id = circuit_builder.create_witin(|| "rs1_id")?;
+        let rd_id = circuit_builder.create_witin(|| "rd_id")?;
+
+        // Fetch the instruction.
+        circuit_builder.lk_fetch(&InsnRecord::new(
+            pc.expr(),
+            (insn_kind.codes().opcode as usize).into(),
+            rd_id.expr(),
+            (insn_kind.codes().func3 as usize).into(),
+            rs1_id.expr(),
+            0.into(),
+            imm.clone(),
+        ))?;
+
+        // Register state.
+        let prev_rs1_ts = circuit_builder.create_witin(|| "prev_rs1_ts")?;
+        let prev_rd_ts = circuit_builder.create_witin(|| "prev_rd_ts")?;
+        let prev_rd_value = UInt::new_unchecked(|| "prev_rd_value", circuit_builder)?;
+
+        // Register read and write.
+        let (next_ts, lt_rs1_cfg) = circuit_builder.register_read(
+            || "read_rs1",
+            &rs1_id,
+            prev_rs1_ts.expr(),
+            cur_ts.expr(),
+            rs1_read,
+        )?;
+        let (next_ts, lt_rd_cfg) = circuit_builder.register_write(
+            || "write_rd",
+            &rd_id,
+            prev_rd_ts.expr(),
+            next_ts,
+            prev_rd_value.register_expr(),
+            rd_written,
+        )?;
+
+        // State out.
+        let next_pc = pc.expr() + PC_STEP_SIZE.into();
+        circuit_builder.state_out(next_pc, next_ts)?;
+
+        Ok(IInstructionConfig {
+            pc,
+            ts: cur_ts,
+            rs1_id,
+            rd_id,
+            prev_rd_value,
+            prev_rs1_ts,
+            prev_rd_ts,
+            lt_rs1_cfg,
+            lt_rd_cfg,
+        })
+    }
+
+    pub fn assign_instance(
+        &self,
+        instance: &mut [MaybeUninit<<E as ExtensionField>::BaseField>],
+        lk_multiplicity: &mut LkMultiplicity,
+        step: &StepRecord,
+    ) -> Result<(), ZKVMError> {
+        // State in.
+        set_val!(instance, self.pc, step.pc().before.0 as u64);
+        set_val!(instance, self.ts, step.cycle());
+
+        // Register indexes.
+        set_val!(instance, self.rs1_id, step.insn().rs1() as u64);
+        set_val!(instance, self.rd_id, step.insn().rd() as u64);
+
+        // Fetch the instruction.
+        lk_multiplicity.fetch(step.pc().before.0);
+
+        // Register state.
+        set_val!(
+            instance,
+            self.prev_rs1_ts,
+            step.rs1().unwrap().previous_cycle
+        );
+        set_val!(instance, self.prev_rd_ts, step.rd().unwrap().previous_cycle);
+        self.prev_rd_value.assign_limbs(
+            instance,
+            Value::new_unchecked(step.rd().unwrap().value.before).u16_fields(),
+        );
+
+        // Register read and write.
+        ExprLtInput {
+            lhs: step.rs1().unwrap().previous_cycle,
+            rhs: step.cycle(),
+        }
+        .assign(instance, &self.lt_rs1_cfg, lk_multiplicity);
+        ExprLtInput {
+            lhs: step.rd().unwrap().previous_cycle,
+            rhs: step.cycle() + 1,
+        }
+        .assign(instance, &self.lt_rd_cfg, lk_multiplicity);
+
+        Ok(())
+    }
+}

ceno_zkvm/src/instructions/riscv/shift_imm.rs

@@ -0,0 +1,12 @@
+use super::RIVInstruction;
+
+mod shift_imm_circuit;
+
+#[cfg(test)]
+mod test;
+
+pub struct SrliOp;
+
+impl RIVInstruction for SrliOp {
+    const INST_KIND: ceno_emul::InsnKind = ceno_emul::InsnKind::SRLI;
+}

ceno_zkvm/src/instructions/riscv/shift_imm/shift_imm_circuit.rs

@@ -0,0 +1,116 @@
+use std::{marker::PhantomData, mem::MaybeUninit};
+
+use ceno_emul::StepRecord;
+use ff_ext::ExtensionField;
+
+use crate::{
+    circuit_builder::CircuitBuilder,
+    error::ZKVMError,
+    instructions::{
+        riscv::{constants::UInt, i_insn::IInstructionConfig, RIVInstruction},
+        Instruction,
+    },
+    witness::LkMultiplicity,
+    Value,
+};
+
+pub struct ShiftImmInstruction<E, I>(PhantomData<(E, I)>);
+
+pub struct InstructionConfig<E: ExtensionField> {
+    i_insn: IInstructionConfig<E>,
+
+    imm: UInt<E>,
+    rd_written: UInt<E>,
+    remainder: UInt<E>,
+    rd_imm_mul: UInt<E>,
+    rd_imm_rem_add: UInt<E>,
+}
+
+impl<E: ExtensionField, I: RIVInstruction> Instruction<E> for ShiftImmInstruction<E, I> {
+    type InstructionConfig = InstructionConfig<E>;
+
+    fn name() -> String {
+        format!("{:?}", I::INST_KIND)
+    }
+
+    fn construct_circuit(
+        circuit_builder: &mut CircuitBuilder<E>,
+    ) -> Result<Self::InstructionConfig, ZKVMError> {
+        let mut imm = UInt::new(|| "imm", circuit_builder)?;
+        let mut rd_written = UInt::new_unchecked(|| "rd_written", circuit_builder)?;
+
+        // Note: `imm` is set to 2**imm (upto 32 bit) just for SRLI for efficient verification
+        // Goal is to constrain:
+        // rs1_read == rd_written * imm + remainder
+        let remainder = UInt::new(|| "remainder", circuit_builder)?;
+        let (rd_imm_mul, rd_imm_rem_add) = rd_written.mul_add(
+            || "rd_written * imm +remainder ",
+            circuit_builder,
+            &mut imm,
+            &remainder,
+            true,
+        )?;
+
+        let i_insn = IInstructionConfig::<E>::construct_circuit(
+            circuit_builder,
+            I::INST_KIND,
+            &imm.value(),
+            rd_imm_rem_add.register_expr(),
+            rd_written.register_expr(),
+        )?;
+
+        Ok(InstructionConfig {
+            i_insn,
+            imm,
+            rd_written,
+            remainder,
+            rd_imm_mul,
+            rd_imm_rem_add,
+        })
+    }
+
+    fn assign_instance(
+        config: &Self::InstructionConfig,
+        instance: &mut [MaybeUninit<<E as ExtensionField>::BaseField>],
+        lk_multiplicity: &mut LkMultiplicity,
+        step: &StepRecord,
+    ) -> Result<(), ZKVMError> {
+        // We need to calculate result and remainder.
+        let rs1_read = step.rs1().unwrap().value;
+        let rd_written = step.rd().unwrap().value.after;
+        let imm = step.insn().imm_or_funct7();
+        let result = rs1_read.wrapping_div(imm);
+        let remainder = rs1_read.wrapping_sub(result * imm);
+        assert_eq!(result, rd_written, "SRLI: result mismatch");
+
+        // Assignment.
+        let rd_written = Value::new_unchecked(rd_written);
+        let imm = Value::new(imm, lk_multiplicity);
+        let remainder = Value::new(remainder, lk_multiplicity);
+
+        let rd_imm_mul = rd_written.mul(&imm, lk_multiplicity, true);
+        let rd_imm = Value::from_limb_slice_unchecked(&rd_imm_mul.0);
+        config
+            .rd_imm_mul
+            .assign_limb_with_carry(instance, &rd_imm_mul);
+
+        let rd_imm_rem_add = rd_imm.add(&remainder, lk_multiplicity, true);
+        debug_assert_eq!(
+            Value::<u32>::from_limb_slice_unchecked(&rd_imm_rem_add.0).as_u64(),
+            rs1_read as u64,
+            "SRLI: rd_imm_rem_add mismatch"
+        );
+        config
+            .rd_imm_rem_add
+            .assign_limb_with_carry(instance, &rd_imm_rem_add);
+
+        config
+            .i_insn
+            .assign_instance(instance, lk_multiplicity, step)?;
+        config.imm.assign_value(instance, imm);
+        config.rd_written.assign_value(instance, rd_written);
+        config.remainder.assign_value(instance, remainder);
+
+        Ok(())
+    }
+}