From 87e6a2c0b2ac6591430dadae28132abcca234517 Mon Sep 17 00:00:00 2001
From: steviez <steven@anza.xyz>
Date: Mon, 25 Mar 2024 12:16:20 -0500
Subject: [PATCH] Fixup poh_bench

---
 Cargo.lock            |  1 +
 poh-bench/Cargo.toml  |  1 +
 poh-bench/src/main.rs | 23 ++++++++++++++++-------
 3 files changed, 18 insertions(+), 7 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index 6a2abab7f46afc..890e060c9edb7d 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -6621,6 +6621,7 @@ dependencies = [
  "solana-logger",
  "solana-measure",
  "solana-perf",
+ "solana-rayon-threadlimit",
  "solana-sdk",
  "solana-version",
 ]
diff --git a/poh-bench/Cargo.toml b/poh-bench/Cargo.toml
index fb44c0cb81d966..8cd3979b17c79b 100644
--- a/poh-bench/Cargo.toml
+++ b/poh-bench/Cargo.toml
@@ -17,6 +17,7 @@ solana-entry = { workspace = true }
 solana-logger = { workspace = true }
 solana-measure = { workspace = true }
 solana-perf = { workspace = true }
+solana-rayon-threadlimit = { workspace = true }
 solana-sdk = { workspace = true }
 solana-version = { workspace = true }
 
diff --git a/poh-bench/src/main.rs b/poh-bench/src/main.rs
index d835bac05a3ff9..d0a88f4db8fe8d 100644
--- a/poh-bench/src/main.rs
+++ b/poh-bench/src/main.rs
@@ -7,6 +7,7 @@ use {
     clap::{crate_description, crate_name, Arg, Command},
     solana_measure::measure::Measure,
     solana_perf::perf_libs,
+    solana_rayon_threadlimit::get_max_thread_count,
     solana_sdk::hash::hash,
 };
 
@@ -73,6 +74,14 @@ fn main() {
     let start_hash = hash(&[1, 2, 3, 4]);
     let ticks = create_ticks(max_num_entries, hashes_per_tick, start_hash);
     let mut num_entries = start_num_entries as usize;
+    let num_threads = matches
+        .value_of_t("num_threads")
+        .unwrap_or(get_max_thread_count());
+    let threadpool = rayon::ThreadPoolBuilder::new()
+        .num_threads(num_threads)
+        .thread_name(|i| format!("solPohBench{i:02}"))
+        .build()
+        .expect("rayon threadpool");
     if matches.is_present("cuda") {
         perf_libs::init_cuda();
     }
@@ -81,8 +90,8 @@ fn main() {
         let mut time = Measure::start("time");
         for _ in 0..iterations {
             assert!(ticks[..num_entries]
-                .verify_cpu_generic(&start_hash)
-                .finish_verify());
+                .verify_cpu_generic(&start_hash, &threadpool)
+                .finish_verify(&threadpool));
         }
         time.stop();
         println!(
@@ -100,7 +109,7 @@ fn main() {
                 let mut time = Measure::start("time");
                 for _ in 0..iterations {
                     assert!(ticks[..num_entries]
-                        .verify_cpu_x86_simd(&start_hash, 8)
+                        .verify_cpu_x86_simd(&start_hash, 8, &threadpool)
                         .finish_verify());
                 }
                 time.stop();
@@ -115,8 +124,8 @@ fn main() {
                 let mut time = Measure::start("time");
                 for _ in 0..iterations {
                     assert!(ticks[..num_entries]
-                        .verify_cpu_x86_simd(&start_hash, 16)
-                        .finish_verify());
+                        .verify_cpu_x86_simd(&start_hash, 16, &threadpool)
+                        .finish_verify(&threadpool));
                 }
                 time.stop();
                 println!(
@@ -132,8 +141,8 @@ fn main() {
             let recyclers = VerifyRecyclers::default();
             for _ in 0..iterations {
                 assert!(ticks[..num_entries]
-                    .start_verify(&start_hash, recyclers.clone())
-                    .finish_verify());
+                    .start_verify(&start_hash, &threadpool, recyclers.clone())
+                    .finish_verify(&threadpool));
             }
             time.stop();
             println!(