Tpetra_Distribution1D.hpp

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

//  1D Row-based Distribution class
//  Assumes square matrix
//  Karen Devine, SNL
//

#ifndef __TPETRA_DISTRIBUTION1D_HPP
#define __TPETRA_DISTRIBUTION1D_HPP

namespace Tpetra {

// Forward definition
template <typename gno_t, typename scalar_t>
class DistributionLowerTriangularBlock;

template <typename gno_t, typename scalar_t>
class Distribution1D : public Distribution<gno_t, scalar_t> {
  // 1D row-wise distribution of matrix and vector entries
  // Rows and vector entries may be linearly or randomly distributed, or
  // read from a file.
  // Row map and vector map are identical

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

  Distribution1D(size_t nrows_,
                 const Teuchos::RCP<const Teuchos::Comm<int> > &comm_,
                 const Teuchos::ParameterList &params)
    : Distribution<gno_t, scalar_t>(nrows_, comm_, params) {
    int npRow = -1;  // Number of processors among which to distribute rows;
                     // Will compute if not set by user
    const Teuchos::ParameterEntry *pe = params.getEntryPtr("nProcessorRows");
    if (pe != NULL) npRow = pe->getValue<int>(&npRow);

    TEUCHOS_TEST_FOR_EXCEPTION(npRow != -1 && npRow != np, std::logic_error,
                               " nProcessorRows " << npRow << " must equal"
                                                  << " nProcessors " << np << " for 1D distribution");

    if (me == 0) std::cout << "\n 1D Distribution: "
                           << "\n     np     = " << np << std::endl;
  }

  // Return whether this rank owns vector entry i.
  virtual bool VecMine(gno_t i) = 0;

  // Return whether this rank owns nonzero (i,j)
  // Vector map and row map are the same in 1D distribution.
  inline bool Mine(gno_t i, gno_t j) { return VecMine(i); }
  inline bool Mine(gno_t i, gno_t j, int p) { return VecMine(i); }
};

template <typename gno_t, typename scalar_t>
class Distribution1DLinear : public Distribution1D<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;

  Distribution1DLinear(size_t nrows_,
                       const Teuchos::RCP<const Teuchos::Comm<int> > &comm_,
                       const Teuchos::ParameterList &params)
    : Distribution1D<gno_t, scalar_t>(nrows_, comm_, params) {
    gno_t nMyRows = getNumRow(me);
    myFirstRow    = getFirstRow(me);
    myLastRow     = myFirstRow + nMyRows - 1;
  }

  inline enum DistributionType DistType() { return OneDLinear; }

  inline bool VecMine(gno_t i) { return (i >= myFirstRow && i <= myLastRow); }

 private:
  gno_t myFirstRow;
  gno_t myLastRow;

  inline size_t getNumRow(int p) {
    return (nrows / np + (int(nrows % np) > p));
  }

  inline gno_t getFirstRow(int p) {
    return (p * (nrows / np) + std::min<int>(int(nrows % np), p));
  }

  // DistributionLowerTriangularBlock class needs a 1DLinear distribution
  friend class DistributionLowerTriangularBlock<gno_t, scalar_t>;
};

template <typename gno_t, typename scalar_t>
class Distribution1DRandom : public Distribution1D<gno_t, scalar_t> {
 public:
  using Distribution<gno_t, scalar_t>::me;

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

  inline enum DistributionType DistType() { return OneDRandom; }

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

template <typename gno_t, typename scalar_t>
class Distribution1DVec : public Distribution1D<gno_t, scalar_t> {
  // Distribution of nonzeros is determined by the distribution of the
  // vector entries, as read from a file.
  //
  // Assumptions include:
  // -  Distribution file containing the vector part assignments (N lines)
  //    is provided.  This file is read during the constructor.
  //    Format for an NxN matrix:
  //        line 1 to N:  0-based part assignment of 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 Distribution<gno_t, scalar_t>::Mine;

  Distribution1DVec(size_t nrows_,
                    const Teuchos::RCP<const Teuchos::Comm<int> > &comm_,
                    const Teuchos::ParameterList &params,
                    std::string &distributionfile)
    : Distribution1D<gno_t, scalar_t>(nrows_, comm_, params) {
    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<int, int>(*comm, 0, cnt, &(vecpart[start]));
        start += cnt;
        cnt = 0;
      }
    }

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

  ~Distribution1DVec() { delete[] vecpart; }

  inline enum DistributionType DistType() { return OneDVec; }

  // Vector distribution was read in.
  inline bool VecMine(gno_t i) { return (vecpart[i] == me); }

 protected:
  int *vecpart;  // part assignment of vector entries
};

}  // namespace Tpetra

#endif