fft.hpp

Go to the documentation of this file.

// SPDX-FileCopyrightText: 2025 VTT Technical Research Centre of Finland Ltd
// SPDX-License-Identifier: AGPL-3.0-or-later
 
#pragma once
 
#include "core/decomposition.hpp"
#include "openpfc/backends/heffte_adapter.hpp" // Ensure this is included for the conversion operator
#include "openpfc/core/backend_tags.hpp"
#include "openpfc/core/databuffer.hpp"
#include "openpfc/core/world.hpp"
#include "openpfc/fft/kspace.hpp"
 
#include <heffte.h>
#include <iostream>
#include <memory>
#include <mpi.h>
 
namespace pfc {
namespace fft {
 
namespace layout {
 
using box3di = heffte::box3d<int>;
using Decomposition = pfc::decomposition::Decomposition;
using pfc::types::Int3;
 
struct FFTLayout {
  const Decomposition m_decomposition; 
  const int m_r2c_direction = 0;       
  const std::vector<heffte::box3d<int>> m_real_boxes;    
  const std::vector<heffte::box3d<int>> m_complex_boxes; 
};
 
const FFTLayout create(const Decomposition &decomposition, int r2c_direction);
 
inline const auto &get_real_box(const FFTLayout &layout, int i) {
  return layout.m_real_boxes.at(i);
}
 
inline const auto &get_complex_box(const FFTLayout &layout, int i) {
  return layout.m_complex_boxes.at(i);
}
 
inline auto get_r2c_direction(const FFTLayout &layout) {
  return layout.m_r2c_direction;
}
 
} // namespace layout
 
using pfc::types::Int3;
using pfc::types::Real3;
 
using Decomposition = pfc::decomposition::Decomposition;
using RealVector = std::vector<double>;
using ComplexVector = std::vector<std::complex<double>>;
using box3di = heffte::box3d<int>; 
 
// Backend-aware DataBuffer type aliases
using RealDataBuffer = core::DataBuffer<backend::CpuTag, double>;
using ComplexDataBuffer = core::DataBuffer<backend::CpuTag, std::complex<double>>;
#if defined(OpenPFC_ENABLE_CUDA)
using RealDataBufferCUDA = core::DataBuffer<backend::CudaTag, double>;
using ComplexDataBufferCUDA =
    core::DataBuffer<backend::CudaTag, std::complex<double>>;
#endif
 
enum class Backend {
  FFTW, 
  CUDA  
};
 
// Type aliases for different FFT backends
using fft_r2c = heffte::fft3d_r2c<heffte::backend::fftw>; 
#if defined(OpenPFC_ENABLE_CUDA)
using fft_r2c_cuda =
    heffte::fft3d_r2c<heffte::backend::cufft>; 
#endif
 
struct IFFT {
  virtual ~IFFT() = default;
 
  virtual void forward(const RealVector &in, ComplexVector &out) = 0;
 
  virtual void backward(const ComplexVector &in, RealVector &out) = 0;
 
  virtual void reset_fft_time() = 0;
  virtual double get_fft_time() const = 0;
 
  virtual size_t size_inbox() const = 0;
  virtual size_t size_outbox() const = 0;
  virtual size_t size_workspace() const = 0;
 
