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

Author: majosm

arraycontext/context.py

@@ -175,9 +175,8 @@ class SparseMatrix(ABC):
     axes: tuple[ToTagSetConvertible, ...] = dataclasses.field(kw_only=True)
     _actx: ArrayContext = dataclasses.field(kw_only=True)
 
-    @abstractmethod
     def __matmul__(self, other: ArrayOrContainer) -> ArrayOrContainer:
-        ...
+        return self._actx.sparse_matmul(self, other)
 
 
 @dataclasses.dataclass(frozen=True, eq=False, repr=False)
@@ -186,21 +185,6 @@ class CSRMatrix(SparseMatrix):
     elem_col_indices: Array
     row_starts: Array
 
-    @override
-    def __matmul__(self, other: ArrayOrContainer) -> ArrayOrContainer:
-        def _matmul(ary: ArrayOrScalar) -> ArrayOrScalar:
-            assert self._actx.is_array_type(ary)
-            prg = self._actx._get_csr_matmul_prg(len(ary.shape))
-            out_ary = self._actx.call_loopy(
-                prg, elem_values=self.elem_values,
-                elem_col_indices=self.elem_col_indices,
-                row_starts=self.row_starts, array=ary)["out"]
-            # FIXME
-            # return self.tag(tagged, out_ary)
-            return out_ary
-
-        return cast("ArrayOrContainer", rec_map_container(_matmul, other))
-
 
 # {{{ ArrayContext
 
