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>
#include <TiledArray/math/eigen.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