Homec4science
Diffusion tamaas (feature/mpi_fft_engine) Browse

tamaas/src/mpi/fftw_mpi_engine.cpp
0670ef204b92feature/mpi_fft_engine

Last Change
Raw File

fftw_mpi_engine.cpp
View Options

/**
* @file
* @section LICENSE
*
* Copyright (©) 2016-2020 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 "fftw_mpi_engine.hh"
#include "fftw_mpi_interface.hh"
#include "logger.hh"
#include "mpi_interface.hh"
#include "partitioner.hh"
#include <algorithm>
/* -------------------------------------------------------------------------- */
namespace tamaas {
auto FFTWMPIEngine::local_size(const key_t& key) {
const auto n = [&key]() {
std::array<UInt, 2> sizes;
std::copy_n(key.cbegin(), 2, sizes.begin());
sizes[1] = sizes[1] / 2 + 1; // hermitian size
return sizes;
}();
const auto sizes = Partitioner<2>::local_size(n);
return std::make_pair(Partitioner<2>::alloc_size(n, key[2]),
std::get<0>(sizes));
}
void FFTWMPIEngine::forward(Grid<Real, 2>& real,
GridHermitian<Real, 2>& spectral) {
const auto key = make_key(real, spectral);
auto& plans = getPlans(key);
auto& workspace = workspaces[key];
// Ensure enough space for MPI overhead
const auto sizes = local_size(key);
spectral.reserve(sizes.first);
// Can't use iterators here because of size mistmatch in workspace
#pragma omp parallel for collapse(3)
for (UInt i = 0; i < real.sizes()[0]; ++i)
for (UInt j = 0; j < real.sizes()[1]; ++j)
for (UInt k = 0; k < real.getNbComponents(); ++k)
workspace(i, j, k) = real(i, j, k);
fftw_mpi_execute_dft_r2c(plans.first, workspace.getInternalData(),
cast(spectral.getInternalData()));
}
void FFTWMPIEngine::backward(Grid<Real, 2>& real,
GridHermitian<Real, 2>& spectral) {
auto key = make_key(real, spectral);
auto& plans = getPlans(key);
auto& workspace = workspaces[key];
// Ensure enough space for MPI overhead
const auto sizes = local_size(key);
spectral.reserve(sizes.first);
fftw_mpi_execute_dft_c2r(plans.second, cast(spectral.getInternalData()),
workspace.getInternalData());
// Can't use iterators here because of size mistmatch in worspace
#pragma omp parallel for collapse(3)
for (UInt i = 0; i < real.sizes()[0]; ++i)
for (UInt j = 0; j < real.sizes()[1]; ++j)
for (UInt k = 0; k < real.getNbComponents(); ++k)
real(i, j, k) = workspace(i, j, k);
real *= (1. / real.getGlobalNbPoints());
}
FFTWMPIEngine::key_t FFTWMPIEngine::make_key(Grid<Real, 2>& real,
GridHermitian<Real, 2>& spectral) {
auto key{FFTEngine::make_key(real, spectral)};
// Reduce first dimension for total size
key.front() = mpi::allreduce<operation::plus>(key.front());
return key;
}
FFTWMPIEngine::plan_t& FFTWMPIEngine::getPlans(key_t key) {
if (plans.find(key) != plans.end())
return plans[key];
const int rank = key.size() - 3; // dimension of fft
const int howmany = key[rank];
std::vector<std::ptrdiff_t> n(rank); // size of individual fft
std::copy_n(key.begin(), rank, n.begin());
fftw::ptr<Real> in{nullptr};
fftw::ptr<complex_t> out{nullptr};
const auto sizes = local_size(key);
workspaces.emplace(key, Grid<Real, 2>({static_cast<UInt>(sizes.second),
2 * (key[1] / 2 + 1)},
howmany));
auto& workspace = workspaces[key];
workspace.reserve(2 * sizes.first);
out._ptr = fftw::helper<Real>::alloc_complex(sizes.first);
fftw::plan<Real> forward{fftw_mpi_plan_many_dft_r2c(
rank, n.data(), howmany, FFTW_MPI_DEFAULT_BLOCK, FFTW_MPI_DEFAULT_BLOCK,
workspace.getInternalData(), out, mpi::comm::world(), flags())};
fftw::plan<Real> backward{fftw_mpi_plan_many_dft_c2r(
rank, n.data(), howmany, FFTW_MPI_DEFAULT_BLOCK, FFTW_MPI_DEFAULT_BLOCK,
out, workspace.getInternalData(), mpi::comm::world(), flags())};
plans[key] = std::make_pair(std::move(forward), std::move(backward));
return plans[key];
}
} // namespace tamaas
c4science · Help