Extend and fix Richardson test

Author: thoasm

reference/test/solver/gmres_kernels.cpp

@@ -207,6 +207,7 @@ TYPED_TEST(Gmres, KernelRestart)
                         r<value_type>::value);
 }
 
+
 TYPED_TEST(Gmres, KernelRestartRgs)
 {
     using value_type = typename TestFixture::value_type;
@@ -280,6 +281,7 @@ TYPED_TEST(Gmres, KernelRestartRgs)
                         r<value_type>::value);
 }
 
+
 TYPED_TEST(Gmres, KernelRichardsonLsq)
 {
     using T = typename TestFixture::value_type;
@@ -291,32 +293,6 @@ TYPED_TEST(Gmres, KernelRichardsonLsq)
     const gko::size_type num_rhs = 1;
     const gko::size_type k_rows = 3;
     const gko::size_type krylov_dim = 4;
-    const gko::size_type iter = 2;
-    auto hessenberg_iter =
-        Mtx::create(this->exec, gko::dim<2>{iter + 2, num_rhs});
-    hessenberg_iter->fill(nan);
-    auto exp_hessenberg_iter =
-        Mtx::create(this->exec, gko::dim<2>{iter + 2, num_rhs});
-    for (gko::size_type i = iter + 1; i < iter + 2; ++i) {
-        for (gko::size_type r = 0; r < num_rhs; ++r) {
-            // These entries should not be touched!
-            hessenberg_iter->at(i, r) = 42.;
-            exp_hessenberg_iter->at(i, r) = 42.;
-        }
-    }
-    auto d_hessenberg_iter =
-        Mtx::create(this->exec, gko::dim<2>{krylov_dim + 1, num_rhs});
-    auto exp_d_hessenberg_iter = gko::clone(d_hessenberg_iter);
-    auto sketched_next_krylov2 =
-        Mtx::create(this->exec, gko::dim<2>{k_rows, num_rhs});
-    auto exp_sketched_next_krylov2 = gko::clone(sketched_next_krylov2);
-    sketched_next_krylov2->fill(nan);
-    d_hessenberg_iter->fill(nan);
-    for (gko::size_type i = k_rows; i < krylov_dim + 1; ++i) {
-        // These entries should not be touched by kernel
-        d_hessenberg_iter->at(i, 0) = 13;
-        exp_d_hessenberg_iter->at(i, 0) = 13;
-    }
     // size of sketched_krylov_bases: gko::dim<2>{k_rows * (krylov_dim + 1),
     // num_rhs}
     auto sketched_krylov_bases = gko::initialize<Mtx>(
@@ -326,7 +302,7 @@ TYPED_TEST(Gmres, KernelRichardsonLsq)
                .5, -.5, 2., 1.5, -1.,
             }, this->exec);
     // clang-format on
-    // Normalize vectors
+    // Normalize all vectors
     for (gko::size_type i = 0; i < krylov_dim + 1; ++i) {
         gko::remove_complex<T> norm{0};
         for (gko::size_type k = 0; k < k_rows; ++k) {
@@ -338,57 +314,76 @@ TYPED_TEST(Gmres, KernelRichardsonLsq)
             sketched_krylov_bases->at(k + i * k_rows, 0) /= norm;
         }
     }
