File fftlibraryselector.h
File List > code_source > templat > include > TempLat > fft > fftlibraryselector.h
Go to the documentation of this file
#ifndef TEMPLAT_FFT_FFTLIBRARYSELECTOR_H
#define TEMPLAT_FFT_FFTLIBRARYSELECTOR_H
/* This file is part of TempLat, available at https://cosmolattice.github.io/templat .
Copyright 2021-2026 The TempLat authors, see AUTHORS.md.
Released under the MIT license, see LICENSE.md. */
// File info: Main contributor(s): Wessel Valkenburg, Franz R. Sattler, Year: 2025
#include "TempLat/fft/external/fftw/fftwinterface.h"
#include "TempLat/fft/fftdecomposition.h"
#include "TempLat/fft/ffttopology.h"
#ifndef NOFFT
#ifdef HAVE_MPI
#ifdef HAVE_PARAFAFT
#include "TempLat/fft/external/parafaft/parafaftinterface.h"
#endif
#endif
#ifdef HAVE_KOKKOSFFT
#include "TempLat/fft/external/kokkosfft/kokkosfftinterface.h"
#endif
#endif
namespace TempLat
{
MakeException(FFTLibraryDoubleInitializationException);
MakeException(FFTLibraryDecompositionMismatchException);
struct holdStaticGuard {
static bool getSessionGuardsWasCalledOnce()
{
static bool wasOnce = false;
bool result = wasOnce;
wasOnce = true;
return result;
}
};
static inline std::vector<std::shared_ptr<FFTSessionGuard>> getFFTSessionGuards(bool pVerbose = true)
{
if (holdStaticGuard::getSessionGuardsWasCalledOnce())
throw FFTLibraryDoubleInitializationException("You can only call getSessionGuards once.");
std::vector<std::shared_ptr<FFTSessionGuard>> result;
/* add your guards here. */
/* Note: the standard guarantees that the destructors are called in the
inverse order of appearance in the vector. So if you must construct after FFTW and
destruct before FFTW, you should be safe if you add your thing after FFTW.
*/
#ifndef NOFFT
result.push_back(getFFTWSessionGuard(pVerbose));
#ifdef HAVE_MPI
#endif
#endif
return result;
}
enum class FFTBackendTag { Parafaft, KokkosFFT, FFTW };
template <size_t NDim> class FFTLibrarySelector
{
public:
// Put public methods here. These should change very little over time.
FFTLibrarySelector(MPICartesianGroup group, const device::IdxArray<NDim> &nGridPoints,
bool forbidTransposition = false)
: mGroup(group), mNGridPoints(nGridPoints), theLibrary(nullptr), mLayout(mNGridPoints), madePlansFloat(false),
madePlansDouble(false), verbose(false)
{
std::tie(theLibrary, backend) = selectBackend(group.size());
if (mGroup.getRank() == 0) sayShort << "Using " << backend << " backend for FFTs.\n";
verifyDecompositionMatchesBackend(group, nGridPoints);
mLayout = theLibrary->computeLocalSizes(mGroup, mNGridPoints, forbidTransposition);
}
const auto &getLayout() { return mLayout; }
const std::string &getBackend() const { return backend; }
void setVerbose() { verbose = true; }
static FFTTopology<NDim> topology(MPICommReference baseComm, const device::IdxArray<NDim> &nGridPoints)
{
switch (selectBackendTag(baseComm.size())) {
#ifdef HAVE_MPI
#ifdef HAVE_PARAFAFT
case FFTBackendTag::Parafaft:
if constexpr (NDim >= 2)
return ParafaftInterface<NDim>::topology(baseComm, nGridPoints);
else
break;
#endif
#endif
#ifdef HAVE_KOKKOSFFT
case FFTBackendTag::KokkosFFT:
if constexpr (NDim <= 3)
return KokkosFFTInterface<NDim>::topology(baseComm, nGridPoints);
else
break;
#endif
default: break;
}
return FFTWInterface<NDim>::topology(baseComm, nGridPoints);
}
static FFTDecomposition<NDim> decomposition(MPICommReference baseComm, const device::IdxArray<NDim> &nGridPoints)
{
switch (selectBackendTag(baseComm.size())) {
#ifdef HAVE_MPI
#ifdef HAVE_PARAFAFT
case FFTBackendTag::Parafaft:
if constexpr (NDim >= 2)
return ParafaftInterface<NDim>::decomposition(baseComm, nGridPoints);
else
break;
#endif
#endif
#ifdef HAVE_KOKKOSFFT
case FFTBackendTag::KokkosFFT:
if constexpr (NDim <= 3)
return KokkosFFTInterface<NDim>::decomposition(baseComm, nGridPoints);
else
break;
#endif
default: break;
}
return FFTWInterface<NDim>::decomposition(baseComm, nGridPoints);
}
void r2c(MemoryBlock<double, NDim> &mBlock)
{
getPlans_double();
if (verbose) sayMPI << "Going to perform double r2c.\n";
mPlansDouble->r2c(mBlock);
mBlock.flagHostMirrorOutdated();
}
void r2c(MemoryBlock<float, NDim> &mBlock)
{
getPlans_float();
if (verbose) sayMPI << "Going to perform float r2c.\n";
mPlansFloat->r2c(mBlock);
mBlock.flagHostMirrorOutdated();
}
void c2r(MemoryBlock<double, NDim> &mBlock)
{
getPlans_double();
if (verbose) sayMPI << "Going to perform double c2r.\n";
mPlansDouble->c2r(mBlock);
mBlock.flagHostMirrorOutdated();
}
void c2r(MemoryBlock<float, NDim> &mBlock)
{
getPlans_float();
if (verbose) sayMPI << "Going to perform float c2r.\n";
mPlansFloat->c2r(mBlock);
mBlock.flagHostMirrorOutdated();
}
void getPlans_float()
{
if (!madePlansFloat) {
if (verbose) sayMPI << "Going to prepare float FFT plans.\n";
madePlansFloat = true;
mPlansFloat = theLibrary->getPlans_float(mGroup, mLayout);
}
}
void getPlans_double()
{
if (!madePlansDouble) {
if (verbose) sayMPI << "Going to prepare double FFT plans.\n";
madePlansDouble = true;
mPlansDouble = theLibrary->getPlans_double(mGroup, mLayout);
}
}
private:
static FFTBackendTag selectBackendTag([[maybe_unused]] device::Idx nProcesses)
{
#ifdef HAVE_MPI
#ifdef HAVE_PARAFAFT
if constexpr (NDim >= 2) {
if (nProcesses > 1) return FFTBackendTag::Parafaft;
}
#endif
#endif
#ifdef HAVE_KOKKOSFFT
if (nProcesses == 1) {
if constexpr (NDim <= 3) return FFTBackendTag::KokkosFFT;
}
#endif
return FFTBackendTag::FFTW;
}
void verifyDecompositionMatchesBackend(const MPICartesianGroup &group,
const device::IdxArray<NDim> &nGridPoints) const
{
const FFTDecomposition<NDim> expected = decomposition(group.getBaseComm(), nGridPoints);
const std::vector<int> &actual = group.getDecomposition();
if (actual.size() != NDim)
throw FFTLibraryDecompositionMismatchException("Cartesian group decomposition has ", actual.size(),
" entries, expected ", NDim, ".");
bool explicitGrid = expected.nDimsToSplit > 0;
for (device::Idx i = 0; i < expected.nDimsToSplit; ++i)
if (expected.dims[i] <= 0) {
explicitGrid = false;
break;
}
if (explicitGrid) {
for (size_t i = 0; i < NDim; ++i) {
if (actual[i] != expected.dims[i])
throw FFTLibraryDecompositionMismatchException(
"FFT backend ", backend, " requires MPI decomposition {", expectedDimsToString(expected),
"} but the provided MPICartesianGroup has {", actualDimsToString(actual),
"}. Use FFTMPIDomainSplit::makeMPIGroup to build a matching group.");
}
} else {
device::Idx actualSplits = 0;
for (size_t i = 0; i < NDim; ++i)
if (actual[i] > 1) ++actualSplits;
if (actualSplits > expected.nDimsToSplit)
throw FFTLibraryDecompositionMismatchException(
"FFT backend ", backend, " supports splitting at most ", expected.nDimsToSplit,
" dimension(s) but the provided MPICartesianGroup splits ", actualSplits,
". Use FFTMPIDomainSplit::makeMPIGroup to build a matching group.");
}
}
static std::string expectedDimsToString(const FFTDecomposition<NDim> &d)
{
std::string s;
for (size_t i = 0; i < NDim; ++i) {
if (i) s += ",";
s += std::to_string(d.dims[i]);
}
return s;
}
static std::string actualDimsToString(const std::vector<int> &v)
{
std::string s;
for (size_t i = 0; i < v.size(); ++i) {
if (i) s += ",";
s += std::to_string(v[i]);
}
return s;
}
static std::pair<std::shared_ptr<FFTLibraryInterface<NDim>>, std::string> selectBackend(device::Idx nProcesses)
{
switch (selectBackendTag(nProcesses)) {
#ifdef HAVE_MPI
#ifdef HAVE_PARAFAFT
case FFTBackendTag::Parafaft:
if constexpr (NDim >= 2)
return {std::make_shared<ParafaftInterface<NDim>>(), "Parafaft"};
else
break;
#endif
#endif
#ifdef HAVE_KOKKOSFFT
case FFTBackendTag::KokkosFFT:
if constexpr (NDim <= 3)
return {std::make_shared<KokkosFFTInterface<NDim>>(), "KokkosFFT"};
else
break;
#endif
default: break;
}
return {std::make_shared<FFTWInterface<NDim>>(), "FFTW"};
}
/* Put all member variables and private methods here. These may change arbitrarily. */
MPICartesianGroup mGroup;
device::IdxArray<NDim> mNGridPoints;
std::shared_ptr<FFTLibraryInterface<NDim>> theLibrary;
FFTLayoutStruct<NDim> mLayout;
bool madePlansFloat;
bool madePlansDouble;
std::shared_ptr<FFTPlanInterface<float, NDim>> mPlansFloat;
std::shared_ptr<FFTPlanInterface<double, NDim>> mPlansDouble;
bool verbose;
std::string backend;
};
} // namespace TempLat
#endif