beq/bne circuit to limb (#1155)

hero78119 · web-flow · commit c2b73d692d79 · 2025-11-27T05:13:49.000Z
some left-over from previous migration
diff --git a/ceno_zkvm/src/instructions/riscv/branch.rs b/ceno_zkvm/src/instructions/riscv/branch.rs
@@ -1,7 +1,9 @@
 use super::RIVInstruction;
 use ceno_emul::InsnKind;
 
+#[cfg(not(feature = "u16limb_circuit"))]
 mod branch_circuit;
+#[cfg(feature = "u16limb_circuit")]
 mod branch_circuit_v2;
 #[cfg(test)]
 mod test;
@@ -10,12 +12,18 @@ pub struct BeqOp;
 impl RIVInstruction for BeqOp {
     const INST_KIND: InsnKind = InsnKind::BEQ;
 }
+#[cfg(feature = "u16limb_circuit")]
+pub type BeqInstruction<E> = branch_circuit_v2::BranchCircuit<E, BeqOp>;
+#[cfg(not(feature = "u16limb_circuit"))]
 pub type BeqInstruction<E> = branch_circuit::BranchCircuit<E, BeqOp>;
 
 pub struct BneOp;
 impl RIVInstruction for BneOp {
     const INST_KIND: InsnKind = InsnKind::BNE;
 }
+#[cfg(feature = "u16limb_circuit")]
+pub type BneInstruction<E> = branch_circuit_v2::BranchCircuit<E, BneOp>;
+#[cfg(not(feature = "u16limb_circuit"))]
 pub type BneInstruction<E> = branch_circuit::BranchCircuit<E, BneOp>;
 
 pub struct BltuOp;
diff --git a/ceno_zkvm/src/instructions/riscv/branch/branch_circuit_v2.rs b/ceno_zkvm/src/instructions/riscv/branch/branch_circuit_v2.rs
@@ -6,15 +6,21 @@ use crate::{
     gadgets::{UIntLimbsLT, UIntLimbsLTConfig},
     instructions::{
         Instruction,
-        riscv::{RIVInstruction, b_insn::BInstructionConfig, constants::UInt},
+        riscv::{
+            RIVInstruction,
+            b_insn::BInstructionConfig,
+            constants::{UINT_LIMBS, UInt},
+        },
     },
     structs::ProgramParams,
     witness::LkMultiplicity,
 };
 use ceno_emul::{InsnKind, StepRecord};
-use ff_ext::ExtensionField;
-use multilinear_extensions::Expression;
-use std::marker::PhantomData;
+use ff_ext::{ExtensionField, FieldInto};
+use multilinear_extensions::{Expression, ToExpr, WitIn};
+use p3::field::{Field, FieldAlgebra};
+use std::{array, marker::PhantomData};
+use witness::set_val;
 
 pub struct BranchCircuit<E, I>(PhantomData<(E, I)>);
 
@@ -23,7 +29,11 @@ pub struct BranchConfig<E: ExtensionField> {
     pub read_rs1: UInt<E>,
     pub read_rs2: UInt<E>,
 
-    pub uint_lt_config: UIntLimbsLTConfig<E>,
+    // for non eq opcode config
+    pub uint_lt_config: Option<UIntLimbsLTConfig<E>>,
+    // for beq/bne
+    pub eq_diff_inv_marker: Option<[WitIn; UINT_LIMBS]>,
+    pub eq_branch_taken_bit: Option<WitIn>,
     phantom: PhantomData<E>,
 }
 
@@ -41,28 +51,82 @@ impl<E: ExtensionField, I: RIVInstruction> Instruction<E> for BranchCircuit<E, I
         let read_rs1 = UInt::new_unchecked(|| "rs1_limbs", circuit_builder)?;
         let read_rs2 = UInt::new_unchecked(|| "rs2_limbs", circuit_builder)?;
 
-        let is_signed = matches!(I::INST_KIND, InsnKind::BLT | InsnKind::BGE);
-        let is_ge = matches!(I::INST_KIND, InsnKind::BGEU | InsnKind::BGE);
-        let uint_lt_config =
-            UIntLimbsLT::<E>::construct_circuit(circuit_builder, &read_rs1, &read_rs2, is_signed)?;
-        let branch_taken_bit = if is_ge {
-            Expression::ONE - uint_lt_config.is_lt()
-        } else {
-            uint_lt_config.is_lt()
-        };
+        let (branch_taken_bit_expr, eq_branch_taken_bit, eq_diff_inv_marker, uint_lt_config) =
+            if matches!(I::INST_KIND, InsnKind::BEQ | InsnKind::BNE) {
+                let branch_taken_bit = circuit_builder.create_bit(|| "branch_taken_bit")?;
+                let eq_diff_inv_marker = array::from_fn(|i| {
+                    circuit_builder.create_witin(|| format!("eq_diff_inv_marker_{i}"))
+                });
+
+                // 1 if cmp_result indicates a and b are EQUAL, 0 otherwise
+                let cmp_eq = match I::INST_KIND {
+                    InsnKind::BEQ => branch_taken_bit.expr(),
+                    InsnKind::BNE => Expression::ONE - branch_taken_bit.expr(),
+                    _ => unreachable!(),
+                };
+                let mut sum = cmp_eq.expr();
+
+                // For BEQ, inv_marker is used to check equality of a and b:
+                // - If a == b, all inv_marker values must be 0 (sum = 0)
+                // - If a != b, inv_marker contains 0s for all positions except ONE position i where a[i] !=
+                //   b[i]
+                // - At this position, inv_marker[i] contains the multiplicative inverse of (a[i] - b[i])
+                // - This ensures inv_marker[i] * (a[i] - b[i]) = 1, making the sum = 1
+                // Note: There might be multiple valid inv_marker if a != b.
+                // But as long as the trace can provide at least one, that’s sufficient to prove a != b.
+                //
+                // Note:
+                // - If cmp_eq == 0, then it is impossible to have sum != 0 if a == b.
+                // - If cmp_eq == 1, then it is impossible for a[i] - b[i] == 0 to pass for all i if a != b.
+                #[allow(clippy::needless_range_loop)]
+                for i in 0..UINT_LIMBS {
+                    sum += (read_rs1.limbs[i].expr() - read_rs2.limbs[i].expr())
+                        * eq_diff_inv_marker[i].expr();
+                    circuit_builder.require_zero(
+                        || "require_zero",
+                        cmp_eq.expr() * (read_rs1.limbs[i].expr() - read_rs2.limbs[i].expr()),
+                    )?
+                }
+                circuit_builder.require_one(|| "sum", sum)?;
+
+                (
+                    branch_taken_bit.expr(),
+                    Some(branch_taken_bit),
+                    Some(eq_diff_inv_marker),
+                    None,
+                )
+            } else {
+                let is_signed = matches!(I::INST_KIND, InsnKind::BLT | InsnKind::BGE);
+                let is_ge = matches!(I::INST_KIND, InsnKind::BGEU | InsnKind::BGE);
+                let uint_lt_config = UIntLimbsLT::<E>::construct_circuit(
+                    circuit_builder,
+                    &read_rs1,
+                    &read_rs2,
+                    is_signed,
+                )?;
+                let branch_taken_bit = if is_ge {
+                    Expression::ONE - uint_lt_config.is_lt()
+                } else {
+                    uint_lt_config.is_lt()
+                };
+                (branch_taken_bit, None, None, Some(uint_lt_config))
+            };
+
         let b_insn = BInstructionConfig::construct_circuit(
             circuit_builder,
             I::INST_KIND,
             read_rs1.register_expr(),
             read_rs2.register_expr(),
-            branch_taken_bit,
+            branch_taken_bit_expr,
         )?;
 
         Ok(BranchConfig {
             b_insn,
             read_rs1,
             read_rs2,
             uint_lt_config,
+            eq_branch_taken_bit,
+            eq_diff_inv_marker,
             phantom: Default::default(),
         })
     }
