在 Pybind11 中是否可以在 Python 端使用 mpi4py,然后将通信器交给 C++ 端?
如果是这样,它将如何工作?
如果没有,是否可以使用Boost?如果是这样,它将如何完成?
我在网上搜索了几个小时,但没有找到任何东西。
This is indeed possible. As was pointed out in the comments by John Zwinck, MPI_COMM_WORLD
will automatically point to the correct communicator, so nothing has to be passed from python to the C++ side.
First we have a simple pybind11 module that does expose a single function which simple prints some MPI information (taken from one of the many online tutorials). To compile the module see here pybind11 cmake example.
#include <pybind11/pybind11.h>
#include <mpi.h>
#include <stdio.h>
void say_hi()
{
int world_size;
MPI_Comm_size(MPI_COMM_WORLD, &world_size);
int world_rank;
MPI_Comm_rank(MPI_COMM_WORLD, &world_rank);
char processor_name[MPI_MAX_PROCESSOR_NAME];
int name_len;
MPI_Get_processor_name(processor_name, &name_len);
printf("Hello world from processor %s, rank %d out of %d processors\n",
processor_name,
world_rank,
world_size);
}
PYBIND11_MODULE(mpi_lib, pybind_module)
{
constexpr auto MODULE_DESCRIPTION = "Just testing out mpi with python.";
pybind_module.doc() = MODULE_DESCRIPTION;
pybind_module.def("say_hi", &say_hi, "Each process is allowed to say hi");
}
Next the python side. Here I reuse the example from this post: Hiding MPI in Python and simply put in the pybind11 library. So first the python script that will call the MPI python script:
import sys
import numpy as np
from mpi4py import MPI
def parallel_fun():
comm = MPI.COMM_SELF.Spawn(
sys.executable,
args = ['child.py'],
maxprocs=4)
N = np.array(0, dtype='i')
comm.Reduce(None, [N, MPI.INT], op=MPI.SUM, root=MPI.ROOT)
print(f'We got the magic number {N}')
And the child process file. Here we simply call the library function and it just works.
from mpi4py import MPI
import numpy as np
from mpi_lib import say_hi
comm = MPI.Comm.Get_parent()
N = np.array(comm.Get_rank(), dtype='i')
say_hi()
comm.Reduce([N, MPI.INT], None, op=MPI.SUM, root=0)
The end result is:
from prog import parallel_fun
parallel_fun()
# Hello world from processor arch_zero, rank 1 out of 4 processors
# Hello world from processor arch_zero, rank 2 out of 4 processors
# Hello world from processor arch_zero, rank 0 out of 4 processors
# Hello world from processor arch_zero, rank 3 out of 4 processors
# We got the magic number 6
使用 pybind11将mpi4py通信器传递给 C++可以使用 mpi4py C-API 完成。可以使用以下 Python 代码找到相应的头文件:
import mpi4py
print(mpi4py.get_include())
为了方便地在 Python 和 C++ 之间传递通信器, 可以实现一个自定义的 pybind11 类型的 caster 。为此,我们从典型的序言开始。
// 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
C void*
++ 编译器无法将其与这些类型的常规使用区分开来。为了创建一个不同的类型,我们定义了一个包装类,MPI_Comm
它隐式地转换为 和 from MPI_Comm
。
struct mpi4py_comm {
mpi4py_comm() = default;
mpi4py_comm(MPI_Comm value) : value(value) {}
operator MPI_Comm () { return value; }
MPI_Comm value;
};
然后类型 caster 实现如下:
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
下面是使用类型转换的示例模块的代码。请注意,我们在暴露给 pybind11 的函数定义中使用mpi4py_comm
而不是。MPI_Comm
但是,由于隐式转换,我们可以将这些变量用作常规MPI_Comm
变量。特别是,它们可以传递给任何需要 type 参数的函数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