diff --git a/taichi/program/sparse_solver.cpp b/taichi/program/sparse_solver.cpp index 5e4444d29accb..caa11ac5e7d08 100644 --- a/taichi/program/sparse_solver.cpp +++ b/taichi/program/sparse_solver.cpp @@ -4,23 +4,80 @@ #include -#define MAKE_SOLVER(dt, type, order) \ - { \ - {#dt, #type, #order}, []() -> std::unique_ptr { \ - using T = Eigen::Simplicial##type, Eigen::Lower, \ - Eigen::order##Ordering>; \ - return std::make_unique< \ - EigenSparseSolver>>(); \ - } \ - } +namespace taichi::lang { +#define EIGEN_LLT_SOLVER_INSTANTIATION(dt, type, order) \ + template class EigenSparseSolver< \ + Eigen::Simplicial##type, Eigen::Lower, \ + Eigen::order##Ordering>, \ + Eigen::SparseMatrix
>; +#define EIGEN_LU_SOLVER_INSTANTIATION(dt, type, order) \ + template class EigenSparseSolver< \ + Eigen::Sparse##type, \ + Eigen::order##Ordering>, \ + Eigen::SparseMatrix
>; +// Explicit instantiation of EigenSparseSolver +EIGEN_LLT_SOLVER_INSTANTIATION(float32, LLT, AMD); +EIGEN_LLT_SOLVER_INSTANTIATION(float32, LLT, COLAMD); +EIGEN_LLT_SOLVER_INSTANTIATION(float32, LDLT, AMD); +EIGEN_LLT_SOLVER_INSTANTIATION(float32, LDLT, COLAMD); +EIGEN_LU_SOLVER_INSTANTIATION(float32, LU, AMD); +EIGEN_LU_SOLVER_INSTANTIATION(float32, LU, COLAMD); +EIGEN_LLT_SOLVER_INSTANTIATION(float64, LLT, AMD); +EIGEN_LLT_SOLVER_INSTANTIATION(float64, LLT, COLAMD); +EIGEN_LLT_SOLVER_INSTANTIATION(float64, LDLT, AMD); +EIGEN_LLT_SOLVER_INSTANTIATION(float64, LDLT, COLAMD); +EIGEN_LU_SOLVER_INSTANTIATION(float64, LU, AMD); +EIGEN_LU_SOLVER_INSTANTIATION(float64, LU, COLAMD); +} // namespace taichi::lang -#define INSTANTIATE_SOLVER(dt, type, order) \ - using dt##type##order = \ +// Explicit instantiation of the template class EigenSparseSolver::solve +#define EIGEN_LLT_SOLVE_INSTANTIATION(dt, type, order, df) \ + using T##dt = Eigen::VectorX##df; \ + using S##dt##type##order = \ Eigen::Simplicial##type, Eigen::Lower, \ Eigen::order##Ordering>; \ - template void \ - EigenSparseSolver>::solve_rf( \ - Program *prog, const SparseMatrix &sm, const Ndarray &b, Ndarray &x); + template T##dt \ + EigenSparseSolver>::solve( \ + const T##dt &b); +#define EIGEN_LU_SOLVE_INSTANTIATION(dt, type, order, df) \ + using LUT##dt = Eigen::VectorX##df; \ + using LUS##dt##type##order = \ + Eigen::Sparse##type, \ + Eigen::order##Ordering>; \ + template LUT##dt \ + EigenSparseSolver>::solve( \ + const LUT##dt &b); + +// Explicit instantiation of the template class EigenSparseSolver::solve_rf +#define INSTANTIATE_LLT_SOLVE_RF(dt, type, order, df) \ + using llt##dt##type##order = \ + Eigen::Simplicial##type, Eigen::Lower, \ + Eigen::order##Ordering>; \ + template void EigenSparseSolver>::solve_rf( \ + Program * prog, const SparseMatrix &sm, const Ndarray &b, \ + const Ndarray &x); + +#define INSTANTIATE_LU_SOLVE_RF(dt, type, order, df) \ + using lu##dt##type##order = \ + Eigen::Sparse##type, \ + Eigen::order##Ordering>; \ + template void EigenSparseSolver>::solve_rf( \ + Program * prog, const SparseMatrix &sm, const Ndarray &b, \ + const Ndarray &x); + +#define MAKE_EIGEN_SOLVER(dt, type, order) \ + std::make_unique() + +#define MAKE_SOLVER(dt, type, order) \ + { \ + {#dt, #type, #order}, []() -> std::unique_ptr { \ + return MAKE_EIGEN_SOLVER(dt, type, order); \ + } \ + } using Triplets = std::tuple; namespace { @@ -70,32 +127,53 @@ void EigenSparseSolver::factorize( } template -Eigen::VectorXf EigenSparseSolver::solve( - const Eigen::Ref &b) { +template +T EigenSparseSolver::solve(const T &b) { return solver_.solve(b); } +EIGEN_LLT_SOLVE_INSTANTIATION(float32, LLT, AMD, f); +EIGEN_LLT_SOLVE_INSTANTIATION(float32, LLT, COLAMD, f); +EIGEN_LLT_SOLVE_INSTANTIATION(float32, LDLT, AMD, f); +EIGEN_LLT_SOLVE_INSTANTIATION(float32, LDLT, COLAMD, f); +EIGEN_LU_SOLVE_INSTANTIATION(float32, LU, AMD, f); +EIGEN_LU_SOLVE_INSTANTIATION(float32, LU, COLAMD, f); +EIGEN_LLT_SOLVE_INSTANTIATION(float64, LLT, AMD, d); +EIGEN_LLT_SOLVE_INSTANTIATION(float64, LLT, COLAMD, d); +EIGEN_LLT_SOLVE_INSTANTIATION(float64, LDLT, AMD, d); +EIGEN_LLT_SOLVE_INSTANTIATION(float64, LDLT, COLAMD, d); +EIGEN_LU_SOLVE_INSTANTIATION(float64, LU, AMD, d); +EIGEN_LU_SOLVE_INSTANTIATION(float64, LU, COLAMD, d); + template bool EigenSparseSolver::info() { return solver_.info() == Eigen::Success; } template +template void EigenSparseSolver::solve_rf( Program *prog, const SparseMatrix &sm, const Ndarray &b, - Ndarray &x) { + const Ndarray &x) { size_t db = prog->get_ndarray_data_ptr_as_int(&b); size_t dX = prog->get_ndarray_data_ptr_as_int(&x); - Eigen::Map((float *)dX, rows_) = - solver_.solve(Eigen::Map((float *)db, cols_)); + Eigen::Map((V *)dX, rows_) = solver_.solve(Eigen::Map((V *)db, cols_)); } -INSTANTIATE_SOLVER(float32, LLT, COLAMD) -INSTANTIATE_SOLVER(float32, LDLT, COLAMD) -INSTANTIATE_SOLVER(float32, LLT, AMD) -INSTANTIATE_SOLVER(float32, LDLT, AMD) +INSTANTIATE_LLT_SOLVE_RF(float32, LLT, COLAMD, Eigen::VectorXf) +INSTANTIATE_LLT_SOLVE_RF(float32, LDLT, COLAMD, Eigen::VectorXf) +INSTANTIATE_LLT_SOLVE_RF(float32, LLT, AMD, Eigen::VectorXf) +INSTANTIATE_LLT_SOLVE_RF(float32, LDLT, AMD, Eigen::VectorXf) +INSTANTIATE_LU_SOLVE_RF(float32, LU, AMD, Eigen::VectorXf) +INSTANTIATE_LU_SOLVE_RF(float32, LU, COLAMD, Eigen::VectorXf) +INSTANTIATE_LLT_SOLVE_RF(float64, LLT, COLAMD, Eigen::VectorXd) +INSTANTIATE_LLT_SOLVE_RF(float64, LDLT, COLAMD, Eigen::VectorXd) +INSTANTIATE_LLT_SOLVE_RF(float64, LLT, AMD, Eigen::VectorXd) +INSTANTIATE_LLT_SOLVE_RF(float64, LDLT, AMD, Eigen::VectorXd) +INSTANTIATE_LU_SOLVE_RF(float64, LU, AMD, Eigen::VectorXd) +INSTANTIATE_LU_SOLVE_RF(float64, LU, COLAMD, Eigen::VectorXd) CuSparseSolver::CuSparseSolver() { #if defined(TI_WITH_CUDA) @@ -184,7 +262,7 @@ void CuSparseSolver::factorize(const SparseMatrix &sm) { void CuSparseSolver::solve_cu(Program *prog, const SparseMatrix &sm, const Ndarray &b, - Ndarray &x) { + const Ndarray &x) { #ifdef TI_WITH_CUDA cusparseHandle_t cusparseHandle = nullptr; CUSPARSEDriver::get_instance().cpCreate(&cusparseHandle); @@ -349,7 +427,7 @@ void CuSparseSolver::solve_cu(Program *prog, void CuSparseSolver::solve_rf(Program *prog, const SparseMatrix &sm, const Ndarray &b, - Ndarray &x) { + const Ndarray &x) { #if defined(TI_WITH_CUDA) if (is_analyzed_ == false) { analyze_pattern(sm); @@ -383,8 +461,10 @@ std::unique_ptr make_sparse_solver(DataType dt, using func_type = std::unique_ptr (*)(); static const std::unordered_map solver_factory = { - MAKE_SOLVER(float32, LLT, AMD), MAKE_SOLVER(float32, LLT, COLAMD), - MAKE_SOLVER(float32, LDLT, AMD), MAKE_SOLVER(float32, LDLT, COLAMD)}; + MAKE_SOLVER(float32, LLT, AMD), MAKE_SOLVER(float32, LLT, COLAMD), + MAKE_SOLVER(float32, LDLT, AMD), MAKE_SOLVER(float32, LDLT, COLAMD), + MAKE_SOLVER(float64, LLT, AMD), MAKE_SOLVER(float64, LLT, COLAMD), + MAKE_SOLVER(float64, LDLT, AMD), MAKE_SOLVER(float64, LDLT, COLAMD)}; static const std::unordered_map dt_map = { {"f32", "float32"}, {"f64", "float64"}}; auto it = dt_map.find(taichi::lang::data_type_name(dt)); @@ -397,9 +477,17 @@ std::unique_ptr make_sparse_solver(DataType dt, auto solver_func = solver_factory.at(solver_key); return solver_func(); } else if (solver_type == "LU") { - using EigenMatrix = Eigen::SparseMatrix; - using LU = Eigen::SparseLU; - return std::make_unique>(); + if (it->first == "f32") { + using EigenMatrix = Eigen::SparseMatrix; + using LU = Eigen::SparseLU; + return std::make_unique>(); + } else if (it->first == "f64") { + using EigenMatrix = Eigen::SparseMatrix; + using LU = Eigen::SparseLU; + return std::make_unique>(); + } else { + TI_ERROR("Not supported sparse solver data type: {}", it->second); + } } else TI_ERROR("Not supported sparse solver type: {}", solver_type); } diff --git a/taichi/program/sparse_solver.h b/taichi/program/sparse_solver.h index 876452cf93161..67f7796784f1d 100644 --- a/taichi/program/sparse_solver.h +++ b/taichi/program/sparse_solver.h @@ -1,10 +1,23 @@ #pragma once +#include "sparse_matrix.h" + #include "taichi/ir/type.h" #include "taichi/rhi/cuda/cuda_driver.h" #include "taichi/program/program.h" -#include "sparse_matrix.h" +#define DECLARE_EIGEN_LLT_SOLVER(dt, type, order) \ + typedef EigenSparseSolver< \ + Eigen::Simplicial##type, Eigen::Lower, \ + Eigen::order##Ordering>, \ + Eigen::SparseMatrix
> \ + EigenSparseSolver##dt##type##order; + +#define DECLARE_EIGEN_LU_SOLVER(dt, type, order) \ + typedef EigenSparseSolver, \ + Eigen::order##Ordering>, \ + Eigen::SparseMatrix
> \ + EigenSparseSolver##dt##type##order; namespace taichi::lang { @@ -25,15 +38,6 @@ class SparseSolver { virtual bool compute(const SparseMatrix &sm) = 0; virtual void analyze_pattern(const SparseMatrix &sm) = 0; virtual void factorize(const SparseMatrix &sm) = 0; - virtual Eigen::VectorXf solve(const Eigen::Ref &b) = 0; - virtual void solve_rf(Program *prog, - const SparseMatrix &sm, - const Ndarray &b, - Ndarray &x) = 0; - virtual void solve_cu(Program *prog, - const SparseMatrix &sm, - const Ndarray &b, - Ndarray &x) = 0; virtual bool info() = 0; }; @@ -47,21 +51,30 @@ class EigenSparseSolver : public SparseSolver { bool compute(const SparseMatrix &sm) override; void analyze_pattern(const SparseMatrix &sm) override; void factorize(const SparseMatrix &sm) override; - Eigen::VectorXf solve(const Eigen::Ref &b) override; - void solve_cu(Program *prog, - const SparseMatrix &sm, - const Ndarray &b, - Ndarray &x) override { - TI_NOT_IMPLEMENTED; - }; + template + T solve(const T &b); + + template void solve_rf(Program *prog, const SparseMatrix &sm, const Ndarray &b, - Ndarray &x) override; - + const Ndarray &x); bool info() override; }; +DECLARE_EIGEN_LLT_SOLVER(float32, LLT, AMD); +DECLARE_EIGEN_LLT_SOLVER(float32, LLT, COLAMD); +DECLARE_EIGEN_LLT_SOLVER(float32, LDLT, AMD); +DECLARE_EIGEN_LLT_SOLVER(float32, LDLT, COLAMD); +DECLARE_EIGEN_LU_SOLVER(float32, LU, AMD); +DECLARE_EIGEN_LU_SOLVER(float32, LU, COLAMD); +DECLARE_EIGEN_LLT_SOLVER(float64, LLT, AMD); +DECLARE_EIGEN_LLT_SOLVER(float64, LLT, COLAMD); +DECLARE_EIGEN_LLT_SOLVER(float64, LDLT, AMD); +DECLARE_EIGEN_LLT_SOLVER(float64, LDLT, COLAMD); +DECLARE_EIGEN_LU_SOLVER(float64, LU, AMD); +DECLARE_EIGEN_LU_SOLVER(float64, LU, COLAMD); + class CuSparseSolver : public SparseSolver { private: csrcholInfo_t info_{nullptr}; @@ -81,17 +94,14 @@ class CuSparseSolver : public SparseSolver { void analyze_pattern(const SparseMatrix &sm) override; void factorize(const SparseMatrix &sm) override; - Eigen::VectorXf solve(const Eigen::Ref &b) override { - TI_NOT_IMPLEMENTED; - }; void solve_cu(Program *prog, const SparseMatrix &sm, const Ndarray &b, - Ndarray &x) override; + const Ndarray &x); void solve_rf(Program *prog, const SparseMatrix &sm, const Ndarray &b, - Ndarray &x) override; + const Ndarray &x); bool info() override { TI_NOT_IMPLEMENTED; }; diff --git a/taichi/python/export_lang.cpp b/taichi/python/export_lang.cpp index f1f6e630a828d..e944cc4a7a23f 100644 --- a/taichi/python/export_lang.cpp +++ b/taichi/python/export_lang.cpp @@ -1259,11 +1259,43 @@ void export_lang(py::module &m) { .def("compute", &SparseSolver::compute) .def("analyze_pattern", &SparseSolver::analyze_pattern) .def("factorize", &SparseSolver::factorize) - .def("solve", &SparseSolver::solve) - .def("solve_cu", &SparseSolver::solve_cu) - .def("solve_rf", &SparseSolver::solve_rf) .def("info", &SparseSolver::info); +#define REGISTER_EIGEN_SOLVER(dt, type, order, fd) \ + py::class_( \ + m, "EigenSparseSolver" #dt #type #order) \ + .def("compute", &EigenSparseSolver##dt##type##order::compute) \ + .def("analyze_pattern", \ + &EigenSparseSolver##dt##type##order::analyze_pattern) \ + .def("factorize", &EigenSparseSolver##dt##type##order::factorize) \ + .def("solve", \ + &EigenSparseSolver##dt##type##order::solve) \ + .def("solve_rf", \ + &EigenSparseSolver##dt##type##order::solve_rf) \ + .def("info", &EigenSparseSolver##dt##type##order::info); + + REGISTER_EIGEN_SOLVER(float32, LLT, AMD, f) + REGISTER_EIGEN_SOLVER(float32, LLT, COLAMD, f) + REGISTER_EIGEN_SOLVER(float32, LDLT, AMD, f) + REGISTER_EIGEN_SOLVER(float32, LDLT, COLAMD, f) + REGISTER_EIGEN_SOLVER(float32, LU, AMD, f) + REGISTER_EIGEN_SOLVER(float32, LU, COLAMD, f) + REGISTER_EIGEN_SOLVER(float64, LLT, AMD, d) + REGISTER_EIGEN_SOLVER(float64, LLT, COLAMD, d) + REGISTER_EIGEN_SOLVER(float64, LDLT, AMD, d) + REGISTER_EIGEN_SOLVER(float64, LDLT, COLAMD, d) + REGISTER_EIGEN_SOLVER(float64, LU, AMD, d) + REGISTER_EIGEN_SOLVER(float64, LU, COLAMD, d) + + py::class_(m, "CuSparseSolver") + .def("compute", &CuSparseSolver::compute) + .def("analyze_pattern", &CuSparseSolver::analyze_pattern) + .def("factorize", &CuSparseSolver::factorize) + .def("solve_rf", &CuSparseSolver::solve_rf) + .def("solve_cu", &CuSparseSolver::solve_cu) + .def("info", &CuSparseSolver::info); + m.def("make_sparse_solver", &make_sparse_solver); m.def("make_cusparse_solver", &make_cusparse_solver); diff --git a/tests/python/test_sparse_linear_solver.py b/tests/python/test_sparse_linear_solver.py index b0e79b66d3bd1..7408c8da1981c 100644 --- a/tests/python/test_sparse_linear_solver.py +++ b/tests/python/test_sparse_linear_solver.py @@ -5,16 +5,19 @@ from tests import test_utils -@pytest.mark.parametrize("dtype", [ti.f32]) +@pytest.mark.parametrize("dtype", [ti.f32, ti.f64]) @pytest.mark.parametrize("solver_type", ["LLT", "LDLT", "LU"]) @pytest.mark.parametrize("ordering", ["AMD", "COLAMD"]) -@test_utils.test(arch=ti.cpu) +@test_utils.test(arch=ti.x64) def test_sparse_LLT_solver(dtype, solver_type, ordering): n = 10 A = np.random.rand(n, n) A_psd = np.dot(A, A.transpose()) - Abuilder = ti.linalg.SparseMatrixBuilder(n, n, max_num_triplets=300) - b = ti.field(ti.f32, shape=n) + Abuilder = ti.linalg.SparseMatrixBuilder(n, + n, + max_num_triplets=100, + dtype=dtype) + b = ti.field(dtype=dtype, shape=n) @ti.kernel def fill(Abuilder: ti.types.sparse_matrix_builder(),