heffte_backend_fftw.h

/*
    -- heFFTe --
       Univ. of Tennessee, Knoxville
       @date
*/
 
#ifndef HEFFTE_BACKEND_FFTW_H
#define HEFFTE_BACKEND_FFTW_H
 
#include "heffte_r2r_executor.h"
 
#ifdef Heffte_ENABLE_FFTW
 
#include "fftw3.h"
 
namespace heffte{
 
namespace backend{
    template<> struct is_enabled<fftw> : std::true_type{};
 
    template<> struct is_enabled<fftw_cos> : std::true_type{};
    template<> struct is_enabled<fftw_sin> : std::true_type{};
 
// Specialization is not necessary since the default behavior assumes CPU parameters.
//     template<>
//     struct buffer_traits<fftw>{
//         using location = tag::cpu;
//         template<typename T> using container = std::vector<T>;
//     };
#ifndef Heffte_ENABLE_MKL
    template<> struct default_backend<tag::cpu> {
        using type = fftw;
    };
#endif
}
 
template<> struct is_ccomplex<fftwf_complex> : std::true_type{};
template<> struct is_zcomplex<fftw_complex> : std::true_type{};
 
template<typename, direction> struct plan_fftw{};
 
template<direction dir>
struct plan_fftw<std::complex<float>, dir>{
    plan_fftw(int size, int howmanyffts, int stride, int dist) :
        plan(fftwf_plan_many_dft(1, &size, howmanyffts, nullptr, nullptr, stride, dist,
                                                    nullptr, nullptr, stride, dist,
                                                    (dir == direction::forward) ? FFTW_FORWARD : FFTW_BACKWARD, FFTW_ESTIMATE
                                ))
        {}
    plan_fftw(int size1, int size2, std::array<int, 2> const &embed, int howmanyffts, int stride, int dist){
        std::array<int, 2> size = {size2, size1};
 
        if (embed[0] == 0 and embed[1] == 0){
            plan = fftwf_plan_many_dft(2, size.data(), howmanyffts, nullptr, nullptr, stride, dist,
                                                    nullptr, nullptr, stride, dist,
                                                    (dir == direction::forward) ? FFTW_FORWARD : FFTW_BACKWARD, FFTW_ESTIMATE
                                      );
        }else{
            plan = fftwf_plan_many_dft(2, size.data(), howmanyffts, nullptr, embed.data(), stride, dist,
                                                    nullptr, embed.data(), stride, dist,
                                                    (dir == direction::forward) ? FFTW_FORWARD : FFTW_BACKWARD, FFTW_ESTIMATE
                                      );
        }
    }
    plan_fftw(int size1, int size2, int size3){
        std::array<int, 3> size = {size3, size2, size1};
        plan = fftwf_plan_many_dft(3, size.data(), 1, nullptr, nullptr, 1, 1, nullptr, nullptr, 1, 1,
                                   (dir == direction::forward) ? FFTW_FORWARD : FFTW_BACKWARD, FFTW_ESTIMATE);
    }
    ~plan_fftw(){ fftwf_destroy_plan(plan); }
    operator fftwf_plan() const{ return plan; }
    fftwf_plan plan;
};
template<direction dir>
struct plan_fftw<std::complex<double>, dir>{
    plan_fftw(int size, int howmanyffts, int stride, int dist) :
        plan(fftw_plan_many_dft(1, &size, howmanyffts, nullptr, nullptr, stride, dist,
                                                   nullptr, nullptr, stride, dist,
                                                   (dir == direction::forward) ? FFTW_FORWARD : FFTW_BACKWARD, FFTW_ESTIMATE
                               ))
        {}
    plan_fftw(int size1, int size2, std::array<int, 2> const &embed, int howmanyffts, int stride, int dist){
        std::array<int, 2> size = {size2, size1};
 
        if (embed[0] == 0 and embed[1] == 0){
            plan = fftw_plan_many_dft(2, size.data(), howmanyffts, nullptr, nullptr, stride, dist,
                                                      nullptr, nullptr, stride, dist,
                                                      (dir == direction::forward) ? FFTW_FORWARD : FFTW_BACKWARD, FFTW_ESTIMATE
                                     );
        }else{
            plan = fftw_plan_many_dft(2, size.data(), howmanyffts, nullptr, embed.data(), stride, dist,
                                                      nullptr, embed.data(), stride, dist,
                                                      (dir == direction::forward) ? FFTW_FORWARD : FFTW_BACKWARD, FFTW_ESTIMATE
                                     );
        }
    }
    plan_fftw(int size1, int size2, int size3){
        std::array<int, 3> size = {size3, size2, size1};
        plan = fftw_plan_many_dft(3, size.data(), 1, nullptr, nullptr, 1, 1, nullptr, nullptr, 1, 1,
                                  (dir == direction::forward) ? FFTW_FORWARD : FFTW_BACKWARD, FFTW_ESTIMATE);
    }
    ~plan_fftw(){ fftw_destroy_plan(plan); }
    operator fftw_plan() const{ return plan; }
    fftw_plan plan;
};
 
class fftw_executor : public executor_base{
public:
    using executor_base::forward;
    using executor_base::backward;
    using executor_base::complex_size;
    template<typename index>
    fftw_executor(void*, box3d<index> const box, int dimension) :
        size(box.size[dimension]), size2(0),
        howmanyffts(fft1d_get_howmany(box, dimension)),
        stride(fft1d_get_stride(box, dimension)),
        dist((dimension == box.order[0]) ? size : 1),
        blocks((dimension == box.order[1]) ? box.osize(2) : 1),
        block_stride(box.osize(0) * box.osize(1)),
        total_size(box.count()),
        embed({0, 0})
    {}
    template<typename index>
    fftw_executor(void*, box3d<index> const box, int dir1, int dir2) :
        size(box.size[std::min(dir1, dir2)]), size2(box.size[std::max(dir1, dir2)]),
        blocks(1), block_stride(0), total_size(box.count()), embed({0, 0})
    {
        int odir1 = box.find_order(dir1);
        int odir2 = box.find_order(dir2);
 
        if (std::min(odir1, odir2) == 0 and std::max(odir1, odir2) == 1){
            stride = 1;
            dist = size * size2;
            howmanyffts = box.size[2];
        }else if (std::min(odir1, odir2) == 1 and std::max(odir1, odir2) == 2){
            stride = box.size[0];
            dist = 1;
            howmanyffts = box.size[0];
        }else{ // case of directions (0, 2)
            stride = 1;
            dist = size;
            embed = {static_cast<int>(box.size[2]), static_cast<int>(box.size[1] * box.size[0])};
            howmanyffts = box.size[1];
        }
    }
    template<typename index>
    fftw_executor(void*, box3d<index> const box) :
        size(box.size[0]), size2(box.size[1]), howmanyffts(box.size[2]),
        stride(0), dist(0),
        blocks(1), block_stride(0),
        total_size(box.count()),
        embed({0, 0})
    {}
 
