AnyFFT.H

Go to the documentation of this file.

/* Copyright 2019-2020
 *
 * This file is part of WarpX.
 *
 * License: BSD-3-Clause-LBNL
 */
 
#ifndef ANYFFT_H_
#define ANYFFT_H_
 
#include <AMReX_Config.H>
#include <AMReX_LayoutData.H>
 
#if defined(AMREX_USE_CUDA)
#  include <cufft.h>
#elif defined(AMREX_USE_HIP)
#  if __has_include(<rocfft/rocfft.h>)  // ROCm 5.3+
#    include <rocfft/rocfft.h>
#  else
#    include <rocfft.h>
#  endif
#else
#  include <fftw3.h>
#endif
 
namespace AnyFFT
{
    // First, define library-dependent types (complex, FFT plan)
 
#if defined(AMREX_USE_CUDA)
#  ifdef AMREX_USE_FLOAT
    using Complex = cuComplex;
#  else
    using Complex = cuDoubleComplex;
#  endif
#elif defined(AMREX_USE_HIP)
#  ifdef AMREX_USE_FLOAT
    using Complex = float2;
#  else
    using Complex = double2;
#  endif
#else
#  ifdef AMREX_USE_FLOAT
    using Complex = fftwf_complex;
#  else
    using Complex = fftw_complex;
#  endif
#endif
 
#if defined(AMREX_USE_CUDA)
    using VendorFFTPlan = cufftHandle;
#elif defined(AMREX_USE_HIP)
    using VendorFFTPlan = rocfft_plan;
#else
#  ifdef AMREX_USE_FLOAT
    using VendorFFTPlan = fftwf_plan;
#  else
    using VendorFFTPlan = fftw_plan;
#  endif
#endif
 
    // Second, define library-independent API
 
    enum struct direction {R2C, C2R};
 
    struct FFTplan
    {
        amrex::Real* m_real_array; 
        Complex* m_complex_array; 
        VendorFFTPlan m_plan; 
        direction m_dir;  
        int m_dim; 
    };
 
    using FFTplans = amrex::LayoutData<FFTplan>;
 
    FFTplan CreatePlan(const amrex::IntVect& real_size, amrex::Real* real_array,
                       Complex* complex_array, direction dir, int dim);
 
    void DestroyPlan(FFTplan& fft_plan);
 
    void Execute(FFTplan& fft_plan);
}
 
#endif // ANYFFT_H_