heffte_plan_logic.h

/*
    -- heFFTe --
       Univ. of Tennessee, Knoxville
       @date
*/
 
#ifndef HEFFTE_PLAN_LOGIC_H
#define HEFFTE_PLAN_LOGIC_H
 
#include "heffte_common.h"
 
namespace heffte {
 
enum class reshape_algorithm{
    alltoallv = 0,
    alltoall = 3,
    p2p_plined = 1,
    p2p = 2
};
 
struct plan_options{
    template<typename backend_tag> plan_options(backend_tag const)
        : use_reorder(default_plan_options<backend_tag>::use_reorder),
          algorithm(reshape_algorithm::alltoallv),
          use_pencils(true),
          use_gpu_aware(true),
          num_sub(-1),
          subcomm(MPI_COMM_NULL)
    {}
    plan_options(bool reorder, reshape_algorithm alg, bool pencils)
        : use_reorder(reorder), algorithm(alg), use_pencils(pencils), use_gpu_aware(true), num_sub(-1), subcomm(MPI_COMM_NULL)
    {}
    bool use_reorder;
    reshape_algorithm algorithm;
    bool use_pencils;
    bool use_gpu_aware;
    void use_num_subranks(int num_subranks){ num_sub = num_subranks; }
    void use_subcomm(MPI_Comm comm){
        num_sub = 1;
        subcomm = comm;
    }
    int get_subranks() const{ return num_sub; }
private:
    int num_sub;
    MPI_Comm subcomm;
};
 
inline std::ostream & operator << (std::ostream &os, plan_options const options){
    std::string algorithm = "";
    switch (options.algorithm){
        case reshape_algorithm::alltoallv  : algorithm = "mpi:alltoallv"; break;
        case reshape_algorithm::alltoall   : algorithm = "mpi:alltoall"; break;
        case reshape_algorithm::p2p_plined : algorithm = "mpi:point-to-point-pipelined"; break;
        case reshape_algorithm::p2p        : algorithm = "mpi:point-to-point"; break;
    };
    os << "options = ("
       << ((options.use_reorder) ? "fft1d:contiguous" : "fft1d:strided") << ", "
       << algorithm << ", "
       << ((options.use_pencils) ? "decomposition:pencil" : "decomposition:slab") << ", "
       << ((options.use_gpu_aware) ? "mpi:from-gpu" : "mpi:from-cpu") << ")";
    return os;
}
 
template<typename backend_tag, bool use_r2c = false>
plan_options set_options(plan_options opts){
    if (std::is_same<backend_tag, backend::stock_cos>::value
        or std::is_same<backend_tag, backend::fftw_cos>::value
        or std::is_same<backend_tag, backend::mkl_cos>::value
        or std::is_same<backend_tag, backend::cufft_cos>::value
        or std::is_same<backend_tag, backend::rocfft_cos>::value
        or std::is_same<backend_tag, backend::onemkl_cos>::value
        or std::is_same<backend_tag, backend::stock_sin>::value
        or std::is_same<backend_tag, backend::fftw_sin>::value
        or std::is_same<backend_tag, backend::mkl_sin>::value
        or std::is_same<backend_tag, backend::cufft_sin>::value
        or std::is_same<backend_tag, backend::rocfft_sin>::value
        or std::is_same<backend_tag, backend::onemkl_sin>::value
    ){
        // currently the cosine options work only with reorder.
        opts.use_reorder = true;
        return opts;
    }else if (use_r2c and std::is_same<backend_tag, backend::rocfft>::value){
        // the rocfft backend with r2c requires the reorder (problem with the strides)
        opts.use_reorder = true;
        return opts;
    }else{
        return opts; // all options are supported for this backend
    }
}
 
inline plan_options force_reorder(plan_options opts){
    opts.use_reorder = true;
    return opts;
}
 
template<typename backend_tag>
plan_options default_options(){
    return plan_options(backend_tag());
}
 
template<typename index>
struct logic_plan3d{
    std::vector<box3d<index>> in_shape[4];
    std::vector<box3d<index>> out_shape[4];
    std::array<int, 3> fft_direction;
    long long index_count;
    plan_options const options;
    int const mpi_rank;
};
 
template<typename index>
inline std::array<bool, 3> pencil_directions(box3d<index> const world, std::vector<box3d<index>> const &boxes){
    std::array<bool, 3> is_pencil = {true, true, true};
    for(auto const &b : boxes){
        for(int i=0; i<3; i++)
            is_pencil[i] = is_pencil[i] and (world.size[i] == b.size[i]);
    }
    return is_pencil;
}
 
template<typename index>
logic_plan3d<index> plan_operations(ioboxes<index> const &boxes, int r2c_direction, plan_options const options, int const mpi_rank);
 
template<typename index>
std::vector<std::array<int, 3>> compute_grids(logic_plan3d<index> const &plan);
 
}
 
#endif