    void forward(std::complex<float> data[], std::complex<float>*) const override{
        make_plan(cforward);
        for(int i=0; i<blocks; i++){
            fftwf_complex* block_data = reinterpret_cast<fftwf_complex*>(data + i * block_stride);
            fftwf_execute_dft(*cforward, block_data, block_data);
        }
    }
    void backward(std::complex<float> data[], std::complex<float>*) const override{
        make_plan(cbackward);
        for(int i=0; i<blocks; i++){
            fftwf_complex* block_data = reinterpret_cast<fftwf_complex*>(data + i * block_stride);
            fftwf_execute_dft(*cbackward, block_data, block_data);
        }
    }
    void forward(std::complex<double> data[], std::complex<double>*) const override{
        make_plan(zforward);
        for(int i=0; i<blocks; i++){
            fftw_complex* block_data = reinterpret_cast<fftw_complex*>(data + i * block_stride);
            fftw_execute_dft(*zforward, block_data, block_data);
        }
    }
    void backward(std::complex<double> data[], std::complex<double>*) const override{
        make_plan(zbackward);
        for(int i=0; i<blocks; i++){
            fftw_complex* block_data = reinterpret_cast<fftw_complex*>(data + i * block_stride);
            fftw_execute_dft(*zbackward, block_data, block_data);
        }
    }
 
