diff --git a/doc/codegen.rst b/doc/codegen.rst
new file mode 100644
index 0000000000000000000000000000000000000000..917751c84b8cdc5c96c039bbd57985f570915e60
--- /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 18dd4da2c519885b2fb0c93120df7f7db82d923b..187dda3a57237a8239734674b06f72e8c0198b6f 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 4cef6a4a58a5fbacf76340cefe8574fdb45446de..dcfaf28c87db410713c6d97ba67fa8565baeaf62 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 e265fbeff8b82724f99c5e77c2a1451a1067e5cc..b3bf553369d7ffa18ed76f06adc6a92da8621b7f 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 4a3660202611a9bbfd08ed766534c9cec25f1fb2..4e0506fba37c43e61f7835e19846d326d33032ea 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 82531cc621a422514f005fc22fa12acb68e3b25d..2fe9afc4da9c6a2560b4c8c2758cdf98ffc6748f 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 7c3f3d4f4d3899bdd0e89aae18a5625b2e8ac334..b0d0924987b6cad885c759657f6fbc171aa59c4f 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 0000000000000000000000000000000000000000..e93200d922f1ec3c3dc14bd298fd5a3e978891dc
--- /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