diff --git a/doc/python/hardware.rst b/doc/python/hardware.rst index df512c75a4555d4363b35aeed54dfa230ea00382..76969f735963de95d28c806008967d290c2ea7de 100644 --- a/doc/python/hardware.rst +++ b/doc/python/hardware.rst @@ -68,6 +68,20 @@ Helper functions for checking cmake or environment variables, as well as configu Check if MPI is finalized. +Env: Helper functions +--------------------- + +The ``arbor.env`` module collects helper functions for interacting with the environment. + +.. function:: find_private_gpu(comm) + + Requires GPU and MPI. Will return an integer id of a GPU such that each GPU + is mapped to at most one MPI task (on the same node as the GPU). Raises an + exception if + - not built with GPU or MPI support + - unable to satisfy the constraints above + - handed an invalid or unknown MPI communicator object + Prescribed resources --------------------- diff --git a/python/CMakeLists.txt b/python/CMakeLists.txt index 10e7e246e8285b2f5d92aa9eb4b2aa53e4202060..e2d6bdb54de3d36ceaadedd381d292acd1e24377 100644 --- a/python/CMakeLists.txt +++ b/python/CMakeLists.txt @@ -38,6 +38,7 @@ set(pyarb_source schedule.cpp simulation.cpp single_cell_model.cpp + env.cpp ) # compile the pyarb sources into an object library that will be diff --git a/python/env.cpp b/python/env.cpp new file mode 100644 index 0000000000000000000000000000000000000000..e194078169783d9ad49497664263146f2dd80ce1 --- /dev/null +++ b/python/env.cpp @@ -0,0 +1,35 @@ +#include <pybind11/pybind11.h> + +#include <arborenv/gpu_env.hpp> + +#include "mpi.hpp" +#include "error.hpp" + +namespace pyarb { + + void register_arborenv(pybind11::module& m) { + auto s = m.def_submodule("env", "Wrappers for arborenv."); + s.def("find_private_gpu", + [] (pybind11::object mpi) { +#ifndef ARB_GPU_ENABLED + throw pyarb_error("Private GPU: Arbor is not configured with GPU support."); +#else +#ifndef ARB_MPI_ENABLED + throw pyarb_error("Private GPU: Arbor is not configured with MPI."); +#else + auto err = ""Private GPU: Invalid MPI Communicator.""; + if (can_convert_to_mpi_comm(mpi)) { + return arbenv::find_private_gpu(can_convert_to_mpi_comm(mpi)); + } + else if (auto c = py2optional<mpi_comm_shim>(mpi, err)) { + return arbenv::find_private_gpu(c->comm); + } else { + throw pyarb_error(err); + } +#endif +#endif + }, + "Identify a private GPU id per node, only available if built with GPU and MPI.\n" + " mpi: The MPI communicator."); + } +} diff --git a/python/pyarb.cpp b/python/pyarb.cpp index c71293e8a0db3613d04e764b7a7cf6e55473170c..57e503b07afb19bb966fbbff78c3fc978a417774 100644 --- a/python/pyarb.cpp +++ b/python/pyarb.cpp @@ -29,6 +29,7 @@ void register_recipe(pybind11::module& m); void register_schedules(pybind11::module& m); void register_simulation(pybind11::module& m, pyarb_global_ptr); void register_single_cell(pybind11::module& m); +void register_arborenv(pybind11::module& m); #ifdef ARB_MPI_ENABLED void register_mpi(pybind11::module& m); @@ -61,6 +62,7 @@ PYBIND11_MODULE(_arbor, m) { pyarb::register_schedules(m); pyarb::register_simulation(m, global_ptr); pyarb::register_single_cell(m); + pyarb::register_arborenv(m); // This is the fallback. All specific translators take precedence by being // registered *later*.