    void forward(float const indata[], std::complex<float> outdata[], std::complex<float> *workspace) const override{
        for(int i=0; i<total_size; i++) outdata[i] = std::complex<float>(indata[i]);
        forward(outdata, workspace);
    }
    void backward(std::complex<float> indata[], float outdata[], std::complex<float> *workspace) const override{
        backward(indata, workspace);
        for(int i=0; i<total_size; i++) outdata[i] = std::real(indata[i]);
    }
    void forward(double const indata[], std::complex<double> outdata[], std::complex<double> *workspace) const override{
        for(int i=0; i<total_size; i++) outdata[i] = std::complex<double>(indata[i]);
        forward(outdata, workspace);
    }
    void backward(std::complex<double> indata[], double outdata[], std::complex<double> *workspace) const override{
        backward(indata, workspace);
        for(int i=0; i<total_size; i++) outdata[i] = std::real(indata[i]);
    }
 
    int box_size() const override{ return total_size; }
    size_t workspace_size() const override{ return 0; }
 
private:
    template<typename scalar_type, direction dir>
    void make_plan(std::unique_ptr<plan_fftw<scalar_type, dir>> &plan) const{
        if (not plan){
            if (dist == 0)
                plan = std::unique_ptr<plan_fftw<scalar_type, dir>>(new plan_fftw<scalar_type, dir>(size, size2, howmanyffts));
            else if (size2 == 0)
                plan = std::unique_ptr<plan_fftw<scalar_type, dir>>(new plan_fftw<scalar_type, dir>(size, howmanyffts, stride, dist));
            else
                plan = std::unique_ptr<plan_fftw<scalar_type, dir>>(new plan_fftw<scalar_type, dir>(size, size2, embed, howmanyffts, stride, dist));
        }
    }
 
    int size, size2, howmanyffts, stride, dist, blocks, block_stride, total_size;
    std::array<int, 2> embed;
    mutable std::unique_ptr<plan_fftw<std::complex<float>, direction::forward>> cforward;
    mutable std::unique_ptr<plan_fftw<std::complex<float>, direction::backward>> cbackward;
    mutable std::unique_ptr<plan_fftw<std::complex<double>, direction::forward>> zforward;
    mutable std::unique_ptr<plan_fftw<std::complex<double>, direction::backward>> zbackward;
};
 
template<direction dir>
struct plan_fftw<float, dir>{
    plan_fftw(int size, int howmanyffts, int stride, int rdist, int cdist) :
        plan((dir == direction::forward) ?
             fftwf_plan_many_dft_r2c(1, &size, howmanyffts, nullptr, nullptr, stride, rdist,
                                                   nullptr, nullptr, stride, cdist,
                                                   FFTW_ESTIMATE
                                   ) :
             fftwf_plan_many_dft_c2r(1, &size, howmanyffts, nullptr, nullptr, stride, cdist,
                                                   nullptr, nullptr, stride, rdist,
                                                   FFTW_ESTIMATE
                                   ))
        {}
    ~plan_fftw(){ fftwf_destroy_plan(plan); }
    operator fftwf_plan() const{ return plan; }
    fftwf_plan plan;
};
template<direction dir>
struct plan_fftw<double, dir>{
    plan_fftw(int size, int howmanyffts, int stride, int rdist, int cdist) :
        plan((dir == direction::forward) ?
             fftw_plan_many_dft_r2c(1, &size, howmanyffts, nullptr, nullptr, stride, rdist,
                                                   nullptr, nullptr, stride, cdist,
                                                   FFTW_ESTIMATE
                                   ) :
             fftw_plan_many_dft_c2r(1, &size, howmanyffts, nullptr, nullptr, stride, cdist,
                                                   nullptr, nullptr, stride, rdist,
                                                   FFTW_ESTIMATE
                                   ))
        {}
    ~plan_fftw(){ fftw_destroy_plan(plan); }
    operator fftw_plan() const{ return plan; }
    fftw_plan plan;
};
 
class fftw_executor_r2c : public executor_base{
public:
    using executor_base::forward;
    using executor_base::backward;
    template<typename index>
    fftw_executor_r2c(void*, box3d<index> const box, int dimension) :
        size(box.size[dimension]),
        howmanyffts(fft1d_get_howmany(box, dimension)),
        stride(fft1d_get_stride(box, dimension)),
        blocks((dimension == box.order[1]) ? box.osize(2) : 1),
        rdist((dimension == box.order[0]) ? size : 1),
        cdist((dimension == box.order[0]) ? size/2 + 1 : 1),
        rblock_stride(box.osize(0) * box.osize(1)),
        cblock_stride(box.osize(0) * (box.osize(1)/2 + 1)),
        rsize(box.count()),
        csize(box.r2c(dimension).count())
    {}
 
