rebase and use diagonal matrix format as output to extract_diagonal

ginkgo-project · Aug 19, 2020 · 302beb8 · 302beb8
1 parent 48e1189
commit 302beb8
Show file tree

Hide file tree

Showing 59 changed files with 145 additions and 74 deletions.
diff --git a/core/matrix/coo.cpp b/core/matrix/coo.cpp
@@ -219,13 +219,13 @@ void Coo<ValueType, IndexType>::write(mat_data &data) const
 
 
 template <typename ValueType, typename IndexType>
-std::unique_ptr<Dense<ValueType>> Coo<ValueType, IndexType>::extract_diagonal()
-    const
+std::unique_ptr<Diagonal<ValueType>>
+Coo<ValueType, IndexType>::extract_diagonal() const
 {
     auto exec = this->get_executor();
 
     const auto diag_size = std::min(this->get_size()[0], this->get_size()[1]);
-    auto diag = Dense<ValueType>::create(exec, dim<2>(diag_size, 1));
+    auto diag = Diagonal<ValueType>::create(exec, diag_size);
     exec->run(coo::make_fill_array(diag->get_values(), diag->get_size()[0],
                                    zero<ValueType>()));
     exec->run(coo::make_extract_diagonal(this, lend(diag)));

diff --git a/core/matrix/coo_kernels.hpp b/core/matrix/coo_kernels.hpp
@@ -40,6 +40,7 @@ OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 #include <ginkgo/core/base/types.hpp>
 #include <ginkgo/core/matrix/csr.hpp>
 #include <ginkgo/core/matrix/dense.hpp>
+#include <ginkgo/core/matrix/diagonal.hpp>
 
 
 namespace gko {
@@ -84,7 +85,7 @@ namespace kernels {
 #define GKO_DECLARE_COO_EXTRACT_DIAGONAL_KERNEL(ValueType, IndexType)    \
     void extract_diagonal(std::shared_ptr<const DefaultExecutor> exec,   \
                           const matrix::Coo<ValueType, IndexType> *orig, \
-                          matrix::Dense<ValueType> *diag)
+                          matrix::Diagonal<ValueType> *diag)
 
 #define GKO_DECLARE_ALL_AS_TEMPLATES                               \
     template <typename ValueType, typename IndexType>              \

diff --git a/core/matrix/csr.cpp b/core/matrix/csr.cpp
@@ -468,13 +468,13 @@ bool Csr<ValueType, IndexType>::is_sorted_by_column_index() const
 
 
 template <typename ValueType, typename IndexType>
-std::unique_ptr<Dense<ValueType>> Csr<ValueType, IndexType>::extract_diagonal()
-    const
+std::unique_ptr<Diagonal<ValueType>>
+Csr<ValueType, IndexType>::extract_diagonal() const
 {
     auto exec = this->get_executor();
 
     const auto diag_size = std::min(this->get_size()[0], this->get_size()[1]);
-    auto diag = Dense<ValueType>::create(exec, dim<2>(diag_size, 1));
+    auto diag = Diagonal<ValueType>::create(exec, diag_size);
     exec->run(csr::make_fill_array(diag->get_values(), diag->get_size()[0],
                                    zero<ValueType>()));
     exec->run(csr::make_extract_diagonal(this, lend(diag)));

diff --git a/core/matrix/csr_kernels.hpp b/core/matrix/csr_kernels.hpp
@@ -41,6 +41,7 @@ OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 #include <ginkgo/core/base/types.hpp>
 #include <ginkgo/core/matrix/coo.hpp>
 #include <ginkgo/core/matrix/dense.hpp>
+#include <ginkgo/core/matrix/diagonal.hpp>
 #include <ginkgo/core/matrix/ell.hpp>
 #include <ginkgo/core/matrix/hybrid.hpp>
 #include <ginkgo/core/matrix/sellp.hpp>
@@ -177,7 +178,7 @@ namespace kernels {
 #define GKO_DECLARE_CSR_EXTRACT_DIAGONAL(ValueType, IndexType)           \
     void extract_diagonal(std::shared_ptr<const DefaultExecutor> exec,   \
                           const matrix::Csr<ValueType, IndexType> *orig, \
-                          matrix::Dense<ValueType> *diag)
+                          matrix::Diagonal<ValueType> *diag)
 
 #define GKO_DECLARE_ALL_AS_TEMPLATES                                         \
     template <typename ValueType, typename IndexType>                        \

diff --git a/core/matrix/dense.cpp b/core/matrix/dense.cpp
@@ -741,12 +741,12 @@ std::unique_ptr<LinOp> Dense<ValueType>::inverse_column_permute(
 
 
 template <typename ValueType>
-std::unique_ptr<Dense<ValueType>> Dense<ValueType>::extract_diagonal() const
+std::unique_ptr<Diagonal<ValueType>> Dense<ValueType>::extract_diagonal() const
 {
     auto exec = this->get_executor();
 
     const auto diag_size = std::min(this->get_size()[0], this->get_size()[1]);
-    auto diag = Dense<ValueType>::create(exec, dim<2>(diag_size, 1));
+    auto diag = Diagonal<ValueType>::create(exec, diag_size);
     exec->run(dense::make_extract_diagonal(this, lend(diag)));
     return diag;
 }

diff --git a/core/matrix/dense_kernels.hpp b/core/matrix/dense_kernels.hpp
@@ -39,6 +39,7 @@ OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 
 #include <ginkgo/core/base/math.hpp>
 #include <ginkgo/core/base/types.hpp>
+#include <ginkgo/core/matrix/diagonal.hpp>
 
 
 namespace gko {
@@ -169,7 +170,7 @@ namespace kernels {
 #define GKO_DECLARE_EXTRACT_DIAGONAL_KERNEL(_vtype)                    \
     void extract_diagonal(std::shared_ptr<const DefaultExecutor> exec, \
                           const matrix::Dense<_vtype> *orig,           \
-                          matrix::Dense<_vtype> *diag)
+                          matrix::Diagonal<_vtype> *diag)
 
 #define GKO_DECLARE_ALL_AS_TEMPLATES                                        \
     template <typename ValueType>                                           \

diff --git a/core/matrix/ell.cpp b/core/matrix/ell.cpp
@@ -237,13 +237,13 @@ void Ell<ValueType, IndexType>::write(mat_data &data) const
 
 
 template <typename ValueType, typename IndexType>
-std::unique_ptr<Dense<ValueType>> Ell<ValueType, IndexType>::extract_diagonal()
-    const
+std::unique_ptr<Diagonal<ValueType>>
+Ell<ValueType, IndexType>::extract_diagonal() const
 {
     auto exec = this->get_executor();
 
     const auto diag_size = std::min(this->get_size()[0], this->get_size()[1]);
-    auto diag = Dense<ValueType>::create(exec, dim<2>(diag_size, 1));
+    auto diag = Diagonal<ValueType>::create(exec, diag_size);
     exec->run(ell::make_fill_array(diag->get_values(), diag->get_size()[0],
                                    zero<ValueType>()));
     exec->run(ell::make_extract_diagonal(this, lend(diag)));

diff --git a/core/matrix/ell_kernels.hpp b/core/matrix/ell_kernels.hpp
@@ -39,6 +39,7 @@ OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 
 #include <ginkgo/core/matrix/csr.hpp>
 #include <ginkgo/core/matrix/dense.hpp>
+#include <ginkgo/core/matrix/diagonal.hpp>
 
 
 namespace gko {
@@ -83,7 +84,7 @@ namespace kernels {
 #define GKO_DECLARE_ELL_EXTRACT_DIAGONAL_KERNEL(ValueType, IndexType)    \
     void extract_diagonal(std::shared_ptr<const DefaultExecutor> exec,   \
                           const matrix::Ell<ValueType, IndexType> *orig, \
-                          matrix::Dense<ValueType> *diag)
+                          matrix::Diagonal<ValueType> *diag)
 
 #define GKO_DECLARE_ALL_AS_TEMPLATES                                         \
     template <typename ValueType, typename IndexType>                        \

diff --git a/core/matrix/hybrid.cpp b/core/matrix/hybrid.cpp
@@ -270,13 +270,13 @@ void Hybrid<ValueType, IndexType>::write(mat_data &data) const
 
 
 template <typename ValueType, typename IndexType>
-std::unique_ptr<Dense<ValueType>>
+std::unique_ptr<Diagonal<ValueType>>
 Hybrid<ValueType, IndexType>::extract_diagonal() const
 {
     auto exec = this->get_executor();
 
     const auto diag_size = std::min(this->get_size()[0], this->get_size()[1]);
-    auto diag = Dense<ValueType>::create(exec, dim<2>(diag_size, 1));
+    auto diag = Diagonal<ValueType>::create(exec, diag_size);
     exec->run(hybrid::make_fill_array(diag->get_values(), diag->get_size()[0],
                                       zero<ValueType>()));
     exec->run(hybrid::make_extract_ell_diagonal(this->get_ell(), lend(diag)));

diff --git a/core/matrix/sellp.cpp b/core/matrix/sellp.cpp
@@ -290,13 +290,13 @@ void Sellp<ValueType, IndexType>::write(mat_data &data) const
 
 
 template <typename ValueType, typename IndexType>
-std::unique_ptr<Dense<ValueType>>
+std::unique_ptr<Diagonal<ValueType>>
 Sellp<ValueType, IndexType>::extract_diagonal() const
 {
     auto exec = this->get_executor();
 
     const auto diag_size = std::min(this->get_size()[0], this->get_size()[1]);
-    auto diag = Dense<ValueType>::create(exec, dim<2>(diag_size, 1));
+    auto diag = Diagonal<ValueType>::create(exec, diag_size);
     exec->run(sellp::make_fill_array(diag->get_values(), diag->get_size()[0],
                                      zero<ValueType>()));
     exec->run(sellp::make_extract_diagonal(this, lend(diag)));

diff --git a/core/matrix/sellp_kernels.hpp b/core/matrix/sellp_kernels.hpp
@@ -39,6 +39,7 @@ OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 
 #include <ginkgo/core/matrix/csr.hpp>
 #include <ginkgo/core/matrix/dense.hpp>
+#include <ginkgo/core/matrix/diagonal.hpp>
 
 
 namespace gko {
@@ -76,7 +77,7 @@ namespace kernels {
 #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)
+                          matrix::Diagonal<ValueType> *diag)
 
 #define GKO_DECLARE_ALL_AS_TEMPLATES                                 \
     template <typename ValueType, typename IndexType>                \

diff --git a/cuda/matrix/coo_kernels.cu b/cuda/matrix/coo_kernels.cu
@@ -249,7 +249,7 @@ GKO_INSTANTIATE_FOR_EACH_VALUE_AND_INDEX_TYPE(
 template <typename ValueType, typename IndexType>
 void extract_diagonal(std::shared_ptr<const CudaExecutor> exec,
                       const matrix::Coo<ValueType, IndexType> *orig,
-                      matrix::Dense<ValueType> *diag)
+                      matrix::Diagonal<ValueType> *diag)
 {
     const auto nnz = orig->get_num_stored_elements();
     const auto diag_size = diag->get_size()[0];

diff --git a/cuda/matrix/csr_kernels.cu b/cuda/matrix/csr_kernels.cu
@@ -1358,7 +1358,7 @@ GKO_INSTANTIATE_FOR_EACH_VALUE_AND_INDEX_TYPE(
 template <typename ValueType, typename IndexType>
 void extract_diagonal(std::shared_ptr<const CudaExecutor> exec,
                       const matrix::Csr<ValueType, IndexType> *orig,
-                      matrix::Dense<ValueType> *diag)
+                      matrix::Diagonal<ValueType> *diag)
 {
     const auto nnz = orig->get_num_stored_elements();
     const auto diag_size = diag->get_size()[0];

diff --git a/cuda/matrix/dense_kernels.cu b/cuda/matrix/dense_kernels.cu
@@ -690,7 +690,7 @@ GKO_INSTANTIATE_FOR_EACH_VALUE_AND_INDEX_TYPE(
 template <typename ValueType>
 void extract_diagonal(std::shared_ptr<const CudaExecutor> exec,
                       const matrix::Dense<ValueType> *orig,
-                      matrix::Dense<ValueType> *diag)
+                      matrix::Diagonal<ValueType> *diag)
 {
     const dim3 grid_dim = ceildiv(diag->get_size()[0], default_block_size);
     kernel::extract_diagonal<<<grid_dim, default_block_size>>>(

diff --git a/cuda/matrix/ell_kernels.cu b/cuda/matrix/ell_kernels.cu
@@ -366,7 +366,7 @@ GKO_INSTANTIATE_FOR_EACH_VALUE_AND_INDEX_TYPE(
 template <typename ValueType, typename IndexType>
 void extract_diagonal(std::shared_ptr<const CudaExecutor> exec,
                       const matrix::Ell<ValueType, IndexType> *orig,
-                      matrix::Dense<ValueType> *diag)
+                      matrix::Diagonal<ValueType> *diag)
 {
     const auto max_nnz_per_row = orig->get_num_stored_elements_per_row();
     const auto orig_stride = orig->get_stride();

diff --git a/cuda/matrix/sellp_kernels.cu b/cuda/matrix/sellp_kernels.cu
@@ -233,7 +233,7 @@ GKO_INSTANTIATE_FOR_EACH_VALUE_AND_INDEX_TYPE(
 template <typename ValueType, typename IndexType>
 void extract_diagonal(std::shared_ptr<const CudaExecutor> exec,
                       const matrix::Sellp<ValueType, IndexType> *orig,
-                      matrix::Dense<ValueType> *diag)
+                      matrix::Diagonal<ValueType> *diag)
 {
     const auto diag_size = diag->get_size()[0];
     const auto diag_stride = diag->get_stride();

diff --git a/cuda/test/matrix/coo_kernels.cpp b/cuda/test/matrix/coo_kernels.cpp
@@ -44,6 +44,7 @@ OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 #include <ginkgo/core/base/executor.hpp>
 #include <ginkgo/core/matrix/csr.hpp>
 #include <ginkgo/core/matrix/dense.hpp>
+#include <ginkgo/core/matrix/diagonal.hpp>
 
 
 #include "core/matrix/coo_kernels.hpp"

diff --git a/cuda/test/matrix/csr_kernels.cpp b/cuda/test/matrix/csr_kernels.cpp
@@ -43,6 +43,7 @@ OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 #include <ginkgo/core/base/executor.hpp>
 #include <ginkgo/core/matrix/coo.hpp>
 #include <ginkgo/core/matrix/dense.hpp>
+#include <ginkgo/core/matrix/diagonal.hpp>
 #include <ginkgo/core/matrix/ell.hpp>
 #include <ginkgo/core/matrix/hybrid.hpp>
 #include <ginkgo/core/matrix/identity.hpp>

diff --git a/cuda/test/matrix/ell_kernels.cpp b/cuda/test/matrix/ell_kernels.cpp
@@ -44,6 +44,7 @@ OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 #include <ginkgo/core/base/executor.hpp>
 #include <ginkgo/core/matrix/csr.hpp>
 #include <ginkgo/core/matrix/dense.hpp>
+#include <ginkgo/core/matrix/diagonal.hpp>
 
 
 #include "core/matrix/ell_kernels.hpp"

diff --git a/cuda/test/matrix/hybrid_kernels.cpp b/cuda/test/matrix/hybrid_kernels.cpp
@@ -43,6 +43,7 @@ OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 #include <ginkgo/core/base/exception_helpers.hpp>
 #include <ginkgo/core/base/executor.hpp>
 #include <ginkgo/core/matrix/dense.hpp>
+#include <ginkgo/core/matrix/diagonal.hpp>
 
 
 #include "core/matrix/hybrid_kernels.hpp"

diff --git a/cuda/test/matrix/sellp_kernels.cpp b/cuda/test/matrix/sellp_kernels.cpp
@@ -44,6 +44,7 @@ OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 #include <ginkgo/core/base/executor.hpp>
 #include <ginkgo/core/matrix/csr.hpp>
 #include <ginkgo/core/matrix/dense.hpp>
+#include <ginkgo/core/matrix/diagonal.hpp>
 
 
 #include "core/matrix/sellp_kernels.hpp"

diff --git a/hip/matrix/coo_kernels.hip.cpp b/hip/matrix/coo_kernels.hip.cpp
@@ -261,7 +261,7 @@ GKO_INSTANTIATE_FOR_EACH_VALUE_AND_INDEX_TYPE(
 template <typename ValueType, typename IndexType>
 void extract_diagonal(std::shared_ptr<const HipExecutor> exec,
                       const matrix::Coo<ValueType, IndexType> *orig,
-                      matrix::Dense<ValueType> *diag)
+                      matrix::Diagonal<ValueType> *diag)
 {
     const auto nnz = orig->get_num_stored_elements();
     const auto diag_size = diag->get_size()[0];

diff --git a/hip/matrix/csr_kernels.hip.cpp b/hip/matrix/csr_kernels.hip.cpp
@@ -1184,7 +1184,7 @@ GKO_INSTANTIATE_FOR_EACH_VALUE_AND_INDEX_TYPE(
 template <typename ValueType, typename IndexType>
 void extract_diagonal(std::shared_ptr<const HipExecutor> exec,
                       const matrix::Csr<ValueType, IndexType> *orig,
-                      matrix::Dense<ValueType> *diag)
+                      matrix::Diagonal<ValueType> *diag)
 {
     const auto nnz = orig->get_num_stored_elements();
     const auto diag_size = diag->get_size()[0];

diff --git a/hip/matrix/dense_kernels.hip.cpp b/hip/matrix/dense_kernels.hip.cpp
@@ -717,7 +717,7 @@ GKO_INSTANTIATE_FOR_EACH_VALUE_AND_INDEX_TYPE(
 template <typename ValueType>
 void extract_diagonal(std::shared_ptr<const HipExecutor> exec,
                       const matrix::Dense<ValueType> *orig,
-                      matrix::Dense<ValueType> *diag)
+                      matrix::Diagonal<ValueType> *diag)
 {
     const dim3 grid_dim = ceildiv(diag->get_size()[0], default_block_size);
     hipLaunchKernelGGL(kernel::extract_diagonal, dim3(grid_dim),

diff --git a/hip/matrix/ell_kernels.hip.cpp b/hip/matrix/ell_kernels.hip.cpp
@@ -375,7 +375,7 @@ GKO_INSTANTIATE_FOR_EACH_VALUE_AND_INDEX_TYPE(
 template <typename ValueType, typename IndexType>
 void extract_diagonal(std::shared_ptr<const HipExecutor> exec,
                       const matrix::Ell<ValueType, IndexType> *orig,
-                      matrix::Dense<ValueType> *diag)
+                      matrix::Diagonal<ValueType> *diag)
 {
     const auto max_nnz_per_row = orig->get_num_stored_elements_per_row();
     const auto orig_stride = orig->get_stride();

diff --git a/hip/matrix/sellp_kernels.hip.cpp b/hip/matrix/sellp_kernels.hip.cpp
@@ -239,7 +239,7 @@ GKO_INSTANTIATE_FOR_EACH_VALUE_AND_INDEX_TYPE(
 template <typename ValueType, typename IndexType>
 void extract_diagonal(std::shared_ptr<const HipExecutor> exec,
                       const matrix::Sellp<ValueType, IndexType> *orig,
-                      matrix::Dense<ValueType> *diag)
+                      matrix::Diagonal<ValueType> *diag)
 {
     const auto diag_size = diag->get_size()[0];
     const auto diag_stride = diag->get_stride();

diff --git a/hip/test/matrix/coo_kernels.hip.cpp b/hip/test/matrix/coo_kernels.hip.cpp
@@ -44,6 +44,7 @@ OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 #include <ginkgo/core/base/executor.hpp>
 #include <ginkgo/core/matrix/csr.hpp>
 #include <ginkgo/core/matrix/dense.hpp>
+#include <ginkgo/core/matrix/diagonal.hpp>
 
 
 #include "core/matrix/coo_kernels.hpp"

diff --git a/hip/test/matrix/csr_kernels.hip.cpp b/hip/test/matrix/csr_kernels.hip.cpp
@@ -43,6 +43,7 @@ OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 #include <ginkgo/core/base/executor.hpp>
 #include <ginkgo/core/matrix/coo.hpp>
 #include <ginkgo/core/matrix/dense.hpp>
+#include <ginkgo/core/matrix/diagonal.hpp>
 #include <ginkgo/core/matrix/ell.hpp>
 #include <ginkgo/core/matrix/hybrid.hpp>
 #include <ginkgo/core/matrix/identity.hpp>

diff --git a/hip/test/matrix/ell_kernels.hip.cpp b/hip/test/matrix/ell_kernels.hip.cpp
@@ -44,6 +44,7 @@ OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 #include <ginkgo/core/base/executor.hpp>
 #include <ginkgo/core/matrix/csr.hpp>
 #include <ginkgo/core/matrix/dense.hpp>
+#include <ginkgo/core/matrix/diagonal.hpp>
 
 
 #include "core/matrix/ell_kernels.hpp"

diff --git a/hip/test/matrix/hybrid_kernels.hip.cpp b/hip/test/matrix/hybrid_kernels.hip.cpp
@@ -43,6 +43,7 @@ OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 #include <ginkgo/core/base/exception_helpers.hpp>
 #include <ginkgo/core/base/executor.hpp>
 #include <ginkgo/core/matrix/dense.hpp>
+#include <ginkgo/core/matrix/diagonal.hpp>
 
 
 #include "core/matrix/hybrid_kernels.hpp"

diff --git a/hip/test/matrix/sellp_kernels.hip.cpp b/hip/test/matrix/sellp_kernels.hip.cpp
@@ -44,6 +44,7 @@ OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 #include <ginkgo/core/base/executor.hpp>
 #include <ginkgo/core/matrix/csr.hpp>
 #include <ginkgo/core/matrix/dense.hpp>
+#include <ginkgo/core/matrix/diagonal.hpp>
 
 
 #include "core/matrix/sellp_kernels.hpp"

diff --git a/include/ginkgo/core/base/lin_op.hpp b/include/ginkgo/core/base/lin_op.hpp
@@ -54,7 +54,7 @@ namespace matrix {
 
 
 template <typename ValueType>
-class Dense;
+class Diagonal;
 
 
 }
@@ -618,7 +618,7 @@ class DiagonalExtractable {
      *
      * @param diag  the vector into which the diagonal will be written
      */
-    virtual std::unique_ptr<matrix::Dense<ValueType>> extract_diagonal()
+    virtual std::unique_ptr<matrix::Diagonal<ValueType>> extract_diagonal()
         const = 0;
 };
 

diff --git a/include/ginkgo/core/matrix/coo.hpp b/include/ginkgo/core/matrix/coo.hpp
@@ -115,7 +115,7 @@ class Coo : public EnableLinOp<Coo<ValueType, IndexType>>,
 
     void write(mat_data &data) const override;
 
-    std::unique_ptr<Dense<ValueType>> extract_diagonal() const override;
+    std::unique_ptr<Diagonal<ValueType>> extract_diagonal() const override;
 
     /**
      * Returns the values of the matrix.

diff --git a/include/ginkgo/core/matrix/csr.hpp b/include/ginkgo/core/matrix/csr.hpp
@@ -699,7 +699,7 @@ class Csr : public EnableLinOp<Csr<ValueType, IndexType>>,
     std::unique_ptr<LinOp> inverse_column_permute(
         const Array<IndexType> *inverse_permutation_indices) const override;
 
-    std::unique_ptr<Dense<ValueType>> extract_diagonal() const override;
+    std::unique_ptr<Diagonal<ValueType>> extract_diagonal() const override;
 
     /**
      * Sorts all (value, col_idx) pairs in each row by column index

diff --git a/include/ginkgo/core/matrix/dense.hpp b/include/ginkgo/core/matrix/dense.hpp
@@ -243,7 +243,7 @@ class Dense : public EnableLinOp<Dense<ValueType>>,
     std::unique_ptr<LinOp> inverse_column_permute(
         const Array<int64> *inverse_permutation_indices) const override;
 
-    std::unique_ptr<Dense<ValueType>> extract_diagonal() const override;
+    std::unique_ptr<Diagonal<ValueType>> extract_diagonal() const override;
 
 
     /**