ParMetisDistribution.hpp

/*
 * ParMetisDistribution.hpp
 *
 *  Created on: Nov 19, 2015
 *      Author: Antonio Leo
 */


#ifndef SRC_DECOMPOSITION_PARMETISDISTRIBUTION_HPP_
#define SRC_DECOMPOSITION_PARMETISDISTRIBUTION_HPP_


#include "SubdomainGraphNodes.hpp"
#include "parmetis_util.hpp"
#include "Graph/ids.hpp"
#include "Graph/CartesianGraphFactory.hpp"

#define PARMETIS_DISTRIBUTION_ERROR 100002

template<unsigned int dim, typename T>
class ParMetisDistribution
{
    bool is_distributed = false;

    Vcluster & v_cl;

    grid_sm<dim, void> gr;

    Box<dim, T> domain;

    Graph_CSR<nm_v, nm_e> gp;

    Parmetis<Graph_CSR<nm_v, nm_e>> parmetis_graph;

    openfpm::vector<size_t> sub_sub_owner;

    //
    // vtxdist is a common array across processor, it indicate how
    // vertex are distributed across processors
    //
    // Example we have 3 processors
    //
    // processor 0 has 3 vertices
    // processor 1 has 5 vertices
    // processor 2 has 4 vertices
    //
    // vtxdist contain, 0,3,8,12
    //
    // vtx dist is the unique global-id of the vertices
    //
    openfpm::vector<rid> vtxdist;

    openfpm::vector<openfpm::vector<idx_t>> partitions;

    openfpm::vector<openfpm::vector<gid>> v_per_proc;

    std::unordered_map<rid, gid> m2g;

    bool verticesGotWeights = false;

    void updateGraphs()
    {
        sub_sub_owner.clear();

        size_t Np = v_cl.getProcessingUnits();

        // Init n_vtxdist to gather informations about the new decomposition
        openfpm::vector<rid> n_vtxdist(Np + 1);
        for (size_t i = 0; i <= Np; i++)
            n_vtxdist.get(i).id = 0;

        // Update the main graph with received data from processor i
        for (size_t i = 0; i < Np; i++)
        {
            size_t ndata = partitions.get(i).size();
            size_t k = 0;

            // Update the main graph with the received informations
            for (rid l = vtxdist.get(i); k < ndata && l < vtxdist.get(i + 1); k++, ++l)
            {
                // Create new n_vtxdist (just count processors vertices)
                ++n_vtxdist.get(partitions.get(i).get(k) + 1);

                // vertex id from vtx to grobal id
                auto v_id = m2g.find(l)->second.id;

                // Update proc id in the vertex (using the old map)
                gp.template vertex_p<nm_v::proc_id>(v_id) = partitions.get(i).get(k);

                if (partitions.get(i).get(k) == (long int)v_cl.getProcessUnitID())
                    sub_sub_owner.add(v_id);

                // Add vertex to temporary structure of distribution (needed to update main graph)
                v_per_proc.get(partitions.get(i).get(k)).add(getVertexGlobalId(l));
            }
        }

        // Create new n_vtxdist (accumulate the counters)
        for (size_t i = 2; i <= Np; i++)
            n_vtxdist.get(i) += n_vtxdist.get(i - 1);

        // Copy the new decomposition in the main vtxdist
        for (size_t i = 0; i <= Np; i++)
            vtxdist.get(i) = n_vtxdist.get(i);

        openfpm::vector<size_t> cnt;
        cnt.resize(Np);

        for (size_t i = 0 ; i < gp.getNVertex(); ++i)
        {
            size_t pid = gp.template vertex_p<nm_v::proc_id>(i);

            rid j = rid(vtxdist.get(pid).id + cnt.get(pid));
            gid gi = gid(i);

            gp.template vertex_p<nm_v::id>(i) = j.id;
            cnt.get(pid)++;

            setMapId(j,gi);
        }
    }

