assigner.hpp

#ifndef DIY_ASSIGNER_HPP
#define DIY_ASSIGNER_HPP

#include <vector>

namespace diy
{
  // Derived types should define
  //   int rank(int gid) const
  // that converts a global block id to a rank that it's assigned to.
  class Assigner
  {
    public:
                    Assigner(int size__,     
                             int nblocks__   
                             ):
                      size_(size__), nblocks_(nblocks__)  {}

      int           size() const                        { return size_; }
      int           nblocks() const                     { return nblocks_; }
      void          set_nblocks(int nblocks__)          { nblocks_ = nblocks__; }
      virtual void  local_gids(int rank, std::vector<int>& gids) const   =0;
      virtual int   rank(int gid) const     =0;

    private:
      int           size_;      // total number of ranks
      int           nblocks_;   // total number of blocks
  };

  class ContiguousAssigner: public Assigner
  {
    public:
            ContiguousAssigner(int size__,   
                               int nblocks__ 
                               ):
              Assigner(size__, nblocks__)      {}

      using Assigner::size;
      using Assigner::nblocks;

      int   rank(int gid) const override
      {
          int div = nblocks() / size();
          int mod = nblocks() % size();
          int r = gid / (div + 1);
          if (r < mod)
          {
              return r;
          } else
          {
              return mod + (gid - (div + 1)*mod)/div;
          }
      }
      inline
      void  local_gids(int rank, std::vector<int>& gids) const override;
  };

  class RoundRobinAssigner: public Assigner
  {
    public:
            RoundRobinAssigner(int size__,   
                               int nblocks__ 
                               ):
              Assigner(size__, nblocks__)         {}

      using Assigner::size;
      using Assigner::nblocks;

      int   rank(int gid) const override        { return gid % size(); }
      inline
      void  local_gids(int rank, std::vector<int>& gids) const override;
  };
}

void
diy::ContiguousAssigner::
local_gids(int rank_, std::vector<int>& gids) const
{
  int div = nblocks() / size();
  int mod = nblocks() % size();

  int from, to;
  if (rank_ < mod)
      from = rank_ * (div + 1);
  else
      from = mod * (div + 1) + (rank_ - mod) * div;

  if (rank_ + 1 < mod)
      to = (rank_ + 1) * (div + 1);
  else
      to = mod * (div + 1) + (rank_ + 1 - mod) * div;

  for (int gid = from; gid < to; ++gid)
    gids.push_back(gid);
}

void
diy::RoundRobinAssigner::
local_gids(int rank_, std::vector<int>& gids) const
{
  int cur = rank_;
  while (cur < nblocks())
  {
    gids.push_back(cur);
    cur += size();
  }
}

#endif