parti_zoltan.cpp

// FEAT3: Finite Element Analysis Toolbox, Version 3
// Copyright (C) 2010 by Stefan Turek & the FEAT group
// FEAT3 is released under the GNU General Public License version 3,
// see the file 'copyright.txt' in the top level directory for details.
 
#include <kernel/geometry/parti_zoltan.hpp>
 
#ifdef FEAT_HAVE_ZOLTAN
#include <zoltan.h>
 
using namespace FEAT;
using namespace Geometry;
 
extern "C" int feat_zoltan_num_elems(void* data, int* ierr)
{
  *ierr = ZOLTAN_OK;
  const Adjacency::Graph& grp = reinterpret_cast<const PartiZoltan::Hypergraph*>(data)->sub_graph;
  return int(grp.get_num_nodes_domain());
}
 
extern "C" void feat_zoltan_elem_list(void* data, int size_gid, int size_lid,
  ZOLTAN_ID_PTR global_id, ZOLTAN_ID_PTR DOXY(local_id), int wgt_dim, float* obj_wgts, int* ierr)
{
  XASSERT(size_gid == 1);
  XASSERT(size_lid == 0);
  //XASSERT(wgt_dim == 0);
  const PartiZoltan::Hypergraph& hyg = *reinterpret_cast<const PartiZoltan::Hypergraph*>(data);
  const Index n = hyg.sub_graph.get_num_nodes_domain();
  for(Index i(0); i < n; ++i)
    global_id[i] = ZOLTAN_ID_TYPE(hyg.first_elem + i);
  if(wgt_dim > 0)
  {
    for(Index i(0); i < n; ++i)
      obj_wgts[i] = hyg.weights[i];
  }
  *ierr = ZOLTAN_OK;
}
 
extern "C" void feat_zoltan_hypergraph_size(void* data, int* num_lists, int* num_nonzeroes, int* format, int* ierr)
{
  *ierr = ZOLTAN_OK;
  const Adjacency::Graph& grp = reinterpret_cast<const PartiZoltan::Hypergraph*>(data)->sub_graph;
  *num_lists = int(grp.get_num_nodes_domain());
  *num_nonzeroes = int(grp.get_num_indices());
  *format = ZOLTAN_COMPRESSED_EDGE;
  *ierr = ZOLTAN_OK;
}
 
extern "C" void feat_zoltan_hypergraph_data(void* data, int size_gid, int num_edges, int num_nonzeroes,
  int format, ZOLTAN_ID_PTR edgeGID, int* vtxPtr, ZOLTAN_ID_PTR vtxGID, int* ierr)
{
  const Adjacency::Graph& grp = reinterpret_cast<const PartiZoltan::Hypergraph*>(data)->sub_graph;
  XASSERT(size_gid == 1);
  XASSERT(num_edges == int(grp.get_num_nodes_domain()));
  XASSERT(num_nonzeroes == int(grp.get_num_indices()));
  XASSERT(format == ZOLTAN_COMPRESSED_EDGE);
 
  const Index* dom_ptr = grp.get_domain_ptr();
  const Index* img_idx = grp.get_image_idx();
 
  // copy domain ptr = hyperedge ptr
  for(int i(0); i < num_edges; ++i)
  {
    edgeGID[i] = ZOLTAN_ID_TYPE(i);
    vtxPtr[i] = int(dom_ptr[i]);
  }
 
  // copy image idx = vertex idx
  for(int i(0); i < num_nonzeroes; ++i)
  {
    vtxGID[i] = ZOLTAN_ID_TYPE(img_idx[i]);
  }
  *ierr = ZOLTAN_OK;
}
 
namespace FEAT
{
  namespace Geometry
  {
    PartiZoltan::PartiZoltan(const Dist::Comm& comm, Index min_elems, int max_procs) :
      _comm(comm),
      _zoltan_comm(Dist::Comm::null()),
      _max_procs(max_procs),
      _min_elems(min_elems),
      _num_elems(0u),
      _num_parts(0u),
      _hypergraph(),
      _coloring(),
      _zz(nullptr)
    {
    }
 
    PartiZoltan::~PartiZoltan()
    {
      if(_zz != nullptr)
      {
        Zoltan_Destroy(reinterpret_cast<Zoltan_Struct**>(&_zz));
        _zz = nullptr;
      }
    }
 
