Zoltan2_XpetraCrsMatrixAdapter.hpp

Go to the documentation of this file.

// @HEADER
// *****************************************************************************
//   Zoltan2: A package of combinatorial algorithms for scientific computing
//
// Copyright 2012 NTESS and the Zoltan2 contributors.
// SPDX-License-Identifier: BSD-3-Clause
// *****************************************************************************
// @HEADER

#ifndef _ZOLTAN2_XPETRACRSMATRIXADAPTER_HPP_
#define _ZOLTAN2_XPETRACRSMATRIXADAPTER_HPP_

#include <Zoltan2_MatrixAdapter.hpp>
#include <Zoltan2_StridedData.hpp>
#include <Zoltan2_XpetraTraits.hpp>
#include <Zoltan2_PartitioningHelpers.hpp>

#include <Xpetra_CrsMatrix.hpp>

#include <iostream>
#include <cassert>

namespace Zoltan2 {


template <typename User, typename UserCoord=User>
  class XpetraCrsMatrixAdapter : public MatrixAdapter<User, UserCoord> {
public:

#ifndef DOXYGEN_SHOULD_SKIP_THIS
  using scalar_t = typename InputTraits<User>::scalar_t;
  using lno_t = typename InputTraits<User>::lno_t;
  using gno_t = typename InputTraits<User>::gno_t;
  using part_t = typename InputTraits<User>::part_t;
  using node_t = typename InputTraits<User>::node_t;
  using offset_t = typename InputTraits<User>::offset_t;
  using xmatrix_t = Xpetra::CrsMatrix<scalar_t, lno_t, gno_t, node_t>;

  using userCoord_t = UserCoord;
  using user_t = User;
#endif

  ~XpetraCrsMatrixAdapter() { }

  XpetraCrsMatrixAdapter(const RCP<const User> &inmatrix,
                         int nWeightsPerRow=0);

  void setWeights(const scalar_t *weightVal, int stride, int idx = 0);

  void setRowWeights(const scalar_t *weightVal, int stride, int idx = 0);

  void setWeightIsDegree(int idx);

  void setRowWeightIsNumberOfNonZeros(int idx);

  // The MatrixAdapter interface.

  size_t getLocalNumRows() const {
    return matrix_->getLocalNumRows();
  }

  size_t getLocalNumColumns() const {
    return matrix_->getLocalNumCols();
  }

  size_t getLocalNumEntries() const {
    return matrix_->getLocalNumEntries();
  }

  void getRowIDsView(const gno_t *&rowIds) const
  {
    ArrayView<const gno_t> rowView = rowMap_->getLocalElementList();
    rowIds = rowView.getRawPtr();
  }

  void getColumnIDsView(const gno_t *&colIds) const
  {
    ArrayView<const gno_t> colView = colMap_->getLocalElementList();
    colIds = colView.getRawPtr();
  }

  void getCRSView(ArrayRCP<const offset_t> &offsets, ArrayRCP<const gno_t> &colIds) const
  {
    ArrayRCP< const lno_t > localColumnIds;
    ArrayRCP<const scalar_t> values;
    matrix_->getAllValues(offsets,localColumnIds,values);
    colIds = columnIds_;
  }

  bool CRSViewAvailable() const { return true; }

  void getCRSView(ArrayRCP<const offset_t> &offsets,
                  ArrayRCP<const gno_t> &colIds,
                  ArrayRCP<const scalar_t> &values) const {
    ArrayRCP< const lno_t > localColumnIds;
    matrix_->getAllValues(offsets,localColumnIds,values);
    colIds = columnIds_;
  }

