Merge pull request ornladios#3976 from anagainaru/gpu-layout-simpler

A simpler way to deal with the mismatching layouts between the user code and the memory space

bindings/CXX11/adios2/cxx11/Engine.cpp

@@ -17,14 +17,6 @@
 namespace adios2
 {
 
-#ifdef ADIOS2_HAVE_GPU_SUPPORT
-void Engine::CheckMemorySpace(MemorySpace variableMem, MemorySpace bufferMem)
-{
-    if (variableMem != MemorySpace::Detect && variableMem != bufferMem)
-        helper::Throw<std::runtime_error>("CXX-Bindings", "Engine", "Put", "Memory space mismatch");
-}
-#endif
-
 Engine::operator bool() const noexcept
 {
     if (m_Engine == nullptr)

bindings/CXX11/adios2/cxx11/Engine.h

@@ -38,10 +38,6 @@ class Engine
     friend class IO;
     friend class QueryWorker;
 
-#ifdef ADIOS2_HAVE_GPU_SUPPORT
-    void CheckMemorySpace(MemorySpace variableMem, MemorySpace bufferMem);
-#endif
-
 public:
     /**
      * Empty (default) constructor, use it as a placeholder for future
@@ -212,12 +208,12 @@ class Engine
     void Put(Variable<T> variable, U const &data, const Mode launch = Mode::Deferred)
     {
         auto bufferView = static_cast<AdiosView<U>>(data);
+#if defined(ADIOS2_HAVE_KOKKOS) || defined(ADIOS2_HAVE_GPU_SUPPORT)
         auto bufferMem = bufferView.memory_space();
-#ifdef ADIOS2_HAVE_GPU_SUPPORT
-        auto variableMem = variable.GetMemorySpace();
-        CheckMemorySpace(variableMem, bufferMem);
-#endif
+        auto bufferLayout = bufferView.layout();
         variable.SetMemorySpace(bufferMem);
+        variable.SetArrayLayout(bufferLayout);
+#endif
         Put(variable, bufferView.data(), launch);
     }
 
@@ -417,10 +413,14 @@ class Engine
               class = typename std::enable_if<std::is_convertible<U, AdiosView<U>>::value>::type>
     void Get(Variable<T> variable, U const &data, const Mode launch = Mode::Deferred)
     {
-        auto adios_data = static_cast<AdiosView<U>>(data);
-        auto mem_space = adios_data.memory_space();
-        variable.SetMemorySpace(mem_space);
-        Get(variable, adios_data.data(), launch);
+        auto bufferView = static_cast<AdiosView<U>>(data);
+#if defined(ADIOS2_HAVE_KOKKOS) || defined(ADIOS2_HAVE_GPU_SUPPORT)
+        auto bufferMem = bufferView.memory_space();
+        auto bufferLayout = bufferView.layout();
+        variable.SetMemorySpace(bufferMem);
+        variable.SetArrayLayout(bufferLayout);
+#endif
+        Get(variable, bufferView.data(), launch);
     }
 
     /** Perform all Get calls in Deferred mode up to this point */

bindings/CXX11/adios2/cxx11/KokkosView.h

@@ -28,6 +28,23 @@ struct memspace_kokkos_to_adios2
 };
 #endif
 
+template <typename T>
+struct layout_kokkos_to_adios2
+{
+    static constexpr adios2::ArrayOrdering value = adios2::ArrayOrdering::RowMajor;
+};
+
+template <>
+struct layout_kokkos_to_adios2<Kokkos::LayoutLeft>
+{
+    static constexpr adios2::ArrayOrdering value = adios2::ArrayOrdering::ColumnMajor;
+};
+
+template <>
+struct layout_kokkos_to_adios2<Kokkos::LayoutRight>
+{
+    static constexpr adios2::ArrayOrdering value = adios2::ArrayOrdering::RowMajor;
+};
 } // namespace detail
 
 template <class T, class... Parameters>
@@ -36,6 +53,7 @@ class AdiosView<Kokkos::View<T, Parameters...>>
     using data_type = typename Kokkos::View<T, Parameters...>::value_type;
     data_type *pointer;
     adios2::MemorySpace mem_space;
