Setup: add option to build pykokkos-base for multiple GPUs

.gitignore

@@ -61,3 +61,6 @@ timing_report*
 
 # stashed source tree
 /.stash
+
+# multi gpu directories
+/gpu*

CMakeLists.txt

@@ -18,6 +18,19 @@ ENDIF()
 INCLUDE(KokkosPythonUtilities)  # miscellaneous macros and functions
 
 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_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/*")
+    FOREACH(_GPU_ID RANGE ${LAST_GPU_ID})
+        LIST(APPEND GPU_DIR_LIST "gpu${_GPU_ID}")
+        FILE(COPY ${TO_COPY} DESTINATION "${CMAKE_CURRENT_LIST_DIR}/gpu${_GPU_ID}")
+    ENDFOREACH()
+ENDIF()
+
 # force to release if not specified
 IF("${CMAKE_BUILD_TYPE}" STREQUAL "")
     SET(CMAKE_BUILD_TYPE Release CACHE STRING "Build type" FORCE)
@@ -79,12 +92,20 @@ TARGET_LINK_LIBRARIES(libpykokkos PRIVATE
     libpykokkos::precompiled-headers
     libpykokkos::build-options)
 
+SET(Kokkos_GPU_INSTALL_LIBDIRS "")
+
 IF(SKBUILD)
     SET(Kokkos_INSTALL_PYTHONDIR ${CMAKE_INSTALL_PREFIX})
     SET(Kokkos_INSTALL_LIBDIR    ${CMAKE_INSTALL_PREFIX}/kokkos)
+    FOREACH(_GPU_DIR ${GPU_DIR_LIST})
+        LIST(APPEND Kokkos_GPU_INSTALL_LIBDIRS "${CMAKE_INSTALL_PREFIX}/${_GPU_DIR}")
+    ENDFOREACH()
 ELSE()
     SET(Kokkos_INSTALL_PYTHONDIR ${Python3_SITEARCH}/kokkos)
     SET(Kokkos_INSTALL_LIBDIR    ${Python3_SITEARCH}/kokkos)
+    FOREACH(_GPU_DIR ${GPU_DIR_LIST})
+        LIST(APPEND Kokkos_GPU_INSTALL_LIBDIRS "${Python3_SITEARCH}/${_GPU_DIR}")
+    ENDFOREACH()
 ENDIF()
 
 # figure out if we can install to Python3_SITEARCH
@@ -155,14 +176,19 @@ ENDIF()
 INSTALL(TARGETS libpykokkos
     DESTINATION ${Kokkos_INSTALL_LIBDIR})
 
+FOREACH(_FILE ${Kokkos_GPU_INSTALL_LIBDIRS})
+    INSTALL(TARGETS libpykokkos
+        DESTINATION ${_FILE})
+ENDFOREACH()
+
 INSTALL(FILES ${PROJECT_BINARY_DIR}/kokkos/__init__.py
     DESTINATION ${Kokkos_INSTALL_PYTHONDIR})
 
 INSTALL(FILES ${PROJECT_BINARY_DIR}/pytest.ini
     DESTINATION ${Kokkos_INSTALL_PYTHONDIR})
 
 # glob any python package files
-FILE(GLOB_RECURSE PYPACKAGE_FILES ${CMAKE_CURRENT_LIST_DIR}/kokkos/*.py*)
+FILE(GLOB_RECURSE PYPACKAGE_FILES ${CMAKE_CURRENT_LIST_DIR}/kokkos/*.py* ${CMAKE_CURRENT_LIST_DIR}/gpu*/*.py*)
 FOREACH(_FILE ${PYPACKAGE_FILES})
     # make it a relative path
     STRING(REPLACE "${CMAKE_CURRENT_LIST_DIR}/" "" _OUT_NAME "${_FILE}")

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

include/libpykokkos.hpp

@@ -50,7 +50,9 @@
 
 namespace py = pybind11;
 
+#ifndef ENABLE_MULTI_GPU
 void generate_tools(py::module& kokkos);
+#endif
 void generate_available(py::module& kokkos);
 void generate_enumeration(py::module& kokkos);
 void generate_view_variants(py::module& kokkos);

include/pools.hpp

