Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

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

Open
wants to merge 5 commits into
base: main
Choose a base branch
from

Conversation

NaderAlAwar
Copy link
Contributor

@NaderAlAwar NaderAlAwar commented Sep 9, 2022

This PR adds multiple copies of libpykokkos and kokkos to allow importing multiple copies in PyKokkos. This will allow us to initialize each copy to on a different GPU (this idea is borrowed from https://github.com/ut-parla/Parla.py/tree/main/examples/kokkos). Some notes and action items to resolve before merging:

  • Added a CMake option ENABLE_MULTI_GPU=num_gpus. I think this also needs to be added as an option to setup.py
  • After setup is complete, run python multi_gpu_copy.py to create multiple copies of the kokkos lib/ or lib64/ directory.
  • The duplicate gpu packages are called gpu0, gpu1, etc. We probably need better names for them.
  • Everything under tools cannot be imported multiple times even with py::module_local().

@NaderAlAwar NaderAlAwar requested a review from jrmadsen September 9, 2022 14:37
@@ -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")
Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@jrmadsen Should I be adding this option somewhere else?

@jrmadsen
Copy link
Contributor

I am pretty sure the vast majority of this is unnecessary. You don't need to copy everything around in the build tree, you can just make multiple copies when installing:

    # approx line 178 of CMakeLists.txt (or line 204 in modifications)
    INSTALL(FILES ${_OUT_FILE} DESTINATION ${_OUT_PATH})
    # add below
    foreach(_GPU_DIR ${GPU_DIR_LIST})
        INSTALL(FILES ${_OUT_FILE} DESTINATION ${_OUT_PATH}/${_GPU_DIR})
    endforeach()

and INSTALL(...) takes a INSTALL_RPATH field which you can modify -- no need for all that patchelf stuff.

I am not terribly familiar with the py::module_local() behavior but from what I've briefly read about it, it seems like it is going to break other packages returning a Kokkos::View from their python bindings since outside of the pykokkos-base module, those Python type translations are no longer global, but maybe I'm wrong?

@jrmadsen
Copy link
Contributor

Just to reiterate, since I saw you pushed to the repo after my last comment, the entire Python script for making copies and changing the library names and manipulating the RPATH can not only be done in CMake but it is much better to do it in CMake. CMake is the literally the one generating the linker statements that you are manipulating in the post-processing in the first place.

@NaderAlAwar
Copy link
Contributor Author

Didn't see your earlier comment until after I pushed my commit, I'll try to move everything to CMake.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

2 participants