// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
// SPDX-FileCopyrightText: Copyright Contributors to the Kokkos project
#ifndef KOKKOSBATCHED_SPMV_HPP
#define KOKKOSBATCHED_SPMV_HPP

/// \author Kim Liegeois (knliege@sandia.gov)

#include "KokkosBatched_Util.hpp"
#include "KokkosBatched_Vector.hpp"

namespace KokkosBatched {

/// \brief Serial Batched SPMV:
///   y_l <- alpha_l * A_l * x_l + beta_l * y_l for all l = 1, ..., N
/// where:
///   * N is the number of matrices,
///   * A_1, ..., A_N are N sparse matrices which share the same sparsity
///   pattern,
///   * x_1, ..., x_N are the N input vectors,
///   * y_1, ..., y_N are the N output vectors,
///   * alpha_1, ..., alpha_N are N scaling factors for x_1, ..., x_N,
///   * beta_1, ..., beta_N are N scaling factors for y_1, ..., y_N.
///
/// \tparam ValuesViewType: Input type for the values of the batched crs matrix,
/// needs to be a 2D view \tparam IntView: Input type for row offset array and
/// column-index array, needs to be a 1D view \tparam xViewType: Input type for
/// X, needs to be a 2D view \tparam yViewType: Input type for Y, needs to be a
/// 2D view \tparam alphaViewType: Input type for alpha, needs to be a 1D view
/// \tparam betaViewType: Input type for beta, needs to be a 1D view
/// \tparam dobeta: Int which sepcifies if beta_l * y_l is used or not (if
/// dobeta == 0, beta_l * y_l is not added to the result of alpha_l * A_l * x_l)
///
/// \param alpha [in]: input coefficient for X, a rank 1 view
/// \param values [in]: values of the batched crs matrix, a rank 2 view
/// \param row_ptr [in]: row offset array of the batched crs matrix, a rank 1
/// view \param colIndices [in]: column-index array of the batched crs matrix, a
/// rank 1 view \param X [in]: Input vector X, a rank 2 view \param beta [in]:
/// input coefficient for Y (if dobeta != 0), a rank 1 view \param Y [in/out]:
/// Output vector Y, a rank 2 view
///
/// The matrices are represented using a Compressed Row Storage (CRS) format and
/// the shared sparsity pattern is reused from one matrix to the others.
///
/// Concretely, instead of providing an array of N matrices to the batched SPMV
/// kernel, the user provides one row offset array (1D view), one column-index
/// array (1D view), and one value array (2D view, one dimension for the
/// non-zero indices and one for the matrix indices).
///
/// No nested parallel_for is used inside of the function.
///

template <typename ArgTrans = Trans::NoTranspose>
struct SerialSpmv {
  template <typename ValuesViewType, typename IntView, typename xViewType, typename yViewType, typename alphaViewType,
            typename betaViewType, int dobeta>
  KOKKOS_INLINE_FUNCTION static int invoke(const alphaViewType &alpha, const ValuesViewType &values,
                                           const IntView &row_ptr, const IntView &colIndices, const xViewType &x,
                                           const betaViewType &beta, const yViewType &Y);

