perf[next-dace]: cheap checks before is_parallel in horizontal map fusion

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

@@ -14,6 +14,7 @@
 import dace
 from dace import nodes as dace_nodes, properties as dace_properties
 from dace.transformation import dataflow as dace_dftrans
+from dace.transformation.dataflow import map_fusion_helper as dace_mfhelper
 
 
 VerticalMapFusionCallback: TypeAlias = Callable[
@@ -132,9 +133,38 @@ def can_be_applied(
         sdfg: dace.SDFG,
         permissive: bool = False,
     ) -> bool:
+        first_map_entry = self.first_parallel_map_entry
+        second_map_entry = self.second_parallel_map_entry
         if self._check_fusion_callback is not None:
             if not self._check_fusion_callback(
-                self, self.first_parallel_map_entry, self.second_parallel_map_entry, graph, sdfg
+                self, first_map_entry, second_map_entry, graph, sdfg
             ):
                 return False
+
+        # The pattern matches every pair of Maps in the state. The base class performs these
+        #  checks too, but only after `is_parallel()`, which traverses the state twice. It also
+        #  rejects a contradictory scope selection, which has to keep happening.
+        if self.only_inner_maps and self.only_toplevel_maps:
+            raise ValueError(
+                "Only one of `only_inner_maps` and `only_toplevel_maps` is allowed per"
+                f" `{type(self).__name__}` instance."
+            )
+        if first_map_entry.map.schedule != second_map_entry.map.schedule:
+            return False
+        scope_dict = graph.scope_dict()
+        map_scope = scope_dict[first_map_entry]
+        if scope_dict[second_map_entry] != map_scope:
+            return False
+        if self.only_toplevel_maps and map_scope is not None:
+            return False
+        if self.only_inner_maps and map_scope is None:
+            return False
+        if (
+            dace_mfhelper.find_parameter_remapping(
+                first_map=first_map_entry.map, second_map=second_map_entry.map
+            )
+            is None
+        ):
+            return False
+
         return super().can_be_applied(graph, expr_index, sdfg, permissive)

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

@@ -442,39 +442,39 @@ def can_be_applied(
 
         # Ensure that the Maps are in the same scope.
         scope_dict = graph.scope_dict()
-        if scope_dict[self.first_map_entry] is not scope_dict[self.second_map_entry]:
+        if scope_dict[first_map_entry] is not scope_dict[second_map_entry]:
             return False
 
         # Test if the map is in the right scope.
-        map_scope: Union[dace_nodes.Node, None] = scope_dict[self.first_map_entry]
+        map_scope: Union[dace_nodes.Node, None] = scope_dict[first_map_entry]
         if self.only_toplevel_maps and (map_scope is not None):
             return False
 
         first_map_src_data = {
             iedge.src.label
-            for iedge in graph.in_edges(self.first_map_entry)
+            for iedge in graph.in_edges(first_map_entry)
             if isinstance(iedge.src, dace_nodes.AccessNode)
         }
         second_map_src_data = {
             iedge.src.label
-            for iedge in graph.in_edges(self.second_map_entry)
+            for iedge in graph.in_edges(second_map_entry)
             if isinstance(iedge.src, dace_nodes.AccessNode)
         }
 
-        # Test if the Maps are parallel.
-        if not dace_mfhelper.is_parallel(
-            graph=graph, node1=first_map_entry, node2=second_map_entry
-        ):
-            return False
-
+        # The pattern matches every pair of Maps in the state, so the cheap checks come
+        #  before `is_parallel()`, which traverses the state twice.
         if len(first_map_src_data.intersection(second_map_src_data)) == 0:
-            # no common source access node
             return False
 
         splitted_range = gtx_mfutils.split_overlapping_map_range(first_map, second_map)
         if splitted_range is None:
             return False
 
+        if not dace_mfhelper.is_parallel(
+            graph=graph, node1=first_map_entry, node2=second_map_entry
+        ):
+            return False
+
         return True
 
     def apply(

tests/next_tests/unit_tests/program_processor_tests/runners_tests/dace_tests/transformation_tests/test_horizontal_map_split_fusion.py

@@ -295,3 +295,17 @@ def test_horizontal_map_fusion(run_map_fusion: bool):
                         f"Found maps with overlapping ranges: {map_entry_i.label} and {map_entry_j.label} "
                         f"[{sdfg.state(0).in_edges(map_entry_i)[0].src.label}]"
                     )
+
+
+def test_horizontal_map_fusion_rejects_contradictory_map_scopes():
+    """`can_be_applied()` tests the scope itself, so it has to reject this like DaCe does."""
+    sdfg, _ = _make_sdfg_with_multiple_maps_that_share_inputs(10)
+
+    with pytest.raises(ValueError, match="only_inner_maps"):
+        sdfg.apply_transformations_repeated(
+            gtx_transformations.MapFusionHorizontal(
+                only_inner_maps=True,
+                only_toplevel_maps=True,
+            ),
+            validate=False,
+        )