Remove the raising to high level operator within Unify Axis

pytato/scalar_expr.py

@@ -163,7 +163,9 @@ def map_reduce(self, expr: Reduce) -> ScalarExpression:
                            for name, bound in expr.bounds.items()}))
 
 
-IDX_LAMBDA_RE = re.compile(r"_r?(0|([1-9][0-9]*))")
+IDX_LAMBDA_RE = re.compile(r"^(_r?(0|([1-9][0-9]*)))$")
+IDX_LAMBDA_INAME = re.compile(r"^(_(0|([1-9][0-9]*)))$")
+IDX_LAMBDA_JUST_REDUCTIONS = re.compile(r"^(_r(0|([1-9][0-9]*)))$")
 
 
 class DependencyMapper(DependencyMapperBase[P]):

pytato/transform/metadata.py

@@ -38,25 +38,27 @@
 THE SOFTWARE.
 """
 
-
 import logging
+import re
 from typing import (
     TYPE_CHECKING,
     Any,
+    ParamSpec,
+    TypeAlias,
     TypeVar,
     cast,
 )
 
 from bidict import bidict
 
+import pymbolic.primitives as prim
 from pytools import UniqueNameGenerator
 from pytools.tag import Tag
 
 from pytato.array import (
     AbstractResultWithNamedArrays,
     AdvancedIndexInContiguousAxes,
     Array,
-    ArrayOrScalar,
     AxisPermutation,
     BasicIndex,
     Concatenate,
@@ -70,18 +72,12 @@
     Reshape,
     Stack,
 )
-from pytato.diagnostic import UnknownIndexLambdaExpr
 from pytato.distributed.nodes import DistributedRecv, DistributedSendRefHolder
-from pytato.raising import (
-    BinaryOp,
-    BroadcastOp,
-    C99CallOp,
-    FullOp,
-    ReduceOp,
-    WhereOp,
-    index_lambda_to_high_level_op,
+from pytato.function import NamedCallResult
+from pytato.scalar_expr import (
+    IDX_LAMBDA_INAME,
+    CombineMapper,
 )
-from pytato.scalar_expr import SCALAR_CLASSES
 from pytato.transform import ArrayOrNames, CopyMapper, Mapper
 from pytato.utils import are_shape_components_equal, are_shapes_equal
 
@@ -90,17 +86,63 @@
 
 
 if TYPE_CHECKING:
-    from collections.abc import Collection, Mapping
+    from collections.abc import Collection, Iterable, Mapping
 
     from pytato.function import NamedCallResult
     from pytato.loopy import LoopyCall
 
 
 GraphNodeT = TypeVar("GraphNodeT")
 
+BindingName: TypeAlias = str
+P = ParamSpec("P")
+
+BINDING_NAME_RESERVED_PATTERN = re.compile(r"^(_in?(0|([1-9][0-9]*)))$")
+
+
+# {{{ BindingSubscriptsCollector
+
+
+class BindingSubscriptsCollector(CombineMapper[dict[BindingName,
+                                               set[tuple[prim.Expression, ...]]],
+                                                      []]):
+    """
+    Return all the subscript expressions used by a variable specified by BindingName.
+    Ex:
+    _in1[_0,_1] would result in an dictionary entry {"_in1": ("_0", "_1")}.
+    """
+    def combine(self,
+                values: Iterable[dict[BindingName,
+                                         set[tuple[prim.Expression, ...]]]]) \
+                        -> dict[BindingName, set[tuple[prim.Expression, ...]]]:
+        out: dict[BindingName, set[tuple[prim.Expression, ...]]] = {}
+        from functools import reduce
+        return reduce(lambda x, y: x | y, values, out)
+
+    def map_subscript(self, expr: prim.Subscript) -> dict[BindingName,
+                                                    set[tuple[prim.Expression, ...]]]:
+        """
+        Record the indexing expression if the Subscript expression has a prim.Variable
+        as its aggregate.
+        """
+
+        if isinstance(expr.aggregate, prim.Variable):
+            return {expr.aggregate.name: {expr.index_tuple}}
+        return {}
+
+    def map_algebraic_leaf(self, expr: prim.Expression) -> dict[BindingName,
+                                                    set[tuple[prim.Expression, ...]]]:
+
+        return {}
+
+    def map_constant(self, expr: object) -> dict[BindingName,
+                                                    set[tuple[prim.Expression, ...]]]:
+        return {}
+# }}}
 
 # {{{ AxesTagsEquationCollector
 
+
 class AxesTagsEquationCollector(Mapper[None, []]):
     r"""
     Records equations arising from operand/output axes equivalence for an array
@@ -156,7 +198,7 @@ def __init__(self, tag_t: type[Tag]) -> None:
 
         # axis_to_var: mapping from (array, iaxis) to the variable to be
         # used for unification.