  template <typename ValuesViewType, typename IntView, typename xViewType, typename yViewType, int dobeta>
  KOKKOS_INLINE_FUNCTION static int invoke(
      const typename KokkosKernels::ArithTraits<typename ValuesViewType::non_const_value_type>::mag_type &alpha,
      const ValuesViewType &values, const IntView &row_ptr, const IntView &colIndices, const xViewType &X,
      const typename KokkosKernels::ArithTraits<typename ValuesViewType::non_const_value_type>::mag_type &beta,
      const yViewType &Y);
};

/// \brief Team Batched SPMV:
///   y_l <- alpha_l * A_l * x_l + beta_l * y_l for all l = 1, ..., N
/// where:
///   * N is the number of matrices,
///   * A_1, ..., A_N are N sparse matrices which share the same sparsity
///   pattern,
///   * x_1, ..., x_N are the N input vectors,
///   * y_1, ..., y_N are the N output vectors,
///   * alpha_1, ..., alpha_N are N scaling factors for x_1, ..., x_N,
///   * beta_1, ..., beta_N are N scaling factors for y_1, ..., y_N.
///
/// \tparam ValuesViewType: Input type for the values of the batched crs matrix,
/// needs to be a 2D view \tparam IntView: Input type for row offset array and
/// column-index array, needs to be a 1D view \tparam xViewType: Input type for
/// X, needs to be a 2D view \tparam yViewType: Input type for Y, needs to be a
/// 2D view \tparam alphaViewType: Input type for alpha, needs to be a 1D view
/// \tparam betaViewType: Input type for beta, needs to be a 1D view
/// \tparam dobeta: Int which sepcifies if beta_l * y_l is used or not (if
/// dobeta == 0, beta_l * y_l is not added to the result of alpha_l * A_l * x_l)
///
/// \param member [in]: TeamPolicy member
/// \param alpha [in]: input coefficient for X, a rank 1 view
/// \param values [in]: values of the batched crs matrix, a rank 2 view
/// \param row_ptr [in]: row offset array of the batched crs matrix, a rank 1
/// view \param colIndices [in]: column-index array of the batched crs matrix, a
/// rank 1 view \param X [in]: Input vector X, a rank 2 view \param beta [in]:
/// input coefficient for Y (if dobeta != 0), a rank 1 view \param Y [in/out]:
/// Output vector Y, a rank 2 view
///
/// The matrices are represented using a Compressed Row Storage (CRS) format and
/// the shared sparsity pattern is reused from one matrix to the others.
///
/// Concretely, instead of providing an array of N matrices to the batched SPMV
/// kernel, the user provides one row offset array (1D view), one column-index
/// array (1D view), and one value array (2D view, one dimension for the
/// non-zero indices and one for the matrix indices).
///
/// A nested parallel_for with TeamThreadRange is used.
///

template <typename MemberType, typename ArgTrans = Trans::NoTranspose>
struct TeamSpmv {
  template <typename ValuesViewType, typename IntView, typename xViewType, typename yViewType, typename alphaViewType,
            typename betaViewType, int dobeta>
  KOKKOS_INLINE_FUNCTION static int invoke(const MemberType &member, const alphaViewType &alpha,
                                           const ValuesViewType &values, const IntView &row_ptr,
                                           const IntView &colIndices, const xViewType &x, const betaViewType &beta,
                                           const yViewType &y);

