TeamSpec.h

#ifndef DASH__TEAM_SPEC_H_
#define DASH__TEAM_SPEC_H_

#include <dash/Types.h>
#include <dash/Dimensional.h>
#include <dash/Exception.h>
#include <dash/Cartesian.h>
#include <dash/internal/Logging.h>

#include <array>
#include <set>
#include <algorithm>
#include <sstream>
#include <iostream>
#include <cstring>
#include <type_traits>
#include <initializer_list>

namespace dash {

template<
  dim_t    MaxDimensions,
  typename IndexType = dash::default_index_t>
class TeamSpec :
  public CartesianIndexSpace<MaxDimensions, ROW_MAJOR, IndexType>
{
private:
  typedef typename std::make_unsigned<IndexType>::type
    SizeType;
  typedef TeamSpec<MaxDimensions, IndexType>
    self_t;
  typedef CartesianIndexSpace<MaxDimensions, ROW_MAJOR, IndexType>
    parent_t;

public:
  TeamSpec(
    Team & team = dash::Team::All())
  : _is_linear(true),
    _myid(team.myid())
  {
    DASH_LOG_TRACE_VAR("TeamSpec(t)", team.is_null());
    auto team_size = team.is_null() ? 0 : team.size();
    _rank = 1;
    this->_extents[0] = team_size;
    for (auto d = 1; d < MaxDimensions; ++d) {
      this->_extents[d] = 1;
    }
    this->resize(this->_extents);
  }

  TeamSpec(
    const self_t & other,
    const DistributionSpec<MaxDimensions> & distribution,
    Team & team = dash::Team::All())
  : CartesianIndexSpace<MaxDimensions, ROW_MAJOR, IndexType>(
      other.extents()),
      _myid(team.myid())
  {
    DASH_LOG_TRACE_VAR("TeamSpec(ts, dist, t)", team.is_null());
#if 0
    if (this->size() != team.size()) {
      DASH_THROW(
        dash::exception::InvalidArgument,
        "Size of team "     << team.size()  << " differs from " <<
        "size of teamspec " << this->size() << " in TeamSpec()");
    }
#endif
    // Test if other teamspec has been default-constructed and has
    // to be rearranged for a distribution with higher rank:
    if (other._is_linear && distribution.rank() > 1) {
      // Set extent of teamspec in the dimension the distribution is
      // different from NONE:
      if (distribution.is_tiled()) {
        bool major_tiled_dim_set = false;
        for (auto d = 0; d < MaxDimensions; ++d) {
          this->_extents[d] = 1;
          if (!major_tiled_dim_set &&
              distribution[d].type == dash::internal::DIST_TILE) {
            this->_extents[d] = team.size();
            major_tiled_dim_set = true;
          }
        }
      } else {
        for (auto d = 0; d < MaxDimensions; ++d) {
          if (distribution[d].type == dash::internal::DIST_NONE) {
            this->_extents[d] = 1;
          } else {
            // Use size of given team; possibly different from size
            // of default-constructed team spec:
            this->_extents[d] = team.size();
          }
        }
      }
    }
    update_rank();
    DASH_LOG_TRACE_VAR("TeamSpec(ts, dist, t)", this->_extents);
    this->resize(this->_extents);
    DASH_LOG_TRACE_VAR("TeamSpec(ts, dist, t)", this->size());
  }

  TeamSpec(
    const DistributionSpec<MaxDimensions> & distribution,
    Team & team = dash::Team::All())
  : _myid(team.myid())
  {
    DASH_LOG_TRACE_VAR("TeamSpec(dist, t)", team.is_null());
    bool distrib_dim_set = false;
    if (distribution.is_tiled()) {
      bool major_tiled_dim_set = false;
      for (auto d = 0; d < MaxDimensions; ++d) {
        this->_extents[d] = 1;
        if (!major_tiled_dim_set &&
            distribution[d].type == dash::internal::DIST_TILE) {
          this->_extents[d] = team.size();
          major_tiled_dim_set = true;
        }
      }
    } else {
      for (auto d = 0; d < MaxDimensions; ++d) {
        if (distribution[d].type == dash::internal::DIST_NONE) {
          this->_extents[d] = 1;
        } else {
          this->_extents[d] = team.size();
          if (distrib_dim_set) {
            DASH_THROW(
              dash::exception::InvalidArgument,
              "TeamSpec(DistributionSpec, Team) only allows "
              "one distributed dimension");
          }
          distrib_dim_set = true;
        }
      }
    }
    update_rank();
    DASH_LOG_TRACE_VAR("TeamSpec(dist, t)", this->_extents);
    this->resize(this->_extents);
    DASH_LOG_TRACE_VAR("TeamSpec(dist, t)", this->size());
  }

  template<typename ... Types>
  TeamSpec(SizeType value, Types ... values)
  : CartesianIndexSpace<MaxDimensions, ROW_MAJOR, IndexType>::
      CartesianIndexSpace(value, values...),
      _myid(dash::Team::All().myid())
  {
    update_rank();
    this->resize(this->_extents);
    DASH_LOG_TRACE_VAR("TeamSpec(size,...)", this->_extents);
  }

  TeamSpec(const std::array<SizeType, MaxDimensions> & extents)
  : CartesianIndexSpace<MaxDimensions, ROW_MAJOR, IndexType>::
      CartesianIndexSpace(extents),
      _myid(dash::Team::All().myid())
  {
    update_rank();
    this->resize(this->_extents);
    DASH_LOG_TRACE_VAR("TeamSpec({extents})", this->_extents);
  }

  void balance_extents()
  {
    DASH_LOG_TRACE_VAR("TeamSpec.balance_extents()", this->_extents);
    DASH_LOG_TRACE_VAR("TeamSpec.balance_extents()", this->size());
    if(MaxDimensions <= 1) {
      return;
    }

    std::multiset<SizeType> new_extents;

    SizeType num_units = 1;
    for (auto d = 0; d < MaxDimensions; ++d) {
      num_units *= this->_extents[d];
      this->_extents[d] = 1;
      new_extents.insert(1);
    }
    _is_linear = false;

    // Find best surface-to-volume:
    auto teamsize_prime_factors = dash::math::factorize(num_units);
    // Equally distribute factors to extents.
    // Start with the largest factors and multiply them onto the lowest
    // value.
    for (auto it = teamsize_prime_factors.rbegin();
         it != teamsize_prime_factors.rend();
         ++it) {
      DASH_LOG_TRACE("TeamSpec.balance_extents()",
                     "factor:", it->first, "x", it->second);
      for (auto i = 1; i < it->second + 1; ++i) {
        new_extents.insert(it->first * *new_extents.begin());
        new_extents.erase(new_extents.begin());
      }
    }

    int d = 0;
    for (auto it = new_extents.rbegin();
         it != new_extents.rend();
         ++it, ++d) {
      this->_extents[d] = *it;
    }

    this->resize(this->_extents);
    update_rank();
    DASH_LOG_TRACE_VAR("TeamSpec.balance_extents() ->", this->_extents);
  }

  team_unit_t neighbor(std::initializer_list<int> offsets) const
  {
    auto neighbor_coords = this->coords(_myid);
    dim_t d = 0;
    for (auto offset_d : offsets) {
      neighbor_coords[d] += offset_d;
      if (neighbor_coords[d] < 0 ||
          neighbor_coords[d] >= this->_extents[d]) {
        return UNDEFINED_TEAM_UNIT_ID;
      }
      ++d;
    }
    return team_unit_t(static_cast<dart_unit_t>(this->at(neighbor_coords)));
  }

  team_unit_t neighbor(const std::array<int, MaxDimensions>&  offsets) const
  {
    auto neighbor_coords = this->coords(_myid);
    for (dim_t d = 0; d < MaxDimensions; ++d) {
      neighbor_coords[d] += offsets[d];
      if (neighbor_coords[d] < 0 ||
          neighbor_coords[d] >= this->_extents[d]) {
        return UNDEFINED_TEAM_UNIT_ID;
      }
    }
    return team_unit_t(static_cast<dart_unit_t>(this->at(neighbor_coords)));
  }

  team_unit_t periodic_neighbor(std::initializer_list<int> offsets) const
  {
    auto neighbor_coords = this->coords(_myid);
    dim_t d = 0;
    for (auto offset_d : offsets) {
      neighbor_coords[d] += offset_d;
      if (neighbor_coords[d] < 0) {
        neighbor_coords[d] = this->_extents[d] + offset_d;
      } else if(neighbor_coords[d] >= this->_extents[d]) {
        neighbor_coords[d] %= this->_extents[d];
      }
      ++d;
    }
    return team_unit_t(static_cast<dart_unit_t>(this->at(neighbor_coords)));
  }

  team_unit_t periodic_neighbor(const std::array<int, MaxDimensions>& offsets) const
  {
    auto neighbor_coords = this->coords(_myid);
    
    for (dim_t d = 0; d < MaxDimensions; ++d) {
      neighbor_coords[d] += offsets[d];
      if (neighbor_coords[d] < 0) {
        neighbor_coords[d] = this->_extents[d] + offsets[d];
      } else if(neighbor_coords[d] >= this->_extents[d]) {
        neighbor_coords[d] %= this->_extents[d];
      }
    }
    return team_unit_t(static_cast<dart_unit_t>(this->at(neighbor_coords)));
  }

  bool includes_index(
    IndexType index,
    dim_t dimension,
    IndexType dim_offset) const
  {
    if (_rank == 1) {
      // Shortcut for trivial case
      return (index >= 0 && index < this->size());
    }
    return parent_t::includes_index(index, dimension, dim_offset);
  }

  SizeType num_units(dim_t dimension) const
  {
    return parent_t::extent(dimension);
  }

  template<typename... Args>
  void resize(SizeType arg, Args... args)
  {
    static_assert(
      sizeof...(Args) == MaxDimensions-1,
      "Invalid number of arguments");
    std::array<SizeType, MaxDimensions> extents =
      {{ arg, (SizeType)(args)... }};
    resize(extents);
  }

  template<typename SizeType_>
  void resize(const std::array<SizeType_, MaxDimensions> & extents)
  {
    _is_linear = false;
    parent_t::resize(extents);
    update_rank();
  }

  void resize_dim(dim_t dim, SizeType extent)
  {
    this->_extents[dim] = extent;
    resize(this->_extents);
  }

  dim_t rank() const
  {
    return _rank;
  }

private:
  void update_rank()
  {
    _rank = 0;
    for (auto d = 0; d < MaxDimensions; ++d) {
      if (this->_extents[d] > 1) {
        ++_rank;
      }
    }
    if (_rank == 0) _rank = 1;
  }

protected:
  dim_t       _rank       = 0;
  bool        _is_linear  = false;
  team_unit_t _myid;

}; // class TeamSpec

} // namespace dash

#endif // DASH__TEAM_SPEC_H_