do view of views correctly for Kokkos 4.5.1

parthenon-hpc-lab · Nov 26, 2024 · e95f691 · e95f691
1 parent 9ae28c8
commit e95f691
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 6 deletions.
diff --git a/src/bvals/comms/coalesced_buffers.cpp b/src/bvals/comms/coalesced_buffers.cpp
@@ -46,7 +46,7 @@ void CoalescedBuffer::AllocateCoalescedBuffer() {
 }
 
 //----------------------------------------------------------------------------------------
-ParArray1D<BndId> &CoalescedBuffer::GetBndIdsOnDevice(const std::set<Uid_t> &vars,
+ParArray1DRaw<BndId> &CoalescedBuffer::GetBndIdsOnDevice(const std::set<Uid_t> &vars,
                                                       int *pcomb_size) {
   const auto &var_set = vars.size() == 0 ? all_vars : vars;
   auto &bnd_ids_device = bnd_ids_device_map[var_set];
@@ -75,8 +75,8 @@ ParArray1D<BndId> &CoalescedBuffer::GetBndIdsOnDevice(const std::set<Uid_t> &var
   }
 
   if (nbnd_id != bnd_ids_device.size()) {
-    bnd_ids_device = ParArray1D<BndId>("bnd_id", nbnd_id);
-    bnd_ids_host = Kokkos::create_mirror_view(bnd_ids_device);
+    bnd_ids_device = ParArray1DRaw<BndId>(parthenon::ViewOfViewAlloc("bnd_id"), nbnd_id);
+    bnd_ids_host = create_view_of_view_mirror(bnd_ids_device);
     updated = true;
   }
 

diff --git a/src/bvals/comms/coalesced_buffers.hpp b/src/bvals/comms/coalesced_buffers.hpp
@@ -61,8 +61,8 @@ struct CoalescedBuffer {
   using var_buf_t = CommBuffer<buf_pool_t<Real>::owner_t>;
   std::map<Uid_t, std::vector<std::pair<BndId, var_buf_t *>>> coalesced_info_buf;
   std::set<Uid_t> all_vars;
-  std::unordered_map<std::set<Uid_t>, ParArray1D<BndId>, uid_set_hash> bnd_ids_device_map;
-  std::unordered_map<std::set<Uid_t>, ParArray1D<BndId>::host_mirror_type, uid_set_hash>
+  std::unordered_map<std::set<Uid_t>, ParArray1DRaw<BndId>, uid_set_hash> bnd_ids_device_map;
+  std::unordered_map<std::set<Uid_t>, ParArray1DRaw<BndId>::host_mirror_type, uid_set_hash>
       bnd_ids_host_map;
   CommBuffer<buf_t> coalesced_comm_buffer;
   CommBuffer<std::vector<int>> sparse_status_buffer;
@@ -92,7 +92,7 @@ struct CoalescedBuffer {
            coalesced_comm_buffer.IsAvailableForWrite();
   }
 
-  ParArray1D<BndId> &GetBndIdsOnDevice(const std::set<Uid_t> &vars,
+  ParArray1DRaw<BndId> &GetBndIdsOnDevice(const std::set<Uid_t> &vars,
                                        int *pcomb_size = nullptr);
 
   void PackAndSend(const std::set<Uid_t> &vars);
@@ -176,6 +176,7 @@ struct CoalescedComms {
     do {
       can_delete = true;
       for (auto &[p, cbr] : coalesced_send_buffers) {
+        can_delete = cbr.message.IsAvailableForWrite() && can_delete;
         for (auto &[r, cbrp] : cbr.coalesced_bufs) {
           can_delete = cbrp.IsAvailableForWrite() && can_delete;
         }