Merge pull request #3918 from anagainaru/sst-with-kokkos

SST with GPU support

examples/hello/CMakeLists.txt

@@ -25,6 +25,9 @@ if(ADIOS2_HAVE_Kokkos_HIP)
 endif()
 if(ADIOS2_HAVE_Kokkos)
   add_subdirectory(bpStepsWriteReadKokkos)
+  if(ADIOS2_HAVE_SST)
+    add_subdirectory(sstKokkos)
+  endif()
 endif()
 
 add_subdirectory(bpThreadWrite)

examples/hello/sstKokkos/CMakeLists.txt

@@ -0,0 +1,37 @@
+#-- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- #
+#Distributed under the OSI - approved Apache License, Version 2.0. See
+#accompanying file Copyright.txt for details.
+#-- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- #
+
+cmake_minimum_required(VERSION 3.12)
+project(ADIOS2HelloSSTKokkosExample)
+
+#CXX Compiler settings only in for this example
+set(CMAKE_CXX_STANDARD 17)
+set(CMAKE_CXX_STANDARD_REQUIRED ON)
+
+if (NOT TARGET adios2_core)
+  set(_components CXX)
+
+  find_package(Kokkos 3.7 QUIET)
+  if (Kokkos_FOUND AND DEFINED Kokkos_CXX_COMPILER)
+      set(CMAKE_CXX_COMPILER "${Kokkos_CXX_COMPILER}")
+  endif()
+
+  find_package(ADIOS2 REQUIRED COMPONENTS ${_components})
+else()
+  if (DEFINED Kokkos_CXX_COMPILER)
+      set(CMAKE_CXX_COMPILER "${Kokkos_CXX_COMPILER}")
+  endif()
+endif()
+
+if (ADIOS2_HAVE_Kokkos)
+  add_executable(adios2_hello_sstWriterKokkos sstWriterKokkos.cpp)
+  add_executable(adios2_hello_sstReaderKokkos sstReaderKokkos.cpp)
+  kokkos_compilation(SOURCE sstWriterKokkos.cpp)
+  kokkos_compilation(SOURCE sstReaderKokkos.cpp)
+  target_link_libraries(adios2_hello_sstWriterKokkos adios2::cxx11 Kokkos::kokkos)
+  install(TARGETS adios2_hello_sstWriterKokkos RUNTIME DESTINATION ${CMAKE_INSTALL_BINDIR})
+  target_link_libraries(adios2_hello_sstReaderKokkos adios2::cxx11 Kokkos::kokkos)
+  install(TARGETS adios2_hello_sstReaderKokkos RUNTIME DESTINATION ${CMAKE_INSTALL_BINDIR})
+endif()

examples/hello/sstKokkos/sstReaderKokkos.cpp

@@ -0,0 +1,88 @@
+/*
+ * Distributed under the OSI-approved Apache License, Version 2.0.  See
+ * accompanying file Copyright.txt for details.
+ *
+ * sstReaderKokkos.cpp  Simple example of reading bpFloats through ADIOS2 SST
+ * engine with multiple simulations steps for every IO step using Kokkos
+ */
+#include <ios>
+#include <iostream>
+#include <vector>
+
+#include <adios2.h>
+#include <adios2/cxx11/KokkosView.h>
+
+#include <Kokkos_Core.hpp>
+
+template <class MemSpace, class ExecSpace>
+int BPRead(adios2::ADIOS &adios, const std::string fname, const size_t Nx, const size_t Ny,
+           const size_t nSteps, const std::string engine)
+{
+    adios2::IO io = adios.DeclareIO("ReadIO");
+    io.SetEngine(engine);
+
+    ExecSpace exe_space;
+    std::cout << "Read on memory space: " << exe_space.name() << std::endl;
+
+    adios2::Engine bpReader = io.Open(fname, adios2::Mode::Read);
+
+    unsigned int step = 0;
+    bool correctValues = true;
+    Kokkos::View<float **, MemSpace> gpuSimData("simBuffer", Nx, Ny);
+    for (; bpReader.BeginStep() == adios2::StepStatus::OK; ++step)
+    {
+        auto data = io.InquireVariable<float>("bpFloats");
+        const adios2::Dims start{0, 0};
+        const adios2::Dims count{Nx, Ny};
+        const adios2::Box<adios2::Dims> sel(start, count);
+        data.SetSelection(sel);
+
+        // var.SetMemorySpace(adios2::MemorySpace::GPU);
+        bpReader.Get(data, gpuSimData);
+        bpReader.EndStep();
+
+        auto cpuData = Kokkos::create_mirror_view_and_copy(Kokkos::HostSpace{}, gpuSimData);
+        if (cpuData(0, 0) != step * 10)
+        {
+            std::cout << "Value mismatch at step " << step << std::endl;
+            correctValues = false;
+            break;
+        }
+    }
+    if (correctValues)
+        std::cout << "Read " << step << " steps successfully" << std::endl;
+
+    bpReader.Close();
+    return 0;
+}
+
+int main(int argc, char **argv)
+{
+    const std::string engine = argv[1] ? argv[1] : "SST";
+    std::cout << "Using engine " << engine << std::endl;
+    const size_t Nx = 600, Ny = 100, nSteps = 2;
+    const std::string memorySpace = "Device";
+
+    const std::string filename = engine + "StepsWriteReadKokkos";
+    Kokkos::initialize(argc, argv);
+    {
+        adios2::ADIOS adios;
+
+        std::cout << "Using engine " << engine << std::endl;
+        if (memorySpace == "Device")
+        {
+            using mem_space = Kokkos::DefaultExecutionSpace::memory_space;
+            std::cout << "Memory space: DefaultMemorySpace" << std::endl;
+            BPRead<mem_space, Kokkos::DefaultExecutionSpace>(adios, filename + "_DD.bp", Nx, Ny,
+                                                             nSteps, engine);
+        }
+        else
+        {
+            std::cout << "Memory space: HostSpace" << std::endl;
+            BPRead<Kokkos::HostSpace, Kokkos::Serial>(adios, filename + "_HH.bp", Nx, Ny, nSteps,
+                                                      engine);
+        }
+    }
+    Kokkos::finalize();
+    return 0;
+}

