arraycontext/impl/numpy/fake_numpy.py

+from __future__ import annotations
+
+
+__copyright__ = """
+Copyright (C) 2021 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 typing import cast
+
+import numpy as np
+
+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 NumpyFakeNumpyLinalgNamespace(BaseFakeNumpyLinalgNamespace):
+    # Everything is implemented in the base class for now.
+    pass
+
+
+_NUMPY_UFUNCS = frozenset({"concatenate", "reshape", "transpose",
+                 "ones_like", "where",
+                 *BaseFakeNumpyNamespace._numpy_math_functions
+                 })
+
+
+class NumpyFakeNumpyNamespace(BaseFakeNumpyNamespace):
+    """
+    A :mod:`numpy` mimic for :class:`NumpyArrayContext`.
+    """
+    def _get_fake_numpy_linalg_namespace(self):
+        return NumpyFakeNumpyLinalgNamespace(self._array_context)
+
+    def zeros(self, shape, dtype):
+        return np.zeros(shape, dtype)
+
+    def __getattr__(self, name):
+
+        if name in _NUMPY_UFUNCS:
+            from functools import partial
+            return partial(rec_multimap_array_container,
+                           getattr(np, name))
+
+        raise AttributeError(name)
+
+    def sum(self, a, axis=None, dtype=None):
+        return rec_map_reduce_array_container(sum, partial(np.sum,
+                                                           axis=axis,
+                                                           dtype=dtype),
+                                              a)
+
+    def min(self, a, axis=None):
+        return rec_map_reduce_array_container(
+                partial(reduce, np.minimum), partial(np.amin, axis=axis), a)
+
+    def max(self, a, axis=None):
+        return rec_map_reduce_array_container(
+                partial(reduce, np.maximum), partial(np.amax, axis=axis), a)
+
+    def stack(self, arrays, axis=0):
+        return rec_multimap_array_container(
+                lambda *args: np.stack(arrays=args, axis=axis),
+                *arrays)
+
+    def broadcast_to(self, array, shape):
+        return rec_map_array_container(partial(np.broadcast_to, shape=shape), array)
+
+    # {{{ relational operators
+
+    def equal(self, x, y):
+        return rec_multimap_array_container(np.equal, x, y)
+
+    def not_equal(self, x, y):
+        return rec_multimap_array_container(np.not_equal, x, y)
+
+    def greater(self, x, y):
+        return rec_multimap_array_container(np.greater, x, y)
+
+    def greater_equal(self, x, y):
+        return rec_multimap_array_container(np.greater_equal, x, y)
+
+    def less(self, x, y):
+        return rec_multimap_array_container(np.less, x, y)
+
+    def less_equal(self, x, y):
+        return rec_multimap_array_container(np.less_equal, x, y)
+
+    # }}}
+
+    def ravel(self, a, order="C"):
+        return rec_map_array_container(partial(np.ravel, order=order), a)
+
+    def vdot(self, x, y):
+        return rec_multimap_reduce_array_container(sum, np.vdot, x, y)
+
+    def any(self, a):
+        return rec_map_reduce_array_container(partial(reduce, np.logical_or),
+                                              lambda subary: np.any(subary), a)
+
+    def all(self, a):
+        return rec_map_reduce_array_container(partial(reduce, np.logical_and),
+                                              lambda subary: np.all(subary), a)
+
+    def array_equal(self, a: ArrayOrContainer, b: ArrayOrContainer) -> Array:
+        false_ary = np.array(False)
+        true_ary = np.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, np.ndarray)
+            assert isinstance(b, np.ndarray)
+            return np.array(np.array_equal(a, b))
+        else:
+            if len(serialized_x) != len(serialized_y):
+                return false_ary
+            return np.logical_and.reduce(
+                    [(true_ary if kx_i == ky_i else false_ary)
+                        and cast(np.ndarray, self.array_equal(x_i, y_i))
+                        for (kx_i, x_i), (ky_i, y_i)
+                        in zip(serialized_x, serialized_y, strict=True)],
+                    initial=true_ary)
+
+    def arange(self, *args, **kwargs):
+        return np.arange(*args, **kwargs)
+
+    def linspace(self, *args, **kwargs):
+        return np.linspace(*args, **kwargs)
+
+    def zeros_like(self, ary):
+        return rec_map_array_container(np.zeros_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

arraycontext/loopy.py

+"""
+.. currentmodule:: arraycontext
+.. autofunction:: make_loopy_program
+"""
+from __future__ import annotations
+
+
+__copyright__ = """
+Copyright (C) 2020-1 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 collections.abc import Mapping
+from typing import ClassVar
+
+import numpy as np
+
+import loopy as lp
+from loopy.version import MOST_RECENT_LANGUAGE_VERSION
+from pytools import memoize_in
+
+from arraycontext.container.traversal import multimapped_over_array_containers
+from arraycontext.fake_numpy import BaseFakeNumpyNamespace
+
+
+# {{{ loopy
+
+_DEFAULT_LOOPY_OPTIONS = lp.Options(
+        no_numpy=True,
+        return_dict=True)
+
+
+def make_loopy_program(domains, statements, kernel_data=None,
+        name="mm_actx_kernel", tags=None):
+    """Return a :class:`loopy.LoopKernel` suitable for use with
+    :meth:`ArrayContext.call_loopy`.
+    """
+    if kernel_data is None:
+        kernel_data = ["..."]
+
+    return lp.make_kernel(
+            domains,
+            statements,
+            kernel_data=kernel_data,
+            options=_DEFAULT_LOOPY_OPTIONS,
+            default_offset=lp.auto,
+            name=name,
+            lang_version=MOST_RECENT_LANGUAGE_VERSION,
+            tags=tags)
+
+
+def get_default_entrypoint(t_unit):
+    try:
+        # main and "kernel callables" branch
+        return t_unit.default_entrypoint
+    except AttributeError:
+        try:
+            return t_unit.root_kernel
+        except AttributeError as err:
+            raise TypeError("unable to find default entry point for loopy "
+                    "translation unit") from err
+
+
+def _get_scalar_func_loopy_program(actx, c_name, nargs, naxes):
+    @memoize_in(actx, _get_scalar_func_loopy_program)
+    def get(c_name, nargs, naxes):
+        from pymbolic.primitives import Subscript, Variable
+
+        var_names = [f"i{i}" for i in range(naxes)]
+        size_names = [f"n{i}" for i in range(naxes)]
+        subscript = tuple(Variable(vname) for vname in var_names)
+
+        from islpy import make_zero_and_vars
+
+        v = make_zero_and_vars(var_names, params=size_names)
+        domain = v[0].domain()
+        for vname, sname in zip(var_names, size_names, strict=True):
+            domain = domain & v[0].le_set(v[vname]) & v[vname].lt_set(v[sname])
+
+        domain_bset, = domain.get_basic_sets()
+
+        import loopy as lp
+
+        from arraycontext.transform_metadata import ElementwiseMapKernelTag
+
+        def sub(name: str) -> Variable | Subscript:
+            return Subscript(Variable(name), subscript) if subscript else Variable(name)
+
+        return make_loopy_program(
+                [domain_bset], [
+                    lp.Assignment(
+                        sub("out"),
+                        Variable(c_name)(*[sub(f"inp{i}") for i in range(nargs)]))
+                ], [
+                    lp.GlobalArg("out", dtype=None, shape=lp.auto, offset=lp.auto)
+                ] + [
+                    lp.GlobalArg(f"inp{i}", dtype=None, shape=lp.auto, offset=lp.auto)
+                    for i in range(nargs)
+                ] + [...],
+                name=f"actx_special_{c_name}",
+                tags=(ElementwiseMapKernelTag(),))
+
+    return get(c_name, nargs, naxes)
+
+
+class LoopyBasedFakeNumpyNamespace(BaseFakeNumpyNamespace):
+    _numpy_to_c_arc_functions: ClassVar[Mapping[str, str]] = {
+            "arcsin": "asin",
+            "arccos": "acos",
+            "arctan": "atan",
+            "arctan2": "atan2",
+
+            "arcsinh": "asinh",
+            "arccosh": "acosh",
+            "arctanh": "atanh",
+            }
+
+    _c_to_numpy_arc_functions: ClassVar[Mapping[str, str]] = {c_name: numpy_name
+            for numpy_name, c_name in _numpy_to_c_arc_functions.items()}
+
+    def __getattr__(self, name):
+        def loopy_implemented_elwise_func(*args):
+            if all(np.isscalar(ary) for ary in args):
+                return getattr(
+                         np, self._c_to_numpy_arc_functions.get(name, name)
+                         )(*args)
+            actx = self._array_context
+            prg = _get_scalar_func_loopy_program(actx,
+                    c_name, nargs=len(args), naxes=len(args[0].shape))
+            outputs = actx.call_loopy(prg,
+                    **{f"inp{i}": arg for i, arg in enumerate(args)})
+            return outputs["out"]
+
+        if name in self._c_to_numpy_arc_functions:
+            raise RuntimeError(f"'{name}' in ArrayContext.np has been removed. "
+                    f"Use '{self._c_to_numpy_arc_functions[name]}' as in numpy. ")
+
+        # normalize to C names anyway
+        c_name = self._numpy_to_c_arc_functions.get(name, name)
+
+        # limit which functions we try to hand off to loopy
+        if (name in self._numpy_math_functions
+                or name in self._c_to_numpy_arc_functions):
+            return multimapped_over_array_containers(loopy_implemented_elwise_func)
+        else:
+            raise AttributeError(
+                    f"'{type(self._array_context).__name__}.np' object "
+                    f"has no attribute '{name}'")
+
+# }}}
+
+
+# vim: foldmethod=marker

doc/Makefile

@@ -4,7 +4,7 @@
 # You can set these variables from the command line, and also
 # from the environment for the first two.
 SPHINXOPTS    ?=
-SPHINXBUILD   ?= sphinx-build
+SPHINXBUILD   ?= python $(shell which sphinx-build)
 SOURCEDIR     = .
 BUILDDIR      = _build
 

doc/array_context.rst

+The Array Context Abstraction
+=============================
+
+.. automodule:: arraycontext
+
+.. automodule:: arraycontext.context