add extract_diagonal for Hybrid

common/matrix/hybrid_kernels.hpp.inc

@@ -139,4 +139,21 @@ __global__ __launch_bounds__(default_block_size) void add(
 }
 
 
-}  // namespace kernel
+template <typename ValueType, typename IndexType>
+__global__ __launch_bounds__(default_block_size) void coo_extract_diagonal(
+    size_type nnz, const ValueType *__restrict__ orig_values,
+    const IndexType *__restrict__ orig_row_idxs,
+    const IndexType *__restrict__ orig_col_idxs, size_type diag_stride,
+    ValueType *__restrict__ diag)
+{
+    const auto tidx = thread::get_thread_id_flat();
+
+    if (tidx < nnz) {
+        if (orig_row_idxs[tidx] == orig_col_idxs[tidx]) {
+            diag[diag_stride * orig_row_idxs[tidx]] = orig_values[tidx];
+        }
+    }
+}
+
+
+}  // namespace kernel

core/device_hooks/common_kernels.inc.cpp

@@ -730,6 +730,12 @@ GKO_NOT_COMPILED(GKO_HOOK_MODULE);
 GKO_INSTANTIATE_FOR_EACH_VALUE_AND_INDEX_TYPE(
     GKO_DECLARE_HYBRID_COUNT_NONZEROS_KERNEL);
 
+template <typename ValueType, typename IndexType>
+GKO_DECLARE_HYBRID_EXTRACT_DIAGONAL_KERNEL(ValueType, IndexType)
+GKO_NOT_COMPILED(GKO_HOOK_MODULE);
+GKO_INSTANTIATE_FOR_EACH_VALUE_AND_INDEX_TYPE(
+    GKO_DECLARE_HYBRID_EXTRACT_DIAGONAL_KERNEL);
+
 
 }  // namespace hybrid
 

core/matrix/coo.cpp

@@ -217,7 +217,7 @@ void Coo<ValueType, IndexType>::write(mat_data &data) const
 
 
 template <typename ValueType, typename IndexType>
