Optimize reshape idx lmbda lowering

.basedpyright/baseline.json

@@ -7623,46 +7623,6 @@
                     "lineCount": 1
                 }
             },
-            {
-                "code": "reportReturnType",
-                "range": {
-                    "startColumn": 11,
-                    "endColumn": 87,
-                    "lineCount": 4
-                }
-            },
-            {
-                "code": "reportUnknownVariableType",
-                "range": {
-                    "startColumn": 11,
-                    "endColumn": 87,
-                    "lineCount": 4
-                }
-            },
-            {
-                "code": "reportOperatorIssue",
-                "range": {
-                    "startColumn": 8,
-                    "endColumn": 46,
-                    "lineCount": 1
-                }
-            },
-            {
-                "code": "reportOperatorIssue",
-                "range": {
-                    "startColumn": 8,
-                    "endColumn": 60,
-                    "lineCount": 1
-                }
-            },
-            {
-                "code": "reportUnknownArgumentType",
-                "range": {
-                    "startColumn": 8,
-                    "endColumn": 86,
-                    "lineCount": 2
-                }
-            },
             {
                 "code": "reportUnknownMemberType",
                 "range": {

.gitignore

@@ -18,6 +18,7 @@ distribute*tar.gz
 doc/_build
 
 .mypy_cache
+.venv
 
 *.dot
 *.svg

pytato/scalar_expr.py

@@ -73,6 +73,7 @@
 )
 from pymbolic.mapper.distributor import DistributeMapper as DistributeMapperBase
 from pymbolic.mapper.evaluator import EvaluationMapper as EvaluationMapperBase
+from pymbolic.mapper.flattener import FlattenMapper as FlattenMapperBase
 from pymbolic.mapper.stringifier import StringifyMapper as StringifyMapperBase
 from pymbolic.mapper.substitutor import SubstitutionMapper as SubstitutionMapperBase
 from pymbolic.typing import Integer
@@ -212,18 +213,20 @@ def map_reduce(self, expr: Reduce) -> Never:
 
 
 class DistributeMapper(DistributeMapperBase):
-
     def map_reduce(self, expr: Reduce) -> None:
         # TODO: not trivial to distribute symbolic reduction nodes
         raise NotImplementedError()
 
 
 class TermCollector(TermCollectorBase):
-
     def map_reduce(self, expr: Reduce) -> None:
         raise NotImplementedError()
 
 
+class FlattenMapper(FlattenMapperBase, IdentityMapper[[]]):
+    pass
+
+
 class StringifyMapper(StringifyMapperBase[P]):
     def map_reduce(self,
                    expr: Any, enclosing_prec: int, *args: P.args, **kwargs: P.kwargs

pytato/transform/lower_to_index_lambda.py

@@ -36,7 +36,7 @@
 
 import pymbolic.primitives as prim
 from pymbolic import ArithmeticExpression
-from pytools import UniqueNameGenerator
+from pytools import UniqueNameGenerator, product
 
 from pytato.array import (
     AbstractResultWithNamedArrays,
@@ -61,7 +61,7 @@
     _get_einsum_access_descr_to_axis_len,
 )
 from pytato.diagnostic import CannotBeLoweredToIndexLambda
-from pytato.scalar_expr import INT_CLASSES, ScalarExpression
+from pytato.scalar_expr import INT_CLASSES, FlattenMapper, ScalarExpression
 from pytato.tags import AssumeNonNegative
 from pytato.transform import (
     Mapper,
@@ -127,13 +127,31 @@ def _generate_index_expressions(
         new_strides = new_strides[::-1]
         old_size_tills = old_size_tills[::-1]
 
-    flattened_index_expn = sum(
+    flattened_index_expn = FlattenMapper()(sum(
         index_var*new_stride
-        for index_var, new_stride in zip(index_vars, new_strides, strict=True))
+        for index_var, new_stride in zip(index_vars, new_strides, strict=True)))
+
+    old_size = product(old_shape)
+
+    def _mod(
+        num: ArithmeticExpression, denom: ArithmeticExpression
+    ) -> ArithmeticExpression:
+        from pymbolic.typing import Integer
+        if isinstance(old_size, Integer) and denom == old_size and denom != 0:
+            return num
+        # Pyright has a point: complex numbers don't support '%'.
+        return num % denom  # pyright: ignore[reportOperatorIssue,reportUnknownVariableType]
+
+    def _floordiv(
+        num: ArithmeticExpression, denom: ArithmeticExpression
+    ) -> ArithmeticExpression:
+        if denom == 1:
+            return num
+        # pyright has a point: complex numbers don't support '//'.
+        return num // denom  # pyright: ignore[reportOperatorIssue,reportUnknownVariableType]
 
     return tuple(
-        # Mypy has a point: complex numbers don't support '//'.
-        (flattened_index_expn % old_size_till) // old_stride  # type: ignore[operator]
+        _floordiv(_mod(flattened_index_expn, old_size_till), old_stride)  # pyright: ignore[reportArgumentType]
         for old_size_till, old_stride in zip(old_size_tills, old_strides, strict=True))
 
 

test/test_pytato.py

@@ -1561,6 +1561,23 @@ def test_lower_to_index_lambda():
     assert idx_lambda.expr.index_tuple[4] == Variable("_1")
 
 
+def test_lower_to_index_lambda_flatten_reshape():
+    # Before commit<=577cb39, pytato would generate redundant floor div
+    # and modulo computations. Specifically, in this case we would get:
+    # out[_0] = x[((0 + _0*1) % 40) // 4, ((0 + _0*1) % 4) // 1].
+
+    from pymbolic import parse
+
+    from pytato.array import IndexLambda
+    x = pt.make_placeholder(name="x", dtype=float, shape=(10, 4))
+    idx_lambda = pt.to_index_lambda(x.reshape(-1))
+    assert isinstance(idx_lambda, IndexLambda)
+    assert idx_lambda.expr.index_tuple == (
+        parse("_0 // 4"),
+        parse("_0 % 4"),
+    )
+
+
 def test_reserved_binding_name_patterns():
     from pytato.transform.metadata import BINDING_NAME_RESERVED_PATTERN