From b31463f3b5cfb8b716a3f80b30c57440ffaed25e Mon Sep 17 00:00:00 2001 From: Matthew Rocklin Date: Mon, 13 Nov 2023 18:56:29 -0600 Subject: [PATCH 1/3] Add nogil declaration to _run_paradiso function This releases the Python GIL around the low level paradiso solver call. Hopefully this allows this library and dependent libraries (like SimPEG) to operate more smoothly in a multi-threaded environment. This should be safe from a Python object perspective because no Python objects are created/destroyed/inc-reffed during this call. It could still be that multiple concurrent calls of this function on the same data cause undesried results, but hopefully downstream users can apply concurrency primitives (like locks) as needed. --- pydiso/mkl_solver.pyx | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pydiso/mkl_solver.pyx b/pydiso/mkl_solver.pyx index 2239101..3bbba7a 100644 --- a/pydiso/mkl_solver.pyx +++ b/pydiso/mkl_solver.pyx @@ -530,7 +530,7 @@ cdef class MKLPardisoSolver: if err!=0: raise PardisoError("Solve step error, "+_err_messages[err]) - cdef int _run_pardiso(self, int_t phase, void* b=NULL, void* x=NULL, int_t nrhs=0): + cdef int _run_pardiso(self, int_t phase, void* b=NULL, void* x=NULL, int_t nrhs=0) nogil: cdef int_t error=0 cdef long_t error64=0, phase64=phase, nrhs64=nrhs From 86dc5752b562de5a23333149d21388872052b558 Mon Sep 17 00:00:00 2001 From: Matthew Rocklin Date: Mon, 13 Nov 2023 19:30:05 -0600 Subject: [PATCH 2/3] add nogil annotation to function headers --- pydiso/mkl_solver.pyx | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/pydiso/mkl_solver.pyx b/pydiso/mkl_solver.pyx index 3bbba7a..c4bb758 100644 --- a/pydiso/mkl_solver.pyx +++ b/pydiso/mkl_solver.pyx @@ -42,12 +42,12 @@ cdef extern from 'mkl.h': void pardiso(_MKL_DSS_HANDLE_t, const int*, const int*, const int*, const int *, const int *, const void *, const int *, const int *, int *, const int *, int *, - const int *, void *, void *, int *) + const int *, void *, void *, int *) nogil void pardiso_64(_MKL_DSS_HANDLE_t, const long_t *, const long_t *, const long_t *, const long_t *, const long_t *, const void *, const long_t *, const long_t *, long_t *, const long_t *, long_t *, - const long_t *, void *, void *, long_t *) + const long_t *, void *, void *, long_t *) nogil #call pardiso (pt, maxfct, mnum, mtype, phase, n, a, ia, ja, perm, nrhs, iparm, msglvl, b, x, error) From 35d4d5de04a9d15664029cb0c85aa357e3d7c39c Mon Sep 17 00:00:00 2001 From: Joseph Capriotti Date: Tue, 14 Nov 2023 13:45:45 -0800 Subject: [PATCH 3/3] Add locks to guard against non-threadsafe solver behavior --- pydiso/mkl_solver.pyx | 19 ++++++++++++++++++- tests/test_pydiso.py | 23 +++++++++++++++++++++++ 2 files changed, 41 insertions(+), 1 deletion(-) diff --git a/pydiso/mkl_solver.pyx b/pydiso/mkl_solver.pyx index b1dc21e..e1d8c7f 100644 --- a/pydiso/mkl_solver.pyx +++ b/pydiso/mkl_solver.pyx @@ -2,6 +2,13 @@ #cython: linetrace=True cimport numpy as np from cython cimport numeric +from cpython.pythread cimport ( + PyThread_type_lock, + PyThread_allocate_lock, + PyThread_acquire_lock, + PyThread_release_lock, + PyThread_free_lock +) import warnings import numpy as np @@ -184,7 +191,7 @@ cdef class MKLPardisoSolver: cdef int_t _factored cdef size_t shape[2] cdef int_t _initialized - + cdef PyThread_type_lock lock cdef void * a cdef object _data_type @@ -253,6 +260,9 @@ cdef class MKLPardisoSolver: raise ValueError("Matrix is not square") self.shape = n_row, n_col + # allocate the lock + self.lock = PyThread_allocate_lock() + self._data_type = A.dtype if matrix_type is None: if np.issubdtype(self._data_type, np.complexfloating): @@ -496,6 +506,7 @@ cdef class MKLPardisoSolver: cdef long_t phase64=-1, nrhs64=0, error64=0 if self._initialized: + PyThread_acquire_lock(self.lock, 1) if self._is_32: pardiso( self.handle, &self._par.maxfct, &self._par.mnum, &self._par.mtype, @@ -508,9 +519,12 @@ cdef class MKLPardisoSolver: &phase64, &self._par64.n, self.a, NULL, NULL, NULL, &nrhs64, self._par64.iparm, &self._par64.msglvl, NULL, NULL, &error64 ) + PyThread_release_lock(self.lock) err = error or error64 if err!=0: raise PardisoError("Memmory release error "+_err_messages[err]) + #dealloc lock + PyThread_free_lock(self.lock) cdef _analyze(self): #phase = 11 @@ -540,13 +554,16 @@ cdef class MKLPardisoSolver: cdef int_t error=0 cdef long_t error64=0, phase64=phase, nrhs64=nrhs + PyThread_acquire_lock(self.lock, 1) if self._is_32: pardiso(self.handle, &self._par.maxfct, &self._par.mnum, &self._par.mtype, &phase, &self._par.n, self.a, &self._par.ia[0], &self._par.ja[0], &self._par.perm[0], &nrhs, self._par.iparm, &self._par.msglvl, b, x, &error) + PyThread_release_lock(self.lock) return error else: pardiso_64(self.handle, &self._par64.maxfct, &self._par64.mnum, &self._par64.mtype, &phase64, &self._par64.n, self.a, &self._par64.ia[0], &self._par64.ja[0], &self._par64.perm[0], &nrhs64, self._par64.iparm, &self._par64.msglvl, b, x, &error64) + PyThread_release_lock(self.lock) return error64 diff --git a/tests/test_pydiso.py b/tests/test_pydiso.py index 4aa6d0c..d5c715a 100644 --- a/tests/test_pydiso.py +++ b/tests/test_pydiso.py @@ -8,6 +8,7 @@ set_mkl_threads, set_mkl_pardiso_threads, ) +from concurrent.futures import ThreadPoolExecutor import pytest import sys @@ -147,3 +148,25 @@ def test_rhs_size_error(): solver.solve(b_bad) with pytest.raises(ValueError): solver.solve(b, x_bad) + +def test_threading(): + """ + Here we test that calling the solver is safe from multiple threads. + There isn't actually any speedup because it acquires a lock on each call + to pardiso internally (because those calls are not thread safe). + """ + n = 200 + n_rhs = 75 + A = sp.diags([-1, 2, -1], (-1, 0, 1), shape=(n, n), format='csr') + Ainv = Solver(A) + + x_true = np.random.rand(n, n_rhs) + rhs = A @ x_true + + with ThreadPoolExecutor() as pool: + x_sol = np.stack( + list(pool.map(lambda i: Ainv.solve(rhs[:, i]), range(n_rhs))), + axis=1 + ) + + np.testing.assert_allclose(x_true, x_sol)