使用 pybind11 共享 MPI 通信器

Posted 2023-03-27

【中文标题】使用 pybind11 共享 MPI 通信器

【英文标题】：Sharing an MPI communicator using pybind11

【发布时间】：2019-03-10 11:06:30

【问题描述】：

假设我已经为 MPI 通信器创建了一个包装器：

class Communicator 
   public:
      Communicator() : comm(MPI_COMM_WORLD) 

      Communicator(int const color, int const key) 
        MPI_Comm_split(MPI_COMM_WORLD, color, key, &comm);
      

      Communicator(MPI_Comm comm) : comm(comm) 

      MPI_Comm GetComm() const  return comm; 
    private:
      MPI_Comm comm;
; 

我想使用 pybind11 围绕这个对象创建一个 python 包装器，看起来像这样：

void CommunicatorWrapper(pybind11::module &m) 
   py::class_<Communicator, std::shared_ptr<Communicator> > commWrap(m, "Communicator");

   commWrap.def(py::init( []()  return new Communicator();  ));
   commWrap.def(py::init( [](int const color, int const key)  return new Communicator(color, key);  ));
   commWrap.def(py::init( [](MPI_Comm comm)  return new Communicator(comm);  ));
   commWrap.def("GetComm", &Communicator::GetComm);

但是，我希望 python 看到的 MPI_Comm 类型是 mpi4py.MPI.Comm。这可能吗？如果有，怎么做？

上述（天真的）实现会导致以下行为：

comm = Communicator(MPI.COMM_WORLD)

错误：

TypeError: __init__(): incompatible constructor arguments. The following argument types are supported:
1. Communicator()
2. Communicator(arg0: int, arg1: int)
3. Communicator(arg0: int)

和

comm = Communicator()
print(comm.GetComm())

打印-2080374784。考虑到 MPI_Comm 是什么，这种行为是有意义的，但显然不是我需要的功能。

【问题讨论】：

***.com/questions/51557135/… 似乎相关，但 ob_mpi 对象似乎并不实际存在。

【参考方案1】：

我通过改变包装解决了这个问题

#include <mpi4py/mpi4py.h>

pybind11::handle CallGetComm(Communicator *comm) 
    const int rc = import_mpi4py();
    return pybind11::handle(PyMPIComm_New(comm->GetComm()));;


void CommunicatorWrapper(pybind11::module &m) 
   py::class_<Communicator, std::shared_ptr<Communicator> > commWrap(m, "Communicator");

   commWrap.def(py::init( []()  return new Communicator();  ));
   commWrap.def(py::init( [](int const color, int const key)  return new Communicator(color, key);  ));
   commWrap.def(py::init( [](pybind11::handle const& comm) 
     const int rc = import_mpi4py();
     assert(rc==0);
     return new Communicator(*PyMPIComm_Get(comm.ptr()));
     ));
   commWrap.def("GetComm", &CallGetComm);