#1490: vrt_coll: add some docs, cover more complex case with elements…

… off-home
DARMA-tasking · Jul 28, 2021 · 7de48d2 · 7de48d2
1 parent e9a045b
commit 7de48d2
Show file tree

Hide file tree

Showing 2 changed files with 36 additions and 1 deletion.
diff --git a/src/vt/vrt/collection/manager.h b/src/vt/vrt/collection/manager.h
@@ -1871,6 +1871,7 @@ struct CollectionManager
     typename ColT::IndexType range, std::string const& file_base
   );
 
+private:
   template <typename ColT, typename IdxT = typename ColT::IndexType>
   struct RestoreMigrateMsg : vt::Message {
     RestoreMigrateMsg() = default;
@@ -1885,9 +1886,27 @@ struct CollectionManager
     CollectionProxyWrapType<ColT> proxy_;
   };
 
+  template <typename ColT, typename IdxT = typename ColT::IndexType>
+  struct RestoreMigrateColMsg : RestoreMigrateMsg<ColT, IdxT>, CollectionMessage<ColT> {
+    RestoreMigrateColMsg() = default;
+    RestoreMigrateColMsg(NodeType in_to_node, IdxT in_idx, CollectionProxyWrapType<ColT> in_proxy)
+      : RestoreMigrateMsg<ColT>(in_to_node, in_idx, in_proxy)
+    { }
+  };
+
   template <typename ColT>
   static void restoreHandler(RestoreMigrateMsg<ColT>* msg);
 
+public:
+  /**
+   * \brief Restore the collection (collective) from file on top of an existing
+   * collection. Migrates collection elements to the rank saved from the
+   * checkpoint.
+   *
+   * \param[in] proxy the collection proxy
+   * \param[in] range the range of the collection to restore
+   * \param[in] file_base the base file name for the files to read
+   */
   template <typename ColT>
   void restoreFromFileInPlace(
     CollectionProxyWrapType<ColT> proxy, typename ColT::IndexType range,

diff --git a/src/vt/vrt/collection/manager.impl.h b/src/vt/vrt/collection/manager.impl.h
@@ -3160,13 +3160,29 @@ void CollectionManager::checkpointToFile(
   checkpoint::serializeToFile(directory, directory_name);
 }
 
+namespace detail {
 template <typename ColT>
-/*static*/ void CollectionManager::restoreHandler(RestoreMigrateMsg<ColT>* msg) {
+inline void restoreOffHomeElement(CollectionManager::RestoreMigrateColMsg<ColT>* msg, ColT*) {
   auto idx = msg->idx_;
   auto node = msg->to_node_;
   auto proxy = msg->proxy_;
   theCollection()->migrate(proxy(idx), node);
 }
+} /* end namespace detail */
+
+template <typename ColT>
+/*static*/ void CollectionManager::restoreHandler(RestoreMigrateMsg<ColT>* msg) {
+  auto idx = msg->idx_;
+  auto node = msg->to_node_;
+  auto proxy = msg->proxy_;
+  if (proxy(idx).tryGetLocalPtr() != nullptr) {
+    theCollection()->migrate(proxy(idx), node);
+  } else {
+    proxy(idx).template send<
+      RestoreMigrateColMsg<ColT>, detail::restoreOffHomeElement<ColT>
+    >(node, idx, proxy);
+  }
+}
 
 template <typename ColT>
 void CollectionManager::restoreFromFileInPlace(