-        self.axis_to_var: bidict[tuple[Array, int], str] = bidict()
+        self.axis_to_var: bidict[tuple[Array, int | str], str] = bidict()
         self.known_tag_to_var: dict[Tag, str] = {}
 
         self.equations: list[tuple[str, str]] = []
@@ -165,7 +207,7 @@ def __init__(self, tag_t: type[Tag]) -> None:
 
     # {{{ unification helpers
 
-    def get_var_for_axis(self, ary: Array, iaxis: int) -> str:
+    def get_var_for_axis(self, ary: Array, iaxis: int | str) -> str:
         key = (ary, iaxis)
 
         try:
@@ -227,75 +269,34 @@ def _map_input_base(self, expr: InputArgumentBase) -> None:
 
     def map_index_lambda(self, expr: IndexLambda) -> None:
         """
-        The propagation semantics for a :class:`~pytato.IndexLambda` are
-        implemented only for operations that can be raised to a
-        :class:`~pytato.raising.HighLevelOp`. In such cases, an equality
-        equation is recorded for every non-broadcasted axis of an operand and
-        its corresponding axis of *expr*.
+        Equality conditions are added between an axis of the operands which is indexed
+        by a :class:`~pymbolic.Variable` which has a name that follows the reserved
+        iname format, "_[0-9]+", and the axis of the output specified by the iname.
         """
         for bnd in expr.bindings.values():
             self.rec(bnd)
 
-        try:
-            hlo = index_lambda_to_high_level_op(expr)
-        except UnknownIndexLambdaExpr:
-            from warnings import warn
-            warn(f"'{expr}' is an unknown index lambda type"
-                 " no tags were propagated across it.", stacklevel=1)
-            # no propagation semantics implemented for such cases
-            return
-
-        if isinstance(hlo, BinaryOp):
-            subexprs: tuple[ArrayOrScalar, ...] = (hlo.x1, hlo.x2)
-        elif isinstance(hlo, WhereOp):
-            subexprs = (hlo.condition, hlo.then, hlo.else_)
-        elif isinstance(hlo, FullOp):
-            # A full-op does not impose any equations
-            subexprs = ()
-        elif isinstance(hlo, BroadcastOp):
-            subexprs = (hlo.x,)
-        elif isinstance(hlo, C99CallOp):
-            subexprs = hlo.args
-        elif isinstance(hlo, ReduceOp):
-
-            # {{{ ReduceOp doesn't quite involve broadcasting
-
-            i_out_axis = 0
-            for i_in_axis in range(hlo.x.ndim):
-                if i_in_axis not in hlo.axes:
-                    self.record_equation(
-                        self.get_var_for_axis(hlo.x, i_in_axis),
-                        self.get_var_for_axis(expr, i_out_axis)
-                    )
-                    i_out_axis += 1
+        index_expr_used = BindingSubscriptsCollector()(expr.expr)
 
-            assert i_out_axis == expr.ndim
+        assert len(expr.shape) == expr.ndim
 
-            # }}}
+        for vname, set_of_ind_tuple in index_expr_used.items():
+            for ind_tuple in set_of_ind_tuple:
+                for axis_ind, var_ind_name in enumerate(ind_tuple):
+                    if isinstance(var_ind_name, prim.Variable):
+                        lhs: str = self.get_var_for_axis(expr.bindings[vname],
+                                                         axis_ind)
+                        if IDX_LAMBDA_INAME.fullmatch(var_ind_name.name):
+                            # matched with an iname.
+                            inum = int(var_ind_name.name[1:])
+                            rhs: str = self.get_var_for_axis(expr, inum)
+                            self.record_equation(lhs, rhs)
+                        elif var_ind_name.name in expr.var_to_reduction_descr.keys():
+                            # matched with a reduction iname.
+                            rhs = self.get_var_for_axis(expr, var_ind_name.name)
+                            self.record_equation(lhs, rhs)
 
-            return
-
-        else:
-            raise NotImplementedError(type(hlo))
-
-        for subexpr in subexprs:
-            if isinstance(subexpr, Array):
-                for i_in_axis, i_out_axis in zip(
-                        range(subexpr.ndim),
-                        range(expr.ndim-subexpr.ndim, expr.ndim),
-                        strict=True):
-                    in_dim = subexpr.shape[i_in_axis]
-                    out_dim = expr.shape[i_out_axis]
-                    if are_shape_components_equal(in_dim, out_dim):
-                        self.record_equation(
-                            self.get_var_for_axis(subexpr, i_in_axis),
-                            self.get_var_for_axis(expr, i_out_axis)
-                        )
-                    else:
-                        # i_in_axis is broadcasted => do not propagate
-                        assert are_shape_components_equal(in_dim, 1)
-            else:
-                assert isinstance(subexpr, SCALAR_CLASSES)
+        return
 
     def map_stack(self, expr: Stack) -> None:
         """
@@ -594,10 +595,11 @@ class AxisTagAttacher(CopyMapper):
     A mapper that tags the axes in a DAG as prescribed by *axis_to_tags*.
     """
     def __init__(self,
-                 axis_to_tags: Mapping[tuple[Array, int], Collection[Tag]],
+                 axis_to_tags: Mapping[tuple[Array, int | str], Collection[Tag]],
                  tag_corresponding_redn_descr: bool):
         super().__init__()
