Tpetra_DistributionLowerTriangularBlock.hpp

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

#ifndef __TPETRA_DISTRIBUTORLOWERTRIANGULARBLOCK_HPP
#define __TPETRA_DISTRIBUTORLOWERTRIANGULARBLOCK_HPP

// Needed by DistributionLowerTriangularBlock
#include "Tpetra_Distributor.hpp"

// Needed by LowerTriangularBlock operator
#include "Tpetra_Core.hpp"
#include "Tpetra_Map.hpp"
#include "Tpetra_Operator.hpp"
#include "Tpetra_Vector.hpp"
#include "Tpetra_CrsMatrix.hpp"

namespace Tpetra {

template <typename gno_t, typename scalar_t>
class DistributionLowerTriangularBlock : public Distribution<gno_t, scalar_t> {
  // Seher Acer's lower-triangular block decomposition for triangle counting
  // See also:  LowerTriangularBlockOperator below that allows this distribution
  // to be used in Tpetra SpMV.
  //
  // Requirements:
  //    Matrix must be square (undirected graph)
  //    Number of processors np = q(q+1)/2 for some q.
  //
  // Only the lower triangular portion of the matrix is stored.
  // Processors are arranged logically as follows:
  //    0
  //    1  2
  //    3  4  5
  //    ...
  //
  // The lower triangular part of the matrix is stored in a 2D distribution.
  // For example, the dense 7x7 lower triangular matrix below would be assigned
  // to processors according to numbers shown as follows:
  //    0 |   |
  //    00|   |
  //    ---------
  //    11|2  |
  //    11|22 |
  //    11|222|
  //    ---------
  //    33|444|5
  //    33|444|55
  //    ...
  // (Note that we expect the matrix to be sparse.  For dense matrices,
  // CrsMatrix is the wrong tool.)
  //
  // Matrix rows are assigned to processor rows greedily to roughly balance
  //   (# nonzeros in processor row / # processors in processor row)
  // across processor rows.
  // The same cuts are used to divide rows and columns among processors
  // (that is, all processors have a square block).
  //
  // The lower triangular algorithm:
  // 1. distribute all matrix entries via 1D linear distribution
  //    (this initial distribution is needed to avoid storing the entire
  //    matrix on one processor, while providing info about the nonzeros per row
  //    needed in step 2.
  // 2. (optional) sort rows in decreasing order wrt the number of nonzeros
  //    per row
  // 3. find "chunk cuts":  divisions in row assignments such that
  //    (# nonzeros in processor row / # processors in processor row) is
  //    roughly equal for all processor rows
  // 4. send nonzeros to their new processor assignment
  //
  // Known issues:  (TODO)
  // - The sorting in Step 2 and computation of chunk cuts in step 3 are
  //   currently done in serial and requires O(number of rows) storage each
  //   processor.  More effort could parallelize this computation, but parallel
  //   load balancing algorithms are more appropriate in Zoltan2 than Tpetra.
  // - The sorting in Step 2 renumbers the rows (assigns new Global Ordinals to
  //   the rows) to make them contiguous, as needed in Acer's triangle counting
  //   algorithm.
  //   (Acer's algorithm relies on local indexing from the chunk boundaries to
  //   find neighbors needed for communication.)
  //   The class currently provides a permutation matrix P describing the
  //   reordering.  Thus, the matrix stored in the lower triangular block
  //   distribution is actually P A P -- the row and column permutation of
  //   matrix A in the Matrix Market file.
  //   The fact that a permuted matrix is stored complicates use of the matrix
  //   in algorithms other than Acer's triangle counting.  For SpMV with the
  //   vector numbered according to the MatrixMarket numbering, for example,
  //   P^T must be applied to the vector before SpMV, and P^T must be applied to
  //   the result of SpMV.  See LowerTriangularBlockOperator to see how this
  //   permutation matrix is used.
  //
  // Before addressing these issues, we will decide (TODO)
  // -  Is this Distribution general enough to be in Tpetra?
  // -  Should we, instead, have a separate package for distributions (that could
  //    use Zoltan2 and Tpetra without circular dependence)?
  // -  Or should we allow users (such as the triangle counting algorithm) to
  //    provide their own distributions (e.g., LowerTriangularBlock) that
  //    inherit from Tpetra's Distribution class?
  // For now, we will push this Distribution into Tpetra, but we will revisit
  // this decision.

 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 typename Distribution<gno_t, scalar_t>::NZindex_t;
  using typename Distribution<gno_t, scalar_t>::LocalNZmap_t;

