Higher-order functions hang

[MarcelFerrari]

Hi all,

I am encountering a strange behaviour when calling a Nanobind function that accepts a stateful function argument from python.

I have the following function that is a member of an iterative linear system solver class. The solver accepts a MatrixReplacement object which wraps a std::function defining the matrix vector operation implicitly.

 .def("solve", [](Solver& solver, const MatrixOp& op, int rows, int cols,
                       const Eigen::Ref<const Vector>& rhs) -> Vector {

            MatrixReplacement op_wrapper(op, rows, cols);
            solver.compute(op_wrapper); 
            return solver.solve(rhs); 

        }, "f"_a, "rows"_a, "cols"_a, "rhs"_a, "Solve the linear system Ax = b.")

where:

MatrixOp = std::function<void(const Eigen::Ref<const Eigen::Vector<Scalar>>&, 
                                     Eigen::Ref<Eigen::Vector<Scalar>>)>

This allows me to do something like this from Python:

@nb.njit(cache=True, parallel=True)
def poisson(u, u_new, nx, ny, dx, dy):
    # Reshape u and u_new to 2D arrays
    u = u.reshape((nx, ny))
    u_new = u_new.reshape((nx, ny))

    # Apply the discrete Poisson operator
    for i in nb.prange(1, nx - 1):
        for j in nb.prange(1, ny - 1):
            u_new[i, j] = (u[i + 1, j] - 2*u[i, j] + u[i - 1, j])/dx**2 + \
                          (u[i, j + 1] - 2*u[i, j] + u[i, j - 1])/dy**2

    # Boundary values remain unchanged
    u_new[0, :] = u[0, :]
    u_new[-1, :] = u[-1, :]
    u_new[:, 0] = u[:, 0]
    u_new[:, -1] = u[:, -1]

def closure(u, u_new):
    poisson(u, u_new, nx, ny, dx, dy)

solver = ph.BiCGSTAB()
u = solver.solve(closure, rows, cols, rhs)

This works absolutely great in a matrix-free context and has good performance.

The same Nanobind module also exposes a SparseMatrix object that I can assemble explicitly.

However, if I define a closure that captures in any way one of these SparseMatrix objects, then the C++ bindings hang and start consuming a large amount of memory. No explicit errors are thrown.

A = SparseMatrix(*assemble(...))

def closure(u, u_new):
    u_new[:] = A*u

u = solver.solve(closure, rows, cols, rhs)

This only happens for objects of type SparseMatrix. Everything else is fine.

I feel like I am doing something fundamentally wrong, but I cannot seem to figure out what exactly.

Any inputs are much appreciated!

Thanks in advance!

[wjakob]

Not sure. It could be a GIL issue. Did you try to break in a debugger to see where the code is spending its time? In general, this issue has too many moving parts to help debug at the nanobind level. You will have to reduce it more if you'd like to open an issue ticket.