perf[next-dace]: RemoveScalarCopies and FuseHorizontalConditionBlocks transformations #2469

iomaganaris · 2026-02-02T16:19:21Z

Added transformations related to Graupel granule optimization
Main point of this PR is the FuseHorizontalConditionBlocks transformation that groups together ConditionalBlocks that share the same condition (__cond)
- To make sure that this can happen we have to remove the aliasing scalar AccessNodes so that we can find which ConditionalBlocks correspond to the same condition whose value is saved to an AccessNode
TODOs
- Merge first feat[next-dace]: Enable setting gpu_maxnreg attribute in maps #2464
- Figure out if there's any regression in greenline

philip-paul-mueller

It looks like a lot but most things are actually small.

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

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

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

...t4py/next/program_processors/runners/dace/transformations/fuse_horizontal_conditionblocks.py

…_name to gtx_transformations utils

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

…l_group_ifs

…emoval of conditional block

…d sdfgs

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

philip-paul-mueller

There are some things that need some modifications.

...t4py/next/program_processors/runners/dace/transformations/fuse_horizontal_conditionblocks.py

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

...t4py/next/program_processors/runners/dace/transformations/fuse_horizontal_conditionblocks.py

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

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

philip-paul-mueller · 2026-02-10T09:47:52Z

...t4py/next/program_processors/runners/dace/transformations/fuse_horizontal_conditionblocks.py

+        first_conditional_block_state_names = [
+            state.name for state in first_conditional_block.all_states()
+        ]
+        second_conditional_block_state_names = [
+            state.name for state in second_conditional_block.all_states()
+        ]
+        if not (
+            any("true_branch" in name for name in first_conditional_block_state_names)
+            and any("false_branch" in name for name in first_conditional_block_state_names)
+            and any("true_branch" in name for name in second_conditional_block_state_names)
+            and any("false_branch" in name for name in second_conditional_block_state_names)
+        ):
+            return False


I think this is kind of okay, but I would look at the conditions this would make it possible that there are multiple states inside the region.
However, I fine with a TODO in that regard.

I don't think atm there's any possibility to get multiple states that are not true/false_branch so this check is mostly for sanity

Then a TODO should be sufficient.

I didn't understand, would you like to remove this check or would you like to check if there can be more that true/false branches inside?
The latter is checked above:

if not ( isinstance(first_conditional_block, dace.sdfg.state.ConditionalBlock) and len(first_conditional_block.sub_regions()) == 2 and isinstance(second_conditional_block, dace.sdfg.state.ConditionalBlock) and len(second_conditional_block.sub_regions()) == 2 ):

philip-paul-mueller · 2026-02-10T09:50:35Z

...t4py/next/program_processors/runners/dace/transformations/fuse_horizontal_conditionblocks.py

+        ):
+            return False
+
+        # TODO(iomaganaris): Currently we do not handle NestedSDFGs inside the conditional blocks


I am not 100% sure, but very positive that you can remove this restriction, because of the check about the compatibility of the symbol mapping.

I would need to add a test though then and I'm not sure it's worth the time 🙈

I am fine without a test, because a Tasklet is, on the DF level a self contained object, as such it it not different from a Tasklet.

...t4py/next/program_processors/runners/dace/transformations/fuse_horizontal_conditionblocks.py

philip-paul-mueller · 2026-02-10T12:31:16Z

...t4py/next/program_processors/runners/dace/transformations/fuse_horizontal_conditionblocks.py

+        second_to_first_connections: dict[str, str] = {}
+        for node in nodes_renamed_map:
+            if isinstance(node, dace_nodes.AccessNode):
+                second_to_first_connections[node.data] = nodes_renamed_map[node].data
+        for first_inner_state in first_conditional_block.all_states():
+            corresponding_state_in_second = corresponding_states_first_to_second[first_inner_state]
+            for edge in second_state_edges_to_add[corresponding_state_in_second]:
+                new_memlet = copy.deepcopy(edge.data)
+                if edge.data.data in second_to_first_connections:
+                    new_memlet.data = second_to_first_connections[edge.data.data]
+                first_inner_state.add_edge(
+                    nodes_renamed_map[edge.src],
+                    nodes_renamed_map[edge.src].data
+                    if isinstance(edge.src, dace_nodes.AccessNode) and edge.src_conn
+                    else edge.src_conn,
+                    nodes_renamed_map[edge.dst],
+                    second_to_first_connections[edge.dst.data]
+                    if isinstance(edge.dst, dace_nodes.AccessNode) and edge.dst_conn
+                    else edge.dst_conn,
+                    new_memlet,
+                )