-        self.axis_to_tags: Mapping[tuple[Array, int], Collection[Tag]] = axis_to_tags
+        self.axis_to_tags: Mapping[tuple[Array, int | str],
+                                   Collection[Tag]] = axis_to_tags
         self.tag_corresponding_redn_descr: bool = tag_corresponding_redn_descr
 
     def rec(self, expr: ArrayOrNames) -> Any:
@@ -634,18 +636,15 @@ def rec(self, expr: ArrayOrNames) -> Any:
 
                     if isinstance(expr, IndexLambda):
                         assert isinstance(expr_copy, IndexLambda)
-                        try:
-                            hlo = index_lambda_to_high_level_op(expr)
-                        except UnknownIndexLambdaExpr:
-                            pass
-                        else:
-                            if isinstance(hlo, ReduceOp):
-                                for iaxis, redn_var in hlo.axes.items():
-                                    expr_copy = expr_copy.with_tagged_reduction(
+                        if expr_copy.var_to_reduction_descr:
+                            # This is a reduction operation.
+                            # We need to find the axes that are reduced over
+                            # and update the tag/tag them appropriately.
+                            for redn_var in expr.var_to_reduction_descr.keys():
+                                expr_copy = expr_copy.with_tagged_reduction(
                                         redn_var,
-                                        self.axis_to_tags.get((hlo.x, iaxis), [])
+                                        self.axis_to_tags.get((expr, redn_var), [])
                                     )
-
                 # }}}
 
                 self._cache[key] = expr_copy
@@ -710,7 +709,7 @@ def unify_axes_tags(
     )
 
     known_tag_vars = frozenset(equations_collector.known_tag_to_var.values())
-    axis_to_solved_tags: dict[tuple[Array, int], set[Tag]] = {}
+    axis_to_solved_tags: dict[tuple[Array, int | str], set[Tag]] = {}
 
     propagation_graph = undirected_graph_from_edges(
         equations_collector.equations
@@ -721,10 +720,10 @@ def unify_axes_tags(
         if isinstance(tag, AxisIgnoredForPropagationTag)
     })
 
-    ignored_vars.update({
-        ax_var for (ary, ax), ax_var in equations_collector.axis_to_var.items()
-        if ary.axes[ax].tags_of_type(AxisIgnoredForPropagationTag)
-    })
+    for (ary, ax), ax_var in equations_collector.axis_to_var.items():
+        if isinstance(ax, int):
+            if ary.axes[ax].tags_of_type(AxisIgnoredForPropagationTag):
+                ignored_vars.update({ax_var})
 
     for tag, var in equations_collector.known_tag_to_var.items():
         reachable_nodes = get_reachable_nodes(propagation_graph, var,

test/test_codegen.py

@@ -916,11 +916,11 @@ def test_einsum_with_parameterized_shapes(ctx_factory):
     m_in = np.random.randint(2, 20)
     n_in = np.random.randint(2, 20)
 
-    def _get_a_shape(_m, _n):
-        return (2*_m+1, 3*_n+7)
+    def _get_a_shape(m_, n_):
+        return (2*m_+1, 3*n_+7)
 
-    def _get_x_shape(_m, _n):
-        return (3*_n+7, )
+    def _get_x_shape(_m, n_):
+        return (3*n_+7, )
 
     A_in = np.random.rand(*_get_a_shape(m_in, n_in))  # noqa: N806
     x_in = np.random.rand(*_get_x_shape(m_in, n_in))
@@ -1570,16 +1570,16 @@ def test_regression_reduction_in_conditional(ctx_factory):
     ctx = ctx_factory()
     cq = cl.CommandQueue(ctx)
 
-    def kernel(usr_np, _pt_data_9):
-        pt_tmp_53 = _pt_data_9 @ _pt_data_9
+    def kernel(usr_np, pt_data_9):
+        pt_tmp_53 = pt_data_9 @ pt_data_9
         pt_tmp_42 = usr_np.maximum(pt_tmp_53, pt_tmp_53)
         pt_tmp_27 = usr_np.sum(pt_tmp_42)
         pt_tmp_0 = usr_np.maximum(pt_tmp_27, pt_tmp_53)
         return pt_tmp_0
 
     def get_np_input_args():
         return {
-            "_pt_data_9": np.ones((2, 2)),
+            "pt_data_9": np.ones((2, 2)),
         }
 
     np_inputs = get_np_input_args()