Homec4science
Diffusion tamaas (features/netcdf-mpi) Browse

tamaas/src/core/partitioner.hh
7d1b9534fac0features/netcdf-mpi

Last Change
Raw File

partitioner.hh
View Options

/*
* SPDX-License-Indentifier: AGPL-3.0-or-later
*
* Copyright (©) 2016-2023 EPFL (École Polytechnique Fédérale de Lausanne),
* Laboratory (LSMS - Laboratoire de Simulation en Mécanique des Solides)
* Copyright (©) 2020-2023 Lucas Frérot
*
* 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/>.
*
*/
/* -------------------------------------------------------------------------- */
#ifndef PARTITIONER_HH
#define PARTITIONER_HH
/* -------------------------------------------------------------------------- */
#include "fftw/interface.hh"
#include "grid.hh"
#include "mpi_interface.hh"
#include "tamaas.hh"
#include <algorithm>
#include <array>
#include <cstdlib>
#include <functional>
/* -------------------------------------------------------------------------- */
namespace tamaas {
template <UInt dim>
struct Partitioner {
template <typename Container>
static decltype(auto) global_size(Container local) {
local.front() = mpi::allreduce<operation::plus>(local.front());
return local;
}
template <typename T>
static decltype(auto) global_size(const Grid<T, dim>& grid) {
return global_size(grid.sizes());
}
template <typename Container>
static decltype(auto) local_size(Container global) {
if (dim == 1)
return global;
auto tup = fftw::mpi::local_size_many(dim, cast_size(global).data(), 1);
global.front() = static_cast<UInt>(std::get<1>(tup));
return global;
}
template <typename T>
static decltype(auto) local_size(const Grid<T, dim>& grid) {
return local_size(grid.sizes());
}
static decltype(auto) local_size(std::initializer_list<UInt> list) {
std::array<UInt, dim> global;
std::copy_n(list.begin(), dim, global.begin());
return local_size(global);
}
template <typename Container>
static decltype(auto) local_offset(const Container& global) {
if (dim == 1)
return std::ptrdiff_t{0};
auto tup = fftw::mpi::local_size_many(dim, cast_size(global).data(), 1);
return std::ptrdiff_t{std::get<2>(tup)};
}
template <typename T>
static decltype(auto) local_offset(const Grid<T, dim>& grid) {
auto offset = local_offset(global_size(grid.sizes()));
return offset * grid.getStrides().front();
}
static decltype(auto) local_offset(std::initializer_list<UInt> list) {
std::array<UInt, dim> global;
std::copy_n(list.begin(), dim, global.begin());
return local_offset(global);
}
template <typename Container>
static decltype(auto) cast_size(const Container& s) {
std::array<std::ptrdiff_t, dim> n;
std::copy_n(s.cbegin(), dim, n.begin());
return n;
}
static decltype(auto) alloc_size(const std::array<UInt, dim>& global,
UInt howmany) {
if (dim == 1)
return std::accumulate(global.begin(), global.end(), std::size_t{1},
std::multiplies<void>());
auto tup =
fftw::mpi::local_size_many(dim, cast_size(global).data(), howmany);
return static_cast<std::size_t>(std::get<0>(tup));
}
template <typename T>
static Grid<T, dim> gather(const Grid<T, dim>& send) {
Grid<T, dim> result(global_size(send), send.getNbComponents());
mpi::gather(send.getInternalData(), result.getInternalData(),
send.dataSize());
return result;
}
template <typename T>
static Grid<T, dim> scatter(const Grid<T, dim>& send) {
auto global = send.sizes();
auto components = send.getNbComponents();
auto n = send.getStrides().front();
mpi::bcast(global.data(), global.size());
mpi::bcast(&components, 1);
mpi::bcast(&n, 1);
auto tup =
fftw::mpi::local_size_many(dim, cast_size(global).data(), components);
std::vector<int> disp(mpi::size()), sendcounts(mpi::size());
sendcounts[0] = std::get<1>(tup) * n;
disp[0] = std::get<2>(tup) * n;
mpi::gather(sendcounts.data(), sendcounts.data(), 1);
mpi::gather(disp.data(), disp.data(), 1);
Grid<T, dim> result(local_size(global), components);
mpi::scatterv(send.getInternalData(), sendcounts, disp,
result.getInternalData(), result.dataSize());
return result;
}
};
} // namespace tamaas
#endif
c4science · Help