  using map_t    = Tpetra::Map<>;
  using matrix_t = Tpetra::CrsMatrix<scalar_t>;

  DistributionLowerTriangularBlock(size_t nrows_,
                                   const Teuchos::RCP<const Teuchos::Comm<int> > &comm_,
                                   const Teuchos::ParameterList &params)
    : Distribution<gno_t, scalar_t>(nrows_, comm_, params)
    , initialDist(nrows_, comm_, params)
    , sortByDegree(false)
    , permMatrix(Teuchos::null)
    , redistributed(false)
    , chunksComputed(false)
    , nChunks(0) {
    int npnp = 2 * np;
    nChunks  = int(std::sqrt(float(npnp)));
    while (nChunks * (nChunks + 1) < npnp) nChunks++;

    TEUCHOS_TEST_FOR_EXCEPTION(nChunks * (nChunks + 1) != npnp, std::logic_error,
                               "Number of processors np = " << np << " must satisfy np = q(q+1)/2 for some q"
                                                            << " for LowerTriangularBlock distribution");
    nChunksPerRow = double(nChunks) / double(nrows);

    const Teuchos::ParameterEntry *pe = params.getEntryPtr("sortByDegree");
    if (pe != NULL) sortByDegree = pe->getValue<bool>(&sortByDegree);

    pe = params.getEntryPtr("readPerProcess");
    if (pe != NULL) redistributed = pe->getValue<bool>(&redistributed);

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

  enum DistributionType DistType() { return LowerTriangularBlock; }

  bool areChunksComputed() { return chunksComputed; }

  Teuchos::Array<gno_t> getChunkCuts() {
    if (chunksComputed)
      return chunkCuts;
    else {
      throw std::runtime_error("Error:  Requested chunk cuts have not been computed yet.");
    }
  }

  // Return whether this rank owns vector entry i.
  // TODO:  for now, use same vector dist as 1DLinear;
  // TODO:  think about best distribution of Vectors
  inline bool VecMine(gno_t i) { return initialDist.VecMine(i); }

  // Return whether this rank owns nonzero (i,j)
  // Vector map and row map are the same in 1D distribution.
  // But keep only the lower Triangular entries
  bool Mine(gno_t i, gno_t j) {
    if (redistributed) {
      if (j > i)
        return false;  // Don't keep any upper triangular entries
      else
        return (procFromChunks(i, j) == me);
    } else
      return initialDist.Mine(i, j);
  }

  inline bool Mine(gno_t i, gno_t j, int p) { return Mine(i, j); }

  // How to redistribute according to chunk-based row distribution
  void Redistribute(LocalNZmap_t &localNZ) {
    // Going to do chunking and sorting serially for now;
    // need to gather per-row information from each processor
    // TODO:  think about a parallel implementation

    gno_t myFirstRow = initialDist.getFirstRow(me);
    gno_t nMyRows    = initialDist.getNumRow(me);
    Teuchos::Array<gno_t> nnzPerRow(nMyRows, 0);

    Teuchos::Array<int> rcvcnt(np);
    Teuchos::Array<int> disp(np);
    for (int sum = 0, p = 0; p < np; p++) {
      int prows = initialDist.getNumRow(p);
      rcvcnt[p] = prows;
      disp[p]   = sum;
      sum += prows;
    }

    // If desire sortByDegree, first need to sort with respect to ALL entries
    // in matrix (not lower-triangular entries);
    // decreasing sort by number of entries per row in global matrix.
    // Generate permuteIndex for the sorted rows

    Teuchos::Array<gno_t> permuteIndex;  // This is the inverse permutation
    Teuchos::Array<gno_t> sortedOrder;   // This is the original permutation

    Teuchos::Array<gno_t> globalRowBuf;
    // TODO Dunno why there isn't a Teuchos::gatherAllv;
    // TODO for now, compute and broadcast
    if (me == 0) {
      globalRowBuf.resize(nrows, 0);  // TODO:  Ick!  Need parallel
    }

    if (sortByDegree) {
      // Compute nnzPerRow; distribution is currently 1D and includes all nz
      for (auto it = localNZ.begin(); it != localNZ.end(); it++) {
        gno_t I = it->first.first;
        nnzPerRow[I - myFirstRow]++;
      }

      Teuchos::gatherv<int, gno_t>(nnzPerRow.getRawPtr(), nMyRows,
                                   globalRowBuf.getRawPtr(),
                                   rcvcnt.getRawPtr(), disp.getRawPtr(),
                                   0, *comm);

      permuteIndex.resize(nrows);  // TODO:  Ick!  Need parallel
      sortedOrder.resize(nrows);   // TODO:  Ick!  Need parallel

      if (me == 0) {  // TODO:  do on all procs once have allgatherv

        for (size_t i = 0; i != nrows; i++) sortedOrder[i] = i;

        std::sort(sortedOrder.begin(), sortedOrder.end(),
                  [&](const size_t &a, const size_t &b) {
                    return (globalRowBuf[a] > globalRowBuf[b]);
                  });

        // Compute inverse permutation; it is more useful for our needs
        for (size_t i = 0; i < nrows; i++) {
          permuteIndex[sortedOrder[i]] = i;
        }
      }

      Teuchos::broadcast<int, gno_t>(*comm, 0, permuteIndex(0, nrows));
      // Ick!  Use a directory  TODO

      // Sorting is changing the global IDs associated
      // with rows/columns.  To make this distribution applicable beyond
      // triangle counting (e.g., in a Tpetra operator), we need a way
      // to map from the original global IDs and back again.
      // Create a permutation matrix for use in the operator; use
      // default Tpetra layout.
      Teuchos::Array<gno_t> myRows;
      for (size_t i = 0; i < nrows; i++) {
        if (VecMine(i)) myRows.push_back(i);
      }

      Tpetra::global_size_t dummy =
          Teuchos::OrdinalTraits<Tpetra::global_size_t>::invalid();
      Teuchos::RCP<const map_t> permMap =
          rcp(new map_t(dummy, myRows(), 0, comm));

      permMatrix = rcp(new matrix_t(permMap, 1));  // one nz / row in permMatrix

      Teuchos::Array<gno_t> cols(1);
      Teuchos::Array<scalar_t> vals(1);
      vals[0] = 1.;

      for (size_t i = 0; i < permMap->getLocalNumElements(); i++) {
        gno_t gid = permMap->getGlobalElement(i);
        cols[0]   = permuteIndex[gid];
        permMatrix->insertGlobalValues(gid, cols(), vals());
      }

      permMatrix->fillComplete(permMap, permMap);
    }

    // Now, to determine the chunks, we care only about the number of
    // nonzeros in the lower triangular matrix.
    // Compute nnzPerRow; distribution is currently 1D
    nnzPerRow.assign(nMyRows, 0);
    size_t nnz = 0;
    for (auto it = localNZ.begin(); it != localNZ.end(); it++) {
      gno_t I = (sortByDegree ? permuteIndex[it->first.first]
                              : it->first.first);
      gno_t J = (sortByDegree ? permuteIndex[it->first.second]
                              : it->first.second);
      if (J <= I) {  // Lower-triangular part
        nnzPerRow[it->first.first - myFirstRow]++;
        nnz++;
      }
    }

    // TODO Dunno why there isn't a Teuchos::gatherAllv;
    // TODO for now, compute and broadcast

    Teuchos::gatherv<int, gno_t>(nnzPerRow.getRawPtr(), nMyRows,
                                 globalRowBuf.getRawPtr(),
                                 rcvcnt.getRawPtr(), disp.getRawPtr(),
                                 0, *comm);

    Teuchos::Array<int>().swap(rcvcnt);  // no longer needed
    Teuchos::Array<int>().swap(disp);    // no longer needed

    size_t gNnz;
    Teuchos::reduceAll(*comm, Teuchos::REDUCE_SUM, 1, &nnz, &gNnz);

    chunkCuts.resize(nChunks + 1, 0);

    if (me == 0) {  // TODO:  when have allgatherv, can do on all procs
                    // TODO:  or better, implement parallel version

      // Determine chunk cuts
      size_t target              = gNnz / np;  // target nnz per processor
      size_t targetRunningTotal  = 0;
      size_t currentRunningTotal = 0;
      gno_t I                    = gno_t(0);
      for (int chunkCnt = 0; chunkCnt < nChunks; chunkCnt++) {
        targetRunningTotal  = (target * (chunkCnt + 1));
        currentRunningTotal = 0;
        while (I < static_cast<gno_t>(nrows)) {
          size_t nextNnz = (sortByDegree ? globalRowBuf[sortedOrder[I]]
                                         : globalRowBuf[I]);
          if (currentRunningTotal + nextNnz <= targetRunningTotal) {
            currentRunningTotal += nextNnz;
            I++;
          } else
            break;
        }
        chunkCuts[chunkCnt + 1] = I;
      }
      chunkCuts[nChunks] = static_cast<gno_t>(nrows);
    }

    // Free memory associated with globalRowBuf
    Teuchos::Array<gno_t>().swap(globalRowBuf);

    Teuchos::broadcast<int, gno_t>(*comm, 0, chunkCuts(0, nChunks + 1));
    chunksComputed = true;

    // Determine new owner of each nonzero; buffer for sending
    Kokkos::View<gno_t *, Kokkos::HostSpace> iOut("iOut", localNZ.size());
    Kokkos::View<gno_t *, Kokkos::HostSpace> jOut("jOut", localNZ.size());
    Kokkos::View<scalar_t *, Kokkos::HostSpace> vOut("vOut", localNZ.size());
    Teuchos::Array<int> pOut(localNZ.size());

    size_t sendCnt = 0;
    for (auto it = localNZ.begin(); it != localNZ.end(); it++) {
      iOut[sendCnt] = (sortByDegree ? permuteIndex[it->first.first]
                                    : it->first.first);
      jOut[sendCnt] = (sortByDegree ? permuteIndex[it->first.second]
                                    : it->first.second);
      if (jOut[sendCnt] <= iOut[sendCnt]) {  // keep only lower diagonal entries
        vOut[sendCnt] = it->second;
        pOut[sendCnt] = procFromChunks(iOut[sendCnt], jOut[sendCnt]);

        sendCnt++;
      }
    }

    // Free memory associated with localNZ and permuteIndex
    LocalNZmap_t().swap(localNZ);
    if (sortByDegree) Teuchos::Array<gno_t>().swap(permuteIndex);

    // Use a Distributor to send nonzeros to new processors.
    Tpetra::Distributor plan(comm);
    size_t nrecvs = plan.createFromSends(pOut(0, sendCnt));
    Kokkos::View<gno_t *, Kokkos::HostSpace> iIn("iIn", nrecvs);
    Kokkos::View<gno_t *, Kokkos::HostSpace> jIn("jIn", nrecvs);
    Kokkos::View<scalar_t *, Kokkos::HostSpace> vIn("vIn", nrecvs);

    // TODO:  With more clever packing, could do only one round of communication
    auto sendIndices = std::make_pair(static_cast<size_t>(0), sendCnt);
    plan.doPostsAndWaits(Kokkos::subview(iOut, sendIndices), 1, iIn);
    plan.doPostsAndWaits(Kokkos::subview(jOut, sendIndices), 1, jIn);
    plan.doPostsAndWaits(Kokkos::subview(vOut, sendIndices), 1, vIn);

    // Put received nonzeros in map
    for (size_t n = 0; n < nrecvs; n++) {
      NZindex_t nz(iIn[n], jIn[n]);
      localNZ[nz] = vIn[n];
    }

    redistributed = true;
  }

  Teuchos::RCP<matrix_t> getPermutationMatrix() const { return permMatrix; }

 private:
  // Initially distribute nonzeros with a 1D linear distribution
  Distribution1DLinear<gno_t, scalar_t> initialDist;

  // Flag indicating whether matrix should be reordered and renumbered
  // in decreasing sort order of number of nonzeros per row in full matrix
  bool sortByDegree;

  // Column permutation matrix built only when sortByDegree = true;
  Teuchos::RCP<matrix_t> permMatrix;

  // Flag whether redistribution has occurred yet
  // This is true
  //   i) after Tpetra performs the redistribution or
  //  ii) when Tpetra reads already-distributed nonzeros by readPerProcess function
  bool redistributed;

  // If we read the already-distributed nonzeros from per-process files,
  //  this will remain false until a triangle counting code actually computes
  //  the chunks when the need arises.
  bool chunksComputed;

  int nChunks;  // in np = q(q+1)/2 layout, nChunks = q
  double nChunksPerRow;
  Teuchos::Array<gno_t> chunkCuts;

  int findIdxInChunks(gno_t I) {
    int m = I * nChunksPerRow;
    while (I < chunkCuts[m]) m--;
    while (I >= chunkCuts[m + 1]) m++;
    return m;
  }

  int procFromChunks(gno_t I, gno_t J) {
    int m = findIdxInChunks(I);
    int n = findIdxInChunks(J);
    int p = m * (m + 1) / 2 + n;
    return p;
  }
};

// Tpetra::Operator that works with the DistributionLowerTriangularBlock

template <typename scalar_t,
          class Node = ::Tpetra::Details::DefaultTypes::node_type>
class LowerTriangularBlockOperator : public Tpetra::Operator<scalar_t, Tpetra::Map<>::local_ordinal_type,
                                                             Tpetra::Map<>::global_ordinal_type,
                                                             Node> {
 public:
  using lno_t     = Tpetra::Map<>::local_ordinal_type;
  using gno_t     = Tpetra::Map<>::global_ordinal_type;
  using map_t     = Tpetra::Map<>;
  using import_t  = Tpetra::Import<>;
  using export_t  = Tpetra::Export<>;
  using vector_t  = Tpetra::Vector<scalar_t>;
  using mvector_t = Tpetra::MultiVector<scalar_t>;
  using matrix_t  = Tpetra::CrsMatrix<scalar_t, lno_t, gno_t, Node>;
  using dist_t    = Tpetra::DistributionLowerTriangularBlock<gno_t, scalar_t>;

  LowerTriangularBlockOperator(
      const Teuchos::RCP<const matrix_t> &lowerTriangularMatrix_,
      const dist_t &dist)
    : lowerTriangularMatrix(lowerTriangularMatrix_)
    , permMatrix(dist.getPermutationMatrix()) {
    // LowerTriangularBlockOperator requires the range map and domain map
    // to be the same.  Check it here.
    TEUCHOS_TEST_FOR_EXCEPTION(
        !lowerTriangularMatrix->getRangeMap()->isSameAs(
            *lowerTriangularMatrix->getDomainMap()),
        std::logic_error,
        "The Domain and Range maps of the LowerTriangularBlock matrix "
        "must be the same");

    // Extract diagonals

    vector_t diagByRowMap(lowerTriangularMatrix->getRowMap());
    lowerTriangularMatrix->getLocalDiagCopy(diagByRowMap);
    diag = Teuchos::rcp(new vector_t(lowerTriangularMatrix->getRangeMap()));
    Tpetra::Export<> exporter(lowerTriangularMatrix->getRowMap(),
                              lowerTriangularMatrix->getRangeMap());
    diag->doExport(diagByRowMap, exporter, Tpetra::ADD);
  }

  void apply(const mvector_t &x, mvector_t &y, Teuchos::ETransp mode,
             scalar_t alpha, scalar_t beta) const {
    scalar_t ZERO = Teuchos::ScalarTraits<scalar_t>::zero();
    scalar_t ONE  = Teuchos::ScalarTraits<scalar_t>::one();
    if (alpha == ZERO) {
      if (beta == ZERO)
        y.putScalar(ZERO);
      else
        y.scale(beta);
      return;
    }

    if (permMatrix == Teuchos::null) {
      // Multiply lower triangular
      lowerTriangularMatrix->apply(x, y, Teuchos::NO_TRANS, alpha, beta);

      // Multiply upper triangular
      lowerTriangularMatrix->apply(x, y, Teuchos::TRANS, alpha, ONE);

      // Subtract out duplicate diagonal terms
      y.elementWiseMultiply(-alpha, *diag, x, ONE);
    } else {
      // With sorting, the LowerTriangularBlock distribution stores (P^T A P)
      // in the CrsMatrix, for permutation matrix P.
      // Thus, apply must compute
      // y = P (beta (P^T y) + alpha (P^T A P) (P^T x))

      vector_t xtmp(x.getMap(), x.getNumVectors());
      vector_t ytmp(y.getMap(), y.getNumVectors());

      permMatrix->apply(x, xtmp, Teuchos::TRANS);
      if (beta != ZERO) permMatrix->apply(y, ytmp, Teuchos::TRANS);

      // Multiply lower triangular
      lowerTriangularMatrix->apply(xtmp, ytmp, Teuchos::NO_TRANS, alpha, beta);

      // Multiply upper triangular
      lowerTriangularMatrix->apply(xtmp, ytmp, Teuchos::TRANS, alpha, ONE);

      // Subtract out duplicate diagonal terms
      ytmp.elementWiseMultiply(-alpha, *diag, xtmp, ONE);

      permMatrix->apply(ytmp, y, Teuchos::NO_TRANS);
    }
  }

  Teuchos::RCP<const map_t> getDomainMap() const {
    return lowerTriangularMatrix->getDomainMap();
  }

  Teuchos::RCP<const map_t> getRangeMap() const {
    return lowerTriangularMatrix->getRangeMap();
  }

  bool hasTransposeApply() const { return true; }  // Symmetric matrix

 private:
  const Teuchos::RCP<const matrix_t> lowerTriangularMatrix;
  const Teuchos::RCP<const matrix_t> permMatrix;
  Teuchos::RCP<vector_t> diag;
};

}  // namespace Tpetra
#endif