1051: Add local invoke to proxies that trace and record LB

src/CMakeLists.txt

@@ -76,6 +76,7 @@ set(
       vrt/collection/dispatch vrt/collection/gettable vrt/collection/rdmaable
       vrt/collection/staged_token
       vrt/collection/listener
+      vrt/collection/invoke
       vrt/collection/balance
         vrt/collection/balance/baselb
         vrt/collection/balance/hierarchicallb

src/vt/objgroup/manager.fwd.h

@@ -57,6 +57,8 @@ void dispatchObjGroup(MsgSharedPtr<ShortMessage> msg, HandlerType han);
 template <typename MsgT>
 void send(MsgSharedPtr<MsgT> msg, HandlerType han, NodeType node);
 template <typename MsgT>
+void invoke(messaging::MsgPtrThief<MsgT> msg, HandlerType han, NodeType node);
+template <typename MsgT>
 void broadcast(MsgSharedPtr<MsgT> msg, HandlerType han);
 void scheduleMsg(
   MsgSharedPtr<ShortMessage> msg, HandlerType han, EpochType epoch

src/vt/objgroup/manager.h

@@ -215,6 +215,26 @@ struct ObjGroupManager : runtime::component::Component<ObjGroupManager> {
   template <typename ObjT, typename MsgT, ActiveObjType<MsgT, ObjT> fn>
   void send(ProxyElmType<ObjT> proxy, MsgSharedPtr<MsgT> msg);
 
+  /**
+   * \internal \brief Invoke message handler on an element of the object group
+   * The message handler will be invoked inline without going through scheduler
+   *
+   * \param[in] proxy proxy to the object group
+   * \param[in] msg message
+   */
+  template <typename ObjT, typename MsgT, ActiveObjType<MsgT, ObjT> fn>
+  void invoke(ProxyElmType<ObjT> proxy, messaging::MsgPtrThief<MsgT> msg);
+
+  /**
+   * \internal \brief Invoke function 'f' on an element of the object group
+   * The function will be invoked inline without going through scheduler
+   *
+   * \param[in] proxy proxy to the object group
+   * \param[in] args function arguments
+   */
+  template <typename ObjT, typename Type, Type f, typename... Args>
+  decltype(auto) invoke(ProxyElmType<ObjT> proxy, Args&&... args);
+
   /**
    * \internal \brief Broadcast a message to all nodes in object group
    *
@@ -316,6 +336,17 @@ struct ObjGroupManager : runtime::component::Component<ObjGroupManager> {
   template <typename MsgT>
   void send(MsgSharedPtr<MsgT> msg, HandlerType han, NodeType node);
 
+  /**
+   * \internal \brief Invoke a message handler on an objgroup
+   * The message handler will be invoked inline without going through scheduler
+   *
+   * \param[in] msg message
+   * \param[in] han handler to invoke
+   * \param[in] node node to invoke the handler on
+   */
+  template <typename MsgT>
+  void invoke(messaging::MsgPtrThief<MsgT> msg, HandlerType han, NodeType node);
+
   /**
    * \internal \brief Broadcast message to an objgroup
    *

src/vt/objgroup/manager.impl.h

@@ -58,6 +58,7 @@
 #include "vt/registry/auto/auto_registry.h"
 #include "vt/collective/collective_alg.h"
 #include "vt/messaging/active.h"
+#include "vt/runnable/invoke.h"
 
 #include <memory>
 
@@ -262,17 +263,55 @@ void ObjGroupManager::send(ProxyElmType<ObjT> proxy, MsgSharedPtr<MsgT> msg) {
   send<MsgT>(msg,han,dest_node);
 }
 
+template <typename ObjT, typename MsgT, ActiveObjType<MsgT, ObjT> fn>
+void ObjGroupManager::invoke(
+  ProxyElmType<ObjT> proxy, messaging::MsgPtrThief<MsgT> msg
+) {
+  auto const proxy_bits = proxy.getProxy();
+  auto const dest_node = proxy.getNode();
+  auto const ctrl = proxy::ObjGroupProxy::getID(proxy_bits);
+  auto const han = auto_registry::makeAutoHandlerObjGroup<ObjT, MsgT, fn>(ctrl);
+
+  vt_debug_print(
+    objgroup, node,
+    "ObjGroupManager::invoke: proxy={:x}, node={}, ctrl={:x}, han={:x}\n",
+    proxy_bits, dest_node, ctrl, han
+  );
+
+  invoke<MsgT>(msg, han, dest_node);
+}
+
+template <typename ObjT, typename Type, Type f, typename... Args>
+decltype(auto)
+ObjGroupManager::invoke(ProxyElmType<ObjT> proxy, Args&&... args) {
+  auto const dest_node = proxy.getNode();
+  auto const this_node = theContext()->getNode();
+
+  vtAssert(
+    dest_node == this_node,
+    fmt::format(
+      "Attempting to invoke handler on node:{} instead of node:{}!\n", this_node,
+      dest_node
+    )
+  );
+
+  return runnable::invoke<Type, f>(get(proxy), std::forward<Args>(args)...);
+}
+
+
 template <typename ObjT, typename MsgT, ActiveObjType<MsgT, ObjT> fn>
 void ObjGroupManager::broadcast(ProxyType<ObjT> proxy, MsgSharedPtr<MsgT> msg) {
   auto const proxy_bits = proxy.getProxy();
   auto const ctrl = proxy::ObjGroupProxy::getID(proxy_bits);
   auto const han = auto_registry::makeAutoHandlerObjGroup<ObjT,MsgT,fn>(ctrl);
+
   vt_debug_print(
     objgroup, node,
     "ObjGroupManager::broadcast: proxy={:x}, ctrl={:x}, han={:x}\n",
     proxy_bits, ctrl, han
   );
-  broadcast<MsgT>(msg,han);
+
+  broadcast<MsgT>(msg, han);
 }
 
 template <typename MsgT>
@@ -282,6 +321,13 @@ void ObjGroupManager::send(
   return objgroup::send(msg,han,dest_node);
 }
 
+template <typename MsgT>
+void ObjGroupManager::invoke(
+  messaging::MsgPtrThief<MsgT> msg, HandlerType han, NodeType dest_node
+) {
+  objgroup::invoke(msg, han, dest_node);
+}
+
 template <typename MsgT>
 void ObjGroupManager::broadcast(MsgSharedPtr<MsgT> msg, HandlerType han) {
   return objgroup::broadcast(msg,han);

src/vt/objgroup/manager.static.h

@@ -66,6 +66,21 @@ void send(MsgSharedPtr<MsgT> msg, HandlerType han, NodeType dest_node) {
   }
 }
 
+template <typename MsgT>
+void invoke(messaging::MsgPtrThief<MsgT> msg, HandlerType han, NodeType dest_node) {
+  auto const this_node = theContext()->getNode();
+
+  vtAssert(
+    dest_node == this_node,
+    fmt::format(
+      "Attempting to invoke handler on node:{} instead of node:{}!", this_node,
+      dest_node
+    )
+  );
+
+  runnable::Runnable<MsgT>::run(han, nullptr, msg.msg_.get(), this_node);
+}
+
 template <typename MsgT>
 void broadcast(MsgSharedPtr<MsgT> msg, HandlerType han) {
   // Get the current epoch for the message

src/vt/objgroup/proxy/proxy_objgroup_elm.h

@@ -127,6 +127,24 @@ struct ProxyElm {
   template <typename MsgT, ActiveObjType<MsgT, ObjT> fn, typename... Args>
   void send(Args&&... args) const;
 
+  /**
+   * \brief Invoke locally a message handler on the node/element indexed by this proxy.
+   * The message handler will be invoked inline without going through scheduler
+   *
+   * \param[in] args args to pass to the message constructor
+   */
+  template <typename MsgT, ActiveObjType<MsgT, ObjT> fn, typename... Args>
+  void invoke(Args&&... args) const;
+
+  /**
+   * \brief Invoke locally a function 'f' on the node/element indexed by this proxy.
+   * The function will be invoked inline without going through scheduler
+   *
+   * \param[in] args function arguments
+   */
+  template <typename Type, Type f, typename... Args>
+  decltype(auto) invoke(Args&&... args) const;
+
   /**
    * \brief Update the local object instance pointer on this node. Must be run
    * on local node.

src/vt/objgroup/proxy/proxy_objgroup_elm.impl.h

@@ -78,6 +78,26 @@ void ProxyElm<ObjT>::send(Args&&... args) const {
   return sendMsg<MsgT,fn>(makeMessage<MsgT>(std::forward<Args>(args)...));
 }
 
+template <typename ObjT>
+template <typename MsgT, ActiveObjType<MsgT, ObjT> fn, typename... Args>
+void ProxyElm<ObjT>::invoke(Args&&... args) const {
+  auto proxy = ProxyElm<ObjT>(*this);
+  theObjGroup()->invoke<ObjT, MsgT, fn>(
+    proxy, makeMessage<MsgT>(std::forward<Args>(args)...)
+  );
+}
+
+template <typename ObjT>
+template <typename Type, Type f, typename... Args>
+decltype(auto) ProxyElm<ObjT>::invoke(
+  Args&&... args
+) const
+{
+  auto proxy = ProxyElm<ObjT>(*this);
+  return theObjGroup()->invoke<ObjT, Type, f>(
+    proxy, std::forward<Args>(args)...);
+}
+
 template <typename ObjT>
 template <typename SerializerT>
 void ProxyElm<ObjT>::serialize(SerializerT& s) {