pykokkos-base icon indicating copy to clipboard operation
pykokkos-base copied to clipboard

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

Open NaderAlAwar opened this issue 3 years ago • 3 comments

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 avatar Sep 09 '22 14:09 NaderAlAwar

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 avatar Sep 19 '22 01:09 jrmadsen

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.

jrmadsen avatar Sep 19 '22 19:09 jrmadsen

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

NaderAlAwar avatar Sep 19 '22 19:09 NaderAlAwar