Tpetra_Distribution2D.hpp

// @HEADER
// *****************************************************************************
//          Tpetra: Templated Linear Algebra Services Package
//
// Copyright 2008 NTESS and the Tpetra contributors.
// SPDX-License-Identifier: BSD-3-Clause
// *****************************************************************************
// @HEADER

//  2D matrix distribution
//  Assumes square matrix
//  Karen Devine, SNL 
//

#ifndef __TPETRA_DISTRIBUTION2D_HPP
#define __TPETRA_DISTRIBUTION2D_HPP

namespace Tpetra
{

template <typename gno_t, typename scalar_t>
class Distribution2D : public Distribution<gno_t,scalar_t> {
// Processors will be laid out logically first down columns then 
// across rows.  For example, assume np = 8, npRows=2, npCols=4.
// Then the processors will be laid out in 2D as
//   0  2  4  6
//   1  3  5  7
//
// The matrix will be distributed using np=8 row blocks:
//   0  2  4  6
//   1  3  5  7
//   0  2  4  6
//   1  3  5  7
//   0  2  4  6
//   1  3  5  7
//   0  2  4  6
//   1  3  5  7
// 
// The vector will be distributed linearly or randomly.  The row and 
// column maps will be built to allow only row- or column-based 
// communication in the matvec.

public:
  using Distribution<gno_t,scalar_t>::me;
  using Distribution<gno_t,scalar_t>::np;
  using Distribution<gno_t,scalar_t>::comm;
  using Distribution<gno_t,scalar_t>::nrows;
  using Distribution<gno_t,scalar_t>::Mine;

  Distribution2D(size_t nrows_, 
                 const Teuchos::RCP<const Teuchos::Comm<int> > &comm_, 
                 const Teuchos::ParameterList &params) :
                 Distribution<gno_t,scalar_t>(nrows_, comm_, params),
                 npRows(-1), npCols(-1)
  {

    {
      const Teuchos::ParameterEntry *pe = params.getEntryPtr("nProcessorRows");
      if (pe != NULL)
        npRows = pe->getValue<int>(&npRows);
    }

    {
      const Teuchos::ParameterEntry *pe = params.getEntryPtr("nProcessorCols");
      if (pe != NULL)
        npCols = pe->getValue<int>(&npCols);
    }

    // Compute the processor configuration npRows * npCols

    if (npRows == -1 && npCols == -1) { // Compute both npRows and npCols 
      // First guess
      npRows = (int)(sqrt(np));   
      npCols = np / npRows;
      // Adjust npRows so that npRows * npCols == np
      while (npRows * npCols != np) {
        npRows++;
        npCols = np / npRows;
      }
    }
    else {  // User specified either npRows or npCols
      if (npRows == -1) // npCols specified; compute npRows
        npRows = np / npCols;
      else if (npCols == -1) // npRows specified; compute npCols
        npCols = np / npRows;

      if (npCols * npRows != np) {
        TEUCHOS_TEST_FOR_EXCEPTION(npRows * npCols != np, std::logic_error,
                                   "nProcessorCols " << npCols <<
                                   " * nProcessorRows " << npRows <<
                                   " = " << npCols * npRows <<
                                   " must equal nProcessors " << np <<
                                   " for 2D distribution");
      }
    }
    if (me == 0) 
      std::cout << "\n2D Distribution: " 
                << "\n    npRows = " << npRows
                << "\n    npCols = " << npCols
                << "\n    np     = " << np << std::endl;

    mypCol = this->TWODPCOL(me);
    mypRow = this->TWODPROW(me);
  }

  virtual ~Distribution2D() {};

protected:

  // Return the processor row for rank p
  inline int TWODPROW(int p) {return (p % npRows);}

  // Return the processor column for rank p
  inline int TWODPCOL(int p) {return (p / npRows);}

  // Return the rank for processor row i and processor column j
  inline int TWODPRANK(int i, int j) {return (j * npRows + (j+i) % npRows);}

  int npRows;  // Number of processor rows
  int npCols;  // Number of processor columns
  int mypRow;  // This rank's processor row
  int mypCol;  // This rank's processor column
};

template <typename gno_t, typename scalar_t>
class Distribution2DLinear : public Distribution2D<gno_t,scalar_t> {

public:
  using Distribution<gno_t,scalar_t>::me;
  using Distribution<gno_t,scalar_t>::np;
  using Distribution<gno_t,scalar_t>::nrows;
  using Distribution2D<gno_t,scalar_t>::npRows;
  using Distribution2D<gno_t,scalar_t>::npCols;
  using Distribution2D<gno_t,scalar_t>::mypRow;
  using Distribution2D<gno_t,scalar_t>::mypCol;