    inline auto vertexByMapId(rid id) -> decltype( gp.vertex(m2g.find(id)->second.id) )
    {
        return gp.vertex(m2g.find(id)->second.id);
    }

    inline void setMapId(rid n, gid g)
    {
        m2g[n] = g;
    }

    gid getVertexGlobalId(rid n)
    {
        return m2g.find(n)->second;
    }

    void initLocalToGlobalMap()
    {
        gid g;
        rid i;
        i.id = 0;

        m2g.clear();
        for ( ; (size_t)i.id < gp.getNVertex(); ++i)
        {
            g.id = i.id;

            m2g.insert( { i, g });
        }
    }

    static void * message_receive(size_t msg_i, size_t total_msg, size_t total_p, size_t i, size_t ri, void * ptr)
    {
        openfpm::vector < openfpm::vector < idx_t >> *v = static_cast<openfpm::vector<openfpm::vector<idx_t>> *>(ptr);

        v->get(i).resize(msg_i / sizeof(idx_t));

        return &(v->get(i).get(0));
    }

    void postDecomposition()
    {
        size_t p_id = v_cl.getProcessUnitID();

        size_t Np = v_cl.getProcessingUnits();

        // Number of local vertex
        size_t nl_vertex = vtxdist.get(p_id+1).id - vtxdist.get(p_id).id;

        idx_t * partition = parmetis_graph.getPartition();

        partitions.get(p_id).resize(nl_vertex);
        std::copy(partition, partition + nl_vertex, &partitions.get(p_id).get(0));

        // Reset data structure to keep trace of new vertices distribution in processors (needed to update main graph)
        for (size_t i = 0; i < Np; ++i)
        {
            v_per_proc.get(i).clear();
        }

        // Communicate the local distribution to the other processors
        // to reconstruct individually the global graph
        openfpm::vector<size_t> prc;
        openfpm::vector<size_t> sz;
        openfpm::vector<void *> ptr;

        for (size_t i = 0; i < Np; i++)
        {
            if (i != v_cl.getProcessUnitID())
            {
                partitions.get(i).clear();
                prc.add(i);
                sz.add(nl_vertex * sizeof(idx_t));
                ptr.add(partitions.get(p_id).getPointer());
            }
        }

        if (prc.size() == 0)
            v_cl.sendrecvMultipleMessagesNBX(0, NULL, NULL, NULL, message_receive, &partitions,NONE);
        else
            v_cl.sendrecvMultipleMessagesNBX(prc.size(), &sz.get(0), &prc.get(0), &ptr.get(0), message_receive, &partitions,NONE);

        // Update graphs with the received data
        updateGraphs();
    }


public:

    ParMetisDistribution(Vcluster & v_cl)
    :is_distributed(false),v_cl(v_cl), parmetis_graph(v_cl, v_cl.getProcessingUnits()), vtxdist(v_cl.getProcessingUnits() + 1), partitions(v_cl.getProcessingUnits()), v_per_proc(v_cl.getProcessingUnits())
    {
    }

    ParMetisDistribution(const ParMetisDistribution<dim,T> & pm)
    :v_cl(pm.v_cl),parmetis_graph(v_cl, v_cl.getProcessingUnits())
    {
        this->operator=(pm);
    }

    ParMetisDistribution(ParMetisDistribution<dim,T> && pm)
    :v_cl(pm.v_cl)
    {
        this->operator=(pm);
    }

