#2102: collection: finish the final fix for this bug

DARMA-tasking · Mar 15, 2023 · b4d56b1 · b4d56b1
1 parent 656a50e
commit b4d56b1
Show file tree

Hide file tree

Showing 4 changed files with 45 additions and 7 deletions.
diff --git a/src/vt/registry/auto/auto_registry_common.h b/src/vt/registry/auto/auto_registry_common.h
@@ -58,6 +58,13 @@
 #include <cstdlib>
 #include <functional>
 
+namespace vt::vrt::collection {
+
+template <typename ColT, typename UserMsgT, typename BaseMsgT>
+struct ColMsgWrap;
+
+} /* end namespace vt::vrt::collection */
+
 namespace vt { namespace auto_registry {
 
 struct SentinelObject {};
@@ -74,6 +81,18 @@ struct HandlersDispatcher final : BaseHandlersDispatcher {
 
   explicit HandlersDispatcher(HandlerT in_fn_ptr) : fp(in_fn_ptr) { }
 
+  template <typename U>
+  using ColMsgTrait = typename U::IsCollectionMessage;
+  template <typename U>
+  using IsColMsgTrait =
+    detection::is_detected_convertible<std::true_type, ColMsgTrait, U>;
+
+  template <typename U>
+  using ColTrait = typename U::IsCollectionType;
+  template <typename U>
+  using IsColTrait =
+    detection::is_detected_convertible<std::true_type, ColTrait, U>;
+
 public:
   void dispatch(messaging::BaseMsg* base_msg, void* object) const override {
     using T = HandlerT;
@@ -86,9 +105,23 @@ struct HandlersDispatcher final : BaseHandlersDispatcher {
     } else if constexpr (std::is_same_v<T, ActiveTypedFnType<MsgT>*>) {
       fp(msg);
     } else if constexpr (std::is_same_v<T, ColTypedFnType*>) {
-      fp(elm, msg);
+      if constexpr (IsColMsgTrait<MsgT>::value or not IsColTrait<ObjT>::value) {
+        fp(elm, msg);
+      } else {
+        auto wrap_msg = static_cast<
+          vrt::collection::ColMsgWrap<ObjT, MsgT, vt::Message>*
+        >(base_msg);
+        fp(elm, &wrap_msg->getMsg());
+      }
     } else if constexpr (std::is_same_v<T, ColMemberTypedFnType>) {
-      (elm->*fp)(msg);
+      if constexpr (IsColMsgTrait<MsgT>::value or not IsColTrait<ObjT>::value) {
+        (elm->*fp)(msg);
+      } else {
+        auto wrap_msg = static_cast<
+          vrt::collection::ColMsgWrap<ObjT, MsgT, vt::Message>*
+        >(base_msg);
+        (elm->*fp)(&wrap_msg->getMsg());
+      }
     } else if constexpr (std::is_same_v<ObjT, SentinelObject>) {
       std::apply(fp, msg->getTuple());
     } else {

diff --git a/src/vt/vrt/collection/manager.impl.h b/src/vt/vrt/collection/manager.impl.h
@@ -210,8 +210,11 @@ CollectionManager::collectionAutoMsgDeliver(
   trace::TraceEventIDType event, bool immediate
 ) {
   // Reference because it's a inner message that should *never* be deallocated
-  messageRef(&msg->getMsg());
-  MsgSharedPtr<UserMsgT> user_msg{&msg->getMsg()};
+  // messageRef(&msg->getMsg());
+  // messageRef(&msg->getMsg());
+  // MsgSharedPtr<UserMsgT> user_msg{&msg->getMsg()};
+
+  auto m = promoteMsg(msg);
 
   // Expand out the index for tracing purposes; Projections takes up to
   // 4-dimensions
@@ -223,7 +226,7 @@ CollectionManager::collectionAutoMsgDeliver(
   uint64_t const idx4 = idx.ndims() > 3 ? idx[3] : 0;
 #endif
 
-  runnable::makeRunnable(user_msg, true, han, from)
+  runnable::makeRunnable(m, true, han, from)
     .withTDEpoch(theMsg()->getEpochContextMsg(msg))
     .withCollection(base)
 #if vt_check_enabled(trace_enabled)

diff --git a/src/vt/vrt/collection/types/base.h b/src/vt/vrt/collection/types/base.h
@@ -58,6 +58,7 @@ namespace vt { namespace vrt { namespace collection {
 
 template <typename ColT, typename IndexT>
 struct CollectionBase : Indexable<IndexT> {
+  using IsCollectionType = std::true_type;
   using ProxyType = VirtualElmProxyType<ColT, IndexT>;
   using CollectionProxyType = CollectionProxy<ColT, IndexT>;
 

diff --git a/tests/unit/collection/test_promote.cc b/tests/unit/collection/test_promote.cc
@@ -83,8 +83,9 @@ TEST_F(TestCollectionPromoteMsg, test_collection_promote_1) {
   auto proxy = theCollection()->constructCollective<Hello>(
     num_elems, "test_collection_promote_1"
   );
-  proxy.broadcast<TestMsg,&Hello::doWork>("hello there");
-
+  if (this_node == 0) {
+    proxy.broadcast<TestMsg,&Hello::doWork>("hello there");
+  }
 }
 
 }}}} // end namespace vt::tests::unit::invoke