change is_materializable -> is_materialized / has_taggable_materialization

Author: majosm

pytato/analysis/__init__.py

@@ -51,6 +51,7 @@
     ShapeType,
     Stack,
 )
+from pytato.diagnostic import CannotBeLoweredToIndexLambda
 from pytato.function import Call, FunctionDefinition, NamedCallResult
 from pytato.scalar_expr import (
     FlopCounter as ScalarFlopCounter,
@@ -66,7 +67,7 @@
     map_and_copy,
 )
 from pytato.transform.lower_to_index_lambda import to_index_lambda
-from pytato.utils import is_materializable
+from pytato.utils import has_taggable_materialization, is_materialized
 
 
 if TYPE_CHECKING:
@@ -758,18 +759,6 @@ def update_for_Array(self, key_hash: Any, key: Any) -> None:
 
 # {{{ flop counting
 
-def _is_materialized(expr: ArrayOrNames | FunctionDefinition) -> bool:
-    return (
-        is_materializable(expr)
-        and bool(expr.tags_of_type(ImplStored)))
-
-
-def _is_unmaterialized(expr: ArrayOrNames | FunctionDefinition) -> bool:
-    return (
-        is_materializable(expr)
-        and not bool(expr.tags_of_type(ImplStored)))
-
-
 @dataclass
 class UndefinedOpFlopCountError(ValueError):
     op_name: str
@@ -787,7 +776,10 @@ def combine(self, *args: int) -> int:
         return sum(args)
 
     def _get_own_flop_count(self, expr: Array) -> int:
-        nflops = self.scalar_flop_counter(to_index_lambda(expr).expr)
+        try:
+            nflops = self.scalar_flop_counter(to_index_lambda(expr).expr)
+        except CannotBeLoweredToIndexLambda:
+            nflops = 0
         if not isinstance(nflops, int):
             from pytato.scalar_expr import InputGatherer as ScalarInputGatherer
             var_names: set[str] = set(ScalarInputGatherer()(nflops))
@@ -805,8 +797,7 @@ def rec(self, expr: ArrayOrNames) -> int:
             return self._cache_retrieve(inputs)
         except KeyError:
             result: int
