#2389: vrt_coll: implement traced invoke calls on collection elements

Author: lifflander

src/vt/context/runnable_context/trace.cc

@@ -65,6 +65,22 @@ Trace::Trace(
     handler_(in_handler)
 { }
 
+Trace::Trace(
+  trace::TraceEventIDType event, HandlerType const in_handler,
+  NodeType const in_from_node, std::size_t msg_size,
+  uint64_t in_idx1, uint64_t in_idx2, uint64_t in_idx3, uint64_t in_idx4
+) : is_collection_(false),
+    event_(event),
+    msg_size_(msg_size),
+    is_traced_(HandlerManager::isHandlerTrace(in_handler)),
+    from_node_(in_from_node),
+    handler_(in_handler),
+    idx1_(in_idx1),
+    idx2_(in_idx2),
+    idx3_(in_idx3),
+    idx4_(in_idx4)
+{ }
+
 void Trace::start(TimeType time) {
   if (not is_traced_) {
     return;

src/vt/context/runnable_context/trace.h

@@ -75,6 +75,26 @@ struct Trace {
     NodeType const in_from_node, std::size_t msg_size
   );
 
+  /**
+   * \brief Construct a new trace context (collection processing) without a
+   * message directly
+   *
+   * \param[in] event the trace event
+   * \param[in] in_handler the handler
+   * \param[in] in_from_node from node
+   * \param[in] msg_size size of message/payload (zero of not serializable
+   * \param[in] in_idx1 1-dimension index
+   * \param[in] in_idx2 2-dimension index
+   * \param[in] in_idx3 3-dimension index
+   * \param[in] in_idx4 4-dimension index
+   */
+  Trace(
+    trace::TraceEventIDType event, HandlerType const in_handler,
+    NodeType const in_from_node, std::size_t msg_size,
+    uint64_t in_idx1, uint64_t in_idx2, uint64_t in_idx3, uint64_t in_idx4
+  );
+
+
   /**
    * \brief Construct a new trace context (basic processing event)
    *

src/vt/runnable/make_runnable.h

@@ -399,6 +399,44 @@ inline RunnableMaker<BaseMsgType> makeRunnableVoidTraced(
   return RunnableMaker<BaseMsgType>{r, nullptr, handler, from};
 }
 
+/**
+ * \brief Make a new runnable without a message (void handler) with tracing for
+ * collections
+ *
+ * \param[in] is_threaded whether it is threaded
+ * \param[in] handler the handler bits
+ * \param[in] from the node that caused this runnable to execute
+ * \param[in] idx1 1-dimension index
+ * \param[in] idx2 2-dimension index
+ * \param[in] idx3 3-dimension index
+ * \param[in] idx4 4-dimension index
+ *
+ * \return the maker for further customization
+ */
+inline RunnableMaker<BaseMsgType> makeRunnableVoidTraced(
+  bool is_threaded, HandlerType handler, NodeType from,
+  [[maybe_unused]] trace::TraceEventIDType trace_event,
+  [[maybe_unused]] std::size_t msg_size,
+  uint64_t idx1, uint64_t idx2, uint64_t idx3, uint64_t idx4
+) {
+  // These are currently only types of registry entries that can be void
+  auto r = new RunnableNew(is_threaded);
+  r->addContextSetContext(r, from);
+
+#if vt_check_enabled(trace_enabled)
+  auto const han_type = HandlerManager::getHandlerRegistryType(handler);
+  if (han_type == auto_registry::RegistryTypeEnum::RegVrtCollection or
+      han_type == auto_registry::RegistryTypeEnum::RegVrtCollectionMember) {
+    r->addContextTrace(
+      trace_event, handler, from, msg_size, idx1, idx2, idx3, idx4
+    );
+  }
+#endif
+
+  return RunnableMaker<BaseMsgType>{r, nullptr, handler, from};
+}
+
+
 }} /* end namespace vt::runnable */
 
 #include "vt/runnable/make_runnable.impl.h"

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

@@ -465,10 +465,33 @@ auto CollectionManager::invoke(
 
   auto const this_node = theContext()->getNode();
 
+#if vt_check_enabled(trace_enabled)
+  auto const han = auto_registry::makeAutoHandlerCollectionMemParam<
+     ColT, decltype(f), f, void
+  >();
+  auto const trace_event = theMsg()->makeTraceCreationSend(han, 0, false);
+
+#if vt_check_enabled(trace_enabled)
+  auto idx = ptr->getIndex();
+  uint64_t const idx1 = idx.ndims() > 0 ? idx[0] : 0;
+  uint64_t const idx2 = idx.ndims() > 1 ? idx[1] : 0;
+  uint64_t const idx3 = idx.ndims() > 2 ? idx[2] : 0;
+  uint64_t const idx4 = idx.ndims() > 3 ? idx[3] : 0;
+#endif
+
+  return runnable::makeRunnableVoidTraced(
+    false, han, this_node, trace_event, 0, idx1, idx2, idx3, idx4
+  )
+    .withCollection(ptr)
+    .withLBDataVoidMsg(ptr)
+    .runLambda(f, ptr, std::forward<Args>(args)...);
+
+#else
   return runnable::makeRunnableVoid(false, uninitialized_handler, this_node)
     .withCollection(ptr)
     .withLBDataVoidMsg(ptr)
     .runLambda(f, ptr, std::forward<Args>(args)...);
+#endif
 }
 
 template <