add checked_lookup into LU

common/cuda_hip/factorization/lu_kernels.cpp

@@ -85,7 +85,7 @@ __global__ __launch_bounds__(default_block_size) void initialize(
 }
 
 
-template <typename ValueType, typename IndexType>
+template <bool checked_lookup, typename ValueType, typename IndexType>
 __global__ __launch_bounds__(default_block_size) void factorize(
     const IndexType* __restrict__ row_ptrs, const IndexType* __restrict__ cols,
     const IndexType* __restrict__ storage_offsets,
@@ -130,8 +130,16 @@ __global__ __launch_bounds__(default_block_size) void factorize(
              upper_nz += config::warp_size) {
             const auto upper_col = cols[upper_nz];
             const auto upper_val = vals[upper_nz];
-            const auto output_pos = lookup.lookup_unsafe(upper_col) + row_begin;
-            vals[output_pos] -= scale * upper_val;
+            if (checked_lookup) {
+                const auto pos = lookup[upper_col];
+                if (pos != invalid_index<IndexType>()) {
+                    vals[row_begin + pos] -= scale * upper_val;
+                }
+            } else {
+                const auto output_pos =
+                    lookup.lookup_unsafe(upper_col) + row_begin;
+                vals[output_pos] -= scale * upper_val;
+            }
         }
     }
     scheduler.mark_ready();
@@ -252,19 +260,29 @@ template <typename ValueType, typename IndexType>
 void factorize(std::shared_ptr<const DefaultExecutor> exec,
                const IndexType* lookup_offsets, const int64* lookup_descs,
                const int32* lookup_storage, const IndexType* diag_idxs,
-               matrix::Csr<ValueType, IndexType>* factors,
+               matrix::Csr<ValueType, IndexType>* factors, bool checked_lookup,
                array<int>& tmp_storage)
 {
     const auto num_rows = factors->get_size()[0];
     if (num_rows > 0) {
         syncfree_storage storage(exec, tmp_storage, num_rows);
         const auto num_blocks =
             ceildiv(num_rows, default_block_size / config::warp_size);
-        kernel::factorize<<<num_blocks, default_block_size, 0,
-                            exec->get_stream()>>>(
-            factors->get_const_row_ptrs(), factors->get_const_col_idxs(),
-            lookup_offsets, lookup_storage, lookup_descs, diag_idxs,
-            as_device_type(factors->get_values()), storage, num_rows);
+        if (checked_lookup) {
+            kernel::factorize<true>
+                <<<num_blocks, default_block_size, 0, exec->get_stream()>>>(
+                    factors->get_const_row_ptrs(),
+                    factors->get_const_col_idxs(), lookup_offsets,
+                    lookup_storage, lookup_descs, diag_idxs,
+                    as_device_type(factors->get_values()), storage, num_rows);
+        } else {
+            kernel::factorize<false>
+                <<<num_blocks, default_block_size, 0, exec->get_stream()>>>(
+                    factors->get_const_row_ptrs(),
+                    factors->get_const_col_idxs(), lookup_offsets,
+                    lookup_storage, lookup_descs, diag_idxs,
+                    as_device_type(factors->get_values()), storage, num_rows);
+        }
     }
 }
 

core/factorization/lu.cpp

@@ -69,6 +69,9 @@ Lu<ValueType, IndexType>::parse(const config::pnode& config,
     if (auto& obj = config.get("skip_sorting")) {
         params.with_skip_sorting(config::get_value<bool>(obj));
     }
+    if (auto& obj = config.get("checked_lookup")) {
+        params.with_checked_lookup(config::get_value<bool>(obj));
+    }
 
     return params;
 }
@@ -160,10 +163,10 @@ std::unique_ptr<LinOp> Lu<ValueType, IndexType>::generate_impl(
         storage.get_const_data(), diag_idxs.get_data(), factors.get()));
     // run numerical factorization
     array<int> tmp{exec};
-    exec->run(make_factorize(storage_offsets.get_const_data(),
-                             row_descs.get_const_data(),
-                             storage.get_const_data(),
-                             diag_idxs.get_const_data(), factors.get(), tmp));
+    exec->run(make_factorize(
+        storage_offsets.get_const_data(), row_descs.get_const_data(),
+        storage.get_const_data(), diag_idxs.get_const_data(), factors.get(),
+        parameters_.checked_lookup, tmp));
     return factorization_type::create_from_combined_lu(std::move(factors));
 }
 

core/factorization/lu_kernels.hpp

