perf[next-dace]: Enhance MoveDataflowIntoIfBody transformation by iomaganaris · Pull Request #2514 · GridTools/gt4py

iomaganaris · 2026-03-05T15:50:53Z

Move nodes into the state of a ConditionalBlock even if they are marked by multiple Connectors

Needs:

Rebase on latest main

edopao

LGTM, just one minor comment.

src/gt4py/next/program_processors/runners/dace/transformations/move_dataflow_into_if_body.py

edopao

LGTM

philip-paul-mueller

See my comment.

src/gt4py/next/program_processors/runners/dace/transformations/move_dataflow_into_if_body.py

philip-paul-mueller

I have now looked at the changes and I have some suggestions.
A part of my comments are "adding todos for me" and suggestions about improving the comments, so it is not as much as it looks like.

My main concern is that some nodes are replicated (added to the nested state) multiple times.
However, I am not fully sure why the unit test passes, it is probably that I do not see the part that does it.

src/gt4py/next/program_processors/runners/dace/transformations/move_dataflow_into_if_body.py

philip-paul-mueller · 2026-03-09T06:11:08Z

src/gt4py/next/program_processors/runners/dace/transformations/move_dataflow_into_if_body.py

+        # Gather all the already moved nodes to avoid that we move the same node multiple times
+        already_moved_nodes: set[dace_nodes.Node] = set()
        # Finally relocate the dataflow


Suggested change

# Gather all the already moved nodes to avoid that we move the same node multiple times

already_moved_nodes: set[dace_nodes.Node] = set()

# Finally relocate the dataflow

# Relocate the dataflow, because the node sets listed in `relocatable_dataflow` are not disjoint we have to make sure that we relocate the nodes only once.

already_moved_nodes: set[dace_nodes.Node] = set()

src/gt4py/next/program_processors/runners/dace/transformations/move_dataflow_into_if_body.py

philip-paul-mueller · 2026-03-09T06:15:38Z

src/gt4py/next/program_processors/runners/dace/transformations/move_dataflow_into_if_body.py

        sdfg: dace.SDFG,
    ) -> None:
        if_block: dace_nodes.NestedSDFG = self.if_block
        if_block_spec = self._partition_if_block(if_block)


I kind of have the feeling that this function should compute the mapping you compute in conn_name_to_access_node_map since it has to look at all nodes anyway.

However, I think a TODO in _partition_if_block()'s docstring should be sufficient.

It needs the relocatable_dataflow though so not sure if it's possible to do that there unless we do it for every node

Not necessarily.
Having relocatable_dataflow would just allow you skip some instances that are not needed, but I think it is still a net gain if we do it in one go and then reuse, because we have to look at every node anyway.
But it is for sure something for later.
I think I will create a PR that addresses this and some indeterminism.

src/gt4py/next/program_processors/runners/dace/transformations/move_dataflow_into_if_body.py

philip-paul-mueller · 2026-03-09T07:34:27Z

src/gt4py/next/program_processors/runners/dace/transformations/move_dataflow_into_if_body.py

Are you sure that here no filtering is needed?
This was possible because nodes_to_move where disjoint, which is no longer the case, thus I would expect that there is some filtering needed, when you create the nodes inside the branch, but not when you create the new_nodes mapping.
The same applies to where you create the new data descriptors.

I have added the old_to_new_nodes_map that takes care of copying the access nodes per branch only once

src/gt4py/next/program_processors/runners/dace/transformations/move_dataflow_into_if_body.py

philip-paul-mueller · 2026-03-09T07:42:13Z

src/gt4py/next/program_processors/runners/dace/transformations/move_dataflow_into_if_body.py

+                        existing_connector = conn_name_to_access_node_map[oedge.dst_conn][1].data
+                        if not inner_sdfg.arrays[existing_connector].transient:
+                            inner_sdfg.arrays[existing_connector].transient = True
+                            if_block.remove_in_connector(existing_connector)


I do not like this.
It is not "uniform" meaning now you remove the in connector in two different locations.
I would prefer if it only at one location, have you tried to just remove it here and then play the other thing out.

I think it would actually complicate things if we removed it elsewhere because we would have to keep track of which connectors have to be removed from here or look for which connectors don't have any edges assigned to them

Some changes were requested by Philip, removing my approval.

Co-authored-by: Philip Müller <philip.mueller@cscs.ch>

src/gt4py/next/program_processors/runners/dace/transformations/move_dataflow_into_if_body.py

philip-paul-mueller

LGTM

src/gt4py/next/program_processors/runners/dace/transformations/move_dataflow_into_if_body.py

philip-paul-mueller · 2026-03-09T14:07:49Z

