Merge pull request #4346 from anagainaru/memSelLayout

Update the memory selection layout when there is a layout mismatch
ornladios · Oct 7, 2024 · 24c3860 · 24c3860
2 parents 2b40572 + 04aac83
commit 24c3860
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 4 deletions.
diff --git a/bindings/CXX11/adios2/cxx11/Engine.h b/bindings/CXX11/adios2/cxx11/Engine.h
@@ -218,9 +218,7 @@ class Engine
         auto bufferView = static_cast<AdiosView<U>>(data);
 #if defined(ADIOS2_HAVE_KOKKOS) || defined(ADIOS2_HAVE_GPU_SUPPORT)
         auto bufferMem = bufferView.memory_space();
-        auto bufferLayout = bufferView.layout();
         variable.SetMemorySpace(bufferMem);
-        variable.SetArrayLayout(bufferLayout);
 #endif
         Put(variable, bufferView.data(), launch);
     }
@@ -424,9 +422,7 @@ class Engine
         auto bufferView = static_cast<AdiosView<U>>(data);
 #if defined(ADIOS2_HAVE_KOKKOS) || defined(ADIOS2_HAVE_GPU_SUPPORT)
         auto bufferMem = bufferView.memory_space();
-        auto bufferLayout = bufferView.layout();
         variable.SetMemorySpace(bufferMem);
-        variable.SetArrayLayout(bufferLayout);
 #endif
         Get(variable, bufferView.data(), launch);
     }

diff --git a/source/adios2/core/VariableBase.cpp b/source/adios2/core/VariableBase.cpp
@@ -57,6 +57,8 @@ void VariableBase::SetArrayLayout(const ArrayOrdering layout)
         UpdateLayout(m_Shape);
         UpdateLayout(m_Count);
         UpdateLayout(m_Start);
+        UpdateLayout(m_MemoryStart);
+        UpdateLayout(m_MemoryCount);
         return;
     }
     if (m_ArrayLayout != layout)
@@ -103,6 +105,9 @@ MemorySpace VariableBase::GetMemorySpace(const void *ptr)
 
 void VariableBase::SetMemorySpace(const MemorySpace mem)
 {
+#if defined(ADIOS2_HAVE_KOKKOS) || defined(ADIOS2_HAVE_GPU_SUPPORT)
+    ArrayOrdering layout = m_BaseLayout;
+#endif
 #ifdef ADIOS2_HAVE_GPU_SUPPORT
     if (m_MemSpace != MemorySpace::Detect && m_MemSpace != mem)
     {
@@ -117,6 +122,14 @@ void VariableBase::SetMemorySpace(const MemorySpace mem)
                                                  ExistingMemSpace + " and cannot received a " +
                                                  NewMemSpace + " buffer");
     }
+    if (mem == MemorySpace::GPU)
+        layout = ArrayOrdering::ColumnMajor;
+#endif
+#if defined(ADIOS2_HAVE_KOKKOS) || defined(ADIOS2_HAVE_GPU_SUPPORT)
+    // set the layout based on the buffer memory space
+    // skipping throwing an exception for a mismatch
+    if (m_ArrayLayout == ArrayOrdering::Auto)
+        SetArrayLayout(layout);
 #endif
     m_MemSpace = mem;
 }