Throw a runtime error for incorrect memory space when Kokkos::Views a…

…re used
ornladios · Dec 28, 2022 · 95d0a09 · 95d0a09
1 parent 0e7c650
commit 95d0a09
Show file tree

Hide file tree

Showing 6 changed files with 32 additions and 39 deletions.
diff --git a/bindings/CXX11/adios2/cxx11/Engine.h b/bindings/CXX11/adios2/cxx11/Engine.h
@@ -18,6 +18,7 @@
 
 #include "adios2/common/ADIOSMacros.h"
 #include "adios2/common/ADIOSTypes.h"
+#include "adios2/helper/adiosFunctions.h"
 
 namespace adios2
 {
@@ -219,7 +220,14 @@ class Engine
     {
         auto adios_data = static_cast<AdiosView<U>>(data);
         auto mem_space = adios_data.memory_space();
-        variable.SetMemorySpace(mem_space);
+#ifdef ADIOS2_HAVE_GPU_SUPPORT
+        if (variable.m_MemorySpace != MemorySpace::Detect &&
+            variable.m_MemorySpace != mem_space)
+            helper::Throw<std::runtime_error>(
+                "CXX-Bindings", "Engine", "Put",
+                "Memory space mismatch between the variable and the "
+                "Kokkos::View");
+#endif
         Put(variable, adios_data.data(), launch);
     }
 

diff --git a/bindings/CXX11/adios2/cxx11/KokkosView.h b/bindings/CXX11/adios2/cxx11/KokkosView.h
@@ -17,14 +17,26 @@ struct memspace_kokkos_to_adios2<Kokkos::HostSpace>
     static constexpr adios2::MemorySpace value = adios2::MemorySpace::Host;
 };
 
-#ifdef KOKKOS_ENABLE_CUDA
+#if defined(KOKKOS_ENABLE_CUDA) && defined(ADIOS2_HAVE_CUDA)
 
 template <>
 struct memspace_kokkos_to_adios2<Kokkos::CudaSpace>
 {
     static constexpr adios2::MemorySpace value = adios2::MemorySpace::CUDA;
 };
 
+template <>
+struct memspace_kokkos_to_adios2<Kokkos::CudaUVMSpace>
+{
+    static constexpr adios2::MemorySpace value = adios2::MemorySpace::CUDA;
+};
+
+template <>
+struct memspace_kokkos_to_adios2<Kokkos::CudaHostPinnedSpace>
+{
+    static constexpr adios2::MemorySpace value = adios2::MemorySpace::CUDA;
+};
+
 #endif
 
 } // namespace detail

diff --git a/source/adios2/core/Variable.cpp b/source/adios2/core/Variable.cpp
@@ -49,7 +49,7 @@ namespace core
         info.StepsCount = stepsCount;                                          \
         info.Data = const_cast<T *>(data);                                     \
         info.Operations = m_Operations;                                        \
-        info.MemSpace = DetectMemorySpace((void *)data);                       \
+        info.MemSpace = GetMemorySpace((void *)data);                          \
         m_BlocksInfo.push_back(info);                                          \
         return m_BlocksInfo.back();                                            \
     }                                                                          \

diff --git a/source/adios2/core/VariableBase.cpp b/source/adios2/core/VariableBase.cpp
@@ -48,13 +48,12 @@ size_t VariableBase::TotalSize() const noexcept
     return helper::GetTotalSize(m_Count);
 }
 
-MemorySpace VariableBase::DetectMemorySpace(const void *ptr)
+MemorySpace VariableBase::GetMemorySpace(const void *ptr)
 {
 #ifdef ADIOS2_HAVE_GPU_SUPPORT
-    if (m_MemSpaceRequested != MemorySpace::Detect)
+    if (m_MemSpace != MemorySpace::Detect)
     {
-        m_MemSpaceDetected = m_MemSpaceRequested;
-        return m_MemSpaceRequested;
+        return m_MemSpace;
     }
 #endif
 
@@ -63,32 +62,13 @@ MemorySpace VariableBase::DetectMemorySpace(const void *ptr)
     cudaPointerGetAttributes(&attr, ptr);
     if (attr.type == cudaMemoryTypeDevice)
     {
-        m_MemSpaceDetected = MemorySpace::CUDA;
         return MemorySpace::CUDA;
     }
-#endif
-    m_MemSpaceDetected = MemorySpace::Host;
-    return MemorySpace::Host;
-}
-
-MemorySpace VariableBase::GetMemorySpace(const void *ptr)
-{
-#ifdef ADIOS2_HAVE_GPU_SUPPORT
-    if (m_MemSpaceDetected == MemorySpace::Detect)
-        m_MemSpaceDetected = DetectMemorySpace(ptr);
-    return m_MemSpaceDetected;
 #endif
     return MemorySpace::Host;
 }
 
-void VariableBase::SetMemorySpace(const MemorySpace mem)
-{
-    m_MemSpaceRequested = mem;
-    // reset the detected memory space
-#ifdef ADIOS2_HAVE_GPU_SUPPORT
-    m_MemSpaceDetected = MemorySpace::Detect;
-#endif
-}
+void VariableBase::SetMemorySpace(const MemorySpace mem){ m_MemSpace = mem; }
 
 void VariableBase::SetShape(const adios2::Dims &shape)
 {

diff --git a/source/adios2/core/VariableBase.h b/source/adios2/core/VariableBase.h
@@ -50,13 +50,12 @@ class VariableBase
     /** Variable -> sizeof(T),
      *  VariableStruct -> from constructor sizeof(struct) */
     const size_t m_ElementSize;
-    /* User requested memory space and the detected memory space */
+
+    /* User requested memory space */
 #ifdef ADIOS2_HAVE_GPU_SUPPORT
-    MemorySpace m_MemSpaceRequested = MemorySpace::Detect;
-    MemorySpace m_MemSpaceDetected = MemorySpace::Detect;
+    MemorySpace m_MemSpace = MemorySpace::Detect;
 #else
-    MemorySpace m_MemSpaceRequested = MemorySpace::Host;
-    MemorySpace m_MemSpaceDetected = MemorySpace::Host;
+    MemorySpace m_MemSpace = MemorySpace::Host;
 #endif
 
     ShapeID m_ShapeID = ShapeID::Unknown; ///< see shape types in ADIOSTypes.h
@@ -123,12 +122,6 @@ class VariableBase
      */
     size_t TotalSize() const noexcept;
 
-    /**
-     * Detect the memory space where a buffer was allocated and return it
-     * @param pointer to the user data
-     */
-    MemorySpace DetectMemorySpace(const void *ptr);
-
     /**
      * Get the memory space where a given buffers was allocated
      * @param pointer to the user data

diff --git a/source/adios2/engine/bp5/BP5Writer.cpp b/source/adios2/engine/bp5/BP5Writer.cpp
@@ -1737,7 +1737,7 @@ void BP5Writer::PutCommon(VariableBase &variable, const void *values, bool sync)
     }
 
     // if the user buffer is allocated on the GPU always use sync mode
-    if (variable.DetectMemorySpace(values) != MemorySpace::Host)
+    if (variable.GetMemorySpace(values) != MemorySpace::Host)
         sync = true;
 
     size_t *Shape = NULL;