Add sparse matrix interface (#349)

* add sparse matrix interface

* add tentative numpy actx support

* add tests

* fix some typing

* put actx-specific matrix multiplication implementations inside sparse_matmul instead of using derived matrix classes

* remove FIXMEs about tags

* add optional scipy dependency for numpy sparse

* update basedpyright baseline for loopy stuff

* fix make_csr_matrix docstring

* omit dtype specification for floating point CSR matmul args

Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com>

* mention arraycontext[sparse] in exception message if trying to use numpy array context without scipy

* don't test CSR on numpy array context if scipy isn't installed

* install scipy for pytest CI

* install scipy for pytest CI (gitlab)

* don't use global address space for reduction bound temps

---------

Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com>

Author: majosm

.basedpyright/baseline.json

@@ -954,6 +954,54 @@
                     "endColumn": 29,
                     "lineCount": 1
                 }
+            },
+            {
+                "code": "reportUnknownMemberType",
+                "range": {
+                    "startColumn": 16,
+                    "endColumn": 28,
+                    "lineCount": 1
+                }
+            },
+            {
+                "code": "reportUnknownMemberType",
+                "range": {
+                    "startColumn": 16,
+                    "endColumn": 28,
+                    "lineCount": 1
+                }
+            },
+            {
+                "code": "reportUnknownMemberType",
+                "range": {
+                    "startColumn": 16,
+                    "endColumn": 28,
+                    "lineCount": 1
+                }
+            },
+            {
+                "code": "reportUnknownMemberType",
+                "range": {
+                    "startColumn": 16,
+                    "endColumn": 28,
+                    "lineCount": 1
+                }
+            },
+            {
+                "code": "reportUnknownMemberType",
+                "range": {
+                    "startColumn": 16,
+                    "endColumn": 28,
+                    "lineCount": 1
+                }
+            },
+            {
+                "code": "reportUnknownMemberType",
+                "range": {
+                    "startColumn": 15,
+                    "endColumn": 38,
+                    "lineCount": 1
+                }
             }
         ],
         "./arraycontext/fake_numpy.py": [

.github/workflows/ci.yml

@@ -38,7 +38,7 @@ jobs:
         -   uses: actions/checkout@v6
         -   name: "Main Script"
             run: |
-                EXTRA_INSTALL="pytest types-colorama types-Pygments"
+                EXTRA_INSTALL="pytest types-colorama types-Pygments scipy-stubs"
                 curl -L -O https://tiker.net/ci-support-v0
                 . ./ci-support-v0
 
@@ -54,6 +54,7 @@ jobs:
         -   uses: actions/checkout@v6
         -   name: "Main Script"
             run: |
+                EXTRA_INSTALL="scipy"
                 curl -L -O https://tiker.net/ci-support-v0
                 . ./ci-support-v0
                 build_py_project_in_conda_env
@@ -75,6 +76,7 @@ jobs:
                 export PYOPENCL_TEST=intel
                 source /opt/enable-intel-cl.sh
 
+                EXTRA_INSTALL="scipy"
                 curl -L -O https://tiker.net/ci-support-v0
                 . ./ci-support-v0
                 build_py_project_in_conda_env

.gitlab-ci.yml

@@ -1,7 +1,7 @@
 Python 3 POCL:
   script: |
     export PYOPENCL_TEST=portable:cpu
-    export EXTRA_INSTALL="jax[cpu]"
+    export EXTRA_INSTALL="jax[cpu] scipy"
     export JAX_PLATFORMS=cpu
     curl -L -O https://gitlab.tiker.net/inducer/ci-support/raw/main/build-and-test-py-project.sh
     . ./build-and-test-py-project.sh
@@ -20,6 +20,7 @@ Python 3 Nvidia Titan V:
     curl -L -O https://tiker.net/ci-support-v0
     . ./ci-support-v0
     export PYOPENCL_TEST=nvi:titan
+    export EXTRA_INSTALL="scipy"
     build_py_project_in_venv
     pip install --upgrade "jax[cuda]" -f https://storage.googleapis.com/jax-releases/jax_releases.html
     test_py_project
@@ -38,6 +39,7 @@ Python 3 POCL Nvidia Titan V:
     curl -L -O https://tiker.net/ci-support-v0
     . ./ci-support-v0
     export PYOPENCL_TEST=port:titan
+    export EXTRA_INSTALL="scipy"
     build_py_project_in_venv
     test_py_project
 
@@ -66,6 +68,7 @@ Python 3 POCL Examples:
 Python 3 Conda:
   script: |
     export PYOPENCL_TEST=portable:cpu
+    export EXTRA_INSTALL="scipy"
 
     # Avoid crashes like https://gitlab.tiker.net/inducer/arraycontext/-/jobs/536021
     sed -i 's/jax/jax !=0.4.6/' .test-conda-env-py3.yml

arraycontext/__init__.py

@@ -73,6 +73,8 @@
 from .context import (
     ArrayContext,
     ArrayContextFactory,
+    CSRMatrix,
+    SparseMatrix,
     tag_axes,
 )
 from .impl.jax import EagerJAXArrayContext
@@ -129,6 +131,7 @@
     "ArrayOrScalarT",
     "ArrayT",
     "BcastUntilActxArray",
+    "CSRMatrix",
     "CommonSubexpressionTag",
     "ContainerOrScalarT",
     "EagerJAXArrayContext",
@@ -144,6 +147,7 @@
     "ScalarLike",
     "SerializationKey",
     "SerializedContainer",
+    "SparseMatrix",
     "dataclass_array_container",
     "deserialize_container",
     "flat_size_and_dtype",

arraycontext/context.py

@@ -78,6 +78,9 @@
 
 .. autoclass:: ArrayContext
 
+.. autoclass:: SparseMatrix
+.. autoclass:: CSRMatrix
+
 .. autofunction:: tag_axes
 
 .. class:: P
@@ -114,13 +117,15 @@
 """
 
 
+import dataclasses
 from abc import ABC, abstractmethod
 from collections.abc import Callable, Hashable, Mapping
 from typing import (
     TYPE_CHECKING,
     Any,
     ParamSpec,
     TypeAlias,
+    cast,
     overload,
 )
 from warnings import warn
@@ -129,21 +134,27 @@
 
 from pytools import memoize_method
 
+from arraycontext.container.traversal import (
+    rec_map_container,
+)
+
 
 if TYPE_CHECKING:
     import numpy as np
     from numpy.typing import DTypeLike
 
     import loopy
-    from pytools.tag import ToTagSetConvertible
+    from pytools.tag import Tag, ToTagSetConvertible
 
     from .fake_numpy import BaseFakeNumpyNamespace
     from .typing import (
         Array,
         ArrayContainerT,
         ArrayOrArithContainerOrScalarT,
+        ArrayOrContainer,
         ArrayOrContainerOrScalar,
         ArrayOrContainerOrScalarT,
+        ArrayOrScalar,
         ContainerOrScalarT,
         NumpyOrContainerOrScalar,
         ScalarLike,
@@ -152,6 +163,26 @@
 
 P = ParamSpec("P")
 
+_EMPTY_TAG_SET: frozenset[Tag] = frozenset()
+
+
+@dataclasses.dataclass(frozen=True, eq=False, repr=False)
+class SparseMatrix(ABC):
+    shape: tuple[int, int]
+    tags: ToTagSetConvertible = dataclasses.field(kw_only=True)
+    axes: tuple[ToTagSetConvertible, ...] = dataclasses.field(kw_only=True)
+    _actx: ArrayContext = dataclasses.field(kw_only=True)
+
+    def __matmul__(self, other: ArrayOrContainer) -> ArrayOrContainer:
+        return self._actx.sparse_matmul(self, other)
+
+
+@dataclasses.dataclass(frozen=True, eq=False, repr=False)
+class CSRMatrix(SparseMatrix):
+    elem_values: Array
+    elem_col_indices: Array
+    row_starts: Array
+
 
 # {{{ ArrayContext
 
@@ -169,6 +200,8 @@ class ArrayContext(ABC):
     .. automethod:: to_numpy
     .. automethod:: call_loopy
     .. automethod:: einsum
+    .. automethod:: make_csr_matrix
+    .. automethod:: sparse_matmul
     .. attribute:: np
 
          Provides access to a namespace that serves as a work-alike to
@@ -421,6 +454,166 @@ def einsum(self,
         )["out"]
         return self.tag(tagged, out_ary)
 
+    def make_csr_matrix(
+            self,
+            shape: tuple[int, int],
+            elem_values: Array,
+            elem_col_indices: Array,
+            row_starts: Array,
+            *,
+            tags: ToTagSetConvertible = _EMPTY_TAG_SET,
+            axes: tuple[ToTagSetConvertible, ...] | None = None) -> CSRMatrix:
+        """Return a sparse matrix in compressed sparse row (CSR) format, to be used
+        with :meth:`sparse_matmul`.
+
+        :arg shape: the (two-dimensional) shape of the matrix
+        :arg elem_values: a one-dimensional array containing the values of all of the
+            nonzero entries of the matrix, grouped by row.
+        :arg elem_col_indices: a one-dimensional array containing the column index
+            values corresponding to each entry in *elem_values*.
+        :arg row_starts: a one-dimensional array of length ``nrows+1``, where each entry
+            gives the starting index in *elem_values* and *elem_col_indices* for the
+            given row, with the last entry being equal to ``len(elem_values)``.
+        """
+        if axes is None:
+            axes = (frozenset(), frozenset())
+
+        return CSRMatrix(
+            shape, elem_values, elem_col_indices, row_starts,
+            tags=tags, axes=axes,
+            _actx=self)
+
+    @memoize_method
+    def _get_csr_matmul_prg(self, out_ndim: int) -> loopy.TranslationUnit:
+        import loopy as lp
+
+        out_extra_inames = tuple(f"i{n}" for n in range(1, out_ndim))
+        out_inames = ("irow", *out_extra_inames)
+        out_inames_set = frozenset(out_inames)
+
+        out_extra_shape_comp_names = tuple(f"n{n}" for n in range(1, out_ndim))
+        out_shape_comp_names = ("nrows", *out_extra_shape_comp_names)
+
+        domains: list[str] = []
+        domains.append(
+            "{ [" + ",".join(out_inames) + "] : "
+            + " and ".join(
+                f"0 <= {iname} < {shape_comp_name}"
+                for iname, shape_comp_name in zip(
+                    out_inames, out_shape_comp_names, strict=True))
+            + " }")
+        domains.append(
+            "{ [iel] : iel_lbound <= iel < iel_ubound }")
+
+        temporary_variables: Mapping[str, lp.TemporaryVariable] = {
+            "iel_lbound": lp.TemporaryVariable(
+                "iel_lbound",
+                shape=(),
+                ),
+            "iel_ubound": lp.TemporaryVariable(
+                "iel_ubound",
+                shape=(),
+                )}
+
+        from loopy.kernel.instruction import make_assignment
+        from pymbolic import var
+        instructions: list[lp.Assignment | lp.CallInstruction] = [
+            make_assignment(
+                (var("iel_lbound"),),
+                var("row_starts")[var("irow")],
+                id="insn0",
+                within_inames=out_inames_set),
+            make_assignment(
+                (var("iel_ubound"),),
+                var("row_starts")[var("irow") + 1],
+                id="insn1",
+                within_inames=out_inames_set),
+            make_assignment(
+                (var("out")[tuple(var(iname) for iname in out_inames)],),
+                lp.Reduction(
+                    "sum",
+                    (var("iel"),),
+                    var("elem_values")[var("iel"),]
+                    * var("array")[(
+                        var("elem_col_indices")[var("iel"),],
+                        *(var(iname) for iname in out_extra_inames))]),
+                id="insn2",
+                within_inames=out_inames_set,
+                depends_on=frozenset({"insn0", "insn1"}))]
+
+        from loopy.version import MOST_RECENT_LANGUAGE_VERSION
+
+        from .loopy import _DEFAULT_LOOPY_OPTIONS
+
+        knl = lp.make_kernel(
+            domains=domains,
+            instructions=instructions,
+            temporary_variables=temporary_variables,
+            kernel_data=[
+                lp.ValueArg("nrows", is_input=True),
+                lp.ValueArg("ncols", is_input=True),
+                lp.ValueArg("nels", is_input=True),
+                *(
+                    lp.ValueArg(shape_comp_name, is_input=True)
+                    for shape_comp_name in out_extra_shape_comp_names),
+                lp.GlobalArg("elem_values", shape=(var("nels"),), is_input=True),
+                lp.GlobalArg("elem_col_indices", shape=(var("nels"),), is_input=True),
+                lp.GlobalArg("row_starts", shape=lp.auto, is_input=True),
+                lp.GlobalArg(
+                    "array",
+                    shape=(
+                        var("ncols"),
+                        *(
+                            var(shape_comp_name)
+                            for shape_comp_name in out_extra_shape_comp_names),),
+                    is_input=True),
+                lp.GlobalArg(
+                    "out",
+                    shape=(
+                        var("nrows"),
+                        *(
+                            var(shape_comp_name)
+                            for shape_comp_name in out_extra_shape_comp_names),),
+                    is_input=False),
+                ...],
+            name="csr_matmul_kernel",
+            lang_version=MOST_RECENT_LANGUAGE_VERSION,
+            options=_DEFAULT_LOOPY_OPTIONS,
+            default_order=lp.auto,
+            default_offset=lp.auto,
+            )
+
+        idx_dtype = knl.default_entrypoint.index_dtype
+
+        return lp.add_and_infer_dtypes(
+            knl,
+            {
+                ",".join([
+                    "ncols", "nrows", "nels",
+                    *out_extra_shape_comp_names]): idx_dtype,
+                "elem_col_indices,row_starts": idx_dtype})
+
+    def sparse_matmul(
+            self, x1: SparseMatrix, x2: ArrayOrContainer) -> ArrayOrContainer:
+        """Multiply a sparse matrix by an array.
+
+        :arg x1: the sparse matrix.
+        :arg x2: the array.
+        """
+        if isinstance(x1, CSRMatrix):
+            def _matmul(ary: ArrayOrScalar) -> ArrayOrScalar:
+                assert self.is_array_type(ary)
+                prg = self._get_csr_matmul_prg(len(ary.shape))
+                return self.call_loopy(
+                    prg, elem_values=x1.elem_values,
+                    elem_col_indices=x1.elem_col_indices,
+                    row_starts=x1.row_starts, array=ary)["out"]
+
+            return cast("ArrayOrContainer", rec_map_container(_matmul, x2))
+
+        else:
+            raise TypeError(f"unrecognized sparse matrix type '{type(x1).__name__}'")
+
     @abstractmethod
     def clone(self) -> Self:
         """If possible, return a version of *self* that is semantically