Add BGEU Chip (#337)

Summary:

Fixes #304

Test Plan:

Co-authored-by: duc-nx <>

Author: duc-nx

prover/src/chips/bgeu.rs

@@ -0,0 +1,255 @@
+use num_traits::One;
+use stwo_prover::constraint_framework::{logup::LookupElements, EvalAtRow};
+
+use nexus_vm::{riscv::BuiltinOpcode, WORD_SIZE};
+
+use crate::{
+    column::Column::{self, *},
+    components::MAX_LOOKUP_TUPLE_SIZE,
+    trace::{
+        eval::{trace_eval, TraceEval},
+        sidenote::SideNote,
+        BoolWord, ProgramStep, Traces, Word,
+    },
+    traits::{ExecuteChip, MachineChip},
+};
+
+use super::add::{self};
+
+pub struct ExecutionResult {
+    pub diff_bytes: Word,
+    pub borrow_bits: BoolWord,
+    pub pc_next: Word,
+    pub carry_bits: BoolWord,
+}
+
+pub struct BgeuChip;
+
+impl ExecuteChip for BgeuChip {
+    type ExecutionResult = ExecutionResult;
+
+    fn execute(program_step: &ProgramStep) -> Self::ExecutionResult {
+        let value_a = program_step.get_value_a();
+        let value_b = program_step.get_value_b();
+        let imm = program_step.get_value_c().0;
+        let pc = program_step.step.pc.to_le_bytes();
+
+        let (diff_bytes, borrow_bits) = super::sub::subtract_with_borrow(value_a, value_b);
+
+        // ltu_flag is equal to borrow_bit[3]
+        let (pc_next, carry_bits) = if borrow_bits[3] {
+            // a < b is true: pc_next = pc + 4
+            add::add_with_carries(pc, 4u32.to_le_bytes())
+        } else {
+            // a >= b is true: pc_next = pc + imm
+            add::add_with_carries(pc, imm)
+        };
+
+        ExecutionResult {
+            diff_bytes,
+            borrow_bits,
+            pc_next,
+            carry_bits,
+        }
+    }
+}
+
+impl MachineChip for BgeuChip {
+    fn fill_main_trace(
+        traces: &mut Traces,
+        row_idx: usize,
+        vm_step: &ProgramStep,
+        _side_note: &mut SideNote,
+    ) {
+        if !matches!(
+            vm_step.step.instruction.opcode.builtin(),
+            Some(BuiltinOpcode::BGEU)
+        ) {
+            return;
+        }
+
+        let ExecutionResult {
+            diff_bytes,
+            borrow_bits,
+            pc_next,
+            carry_bits,
+        } = Self::execute(vm_step);
+
+        traces.fill_columns(row_idx, diff_bytes, Column::Helper1);
+        traces.fill_columns(row_idx, borrow_bits, Column::BorrowFlag);
+
+        // Fill valueA
+        traces.fill_columns(row_idx, vm_step.get_value_a(), Column::ValueA);
+
+        // Fill PcNext and CarryFlag, since Pc and Immediate are filled to the main trace in CPU.
+        traces.fill_columns(row_idx, pc_next, Column::PcNext);
+        traces.fill_columns(row_idx, carry_bits, Column::CarryFlag);
+    }
+
+    fn add_constraints<E: EvalAtRow>(
+        eval: &mut E,
+        trace_eval: &TraceEval<E>,
+        _lookup_elements: &LookupElements<MAX_LOOKUP_TUPLE_SIZE>,
+    ) {
+        let modulus = E::F::from(256u32.into());
+        let (value_a, _) = trace_eval!(trace_eval, ValueA);
+        let (value_b, _) = trace_eval!(trace_eval, ValueB);
+        let (value_c, _) = trace_eval!(trace_eval, ValueC);
+        let (pc, _) = trace_eval!(trace_eval, Column::Pc);
+        let (carry_bits, _) = trace_eval!(trace_eval, Column::CarryFlag);
+        let (borrow_bits, _) = trace_eval!(trace_eval, Column::BorrowFlag);
+        let (diff_bytes, _) = trace_eval!(trace_eval, Column::Helper1);
+        let (pc_next, _) = trace_eval!(trace_eval, Column::PcNext);
+        let ([is_bgeu], _) = trace_eval!(trace_eval, Column::IsBgeu);
+        let ltu_flag = borrow_bits[3].clone();
+
+        // is_bgeu・(a_val_1 - b_val_1 - h1_1 + borrow_1・2^8) = 0
+        // is_bgeu・(a_val_2 - b_val_2 - h1_2 + borrow_2・2^8 - borrow_1) = 0
+        // is_bgeu・(a_val_3 - b_val_3 - h1_3 + borrow_3・2^8 - borrow_2) = 0
+        // is_bgeu・(a_val_4 - b_val_4 - h1_4 + ltu_flag・2^8 - borrow_3) = 0
+        eval.add_constraint(
+            is_bgeu.clone()
+                * (value_a[0].clone() - value_b[0].clone() - diff_bytes[0].clone()
+                    + borrow_bits[0].clone() * modulus.clone()),
+        );
+        for i in 1..WORD_SIZE {
+            eval.add_constraint(
+                is_bgeu.clone()
+                    * (value_a[i].clone() - value_b[i].clone() - diff_bytes[i].clone()
+                        + borrow_bits[i].clone() * modulus.clone()
+                        - borrow_bits[i - 1].clone()),
+            );
+        }
+
+        // is_bgeu・( (1 - ltu_flag)・c_val_1 + ltu_flag・4 + pc_1 - carry_1·2^8 - pc_next_1) =0
+        // is_bgeu・( (1 - ltu_flag)・c_val_2 + pc_2 + carry_1 - carry_2·2^8 - pc_next_2) = 0
+        // is_bgeu・( (1 - ltu_flag)・c_val_3 + pc_3 + carry_2 - carry_3·2^8 - pc_next_3) = 0
+        // is_bgeu・( (1 - ltu_flag)・c_val_4 + pc_4 + carry_3 - carry_4·2^8 - pc_next_4) = 0
+        eval.add_constraint(
+            is_bgeu.clone()
+                * ((E::F::one() - ltu_flag.clone()) * value_c[0].clone()
+                    + ltu_flag.clone() * E::F::from(4u32.into())
+                    + pc[0].clone()
+                    - carry_bits[0].clone() * modulus.clone()
+                    - pc_next[0].clone()),
+        );
+        for i in 1..WORD_SIZE {
+            eval.add_constraint(
+                is_bgeu.clone()
+                    * ((E::F::one() - ltu_flag.clone()) * value_c[i].clone()
+                        + pc[i].clone()
+                        + carry_bits[i - 1].clone()
+                        - carry_bits[i].clone() * modulus.clone()
+                        - pc_next[i].clone()),
+            );
+        }
+    }
+}
+
+#[cfg(test)]
+mod test {
+    use crate::{
+        chips::{AddChip, CpuChip, RegisterMemCheckChip, SubChip},
+        test_utils::assert_chip,
+        trace::{program::iter_program_steps, PreprocessedTraces},
+    };
+
+    use super::*;
+    use nexus_vm::{
+        riscv::{BasicBlock, BuiltinOpcode, Instruction, InstructionType, Opcode},
+        trace::k_trace_direct,
+    };
+
+    const LOG_SIZE: u32 = PreprocessedTraces::MIN_LOG_SIZE;
+
+    #[rustfmt::skip]
+    fn setup_basic_block_ir() -> Vec<BasicBlock> {
+        let basic_block = BasicBlock::new(vec![
+            // Set x10 = 1
+            Instruction::new(Opcode::from(BuiltinOpcode::ADDI), 10, 0, 1, InstructionType::IType),
+            // Set x1 = 10
+            Instruction::new(Opcode::from(BuiltinOpcode::ADDI), 1, 0, 10, InstructionType::IType),
+            // Set x2 = 20
+            Instruction::new(Opcode::from(BuiltinOpcode::ADDI), 2, 0, 20, InstructionType::IType),
+            // Set x3 = 10 (same as x1)
+            Instruction::new(Opcode::from(BuiltinOpcode::ADDI), 3, 0, 10, InstructionType::IType),
+            // Set x4 = -10
+            Instruction::new(Opcode::from(BuiltinOpcode::SUB), 4, 0, 1, InstructionType::RType),
+            // Set x5 = 0xFFFFFFFF (max unsigned value)
+            Instruction::new(Opcode::from(BuiltinOpcode::SUB), 5, 0, 10, InstructionType::RType),
+    
+            // Case 1: BGEU with equal values (should branch)
+            // BGEU x1, x3, 0xff (should branch as x1 >= x3 is true)
+            Instruction::new(Opcode::from(BuiltinOpcode::BGEU), 1, 3, 12, InstructionType::BType),
+            
+            // Unimpl instructions to fill the gap (trigger error when executed)
+            Instruction::unimpl(),
+            Instruction::unimpl(),
+
+            // Case 2: BGEU with different values (should not branch)
+            // BGEU x1, x2, 12 (should not branch as x1 >= x2 is false)
+            Instruction::new(Opcode::from(BuiltinOpcode::BGEU), 1, 2, 0xff, InstructionType::BType),
+
+
+            // Case 3: BGEU with zero and non-zero (should not branch)
+            // BGEU x0, x1, 8 (should not branch as x0 >= x1 is false)
+            Instruction::new(Opcode::from(BuiltinOpcode::BGEU), 0, 1, 0xff, InstructionType::BType),
+
+            // Case 4: BGEU with zero and zero (should branch)
+            // BGEU x0, x0, 0xff (should branch as x0 >= x0 is true)
+            Instruction::new(Opcode::from(BuiltinOpcode::BGEU), 0, 0, 12, InstructionType::BType),
+
+            // Unimpl instructions to fill the gap (trigger error when executed)
+            Instruction::unimpl(),
+            Instruction::unimpl(),
+
+            // Case 5: BGEU with negative and positive values (should branch)
+            // BGEU x4, x1, 0xff (should branch as 0xfffffff6 >= 10 unsigned)
+            Instruction::new(Opcode::from(BuiltinOpcode::BGEU), 4, 1, 12, InstructionType::BType),
+
+            // Unimpl instructions to fill the gap (trigger error when executed)
+            Instruction::unimpl(),
+            Instruction::unimpl(),
+
+            // Case 6: BGEU with max unsigned value and zero (should branch)
+            // BGEU x5, x0, 0xff (should branch as 0xFFFFFFFF >= 0)
+            Instruction::new(Opcode::from(BuiltinOpcode::BGEU), 5, 0, 12, InstructionType::BType),
+
+            // Unimpl instructions to fill the gap (trigger error when executed)
+            Instruction::unimpl(),
+            Instruction::unimpl(),
+
+            // Case 7: BGEU with zero and max unsigned value (should not branch)
+            // BGEU x0, x5, 12 (should not branch as 0 >= 0xFFFFFFFF is false)
+            Instruction::new(Opcode::from(BuiltinOpcode::BGEU), 0, 5, 0xff, InstructionType::BType),
+        ]);
+        vec![basic_block]
+    }
+
+    #[test]
+    fn test_k_trace_constrained_bgeu_instructions() {
+        type Chips = (CpuChip, AddChip, SubChip, BgeuChip, RegisterMemCheckChip);
+        let basic_block = setup_basic_block_ir();
+        let k = 1;
+
+        // Get traces from VM K-Trace interface
+        let vm_traces = k_trace_direct(&basic_block, k).expect("Failed to create trace");
+
+        // Trace circuit
+        let mut traces = Traces::new(LOG_SIZE);
+        let mut side_note = SideNote::default();
+        let program_steps = iter_program_steps(&vm_traces, traces.num_rows());
+
+        // We iterate each block in the trace for each instruction
+        for (row_idx, program_step) in program_steps.enumerate() {
+            Chips::fill_main_trace(&mut traces, row_idx, &program_step, &mut side_note);
+        }
+
+        let mut preprocessed_column = PreprocessedTraces::empty(LOG_SIZE);
+        preprocessed_column.fill_is_first();
+        preprocessed_column.fill_is_first32();
+        preprocessed_column.fill_row_idx();
+        preprocessed_column.fill_timestamps();
+        assert_chip::<Chips>(traces, Some(preprocessed_column));
+    }
+}

prover/src/chips/cpu.rs

@@ -82,6 +82,9 @@ impl MachineChip for CpuChip {
             Some(BuiltinOpcode::BLTU) => {
                 traces.fill_columns(row_idx, true, IsBltu);
             }
+            Some(BuiltinOpcode::BGEU) => {
+                traces.fill_columns(row_idx, true, IsBgeu);
+            }
             _ => {
                 if !step.is_padding {
                     panic!(
@@ -254,6 +257,7 @@ impl MachineChip for CpuChip {
         let ([is_bne], _) = trace_eval!(trace_eval, IsBne);
         let ([is_beq], _) = trace_eval!(trace_eval, IsBeq);
         let ([is_bltu], _) = trace_eval!(trace_eval, IsBltu);
+        let ([is_bgeu], _) = trace_eval!(trace_eval, IsBgeu);
         let ([is_padding], _) = trace_eval!(trace_eval, IsPadding);
         eval.add_constraint(
             is_add.clone()
@@ -266,6 +270,7 @@ impl MachineChip for CpuChip {
                 + is_bne.clone()
                 + is_beq.clone()
                 + is_bltu.clone()
+                + is_bgeu.clone()
                 + is_padding
                 - E::F::one(),
         );
@@ -316,7 +321,7 @@ impl MachineChip for CpuChip {
         eval.add_constraint((is_type_r + is_type_i) * (op_a.clone() - reg3_address.clone()));
 
         // is_type_b = is_beq + is_bne + is_blt + is_bge + is_bltu + is_bgeu
-        let is_type_b = is_beq + is_bne + is_bltu; // TODO: add more flags when they are available
+        let is_type_b = is_beq + is_bne + is_bltu + is_bgeu; // TODO: add more flags when they are available
 
         // is_type_s = is_sb + is_sh + is_sw
         // TODO: define is_type_s when flags are available

prover/src/chips/mod.rs

@@ -1,5 +1,6 @@
 mod add;
 mod beq;
+mod bgeu;
 mod bit_op;
 mod bltu;
 mod bne;
@@ -17,8 +18,9 @@ mod sltu;
 mod timestamp;
 
 pub use self::{
-    add::AddChip, beq::BeqChip, bit_op::BitOpChip, bltu::BltuChip, bne::BneChip, cpu::CpuChip,
-    prog_mem_check::ProgramMemCheckChip, range128::Range128Chip, range256::Range256Chip,
-    range32::Range32Chip, range_bool::RangeBoolChip, register_mem_check::RegisterMemCheckChip,
-    slt::SltChip, sltu::SltuChip, sub::SubChip, timestamp::TimestampChip,
+    add::AddChip, beq::BeqChip, bgeu::BgeuChip, bit_op::BitOpChip, bltu::BltuChip, bne::BneChip,
+    cpu::CpuChip, prog_mem_check::ProgramMemCheckChip, range128::Range128Chip,
+    range256::Range256Chip, range32::Range32Chip, range_bool::RangeBoolChip,
+    register_mem_check::RegisterMemCheckChip, slt::SltChip, sltu::SltuChip, sub::SubChip,
+    timestamp::TimestampChip,
 };

prover/src/chips/range_bool.rs

@@ -7,7 +7,7 @@ use num_traits::One;
 use crate::{
     column::Column::{
         self, BorrowFlag, CH1Minus, CH2Minus, CH3Minus, CarryFlag, ImmB, ImmC, IsAdd, IsAnd,
-        IsBltu, IsOr, IsPadding, IsSlt, IsSltu, IsSub, IsXor, PrgMemoryFlag, Ram1Accessed,
+        IsBgeu, IsBltu, IsOr, IsPadding, IsSlt, IsSltu, IsSub, IsXor, PrgMemoryFlag, Ram1Accessed,
         Ram2Accessed, Ram3Accessed, Ram4Accessed, Reg1Accessed, Reg2Accessed, Reg3Accessed, SgnB,
         SgnC,
     },
@@ -23,7 +23,7 @@ use crate::{
 
 pub struct RangeBoolChip;
 
-const CHECKED_SINGLE: [Column; 21] = [
+const CHECKED_SINGLE: [Column; 22] = [
     ImmB,
     ImmC,
     IsAdd,
@@ -34,6 +34,7 @@ const CHECKED_SINGLE: [Column; 21] = [
     IsSltu,
     IsSlt,
     IsBltu,
+    IsBgeu,
     IsPadding,
     SgnB,
     SgnC,

prover/src/column.rs

@@ -95,6 +95,9 @@ pub enum Column {
     /// Boolean flag on whether the row is a BLTU.
     #[size = 1]
     IsBltu,
+    /// Boolean flag on whether the row is a BGEU.
+    #[size = 1]
+    IsBgeu,
     /// Boolean flag on whether the row is a padding.
     #[size = 1]
     IsPadding,

prover/src/lib.rs

@@ -28,7 +28,7 @@ mod test_utils;
 pub(crate) use nexus_vm::WORD_SIZE;
 
 use chips::{
-    AddChip, BeqChip, BitOpChip, BltuChip, BneChip, CpuChip, Range128Chip, Range256Chip,
+    AddChip, BeqChip, BgeuChip, BitOpChip, BltuChip, BneChip, CpuChip, Range128Chip, Range256Chip,
     Range32Chip, RangeBoolChip, SltChip, SltuChip, SubChip, TimestampChip,
 };
 use components::{MachineComponent, MachineEval, LOG_CONSTRAINT_DEGREE};
@@ -44,6 +44,7 @@ pub type Components = (
     BneChip,
     BeqChip,
     BltuChip,
+    BgeuChip,
     TimestampChip,
     // Range checks must be positioned at the end. They use values filled by instruction chips.
     RangeBoolChip,