Add extract_diagonal for SELLP

ginkgo-project · Jun 9, 2020 · 385e508 · 385e508
1 parent 9fe6fa4
commit 385e508
Show file tree

Hide file tree

Showing 13 changed files with 341 additions and 6 deletions.
diff --git a/common/matrix/sellp_kernels.hpp.inc b/common/matrix/sellp_kernels.hpp.inc
@@ -196,4 +196,33 @@ __global__ __launch_bounds__(default_block_size) void fill_in_csr(
 }
 
 
-}  // namespace kernel
+template <typename ValueType, typename IndexType>
+__global__ __launch_bounds__(default_block_size) void extract_diagonal(
+    size_type diag_size, size_type slice_size,
+    const size_type *__restrict__ orig_slice_sets,
+    const ValueType *__restrict__ orig_values,
+    const IndexType *__restrict__ orig_col_idxs, size_type diag_stride,
+    ValueType *__restrict__ diag)
+{
+    constexpr auto warp_size = config::warp_size;
+    auto warp_tile =
+        group::tiled_partition<warp_size>(group::this_thread_block());
+    const auto slice_id = thread::get_subwarp_id_flat<warp_size>();
+    const auto tid_in_warp = warp_tile.thread_rank();
+
+    for (size_type sellp_ind =
+             orig_slice_sets[slice_id] * slice_size + tid_in_warp;
+         sellp_ind < orig_slice_sets[slice_id + 1] * slice_size;
+         sellp_ind += warp_size) {
+        auto global_row = slice_id * slice_size + sellp_ind % slice_size;
+        if (global_row < diag_size) {
+            if (orig_col_idxs[sellp_ind] == global_row &&
+                orig_values[sellp_ind] != zero<ValueType>()) {
+                diag[global_row * diag_stride] = orig_values[sellp_ind];
+            }
+        }
+    }
+}
+
+
+}  // namespace kernel
diff --git a/core/device_hooks/common_kernels.inc.cpp b/core/device_hooks/common_kernels.inc.cpp
@@ -766,6 +766,12 @@ GKO_NOT_COMPILED(GKO_HOOK_MODULE);
 GKO_INSTANTIATE_FOR_EACH_VALUE_AND_INDEX_TYPE(
     GKO_DECLARE_SELLP_COUNT_NONZEROS_KERNEL);
 
+template <typename ValueType, typename IndexType>
+GKO_DECLARE_SELLP_EXTRACT_DIAGONAL_KERNEL(ValueType, IndexType)
+GKO_NOT_COMPILED(GKO_HOOK_MODULE);
+GKO_INSTANTIATE_FOR_EACH_VALUE_AND_INDEX_TYPE(
+    GKO_DECLARE_SELLP_EXTRACT_DIAGONAL_KERNEL);
+
 
 }  // namespace sellp
 

diff --git a/core/matrix/sellp.cpp b/core/matrix/sellp.cpp
@@ -55,6 +55,7 @@ GKO_REGISTER_OPERATION(advanced_spmv, sellp::advanced_spmv);
 GKO_REGISTER_OPERATION(convert_to_dense, sellp::convert_to_dense);
 GKO_REGISTER_OPERATION(convert_to_csr, sellp::convert_to_csr);
 GKO_REGISTER_OPERATION(count_nonzeros, sellp::count_nonzeros);
+GKO_REGISTER_OPERATION(extract_diagonal, sellp::extract_diagonal);
 
 
 }  // namespace sellp
@@ -285,6 +286,18 @@ void Sellp<ValueType, IndexType>::write(mat_data &data) const
 }
 
 
+template <typename ValueType, typename IndexType>
+void Sellp<ValueType, IndexType>::extract_diagonal(Dense<ValueType> *diag)
+{
+    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(sellp::make_extract_diagonal(this, diag));
+}
+
+
 #define GKO_DECLARE_SELLP_MATRIX(ValueType, IndexType) \
     class Sellp<ValueType, IndexType>
 GKO_INSTANTIATE_FOR_EACH_VALUE_AND_INDEX_TYPE(GKO_DECLARE_SELLP_MATRIX);

diff --git a/core/matrix/sellp_kernels.hpp b/core/matrix/sellp_kernels.hpp
@@ -73,6 +73,11 @@ namespace kernels {
                         const matrix::Sellp<ValueType, IndexType> *source, \
                         size_type *result)
 
+#define GKO_DECLARE_SELLP_EXTRACT_DIAGONAL_KERNEL(ValueType, IndexType)    \
+    void extract_diagonal(std::shared_ptr<const DefaultExecutor> exec,     \
+                          const matrix::Sellp<ValueType, IndexType> *orig, \
+                          matrix::Dense<ValueType> *diag)
+
 #define GKO_DECLARE_ALL_AS_TEMPLATES                                 \
     template <typename ValueType, typename IndexType>                \
     GKO_DECLARE_SELLP_SPMV_KERNEL(ValueType, IndexType);             \