@@ -33,7 +33,7 @@ namespace kernels {
                    const IndexType* lookup_offsets, const int64* lookup_descs, \
                    const int32* lookup_storage, const IndexType* diag_idxs,    \
                    matrix::Csr<ValueType, IndexType>* factors,                 \
-                   array<int>& tmp_storage)
+                   bool checked_lookup, array<int>& tmp_storage)
 
 
 #define GKO_DECLARE_LU_SYMMETRIC_FACTORIZE_SIMPLE(IndexType)                  \

core/test/config/factorization.cpp

@@ -178,6 +178,8 @@ struct Lu : FactorizationConfigTest<
             gko::experimental::factorization::symbolic_type::near_symmetric);
         config_map["skip_sorting"] = pnode{true};
         param.with_skip_sorting(true);
+        config_map["checked_lookup"] = pnode{true};
+        param.with_checked_lookup(true);
     }
 
     template <typename AnswerType>
@@ -190,6 +192,7 @@ struct Lu : FactorizationConfigTest<
                   ans_param.symbolic_factorization);
         ASSERT_EQ(res_param.symbolic_algorithm, ans_param.symbolic_algorithm);
         ASSERT_EQ(res_param.skip_sorting, ans_param.skip_sorting);
+        ASSERT_EQ(res_param.checked_lookup, ans_param.checked_lookup);
     }
 };
 

dpcpp/factorization/lu_kernels.dp.cpp

@@ -39,7 +39,7 @@ template <typename ValueType, typename IndexType>
 void factorize(std::shared_ptr<const DefaultExecutor> exec,
                const IndexType* lookup_offsets, const int64* lookup_descs,
                const int32* lookup_storage, const IndexType* diag_idxs,
-               matrix::Csr<ValueType, IndexType>* factors,
+               matrix::Csr<ValueType, IndexType>* factors, bool checked_lookup,
                array<int>& tmp_storage) GKO_NOT_IMPLEMENTED;
 
 GKO_INSTANTIATE_FOR_EACH_VALUE_AND_INDEX_TYPE(GKO_DECLARE_LU_FACTORIZE);

include/ginkgo/core/factorization/lu.hpp

@@ -97,6 +97,15 @@ class Lu
          * incorrect results or crash.
          */
         bool GKO_FACTORY_PARAMETER_SCALAR(skip_sorting, false);
