#1983: location: refactor to reduce PendingSend allocation

src/vt/topos/location/location.h

@@ -224,6 +224,12 @@ struct EntityLocationCoord : LocationCoord {
     EntityID const& id, NodeType const& home_node, NodeActionType const& action
   );
 
+  template <typename MessageT, ActiveTypedFnType<MessageT> *f>
+  void setupMessageForRouting(
+    EntityID const& id, NodeType const& home_node,
+    MsgSharedPtr<MessageT> const& msg
+  );
+
   /**
    * \brief Route a message with a custom handler
    *
@@ -237,6 +243,22 @@ struct EntityLocationCoord : LocationCoord {
     MsgSharedPtr<MessageT> const& msg
   );
 
+  /**
+   * \brief Route a message with a custom handler
+   *
+   * \param[in] m message shared pointer
+   */
+  template <typename MessageT>
+  void routePreparedMsgHandler(MsgSharedPtr<MessageT> const& msg);
+
+  /**
+   * \brief Route a message with a custom handler
+   *
+   * \param[in] m message shared pointer
+   */
+  template <typename MessageT>
+  void routePreparedMsg(MsgSharedPtr<MessageT> const& msg);
+
   /**
    * \brief Route a message with a custom handler where the element is local
    *

src/vt/topos/location/location.impl.h

@@ -535,7 +535,7 @@ void EntityLocationCoord<EntityID>::routeMsgNode(
     // set the instance on the message to deliver to the correct manager
     msg->setLocInst(this_inst);
 
-    auto m = msg;
+    auto m = msg; //copy for msg thief
     // send to the node discovered by the location manager
     theMsg()->sendMsg<MessageT, &EntityLocationCoord<EntityID>::routedHandler>(to_node, m);
   } else {
@@ -656,6 +656,35 @@ template <typename MessageT, ActiveTypedFnType<MessageT> *f>
 void EntityLocationCoord<EntityID>::routeMsgHandler(
   EntityID const& id, NodeType const& home_node,
   MsgSharedPtr<MessageT> const& msg
+) {
+  setupMessageForRouting<MessageT, f>(id, home_node, msg);
+
+  routePreparedMsgHandler(msg);
+
+  if (local_registered_.find(id) == local_registered_.end()) {
+    return routeMsg<MessageT>(id,home_node,msg);
+  } else {
+    return routeMsgHandlerLocal(msg);
+  }
+}
+
+template <typename EntityID>
+template <typename MessageT>
+void EntityLocationCoord<EntityID>::routePreparedMsgHandler(
+  MsgSharedPtr<MessageT> const& msg
+) {
+  if (local_registered_.find(msg->getEntity()) == local_registered_.end()) {
+    return routePreparedMsg(msg);
+  } else {
+    return routeMsgHandlerLocal(msg);
+  }
+}
+
+template <typename EntityID>
+template <typename MessageT, ActiveTypedFnType<MessageT> *f>
+void EntityLocationCoord<EntityID>::setupMessageForRouting(
+  EntityID const& id, NodeType const& home_node,
+  MsgSharedPtr<MessageT> const& msg
 ) {
   using auto_registry::HandlerManagerType;
 
@@ -668,12 +697,10 @@ void EntityLocationCoord<EntityID>::routeMsgHandler(
 # endif
 
   msg->setHandler(handler);
-
-  if (local_registered_.find(id) == local_registered_.end()) {
-    return routeMsg<MessageT>(id,home_node,msg);
-  } else {
-    return routeMsgHandlerLocal(msg);
-  }
+  msg->setEntity(id);
+  msg->setHomeNode(home_node);
+  msg->setLocFromNode(theContext()->getNode());
+  msg->setLocInst(this_inst);
 }
 
 template <typename EntityID>
@@ -688,50 +715,58 @@ void EntityLocationCoord<EntityID>::routeMsgHandlerLocal(
 
 template <typename EntityID>
 template <typename MessageT>
-void EntityLocationCoord<EntityID>::routeMsg(
-  EntityID const& id, NodeType const& home_node,
-  MsgSharedPtr<MessageT> const& msg, NodeType from_node
+void EntityLocationCoord<EntityID>::routePreparedMsg(
+  MsgSharedPtr<MessageT> const& msg
 ) {
-  auto const from =
-    from_node == uninitialized_destination ? theContext()->getNode() :
-    from_node;
-
-  // set field for location routed message
-  msg->setEntity(id);
-  msg->setHomeNode(home_node);
-  msg->setLocFromNode(from);
-
   auto const msg_size = sizeof(*msg);
   bool const use_eager = useEagerProtocol(msg);
   auto const epoch = theMsg()->getEpochContextMsg(msg);
 
   vt_debug_print(
     verbose, location,
     "routeMsg: inst={}, home={}, msg_size={}, is_large_msg={}, eager={}, "
-    "in_from={}, from={}, msg{}, msg from={}, epoch={:x}\n",
-    this_inst, home_node, msg_size, msg_size > small_msg_max_size, use_eager,
-    from_node, from, print_ptr(msg.get()), msg->getLocFromNode(),
+    "msg{}, msg from={}, epoch={:x}\n",
+    this_inst, msg->getHomeNode(), msg_size, msg_size > small_msg_max_size, use_eager,
+    print_ptr(msg.get()), msg->getLocFromNode(),
     epoch
   );
 
-  msg->setLocInst(this_inst);
-
   if (use_eager) {
     theMsg()->pushEpoch(epoch);
-    routeMsgEager<MessageT>(id, home_node, msg);
+    routeMsgEager<MessageT>(msg->getEntity(), msg->getHomeNode(), msg);
     theMsg()->popEpoch(epoch);
   } else {
     theTerm()->produce(epoch);
     // non-eager protocol: get location first then send message after resolution
-    getLocation(id, home_node, [=](NodeType node) {
+    getLocation(msg->getEntity(), msg->getHomeNode(), [=](NodeType node) {
       theMsg()->pushEpoch(epoch);
-      routeMsgNode<MessageT>(id, home_node, node, msg);
+      routeMsgNode<MessageT>(
+        msg->getEntity(), msg->getHomeNode(), node, msg
+      );
       theMsg()->popEpoch(epoch);
       theTerm()->consume(epoch);
     });
   }
 }
 
+template <typename EntityID>
+template <typename MessageT>
+void EntityLocationCoord<EntityID>::routeMsg(
+  EntityID const& id, NodeType const& home_node,
+  MsgSharedPtr<MessageT> const& msg, NodeType from_node
+) {
+  auto const from =
+    from_node == uninitialized_destination ? theContext()->getNode() :
+    from_node;
+
+  // set field for location routed message
+  msg->setEntity(id);
+  msg->setHomeNode(home_node);
+  msg->setLocFromNode(from);
+
+  routePreparedMsg(msg);
+}
+
 template <typename EntityID>
 void EntityLocationCoord<EntityID>::updatePendingRequest(
   LocEventID const& event_id, EntityID const& id,

src/vt/vrt/collection/manager.impl.h

@@ -1187,6 +1187,7 @@ messaging::PendingSend CollectionManager::sendMsgUntypedHandler(
   msg->setFromNode(theContext()->getNode());
   msg->setVrtHandler(handler);
   msg->setProxy(toProxy);
+  theMsg()->markAsCollectionMessage(msg);
 
   auto idx = elm_proxy.getIndex();
   vt_debug_print(
@@ -1196,18 +1197,19 @@ messaging::PendingSend CollectionManager::sendMsgUntypedHandler(
     col_proxy, cur_epoch, idx, handler, imm_context
   );
 
+  auto home_node = theCollection()->getMappedNode<ColT>(col_proxy, idx);
+  // route the message to the destination using the location manager
+  auto lm = theLocMan()->getCollectionLM<IdxT>(col_proxy);
+  vtAssert(lm != nullptr, "LM must exist");
+  lm->template setupMessageForRouting<
+    MsgT, collectionMsgTypedHandler<ColT,IdxT,MsgT>
+  >(idx, home_node, msg);
+
   return messaging::PendingSend{
-    msg, [=](MsgSharedPtr<BaseMsgType>& inner_msg){
-      theMsg()->pushEpoch(cur_epoch);
-      auto home_node = theCollection()->getMappedNode<ColT>(col_proxy, idx);
-      // route the message to the destination using the location manager
-      auto lm = theLocMan()->getCollectionLM<IdxT>(col_proxy);
-      vtAssert(lm != nullptr, "LM must exist");
-      theMsg()->markAsCollectionMessage(msg);
-      lm->template routeMsgHandler<
-        MsgT, collectionMsgTypedHandler<ColT,IdxT,MsgT>
-      >(idx, home_node, msg);
-      theMsg()->popEpoch(cur_epoch);
+    msg, [](MsgSharedPtr<BaseMsgType>& inner_msg){
+      auto typed_msg = inner_msg.template to<MsgT>();
+      auto lm2 = theLocMan()->getCollectionLM<IdxT>(typed_msg->getLocInst());
+      lm2->template routePreparedMsgHandler<MsgT>(typed_msg);
     }
   };
 }