@@ -83,7 +88,9 @@ namespace kernels {
     template <typename ValueType, typename IndexType>                \
     GKO_DECLARE_SELLP_CONVERT_TO_CSR_KERNEL(ValueType, IndexType);   \
     template <typename ValueType, typename IndexType>                \
-    GKO_DECLARE_SELLP_COUNT_NONZEROS_KERNEL(ValueType, IndexType)
+    GKO_DECLARE_SELLP_COUNT_NONZEROS_KERNEL(ValueType, IndexType);   \
+    template <typename ValueType, typename IndexType>                \
+    GKO_DECLARE_SELLP_EXTRACT_DIAGONAL_KERNEL(ValueType, IndexType)
 
 
 namespace omp {

diff --git a/cuda/matrix/sellp_kernels.cu b/cuda/matrix/sellp_kernels.cu
@@ -215,6 +215,38 @@ GKO_INSTANTIATE_FOR_EACH_VALUE_AND_INDEX_TYPE(
     GKO_DECLARE_SELLP_COUNT_NONZEROS_KERNEL);
 
 
+template <typename ValueType, typename IndexType>
+void extract_diagonal(std::shared_ptr<const CudaExecutor> exec,
+                      const matrix::Sellp<ValueType, IndexType> *orig,
+                      matrix::Dense<ValueType> *diag)
+{
+    const auto diag_size = diag->get_size()[0];
+    const auto diag_stride = diag->get_stride();
+    const auto slice_size = orig->get_slice_size();
+    const auto slice_num = ceildiv(orig->get_size()[0], slice_size);
+    auto num_blocks = ceildiv(diag_size, default_block_size);
+
+    const auto orig_slice_sets = orig->get_const_slice_sets();
+    const auto orig_values = orig->get_const_values();
+    const auto orig_col_idxs = orig->get_const_col_idxs();
+    auto diag_values = diag->get_values();
+
+    kernel::initialize_zero_dense<<<num_blocks, default_block_size>>>(
+        diag->get_size()[0], diag->get_size()[1], diag_stride,
+        as_cuda_type(diag_values));
+
+    num_blocks = ceildiv(slice_num * config::warp_size, default_block_size);
+
+    kernel::extract_diagonal<<<num_blocks, default_block_size>>>(
+        diag_size, slice_size, as_cuda_type(orig_slice_sets),
+        as_cuda_type(orig_values), as_cuda_type(orig_col_idxs), diag_stride,
+        as_cuda_type(diag_values));
+}
+
+GKO_INSTANTIATE_FOR_EACH_VALUE_AND_INDEX_TYPE(
+    GKO_DECLARE_SELLP_EXTRACT_DIAGONAL_KERNEL);
+
+
 }  // namespace sellp
 }  // namespace cuda
 }  // namespace kernels

diff --git a/cuda/test/matrix/sellp_kernels.cpp b/cuda/test/matrix/sellp_kernels.cpp
@@ -295,4 +295,38 @@ TEST_F(Sellp, CountNonzerosIsEquivalentToRef)
 }
 
 
+TEST_F(Sellp, ExtractDiagonalIsquivalentToRef)
+{
+    set_up_apply_matrix();
+
+    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);
+}
+
+
+TEST_F(Sellp, ExtractDiagonalWithSliceSizeAndStrideFactorIsquivalentToRef)
+{
+    set_up_apply_matrix(32, 2);
+
+    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
diff --git a/hip/matrix/sellp_kernels.hip.cpp b/hip/matrix/sellp_kernels.hip.cpp
@@ -221,6 +221,40 @@ GKO_INSTANTIATE_FOR_EACH_VALUE_AND_INDEX_TYPE(
     GKO_DECLARE_SELLP_COUNT_NONZEROS_KERNEL);
 
 
+template <typename ValueType, typename IndexType>
+void extract_diagonal(std::shared_ptr<const HipExecutor> exec,
+                      const matrix::Sellp<ValueType, IndexType> *orig,
+                      matrix::Dense<ValueType> *diag)
+{
+    const auto diag_size = diag->get_size()[0];
+    const auto diag_stride = diag->get_stride();
+    const auto slice_size = orig->get_slice_size();
+    const auto slice_num = ceildiv(orig->get_size()[0], slice_size);
+    auto num_blocks = ceildiv(diag_size, default_block_size);
+
+    const auto orig_slice_sets = orig->get_const_slice_sets();
+    const auto orig_values = orig->get_const_values();
+    const auto orig_col_idxs = orig->get_const_col_idxs();
+    auto diag_values = diag->get_values();
+
+    hipLaunchKernelGGL(kernel::initialize_zero_dense, dim3(num_blocks),
+                       dim3(default_block_size), 0, 0, diag->get_size()[0],
+                       diag->get_size()[1], diag_stride,
+                       as_hip_type(diag_values));
+
+    num_blocks = ceildiv(slice_num * config::warp_size, default_block_size);
+
+    hipLaunchKernelGGL(kernel::extract_diagonal, dim3(num_blocks),
+                       dim3(default_block_size), 0, 0, diag_size, slice_size,
+                       as_cuda_type(orig_slice_sets), as_cuda_type(orig_values),
+                       as_cuda_type(orig_col_idxs), diag_stride,
+                       as_cuda_type(diag_values));
+}
+
+GKO_INSTANTIATE_FOR_EACH_VALUE_AND_INDEX_TYPE(
+    GKO_DECLARE_SELLP_EXTRACT_DIAGONAL_KERNEL);
+
+
 }  // namespace sellp
 }  // namespace hip
 }  // namespace kernels

