File parafaftplanholder.h

File List > code_source > templat > include > TempLat > fft > external > parafaft > parafaftplanholder.h

Go to the documentation of this file

#ifndef TEMPLAT_FFT_EXTERNAL_PARAFAFT_PARAFAFTPLANHOLDER_H
#define TEMPLAT_FFT_EXTERNAL_PARAFAFT_PARAFAFTPLANHOLDER_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): Adrien Florio, Year: 2026

#ifndef NOFFT
#ifdef HAVE_MPI
#ifdef HAVE_PARAFAFT
#include <parafaft_r2c.hpp>
#endif
#endif
#endif

#include "TempLat/util/exception.h"
#include "TempLat/fft/fftlibraryinterface.h"
#include "TempLat/parallel/mpi/cartesian/mpicartesiangroup.h"
#include "TempLat/lattice/memory/memoryblock.h"

#include <memory>
#include "TempLat/parallel/device.h"
#include "TempLat/parallel/device_iteration.h"

namespace TempLat
{
  // Depending on the device, we alias ParaFaFT_Backend to the appropriate backend type. Parafaft supports multiple
  // backends (FFTW, cuFFT, HIPFFT), and we select the one based on compile-time definitions set by CMake when detecting
  // the device and available libraries. This allows us to write device-agnostic code in the ParafaftPlanHolder, while
  // still leveraging the performance benefits of the appropriate backend for the target platform. Each backend is
  // itself a class template parameterized on the floating-point precision, so the alias is a template alias.
#ifdef DEVICE_CUDA
  template <typename T> using ParaFaFT_Backend = parafaft::CuFFTBackend<T>;
#elif defined(DEVICE_HIP)
  template <typename T> using ParaFaFT_Backend = parafaft::HipFFTBackend<T>;
#else
  template <typename T> using ParaFaFT_Backend = parafaft::FFTWBackend<T>;
#endif

  MakeException(ParafaftPlanHolderException);
  MakeException(ParafaftCompiledWithoutSinglePrecisionSupport);

  template <typename T, size_t NDim> class ParafaftPlanHolder : public FFTPlanInterface<T, NDim>
  {
  public:
    using Complex = complex<T>;

#ifdef HAVE_MPI
#ifdef HAVE_PARAFAFT
    ParafaftPlanHolder(MPICartesianGroup group,
                       std::shared_ptr<parafaft::ParaFaFT_R2C<NDim, ParaFaFT_Backend<T>>> parafaftObj)
        : mGroup(group), mParafaft(parafaftObj)
    {
    }
#endif
#endif

    virtual ~ParafaftPlanHolder()
    {
      // Parafaft cleans up internally via its destructor
    }

    virtual void c2r(MemoryBlock<T, NDim> &mBlock) override { execute_c2r(mBlock); }

    virtual void r2c(MemoryBlock<T, NDim> &mBlock) override { execute_r2c(mBlock); }

  private:
#ifdef HAVE_MPI
#ifdef HAVE_PARAFAFT
    // Keep group alive for MPI communicator lifetime
    MPICartesianGroup mGroup;

    // Shared pointer to parafaft object
    std::shared_ptr<parafaft::ParaFaFT_R2C<NDim, ParaFaFT_Backend<T>>> mParafaft;

#endif
#endif

    void execute_r2c(MemoryBlock<T, NDim> &mBlock)
    {
#ifdef HAVE_MPI
#ifdef HAVE_PARAFAFT
      // ParaFaFT reads/writes mBlock.data() through its own (cuFFT/transpose)
      // streams, which are not ordered against the Kokkos execution space that
      // produced the buffer. Fence before so any pending device writes to the
      // buffer are complete before ParaFaFT reads it, and after so ParaFaFT's
      // writes are complete before downstream Kokkos work reads them. Without
      // this, the async fill and the FFT race (host-serial backend: no-op).
      // Mirrors KokkosFFTPlanHolder.
      device::iteration::fence();
      // Parafaft's forward_in_place accepts padded buffers directly
      // Buffer layout: [N0_local][N1_local][2*(N/2+1)] - matches CosmoLattice
      mParafaft->forward_in_place(mBlock.data());
      device::iteration::fence();
#endif
#endif
    }

    void execute_c2r(MemoryBlock<T, NDim> &mBlock)
    {
#ifdef HAVE_MPI
#ifdef HAVE_PARAFAFT
      // See execute_r2c: fence around the ParaFaFT call to order it against the
      // Kokkos execution space (no-op on the host-serial backend).
      device::iteration::fence();
      // Parafaft's backward_in_place accepts padded buffers directly
      // Buffer layout: [N0_local][N1_local][2*(N/2+1)] - matches CosmoLattice
      mParafaft->backward_in_place(mBlock.data());
      device::iteration::fence();
#endif
#endif
    }
  };
} // namespace TempLat

#endif