fix bug that made benchmark_quda hang randomly

Grid/Benchmark_Grid.cpp

@@ -73,6 +73,8 @@ class Benchmark
         {local[0] * mpi[0], local[1] * mpi[1], local[2] * mpi[2], local[3] * mpi[3]});
     GridCartesian *TmpGrid = SpaceTimeGrid::makeFourDimGrid(
         latt4, GridDefaultSimd(Nd, vComplex::Nsimd()), GridDefaultMpi());
+    Grid::Coordinate shm;
+    GlobalSharedMemory::GetShmDims(mpi, shm);
 
     uint64_t NP = TmpGrid->RankCount();
     uint64_t NN = TmpGrid->NodeCount();
@@ -85,7 +87,9 @@ class Benchmark
     std::cout << GridLogMessage << "* OpenMP threads : " << GridThread::GetThreads()
               << std::endl;
 
-    std::cout << GridLogMessage << "* MPI tasks      : " << GridCmdVectorIntToString(mpi)
+    std::cout << GridLogMessage << "* MPI layout     : " << GridCmdVectorIntToString(mpi)
+              << std::endl;
+    std::cout << GridLogMessage << "* Shm layout     : " << GridCmdVectorIntToString(shm)
               << std::endl;
 
     std::cout << GridLogMessage << "* vReal          : " << sizeof(vReal) * 8 << "bits ; "
@@ -118,6 +122,7 @@ class Benchmark
     for (unsigned int i = 0; i < mpi.size(); ++i)
     {
       tmp["mpi"].push_back(mpi[i]);
+      tmp["shm"].push_back(shm[i]);
     }
     tmp["ranks"] = NP;
     tmp["nodes"] = NN;
@@ -132,6 +137,8 @@ class Benchmark
 
     Coordinate simd_layout = GridDefaultSimd(Nd, vComplexD::Nsimd());
     Coordinate mpi_layout = GridDefaultMpi();
+    Coordinate shm_layout;
+    GlobalSharedMemory::GetShmDims(mpi_layout, shm_layout);
 
     for (int mu = 0; mu < Nd; mu++)
       if (mpi_layout[mu] > 1)
@@ -143,8 +150,8 @@ class Benchmark
     std::cout << GridLogMessage << "Benchmarking threaded STENCIL halo exchange in "
               << nmu << " dimensions" << std::endl;
     grid_small_sep();
