libtiledarray-dev 0.4.4-1 / usr / include / TiledArray / proc_grid.h

/*
 *  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 {

    /// A 2D processor grid

    /// ProcGrid attempts to create a near optimal 2D grid of P processes for
    /// an MxN grid of tiles. The size of the grid is optimized such that the
    /// total communication time required for SUMMA and the number of unused
    /// processes is minimized. The total communication time of SUMMA is given
    /// by:
    /// \f[
    ///   T = \frac{MK}{P_{\rm{row}}} \left(\alpha + \frac{mk}{\beta}\right)
    ///     \left((P/P_{\rm{row}}) - 1\right) + \frac{KN}{P/P_{\rm{row}}}
    ///     \left(\alpha + \frac{kn}{\beta}\right) \left(P_{\rm{row}} - 1\right)
    /// \f]
    /// where \f$P_{\rm{row}}\f$ is the number of process rows; \f$M\f$,
    /// \f$N\f$, and \f$K\f$ are the number of tile rows and columns in a matrix
    /// product with average tiles sizes of \f$m\f$, \f$n\f$, and \f$k\f$,
    /// respectively; \f$P\f$ is the number or available processes; \f$\alpha\f$
    /// is the message latency; and \f$\beta\f$ is the message data rate. If we
    /// evaluate \f$dT/dP_{\rm{row}} = 0\f$ and assume that
    /// \f$\alpha \approx 0\f$, the expression simplifies to:
    /// \f[
    ///   Nn(2 P_{\rm{row}}^4 - P_{\rm{row}}^3) + Mm(P P_{\rm{row}} - P^2) = 0
    /// \f]
    /// where the positive, real root of \f$P_{\rm{row}}\f$ give the optimal
    /// optimal communication time.
    class ProcGrid {
    public:
      typedef uint_fast32_t size_type;

    private:
      World* world_; ///< The world where this process grid lives
      size_type rows_; ///< Number of element rows
      size_type cols_; ///< Number of element columns
      size_type size_; ///< Number of elements
      size_type proc_rows_; ///< Number of rows in the process grid
      size_type proc_cols_; ///< Number of columns in the process grid
      size_type proc_size_; ///< Number of processes in the process grid. This
                         ///<  may be less than the number of processes in world.
      ProcessID rank_row_; ///< This process's row in the process grid
      ProcessID rank_col_; ///< This process's column in the process grid
      size_type local_rows_; ///< The number of local element rows
      size_type local_cols_; ///< The number of local element columns
      size_type local_size_; ///< Number of local elements


      /// Compute the number of process rows that minimizes communication

      /// This function computes the optimal number of process row such that the
      /// communication time of a single SUMMA iteration is minimum.
      /// \param nprocs The number of processes
      /// \param Mm The number of row elements
      /// \param Nn The number of column elements
      /// \return The number of process rows that minimizes communication time
      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;
      }

      /// Search for optimal values of x and y

      /// This function will search for values of x and y such that minimize the
      /// number of unused processes, subject to the constraint that
      /// <tt>x*y <= nprocs</tt>. When the number of unused processes is equal,
      /// the solution that is closest to the initial guess for x and y will be
      /// used, which is also the solution with lower communication cost.
      /// \param[in,out] x The initial guess for the number of rows
      /// \param[in,out] y The initial guess for the number of columns
      /// \param[in] nprocs The number of available processes
      /// \param[in] min_x The minimum valid value for x
      /// \param[in] max_x The maximum valid value for x
      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;
          }
        }
      }

      /// Member variable initialization

      /// This function initializes the member variables with with the optimal
      /// sizes.
      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:
      /// Default constructor

      /// All sizes are initialized to zero.
      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)
      { }

      /// Construct a process 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.
      /// \param world The world where the process grid will live
      /// \param rows The number of tile rows
      /// \param cols The number of tile columns
      /// \param row_size The number of element rows
      /// \param col_size The number of element columns
      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.


      /// Construct a process 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.
      /// \param world The world where the process grid will live
      /// \param test_rank Test rank
      /// \param test_nprocs Test number of procs
      /// \param rows The number of tile rows
      /// \param cols The number of tile columns
      /// \param row_size The number of element rows
      /// \param col_size The number of element columns
      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

      /// Copy constructor

      // 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.
      /// \param other The other process grid to be copied
      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_)
      { }

      /// Copy assignment operator

      /// \param other The other process grid to be copied
      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;
      }

      /// Element row count accessor

      /// \return The number of element rows
      size_type rows() const { return rows_; }

      /// Element column count accessor

      /// \return The number of element columns
      size_type cols() const { return cols_; }

      /// Element count accessor

      /// \return The number of element
      size_type size() const { return size_; }

      /// Local element row count accessor

      /// \return The number of element rows
      size_type local_rows() const { return local_rows_; }

      /// Local element column count accessor

      /// \return The number of element columns
      size_type local_cols() const { return local_cols_; }

      /// Local element count accessor

      /// \return The number of element
      size_type local_size() const { return local_size_; }

      /// Rank row accessor

      /// \return The row of this process in the process grid
      ProcessID rank_row() const { return rank_row_; }

      /// Rank row accessor

      /// \return The column of this process in the process grid
      ProcessID rank_col() const { return rank_col_; }

      /// Process row count accessor

      /// \return The number of rows in the process grid
      size_type proc_rows() const { return proc_rows_; }

      /// Process column count accessor

      /// \return The number of columns in the process grid
      size_type proc_cols() const { return proc_cols_; }

      /// Process grid size accessor

      /// \return The number of processes included in the process grid (may be
      /// less than the number of process in world).
      size_type proc_size() const { return proc_size_; }


      /// Construct a row group

      /// \param did The distributed id for the result group
      /// \return A \c Group object that includes all processes in \c rank_row
      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;
      }

      /// Construct a column group

      /// \param did The distributed id for the result group
      /// \return A \c Group object that includes all processes in \c rank_col
      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;
      }

      /// Map a row to the process in this process's column

      /// \param row The row to be mapped
      /// \return The process the corresponds to the process coordinate \c (row,rank_col)
      ProcessID map_row(const size_type row) const {
        TA_ASSERT(row < proc_rows_);
        return rank_col_ + row * proc_cols_;
      }

      /// Map a column to the process in this process's row

      /// \param col The column to be mapped
      /// \return The process the corresponds to the process coordinate \c (rank_row,col)
      ProcessID map_col(const size_type col) const {
        TA_ASSERT(col < proc_cols_);
        return rank_row_ * proc_cols_ + col;
      }

      /// Construct a cyclic process

      /// Construct a cyclic process map with the same phase as the process grid.
      /// \return Cyclic process map
      std::shared_ptr<Pmap> make_pmap() const {
        TA_ASSERT(world_);

        return std::shared_ptr<Pmap>(new CyclicPmap(*world_, rows_, cols_, proc_rows_, proc_cols_));
      }

      /// Construct column phased a cyclic process

      /// Construct a cyclic process map where the column phase of the process
      /// matches that of this process grid.
      /// \param rows The number of rows in the process map
      /// \return Cyclic process map with matching column phase
      std::shared_ptr<Pmap> make_col_phase_pmap(const size_type rows) const {
        TA_ASSERT(world_);

        return std::shared_ptr<Pmap>(new CyclicPmap(*world_, rows, cols_, proc_rows_, proc_cols_));
      }

      /// Construct row phased a cyclic process

      /// Construct a cyclic process map where the column phase of the process
      /// matches that of this process grid.
      /// \param cols The number of columns in the process map
      /// \return Cyclic process map with matching column phase
      std::shared_ptr<Pmap> make_row_phase_pmap(const size_type cols) const {
        TA_ASSERT(world_);

        return std::shared_ptr<Pmap>(new CyclicPmap(*world_, rows_, cols, proc_rows_, proc_cols_));
      }
    }; // class Grid

  } // namespace detail
} // namespace TiledArray

#endif // TILEDARRAY_GRID_H__INCLUDED