    bool PartiZoltan::execute(const Adjacency::Graph& faces_at_elem, const Index num_parts, const std::vector<Real>& weights)
    {
      // set number of desired partitions and total number of elements
      this->_num_parts = num_parts;
      this->_num_elems = faces_at_elem.get_num_nodes_domain();
 
      // first of all, determine how many processes we want to use for Zoltan
      int num_procs = Math::max(1, int(faces_at_elem.get_num_nodes_domain() / _min_elems));
      if(_max_procs > 0)
        Math::mini(num_procs, _max_procs);
      Math::mini(num_procs, int(num_parts));
      Math::mini(num_procs, _comm.size());
 
      // do we need to create a sub-communicator for Zoltan?
      if(num_procs < _comm.size())
        _zoltan_comm = _comm.comm_create_range_incl(num_procs);
      else
        _zoltan_comm = _comm.comm_dup();
 
      // get my rank and size
      //const int z_rank = _zoltan_comm.rank();
      const int z_size = _zoltan_comm.size();
 
      // make sure that the number of desired parts is not smaller than the size of the communicator
      XASSERTM(Index(z_size) <= num_parts, "thou shall not create less partitions than thou hast ranks");
 
      // if the size of the Zoltan communicator is zero, then this process does not participate
      // in the actual partitioning process and it only needs to participate in the broadcast
      bool is_ok = true;
      if(z_size > 0)
      {
        // create our sub-hypergraph
        this->_create_hypergraph(faces_at_elem, weights);
 
        // apply Zoltan partitioner
        is_ok = this->_apply_zoltan();
      }
 
      // broadcast coloring from root rank and return
      return this->_broadcast_coloring(is_ok);
    }
 
    void PartiZoltan::_create_hypergraph(const Adjacency::Graph& faces_at_elem, const std::vector<Real>& weights)
    {
      // get number of elements and faces
      const Index num_elems = faces_at_elem.get_num_nodes_domain();
      const Index num_faces = faces_at_elem.get_num_nodes_image();
 
      // get my rank and size
      const int z_rank = _zoltan_comm.rank();
      const int z_size = _zoltan_comm.size();
 
      // this should not happen
      XASSERT(z_size > 0);
 
      // determine the first and the last element of this rank
      Index first_elem = (Index(z_rank) * num_elems) / Index(z_size);
      Index end_elem = (Index(z_rank+1) * num_elems) / Index(z_size);
      Index my_num_elems = end_elem - first_elem;
 
      const Index* dom_ptr = faces_at_elem.get_domain_ptr();
      const Index* img_idx = faces_at_elem.get_image_idx();
 
      // create a sub-graph for this rank
      const Index first_ptr = dom_ptr[first_elem];
      const Index my_num_idx = dom_ptr[end_elem] - first_ptr;
      Adjacency::Graph faces_at_my_elem(my_num_elems, num_faces, my_num_idx);
 
      Index* my_dom_ptr = faces_at_my_elem.get_domain_ptr();
      Index* my_img_idx = faces_at_my_elem.get_image_idx();
      for(Index i(0); i <= my_num_elems; ++i)
        my_dom_ptr[i] = dom_ptr[first_elem + i] - first_ptr;
      for(Index i(0); i < my_num_idx; ++i)
        my_img_idx[i] = img_idx[first_ptr + i];
 
      // transpose graph for this rank
      Adjacency::Graph elems_at_faces(Adjacency::RenderType::transpose, faces_at_elem);
 
      // combine graph into our sub-hypergraph
      this->_hypergraph.first_elem = first_elem;
      this->_hypergraph.sub_graph = Adjacency::Graph(Adjacency::RenderType::injectify_sorted, faces_at_my_elem, elems_at_faces);
 
      // extract weights if desired
      if(!weights.empty())
      {
        XASSERT(end_elem <= weights.size());
        this->_hypergraph.weights.resize(my_num_elems);
        for(Index i(0); i < my_num_elems; ++i)
          this->_hypergraph.weights[i] = float(weights[first_elem + i]);
      }
    }
 
