Feat: Ecall/Halt

ceno_emul/src/vm_state.rs

@@ -93,7 +93,10 @@ impl EmuContext for VMState {
     fn ecall(&mut self) -> Result<bool> {
         let function = self.load_register(self.platform.reg_ecall())?;
         if function == self.platform.ecall_halt() {
-            let _exit_code = self.load_register(self.platform.reg_arg0())?;
+            let exit_code = self.load_register(self.platform.reg_arg0())?;
+            tracing::debug!("halt with exit_code={}", exit_code);
+
+            self.set_pc(ByteAddr(0));
             self.halted = true;
             Ok(true)
         } else {

ceno_zkvm/examples/riscv_opcodes.rs

@@ -10,10 +10,11 @@ use const_env::from_env;
 
 use ceno_emul::{
     ByteAddr,
-    InsnKind::{ADD, BLT},
+    InsnKind::{ADD, BLT, EANY},
     StepRecord, VMState, CENO_PLATFORM,
 };
 use ceno_zkvm::{
+    instructions::riscv::ecall::HaltInstruction,
     scheme::{constants::MAX_NUM_VARIABLES, verifier::ZKVMVerifier},
     structs::{ZKVMConstraintSystem, ZKVMFixedTraces, ZKVMWitnesses},
     tables::{AndTableCircuit, LtuTableCircuit, U16TableCircuit},
@@ -104,6 +105,7 @@ fn main() {
     // opcode circuits
     let add_config = zkvm_cs.register_opcode_circuit::<AddInstruction<E>>();
     let blt_config = zkvm_cs.register_opcode_circuit::<BltInstruction>();
+    let halt_config = zkvm_cs.register_opcode_circuit::<HaltInstruction<E>>();
     // tables
     let u16_range_config = zkvm_cs.register_table_circuit::<U16TableCircuit<E>>();
     // let u1_range_config = zkvm_cs.register_table_circuit::<U1TableCircuit<E>>();
@@ -120,6 +122,7 @@ fn main() {
     let mut zkvm_fixed_traces = ZKVMFixedTraces::default();
     zkvm_fixed_traces.register_opcode_circuit::<AddInstruction<E>>(&zkvm_cs);
     zkvm_fixed_traces.register_opcode_circuit::<BltInstruction>(&zkvm_cs);
+    zkvm_fixed_traces.register_opcode_circuit::<HaltInstruction<E>>(&zkvm_cs);
 
     zkvm_fixed_traces.register_table_circuit::<U16TableCircuit<E>>(
         &zkvm_cs,
@@ -179,12 +182,18 @@ fn main() {
             .collect::<Vec<_>>();
         let mut add_records = Vec::new();
         let mut blt_records = Vec::new();
-        all_records.iter().for_each(|record| {
+        let mut halt_records = Vec::new();
+        all_records.into_iter().for_each(|record| {
             let kind = record.insn().kind().1;
-            if kind == ADD {
-                add_records.push(record.clone());
-            } else if kind == BLT {
-                blt_records.push(record.clone());
+            match kind {
+                ADD => add_records.push(record),
+                BLT => blt_records.push(record),
+                EANY => {
+                    if record.rs1().unwrap().value == CENO_PLATFORM.ecall_halt() {
+                        halt_records.push(record);
+                    }
+                }
+                _ => {}
             }
         });
 
@@ -202,6 +211,10 @@ fn main() {
         zkvm_witness
             .assign_opcode_circuit::<BltInstruction>(&zkvm_cs, &blt_config, blt_records)
             .unwrap();
+        zkvm_witness
+            .assign_opcode_circuit::<HaltInstruction<E>>(&zkvm_cs, &halt_config, halt_records)
+            .unwrap();
+
         zkvm_witness.finalize_lk_multiplicities();
         // assign table circuits
         zkvm_witness

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,
 };
 
@@ -16,6 +16,13 @@ pub trait GlobalStateRegisterMachineChipOperations<E: ExtensionField> {
     fn state_in(&mut self, pc: Expression<E>, ts: Expression<E>) -> Result<(), ZKVMError>;
 
     fn state_out(&mut self, pc: Expression<E>, ts: Expression<E>) -> Result<(), ZKVMError>;
+
+    fn state_out_with_exit_code(
+        &mut self,
+        pc: Expression<E>,
+        ts: Expression<E>,
+        exit_code: Expression<E>,
+    ) -> Result<(), ZKVMError>;
 }
 
 /// The common representation of a register value.
@@ -26,7 +33,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 +43,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/global_state.rs

@@ -12,6 +12,7 @@ impl<'a, E: ExtensionField> GlobalStateRegisterMachineChipOperations<E> for Circ
             Expression::Constant(E::BaseField::from(RAMType::GlobalState as u64)),
             pc,
             ts,
+            0.into(),
         ];
 
         let rlc_record = self.rlc_chip_record(items);
@@ -23,8 +24,25 @@ impl<'a, E: ExtensionField> GlobalStateRegisterMachineChipOperations<E> for Circ
             Expression::Constant(E::BaseField::from(RAMType::GlobalState as u64)),
             pc,
             ts,
+            0.into(),
         ];
         let rlc_record = self.rlc_chip_record(items);
         self.write_record(|| "state_out", rlc_record)
     }
+
+    fn state_out_with_exit_code(
+        &mut self,
+        pc: Expression<E>,
+        ts: Expression<E>,
+        exit_code: Expression<E>,
+    ) -> Result<(), ZKVMError> {
+        let items: Vec<Expression<E>> = vec![
+            Expression::Constant(E::BaseField::from(RAMType::GlobalState as u64)),
+            pc,
+            ts,
+            exit_code,
+        ];
+        let rlc_record = self.rlc_chip_record(items);
+        self.write_record(|| "state_out_with_exit_code", rlc_record)
+    }
 }

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,3 @@
+mod halt;
+
+pub use halt::HaltInstruction;

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

@@ -0,0 +1,106 @@
+use crate::{
+    chip_handler::RegisterChipOperations,
+    circuit_builder::CircuitBuilder,
+    error::ZKVMError,
+    expression::{ToExpr, WitIn},
+    instructions::{
+        riscv::{
+            config::{ExprLtConfig, ExprLtInput},
+            constants::{UInt, ECALL_HALT},
+            ecall_insn::EcallInstructionConfig,
+        },
+        Instruction,
+    },
+    set_val,
+    witness::LkMultiplicity,
+    Value,
+};
+use ceno_emul::StepRecord;
+use ff_ext::ExtensionField;
+use std::{marker::PhantomData, mem::MaybeUninit};
+
+pub struct HaltConfig<E: ExtensionField> {
+    ecall_cfg: EcallInstructionConfig,
+    prev_x10_value: UInt<E>,
+    prev_x10_ts: WitIn,
+    lt_x10_cfg: ExprLtConfig,
+}
+
+pub struct HaltInstruction<E>(PhantomData<E>);
+
+impl<E: ExtensionField> Instruction<E> for HaltInstruction<E> {
+    type InstructionConfig = HaltConfig<E>;
+
+    fn name() -> String {
+        "ECALL_HALT".into()
+    }
+
+    fn construct_circuit(cb: &mut CircuitBuilder<E>) -> Result<Self::InstructionConfig, ZKVMError> {
+        let prev_x10_ts = cb.create_witin(|| "prev_x10_ts")?;
+        let prev_x10_value = UInt::new_unchecked(|| "prev_x10_value", cb)?;
+
+        let ecall_cfg = EcallInstructionConfig::construct_circuit(
+            cb,
+            [ECALL_HALT[0].into(), ECALL_HALT[1].into()],
+            None,
+            Some(0.into()),
+            Some(prev_x10_value.register_expr()),
+        )?;
+
+        // read exit_code from arg0 (X10 register) and write it to global state
+        let (_, lt_x10_cfg) = cb.register_read(
+            || "read x10",
+            E::BaseField::from(ceno_emul::CENO_PLATFORM.reg_arg0() as u64),
+            prev_x10_ts.expr(),
+            ecall_cfg.ts.expr(),
+            prev_x10_value.register_expr(),
+        )?;
+
+        Ok(HaltConfig {
+            ecall_cfg,
+            prev_x10_value,
+            prev_x10_ts,
+            lt_x10_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,
+            "pc after ecall/halt {:x}",
+            step.pc().after.0
+        );
+
+        set_val!(
+            instance,
+            config.prev_x10_ts,
+            step.rs2().unwrap().previous_cycle
+        );
+        config.prev_x10_value.assign_limbs(
+            instance,
+            Value::new_unchecked(step.rs2().unwrap().value).u16_fields(),
+        );
+
+        ExprLtInput {
+            lhs: step.rs2().unwrap().previous_cycle,
+            rhs: step.cycle(),
+        }
+        .assign(instance, &config.lt_x10_cfg, lk_multiplicity);
+
+        config
+            .ecall_cfg
+            .assign_instance::<E>(instance, lk_multiplicity, step)?;
+
+        Ok(())
+    }
+}