implements inlined; codegen works now

Author: kaushikcfd

pytato/__init__.py

@@ -90,6 +90,7 @@ def set_debug_enabled(flag: bool) -> None:
 from pytato.visualization import (get_dot_graph, show_dot_graph,
                                   get_ascii_graph, show_ascii_graph,
                                   get_dot_graph_from_partition)
+from pytato.transform.calls import tag_all_calls_to_be_inlined, inline_calls
 import pytato.analysis as analysis
 import pytato.tags as tags
 import pytato.tracing as tracing
@@ -154,6 +155,8 @@ def set_debug_enabled(flag: bool) -> None:
         "make_distributed_recv", "make_distributed_send", "DistributedRecv",
         "DistributedSend", "staple_distributed_send", "DistributedSendRefHolder",
 
+        "tag_all_calls_to_be_inlined", "inline_calls",
+
         "find_distributed_partition",
         "number_distributed_tags",
         "execute_distributed_partition",

pytato/array.py

@@ -767,6 +767,17 @@ class AbstractResultWithNamedArrays(Mapping[str, NamedArray], Taggable, ABC):
     tags: FrozenSet[Tag] = attrs.field(kw_only=True)
     _mapper_method: ClassVar[str]
 
+    def _is_eq_valid(self) -> bool:
+        return self.__class__.__eq__ is AbstractResultWithNamedArrays.__eq__
+
+    def __post_init__(self) -> None:
+        # ensure that a developer does not uses dataclass' "__eq__"
+        # or "__hash__" implementation as they have exponential complexity.
+        assert self._is_eq_valid()
+
+    def __attrs_post_init__(self) -> None:
+        return self.__post_init__()
+
     @abstractmethod
     def __contains__(self, name: object) -> bool:
         pass
@@ -779,6 +790,13 @@ def __getitem__(self, name: str) -> NamedArray:
     def __len__(self) -> int:
         pass
 
+    def __eq__(self, other: Any) -> bool:
+        if self is other:
+            return True
+
+        from pytato.equality import EqualityComparer
+        return EqualityComparer()(self, other)
+
 
 @attrs.define(frozen=True, eq=False, init=False)
 class DictOfNamedArrays(AbstractResultWithNamedArrays):
@@ -825,13 +843,6 @@ def __len__(self) -> int:
     def __iter__(self) -> Iterator[str]:
         return iter(self._data)
 
-    def __eq__(self, other: Any) -> bool:
-        if self is other:
-            return True
-
-        from pytato.equality import EqualityComparer
-        return EqualityComparer()(self, other)
-
     def __repr__(self) -> str:
         return "DictOfNamedArrays(" + str(self._data) + ")"
 

pytato/codegen.py

@@ -118,6 +118,11 @@ def __init__(self, target: Target) -> None:
         self.target = target
         self.kernels_seen: Dict[str, lp.LoopKernel] = {}
 
+    def clone(self) -> CodeGenPreprocessor:
+        new_mapper = CodeGenPreprocessor(self.target)
+        new_mapper.kernels_seen = self.kernels_seen
+        return new_mapper
+
     def map_size_param(self, expr: SizeParam) -> Array:
         name = expr.name
         assert name is not None
@@ -602,6 +607,7 @@ def map_non_contiguous_advanced_index(self,
                            var_to_reduction_descr=Map(),
                            tags=expr.tags,
                            )
+
 # }}}
 
 
@@ -697,11 +703,11 @@ def preprocess(outputs: DictOfNamedArrays, target: Target) -> PreprocessResult:
 
     # }}}
 
-    mapper = CodeGenPreprocessor(target)
-
-    new_outputs = copy_dict_of_named_arrays(outputs, mapper)
     new_outputs = inline_calls(outputs)
 
+    mapper = CodeGenPreprocessor(target)
+    new_outputs = copy_dict_of_named_arrays(new_outputs, mapper)
+
     return PreprocessResult(outputs=new_outputs,
             compute_order=tuple(output_order),
             bound_arguments=mapper.bound_arguments)

pytato/equality.py

@@ -31,8 +31,7 @@
                           IndexBase, IndexLambda, NamedArray,
                           Reshape, Roll, Stack, AbstractResultWithNamedArrays,
                           Array, DictOfNamedArrays, Placeholder, SizeParam)
-from pytato.tracing import (TracePlaceholder, Call, NamedCallResult,
-                            FunctionDefinition)
+from pytato.tracing import Call, NamedCallResult, FunctionDefinition
 
 if TYPE_CHECKING:
     from pytato.loopy import LoopyCall, LoopyCallResult
@@ -276,10 +275,7 @@ def map_function_definition(self, expr1: FunctionDefinition, expr2: Any
                                 ) -> bool:
         return (expr1.__class__ is expr2.__class__
                 and expr1.identifier == expr2.identifier
-                and len(expr1.parameters.keys()) == len(expr2.parameters.keys())
-                and all(self.rec(expr1_pl, expr2_pl)
-                        for expr1_pl, expr2_pl in zip(expr1.parameters,
-                                                      expr2.parameters))
+                and expr1.parameters == expr2.parameters
                 and (set(expr1.returns.keys()) == set(expr2.returns.keys()))
                 and all(self.rec(expr1.returns[k], expr2.returns[k])
                         for k in expr1.returns)
@@ -299,16 +295,6 @@ def map_named_call_result(self, expr1: NamedCallResult, expr2: Any) -> bool:
                 and expr1.name == expr2.name
                 and self.rec(expr1._container, expr2._container))
 
