Merge Fix inconsistent behavior of scalar Jacobi

This PR fixes inconsistent behavior of scalar Jacobi between reference and the others Related PR: #1642
ginkgo-project · Jul 6, 2024 · 61bbba8 · 61bbba8
2 parents 70db4f7 + 7cf1b31
commit 61bbba8
Show file tree

Hide file tree

Showing 4 changed files with 69 additions and 1 deletion.
diff --git a/common/unified/preconditioner/jacobi_kernels.cpp b/common/unified/preconditioner/jacobi_kernels.cpp
@@ -42,7 +42,9 @@ void invert_diagonal(std::shared_ptr<const DefaultExecutor> exec,
     run_kernel(
         exec,
         [] GKO_KERNEL(auto elem, auto diag, auto inv_diag) {
-            inv_diag[elem] = safe_divide(one(diag[elem]), diag[elem]);
+            // if the diagonal is zero, we use 1 for in the inverted result.
+            inv_diag[elem] = is_zero(diag[elem]) ? one(diag[elem])
+                                                 : one(diag[elem]) / diag[elem];
         },
         diag.get_size(), diag, inv_diag);
 }

diff --git a/dpcpp/test/preconditioner/jacobi_kernels.dp.cpp b/dpcpp/test/preconditioner/jacobi_kernels.dp.cpp
@@ -905,4 +905,27 @@ TEST_F(
 }
 
 
+TEST_F(Jacobi, ScalarJacobiHandleZero)
+{
+    auto mtx = gko::share(
+        gko::initialize<Vec>({{0, 0, 0}, {0, 2, 0}, {0, 0, 0}}, ref));
+    auto b = gko::initialize<Vec>({1, 2, 3}, ref);
+    auto x = Vec::create(ref, gko::dim<2>(3, 1));
+    auto d_b = b->clone(dpcpp);
+    auto d_x = x->clone(dpcpp);
+    auto d_mtx = gko::share(mtx->clone(dpcpp));
+
+    auto jacobi = Bj::build().with_max_block_size(1u).on(ref)->generate(mtx);
+    // Must generate from scratch because the clone copies the inverted
+    // information.
+    auto d_jacobi =
+        Bj::build().with_max_block_size(1u).on(dpcpp)->generate(d_mtx);
+
+    // Jacobi uses 1 as the result when diagonal value is zero.
+    jacobi->apply(b, x);
+    d_jacobi->apply(d_b, d_x);
+    GKO_ASSERT_MTX_NEAR(d_x, x, 0.0);
+}
+
+
 }  // namespace
diff --git a/reference/test/preconditioner/jacobi.cpp b/reference/test/preconditioner/jacobi.cpp
@@ -494,4 +494,24 @@ TYPED_TEST(Jacobi, ScalarJacobiGeneratesOnDifferentPrecision)
 }
 
 
+TYPED_TEST(Jacobi, ScalarJacobiHandleZero)
+{
+    using value_type = typename TestFixture::value_type;
+    using Vec = typename TestFixture::Vec;
+    using Bj = typename TestFixture::Bj;
+    auto mtx = gko::share(
+        gko::initialize<Vec>({{0, 0, 0}, {0, 2, 0}, {0, 0, 0}}, this->exec));
+    auto b = gko::initialize<Vec>({1, 2, 3}, this->exec);
+    auto x = Vec::create(this->exec, gko::dim<2>(3, 1));
+
+    auto jacobi = this->scalar_j_factory->generate(mtx);
+
+    // Jacobi uses 1 as the result when diagonal value is zero.
+    jacobi->apply(b, x);
+    ASSERT_EQ(x->at(0, 0), value_type{1.0});
+    ASSERT_EQ(x->at(1, 0), value_type{1.0});
+    ASSERT_EQ(x->at(2, 0), value_type{3.0});
+}
+
+
 }  // namespace
diff --git a/test/preconditioner/jacobi_kernels.cpp b/test/preconditioner/jacobi_kernels.cpp
@@ -887,3 +887,26 @@ TEST_F(
 
     GKO_ASSERT_MTX_NEAR(d_x, x, 1e-6);
 }
+
+
+TEST_F(Jacobi, ScalarJacobiHandleZero)
+{
+    auto mtx = gko::share(
+        gko::initialize<Vec>({{0, 0, 0}, {0, 2, 0}, {0, 0, 0}}, ref));
+    auto b = gko::initialize<Vec>({1, 2, 3}, ref);
+    auto x = Vec::create(ref, gko::dim<2>(3, 1));
+    auto d_b = b->clone(exec);
+    auto d_x = x->clone(exec);
+    auto d_mtx = gko::share(mtx->clone(exec));
+
+    auto jacobi = Bj::build().with_max_block_size(1u).on(ref)->generate(mtx);
+    // Must generate from scratch because the clone copies the inverted
+    // information.
+    auto d_jacobi =
+        Bj::build().with_max_block_size(1u).on(exec)->generate(d_mtx);
+
+    // Jacobi uses 1 as the result when diagonal value is zero.
+    jacobi->apply(b, x);
+    d_jacobi->apply(d_b, d_x);
+    GKO_ASSERT_MTX_NEAR(d_x, x, 0.0);
+}