fix map fusion and loop blocking for apply_diffusion_to_vn

GridTools · Feb 10, 2025 · 3626e17 · 3626e17
1 parent 4bfa5b0
commit 3626e17
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 3 deletions.
diff --git a/src/gt4py/next/program_processors/runners/dace/transformations/loop_blocking.py b/src/gt4py/next/program_processors/runners/dace/transformations/loop_blocking.py
@@ -380,6 +380,14 @@ def _classify_node(
             ):
                 return False
 
+        elif isinstance(node_to_classify, dace.nodes.NestedSDFG):
+            # Similar checks apply to a nested SDFG node
+            if not all(
+                isinstance(out_edge.dst, dace_nodes.AccessNode)
+                for out_edge in state.out_edges(node_to_classify)
+            ):
+                return False
+
         elif isinstance(node_to_classify, dace_nodes.AccessNode):
             # AccessNodes need to have some special properties.
             node_desc: dace.data.Data = node_to_classify.desc(sdfg)

diff --git a/src/gt4py/next/program_processors/runners/dace/transformations/map_fusion_serial.py b/src/gt4py/next/program_processors/runners/dace/transformations/map_fusion_serial.py
@@ -589,9 +589,6 @@ def partition_first_outputs(
                 ):
                     if inner_consumer_edge.data.src_subset is None:
                         return None
-                    if inner_consumer_edge.data.dynamic:
-                        # TODO(phimuell): Is this restriction necessary, I am not sure.
-                        return None
                     consumer_subsets.append(inner_consumer_edge.data.src_subset)
             assert (
                 found_second_map