From 7a021793ee1ca70b4744ce036c7785b0d2e97d13 Mon Sep 17 00:00:00 2001
From: Raph Levien <raph.levien@gmail.com>
Date: Thu, 11 Nov 2021 07:26:32 -0800
Subject: [PATCH] Configure number of iterations

---
 tests/src/clear.rs       |  9 ++-------
 tests/src/config.rs      | 11 ++++++++++-
 tests/src/prefix.rs      |  2 +-
 tests/src/prefix_tree.rs |  2 +-
 4 files changed, 14 insertions(+), 10 deletions(-)

diff --git a/tests/src/clear.rs b/tests/src/clear.rs
index a7934d1..f691928 100644
--- a/tests/src/clear.rs
+++ b/tests/src/clear.rs
@@ -43,14 +43,12 @@ pub struct ClearBinding {
 
 pub unsafe fn run_clear_test(runner: &mut Runner, config: &Config) -> TestResult {
     let mut result = TestResult::new("clear buffers");
-    // This will be configurable.
     let n_elements: u64 = config.size.choose(1 << 12, 1 << 20, 1 << 24);
     let out_buf = runner.buf_down(n_elements * 4);
     let code = ClearCode::new(runner);
     let stage = ClearStage::new_with_value(runner, n_elements, 0x42);
     let binding = stage.bind(runner, &code, &out_buf.dev_buf);
-    // Also will be configurable of course.
-    let n_iter = 1000;
+    let n_iter = config.n_iter;
     let mut total_elapsed = 0.0;
     for i in 0..n_iter {
         let mut commands = runner.commands();
@@ -79,10 +77,7 @@ impl ClearCode {
         let code = include_shader!(&runner.session, "../shader/gen/Clear");
         let pipeline = runner
             .session
-            .create_compute_pipeline(
-                code,
-                &[BindType::BufReadOnly, BindType::Buffer],
-            )
+            .create_compute_pipeline(code, &[BindType::BufReadOnly, BindType::Buffer])
             .unwrap();
         ClearCode { pipeline }
     }
diff --git a/tests/src/config.rs b/tests/src/config.rs
index 1ead3bd..edc1140 100644
--- a/tests/src/config.rs
+++ b/tests/src/config.rs
@@ -21,6 +21,7 @@ use clap::ArgMatches;
 pub struct Config {
     pub groups: Groups,
     pub size: Size,
+    pub n_iter: u64,
 }
 
 pub struct Groups(String);
@@ -35,7 +36,15 @@ impl Config {
     pub fn from_matches(matches: &ArgMatches) -> Config {
         let groups = Groups::from_str(matches.value_of("groups").unwrap_or("all"));
         let size = Size::from_str(matches.value_of("size").unwrap_or("m"));
-        Config { groups, size }
+        let n_iter = matches
+            .value_of("n_iter")
+            .and_then(|s| s.parse().ok())
+            .unwrap_or(1000);
+        Config {
+            groups,
+            size,
+            n_iter,
+        }
     }
 }
 
diff --git a/tests/src/prefix.rs b/tests/src/prefix.rs
index d431480..b668fac 100644
--- a/tests/src/prefix.rs
+++ b/tests/src/prefix.rs
@@ -69,7 +69,7 @@ pub unsafe fn run_prefix_test(runner: &mut Runner, config: &Config) -> TestResul
     let stage = PrefixStage::new(runner, &code, n_elements);
     let binding = stage.bind(runner, &code, &data_buf, &out_buf.dev_buf);
     // Also will be configurable of course.
-    let n_iter = 1000;
+    let n_iter = config.n_iter;
     let mut total_elapsed = 0.0;
     for i in 0..n_iter {
         let mut commands = runner.commands();
diff --git a/tests/src/prefix_tree.rs b/tests/src/prefix_tree.rs
index 762772e..80a332f 100644
--- a/tests/src/prefix_tree.rs
+++ b/tests/src/prefix_tree.rs
@@ -57,7 +57,7 @@ pub unsafe fn run_prefix_test(runner: &mut Runner, config: &Config) -> TestResul
     let stage = PrefixTreeStage::new(runner, n_elements);
     let binding = stage.bind(runner, &code, &out_buf.dev_buf);
     // Also will be configurable of course.
-    let n_iter = 1000;
+    let n_iter = config.n_iter;
     let mut total_elapsed = 0.0;
     for i in 0..n_iter {
         let mut commands = runner.commands();