#2240: Fix non power of 2 for new allreduce

DARMA-tasking · Apr 10, 2024 · 193f6ed · 193f6ed
1 parent 96a0eda
commit 193f6ed
Show file tree

Hide file tree

Showing 2 changed files with 20 additions and 0 deletions.
diff --git a/src/vt/collective/reduce/allreduce/allreduce.h b/src/vt/collective/reduce/allreduce/allreduce.h
@@ -360,6 +360,18 @@ struct Allreduce {
     }
   }
 
+  void partFour() {
+    if (is_part_of_adjustment_group_ and is_even_) {
+      if constexpr (debug) {
+        fmt::print(
+          "[{}] Part4 : Sending to Node {}  \n", this_node_, this_node_ + 1);
+      }
+      proxy_[this_node_ + 1].template send<&Allreduce::partFourHandler>(val_);
+    }
+  }
+
+  void partFourHandler(AllreduceMsg<DataT>* msg) { val_ = msg->val_; }
+
   NodeType this_node_ = {};
   bool is_even_ = false;
   vt::objgroup::proxy::Proxy<Allreduce> proxy_ = {};

diff --git a/src/vt/objgroup/manager.impl.h b/src/vt/objgroup/manager.impl.h
@@ -274,6 +274,10 @@ ObjGroupManager::allreduce_r(ProxyType<ObjT> proxy, const DataT& data) {
   auto const this_node = vt::theContext()->getNode();
   auto const num_nodes = theContext()->getNumNodes();
 
+  if(num_nodes < 2){
+    return PendingSendType{nullptr};
+  }
+
   using Reducer = collective::reduce::allreduce::Allreduce<DataT>;
 
   auto grp_proxy =
@@ -294,6 +298,10 @@ ObjGroupManager::allreduce_r(ProxyType<ObjT> proxy, const DataT& data) {
     grp_proxy[this_node].template invoke<&Reducer::partThree>();
   });
 
+  vt::runInEpochCollective([=] {
+    grp_proxy[this_node].template invoke<&Reducer::partFour>();
+  });
+
   proxy[this_node].template invoke<f>(grp_proxy.get()->val_);
 
   return PendingSendType{nullptr};