From 7e21160206ff7ffd087dc4d3bbd956f4521cb058 Mon Sep 17 00:00:00 2001
From: Shahar Papini <spapini@starkware.co>
Date: Mon, 18 Mar 2024 18:43:31 +0200
Subject: [PATCH] parallel wide fib test

---
 Cargo.lock                   |  1 +
 Cargo.toml                   |  1 +
 src/examples/wide_fib/mod.rs | 14 +++++++++-----
 3 files changed, 11 insertions(+), 5 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index 9d93e03c0..68af89fb5 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -676,6 +676,7 @@ dependencies = [
  "merging-iterator",
  "num-traits",
  "rand",
+ "rayon",
  "test-log",
  "thiserror",
  "tracing",
diff --git a/Cargo.toml b/Cargo.toml
index b2fc0b5ee..0d5c81f0e 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -15,6 +15,7 @@ thiserror = "1.0.56"
 merging-iterator = "1.3.0"
 bytemuck = { version = "1.14.3", features = ["derive"] }
 tracing = "0.1.40"
+rayon = "1.9.0"
 
 [dev-dependencies]
 criterion = { version = "0.5.1", features = ["html_reports"] }
diff --git a/src/examples/wide_fib/mod.rs b/src/examples/wide_fib/mod.rs
index 77262faaf..eadd183f3 100644
--- a/src/examples/wide_fib/mod.rs
+++ b/src/examples/wide_fib/mod.rs
@@ -148,6 +148,8 @@ impl Component<AVX512Backend> for WideFibComponent {
 #[cfg(all(target_arch = "x86_64", target_feature = "avx512f"))]
 #[cfg(test)]
 mod tests {
+    use rayon::prelude::*;
+
     use crate::commitment_scheme::blake2_hash::Blake2sHasher;
     use crate::commitment_scheme::hasher::Hasher;
     use crate::core::channel::{Blake2sChannel, Channel};
@@ -165,10 +167,12 @@ mod tests {
 
         // TODO(spapini): Increase to 20, to get 1GB.
         const LOG_SIZE: u32 = 12;
-        let component = WideFibComponent { log_size: LOG_SIZE };
-        let air = WideFibAir { component };
-        let trace = gen_trace(LOG_SIZE as usize);
-        let channel = &mut Blake2sChannel::new(Blake2sHasher::hash(BaseField::into_slice(&[])));
-        prove(&air, channel, trace);
+        (0..8).into_par_iter().for_each(|_| {
+            let component = WideFibComponent { log_size: LOG_SIZE };
+            let air = WideFibAir { component };
+            let trace = gen_trace(LOG_SIZE as usize);
+            let channel = &mut Blake2sChannel::new(Blake2sHasher::hash(BaseField::into_slice(&[])));
+            prove(&air, channel, trace);
+        });
     }
 }