proc_grid.h

Go to the documentation of this file.

/*
 *  This file is a part of TiledArray.
 *  Copyright (C) 2013  Virginia Tech
 *
 *  This program is free software: you can redistribute it and/or modify
 *  it under the terms of the GNU General Public License as published by
 *  the Free Software Foundation, either version 3 of the License, or
 *  (at your option) any later version.
 *
 *  This program is distributed in the hope that it will be useful,
 *  but WITHOUT ANY WARRANTY; without even the implied warranty of
 *  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 *  GNU General Public License for more details.
 *
 *  You should have received a copy of the GNU General Public License
 *  along with this program.  If not, see <http://www.gnu.org/licenses/>.
 *
 *  Justus Calvin
 *  Department of Chemistry, Virginia Tech
 *
 *  proc_grid.h
 *  Nov 6, 2013
 *
 */
 
#ifndef TILEDARRAY_GRID_H__INCLUDED
#define TILEDARRAY_GRID_H__INCLUDED
 
#include <TiledArray/pmap/cyclic_pmap.h>
 
namespace TiledArray {
namespace detail {
 
 
class ProcGrid {
 public:
  typedef uint_fast32_t size_type;
 
 private:
  World* world_;         
  size_type rows_;       
  size_type cols_;       
  size_type size_;       
  size_type proc_rows_;  
  size_type proc_cols_;  
  size_type
      proc_size_;       
  ProcessID rank_row_;  
  ProcessID rank_col_;  
  size_type local_rows_;  
  size_type local_cols_;  
  size_type local_size_;  
 
 
  static size_type optimal_proc_row(const double nprocs, const double Mm,
                                    const double Nn) {
    // Compute the initial guess for P_row. This is the optimal guess when
    // Mm is equal to Nn, and the ideal solution.
    double P_row_estimate = std::sqrt(nprocs);
 
    // Here we want to find the positive, real root of the polynomial:
    //   Nn(2x^4 - x^3) + Mm(Px - 2P^2) = 0
    // using a Newton-Raphson algorithm.
 
    // Precompute some constants
    const double PMm = nprocs * Mm;
    const double two_P = nprocs + nprocs;
 
    const unsigned int max_it = 21u;
    unsigned int it = 0u;
    double r = 0.0;
    do {
      // Precompute P_row squared
      const double P_row2 = P_row_estimate * P_row_estimate;
      const double NnP_row2 = Nn * P_row2;
 
      // Compute the value of f(P_row_estimate) and df(P_row_estimate)
      const double f = NnP_row2 * (2.0 * P_row2 - P_row_estimate) +
                       PMm * (P_row_estimate - two_P);
      const double df = NnP_row2 * (8.0 * P_row_estimate - 3.0) + PMm;
 
      // Compute a new guess for P_row
      const double P_row_n1 = P_row_estimate - (f / df);
 
      // Compute the residual for this iteration
      r = std::abs(P_row_n1 - P_row_estimate);
 
      // Update the guess
      P_row_estimate = P_row_n1;
 
    } while ((r > 0.1) && ((++it) < max_it));
 
    return P_row_estimate + 0.5;
  }
 
 
  void minimize_unused_procs(size_type& x, size_type& y, const size_type nprocs,
                             const size_type min_x, const size_type max_x) {
    // Check for the quick exit
    size_type unused = x * y;
    if (unused == 0u) return;
 
    // Compute the range of values for x to be tested.
    const size_type delta = std::max<size_type>(1ul, std::log2(nprocs));
 
    const size_type optimal_x = x;
    size_type diff = 0ul;
    const size_type min_test_x =
        std::max<int_fast32_t>(min_x, int_fast32_t(x) - delta);
    size_type test_x = std::min(x + delta, max_x);
 
    for (; test_x >= min_test_x; --test_x) {
      const size_type test_y = nprocs / test_x;
      const size_type test_unused = nprocs - test_x * test_y;
      const size_type test_diff = std::abs(long(optimal_x) - long(test_x));
 
      if ((test_unused < unused) ||
          ((test_unused == unused) && (test_diff < diff))) {
        x = test_x;
        y = test_y;
        unused = test_unused;
        diff = test_diff;
      }
    }
  }
 
 
  void init(const size_type rank, const size_type nprocs,
            const std::size_t row_size, const std::size_t col_size) {
    // Check for the simple cases first ...
    if (nprocs == 1u) {  // Only one process
 
      // Set process grid sizes
      proc_rows_ = 1u;
      proc_cols_ = 1u;
      proc_size_ = 1u;
 
      // Set this process rank
      rank_row_ = 0;
      rank_col_ = 0;
 
      // Set local counts
      local_rows_ = rows_;
      local_cols_ = cols_;
      local_size_ = size_;
 
    } else if (size_ <= nprocs) {  // Max one tile per process
 
      // Set process grid sizes
      proc_rows_ = rows_;
      proc_cols_ = cols_;
      proc_size_ = size_;
 
      if (rank < proc_size_) {
        // Set this process rank
        rank_row_ = rank / proc_cols_;
        rank_col_ = rank % proc_cols_;
 
        // Set local counts
        local_rows_ = 1u;
        local_cols_ = 1u;
        local_size_ = 1u;
      }
 
    } else {  // The not so simple case
 
      // Compute the limits for process rows
      const size_type min_proc_rows =
          std::max<size_type>(((nprocs + cols_ - 1ul) / cols_), 1ul);
      const size_type max_proc_rows = std::min<size_type>(nprocs, rows_);
 
      // Compute optimal the number of process rows and columns in terms of
      // communication time.
      proc_rows_ = std::max<size_type>(
          min_proc_rows,
          std::min<size_type>(optimal_proc_row(nprocs, row_size, col_size),
                              max_proc_rows));
      proc_cols_ = nprocs / proc_rows_;
 
      if ((proc_rows_ > min_proc_rows) && (proc_rows_ < max_proc_rows)) {
        // Search for the values of proc_rows_ and proc_cols_ that minimizes
        // the number of unused processes in the process grid.
        minimize_unused_procs(proc_rows_, proc_cols_, nprocs, min_proc_rows,
                              max_proc_rows);
      }
 
      proc_size_ = proc_rows_ * proc_cols_;
 
      if (rank < proc_size_) {
        // Set this process rank
        rank_row_ = rank / proc_cols_;
        rank_col_ = rank % proc_cols_;
 
        // Set local counts
        local_rows_ = (rows_ / proc_rows_) +
                      (size_type(rank_row_) < (rows_ % proc_rows_) ? 1u : 0u);
        local_cols_ = (cols_ / proc_cols_) +
                      (size_type(rank_col_) < (cols_ % proc_cols_) ? 1u : 0u);
        local_size_ = local_rows_ * local_cols_;
      }
    }
  }
 
