feat: plonk circuit optimizations

primitives/src/lib.rs

@@ -6,6 +6,8 @@ use p3_baby_bear::{BabyBear, DiffusionMatrixBabyBear};
 use p3_field::AbstractField;
 use p3_poseidon2::{Poseidon2, Poseidon2ExternalMatrixGeneral};
 
+pub mod types;
+
 lazy_static! {
     // These constants are created by a RNG.
 

primitives/src/types.rs

@@ -0,0 +1,8 @@
+#[derive(Debug, Clone, Copy)]
+pub enum RecursionProgramType {
+    Core,
+    Deferred,
+    Compress,
+    Shrink,
+    Wrap,
+}

prover/src/lib.rs

@@ -39,6 +39,7 @@ use sp1_core::{
     utils::{BabyBearPoseidon2, SP1CoreProverError},
 };
 use sp1_primitives::hash_deferred_proof;
+use sp1_primitives::types::RecursionProgramType;
 use sp1_recursion_circuit::witness::Witnessable;
 use sp1_recursion_compiler::config::InnerConfig;
 use sp1_recursion_compiler::ir::Witness;
@@ -164,14 +165,20 @@ impl SP1Prover {
         let (compress_pk, compress_vk) = compress_machine.setup(&compress_program);
 
         // Get the compress program, machine, and keys.
-        let shrink_program =
-            SP1RootVerifier::<InnerConfig, _, _>::build(&compress_machine, &compress_vk, true);
+        let shrink_program = SP1RootVerifier::<InnerConfig, _, _>::build(
+            &compress_machine,
+            &compress_vk,
+            RecursionProgramType::Shrink,
+        );
         let shrink_machine = CompressAir::wrap_machine_dyn(InnerSC::compressed());
         let (shrink_pk, shrink_vk) = shrink_machine.setup(&shrink_program);
 
         // Get the wrap program, machine, and keys.
-        let wrap_program =
-            SP1RootVerifier::<InnerConfig, _, _>::build(&shrink_machine, &shrink_vk, false);
+        let wrap_program = SP1RootVerifier::<InnerConfig, _, _>::build(
+            &shrink_machine,
+            &shrink_vk,
+            RecursionProgramType::Wrap,
+        );
         let wrap_machine = WrapAir::wrap_machine(OuterSC::default());
         let (wrap_pk, wrap_vk) = wrap_machine.setup(&wrap_program);
 

recursion/circuit/src/fri.rs

@@ -122,13 +122,14 @@ pub fn verify_two_adic_pcs<C: Config>(
                     let x: Felt<_> = builder.eval(g * two_adic_generator_exp);
 
                     for (z, ps_at_z) in izip!(mat_points, mat_values) {
+                        let mut acc: Ext<C::F, C::EF> =
+                            builder.eval(SymbolicExt::from_f(C::EF::zero()));
                         for (p_at_x, &p_at_z) in izip!(mat_opening.clone(), ps_at_z) {
-                            let quotient: SymbolicExt<C::F, C::EF> =
-                                (p_at_z - p_at_x[0]) / (*z - x);
-                            ro[log_height] =
-                                builder.eval(ro[log_height] + alpha_pow[log_height] * quotient);
+                            acc =
+                                builder.eval(acc + (alpha_pow[log_height] * (p_at_z - p_at_x[0])));
                             alpha_pow[log_height] = builder.eval(alpha_pow[log_height] * alpha);
                         }
+                        ro[log_height] = builder.eval(ro[log_height] + acc / (*z - x));
                     }
                 }
             }

recursion/compiler/Cargo.toml

@@ -5,20 +5,23 @@ edition = "2021"
 
 [dependencies]
 p3-air = { workspace = true }
-p3-field = { workspace = true }
+p3-bn254-fr = { workspace = true }
+p3-baby-bear = { workspace = true }
 p3-commit = { workspace = true }
+p3-field = { workspace = true }
 p3-fri = { workspace = true }
 p3-matrix = { workspace = true }
+p3-poseidon2 = { workspace = true }
 p3-symmetric = { workspace = true }
 p3-util = { workspace = true }
-sp1-recursion-core = { path = "../core" }
+
 sp1-core = { path = "../../core" }
+sp1-primitives = { path = "../../primitives" }
+sp1-recursion-core = { path = "../core" }
+sp1-recursion-derive = { path = "../derive" }
+
 itertools = "0.13.0"
 serde = { version = "1.0.201", features = ["derive"] }
-sp1-recursion-derive = { path = "../derive" }
-p3-bn254-fr = { workspace = true }
-p3-baby-bear = { workspace = true }
-p3-poseidon2 = { workspace = true }
 backtrace = "0.3.71"
 tracing = "0.1.40"
 

recursion/compiler/src/ir/builder.rs

@@ -2,6 +2,7 @@ use std::{iter::Zip, vec::IntoIter};
 
 use backtrace::Backtrace;
 use p3_field::AbstractField;
+use sp1_primitives::types::RecursionProgramType;
 
 use super::{
     Array, Config, DslIr, Ext, Felt, FromConstant, SymbolicExt, SymbolicFelt, SymbolicUsize,
@@ -103,10 +104,17 @@ pub struct Builder<C: Config> {
     pub(crate) p2_hash_num: Var<C::N>,
     pub(crate) debug: bool,
     pub(crate) is_sub_builder: bool,
+    pub program_type: RecursionProgramType,
 }
 
 impl<C: Config> Default for Builder<C> {
     fn default() -> Self {
+        Self::new(RecursionProgramType::Core)
+    }
+}
+
+impl<C: Config> Builder<C> {
+    pub fn new(program_type: RecursionProgramType) -> Self {
         // We need to create a temporary placeholder for the p2_hash_num variable.
         let placeholder_p2_hash_num = Var::new(0);
 
@@ -122,14 +130,13 @@ impl<C: Config> Default for Builder<C> {
             p2_hash_num: placeholder_p2_hash_num,
             debug: false,
             is_sub_builder: false,
+            program_type,
         };
 
         new_builder.p2_hash_num = new_builder.uninit();
         new_builder
     }
-}
 
-impl<C: Config> Builder<C> {
     /// Creates a new builder with a given number of counts for each type.
     pub fn new_sub_builder(
         var_count: u32,
@@ -138,6 +145,7 @@ impl<C: Config> Builder<C> {
         nb_public_values: Option<Var<C::N>>,
         p2_hash_num: Var<C::N>,
         debug: bool,
+        program_type: RecursionProgramType,
     ) -> Self {
         Self {
             felt_count,
@@ -153,6 +161,7 @@ impl<C: Config> Builder<C> {
             p2_hash_num,
             debug,
             is_sub_builder: true,
+            program_type,
         }
     }
 
@@ -546,6 +555,7 @@ impl<'a, C: Config> IfBuilder<'a, C> {
             self.builder.nb_public_values,
             self.builder.p2_hash_num,
             self.builder.debug,
+            self.builder.program_type,
         );
         f(&mut f_builder);
         self.builder.p2_hash_num = f_builder.p2_hash_num;
@@ -597,6 +607,7 @@ impl<'a, C: Config> IfBuilder<'a, C> {
             self.builder.nb_public_values,
             self.builder.p2_hash_num,
             self.builder.debug,
+            self.builder.program_type,
         );
 
         // Execute the `then` and `else_then` blocks and collect the instructions.
@@ -612,6 +623,7 @@ impl<'a, C: Config> IfBuilder<'a, C> {
             self.builder.nb_public_values,
             self.builder.p2_hash_num,
             self.builder.debug,
+            self.builder.program_type,
         );
         else_f(&mut else_builder);
         self.builder.p2_hash_num = else_builder.p2_hash_num;
@@ -749,6 +761,7 @@ impl<'a, C: Config> RangeBuilder<'a, C> {
             self.builder.nb_public_values,
             self.builder.p2_hash_num,
             self.builder.debug,
+            self.builder.program_type,
         );
 
         f(loop_variable, &mut loop_body_builder);

recursion/core/src/poseidon2_wide/columns/mod.rs

@@ -206,7 +206,6 @@ pub struct Poseidon2Degree9<T: Copy> {
     pub memory: Memory<T>,
     pub opcode_specific_cols: OpcodeWorkspace<T>,
     pub permutation_cols: PermutationNoSbox<T>,
-    pub state_cursor: [T; WIDTH / 2], // Only used for absorb
 }
 
 impl<'a, T: Copy + 'a> Poseidon2<'a, T> for Poseidon2Degree9<T> {

recursion/core/src/stark/mod.rs

@@ -118,22 +118,16 @@ impl<F: PrimeField32 + BinomiallyExtendable<D>, const DEGREE: usize> RecursionAi
     pub fn get_wrap_all() -> Vec<Self> {
         once(RecursionAir::Program(ProgramChip))
             .chain(once(RecursionAir::Cpu(CpuChip {
-                fixed_log2_rows: Some(19),
+                fixed_log2_rows: Some(20),
                 _phantom: PhantomData,
             })))
             .chain(once(RecursionAir::MemoryGlobal(MemoryGlobalChip {
-                fixed_log2_rows: Some(20),
+                fixed_log2_rows: Some(19),
             })))
             .chain(once(RecursionAir::Multi(MultiChip {
                 fixed_log2_rows: Some(17),
             })))
             .chain(once(RecursionAir::RangeCheck(RangeCheckChip::default())))
-            .chain(once(RecursionAir::ExpReverseBitsLen(
-                ExpReverseBitsLenChip::<DEGREE> {
-                    fixed_log2_rows: None,
-                    pad: true,
-                },
-            )))
             .collect()
     }
 }

recursion/core/src/stark/utils.rs

@@ -17,7 +17,7 @@ pub enum TestConfig {
     All,
     WideDeg3,
     SkinnyDeg7,
-    SkinnyDeg7Wrap,
+    WideDeg17Wrap,
 }
 
 type Val = <BabyBearPoseidon2 as StarkGenericConfig>::Val;
@@ -63,7 +63,7 @@ pub fn run_test_recursion(
         }
     }
 
-    if test_config == TestConfig::All || test_config == TestConfig::SkinnyDeg7Wrap {
+    if test_config == TestConfig::All || test_config == TestConfig::WideDeg17Wrap {
         let machine = RecursionAirWideDeg9::wrap_machine(BabyBearPoseidon2::compressed());
         let (pk, vk) = machine.setup(&program);
         let record = runtime.record.clone();

recursion/program/Cargo.toml

@@ -20,6 +20,7 @@ p3-poseidon2 = { workspace = true }
 sp1-recursion-core = { path = "../core" }
 sp1-recursion-compiler = { path = "../compiler" }
 sp1-core = { path = "../../core" }
+sp1-primitives = { path = "../../primitives" }
 itertools = "0.13.0"
 serde = { version = "1.0.201", features = ["derive"] }
 rand = "0.8.5"

recursion/program/src/constraints.rs

@@ -383,7 +383,10 @@ mod tests {
 
         let program = builder.compile_program();
 
-        run_test_recursion(program, None, TestConfig::All);
+        // We don't test with the config TestConfig::WideDeg17Wrap, since it doesn't have the
+        // `ExpReverseBitsLen` chip.
+        run_test_recursion(program.clone(), None, TestConfig::WideDeg3);
+        run_test_recursion(program, None, TestConfig::SkinnyDeg7);
     }
 
     #[test]

recursion/program/src/fri/mod.rs

@@ -4,6 +4,7 @@ pub mod two_adic_pcs;
 pub mod types;
 
 pub use domain::*;
+use sp1_primitives::types::RecursionProgramType;
 use sp1_recursion_compiler::ir::ExtensionOperand;
 use sp1_recursion_compiler::ir::Ptr;
 use sp1_recursion_core::runtime::DIGEST_SIZE;
@@ -141,7 +142,11 @@ where
     let folded_eval: Ext<C::F, C::EF> = builder.eval(C::F::zero());
     let two_adic_generator_f = config.get_two_adic_generator(builder, log_max_height);
 
-    let x = builder.exp_reverse_bits_len_fast(two_adic_generator_f, index_bits, log_max_height);
+    let x = if matches!(builder.program_type, RecursionProgramType::Wrap) {
+        builder.exp_reverse_bits_len(two_adic_generator_f, index_bits, log_max_height)
+    } else {
+        builder.exp_reverse_bits_len_fast(two_adic_generator_f, index_bits, log_max_height)
+    };
 
     let log_max_height = log_max_height.materialize(builder);
     builder

recursion/program/src/fri/two_adic_pcs.rs

@@ -2,6 +2,7 @@ use p3_commit::TwoAdicMultiplicativeCoset;
 use p3_field::AbstractField;
 use p3_field::TwoAdicField;
 use p3_symmetric::Hash;
+use sp1_primitives::types::RecursionProgramType;
 use sp1_recursion_compiler::prelude::*;
 use sp1_recursion_core::runtime::DIGEST_SIZE;
 
@@ -120,11 +121,22 @@ pub fn verify_two_adic_pcs<C: Config>(
 
                         let two_adic_generator = config.get_two_adic_generator(builder, log_height);
                         builder.cycle_tracker("exp_reverse_bits_len");
-                        let two_adic_generator_exp = builder.exp_reverse_bits_len_fast(
-                            two_adic_generator,
-                            &index_bits_shifted,
-                            log_height,
-                        );
+
+                        let two_adic_generator_exp: Felt<C::F> =
+                            if matches!(builder.program_type, RecursionProgramType::Wrap) {
+                                builder.exp_reverse_bits_len(
+                                    two_adic_generator,
+                                    &index_bits_shifted,
+                                    log_height,
+                                )
+                            } else {
+                                builder.exp_reverse_bits_len_fast(
+                                    two_adic_generator,
+                                    &index_bits_shifted,
+                                    log_height,
+                                )
+                            };
+
                         builder.cycle_tracker("exp_reverse_bits_len");
                         let x: Felt<C::F> = builder.eval(two_adic_generator_exp * g);
 

recursion/program/src/machine/compress.rs

@@ -14,6 +14,7 @@ use sp1_core::air::{Word, POSEIDON_NUM_WORDS, PV_DIGEST_NUM_WORDS};
 use sp1_core::stark::StarkMachine;
 use sp1_core::stark::{Com, ShardProof, StarkGenericConfig, StarkVerifyingKey};
 use sp1_core::utils::BabyBearPoseidon2;
+use sp1_primitives::types::RecursionProgramType;
 use sp1_recursion_compiler::config::InnerConfig;
 use sp1_recursion_compiler::ir::{Array, Builder, Config, Felt, Var};
 use sp1_recursion_compiler::prelude::DslVariable;
@@ -81,7 +82,7 @@ where
         recursive_vk: &StarkVerifyingKey<BabyBearPoseidon2>,
         deferred_vk: &StarkVerifyingKey<BabyBearPoseidon2>,
     ) -> RecursionProgram<BabyBear> {
-        let mut builder = Builder::<InnerConfig>::default();
+        let mut builder = Builder::<InnerConfig>::new(RecursionProgramType::Compress);
 
         let input: SP1ReduceMemoryLayoutVariable<_> = builder.uninit();
         SP1ReduceMemoryLayout::<BabyBearPoseidon2, A>::witness(&input, &mut builder);

recursion/program/src/machine/core.rs

@@ -11,6 +11,7 @@ use sp1_core::air::{Word, POSEIDON_NUM_WORDS, PV_DIGEST_NUM_WORDS};
 use sp1_core::stark::StarkMachine;
 use sp1_core::stark::{Com, RiscvAir, ShardProof, StarkGenericConfig, StarkVerifyingKey};
 use sp1_core::utils::BabyBearPoseidon2;
+use sp1_primitives::types::RecursionProgramType;
 use sp1_recursion_compiler::config::InnerConfig;
 use sp1_recursion_compiler::ir::{Array, Builder, Config, Ext, ExtConst, Felt, Var};
 use sp1_recursion_compiler::prelude::DslVariable;
@@ -64,7 +65,7 @@ impl SP1RecursiveVerifier<InnerConfig, BabyBearPoseidon2> {
     pub fn build(
         machine: &StarkMachine<BabyBearPoseidon2, RiscvAir<BabyBear>>,
     ) -> RecursionProgram<BabyBear> {
-        let mut builder = Builder::<InnerConfig>::default();
+        let mut builder = Builder::<InnerConfig>::new(RecursionProgramType::Core);
 
         let input: SP1RecursionMemoryLayoutVariable<_> = builder.uninit();
         SP1RecursionMemoryLayout::<BabyBearPoseidon2, RiscvAir<_>>::witness(&input, &mut builder);

recursion/program/src/machine/deferred.rs

@@ -11,6 +11,7 @@ use sp1_core::air::{Word, POSEIDON_NUM_WORDS, PV_DIGEST_NUM_WORDS};
 use sp1_core::stark::StarkMachine;
 use sp1_core::stark::{Com, RiscvAir, ShardProof, StarkGenericConfig, StarkVerifyingKey};
 use sp1_core::utils::BabyBearPoseidon2;
+use sp1_primitives::types::RecursionProgramType;
 use sp1_recursion_compiler::config::InnerConfig;
 use sp1_recursion_compiler::ir::{Array, Builder, Config, Felt, Var};
 use sp1_recursion_compiler::prelude::DslVariable;
@@ -83,7 +84,7 @@ where
 {
     /// Create a new instance of the program for the [BabyBearPoseidon2] config.
     pub fn build(machine: &StarkMachine<BabyBearPoseidon2, A>) -> RecursionProgram<BabyBear> {
-        let mut builder = Builder::<InnerConfig>::default();
+        let mut builder = Builder::<InnerConfig>::new(RecursionProgramType::Deferred);
         let input: SP1DeferredMemoryLayoutVariable<_> = builder.uninit();
         SP1DeferredMemoryLayout::<BabyBearPoseidon2, A>::witness(&input, &mut builder);