From db887a9e4d5a986fa8f03c80dcf71eaf140588ea Mon Sep 17 00:00:00 2001
From: Thorsten Hater <24411438+thorstenhater@users.noreply.github.com>
Date: Thu, 11 Aug 2022 20:43:43 +0200
Subject: [PATCH] Expose find_private_gpu to python (#1943)
* Add find_private_gpu and arbor.env.
---
doc/python/hardware.rst | 14 ++++++++++++++
python/CMakeLists.txt | 1 +
python/env.cpp | 35 +++++++++++++++++++++++++++++++++++
python/pyarb.cpp | 2 ++
4 files changed, 52 insertions(+)
create mode 100644 python/env.cpp
diff --git a/doc/python/hardware.rst b/doc/python/hardware.rst
index df512c75..76969f73 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 10e7e246..e2d6bdb5 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 00000000..e1940781
--- /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 c71293e8..57e503b0 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*.
--
GitLab