Skip to content
Snippets Groups Projects
Select Git revision
  • 148e8d658c58936d20050814f39b40ca6961c790
  • master default protected
  • tut_ring_allen
  • docs_furo
  • docs_reorder_cable_cell
  • docs_graphviz
  • docs_rtd_dev
  • ebrains_mirror
  • doc_recat
  • docs_spike_source
  • docs_sim_sample_clar
  • docs_pip_warn
  • github_template_updates
  • docs_fix_link
  • cv_default_and_doc_clarification
  • docs_add_numpy_req
  • readme_zenodo_05
  • install_python_fix
  • install_require_numpy
  • typofix_propetries
  • docs_recipe_lookup
  • v0.10.0
  • v0.10.1
  • v0.10.0-rc5
  • v0.10.0-rc4
  • v0.10.0-rc3
  • v0.10.0-rc2
  • v0.10.0-rc
  • v0.9.0
  • v0.9.0-rc
  • v0.8.1
  • v0.8
  • v0.8-rc
  • v0.7
  • v0.6
  • v0.5.2
  • v0.5.1
  • v0.5
  • v0.4
  • v0.3
  • v0.2.2
41 results

profiler.cpp

Blame
  • mpi.cpp 3.12 KiB
    #include <sstream>
    #include <string>
    
    #include <arbor/version.hpp>
    
    #include <pybind11/pybind11.h>
    
    #ifdef ARB_MPI_ENABLED
    #include <arbor/communication/mpi_error.hpp>
    
    #include <mpi.h>
    
    #include "mpi.hpp"
    
    #ifdef ARB_WITH_MPI4PY
    #include <mpi4py/mpi4py.h>
    #endif
    #endif
    
    namespace pyarb {
    
    // Some helper functions to determine how Arbor was compiled
    bool mpi_compiled() {
        #ifdef ARB_MPI_ENABLED
        return true;
        #else
        return false;
        #endif
    }
    
    bool mpi4py_compiled() {
        #ifdef ARB_WITH_MPI4PY
        return true;
        #else
        return false;
        #endif
    }
    
    void register_queryenvvars(pybind11::module& m) {
        using namespace std::string_literals;
        
        m.def("mpi_compiled", &mpi_compiled, "Check if Arbor was compiled with MPI.");
        m.def("mpi4py_compiled", &mpi4py_compiled, "Check if Arbor was compiled with MPI4PY.");
    }
        
    #ifdef ARB_MPI_ENABLED
    #ifdef ARB_WITH_MPI4PY
    
    mpi_comm_shim comm_from_mpi4py(pybind11::object& o) {
        import_mpi4py();
    
        // If object is already a mpi4py communicator, return
        if (PyObject_TypeCheck(o.ptr(), &PyMPIComm_Type)) {
            return mpi_comm_shim(*PyMPIComm_Get(o.ptr()));
        }
        throw arb::mpi_error(MPI_ERR_OTHER, "The argument is not an mpi4py communicator");
    }
    
    #endif
    
    // Some helper functions for initializing and finalizing MPI.
    // Arbor requires at least MPI_THREAD_SERIALIZED, because the communication task
    // that uses MPI can be run on any thread, and there will never be more than one
    // concurrent communication task.
    
    void mpi_init() {
        int provided = MPI_THREAD_SINGLE;
        int ev = MPI_Init_thread(nullptr, nullptr, MPI_THREAD_SERIALIZED, &provided);
        if (ev) {
            throw arb::mpi_error(ev, "MPI_Init_thread");
        }
        else if (provided<MPI_THREAD_SERIALIZED) {
            throw arb::mpi_error(MPI_ERR_OTHER, "MPI_Init_thread: MPI_THREAD_SERIALIZED unsupported");
        }
    }
    
    void mpi_finalize() {
        MPI_Finalize();
    }
    
    int mpi_is_initialized() {
        int initialized;
        MPI_Initialized(&initialized);
        return initialized; 
    }
    
    int mpi_is_finalized() {
        int finalized;
        MPI_Finalized(&finalized);
        return finalized;
    }
    // Define the stringifier for mpi_comm_shim here, to minimise the ifdefication
    // elsewhere in this wrapper code.
    
    std::string mpi_comm_string(const mpi_comm_shim& c) {
        std::stringstream s;
    
        s << "<mpi_comm: ";
        if (c.comm==MPI_COMM_WORLD) s << "MPI_COMM_WORLD>";
        else s << c.comm << ">";
        return s.str();
    }
    
    void register_mpi(pybind11::module& m) {
        using namespace std::string_literals;
    
        pybind11::class_<mpi_comm_shim> mpi_comm(m, "mpi_comm");
        mpi_comm
            .def(pybind11::init<>())
            .def("__str__", &mpi_comm_string)
            .def("__repr__", &mpi_comm_string);
    
        m.def("mpi_init", &mpi_init, "Initialize MPI with MPI_THREAD_SINGLE, as required by Arbor.");
        m.def("mpi_finalize", &mpi_finalize, "Finalize MPI (calls MPI_Finalize)");
        m.def("mpi_is_initialized", &mpi_is_initialized, "Check if MPI is initialized.");
        m.def("mpi_is_finalized", &mpi_is_finalized, "Check if MPI is finalized.");
    
        #ifdef ARB_WITH_MPI4PY
        m.def("mpi_comm_from_mpi4py", comm_from_mpi4py);
        #endif
    }
    #endif
    } // namespace pyarb