Optimize Static Verifier alpha pow computation

openvm-org · Jan 22, 2025 · 2c8b3b6 · 2c8b3b6
1 parent acdb0e2
commit 2c8b3b6
Show file tree

Hide file tree

Showing 5 changed files with 50 additions and 4 deletions.
diff --git a/extensions/native/compiler/src/constraints/halo2/baby_bear.rs b/extensions/native/compiler/src/constraints/halo2/baby_bear.rs
@@ -80,6 +80,15 @@ impl BabyBearChip {
         r
     }
 
+    /// Reduce max_bits if possible. This function doesn't guarantee that the actual value is within BabyBear.
+    pub fn reduce_max_bits(&self, ctx: &mut Context<Fr>, a: AssignedBabyBear) -> AssignedBabyBear {
+        if a.max_bits > BABYBEAR_MAX_BITS {
+            self.reduce(ctx, a)
+        } else {
+            a
+        }
+    }
+
     pub fn add(
         &self,
         ctx: &mut Context<Fr>,
@@ -527,4 +536,18 @@ impl BabyBearExt4Chip {
         );
         c
     }
+
+    pub fn reduce_max_bits(
+        &self,
+        ctx: &mut Context<Fr>,
+        a: AssignedBabyBearExt4,
+    ) -> AssignedBabyBearExt4 {
+        AssignedBabyBearExt4(
+            a.0.into_iter()
+                .map(|x| self.base.reduce_max_bits(ctx, x))
+                .collect::<Vec<_>>()
+                .try_into()
+                .unwrap(),
+        )
+    }
 }
diff --git a/extensions/native/compiler/src/constraints/halo2/compiler.rs b/extensions/native/compiler/src/constraints/halo2/compiler.rs
@@ -408,6 +408,14 @@ impl<C: Config + Debug> Halo2ConstraintCompiler<C> {
                         );
                         exts.insert(b.0, x);
                     }
+                    DslIr::CircuitFeltReduce(a) => {
+                        let x = f_chip.reduce_max_bits(ctx, felts[&a.0]);
+                        felts.insert(a.0, x);
+                    }
+                    DslIr::CircuitExtReduce(a) => {
+                        let x = ext_chip.reduce_max_bits(ctx, exts[&a.0]);
+                        exts.insert(a.0, x);
+                    }
                     DslIr::CycleTrackerStart(_name) => {
                         #[cfg(feature = "bench-metrics")]
                         cell_tracker.start(_name);
@@ -482,6 +490,8 @@ fn is_babybear_ir<C: Config>(ir: &DslIr<C>) -> bool {
             | DslIr::AssertEqFI(_, _)
             | DslIr::WitnessFelt(_, _)
             | DslIr::CircuitFelts2Ext(_, _)
+            | DslIr::CircuitFeltReduce(_)
+            | DslIr::CircuitExtReduce(_)
             | DslIr::ImmE(_, _)
             | DslIr::AddE(_, _, _)
             | DslIr::AddEF(_, _, _)

diff --git a/extensions/native/compiler/src/ir/instructions.rs b/extensions/native/compiler/src/ir/instructions.rs
@@ -248,6 +248,10 @@ pub enum DslIr<C: Config> {
     CircuitExt2Felt([Felt<C::F>; 4], Ext<C::F, C::EF>),
     /// Converts a slice of felts to an ext. Should only be used when target is a circuit.
     CircuitFelts2Ext([Felt<C::F>; 4], Ext<C::F, C::EF>),
+    /// Halo2 only. Reduce a Felt so later computation becomes cheaper.
+    CircuitFeltReduce(Felt<C::F>),
+    /// Halo2 only. Reduce an Ext so later computation becomes cheaper.
+    CircuitExtReduce(Ext<C::F, C::EF>),
     /// FriReducedOpening(alpha, curr_alpha_pow, at_x_array, at_z_array, result)
     FriReducedOpening(
         Ext<C::F, C::EF>,

diff --git a/extensions/native/compiler/src/ir/utils.rs b/extensions/native/compiler/src/ir/utils.rs
@@ -175,4 +175,10 @@ impl<C: Config> Builder<C> {
             .push(DslIr::CircuitExt2Felt([a, b, c, d], value));
         [a, b, c, d]
     }
+    pub fn felt_reduce_circuit(&mut self, value: Felt<C::F>) {
+        self.operations.push(DslIr::CircuitFeltReduce(value));
+    }
+    pub fn ext_reduce_circuit(&mut self, value: Ext<C::F, C::EF>) {
+        self.operations.push(DslIr::CircuitExtReduce(value));
+    }
 }
diff --git a/extensions/native/recursion/src/fri/two_adic_pcs.rs b/extensions/native/recursion/src/fri/two_adic_pcs.rs
@@ -232,7 +232,7 @@ pub fn verify_two_adic_pcs<C: Config>(
                                     let p_at_z = builder.get(&ps_at_z, t);
                                     builder.assign(&n, n * alpha + (p_at_z - p_at_x));
                                 }
-                                builder.assign(&cur_ro, cur_ro + cur_alpha_pow * n / (z - x));
+                                builder.assign(&cur_ro, cur_ro + n / (z - x) * cur_alpha_pow);
                                 builder.assign(&cur_alpha_pow, cur_alpha_pow * mat_alpha_pow);
                             } else {
                                 // TODO: this is just for testing the correctness. Will remove later.
@@ -396,13 +396,16 @@ fn compute_round_alpha_pows<C: Config>(
             let mat_alpha_pow: Ext<_, _> = if builder.flags.static_only {
                 let width = width.value();
                 assert!(width < 1 << MAX_LOG_WIDTH);
-                let mut ret = C::EF::ONE.cons();
+                let mut expr = C::EF::ONE.cons();
                 for i in 0..MAX_LOG_WIDTH {
                     if width & (1 << i) != 0 {
-                        ret *= builder.get(&pow_of_alpha, i);
+                        expr *= builder.get(&pow_of_alpha, i);
                     }
                 }
-                builder.eval(ret)
+                let ret: Ext<_, _> = builder.eval(expr);
+                // Minimize max_bits so following computation becomes cheaper.
+                builder.ext_reduce_circuit(ret);
+                ret
             } else {
                 let width = width.get_var();
                 // This is dynamic only so safe to cast.