@@ -85,15 +149,54 @@ impl<E: ExtensionField, I: RIVInstruction> Instruction<E> for BranchCircuit<E, I
         config.read_rs1.assign_limbs(instance, rs1_limbs);
         config.read_rs2.assign_limbs(instance, rs2_limbs);
 
-        let is_signed = matches!(step.insn().kind, InsnKind::BLT | InsnKind::BGE);
-        UIntLimbsLT::<E>::assign(
-            &config.uint_lt_config,
-            instance,
-            lk_multiplicity,
-            rs1_limbs,
-            rs2_limbs,
-            is_signed,
-        )?;
+        if matches!(I::INST_KIND, InsnKind::BEQ | InsnKind::BNE) {
+            // Returns (branch_taken, diff_idx, x[diff_idx] - y[diff_idx])
+            #[inline(always)]
+            fn run_eq<F, const NUM_LIMBS: usize>(
+                is_beq: bool,
+                x: &[u16],
+                y: &[u16],
+            ) -> (bool, usize, F)
+            where
+                F: FieldAlgebra + Field,
+            {
+                for i in 0..NUM_LIMBS {
+                    if x[i] != y[i] {
+                        return (
+                            !is_beq,
+                            i,
+                            (F::from_canonical_u16(x[i]) - F::from_canonical_u16(y[i])).inverse(),
+                        );
+                    }
+                }
+                (is_beq, 0, F::ZERO)
+            }
+            let (branch_taken, diff_idx, diff_inv_val) = run_eq::<E::BaseField, UINT_LIMBS>(
+                matches!(I::INST_KIND, InsnKind::BEQ),
+                rs1_limbs,
+                rs2_limbs,
+            );
+            set_val!(
+                instance,
+                config.eq_branch_taken_bit.as_ref().unwrap(),
+                E::BaseField::from_bool(branch_taken)
+            );
+            set_val!(
+                instance,
+                config.eq_diff_inv_marker.as_ref().unwrap()[diff_idx],
+                diff_inv_val
+            );
+        } else {
+            let is_signed = matches!(step.insn().kind, InsnKind::BLT | InsnKind::BGE);
+            UIntLimbsLT::<E>::assign(
+                config.uint_lt_config.as_ref().unwrap(),
+                instance,
+                lk_multiplicity,
+                rs1_limbs,
+                rs2_limbs,
+                is_signed,
+            )?;
+        }
         Ok(())
     }
 }
