heffte_geometry.h

/*
    -- heFFTe --
       Univ. of Tennessee, Knoxville
       @date
*/
 
#ifndef HEFFTE_GEOMETRY_H
#define HEFFTE_GEOMETRY_H
 
#include <array>
#include <valarray>
#include "heffte_utils.h"
 
namespace heffte {
 
template<typename index = int>
struct box3d{
    box3d(std::array<index, 3> clow, std::array<index, 3> chigh) :
        low(clow), high(chigh), size({high[0] - low[0] + 1, high[1] - low[1] + 1, high[2] - low[2] + 1}), order({0, 1, 2})
    {}
    box3d(std::array<index, 3> clow, std::array<index, 3> chigh, std::array<int, 3> corder) :
        low(clow), high(chigh), size({high[0] - low[0] + 1, high[1] - low[1] + 1, high[2] - low[2] + 1}), order(corder)
    {}
    box3d(std::array<index, 2> clow, std::array<index, 2> chigh) :
        box3d(std::array<index, 3>{clow[0], clow[1], 0}, std::array<index, 3>{chigh[0], chigh[1], 0}){}
    box3d(std::array<index, 2> clow, std::array<index, 2> chigh, std::array<int, 2> corder) :
        box3d(std::array<index, 3>{clow[0], clow[1], 0}, std::array<index, 3>{chigh[0], chigh[1], 0},
              std::array<index, 3>{corder[0], corder[1], 2}){}
    bool empty() const{ return (size[0] <= 0 or size[1] <= 0 or size[2] <= 0); }
    long long count() const{ return (empty()) ? 0 :
        (static_cast<long long>(size[0]) * static_cast<long long>(size[1]) * static_cast<long long>(size[2])); }
    box3d collide(box3d const other) const{
        if (empty() or other.empty()) return box3d({0, 0, 0}, {-1, -1, -1});
        return box3d({std::max(low[0], other.low[0]), std::max(low[1], other.low[1]), std::max(low[2], other.low[2])},
                     {std::min(high[0], other.high[0]), std::min(high[1], other.high[1]), std::min(high[2], other.high[2])}, order);
    }
    box3d r2c(int dimension) const{
        if (empty()) return box3d({0, 0, 0}, {-1, -1, -1});
        switch(dimension){
            case 0: return box3d(low, {low[0] + size[0] / 2, high[1], high[2]}, order);
            case 1: return box3d(low, {high[0], low[1] + size[1] / 2, high[2]}, order);
            default: // dimension == 2
                return box3d(low, {high[0], high[1], low[2] + size[2] / 2}, order);
        }
    }
    bool is2d() const{ return (size[0] == 1 or size[1] == 1 or size[2] == 1); }
    bool operator == (box3d const &other) const{
        return not (*this != other);
    }
    bool operator != (box3d const &other) const{
        for(int i=0; i<3; i++)
            if (low[i] != other.low[i] or high[i] != other.high[i]) return true;
        return false;
    }
    bool ordered_same_as(box3d const &other) const{
        return (order[0] == other.order[0]) and (order[1] == other.order[1]) and (order[2] == other.order[2]);
    }
    int find_order(int dir) const{ return ((dir == order[0]) ? 0 : ((dir == order[1]) ? 1 : 2)); }
    index osize(int dimension) const{ return size[order[dimension]]; }
    std::array<index, 3> const low;
    std::array<index, 3> const high;
    std::array<index, 3> const size;
    std::array<int, 3> const order;
};
 
template<typename index = int>
using box2d = box3d<index>;
 
template<typename index>
inline std::ostream & operator << (std::ostream &os, box3d<index> const box){
    for(int i=0; i<3; i++)
        os << box.low[i] << "  " << box.high[i] << "  (" << box.size[i] << ")\n";
    os << "(" << box.order[0] << "," << box.order[1] << "," << box.order[2] << ")\n";
    os << "\n";
    return os;
}
 
template<typename index>
inline int fft1d_get_howmany(box3d<index> const box, int const dimension){
    if (dimension == box.order[0]) return box.osize(1) * box.osize(2);
    if (dimension == box.order[1]) return box.osize(0);
    return box.osize(0) * box.osize(1);
}
template<typename index>
inline int fft1d_get_stride(box3d<index> const box, int const dimension){
    if (dimension == box.order[0]) return 1;
    if (dimension == box.order[1]) return box.osize(0);
    return box.osize(0) * box.osize(1);
}
 
struct rank_remap{
    rank_remap() : mpi_rank(-1), size_subcomm(0){}
    rank_remap(int rank) : mpi_rank(rank), size_subcomm(0){}
    int const mpi_rank;
    std::vector<int> map;
    void set_subranks(size_t all_ranks, int sub_ranks){
        size_subcomm = sub_ranks;
        map = std::vector<int>(all_ranks, -1);
        for(int i=0; i<sub_ranks; i++) map[i] = i;
    }
    void set_subranks(MPI_Comm global, MPI_Comm subcomm){
        int subcomm_rank = (subcomm == MPI_COMM_NULL) ? -1 : mpi::comm_rank(subcomm);
        map = std::vector<int>(mpi::comm_size(global));
        MPI_Allgather(&subcomm_rank, 1, MPI_INT, map.data(), 1, MPI_INT, global);
    }
    int size_subcomm;
    bool empty() const{ return map.empty(); }
};
 
template<typename index = int>
struct ioboxes{
    std::vector<box3d<index>> in;
    std::vector<box3d<index>> out;
};
 
template<typename index>
inline box3d<index> find_world(std::vector<box3d<index>> const &boxes){
    std::array<index, 3> low  = boxes[0].low;
    std::array<index, 3> high = boxes[0].high;
    for(auto b : boxes){
        for(index i=0; i<3; i++)
            low[i] = std::min(low[i], b.low[i]);
        for(index i=0; i<3; i++)
            high[i] = std::max(high[i], b.high[i]);
    }
    return {low, high};
}
 
template<typename index>
inline bool match(std::vector<box3d<index>> const &shape0, std::vector<box3d<index>> const &shape1){
    if (shape0.size() != shape1.size()) return false;
    for(size_t i=0; i<shape0.size(); i++)
        if (shape0[i] != shape1[i])
            return false;
    return true;
}
 
template<typename index>
inline bool world_complete(std::vector<box3d<index>> const &boxes, box3d<index> const world){
    long long wsize = 0;
    for(auto b : boxes) wsize += b.count();
    if (wsize < world.count())
        throw std::invalid_argument("The provided input boxes do not fill the world box!");
 
    for(size_t i=0; i<3; i++)
        if (world.low[i] != 0)
            throw std::invalid_argument("Global box indexing must start from 0!");
 
    for(size_t i=0; i<boxes.size(); i++)
        for(size_t j=0; j<boxes.size(); j++)
            if (i != j and not boxes[i].collide(boxes[j]).empty())
                throw std::invalid_argument("Input boxes cannot overlap!");
 
    return true;
}
 
inline std::vector<std::array<int, 2>> get_factors(int const n){
    std::vector<std::array<int, 2>> result;
    for(int i=1; i<=n; i++){
        if (n % i == 0)
            result.push_back({i, n / i});
    }
    if (n == 1) {
        result.push_back({1, 1});
    }
    return result;
}
 
inline int get_area(std::array<int, 2> const &dims){ return dims[0] * dims[1] + dims[0] + dims[1]; }
 
 
inline std::array<int, 2> make_procgrid(int const num_procs){
    auto factors = get_factors(num_procs);
    std::array<int, 2> min_array = factors.front();
    int min_area = get_area(min_array);
    for(auto f : factors){
        int farea = get_area(f);
        if (farea < min_area){
            min_array = f;
            min_area = farea;
        }
    }
    return min_array;
}
 
template<typename index>
inline std::array<int, 3> make_procgrid2d(box3d<index> const world, int direction_1d, std::array<int, 2> const candidate_grid){
    auto make_grid = [&](std::array<int, 2> const &grid)
                     ->std::array<int, 3>{
                         return  (direction_1d == 0) ?  std::array<int, 3>{1, grid[0], grid[1]} :
                                ((direction_1d == 1) ? std::array<int, 3>{grid[0], 1, grid[1]} :
                                                       std::array<int, 3>{grid[0], grid[1], 1});
                     };
    std::array<int, 3> result = make_grid(candidate_grid);
 
    auto valid = [&](std::array<int, 3> const &grid)
                 ->bool{
                     for(int i=0; i<3; i++) if (grid[i] > world.size[i]) return false;
                     return true;
                 };
    if (valid(result)) return result; // if the first constraint is OK
 
    // otherwise seek a new factorization
    auto factors = get_factors(candidate_grid[0] * candidate_grid[1]);
 
    int min_area = get_area(factors.front());
    // a bit misleading, but here we initialize min_area with the largest possible area
    for(auto const &g : factors) min_area = std::max(min_area, get_area(g));
 
    for(auto const &g : factors){
        if (valid(make_grid(g)) and get_area(g) <= min_area){
            result = make_grid(g);
            min_area = get_area(g);
        }
    }
 
    if (not valid(result)) throw std::runtime_error("Cannot split the given number of indexes into the given set of mpi-ranks. Most liklely, the number of indexes is too small compared to the number of mpi-ranks.");
 
    return result;
}
 
template<typename index>
inline std::vector<box3d<index>> split_world(box3d<index> const world, std::array<int, 3> const proc_grid, rank_remap const &remap = rank_remap()){
 
    auto fast = [=](index i)->index{ return world.low[0] + i * (world.size[0] / proc_grid[0]) + std::min(i, (world.size[0] % proc_grid[0])); };
    auto mid  = [=](index i)->index{ return world.low[1] + i * (world.size[1] / proc_grid[1]) + std::min(i, (world.size[1] % proc_grid[1])); };
    auto slow = [=](index i)->index{ return world.low[2] + i * (world.size[2] / proc_grid[2]) + std::min(i, (world.size[2] % proc_grid[2])); };
 
    std::vector<box3d<index>> result;
    result.reserve(proc_grid[0] * proc_grid[1] * proc_grid[2]);
    for(index k = 0; k < proc_grid[2]; k++){
        for(index j = 0; j < proc_grid[1]; j++){
            for(index i = 0; i < proc_grid[0]; i++){
                result.push_back(box3d<index>({fast(i), mid(j), slow(k)}, {fast(i+1)-1, mid(j+1)-1, slow(k+1)-1}, world.order));
            }
        }
    }
    if (remap.empty()){
        return result;
    }else{
        std::vector<box3d<index>> remapped_result;
        for(size_t i=0; i<remap.map.size(); i++)
            remapped_result.push_back( (remap.map[i] == -1) ?
                                       box3d<index>(std::array<index, 3>{0, 0, 0}, std::array<index, 3>{-1, -1, -1}, world.order) :
                                       result[remap.map[i]]
                            );
        return remapped_result;
    }
}
 
template<typename index>
inline bool is_pencils(box3d<index> const world, std::vector<box3d<index>> const &shape, int direction){
    for(auto s : shape)
        if (s.size[direction] != world.size[direction])
            return false;
    return true;
}
 
template<typename index>
inline bool is_slab(box3d<index> const world, std::vector<box3d<index>> const &shape, int direction1, int direction2){
    for(auto s : shape)
        if (s.size[direction1] != world.size[direction1] or s.size[direction2] != world.size[direction2])
            return false;
    return true;
}
 
template<typename index>
inline std::vector<box3d<index>> reorder(std::vector<box3d<index>> const &shape, std::array<int, 3> order){
    std::vector<box3d<index>> result;
    result.reserve(shape.size());
    for(auto const &b : shape)
        result.push_back(box3d<index>(b.low, b.high, order));
    return result;
}
 
template<typename index>
inline std::vector<box3d<index>> maximize_overlap(std::vector<box3d<index>> const &new_boxes,
                                                  std::vector<box3d<index>> const &old_boxes,
                                                  std::array<int, 3> const order,
                                                  rank_remap const &remap){
    std::vector<box3d<index>> result;
    result.reserve(new_boxes.size());
    std::vector<bool> taken(new_boxes.size(), false);
    if (not remap.empty()){
        for(size_t i=0; i<remap.map.size(); i++)
            if (remap.map[i] == -1) taken[i] = true;
    }
 
    for(size_t i=0; i<new_boxes.size(); i++){
        if (not remap.empty() and remap.map[i] == -1){
            result.push_back(box3d<index>(new_boxes[i].low, new_boxes[i].high, order));
            continue;
        }
        // for each box in the result, find the box among the new_boxes
        // that has not been taken and has the largest overlap with the corresponding old box
        int max_overlap = -1;
        size_t max_index = new_boxes.size();
        for(size_t j=0; j<new_boxes.size(); j++){
            int overlap = old_boxes[i].collide(new_boxes[j]).count();
            if (not taken[j] and overlap > max_overlap){
                max_overlap = overlap;
                max_index = j;
            }
        }
        assert( max_index < new_boxes.size() ); // if we found a box
        taken[max_index] = true;
        result.push_back(box3d<index>(new_boxes[max_index].low, new_boxes[max_index].high, order));
    }
 
    return result;
}
 
template<typename index>
inline long long count_connections(std::vector<box3d<index>> const &new_boxes, std::vector<box3d<index>> const &old_boxes){
    long long count = 0;
 
    for(auto &nbox : new_boxes)
        for(auto &obox : old_boxes)
            if (not nbox.collide(obox).empty())
                count++;
 
    return count;
}
 
template<typename index>
inline std::vector<box3d<index>> make_pencils(box3d<index> const world,
                                              std::array<int, 2> const proc_grid,
                                              int const dimension,
                                              std::vector<box3d<index>> const &source,
                                              std::array<int, 3> const order,
                                              rank_remap const &remap = rank_remap()
                                             ){
    // trivial case, the grid is already in a pencil format
    if (is_pencils(world, source, dimension))
        return reorder(source, order);
 
    // create a list of boxes ordered in column major format (following the proc_grid box)
    // using two boxes corresponding to the two dimensions of proc-grid in both variants
    std::vector<box3d<index>> pencilsA =
        maximize_overlap(
            split_world(world, make_procgrid2d(world, dimension, proc_grid), remap), source, order, remap);
 
    std::vector<box3d<index>> pencilsB =
        maximize_overlap(
            split_world(world, make_procgrid2d(world, dimension, std::array<int, 2>{proc_grid[1], proc_grid[0]}), remap), source, order, remap);
 
    return (count_connections(pencilsB, source) < count_connections(pencilsA, source)) ? pencilsB : pencilsA;
}
 
template<typename index>
inline std::vector<box3d<index>> make_slabs(box3d<index> const world, int num_slabs,
                                            int const dimension1, int const dimension2,
                                            std::vector<box3d<index>> const &source,
                                            std::array<int, 3> const order,
                                            rank_remap const &remap
                                           ){
    assert( dimension1 != dimension2 );
    std::vector<box3d<index>> slabs;
    if (dimension1 == 0){
        if (dimension2 == 1){
            slabs = split_world(world, {1, 1, num_slabs}, remap);
        }else{
            slabs = split_world(world, {1, num_slabs, 1}, remap);
        }
    }else if (dimension1 == 1){
        if (dimension2 == 0){
            slabs = split_world(world, {1, 1, num_slabs}, remap);
        }else{
            slabs = split_world(world, {num_slabs, 1, 1}, remap);
        }
    }else{ // third dimension
        if (dimension2 == 0){
            slabs = split_world(world, {1, num_slabs, 1}, remap);
        }else{
            slabs = split_world(world, {num_slabs, 1, 1}, remap);
        }
    }
 
    return maximize_overlap(slabs, source, order, remap);
}
 
template<typename index>
inline std::array<int, 3> proc_setup_min_surface(box3d<index> const world, int num_procs){
    assert(world.count() > 0); // make sure the world is not empty
 
    // using valarrays that work much like vectors, but can perform basic
    // point-wise operations such as addition, multiply, and division
    std::valarray<index> all_indexes = {world.size[0], world.size[1], world.size[2]};
    // set initial guess, probably the worst grid but a valid one
    std::valarray<index> best_grid = {1, 1, num_procs};
 
    // internal helper method to compute the surface
    auto surface = [&](std::valarray<index> const &proc_grid)->
        index{
            auto box_size = all_indexes / proc_grid;
            return ( box_size * box_size.cshift(1) ).sum();
        };
 
    index best_surface = surface({1, 1, num_procs});
 
    for(int i=1; i<=num_procs; i++){
        if (num_procs % i == 0){
            int const remainder = num_procs / i;
            for(int j=1; j<=remainder; j++){
                if (remainder % j == 0){
                    std::valarray<index> candidate_grid = {i, j, remainder / j};
                    index const candidate_surface = surface(candidate_grid);
                    if (candidate_surface < best_surface){
                        best_surface = candidate_surface;
                        best_grid    = candidate_grid;
                    }
                }
            }
        }
    }
 
    assert(best_grid[0] * best_grid[1] * best_grid[2] == num_procs);
 
    return {static_cast<int>(best_grid[0]), static_cast<int>(best_grid[1]), static_cast<int>(best_grid[2])};
}
 
namespace mpi {
template<typename index>
inline ioboxes<index> gather_boxes(box3d<index> const my_inbox, box3d<index> const my_outbox, MPI_Comm const comm){
    std::array<box3d<index>, 2> my_data = {my_inbox, my_outbox};
    std::vector<box3d<index>> all_boxes(2 * mpi::comm_size(comm), box3d<index>({0, 0, 0}, {0, 0, 0}));
    MPI_Allgather(&my_data, 2 * sizeof(box3d<index>), MPI_BYTE, all_boxes.data(), 2 * sizeof(box3d<index>), MPI_BYTE, comm);
    ioboxes<index> result;
    for(auto i = all_boxes.begin(); i < all_boxes.end(); i += 2){
        result.in.push_back(*i);
        result.out.push_back(*(i+1));
    }
    return result;
}
 
}
 
}
 
 
#endif /* HEFFTE_GEOMETRY_H */