lattice.h

#ifndef HILA_LATTICE_H
#define HILA_LATTICE_H
 
#include <sstream>
#include <iostream>
#include <fstream>
#include <array>
#include <vector>
 
// SUBNODE_LAYOUT is now defined in main.mk
// #define SUBNODE_LAYOUT
 
#include "plumbing/defs.h"
#include "plumbing/coordinates.h"
#include "plumbing/timing.h"
 
#ifdef SUBNODE_LAYOUT
#ifndef VECTOR_SIZE
#if defined(CUDA) || defined(HIP)
#define VECTOR_SIZE 8 // Size of float, length 1 vectors
#else
#define VECTOR_SIZE (256 / 8) // this is for AVX2
#endif
#endif
// This is the vector size used to determine the layout
constexpr unsigned number_of_subnodes = VECTOR_SIZE / sizeof(float);
#endif
 
namespace hila {
/// list of field boundary conditions - used only if SPECIAL_BOUNDARY_CONDITIONS defined
enum class bc { PERIODIC, ANTIPERIODIC, DIRICHLET };
 
/// False if we have b.c. which does not require communication
inline bool bc_need_communication(hila::bc bc) {
    if (bc == hila::bc::DIRICHLET) {
        return false;
    } else {
        return true;
    }
}
 
} // namespace hila
 
void test_std_gathers();
void report_too_large_node(); // report on too large node
 
/// useful information about a node
struct node_info {
    CoordinateVector min, size;
    unsigned evensites, oddsites;
};
 
namespace hila {
extern int64_t n_gather_done, n_gather_avoided;
}
 
/// Some backends need specialized lattice data
/// in loops. Forward declaration here and
/// implementations in backend headers.
/// Loops generated by hilapp can access
/// this through lattice.backend_lattice.
struct backend_lattice_struct;
 
/// The lattice struct defines the lattice geometry ans sets up MPI communication
/// patterns
class lattice_struct {
  public:
    CoordinateVector l_size;
    size_t l_volume = 0; // use this to flag initialization
 
    int l_label; // running number, identification of the lattice (TODO)
 
    // Guarantee 64 bits for these - 32 can overflow!
    MPI_Comm mpi_comm_lat;
 
    // is this lattice derived from another, through e.g. .block()?  Pointer to parent lattice
    lattice_struct *parent;
 
    /// Information about the node stored on this process
    struct node_struct {
        lattice_struct *parent; // parent lattice, for methods
        size_t volume, evensites, oddsites;
        size_t field_alloc_size;    // how many sites/node in allocations
        CoordinateVector min, size; // node local coordinate ranges
        int rank;                   // rank of this node
        bool first_site_even;       // is location min even or odd?
 
#ifdef EVEN_SITES_FIRST
        std::vector<CoordinateVector> coordinates;
#endif
 
        Vector<NDIM, unsigned> size_factor; // components: 1, size[0], size[0]*size[1], ...
 
        void setup(node_info &ni, lattice_struct &lattice);
 
#ifdef SUBNODE_LAYOUT
        /// If we have vectorized-style layout, we introduce "subnodes"
        /// size is mynode.size/subnodes.divisions, which is not
        /// constant across nodes
        struct subnode_struct {
            CoordinateVector divisions, size; // div to subnodes to directions, size
            size_t sites, evensites, oddsites;
            Direction merged_subnodes_dir;
 
            void setup(const node_struct &tn);
        } subnodes;
#endif
 
        /// true if this node is on the edge of the lattice to dir d
        bool is_on_edge(Direction d) const {
            return (is_up_dir(d) && min[d] + size[d] == parent->l_size[d]) ||
                   (!is_up_dir(d) && min[-d] == 0);
        }
 
    } mynode;
 
    /// information about all nodes
    struct allnodes {
        int number;                   // number of nodes
        CoordinateVector n_divisions; // number of node divisions to dir
        // lattice division: divisors[d] will have n_divisions[d]+1 elements, last size to d
        std::vector<unsigned> divisors[NDIM];
        CoordinateVector max_size; // size of largest node
 
        // std::vector<node_info> nodelist;
        node_info nodeinfo(int i) const;
 
        int *RESTRICT map_array;        // mapping (optional)
        int *RESTRICT map_inverse;      // inv of it
        void create_remap();            // create remap_node
        int remap(int i) const;         // use remap
        int inverse_remap(int i) const; // inverse remap
 
    } nodes;
 
