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)
// cstddef: work-around for ROCm/rocFFT <=4.3.0
// https://github.com/ROCmSoftwarePlatform/rocFFT/blob/rocm-4.3.0/library/include/rocfft.h#L36-L42
#  include <cstddef>
#  include <rocfft.h>
#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 * const real_array,
                       Complex * const complex_array, const direction dir, const int dim);

    void DestroyPlan(FFTplan& fft_plan);

    void Execute(FFTplan& fft_plan);
}

#endif // ANYFFT_H_