Merge pull request #1225 from masterleinad/query_memory_sycl

Allow querying free/total memory for SYCL
kokkos · Jan 19, 2022 · 701fb05 · 701fb05
2 parents a367801 + 5a8e014
commit 701fb05
Show file tree

Hide file tree

Showing 4 changed files with 71 additions and 3 deletions.
diff --git a/CMakeLists.txt b/CMakeLists.txt
@@ -47,6 +47,8 @@ ENDIF()
 INCLUDE(cmake/fake_tribits.cmake)
 INCLUDE(cmake/kokkoskernels_tribits.cmake)
 
+OPTION(BUILD_SHARED_LIBS "Build shared libraries" OFF)
+
 KOKKOSKERNELS_PACKAGE()
 
 IF (NOT KOKKOSKERNELS_HAS_TRILINOS)

diff --git a/example/wiki/sparse/CMakeLists.txt b/example/wiki/sparse/CMakeLists.txt
@@ -25,13 +25,10 @@ KOKKOSKERNELS_ADD_EXECUTABLE_AND_TEST(
   SOURCES KokkosSparse_wiki_spadd.cpp
   )
 
-# FIXME_SYCL SYCL does not support querying free/total memory
-IF (NOT KOKKOS_ENABLE_SYCL)
 KOKKOSKERNELS_ADD_EXECUTABLE_AND_TEST(
   wiki_spgemm
   SOURCES KokkosSparse_wiki_spgemm.cpp
   )
-ENDIF()
 
 KOKKOSKERNELS_ADD_EXECUTABLE_AND_TEST(
   wiki_gauss_seidel

diff --git a/src/CMakeLists.txt b/src/CMakeLists.txt
@@ -459,6 +459,20 @@ ELSE()
     $<INSTALL_INTERFACE:${KOKKOSKERNELS_HEADER_INSTALL_DIR}>)
 ENDIF()
 
+IF (KOKKOS_ENABLE_SYCL)
+  SET(KOKKOSKERNELS_INTEL_ARCHS ${Kokkos_ARCH})
+  LIST(FILTER KOKKOSKERNELS_INTEL_ARCHS INCLUDE REGEX ".*INTEL.*")
+  LIST(LENGTH KOKKOSKERNELS_INTEL_ARCHS KOKKOSKERNELS_INTEL_ARCHS_NUM)
+  IF(KOKKOSKERNELS_INTEL_ARCHS_NUM GREATER_EQUAL 1)
+    IF (NOT BUILD_SHARED_LIBS)
+      MESSAGE(SEND_ERROR
+        "At the moment, we require KokkosKernels (and Kokkos) to be built as "
+        "shared libraries to allow querying free and total device memory!"
+      )
+    ENDIF()
+    TARGET_LINK_LIBRARIES(kokkoskernels PUBLIC ze_loader)
+  ENDIF()
+ENDIF()
 
 KOKKOSKERNELS_LINK_TPL(kokkoskernels PUBLIC LAPACK)
 KOKKOSKERNELS_LINK_TPL(kokkoskernels PUBLIC BLAS)

diff --git a/src/common/KokkosKernels_ExecSpaceUtils.hpp b/src/common/KokkosKernels_ExecSpaceUtils.hpp
@@ -45,6 +45,11 @@
 #include "Kokkos_Core.hpp"
 #include "Kokkos_Atomic.hpp"
 
+#if defined(KOKKOS_ENABLE_SYCL) && defined(KOKKOS_ARCH_INTEL_GPU)
+#include <level_zero/zes_api.h>
+#include <CL/sycl/backend/level_zero.hpp>
+#endif
+
 #ifndef _KOKKOSKERNELSUTILSEXECSPACEUTILS_HPP
 #define _KOKKOSKERNELSUTILSEXECSPACEUTILS_HPP
 
@@ -205,6 +210,56 @@ inline void kk_get_free_total_memory<Kokkos::Experimental::HIPSpace>(
 }
 #endif
 
+// FIXME_SYCL Use compiler extension instead of low level interface when
+// available. Also, we assume to query memory associated with the default queue.
+#if defined(KOKKOS_ENABLE_SYCL) && defined(KOKKOS_ARCH_INTEL_GPU)
+template <>
+inline void kk_get_free_total_memory<Kokkos::Experimental::SYCLDeviceUSMSpace>(
+    size_t& free_mem, size_t& total_mem) {
+  sycl::queue queue;
+  sycl::device device = queue.get_device();
+  auto level_zero_handle =
+      sycl::get_native<sycl::backend::ext_oneapi_level_zero>(device);
+
+  uint32_t n_memory_modules = 0;
+  zesDeviceEnumMemoryModules(level_zero_handle, &n_memory_modules, nullptr);
+
+  if (n_memory_modules != 1) {
+    std::ostringstream oss;
+    oss << "Error: number of memory modules for the SYCL backend: "
+        << n_memory_modules
+        << ". We only support querying free/total memory if exactly one memory "
+           "module was found. Make sure that ZES_ENABLE_SYSMAN=1 is set at run "
+           "time if no memeory modules were found!";
+    throw std::runtime_error(oss.str());
+  }
+
+  zes_mem_handle_t memory_module_handle;
+  zesDeviceEnumMemoryModules(level_zero_handle, &n_memory_modules,
+                             &memory_module_handle);
+  zes_mem_state_t memory_properties{
+      ZES_STRUCTURE_TYPE_MEM_PROPERTIES,
+  };
+  zesMemoryGetState(memory_module_handle, &memory_properties);
+  total_mem = memory_properties.size;
+  free_mem  = memory_properties.free;
+}
+
+template <>
+inline void kk_get_free_total_memory<Kokkos::Experimental::SYCLHostUSMSpace>(
+    size_t& free_mem, size_t& total_mem) {
+  kk_get_free_total_memory<Kokkos::Experimental::SYCLDeviceUSMSpace>(free_mem,
+                                                                     total_mem);
+}
+
+template <>
+inline void kk_get_free_total_memory<Kokkos::Experimental::SYCLSharedUSMSpace>(
+    size_t& free_mem, size_t& total_mem) {
+  kk_get_free_total_memory<Kokkos::Experimental::SYCLDeviceUSMSpace>(free_mem,
+                                                                     total_mem);
+}
+#endif
+
 inline int kk_get_suggested_vector_size(const size_t nr, const size_t nnz,
                                         const ExecSpaceType exec_space) {
   int suggested_vector_size_ = 1;