Merge pull request kokkos#50 from janciesko/fix_offset_mpiwin

jeffhammond · Aug 1, 2022 · 35f4cb5 · 35f4cb5
2 parents a3af296 + daaaf8b
commit 35f4cb5
Show file tree

Hide file tree

Showing 3 changed files with 26 additions and 19 deletions.
diff --git a/src/core/Kokkos_RemoteSpaces_ViewMapping.hpp b/src/core/Kokkos_RemoteSpaces_ViewMapping.hpp
@@ -302,12 +302,23 @@ class ViewMapping<
     dst.m_offset_remote_dim = extents.domain_offset(0);
     dst.dim0_is_pe          = R0;
 
+#ifdef KOKKOS_ENABLE_MPISPACE
+    // Subviews propagate MPI_Window of the original view
+    dst.m_handle = ViewDataHandle<DstTraits>::assign(
+        src.m_handle,
+        src.m_offset(0, extents.domain_offset(1), extents.domain_offset(2),
+                     extents.domain_offset(3), extents.domain_offset(4),
+                     extents.domain_offset(5), extents.domain_offset(6),
+                     extents.domain_offset(7)),
+        src.m_handle.win);
+#else
     dst.m_handle = ViewDataHandle<DstTraits>::assign(
         src.m_handle,
         src.m_offset(0, extents.domain_offset(1), extents.domain_offset(2),
                      extents.domain_offset(3), extents.domain_offset(4),
                      extents.domain_offset(5), extents.domain_offset(6),
                      extents.domain_offset(7)));
+#endif
   }
 };
 

diff --git a/src/impl/mpispace/Kokkos_MPISpace_DataHandle.hpp b/src/impl/mpispace/Kokkos_MPISpace_DataHandle.hpp
@@ -52,21 +52,21 @@ template <class T, class Traits>
 struct MPIDataHandle {
   T *ptr;
   mutable MPI_Win win;
+  size_t win_offset;
   KOKKOS_INLINE_FUNCTION
-  MPIDataHandle() : ptr(NULL), win(MPI_WIN_NULL) {}
+  MPIDataHandle() : ptr(NULL), win(MPI_WIN_NULL), win_offset(0) {}
   KOKKOS_INLINE_FUNCTION
-  MPIDataHandle(T *ptr_, MPI_Win &win_) : ptr(ptr_), win(win_) {}
+  MPIDataHandle(T *ptr_, MPI_Win &win_, size_t offset_ = 0)
+      : ptr(ptr_), win(win_), win_offset(offset_) {}
   KOKKOS_INLINE_FUNCTION
   MPIDataHandle(MPIDataHandle<T, Traits> const &arg)
-      : ptr(arg.ptr), win(arg.win) {}
-  KOKKOS_INLINE_FUNCTION
-  MPIDataHandle(T *ptr_) : ptr(ptr_), win(MPI_WIN_NULL) {}
+      : ptr(arg.ptr), win(arg.win), win_offset(arg.win_offset) {}
 
   template <typename iType>
   KOKKOS_INLINE_FUNCTION MPIDataElement<T, Traits> operator()(
       const int &pe, const iType &i) const {
     assert(win != MPI_WIN_NULL);
-    MPIDataElement<T, Traits> element(&win, pe, i);
+    MPIDataElement<T, Traits> element(&win, pe, i + win_offset);
     return element;
   }
 
@@ -95,10 +95,8 @@ struct ViewDataHandle<
 
   template <class SrcHandleType>
   KOKKOS_INLINE_FUNCTION static handle_type assign(
-      SrcHandleType const arg_data_ptr, size_t offset) {
-    // FIXME: Invocation of handle_type constructor sets win to MPI_WIN_NULL
-    // This is invoked by subview ViewMapping so subviews will likely fail
-    return handle_type(arg_data_ptr + offset);
+      SrcHandleType const arg_data_ptr, size_t offset, MPI_Win &win) {
+    return handle_type(arg_data_ptr + offset, win, offset);
   }
 };
 

diff --git a/unit_tests/Test_PartitionedSubview.cpp b/unit_tests/Test_PartitionedSubview.cpp
@@ -72,8 +72,7 @@ void test_partitioned_subview1D(int i1, int i2, int sub1, int sub2) {
   ViewHost_3D_t v_h("HostView", 1, i1, i2);
 
   // Init
-  for (int i = 0; i < i1; ++i)
-    for (int j = 0; j < i2; ++j) v_h(0, i, j) = VAL;
+  deep_copy(v_h, VAL);
 
   auto v_sub = Kokkos::subview(v, std::make_pair(my_rank, my_rank + 1),
                                Kokkos::ALL, Kokkos::ALL);
@@ -93,10 +92,11 @@ void test_partitioned_subview1D(int i1, int i2, int sub1, int sub2) {
 
   for (int i = 0; i < i1; ++i)
     for (int j = 0; j < i2; ++j)
-      if (i == sub1 && j == sub2)
+      if (i == sub1 && j == sub2) {
         ASSERT_EQ(v_h(0, i, j), VAL + 2);
-      else
+      } else {
         ASSERT_EQ(v_h(0, i, j), VAL);
+      }
 }
 
 template <class Data_t>
@@ -117,8 +117,7 @@ void test_partitioned_subview2D(int i1, int i2, int sub1) {
   ViewHost_3D_t v_h("HostView", 1, i1, i2);
 
   // Init
-  for (int i = 0; i < i1; ++i)
-    for (int j = 0; j < i2; ++j) v_h(0, i, j) = VAL;
+  deep_copy(v_h, VAL);
 
   auto v_sub = Kokkos::subview(v, std::make_pair(my_rank, my_rank + 1),
                                Kokkos::ALL, Kokkos::ALL);
@@ -162,8 +161,7 @@ void test_partitioned_subview3D(int i1, int i2, int sub1, int sub2) {
   ViewHost_3D_t v_h("HostView", 1, i1, i2);
 
   // Init
-  for (int i = 0; i < i1; ++i)
-    for (int j = 0; j < i2; ++j) v_h(0, i, j) = VAL;
+  deep_copy(v_h, VAL);
 
   auto v_sub = Kokkos::subview(v, std::make_pair(my_rank, my_rank + 1),
                                Kokkos::ALL, Kokkos::ALL);
@@ -249,4 +247,4 @@ TEST(TEST_CATEGORY, test_partitioned_subview) {
   test_partitioned_subview2D_byRank<int>(50, 77);
 }
 
-#endif /* TEST_PARTITIONED_SUBVIEW_HPP_ */
+#endif /* TEST_PARTITIONED_SUBVIEW_HPP_ */