@@ -57,7 +57,11 @@ void generate_pool(py::module &_mod, const std::string &_name,
 
   // using PoolT = Kokkos::Random_XorShift64_Pool<Kokkos::Cuda>;
   // class decl
+#ifdef ENABLE_MULTI_GPU
+  py::class_<PoolT> _pool(_mod, _name.c_str(), py::module_local());
+#else
   py::class_<PoolT> _pool(_mod, _name.c_str());
+#endif
 
   // default initializer
   _pool.def(py::init([]() { return new PoolT{}; }));

include/views.hpp

@@ -308,8 +308,11 @@ void generate_view(py::module &_mod, const std::string &_name,
               << "'..." << std::endl;
 
   // class decl
+#ifdef ENABLE_MULTI_GPU
+  py::class_<ViewT> _view(_mod, _name.c_str(), py::buffer_protocol(), py::module_local());
+#else
   py::class_<ViewT> _view(_mod, _name.c_str(), py::buffer_protocol());
-
+#endif
   // default initializer
   _view.def(py::init([]() { return new ViewT{}; }));
 

multi_gpu_copy.py

@@ -0,0 +1,87 @@
+# This should only be run after kokkos is installed
+
+import glob
+from pathlib import Path
+import shutil
+import subprocess
+import sys
+
+import kokkos
+
+if shutil.which("patchelf") is None:
+    sys.exit("ERROR: Cannot run multi_gpu_copy.py without 'patchelf'")
+
+kokkos_path = kokkos.__path__[0]
+base_path = Path(kokkos_path).parent
+
+lib_path = None
+if (base_path / "lib").is_dir():
+    lib_path = base_path / "lib"
+elif (base_path / "lib64").is_dir():
+    lib_path = base_path / "lib64"
+
+assert(lib_path is not None)
+package_paths = [kokkos_path] + glob.glob(f"{str(base_path)}/gpu*")
+
+for package in package_paths:
+    package_path = Path(package)
+    package_lib_path = package_path / "lib"
+    shutil.copytree(lib_path, package_lib_path, dirs_exist_ok=True)
+
+    if package.endswith("kokkos"):
+        continue
+
+    package_id = package[-1]
+
+    libkokkoscore_remove = None
+    libkokkoscore_add = None
+    libkokkoscontainers_remove = None
+    libkokkoscontainers_add = None
+
+    for lib in package_lib_path.iterdir():
+        if lib.name == "cmake":
+            continue
+
+        # Add the suffix to the end of each copy
+        suffix = lib.name.split(".")
+        lib_name = suffix[0]
+        suffix = suffix[1:]
+        suffix = ".".join(suffix)
+
+        new_name = f"{lib_name}_{package_id}.{suffix}"
+
+        if new_name.count(".") == 3: # this is the library that's listed as a dependency
+            if "libkokkoscore" in new_name:
+                libkokkoscore_remove = lib.name
+                libkokkoscore_add = new_name
+            elif "libkokkoscontainers" in new_name:
+                libkokkoscontainers_remove = lib.name
+                libkokkoscontainers_add = new_name
+
+        new_lib_path = Path(lib.parent) / new_name
+        lib.rename(new_lib_path)
+
+        # Add the suffix to the end of the SONAME
+        so_name = subprocess.run(["patchelf", "--print-soname", new_lib_path], capture_output=True).stdout.decode("utf-8")
+        so_name = so_name.split(".")
+        lib_name = so_name[0]
+        suffix = so_name[1:]
+        suffix = [s.strip() for s in suffix]
+        suffix = ".".join(suffix)
+
+        new_so_name = f"{lib_name}_{package_id}.{suffix}"
+        subprocess.run(["patchelf", "--set-soname", new_so_name, new_lib_path])
+
+    for file in package_path.iterdir():
+        if "libpykokkos" in file.name:
+            libpykokkos_path = file
+            break
+
+    assert(libkokkoscore_remove is not None)
+    assert(libkokkoscore_add is not None)
+    assert(libkokkoscontainers_remove is not None)
+    assert(libkokkoscontainers_add is not None)
+
+    subprocess.run(["patchelf", "--replace-needed", libkokkoscore_remove, libkokkoscore_add, libpykokkos_path])
+    subprocess.run(["patchelf", "--replace-needed", libkokkoscontainers_remove, libkokkoscontainers_add, libpykokkos_path])
+    subprocess.run(["patchelf", "--set-rpath", package_lib_path, libpykokkos_path])

setup.py

@@ -5,6 +5,9 @@
 import argparse
 import warnings
 import platform
+from pathlib import Path
+import shutil
+import glob
 from skbuild import setup
 
 # some Cray systems default to static libraries and the build

src/enumeration.cpp

@@ -129,8 +129,13 @@ void generate_enumeration(py::module &kokkos) {
   //                                execution spaces
   //
   //----------------------------------------------------------------------------//
+#ifdef ENABLE_MULTI_GPU
+  py::enum_<KokkosExecutionSpace> _device(kokkos, "device",
+                                          "Device execution spaces", py::module_local());
+#else
   py::enum_<KokkosExecutionSpace> _device(kokkos, "device",
                                           "Device execution spaces");
+#endif
   generate_enumeration<ExecutionSpaceSpecialization>(
       _device, std::make_index_sequence<ExecutionSpacesEnd>{});
 
@@ -164,7 +169,11 @@ void generate_enumeration(py::module &kokkos) {
   //
   //----------------------------------------------------------------------------//
   // an enumeration of the data types for views
+#ifdef ENABLE_MULTI_GPU
+  py::enum_<KokkosViewDataType> _dtype(kokkos, "dtype", "View data types", py::module_local());
+#else
   py::enum_<KokkosViewDataType> _dtype(kokkos, "dtype", "View data types");
+#endif
   generate_enumeration<ViewDataTypeSpecialization>(
       _dtype, std::make_index_sequence<ViewDataTypesEnd>{});
   _dtype.export_values();
@@ -191,8 +200,13 @@ void generate_enumeration(py::module &kokkos) {
   //
   //----------------------------------------------------------------------------//
   // an enumeration of the memory spaces for views
+#ifdef ENABLE_MULTI_GPU
+  py::enum_<KokkosMemorySpace> _memspace(kokkos, "memory_space",
+                                         "View memory spaces", py::module_local());
+#else
   py::enum_<KokkosMemorySpace> _memspace(kokkos, "memory_space",
                                          "View memory spaces");
+#endif
   generate_enumeration<MemorySpaceSpecialization>(
       _memspace, std::make_index_sequence<MemorySpacesEnd>{});
   _memspace.value(
@@ -227,8 +241,13 @@ void generate_enumeration(py::module &kokkos) {
   //
   //----------------------------------------------------------------------------//
   // an enumeration of the layout types for views
+#ifdef ENABLE_MULTI_GPU
+  py::enum_<KokkosMemoryLayoutType> _ltype(kokkos, "layout",
+                                           "View layout types", py::module_local());
+#else
   py::enum_<KokkosMemoryLayoutType> _ltype(kokkos, "layout",
                                            "View layout types");
+#endif
   generate_enumeration<MemoryLayoutSpecialization>(
       _ltype, std::make_index_sequence<MemoryLayoutEnd>{});
   _ltype.export_values();
@@ -255,8 +274,13 @@ void generate_enumeration(py::module &kokkos) {
   //
   //----------------------------------------------------------------------------//
   // an enumeration of the memory traits for views
+#ifdef ENABLE_MULTI_GPU
+  py::enum_<KokkosMemoryTrait> _memtrait(kokkos, "memory_trait",
+                                         "View memory traits", py::module_local());
+#else
   py::enum_<KokkosMemoryTrait> _memtrait(kokkos, "memory_trait",
                                          "View memory traits");
+#endif
   generate_enumeration<MemoryTraitSpecialization>(
       _memtrait, std::make_index_sequence<MemoryTraitEnd>{});
   _memtrait.export_values();

src/libpykokkos.cpp

@@ -107,7 +107,9 @@ PYBIND11_MODULE(libpykokkos, kokkos) {
   kokkos.def("initialize", _initialize, "Initialize Kokkos");
   kokkos.def("finalize", _finalize, "Finalize Kokkos");
 
+#ifndef ENABLE_MULTI_GPU
   generate_tools(kokkos);
+#endif
   generate_available(kokkos);
   generate_enumeration(kokkos);
   generate_view_variants(kokkos);

src/tools.cpp

@@ -313,6 +313,7 @@ void internal_setup();
 //
 void destroy_callbacks() { callbacks.reset(); }
 
+#ifndef ENABLE_MULTI_GPU
 void generate_tools(py::module& kokkos) {
   //--------------------------------------------------------------------//
   //
@@ -529,3 +530,4 @@ void internal_test() {
   get_src_view().reset();
   get_dst_view().reset();
 }
+#endif