Skip to content
Snippets Groups Projects
Select Git revision
  • 27e0d16eef2817956977adc2c78ebbefed314e6c
  • master default protected
  • update-py-siibra
  • update-libneuroml
  • update-pyunicore
  • update-py-netpyne
  • update-py-ebrains-drive
  • update-bluebrain-packages
  • experimental_rel
  • fixup_fix_wf-brainscales2-demos
  • feat_arbor_install_all_binaries
  • ebrains-23.09-jsc-site-config
  • spack-v0.20.0
  • create-module-file
  • ebrains-23-09-spack-v0.19.2
  • ebrains-23-09
  • nestml-source-distribution
  • ebrains-23-06
  • ebrains-23-02
  • ebrains-22-10
  • ebrains-22-07
21 results

create_JupyterLab_kernel.sh

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