perf[next-dace]: Enhance LoopBlocking pass by iomaganaris · Pull Request #2578 · GridTools/gt4py

iomaganaris · 2026-04-17T17:50:43Z

Added the following features to the LoopBlocking pass:

Memlets that have independent data are promoted in the outer Map so that they only have to be read once
Added option blocking_independent_node_threshold. LoopBlocking is only going to be applied to Maps that have more than the threshold number independent variables

philip-paul-mueller

I like the changes but there some things that needs to be addressed.

philip-paul-mueller · 2026-04-19T10:37:53Z

+                            "Independent memlets should only be inputs to maps that have a single parameter. "
+                            "Those should always be neighbor reductions."
+                        )
+                        edge.data.subset = next(iter(original_dst_of_in_edge.params))


It is correct that you have to update subset here, but does does not make much sense to me.

Not sure I understood your comment. Should I change something?

I mean original_dst_of_in_edge is a MapEntry and it does not have an attribute called params, the Map has this, i.e. original_dst_of_in_edge.map.params would exists.
Map::params stores the iteration variables (ordered accordingly to this function).
However, now you take the first iteration variable, which at this point should be horizontal dimension.
So I do not understand the logic that is applied here.

original_dst_of_in_edge is a MapEntry indeed but it has a params attribute.
You're right that next(iter(original_dst_of_in_edge.params)) is wrong. Instead it should be I believe:

edge.data.subset = dace_subsets.Range.from_indices(original_dst_of_in_edge.params)

Maybe I miss something, but here you modifying the (downstream) Memlets such that instead reading from the global data they read from promoted_anode.
You do this because you know that the read (expressed by the Memlet you promoted) does not depend on K.
However, has far as I understand original_dst_of_in_edge.params contains K?

original_dst_of_in_edge is the dst of the edge before any changes. In case this is a Map, I have the impression that based on our lowering it can only be a neighbor reduction. In case it's a neighbor reduction it shouldn't have as parameter a vertical or horizontal index. I think in general in our lowering there should be a Map inside another Map that refers to vertical or horizontal indexes. In case there is such map it should be one that was created by the LoopBlocking pass. However we only call this pass once and even if we called it more than once then the inner Map would be dependent and also after the latest changes I added, the edges pointing to such map shouldn't be considered for promotion.
Let me know if the above is right and makes sense

philip-paul-mueller · 2026-04-19T10:40:36Z

+            elif isinstance(original_dst_of_in_edge, dace_nodes.NestedSDFG):
+                raise NotImplementedError("Promotion of memlets to NestedSDFG not implemented yet.")
+            elif isinstance(original_dst_of_in_edge, dace_nodes.LibraryNode):
+                raise NotImplementedError(
+                    "Promotion of memlets to LibraryNode not implemented yet."
+                )


I think you can remove them, since original_dst_of_in_edge is always classified as a dependent node.

Not sure why I could remove these checks. Couldn't we have a dependent NestedSDFG or LibraryNode?

I am not fully sure what I meant with it.
But thinking about it now, I think the can_be_applied() should make sure that you do not hit that case.

I think that in that case we actually have to implement the handling of those cases instead of silently avoid applying the transformation but I haven't seen any of these cases yet so that's why I left it as it is 🙈

philip-paul-mueller · 2026-04-21T08:22:26Z

+            for subset_range in in_edge.data.subset.ranges:
+                if subset_range not in independent_outer_map_as_range.ranges:
+                    new_subset.append(subset_range)


This looks a bit brittle, because you look at sizes.
However, currently I do not have a better idea, beside looking at the subset and checking if it contains the blocking parameter, but I am not sure if this is better.

In case we have parameters in the ranges I think we actually check the parameters and not the sizes. For example:

(Pdb) p in_edge.data.subset.ranges [(__i0, __i0, 1), (0, 2, 1)] (Pdb) p independent_outer_map_as_range.ranges [(__i0, __i0, 1)] (Pdb) p new_subset [(0, 2, 1)]

philip-paul-mueller · 2026-04-22T07:07:48Z

+            )
+
+            if isinstance(original_dst_of_in_edge, dace_nodes.MapEntry):
+                for edge in state.out_edges(original_dst_of_in_edge):


I did not saw this before, but iterating over the out edges is not enough, as there could be nested Maps.
I think there is something in utils, the reroute or so that can help you do it or at least give you some hints on how to do it.

