Merge Fix scalar jacobi generate on diff precision, add DiagonalLinOp…

…Extractable and precision convert in Diagonal

This pr fixes scalar jacobi generation on different precision, add `DiagonalLinOpExtractable` to return `LinOp` without type and precision convert in `Diagonal`

Related PR: #834

core/matrix/diagonal.cpp

@@ -146,6 +146,22 @@ std::unique_ptr<LinOp> Diagonal<ValueType>::conj_transpose() const
 }
 
 
+template <typename ValueType>
+void Diagonal<ValueType>::convert_to(
+    Diagonal<next_precision<ValueType>> *result) const
+{
+    result->values_ = this->values_;
+    result->set_size(this->get_size());
+}
+
+
+template <typename ValueType>
+void Diagonal<ValueType>::move_to(Diagonal<next_precision<ValueType>> *result)
+{
+    this->convert_to(result);
+}
+
+
 template <typename ValueType>
 void Diagonal<ValueType>::convert_to(Csr<ValueType, int32> *result) const
 {
@@ -302,4 +318,21 @@ GKO_INSTANTIATE_FOR_EACH_VALUE_TYPE(GKO_DECLARE_DIAGONAL_MATRIX);
 
 
 }  // namespace matrix
+
+
+// Implement DiagonalExtractable for LinOp when Diagonal is complete class
+template <typename ValueType>
+std::unique_ptr<LinOp> DiagonalExtractable<ValueType>::extract_diagonal_linop()
+    const
+{
+    return this->extract_diagonal();
+}
+
+
+#define GKO_DECLARE_DIAGONAL_EXTRACTABLE(value_type) \
+    std::unique_ptr<LinOp>                           \
+    DiagonalExtractable<value_type>::extract_diagonal_linop() const
+GKO_INSTANTIATE_FOR_EACH_VALUE_TYPE(GKO_DECLARE_DIAGONAL_EXTRACTABLE);
+
+
 }  // namespace gko

core/preconditioner/jacobi.cpp

@@ -40,6 +40,7 @@ OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 #include <ginkgo/core/base/executor.hpp>
 #include <ginkgo/core/base/math.hpp>
 #include <ginkgo/core/base/precision_dispatch.hpp>
+#include <ginkgo/core/base/temporary_conversion.hpp>
 #include <ginkgo/core/base/utils.hpp>
 #include <ginkgo/core/matrix/csr.hpp>
 #include <ginkgo/core/matrix/dense.hpp>