  void getCCSView(ArrayRCP<const offset_t> &offsets,
                  ArrayRCP<const gno_t> &rowIds) const override {
    ArrayRCP<const offset_t> crsOffsets;
    ArrayRCP<const lno_t> crsLocalColumnIds;
    ArrayRCP<const scalar_t> values;
    matrix_->getAllValues(crsOffsets, crsLocalColumnIds, values);

    const auto localRowIds = rowMap_->getLocalElementList();
    const auto numLocalCols = colMap_->getLocalNumElements();

    // Lambda used to compute local row based on column index from CRS view
    auto determineRow = [&crsOffsets, &localRowIds](const int columnIdx) {
      int curLocalRow = 0;
      for (int rowIdx = 0; rowIdx < localRowIds.size(); ++rowIdx) {
        if (rowIdx < (localRowIds.size() - 1)) {
          if (static_cast<offset_t>(columnIdx) < crsOffsets[rowIdx + 1]) {
            return curLocalRow;
          }
          ++curLocalRow;
        } else {
          return curLocalRow;
        }
      }

      return -1;
    };

    // Vector of global rows per each local column
    std::vector<std::vector<gno_t>> rowIDsPerCol(numLocalCols);

    for (int colIdx = 0; colIdx < crsLocalColumnIds.size(); ++colIdx) {
      const auto colID = crsLocalColumnIds[colIdx];
      const auto globalRow = rowMap_->getGlobalElement(determineRow(colIdx));

      rowIDsPerCol[colID].push_back(globalRow);
    }

    size_t offsetWrite = 0;
    ArrayRCP<gno_t> ccsRowIds(values.size());
    ArrayRCP<offset_t> ccsOffsets(colMap_->getLocalNumElements() + 1);

    ccsOffsets[0] = 0;
    for (int64_t colID = 1; colID < ccsOffsets.size(); ++colID) {
      const auto &rowIDs = rowIDsPerCol[colID - 1];

      if (not rowIDs.empty()) {
        std::copy(rowIDs.begin(), rowIDs.end(),
                  ccsRowIds.begin() + offsetWrite);
        offsetWrite += rowIDs.size();
      }

      ccsOffsets[colID] = offsetWrite;
    }

    ccsOffsets[numLocalCols] = crsLocalColumnIds.size();

    rowIds = ccsRowIds;
    offsets = ccsOffsets;
  }

  int getNumWeightsPerRow() const { return nWeightsPerRow_; }

  void getRowWeightsView(const scalar_t *&weights, int &stride,
                           int idx = 0) const
  {
    if(idx<0 || idx >= nWeightsPerRow_)
    {
      std::ostringstream emsg;
      emsg << __FILE__ << ":" << __LINE__
           << "  Invalid row weight index " << idx << std::endl;
      throw std::runtime_error(emsg.str());
    }

    size_t length;
    rowWeights_[idx].getStridedList(length, weights, stride);
  }

  bool useNumNonzerosAsRowWeight(int idx) const { return numNzWeight_[idx];}

  template <typename Adapter>
    void applyPartitioningSolution(const User &in, User *&out,
         const PartitioningSolution<Adapter> &solution) const;

  template <typename Adapter>
    void applyPartitioningSolution(const User &in, RCP<User> &out,
         const PartitioningSolution<Adapter> &solution) const;

private:

  RCP<const User> inmatrix_;
  RCP<const xmatrix_t> matrix_;
  RCP<const Xpetra::Map<lno_t, gno_t, node_t> > rowMap_;
  RCP<const Xpetra::Map<lno_t, gno_t, node_t> > colMap_;
  lno_t base_;
  ArrayRCP<gno_t> columnIds_;  // TODO:  Refactor adapter to localColumnIds_

  int nWeightsPerRow_;
  ArrayRCP<StridedData<lno_t, scalar_t> > rowWeights_;
  ArrayRCP<bool> numNzWeight_;

