[perf] Add CircuitVarTo64BitsF for split_32 (#80)

nyunyunyunyu · jonathanpwang · commit bff6d573ce7e · 2025-03-09T03:00:21.000-04:00
Reth-benchmark: https://github.com/axiom-crypto/openvm-reth-benchmark/actions/runs/13730024145 Static verifier performance is back to before #5 Used cells: 57M(before #5) v.s. 62M(after this PR) It seems that the number of columns are same so the actual proof times are close. This actually fixes a soundness issue. `CircuitNum2BitsV` cannot decompose a `Var` into unique limbs.
diff --git a/extensions/native/compiler/src/constraints/halo2/compiler.rs b/extensions/native/compiler/src/constraints/halo2/compiler.rs
@@ -14,13 +14,14 @@ use openvm_stark_sdk::{p3_baby_bear::BabyBear, p3_bn254_fr::Bn254Fr};
 use snark_verifier_sdk::snark_verifier::{
     halo2_base::{
         gates::{
-            circuit::builder::BaseCircuitBuilder, GateInstructions, RangeChip, RangeInstructions,
+            circuit::builder::BaseCircuitBuilder, GateChip, GateInstructions, RangeChip,
+            RangeInstructions,
         },
         halo2_proofs::halo2curves::bn256::Fr,
-        utils::{biguint_to_fe, ScalarField},
-        Context,
+        utils::{biguint_to_fe, decompose_fe_to_u64_limbs, ScalarField},
+        AssignedValue, Context, QuantumCell,
     },
-    util::arithmetic::PrimeField as _,
+    util::arithmetic::{Field as _, PrimeField as _},
 };
 
 use super::stats::Halo2Stats;
@@ -314,15 +315,6 @@ impl<C: Config + Debug> Halo2ConstraintCompiler<C> {
                         let reduced_felt = f_chip.reduce(ctx, felt);
                         vars.insert(a.0, reduced_felt.value);
                     }
-                    DslIr::CircuitNum2BitsV(value, bits, output) => {
-                        let shortened_bits = bits.min(Fr::NUM_BITS as usize);
-                        let mut x = gate.num_to_bits(ctx, vars[&value.0], shortened_bits);
-                        let zero = ctx.load_zero();
-                        x.resize(bits, zero);
-                        for (o, x) in output.into_iter().zip_eq(x) {
-                            vars.insert(o.0, x);
-                        }
-                    }
                     DslIr::CircuitNum2BitsF(value, output) => {
                         let val = f_chip.reduce(ctx, felts[&value.0]);
                         let x = gate.num_to_bits(ctx, val.value, 32); // C::F::bits());
@@ -331,6 +323,13 @@ impl<C: Config + Debug> Halo2ConstraintCompiler<C> {
                             vars.insert(o.0, x);
                         }
                     }
+                    DslIr::CircuitVarTo64BitsF(value, output) => {
+                        let x = vars[&value.0];
+                        let limbs = var_to_u64_limbs(ctx, &range, gate, x);
+                        for (o, l) in output.into_iter().zip(limbs) {
+                            felts.insert(o.0, l);
+                        }
+                    }
                     DslIr::CircuitPoseidon2Permute(state_vars) => {
                         let mut state =
                             Poseidon2State::<Fr, POSEIDON2_T>::new(state_vars.map(|x| vars[&x.0]));
@@ -526,10 +525,79 @@ fn is_babybear_ir<C: Config>(ir: &DslIr<C>) -> bool {
     )
 }
 
-#[allow(dead_code)]
-fn is_num2bits_ir<C: Config>(ir: &DslIr<C>) -> bool {
-    matches!(
-        ir,
-        DslIr::CircuitNum2BitsV(_, _, _) | DslIr::CircuitNum2BitsF(_, _)
-    )
+fn fr_to_u64_limbs(fr: &Fr) -> [u64; 4] {
+    // We need 64-bit limbs but `decompose_fe_to_u64_limbs` only support `bit_len < 64`.
+    let limbs = decompose_fe_to_u64_limbs(fr, 8, 32);
+    std::array::from_fn(|i| limbs[2 * i] + limbs[2 * i + 1] * (1 << 32))
 }
+
+fn var_to_u64_limbs(
+    ctx: &mut Context<Fr>,
+    range: &RangeChip<Fr>,
+    gate: &GateChip<Fr>,
+    x: AssignedValue<Fr>,
+) -> [AssignedBabyBear; 4] {
+    let limbs = fr_to_u64_limbs(x.value()).map(|limb| ctx.load_witness(Fr::from(limb)));
+    let factors = [
+        Fr::from([1, 0, 0, 0]),
+        Fr::from([0, 1, 0, 0]),
+        Fr::from([0, 0, 1, 0]),
+        Fr::from([0, 0, 0, 1]),
+    ];
+    let sum = gate.inner_product(ctx, limbs, factors.map(QuantumCell::Constant));
+    ctx.constrain_equal(&sum, &x);
+    let fr_bound_limbs = fr_to_u64_limbs(&(Fr::ZERO - Fr::ONE));
+    let ret = std::array::from_fn(|i| {
+        let limb = limbs[i];
+        let bits = if i < 3 {
+            range.range_check(ctx, limb, 64);
+            64
+        } else {
+            range.check_less_than_safe(ctx, limbs[3], fr_bound_limbs[3] + 1);
+            (Fr::NUM_BITS - 3 * 64) as usize
+        };
+        AssignedBabyBear {
+            value: limb,
+            max_bits: bits,
+        }
+    });
+    // Constraint decomposition doesn't overflow.
+    // Whether limbs[i] == fr_bound_limbs[i] so far
+    let mut on_bound = gate.is_equal(
+        ctx,
+        limbs[3],
+        QuantumCell::Constant(Fr::from(fr_bound_limbs[3])),
+    );
+    for i in (0..3).rev() {
+        // limbs[i] > fr_bound_limbs[i]
+        let li_gt_bd = range.is_less_than(
+            ctx,
+            QuantumCell::Constant(Fr::from(fr_bound_limbs[i])),
+            limbs[i],
+            64,
+        );
+        let li_out_bd = gate.add(ctx, on_bound, li_gt_bd);
+        // on_bound  li_gt_bd  result
+        //    1         1       fail
+        //    1         0       pass
+        //    0         1       pass
+        //    0         0       pass
+        gate.assert_bit(ctx, li_out_bd);
+        // Update on_bound except the last limb
+        if i > 0 {
+            debug_assert_ne!(fr_bound_limbs[i], 0, "This should never happen for Bn254Fr");
+            // on_bound && limbs[i] - fr_bound_limbs[i] == 0
+            let diff = gate.sub_mul(
+                ctx,
+                QuantumCell::Constant(Fr::from(fr_bound_limbs[i])),
+                on_bound,
+                limbs[i],
+            );
+            on_bound = gate.is_zero(ctx, diff);
+        }
+    }
+    ret
+}
+
+#[test]
+fn test_var_to_u64_limbs() {}
diff --git a/extensions/native/compiler/src/constraints/mod.rs b/extensions/native/compiler/src/constraints/mod.rs
@@ -236,14 +236,6 @@ impl<C: Config + Debug> ConstraintCompiler<C> {
                     opcode: ConstraintOpcode::NegE,
                     args: vec![vec![a.id()], vec![b.id()]],
                 }),
-                DslIr::CircuitNum2BitsV(value, bits, output) => constraints.push(Constraint {
-                    opcode: ConstraintOpcode::Num2BitsV,
-                    args: vec![
-                        output.iter().map(|x| x.id()).collect(),
-                        vec![value.id()],
-                        vec![bits.to_string()],
-                    ],
-                }),
                 DslIr::CircuitNum2BitsF(value, output) => constraints.push(Constraint {
                     opcode: ConstraintOpcode::Num2BitsF,
                     args: vec![output.iter().map(|x| x.id()).collect(), vec![value.id()]],
diff --git a/extensions/native/compiler/src/ir/bits.rs b/extensions/native/compiler/src/ir/bits.rs
@@ -1,23 +1,11 @@
-use std::any::TypeId;
+use std::{any::TypeId, array};
 
 use openvm_stark_backend::p3_field::FieldAlgebra;
 use openvm_stark_sdk::p3_baby_bear::BabyBear;
 
 use super::{Array, Builder, Config, DslIr, Felt, MemIndex, Var};
 
 impl<C: Config> Builder<C> {
-    /// Converts a variable to bits inside a circuit.
-    pub fn num2bits_v_circuit(&mut self, num: Var<C::N>, bits: usize) -> Vec<Var<C::N>> {
-        let mut output = Vec::new();
-        for _ in 0..bits {
-            output.push(self.uninit());
-        }
-
-        self.push(DslIr::CircuitNum2BitsV(num, bits, output.clone()));
-
-        output
-    }
-
     /// Converts a felt to bits. Will result in a failed assertion if `num` has more than `num_bits` bits.
     /// Only works for C::F = BabyBear
     pub fn num2bits_f(&mut self, num: Felt<C::F>, num_bits: u32) -> Array<C, Var<C::N>> {
@@ -94,4 +82,11 @@ impl<C: Config> Builder<C> {
         }
         result
     }
+
+    /// Decompose a Var into 64-bit Felt limbs.
+    pub fn var_to_64bits_f_circuit(&mut self, value: Var<C::N>) -> [Felt<C::F>; 4] {
+        let ret = array::from_fn(|_| self.uninit());
+        self.push(DslIr::CircuitVarTo64BitsF(value, ret));
+        ret
+    }
 }
diff --git a/extensions/native/compiler/src/ir/instructions.rs b/extensions/native/compiler/src/ir/instructions.rs
@@ -179,10 +179,10 @@ pub enum DslIr<C: Config> {
     StoreHeapPtr(Ptr<C::N>),
 
     // Bits.
-    /// Decompose a variable into size bits (bits = num2bits(var, size)). Should only be used when target is a circuit.
-    CircuitNum2BitsV(Var<C::N>, usize, Vec<Var<C::N>>),
     /// Decompose a field element into bits (bits = num2bits(felt)). Should only be used when target is a circuit.
     CircuitNum2BitsF(Felt<C::F>, Vec<Var<C::N>>),
+    /// Decompose a Var into 16-bit limbs.
+    CircuitVarTo64BitsF(Var<C::N>, [Felt<C::F>; 4]),
 
     // Hashing.
     /// Permutes an array of baby bear elements using Poseidon2 (output = p2_permute(array)).
diff --git a/extensions/native/recursion/src/halo2/tests/mod.rs b/extensions/native/recursion/src/halo2/tests/mod.rs
@@ -93,28 +93,6 @@ fn test_publish() {
     assert_eq!(pis, vec![vec![convert_fr(&value_fr)]]);
 }
 
-#[test]
-fn test_num2bits_v_circuit() {
-    let mut builder = Builder::<OuterConfig>::default();
-    builder.flags.static_only = true;
-    let mut value_u32 = 1345237507;
-    let value = builder.eval(Bn254Fr::from_canonical_u32(value_u32));
-    let result = builder.num2bits_v_circuit(value, 32);
-    for r in result {
-        builder.assert_var_eq(r, Bn254Fr::from_canonical_u32(value_u32 & 1));
-        value_u32 >>= 1;
-    }
-
-    Halo2Prover::mock::<OuterConfig>(
-        10,
-        DslOperations {
-            operations: builder.operations,
-            num_public_values: 0,
-        },
-        Witness::default(),
-    );
-}
-
 #[test]
 fn test_reduce_32() {
     let value_1 = BabyBear::from_canonical_u32(1345237507);
@@ -140,27 +118,32 @@ fn test_reduce_32() {
 
 #[test]
 fn test_split_32() {
-    let value = Bn254Fr::from_canonical_u32(1345237507);
-    let gt: Vec<BabyBear> = split_32_gt(value, 3);
-    dbg!(&gt);
-
-    let mut builder = Builder::<OuterConfig>::default();
-    builder.flags.static_only = true;
-    let value = builder.eval(value);
-    let result = split_32(&mut builder, value, 3);
-
-    builder.assert_felt_eq(result[0], gt[0]);
-    builder.assert_felt_eq(result[1], gt[1]);
-    builder.assert_felt_eq(result[2], gt[2]);
-
-    Halo2Prover::mock::<OuterConfig>(
-        10,
-        DslOperations {
-            operations: builder.operations,
-            num_public_values: 0,
-        },
-        Witness::default(),
-    );
+    let f = |value| {
+        let gt: Vec<BabyBear> = split_32_gt(value, 3);
+        dbg!(&gt);
+
+        let mut builder = Builder::<OuterConfig>::default();
+        builder.flags.static_only = true;
+        let value = builder.eval(value);
+        let result = split_32(&mut builder, value, 3);
+
+        builder.assert_felt_eq(result[0], gt[0]);
+        builder.assert_felt_eq(result[1], gt[1]);
+        builder.assert_felt_eq(result[2], gt[2]);
+
+        Halo2Prover::mock::<OuterConfig>(
+            10,
+            DslOperations {
+                operations: builder.operations,
+                num_public_values: 0,
+            },
+            Witness::default(),
+        );
+    };
+    let modulus = Bn254Fr::ZERO - Bn254Fr::ONE;
+    f(Bn254Fr::from_canonical_u32(1345237507));
+    f(Bn254Fr::ZERO);
+    f(modulus);
 }
 
 #[test]
diff --git a/extensions/native/recursion/src/utils.rs b/extensions/native/recursion/src/utils.rs
@@ -52,19 +52,9 @@ pub fn reduce_32<C: Config>(builder: &mut Builder<C>, vals: &[Felt<C::F>]) -> Va
 
 /// Reference: <https://github.com/Plonky3/Plonky3/blob/622375885320ac6bf3c338001760ed8f2230e3cb/field/src/helpers.rs#L149>
 pub fn split_32<C: Config>(builder: &mut Builder<C>, val: Var<C::N>, n: usize) -> Vec<Felt<C::F>> {
-    let bits = builder.num2bits_v_circuit(val, 256);
-    let mut results = Vec::new();
-    for i in 0..n {
-        let result: Felt<C::F> = builder.eval(C::F::ZERO);
-        for j in 0..64 {
-            let bit = bits[i * 64 + j];
-            let t = builder.eval(result + C::F::from_wrapped_u64(1 << j));
-            let z = builder.select_f(bit, t, result);
-            builder.assign(&result, z);
-        }
-        results.push(result);
-    }
-    results
+    let felts = builder.var_to_64bits_f_circuit(val);
+    assert!(n <= felts.len());
+    felts[0..n].to_vec()
 }
 
 /// Eval two expressions, return in the reversed order if cond == 1. Otherwise, return in the original order.
