Implements TaggableCLArray

arraycontext/impl/pytato/__init__.py

@@ -112,13 +112,13 @@ def to_numpy(self, array):
         return cl_array.get(queue=self.queue)
 
     def call_loopy(self, program, **kwargs):
-        import pyopencl.array as cla
         from pytato.loopy import call_loopy
+        from .utils import TaggableCLArray
 
         entrypoint = program.default_entrypoint.name
 
         # thaw frozen arrays
-        kwargs = {kw: (self.thaw(arg) if isinstance(arg, cla.Array) else arg)
+        kwargs = {kw: (self.thaw(arg) if isinstance(arg, TaggableCLArray) else arg)
                   for kw, arg in kwargs.items()}
 
         return call_loopy(program, kwargs, entrypoint)
@@ -127,8 +127,11 @@ def freeze(self, array):
         import pytato as pt
         import pyopencl.array as cla
         import loopy as lp
+        from arraycontext.impl.pytato.utils import (_normalize_pt_expr,
+                                                    to_tagged_cl_array,
+                                                    TaggableCLArray)
 
-        if isinstance(array, cla.Array):
+        if isinstance(array, TaggableCLArray):
             return array.with_queue(None)
         if not isinstance(array, pt.Array):
             raise TypeError("PytatoPyOpenCLArrayContext.freeze invoked with "
@@ -137,14 +140,16 @@ def freeze(self, array):
         # {{{ early exit for 0-sized arrays
 
         if array.size == 0:
-            return cla.empty(self.queue.context,
-                             shape=array.shape,
-                             dtype=array.dtype,
-                             allocator=self.allocator)
+            return to_tagged_cl_array(
+                cla.empty(self.queue.context,
+                          shape=array.shape,
+                          dtype=array.dtype,
+                          allocator=self.allocator),
+                array.axes,
+                array.tags)
 
         # }}}
 
-        from arraycontext.impl.pytato.utils import _normalize_pt_expr
         pt_dict_of_named_arrays = pt.make_dict_of_named_arrays(
                 {"_actx_out": array})
 
@@ -165,17 +170,21 @@ def freeze(self, array):
         evt, out_dict = pt_prg(self.queue, **bound_arguments)
         evt.wait()
 
-        return out_dict["_actx_out"].with_queue(None)
+        return to_tagged_cl_array(out_dict["_actx_out"].with_queue(None),
+                                  array.axes,
+                                  array.tags)
 
     def thaw(self, array):
         import pytato as pt
-        import pyopencl.array as cla
+        from .utils import TaggableCLArray
 
-        if not isinstance(array, cla.Array):
-            raise TypeError("PytatoPyOpenCLArrayContext.thaw expects CL arrays, got "
-                    f"{type(array)}")
+        if not isinstance(array, TaggableCLArray):
+            raise TypeError("PytatoPyOpenCLArrayContext.thaw expects "
+                            f"TaggableCLArray, got {type(array)}.")
 
-        return pt.make_data_wrapper(array.with_queue(self.queue))
+        return pt.make_data_wrapper(array.with_queue(self.queue),
+                                    axes=array.axes,
+                                    tags=array.tags)
 
     # }}}
 
@@ -206,22 +215,18 @@ def tag(self, tags: Union[Sequence[Tag], Tag], array):
         return array.tagged(tags)
 
     def tag_axis(self, iaxis, tags: Union[Sequence[Tag], Tag], array):
-        # TODO
-        from warnings import warn
-        warn("tagging PytatoPyOpenCLArrayContext's array axes: not yet implemented",
-             stacklevel=2)
-        return array
+        return array.with_tagged_axis(iaxis, tags)
 
     def einsum(self, spec, *args, arg_names=None, tagged=()):
-        import pyopencl.array as cla
+        from .utils import TaggableCLArray
         import pytato as pt
         if arg_names is not None:
             from warnings import warn
             warn("'arg_names' don't bear any significance in "
                  "PytatoPyOpenCLArrayContext.", stacklevel=2)
 
         def preprocess_arg(arg):
-            if isinstance(arg, cla.Array):
+            if isinstance(arg, TaggableCLArray):
                 return self.thaw(arg)
             else:
                 assert isinstance(arg, pt.Array)

arraycontext/impl/pytato/compile.py

@@ -320,6 +320,7 @@ def __call__(self, arg_id_to_arg) -> ArrayContainer:
             representation.
         """
         from arraycontext.container.traversal import rec_keyed_map_array_container
+        from .utils import TaggableCLArray, to_tagged_cl_array
 
         input_kwargs_to_loopy = {}
 
@@ -332,7 +333,7 @@ def __call__(self, arg_id_to_arg) -> ArrayContainer:
             elif isinstance(arg, pt.array.DataWrapper):
                 # got a Datwwrapper => simply gets its data
                 arg = arg.data
-            elif isinstance(arg, cla.Array):
+            elif isinstance(arg, TaggableCLArray):
                 # got a frozen array  => do nothing
                 pass
             elif isinstance(arg, pt.Array):
@@ -354,7 +355,11 @@ def __call__(self, arg_id_to_arg) -> ArrayContainer:
         # }}}
 
         def to_output_template(keys, _):
-            return self.actx.thaw(out_dict[self.output_id_to_name_in_program[keys]])
+            # TODO: What should be done about the tags here?
+            return self.actx.thaw(to_tagged_cl_array(
+                out_dict[self.output_id_to_name_in_program[keys]],
+                None,
+                frozenset()))
 
         return rec_keyed_map_array_container(to_output_template,
                                              self.output_template)

arraycontext/impl/pytato/utils.py

@@ -23,11 +23,88 @@
 """
 
 
-from typing import Any, Dict, Set, Tuple, Mapping
-from pytato.array import SizeParam, Placeholder, make_placeholder
+from typing import Any, Dict, Set, Tuple, Mapping, FrozenSet, Optional
+from pytato.array import SizeParam, Placeholder, make_placeholder, Axis
 from pytato.array import Array, DataWrapper, DictOfNamedArrays
 from pytato.transform import CopyMapper
 from pytools import UniqueNameGenerator
+from pytools.tag import Taggable, Tag
+import pyopencl.array as cl_array
+
+
+class TaggableCLArray(cl_array.Array, Taggable):
+    """
+    A :class:`cl_array.Array` that can be tagged.
+    """
+    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)
+
+        self.tags = tags
+        self.axes = axes
+
+    def copy(self, queue=cl_array._copy_queue, tags=None, axes=None):
+        if tags is not None or axes is not None:
+            if queue is not cl_array._copy_queue:
+                raise ValueError("Cannot change both 'tags'/'axes' and 'queue'"
+                                 " at once.")
+            tags = self.tags if tags is None else tags
+            axes = self.axes if axes is None else axes
+            return self.__class__(None, self.shape, self.dtype,
+                                  allocator=self.allocator,
+                                  strides=self.strides, data=self.base_data,
+                                  offset=self.offset, events=self.events,
+                                  _fast=True, _context=self.context,
+                                  _queue=self.queue, _size=self.size,
+                                  tags=tags, axes=axes)
+        else:
+            new_with_queue = super().copy(queue=queue)
+            return self.__class__(None, new_with_queue.shape,
+                                  new_with_queue.dtype,
+                                  allocator=new_with_queue.allocator,
+                                  strides=new_with_queue.strides,
+                                  data=new_with_queue.base_data,
+                                  offset=new_with_queue.offset,
+                                  events=new_with_queue.events, _fast=True,
+                                  _context=new_with_queue.context,
+                                  _queue=queue, _size=new_with_queue.size,
+                                  tags=new_with_queue.tags,
+                                  axes=new_with_queue.axes)
+
+
+def to_tagged_cl_array(ary: cl_array.Array,
+                       axes: Optional[Tuple[Axis, ...]],
+                       tags: FrozenSet[Tag]) -> TaggableCLArray:
+    """
+    Converts a *ary* to a :class:`TaggableCLArray` with *tags* attached to it.
+
+    :arg axes: An instance of :class:`pytato.Axis` for each dimension of the
+        array. If passed *None*, then initialized to a :class:`pytato.Axis`
+        with no tags attached for each dimension.
+    """
+    axes = axes if axes is not None else tuple(Axis(frozenset())
+                                               for _ in ary.shape)
+
+    return TaggableCLArray(None, ary.shape,
+                           ary.dtype,
+                           allocator=ary.allocator,
+                           strides=ary.strides,
+                           data=ary.base_data,
+                           offset=ary.offset,
+                           events=ary.events, _fast=True,
+                           _context=ary.context,
+                           _queue=ary.queue, _size=ary.size,
+                           axes=axes,
+                           tags=tags)
 
 
 class _DatawrapperToBoundPlaceholderMapper(CopyMapper):
@@ -57,6 +134,7 @@ def map_data_wrapper(self, expr: DataWrapper) -> Array:
                     shape=tuple(self.rec(s) if isinstance(s, Array) else s
                                 for s in expr.shape),
                     dtype=expr.dtype,
+                    axes=expr.axes,
                     tags=expr.tags)
 
     def map_size_param(self, expr: SizeParam) -> Array:

test/test_pytato_arraycontext.py

@@ -0,0 +1,90 @@
+""" PytatoArrayContext specific tests"""
+
+__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 arraycontext import (freeze, thaw, PytatoPyOpenCLArrayContext)
+from arraycontext import pytest_generate_tests_for_array_contexts
+from arraycontext.pytest import _PytestPytatoPyOpenCLArrayContextFactory
+from pytools.tag import Tag
+
+
+import logging
+logger = logging.getLogger(__name__)
+
+
+# {{{ pytato-array context fixture
+
+class _PytatoPyOpenCLArrayContextForTests(PytatoPyOpenCLArrayContext):
+    """Like :class:`PytatoPyOpenCLArrayContext`, but applies no program
+    transformations whatsoever. Only to be used for testing internal to
+    :mod:`arraycontext`.
+    """
+
+    def transform_loopy_program(self, t_unit):
+        return t_unit
+
+
+class _PytatoPyOpenCLArrayContextForTestsFactory(
+        _PytestPytatoPyOpenCLArrayContextFactory):
+    actx_class = _PytatoPyOpenCLArrayContextForTests
+
+
+pytest_generate_tests = pytest_generate_tests_for_array_contexts([
+    _PytatoPyOpenCLArrayContextForTestsFactory,
+    ])
+
+# }}}
+
+
+# {{{ dummy tag types
+
+class FooTag(Tag):
+    """
+    Foo
+    """
+
+# }}}
+
+
+def test_tags_preserved_after_freeze(actx_factory):
+    from numpy.random import default_rng
+    rng = default_rng()
+
+    actx = actx_factory()
+    foo = thaw(freeze(actx
+                      .from_numpy(rng.random((10, 4)))
+                      .tagged(FooTag()),
+                      actx),
+               actx)
+    assert foo.tags_of_type(FooTag)
+
+
+if __name__ == "__main__":
+    import sys
+    if len(sys.argv) > 1:
+        exec(sys.argv[1])
+    else:
+        from pytest import main
+        main([__file__])
+
+# vim: fdm=marker