+    adios2::ArrayOrdering m_layout;
 
 public:
     template <class... P>
@@ -44,11 +62,14 @@ class AdiosView<Kokkos::View<T, Parameters...>>
         pointer = v.data();
         mem_space =
             detail::memspace_kokkos_to_adios2<typename Kokkos::View<T, P...>::memory_space>::value;
+        m_layout =
+            detail::layout_kokkos_to_adios2<typename Kokkos::View<T, P...>::array_layout>::value;
     }
 
     data_type const *data() const { return pointer; }
     data_type *data() { return pointer; }
     adios2::MemorySpace memory_space() const { return mem_space; }
+    adios2::ArrayOrdering layout() const { return m_layout; }
 };
 
 }

bindings/CXX11/adios2/cxx11/Variable.cpp

@@ -126,6 +126,20 @@ namespace adios2
     }                                                                                              \
                                                                                                    \
     template <>                                                                                    \
+    Dims Variable<T>::Shape(const ArrayOrdering layout, const size_t step) const                   \
+    {                                                                                              \
+        helper::CheckForNullptr(m_Variable, "in call to Variable<T>::Shape");                      \
+        return m_Variable->Shape(step, MemorySpace::Host, layout);                                 \
+    }                                                                                              \
+                                                                                                   \
+    template <>                                                                                    \
+    Dims Variable<T>::Shape(const MemorySpace memSpace, const size_t step) const                   \
+    {                                                                                              \
+        helper::CheckForNullptr(m_Variable, "in call to Variable<T>::Shape");                      \
+        return m_Variable->Shape(step, memSpace);                                                  \
+    }                                                                                              \
+                                                                                                   \
+    template <>                                                                                    \
     Dims Variable<T>::Start() const                                                                \
     {                                                                                              \
         helper::CheckForNullptr(m_Variable, "in call to Variable<T>::Start");                      \
@@ -266,6 +280,14 @@ ADIOS2_FOREACH_TYPE_1ARG(declare_template_instantiation)
 ADIOS2_FOREACH_PRIMITIVE_TYPE_1ARG(declare_template_instantiation)
 #undef declare_template_instantiation
 
+#if defined(ADIOS2_HAVE_KOKKOS) || defined(ADIOS2_HAVE_GPU_SUPPORT)
+#define declare_layout_template_instantiation(T)                                                   \
+    template void Variable<T>::SetArrayLayout(const ArrayOrdering layout);                         \
+    template ArrayOrdering Variable<T>::GetArrayLayout();
+ADIOS2_FOREACH_TYPE_1ARG(declare_layout_template_instantiation)
+#undef declare_layout_template_instantiation
+#endif
+
 #define declare_template_instantiation(T)                                                          \
     template std::vector<typename Variable<T>::Info> Variable<T>::ToBlocksInfoMin(                 \
         const MinVarInfo *coreVarInfo) const;

bindings/CXX11/adios2/cxx11/Variable.h

@@ -160,6 +160,20 @@ class Variable
      */
     MemorySpace GetMemorySpace();
 
+#if defined(ADIOS2_HAVE_KOKKOS) || defined(ADIOS2_HAVE_GPU_SUPPORT)
+    /**
+     * Sets the  for all following Puts
+     * to either host (default) or device (currently only CUDA supported)
+     * @param mem memory space where Put buffers are allocated
+     */
+    void SetArrayLayout(const adios2::ArrayOrdering layout);
+    /**
+     * Get the memory space that was set by the application
+     * @return the memory space stored in the Variable object
+     */
+    adios2::ArrayOrdering GetArrayLayout();
+#endif
+
     /**
      * Set new shape, care must be taken when reading back the variable for
      * different steps. Only applies to Global arrays.
@@ -255,6 +269,10 @@ class Variable
      * @return shape vector
      */
     adios2::Dims Shape(const size_t step = adios2::EngineCurrentStep) const;
+    adios2::Dims Shape(const ArrayOrdering layout,
+                       const size_t step = adios2::EngineCurrentStep) const;
+    adios2::Dims Shape(const MemorySpace memSpace,
+                       const size_t step = adios2::EngineCurrentStep) const;
 
     /**
      * Inspects current start point

bindings/CXX11/adios2/cxx11/Variable.tcc

@@ -215,6 +215,20 @@ std::string ToString(const Variable<T> &variable)
     return std::string("Variable<") + variable.Type() + ">(Name: \"" + variable.Name() + "\")";
 }
 
+#if defined(ADIOS2_HAVE_KOKKOS) || defined(ADIOS2_HAVE_GPU_SUPPORT)
+template <class T>
+void Variable<T>::SetArrayLayout(const ArrayOrdering layout)
+{
+    m_Variable->SetArrayLayout(layout);
+}
+
+template <class T>
+adios2::ArrayOrdering Variable<T>::GetArrayLayout()
+{
+    return m_Variable->GetArrayLayout();
+}
+#endif
+
 namespace detail
 {
 // Span

examples/hello/bpStepsWriteReadKokkos/CMakeLists.txt

@@ -18,16 +18,62 @@ if(NOT TARGET adios2_core)
     set(CMAKE_CXX_COMPILER "${Kokkos_CXX_COMPILER}")
   endif()
 
+  include(CheckLanguage)
+  check_language(Fortran)
+  if(CMAKE_Fortran_COMPILER)
+    enable_language(Fortran)
+  endif()
+  if(CMAKE_Fortran_COMPILER_LOADED)
+    list(APPEND _components Fortran)
+  endif()
+
+  find_package(MPI QUIET COMPONENTS ${_components})
+  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(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_bpStepsWriteReadKokkos bpStepsWriteReadKokkos.cpp)
-  kokkos_compilation(SOURCE bpStepsWriteReadKokkos.cpp)
+# flcl package needed for the Kokkos Fortran example
+find_package(flcl QUIET)
+
+# C++ Kokkos example using default layouts and different memory spaces
+add_executable(adios2_hello_bpStepsWriteReadKokkos bpStepsWriteReadKokkos.cpp)
+kokkos_compilation(SOURCE bpStepsWriteReadKokkos.cpp)
+if(ADIOS2_HAVE_MPI)
+  target_link_libraries(adios2_hello_bpStepsWriteReadKokkos adios2::cxx11_mpi MPI::MPI_C Kokkos::kokkos)
+else()
   target_link_libraries(adios2_hello_bpStepsWriteReadKokkos adios2::cxx11 Kokkos::kokkos)
-  install(TARGETS adios2_hello_bpStepsWriteReadKokkos RUNTIME DESTINATION ${CMAKE_INSTALL_BINDIR})
+endif()
+install(TARGETS adios2_hello_bpStepsWriteReadKokkos RUNTIME DESTINATION ${CMAKE_INSTALL_BINDIR})
+
+# C++ Kokkos example using default memory space and layouts
+add_executable(adios2_hello_bpWriteReadKokkosView bpWriteReadKokkosView.cpp)
+kokkos_compilation(SOURCE bpStepsWriteReadKokkos.cpp)
+if(ADIOS2_HAVE_MPI)
+  target_link_libraries(adios2_hello_bpWriteReadKokkosView adios2::cxx11_mpi MPI::MPI_C Kokkos::kokkos)
+else()
+  target_link_libraries(adios2_hello_bpWriteReadKokkosView adios2::cxx11 Kokkos::kokkos)
+endif()
+install(TARGETS adios2_hello_bpWriteReadKokkosView RUNTIME DESTINATION ${CMAKE_INSTALL_BINDIR})
+
+# Fortran Kokkos example using default layouts and different memory spaces
+if(ADIOS2_HAVE_MPI AND ADIOS2_HAVE_Fortran AND flcl_FOUND)
+  add_executable(adios2_hello_bpStepsWriteReadKokkos_f bpStepsWriteReadKokkos.F90 view-f.f90 view-cxx.cc)
+  target_link_libraries(adios2_hello_bpStepsWriteReadKokkos_f adios2::fortran_mpi MPI::MPI_Fortran flcl::flcl)
+  if (CMAKE_Fortran_COMPILER_ID STREQUAL "XL")
+    target_link_options(adios2_hello_bpStepsWriteReadKokkos_f PRIVATE LINKER:-lxlf90_r)
+  endif()
+  if (CMAKE_Fortran_COMPILER_ID STREQUAL "Intel" OR CMAKE_Fortran_COMPILER_ID STREQUAL "GNU")
+    set_target_properties(adios2_hello_bpStepsWriteReadKokkos_f PROPERTIES LINKER_LANGUAGE Fortran)
+  endif()
+  install(TARGETS adios2_hello_bpStepsWriteReadKokkos_f RUNTIME DESTINATION ${CMAKE_INSTALL_BINDIR})
 endif()

examples/hello/bpStepsWriteReadKokkos/bpStepsWriteReadKokkos.F90

@@ -0,0 +1,80 @@
+program TestBPWriteReadHeatMap2D use mpi use adios2
+
+    implicit none
+
+        integer(kind = 8)::sum_i1,
+    sum_i2 type(adios2_adios)::adios type(adios2_io)::ioPut, ioGet type(adios2_engine)::bpWriter,
+    bpReader type(adios2_variable), dimension(1)::var_g,
+    var_gIn
+
+    integer(kind = 2),
+    dimension(
+        :,
+        :),
+    allocatable ::g, &sel_g integer(kind = 8),
+    dimension(2)::ishape, istart, icount integer(kind = 8),
+    dimension(2)::sel_start, sel_count integer ::ierr, irank, isize,
+    step_status integer ::in1, in2 integer ::i1,
+    i2
+
+    call MPI_INIT(ierr) call MPI_COMM_RANK(MPI_COMM_WORLD, irank, ierr) call
+    MPI_COMM_SIZE(MPI_COMM_WORLD, isize, ierr)
+
+        in1 = 3 in2 = 4 icount = (/ in1, in2 /) istart = (/ 0, in2 *irank /) ishape = (/ in1,
+                                                                                       in2 *isize /)
+
+            allocate(g(in1, in2)) do i2 = 1,
+        in2 do i1 = 1,
+        in1 g(i1, i2) = irank + i1 end do end do
+
+                        !Start adios2 Writer call adios2_init(adios, MPI_COMM_WORLD, ierr) call
+                        adios2_declare_io(ioPut, adios, 'WriteIO', ierr)
+
+                            call adios2_define_variable(var_g(1), ioPut, &'bpFloats',
+                                                        adios2_type_integer2, &2, ishape, istart,
+                                                        icount, &adios2_constant_dims, ierr)
+
+                                call
+                        adios2_open(bpWriter, ioPut, 'BPFortranKokkos.bp', adios2_mode_write, &ierr)
+
+                            call adios2_put(bpWriter, var_g(1), g, ierr)
+
+                                call adios2_close(bpWriter, ierr)
+
+                                    if (allocated(g)) deallocate(g)
+
+                                        !Start adios2 Reader in rank 0 if (irank == 0) then
+
+                        call adios2_declare_io(ioGet, adios, 'ReadIO', ierr)
+
+                            call adios2_open(bpReader, ioGet, 'BPFortranKokkos.bp',
+                                             &adios2_mode_read, MPI_COMM_SELF, ierr)
+
+                                call
+                        adios2_begin_step(bpReader, adios2_step_mode_read, -1., &step_status, ierr)
+
+                            call adios2_inquire_variable(var_gIn(1), ioGet, &'bpFloats', ierr)
+
+                                sel_start = (/ 0, 0 /) sel_count = (/ ishape(1), ishape(2) /)
+
+                            allocate(sel_g(ishape(1), ishape(2))) sel_g = 0
+
+        call adios2_set_selection(var_gIn(1), 2, sel_start, sel_count, &ierr) call
+        adios2_get(bpReader, var_gIn(1), sel_g, ierr)
+
+            call adios2_end_step(bpReader, ierr)
+
+                call adios2_close(bpReader, ierr)
+
+                    do i2 = 1,
+        INT(sel_count(2), 4) do i1 = 1,
+        INT(sel_count(1), 4) write(6, "(i8)", advance = "no") sel_g(i1, i2) end do write(6, *) end
+    do
+
+    if (allocated(sel_g)) deallocate(sel_g)
+
+        end if
+
+    call adios2_finalize(adios, ierr) call MPI_Finalize(ierr)
+
+        end program TestBPWriteReadHeatMap2D