round_robin_pmap.h

Go to the documentation of this file.

//
// Created by Karl Pierce on 7/26/20.
//
 
#ifndef TILEDARRAY_PMAP_ROUND_ROBIN_PMAP_H__INCLUDED
#define TILEDARRAY_PMAP_ROUND_ROBIN_PMAP_H__INCLUDED
 
#include <TiledArray/pmap/pmap.h>
 
namespace TiledArray {
namespace detail {
 
 
class RoundRobinPmap : public Pmap {
 protected:
  // Import Pmap protected variables
  using Pmap::procs_;  
  using Pmap::rank_;   
  using Pmap::size_;   
 
 private:
  const size_type remainder_;  
 
 public:
  typedef Pmap::size_type size_type;  
 
 
  RoundRobinPmap(World& world, size_type size)
      : Pmap(world, size), remainder_(size_ % procs_) {
    auto num_tiles_per_proc = size / procs_;
    if (remainder_ == 0 || rank_ >= remainder_)
      this->local_size_ = num_tiles_per_proc;
    else
      this->local_size_ = num_tiles_per_proc + 1;
  }
 
  virtual ~RoundRobinPmap() {}
 
 
  virtual size_type owner(const size_type tile) const {
    TA_ASSERT(tile < size_);
    return (tile % procs_);
  }
 
 
  virtual bool is_local(const size_type tile) const {
    return (tile % procs_ == rank_);
  }
 
};  // class RoundRobinPmap
 
}  // namespace detail
}  // namespace TiledArray
 
#endif  // TILEDARRAY_PMAP_ROUND_ROBIN_PMAP_H__INCLUDED