Allow Put/Get functions to accept Kokkos Views

CMakeLists.txt

@@ -146,6 +146,7 @@ adios_option(LIBPRESSIO "Enable support for LIBPRESSIO transforms" AUTO)
 adios_option(MGARD      "Enable support for MGARD transforms" AUTO)
 adios_option(PNG        "Enable support for PNG transforms" AUTO)
 adios_option(CUDA       "Enable support for Cuda" AUTO)
+adios_option(Kokkos     "Enable support for Kokkos" AUTO)
 adios_option(MPI        "Enable support for MPI" AUTO)
 adios_option(DAOS       "Enable support for DAOS" AUTO)
 adios_option(DataMan    "Enable support for DataMan" AUTO)
@@ -222,7 +223,7 @@ endif()
 
 
 set(ADIOS2_CONFIG_OPTS
-    BP5 DataMan DataSpaces HDF5 HDF5_VOL MHS SST CUDA Fortran MPI Python Blosc BZip2 LIBPRESSIO MGARD PNG SZ ZFP DAOS IME O_DIRECT Sodium SysVShMem ZeroMQ Profiling Endian_Reverse
+    BP5 DataMan DataSpaces HDF5 HDF5_VOL MHS SST Fortran MPI Python Blosc BZip2 LIBPRESSIO MGARD PNG SZ ZFP CUDA Kokkos DAOS IME O_DIRECT Sodium SysVShMem ZeroMQ Profiling Endian_Reverse
 )
 
 GenerateADIOSHeaderConfig(${ADIOS2_CONFIG_OPTS})

bindings/CXX11/CMakeLists.txt

