Expose find_private_gpu to python (#1943)

* Add find_private_gpu and arbor.env.

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
 ---------------------
 

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

python/env.cpp

+#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.");
+    }
+}

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*.