    void createCartGraph(grid_sm<dim, void> & grid, Box<dim, T> dom)
    {
        size_t bc[dim];

        for (size_t i = 0 ; i < dim ; i++)
            bc[i] = NON_PERIODIC;

        // Set grid and domain
        gr = grid;
        domain = dom;

        // Create a cartesian grid graph
        CartesianGraphFactory<dim, Graph_CSR<nm_v, nm_e>> g_factory_part;
        gp = g_factory_part.template construct<NO_EDGE, nm_v::id, T, dim - 1, 0>(gr.getSize(), domain, bc);
        initLocalToGlobalMap();

        size_t Np = v_cl.getProcessingUnits();

        size_t mod_v = gr.size() % Np;
        size_t div_v = gr.size() / Np;

        for (size_t i = 0; i <= Np; i++)
        {
            if (i < mod_v)
                vtxdist.get(i).id = (div_v + 1) * i;
            else
                vtxdist.get(i).id = (div_v) * i + mod_v;
        }

        // Init to 0.0 axis z (to fix in graphFactory)
        if (dim < 3)
        {
            for (size_t i = 0; i < gp.getNVertex(); i++)
            {
                gp.vertex(i).template get<nm_v::x>()[2] = 0.0;
            }
        }
        for (size_t i = 0; i < gp.getNVertex(); i++)
        {
            gp.vertex(i).template get<nm_v::global_id>() = i;
        }

    }

    Graph_CSR<nm_v, nm_e> & getGraph()
    {
        return gp;
    }

    void decompose()
    {
        if (is_distributed == false)
            parmetis_graph.initSubGraph(gp, vtxdist, m2g, verticesGotWeights);
        else
            parmetis_graph.reset(gp, vtxdist, m2g, verticesGotWeights);

        parmetis_graph.decompose(vtxdist);

        // update after decomposition
        postDecomposition();

        is_distributed = true;
    }

    void refine()
    {
        // Reset parmetis graph and reconstruct it
        parmetis_graph.reset(gp, vtxdist, m2g, verticesGotWeights);

        // Refine
        parmetis_graph.refine(vtxdist);

        postDecomposition();
    }

    void redecompose()
    {
        // Reset parmetis graph and reconstruct it
        parmetis_graph.reset(gp, vtxdist, m2g, verticesGotWeights);

        // Refine
        parmetis_graph.redecompose(vtxdist);

        postDecomposition();
    }

    float getUnbalance()
    {
        long t_cost = 0;

        long min, max, sum;
        float unbalance;

        t_cost = getProcessorLoad();

        min = t_cost;
        max = t_cost;
        sum = t_cost;

        v_cl.min(min);
        v_cl.max(max);
        v_cl.sum(sum);
        v_cl.execute();

        unbalance = ((float) (max - min)) / (float) (sum / v_cl.getProcessingUnits());

        return unbalance * 100;
    }

    void getSubSubDomainPosition(size_t id, T (&pos)[dim])
    {
#ifdef SE_CLASS1
        if (id >= gp.getNVertex())
            std::cerr << __FILE__ << ":" << __LINE__ << "Such vertex doesn't exist (id = " << id << ", " << "total size = " << gp.getNVertex() << ")\n";
#endif

        // Copy the geometrical informations inside the pos vector
        pos[0] = gp.vertex(id).template get<nm_v::x>()[0];
        pos[1] = gp.vertex(id).template get<nm_v::x>()[1];
        if (dim == 3)
            pos[2] = gp.vertex(id).template get<nm_v::x>()[2];
    }

    inline void setComputationCost(size_t id, size_t weight)
    {
        if (!verticesGotWeights)
            verticesGotWeights = true;

#ifdef SE_CLASS1
        if (id >= gp.getNVertex())
            std::cerr << __FILE__ << ":" << __LINE__ << "Such vertex doesn't exist (id = " << id << ", " << "total size = " << gp.getNVertex() << ")\n";
#endif

        // Update vertex in main graph
        gp.vertex(id).template get<nm_v::computation>() = weight;
    }

    bool weightsAreUsed()
    {
        return verticesGotWeights;
    }

    size_t getSubSubDomainComputationCost(size_t id)
    {
#ifdef SE_CLASS1
        if (id >= gp.getNVertex())
            std::cerr << __FILE__ << ":" << __LINE__ << "Such vertex doesn't exist (id = " << id << ", " << "total size = " << gp.getNVertex() << ")\n";
#endif

        return gp.vertex(id).template get<nm_v::computation>();
    }

