feature[next]: Non-tree-size-increasing collapse tuple on ifs #1762

tehrengruber · 2024-12-01T18:58:39Z

Removing the tuple expressions across if_ calls on ITIR has been a pain point in the past. While the PROPAGATE_TO_IF_ON_TUPLES option of the CollapseTuplePass works very reliably, the resulting increase in the tree size has been prohibitive. With the refactoring to GTIR this problem became much less pronounced, as we could restrict the propagation to field-level, i.e., outside of stencils, but the tree still grew exponentially in the number of references to boolean arguments used inside if_ conditions. This PR adds an additional option PROPAGATE_TO_IF_ON_TUPLES_CPS to the CollapseTuplePass, which is similar to the existing PROPAGATE_TO_IF_ON_TUPLES, but propagates in the opposite direction, i.e. into the tree. This allows removal of tuple expressions across if_ calls without increasing the size of the tree. This is particularly important for if statements in the frontend, where outwards propagation can have devastating effects on the tree size, without any gained optimization potential. For example

complex_lambda(if cond1
  if cond2
    {...}
  else:
    {...}
else
  {...})

is problematic, since PROPAGATE_TO_IF_ON_TUPLES would propagate, and hence duplicate, complex_lambda three times, while we only want to get rid of the tuple expressions inside of the if_s. Note that this transformation is not mutually exclusive to PROPAGATE_TO_IF_ON_TUPLES.

tehrengruber · 2024-12-05T10:11:43Z

src/gt4py/next/iterator/type_system/inference.py

+    def visit_OffsetLiteral(
+        self, node: itir.OffsetLiteral, **kwargs
+    ) -> it_ts.OffsetLiteralType | ts.DeferredType:
+        if self.reinfer:


Suggested change

if self.reinfer:

# `self.dimensions` not available in reinference mode. Skip since we don't care anyway.

if self.reinfer:

havogt

Need more documentation of the new trafo or additionally a more detailed explanation on how it works.

havogt · 2024-12-05T12:21:45Z

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

+    if cpm.is_call_to(node, "make_tuple"):
+        return all(_is_trivial_or_tuple_thereof_expr(arg) for arg in node.args)
+    if cpm.is_call_to(node, "tuple_get"):
+        return _is_trivial_or_tuple_thereof_expr(node.args[1])
+    if isinstance(node, (ir.SymRef, ir.Literal)):
+        return True


Suggested change

if cpm.is_call_to(node, "make_tuple"):

return all(_is_trivial_or_tuple_thereof_expr(arg) for arg in node.args)

if cpm.is_call_to(node, "tuple_get"):

return _is_trivial_or_tuple_thereof_expr(node.args[1])

if isinstance(node, (ir.SymRef, ir.Literal)):

return True

if isinstance(node, (ir.SymRef, ir.Literal)):

return True

if cpm.is_call_to(node, "make_tuple"):

return all(_is_trivial_or_tuple_thereof_expr(arg) for arg in node.args)

if cpm.is_call_to(node, "tuple_get"):

return _is_trivial_or_tuple_thereof_expr(node.args[1])

let's move the definition of trivial to the top

havogt · 2024-12-05T12:22:17Z

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

@@ -47,6 +48,32 @@ def _is_trivial_make_tuple_call(node: ir.Expr):
    return True


