FieldLayout.h

//
// Class FieldLayout
//   FieldLayout describes how a given index space (represented by an NDIndex
//   object) is distributed among MPI ranks. It performs the initial
//   partitioning. The user may request that a particular dimension not be
//   partitioned by flagging that axis as 'SERIAL' (instead of 'PARALLEL').
//
#ifndef IPPL_FIELD_LAYOUT_H
#define IPPL_FIELD_LAYOUT_H
 
#include <array>
#include <iostream>
#include <map>
#include <vector>
 
#include "Types/ViewTypes.h"
 
#include "Communicate/Communicator.h"
#include "Index/NDIndex.h"
#include "Partition/Partitioner.h"
 
namespace ippl {
 
    template <unsigned Dim>
    class FieldLayout;
 
    template <unsigned Dim>
    std::ostream& operator<<(std::ostream&, const FieldLayout<Dim>&);
 
    // enumeration used to describe a hypercube's relation to
    // a particular axis in a given bounded domain
    enum e_cube_tag {
        UPPER       = 0,
        LOWER       = 1,
        IS_PARALLEL = 2
    };
 
    namespace detail {
        KOKKOS_INLINE_FUNCTION constexpr unsigned int countHypercubes(unsigned int dim) {
            unsigned int ret = 1;
            for (unsigned int d = 0; d < dim; d++) {
                ret *= 3;
            }
            return ret;
        }
 
        constexpr unsigned int factorial(unsigned x) {
            return x == 0 ? 1 : x * factorial(x - 1);
        }
 
        constexpr unsigned int binomialCoefficient(unsigned a, unsigned b) {
            return factorial(a) / (factorial(b) * factorial(a - b));
        }
 
        template <unsigned Dim>
        constexpr unsigned int countCubes(unsigned m) {
            return (1 << (Dim - m)) * binomialCoefficient(Dim, m);
        }
 
        bool isUpper(unsigned int face);
 
        unsigned int getFaceDim(unsigned int face);
 
        template <unsigned Dim>
        unsigned int indexToFace(unsigned int index) {
            // facets are group low/high by axis
            unsigned int axis = index / 2;
            // the digit to subtract is determined by whether the index describes an upper
            // face (even index) or lower face (odd index) and that digit's position is
            // determined by the axis of the face
            unsigned int toRemove = (2 - index % 2) * countHypercubes(axis);
            // start with all 2s (in base 3) and change the correct digit to get the encoded face
            return countHypercubes(Dim) - 1 - toRemove;
        }
 
        template <
            unsigned Dim, typename... CubeTags,
            typename = std::enable_if_t<sizeof...(CubeTags) == Dim - 1>,
            typename = std::enable_if_t<std::conjunction_v<std::is_same<e_cube_tag, CubeTags>...>>>
        unsigned int getCube(e_cube_tag tag, CubeTags... tags) {
            if constexpr (Dim == 1) {
                return tag;
            } else {
                return tag + 3 * getCube<Dim - 1>(tags...);
            }
        }
 
        template <size_t... Idx>
        unsigned int getFace_impl(const std::array<e_cube_tag, sizeof...(Idx)>& args,
                                  const std::index_sequence<Idx...>&) {
            return getCube<sizeof...(Idx)>(args[Idx]...);
        }
 
        template <unsigned Dim>
        unsigned int getFace(unsigned int axis, e_cube_tag side) {
            std::array<e_cube_tag, Dim> args;
            args.fill(IS_PARALLEL);
            args[axis] = side;
            return getFace_impl(args, std::make_index_sequence<Dim>{});
        }
    }  // namespace detail
 
    template <unsigned Dim>
    class FieldLayout {
    public:
        using NDIndex_t        = NDIndex<Dim>;
        using view_type        = typename detail::ViewType<NDIndex_t, 1>::view_type;
        using host_mirror_type = typename view_type::host_mirror_type;
 
        struct bound_type {
            // lower bounds (ordering: x, y, z, ...)
            std::array<long, Dim> lo;
            // upper bounds (ordering: x, y, z, ...)
            std::array<long, Dim> hi;
 
            long size() const {
                long total = 1;
                for (unsigned d = 0; d < Dim; d++) {
                    total *= hi[d] - lo[d];
                }
                return total;
            }
        };
 
        using rank_list   = std::vector<int>;
        using bounds_list = std::vector<bound_type>;
 
        using neighbor_list       = std::array<rank_list, detail::countHypercubes(Dim) - 1>;
        using neighbor_range_list = std::array<bounds_list, detail::countHypercubes(Dim) - 1>;
 
        FieldLayout(const mpi::Communicator& = MPI_COMM_WORLD);
 
