Implement CupyArrayContext

.github/workflows/ci.yml

@@ -41,6 +41,9 @@ jobs:
                 curl -L -O https://tiker.net/ci-support-v0
                 . ./ci-support-v0
 
+                CONDA_ENVIRONMENT=.test-conda-env-py3.yml
+                echo "- cupy" >> "$CONDA_ENVIRONMENT"
+
                 build_py_project_in_conda_env
                 cipip install basedpyright
 

.gitlab-ci.yml

@@ -33,6 +33,25 @@ Python 3 Nvidia Titan V:
     reports:
       junit: test/pytest.xml
 
+Python 3 CuPy Nvidia Titan V:
+  script: |
+    curl -L -O https://tiker.net/ci-support-v0
+    . ./ci-support-v0
+    CONDA_ENVIRONMENT=.test-conda-env-py3.yml
+    echo "- cupy" >> "$CONDA_ENVIRONMENT"
+    export PYOPENCL_TEST=port:cpu
+    build_py_project_in_conda_env
+    test_py_project
+
+  tags:
+  - python3
+  - nvidia-titan-v
+  except:
+  - tags
+  artifacts:
+    reports:
+      junit: test/pytest.xml
+
 Python 3 POCL Nvidia Titan V:
   script: |
     curl -L -O https://tiker.net/ci-support-v0

README.rst

@@ -15,7 +15,8 @@ GPU arrays? Deferred-evaluation arrays? Just plain ``numpy`` arrays? You'd like
 code to work with all of them? No problem! Comes with pre-made array context
 implementations for:
 
-- numpy
+- `numpy <https://numpy.org>`__
+- `cupy <https://cupy.dev>`__
 - `PyOpenCL <https://documen.tician.de/pyopencl/array.html>`__
 - `JAX <https://jax.readthedocs.io/en/latest/>`__
 - `Pytato <https://documen.tician.de/pytato>`__ (for lazy/deferred evaluation)
@@ -24,7 +25,7 @@ implementations for:
 - Profiling
 
 ``arraycontext`` started life as an array abstraction for use with the
-`meshmode <https://documen.tician.de/meshmode/>`__ unstrucuted discretization
+`meshmode <https://documen.tician.de/meshmode/>`__ unstructured discretization
 package.
 
 Distributed under the MIT license.

arraycontext/__init__.py

@@ -88,6 +88,7 @@
     ScalarLike,
     tag_axes,
 )
+from .impl.cupy import CupyArrayContext
 from .impl.jax import EagerJAXArrayContext
 from .impl.numpy import NumpyArrayContext
 from .impl.pyopencl import PyOpenCLArrayContext
@@ -118,6 +119,7 @@
     "ArrayT",
     "BcastUntilActxArray",
     "CommonSubexpressionTag",
+    "CupyArrayContext",
     "EagerJAXArrayContext",
     "ElementwiseMapKernelTag",
     "NotAnArrayContainerError",

arraycontext/container/__init__.py

@@ -94,6 +94,8 @@
 
 
 if TYPE_CHECKING:
+    from typing import Any
+
     from pymbolic.geometric_algebra import MultiVector
 
     from arraycontext import ArrayOrContainer
@@ -283,18 +285,18 @@
 
 @serialize_container.register(np.ndarray)
 def _serialize_ndarray_container(
-        ary: numpy.ndarray) -> SerializedContainer:
+        ary: numpy.ndarray[Any, Any]) -> SerializedContainer:
     if ary.dtype.char != "O":
         raise NotAnArrayContainerError(
                 f"cannot serialize '{type(ary).__name__}' with dtype '{ary.dtype}'")
 
     # special-cased for speed
     if ary.ndim == 1:
         return [(i, ary[i]) for i in range(ary.shape[0])]
     elif ary.ndim == 2:
         return [((i, j), ary[i, j])
                 for i in range(ary.shape[0])
                 for j in range(ary.shape[1])
                 ]
     else:
         return list(np.ndenumerate(ary))
@@ -303,13 +305,13 @@
 @deserialize_container.register(np.ndarray)
 # https://github.com/python/mypy/issues/13040
 def _deserialize_ndarray_container(  # type: ignore[misc]
-        template: numpy.ndarray,
-        serialized: SerializedContainer) -> numpy.ndarray:
+        template: numpy.ndarray[Any, Any],
+        serialized: SerializedContainer) -> numpy.ndarray[Any, Any]:
     # disallow subclasses
     assert type(template) is np.ndarray
     assert template.dtype.char == "O"
 
     result = type(template)(template.shape, dtype=object)
     for i, subary in serialized:
         # FIXME: numpy annotations don't seem to handle object arrays very well
         result[i] = subary  # type: ignore[call-overload]

arraycontext/container/traversal.py

