From e3fa49b58a5d535e55d763a6d7eaac0cc5a68815 Mon Sep 17 00:00:00 2001
From: Jacob Domagala <domagala.jacob@gmail.com>
Date: Fri, 26 Apr 2024 11:53:32 +0200
Subject: [PATCH] #2240: Improve accuracy of timing allreduce algorithms in
 allreduce.cc

---
 tests/perf/allreduce.cc | 38 +++++++++++++++++++++-----------------
 1 file changed, 21 insertions(+), 17 deletions(-)

diff --git a/tests/perf/allreduce.cc b/tests/perf/allreduce.cc
index 00f31773a1..a5d9d62d93 100644
--- a/tests/perf/allreduce.cc
+++ b/tests/perf/allreduce.cc
@@ -45,23 +45,20 @@
 #include "vt/collective/reduce/operators/functors/plus_op.h"
 #include "vt/configs/error/config_assert.h"
 #include "vt/context/context.h"
-#include <unordered_map>
+#include "vt/scheduler/scheduler.h"
 #include <vt/collective/collective_ops.h>
 #include <vt/objgroup/manager.h>
 #include <vt/messaging/active.h>
 #include <vt/collective/reduce/allreduce/rabenseifner.h>
 #include <vt/collective/reduce/allreduce/recursive_doubling.h>
 
-#include <fmt-vt/core.h>
-
 using namespace vt;
 using namespace vt::tests::perf::common;
 
-static constexpr int num_iters = 1;
 struct MyTest : PerfTestHarness {
   void SetUp() override {
     PerfTestHarness::SetUp();
-    data.resize(1 << 16);
+    data.resize(1 << 2);
     for (auto& val : data) {
       val = theContext()->getNode() + 1;
     }
@@ -71,7 +68,7 @@ struct MyTest : PerfTestHarness {
 };
 
 struct NodeObj {
-  explicit NodeObj(MyTest* test_obj) : test_obj_(test_obj) { }
+  explicit NodeObj(MyTest* test_obj, const std::string& name) : test_obj_(test_obj), timer_name_(name) { }
 
   void initialize() {
     proxy_ = vt::theObjGroup()->getProxy<NodeObj>(this);
@@ -102,6 +99,7 @@ struct NodeObj {
     // for (auto val : in) {
     //   vtAssert(val == expected, "FAILURE!");
     // }
+    test_obj_->StopTimer(timer_name_);
   }
 
   void newReduceComplete(std::vector<int32_t> in) {
@@ -127,6 +125,7 @@ struct NodeObj {
     // for (auto val : in) {
     //   vtAssert(val == expected, "FAILURE!");
     // }
+    test_obj_->StopTimer(timer_name_);
   }
 
   void reduceComplete(std::vector<int32_t> in) {
@@ -137,25 +136,26 @@ struct NodeObj {
     // }
 
     // fmt::print("\n");
+    test_obj_->StopTimer(timer_name_);
   }
 
-private:
+  std::string timer_name_ = {};
   MyTest* test_obj_ = nullptr;
   vt::objgroup::proxy::Proxy<NodeObj> proxy_ = {};
 };
 
 VT_PERF_TEST(MyTest, test_reduce) {
   auto grp_proxy =
-    vt::theObjGroup()->makeCollective<NodeObj>("test_allreduce", this);
+    vt::theObjGroup()->makeCollective<NodeObj>("test_allreduce", this, "Reduce -> Bcast");
 
-  vt::runInEpochCollective([=] {
-    grp_proxy.allreduce<&NodeObj::reduceComplete, collective::PlusOp>(data);
-  });
+  theCollective()->barrier();
+  StartTimer(grp_proxy[theContext()->getNode()].get()->timer_name_);
+  grp_proxy.allreduce<&NodeObj::reduceComplete, collective::PlusOp>(data);
 }
 
 VT_PERF_TEST(MyTest, test_allreduce_rabenseifner) {
   auto proxy =
-    vt::theObjGroup()->makeCollective<NodeObj>("test_allreduce_new", this);
+    vt::theObjGroup()->makeCollective<NodeObj>("test_allreduce_new", this, "Rabenseifner");
 
   using DataT = decltype(data);
   using Reducer = collective::reduce::allreduce::Rabenseifner<
@@ -164,13 +164,15 @@ VT_PERF_TEST(MyTest, test_allreduce_rabenseifner) {
   auto grp_proxy = vt::theObjGroup()->makeCollective<Reducer>(
     "allreduce_rabenseifner", proxy, num_nodes_, data);
   grp_proxy[my_node_].get()->proxy_ = grp_proxy;
-  vt::runInEpochCollective(
-    [=] { grp_proxy[my_node_].template invoke<&Reducer::allreduce>(); });
+
+  theCollective()->barrier();
+  StartTimer(proxy[theContext()->getNode()].get()->timer_name_);
+  grp_proxy[my_node_].template invoke<&Reducer::allreduce>();
 }
 
 VT_PERF_TEST(MyTest, test_allreduce_recursive_doubling) {
   auto proxy =
-    vt::theObjGroup()->makeCollective<NodeObj>("test_allreduce_new_2", this);
+    vt::theObjGroup()->makeCollective<NodeObj>("test_allreduce_new_2", this, "Recursive doubling");
 
   using DataT = decltype(data);
   using Reducer = collective::reduce::allreduce::DistanceDoubling<
@@ -179,8 +181,10 @@ VT_PERF_TEST(MyTest, test_allreduce_recursive_doubling) {
   auto grp_proxy = vt::theObjGroup()->makeCollective<Reducer>(
     "allreduce_recursive_doubling", proxy, num_nodes_, data);
   grp_proxy[my_node_].get()->proxy_ = grp_proxy;
-  vt::runInEpochCollective(
-    [=] { grp_proxy[my_node_].template invoke<&Reducer::allreduce>(); });
+
+  theCollective()->barrier();
+  StartTimer(proxy[theContext()->getNode()].get()->timer_name_);
+  grp_proxy[my_node_].template invoke<&Reducer::allreduce>();
 }
 
 VT_PERF_TEST_MAIN()