Repository: GridTools/gt4py
Pull request #1810 (2025-01-21 11:12:43)

feat[next]: Extend and refactor constant folding
The following new transformations are introduced:
- _CANONICALIZE_OP_FUNCALL_SYMREF_LITERAL_:
`literal, symref` -> `symref, literal`, prerequisite for FOLD_FUNCALL_LITERAL, FOLD_NEUTRAL_OP
`literal, funcall` -> `funcall, literal`, prerequisite for FOLD_FUNCALL_LITERAL, FOLD_NEUTRAL_OP
`funcall, op` -> `op, funcall` for s[0] + (s[0] + 1), prerequisite for FOLD_MIN_MAX_PLUS
- _CANONICALIZE_MINUS_: `a - b` -> `a + (-b)`, prerequisite for FOLD_MIN_MAX_PLUS
- _CANONICALIZE_MIN_MAX_: `maximum(a, maximum(...))` -> `maximum(maximum(...), a)`, prerequisite for
- _FOLD_FUNCALL_LITERAL_:  `(a + 1) + 1` -> `a + (1 + 1)`
- _FOLD_MIN_MAX_:
`maximum(maximum(a, 1), a)` -> `maximum(a, 1)`
`maximum(maximum(a, 1), 1)` -> `maximum(a, 1)`
- FOLD_MIN_MAX_PLUS: `maximum(plus(a, 1), a)` -> `plus(a, 1)`
`maximum(plus(a, 1), plus(a, -1))` -> `plus(a, maximum(1, -1))`
- FOLD_NEUTRAL_OP: `a + 0` -> `a`, `a * 1` -> `a`

In the end, unary minuses are transformed back to `minus`-calls where possible by `UndoCanonicalizeMinus`:
`a + (-b)` -> `a - b` , `-a + b` -> `b - a`, `-a + (-b)` -> `-a - b`

