PyBind11 Sparse Solver - python-3.x

Question

I try to build a fast sparse solver with Eigen and OpenMP for Python. For the interface between this solver and Python I use the PyBind11 package. Basically, the solver works fine, but unfortunately it only runs on one core and I cannot figure out how to use all cores of my cpu.
Although the OpenMP test functions uses the entire cpu.
Here is the C++ code:
#include <iostream>
#include <cmath>
#include <omp.h>
#include <unistd.h>
#include <pybind11/pybind11.h>
#include <pybind11/numpy.h>
namespace py = pybind11;
#include <Eigen/Sparse>
#include <Eigen/IterativeLinearSolvers>
void openmp_test()
{
int N;
N = 8;
// omp_set_num_threads(4);
#pragma omp parallel for
for (int i=0; i<N; i=i+1)
{
sleep(10);
}
}
void eigen_test(int N)
{
Eigen::SparseMatrix<double> A(N, N);
Eigen::SparseMatrix<double> b(N, 1);
Eigen::SparseMatrix<double> x(N, 1);
Eigen::BiCGSTAB<Eigen::SparseMatrix<double>> solver;
A.reserve(5*N);
b.reserve(N);
x.reserve(N);
for(int i=0; i<N; i++)
{
b.insert(i, 0) = 20.0;
for(int j=(i-2); j<=(i+2); j++)
{
if(j == i)
{
A.insert(j, i) = 10.0;
}
else if((j >= 0) && (j < N))
{
A.insert(j, i) = 5.0;
}
}
}
solver.compute(A);
x = solver.solve(b);
}
PYBIND11_MODULE(mytest, m)
{
m.def("openmp_test", &openmp_test, py::call_guard<py::gil_scoped_release>());
m.def("eigen_test", &eigen_test, py::call_guard<py::gil_scoped_release>());
}
Here is the Compiler code:
g++ \
-O3 \
-Wall \
-shared \
-std=c++14 \
-fopenmp \
-fPIC \
-I /usr/local/lib/python3.10/dist-packages/pybind11/include \
-I /usr/include/python3.10 \
-I /workspaces/pybind11/external/eigen \
mytest.cpp \
-o mytest.so
And finally here the Python code:
from time import time
import mytest
N = 10 * 10**3
start = time()
mytest.openmp_test()
print("runtime: {:.3f} s".format(time() - start))
start = time()
mytest.eigen_test(N)
print("runtime: {:.3e} s".format(time() - start))
Has anyone an idea how to fix this problem?
Thanks a lot.

Answer 1

See documentation:
Currently, the following algorithms can make use of multi-threading:
general dense matrix - matrix products
PartialPivLU
row-major-sparse * dense vector/matrix products
ConjugateGradient with Lower|Upper as the UpLo template parameter.
BiCGSTAB with a row-major sparse matrix format.
LeastSquaresConjugateGradient
Only the row-major sparse matrix format is supported for multi-threading. Therefore, you will need to have:
Eigen::SparseMatrix<double, Eigen::RowMajor> A(N, N);
Eigen::BiCGSTAB<Eigen::SparseMatrix<double, Eigen::RowMajor>> solver;
I have noticed that inserting elements to the matrix takes longer than calculating the solution. Since the prior is sequential, it may not be feasible to observe the parallelism. You can increase the maximum number of solver iterations (i.e solver.setMaxIterations(1e9)) to force solution to take longer (thus easier to observe CPU occupation). You can also print checkpoints to observe which part of your code is executing at that moment. For example:
void eigen_test(int N)
{
Eigen::SparseMatrix<double, Eigen::RowMajor> A(N, N);
Eigen::SparseMatrix<double> b(N, 1);
Eigen::SparseMatrix<double> x(N, 1);
Eigen::BiCGSTAB<Eigen::SparseMatrix<double, Eigen::RowMajor>> solver;
solver.setMaxIterations(1e9);
A.reserve(5*N);
b.reserve(N);
x.reserve(N);
std::cout << "checkpoint: insert elems..." << std::endl;
for(int i=0; i<N; i++)
{
b.insert(i, 0) = 20.0;
for(int j=(i-2); j<=(i+2); j++)
{
if(j == i)
{
A.insert(j, i) = 10.0;
}
else if((j >= 0) && (j < N))
{
A.insert(j, i) = 5.0;
}
}
}
std::cout << "checkpoint: find solution..." << std::endl;
solver.compute(A);
x = solver.solve(b);
std::cout << "checkpoint: done!" << std::endl;
}