multi_gpu: mark execution spaces as module_local and fix cmake check …

…for number of gpus
kokkos · Oct 26, 2023 · a3a76af · a3a76af
1 parent aed4237
commit a3a76af
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 1 deletion.
diff --git a/CMakeLists.txt b/CMakeLists.txt
@@ -21,7 +21,7 @@ ADD_OPTION(BUILD_SHARED_LIBS "Build shared libraries" ON)
 SET(ENABLE_MULTI_GPU "0" CACHE STRING "Build multiple copies to enable multi GPU")
 SET(GPU_DIR_LIST "")
 
-IF(${ENABLE_MULTI_GPU} GREATER 2)
+IF(${ENABLE_MULTI_GPU} GREATER_EQUAL 2)
     ADD_COMPILE_DEFINITIONS(ENABLE_MULTI_GPU)
     MATH(EXPR LAST_GPU_ID "${ENABLE_MULTI_GPU} - 1")
     FILE(GLOB TO_COPY "${CMAKE_CURRENT_LIST_DIR}/kokkos/*")

diff --git a/include/execution_spaces.hpp b/include/execution_spaces.hpp
@@ -93,7 +93,11 @@ void generate_execution_space(py::module &_mod, const std::string &_name,
     std::cerr << "Registering " << _msg << " as python class '" << _name
               << "'..." << std::endl;
 
+#ifdef ENABLE_MULTI_GPU
+  py::class_<Sp> _space(_mod, _name.c_str(), py::module_local());
+#else
   py::class_<Sp> _space(_mod, _name.c_str());
+#endif
   _space.def(py::init([]() { return new Sp{}; }));
 
   // Add other constructors with arguments if they exist