Merge pyDIY changes to MPI communicators

bindings/python/diy/__init__.py

@@ -2,15 +2,19 @@ from ._diy import *
 
 # Monkey-patch the constructor to accept a normal mpi4py communicator
 
-init = Master.__init__
-def convert_mpi_comm(self, comm, *args, **kwargs):
-    if not isinstance(comm, mpi.MPIComm):
-        from mpi4py import MPI
-        a = MPI._addressof(comm)
+init = mpi.MPIComm.__init__
+def convert_mpi_comm(self, *args, **kwargs):
+    if len(args) == 0:
+        init(self, *args, **kwargs)
     else:
-        a = comm.comm()
-    init(self, a, *args, **kwargs)
-Master.__init__ = convert_mpi_comm
+        comm = args[0]
+
+        if not isinstance(comm, mpi.MPIComm):
+            from mpi4py import MPI
+            comm = MPI._addressof(comm)
+
+        init(self, comm, *args[1:], **kwargs)
+mpi.MPIComm.__init__ = convert_mpi_comm
 
 class InitFinalize:
     def __init__(self):

bindings/python/examples/mpi.py

@@ -10,7 +10,8 @@ import diy
 # can get comm either from diy or mpi4py, in this case exercising mpi4py
 comm = MPI.COMM_WORLD
 
-m = diy.Master(comm)               # master
+c = diy.mpi.MPIComm(comm)
+m = diy.Master(c)               # master
 
 # try some more mpi4py to get my rank and size of comm
 rank = comm.Get_rank()

bindings/python/src/assigner.cpp

@@ -5,6 +5,7 @@
 namespace py = pybind11;
 
 #include <diy/assigner.hpp>
+#include <diy/resolve.hpp>
 using namespace diy;
 
 void init_assigner(py::module& m)
@@ -50,4 +51,8 @@ void init_assigner(py::module& m)
         .def("set_ranks",   &DynamicAssigner::set_ranks)
         .def("rank_offset", &DynamicAssigner::rank_offset)
     ;
+
+    m.def("fix_links", &fix_links);
+    m.def("record_local_gids", &record_local_gids);
+    m.def("update_links", &update_links);
 }

bindings/python/src/master.cpp

@@ -6,6 +6,7 @@ namespace py = pybind11;
 using namespace diy;
 
 #include "serialization.h"
+#include "mpi.hpp"
 
 void init_master(py::module& m)
 {
@@ -23,7 +24,7 @@ void init_master(py::module& m)
     //py::class_<FileStorage>(m, "FileStorage");
 
     py::class_<Master>(m, "Master")
-      .def(py::init([](long                 comm_,
+      .def(py::init([](mpi::communicator    comm,
                        int                  threads,
                        int                  limit,
                        Create               create,
@@ -32,7 +33,6 @@ void init_master(py::module& m)
                        Load                 load)
                        //FileStorage*         storage)
                        {
-                           MPI_Comm comm = *static_cast<MPI_Comm*>(reinterpret_cast<void*>(comm_));
                            return new Master(comm,
                                              threads,
                                              limit,
@@ -48,7 +48,7 @@ void init_master(py::module& m)
                        }),
                        "comm"_a, "threads"_a = 1, "limit"_a = -1,
                        "create"_a = Create([]() -> py::object { return py::none(); }),
-                       "destroy"_a = Destroy([](py::object* b) { }),
+                       "destroy"_a = Destroy([](py::object*) { }),
                        "save"_a = Save([](const py::object* b, diy::BinaryBuffer* bb) { diy::save(*bb, *b); }),
                        "load"_a = Load([](diy::BinaryBuffer* bb) -> py::object
                                        {

bindings/python/src/mpi.cpp

@@ -4,27 +4,21 @@ namespace py = pybind11;
 
 #include <mpi.h>
 
-struct PyMPIComm
-{
-                PyMPIComm(MPI_Comm comm_ = MPI_COMM_WORLD):
-                    comm(comm_)
-    {
-        MPI_Comm_size(comm, &size);
-        MPI_Comm_rank(comm, &rank);
-    }
-
-    MPI_Comm    comm;
-    int         size;
-    int         rank;
-};
+#include <diy/mpi/communicator.hpp>
 
 void init_mpi(py::module& m)
 {
+    using PyMPIComm = diy::mpi::communicator;
+
     py::class_<PyMPIComm>(m, "MPIComm")
         .def(py::init<>())
-        .def_readonly("size", &PyMPIComm::size)
-        .def_readonly("rank", &PyMPIComm::rank)
-        .def("comm",          [](const PyMPIComm& comm) { return (long) &comm.comm; })
+        .def(py::init([](long comm_)
+             {
+                return new diy::mpi::communicator(*static_cast<MPI_Comm*>(reinterpret_cast<void*>(comm_)));
+             }))
+        .def_property_readonly("size", &PyMPIComm::size)
+        .def_property_readonly("rank", &PyMPIComm::rank)
+        .def_property_readonly("comm", &PyMPIComm::handle)
     ;
 
     m.def("init",       []() { int argc = 0; char** argv = 0; MPI_Init(&argc, &argv); });