refactor signed imm to support limb based constraints BabyBear

ceno_emul/Cargo.toml

@@ -30,3 +30,4 @@ tracing.workspace = true
 [features]
 default = ["forbid_overflow"]
 forbid_overflow = []
+u16limb_circuit = []

ceno_emul/src/disassemble/mod.rs

@@ -8,6 +8,7 @@ use rrs_lib::{
 
 /// A transpiler that converts the 32-bit encoded instructions into instructions.
 pub(crate) struct InstructionTranspiler {
+    #[allow(dead_code)]
     pc: u32,
     word: u32,
 }
@@ -249,59 +250,84 @@ impl InstructionProcessor for InstructionTranspiler {
         }
     }
 
-    /// Convert LUI to ADDI.
-    ///
-    /// RiscV's load-upper-immediate instruction is necessary to build arbitrary constants,
-    /// because its ADDI can only have a relatively small immediate value: there's just not
-    /// enough space in the 32 bits for more.
-    ///
-    /// Our internal ADDI does not have this limitation, so we can convert LUI to ADDI.
-    /// See [`InstructionTranspiler::process_auipc`] for more background on the conversion.
     fn process_lui(&mut self, dec_insn: UType) -> Self::InstructionResult {
         // Verify assumption that the immediate is already shifted left by 12 bits.
         assert_eq!(dec_insn.imm & 0xfff, 0);
-        Instruction {
-            kind: InsnKind::ADDI,
-            rd: dec_insn.rd,
-            rs1: 0,
-            rs2: 0,
-            imm: dec_insn.imm,
-            raw: self.word,
+        #[cfg(not(feature = "u16limb_circuit"))]
+        {
+            // Convert LUI to ADDI.
+            //
+            // RiscV's load-upper-immediate instruction is necessary to build arbitrary constants,
+            // because its ADDI can only have a relatively small immediate value: there's just not
+            // enough space in the 32 bits for more.
+            //
+            // Our internal ADDI does not have this limitation, so we can convert LUI to ADDI.
+            // See [`InstructionTranspiler::process_auipc`] for more background on the conversion.
+            Instruction {
+                kind: InsnKind::ADDI,
+                rd: dec_insn.rd,
+                rs1: 0,
+                rs2: 0,
+                imm: dec_insn.imm,
+                raw: self.word,
+            }
+        }
+        #[cfg(feature = "u16limb_circuit")]
+        {
+            Instruction {
+                kind: InsnKind::LUI,
+                rd: dec_insn.rd,
+                rs1: 0,
+                rs2: 0,
+                imm: dec_insn.imm,
+                raw: self.word,
+            }
         }
     }
 
-    /// Convert AUIPC to ADDI.
-    ///
-    /// RiscV's instructions are designed to be (mosty) position-independent.  AUIPC is used
-    /// to get access to the current program counter, even if the code has been moved around
-    /// by the linker.
-    ///
-    /// Our conversion here happens after the linker has done its job, so we can safely hardcode
-    /// the current program counter into the immediate value of our internal ADDI.
-    ///
-    /// Note that our internal ADDI can have arbitrary intermediate values, not just 12 bits.
-    ///
-    /// ADDI is slightly more general than LUI or AUIPC, because you can also specify an
-    /// input register rs1.  That generality might cost us sligthtly in the non-recursive proof,
-    /// but we suspect decreasing the total number of different instruction kinds will speed up
-    /// the recursive proof.
-    ///
-    /// In any case, AUIPC and LUI together make up ~0.1% of instructions executed in typical
-    /// real world scenarios like a `reth` run.
-    ///
-    /// TODO(Matthias): run benchmarks to verify the impact on recursion, once we have a working
-    /// recursion.
     fn process_auipc(&mut self, dec_insn: UType) -> Self::InstructionResult {
-        let pc = self.pc;
         // Verify our assumption that the immediate is already shifted left by 12 bits.
         assert_eq!(dec_insn.imm & 0xfff, 0);
-        Instruction {
-            kind: InsnKind::ADDI,
-            rd: dec_insn.rd,
-            rs1: 0,
-            rs2: 0,
-            imm: dec_insn.imm.wrapping_add(pc as i32),
-            raw: self.word,
+        #[cfg(not(feature = "u16limb_circuit"))]
+        {
+            let pc = self.pc;
+            // Convert AUIPC to ADDI.
+            //
+            // RiscV's instructions are designed to be (mosty) position-independent.  AUIPC is used
+            // to get access to the current program counter, even if the code has been moved around
+            // by the linker.
+            //
+            // Our conversion here happens after the linker has done its job, so we can safely hardcode
+            // the current program counter into the immediate value of our internal ADDI.
+            //
+            // Note that our internal ADDI can have arbitrary intermediate values, not just 12 bits.
+            //
+            // ADDI is slightly more general than LUI or AUIPC, because you can also specify an
+            // input register rs1.  That generality might cost us sligthtly in the non-recursive proof,
+            // but we suspect decreasing the total number of different instruction kinds will speed up
+            // the recursive proof.
+            //
+            // In any case, AUIPC and LUI together make up ~0.1% of instructions executed in typical
+            // real world scenarios like a `reth` run.
+            Instruction {
+                kind: InsnKind::ADDI,
+                rd: dec_insn.rd,
+                rs1: 0,
+                rs2: 0,
+                imm: dec_insn.imm.wrapping_add(pc as i32),
+                raw: self.word,
+            }
+        }
+        #[cfg(feature = "u16limb_circuit")]
+        {
+            Instruction {
+                kind: InsnKind::AUIPC,
+                rd: dec_insn.rd,
+                rs1: 0,
+                rs2: 0,
+                imm: dec_insn.imm,
+                raw: self.word,
+            }
         }
     }
 

ceno_emul/src/rv32im.rs

@@ -195,6 +195,10 @@ pub enum InsnKind {
     LW,
     LBU,
     LHU,
+    #[cfg(feature = "u16limb_circuit")]
+    LUI,
+    #[cfg(feature = "u16limb_circuit")]
+    AUIPC,
     SB,
     SH,
     SW,
@@ -216,6 +220,8 @@ impl From<InsnKind> for InsnCategory {
             LB | LH | LW | LBU | LHU => Load,
             SB | SH | SW => Store,
             ECALL => System,
+            #[cfg(feature = "u16limb_circuit")]
+            LUI | AUIPC => Compute,
         }
     }
 }
@@ -234,6 +240,8 @@ impl From<InsnKind> for InsnFormat {
             SB | SH | SW => S,
             ECALL => I,
             INVALID => I,
+            #[cfg(feature = "u16limb_circuit")]
+            LUI | AUIPC => U,
         }
     }
 }