diff --git a/hip/test/matrix/sellp_kernels.hip.cpp b/hip/test/matrix/sellp_kernels.hip.cpp
@@ -294,4 +294,38 @@ TEST_F(Sellp, CountNonzerosIsEquivalentToRef)
 }
 
 
+TEST_F(Sellp, ExtractDiagonalIsquivalentToRef)
+{
+    set_up_apply_matrix();
+
+    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);
+}
+
+
+TEST_F(Sellp, ExtractDiagonalWithSliceSizeAndStrideFactorIsquivalentToRef)
+{
+    set_up_apply_matrix(32, 2);
+
+    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
diff --git a/include/ginkgo/core/matrix/sellp.hpp b/include/ginkgo/core/matrix/sellp.hpp
@@ -235,8 +235,8 @@ class Sellp : public EnableLinOp<Sellp<ValueType, IndexType>>,
     /**
      * @copydoc Sellp::val_at(size_type, size_type, size_type)
      */
-    value_type val_at(size_type row, size_type slice_set, size_type idx) const
-        noexcept
+    value_type val_at(size_type row, size_type slice_set,
+                      size_type idx) const noexcept
     {
         return values_
             .get_const_data()[this->linearize_index(row, slice_set, idx)];
@@ -263,13 +263,20 @@ class Sellp : public EnableLinOp<Sellp<ValueType, IndexType>>,
     /**
      * @copydoc Sellp::col_at(size_type, size_type, size_type)
      */
-    index_type col_at(size_type row, size_type slice_set, size_type idx) const
-        noexcept
+    index_type col_at(size_type row, size_type slice_set,
+                      size_type idx) const noexcept
     {
         return this
             ->get_const_col_idxs()[this->linearize_index(row, slice_set, idx)];
     }
 
+    /**
+     * 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);
+
 protected:
     /**
      * Creates an uninitialized Sellp matrix of the specified size.

diff --git a/omp/matrix/sellp_kernels.cpp b/omp/matrix/sellp_kernels.cpp
@@ -178,6 +178,43 @@ GKO_INSTANTIATE_FOR_EACH_VALUE_AND_INDEX_TYPE(
     GKO_DECLARE_SELLP_COUNT_NONZEROS_KERNEL);
 
 
+template <typename ValueType, typename IndexType>
+void extract_diagonal(std::shared_ptr<const OmpExecutor> exec,
+                      const matrix::Sellp<ValueType, IndexType> *orig,
+                      matrix::Dense<ValueType> *diag)
+{
+    const auto diag_size = diag->get_size()[0];
+    const auto slice_size = orig->get_slice_size();
+    const auto slice_num = ceildiv(orig->get_size()[0], slice_size);
+
+    const auto orig_values = orig->get_const_values();
+    const auto orig_slice_sets = orig->get_const_slice_sets();
+    const auto orig_slice_lengths = orig->get_const_slice_lengths();
+    const auto orig_col_idxs = orig->get_const_col_idxs();
+
+#pragma omp parallel for
+    for (size_type slice = 0; slice < slice_num; slice++) {
+        for (size_type row = 0;
+             row < slice_size && slice_size * slice + row < diag_size; row++) {
+            auto global_row = slice_size * slice + row;
+            diag->at(global_row, 0) = zero<ValueType>();
+            for (size_type i = 0; i < orig_slice_lengths[slice]; i++) {
+                if (orig->col_at(row, orig_slice_sets[slice], i) ==
+                        global_row &&
+                    orig->val_at(row, orig_slice_sets[slice], i) !=
+                        zero<ValueType>()) {
+                    diag->at(global_row, 0) =
+                        orig->val_at(row, orig_slice_sets[slice], i);
+                }
+            }
+        }
+    }
+}
+
+GKO_INSTANTIATE_FOR_EACH_VALUE_AND_INDEX_TYPE(
+    GKO_DECLARE_SELLP_EXTRACT_DIAGONAL_KERNEL);
+
+
 }  // namespace sellp
 }  // namespace omp
 }  // namespace kernels

diff --git a/omp/test/matrix/sellp_kernels.cpp b/omp/test/matrix/sellp_kernels.cpp
@@ -243,4 +243,38 @@ TEST_F(Sellp, MoveToDenseIsEquivalentToRef)
 }
 
 
+TEST_F(Sellp, 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);
+}
+
+
+TEST_F(Sellp, ExtractDiagonalWithSliceSizeAndStrideFactorIsquivalentToRef)
+{
+    set_up_apply_data(32, 4, 0, 3);
+
+    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