+
+        /**
+         * The symbolic factoization should contains the fill-in information. If
+         * it is not the case (like Ilu), users might face hang or illegal
+         * access issue. Please enable this option when the symbolic
+         * factorization does not contain the full fill-in information. Symbolic
+         * factorization must still contain the entry for the original matrix.
+         */
+        bool GKO_FACTORY_PARAMETER_SCALAR(checked_lookup, false);
     };
 
     /**

omp/factorization/lu_kernels.cpp

@@ -66,7 +66,7 @@ template <typename ValueType, typename IndexType>
 void factorize(std::shared_ptr<const DefaultExecutor> exec,
                const IndexType* lookup_offsets, const int64* lookup_descs,
                const int32* lookup_storage, const IndexType* diag_idxs,
-               matrix::Csr<ValueType, IndexType>* factors,
+               matrix::Csr<ValueType, IndexType>* factors, bool checked_lookup,
                array<int>& tmp_storage)
 {
     const auto num_rows = factors->get_size()[0];
@@ -89,8 +89,15 @@ void factorize(std::shared_ptr<const DefaultExecutor> exec,
             for (auto dep_nz = dep_diag_idx + 1; dep_nz < dep_end; dep_nz++) {
                 const auto col = cols[dep_nz];
                 const auto val = vals[dep_nz];
-                const auto nz = row_begin + lookup.lookup_unsafe(col);
-                vals[nz] -= scale * val;
+                if (checked_lookup) {
+                    const auto idx = lookup[col];
+                    if (idx != invalid_index<IndexType>()) {
+                        vals[row_begin + idx] -= scale * val;
+                    }
+                } else {
+                    const auto nz = row_begin + lookup.lookup_unsafe(col);
+                    vals[nz] -= scale * val;
+                }
             }
         }
     }
@@ -185,4 +192,4 @@ GKO_INSTANTIATE_FOR_EACH_INDEX_TYPE(
 }  // namespace lu_factorization
 }  // namespace omp
 }  // namespace kernels
-}  // namespace gko
+}  // namespace gko

reference/factorization/lu_kernels.cpp

@@ -65,7 +65,7 @@ template <typename ValueType, typename IndexType>
 void factorize(std::shared_ptr<const DefaultExecutor> exec,
                const IndexType* lookup_offsets, const int64* lookup_descs,
                const int32* lookup_storage, const IndexType* diag_idxs,
-               matrix::Csr<ValueType, IndexType>* factors,
+               matrix::Csr<ValueType, IndexType>* factors, bool checked_lookup,
                array<int>& tmp_storage)
 {
     const auto num_rows = factors->get_size()[0];
@@ -87,8 +87,15 @@ void factorize(std::shared_ptr<const DefaultExecutor> exec,
             for (auto dep_nz = dep_diag_idx + 1; dep_nz < dep_end; dep_nz++) {
                 const auto col = cols[dep_nz];
                 const auto val = vals[dep_nz];
-                const auto nz = row_begin + lookup.lookup_unsafe(col);
-                vals[nz] -= scale * val;
+                if (checked_lookup) {
+                    const auto idx = lookup[col];
+                    if (idx != invalid_index<IndexType>()) {
+                        vals[row_begin + idx] -= scale * val;
+                    }
+                } else {
+                    const auto nz = row_begin + lookup.lookup_unsafe(col);
+                    vals[nz] -= scale * val;
+                }
             }
         }
     }
@@ -182,4 +189,4 @@ GKO_INSTANTIATE_FOR_EACH_INDEX_TYPE(
 }  // namespace lu_factorization
 }  // namespace reference
 }  // namespace kernels
-}  // namespace gko
+}  // namespace gko

reference/test/factorization/lu_kernels.cpp

@@ -218,7 +218,7 @@ TYPED_TEST(Lu, KernelFactorizeWorks)
         gko::kernels::reference::lu_factorization::factorize(
             this->ref, this->storage_offsets.get_const_data(),
             this->row_descs.get_const_data(), this->storage.get_const_data(),
-            diag_idxs.get_const_data(), this->mtx_lu.get(), tmp);
+            diag_idxs.get_const_data(), this->mtx_lu.get(), false, tmp);
 
         GKO_ASSERT_MTX_NEAR(this->mtx_lu, mtx_lu_ref,
                             15 * r<value_type>::value);
@@ -356,6 +356,7 @@ TYPED_TEST(Lu, GenerateIluWithBitmapIsEquivalentToRef)
     auto factory =
         gko::experimental::factorization::Lu<value_type, index_type>::build()
             .with_symbolic_factorization(sparsity)
+            .with_checked_lookup(true)
             .on(this->ref);
 
     auto lu = factory->generate(mtx);
@@ -389,6 +390,7 @@ TYPED_TEST(Lu, GenerateIluWithHashmapIsEquivalentToRef)
     auto factory =
         gko::experimental::factorization::Lu<value_type, index_type>::build()
             .with_symbolic_factorization(sparsity)
+            .with_checked_lookup(true)
             .on(this->ref);
 
     auto lu = factory->generate(mtx);

test/factorization/lu_kernels.cpp

@@ -198,11 +198,11 @@ TYPED_TEST(Lu, KernelFactorizeIsEquivalentToRef)
         gko::kernels::reference::lu_factorization::factorize(
             this->ref, this->storage_offsets.get_const_data(),
             this->row_descs.get_const_data(), this->storage.get_const_data(),
-            diag_idxs.get_const_data(), this->mtx_lu.get(), tmp);
+            diag_idxs.get_const_data(), this->mtx_lu.get(), false, tmp);
         gko::kernels::GKO_DEVICE_NAMESPACE::lu_factorization::factorize(
             this->exec, this->dstorage_offsets.get_const_data(),
             this->drow_descs.get_const_data(), this->dstorage.get_const_data(),
-            ddiag_idxs.get_const_data(), this->dmtx_lu.get(), dtmp);
+            ddiag_idxs.get_const_data(), this->dmtx_lu.get(), false, dtmp);
 
         GKO_ASSERT_MTX_NEAR(this->mtx_lu, this->dmtx_lu, r<value_type>::value);
     });
@@ -376,10 +376,12 @@ TYPED_TEST(Lu, GenerateIluWithBitmapIsEquivalentToRef)
     auto factory =
         gko::experimental::factorization::Lu<value_type, index_type>::build()
             .with_symbolic_factorization(sparsity)
+            .with_checked_lookup(true)
             .on(this->ref);
     auto dfactory =
         gko::experimental::factorization::Lu<value_type, index_type>::build()
             .with_symbolic_factorization(dsparsity)
+            .with_checked_lookup(true)
             .on(this->exec);
 
     auto lu = factory->generate(mtx);
@@ -417,10 +419,12 @@ TYPED_TEST(Lu, GenerateIluWithHashmapIsEquivalentToRef)
     auto factory =
         gko::experimental::factorization::Lu<value_type, index_type>::build()
             .with_symbolic_factorization(sparsity)
+            .with_checked_lookup(true)
             .on(this->ref);
     auto dfactory =
         gko::experimental::factorization::Lu<value_type, index_type>::build()
             .with_symbolic_factorization(dsparsity)
+            .with_checked_lookup(true)
             .on(this->exec);
 
     auto lu = factory->generate(mtx);