From ce50e7a83a5a14ea3f08f7c2eb2063fc17505b66 Mon Sep 17 00:00:00 2001 From: Sam Mish Date: Fri, 5 Jan 2024 11:37:28 -0800 Subject: [PATCH] add optional argument for specifying factorization strategy --- src/cholesky_solver.cpp | 9 ++++----- src/cholesky_solver.h | 5 +++-- src/docstr.h | 2 ++ src/main.cpp | 8 +++++--- 4 files changed, 14 insertions(+), 10 deletions(-) diff --git a/src/cholesky_solver.cpp b/src/cholesky_solver.cpp index b2d0057..9b68e60 100644 --- a/src/cholesky_solver.cpp +++ b/src/cholesky_solver.cpp @@ -149,8 +149,7 @@ void csc_sum_duplicates(int n_rows, int &m_nnz, int **col_ptr, int **rows, doubl } template -CholeskySolver::CholeskySolver(int n_rows, int nnz, int *ii, int *jj, double *x, MatrixType type, bool cpu) : m_n(n_rows), m_nnz(nnz), m_cpu(cpu) { - +CholeskySolver::CholeskySolver(int n_rows, int nnz, int *ii, int *jj, double *x, MatrixType type, bool cpu, int strategy) : m_n(n_rows), m_nnz(nnz), m_cpu(cpu) { // Placeholders for the CSC matrix data int *col_ptr, *rows; @@ -188,7 +187,7 @@ CholeskySolver::CholeskySolver(int n_rows, int nnz, int *ii, int *jj, dou } // Run the Cholesky factorization through CHOLMOD and run the analysis - factorize(col_ptr, rows, data); + factorize(col_ptr, rows, data, strategy); if (type != MatrixType::CSC) { free(col_ptr); @@ -198,12 +197,12 @@ CholeskySolver::CholeskySolver(int n_rows, int nnz, int *ii, int *jj, dou } template -void CholeskySolver::factorize(int *col_ptr, int *rows, double *data) { +void CholeskySolver::factorize(int *col_ptr, int *rows, double *data, int strategy) { cholmod_sparse *A; cholmod_start(&m_common); - m_common.supernodal = CHOLMOD_SIMPLICIAL; + m_common.supernodal = strategy; m_common.final_ll = 1; // compute LL' factorization instead of LDL´ (default for simplicial) m_common.nmethods = 1; m_common.method[0].ordering = CHOLMOD_NESDIS; diff --git a/src/cholesky_solver.h b/src/cholesky_solver.h index 1c98651..f7c1660 100644 --- a/src/cholesky_solver.h +++ b/src/cholesky_solver.h @@ -32,8 +32,9 @@ class CholeskySolver { * @param x Array of nonzero entries * @param type The type of the matrix representation. Can be COO, CSC or CSR * @param cpu Whether or not to run the CPU version of the solver. + * @param strategy 0: SIMPLICIAL, 1: AUTO, 2: SUPERNODAL */ - CholeskySolver(int n_rows, int nnz, int *ii, int *jj, double *x, MatrixType type, bool cpu); + CholeskySolver(int n_rows, int nnz, int *ii, int *jj, double *x, MatrixType type, bool cpu, int strategy); ~CholeskySolver(); @@ -49,7 +50,7 @@ class CholeskySolver { private: // Factorize the CSC matrix using CHOLMOD - void factorize(int *col_ptr, int *rows, double *data); + void factorize(int *col_ptr, int *rows, double *data, int strategy); // Run the analysis of a triangular matrix obtained through Cholesky void analyze_cuda(int n_rows, int n_entries, void *csr_rows, void *csr_cols, Float *csr_data, bool lower); diff --git a/src/docstr.h b/src/docstr.h index 872bee3..4cf297c 100644 --- a/src/docstr.h +++ b/src/docstr.h @@ -16,6 +16,8 @@ Parameters - x - The array of nonzero entries. - type - The matrix representation type, of type MatrixType. Available types are MatrixType.COO, MatrixType.CSC and MatrixType.CSR. +- strategy - the factorization strategy to be used by CHOLMOD. Available options are + 0 (Simplicial), 1 (automatic selectiong, default), 2 (Supernodal). )"; const char *doc_matrix_type = diff --git a/src/main.cpp b/src/main.cpp index 8fd137c..58e0640 100644 --- a/src/main.cpp +++ b/src/main.cpp @@ -23,7 +23,8 @@ void declare_cholesky(nb::module_ &m, const std::string &typestr, const char *do nb::ndarray, nb::c_contig> ii, nb::ndarray, nb::c_contig> jj, nb::ndarray, nb::c_contig> x, - MatrixType type) { + MatrixType type, + int strategy) { if (type == MatrixType::COO){ if (ii.shape(0) != jj.shape(0)) @@ -60,14 +61,14 @@ void declare_cholesky(nb::module_ &m, const std::string &typestr, const char *do cuda_check(cuMemcpyAsync((CUdeviceptr) indices_b, (CUdeviceptr) jj.data(), jj.shape(0)*sizeof(int), 0)); cuda_check(cuMemcpyAsync((CUdeviceptr) data, (CUdeviceptr) x.data(), x.shape(0)*sizeof(double), 0)); - new (self) Class(n_rows, x.shape(0), indices_a, indices_b, data, type, false); + new (self) Class(n_rows, x.shape(0), indices_a, indices_b, data, type, false, strategy); free(indices_a); free(indices_b); free(data); } else if (ii.device_type() == nb::device::cpu::value) { // CPU init - new (self) Class(n_rows, x.shape(0), (int *) ii.data(), (int *) jj.data(), (double *) x.data(), type, true); + new (self) Class(n_rows, x.shape(0), (int *) ii.data(), (int *) jj.data(), (double *) x.data(), type, true, strategy); } else throw std::invalid_argument("Unsupported input device! Only CPU and CUDA arrays are supported."); }, @@ -76,6 +77,7 @@ void declare_cholesky(nb::module_ &m, const std::string &typestr, const char *do nb::arg("jj"), nb::arg("x"), nb::arg("type"), + nb::arg("strategy") = 1, doc_constructor) .def("solve", [](Class &self, nb::ndarray b,