  template <typename ValuesViewType, typename IntView, typename xViewType, typename yViewType, int dobeta>
  KOKKOS_INLINE_FUNCTION static int invoke(
      const MemberType &member,
      const typename KokkosKernels::ArithTraits<typename ValuesViewType::non_const_value_type>::mag_type &alpha,
      const ValuesViewType &values, const IntView &row_ptr, const IntView &colIndices, const xViewType &x,
      const typename KokkosKernels::ArithTraits<typename ValuesViewType::non_const_value_type>::mag_type &beta,
      const yViewType &y);
};

/// \brief TeamVector Batched SPMV:
///   y_l <- alpha_l * A_l * x_l + beta_l * y_l for all l = 1, ..., N
/// where:
///   * N is the number of matrices,
///   * A_1, ..., A_N are N sparse matrices which share the same sparsity
///   pattern,
///   * x_1, ..., x_N are the N input vectors,
///   * y_1, ..., y_N are the N output vectors,
///   * alpha_1, ..., alpha_N are N scaling factors for x_1, ..., x_N,
///   * beta_1, ..., beta_N are N scaling factors for y_1, ..., y_N.
///
/// \tparam ValuesViewType: Input type for the values of the batched crs matrix,
/// needs to be a 2D view \tparam IntView: Input type for row offset array and
/// column-index array, needs to be a 1D view \tparam xViewType: Input type for
/// X, needs to be a 2D view \tparam yViewType: Input type for Y, needs to be a
/// 2D view \tparam alphaViewType: Input type for alpha, needs to be a 1D view
/// \tparam betaViewType: Input type for beta, needs to be a 1D view
/// \tparam dobeta: Int which sepcifies if beta_l * y_l is used or not (if
/// dobeta == 0, beta_l * y_l is not added to the result of alpha_l * A_l * x_l)
///
/// \param member [in]: TeamPolicy member
/// \param alpha [in]: input coefficient for X, a rank 1 view
/// \param values [in]: values of the batched crs matrix, a rank 2 view
/// \param row_ptr [in]: row offset array of the batched crs matrix, a rank 1
/// view \param colIndices [in]: column-index array of the batched crs matrix, a
/// rank 1 view \param X [in]: Input vector X, a rank 2 view \param beta [in]:
/// input coefficient for Y (if dobeta != 0), a rank 1 view \param Y [in/out]:
/// Output vector Y, a rank 2 view
///
/// The matrices are represented using a Compressed Row Storage (CRS) format and
/// the shared sparsity pattern is reused from one matrix to the others.
///
/// Concretely, instead of providing an array of N matrices to the batched SPMV
/// kernel, the user provides one row offset array (1D view), one column-index
/// array (1D view), and one value array (2D view, one dimension for the
/// non-zero indices and one for the matrix indices).
///
/// Two nested parallel_for with both TeamThreadRange and ThreadVectorRange
/// (or one with TeamVectorRange) are used inside.
///

template <typename MemberType, typename ArgTrans = Trans::NoTranspose, unsigned N_team = 1>
struct TeamVectorSpmv {
  template <typename ValuesViewType, typename IntView, typename xViewType, typename yViewType, typename alphaViewType,
            typename betaViewType, int dobeta>
  KOKKOS_INLINE_FUNCTION static int invoke(const MemberType &member, const alphaViewType &alpha,
                                           const ValuesViewType &values, const IntView &row_ptr,
                                           const IntView &colIndices, const xViewType &x, const betaViewType &beta,
                                           const yViewType &y);

