Fix multiple calls to MPISpace::deallocate on View dealloc

kokkos · May 6, 2024 · 7612edf · 7612edf
1 parent eb5acd5
commit 7612edf
Show file tree

Hide file tree

Showing 3 changed files with 19 additions and 26 deletions.
diff --git a/src/impl/mpispace/Kokkos_MPISpace.cpp b/src/impl/mpispace/Kokkos_MPISpace.cpp
@@ -27,6 +27,7 @@ namespace Experimental {
 
 MPI_Win MPISpace::current_win;
 std::vector<MPI_Win> MPISpace::mpi_windows;
+std::mutex internal_mpi_backend_mutex;
 
 /* Default allocation mechanism */
 MPISpace::MPISpace() : allocation_mode(Kokkos::Experimental::Symmetric) {}
@@ -80,19 +81,24 @@ void *MPISpace::impl_allocate(
       assert(ptr != nullptr);
       assert(current_win != MPI_WIN_NULL);
 
+      assert(ptr != nullptr);
+      assert(current_win != MPI_WIN_NULL);
+
       int ret = MPI_Win_lock_all(MPI_MODE_NOCHECK, current_win);
       if (ret != MPI_SUCCESS) {
         Kokkos::abort("MPI window lock all failed.");
       }
       int i;
+
+      internal_mpi_backend_mutex.lock();
       for (i = 0; i < mpi_windows.size(); ++i) {
         if (mpi_windows[i] == MPI_WIN_NULL) break;
       }
-
       if (i == mpi_windows.size())
         mpi_windows.push_back(current_win);
       else
         mpi_windows[i] = current_win;
+      internal_mpi_backend_mutex.unlock();
     } else {
       Kokkos::abort("MPISpace only supports symmetric allocation policy.");
     }
@@ -103,11 +109,9 @@ void *MPISpace::impl_allocate(
   using MemAllocFailureMode = Kokkos::Impl::Experimental::
       RemoteSpacesMemoryAllocationFailure::FailureMode;
 
-  if ((ptr == nullptr) ||
-      (reinterpret_cast<uintptr_t>(ptr) == ~uintptr_t(0))
-      // MPI_Win_allocate may allocate non-alligned to
-      // Kokkos::Impl::MEMORY_ALIGNMENT
-      || (reinterpret_cast<uintptr_t>(ptr) & alignment_mask)) {
+  if ((ptr == nullptr)
+      /* Uncomment once MPI makes the alignement via info keys work */
+      /*|| (reinterpret_cast<uintptr_t>(ptr) & alignment_mask)*/) {
     MemAllocFailureMode failure_mode =
         MemAllocFailureMode::AllocationNotAligned;
     if (ptr == nullptr) {
@@ -151,15 +155,15 @@ void MPISpace::impl_deallocate(
                                         reported_size);
     }
 
+    internal_mpi_backend_mutex.lock();
     int last_valid;
     for (last_valid = 0; last_valid < mpi_windows.size(); ++last_valid) {
       if (mpi_windows[last_valid] == MPI_WIN_NULL) break;
     }
 
     last_valid--;
     for (int i = 0; i < mpi_windows.size(); ++i) {
-      if (mpi_windows[i] == current_win &&
-          current_win != mpi_windows[last_valid]) {
+      if (mpi_windows[i] == current_win) {
         mpi_windows[i]          = mpi_windows[last_valid];
         mpi_windows[last_valid] = MPI_WIN_NULL;
         break;
@@ -170,28 +174,21 @@ void MPISpace::impl_deallocate(
     MPI_Win_unlock_all(current_win);
     MPI_Win_free(&current_win);
 
-    // We pass a mempory space instance do multiple Views thus
-    // setting "current_win = MPI_WIN_NULL;" will result in a wrong handle if
-    // subsequent view runs out of scope
-    // Fixme: The following only works when views are allocated sequentially
-    // We need a thread-safe map to associate views and windows
-
     if (last_valid != 0)
       current_win = mpi_windows[last_valid - 1];
     else
       current_win = MPI_WIN_NULL;
+
+    internal_mpi_backend_mutex.unlock();
   }
 }
 
 void MPISpace::fence() {
-  Kokkos::fence();
+  internal_mpi_backend_mutex.lock();
   for (int i = 0; i < mpi_windows.size(); i++) {
-    if (mpi_windows[i] != MPI_WIN_NULL) {
-      MPI_Win_flush_all(mpi_windows[i]);
-    } else {
-      break;
-    }
+    if (mpi_windows[i] != MPI_WIN_NULL) MPI_Win_flush_all(mpi_windows[i]);
   }
+  internal_mpi_backend_mutex.unlock();
   MPI_Barrier(MPI_COMM_WORLD);
 }
 
@@ -200,7 +197,6 @@ size_t get_num_pes() {
   MPI_Comm_size(MPI_COMM_WORLD, &n_ranks);
   return n_ranks;
 }
-
 size_t get_my_pe() {
   int rank;
   MPI_Comm_rank(MPI_COMM_WORLD, &rank);

diff --git a/src/impl/mpispace/Kokkos_MPISpace_AllocationRecord.cpp b/src/impl/mpispace/Kokkos_MPISpace_AllocationRecord.cpp
@@ -31,11 +31,7 @@ SharedAllocationRecord<void, void>
 
 SharedAllocationRecord<Kokkos::Experimental::MPISpace,
                        void>::~SharedAllocationRecord() {
-  m_space.deallocate(m_label.c_str(),
-                     SharedAllocationRecord<void, void>::m_alloc_ptr,
-                     SharedAllocationRecord<void, void>::m_alloc_size,
-                     (SharedAllocationRecord<void, void>::m_alloc_size -
-                      sizeof(SharedAllocationHeader)));
+  // Let SharedAllocationRecordCommon do the deallocation
 }
 
 SharedAllocationHeader *_do_allocation(

diff --git a/src/impl/mpispace/Kokkos_MPISpace_AllocationRecord.hpp b/src/impl/mpispace/Kokkos_MPISpace_AllocationRecord.hpp
@@ -74,6 +74,7 @@ class SharedAllocationRecord<Kokkos::Experimental::MPISpace, void>
     this->base_t::_fill_host_accessible_header_info(*RecordBase::m_alloc_ptr,
                                                     arg_label);
 #endif
+    win = m_space.current_win;
   }
 
   SharedAllocationRecord(