In addition the pass now transforms until a fixed point is reached (like in the `CollapseTuple` pass) instead of just transforming once. The `FixedPointTransformation` class from  [PR 1826](https://github.com/GridTools/gt4py/pull/1826) is used here.

Previously, large nested maximum expressions like `maximum(maximum(maximum(maximum(sym, 1),...), 1), maximum(maximum(sym, 1), 1))` caused timeouts in PMAP-GO as the runtime of the domain inference increased significantly due to the large domain expressions that could not be simplified.

This replaces [this PR](https://github.com/tehrengruber/gt4py/pull/4).


Patch:
diff --git a/src/gt4py/next/iterator/transforms/constant_folding.py b/src/gt4py/next/iterator/transforms/constant_folding.py
index 0dc324f9..fdbfec99 100644
--- a/src/gt4py/next/iterator/transforms/constant_folding.py
+++ b/src/gt4py/next/iterator/transforms/constant_folding.py
@@ -6,57 +6,226 @@
 # Please, refer to the LICENSE file in the root directory.
 # SPDX-License-Identifier: BSD-3-Clause
 
-from gt4py.eve import NodeTranslator, PreserveLocationVisitor
+from __future__ import annotations
+
+import dataclasses
+import enum
+import functools
+import operator
+from typing import Optional
+
+from gt4py import eve
 from gt4py.next.iterator import builtins, embedded, ir
-from gt4py.next.iterator.ir_utils import ir_makers as im
+from gt4py.next.iterator.ir_utils import common_pattern_matcher as cpm, ir_makers as im
+from gt4py.next.iterator.transforms import fixed_point_transformation
+
 
+def _value_from_literal(literal: ir.Literal):
+    return getattr(embedded, str(literal.type))(literal.value)
 
-class ConstantFolding(PreserveLocationVisitor, NodeTranslator):
+
+class UndoCanonicalizeMinus(eve.NodeTranslator):
     PRESERVED_ANNEX_ATTRS = (
         "type",
         "domain",
     )
 
+    def visit_FunCall(self, node: ir.FunCall, **kwargs) -> ir.Node:
+        node = super().generic_visit(node, **kwargs)
+        # `a + (-b)` -> `a - b` , `-a + b` -> `b - a`, `-a + (-b)` -> `-a - b`
+        if cpm.is_call_to(node, "plus"):
+            a, b = node.args
+            if cpm.is_call_to(b, "neg"):
+                return im.minus(a, b.args[0])
+            if isinstance(b, ir.Literal) and _value_from_literal(b) < 0:
+                return im.minus(a, -_value_from_literal(b))
+            if cpm.is_call_to(a, "neg"):
+                return im.minus(b, a.args[0])
+            if isinstance(a, ir.Literal) and _value_from_literal(a) < 0:
+                return im.minus(b, -_value_from_literal(a))
+        return node
+
+
+_COMMUTATIVE_OPS = ("plus", "multiplies", "minimum", "maximum")
+
+
+@dataclasses.dataclass(frozen=True, kw_only=True)
+class ConstantFolding(
+    fixed_point_transformation.FixedPointTransformation, eve.PreserveLocationVisitor
+):
+    PRESERVED_ANNEX_ATTRS = (
+        "type",
+        "domain",
+    )
+
+    class Transformation(enum.Flag):
+        # `1 + a` -> `a + 1`, prerequisite for FOLD_FUNCALL_LITERAL, FOLD_NEUTRAL_OP
+        # `1 + f(...)` -> `f(...) + 1`, prerequisite for FOLD_FUNCALL_LITERAL, FOLD_NEUTRAL_OP
+        # `f(...) + (expr1 + expr2)` -> `(expr1 + expr2) + f(...)`, for `s[0] + (s[0] + 1)`, prerequisite for FOLD_MIN_MAX_PLUS
+        CANONICALIZE_OP_FUNCALL_SYMREF_LITERAL = enum.auto()
+
+        # `a - b` -> `a + (-b)`, prerequisite for FOLD_MIN_MAX_PLUS
+        CANONICALIZE_MINUS = enum.auto()
+
+        # `maximum(a, maximum(...))` -> `maximum(maximum(...), a)`, prerequisite for FOLD_MIN_MAX
+        CANONICALIZE_MIN_MAX = enum.auto()
+
+        # `(a + 1) + 1` -> `a + (1 + 1)`
+        FOLD_FUNCALL_LITERAL = enum.auto()
+
+        # `maximum(maximum(a, 1), a)` -> `maximum(a, 1)`
+        # `maximum(maximum(a, 1), 1)` -> `maximum(a, 1)`
+        FOLD_MIN_MAX = enum.auto()
+
+        # `maximum(a + 1), a)` -> `a + 1`
+        # `maximum(a + 1, a + (-1))` -> `a + maximum(1, -1)`
+        FOLD_MIN_MAX_PLUS = enum.auto()
+
+        # `a + 0` -> `a`, `a * 1` -> `a`
+        FOLD_NEUTRAL_OP = enum.auto()
+
+        # `1 + 1` -> `2`
+        FOLD_ARITHMETIC_BUILTINS = enum.auto()
+
+        # `minimum(a, a)` -> `a`
+        FOLD_MIN_MAX_LITERALS = enum.auto()
+
+        # `if_(True, true_branch, false_branch)` -> `true_branch`
+        FOLD_IF = enum.auto()
+
+        @classmethod
+        def all(self) -> ConstantFolding.Transformation:
+            return functools.reduce(operator.or_, self.__members__.values())
+
+    enabled_transformations: Transformation = Transformation.all()  # noqa: RUF009 [function-call-in-dataclass-default-argument]
+
     @classmethod
     def apply(cls, node: ir.Node) -> ir.Node:
-        return cls().visit(node)
+        node = cls().visit(node)
+        return UndoCanonicalizeMinus().visit(node)
+
+    def transform_canonicalize_op_funcall_symref_literal(
+        self, node: ir.FunCall, **kwargs
+    ) -> Optional[ir.Node]:
+        # `op(literal, symref|funcall)` -> `op(symref|funcall, literal)`
+        # `op1(funcall, op2(...))` -> `op1(op2(...), funcall)` for `s[0] + (s[0] + 1)`
+        if cpm.is_call_to(node, _COMMUTATIVE_OPS):
+            a, b = node.args
+            if (isinstance(a, ir.Literal) and not isinstance(b, ir.Literal)) or (
+                not cpm.is_call_to(a, _COMMUTATIVE_OPS) and cpm.is_call_to(b, _COMMUTATIVE_OPS)
+            ):
+                return im.call(node.fun)(b, a)
+        return None
 
-    def visit_FunCall(self, node: ir.FunCall):
-        # visit depth-first such that nested constant expressions (e.g. `(1+2)+3`) are properly folded
-        new_node = self.generic_visit(node)
+    def transform_canonicalize_minus(self, node: ir.FunCall, **kwargs) -> Optional[ir.Node]:
+        # `a - b` -> `a + (-b)`
+        if cpm.is_call_to(node, "minus"):
+            return im.plus(node.args[0], self.fp_transform(im.call("neg")(node.args[1])))
+        return None
 
+    def transform_canonicalize_min_max(self, node: ir.FunCall, **kwargs) -> Optional[ir.Node]:
+        # `maximum(a, maximum(...))` -> `maximum(maximum(...), a)`
+        if cpm.is_call_to(node, ("maximum", "minimum")):
+            op = node.fun.id  # type: ignore[attr-defined] # assured by if above
+            if cpm.is_call_to(node.args[1], op) and not cpm.is_call_to(node.args[0], op):
+                return im.call(op)(node.args[1], node.args[0])
+        return None
+
+    def transform_fold_funcall_literal(self, node: ir.FunCall, **kwargs) -> Optional[ir.Node]:
+        # `(a + 1) + 1` -> `a + (1 + 1)`
+        if cpm.is_call_to(node, "plus"):
+            if cpm.is_call_to(node.args[0], "plus") and isinstance(node.args[1], ir.Literal):
+                (expr, lit1), lit2 = node.args[0].args, node.args[1]
+                if isinstance(expr, (ir.SymRef, ir.FunCall)) and isinstance(lit1, ir.Literal):
+                    return im.plus(
+                        expr,
+                        self.fp_transform(im.plus(lit1, lit2)),
+                    )
+        return None
+
+    def transform_fold_min_max(self, node: ir.FunCall, **kwargs) -> Optional[ir.Node]:
+        # `maximum(maximum(a, 1), a)` -> `maximum(a, 1)`
+        # `maximum(maximum(a, 1), 1)` -> `maximum(a, 1)`
+        if cpm.is_call_to(node, ("minimum", "maximum")):
+            op = node.fun.id  # type: ignore[attr-defined] # assured by if above
+            if cpm.is_call_to(node.args[0], op):
+                fun_call, arg1 = node.args
+                if arg1 in fun_call.args:  # type: ignore[attr-defined] # assured by if above
+                    return fun_call
+        return None
+
+    def transform_fold_min_max_plus(self, node: ir.FunCall, **kwargs) -> Optional[ir.Node]:
         if (
-            isinstance(new_node.fun, ir.SymRef)
-            and new_node.fun.id in ["minimum", "maximum"]
-            and new_node.args[0] == new_node.args[1]
-        ):  # `minimum(a, a)` -> `a`
-            return new_node.args[0]
+            isinstance(node, ir.FunCall)
+            and isinstance(node.fun, ir.SymRef)
+            and cpm.is_call_to(node, ("minimum", "maximum"))
+        ):
+            arg0, arg1 = node.args
+            # `maximum(a + 1, a)` -> `a + 1`
+            if cpm.is_call_to(arg0, "plus"):
+                if arg0.args[0] == arg1:
+                    return im.plus(
+                        arg0.args[0], self.fp_transform(im.call(node.fun.id)(arg0.args[1], 0))
+                    )
+            # `maximum(a + 1, a + (-1))` -> `a + maximum(1, -1)`
+            if cpm.is_call_to(arg0, "plus") and cpm.is_call_to(arg1, "plus"):
+                if arg0.args[0] == arg1.args[0]:
+                    return im.plus(
+                        arg0.args[0],
+                        self.fp_transform(im.call(node.fun.id)(arg0.args[1], arg1.args[1])),
+                    )
+
+        return None
 
+    def transform_fold_neutral_op(self, node: ir.FunCall, **kwargs) -> Optional[ir.Node]:
+        # `a + 0` -> `a`, `a * 1` -> `a`
         if (
-            isinstance(new_node.fun, ir.SymRef)
-            and new_node.fun.id == "if_"
-            and isinstance(new_node.args[0], ir.Literal)
-        ):  # `if_(True, true_branch, false_branch)` -> `true_branch`
-            if new_node.args[0].value == "True":
-                new_node = new_node.args[1]
-            else:
-                new_node = new_node.args[2]
+            cpm.is_call_to(node, "plus")
+            and isinstance(node.args[1], ir.Literal)
+            and node.args[1].value.isdigit()
+            and int(node.args[1].value) == 0
+        ) or (
+            cpm.is_call_to(node, "multiplies")
+            and isinstance(node.args[1], ir.Literal)
+            and node.args[1].value.isdigit()
+            and int(node.args[1].value) == 1
+        ):
+            return node.args[0]
+        return None
 