  template <typename ValuesViewType, typename IntView, typename xViewType, typename yViewType, int dobeta>
  KOKKOS_INLINE_FUNCTION static int invoke(
      const MemberType &member,
      const typename KokkosKernels::ArithTraits<typename ValuesViewType::non_const_value_type>::mag_type &alpha,
      const ValuesViewType &values, const IntView &row_ptr, const IntView &colIndices, const xViewType &x,
      const typename KokkosKernels::ArithTraits<typename ValuesViewType::non_const_value_type>::mag_type &beta,
      const yViewType &y);
};

/// \brief Batched SPMV: Selective Interface
///   y_l <- alpha_l * A_l * x_l + beta_l * y_l for all l = 1, ..., N
/// where:
///   * N is the number of matrices,
///   * A_1, ..., A_N are N sparse matrices which share the same sparsity
///   pattern,
///   * x_1, ..., x_N are the N input vectors,
///   * y_1, ..., y_N are the N output vectors,
///   * alpha_1, ..., alpha_N are N scaling factors for x_1, ..., x_N,
///   * beta_1, ..., beta_N are N scaling factors for y_1, ..., y_N.
///
/// \tparam ValuesViewType: Input type for the values of the batched crs matrix,
/// needs to be a 2D view \tparam IntView: Input type for row offset array and
/// column-index array, needs to be a 1D view \tparam xViewType: Input type for
/// X, needs to be a 2D view \tparam yViewType: Input type for Y, needs to be a
/// 2D view \tparam alphaViewType: Input type for alpha, needs to be a 1D view
/// \tparam betaViewType: Input type for beta, needs to be a 1D view
/// \tparam dobeta: Int which sepcifies if beta_l * y_l is used or not (if
/// dobeta == 0, beta_l * y_l is not added to the result of alpha_l * A_l * x_l)
///
/// \param member [in]: TeamPolicy member
/// \param alpha [in]: input coefficient for X, a rank 1 view
/// \param values [in]: values of the batched crs matrix, a rank 2 view
/// \param row_ptr [in]: row offset array of the batched crs matrix, a rank 1
/// view \param colIndices [in]: column-index array of the batched crs matrix, a
/// rank 1 view \param X [in]: Input vector X, a rank 2 view \param beta [in]:
/// input coefficient for Y (if dobeta != 0), a rank 1 view \param Y [in/out]:
/// Output vector Y, a rank 2 view

template <typename MemberType, typename ArgTrans, typename ArgMode>
struct Spmv {
  template <typename ValuesViewType, typename IntView, typename xViewType, typename yViewType, typename alphaViewType,
            typename betaViewType, int dobeta>
  KOKKOS_INLINE_FUNCTION static int invoke(const MemberType &member, const alphaViewType &alpha,
                                           const ValuesViewType &values, const IntView &row_ptr,
                                           const IntView &colIndices, const xViewType &x, const betaViewType &beta,
                                           const yViewType &y) {
    int r_val = 0;
    if (std::is_same<ArgMode, Mode::Serial>::value) {
      r_val =
          SerialSpmv<ArgTrans>::template invoke<ValuesViewType, IntView, xViewType, yViewType, alphaViewType,
                                                betaViewType, dobeta>(alpha, values, row_ptr, colIndices, x, beta, y);
    } else if (std::is_same<ArgMode, Mode::Team>::value) {
      r_val = TeamSpmv<MemberType, ArgTrans>::template invoke<ValuesViewType, IntView, xViewType, yViewType,
                                                              alphaViewType, betaViewType, dobeta>(
          member, alpha, values, row_ptr, colIndices, x, beta, y);
    } else if (std::is_same<ArgMode, Mode::TeamVector>::value) {
      r_val = TeamVectorSpmv<MemberType, ArgTrans>::template invoke<ValuesViewType, IntView, xViewType, yViewType,
                                                                    alphaViewType, betaViewType, dobeta>(
          member, alpha, values, row_ptr, colIndices, x, beta, y);
    }
    return r_val;
  }

  template <typename ValuesViewType, typename IntView, typename xViewType, typename yViewType, int dobeta>
  KOKKOS_INLINE_FUNCTION static int invoke(
      const MemberType &member,
      const typename KokkosKernels::ArithTraits<typename ValuesViewType::non_const_value_type>::mag_type &alpha,
      const ValuesViewType &values, const IntView &row_ptr, const IntView &colIndices, const xViewType &x,
      const typename KokkosKernels::ArithTraits<typename ValuesViewType::non_const_value_type>::mag_type &beta,
      const yViewType &y) {
    int r_val = 0;
    if (std::is_same<ArgMode, Mode::Serial>::value) {
      r_val = SerialSpmv<ArgTrans>::template invoke<ValuesViewType, IntView, xViewType, yViewType, dobeta>(
          alpha, values, row_ptr, colIndices, x, beta, y);
    } else if (std::is_same<ArgMode, Mode::Team>::value) {
      r_val = TeamSpmv<MemberType, ArgTrans>::template invoke<ValuesViewType, IntView, xViewType, yViewType, dobeta>(
          member, alpha, values, row_ptr, colIndices, x, beta, y);
    } else if (std::is_same<ArgMode, Mode::TeamVector>::value) {
      r_val =
          TeamVectorSpmv<MemberType, ArgTrans>::template invoke<ValuesViewType, IntView, xViewType, yViewType, dobeta>(
              member, alpha, values, row_ptr, colIndices, x, beta, y);
    }
    return r_val;
  }
};
}  // namespace KokkosBatched

#include "KokkosBatched_Spmv_Serial_Impl.hpp"
#include "KokkosBatched_Spmv_Team_Impl.hpp"
#include "KokkosBatched_Spmv_TeamVector_Impl.hpp"
#endif