If isinstance(edge.dst, dace_nodes.MapEntry) would it help if we called

dace_sdutils.canonicalize_memlet_trees_for_map(state=state, map_node=edge.dst) dace_propagation.propagate_memlets_map_scope(sdfg, state, edge.dst)

?

No.
What I meant is that one level of Memlets might not be enough, as they could "continue".
Thus you must use something like here, i.e. you iterate over the Memlet tree.
Ideally you could use reconfigure_dataflow_after_rerouting(), however, that function can only handle simple shifts.
You do not have simple shifts because in line 726 you conditionally add.
What you could do is adding dummy dimensions, i.e. size 0 dimensions.

I tried traversing the memlet tree as you mentioned. Hopefully that should be enough? 🙈 Even if there some indirection using another AccessNode, based on what I've seen until now, accessing this AccessNode will be done using only local indexes which we don't touch

…be_applied

…contains the blocking parameter

…around warnings

…king

philip-paul-mueller

Some additional comments.

philip-paul-mueller · 2026-04-30T06:37:32Z

 import enum
 import warnings
-from typing import Any, Callable, Optional, Sequence, TypeAlias, Union
+from typing import Any, Callable, List, Optional, Sequence, TypeAlias, Union


Suggested change

from typing import Any, Callable, List, Optional, Sequence, TypeAlias, Union

from typing import Any, Callable, Optional, Sequence, TypeAlias, Union

philip-paul-mueller · 2026-04-30T06:37:51Z

    gpu_block_size_3d: Optional[Sequence[int | str] | str] = None,
    gpu_maxnreg: Optional[int] = None,
-    blocking_dim: Optional[gtx_common.Dimension] = None,
+    blocking_dims: Optional[List[gtx_common.Dimension]] = None,


Suggested change

blocking_dims: Optional[List[gtx_common.Dimension]] = None,

blocking_dims: Optional[Sequence[gtx_common.Dimension]] = None,

philip-paul-mueller · 2026-04-30T06:39:03Z


        gpu_map.gpu_block_size = tuple(block_size)
-        if self.maxnreg is not None:
+        if self.maxnreg is not None and gpu_map.gpu_maxnreg == 0:


Just for my curiosity, what is the intention behind this change?

philip-paul-mueller · 2026-04-30T06:40:03Z