  Distribution2DLinear(size_t nrows_,
                       const Teuchos::RCP<const Teuchos::Comm<int> > &comm_, 
                       const Teuchos::ParameterList &params) :
                       Distribution2D<gno_t,scalar_t>(nrows_, comm_, params)
  {
    // Build vector describing distribution of vector entries across ranks
    entries.assign(np+1, nrows / np);
    int nExtraEntries = nrows % np;

    // Distribute the extra entries evenly among processors.
    // To evenly distribute them extra entries among processor rows and 
    // columns, we distribute them along diagonals of the matrix distribution.
    // For our example, assume we have seven extra values (the max possible
    // with np=8).  Then we give one extra entry each to ranks 
    // [0, 3, 4, 7, 1, 2, 5].  For fewer extra entries, we follow the same
    // order of assignment, and just stop early.
  
    for (int cnt = 0, i = 0; (cnt < nExtraEntries) && (i < npRows); i++) {
      for (int j = 0; (cnt < nExtraEntries) && (j < npCols); cnt++, j++) {
        int rankForExtra = Distribution2D<gno_t,scalar_t>::TWODPRANK(i, j); 
        entries[rankForExtra+1]++;  // Store in rankForExtra+1 to simplify 
                                    // prefix sum.
      }
    }

    // Perform prefix sum of entries.
    entries[0] = 0;
    for (int i = 1; i <= np; i++)
      entries[i] = entries[i-1] + entries[i];
    // Now entries contains the first vector entry for each rank.

    // Column map:  Easy; consecutive entries for all ranks in column.
    int firstRank = mypCol * npRows;  // First rank in my column
    myFirstCol = entries[firstRank];

    gno_t nMyCols = 0;
    for (int i = firstRank; i < firstRank + npRows; i++)
      nMyCols += entries[i+1] - entries[i];
    myLastCol = myFirstCol + nMyCols - 1;
  }

  inline enum DistributionType DistType() { return TwoDLinear; }

  bool Mine(gno_t i, gno_t j) { 
    int idx = int(float(i) * float(np) / float(nrows));
    while (i < entries[idx]) idx--;
    while (i >= entries[idx+1]) idx++;
    return ((mypRow == Distribution2D<gno_t,scalar_t>::TWODPROW(idx)) // RowMine
            && (j >= myFirstCol && j <= myLastCol));  // ColMine
  }
  inline bool Mine(gno_t i, gno_t j, int p) {return Mine(i,j);}

  inline bool VecMine(gno_t i) {
    return(i >= entries[me] && i < entries[me+1]); 
  }


private:
  std::vector<gno_t> entries; // Describes vector entries' distribution to ranks
                              // Organized like vtxdist
  gno_t myFirstCol;       // First column owned by this rank
  gno_t myLastCol;        // Last column owned by this rank
};

template <typename gno_t, typename scalar_t>
class Distribution2DRandom : public Distribution2D<gno_t,scalar_t> {

public:
  using Distribution<gno_t,scalar_t>::me;
  using Distribution2D<gno_t,scalar_t>::mypRow;
  using Distribution2D<gno_t,scalar_t>::mypCol;

  Distribution2DRandom(size_t nrows_,
                       const Teuchos::RCP<const Teuchos::Comm<int> > &comm_, 
                       const Teuchos::ParameterList &params) :
                       Distribution2D<gno_t,scalar_t>(nrows_, comm_, params)
  { if (me == 0) std::cout << "    randomize = true" << std::endl; }

  inline enum DistributionType DistType() { return TwoDRandom; }

  inline bool Mine(gno_t i, gno_t j) { 
    return ((mypRow == this->TWODPROW(this->HashToProc(i))) &&  // RowMine
            (mypCol == this->TWODPCOL(this->HashToProc(j))));   // ColMine
  }
  inline bool Mine(gno_t i, gno_t j, int p) {return Mine(i,j);}

  inline bool VecMine(gno_t i) { return (me == this->HashToProc(i)); }

};


template <typename gno_t, typename scalar_t>
class Distribution2DVec : public Distribution2D<gno_t,scalar_t>
{
// Distribute non-zeros in a 2D manner based on the vector distribution
// and the nprows x npcols configuration;
// rows are assigned to same process owning the vector entry.
public:
  using Distribution<gno_t,scalar_t>::me;
  using Distribution<gno_t,scalar_t>::np;
  using Distribution<gno_t,scalar_t>::comm;
  using Distribution<gno_t,scalar_t>::nrows;
  using Distribution2D<gno_t,scalar_t>::npRows;
  using Distribution2D<gno_t,scalar_t>::npCols;

  Distribution2DVec(size_t nrows_, 
                    const Teuchos::RCP<const Teuchos::Comm<int> > &comm_, 
                    const Teuchos::ParameterList &params,
                    std::string &distributionfile) :
                    Distribution2D<gno_t,scalar_t>(nrows_, comm_, params)
  {
    if (me == 0) std::cout << "\n 2DVec Distribution: "
                           << "\n      np     = " << np << std::endl;
    std::ifstream fpin;
    if (me == 0) {
      fpin.open(distributionfile.c_str(), std::ios::in);
      if (!fpin.is_open()) {
        std::cout << "Error:  distributionfile " << distributionfile
             << " not found" << std::endl;
        exit(-1);
      }
    }

    // Read the vector part assignment and broadcast it to all processes.
    // Broadcast in chunks of bcastsize values.
    // TODO:  Make the vector part assignment more scalable instead of 
    // TODO:  storing entire vector on every process.

    vecpart = new int[nrows];

    const int bcastsize = 1000000;

    gno_t start = 0;
    int cnt = 0;
    for (size_t i = 0; i < nrows; i++) {
      if (me == 0) fpin >> vecpart[i];
      cnt++;
      if (cnt == bcastsize || i == nrows-1) {
        Teuchos::broadcast(*comm, 0, cnt, &(vecpart[start]));
        start += cnt;
        cnt = 0;
      }
    }

    if (me == 0) fpin.close();
  }

  ~Distribution2DVec() {delete [] vecpart;}

  inline enum DistributionType DistType() { return TwoDVec; }

  bool Mine(gno_t i, gno_t j) {
    return (me == (vecpart[i] % npRows + (vecpart[j] / npRows) * npRows));
  }
  inline bool Mine(gno_t i, gno_t j, int p) {return Mine(i,j);}

  inline bool VecMine(gno_t i) { return(vecpart[i] == me); }

private:
  int *vecpart;

};

}
#endif