diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 28923da7..fe75e8e2 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -40,6 +40,10 @@ jobs: run: | USE_CONDA_BUILD=1 curl -L -O https://gitlab.tiker.net/inducer/ci-support/raw/master/prepare-and-run-pylint.sh + + CONDA_ENVIRONMENT=.test-conda-env-py3.yml + echo "- cupy" >> "$CONDA_ENVIRONMENT" + . ./prepare-and-run-pylint.sh "$(basename $GITHUB_REPOSITORY)" examples/*.py test/test_*.py mypy: @@ -52,6 +56,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 python -m pip install mypy pytest ./run-mypy.sh diff --git a/.gitlab-ci.yml b/.gitlab-ci.yml index 313b7633..e2fcd3cf 100644 --- a/.gitlab-ci.yml +++ b/.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 diff --git a/README.rst b/README.rst index a704c122..1206093e 100644 --- a/README.rst +++ b/README.rst @@ -16,6 +16,7 @@ code to work with all of them? No problem! Comes with pre-made array context implementations for: - numpy +- cupy - `PyOpenCL `__ - `JAX `__ - `Pytato `__ (for lazy/deferred evaluation) diff --git a/arraycontext/__init__.py b/arraycontext/__init__.py index 74adae96..2e841b7b 100644 --- a/arraycontext/__init__.py +++ b/arraycontext/__init__.py @@ -87,6 +87,7 @@ ScalarLike, tag_axes, ) +from .impl.cupy import CupyArrayContext from .impl.jax import EagerJAXArrayContext from .impl.numpy import NumpyArrayContext from .impl.pyopencl import PyOpenCLArrayContext @@ -116,6 +117,7 @@ "ArrayOrContainerT", "ArrayT", "CommonSubexpressionTag", + "CupyArrayContext", "EagerJAXArrayContext", "ElementwiseMapKernelTag", "NotAnArrayContainerError", diff --git a/arraycontext/impl/cupy/__init__.py b/arraycontext/impl/cupy/__init__.py new file mode 100644 index 00000000..c8c8e91a --- /dev/null +++ b/arraycontext/impl/cupy/__init__.py @@ -0,0 +1,207 @@ +""" +.. 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, + UntransformedCodeWarning, +) + + +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. + + .. automethod:: __init__ + """ + + _loopy_transform_cache: dict[lp.TranslationUnit, lp.ExecutorBase] + + def __init__(self, device: int | None = None) -> None: + super().__init__() + self._loopy_transform_cache = {} + + if device is not None: + import cupy as cp + cp.cuda.runtime.setDevice(device) + + 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) -> Array: + ... + + @overload + def from_numpy(self, array: ContainerOrScalarT) -> ContainerOrScalarT: + ... + + def from_numpy(self, + array: NumpyOrContainerOrScalar + ) -> ArrayOrContainerOrScalar: + import cupy as cp + + def _from_numpy(ary): + return cp.array(ary) + + return with_array_context(rec_map_array_container(_from_numpy, array), + actx=self) + + @overload + def to_numpy(self, array: Array) -> np.ndarray: + ... + + @overload + def to_numpy(self, array: ContainerOrScalarT) -> ContainerOrScalarT: + ... + + def to_numpy(self, + array: ArrayOrContainerOrScalar + ) -> NumpyOrContainerOrScalar: + import cupy as cp + + def _to_numpy(ary): + return cp.asnumpy(ary) + + return with_array_context(rec_map_array_container(_to_numpy, array), + actx=None) + + def call_loopy( + self, + t_unit: lp.TranslationUnit, **kwargs: Any + ) -> dict[str, Array]: + t_unit = t_unit.copy(target=lp.ExecutableCTarget()) + try: + executor = self._loopy_transform_cache[t_unit] + except KeyError: + executor = self.transform_loopy_program(t_unit).executor() + self._loopy_transform_cache[t_unit] = executor + + _, result = executor(**kwargs) + + return result + + def freeze(self, array): + import cupy as cp + + def _freeze(ary): + return cp.asnumpy(ary) + + return with_array_context(rec_map_array_container(_freeze, array), actx=None) + + def thaw(self, array): + import cupy as cp + + def _thaw(ary): + return cp.array(ary) + + return with_array_context(rec_map_array_container(_thaw, array), actx=self) + + # }}} + + def transform_loopy_program(self, t_unit): + from warnings import warn + warn("Using the base " + f"{type(self).__name__}.transform_loopy_program " + "to transform a translation unit. " + "This is a no-op and will result in unoptimized C code for" + "the requested optimization, all in a single statement." + "This will work, but is unlikely to be performant." + f"Instead, subclass {type(self).__name__} and implement " + "the specific transform logic required to transform the program " + "for your package or application. Check higher-level packages " + "(e.g. meshmode), which may already have subclasses you may want " + "to build on.", + UntransformedCodeWarning, stacklevel=2) + + return t_unit + + 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 diff --git a/arraycontext/impl/cupy/fake_numpy.py b/arraycontext/impl/cupy/fake_numpy.py new file mode 100644 index 00000000..e45ca8a8 --- /dev/null +++ b/arraycontext/impl/cupy/fake_numpy.py @@ -0,0 +1,201 @@ +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 functools import partial, reduce + +from arraycontext.container import NotAnArrayContainerError, serialize_container +from arraycontext.container.traversal import ( + rec_map_array_container, + rec_map_reduce_array_container, + rec_multimap_array_container, + rec_multimap_reduce_array_container, +) +from arraycontext.context import Array, ArrayOrContainer +from arraycontext.fake_numpy import ( + BaseFakeNumpyLinalgNamespace, + BaseFakeNumpyNamespace, +) + + +class CupyFakeNumpyLinalgNamespace(BaseFakeNumpyLinalgNamespace): + # Everything is implemented in the base class for now. + pass + + +_NUMPY_UFUNCS = frozenset({"concatenate", "reshape", "transpose", + "where", + *BaseFakeNumpyNamespace._numpy_math_functions + }) + + +class CupyFakeNumpyNamespace(BaseFakeNumpyNamespace): + """ + A :mod:`cupy` mimic for :class:`CupyArrayContext`. + """ + def _get_fake_numpy_linalg_namespace(self): + return CupyFakeNumpyLinalgNamespace(self._array_context) + + def zeros(self, shape, dtype): + import cupy as cp # type: ignore[import-untyped] + return cp.zeros(shape, dtype) + + def __getattr__(self, name): + import cupy as cp + + if name in _NUMPY_UFUNCS: + from functools import partial + return partial(rec_multimap_array_container, + getattr(cp, name)) + + raise AttributeError(name) + + def sum(self, a, axis=None, dtype=None): + import cupy as cp + return rec_map_reduce_array_container(sum, partial(cp.sum, + axis=axis, + dtype=dtype), + a) + + def min(self, a, axis=None): + import cupy as cp + return rec_map_reduce_array_container( + partial(reduce, cp.minimum), partial(cp.amin, axis=axis), a) + + def max(self, a, axis=None): + import cupy as cp + return rec_map_reduce_array_container( + partial(reduce, cp.maximum), partial(cp.amax, axis=axis), a) + + def stack(self, arrays, axis=0): + import cupy as cp + return rec_multimap_array_container( + lambda *args: cp.stack(args, axis=axis), + *arrays) + + def broadcast_to(self, array, shape): + import cupy as cp + return rec_map_array_container(partial(cp.broadcast_to, shape=shape), array) + + # {{{ relational operators + + def equal(self, x, y): + import cupy as cp + return rec_multimap_array_container(cp.equal, x, y) + + def not_equal(self, x, y): + import cupy as cp + return rec_multimap_array_container(cp.not_equal, x, y) + + def greater(self, x, y): + import cupy as cp + return rec_multimap_array_container(cp.greater, x, y) + + def greater_equal(self, x, y): + import cupy as cp + return rec_multimap_array_container(cp.greater_equal, x, y) + + def less(self, x, y): + import cupy as cp + return rec_multimap_array_container(cp.less, x, y) + + def less_equal(self, x, y): + import cupy as cp + return rec_multimap_array_container(cp.less_equal, x, y) + + # }}} + + def ravel(self, a, order="C"): + import cupy as cp + return rec_map_array_container(partial(cp.ravel, order=order), a) + + def vdot(self, x, y): + import cupy as cp + return rec_multimap_reduce_array_container(sum, cp.vdot, x, y) + + def any(self, a): + import cupy as cp + return rec_map_reduce_array_container(partial(reduce, cp.logical_or), + lambda subary: cp.any(subary), a) + + def all(self, a): + import cupy as cp + return rec_map_reduce_array_container(partial(reduce, cp.logical_and), + lambda subary: cp.all(subary), a) + + def array_equal(self, a: ArrayOrContainer, b: ArrayOrContainer) -> Array: + import cupy as cp + false_ary = cp.array(False) + true_ary = cp.array(True) + if type(a) is not type(b): + return false_ary + + try: + serialized_x = serialize_container(a) + serialized_y = serialize_container(b) + except NotAnArrayContainerError: + assert isinstance(a, cp.ndarray) + assert isinstance(b, cp.ndarray) + return cp.array(cp.array_equal(a, b)) + else: + if len(serialized_x) != len(serialized_y): + return false_ary + return reduce( + cp.logical_and, + [(true_ary if kx_i == ky_i else false_ary) + and self.array_equal(x_i, y_i) + for (kx_i, x_i), (ky_i, y_i) + in zip(serialized_x, serialized_y, strict=True)], + true_ary) + + def arange(self, *args, **kwargs): + import cupy as cp + return cp.arange(*args, **kwargs) + + def linspace(self, *args, **kwargs): + import cupy as cp + return cp.linspace(*args, **kwargs) + + def zeros_like(self, ary): + import cupy as cp + if isinstance(ary, int | float | complex): + # Cupy does not support zeros_like with scalar arguments + ary = cp.array(ary) + return rec_map_array_container(cp.zeros_like, ary) + + def ones_like(self, ary): + import cupy as cp + if isinstance(ary, int | float | complex): + # Cupy does not support ones_like with scalar arguments + ary = cp.array(ary) + return rec_map_array_container(cp.ones_like, ary) + + def reshape(self, a, newshape, order="C"): + return rec_map_array_container( + lambda ary: ary.reshape(newshape, order=order), + a) + + +# vim: fdm=marker diff --git a/arraycontext/impl/numpy/__init__.py b/arraycontext/impl/numpy/__init__.py index f9d6c541..6fcc88c7 100644 --- a/arraycontext/impl/numpy/__init__.py +++ b/arraycontext/impl/numpy/__init__.py @@ -63,7 +63,7 @@ class NumpyNonObjectArray(metaclass=NumpyNonObjectArrayMetaclass): class NumpyArrayContext(ArrayContext): """ - A :class:`ArrayContext` that uses :class:`numpy.ndarray` to represent arrays. + An :class:`ArrayContext` that uses :class:`numpy.ndarray` to represent arrays. .. automethod:: __init__ """ diff --git a/arraycontext/impl/pyopencl/__init__.py b/arraycontext/impl/pyopencl/__init__.py index 84d5f483..190599e2 100644 --- a/arraycontext/impl/pyopencl/__init__.py +++ b/arraycontext/impl/pyopencl/__init__.py @@ -58,7 +58,7 @@ class PyOpenCLArrayContext(ArrayContext): """ - A :class:`ArrayContext` that uses :class:`pyopencl.array.Array` instances + An :class:`ArrayContext` that uses :class:`pyopencl.array.Array` instances for its base array class. .. attribute:: context diff --git a/arraycontext/pytest.py b/arraycontext/pytest.py index 760fc103..c85c7a22 100644 --- a/arraycontext/pytest.py +++ b/arraycontext/pytest.py @@ -227,6 +227,26 @@ def __str__(self): return "" +class _PytestCupyArrayContextFactory(PytestArrayContextFactory): + @classmethod + def is_available(cls) -> bool: + try: + import cupy # type: ignore[import-untyped] # noqa: F401 + return True + except ImportError: + return False + + def __call__(self): + from arraycontext import CupyArrayContext + return CupyArrayContext() + + def __str__(self): + import cupy # pylint: disable=import-error + d = cupy.cuda.runtime.getDeviceProperties(cupy.cuda.Device()) + name = d["name"].decode("utf-8") + return f" on {cupy.cuda.Device()}:{name}" + + # {{{ _PytestArrayContextFactory class _NumpyArrayContextForTests(NumpyArrayContext): @@ -253,6 +273,7 @@ def __str__(self): "pytato:jax": _PytestPytatoJaxArrayContextFactory, "eagerjax": _PytestEagerJaxArrayContextFactory, "numpy": _PytestNumpyArrayContextFactory, + "cupy": _PytestCupyArrayContextFactory, } diff --git a/doc/Makefile b/doc/Makefile index d0ac5f2f..0568a00c 100644 --- a/doc/Makefile +++ b/doc/Makefile @@ -3,7 +3,7 @@ # You can set these variables from the command line, and also # from the environment for the first two. -SPHINXOPTS ?= +SPHINXOPTS ?= -W -n SPHINXBUILD ?= python $(shell which sphinx-build) SOURCEDIR = . BUILDDIR = _build diff --git a/doc/conf.py b/doc/conf.py index 0042ae57..04b3971b 100644 --- a/doc/conf.py +++ b/doc/conf.py @@ -23,6 +23,7 @@ "pytest": ("https://docs.pytest.org/en/latest/", None), "python": ("https://docs.python.org/3/", None), "pytools": ("https://documen.tician.de/pytools", None), + "cupy": ("https://docs.cupy.dev/en/stable/", None), } # Some modules need to import things just so that sphinx can resolve symbols in diff --git a/doc/implementations.rst b/doc/implementations.rst index 4023e37c..2e6344e8 100644 --- a/doc/implementations.rst +++ b/doc/implementations.rst @@ -13,6 +13,12 @@ Array context based on :mod:`numpy` .. automodule:: arraycontext.impl.numpy + +Array context based on :mod:`cupy` +-------------------------------------------- + +.. automodule:: arraycontext.impl.cupy + Array context based on :mod:`pyopencl.array` -------------------------------------------- diff --git a/doc/index.rst b/doc/index.rst index d3f9854b..2445d2e2 100644 --- a/doc/index.rst +++ b/doc/index.rst @@ -6,6 +6,7 @@ code to work with all of them? No problem! Comes with pre-made array context implementations for: - :mod:`numpy` +- :mod:`cupy` - :mod:`pyopencl` - :mod:`jax.numpy` - :mod:`pytato` (for lazy/deferred evaluation) @@ -13,7 +14,7 @@ implementations for: - Profiling :mod:`arraycontext` started life as an array abstraction for use with the -:mod:`meshmode` unstrucuted discretization package. +:mod:`meshmode` unstructured discretization package. Design Guidelines ----------------- diff --git a/test/test_arraycontext.py b/test/test_arraycontext.py index 11ccbb1f..7d7e5a21 100644 --- a/test/test_arraycontext.py +++ b/test/test_arraycontext.py @@ -34,6 +34,7 @@ from pytools.tag import Tag from arraycontext import ( + CupyArrayContext, EagerJAXArrayContext, NumpyArrayContext, PyOpenCLArrayContext, @@ -45,6 +46,7 @@ with_container_arithmetic, ) from arraycontext.pytest import ( + _PytestCupyArrayContextFactory, _PytestEagerJaxArrayContextFactory, _PytestNumpyArrayContextFactory, _PytestPyOpenCLArrayContextFactoryWithClass, @@ -99,6 +101,7 @@ class _PytatoPyOpenCLArrayContextForTestsFactory( _PytatoPyOpenCLArrayContextForTestsFactory, _PytestEagerJaxArrayContextFactory, _PytestPytatoJaxArrayContextFactory, + _PytestCupyArrayContextFactory, _PytestNumpyArrayContextFactory, ]) @@ -1027,7 +1030,7 @@ def test_numpy_conversion(actx_factory): assert np.allclose(ac.mass, ac_roundtrip.mass) assert np.allclose(ac.momentum[0], ac_roundtrip.momentum[0]) - if not isinstance(actx, NumpyArrayContext): + if not isinstance(actx, NumpyArrayContext | CupyArrayContext): from dataclasses import replace ac_with_cl = replace(ac, enthalpy=ac_actx.mass) with pytest.raises(TypeError): @@ -1382,7 +1385,7 @@ def test_to_numpy_on_frozen_arrays(actx_factory): def test_tagging(actx_factory): actx = actx_factory() - if isinstance(actx, NumpyArrayContext | EagerJAXArrayContext): + if isinstance(actx, NumpyArrayContext | EagerJAXArrayContext | CupyArrayContext): pytest.skip(f"{type(actx)} has no tagging support") from pytools.tag import Tag @@ -1431,6 +1434,9 @@ def test_linspace(actx_factory, args, kwargs): actx = actx_factory() + if isinstance(actx, CupyArrayContext) and kwargs.get("dtype") == np.complex128: + pytest.skip("CupyArrayContext does not support complex args to linspace") + actx_linspace = actx.to_numpy(actx.np.linspace(*args, **kwargs)) np_linspace = np.linspace(*args, **kwargs)