Fix missing fence and incorrect ptr calculation in mpi_block_type{put…

…,get}
kokkos · Sep 7, 2023 · ced79a9 · ced79a9
1 parent abfa254
commit ced79a9
Show file tree

Hide file tree

Showing 4 changed files with 26 additions and 38 deletions.
diff --git a/src/core/Kokkos_RemoteSpaces_ViewMapping.hpp b/src/core/Kokkos_RemoteSpaces_ViewMapping.hpp
@@ -1057,15 +1057,10 @@ class ViewMapping<Traits, Kokkos::Experimental::RemoteSpaceSpecializeTag> {
       Kokkos::Impl::ViewCtorProp<P...> const &arg_prop,
       typename Traits::array_layout const &arg_layout)
       : m_offset_remote_dim(0),
-#ifdef KRS_ENABLE_MPISPACE
         m_handle(
             ((Kokkos::Impl::ViewCtorProp<void, pointer_type> const &)arg_prop)
                 .value)
-#else
-        m_handle(
-            ((Kokkos::Impl::ViewCtorProp<void, pointer_type> const &)arg_prop)
-                .value)
-#endif
+
   {
     typedef typename Traits::value_type value_type;
     typedef std::integral_constant<

diff --git a/src/impl/mpispace/Kokkos_MPISpace_BlockOps.hpp b/src/impl/mpispace/Kokkos_MPISpace_BlockOps.hpp
@@ -33,9 +33,10 @@ namespace Impl {
     int _typesize;                                                             \
     MPI_Request request;                                                       \
     MPI_Type_size(mpi_type, &_typesize);                                       \
-    MPI_Rput(ptr, nelems, mpi_type, pe,                                        \
-             sizeof(SharedAllocationHeader) + offset * _typesize, nelems,      \
-             mpi_type, win, &request);                                         \
+    const void *src_adr = ptr + offset;                                        \
+    size_t win_offset   = sizeof(SharedAllocationHeader) + offset * _typesize; \
+    MPI_Rput(src_adr, nelems, mpi_type, pe, win_offset, nelems, mpi_type, win, \
+             &request);                                                        \
     MPI_Wait(&request, MPI_STATUS_IGNORE);                                     \
   }
 
@@ -54,18 +55,19 @@ KOKKOS_REMOTESPACES_PUT(double, MPI_DOUBLE)
 
 #undef KOKKOS_REMOTESPACES_PUT
 
-#define KOKKOS_REMOTESPACES_GET(type, mpi_type)                           \
-  static KOKKOS_INLINE_FUNCTION void mpi_block_type_get(                  \
-      type *ptr, const size_t offset, const size_t nelems, const int pe,  \
-      const MPI_Win &win) {                                               \
-    assert(win != MPI_WIN_NULL);                                          \
-    int _typesize;                                                        \
-    MPI_Request request;                                                  \
-    MPI_Type_size(mpi_type, &_typesize);                                  \
-    MPI_Rget(ptr, nelems, mpi_type, pe,                                   \
-             sizeof(SharedAllocationHeader) + offset * _typesize, nelems, \
-             mpi_type, win, &request);                                    \
-    MPI_Wait(&request, MPI_STATUS_IGNORE);                                \
+#define KOKKOS_REMOTESPACES_GET(type, mpi_type)                                \
+  static KOKKOS_INLINE_FUNCTION void mpi_block_type_get(                       \
+      type *ptr, const size_t offset, const size_t nelems, const int pe,       \
+      const MPI_Win &win) {                                                    \
+    assert(win != MPI_WIN_NULL);                                               \
+    int _typesize;                                                             \
+    MPI_Request request;                                                       \
+    MPI_Type_size(mpi_type, &_typesize);                                       \
+    void *dst_adr     = ptr + offset;                                          \
+    size_t win_offset = sizeof(SharedAllocationHeader) + offset * _typesize;   \
+    MPI_Rget(dst_adr, nelems, mpi_type, pe, win_offset, nelems, mpi_type, win, \
+             &request);                                                        \
+    MPI_Wait(&request, MPI_STATUS_IGNORE);                                     \
   }
 
 KOKKOS_REMOTESPACES_GET(char, MPI_SIGNED_CHAR)

diff --git a/src/impl/mpispace/Kokkos_MPISpace_DataHandle.hpp b/src/impl/mpispace/Kokkos_MPISpace_DataHandle.hpp
@@ -47,45 +47,36 @@ struct MPIDataHandle {
   }
 
   KOKKOS_INLINE_FUNCTION
-  MPIDataHandle *operator+(size_t &offset) {
-    ptr += offset;
-    loc.offset += offset;
+  MPIDataHandle operator+(size_t &offset) {
+    return MPIDataHandle(ptr += offset, loc.offset += offset);
   }
 };
 
 template <class T, class Traits>
 struct BlockDataHandle {
   T *ptr;
-  MPIAccessLocation remote_loc;
+  MPIAccessLocation loc;
   size_t pe;
   size_t elems;
 
-  KOKKOS_INLINE_FUNCTION
-  BlockDataHandle() : elems(0) {}
-
   KOKKOS_INLINE_FUNCTION
   BlockDataHandle(T *ptr_, MPI_Win win_, size_t offset_, size_t elems_,
                   size_t pe_)
-      : ptr(ptr_), remote_loc(win_, offset_), elems(elems_), pe(pe_) {}
+      : ptr(ptr_), loc(win_, offset_), elems(elems_), pe(pe_) {}
 
   KOKKOS_INLINE_FUNCTION
   BlockDataHandle(BlockDataHandle<T, Traits> const &arg)
-      : ptr(arg.ptr),
-        remote_loc(arg.remote_loc),
-        elems(arg.elems),
-        pe(arg.pe) {}
+      : ptr(arg.ptr), loc(arg.loc), elems(arg.elems), pe(arg.pe) {}
 
   KOKKOS_INLINE_FUNCTION
   void get() {
-    MPIBlockDataElement<T, Traits> element(ptr, remote_loc.win, pe,
-                                           remote_loc.offset, elems);
+    MPIBlockDataElement<T, Traits> element(ptr, loc.win, pe, loc.offset, elems);
     element.get();
   }
 
   KOKKOS_INLINE_FUNCTION
   void put() {
-    MPIBlockDataElement<T, Traits> element(ptr, remote_loc.win, pe,
-                                           remote_loc.offset, elems);
+    MPIBlockDataElement<T, Traits> element(ptr, loc.win, pe, loc.offset, elems);
     element.put();
   }
 };

diff --git a/unit_tests/Test_Reduction.cpp b/unit_tests/Test_Reduction.cpp
@@ -117,7 +117,7 @@ void test_scalar_reduce_partitioned_1D(int dim1) {
         v(my_rank, i) = static_cast<Data_t>(start + i);
       });
 
-  Kokkos::fence();
+  RemoteSpace_t().fence();
 
   Data_t gsum = 0;
   Kokkos::parallel_reduce(