Zoltan2_TpetraCrsMatrixAdapter.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_TPETRACRSMATRIXADAPTER_HPP_
#define _ZOLTAN2_TPETRACRSMATRIXADAPTER_HPP_
 
#include <Zoltan2_MatrixAdapter.hpp>
#include <Zoltan2_TpetraRowMatrixAdapter.hpp>
#include <Zoltan2_StridedData.hpp>
#include <Zoltan2_InputTraits.hpp>
#include <Zoltan2_PartitioningHelpers.hpp>
 
#include <Tpetra_CrsMatrix.hpp>
 
#include <vector>
 
namespace Zoltan2 {
 
 
template <typename User, typename UserCoord = User>
  class TpetraCrsMatrixAdapter : public TpetraRowMatrixAdapter<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 tmatrix_t = Tpetra::CrsMatrix<scalar_t, lno_t, gno_t, node_t>;
  using device_t = typename node_t::device_type;
  using host_t = typename Kokkos::HostSpace::memory_space;
  using user_t = User;
  using userCoord_t = UserCoord;
 
  using Base = MatrixAdapter<User, UserCoord>;
  using RowMatrix = TpetraRowMatrixAdapter<User, UserCoord>;
  #endif
 
  TpetraCrsMatrixAdapter(const RCP<const User> &inmatrix,
                         int nWeightsPerRow=0);
 
  void init(const RCP<const User> &inmatrix);
 
  RCP<const User> getUserMatrix() const { return this->matrix_; }
 
  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;
};
 
// Definitions
 
  template <typename User, typename UserCoord>
  void TpetraCrsMatrixAdapter<User,UserCoord>::init(const RCP<const User> &inmatrix) {
    auto colIdsDevice = inmatrix->getLocalIndicesDevice();
 
    auto colIdsGlobalDevice =
      typename TpetraCrsMatrixAdapter<User, UserCoord>::Base::IdsDeviceView("colIdsGlobalDevice", colIdsDevice.extent(0));
    auto colMap = inmatrix->getColMap();
    auto lclColMap = colMap->getLocalMap();
 
    // Convert to global IDs using Tpetra::Map
    Kokkos::parallel_for("colIdsGlobalDevice",
                         Kokkos::RangePolicy<typename User::node_type::execution_space>(
                                                                                        0, colIdsGlobalDevice.extent(0)),
                         KOKKOS_LAMBDA(const int i) {
      colIdsGlobalDevice(i) =
        lclColMap.getGlobalElement(colIdsDevice(i));
    });
 
    this->colIdsDevice_ =   colIdsGlobalDevice;
    this->offsDevice_ = inmatrix->getLocalRowPtrsDevice();
  }
 
  template <typename User, typename UserCoord>
  TpetraCrsMatrixAdapter<User,UserCoord>::TpetraCrsMatrixAdapter(
    const RCP<const User> &inmatrix, int nWeightsPerRow):
      RowMatrix(nWeightsPerRow, inmatrix) {
 
  init(inmatrix);
 
  if (this->nWeightsPerRow_ > 0) {
 
    this->rowWeightsDevice_ = typename Base::WeightsDeviceView(
        "rowWeightsDevice_", inmatrix->getLocalNumRows(),
        this->nWeightsPerRow_);
 
    this->numNzWeight_ =  Kokkos::View<bool *, host_t>(
        "numNzWeight_", this->nWeightsPerRow_);
 
    for (int i = 0; i < this->nWeightsPerRow_; ++i) {
      this->numNzWeight_(i) = false;
    }
  }
}
 
template <typename User, typename UserCoord>
  template <typename Adapter>
    void TpetraCrsMatrixAdapter<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,
                                        TpetraCrsMatrixAdapter<User,UserCoord> >
                                       (solution, this, importList);
  }
  Z2_FORWARD_EXCEPTIONS;
 
  // Move the rows, creating a new matrix.
  RCP<User> outPtr = this->doMigration(in, numNewRows,importList.getRawPtr());
  out = const_cast<User *>(outPtr.get());
  outPtr.release();
}
 
template <typename User, typename UserCoord>
  template <typename Adapter>
    void TpetraCrsMatrixAdapter<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,
                                        TpetraCrsMatrixAdapter<User,UserCoord> >
                                       (solution, this, importList);
  }
  Z2_FORWARD_EXCEPTIONS;
 
  // Move the rows, creating a new matrix.
  out = this->doMigration(in, numNewRows, importList.getRawPtr());
}
 
}  //namespace Zoltan2
 
#endif