Example using DataMan with Kokkos buffers

ornladios · Nov 18, 2023 · beabde3 · beabde3
1 parent 4d2ac4b
commit beabde3
Show file tree

Hide file tree

Showing 4 changed files with 209 additions and 0 deletions.
diff --git a/examples/hello/CMakeLists.txt b/examples/hello/CMakeLists.txt
@@ -28,6 +28,9 @@ if(ADIOS2_HAVE_Kokkos)
   if(ADIOS2_HAVE_SST)
     add_subdirectory(sstKokkos)
   endif()
+  if(ADIOS2_HAVE_DataMan)
+    add_subdirectory(datamanKokkos)
+  endif()
 endif()
 
 add_subdirectory(bpThreadWrite)

diff --git a/examples/hello/datamanKokkos/CMakeLists.txt b/examples/hello/datamanKokkos/CMakeLists.txt
@@ -0,0 +1,33 @@
+#------------------------------------------------------------------------------#
+# Distributed under the OSI-approved Apache License, Version 2.0.  See
+# accompanying file Copyright.txt for details.
+#------------------------------------------------------------------------------#
+
+cmake_minimum_required(VERSION 3.12)
+project(ADIOS2HelloDataManKokkosExample)
+
+if(NOT TARGET adios2_core)
+  set(_components CXX)
+
+  find_package(MPI COMPONENTS C)
+  if(MPI_FOUND)
+    # Workaround for various MPI implementations forcing the link of C++ bindings
+    add_definitions(-DOMPI_SKIP_MPICXX -DMPICH_SKIP_MPICXX)
+
+    list(APPEND _components MPI)
+  endif()
+
+  find_package(ZeroMQ 4.1 QUIET)
+
+  find_package(ADIOS2 REQUIRED COMPONENTS ${_components})
+endif()
+
+if(ADIOS2_HAVE_MPI AND ADIOS2_HAVE_DataMan)
+  add_executable(adios2_hello_datamanWriterKokkos dataManWriterKokkos.cpp)
+  target_link_libraries(adios2_hello_datamanWriterKokkos adios2::cxx11_mpi MPI::MPI_C Kokkos::kokkos)
+  install(TARGETS adios2_hello_datamanWriterKokkos RUNTIME DESTINATION ${CMAKE_INSTALL_BINDIR})
+
+  add_executable(adios2_hello_datamanReaderKokkos dataManReaderKokkos.cpp)
+  target_link_libraries(adios2_hello_datamanReaderKokkos adios2::cxx11_mpi MPI::MPI_C Kokkos::kokkos)
+  install(TARGETS adios2_hello_datamanReaderKokkos RUNTIME DESTINATION ${CMAKE_INSTALL_BINDIR})
+endif()
diff --git a/examples/hello/datamanKokkos/dataManReaderKokkos.cpp b/examples/hello/datamanKokkos/dataManReaderKokkos.cpp
@@ -0,0 +1,76 @@
+#include <adios2.h>
+#include <chrono>
+#include <iostream>
+#include <mpi.h>
+#include <numeric>
+#include <thread>
+#include <vector>
+
+#include <adios2/cxx11/KokkosView.h>
+
+#include <Kokkos_Core.hpp>
+
+int mpiRank, mpiSize;
+
+template <class T, class MemSpace>
+void PrintData(Kokkos::View<T *, MemSpace> &gpuData, const size_t step)
+{
+    auto data = Kokkos::create_mirror_view_and_copy(Kokkos::HostSpace{}, gpuData);
+    std::cout << "Rank: " << mpiRank << " Step: " << step << " [";
+    for (int i = 0; i < data.extent_int(0); ++i)
+    {
+        std::cout << data(i) << " ";
+    }
+    std::cout << "]" << std::endl;
+}
+
+int main(int argc, char *argv[])
+{
+    // initialize MPI
+    MPI_Init(&argc, &argv);
+    MPI_Comm_rank(MPI_COMM_WORLD, &mpiRank);
+    MPI_Comm_size(MPI_COMM_WORLD, &mpiSize);
+
+    // initialize adios2
+    adios2::ADIOS adios(MPI_COMM_WORLD);
+    adios2::IO dataManIO = adios.DeclareIO("whatever");
+    dataManIO.SetEngine("DataMan");
+    dataManIO.SetParameters({{"IPAddress", "127.0.0.1"}, {"Port", "12306"}, {"Timeout", "5"}});
+
+    // open stream
+    adios2::Engine dataManReader = dataManIO.Open("HelloDataMan", adios2::Mode::Read);
+
+    // define variable
+    adios2::Variable<float> floatArrayVar;
+
+    Kokkos::DefaultExecutionSpace exe_space;
+    std::cout << "Read on memory space: " << exe_space.name() << std::endl;
+    // read data
+    while (true)
+    {
+        auto status = dataManReader.BeginStep();
+        if (status == adios2::StepStatus::OK)
+        {
+            floatArrayVar = dataManIO.InquireVariable<float>("FloatArray");
+            auto shape = floatArrayVar.Shape();
+            size_t datasize =
+                std::accumulate(shape.begin(), shape.end(), 1, std::multiplies<size_t>());
+            Kokkos::View<float *, Kokkos::DefaultExecutionSpace::memory_space> floatVector(
+                "simBuffer", datasize);
+            dataManReader.Get<float>(floatArrayVar, floatVector, adios2::Mode::Sync);
+            dataManReader.EndStep();
+            PrintData(floatVector, dataManReader.CurrentStep());
+        }
+        else if (status == adios2::StepStatus::EndOfStream)
+        {
+            std::cout << "End of stream" << std::endl;
+            break;
+        }
+    }
+
+    // clean up
+    dataManReader.Close();
+    MPI_Finalize();
+
+    return 0;
+}
diff --git a/examples/hello/datamanKokkos/dataManWriterKokkos.cpp b/examples/hello/datamanKokkos/dataManWriterKokkos.cpp
@@ -0,0 +1,97 @@
+/*
+ * Distributed under the OSI-approved Apache License, Version 2.0.  See
+ * accompanying file Copyright.txt for details.
+ *
+ * datamanWriterKokkos.cpp  Simple example of writing multiple steps of a 2D float Kokkos::View
+ * through ADIOS2 DataMan
+ */
+#include <adios2.h>
+#include <adios2/cxx11/KokkosView.h>
+#include <iostream>
+#include <mpi.h>
+#include <numeric>
+#include <thread>
+#include <vector>
+
+#include <Kokkos_Core.hpp>
+
+size_t Nx = 10;
+size_t Ny = 10;
+size_t steps = 2;
+adios2::Dims shape;
+adios2::Dims start;
+adios2::Dims count;
+
+int mpiRank, mpiSize;
+
+template <class T, class MemSpace>
+void PrintData(Kokkos::View<T **, MemSpace> &gpuData, const size_t step)
+{
+    auto data = Kokkos::create_mirror_view_and_copy(Kokkos::HostSpace{}, gpuData);
+    std::cout << "Rank: " << mpiRank << " Step: " << step << " [";
+    for (int i = 0; i < data.extent_int(0); ++i)
+        for (int j = 0; j < data.extent_int(1); ++j)
+            std::cout << data(i, j) << " ";
+    std::cout << "]" << std::endl;
+}
+
+template <class T, class MemSpace, class ExecSpace>
+Kokkos::View<T **, MemSpace> GenerateData(const size_t step, const size_t Ny, const size_t mpiRank)
+{
+    Kokkos::View<T **, 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 * Ny + j) + mpiRank * 10000 + step;
+        });
+    Kokkos::fence();
+    ExecSpace exe_space;
+    std::cout << "Create data for step " << step << " on memory space: " << exe_space.name()
+              << std::endl;
+    return gpuSimData;
+}
+
+int main(int argc, char *argv[])
+{
+    // initialize MPI
+    MPI_Init(&argc, &argv);
+    MPI_Comm_rank(MPI_COMM_WORLD, &mpiRank);
+    MPI_Comm_size(MPI_COMM_WORLD, &mpiSize);
+
+    // initialize data dimensions
+    count = {Nx, Ny};
+    start = {mpiRank * Nx, 0};
+    shape = {mpiSize * Nx, Ny};
+
+    // initialize adios2
+    adios2::ADIOS adios(MPI_COMM_WORLD);
+    adios2::IO dataManIO = adios.DeclareIO("whatever");
+    dataManIO.SetEngine("DataMan");
+    dataManIO.SetParameters({{"IPAddress", "127.0.0.1"},
+                             {"Port", "12306"},
+                             {"Timeout", "5"},
+                             {"RendezvousReaderCount", "1"}});
+
+    // open stream
+    adios2::Engine dataManWriter = dataManIO.Open("HelloDataMan", adios2::Mode::Write);
+
+    // define variable
+    auto floatArrayVar = dataManIO.DefineVariable<float>("FloatArray", shape, start, count);
+
+    // write data
+    for (size_t i = 0; i < steps; ++i)
+    {
+        auto floatVector = GenerateData<float, Kokkos::DefaultExecutionSpace::memory_space,
+                                        Kokkos::DefaultExecutionSpace>(i, Ny, mpiRank);
+        dataManWriter.BeginStep();
+        dataManWriter.Put(floatArrayVar, floatVector, adios2::Mode::Sync);
+        PrintData(floatVector, dataManWriter.CurrentStep());
+        dataManWriter.EndStep();
+    }
+
+    dataManWriter.Close();
+    MPI_Finalize();
+
+    return 0;
+}