    void forward(float const indata[], std::complex<float> outdata[], std::complex<float>*) const override{
        make_plan(sforward);
        for(int i=0; i<blocks; i++){
            float *rdata = const_cast<float*>(indata + i * rblock_stride);
            fftwf_complex* cdata = reinterpret_cast<fftwf_complex*>(outdata + i * cblock_stride);
            fftwf_execute_dft_r2c(*sforward, rdata, cdata);
        }
    }
    void backward(std::complex<float> indata[], float outdata[], std::complex<float>*) const override{
        make_plan(sbackward);
        for(int i=0; i<blocks; i++){
            fftwf_complex* cdata = const_cast<fftwf_complex*>(reinterpret_cast<fftwf_complex const*>(indata + i * cblock_stride));
            fftwf_execute_dft_c2r(*sbackward, cdata, outdata + i * rblock_stride);
        }
    }
    void forward(double const indata[], std::complex<double> outdata[], std::complex<double>*) const override{
        make_plan(dforward);
        for(int i=0; i<blocks; i++){
            double *rdata = const_cast<double*>(indata + i * rblock_stride);
            fftw_complex* cdata = reinterpret_cast<fftw_complex*>(outdata + i * cblock_stride);
            fftw_execute_dft_r2c(*dforward, rdata, cdata);
        }
    }
    void backward(std::complex<double> indata[], double outdata[], std::complex<double>*) const override{
        make_plan(dbackward);
        for(int i=0; i<blocks; i++){
            fftw_complex* cdata = const_cast<fftw_complex*>(reinterpret_cast<fftw_complex const*>(indata + i * cblock_stride));
            fftw_execute_dft_c2r(*dbackward, cdata, outdata + i * rblock_stride);
        }
    }
 
    int box_size() const override{ return rsize; }
    int complex_size() const override{ return csize; }
    size_t workspace_size() const override{ return 0; }
 
private:
    template<typename scalar_type, direction dir>
    void make_plan(std::unique_ptr<plan_fftw<scalar_type, dir>> &plan) const{
        if (!plan) plan = std::unique_ptr<plan_fftw<scalar_type, dir>>(new plan_fftw<scalar_type, dir>(size, howmanyffts, stride, rdist, cdist));
    }
 
    int size, howmanyffts, stride, blocks;
    int rdist, cdist, rblock_stride, cblock_stride, rsize, csize;
    mutable std::unique_ptr<plan_fftw<float, direction::forward>> sforward;
    mutable std::unique_ptr<plan_fftw<double, direction::forward>> dforward;
    mutable std::unique_ptr<plan_fftw<float, direction::backward>> sbackward;
    mutable std::unique_ptr<plan_fftw<double, direction::backward>> dbackward;
};
 
template<> struct one_dim_backend<backend::fftw>{
    using executor = fftw_executor;
    using executor_r2c = fftw_executor_r2c;
};
 
template<> struct one_dim_backend<backend::fftw_cos>{
    using executor = real2real_executor<backend::fftw, cpu_cos_pre_pos_processor>;
    using executor_r2c = void;
};
template<> struct one_dim_backend<backend::fftw_sin>{
    using executor = real2real_executor<backend::fftw, cpu_sin_pre_pos_processor>;
    using executor_r2c = void;
};
 
template<> struct default_plan_options<backend::fftw>{
    static const bool use_reorder = true;
};
 
template<> struct default_plan_options<backend::fftw_cos>{
    static const bool use_reorder = true;
};
template<> struct default_plan_options<backend::fftw_sin>{
    static const bool use_reorder = true;
};
 
}
 
#endif
 
#endif   /* HEFFTE_BACKEND_FFTW_H */