Tpetra_Details_crsMatrixAssembleElement.hpp

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

#ifndef TPETRA_DETAILS_CRSMATRIXASSEMBLEELEMENT_HPP
#define TPETRA_DETAILS_CRSMATRIXASSEMBLEELEMENT_HPP

#include "KokkosSparse_CrsMatrix.hpp"
#include "Tpetra_Details_shortSort.hpp"
#include <type_traits>

namespace Tpetra {
namespace Details {

template <class SparseMatrixType,
          class ValsViewType>
KOKKOS_FUNCTION
    typename SparseMatrixType::ordinal_type
    crsMatrixSumIntoValues_sortedSortedLinear(const SparseMatrixType& A,
                                              const typename SparseMatrixType::ordinal_type lclRow,
                                              const typename SparseMatrixType::ordinal_type lclColInds[],
                                              const typename SparseMatrixType::ordinal_type sortPerm[],
                                              const ValsViewType& vals,
                                              const typename SparseMatrixType::ordinal_type numEntInInput,
                                              const bool forceAtomic =
#ifdef KOKKOS_ENABLE_SERIAL
                                                  !std::is_same<typename SparseMatrixType::device_type::execution_space, Kokkos::Serial>::type,
#else   // NOT KOKKOS_ENABLE_SERIAL
                                                  false,
#endif  // KOKKOS_ENABLE_SERIAL
                                              const bool checkInputIndices = true) {
  typedef typename std::remove_const<typename SparseMatrixType::value_type>::type
      matrix_scalar_type;
  static_assert(std::is_same<matrix_scalar_type,
                             typename SparseMatrixType::value_type>::value,
                "The matrix's entries must have a nonconst type.");
  // static_assert (std::is_assignable<matrix_scalar_type,
  //                typename std::decay< decltype (A.values[0] + vals[0]) >::type>::value,
  //                "The result of adding a matrix entry and an entry of vals "
  //                "MUST be assignable to a matrix entry.");
  typedef typename SparseMatrixType::ordinal_type LO;
  static_assert(std::is_integral<LO>::value,
                "SparseMatrixType::ordinal_type "
                "must be a built-in integer type.");

  // If lclRow is NOT a valid row index, this will return a view of
  // zero entries.  If checkInputIndices is true, thus, then none of
  // the input indices will be valid in that case.
  auto row_view        = A.row(lclRow);
  const LO numEntInRow = static_cast<LO>(row_view.length);
  // Number of valid local column indices found, that is, the number
  // of input indices that are valid column indices found in row
  // lclRow of the matrix.  If not checking, we just return the number
  // of input indices.
  LO numValid = checkInputIndices ? static_cast<LO>(0) : numEntInRow;

  // Since both the matrix row and the input (after permutation) are
  // sorted, we only need to pass once over the matrix row.  'offset'
  // tells us the current search position in the matrix row.
  LO offset = 0;
  for (LO j = 0; j < numEntInInput; ++j) {
    const LO perm_index = sortPerm[j];
    const LO lclColInd  = lclColInds[perm_index];
    // Search linearly in the matrix row for the current index.
    // If we ever want binary search, this would be the place.
    while (row_view.colidx(offset) != lclColInd) {
      ++offset;
    }

    // If we could make checkInputIndices a compile-time constant,
    // then the compiler might not need to insert a branch here.  This
    // should help vectorization, if vectorization is possible.
    if (checkInputIndices) {
      if (offset != numEntInRow) {
        // If we could make forceAtomic a compile-time constant, then
        // the compiler might not need to insert a branch here.  This
        // should help vectorization, if vectorization is possible.
        if (forceAtomic) {
          Kokkos::atomic_add(&(row_view.value(offset)), vals[perm_index]);
        } else {
          row_view.value(offset) += vals[perm_index];
        }
        ++numValid;
      }
    } else {  // don't check input indices; assume they are in the row
      // See above note on forceAtomic.
      if (forceAtomic) {
        Kokkos::atomic_add(&(row_view.value(offset)), vals[perm_index]);
      } else {
        row_view.value(offset) += vals[perm_index];
      }
    }
  }

  return numValid;
}

template <class SparseMatrixType,
          class ValsViewType>
KOKKOS_FUNCTION
    typename SparseMatrixType::ordinal_type
    crsMatrixReplaceValues_sortedSortedLinear(const SparseMatrixType& A,
                                              const typename SparseMatrixType::ordinal_type lclRow,
                                              const typename SparseMatrixType::ordinal_type lclColInds[],
                                              const typename SparseMatrixType::ordinal_type sortPerm[],
                                              const ValsViewType& vals,
                                              const typename SparseMatrixType::ordinal_type numEntInInput,
                                              const bool forceAtomic =
#ifdef KOKKOS_ENABLE_SERIAL
                                                  !std::is_same<typename SparseMatrixType::device_type::execution_space, Kokkos::Serial>::type,
#else   // NOT KOKKOS_ENABLE_SERIAL
                                                  false,
#endif  // KOKKOS_ENABLE_SERIAL
                                              const bool checkInputIndices = true) {
  typedef typename std::remove_const<typename SparseMatrixType::value_type>::type
      matrix_scalar_type;
  static_assert(std::is_same<matrix_scalar_type,
                             typename SparseMatrixType::value_type>::value,
                "The matrix's entries must have a nonconst type.");
  static_assert(std::is_assignable<matrix_scalar_type,
                                   typename std::decay<decltype(A.values[0] + vals[0])>::type>::value,
                "The result of adding a matrix entry and an entry of vals "
                "MUST be assignable to a matrix entry.");
  typedef typename SparseMatrixType::ordinal_type LO;
  static_assert(std::is_integral<LO>::value,
                "SparseMatrixType::ordinal_type "
                "must be a built-in integer type.");

  // If lclRow is NOT a valid row index, this will return a view of
  // zero entries.  If checkInputIndices is true, thus, then none of
  // the input indices will be valid in that case.
  auto row_view        = A.row(lclRow);
  const LO numEntInRow = static_cast<LO>(row_view.length);
  // Number of valid local column indices found, that is, the number
  // of input indices that are valid column indices found in row
  // lclRow of the matrix.  If not checking, we just return the number
  // of input indices.
  LO numValid = checkInputIndices ? static_cast<LO>(0) : numEntInRow;

  // Since both the matrix row and the input (after permutation) are
  // sorted, we only need to pass once over the matrix row.  'offset'
  // tells us the current search position in the matrix row.
  LO offset = 0;
  for (LO j = 0; j < numEntInInput; ++j) {
    const LO perm_index = sortPerm[j];
    const LO lclColInd  = lclColInds[perm_index];
    // Search linearly in the matrix row for the current index.
    // If we ever want binary search, this would be the place.
    while (row_view.colidx(offset) != lclColInd) {
      ++offset;
    }

    // If checkInputIndices were a compile-time constant, then the
    // compiler might not need to insert a branch here.  This should
    // help vectorization, if vectorization is possible at all.
    if (checkInputIndices) {
      if (offset != numEntInRow) {
        // If forceAtomic were a compile-time constant, then the
        // compiler might not need to insert a branch here.  This
        // could help vectorization, if vectorization is possible.
        if (forceAtomic) {
          Kokkos::atomic_store(&(row_view.value(offset)), vals[perm_index]);
        } else {
          row_view.value(offset) += vals[perm_index];
        }
        ++numValid;
      }
    } else {  // don't check input indices; assume they are in the row
      // See above note on forceAtomic.
      if (forceAtomic) {
        Kokkos::atomic_add(&(row_view.value(offset)), vals[perm_index]);
      } else {
        row_view.value(offset) += vals[perm_index];
      }
    }
  }

  return numValid;
}

template <class SparseMatrixType,
          class VectorViewType,
          class RhsViewType,
          class LhsViewType>
KOKKOS_FUNCTION
    typename SparseMatrixType::ordinal_type
    crsMatrixAssembleElement_sortedLinear(const SparseMatrixType& A,
                                          const VectorViewType& x,
                                          typename SparseMatrixType::ordinal_type lids[],
                                          typename SparseMatrixType::ordinal_type sortPerm[],
                                          const RhsViewType& rhs,
                                          const LhsViewType& lhs,
                                          const bool forceAtomic =
#ifdef KOKKOS_ENABLE_SERIAL
                                              !std::is_same<typename SparseMatrixType::device_type::execution_space, Kokkos::Serial>::type,
#else   // NOT KOKKOS_ENABLE_SERIAL
                                              false,
#endif  // KOKKOS_ENABLE_SERIAL
                                          const bool checkInputIndices = true) {
  typedef typename std::remove_const<typename SparseMatrixType::value_type>::type
      matrix_scalar_type;
  typedef typename std::remove_const<typename VectorViewType::value_type>::type
      vector_scalar_type;
  static_assert(std::is_same<matrix_scalar_type,
                             typename SparseMatrixType::value_type>::value,
                "The sparse output matrix A's entries must have a nonconst type.");
  static_assert(std::is_same<vector_scalar_type,
                             typename VectorViewType::value_type>::value,
                "The dense output vector x's entries must have a nonconst type.");
  // static_assert (std::is_assignable<matrix_scalar_type,
  //                typename std::decay< decltype (A.values[0] + lhs(0,0)) >::type>::value,
  //                "The result of adding a sparse matrix entry and an entry of "
  //                "lhs (the dense element matrix) "
  //                "MUST be assignable to a matrix entry.");
  // static_assert (std::is_assignable<vector_scalar_type,
  //                typename std::decay< decltype (x[0] + rhs[0]) >::type>::value,
  //                "The result of adding a vector entry and an entry of "
  //                "rhs (the dense element vector) "
  //                "MUST be assignable to a vector entry.");
  typedef typename SparseMatrixType::ordinal_type LO;
  static_assert(std::is_integral<LO>::value,
                "SparseMatrixType::ordinal_type "
                "must be a built-in integer type.");

  const LO eltDim = rhs.extent(0);

  // Generate sort permutation
  for (LO i = 0; i < eltDim; ++i) {
    sortPerm[i] = i;
  }
  shellSortKeysAndValues(lids, sortPerm, eltDim);

  LO totalNumValid = 0;
  for (LO r = 0; r < eltDim; ++r) {
    const LO lid = lids[r];
    // auto lhs_r = Kokkos::subview (lhs, sortPerm[r], Kokkos::ALL ());
    auto lhs_r = Kokkos::subview(lhs, r, Kokkos::ALL());

    // This assumes that lid is always a valid row in the sparse
    // matrix, and that the local indices in each row of the matrix
    // are always sorted.
    const LO curNumValid =
        crsMatrixSumIntoValues_sortedSortedLinear(A, lid, lids, sortPerm, lhs_r,
                                                  eltDim, forceAtomic,
                                                  checkInputIndices);
    if (forceAtomic) {
      Kokkos::atomic_add(&x(lid), rhs(sortPerm[r]));
    } else {
      x(lid) += rhs(sortPerm[r]);
    }
    totalNumValid += curNumValid;
  }
  return totalNumValid;
}

}  // namespace Details
}  // namespace Tpetra

#endif  // TPETRA_DETAILS_CRSMATRIXASSEMBLEELEMENT_HPP