Adding Kokkos specific functions

source/adios2/CMakeLists.txt

@@ -124,6 +124,11 @@ if(ADIOS2_HAVE_CUDA)
   set(maybe_adios2_core_cuda adios2_core_cuda)
 endif()
 
+if(ADIOS2_HAVE_Kokkos)
+  target_sources(adios2_core PRIVATE helper/adiosKokkos.h helper/adiosKokkos.cpp)
+  target_link_libraries(adios2_core PRIVATE Kokkos::kokkos)
+endif()
+
 target_include_directories(adios2_core
   PUBLIC
     $<BUILD_INTERFACE:${ADIOS2_SOURCE_DIR}/source>

source/adios2/core/ADIOS.cpp

@@ -66,6 +66,13 @@ class ADIOS::GlobalServices
             Aws::ShutdownAPI(options);
             isAWSInitialized = false;
         }
+#endif
+#ifdef ADIOS2_HAVE_KOKKOS
+        if (isKokkosInitialized && !isKokkosFinalized)
+        {
+            std::atexit(helper::KokkosFinalize);
+            isKokkosFinalized = true;
+        }
 #endif
         wasGlobalShutdown = true;
     }
@@ -85,6 +92,20 @@ class ADIOS::GlobalServices
     bool isAWSInitialized = false;
 #endif
 
+#ifdef ADIOS2_HAVE_KOKKOS
+    void Init_Kokkos_API()
+    {
+        if (helper::KokkosIsInitialized())
+            return;
+        if (!isKokkosInitialized)
+        {
+            helper::KokkosInit();
+            isKokkosInitialized = true;
+        }
+    }
+    bool isKokkosInitialized = false;
+    bool isKokkosFinalized = false;
+#endif
     bool wasGlobalShutdown = false;
 };
 
@@ -129,6 +150,9 @@ ADIOS::ADIOS(const std::string configFile, helper::Comm comm,
             YAMLInit(configFile);
         }
     }
+#ifdef ADIOS2_HAVE_KOKKOS
+    m_GlobalServices.Init_Kokkos_API();
+#endif
 }
 
 ADIOS::ADIOS(const std::string configFile, const std::string hostLanguage)

source/adios2/helper/adiosGPUFunctions.h

@@ -0,0 +1,12 @@
+#ifndef ADIOS2_HELPER_ADIOSGPUFUNCTIONS_H_
+#define ADIOS2_HELPER_ADIOSGPUFUNCTIONS_H_
+
+#ifdef ADIOS2_HAVE_CUDA
+#include "adios2/helper/adiosCUDA.h"
+#endif
+
+#ifdef ADIOS2_HAVE_KOKKOS
+#include "adios2/helper/adiosKokkos.h"
+#endif
+
+#endif /* ADIOS2_HELPER_ADIOSGPUFUNCTIONS_H_ */

source/adios2/helper/adiosKokkos.cpp