+    @classmethod
+    def transform_fold_arithmetic_builtins(self, node: ir.FunCall, **kwargs) -> Optional[ir.Node]:
+        # `1 + 1` -> `2`
         if (
-            isinstance(new_node, ir.FunCall)
-            and isinstance(new_node.fun, ir.SymRef)
-            and len(new_node.args) > 0
-            and all(isinstance(arg, ir.Literal) for arg in new_node.args)
-        ):  # `1 + 1` -> `2`
+            isinstance(node, ir.FunCall)
+            and isinstance(node.fun, ir.SymRef)
+            and len(node.args) > 0
+            and all(isinstance(arg, ir.Literal) for arg in node.args)
+        ):
             try:
-                if new_node.fun.id in builtins.ARITHMETIC_BUILTINS:
-                    fun = getattr(embedded, str(new_node.fun.id))
+                if node.fun.id in builtins.ARITHMETIC_BUILTINS:
+                    fun = getattr(embedded, str(node.fun.id))
                     arg_values = [
-                        getattr(embedded, str(arg.type))(arg.value)  # type: ignore[attr-defined] # arg type already established in if condition
-                        for arg in new_node.args
+                        _value_from_literal(arg)  # type: ignore[arg-type] # arg type already established in if condition
+                        for arg in node.args
                     ]