    /// Information necessary to communicate with a node
    struct comm_node_struct {
        unsigned rank; // rank of communicated with node
        size_t sites, evensites, oddsites;
        size_t buffer; // offset from the start of field array
        unsigned *sitelist;
 
        // Get a vector containing the sites of parity par and number of elements
        const unsigned *RESTRICT get_sitelist(Parity par, int &size) const {
            if (par == ALL) {
                size = sites;
                return sitelist;
            } else if (par == EVEN) {
                size = evensites;
                return sitelist;
            } else {
                size = oddsites;
                return sitelist + evensites;
            }
        }
 
        // The number of sites that need to be communicated
        unsigned n_sites(Parity par) const {
            if (par == ALL) {
                return sites;
            } else if (par == EVEN) {
                return evensites;
            } else {
                return oddsites;
            }
        }
 
        // The local index of a site that is sent to neighbour
        unsigned site_index(unsigned site, Parity par) const {
            if (par == ODD) {
                return sitelist[evensites + site];
            } else {
                return sitelist[site];
            }
        }
 
        // The offset of the halo from the start of the field array
        unsigned offset(Parity par) const {
            if (par == ODD) {
                return buffer + evensites;
            } else {
                return buffer;
            }
        }
 
        void init() {
            sites = evensites = oddsites = 0;
            buffer = 0;
            sitelist = nullptr;
            rank = hila::myrank();
        }
    };
 
    /// nn-communication has only 1 node to talk to
    struct nn_comminfo_struct {
        unsigned *index;
        comm_node_struct from_node, to_node;
        unsigned receive_buf_size; // only for general gathers
    };
 
    /// general communication
    struct gen_comminfo_struct {
        unsigned *index;
        std::vector<comm_node_struct> from_node;
        std::vector<comm_node_struct> to_node;
        size_t receive_buf_size;
    };
 
    /// nearest neighbour comminfo struct
    std::array<nn_comminfo_struct, NDIRS> nn_comminfo;
 
    /// Main neighbour index array
    unsigned *RESTRICT neighb[NDIRS];
 
    /// implement waiting using mask_t - unsigned char is good for up to 4 dim.
    dir_mask_t *RESTRICT wait_arr_;
 
#ifdef SPECIAL_BOUNDARY_CONDITIONS
    /// special boundary pointers are needed only in cases neighbour
    /// pointers must be modified (new halo elements). That is known only during
    /// runtime.
    struct special_boundary_struct {
        unsigned *neighbours;
        unsigned *move_index;
        size_t offset, n_even, n_odd, n_total;
        bool is_needed;
    };
    // holder for nb ptr info
    special_boundary_struct special_boundaries[NDIRS];
#endif
 
#ifndef VANILLA
    backend_lattice_struct *backend_lattice;
#endif
 
    void setup_base_lattice(const CoordinateVector &siz);
 
    void setup_layout();
    void setup_nodes();
    void setup_node_divisors();
 
    // std::vector<node_info> nodelist() { return nodes.nodelist; }
    // CoordinateVector min_coordinate() const { return mynode.min; }
    // int min_coordinate(Direction d) const { return mynode.min[d]; }
 
    bool is_on_mynode(const CoordinateVector &c) const;
    int node_rank(const CoordinateVector &c) const;
    unsigned site_index(const CoordinateVector &c) const;
    unsigned site_index(const CoordinateVector &c, const unsigned node) const;
 
    void create_std_gathers();
    gen_comminfo_struct create_general_gather(const CoordinateVector &r);
    std::vector<comm_node_struct> create_comm_node_vector(CoordinateVector offset, unsigned *index,
                                                          bool receive);
 
    // bool first_site_even() const {
    //     return mynode.first_site_even;
    // };
 
#ifdef SPECIAL_BOUNDARY_CONDITIONS
    void init_special_boundaries();
    void setup_special_boundary_array(Direction d);
 
    const unsigned *get_neighbour_array(Direction d, hila::bc bc);
#else
    const unsigned *get_neighbour_array(Direction d, hila::bc bc) {
        return neighb[d];
    }
#endif
 
 
#ifdef EVEN_SITES_FIRST
    unsigned loop_begin(Parity P) const {
        if (P == ODD) {
            return mynode.evensites;
        } else {
            return 0;
        }
    }
    unsigned loop_end(Parity P) const {
        if (P == EVEN) {
            return mynode.evensites;
        } else {
            return mynode.volume;
        }
    }
 
