From fee9ac5a50f93b5f5eafd82f775e5a27b7021b2a Mon Sep 17 00:00:00 2001 From: Kaushik Kulkarni Date: Mon, 1 Nov 2021 13:45:35 -0500 Subject: [PATCH 1/6] recursively add tags to all arrays in the container --- arraycontext/impl/pytato/__init__.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/arraycontext/impl/pytato/__init__.py b/arraycontext/impl/pytato/__init__.py index 744fcbc4..42619e8e 100644 --- a/arraycontext/impl/pytato/__init__.py +++ b/arraycontext/impl/pytato/__init__.py @@ -42,6 +42,7 @@ """ from arraycontext.context import ArrayContext, _ScalarLike +from arraycontext.container.traversal import rec_map_array_container import numpy as np from typing import Any, Callable, Union, Sequence, TYPE_CHECKING from pytools.tag import Tag @@ -207,7 +208,8 @@ def transform_dag(self, dag: "pytato.DictOfNamedArrays" return dag def tag(self, tags: Union[Sequence[Tag], Tag], array): - return array.tagged(tags) + return rec_map_array_container(lambda x: x.tagged(tags), + array) def tag_axis(self, iaxis, tags: Union[Sequence[Tag], Tag], array): # TODO From b74e2364dc06d1ae6495c3cbffc0dbf99b444952 Mon Sep 17 00:00:00 2001 From: Kaushik Kulkarni Date: Sun, 31 Oct 2021 23:35:08 -0500 Subject: [PATCH 2/6] tag pt arrays' axes --- arraycontext/impl/pytato/__init__.py | 8 +++----- 1 file changed, 3 insertions(+), 5 deletions(-) diff --git a/arraycontext/impl/pytato/__init__.py b/arraycontext/impl/pytato/__init__.py index 42619e8e..f7592520 100644 --- a/arraycontext/impl/pytato/__init__.py +++ b/arraycontext/impl/pytato/__init__.py @@ -212,11 +212,9 @@ def tag(self, tags: Union[Sequence[Tag], Tag], array): array) 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 rec_map_array_container(lambda x: x.with_tagged_axis(iaxis, + tags), + array) def einsum(self, spec, *args, arg_names=None, tagged=()): import pyopencl.array as cla From 611cb2236e568a82f7b61ea435ca52b1a3a2b2b0 Mon Sep 17 00:00:00 2001 From: Kaushik Kulkarni Date: Tue, 19 Oct 2021 09:23:33 -0500 Subject: [PATCH 3/6] define TaggableCLArray MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Co-authored-by: Andreas Klöckner --- arraycontext/impl/pyopencl/__init__.py | 87 ++++++++++-- .../impl/pyopencl/taggable_cl_array.py | 129 ++++++++++++++++++ arraycontext/impl/pytato/utils.py | 11 +- 3 files changed, 212 insertions(+), 15 deletions(-) create mode 100644 arraycontext/impl/pyopencl/taggable_cl_array.py diff --git a/arraycontext/impl/pyopencl/__init__.py b/arraycontext/impl/pyopencl/__init__.py index 7336aa9b..4e6cc384 100644 --- a/arraycontext/impl/pyopencl/__init__.py +++ b/arraycontext/impl/pyopencl/__init__.py @@ -1,6 +1,7 @@ """ .. currentmodule:: arraycontext .. autoclass:: PyOpenCLArrayContext +.. automodule:: arraycontext.impl.pyopencl.taggable_cl_array """ __copyright__ = """ @@ -35,6 +36,7 @@ from pytools.tag import Tag from arraycontext.context import ArrayContext, _ScalarLike +from arraycontext.container.traversal import rec_map_array_container if TYPE_CHECKING: @@ -65,6 +67,8 @@ class PyOpenCLArrayContext(ArrayContext): of arrays are created (e.g. as results of computation), the associated cost may become significant. Using e.g. :class:`pyopencl.tools.MemoryPool` as the allocator can help avoid this cost. + + .. automethod:: transform_loopy_program """ def __init__(self, @@ -109,7 +113,7 @@ def __init__(self, DeprecationWarning, stacklevel=2) import pyopencl as cl - import pyopencl.array as cla + import pyopencl.array as cl_array super().__init__() self.context = queue.context @@ -138,7 +142,9 @@ def __init__(self, self._loopy_transform_cache: \ Dict["lp.TranslationUnit", "lp.TranslationUnit"] = {} - self.array_types = (cla.Array,) + # TODO: Ideally this should only be `(TaggableCLArray,)`, but + # that would break the logic in the downstream users. + self.array_types = (cl_array.Array,) def _get_fake_numpy_namespace(self): from arraycontext.impl.pyopencl.fake_numpy import PyOpenCLFakeNumpyNamespace @@ -147,18 +153,27 @@ def _get_fake_numpy_namespace(self): # {{{ ArrayContext interface def empty(self, shape, dtype): - import pyopencl.array as cl_array - return cl_array.empty(self.queue, shape=shape, dtype=dtype, - allocator=self.allocator) + from arraycontext.impl.pyopencl.taggable_cl_array import TaggableCLArray + + return TaggableCLArray(self.queue, shape=shape, dtype=dtype, + allocator=self.allocator) def zeros(self, shape, dtype): import pyopencl.array as cl_array - return cl_array.zeros(self.queue, shape=shape, dtype=dtype, - allocator=self.allocator) + from arraycontext.impl.pyopencl.taggable_cl_array import to_tagged_cl_array + return to_tagged_cl_array(cl_array.zeros(self.queue, shape=shape, + dtype=dtype, + allocator=self.allocator), + axes=None, tags=frozenset()) def from_numpy(self, array: Union[np.ndarray, _ScalarLike]): import pyopencl.array as cl_array - return cl_array.to_device(self.queue, array, allocator=self.allocator) + from arraycontext.impl.pyopencl.taggable_cl_array import to_tagged_cl_array + return to_tagged_cl_array(cl_array + .to_device(self.queue, + array, + allocator=self.allocator), + axes=None, tags=frozenset()) def to_numpy(self, array): if np.isscalar(array): @@ -186,14 +201,33 @@ def call_loopy(self, t_unit, **kwargs): if len(wait_event_queue) > self._wait_event_queue_length: wait_event_queue.pop(0).wait() - return result + from arraycontext.impl.pyopencl.taggable_cl_array import to_tagged_cl_array + # FIXME: Inherit loopy tags for these arrays + return {name: to_tagged_cl_array(ary, axes=None, tags=frozenset()) + for name, ary in result.items()} def freeze(self, array): array.finish() return array.with_queue(None) def thaw(self, array): - return array.with_queue(self.queue) + from arraycontext.impl.pyopencl.taggable_cl_array import (TaggableCLArray, + to_tagged_cl_array) + import pyopencl.array as cl_array + + if isinstance(array, TaggableCLArray): + return array.with_queue(self.queue) + elif isinstance(array, cl_array.Array): + from warnings import warn + warn("Invoking PyOpenCLArrayContext.thaw with pyopencl.Array" + " will be unsupported in 2023. Use `to_tagged_cl_array`" + " to convert instances of pyopencl.Array to TaggableCLArray.", + DeprecationWarning, stacklevel=2) + return (to_tagged_cl_array(array, axes=None, tags=frozenset()) + .with_queue(self.queue)) + else: + raise ValueError("array should be a cl.array.Array," + f" got '{type(array)}'") # }}} @@ -268,12 +302,37 @@ def transform_loopy_program(self, t_unit): return t_unit def tag(self, tags: Union[Sequence[Tag], Tag], array): - # Sorry, not capable. - return array + import pyopencl.array as cl_array + from arraycontext.impl.pyopencl.taggable_cl_array import (TaggableCLArray, + to_tagged_cl_array) + + def _rec_tagged(ary): + if isinstance(ary, TaggableCLArray): + return ary.tagged(tags) + elif isinstance(ary, cl_array.Array): + return to_tagged_cl_array(ary, axes=None, tags=tags) + else: + raise ValueError("array should be a cl.array.Array," + f" got '{type(ary)}'") + + return rec_map_array_container(_rec_tagged, array) def tag_axis(self, iaxis, tags: Union[Sequence[Tag], Tag], array): - # Sorry, not capable. - return array + import pyopencl.array as cl_array + from arraycontext.impl.pyopencl.taggable_cl_array import (TaggableCLArray, + to_tagged_cl_array) + + def _rec_tagged(ary): + if isinstance(ary, TaggableCLArray): + return ary.with_tagged_axis(iaxis, tags) + elif isinstance(ary, cl_array.Array): + return (to_tagged_cl_array(ary, axes=None, tags=tags) + .with_tagged_axis(iaxis, tags)) + else: + raise ValueError("array should be a cl.array.Array," + f" got '{type(ary)}'") + + return rec_map_array_container(_rec_tagged, array) def clone(self): return type(self)(self.queue, self.allocator, diff --git a/arraycontext/impl/pyopencl/taggable_cl_array.py b/arraycontext/impl/pyopencl/taggable_cl_array.py new file mode 100644 index 00000000..ecde87d1 --- /dev/null +++ b/arraycontext/impl/pyopencl/taggable_cl_array.py @@ -0,0 +1,129 @@ +""" +.. autoclass:: TaggableCLArray +.. autoclass:: Axis + +.. autofunction:: to_tagged_cl_array +""" + +import pyopencl.array as cla +from typing import FrozenSet, Union, Sequence, Optional, Tuple +from pytools.tag import Taggable, Tag +from dataclasses import dataclass +from pytools import memoize + + +@dataclass(frozen=True, eq=True) +class Axis(Taggable): + """ + Records the tags corresponding to a dimensions of :class:`TaggableCLArray`. + """ + tags: FrozenSet[Tag] + + def copy(self, **kwargs): + from dataclasses import replace + return replace(self, **kwargs) + + +@memoize +def _construct_untagged_axes(ndim: int) -> Tuple[Axis, ...]: + return tuple(Axis(frozenset()) for _ in range(ndim)) + + +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) + + self.tags = tags + axes = axes if axes is not None else _construct_untagged_axes(len(self + .shape)) + self.axes = axes + + def copy(self, queue=cla._copy_queue, tags=None, axes=None, _new_class=None): + """ + :arg _new_class: The class of the copy. :func:`to_tagged_cl_array` is + sets this to convert instances of :class:`pyopencl.array.Array` to + :class:`TaggableCLArray`. If not provided, defaults to + ``self.__class__``. + """ + _new_class = self.__class__ if _new_class is None else _new_class + + if queue is not cla._copy_queue: + # Copying command queue is an involved operation, use super-class' + # implementation. + base_instance = super().copy(queue=queue) + else: + base_instance = self + + if tags is None and axes is None and _new_class is self.__class__: + # early exit + return base_instance + + tags = getattr(base_instance, "tags", frozenset()) if tags is None else tags + axes = getattr(base_instance, "axes", None) if axes is None else axes + + return _new_class(None, + base_instance.shape, + base_instance.dtype, + allocator=base_instance.allocator, + strides=base_instance.strides, + data=base_instance.base_data, + offset=base_instance.offset, + events=base_instance.events, _fast=True, + _context=base_instance.context, + _queue=base_instance.queue, + _size=base_instance.size, + tags=tags, + axes=axes, + ) + + def with_tagged_axis(self, iaxis: int, + tags: Union[Sequence[Tag], Tag]) -> "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 self.copy(axes=new_axes) + + +def to_tagged_cl_array(ary: cla.Array, + axes: Optional[Tuple[Axis, ...]], + tags: FrozenSet[Tag]) -> TaggableCLArray: + """ + Returns a :class:`TaggableCLArray` that is constructed from the data in + *ary* along with the metadata from *axes* and *tags*. + + :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. + """ + return TaggableCLArray.copy(ary, axes=axes, tags=tags, + _new_class=TaggableCLArray) diff --git a/arraycontext/impl/pytato/utils.py b/arraycontext/impl/pytato/utils.py index f69a1961..f14d166e 100644 --- a/arraycontext/impl/pytato/utils.py +++ b/arraycontext/impl/pytato/utils.py @@ -24,10 +24,11 @@ from typing import Any, Dict, Set, Tuple, Mapping -from pytato.array import SizeParam, Placeholder, make_placeholder +from pytato.array import SizeParam, Placeholder, make_placeholder, Axis as PtAxis from pytato.array import Array, DataWrapper, DictOfNamedArrays from pytato.transform import CopyMapper from pytools import UniqueNameGenerator +from arraycontext.impl.pyopencl.taggable_cl_array import Axis as ClAxis class _DatawrapperToBoundPlaceholderMapper(CopyMapper): @@ -81,3 +82,11 @@ def _normalize_pt_expr(expr: DictOfNamedArrays) -> Tuple[DictOfNamedArrays, normalize_mapper = _DatawrapperToBoundPlaceholderMapper() normalized_expr = normalize_mapper(expr) return normalized_expr, normalize_mapper.bound_arguments + + +def get_pt_axes_from_cl_axes(axes: Tuple[ClAxis, ...]) -> Tuple[PtAxis, ...]: + return tuple(PtAxis(axis.tags) for axis in axes) + + +def get_cl_axes_from_pt_axes(axes: Tuple[PtAxis, ...]) -> Tuple[ClAxis, ...]: + return tuple(ClAxis(axis.tags) for axis in axes) From 21c38824bf49dda012e8a62bdcc6b7fcb6e1ab66 Mon Sep 17 00:00:00 2001 From: Kaushik Kulkarni Date: Tue, 19 Oct 2021 09:26:34 -0500 Subject: [PATCH 4/6] make the frozen type of PytatoPyOpenCLArrayContext to be TaggableCLArrays --- arraycontext/impl/pytato/__init__.py | 108 +++++++++++++++++++++------ arraycontext/impl/pytato/compile.py | 56 +++++++++++--- arraycontext/impl/pytato/utils.py | 1 + 3 files changed, 134 insertions(+), 31 deletions(-) diff --git a/arraycontext/impl/pytato/__init__.py b/arraycontext/impl/pytato/__init__.py index f7592520..730133e2 100644 --- a/arraycontext/impl/pytato/__init__.py +++ b/arraycontext/impl/pytato/__init__.py @@ -79,6 +79,7 @@ def __init__(self, queue, allocator=None): self.allocator = allocator self.array_types = (pt.Array, ) self._freeze_prg_cache = {} + self._dag_transform_cache = {} # unused, but necessary to keep the context alive self.context = self.queue.context @@ -113,24 +114,56 @@ 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.scalar_expr import SCALAR_CLASSES from pytato.loopy import call_loopy + from arraycontext.impl.pyopencl.taggable_cl_array import TaggableCLArray entrypoint = program.default_entrypoint.name - # thaw frozen arrays - kwargs = {kw: (self.thaw(arg) if isinstance(arg, cla.Array) else arg) - for kw, arg in kwargs.items()} + # {{{ preprocess args + + processed_kwargs = {} + + for kw, arg in sorted(kwargs.items()): + if isinstance(arg, self.array_types + SCALAR_CLASSES): + pass + elif isinstance(arg, TaggableCLArray): + arg = self.thaw(arg) + else: + raise ValueError(f"call_loopy argument '{kw}' expected to be an" + " instance of 'pytato.Array', 'Number' or" + f"'TaggableCLArray', got '{type(arg)}'") + + processed_kwargs[kw] = arg + + # }}} - return call_loopy(program, kwargs, entrypoint) + return call_loopy(program, processed_kwargs, entrypoint) 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, + get_cl_axes_from_pt_axes) + from arraycontext.impl.pyopencl.taggable_cl_array import (to_tagged_cl_array, + TaggableCLArray) - if isinstance(array, cla.Array): + if isinstance(array, TaggableCLArray): return array.with_queue(None) + if isinstance(array, cla.Array): + from warnings import warn + warn("Freezing pyopencl.array.Array will be deprecated in 2023." + " Use `to_tagged_cl_array` to convert the array to" + " TaggableCLArray", DeprecationWarning, stacklevel=2) + return to_tagged_cl_array(array.with_queue(None), + axes=None, + tags=frozenset()) + if isinstance(array, pt.DataWrapper): + # trivial freeze. + return to_tagged_cl_array(array.data.with_queue(None), + axes=get_cl_axes_from_pt_axes(array.axes), + tags=array.tags) if not isinstance(array, pt.Array): raise TypeError("PytatoPyOpenCLArrayContext.freeze invoked with " f"non-pytato array of type '{type(array)}'") @@ -138,14 +171,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), + get_cl_axes_from_pt_axes(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}) @@ -155,7 +190,13 @@ def freeze(self, array): try: pt_prg = self._freeze_prg_cache[normalized_expr] except KeyError: - pt_prg = pt.generate_loopy(self.transform_dag(normalized_expr), + if normalized_expr in self._dag_transform_cache: + transformed_dag = self._dag_transform_cache[normalized_expr] + else: + transformed_dag = self.transform_dag(normalized_expr) + self._dag_transform_cache[normalized_expr] = transformed_dag + + pt_prg = pt.generate_loopy(transformed_dag, options=lp.Options(return_dict=True, no_numpy=True), cl_device=self.queue.device) @@ -166,17 +207,31 @@ 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), + get_cl_axes_from_pt_axes( + self._dag_transform_cache[normalized_expr]["_actx_out"].expr.axes), + array.tags) def thaw(self, array): import pytato as pt - import pyopencl.array as cla - - if not isinstance(array, cla.Array): - raise TypeError("PytatoPyOpenCLArrayContext.thaw expects CL arrays, got " - f"{type(array)}") - - return pt.make_data_wrapper(array.with_queue(self.queue)) + from .utils import get_pt_axes_from_cl_axes + from arraycontext.impl.pyopencl.taggable_cl_array import (TaggableCLArray, + to_tagged_cl_array) + import pyopencl.array as cl_array + + if isinstance(array, TaggableCLArray): + pass + elif isinstance(array, cl_array.Array): + array = to_tagged_cl_array(array, axes=None, tags=frozenset()) + else: + raise TypeError("PytatoPyOpenCLArrayContext.thaw expects " + "'TaggableCLArray' or 'cl.array.Array' got " + f"{type(array)}.") + + return pt.make_data_wrapper(array.with_queue(self.queue), + axes=get_pt_axes_from_cl_axes(array.axes), + tags=array.tags) # }}} @@ -219,12 +274,23 @@ def tag_axis(self, iaxis, tags: Union[Sequence[Tag], Tag], array): def einsum(self, spec, *args, arg_names=None, tagged=()): import pyopencl.array as cla import pytato as pt + from arraycontext.impl.pyopencl.taggable_cl_array import (TaggableCLArray, + to_tagged_cl_array) if arg_names is None: arg_names = (None,) * len(args) def preprocess_arg(name, arg): - if isinstance(arg, cla.Array): + if isinstance(arg, TaggableCLArray): ary = self.thaw(arg) + elif isinstance(arg, cla.Array): + from warnings import warn + warn("Passing pyopencl.array.Array to einsum will be " + "deprecated in 2023." + " Use `to_tagged_cl_array` to convert the array to" + " TaggableCLArray.", DeprecationWarning, stacklevel=2) + ary = self.thaw(to_tagged_cl_array(arg, + axes=None, + tags=frozenset())) else: assert isinstance(arg, pt.Array) ary = arg diff --git a/arraycontext/impl/pytato/compile.py b/arraycontext/impl/pytato/compile.py index 71f98a8c..d83e376a 100644 --- a/arraycontext/impl/pytato/compile.py +++ b/arraycontext/impl/pytato/compile.py @@ -34,7 +34,7 @@ import abc import numpy as np -from typing import Any, Callable, Tuple, Dict, Mapping +from typing import Any, Callable, Tuple, Dict, Mapping, FrozenSet from dataclasses import dataclass, field from pyrsistent import pmap, PMap @@ -169,7 +169,11 @@ def _get_f_placeholder_args(arg, kw, arg_id_to_name): elif is_array_container_type(arg.__class__): def _rec_to_placeholder(keys, ary): name = arg_id_to_name[(kw,) + keys] - return pt.make_placeholder(name, ary.shape, ary.dtype) + return pt.make_placeholder(name, + ary.shape, + ary.dtype, + axes=ary.axes, + tags=ary.tags) return rec_keyed_map_array_container(_rec_to_placeholder, arg) else: @@ -204,6 +208,13 @@ def _dag_to_transformed_loopy_prg(self, dict_of_named_arrays): with ProcessLogger(logger, "transform_dag"): pt_dict_of_named_arrays = self.actx.transform_dag(dict_of_named_arrays) + name_in_program_to_tags = { + name: out.tags + for name, out in pt_dict_of_named_arrays._data.items()} + name_in_program_to_axes = { + name: out.axes + for name, out in pt_dict_of_named_arrays._data.items()} + with ProcessLogger(logger, "generate_loopy"): pytato_program = pt.generate_loopy(pt_dict_of_named_arrays, options=lp.Options( @@ -225,7 +236,7 @@ def _dag_to_transformed_loopy_prg(self, dict_of_named_arrays): .actx .transform_loopy_program)) - return pytato_program + return pytato_program, name_in_program_to_tags, name_in_program_to_axes def _dag_to_compiled_func(self, ary_or_dict_of_named_arrays, input_id_to_name_in_program, output_id_to_name_in_program, @@ -234,18 +245,23 @@ def _dag_to_compiled_func(self, ary_or_dict_of_named_arrays, output_id = "_pt_out" dict_of_named_arrays = pt.make_dict_of_named_arrays( {output_id: ary_or_dict_of_named_arrays}) - pytato_program = self._dag_to_transformed_loopy_prg(dict_of_named_arrays) + pytato_program, name_in_program_to_tags, name_in_program_to_axes = ( + self._dag_to_transformed_loopy_prg(dict_of_named_arrays)) return CompiledFunctionReturningArray( self.actx, pytato_program, input_id_to_name_in_program=input_id_to_name_in_program, - output_name_in_program=output_id) + output_tags=name_in_program_to_tags[output_id], + output_axes=name_in_program_to_axes[output_id], + output_name=output_id) elif isinstance(ary_or_dict_of_named_arrays, pt.DictOfNamedArrays): - pytato_program = self._dag_to_transformed_loopy_prg( - ary_or_dict_of_named_arrays) + pytato_program, name_in_program_to_tags, name_in_program_to_axes = ( + self._dag_to_transformed_loopy_prg(ary_or_dict_of_named_arrays)) return CompiledFunctionReturningArrayContainer( self.actx, pytato_program, input_id_to_name_in_program=input_id_to_name_in_program, output_id_to_name_in_program=output_id_to_name_in_program, + name_in_program_to_tags=name_in_program_to_tags, + name_in_program_to_axes=name_in_program_to_axes, output_template=output_template) else: raise NotImplementedError(type(ary_or_dict_of_named_arrays)) @@ -312,6 +328,8 @@ def _as_dict_of_named_arrays(keys, ary): def _args_to_cl_buffers(actx, input_id_to_name_in_program, arg_id_to_arg): + from arraycontext.impl.pyopencl.taggable_cl_array import TaggableCLArray + input_kwargs_for_loopy = {} for arg_id, arg in arg_id_to_arg.items(): @@ -320,7 +338,7 @@ def _args_to_cl_buffers(actx, input_id_to_name_in_program, arg_id_to_arg): 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): @@ -383,9 +401,14 @@ class CompiledFunctionReturningArrayContainer(CompiledFunction): pytato_program: pt.target.BoundProgram input_id_to_name_in_program: Mapping[Tuple[Any, ...], str] output_id_to_name_in_program: Mapping[Tuple[Any, ...], str] + name_in_program_to_tags: Mapping[str, FrozenSet[Tag]] + name_in_program_to_axes: Mapping[str, Tuple[pt.Axis, ...]] output_template: ArrayContainer def __call__(self, arg_id_to_arg) -> ArrayContainer: + from arraycontext.impl.pyopencl.taggable_cl_array import to_tagged_cl_array + from .utils import get_cl_axes_from_pt_axes + input_kwargs_for_loopy = _args_to_cl_buffers( self.actx, self.input_id_to_name_in_program, arg_id_to_arg) @@ -399,7 +422,12 @@ def __call__(self, arg_id_to_arg) -> ArrayContainer: evt.wait() def to_output_template(keys, _): - return self.actx.thaw(out_dict[self.output_id_to_name_in_program[keys]]) + name_in_program = self.output_id_to_name_in_program[keys] + return self.actx.thaw(to_tagged_cl_array( + out_dict[name_in_program], + axes=get_cl_axes_from_pt_axes( + self.name_in_program_to_axes[name_in_program]), + tags=self.name_in_program_to_tags[name_in_program])) return rec_keyed_map_array_container(to_output_template, self.output_template) @@ -415,9 +443,14 @@ class CompiledFunctionReturningArray(CompiledFunction): actx: PytatoPyOpenCLArrayContext pytato_program: pt.target.BoundProgram input_id_to_name_in_program: Mapping[Tuple[Any, ...], str] + output_tags: FrozenSet[Tag] + output_axes: Tuple[pt.Axis, ...] output_name: str def __call__(self, arg_id_to_arg) -> ArrayContainer: + from arraycontext.impl.pyopencl.taggable_cl_array import to_tagged_cl_array + from .utils import get_cl_axes_from_pt_axes + input_kwargs_for_loopy = _args_to_cl_buffers( self.actx, self.input_id_to_name_in_program, arg_id_to_arg) @@ -430,4 +463,7 @@ def __call__(self, arg_id_to_arg) -> ArrayContainer: # running out of memory. This mitigates that risk a bit, for now. evt.wait() - return self.actx.thaw(out_dict[self.output_name]) + return self.actx.thaw(to_tagged_cl_array(out_dict[self.output_name], + axes=get_cl_axes_from_pt_axes( + self.output_axes), + tags=self.output_tags)) diff --git a/arraycontext/impl/pytato/utils.py b/arraycontext/impl/pytato/utils.py index f14d166e..2babd559 100644 --- a/arraycontext/impl/pytato/utils.py +++ b/arraycontext/impl/pytato/utils.py @@ -58,6 +58,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 cf58281d36273d6ddb2bd127d6ce242217bb8481 Mon Sep 17 00:00:00 2001 From: Kaushik Kulkarni Date: Tue, 19 Oct 2021 09:28:04 -0500 Subject: [PATCH 5/6] adds tests specific to PytatoArrayContexts --- test/test_pytato_arraycontext.py | 106 +++++++++++++++++++++++++++++++ 1 file changed, 106 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..b71f7955 --- /dev/null +++ b/test/test_pytato_arraycontext.py @@ -0,0 +1,106 @@ +""" 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 + """ + + +class BarTag(Tag): + """ + Bar + """ + + +class BazTag(Tag): + """ + Baz + """ + +# }}} + + +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()) + .with_tagged_axis(0, BarTag()) + .with_tagged_axis(1, BazTag()), + actx), + actx) + assert foo.tags_of_type(FooTag) + assert foo.axes[0].tags_of_type(BarTag) + assert foo.axes[1].tags_of_type(BazTag) + + +if __name__ == "__main__": + import sys + if len(sys.argv) > 1: + exec(sys.argv[1]) + else: + from pytest import main + main([__file__]) + +# vim: fdm=marker From b8631d3f2591dfdc7b098318510b76760c3c7b66 Mon Sep 17 00:00:00 2001 From: Kaushik Kulkarni Date: Sat, 12 Feb 2022 12:15:03 -0600 Subject: [PATCH 6/6] PyOpenCLArrayContext.make_einsum: use `tagged` tags for the output array as well --- arraycontext/context.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/arraycontext/context.py b/arraycontext/context.py index 0c9232d8..2bff0ac2 100644 --- a/arraycontext/context.py +++ b/arraycontext/context.py @@ -328,9 +328,10 @@ def einsum(self, spec, *args, arg_names=None, tagged=()): arg_names = tuple("arg%d" % i for i in range(len(args))) prg = self._get_einsum_prg(spec, arg_names, tagged) - return self.call_loopy( + out_ary = self.call_loopy( prg, **{arg_names[i]: arg for i, arg in enumerate(args)} )["out"] + return self.tag(tagged, out_ary) @abstractmethod def clone(self):