        FieldLayout(mpi::Communicator, const NDIndex<Dim>& domain, std::array<bool, Dim> decomp,
                    bool isAllPeriodic = false);
 
        // Destructor: Everything deletes itself automatically ... the base
        // class destructors inform all the FieldLayoutUser's we're going away.
        virtual ~FieldLayout() = default;
 
        // Initialization functions, only to be called by the user of FieldLayout
        // objects when the FieldLayout was created using the default constructor;
        // otherwise these are only called internally by the various non-default
        // FieldLayout constructors:
 
        void initialize(const NDIndex<Dim>& domain, std::array<bool, Dim> decomp,
                        bool isAllPeriodic = false);
 
        // Return the domain.
        const NDIndex<Dim>& getDomain() const { return gDomain_m; }
 
        // Compare FieldLayouts to see if they represent the same domain; if
        // dimensionalities are different, the NDIndex operator==() will return
        // false:
        template <unsigned Dim2>
        bool operator==(const FieldLayout<Dim2>& x) const {
 
            // Throw exception if the domains are not the same
            if (gDomain_m != x.getDomain()) {
                throw std::runtime_error("FieldLayout: only FieldLayouts with the same global domain should be compared");
            }
 
            return gDomain_m == x.getDomain();
        }
 
        bool operator==(const FieldLayout<Dim>& x) const {
 
            // Throw exception if the domains are not the same
            if (gDomain_m != x.getDomain()) {
                throw std::runtime_error("FieldLayout: only FieldLayouts with the same global domain should be compared");
            }
 
            for (unsigned int i = 0; i < Dim; ++i) {
                if (hLocalDomains_m(comm.rank())[i] != x.getLocalNDIndex()[i]) {
                    return false;
                }
            }
            return true;
        }
 
        // for the requested dimension, report if the distribution is
        // SERIAL or PARALLEL
        bool getDistribution(unsigned int d) const {
            return minWidth_m[d] == (unsigned int)gDomain_m[d].length();
        }
 
        // for the requested dimension, report if the distribution was requested to
        // be SERIAL or PARALLEL
        std::array<bool, Dim> isParallel() const { return isParallelDim_m; }
 
        // Get the local domain for the current rank.
        const NDIndex_t& getLocalNDIndex() const;
 
        // Get the local domain for a specific rank.
        const NDIndex_t& getLocalNDIndex(int rank) const;
 
        const host_mirror_type getHostLocalDomains() const;
 
        const view_type getDeviceLocalDomains() const;
 
        const neighbor_list& getNeighbors() const;
 
        const neighbor_range_list& getNeighborsSendRange() const;
 
        const neighbor_range_list& getNeighborsRecvRange() const;
 
        static int getMatchingIndex(int index);
 
        void findPeriodicNeighbors(const int nghost, const NDIndex<Dim>& localDomain,
                                   NDIndex<Dim>& grown, NDIndex<Dim>& neighborDomain,
                                   const int rank, std::map<unsigned int, int>& offsets,
                                   unsigned d0 = 0, unsigned codim = 0);
 
        void findNeighbors(int nghost = 1);
 
        void addNeighbors(const NDIndex_t& gnd, const NDIndex_t& nd, const NDIndex_t& ndNeighbor,
                          const NDIndex_t& intersect, int nghost, int rank);
 
        void write(std::ostream& = std::cout) const;
 
        void updateLayout(const std::vector<NDIndex_t>& domains);
 
        bool isAllPeriodic_m;
 
        mpi::Communicator comm;
 
    private:
        bound_type getBounds(const NDIndex_t& nd1, const NDIndex_t& nd2, const NDIndex_t& offset,
                             int nghost);
 
        int getPeriodicOffset(const NDIndex_t& nd, const unsigned int d, const int k);
 
        // List of all the neighboring ranks, arranged by ternary encoding
        neighbor_list neighbors_m;
 
        // Lower and upper bounds of the regions to send and receive to and from neighbors,
        // arranged by ternary encoding
        neighbor_range_list neighborsSendRange_m, neighborsRecvRange_m;
 
    protected:
        NDIndex_t gDomain_m;
 
        view_type dLocalDomains_m;
 
        host_mirror_type hLocalDomains_m;
 
        std::array<bool, Dim> isParallelDim_m;
 
        // Minimum width of all the local domains for each dimension
        unsigned int minWidth_m[Dim];
 
        void calcWidths();
    };
 
    template <unsigned Dim>
    inline std::ostream& operator<<(std::ostream& out, const FieldLayout<Dim>& f) {
        f.write(out);
        return out;
    }
}  // namespace ippl
 
#include "FieldLayout/FieldLayout.hpp"
 
#endif