kokkos / pykokkos-base

Python bindings for data interoperability with Kokkos (View, DynRankView)
Other
26 stars 9 forks source link

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

Open NaderAlAwar opened 2 years ago

NaderAlAwar commented 2 years ago

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:

jrmadsen commented 2 years ago

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 commented 2 years ago

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 commented 2 years ago

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