+def _is_trivial_or_tuple_thereof_expr(node: ir.Node) -> bool:
+    """
+    Return `true` if the expr is a trivial expression or tuple thereof.


Suggested change

Return `true` if the expr is a trivial expression or tuple thereof.

Return `true` if the expr is a trivial expression (`SymRef` or `Literal`) or tuple thereof.

havogt · 2024-12-05T12:25:45Z

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

@@ -185,6 +229,8 @@ def transform(self, node: ir.Node, **kwargs) -> Optional[ir.Node]:
                method = getattr(self, f"transform_{transformation.name.lower()}")
                result = method(node, **kwargs)
                if result is not None:
+                    assert result is not node


Why is this assert useful?

havogt · 2024-12-05T12:31:20Z

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

@@ -126,7 +137,10 @@ def apply_common_transforms(
    # only run the unconditional version here instead of in the loop above.
    if unconditionally_collapse_tuples:
        ir = CollapseTuple.apply(
-            ir, ignore_tuple_size=True, offset_provider_type=offset_provider_type
+            ir,


why don't we need to exclude one of [PROPAGATE_TO_IF_ON_TUPLES, PROPAGATE_TO_IF_ON_TUPLES_CPS]

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

havogt · 2024-12-05T13:03:59Z

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

+        self, node: ir.FunCall, **kwargs
+    ) -> Optional[ir.Node]:
+        if not cpm.is_call_to(node, "if_"):
+            for i, arg in enumerate(node.args):


Here we iterate over any functions args or do we know more?

havogt · 2024-12-05T13:04:16Z

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

@@ -312,6 +358,73 @@ def transform_propagate_to_if_on_tuples(self, node: ir.FunCall, **kwargs) -> Opt
                    return im.if_(cond, new_true_branch, new_false_branch)
        return None

+    def transform_propagate_to_if_on_tuples_cps(


Can you break this function into smaller pieces, I am completely lost...

havogt · 2024-12-05T13:07:11Z

src/gt4py/next/iterator/type_system/inference.py

+        Contrary to the regular inference, this method does not descend into already typed sub-nodes
+        and can be used as a lightweight way to restore type information during a pass.
+
+        Note that this function is stateful, which is usually desired, and more performant.


Where is the statefulness?

havogt · 2024-12-05T13:08:28Z

src/gt4py/next/iterator/type_system/inference.py

+        # we found a node that is typed, do not descend into children
+        if self.reinfer and isinstance(node, itir.Node) and node.type:
+            if isinstance(node.type, ts.FunctionType):
+                return _type_synthesizer_from_function_type(node.type)


What does this do?

havogt · 2024-12-05T13:14:52Z

tests/next_tests/unit_tests/iterator_tests/transforms_tests/test_collapse_tuple.py

@@ -239,3 +240,48 @@ def test_tuple_get_on_untyped_ref():

    actual = CollapseTuple.apply(testee, allow_undeclared_symbols=True, within_stencil=False)
    assert actual == testee
+
+
+def test_if_make_tuple_reorder_cps():


Can you add a test with ifs in more than one arg?

tehrengruber added 15 commits December 1, 2024 19:58

Non-tree-size-increasing collapse tuple on ifs

34d6040

Fix typos

48abc08

Fix typos

3790944

Disable PROPAGATE_TO_IF_ON_TUPLES by default in pass manager

a8a63bf

Improve doc

2c44ffc

Improve typo

42b5817

Improve typo

9da19a2

Improve typo

bcd9e48

Improve typo

0a212bd

Fix typo

70562fe

Improve doc

9cee650

Format

43f5741

Fix type synthesizer for partially typed arithmetic ops

7b37f1c

Fix test

a0341a6

Fix test

5a892f3

tehrengruber mentioned this pull request Dec 2, 2024

feat[next]: Improve fieldop fusion #1764

Draft

tehrengruber commented Dec 5, 2024

View reviewed changes

havogt requested changes Dec 5, 2024

View reviewed changes

tehrengruber added 3 commits December 6, 2024 10:19

Address review comments

914a9e5

Add test and fix nested transformation on nested ifs

fc46edf

Merge origin/main

4e12195

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

feature[next]: Non-tree-size-increasing collapse tuple on ifs #1762

feature[next]: Non-tree-size-increasing collapse tuple on ifs #1762

tehrengruber commented Dec 1, 2024 •

edited

Loading

tehrengruber Dec 5, 2024 •

edited

Loading

havogt left a comment

havogt Dec 5, 2024

havogt Dec 5, 2024

havogt Dec 5, 2024

havogt Dec 5, 2024

havogt Dec 5, 2024

havogt Dec 5, 2024

havogt Dec 5, 2024

havogt Dec 5, 2024

havogt Dec 5, 2024

	if self.reinfer:
	# `self.dimensions` not available in reinference mode. Skip since we don't care anyway.
	if self.reinfer:

	Return `true` if the expr is a trivial expression or tuple thereof.
	Return `true` if the expr is a trivial expression (`SymRef` or `Literal`) or tuple thereof.

feature[next]: Non-tree-size-increasing collapse tuple on ifs #1762

Are you sure you want to change the base?

feature[next]: Non-tree-size-increasing collapse tuple on ifs #1762

Conversation

tehrengruber commented Dec 1, 2024 • edited Loading

tehrengruber Dec 5, 2024 • edited Loading

Choose a reason for hiding this comment

havogt left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

tehrengruber commented Dec 1, 2024 •

edited

Loading

tehrengruber Dec 5, 2024 •

edited

Loading