diff --git a/ceno_zkvm/src/instructions/riscv/branch/test.rs b/ceno_zkvm/src/instructions/riscv/branch/test.rs
@@ -13,16 +13,14 @@ use crate::{
     structs::ProgramParams,
 };
 
-const A: Word = 0xbead1010;
-const B: Word = 0xef552020;
-
 #[test]
 fn test_opcode_beq() {
-    impl_opcode_beq(false);
-    impl_opcode_beq(true);
+    impl_opcode_beq(false, 0xbead1010, 0xef552020);
+    impl_opcode_beq(true, 0xef552020, 0xef552020);
+    impl_opcode_beq(true, 0xffffffff, 0xffffffff);
 }
 
-fn impl_opcode_beq(equal: bool) {
+fn impl_opcode_beq(take_branch: bool, a: u32, b: u32) {
     let mut cs = ConstraintSystem::<GoldilocksExt2>::new(|| "riscv");
     let mut cb = CircuitBuilder::new(&mut cs);
     let config = cb
@@ -37,7 +35,7 @@ fn impl_opcode_beq(equal: bool) {
         .unwrap();
 
     let insn_code = encode_rv32(InsnKind::BEQ, 2, 3, 0, 8);
-    let pc_offset = if equal { 8 } else { PC_STEP_SIZE };
+    let pc_offset = if take_branch { 8 } else { PC_STEP_SIZE };
     let (raw_witin, lkm) = BeqInstruction::assign_instances(
         &config,
         &mut ShardContext::default(),
@@ -47,8 +45,8 @@ fn impl_opcode_beq(equal: bool) {
             3,
             Change::new(MOCK_PC_START, MOCK_PC_START + pc_offset),
             insn_code,
-            A,
-            if equal { A } else { B },
+            a as Word,
+            b as Word,
             0,
         )],
     )
@@ -59,11 +57,12 @@ fn impl_opcode_beq(equal: bool) {
 
 #[test]
 fn test_opcode_bne() {
-    impl_opcode_bne(false);
-    impl_opcode_bne(true);
+    impl_opcode_bne(true, 0xbead1010, 0xef552020);
+    impl_opcode_bne(false, 0xef552020, 0xef552020);
+    impl_opcode_bne(false, 0xffffffff, 0xffffffff);
 }
 
-fn impl_opcode_bne(equal: bool) {
+fn impl_opcode_bne(take_branch: bool, a: u32, b: u32) {
     let mut cs = ConstraintSystem::<GoldilocksExt2>::new(|| "riscv");
     let mut cb = CircuitBuilder::new(&mut cs);
     let config = cb
@@ -78,7 +77,7 @@ fn impl_opcode_bne(equal: bool) {
         .unwrap();
 
     let insn_code = encode_rv32(InsnKind::BNE, 2, 3, 0, 8);
-    let pc_offset = if equal { PC_STEP_SIZE } else { 8 };
+    let pc_offset = if take_branch { 8 } else { PC_STEP_SIZE };
     let (raw_witin, lkm) = BneInstruction::assign_instances(
         &config,
         &mut ShardContext::default(),
@@ -88,8 +87,8 @@ fn impl_opcode_bne(equal: bool) {
             3,
             Change::new(MOCK_PC_START, MOCK_PC_START + pc_offset),
             insn_code,
-            A,
-            if equal { A } else { B },
+            a as Word,
+            b as Word,
             0,
         )],
     )
