add sparse matrix interface

Author: majosm

arraycontext/context.py

@@ -114,6 +114,7 @@
 """
 
 
+import dataclasses
 from abc import ABC, abstractmethod
 from collections.abc import Callable, Hashable, Mapping
 from typing import (
@@ -138,7 +139,7 @@
     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 (
@@ -156,6 +157,22 @@
 P = ParamSpec("P")
 
 
+@dataclasses.dataclass(frozen=True, eq=False, repr=False)
+class _CSRMatrix:
+    # FIXME: Type for shape?
+    shape: Any
+    elem_values: Array
+    elem_col_indices: Array
+    row_starts: Array
+    tags: frozenset[Tag] = dataclasses.field(kw_only=True)
+    axes: tuple[ToTagSetConvertible, ...] = dataclasses.field(kw_only=True)
+    _matmul_func: Callable[[_CSRMatrix, Array], Array] = \
+        dataclasses.field(kw_only=True)
+
+    def __matmul__(self, other: Array) -> Array:
+        return self._matmul_func(self, other)
+
+
 # {{{ ArrayContext
 
 class ArrayContext(ABC):
@@ -172,6 +189,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
@@ -424,6 +443,177 @@ def einsum(self,
         )["out"]
         return self.tag(tagged, out_ary)
 
+    # FIXME: Not sure what type annotations to use for shape and result
+    def make_csr_matrix(
+            self,
+            shape,
+            elem_values: Array,
+            elem_col_indices: Array,
+            row_starts: Array,
+            *,
+            tags: ToTagSetConvertible = frozenset(),
+            axes: tuple[ToTagSetConvertible, ...] | None = None) -> Any:
+        """Return a context-dependent object that represents a sparse matrix in
+        compressed sparse row (CSR) format. Result is suitable for passing to
+        :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 `nrows`.
+        """
+        if axes is None:
+            axes = (frozenset(), frozenset())
+
+        return _CSRMatrix(
+            shape, elem_values, elem_col_indices, row_starts,
+            tags=tags, axes=axes,
+            _matmul_func=lambda x1, x2: self.sparse_matmul(x1, x2))
+
+    @memoize_method
+    def _get_csr_matmul_prg(self, out_ndim: int) -> loopy.TranslationUnit:
+        import numpy as np
+
+        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=(),
+                address_space=lp.AddressSpace.GLOBAL,
+                # FIXME: Need to do anything with tags?
+                ),
+            "iel_ubound": lp.TemporaryVariable(
+                "iel_ubound",
+                shape=(),
+                address_space=lp.AddressSpace.GLOBAL,
+                # FIXME: Need to do anything with tags?
+                )}
+
+        from loopy.kernel.instruction import make_assignment
+        from pymbolic import var
+        # FIXME: Need tags for any of these?
+        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),),
+                    # order="C",
+                    is_input=True),
+                lp.GlobalArg(
+                    "out",
+                    shape=(
+                        var("nrows"),
+                        *(
+                            var(shape_comp_name)
+                            for shape_comp_name in out_extra_shape_comp_names),),
+                    # order="C",
+                    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,
+            # FIXME: Need to do anything with tags?
+            )
+
+        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_values,array,out": np.float64,
+                "elem_col_indices,row_starts": idx_dtype})
+
+    # FIXME: Not sure what type annotation to use for x1
+    def sparse_matmul(self, x1, x2: Array) -> Array:
+        """Multiply a sparse matrix by an array.
+
+        :arg x1: the sparse matrix.
+        :arg x2: the array.
+        """
+        if isinstance(x1, _CSRMatrix):
+            prg = self._get_csr_matmul_prg(len(x2.shape))
+            out_ary = self.call_loopy(
+                prg, elem_values=x1.elem_values,
+                elem_col_indices=x1.elem_col_indices,
+                row_starts=x1.row_starts, array=x2)["out"]
+            # FIXME
+            # return self.tag(tagged, out_ary)
+            return out_ary
+        else:
+            raise TypeError(f"unrecognized matrix type '{type(x1).__name__}'.")
+
     @abstractmethod
     def clone(self) -> Self:
         """If possible, return a version of *self* that is semantically

arraycontext/impl/pytato/__init__.py

@@ -833,6 +833,7 @@ def transform_dag(self, dag: pytato.AbstractResultWithNamedArrays
         dag = pt.transform.materialize_with_mpms(dag)
         return dag
 
+    @override
     def einsum(self, spec, *args, arg_names=None, tagged=()):
         import pytato as pt
 
@@ -876,6 +877,32 @@ def preprocess_arg(name, arg):
             for name, arg in zip(arg_names, args, strict=True)
             ]).tagged(_preprocess_array_tags(tagged))
 
+    # FIXME: Not sure what type annotations to use for shape and result
+    @override
+    def make_csr_matrix(
+            self,
+            shape,
+            elem_values: Array,
+            elem_col_indices: Array,
+            row_starts: Array,
+            *,
+            tags: ToTagSetConvertible = _EMPTY_TAG_SET,
+            axes: tuple[ToTagSetConvertible, ...] | None = None) -> Any:
+        import pytato as pt
+        assert isinstance(elem_values, pt.Array)
+        assert isinstance(elem_col_indices, pt.Array)
+        assert isinstance(row_starts, pt.Array)
+        return pt.make_csr_matrix(
+            shape, elem_values, elem_col_indices, row_starts,
+            # FIXME: Do I need to call _preprocess_array_tags on axes?
+            tags=_preprocess_array_tags(tags), axes=axes)
+
+    # FIXME: Not sure what type annotation to use for x1
+    @override
+    def sparse_matmul(self, x1, x2: Array) -> Array:
+        import pytato as pt
+        return pt.sparse_matmul(x1, x2)
+
     def clone(self):
         return type(self)(self.queue, self.allocator)
 
@@ -1115,6 +1142,24 @@ def preprocess_arg(name: str | None, arg: Array):
             for name, arg in zip(arg_names, args, strict=True)
             ]).tagged(_preprocess_array_tags(tagged)))
 
+    # FIXME: Not sure what type annotations to use for shape and result
+    @override
+    def make_csr_matrix(
+            self,
+            shape,
+            elem_values: Array,
+            elem_col_indices: Array,
+            row_starts: Array,
+            *,
+            tags: ToTagSetConvertible = frozenset(),
+            axes: tuple[ToTagSetConvertible, ...] | None = None) -> Any:
+        raise NotImplementedError
+
+    # FIXME: Not sure what type annotation to use for x1
+    @override
+    def sparse_matmul(self, x1, x2: Array) -> Array:
+        raise NotImplementedError
+
     @override
     def clone(self):
         return type(self)()