ProcStats: refactor addProcStats to de-template, taking the common Mi…

…gratable base instead

src/vt/vrt/collection/balance/elm_stats.impl.h

@@ -94,7 +94,7 @@ template <typename ColT>
   auto const& idx = col->getIndex();
   auto const& elm_proxy = proxy[idx];
 
-  theProcStats()->addProcStats<ColT>(elm_proxy, col, cur_phase, total_load, subphase_loads, comm);
+  theProcStats()->addProcStats(col, cur_phase, total_load, subphase_loads, comm);
 
   auto const before_ready = theCollection()->numReadyCollections();
   theCollection()->makeCollectionReady(untyped_proxy);

src/vt/vrt/collection/balance/proc_stats.cc

@@ -252,4 +252,59 @@ void ProcStats::outputStatsFile() {
   closeStatsFile();
 }
 
+ElementIDType ProcStats::addProcStats(
+  Migratable* col_elm,
+  PhaseType const& phase, TimeType const& time,
+  std::vector<TimeType> const& subphase_time, CommMapType const& comm
+) {
+  // A new temp ID gets assigned when a object is migrated into a node
+
+  auto const temp_id = col_elm->temp_elm_id_;
+  auto const perm_id = col_elm->stats_elm_id_;
+
+  debug_print(
+    lb, node,
+    "ProcStats::addProcStats: temp_id={}, perm_id={}, phase={}, subphases={}, load={}\n",
+    temp_id, perm_id, phase, subphase_time.size(), time
+  );
+
+  proc_data_.resize(phase + 1);
+  auto elm_iter = proc_data_.at(phase).find(temp_id);
+  vtAssert(elm_iter == proc_data_.at(phase).end(), "Must not exist");
+  proc_data_.at(phase).emplace(
+    std::piecewise_construct,
+    std::forward_as_tuple(temp_id),
+    std::forward_as_tuple(time)
+  );
+
+  proc_subphase_data_.resize(phase + 1);
+  auto elm_subphase_iter = proc_subphase_data_.at(phase).find(temp_id);
+  vtAssert(elm_subphase_iter == proc_subphase_data_.at(phase).end(), "Must not exist");
+  proc_subphase_data_.at(phase).emplace(
+    std::piecewise_construct,
+    std::forward_as_tuple(temp_id),
+    std::forward_as_tuple(subphase_time)
+  );
+
+  proc_comm_.resize(phase + 1);
+  for (auto&& c : comm) {
+    proc_comm_.at(phase)[c.first] += c.second;
+  }
+
+  proc_temp_to_perm_[temp_id] = perm_id;
+  proc_perm_to_temp_[perm_id] = temp_id;
+
+  auto migrate_iter = proc_migrate_.find(temp_id);
+  if (migrate_iter == proc_migrate_.end()) {
+    proc_migrate_.emplace(
+      std::piecewise_construct,
+      std::forward_as_tuple(temp_id),
+      std::forward_as_tuple([col_elm](NodeType node){
+        col_elm->migrate(node);
+      })
+    );
+  }
+  return temp_id;
+}
+
 }}}} /* end namespace vt::vrt::collection::balance */

src/vt/vrt/collection/balance/proc_stats.h

@@ -50,6 +50,7 @@
 #include "vt/vrt/collection/balance/lb_comm.h"
 #include "vt/vrt/collection/balance/phase_msg.h"
 #include "vt/vrt/collection/balance/stats_msg.h"
+#include "vt/vrt/collection/types/migratable.h"
 #include "vt/runtime/component/component_pack.h"
 #include "vt/timing/timing.h"
 #include "vt/objgroup/proxy/proxy_objgroup.h"
@@ -88,17 +89,15 @@ struct ProcStats : runtime::component::Component<ProcStats> {
   /**
    * \internal \brief Add processor statistics for local object
    *
-   * \param[in] elm_proxy the element proxy to the object
    * \param[in] col_elm the collection element pointer
    * \param[in] phase the current phase
    * \param[in] time the time the object took
    * \param[in] comm the comm graph for the object
    *
    * \return the temporary ID for the object assigned for this phase
    */
-  template <typename ColT>
   ElementIDType addProcStats(
-    VirtualElmProxyType<ColT> const& elm_proxy, ColT* col_elm,
+    Migratable* col_elm,
     PhaseType const& phase, TimeType const& time,
     std::vector<TimeType> const& subphase_time, CommMapType const& comm
   );
@@ -226,6 +225,4 @@ extern vrt::collection::balance::ProcStats* theProcStats();
 
 } /* end namespace vt */
 
-#include "vt/vrt/collection/balance/proc_stats.impl.h"
-
 #endif /*INCLUDED_VRT_COLLECTION_BALANCE_PROC_STATS_H*/

src/vt/vrt/collection/types/migratable.h

@@ -52,10 +52,14 @@
 #include "vt/vrt/collection/types/migratable.fwd.h"
 #include "vt/vrt/collection/balance/lb_common.h"
 #include "vt/vrt/collection/balance/elm_stats.h"
-#include "vt/vrt/collection/balance/proc_stats.h"
 
 namespace vt { namespace vrt { namespace collection {
 
+// Forward declaration for friend declaration below
+namespace balance {
+  struct ProcStats;
+}
+
 struct Migratable : MigrateHookBase {
 
   Migratable();