A couple FRI PoW improvements (Plonky3#222)

- Make proof-of-work bits configurable
- Parallelize search for the proof-of-work witness

challenger/Cargo.toml

@@ -6,7 +6,9 @@ license = "MIT OR Apache-2.0"
 
 [dependencies]
 p3-field = { path = "../field" }
+p3-maybe-rayon = { path = "../maybe-rayon" }
 p3-symmetric = { path = "../symmetric" }
+tracing = "0.1.37"
 
 [dev-dependencies]
 p3-goldilocks = { path = "../goldilocks" }

challenger/src/grinding_challenger.rs

@@ -1,23 +1,22 @@
 use p3_field::PrimeField64;
+use p3_maybe_rayon::prelude::*;
 use p3_symmetric::CryptographicPermutation;
+use tracing::instrument;
 
 use crate::{DuplexChallenger, FieldChallenger};
 
 pub trait GrindingChallenger<F: PrimeField64>: FieldChallenger<F> + Clone {
     // Can be overridden for more efficient methods not involving cloning, depending on the
     // internals of the challenger.
+    #[instrument(name = "grind for proof-of-work witness", skip_all)]
     fn grind(&mut self, bits: usize) -> F {
-        for i in 0..F::ORDER_U64 {
-            let witness = F::from_canonical_u64(i);
-            let mut forked = self.clone();
-
-            if forked.check_witness(bits, witness) {
-                assert!(self.check_witness(bits, witness));
-                return witness;
-            }
-        }
-
-        panic!("failed to find witness")
+        let witness = (0..F::ORDER_U64)
+            .into_par_iter()
+            .map(|i| F::from_canonical_u64(i))
+            .find_any(|witness| self.clone().check_witness(bits, *witness))
+            .expect("failed to find witness");
+        assert!(self.check_witness(bits, witness));
+        witness
     }
 
     #[must_use]

challenger/src/lib.rs

@@ -45,7 +45,9 @@ pub trait CanSampleBits<T> {
     fn sample_bits(&mut self, bits: usize) -> T;
 }
 
-pub trait FieldChallenger<F: Field>: CanObserve<F> + CanSample<F> + CanSampleBits<usize> {
+pub trait FieldChallenger<F: Field>:
+    CanObserve<F> + CanSample<F> + CanSampleBits<usize> + Sync
+{
     fn observe_ext_element<EF: AbstractExtensionField<F>>(&mut self, ext: EF) {
         self.observe_slice(ext.as_base_slice());
     }

fri/src/config.rs

@@ -30,18 +30,25 @@ pub trait FriConfig {
 pub struct FriConfigImpl<Val, Challenge, InputMmcs, CommitPhaseMmcs, Challenger> {
     log_blowup: usize,
     num_queries: usize,
+    proof_of_work_bits: usize,
     commit_phase_mmcs: CommitPhaseMmcs,
     _phantom: PhantomData<(Val, Challenge, InputMmcs, Challenger)>,
 }
 
 impl<Val, Challenge, InputMmcs, CommitPhaseMmcs, Challenger>
     FriConfigImpl<Val, Challenge, InputMmcs, CommitPhaseMmcs, Challenger>
 {
-    pub fn new(log_blowup: usize, num_queries: usize, commit_phase_mmcs: CommitPhaseMmcs) -> Self {
+    pub fn new(
+        log_blowup: usize,
+        num_queries: usize,
+        proof_of_work_bits: usize,
+        commit_phase_mmcs: CommitPhaseMmcs,
+    ) -> Self {
         Self {
             log_blowup,
             num_queries,
             commit_phase_mmcs,
+            proof_of_work_bits,
             _phantom: PhantomData,
         }
     }
@@ -76,6 +83,6 @@ where
     }
 
     fn proof_of_work_bits(&self) -> usize {
-        16 // TODO: should make this configurable too
+        self.proof_of_work_bits
     }
 }

fri/tests/fri.rs

@@ -36,7 +36,7 @@ fn get_ldt_for_testing<R: Rng>(rng: &mut R) -> (Perm, ValMmcs, FriLdt<MyFriConfi
     let compress = MyCompress::new(perm.clone());
     let val_mmcs = ValMmcs::new(hash, compress);
     let challenge_mmcs = ChallengeMmcs::new(val_mmcs.clone());
-    let fri_config = MyFriConfig::new(1, 10, challenge_mmcs);
+    let fri_config = MyFriConfig::new(1, 10, 8, challenge_mmcs);
     (perm, val_mmcs, FriLdt { config: fri_config })
 }
 

keccak-air/examples/prove_baby_bear_keccak.rs

@@ -62,7 +62,7 @@ fn main() -> Result<(), VerificationError> {
 
     type Quotient = QuotientMmcs<Domain, Challenge, ValMmcs>;
     type MyFriConfig = FriConfigImpl<Val, Challenge, Quotient, ChallengeMmcs, Challenger>;
-    let fri_config = MyFriConfig::new(1, 100, challenge_mmcs);
+    let fri_config = MyFriConfig::new(1, 100, 16, challenge_mmcs);
     let ldt = FriLdt { config: fri_config };
 
     type Pcs = FriBasedPcs<MyFriConfig, ValMmcs, Dft, Challenger>;

keccak-air/examples/prove_baby_bear_poseidon2.rs

@@ -61,7 +61,7 @@ fn main() -> Result<(), VerificationError> {
 
     type Quotient = QuotientMmcs<Domain, Challenge, ValMmcs>;
     type MyFriConfig = FriConfigImpl<Val, Challenge, Quotient, ChallengeMmcs, Challenger>;
-    let fri_config = MyFriConfig::new(1, 100, challenge_mmcs);
+    let fri_config = MyFriConfig::new(1, 100, 16, challenge_mmcs);
     let ldt = FriLdt { config: fri_config };
 
     type Pcs = FriBasedPcs<MyFriConfig, ValMmcs, Dft, Challenger>;

keccak-air/examples/prove_goldilocks_keccak.rs

@@ -61,7 +61,7 @@ fn main() -> Result<(), VerificationError> {
 
     type Quotient = QuotientMmcs<Domain, Challenge, ValMmcs>;
     type MyFriConfig = FriConfigImpl<Val, Challenge, Quotient, ChallengeMmcs, Challenger>;
-    let fri_config = MyFriConfig::new(1, 100, challenge_mmcs);
+    let fri_config = MyFriConfig::new(1, 100, 16, challenge_mmcs);
     let ldt = FriLdt { config: fri_config };
 
     type Pcs = FriBasedPcs<MyFriConfig, ValMmcs, Dft, Challenger>;

rescue/src/sbox.rs

@@ -4,7 +4,7 @@ use p3_field::{AbstractField, PrimeField, PrimeField64};
 
 use crate::util::get_inverse;
 
-pub trait SboxLayers<AF, const WIDTH: usize>: Clone
+pub trait SboxLayers<AF, const WIDTH: usize>: Clone + Sync
 where
     AF: AbstractField,
     AF::F: PrimeField,

symmetric/src/permutation.rs

@@ -1,5 +1,5 @@
 /// A permutation in the mathematical sense.
-pub trait Permutation<T: Clone>: Clone {
+pub trait Permutation<T: Clone>: Clone + Sync {
     fn permute(&self, mut input: T) -> T {
         self.permute_mut(&mut input);
         input

uni-stark/tests/mul_air.rs

@@ -102,7 +102,7 @@ fn test_prove_baby_bear() -> Result<(), VerificationError> {
 
     type Quotient = QuotientMmcs<Domain, Challenge, ValMmcs>;
     type MyFriConfig = FriConfigImpl<Val, Challenge, Quotient, ChallengeMmcs, Challenger>;
-    let fri_config = MyFriConfig::new(1, 40, challenge_mmcs);
+    let fri_config = MyFriConfig::new(1, 40, 8, challenge_mmcs);
     let ldt = FriLdt { config: fri_config };
 
     type Pcs = FriBasedPcs<MyFriConfig, ValMmcs, Dft, Challenger>;