From 80d67125b7ec497def363d2717a10a04b42a4d7a Mon Sep 17 00:00:00 2001
From: Jacob Domagala <domagala.jacob@gmail.com>
Date: Tue, 21 May 2024 18:02:03 +0200
Subject: [PATCH] #2240: Add unit tests for new allreduce and cleanup code

---
 .../reduce/allreduce/rabenseifner.h           |  37 ++++--
 .../reduce/allreduce/recursive_doubling.h     |  14 ++-
 src/vt/objgroup/manager.h                     |   6 +
 src/vt/objgroup/manager.impl.h                | 105 ++++++++++--------
 tests/perf/allreduce.cc                       |   6 +-
 tests/unit/objgroup/test_objgroup.cc          |  43 +++++++
 tests/unit/objgroup/test_objgroup_common.h    |  28 +++++
 7 files changed, 178 insertions(+), 61 deletions(-)

diff --git a/src/vt/collective/reduce/allreduce/rabenseifner.h b/src/vt/collective/reduce/allreduce/rabenseifner.h
index fa5ea0a557..c47f5ff0f0 100644
--- a/src/vt/collective/reduce/allreduce/rabenseifner.h
+++ b/src/vt/collective/reduce/allreduce/rabenseifner.h
@@ -41,6 +41,7 @@
 //@HEADER
 */
 
+
 #if !defined INCLUDED_VT_COLLECTIVE_REDUCE_ALLREDUCE_RABENSEIFNER_H
 #define INCLUDED_VT_COLLECTIVE_REDUCE_ALLREDUCE_RABENSEIFNER_H
 
@@ -48,6 +49,8 @@
 #include "vt/context/context.h"
 #include "vt/messaging/message/message.h"
 #include "vt/objgroup/proxy/proxy_objgroup.h"
+#include "vt/registry/auto/auto_registry.h"
+#include "vt/pipe/pipe_manager.h"
 
 #include <tuple>
 #include <cstdint>
@@ -95,7 +98,6 @@ struct Rabenseifner {
     vt::objgroup::proxy::Proxy<ObjT> parentProxy, NodeType num_nodes,
     Args&&... args)
     : parent_proxy_(parentProxy),
-      val_(std::forward<Args>(args)...),
       num_nodes_(num_nodes),
       this_node_(vt::theContext()->getNode()),
       is_even_(this_node_ % 2 == 0),
