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/impl/pyopencl/taggable_cl_array.py

+"""
+.. autoclass:: TaggableCLArray
+.. autoclass:: Axis
+
+.. autofunction:: to_tagged_cl_array
+"""
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Any
+
+import numpy as np
+
+import pyopencl.array as cla
+from pytools import memoize
+from pytools.tag import Tag, Taggable, ToTagSetConvertible
+
+
+# {{{ utils
+
+@dataclass(frozen=True, eq=True)
+class Axis(Taggable):
+    """
+    Records the tags corresponding to a dimension of :class:`TaggableCLArray`.
+    """
+
+    tags: frozenset[Tag]
+
+    def _with_new_tags(self, tags: frozenset[Tag]) -> Axis:
+        from dataclasses import replace
+        return replace(self, tags=tags)
+
+
+@memoize
+def _construct_untagged_axes(ndim: int) -> tuple[Axis, ...]:
+    return tuple(Axis(frozenset()) for _ in range(ndim))
+
+
+def _unwrap_cl_array(ary: cla.Array) -> dict[str, Any]:
+    return {
+        "shape": ary.shape,
+        "dtype": ary.dtype,
+        "allocator": ary.allocator,
+        "strides": ary.strides,
+        "data": ary.base_data,
+        "offset": ary.offset,
+        "events": ary.events,
+        "_context": ary.context,
+        "_queue": ary.queue,
+        "_size": ary.size,
+        "_fast": True,
+        }
+
+# }}}
+
+
+# {{{ TaggableCLArray
+
+class TaggableCLArray(cla.Array, Taggable):
+    """
+    A :class:`pyopencl.array.Array` with additional metadata. This is used by
+    :class:`~arraycontext.PytatoPyOpenCLArrayContext` to preserve tags for data
+    while frozen, and also in a similar capacity by
+    :class:`~arraycontext.PyOpenCLArrayContext`.
+
+    .. attribute:: axes
+
+       A :class:`tuple` of instances of :class:`Axis`, with one :class:`Axis`
+       for each dimension of the array.
+
+    .. attribute:: tags
+
+        A :class:`frozenset` of :class:`pytools.tag.Tag`. Typically intended to
+        record application-specific metadata to drive the optimizations in
+        :meth:`arraycontext.PyOpenCLArrayContext.transform_loopy_program`.
+    """
+    def __init__(self, cq, shape, dtype, order="C", allocator=None,
+                 data=None, offset=0, strides=None, events=None, _flags=None,
+                 _fast=False, _size=None, _context=None, _queue=None,
+                 axes=None, tags=frozenset()):
+
+        super().__init__(cq=cq, shape=shape, dtype=dtype,
+                         order=order, allocator=allocator,
+                         data=data, offset=offset,
+                         strides=strides, events=events,
+                         _flags=_flags, _fast=_fast,
+                         _size=_size, _context=_context,
+                         _queue=_queue)
+
+        if __debug__:
+            if not isinstance(tags, frozenset):
+                raise TypeError("tags are not a frozenset")
+
+            if axes is not None and len(axes) != self.ndim:
+                raise ValueError("axes length does not match array dimension: "
+                                 f"got {len(axes)} axes for {self.ndim}d array")
+
+        if axes is None:
+            axes = _construct_untagged_axes(self.ndim)
+
+        self.tags = tags
+        self.axes = axes
+
+    def __repr__(self) -> str:
+        return (f"{type(self).__name__}(shape={self.shape}, dtype={self.dtype}, "
+                f"tags={self.tags}, axes={self.axes})")
+
+    def copy(self, queue=cla._copy_queue):
+        ary = super().copy(queue=queue)
+        return type(self)(None, tags=self.tags, axes=self.axes,
+                          **_unwrap_cl_array(ary))
+
+    def _with_new_tags(self, tags: frozenset[Tag]) -> TaggableCLArray:
+        return type(self)(None, tags=tags, axes=self.axes,
+                          **_unwrap_cl_array(self))
+
+    def with_tagged_axis(self, iaxis: int,
+                         tags: ToTagSetConvertible) -> TaggableCLArray:
+        """
+        Returns a copy of *self* with *iaxis*-th axis tagged with *tags*.
+        """
+        new_axes = (*self.axes[:iaxis],
+                    self.axes[iaxis].tagged(tags),
+                    *self.axes[iaxis + 1:])
+
+        return type(self)(None, tags=self.tags, axes=new_axes,
+                          **_unwrap_cl_array(self))
+
+
+def to_tagged_cl_array(ary: cla.Array,
+                       axes: tuple[Axis, ...] | None = None,
+                       tags: frozenset[Tag] = frozenset()) -> TaggableCLArray:
+    """
+    Returns a :class:`TaggableCLArray` that is constructed from the data in
+    *ary* along with the metadata from *axes* and *tags*. If *ary* is already a
+    :class:`TaggableCLArray`, the new *tags* and *axes* are added to the
+    existing ones.
+
+    :arg axes: An instance of :class:`Axis` for each dimension of the
+        array. If passed *None*, then initialized to a :class:`pytato.Axis`
+        with no tags attached for each dimension.
+    """
+    if axes is not None and len(axes) != ary.ndim:
+        raise ValueError("axes length does not match array dimension: "
+                         f"got {len(axes)} axes for {ary.ndim}d array")
+
+    from pytools.tag import normalize_tags
+    tags = normalize_tags(tags)
+
+    if isinstance(ary, TaggableCLArray):
+        if axes is not None:
+            for i, axis in enumerate(axes):
+                ary = ary.with_tagged_axis(i, axis.tags)
+
+        if tags:
+            ary = ary.tagged(tags)
+
+        return ary
+    elif isinstance(ary, cla.Array):
+        return TaggableCLArray(None, tags=tags, axes=axes,
+                               **_unwrap_cl_array(ary))
+    else:
+        raise TypeError(f"unsupported array type: '{type(ary).__name__}'")
+
+# }}}
+
+
+# {{{ creation
+
+def empty(queue, shape, dtype=float, *,
+        axes: tuple[Axis, ...] | None = None,
+        tags: frozenset[Tag] = frozenset(),
+        order: str = "C",
+        allocator=None) -> TaggableCLArray:
+    if dtype is not None:
+        dtype = np.dtype(dtype)
+
+    return TaggableCLArray(
+        queue, shape, dtype,
+        axes=axes, tags=tags,
+        order=order, allocator=allocator)
+
+
+def zeros(queue, shape, dtype=float, *,
+        axes: tuple[Axis, ...] | None = None,
+        tags: frozenset[Tag] = frozenset(),
+        order: str = "C",
+        allocator=None) -> TaggableCLArray:
+    result = empty(
+        queue, shape, dtype=dtype, axes=axes, tags=tags,
+        order=order, allocator=allocator)
+    result._zero_fill()
+
+    return result
+
+
+def to_device(queue, ary, *,
+        axes: tuple[Axis, ...] | None = None,
+        tags: frozenset[Tag] = frozenset(),
+        allocator=None):
+    return to_tagged_cl_array(
+        cla.to_device(queue, ary, allocator=allocator),
+        axes=axes, tags=tags)
+
+# }}}