Suggested change

second_to_first_connections: dict[str, str] = {}

for node in nodes_renamed_map:

if isinstance(node, dace_nodes.AccessNode):

second_to_first_connections[node.data] = nodes_renamed_map[node].data

for first_inner_state in first_conditional_block.all_states():

corresponding_state_in_second = corresponding_states_first_to_second[first_inner_state]

for edge in second_state_edges_to_add[corresponding_state_in_second]:

new_memlet = copy.deepcopy(edge.data)

if edge.data.data in second_to_first_connections:

new_memlet.data = second_to_first_connections[edge.data.data]

first_inner_state.add_edge(

nodes_renamed_map[edge.src],

nodes_renamed_map[edge.src].data

if isinstance(edge.src, dace_nodes.AccessNode) and edge.src_conn

else edge.src_conn,

nodes_renamed_map[edge.dst],

second_to_first_connections[edge.dst.data]

if isinstance(edge.dst, dace_nodes.AccessNode) and edge.dst_conn

else edge.dst_conn,

new_memlet,

)

for edge_to_copy in edges_to_copy:

new_edge = corresponding_state_in_second.add_edge(

nodes_renamed_map[edge_to_copy.src],

edge_to_copy.src_conn,

nodes_renamed_map[edge_to_copy.dst],

edge_to_copy.dst_conn,

edge_to_copy.data,

)

if (not new_edge.data.is_empty()) and new_edge.data.data in second_arrays_rename_map:

new_edge.data.data = second_arrays_rename_map[new_edge.data.data]

However, this does not handle symbol renaming in the subsets.
See also comment above.

The subsets can only have symbols or something else as well?

philip-paul-mueller · 2026-02-10T12:39:26Z

...t4py/next/program_processors/runners/dace/transformations/fuse_horizontal_conditionblocks.py

+            second_arrays_rename_map[data_name] = new_data_name
+            data_desc_renamed = copy.deepcopy(data_desc)
+            first_cb.sdfg.add_datadesc(new_data_name, data_desc_renamed)
+


There is something missing here.
While in can_be_applied() you check if there are no symbol conflicts you never copy the symbols that are exclusive to the second conditional block.
Put it differently you detect the case where SDFG 1 has n := a + 1 and SDFG 2 has n := a - 1, but you do not handle the case where SDFG 1 has m := a + 1 and SDFG 2 has o := a + 1.
So you must copy them as well, for this you have to do something like:

missing_symbols = {sym2: val2 for sym2, val2 in second_cb.symbol_mapping.items() if sym2 not in first_cb.symbol_mapping} for missing_symb, symb_def in missing_symbols.items(): first_cb.symbol_mapping[missing_symb] = symb_def first_cb.add_symbol(missing_symb, second_cb.sdfg.symbols[missing_symb], find_new_name=False)

I have to figure out how to add this to a test 👍

...t4py/next/program_processors/runners/dace/transformations/fuse_horizontal_conditionblocks.py

philip-paul-mueller · 2026-02-10T12:42:08Z

...t4py/next/program_processors/runners/dace/transformations/fuse_horizontal_conditionblocks.py

+        for node in nodes_renamed_map:
+            if isinstance(node, dace_nodes.AccessNode):
+                second_to_first_connections[node.data] = nodes_renamed_map[node].data
+        for first_inner_state in first_conditional_block.all_states():


Suggested change

for node in nodes_renamed_map:

if isinstance(node, dace_nodes.AccessNode):

second_to_first_connections[node.data] = nodes_renamed_map[node].data

for first_inner_state in first_conditional_block.all_states():

I am pretty sure this information is already contained in second_arrays_rename_map.

I need to keep the reference to the node object though so I can find the corresponding new node in

