#276: Synthesize Msg type in vt/collection

DARMA-tasking · Feb 10, 2023 · 57c81e5 · 57c81e5
1 parent f40ee4d
commit 57c81e5
Show file tree

Hide file tree

Showing 7 changed files with 196 additions and 10 deletions.
diff --git a/src/vt/vrt/collection/broadcast/broadcastable.h b/src/vt/vrt/collection/broadcast/broadcastable.h
@@ -195,6 +195,94 @@ struct Broadcastable : BaseProxyT {
     typename... Args
   >
   void invokeCollective(Args&&... args) const;
+
+
+  template <typename Return, typename... Args>
+  struct FunctionTraits;
+
+  template <typename Return, typename Msg>
+  struct FunctionTraits<Return(*)(Msg*, ColT*)> {
+    using MsgT = Msg;
+    using ReturnT = Return;
+  };
+
+  template <typename Return, typename Msg>
+  struct FunctionTraits<Return(ColT::*)(Msg*)> {
+    using MsgT = Msg;
+    using ReturnT = Return;
+  };
+
+  /**
+   * \brief Rooted broadcast with action function handler
+   * \note Takes ownership of the supplied message
+   *
+   * \param[in] msg the message
+   *
+   * \return a pending send
+   */
+  template <auto f>
+  messaging::PendingSend broadcastMsg(
+    messaging::MsgPtrThief<typename FunctionTraits<decltype(f)>::MsgT> msg
+  ) const {
+    using MsgT = typename FunctionTraits<decltype(f)>::MsgT;
+    return broadcastMsg<MsgT, f>(msg);
+  }
+
+  /**
+   * \brief Rooted broadcast with action function handler
+   *
+   * \param[in] args arguments needed for creteating the message
+   *
+   * \return a pending send
+   */
+  template <auto f, typename... Args>
+  messaging::PendingSend broadcast(Args&&... args) const {
+    using MsgT = typename FunctionTraits<decltype(f)>::MsgT;
+    return broadcast<MsgT, f>(std::forward<Args>(args)...);
+  }
+
+  /**
+   * \brief Collective broadcast with action function handler
+   * \note Takes ownership of the supplied message
+   *
+   * \param[in] msg the message
+   *
+   * \return a pending send
+   */
+  template <auto f>
+  messaging::PendingSend broadcastCollectiveMsg(
+    messaging::MsgPtrThief<typename FunctionTraits<decltype(f)>::MsgT> msg
+  ) const {
+    using MsgT = typename FunctionTraits<decltype(f)>::MsgT;
+    return broadcastCollectiveMsg<MsgT, f>(msg);
+  }
+
+  /**
+   * \brief Create message (with action function handler) and broadcast it in a
+   * collective manner to the collection
+   *
+   * \param[in] args arguments needed for creteating the message
+   *
+   * \return a pending send
+   */
+  template <auto f, typename... Args>
+  messaging::PendingSend broadcastCollective(Args&&... args) const {
+    using MsgT = typename FunctionTraits<decltype(f)>::MsgT;
+    return broadcastCollectiveMsg<MsgT, f>(std::forward<Args>(args)...);
+  }
+
+  /**
+   * \brief Invoke member message handler on all collection elements
+   * The function will be invoked inline without going through scheduler
+   *
+   * \param[in] args arguments for creating the message
+   */
+  template <auto f, typename... Args>
+  void invokeCollective(Args&&... args) const {
+    using MsgT = typename FunctionTraits<decltype(f)>::MsgT;
+    return invokeCollective<MsgT, f>(std::forward<Args>(args)...);
+  }
+
 };
 
 }}} /* end namespace vt::vrt::collection */

diff --git a/src/vt/vrt/collection/manager.h b/src/vt/vrt/collection/manager.h
@@ -493,6 +493,40 @@ struct CollectionManager
     VirtualElmProxyType<typename MsgT::CollectionType> const& proxy, MsgT *msg
   );
 
