Homec4science
Diffusion tamaas (ci-fix-static-checks) Browse

tamaas/python/wrap/mpi.cpp
9f0fc6dcf47fci-fix-static-checks

Last Change
Raw File

mpi.cpp
View Options

/*
* SPDX-License-Indentifier: AGPL-3.0-or-later
*
* Copyright (©) 2016-2022 EPFL (École Polytechnique Fédérale de Lausanne),
* Laboratory (LSMS - Laboratoire de Simulation en Mécanique des Solides)
*
* This program is free software: you can redistribute it and/or modify
* it under the terms of the GNU Affero General Public License as published
* by the Free Software Foundation, either version 3 of the License, or
* (at your option) any later version.
*
* This program 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 Affero General Public License for more details.
*
* You should have received a copy of the GNU Affero General Public License
* along with this program. If not, see <https://www.gnu.org/licenses/>.
*
*/
/* -------------------------------------------------------------------------- */
#include "mpi_interface.hh"
#include "partitioner.hh"
#include "wrap.hh"
#include <pybind11/stl.h>
/* -------------------------------------------------------------------------- */
namespace tamaas {
/* -------------------------------------------------------------------------- */
namespace wrap {
using namespace py::literals;
/* -------------------------------------------------------------------------- */
void wrapMPI(py::module& mod) {
auto mpi_mod = mod.def_submodule("mpi");
py::class_<mpi::sequential>(mpi_mod, "sequential")
.def(py::init<>())
.def("__enter__",
[](mpi::sequential& /*obj*/) { mpi::sequential::enter(); })
.def("__exit__", [](mpi::sequential& /*obj*/, py::object /*unused*/,
py::object /*unused*/,
py::object /*unused*/) { mpi::sequential::exit(); });
mpi_mod.def(
"local_shape",
[](std::vector<UInt> global) {
switch (global.size()) {
case 1:
return Partitioner<1>::local_size(global);
case 2:
return Partitioner<2>::local_size(global);
default:
TAMAAS_EXCEPTION("Please provide a 1D/2D shape");
}
},
"global_shape"_a, "Gives the local size of a 1D/2D global shape");
mpi_mod.def(
"global_shape",
[](std::vector<UInt> local) {
switch (local.size()) {
case 1:
return Partitioner<1>::global_size(local);
case 2:
return Partitioner<2>::global_size(local);
default:
TAMAAS_EXCEPTION("Please provide a 1D/2D shape");
}
},
"local_shape"_a, "Gives the global shape of a 1D/2D local shape");
mpi_mod.def(
"local_offset",
[](std::vector<UInt> global) {
switch (global.size()) {
case 1:
return Partitioner<1>::local_offset(global);
case 2:
return Partitioner<2>::local_offset(global);
default:
TAMAAS_EXCEPTION("Please provide a 1D/2D shape");
}
},
"global_shape"_a, "Gives the local offset of a 1D/2D global shape");
mpi_mod.def("gather", &Partitioner<2>::gather<Real>,
py::return_value_policy::move, "Gather 2D surfaces");
mpi_mod.def("scatter", &Partitioner<2>::scatter<Real>,
py::return_value_policy::move, "Scatter 2D surfaces");
mpi_mod.def(
"size", []() { return mpi::size(); },
"Returns the number of MPI processes");
mpi_mod.def(
"rank", []() { return mpi::rank(); },
"Returns the rank of the local process");
}
} // namespace wrap
/* -------------------------------------------------------------------------- */
} // namespace tamaas
c4science · Help