-        dtype=str,
+    blocking_parameters = dace_properties.Property(
+        dtype=list,
        allow_none=True,


Suggested change

allow_none=True,

blocking_parameters = dace_properties.ListProperty(

dtype=str,

philip-paul-mueller · 2026-04-30T06:40:34Z

@@ -52,6 +54,8 @@ class LoopBlocking(dace_transformation.SingleStateTransformation):
        blocking_parameter: On which parameter should we block.


Needs to be updated as well.

philip-paul-mueller · 2026-04-30T06:59:35Z

+
+        _ = sdfg.reset_cfg_list()
+        dace_sdutils.canonicalize_memlet_trees_for_map(state=state, map_node=outer_map_entry)
+        dace_propagation.propagate_memlets_map_scope(sdfg, state, outer_map_entry)


I am not sure if Memlet propagation is needed here.

philip-paul-mueller · 2026-04-30T07:00:21Z

+
+        self._populate_memlet_to_promote(matched_blocking_var, state, outer_map_entry)
+        # Below checks are necessary for MyPy
+        if self._memlet_to_promote and len(self._memlet_to_promote) == 0:


Suggested change

if self._memlet_to_promote and len(self._memlet_to_promote) == 0:

if len(self._memlet_to_promote) == 0:

philip-paul-mueller · 2026-04-30T07:01:12Z

+
+        for in_edge in self._memlet_to_promote:
+            if isinstance(in_edge.dst, dace_nodes.AccessNode):
+                raise NotImplementedError(


This case should be rejected already in the populate_memlet_to_promote() function.

philip-paul-mueller · 2026-04-30T07:04:09Z

+            corresponding_inner_map_out_edges = list(
+                state.out_edges_by_connector(in_edge.dst, "OUT_" + in_edge.dst_conn[3:])
+            )


You only have connectors starting with IN_ at Map nodes.
It can even be None.
So doing this operation, dst_conn[3:], unprotected by an if isinstance(node, MapEntry) is most certainly wrong.

This function should probably look like:

inode = add_new_intermediate_storage() add_edge(outer_map, inode, modified_memlet1()) e = add_edge(inode, original_edge.dst, dcpy(original_edge.data)) for mtree in state.memlet_tree(e).traverse(include_self=True): apply_correction(mtree.edge)

I think that the best idea would be to copy the logic of MapFusionVertical::compute_reduced_intermediate().

philip-paul-mueller · 2026-04-30T11:48:21Z

+            return
+        assert self._memlet_to_promote is not None
+
+        for in_edge in self._memlet_to_promote:


I took the liberty and tried to come up with something different: iomaganaris#2
Feel free to modify.

iomaganaris and others added 25 commits November 21, 2025 16:18

Fix k blocking for non independent option

32476d7

PoC for moving temporary outside kloop

8547aba

Working version

e265480

Enable maxnreg setting

0cb5f5b

Make sure that the inner loop doesn't get unrolled

dcc74d4

Enable promote_independent_memlets option

d0d567f

Enable loop blocking if there's any independent memlet to promote

ec11475

Added initial test and fixed pass

0846ba8

Fixing test

8c7ce09

If we don't require_independent_nodes always apply

2ac4aa7

Fix most of the tests

b8686f9

Fix tests

21f053b

Fix memlet promotion number

b80dc12

Added option for independent node promotion threshold

56e390f

Extended the tests and fixes

93e84ab

Skip maps with single sizes

0305aaa

Make formatting happy

da784a2

Set better block size and gpu_maxnreg for kblocking

329803d

Make formatting happy

cf87306

Improve some hacks for subsets

b1cce65

Merge remote-tracking branch 'origin/main' into extend_loopblocking

0a12898

Merge remote-tracking branch 'origin/main' into extend_loopblocking

729dca7

Fix unique_name call

180d777

Merge remote-tracking branch 'origin/main' into extend_loopblocking

5391431

Remove NVIDIA related options for the loop in kblocking

73a8ac0

iomaganaris requested review from edopao and philip-paul-mueller April 17, 2026 17:50

philip-paul-mueller reviewed Apr 19, 2026

View reviewed changes

iomaganaris added 2 commits April 19, 2026 18:40

Don't change the maxnreg of a map if it's already set

87fc4bf

Handling comments from Philip

eabbbb1

philip-paul-mueller reviewed Apr 21, 2026

View reviewed changes

Comment thread src/gt4py/next/program_processors/runners/dace/transformations/loop_blocking.py Outdated

philip-paul-mueller reviewed Apr 22, 2026

View reviewed changes

Comment thread src/gt4py/next/program_processors/runners/dace/transformations/loop_blocking.py Outdated

iomaganaris and others added 11 commits April 27, 2026 15:44

Extend test and address more comments

7e07bff

Avoid calling self.outer_entry in functions apart from apply and can_…

dbb2ace

…be_applied

Don't promote independent memlets if their destination is a map that …

ca13254

…contains the blocking parameter

Handle better adjustment of edges in inner map

794ff4a

Add more comments for handling the out edges from inner map and move …

c147920

…around warnings

Fix how subsets are set to memlets

bdb24d0

Fixed unit test with empty memlet

197a4f1

Instead of a single dimension pass a list of dimensions for loop bloc…

f77078e

…king

Use unroll length same as blocking size

8488fb9

Avoid loop blocking for maps that have scans

2328a2d

Fix disabling normal loop blocking for scans

d2947cc

philip-paul-mueller force-pushed the extend_loopblocking branch from 37a6568 to d2947cc Compare April 30, 2026 11:46

philip-paul-mueller reviewed Apr 30, 2026

View reviewed changes

	from typing import Any, Callable, List, Optional, Sequence, TypeAlias, Union
	from typing import Any, Callable, Optional, Sequence, TypeAlias, Union

	blocking_dims: Optional[List[gtx_common.Dimension]] = None,
	blocking_dims: Optional[Sequence[gtx_common.Dimension]] = None,

	allow_none=True,
	blocking_parameters = dace_properties.ListProperty(
	dtype=str,

		@@ -52,6 +54,8 @@ class LoopBlocking(dace_transformation.SingleStateTransformation):
		blocking_parameter: On which parameter should we block.

	if self._memlet_to_promote and len(self._memlet_to_promote) == 0:
	if len(self._memlet_to_promote) == 0:

Conversation

iomaganaris commented Apr 17, 2026

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!

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!

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!

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!

philip-paul-mueller left a comment

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!

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!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants