Csr/HIP/Morpheus_Multiply_Impl.hpp

 
#ifndef MORPHEUS_CSR_HIP_MULTIPLY_IMPL_HPP
#define MORPHEUS_CSR_HIP_MULTIPLY_IMPL_HPP
 
#include <Morpheus_Macros.hpp>
#if defined(MORPHEUS_ENABLE_HIP)
 
#include <Morpheus_SpaceTraits.hpp>
#include <Morpheus_FormatTraits.hpp>
#include <Morpheus_FormatTags.hpp>
#include <Morpheus_Spaces.hpp>
 
#include <impl/Morpheus_HIPUtils.hpp>
#include <impl/Csr/Kernels/Morpheus_Multiply_Impl.hpp>
 
namespace Morpheus {
namespace Impl {
 
// forward decl
template <typename Matrix, typename Vector>
void __spmv_csr_vector(const Matrix& A, const Vector& x, Vector& y,
                       const bool init);
 
template <typename Matrix, typename Vector>
void __spmv_csr_scalar(const Matrix& A, const Vector& x, Vector& y,
                       const bool init);
 
template <typename ExecSpace, typename Matrix, typename Vector>
inline void multiply(
    const Matrix& A, const Vector& x, Vector& y, const bool init,
    typename std::enable_if_t<
        Morpheus::is_csr_matrix_format_container_v<Matrix> &&
        Morpheus::is_dense_vector_format_container_v<Vector> &&
        Morpheus::has_custom_backend_v<ExecSpace> &&
        Morpheus::has_hip_execution_space_v<ExecSpace> &&
        Morpheus::has_access_v<ExecSpace, Matrix, Vector>>* = nullptr) {
  switch (A.options()) {
    case MATOPT_SHORT_ROWS: __spmv_csr_scalar(A, x, y, init); break;
    default: __spmv_csr_vector(A, x, y, init);
  }
}
 
template <typename Matrix, typename Vector>
void __spmv_csr_scalar(const Matrix& A, const Vector& x, Vector& y,
                       const bool init) {
  using size_type  = typename Matrix::size_type;
  using index_type = typename Matrix::index_type;
  using value_type = typename Matrix::value_type;
 
  const size_type BLOCK_SIZE = 256;
  const size_type NUM_BLOCKS = Impl::ceil_div<size_type>(A.nrows(), BLOCK_SIZE);
 
  const index_type* I = A.crow_offsets().data();
  const index_type* J = A.ccolumn_indices().data();
  const value_type* V = A.cvalues().data();
 
  const value_type* x_ptr = x.data();
  value_type* y_ptr       = y.data();
 
  if (init) {
    y.assign(y.size(), 0);
  }
 
  Morpheus::Impl::Kernels::spmv_csr_scalar_kernel<size_type, index_type,
                                                  value_type>
      <<<NUM_BLOCKS, BLOCK_SIZE, 0>>>(A.nrows(), I, J, V, x_ptr, y_ptr);
#if defined(DEBUG) || defined(MORPHEUS_DEBUG)
  getLastCudaError("spmv_csr_scalar_kernel: Kernel execution failed");
#endif
}
 
template <size_t THREADS_PER_VECTOR, typename Matrix, typename Vector>
void __spmv_csr_vector_dispatch(const Matrix& A, const Vector& x, Vector& y,
                                const bool init) {
  using size_type  = typename Matrix::size_type;
  using index_type = typename Matrix::index_type;
  using value_type = typename Matrix::value_type;
 
  const index_type* I = A.crow_offsets().data();
  const index_type* J = A.ccolumn_indices().data();
  const value_type* V = A.cvalues().data();
 
  const value_type* x_ptr = x.data();
  value_type* y_ptr       = y.data();
 
  const size_type THREADS_PER_BLOCK = 128;
  const size_type VECTORS_PER_BLOCK = THREADS_PER_BLOCK / THREADS_PER_VECTOR;
 
  if (init) {
    y.assign(y.size(), 0);
  }
 
  const size_type MAX_BLOCKS = max_active_blocks(
      Kernels::spmv_csr_vector_kernel<size_type, index_type, value_type,
                                      VECTORS_PER_BLOCK, THREADS_PER_VECTOR>,
      THREADS_PER_BLOCK, 0);
 
  const size_type NUM_BLOCKS = std::min<size_type>(
      MAX_BLOCKS, Impl::ceil_div<size_type>(A.nrows(), VECTORS_PER_BLOCK));
 
  Kernels::spmv_csr_vector_kernel<size_type, index_type, value_type,
                                  VECTORS_PER_BLOCK, THREADS_PER_VECTOR>
      <<<NUM_BLOCKS, THREADS_PER_BLOCK, 0>>>(A.nrows(), I, J, V, x_ptr, y_ptr);
 
#if defined(DEBUG) || defined(MORPHEUS_DEBUG)
  getLastCudaError("spmv_csr_vector_kernel: Kernel execution failed");
#endif
}
 
template <typename Matrix, typename Vector>
void __spmv_csr_vector(const Matrix& A, const Vector& x, Vector& y,
                       const bool init) {
  using size_type = typename Matrix::size_type;
 
  const size_type nnz_per_row = A.nnnz() / A.nrows();
 
  if (nnz_per_row <= 2) {
    __spmv_csr_vector_dispatch<2>(A, x, y, init);
    return;
  }
  if (nnz_per_row <= 4) {
    __spmv_csr_vector_dispatch<4>(A, x, y, init);
    return;
  }
  if (nnz_per_row <= 8) {
    __spmv_csr_vector_dispatch<8>(A, x, y, init);
    return;
  }
  if (nnz_per_row <= 16) {
    __spmv_csr_vector_dispatch<16>(A, x, y, init);
    return;
  }
 
  if (nnz_per_row <= 32) {
    __spmv_csr_vector_dispatch<32>(A, x, y, init);
    return;
  }
 
  __spmv_csr_vector_dispatch<64>(A, x, y, init);
}
 
}  // namespace Impl
}  // namespace Morpheus
 
#endif  // MORPHEUS_ENABLE_HIP
#endif  // MORPHEUS_CSR_HIP_MULTIPLY_IMPL_HPP