-                    new_node = im.literal_from_value(fun(*arg_values))
+                    return im.literal_from_value(fun(*arg_values))
             except ValueError:
                 pass  # happens for inf and neginf
+        return None
 
-        return new_node
+    def transform_fold_min_max_literals(self, node: ir.FunCall, **kwargs) -> Optional[ir.Node]:
+        # `minimum(a, a)` -> `a`
+        if cpm.is_call_to(node, ("minimum", "maximum")):
+            if node.args[0] == node.args[1]:
+                return node.args[0]
+        return None
+
+    def transform_fold_if(self, node: ir.FunCall, **kwargs) -> Optional[ir.Node]:
+        # `if_(True, true_branch, false_branch)` -> `true_branch`
+        if cpm.is_call_to(node, "if_") and isinstance(node.args[0], ir.Literal):
+            if node.args[0].value == "True":
+                return node.args[1]
+            else:
+                return node.args[2]
+        return None


Test patch:
diff --git a/tests/next_tests/unit_tests/iterator_tests/transforms_tests/test_constant_folding.py b/tests/next_tests/unit_tests/iterator_tests/transforms_tests/test_constant_folding.py
index cf325c2d..1da2b8ce 100644
--- a/tests/next_tests/unit_tests/iterator_tests/transforms_tests/test_constant_folding.py
+++ b/tests/next_tests/unit_tests/iterator_tests/transforms_tests/test_constant_folding.py
@@ -9,54 +9,155 @@
 from gt4py.next.iterator.ir_utils import ir_makers as im
 from gt4py.next.iterator.transforms.constant_folding import ConstantFolding
 
