Pybind11：可以使用 mpi4py 吗？

Question

是否可以在Pybind11中在Python端使用mpi4py，然后将communicator交给C++端？

如果是这样，它是如何工作的？

如果不能，是否可以使用 Boost？如果是这样，它将如何完成？

我在网上搜索了几个小时，但没有找到任何东西。

Answer 1

这确实是可以的。正如 John Zwinck 在评论中指出的那样，MPI_COMM_WORLD 将自动指向正确的通信器，因此无需从 python 向 C++ 端传递任何内容。

例子

首先，我们有一个简单的 pybind11 模块，它确实公开了一个函数，该函数简单地打印了一些 MPI 信息（取自许多在线教程之一）。要编译模块，请参见此处 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");
}

接下来 python 一边。在这里，我重复使用 post: Hiding MPI in Python 中的示例，并简单地放入 pybind11 库中。所以首先调用 MPI python 脚本的 python 脚本：

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}')

和子进程文件。这里我们简单调用库函数就可以了

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)

最后的结果是：

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

Answer 2

将 mpi4py 通信器传递给 C++ 使用 pybind11 可以使用 mpi4py C-API。相应的头文件可以使用以下 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 定义了一个包装器 class，它隐式转换为 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

Pybind11：可以使用 mpi4py 吗？

Pybind11: Possible to use mpi4py?

boost

mpi4py

pybind11

例子