@@ -0,0 +1,117 @@
+/*
+ * Distributed under the OSI-approved Apache License, Version 2.0.  See
+ * accompanying file Copyright.txt for details.
+ */
+
+#ifndef ADIOS2_HELPER_ADIOSKokkos_CPP_
+#define ADIOS2_HELPER_ADIOSKokkos_CPP_
+
+#include "adiosKokkos.h"
+#include "adios2/common/ADIOSMacros.h"
+
+#include <Kokkos_Core.hpp>
+
+namespace
+{
+template <class MemSpace>
+void KokkosDeepCopy(const char *src, char *dst, size_t byteCount)
+{
+    Kokkos::View<const char *, MemSpace,
+                 Kokkos::MemoryTraits<Kokkos::Unmanaged>>
+        srcView(src, byteCount);
+    Kokkos::View<char *, Kokkos::HostSpace,
+                 Kokkos::MemoryTraits<Kokkos::Unmanaged>>
+        dstView(dst, byteCount);
+    Kokkos::deep_copy(dstView, srcView);
+}
+
+template <class T>
+void KokkosMinMaxImpl(const T *data, const size_t size, T &min, T &max)
+{
+    Kokkos::parallel_reduce(size,
+                            KOKKOS_LAMBDA(int i, T &lmax, T &lmin) {
+                                if (lmax < data[i])
+                                    lmax = data[i];
+                                if (lmin > data[i])
+                                    lmin = data[i];
+                            },
+                            Kokkos::Max<T>(max), Kokkos::Min<T>(min));
+}
+
+// types non supported on the device
+void KokkosMinMaxImpl(const char * /*values*/, const size_t /*size*/,
+                      char & /*min*/, char & /*max*/)
+{
+}
+void KokkosMinMaxImpl(const std::complex<float> * /*values*/,
+                      const size_t /*size*/, std::complex<float> & /*min*/,
+                      std::complex<float> & /*max*/)
+{
+}
+void KokkosMinMaxImpl(const std::complex<double> * /*values*/,
+                      const size_t /*size*/, std::complex<double> & /*min*/,
+                      std::complex<double> & /*max*/)
+{
+}
+
+}
+
+namespace adios2
+{
+namespace helper
+{
+void MemcpyGPUToBuffer(char *dst, const char *GPUbuffer, size_t byteCount)
+{
+#ifdef KOKKOS_ENABLE_CUDA
+    KokkosDeepCopy<Kokkos::CudaSpace>(GPUbuffer, dst, byteCount);
+#endif
+#ifdef KOKKOS_ENABLE_HIP
+    KokkosDeepCopy<Kokkos::Experimental::HIPSpace>(GPUbuffer, dst, byteCount);
+#endif
+}
+
+void MemcpyBufferToGPU(char *GPUbuffer, const char *src, size_t byteCount)
+{
+#ifdef KOKKOS_ENABLE_CUDA
+    KokkosDeepCopy<Kokkos::CudaSpace>(src, GPUbuffer, byteCount);
+#endif
+#ifdef KOKKOS_ENABLE_HIP
+    KokkosDeepCopy<Kokkos::Experimental::HIPSpace>(src, GPUbuffer, byteCount);
+#endif
+}
+
+bool IsGPUbuffer(const void *ptr)
+{
+#ifdef KOKKOS_ENABLE_CUDA
+    cudaPointerAttributes attr;
+    cudaPointerGetAttributes(&attr, ptr);
+    if (attr.type == cudaMemoryTypeDevice)
+    {
+        return true;
+    }
+#endif
+    return false;
+}
+
+void KokkosFinalize() { Kokkos::finalize(); }
+
+void KokkosInit() { Kokkos::initialize(); }
+
+bool KokkosIsInitialized() { return Kokkos::is_initialized(); }
+
+template <class T>
+void GPUMinMax(const T *values, const size_t size, T &min, T &max)
+{
+    KokkosMinMaxImpl(values, size, min, max);
+}
+
+}
+}
+
+#define declare_type(T)                                                        \
+    template void adios2::helper::GPUMinMax(                                   \
+        const T *values, const size_t size, T &min, T &max);
+ADIOS2_FOREACH_PRIMITIVE_STDTYPE_1ARG(declare_type)
+#undef declare_type
+
+#endif /* ADIOS2_HELPER_ADIOSKokkos_CPP_ */

source/adios2/helper/adiosKokkos.h

@@ -0,0 +1,35 @@
+/*
+ * Distributed under the OSI-approved Apache License, Version 2.0.  See
+ * accompanying file Copyright.txt for details.
+ */
+
+#ifndef ADIOS2_HELPER_ADIOSKOKKOS_H_
+#define ADIOS2_HELPER_ADIOSKOKKOS_H_
+
+#include <cstddef>
+
+namespace adios2
+{
+namespace helper
+{
+/*
+ * CUDA kernel for computing the min and max from a GPU buffer
+ */
+template <class T>
+void GPUMinMax(const T *values, const size_t size, T &min, T &max);
+
+/**
+ * Wrapper around cudaMemcpy needed for isolating CUDA interface dependency
+ */
+void MemcpyGPUToBuffer(char *dst, const char *GPUbuffer, size_t byteCount);
+void MemcpyBufferToGPU(char *GPUbuffer, const char *src, size_t byteCount);
+
+void KokkosFinalize();
+void KokkosInit();
+bool KokkosIsInitialized();
+bool IsGPUbuffer(const void *ptr);
+
+} // helper
+} // adios2
+
+#endif /* ADIOS2_HELPER_ADIOSKOKKOS_H_ */