-
-def test_constant_folding_boolean():
-    testee = im.not_(im.literal_from_value(True))
-    expected = im.literal_from_value(False)
-
-    actual = ConstantFolding.apply(testee)
-    assert actual == expected
+import pytest
+from gt4py.next.iterator.ir_utils import ir_makers as im
+from gt4py.next.iterator.transforms.constant_folding import ConstantFolding
 
 
-def test_constant_folding_math_op():
-    expected = im.literal_from_value(13)
-    testee = im.plus(
-        im.literal_from_value(4),
-        im.plus(
-            im.literal_from_value(7), im.minus(im.literal_from_value(7), im.literal_from_value(5))
+def test_cases():
+    return (
+        # expr, simplified expr
+        (im.plus(1, 1), 2),
+        (im.not_(True), False),
+        (im.plus(4, im.plus(7, im.minus(7, 5))), 13),
+        (im.if_(True, im.plus(im.ref("a"), 2), im.minus(9, 5)), im.plus("a", 2)),
+        (im.minimum("a", "a"), "a"),
+        (im.maximum(1, 2), 2),
+        # canonicalization
+        (im.plus("a", 1), im.plus("a", 1)),
+        (im.plus(1, "a"), im.plus("a", 1)),
+        # nested plus
+        (im.plus(im.plus("a", 1), 1), im.plus("a", 2)),
+        (im.plus(1, im.plus("a", 1)), im.plus("a", 2)),
+        # nested maximum
+        (im.maximum(im.maximum("a", 1), 1), im.maximum("a", 1)),
+        (im.maximum(im.maximum(1, "a"), 1), im.maximum("a", 1)),
+        (im.maximum("a", im.maximum(1, "a")), im.maximum("a", 1)),
+        (im.maximum(im.maximum(1, "a"), im.maximum(1, "a")), im.maximum("a", 1)),
+        (im.maximum(im.maximum(1, "a"), im.maximum("a", 1)), im.maximum("a", 1)),
+        (im.maximum(im.minimum("a", 1), "a"), im.maximum(im.minimum("a", 1), "a")),
+        # maximum & plus
+        (im.maximum(im.plus("a", 1), im.plus("a", 0)), im.plus("a", 1)),
+        (
+            im.maximum(im.plus("a", 1), im.plus(im.plus("a", 1), 0)),
+            im.plus("a", 1),
+        ),
+        (im.maximum("a", im.plus("a", 1)), im.plus("a", 1)),
+        (im.maximum("a", im.plus("a", im.literal_from_value(-1))), im.ref("a")),
+        (
+            im.plus("a", im.maximum(0, im.literal_from_value(-1))),
+            im.ref("a"),
+        ),
+        # plus & minus
+        (im.minus(im.plus("a", 1), im.plus(1, 1)), im.minus("a", 1)),
+        (im.plus(im.minus("a", 1), 2), im.plus("a", 1)),
+        (im.plus(im.minus(1, "a"), 1), im.minus(2, "a")),
+        # nested plus
+        (im.plus(im.plus("a", 1), im.plus(1, 1)), im.plus("a", 3)),
+        (
+            im.plus(im.plus("a", im.literal_from_value(-1)), im.plus("a", 3)),
+            im.plus(im.minus("a", 1), im.plus("a", 3)),
+        ),
+        # maximum & minus
+        (im.maximum(im.minus("a", 1), "a"), im.ref("a")),
+        (im.maximum("a", im.minus("a", im.literal_from_value(-1))), im.plus("a", 1)),
+        (
+            im.maximum(im.plus("a", im.literal_from_value(-1)), 1),
+            im.maximum(im.minus("a", 1), 1),
+        ),
+        # minimum & plus & minus
+        (im.minimum(im.plus("a", 1), "a"), im.ref("a")),
+        (im.minimum("a", im.plus("a", im.literal_from_value(-1))), im.minus("a", 1)),
+        (im.minimum(im.minus("a", 1), "a"), im.minus("a", 1)),
+        (im.minimum("a", im.minus("a", im.literal_from_value(-1))), im.ref("a")),
+        # nested maximum
+        (im.maximum("a", im.maximum("b", "a")), im.maximum("b", "a")),
+        # maximum & plus on complicated expr (tuple_get)
+        (
+            im.maximum(
+                im.plus(im.tuple_get(1, "a"), 1),
+                im.maximum(im.tuple_get(1, "a"), im.plus(im.tuple_get(1, "a"), 1)),
+            ),
+            im.plus(im.tuple_get(1, "a"), 1),
+        ),
+        # nested maximum & plus
+        (
+            im.maximum(im.maximum(im.plus(1, "a"), 1), im.plus(1, "a")),
+            im.maximum(im.plus("a", 1), 1),
+        ),
+        # sanity check that no strange things happen
+        # complex tests
+        (
+            # 1 - max(max(1, max(1, sym), min(1, sym), sym), 1 + (min(-1, 2) + max(-1, 1 - sym)))
+            im.minus(
+                1,
+                im.maximum(
+                    im.maximum(
+                        im.maximum(1, im.maximum(1, "a")),
+                        im.maximum(im.maximum(1, "a"), "a"),
+                    ),
+                    im.plus(
+                        1,
+                        im.plus(
+                            im.minimum(im.literal_from_value(-1), 2),
+                            im.maximum(im.literal_from_value(-1), im.minus(1, "a")),
+                        ),
+                    ),
+                ),
+            ),
+            # 1 - maximum(maximum(sym, 1), maximum(1 - sym, -1))
+            im.minus(
+                1,
+                im.maximum(
+                    im.maximum("a", 1),
+                    im.maximum(im.minus(1, "a"), im.literal_from_value(-1)),
+                ),
+            ),
+        ),
+        (
+            # maximum(sym, 1 + sym) + (maximum(1, maximum(1, sym)) + (sym - 1 + (1 + (sym + 1) + 1))) - 2
+            im.minus(
+                im.plus(
+                    im.maximum("a", im.plus(1, "a")),
+                    im.plus(
+                        im.maximum(1, im.maximum(1, "a")),
+                        im.plus(im.minus("a", 1), im.plus(im.plus(1, im.plus("a", 1)), 1)),
+                    ),
+                ),
+                2,
+            ),
+            # sym + 1 + (maximum(sym, 1) + (sym - 1 + (sym + 3))) - 2
+            im.minus(
+                im.plus(
+                    im.plus("a", 1),
+                    im.plus(
+                        im.maximum("a", 1),
+                        im.plus(im.minus("a", 1), im.plus("a", 3)),
+                    ),
+                ),
+                2,
+            ),
+        ),
+        (
+            # minimum(1 - sym, 1 + sym) + (maximum(maximum(1 - sym, 1 + sym), 1 - sym) + maximum(1 - sym, 1 - sym))
+            im.plus(
+                im.minimum(im.minus(1, "a"), im.plus(1, "a")),
+                im.plus(
+                    im.maximum(im.maximum(im.minus(1, "a"), im.plus(1, "a")), im.minus(1, "a")),
+                    im.maximum(im.minus(1, "a"), im.minus(1, "a")),
+                ),
+            ),
+            # minimum(1 - sym, sym + 1) + (maximum(1 - sym, sym + 1) + (1 - sym))
+            im.plus(
+                im.minimum(im.minus(1, "a"), im.plus("a", 1)),
+                im.plus(im.maximum(im.minus(1, "a"), im.plus("a", 1)), im.minus(1, "a")),
+            ),
         ),
     )
-    actual = ConstantFolding.apply(testee)
-    assert actual == expected
-
-
-def test_constant_folding_if():
-    expected = im.plus("a", 2)
-    testee = im.if_(
-        im.literal_from_value(True),
-        im.plus(im.ref("a"), im.literal_from_value(2)),
-        im.minus(im.literal_from_value(9), im.literal_from_value(5)),
-    )
-    actual = ConstantFolding.apply(testee)
-    assert actual == expected
-
-
-def test_constant_folding_minimum():
-    testee = im.minimum("a", "a")
-    expected = im.ref("a")
-    actual = ConstantFolding.apply(testee)
-    assert actual == expected
-
-
-def test_constant_folding_literal():
-    testee = im.plus(im.literal_from_value(1), im.literal_from_value(2))
-    expected = im.literal_from_value(3)
-    actual = ConstantFolding.apply(testee)
-    assert actual == expected
 
 
-def test_constant_folding_literal_maximum():
-    testee = im.maximum(im.literal_from_value(1), im.literal_from_value(2))
-    expected = im.literal_from_value(2)
+@pytest.mark.parametrize("test_case", test_cases())
+def test_constant_folding(test_case):
+    testee, expected = test_case
     actual = ConstantFolding.apply(testee)
-    assert actual == expected
+    assert actual == im.ensure_expr(expected)