From 9204fb5d1b24612a5a89f79b506a00d501587974 Mon Sep 17 00:00:00 2001 From: Matt Wala Date: Tue, 2 Jun 2020 22:41:17 -0500 Subject: [PATCH 01/76] Remove scalar_expr.py --- pytato/scalar_expr.py | 46 ------------------------------------------- 1 file changed, 46 deletions(-) delete mode 100644 pytato/scalar_expr.py diff --git a/pytato/scalar_expr.py b/pytato/scalar_expr.py deleted file mode 100644 index c675984..0000000 --- a/pytato/scalar_expr.py +++ /dev/null @@ -1,46 +0,0 @@ -from __future__ import annotations - -__copyright__ = """ -Copyright (C) 2020 Andreas Kloeckner -""" - -__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 pymbolic.mapper import WalkMapper as WalkMapperBase -import pymbolic.primitives as prim - -from numbers import Number -from typing import Union - - -ScalarExpression = Union[Number, prim.Expression] - - -def parse(s: str) -> ScalarExpression: - from pymbolic.parser import Parser - return Parser()(s) - - -class WalkMapper(WalkMapperBase): - pass - - -# vim: foldmethod=marker -- GitLab From 590a79fc6586c1b65a8f6d5bfa007fa37393b5ee Mon Sep 17 00:00:00 2001 From: Matt Wala Date: Tue, 2 Jun 2020 23:13:11 -0500 Subject: [PATCH 02/76] Add, document pytato.symbolic --- doc/conf.py | 1 + doc/reference.rst | 2 + pytato/array.py | 4 +- pytato/program.py | 133 ++++++++++++++++++++++++++++++++++++++ pytato/symbolic.py | 156 +++++++++++++++++++++++++++++++++++++++++++++ 5 files changed, 294 insertions(+), 2 deletions(-) create mode 100644 pytato/program.py create mode 100644 pytato/symbolic.py diff --git a/doc/conf.py b/doc/conf.py index 7bd8b30..96130ca 100644 --- a/doc/conf.py +++ b/doc/conf.py @@ -65,4 +65,5 @@ intersphinx_mapping = { 'http://documen.tician.de/pymbolic/': None, 'http://documen.tician.de/loopy/': None, 'http://documen.tician.de/sumpy/': None, + 'http://documen.tician.de/islpy/': None, } diff --git a/doc/reference.rst b/doc/reference.rst index 522bdd4..c743d58 100644 --- a/doc/reference.rst +++ b/doc/reference.rst @@ -2,3 +2,5 @@ Reference ========= .. automodule:: pytato.array +.. automodule:: pytato.program +.. automodule:: pytato.symbolic diff --git a/pytato/array.py b/pytato/array.py index 046955f..fae2030 100644 --- a/pytato/array.py +++ b/pytato/array.py @@ -87,8 +87,8 @@ Node constructors such as :class:`Placeholder.__init__` and import numpy as np import pymbolic.primitives as prim -import pytato.scalar_expr as scalar_expr -from pytato.scalar_expr import ScalarExpression +import pytato.symbolic as scalar_expr +from pytato.symbolic import ScalarExpression from dataclasses import dataclass from pytools import is_single_valued diff --git a/pytato/program.py b/pytato/program.py new file mode 100644 index 0000000..a3aea61 --- /dev/null +++ b/pytato/program.py @@ -0,0 +1,133 @@ +from __future__ import annotations + +__copyright__ = """Copyright (C) 2020 Matt Wala""" + +__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. +""" + +__doc__ = """ + +.. currentmodule:: pytato.program + +Code Generation Targets +----------------------- + +.. autoclass:: Target +.. autoclass:: PyOpenCLTarget + +Generated Executable Programs +----------------------------- + +.. autoclass:: BoundProgram +.. autoclass:: BoundPyOpenCLProgram +""" + +from dataclasses import dataclass +import typing +from typing import Any, Mapping, Optional + +if typing.TYPE_CHECKING: + # Skip imports for efficiency. FIXME: Neither of these work as type stubs + # are not present. Types are here only as documentation. + import pyopencl as cl + import loopy as lp + + +class Target: + """An abstract code generation target.""" + def get_loopy_target(self) -> "lp.TargetBase": + """Return the corresponding :mod:`loopy` target.""" + raise NotImplementedError + + def bind_program(self, program: "lp.LoopKernel", + bound_arguments: Mapping[str, Any]) -> BoundProgram: + """Create a :class:`BoundProgram` for this code generation target. + + :arg program: the :mod:`loopy` kernel + :arg bound_arguments: a mapping from argument names to outputs + """ + raise NotImplementedError + + +class PyOpenCLTarget(Target): + """A :mod:`pyopencl` code generation target.""" + def __init__(self, queue: Optional["cl.CommandQueue"] = None): + self.queue = queue + + def get_loopy_target(self) -> "lp.PyOpenCLTarget": + import loopy as lp + device = None + if self.queue is not None: + device = self.queue.device + return lp.PyOpenCLTarget(device) + + def bind_program(self, program: "lp.LoopKernel", + bound_arguments: Mapping[str, Any]) -> BoundProgram: + return BoundPyOpenCLProgram(program=program, + queue=self.queue, + bound_arguments=bound_arguments, + target=self) + + +@dataclass(init=True, repr=False, eq=False) +class BoundProgram: + """A wrapper around a :mod:`loopy` kernel for execution. + + .. attribute:: program + + The underlying :class:`loopy.LoopKernel`. + + .. attribute:: target + + The code generation target. + + .. attribute:: bound_arguments + + A map from names to pre-bound kernel arguments. + """ + + program: "lp.LoopKernel" + bound_arguments: Mapping[str, Any] + target: Target + + def __call__(self, *args: Any, **kwargs: Any) -> Any: + raise NotImplementedError + + +@dataclass(init=True, repr=False, eq=False) +class BoundPyOpenCLProgram(BoundProgram): + """A wrapper around a :mod:`loopy` kernel for execution with :mod`pyopencl`. + + .. attribute:: queue + + A :mod:`pyopencl` command queue. + """ + queue: Optional["cl.CommandQueue"] + + def __call__(self, *args: Any, **kwargs: Any) -> Any: + """Convenience function for launching a :mod:`pyopencl` computation.""" + if not self.queue: + raise ValueError("queue must be specified") + + updated_kwargs = dict(self.bound_arguments) + updated_kwargs.update(kwargs) + return self.program(self.queue, *args, **updated_kwargs) + +# vim: foldmethod=marker diff --git a/pytato/symbolic.py b/pytato/symbolic.py new file mode 100644 index 0000000..a6cee3d --- /dev/null +++ b/pytato/symbolic.py @@ -0,0 +1,156 @@ +from __future__ import annotations + +__copyright__ = """ +Copyright (C) 2020 Andreas Kloeckner +""" + +__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 numbers import Number +from typing import Any, Union, Mapping, FrozenSet, Set, Tuple + +import islpy as isl +from pymbolic.mapper import ( + WalkMapper as WalkMapperBase, + IdentityMapper as IdentityMapperBase) +from pymbolic.mapper.substitutor import ( + SubstitutionMapper as SubstitutionMapperBase) +from pymbolic.mapper.dependency import ( + DependencyMapper as DependencyMapperBase) +import pymbolic.primitives as prim + + +__doc__ = """ +.. currentmodule:: pytato.symbolic + +Symbolic Infrastructure +----------------------- + +.. data:: ScalarExpression + + A type alias for ``Union[Number, pymbolic.primitives.Expression]``. + +.. autofunction:: parse +.. autofunction:: get_dependencies +.. autofunction:: substitute +.. autofunction:: domain_for_shape + +""" + + +# {{{ scalar expressions + +ScalarExpression = Union[Number, prim.Expression] + + +def parse(s: str) -> ScalarExpression: + from pymbolic.parser import Parser + return Parser()(s) + +# }}} + + +# {{{ mapper classes + +class WalkMapper(WalkMapperBase): + pass + + +class IdentityMapper(IdentityMapperBase): + pass + + +class SubstitutionMapper(SubstitutionMapperBase): + pass + + +class DependencyMapper(DependencyMapperBase): + pass + +# }}} + + +# {{{ mapper frontends + +def get_dependencies(expression: Any) -> FrozenSet[str]: + """Return the set of variable names in an expression. + + :param expression: A :mod:`pymbolic` expression + """ + mapper = DependencyMapper(composite_leaves=False) + return frozenset(dep.name for dep in mapper(expression)) + + +def substitute(expression: Any, variable_assigments: Mapping[str, Any]) -> Any: + """Perform variable substitution in an expression. + + :param expression: A :mod:`pymbolic` expression + :param variable_assigments: A mapping from variable names to substitutions + """ + from pymbolic.mapper.substitutor import make_subst_func + return SubstitutionMapper(make_subst_func(variable_assigments))(expression) + +# }}} + + +def domain_for_shape(dim_names: Tuple[str, ...], shape: Tuple[ScalarExpression, ...]) -> isl.BasicSet: + """Create a :class:`islpy.BasicSet` that expresses an appropriate index domain + for an array of (potentially symbolic) shape *shape*. + + :param dim_names: A tuple of strings, the names of the axes. These become set + dimensions in the returned domain. + + :param shape: A tuple of constant or quasi-affine :mod:`pymbolic` + expressions. The variables in these expressions become parameter + dimensions in the returned set. Must have the same length as + *dim_names*. + """ + assert len(dim_names) == len(shape) + + # Collect parameters. + param_names_set: Set[str] = set() + for sdep in map(get_dependencies, shape): + param_names_set |= sdep + + set_names = sorted(dim_names) + param_names = sorted(param_names_set) + + # Build domain. + dom = isl.BasicSet.universe( + isl.Space.create_from_names( + isl.DEFAULT_CONTEXT, + set=set_names, + params=param_names)) + + # Add constraints. + from loopy.symbolic import aff_from_expr + affs = isl.affs_from_space(dom.space) + + for iname, dim in zip(dim_names, shape): + dom &= affs[0].le_set(affs[iname]) + dom &= affs[iname].lt_set(aff_from_expr(dom.space, dim)) + + dom, = dom.get_basic_sets() + + return dom + + +# vim: foldmethod=marker -- GitLab From f42fac62b3ab79ff71da9b8c50bc2f555a95580e Mon Sep 17 00:00:00 2001 From: Matt Wala Date: Tue, 2 Jun 2020 23:15:37 -0500 Subject: [PATCH 03/76] Fix typo --- pytato/program.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pytato/program.py b/pytato/program.py index a3aea61..4a36602 100644 --- a/pytato/program.py +++ b/pytato/program.py @@ -113,7 +113,7 @@ class BoundProgram: @dataclass(init=True, repr=False, eq=False) class BoundPyOpenCLProgram(BoundProgram): - """A wrapper around a :mod:`loopy` kernel for execution with :mod`pyopencl`. + """A wrapper around a :mod:`loopy` kernel for execution with :mod:`pyopencl`. .. attribute:: queue -- GitLab From 9d97d44fb30ec9a5c44e7b01d2d72fa87b5fd296 Mon Sep 17 00:00:00 2001 From: Matt Wala Date: Tue, 2 Jun 2020 23:20:50 -0500 Subject: [PATCH 04/76] Linter fixes --- pytato/array.py | 10 ++++------ pytato/symbolic.py | 34 +++++++++++++++++----------------- setup.cfg | 12 +++++------- 3 files changed, 26 insertions(+), 30 deletions(-) diff --git a/pytato/array.py b/pytato/array.py index fae2030..061bf82 100644 --- a/pytato/array.py +++ b/pytato/array.py @@ -87,7 +87,7 @@ Node constructors such as :class:`Placeholder.__init__` and import numpy as np import pymbolic.primitives as prim -import pytato.symbolic as scalar_expr +import pytato.symbolic as sym from pytato.symbolic import ScalarExpression from dataclasses import dataclass @@ -267,7 +267,7 @@ def _check_identifier(s: str, ns: Optional[Namespace] = None) -> bool: return True -class _ShapeChecker(scalar_expr.WalkMapper): +class _ShapeChecker(sym.WalkMapper): def __init__(self, ns: Optional[Namespace] = None): super().__init__() self.ns = ns @@ -285,12 +285,10 @@ def normalize_shape( :param ns: if a namespace is given, extra checks are performed to ensure that expressions are well-defined. """ - from pytato.scalar_expr import parse - def normalize_shape_component( s: ConvertibleToShapeComponent) -> ScalarExpression: if isinstance(s, str): - s = parse(s) + s = sym.parse(s) if isinstance(s, int): if s < 0: @@ -303,7 +301,7 @@ def normalize_shape( return s if isinstance(shape, str): - shape = parse(shape) + shape = sym.parse(shape) from numbers import Number if isinstance(shape, (Number, prim.Expression)): diff --git a/pytato/symbolic.py b/pytato/symbolic.py index a6cee3d..8a6799c 100644 --- a/pytato/symbolic.py +++ b/pytato/symbolic.py @@ -28,16 +28,14 @@ from numbers import Number from typing import Any, Union, Mapping, FrozenSet, Set, Tuple import islpy as isl -from pymbolic.mapper import ( - WalkMapper as WalkMapperBase, - IdentityMapper as IdentityMapperBase) -from pymbolic.mapper.substitutor import ( - SubstitutionMapper as SubstitutionMapperBase) -from pymbolic.mapper.dependency import ( - DependencyMapper as DependencyMapperBase) +from pymbolic.mapper import (WalkMapper as WalkMapperBase, IdentityMapper as + IdentityMapperBase) +from pymbolic.mapper.substitutor import (SubstitutionMapper as + SubstitutionMapperBase) +from pymbolic.mapper.dependency import (DependencyMapper as + DependencyMapperBase) import pymbolic.primitives as prim - __doc__ = """ .. currentmodule:: pytato.symbolic @@ -55,7 +53,6 @@ Symbolic Infrastructure """ - # {{{ scalar expressions ScalarExpression = Union[Number, prim.Expression] @@ -65,11 +62,12 @@ def parse(s: str) -> ScalarExpression: from pymbolic.parser import Parser return Parser()(s) -# }}} +# }}} # {{{ mapper classes + class WalkMapper(WalkMapperBase): pass @@ -85,11 +83,12 @@ class SubstitutionMapper(SubstitutionMapperBase): class DependencyMapper(DependencyMapperBase): pass -# }}} +# }}} # {{{ mapper frontends + def get_dependencies(expression: Any) -> FrozenSet[str]: """Return the set of variable names in an expression. @@ -108,10 +107,12 @@ def substitute(expression: Any, variable_assigments: Mapping[str, Any]) -> Any: from pymbolic.mapper.substitutor import make_subst_func return SubstitutionMapper(make_subst_func(variable_assigments))(expression) + # }}} -def domain_for_shape(dim_names: Tuple[str, ...], shape: Tuple[ScalarExpression, ...]) -> isl.BasicSet: +def domain_for_shape(dim_names: Tuple[str, ...], shape: Tuple[ScalarExpression, + ...]) -> isl.BasicSet: """Create a :class:`islpy.BasicSet` that expresses an appropriate index domain for an array of (potentially symbolic) shape *shape*. @@ -124,7 +125,7 @@ def domain_for_shape(dim_names: Tuple[str, ...], shape: Tuple[ScalarExpression, *dim_names*. """ assert len(dim_names) == len(shape) - + # Collect parameters. param_names_set: Set[str] = set() for sdep in map(get_dependencies, shape): @@ -135,10 +136,9 @@ def domain_for_shape(dim_names: Tuple[str, ...], shape: Tuple[ScalarExpression, # Build domain. dom = isl.BasicSet.universe( - isl.Space.create_from_names( - isl.DEFAULT_CONTEXT, - set=set_names, - params=param_names)) + isl.Space.create_from_names(isl.DEFAULT_CONTEXT, + set=set_names, + params=param_names)) # Add constraints. from loopy.symbolic import aff_from_expr diff --git a/setup.cfg b/setup.cfg index b975cd3..e071c15 100644 --- a/setup.cfg +++ b/setup.cfg @@ -2,23 +2,21 @@ ignore = E126,E127,E128,E123,E226,E241,E242,E265,N802,W503,E402,N814,N817,W504 max-line-length=85 -[mypy-pytato.scalar_expr] +[mypy-pytato.symbolic] disallow_subclassing_any = False -[mypy-pymbolic] +[mypy-islpy] ignore_missing_imports = True -[mypy-pymbolic.primitives] +[mypy-loopy.*] ignore_missing_imports = True -[mypy-pymbolic.mapper] +[mypy-numpy] ignore_missing_imports = True -[mypy-pymbolic.parser] +[mypy-pymbolic.*] ignore_missing_imports = True [mypy-pytools] ignore_missing_imports = True -[mypy-numpy] -ignore_missing_imports = True -- GitLab From f5e1450d17ebb8e89564b98c518f5119b06b7ebc Mon Sep 17 00:00:00 2001 From: Matt Wala Date: Tue, 2 Jun 2020 23:22:55 -0500 Subject: [PATCH 05/76] Mypy fix --- setup.cfg | 3 +++ 1 file changed, 3 insertions(+) diff --git a/setup.cfg b/setup.cfg index e071c15..4d4fdb8 100644 --- a/setup.cfg +++ b/setup.cfg @@ -17,6 +17,9 @@ ignore_missing_imports = True [mypy-pymbolic.*] ignore_missing_imports = True +[mypy-pyopencl] +ignore_missing_imports = True + [mypy-pytools] ignore_missing_imports = True -- GitLab From 8122a2ac89fadae658103c989ff694bbc4e5c95c Mon Sep 17 00:00:00 2001 From: Matt Wala Date: Wed, 3 Jun 2020 00:23:50 -0500 Subject: [PATCH 06/76] Rename symbolic.py to scalar_expr.py --- pytato/array.py | 316 +++++++++++++++++++------ pytato/{symbolic.py => scalar_expr.py} | 19 +- 2 files changed, 255 insertions(+), 80 deletions(-) rename pytato/{symbolic.py => scalar_expr.py} (92%) diff --git a/pytato/array.py b/pytato/array.py index 8429b75..5c2579c 100644 --- a/pytato/array.py +++ b/pytato/array.py @@ -67,6 +67,7 @@ Built-in Expression Nodes .. autoclass:: Reshape .. autoclass:: DataWrapper .. autoclass:: Placeholder +.. autoclass:: Output .. autoclass:: LoopyFunction User-Facing Node Creation @@ -85,15 +86,20 @@ Node constructors such as :class:`Placeholder.__init__` and # }}} +import collections +from functools import partialmethod +from numbers import Number +import operator +from dataclasses import dataclass +from typing import Optional, Dict, Any, MutableMapping, Mapping, Iterator, Tuple, Union, FrozenSet + import numpy as np +import pymbolic.mapper import pymbolic.primitives as prim -import pytato.symbolic as sym -from pytato.symbolic import ScalarExpression - -from dataclasses import dataclass -from pytools import is_single_valued -from typing import Optional, Dict, Any, Mapping, Iterator, Tuple, Union, FrozenSet +from pytools import is_single_valued, memoize_method +import pytato.scalar_expr as scalar_expr +from pytato.scalar_expr import ScalarExpression # {{{ dotted name @@ -138,7 +144,37 @@ class DottedName: # {{{ namespace -class Namespace: +class _NamespaceCopyMapper(sym.IdentityMapper): + + def __call__(self, expr: Array, namespace: Namespace, cache: Dict[Array, Array]) -> Array: + return self.rec(expr, namespace, cache) + + def rec(self, expr: Array, namespace: Namespace, cache: Dict[Array, Array]) -> Array: + if expr in cache: + return cache[expr] + result: Array = super().rec(expr, namespace, cache) + cache[expr] = result + return result + + def map_index_lambda(self, expr: IndexLambda, namespace: Namespace, cache: Dict[Array, Array]) -> Array: + bindings = { + name: self.rec(subexpr, namespace, cache) + for name, subexpr in expr.bindings.items()} + return IndexLambda( + namespace, + expr=expr.expr, + shape=expr.shape, + dtype=expr.dtype, + bindings=bindings) + + def map_placeholder(self, expr: Placeholder, namespace: Namespace, cache: Dict[Array, Array]) -> Array: + return Placeholder(namespace, expr.name, expr.shape, expr.dtype, expr.tags) + + def map_output(self, expr: Output, namespace: Namespace, cache: Dict[Array, Array]) -> Array: + return Output(namespace, expr.name, self.rec(expr.array, namespace, cache), expr.tags) + + +class Namespace(Mapping[str, "Array"]): # Possible future extension: .parent attribute r""" Represents a mapping from :term:`identifier` strings to @@ -149,14 +185,18 @@ class Namespace: .. automethod:: __contains__ .. automethod:: __getitem__ .. automethod:: __iter__ + .. automethod:: __len__ .. automethod:: assign + .. automethod:: copy .. automethod:: ref """ - def __init__(self) -> None: - self._symbol_table: Dict[str, Optional[Array]] = {} + def __init__(self, _symbol_table: Optional[MutableMapping[str, Array]] = None) -> None: + if _symbol_table is None: + _symbol_table = {} + self._symbol_table: MutableMapping[str, Array] = _symbol_table - def __contains__(self, name: str) -> bool: + def __contains__(self, name: object) -> bool: return name in self._symbol_table def __getitem__(self, name: str) -> Array: @@ -167,14 +207,28 @@ class Namespace: def __iter__(self) -> Iterator[str]: return iter(self._symbol_table) - - def assign(self, name: str, - value: Optional[Array]) -> str: + + def __len__(self) -> int: + return len(self._symbol_table) + + def _chain(self) -> Namespace: + return Namespace(collections.ChainMap(dict(), self._symbol_table)) + + def copy(self) -> Namespace: + result = Namespace() + mapper = _NamespaceCopyMapper() + cache: Dict[Array, Array] = {} + for name in self: + val = mapper(self[name], result, cache) + if name not in result: + result.assign(name, val) + return result + + def assign(self, name: str, value: Array) -> str: """Declare a new array. :param name: a Python identifier - :param value: the array object, or None if the assignment is to - just reserve a name + :param value: the array object :returns: *name* """ @@ -239,6 +293,7 @@ class UniqueTag(Tag): Only one instance of this type of tag may be assigned to a single tagged object. """ + pass TagsType = FrozenSet[Tag] @@ -288,7 +343,7 @@ def normalize_shape( def normalize_shape_component( s: ConvertibleToShapeComponent) -> ScalarExpression: if isinstance(s, str): - s = sym.parse(s) + s = scalar_expr.parse(s) if isinstance(s, int): if s < 0: @@ -301,7 +356,7 @@ def normalize_shape( return s if isinstance(shape, str): - shape = sym.parse(shape) + shape = scalar_expr.parse(shape) from numbers import Number if isinstance(shape, (Number, prim.Expression)): @@ -374,21 +429,25 @@ class Array: """ - def __init__(self, namespace: Namespace, - tags: Optional[TagsType] = None): + def __init__(self, namespace: Namespace, shape: ShapeType, dtype: np.dtype, tags: Optional[TagsType] = None): if tags is None: tags = frozenset() self.namespace = namespace self.tags = tags - self.dtype: np.dtype = np.float64 # FIXME + self._shape = shape + self._dtype = dtype def copy(self, **kwargs: Any) -> Array: raise NotImplementedError @property def shape(self) -> ShapeType: - raise NotImplementedError + return self._shape + + @property + def dtype(self) -> np.dtype: + return self._dtype def named(self, name: str) -> Array: return self.namespace.ref(self.namespace.assign(name, self)) @@ -416,8 +475,56 @@ class Array: return self.copy(tags=new_tags) - # TODO: - # - codegen interface + @memoize_method + def __hash__(self) -> int: + raise NotImplementedError + + def __eq__(self, other: Any) -> bool: + raise NotImplementedError + + def __ne__(self, other: Any) -> bool: + return not self.__eq__(other) + + def _join_dtypes(self, *args: np.dtype) -> np.dtype: + result = args[0] + for arg in args[1:]: + result = (np.empty(0, dtype=result) + np.empty(0, dtype=arg)).dtype + return result + + def _binary_op(self, op: Any, + other: Union[Array, Number], + reverse: bool = False) -> Array: + if isinstance(other, Number): + # TODO + raise NotImplementedError + else: + if self.shape != other.shape: + raise ValueError("shapes do not match for binary operator") + + dtype = self._join_dtypes(self.dtype, other.dtype) + + # FIXME: If either *self* or *other* is an IndexLambda, its expression + # could be folded into the output, producing a fused result. + if self.shape == (): + expr = op(prim.Variable("_in0"), prim.Variable("_in1")) + else: + indices = tuple(prim.Variable(f"_{i}") for i in range(self.ndim)) + expr = op( + prim.Variable("_in0")[indices], + prim.Variable("_in1")[indices]) + + first, second = self, other + if reverse: + first, second = second, first + + bindings = dict(_in0=first, _in1=second) + + return IndexLambda( + self.namespace, expr, + shape=self.shape, dtype=dtype, bindings=bindings) + + __mul__ = partialmethod(_binary_op, operator.mul) + __rmul__ = partialmethod(_binary_op, operator.mul, reverse=True) # }}} @@ -536,32 +643,26 @@ class IndexLambda(Array): .. automethod:: is_reference """ - # TODO: write make_index_lambda() that does dtype inference + mapper_method = "map_index_lambda" def __init__( - self, namespace: Namespace, expr: prim.Expression, - shape: ShapeType, dtype: np.dtype, + self, + namespace: Namespace, + expr: prim.Expression, + shape: ShapeType, + dtype: np.dtype, bindings: Optional[Dict[str, Array]] = None, tags: Optional[TagsType] = None): if bindings is None: bindings = {} - super().__init__(namespace, tags=tags) + super().__init__(namespace, shape=shape, dtype=dtype, tags=tags) - self._shape = shape - self._dtype = dtype self.expr = expr self.bindings = bindings - @property - def shape(self) -> ShapeType: - return self._shape - - @property - def dtype(self) -> np.dtype: - return self._dtype - + @memoize_method def is_reference(self) -> bool: # FIXME: Do we want a specific 'reference' node to make all this # checking unnecessary? @@ -592,6 +693,28 @@ class IndexLambda(Array): return True + @memoize_method + def __hash__(self) -> int: + return hash(( + self.expr, + self.shape, + self.dtype, + frozenset(self.bindings.items()), + self.tags)) + + def __eq__(self, other: object) -> bool: + if self is other: + return True + + return ( + isinstance(other, IndexLambda) + and self.namespace is other.namespace + and self.expr == other.expr + and self.shape == other.shape + and self.dtype == other.dtype + and self.bindings == other.bindings + and self.tags == other.tags) + # }}} @@ -636,24 +759,54 @@ class DataWrapper(Array): # TODO: not really Any data def __init__(self, namespace: Namespace, data: Any, tags: Optional[TagsType] = None): - super().__init__(namespace, tags) - + super().__init__(namespace, shape=data.shape, dtype=data.dtype, tags=tags) self.data = data - @property - def shape(self) -> Any: # FIXME - return self.data.shape - - @property - def dtype(self) -> np.dtype: - return self.data.dtype - # }}} # {{{ placeholder -class Placeholder(Array): +class _ArgLike(Array): + + def __init__(self, + namespace: Namespace, + name: str, + shape: ShapeType, + dtype: np.dtype, + tags: Optional[TagsType] = None): + if name is None: + raise ValueError("Must have explicit name") + + # Reserve the name, prevent others from using it. + namespace.assign(name, self) + + super().__init__( + namespace=namespace, shape=shape, dtype=dtype, tags=tags) + + self.name = name + + @memoize_method + def __hash__(self) -> int: + return hash((self.name,)) + + def __eq__(self, other: object) -> bool: + if self is other: + return True + # Uniquely identified by name. + return ( + isinstance(other, _ArgLike) + and self.namespace is other.namespace + and self.name == other.name) + + def tagged(self, tag: Tag) -> Array: + raise ValueError("Cannot modify tags") + + def without_tag(self, tag: Tag, verify_existence: bool = True) -> Array: + raise ValueError("Cannot modify tags") + + +class Placeholder(_ArgLike): """ A named placeholder for an array whose concrete value is supplied by the user during evaluation. @@ -665,29 +818,46 @@ class Placeholder(Array): .. note:: - :attr:`name` is not a :term:`namespace name`. In fact, - it is prohibited from being one. (This has to be the case: Suppose a - :class:`Placeholder` is :meth:`~Array.tagged`, would the namespace name - refer to the tagged or the untagged version?) + Modifying :class:`Placeholder` tags is not supported after + creation. """ - def __init__(self, namespace: Namespace, - name: str, shape: ShapeType, - tags: Optional[TagsType] = None): + mapper_method = "map_placeholder" - # Reserve the name, prevent others from using it. - namespace.assign(name, None) +# }}} - super().__init__(namespace=namespace, tags=tags) - self.name = name - self._shape = shape +# {{{ output - @property - def shape(self) -> ShapeType: - # Matt added this to make Pylint happy. - # Not tied to this, open for discussion about how to implement this. - return self._shape +class Output(_ArgLike): + """A named output of the computation. + + .. attribute:: name + + The name of the output array. + + .. attribute:: array + + The :class:`Array` value that is output. + + .. note:: + + Modifying :class:`Output` tags is not supported after creation. + """ + + mapper_method = "map_output" + + def __init__(self, + namespace: Namespace, + name: str, + array: Array, + tags: Optional[TagsType] = None): + super().__init__(namespace=namespace, + name=name, + shape=array.shape, + dtype=array.dtype, + tags=tags) + self.array = array # }}} @@ -724,20 +894,26 @@ def make_dict_of_named_arrays( def make_placeholder(namespace: Namespace, name: str, shape: ConvertibleToShape, + dtype: np.dtype, tags: Optional[TagsType] = None ) -> Placeholder: """Make a :class:`Placeholder` object. - :param namespace: namespace of the placeholder array - :param shape: shape of the placeholder array - :param tags: implementation tags + :param namespace: namespace of the placeholder array + :param name: name of the placeholder array + :param shape: shape of the placeholder array + :param dtype: dtype of the placeholder array + :param tags: implementation tags """ + if name is None: + raise ValueError("Placeholder instances must have a name") + if not str.isidentifier(name): - raise ValueError(f"{name} is not a Python identifier") + raise ValueError(f"'{name}' is not a Python identifier") shape = normalize_shape(shape, namespace) - return Placeholder(namespace, name, shape, tags) + return Placeholder(namespace, name, shape, dtype, tags) # }}} diff --git a/pytato/symbolic.py b/pytato/scalar_expr.py similarity index 92% rename from pytato/symbolic.py rename to pytato/scalar_expr.py index 8a6799c..bfd948f 100644 --- a/pytato/symbolic.py +++ b/pytato/scalar_expr.py @@ -37,14 +37,17 @@ from pymbolic.mapper.dependency import (DependencyMapper as import pymbolic.primitives as prim __doc__ = """ -.. currentmodule:: pytato.symbolic +.. currentmodule:: pytato.scalar_expr -Symbolic Infrastructure ------------------------ +Scalar-Level Expressions +------------------------ .. data:: ScalarExpression - A type alias for ``Union[Number, pymbolic.primitives.Expression]``. + A :class:`type`, representing a scalar-valued symbolic expression. Such + expressions can be manipulated and composed as :mod:`pymbolic` expressions. + + Concretely, this is an alias for ``Union[Number, pymbolic.primitives.Expression]``. .. autofunction:: parse .. autofunction:: get_dependencies @@ -57,16 +60,14 @@ Symbolic Infrastructure ScalarExpression = Union[Number, prim.Expression] - def parse(s: str) -> ScalarExpression: from pymbolic.parser import Parser return Parser()(s) - # }}} -# {{{ mapper classes +# {{{ mapper classes class WalkMapper(WalkMapperBase): pass @@ -83,11 +84,10 @@ class SubstitutionMapper(SubstitutionMapperBase): class DependencyMapper(DependencyMapperBase): pass - # }}} -# {{{ mapper frontends +# {{{ mapper frontends def get_dependencies(expression: Any) -> FrozenSet[str]: """Return the set of variable names in an expression. @@ -107,7 +107,6 @@ def substitute(expression: Any, variable_assigments: Mapping[str, Any]) -> Any: from pymbolic.mapper.substitutor import make_subst_func return SubstitutionMapper(make_subst_func(variable_assigments))(expression) - # }}} -- GitLab From 3c08afda080deee8f2602e41e39fcd4a922ba20a Mon Sep 17 00:00:00 2001 From: Matt Wala Date: Wed, 3 Jun 2020 00:25:08 -0500 Subject: [PATCH 07/76] Revert "Rename symbolic.py to scalar_expr.py" This reverts commit 8122a2ac89fadae658103c989ff694bbc4e5c95c. --- pytato/array.py | 316 ++++++------------------- pytato/{scalar_expr.py => symbolic.py} | 19 +- 2 files changed, 80 insertions(+), 255 deletions(-) rename pytato/{scalar_expr.py => symbolic.py} (92%) diff --git a/pytato/array.py b/pytato/array.py index 5c2579c..8429b75 100644 --- a/pytato/array.py +++ b/pytato/array.py @@ -67,7 +67,6 @@ Built-in Expression Nodes .. autoclass:: Reshape .. autoclass:: DataWrapper .. autoclass:: Placeholder -.. autoclass:: Output .. autoclass:: LoopyFunction User-Facing Node Creation @@ -86,20 +85,15 @@ Node constructors such as :class:`Placeholder.__init__` and # }}} -import collections -from functools import partialmethod -from numbers import Number -import operator -from dataclasses import dataclass -from typing import Optional, Dict, Any, MutableMapping, Mapping, Iterator, Tuple, Union, FrozenSet - import numpy as np -import pymbolic.mapper import pymbolic.primitives as prim -from pytools import is_single_valued, memoize_method +import pytato.symbolic as sym +from pytato.symbolic import ScalarExpression + +from dataclasses import dataclass +from pytools import is_single_valued +from typing import Optional, Dict, Any, Mapping, Iterator, Tuple, Union, FrozenSet -import pytato.scalar_expr as scalar_expr -from pytato.scalar_expr import ScalarExpression # {{{ dotted name @@ -144,37 +138,7 @@ class DottedName: # {{{ namespace -class _NamespaceCopyMapper(sym.IdentityMapper): - - def __call__(self, expr: Array, namespace: Namespace, cache: Dict[Array, Array]) -> Array: - return self.rec(expr, namespace, cache) - - def rec(self, expr: Array, namespace: Namespace, cache: Dict[Array, Array]) -> Array: - if expr in cache: - return cache[expr] - result: Array = super().rec(expr, namespace, cache) - cache[expr] = result - return result - - def map_index_lambda(self, expr: IndexLambda, namespace: Namespace, cache: Dict[Array, Array]) -> Array: - bindings = { - name: self.rec(subexpr, namespace, cache) - for name, subexpr in expr.bindings.items()} - return IndexLambda( - namespace, - expr=expr.expr, - shape=expr.shape, - dtype=expr.dtype, - bindings=bindings) - - def map_placeholder(self, expr: Placeholder, namespace: Namespace, cache: Dict[Array, Array]) -> Array: - return Placeholder(namespace, expr.name, expr.shape, expr.dtype, expr.tags) - - def map_output(self, expr: Output, namespace: Namespace, cache: Dict[Array, Array]) -> Array: - return Output(namespace, expr.name, self.rec(expr.array, namespace, cache), expr.tags) - - -class Namespace(Mapping[str, "Array"]): +class Namespace: # Possible future extension: .parent attribute r""" Represents a mapping from :term:`identifier` strings to @@ -185,18 +149,14 @@ class Namespace(Mapping[str, "Array"]): .. automethod:: __contains__ .. automethod:: __getitem__ .. automethod:: __iter__ - .. automethod:: __len__ .. automethod:: assign - .. automethod:: copy .. automethod:: ref """ - def __init__(self, _symbol_table: Optional[MutableMapping[str, Array]] = None) -> None: - if _symbol_table is None: - _symbol_table = {} - self._symbol_table: MutableMapping[str, Array] = _symbol_table + def __init__(self) -> None: + self._symbol_table: Dict[str, Optional[Array]] = {} - def __contains__(self, name: object) -> bool: + def __contains__(self, name: str) -> bool: return name in self._symbol_table def __getitem__(self, name: str) -> Array: @@ -207,28 +167,14 @@ class Namespace(Mapping[str, "Array"]): def __iter__(self) -> Iterator[str]: return iter(self._symbol_table) - - def __len__(self) -> int: - return len(self._symbol_table) - - def _chain(self) -> Namespace: - return Namespace(collections.ChainMap(dict(), self._symbol_table)) - - def copy(self) -> Namespace: - result = Namespace() - mapper = _NamespaceCopyMapper() - cache: Dict[Array, Array] = {} - for name in self: - val = mapper(self[name], result, cache) - if name not in result: - result.assign(name, val) - return result - - def assign(self, name: str, value: Array) -> str: + + def assign(self, name: str, + value: Optional[Array]) -> str: """Declare a new array. :param name: a Python identifier - :param value: the array object + :param value: the array object, or None if the assignment is to + just reserve a name :returns: *name* """ @@ -293,7 +239,6 @@ class UniqueTag(Tag): Only one instance of this type of tag may be assigned to a single tagged object. """ - pass TagsType = FrozenSet[Tag] @@ -343,7 +288,7 @@ def normalize_shape( def normalize_shape_component( s: ConvertibleToShapeComponent) -> ScalarExpression: if isinstance(s, str): - s = scalar_expr.parse(s) + s = sym.parse(s) if isinstance(s, int): if s < 0: @@ -356,7 +301,7 @@ def normalize_shape( return s if isinstance(shape, str): - shape = scalar_expr.parse(shape) + shape = sym.parse(shape) from numbers import Number if isinstance(shape, (Number, prim.Expression)): @@ -429,25 +374,21 @@ class Array: """ - def __init__(self, namespace: Namespace, shape: ShapeType, dtype: np.dtype, tags: Optional[TagsType] = None): + def __init__(self, namespace: Namespace, + tags: Optional[TagsType] = None): if tags is None: tags = frozenset() self.namespace = namespace self.tags = tags - self._shape = shape - self._dtype = dtype + self.dtype: np.dtype = np.float64 # FIXME def copy(self, **kwargs: Any) -> Array: raise NotImplementedError @property def shape(self) -> ShapeType: - return self._shape - - @property - def dtype(self) -> np.dtype: - return self._dtype + raise NotImplementedError def named(self, name: str) -> Array: return self.namespace.ref(self.namespace.assign(name, self)) @@ -475,56 +416,8 @@ class Array: return self.copy(tags=new_tags) - @memoize_method - def __hash__(self) -> int: - raise NotImplementedError - - def __eq__(self, other: Any) -> bool: - raise NotImplementedError - - def __ne__(self, other: Any) -> bool: - return not self.__eq__(other) - - def _join_dtypes(self, *args: np.dtype) -> np.dtype: - result = args[0] - for arg in args[1:]: - result = (np.empty(0, dtype=result) + np.empty(0, dtype=arg)).dtype - return result - - def _binary_op(self, op: Any, - other: Union[Array, Number], - reverse: bool = False) -> Array: - if isinstance(other, Number): - # TODO - raise NotImplementedError - else: - if self.shape != other.shape: - raise ValueError("shapes do not match for binary operator") - - dtype = self._join_dtypes(self.dtype, other.dtype) - - # FIXME: If either *self* or *other* is an IndexLambda, its expression - # could be folded into the output, producing a fused result. - if self.shape == (): - expr = op(prim.Variable("_in0"), prim.Variable("_in1")) - else: - indices = tuple(prim.Variable(f"_{i}") for i in range(self.ndim)) - expr = op( - prim.Variable("_in0")[indices], - prim.Variable("_in1")[indices]) - - first, second = self, other - if reverse: - first, second = second, first - - bindings = dict(_in0=first, _in1=second) - - return IndexLambda( - self.namespace, expr, - shape=self.shape, dtype=dtype, bindings=bindings) - - __mul__ = partialmethod(_binary_op, operator.mul) - __rmul__ = partialmethod(_binary_op, operator.mul, reverse=True) + # TODO: + # - codegen interface # }}} @@ -643,26 +536,32 @@ class IndexLambda(Array): .. automethod:: is_reference """ - mapper_method = "map_index_lambda" + # TODO: write make_index_lambda() that does dtype inference def __init__( - self, - namespace: Namespace, - expr: prim.Expression, - shape: ShapeType, - dtype: np.dtype, + self, namespace: Namespace, expr: prim.Expression, + shape: ShapeType, dtype: np.dtype, bindings: Optional[Dict[str, Array]] = None, tags: Optional[TagsType] = None): if bindings is None: bindings = {} - super().__init__(namespace, shape=shape, dtype=dtype, tags=tags) + super().__init__(namespace, tags=tags) + self._shape = shape + self._dtype = dtype self.expr = expr self.bindings = bindings - @memoize_method + @property + def shape(self) -> ShapeType: + return self._shape + + @property + def dtype(self) -> np.dtype: + return self._dtype + def is_reference(self) -> bool: # FIXME: Do we want a specific 'reference' node to make all this # checking unnecessary? @@ -693,28 +592,6 @@ class IndexLambda(Array): return True - @memoize_method - def __hash__(self) -> int: - return hash(( - self.expr, - self.shape, - self.dtype, - frozenset(self.bindings.items()), - self.tags)) - - def __eq__(self, other: object) -> bool: - if self is other: - return True - - return ( - isinstance(other, IndexLambda) - and self.namespace is other.namespace - and self.expr == other.expr - and self.shape == other.shape - and self.dtype == other.dtype - and self.bindings == other.bindings - and self.tags == other.tags) - # }}} @@ -759,54 +636,24 @@ class DataWrapper(Array): # TODO: not really Any data def __init__(self, namespace: Namespace, data: Any, tags: Optional[TagsType] = None): - super().__init__(namespace, shape=data.shape, dtype=data.dtype, tags=tags) - self.data = data - -# }}} - - -# {{{ placeholder - -class _ArgLike(Array): + super().__init__(namespace, tags) - def __init__(self, - namespace: Namespace, - name: str, - shape: ShapeType, - dtype: np.dtype, - tags: Optional[TagsType] = None): - if name is None: - raise ValueError("Must have explicit name") - - # Reserve the name, prevent others from using it. - namespace.assign(name, self) - - super().__init__( - namespace=namespace, shape=shape, dtype=dtype, tags=tags) + self.data = data - self.name = name + @property + def shape(self) -> Any: # FIXME + return self.data.shape - @memoize_method - def __hash__(self) -> int: - return hash((self.name,)) + @property + def dtype(self) -> np.dtype: + return self.data.dtype - def __eq__(self, other: object) -> bool: - if self is other: - return True - # Uniquely identified by name. - return ( - isinstance(other, _ArgLike) - and self.namespace is other.namespace - and self.name == other.name) +# }}} - def tagged(self, tag: Tag) -> Array: - raise ValueError("Cannot modify tags") - - def without_tag(self, tag: Tag, verify_existence: bool = True) -> Array: - raise ValueError("Cannot modify tags") +# {{{ placeholder -class Placeholder(_ArgLike): +class Placeholder(Array): """ A named placeholder for an array whose concrete value is supplied by the user during evaluation. @@ -818,46 +665,29 @@ class Placeholder(_ArgLike): .. note:: - Modifying :class:`Placeholder` tags is not supported after - creation. + :attr:`name` is not a :term:`namespace name`. In fact, + it is prohibited from being one. (This has to be the case: Suppose a + :class:`Placeholder` is :meth:`~Array.tagged`, would the namespace name + refer to the tagged or the untagged version?) """ - mapper_method = "map_placeholder" - -# }}} - - -# {{{ output - -class Output(_ArgLike): - """A named output of the computation. - - .. attribute:: name - - The name of the output array. - - .. attribute:: array - - The :class:`Array` value that is output. + def __init__(self, namespace: Namespace, + name: str, shape: ShapeType, + tags: Optional[TagsType] = None): - .. note:: + # Reserve the name, prevent others from using it. + namespace.assign(name, None) - Modifying :class:`Output` tags is not supported after creation. - """ + super().__init__(namespace=namespace, tags=tags) - mapper_method = "map_output" + self.name = name + self._shape = shape - def __init__(self, - namespace: Namespace, - name: str, - array: Array, - tags: Optional[TagsType] = None): - super().__init__(namespace=namespace, - name=name, - shape=array.shape, - dtype=array.dtype, - tags=tags) - self.array = array + @property + def shape(self) -> ShapeType: + # Matt added this to make Pylint happy. + # Not tied to this, open for discussion about how to implement this. + return self._shape # }}} @@ -894,26 +724,20 @@ def make_dict_of_named_arrays( def make_placeholder(namespace: Namespace, name: str, shape: ConvertibleToShape, - dtype: np.dtype, tags: Optional[TagsType] = None ) -> Placeholder: """Make a :class:`Placeholder` object. - :param namespace: namespace of the placeholder array - :param name: name of the placeholder array - :param shape: shape of the placeholder array - :param dtype: dtype of the placeholder array - :param tags: implementation tags + :param namespace: namespace of the placeholder array + :param shape: shape of the placeholder array + :param tags: implementation tags """ - if name is None: - raise ValueError("Placeholder instances must have a name") - if not str.isidentifier(name): - raise ValueError(f"'{name}' is not a Python identifier") + raise ValueError(f"{name} is not a Python identifier") shape = normalize_shape(shape, namespace) - return Placeholder(namespace, name, shape, dtype, tags) + return Placeholder(namespace, name, shape, tags) # }}} diff --git a/pytato/scalar_expr.py b/pytato/symbolic.py similarity index 92% rename from pytato/scalar_expr.py rename to pytato/symbolic.py index bfd948f..8a6799c 100644 --- a/pytato/scalar_expr.py +++ b/pytato/symbolic.py @@ -37,17 +37,14 @@ from pymbolic.mapper.dependency import (DependencyMapper as import pymbolic.primitives as prim __doc__ = """ -.. currentmodule:: pytato.scalar_expr +.. currentmodule:: pytato.symbolic -Scalar-Level Expressions ------------------------- +Symbolic Infrastructure +----------------------- .. data:: ScalarExpression - A :class:`type`, representing a scalar-valued symbolic expression. Such - expressions can be manipulated and composed as :mod:`pymbolic` expressions. - - Concretely, this is an alias for ``Union[Number, pymbolic.primitives.Expression]``. + A type alias for ``Union[Number, pymbolic.primitives.Expression]``. .. autofunction:: parse .. autofunction:: get_dependencies @@ -60,15 +57,17 @@ Scalar-Level Expressions ScalarExpression = Union[Number, prim.Expression] + def parse(s: str) -> ScalarExpression: from pymbolic.parser import Parser return Parser()(s) -# }}} +# }}} # {{{ mapper classes + class WalkMapper(WalkMapperBase): pass @@ -84,11 +83,12 @@ class SubstitutionMapper(SubstitutionMapperBase): class DependencyMapper(DependencyMapperBase): pass -# }}} +# }}} # {{{ mapper frontends + def get_dependencies(expression: Any) -> FrozenSet[str]: """Return the set of variable names in an expression. @@ -107,6 +107,7 @@ def substitute(expression: Any, variable_assigments: Mapping[str, Any]) -> Any: from pymbolic.mapper.substitutor import make_subst_func return SubstitutionMapper(make_subst_func(variable_assigments))(expression) + # }}} -- GitLab From 3bc2ce88abaf4f0125230d0aebc7f3ce26d967c7 Mon Sep 17 00:00:00 2001 From: Matt Wala Date: Wed, 3 Jun 2020 00:25:20 -0500 Subject: [PATCH 08/76] Revert "Merge branch 'master' into matt-codegen-v4" This reverts commit 962ea5763e4cd5926727e4061c1734cf6820593f, reversing changes made to f5e1450d17ebb8e89564b98c518f5119b06b7ebc. --- pytato/array.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/pytato/array.py b/pytato/array.py index 8429b75..061bf82 100644 --- a/pytato/array.py +++ b/pytato/array.py @@ -332,9 +332,9 @@ class Array: .. attribute:: namespace - A (mutable) instance of :class:`Namespace` containing the - names used in the computation. All arrays in a - computation share the same namespace. + A (mutable) instance of :class:`Namespace` containing the + names used in the computation. All arrays in a + computation share the same namespace. .. attribute:: shape -- GitLab From fbc8c5c708a91f9e13393a0dd3271dae94d64d85 Mon Sep 17 00:00:00 2001 From: Matt Wala Date: Wed, 3 Jun 2020 00:27:46 -0500 Subject: [PATCH 09/76] Rename symbolic to scalar_expr --- doc/reference.rst | 2 +- pytato/array.py | 10 +++++----- pytato/{symbolic.py => scalar_expr.py} | 0 3 files changed, 6 insertions(+), 6 deletions(-) rename pytato/{symbolic.py => scalar_expr.py} (100%) diff --git a/doc/reference.rst b/doc/reference.rst index c743d58..a05ca62 100644 --- a/doc/reference.rst +++ b/doc/reference.rst @@ -3,4 +3,4 @@ Reference .. automodule:: pytato.array .. automodule:: pytato.program -.. automodule:: pytato.symbolic +.. automodule:: pytato.scalar_expr diff --git a/pytato/array.py b/pytato/array.py index 061bf82..f0dd76c 100644 --- a/pytato/array.py +++ b/pytato/array.py @@ -87,8 +87,8 @@ Node constructors such as :class:`Placeholder.__init__` and import numpy as np import pymbolic.primitives as prim -import pytato.symbolic as sym -from pytato.symbolic import ScalarExpression +import pytato.scalar_expr as scalar_expr +from pytato.scalar_expr import ScalarExpression from dataclasses import dataclass from pytools import is_single_valued @@ -267,7 +267,7 @@ def _check_identifier(s: str, ns: Optional[Namespace] = None) -> bool: return True -class _ShapeChecker(sym.WalkMapper): +class _ShapeChecker(scalar_expr.WalkMapper): def __init__(self, ns: Optional[Namespace] = None): super().__init__() self.ns = ns @@ -288,7 +288,7 @@ def normalize_shape( def normalize_shape_component( s: ConvertibleToShapeComponent) -> ScalarExpression: if isinstance(s, str): - s = sym.parse(s) + s = scalar_expr.parse(s) if isinstance(s, int): if s < 0: @@ -301,7 +301,7 @@ def normalize_shape( return s if isinstance(shape, str): - shape = sym.parse(shape) + shape = scalar_expr.parse(shape) from numbers import Number if isinstance(shape, (Number, prim.Expression)): diff --git a/pytato/symbolic.py b/pytato/scalar_expr.py similarity index 100% rename from pytato/symbolic.py rename to pytato/scalar_expr.py -- GitLab From 11fc9c0fa4cdaa8109047eca85a0386a282a4e6d Mon Sep 17 00:00:00 2001 From: Matt Wala Date: Wed, 3 Jun 2020 00:30:11 -0500 Subject: [PATCH 10/76] Fix setup.cfg --- setup.cfg | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/setup.cfg b/setup.cfg index 4d4fdb8..9984342 100644 --- a/setup.cfg +++ b/setup.cfg @@ -2,7 +2,7 @@ ignore = E126,E127,E128,E123,E226,E241,E242,E265,N802,W503,E402,N814,N817,W504 max-line-length=85 -[mypy-pytato.symbolic] +[mypy-pytato.scalar_expr] disallow_subclassing_any = False [mypy-islpy] -- GitLab From 2b44c763491827e3760e925327f0492c888905ca Mon Sep 17 00:00:00 2001 From: Matt Wala Date: Wed, 3 Jun 2020 00:30:58 -0500 Subject: [PATCH 11/76] Style fixes --- pytato/scalar_expr.py | 8 ++------ 1 file changed, 2 insertions(+), 6 deletions(-) diff --git a/pytato/scalar_expr.py b/pytato/scalar_expr.py index 8a6799c..b4a4dca 100644 --- a/pytato/scalar_expr.py +++ b/pytato/scalar_expr.py @@ -57,16 +57,14 @@ Symbolic Infrastructure ScalarExpression = Union[Number, prim.Expression] - def parse(s: str) -> ScalarExpression: from pymbolic.parser import Parser return Parser()(s) - # }}} -# {{{ mapper classes +# {{{ mapper classes class WalkMapper(WalkMapperBase): pass @@ -83,11 +81,10 @@ class SubstitutionMapper(SubstitutionMapperBase): class DependencyMapper(DependencyMapperBase): pass - # }}} -# {{{ mapper frontends +# {{{ mapper frontends def get_dependencies(expression: Any) -> FrozenSet[str]: """Return the set of variable names in an expression. @@ -107,7 +104,6 @@ def substitute(expression: Any, variable_assigments: Mapping[str, Any]) -> Any: from pymbolic.mapper.substitutor import make_subst_func return SubstitutionMapper(make_subst_func(variable_assigments))(expression) - # }}} -- GitLab From 7ad89d9371dd7d6a27fe1efc659c0393148818e7 Mon Sep 17 00:00:00 2001 From: Matt Wala Date: Wed, 3 Jun 2020 00:35:06 -0500 Subject: [PATCH 12/76] Tweak scalar_expr docs --- pytato/scalar_expr.py | 16 +++++++++++----- 1 file changed, 11 insertions(+), 5 deletions(-) diff --git a/pytato/scalar_expr.py b/pytato/scalar_expr.py index b4a4dca..b0dbdf3 100644 --- a/pytato/scalar_expr.py +++ b/pytato/scalar_expr.py @@ -39,12 +39,15 @@ import pymbolic.primitives as prim __doc__ = """ .. currentmodule:: pytato.symbolic -Symbolic Infrastructure ------------------------ +Scalar Expressions +------------------ .. data:: ScalarExpression - A type alias for ``Union[Number, pymbolic.primitives.Expression]``. + A :class:`type` for scalar-valued symbolic expressions. Expressions are + composable and manipulable via :mod:`pymbolic`. + + Concretely, this is an alias for ``Union[Number, pymbolic.primitives.Expression]``. .. autofunction:: parse .. autofunction:: get_dependencies @@ -57,6 +60,7 @@ Symbolic Infrastructure ScalarExpression = Union[Number, prim.Expression] + def parse(s: str) -> ScalarExpression: from pymbolic.parser import Parser return Parser()(s) @@ -89,7 +93,8 @@ class DependencyMapper(DependencyMapperBase): def get_dependencies(expression: Any) -> FrozenSet[str]: """Return the set of variable names in an expression. - :param expression: A :mod:`pymbolic` expression + :param expression: A scalar expression, or an expression derived from such + (e.g., a tuple of scalar expressions) """ mapper = DependencyMapper(composite_leaves=False) return frozenset(dep.name for dep in mapper(expression)) @@ -98,7 +103,8 @@ def get_dependencies(expression: Any) -> FrozenSet[str]: def substitute(expression: Any, variable_assigments: Mapping[str, Any]) -> Any: """Perform variable substitution in an expression. - :param expression: A :mod:`pymbolic` expression + :param expression: A scalar expression, or an expression derived from such + (e.g., a tuple of scalar expressions) :param variable_assigments: A mapping from variable names to substitutions """ from pymbolic.mapper.substitutor import make_subst_func -- GitLab From 4309667ecd02d3c6ee10161e19aeb2a075941c89 Mon Sep 17 00:00:00 2001 From: Matt Wala Date: Wed, 3 Jun 2020 00:36:46 -0500 Subject: [PATCH 13/76] Fix array.py indentation --- pytato/array.py | 12 +++++++----- 1 file changed, 7 insertions(+), 5 deletions(-) diff --git a/pytato/array.py b/pytato/array.py index f0dd76c..4204106 100644 --- a/pytato/array.py +++ b/pytato/array.py @@ -285,10 +285,12 @@ def normalize_shape( :param ns: if a namespace is given, extra checks are performed to ensure that expressions are well-defined. """ + from pytato.scalar_expr import parse + def normalize_shape_component( s: ConvertibleToShapeComponent) -> ScalarExpression: if isinstance(s, str): - s = scalar_expr.parse(s) + s = parse(s) if isinstance(s, int): if s < 0: @@ -301,7 +303,7 @@ def normalize_shape( return s if isinstance(shape, str): - shape = scalar_expr.parse(shape) + shape = parse(shape) from numbers import Number if isinstance(shape, (Number, prim.Expression)): @@ -332,9 +334,9 @@ class Array: .. attribute:: namespace - A (mutable) instance of :class:`Namespace` containing the - names used in the computation. All arrays in a - computation share the same namespace. + A (mutable) instance of :class:`Namespace` containing the + names used in the computation. All arrays in a + computation share the same namespace. .. attribute:: shape -- GitLab From 2650207af7a5a639e610b346a52004bb8da26879 Mon Sep 17 00:00:00 2001 From: Matt Wala Date: Wed, 3 Jun 2020 00:37:18 -0500 Subject: [PATCH 14/76] flake8 fix --- pytato/scalar_expr.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/pytato/scalar_expr.py b/pytato/scalar_expr.py index b0dbdf3..f9ae82a 100644 --- a/pytato/scalar_expr.py +++ b/pytato/scalar_expr.py @@ -47,7 +47,8 @@ Scalar Expressions A :class:`type` for scalar-valued symbolic expressions. Expressions are composable and manipulable via :mod:`pymbolic`. - Concretely, this is an alias for ``Union[Number, pymbolic.primitives.Expression]``. + Concretely, this is an alias for + ``Union[Number, pymbolic.primitives.Expression]``. .. autofunction:: parse .. autofunction:: get_dependencies -- GitLab From 0872498d6a162336eb33415f3f2110ec645bab00 Mon Sep 17 00:00:00 2001 From: Matt Wala Date: Wed, 3 Jun 2020 00:38:04 -0500 Subject: [PATCH 15/76] Fix scalar_expr.py --- pytato/scalar_expr.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pytato/scalar_expr.py b/pytato/scalar_expr.py index f9ae82a..6755135 100644 --- a/pytato/scalar_expr.py +++ b/pytato/scalar_expr.py @@ -37,7 +37,7 @@ from pymbolic.mapper.dependency import (DependencyMapper as import pymbolic.primitives as prim __doc__ = """ -.. currentmodule:: pytato.symbolic +.. currentmodule:: pytato.scalar_expr Scalar Expressions ------------------ -- GitLab From 7d1a72312aab879db194dba6eb2b0e8e935a4fa1 Mon Sep 17 00:00:00 2001 From: Matt Wala Date: Wed, 3 Jun 2020 00:42:52 -0500 Subject: [PATCH 16/76] WIP: Update array.py --- pytato/array.py | 314 +++++++++++++++++++++++++++++++++++++----------- 1 file changed, 244 insertions(+), 70 deletions(-) diff --git a/pytato/array.py b/pytato/array.py index 4204106..9027342 100644 --- a/pytato/array.py +++ b/pytato/array.py @@ -67,6 +67,7 @@ Built-in Expression Nodes .. autoclass:: Reshape .. autoclass:: DataWrapper .. autoclass:: Placeholder +.. autoclass:: Output .. autoclass:: LoopyFunction User-Facing Node Creation @@ -85,16 +86,21 @@ Node constructors such as :class:`Placeholder.__init__` and # }}} +import collections +from functools import partialmethod +from numbers import Number +import operator +from dataclasses import dataclass +from typing import Optional, Dict, Any, MutableMapping, Mapping, Iterator, Tuple, Union, FrozenSet + import numpy as np +import pymbolic.mapper import pymbolic.primitives as prim +from pytools import is_single_valued, memoize_method + import pytato.scalar_expr as scalar_expr from pytato.scalar_expr import ScalarExpression -from dataclasses import dataclass -from pytools import is_single_valued -from typing import Optional, Dict, Any, Mapping, Iterator, Tuple, Union, FrozenSet - - # {{{ dotted name class DottedName: @@ -138,7 +144,37 @@ class DottedName: # {{{ namespace -class Namespace: +class _NamespaceCopyMapper(scalar_expr.IdentityMapper): + + def __call__(self, expr: Array, namespace: Namespace, cache: Dict[Array, Array]) -> Array: + return self.rec(expr, namespace, cache) + + def rec(self, expr: Array, namespace: Namespace, cache: Dict[Array, Array]) -> Array: + if expr in cache: + return cache[expr] + result: Array = super().rec(expr, namespace, cache) + cache[expr] = result + return result + + def map_index_lambda(self, expr: IndexLambda, namespace: Namespace, cache: Dict[Array, Array]) -> Array: + bindings = { + name: self.rec(subexpr, namespace, cache) + for name, subexpr in expr.bindings.items()} + return IndexLambda( + namespace, + expr=expr.expr, + shape=expr.shape, + dtype=expr.dtype, + bindings=bindings) + + def map_placeholder(self, expr: Placeholder, namespace: Namespace, cache: Dict[Array, Array]) -> Array: + return Placeholder(namespace, expr.name, expr.shape, expr.dtype, expr.tags) + + def map_output(self, expr: Output, namespace: Namespace, cache: Dict[Array, Array]) -> Array: + return Output(namespace, expr.name, self.rec(expr.array, namespace, cache), expr.tags) + + +class Namespace(Mapping[str, "Array"]): # Possible future extension: .parent attribute r""" Represents a mapping from :term:`identifier` strings to @@ -149,14 +185,18 @@ class Namespace: .. automethod:: __contains__ .. automethod:: __getitem__ .. automethod:: __iter__ + .. automethod:: __len__ .. automethod:: assign + .. automethod:: copy .. automethod:: ref """ - def __init__(self) -> None: - self._symbol_table: Dict[str, Optional[Array]] = {} + def __init__(self, _symbol_table: Optional[MutableMapping[str, Array]] = None) -> None: + if _symbol_table is None: + _symbol_table = {} + self._symbol_table: MutableMapping[str, Array] = _symbol_table - def __contains__(self, name: str) -> bool: + def __contains__(self, name: object) -> bool: return name in self._symbol_table def __getitem__(self, name: str) -> Array: @@ -168,13 +208,27 @@ class Namespace: def __iter__(self) -> Iterator[str]: return iter(self._symbol_table) - def assign(self, name: str, - value: Optional[Array]) -> str: + def __len__(self) -> int: + return len(self._symbol_table) + + def _chain(self) -> Namespace: + return Namespace(collections.ChainMap(dict(), self._symbol_table)) + + def copy(self) -> Namespace: + result = Namespace() + mapper = _NamespaceCopyMapper() + cache: Dict[Array, Array] = {} + for name in self: + val = mapper(self[name], result, cache) + if name not in result: + result.assign(name, val) + return result + + def assign(self, name: str, value: Array) -> str: """Declare a new array. :param name: a Python identifier - :param value: the array object, or None if the assignment is to - just reserve a name + :param value: the array object :returns: *name* """ @@ -239,6 +293,7 @@ class UniqueTag(Tag): Only one instance of this type of tag may be assigned to a single tagged object. """ + pass TagsType = FrozenSet[Tag] @@ -285,12 +340,10 @@ def normalize_shape( :param ns: if a namespace is given, extra checks are performed to ensure that expressions are well-defined. """ - from pytato.scalar_expr import parse - def normalize_shape_component( s: ConvertibleToShapeComponent) -> ScalarExpression: if isinstance(s, str): - s = parse(s) + s = scalar_expr.parse(s) if isinstance(s, int): if s < 0: @@ -303,7 +356,7 @@ def normalize_shape( return s if isinstance(shape, str): - shape = parse(shape) + shape = scalar_expr.parse(shape) from numbers import Number if isinstance(shape, (Number, prim.Expression)): @@ -376,21 +429,25 @@ class Array: """ - def __init__(self, namespace: Namespace, - tags: Optional[TagsType] = None): + def __init__(self, namespace: Namespace, shape: ShapeType, dtype: np.dtype, tags: Optional[TagsType] = None): if tags is None: tags = frozenset() self.namespace = namespace self.tags = tags - self.dtype: np.dtype = np.float64 # FIXME + self._shape = shape + self._dtype = dtype def copy(self, **kwargs: Any) -> Array: raise NotImplementedError @property def shape(self) -> ShapeType: - raise NotImplementedError + return self._shape + + @property + def dtype(self) -> np.dtype: + return self._dtype def named(self, name: str) -> Array: return self.namespace.ref(self.namespace.assign(name, self)) @@ -418,8 +475,56 @@ class Array: return self.copy(tags=new_tags) - # TODO: - # - codegen interface + @memoize_method + def __hash__(self) -> int: + raise NotImplementedError + + def __eq__(self, other: Any) -> bool: + raise NotImplementedError + + def __ne__(self, other: Any) -> bool: + return not self.__eq__(other) + + def _join_dtypes(self, *args: np.dtype) -> np.dtype: + result = args[0] + for arg in args[1:]: + result = (np.empty(0, dtype=result) + np.empty(0, dtype=arg)).dtype + return result + + def _binary_op(self, op: Any, + other: Union[Array, Number], + reverse: bool = False) -> Array: + if isinstance(other, Number): + # TODO + raise NotImplementedError + else: + if self.shape != other.shape: + raise ValueError("shapes do not match for binary operator") + + dtype = self._join_dtypes(self.dtype, other.dtype) + + # FIXME: If either *self* or *other* is an IndexLambda, its expression + # could be folded into the output, producing a fused result. + if self.shape == (): + expr = op(prim.Variable("_in0"), prim.Variable("_in1")) + else: + indices = tuple(prim.Variable(f"_{i}") for i in range(self.ndim)) + expr = op( + prim.Variable("_in0")[indices], + prim.Variable("_in1")[indices]) + + first, second = self, other + if reverse: + first, second = second, first + + bindings = dict(_in0=first, _in1=second) + + return IndexLambda( + self.namespace, expr, + shape=self.shape, dtype=dtype, bindings=bindings) + + __mul__ = partialmethod(_binary_op, operator.mul) + __rmul__ = partialmethod(_binary_op, operator.mul, reverse=True) # }}} @@ -538,32 +643,26 @@ class IndexLambda(Array): .. automethod:: is_reference """ - # TODO: write make_index_lambda() that does dtype inference + mapper_method = "map_index_lambda" def __init__( - self, namespace: Namespace, expr: prim.Expression, - shape: ShapeType, dtype: np.dtype, + self, + namespace: Namespace, + expr: prim.Expression, + shape: ShapeType, + dtype: np.dtype, bindings: Optional[Dict[str, Array]] = None, tags: Optional[TagsType] = None): if bindings is None: bindings = {} - super().__init__(namespace, tags=tags) + super().__init__(namespace, shape=shape, dtype=dtype, tags=tags) - self._shape = shape - self._dtype = dtype self.expr = expr self.bindings = bindings - @property - def shape(self) -> ShapeType: - return self._shape - - @property - def dtype(self) -> np.dtype: - return self._dtype - + @memoize_method def is_reference(self) -> bool: # FIXME: Do we want a specific 'reference' node to make all this # checking unnecessary? @@ -594,6 +693,28 @@ class IndexLambda(Array): return True + @memoize_method + def __hash__(self) -> int: + return hash(( + self.expr, + self.shape, + self.dtype, + frozenset(self.bindings.items()), + self.tags)) + + def __eq__(self, other: object) -> bool: + if self is other: + return True + + return ( + isinstance(other, IndexLambda) + and self.namespace is other.namespace + and self.expr == other.expr + and self.shape == other.shape + and self.dtype == other.dtype + and self.bindings == other.bindings + and self.tags == other.tags) + # }}} @@ -638,24 +759,54 @@ class DataWrapper(Array): # TODO: not really Any data def __init__(self, namespace: Namespace, data: Any, tags: Optional[TagsType] = None): - super().__init__(namespace, tags) - + super().__init__(namespace, shape=data.shape, dtype=data.dtype, tags=tags) self.data = data - @property - def shape(self) -> Any: # FIXME - return self.data.shape - - @property - def dtype(self) -> np.dtype: - return self.data.dtype - # }}} # {{{ placeholder -class Placeholder(Array): +class _ArgLike(Array): + + def __init__(self, + namespace: Namespace, + name: str, + shape: ShapeType, + dtype: np.dtype, + tags: Optional[TagsType] = None): + if name is None: + raise ValueError("Must have explicit name") + + # Reserve the name, prevent others from using it. + namespace.assign(name, self) + + super().__init__( + namespace=namespace, shape=shape, dtype=dtype, tags=tags) + + self.name = name + + @memoize_method + def __hash__(self) -> int: + return hash((self.name,)) + + def __eq__(self, other: object) -> bool: + if self is other: + return True + # Uniquely identified by name. + return ( + isinstance(other, _ArgLike) + and self.namespace is other.namespace + and self.name == other.name) + + def tagged(self, tag: Tag) -> Array: + raise ValueError("Cannot modify tags") + + def without_tag(self, tag: Tag, verify_existence: bool = True) -> Array: + raise ValueError("Cannot modify tags") + + +class Placeholder(_ArgLike): """ A named placeholder for an array whose concrete value is supplied by the user during evaluation. @@ -667,29 +818,46 @@ class Placeholder(Array): .. note:: - :attr:`name` is not a :term:`namespace name`. In fact, - it is prohibited from being one. (This has to be the case: Suppose a - :class:`Placeholder` is :meth:`~Array.tagged`, would the namespace name - refer to the tagged or the untagged version?) + Modifying :class:`Placeholder` tags is not supported after + creation. """ - def __init__(self, namespace: Namespace, - name: str, shape: ShapeType, - tags: Optional[TagsType] = None): + mapper_method = "map_placeholder" - # Reserve the name, prevent others from using it. - namespace.assign(name, None) +# }}} - super().__init__(namespace=namespace, tags=tags) - self.name = name - self._shape = shape +# {{{ output - @property - def shape(self) -> ShapeType: - # Matt added this to make Pylint happy. - # Not tied to this, open for discussion about how to implement this. - return self._shape +class Output(_ArgLike): + """A named output of the computation. + + .. attribute:: name + + The name of the output array. + + .. attribute:: array + + The :class:`Array` value that is output. + + .. note:: + + Modifying :class:`Output` tags is not supported after creation. + """ + + mapper_method = "map_output" + + def __init__(self, + namespace: Namespace, + name: str, + array: Array, + tags: Optional[TagsType] = None): + super().__init__(namespace=namespace, + name=name, + shape=array.shape, + dtype=array.dtype, + tags=tags) + self.array = array # }}} @@ -726,20 +894,26 @@ def make_dict_of_named_arrays( def make_placeholder(namespace: Namespace, name: str, shape: ConvertibleToShape, + dtype: np.dtype, tags: Optional[TagsType] = None ) -> Placeholder: """Make a :class:`Placeholder` object. - :param namespace: namespace of the placeholder array - :param shape: shape of the placeholder array - :param tags: implementation tags + :param namespace: namespace of the placeholder array + :param name: name of the placeholder array + :param shape: shape of the placeholder array + :param dtype: dtype of the placeholder array + :param tags: implementation tags """ + if name is None: + raise ValueError("Placeholder instances must have a name") + if not str.isidentifier(name): - raise ValueError(f"{name} is not a Python identifier") + raise ValueError(f"'{name}' is not a Python identifier") shape = normalize_shape(shape, namespace) - return Placeholder(namespace, name, shape, tags) + return Placeholder(namespace, name, shape, dtype, tags) # }}} -- GitLab From cbd15f579d65b8caf0f195b53cebfb879b25f19c Mon Sep 17 00:00:00 2001 From: Matt Wala Date: Wed, 3 Jun 2020 00:57:34 -0500 Subject: [PATCH 17/76] Add stubs for pytools --- MANIFEST.in | 1 + pytato/stubs/pytools.pyi | 11 +++++++++++ run-mypy.sh | 1 + setup.cfg | 4 ---- 4 files changed, 13 insertions(+), 4 deletions(-) create mode 100644 MANIFEST.in create mode 100644 pytato/stubs/pytools.pyi diff --git a/MANIFEST.in b/MANIFEST.in new file mode 100644 index 0000000..d887c9f --- /dev/null +++ b/MANIFEST.in @@ -0,0 +1 @@ +recursive-include pytato *.pyi diff --git a/pytato/stubs/pytools.pyi b/pytato/stubs/pytools.pyi new file mode 100644 index 0000000..82531cc --- /dev/null +++ b/pytato/stubs/pytools.pyi @@ -0,0 +1,11 @@ +# FIXME: Should be in pytools +from typing import TypeVar, Iterable, Optional + +T = TypeVar("T") + +def memoize_method(f: T) -> T: ... +def is_single_valued(it: Iterable[T]) -> bool: ... + +class UniqueNameGenerator: + def __init__(self, existing_names: Optional[Iterable[str]], forced_prefix: str=""): ... + def __call__(self, based_on: str = "") -> str: ... diff --git a/run-mypy.sh b/run-mypy.sh index 14730b1..72b4e41 100755 --- a/run-mypy.sh +++ b/run-mypy.sh @@ -1,3 +1,4 @@ #! /bin/bash +export MYPYPATH=$MYPYPATH:pytato/stubs/ mypy --strict pytato diff --git a/setup.cfg b/setup.cfg index 9984342..7b0c351 100644 --- a/setup.cfg +++ b/setup.cfg @@ -19,7 +19,3 @@ ignore_missing_imports = True [mypy-pyopencl] ignore_missing_imports = True - -[mypy-pytools] -ignore_missing_imports = True - -- GitLab From 4802c242bc13d138236708c5a56db393013b2968 Mon Sep 17 00:00:00 2001 From: Matt Wala Date: Wed, 3 Jun 2020 01:22:41 -0500 Subject: [PATCH 18/76] Implement array_expr --- doc/reference.rst | 1 + pytato/array_expr.py | 106 +++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 107 insertions(+) create mode 100644 pytato/array_expr.py diff --git a/doc/reference.rst b/doc/reference.rst index a05ca62..a991d28 100644 --- a/doc/reference.rst +++ b/doc/reference.rst @@ -3,4 +3,5 @@ Reference .. automodule:: pytato.array .. automodule:: pytato.program +.. automodule:: pytato.array_expr .. automodule:: pytato.scalar_expr diff --git a/pytato/array_expr.py b/pytato/array_expr.py new file mode 100644 index 0000000..4c68081 --- /dev/null +++ b/pytato/array_expr.py @@ -0,0 +1,106 @@ +from __future__ import annotations + +__copyright__ = """ +Copyright (C) 2020 Matt Wala +""" + +__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 typing import Dict + +from pymbolic.mapper import Mapper as MapperBase + +from pytato.array import Array, IndexLambda, Namespace, Output, Placeholder + +__doc__ = """ +.. currentmodule:: pytato.array_expr + +Tools for Array Expressions +--------------------------- + +.. autoclass:: CopyMapper +.. autofunction:: copy_namespace + +""" + + +# {{{ mapper classes + +class Mapper(MapperBase): + pass + + +class CopyMapper(Mapper): + namespace: Namespace + + def __init__(self, new_namespace: Namespace): + self.namespace = new_namespace + self.cache: Dict[Array, Array] = {} + + def __call__(self, expr: Array) -> Array: + return self.rec(expr) + + def rec(self, expr: Array) -> Array: + if expr in self.cache: + return self.cache[expr] + result: Array = super().rec(expr) + self.cache[expr] = result + return result + + def map_index_lambda(self, expr: IndexLambda) -> Array: + bindings = { + name: self.rec(subexpr) + for name, subexpr in expr.bindings.items()} + return IndexLambda(self.namespace, + expr=expr.expr, + shape=expr.shape, + dtype=expr.dtype, + bindings=bindings) + + def map_placeholder(self, expr: Placeholder) -> Array: + return Placeholder(self.namespace, expr.name, expr.shape, expr.dtype, + expr.tags) + + def map_output(self, expr: Output) -> Array: + return Output(self.namespace, expr.name, self.rec(expr.array), + expr.tags) + +# }}} + + +# {{{ mapper frontends + +def copy_namespace(namespace: Namespace, copy_mapper: CopyMapper) -> Namespace: + """Copy the elements of *namespace* into a new namespace. + + :param namespace: The original namespace + :param mapper: A mapper that performs copies into a new namespace + :returns: The new namespace + """ + for name, val in namespace.items(): + mapped_val = copy_mapper(val) + if name not in copy_mapper.namespace: + copy_mapper.namespace.assign(name, mapped_val) + return copy_mapper.namespace + +# }}} + +# vim: foldmethod=marker -- GitLab From c9a4aad9cd4c34acaa295f3d337b38c59999956b Mon Sep 17 00:00:00 2001 From: Matt Wala Date: Wed, 3 Jun 2020 01:26:15 -0500 Subject: [PATCH 19/76] Fix array and array_expr --- pytato/array.py | 59 +++++++------------------------------------- pytato/array_expr.py | 4 +-- setup.cfg | 3 +++ 3 files changed, 14 insertions(+), 52 deletions(-) diff --git a/pytato/array.py b/pytato/array.py index 9027342..15237ce 100644 --- a/pytato/array.py +++ b/pytato/array.py @@ -86,21 +86,20 @@ Node constructors such as :class:`Placeholder.__init__` and # }}} -import collections from functools import partialmethod from numbers import Number import operator from dataclasses import dataclass -from typing import Optional, Dict, Any, MutableMapping, Mapping, Iterator, Tuple, Union, FrozenSet +from typing import Optional, Dict, Any, Mapping, Iterator, Tuple, Union, FrozenSet import numpy as np -import pymbolic.mapper import pymbolic.primitives as prim from pytools import is_single_valued, memoize_method import pytato.scalar_expr as scalar_expr from pytato.scalar_expr import ScalarExpression + # {{{ dotted name class DottedName: @@ -144,36 +143,6 @@ class DottedName: # {{{ namespace -class _NamespaceCopyMapper(scalar_expr.IdentityMapper): - - def __call__(self, expr: Array, namespace: Namespace, cache: Dict[Array, Array]) -> Array: - return self.rec(expr, namespace, cache) - - def rec(self, expr: Array, namespace: Namespace, cache: Dict[Array, Array]) -> Array: - if expr in cache: - return cache[expr] - result: Array = super().rec(expr, namespace, cache) - cache[expr] = result - return result - - def map_index_lambda(self, expr: IndexLambda, namespace: Namespace, cache: Dict[Array, Array]) -> Array: - bindings = { - name: self.rec(subexpr, namespace, cache) - for name, subexpr in expr.bindings.items()} - return IndexLambda( - namespace, - expr=expr.expr, - shape=expr.shape, - dtype=expr.dtype, - bindings=bindings) - - def map_placeholder(self, expr: Placeholder, namespace: Namespace, cache: Dict[Array, Array]) -> Array: - return Placeholder(namespace, expr.name, expr.shape, expr.dtype, expr.tags) - - def map_output(self, expr: Output, namespace: Namespace, cache: Dict[Array, Array]) -> Array: - return Output(namespace, expr.name, self.rec(expr.array, namespace, cache), expr.tags) - - class Namespace(Mapping[str, "Array"]): # Possible future extension: .parent attribute r""" @@ -191,10 +160,8 @@ class Namespace(Mapping[str, "Array"]): .. automethod:: ref """ - def __init__(self, _symbol_table: Optional[MutableMapping[str, Array]] = None) -> None: - if _symbol_table is None: - _symbol_table = {} - self._symbol_table: MutableMapping[str, Array] = _symbol_table + def __init__(self) -> None: + self._symbol_table: Dict[str, Array] = {} def __contains__(self, name: object) -> bool: return name in self._symbol_table @@ -211,18 +178,9 @@ class Namespace(Mapping[str, "Array"]): def __len__(self) -> int: return len(self._symbol_table) - def _chain(self) -> Namespace: - return Namespace(collections.ChainMap(dict(), self._symbol_table)) - def copy(self) -> Namespace: - result = Namespace() - mapper = _NamespaceCopyMapper() - cache: Dict[Array, Array] = {} - for name in self: - val = mapper(self[name], result, cache) - if name not in result: - result.assign(name, val) - return result + from pytato.array_expr import CopyMapper, copy_namespace + return copy_namespace(self, CopyMapper(Namespace())) def assign(self, name: str, value: Array) -> str: """Declare a new array. @@ -429,7 +387,8 @@ class Array: """ - def __init__(self, namespace: Namespace, shape: ShapeType, dtype: np.dtype, tags: Optional[TagsType] = None): + def __init__(self, namespace: Namespace, shape: ShapeType, dtype: np.dtype, + tags: Optional[TagsType] = None): if tags is None: tags = frozenset() @@ -794,7 +753,7 @@ class _ArgLike(Array): if self is other: return True # Uniquely identified by name. - return ( + return ( isinstance(other, _ArgLike) and self.namespace is other.namespace and self.name == other.name) diff --git a/pytato/array_expr.py b/pytato/array_expr.py index 4c68081..fdd8202 100644 --- a/pytato/array_expr.py +++ b/pytato/array_expr.py @@ -91,8 +91,8 @@ class CopyMapper(Mapper): def copy_namespace(namespace: Namespace, copy_mapper: CopyMapper) -> Namespace: """Copy the elements of *namespace* into a new namespace. - :param namespace: The original namespace - :param mapper: A mapper that performs copies into a new namespace + :param namespace: The source namespace + :param copy_mapper: A mapper that performs copies into a new namespace :returns: The new namespace """ for name, val in namespace.items(): diff --git a/setup.cfg b/setup.cfg index 7b0c351..7c3f3d4 100644 --- a/setup.cfg +++ b/setup.cfg @@ -2,6 +2,9 @@ ignore = E126,E127,E128,E123,E226,E241,E242,E265,N802,W503,E402,N814,N817,W504 max-line-length=85 +[mypy-pytato.array_expr] +disallow_subclassing_any = False + [mypy-pytato.scalar_expr] disallow_subclassing_any = False -- GitLab From 384fde39424880451f3aab49d678c4cb861d6d35 Mon Sep 17 00:00:00 2001 From: Matt Wala Date: Wed, 3 Jun 2020 01:57:10 -0500 Subject: [PATCH 20/76] Formatting in array.py --- pytato/array.py | 105 +++++++-------- pytato/codegen.py | 318 ++++++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 373 insertions(+), 50 deletions(-) create mode 100644 pytato/codegen.py diff --git a/pytato/array.py b/pytato/array.py index 15237ce..096cfc2 100644 --- a/pytato/array.py +++ b/pytato/array.py @@ -386,8 +386,10 @@ class Array: .. attribute:: ndim """ - - def __init__(self, namespace: Namespace, shape: ShapeType, dtype: np.dtype, + def __init__(self, + namespace: Namespace, + shape: ShapeType, + dtype: np.dtype, tags: Optional[TagsType] = None): if tags is None: tags = frozenset() @@ -425,9 +427,7 @@ class Array: return self.copy(tags=self.tags | frozenset([tag])) def without_tag(self, tag: Tag, verify_existence: bool = True) -> Array: - new_tags = tuple( - t for t in self.tags - if t != tag) + new_tags = tuple(t for t in self.tags if t != tag) if verify_existence and len(new_tags) == len(self.tags): raise ValueError(f"tag '{tag}' was not present") @@ -450,9 +450,10 @@ class Array: result = (np.empty(0, dtype=result) + np.empty(0, dtype=arg)).dtype return result - def _binary_op(self, op: Any, - other: Union[Array, Number], - reverse: bool = False) -> Array: + def _binary_op(self, + op: Any, + other: Union[Array, Number], + reverse: bool = False) -> Array: if isinstance(other, Number): # TODO raise NotImplementedError @@ -467,7 +468,8 @@ class Array: if self.shape == (): expr = op(prim.Variable("_in0"), prim.Variable("_in1")) else: - indices = tuple(prim.Variable(f"_{i}") for i in range(self.ndim)) + indices = tuple( + prim.Variable(f"_{i}") for i in range(self.ndim)) expr = op( prim.Variable("_in0")[indices], prim.Variable("_in1")[indices]) @@ -478,13 +480,16 @@ class Array: bindings = dict(_in0=first, _in1=second) - return IndexLambda( - self.namespace, expr, - shape=self.shape, dtype=dtype, bindings=bindings) + return IndexLambda(self.namespace, + expr, + shape=self.shape, + dtype=dtype, + bindings=bindings) __mul__ = partialmethod(_binary_op, operator.mul) __rmul__ = partialmethod(_binary_op, operator.mul, reverse=True) + # }}} @@ -604,8 +609,7 @@ class IndexLambda(Array): mapper_method = "map_index_lambda" - def __init__( - self, + def __init__(self, namespace: Namespace, expr: prim.Expression, shape: ShapeType, @@ -654,25 +658,20 @@ class IndexLambda(Array): @memoize_method def __hash__(self) -> int: - return hash(( - self.expr, - self.shape, - self.dtype, - frozenset(self.bindings.items()), - self.tags)) + return hash((self.expr, self.shape, self.dtype, + frozenset(self.bindings.items()), self.tags)) def __eq__(self, other: object) -> bool: if self is other: return True - return ( - isinstance(other, IndexLambda) - and self.namespace is other.namespace - and self.expr == other.expr - and self.shape == other.shape - and self.dtype == other.dtype - and self.bindings == other.bindings - and self.tags == other.tags) + return (isinstance(other, IndexLambda) + and self.namespace is other.namespace + and self.expr == other.expr + and self.shape == other.shape + and self.dtype == other.dtype + and self.bindings == other.bindings + and self.tags == other.tags) # }}} @@ -683,15 +682,17 @@ class Einsum(Array): """ """ -# }}} +# }}} # {{{ reshape + class Reshape(Array): """ """ + # }}} @@ -716,9 +717,14 @@ class DataWrapper(Array): """ # TODO: not really Any data - def __init__(self, namespace: Namespace, data: Any, - tags: Optional[TagsType] = None): - super().__init__(namespace, shape=data.shape, dtype=data.dtype, tags=tags) + def __init__(self, + namespace: Namespace, + data: Any, + tags: Optional[TagsType] = None): + super().__init__(namespace, + shape=data.shape, + dtype=data.dtype, + tags=tags) self.data = data # }}} @@ -727,34 +733,34 @@ class DataWrapper(Array): # {{{ placeholder class _ArgLike(Array): - def __init__(self, - namespace: Namespace, - name: str, - shape: ShapeType, - dtype: np.dtype, - tags: Optional[TagsType] = None): + namespace: Namespace, + name: str, + shape: ShapeType, + dtype: np.dtype, + tags: Optional[TagsType] = None): if name is None: raise ValueError("Must have explicit name") # Reserve the name, prevent others from using it. namespace.assign(name, self) - super().__init__( - namespace=namespace, shape=shape, dtype=dtype, tags=tags) + super().__init__(namespace=namespace, + shape=shape, + dtype=dtype, + tags=tags) self.name = name @memoize_method def __hash__(self) -> int: - return hash((self.name,)) + return hash((self.name, )) def __eq__(self, other: object) -> bool: if self is other: return True # Uniquely identified by name. - return ( - isinstance(other, _ArgLike) + return (isinstance(other, _ArgLike) and self.namespace is other.namespace and self.name == other.name) @@ -832,13 +838,13 @@ class LoopyFunction(DictOfNamedArrays): name. """ + # }}} # {{{ end-user-facing -def make_dict_of_named_arrays( - data: Dict[str, Array]) -> DictOfNamedArrays: +def make_dict_of_named_arrays(data: Dict[str, Array]) -> DictOfNamedArrays: """Make a :class:`DictOfNamedArrays` object and ensure that all arrays share the same namespace. @@ -851,11 +857,10 @@ def make_dict_of_named_arrays( def make_placeholder(namespace: Namespace, - name: str, - shape: ConvertibleToShape, - dtype: np.dtype, - tags: Optional[TagsType] = None - ) -> Placeholder: + name: str, + shape: ConvertibleToShape, + dtype: np.dtype, + tags: Optional[TagsType] = None) -> Placeholder: """Make a :class:`Placeholder` object. :param namespace: namespace of the placeholder array diff --git a/pytato/codegen.py b/pytato/codegen.py new file mode 100644 index 0000000..e265fbe --- /dev/null +++ b/pytato/codegen.py @@ -0,0 +1,318 @@ +from __future__ import annotations + +# Codegen output class. + + +import collections +import dataclasses +from typing import cast, Any, ContextManager, Callable, Union, Optional, Mapping, Iterator, Dict, Tuple, Set, FrozenSet +import typing + +import contextlib +import loopy as lp +import numpy as np +import pymbolic.mapper +import pymbolic.primitives as prim +import pytools + +from pytato.array import Array, DictOfNamedArrays, Placeholder, Output, Namespace, ShapeType, IndexLambda +from pytato.program import BoundProgram, Target, PyOpenCLTarget +import pytato.symbolic as sym +from pytato.symbolic import ScalarExpression + + +# These are semantically distinct but identical at the type level. +SymbolicIndex = ShapeType + + +# {{{ nodal result + +class NodalResult(object): + + def __init__(self, shape: ShapeType, dtype: np.dtype): + self.shape = shape + self.dtype = dtype + + @property + def ndim(self) -> int: + return len(self.shape) + + def to_loopy_expression(self, indices: SymbolicIndex, context: ExpressionContext, reduction_inames: Optional[Tuple[str, ...]] = None) -> ScalarExpression: + raise NotImplementedError + + +class ArrayResult(NodalResult): + + def __init__(self, name: str, shape: ShapeType, dtype: np.dtype): + # TODO: Stuff dependencies in here. + super().__init__(shape, dtype) + self.name = name + + def to_loopy_expression(self, indices: SymbolicIndex, context: ExpressionContext, reduction_inames: Optional[Tuple[str, ...]] = None) -> ScalarExpression: + if indices == (): + return prim.Variable(self.name) + else: + return prim.Variable(self.name)[indices] + + +class ExpressionResult(NodalResult): + + def __init__(self, expr: ScalarExpression, shape: ShapeType, dtype: np.dtype): + super().__init__(shape, dtype) + self.expr = expr + + def to_loopy_expression(self, indices: SymbolicIndex, context: ExpressionContext, reduction_inames: Optional[Tuple[str, ...]] = None) -> ScalarExpression: + return sym.substitute( + self.expr, + dict(zip( + (f"_{d}" for d in range(self.ndim)), + indices))) + + +class SubstitutionRuleResult(NodalResult): + # TODO: implement + pass + +# }}} + + +class ExpressionGenMapper(sym.IdentityMapper): + """A mapper for generating :mod:`loopy` expressions. + + The inputs to this mapper are :class:`IndexLambda` expressions, or + expressions that are closely equivalent (e.g., shape expressions). In + particular + """ + codegen_mapper: CodeGenMapper + + def __init__(self, codegen_mapper: CodeGenMapper): + self.codegen_mapper = codegen_mapper + + def __call__(self, + expr: ScalarExpression, + indices: Tuple[ScalarExpression, ...], + context: ExpressionContext) -> ScalarExpression: + return self.rec(expr, indices, context) + + def map_subscript(self, expr: prim.Subscript, indices: SymbolicIndex, context: ExpressionContext) -> ScalarExpression: + assert isinstance(expr.aggregate, prim.Variable) + result: NodalResult = self.codegen_mapper( + context.namespace[expr.aggregate.name], + context.state) + assert len(expr.index) == len(indices) + mapped_indices = sym.substitute( + expr.index, + dict(zip( + (f"_{d}" for d in range(len(indices))), + indices))) + return result.to_loopy_expression(mapped_indices, context) + + # TODO: map_reduction() + + def map_variable(self, expr: prim.Variable, indices: SymbolicIndex, context: ExpressionContext) -> ScalarExpression: + result: NodalResult = self.codegen_mapper( + context.namespace[expr.name], + context.state) + return result.to_loopy_expression((), context) + + +class CodeGenMapper(pymbolic.mapper.Mapper): + """A mapper for generating code for nodes in the computation graph. + """ + exprgen_mapper: ExpressionGenMapper + + def __init__(self) -> None: + self.exprgen_mapper = ExpressionGenMapper(self) + + def map_placeholder(self, expr: Placeholder, state: CodeGenState) -> NodalResult: + if expr in state.results: + return state.results[expr] + + arg = lp.GlobalArg(expr.name, shape=expr.shape, dtype=expr.dtype, order="C") + kernel = state.kernel.copy(args=state.kernel.args + [arg]) + state.update_kernel(kernel) + + result = ArrayResult(expr.name, expr.dtype, expr.shape) + state.results[expr] = result + return result + + def map_output(self, expr: Output, state: CodeGenState) -> NodalResult: + if expr in state.results: + return state.results[expr] + + # FIXE: Scalar outputs are not supported yet. + assert expr.shape != () + + inner_result = self.rec(expr.array, state) + + inames = tuple( + state.var_name_gen(f"{expr.name}_dim{d}") + for d in range(expr.ndim)) + domain = sym.domain_for_shape(inames, expr.shape) + + arg = lp.GlobalArg( + expr.name, + shape=expr.shape, + dtype=expr.dtype, + order="C", + is_output_only=True) + + indices = tuple(prim.Variable(iname) for iname in inames) + context = state.make_expression_context() + copy_expr = inner_result.to_loopy_expression(indices, context) + # TODO: Context data supported yet. + assert not context.reduction_bounds + assert not context.depends_on, context.depends_on + + from loopy.kernel.instruction import make_assignment + insn = make_assignment( + (prim.Variable(expr.name)[indices],), + copy_expr, + id=state.insn_id_gen(f"{expr.name}_copy"), + within_inames=frozenset(inames), + depends_on=context.depends_on) + + kernel = state.kernel + kernel = kernel.copy( + args=kernel.args + [arg], + instructions=kernel.instructions + [insn], + domains=kernel.domains + [domain]) + state.update_kernel(kernel) + + result = ArrayResult(expr.name, expr.dtype, expr.shape) + state.results[expr] = result + return result + + def map_index_lambda(self, expr: IndexLambda, state: CodeGenState) -> NodalResult: + if expr in state.results: + return state.results[expr] + + # TODO: tags + + with state.chain_namespaces(expr.bindings) as chained_state: + expr_context = chained_state.make_expression_context() + indices = tuple(prim.Variable(f"_{d}") for d in range(expr.ndim)) + generated_expr = self.exprgen_mapper(expr.expr, indices, expr_context) + + result = ExpressionResult(generated_expr, expr.shape, expr.dtype) + state.results[expr] = result + return result + + +@dataclasses.dataclass(init=True, repr=False, eq=False) +class CodeGenState: + """ + This data is threaded through :class:`CodeGenMapper`. + + .. attribute:: namespace + .. attribute:: kernel + .. attribute:: results + .. attribute:: var_name_gen + .. attribute:: insn_id_gen + """ + namespace: typing.ChainMap[str, Array] + _kernel: lp.LoopKernel + results: Dict[Array, NodalResult] + # Both of these have type Callable[[str], str], but mypy's support for that is broken. + var_name_gen: Any = dataclasses.field(init=False) + insn_id_gen: Any = dataclasses.field(init=False) + + def __post_init__(self) -> None: + self.var_name_gen = self._kernel.get_var_name_generator() + self.insn_id_gen = self._kernel.get_var_name_generator() + + @property + def kernel(self) -> lp.LoopKernel: + return self._kernel + + def update_kernel(self, kernel: lp.LoopyKernel) -> None: + self._kernel = kernel + + @contextlib.contextmanager + def chain_namespaces(self, local_namespace: Mapping[str, Array]) -> Iterator[CodeGenState]: + self.namespace.maps.insert(0, local_namespace) + yield self + self.namespace.maps.pop(0) + + def make_expression_context(self, depends_on: FrozenSet[str] = frozenset(), reduction_bounds: Optional[ReductionBounds] = None) -> ExpressionContext: + if reduction_bounds is None: + reduction_bounds = {} + return ExpressionContext(self, _depends_on=depends_on, reduction_bounds=reduction_bounds) + + +ReductionBounds = Dict[str, Tuple[ScalarExpression, ScalarExpression]] + + +@dataclasses.dataclass(init=True, repr=False, eq=False) +class ExpressionContext(object): + """Contextual data for generating :mod:`loopy` expressions. + + This data is threaded through :class:`ExpressionGenMapper`. + + .. attribute:: state + .. attribute:: _depends_on + .. attribute:: reduction_bounds + """ + state: CodeGenState + _depends_on: FrozenSet[str] + reduction_bounds: Dict[str, Tuple[ScalarExpression, ScalarExpression]] + + @property + def namespace(self) -> typing.ChainMap[str, Array]: + return self.state.namespace + + @property + def depends_on(self) -> FrozenSet[str]: + return self._depends_on + + def update_depends_on(self, other: FrozenSet[str]) -> None: + self._depends_on = self._depends_on | other + + +def generate_loopy(result: Union[Namespace, Array, DictOfNamedArrays], + target: Optional[Target] = None) -> BoundProgram: + # {{{ get namespace + + if isinstance(result, Array): + if isinstance(result, Output): + result = result.namespace + else: + result = DictOfNamedArrays({"_out": result}) + + if isinstance(result, DictOfNamedArrays): + namespace = result.namespace._chain() + + # Augment with Output nodes. + name_gen = pytools.UniqueNameGenerator(set(namespace)) + for name, val in result.items(): + out_name = name_gen(name) + Output(namespace, out_name, val) + + result = namespace.copy() + + assert isinstance(result, Namespace) + + # Make an internal copy. + result = result.copy() + + # }}} + + if target is None: + target = PyOpenCLTarget() + + # Set up codegen state. + kernel = lp.make_kernel( + "{:}", [], target=target.get_loopy_target(), + lang_version=lp.MOST_RECENT_LANGUAGE_VERSION) + + state = CodeGenState( + namespace=collections.ChainMap(result), + _kernel=kernel, + results=dict()) + + # Generate code for graph nodes. + mapper = CodeGenMapper() + for name, val in result.items(): + _ = mapper(val, state) + + return target.bind_program(program=state.kernel, bound_arguments=dict()) -- GitLab From 50d570f8d67ee6c5d9160b90b3f0ed56fbcc3bc4 Mon Sep 17 00:00:00 2001 From: Matt Wala Date: Wed, 3 Jun 2020 02:14:13 -0500 Subject: [PATCH 21/76] Rename array_expr to transform --- doc/reference.rst | 2 +- pytato/array.py | 2 +- pytato/{array_expr.py => transform.py} | 6 +++--- 3 files changed, 5 insertions(+), 5 deletions(-) rename pytato/{array_expr.py => transform.py} (97%) diff --git a/doc/reference.rst b/doc/reference.rst index a991d28..2a14d6a 100644 --- a/doc/reference.rst +++ b/doc/reference.rst @@ -3,5 +3,5 @@ Reference .. automodule:: pytato.array .. automodule:: pytato.program -.. automodule:: pytato.array_expr .. automodule:: pytato.scalar_expr +.. automodule:: pytato.transform diff --git a/pytato/array.py b/pytato/array.py index 096cfc2..54f6e32 100644 --- a/pytato/array.py +++ b/pytato/array.py @@ -179,7 +179,7 @@ class Namespace(Mapping[str, "Array"]): return len(self._symbol_table) def copy(self) -> Namespace: - from pytato.array_expr import CopyMapper, copy_namespace + from pytato.transform import CopyMapper, copy_namespace return copy_namespace(self, CopyMapper(Namespace())) def assign(self, name: str, value: Array) -> str: diff --git a/pytato/array_expr.py b/pytato/transform.py similarity index 97% rename from pytato/array_expr.py rename to pytato/transform.py index fdd8202..76a1b49 100644 --- a/pytato/array_expr.py +++ b/pytato/transform.py @@ -31,10 +31,10 @@ from pymbolic.mapper import Mapper as MapperBase from pytato.array import Array, IndexLambda, Namespace, Output, Placeholder __doc__ = """ -.. currentmodule:: pytato.array_expr +.. currentmodule:: pytato.transform -Tools for Array Expressions ---------------------------- +Transforming Computations +------------------------- .. autoclass:: CopyMapper .. autofunction:: copy_namespace -- GitLab From 6092000d6272d6d5f69e36c8539faefc01146669 Mon Sep 17 00:00:00 2001 From: Matt Wala Date: Wed, 3 Jun 2020 03:44:08 -0500 Subject: [PATCH 22/76] Add codegen infrastructure --- doc/codegen.rst | 4 + doc/index.rst | 1 + pytato/__init__.py | 12 +- pytato/codegen.py | 394 ++++++++++++++++++++++++--------------- pytato/program.py | 23 ++- pytato/stubs/pytools.pyi | 4 +- setup.cfg | 2 +- test/test_codegen.py | 60 ++++++ 8 files changed, 345 insertions(+), 155 deletions(-) create mode 100644 doc/codegen.rst create mode 100755 test/test_codegen.py diff --git a/doc/codegen.rst b/doc/codegen.rst new file mode 100644 index 0000000..917751c --- /dev/null +++ b/doc/codegen.rst @@ -0,0 +1,4 @@ +Generating Code +=============== + +.. automodule:: pytato.codegen diff --git a/doc/index.rst b/doc/index.rst index 18dd4da..187dda3 100644 --- a/doc/index.rst +++ b/doc/index.rst @@ -6,6 +6,7 @@ Welcome to Pytato's documentation! :caption: Contents: reference + codegen design misc diff --git a/pytato/__init__.py b/pytato/__init__.py index 4cef6a4..dcfaf28 100644 --- a/pytato/__init__.py +++ b/pytato/__init__.py @@ -29,5 +29,13 @@ from pytato.array import ( DottedName, Placeholder, make_placeholder, ) -__all__ = ("DottedName", "Namespace", "Array", "DictOfNamedArrays", - "Tag", "UniqueTag", "Placeholder", "make_placeholder") +from pytato.codegen import generate_loopy +from pytato.program import Target, PyOpenCLTarget + +__all__ = ( + "DottedName", "Namespace", "Array", "DictOfNamedArrays", + "Tag", "UniqueTag", "Placeholder", "make_placeholder", + + "generate_loopy", + "Target", "PyOpenCLTarget", +) diff --git a/pytato/codegen.py b/pytato/codegen.py index e265fbe..b3bf553 100644 --- a/pytato/codegen.py +++ b/pytato/codegen.py @@ -1,34 +1,82 @@ from __future__ import annotations -# Codegen output class. - +__copyright__ = """Copyright (C) 2020 Matt Wala""" + +__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. +""" import collections import dataclasses -from typing import cast, Any, ContextManager, Callable, Union, Optional, Mapping, Iterator, Dict, Tuple, Set, FrozenSet +from typing import Any, Union, Optional, Mapping, Iterator, Dict, Tuple, FrozenSet import typing import contextlib import loopy as lp import numpy as np -import pymbolic.mapper import pymbolic.primitives as prim import pytools -from pytato.array import Array, DictOfNamedArrays, Placeholder, Output, Namespace, ShapeType, IndexLambda +from pytato.array import ( + Array, DictOfNamedArrays, Placeholder, Output, Namespace, ShapeType, + IndexLambda) from pytato.program import BoundProgram, Target, PyOpenCLTarget -import pytato.symbolic as sym -from pytato.symbolic import ScalarExpression +import pytato.scalar_expr as scalar_expr +from pytato.scalar_expr import ScalarExpression +import pytato.transform -# These are semantically distinct but identical at the type level. -SymbolicIndex = ShapeType +__doc__ = """ + +.. currentmodule:: pytato + +.. autofunction:: generate_loopy + +Code Generation Internals +------------------------- + +.. currentmodule:: pytato.codegen + +.. autoclass:: GeneratedResult +.. autoclass:: ArrayResult +.. autoclass:: LoopyExpressionResult +.. autoclass:: SubstitutionRuleResult + +.. autoclass:: CodeGenState +.. autoclass:: CodeGenMapper +.. autoclass:: LoopyExpressionContext +.. autoclass:: LoopyExpressionGenMapper -# {{{ nodal result +""" -class NodalResult(object): +# {{{ generated array expressions + +# These are semantically distinct but identical at the type level. +SymbolicIndex = ShapeType + + +class GeneratedResult(object): + """Generated code for a node in the computation graph (i.e., an array + expression). + """ def __init__(self, shape: ShapeType, dtype: np.dtype): self.shape = shape self.dtype = dtype @@ -37,98 +85,133 @@ class NodalResult(object): def ndim(self) -> int: return len(self.shape) - def to_loopy_expression(self, indices: SymbolicIndex, context: ExpressionContext, reduction_inames: Optional[Tuple[str, ...]] = None) -> ScalarExpression: + def to_loopy_expression(self, indices: SymbolicIndex, + context: LoopyExpressionContext) -> ScalarExpression: + """Return a :mod:`loopy` expression for this result.""" raise NotImplementedError -class ArrayResult(NodalResult): - +class ArrayResult(GeneratedResult): + """An array expression generated as a :mod:`loopy` array.""" def __init__(self, name: str, shape: ShapeType, dtype: np.dtype): - # TODO: Stuff dependencies in here. super().__init__(shape, dtype) self.name = name - def to_loopy_expression(self, indices: SymbolicIndex, context: ExpressionContext, reduction_inames: Optional[Tuple[str, ...]] = None) -> ScalarExpression: + # TODO: Handle dependencies. + def to_loopy_expression(self, indices: SymbolicIndex, + context: LoopyExpressionContext) -> ScalarExpression: if indices == (): return prim.Variable(self.name) else: return prim.Variable(self.name)[indices] -class ExpressionResult(NodalResult): - - def __init__(self, expr: ScalarExpression, shape: ShapeType, dtype: np.dtype): +class LoopyExpressionResult(GeneratedResult): + """An array expression generated as a :mod:`loopy` expression.""" + def __init__( + self, expr: ScalarExpression, shape: ShapeType, dtype: np.dtype): super().__init__(shape, dtype) self.expr = expr - def to_loopy_expression(self, indices: SymbolicIndex, context: ExpressionContext, reduction_inames: Optional[Tuple[str, ...]] = None) -> ScalarExpression: - return sym.substitute( + # TODO: Handle dependencies and reduction domains. + def to_loopy_expression(self, indices: SymbolicIndex, + context: LoopyExpressionContext) -> ScalarExpression: + return scalar_expr.substitute( self.expr, - dict(zip( - (f"_{d}" for d in range(self.ndim)), - indices))) + dict(zip((f"_{d}" for d in range(self.ndim)), indices))) -class SubstitutionRuleResult(NodalResult): +class SubstitutionRuleResult(GeneratedResult): # TODO: implement pass # }}} -class ExpressionGenMapper(sym.IdentityMapper): - """A mapper for generating :mod:`loopy` expressions. +# {{{ codegen + +@dataclasses.dataclass(init=True, repr=False, eq=False) +class CodeGenState: + """Data threaded through :class:`CodeGenMapper`. + + .. attribute:: namespace + + The namespace + + .. attribute:: kernel - The inputs to this mapper are :class:`IndexLambda` expressions, or - expressions that are closely equivalent (e.g., shape expressions). In - particular + The partial kernel + + .. attribute:: results + + A mapping from arrays to code generation results + + .. attribute:: var_name_gen + .. attribute:: insn_id_gen + + .. automethod:: update_kernel + .. automethod:: chain_namespaces + .. automethod:: make_expression_context """ - codegen_mapper: CodeGenMapper + namespace: typing.ChainMap[str, Array] + _kernel: lp.LoopKernel + results: Dict[Array, GeneratedResult] - def __init__(self, codegen_mapper: CodeGenMapper): - self.codegen_mapper = codegen_mapper + # Both of these have type Callable[[str], str], but mypy's support for that + # is broken (https://github.com/python/mypy/issues/6910) + var_name_gen: Any = dataclasses.field(init=False) + insn_id_gen: Any = dataclasses.field(init=False) - def __call__(self, - expr: ScalarExpression, - indices: Tuple[ScalarExpression, ...], - context: ExpressionContext) -> ScalarExpression: - return self.rec(expr, indices, context) + def __post_init__(self) -> None: + self.var_name_gen = self._kernel.get_var_name_generator() + self.insn_id_gen = self._kernel.get_var_name_generator() - def map_subscript(self, expr: prim.Subscript, indices: SymbolicIndex, context: ExpressionContext) -> ScalarExpression: - assert isinstance(expr.aggregate, prim.Variable) - result: NodalResult = self.codegen_mapper( - context.namespace[expr.aggregate.name], - context.state) - assert len(expr.index) == len(indices) - mapped_indices = sym.substitute( - expr.index, - dict(zip( - (f"_{d}" for d in range(len(indices))), - indices))) - return result.to_loopy_expression(mapped_indices, context) + @property + def kernel(self) -> lp.LoopKernel: + return self._kernel - # TODO: map_reduction() + def update_kernel(self, kernel: lp.LoopKernel) -> None: + self._kernel = kernel - def map_variable(self, expr: prim.Variable, indices: SymbolicIndex, context: ExpressionContext) -> ScalarExpression: - result: NodalResult = self.codegen_mapper( - context.namespace[expr.name], - context.state) - return result.to_loopy_expression((), context) + @contextlib.contextmanager + def chain_namespaces( + self, + local_namespace: Mapping[str, Array]) -> Iterator[CodeGenState]: + """A context manager for overriding with a local scope.""" + self.namespace.maps.insert(0, local_namespace) + yield self + self.namespace.maps.pop(0) + + def make_expression_context( + self, + depends_on: FrozenSet[str] = frozenset(), + reduction_bounds: Optional[ReductionBounds] = None + ) -> LoopyExpressionContext: + """Get a new :class:`LoopyExpressionContext`.""" + if reduction_bounds is None: + reduction_bounds = {} + return LoopyExpressionContext(self, + _depends_on=depends_on, + reduction_bounds=reduction_bounds) -class CodeGenMapper(pymbolic.mapper.Mapper): +class CodeGenMapper(pytato.transform.Mapper): """A mapper for generating code for nodes in the computation graph. """ - exprgen_mapper: ExpressionGenMapper + exprgen_mapper: LoopyExpressionGenMapper def __init__(self) -> None: - self.exprgen_mapper = ExpressionGenMapper(self) + self.exprgen_mapper = LoopyExpressionGenMapper(self) - def map_placeholder(self, expr: Placeholder, state: CodeGenState) -> NodalResult: + def map_placeholder(self, expr: Placeholder, + state: CodeGenState) -> GeneratedResult: if expr in state.results: return state.results[expr] - arg = lp.GlobalArg(expr.name, shape=expr.shape, dtype=expr.dtype, order="C") + arg = lp.GlobalArg(expr.name, + shape=expr.shape, + dtype=expr.dtype, + order="C") kernel = state.kernel.copy(args=state.kernel.args + [arg]) state.update_kernel(kernel) @@ -136,7 +219,7 @@ class CodeGenMapper(pymbolic.mapper.Mapper): state.results[expr] = result return result - def map_output(self, expr: Output, state: CodeGenState) -> NodalResult: + def map_output(self, expr: Output, state: CodeGenState) -> GeneratedResult: if expr in state.results: return state.results[expr] @@ -144,14 +227,13 @@ class CodeGenMapper(pymbolic.mapper.Mapper): assert expr.shape != () inner_result = self.rec(expr.array, state) - + inames = tuple( state.var_name_gen(f"{expr.name}_dim{d}") for d in range(expr.ndim)) - domain = sym.domain_for_shape(inames, expr.shape) + domain = scalar_expr.domain_for_shape(inames, expr.shape) - arg = lp.GlobalArg( - expr.name, + arg = lp.GlobalArg(expr.name, shape=expr.shape, dtype=expr.dtype, order="C", @@ -160,21 +242,20 @@ class CodeGenMapper(pymbolic.mapper.Mapper): indices = tuple(prim.Variable(iname) for iname in inames) context = state.make_expression_context() copy_expr = inner_result.to_loopy_expression(indices, context) - # TODO: Context data supported yet. + + # TODO: Contextual data not supported yet. assert not context.reduction_bounds - assert not context.depends_on, context.depends_on + assert not context.depends_on from loopy.kernel.instruction import make_assignment - insn = make_assignment( - (prim.Variable(expr.name)[indices],), + insn = make_assignment((prim.Variable(expr.name)[indices], ), copy_expr, id=state.insn_id_gen(f"{expr.name}_copy"), within_inames=frozenset(inames), depends_on=context.depends_on) kernel = state.kernel - kernel = kernel.copy( - args=kernel.args + [arg], + kernel = kernel.copy(args=kernel.args + [arg], instructions=kernel.instructions + [insn], domains=kernel.domains + [domain]) state.update_kernel(kernel) @@ -183,84 +264,55 @@ class CodeGenMapper(pymbolic.mapper.Mapper): state.results[expr] = result return result - def map_index_lambda(self, expr: IndexLambda, state: CodeGenState) -> NodalResult: + def map_index_lambda(self, expr: IndexLambda, + state: CodeGenState) -> GeneratedResult: if expr in state.results: return state.results[expr] - # TODO: tags + # TODO: Respect tags. with state.chain_namespaces(expr.bindings) as chained_state: expr_context = chained_state.make_expression_context() - indices = tuple(prim.Variable(f"_{d}") for d in range(expr.ndim)) - generated_expr = self.exprgen_mapper(expr.expr, indices, expr_context) + loopy_expr = self.exprgen_mapper(expr.expr, expr_context) - result = ExpressionResult(generated_expr, expr.shape, expr.dtype) + result = LoopyExpressionResult(loopy_expr, expr.shape, expr.dtype) state.results[expr] = result return result +# }}} -@dataclasses.dataclass(init=True, repr=False, eq=False) -class CodeGenState: - """ - This data is threaded through :class:`CodeGenMapper`. - - .. attribute:: namespace - .. attribute:: kernel - .. attribute:: results - .. attribute:: var_name_gen - .. attribute:: insn_id_gen - """ - namespace: typing.ChainMap[str, Array] - _kernel: lp.LoopKernel - results: Dict[Array, NodalResult] - # Both of these have type Callable[[str], str], but mypy's support for that is broken. - var_name_gen: Any = dataclasses.field(init=False) - insn_id_gen: Any = dataclasses.field(init=False) - - def __post_init__(self) -> None: - self.var_name_gen = self._kernel.get_var_name_generator() - self.insn_id_gen = self._kernel.get_var_name_generator() - - @property - def kernel(self) -> lp.LoopKernel: - return self._kernel - - def update_kernel(self, kernel: lp.LoopyKernel) -> None: - self._kernel = kernel - - @contextlib.contextmanager - def chain_namespaces(self, local_namespace: Mapping[str, Array]) -> Iterator[CodeGenState]: - self.namespace.maps.insert(0, local_namespace) - yield self - self.namespace.maps.pop(0) - - def make_expression_context(self, depends_on: FrozenSet[str] = frozenset(), reduction_bounds: Optional[ReductionBounds] = None) -> ExpressionContext: - if reduction_bounds is None: - reduction_bounds = {} - return ExpressionContext(self, _depends_on=depends_on, reduction_bounds=reduction_bounds) +# {{{ loopy expression gen mapper ReductionBounds = Dict[str, Tuple[ScalarExpression, ScalarExpression]] @dataclasses.dataclass(init=True, repr=False, eq=False) -class ExpressionContext(object): +class LoopyExpressionContext(object): """Contextual data for generating :mod:`loopy` expressions. - This data is threaded through :class:`ExpressionGenMapper`. + This data is threaded through :class:`LoopyExpressionGenMapper`. .. attribute:: state + + The :class:`CodeGenState`. + .. attribute:: _depends_on + + The set of dependencies associated with the expression. + .. attribute:: reduction_bounds + + A mapping from inames to reduction bounds in the expression. """ state: CodeGenState _depends_on: FrozenSet[str] - reduction_bounds: Dict[str, Tuple[ScalarExpression, ScalarExpression]] + reduction_bounds: ReductionBounds @property def namespace(self) -> typing.ChainMap[str, Array]: return self.state.namespace - + @property def depends_on(self) -> FrozenSet[str]: return self._depends_on @@ -269,31 +321,82 @@ class ExpressionContext(object): self._depends_on = self._depends_on | other -def generate_loopy(result: Union[Namespace, Array, DictOfNamedArrays], - target: Optional[Target] = None) -> BoundProgram: - # {{{ get namespace +class LoopyExpressionGenMapper(scalar_expr.IdentityMapper): + """A mapper for generating :mod:`loopy` expressions. - if isinstance(result, Array): - if isinstance(result, Output): - result = result.namespace - else: - result = DictOfNamedArrays({"_out": result}) + The inputs to this mapper are scalar expression as found in + :class:`pytato.IndexLambda`, or expressions that are compatible (e.g., shape + expressions). - if isinstance(result, DictOfNamedArrays): - namespace = result.namespace._chain() + The outputs of this mapper are scalar expressions suitable for wrapping in + :class:`LoopyExpressionResult`. + """ + codegen_mapper: CodeGenMapper - # Augment with Output nodes. - name_gen = pytools.UniqueNameGenerator(set(namespace)) - for name, val in result.items(): - out_name = name_gen(name) - Output(namespace, out_name, val) + def __init__(self, codegen_mapper: CodeGenMapper): + self.codegen_mapper = codegen_mapper - result = namespace.copy() + def __call__(self, expr: ScalarExpression, + context: LoopyExpressionContext) -> ScalarExpression: + return self.rec(expr, context) - assert isinstance(result, Namespace) - - # Make an internal copy. - result = result.copy() + def map_subscript(self, expr: prim.Subscript, + context: LoopyExpressionContext) -> ScalarExpression: + assert isinstance(expr.aggregate, prim.Variable) + result: GeneratedResult = self.codegen_mapper( + context.namespace[expr.aggregate.name], context.state) + return result.to_loopy_expression(expr.index, context) + + # TODO: map_reduction() + + def map_variable(self, expr: prim.Variable, + context: LoopyExpressionContext) -> ScalarExpression: + result: GeneratedResult = self.codegen_mapper( + context.namespace[expr.name], + context.state) + return result.to_loopy_expression((), context) + +# }}} + + +def _promote_named_arrays_to_outputs(arrays: DictOfNamedArrays) -> Namespace: + # Turns named arrays into Output nodes, returning a new namespace. + copy_mapper = pytato.transform.CopyMapper(Namespace()) + result = pytato.transform.copy_namespace(arrays.namespace, copy_mapper) + + name_gen = pytools.UniqueNameGenerator(set(result)) + for name, val in arrays.items(): + Output(result, name_gen(name), copy_mapper(val)) + + return result + + +def generate_loopy( + result_or_namespace: Union[Namespace, Array, DictOfNamedArrays], + target: Optional[Target] = None) -> BoundProgram: + """Code generation entry point. + + :param result_or_namespace: Either a :class:`pytato.Namespace`, a single + :class:`pytato.Array`, or a :class:`pytato.DictOfNamedArrays`. In the + latter two cases, code generation treats the node(s) as outputs of the + computation. + + :param target: The target for code generation + + :returns: A wrapped generated :mod:`loopy` kernel + """ + # {{{ get namespace + + if isinstance(result_or_namespace, Array): + result_or_namespace = DictOfNamedArrays({"out": result_or_namespace}) + + if isinstance(result_or_namespace, DictOfNamedArrays): + result_or_namespace = _promote_named_arrays_to_outputs( + result_or_namespace) + + assert isinstance(result_or_namespace, Namespace) + namespace = result_or_namespace + del result_or_namespace # }}} @@ -301,18 +404,17 @@ def generate_loopy(result: Union[Namespace, Array, DictOfNamedArrays], target = PyOpenCLTarget() # Set up codegen state. - kernel = lp.make_kernel( - "{:}", [], target=target.get_loopy_target(), + kernel = lp.make_kernel("{:}", [], + target=target.get_loopy_target(), lang_version=lp.MOST_RECENT_LANGUAGE_VERSION) - - state = CodeGenState( - namespace=collections.ChainMap(result), + + state = CodeGenState(namespace=collections.ChainMap(namespace), _kernel=kernel, results=dict()) # Generate code for graph nodes. mapper = CodeGenMapper() - for name, val in result.items(): + for name, val in namespace.items(): _ = mapper(val, state) return target.bind_program(program=state.kernel, bound_arguments=dict()) diff --git a/pytato/program.py b/pytato/program.py index 4a36602..4e0506f 100644 --- a/pytato/program.py +++ b/pytato/program.py @@ -51,7 +51,12 @@ if typing.TYPE_CHECKING: class Target: - """An abstract code generation target.""" + """An abstract code generation target. + + .. automethod:: get_loopy_target + .. automethod:: bind_program + """ + def get_loopy_target(self) -> "lp.TargetBase": """Return the corresponding :mod:`loopy` target.""" raise NotImplementedError @@ -60,14 +65,20 @@ class Target: bound_arguments: Mapping[str, Any]) -> BoundProgram: """Create a :class:`BoundProgram` for this code generation target. - :arg program: the :mod:`loopy` kernel - :arg bound_arguments: a mapping from argument names to outputs + :param program: the :mod:`loopy` kernel + :param bound_arguments: a mapping from argument names to outputs """ raise NotImplementedError class PyOpenCLTarget(Target): - """A :mod:`pyopencl` code generation target.""" + """A :mod:`pyopencl` code generation target. + + .. attribute:: queue + + The :mod:`pyopencl` command queue, or *None*. + """ + def __init__(self, queue: Optional["cl.CommandQueue"] = None): self.queue = queue @@ -101,6 +112,8 @@ class BoundProgram: .. attribute:: bound_arguments A map from names to pre-bound kernel arguments. + + .. automethod:: __call__ """ program: "lp.LoopKernel" @@ -118,6 +131,8 @@ class BoundPyOpenCLProgram(BoundProgram): .. attribute:: queue A :mod:`pyopencl` command queue. + + .. automethod:: __call__ """ queue: Optional["cl.CommandQueue"] diff --git a/pytato/stubs/pytools.pyi b/pytato/stubs/pytools.pyi index 82531cc..2fe9afc 100644 --- a/pytato/stubs/pytools.pyi +++ b/pytato/stubs/pytools.pyi @@ -1,5 +1,5 @@ # FIXME: Should be in pytools -from typing import TypeVar, Iterable, Optional +from typing import TypeVar, Iterable, Set, Optional T = TypeVar("T") @@ -7,5 +7,5 @@ def memoize_method(f: T) -> T: ... def is_single_valued(it: Iterable[T]) -> bool: ... class UniqueNameGenerator: - def __init__(self, existing_names: Optional[Iterable[str]], forced_prefix: str=""): ... + def __init__(self, existing_names: Optional[Set[str]], forced_prefix: str=""): ... def __call__(self, based_on: str = "") -> str: ... diff --git a/setup.cfg b/setup.cfg index 7c3f3d4..b0d0924 100644 --- a/setup.cfg +++ b/setup.cfg @@ -2,7 +2,7 @@ ignore = E126,E127,E128,E123,E226,E241,E242,E265,N802,W503,E402,N814,N817,W504 max-line-length=85 -[mypy-pytato.array_expr] +[mypy-pytato.transform] disallow_subclassing_any = False [mypy-pytato.scalar_expr] diff --git a/test/test_codegen.py b/test/test_codegen.py new file mode 100755 index 0000000..e93200d --- /dev/null +++ b/test/test_codegen.py @@ -0,0 +1,60 @@ +#!/usr/bin/env python + +__copyright__ = "Copyright (C) 2020 Andreas Kloeckner" + +__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. +""" + +import sys + +import numpy as np +import pyopencl as cl +import pyopencl.array as cl_array # noqa +import pyopencl.cltypes as cltypes # noqa +import pyopencl.tools as cl_tools # noqa +from pyopencl.tools import ( # noqa + pytest_generate_tests_for_pyopencl as pytest_generate_tests) +import pytest # noqa + +import pytato as pt + + +def test_basic_codegen(ctx_factory): + ctx = ctx_factory() + queue = cl.CommandQueue(ctx) + + namespace = pt.Namespace() + x = pt.Placeholder(namespace, "x", (5,), np.int) + prog = pt.generate_loopy(x * x, target=pt.PyOpenCLTarget(queue)) + x_in = np.array([1, 2, 3, 4, 5]) + _, (out,) = prog(x=x_in) + assert (out == x_in * x_in).all() + + +if __name__ == "__main__": + # make sure that import failures get reported, instead of skipping the + # tests. + if len(sys.argv) > 1: + exec(sys.argv[1]) + else: + from pytest import main + main([__file__]) + +# vim: filetype=pyopencl:fdm=marker -- GitLab From 3fed8384da4a00c706c3f031b4e1b6ee223e629f Mon Sep 17 00:00:00 2001 From: Matt Wala Date: Wed, 3 Jun 2020 04:20:26 -0500 Subject: [PATCH 23/76] Fix fold whitespace --- pytato/array.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/pytato/array.py b/pytato/array.py index 54f6e32..0fdc0ad 100644 --- a/pytato/array.py +++ b/pytato/array.py @@ -682,17 +682,15 @@ class Einsum(Array): """ """ - # }}} -# {{{ reshape +# {{{ reshape class Reshape(Array): """ """ - # }}} -- GitLab From 0fe8473051dec156a8a57b50222b56a7e63686c6 Mon Sep 17 00:00:00 2001 From: Matt Wala Date: Wed, 3 Jun 2020 04:29:23 -0500 Subject: [PATCH 24/76] Formatting --- pytato/scalar_expr.py | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/pytato/scalar_expr.py b/pytato/scalar_expr.py index 6755135..df0f691 100644 --- a/pytato/scalar_expr.py +++ b/pytato/scalar_expr.py @@ -114,8 +114,8 @@ def substitute(expression: Any, variable_assigments: Mapping[str, Any]) -> Any: # }}} -def domain_for_shape(dim_names: Tuple[str, ...], shape: Tuple[ScalarExpression, - ...]) -> isl.BasicSet: +def domain_for_shape(dim_names: Tuple[str, ...], + shape: Tuple[ScalarExpression, ...]) -> isl.BasicSet: """Create a :class:`islpy.BasicSet` that expresses an appropriate index domain for an array of (potentially symbolic) shape *shape*. @@ -155,5 +155,4 @@ def domain_for_shape(dim_names: Tuple[str, ...], shape: Tuple[ScalarExpression, return dom - # vim: foldmethod=marker -- GitLab From 79bd4469a75c88dbaf17aa79cde0f732ee9c1005 Mon Sep 17 00:00:00 2001 From: Matt Wala Date: Wed, 3 Jun 2020 04:40:35 -0500 Subject: [PATCH 25/76] Improve comments --- pytato/codegen.py | 10 ++++++++-- 1 file changed, 8 insertions(+), 2 deletions(-) diff --git a/pytato/codegen.py b/pytato/codegen.py index b3bf553..138f1e4 100644 --- a/pytato/codegen.py +++ b/pytato/codegen.py @@ -92,7 +92,10 @@ class GeneratedResult(object): class ArrayResult(GeneratedResult): - """An array expression generated as a :mod:`loopy` array.""" + """An array expression generated as a :mod:`loopy` array. + + See also: :class:`pytato.array.ImplStored`. + """ def __init__(self, name: str, shape: ShapeType, dtype: np.dtype): super().__init__(shape, dtype) self.name = name @@ -107,7 +110,10 @@ class ArrayResult(GeneratedResult): class LoopyExpressionResult(GeneratedResult): - """An array expression generated as a :mod:`loopy` expression.""" + """An array expression generated as a :mod:`loopy` expression. + + See also: :class:`pytato.array.ImplInlined`. + """ def __init__( self, expr: ScalarExpression, shape: ShapeType, dtype: np.dtype): super().__init__(shape, dtype) -- GitLab From 43e290cef0d703d2bc6a7c8cce619d7d2eca0c36 Mon Sep 17 00:00:00 2001 From: Matt Wala Date: Wed, 3 Jun 2020 04:49:11 -0500 Subject: [PATCH 26/76] Remove mapper.py --- pytato/mapper.py | 56 ------------------------------------------------ 1 file changed, 56 deletions(-) delete mode 100644 pytato/mapper.py diff --git a/pytato/mapper.py b/pytato/mapper.py deleted file mode 100644 index afeec20..0000000 --- a/pytato/mapper.py +++ /dev/null @@ -1,56 +0,0 @@ -__copyright__ = """ -Copyright (C) 2020 Andreas Kloeckner -Copyright (C) 2020 Matt Wala -Copyright (C) 2020 Xiaoyu Wei -""" - -__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 typing import Any - - -class Mapper: - pass - - -class IdentityMapper: - """ - Graph transformations subclass this - """ - pass - - -class StringifyMapper: - pass - - -class ToLoopyMapper: - pass - - # {{{ - - def _stringify(self) -> str: - pass - - def _generate_code(self) -> Any: - pass - - # }}} -- GitLab From 368e81f7c5230351f37affeef85e63579821d268 Mon Sep 17 00:00:00 2001 From: Matt Wala Date: Wed, 3 Jun 2020 12:01:15 +0200 Subject: [PATCH 27/76] Formatting --- pytato/stubs/pytools.pyi | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pytato/stubs/pytools.pyi b/pytato/stubs/pytools.pyi index 2fe9afc..c9ee9bc 100644 --- a/pytato/stubs/pytools.pyi +++ b/pytato/stubs/pytools.pyi @@ -7,5 +7,5 @@ def memoize_method(f: T) -> T: ... def is_single_valued(it: Iterable[T]) -> bool: ... class UniqueNameGenerator: - def __init__(self, existing_names: Optional[Set[str]], forced_prefix: str=""): ... + def __init__(self, existing_names: Optional[Set[str]], forced_prefix: str = ""): ... def __call__(self, based_on: str = "") -> str: ... -- GitLab From a6362613a2547bb924c2bf218075f78564ce549d Mon Sep 17 00:00:00 2001 From: Matt Wala Date: Wed, 3 Jun 2020 23:27:38 -0500 Subject: [PATCH 28/76] Remove pytools type stub now that pytools has type annotation support --- pytato/stubs/pytools.pyi | 11 ----------- requirements.txt | 1 + 2 files changed, 1 insertion(+), 11 deletions(-) delete mode 100644 pytato/stubs/pytools.pyi diff --git a/pytato/stubs/pytools.pyi b/pytato/stubs/pytools.pyi deleted file mode 100644 index 2fe9afc..0000000 --- a/pytato/stubs/pytools.pyi +++ /dev/null @@ -1,11 +0,0 @@ -# FIXME: Should be in pytools -from typing import TypeVar, Iterable, Set, Optional - -T = TypeVar("T") - -def memoize_method(f: T) -> T: ... -def is_single_valued(it: Iterable[T]) -> bool: ... - -class UniqueNameGenerator: - def __init__(self, existing_names: Optional[Set[str]], forced_prefix: str=""): ... - def __call__(self, based_on: str = "") -> str: ... diff --git a/requirements.txt b/requirements.txt index 22617c7..8900766 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1 +1,2 @@ +git+https://github.com/inducer/pytools.git git+https://github.com/inducer/loopy.git -- GitLab From 944513fb808b3005f41c3c42524280215b2bd9ae Mon Sep 17 00:00:00 2001 From: Matt Wala Date: Wed, 3 Jun 2020 23:28:58 -0500 Subject: [PATCH 29/76] Fix run-mypy.sh --- run-mypy.sh | 1 - 1 file changed, 1 deletion(-) diff --git a/run-mypy.sh b/run-mypy.sh index 72b4e41..14730b1 100755 --- a/run-mypy.sh +++ b/run-mypy.sh @@ -1,4 +1,3 @@ #! /bin/bash -export MYPYPATH=$MYPYPATH:pytato/stubs/ mypy --strict pytato -- GitLab From fdc7bcd401e0298c805d8e778d34556ba4058e7f Mon Sep 17 00:00:00 2001 From: Matt Wala Date: Thu, 4 Jun 2020 22:59:51 -0500 Subject: [PATCH 30/76] Use numpy.result_type --- pytato/array.py | 8 +------- 1 file changed, 1 insertion(+), 7 deletions(-) diff --git a/pytato/array.py b/pytato/array.py index 0fdc0ad..0240cf5 100644 --- a/pytato/array.py +++ b/pytato/array.py @@ -444,12 +444,6 @@ class Array: def __ne__(self, other: Any) -> bool: return not self.__eq__(other) - def _join_dtypes(self, *args: np.dtype) -> np.dtype: - result = args[0] - for arg in args[1:]: - result = (np.empty(0, dtype=result) + np.empty(0, dtype=arg)).dtype - return result - def _binary_op(self, op: Any, other: Union[Array, Number], @@ -461,7 +455,7 @@ class Array: if self.shape != other.shape: raise ValueError("shapes do not match for binary operator") - dtype = self._join_dtypes(self.dtype, other.dtype) + dtype = np.result_type(self.dtype, other.dtype) # FIXME: If either *self* or *other* is an IndexLambda, its expression # could be folded into the output, producing a fused result. -- GitLab From a31d27492b9d796ef0fbac2875e1a0092eef2825 Mon Sep 17 00:00:00 2001 From: Matt Wala Date: Thu, 4 Jun 2020 23:34:34 -0500 Subject: [PATCH 31/76] transform.py: Don't inherit from pymbolic --- pytato/transform.py | 46 ++++++++++++++++++++++++++++++++++----------- 1 file changed, 35 insertions(+), 11 deletions(-) diff --git a/pytato/transform.py b/pytato/transform.py index 76a1b49..24959a5 100644 --- a/pytato/transform.py +++ b/pytato/transform.py @@ -24,9 +24,7 @@ OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. """ -from typing import Dict - -from pymbolic.mapper import Mapper as MapperBase +from typing import Any, Dict from pytato.array import Array, IndexLambda, Namespace, Output, Placeholder @@ -41,22 +39,46 @@ Transforming Computations """ - # {{{ mapper classes -class Mapper(MapperBase): +class UnsupportedArrayError(ValueError): pass +class Mapper: + def handle_unsupported_array(self, expr: Array, *args: Any, + **kwargs: Any) -> Any: + """Mapper method that is invoked for + :class:`pytato.Array` subclasses for which a mapper + method does not exist in this mapper. + """ + raise UnsupportedArrayError("%s cannot handle expressions of type %s" + % (type(self).__name__, type(expr))) + + def map_foreign(self, expr: Any, *args: Any, **kwargs: Any) -> Any: + raise ValueError("%s encountered invalid foreign object: %s" + % (type(self).__name__, repr(expr))) + + def __call__(self, expr: Array, *args: Any, **kwargs: Any) -> Array: + try: + method = getattr(self, expr.mapper_method) + except AttributeError: + if isinstance(expr, Array): + return self.handle_unsupported_array(expr, *args, **kwargs) + else: + return self.map_foreign(expr, *args, **kwargs) + + return method(expr, *args, **kwargs) + + rec = __call__ + + class CopyMapper(Mapper): namespace: Namespace - def __init__(self, new_namespace: Namespace): - self.namespace = new_namespace - self.cache: Dict[Array, Array] = {} - - def __call__(self, expr: Array) -> Array: - return self.rec(expr) + def __init__(self, namespace): + self.namespace = namespace + self.cache = {} def rec(self, expr: Array) -> Array: if expr in self.cache: @@ -65,6 +87,8 @@ class CopyMapper(Mapper): self.cache[expr] = result return result + __call__ = rec + def map_index_lambda(self, expr: IndexLambda) -> Array: bindings = { name: self.rec(subexpr) -- GitLab From d2cfca23d471fd7f59bd282fac35f225a5fbfc99 Mon Sep 17 00:00:00 2001 From: Matt Wala Date: Thu, 4 Jun 2020 23:51:48 -0500 Subject: [PATCH 32/76] Make Mypy happy with transform.py --- pytato/array.py | 6 +++++- pytato/transform.py | 15 +++++++++------ 2 files changed, 14 insertions(+), 7 deletions(-) diff --git a/pytato/array.py b/pytato/array.py index 0240cf5..2f0aed4 100644 --- a/pytato/array.py +++ b/pytato/array.py @@ -90,7 +90,9 @@ from functools import partialmethod from numbers import Number import operator from dataclasses import dataclass -from typing import Optional, Dict, Any, Mapping, Iterator, Tuple, Union, FrozenSet +from typing import ( + Optional, ClassVar, Dict, Any, Mapping, Iterator, Tuple, Union, + FrozenSet) import numpy as np import pymbolic.primitives as prim @@ -386,6 +388,8 @@ class Array: .. attribute:: ndim """ + mapper_method: ClassVar[str] + def __init__(self, namespace: Namespace, shape: ShapeType, diff --git a/pytato/transform.py b/pytato/transform.py index 24959a5..1e9f161 100644 --- a/pytato/transform.py +++ b/pytato/transform.py @@ -24,7 +24,7 @@ OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. """ -from typing import Any, Dict +from typing import Any, Callable, Dict, no_type_check from pytato.array import Array, IndexLambda, Namespace, Output, Placeholder @@ -59,7 +59,9 @@ class Mapper: raise ValueError("%s encountered invalid foreign object: %s" % (type(self).__name__, repr(expr))) - def __call__(self, expr: Array, *args: Any, **kwargs: Any) -> Array: + def __call__(self, expr: Array, *args: Any, **kwargs: Any) -> Any: + method: Callable[..., Array] + try: method = getattr(self, expr.mapper_method) except AttributeError: @@ -76,18 +78,19 @@ class Mapper: class CopyMapper(Mapper): namespace: Namespace - def __init__(self, namespace): + def __init__(self, namespace: Namespace): self.namespace = namespace - self.cache = {} + self.cache: Dict[Array, Array] = {} - def rec(self, expr: Array) -> Array: + def rec(self, expr: Array) -> Array: # type: ignore[override] if expr in self.cache: return self.cache[expr] result: Array = super().rec(expr) self.cache[expr] = result return result - __call__ = rec + def __call__(self, expr: Array) -> Array: # type: ignore[override] + return self.rec(expr) def map_index_lambda(self, expr: IndexLambda) -> Array: bindings = { -- GitLab From 8e0b07779da372d1620880f85c27c3a24e1760be Mon Sep 17 00:00:00 2001 From: Matt Wala Date: Thu, 4 Jun 2020 23:53:31 -0500 Subject: [PATCH 33/76] flake8 fixes --- pytato/transform.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/pytato/transform.py b/pytato/transform.py index 1e9f161..4582c5d 100644 --- a/pytato/transform.py +++ b/pytato/transform.py @@ -24,7 +24,7 @@ OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. """ -from typing import Any, Callable, Dict, no_type_check +from typing import Any, Callable, Dict from pytato.array import Array, IndexLambda, Namespace, Output, Placeholder @@ -39,6 +39,7 @@ Transforming Computations """ + # {{{ mapper classes class UnsupportedArrayError(ValueError): -- GitLab From 9d3fcdd990bf67c348be2c440a37942e25319cc4 Mon Sep 17 00:00:00 2001 From: Matt Wala Date: Thu, 4 Jun 2020 23:55:44 -0500 Subject: [PATCH 34/76] Move domain_for_shape to codegen.py --- pytato/codegen.py | 51 ++++++++++++++++++++++++++++++++++++++++++- pytato/scalar_expr.py | 42 ----------------------------------- 2 files changed, 50 insertions(+), 43 deletions(-) diff --git a/pytato/codegen.py b/pytato/codegen.py index 138f1e4..7665944 100644 --- a/pytato/codegen.py +++ b/pytato/codegen.py @@ -23,11 +23,12 @@ THE SOFTWARE. """ import collections +import contextlib import dataclasses from typing import Any, Union, Optional, Mapping, Iterator, Dict, Tuple, FrozenSet import typing -import contextlib +import islpy as isl import loopy as lp import numpy as np import pymbolic.primitives as prim @@ -64,6 +65,8 @@ Code Generation Internals .. autoclass:: LoopyExpressionContext .. autoclass:: LoopyExpressionGenMapper +.. autofunction:: domain_for_shape + """ @@ -365,6 +368,52 @@ class LoopyExpressionGenMapper(scalar_expr.IdentityMapper): # }}} +# {{{ utils + +def domain_for_shape( + dim_names: Tuple[str, ...], shape: ShapeType) -> isl.BasicSet: + """Create a :class:`islpy.BasicSet` that expresses an appropriate index domain + for an array of (potentially symbolic) shape *shape*. + + :param dim_names: A tuple of strings, the names of the axes. These become set + dimensions in the returned domain. + + :param shape: A tuple of constant or quasi-affine :mod:`pymbolic` + expressions. The variables in these expressions become parameter + dimensions in the returned set. Must have the same length as + *dim_names*. + """ + assert len(dim_names) == len(shape) + + # Collect parameters. + param_names_set: Set[str] = set() + for sdep in map(get_dependencies, shape): + param_names_set |= sdep + + set_names = sorted(dim_names) + param_names = sorted(param_names_set) + + # Build domain. + dom = isl.BasicSet.universe( + isl.Space.create_from_names(isl.DEFAULT_CONTEXT, + set=set_names, + params=param_names)) + + # Add constraints. + from loopy.symbolic import aff_from_expr + affs = isl.affs_from_space(dom.space) + + for iname, dim in zip(dim_names, shape): + dom &= affs[0].le_set(affs[iname]) + dom &= affs[iname].lt_set(aff_from_expr(dom.space, dim)) + + dom, = dom.get_basic_sets() + + return dom + +# }}} + + def _promote_named_arrays_to_outputs(arrays: DictOfNamedArrays) -> Namespace: # Turns named arrays into Output nodes, returning a new namespace. copy_mapper = pytato.transform.CopyMapper(Namespace()) diff --git a/pytato/scalar_expr.py b/pytato/scalar_expr.py index df0f691..1d2810f 100644 --- a/pytato/scalar_expr.py +++ b/pytato/scalar_expr.py @@ -53,7 +53,6 @@ Scalar Expressions .. autofunction:: parse .. autofunction:: get_dependencies .. autofunction:: substitute -.. autofunction:: domain_for_shape """ @@ -114,45 +113,4 @@ def substitute(expression: Any, variable_assigments: Mapping[str, Any]) -> Any: # }}} -def domain_for_shape(dim_names: Tuple[str, ...], - shape: Tuple[ScalarExpression, ...]) -> isl.BasicSet: - """Create a :class:`islpy.BasicSet` that expresses an appropriate index domain - for an array of (potentially symbolic) shape *shape*. - - :param dim_names: A tuple of strings, the names of the axes. These become set - dimensions in the returned domain. - - :param shape: A tuple of constant or quasi-affine :mod:`pymbolic` - expressions. The variables in these expressions become parameter - dimensions in the returned set. Must have the same length as - *dim_names*. - """ - assert len(dim_names) == len(shape) - - # Collect parameters. - param_names_set: Set[str] = set() - for sdep in map(get_dependencies, shape): - param_names_set |= sdep - - set_names = sorted(dim_names) - param_names = sorted(param_names_set) - - # Build domain. - dom = isl.BasicSet.universe( - isl.Space.create_from_names(isl.DEFAULT_CONTEXT, - set=set_names, - params=param_names)) - - # Add constraints. - from loopy.symbolic import aff_from_expr - affs = isl.affs_from_space(dom.space) - - for iname, dim in zip(dim_names, shape): - dom &= affs[0].le_set(affs[iname]) - dom &= affs[iname].lt_set(aff_from_expr(dom.space, dim)) - - dom, = dom.get_basic_sets() - - return dom - # vim: foldmethod=marker -- GitLab From 0f80830f00fc9b93213f07d91045b3aad3d2dc51 Mon Sep 17 00:00:00 2001 From: Matt Wala Date: Thu, 4 Jun 2020 23:57:37 -0500 Subject: [PATCH 35/76] Remove stray testing comment --- test/test_codegen.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/test/test_codegen.py b/test/test_codegen.py index e93200d..ee2d65d 100755 --- a/test/test_codegen.py +++ b/test/test_codegen.py @@ -49,8 +49,6 @@ def test_basic_codegen(ctx_factory): if __name__ == "__main__": - # make sure that import failures get reported, instead of skipping the - # tests. if len(sys.argv) > 1: exec(sys.argv[1]) else: -- GitLab From f22ec1dc6bca294ec657594a9a4b6b6a258a970b Mon Sep 17 00:00:00 2001 From: Matt Wala Date: Fri, 5 Jun 2020 00:09:32 -0500 Subject: [PATCH 36/76] Fix domain_for_shape invocations --- pytato/codegen.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/pytato/codegen.py b/pytato/codegen.py index 7665944..e482ed9 100644 --- a/pytato/codegen.py +++ b/pytato/codegen.py @@ -240,7 +240,7 @@ class CodeGenMapper(pytato.transform.Mapper): inames = tuple( state.var_name_gen(f"{expr.name}_dim{d}") for d in range(expr.ndim)) - domain = scalar_expr.domain_for_shape(inames, expr.shape) + domain = domain_for_shape(inames, expr.shape) arg = lp.GlobalArg(expr.name, shape=expr.shape, @@ -387,7 +387,7 @@ def domain_for_shape( # Collect parameters. param_names_set: Set[str] = set() - for sdep in map(get_dependencies, shape): + for sdep in map(scalar_expr.get_dependencies, shape): param_names_set |= sdep set_names = sorted(dim_names) -- GitLab From dbe3624970e3308f9f402128ee08b09ea5f875f1 Mon Sep 17 00:00:00 2001 From: Matt Wala Date: Fri, 5 Jun 2020 00:24:32 -0500 Subject: [PATCH 37/76] Kill Output node --- pytato/array.py | 73 ++++++-------------------- pytato/codegen.py | 122 +++++++++++++++++++++----------------------- pytato/transform.py | 6 +-- 3 files changed, 77 insertions(+), 124 deletions(-) diff --git a/pytato/array.py b/pytato/array.py index 2f0aed4..ec2cbd2 100644 --- a/pytato/array.py +++ b/pytato/array.py @@ -67,7 +67,6 @@ Built-in Expression Nodes .. autoclass:: Reshape .. autoclass:: DataWrapper .. autoclass:: Placeholder -.. autoclass:: Output .. autoclass:: LoopyFunction User-Facing Node Creation @@ -728,7 +727,23 @@ class DataWrapper(Array): # {{{ placeholder -class _ArgLike(Array): +class Placeholder(Array): + """ + A named placeholder for an array whose concrete value + is supplied by the user during evaluation. + + .. attribute:: name + + The name by which a value is supplied + for the placeholder once computation begins. + + .. note:: + + Modifying :class:`Placeholder` tags is not supported after + creation. + """ + mapper_method = "map_placeholder" + def __init__(self, namespace: Namespace, name: str, @@ -766,60 +781,6 @@ class _ArgLike(Array): def without_tag(self, tag: Tag, verify_existence: bool = True) -> Array: raise ValueError("Cannot modify tags") - -class Placeholder(_ArgLike): - """ - A named placeholder for an array whose concrete value - is supplied by the user during evaluation. - - .. attribute:: name - - The name by which a value is supplied - for the placeholder once computation begins. - - .. note:: - - Modifying :class:`Placeholder` tags is not supported after - creation. - """ - - mapper_method = "map_placeholder" - -# }}} - - -# {{{ output - -class Output(_ArgLike): - """A named output of the computation. - - .. attribute:: name - - The name of the output array. - - .. attribute:: array - - The :class:`Array` value that is output. - - .. note:: - - Modifying :class:`Output` tags is not supported after creation. - """ - - mapper_method = "map_output" - - def __init__(self, - namespace: Namespace, - name: str, - array: Array, - tags: Optional[TagsType] = None): - super().__init__(namespace=namespace, - name=name, - shape=array.shape, - dtype=array.dtype, - tags=tags) - self.array = array - # }}} diff --git a/pytato/codegen.py b/pytato/codegen.py index e482ed9..1df97a5 100644 --- a/pytato/codegen.py +++ b/pytato/codegen.py @@ -35,7 +35,7 @@ import pymbolic.primitives as prim import pytools from pytato.array import ( - Array, DictOfNamedArrays, Placeholder, Output, Namespace, ShapeType, + Array, DictOfNamedArrays, Placeholder, Namespace, ShapeType, IndexLambda) from pytato.program import BoundProgram, Target, PyOpenCLTarget import pytato.scalar_expr as scalar_expr @@ -66,6 +66,7 @@ Code Generation Internals .. autoclass:: LoopyExpressionGenMapper .. autofunction:: domain_for_shape +.. autofunction:: add_output """ @@ -228,51 +229,6 @@ class CodeGenMapper(pytato.transform.Mapper): state.results[expr] = result return result - def map_output(self, expr: Output, state: CodeGenState) -> GeneratedResult: - if expr in state.results: - return state.results[expr] - - # FIXE: Scalar outputs are not supported yet. - assert expr.shape != () - - inner_result = self.rec(expr.array, state) - - inames = tuple( - state.var_name_gen(f"{expr.name}_dim{d}") - for d in range(expr.ndim)) - domain = domain_for_shape(inames, expr.shape) - - arg = lp.GlobalArg(expr.name, - shape=expr.shape, - dtype=expr.dtype, - order="C", - is_output_only=True) - - indices = tuple(prim.Variable(iname) for iname in inames) - context = state.make_expression_context() - copy_expr = inner_result.to_loopy_expression(indices, context) - - # TODO: Contextual data not supported yet. - assert not context.reduction_bounds - assert not context.depends_on - - from loopy.kernel.instruction import make_assignment - insn = make_assignment((prim.Variable(expr.name)[indices], ), - copy_expr, - id=state.insn_id_gen(f"{expr.name}_copy"), - within_inames=frozenset(inames), - depends_on=context.depends_on) - - kernel = state.kernel - kernel = kernel.copy(args=kernel.args + [arg], - instructions=kernel.instructions + [insn], - domains=kernel.domains + [domain]) - state.update_kernel(kernel) - - result = ArrayResult(expr.name, expr.dtype, expr.shape) - state.results[expr] = result - return result - def map_index_lambda(self, expr: IndexLambda, state: CodeGenState) -> GeneratedResult: if expr in state.results: @@ -411,19 +367,50 @@ def domain_for_shape( return dom -# }}} +def add_output(name: str, expr: Array, state: CodeGenState, + mapper: CodeGenMapper) -> None: + """Add an output argument to the kernel. + """ + # FIXE: Scalar outputs are not supported yet. + assert expr.shape != () -def _promote_named_arrays_to_outputs(arrays: DictOfNamedArrays) -> Namespace: - # Turns named arrays into Output nodes, returning a new namespace. - copy_mapper = pytato.transform.CopyMapper(Namespace()) - result = pytato.transform.copy_namespace(arrays.namespace, copy_mapper) + result = mapper(expr, state) + name = state.var_name_gen(name) - name_gen = pytools.UniqueNameGenerator(set(result)) - for name, val in arrays.items(): - Output(result, name_gen(name), copy_mapper(val)) + inames = tuple( + state.var_name_gen(f"{name}_dim{d}") + for d in range(expr.ndim)) + domain = domain_for_shape(inames, expr.shape) + + arg = lp.GlobalArg(name, + shape=expr.shape, + dtype=expr.dtype, + order="C", + is_output_only=True) + + indices = tuple(prim.Variable(iname) for iname in inames) + context = state.make_expression_context() + copy_expr = result.to_loopy_expression(indices, context) + + # TODO: Contextual data not supported yet. + assert not context.reduction_bounds + assert not context.depends_on + + from loopy.kernel.instruction import make_assignment + insn = make_assignment((prim.Variable(name)[indices],), + copy_expr, + id=state.insn_id_gen(f"{name}_copy"), + within_inames=frozenset(inames), + depends_on=context.depends_on) + + kernel = state.kernel + kernel = kernel.copy(args=kernel.args + [arg], + instructions=kernel.instructions + [insn], + domains=kernel.domains + [domain]) + state.update_kernel(kernel) - return result +# }}} def generate_loopy( @@ -440,17 +427,22 @@ def generate_loopy( :returns: A wrapped generated :mod:`loopy` kernel """ - # {{{ get namespace + # {{{ get namespace and outputs - if isinstance(result_or_namespace, Array): - result_or_namespace = DictOfNamedArrays({"out": result_or_namespace}) + outputs: DictOfNamedArrays + namespace: Namespace - if isinstance(result_or_namespace, DictOfNamedArrays): - result_or_namespace = _promote_named_arrays_to_outputs( - result_or_namespace) + if isinstance(result_or_namespace, Array): + outputs = DictOfNamedArrays({"out": result_or_namespace}) + namespace = outputs.namespace + elif isinstance(result_or_namespace, DictOfNamedArrays): + outputs = result_or_namespace + namespace = outputs.namespace + else: + assert isinstance(result_or_namespace, Namespace) + outputs = DictOfNamedArrays() + namespace = result_or_namespace - assert isinstance(result_or_namespace, Namespace) - namespace = result_or_namespace del result_or_namespace # }}} @@ -472,4 +464,8 @@ def generate_loopy( for name, val in namespace.items(): _ = mapper(val, state) + # Generate code for outputs. + for name, expr in outputs.items(): + add_output(name, expr, state, mapper) + return target.bind_program(program=state.kernel, bound_arguments=dict()) diff --git a/pytato/transform.py b/pytato/transform.py index 4582c5d..2b2846d 100644 --- a/pytato/transform.py +++ b/pytato/transform.py @@ -26,7 +26,7 @@ THE SOFTWARE. from typing import Any, Callable, Dict -from pytato.array import Array, IndexLambda, Namespace, Output, Placeholder +from pytato.array import Array, IndexLambda, Namespace, Placeholder __doc__ = """ .. currentmodule:: pytato.transform @@ -107,10 +107,6 @@ class CopyMapper(Mapper): return Placeholder(self.namespace, expr.name, expr.shape, expr.dtype, expr.tags) - def map_output(self, expr: Output) -> Array: - return Output(self.namespace, expr.name, self.rec(expr.array), - expr.tags) - # }}} -- GitLab From 984ccc31e222d57a726e2067f89891b985ba460a Mon Sep 17 00:00:00 2001 From: Matt Wala Date: Fri, 5 Jun 2020 00:26:21 -0500 Subject: [PATCH 38/76] Make fields registrable --- pytato/array.py | 54 +++++++++++++++++++------------------------------ 1 file changed, 21 insertions(+), 33 deletions(-) diff --git a/pytato/array.py b/pytato/array.py index ec2cbd2..99dedf7 100644 --- a/pytato/array.py +++ b/pytato/array.py @@ -388,6 +388,9 @@ class Array: """ mapper_method: ClassVar[str] + # A tuple of field names. Fields must be equality comparable and + # hashable. Dicts of hashable keys and values are also permitted. + fields: ClassVar[Tuple[str, ...]] = ("shape", "dtype", "tags") def __init__(self, namespace: Namespace, @@ -439,10 +442,23 @@ class Array: @memoize_method def __hash__(self) -> int: - raise NotImplementedError + attrs = [] + for field in self.fields: + attr = getattr(self, field) + if isinstance(attr, dict): + attr = frozenset(attr.items()) + attrs.append(attr) + return hash(tuple(attrs)) def __eq__(self, other: Any) -> bool: - raise NotImplementedError + if self is other: + return True + return ( + isinstance(other, type(self)) + and self.namespace is other.namespace + and all( + getattr(self, field) == getattr(other, field) + for field in self.fields)) def __ne__(self, other: Any) -> bool: return not self.__eq__(other) @@ -603,7 +619,7 @@ class IndexLambda(Array): .. automethod:: is_reference """ - + fields = Array.fields + ("expr", "bindings") mapper_method = "map_index_lambda" def __init__(self, @@ -653,23 +669,6 @@ class IndexLambda(Array): return True - @memoize_method - def __hash__(self) -> int: - return hash((self.expr, self.shape, self.dtype, - frozenset(self.bindings.items()), self.tags)) - - def __eq__(self, other: object) -> bool: - if self is other: - return True - - return (isinstance(other, IndexLambda) - and self.namespace is other.namespace - and self.expr == other.expr - and self.shape == other.shape - and self.dtype == other.dtype - and self.bindings == other.bindings - and self.tags == other.tags) - # }}} @@ -743,7 +742,8 @@ class Placeholder(Array): creation. """ mapper_method = "map_placeholder" - + fields = Array.fields + ("name",) + def __init__(self, namespace: Namespace, name: str, @@ -763,18 +763,6 @@ class Placeholder(Array): self.name = name - @memoize_method - def __hash__(self) -> int: - return hash((self.name, )) - - def __eq__(self, other: object) -> bool: - if self is other: - return True - # Uniquely identified by name. - return (isinstance(other, _ArgLike) - and self.namespace is other.namespace - and self.name == other.name) - def tagged(self, tag: Tag) -> Array: raise ValueError("Cannot modify tags") -- GitLab From 1d2b5636ae341d5b10db8fcbf5b75b02d1cd922a Mon Sep 17 00:00:00 2001 From: Matt Wala Date: Fri, 5 Jun 2020 00:43:13 -0500 Subject: [PATCH 39/76] Fix attribute name --- pytato/codegen.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pytato/codegen.py b/pytato/codegen.py index 1df97a5..f04cabd 100644 --- a/pytato/codegen.py +++ b/pytato/codegen.py @@ -262,7 +262,7 @@ class LoopyExpressionContext(object): The :class:`CodeGenState`. - .. attribute:: _depends_on + .. attribute:: depends_on The set of dependencies associated with the expression. -- GitLab From 66c3a5c66994a17c016a52295671c3c84a17460f Mon Sep 17 00:00:00 2001 From: Matt Wala Date: Fri, 5 Jun 2020 00:44:14 -0500 Subject: [PATCH 40/76] Fix mypy issues --- pytato/codegen.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/pytato/codegen.py b/pytato/codegen.py index f04cabd..6d78d81 100644 --- a/pytato/codegen.py +++ b/pytato/codegen.py @@ -25,7 +25,9 @@ THE SOFTWARE. import collections import contextlib import dataclasses -from typing import Any, Union, Optional, Mapping, Iterator, Dict, Tuple, FrozenSet +from typing import ( + Any, Union, Optional, Mapping, Iterator, Dict, Tuple, FrozenSet, + Set) import typing import islpy as isl @@ -440,7 +442,7 @@ def generate_loopy( namespace = outputs.namespace else: assert isinstance(result_or_namespace, Namespace) - outputs = DictOfNamedArrays() + outputs = DictOfNamedArrays({}) namespace = result_or_namespace del result_or_namespace -- GitLab From 3d3b9597744ff6e20d41cfa870e8ee184c6edf4f Mon Sep 17 00:00:00 2001 From: Matt Wala Date: Fri, 5 Jun 2020 01:00:38 -0500 Subject: [PATCH 41/76] flake8 fixes --- pytato/array.py | 2 +- pytato/codegen.py | 3 +-- pytato/scalar_expr.py | 3 +-- 3 files changed, 3 insertions(+), 5 deletions(-) diff --git a/pytato/array.py b/pytato/array.py index 99dedf7..4895931 100644 --- a/pytato/array.py +++ b/pytato/array.py @@ -743,7 +743,7 @@ class Placeholder(Array): """ mapper_method = "map_placeholder" fields = Array.fields + ("name",) - + def __init__(self, namespace: Namespace, name: str, diff --git a/pytato/codegen.py b/pytato/codegen.py index 6d78d81..64b8b00 100644 --- a/pytato/codegen.py +++ b/pytato/codegen.py @@ -34,7 +34,6 @@ import islpy as isl import loopy as lp import numpy as np import pymbolic.primitives as prim -import pytools from pytato.array import ( Array, DictOfNamedArrays, Placeholder, Namespace, ShapeType, @@ -382,7 +381,7 @@ def add_output(name: str, expr: Array, state: CodeGenState, inames = tuple( state.var_name_gen(f"{name}_dim{d}") - for d in range(expr.ndim)) + for d in range(expr.ndim)) domain = domain_for_shape(inames, expr.shape) arg = lp.GlobalArg(name, diff --git a/pytato/scalar_expr.py b/pytato/scalar_expr.py index 1d2810f..a09e9db 100644 --- a/pytato/scalar_expr.py +++ b/pytato/scalar_expr.py @@ -25,9 +25,8 @@ THE SOFTWARE. """ from numbers import Number -from typing import Any, Union, Mapping, FrozenSet, Set, Tuple +from typing import Any, Union, Mapping, FrozenSet -import islpy as isl from pymbolic.mapper import (WalkMapper as WalkMapperBase, IdentityMapper as IdentityMapperBase) from pymbolic.mapper.substitutor import (SubstitutionMapper as -- GitLab From 4074c3fd404ff4a640af67f9b0891c6c48445bed Mon Sep 17 00:00:00 2001 From: Matt Wala Date: Fri, 5 Jun 2020 01:24:31 -0500 Subject: [PATCH 42/76] Remove MANIFEST.in --- MANIFEST.in | 1 - 1 file changed, 1 deletion(-) delete mode 100644 MANIFEST.in diff --git a/MANIFEST.in b/MANIFEST.in deleted file mode 100644 index d887c9f..0000000 --- a/MANIFEST.in +++ /dev/null @@ -1 +0,0 @@ -recursive-include pytato *.pyi -- GitLab From 3dcc50b0e5cc37a1577bed47fa53c4d4ae8765b4 Mon Sep 17 00:00:00 2001 From: Matt Wala Date: Fri, 12 Jun 2020 00:03:04 +0200 Subject: [PATCH 43/76] Apply suggestion to pytato/array.py --- pytato/array.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pytato/array.py b/pytato/array.py index 4895931..6afad12 100644 --- a/pytato/array.py +++ b/pytato/array.py @@ -472,7 +472,7 @@ class Array: raise NotImplementedError else: if self.shape != other.shape: - raise ValueError("shapes do not match for binary operator") + raise NotImplementedError("broadcasting") dtype = np.result_type(self.dtype, other.dtype) -- GitLab From c5b994183d79fa85232d8fe224f6cd942ea00d90 Mon Sep 17 00:00:00 2001 From: Matt Wala Date: Fri, 12 Jun 2020 00:03:14 +0200 Subject: [PATCH 44/76] Apply suggestion to pytato/array.py --- pytato/array.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/pytato/array.py b/pytato/array.py index 6afad12..5b457aa 100644 --- a/pytato/array.py +++ b/pytato/array.py @@ -735,6 +735,8 @@ class Placeholder(Array): The name by which a value is supplied for the placeholder once computation begins. + The name is also implicitly :meth:`~Namespace.assign`\ ed + in the :class:`Namespace`. .. note:: -- GitLab From 14ab7ed2688d7ba5788a344d57e567d0d0048a02 Mon Sep 17 00:00:00 2001 From: Matt Wala Date: Fri, 12 Jun 2020 00:03:18 +0200 Subject: [PATCH 45/76] Apply suggestion to pytato/array.py --- pytato/array.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pytato/array.py b/pytato/array.py index 5b457aa..ec6eaae 100644 --- a/pytato/array.py +++ b/pytato/array.py @@ -755,7 +755,7 @@ class Placeholder(Array): if name is None: raise ValueError("Must have explicit name") - # Reserve the name, prevent others from using it. + # Publish our name to the namespace namespace.assign(name, self) super().__init__(namespace=namespace, -- GitLab From 26b2dca73fbc58ab76af763b7c44b420d1c52cf1 Mon Sep 17 00:00:00 2001 From: Matt Wala Date: Fri, 12 Jun 2020 00:03:54 +0200 Subject: [PATCH 46/76] Apply suggestion to pytato/array.py --- pytato/array.py | 1 + 1 file changed, 1 insertion(+) diff --git a/pytato/array.py b/pytato/array.py index ec6eaae..4594ebd 100644 --- a/pytato/array.py +++ b/pytato/array.py @@ -442,6 +442,7 @@ class Array: @memoize_method def __hash__(self) -> int: + # TODO Cache computed hash value attrs = [] for field in self.fields: attr = getattr(self, field) -- GitLab From 38ebc911cfc7370f18facb136a5e621f642a0265 Mon Sep 17 00:00:00 2001 From: Matt Wala Date: Fri, 12 Jun 2020 00:07:13 +0200 Subject: [PATCH 47/76] Fix escaping --- pytato/array.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pytato/array.py b/pytato/array.py index 4594ebd..70101cb 100644 --- a/pytato/array.py +++ b/pytato/array.py @@ -728,7 +728,7 @@ class DataWrapper(Array): # {{{ placeholder class Placeholder(Array): - """ + r""" A named placeholder for an array whose concrete value is supplied by the user during evaluation. -- GitLab From c5de67ea9477b9a587bff7ed78045124300e7640 Mon Sep 17 00:00:00 2001 From: Matt Wala Date: Thu, 11 Jun 2020 17:44:04 -0500 Subject: [PATCH 48/76] Array: Don't store shape and dtype, implement _SuppliedShapeAndDtypeMixin --- pytato/array.py | 58 ++++++++++++++++++++++++++++++++++++------------- 1 file changed, 43 insertions(+), 15 deletions(-) diff --git a/pytato/array.py b/pytato/array.py index 70101cb..73cac6f 100644 --- a/pytato/array.py +++ b/pytato/array.py @@ -91,7 +91,7 @@ import operator from dataclasses import dataclass from typing import ( Optional, ClassVar, Dict, Any, Mapping, Iterator, Tuple, Union, - FrozenSet) + FrozenSet, Protocol) import numpy as np import pymbolic.primitives as prim @@ -394,27 +394,23 @@ class Array: def __init__(self, namespace: Namespace, - shape: ShapeType, - dtype: np.dtype, tags: Optional[TagsType] = None): if tags is None: tags = frozenset() self.namespace = namespace self.tags = tags - self._shape = shape - self._dtype = dtype def copy(self, **kwargs: Any) -> Array: raise NotImplementedError @property def shape(self) -> ShapeType: - return self._shape + raise NotImplementedError @property def dtype(self) -> np.dtype: - return self._dtype + raise NotImplementedError def named(self, name: str) -> Array: return self.namespace.ref(self.namespace.assign(name, self)) @@ -504,6 +500,29 @@ class Array: __rmul__ = partialmethod(_binary_op, operator.mul, reverse=True) +class _SuppliedShapeAndDtypeMixin(object): + """A mixin class for when an array must store its own *shape* and *dtype*, + rather than when it can derive them easily from inputs. + """ + + def __init__(self, + namespace: Namespace, + shape: ShapeType, + dtype: np.dtype, + **kwargs: Any): + # https://github.com/python/mypy/issues/5887 + super().__init__(namespace, **kwargs) # type: ignore + self._shape = shape + self._dtype = dtype + + @property + def shape(self) -> ShapeType: + return self._shape + + @property + def dtype(self) -> np.dtype: + return self._dtype + # }}} @@ -595,7 +614,7 @@ class DictOfNamedArrays(Mapping[str, Array]): # {{{ index lambda -class IndexLambda(Array): +class IndexLambda(_SuppliedShapeAndDtypeMixin, Array): """ .. attribute:: expr @@ -693,6 +712,11 @@ class Reshape(Array): # {{{ data wrapper +class DataInterface(Protocol): + shape: ShapeType + dtype: np.dtype + + class DataWrapper(Array): # TODO: Name? """ @@ -711,23 +735,27 @@ class DataWrapper(Array): this array may not be updated in-place. """ - # TODO: not really Any data def __init__(self, namespace: Namespace, - data: Any, + data: DataInterface, tags: Optional[TagsType] = None): - super().__init__(namespace, - shape=data.shape, - dtype=data.dtype, - tags=tags) + super().__init__(namespace, tags=tags) self.data = data + @property + def shape(self) -> ShapeType: + return self.data.shape + + @property + def dtype(self) -> np.dtype: + return self.data.dtype + # }}} # {{{ placeholder -class Placeholder(Array): +class Placeholder(_SuppliedShapeAndDtypeMixin, Array): r""" A named placeholder for an array whose concrete value is supplied by the user during evaluation. -- GitLab From cf916437c71922420e99a49db73ea8213df0b196 Mon Sep 17 00:00:00 2001 From: Matt Wala Date: Thu, 11 Jun 2020 22:55:24 -0500 Subject: [PATCH 49/76] Remove Namespace.copy() --- pytato/array.py | 3 +-- pytato/transform.py | 56 +++------------------------------------------ 2 files changed, 4 insertions(+), 55 deletions(-) diff --git a/pytato/array.py b/pytato/array.py index 73cac6f..737ea46 100644 --- a/pytato/array.py +++ b/pytato/array.py @@ -180,8 +180,7 @@ class Namespace(Mapping[str, "Array"]): return len(self._symbol_table) def copy(self) -> Namespace: - from pytato.transform import CopyMapper, copy_namespace - return copy_namespace(self, CopyMapper(Namespace())) + raise NotImplementedError def assign(self, name: str, value: Array) -> str: """Declare a new array. diff --git a/pytato/transform.py b/pytato/transform.py index 2b2846d..1354665 100644 --- a/pytato/transform.py +++ b/pytato/transform.py @@ -24,9 +24,9 @@ OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. """ -from typing import Any, Callable, Dict +from typing import Any, Callable -from pytato.array import Array, IndexLambda, Namespace, Placeholder +from pytato.array import Array __doc__ = """ .. currentmodule:: pytato.transform @@ -34,8 +34,7 @@ __doc__ = """ Transforming Computations ------------------------- -.. autoclass:: CopyMapper -.. autofunction:: copy_namespace +.. autoclass:: Mapper """ @@ -75,56 +74,7 @@ class Mapper: rec = __call__ - -class CopyMapper(Mapper): - namespace: Namespace - - def __init__(self, namespace: Namespace): - self.namespace = namespace - self.cache: Dict[Array, Array] = {} - - def rec(self, expr: Array) -> Array: # type: ignore[override] - if expr in self.cache: - return self.cache[expr] - result: Array = super().rec(expr) - self.cache[expr] = result - return result - - def __call__(self, expr: Array) -> Array: # type: ignore[override] - return self.rec(expr) - - def map_index_lambda(self, expr: IndexLambda) -> Array: - bindings = { - name: self.rec(subexpr) - for name, subexpr in expr.bindings.items()} - return IndexLambda(self.namespace, - expr=expr.expr, - shape=expr.shape, - dtype=expr.dtype, - bindings=bindings) - - def map_placeholder(self, expr: Placeholder) -> Array: - return Placeholder(self.namespace, expr.name, expr.shape, expr.dtype, - expr.tags) - # }}} -# {{{ mapper frontends - -def copy_namespace(namespace: Namespace, copy_mapper: CopyMapper) -> Namespace: - """Copy the elements of *namespace* into a new namespace. - - :param namespace: The source namespace - :param copy_mapper: A mapper that performs copies into a new namespace - :returns: The new namespace - """ - for name, val in namespace.items(): - mapped_val = copy_mapper(val) - if name not in copy_mapper.namespace: - copy_mapper.namespace.assign(name, mapped_val) - return copy_mapper.namespace - -# }}} - # vim: foldmethod=marker -- GitLab From 72531d5e9a3796a7dc9dbf19fabfe990acc36a8a Mon Sep 17 00:00:00 2001 From: Andreas Kloeckner Date: Sun, 14 Jun 2020 23:51:24 -0500 Subject: [PATCH 50/76] Codegen doc tweaks --- pytato/codegen.py | 23 +++++++++++------------ 1 file changed, 11 insertions(+), 12 deletions(-) diff --git a/pytato/codegen.py b/pytato/codegen.py index 64b8b00..4fc0e7d 100644 --- a/pytato/codegen.py +++ b/pytato/codegen.py @@ -143,7 +143,7 @@ class SubstitutionRuleResult(GeneratedResult): @dataclasses.dataclass(init=True, repr=False, eq=False) class CodeGenState: - """Data threaded through :class:`CodeGenMapper`. + """A container for data kept by :class:`CodeGenMapper`. .. attribute:: namespace @@ -151,11 +151,12 @@ class CodeGenState: .. attribute:: kernel - The partial kernel + The partial :class:`loopy.LoopKernel` being built. .. attribute:: results - A mapping from arrays to code generation results + A mapping from :class:`pytato.array.Array` instances to + instances of :class:`GeneratedResult`. .. attribute:: var_name_gen .. attribute:: insn_id_gen @@ -257,7 +258,8 @@ ReductionBounds = Dict[str, Tuple[ScalarExpression, ScalarExpression]] class LoopyExpressionContext(object): """Contextual data for generating :mod:`loopy` expressions. - This data is threaded through :class:`LoopyExpressionGenMapper`. + This data is passed through :class:`LoopyExpressionGenMapper` + via arguments. .. attribute:: state @@ -265,7 +267,8 @@ class LoopyExpressionContext(object): .. attribute:: depends_on - The set of dependencies associated with the expression. + The set of statement IDs that need to be included in + :attr:`loopy.kernel.data.instruction.InstructionBase.depends_on`. .. attribute:: reduction_bounds @@ -417,14 +420,10 @@ def add_output(name: str, expr: Array, state: CodeGenState, def generate_loopy( result_or_namespace: Union[Namespace, Array, DictOfNamedArrays], target: Optional[Target] = None) -> BoundProgram: - """Code generation entry point. + r"""Code generation entry point. - :param result_or_namespace: Either a :class:`pytato.Namespace`, a single - :class:`pytato.Array`, or a :class:`pytato.DictOfNamedArrays`. In the - latter two cases, code generation treats the node(s) as outputs of the - computation. - - :param target: The target for code generation + :param result_or_namespace: For :term:`array result`\ s, code generation treats + the node(s) as outputs of the computation. :returns: A wrapped generated :mod:`loopy` kernel """ -- GitLab From 05fd38fa15348942ab83c90bbb58f14e6e3d518e Mon Sep 17 00:00:00 2001 From: Matt Wala Date: Mon, 15 Jun 2020 13:20:14 -0500 Subject: [PATCH 51/76] Use a dict comprehension --- pytato/codegen.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pytato/codegen.py b/pytato/codegen.py index 64b8b00..83145d6 100644 --- a/pytato/codegen.py +++ b/pytato/codegen.py @@ -129,7 +129,7 @@ class LoopyExpressionResult(GeneratedResult): context: LoopyExpressionContext) -> ScalarExpression: return scalar_expr.substitute( self.expr, - dict(zip((f"_{d}" for d in range(self.ndim)), indices))) + {f"_{d}": i for d, i in zip(range(self.ndim), indices)}) class SubstitutionRuleResult(GeneratedResult): -- GitLab From 44455cae90ad86229f23fcc7ab88b3c691e9efac Mon Sep 17 00:00:00 2001 From: Matt Wala Date: Mon, 15 Jun 2020 13:26:47 -0500 Subject: [PATCH 52/76] s/get_var_name_generator/get_instruction_id_generator/ --- pytato/codegen.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pytato/codegen.py b/pytato/codegen.py index 83145d6..d8ec7fe 100644 --- a/pytato/codegen.py +++ b/pytato/codegen.py @@ -175,7 +175,7 @@ class CodeGenState: def __post_init__(self) -> None: self.var_name_gen = self._kernel.get_var_name_generator() - self.insn_id_gen = self._kernel.get_var_name_generator() + self.insn_id_gen = self._kernel.get_instruction_id_generator() @property def kernel(self) -> lp.LoopKernel: -- GitLab From 94c989ca8aefa3b97c80aeea5c14fed39a2656b0 Mon Sep 17 00:00:00 2001 From: Matt Wala Date: Mon, 15 Jun 2020 13:33:43 -0500 Subject: [PATCH 53/76] s/context/expr_context/ --- pytato/codegen.py | 34 +++++++++++++++++----------------- 1 file changed, 17 insertions(+), 17 deletions(-) diff --git a/pytato/codegen.py b/pytato/codegen.py index 1558b52..fc45042 100644 --- a/pytato/codegen.py +++ b/pytato/codegen.py @@ -91,7 +91,7 @@ class GeneratedResult(object): return len(self.shape) def to_loopy_expression(self, indices: SymbolicIndex, - context: LoopyExpressionContext) -> ScalarExpression: + expr_context: LoopyExpressionContext) -> ScalarExpression: """Return a :mod:`loopy` expression for this result.""" raise NotImplementedError @@ -107,7 +107,7 @@ class ArrayResult(GeneratedResult): # TODO: Handle dependencies. def to_loopy_expression(self, indices: SymbolicIndex, - context: LoopyExpressionContext) -> ScalarExpression: + expr_context: LoopyExpressionContext) -> ScalarExpression: if indices == (): return prim.Variable(self.name) else: @@ -126,7 +126,7 @@ class LoopyExpressionResult(GeneratedResult): # TODO: Handle dependencies and reduction domains. def to_loopy_expression(self, indices: SymbolicIndex, - context: LoopyExpressionContext) -> ScalarExpression: + expr_context: LoopyExpressionContext) -> ScalarExpression: return scalar_expr.substitute( self.expr, {f"_{d}": i for d, i in zip(range(self.ndim), indices)}) @@ -306,24 +306,24 @@ class LoopyExpressionGenMapper(scalar_expr.IdentityMapper): self.codegen_mapper = codegen_mapper def __call__(self, expr: ScalarExpression, - context: LoopyExpressionContext) -> ScalarExpression: - return self.rec(expr, context) + expr_context: LoopyExpressionContext) -> ScalarExpression: + return self.rec(expr, expr_context) def map_subscript(self, expr: prim.Subscript, - context: LoopyExpressionContext) -> ScalarExpression: + expr_context: LoopyExpressionContext) -> ScalarExpression: assert isinstance(expr.aggregate, prim.Variable) result: GeneratedResult = self.codegen_mapper( - context.namespace[expr.aggregate.name], context.state) - return result.to_loopy_expression(expr.index, context) + expr_context.namespace[expr.aggregate.name], expr_context.state) + return result.to_loopy_expression(expr.index, expr_context) # TODO: map_reduction() def map_variable(self, expr: prim.Variable, - context: LoopyExpressionContext) -> ScalarExpression: + expr_context: LoopyExpressionContext) -> ScalarExpression: result: GeneratedResult = self.codegen_mapper( - context.namespace[expr.name], - context.state) - return result.to_loopy_expression((), context) + expr_context.namespace[expr.name], + expr_context.state) + return result.to_loopy_expression((), expr_context) # }}} @@ -394,19 +394,19 @@ def add_output(name: str, expr: Array, state: CodeGenState, is_output_only=True) indices = tuple(prim.Variable(iname) for iname in inames) - context = state.make_expression_context() - copy_expr = result.to_loopy_expression(indices, context) + expr_context = state.make_expression_context() + copy_expr = result.to_loopy_expression(indices, expr_context) # TODO: Contextual data not supported yet. - assert not context.reduction_bounds - assert not context.depends_on + assert not expr_context.reduction_bounds + assert not expr_context.depends_on from loopy.kernel.instruction import make_assignment insn = make_assignment((prim.Variable(name)[indices],), copy_expr, id=state.insn_id_gen(f"{name}_copy"), within_inames=frozenset(inames), - depends_on=context.depends_on) + depends_on=expr_context.depends_on) kernel = state.kernel kernel = kernel.copy(args=kernel.args + [arg], -- GitLab From b053bec8ea9ebb57a0a8ed6d272f15c7ad2d2dd2 Mon Sep 17 00:00:00 2001 From: Matt Wala Date: Mon, 15 Jun 2020 13:41:26 -0500 Subject: [PATCH 54/76] GeneratedResult: Don't have own shape, dtype --- pytato/codegen.py | 29 +++++++++++++---------------- 1 file changed, 13 insertions(+), 16 deletions(-) diff --git a/pytato/codegen.py b/pytato/codegen.py index fc45042..1d455d4 100644 --- a/pytato/codegen.py +++ b/pytato/codegen.py @@ -32,7 +32,6 @@ import typing import islpy as isl import loopy as lp -import numpy as np import pymbolic.primitives as prim from pytato.array import ( @@ -81,14 +80,13 @@ SymbolicIndex = ShapeType class GeneratedResult(object): """Generated code for a node in the computation graph (i.e., an array expression). - """ - def __init__(self, shape: ShapeType, dtype: np.dtype): - self.shape = shape - self.dtype = dtype - @property - def ndim(self) -> int: - return len(self.shape) + .. attribute:: array + + The :class:`pytato.Array` associated with this code. + """ + def __init__(self, array: Array): + self.array = array def to_loopy_expression(self, indices: SymbolicIndex, expr_context: LoopyExpressionContext) -> ScalarExpression: @@ -101,8 +99,8 @@ class ArrayResult(GeneratedResult): See also: :class:`pytato.array.ImplStored`. """ - def __init__(self, name: str, shape: ShapeType, dtype: np.dtype): - super().__init__(shape, dtype) + def __init__(self, name: str, array: Array): + super().__init__(array) self.name = name # TODO: Handle dependencies. @@ -119,9 +117,8 @@ class LoopyExpressionResult(GeneratedResult): See also: :class:`pytato.array.ImplInlined`. """ - def __init__( - self, expr: ScalarExpression, shape: ShapeType, dtype: np.dtype): - super().__init__(shape, dtype) + def __init__(self, expr: ScalarExpression, array: Array): + super().__init__(array) self.expr = expr # TODO: Handle dependencies and reduction domains. @@ -129,7 +126,7 @@ class LoopyExpressionResult(GeneratedResult): expr_context: LoopyExpressionContext) -> ScalarExpression: return scalar_expr.substitute( self.expr, - {f"_{d}": i for d, i in zip(range(self.ndim), indices)}) + {f"_{d}": i for d, i in zip(range(self.array.ndim), indices)}) class SubstitutionRuleResult(GeneratedResult): @@ -227,7 +224,7 @@ class CodeGenMapper(pytato.transform.Mapper): kernel = state.kernel.copy(args=state.kernel.args + [arg]) state.update_kernel(kernel) - result = ArrayResult(expr.name, expr.dtype, expr.shape) + result = ArrayResult(expr.name, expr) state.results[expr] = result return result @@ -242,7 +239,7 @@ class CodeGenMapper(pytato.transform.Mapper): expr_context = chained_state.make_expression_context() loopy_expr = self.exprgen_mapper(expr.expr, expr_context) - result = LoopyExpressionResult(loopy_expr, expr.shape, expr.dtype) + result = LoopyExpressionResult(loopy_expr, expr) state.results[expr] = result return result -- GitLab From 8dc8ccac6a30207daad3f59add5179a38956c2c5 Mon Sep 17 00:00:00 2001 From: Matt Wala Date: Mon, 15 Jun 2020 14:12:46 -0500 Subject: [PATCH 55/76] Rename result classes and LoopyExpressionGenMapper. --- pytato/codegen.py | 126 +++++++++++++++++++++++----------------------- 1 file changed, 63 insertions(+), 63 deletions(-) diff --git a/pytato/codegen.py b/pytato/codegen.py index 1d455d4..b89e271 100644 --- a/pytato/codegen.py +++ b/pytato/codegen.py @@ -54,16 +54,16 @@ Code Generation Internals .. currentmodule:: pytato.codegen -.. autoclass:: GeneratedResult -.. autoclass:: ArrayResult -.. autoclass:: LoopyExpressionResult +.. autoclass:: LoopyExpressionContext +.. autoclass:: ImplementedResult +.. autoclass:: StoredResult +.. autoclass:: InlinedResult .. autoclass:: SubstitutionRuleResult .. autoclass:: CodeGenState .. autoclass:: CodeGenMapper -.. autoclass:: LoopyExpressionContext -.. autoclass:: LoopyExpressionGenMapper +.. autoclass:: InlinedExpressionGenMapper .. autofunction:: domain_for_shape .. autofunction:: add_output @@ -73,11 +73,49 @@ Code Generation Internals # {{{ generated array expressions -# These are semantically distinct but identical at the type level. +# SymbolicIndex and ShapeType are semantically distinct but identical at the +# type level. SymbolicIndex = ShapeType +ReductionBounds = Dict[str, Tuple[ScalarExpression, ScalarExpression]] + + +@dataclasses.dataclass(init=True, repr=False, eq=False) +class LoopyExpressionContext(object): + """Contextual data for generating :mod:`loopy` expressions. + + This data is passed through :class:`LoopyExpressionGenMapper` + via arguments. + + .. attribute:: state + + The :class:`CodeGenState`. + + .. attribute:: depends_on + + The set of statement IDs that need to be included in + :attr:`loopy.kernel.data.instruction.InstructionBase.depends_on`. + + .. attribute:: reduction_bounds + + A mapping from inames to reduction bounds in the expression. + """ + state: CodeGenState + _depends_on: FrozenSet[str] + reduction_bounds: ReductionBounds + + @property + def namespace(self) -> typing.ChainMap[str, Array]: + return self.state.namespace + + @property + def depends_on(self) -> FrozenSet[str]: + return self._depends_on + + def update_depends_on(self, other: FrozenSet[str]) -> None: + self._depends_on = self._depends_on | other -class GeneratedResult(object): +class ImplementedResult(object): """Generated code for a node in the computation graph (i.e., an array expression). @@ -94,7 +132,7 @@ class GeneratedResult(object): raise NotImplementedError -class ArrayResult(GeneratedResult): +class StoredResult(ImplementedResult): """An array expression generated as a :mod:`loopy` array. See also: :class:`pytato.array.ImplStored`. @@ -112,8 +150,9 @@ class ArrayResult(GeneratedResult): return prim.Variable(self.name)[indices] -class LoopyExpressionResult(GeneratedResult): - """An array expression generated as a :mod:`loopy` expression. +class InlinedResult(ImplementedResult): + """An array expression generated as :mod:`loopy` expression that can be inlined + within other expressions. See also: :class:`pytato.array.ImplInlined`. """ @@ -129,7 +168,7 @@ class LoopyExpressionResult(GeneratedResult): {f"_{d}": i for d, i in zip(range(self.array.ndim), indices)}) -class SubstitutionRuleResult(GeneratedResult): +class SubstitutionRuleResult(ImplementedResult): # TODO: implement pass @@ -153,7 +192,7 @@ class CodeGenState: .. attribute:: results A mapping from :class:`pytato.array.Array` instances to - instances of :class:`GeneratedResult`. + instances of :class:`ImplementedResult`. .. attribute:: var_name_gen .. attribute:: insn_id_gen @@ -164,7 +203,7 @@ class CodeGenState: """ namespace: typing.ChainMap[str, Array] _kernel: lp.LoopKernel - results: Dict[Array, GeneratedResult] + results: Dict[Array, ImplementedResult] # Both of these have type Callable[[str], str], but mypy's support for that # is broken (https://github.com/python/mypy/issues/6910) @@ -207,13 +246,13 @@ class CodeGenState: class CodeGenMapper(pytato.transform.Mapper): """A mapper for generating code for nodes in the computation graph. """ - exprgen_mapper: LoopyExpressionGenMapper + exprgen_mapper: InlinedExpressionGenMapper def __init__(self) -> None: - self.exprgen_mapper = LoopyExpressionGenMapper(self) + self.exprgen_mapper = InlinedExpressionGenMapper(self) def map_placeholder(self, expr: Placeholder, - state: CodeGenState) -> GeneratedResult: + state: CodeGenState) -> ImplementedResult: if expr in state.results: return state.results[expr] @@ -224,12 +263,12 @@ class CodeGenMapper(pytato.transform.Mapper): kernel = state.kernel.copy(args=state.kernel.args + [arg]) state.update_kernel(kernel) - result = ArrayResult(expr.name, expr) + result = StoredResult(expr.name, expr) state.results[expr] = result return result def map_index_lambda(self, expr: IndexLambda, - state: CodeGenState) -> GeneratedResult: + state: CodeGenState) -> ImplementedResult: if expr in state.results: return state.results[expr] @@ -239,55 +278,16 @@ class CodeGenMapper(pytato.transform.Mapper): expr_context = chained_state.make_expression_context() loopy_expr = self.exprgen_mapper(expr.expr, expr_context) - result = LoopyExpressionResult(loopy_expr, expr) + result = InlinedResult(loopy_expr, expr) state.results[expr] = result return result # }}} -# {{{ loopy expression gen mapper - -ReductionBounds = Dict[str, Tuple[ScalarExpression, ScalarExpression]] - - -@dataclasses.dataclass(init=True, repr=False, eq=False) -class LoopyExpressionContext(object): - """Contextual data for generating :mod:`loopy` expressions. - - This data is passed through :class:`LoopyExpressionGenMapper` - via arguments. - - .. attribute:: state - - The :class:`CodeGenState`. - - .. attribute:: depends_on - - The set of statement IDs that need to be included in - :attr:`loopy.kernel.data.instruction.InstructionBase.depends_on`. - - .. attribute:: reduction_bounds - - A mapping from inames to reduction bounds in the expression. - """ - state: CodeGenState - _depends_on: FrozenSet[str] - reduction_bounds: ReductionBounds - - @property - def namespace(self) -> typing.ChainMap[str, Array]: - return self.state.namespace - - @property - def depends_on(self) -> FrozenSet[str]: - return self._depends_on - - def update_depends_on(self, other: FrozenSet[str]) -> None: - self._depends_on = self._depends_on | other - +# {{{ inlined expression gen mapper -class LoopyExpressionGenMapper(scalar_expr.IdentityMapper): +class InlinedExpressionGenMapper(scalar_expr.IdentityMapper): """A mapper for generating :mod:`loopy` expressions. The inputs to this mapper are scalar expression as found in @@ -295,7 +295,7 @@ class LoopyExpressionGenMapper(scalar_expr.IdentityMapper): expressions). The outputs of this mapper are scalar expressions suitable for wrapping in - :class:`LoopyExpressionResult`. + :class:`InlinedResult`. """ codegen_mapper: CodeGenMapper @@ -309,7 +309,7 @@ class LoopyExpressionGenMapper(scalar_expr.IdentityMapper): def map_subscript(self, expr: prim.Subscript, expr_context: LoopyExpressionContext) -> ScalarExpression: assert isinstance(expr.aggregate, prim.Variable) - result: GeneratedResult = self.codegen_mapper( + result: ImplementedResult = self.codegen_mapper( expr_context.namespace[expr.aggregate.name], expr_context.state) return result.to_loopy_expression(expr.index, expr_context) @@ -317,7 +317,7 @@ class LoopyExpressionGenMapper(scalar_expr.IdentityMapper): def map_variable(self, expr: prim.Variable, expr_context: LoopyExpressionContext) -> ScalarExpression: - result: GeneratedResult = self.codegen_mapper( + result: ImplementedResult = self.codegen_mapper( expr_context.namespace[expr.name], expr_context.state) return result.to_loopy_expression((), expr_context) -- GitLab From 5ee5e4c76b4e0b8a7ca64295e77c7c28386d0099 Mon Sep 17 00:00:00 2001 From: Matt Wala Date: Mon, 15 Jun 2020 14:17:08 -0500 Subject: [PATCH 56/76] Improve docs --- pytato/codegen.py | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/pytato/codegen.py b/pytato/codegen.py index b89e271..565f24c 100644 --- a/pytato/codegen.py +++ b/pytato/codegen.py @@ -151,8 +151,8 @@ class StoredResult(ImplementedResult): class InlinedResult(ImplementedResult): - """An array expression generated as :mod:`loopy` expression that can be inlined - within other expressions. + """An array expression generated as a :mod:`loopy` expression containing inlined + sub-expressions. See also: :class:`pytato.array.ImplInlined`. """ @@ -288,7 +288,8 @@ class CodeGenMapper(pytato.transform.Mapper): # {{{ inlined expression gen mapper class InlinedExpressionGenMapper(scalar_expr.IdentityMapper): - """A mapper for generating :mod:`loopy` expressions. + """A mapper for generating :mod:`loopy` expressions with inlined + sub-expressions. The inputs to this mapper are scalar expression as found in :class:`pytato.IndexLambda`, or expressions that are compatible (e.g., shape -- GitLab From 695adc7fc2bea48c2b315124eaedbc69091cd189 Mon Sep 17 00:00:00 2001 From: Matt Wala Date: Mon, 15 Jun 2020 17:20:35 -0500 Subject: [PATCH 57/76] More doc tweaks --- pytato/codegen.py | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/pytato/codegen.py b/pytato/codegen.py index 565f24c..eff58ae 100644 --- a/pytato/codegen.py +++ b/pytato/codegen.py @@ -83,8 +83,9 @@ ReductionBounds = Dict[str, Tuple[ScalarExpression, ScalarExpression]] class LoopyExpressionContext(object): """Contextual data for generating :mod:`loopy` expressions. - This data is passed through :class:`LoopyExpressionGenMapper` - via arguments. + This data is passed through :class:`LoopyExpressionGenMapper` via arguments, + and is also used by :meth:`ImplementedResult.to_loopy_expression` to + retrieve contextual data. .. attribute:: state @@ -122,6 +123,8 @@ class ImplementedResult(object): .. attribute:: array The :class:`pytato.Array` associated with this code. + + .. automethod:: to_loopy_expression """ def __init__(self, array: Array): self.array = array -- GitLab From da545a49dac0c7c9275f25a20841fe974d05331c Mon Sep 17 00:00:00 2001 From: Matt Wala Date: Mon, 15 Jun 2020 17:21:02 -0500 Subject: [PATCH 58/76] Revert "Apply suggestion to pytato/array.py" This reverts commit 26b2dca73fbc58ab76af763b7c44b420d1c52cf1. --- pytato/array.py | 1 - 1 file changed, 1 deletion(-) diff --git a/pytato/array.py b/pytato/array.py index 737ea46..381a9a7 100644 --- a/pytato/array.py +++ b/pytato/array.py @@ -437,7 +437,6 @@ class Array: @memoize_method def __hash__(self) -> int: - # TODO Cache computed hash value attrs = [] for field in self.fields: attr = getattr(self, field) -- GitLab From 60a05067376ddee1cb2a4386f36a270014ca564d Mon Sep 17 00:00:00 2001 From: Matt Wala Date: Mon, 15 Jun 2020 17:24:47 -0500 Subject: [PATCH 59/76] Remove namespace argument from generate_code() --- pytato/codegen.py | 23 ++++++++--------------- 1 file changed, 8 insertions(+), 15 deletions(-) diff --git a/pytato/codegen.py b/pytato/codegen.py index eff58ae..f4817cf 100644 --- a/pytato/codegen.py +++ b/pytato/codegen.py @@ -418,33 +418,26 @@ def add_output(name: str, expr: Array, state: CodeGenState, # }}} -def generate_loopy( - result_or_namespace: Union[Namespace, Array, DictOfNamedArrays], +def generate_loopy(result: Union[Array, DictOfNamedArrays], target: Optional[Target] = None) -> BoundProgram: r"""Code generation entry point. - :param result_or_namespace: For :term:`array result`\ s, code generation treats - the node(s) as outputs of the computation. - + :param result: Outputs of the computation. :returns: A wrapped generated :mod:`loopy` kernel """ # {{{ get namespace and outputs outputs: DictOfNamedArrays - namespace: Namespace - if isinstance(result_or_namespace, Array): - outputs = DictOfNamedArrays({"out": result_or_namespace}) - namespace = outputs.namespace - elif isinstance(result_or_namespace, DictOfNamedArrays): - outputs = result_or_namespace + if isinstance(result, Array): + outputs = DictOfNamedArrays({"out": result}) namespace = outputs.namespace else: - assert isinstance(result_or_namespace, Namespace) - outputs = DictOfNamedArrays({}) - namespace = result_or_namespace + assert isinstance(result, DictOfNamedArrays) + outputs = result - del result_or_namespace + namespace = outputs.namespace + del result # }}} -- GitLab From 12bf1819aa00c4b615b32216c3f89da5659c74eb Mon Sep 17 00:00:00 2001 From: Matt Wala Date: Tue, 16 Jun 2020 00:27:51 +0200 Subject: [PATCH 60/76] Remove unused import --- pytato/codegen.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/pytato/codegen.py b/pytato/codegen.py index f4817cf..a38cd45 100644 --- a/pytato/codegen.py +++ b/pytato/codegen.py @@ -35,8 +35,7 @@ import loopy as lp import pymbolic.primitives as prim from pytato.array import ( - Array, DictOfNamedArrays, Placeholder, Namespace, ShapeType, - IndexLambda) + Array, DictOfNamedArrays, Placeholder, ShapeType, IndexLambda) from pytato.program import BoundProgram, Target, PyOpenCLTarget import pytato.scalar_expr as scalar_expr from pytato.scalar_expr import ScalarExpression -- GitLab From 571513d13e9283560ae1789db6c6e97ae7be4320 Mon Sep 17 00:00:00 2001 From: Matt Wala Date: Wed, 17 Jun 2020 00:09:51 -0500 Subject: [PATCH 61/76] Tweak doc layout --- doc/codegen.rst | 4 ---- doc/index.rst | 1 - doc/reference.rst | 1 + pytato/codegen.py | 2 ++ pytato/program.py | 1 - 5 files changed, 3 insertions(+), 6 deletions(-) delete mode 100644 doc/codegen.rst diff --git a/doc/codegen.rst b/doc/codegen.rst deleted file mode 100644 index 917751c..0000000 --- a/doc/codegen.rst +++ /dev/null @@ -1,4 +0,0 @@ -Generating Code -=============== - -.. automodule:: pytato.codegen diff --git a/doc/index.rst b/doc/index.rst index 187dda3..18dd4da 100644 --- a/doc/index.rst +++ b/doc/index.rst @@ -6,7 +6,6 @@ Welcome to Pytato's documentation! :caption: Contents: reference - codegen design misc diff --git a/doc/reference.rst b/doc/reference.rst index 2a14d6a..0bd708b 100644 --- a/doc/reference.rst +++ b/doc/reference.rst @@ -5,3 +5,4 @@ Reference .. automodule:: pytato.program .. automodule:: pytato.scalar_expr .. automodule:: pytato.transform +.. automodule:: pytato.codegen diff --git a/pytato/codegen.py b/pytato/codegen.py index f4817cf..4f474d8 100644 --- a/pytato/codegen.py +++ b/pytato/codegen.py @@ -44,6 +44,8 @@ import pytato.transform __doc__ = """ +Generating Code +--------------- .. currentmodule:: pytato diff --git a/pytato/program.py b/pytato/program.py index 4e0506f..358f60e 100644 --- a/pytato/program.py +++ b/pytato/program.py @@ -23,7 +23,6 @@ THE SOFTWARE. """ __doc__ = """ - .. currentmodule:: pytato.program Code Generation Targets -- GitLab From 9b01cb652f57b3249d6230f95ae758c332091960 Mon Sep 17 00:00:00 2001 From: Matt Wala Date: Mon, 22 Jun 2020 05:31:51 +0200 Subject: [PATCH 62/76] Apply suggestion to pytato/codegen.py --- pytato/codegen.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/pytato/codegen.py b/pytato/codegen.py index 39d7694..f092555 100644 --- a/pytato/codegen.py +++ b/pytato/codegen.py @@ -229,7 +229,9 @@ class CodeGenState: def chain_namespaces( self, local_namespace: Mapping[str, Array]) -> Iterator[CodeGenState]: - """A context manager for overriding with a local scope.""" + """A context manager for (locally) adding an overlaid name mapping + to this code generation state. + """ self.namespace.maps.insert(0, local_namespace) yield self self.namespace.maps.pop(0) -- GitLab From 50af86d663aa27451671f950278fbe8c54b31b83 Mon Sep 17 00:00:00 2001 From: Matt Wala Date: Mon, 22 Jun 2020 05:31:55 +0200 Subject: [PATCH 63/76] Apply suggestion to pytato/codegen.py --- pytato/codegen.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pytato/codegen.py b/pytato/codegen.py index f092555..8086881 100644 --- a/pytato/codegen.py +++ b/pytato/codegen.py @@ -226,7 +226,7 @@ class CodeGenState: self._kernel = kernel @contextlib.contextmanager - def chain_namespaces( + def chained_namespaces( self, local_namespace: Mapping[str, Array]) -> Iterator[CodeGenState]: """A context manager for (locally) adding an overlaid name mapping -- GitLab From 399b991bd8f498dc209727d59bc0b58d4fcadbee Mon Sep 17 00:00:00 2001 From: Matt Wala Date: Mon, 22 Jun 2020 08:31:09 +0200 Subject: [PATCH 64/76] Apply suggestion to pytato/array.py --- pytato/array.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/pytato/array.py b/pytato/array.py index 381a9a7..87a8fb7 100644 --- a/pytato/array.py +++ b/pytato/array.py @@ -767,8 +767,8 @@ class Placeholder(_SuppliedShapeAndDtypeMixin, Array): .. note:: - Modifying :class:`Placeholder` tags is not supported after - creation. + Creating multiple instances of a :class:`Placeholder` with the same name + and within the same :class:`Namespace` is not allowed. """ mapper_method = "map_placeholder" fields = Array.fields + ("name",) -- GitLab From 01d4ef868bf3f0acccf89192197a37469d3d67d2 Mon Sep 17 00:00:00 2001 From: Matt Wala Date: Mon, 22 Jun 2020 15:49:14 +0200 Subject: [PATCH 65/76] Apply suggestion to pytato/codegen.py --- pytato/codegen.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/pytato/codegen.py b/pytato/codegen.py index 8086881..5c7de68 100644 --- a/pytato/codegen.py +++ b/pytato/codegen.py @@ -82,7 +82,8 @@ ReductionBounds = Dict[str, Tuple[ScalarExpression, ScalarExpression]] @dataclasses.dataclass(init=True, repr=False, eq=False) class LoopyExpressionContext(object): - """Contextual data for generating :mod:`loopy` expressions. + """Mutable state used while generating :mod:`loopy` expressions. + Wraps :class:`CodeGenState` with more expression-specific information. This data is passed through :class:`LoopyExpressionGenMapper` via arguments, and is also used by :meth:`ImplementedResult.to_loopy_expression` to -- GitLab From 5d12a162958730eea930dc52454600ea98e6436f Mon Sep 17 00:00:00 2001 From: Matt Wala Date: Mon, 22 Jun 2020 08:51:55 -0500 Subject: [PATCH 66/76] chained_namespaces --- pytato/codegen.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/pytato/codegen.py b/pytato/codegen.py index 5c7de68..98eeee5 100644 --- a/pytato/codegen.py +++ b/pytato/codegen.py @@ -203,7 +203,7 @@ class CodeGenState: .. attribute:: insn_id_gen .. automethod:: update_kernel - .. automethod:: chain_namespaces + .. automethod:: chained_namespaces .. automethod:: make_expression_context """ namespace: typing.ChainMap[str, Array] @@ -281,7 +281,7 @@ class CodeGenMapper(pytato.transform.Mapper): # TODO: Respect tags. - with state.chain_namespaces(expr.bindings) as chained_state: + with state.chained_namespaces(expr.bindings) as chained_state: expr_context = chained_state.make_expression_context() loopy_expr = self.exprgen_mapper(expr.expr, expr_context) -- GitLab From 64a0c5620b4ba25853e7cbfbae4c88798c405600 Mon Sep 17 00:00:00 2001 From: Matt Wala Date: Mon, 22 Jun 2020 15:32:41 -0500 Subject: [PATCH 67/76] Create pytato.target, separate from pytato.program --- doc/reference.rst | 3 +- pytato/__init__.py | 2 +- pytato/codegen.py | 3 +- pytato/program.py | 70 +++++----------------------------- pytato/target.py | 94 ++++++++++++++++++++++++++++++++++++++++++++++ 5 files changed, 109 insertions(+), 63 deletions(-) create mode 100644 pytato/target.py diff --git a/doc/reference.rst b/doc/reference.rst index 0bd708b..43f6555 100644 --- a/doc/reference.rst +++ b/doc/reference.rst @@ -2,7 +2,8 @@ Reference ========= .. automodule:: pytato.array -.. automodule:: pytato.program .. automodule:: pytato.scalar_expr .. automodule:: pytato.transform +.. automodule:: pytato.program +.. automodule:: pytato.target .. automodule:: pytato.codegen diff --git a/pytato/__init__.py b/pytato/__init__.py index dcfaf28..123041a 100644 --- a/pytato/__init__.py +++ b/pytato/__init__.py @@ -30,7 +30,7 @@ from pytato.array import ( ) from pytato.codegen import generate_loopy -from pytato.program import Target, PyOpenCLTarget +from pytato.target import Target, PyOpenCLTarget __all__ = ( "DottedName", "Namespace", "Array", "DictOfNamedArrays", diff --git a/pytato/codegen.py b/pytato/codegen.py index 98eeee5..9ea0478 100644 --- a/pytato/codegen.py +++ b/pytato/codegen.py @@ -36,7 +36,8 @@ import pymbolic.primitives as prim from pytato.array import ( Array, DictOfNamedArrays, Placeholder, ShapeType, IndexLambda) -from pytato.program import BoundProgram, Target, PyOpenCLTarget +from pytato.program import BoundProgram +from pytato.target import Target, PyOpenCLTarget import pytato.scalar_expr as scalar_expr from pytato.scalar_expr import ScalarExpression import pytato.transform diff --git a/pytato/program.py b/pytato/program.py index 358f60e..e2e2c86 100644 --- a/pytato/program.py +++ b/pytato/program.py @@ -25,12 +25,6 @@ THE SOFTWARE. __doc__ = """ .. currentmodule:: pytato.program -Code Generation Targets ------------------------ - -.. autoclass:: Target -.. autoclass:: PyOpenCLTarget - Generated Executable Programs ----------------------------- @@ -42,58 +36,14 @@ from dataclasses import dataclass import typing from typing import Any, Mapping, Optional -if typing.TYPE_CHECKING: - # Skip imports for efficiency. FIXME: Neither of these work as type stubs - # are not present. Types are here only as documentation. - import pyopencl as cl - import loopy as lp - - -class Target: - """An abstract code generation target. - .. automethod:: get_loopy_target - .. automethod:: bind_program - """ - - def get_loopy_target(self) -> "lp.TargetBase": - """Return the corresponding :mod:`loopy` target.""" - raise NotImplementedError - - def bind_program(self, program: "lp.LoopKernel", - bound_arguments: Mapping[str, Any]) -> BoundProgram: - """Create a :class:`BoundProgram` for this code generation target. - - :param program: the :mod:`loopy` kernel - :param bound_arguments: a mapping from argument names to outputs - """ - raise NotImplementedError - - -class PyOpenCLTarget(Target): - """A :mod:`pyopencl` code generation target. - - .. attribute:: queue - - The :mod:`pyopencl` command queue, or *None*. - """ - - def __init__(self, queue: Optional["cl.CommandQueue"] = None): - self.queue = queue - - def get_loopy_target(self) -> "lp.PyOpenCLTarget": - import loopy as lp - device = None - if self.queue is not None: - device = self.queue.device - return lp.PyOpenCLTarget(device) - - def bind_program(self, program: "lp.LoopKernel", - bound_arguments: Mapping[str, Any]) -> BoundProgram: - return BoundPyOpenCLProgram(program=program, - queue=self.queue, - bound_arguments=bound_arguments, - target=self) +if typing.TYPE_CHECKING: + # Imports skipped for efficiency. FIXME: Neither of these work as type + # stubs are not present. Types are here only as documentation. + import pyopencl + import loopy + # Imports skipped to avoid circular dependencies. + import pytato.target @dataclass(init=True, repr=False, eq=False) @@ -115,9 +65,9 @@ class BoundProgram: .. automethod:: __call__ """ - program: "lp.LoopKernel" + program: "loopy.LoopKernel" bound_arguments: Mapping[str, Any] - target: Target + target: "pytato.target.Target" def __call__(self, *args: Any, **kwargs: Any) -> Any: raise NotImplementedError @@ -133,7 +83,7 @@ class BoundPyOpenCLProgram(BoundProgram): .. automethod:: __call__ """ - queue: Optional["cl.CommandQueue"] + queue: Optional["pyopencl.CommandQueue"] def __call__(self, *args: Any, **kwargs: Any) -> Any: """Convenience function for launching a :mod:`pyopencl` computation.""" diff --git a/pytato/target.py b/pytato/target.py new file mode 100644 index 0000000..0242c0d --- /dev/null +++ b/pytato/target.py @@ -0,0 +1,94 @@ +from __future__ import annotations + +__copyright__ = """Copyright (C) 2020 Matt Wala""" + +__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. +""" + +__doc__ = """ +.. currentmodule:: pytato.target + +Code Generation Targets +----------------------- + +.. autoclass:: Target +.. autoclass:: PyOpenCLTarget +""" + +import typing +from typing import Any, Mapping, Optional + +from pytato.program import BoundProgram, BoundPyOpenCLProgram + + +if typing.TYPE_CHECKING: + # Skip imports for efficiency. FIXME: Neither of these work as type stubs + # are not present. Types are here only as documentation. + import pyopencl as cl + import loopy as lp + + +class Target: + """An abstract code generation target. + + .. automethod:: get_loopy_target + .. automethod:: bind_program + """ + + def get_loopy_target(self) -> "lp.TargetBase": + """Return the corresponding :mod:`loopy` target.""" + raise NotImplementedError + + def bind_program(self, program: "lp.LoopKernel", + bound_arguments: Mapping[str, Any]) -> BoundProgram: + """Create a :class:`pytato.program.BoundProgram` for this code generation target. + + :param program: the :mod:`loopy` kernel + :param bound_arguments: a mapping from argument names to outputs + """ + raise NotImplementedError + + +class PyOpenCLTarget(Target): + """A :mod:`pyopencl` code generation target. + + .. attribute:: queue + + The :mod:`pyopencl` command queue, or *None*. + """ + + def __init__(self, queue: Optional["cl.CommandQueue"] = None): + self.queue = queue + + def get_loopy_target(self) -> "lp.PyOpenCLTarget": + import loopy as lp + device = None + if self.queue is not None: + device = self.queue.device + return lp.PyOpenCLTarget(device) + + def bind_program(self, program: "lp.LoopKernel", + bound_arguments: Mapping[str, Any]) -> BoundProgram: + return BoundPyOpenCLProgram(program=program, + queue=self.queue, + bound_arguments=bound_arguments, + target=self) + +# vim: foldmethod=marker -- GitLab From b337766cae388ca186d0c9179605da1ac7e5bf5f Mon Sep 17 00:00:00 2001 From: Matt Wala Date: Mon, 22 Jun 2020 15:38:12 -0500 Subject: [PATCH 68/76] Change mind about abbreviations --- pytato/program.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/pytato/program.py b/pytato/program.py index e2e2c86..ded1909 100644 --- a/pytato/program.py +++ b/pytato/program.py @@ -40,8 +40,8 @@ from typing import Any, Mapping, Optional if typing.TYPE_CHECKING: # Imports skipped for efficiency. FIXME: Neither of these work as type # stubs are not present. Types are here only as documentation. - import pyopencl - import loopy + import pyopencl as cl + import loopy as lp # Imports skipped to avoid circular dependencies. import pytato.target @@ -65,7 +65,7 @@ class BoundProgram: .. automethod:: __call__ """ - program: "loopy.LoopKernel" + program: "lp.LoopKernel" bound_arguments: Mapping[str, Any] target: "pytato.target.Target" @@ -83,7 +83,7 @@ class BoundPyOpenCLProgram(BoundProgram): .. automethod:: __call__ """ - queue: Optional["pyopencl.CommandQueue"] + queue: Optional["cl.CommandQueue"] def __call__(self, *args: Any, **kwargs: Any) -> Any: """Convenience function for launching a :mod:`pyopencl` computation.""" -- GitLab From 174ba9901280d950b059e7238364affc96ed623c Mon Sep 17 00:00:00 2001 From: Matt Wala Date: Tue, 23 Jun 2020 14:59:51 -0500 Subject: [PATCH 69/76] Change the way local namespaces are handled to avoid (potentially incorrect/confusing) chaining --- pytato/codegen.py | 60 +++++++++++++++++++++++------------------------ 1 file changed, 30 insertions(+), 30 deletions(-) diff --git a/pytato/codegen.py b/pytato/codegen.py index 9ea0478..06aa17b 100644 --- a/pytato/codegen.py +++ b/pytato/codegen.py @@ -23,12 +23,9 @@ THE SOFTWARE. """ import collections -import contextlib import dataclasses from typing import ( - Any, Union, Optional, Mapping, Iterator, Dict, Tuple, FrozenSet, - Set) -import typing + Any, Union, Optional, Mapping, Dict, Tuple, FrozenSet, Set) import islpy as isl import loopy as lp @@ -94,6 +91,11 @@ class LoopyExpressionContext(object): The :class:`CodeGenState`. + .. attribute:: local_namespace + + A (read-only) local name mapping used for name lookup when generating + code. + .. attribute:: depends_on The set of statement IDs that need to be included in @@ -102,14 +104,24 @@ class LoopyExpressionContext(object): .. attribute:: reduction_bounds A mapping from inames to reduction bounds in the expression. + + .. automethod:: update_depends_on + .. automethod:: lookup + """ state: CodeGenState - _depends_on: FrozenSet[str] - reduction_bounds: ReductionBounds - - @property - def namespace(self) -> typing.ChainMap[str, Array]: - return self.state.namespace + _depends_on: FrozenSet[str] = \ + dataclasses.field(default_factory=frozenset) + local_namespace: Mapping[str, Array] = \ + dataclasses.field(default_factory=dict) + reduction_bounds: ReductionBounds = \ + dataclasses.field(default_factory=dict) + + def lookup(self, name: str) -> Array: + try: + return self.local_namespace[name] + except KeyError: + return self.state.namespace[name] @property def depends_on(self) -> FrozenSet[str]: @@ -189,7 +201,7 @@ class CodeGenState: .. attribute:: namespace - The namespace + The (global) namespace .. attribute:: kernel @@ -204,10 +216,9 @@ class CodeGenState: .. attribute:: insn_id_gen .. automethod:: update_kernel - .. automethod:: chained_namespaces .. automethod:: make_expression_context """ - namespace: typing.ChainMap[str, Array] + namespace: Mapping[str, Array] _kernel: lp.LoopKernel results: Dict[Array, ImplementedResult] @@ -227,17 +238,6 @@ class CodeGenState: def update_kernel(self, kernel: lp.LoopKernel) -> None: self._kernel = kernel - @contextlib.contextmanager - def chained_namespaces( - self, - local_namespace: Mapping[str, Array]) -> Iterator[CodeGenState]: - """A context manager for (locally) adding an overlaid name mapping - to this code generation state. - """ - self.namespace.maps.insert(0, local_namespace) - yield self - self.namespace.maps.pop(0) - def make_expression_context( self, depends_on: FrozenSet[str] = frozenset(), @@ -282,9 +282,9 @@ class CodeGenMapper(pytato.transform.Mapper): # TODO: Respect tags. - with state.chained_namespaces(expr.bindings) as chained_state: - expr_context = chained_state.make_expression_context() - loopy_expr = self.exprgen_mapper(expr.expr, expr_context) + expr_context = LoopyExpressionContext(state, + local_namespace=expr.bindings) + loopy_expr = self.exprgen_mapper(expr.expr, expr_context) result = InlinedResult(loopy_expr, expr) state.results[expr] = result @@ -319,7 +319,7 @@ class InlinedExpressionGenMapper(scalar_expr.IdentityMapper): expr_context: LoopyExpressionContext) -> ScalarExpression: assert isinstance(expr.aggregate, prim.Variable) result: ImplementedResult = self.codegen_mapper( - expr_context.namespace[expr.aggregate.name], expr_context.state) + expr_context.lookup(expr.aggregate.name), expr_context.state) return result.to_loopy_expression(expr.index, expr_context) # TODO: map_reduction() @@ -327,7 +327,7 @@ class InlinedExpressionGenMapper(scalar_expr.IdentityMapper): def map_variable(self, expr: prim.Variable, expr_context: LoopyExpressionContext) -> ScalarExpression: result: ImplementedResult = self.codegen_mapper( - expr_context.namespace[expr.name], + expr_context.lookup(expr.name), expr_context.state) return result.to_loopy_expression((), expr_context) @@ -400,7 +400,7 @@ def add_output(name: str, expr: Array, state: CodeGenState, is_output_only=True) indices = tuple(prim.Variable(iname) for iname in inames) - expr_context = state.make_expression_context() + expr_context = LoopyExpressionContext(state) copy_expr = result.to_loopy_expression(indices, expr_context) # TODO: Contextual data not supported yet. -- GitLab From c20a0a7d6c5e095f622b75b91333df0c4d8ecbbf Mon Sep 17 00:00:00 2001 From: Matt Wala Date: Tue, 23 Jun 2020 15:01:50 -0500 Subject: [PATCH 70/76] Remove make_expression_context() --- pytato/codegen.py | 13 ------------- 1 file changed, 13 deletions(-) diff --git a/pytato/codegen.py b/pytato/codegen.py index 06aa17b..73e7992 100644 --- a/pytato/codegen.py +++ b/pytato/codegen.py @@ -216,7 +216,6 @@ class CodeGenState: .. attribute:: insn_id_gen .. automethod:: update_kernel - .. automethod:: make_expression_context """ namespace: Mapping[str, Array] _kernel: lp.LoopKernel @@ -238,18 +237,6 @@ class CodeGenState: def update_kernel(self, kernel: lp.LoopKernel) -> None: self._kernel = kernel - def make_expression_context( - self, - depends_on: FrozenSet[str] = frozenset(), - reduction_bounds: Optional[ReductionBounds] = None - ) -> LoopyExpressionContext: - """Get a new :class:`LoopyExpressionContext`.""" - if reduction_bounds is None: - reduction_bounds = {} - return LoopyExpressionContext(self, - _depends_on=depends_on, - reduction_bounds=reduction_bounds) - class CodeGenMapper(pytato.transform.Mapper): """A mapper for generating code for nodes in the computation graph. -- GitLab From 406f852af2cf8f0339a53f51af67565547ed722e Mon Sep 17 00:00:00 2001 From: Matt Wala Date: Tue, 23 Jun 2020 15:18:03 -0500 Subject: [PATCH 71/76] Document array.DataInterface --- pytato/array.py | 27 +++++++++++++++++++++------ 1 file changed, 21 insertions(+), 6 deletions(-) diff --git a/pytato/array.py b/pytato/array.py index 87a8fb7..08d3329 100644 --- a/pytato/array.py +++ b/pytato/array.py @@ -40,19 +40,23 @@ __doc__ = """ Array Interface --------------- -.. autoclass :: Namespace -.. autoclass :: Array -.. autoclass :: Tag -.. autoclass :: UniqueTag -.. autoclass :: DictOfNamedArrays +.. autoclass:: Namespace +.. autoclass:: Array +.. autoclass:: Tag +.. autoclass:: UniqueTag +.. autoclass:: DictOfNamedArrays Supporting Functionality ------------------------ -.. autoclass :: DottedName +.. autoclass:: DottedName .. currentmodule:: pytato.array +Concrete Array Data +------------------- +.. autoclass:: DataInterface + Pre-Defined Tags ---------------- @@ -711,6 +715,17 @@ class Reshape(Array): # {{{ data wrapper class DataInterface(Protocol): + """A protocol specifying the minimal interface requirements for concrete + array data supported by :class:`DataWrapper`. + + See :class:`typing.Protocol` for more information about protocols. + + Code generation targets may impose additional restrictions on the kinds of + concrete array data they support. + + .. attribute:: shape + .. attribute:: dtype + """ shape: ShapeType dtype: np.dtype -- GitLab From 29a96b6f2b99833eba224a294e8e62fff85b8989 Mon Sep 17 00:00:00 2001 From: Matt Wala Date: Wed, 24 Jun 2020 13:48:31 -0500 Subject: [PATCH 72/76] Remove unnecessary use of ChainMap --- pytato/codegen.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/pytato/codegen.py b/pytato/codegen.py index 73e7992..758aa79 100644 --- a/pytato/codegen.py +++ b/pytato/codegen.py @@ -22,7 +22,6 @@ OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. """ -import collections import dataclasses from typing import ( Any, Union, Optional, Mapping, Dict, Tuple, FrozenSet, Set) @@ -441,7 +440,7 @@ def generate_loopy(result: Union[Array, DictOfNamedArrays], target=target.get_loopy_target(), lang_version=lp.MOST_RECENT_LANGUAGE_VERSION) - state = CodeGenState(namespace=collections.ChainMap(namespace), + state = CodeGenState(namespace=namespace, _kernel=kernel, results=dict()) -- GitLab From 8ad21b19a11ead1d6737bbdc551de0a49afda2db Mon Sep 17 00:00:00 2001 From: Matt Wala Date: Wed, 24 Jun 2020 14:12:09 -0500 Subject: [PATCH 73/76] Improve input/output name handling * Reserve input/output names before code generation * Make _pt_out the default output name * Add tests for DictOfNamedArrays outputs --- doc/design.rst | 2 ++ pytato/codegen.py | 22 +++++++++++++++------- test/test_codegen.py | 29 +++++++++++++++++++++++++++++ 3 files changed, 46 insertions(+), 7 deletions(-) diff --git a/doc/design.rst b/doc/design.rst index 2309491..fd50da9 100644 --- a/doc/design.rst +++ b/doc/design.rst @@ -119,6 +119,8 @@ Reserved Identifiers - ``_pt_shp``: Used to automatically generate identifiers used in data-dependent shapes. + - ``_pt_out``: The default name of an unnamed output argument + - Identifiers used in index lambdas are also reserved. These include: - Identifiers matching the regular expression ``_[0-9]+``. They are used diff --git a/pytato/codegen.py b/pytato/codegen.py index 758aa79..32f5f32 100644 --- a/pytato/codegen.py +++ b/pytato/codegen.py @@ -29,6 +29,7 @@ from typing import ( import islpy as isl import loopy as lp import pymbolic.primitives as prim +import pytools from pytato.array import ( Array, DictOfNamedArrays, Placeholder, ShapeType, IndexLambda) @@ -220,10 +221,8 @@ class CodeGenState: _kernel: lp.LoopKernel results: Dict[Array, ImplementedResult] - # Both of these have type Callable[[str], str], but mypy's support for that - # is broken (https://github.com/python/mypy/issues/6910) - var_name_gen: Any = dataclasses.field(init=False) - insn_id_gen: Any = dataclasses.field(init=False) + var_name_gen: pytools.UniqueNameGenerator = dataclasses.field(init=False) + insn_id_gen: pytools.UniqueNameGenerator = dataclasses.field(init=False) def __post_init__(self) -> None: self.var_name_gen = self._kernel.get_var_name_generator() @@ -372,7 +371,6 @@ def add_output(name: str, expr: Array, state: CodeGenState, assert expr.shape != () result = mapper(expr, state) - name = state.var_name_gen(name) inames = tuple( state.var_name_gen(f"{name}_dim{d}") @@ -410,10 +408,13 @@ def add_output(name: str, expr: Array, state: CodeGenState, def generate_loopy(result: Union[Array, DictOfNamedArrays], - target: Optional[Target] = None) -> BoundProgram: + target: Optional[Target] = None, + options: Optional[lp.Options] = None) -> BoundProgram: r"""Code generation entry point. :param result: Outputs of the computation. + :param target: Code generation target. + :param options: Code generation options for the kernel. :returns: A wrapped generated :mod:`loopy` kernel """ # {{{ get namespace and outputs @@ -421,7 +422,7 @@ def generate_loopy(result: Union[Array, DictOfNamedArrays], outputs: DictOfNamedArrays if isinstance(result, Array): - outputs = DictOfNamedArrays({"out": result}) + outputs = DictOfNamedArrays({"_pt_out": result}) namespace = outputs.namespace else: assert isinstance(result, DictOfNamedArrays) @@ -438,12 +439,19 @@ def generate_loopy(result: Union[Array, DictOfNamedArrays], # Set up codegen state. kernel = lp.make_kernel("{:}", [], target=target.get_loopy_target(), + options=options, lang_version=lp.MOST_RECENT_LANGUAGE_VERSION) state = CodeGenState(namespace=namespace, _kernel=kernel, results=dict()) + # Reserve names of input and output arguments. + for val in namespace.values(): + if isinstance(val, Placeholder): + state.var_name_gen.add_name(val.name) + state.var_name_gen.add_names(outputs) + # Generate code for graph nodes. mapper = CodeGenMapper() for name, val in namespace.items(): diff --git a/test/test_codegen.py b/test/test_codegen.py index ee2d65d..1e65f46 100755 --- a/test/test_codegen.py +++ b/test/test_codegen.py @@ -24,6 +24,7 @@ THE SOFTWARE. import sys +import loopy as lp import numpy as np import pyopencl as cl import pyopencl.array as cl_array # noqa @@ -48,6 +49,34 @@ def test_basic_codegen(ctx_factory): assert (out == x_in * x_in).all() +def test_DictOfNamedArrays_output(ctx_factory): # noqa + ctx = ctx_factory() + queue = cl.CommandQueue(ctx) + + namespace = pt.Namespace() + x = pt.Placeholder(namespace, "x", (5,), np.int) + y = pt.Placeholder(namespace, "y", (5,), np.int) + x_in = np.array([1, 2, 3, 4, 5]) + y_in = np.array([6, 7, 8, 9, 10]) + + result = pt.DictOfNamedArrays(dict(x_out=x, y_out=y)) + + # Without return_dict. + prog = pt.generate_loopy(result, target=pt.PyOpenCLTarget(queue)) + _, (x_out, y_out) = prog(x=x_in, y=y_in) + assert (x_out == x_in).all() + assert (y_out == y_in).all() + + # With return_dict. + prog = pt.generate_loopy(result, + target=pt.PyOpenCLTarget(queue), + options=lp.Options(return_dict=True)) + + _, outputs = prog(x=x_in, y=y_in) + assert (outputs["x_out"] == x_in).all() + assert (outputs["y_out"] == y_in).all() + + if __name__ == "__main__": if len(sys.argv) > 1: exec(sys.argv[1]) -- GitLab From a4fbfd7d0799f66aa01469aa3976dc5978e47b7e Mon Sep 17 00:00:00 2001 From: Matt Wala Date: Wed, 24 Jun 2020 14:16:18 -0500 Subject: [PATCH 74/76] flake8 fix --- pytato/codegen.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/pytato/codegen.py b/pytato/codegen.py index 32f5f32..5109529 100644 --- a/pytato/codegen.py +++ b/pytato/codegen.py @@ -23,8 +23,7 @@ THE SOFTWARE. """ import dataclasses -from typing import ( - Any, Union, Optional, Mapping, Dict, Tuple, FrozenSet, Set) +from typing import (Union, Optional, Mapping, Dict, Tuple, FrozenSet, Set) import islpy as isl import loopy as lp -- GitLab From 699cd1d0d4a15275a137b2f251a44e1829ac1cb6 Mon Sep 17 00:00:00 2001 From: Matt Wala Date: Wed, 24 Jun 2020 14:38:32 -0500 Subject: [PATCH 75/76] Rename test --- test/test_codegen.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/test/test_codegen.py b/test/test_codegen.py index 1e65f46..0f26bc0 100755 --- a/test/test_codegen.py +++ b/test/test_codegen.py @@ -49,7 +49,7 @@ def test_basic_codegen(ctx_factory): assert (out == x_in * x_in).all() -def test_DictOfNamedArrays_output(ctx_factory): # noqa +def test_codegen_with_DictOfNamedArrays(ctx_factory): # noqa ctx = ctx_factory() queue = cl.CommandQueue(ctx) -- GitLab From c0f7a4a207584380e7a299376ec7ba8da9ce2009 Mon Sep 17 00:00:00 2001 From: Matt Wala Date: Wed, 24 Jun 2020 14:48:43 -0500 Subject: [PATCH 76/76] Documentation improvements --- pytato/array.py | 3 +++ pytato/codegen.py | 12 ++++++------ 2 files changed, 9 insertions(+), 6 deletions(-) diff --git a/pytato/array.py b/pytato/array.py index 08d3329..e9bea17 100644 --- a/pytato/array.py +++ b/pytato/array.py @@ -62,6 +62,9 @@ Pre-Defined Tags .. autoclass:: ImplementAs .. autoclass:: CountNamed +.. autoclass:: ImplStored +.. autoclass:: ImplInlined +.. autoclass:: ImplDefault Built-in Expression Nodes ------------------------- diff --git a/pytato/codegen.py b/pytato/codegen.py index 5109529..5a7a7cb 100644 --- a/pytato/codegen.py +++ b/pytato/codegen.py @@ -82,7 +82,7 @@ class LoopyExpressionContext(object): """Mutable state used while generating :mod:`loopy` expressions. Wraps :class:`CodeGenState` with more expression-specific information. - This data is passed through :class:`LoopyExpressionGenMapper` via arguments, + This data is passed through :class:`InlinedExpressionGenMapper` via arguments, and is also used by :meth:`ImplementedResult.to_loopy_expression` to retrieve contextual data. @@ -98,7 +98,7 @@ class LoopyExpressionContext(object): .. attribute:: depends_on The set of statement IDs that need to be included in - :attr:`loopy.kernel.data.instruction.InstructionBase.depends_on`. + :attr:`loopy.InstructionBase.depends_on`. .. attribute:: reduction_bounds @@ -208,7 +208,7 @@ class CodeGenState: .. attribute:: results - A mapping from :class:`pytato.array.Array` instances to + A mapping from :class:`pytato.Array` instances to instances of :class:`ImplementedResult`. .. attribute:: var_name_gen @@ -284,8 +284,8 @@ class InlinedExpressionGenMapper(scalar_expr.IdentityMapper): sub-expressions. The inputs to this mapper are scalar expression as found in - :class:`pytato.IndexLambda`, or expressions that are compatible (e.g., shape - expressions). + :class:`pytato.array.IndexLambda`, or expressions that are + compatible (e.g., shape expressions). The outputs of this mapper are scalar expressions suitable for wrapping in :class:`InlinedResult`. @@ -322,7 +322,7 @@ class InlinedExpressionGenMapper(scalar_expr.IdentityMapper): def domain_for_shape( dim_names: Tuple[str, ...], shape: ShapeType) -> isl.BasicSet: - """Create a :class:`islpy.BasicSet` that expresses an appropriate index domain + """Create an :class:`islpy.BasicSet` that expresses an appropriate index domain for an array of (potentially symbolic) shape *shape*. :param dim_names: A tuple of strings, the names of the axes. These become set -- GitLab