@@ -231,11 +232,20 @@ void Jacobi<ValueType, IndexType>::generate(const LinOp *system_matrix,
     using csr_type = matrix::Csr<ValueType, IndexType>;
     const auto exec = this->get_executor();
     if (parameters_.max_block_size == 1) {
-        auto diag = as<DiagonalExtractable<ValueType>>(system_matrix)
-                        ->extract_diagonal();
-        auto temp = gko::Array<ValueType>::view(
-            exec, system_matrix->get_size()[0], diag->get_values());
+        auto diag = share(as<DiagonalLinOpExtractable>(system_matrix)
+                              ->extract_diagonal_linop());
+        auto diag_vt =
+            ::gko::detail::temporary_conversion<matrix::Diagonal<ValueType>>::
+                template create<matrix::Diagonal<next_precision<ValueType>>>(
+                    diag.get());
+        if (!diag_vt) {
+            GKO_NOT_SUPPORTED(system_matrix);
+        }
+        auto temp = Array<ValueType>::view(diag_vt->get_executor(),
+                                           diag_vt->get_size()[0],
+                                           diag_vt->get_values());
         this->blocks_ = temp;
+        this->num_blocks_ = diag_vt->get_size()[0];
     } else {
         auto csr_mtx = convert_to_with_sorting<csr_type>(exec, system_matrix,
                                                          skip_sorting);

include/ginkgo/core/base/lin_op.hpp

@@ -654,6 +654,29 @@ class Preconditionable {
 };
 
 
+/**
+ * The diagonal of a LinOp can be extracted. It will be implemented by
+ * DiagonalExtractable<ValueType>, so the class does not need to implement it.
+ * extract_diagonal_linop returns a linop which extracts the elements whose col
+ * and row index are the same and stores the result in a min(nrows, ncols) x 1
+ * dense matrix.
+ *
+ * @ingroup diagonal
+ * @ingroup LinOp
+ */
+class DiagonalLinOpExtractable {
+public:
+    virtual ~DiagonalLinOpExtractable() = default;
+
+    /**
+     * Extracts the diagonal entries of the matrix into a vector.
+     *
+     * @return linop  the linop of diagonal format
+     */
+    virtual std::unique_ptr<LinOp> extract_diagonal_linop() const = 0;
+};
+
+
 /**
  * The diagonal of a LinOp implementing this interface can be extracted.
  * extract_diagonal extracts the elements whose col and row index are the
@@ -662,12 +685,14 @@ class Preconditionable {
  * @ingroup LinOp
  */
 template <typename ValueType>
-class DiagonalExtractable {
+class DiagonalExtractable : public DiagonalLinOpExtractable {
 public:
     using value_type = ValueType;
 
     virtual ~DiagonalExtractable() = default;
 
+    std::unique_ptr<LinOp> extract_diagonal_linop() const override;
+
     /**
      * Extracts the diagonal entries of the matrix into a vector.
      *

include/ginkgo/core/matrix/diagonal.hpp

@@ -70,6 +70,7 @@ class Diagonal
       public EnableCreateMethod<Diagonal<ValueType>>,
       public ConvertibleTo<Csr<ValueType, int32>>,
       public ConvertibleTo<Csr<ValueType, int64>>,
+      public ConvertibleTo<Diagonal<next_precision<ValueType>>>,
       public Transposable,
       public WritableToMatrixData<ValueType, int32>,
       public WritableToMatrixData<ValueType, int64>,
@@ -92,10 +93,16 @@ class Diagonal
     using mat_data32 = gko::matrix_data<ValueType, int32>;
     using absolute_type = remove_complex<Diagonal>;
 
+    friend class Diagonal<next_precision<ValueType>>;
+
     std::unique_ptr<LinOp> transpose() const override;
 
     std::unique_ptr<LinOp> conj_transpose() const override;
 
+    void convert_to(Diagonal<next_precision<ValueType>> *result) const override;
+
+    void move_to(Diagonal<next_precision<ValueType>> *result) override;
+
     void convert_to(Csr<ValueType, int32> *result) const override;
 
     void move_to(Csr<ValueType, int32> *result) override;

reference/test/matrix/diagonal_kernels.cpp

@@ -108,6 +108,46 @@ class Diagonal : public ::testing::Test {
 TYPED_TEST_SUITE(Diagonal, gko::test::ValueTypes);
 
 
+TYPED_TEST(Diagonal, ConvertsToPrecision)
+{
+    using ValueType = typename TestFixture::value_type;
+    using OtherType = typename gko::next_precision<ValueType>;
+    using Diagonal = typename TestFixture::Diag;
+    using OtherDiagonal = gko::matrix::Diagonal<OtherType>;
+    auto tmp = OtherDiagonal::create(this->exec);
+    auto res = Diagonal::create(this->exec);
+    // If OtherType is more precise: 0, otherwise r
+    auto residual = r<OtherType>::value < r<ValueType>::value
+                        ? gko::remove_complex<ValueType>{0}
+                        : gko::remove_complex<ValueType>{r<OtherType>::value};
+
+    this->diag1->convert_to(tmp.get());
+    tmp->convert_to(res.get());
+
+    GKO_ASSERT_MTX_NEAR(this->diag1, res, residual);
+}
+
+
+TYPED_TEST(Diagonal, MovesToPrecision)
+{
+    using ValueType = typename TestFixture::value_type;
+    using OtherType = typename gko::next_precision<ValueType>;
+    using Diagonal = typename TestFixture::Diag;
+    using OtherDiagonal = gko::matrix::Diagonal<OtherType>;
+    auto tmp = OtherDiagonal::create(this->exec);
+    auto res = Diagonal::create(this->exec);
+    // If OtherType is more precise: 0, otherwise r
+    auto residual = r<OtherType>::value < r<ValueType>::value
+                        ? gko::remove_complex<ValueType>{0}
+                        : gko::remove_complex<ValueType>{r<OtherType>::value};
+
+    this->diag1->move_to(tmp.get());
+    tmp->move_to(res.get());
+
+    GKO_ASSERT_MTX_NEAR(this->diag1, res, residual);
+}
+
+
 TYPED_TEST(Diagonal, AppliesToDense)
 {
     using value_type = typename TestFixture::value_type;

reference/test/preconditioner/jacobi.cpp

@@ -65,6 +65,7 @@ class Jacobi : public ::testing::Test {
     Jacobi()
         : exec(gko::ReferenceExecutor::create()),
           bj_factory(Bj::build().with_max_block_size(3u).on(exec)),
+          scalar_j_factory(Bj::build().with_max_block_size(1u).on(exec)),
           block_pointers(exec, 3),
           block_precisions(exec, 2),
           mtx(Mtx::create(exec, gko::dim<2>{5}, 13))
@@ -168,6 +169,7 @@ class Jacobi : public ::testing::Test {
 
     std::shared_ptr<const gko::Executor> exec;
     std::unique_ptr<typename Bj::Factory> bj_factory;
+    std::unique_ptr<typename Bj::Factory> scalar_j_factory;
     std::unique_ptr<typename Bj::Factory> adaptive_bj_factory;
     gko::Array<index_type> block_pointers;
     gko::Array<gko::precision_reduction> block_precisions;
@@ -378,4 +380,25 @@ TYPED_TEST(Jacobi, GeneratesCorrectMatrixDataWithAdaptivePrecision)
 }
 
 
+TYPED_TEST(Jacobi, ScalarJacobiGeneratesOnDifferentPrecision)
+{
+    using value_type = typename TestFixture::value_type;
+    using index_type = typename TestFixture::index_type;
+    using next_type = gko::next_precision<value_type>;
+    using Bj = typename TestFixture::Bj;
+    auto csr =
+        gko::share(gko::matrix::Csr<next_type, index_type>::create(this->exec));
+    csr->copy_from(gko::lend(this->mtx));
+    std::shared_ptr<Bj> bj{};
+
+    ASSERT_NO_THROW(bj = this->scalar_j_factory->generate(csr));
+    ASSERT_EQ(bj->get_num_blocks(), 5u);
+    ASSERT_EQ(bj->get_blocks()[0], value_type{4});
+    ASSERT_EQ(bj->get_blocks()[1], value_type{4});
+    ASSERT_EQ(bj->get_blocks()[2], value_type{4});
+    ASSERT_EQ(bj->get_blocks()[3], value_type{4});
+    ASSERT_EQ(bj->get_blocks()[4], value_type{4});
+}
+
+
 }  // namespace