From 8efa3d911ec3946c30b1a034e3ae895926d58fe8 Mon Sep 17 00:00:00 2001
From: Alejandro Gallo <aamsgallo@gmail.com>
Date: Tue, 6 Dec 2022 20:38:38 +0100
Subject: [PATCH] Add --max-iterations to main bench

---
 bench/main.cxx          | 14 +++++++++-----
 include/atrip/Atrip.hpp |  2 +-
 src/atrip/Atrip.cxx     |  2 ++
 3 files changed, 12 insertions(+), 6 deletions(-)

diff --git a/bench/main.cxx b/bench/main.cxx
index 0e4129e..dd1bb2d 100644
--- a/bench/main.cxx
+++ b/bench/main.cxx
@@ -16,7 +16,7 @@
 int main(int argc, char** argv) {
   MPI_Init(&argc, &argv);
 
-  size_t checkpoint_it;
+  size_t checkpoint_it, max_iterations;
   int no(10), nv(100), itMod(-1), percentageMod(10);
   float checkpoint_percentage;
   bool
@@ -30,6 +30,9 @@ int main(int argc, char** argv) {
   app.add_option("--no", no, "Occupied orbitals");
   app.add_option("--nv", nv, "Virtual orbitals");
   app.add_option("--mod", itMod, "Iteration modifier");
+  app.add_option("--max-iterations",
+                 max_iterations,
+                 "Maximum number of iterations to run");
   app.add_flag("--keep-vppph", keepVppph, "Do not delete Vppph");
   app.add_flag("--nochrono", nochrono, "Do not print chrono");
   app.add_flag("--rank-round-robin", rankRoundRobin, "Do rank round robin");
@@ -48,11 +51,11 @@ int main(int argc, char** argv) {
 #if defined(HAVE_CUDA)
   size_t ooo_threads = 0, ooo_blocks = 0;
   app.add_option("--ooo-blocks",
-		 ooo_blocks,
-		 "CUDA: Number of blocks per block for kernels going through ooo tensors");
+                 ooo_blocks,
+                 "CUDA: Number of blocks per block for kernels going through ooo tensors");
   app.add_option("--ooo-threads",
-		 ooo_threads,
-		 "CUDA: Number of threads per block for kernels going through ooo tensors");
+                 ooo_threads,
+                 "CUDA: Number of threads per block for kernels going through ooo tensors");
 #endif
 
   CLI11_PARSE(app, argc, argv);
@@ -199,6 +202,7 @@ int main(int argc, char** argv) {
        .with_iterationMod(itMod)
        .with_percentageMod(percentageMod)
        .with_tuplesDistribution(tuplesDistribution)
+       .with_maxIterations(max_iterations)
        // checkpoint options
        .with_checkpointAtEveryIteration(checkpoint_it)
        .with_checkpointAtPercentage(checkpoint_percentage)
diff --git a/include/atrip/Atrip.hpp b/include/atrip/Atrip.hpp
index 9567529..b4533dd 100644
--- a/include/atrip/Atrip.hpp
+++ b/include/atrip/Atrip.hpp
@@ -86,7 +86,7 @@ namespace atrip {
       ADD_ATTRIBUTE(bool, rankRoundRobin, false)
       ADD_ATTRIBUTE(bool, chrono, false)
       ADD_ATTRIBUTE(bool, barrier, false)
-      ADD_ATTRIBUTE(int, maxIterations, 0)
+      ADD_ATTRIBUTE(size_t, maxIterations, 0)
       ADD_ATTRIBUTE(int, iterationMod, -1)
       ADD_ATTRIBUTE(int, percentageMod, -1)
       ADD_ATTRIBUTE(TuplesDistribution, tuplesDistribution, NAIVE)
diff --git a/src/atrip/Atrip.cxx b/src/atrip/Atrip.cxx
index d85cf20..a946459 100644
--- a/src/atrip/Atrip.cxx
+++ b/src/atrip/Atrip.cxx
@@ -773,6 +773,8 @@ Atrip::Output Atrip::run(Atrip::Input<F> const& in) {
     Atrip::chrono["iterations"].stop();
     // ITERATION END %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%{{{1
 
+    if (in.maxIterations != 0 && i >= in.maxIterations) break;
+
   }
     // END OF MAIN LOOP