@@ -25,6 +25,10 @@ set_property(TARGET adios2_cxx11 PROPERTY OUTPUT_NAME adios2${ADIOS2_LIBRARY_SUF
 target_link_libraries(adios2_cxx11 PRIVATE adios2_core adios2::thirdparty::pugixml)
 target_compile_features(adios2_cxx11 INTERFACE ${ADIOS2_CXX11_FEATURES})
 
+if(ADIOS2_HAVE_Kokkos)
+    target_link_libraries(adios2_cxx11 PUBLIC Kokkos::kokkos)
+endif()
+
 target_include_directories(adios2_cxx11
   PUBLIC
     $<BUILD_INTERFACE:${CMAKE_CURRENT_SOURCE_DIR}>
@@ -80,6 +84,7 @@ install(
 install(
   FILES adios2/cxx11/ADIOS.h
         adios2/cxx11/ADIOS.inl
+        adios2/cxx11/ADIOSMacros.h
         adios2/cxx11/IO.h
         adios2/cxx11/Group.h
         adios2/cxx11/Variable.h

bindings/CXX11/adios2/cxx11/ADIOSMacros.h

@@ -0,0 +1,37 @@
+/*
+ * Distributed under the OSI-approved Apache License, Version 2.0.  See
+ * accompanying file Copyright.txt for details.
+ *
+ * ADIOSMacros.h : a set of helper macros used by the CXX binding
+ *
+ */
+
+#ifndef ADIOS2_ADIOSMACROS_CXX
+#define ADIOS2_ADIOSMACROS_CXX
+
+#ifdef ADIOS2_HAVE_KOKKOS
+#include <Kokkos_Core.hpp>
+
+#define ADIOS2_KOKKOS_MEMORY_SPACE(MACRO, MEMORY_SPACE)                        \
+    MACRO(int32_t, MEMORY_SPACE)                                               \
+    MACRO(uint32_t, MEMORY_SPACE)                                              \
+    MACRO(int64_t, MEMORY_SPACE)                                               \
+    MACRO(uint64_t, MEMORY_SPACE)                                              \
+    MACRO(float, MEMORY_SPACE)                                                 \
+    MACRO(double, MEMORY_SPACE)
+
+#ifdef KOKKOS_ENABLE_CUDA
+#define ADIOS2_FOREACH_KOKKOS_TYPE_2ARGS(MACRO)                                \
+    ADIOS2_KOKKOS_MEMORY_SPACE(MACRO, Kokkos::HostSpace)                       \
+    ADIOS2_KOKKOS_MEMORY_SPACE(MACRO, Kokkos::CudaSpace)                       \
+    ADIOS2_KOKKOS_MEMORY_SPACE(MACRO, Kokkos::CudaHostPinnedSpace)             \
+    ADIOS2_KOKKOS_MEMORY_SPACE(MACRO, Kokkos::CudaUVMSpace)
+
+#else
+#define ADIOS2_FOREACH_KOKKOS_TYPE_2ARGS(MACRO)                                \
+    ADIOS2_KOKKOS_MEMORY_SPACE(MACRO, Kokkos::HostSpace)
+#endif /* KOKKOS_ENABLE_CUDA */
+
+#endif /* ADIOS2_HAVE_KOKKOS */
+
+#endif /* ADIOS2_ADIOSMACROS_CXX */

bindings/CXX11/adios2/cxx11/Engine.cpp

@@ -9,6 +9,7 @@
  */
 
 #include "Engine.h"
+#include "ADIOSMacros.h"
 #include "Engine.tcc"
 
 #include "adios2/core/Engine.h"
@@ -344,6 +345,18 @@ Engine::AllStepsBlocksInfo(const VariableNT &variable) const
     return ret;
 }
 
+#ifdef ADIOS2_HAVE_KOKKOS
+#define declare_template_instantiation(T, MemSpace)                            \
+    template void Engine::Put<T, MemSpace>(                                    \
+        Variable<T>, Kokkos::View<T *, MemSpace>, const Mode);                 \
+                                                                               \
+    template void Engine::Get<T, MemSpace>(                                    \
+        Variable<T>, Kokkos::View<T *, MemSpace>, const Mode);
+
+ADIOS2_FOREACH_KOKKOS_TYPE_2ARGS(declare_template_instantiation)
+#undef declare_template_instantiation
+#endif
+
 #define declare_template_instantiation(T)                                      \
     template void Engine::Put<T>(Variable<T>, const T *, const Mode);          \
     template void Engine::Put<T>(const std::string &, const T *, const Mode);  \

bindings/CXX11/adios2/cxx11/Engine.h

@@ -18,6 +18,10 @@
 #include "adios2/common/ADIOSMacros.h"
 #include "adios2/common/ADIOSTypes.h"
 
+#ifdef ADIOS2_HAVE_KOKKOS
+#include <Kokkos_Core_fwd.hpp>
+#endif
+
 namespace adios2
 {
 
@@ -211,6 +215,17 @@ class Engine
      * collective call and can only be called between Begin/EndStep pairs. */
     void PerformDataWrite();
 
+#ifdef ADIOS2_HAVE_KOKKOS
+    /** Get and Put functions for Kokkos buffers */
+    template <class T, class MemSpace>
+    void Put(Variable<T> variable, Kokkos::View<T *, MemSpace> data,
+             const Mode launch = Mode::Deferred);
+
+    template <class T, class MemSpace>
+    void Get(Variable<T> variable, Kokkos::View<T *, MemSpace> data,
+             const Mode launch = Mode::Deferred);
+#endif
+
     /**
      * Get data associated with a Variable from the Engine
      * @param variable contains variable metadata information

bindings/CXX11/adios2/cxx11/Engine.tcc

@@ -137,6 +137,20 @@ void Engine::Put(const std::string &variableName, const T &datum,
                   launch);
 }
 
+#ifdef ADIOS2_HAVE_KOKKOS
+template <class T, class MemSpace>
+void Engine::Put(Variable<T> variable, Kokkos::View<T *, MemSpace> data,
+                 const Mode launch)
+{
+    using IOType = typename TypeInfo<T>::IOType;
+    adios2::helper::CheckForNullptr(m_Engine, "in call to Engine::Put");
+    adios2::helper::CheckForNullptr(variable.m_Variable,
+                                    "for variable in call to Engine::Put");
+    m_Engine->Put(*variable.m_Variable,
+                  reinterpret_cast<const IOType *>(data.data()), launch);
+}
+#endif
+
 template <class T>
 void Engine::Get(Variable<T> variable, T *data, const Mode launch)
 {
@@ -234,6 +248,20 @@ void Engine::Get(Variable<T> variable, T **data) const
     return;
 }
 
+#ifdef ADIOS2_HAVE_KOKKOS
+template <class T, class MemSpace>
+void Engine::Get(Variable<T> variable, Kokkos::View<T *, MemSpace> data,
+                 const Mode launch)
+{
+    adios2::helper::CheckForNullptr(variable.m_Variable,
+                                    "for variable in call to Engine::Get");
+    using IOType = typename TypeInfo<T>::IOType;
+    adios2::helper::CheckForNullptr(m_Engine, "in call to Engine::Get");
+    m_Engine->Get(*variable.m_Variable, reinterpret_cast<IOType *>(data.data()),
+                  launch);
+}
+#endif
+
 template <class T>
 std::map<size_t, std::vector<typename Variable<T>::Info>>
 Engine::AllStepsBlocksInfo(const Variable<T> variable) const

cmake/DetectOptions.cmake

@@ -180,6 +180,18 @@ if(CMAKE_CUDA_COMPILER AND CUDAToolkit_FOUND)
   set(ADIOS2_HAVE_CUDA TRUE)
 endif()
 
+# Kokkos
+if(ADIOS2_USE_Kokkos)
+  if(ADIOS2_USE_Kokkos STREQUAL AUTO)
+    find_package(Kokkos QUIET)
+  else()
+    find_package(Kokkos REQUIRED)
+  endif()
+endif()
+if(Kokkos_FOUND)
+  set(ADIOS2_HAVE_Kokkos TRUE)
+endif()
+
 # Fortran
 if(ADIOS2_USE_Fortran STREQUAL AUTO)
   include(CheckLanguage)