Mpi.hpp

/* Copyright 2021 Axel Huebl and Franz Poeschel
 *
 * This file is part of openPMD-api.
 *
 * openPMD-api is free software: you can redistribute it and/or modify
 * it under the terms of of either the GNU General Public License or
 * the GNU Lesser General Public License as published by
 * the Free Software Foundation, either version 3 of the License, or
 * (at your option) any later version.
 *
 * openPMD-api is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License and the GNU Lesser General Public License
 * for more details.
 *
 * You should have received a copy of the GNU General Public License
 * and the GNU Lesser General Public License along with openPMD-api.
 * If not, see <http://www.gnu.org/licenses/>.
 */
 
#pragma once
 
#include "openPMD/config.hpp"
 
#if openPMD_HAVE_MPI
 
#include "openPMD/binding/python/Common.hpp"
 
#include <mpi.h>

struct openPMD_PyMPICommObject
{
    PyObject_HEAD MPI_Comm ob_mpi;
    unsigned int flags;
};
using openPMD_PyMPIIntracommObject = openPMD_PyMPICommObject;
 
struct py_object_to_mpi_comm_error
{
    enum class error_type
    {
        invalid_data,
        is_not_an_mpi_communicator
    };
    using error_type = error_type;
 
    std::string error_msg;
    error_type type;
 
    operator std::string() const
    {
        return error_msg;
    }
};
 
inline std::variant<MPI_Comm, py_object_to_mpi_comm_error>
pythonObjectAsMpiComm(pybind11::object &comm)
{
    namespace py = pybind11;
    // if( import_mpi4py() < 0 ) { here be dragons }
 
    using e = py_object_to_mpi_comm_error;
    using e_t = e::error_type;
 
    if (comm.ptr() == Py_None)
        return e{"MPI communicator cannot be None.", e_t::invalid_data};
    if (comm.ptr() == nullptr)
        return e{"MPI communicator is a nullptr.", e_t::invalid_data};
 
    // check type string to see if this is mpi4py
    //   __str__ (pretty)
    //   __repr__ (unambiguous)
    //   mpi4py: <mpi4py.MPI.Intracomm object at 0x7f998e6e28d0>
    //   pyMPI:  ... (TODO)
    py::str const comm_pystr = py::repr(comm);
    std::string const comm_str = comm_pystr.cast<std::string>();
    if (comm_str.substr(0, 12) != std::string("<mpi4py.MPI."))
        return e{
            "comm is not an mpi4py communicator: " + comm_str,
            e_t::is_not_an_mpi_communicator};
    // only checks same layout, e.g. an `int` in `PyObject` could pass this
    if (!py::isinstance<py::class_<openPMD_PyMPIIntracommObject> >(
            py::type::of(comm)))
        // TODO add mpi4py version from above import check to error message
        return e{
            "comm has unexpected type layout in " + comm_str +
                " (Mismatched MPI at compile vs. runtime? "
                "Breaking mpi4py release?)",
            e_t::invalid_data};
 
    // todo other possible implementations:
    // - pyMPI (inactive since 2008?): import mpi; mpi.WORLD
 
    // reimplementation of mpi4py's:
    // MPI_Comm* mpiCommPtr = PyMPIComm_Get(comm.ptr());
    MPI_Comm *mpiCommPtr =
        &((openPMD_PyMPIIntracommObject *)(comm.ptr()))->ob_mpi;
 
    if (PyErr_Occurred())
        return e{"MPI communicator access error.", e_t::invalid_data};
    if (mpiCommPtr == nullptr)
    {
        return e{
            "MPI communicator cast failed. "
            "(Mismatched MPI at compile vs. runtime?)",
            e_t::invalid_data};
    }
    return {*mpiCommPtr};
}
 
#endif