Feat: Ecall/Halt

ceno_zkvm/src/chip_handler.rs

@@ -2,7 +2,7 @@ use ff_ext::ExtensionField;
 
 use crate::{
     error::ZKVMError,
-    expression::{Expression, WitIn},
+    expression::{Expression, ToExpr, WitIn},
     instructions::riscv::config::ExprLtConfig,
 };
 
@@ -26,7 +26,7 @@ pub trait RegisterChipOperations<E: ExtensionField, NR: Into<String>, N: FnOnce(
     fn register_read(
         &mut self,
         name_fn: N,
-        register_id: &WitIn,
+        register_id: impl ToExpr<E, Output = Expression<E>>,
         prev_ts: Expression<E>,
         ts: Expression<E>,
         value: RegisterExpr<E>,
@@ -36,7 +36,7 @@ pub trait RegisterChipOperations<E: ExtensionField, NR: Into<String>, N: FnOnce(
     fn register_write(
         &mut self,
         name_fn: N,
-        register_id: &WitIn,
+        register_id: impl ToExpr<E, Output = Expression<E>>,
         prev_ts: Expression<E>,
         ts: Expression<E>,
         prev_values: RegisterExpr<E>,

ceno_zkvm/src/chip_handler/register.rs

@@ -3,7 +3,7 @@ use ff_ext::ExtensionField;
 use crate::{
     circuit_builder::CircuitBuilder,
     error::ZKVMError,
-    expression::{Expression, ToExpr, WitIn},
+    expression::{Expression, ToExpr},
     instructions::riscv::config::ExprLtConfig,
     structs::RAMType,
 };
@@ -16,7 +16,7 @@ impl<'a, E: ExtensionField, NR: Into<String>, N: FnOnce() -> NR> RegisterChipOpe
     fn register_read(
         &mut self,
         name_fn: N,
-        register_id: &WitIn,
+        register_id: impl ToExpr<E, Output = Expression<E>>,
         prev_ts: Expression<E>,
         ts: Expression<E>,
         value: RegisterExpr<E>,
@@ -61,12 +61,13 @@ impl<'a, E: ExtensionField, NR: Into<String>, N: FnOnce() -> NR> RegisterChipOpe
     fn register_write(
         &mut self,
         name_fn: N,
-        register_id: &WitIn,
+        register_id: impl ToExpr<E, Output = Expression<E>>,
         prev_ts: Expression<E>,
         ts: Expression<E>,
         prev_values: RegisterExpr<E>,
         value: RegisterExpr<E>,
     ) -> Result<(Expression<E>, ExprLtConfig), ZKVMError> {
+        assert!(register_id.expr().degree() <= 1);
         self.namespace(name_fn, |cb| {
             // READ (a, v, t)
             let read_record = cb.rlc_chip_record(

ceno_zkvm/src/expression.rs

@@ -486,6 +486,13 @@ impl<E: ExtensionField> ToExpr<E> for WitIn {
     }
 }
 
+impl<E: ExtensionField> ToExpr<E> for &WitIn {
+    type Output = Expression<E>;
+    fn expr(&self) -> Expression<E> {
+        Expression::WitIn(self.id)
+    }
+}
+
 impl<F: SmallField, E: ExtensionField<BaseField = F>> ToExpr<E> for F {
     type Output = Expression<E>;
     fn expr(&self) -> Expression<E> {

ceno_zkvm/src/instructions/riscv.rs

@@ -7,8 +7,11 @@ pub mod branch;
 pub mod config;
 pub mod constants;
 pub mod divu;
+pub mod ecall;
 mod i_insn;
 pub mod logic;
+
+mod ecall_insn;
 mod r_insn;
 pub mod shift_imm;
 

ceno_zkvm/src/instructions/riscv/b_insn.rs

@@ -85,14 +85,14 @@ impl BInstructionConfig {
         // Register reads.
         let (ts, lt_rs1_cfg) = circuit_builder.register_read(
             || "read_rs1",
-            &rs1_id,
+            rs1_id,
             prev_rs1_ts.expr(),
             cur_ts.expr(),
             rs1_read,
         )?;
         let (_ts, lt_rs2_cfg) = circuit_builder.register_read(
             || "read_rs2",
-            &rs2_id,
+            rs2_id,
             prev_rs2_ts.expr(),
             ts,
             rs2_read,

ceno_zkvm/src/instructions/riscv/constants.rs

@@ -1,6 +1,7 @@
 use crate::uint::UIntLimbs;
 pub use ceno_emul::PC_STEP_SIZE;
 
+pub const ECALL_HALT: [usize; 2] = [0x00_00, 0x00_00];
 pub const VALUE_BIT_WIDTH: usize = 16;
 
 #[cfg(feature = "riv32")]

ceno_zkvm/src/instructions/riscv/ecall.rs

@@ -0,0 +1 @@
+mod halt;

ceno_zkvm/src/instructions/riscv/ecall/halt.rs

@@ -0,0 +1,58 @@
+use crate::{
+    circuit_builder::CircuitBuilder,
+    error::ZKVMError,
+    instructions::{
+        riscv::{constants::ECALL_HALT, ecall_insn::EcallInstructionConfig},
+        Instruction,
+    },
+    witness::LkMultiplicity,
+};
+use ceno_emul::StepRecord;
+use ff_ext::ExtensionField;
+use std::{marker::PhantomData, mem::MaybeUninit};
+
+pub struct HaltConfig {
+    ecall_cfg: EcallInstructionConfig,
+}
+
+pub struct HaltCircuit<E>(PhantomData<E>);
+
+impl<E: ExtensionField> Instruction<E> for HaltCircuit<E> {
+    type InstructionConfig = HaltConfig;
+
+    fn name() -> String {
+        "ECALL_HALT".into()
+    }
+
+    fn construct_circuit(cb: &mut CircuitBuilder<E>) -> Result<Self::InstructionConfig, ZKVMError> {
+        let ecall_cfg = EcallInstructionConfig::construct_circuit(
+            cb,
+            [ECALL_HALT[0].into(), ECALL_HALT[1].into()],
+            None,
+            Some(0.into()),
+        )?;
+
+        // TODO: read exit_code from arg1 and write it to global state
+
+        Ok(HaltConfig { ecall_cfg })
+    }
+
+    fn assign_instance(
+        config: &Self::InstructionConfig,
+        instance: &mut [MaybeUninit<E::BaseField>],
+        lk_multiplicity: &mut LkMultiplicity,
+        step: &StepRecord,
+    ) -> Result<(), ZKVMError> {
+        assert_eq!(
+            step.rs1().unwrap().value,
+            (ECALL_HALT[0] + (ECALL_HALT[1] << 16)) as u32
+        );
+        assert_eq!(step.pc().after.0, 0);
+
+        config
+            .ecall_cfg
+            .assign_instance::<E>(instance, lk_multiplicity, step)?;
+
+        Ok(())
+    }
+}

ceno_zkvm/src/instructions/riscv/ecall_insn.rs

@@ -0,0 +1,103 @@
+use crate::{
+    chip_handler::{
+        GlobalStateRegisterMachineChipOperations, RegisterChipOperations, RegisterExpr,
+    },
+    circuit_builder::CircuitBuilder,
+    error::ZKVMError,
+    expression::{Expression, ToExpr, WitIn},
+    instructions::riscv::config::{ExprLtConfig, ExprLtInput},
+    set_val,
+    tables::InsnRecord,
+    witness::LkMultiplicity,
+};
+use ceno_emul::{InsnKind::EANY, StepRecord, PC_STEP_SIZE};
+use ff_ext::ExtensionField;
+use std::mem::MaybeUninit;
+
+const X5: u64 = 5;
+#[allow(dead_code)]
+const X10: usize = 10;
+#[allow(dead_code)]
+const X11: usize = 11;
+
+pub struct EcallInstructionConfig {
+    pub pc: WitIn,
+    pub ts: WitIn,
+    prev_x5_ts: WitIn,
+    lt_x5_cfg: ExprLtConfig,
+}
+
+impl EcallInstructionConfig {
+    pub fn construct_circuit<E: ExtensionField>(
+        cb: &mut CircuitBuilder<E>,
+        syscall_id: RegisterExpr<E>,
+        syscall_ret_value: Option<RegisterExpr<E>>,
+        next_pc: Option<Expression<E>>,
+    ) -> Result<Self, ZKVMError> {
+        let pc = cb.create_witin(|| "pc")?;
+        let ts = cb.create_witin(|| "cur_ts")?;
+
+        cb.state_in(pc.expr(), ts.expr())?;
+
+        cb.lk_fetch(&InsnRecord::new(
+            pc.expr(),
+            (EANY.codes().opcode as usize).into(),
+            0.into(),
+            (EANY.codes().func3 as usize).into(),
+            0.into(),
+            0.into(),
+            0.into(), // imm = 0
+        ))?;
+
+        let prev_x5_ts = cb.create_witin(|| "prev_x5_ts")?;
+
+        // read syscall_id from x5 and write return value to x5
+        let (_, lt_x5_cfg) = cb.register_write(
+            || "write x5",
+            E::BaseField::from(X5),
+            prev_x5_ts.expr(),
+            ts.expr(),
+            syscall_id.clone(),
+            syscall_ret_value.map_or(syscall_id, |v| v),
+        )?;
+
+        cb.state_out(
+            next_pc.map_or(pc.expr() + PC_STEP_SIZE.into(), |next_pc| next_pc),
+            ts.expr() + 4.into(),
+        )?;
+
+        Ok(Self {
+            pc,
+            ts,
+            prev_x5_ts,
+            lt_x5_cfg,
+        })
+    }
+
+    pub fn assign_instance<E: ExtensionField>(
+        &self,
+        instance: &mut [MaybeUninit<E::BaseField>],
+        lk_multiplicity: &mut LkMultiplicity,
+        step: &StepRecord,
+    ) -> Result<(), ZKVMError> {
+        set_val!(instance, self.pc, step.pc().before.0 as u64);
+        set_val!(instance, self.ts, step.cycle());
+        lk_multiplicity.fetch(step.pc().before.0);
+
+        // the access of X5 register is stored in rs1()
+        // the access of X10 register is stored in rs2()
+        set_val!(
+            instance,
+            self.prev_x5_ts,
+            step.rs1().unwrap().previous_cycle
+        );
+
+        ExprLtInput {
+            lhs: step.rs1().unwrap().previous_cycle,
+            rhs: step.cycle(),
+        }
+        .assign(instance, &self.lt_x5_cfg, lk_multiplicity);
+
+        Ok(())
+    }
+}

ceno_zkvm/src/instructions/riscv/i_insn.rs

@@ -74,14 +74,14 @@ impl<E: ExtensionField> IInstructionConfig<E> {
         // Register read and write.
         let (next_ts, lt_rs1_cfg) = circuit_builder.register_read(
             || "read_rs1",
-            &rs1_id,
+            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,
+            rd_id,
             prev_rd_ts.expr(),
             next_ts,
             prev_rd_value.register_expr(),

ceno_zkvm/src/instructions/riscv/r_insn.rs

@@ -79,21 +79,21 @@ impl<E: ExtensionField> RInstructionConfig<E> {
         // Register read and write.
         let (next_ts, lt_rs1_cfg) = circuit_builder.register_read(
             || "read_rs1",
-            &rs1_id,
+            rs1_id,
             prev_rs1_ts.expr(),
             cur_ts.expr(),
             rs1_read,
         )?;
         let (next_ts, lt_rs2_cfg) = circuit_builder.register_read(
             || "read_rs2",
-            &rs2_id,
+            rs2_id,
             prev_rs2_ts.expr(),
             next_ts,
             rs2_read,
         )?;
         let (next_ts, lt_prev_ts_cfg) = circuit_builder.register_write(
             || "write_rd",
-            &rd_id,
+            rd_id,
             prev_rd_ts.expr(),
             next_ts,
             prev_rd_value.register_expr(),