+    auto d_hessenberg_iter =
+        Mtx::create(this->exec, gko::dim<2>{krylov_dim + 1, num_rhs});
+    auto exp_d_hessenberg_iter = gko::clone(d_hessenberg_iter);
+    // Test all the different iteration steps
+    for (gko::size_type iter = 1; iter < krylov_dim + 1; ++iter) {
+        SCOPED_TRACE(iter);
+        auto hessenberg_iter =
+            Mtx::create(this->exec, gko::dim<2>{iter + 2, num_rhs});
+        auto exp_hessenberg_iter = gko::clone(hessenberg_iter);
+        auto sketched_next_krylov2 =
+            Mtx::create(this->exec, gko::dim<2>{k_rows, num_rhs});
+        auto exp_sketched_next_krylov2 = gko::clone(sketched_next_krylov2);
+        hessenberg_iter->fill(nan);
+        for (gko::size_type i = iter + 1; i < iter + 2; ++i) {
+            for (gko::size_type r = 0; r < num_rhs; ++r) {
+                // These entries should not be touched!
+                hessenberg_iter->at(i, r) = 42.;
+                exp_hessenberg_iter->at(i, r) = 42.;
+            }
+        }
+        sketched_next_krylov2->fill(nan);
+        d_hessenberg_iter->fill(nan);
+        for (gko::size_type i = iter + 1; i < krylov_dim + 1; ++i) {
+            // These entries should not be touched by kernel
+            d_hessenberg_iter->at(i, 0) = 13;
+            exp_d_hessenberg_iter->at(i, 0) = 13;
+        }
 
-    gko::kernels::reference::gmres::richardson_lsq(
-        this->exec, sketched_krylov_bases.get(), hessenberg_iter.get(),
-        d_hessenberg_iter.get(), sketched_next_krylov2.get(), iter, k_rows);
-
-    // Adjust the dimensions for the GEMV and AXPY operations
-    // Note: since this matrix is stored column-major in our code, it appears as
-    // transposed in its normal state
-    auto shrunk_sketched_krylov_bases_trans = Mtx::create_const(
-        this->exec, gko::dim<2>{iter + 1, k_rows},
-        gko::make_const_array_view<T>(
-            this->exec, sketched_krylov_bases->get_num_stored_elements(),
-            sketched_krylov_bases->get_const_values()),
-        k_rows);
-    auto shrunk_sketched_krylov_bases =
-        gko::as<Mtx>(shrunk_sketched_krylov_bases_trans->transpose());
-
-    auto exp_shrunk_hessenberg = exp_hessenberg_iter->create_submatrix(
-        gko::span{0, iter + 1}, gko::span{0, num_rhs});
-    auto exp_shrunk_d_hessenberg = exp_d_hessenberg_iter->create_submatrix(
-        gko::span{0, iter + 1}, gko::span{0, num_rhs});
-
-    // sketched_krylov2 = sketched_krylov_bases[iter + 1, :];
-    for (gko::size_type k = 0; k < num_rhs; k++) {
-        for (gko::size_type j = 0; j < k_rows; j++)
-            exp_sketched_next_krylov2->at(j, k) =
-                sketched_krylov_bases->at(j + (iter + 1) * k_rows, k);
-    }
-    exp_shrunk_hessenberg->fill(0);
-    for (int i = 0; i < 3; ++i) {
-        // d_hessenberg_iter = Transpose(sketched_krylov_bases) *
-        // sketched_krylov2;
-        shrunk_sketched_krylov_bases_trans->apply(
-            exp_sketched_next_krylov2.get(), exp_shrunk_d_hessenberg.get());
-
-        // sketched_krylov2 = sketched_krylov2 - sketched_krylov_bases *
-        // d_hessenberg_iter;
-        shrunk_sketched_krylov_bases->apply(neg_one_mtx,
-                                            exp_shrunk_d_hessenberg, one_mtx,
-                                            exp_sketched_next_krylov2.get());
-
-        // hessenberg_iter = hessenberg_iter + d_hessenberg_iter;
-        exp_shrunk_hessenberg->add_scaled(one_mtx, exp_shrunk_d_hessenberg);
-    }
+        gko::kernels::reference::gmres::richardson_lsq(
+            this->exec, sketched_krylov_bases.get(), hessenberg_iter.get(),
+            d_hessenberg_iter.get(), sketched_next_krylov2.get(), iter, k_rows);
+        // Adjust the dimensions for the GEMV and AXPY operations
+        // Note: since this matrix is stored column-major in our code, it
+        // appears as transposed in its normal state
+        auto shrunk_sketched_krylov_bases_trans = Mtx::create_const(
+            this->exec, gko::dim<2>{iter + 1, k_rows},
+            gko::make_const_array_view<T>(
+                this->exec, sketched_krylov_bases->get_num_stored_elements(),
+                sketched_krylov_bases->get_const_values()),
+            k_rows);
+        auto shrunk_sketched_krylov_bases =
+            gko::as<Mtx>(shrunk_sketched_krylov_bases_trans->transpose());
+        auto exp_shrunk_hessenberg = exp_hessenberg_iter->create_submatrix(
+            gko::span{0, iter + 1}, gko::span{0, num_rhs});
+        auto exp_shrunk_d_hessenberg = exp_d_hessenberg_iter->create_submatrix(
+            gko::span{0, iter + 1}, gko::span{0, num_rhs});
+        for (gko::size_type k = 0; k < num_rhs; k++) {
+            for (gko::size_type j = 0; j < k_rows; j++)
+                exp_sketched_next_krylov2->at(j, k) =
+                    sketched_krylov_bases->at(j + (iter + 1) * k_rows, k);
+        }
+        exp_shrunk_hessenberg->fill(0);
+        for (int i = 0; i < 3; ++i) {
+            shrunk_sketched_krylov_bases_trans->apply(
+                exp_sketched_next_krylov2.get(), exp_shrunk_d_hessenberg.get());
+            shrunk_sketched_krylov_bases->apply(
+                neg_one_mtx, exp_shrunk_d_hessenberg, one_mtx,
+                exp_sketched_next_krylov2.get());
+            exp_shrunk_hessenberg->add_scaled(one_mtx, exp_shrunk_d_hessenberg);
+        }
 
-    GKO_EXPECT_MTX_NEAR(hessenberg_iter, exp_hessenberg_iter, r<T>::value);
-    GKO_EXPECT_MTX_NEAR(d_hessenberg_iter, exp_d_hessenberg_iter, r<T>::value);
-    GKO_EXPECT_MTX_NEAR(sketched_next_krylov2, exp_sketched_next_krylov2,
-                        r<T>::value);
+        GKO_EXPECT_MTX_NEAR(hessenberg_iter, exp_hessenberg_iter, r<T>::value);
+        GKO_EXPECT_MTX_NEAR(d_hessenberg_iter, exp_d_hessenberg_iter,
+                            r<T>::value);
+        GKO_EXPECT_MTX_NEAR(sketched_next_krylov2, exp_sketched_next_krylov2,
+                            r<T>::value);
+    }
 }
 
+
 TYPED_TEST(Gmres, KernelHessenbergQrIter0)
 {
     using T = typename TestFixture::value_type;
@@ -560,6 +555,7 @@ TYPED_TEST(Gmres, KernelMultiAxpy)
                         r<T>::value);
 }
 
+
 TYPED_TEST(Gmres, KernelMultiDot)
 {
     using T = typename TestFixture::value_type;
@@ -599,6 +595,7 @@ TYPED_TEST(Gmres, KernelMultiDot)
                         r<T>::value);
 }
 
+
 TYPED_TEST(Gmres, SolvesStencilSystem)
 {
     using Mtx = typename TestFixture::Mtx;
@@ -968,6 +965,7 @@ TYPED_TEST(Gmres, SolvesWithPreconditioner)
     }
 }
 
+
 TYPED_TEST(Gmres, SolvesWithRgs)
 {
     using Mtx = typename TestFixture::Mtx;
@@ -1009,6 +1007,7 @@ TYPED_TEST(Gmres, SolvesWithRgs)
         r<value_type>::value * 1e3);
 }
 
+
 TYPED_TEST(Gmres, SolvesWithRgsBigAndPreconditioner)
 {
     using Mtx = typename TestFixture::Mtx;