@@ -627,7 +611,22 @@ def sparse_matmul(
         :arg x1: the sparse matrix.
         :arg x2: the array.
         """
-        return x1 @ x2
+        if isinstance(x1, CSRMatrix):
+            def _matmul(ary: ArrayOrScalar) -> ArrayOrScalar:
+                assert self.is_array_type(ary)
+                prg = self._get_csr_matmul_prg(len(ary.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=ary)["out"]
+                # FIXME
+                # return self.tag(tagged, out_ary)
+                return out_ary
+
+            return cast("ArrayOrContainer", rec_map_container(_matmul, x2))
+
+        else:
+            raise TypeError(f"unrecognized sparse matrix type '{type(x1).__name__}'")
 
     @abstractmethod
     def clone(self) -> Self:

arraycontext/impl/numpy/__init__.py

@@ -33,15 +33,12 @@
 THE SOFTWARE.
 """
 
-from dataclasses import dataclass
-from functools import cached_property
 from typing import TYPE_CHECKING, Any, cast, overload
 
 import numpy as np
 from typing_extensions import override
 
 import loopy as lp
-from pytools.tag import normalize_tags
 
 from arraycontext.container.traversal import (
     rec_map_array_container as rec_map_array_container,
@@ -50,23 +47,23 @@
 )
 from arraycontext.context import (
     ArrayContext,
-    CSRMatrix as _BaseCSRMatrix,
+    CSRMatrix,
+    SparseMatrix,
     UntransformedCodeWarning,
 )
 from arraycontext.typing import (
     Array,
     ArrayOrContainer,
     ArrayOrContainerOrScalar,
     ArrayOrContainerOrScalarT,
+    ArrayOrScalar,
     ContainerOrScalarT,
     NumpyOrContainerOrScalar,
     is_scalar_like,
 )
 
 
 if TYPE_CHECKING:
-    import scipy.sparse
-
     from pymbolic import Scalar
     from pytools.tag import Tag, ToTagSetConvertible
 
@@ -86,29 +83,6 @@ class NumpyNonObjectArray(metaclass=NumpyNonObjectArrayMetaclass):
     pass
 
 
-@dataclass(frozen=True, eq=False, repr=False)
-class CSRMatrix(_BaseCSRMatrix):
-    @cached_property
-    def _np_matrix(self) -> scipy.sparse.csr_matrix:
-        assert isinstance(self.elem_values, np.ndarray)
-        assert isinstance(self.elem_col_indices, np.ndarray)
-        assert isinstance(self.row_starts, np.ndarray)
-        # FIXME: Not sure if the scipy dependency is OK or if it should just use the
-        # call_loopy fallback? Currently getting errors with the loopy version:
-        #     loopy.diagnostic.LoopyError: One of the kernels in the program has
-        #         been preprocessed, cannot modify target now.
-        from scipy.sparse import csr_matrix
-        return csr_matrix(
-            (self.elem_values, self.elem_col_indices, self.row_starts),
-            shape=self.shape)
-
-    @override
-    def __matmul__(self, other: ArrayOrContainer) -> ArrayOrContainer:
-        return cast(
-            "ArrayOrContainer",
-            rec_map_container(lambda ary: self._np_matrix @ ary, other))
-
-
 class NumpyArrayContext(ArrayContext):
     """
     A :class:`ArrayContext` that uses :class:`numpy.ndarray` to represent arrays.
@@ -233,22 +207,30 @@ def einsum(self, spec, *args, arg_names=None, tagged=()):
         return np.einsum(spec, *args, optimize="optimal")
 
     @override
-    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:
-        tags = normalize_tags(tags)
-        if axes is None:
-            axes = (frozenset(), frozenset())
-        return CSRMatrix(
-            shape, elem_values, elem_col_indices, row_starts,
-            tags=tags, axes=axes,
-            _actx=self)
+    def sparse_matmul(
+            self, x1: SparseMatrix, x2: ArrayOrContainer) -> ArrayOrContainer:
+        if isinstance(x1, CSRMatrix):
+            assert isinstance(x1.elem_values, np.ndarray)
+            assert isinstance(x1.elem_col_indices, np.ndarray)
+            assert isinstance(x1.row_starts, np.ndarray)
+
+            # FIXME: Not sure if the scipy dependency is OK or if it should just use
+            # the call_loopy fallback? Currently getting errors with the loopy version:
+            #     loopy.diagnostic.LoopyError: One of the kernels in the program has
+            #         been preprocessed, cannot modify target now.
+            from scipy.sparse import csr_matrix
+            np_matrix = csr_matrix(
+                (x1.elem_values, x1.elem_col_indices, x1.row_starts),
+                shape=x1.shape)
+
+            def _matmul(ary: ArrayOrScalar) -> ArrayOrScalar:
+                assert isinstance(ary, np.ndarray)
+                return np_matrix @ ary
+
+            return cast("ArrayOrContainer", rec_map_container(_matmul, x2))
+
+        else:
+            raise TypeError(f"unrecognized sparse matrix type '{type(x1).__name__}'")
 
     @property
     def permits_inplace_modification(self):

arraycontext/impl/pytato/__init__.py

@@ -54,13 +54,12 @@
 import abc
 import sys
 from dataclasses import dataclass
-from functools import cached_property
 from typing import TYPE_CHECKING, Any, cast
 
 import numpy as np
 from typing_extensions import override
 
-from pytools import memoize_method
+from pytools import memoize_in, memoize_method
 from pytools.tag import Tag, ToTagSetConvertible, normalize_tags
 
 from arraycontext.container.traversal import (
@@ -69,7 +68,7 @@
 )
 from arraycontext.context import (
     ArrayContext,
-    CSRMatrix as _BaseCSRMatrix,
+    CSRMatrix,
     P,
     SparseMatrix,
     UntransformedCodeWarning,
@@ -144,28 +143,6 @@ class _NotOnlyDataWrappers(Exception):  # noqa: N818
     pass
 
 
-@dataclass(frozen=True, eq=False, repr=False)
-class CSRMatrix(_BaseCSRMatrix):
-    @cached_property
-    def _pt_matrix(self) -> pt.CSRMatrix:
-        import pytato as pt
-        assert isinstance(self.elem_values, pt.Array)
-        assert isinstance(self.elem_col_indices, pt.Array)
-        assert isinstance(self.row_starts, pt.Array)
-        return pt.make_csr_matrix(
-            self.shape, self.elem_values, self.elem_col_indices, self.row_starts,
-            tags=_preprocess_array_tags(self.tags), axes=self.axes)
-
-    @override
-    def __matmul__(self, other: ArrayOrContainer) -> ArrayOrContainer:
-        def _matmul(ary: ArrayOrScalar) -> ArrayOrScalar:
-            import pytato as pt
-            assert isinstance(ary, pt.Array)
-            return self._pt_matrix @ ary
-
-        return cast("ArrayOrContainer", rec_map_container(_matmul, other))
-
-
 # {{{ _BasePytatoArrayContext
 
 class _BasePytatoArrayContext(ArrayContext, abc.ABC):
@@ -904,21 +881,30 @@ def preprocess_arg(name, arg):
             ]).tagged(_preprocess_array_tags(tagged))
 
     @override
-    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:
-        if axes is None:
-            axes = (frozenset(), frozenset())
-        return CSRMatrix(
-            shape, elem_values, elem_col_indices, row_starts,
-            tags=tags, axes=axes,
-            _actx=self)
+    def sparse_matmul(
+            self, x1: SparseMatrix, x2: ArrayOrContainer) -> ArrayOrContainer:
+        import pytato as pt
+
+        if isinstance(x1, CSRMatrix):
+            @memoize_in(x1, "pt_matrix")
+            def _get_pt_matrix() -> pt.CSRMatrix:
+                assert isinstance(x1.elem_values, pt.Array)
+                assert isinstance(x1.elem_col_indices, pt.Array)
+                assert isinstance(x1.row_starts, pt.Array)
+                return pt.make_csr_matrix(
+                    x1.shape, x1.elem_values, x1.elem_col_indices, x1.row_starts,
+                    tags=_preprocess_array_tags(x1.tags), axes=x1.axes)
+
+            pt_matrix: pt.CSRMatrix = _get_pt_matrix()
+
+            def _matmul(ary: ArrayOrScalar) -> ArrayOrScalar:
+                assert isinstance(ary, pt.Array)
+                return pt_matrix @ ary
+
+            return cast("ArrayOrContainer", rec_map_container(_matmul, x2))
+
+        else:
+            raise TypeError(f"unrecognized sparse matrix type '{type(x1).__name__}'")
 
     def clone(self):
         return type(self)(self.queue, self.allocator)