diff --git a/.gitignore b/.gitignore index 61d365c..b83aea7 100644 --- a/.gitignore +++ b/.gitignore @@ -61,3 +61,6 @@ timing_report* # stashed source tree /.stash + +# multi gpu directories +/gpu* \ No newline at end of file diff --git a/CMakeLists.txt b/CMakeLists.txt index 80482b7..855a299 100644 --- a/CMakeLists.txt +++ b/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,6 +176,11 @@ 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}) @@ -162,7 +188,7 @@ 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}") diff --git a/include/execution_spaces.hpp b/include/execution_spaces.hpp index 669ed88..c89ba80 100644 --- 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_ _space(_mod, _name.c_str(), py::module_local()); +#else py::class_ _space(_mod, _name.c_str()); +#endif _space.def(py::init([]() { return new Sp{}; })); // Add other constructors with arguments if they exist diff --git a/include/libpykokkos.hpp b/include/libpykokkos.hpp index 18aed0e..fc6d404 100644 --- a/include/libpykokkos.hpp +++ b/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); diff --git a/include/pools.hpp b/include/pools.hpp index 5548a34..f61e393 100644 --- a/include/pools.hpp +++ b/include/pools.hpp @@ -57,7 +57,11 @@ void generate_pool(py::module &_mod, const std::string &_name, // using PoolT = Kokkos::Random_XorShift64_Pool; // class decl +#ifdef ENABLE_MULTI_GPU + py::class_ _pool(_mod, _name.c_str(), py::module_local()); +#else py::class_ _pool(_mod, _name.c_str()); +#endif // default initializer _pool.def(py::init([]() { return new PoolT{}; })); diff --git a/include/views.hpp b/include/views.hpp index b7c65f2..3ffb38a 100644 --- a/include/views.hpp +++ b/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_ _view(_mod, _name.c_str(), py::buffer_protocol(), py::module_local()); +#else py::class_ _view(_mod, _name.c_str(), py::buffer_protocol()); - +#endif // default initializer _view.def(py::init([]() { return new ViewT{}; })); diff --git a/multi_gpu_copy.py b/multi_gpu_copy.py new file mode 100644 index 0000000..56baba6 --- /dev/null +++ b/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]) \ No newline at end of file diff --git a/setup.py b/setup.py index e740ab0..3af8bb8 100644 --- a/setup.py +++ b/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 diff --git a/src/enumeration.cpp b/src/enumeration.cpp index 4e9c7be..ca8967c 100644 --- a/src/enumeration.cpp +++ b/src/enumeration.cpp @@ -129,8 +129,13 @@ void generate_enumeration(py::module &kokkos) { // execution spaces // //----------------------------------------------------------------------------// +#ifdef ENABLE_MULTI_GPU + py::enum_ _device(kokkos, "device", + "Device execution spaces", py::module_local()); +#else py::enum_ _device(kokkos, "device", "Device execution spaces"); +#endif generate_enumeration( _device, std::make_index_sequence{}); @@ -164,7 +169,11 @@ void generate_enumeration(py::module &kokkos) { // //----------------------------------------------------------------------------// // an enumeration of the data types for views +#ifdef ENABLE_MULTI_GPU + py::enum_ _dtype(kokkos, "dtype", "View data types", py::module_local()); +#else py::enum_ _dtype(kokkos, "dtype", "View data types"); +#endif generate_enumeration( _dtype, std::make_index_sequence{}); _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_ _memspace(kokkos, "memory_space", + "View memory spaces", py::module_local()); +#else py::enum_ _memspace(kokkos, "memory_space", "View memory spaces"); +#endif generate_enumeration( _memspace, std::make_index_sequence{}); _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_ _ltype(kokkos, "layout", + "View layout types", py::module_local()); +#else py::enum_ _ltype(kokkos, "layout", "View layout types"); +#endif generate_enumeration( _ltype, std::make_index_sequence{}); _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_ _memtrait(kokkos, "memory_trait", + "View memory traits", py::module_local()); +#else py::enum_ _memtrait(kokkos, "memory_trait", "View memory traits"); +#endif generate_enumeration( _memtrait, std::make_index_sequence{}); _memtrait.export_values(); diff --git a/src/libpykokkos.cpp b/src/libpykokkos.cpp index f3305e2..8bd253b 100644 --- a/src/libpykokkos.cpp +++ b/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); diff --git a/src/tools.cpp b/src/tools.cpp index 9be90a9..bf1ab26 100644 --- a/src/tools.cpp +++ b/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