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: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.

I thought that the two calls are probably needed to make sure we don't have any strange state of memlets from previous transformations

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:

I changed that to:

if not self._memlet_to_promote: return

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.

I left my friend take care of the following based on your changes. I hope it's okay now 😅 Same for the comments above

…proves a bit codegen for AMD

…transformation and setting the blocking size

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

philip-paul-mueller

I have some small comments, but in general it looks okay.

philip-paul-mueller · 2026-07-14T11:00:20Z

+
+    def _get_blocking_parameter(self, map_params: list[str]) -> Optional[str]:
+        if self.blocking_parameters is None:
+            return None


Actually I just realized that you were right and there should be an error here.
However, because of the guard in can_be_applied() we can condense it to assert self.blocking_parameters

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

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

iomaganaris added 2 commits June 10, 2026 15:28

Remove duplicate tmp access nodes outside the inner for loop which im…

a4fb67c

…proves a bit codegen for AMD

Added a comment for future work regarding enabling the loop blocking …

4e82812

…transformation and setting the blocking size

iomaganaris force-pushed the extend_loopblocking branch from 408149d to 4e82812 Compare June 10, 2026 14:57

iomaganaris and others added 4 commits July 10, 2026 18:42

Take care of Philip's comments with the help of a friend

22d3bbe

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

Merge branch 'main' into extend_loopblocking

6da0520

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

69aef7a

Fix issue with util.unique_name

cc112ad

philip-paul-mueller reviewed Jul 14, 2026

View reviewed changes

iomaganaris and others added 4 commits July 16, 2026 17:53

Take care of Philip's comments

a9322ef

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

Merge branch 'main' into extend_loopblocking

1cb5e91

Fix check for independent_nodes

72be0c2

Merge branch 'main' into extend_loopblocking

6386f66

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

Uh oh!

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!

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!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

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

Uh oh!

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!

Choose a reason for hiding this comment

Uh oh!

iomaganaris Jul 10, 2026 •

edited

Loading