    inline const CoordinateVector &coordinates(unsigned idx) const {
        return mynode.coordinates[idx];
    }
 
    inline int coordinate(unsigned idx, Direction d) const {
        return mynode.coordinates[idx][d];
    }
 
    inline Parity site_parity(unsigned idx) const {
        if (idx < mynode.evensites)
            return EVEN;
        else
            return ODD;
    }
 
#else // Now not EVEN_SITES_FIRST
 
    unsigned loop_begin(Parity P) const {
        assert(P == ALL && "Only parity ALL when EVEN_SITES_FIRST is off");
        return 0;
    }
    unsigned loop_end(Parity P) const {
        return mynode.volume;
    }
 
    // Use computation to get coordinates: from fastest
    // to lowest, dir = 0, 1, 2, ...
    // each coordinate is c[d] = (idx/size_factor[d]) % size[d] + min[d], but
    // do it like below to avoid the mod
 
    inline const CoordinateVector coordinates(unsigned idx) const {
        CoordinateVector c;
        unsigned vdiv, ndiv;
 
        vdiv = idx;
        for (int d = 0; d < NDIM - 1; ++d) {
            ndiv = vdiv / mynode.size[d];
            c[d] = vdiv - ndiv * mynode.size[d] + mynode.min[d];
            vdiv = ndiv;
        }
        c[NDIM - 1] = vdiv + mynode.min[NDIM - 1];
 
        return c;
    }
 
    inline int coordinate(unsigned idx, Direction d) const {
        return (idx / mynode.size_factor[d]) % mynode.size[d] + mynode.min[d];
    }
 
    inline Parity site_parity(unsigned idx) const {
        return coordinates(idx).parity();
    }
 
#endif
 
    CoordinateVector local_coordinates(unsigned idx) const {
        return coordinates(idx) - mynode.min;
    }
 
    /* MPI functions and variables. Define here in lattice? */
    void initialize_wait_arrays();
 
 
    /// Return the coordinates of a site, where 1st dim (x) runs fastest etc.
    /// Useful in
    ///   for (int64_t i=0; i<lattice.volume(); i++) {
    ///      auto c = lattice.global_coordinates(i);
 
    CoordinateVector global_coordinates(size_t index) const {
        CoordinateVector site;
        foralldir(dir) {
            site[dir] = index % l_size[dir];
            index /= l_size[dir];
        }
        return site;
    }
 
    int id() const {
        return l_label;
    }
 
 
    bool is_this_odd_boundary(Direction d) const;
 
    lattice_struct *block_by_factor(const CoordinateVector &blocking_factor);
 
    bool can_block_by_factor(const CoordinateVector &blocking_factor) const;
 
    void setup_blocked_lattice(const CoordinateVector &vol, int label,
                               lattice_struct &orig_lattice);
 
    void set_lattice_globals() const;
};
 
 
/**
 * @brief global vector of defined lattice pointers
 */
 
extern std::vector<lattice_struct *> defined_lattices;
 
 
/**
 * @brief main interface class to lattices.  
 * @details Lightweight class, contains only pointer to lattice_struct so
 * can be returned as such from functions.
 */
 
class Lattice {
 
  private:
    // ptr to current lattice_struct
    lattice_struct *lat_ptr = nullptr;
 
  public:
    /**
     * @internal Set the lattice_struct pointer only. Don't use this in normal code!
     */
 
    void set_lattice_pointer(lattice_struct *lat) {
        lat_ptr = lat;
    }
 
    /**
     * @internal check if the base lattice has been set
     */
 
    bool is_initialized() const {
        return lat_ptr != nullptr;
    }
 
    /**
     * @brief lattice.volume() returns lattice volume
     * Can be used inside onsites()-loops
     * @return lattice volume
     */
    inline int64_t volume() const {
        return lat_ptr->l_volume;
    }
 
    /**
     * @brief lattice.size() -> CoordinateVector  or lattice.size(d) -> int returns the
     * dimensions of the lattice, in latter case to direction d
     * Can be used inside onsites()-loops
     */
    inline int size(Direction d) const {
        return lat_ptr->l_size[d];
    }
    inline int size(int d) const {
        return lat_ptr->l_size[d];
    }
    const CoordinateVector &size() const {
        return lat_ptr->l_size;
    }
 