 public:
 
  ProcGrid()
      : world_(NULL),
        rows_(0u),
        cols_(0u),
        size_(0u),
        proc_rows_(0u),
        proc_cols_(0u),
        proc_size_(0u),
        rank_row_(0),
        rank_col_(0),
        local_rows_(0u),
        local_cols_(0u),
        local_size_(0u) {}
 
 
  // This constructor makes a rough estimate of the optimal process
  // dimensions. The goal is for the ratios of proc_rows/proc_cols and
  // rows/cols to be approximately equal.
  ProcGrid(World& world, const size_type rows, const size_type cols,
           const std::size_t row_size, const std::size_t col_size)
      : world_(&world),
        rows_(rows),
        cols_(cols),
        size_(rows_ * cols_),
        proc_rows_(0ul),
        proc_cols_(0ul),
        proc_size_(0ul),
        rank_row_(-1),
        rank_col_(-1),
        local_rows_(0ul),
        local_cols_(0ul),
        local_size_(0ul) {
    // Check for non-zero sizes
    TA_ASSERT(rows_ >= 1u);
    TA_ASSERT(cols_ >= 1u);
    TA_ASSERT(row_size >= 1ul);
    TA_ASSERT(col_size >= 1ul);
 
    init(world_->rank(), world_->size(), row_size, col_size);
  }
 
#ifdef TILEDARRAY_ENABLE_TEST_PROC_GRID
  // Note: The following function is here for testing purposes only. It
  // has the same functionality as the default constructor above, except the
  // rank and number of processes can be specified.
 
 
  // This constructor makes a rough estimate of the optimal process
  // dimensions. The goal is for the ratios of proc_rows/proc_cols and
  // rows/cols to be approximately equal.
  ProcGrid(World& world, const size_type test_rank, size_type test_nprocs,
           const size_type rows, const size_type cols,
           const std::size_t row_size, const std::size_t col_size)
      : world_(&world),
        rows_(rows),
        cols_(cols),
        size_(rows_ * cols_),
        proc_rows_(0u),
        proc_cols_(0u),
        proc_size_(0u),
        rank_row_(-1),
        rank_col_(-1),
        local_rows_(0u),
        local_cols_(0u),
        local_size_(0u) {
    // Check for non-zero sizes
    TA_ASSERT(rows >= 1u);
    TA_ASSERT(cols >= 1u);
    TA_ASSERT(row_size >= 1u);
    TA_ASSERT(col_size >= 1u);
    TA_ASSERT(test_rank < test_nprocs);
 
    init(test_rank, test_nprocs, row_size, col_size);
  }
#endif  // TILEDARRAY_ENABLE_TEST_PROC_GRID
 
 
  // This constructor makes a rough estimate of the optimal process
  // dimensions. The goal is for the ratios of proc_rows/proc_cols and
  // rows/cols to be approximately equal.
  ProcGrid(const ProcGrid& other)
      : world_(other.world_),
        rows_(other.rows_),
        cols_(other.cols_),
        size_(other.size_),
        proc_rows_(other.proc_rows_),
        proc_cols_(other.proc_cols_),
        proc_size_(other.proc_size_),
        rank_row_(other.rank_row_),
        rank_col_(other.rank_col_),
        local_rows_(other.local_rows_),
        local_cols_(other.local_cols_),
        local_size_(other.local_size_) {}
 
 
  ProcGrid& operator=(const ProcGrid& other) {
    world_ = other.world_;
    rows_ = other.rows_;
    cols_ = other.cols_;
    size_ = other.size_;
    proc_rows_ = other.proc_rows_;
    proc_cols_ = other.proc_cols_;
    proc_size_ = other.proc_size_;
    rank_row_ = other.rank_row_;
    rank_col_ = other.rank_col_;
    local_rows_ = other.local_rows_;
    local_cols_ = other.local_cols_;
    local_size_ = other.local_size_;
 
    return *this;
  }
 
 
  size_type rows() const { return rows_; }
 
 
  size_type cols() const { return cols_; }
 
 
  size_type size() const { return size_; }
 
 
  size_type local_rows() const { return local_rows_; }
 
 
  size_type local_cols() const { return local_cols_; }
 
 
  size_type local_size() const { return local_size_; }
 
 
  ProcessID rank_row() const { return rank_row_; }
 
 
  ProcessID rank_col() const { return rank_col_; }
 
 
  size_type proc_rows() const { return proc_rows_; }
 
 
  size_type proc_cols() const { return proc_cols_; }
 
 
  size_type proc_size() const { return proc_size_; }
 
 
  madness::Group make_row_group(const madness::DistributedID& did) const {
    TA_ASSERT(world_);
 
    madness::Group group;
 
    if (local_size_ != 0u) {
      // Construct a vector to hold the
      std::vector<ProcessID> proc_list;
      proc_list.reserve(proc_cols_);
 
      // Populate the row process list
      size_type p = rank_row_ * proc_cols_;
      const size_type row_end = p + proc_cols_;
      for (; p < row_end; ++p) proc_list.push_back(p);
 
      // Construct the group
      group = madness::Group(*world_, proc_list, did);
    }
 
    return group;
  }
 
 
  madness::Group make_col_group(const madness::DistributedID& did) const {
    TA_ASSERT(world_);
 
    madness::Group group;
 
    if (local_size_ != 0u) {
      // Generate the list of processes in rank_row
      std::vector<ProcessID> proc_list;
      proc_list.reserve(proc_rows_);
 
      // Populate the column process list
      for (size_type p = rank_col_; p < proc_size_; p += proc_cols_)
        proc_list.push_back(p);
 
      // Construct the group
      if (proc_list.size() != 0)
        group = madness::Group(*world_, proc_list, did);
    }
 
    return group;
  }
 
 
  ProcessID map_row(const size_type row) const {
    TA_ASSERT(row < proc_rows_);
    return rank_col_ + row * proc_cols_;
  }
 
 
  ProcessID map_col(const size_type col) const {
    TA_ASSERT(col < proc_cols_);
    return rank_row_ * proc_cols_ + col;
  }
 
 
  std::shared_ptr<Pmap> make_pmap() const {
    TA_ASSERT(world_);
 
    return std::make_shared<CyclicPmap>(*world_, rows_, cols_, proc_rows_,
                                        proc_cols_);
  }
 
 
  std::shared_ptr<Pmap> make_col_phase_pmap(const size_type rows) const {
    TA_ASSERT(world_);
 
    return std::make_shared<CyclicPmap>(*world_, rows, cols_, proc_rows_,
                                        proc_cols_);
  }
 
 
  std::shared_ptr<Pmap> make_row_phase_pmap(const size_type cols) const {
    TA_ASSERT(world_);
 
    return std::make_shared<CyclicPmap>(*world_, rows_, cols, proc_rows_,
                                        proc_cols_);
  }
};  // class Grid
 
}  // namespace detail
}  // namespace TiledArray
 
#endif  // TILEDARRAY_GRID_H__INCLUDED