From 32a1bd23eec49b1aaf171cd591999088e13d6e3a Mon Sep 17 00:00:00 2001
From: Jonathan Lifflander <jliffla@sandia.gov>
Date: Wed, 31 May 2023 14:15:59 -0700
Subject: [PATCH] #276: reduce: simplify stamp and message extraction

---
 src/vt/collective/reduce/get_reduce_stamp.h   | 33 ++++++++++++-------
 src/vt/collective/reduce/reduce.h             |  6 +---
 src/vt/collective/reduce/reduce.impl.h        | 20 ++++++++---
 src/vt/objgroup/proxy/proxy_objgroup.impl.h   | 21 ++----------
 .../vrt/collection/reducable/reducable.impl.h | 22 ++-----------
 5 files changed, 44 insertions(+), 58 deletions(-)

diff --git a/src/vt/collective/reduce/get_reduce_stamp.h b/src/vt/collective/reduce/get_reduce_stamp.h
index 03032b2406..1d3783cece 100644
--- a/src/vt/collective/reduce/get_reduce_stamp.h
+++ b/src/vt/collective/reduce/get_reduce_stamp.h
@@ -57,8 +57,12 @@ namespace vt { namespace collective { namespace reduce {
 template <typename enable = void, typename... Args>
 struct GetReduceStamp : std::false_type {
   template <typename MsgT>
-  static auto getMsg(Args&&... args) {
-    return vt::makeMessage<MsgT>(std::tuple{std::forward<Args>(args)...});
+  static auto getStampMsg(Args&&... args) {
+    return
+      std::make_tuple(
+	collective::reduce::ReduceStamp{},
+	vt::makeMessage<MsgT>(std::tuple{std::forward<Args>(args)...})
+      );
   }
 };
 
@@ -67,8 +71,11 @@ struct GetReduceStamp<
   std::enable_if_t<std::is_same_v<void, void>>
 > : std::false_type {
   template <typename MsgT>
-  static auto getMsg() {
-    return vt::makeMessage<MsgT>(std::tuple<>{});
+  static auto getStampMsg() {
+    return std::make_tuple(
+      collective::reduce::ReduceStamp{},
+      vt::makeMessage<MsgT>(std::tuple<>{})
+    );
   }
 };
 
@@ -90,13 +97,17 @@ struct GetReduceStamp<
   }
 
   template <typename MsgT>
-  static auto getMsg(Args&&... args) {
-    return vt::makeMessage<MsgT>(
-      getMsgHelper(
-        std::tie(std::forward<Args>(args)...),
-        std::make_index_sequence<sizeof...(Args) - 1>{}
-      )
-    );
+  static auto getStampMsg(Args&&... args) {
+    auto tp = std::make_tuple(std::forward<Args>(args)...);
+    return
+      std::make_tuple(
+        std::get<sizeof...(Args) - 1>(tp),
+	vt::makeMessage<MsgT>(
+          getMsgHelper(
+            std::move(tp), std::make_index_sequence<sizeof...(Args) - 1>{}
+          )
+	)
+      );
   }
 };
 
diff --git a/src/vt/collective/reduce/reduce.h b/src/vt/collective/reduce/reduce.h
index 1293782edb..37b1905b2f 100644
--- a/src/vt/collective/reduce/reduce.h
+++ b/src/vt/collective/reduce/reduce.h
@@ -181,11 +181,7 @@ struct Reduce : virtual collective::tree::Tree {
     template <typename Arg> class Op = NoneOp,
     typename... Params
   >
-  PendingSendType reduce(Node root, Params&&... params) {
-    using Tuple = typename FuncTraits<decltype(f)>::TupleType;
-    using OpT = Op<Tuple>;
-    return reduce<OpT, f>(root, std::forward<Params>(params)...);
-  }
+  PendingSendType reduce(Node root, Params&&... params);
 
   template <typename Op, auto f, typename... Params>
   PendingSendType reduce(Node root, Params&&... params);
diff --git a/src/vt/collective/reduce/reduce.impl.h b/src/vt/collective/reduce/reduce.impl.h
index 692297efa0..d91de33e12 100644
--- a/src/vt/collective/reduce/reduce.impl.h
+++ b/src/vt/collective/reduce/reduce.impl.h
@@ -50,6 +50,7 @@
 #include "vt/messaging/active.h"
 #include "vt/messaging/message.h"
 #include "vt/runnable/make_runnable.h"
+#include "vt/collective/reduce/get_reduce_stamp.h"
 
 namespace vt { namespace collective { namespace reduce {
 
@@ -83,20 +84,29 @@ void Reduce::reduceRootRecv(MsgT* msg) {
     .run();
 }
 
+template <
+  auto f,
+  template <typename Arg> class Op,
+  typename... Params
+>
+Reduce::PendingSendType Reduce::reduce(Node root, Params&&... params) {
+  using Tuple = typename FuncTraits<decltype(f)>::TupleType;
+  using OpT = Op<Tuple>;
+  return reduce<OpT, f>(root, std::forward<Params>(params)...);
+}
+
 template <typename Op, auto f, typename... Params>
 Reduce::PendingSendType Reduce::reduce(Node root, Params&&... params) {
   using Tuple = typename FuncTraits<decltype(f)>::TupleType;
   using MsgT = ReduceTMsg<Tuple>;
-
-  auto msg = vt::makeMessage<MsgT>(std::tuple{std::forward<Params>(params)...});
-  auto id = detail::ReduceStamp{};
+  using GetReduceStamp = collective::reduce::GetReduceStamp<void, Params...>;
+  auto [stamp, msg] = GetReduceStamp::template getStampMsg<MsgT>(std::forward<Params>(params)...);
   auto han = auto_registry::makeAutoHandlerParam<decltype(f), f, MsgT>();
   msg->root_handler_ = han;
 
-  return reduce<Op, operators::NoCombine, MsgT>(root.get(), msg.get(), id, 1);
+  return reduce<Op, operators::NoCombine, MsgT>(root.get(), msg.get(), stamp, 1);
 }
 
-
 template <typename OpT, typename MsgT, ActiveTypedFnType<MsgT> *f>
 Reduce::PendingSendType Reduce::reduce(
   NodeType const& root, MsgT* msg, Callback<MsgT> cb, detail::ReduceStamp id,
diff --git a/src/vt/objgroup/proxy/proxy_objgroup.impl.h b/src/vt/objgroup/proxy/proxy_objgroup.impl.h
index 21c380608f..50d28298f0 100644
--- a/src/vt/objgroup/proxy/proxy_objgroup.impl.h
+++ b/src/vt/objgroup/proxy/proxy_objgroup.impl.h
@@ -124,12 +124,7 @@ Proxy<ObjT>::allreduce(
   using GetReduceStamp = collective::reduce::GetReduceStamp<void, Args...>;
   auto cb = theCB()->makeBcast<f>(*this);
 
-  ReduceStamp stamp{};
-  if constexpr (GetReduceStamp::value) {
-    stamp = std::get<sizeof...(Args) - 1>(std::tie(std::forward<Args>(args)...));
-  }
-
-  auto msg = GetReduceStamp::template getMsg<MsgT>(std::forward<Args>(args)...);
+  auto [stamp, msg] = GetReduceStamp::template getStampMsg<MsgT>(std::forward<Args>(args)...);
   msg->setCallback(cb);
   auto proxy = Proxy<ObjT>(*this);
   return theObjGroup()->reduce<
@@ -158,12 +153,7 @@ Proxy<ObjT>::reduce(
   using GetReduceStamp = collective::reduce::GetReduceStamp<void, Args...>;
   auto cb = theCB()->makeSend<f>(target);
 
-  ReduceStamp stamp{};
-  if constexpr (GetReduceStamp::value ) {
-    stamp = std::get<sizeof...(Args) - 1>(std::tie(std::forward<Args>(args)...));
-  }
-
-  auto msg = GetReduceStamp::template getMsg<MsgT>(std::forward<Args>(args)...);
+  auto [stamp, msg] = GetReduceStamp::template getStampMsg<MsgT>(std::forward<Args>(args)...);
   msg->setCallback(cb);
   auto proxy = Proxy<ObjT>(*this);
   return theObjGroup()->reduce<
@@ -191,12 +181,7 @@ Proxy<ObjT>::reduce(
   using MsgT = collective::ReduceTMsg<Tuple>;
   using GetReduceStamp = collective::reduce::GetReduceStamp<void, Args...>;
 
-  ReduceStamp stamp{};
-  if constexpr (GetReduceStamp::value ) {
-    stamp = std::get<sizeof...(Args) - 1>(std::tie(std::forward<Args>(args)...));
-  }
-
-  auto msg = GetReduceStamp::template getMsg<MsgT>(std::forward<Args>(args)...);
+  auto [stamp, msg] = GetReduceStamp::template getStampMsg<MsgT>(std::forward<Args>(args)...);
   msg->setCallback(cb);
   auto proxy = Proxy<ObjT>(*this);
   return theObjGroup()->reduce<
diff --git a/src/vt/vrt/collection/reducable/reducable.impl.h b/src/vt/vrt/collection/reducable/reducable.impl.h
index dd0481fe7e..39afc6fca6 100644
--- a/src/vt/vrt/collection/reducable/reducable.impl.h
+++ b/src/vt/vrt/collection/reducable/reducable.impl.h
@@ -67,13 +67,7 @@ messaging::PendingSend Reducable<ColT,IndexT,BaseProxyT>::allreduce(
   using MsgT = collective::ReduceTMsg<Tuple>;
   using GetReduceStamp = collective::reduce::GetReduceStamp<void, Args...>;
   auto cb = theCB()->makeBcast<f>(*this);
-
-  ReduceStamp stamp{};
-  if constexpr (GetReduceStamp::value) {
-    stamp = std::get<sizeof...(Args) - 1>(std::tie(std::forward<Args>(args)...));
-  }
-
-  auto msg = GetReduceStamp::template getMsg<MsgT>(std::forward<Args>(args)...);
+  auto [stamp, msg] = GetReduceStamp::template getStampMsg<MsgT>(std::forward<Args>(args)...);
   msg->setCallback(cb);
   auto const root_node = 0;
   auto const proxy = this->getProxy();
@@ -98,12 +92,7 @@ messaging::PendingSend Reducable<ColT,IndexT,BaseProxyT>::reduce(
 
   auto cb = theCB()->makeSend<f>(target);
 
-  ReduceStamp stamp{};
-  if constexpr (GetReduceStamp::value) {
-    stamp = std::get<sizeof...(Args) - 1>(std::tie(std::forward<Args>(args)...));
-  }
-
-  auto msg = GetReduceStamp::template getMsg<MsgT>(std::forward<Args>(args)...);
+  auto [stamp, msg] = GetReduceStamp::template getStampMsg<MsgT>(std::forward<Args>(args)...);
   msg->setCallback(cb);
   auto const root_node = 0;
   auto const proxy = this->getProxy();
@@ -128,12 +117,7 @@ messaging::PendingSend Reducable<ColT,IndexT,BaseProxyT>::reduce(
   using MsgT = collective::ReduceTMsg<Tuple>;
   using GetReduceStamp = collective::reduce::GetReduceStamp<void, Args...>;
 
-  ReduceStamp stamp{};
-  if constexpr (GetReduceStamp::value) {
-    stamp = std::get<sizeof...(Args) - 1>(std::tie(std::forward<Args>(args)...));
-  }
-
-  auto msg = GetReduceStamp::template getMsg<MsgT>(std::forward<Args>(args)...);
+  auto [stamp, msg] = GetReduceStamp::template getStampMsg<MsgT>(std::forward<Args>(args)...);
   msg->setCallback(cb);
   auto const root_node = 0;
   auto const proxy = this->getProxy();