    /**
     * @brief lattice.setup(CoordinateVector size) - set up the base lattice. Called at the
     * beginning of the program. Can be called only once during the program run
     */
    void setup(const CoordinateVector &siz) {
        assert(lat_ptr != nullptr);
        lat_ptr->setup_base_lattice(siz);
    }
 
    /**
     * @brief lattice-> defines an operator with which detailed lattice_struct fields
     * can be accessed.
     *
     * @details For example:
     *  lattice->mynode.min: CoordinateVector of the min-corner managed by this MPI rank
     *  lattice->mynode.size: CV of the local node dimensions.
     *
     *  For further info, check the class lattice_struct.  This cannot be used
     *  within onsites(), copy the relevant field to local variable first.
     */
    inline const lattice_struct *operator->() const noexcept {
        return lat_ptr;
    }
 
    /**
     * @brief get non-const pointer to lattice_struct (cf. operator ->)
     *
     * @details This can be used to access fields of lattice_struct, this time
     * returning non-const. reference.  NOTE! you should not
     */
 
    inline lattice_struct *ptr() const {
        return lat_ptr;
    }
 
    /**
     * @brief get const ref to lattice_struct, lattice.ref(). is equivalent to lattice->
     */
    inline const lattice_struct &ref() const {
        return *lat_ptr;
    }
 
    /**
     * @brief Equality operator == is true if the two Lattices are the same, i.e. point
     * to the same lattice
     *
     * @param rhs
     */
 
    bool operator==(const Lattice &rhs) {
        return (this->lat_ptr == rhs.ptr());
    }
 
    bool operator!=(const Lattice &rhs) {
        return (this->lat_ptr != rhs.ptr());
    }
 
    /**
     * @brief With a valid lattice_struct pointer, switch this lattice to be active.
     */
    Lattice switch_to(lattice_struct *lat) {
        if (lat_ptr != lat) {
            lat_ptr = lat;
            lat->set_lattice_globals();
        }
        return *this;
    }
 
    /**
     * @brief With a valid Lattice, switch this lattice to be active.
     * @details useful e.g. in switching to lattice where field a belongs to:
     *    lattice.switch_to(a.mylattice());
     *
     */
    Lattice switch_to(Lattice &lat) {
        if (lat_ptr != lat.ptr()) {
            lat_ptr = lat.ptr();
            lat_ptr->set_lattice_globals();
        }
        return *this;
    }
 
    /**
     * @brief Test if lattice can be blocked by factor given in argument.
     * @details lattice.size() must be element-by-element divisible by factor
     *
     * Example: lattice.can_block({2,2,2}) returns true if (3-dim) lattice dimensions are even
     */
    bool can_block(const CoordinateVector &factor) const {
        return lat_ptr->can_block_by_factor(factor);
    }
 
    /**
     * @brief block the lattice by factor, switching to smaller lattice.
     * @details lattice.size() must be element-by-element divisible by factor
     *
     * @example lattice.block({2,2,2}) reduces current lattice size by 2 to each direction.
     *
     * @note Previously used Field variables cannot be used in onsites(). However,
     * their content can be blocked to new Field with Field<T>::block_from(),
     * which copies the content from the blocked (sparse) set of sites;
     *
     * @code{.cpp}
     * Field<double> a;
     * a[ALL] = X.x() + X.y() + X.z();   // in 3d
     *
     * CoordinateVector factor{2,2,2};
     *
     * lattice.block(factor);
     * Field<double> b;
     *
     * b.block_from(a);               // Now b contains 0+0+0, 2+0+0 ...
     * b[ALL] = -b[X];                // Do operations, here flip sign
     * // onsites(ALL) a[X] = 1;      // ERROR, a belongs to original lattice
     * b.unblock_to(a);               // copy content of b back to a on blocked sites,
     *                                // leaving other sites of a
     * lattice.unblock();             // return to original lattice
     *
     * // Now a can be used, it contains -(0+0+0), 1+0+0, -(2+0+0), ...
     * @endcode
     *
     * @note Fields which were not used previously can be used in blocked levels, or
     * their association (and content) can be deleted with .clear():
     *
     * @code{.cpp}
     * Field<double> a,b;
     * a = 1;
     * lattice.block({2,2,1});
     * b = 2;       // OK, b was not used before blocking
     * a = 3;       // ERROR; a belongs to non-blocked lattice
     * a.clear();
     * a = 3;       // OK, now a belongs to blocked lattice
     * @endcode
     *
     *
     * @returns lattice_struct * to blocked lattice
     */
    Lattice block(const CoordinateVector &cv) {
        lat_ptr->block_by_factor(cv);
        return *this;
    }
 
