From c8f3ca16cecac1bea70845bc5c3de425626c037d Mon Sep 17 00:00:00 2001 From: Tobias Ribizel Date: Thu, 19 Dec 2024 21:39:42 +0100 Subject: [PATCH] add benchmark for MST-enhanced symbolic Cholesky --- benchmark/sparse_blas/operations.cpp | 31 +++++++++++++---- core/factorization/symbolic.cpp | 51 ++++++++++++++++++++++++++++ core/factorization/symbolic.hpp | 16 +++++++++ 3 files changed, 92 insertions(+), 6 deletions(-) diff --git a/benchmark/sparse_blas/operations.cpp b/benchmark/sparse_blas/operations.cpp index 30f3b5a80fe..9394e07d7e3 100644 --- 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 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 result_; std::unique_ptr> forest_; }; @@ -822,13 +829,25 @@ const std::map(mtx); }}, + {"symbolic_cholesky_device", + [](const Mtx* mtx) { + return std::make_unique(mtx, true, + false); + }}, + {"symbolic_cholesky_device_symmetric", + [](const Mtx* mtx) { + return std::make_unique(mtx, true, + true); + }}, {"symbolic_cholesky", [](const Mtx* mtx) { - return std::make_unique(mtx, false); + return std::make_unique(mtx, false, + false); }}, {"symbolic_cholesky_symmetric", [](const Mtx* mtx) { - return std::make_unique(mtx, true); + return std::make_unique(mtx, false, + true); }}, {"reorder_rcm", [](const Mtx* mtx) { diff --git a/core/factorization/symbolic.cpp b/core/factorization/symbolic.cpp index 23f6b94cc14..6a8a3f44f2c 100644 --- 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 +void symbolic_cholesky_device( + const matrix::Csr* mtx, bool symmetrize, + std::unique_ptr>& factors, + std::unique_ptr>& forest) +{ + using matrix_type = matrix::Csr; + 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 row_ptrs{exec, num_rows + 1}; + array 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(get_element(row_ptrs, num_rows)); + factors = matrix_type::create( + exec, mtx->get_size(), array{exec, factor_nnz}, + array{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(factors->transpose()); + const auto scalar = + initialize>({one()}, exec); + const auto id = matrix::Identity::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* mtx, bool symmetrize, \ + std::unique_ptr>& factors, \ + std::unique_ptr>& forest) + +GKO_INSTANTIATE_FOR_EACH_VALUE_AND_INDEX_TYPE( + GKO_DECLARE_SYMBOLIC_CHOLESKY_DEVICE); + + template void symbolic_lu_near_symm( const matrix::Csr* mtx, diff --git a/core/factorization/symbolic.hpp b/core/factorization/symbolic.hpp index 096d8c998bc..84478e214e6 100644 --- a/core/factorization/symbolic.hpp +++ b/core/factorization/symbolic.hpp @@ -25,6 +25,22 @@ void symbolic_cholesky( std::unique_ptr>& factors, std::unique_ptr>& 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 +void symbolic_cholesky_device( + const matrix::Csr* mtx, bool symmetrize, + std::unique_ptr>& factors, + std::unique_ptr>& forest); + /** * Computes the symbolic LU factorization of the given matrix. *