@@ -306,6 +314,10 @@ fn step_compute<M: EmuContext>(ctx: &mut M, kind: InsnKind, insn: &Instruction)
 
             match kind {
                 ADDI => rs1.wrapping_add(imm_i),
+                #[cfg(feature = "u16limb_circuit")]
+                LUI => imm_i,
+                #[cfg(feature = "u16limb_circuit")]
+                AUIPC => pc.wrapping_add(imm_i).0,
                 XORI => rs1 ^ imm_i,
                 ORI => rs1 | imm_i,
                 ANDI => rs1 & imm_i,

ceno_zkvm/Cargo.toml

@@ -77,7 +77,7 @@ nightly-features = [
   "witness/nightly-features",
 ]
 sanity-check = ["mpcs/sanity-check"]
-u16limb_circuit = []
+u16limb_circuit = ["ceno_emul/u16limb_circuit"]
 
 [[bench]]
 harness = false

ceno_zkvm/src/gadgets/signed_limbs.rs

@@ -15,6 +15,7 @@ use witness::set_val;
 ///
 /// this configuration structure allows flexible comparison logic depending on whether
 /// the operands should be interpreted as signed or unsigned integers.
+#[derive(Debug)]
 pub struct UIntLimbsLTConfig<E: ExtensionField> {
     // Most significant limb of a and b respectively as a field element, will be range
     // checked to be within [-32768, 32767) if signed and [0, 65536) if unsigned.
@@ -46,7 +47,7 @@ impl<E: ExtensionField> UIntLimbsLT<E> {
         circuit_builder: &mut CircuitBuilder<E>,
         a: &UInt<E>,
         b: &UInt<E>,
-        is_signed: bool,
+        is_sign_comparison: bool,
     ) -> Result<UIntLimbsLTConfig<E>, ZKVMError> {
         // 1 if a < b, 0 otherwise.
         let cmp_lt = circuit_builder.create_bit(|| "cmp_lt")?;
@@ -120,7 +121,7 @@ impl<E: ExtensionField> UIntLimbsLT<E> {
         circuit_builder.assert_ux::<_, _, LIMB_BITS>(
             || "a_msb_f_signed_range_check",
             a_msb_f.expr()
-                + if is_signed {
+                + if is_sign_comparison {
                     E::BaseField::from_canonical_u32(1 << (LIMB_BITS - 1)).expr()
                 } else {
                     Expression::ZERO
@@ -130,7 +131,7 @@ impl<E: ExtensionField> UIntLimbsLT<E> {
         circuit_builder.assert_ux::<_, _, LIMB_BITS>(
             || "b_msb_f_signed_range_check",
             b_msb_f.expr()
-                + if is_signed {
+                + if is_sign_comparison {
                     E::BaseField::from_canonical_u32(1 << (LIMB_BITS - 1)).expr()
                 } else {
                     Expression::ZERO
@@ -153,9 +154,9 @@ impl<E: ExtensionField> UIntLimbsLT<E> {
         lkm: &mut gkr_iop::utils::lk_multiplicity::LkMultiplicity,
         a: &[u16],
         b: &[u16],
-        is_signed: bool,
+        is_sign_comparison: bool,
     ) -> Result<(), CircuitBuilderError> {
-        let (cmp_lt, diff_idx, is_a_neg, is_b_neg) = run_cmp(is_signed, a, b);
+        let (cmp_lt, diff_idx, is_a_neg, is_b_neg) = run_cmp(is_sign_comparison, a, b);
         config
             .diff_marker
             .iter()
@@ -175,7 +176,7 @@ impl<E: ExtensionField> UIntLimbsLT<E> {
         } else {
             (
                 E::BaseField::from_canonical_u16(a[UINT_LIMBS - 1]),
-                a[UINT_LIMBS - 1] + ((is_signed as u16) << (LIMB_BITS - 1)),
+                a[UINT_LIMBS - 1] + ((is_sign_comparison as u16) << (LIMB_BITS - 1)),
             )
         };
         let (b_msb_f, b_msb_range) = if is_b_neg {
@@ -186,7 +187,7 @@ impl<E: ExtensionField> UIntLimbsLT<E> {
         } else {
             (
                 E::BaseField::from_canonical_u16(b[UINT_LIMBS - 1]),
-                b[UINT_LIMBS - 1] + ((is_signed as u16) << (LIMB_BITS - 1)),
+                b[UINT_LIMBS - 1] + ((is_sign_comparison as u16) << (LIMB_BITS - 1)),
             )
         };
 

ceno_zkvm/src/instructions/riscv.rs

@@ -32,7 +32,11 @@ mod r_insn;
 
 mod ecall_insn;
 
+#[cfg(feature = "u16limb_circuit")]
+mod auipc;
 mod im_insn;
+#[cfg(feature = "u16limb_circuit")]
+mod lui;
 mod memory;
 mod s_insn;
 #[cfg(test)]