  virtual size_t get_allocated_memory_bytes() const = 0;
};
 
template <typename BackendTag = heffte::backend::fftw> struct FFT_Impl : IFFT {
 
  // const Decomposition m_decomposition; /**< The Decomposition object. */
  // const box3di m_inbox, m_outbox;      /**< Local inbox and outbox boxes. */
 
  using fft_type = heffte::fft3d_r2c<BackendTag>;
  const fft_type m_fft;    
  double m_fft_time = 0.0; 
  // Backend-aware workspace - precision determined by data types in forward/backward
  // calls Default to double precision workspace (can be overridden per call)
  using workspace_type = typename heffte::fft3d_r2c<
      BackendTag>::template buffer_container<std::complex<double>>;
  workspace_type
      m_wrk; 
  FFT_Impl(fft_type fft) : m_fft(std::move(fft)), m_wrk(m_fft.size_workspace()) {}
 
  // Forward method using DataBuffer (backend-aware, precision-aware via template)
  template <typename RealBackendTag, typename ComplexBackendTag, typename RealType>
  void forward(const core::DataBuffer<RealBackendTag, RealType> &in,
               core::DataBuffer<ComplexBackendTag, std::complex<RealType>> &out) {
    static_assert(std::is_same_v<RealBackendTag, ComplexBackendTag>,
                  "Input and output must use the same backend");
    m_fft_time -= MPI_Wtime();
    // HeFFTe's forward method is templated on input/output types and handles
    // precision automatically Create workspace with matching precision
    auto wrk = typename heffte::fft3d_r2c<BackendTag>::template buffer_container<
        std::complex<RealType>>(m_fft.size_workspace());
    m_fft.forward(in.data(), out.data(), wrk.data());
    m_fft_time += MPI_Wtime();
  }
 
  // Forward method using std::vector (implements IFFT interface)
  // For CPU backend: works directly with std::vector
  // For GPU backend: throws error (must use DataBuffer overload)
  void forward(const RealVector &in, ComplexVector &out) override {
    if constexpr (std::is_same_v<BackendTag, heffte::backend::fftw>) {
      // CPU backend: call HeFFTe directly (no conversion needed)
      m_fft_time -= MPI_Wtime();
      m_fft.forward(in.data(), out.data(), m_wrk.data());
      m_fft_time += MPI_Wtime();
    } else {
      // GPU backend: must use DataBuffer
      throw std::runtime_error(
          "GPU FFT requires DataBuffer, not std::vector. Use forward(DataBuffer, "
          "DataBuffer) instead.");
    }
  }
 
  // Backward method using DataBuffer (backend-aware, precision-aware via template)
  template <typename ComplexBackendTag, typename RealBackendTag, typename RealType>
  void
  backward(const core::DataBuffer<ComplexBackendTag, std::complex<RealType>> &in,
           core::DataBuffer<RealBackendTag, RealType> &out) {
    static_assert(std::is_same_v<ComplexBackendTag, RealBackendTag>,
                  "Input and output must use the same backend");
    m_fft_time -= MPI_Wtime();
    // HeFFTe's backward method is templated on input/output types and handles
    // precision automatically Create workspace with matching precision
    auto wrk = typename heffte::fft3d_r2c<BackendTag>::template buffer_container<
        std::complex<RealType>>(m_fft.size_workspace());
    m_fft.backward(in.data(), out.data(), wrk.data(), heffte::scale::full);
    m_fft_time += MPI_Wtime();
  }
 
  // Backward method using std::vector (implements IFFT interface)
  // For CPU backend: works directly with std::vector
  // For GPU backend: throws error (must use DataBuffer overload)
  void backward(const ComplexVector &in, RealVector &out) override {
    if constexpr (std::is_same_v<BackendTag, heffte::backend::fftw>) {
      // CPU backend: call HeFFTe directly (no conversion needed)
      m_fft_time -= MPI_Wtime();
      m_fft.backward(in.data(), out.data(), m_wrk.data(), heffte::scale::full);
      m_fft_time += MPI_Wtime();
    } else {
      // GPU backend: must use DataBuffer
      throw std::runtime_error(
          "GPU FFT requires DataBuffer, not std::vector. Use backward(DataBuffer, "
          "DataBuffer) instead.");
    }
  }
 
  void reset_fft_time() override { m_fft_time = 0.0; }
 
  double get_fft_time() const override { return m_fft_time; }
 
  // const Decomposition &get_decomposition() { return m_decomposition; }
 
  size_t size_inbox() const override { return m_fft.size_inbox(); }
 
  size_t size_outbox() const override { return m_fft.size_outbox(); }
 
  size_t size_workspace() const override { return m_fft.size_workspace(); }
 
  size_t get_allocated_memory_bytes() const override {
    return m_wrk.size() * sizeof(typename workspace_type::value_type);
  }
};
 
// Type aliases for backward compatibility (defaults to FFTW backend)
// Precision is handled by data types, not template parameters
using FFT = FFT_Impl<heffte::backend::fftw>;
 
// Helper functions
template <typename BackendTag>
inline const auto &get_fft_object(const FFT_Impl<BackendTag> &fft) noexcept {
  return fft.m_fft;
}
 
template <typename BackendTag>
inline auto get_inbox(const FFT_Impl<BackendTag> &fft) noexcept {
  return get_fft_object(fft).inbox();
}
 
template <typename BackendTag>
inline auto get_outbox(const FFT_Impl<BackendTag> &fft) noexcept {
  return get_fft_object(fft).outbox();
}
 
using heffte::plan_options;
using layout::FFTLayout;
 
FFT create(const FFTLayout &fft_layout, int rank_id, plan_options options);
 
FFT create(const Decomposition &decomposition, int rank_id);
 
FFT create(const Decomposition &decomposition);
 
std::unique_ptr<IFFT> create_with_backend(const FFTLayout &fft_layout, int rank_id,
                                          plan_options options, Backend backend);
 
std::unique_ptr<IFFT> create_with_backend(const Decomposition &decomposition,
                                          int rank_id, Backend backend);
 
} // namespace fft
 
using FFT = fft::FFT;                     
using FFTLayout = fft::layout::FFTLayout; 
 
} // namespace pfc