add benchmark for MST-enhanced symbolic Cholesky

ginkgo-project · Dec 19, 2024 · c8f3ca1 · c8f3ca1
1 parent d9c6fe5
commit c8f3ca1
Show file tree

Hide file tree

Showing 3 changed files with 92 additions and 6 deletions.
diff --git a/benchmark/sparse_blas/operations.cpp b/benchmark/sparse_blas/operations.cpp
@@ -645,8 +645,9 @@ class SymbolicLuNearSymmOperation : public BenchmarkOperation {
 
 class SymbolicCholeskyOperation : public BenchmarkOperation {
 public:
-    explicit SymbolicCholeskyOperation(const Mtx* mtx, bool symmetric)
-        : mtx_{mtx}, symmetric_{symmetric}, result_{}
+    explicit SymbolicCholeskyOperation(const Mtx* mtx, bool device,
+                                       bool symmetric)
+        : mtx_{mtx}, symmetric_{symmetric}, device_{device}, result_{}
     {}
 
     std::pair<bool, double> validate() const override
@@ -674,8 +675,13 @@ class SymbolicCholeskyOperation : public BenchmarkOperation {
 
     void run() override
     {
-        gko::factorization::symbolic_cholesky(mtx_, symmetric_, result_,
-                                              forest_);
+        if (device_) {
+            gko::factorization::symbolic_cholesky_device(mtx_, symmetric_,
+                                                         result_, forest_);
+        } else {
+            gko::factorization::symbolic_cholesky(mtx_, symmetric_, result_,
+                                                  forest_);
+        }
     }
 
     void write_stats(json& object) override
@@ -686,6 +692,7 @@ class SymbolicCholeskyOperation : public BenchmarkOperation {
 private:
     const Mtx* mtx_;
     bool symmetric_;
+    bool device_;
     std::unique_ptr<Mtx> result_;
     std::unique_ptr<gko::factorization::elimination_forest<itype>> forest_;
 };
@@ -822,13 +829,25 @@ const std::map<std::string,
          [](const Mtx* mtx) {
              return std::make_unique<SymbolicLuNearSymmOperation>(mtx);
          }},
+        {"symbolic_cholesky_device",
+         [](const Mtx* mtx) {
+             return std::make_unique<SymbolicCholeskyOperation>(mtx, true,
+                                                                false);
+         }},
+        {"symbolic_cholesky_device_symmetric",
+         [](const Mtx* mtx) {
+             return std::make_unique<SymbolicCholeskyOperation>(mtx, true,
+                                                                true);
+         }},
         {"symbolic_cholesky",
          [](const Mtx* mtx) {
-             return std::make_unique<SymbolicCholeskyOperation>(mtx, false);
+             return std::make_unique<SymbolicCholeskyOperation>(mtx, false,
+                                                                false);
          }},
         {"symbolic_cholesky_symmetric",
          [](const Mtx* mtx) {
-             return std::make_unique<SymbolicCholeskyOperation>(mtx, true);
+             return std::make_unique<SymbolicCholeskyOperation>(mtx, false,
+                                                                true);
          }},
         {"reorder_rcm",
          [](const Mtx* mtx) {

diff --git a/core/factorization/symbolic.cpp b/core/factorization/symbolic.cpp
@@ -25,6 +25,7 @@ namespace factorization {
 namespace {
 
 
+GKO_REGISTER_OPERATION(compute_skeleton_tree, cholesky::compute_skeleton_tree);
 GKO_REGISTER_OPERATION(symbolic_count, cholesky::symbolic_count);
 GKO_REGISTER_OPERATION(symbolic, cholesky::symbolic_factorize);
 GKO_REGISTER_OPERATION(build_lookup_offsets, csr::build_lookup_offsets);
@@ -83,6 +84,56 @@ void symbolic_cholesky(
 GKO_INSTANTIATE_FOR_EACH_VALUE_AND_INDEX_TYPE(GKO_DECLARE_SYMBOLIC_CHOLESKY);
 
 
+template <typename ValueType, typename IndexType>
+void symbolic_cholesky_device(
+    const matrix::Csr<ValueType, IndexType>* mtx, bool symmetrize,
+    std::unique_ptr<matrix::Csr<ValueType, IndexType>>& factors,
+    std::unique_ptr<elimination_forest<IndexType>>& forest)
+{
+    using matrix_type = matrix::Csr<ValueType, IndexType>;
+    GKO_ASSERT_IS_SQUARE_MATRIX(mtx);
+    const auto exec = mtx->get_executor();
+    const auto num_rows = mtx->get_size()[0];
+    const auto host_exec = exec->get_master();
+    {
+        const auto skeleton =
+            matrix_type::create(exec, mtx->get_size(), num_rows);
+        exec->run(make_compute_skeleton_tree(
+            mtx->get_const_row_ptrs(), mtx->get_const_col_idxs(), num_rows,
+            skeleton->get_row_ptrs(), skeleton->get_col_idxs()));
+        exec->run(make_compute_elim_forest(skeleton.get(), forest));
+    }
+    array<IndexType> row_ptrs{exec, num_rows + 1};
+    array<IndexType> tmp{exec};
+    exec->run(make_symbolic_count(mtx, *forest, row_ptrs.get_data(), tmp));
+    exec->run(make_prefix_sum_nonnegative(row_ptrs.get_data(), num_rows + 1));
+    const auto factor_nnz =
+        static_cast<size_type>(get_element(row_ptrs, num_rows));
+    factors = matrix_type::create(
+        exec, mtx->get_size(), array<ValueType>{exec, factor_nnz},
+        array<IndexType>{exec, factor_nnz}, std::move(row_ptrs));
+    exec->run(make_symbolic(mtx, *forest, factors.get(), tmp));
+    factors->sort_by_column_index();
+    if (symmetrize) {
+        auto lt_factor = as<matrix_type>(factors->transpose());
+        const auto scalar =
+            initialize<matrix::Dense<ValueType>>({one<ValueType>()}, exec);
+        const auto id = matrix::Identity<ValueType>::create(exec, num_rows);
+        lt_factor->apply(scalar, id, scalar, factors);
+    }
+}
+
+
+#define GKO_DECLARE_SYMBOLIC_CHOLESKY_DEVICE(ValueType, IndexType)     \
+    void symbolic_cholesky_device(                                     \
+        const matrix::Csr<ValueType, IndexType>* mtx, bool symmetrize, \
+        std::unique_ptr<matrix::Csr<ValueType, IndexType>>& factors,   \
+        std::unique_ptr<factorization::elimination_forest<IndexType>>& forest)
+
+GKO_INSTANTIATE_FOR_EACH_VALUE_AND_INDEX_TYPE(
+    GKO_DECLARE_SYMBOLIC_CHOLESKY_DEVICE);
+
+
 template <typename ValueType, typename IndexType>
 void symbolic_lu_near_symm(
     const matrix::Csr<ValueType, IndexType>* mtx,

diff --git a/core/factorization/symbolic.hpp b/core/factorization/symbolic.hpp
@@ -25,6 +25,22 @@ void symbolic_cholesky(
     std::unique_ptr<matrix::Csr<ValueType, IndexType>>& factors,
     std::unique_ptr<elimination_forest<IndexType>>& forest);
 
+
+/**
+ * Computes the symbolic Cholesky factorization of the given matrix on the
+ * device.
+ *
+ * @param mtx  the input matrix
+ * @param symmetrize  output the pattern of L + L^T (true) or just L (false)?
+ * @param factors  the output factor(s)
+ * @param forest  the elimination forest of the input matrix
+ */
+template <typename ValueType, typename IndexType>
+void symbolic_cholesky_device(
+    const matrix::Csr<ValueType, IndexType>* mtx, bool symmetrize,
+    std::unique_ptr<matrix::Csr<ValueType, IndexType>>& factors,
+    std::unique_ptr<elimination_forest<IndexType>>& forest);
+
 /**
  * Computes the symbolic LU factorization of the given matrix.
  *