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)
+
+# }}}

arraycontext/loopy.py

@@ -2,6 +2,8 @@
 .. currentmodule:: arraycontext
 .. autofunction:: make_loopy_program
 """
+from __future__ import annotations
+
 
 __copyright__ = """
 Copyright (C) 2020-1 University of Illinois Board of Trustees
@@ -27,8 +29,17 @@ 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
@@ -64,9 +75,96 @@ def get_default_entrypoint(t_unit):
     except AttributeError:
         try:
             return t_unit.root_kernel
-        except AttributeError:
+        except AttributeError as err:
             raise TypeError("unable to find default entry point for loopy "
-                    "translation unit")
+                    "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}'")
 
 # }}}
 

arraycontext/metadata.py

+"""
+.. autoclass:: NameHint
+"""
+from __future__ import annotations
+
+
 __copyright__ = """
 Copyright (C) 2020-1 University of Illinois Board of Trustees
 """
@@ -22,36 +28,25 @@ OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
 THE SOFTWARE.
 """
 
-import sys
-from pytools.tag import Tag
-from warnings import warn
+from dataclasses import dataclass
 
+from pytools.tag import UniqueTag
 
-# {{{ deprecation handling
 
-try:
-    from meshmode.transform_metadata import FirstAxisIsElementsTag \
-            as _FirstAxisIsElementsTag
-except ImportError:
-    # placeholder in case meshmode is too old to have it.
-    class _FirstAxisIsElementsTag(Tag):  # type: ignore[no-redef]
-        pass
+@dataclass(frozen=True)
+class NameHint(UniqueTag):
+    """A tag acting on arrays or array axes. Express that :attr:`name` is a
+    useful starting point in forming an identifier for the tagged object.
 
+    .. attribute:: name
 
-if sys.version_info >= (3, 7):
-    def __getattr__(name):
-        if name == "FirstAxisIsElementsTag":
-            warn(f"'arraycontext.{name}' is deprecated. "
-                    f"Use 'meshmode.transform_metadata.{name}' instead. "
-                    f"'arraycontext.{name}' will continue to work until 2022.",
-                    DeprecationWarning, stacklevel=2)
-            return _FirstAxisIsElementsTag
-        else:
-            raise AttributeError(name)
-else:
-    FirstAxisIsElementsTag = _FirstAxisIsElementsTag
+        A string. Must be a valid Python identifier. Not necessarily unique.
+    """
+    name: str
 
-# }}}
+    def __post_init__(self):
+        if not self.name.isidentifier():
+            raise ValueError("'name' must be an identifier")
 
 
 # vim: foldmethod=marker