Tpetra_Details_getDiagCopyWithoutOffsets_decl.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_GETDIAGCOPYWITHOUTOFFSETS_DECL_HPP
#define TPETRA_DETAILS_GETDIAGCOPYWITHOUTOFFSETS_DECL_HPP


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

namespace Tpetra {
namespace Details {

template <class DiagType,
          class LocalMapType,
          class CrsMatrixType>
struct CrsMatrixGetDiagCopyFunctor {
  typedef typename LocalMapType::local_ordinal_type LO;   // local ordinal type
  typedef typename LocalMapType::global_ordinal_type GO;  // global ordinal type
  typedef typename CrsMatrixType::device_type device_type;
  typedef typename CrsMatrixType::value_type scalar_type;
  typedef typename CrsMatrixType::size_type offset_type;

  typedef LO value_type;

  CrsMatrixGetDiagCopyFunctor(const DiagType& D,
                              const LocalMapType& rowMap,
                              const LocalMapType& colMap,
                              const CrsMatrixType& A)
    : D_(D)
    , rowMap_(rowMap)
    , colMap_(colMap)
    , A_(A) {}

  KOKKOS_FUNCTION void
  operator()(const LO& lclRowInd, value_type& errCount) const {
    const LO INV = Tpetra::Details::OrdinalTraits<LO>::invalid();
    const scalar_type ZERO =
#if KOKKOS_VERSION >= 40799
        KokkosKernels::ArithTraits<scalar_type>::zero();
#else
        Kokkos::ArithTraits<scalar_type>::zero();
#endif

    // If the row lacks a stored diagonal entry, then its value is zero.
    D_(lclRowInd)      = ZERO;
    const GO gblInd    = rowMap_.getGlobalElement(lclRowInd);
    const LO lclColInd = colMap_.getLocalElement(gblInd);

    if (lclColInd != INV) {
      auto curRow = A_.rowConst(lclRowInd);

      // FIXME (mfh 12 May 2016) Use binary search when the row is
      // long enough.  findRelOffset currently lives in KokkosKernels
      // (in tpetra/kernels/src/Kokkos_Sparse_findRelOffset.hpp).
      LO offset       = 0;
      const LO numEnt = curRow.length;
      for (; offset < numEnt; ++offset) {
        if (curRow.colidx(offset) == lclColInd) {
          break;
        }
      }

      if (offset == numEnt) {
        ++errCount;
      } else {
        D_(lclRowInd) = curRow.value(offset);
      }
    }
  }

 private:
  DiagType D_;
  LocalMapType rowMap_;
  LocalMapType colMap_;
  CrsMatrixType A_;
};

template <class DiagType,
          class LocalMapType,
          class CrsMatrixType>
static typename LocalMapType::local_ordinal_type
getDiagCopyWithoutOffsets(const DiagType& D,
                          const LocalMapType& rowMap,
                          const LocalMapType& colMap,
                          const CrsMatrixType& A) {
  static_assert(Kokkos::is_view<DiagType>::value,
                "DiagType must be a Kokkos::View.");
  static_assert(static_cast<int>(DiagType::rank) == 1,
                "DiagType must be a 1-D Kokkos::View.");
  static_assert(std::is_same<typename DiagType::value_type, typename DiagType::non_const_value_type>::value,
                "DiagType must be a nonconst Kokkos::View.");
  typedef typename LocalMapType::local_ordinal_type LO;
  typedef typename CrsMatrixType::device_type device_type;
  typedef typename device_type::execution_space execution_space;
  typedef Kokkos::RangePolicy<execution_space, LO> policy_type;

  typedef Kokkos::View<typename DiagType::non_const_value_type*,
                       typename DiagType::array_layout,
                       typename DiagType::device_type,
                       Kokkos::MemoryUnmanaged>
      diag_type;
  diag_type D_out = D;
  CrsMatrixGetDiagCopyFunctor<diag_type, LocalMapType, CrsMatrixType>
      functor(D_out, rowMap, colMap, A);
  const LO numRows = static_cast<LO>(D.extent(0));
  LO errCount      = 0;
  Kokkos::parallel_reduce(policy_type(0, numRows), functor, errCount);
  return errCount;
}

template <class SC, class LO, class GO, class NT>
LO getLocalDiagCopyWithoutOffsetsNotFillComplete(::Tpetra::Vector<SC, LO, GO, NT>& diag,
                                                 const ::Tpetra::RowMatrix<SC, LO, GO, NT>& A,
                                                 const bool debug =
#ifdef HAVE_TPETRA_DEBUG
                                                     true);
#else   // ! HAVE_TPETRA_DEBUG
                                                     false);
#endif  // HAVE_TPETRA_DEBUG

}  // namespace Details
}  // namespace Tpetra

#endif  // TPETRA_DETAILS_GETDIAGCOPYWITHOUTOFFSETS_DECL_HPP