    /**
     * @brief Unblock lattice, returning to parent. Current lattice must be a blocked lattice
     * @returns lattice_struct * to new lattice
     */
    Lattice unblock() {
        if (lat_ptr->parent != nullptr) {
            return switch_to(lat_ptr->parent);
        } else {
            hila::out0 << "ERROR : cannot unblock lattice, no parent\n";
            hila::terminate(0);
            return *this;
        }
    }
 
    /**
     * @brief lattice.is_base() is used to check if the current lattice is the original one,
     * i.e. not a blocked one
     * @returns true if current lattice is the base lattice, otherwise false
     */
 
    bool is_base() const {
        return lat_ptr == defined_lattices.front();
    }
 
    /**
     * @brief lattice.switch_to_base() switches the base lattice active (if not already)
     * @returns The Lattice for the base lattice (i.e. the first lattice)
     */
 
    Lattice switch_to_base() {
        return switch_to(defined_lattices.front());
    }
};
 
 
/**
 * @brief global handle to lattice. For methods, see classes Lattice and lattice_struct
 */
extern Lattice lattice;
 
 
#ifdef VANILLA
#include "plumbing/backend_cpu/lattice.h"
#elif defined(CUDA) || defined(HIP)
#include "plumbing/backend_gpu/lattice.h"
#elif defined(VECTORIZED)
#include "plumbing/backend_vector/lattice_vector.h"
#endif
 
 
//////////////////////////////////////////////////////////////////////
// Define looping utilities
// forallcoordinates(cv)  - loops over coordinates in "natural" order
// forcoordinaterange(cv, min, max) - loops over a box subvolume in natural order
// Note - not meant for regular lattice traversal.
 
// clang-format off
#if NDIM == 4
 
#define forallcoordinates(cv) \
for (cv[3] = 0; cv[3] < lattice.size(3); cv[3]++) \
for (cv[2] = 0; cv[2] < lattice.size(2); cv[2]++) \
for (cv[1] = 0; cv[1] < lattice.size(1); cv[1]++) \
for (cv[0] = 0; cv[0] < lattice.size(0); cv[0]++) 
 
#define forcoordinaterange(cv,cmin,cmax) \
for (cv[3] = cmin[3]; cv[3] <= cmax[3]; cv[3]++) \
for (cv[2] = cmin[2]; cv[2] <= cmax[2]; cv[2]++) \
for (cv[1] = cmin[1]; cv[1] <= cmax[1]; cv[1]++) \
for (cv[0] = cmin[0]; cv[0] <= cmax[0]; cv[0]++) 
 
#elif NDIM == 3
 
#define forallcoordinates(cv) \
for (cv[2] = 0; cv[2] < lattice.size(2); cv[2]++) \
for (cv[1] = 0; cv[1] < lattice.size(1); cv[1]++) \
for (cv[0] = 0; cv[0] < lattice.size(0); cv[0]++) 
 
#define forcoordinaterange(cv,cmin,cmax) \
for (cv[2] = cmin[2]; cv[2] <= cmax[2]; cv[2]++) \
for (cv[1] = cmin[1]; cv[1] <= cmax[1]; cv[1]++) \
for (cv[0] = cmin[0]; cv[0] <= cmax[0]; cv[0]++) 
 
#elif NDIM == 2
 
#define forallcoordinates(cv) \
for (cv[1] = 0; cv[1] < lattice.size(1); cv[1]++) \
for (cv[0] = 0; cv[0] < lattice.size(0); cv[0]++) 
 
#define forcoordinaterange(cv,cmin,cmax) \
for (cv[1] = cmin[1]; cv[1] <= cmax[1]; cv[1]++) \
for (cv[0] = cmin[0]; cv[0] <= cmax[0]; cv[0]++) 
 
#elif NDIM == 1
 
#define forallcoordinates(cv) \
for (cv[0] = 0; cv[0] < lattice.size(0); cv[0]++) 
 
#define forcoordinaterange(cv,cmin,cmax) \
for (cv[0] = cmin[0]; cv[0] <= cmax[0]; cv[0]++) 
 
#endif
// clang-format on
 
#endif