@@ -104,7 +106,15 @@ struct Rabenseifner {
       nprocs_rem_(num_nodes_ - nprocs_pof2_),
       gather_step_(num_steps_ - 1),
       gather_mask_(nprocs_pof2_ >> 1),
-      finished_adjustment_part_(nprocs_rem_ == 0) {
+      finished_adjustment_part_(nprocs_rem_ == 0)
+  {
+    initialize(std::forward<Args>(args)...);
+  }
+
+  template <typename... Args>
+  void initialize(Args&&... args) {
+    val_ = DataT(std::forward<Args>(args)...);
+
     is_part_of_adjustment_group_ = this_node_ < (2 * nprocs_rem_);
     if (is_part_of_adjustment_group_) {
       if (is_even_) {
@@ -156,6 +166,13 @@ struct Rabenseifner {
     scatter_steps_recv_.resize(num_steps_, false);
   }
 
+  void executeFinalHan() {
+
+    // theCB()->makeSend<finalHandler>(parent_proxy_[this_node_]).sendTuple(std::make_tuple(val_));
+    parent_proxy_[this_node_].template invoke<finalHandler>(val_);
+    completed_ = true;
+  }
+
   void allreduce() {
     if (nprocs_rem_) {
       adjustForPowerOfTwo();
@@ -181,7 +198,7 @@ struct Rabenseifner {
   }
 
   void adjustForPowerOfTwoRightHalf(AllreduceRbnMsg<DataT>* msg) {
-    for (int i = 0; i < msg->val_.size(); i++) {
+    for (uint32_t i = 0; i < msg->val_.size(); i++) {
       val_[(val_.size() / 2) + i] += msg->val_[i];
     }
 
@@ -192,13 +209,13 @@ struct Rabenseifner {
   }
 
   void adjustForPowerOfTwoLeftHalf(AllreduceRbnMsg<DataT>* msg) {
-    for (int i = 0; i < msg->val_.size(); i++) {
+    for (uint32_t i = 0; i < msg->val_.size(); i++) {
       val_[i] += msg->val_[i];
     }
   }
 
   void adjustForPowerOfTwoFinalPart(AllreduceRbnMsg<DataT>* msg) {
-    for (int i = 0; i < msg->val_.size(); i++) {
+    for (uint32_t i = 0; i < msg->val_.size(); i++) {
       val_[(val_.size() / 2) + i] = msg->val_[i];
     }
 
@@ -243,7 +260,7 @@ struct Rabenseifner {
         [](const auto val) { return val; })) {
       auto& in_msg = scatter_messages_.at(step);
       auto& in_val = in_msg->val_;
-      for (int i = 0; i < in_val.size(); i++) {
+      for (uint32_t i = 0; i < in_val.size(); i++) {
         Op<typename DataT::value_type>()(
           val_[r_index_[in_msg->step_] + i], in_val[i]);
       }
@@ -339,7 +356,7 @@ struct Rabenseifner {
     if (doRed) {
       auto& in_msg = gather_messages_.at(step);
       auto& in_val = in_msg->val_;
-      for (int i = 0; i < in_val.size(); i++) {
+      for (uint32_t i = 0; i < in_val.size(); i++) {
         val_[s_index_[in_msg->step_] + i] = in_val[i];
       }
 
@@ -417,8 +434,7 @@ struct Rabenseifner {
       sendToExcludedNodes();
     }
 
-    parent_proxy_[this_node_].template invoke<finalHandler>(val_);
-    completed_ = true;
+    executeFinalHan();
   }
 
   void sendToExcludedNodes() {
@@ -435,8 +451,7 @@ struct Rabenseifner {
   void sendToExcludedNodesHandler(AllreduceRbnMsg<DataT>* msg) {
     val_ = msg->val_;
 
-    parent_proxy_[this_node_].template invoke<finalHandler>(val_);
-    completed_ = true;
+    executeFinalHan();
   }
 
   vt::objgroup::proxy::Proxy<Rabenseifner> proxy_ = {};
diff --git a/src/vt/collective/reduce/allreduce/recursive_doubling.h b/src/vt/collective/reduce/allreduce/recursive_doubling.h
index 917196a4d3..f283d8499f 100644
--- a/src/vt/collective/reduce/allreduce/recursive_doubling.h
+++ b/src/vt/collective/reduce/allreduce/recursive_doubling.h
@@ -97,7 +97,6 @@ struct DistanceDoubling {
     vt::objgroup::proxy::Proxy<ObjT> parentProxy, NodeType num_nodes,
     Args&&... args)
     : parent_proxy_(parentProxy),
-      val_(std::forward<Args>(args)...),
       num_nodes_(num_nodes),
       this_node_(vt::theContext()->getNode()),
       is_even_(this_node_ % 2 == 0),
@@ -105,6 +104,12 @@ struct DistanceDoubling {
       nprocs_pof2_(1 << num_steps_),
       nprocs_rem_(num_nodes_ - nprocs_pof2_),
       finished_adjustment_part_(nprocs_rem_ == 0) {
+        initialize(std::forward<Args>(args)...);
+  }
+
+  template <typename... Args>
+  void initialize(Args&&... args) {
+    val_ = DataT(std::forward<Args>(args)...);
     is_part_of_adjustment_group_ = this_node_ < (2 * nprocs_rem_);
     if (is_part_of_adjustment_group_) {
       if (is_even_) {
@@ -168,8 +173,8 @@ struct DistanceDoubling {
       [](const auto val) { return val; });
   }
   bool isReady() {
-    return (is_part_of_adjustment_group_ and finished_adjustment_part_) and
-      step_ == 0 or
+    return ((is_part_of_adjustment_group_ and finished_adjustment_part_) and
+      step_ == 0) or
       allMessagesReceived();
   }
 
@@ -279,8 +284,9 @@ struct DistanceDoubling {
   vt::objgroup::proxy::Proxy<ObjT> parent_proxy_ = {};
 
   DataT val_ = {};
-  NodeType this_node_ = {};
   NodeType num_nodes_ = {};
+  NodeType this_node_ = {};
+
   bool is_even_ = false;
   int32_t num_steps_ = {};
   int32_t nprocs_pof2_ = {};
diff --git a/src/vt/objgroup/manager.h b/src/vt/objgroup/manager.h
index 3ea00d7191..e4921dfef9 100644
--- a/src/vt/objgroup/manager.h
+++ b/src/vt/objgroup/manager.h
@@ -41,6 +41,7 @@
 //@HEADER
 */
 
+#include "vt/configs/types/types_type.h"
 #if !defined INCLUDED_VT_OBJGROUP_MANAGER_H
 #define INCLUDED_VT_OBJGROUP_MANAGER_H
 
@@ -291,6 +292,9 @@ struct ObjGroupManager : runtime::component::Component<ObjGroupManager> {
     ProxyType<ObjT> proxy, std::string const& name, std::string const& parent = ""
   );
 
+template <typename Reducer, auto f, typename ObjT, template <typename Arg> class Op, typename DataT>
+ObjGroupManager::PendingSendType allreduce(ProxyType<ObjT> proxy, const DataT& data);
+
 template <auto f, typename ObjT, template <typename Arg> class Op, typename DataT>
 ObjGroupManager::PendingSendType allreduce(ProxyType<ObjT> proxy, const DataT& data);
 
@@ -504,6 +508,8 @@ ObjGroupManager::PendingSendType allreduce(ProxyType<ObjT> proxy, const DataT& d
   std::unordered_map<ObjGroupProxyType, std::vector<ActionType>> pending_;
   /// Map of object groups' labels
   std::unordered_map<ObjGroupProxyType, std::string> labels_;
+
+  std::unordered_map<ObjGroupProxyType, ObjGroupProxyType> reducers_;
 };
 
 }} /* end namespace vt::objgroup */
diff --git a/src/vt/objgroup/manager.impl.h b/src/vt/objgroup/manager.impl.h
index 4fb3d3dac9..72e464335e 100644
--- a/src/vt/objgroup/manager.impl.h
+++ b/src/vt/objgroup/manager.impl.h
@@ -41,6 +41,7 @@
 //@HEADER
 */
 
+#include "vt/configs/types/types_sentinels.h"
 #if !defined INCLUDED_VT_OBJGROUP_MANAGER_IMPL_H
 #define INCLUDED_VT_OBJGROUP_MANAGER_IMPL_H
 
@@ -58,7 +59,10 @@
 #include "vt/messaging/active.h"
 #include "vt/elm/elm_id_bits.h"
 #include "vt/messaging/message/smart_ptr.h"
+#include "vt/collective/reduce/allreduce/rabenseifner.h"
+#include "vt/collective/reduce/allreduce/recursive_doubling.h"
 #include <utility>
+#include <array>
 
 #include <memory>
 
@@ -264,57 +268,70 @@ ObjGroupManager::PendingSendType ObjGroupManager::broadcast(MsgSharedPtr<MsgT> m
   return objgroup::broadcast(msg,han);
 }
 
+
+// Helper trait to detect if a type is a specialization of a given variadic template
+template <template <typename...> class Template, typename T>
+struct is_specialization_of : std::false_type {};
+
+template <template <typename...> class Template, typename... Args>
+struct is_specialization_of<Template, Template<Args...>> : std::true_type {};
+
+// Specialized trait for std::array
+template <typename T>
+struct is_std_array : std::false_type {};
+
+template <typename T, std::size_t N>
+struct is_std_array<std::array<T, N>> : std::true_type {};
+
+// Trait to detect if a type is a standard container (std::vector or std::array in this case)
+template <typename T>
+struct is_std_container : std::integral_constant<bool,
+    is_specialization_of<std::vector, T>::value || is_std_array<T>::value> {};
+
+template <
+  typename Reducer, auto f, typename ObjT, template <typename Arg> class Op, typename DataT>
+ObjGroupManager::PendingSendType ObjGroupManager::allreduce(
+  ProxyType<ObjT> proxy, const DataT& data) {
+  return PendingSendType{
+    theTerm()->getEpoch(), [=] {
+      auto const this_node = vt::theContext()->getNode();
+      auto const num_nodes = theContext()->getNumNodes();
+
+      proxy::Proxy<Reducer> grp_proxy = {};
+
+      if (reducers_.find(proxy.getProxy()) != reducers_.end()) {
+        auto* obj = reinterpret_cast<Reducer*>(
+          objs_[reducers_[proxy.getProxy()]]->getPtr()
+        );
+        obj->initialize(data);
+        grp_proxy = obj->proxy_;
+      } else {
+        grp_proxy = vt::theObjGroup()->makeCollective<Reducer>(
+          "allreduce_rabenseifner", proxy, num_nodes, data);
+        grp_proxy[this_node].get()->proxy_ = grp_proxy;
+      }
+
+      grp_proxy[this_node].template invoke<&Reducer::allreduce>();
+    }};
+}
+
 template <
   auto f, typename ObjT, template <typename Arg> class Op, typename DataT>
 ObjGroupManager::PendingSendType
 ObjGroupManager::allreduce(ProxyType<ObjT> proxy, const DataT& data) {
-  // check payload size and choose appropriate algorithm
-
-  auto const this_node = vt::theContext()->getNode();
-  auto const num_nodes = theContext()->getNumNodes();
-
-  if (num_nodes < 2) {
+  if (theContext()->getNumNodes() < 2) {
     return PendingSendType{nullptr};
   }
 
-  // using Reducer = collective::reduce::allreduce::Rabenseifner<DataT>;
-  // using Reducer = collective::reduce::allreduce::DistanceDoubling<DataT, Op, ObjT, f>;
-
-  return PendingSendType{theTerm()->getEpoch(), [=] {
-    // auto grp_proxy =
-    //   vt::theObjGroup()->makeCollective<Reducer>("allreduce_rabenseifner");
-    // if constexpr (std::is_same_v<
-    //                 Reducer,
-    //                 collective::reduce::allreduce::DistanceDoubling<DataT, Op, ObjT, f>>) {
-    //   grp_proxy[this_node].template invoke<&Reducer::initialize>(
-    //     data, grp_proxy, proxy, num_nodes);
-
-    //   grp_proxy[this_node].template invoke<&Reducer::partOne>();
-
-    // } else if constexpr (std::is_same_v<
-    //                        Reducer,
-    //                        collective::reduce::allreduce::Rabenseifner<
-    //                          DataT, Op, ObjT, f>>) {
-    //   grp_proxy[this_node].template invoke<&Reducer::initialize>(
-    //     data, grp_proxy, num_nodes);
-
-    //   if (grp_proxy.get()->nprocs_rem_) {
-    //     vt::runInEpochCollective(
-    //       [=] { grp_proxy[this_node].template invoke<&Reducer::partOne>(); });
-    //   }
-
-    //   vt::runInEpochCollective(
-    //     [=] { grp_proxy[this_node].template invoke<&Reducer::partTwo>(); });
-
-    //   vt::runInEpochCollective(
-    //     [=] { grp_proxy[this_node].template invoke<&Reducer::partThree>(); });
-
-    //   if (grp_proxy.get()->nprocs_rem_) {
-    //     vt::runInEpochCollective(
-    //       [=] { grp_proxy[this_node].template invoke<&Reducer::partFour>(); });
-    //   }
-    // }
-  }};
+  if constexpr (is_std_container<DataT>::value) {
+    using Reducer =
+      vt::collective::reduce::allreduce::Rabenseifner<DataT, Op, ObjT, f>;
+    return allreduce<Reducer, f, ObjT, Op>(proxy, data);
+  } else {
+    using Reducer =
+      vt::collective::reduce::allreduce::DistanceDoubling<DataT, Op, ObjT, f>;
+    return allreduce<Reducer, f, ObjT, Op>(proxy, data);
+  }
 }
 
 template <typename ObjT, typename MsgT, ActiveTypedFnType<MsgT> *f>
diff --git a/tests/perf/allreduce.cc b/tests/perf/allreduce.cc
index eba84358eb..0d394c4fd5 100644
--- a/tests/perf/allreduce.cc
+++ b/tests/perf/allreduce.cc
@@ -184,8 +184,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()
diff --git a/tests/unit/objgroup/test_objgroup.cc b/tests/unit/objgroup/test_objgroup.cc
index ddfc690cd0..ab916e0d20 100644
--- a/tests/unit/objgroup/test_objgroup.cc
+++ b/tests/unit/objgroup/test_objgroup.cc
@@ -256,6 +256,49 @@ TEST_F(TestObjGroup, test_proxy_reduce) {
   }
 }
 
+TEST_F(TestObjGroup, test_proxy_allreduce) {
+  using namespace vt::collective;
+
+  auto const my_node = vt::theContext()->getNode();
+
+  TestObjGroup::total_verify_expected_ = 0;
+  auto proxy = vt::theObjGroup()->makeCollective<MyObjA>("test_proxy_reduce");
+
+  vt::theCollective()->barrier();
+
+  runInEpochCollective(
+    [&] { proxy.allreduce_h<&MyObjA::verifyAllred<1>, PlusOp>(my_node); }
+  );
+
+  EXPECT_EQ(MyObjA::total_verify_expected_, 1);
+
+  runInEpochCollective(
+    [&] { proxy.allreduce_h<&MyObjA::verifyAllred<2>, PlusOp>(4); }
+  );
+
+  EXPECT_EQ(MyObjA::total_verify_expected_, 2);
+
+  runInEpochCollective(
+    [&] { proxy.allreduce_h<&MyObjA::verifyAllred<3>, MaxOp>(my_node); }
+  );
+
+  EXPECT_EQ(MyObjA::total_verify_expected_, 3);
+
+  runInEpochCollective([&] {
+    using Reducer =
+        vt::collective::reduce::allreduce::Rabenseifner<std::vector<int>, PlusOp, MyObjA, &MyObjA::verifyAllredVec>;
+    std::vector<int> payload(256, my_node);
+    theObjGroup()->allreduce<Reducer, &MyObjA::verifyAllredVec, MyObjA, PlusOp>(
+      proxy, payload
+    );
+    theObjGroup()->allreduce<Reducer, &MyObjA::verifyAllredVec, MyObjA, PlusOp>(
+      proxy, payload
+    );
+  });
+
+  EXPECT_EQ(MyObjA::total_verify_expected_, 5);
+}
+
 TEST_F(TestObjGroup, test_proxy_invoke) {
   auto const& this_node = theContext()->getNode();
 
diff --git a/tests/unit/objgroup/test_objgroup_common.h b/tests/unit/objgroup/test_objgroup_common.h
index 67acfd2a3e..10491562cf 100644
--- a/tests/unit/objgroup/test_objgroup_common.h
+++ b/tests/unit/objgroup/test_objgroup_common.h
@@ -41,6 +41,7 @@
 //@HEADER
 */
 
+#include <vector>
 #if !defined INCLUDED_UNIT_OBJGROUP_TEST_OBJGROUP_COMMON_H
 #define INCLUDED_UNIT_OBJGROUP_TEST_OBJGROUP_COMMON_H
 
@@ -97,6 +98,33 @@ struct MyObjA {
     return i;
   }
 
+  static inline int32_t total_verify_expected_ = 0;
+
+  template <int test>
+  void verifyAllred(int value) {
+    auto const n = vt::theContext()->getNumNodes();
+    switch (test) {
+    case 1: EXPECT_EQ(value, n * (n - 1)/2); break;
+    case 2: EXPECT_EQ(value, n * 4); break;
+    case 3: EXPECT_EQ(value, n - 1); break;
+    default: vtAbort("Failure: should not be reached"); break;
+    }
+    total_verify_expected_++;
+  }
+
+  void verifyAllredVec(std::vector<int> vec) {
+    auto final_size = vec.size();
+    EXPECT_EQ(final_size, 256);
+
+    auto n = vt::theContext()->getNumNodes();
+    auto const total_sum = n * (n - 1)/2;
+    for(auto val : vec){
+      EXPECT_EQ(val, total_sum);
+    }
+
+    total_verify_expected_++;
+  }
+
   int id_ = -1;
   int recv_ = 0;
   static int next_id;