refactor lt config to gadget and blt

ceno_zkvm/src/chip_handler.rs

@@ -2,8 +2,7 @@ use ff_ext::ExtensionField;
 
 use crate::{
     error::ZKVMError,
-    expression::{Expression, WitIn},
-    instructions::riscv::config::ExprLtConfig,
+    expression::{Expression, WitIn}, gadgets::IsLtConfig,
 };
 
 pub mod general;
@@ -30,7 +29,7 @@ pub trait RegisterChipOperations<E: ExtensionField, NR: Into<String>, N: FnOnce(
         prev_ts: Expression<E>,
         ts: Expression<E>,
         value: RegisterExpr<E>,
-    ) -> Result<(Expression<E>, ExprLtConfig), ZKVMError>;
+    ) -> Result<(Expression<E>, IsLtConfig), ZKVMError>;
 
     #[allow(clippy::too_many_arguments)]
     fn register_write(
@@ -41,7 +40,7 @@ pub trait RegisterChipOperations<E: ExtensionField, NR: Into<String>, N: FnOnce(
         ts: Expression<E>,
         prev_values: RegisterExpr<E>,
         value: RegisterExpr<E>,
-    ) -> Result<(Expression<E>, ExprLtConfig), ZKVMError>;
+    ) -> Result<(Expression<E>, IsLtConfig), ZKVMError>;
 }
 
 /// The common representation of a memory value.
@@ -57,7 +56,7 @@ pub trait MemoryChipOperations<E: ExtensionField, NR: Into<String>, N: FnOnce()
         prev_ts: Expression<E>,
         ts: Expression<E>,
         value: crate::chip_handler::MemoryExpr<E>,
-    ) -> Result<(Expression<E>, ExprLtConfig), ZKVMError>;
+    ) -> Result<(Expression<E>, IsLtConfig), ZKVMError>;
 
     #[allow(clippy::too_many_arguments)]
     #[allow(dead_code)]
@@ -69,5 +68,5 @@ pub trait MemoryChipOperations<E: ExtensionField, NR: Into<String>, N: FnOnce()
         ts: Expression<E>,
         prev_values: crate::chip_handler::MemoryExpr<E>,
         value: crate::chip_handler::MemoryExpr<E>,
-    ) -> Result<(Expression<E>, ExprLtConfig), ZKVMError>;
+    ) -> Result<(Expression<E>, IsLtConfig), ZKVMError>;
 }

ceno_zkvm/src/chip_handler/general.rs

@@ -1,14 +1,12 @@
 use std::fmt::Display;
 
 use ff_ext::ExtensionField;
-use itertools::Itertools;
 
 use crate::{
-    chip_handler::utils::pows_expr,
     circuit_builder::{CircuitBuilder, ConstraintSystem},
     error::ZKVMError,
     expression::{Expression, Fixed, ToExpr, WitIn},
-    instructions::riscv::config::ExprLtConfig,
+    gadgets::IsLtConfig,
     structs::ROMType,
     tables::InsnRecord,
 };
