From 6df1722b2dfe9688604ac01260ae43de7773ccf2 Mon Sep 17 00:00:00 2001
From: "Yu-Hsiang M. Tsai" <yhmtsai@gmail.com>
Date: Wed, 25 Sep 2024 16:40:08 +0200
Subject: [PATCH] fix infinite loop of lookup_hash_unsafe and add test in
 reference

---
 core/matrix/csr_lookup.hpp                  | 14 +++--
 reference/test/factorization/lu_kernels.cpp | 68 +++++++++++++++++++++
 2 files changed, 77 insertions(+), 5 deletions(-)

diff --git a/core/matrix/csr_lookup.hpp b/core/matrix/csr_lookup.hpp
index a7b687c3618..50d9c087b68 100644
--- a/core/matrix/csr_lookup.hpp
+++ b/core/matrix/csr_lookup.hpp
@@ -183,7 +183,8 @@ struct device_sparsity_lookup {
             result = lookup_search_unsafe(col);
             break;
         }
-        GKO_ASSERT(local_cols[result] == col);
+        GKO_ASSERT(result >= 0 && result < row_nnz &&
+                   local_cols[result] == col);
         return result;
     }
 
@@ -230,7 +231,8 @@ struct device_sparsity_lookup {
         const auto out_idx =
             block_bases[block] +
             gko::detail::popcount(block_bitmaps[block] & prefix_mask);
-        GKO_ASSERT(local_cols[out_idx] == col);
+        GKO_ASSERT(out_idx >= 0 && out_idx < row_nnz &&
+                   local_cols[out_idx] == col);
         return out_idx;
     }
 
@@ -262,15 +264,17 @@ struct device_sparsity_lookup {
             (static_cast<unsigned_index_type>(col) * hash_param) % hashmap_size;
         GKO_ASSERT(hashmap[hash] >= 0);
         GKO_ASSERT(hashmap[hash] < row_nnz);
-        while (local_cols[hashmap[hash]] != col) {
+        auto out_idx = hashmap[hash];
+        // linear probing with invalid_index sentinel to avoid infinite loop
+        while (out_idx >= 0 && local_cols[out_idx] != col) {
             hash++;
             if (hash >= hashmap_size) {
                 hash = 0;
             }
-            GKO_ASSERT(hashmap[hash] >= 0);
+            out_idx = hashmap[hash];
             GKO_ASSERT(hashmap[hash] < row_nnz);
         }
-        const auto out_idx = hashmap[hash];
+        // out_idx is either correct or invalid_index, the hashmap sentinel
         return out_idx;
     }
 
diff --git a/reference/test/factorization/lu_kernels.cpp b/reference/test/factorization/lu_kernels.cpp
index f4a8b240b38..d027968e97f 100644
--- a/reference/test/factorization/lu_kernels.cpp
+++ b/reference/test/factorization/lu_kernels.cpp
@@ -36,6 +36,8 @@ class Lu : public ::testing::Test {
     using index_type =
         typename std::tuple_element<1, decltype(ValueIndexType())>::type;
     using matrix_type = gko::matrix::Csr<value_type, index_type>;
+    using sparsity_pattern_type =
+        gko::matrix::SparsityCsr<value_type, index_type>;
 
     Lu()
         : ref(gko::ReferenceExecutor::create()),
@@ -329,3 +331,69 @@ TYPED_TEST(Lu, FactorizeWithKnownSparsityWorks)
         ASSERT_EQ(lu->get_diagonal(), nullptr);
     });
 }
+
+
+TYPED_TEST(Lu, GenerateIluWithBitmapIsEquivalentToRef)
+{
+    using value_type = typename TestFixture::value_type;
+    using index_type = typename TestFixture::index_type;
+    using matrix_type = typename TestFixture::matrix_type;
+    using sparsity_pattern_type = typename TestFixture::sparsity_pattern_type;
+    // diag + full first row and column
+    // the third and forth row use bitmap for lookup table
+    auto mtx = gko::share(gko::initialize<matrix_type>({{1.0, 1.0, 1.0, 1.0},
+                                                        {1.0, 1.0, 0.0, 0.0},
+                                                        {1.0, 0.0, 1.0, 0.0},
+                                                        {1.0, 0.0, 0.0, 1.0}},
+                                                       this->ref));
+    auto sparsity = gko::share(sparsity_pattern_type::create(this->ref));
+    mtx->convert_to(sparsity);
+    auto result = gko::initialize<matrix_type>({{1.0, 1.0, 1.0, 1.0},
+                                                {1.0, 0.0, 0.0, 0.0},
+                                                {1.0, 0.0, 0.0, 0.0},
+                                                {1.0, 0.0, 0.0, 0.0}},
+                                               this->ref);
+    auto factory =
+        gko::experimental::factorization::Lu<value_type, index_type>::build()
+            .with_symbolic_factorization(sparsity)
+            .on(this->ref);
+
+    auto lu = factory->generate(mtx);
+
+    GKO_ASSERT_MTX_EQ_SPARSITY(lu->get_combined(), mtx);
+    GKO_ASSERT_MTX_NEAR(lu->get_combined(), result, r<value_type>::value);
+}
+
+
+TYPED_TEST(Lu, GenerateIluWithHashmapIsEquivalentToRef)
+{
+    using value_type = typename TestFixture::value_type;
+    using index_type = typename TestFixture::index_type;
+    using matrix_type = typename TestFixture::matrix_type;
+    using sparsity_pattern_type = typename TestFixture::sparsity_pattern_type;
+    int n = 68;
+    // the first row and second last row use hashmap for lookup table
+    gko::matrix_data<value_type, index_type> data(gko::dim<2>(n, n));
+    for (int i = 0; i < n; i++) {
+        data.nonzeros.emplace_back(i, i, gko::one<value_type>());
+    }
+    // add dependence
+    data.nonzeros.emplace_back(n - 3, 0, gko::one<value_type>());
+    // add a entry whose col idx is not shown in the above row
+    data.nonzeros.emplace_back(0, n - 2, gko::one<value_type>());
+    data.sort_row_major();
+    auto mtx = gko::share(matrix_type::create(this->ref));
+    mtx->read(data);
+    auto sparsity = gko::share(sparsity_pattern_type::create(this->ref));
+    mtx->convert_to(sparsity);
+    auto factory =
+        gko::experimental::factorization::Lu<value_type, index_type>::build()
+            .with_symbolic_factorization(sparsity)
+            .on(this->ref);
+
+    auto lu = factory->generate(mtx);
+
+    // the result combined matrix is the same as the original matrix
+    GKO_ASSERT_MTX_EQ_SPARSITY(lu->get_combined(), mtx);
+    GKO_ASSERT_MTX_NEAR(lu->get_combined(), mtx, r<value_type>::value);
+}
\ No newline at end of file