Merge pull request ornladios#3446 from anagainaru/gpu-reorg-kokkos

Kokkos backend in ADIOS2
vicentebolea · Feb 21, 2023 · 4511299 · 4511299
2 parents a5d7972 + f30714f
commit 4511299
Show file tree

Hide file tree

Showing 43 changed files with 698 additions and 163 deletions.
diff --git a/.github/workflows/everything.yml b/.github/workflows/everything.yml
@@ -97,7 +97,7 @@ jobs:
       image: ornladios/adios2:ci-spack-el8-${{ matrix.compiler }}-${{ matrix.parallel }}
       options: --shm-size=1g
       env:
-        GH_YML_JOBNAME: ${{ matrix.os }}-${{ matrix.compiler }}-${{ matrix.parallel }}
+        GH_YML_JOBNAME: ${{ matrix.os }}-${{ matrix.gpu_backend }}${{ matrix.compiler }}-${{ matrix.parallel }}
         GH_YML_BASE_OS: Linux
         GH_YML_MATRIX_OS: ${{ matrix.os }}
         GH_YML_MATRIX_COMPILER: ${{ matrix.compiler }}
@@ -114,6 +114,11 @@ jobs:
             compiler: cuda
             parallel: serial
             constrains: build_only
+          - os: el8
+            compiler: cuda
+            parallel: serial
+            gpu_backend: kokkos
+            constrains: build_only
           - os: el8
             compiler: gcc10
             parallel: mpich

diff --git a/CMakeLists.txt b/CMakeLists.txt
@@ -146,7 +146,8 @@ adios_option(SZ         "Enable support for SZ transforms" AUTO)
 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(CUDA       "Enable support for Cuda" OFF)
+adios_option(Kokkos     "Enable support for Kokkos" OFF)
 adios_option(MPI        "Enable support for MPI" AUTO)
 adios_option(DAOS       "Enable support for DAOS" AUTO)
 adios_option(DataMan    "Enable support for DataMan" AUTO)
@@ -226,9 +227,9 @@ endif()
 
 
 set(ADIOS2_CONFIG_OPTS
-    BP5 DataMan DataSpaces HDF5 HDF5_VOL MHS SST CUDA Fortran MPI Python Blosc2 
+    BP5 DataMan DataSpaces HDF5 HDF5_VOL MHS SST Fortran MPI Python Blosc2
     BZip2 LIBPRESSIO MGARD PNG SZ ZFP DAOS IME O_DIRECT Sodium Catalyst SysVShMem UCX ZeroMQ 
-    Profiling Endian_Reverse GPU_Support AWSSDK
+    Profiling Endian_Reverse AWSSDK GPU_Support CUDA Kokkos Kokkos_CUDA Kokkos_HIP
 )
 
 GenerateADIOSHeaderConfig(${ADIOS2_CONFIG_OPTS})