    bool PartiZoltan::_apply_zoltan()
    {
      // create Zoltan structure
      Zoltan_Struct* zz = Zoltan_Create(_zoltan_comm.mpi_comm());
      this->_zz = zz;
 
      // define Zoltan parameters
      Zoltan_Set_Param(zz, "DEBUG_LEVEL", "0");
      Zoltan_Set_Param(zz, "LB_METHOD", "HYPERGRAPH");
      Zoltan_Set_Param(zz, "HYPERGRAPH_PACKAGE", "PHG");
      Zoltan_Set_Param(zz, "LB_APPROACH", "PARTITION");
      Zoltan_Set_Param(zz, "NUM_GID_ENTRIES", "1");
      Zoltan_Set_Param(zz, "NUM_LID_ENTRIES", "0");
      Zoltan_Set_Param(zz, "RETURN_LISTS", "PARTS");
      Zoltan_Set_Param(zz, "NUM_GLOBAL_PARTS", stringify(this->_num_parts).c_str());
      Zoltan_Set_Param(zz, "OBJ_WEIGHT_DIM", this->_hypergraph.weights.empty() ? "0" : "1");
      Zoltan_Set_Param(zz, "EDGE_WEIGHT_DIM", "0");
      //Zoltan_Set_Param(zz, "CHECK_GRAPH", "2");
      Zoltan_Set_Param(zz, "PHG_EDGE_SIZE_THRESHOLD", "1.0"); // Zoltan keeps complaining for values < 1
 
      // set hypergraph interface callbacks
      Zoltan_Set_Num_Obj_Fn(zz, feat_zoltan_num_elems, &this->_hypergraph);
      Zoltan_Set_Obj_List_Fn(zz, feat_zoltan_elem_list, &this->_hypergraph);
      Zoltan_Set_HG_Size_CS_Fn(zz, feat_zoltan_hypergraph_size, &this->_hypergraph);
      Zoltan_Set_HG_CS_Fn(zz, feat_zoltan_hypergraph_data, &this->_hypergraph);
 
      int changes(0), num_gid_entries(0), num_lid_entries(0), num_import(0), num_export(0);
      ZOLTAN_ID_PTR import_global_ids(nullptr);
      ZOLTAN_ID_PTR import_local_ids(nullptr);
      int* import_procs(nullptr);
      int* import_parts(nullptr);
      ZOLTAN_ID_PTR export_global_ids(nullptr);
      ZOLTAN_ID_PTR export_local_ids(nullptr);
      int* export_procs(nullptr);
      int* export_parts(nullptr);
 
      int rtn = Zoltan_LB_Partition(zz, &changes, &num_gid_entries, &num_lid_entries,
        &num_import, &import_global_ids, &import_local_ids, &import_procs, &import_parts,
        &num_export, &export_global_ids, &export_local_ids, &export_procs, &export_parts);
 
      // create coloring from Zoltan partitioning
      bool gather_ok = true;
      if(rtn == ZOLTAN_OK)
        gather_ok = this->_gather_coloring(num_export, export_parts);
 
      Zoltan_LB_Free_Part(&import_global_ids, &import_local_ids, &import_procs, &import_parts);
      Zoltan_LB_Free_Part(&export_global_ids, &export_local_ids, &export_procs, &export_parts);
 
      Zoltan_Destroy(&zz);
      this->_zz = nullptr;
 
      return gather_ok && (rtn == ZOLTAN_OK);
    }
 
    bool PartiZoltan::_gather_coloring(const int num_export, const int* export_parts)
    {
      // get my rank and size
      const Index z_rank = Index(_zoltan_comm.rank());
      const Index z_size = Index(_zoltan_comm.size());
 
      // determine maximum export size
      Index max_export = Index(num_export);
      _zoltan_comm.allreduce(&max_export, &max_export, 1u, Dist::op_max);
      max_export += 1; // for actual size
 
      // allocate send buffer
      std::vector<int> send_buf(max_export);
      send_buf[0] = num_export;
      for(Index i(0); i < Index(num_export); ++i)
        send_buf[i+1] = export_parts[i];
 
      // allocate receive buffer and gather coloring at rank 0
      std::vector<int> recv_buf(z_rank == 0 ? max_export * z_size : Index(0));
      _zoltan_comm.gather(send_buf.data(), max_export, recv_buf.data(), max_export, 0);
 
      // all ranks > 0 are done here
      if(z_rank > 0)
        return true;
 
      // compute total element count
      Index num_elems(0);
      for(Index i(0); i < z_size; ++i)
        num_elems += Index(recv_buf[i*max_export]);
 
      // allocate ranks vector
      this->_coloring = Adjacency::Coloring(num_elems, this->_num_parts);
      Index* col = this->_coloring.get_coloring();
 
      // build ranks vector
      for(Index i(0), k(0); i < z_size; ++i)
      {
        Index off = i*max_export;
        Index n = Index(recv_buf[off]);
        for(Index j(1); j <= n; ++j, ++k)
          col[k] = Index(recv_buf[off+j]);
      }
 
      // build color counts
      std::vector<int> color_counts(this->_num_parts, 0);
      for(Index i(0); i < this->_num_elems; ++i)
        ++color_counts[col[i]];
 
      // make sure we have at least 1 element per color/partition
      for(Index i(0); i < this->_num_parts; ++i)
      {
        if(color_counts[i] == 0)
          return false;
      }
 
      // okay
      return true;
    }
 
    bool PartiZoltan::_broadcast_coloring(bool zoltan_ok)
    {
      // first of all, let's compare our Zoltan return codes;
      // all processes which did not participate in the partitioning have a 'true' status
      int was_ok(zoltan_ok ? 0 : 1),  was_all_ok(0);
      this->_comm.allreduce(&was_ok, &was_all_ok, std::size_t(1), Dist::op_sum);
      if(was_all_ok > 0)
        return false; // at least 1 process failed
 
      // allocate coloring if not on root
      if(this->_comm.rank() > 0)
        this->_coloring = Adjacency::Coloring(this->_num_elems, this->_num_parts);
 
      // broadcast coloring
      this->_comm.bcast(this->_coloring.get_coloring(), this->_num_elems, 0);
 
      // okay
      return true;
    }
  } // namespace Geometry
} // namespace FEAT
 
#else // no FEAT_HAVE_ZOLTAN
 
// dummy function to suppress linker warnings
void zoltan_did_not_make_it_to_the_feat_parti() {}
 
#endif // FEAT_HAVE_ZOLTAN