-    grid_printf("%5s %5s %15s %15s %15s %15s %15s\n", "L", "dir", "payload (B)",
-                "time (usec)", "rate (GB/s/node)", "std dev", "max");
+    grid_printf("%5s %5s %7s %15s %15s %15s %15s %15s\n", "L", "dir", "shm",
+                "payload (B)", "time (usec)", "rate (GB/s/node)", "std dev", "max");
 
     for (int lat = 16; lat <= maxlat; lat += 8)
     {
@@ -173,74 +180,80 @@ class Benchmark
       for (int dir = 0; dir < 8; dir++)
       {
         int mu = dir % 4;
-        if (mpi_layout[mu] > 1)
+        if (mpi_layout[mu] == 1) // skip directions that are not distributed
+          continue;
+        bool is_shm = mpi_layout[mu] == shm_layout[mu];
+        bool is_partial_shm = !is_shm && shm_layout[mu] != 1;
+
+        std::vector<double> times(Nloop);
+        for (int i = 0; i < NWARMUP; i++)
+        {
+          int xmit_to_rank;
+          int recv_from_rank;
+
+          if (dir == mu)
+          {
+            int comm_proc = 1;
+            Grid.ShiftedRanks(mu, comm_proc, xmit_to_rank, recv_from_rank);
+          }
+          else
+          {
+            int comm_proc = mpi_layout[mu] - 1;
+            Grid.ShiftedRanks(mu, comm_proc, xmit_to_rank, recv_from_rank);
+          }
+          Grid.SendToRecvFrom((void *)&xbuf[dir][0], xmit_to_rank, (void *)&rbuf[dir][0],
+                              recv_from_rank, bytes);
+        }
+        for (int i = 0; i < Nloop; i++)
         {
 
-          std::vector<double> times(Nloop);
-          for (int i = 0; i < NWARMUP; i++)
+          dbytes = 0;
+          double start = usecond();
+          int xmit_to_rank;
+          int recv_from_rank;
+
+          if (dir == mu)
           {
-            int xmit_to_rank;
-            int recv_from_rank;
-
-            if (dir == mu)
-            {
-              int comm_proc = 1;
-              Grid.ShiftedRanks(mu, comm_proc, xmit_to_rank, recv_from_rank);
-            }
-            else
-            {
-              int comm_proc = mpi_layout[mu] - 1;
-              Grid.ShiftedRanks(mu, comm_proc, xmit_to_rank, recv_from_rank);
-            }
-            Grid.SendToRecvFrom((void *)&xbuf[dir][0], xmit_to_rank,
-                                (void *)&rbuf[dir][0], recv_from_rank, bytes);
+            int comm_proc = 1;
+            Grid.ShiftedRanks(mu, comm_proc, xmit_to_rank, recv_from_rank);
           }
-          for (int i = 0; i < Nloop; i++)
+          else
           {
-
-            dbytes = 0;
-            double start = usecond();
-            int xmit_to_rank;
-            int recv_from_rank;
-
-            if (dir == mu)
-            {
-              int comm_proc = 1;
-              Grid.ShiftedRanks(mu, comm_proc, xmit_to_rank, recv_from_rank);
-            }
-            else
-            {
-              int comm_proc = mpi_layout[mu] - 1;
-              Grid.ShiftedRanks(mu, comm_proc, xmit_to_rank, recv_from_rank);
-            }
-            Grid.SendToRecvFrom((void *)&xbuf[dir][0], xmit_to_rank,
-                                (void *)&rbuf[dir][0], recv_from_rank, bytes);
-            dbytes += bytes;
-
-            double stop = usecond();
-            t_time[i] = stop - start; // microseconds
+            int comm_proc = mpi_layout[mu] - 1;
+            Grid.ShiftedRanks(mu, comm_proc, xmit_to_rank, recv_from_rank);
           }
-          timestat.statistics(t_time);
+          Grid.SendToRecvFrom((void *)&xbuf[dir][0], xmit_to_rank, (void *)&rbuf[dir][0],
+                              recv_from_rank, bytes);
+          dbytes += bytes;
 
-          dbytes = dbytes * ppn;
-          double bidibytes = 2. * dbytes;
-          double rate = bidibytes / (timestat.mean / 1.e6) / 1024. / 1024. / 1024.;
-          double rate_err = rate * timestat.err / timestat.mean;
-          double rate_max = rate * timestat.mean / timestat.min;
-          grid_printf("%5d %5d %15d %15.2f %15.2f %15.1f %15.2f\n", lat, dir, bytes,
-                      timestat.mean, rate, rate_err, rate_max);
-          nlohmann::json tmp;
-          nlohmann::json tmp_rate;
-          tmp["L"] = lat;
-          tmp["dir"] = dir;
-          tmp["bytes"] = bytes;
-          tmp["time_usec"] = timestat.mean;
-          tmp_rate["mean"] = rate;
-          tmp_rate["error"] = rate_err;
-          tmp_rate["max"] = rate_max;
-          tmp["rate_GBps"] = tmp_rate;
-          json_results["comms"].push_back(tmp);
+          double stop = usecond();
+          t_time[i] = stop - start; // microseconds
         }
+        timestat.statistics(t_time);
+
+        dbytes = dbytes * ppn;
+        double bidibytes = 2. * dbytes;
+        double rate = bidibytes / (timestat.mean / 1.e6) / 1024. / 1024. / 1024.;
+        double rate_err = rate * timestat.err / timestat.mean;
+        double rate_max = rate * timestat.mean / timestat.min;
+        grid_printf("%5d %5d %7s %15d %15.2f %15.2f %15.1f %15.2f\n", lat, dir,
+                    is_shm           ? "yes"
+                    : is_partial_shm ? "partial"
+                                     : "no",
+                    bytes, timestat.mean, rate, rate_err, rate_max);
+        nlohmann::json tmp;
+        nlohmann::json tmp_rate;
+        tmp["L"] = lat;
+        tmp["dir"] = dir;
+        tmp["shared_mem"] = is_shm;
+        tmp["partial_shared_mem"] = is_partial_shm;
+        tmp["bytes"] = bytes;
+        tmp["time_usec"] = timestat.mean;
+        tmp_rate["mean"] = rate;
+        tmp_rate["error"] = rate_err;
+        tmp_rate["max"] = rate_max;
+        tmp["rate_GBps"] = tmp_rate;
+        json_results["comms"].push_back(tmp);
       }
       for (int d = 0; d < 8; d++)
       {

Quda/Benchmark_Quda.cpp

@@ -2,23 +2,68 @@
 #include <array>
 #include <blas_quda.h>
 #include <cassert>
+#include <chrono>
 #include <color_spinor_field.h>
+#include <communicator_quda.h>
 #include <dirac_quda.h>
+#include <fstream>
 #include <gauge_tools.h>
 #include <memory>
 #include <mpi.h>
 #include <stdio.h>
 #include <stdlib.h>
 
-using namespace quda;
-
 // remove to use QUDA's own flop counting instead of Grid's convention
 #define FLOP_COUNTING_GRID
 
+#include "json.hpp"
+using nlohmann::json;
+json json_results;
+
+using namespace quda;
+
+// timestamp = seconds since program start.
+// these are written to the json output with the goal of later matching them against
+// power-measurments to determine energy efficiency.
+using Clock = std::chrono::steady_clock;
+Clock::time_point program_start_time = Clock::now();
+double get_timestamp()
+{
+  auto dur = Clock::now() - program_start_time;
+  return std::chrono::duration_cast<std::chrono::microseconds>(dur).count() * 1.0e-6;
+}
+
 // This is the MPI grid, i.e. the layout of ranks
 int nranks = -1;
 std::array<int, 4> mpi_grid = {1, 1, 1, 1};
 
+// run f() in a loop for roughly target_time seconds
+// returns seconds per iteration it took
+template <class F> double bench(F const &f, double target_time, int niter_warmup = 5)
+{
+  device_timer_t timer;
+  timer.start();
+  for (int iter = 0; iter < niter_warmup; ++iter)
+    f();
+  timer.stop();
+
+  double secs = timer.last() / niter_warmup;
+  int niter = std::max(1, int(target_time / secs));
+  // niter = std::min(1000, niter);
+  // printfQuda("during warmup took %f s/iter, deciding on %d iters\n", secs, niter);
+
+  // important: each rank has its own timer, so their measurements can slightly vary. But
+  // 'niter' needs to be consistent (bug took me a couple hours to track down)
+  comm_broadcast_global(&niter, sizeof(niter), 0);
+
+  timer.reset(__FUNCTION__, __FILE__, __LINE__);
+  timer.start();
+  for (int iter = 0; iter < niter; ++iter)
+    f();
+  timer.stop();
+  return timer.last() / niter;
+}
+
 void initComms(int argc, char **argv)
 {
   // init MPI communication
@@ -43,6 +88,9 @@ void initComms(int argc, char **argv)
   for (int d = 0; d < 4; d++)
     if (mpi_grid[d] > 1)
       commDimPartitionedSet(d);
+
+  json_results["geometry"]["ranks"] = nranks;
+  json_results["geometry"]["mpi"] = mpi_grid;
 }
 
 // creates a random gauge field. L = local(!) size
@@ -149,11 +197,8 @@ ColorSpinorField make_source(int L, int Ls = 1)
   return src;
 }
 
-void benchmark_wilson()
+void benchmark_wilson(std::vector<int> const &L_list, double target_time)
 {
-  int niter = 20;
-  int niter_warmup = 10;
-
   printfQuda("==================== wilson dirac operator ====================\n");
 #ifdef FLOP_COUNTING_GRID
   printfQuda("IMPORTANT: flop counting as in Benchmark_Grid\n");
@@ -163,8 +208,10 @@ void benchmark_wilson()
 #endif
   printfQuda("%5s %15s %15s\n", "L", "time (usec)", "Gflop/s/rank");
 
-  for (int L : {8, 12, 16, 24, 32, 48})
+  for (int L : L_list)
   {
+    // printfQuda("starting wilson L=%d\n", L);
+
     auto U = make_gauge_field(L);
     auto src = make_source(L);
 
@@ -179,44 +226,41 @@ void benchmark_wilson()
     // (the additional nullptr's are for smeared links and fancy preconditioners and such.
     // Not used for simple Wilson fermions)
     dirac.updateFields(&U, nullptr, nullptr, nullptr);
+    auto res = ColorSpinorField(ColorSpinorParam(src));
+    auto f = [&]() { dirac.Dslash(res, src, QUDA_EVEN_PARITY); };
 
-    auto tmp = ColorSpinorField(ColorSpinorParam(src));
-
-    // couple iterations without timing to warm up
-    for (int iter = 0; iter < niter_warmup; ++iter)
-      dirac.Dslash(tmp, src, QUDA_EVEN_PARITY);
-
-    // actual benchmark with timings
+    // first run to get the quda tuning out of the way
     dirac.Flops(); // reset flops counter
-    device_timer_t device_timer;
-    device_timer.start();
-    for (int iter = 0; iter < niter; ++iter)
-      dirac.Dslash(tmp, src, QUDA_EVEN_PARITY);
-    device_timer.stop();
+    f();
+    double flops = 1.0 * dirac.Flops();
 
-    double secs = device_timer.last() / niter;
+    // actual benchmarking
+    double start_time = get_timestamp();
+    double secs = bench(f, target_time);
+    double end_time = get_timestamp();
 
 #ifdef FLOP_COUNTING_GRID
     // this is the flop counting from Benchmark_Grid
     double Nc = 3;
     double Nd = 4;
     double Ns = 4;
-    double flops =
-        (Nc * (6 + (Nc - 1) * 8) * Ns * Nd + 2 * Nd * Nc * Ns + 2 * Nd * Nc * Ns * 2);
+    flops = (Nc * (6 + (Nc - 1) * 8) * Ns * Nd + 2 * Nd * Nc * Ns + 2 * Nd * Nc * Ns * 2);
     flops *= L * L * L * L / 2.0;
-#else
-    double flops = 1.0 * dirac.Flops() / niter;
 #endif
 
     printfQuda("%5d %15.2f %15.2f\n", L, secs * 1e6, flops / secs * 1e-9);
+
+    json tmp;
+    tmp["L"] = L;
+    tmp["Gflops_wilson"] = flops / secs * 1e-9;
+    tmp["start_time"] = start_time;
+    tmp["end_time"] = end_time;
+    json_results["flops"]["results"].push_back(tmp);
   }
 }
 
-void benchmark_dwf()
+void benchmark_dwf(std::vector<int> const &L_list, double target_time)
 {
-  int niter = 20;
-  int niter_warmup = 10;
-
   printfQuda("==================== domain wall dirac operator ====================\n");
 #ifdef FLOP_COUNTING_GRID
   printfQuda("IMPORTANT: flop counting as in Benchmark_Grid\n");
@@ -226,8 +270,9 @@ void benchmark_dwf()
 #endif
   printfQuda("%5s %15s %15s\n", "L", "time (usec)", "Gflop/s/rank");
   int Ls = 12;
-  for (int L : {8, 12, 16, 24})
+  for (int L : L_list)
   {
+    // printfQuda("starting dwf L=%d\n", L);
     auto U = make_gauge_field(L);
     auto src = make_source(L, Ls);
 
@@ -243,45 +288,43 @@ void benchmark_dwf()
     // insert gauge field into the dirac operator
     // (the additional nullptr's are for smeared links and fancy preconditioners and such)
     dirac.updateFields(&U, nullptr, nullptr, nullptr);
+    auto res = ColorSpinorField(ColorSpinorParam(src));
+    auto f = [&]() { dirac.Dslash(res, src, QUDA_EVEN_PARITY); };
 
-    auto tmp = ColorSpinorField(ColorSpinorParam(src));
-
-    // couple iterations without timing to warm up
-    for (int iter = 0; iter < niter_warmup; ++iter)
-      dirac.Dslash(tmp, src, QUDA_EVEN_PARITY);
-
-    // actual benchmark with timings
+    // first run to get the quda tuning out of the way
     dirac.Flops(); // reset flops counter
-    device_timer_t device_timer;
-    device_timer.start();
-    for (int iter = 0; iter < niter; ++iter)
-      dirac.Dslash(tmp, src, QUDA_EVEN_PARITY);
-    device_timer.stop();
+    f();
+    double flops = 1.0 * dirac.Flops();
 
-    double secs = device_timer.last() / niter;
+    // actual benchmarking
+    double start_time = get_timestamp();
+    double secs = bench(f, target_time);
+    double end_time = get_timestamp();
 
 #ifdef FLOP_COUNTING_GRID
     // this is the flop counting from Benchmark_Grid
     double Nc = 3;
     double Nd = 4;
     double Ns = 4;
-    double flops =
-        (Nc * (6 + (Nc - 1) * 8) * Ns * Nd + 2 * Nd * Nc * Ns + 2 * Nd * Nc * Ns * 2);
+    flops = (Nc * (6 + (Nc - 1) * 8) * Ns * Nd + 2 * Nd * Nc * Ns + 2 * Nd * Nc * Ns * 2);
     flops *= L * L * L * L * Ls / 2.0;
-#else
-    double flops = 1.0 * dirac.Flops() / niter;
 #endif
 
     printfQuda("%5d %15.2f %15.2f\n", L, secs * 1e6, flops / secs * 1e-9);
+    json tmp;
+    tmp["L"] = L;
+    tmp["Gflops_dwf4"] = flops / secs * 1e-9;
+    tmp["start_time"] = start_time;
+    tmp["end_time"] = end_time;
+    json_results["flops"]["results"].push_back(tmp);
   }
 }
 
-void benchmark_axpy()
+void benchmark_axpy(std::vector<int> const &L_list, double target_time)
 {
   // number of iterations for warmup / measurement
   // (feel free to change for noise/time tradeoff)
-  constexpr int niter_warmup = 10;
-  constexpr int niter = 20;
+  constexpr int niter_warmup = 5;
 
   printfQuda("==================== axpy / memory ====================\n");
 
@@ -305,11 +348,11 @@ void benchmark_axpy()
 
   printfQuda("%5s %15s %15s %15s %15s\n", "L", "size (MiB/rank)", "time (usec)",
              "GiB/s/rank", "Gflop/s/rank");
-  std::vector L_list = {8, 12, 16, 24, 32, 48};
   for (int L : L_list)
   {
-    // IMPORTANT: all of `param.x`, `field_elements`, `field.Bytes()`
-    //            are LOCAL, i.e. per rank / per GPU
+    // printfQuda("starting axpy L=%d\n", L);
+    //  IMPORTANT: all of `param.x`, `field_elements`, `field.Bytes()`
+    //             are LOCAL, i.e. per rank / per GPU
 
     param.x[0] = L;
     param.x[1] = L;
@@ -336,26 +379,41 @@ void benchmark_axpy()
     double flops = 2 * field_elements;
     double memory = 3 * sizeof(float) * field_elements;
 
-    // do some iterations to to let QUDA do its internal tuning and also stabilize cache
-    // behaviour and such
-    for (int iter = 0; iter < niter_warmup; ++iter)
-      blas::axpy(1.234, fieldA, fieldB);
+    auto f = [&]() { blas::axpy(1.234, fieldA, fieldB); };
 
-    // running the actual benchmark
-    device_timer_t device_timer;
-    device_timer.start();
-    for (int iter = 0; iter < niter; ++iter)
-      blas::axpy(1.234, fieldA, fieldB);
-    device_timer.stop();
-    double secs = device_timer.last() / niter; // seconds per iteration
+    // first run to get the quda tuning out of the way
+    f();
 
-    printfQuda("%5d %15.2f %15.2f %15.2f %15.2f\n", L, memory / 1024. / 1024., secs * 1e6,
-               memory / secs / 1024. / 1024. / 1024., flops / secs * 1e-9);
+    // actual benchmarking
+    double start_time = get_timestamp();
+    double secs = bench(f, target_time);
+    double end_time = get_timestamp();
+
+    double mem_MiB = memory / 1024. / 1024.;
+    double GBps = mem_MiB / 1024 / secs;
+    printfQuda("%5d %15.2f %15.2f %15.2f %15.2f\n", L, mem_MiB, secs * 1e6, GBps,
+               flops / secs * 1e-9);
+
+    json tmp;
+    tmp["L"] = L;
+    tmp["size_MB"] = mem_MiB;
+    tmp["GBps"] = GBps;
+    tmp["GFlops"] = flops / secs * 1e-9;
+    tmp["start_time"] = start_time;
+    tmp["end_time"] = end_time;
+    json_results["axpy"].push_back(tmp);
   }
 }
 
 int main(int argc, char **argv)
 {
+  std::string json_filename = ""; // empty indicates no json output
+  for (int i = 0; i < argc; i++)
+  {
+    if (std::string(argv[i]) == "--json-out")
+      json_filename = argv[i + 1];
+  }
+
   initComms(argc, argv);
 
   initQuda(-1); // -1 for multi-gpu. otherwise this selects the device to be used
@@ -367,14 +425,28 @@ int main(int argc, char **argv)
   printfQuda("MPI layout = %d %d %d %d\n", mpi_grid[0], mpi_grid[1], mpi_grid[2],
              mpi_grid[3]);
 
-  benchmark_axpy();
+  benchmark_axpy({8, 12, 16, 24, 32, 48}, 1.0);
 
   setVerbosity(QUDA_SILENT);
-  benchmark_wilson();
-  benchmark_dwf();
+  benchmark_wilson({8, 12, 16, 24, 32, 48}, 1.0);
+  benchmark_dwf({8, 12, 16, 24, 32}, 1.0);
   setVerbosity(QUDA_SUMMARIZE);
 
   printfQuda("==================== done with all benchmarks ====================\n");
+
+  if (!json_filename.empty())
+  {
+    printfQuda("writing benchmark results to %s\n", json_filename.c_str());
+
+    int me = 0;
+    MPI_Comm_rank(MPI_COMM_WORLD, &me);
+    if (me == 0)
+    {
+      std::ofstream json_file(json_filename);
+      json_file << std::setw(2) << json_results;
+    }
+  }
+
   endQuda();
   quda::comm_finalize();
   MPI_Finalize();

Quda/build-benchmark.sh

@@ -28,5 +28,5 @@ mkdir -p "${PREFIX_DIR}"
 
 LINK_FLAGS="-Wl,-rpath,$QUDA_DIR/lib: $QUDA_DIR/lib/libquda.so $EXTRA_LIBS -lpthread -lmpi"
 
-g++ $BUILD_FLAGS  -I$QUDA_DIR/include   -c -o $BUILD_DIR/Benchmark_Quda.o  $script_dir/Benchmark_Quda.cpp
+g++ $BUILD_FLAGS -I$QUDA_DIR/include/targets/cuda  -I$QUDA_DIR/include   -c -o $BUILD_DIR/Benchmark_Quda.o  $script_dir/Benchmark_Quda.cpp
 g++ -g -O3 $BUILD_DIR/Benchmark_Quda.o -o $PREFIX_DIR/Benchmark_Quda $LINK_FLAGS -lmpi