From fe16c3d4221b17ce142d0788a6d60ba2389db18a Mon Sep 17 00:00:00 2001 From: Kaushik Kulkarni Date: Sun, 31 Oct 2021 23:35:08 -0500 Subject: [PATCH 1/4] tag pt arrays' axes --- arraycontext/impl/pytato/__init__.py | 6 +----- 1 file changed, 1 insertion(+), 5 deletions(-) diff --git a/arraycontext/impl/pytato/__init__.py b/arraycontext/impl/pytato/__init__.py index a7a0b1b2..19ba219e 100644 --- a/arraycontext/impl/pytato/__init__.py +++ b/arraycontext/impl/pytato/__init__.py @@ -206,11 +206,7 @@ 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 54ae38a75fc50affb211c2c7b13edb48ec963ccc Mon Sep 17 00:00:00 2001 From: Kaushik Kulkarni Date: Tue, 19 Oct 2021 09:23:33 -0500 Subject: [PATCH 2/4] define TaggableCLArray --- arraycontext/impl/pytato/utils.py | 80 ++++++++++++++++++++++++++++++- 1 file changed, 78 insertions(+), 2 deletions(-) diff --git a/arraycontext/impl/pytato/utils.py b/arraycontext/impl/pytato/utils.py index f69a1961..f23c344c 100644 --- a/arraycontext/impl/pytato/utils.py +++ b/arraycontext/impl/pytato/utils.py @@ -23,11 +23,87 @@ """ -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, + tags=tags) class _DatawrapperToBoundPlaceholderMapper(CopyMapper): From d25652d272a3ac2db11322b4ab76cdc1ba4ef206 Mon Sep 17 00:00:00 2001 From: Kaushik Kulkarni Date: Tue, 19 Oct 2021 09:26:34 -0500 Subject: [PATCH 3/4] make the frozen type of PytatoPyOpenCLArrayContext to be TaggableCLArrays --- arraycontext/impl/pytato/__init__.py | 41 +++++++++++++++++----------- arraycontext/impl/pytato/compile.py | 9 ++++-- arraycontext/impl/pytato/utils.py | 2 ++ 3 files changed, 34 insertions(+), 18 deletions(-) diff --git a/arraycontext/impl/pytato/__init__.py b/arraycontext/impl/pytato/__init__.py index 19ba219e..2bcf90bc 100644 --- a/arraycontext/impl/pytato/__init__.py +++ b/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) # }}} @@ -209,7 +218,7 @@ def tag_axis(self, iaxis, tags: Union[Sequence[Tag], Tag], 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 @@ -217,7 +226,7 @@ def einsum(self, spec, *args, arg_names=None, tagged=()): "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) diff --git a/arraycontext/impl/pytato/compile.py b/arraycontext/impl/pytato/compile.py index a772a85a..2c1e3a95 100644 --- a/arraycontext/impl/pytato/compile.py +++ b/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) diff --git a/arraycontext/impl/pytato/utils.py b/arraycontext/impl/pytato/utils.py index f23c344c..1463fb71 100644 --- a/arraycontext/impl/pytato/utils.py +++ b/arraycontext/impl/pytato/utils.py @@ -103,6 +103,7 @@ def to_tagged_cl_array(ary: cl_array.Array, events=ary.events, _fast=True, _context=ary.context, _queue=ary.queue, _size=ary.size, + axes=axes, tags=tags) @@ -133,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: From d47f8911182c852b7e82362bcc30e1dda01a890e Mon Sep 17 00:00:00 2001 From: Kaushik Kulkarni Date: Tue, 19 Oct 2021 09:28:04 -0500 Subject: [PATCH 4/4] adds tests specific to PytatoArrayContexts --- test/test_pytato_arraycontext.py | 90 ++++++++++++++++++++++++++++++++ 1 file changed, 90 insertions(+) create mode 100644 test/test_pytato_arraycontext.py diff --git a/test/test_pytato_arraycontext.py b/test/test_pytato_arraycontext.py new file mode 100644 index 00000000..9768f6ad --- /dev/null +++ b/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