    size_t getProcessorLoad()
    {
        size_t load = 0;

        // Processor id
        size_t p_id = v_cl.getProcessUnitID();


        for (rid i = vtxdist.get(p_id); i < vtxdist.get(p_id+1) ; ++i)
            load += gp.vertex(m2g.find(i)->second.id).template get<nm_v::computation>();

        //std::cout << v_cl.getProcessUnitID() << " weight " << load << " size " << sub_g.getNVertex() << "\n";
        return load;
    }

    void setMigrationCost(size_t id, size_t migration)
    {
#ifdef SE_CLASS1
        if (id >= gp.getNVertex())
            std::cerr << __FILE__ << ":" << __LINE__ << "Such vertex doesn't exist (id = " << id << ", " << "total size = " << gp.getNVertex() << ")\n";
#endif

        gp.vertex(id).template get<nm_v::migration>() = migration;
    }

    void setCommunicationCost(size_t v_id, size_t e, size_t communication)
    {
#ifdef SE_CLASS1

        size_t e_id = v_id + e;

        if (e_id >= gp.getNEdge())
            std::cerr << "Such edge doesn't exist (id = " << e_id << ", " << "total size = " << gp.getNEdge() << ")\n";
#endif

        gp.getChildEdge(v_id, e).template get<nm_e::communication>() = communication;
    }

    size_t getNSubSubDomains() const
    {
        return gp.getNVertex();
    }

    size_t getNOwnerSubSubDomains() const
    {
        return sub_sub_owner.size();
    }

    size_t getOwnerSubSubDomain(size_t id) const
    {
        return sub_sub_owner.get(id);
    }

    size_t getNSubSubDomainNeighbors(size_t id)
    {
#ifdef SE_CLASS1
        if (id >= gp.getNVertex())
            std::cerr << __FILE__ << ":" << __LINE__ << "Such vertex doesn't exist (id = " << id << ", " << "total size = " << gp.getNVertex() << ")\n";
#endif

        return gp.getNChilds(id);
    }

    void write(const std::string & file)
    {
        VTKWriter<Graph_CSR<nm_v, nm_e>, VTK_GRAPH> gv2(gp);
        gv2.write(std::to_string(v_cl.getProcessUnitID()) + "_" + file + ".vtk");
    }

    const ParMetisDistribution<dim,T> & operator=(const ParMetisDistribution<dim,T> & dist)
    {
        is_distributed = dist.is_distributed;
        gr = dist.gr;
        domain = dist.domain;
        gp = dist.gp;
        vtxdist = dist.vtxdist;
        partitions = dist.partitions;
        v_per_proc = dist.v_per_proc;
        verticesGotWeights = dist.verticesGotWeights;
        sub_sub_owner = dist.sub_sub_owner;
        m2g = dist.m2g;
        parmetis_graph = dist.parmetis_graph;

        return *this;
    }

    const ParMetisDistribution<dim,T> & operator=(ParMetisDistribution<dim,T> && dist)
    {
        is_distributed = dist.is_distributed;
        v_cl = dist.v_cl;
        gr = dist.gr;
        domain = dist.domain;
        gp.swap(dist.gp);
        vtxdist.swap(dist.vtxdist);
        partitions.swap(dist.partitions);
        v_per_proc.swap(dist.v_per_proc);
        verticesGotWeights = dist.verticesGotWeights;
        sub_sub_owner.swap(dist.sub_sub_owner);
        m2g.swap(dist.m2g);
        parmetis_graph = dist.parmetis_graph;

        return *this;
    }

    size_t get_ndec()
    {
        return parmetis_graph.get_ndec();
    }

    void setDistTol(double tol)
    {
        parmetis_graph.setDistTol(tol);
    }
};

#endif /* SRC_DECOMPOSITION_PARMETISDISTRIBUTION_HPP_ */