Skip to content

Add sparse matrix interface#349

Open
majosm wants to merge 4 commits intoinducer:mainfrom
majosm:sparse-matrix
Open

Add sparse matrix interface#349
majosm wants to merge 4 commits intoinducer:mainfrom
majosm:sparse-matrix

Conversation

@majosm
Copy link
Collaborator

@majosm majosm commented Jan 27, 2026

Adds a generic interface for operations involving sparse CSR matrices.

cc @lukeolson

@majosm
Copy link
Collaborator Author

majosm commented Jan 27, 2026

@inducer I added a loopy kernel fallback for eager (code), and it's working when the array being multiplied is 1D, but I'm running into some trouble with array strides for higher dimensions. I set up the input array like this:

    u1 = actx.np.where(actx.np.less_equal(x, (a+b)/2), 0*x + 1, 0*x)
    u2 = actx.np.where(actx.np.greater_equal(x, (a+b)/2), 0*x + 1, 0*x)
    # u = actx.freeze_thaw(u1)  # 1D
    u = actx.freeze_thaw(actx.np.stack([u1, u2]).T)  # 2D

The strides are no longer C-like after transposing. Lazy doesn't seem to care about this, but eager doesn't like it; it makes my 1D diffusion solver blow up instantly. If I add order="F" to the GlobalArg(...) calls for "array" and "out", then it works fine. Is there a way I can fix this in general?

Edit: I guess the strides of u are only non-C-like in the eager case.

@majosm
Copy link
Collaborator Author

majosm commented Jan 29, 2026

FWIW, using a dense matrix + actx.einsum instead shows the same behavior.

@majosm
Copy link
Collaborator Author

majosm commented Feb 2, 2026

Upon closer inspection, it looks like the kernels may actually be OK. The problem occurs later in the test driver, when I do this:

        u = actx.freeze_thaw(u + dt*compiled_rhs(t, u))

(here u has order='F' and dt*compiled_rhs(t, u) has order='C' in the eager case). They don't add together correctly in eager, which seems to be due to inducer/pyopencl#681. So I guess the solution is just "don't mix orders".

@majosm
Copy link
Collaborator Author

majosm commented Feb 5, 2026

I added some limited support for applying sparse matrices to array containers. Currently it's restricted to the cases of mat @ container. I thought about also allowing matrices to live inside containers (so one could do container_of_mat @ array or container_of_mat @ container), but decided against it for now because it has some complications:

  1. Matrices would need to be arrays. I tried doing this in pytato, and it's not too horrible; I'm not sure if it would work for numpy though (since scipy.sparse.csr_matrix is not a numpy array).
  2. @ operations might be ambiguous in some cases (e.g., object arrays define @).

@inducer I think this is ready for a first look. I still have some FIXMEs scattered throughout the changes for things I could use some input on.

@majosm majosm marked this pull request as ready for review February 5, 2026 22:00
@majosm majosm requested a review from inducer February 5, 2026 22:00
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant