diff --git a/src/core/fft_engine.cpp b/src/core/fft_engine.cpp index 33ce190..81a7f20 100644 --- a/src/core/fft_engine.cpp +++ b/src/core/fft_engine.cpp @@ -1,59 +1,96 @@ /* * 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 . * */ /* -------------------------------------------------------------------------- */ #include "fft_engine.hh" #if TAMAAS_USE_CUDA #include "cuda/cufft_engine.hh" #else #include "fftw/fftw_engine.hh" #ifdef TAMAAS_USE_MPI #include "fftw/mpi/fftw_mpi_engine.hh" #include "mpi_interface.hh" #endif #endif /* -------------------------------------------------------------------------- */ namespace tamaas { std::unique_ptr FFTEngine::makeEngine(unsigned int flags) { #define inst(x) \ do { \ Logger().get(LogLevel::debug) << TAMAAS_DEBUG_MSG("[" #x "] Init"); \ return std::make_unique(flags); \ } while (0) #ifdef TAMAAS_USE_MPI if (mpi::size() != 1) inst(FFTWMPIEngine); else inst(FFTWEngine); #elif TAMAAS_USE_CUDA inst(CuFFTEngine); #else inst(FFTWEngine); #endif #undef inst } +template <> +std::vector> +FFTEngine::realCoefficients<2>(const std::array& local_sizes) { + constexpr UInt dim = 2; + std::vector> indices; + indices.reserve(dim * dim); + + auto n = Partitioner::global_size(local_sizes); + std::array even{{n[0] % 2 == 0, n[1] % 2 == 0}}; + + indices.push_back({{0, 0}}); + + if (even[0]) + indices.push_back({{n[0] / 2, 0}}); + + if (even[1]) + indices.push_back({{0, n[1] / 2}}); + + if (even[0] and even[1]) + indices.push_back({{n[0] / 2, n[1] / 2}}); + + return indices; +} + +template <> +std::vector> +FFTEngine::realCoefficients<1>(const std::array& local_sizes) { + constexpr UInt dim = 1; + std::vector> indices; + indices.push_back({{0}}); + + if (local_sizes[0] % 2 == 0) + indices.push_back({{local_sizes[0] / 2}}); + + return indices; +} + } // namespace tamaas diff --git a/src/core/fft_engine.hh b/src/core/fft_engine.hh index 0f87fae..6cdece3 100644 --- a/src/core/fft_engine.hh +++ b/src/core/fft_engine.hh @@ -1,136 +1,141 @@ /* * 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 . * */ /* -------------------------------------------------------------------------- */ #ifndef FFT_ENGINE_H #define FFT_ENGINE_H /* -------------------------------------------------------------------------- */ #include "grid.hh" #include "grid_base.hh" #include "grid_hermitian.hh" #include "partitioner.hh" #include #include #include #include #include /* -------------------------------------------------------------------------- */ namespace tamaas { class FFTEngine { protected: using key_t = std::basic_string; public: virtual ~FFTEngine() noexcept = default; /// Execute a forward plan on real and spectral 1D virtual void forward(const Grid& real, GridHermitian& spectral) = 0; /// Execute a forward plan on real and spectral 2D virtual void forward(const Grid& real, GridHermitian& spectral) = 0; /// Execute a backward plan on real and spectral 1D virtual void backward(Grid& real, GridHermitian& spectral) = 0; /// Execute a backward plan on real and spectral 2D virtual void backward(Grid& real, GridHermitian& spectral) = 0; /// Fill a grid with wavevector values in appropriate ordering template static Grid computeFrequencies(const std::array& sizes); + /// Return the grid indices that contain real coefficients (see R2C transform) + template + static std::vector> + realCoefficients(const std::array& sizes); + /// Instanciate an appropriate FFTEngine subclass static std::unique_ptr makeEngine(unsigned int flags = FFTW_ESTIMATE); protected: /// Make a transform signature from a pair of grids template static key_t make_key(const Grid& real, const GridHermitian& spectral); }; template FFTEngine::key_t FFTEngine::make_key(const Grid& real, const GridHermitian& spectral) { if (real.getNbComponents() != spectral.getNbComponents()) TAMAAS_EXCEPTION("Components do not match"); auto hermitian_dims = GridHermitian::hermitianDimensions(real.sizes()); if (not std::equal(hermitian_dims.begin(), hermitian_dims.end(), spectral.sizes().begin())) TAMAAS_EXCEPTION("Spectral grid does not have hermitian size"); // Reserve space for dimensions + components + both last strides key_t key(real.getDimension() + 3, 0); std::copy_n(real.sizes().begin(), dim, key.begin()); key[dim] = real.getNbComponents(); key[dim + 1] = real.getStrides().back(); key[dim + 2] = spectral.getStrides().back(); return key; } template Grid FFTEngine::computeFrequencies(const std::array& local_sizes) { // If hermitian is true, we suppose the dimensions of freq are // reduced based on hermitian symetry and that it has dim components Grid freq(local_sizes, dim); constexpr Int dmax = dim - static_cast(hermitian); // For MPI const auto& n = Partitioner::global_size(local_sizes); const auto offset = Partitioner::local_offset(freq) / dim; #pragma omp parallel for // to get rid of a compilation warning from nvcc for (UInt i = 0; i < freq.getNbPoints(); ++i) { UInt index = i + offset; VectorProxy wavevector(freq(i * freq.getNbComponents())); std::array tuple{{0}}; /// Computing tuple from index for (Int d = dim - 1; d >= 0; d--) { tuple[d] = index % n[d]; index -= tuple[d]; index /= n[d]; } if (hermitian) wavevector(dim - 1) = tuple[dim - 1]; for (Int d = 0; d < dmax; d++) { // Type conversion T td = tuple[d]; T nd = n[d]; T q = (tuple[d] < n[d] / 2) ? td : td - nd; wavevector(d) = q; } } return freq; } } // namespace tamaas #endif diff --git a/tests/test_fftfreq.cpp b/tests/test_fftfreq.cpp index bb7c2fc..642ccb6 100644 --- a/tests/test_fftfreq.cpp +++ b/tests/test_fftfreq.cpp @@ -1,89 +1,117 @@ /* * 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 . * */ /* -------------------------------------------------------------------------- */ #include "fft_engine.hh" #include "grid.hh" #include "grid_hermitian.hh" #include "grid_view.hh" #include "mpi_interface.hh" #include "partitioner.hh" #include "test.hh" #include #include #include using namespace tamaas; namespace py = pybind11; /* -------------------------------------------------------------------------- */ TEST(TestFFTEngine, Frequencies1D) { mpi::sequential_guard guard{}; std::array sizes = {{10}}; auto freq = FFTEngine::computeFrequencies(sizes); py::module fftfreq = py::module::import("fftfreq"); std::vector reference(freq.dataSize()); py::array py_arg(reference.size(), reference.data(), py::none()); fftfreq.attr("frequencies1D")(py_arg); EXPECT_TRUE(compare(reference, freq, AreFloatEqual())) << "Non hermitian frequencies are wrong"; auto hfreq = FFTEngine::computeFrequencies(sizes); std::iota(reference.begin(), reference.end(), 0); EXPECT_TRUE(compare(reference, hfreq, AreFloatEqual())) << "Hermitian frequencies are wrong"; } /* -------------------------------------------------------------------------- */ TEST(TestFFTEngine, Frequencies2D) { std::array sizes = {{10, 10}}; auto global_sizes = Partitioner<2>::global_size(sizes); auto freq = FFTEngine::computeFrequencies(sizes); auto hfreq = FFTEngine::computeFrequencies(sizes); auto gathered = Partitioner<2>::gather(freq); auto hgathered = Partitioner<2>::gather(hfreq); if (mpi::rank() != 0) return; py::module fftfreq = py::module::import("fftfreq"); std::vector reference(gathered.dataSize()); py::array py_arg({global_sizes[0], global_sizes[1], 2u}, reference.data(), py::none()); fftfreq.attr("frequencies2D")(py_arg); EXPECT_TRUE(compare(reference, gathered, AreFloatEqual())) << "Non hermitian frequencies are wrong"; fftfreq.attr("hfrequencies2D")(py_arg); EXPECT_TRUE(compare(reference, hgathered, AreFloatEqual())) << "Hermitian frequencies are wrong"; } + +/* -------------------------------------------------------------------------- */ +TEST(TestFFTEngine, RealComponents) { + mpi::sequential_guard guard; + + auto even_even = FFTEngine::realCoefficients<2>({{10, 10}}); + auto even_odd = FFTEngine::realCoefficients<2>({{10, 11}}); + auto odd_even = FFTEngine::realCoefficients<2>({{11, 10}}); + auto odd_odd = FFTEngine::realCoefficients<2>({{11, 11}}); + + std::vector even_even_ref{{0, 0, 5, 0, 0, 5, 5, 5}}; + std::vector even_odd_ref{{0, 0, 5, 0}}; + std::vector odd_even_ref{{0, 0, 0, 5}}; + std::vector odd_odd_ref{{0, 0}}; + + auto flatten = [](auto v) { + std::vector flat; + for (auto&& tuple : v) + for (auto i : tuple) + flat.push_back(i); + return flat; + }; + + compare(flatten(even_even), even_even_ref); + compare(flatten(even_odd), even_odd_ref); + compare(flatten(odd_even), odd_even_ref); + compare(flatten(odd_odd), odd_odd_ref); +}