  bool mayHaveDiagonalEntries;
};

// Definitions

template <typename User, typename UserCoord>
  XpetraCrsMatrixAdapter<User,UserCoord>::XpetraCrsMatrixAdapter(
    const RCP<const User> &inmatrix, int nWeightsPerRow):
      inmatrix_(inmatrix), matrix_(), rowMap_(), colMap_(),
      columnIds_(),
      nWeightsPerRow_(nWeightsPerRow), rowWeights_(), numNzWeight_(),
      mayHaveDiagonalEntries(true)
{
  typedef StridedData<lno_t,scalar_t> input_t;
  try {
    matrix_ = rcp_const_cast<const xmatrix_t>(
           XpetraTraits<User>::convertToXpetra(rcp_const_cast<User>(inmatrix)));
  }
  Z2_FORWARD_EXCEPTIONS

  rowMap_ = matrix_->getRowMap();
  colMap_ = matrix_->getColMap();

  size_t nrows = matrix_->getLocalNumRows();
  size_t nnz = matrix_->getLocalNumEntries();

  // Get ArrayRCP pointers to the structures in the underlying matrix
  ArrayRCP< const offset_t > offset;
  ArrayRCP< const lno_t > localColumnIds;
  ArrayRCP< const scalar_t > values;
  matrix_->getAllValues(offset,localColumnIds,values);
  columnIds_.resize(nnz, 0);

  for (offset_t i = 0; i < offset[nrows]; i++) {
    columnIds_[i] = colMap_->getGlobalElement(localColumnIds[i]);
  }

  if (nWeightsPerRow_ > 0){
    rowWeights_ = arcp(new input_t [nWeightsPerRow_], 0, nWeightsPerRow_, true);
    numNzWeight_ = arcp(new bool [nWeightsPerRow_], 0, nWeightsPerRow_, true);
    for (int i=0; i < nWeightsPerRow_; i++)
      numNzWeight_[i] = false;
  }
}

template <typename User, typename UserCoord>
  void XpetraCrsMatrixAdapter<User,UserCoord>::setWeights(
    const scalar_t *weightVal, int stride, int idx)
{
  if (this->getPrimaryEntityType() == MATRIX_ROW)
    setRowWeights(weightVal, stride, idx);
  else {
    // TODO:  Need to allow weights for columns and/or nonzeros
    std::ostringstream emsg;
    emsg << __FILE__ << "," << __LINE__
         << " error:  setWeights not yet supported for"
         << " columns or nonzeros."
         << std::endl;
    throw std::runtime_error(emsg.str());
  }
}

template <typename User, typename UserCoord>
  void XpetraCrsMatrixAdapter<User,UserCoord>::setRowWeights(
    const scalar_t *weightVal, int stride, int idx)
{
  typedef StridedData<lno_t,scalar_t> input_t;
  if(idx<0 || idx >= nWeightsPerRow_)
  {
      std::ostringstream emsg;
      emsg << __FILE__ << ":" << __LINE__
           << "  Invalid row weight index " << idx << std::endl;
      throw std::runtime_error(emsg.str());
  }

  size_t nvtx = getLocalNumRows();
  ArrayRCP<const scalar_t> weightV(weightVal, 0, nvtx*stride, false);
  rowWeights_[idx] = input_t(weightV, stride);
}

template <typename User, typename UserCoord>
  void XpetraCrsMatrixAdapter<User,UserCoord>::setWeightIsDegree(
    int idx)
{
  if (this->getPrimaryEntityType() == MATRIX_ROW)
    setRowWeightIsNumberOfNonZeros(idx);
  else {
    // TODO:  Need to allow weights for columns and/or nonzeros
    std::ostringstream emsg;
    emsg << __FILE__ << "," << __LINE__
         << " error:  setWeightIsNumberOfNonZeros not yet supported for"
         << " columns" << std::endl;
    throw std::runtime_error(emsg.str());
  }
}

template <typename User, typename UserCoord>
  void XpetraCrsMatrixAdapter<User,UserCoord>::setRowWeightIsNumberOfNonZeros(
    int idx)
{
  if(idx<0 || idx >= nWeightsPerRow_)
  {
      std::ostringstream emsg;
      emsg << __FILE__ << ":" << __LINE__
           << "  Invalid row weight index " << idx << std::endl;
      throw std::runtime_error(emsg.str());
  }


  numNzWeight_[idx] = true;
}

template <typename User, typename UserCoord>
  template <typename Adapter>
    void XpetraCrsMatrixAdapter<User,UserCoord>::applyPartitioningSolution(
      const User &in, User *&out,
      const PartitioningSolution<Adapter> &solution) const
{
  // Get an import list (rows to be received)
  size_t numNewRows;
  ArrayRCP<gno_t> importList;
  try{
    numNewRows = Zoltan2::getImportList<Adapter,
                                        XpetraCrsMatrixAdapter<User,UserCoord> >
                                       (solution, this, importList);
  }
  Z2_FORWARD_EXCEPTIONS;

  // Move the rows, creating a new matrix.
  RCP<User> outPtr = XpetraTraits<User>::doMigration(in, numNewRows,
                                                     importList.getRawPtr());
  out = const_cast<User *>(outPtr.get());
  outPtr.release();
}

template <typename User, typename UserCoord>
  template <typename Adapter>
    void XpetraCrsMatrixAdapter<User,UserCoord>::applyPartitioningSolution(
      const User &in, RCP<User> &out,
      const PartitioningSolution<Adapter> &solution) const
{
  // Get an import list (rows to be received)
  size_t numNewRows;
  ArrayRCP<gno_t> importList;
  try{
    numNewRows = Zoltan2::getImportList<Adapter,
                                        XpetraCrsMatrixAdapter<User,UserCoord> >
                                       (solution, this, importList);
  }
  Z2_FORWARD_EXCEPTIONS;

  // Move the rows, creating a new matrix.
  out = XpetraTraits<User>::doMigration(in, numNewRows,
                                        importList.getRawPtr());
}

}  //namespace Zoltan2

#endif