+  template <typename Return, typename... Args>
+  struct FunctionTraits;
+
+  template <typename Return, typename Msg, typename Col>
+  struct FunctionTraits<Return(*)(Msg*, Col*)> {
+    using MsgT = Msg;
+    using ColT = Col;
+    using ReturnT = Return;
+  };
+
+  template <typename Return, typename Msg, typename Col>
+  struct FunctionTraits<Return(Col::*)(Msg*)> {
+    using ColT = Col;
+    using MsgT = Msg;
+    using ReturnT = Return;
+  };
+
+  /**
+   * \brief Send collection element a message from active function handler
+   *
+   * \param[in] proxy the collection proxy
+   * \param[in] msg the message
+   *
+   * \return a pending send
+   */
+  template <auto f>
+  messaging::PendingSend sendMsg(
+    VirtualElmProxyType<typename FunctionTraits<decltype(f)>::MsgT::CollectionType> const& proxy,
+    typename FunctionTraits<decltype(f)>::MsgT *msg
+  ) {
+    using MsgT = typename FunctionTraits<decltype(f)>::MsgT;
+    return sendMsg<MsgT, f>(proxy, msg);
+  }
+
   /**
    * \brief Send collection element a message from active member handler
    *
@@ -965,6 +999,26 @@ struct CollectionManager
     MsgT *msg, bool instrument = true
   );
 
+  /**
+   * \brief Broadcast a message with action function handler
+   *
+   * \param[in] proxy the collection proxy
+   * \param[in] msg the message
+   * \param[in] instrument whether to instrument the broadcast for load
+   * balancing (some system calls use this to disable instrumentation)
+   *
+   * \return a pending send
+   */
+  template < auto f>
+  messaging::PendingSend broadcastMsg(
+    CollectionProxyWrapType<typename FunctionTraits<decltype(f)>::MsgT::CollectionType> const& proxy,
+    typename FunctionTraits<decltype(f)>::MsgT *msg,
+    bool instrument = true
+  ) {
+    using MsgT = typename FunctionTraits<decltype(f)>::MsgT;
+    return broadcastMsg<MsgT, f>(proxy, msg, instrument);
+  }
+
   /**
    * \brief Broadcast a message with action member handler
    *

diff --git a/src/vt/vrt/collection/send/sendable.h b/src/vt/vrt/collection/send/sendable.h
@@ -121,6 +121,51 @@ struct Sendable : BaseProxyT {
     typename MsgT, ActiveColMemberTypedFnType<MsgT,ColT> f, typename... Args
   >
   messaging::PendingSend send(Args&&... args) const;
+
+
+  template <typename Return, typename... Args>
+  struct FunctionTraits;
+
+  template <typename Return, typename Msg>
+  struct FunctionTraits<Return(*)(Msg*, ColT*)> {
+    using MsgT = Msg;
+    using ReturnT = Return;
+  };
+
+  template <typename Return, typename Msg>
+  struct FunctionTraits<Return(ColT::*)(Msg*)> {
+    using MsgT = Msg;
+    using ReturnT = Return;
+  };
+
+  /**
+   * \brief Create message (with action function handler) and send to collection element
+   *
+   * \param[in] args arguments needed for creteating the message
+   *
+   * \return a pending send
+   */
+  template <auto f>
+  messaging::PendingSend sendMsg(
+    messaging::MsgPtrThief<typename FunctionTraits<decltype(f)>::MsgT> msg
+  ) const {
+    using MsgT = typename FunctionTraits<decltype(f)>::MsgT;
+    return sendMsg<MsgT, f>(msg);
+  }
+
+  /**
+   * \brief Create message (with action function handler) and send to collection element
+   *
+   * \param[in] args arguments needed for creteating the message
+   *
+   * \return a pending send
+   */
+  template <auto f, typename... Args>
+  messaging::PendingSend send(Args&&... args) const {
+    using MsgT = typename FunctionTraits<decltype(f)>::MsgT;
+    return send<MsgT, f>(std::forward<Args>(args)...);
+  }
+
 };
 
 }}} /* end namespace vt::vrt::collection */

