[chore] Optimize reduce_32 (#918)

* Optimize reduce_32 * safety: static compiler CastFV max_bits bound * opt: handle equal bits case --------- Co-authored-by: Jonathan Wang <[email protected]>
openvm-org · Dec 3, 2024 · ddabf59 · ddabf59
1 parent 23ed04a
commit ddabf59
Show file tree

Hide file tree

Showing 3 changed files with 17 additions and 6 deletions.
diff --git a/extensions/native/compiler/Cargo.toml b/extensions/native/compiler/Cargo.toml
@@ -45,7 +45,7 @@ rand = "0.8.5"
 axvm-circuit = { workspace = true, features = ["test-utils"] }
 
 [features]
-default = ["parallel"]
+default = ["parallel", "halo2-compiler"]
 halo2-compiler = ["dep:snark-verifier-sdk"]
 parallel = ["axvm-circuit/parallel"]
 bench-metrics = ["dep:metrics", "axvm-circuit/bench-metrics"]
diff --git a/extensions/native/compiler/src/constraints/halo2/compiler.rs b/extensions/native/compiler/src/constraints/halo2/compiler.rs
@@ -10,13 +10,15 @@ use axvm_circuit::metrics::cycle_tracker::CycleTracker;
 use itertools::Itertools;
 use p3_baby_bear::BabyBear;
 use p3_bn254_fr::Bn254Fr;
-use p3_field::{ExtensionField, PrimeField};
+use p3_field::{ExtensionField, PrimeField, PrimeField32};
 use snark_verifier_sdk::snark_verifier::{
     halo2_base::{
-        gates::{circuit::builder::BaseCircuitBuilder, GateInstructions, RangeChip},
+        gates::{
+            circuit::builder::BaseCircuitBuilder, GateInstructions, RangeChip, RangeInstructions,
+        },
         halo2_proofs::halo2curves::bn256::Fr,
         utils::{biguint_to_fe, ScalarField},
-        Context,
+        Context, QuantumCell,
     },
     util::arithmetic::PrimeField as _,
 };
@@ -249,8 +251,18 @@ impl<C: Config + Debug> Halo2ConstraintCompiler<C> {
                     }
                     DslIr::CastFV(a, b) => {
                         let felt = felts[&b.0];
+                        #[allow(clippy::comparison_chain)]
                         let reduced_felt = if felt.max_bits > BABYBEAR_MAX_BITS {
                             f_chip.reduce(ctx, felt)
+                        } else if felt.max_bits == BABYBEAR_MAX_BITS {
+                            // Ensure cast is canonical
+                            f_chip.range.check_less_than(
+                                ctx,
+                                felt.value,
+                                QuantumCell::Constant(Fr::from(BabyBear::ORDER_U32 as u64)),
+                                BABYBEAR_MAX_BITS,
+                            );
+                            felt
                         } else {
                             felt
                         };

diff --git a/extensions/native/recursion/src/utils.rs b/extensions/native/recursion/src/utils.rs
@@ -40,8 +40,7 @@ pub fn reduce_32<C: Config>(builder: &mut Builder<C>, vals: &[Felt<C::F>]) -> Va
     let mut power = C::N::ONE;
     let result: Var<C::N> = builder.eval(C::N::ZERO);
     for val in vals.iter() {
-        let bits = builder.num2bits_f_circuit(*val);
-        let val = builder.bits2num_v_circuit(&bits);
+        let val = builder.cast_felt_to_var(*val);
         builder.assign(&result, result + val * power);
         power *= C::N::from_canonical_usize(1usize << 32);
     }