Csr/Kernels/Morpheus_MatrixOperations_Impl.hpp

 
#ifndef MORPHEUS_CSR_KERNELS_MATRIXOPERATIONS_IMPL_HPP
#define MORPHEUS_CSR_KERNELS_MATRIXOPERATIONS_IMPL_HPP
 
#include <Morpheus_Macros.hpp>
#if defined(MORPHEUS_ENABLE_CUDA) || defined(MORPHEUS_ENABLE_HIP)
 
namespace Morpheus {
namespace Impl {
 
namespace Kernels {
template <typename ValueType, typename IndexType, typename SizeType>
__global__ void update_csr_diagonal_kernel(const SizeType nrows,
                                           const IndexType* Ap,
                                           const IndexType* Aj, ValueType* Ax,
                                           const ValueType* diagonal) {
  const SizeType thread_id = blockDim.x * blockIdx.x + threadIdx.x;
  const SizeType grid_size = gridDim.x * blockDim.x;
 
  for (SizeType row = thread_id; row < nrows; row += grid_size) {
    const IndexType row_start = Ap[row];
    const IndexType row_end   = Ap[row + 1];
 
    for (IndexType jj = row_start; jj < row_end; jj++) {
      if (Aj[jj] == row) {
        Ax[jj] = diagonal[row];
        break;
      }
    }
  }
}
 
}  // namespace Kernels
}  // namespace Impl
}  // namespace Morpheus
 
#endif  // MORPHEUS_ENABLE_CUDA || MORPHEUS_ENABLE_HIP
#endif  // MORPHEUS_CSR_KERNELS_MATRIXOPERATIONS_IMPL_HPP