draft srli opcode

ceno_zkvm/src/instructions/riscv.rs

@@ -7,6 +7,7 @@ pub mod addsub;
 pub mod blt;
 pub mod config;
 pub mod constants;
+pub mod srli;
 
 #[cfg(test)]
 mod test;

ceno_zkvm/src/instructions/riscv/srli.rs

@@ -0,0 +1,169 @@
+use std::mem::MaybeUninit;
+
+use ff_ext::ExtensionField;
+use goldilocks::SmallField;
+
+use crate::{
+    chip_handler::{GlobalStateRegisterMachineChipOperations, RegisterChipOperations},
+    circuit_builder::CircuitBuilder,
+    error::ZKVMError,
+    expression::{ToExpr, WitIn},
+    instructions::Instruction,
+    set_val,
+};
+
+use super::{
+    config::ExprLtConfig,
+    constants::{RegUInt, PC_STEP_SIZE},
+};
+
+pub struct SrliInstruction;
+
+pub struct InstructionConfig {
+    pub pc: WitIn,
+    pub ts: WitIn,
+    pub imm: WitIn,
+    pub rs1_id: WitIn,
+    pub rd_id: WitIn,
+    pub prev_rs1_ts: WitIn,
+    pub prev_rd_ts: WitIn,
+    pub lt_rs1_cfg: ExprLtConfig,
+    pub lt_rd_cfg: ExprLtConfig,
+}
+
+pub struct SrliInput {
+    pub pc: u16,
+    pub ts: u16,
+    pub imm: i16,
+    pub rs1_id: u8,
+    pub rd_id: u8,
+    pub prev_rs1_ts: u16,
+    pub prev_rd_ts: u16,
+}
+
+impl SrliInput {
+    pub fn assign<F: SmallField, E: ExtensionField<BaseField = F>>(
+        &self,
+        config: &InstructionConfig,
+        instance: &mut [MaybeUninit<F>],
+    ) {
+        set_val!(instance, config.pc, self.pc as u64);
+        set_val!(instance, config.ts, self.ts as u64);
+        set_val!(instance, config.imm, self.imm as u64);
+        set_val!(instance, config.rs1_id, self.rs1_id as u64);
+        set_val!(instance, config.rd_id, self.rd_id as u64);
+        set_val!(instance, config.prev_rs1_ts, self.prev_rs1_ts as u64);
+        set_val!(instance, config.prev_rd_ts, self.prev_rd_ts as u64);
+    }
+
+    pub fn random() -> Self {
+        use ark_std::{rand::Rng, test_rng};
+        let mut rng = test_rng();
+
+        // hack to generate valid inputs
+        let ts_bound: u16 = rng.gen_range(100..1000);
+        let pc_bound: u16 = rng.gen_range(100..1000);
+
+        Self {
+            pc: rng.gen_range(pc_bound..(1 << 15)),
+            ts: rng.gen_range(ts_bound..(1 << 15)),
+            imm: rng.gen_range(-(pc_bound as i16)..2047),
+            rs1_id: rng.gen(),
+            rd_id: rng.gen(),
+            prev_rs1_ts: rng.gen_range(0..ts_bound),
+            prev_rd_ts: rng.gen_range(0..ts_bound),
+        }
+    }
+}
+
+fn srli_gadget<E: ExtensionField>(
+    circuit_builder: &mut CircuitBuilder<E>,
+) -> Result<InstructionConfig, ZKVMError> {
+    let pc = circuit_builder.create_witin(|| "pc")?;
+    let imm = circuit_builder.create_witin(|| "imm")?;
+    let cur_ts = circuit_builder.create_witin(|| "ts")?;
+
+    // state in
+    circuit_builder.state_in(pc.expr(), cur_ts.expr())?;
+
+    let rs1_id = circuit_builder.create_witin(|| "rs1_id")?;
+    let prev_rs1_ts = circuit_builder.create_witin(|| "prev_rs1_ts")?;
+
+    let rd_id = circuit_builder.create_witin(|| "rd_id")?;
+    let prev_rd_ts = circuit_builder.create_witin(|| "prev_rd_ts")?;
+    let prev_rd_value = RegUInt::new_unchecked(|| "prev_rd_value", circuit_builder)?;
+
+    // TODO constrain 2 ** imm == two_pow_imm using lookup
+    let operand = RegUInt::new_unchecked(|| "operand", circuit_builder)?;
+    let mut two_pow_imm = RegUInt::new_unchecked(|| "two_pow_imm", circuit_builder)?;
+    let remainder = RegUInt::new_unchecked(|| "remainder", circuit_builder)?;
+    let mut result = RegUInt::new_unchecked(|| "result", circuit_builder)?;
+
+    // read operand from rs1
+    let (ts, lt_rs1_cfg) = circuit_builder.register_read(
+        || "read from rs1",
+        &rs1_id,
+        prev_rs1_ts.expr(),
+        cur_ts.expr(),
+        &operand,
+    )?;
+
+    // for: operand >> imm == result
+    // constraint: operand == (result << imm) + remainder
+    let rhs = result
+        .mul(|| "hey", circuit_builder, &mut two_pow_imm, false)?
+        .add(|| "add", circuit_builder, &remainder, false)?;
+    operand.eq(|| "srli check", circuit_builder, &rhs)?;
+
+    // write to register destination rd
+    let (ts, lt_rd_cfg) = circuit_builder.register_write(
+        || "write to rd",
+        &rd_id,
+        prev_rd_ts.expr(),
+        ts,
+        &prev_rd_value,
+        &result,
+    )?;
+
+    // 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(InstructionConfig {
+        pc,
+        ts: cur_ts,
+        imm,
+        rs1_id,
+        rd_id,
+        prev_rs1_ts,
+        prev_rd_ts,
+        lt_rs1_cfg,
+        lt_rd_cfg,
+    })
+}
+
+impl<E: ExtensionField> Instruction<E> for SrliInstruction {
+    type InstructionConfig = InstructionConfig;
+
+    fn name() -> String {
+        "SRLI".into()
+    }
+
+    fn construct_circuit(
+        circuit_builder: &mut crate::circuit_builder::CircuitBuilder<E>,
+    ) -> Result<Self::InstructionConfig, crate::error::ZKVMError> {
+        srli_gadget(circuit_builder)
+    }
+
+    fn assign_instance(
+        config: &Self::InstructionConfig,
+        instance: &mut [MaybeUninit<<E as ExtensionField>::BaseField>],
+        _lk_multiplicity: &mut crate::witness::LkMultiplicity,
+        _step: &ceno_emul::StepRecord,
+    ) -> Result<(), crate::error::ZKVMError> {
+        let input = SrliInput::random();
+        input.assign::<_, E>(config, instance);
+        Ok(())
+    }
+}