Merge improve_collapse_tuple

Author: tehrengruber

src/gt4py/next/iterator/ir_utils/common_pattern_matcher.py

@@ -27,11 +27,11 @@ def is_applied_lift(arg: itir.Node) -> TypeGuard[itir.FunCall]:
     )
 
 
-def is_let(node: itir.Node) -> bool:
+def is_let(node: itir.Node) -> TypeGuard[itir.FunCall]:
     """Match expression of the form `(λ(...) → ...)(...)`."""
     return isinstance(node, itir.FunCall) and isinstance(node.fun, itir.Lambda)
 
 
-def is_if_call(node: itir.Expr):
+def is_if_call(node: itir.Expr) -> TypeGuard[itir.FunCall]:
     """Match expression of the form `if_(cond, true_branch, false_branch)`."""
     return isinstance(node, itir.FunCall) and node.fun == im.ref("if_")

src/gt4py/next/iterator/ir_utils/ir_makers.py

@@ -250,12 +250,10 @@ class let:
     """
 
     @typing.overload
-    def __init__(self, var: str | itir.Sym, init_form: itir.Expr):
-        ...
+    def __init__(self, var: str | itir.Sym, init_form: itir.Expr): ...
 
     @typing.overload
-    def __init__(self, *args: Iterable[tuple[str | itir.Sym, itir.Expr]]):
-        ...
+    def __init__(self, *args: Iterable[tuple[str | itir.Sym, itir.Expr]]): ...
 
     def __init__(self, *args):
         if all(isinstance(arg, tuple) and len(arg) == 2 for arg in args):
@@ -268,7 +266,7 @@ def __init__(self, *args):
             self.init_forms = [args[1]]
         else:
             raise TypeError(
-                "Invalid arguments. Expected a variable name and an init form or a list thereof."
+                "Invalid arguments: expected a variable name and an init form or a list thereof."
             )
 
     def __call__(self, form):

src/gt4py/next/iterator/ir_utils/misc.py

@@ -57,13 +57,16 @@ def visit_SymRef(self, node: itir.SymRef, *, sym_map: dict[str, str]):
         return im.ref(sym_map[node.id]) if node.id in sym_map else node
 
 
-def is_provable_equal(a: itir.Expr, b: itir.Expr):
+def is_equal(a: itir.Expr, b: itir.Expr):
     """
-    Return true if, bot not only if, two expression (with equal scope) have the same value.
+    Return true if, but not only if, two expression (with equal scope) have the same value.
+
+    Be aware that this function might return false even though the two expression have the same
+    value.
 
     >>> testee1 = im.lambda_("a")(im.plus("a", "b"))
     >>> testee2 = im.lambda_("c")(im.plus("c", "b"))