diff --git a/tests/perf/collection_local_send.cc b/tests/perf/collection_local_send.cc
@@ -94,7 +94,7 @@ struct NodeObj {
   void perfRunBenchmark() {
     for (int i = 0; i < num_iters; i++) {
       auto m = makeMessage<TestCol::ColMsg>();
-      col_proxy[0].template sendMsg<TestCol::ColMsg, &TestCol::han>(m);
+      col_proxy[0].template sendMsg<&TestCol::han>(m);
       vt::theSched()->runSchedulerOnceImpl();
     }
   }
@@ -117,7 +117,7 @@ VT_PERF_TEST(MyTest, test_collection_local_send) {
   grp_proxy[my_node_].invoke<&NodeObj::initialize>();
 
   if (theContext()->getNode() == 0) {
-    grp_proxy[my_node_].send<MyMsg, &NodeObj::perfMakeRunnable>();
+    grp_proxy[my_node_].send<&NodeObj::perfMakeRunnable>();
   }
 }
 

diff --git a/tests/perf/collection_local_send_prealloc.cc b/tests/perf/collection_local_send_prealloc.cc
@@ -98,7 +98,7 @@ struct NodeObj {
   void perfRunBenchmark() {
     for (int i = 0; i < num_iters; i++) {
       auto m = msgs[i];
-      col_proxy[0].template sendMsg<TestCol::ColMsg, &TestCol::han>(m);
+      col_proxy[0].template sendMsg<&TestCol::han>(m);
       vt::theSched()->runSchedulerOnceImpl();
     }
   }

diff --git a/tests/unit/collection/test_broadcast.h b/tests/unit/collection/test_broadcast.h
@@ -131,13 +131,12 @@ void test_broadcast_1(std::string const& label) {
     auto proxy = theCollection()->construct<ColType>(range, label);
 
     proxy.template broadcast<
-      MsgType,
       BroadcastHandlers<ColType>::handler
     >(args);
 
     auto msg = makeMessage<MsgType>(args);
     theCollection()->broadcastMsg<
-      MsgType,BroadcastHandlers<ColType>::handler
+      BroadcastHandlers<ColType>::handler
     >(proxy, msg.get());
   }
 }

diff --git a/tests/unit/collection/test_send.h b/tests/unit/collection/test_send.h
@@ -195,9 +195,9 @@ void test_collection_send_1(std::string const& label) {
       auto msg = makeMessage<MsgType>(args);
       EXPECT_EQ(msg.size(), sizeof(MsgType));
       if (i % 2 == 0) {
-        proxy[i].template sendMsg<MsgType,SendHandlers<ColType>::handler>(msg.get());
+        proxy[i].template sendMsg<SendHandlers<ColType>::handler>(msg.get());
       } else {
-        theCollection()->sendMsg<MsgType,SendHandlers<ColType>::handler>(
+        theCollection()->sendMsg<SendHandlers<ColType>::handler>(
           proxy[i], msg.get()
         );
       }
@@ -222,7 +222,7 @@ void test_collection_send_sz_1() {
       EXPECT_EQ(msg.size(), sizeof(MsgType) + sizeof(PayloadType));
       msg->buff_size = sizeof(PayloadType);
       std::memcpy(reinterpret_cast<void *>(msg->payload()), &args, msg->buff_size);
-      proxy[i].template sendMsg<MsgType,SendSzHandlers<ColType>::handler>(msg);
+      proxy[i].template sendMsg<SendSzHandlers<ColType>::handler>(msg);
     }
   }
 }
@@ -242,9 +242,9 @@ void test_collection_send_ptm_1(std::string const& label) {
       auto msg = makeMessage<MsgType>(args);
       //proxy[i].template send<MsgType,SendHandlers<ColType>::handler>(msg);
       if (i % 2 == 0) {
-        proxy[i].template sendMsg<MsgType,&ColType::handler>(msg.get());
+        proxy[i].template sendMsg<&ColType::handler>(msg.get());
       } else {
-        theCollection()->sendMsg<MsgType,&ColType::handler>(
+        theCollection()->sendMsg<&ColType::handler>(
           proxy[i], msg.get()
         );
       }