Porting to GHEX 0.2.0 + Fixes for Multiple Devices (#77)

- `ghex_comm` runtime ported to ghex 0.2.0, using oomph as transport library.
- fixed an issue that caused `detected error differences between the runs` when using the `ghex_comm` runtime on multiple devices. It was due both to streams synchronization issues and to the choice of the default set-device strategy.
- GHEX code does not depend anymore on the `gridtools::_impl` namespace.
- NOTE: semantics of the `--device-mapping` option has changed: it now accepts only one device id per rank, meaning that different threads on the same rank must use the same device.
- AMD GPUs fixes are not included, and will follow shortly.

Author: bettiolm

.github/workflows/tests.yml

@@ -19,9 +19,6 @@ jobs:
       matrix:
         backend: [cpu_ifirst, cpu_kfirst, cuda, hip]
         runtime: [single_node, simple_mpi, gcl, ghex_comm]
-        exclude:
-          - backend: hip
-            runtime: ghex_comm
     steps:
     - uses: actions/checkout@v2
       name: checkout

CMakeLists.txt

@@ -65,7 +65,7 @@ endif()
 # Helper functions
 function(compile_as_cuda)
   if(_gtbench_cuda_enabled)
-    set_source_files_properties(${ARGN} PROPERTIES LANGUAGE CUDA)
+    set_source_files_properties(${ARGN} PROPERTIES LANGUAGE CUDA COMPILE_OPTIONS "--default-stream=per-thread")
   endif()
 endfunction()
 

Dockerfile

@@ -10,7 +10,8 @@ RUN apt-get update -qq && \
     libmpich-dev \
     tar \
     software-properties-common \
-    wget && \
+    wget \
+    libnuma-dev && \
     rm -rf /var/lib/apt/lists/*
 
 ARG CMAKE_VERSION=3.18.4
@@ -34,7 +35,7 @@ COPY . /gtbench
 RUN cd /gtbench && \
     mkdir -p build && \
     cd build && \
-    if [ -d /opt/rocm ]; then export CXX=/opt/rocm/bin/hipcc; fi && \
+    if [ -d /opt/rocm ]; then export ROCM_PATH=/opt/rocm; export PATH=${ROCM_PATH}/bin:${PATH}; export CXX=${ROCM_PATH}/bin/hipcc; fi && \
     cmake \
     -DCMAKE_BUILD_TYPE=Release \
     -DGTBENCH_BACKEND=${GTBENCH_BACKEND} \
@@ -43,5 +44,6 @@ RUN cd /gtbench && \
     .. && \
     make -j $(nproc) install && \
     rm -rf /gtbench/build
+ENV LD_LIBRARY_PATH=/usr/local/lib64:/usr/local/lib:${LD_LIBRARY_PATH}
 
 CMD ["convergence_tests"]

include/gtbench/runtime/ghex_comm/run.hpp

@@ -40,6 +40,7 @@ struct runtime {
   std::array<int, 2> m_cart_dims;
   std::array<int, 2> m_thread_cart_dims;
   std::vector<int> m_device_mapping;
+  int m_device;
   std::string m_output_filename;
 };
 
@@ -74,7 +75,7 @@ result runtime_solve(runtime &rt, Analytical analytical, Stepper stepper,
 
   std::vector<result> results(rt.m_num_threads);
   auto execution_func = [&](int id = 0) {
-    set_device(rt.m_device_mapping[id]);
+    set_device(rt.m_device);
     auto sub_grid = comm_grid[id];
     const auto exact = discrete_analytical::discretize(
         analytical, global_resolution, sub_grid.m_local_resolution,
@@ -116,7 +117,7 @@ result runtime_solve(runtime &rt, Analytical analytical, Stepper stepper,
   threads.reserve(rt.m_num_threads - 1);
   for (int i = 1; i < rt.m_num_threads; ++i)
     threads.emplace_back(execution_func, i);
-  set_device(rt.m_device_mapping[0]);
+  set_device(rt.m_device);
   execution_func(0);
 
   for (auto &thread : threads)

src/runtime/CMakeLists.txt

@@ -2,5 +2,5 @@ add_library(runtime)
 target_link_libraries(runtime PUBLIC gtbench_common)
 target_link_libraries(gtbench PUBLIC runtime)
 
-add_subdirectory(${GTBENCH_RUNTIME})
 add_subdirectory(device)
+add_subdirectory(${GTBENCH_RUNTIME})

src/runtime/device/set_device.cpp

@@ -9,8 +9,17 @@
  */
 #include <stdexcept>
 
+#include <gridtools/common/defs.hpp>
 #include <gtbench/runtime/device/set_device.hpp>
 
+#ifdef GT_CUDACC
+#include <gridtools/common/cuda_runtime.hpp>
+#if defined(__HIP__)
+#define cudaDeviceCanAccessPeer hipDeviceCanAccessPeer
+#define cudaDeviceEnablePeerAccess hipDeviceEnablePeerAccess
+#endif
+#endif
+
 namespace gtbench {
 namespace runtime {
 

src/runtime/ghex_comm/CMakeLists.txt

@@ -1,9 +1,12 @@
+compile_as_cuda(run.cpp TARGET_DIRECTORY runtime)
+target_sources(runtime PRIVATE run.cpp)
+
 if(NOT _ghex_already_fetched)
   find_package(GHEX QUIET)
 endif()
 if(NOT GHEX_FOUND)
-  set(_ghex_repository "https://github.com/GridTools/GHEX.git")
-  set(_ghex_tag        "4d48f3145349064942941f191021b970db3cb36e")
+  set(_ghex_repository "https://github.com/ghex-org/GHEX.git")
+  set(_ghex_tag        "b9adb65a28303a05d607a34d6a75de4f36305ae0")
   if(NOT _ghex_already_fetched)
     message(STATUS "Fetching GHEX ${_ghex_tag} from ${_ghex_repository}")
   endif()
@@ -13,21 +16,25 @@ if(NOT GHEX_FOUND)
     GIT_REPOSITORY ${_ghex_repository}
     GIT_TAG        ${_ghex_tag}
   )
+  if(GTBENCH_BACKEND STREQUAL "gpu")
+    set(GHEX_USE_GPU ON CACHE INTERNAL "")
+    if(_gtbench_cuda_enabled)
+      set(GHEX_GPU_TYPE "NVIDIA" CACHE INTERNAL "")
+    endif()
+  endif()
   FetchContent_MakeAvailable(ghex)
   set(_ghex_already_fetched ON CACHE INTERNAL "")
 endif()
 
-compile_as_cuda(run.cpp TARGET_DIRECTORY runtime)
-target_sources(runtime PRIVATE run.cpp)
-target_link_libraries(runtime PUBLIC GHEX::ghexlib)
-target_compile_options(runtime PUBLIC "$<$<COMPILE_LANGUAGE:CUDA>:SHELL:--default-stream per-thread>")
-
-if(GHEX_USE_UCP)
-  target_compile_definitions(runtime PRIVATE GTBENCH_USE_GHEX_UCP)
-  if(GHEX_USE_PMIX)
-    target_compile_definitions(runtime PRIVATE GTBENCH_USE_GHEX_PMIX)
-  endif()
+target_link_libraries(runtime PUBLIC GHEX::lib)
+if(GHEX_TRANSPORT_BACKEND STREQUAL "LIBFABRIC")
+  target_link_libraries(runtime PUBLIC oomph::libfabric)
+elseif(GHEX_TRANSPORT_BACKEND STREQUAL "UCX")
+  target_link_libraries(runtime PUBLIC oomph::ucx)
+elseif(GHEX_TRANSPORT_BACKEND STREQUAL "MPI")
+  target_link_libraries(runtime PUBLIC oomph::mpi)
 endif()
+
 if(GHEX_USE_XPMEM)
     target_compile_definitions(runtime PRIVATE GHEX_USE_XPMEM)
 endif()

src/runtime/ghex_comm/run.cpp

@@ -11,21 +11,13 @@
 #include <numeric>
 #include <regex>
 
+#include <ghex/barrier.hpp>
 #include <ghex/bulk_communication_object.hpp>
 #include <ghex/glue/gridtools/field.hpp>
 #include <ghex/structured/grid.hpp>
 #include <ghex/structured/pattern.hpp>
 #include <ghex/structured/rma_range_generator.hpp>
 
-#ifdef GTBENCH_USE_GHEX_UCP
-#include <ghex/transport_layer/ucx/context.hpp>
-using transport = gridtools::ghex::tl::ucx_tag;
-#else
-#include <ghex/transport_layer/mpi/context.hpp>
-using transport = gridtools::ghex::tl::mpi_tag;
-#endif
-#include <ghex/transport_layer/util/barrier.hpp>
-
 #include <gtbench/runtime/ghex_comm/factorize.hpp>
 #include <gtbench/runtime/ghex_comm/run.hpp>
 
@@ -49,7 +41,7 @@ runtime::runtime(int num_threads, std::array<int, 2> cart_dims,
           },
           MPI_Finalize),
       m_num_threads(num_threads), m_cart_dims(cart_dims),
-      m_thread_cart_dims(thread_cart_dims), m_device_mapping(num_threads, 0),
+      m_thread_cart_dims(thread_cart_dims), m_device_mapping(), m_device(0),
       m_output_filename(output_filename) {
   int size, rank;
   MPI_Comm_size(MPI_COMM_WORLD, &size);
@@ -91,22 +83,20 @@ runtime::runtime(int num_threads, std::array<int, 2> cart_dims,
   MPI_Comm_rank(shmem_comm, &shmem_rank);
   MPI_Comm_free(&shmem_comm);
   if (!device_mapping.empty()) {
-    if (device_mapping.size() != shmem_size * num_threads)
+    if (device_mapping.size() != shmem_size)
       throw std::runtime_error("device mapping has wrong size");
     m_device_mapping = device_mapping;
   } else {
-    m_device_mapping.resize(shmem_size * m_num_threads);
+    m_device_mapping.resize(shmem_size);
     std::iota(m_device_mapping.begin(), m_device_mapping.end(), 0);
   }
-  m_device_mapping = std::vector<int>(
-      m_device_mapping.begin() + shmem_rank * num_threads,
-      m_device_mapping.begin() + (shmem_rank + 1) * num_threads);
+  m_device = m_device_mapping[shmem_rank];
 #endif
 }
 
 using domain_id_t = int;
 using dimension_t = std::integral_constant<int, 3>;
-using coordinate_t = gt::ghex::coordinate<std::array<int, 3>>;
+using coordinate_t = ghex::coordinate<std::array<int, 3>>;
 
 struct local_domain {
   using domain_id_type = domain_id_t;
@@ -122,12 +112,11 @@ struct local_domain {
   coordinate_t const &last() const { return m_last; }
 };
 
-using context_t =
-    typename gt::ghex::tl::context_factory<transport>::context_type;
+using context_t = ghex::context;
 using communicator_t = context_t::communicator_type;
-using grid_t = gt::ghex::structured::grid::type<local_domain>;
-using patterns_t =
-    gt::ghex::pattern_container<communicator_t, grid_t, domain_id_t>;
+using barrier_t = ghex::barrier;
+using grid_t = ghex::structured::grid::type<local_domain>;
+using patterns_t = ghex::pattern_container<grid_t, domain_id_t>;
 
 struct halo_generator {
   using domain_type = local_domain;
@@ -214,7 +203,7 @@ class grid::impl {
   domain_vec m_domains;
   context_ptr_t m_context;
   patterns_ptr_t m_patterns;
-  gridtools::ghex::tl::barrier_t m_barrier;
+  barrier_t m_barrier;
 
 public:
   impl(vec<std::size_t, 3> const &global_resolution, int num_sub_domains,
@@ -224,7 +213,9 @@ class grid::impl {
                                 (int)global_resolution.y - 1,
                                 (int)global_resolution.z - 1}},
         m_global_resolution{global_resolution.x, global_resolution.y},
-        m_barrier(num_sub_domains) {
+        m_context{
+            std::make_unique<context_t>(MPI_COMM_WORLD, (num_sub_domains > 1))},
+        m_barrier{*m_context, static_cast<std::size_t>(num_sub_domains)} {
     MPI_Comm_size(MPI_COMM_WORLD, &m_size);
     MPI_Comm_rank(MPI_COMM_WORLD, &m_rank);
 
@@ -271,20 +262,17 @@ class grid::impl {
                                          (int)global_resolution.z - 1}});
       }
     }
-    m_context =
-        gt::ghex::tl::context_factory<transport>::create(MPI_COMM_WORLD);
     m_patterns = std::make_unique<patterns_type>(
-        gt::ghex::make_pattern<gt::ghex::structured::grid>(*m_context, m_hg,
-                                                           m_domains));
+        ghex::make_pattern<ghex::structured::grid>(*m_context, m_hg,
+                                                   m_domains));
   }
 
   impl(impl const &) = delete;
   impl &operator=(impl const &) = delete;
 
   sub_grid operator[](unsigned int i) {
     const auto &dom = m_domains[i];
-    auto comm = m_context->get_communicator();
-    m_barrier(comm);
+    m_barrier();
 
     vec<std::size_t, 3> local_resolution = {
         (std::size_t)(dom.last()[0] - dom.first()[0] + 1),
@@ -295,29 +283,30 @@ class grid::impl {
                                         (std::size_t)dom.first()[2]};
 
     auto b_comm_obj_map = std::make_shared<
-        std::map<void *, gt::ghex::generic_bulk_communication_object>>();
+        std::map<void *, ghex::generic_bulk_communication_object>>();
 
-    auto halo_exchange = [b_comm_obj_map = std::move(b_comm_obj_map), comm,
-                          domain = dom,
+    auto halo_exchange = [b_comm_obj_map = std::move(b_comm_obj_map),
+                          &context = m_context, domain = dom,
                           &patterns = *m_patterns](storage_t &storage) mutable {
 #ifdef GTBENCH_BACKEND_GPU
-      using arch_t = gt::ghex::gpu;
+      using arch_t = ghex::gpu;
 #else
-      using arch_t = gt::ghex::cpu;
+      using arch_t = ghex::cpu;
 #endif
-      auto field =
-          gt::ghex::wrap_gt_field<arch_t>(domain, storage, {halo, halo, 0});
+      auto field = ghex::wrap_gt_field<arch_t>(
+          domain, storage,
+          {halo, halo, 0}); // device_id is initialized to the current device id
+                            // by default in GHEX
       auto it = b_comm_obj_map->find(field.data());
       if (it == b_comm_obj_map->end()) {
-        auto sbco = gt::ghex::bulk_communication_object<
-            gt::ghex::structured::rma_range_generator, patterns_type,
-            decltype(field)>(comm);
+        auto sbco = ghex::bulk_communication_object<
+            ghex::structured::rma_range_generator, patterns_type,
+            decltype(field)>(*context);
         sbco.add_field(patterns(field));
         it = b_comm_obj_map
-                 ->insert(
-                     std::make_pair((void *)field.data(),
-                                    gt::ghex::generic_bulk_communication_object(
-                                        std::move(sbco))))
+                 ->insert(std::make_pair(
+                     (void *)field.data(),
+                     ghex::generic_bulk_communication_object(std::move(sbco))))
                  .first;
       }
       auto &bco = it->second;
@@ -372,10 +361,10 @@ void runtime_register_options(ghex_comm, options &options) {
           "TX TY", 2);
 #ifdef GT_CUDACC
   options("device-mapping",
-          "node device mapping: device id per sub-domain in the format "
+          "node device mapping: device id per rank in the format "
           "I_0:I_1:...:I_(N-1) "
           "where I_i are cuda device ids "
-          "and N = #ranks-per-node x S",
+          "and N = #ranks-per-node",
           "M");
 #endif
 }