From 64a023ec203925c626d844b20739d73b8c2d09da Mon Sep 17 00:00:00 2001
From: jdolence <jcdolence@gmail.com>
Date: Thu, 2 Nov 2023 17:24:20 -0600
Subject: [PATCH] add more things

---
 src/bvals/comms/boundary_communication.cpp | 12 ++++++++++++
 src/mesh/mesh.cpp                          |  8 ++++++++
 2 files changed, 20 insertions(+)

diff --git a/src/bvals/comms/boundary_communication.cpp b/src/bvals/comms/boundary_communication.cpp
index fc9ab8dfdf6e..15577b353055 100644
--- a/src/bvals/comms/boundary_communication.cpp
+++ b/src/bvals/comms/boundary_communication.cpp
@@ -51,6 +51,14 @@ TaskStatus SendBoundBufs(std::shared_ptr<MeshData<Real>> &md) {
 
   Mesh *pmesh = md->GetMeshPointer();
   auto &block_cost = pmesh->GetBlockCost();
+  int min_lid = 9999999;
+  int max_lid = 0;
+  for (int b = 0; b < md->NumBlocks(); b++) {
+    int lid = md->GetBlockData(b)->GetBlockPointer()->lid;
+    min_lid = lid < min_lid ? lid : min_lid;
+    max_lid = lid > max_lid ? lid : max_lid;
+  }
+  BlockTimerHost host_timer(pmesh->block_cost_host, min_lid, max_lid);
   auto &cache = md->GetBvarsCache().GetSubCache(bound_type, true);
 
   if (cache.buf_vec.size() == 0)
@@ -83,6 +91,8 @@ TaskStatus SendBoundBufs(std::shared_ptr<MeshData<Real>> &md) {
   auto &sending_nonzero_flags = cache.sending_non_zero_flags;
   auto &sending_nonzero_flags_h = cache.sending_non_zero_flags_h;
 
+  host_timer.Stop();
+
   Kokkos::parallel_for(
       PARTHENON_AUTO_LABEL,
       Kokkos::TeamPolicy<>(parthenon::DevExecSpace(), nbound, Kokkos::AUTO),
@@ -131,6 +141,7 @@ TaskStatus SendBoundBufs(std::shared_ptr<MeshData<Real>> &md) {
         });
       });
 
+  BlockTimerHost host_timer2(pmesh->block_cost_host, min_lid, max_lid);
   // Send buffers
   if (Globals::sparse_config.enabled)
     Kokkos::deep_copy(sending_nonzero_flags_h, sending_nonzero_flags);
@@ -146,6 +157,7 @@ TaskStatus SendBoundBufs(std::shared_ptr<MeshData<Real>> &md) {
     else
       buf.SendNull();
   }
+  host_timer2.Stop();
 
   return TaskStatus::complete;
 }
diff --git a/src/mesh/mesh.cpp b/src/mesh/mesh.cpp
index 8eb073ea11f1..d759930b3865 100644
--- a/src/mesh/mesh.cpp
+++ b/src/mesh/mesh.cpp
@@ -1142,6 +1142,14 @@ void Mesh::Initialize(bool init_problem, ParameterInput *pin, ApplicationInput *
       init_done = false;
       // caching nbtotal the private variable my be updated in the following function
       const int nb_before_loadbalance = nbtotal;
+      // don't trust costs during problem initialization
+      for (int i = 0; i < nmb; ++i)
+        block_cost_host[i] = 1.0;
+#ifdef ENABLE_LB_TIMERS
+      parthenon::par_for(
+          DEFAULT_LOOP_PATTERN, PARTHENON_AUTO_LABEL, DevExecSpace(), 0, nmb - 1,
+          KOKKOS_LAMBDA(const int b) { block_cost(b) = 1.0; });
+#endif
       LoadBalancingAndAdaptiveMeshRefinement(pin, app_in);
       if (nbtotal == nb_before_loadbalance) {
         init_done = true;