src/gt4py/next/program_processors/runners/dace/transformations/move_dataflow_into_if_body.py

+                        ][1].data
+                        if not inner_sdfg.arrays[inner_access_node_of_connector_name].transient:
+                            inner_sdfg.arrays[inner_access_node_of_connector_name].transient = True
+                            if_block.remove_in_connector(inner_access_node_of_connector_name)


Have you tried to remove just this line.
I think it should still work, because this function is called for the connector anyway and at the end it is removed.
//Currently remove_in_connector() does not check if the connector exits or not and ALWAYS returns True.

This doesn't work because in the following case __arg1 will have no relocated_nodes due to the filter_nodes function that checks if all the outputs of the dataflow go to other nodes on the same dataflow

Actually I had a closer look and figure out that if we add another tasklet after the a2 AccessNode then the transformation doesn't move all the AccessNodes and Tasklets it could into the if. See the following SDFG.
The blocker is that _partition_if_block expects in the states of the ConditionalBlocks only nodes that are AccessNodes that have the names of the in/out connectors of the ConditionalBlock NestedSDFG.
Then _check_for_data_and_symbol_conflicts fails because it finds a2 inside the NestedSDFG.
I think the first blocker can be easily handled by relaxing the requirements. The second one however requires more work because we have to rename the AccessNodes in case we have to move them inside and there's an existing AccessNode with the same name. I think that this situation may arise only if we try to move inside the NestedSDFG of the ConditionalBlock an AccessNode that has the same name as one of the input Connectors and thus an AccessNode inside the NestedSDFG as well.
Currently thankfully we don't come across this case in the graupel case so we can avoid fixing this but I don't know if we should actually take care of this in this PR

src/gt4py/next/program_processors/runners/dace/transformations/move_dataflow_into_if_body.py

iomaganaris · 2026-03-10T12:55:34Z

I discussed with @philip-paul-mueller to get this merged and the added test with xfail added to this PR can be handled in a subsequent PR

iomaganaris requested a review from edopao March 5, 2026 15:50

edopao reviewed Mar 6, 2026

View reviewed changes

src/gt4py/next/program_processors/runners/dace/transformations/move_dataflow_into_if_body.py Outdated Show resolved Hide resolved

src/gt4py/next/program_processors/runners/dace/transformations/move_dataflow_into_if_body.py Outdated Show resolved Hide resolved

iomaganaris added 6 commits March 6, 2026 14:43

Add failing test

d0e3ae4

Fix test and pass

b90393a

Remove commented code and added explanations

d7d05bd

Make formatting happy

2b7a501

We can also have only a false/true branch instead of both

cd6f596

Address review comments

3ab44a0

iomaganaris force-pushed the enhance_move_dataflow_into_if branch from 0cfd5fb to 3ab44a0 Compare March 6, 2026 13:56

iomaganaris marked this pull request as ready for review March 6, 2026 13:57

edopao previously approved these changes Mar 6, 2026

View reviewed changes

philip-paul-mueller requested changes Mar 9, 2026

View reviewed changes

src/gt4py/next/program_processors/runners/dace/transformations/move_dataflow_into_if_body.py Outdated Show resolved Hide resolved

philip-paul-mueller reviewed Mar 9, 2026

View reviewed changes

iomaganaris and others added 4 commits March 9, 2026 13:26

Improve new node creation

16926ce

Addressed Philip's comments

89bf522

Co-authored-by: Philip Müller <philip.mueller@cscs.ch>

Merge branch 'main' into enhance_move_dataflow_into_if

4c54704

Updated docstring

a609bc5

philip-paul-mueller reviewed Mar 9, 2026

View reviewed changes

src/gt4py/next/program_processors/runners/dace/transformations/move_dataflow_into_if_body.py Show resolved Hide resolved

philip-paul-mueller approved these changes Mar 9, 2026

View reviewed changes

iomaganaris added 2 commits March 9, 2026 18:37

Handle remaining comments

af5051d

Added test with xfail that shows restriction of current transformation

8eadaf9

iomaganaris merged commit b088235 into main Mar 10, 2026
23 checks passed

philip-paul-mueller mentioned this pull request Mar 26, 2026

feat[next-dace]: Updated MoveDataflowIntoIfBody #2531

Open

1 task

Conversation

iomaganaris commented Mar 5, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

edopao left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

edopao left a comment

Choose a reason for hiding this comment

Uh oh!

philip-paul-mueller left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

philip-paul-mueller left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

philip-paul-mueller left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

iomaganaris commented Mar 10, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

iomaganaris commented Mar 5, 2026 •

edited

Loading