File kokkosfftplanholder.h

File List > code_source > templat > include > TempLat > fft > external > kokkosfft > kokkosfftplanholder.h

Go to the documentation of this file

#ifndef TEMPLAT_FFT_EXTERNAL_KOKKOSFFT_KOKKOSFFTPLANHOLDER_H
#define TEMPLAT_FFT_EXTERNAL_KOKKOSFFT_KOKKOSFFTPLANHOLDER_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): Franz R. Sattler, Year: 2025

#ifndef NOFFT
#ifdef HAVE_KOKKOSFFT
#include <KokkosFFT.hpp>
#endif
#endif

#include "TempLat/fft/fftlibraryinterface.h"
#include "TempLat/parallel/mpi/cartesian/mpicartesiangroup.h"

#include "TempLat/parallel/device_iteration.h"

namespace TempLat
{
  template <typename T, size_t NDim> class KokkosFFTPlanHolder : public FFTPlanInterface<T, NDim>
  {
  public:
    // Put public methods here. These should change very little over time.

    using PlanType_c2r =
        typename KokkosFFT::Plan<Kokkos::DefaultExecutionSpace, device::memory::NDViewUnmanaged<complex<T>, NDim>,
                                 device::memory::NDViewUnmanaged<T, NDim>, NDim>;
    using PlanType_r2c =
        typename KokkosFFT::Plan<Kokkos::DefaultExecutionSpace, device::memory::NDViewUnmanaged<T, NDim>,
                                 device::memory::NDViewUnmanaged<complex<T>, NDim>, NDim>;

    struct Plans {
      std::shared_ptr<PlanType_c2r> c2rPlan;
      std::shared_ptr<PlanType_r2c> r2cPlan;

      void execute_c2r(const auto &src, const auto &dest)
      {
        KokkosFFT::execute(*c2rPlan, src, dest, KokkosFFT::Normalization::none);
      }

      void execute_r2c(const auto &src, const auto &dest)
      {
        KokkosFFT::execute(*r2cPlan, src, dest, KokkosFFT::Normalization::none);
      }

      device::array<int, NDim> configSizes;
      device::array<int, NDim> fourierSizes;
    };

    KokkosFFTPlanHolder(MPICartesianGroup group, const Plans &plans) : mGroup(group), mPlans(plans) {}

    virtual ~KokkosFFTPlanHolder() = default;

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

  private:
    /* Put all member variables and private methods here. These may change arbitrarily. */
    MPICartesianGroup mGroup;
    Plans mPlans;

    void execute_r2c(MemoryBlock<T, NDim> &mBlock)
    {
      device::iteration::fence();
      auto fourier_view = device::apply(
          [&](auto &&...args) {
            return device::memory::NDViewUnmanaged<complex<T>, NDim>(reinterpret_cast<complex<T> *>(mBlock.data()),
                                                                     args...);
          },
          mPlans.fourierSizes);
      auto config_view = device::apply(
          [&](auto &&...args) { return device::memory::NDViewUnmanaged<T, NDim>(mBlock.data(), args...); },
          mPlans.configSizes);

      mPlans.execute_r2c(config_view, fourier_view);
      device::iteration::fence();
    }

    void execute_c2r(MemoryBlock<T, NDim> &mBlock)
    {
      device::iteration::fence();
      auto fourier_view = device::apply(
          [&](auto &&...args) {
            return device::memory::NDViewUnmanaged<complex<T>, NDim>(reinterpret_cast<complex<T> *>(mBlock.data()),
                                                                     args...);
          },
          mPlans.fourierSizes);
      auto config_view = device::apply(
          [&](auto &&...args) { return device::memory::NDViewUnmanaged<T, NDim>(mBlock.data(), args...); },
          mPlans.configSizes);

      mPlans.execute_c2r(fourier_view, config_view);
      device::iteration::fence();
    }
  };

} // namespace TempLat

#endif