Tpetra_Details_rightScaleLocalCrsMatrix.hpp

Go to the documentation of this file.

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

#ifndef TPETRA_DETAILS_RIGHTSCALELOCALCRSMATRIX_HPP
#define TPETRA_DETAILS_RIGHTSCALELOCALCRSMATRIX_HPP


#include "TpetraCore_config.h"
#include "Kokkos_Core.hpp"
#if KOKKOS_VERSION >= 40799
#include "KokkosKernels_ArithTraits.hpp"
#else
#include "Kokkos_ArithTraits.hpp"
#endif
#include <type_traits>

namespace Tpetra {
namespace Details {

template <class LocalSparseMatrixType,
          class ScalingFactorsViewType,
          const bool divide>
class RightScaleLocalCrsMatrix {
 public:
  using val_type =
      typename std::remove_const<typename LocalSparseMatrixType::value_type>::type;
  using mag_type = typename ScalingFactorsViewType::non_const_value_type;
  static_assert(ScalingFactorsViewType::rank == 1,
                "scalingFactors must be a rank-1 Kokkos::View.");
  using device_type = typename LocalSparseMatrixType::device_type;
  using LO          = typename LocalSparseMatrixType::ordinal_type;
  using policy_type = Kokkos::TeamPolicy<typename device_type::execution_space, LO>;

  RightScaleLocalCrsMatrix(const LocalSparseMatrixType& A_lcl,
                           const ScalingFactorsViewType& scalingFactors,
                           const bool assumeSymmetric)
    : A_lcl_(A_lcl)
    , scalingFactors_(scalingFactors)
    , assumeSymmetric_(assumeSymmetric) {}

  KOKKOS_INLINE_FUNCTION void
  operator()(const typename policy_type::member_type& team) const {
#if KOKKOS_VERSION >= 40799
    using KAM = KokkosKernels::ArithTraits<mag_type>;
#else
    using KAM = Kokkos::ArithTraits<mag_type>;
#endif

    const LO lclRow = team.league_rank();
    auto curRow     = A_lcl_.row(lclRow);
    const LO numEnt = curRow.length;
    Kokkos::parallel_for(Kokkos::TeamThreadRange(team, numEnt), [&](const LO k) {
      const LO lclColInd        = curRow.colidx(k);
      const mag_type curColNorm = scalingFactors_(lclColInd);
      // Users are responsible for any divisions or multiplications by
      // zero.
      const mag_type scalingFactor = assumeSymmetric_ ? KAM::sqrt(curColNorm) : curColNorm;
      if (divide) {  // constexpr, so should get compiled out
        curRow.value(k) = curRow.value(k) / scalingFactor;
      } else {
        curRow.value(k) = curRow.value(k) * scalingFactor;
      }
    });
  }

 private:
  LocalSparseMatrixType A_lcl_;
  typename ScalingFactorsViewType::const_type scalingFactors_;
  bool assumeSymmetric_;
};

template <class LocalSparseMatrixType, class ScalingFactorsViewType>
void rightScaleLocalCrsMatrix(const LocalSparseMatrixType& A_lcl,
                              const ScalingFactorsViewType& scalingFactors,
                              const bool assumeSymmetric,
                              const bool divide = true) {
  using device_type     = typename LocalSparseMatrixType::device_type;
  using execution_space = typename device_type::execution_space;
  using LO              = typename LocalSparseMatrixType::ordinal_type;
  using policy_type     = Kokkos::TeamPolicy<execution_space, LO>;

  const LO lclNumRows = A_lcl.numRows();
  if (divide) {
    using functor_type =
        RightScaleLocalCrsMatrix<LocalSparseMatrixType,
                                 typename ScalingFactorsViewType::const_type, true>;
    functor_type functor(A_lcl, scalingFactors, assumeSymmetric);
    Kokkos::parallel_for("rightScaleLocalCrsMatrix",
                         policy_type(lclNumRows, Kokkos::AUTO), functor);
  } else {
    using functor_type =
        RightScaleLocalCrsMatrix<LocalSparseMatrixType,
                                 typename ScalingFactorsViewType::const_type, false>;
    functor_type functor(A_lcl, scalingFactors, assumeSymmetric);
    Kokkos::parallel_for("rightScaleLocalCrsMatrix",
                         policy_type(lclNumRows, Kokkos::AUTO), functor);
  }
}

}  // namespace Details
}  // namespace Tpetra

#endif  // TPETRA_DETAILS_RIGHTSCALELOCALCRSMATRIX_HPP