@@ -322,64 +320,12 @@ impl<'a, E: ExtensionField> CircuitBuilder<'a, E> {
         lhs: Expression<E>,
         rhs: Expression<E>,
         assert_less_than: Option<bool>,
-    ) -> Result<ExprLtConfig, ZKVMError>
+    ) -> Result<IsLtConfig, ZKVMError>
     where
         NR: Into<String> + Display + Clone,
         N: FnOnce() -> NR,
     {
-        #[cfg(feature = "riv64")]
-        panic!("less_than is not supported for riv64 yet");
-
-        #[cfg(feature = "riv32")]
-        self.namespace(
-            || "less_than",
-            |cb| {
-                let name = name_fn();
-                let (is_lt, is_lt_expr) = if let Some(lt) = assert_less_than {
-                    (
-                        None,
-                        if lt {
-                            Expression::ONE
-                        } else {
-                            Expression::ZERO
-                        },
-                    )
-                } else {
-                    let is_lt = cb.create_witin(|| format!("{name} is_lt witin"))?;
-                    (Some(is_lt), is_lt.expr())
-                };
-
-                let mut witin_u16 = |var_name: String| -> Result<WitIn, ZKVMError> {
-                    cb.namespace(
-                        || format!("var {var_name}"),
-                        |cb| {
-                            let witin = cb.create_witin(|| var_name.to_string())?;
-                            cb.assert_ux::<_, _, 16>(|| name.clone(), witin.expr())?;
-                            Ok(witin)
-                        },
-                    )
-                };
-
-                let diff = (0..2)
-                    .map(|i| witin_u16(format!("diff_{i}")))
-                    .collect::<Result<Vec<WitIn>, _>>()?;
-
-                let pows = pows_expr((1 << u16::BITS).into(), diff.len());
-
-                let diff_expr = diff
-                    .iter()
-                    .zip_eq(pows)
-                    .map(|(record, beta)| beta * record.expr())
-                    .reduce(|a, b| a + b)
-                    .expect("reduce error");
-
-                let range = (1 << u32::BITS).into();
-
-                cb.require_equal(|| name.clone(), lhs - rhs, diff_expr - is_lt_expr * range)?;
-
-                Ok(ExprLtConfig { is_lt, diff })
-            },
-        )
+        IsLtConfig::construct_circuit(self, name_fn, lhs, rhs, assert_less_than)
     }
 
     pub(crate) fn is_equal(

ceno_zkvm/src/chip_handler/memory.rs

@@ -3,7 +3,7 @@ use crate::{
     circuit_builder::CircuitBuilder,
     error::ZKVMError,
     expression::{Expression, ToExpr, WitIn},
-    instructions::riscv::config::ExprLtConfig,
+    gadgets::IsLtConfig,
     structs::RAMType,
 };
 use ff_ext::ExtensionField;
@@ -19,7 +19,7 @@ impl<'a, E: ExtensionField, NR: Into<String>, N: FnOnce() -> NR> MemoryChipOpera
         prev_ts: Expression<E>,
         ts: Expression<E>,
         value: MemoryExpr<E>,
-    ) -> Result<(Expression<E>, ExprLtConfig), ZKVMError> {
+    ) -> Result<(Expression<E>, IsLtConfig), ZKVMError> {
         self.namespace(name_fn, |cb| {
             // READ (a, v, t)
             let read_record = cb.rlc_chip_record(
@@ -66,7 +66,7 @@ impl<'a, E: ExtensionField, NR: Into<String>, N: FnOnce() -> NR> MemoryChipOpera
         ts: Expression<E>,
         prev_values: MemoryExpr<E>,
         value: MemoryExpr<E>,
-    ) -> Result<(Expression<E>, ExprLtConfig), ZKVMError> {
+    ) -> Result<(Expression<E>, IsLtConfig), ZKVMError> {
         self.namespace(name_fn, |cb| {
             // READ (a, v, t)
             let read_record = cb.rlc_chip_record(

ceno_zkvm/src/chip_handler/register.rs

@@ -4,7 +4,7 @@ use crate::{
     circuit_builder::CircuitBuilder,
     error::ZKVMError,
     expression::{Expression, ToExpr, WitIn},
-    instructions::riscv::config::ExprLtConfig,
+    gadgets::IsLtConfig,
     structs::RAMType,
 };
 
@@ -20,7 +20,7 @@ impl<'a, E: ExtensionField, NR: Into<String>, N: FnOnce() -> NR> RegisterChipOpe
         prev_ts: Expression<E>,
         ts: Expression<E>,
         value: RegisterExpr<E>,
-    ) -> Result<(Expression<E>, ExprLtConfig), ZKVMError> {
+    ) -> Result<(Expression<E>, IsLtConfig), ZKVMError> {
         self.namespace(name_fn, |cb| {
             // READ (a, v, t)
             let read_record = cb.rlc_chip_record(
@@ -66,7 +66,7 @@ impl<'a, E: ExtensionField, NR: Into<String>, N: FnOnce() -> NR> RegisterChipOpe
         ts: Expression<E>,
         prev_values: RegisterExpr<E>,
         value: RegisterExpr<E>,
-    ) -> Result<(Expression<E>, ExprLtConfig), ZKVMError> {
+    ) -> Result<(Expression<E>, IsLtConfig), ZKVMError> {
         self.namespace(name_fn, |cb| {
             // READ (a, v, t)
             let read_record = cb.rlc_chip_record(

ceno_zkvm/src/gadgets/is_zero.rs

@@ -20,20 +20,23 @@ impl IsZeroConfig {
         self.is_zero.expr()
     }
 
-    pub fn construct_circuit<E: ExtensionField>(
+    pub fn construct_circuit<E: ExtensionField, NR: Into<String>, N: FnOnce() -> NR>(
         cb: &mut CircuitBuilder<E>,
+        name_fn: N,
         x: Expression<E>,
     ) -> Result<Self, ZKVMError> {
-        let is_zero = cb.create_witin(|| "is_zero")?;
-        let inverse = cb.create_witin(|| "inv")?;
+        cb.namespace(name_fn, |cb| {
+            let is_zero = cb.create_witin(|| "is_zero")?;
+            let inverse = cb.create_witin(|| "inv")?;
 
-        // x==0 => is_zero=1
-        cb.require_one(|| "is_zero_1", is_zero.expr() + x.clone() * inverse.expr())?;
+            // x==0 => is_zero=1
+            cb.require_one(|| "is_zero_1", is_zero.expr() + x.clone() * inverse.expr())?;
 
-        // x!=0 => is_zero=0
-        cb.require_zero(|| "is_zero_0", is_zero.expr() * x.clone())?;
+            // x!=0 => is_zero=0
+            cb.require_zero(|| "is_zero_0", is_zero.expr() * x.clone())?;
 
-        Ok(IsZeroConfig { is_zero, inverse })
+            Ok(IsZeroConfig { is_zero, inverse })
+        })
     }
 
     pub fn assign_instance<F: SmallField>(
@@ -61,12 +64,17 @@ impl IsEqualConfig {
         self.0.expr()
     }
 
-    pub fn construct_circuit<E: ExtensionField>(
+    pub fn construct_circuit<E: ExtensionField, NR: Into<String>, N: FnOnce() -> NR>(
         cb: &mut CircuitBuilder<E>,
+        name_fn: N,
         a: Expression<E>,
         b: Expression<E>,
     ) -> Result<Self, ZKVMError> {
-        Ok(IsEqualConfig(IsZeroConfig::construct_circuit(cb, a - b)?))
+        Ok(IsEqualConfig(IsZeroConfig::construct_circuit(
+            cb,
+            name_fn,
+            a - b,
+        )?))
     }
 
     pub fn assign_instance<F: SmallField>(

ceno_zkvm/src/gadgets/mod.rs

@@ -1,2 +1,4 @@
+mod is_lt;
 mod is_zero;
+pub use is_lt::IsLtConfig;
 pub use is_zero::{IsEqualConfig, IsZeroConfig};

ceno_zkvm/src/instructions/riscv/b_insn.rs

@@ -3,15 +3,15 @@
 use ceno_emul::{InsnKind, StepRecord};
 use ff_ext::ExtensionField;
 
-use super::{config::ExprLtConfig, constants::PC_STEP_SIZE};
+use super::constants::PC_STEP_SIZE;
 use crate::{
     chip_handler::{
         GlobalStateRegisterMachineChipOperations, RegisterChipOperations, RegisterExpr,
     },
     circuit_builder::CircuitBuilder,
     error::ZKVMError,
     expression::{Expression, ToExpr, WitIn},
-    instructions::riscv::config::ExprLtInput,
+    gadgets::IsLtConfig,
     set_val,
     tables::InsnRecord,
     witness::LkMultiplicity,
@@ -45,8 +45,8 @@ pub struct BInstructionConfig {
     imm: WitIn,
     prev_rs1_ts: WitIn,
     prev_rs2_ts: WitIn,
-    lt_rs1_cfg: ExprLtConfig,
-    lt_rs2_cfg: ExprLtConfig,
+    lt_rs1_cfg: IsLtConfig,
+    lt_rs2_cfg: IsLtConfig,
 }
 
 impl BInstructionConfig {
@@ -160,16 +160,18 @@ impl BInstructionConfig {
         );
 
         // 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.rs2().unwrap().previous_cycle,
-            rhs: step.cycle() + 1,
-        }
-        .assign(instance, &self.lt_rs2_cfg, lk_multiplicity);
+        self.lt_rs1_cfg.assign_instance(
+            instance,
+            lk_multiplicity,
+            step.rs1().unwrap().previous_cycle,
+            step.cycle(),
+        )?;
+        self.lt_rs2_cfg.assign_instance(
+            instance,
+            lk_multiplicity,
+            step.rs2().unwrap().previous_cycle,
+            step.cycle() + 1,
+        )?;
 
         Ok(())
     }

ceno_zkvm/src/instructions/riscv/blt.rs

@@ -3,26 +3,19 @@ use ceno_emul::InsnKind;
 use ff_ext::ExtensionField;
 
 use crate::{
-    circuit_builder::CircuitBuilder,
-    error::ZKVMError,
-    expression::ToExpr,
-    instructions::{
-        riscv::config::{UIntLtConfig, UIntLtInput},
-        Instruction,
-    },
-    utils::{i64_to_base, split_to_u8},
-    witness::LkMultiplicity,
+    circuit_builder::CircuitBuilder, error::ZKVMError, gadgets::IsLtConfig,
+    instructions::Instruction, witness::LkMultiplicity, Value,
 };
 
-use super::{b_insn::BInstructionConfig, constants::UInt8, RIVInstruction};
+use super::{b_insn::BInstructionConfig, constants::UInt, RIVInstruction};
 
 pub struct BltInstruction;
 
 pub struct InstructionConfig<E: ExtensionField> {
     pub b_insn: BInstructionConfig,
-    pub read_rs1: UInt8<E>,
-    pub read_rs2: UInt8<E>,
-    pub is_lt: UIntLtConfig,
+    pub read_rs1: UInt<E>,
+    pub read_rs2: UInt<E>,
+    pub is_lt: IsLtConfig,
 }
 
 impl RIVInstruction for BltInstruction {
@@ -38,16 +31,17 @@ impl<E: ExtensionField> Instruction<E> for BltInstruction {
     fn construct_circuit(
         circuit_builder: &mut CircuitBuilder<E>,
     ) -> Result<InstructionConfig<E>, ZKVMError> {
-        let read_rs1 = UInt8::new_unchecked(|| "rs1_limbs", circuit_builder)?;
-        let read_rs2 = UInt8::new_unchecked(|| "rs2_limbs", circuit_builder)?;
-        let is_lt = read_rs1.lt_limb8(circuit_builder, &read_rs2)?;
+        let read_rs1 = UInt::new_unchecked(|| "rs1_limbs", circuit_builder)?;
+        let read_rs2 = UInt::new_unchecked(|| "rs2_limbs", circuit_builder)?;
+        let is_lt =
+            circuit_builder.less_than(|| "rs1<rs2", read_rs1.value(), read_rs2.value(), None)?;
 
         let b_insn = BInstructionConfig::construct_circuit(
             circuit_builder,
             Self::INST_KIND,
             read_rs1.register_expr(),
             read_rs2.register_expr(),
-            is_lt.is_lt.expr(),
+            is_lt.expr(),
         )?;
 
         Ok(InstructionConfig {
@@ -64,25 +58,16 @@ impl<E: ExtensionField> Instruction<E> for BltInstruction {
         lk_multiplicity: &mut LkMultiplicity,
         step: &ceno_emul::StepRecord,
     ) -> Result<(), ZKVMError> {
-        let rs1_limbs = split_to_u8(step.rs1().unwrap().value);
-        let rs2_limbs = split_to_u8(step.rs2().unwrap().value);
-        config.read_rs1.assign_limbs(instance, {
-            rs1_limbs
-                .iter()
-                .map(|&limb| i64_to_base::<E::BaseField>(limb as i64))
-                .collect()
-        });
-        config.read_rs2.assign_limbs(instance, {
-            rs2_limbs
-                .iter()
-                .map(|&limb| i64_to_base::<E::BaseField>(limb as i64))
-                .collect()
-        });
-        let lt_input = UIntLtInput {
-            lhs_limbs: &rs1_limbs,
-            rhs_limbs: &rs2_limbs,
-        };
-        lt_input.assign(instance, &config.is_lt, lk_multiplicity);
+        let rs1 = Value::new_unchecked(step.rs1().unwrap().value);
+        let rs2 = Value::new_unchecked(step.rs2().unwrap().value);
+        config.read_rs1.assign_limbs(instance, rs1.u16_fields());
+        config.read_rs2.assign_limbs(instance, rs2.u16_fields());
+        config.is_lt.assign_instance(
+            instance,
+            lk_multiplicity,
+            step.rs1().unwrap().value as u64,
+            step.rs2().unwrap().value as u64,
+        )?;
 
         config
             .b_insn