working ishmem_init & ishmem_finalize on borealis (#599)

oneapi-src · Oct 31, 2023 · b35f0d8 · b35f0d8
1 parent 177d231
commit b35f0d8
Show file tree

Hide file tree

Showing 3 changed files with 28 additions and 3 deletions.
diff --git a/CMakeLists.txt b/CMakeLists.txt
@@ -194,6 +194,16 @@ if(ENABLE_ISHMEM)
 
   add_custom_target(shmem DEPENDS ishmem)
 
+  function(add_shmem_test test_name name processes)
+    add_test(
+      NAME ${test_name}
+      # FIXME, can not use ${MPIEXEC_EXECUTABLE} on Borealis because it gives
+      # two ranks 0 and failure in ISHMEM init
+      COMMAND
+        mpiexec ${MPIEXEC_NUMPROC_FLAG} ${processes} ${MPIEXEC_PREFLAGS}
+        ${CMAKE_BINARY_DIR}/bin/ishmrun ./${name} ${ARGN} COMMAND_EXPAND_LISTS)
+  endfunction()
+
 endif()
 
 function(add_mhp_ctest test_name name processes)

diff --git a/include/dr/mhp/global.hpp b/include/dr/mhp/global.hpp
@@ -5,7 +5,9 @@
 #pragma once
 
 #include <unistd.h>
-
+#ifdef DRISHMEM
+#include <ishmem.h>
+#endif
 #include <dr/detail/sycl_utils.hpp>
 #include <dr/mhp/sycl_support.hpp>
 
@@ -71,6 +73,10 @@ inline void initialize_mpi() {
     MPI_Init(nullptr, nullptr);
     we_initialized_mpi_ = true;
   }
+
+#ifdef DRISHMEM
+  ishmem_init();
+#endif
 }
 
 // Finalize MPI *if* we initialized it and it has not been finalized.
@@ -81,6 +87,10 @@ inline void finalize_mpi() {
   if (we_initialized_mpi_ && !finalized) {
     MPI_Finalize();
   }
+
+#ifdef DRISHMEM
+  ishmem_finalize();
+#endif
 }
 
 } // namespace __detail

diff --git a/test/gtest/mhp/CMakeLists.txt b/test/gtest/mhp/CMakeLists.txt
@@ -68,9 +68,14 @@ endforeach()
 
 if(ENABLE_ISHMEM)
   add_executable(mhp-shmem-test mhp-tests.cpp ../common/zip.cpp)
+  target_link_libraries(
+    mhp-shmem-test GTest::gtest_main cxxopts DR::mpi ze_loader fabric
+    pmi_simple sma)
+  target_link_libraries(mhp-shmem-test ${CMAKE_BINARY_DIR}/lib/libishmem.a)
+  target_compile_definitions(mhp-shmem-test PRIVATE DRISHMEM)
 
-  target_link_libraries(mhp-shmem-test GTest::gtest_main cxxopts DR::mpi
-                        ze_loader fabric)
+  add_shmem_test(mhp-shmem-test-1 mhp-shmem-test 1 --sycl)
+  add_shmem_test(mhp-shmem-test-2 mhp-shmem-test 2 --sycl)
 endif()
 
 add_mhp_ctest(mhp-quick-test-1 mhp-quick-test 1)