Mapper typing: Array->Array, Names->Names (#610)

* Mapper typing: Array->Array, Names->Names

* more typing work

* More type fixes

* MappedT -> ArrayOrNamesT

* handle DictOfNamedArrays in materialize_with_mpms outside of the mapper

* a few more tweaks

* add deduplicate

* fix doc

* Update baseline

---------

Co-authored-by: Matthew Smith <mjsmith6@illinois.edu>

Author: inducer

.basedpyright/baseline.json

@@ -156,7 +156,7 @@ def test_advection_convergence(order, flux_type):
         op = AdvectionOperator(discr, c=1, flux_type=flux_type,
                                dg_ops=dg_ops)
         result = op.apply(u)
-        result = pt.transform.Deduplicator()(result)
+        result = pt.transform.deduplicate(result)
 
         prog = pt.generate_loopy(result, cl_device=queue.device)
 

examples/advection.py

@@ -157,6 +157,8 @@ reportImplicitStringConcatenation = "none"
 reportUnnecessaryIsInstance = "none"
 reportUnusedCallResult = "none"
 reportExplicitAny = "none"
+reportPrivateUsage = "hint"
+reportUnusedParameter = "hint"
 
 # This reports even cycles that are qualified by 'if TYPE_CHECKING'. Not what
 # we care about at this moment.
@@ -189,6 +191,8 @@ reportArgumentType = "hint"
 reportUnknownMemberType = "hint"
 reportUnknownParameterType = "hint"
 reportAny = "none"
+reportPrivateUsage = "hint"
+reportUnusedParameter = "hint"
 
 [[tool.basedpyright.executionEnvironments]]
 root = "examples"
@@ -201,6 +205,8 @@ reportArgumentType = "hint"
 reportUnknownMemberType = "hint"
 reportUnknownParameterType = "hint"
 reportAny = "none"
+reportPrivateUsage = "hint"
+reportUnusedParameter = "hint"
 
 [tool.typos.default]
 extend-ignore-re = [

pyproject.toml

@@ -211,13 +211,11 @@ def map_named_call_result(self, expr: NamedCallResult) -> None:
 # }}}
 
 
-def get_nusers(outputs: Array | DictOfNamedArrays) -> Mapping[Array, int]:
+def get_nusers(outputs: ArrayOrNames) -> Mapping[Array, int]:
     """
-    For the DAG *outputs*, returns the mapping from each node to the number of
+    For the DAG *outputs*, returns the mapping from each array node to the number of
     nodes using its value within the DAG given by *outputs*.
     """
-    from pytato.codegen import normalize_outputs
-    outputs = normalize_outputs(outputs)
     nuser_collector = NUserCollector()
     nuser_collector(outputs)
     return nuser_collector.nusers
@@ -508,7 +506,7 @@ def post_visit(self, expr: Any) -> None:
 
 
 def get_node_type_counts(
-        outputs: Array | DictOfNamedArrays,
+        outputs: ArrayOrNames,
         count_duplicates: bool = False
         ) -> dict[type[Any], int]:
     """
@@ -518,17 +516,14 @@ def get_node_type_counts(
     Instances of `DictOfNamedArrays` are excluded from counting.
     """
 
-    from pytato.codegen import normalize_outputs
-    outputs = normalize_outputs(outputs)
-
     ncm = NodeCountMapper(count_duplicates)
     ncm(outputs)
 
     return ncm.expr_type_counts
 
 
 def get_num_nodes(
-        outputs: Array | DictOfNamedArrays,
+        outputs: ArrayOrNames,
         count_duplicates: bool | None = None
         ) -> int:
     """
@@ -544,9 +539,6 @@ def get_num_nodes(
             DeprecationWarning, stacklevel=2)
         count_duplicates = True
 
-    from pytato.codegen import normalize_outputs
-    outputs = normalize_outputs(outputs)
-
     ncm = NodeCountMapper(count_duplicates)
     ncm(outputs)
 
@@ -586,14 +578,10 @@ def post_visit(self, expr: Any) -> None:
             self.expr_multiplicity_counts[expr] += 1
 
 
-def get_node_multiplicities(
-        outputs: Array | DictOfNamedArrays) -> dict[Array, int]:
+def get_node_multiplicities(outputs: ArrayOrNames) -> dict[Array, int]:
     """
     Returns the multiplicity per `expr`.
     """
-    from pytato.codegen import normalize_outputs
-    outputs = normalize_outputs(outputs)
-
     nmm = NodeMultiplicityMapper()
     nmm(outputs)
 
@@ -640,7 +628,7 @@ def post_visit(self, expr: Any) -> None:
             self.count += 1
 
 
-def get_num_call_sites(outputs: Array | DictOfNamedArrays) -> int:
+def get_num_call_sites(outputs: ArrayOrNames) -> int:
     """Returns the number of nodes in DAG *outputs*."""
 
     from pytato.codegen import normalize_outputs
@@ -700,7 +688,7 @@ def rec(self, expr: ArrayOrNames) -> int:
 
 
 def get_num_tags_of_type(
-        outputs: Array | DictOfNamedArrays,
+        outputs: ArrayOrNames,
         tag_types: type[pytools.tag.Tag] | Iterable[type[pytools.tag.Tag]]) -> int:
     """Returns the number of nodes in DAG *outputs* that are tagged with
     all the tag types in *tag_types*."""

pytato/analysis/__init__.py

@@ -10,10 +10,6 @@
 
 from __future__ import annotations
 
-from typing import TypeAlias
-
-from typing_extensions import TypeIs
-
 
 __copyright__ = """Copyright (C) 2020 Matt Wala"""
 
@@ -38,15 +34,17 @@
 """
 
 import dataclasses
-from typing import TYPE_CHECKING, Any
+from typing import TYPE_CHECKING, Any, TypeAlias
 
 from immutabledict import immutabledict
+from typing_extensions import TypeIs
 
 import loopy as lp
 from pymbolic.mapper.optimize import optimize_mapper
 from pytools import UniqueNameGenerator
 
 from pytato.array import (
+    AbstractResultWithNamedArrays,
     Array,
     DataInterface,
     DataWrapper,
@@ -241,24 +239,19 @@ def map_named_call_result(self, expr: NamedCallResult) -> Array:
 
 
 def normalize_outputs(
-            result: Array | DictOfNamedArrays | dict[str, Array]
-        ) -> DictOfNamedArrays:
+            result: ArrayOrNames | dict[str, Array]
+        ) -> AbstractResultWithNamedArrays:
     """Convert outputs of a computation to the canonical form.
 
     Performs a conversion to :class:`~pytato.DictOfNamedArrays` if necessary.
 
     :param result: Outputs of the computation.
     """
-    if not isinstance(result, Array | DictOfNamedArrays | dict):
-        raise TypeError("outputs of the computation should be "
-                "either an Array or a DictOfNamedArrays")
-
     if isinstance(result, Array):
         outputs = make_dict_of_named_arrays({"_pt_out": result})
     elif isinstance(result, dict):
         outputs = make_dict_of_named_arrays(result)
     else:
-        assert isinstance(result, DictOfNamedArrays)
         outputs = result
 
     return outputs

pytato/codegen.py

@@ -41,6 +41,7 @@
 import pytato.reductions as red
 import pytato.scalar_expr as scalar_expr
 from pytato.array import (
+    AbstractResultWithNamedArrays,
     Array,
     DataWrapper,
     DictOfNamedArrays,
@@ -1036,7 +1037,7 @@ def get_initial_codegen_state(target: LoopyTarget,
 
 # {{{ generate_loopy
 
-def generate_loopy(result: Array | DictOfNamedArrays | dict[str, Array],
+def generate_loopy(result: Array | AbstractResultWithNamedArrays | dict[str, Array],
                    target: LoopyTarget | None = None,
                    options: lp.Options | None = None,
                    *,
@@ -1083,7 +1084,11 @@ def generate_loopy(result: Array | DictOfNamedArrays | dict[str, Array],
     """
 
     result_is_dict = isinstance(result, dict | DictOfNamedArrays)
-    orig_outputs: DictOfNamedArrays = normalize_outputs(result)
+    orig_outputs: AbstractResultWithNamedArrays = normalize_outputs(result)
+
+    if not isinstance(orig_outputs, DictOfNamedArrays):
+        raise NotImplementedError(
+            f"not implemented for {type(result).__name__}.")
 
     del result
 

pytato/target/loopy/codegen.py

@@ -42,7 +42,7 @@
 
 import numpy as np
 from immutabledict import immutabledict
-from typing_extensions import Never, Self
+from typing_extensions import Never, Self, override
 
 from pymbolic.mapper.optimize import optimize_mapper
 
@@ -90,8 +90,11 @@
 
 
 ArrayOrNames: TypeAlias = Array | AbstractResultWithNamedArrays
-MappedT = TypeVar("MappedT",
-                  Array, AbstractResultWithNamedArrays, ArrayOrNames)
+ArrayOrNamesTc = TypeVar("ArrayOrNamesTc",
+                  Array, AbstractResultWithNamedArrays, DictOfNamedArrays)
+ArrayOrNamesOrFunctionDefTc = TypeVar("ArrayOrNamesOrFunctionDefTc",
+                  Array, AbstractResultWithNamedArrays, DictOfNamedArrays,
+                  FunctionDefinition)
 IndexOrShapeExpr = TypeVar("IndexOrShapeExpr")
 R = frozenset[Array]
 
@@ -116,6 +119,7 @@
 .. autoclass:: TopoSortMapper
 .. autoclass:: CachedMapAndCopyMapper
 .. autofunction:: copy_dict_of_named_arrays
+.. autofunction:: deduplicate
 .. autofunction:: get_dependencies
 .. autofunction:: map_and_copy
 .. autofunction:: materialize_with_mpms
@@ -145,9 +149,15 @@
 
 .. class:: ArrayOrNames
 
-.. class:: MappedT
+.. class:: ArrayOrNamesTc
 
-    A type variable representing the input type of a :class:`Mapper`.
+    A type variable representing the input type of a :class:`Mapper`, excluding
+    functions.
+
+.. class:: ArrayOrNamesOrFunctionDefTc
+
+    A type variable representing the input type of a :class:`Mapper`, including
+    functions.
 
 .. class:: ResultT
 
@@ -702,6 +712,22 @@ def clone_for_callee(self, function: FunctionDefinition) -> Self:
             err_on_created_duplicate=function_cache.err_on_created_duplicate,
             _function_cache=function_cache)
 
+    @override
+    # This overrides incompatibly on purpose, in order to convey stronger
+    # guarantees. We're not trying to be very mapper-polymorphic, so
+    # IMO this inconsistency is "worth it(tm)". -AK, 2025-06-16
+    def __call__(  # pyright: ignore[reportIncompatibleMethodOverride]
+            self,
+            expr: ArrayOrNamesOrFunctionDefTc,
+            ) -> ArrayOrNamesOrFunctionDefTc:
+        """Handle the mapping of *expr*."""
+        if isinstance(expr, Array):
+            return cast("Array", self.rec(expr))
+        elif isinstance(expr, AbstractResultWithNamedArrays):
+            return cast("AbstractResultWithNamedArrays", self.rec(expr))
+        else:
+            return self.rec_function_definition(expr)
+
 # }}}
 
 
@@ -1116,7 +1142,7 @@ def map_named_call_result(self, expr: NamedCallResult,
 # }}}
 
 
-# {{{ Deduplicator
+# {{{ deduplicate
 
 class Deduplicator(CopyMapper):
     """Removes duplicate nodes from an expression."""
@@ -1135,6 +1161,20 @@ def clone_for_callee(self, function: FunctionDefinition) -> Self:
             _function_cache=cast(
                 "TransformMapperCache[FunctionDefinition, []]", self._function_cache))
 
+
+def deduplicate(
+        expr: ArrayOrNamesOrFunctionDefTc
+        ) -> ArrayOrNamesOrFunctionDefTc:
+    """
+    Remove duplicate nodes from an expression.
+
+    .. note::
+        Does not remove distinct instances of data wrappers that point to the same
+        data (as they will not hash the same). For a utility that does that, see
+        :func:`deduplicate_data_wrappers`.
+    """
+    return Deduplicator()(expr)
+
 # }}}
 
 
@@ -2061,9 +2101,9 @@ def get_dependencies(expr: DictOfNamedArrays) -> dict[str, frozenset[Array]]:
     return {name: dep_mapper(val.expr) for name, val in expr.items()}
 
 
-def map_and_copy(expr: MappedT,
+def map_and_copy(expr: ArrayOrNamesTc,
                  map_fn: Callable[[ArrayOrNames], ArrayOrNames]
-                 ) -> MappedT:
+                 ) -> ArrayOrNamesTc:
     """
     Returns a copy of *expr* with every array expression reachable from *expr*
     mapped via *map_fn*.
@@ -2073,10 +2113,10 @@ def map_and_copy(expr: MappedT,
         Uses :class:`CachedMapAndCopyMapper` under the hood and because of its
         caching nature each node is mapped exactly once.
     """
-    return cast("MappedT", CachedMapAndCopyMapper(map_fn)(expr))
+    return CachedMapAndCopyMapper(map_fn)(expr)
 
 
-def materialize_with_mpms(expr: DictOfNamedArrays) -> DictOfNamedArrays:
+def materialize_with_mpms(expr: ArrayOrNamesTc) -> ArrayOrNamesTc:
     r"""
     Materialize nodes in *expr* with MPMS materialization strategy.
     MPMS stands for Multiple-Predecessors, Multiple-Successors.
@@ -2128,11 +2168,18 @@ def materialize_with_mpms(expr: DictOfNamedArrays) -> DictOfNamedArrays:
     """
     from pytato.analysis import get_num_nodes, get_num_tags_of_type, get_nusers
     materializer = MPMSMaterializer(get_nusers(expr))
-    new_data = {}
-    for name, ary in expr.items():
-        new_data[name] = materializer(ary.expr).expr
 
-    res = DictOfNamedArrays(new_data, tags=expr.tags)
+    if isinstance(expr, Array):
+        res = materializer(expr).expr
+        assert isinstance(res, Array)
+    elif isinstance(expr, DictOfNamedArrays):
+        res = expr.replace_if_different(
+            data={
+                name: _verify_is_array(materializer(ary).expr)
+                for name, ary, in expr._data.items()})
+        assert isinstance(res, DictOfNamedArrays)
+    else:
+        raise NotImplementedError("not implemented for {type(expr).__name__}.")
 
     from pytato import DEBUG_ENABLED
     if DEBUG_ENABLED:
@@ -2408,7 +2455,9 @@ def clone_for_callee(self, function: FunctionDefinition) -> Self:
                 "TransformMapperCache[FunctionDefinition, []]", self._function_cache))
 
 
-def deduplicate_data_wrappers(array_or_names: ArrayOrNames) -> ArrayOrNames:
+def deduplicate_data_wrappers(
+            array_or_names: ArrayOrNamesOrFunctionDefTc
+        ) -> ArrayOrNamesOrFunctionDefTc:
     """For the expression graph given as *array_or_names*, replace all
     :class:`pytato.array.DataWrapper` instances containing identical data
     with a single instance.