-    def map_trace_placeholder(self, expr1: TracePlaceholder, expr2: Any) -> bool:
-        return (expr1.__class__ is expr2.__class__
-                and expr1.name == expr2.name
-                and expr1.identifier == expr2.identifier
-                and expr1.shape == expr2.shape
-                and expr1.dtype == expr2.dtype
-                and expr1.tags == expr2.tags
-                and expr1.axes == expr2.axes
-                )
-
 # }}}
 
 # vim: fdm=marker

pytato/loopy.py

@@ -118,19 +118,6 @@ def __len__(self) -> int:
     def __iter__(self) -> Iterator[str]:
         return iter(self._result_names)
 
-    def __eq__(self, other: Any) -> bool:
-        if self is other:
-            return True
-
-        if not isinstance(other, LoopyCall):
-            return False
-
-        if ((self.entrypoint == other.entrypoint)
-             and (self.bindings == other.bindings)
-             and (self.translation_unit == other.translation_unit)):
-            return True
-        return False
-
 
 class LoopyCallResult(NamedArray):
     """

pytato/tracing.py

@@ -38,8 +38,8 @@
 
 import attrs
 
-from typing import (Callable, Dict, FrozenSet, Tuple, Union, TypeVar, Optional, Hashable,
-                    Sequence, ClassVar)
+from typing import (Callable, Dict, FrozenSet, Tuple, Union, TypeVar, Optional,
+                    Hashable, Sequence, ClassVar)
 from immutables import Map
 from pytato.array import (Array,  AbstractResultWithNamedArrays, AxesT,
                           Placeholder, NamedArray, ShapeType, _dtype_any)
@@ -56,7 +56,7 @@
 
 # {{{ Call/NamedCallResult
 
-@attrs.define(frozen=True, eq=False, repr=False)
+@attrs.define(frozen=True, repr=False, eq=False)
 class FunctionDefinition(AbstractResultWithNamedArrays):
     r"""
     A function definition that represents its outputs as instances of
@@ -67,9 +67,7 @@ class FunctionDefinition(AbstractResultWithNamedArrays):
 
     .. attribute:: parameters
 
-        The inputs to the function node. A mapping from the names of the
-        function's parameters to its corresponding placeholder in the
-        outputs' expression graph.
+        Names of the inputs to the function node.
 
     .. attribute:: returns
 
@@ -85,16 +83,11 @@ class FunctionDefinition(AbstractResultWithNamedArrays):
           :class:`pytato.Array`\ s, then *returns* uses the same mapping.
     """
     identifier: Hashable
-    parameters: Map[str, Placeholder]
+    parameters: FrozenSet[str]
     returns: Map[str, Array]
 
     _mapper_method: ClassVar[str] = "map_function_definition"
 
-    def __post_init__(self):
-        if __debug__:
-            assert all(name == param.name
-                       for name, param in self.parameters.items())
-
     def __contains__(self, name: object) -> bool:
         return name in self.returns
 
@@ -189,15 +182,17 @@ class Call(AbstractResultWithNamedArrays):
     result_tags: Map[str, FrozenSet[Tag]] = attrs.field(kw_only=True)
     result_axes: Map[str, AxesT] = attrs.field(kw_only=True)
 
+    _mapper_method: ClassVar[str] = "map_call"
+
+    copy = attrs.evolve
+
     def __post_init__(self):
         if __debug__:
             # check that the invocation parameters and the function definition
             # parameters agree with each other.
-            assert set(self.bindings) == set(self.function.parameters.keys())
-            assert set(self.tags.keys()) == set(self.function.returns.keys())
-            assert set(self.axes.keys()) == set(self.function.returns.keys())
-
-    _mapper_method: ClassVar[str] = "map_call"
+            assert frozenset(self.bindings) == self.function.parameters
+            assert set(self.result_tags.keys()) == set(self.function.returns.keys())
+            assert set(self.result_axes.keys()) == set(self.function.returns.keys())
 
     def __contains__(self, name: object) -> bool:
         return name in self.function.returns
@@ -211,8 +206,8 @@ def __getitem__(self, name: str) -> NamedCallResult:
     def __len__(self) -> int:
         return len(self.function.returns)
 
-    def with_tagged_axis(self, name: str, iaxis: int,
-                         tags: Union[Sequence[Tag], Tag]) -> Self:
+    def with_result_axis_tagged(self, name: str, iaxis: int,
+                                tags: Union[Sequence[Tag], Tag]) -> Self:
         """
         Returns a copy of *self* with the result corresponding to *name*\'s
         *iaxis*-th axis tagged with *tags*. Also, see
@@ -230,14 +225,15 @@ def with_tagged_axis(self, name: str, iaxis: int,
         return replace(self,
                        result_axes=self.result_axes.set(name, new_axes))
 
-    def tagged(self, name: str, tags: Union[Sequence[Tag], Tag]) -> Self:
+    def with_result_tagged(self, name: str, tags: Union[Sequence[Tag], Tag]) -> Self:
         """
         Returns a copy of *self* with the result corresponding to *name*\'s
         tagged with *tags*. Also, see :meth:`pytato.Array.tagged`.
         """
         from attrs import evolve as replace
         from pytools.tag import check_tag_uniqueness, normalize_tags
-        new_tags = check_tag_uniqueness(normalize_tags(tags) | self.result_tags[name])
+        new_tags = check_tag_uniqueness(normalize_tags(tags)
+                                        | self.result_tags[name])
 
         return replace(self, result_tags=self.result_tags.set(name, new_tags))
 
@@ -299,8 +295,7 @@ def trace_call(f: Callable[..., ReturnT],
     # construct the function
     function = FunctionDefinition(
         identifier,
-        (Map({pl.name: pl for pl in pl_args})
-         .update(Map({pl.name: pl for pl in pl_kwargs.values()}))),
+        frozenset(pl_arg.name for pl_arg in pl_args) | frozenset(pl_kwargs),
         Map(returns),
         tags=_get_default_tags()
     )