#2240: Add unit tests for new allreduce and cleanup code

Author: JacobDomagala

src/vt/collective/reduce/allreduce/rabenseifner.h

@@ -41,13 +41,16 @@
 //@HEADER
 */
 
+
 #if !defined INCLUDED_VT_COLLECTIVE_REDUCE_ALLREDUCE_RABENSEIFNER_H
 #define INCLUDED_VT_COLLECTIVE_REDUCE_ALLREDUCE_RABENSEIFNER_H
 
 #include "vt/config.h"
 #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_ = {};

src/vt/collective/reduce/allreduce/recursive_doubling.h

@@ -97,14 +97,19 @@ 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),
       num_steps_(static_cast<int32_t>(log2(num_nodes_))),
       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_ = {};

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 */

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>

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()

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();