Coo/Kernels/Morpheus_MatrixOperations_Impl.hpp

 
#ifndef MORPHEUS_COO_KERNELS_MATRIXOPERATIONS_IMPL_HPP
#define MORPHEUS_COO_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_coo_diagonal_kernel(const SizeType nnnz,
                                           const IndexType* I,
                                           const IndexType* J, ValueType* V,
                                           const ValueType* diagonal) {
  const SizeType tid = blockDim.x * blockIdx.x + threadIdx.x;
  if (tid >= nnnz) return;
 
  if (I[tid] == J[tid]) V[tid] = diagonal[J[tid]];
}
 
}  // namespace Kernels
}  // namespace Impl
}  // namespace Morpheus
 
#endif  // MORPHEUS_ENABLE_CUDA || MORPHEUS_ENABLE_HIP
#endif  // MORPHEUS_COO_KERNELS_MATRIXOPERATIONS_IMPL_HPP