-    >>> assert is_provable_equal(testee1, testee2)
+    >>> assert is_equal(testee1, testee2)
     """
     return a == b or (
         CannonicalizeBoundSymbolNames.apply(a) == CannonicalizeBoundSymbolNames.apply(b)

src/gt4py/next/iterator/transforms/collapse_tuple.py

@@ -137,8 +137,6 @@ class InlineLambdas(PreserveLocationVisitor, NodeTranslator):
 
     force_inline_trivial_lift_args: bool
 
-    force_inline_lambda_args: bool
-
     @classmethod
     def apply(
         cls,

src/gt4py/next/iterator/transforms/inline_lambdas.py

@@ -165,15 +165,17 @@ def main_transforms(
         inlined = ConstantFolding.apply(inlined)
         # This pass is required to be in the loop such that when an `if_` call with tuple arguments
         # is constant-folded the surrounding tuple_get calls can be removed.
-        if stage == 1:
-            inlined = CollapseTuple.apply(
-                inlined,
-                # to limit number of times global type inference is executed, only in the last iterations.
-                # use_global_type_inference=inlined == ir,
-                ignore_tuple_size=True,  # possibly dangerous
-                use_global_type_inference=False,
-            )
-        inlined = PropagateDeref.apply(inlined)  # todo: document
+        inlined = CollapseTuple.apply(
+            inlined,
+            # to limit number of times global type inference is executed, only in the last iterations.
+            # use_global_type_inference=inlined == ir,
+            ignore_tuple_size=True,  # possibly dangerous
+            use_global_type_inference=False,
+        )
+        # This pass is required such that a deref outside of a
+        # `tuple_get(make_tuple(let(...), ...))` call is propagated into the let after the
+        # `tuple_get` is removed by the `CollapseTuple` pass.
+        inlined = PropagateDeref.apply(inlined)
 
         if inlined == ir:
             stage += 1

src/gt4py/next/iterator/transforms/pass_manager.py

@@ -26,7 +26,7 @@
 from gt4py.next.otf import languages, recipes, stages, step_types, workflow
 from gt4py.next.otf.binding import nanobind
 from gt4py.next.otf.compilation import cache, compiler
-from gt4py.next.otf.compilation.build_systems import cmake, compiledb
+from gt4py.next.otf.compilation.build_systems import compiledb
 from gt4py.next.program_processors import otf_compile_executor
 from gt4py.next.program_processors.codegens.gtfn import gtfn_module
 from gt4py.next.type_system.type_translation import from_value
@@ -130,7 +130,7 @@ def compilation_hash(otf_closure: stages.ProgramCall) -> int:
 )
 
 GTFN_DEFAULT_COMPILE_STEP: step_types.CompilationStep = compiler.Compiler(
-    cache_strategy=cache.Strategy.PERSISTENT, builder_factory=cmake.CMakeFactory()
+    cache_strategy=cache.Strategy.SESSION, builder_factory=compiledb.CompiledbFactory()
 )
 
 

src/gt4py/next/program_processors/runners/gtfn.py

@@ -148,7 +148,7 @@ class ProgramFormatterId(_PythonObjectIdMixin, str, enum.Enum):
 COMMON_SKIP_TEST_LIST = [
     (REQUIRES_ATLAS, XFAIL, BINDINGS_UNSUPPORTED_MESSAGE),
     (USES_APPLIED_SHIFTS, XFAIL, UNSUPPORTED_MESSAGE),
-    # (USES_IF_STMTS, XFAIL, UNSUPPORTED_MESSAGE),
+    (USES_IF_STMTS, XFAIL, UNSUPPORTED_MESSAGE),
     (USES_NEGATIVE_MODULO, XFAIL, UNSUPPORTED_MESSAGE),
     (USES_REDUCTION_WITH_ONLY_SPARSE_FIELDS, XFAIL, REDUCTION_WITH_ONLY_SPARSE_FIELDS_MESSAGE),
     (USES_SPARSE_FIELDS_AS_OUTPUT, XFAIL, UNSUPPORTED_MESSAGE),

tests/next_tests/definitions.py

@@ -20,7 +20,11 @@ def test_simple_make_tuple_tuple_get():
     tuple_of_size_2 = im.make_tuple("first", "second")
     testee = im.make_tuple(im.tuple_get(0, tuple_of_size_2), im.tuple_get(1, tuple_of_size_2))
 
-    actual = CollapseTuple.apply(testee, flags=CollapseTuple.Flag.COLLAPSE_MAKE_TUPLE_TUPLE_GET)
+    actual = CollapseTuple.apply(
+        testee,
+        remove_letified_make_tuple_elements=False,
+        flags=CollapseTuple.Flag.COLLAPSE_MAKE_TUPLE_TUPLE_GET,
+    )
 
     expected = tuple_of_size_2
     assert actual == expected
@@ -32,7 +36,11 @@ def test_nested_make_tuple_tuple_get():
         im.tuple_get(0, tup_of_size2_from_lambda), im.tuple_get(1, tup_of_size2_from_lambda)
     )
 
-    actual = CollapseTuple.apply(testee, flags=CollapseTuple.Flag.COLLAPSE_MAKE_TUPLE_TUPLE_GET)
+    actual = CollapseTuple.apply(
+        testee,
+        remove_letified_make_tuple_elements=False,
+        flags=CollapseTuple.Flag.COLLAPSE_MAKE_TUPLE_TUPLE_GET,
+    )
 
     assert actual == tup_of_size2_from_lambda
 
@@ -42,21 +50,33 @@ def test_different_tuples_make_tuple_tuple_get():
     t1 = im.make_tuple("foo1", "bar1")
     testee = im.make_tuple(im.tuple_get(0, t0), im.tuple_get(1, t1))
 
-    actual = CollapseTuple.apply(testee, flags=CollapseTuple.Flag.COLLAPSE_MAKE_TUPLE_TUPLE_GET)
+    actual = CollapseTuple.apply(
+        testee,
+        remove_letified_make_tuple_elements=False,
+        flags=CollapseTuple.Flag.COLLAPSE_MAKE_TUPLE_TUPLE_GET,
+    )
 
     assert actual == testee  # did nothing
 
 
 def test_incompatible_order_make_tuple_tuple_get():
     tuple_of_size_2 = im.make_tuple("first", "second")
     testee = im.make_tuple(im.tuple_get(1, tuple_of_size_2), im.tuple_get(0, tuple_of_size_2))
-    actual = CollapseTuple.apply(testee, flags=CollapseTuple.Flag.COLLAPSE_MAKE_TUPLE_TUPLE_GET)
+    actual = CollapseTuple.apply(
+        testee,
+        remove_letified_make_tuple_elements=False,
+        flags=CollapseTuple.Flag.COLLAPSE_MAKE_TUPLE_TUPLE_GET,
+    )
     assert actual == testee  # did nothing
 
 
 def test_incompatible_size_make_tuple_tuple_get():
     testee = im.make_tuple(im.tuple_get(0, im.make_tuple("first", "second")))
-    actual = CollapseTuple.apply(testee, flags=CollapseTuple.Flag.COLLAPSE_MAKE_TUPLE_TUPLE_GET)
+    actual = CollapseTuple.apply(
+        testee,
+        remove_letified_make_tuple_elements=False,
+        flags=CollapseTuple.Flag.COLLAPSE_MAKE_TUPLE_TUPLE_GET,
+    )
     assert actual == testee  # did nothing
 
 
@@ -71,14 +91,22 @@ def test_merged_with_smaller_outer_size_make_tuple_tuple_get():
 def test_simple_tuple_get_make_tuple():
     expected = im.ref("bar")
     testee = im.tuple_get(1, im.make_tuple("foo", expected))
-    actual = CollapseTuple.apply(testee, flags=CollapseTuple.Flag.COLLAPSE_TUPLE_GET_MAKE_TUPLE)
+    actual = CollapseTuple.apply(
+        testee,
+        remove_letified_make_tuple_elements=False,
+        flags=CollapseTuple.Flag.COLLAPSE_TUPLE_GET_MAKE_TUPLE,
+    )
     assert expected == actual
 
 
 def test_propagate_tuple_get():
     expected = im.let(("el1", 1), ("el2", 2))(im.tuple_get(0, im.make_tuple("el1", "el2")))
     testee = im.tuple_get(0, im.let(("el1", 1), ("el2", 2))(im.make_tuple("el1", "el2")))
-    actual = CollapseTuple.apply(testee, flags=CollapseTuple.Flag.PROPAGATE_TUPLE_GET)
+    actual = CollapseTuple.apply(
+        testee,
+        remove_letified_make_tuple_elements=False,
+        flags=CollapseTuple.Flag.PROPAGATE_TUPLE_GET,
+    )
     assert expected == actual
 
 
@@ -89,23 +117,35 @@ def test_letify_make_tuple_elements():
         im.make_tuple("_tuple_el_1", "_tuple_el_2")
     )
 
-    actual = CollapseTuple.apply(testee, flags=CollapseTuple.Flag.LETIFY_MAKE_TUPLE_ELEMENTS)
+    actual = CollapseTuple.apply(
+        testee,
+        remove_letified_make_tuple_elements=False,
+        flags=CollapseTuple.Flag.LETIFY_MAKE_TUPLE_ELEMENTS,
+    )
     assert actual == expected
 
 
 def test_letify_make_tuple_with_trivial_elements():
     testee = im.let(("a", 1), ("b", 2))(im.make_tuple("a", "b"))
     expected = testee  # did nothing
 
-    actual = CollapseTuple.apply(testee, flags=CollapseTuple.Flag.LETIFY_MAKE_TUPLE_ELEMENTS)
+    actual = CollapseTuple.apply(
+        testee,
+        remove_letified_make_tuple_elements=False,
+        flags=CollapseTuple.Flag.LETIFY_MAKE_TUPLE_ELEMENTS,
+    )
     assert actual == expected
 
 
 def test_inline_trivial_make_tuple():
     testee = im.let("tup", im.make_tuple("a", "b"))("tup")
     expected = im.make_tuple("a", "b")
 
-    actual = CollapseTuple.apply(testee, flags=CollapseTuple.Flag.INLINE_TRIVIAL_MAKE_TUPLE)
+    actual = CollapseTuple.apply(
+        testee,
+        remove_letified_make_tuple_elements=False,
+        flags=CollapseTuple.Flag.INLINE_TRIVIAL_MAKE_TUPLE,
+    )
     assert actual == expected
 
 
@@ -114,7 +154,11 @@ def test_propagate_to_if_on_tuples():
     expected = im.if_(
         "cond", im.tuple_get(0, im.make_tuple(1, 2)), im.tuple_get(0, im.make_tuple(3, 4))
     )
-    actual = CollapseTuple.apply(testee, flags=CollapseTuple.Flag.PROPAGATE_TO_IF_ON_TUPLES)
+    actual = CollapseTuple.apply(
+        testee,
+        remove_letified_make_tuple_elements=False,
+        flags=CollapseTuple.Flag.PROPAGATE_TO_IF_ON_TUPLES,
+    )
     assert actual == expected
 
 
@@ -127,28 +171,28 @@ def test_propagate_to_if_on_tuples_with_let():
     )
     actual = CollapseTuple.apply(
         testee,
+        remove_letified_make_tuple_elements=True,
         flags=CollapseTuple.Flag.PROPAGATE_TO_IF_ON_TUPLES
-        | CollapseTuple.Flag.LETIFY_MAKE_TUPLE_ELEMENTS
-        | CollapseTuple.Flag.REMOVE_LETIFIED_MAKE_TUPLE_ELEMENTS,
+        | CollapseTuple.Flag.LETIFY_MAKE_TUPLE_ELEMENTS,
     )
     assert actual == expected
 
 
 def test_propagate_nested_lift():
     testee = im.let("a", im.let("b", 1)("a_val"))("a")
     expected = im.let("b", 1)(im.let("a", "a_val")("a"))
-    actual = CollapseTuple.apply(testee, flags=CollapseTuple.Flag.PROPAGATE_NESTED_LET)
+    actual = CollapseTuple.apply(
+        testee,
+        remove_letified_make_tuple_elements=False,
+        flags=CollapseTuple.Flag.PROPAGATE_NESTED_LET,
+    )
     assert actual == expected
 
 
-def test_collapse_complicated_():
-    # TODO: fuse with test_propagate_to_if_on_tuples_with_let
+def test_if_on_tuples_with_let():
     testee = im.let("val", im.if_("cond", im.make_tuple(1, 2), im.make_tuple(3, 4)))(
         im.tuple_get(0, "val")
     )
     expected = im.if_("cond", 1, 3)
-    actual = CollapseTuple.apply(
-        testee,
-        # flags=CollapseTuple.Flag.PROPAGATE_TO_IF_ON_TUPLES
-    )
+    actual = CollapseTuple.apply(testee, remove_letified_make_tuple_elements=False)
     assert actual == expected