diff --git a/CTestCustom.cmake.in b/CTestCustom.cmake.in
@@ -21,6 +21,7 @@ list(APPEND CTEST_CUSTOM_WARNING_EXCEPTION
   "warnings generated"
   "warning: template parameter ... is not used in declaring the parameter types of function template"
   "warning: command-line option '.*' is valid for Fortran but not for C"
+  "Warning #20208-D: '.*' is treated as '.*' in device code"
 )
 list(APPEND CTEST_CUSTOM_COVERAGE_EXCLUDE
   ".*/thirdparty/.*"

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

diff --git a/cmake/DetectOptions.cmake b/cmake/DetectOptions.cmake
@@ -170,7 +170,34 @@ endif()
 
 set(mpi_find_components C)
 
-# Cuda
+if(ADIOS_USE_Kokkos AND ADIOS_USE_CUDA)
+  message(FATAL_ERROR "ADIOS2_USE_Kokkos is incompatible with ADIOS_USE_CUDA")
+endif()
+
+# Kokkos
+if(ADIOS2_USE_Kokkos)
+  if(ADIOS2_USE_Kokkos STREQUAL AUTO)
+    find_package(Kokkos 3.7 QUIET)
+  else()
+    find_package(Kokkos 3.7 REQUIRED)
+  endif()
+  if(Kokkos_FOUND)
+    set(ADIOS2_HAVE_Kokkos TRUE)
+    if(Kokkos_ENABLE_CUDA OR Kokkos_ENABLE_HIP OR Kokkos_ENABLE_SYCL)
+      if(Kokkos_ENABLE_CUDA)
+        set(ADIOS2_HAVE_Kokkos_CUDA TRUE)
+        enable_language(CUDA)
+      endif()
+      if(Kokkos_ENABLE_HIP)
+        set(ADIOS2_HAVE_Kokkos_HIP TRUE)
+        enable_language(HIP)
+      endif()
+      set(ADIOS2_HAVE_GPU_Support TRUE)
+    endif()
+  endif()
+endif()
+
+# CUDA
 if(ADIOS2_USE_CUDA)
   include(CheckLanguage)
   check_language(CUDA)
@@ -179,11 +206,15 @@ if(ADIOS2_USE_CUDA)
   else()
     find_package(CUDAToolkit REQUIRED)
   endif()
+  if(CMAKE_CUDA_COMPILER AND CUDAToolkit_FOUND)
+    enable_language(CUDA)
+    set(ADIOS2_HAVE_CUDA TRUE)
+    set(ADIOS2_HAVE_GPU_Support TRUE)
+  endif()
 endif()
-if(CMAKE_CUDA_COMPILER AND CUDAToolkit_FOUND)
-  enable_language(CUDA)
-  set(ADIOS2_HAVE_CUDA TRUE)
-  set(ADIOS2_HAVE_GPU_Support TRUE)
+
+if(ADIOS_HAVE_Kokkos AND ADIOS_HAVE_CUDA)
+    message(FATAL_ERROR "The Kokkos and CUDA backends cannot be active concurrently")
 endif()
 
 # Fortran

diff --git a/examples/CMakeLists.txt b/examples/CMakeLists.txt
@@ -11,11 +11,6 @@ add_subdirectory(inlineMWE)
 add_subdirectory(plugins)
 add_subdirectory(fides)
 
-find_package(Kokkos QUIET)
-if(Kokkos_FOUND)
-  add_subdirectory(kokkos)
-endif()
-
 if(ADIOS2_HAVE_MPI)
   add_subdirectory(heatTransfer)
 endif()
@@ -24,6 +19,15 @@ if(ADIOS2_BUILD_EXAMPLES_EXPERIMENTAL)
   add_subdirectory(experimental)
 endif()
 
-if(ADIOS2_HAVE_CUDA)
+if(ADIOS2_HAVE_CUDA OR ADIOS2_HAVE_Kokkos_CUDA)
   add_subdirectory(cuda)
 endif()
+
+if(ADIOS2_HAVE_Kokkos)
+  add_subdirectory(kokkos)
+elseif(ADIOS2_HAVE_CUDA)
+  find_package(Kokkos QUIET)
+  if(Kokkos_FOUND)
+    add_subdirectory(kokkos)
+  endif()
+endif()
diff --git a/examples/cuda/cudaBP4WriteRead.cu b/examples/cuda/cudaBP4WriteRead.cu
@@ -44,7 +44,7 @@ int BPWrite(const std::string fname, const size_t N, int nSteps)
 
         // Start IO step every write step
         bpWriter.BeginStep();
-        data.SetMemorySpace(adios2::MemorySpace::CUDA);
+        data.SetMemorySpace(adios2::MemorySpace::GPU);
         bpWriter.Put(data, gpuSimData);
         bpWriter.EndStep();
 
@@ -85,7 +85,7 @@ int BPRead(const std::string fname, const size_t N, int nSteps)
     for (size_t step = 0; step < write_step; step++)
     {
         data.SetStepSelection({step, 1});
-        data.SetMemorySpace(adios2::MemorySpace::CUDA);
+        data.SetMemorySpace(adios2::MemorySpace::GPU);
         bpReader.Get(data, gpuSimData, adios2::Mode::Deferred);
         bpReader.PerformGets();
         cudaMemcpy(simData.data(), gpuSimData, N * sizeof(float),

diff --git a/examples/cuda/cudaBP5WriteRead.cu b/examples/cuda/cudaBP5WriteRead.cu
@@ -44,7 +44,7 @@ int BPWrite(const std::string fname, const size_t N, int nSteps)
 
         // Start IO step every write step
         bpWriter.BeginStep();
-        data.SetMemorySpace(adios2::MemorySpace::CUDA);
+        data.SetMemorySpace(adios2::MemorySpace::GPU);
         bpWriter.Put(data, gpuSimData);
         bpWriter.EndStep();
 
@@ -78,7 +78,7 @@ int BPRead(const std::string fname, const size_t N, int nSteps)
         const adios2::Box<adios2::Dims> sel(start, count);
         data.SetSelection(sel);
 
-        data.SetMemorySpace(adios2::MemorySpace::CUDA);
+        data.SetMemorySpace(adios2::MemorySpace::GPU);
         bpReader.Get(data, gpuSimData); //, adios2::Mode::Deferred);
         bpReader.EndStep();
         cudaMemcpy(simData.data(), gpuSimData, N * sizeof(float),

diff --git a/examples/kokkos/CMakeLists.txt b/examples/kokkos/CMakeLists.txt
@@ -5,3 +5,9 @@
 
 add_executable(KokkosWriteRead kokkosWriteRead.cpp)
 target_link_libraries(KokkosWriteRead PUBLIC adios2::cxx11 Kokkos::kokkos)
+
+if(Kokkos_ENABLE_CUDA)
+  set_property(TARGET KokkosWriteRead PROPERTY CUDA_STANDARD 17)
+  set_property(SOURCE kokkosWriteRead.cpp PROPERTY LANGUAGE CUDA)
+  set_property(SOURCE kokkosWriteRead.cpp APPEND PROPERTY COMPILE_FLAGS "--extended-lambda")
+endif()
diff --git a/scripts/ci/cmake-v2/ci-el8-kokkoscuda-serial.cmake b/scripts/ci/cmake-v2/ci-el8-kokkoscuda-serial.cmake
@@ -0,0 +1,26 @@
+# Client maintainer: vicente.bolea@kitware.com
+
+set(ENV{CC}  gcc)
+set(ENV{CXX} g++)
+set(ENV{FC}  gfortran)
+
+set(dashboard_cache "
+ADIOS2_USE_BZip2:BOOL=ON
+ADIOS2_USE_Blosc:BOOL=ON
+ADIOS2_USE_DataMan:BOOL=ON
+ADIOS2_USE_Fortran:BOOL=ON
+ADIOS2_USE_HDF5:BOOL=ON
+ADIOS2_USE_Python:BOOL=ON
+ADIOS2_USE_SZ:BOOL=ON
+ADIOS2_USE_ZeroMQ:STRING=ON
+ADIOS2_USE_ZFP:BOOL=ON
+ADIOS2_USE_Kokkos:BOOL=ON
+ADIOS2_USE_MPI:BOOL=OFF
+CMAKE_C_FLAGS:STRING=-Wall
+CMAKE_CXX_FLAGS:STRING=-Wall
+CMAKE_Fortran_FLAGS:STRING=-Wall
+")
+
+set(CTEST_CMAKE_GENERATOR "Ninja")
+list(APPEND CTEST_UPDATE_NOTES_FILES "${CMAKE_CURRENT_LIST_FILE}")
+include(${CMAKE_CURRENT_LIST_DIR}/ci-common.cmake)
diff --git a/source/adios2/CMakeLists.txt b/source/adios2/CMakeLists.txt
@@ -124,6 +124,33 @@ if(ADIOS2_HAVE_CUDA)
   set(maybe_adios2_core_cuda adios2_core_cuda)
 endif()
 
+set(maybe_adios2_core_kokkos)
+if(ADIOS2_HAVE_Kokkos)
+  add_library(adios2_core_kokkos helper/adiosKokkos.h helper/adiosKokkos.cpp)
+
+  set_target_properties(adios2_core_kokkos PROPERTIES
+    VISIBILITY_INLINES_HIDDEN ON
+    INCLUDE_DIRECTORIES "$<BUILD_INTERFACE:${ADIOS2_SOURCE_DIR}/source>;$<BUILD_INTERFACE:${ADIOS2_BINARY_DIR}/source>"
+    EXPORT_NAME core_kokkos
+    OUTPUT_NAME adios2${ADIOS2_LIBRARY_SUFFIX}_core_kokkos
+    )
+
+  kokkos_compilation(SOURCE helper/adiosKokkos.cpp)
+  if(Kokkos_ENABLE_CUDA)
+    set_property(SOURCE helper/adiosKokkos.cpp PROPERTY LANGUAGE CUDA)
+    set_property(SOURCE helper/adiosKokkos.cpp APPEND PROPERTY COMPILE_FLAGS "--extended-lambda")
+    set_target_properties(adios2_core_kokkos PROPERTIES
+      CUDA_VISIBILITY_PRESET hidden
+      )
+    target_compile_features(adios2_core_kokkos PRIVATE cuda_std_17)
+  endif()
+
+  target_link_libraries(adios2_core_kokkos PRIVATE Kokkos::kokkos)
+  target_link_libraries(adios2_core PRIVATE adios2_core_kokkos)
+
+  set(maybe_adios2_core_kokkos adios2_core_kokkos)
+endif()
+
 target_include_directories(adios2_core
   PUBLIC
     $<BUILD_INTERFACE:${ADIOS2_SOURCE_DIR}/source>
@@ -412,7 +439,7 @@ install(DIRECTORY toolkit/
 )
 
 # Library installation
-install(TARGETS adios2_core ${maybe_adios2_core_mpi} ${maybe_adios2_core_cuda} EXPORT adios2Exports
+install(TARGETS adios2_core ${maybe_adios2_core_mpi} ${maybe_adios2_core_cuda} ${maybe_adios2_core_kokkos} EXPORT adios2Exports
   RUNTIME DESTINATION ${CMAKE_INSTALL_BINDIR} COMPONENT adios2_core-runtime
   LIBRARY DESTINATION ${CMAKE_INSTALL_LIBDIR} COMPONENT adios2_core-libraries NAMELINK_COMPONENT adios2_core-development
   ARCHIVE DESTINATION ${CMAKE_INSTALL_LIBDIR} COMPONENT adios2_core-development

diff --git a/source/adios2/common/ADIOSTypes.h b/source/adios2/common/ADIOSTypes.h
@@ -39,8 +39,8 @@ enum class MemorySpace
     Detect, ///< Detect the memory space automatically
 #endif
     Host, ///< Host memory space
-#ifdef ADIOS2_HAVE_CUDA
-    CUDA ///< CUDA memory spaces
+#ifdef ADIOS2_HAVE_GPU_SUPPORT
+    GPU ///< GPU memory space
 #endif
 };
 

diff --git a/source/adios2/core/ADIOS.cpp b/source/adios2/core/ADIOS.cpp
@@ -85,6 +85,19 @@ class ADIOS::GlobalServices
     bool isAWSInitialized = false;
 #endif
 
+#ifdef ADIOS2_HAVE_KOKKOS
+    void Init_Kokkos_API()
+    {
+        if (isKokkosInitialized)
+            return;
+        if (helper::KokkosIsInitialized())
+            return;
+        helper::KokkosInit();
+        std::atexit(helper::KokkosFinalize);
+        isKokkosInitialized = true;
+    }
+    bool isKokkosInitialized = false;
+#endif
     bool wasGlobalShutdown = false;
 };
 
@@ -129,6 +142,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)

diff --git a/source/adios2/core/VariableBase.cpp b/source/adios2/core/VariableBase.cpp
@@ -24,9 +24,7 @@
 #include "adios2/helper/adiosString.h"
 #include "adios2/operator/OperatorFactory.h"
 
-#ifdef ADIOS2_HAVE_CUDA
-#include <cuda_runtime.h>
-#endif
+#include "adios2/helper/adiosGPUFunctions.h"
 
 namespace adios2
 {
@@ -55,14 +53,10 @@ MemorySpace VariableBase::GetMemorySpace(const void *ptr)
     {
         return m_MemSpace;
     }
-#endif
 
-#ifdef ADIOS2_HAVE_CUDA
-    cudaPointerAttributes attr;
-    cudaPointerGetAttributes(&attr, ptr);
-    if (attr.type == cudaMemoryTypeDevice)
+    if (helper::IsGPUbuffer(ptr))
     {
-        return MemorySpace::CUDA;
+        return MemorySpace::GPU;
     }
 #endif
     return MemorySpace::Host;

diff --git a/source/adios2/engine/bp5/BP5Writer.cpp b/source/adios2/engine/bp5/BP5Writer.cpp
@@ -1756,6 +1756,11 @@ void BP5Writer::PutCommon(VariableBase &variable, const void *values, bool sync)
     {
         Count = variable.m_Count.data();
     }
+    else if (variable.m_ShapeID == ShapeID::JoinedArray)
+    {
+        Count = variable.m_Count.data();
+        Shape = variable.m_Shape.data();
+    }
 
     size_t ObjSize;
     if (variable.m_Type == DataType::Struct)

diff --git a/source/adios2/engine/bp5/BP5Writer.tcc b/source/adios2/engine/bp5/BP5Writer.tcc
@@ -41,6 +41,12 @@ void BP5Writer::PutCommonSpan(Variable<T> &variable,
         Start = variable.m_Start.data();
         Count = variable.m_Count.data();
     }
+    else if (variable.m_ShapeID == ShapeID::JoinedArray)
+    {
+        Shape = variable.m_Shape.data();
+        DimCount = variable.m_Count.size();
+        Count = variable.m_Count.data();
+    }
     else if (variable.m_ShapeID == ShapeID::LocalArray)
     {
         DimCount = variable.m_Count.size();

diff --git a/source/adios2/engine/sst/SstWriter.tcc b/source/adios2/engine/sst/SstWriter.tcc
@@ -52,6 +52,12 @@ void SstWriter::PutSyncCommon(Variable<T> &variable, const T *values)
             Start = variable.m_Start.data();
             Count = variable.m_Count.data();
         }
+        else if (variable.m_ShapeID == ShapeID::JoinedArray)
+        {
+            DimCount = variable.m_Shape.size();
+            Shape = variable.m_Shape.data();
+            Count = variable.m_Count.data();
+        }
         else if (variable.m_ShapeID == ShapeID::LocalArray)
         {
             DimCount = variable.m_Count.size();