-void Coo<ValueType, IndexType>::extract_diagonal(Dense<ValueType> *diag)
+void Coo<ValueType, IndexType>::extract_diagonal(Dense<ValueType> *diag) const
 {
     GKO_ASSERT_EQ(std::min(this->get_size()[0], this->get_size()[1]),
                   diag->get_size()[0]);

core/matrix/csr.cpp

@@ -500,7 +500,7 @@ bool Csr<ValueType, IndexType>::is_sorted_by_column_index() const
 
 
 template <typename ValueType, typename IndexType>
-void Csr<ValueType, IndexType>::extract_diagonal(Dense<ValueType> *diag)
+void Csr<ValueType, IndexType>::extract_diagonal(Dense<ValueType> *diag) const
 {
     GKO_ASSERT_EQ(std::min(this->get_size()[0], this->get_size()[1]),
                   diag->get_size()[0]);

core/matrix/dense.cpp

@@ -739,7 +739,7 @@ std::unique_ptr<LinOp> Dense<ValueType>::inverse_column_permute(
 
 
 template <typename ValueType>
-void Dense<ValueType>::extract_diagonal(Dense<ValueType> *diag)
+void Dense<ValueType>::extract_diagonal(Dense<ValueType> *diag) const
 {
     GKO_ASSERT_EQ(std::min(this->get_size()[0], this->get_size()[1]),
                   diag->get_size()[0]);

core/matrix/ell.cpp

@@ -235,7 +235,7 @@ void Ell<ValueType, IndexType>::write(mat_data &data) const
 
 
 template <typename ValueType, typename IndexType>
-void Ell<ValueType, IndexType>::extract_diagonal(Dense<ValueType> *diag)
+void Ell<ValueType, IndexType>::extract_diagonal(Dense<ValueType> *diag) const
 {
     GKO_ASSERT_EQ(std::min(this->get_size()[0], this->get_size()[1]),
                   diag->get_size()[0]);

core/matrix/hybrid.cpp

@@ -54,6 +54,7 @@ namespace hybrid {
 GKO_REGISTER_OPERATION(convert_to_dense, hybrid::convert_to_dense);
 GKO_REGISTER_OPERATION(convert_to_csr, hybrid::convert_to_csr);
 GKO_REGISTER_OPERATION(count_nonzeros, hybrid::count_nonzeros);
+GKO_REGISTER_OPERATION(extract_diagonal, hybrid::extract_diagonal);
 
 
 }  // namespace hybrid
@@ -263,6 +264,19 @@ void Hybrid<ValueType, IndexType>::write(mat_data &data) const
 }
 
 
+template <typename ValueType, typename IndexType>
+void Hybrid<ValueType, IndexType>::extract_diagonal(
+    Dense<ValueType> *diag) const
+{
+    GKO_ASSERT_EQ(std::min(this->get_size()[0], this->get_size()[1]),
+                  diag->get_size()[0]);
+    GKO_ASSERT_EQ(diag->get_size()[1], 1);
+
+    auto exec = this->get_executor();
+    exec->run(hybrid::make_extract_diagonal(this, diag));
+}
+
+
 #define GKO_DECLARE_HYBRID_MATRIX(ValueType, IndexType) \
     class Hybrid<ValueType, IndexType>
 GKO_INSTANTIATE_FOR_EACH_VALUE_AND_INDEX_TYPE(GKO_DECLARE_HYBRID_MATRIX);

core/matrix/hybrid_kernels.hpp

@@ -59,13 +59,20 @@ namespace kernels {
                         const matrix::Hybrid<ValueType, IndexType> *source, \
                         size_type *result)
 
+#define GKO_DECLARE_HYBRID_EXTRACT_DIAGONAL_KERNEL(ValueType, IndexType)    \
+    void extract_diagonal(std::shared_ptr<const DefaultExecutor> exec,      \
+                          const matrix::Hybrid<ValueType, IndexType> *orig, \
+                          matrix::Dense<ValueType> *diag)
+
 #define GKO_DECLARE_ALL_AS_TEMPLATES                                  \
     template <typename ValueType, typename IndexType>                 \
     GKO_DECLARE_HYBRID_CONVERT_TO_DENSE_KERNEL(ValueType, IndexType); \
     template <typename ValueType, typename IndexType>                 \
     GKO_DECLARE_HYBRID_CONVERT_TO_CSR_KERNEL(ValueType, IndexType);   \
     template <typename ValueType, typename IndexType>                 \
-    GKO_DECLARE_HYBRID_COUNT_NONZEROS_KERNEL(ValueType, IndexType)
+    GKO_DECLARE_HYBRID_COUNT_NONZEROS_KERNEL(ValueType, IndexType);   \
+    template <typename ValueType, typename IndexType>                 \
+    GKO_DECLARE_HYBRID_EXTRACT_DIAGONAL_KERNEL(ValueType, IndexType)
 
 
 namespace omp {

core/matrix/sellp.cpp

@@ -287,7 +287,7 @@ void Sellp<ValueType, IndexType>::write(mat_data &data) const
 
 
 template <typename ValueType, typename IndexType>
-void Sellp<ValueType, IndexType>::extract_diagonal(Dense<ValueType> *diag)
+void Sellp<ValueType, IndexType>::extract_diagonal(Dense<ValueType> *diag) const
 {
     GKO_ASSERT_EQ(std::min(this->get_size()[0], this->get_size()[1]),
                   diag->get_size()[0]);

cuda/matrix/coo_kernels.cu

@@ -250,9 +250,9 @@ void extract_diagonal(std::shared_ptr<const CudaExecutor> exec,
                       const matrix::Coo<ValueType, IndexType> *orig,
                       matrix::Dense<ValueType> *diag)
 {
-    auto nnz = orig->get_num_stored_elements();
-    auto diag_size = diag->get_size()[0];
-    auto diag_stride = diag->get_stride();
+    const auto nnz = orig->get_num_stored_elements();
+    const auto diag_size = diag->get_size()[0];
+    const auto diag_stride = diag->get_stride();
     auto num_blocks = ceildiv(diag_size, default_block_size);
 
     const auto orig_values = orig->get_const_values();

cuda/matrix/csr_kernels.cu

@@ -1035,9 +1035,9 @@ void extract_diagonal(std::shared_ptr<const CudaExecutor> exec,
                       const matrix::Csr<ValueType, IndexType> *orig,
                       matrix::Dense<ValueType> *diag)
 {
-    auto nnz = orig->get_num_stored_elements();
-    auto diag_size = diag->get_size()[0];
-    auto diag_stride = diag->get_stride();
+    const auto nnz = orig->get_num_stored_elements();
+    const auto diag_size = diag->get_size()[0];
+    const auto diag_stride = diag->get_stride();
     auto num_blocks =
         ceildiv(config::warp_size * diag_size, default_block_size);
 

cuda/matrix/ell_kernels.cu

@@ -367,10 +367,10 @@ void extract_diagonal(std::shared_ptr<const CudaExecutor> exec,
                       const matrix::Ell<ValueType, IndexType> *orig,
                       matrix::Dense<ValueType> *diag)
 {
-    auto max_nnz_per_row = orig->get_num_stored_elements_per_row();
-    auto orig_stride = orig->get_stride();
-    auto diag_size = diag->get_size()[0];
-    auto diag_stride = diag->get_stride();
+    const auto max_nnz_per_row = orig->get_num_stored_elements_per_row();
+    const auto orig_stride = orig->get_stride();
+    const auto diag_size = diag->get_size()[0];
+    const auto diag_stride = diag->get_stride();
     auto num_blocks = ceildiv(diag_size, default_block_size);
 
     const auto orig_values = orig->get_const_values();

cuda/matrix/hybrid_kernels.cu

@@ -179,6 +179,31 @@ GKO_INSTANTIATE_FOR_EACH_VALUE_AND_INDEX_TYPE(
     GKO_DECLARE_HYBRID_COUNT_NONZEROS_KERNEL);
 
 
+template <typename ValueType, typename IndexType>
+void extract_diagonal(std::shared_ptr<const CudaExecutor> exec,
+                      const matrix::Hybrid<ValueType, IndexType> *orig,
+                      matrix::Dense<ValueType> *diag)
+{
+    gko::kernels::cuda::ell::extract_diagonal(exec, orig->get_ell(), diag);
+
+    const auto coo_row_idxs = orig->get_const_coo_row_idxs();
+    const auto coo_col_idxs = orig->get_const_coo_col_idxs();
+    const auto coo_values = orig->get_const_coo_values();
+    const auto coo_nnz = orig->get_coo_num_stored_elements();
+
+    const auto diag_stride = diag->get_stride();
+    const auto num_blocks = ceildiv(coo_nnz, default_block_size);
+    auto diag_values = diag->get_values();
+
+    kernel::coo_extract_diagonal<<<num_blocks, default_block_size>>>(
+        coo_nnz, as_cuda_type(coo_values), as_cuda_type(coo_row_idxs),
+        as_cuda_type(coo_col_idxs), diag_stride, as_cuda_type(diag_values));
+}
+
+GKO_INSTANTIATE_FOR_EACH_VALUE_AND_INDEX_TYPE(
+    GKO_DECLARE_HYBRID_EXTRACT_DIAGONAL_KERNEL);
+
+
 }  // namespace hybrid
 }  // namespace cuda
 }  // namespace kernels

cuda/test/matrix/hybrid_kernels.cpp

@@ -219,4 +219,21 @@ TEST_F(Hybrid, MoveToCsrIsEquivalentToRef)
 }
 
 
+TEST_F(Hybrid, ExtractDiagonalIsquivalentToRef)
+{
+    set_up_apply_data();
+
+    auto diag_size = std::min(mtx->get_size()[0], mtx->get_size()[1]);
+
+    auto diag = gko::matrix::Dense<>::create(mtx->get_executor(),
+                                             gko::dim<2>(diag_size, 1));
+    auto ddiag = gko::matrix::Dense<>::create(dmtx->get_executor(),
+                                              gko::dim<2>(diag_size, 1));
+    mtx->extract_diagonal(lend(diag));
+    dmtx->extract_diagonal(lend(ddiag));
+
+    GKO_ASSERT_MTX_NEAR(diag.get(), ddiag.get(), 0);
+}
+
+
 }  // namespace

hip/matrix/coo_kernels.hip.cpp

@@ -262,9 +262,9 @@ void extract_diagonal(std::shared_ptr<const HipExecutor> exec,
                       const matrix::Coo<ValueType, IndexType> *orig,
                       matrix::Dense<ValueType> *diag)
 {
-    auto nnz = orig->get_num_stored_elements();
-    auto diag_size = diag->get_size()[0];
-    auto diag_stride = diag->get_stride();
+    const auto nnz = orig->get_num_stored_elements();
+    const auto diag_size = diag->get_size()[0];
+    const auto diag_stride = diag->get_stride();
     auto num_blocks = ceildiv(diag_size, default_block_size);
 
     const auto orig_values = orig->get_const_values();

hip/matrix/csr_kernels.hip.cpp

@@ -1147,10 +1147,10 @@ void extract_diagonal(std::shared_ptr<const HipExecutor> exec,
                       const matrix::Csr<ValueType, IndexType> *orig,
                       matrix::Dense<ValueType> *diag)
 {
-    auto nnz = orig->get_num_stored_elements();
-    auto diag_size = diag->get_size()[0];
-    auto diag_stride = diag->get_stride();
-    auto num_blocks =
+    const auto nnz = orig->get_num_stored_elements();
+    const auto diag_size = diag->get_size()[0];
+    const auto diag_stride = diag->get_stride();
+    const auto num_blocks =
         ceildiv(config::warp_size * diag_size, default_block_size);
 
     const auto orig_values = orig->get_const_values();

hip/matrix/ell_kernels.hip.cpp

@@ -376,10 +376,10 @@ void extract_diagonal(std::shared_ptr<const HipExecutor> exec,
                       const matrix::Ell<ValueType, IndexType> *orig,
                       matrix::Dense<ValueType> *diag)
 {
-    auto max_nnz_per_row = orig->get_num_stored_elements_per_row();
-    auto orig_stride = orig->get_stride();
-    auto diag_size = diag->get_size()[0];
-    auto diag_stride = diag->get_stride();
+    const auto max_nnz_per_row = orig->get_num_stored_elements_per_row();
+    const auto orig_stride = orig->get_stride();
+    const auto diag_size = diag->get_size()[0];
+    const auto diag_stride = diag->get_stride();
     auto num_blocks = ceildiv(diag_size, default_block_size);
 
     const auto orig_values = orig->get_const_values();

hip/matrix/hybrid_kernels.hip.cpp

@@ -186,6 +186,32 @@ GKO_INSTANTIATE_FOR_EACH_VALUE_AND_INDEX_TYPE(
     GKO_DECLARE_HYBRID_COUNT_NONZEROS_KERNEL);
 
 
+template <typename ValueType, typename IndexType>
+void extract_diagonal(std::shared_ptr<const HipExecutor> exec,
+                      const matrix::Hybrid<ValueType, IndexType> *orig,
+                      matrix::Dense<ValueType> *diag)
+{
+    gko::kernels::hip::ell::extract_diagonal(exec, orig->get_ell(), diag);
+
+    const auto coo_row_idxs = orig->get_const_coo_row_idxs();
+    const auto coo_col_idxs = orig->get_const_coo_col_idxs();
+    const auto coo_values = orig->get_const_coo_values();
+    const auto coo_nnz = orig->get_coo_num_stored_elements();
+
+    const auto diag_stride = diag->get_stride();
+    const auto num_blocks = ceildiv(coo_nnz, default_block_size);
+    auto diag_values = diag->get_values();
+
+    hipLaunchKernelGGL(num_blocks, default_block_size, 0, 0, coo_nnz,
+                       as_cuda_type(coo_values), as_cuda_type(coo_row_idxs),
+                       as_cuda_type(coo_col_idxs), diag_stride,
+                       as_cuda_type(diag_values));
+}
+
+GKO_INSTANTIATE_FOR_EACH_VALUE_AND_INDEX_TYPE(
+    GKO_DECLARE_HYBRID_EXTRACT_DIAGONAL_KERNEL);
+
+
 }  // namespace hybrid
 }  // namespace hip
 }  // namespace kernels

hip/test/matrix/hybrid_kernels.hip.cpp

@@ -219,4 +219,21 @@ TEST_F(Hybrid, MoveToCsrIsEquivalentToRef)
 }
 
 
+TEST_F(Hybrid, ExtractDiagonalIsquivalentToRef)
+{
+    set_up_apply_data();
+
+    auto diag_size = std::min(mtx->get_size()[0], mtx->get_size()[1]);
+
+    auto diag = gko::matrix::Dense<>::create(mtx->get_executor(),
+                                             gko::dim<2>(diag_size, 1));
+    auto ddiag = gko::matrix::Dense<>::create(dmtx->get_executor(),
+                                              gko::dim<2>(diag_size, 1));
+    mtx->extract_diagonal(lend(diag));
+    dmtx->extract_diagonal(lend(ddiag));
+
+    GKO_ASSERT_MTX_NEAR(diag.get(), ddiag.get(), 0);
+}
+
+
 }  // namespace

include/ginkgo/core/matrix/coo.hpp

@@ -251,7 +251,7 @@ class Coo : public EnableLinOp<Coo<ValueType, IndexType>>,
      *
      * @param diag  the vector into which the diagonal will be written
      */
-    void extract_diagonal(Dense<value_type> *diag);
+    void extract_diagonal(Dense<value_type> *diag) const;
 
 protected:
     /**

include/ginkgo/core/matrix/csr.hpp

@@ -748,7 +748,7 @@ class Csr : public EnableLinOp<Csr<ValueType, IndexType>>,
      *
      * @param diag  the vector into which the diagonal will be written
      */
-    void extract_diagonal(Dense<value_type> *diag);
+    void extract_diagonal(Dense<value_type> *diag) const;
 
 protected:
     /**

include/ginkgo/core/matrix/dense.hpp

@@ -262,7 +262,7 @@ class Dense : public EnableLinOp<Dense<ValueType>>,
      *
      * @param diag  the vector into which the diagonal will be written
      */
-    void extract_diagonal(Dense<ValueType> *diag);
+    void extract_diagonal(Dense<ValueType> *diag) const;
 
 
     /**

include/ginkgo/core/matrix/ell.hpp

@@ -223,7 +223,7 @@ class Ell : public EnableLinOp<Ell<ValueType, IndexType>>,
      *
      * @param diag  the vector into which the diagonal will be written
      */
-    void extract_diagonal(Dense<value_type> *diag);
+    void extract_diagonal(Dense<value_type> *diag) const;
 
 protected:
     /**

include/ginkgo/core/matrix/hybrid.hpp

@@ -583,6 +583,13 @@ class Hybrid
         return *this;
     }
 
+    /**
+     * Extracts the diagonal entries of the matrix into a vector.
+     *
+     * @param diag  the vector into which the diagonal will be written
+     */
+    void extract_diagonal(Dense<value_type> *diag) const;
+
 protected:
     /**
      * Creates an uninitialized Hybrid matrix of specified method.

include/ginkgo/core/matrix/sellp.hpp

@@ -275,7 +275,7 @@ class Sellp : public EnableLinOp<Sellp<ValueType, IndexType>>,
      *
      * @param diag  the vector into which the diagonal will be written
      */
-    void extract_diagonal(Dense<value_type> *diag);
+    void extract_diagonal(Dense<value_type> *diag) const;
 
 protected:
     /**