@@ -917,7 +917,7 @@ def _flat_size(subary: ArrayOrContainer) -> Array | Integer:
 # {{{ numpy conversion
 
 def from_numpy(
-        ary: np.ndarray | ScalarLike,
+        ary: np.ndarray[Any, Any] | ScalarLike,
         actx: ArrayContext) -> ArrayOrContainerOrScalar:
     """Convert all :mod:`numpy` arrays in the :class:`~arraycontext.ArrayContainer`
     to the base array type of :class:`~arraycontext.ArrayContext`.

arraycontext/context.py

@@ -282,7 +282,7 @@ def __rtruediv__(self, other: Self | ScalarLike) -> Array: ...
 ContainerOrScalarT = TypeVar("ContainerOrScalarT", bound="ArrayContainer | ScalarLike")
 
 
-NumpyOrContainerOrScalar = Union[np.ndarray, "ArrayContainer", ScalarLike]
+NumpyOrContainerOrScalar = Union[np.ndarray[Any, Any], "ArrayContainer", ScalarLike]
 
 # }}}
 
@@ -358,7 +358,7 @@ def zeros(self,
         return self.np.zeros(shape, dtype)
 
     @overload
-    def from_numpy(self, array: np.ndarray) -> Array:
+    def from_numpy(self, array: np.ndarray[Any, Any]) -> Array:
         ...
 
     @overload
@@ -379,7 +379,7 @@ def from_numpy(self,
         """
 
     @overload
-    def to_numpy(self, array: Array) -> np.ndarray:
+    def to_numpy(self, array: Array) -> np.ndarray[Any, Any]:
         ...
 
     @overload

arraycontext/impl/cupy/__init__.py

@@ -0,0 +1,155 @@
+"""
+.. currentmodule:: arraycontext
+
+A :mod:`cupy`-based array context.
+
+.. autoclass:: CupyArrayContext
+"""
+
+from __future__ import annotations
+
+
+__copyright__ = """
+Copyright (C) 2024 University of Illinois Board of Trustees
+"""
+
+__license__ = """
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in
+all copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
+THE SOFTWARE.
+"""
+
+from typing import Any, overload
+
+import numpy as np
+
+import loopy as lp
+from pytools.tag import ToTagSetConvertible
+
+from arraycontext.container.traversal import rec_map_array_container, with_array_context
+from arraycontext.context import (
+    Array,
+    ArrayContext,
+    ArrayOrContainerOrScalar,
+    ArrayOrContainerOrScalarT,
+    ContainerOrScalarT,
+    NumpyOrContainerOrScalar,
+)
+
+
+class CupyNonObjectArrayMetaclass(type):
+    def __instancecheck__(cls, instance: Any) -> bool:
+        import cupy as cp  # type: ignore[import-untyped]
+        return isinstance(instance, cp.ndarray) and instance.dtype != object
+
+
+class CupyNonObjectArray(metaclass=CupyNonObjectArrayMetaclass):
+    pass
+
+
+class CupyArrayContext(ArrayContext):
+    """An :class:`ArrayContext` that uses :class:`cupy.ndarray` to represent arrays."""
+
+    array_types = (CupyNonObjectArray,)
+
+    def _get_fake_numpy_namespace(self):
+        from .fake_numpy import CupyFakeNumpyNamespace
+        return CupyFakeNumpyNamespace(self)
+
+    # {{{ ArrayContext interface
+
+    def clone(self):
+        return type(self)()
+
+    @overload
+    def from_numpy(self, array: np.ndarray[Any, Any]) -> Array:
+        ...
+
+    @overload
+    def from_numpy(self, array: ContainerOrScalarT) -> ContainerOrScalarT:
+        ...
+
+    def from_numpy(self,
+                   array: NumpyOrContainerOrScalar
+                   ) -> ArrayOrContainerOrScalar:
+        import cupy as cp
+        return with_array_context(rec_map_array_container(cp.array, array),
+                                  actx=self)
+
+    @overload
+    def to_numpy(self, array: Array) -> np.ndarray[Any, Any]:
+        ...
+
+    @overload
+    def to_numpy(self, array: ContainerOrScalarT) -> ContainerOrScalarT:
+        ...
+
+    def to_numpy(self,
+                 array: ArrayOrContainerOrScalar
+                 ) -> NumpyOrContainerOrScalar:
+        import cupy as cp
+        return with_array_context(rec_map_array_container(cp.asnumpy, array),
+                                  actx=None)
+
+    def call_loopy(
+                self,
+                t_unit: lp.TranslationUnit, **kwargs: Any
+            ) -> dict[str, Array]:
+        raise NotImplementedError(
+            "Calling loopy on CuPy arrays is not supported. Maybe rewrite"
+            " the loopy kernel as numpy-flavored array operations using"
+            " ArrayContext.np.")
+
+    def freeze(self, array):
+        import cupy as cp
+        # Note that we could use a non-blocking version of cp.asnumpy here, but
+        # it appears to have very little impact on performance.
+        return with_array_context(rec_map_array_container(cp.asnumpy, array), actx=None)
+
+    def thaw(self, array):
+        import cupy as cp
+        return with_array_context(rec_map_array_container(cp.array, array), actx=self)
+
+    # }}}
+
+    def tag(self,
+            tags: ToTagSetConvertible,
+            array: ArrayOrContainerOrScalarT) -> ArrayOrContainerOrScalarT:
+        # Cupy (like numpy) doesn't support tagging
+        return array
+
+    def tag_axis(self,
+                 iaxis: int, tags: ToTagSetConvertible,
+                 array: ArrayOrContainerOrScalarT) -> ArrayOrContainerOrScalarT:
+        # Cupy (like numpy) doesn't support tagging
+        return array
+
+    def einsum(self, spec, *args, arg_names=None, tagged=()):
+        import cupy as cp
+        return cp.einsum(spec, *args)
+
+    @property
+    def permits_inplace_modification(self):
+        return True
+
+    @property
+    def supports_nonscalar_broadcasting(self):
+        return True
+
+    @property
+    def permits_advanced_indexing(self):
+        return True