-            if _is_unmaterialized(expr):
-                assert isinstance(expr, Array)
+            if isinstance(expr, Array) and not is_materialized(expr):
                 result = (
                     self._get_own_flop_count(expr)
                     # Intentionally going to Mapper instead of super() to avoid
@@ -866,14 +857,16 @@ def map_call(self, expr: Call) -> None:
 
     @override
     def post_visit(self, expr: ArrayOrNames | FunctionDefinition) -> None:
-        if not _is_materialized(expr):
+        if not is_materialized(expr):
             return
         assert isinstance(expr, Array)
-        unmaterialized_expr = expr.without_tags(ImplStored())
-        self._per_entry_flop_counter(unmaterialized_expr)
-        self.materialized_node_to_nflops[expr] = (
-            product(expr.shape)
-            * self._per_entry_flop_counter.node_to_nflops[unmaterialized_expr])
+        if has_taggable_materialization(expr):
+            unmaterialized_expr = expr.without_tags(ImplStored())
+            self.materialized_node_to_nflops[expr] = (
+                product(expr.shape)
+                * self._per_entry_flop_counter(unmaterialized_expr))
+        else:
+            self.materialized_node_to_nflops[expr] = 0
 
 
 class _UnmaterializedSubexpressionUseCounter(CombineMapper[dict[Array, int], Never]):
@@ -892,8 +885,7 @@ def rec(self, expr: ArrayOrNames) -> dict[Array, int]:
             return self._cache_retrieve(inputs)
         except KeyError:
             result: dict[Array, int]
-            if _is_unmaterialized(expr):
-                assert isinstance(expr, Array)
+            if isinstance(expr, Array) and not is_materialized(expr):
                 # Intentionally going to Mapper instead of super() to avoid
                 # double caching when subclasses of CachedMapper override rec,
                 # see https://github.com/inducer/pytato/pull/585
@@ -959,7 +951,7 @@ def map_call(self, expr: Call) -> None:
 
     @override
     def post_visit(self, expr: ArrayOrNames | FunctionDefinition) -> None:
-        if not _is_materialized(expr):
+        if not is_materialized(expr) or not has_taggable_materialization(expr):
             return
         assert isinstance(expr, Array)
         unmaterialized_expr = expr.without_tags(ImplStored())
@@ -985,7 +977,10 @@ def _normalize_materialization(expr: ArrayOrNamesTc) -> ArrayOrNamesTc:
     # Make sure outputs are materialized
     if isinstance(expr, DictOfNamedArrays):
         output_to_materialized_output: dict[Array, Array] = {
-            ary: ary.tagged(ImplStored()) if is_materializable(ary) else ary
+            ary: (
+                ary.tagged(ImplStored())
+                if has_taggable_materialization(ary)
+                else ary)
             for ary in expr._data.values()}
 
         def replace_with_materialized(ary: ArrayOrNames) -> ArrayOrNames:

pytato/utils.py

@@ -82,7 +82,8 @@
 .. autofunction:: dim_to_index_lambda_components
 .. autofunction:: get_common_dtype_of_ary_or_scalars
 .. autofunction:: get_einsum_subscript_str
-.. autofunction:: is_materializable
+.. autofunction:: is_materialized
+.. autofunction:: has_taggable_materialization
 
 References
 ^^^^^^^^^^
@@ -739,18 +740,40 @@ def get_einsum_specification(expr: Einsum) -> str:
     return f"{','.join(input_specs)}->{output_spec}"
 
 
-def is_materializable(expr: ArrayOrNames | FunctionDefinition) -> bool:
+def is_materialized(expr: ArrayOrNames | FunctionDefinition) -> bool:
+    """Returns *True* if *expr* is materialized."""
+    from pytato.array import InputArgumentBase
+    from pytato.distributed.nodes import DistributedRecv
+    from pytato.tags import ImplStored
+    return (
+        (
+            isinstance(expr, Array)
+            and bool(expr.tags_of_type(ImplStored)))
+        or isinstance(
+            expr,
+            (
+                # FIXME: Is there a nice way to generalize this?
+                InputArgumentBase,
+                DistributedRecv)))
+
+
+def has_taggable_materialization(expr: ArrayOrNames | FunctionDefinition) -> bool:
     """
-    Returns *True* if *expr* is an instance of an array type that can be materialized.
+    Returns *True* if *expr* uses the :class:`pytato.tags.ImplStored` tag to
+    determine whether or not it is materialized.
     """
     from pytato.array import InputArgumentBase, NamedArray
     from pytato.distributed.nodes import DistributedRecv, DistributedSendRefHolder
     return (
         isinstance(expr, Array)
-        and not isinstance(expr, (
-            # FIXME: Is there a nice way to generalize this?
-            InputArgumentBase, NamedArray, DistributedRecv,
-            DistributedSendRefHolder)))
+        and not isinstance(
+            expr,
+            (
+                # FIXME: Is there a nice way to generalize this?
+                InputArgumentBase,
+                DistributedRecv,
+                NamedArray,
+                DistributedSendRefHolder)))
 
 
 # vim: fdm=marker

test/test_pytato.py

@@ -980,9 +980,13 @@ def test_materialized_node_flop_counts():
 
     #    z[i, j] = x[i, j] + y[i, j]
     # expr[i, j] = 2*z[i, j] + (-1)*3*z[i, j]
-    assert len(materialized_node_to_flop_count) == 2
+    assert len(materialized_node_to_flop_count) == 4
+    assert x in materialized_node_to_flop_count
+    assert y in materialized_node_to_flop_count
     assert z in materialized_node_to_flop_count
     assert expr.tagged(ImplStored()) in materialized_node_to_flop_count
+    assert materialized_node_to_flop_count[x] == 0
+    assert materialized_node_to_flop_count[y] == 0
     assert materialized_node_to_flop_count[z] == 40
     assert materialized_node_to_flop_count[expr.tagged(ImplStored())] == 40*4