examples/hello/sstKokkos/sstWriterKokkos.cpp

@@ -0,0 +1,96 @@
+/*
+ * Distributed under the OSI-approved Apache License, Version 2.0.  See
+ * accompanying file Copyright.txt for details.
+ *
+ * sstWriterKokkos.cpp  Simple example of writing bpFloats through ADIOS2 SST
+ * engine with multiple simulations steps for every IO step using Kokkos
+ */
+#include <ios>
+#include <iostream>
+#include <vector>
+
+#include <adios2.h>
+#include <adios2/cxx11/KokkosView.h>
+
+#include <Kokkos_Core.hpp>
+
+template <class MemSpace, class ExecSpace>
+int BPWrite(adios2::ADIOS &adios, const std::string fname, const size_t Nx, const size_t Ny,
+            const size_t nSteps, const std::string engine)
+{
+    // Initialize the simulation data
+    Kokkos::View<float **, MemSpace> gpuSimData("simBuffer", Nx, Ny);
+    static_assert(Kokkos::SpaceAccessibility<ExecSpace, MemSpace>::accessible, "");
+    Kokkos::parallel_for(
+        "initBuffer", Kokkos::RangePolicy<ExecSpace>(0, Nx), KOKKOS_LAMBDA(int i) {
+            for (int j = 0; j < Ny; j++)
+                gpuSimData(i, j) = static_cast<float>(i);
+        });
+    Kokkos::fence();
+
+    adios2::IO io = adios.DeclareIO("WriteIO");
+    io.SetEngine(engine);
+
+    const adios2::Dims shape{Nx, Ny};
+    const adios2::Dims start{0, 0};
+    const adios2::Dims count{Nx, Ny};
+    auto data = io.DefineVariable<float>("bpFloats", shape, start, count);
+
+    adios2::Engine bpWriter = io.Open(fname, adios2::Mode::Write);
+
+    // Simulation steps
+    for (int step = 0; step < nSteps; ++step)
+    {
+        adios2::Box<adios2::Dims> sel({0, 0}, {Nx, Ny});
+        data.SetSelection(sel);
+
+        bpWriter.BeginStep();
+        // var.SetMemorySpace(adios2::MemorySpace::GPU);
+        bpWriter.Put(data, gpuSimData);
+        bpWriter.EndStep();
+
+        // Update values in the simulation data
+        Kokkos::parallel_for(
+            "updateBuffer", Kokkos::RangePolicy<ExecSpace>(0, Nx), KOKKOS_LAMBDA(int i) {
+                for (int j = 0; j < Ny; j++)
+                    gpuSimData(i, j) += 10;
+            });
+        Kokkos::fence();
+    }
+
+    bpWriter.Close();
+    ExecSpace exe_space;
+    std::cout << "Done writing on memory space: " << exe_space.name() << std::endl;
+    return 0;
+}
+
+int main(int argc, char **argv)
+{
+    const std::string engine = argv[1] ? argv[1] : "SST";
+    std::cout << "Using engine " << engine << std::endl;
+    const size_t Nx = 600, Ny = 100, nSteps = 2;
+    const std::string memorySpace = "Device";
+
+    const std::string filename = engine + "StepsWriteReadKokkos";
+    Kokkos::initialize(argc, argv);
+    {
+        adios2::ADIOS adios;
+
+        std::cout << "Using engine " << engine << std::endl;
+        if (memorySpace == "Device")
+        {
+            using mem_space = Kokkos::DefaultExecutionSpace::memory_space;
+            std::cout << "Memory space: DefaultMemorySpace" << std::endl;
+            BPWrite<mem_space, Kokkos::DefaultExecutionSpace>(adios, filename + "_DD.bp", Nx, Ny,
+                                                              nSteps, engine);
+        }
+        else
+        {
+            std::cout << "Memory space: HostSpace" << std::endl;
+            BPWrite<Kokkos::HostSpace, Kokkos::Serial>(adios, filename + "_HH.bp", Nx, Ny, nSteps,
+                                                       engine);
+        }
+    }
+    Kokkos::finalize();
+    return 0;
+}

source/adios2/toolkit/format/buffer/malloc/MallocV.cpp

@@ -7,6 +7,7 @@
  */
 
 #include "MallocV.h"
+#include "adios2/helper/adiosFunctions.h"
 #include "adios2/toolkit/format/buffer/BufferV.h"
 
 #include <algorithm>
@@ -74,7 +75,12 @@ size_t MallocV::AddToVec(const size_t size, const void *buf, size_t align, bool
             m_InternalBlock = (char *)realloc(m_InternalBlock, NewSize);
             m_AllocatedSize = NewSize;
         }
-        memcpy(m_InternalBlock + m_internalPos, buf, size);
+#ifdef ADIOS2_HAVE_GPU_SUPPORT
+        if (MemSpace == MemorySpace::GPU)
+            helper::CopyFromGPUToBuffer(m_InternalBlock, m_internalPos, buf, MemSpace, size);
+#endif
+        if (MemSpace == MemorySpace::Host)
+            memcpy(m_InternalBlock + m_internalPos, buf, size);
 
         if (DataV.size() && !DataV.back().External &&
             (m_internalPos == (DataV.back().Offset + DataV.back().Size)))