Передача коммуникатора mpi4py в C++ с использованием pybind11 можно выполнить с помощью C-API mpi4py. Соответствующие заголовочные файлы можно найти с помощью следующего кода Python:
import mpi4py
print(mpi4py.get_include())
Для удобного обмена коммуникаторами между Python и C++ используется настраиваемый генератор типов pybind11 а> можно реализовать. Для этого начнем с типичной преамбулы.
// native.cpp
#include <pybind11/pybind11.h>
#include <mpi.h>
#include <mpi4py/mpi4py.h>
namespace py = pybind11;
Чтобы pybind11 автоматически преобразовывал тип Python в тип C++, нам нужен отдельный тип, который может распознать компилятор C++. К сожалению, стандарт MPI не указывает тип для MPI_comm
. Хуже того, в обычных реализациях MPI MPI_comm
может быть определено как int
или void*
, что компилятор C++ не может отличить от обычного использования этих типов. Чтобы создать отдельный тип, мы определяем класс-оболочку для MPI_Comm
, который неявно преобразует в MPI_Comm
и обратно.
struct mpi4py_comm {
mpi4py_comm() = default;
mpi4py_comm(MPI_Comm value) : value(value) {}
operator MPI_Comm () { return value; }
MPI_Comm value;
};
Затем заклинатель типов реализуется следующим образом:
namespace pybind11 { namespace detail {
template <> struct type_caster<mpi4py_comm> {
public:
PYBIND11_TYPE_CASTER(mpi4py_comm, _("mpi4py_comm"));
// Python -> C++
bool load(handle src, bool) {
PyObject *py_src = src.ptr();
// Check that we have been passed an mpi4py communicator
if (PyObject_TypeCheck(py_src, &PyMPIComm_Type)) {
// Convert to regular MPI communicator
value.value = *PyMPIComm_Get(py_src);
} else {
return false;
}
return !PyErr_Occurred();
}
// C++ -> Python
static handle cast(mpi4py_comm src,
return_value_policy /* policy */,
handle /* parent */)
{
// Create an mpi4py handle
return PyMPIComm_New(src.value);
}
};
}} // namespace pybind11::detail
Ниже приведен код примера модуля, в котором используется заклинатель типов. Обратите внимание, что мы используем mpi4py_comm
вместо MPI_Comm
в определениях функций, предоставляемых pybind11. Однако из-за неявного преобразования мы можем использовать эти переменные как обычные переменные MPI_Comm
. В частности, их можно передать любой функции, ожидающей аргумент типа MPI_Comm
.
// recieve a communicator and check if it equals MPI_COMM_WORLD
void print_comm(mpi4py_comm comm)
{
if (comm == MPI_COMM_WORLD) {
std::cout << "Received the world." << std::endl;
} else {
std::cout << "Received something else." << std::endl;
}
}
mpi4py_comm get_comm()
{
return MPI_COMM_WORLD; // Just return MPI_COMM_WORLD for demonstration
}
PYBIND11_MODULE(native, m)
{
// import the mpi4py API
if (import_mpi4py() < 0) {
throw std::runtime_error("Could not load mpi4py API.");
}
// register the test functions
m.def("print_comm", &print_comm, "Do something with the mpi4py communicator.");
m.def("get_comm", &get_comm, "Return some communicator.");
}
Модуль может быть скомпилирован, например, с использованием
mpicxx -O3 -Wall -shared -std=c++14 -fPIC \
$(python3 -m pybind11 --includes) \
-I$(python3 -c 'import mpi4py; print(mpi4py.get_include())') \
native.cpp -o native$(python3-config --extension-suffix)
и протестировано с использованием
import native
from mpi4py import MPI
import math
native.print_comm(MPI.COMM_WORLD)
# Create a cart communicator for testing
# (MPI_COMM_WORLD.size has to be a square number)
d = math.sqrt(MPI.COMM_WORLD.size)
cart_comm = MPI.COMM_WORLD.Create_cart([d,d], [1,1], False)
native.print_comm(cart_comm)
print(f'native.get_comm() == MPI.COMM_WORLD '
f'-> {native.get_comm() == MPI.COMM_WORLD}')
Вывод должен быть:
Received the world.
Received something else.
native.get_comm() == MPI.COMM_WORLD -> True
person
H. Rittich
schedule
18.06.2020