From d0f212ff2991006e2fb0d8444a4d2613db670fc8 Mon Sep 17 00:00:00 2001
From: Carl Pearson <cwpears@sandia.gov>
Date: Mon, 18 Dec 2023 13:17:16 -0700
Subject: [PATCH] add state parameter to do_iteration, sendrecv perf test

---
 perf_tests/test_sendrecv.cpp       | 33 +++++++++++++++++++++---------
 perf_tests/test_utils.hpp          |  2 +-
 src/KokkosComm_request.hpp         | 18 ++++++----------
 src/impl/KokkosComm_pack.hpp       | 10 ++++-----
 src/impl/KokkosComm_packtraits.hpp |  4 +---
 src/impl/KokkosComm_unpack.hpp     |  8 ++++----
 6 files changed, 39 insertions(+), 36 deletions(-)

diff --git a/perf_tests/test_sendrecv.cpp b/perf_tests/test_sendrecv.cpp
index f52c9f05..fc727c42 100644
--- a/perf_tests/test_sendrecv.cpp
+++ b/perf_tests/test_sendrecv.cpp
@@ -1,22 +1,35 @@
 #include "test_utils.hpp"
 
-#include <thread>
+#include "KokkosComm.hpp"
 
-void i_am_sleepy(MPI_Comm comm, int macsleepface) {
-  // Pretend to work ...
-  std::this_thread::sleep_for(std::chrono::milliseconds(macsleepface));
-  // ... as a team
-  MPI_Barrier(MPI_COMM_WORLD);
+template <typename Space, typename View>
+void send_recv(benchmark::State &state, MPI_Comm comm, const Space &space, int rank, const View &v) {
+  if (0 == rank) {
+    KokkosComm::send(space, v, 1, 0, MPI_COMM_WORLD);
+    KokkosComm::recv(space, v, 1, 0, MPI_COMM_WORLD);
+  } else {
+    KokkosComm::recv(space, v, 0, 0, MPI_COMM_WORLD);
+    KokkosComm::send(space, v, 0, 0, MPI_COMM_WORLD);
+  }
 }
 
 
 
-void mpi_benchmark(benchmark::State &state) {
-  int rank;
+void benchmark_sendrecv(benchmark::State &state) {
+  int rank, size;
   MPI_Comm_rank(MPI_COMM_WORLD, &rank);
+  MPI_Comm_size(MPI_COMM_WORLD, &size);
+  if (size < 2) {
+    state.SkipWithError("benchmark_sendrecv needs at least 2 ranks");
+  }
+
+  auto space = Kokkos::DefaultExecutionSpace();
+  using view_type = Kokkos::View<double *>;
+  view_type a("", 1000000);
+
   while(state.KeepRunning()) {
-    do_iteration(state, MPI_COMM_WORLD, i_am_sleepy, rank % 5);
+    do_iteration(state, MPI_COMM_WORLD, send_recv<Kokkos::DefaultExecutionSpace, view_type>, space, rank, a);
   }
 }
 
-BENCHMARK(mpi_benchmark)->UseManualTime()->Unit(benchmark::kMillisecond);;
\ No newline at end of file
+BENCHMARK(benchmark_sendrecv)->UseManualTime()->Unit(benchmark::kMillisecond);
\ No newline at end of file
diff --git a/perf_tests/test_utils.hpp b/perf_tests/test_utils.hpp
index 288d64e3..88a88139 100644
--- a/perf_tests/test_utils.hpp
+++ b/perf_tests/test_utils.hpp
@@ -14,7 +14,7 @@ void do_iteration(benchmark::State &state, MPI_Comm comm, F &&func, Args... args
     using Duration = std::chrono::duration<double>;
 
     auto start = Clock::now();
-    func(comm, args...);
+    func(state, comm, args...);
     Duration elapsed = Clock::now() - start;
 
     double max_elapsed_second;
diff --git a/src/KokkosComm_request.hpp b/src/KokkosComm_request.hpp
index 8e298e4c..8080184c 100644
--- a/src/KokkosComm_request.hpp
+++ b/src/KokkosComm_request.hpp
@@ -6,19 +6,14 @@
 
 namespace KokkosComm {
 
-
-
-
-
 class Req {
 private:
-
-  // a type-erased view. Request uses these to keep temporary views alive for the lifetime of "Immediate" MPI operations
+  // a type-erased view. Request uses these to keep temporary views alive for
+  // the lifetime of "Immediate" MPI operations
   struct ViewHolderBase {
-    virtual ~ViewHolderBase(){}
+    virtual ~ViewHolderBase() {}
   };
-  template <typename V>
-  struct ViewHolder : ViewHolderBase {
+  template <typename V> struct ViewHolder : ViewHolderBase {
     ViewHolder(const V &v) : v_(v) {}
     V v_;
   };
@@ -28,14 +23,13 @@ class Req {
 
   MPI_Request &mpi_req() { return req_; }
 
-  void wait() { 
+  void wait() {
     MPI_Wait(&req_, MPI_STATUS_IGNORE);
     until_waits_.clear(); // drop any views we're keeping alive until wait()
   }
 
   // keep a reference to this view around until wait() is called
-  template <typename View>
-  void keep_until_wait(const View &v) {
+  template <typename View> void keep_until_wait(const View &v) {
     until_waits_.push_back(std::make_shared<ViewHolder<View>>(v));
   }
 
diff --git a/src/impl/KokkosComm_pack.hpp b/src/impl/KokkosComm_pack.hpp
index 8ea3a005..4de2401f 100644
--- a/src/impl/KokkosComm_pack.hpp
+++ b/src/impl/KokkosComm_pack.hpp
@@ -3,13 +3,11 @@
 #include <Kokkos_Core.hpp>
 
 // impl
-#include "KokkosComm_packtraits.hpp" 
+#include "KokkosComm_packtraits.hpp"
 
 template <typename Dst, typename Src, typename ExecSpace>
 void pack(const ExecSpace &space, const Dst &dst, const Src &src) {
-    Kokkos::Tools::pushRegion("KokkosComm::pack");
-    Kokkos::deep_copy(space, dst, src);
-    Kokkos::Tools::popRegion();
+  Kokkos::Tools::pushRegion("KokkosComm::pack");
+  Kokkos::deep_copy(space, dst, src);
+  Kokkos::Tools::popRegion();
 }
-
-
diff --git a/src/impl/KokkosComm_packtraits.hpp b/src/impl/KokkosComm_packtraits.hpp
index 1a3468c1..f6da942a 100644
--- a/src/impl/KokkosComm_packtraits.hpp
+++ b/src/impl/KokkosComm_packtraits.hpp
@@ -1,5 +1,3 @@
 #pragma once
 
-template <typename ViewType>
-struct PackTraits {
-};
\ No newline at end of file
+template <typename ViewType> struct PackTraits {};
\ No newline at end of file
diff --git a/src/impl/KokkosComm_unpack.hpp b/src/impl/KokkosComm_unpack.hpp
index 6caf7abd..5cd08009 100644
--- a/src/impl/KokkosComm_unpack.hpp
+++ b/src/impl/KokkosComm_unpack.hpp
@@ -3,11 +3,11 @@
 #include <Kokkos_Core.hpp>
 
 // impl
-#include "KokkosComm_packtraits.hpp" 
+#include "KokkosComm_packtraits.hpp"
 
 template <typename Dst, typename Src, typename ExecSpace>
 void unpack(const ExecSpace &space, const Dst &dst, const Src &src) {
-    Kokkos::Tools::pushRegion("KokkosComm::unpack");
-    Kokkos::deep_copy(space, dst, src);
-    Kokkos::Tools::popRegion();
+  Kokkos::Tools::pushRegion("KokkosComm::unpack");
+  Kokkos::deep_copy(space, dst, src);
+  Kokkos::Tools::popRegion();
 }
\ No newline at end of file