AnyFFT.H

Go to the documentation of this file.

/* Copyright 2019-2023
 *
 * This file is part of ABLASTR.
 *
 * License: BSD-3-Clause-LBNL
 */
 
#ifndef ABLASTR_ANYFFT_H_
#define ABLASTR_ANYFFT_H_
 
#ifdef ABLASTR_USE_FFT
#   include <AMReX_Config.H>
#   include <AMReX_GpuComplex.H>
#   include <AMReX_LayoutData.H>
 
#   if defined(AMREX_USE_CUDA)
#       include <cufft.h>
#       include <cuComplex.h>
#   elif defined(AMREX_USE_HIP)
#       if __has_include(<rocfft/rocfft.h>)  // ROCm 5.3+
#           include <rocfft/rocfft.h>
#       else
#           include <rocfft.h>
#       endif
#       include <hip/hip_complex.h>
#   elif defined(AMREX_USE_SYCL)
#       include <oneapi/mkl/dft.hpp>
#   else
#       include <fftw3.h>
#   endif
#endif
 

namespace ablastr::math::anyfft
{

    void setup();

    void cleanup();
 
#ifdef ABLASTR_USE_FFT
 
    // 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
#   elif defined(AMREX_USE_SYCL)
        using Complex = amrex::GpuComplex<amrex::Real>;
#   else
#       ifdef AMREX_USE_FLOAT
            using Complex = fftwf_complex;
#       else
            using Complex = fftw_complex;
#       endif
#   endif
 
    /* Library-dependent multiply helpers */
#   if defined(AMREX_USE_CUDA)
#       ifdef AMREX_USE_FLOAT
    AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE void multiply (Complex & c, Complex const & a, Complex const & b) { c = cuCmulf(a, b); }
#       else
    AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE void multiply (Complex & c, Complex const & a, Complex const & b) { c = cuCmul(a, b); }
#       endif
#   elif defined(AMREX_USE_HIP)
    #   ifdef AMREX_USE_FLOAT
    AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE void multiply (Complex & c, Complex const & a, Complex const & b) { c = hipCmulf(a, b); }
#       else
    AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE void multiply (Complex & c, Complex const & a, Complex const & b) { c = hipCmul(a, b); }
#       endif
#   elif defined(AMREX_USE_SYCL)
    AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE void multiply (Complex & c, Complex const & a, Complex const & b) { c = a * b; }
#   else
    AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE void multiply (Complex & c, Complex const & a, Complex const & b) {
        c[0] = a[0] * b[0] - a[1] * b[1];
        c[1] = a[0] * b[1] + a[1] * b[0];
    }
#   endif

#   if defined(AMREX_USE_CUDA)
        using VendorFFTPlan = cufftHandle;
#   elif defined(AMREX_USE_HIP)
        using VendorFFTPlan = rocfft_plan;
#   elif defined(AMREX_USE_SYCL)
        using VendorFFTPlan = oneapi::mkl::dft::descriptor<
#       ifdef AMREX_USE_FLOAT
            oneapi::mkl::dft::precision::SINGLE,
#       else
            oneapi::mkl::dft::precision::DOUBLE,
#       endif
            oneapi::mkl::dft::domain::REAL> *;
        // dft::descriptor has no default ctor, so we have to use ptr.
#   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; 
#ifdef AMREX_USE_SYCL
        amrex::gpuStream_t m_stream;
#endif
    };

    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
 
}
 
#endif // ABLASTR_ANYFFT_H_