new_edge = first_inner_state.add_edge( nodes_renamed_map[edge_to_copy.src], edge_to_copy.src_conn, nodes_renamed_map[edge_to_copy.dst], edge_to_copy.dst_conn, edge_to_copy.data, )

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

iomaganaris requested review from edopao and philip-paul-mueller February 2, 2026 16:19

iomaganaris changed the base branch from muphys_staging to set_gpu_maxnreg February 2, 2026 16:19

philip-paul-mueller reviewed Feb 3, 2026

View reviewed changes

iomaganaris and others added 19 commits February 3, 2026 09:33

Enable maxnreg setting

923e3cd

Added test for gpu_maxnreg

0d2b3e1

Fix formating

215e775

Mention that maxnreg takes precedence over launch bounds

d58ade9

Remove scalar copies

956282f

Rename if statements in _make_if_block

03c8cbb

[WIP] Added fuse condition block transformation

e10b2ee

[WIP] Copy access nodes between condition blocks

a2029f3

Fix branch selection

8961d5e

Moved kill aliasing transformation and updates to the fuse conditionals

49b7619

Rename properly second map arrays

eb462f6

Enable if grouping in proper place

73dc4c1

Rename kill aliasing scalars to remove

08ce68c

Cleared a bit FuseConditionalBlocks

e110ab2

Fixed imports

15ae54d

Remove sdfg.save

8f44eb5

Fix issues after cherry-picking

6a8b1b4

Applied suggestions to FuseHorizontalConditionBlocks and moved unique…

0ce4337

…_name to gtx_transformations utils

Apply review comments on RemoveScalarCopies

3319c00

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

iomaganaris force-pushed the graupel_group_ifs_updated_main branch from 2e18f65 to 3319c00 Compare February 3, 2026 13:37

iomaganaris added 2 commits February 3, 2026 14:37

Remove _dacegraphs

202aac6

Merge remote-tracking branch 'origin/set_gpu_maxnreg' into new_graupe…

0cf2be2

…l_group_ifs

iomaganaris changed the title ~~perf[next-dace]: RemoveAliasingScalars and FuseHorizontalConditionBlocks transformations~~ perf[next-dace]: RemoveScalarCopies and FuseHorizontalConditionBlocks transformations Feb 3, 2026

iomaganaris added 3 commits February 3, 2026 14:52

Added more elaborative comment in FuseHorizontalConditionBlocks for r…

9e6671c

…emoval of conditional block

make sure that the symbol mapping is the same between the fused neste…

c55a04a

…d sdfgs

Added check for NestedSDFGs in FuseHorizontalConditionBlocks

d232b73

iomaganaris and others added 2 commits February 3, 2026 15:32

Handle single use data check only for second AccessNode in RemoveScal…

ed7835d

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

Fix issues with true/false_branch_x_x_x

678f18f

edopao mentioned this pull request Feb 4, 2026

feat[next]: Muphys staging #2462

Draft

philip-paul-mueller reviewed Feb 5, 2026

View reviewed changes

iomaganaris commented Feb 5, 2026

View reviewed changes

...t4py/next/program_processors/runners/dace/transformations/fuse_horizontal_conditionblocks.py Show resolved Hide resolved

iomaganaris commented Feb 5, 2026

View reviewed changes

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

Address Philip's comments

562742b

iomaganaris marked this pull request as ready for review February 10, 2026 09:41

Merge branch 'set_gpu_maxnreg' into graupel_group_ifs_updated_main

d237219

iomaganaris requested a review from philip-paul-mueller February 10, 2026 09:42

philip-paul-mueller reviewed Feb 10, 2026

View reviewed changes

iomaganaris and others added 3 commits February 10, 2026 14:08

Fix issues

636f4c5

Address Philip's comments

caf69a3

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

Removed check for nestedsdfg

4f03b21

perf[next-dace]: RemoveScalarCopies and FuseHorizontalConditionBlocks transformations #2469

Are you sure you want to change the base?

perf[next-dace]: RemoveScalarCopies and FuseHorizontalConditionBlocks transformations #2469

Uh oh!

Conversation

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

Uh oh!

Uh oh!

philip-paul-mueller left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

philip-paul-mueller left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

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!

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!

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!

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!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

iomaganaris commented Feb 2, 2026 •

edited

Loading